Add Refresh TensorBoard command (#16068)

microsoft · May 4, 2021 · 812be6c · 812be6c
1 parent eef545b
commit 812be6c
Show file tree

Hide file tree

Showing 6 changed files with 168 additions and 64 deletions.
diff --git a/news/1 Enhancements/16053.md b/news/1 Enhancements/16053.md
@@ -0,0 +1 @@
+Add Python: Refresh TensorBoard command, keybinding and editor title button to reload TensorBoard (equivalent to browser refresh).
diff --git a/package.json b/package.json
@@ -100,9 +100,22 @@
                 "command": "python.execSelectionInTerminal",
                 "key": "shift+enter",
                 "when": "editorTextFocus && editorLangId == python && !findInputFocussed && !replaceInputFocussed && !jupyter.ownsSelection && !notebookEditorFocused"
+            },
+            {
+                "command": "python.refreshTensorBoard",
+                "key": "ctrl+r",
+                "mac": "cmd+r",
+                "when": "python.hasActiveTensorBoardSession"
             }
         ],
         "commands": [
+            {
+                "command": "python.refreshTensorBoard",
+                "title": "%python.command.python.refreshTensorBoard.title%",
+                "category": "Python",
+                "enablement": "python.hasActiveTensorBoardSession",
+                "icon": "$(refresh)"
+            },
             {
                 "command": "python.clearPersistentStorage",
                 "title": "%python.command.python.clearPersistentStorage.title%",
@@ -419,6 +432,13 @@
                     "when": "resourceLangId == python && !isInDiffEditor"
                 }
             ],
