Implement jump to source from PyTorch profiler plugin in integrated TensorBoard session (#15631)

joyceerhl · web-flow · commit 6d1a75a2b17c · 2021-03-18T13:41:22.000-07:00
diff --git a/news/1 Enhancements/15641.md b/news/1 Enhancements/15641.md
@@ -0,0 +1 @@
+Add jump to source integration with the PyTorch profiler TensorBoard plugin during TensorBoard sessions.
diff --git a/src/client/tensorBoard/tensorBoardSession.ts b/src/client/tensorBoard/tensorBoardSession.ts
@@ -1,18 +1,22 @@
 // Copyright (c) Microsoft Corporation. All rights reserved.
 // Licensed under the MIT License.
-
+import * as fs from 'fs-extra';
 import { ChildProcess } from 'child_process';
 import * as path from 'path';
 import {
     CancellationToken,
     CancellationTokenSource,
+    Position,
     Progress,
     ProgressLocation,
     ProgressOptions,
     QuickPickItem,
+    Selection,
+    Uri,
     ViewColumn,
     WebviewPanel,
     window,
+    workspace,
 } from 'vscode';
 import { IApplicationShell, ICommandManager, IWorkspaceService } from '../common/application/types';
 import { createPromiseFromCancellation } from '../common/cancellation';
@@ -29,6 +33,10 @@ import { EventName } from '../telemetry/constants';
 import { ImportTracker } from '../telemetry/importTracker';
 import { TensorBoardPromptSelection, TensorBoardSessionStartResult } from './constants';
 
+enum Messages {
+    JumpToSource = 'jump_to_source',
+}
+
 /**
  * Manages the lifecycle of a TensorBoard session.
  * Specifically, it:
@@ -372,6 +380,7 @@ export class TensorBoardSession {
                         this.url = match[1];
                         urlThatTensorBoardIsRunningAt.resolve('success');
                     }
+                    traceInfo(output.out);
                 } else if (output.source === 'stderr') {
                     traceError(output.out);
                 }
@@ -405,15 +414,26 @@ export class TensorBoardSession {
             </head>
             <body>
                 <script type="text/javascript">
+                    const vscode = acquireVsCodeApi();
                     function resizeFrame() {
                         var f = window.document.getElementById('vscode-tensorboard-iframe');
                         if (f) {
-                            f.style.height = window.innerHeight / 0.7 + "px";
-                            f.style.width = window.innerWidth / 0.7 + "px";
+                            f.style.height = window.innerHeight / 0.8 + "px";
+                            f.style.width = window.innerWidth / 0.8 + "px";
                         }
                     }
                     resizeFrame();
+                    window.onload = function() {
+                        resizeFrame();
+                    }
                     window.addEventListener('resize', resizeFrame);
+                    window.addEventListener('message', (event) => {
+                        if (!"${this.url}".startsWith(event.origin) || !event.data || !event.data.filename || !event.data.line) {
+                            return;
+                        }
+                        const args = { filename: event.data.filename, line: event.data.line };
+                        vscode.postMessage({ command: '${Messages.JumpToSource}', args: args });
+                    });
                 </script>
                 <iframe
                     id="vscode-tensorboard-iframe"
@@ -426,7 +446,7 @@ export class TensorBoardSession {
                 ></iframe>
                 <style>
                     .responsive-iframe {
-                        transform: scale(0.7);
+                        transform: scale(0.8);
                         transform-origin: 0 0;
                         position: absolute;
                         top: 0;
@@ -447,6 +467,18 @@ export class TensorBoardSession {
                 this.process = undefined;
             }),
         );
+        this.disposables.push(
+            webviewPanel.webview.onDidReceiveMessage((message) => {
+                // Handle messages posted from the webview
+                switch (message.command) {
+                    case Messages.JumpToSource:
+                        jumpToSource(message.args.filename, message.args.line);
+                        break;
+                    default:
+                        break;
+                }
+            }),
+        );
         return webviewPanel;
     }
 
@@ -461,3 +493,21 @@ 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);
+                    editor.selection = new Selection(position, editor.document.lineAt(line).range.end);
+                }
+            });
+    } else {
+        traceError(`Requested jump to source filepath ${fsPath} does not exist`);
+    }
+}

Original file line number	Diff line number	Diff line change
`@@ -0,0 +1 @@`
	`1`	`+Add jump to source integration with the PyTorch profiler TensorBoard plugin during TensorBoard sessions.`