Thanks to visit codestin.com
Credit goes to github.com

Skip to content

Recommend >= 0.2.0 of torch-tb-profiler #16361

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 4 commits into from
Jun 16, 2021
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
3 changes: 2 additions & 1 deletion CHANGELOG.md
Original file line number Diff line number Diff line change
Expand Up @@ -10,9 +10,10 @@
([#16102](https://github.com/Microsoft/vscode-python/issues/16102))
1. Add an `enumDescriptions` key under the `python.languageServer` setting to describe all language server options.
([#16141](https://github.com/Microsoft/vscode-python/issues/16141))
1. Ensure users upgrade to v0.2.0 of the torch-tb-profiler TensorBoard plugin to access jump-to-source functionality.
([#16330](https://github.com/Microsoft/vscode-python/issues/16330))

### Fixes

1. Fixes a bug in the bandit linter where messages weren't being propagated to the editor.
(thanks [Anthony Shaw](https://github.com/tonybaloney))
([#15561](https://github.com/Microsoft/vscode-python/issues/15561))
Expand Down
4 changes: 2 additions & 2 deletions src/client/common/utils/localize.ts
Original file line number Diff line number Diff line change
Expand Up @@ -177,11 +177,11 @@ export namespace TensorBoard {
);
export const installTensorBoardAndProfilerPluginPrompt = localize(
'TensorBoard.installTensorBoardAndProfilerPluginPrompt',
'TensorBoard >= 2.4.1 and the PyTorch Profiler TensorBoard plugin are required. Would you like to install these packages?',
'TensorBoard >= 2.4.1 and the PyTorch Profiler TensorBoard plugin >= 0.2.0 are required. Would you like to install these packages?',
);
export const installProfilerPluginPrompt = localize(
'TensorBoard.installProfilerPluginPrompt',
'We recommend installing the PyTorch Profiler TensorBoard plugin. Would you like to install the package?',
'We recommend installing version >= 0.2.0 of the PyTorch Profiler TensorBoard plugin. Would you like to install the package?',
);
export const upgradePrompt = localize(
'TensorBoard.upgradePrompt',
Expand Down
50 changes: 36 additions & 14 deletions src/client/tensorBoard/tensorBoardSession.ts
Original file line number Diff line number Diff line change
Expand Up @@ -50,6 +50,8 @@ import { ModuleInstallFlags } from '../common/installer/types';
enum Messages {
JumpToSource = 'jump_to_source',
}
const TensorBoardSemVerRequirement = '>= 2.4.1';
const TorchProfilerSemVerRequirement = '>= 0.2.0';

/**
* Manages the lifecycle of a TensorBoard session.
Expand Down Expand Up @@ -143,18 +145,21 @@ export class TensorBoardSession {

private async promptToInstall(
tensorBoardInstallStatus: ProductInstallStatus,
shouldInstallProfilerPlugin: boolean,
profilerPluginInstallStatus: ProductInstallStatus,
) {
sendTelemetryEvent(EventName.TENSORBOARD_INSTALL_PROMPT_SHOWN);
const yes = Common.bannerLabelYes();
const no = Common.bannerLabelNo();
const isUpgrade = tensorBoardInstallStatus === ProductInstallStatus.NeedsUpgrade;
let message;

if (tensorBoardInstallStatus === ProductInstallStatus.Installed && shouldInstallProfilerPlugin) {
if (
tensorBoardInstallStatus === ProductInstallStatus.Installed &&
profilerPluginInstallStatus !== ProductInstallStatus.Installed
) {
// PyTorch user already has TensorBoard, just ask if they want the profiler plugin
message = TensorBoard.installProfilerPluginPrompt();
} else if (shouldInstallProfilerPlugin) {
} else if (profilerPluginInstallStatus !== ProductInstallStatus.Installed) {
// PyTorch user doesn't have compatible TensorBoard or the profiler plugin
message = TensorBoard.installTensorBoardAndProfilerPluginPrompt();
} else if (isUpgrade) {
Expand Down Expand Up @@ -193,15 +198,19 @@ export class TensorBoardSession {

// First see what dependencies we're missing
let [tensorboardInstallStatus, profilerPluginInstallStatus] = await Promise.all([
this.installer.isProductVersionCompatible(Product.tensorboard, '>= 2.4.1', interpreter),
this.installer.isInstalled(Product.torchProfilerImportName, interpreter),
this.installer.isProductVersionCompatible(Product.tensorboard, TensorBoardSemVerRequirement, interpreter),
this.installer.isProductVersionCompatible(
Product.torchProfilerImportName,
TorchProfilerSemVerRequirement,
interpreter,
),
]);
const isTorchUserAndInExperiment = ImportTracker.hasModuleImport('torch') && this.isInTorchProfilerExperiment;
const needsTensorBoardInstall = tensorboardInstallStatus !== ProductInstallStatus.Installed;
const needsProfilerPluginInstall = isTorchUserAndInExperiment && profilerPluginInstallStatus !== true;
const needsProfilerPluginInstall = profilerPluginInstallStatus !== ProductInstallStatus.Installed;
if (
// PyTorch user, in profiler install experiment, TensorBoard and profiler plugin already installed
(isTorchUserAndInExperiment && !needsTensorBoardInstall && profilerPluginInstallStatus === true) ||
(isTorchUserAndInExperiment && !needsTensorBoardInstall && !needsProfilerPluginInstall) ||
// Not PyTorch user or not in profiler install experiment, so no need for profiler plugin,
// and TensorBoard is already installed
(!isTorchUserAndInExperiment && tensorboardInstallStatus === ProductInstallStatus.Installed)
Expand All @@ -212,7 +221,7 @@ export class TensorBoardSession {
// Ask the user if they want to install packages to start a TensorBoard session
const selection = await this.promptToInstall(
tensorboardInstallStatus,
isTorchUserAndInExperiment && !profilerPluginInstallStatus,
isTorchUserAndInExperiment ? profilerPluginInstallStatus : ProductInstallStatus.Installed,
);
if (selection !== Common.bannerLabelYes() && !needsTensorBoardInstall) {
return true;
Expand Down Expand Up @@ -243,26 +252,39 @@ export class TensorBoardSession {
),
);
}
if (needsProfilerPluginInstall) {
installPromises.push(this.installer.install(Product.torchProfilerInstallName, interpreter, installerToken));
if (isTorchUserAndInExperiment && needsProfilerPluginInstall) {
installPromises.push(
this.installer.install(
Product.torchProfilerInstallName,
interpreter,
installerToken,
profilerPluginInstallStatus === ProductInstallStatus.NeedsUpgrade
? ModuleInstallFlags.upgrade
: undefined,
),
);
}
await Promise.race([...installPromises, cancellationPromise]);

// Check install status again after installing
[tensorboardInstallStatus, profilerPluginInstallStatus] = await Promise.all([
this.installer.isProductVersionCompatible(Product.tensorboard, '>= 2.4.1', interpreter),
this.installer.isInstalled(Product.torchProfilerImportName, interpreter),
this.installer.isProductVersionCompatible(Product.tensorboard, TensorBoardSemVerRequirement, interpreter),
this.installer.isProductVersionCompatible(
Product.torchProfilerImportName,
TorchProfilerSemVerRequirement,
interpreter,
),
]);
// Send telemetry regarding results of install
sendTelemetryEvent(EventName.TENSORBOARD_PACKAGE_INSTALL_RESULT, undefined, {
wasTensorBoardAttempted: needsTensorBoardInstall,
wasProfilerPluginAttempted: needsProfilerPluginInstall,
wasTensorBoardInstalled: tensorboardInstallStatus === ProductInstallStatus.Installed,
wasProfilerPluginInstalled: profilerPluginInstallStatus,
wasProfilerPluginInstalled: profilerPluginInstallStatus === ProductInstallStatus.Installed,
});
// Profiler plugin is not required to start TensorBoard. If it failed, note that it failed
// in the log, but report success only based on TensorBoard package install status.
if (isTorchUserAndInExperiment && profilerPluginInstallStatus === false) {
if (isTorchUserAndInExperiment && profilerPluginInstallStatus !== ProductInstallStatus.Installed) {
traceError(`Failed to install torch-tb-plugin. Profiler plugin will not appear in TensorBoard session.`);
}
return tensorboardInstallStatus === ProductInstallStatus.Installed;
Expand Down
69 changes: 46 additions & 23 deletions src/test/tensorBoard/tensorBoardSession.test.ts
Original file line number Diff line number Diff line change
Expand Up @@ -43,6 +43,12 @@ const info: PythonEnvironment = {
sysVersion: '',
};

const interpreter: PythonEnvironment = {
...info,
envType: EnvironmentType.Unknown,
path: PYTHON_PATH,
};

suite('TensorBoard session creation', async () => {
let serviceManager: IServiceManager;
let errorMessageStub: Sinon.SinonStub;
Expand All @@ -67,12 +73,6 @@ suite('TensorBoard session creation', async () => {
sandbox.stub(ExperimentHelpers, 'inDiscoveryExperiment').resolves(false);
experimentService = serviceManager.get<IExperimentService>(IExperimentService);

// Ensure we use CI Python
const interpreter: PythonEnvironment = {
...info,
envType: EnvironmentType.Unknown,
path: PYTHON_PATH,
};
const interpreterService = serviceManager.get<IInterpreterService>(IInterpreterService);
sandbox.stub(interpreterService, 'getActiveInterpreter').resolves(interpreter);

Expand All @@ -94,16 +94,21 @@ suite('TensorBoard session creation', async () => {
isInTorchProfilerExperiment: boolean,
hasTorchImports: boolean,
tensorBoardInstallStatus: ProductInstallStatus,
isTorchProfilerPackageInstalled: boolean,
torchProfilerPackageInstallStatus: ProductInstallStatus,
installPromptSelection: 'Yes' | 'No',
) {
sandbox
.stub(experimentService, 'inExperiment')
.withArgs(TorchProfiler.experiment)
.resolves(isInTorchProfilerExperiment);
sandbox.stub(ImportTracker, 'hasModuleImport').withArgs('torch').returns(hasTorchImports);
sandbox.stub(installer, 'isProductVersionCompatible').resolves(tensorBoardInstallStatus);
sandbox.stub(installer, 'isInstalled').resolves(isTorchProfilerPackageInstalled);
const isProductVersionCompatible = sandbox.stub(installer, 'isProductVersionCompatible');
isProductVersionCompatible
.withArgs(Product.tensorboard, '>= 2.4.1', interpreter)
.resolves(tensorBoardInstallStatus);
isProductVersionCompatible
.withArgs(Product.torchProfilerImportName, '>= 0.2.0', interpreter)
.resolves(torchProfilerPackageInstallStatus);
errorMessageStub = sandbox.stub(applicationShell, 'showErrorMessage');
errorMessageStub.resolves(installPromptSelection);
}
Expand Down Expand Up @@ -170,7 +175,7 @@ suite('TensorBoard session creation', async () => {
async function runTest(expectTensorBoardUpgrade: boolean) {
const installStub = sandbox.stub(installer, 'install').resolves(InstallerResponse.Installed);
await createSessionAndVerifyMessage(TensorBoard.installTensorBoardAndProfilerPluginPrompt());
assert.ok(installStub.calledTwice, 'Did not install anything');
assert.ok(installStub.calledTwice, `Expected 2 installs but got ${installStub.callCount} calls`);
assert.ok(installStub.calledWith(Product.torchProfilerInstallName));
assert.ok(
installStub.calledWith(
Expand All @@ -182,17 +187,17 @@ suite('TensorBoard session creation', async () => {
);
}
test('In experiment: true, has torch imports: true, is profiler package installed: false, TensorBoard needs upgrade', async () => {
configureStubs(true, true, ProductInstallStatus.NeedsUpgrade, false, 'Yes');
configureStubs(true, true, ProductInstallStatus.NeedsUpgrade, ProductInstallStatus.NotInstalled, 'Yes');
await runTest(true);
});
test('In experiment: true, has torch imports: true, is profiler package installed: false, TensorBoard not installed', async () => {
configureStubs(true, true, ProductInstallStatus.NotInstalled, false, 'Yes');
configureStubs(true, true, ProductInstallStatus.NotInstalled, ProductInstallStatus.NotInstalled, 'Yes');
await runTest(false);
});
});
suite('Install profiler only', async () => {
test('In experiment: true, has torch imports: true, is profiler package installed: false, TensorBoard installed', async () => {
configureStubs(true, true, ProductInstallStatus.Installed, false, 'Yes');
configureStubs(true, true, ProductInstallStatus.Installed, ProductInstallStatus.NotInstalled, 'Yes');
sandbox
.stub(applicationShell, 'showQuickPick')
.resolves({ label: TensorBoard.useCurrentWorkingDirectory() });
Expand Down Expand Up @@ -222,14 +227,20 @@ suite('TensorBoard session creation', async () => {
suite('Install tensorboard only', async () => {
[false, true].forEach(async (inExperiment) => {
[false, true].forEach(async (hasTorchImports) => {
[false, true].forEach(async (isTorchProfilerPackageInstalled) => {
[
ProductInstallStatus.Installed,
ProductInstallStatus.NotInstalled,
ProductInstallStatus.NeedsUpgrade,
].forEach(async (torchProfilerInstallStatus) => {
const isTorchProfilerPackageInstalled =
torchProfilerInstallStatus === ProductInstallStatus.Installed;
if (!(inExperiment && hasTorchImports && !isTorchProfilerPackageInstalled)) {
test(`In experiment: ${inExperiment}, has torch imports: ${hasTorchImports}, is profiler package installed: ${isTorchProfilerPackageInstalled}, TensorBoard not installed`, async () => {
configureStubs(
inExperiment,
hasTorchImports,
ProductInstallStatus.NotInstalled,
isTorchProfilerPackageInstalled,
torchProfilerInstallStatus,
'No',
);
await createSessionAndVerifyMessage(TensorBoard.installPrompt());
Expand All @@ -244,7 +255,7 @@ suite('TensorBoard session creation', async () => {
const installStub = sandbox.stub(installer, 'install').resolves(InstallerResponse.Installed);
await createSessionAndVerifyMessage(TensorBoard.upgradePrompt());

assert.ok(installStub.calledOnce, 'Did not install anything');
assert.ok(installStub.calledOnce, `Expected 1 install but got ${installStub.callCount} installs`);
assert.ok(installStub.args[0][0] === Product.tensorboard, 'Did not install tensorboard');
assert.ok(
installStub.args.filter((argsList) => argsList[0] === Product.torchProfilerInstallName).length ===
Expand All @@ -254,14 +265,20 @@ suite('TensorBoard session creation', async () => {
}
[false, true].forEach(async (inExperiment) => {
[false, true].forEach(async (hasTorchImports) => {
[false, true].forEach(async (isTorchProfilerPackageInstalled) => {
[
ProductInstallStatus.Installed,
ProductInstallStatus.NotInstalled,
ProductInstallStatus.NeedsUpgrade,
].forEach(async (torchProfilerInstallStatus) => {
const isTorchProfilerPackageInstalled =
torchProfilerInstallStatus === ProductInstallStatus.Installed;
if (!(inExperiment && hasTorchImports && !isTorchProfilerPackageInstalled)) {
test(`In experiment: ${inExperiment}, has torch imports: ${hasTorchImports}, is profiler package installed: ${isTorchProfilerPackageInstalled}, TensorBoard needs upgrade`, async () => {
configureStubs(
inExperiment,
hasTorchImports,
ProductInstallStatus.NeedsUpgrade,
isTorchProfilerPackageInstalled,
torchProfilerInstallStatus,
'Yes',
);
await runTest();
Expand All @@ -285,14 +302,20 @@ suite('TensorBoard session creation', async () => {
}
[false, true].forEach(async (inExperiment) => {
[false, true].forEach(async (hasTorchImports) => {
[false, true].forEach(async (isTorchProfilerPackageInstalled) => {
[
ProductInstallStatus.Installed,
ProductInstallStatus.NotInstalled,
ProductInstallStatus.NeedsUpgrade,
].forEach(async (torchProfilerInstallStatus) => {
const isTorchProfilerPackageInstalled =
torchProfilerInstallStatus === ProductInstallStatus.Installed;
if (!(inExperiment && hasTorchImports && !isTorchProfilerPackageInstalled)) {
test(`In experiment: ${inExperiment}, has torch imports: ${hasTorchImports}, is profiler package installed: ${isTorchProfilerPackageInstalled}, TensorBoard installed`, async () => {
configureStubs(
inExperiment,
hasTorchImports,
ProductInstallStatus.Installed,
isTorchProfilerPackageInstalled,
torchProfilerInstallStatus,
'Yes',
);
await runTest();
Expand Down Expand Up @@ -335,7 +358,7 @@ suite('TensorBoard session creation', async () => {
assert.ok(quickPickStub.notCalled, 'User opted not to upgrade and we proceeded to create session');
});
test('If TensorBoard is not installed and user chooses not to install, do not show error', async () => {
configureStubs(true, true, ProductInstallStatus.NotInstalled, false, 'Yes');
configureStubs(true, true, ProductInstallStatus.NotInstalled, ProductInstallStatus.NotInstalled, 'Yes');
sandbox.stub(installer, 'install').resolves(InstallerResponse.Ignore);

await commandManager.executeCommand(
Expand Down Expand Up @@ -385,7 +408,7 @@ suite('TensorBoard session creation', async () => {
assert.ok(errorMessageStub.called, 'TensorBoard timed out but no error was shown');
});
test('If installing the profiler package fails, do not show error, continue to create session', async () => {
configureStubs(true, true, ProductInstallStatus.Installed, false, 'Yes');
configureStubs(true, true, ProductInstallStatus.Installed, ProductInstallStatus.NotInstalled, 'Yes');
sandbox
.stub(applicationShell, 'showQuickPick')
.resolves({ label: TensorBoard.useCurrentWorkingDirectory() });
Expand All @@ -404,7 +427,7 @@ suite('TensorBoard session creation', async () => {
assert.ok(session.panel?.visible, 'Webview panel not shown, expected successful session creation');
});
test('If user opts not to install profiler package and tensorboard is already installed, continue to create session', async () => {
configureStubs(true, true, ProductInstallStatus.Installed, false, 'No');
configureStubs(true, true, ProductInstallStatus.Installed, ProductInstallStatus.NotInstalled, 'No');
sandbox
.stub(applicationShell, 'showQuickPick')
.resolves({ label: TensorBoard.useCurrentWorkingDirectory() });
Expand Down