microsoft · joyceerhl · Apr 21, 2021 · Apr 20, 2021 · Apr 20, 2021 · Apr 20, 2021
diff --git a/news/1 Enhancements/15695.md b/news/1 Enhancements/15695.md
@@ -0,0 +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.
diff --git a/package.nls.json b/package.nls.json
@@ -231,6 +231,9 @@
     "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>",
     "StartPage.badWebPanelFormatString": "<html><body><h1>{0} is not a valid file name</h1></body></html>",
     "Jupyter.extensionRequired": "The Jupyter extension is required to perform that task. Click Yes to open the Jupyter extension installation page.",
+    "TensorBoard.missingSourceFile": "We could not locate the requested source file on disk. Please manually specify the file.",
+    "TensorBoard.selectMissingSourceFile": "Choose File",
+    "TensorBoard.selectMissingSourceFileDescription": "The source file's contents may not match the original contents in the trace.",
     "TensorBoard.useCurrentWorkingDirectory": "Use current working directory",
     "TensorBoard.currentDirectory": "Current: {0}",
     "TensorBoard.logDirectoryPrompt": "Select a log directory to start TensorBoard with",

diff --git a/pythonFiles/tensorboard_launcher.py b/pythonFiles/tensorboard_launcher.py
@@ -7,10 +7,7 @@
 
 def main(logdir):
     os.environ["VSCODE_TENSORBOARD_LAUNCH"] = "1"
-    tb = program.TensorBoard(
-        default.get_plugins(),
-        program.get_default_assets_zip_provider(),
-    )
+    tb = program.TensorBoard()
     tb.configure(bind_all=False, logdir=logdir)
     url = tb.launch()
     sys.stdout.write("TensorBoard started at %s\n" % (url))

diff --git a/src/client/common/utils/localize.ts b/src/client/common/utils/localize.ts
@@ -196,6 +196,15 @@ export namespace TensorBoard {
         'TensorBoard.launchNativeTensorBoardSessionCodeAction',
         'Launch TensorBoard session',
     );
