Add support for a tensorboard experiment (#22215)

package.json

@@ -537,15 +537,17 @@
                             "pythonPromptNewToolsExt",
                             "pythonTerminalEnvVarActivation",
                             "pythonTestAdapter",
-                            "pythonREPLSmartSend"
+                            "pythonREPLSmartSend",
+                            "pythonRecommendTensorboardExt"
                         ],
                         "enumDescriptions": [
                             "%python.experiments.All.description%",
                             "%python.experiments.pythonSurveyNotification.description%",
                             "%python.experiments.pythonPromptNewToolsExt.description%",
                             "%python.experiments.pythonTerminalEnvVarActivation.description%",
                             "%python.experiments.pythonTestAdapter.description%",
-                            "%python.experiments.pythonREPLSmartSend.description%"
+                            "%python.experiments.pythonREPLSmartSend.description%",
+                            "%python.experiments.pythonRecommendTensorboardExt.description%"
                         ]
                     },
                     "scope": "machine",

package.nls.json

@@ -42,6 +42,7 @@
     "python.experiments.pythonTerminalEnvVarActivation.description": "Enables use of environment variables to activate terminals instead of sending activation commands.",
     "python.experiments.pythonTestAdapter.description": "Denotes the Python Test Adapter experiment.",
     "python.experiments.pythonREPLSmartSend.description": "Denotes the Python REPL Smart Send experiment.",
+    "python.experiments.pythonRecommendTensorboardExt.description": "Denotes the Tensorboard Extension recommendation experiment.",
     "python.globalModuleInstallation.description": "Whether to install Python modules globally when not using an environment.",
     "python.languageServer.description": "Defines type of the language server.",
     "python.languageServer.defaultDescription": "Automatically select a language server: Pylance if installed and available, otherwise fallback to Jedi.",

src/client/common/experiments/groups.ts

