Skip to content

Commit 359bc18

Browse files
joyceerhljoyceerhl
authored
Allow user to select TensorBoard plugin missing profiler source files using file picker (microsoft#15993)
1 parent 17fb197 commit 359bc18

7 files changed

Lines changed: 153 additions & 51 deletions

File tree

‎news/1 Enhancements/15695.md‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1 @@
1+
In an integrated TensorBoard session, if the jump to source request is for a file that does not exist on disk, allow the user to manually specify the file using the system file picker.

‎package.nls.json‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -231,6 +231,9 @@
231231
"StartPage.folderDesc": "- Open a <div class=\"link\" role=\"button\" onclick={0}>Folder</div><br /> - Open a <div class=\"link\" role=\"button\" onclick={1}>Workspace</div>",
232232
"StartPage.badWebPanelFormatString": "<html><body><h1>{0} is not a valid file name</h1></body></html>",
233233
"Jupyter.extensionRequired": "The Jupyter extension is required to perform that task. Click Yes to open the Jupyter extension installation page.",
234+
"TensorBoard.missingSourceFile": "We could not locate the requested source file on disk. Please manually specify the file.",
235+
"TensorBoard.selectMissingSourceFile": "Choose File",
236+
"TensorBoard.selectMissingSourceFileDescription": "The source file's contents may not match the original contents in the trace.",
234237
"TensorBoard.useCurrentWorkingDirectory": "Use current working directory",
235238
"TensorBoard.currentDirectory": "Current: {0}",
236239
"TensorBoard.logDirectoryPrompt": "Select a log directory to start TensorBoard with",

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

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -196,6 +196,15 @@ export namespace TensorBoard {
196196
'TensorBoard.launchNativeTensorBoardSessionCodeAction',
197197
'Launch TensorBoard session',
198198
);
199+
export const missingSourceFile = localize(
200+
'TensorBoard.missingSourceFile',
201+
'We could not locate the requested source file on disk. Please manually specify the file.',
202+
);
203+
export const selectMissingSourceFile = localize('TensorBoard.selectMissingSourceFile', 'Choose File');
204+
export const selectMissingSourceFileDescription = localize(
205+
'TensorBoard.selectMissingSourceFileDescription',
206+
"The source file's contents may not match the original contents in the trace.",
207+
);
199208
}
200209