+            "editor/title": [
+                {
+                    "command": "python.refreshTensorBoard",
+                    "group": "navigation@0",
+                    "when": "python.hasActiveTensorBoardSession"
+                }
+            ],
             "explorer/context": [
                 {
                     "when": "resourceLangId == python && !busyTests && !notebookEditorFocused",

diff --git a/package.nls.json b/package.nls.json
@@ -39,6 +39,7 @@
     "python.command.python.analysis.clearCache.title": "Clear Module Analysis Cache",
     "python.command.python.analysis.restartLanguageServer.title": "Restart Language Server",
     "python.command.python.launchTensorBoard.title": "Launch TensorBoard",
+    "python.command.python.refreshTensorBoard.title": "Refresh TensorBoard",
     "python.snippet.launch.standard.label": "Python: Current File",
     "python.snippet.launch.module.label": "Python: Module",
     "python.snippet.launch.module.default": "enter-your-module-name",

diff --git a/src/client/common/constants.ts b/src/client/common/constants.ts
@@ -68,6 +68,7 @@ export namespace Commands {
     export const ResetInterpreterSecurityStorage = 'python.resetInterpreterSecurityStorage';
     export const OpenStartPage = 'python.startPage.open';
     export const LaunchTensorBoard = 'python.launchTensorBoard';
+    export const RefreshTensorBoard = 'python.refreshTensorBoard';
 }
 export namespace Octicons {
     export const Test_Pass = '$(check)';

diff --git a/src/client/tensorBoard/tensorBoardSession.ts b/src/client/tensorBoard/tensorBoardSession.ts
@@ -7,6 +7,8 @@ import {
     CancellationToken,
     CancellationTokenSource,
     env,
+    Event,
+    EventEmitter,
     Position,
     Progress,
     ProgressLocation,
@@ -67,14 +69,18 @@ export class TensorBoardSession {
         return this.process;
     }
 
-    private active = false;
+    private _active = false;
 
     private webviewPanel: WebviewPanel | undefined;
 
     private url: string | undefined;
 
     private process: ChildProcess | undefined;
 
+    private onDidChangeViewStateEventEmitter = new EventEmitter<void>();
+
+    private onDidDisposeEventEmitter = new EventEmitter<TensorBoardSession>();
+
     // This tracks the total duration of time that the user kept the TensorBoard panel open
     private sessionDurationStopwatch: StopWatch | undefined;
 
@@ -91,6 +97,26 @@ export class TensorBoardSession {
         private readonly multiStepFactory: IMultiStepInputFactory,
     ) {}
 
+    public get onDidDispose(): Event<TensorBoardSession> {
+        return this.onDidDisposeEventEmitter.event;
+    }
+
+    public get onDidChangeViewState(): Event<void> {
+        return this.onDidChangeViewStateEventEmitter.event;
+    }
+
+    public get active(): boolean {
+        return this._active;
+    }
+
+    public async refresh(): Promise<void> {
+        if (!this.webviewPanel) {
+            return;
+        }
+        this.webviewPanel.webview.html = '';
+        this.webviewPanel.webview.html = await this.getHtml();
+    }
+
     public async initialize(): Promise<void> {
         const e2eStartupDurationStopwatch = new StopWatch();
         const tensorBoardWasInstalled = await this.ensurePrerequisitesAreInstalled();
@@ -415,67 +441,15 @@ export class TensorBoardSession {
         traceInfo('Showing TensorBoard panel');
         const panel = this.webviewPanel || (await this.createPanel());
         panel.reveal();
-        this.active = true;
+        this._active = true;
+        this.onDidChangeViewStateEventEmitter.fire();
     }
 
     private async createPanel() {
         const webviewPanel = window.createWebviewPanel('tensorBoardSession', 'TensorBoard', this.globalMemento.value, {
             enableScripts: true,
         });
-        const fullWebServerUri = await env.asExternalUri(Uri.parse(this.url!));
-        webviewPanel.webview.html = `<!DOCTYPE html>
-        <html lang="en">
-            <head>
-                <meta charset="UTF-8">
-                <meta http-equiv="Content-Security-Policy" content="default-src 'unsafe-inline'; frame-src ${fullWebServerUri} http: https:;">
-                <meta name="viewport" content="width=device-width, initial-scale=1.0">
-                <title>TensorBoard</title>
-            </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.8 + "px";
-                            f.style.width = window.innerWidth / 0.8 + "px";
-                        }
-                    }
-                    resizeFrame();
-                    window.onload = function() {
-                        resizeFrame();
-                    }
-                    window.addEventListener('resize', resizeFrame);
-                    window.addEventListener('message', (event) => {
-                        if (!"${fullWebServerUri}".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"
-                    class="responsive-iframe"
-                    sandbox="allow-scripts allow-forms allow-same-origin allow-pointer-lock"
-                    src="${fullWebServerUri}"
-                    frameborder="0"
-                    border="0"
-                    allowfullscreen
-                ></iframe>
-                <style>
-                    .responsive-iframe {
-                        transform: scale(0.8);
-                        transform-origin: 0 0;
-                        position: absolute;
-                        top: 0;
-                        left: 0;
-                        overflow: hidden;
-                        display: block;
-                    }
-                </style>
-            </body>
-        </html>`;
+        webviewPanel.webview.html = await this.getHtml();
         this.webviewPanel = webviewPanel;
         this.disposables.push(
             webviewPanel.onDidDispose(() => {
@@ -484,6 +458,8 @@ export class TensorBoardSession {
                 this.process?.kill();
                 sendTelemetryEvent(EventName.TENSORBOARD_SESSION_DURATION, this.sessionDurationStopwatch?.elapsedTime);
                 this.process = undefined;
+                this._active = false;
+                this.onDidDisposeEventEmitter.fire(this);
             }),
         );
         this.disposables.push(
@@ -492,7 +468,8 @@ export class TensorBoardSession {
                 if (this.active && args.webviewPanel.active) {
                     await this.globalMemento.updateValue(webviewPanel.viewColumn ?? ViewColumn.Active);
                 }
-                this.active = args.webviewPanel.active;
+                this._active = args.webviewPanel.active;
+                this.onDidChangeViewStateEventEmitter.fire();
             }),
         );
         this.disposables.push(
@@ -574,4 +551,69 @@ export class TensorBoardSession {
             editor.revealRange(selection, TextEditorRevealType.InCenterIfOutsideViewport);
         }
     }
+
+    private async getHtml() {
+        // We cannot cache the result of calling asExternalUri, so regenerate
+        // it each time. From docs: "Note that extensions should not cache the
+        // result of asExternalUri as the resolved uri may become invalid due
+        // to a system or user action — for example, in remote cases, a user may
+        // close a port forwarding tunnel that was opened by asExternalUri."
+        const fullWebServerUri = await env.asExternalUri(Uri.parse(this.url!));
+        return `<!DOCTYPE html>
+        <html lang="en">
+            <head>
+                <meta charset="UTF-8">
+                <meta http-equiv="Content-Security-Policy" content="default-src 'unsafe-inline'; frame-src ${fullWebServerUri} http: https:;">
+                <meta name="viewport" content="width=device-width, initial-scale=1.0">
+                <title>TensorBoard</title>
+            </head>
+            <body>
+                <script type="text/javascript">
+                    (function() {
+                        const vscode = acquireVsCodeApi();
+                        function resizeFrame() {
+                            var f = window.document.getElementById('vscode-tensorboard-iframe');
+                            if (f) {
+                                f.style.height = window.innerHeight / 0.8 + "px";
+                                f.style.width = window.innerWidth / 0.8 + "px";
+                            }
+                        }
+                        window.onload = function() {
+                            resizeFrame();
+                        }
+                        window.addEventListener('resize', resizeFrame);
+                        window.addEventListener('message', (event) => {
+                            if (!"${fullWebServerUri}".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"
+                    class="responsive-iframe"
+                    sandbox="allow-scripts allow-forms allow-same-origin allow-pointer-lock"
+                    src="${fullWebServerUri}"
+                    frameborder="0"
+                    border="0"
+                    allowfullscreen
+                ></iframe>
+                <style>
+                    .responsive-iframe {
+                        transform: scale(0.8);
+                        transform-origin: 0 0;
+                        position: absolute;
+                        top: 0;
+                        left: 0;
+                        overflow: hidden;
+                        display: block;
+                        width: 100%;
+                        height: 100%;
+                    }
+                </style>
+            </body>
+        </html>`;
+    }
 }
diff --git a/src/client/tensorBoard/tensorBoardSessionProvider.ts b/src/client/tensorBoard/tensorBoardSessionProvider.ts
@@ -6,10 +6,17 @@ import { ViewColumn } from 'vscode';
 import { IExtensionSingleActivationService } from '../activation/types';
 import { IApplicationShell, ICommandManager, IWorkspaceService } from '../common/application/types';
 import { Commands } from '../common/constants';
+import { ContextKey } from '../common/contextKey';
 import { TorchProfiler } from '../common/experiments/groups';
 import { traceError, traceInfo } from '../common/logger';
 import { IProcessServiceFactory } from '../common/process/types';
-import { IDisposableRegistry, IExperimentService, IInstaller, IPersistentStateFactory } from '../common/types';
+import {
+    IDisposableRegistry,
+    IExperimentService,
+    IInstaller,
+    IPersistentState,
+    IPersistentStateFactory,
+} from '../common/types';
 import { TensorBoard } from '../common/utils/localize';
 import { IMultiStepInputFactory } from '../common/utils/multiStepInput';
 import { IInterpreterService } from '../interpreter/contracts';
@@ -22,6 +29,12 @@ const PREFERRED_VIEWGROUP = 'PythonTensorBoardWebviewPreferredViewGroup';
 
 @injectable()
 export class TensorBoardSessionProvider implements IExtensionSingleActivationService {
+    private knownSessions: TensorBoardSession[] = [];
+
+    private preferredViewGroupMemento: IPersistentState<ViewColumn>;
+
+    private hasActiveTensorBoardSessionContext: ContextKey;
+
     constructor(
         @inject(IInstaller) private readonly installer: IInstaller,
         @inject(IInterpreterService) private readonly interpreterService: IInterpreterService,
@@ -33,7 +46,16 @@ export class TensorBoardSessionProvider implements IExtensionSingleActivationSer
         @inject(IExperimentService) private readonly experimentService: IExperimentService,
         @inject(IPersistentStateFactory) private stateFactory: IPersistentStateFactory,
         @inject(IMultiStepInputFactory) private readonly multiStepFactory: IMultiStepInputFactory,
-    ) {}
+    ) {
+        this.preferredViewGroupMemento = this.stateFactory.createGlobalPersistentState<ViewColumn>(
+            PREFERRED_VIEWGROUP,
+            ViewColumn.Active,
+        );
+        this.hasActiveTensorBoardSessionContext = new ContextKey(
+            'python.hasActiveTensorBoardSession',
+            this.commandManager,
+        );
+    }
 
     public async activate(): Promise<void> {
         this.disposables.push(
@@ -50,16 +72,30 @@ export class TensorBoardSessionProvider implements IExtensionSingleActivationSer
                     return this.createNewSession();
                 },
             ),
+            this.commandManager.registerCommand(Commands.RefreshTensorBoard, () =>
+                this.knownSessions.map((w) => w.refresh()),
+            ),
         );
     }
 
+    private async updateTensorBoardSessionContext() {
+        let hasActiveTensorBoardSession = false;
+        this.knownSessions.forEach((viewer) => {
+            if (viewer.active) {
+                hasActiveTensorBoardSession = true;
+            }
+        });
+        await this.hasActiveTensorBoardSessionContext.set(hasActiveTensorBoardSession);
+    }
+
+    private async didDisposeSession(session: TensorBoardSession) {
+        this.knownSessions = this.knownSessions.filter((s) => s !== session);
+        this.updateTensorBoardSessionContext();
+    }
+
     private async createNewSession(): Promise<TensorBoardSession | undefined> {
         traceInfo('Starting new TensorBoard session...');
         try {
-            const memento = this.stateFactory.createGlobalPersistentState<ViewColumn>(
-                PREFERRED_VIEWGROUP,
-                ViewColumn.Active,
-            );
             const newSession = new TensorBoardSession(
                 this.installer,
                 this.interpreterService,
@@ -69,9 +105,12 @@ export class TensorBoardSessionProvider implements IExtensionSingleActivationSer
                 this.disposables,
                 this.applicationShell,
                 await this.experimentService.inExperiment(TorchProfiler.experiment),
-                memento,
+                this.preferredViewGroupMemento,
                 this.multiStepFactory,
             );
+            newSession.onDidChangeViewState(() => this.updateTensorBoardSessionContext(), this, this.disposables);
+            newSession.onDidDispose((e) => this.didDisposeSession(e), this, this.disposables);
+            this.knownSessions.push(newSession);
             await newSession.initialize();
             return newSession;
         } catch (e) {
Original file line number	Diff line number	Diff line change
		@@ -0,0 +1 @@
		Add Python: Refresh TensorBoard command, keybinding and editor title button to reload TensorBoard (equivalent to browser refresh).