@@ -22,3 +22,8 @@ export enum EnableTestAdapterRewrite {
 export enum EnableREPLSmartSend {
     experiment = 'pythonREPLSmartSend',
 }
+
+// Experiment to recommend installing the tensorboard extension.
+export enum RecommendTensobardExtension {
+    experiment = 'pythonRecommendTensorboardExt',
+}

src/client/tensorBoard/nbextensionCodeLensProvider.ts

@@ -3,34 +3,46 @@
 
 import { inject, injectable } from 'inversify';
 import { once } from 'lodash';
-import { CancellationToken, CodeLens, Command, languages, Position, Range, TextDocument } from 'vscode';
+import { CancellationToken, CodeLens, Command, Disposable, languages, Position, Range, TextDocument } from 'vscode';
 import { IExtensionSingleActivationService } from '../activation/types';
 import { Commands, NotebookCellScheme, PYTHON_LANGUAGE } from '../common/constants';
-import { IDisposableRegistry } from '../common/types';
+import { IDisposable, IDisposableRegistry } from '../common/types';
 import { TensorBoard } from '../common/utils/localize';
 import { sendTelemetryEvent } from '../telemetry';
 import { EventName } from '../telemetry/constants';
 import { TensorBoardEntrypoint, TensorBoardEntrypointTrigger } from './constants';
 import { containsNotebookExtension } from './helpers';
-import { useNewTensorboardExtension } from './tensorboarExperiment';
+import { TensorboardExperiment } from './tensorboarExperiment';
 
 @injectable()
 export class TensorBoardNbextensionCodeLensProvider implements IExtensionSingleActivationService {
     public readonly supportedWorkspaceTypes = { untrustedWorkspace: false, virtualWorkspace: false };
 
+    private readonly disposables: IDisposable[] = [];
+
     private sendTelemetryOnce = once(
         sendTelemetryEvent.bind(this, EventName.TENSORBOARD_ENTRYPOINT_SHOWN, undefined, {
             trigger: TensorBoardEntrypointTrigger.nbextension,
             entrypoint: TensorBoardEntrypoint.codelens,
         }),
     );
 
-    constructor(@inject(IDisposableRegistry) private disposables: IDisposableRegistry) {}
+    constructor(
+        @inject(IDisposableRegistry) disposables: IDisposableRegistry,
+        @inject(TensorboardExperiment) private readonly experiment: TensorboardExperiment,
+    ) {
+        disposables.push(this);
+    }
+
+    public dispose(): void {
+        Disposable.from(...this.disposables).dispose();
+    }
 
     public async activate(): Promise<void> {
-        if (useNewTensorboardExtension()) {
+        if (TensorboardExperiment.isTensorboardExtensionInstalled) {
             return;
         }
+        this.experiment.disposeOnInstallingTensorboard(this);
         this.activateInternal().ignoreErrors();
     }
 

src/client/tensorBoard/serviceRegistry.ts

@@ -11,6 +11,7 @@ import { TensorBoardSessionProvider } from './tensorBoardSessionProvider';
 import { TensorBoardNbextensionCodeLensProvider } from './nbextensionCodeLensProvider';
 import { TerminalWatcher } from './terminalWatcher';
 import { TensorboardDependencyChecker } from './tensorboardDependencyChecker';
+import { TensorboardExperiment } from './tensorboarExperiment';
 
 export function registerTypes(serviceManager: IServiceManager): void {
     serviceManager.addSingleton<TensorBoardSessionProvider>(TensorBoardSessionProvider, TensorBoardSessionProvider);
@@ -34,4 +35,5 @@ export function registerTypes(serviceManager: IServiceManager): void {
     serviceManager.addBinding(TensorBoardNbextensionCodeLensProvider, IExtensionSingleActivationService);
     serviceManager.addSingleton(IExtensionSingleActivationService, TerminalWatcher);
     serviceManager.addSingleton(TensorboardDependencyChecker, TensorboardDependencyChecker);
+    serviceManager.addSingleton(TensorboardExperiment, TensorboardExperiment);
 }

src/client/tensorBoard/tensorBoardFileWatcher.ts

@@ -2,13 +2,13 @@
 // Licensed under the MIT License.
 
 import { inject, injectable } from 'inversify';
-import { FileSystemWatcher, RelativePattern, WorkspaceFolder, WorkspaceFoldersChangeEvent } from 'vscode';
+import { Disposable, FileSystemWatcher, RelativePattern, WorkspaceFolder, WorkspaceFoldersChangeEvent } from 'vscode';
 import { IExtensionSingleActivationService } from '../activation/types';
 import { IWorkspaceService } from '../common/application/types';
-import { IDisposableRegistry } from '../common/types';
+import { IDisposable, IDisposableRegistry } from '../common/types';
 import { TensorBoardEntrypointTrigger } from './constants';
 import { TensorBoardPrompt } from './tensorBoardPrompt';
-import { useNewTensorboardExtension } from './tensorboarExperiment';
+import { TensorboardExperiment } from './tensorboarExperiment';
 
 @injectable()
 export class TensorBoardFileWatcher implements IExtensionSingleActivationService {
@@ -18,16 +18,26 @@ export class TensorBoardFileWatcher implements IExtensionSingleActivationService
 
     private globPatterns = ['*tfevents*', '*/*tfevents*', '*/*/*tfevents*'];
 
+    private readonly disposables: IDisposable[] = [];
+
     constructor(
         @inject(IWorkspaceService) private workspaceService: IWorkspaceService,
         @inject(TensorBoardPrompt) private tensorBoardPrompt: TensorBoardPrompt,
-        @inject(IDisposableRegistry) private readonly disposables: IDisposableRegistry,
-    ) {}
+        @inject(IDisposableRegistry) disposables: IDisposableRegistry,
+        @inject(TensorboardExperiment) private readonly experiment: TensorboardExperiment,
+    ) {
+        disposables.push(this);
+    }
+
+    public dispose(): void {
+        Disposable.from(...this.disposables).dispose();
+    }
 
     public async activate(): Promise<void> {
-        if (useNewTensorboardExtension()) {
+        if (TensorboardExperiment.isTensorboardExtensionInstalled) {
             return;
         }
+        this.experiment.disposeOnInstallingTensorboard(this);
         this.activateInternal().ignoreErrors();
     }
 

src/client/tensorBoard/tensorBoardImportCodeLensProvider.ts

@@ -3,16 +3,16 @@
 
 import { inject, injectable } from 'inversify';
 import { once } from 'lodash';
-import { CancellationToken, CodeLens, Command, languages, Position, Range, TextDocument } from 'vscode';
+import { CancellationToken, CodeLens, Command, Disposable, languages, Position, Range, TextDocument } from 'vscode';
 import { IExtensionSingleActivationService } from '../activation/types';
 import { Commands, PYTHON } from '../common/constants';
-import { IDisposableRegistry } from '../common/types';
+import { IDisposable, IDisposableRegistry } from '../common/types';
 import { TensorBoard } from '../common/utils/localize';
 import { sendTelemetryEvent } from '../telemetry';
 import { EventName } from '../telemetry/constants';
 import { TensorBoardEntrypoint, TensorBoardEntrypointTrigger } from './constants';
 import { containsTensorBoardImport } from './helpers';
-import { useNewTensorboardExtension } from './tensorboarExperiment';
+import { TensorboardExperiment } from './tensorboarExperiment';
 
 @injectable()
 export class TensorBoardImportCodeLensProvider implements IExtensionSingleActivationService {
@@ -25,12 +25,24 @@ export class TensorBoardImportCodeLensProvider implements IExtensionSingleActiva
         }),
     );
 
-    constructor(@inject(IDisposableRegistry) private disposables: IDisposableRegistry) {}
+    private readonly disposables: IDisposable[] = [];
+
+    constructor(
+        @inject(IDisposableRegistry) disposables: IDisposableRegistry,
+        @inject(TensorboardExperiment) private readonly experiment: TensorboardExperiment,
+    ) {
+        disposables.push(this);
+    }
+
+    public dispose(): void {
+        Disposable.from(...this.disposables).dispose();
+    }
 
     public async activate(): Promise<void> {
-        if (useNewTensorboardExtension()) {
+        if (TensorboardExperiment.isTensorboardExtensionInstalled) {
             return;
         }
+        this.experiment.disposeOnInstallingTensorboard(this);
         this.activateInternal().ignoreErrors();
     }
 

src/client/tensorBoard/tensorBoardSessionProvider.ts

@@ -2,7 +2,7 @@
 // Licensed under the MIT License.
 
 import { inject, injectable } from 'inversify';
-import { l10n, ViewColumn } from 'vscode';
+import { Disposable, l10n, ViewColumn } from 'vscode';
 import { IExtensionSingleActivationService } from '../activation/types';
 import { IApplicationShell, ICommandManager, IWorkspaceService } from '../common/application/types';
 import { Commands } from '../common/constants';
@@ -14,6 +14,7 @@ import {
     IPersistentState,
     IPersistentStateFactory,
     IConfigurationService,
+    IDisposable,
 } from '../common/types';
 import { IMultiStepInputFactory } from '../common/utils/multiStepInput';
 import { IInterpreterService } from '../interpreter/contracts';
@@ -22,7 +23,7 @@ import { sendTelemetryEvent } from '../telemetry';
 import { EventName } from '../telemetry/constants';
 import { TensorBoardEntrypoint, TensorBoardEntrypointTrigger } from './constants';
 import { TensorBoardSession } from './tensorBoardSession';
-import { useNewTensorboardExtension } from './tensorboarExperiment';
+import { TensorboardExperiment } from './tensorboarExperiment';
 
 export const PREFERRED_VIEWGROUP = 'PythonTensorBoardWebviewPreferredViewGroup';
 
@@ -36,18 +37,22 @@ export class TensorBoardSessionProvider implements IExtensionSingleActivationSer
 
     private hasActiveTensorBoardSessionContext: ContextKey;
 
+    private readonly disposables: IDisposable[] = [];
+
     constructor(
         @inject(IInstaller) private readonly installer: IInstaller,
         @inject(IInterpreterService) private readonly interpreterService: IInterpreterService,
         @inject(IApplicationShell) private readonly applicationShell: IApplicationShell,
         @inject(IWorkspaceService) private readonly workspaceService: IWorkspaceService,
         @inject(ICommandManager) private readonly commandManager: ICommandManager,
-        @inject(IDisposableRegistry) private readonly disposables: IDisposableRegistry,
+        @inject(IDisposableRegistry) disposables: IDisposableRegistry,
         @inject(IPythonExecutionFactory) private readonly pythonExecFactory: IPythonExecutionFactory,
         @inject(IPersistentStateFactory) private stateFactory: IPersistentStateFactory,
         @inject(IMultiStepInputFactory) private readonly multiStepFactory: IMultiStepInputFactory,
         @inject(IConfigurationService) private readonly configurationService: IConfigurationService,
+        @inject(TensorboardExperiment) private readonly experiment: TensorboardExperiment,
     ) {
+        disposables.push(this);
         this.preferredViewGroupMemento = this.stateFactory.createGlobalPersistentState<ViewColumn>(
             PREFERRED_VIEWGROUP,
             ViewColumn.Active,
@@ -58,27 +63,36 @@ export class TensorBoardSessionProvider implements IExtensionSingleActivationSer
         );
     }
 
+    public dispose(): void {
+        Disposable.from(...this.disposables).dispose();
+    }
+
     public async activate(): Promise<void> {
-        if (useNewTensorboardExtension()) {
+        if (TensorboardExperiment.isTensorboardExtensionInstalled) {
             return;
         }
+        this.experiment.disposeOnInstallingTensorboard(this);
 
         this.disposables.push(
             this.commandManager.registerCommand(
                 Commands.LaunchTensorBoard,
                 (
                     entrypoint: TensorBoardEntrypoint = TensorBoardEntrypoint.palette,
                     trigger: TensorBoardEntrypointTrigger = TensorBoardEntrypointTrigger.palette,
-                ) => {
+                ): void => {
                     sendTelemetryEvent(EventName.TENSORBOARD_SESSION_LAUNCH, undefined, {
                         trigger,
                         entrypoint,
                     });
-                    return this.createNewSession();
+                    if (this.experiment.recommendAndUseNewExtension() === 'continueWithPythonExtension') {
+                        void this.createNewSession();
+                    }
                 },
             ),
             this.commandManager.registerCommand(Commands.RefreshTensorBoard, () =>
-                this.knownSessions.map((w) => w.refresh()),
+                this.experiment.recommendAndUseNewExtension() === 'continueWithPythonExtension'
+                    ? this.knownSessions.map((w) => w.refresh())
+                    : undefined,
             ),
         );
     }

src/client/tensorBoard/tensorBoardUsageTracker.ts

@@ -3,7 +3,7 @@
 
 import { inject, injectable } from 'inversify';
 import * as path from 'path';
-import { TextEditor } from 'vscode';
+import { Disposable, TextEditor } from 'vscode';
 import { IExtensionSingleActivationService } from '../activation/types';
 import { IDocumentManager } from '../common/application/types';
 import { isTestExecution } from '../common/constants';
@@ -12,7 +12,7 @@ import { getDocumentLines } from '../telemetry/importTracker';
 import { TensorBoardEntrypointTrigger } from './constants';
 import { containsTensorBoardImport } from './helpers';
 import { TensorBoardPrompt } from './tensorBoardPrompt';
-import { useNewTensorboardExtension } from './tensorboarExperiment';
+import { TensorboardExperiment } from './tensorboarExperiment';
 
 const testExecution = isTestExecution();
 
@@ -26,12 +26,20 @@ export class TensorBoardUsageTracker implements IExtensionSingleActivationServic
         @inject(IDocumentManager) private documentManager: IDocumentManager,
         @inject(IDisposableRegistry) private disposables: IDisposableRegistry,
         @inject(TensorBoardPrompt) private prompt: TensorBoardPrompt,
-    ) {}
+        @inject(TensorboardExperiment) private readonly experiment: TensorboardExperiment,
+    ) {
+        disposables.push(this);
+    }
+
+    public dispose(): void {
+        Disposable.from(...this.disposables).dispose();
+    }
 
     public async activate(): Promise<void> {
-        if (useNewTensorboardExtension()) {
+        if (TensorboardExperiment.isTensorboardExtensionInstalled) {
             return;
         }
+        this.experiment.disposeOnInstallingTensorboard(this);
         if (testExecution) {
             await this.activateInternal();
         } else {

src/client/tensorBoard/tensorboarExperiment.ts

@@ -1,8 +1,67 @@
 // Copyright (c) Microsoft Corporation.
 // Licensed under the MIT License.
 
-import { extensions } from 'vscode';
+import { Disposable, EventEmitter, commands, extensions, l10n, window } from 'vscode';
+import { inject, injectable } from 'inversify';
+import { IDisposable, IDisposableRegistry, IExperimentService } from '../common/types';
+import { RecommendTensobardExtension } from '../common/experiments/groups';
+import { TENSORBOARD_EXTENSION_ID } from '../common/constants';
 
-export function useNewTensorboardExtension(): boolean {
-    return !!extensions.getExtension('ms-toolsai.tensorboard');
+@injectable()
+export class TensorboardExperiment {
+    private readonly _onDidChange = new EventEmitter<void>();
+
+    public readonly onDidChange = this._onDidChange.event;
+
+    private readonly toDisposeWhenTensobardIsInstalled: IDisposable[] = [];
+
+    public static get isTensorboardExtensionInstalled(): boolean {
+        return !!extensions.getExtension(TENSORBOARD_EXTENSION_ID);
+    }
+
+    private readonly isExperimentEnabled: boolean;
+
+    constructor(
+        @inject(IDisposableRegistry) disposables: IDisposableRegistry,
+        @inject(IExperimentService) experiments: IExperimentService,
+    ) {
+        this.isExperimentEnabled = experiments.inExperimentSync(RecommendTensobardExtension.experiment);
+        disposables.push(this._onDidChange);
+        extensions.onDidChange(
+            () =>
+                TensorboardExperiment.isTensorboardExtensionInstalled
+                    ? Disposable.from(...this.toDisposeWhenTensobardIsInstalled).dispose()
+                    : undefined,
+            this,
+            disposables,
+        );
+    }
+
+    public recommendAndUseNewExtension(): 'continueWithPythonExtension' | 'usingTensorboardExtension' {
+        if (!this.isExperimentEnabled) {
+            return 'continueWithPythonExtension';
+        }
+        if (TensorboardExperiment.isTensorboardExtensionInstalled) {
+            return 'usingTensorboardExtension';
+        }
+        const install = l10n.t('Install Tensorboard Extension');
+        window
+            .showInformationMessage(
+                l10n.t(
+                    'Install the TensorBoard extension to use the this functionality. Once installed, select the command `Launch Tensorboard`.',
+                ),
+                { modal: true },
+                install,
+            )
+            .then((result): void => {
+                if (result === install) {
+                    void commands.executeCommand('workbench.extensions.installExtension', TENSORBOARD_EXTENSION_ID);
+                }
+            });
+        return 'usingTensorboardExtension';
+    }
+
+    public disposeOnInstallingTensorboard(disposabe: IDisposable): void {
+        this.toDisposeWhenTensobardIsInstalled.push(disposabe);
+    }
 }