Skip to content

Commit c1018da

Browse files
joyceerhlkimadeline
joyceerhl
and
kimadeline
authored
Recommend >= 0.2.0 of torch-tb-profiler (microsoft#16361)
* Recommend >= 0.2.0 of `torch-tb-profiler` (microsoft#16440) * Recommend >= 0.2.0 of torch-tb-profiler * Update CHANGELOG.md and delete news * Extract semver requirements into constants * Update CHANGELOG.md Co-authored-by: Kim-Adeline Miguel <[email protected]> Co-authored-by: Kim-Adeline Miguel <[email protected]>
1 parent 23de127 commit c1018da

File tree

4 files changed

+86
-40
lines changed

4 files changed

+86
-40
lines changed

‎CHANGELOG.md‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -10,9 +10,10 @@
1010
([#16102](https://github.com/Microsoft/vscode-python/issues/16102))
1111
1. Add an `enumDescriptions` key under the `python.languageServer` setting to describe all language server options.
1212
([#16141](https://github.com/Microsoft/vscode-python/issues/16141))
13+
1. Ensure users upgrade to v0.2.0 of the torch-tb-profiler TensorBoard plugin to access jump-to-source functionality.
14+
([#16330](https://github.com/Microsoft/vscode-python/issues/16330))
1315

1416
### Fixes
15-
1617
1. Fixes a bug in the bandit linter where messages weren't being propagated to the editor.
1718
(thanks [Anthony Shaw](https://github.com/tonybaloney))
1819
([#15561](https://github.com/Microsoft/vscode-python/issues/15561))

‎src/client/common/utils/localize.ts‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -177,11 +177,11 @@ export namespace TensorBoard {
177177
);
178178
export const installTensorBoardAndProfilerPluginPrompt = localize(
179179
'TensorBoard.installTensorBoardAndProfilerPluginPrompt',
180-
'TensorBoard >= 2.4.1 and the PyTorch Profiler TensorBoard plugin are required. Would you like to install these packages?',
180+
'TensorBoard >= 2.4.1 and the PyTorch Profiler TensorBoard plugin >= 0.2.0 are required. Would you like to install these packages?',
181181
);
182182
export const installProfilerPluginPrompt = localize(
183183
'TensorBoard.installProfilerPluginPrompt',
184-
'We recommend installing the PyTorch Profiler TensorBoard plugin. Would you like to install the package?',
184+
'We recommend installing version >= 0.2.0 of the PyTorch Profiler TensorBoard plugin. Would you like to install the package?',
185185
);
186186
export const upgradePrompt = localize(
187187
'TensorBoard.upgradePrompt',

‎src/client/tensorBoard/tensorBoardSession.ts‎

Lines changed: 36 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -50,6 +50,8 @@ import { ModuleInstallFlags } from '../common/installer/types';
5050
enum Messages {
5151
JumpToSource = 'jump_to_source',
5252
}
53+
const TensorBoardSemVerRequirement = '>= 2.4.1';
54+
const TorchProfilerSemVerRequirement = '>= 0.2.0';
5355

5456
/**
5557
* Manages the lifecycle of a TensorBoard session.
@@ -143,18 +145,21 @@ export class TensorBoardSession {
143145

144146
private async promptToInstall(
145147
tensorBoardInstallStatus: ProductInstallStatus,
146-
shouldInstallProfilerPlugin: boolean,
148+
profilerPluginInstallStatus: ProductInstallStatus,
147149
) {
148150
sendTelemetryEvent(EventName.TENSORBOARD_INSTALL_PROMPT_SHOWN);
149151
const yes = Common.bannerLabelYes();
150152
const no = Common.bannerLabelNo();
151153
const isUpgrade = tensorBoardInstallStatus === ProductInstallStatus.NeedsUpgrade;
152154
let message;
153155

154-
if (tensorBoardInstallStatus === ProductInstallStatus.Installed && shouldInstallProfilerPlugin) {
156+
if (
157+
tensorBoardInstallStatus === ProductInstallStatus.Installed &&
158+
profilerPluginInstallStatus !== ProductInstallStatus.Installed
159+
) {
155160
// PyTorch user already has TensorBoard, just ask if they want the profiler plugin
156161
message = TensorBoard.installProfilerPluginPrompt();
157-
} else if (shouldInstallProfilerPlugin) {
162+
} else if (profilerPluginInstallStatus !== ProductInstallStatus.Installed) {
158163
// PyTorch user doesn't have compatible TensorBoard or the profiler plugin
159164
message = TensorBoard.installTensorBoardAndProfilerPluginPrompt();
160165
} else if (isUpgrade) {
@@ -193,15 +198,19 @@ export class TensorBoardSession {
193198

194199
// First see what dependencies we're missing
195200
let [tensorboardInstallStatus, profilerPluginInstallStatus] = await Promise.all([
196-
this.installer.isProductVersionCompatible(Product.tensorboard, '>= 2.4.1', interpreter),
197-
this.installer.isInstalled(Product.torchProfilerImportName, interpreter),
201+
this.installer.isProductVersionCompatible(Product.tensorboard, TensorBoardSemVerRequirement, interpreter),
202+
this.installer.isProductVersionCompatible(
203+
Product.torchProfilerImportName,
204+
TorchProfilerSemVerRequirement,
205+
interpreter,
206+
),
198207
]);
199208
const isTorchUserAndInExperiment = ImportTracker.hasModuleImport('torch') && this.isInTorchProfilerExperiment;
200209
const needsTensorBoardInstall = tensorboardInstallStatus !== ProductInstallStatus.Installed;
201-
const needsProfilerPluginInstall = isTorchUserAndInExperiment && profilerPluginInstallStatus !== true;
210+
const needsProfilerPluginInstall = profilerPluginInstallStatus !== ProductInstallStatus.Installed;
202211
if (
203212
// PyTorch user, in profiler install experiment, TensorBoard and profiler plugin already installed
204-
(isTorchUserAndInExperiment && !needsTensorBoardInstall && profilerPluginInstallStatus === true) ||
213+
(isTorchUserAndInExperiment && !needsTensorBoardInstall && !needsProfilerPluginInstall) ||
205214
// Not PyTorch user or not in profiler install experiment, so no need for profiler plugin,
206215
// and TensorBoard is already installed
207216
(!isTorchUserAndInExperiment && tensorboardInstallStatus === ProductInstallStatus.Installed)
@@ -212,7 +221,7 @@ export class TensorBoardSession {
212221
// Ask the user if they want to install packages to start a TensorBoard session
213222
const selection = await this.promptToInstall(
214223
tensorboardInstallStatus,
215-
isTorchUserAndInExperiment && !profilerPluginInstallStatus,
224+
isTorchUserAndInExperiment ? profilerPluginInstallStatus : ProductInstallStatus.Installed,
216225
);
217226
if (selection !== Common.bannerLabelYes() && !needsTensorBoardInstall) {
218227
return true;
@@ -243,26 +252,39 @@ export class TensorBoardSession {
243252
),
244253
);
245254
}
246-
if (needsProfilerPluginInstall) {
247-
installPromises.push(this.installer.install(Product.torchProfilerInstallName, interpreter, installerToken));
255+
if (isTorchUserAndInExperiment && needsProfilerPluginInstall) {
256+
installPromises.push(
257+
this.installer.install(
258+
Product.torchProfilerInstallName,
259+
interpreter,
260+
installerToken,
261+
profilerPluginInstallStatus === ProductInstallStatus.NeedsUpgrade
262+
? ModuleInstallFlags.upgrade
263+
: undefined,
264+
),
265+
);
248266
}
249267
await Promise.race([...installPromises, cancellationPromise]);
250268

251269
// Check install status again after installing
252270
[tensorboardInstallStatus, profilerPluginInstallStatus] = await Promise.all([
253-
this.installer.isProductVersionCompatible(Product.tensorboard, '>= 2.4.1', interpreter),
254-
this.installer.isInstalled(Product.torchProfilerImportName, interpreter),
271+
this.installer.isProductVersionCompatible(Product.tensorboard, TensorBoardSemVerRequirement, interpreter),
272+
this.installer.isProductVersionCompatible(
273+
Product.torchProfilerImportName,
274+
TorchProfilerSemVerRequirement,
275+
interpreter,
276+
),
255277
]);
256278
// Send telemetry regarding results of install
257279
sendTelemetryEvent(EventName.TENSORBOARD_PACKAGE_INSTALL_RESULT, undefined, {
258280
wasTensorBoardAttempted: needsTensorBoardInstall,
259281
wasProfilerPluginAttempted: needsProfilerPluginInstall,
260282
wasTensorBoardInstalled: tensorboardInstallStatus === ProductInstallStatus.Installed,
261-
wasProfilerPluginInstalled: profilerPluginInstallStatus,
283+
wasProfilerPluginInstalled: profilerPluginInstallStatus === ProductInstallStatus.Installed,
262284
});
263285
// Profiler plugin is not required to start TensorBoard. If it failed, note that it failed
264286
// in the log, but report success only based on TensorBoard package install status.
265-
if (isTorchUserAndInExperiment && profilerPluginInstallStatus === false) {
287+
if (isTorchUserAndInExperiment && profilerPluginInstallStatus !== ProductInstallStatus.Installed) {
266288
traceError(`Failed to install torch-tb-plugin. Profiler plugin will not appear in TensorBoard session.`);
267289
}
268290
return tensorboardInstallStatus === ProductInstallStatus.Installed;

‎src/test/tensorBoard/tensorBoardSession.test.ts‎

Lines changed: 46 additions & 23 deletions
Original file line numberDiff line numberDiff line change
@@ -43,6 +43,12 @@ const info: PythonEnvironment = {
4343
sysVersion: '',
4444
};
4545

46+
const interpreter: PythonEnvironment = {
47+
...info,
48+
envType: EnvironmentType.Unknown,
49+
path: PYTHON_PATH,
50+
};
51+
4652
suite('TensorBoard session creation', async () => {
4753
let serviceManager: IServiceManager;
4854
let errorMessageStub: Sinon.SinonStub;
@@ -67,12 +73,6 @@ suite('TensorBoard session creation', async () => {
6773
sandbox.stub(ExperimentHelpers, 'inDiscoveryExperiment').resolves(false);
6874
experimentService = serviceManager.get<IExperimentService>(IExperimentService);
6975

70-
// Ensure we use CI Python
71-
const interpreter: PythonEnvironment = {
72-
...info,
73-
envType: EnvironmentType.Unknown,
74-
path: PYTHON_PATH,
75-
};
7676
const interpreterService = serviceManager.get<IInterpreterService>(IInterpreterService);
7777
sandbox.stub(interpreterService, 'getActiveInterpreter').resolves(interpreter);
7878

@@ -94,16 +94,21 @@ suite('TensorBoard session creation', async () => {
9494
isInTorchProfilerExperiment: boolean,
9595
hasTorchImports: boolean,
9696
tensorBoardInstallStatus: ProductInstallStatus,
97-
isTorchProfilerPackageInstalled: boolean,
97+
torchProfilerPackageInstallStatus: ProductInstallStatus,
9898
installPromptSelection: 'Yes' | 'No',
9999
) {
100100
sandbox
101101
.stub(experimentService, 'inExperiment')
102102
.withArgs(TorchProfiler.experiment)
103103
.resolves(isInTorchProfilerExperiment);
104104
sandbox.stub(ImportTracker, 'hasModuleImport').withArgs('torch').returns(hasTorchImports);
105-
sandbox.stub(installer, 'isProductVersionCompatible').resolves(tensorBoardInstallStatus);
106-
sandbox.stub(installer, 'isInstalled').resolves(isTorchProfilerPackageInstalled);
105+
const isProductVersionCompatible = sandbox.stub(installer, 'isProductVersionCompatible');
106+
isProductVersionCompatible
107+
.withArgs(Product.tensorboard, '>= 2.4.1', interpreter)
108+
.resolves(tensorBoardInstallStatus);
109+
isProductVersionCompatible
110+
.withArgs(Product.torchProfilerImportName, '>= 0.2.0', interpreter)
111+
.resolves(torchProfilerPackageInstallStatus);
107112
errorMessageStub = sandbox.stub(applicationShell, 'showErrorMessage');
108113
errorMessageStub.resolves(installPromptSelection);
109114
}
@@ -170,7 +175,7 @@ suite('TensorBoard session creation', async () => {
170175
async function runTest(expectTensorBoardUpgrade: boolean) {
171176
const installStub = sandbox.stub(installer, 'install').resolves(InstallerResponse.Installed);
172177
await createSessionAndVerifyMessage(TensorBoard.installTensorBoardAndProfilerPluginPrompt());
173-
assert.ok(installStub.calledTwice, 'Did not install anything');
178+
assert.ok(installStub.calledTwice, `Expected 2 installs but got ${installStub.callCount} calls`);
174179
assert.ok(installStub.calledWith(Product.torchProfilerInstallName));
175180
assert.ok(
176181
installStub.calledWith(
@@ -182,17 +187,17 @@ suite('TensorBoard session creation', async () => {
182187
);
183188
}
184189
test('In experiment: true, has torch imports: true, is profiler package installed: false, TensorBoard needs upgrade', async () => {
185-
configureStubs(true, true, ProductInstallStatus.NeedsUpgrade, false, 'Yes');
190+
configureStubs(true, true, ProductInstallStatus.NeedsUpgrade, ProductInstallStatus.NotInstalled, 'Yes');
186191
await runTest(true);
187192
});
188193
test('In experiment: true, has torch imports: true, is profiler package installed: false, TensorBoard not installed', async () => {
189-
configureStubs(true, true, ProductInstallStatus.NotInstalled, false, 'Yes');
194+
configureStubs(true, true, ProductInstallStatus.NotInstalled, ProductInstallStatus.NotInstalled, 'Yes');
190195
await runTest(false);
191196
});
192197
});
193198
suite('Install profiler only', async () => {
194199
test('In experiment: true, has torch imports: true, is profiler package installed: false, TensorBoard installed', async () => {
195-
configureStubs(true, true, ProductInstallStatus.Installed, false, 'Yes');
200+
configureStubs(true, true, ProductInstallStatus.Installed, ProductInstallStatus.NotInstalled, 'Yes');
196201
sandbox
197202
.stub(applicationShell, 'showQuickPick')
198203
.resolves({ label: TensorBoard.useCurrentWorkingDirectory() });
@@ -222,14 +227,20 @@ suite('TensorBoard session creation', async () => {
222227
suite('Install tensorboard only', async () => {
223228
[false, true].forEach(async (inExperiment) => {
224229
[false, true].forEach(async (hasTorchImports) => {
225-
[false, true].forEach(async (isTorchProfilerPackageInstalled) => {
230+
[
231+
ProductInstallStatus.Installed,
232+
ProductInstallStatus.NotInstalled,
233+
ProductInstallStatus.NeedsUpgrade,
234+
].forEach(async (torchProfilerInstallStatus) => {
235+
const isTorchProfilerPackageInstalled =
236+
torchProfilerInstallStatus === ProductInstallStatus.Installed;
226237
if (!(inExperiment && hasTorchImports && !isTorchProfilerPackageInstalled)) {
227238
test(`In experiment: ${inExperiment}, has torch imports: ${hasTorchImports}, is profiler package installed: ${isTorchProfilerPackageInstalled}, TensorBoard not installed`, async () => {
228239
configureStubs(
229240
inExperiment,
230241
hasTorchImports,
231242
ProductInstallStatus.NotInstalled,
232-
isTorchProfilerPackageInstalled,
243+
torchProfilerInstallStatus,
233244
'No',
234245
);
235246
await createSessionAndVerifyMessage(TensorBoard.installPrompt());
@@ -244,7 +255,7 @@ suite('TensorBoard session creation', async () => {
244255
const installStub = sandbox.stub(installer, 'install').resolves(InstallerResponse.Installed);
245256
await createSessionAndVerifyMessage(TensorBoard.upgradePrompt());
246257

247-
assert.ok(installStub.calledOnce, 'Did not install anything');
258+
assert.ok(installStub.calledOnce, `Expected 1 install but got ${installStub.callCount} installs`);
248259
assert.ok(installStub.args[0][0] === Product.tensorboard, 'Did not install tensorboard');
249260
assert.ok(
250261
installStub.args.filter((argsList) => argsList[0] === Product.torchProfilerInstallName).length ===
@@ -254,14 +265,20 @@ suite('TensorBoard session creation', async () => {
254265
}
255266
[false, true].forEach(async (inExperiment) => {
256267
[false, true].forEach(async (hasTorchImports) => {
257-
[false, true].forEach(async (isTorchProfilerPackageInstalled) => {
268+
[
269+
ProductInstallStatus.Installed,
270+
ProductInstallStatus.NotInstalled,
271+
ProductInstallStatus.NeedsUpgrade,
272+
].forEach(async (torchProfilerInstallStatus) => {
273+
const isTorchProfilerPackageInstalled =
274+
torchProfilerInstallStatus === ProductInstallStatus.Installed;
258275
if (!(inExperiment && hasTorchImports && !isTorchProfilerPackageInstalled)) {
259276
test(`In experiment: ${inExperiment}, has torch imports: ${hasTorchImports}, is profiler package installed: ${isTorchProfilerPackageInstalled}, TensorBoard needs upgrade`, async () => {
260277
configureStubs(
261278
inExperiment,
262279
hasTorchImports,
263280
ProductInstallStatus.NeedsUpgrade,
264-
isTorchProfilerPackageInstalled,
281+
torchProfilerInstallStatus,
265282
'Yes',
266283
);
267284
await runTest();
@@ -285,14 +302,20 @@ suite('TensorBoard session creation', async () => {
285302
}
286303
[false, true].forEach(async (inExperiment) => {
287304
[false, true].forEach(async (hasTorchImports) => {
288-
[false, true].forEach(async (isTorchProfilerPackageInstalled) => {
305+
[
306+
ProductInstallStatus.Installed,
307+
ProductInstallStatus.NotInstalled,
308+
ProductInstallStatus.NeedsUpgrade,
309+
].forEach(async (torchProfilerInstallStatus) => {
310+
const isTorchProfilerPackageInstalled =
311+
torchProfilerInstallStatus === ProductInstallStatus.Installed;
289312
if (!(inExperiment && hasTorchImports && !isTorchProfilerPackageInstalled)) {
290313
test(`In experiment: ${inExperiment}, has torch imports: ${hasTorchImports}, is profiler package installed: ${isTorchProfilerPackageInstalled}, TensorBoard installed`, async () => {
291314
configureStubs(
292315
inExperiment,
293316
hasTorchImports,
294317
ProductInstallStatus.Installed,
295-
isTorchProfilerPackageInstalled,
318+
torchProfilerInstallStatus,
296319
'Yes',
297320
);
298321
await runTest();
@@ -335,7 +358,7 @@ suite('TensorBoard session creation', async () => {
335358
assert.ok(quickPickStub.notCalled, 'User opted not to upgrade and we proceeded to create session');
336359
});
337360
test('If TensorBoard is not installed and user chooses not to install, do not show error', async () => {
338-
configureStubs(true, true, ProductInstallStatus.NotInstalled, false, 'Yes');
361+
configureStubs(true, true, ProductInstallStatus.NotInstalled, ProductInstallStatus.NotInstalled, 'Yes');
339362
sandbox.stub(installer, 'install').resolves(InstallerResponse.Ignore);
340363

341364
await commandManager.executeCommand(
@@ -385,7 +408,7 @@ suite('TensorBoard session creation', async () => {
385408
assert.ok(errorMessageStub.called, 'TensorBoard timed out but no error was shown');
386409
});
387410
test('If installing the profiler package fails, do not show error, continue to create session', async () => {
388-
configureStubs(true, true, ProductInstallStatus.Installed, false, 'Yes');
411+
configureStubs(true, true, ProductInstallStatus.Installed, ProductInstallStatus.NotInstalled, 'Yes');
389412
sandbox
390413
.stub(applicationShell, 'showQuickPick')
391414
.resolves({ label: TensorBoard.useCurrentWorkingDirectory() });
@@ -404,7 +427,7 @@ suite('TensorBoard session creation', async () => {
404427
assert.ok(session.panel?.visible, 'Webview panel not shown, expected successful session creation');
405428
});
406429
test('If user opts not to install profiler package and tensorboard is already installed, continue to create session', async () => {
407-
configureStubs(true, true, ProductInstallStatus.Installed, false, 'No');
430+
configureStubs(true, true, ProductInstallStatus.Installed, ProductInstallStatus.NotInstalled, 'No');
408431
sandbox
409432
.stub(applicationShell, 'showQuickPick')
410433
.resolves({ label: TensorBoard.useCurrentWorkingDirectory() });

0 commit comments

Comments
 (0)