201210
export namespace LanguageService {

‎src/client/tensorBoard/tensorBoardSession.ts‎

Lines changed: 54 additions & 19 deletions
Original file line numberDiff line numberDiff line change
@@ -42,6 +42,7 @@ import { sendTelemetryEvent } from '../telemetry';
4242
import { EventName } from '../telemetry/constants';
4343
import { ImportTracker } from '../telemetry/importTracker';
4444
import { TensorBoardPromptSelection, TensorBoardSessionStartResult } from './constants';
45+
import { IMultiStepInputFactory } from '../common/utils/multiStepInput';
4546

4647
enum Messages {
4748
JumpToSource = 'jump_to_source',
@@ -87,6 +88,7 @@ export class TensorBoardSession {
8788
private readonly applicationShell: IApplicationShell,
8889
private readonly isInTorchProfilerExperiment: boolean,
8990
private readonly globalMemento: IPersistentState<ViewColumn>,
91+
private readonly multiStepFactory: IMultiStepInputFactory,
9092
) {}
9193

9294
public async initialize(): Promise<void> {
@@ -419,7 +421,6 @@ export class TensorBoardSession {
419421
private async createPanel() {
420422
const webviewPanel = window.createWebviewPanel('tensorBoardSession', 'TensorBoard', this.globalMemento.value, {
421423
enableScripts: true,
422-
retainContextWhenHidden: true,
423424
});
424425
const fullWebServerUri = await env.asExternalUri(Uri.parse(this.url!));
425426
webviewPanel.webview.html = `<!DOCTYPE html>
@@ -499,7 +500,7 @@ export class TensorBoardSession {
499500
// Handle messages posted from the webview
500501
switch (message.command) {
501502
case Messages.JumpToSource:
502-
jumpToSource(message.args.filename, message.args.line);
503+
void this.jumpToSource(message.args.filename, message.args.line);
503504
break;
504505
default:
505506
break;
@@ -519,24 +520,58 @@ export class TensorBoardSession {
519520
}
520521
return undefined;
521522
}
522-
}
523523

524-
function jumpToSource(fsPath: string, line: number) {
525-
if (fs.existsSync(fsPath)) {
526-
const uri = Uri.file(fsPath);
527-
workspace
528-
.openTextDocument(uri)
529-
.then((doc) => window.showTextDocument(doc, ViewColumn.Beside))
530-
.then((editor) => {
531-
// Select the line if it exists in the document
532-
if (line < editor.document.lineCount) {
533-
const position = new Position(line, 0);
534-
const selection = new Selection(position, editor.document.lineAt(line).range.end);
535-
editor.selection = selection;
536-
editor.revealRange(selection, TextEditorRevealType.InCenterIfOutsideViewport);
524+
private async jumpToSource(fsPath: string, line: number) {
525+
let uri: Uri | undefined;
526+
if (fs.existsSync(fsPath)) {
527+
uri = Uri.file(fsPath);
528+
} else {
529+
traceError(
530+
`Requested jump to source filepath ${fsPath} does not exist. Prompting user to select source file...`,
531+
);
532+
// Prompt the user to pick the file on disk
533+
const items: QuickPickItem[] = [
534+
{
535+
label: TensorBoard.selectMissingSourceFile(),
536+
description: TensorBoard.selectMissingSourceFileDescription(),
537+
},
538+
];
539+
// Using a multistep so that we can add a title to the quickpick
540+
const multiStep = this.multiStepFactory.create<unknown>();
541+
await multiStep.run(async (input) => {
542+
const selection = await input.showQuickPick({
543+
items,
544+
title: TensorBoard.missingSourceFile(),
545+
placeholder: fsPath,
546+
});
547+
switch (selection?.label) {
548+
case TensorBoard.selectMissingSourceFile(): {
549+
const filePickerSelection = await this.applicationShell.showOpenDialog({
550+
canSelectFiles: true,
551+
canSelectFolders: false,
552+
canSelectMany: false,
553+
});
554+
if (filePickerSelection !== undefined) {
555+
[uri] = filePickerSelection;
556+
}
557+
break;
558+
}
559+
default:
560+
break;
537561
}
538-
});
539-
} else {
540-
traceError(`Requested jump to source filepath ${fsPath} does not exist`);
562+
}, {});
563+
}
564+
if (uri === undefined) {
565+
return;
566+
}
567+
const document = await workspace.openTextDocument(uri);
568+
const editor = await window.showTextDocument(document, ViewColumn.Beside);
569+
// Select the line if it exists in the document
570+
if (line < editor.document.lineCount) {
571+
const position = new Position(line, 0);
572+
const selection = new Selection(position, editor.document.lineAt(line).range.end);
573+
editor.selection = selection;
574+
editor.revealRange(selection, TextEditorRevealType.InCenterIfOutsideViewport);
575+
}
541576
}
542577
}

‎src/client/tensorBoard/tensorBoardSessionProvider.ts‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -11,6 +11,7 @@ import { traceError, traceInfo } from '../common/logger';
1111
import { IProcessServiceFactory } from '../common/process/types';
1212
import { IDisposableRegistry, IExperimentService, IInstaller, IPersistentStateFactory } from '../common/types';
1313
import { TensorBoard } from '../common/utils/localize';
14+
import { IMultiStepInputFactory } from '../common/utils/multiStepInput';
1415
import { IInterpreterService } from '../interpreter/contracts';
1516
import { sendTelemetryEvent } from '../telemetry';
1617
import { EventName } from '../telemetry/constants';
@@ -31,6 +32,7 @@ export class TensorBoardSessionProvider implements IExtensionSingleActivationSer
3132
@inject(IProcessServiceFactory) private readonly processServiceFactory: IProcessServiceFactory,
3233
@inject(IExperimentService) private readonly experimentService: IExperimentService,
3334
@inject(IPersistentStateFactory) private stateFactory: IPersistentStateFactory,
35+
@inject(IMultiStepInputFactory) private readonly multiStepFactory: IMultiStepInputFactory,
3436
) {}
3537

3638
public async activate(): Promise<void> {
@@ -68,6 +70,7 @@ export class TensorBoardSessionProvider implements IExtensionSingleActivationSer
6870
this.applicationShell,
6971
await this.experimentService.inExperiment(TorchProfiler.experiment),
7072
memento,
73+
this.multiStepFactory,
7174
);
7275
await newSession.initialize();
7376
return newSession;

‎src/test/pythonFiles/tensorBoard/sourcefile.py‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1 @@
1+
from torch.utils.tensorboard import SummaryWriter

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

Lines changed: 82 additions & 32 deletions
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,8 @@
1+
import * as path from 'path';
12
import { assert } from 'chai';
23
import Sinon, * as sinon from 'sinon';
34
import { SemVer } from 'semver';
4-
import { ViewColumn, workspace, WorkspaceConfiguration } from 'vscode';
5+
import { Uri, ViewColumn, window, workspace, WorkspaceConfiguration } from 'vscode';
56
import {
67
IExperimentService,
78
IInstaller,
@@ -14,14 +15,20 @@ import { IApplicationShell, ICommandManager } from '../../client/common/applicat
1415
import { IServiceManager } from '../../client/ioc/types';
1516
import { TensorBoardEntrypoint, TensorBoardEntrypointTrigger } from '../../client/tensorBoard/constants';
1617
import { TensorBoardSession } from '../../client/tensorBoard/tensorBoardSession';
17-
import { closeActiveWindows, initialize } from '../initialize';
18+
import { closeActiveWindows, EXTENSION_ROOT_DIR_FOR_TESTS, initialize } from '../initialize';
1819
import * as ExperimentHelpers from '../../client/common/experiments/helpers';
1920
import { IInterpreterService } from '../../client/interpreter/contracts';
2021
import { Architecture } from '../../client/common/utils/platform';
2122
import { PythonEnvironment, EnvironmentType } from '../../client/pythonEnvironments/info';
2223
import { PYTHON_PATH } from '../common';
2324
import { TorchProfiler } from '../../client/common/experiments/groups';
2425
import { ImportTracker } from '../../client/telemetry/importTracker';
26+
import { IMultiStepInput, IMultiStepInputFactory } from '../../client/common/utils/multiStepInput';
27+
28+
// Class methods exposed just for testing purposes
29+
interface ITensorBoardSessionTestAPI {
30+
jumpToSource(fsPath: string, line: number): Promise<void>;
31+
}
2532

2633
const info: PythonEnvironment = {
2734
architecture: Architecture.Unknown,
@@ -99,47 +106,35 @@ suite('TensorBoard session creation', async () => {
99106
errorMessageStub = sandbox.stub(applicationShell, 'showErrorMessage');
100107
errorMessageStub.resolves(installPromptSelection);
101108
}
102-
suite('Core functionality', async () => {
103-
test('Golden path: TensorBoard session starts successfully and webview is shown', async () => {
104-
sandbox.stub(experimentService, 'inExperiment').resolves(true);
105-
errorMessageStub = sandbox.stub(applicationShell, 'showErrorMessage');
106-
// Stub user selections
107-
sandbox
108-
.stub(applicationShell, 'showQuickPick')
109-
.resolves({ label: TensorBoard.useCurrentWorkingDirectory() });
109+
async function createSession() {
110+
errorMessageStub = sandbox.stub(applicationShell, 'showErrorMessage');
111+
// Stub user selections
112+
sandbox.stub(applicationShell, 'showQuickPick').resolves({ label: TensorBoard.useCurrentWorkingDirectory() });
110113

111-
const session = (await commandManager.executeCommand(
112-
'python.launchTensorBoard',
113-
TensorBoardEntrypoint.palette,
114-
TensorBoardEntrypointTrigger.palette,
115-
)) as TensorBoardSession;
114+
const session = (await commandManager.executeCommand(
115+
'python.launchTensorBoard',
116+
TensorBoardEntrypoint.palette,
117+
TensorBoardEntrypointTrigger.palette,
118+
)) as TensorBoardSession;
116119

117-
assert.ok(session.panel?.viewColumn === ViewColumn.One, 'Panel opened in wrong group');
118-
assert.ok(session.panel?.visible, 'Webview panel not shown on session creation golden path');
119-
assert.ok(errorMessageStub.notCalled, 'Error message shown on session creation golden path');
120+
assert.ok(session.panel?.viewColumn === ViewColumn.One, 'Panel opened in wrong group');
121+
assert.ok(session.panel?.visible, 'Webview panel not shown on session creation golden path');
122+
assert.ok(errorMessageStub.notCalled, 'Error message shown on session creation golden path');
123+
return session;
124+
}
125+
suite('Core functionality', async () => {
126+
test('Golden path: TensorBoard session starts successfully and webview is shown', async () => {
127+
await createSession();
120128
});
121129
test('When webview is closed, session is killed', async () => {
122-
sandbox.stub(experimentService, 'inExperiment').resolves(true);
123-
errorMessageStub = sandbox.stub(applicationShell, 'showErrorMessage');
124-
// Stub user selections
125-
sandbox
126-
.stub(applicationShell, 'showQuickPick')
127-
.resolves({ label: TensorBoard.useCurrentWorkingDirectory() });
128-
129-
const session = (await commandManager.executeCommand(
130-
'python.launchTensorBoard',
131-
TensorBoardEntrypoint.palette,
132-
TensorBoardEntrypointTrigger.palette,
133-
)) as TensorBoardSession;
134-
130+
const session = await createSession();
135131
const { daemon, panel } = session;
136132
assert.ok(panel?.visible, 'Webview panel not shown');
137133
panel?.dispose();
138134
assert.ok(session.panel === undefined, 'Webview still visible');
139135
assert.ok(daemon?.killed, 'TensorBoard session process not killed after webview closed');
140136
});
141137
test('When user selects file picker, display file picker', async () => {
142-
sandbox.stub(experimentService, 'inExperiment').resolves(true);
143138
errorMessageStub = sandbox.stub(applicationShell, 'showErrorMessage');
144139
// Stub user selections
145140
sandbox.stub(applicationShell, 'showQuickPick').resolves({ label: TensorBoard.selectAnotherFolder() });
@@ -440,4 +435,59 @@ suite('TensorBoard session creation', async () => {
440435
'Prompted user to select log directory although setting was specified',
441436
);
442437
});
438+
suite('Jump to source', async () => {
439+
// We can't test a full E2E scenario with the TB profiler plugin because we can't
440+
// accurately target simulated clicks at iframed content. This only tests
441+
// code from the moment that the VS Code webview posts a message back
442+
// to the extension.
443+
const fsPath = path.join(
444+
EXTENSION_ROOT_DIR_FOR_TESTS,
445+
'src',
446+
'test',
447+
'pythonFiles',
448+
'tensorBoard',
449+
'sourcefile.py',
450+
);
451+
teardown(() => {
452+
sandbox.restore();
453+
});
454+
function setupStubsForMultiStepInput() {
455+
// Stub the factory to return our stubbed multistep input when it's asked to create one
456+
const multiStepFactory = serviceManager.get<IMultiStepInputFactory>(IMultiStepInputFactory);
457+
const inputInstance = multiStepFactory.create();
458+
// Create a multistep input with stubs for methods
459+
const showQuickPickStub = sandbox.stub(inputInstance, 'showQuickPick').resolves({
460+
label: TensorBoard.selectMissingSourceFile(),
461+
description: TensorBoard.selectMissingSourceFileDescription(),
462+
});
463+
const createInputStub = sandbox
464+
.stub(multiStepFactory, 'create')
465+
.returns(inputInstance as IMultiStepInput<unknown>);
466+
// Stub the system file picker
467+
const filePickerStub = sandbox.stub(applicationShell, 'showOpenDialog').resolves([Uri.file(fsPath)]);
468+
return [showQuickPickStub, createInputStub, filePickerStub];
469+
}
470+
test('Resolves filepaths without displaying prompt', async () => {
471+
const session = ((await createSession()) as unknown) as ITensorBoardSessionTestAPI;
472+
const stubs = setupStubsForMultiStepInput();
473+
await session.jumpToSource(fsPath, 0);
474+
assert.ok(window.activeTextEditor !== undefined, 'Source file not resolved');
475+
assert.ok(window.activeTextEditor?.document.uri.fsPath === fsPath, 'Wrong source file opened');
476+
assert.ok(
477+
stubs.reduce((prev, current) => current.notCalled && prev, true),
478+
'Stubs were called when file is present',
479+
);
480+
});
481+
test('Display quickpick to user if filepath is not on disk', async () => {
482+
const session = ((await createSession()) as unknown) as ITensorBoardSessionTestAPI;
483+
const stubs = setupStubsForMultiStepInput();
484+
await session.jumpToSource('/nonexistent/file/path.py', 0);
485+
assert.ok(window.activeTextEditor !== undefined, 'Source file not resolved');
486+
assert.ok(window.activeTextEditor?.document.uri.fsPath === fsPath, 'Wrong source file opened');
487+
assert.ok(
488+
stubs.reduce((prev, current) => current.calledOnce && prev, true),
489+
'Stubs called an unexpected number of times',
490+
);
491+
});
492+
});
443493
});

0 commit comments

Comments
 (0)