+    export const missingSourceFile = localize(
+        'TensorBoard.missingSourceFile',
+        'We could not locate the requested source file on disk. Please manually specify the file.',
+    );
+    export const selectMissingSourceFile = localize('TensorBoard.selectMissingSourceFile', 'Choose File');
+    export const selectMissingSourceFileDescription = localize(
+        'TensorBoard.selectMissingSourceFileDescription',
+        "The source file's contents may not match the original contents in the trace.",
+    );
 }
 
 export namespace LanguageService {

diff --git a/src/client/tensorBoard/tensorBoardSession.ts b/src/client/tensorBoard/tensorBoardSession.ts
@@ -41,6 +41,7 @@ import { sendTelemetryEvent } from '../telemetry';
 import { EventName } from '../telemetry/constants';
 import { ImportTracker } from '../telemetry/importTracker';
 import { TensorBoardPromptSelection, TensorBoardSessionStartResult } from './constants';
+import { IMultiStepInputFactory } from '../common/utils/multiStepInput';
 
 enum Messages {
     JumpToSource = 'jump_to_source',
@@ -86,6 +87,7 @@ export class TensorBoardSession {
         private readonly applicationShell: IApplicationShell,
         private readonly isInTorchProfilerExperiment: boolean,
         private readonly globalMemento: IPersistentState<ViewColumn>,
+        private readonly multiStepFactory: IMultiStepInputFactory,
     ) {}
 
     public async initialize(): Promise<void> {
@@ -418,7 +420,6 @@ export class TensorBoardSession {
     private createPanel() {
         const webviewPanel = window.createWebviewPanel('tensorBoardSession', 'TensorBoard', this.globalMemento.value, {
             enableScripts: true,
-            retainContextWhenHidden: true,
         });
         webviewPanel.webview.html = `<!DOCTYPE html>
         <html lang="en">
@@ -497,7 +498,7 @@ export class TensorBoardSession {
                 // Handle messages posted from the webview
                 switch (message.command) {
                     case Messages.JumpToSource:
-                        jumpToSource(message.args.filename, message.args.line);
+                        void this.jumpToSource(message.args.filename, message.args.line);
                         break;
                     default:
                         break;
@@ -517,24 +518,58 @@ export class TensorBoardSession {
         }
         return undefined;
     }
-}
 
-function jumpToSource(fsPath: string, line: number) {
-    if (fs.existsSync(fsPath)) {
-        const uri = Uri.file(fsPath);
-        workspace
-            .openTextDocument(uri)
-            .then((doc) => window.showTextDocument(doc, ViewColumn.Beside))
-            .then((editor) => {
-                // Select the line if it exists in the document
-                if (line < editor.document.lineCount) {
-                    const position = new Position(line, 0);
-                    const selection = new Selection(position, editor.document.lineAt(line).range.end);
-                    editor.selection = selection;
-                    editor.revealRange(selection, TextEditorRevealType.InCenterIfOutsideViewport);
+    private async jumpToSource(fsPath: string, line: number) {
+        let uri: Uri | undefined;
+        if (fs.existsSync(fsPath)) {
+            uri = Uri.file(fsPath);
+        } else {
+            traceError(
+                `Requested jump to source filepath ${fsPath} does not exist. Prompting user to select source file...`,
+            );
+            // Prompt the user to pick the file on disk
+            const items: QuickPickItem[] = [
+                {
+                    label: TensorBoard.selectMissingSourceFile(),
+                    description: TensorBoard.selectMissingSourceFileDescription(),
+                },
+            ];
+            // Using a multistep so that we can add a title to the quickpick
+            const multiStep = this.multiStepFactory.create<unknown>();
+            await multiStep.run(async (input) => {
+                const selection = await input.showQuickPick({
+                    items,
+                    title: TensorBoard.missingSourceFile(),
+                    placeholder: fsPath,
+                });
+                switch (selection?.label) {
+                    case TensorBoard.selectMissingSourceFile(): {
+                        const filePickerSelection = await this.applicationShell.showOpenDialog({
+                            canSelectFiles: true,
+                            canSelectFolders: false,
+                            canSelectMany: false,
+                        });
+                        if (filePickerSelection !== undefined) {
+                            [uri] = filePickerSelection;
+                        }
+                        break;
+                    }
+                    default:
+                        break;
                 }
-            });
-    } else {
-        traceError(`Requested jump to source filepath ${fsPath} does not exist`);
+            }, {});
+        }
+        if (uri === undefined) {
+            return;
+        }
+        const document = await workspace.openTextDocument(uri);
+        const editor = await window.showTextDocument(document, ViewColumn.Beside);
+        // Select the line if it exists in the document
+        if (line < editor.document.lineCount) {
+            const position = new Position(line, 0);
+            const selection = new Selection(position, editor.document.lineAt(line).range.end);
+            editor.selection = selection;
+            editor.revealRange(selection, TextEditorRevealType.InCenterIfOutsideViewport);
+        }
     }
 }
diff --git a/src/client/tensorBoard/tensorBoardSessionProvider.ts b/src/client/tensorBoard/tensorBoardSessionProvider.ts
@@ -11,6 +11,7 @@ import { traceError, traceInfo } from '../common/logger';
 import { IProcessServiceFactory } from '../common/process/types';
 import { IDisposableRegistry, IExperimentService, IInstaller, IPersistentStateFactory } from '../common/types';
 import { TensorBoard } from '../common/utils/localize';
+import { IMultiStepInputFactory } from '../common/utils/multiStepInput';
 import { IInterpreterService } from '../interpreter/contracts';
 import { sendTelemetryEvent } from '../telemetry';
 import { EventName } from '../telemetry/constants';
@@ -31,6 +32,7 @@ export class TensorBoardSessionProvider implements IExtensionSingleActivationSer
         @inject(IProcessServiceFactory) private readonly processServiceFactory: IProcessServiceFactory,
         @inject(IExperimentService) private readonly experimentService: IExperimentService,
         @inject(IPersistentStateFactory) private stateFactory: IPersistentStateFactory,
+        @inject(IMultiStepInputFactory) private readonly multiStepFactory: IMultiStepInputFactory,
     ) {}
 
     public async activate(): Promise<void> {
@@ -68,6 +70,7 @@ export class TensorBoardSessionProvider implements IExtensionSingleActivationSer
                 this.applicationShell,
                 await this.experimentService.inExperiment(TorchProfiler.experiment),
                 memento,
+                this.multiStepFactory,
             );
             await newSession.initialize();
             return newSession;

diff --git a/src/test/pythonFiles/tensorBoard/sourcefile.py b/src/test/pythonFiles/tensorBoard/sourcefile.py
@@ -0,0 +1 @@
+from torch.utils.tensorboard import SummaryWriter
diff --git a/src/test/tensorBoard/tensorBoardSession.test.ts b/src/test/tensorBoard/tensorBoardSession.test.ts
@@ -1,7 +1,8 @@
+import * as path from 'path';
 import { assert } from 'chai';
 import Sinon, * as sinon from 'sinon';
 import { SemVer } from 'semver';
-import { ViewColumn, workspace, WorkspaceConfiguration } from 'vscode';
+import { Uri, ViewColumn, window, workspace, WorkspaceConfiguration } from 'vscode';
 import {
     IExperimentService,
     IInstaller,
@@ -14,14 +15,20 @@ import { IApplicationShell, ICommandManager } from '../../client/common/applicat
 import { IServiceManager } from '../../client/ioc/types';
 import { TensorBoardEntrypoint, TensorBoardEntrypointTrigger } from '../../client/tensorBoard/constants';
 import { TensorBoardSession } from '../../client/tensorBoard/tensorBoardSession';
-import { closeActiveWindows, initialize } from '../initialize';
+import { closeActiveWindows, EXTENSION_ROOT_DIR_FOR_TESTS, initialize } from '../initialize';
 import * as ExperimentHelpers from '../../client/common/experiments/helpers';
 import { IInterpreterService } from '../../client/interpreter/contracts';
 import { Architecture } from '../../client/common/utils/platform';
 import { PythonEnvironment, EnvironmentType } from '../../client/pythonEnvironments/info';
 import { PYTHON_PATH } from '../common';
 import { TorchProfiler } from '../../client/common/experiments/groups';
 import { ImportTracker } from '../../client/telemetry/importTracker';
+import { IMultiStepInput, IMultiStepInputFactory } from '../../client/common/utils/multiStepInput';
+
+// Class methods exposed just for testing purposes
+interface ITensorBoardSessionTestAPI {
+    jumpToSource(fsPath: string, line: number): Promise<void>;
+}
 
 const info: PythonEnvironment = {
     architecture: Architecture.Unknown,
@@ -99,47 +106,35 @@ suite('TensorBoard session creation', async () => {
         errorMessageStub = sandbox.stub(applicationShell, 'showErrorMessage');
         errorMessageStub.resolves(installPromptSelection);
     }
-    suite('Core functionality', async () => {
-        test('Golden path: TensorBoard session starts successfully and webview is shown', async () => {
-            sandbox.stub(experimentService, 'inExperiment').resolves(true);
-            errorMessageStub = sandbox.stub(applicationShell, 'showErrorMessage');
-            // Stub user selections
-            sandbox
-                .stub(applicationShell, 'showQuickPick')
-                .resolves({ label: TensorBoard.useCurrentWorkingDirectory() });
+    async function createSession() {
+        errorMessageStub = sandbox.stub(applicationShell, 'showErrorMessage');
+        // Stub user selections
+        sandbox.stub(applicationShell, 'showQuickPick').resolves({ label: TensorBoard.useCurrentWorkingDirectory() });
 
-            const session = (await commandManager.executeCommand(
-                'python.launchTensorBoard',
-                TensorBoardEntrypoint.palette,
-                TensorBoardEntrypointTrigger.palette,
-            )) as TensorBoardSession;
+        const session = (await commandManager.executeCommand(
+            'python.launchTensorBoard',
+            TensorBoardEntrypoint.palette,
+            TensorBoardEntrypointTrigger.palette,
+        )) as TensorBoardSession;
 
-            assert.ok(session.panel?.viewColumn === ViewColumn.One, 'Panel opened in wrong group');
-            assert.ok(session.panel?.visible, 'Webview panel not shown on session creation golden path');
-            assert.ok(errorMessageStub.notCalled, 'Error message shown on session creation golden path');
+        assert.ok(session.panel?.viewColumn === ViewColumn.One, 'Panel opened in wrong group');
+        assert.ok(session.panel?.visible, 'Webview panel not shown on session creation golden path');
+        assert.ok(errorMessageStub.notCalled, 'Error message shown on session creation golden path');
+        return session;
+    }
+    suite('Core functionality', async () => {
+        test('Golden path: TensorBoard session starts successfully and webview is shown', async () => {
+            await createSession();
         });
         test('When webview is closed, session is killed', async () => {
-            sandbox.stub(experimentService, 'inExperiment').resolves(true);
-            errorMessageStub = sandbox.stub(applicationShell, 'showErrorMessage');
-            // Stub user selections
-            sandbox
-                .stub(applicationShell, 'showQuickPick')
-                .resolves({ label: TensorBoard.useCurrentWorkingDirectory() });
-
-            const session = (await commandManager.executeCommand(
-                'python.launchTensorBoard',
-                TensorBoardEntrypoint.palette,
-                TensorBoardEntrypointTrigger.palette,
-            )) as TensorBoardSession;
-
+            const session = await createSession();
             const { daemon, panel } = session;
             assert.ok(panel?.visible, 'Webview panel not shown');
             panel?.dispose();
             assert.ok(session.panel === undefined, 'Webview still visible');
             assert.ok(daemon?.killed, 'TensorBoard session process not killed after webview closed');
         });
         test('When user selects file picker, display file picker', async () => {
-            sandbox.stub(experimentService, 'inExperiment').resolves(true);
             errorMessageStub = sandbox.stub(applicationShell, 'showErrorMessage');
             // Stub user selections
             sandbox.stub(applicationShell, 'showQuickPick').resolves({ label: TensorBoard.selectAnotherFolder() });
@@ -440,4 +435,59 @@ suite('TensorBoard session creation', async () => {
             'Prompted user to select log directory although setting was specified',
         );
     });
+    suite('Jump to source', async () => {
+        // We can't test a full E2E scenario with the TB profiler plugin because we can't
+        // accurately target simulated clicks at iframed content. This only tests
+        // code from the moment that the VS Code webview posts a message back
+        // to the extension.
+        const fsPath = path.join(
+            EXTENSION_ROOT_DIR_FOR_TESTS,
+            'src',
+            'test',
+            'pythonFiles',
+            'tensorBoard',
+            'sourcefile.py',
+        );
+        teardown(() => {
+            sandbox.restore();
+        });
+        function setupStubsForMultiStepInput() {
+            // Stub the factory to return our stubbed multistep input when it's asked to create one
+            const multiStepFactory = serviceManager.get<IMultiStepInputFactory>(IMultiStepInputFactory);
+            const inputInstance = multiStepFactory.create();
+            // Create a multistep input with stubs for methods
+            const showQuickPickStub = sandbox.stub(inputInstance, 'showQuickPick').resolves({
+                label: TensorBoard.selectMissingSourceFile(),
+                description: TensorBoard.selectMissingSourceFileDescription(),
+            });
+            const createInputStub = sandbox
+                .stub(multiStepFactory, 'create')
+                .returns(inputInstance as IMultiStepInput<unknown>);
+            // Stub the system file picker
+            const filePickerStub = sandbox.stub(applicationShell, 'showOpenDialog').resolves([Uri.file(fsPath)]);
+            return [showQuickPickStub, createInputStub, filePickerStub];
+        }
+        test('Resolves filepaths without displaying prompt', async () => {
+            const session = ((await createSession()) as unknown) as ITensorBoardSessionTestAPI;
+            const stubs = setupStubsForMultiStepInput();
+            await session.jumpToSource(fsPath, 0);
+            assert.ok(window.activeTextEditor !== undefined, 'Source file not resolved');
+            assert.ok(window.activeTextEditor?.document.uri.fsPath === fsPath, 'Wrong source file opened');
+            assert.ok(
+                stubs.reduce((prev, current) => current.notCalled && prev, true),
+                'Stubs were called when file is present',
+            );
+        });
+        test('Display quickpick to user if filepath is not on disk', async () => {
+            const session = ((await createSession()) as unknown) as ITensorBoardSessionTestAPI;
+            const stubs = setupStubsForMultiStepInput();
+            await session.jumpToSource('/nonexistent/file/path.py', 0);
+            assert.ok(window.activeTextEditor !== undefined, 'Source file not resolved');
+            assert.ok(window.activeTextEditor?.document.uri.fsPath === fsPath, 'Wrong source file opened');
+            assert.ok(
+                stubs.reduce((prev, current) => current.calledOnce && prev, true),
+                'Stubs called an unexpected number of times',
+            );
+        });
+    });
 });
Original file line number	Diff line number	Diff line change
		@@ -0,0 +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.
Original file line number	Diff line number	Diff line change
		@@ -0,0 +1 @@
		from torch.utils.tensorboard import SummaryWriter