Skip to content

Commit a4133e4

Browse files
joyceerhljoyceerhl
authored
Recommend >= 0.2.0 of torch-tb-profiler (#16440)
* Recommend >= 0.2.0 of torch-tb-profiler * Update CHANGELOG.md and delete news
1 parent c383e81 commit a4133e4

File tree

4 files changed

+74
-37
lines changed

4 files changed

+74
-37
lines changed

‎CHANGELOG.md‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -10,6 +10,8 @@
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/16339))
1315

1416
### Fixes
1517

‎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: 24 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -143,18 +143,21 @@ export class TensorBoardSession {
143143

144144
private async promptToInstall(
145145
tensorBoardInstallStatus: ProductInstallStatus,
146-
shouldInstallProfilerPlugin: boolean,
146+
profilerPluginInstallStatus: ProductInstallStatus,
147147
) {
148148
sendTelemetryEvent(EventName.TENSORBOARD_INSTALL_PROMPT_SHOWN);
149149
const yes = Common.bannerLabelYes();
150150
const no = Common.bannerLabelNo();
151151
const isUpgrade = tensorBoardInstallStatus === ProductInstallStatus.NeedsUpgrade;
152152
let message;
153153

154-
if (tensorBoardInstallStatus === ProductInstallStatus.Installed && shouldInstallProfilerPlugin) {
154+
if (
155+
tensorBoardInstallStatus === ProductInstallStatus.Installed &&
156+
profilerPluginInstallStatus !== ProductInstallStatus.Installed
157+
) {
155158
// PyTorch user already has TensorBoard, just ask if they want the profiler plugin
156159
message = TensorBoard.installProfilerPluginPrompt();
157-
} else if (shouldInstallProfilerPlugin) {
160+
} else if (profilerPluginInstallStatus !== ProductInstallStatus.Installed) {
158161
// PyTorch user doesn't have compatible TensorBoard or the profiler plugin
159162
message = TensorBoard.installTensorBoardAndProfilerPluginPrompt();
160163
} else if (isUpgrade) {
@@ -194,14 +197,14 @@ export class TensorBoardSession {
194197
// First see what dependencies we're missing
195198
let [tensorboardInstallStatus, profilerPluginInstallStatus] = await Promise.all([
196199
this.installer.isProductVersionCompatible(Product.tensorboard, '>= 2.4.1', interpreter),
197-
this.installer.isInstalled(Product.torchProfilerImportName, interpreter),
200+
this.installer.isProductVersionCompatible(Product.torchProfilerImportName, '>= 0.2.0', interpreter),
198201
]);
199202
const isTorchUserAndInExperiment = ImportTracker.hasModuleImport('torch') && this.isInTorchProfilerExperiment;
200203
const needsTensorBoardInstall = tensorboardInstallStatus !== ProductInstallStatus.Installed;
201-
const needsProfilerPluginInstall = isTorchUserAndInExperiment && profilerPluginInstallStatus !== true;
204+
const needsProfilerPluginInstall = profilerPluginInstallStatus !== ProductInstallStatus.Installed;
202205
if (
203206
// PyTorch user, in profiler install experiment, TensorBoard and profiler plugin already installed
204-
(isTorchUserAndInExperiment && !needsTensorBoardInstall && profilerPluginInstallStatus === true) ||
207+
(isTorchUserAndInExperiment && !needsTensorBoardInstall && !needsProfilerPluginInstall) ||
205208
// Not PyTorch user or not in profiler install experiment, so no need for profiler plugin,
206209
// and TensorBoard is already installed
207210
(!isTorchUserAndInExperiment && tensorboardInstallStatus === ProductInstallStatus.Installed)
@@ -212,7 +215,7 @@ export class TensorBoardSession {
212215
// Ask the user if they want to install packages to start a TensorBoard session
213216
const selection = await this.promptToInstall(
214217
tensorboardInstallStatus,
215-
isTorchUserAndInExperiment && !profilerPluginInstallStatus,
218+
isTorchUserAndInExperiment ? profilerPluginInstallStatus : ProductInstallStatus.Installed,
216219
);
217220
if (selection !== Common.bannerLabelYes() && !needsTensorBoardInstall) {
218221
return true;
@@ -243,26 +246,35 @@ export class TensorBoardSession {
243246
),
244247
);
245248
}
246-
if (needsProfilerPluginInstall) {
247-
installPromises.push(this.installer.install(Product.torchProfilerInstallName, interpreter, installerToken));
249+
if (isTorchUserAndInExperiment && needsProfilerPluginInstall) {
250+
installPromises.push(
251+
this.installer.install(
252+
Product.torchProfilerInstallName,
253+
interpreter,
254+
installerToken,
255+
profilerPluginInstallStatus === ProductInstallStatus.NeedsUpgrade
256+
? ModuleInstallFlags.upgrade
257+
: undefined,
258+
),
259+
);
248260
}
249261
await Promise.race([...installPromises, cancellationPromise]);
250262

251263
// Check install status again after installing
252264
[tensorboardInstallStatus, profilerPluginInstallStatus] = await Promise.all([
253265
this.installer.isProductVersionCompatible(Product.tensorboard, '>= 2.4.1', interpreter),
254-
this.installer.isInstalled(Product.torchProfilerImportName, interpreter),
266+
this.installer.isProductVersionCompatible(Product.torchProfilerImportName, '>= 0.2.0', interpreter),
255267
]);
256268
// Send telemetry regarding results of install
257269
sendTelemetryEvent(EventName.TENSORBOARD_PACKAGE_INSTALL_RESULT, undefined, {
258270
wasTensorBoardAttempted: needsTensorBoardInstall,
259271
wasProfilerPluginAttempted: needsProfilerPluginInstall,
260272
wasTensorBoardInstalled: tensorboardInstallStatus === ProductInstallStatus.Installed,
261-
wasProfilerPluginInstalled: profilerPluginInstallStatus,
273+
wasProfilerPluginInstalled: profilerPluginInstallStatus === ProductInstallStatus.Installed,
262274
});
263275
// Profiler plugin is not required to start TensorBoard. If it failed, note that it failed
264276
// in the log, but report success only based on TensorBoard package install status.
265-
if (isTorchUserAndInExperiment && profilerPluginInstallStatus === false) {
277+
if (isTorchUserAndInExperiment && profilerPluginInstallStatus !== ProductInstallStatus.Installed) {
266278
traceError(`Failed to install torch-tb-plugin. Profiler plugin will not appear in TensorBoard session.`);
267279
}
268280
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)