Move tensorboard support into a separate extension

package.json

@@ -1836,15 +1836,15 @@
                     "category": "Python",
                     "command": "python.launchTensorBoard",
                     "title": "%python.command.python.launchTensorBoard.title%",
-                    "when": "!virtualWorkspace && shellExecutionSupported"
+                    "when": "!virtualWorkspace && shellExecutionSupported && !python.tensorboardExtInstalled"
                 },
                 {
                     "category": "Python",
                     "command": "python.refreshTensorBoard",
                     "enablement": "python.hasActiveTensorBoardSession",
                     "icon": "$(refresh)",
                     "title": "%python.command.python.refreshTensorBoard.title%",
-                    "when": "!virtualWorkspace && shellExecutionSupported"
+                    "when": "!virtualWorkspace && shellExecutionSupported && !python.tensorboardExtInstalled"
                 },
                 {
                     "category": "Python",

src/client/api.ts

@@ -21,6 +21,7 @@ import { IDiscoveryAPI } from './pythonEnvironments/base/locator';
 import { buildEnvironmentApi } from './environmentApi';
 import { ApiForPylance } from './pylanceApi';
 import { getTelemetryReporter } from './telemetry';
+import { TensorboardExtensionIntegration } from './tensorBoard/tensorboardIntegration';
 
 export function buildApi(
     ready: Promise<void>,
@@ -31,7 +32,14 @@ export function buildApi(
     const configurationService = serviceContainer.get<IConfigurationService>(IConfigurationService);
     const interpreterService = serviceContainer.get<IInterpreterService>(IInterpreterService);
     serviceManager.addSingleton<JupyterExtensionIntegration>(JupyterExtensionIntegration, JupyterExtensionIntegration);
+    serviceManager.addSingleton<TensorboardExtensionIntegration>(
+        TensorboardExtensionIntegration,
+        TensorboardExtensionIntegration,
+    );
     const jupyterIntegration = serviceContainer.get<JupyterExtensionIntegration>(JupyterExtensionIntegration);
+    const tensorboardIntegration = serviceContainer.get<TensorboardExtensionIntegration>(
+        TensorboardExtensionIntegration,
+    );
     const outputChannel = serviceContainer.get<ILanguageServerOutputChannel>(ILanguageServerOutputChannel);
 
     const api: PythonExtension & {
@@ -41,6 +49,12 @@ export function buildApi(
         jupyter: {
             registerHooks(): void;
         };
+        /**
+         * Internal API just for Tensorboard, hence don't include in the official types.
+         */
+        tensorboard: {
+            registerHooks(): void;
+        };
     } & {
         /**
          * @deprecated Temporarily exposed for Pylance until we expose this API generally. Will be removed in an
@@ -92,6 +106,9 @@ export function buildApi(
         jupyter: {
             registerHooks: () => jupyterIntegration.integrateWithJupyterExtension(),
         },
+        tensorboard: {
+            registerHooks: () => tensorboardIntegration.integrateWithTensorboardExtension(),
+        },
         debug: {
             async getRemoteLauncherCommand(
                 host: string,

src/client/common/constants.ts

@@ -23,6 +23,7 @@ export const PYTHON_NOTEBOOKS = [
 export const PVSC_EXTENSION_ID = 'ms-python.python';
 export const PYLANCE_EXTENSION_ID = 'ms-python.vscode-pylance';
 export const JUPYTER_EXTENSION_ID = 'ms-toolsai.jupyter';
+export const TENSORBOARD_EXTENSION_ID = 'ms-toolsai.tensorboard';
 export const AppinsightsKey = '0c6ae279ed8443289764825290e4f9e2-1a736e7c-1324-4338-be46-fc2a58ae4d14-7255';
 
 export type Channel = 'stable' | 'insiders';

src/client/tensorBoard/nbextensionCodeLensProvider.ts

@@ -12,6 +12,7 @@ import { sendTelemetryEvent } from '../telemetry';
 import { EventName } from '../telemetry/constants';
 import { TensorBoardEntrypoint, TensorBoardEntrypointTrigger } from './constants';
 import { containsNotebookExtension } from './helpers';
+import { useNewTensorboardExtension } from './tensorboarExperiment';
 
 @injectable()
 export class TensorBoardNbextensionCodeLensProvider implements IExtensionSingleActivationService {
@@ -27,6 +28,9 @@ export class TensorBoardNbextensionCodeLensProvider implements IExtensionSingleA
     constructor(@inject(IDisposableRegistry) private disposables: IDisposableRegistry) {}
 
     public async activate(): Promise<void> {
+        if (useNewTensorboardExtension()) {
+            return;
+        }
         this.activateInternal().ignoreErrors();
     }
 

src/client/tensorBoard/serviceRegistry.ts

@@ -10,6 +10,7 @@ import { TensorBoardPrompt } from './tensorBoardPrompt';
 import { TensorBoardSessionProvider } from './tensorBoardSessionProvider';
 import { TensorBoardNbextensionCodeLensProvider } from './nbextensionCodeLensProvider';
 import { TerminalWatcher } from './terminalWatcher';
+import { TensorboardDependencyChecker } from './tensorboardDependencyChecker';
 
 export function registerTypes(serviceManager: IServiceManager): void {
     serviceManager.addSingleton<TensorBoardSessionProvider>(TensorBoardSessionProvider, TensorBoardSessionProvider);
@@ -32,4 +33,5 @@ export function registerTypes(serviceManager: IServiceManager): void {
     );
     serviceManager.addBinding(TensorBoardNbextensionCodeLensProvider, IExtensionSingleActivationService);
     serviceManager.addSingleton(IExtensionSingleActivationService, TerminalWatcher);
+    serviceManager.addSingleton(TensorboardDependencyChecker, TensorboardDependencyChecker);
 }

src/client/tensorBoard/tensorBoardFileWatcher.ts

@@ -8,6 +8,7 @@ import { IWorkspaceService } from '../common/application/types';
 import { IDisposableRegistry } from '../common/types';
 import { TensorBoardEntrypointTrigger } from './constants';
 import { TensorBoardPrompt } from './tensorBoardPrompt';
+import { useNewTensorboardExtension } from './tensorboarExperiment';
 
 @injectable()
 export class TensorBoardFileWatcher implements IExtensionSingleActivationService {
@@ -24,6 +25,9 @@ export class TensorBoardFileWatcher implements IExtensionSingleActivationService
     ) {}
 
     public async activate(): Promise<void> {
+        if (useNewTensorboardExtension()) {
+            return;
+        }
         this.activateInternal().ignoreErrors();
     }
 

src/client/tensorBoard/tensorBoardImportCodeLensProvider.ts

@@ -12,6 +12,7 @@ import { sendTelemetryEvent } from '../telemetry';
 import { EventName } from '../telemetry/constants';
 import { TensorBoardEntrypoint, TensorBoardEntrypointTrigger } from './constants';
 import { containsTensorBoardImport } from './helpers';
+import { useNewTensorboardExtension } from './tensorboarExperiment';
 
 @injectable()
 export class TensorBoardImportCodeLensProvider implements IExtensionSingleActivationService {
@@ -27,6 +28,9 @@ export class TensorBoardImportCodeLensProvider implements IExtensionSingleActiva
     constructor(@inject(IDisposableRegistry) private disposables: IDisposableRegistry) {}
 
     public async activate(): Promise<void> {
+        if (useNewTensorboardExtension()) {
+            return;
+        }
         this.activateInternal().ignoreErrors();
     }
 

src/client/tensorBoard/tensorBoardPrompt.ts

@@ -84,7 +84,7 @@ export class TensorBoardPrompt {
         }
     }
 
-    private isPromptEnabled(): boolean {
+    public isPromptEnabled(): boolean {
         return this.state.value;
     }
 

src/client/tensorBoard/tensorBoardSession.ts

@@ -100,7 +100,10 @@ export class TensorBoardSession {
         private readonly globalMemento: IPersistentState<ViewColumn>,
         private readonly multiStepFactory: IMultiStepInputFactory,
         private readonly configurationService: IConfigurationService,
-    ) {}
+    ) {
+        this.disposables.push(this.onDidChangeViewStateEventEmitter);
+        this.disposables.push(this.onDidDisposeEventEmitter);
+    }
 
     public get onDidDispose(): Event<TensorBoardSession> {
         return this.onDidDisposeEventEmitter.event;
@@ -189,10 +192,10 @@ export class TensorBoardSession {
     // to start a TensorBoard session. If the user has a torch import in
     // any of their open documents, also try to install the torch-tb-plugin
     // package, but don't block if installing that fails.
-    private async ensurePrerequisitesAreInstalled() {
+    public async ensurePrerequisitesAreInstalled(resource?: Uri): Promise<boolean> {
         traceVerbose('Ensuring TensorBoard package is installed into active interpreter');
         const interpreter =
-            (await this.interpreterService.getActiveInterpreter()) ||
+            (await this.interpreterService.getActiveInterpreter(resource)) ||
             (await this.commandManager.executeCommand('python.setInterpreter'));
         if (!interpreter) {
             return false;

src/client/tensorBoard/tensorBoardSessionProvider.ts

@@ -22,8 +22,9 @@ import { sendTelemetryEvent } from '../telemetry';
 import { EventName } from '../telemetry/constants';
 import { TensorBoardEntrypoint, TensorBoardEntrypointTrigger } from './constants';
 import { TensorBoardSession } from './tensorBoardSession';
+import { useNewTensorboardExtension } from './tensorboarExperiment';
 
-const PREFERRED_VIEWGROUP = 'PythonTensorBoardWebviewPreferredViewGroup';
+export const PREFERRED_VIEWGROUP = 'PythonTensorBoardWebviewPreferredViewGroup';
 
 @injectable()
 export class TensorBoardSessionProvider implements IExtensionSingleActivationService {
@@ -58,6 +59,10 @@ export class TensorBoardSessionProvider implements IExtensionSingleActivationSer
     }
 
     public async activate(): Promise<void> {
+        if (useNewTensorboardExtension()) {
+            return;
+        }
+
         this.disposables.push(
             this.commandManager.registerCommand(
                 Commands.LaunchTensorBoard,

src/client/tensorBoard/tensorBoardUsageTracker.ts

@@ -12,6 +12,7 @@ import { getDocumentLines } from '../telemetry/importTracker';
 import { TensorBoardEntrypointTrigger } from './constants';
 import { containsTensorBoardImport } from './helpers';
 import { TensorBoardPrompt } from './tensorBoardPrompt';
+import { useNewTensorboardExtension } from './tensorboarExperiment';
 
 const testExecution = isTestExecution();
 
@@ -28,6 +29,9 @@ export class TensorBoardUsageTracker implements IExtensionSingleActivationServic
     ) {}
 
     public async activate(): Promise<void> {
+        if (useNewTensorboardExtension()) {
+            return;
+        }
         if (testExecution) {
             await this.activateInternal();
         } else {

src/client/tensorBoard/tensorboarExperiment.ts

@@ -0,0 +1,8 @@
+// Copyright (c) Microsoft Corporation.
+// Licensed under the MIT License.
+
+import { extensions } from 'vscode';
+
+export function useNewTensorboardExtension(): boolean {
+    return !!extensions.getExtension('ms-toolsai.tensorboard');
+}

src/client/tensorBoard/tensorboardDependencyChecker.ts

@@ -0,0 +1,60 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
+import { inject, injectable } from 'inversify';
+import { Uri, ViewColumn } from 'vscode';
+import { IApplicationShell, ICommandManager, IWorkspaceService } from '../common/application/types';
+import { IPythonExecutionFactory } from '../common/process/types';
+import {
+    IInstaller,
+    IPersistentState,
+    IPersistentStateFactory,
+    IConfigurationService,
+    IDisposable,
+} from '../common/types';
+import { IMultiStepInputFactory } from '../common/utils/multiStepInput';
+import { IInterpreterService } from '../interpreter/contracts';
+import { TensorBoardSession } from './tensorBoardSession';
+import { disposeAll } from '../common/utils/resourceLifecycle';
+import { PREFERRED_VIEWGROUP } from './tensorBoardSessionProvider';
+
+@injectable()
+export class TensorboardDependencyChecker {
+    private preferredViewGroupMemento: IPersistentState<ViewColumn>;
+
+    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(IPythonExecutionFactory) private readonly pythonExecFactory: IPythonExecutionFactory,
+        @inject(IPersistentStateFactory) private stateFactory: IPersistentStateFactory,
+        @inject(IMultiStepInputFactory) private readonly multiStepFactory: IMultiStepInputFactory,
+        @inject(IConfigurationService) private readonly configurationService: IConfigurationService,
+    ) {
+        this.preferredViewGroupMemento = this.stateFactory.createGlobalPersistentState<ViewColumn>(
+            PREFERRED_VIEWGROUP,
+            ViewColumn.Active,
+        );
+    }
+
+    public async ensureDependenciesAreInstalled(resource?: Uri): Promise<boolean> {
+        const disposables: IDisposable[] = [];
+        const newSession = new TensorBoardSession(
+            this.installer,
+            this.interpreterService,
+            this.workspaceService,
+            this.pythonExecFactory,
+            this.commandManager,
+            disposables,
+            this.applicationShell,
+            this.preferredViewGroupMemento,
+            this.multiStepFactory,
+            this.configurationService,
+        );
+        const result = await newSession.ensurePrerequisitesAreInstalled(resource);
+        disposeAll(disposables);
+        return result;
+    }
+}

src/client/tensorBoard/tensorboardIntegration.ts

@@ -0,0 +1,102 @@
+/* eslint-disable comma-dangle */
+
+/* eslint-disable implicit-arrow-linebreak */
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
+import { inject, injectable } from 'inversify';
+import { Extension, Uri, commands } from 'vscode';
+import { IWorkspaceService } from '../common/application/types';
+import { TENSORBOARD_EXTENSION_ID } from '../common/constants';
+import { IDisposableRegistry, IExtensions, Resource } from '../common/types';
+import { IEnvironmentActivationService } from '../interpreter/activation/types';
+import { TensorBoardPrompt } from './tensorBoardPrompt';
+import { TensorboardDependencyChecker } from './tensorboardDependencyChecker';
+
+type PythonApiForTensorboardExtension = {
+    /**
+     * Gets activated env vars for the active Python Environment for the given resource.
+     */
+    getActivatedEnvironmentVariables(resource: Resource): Promise<NodeJS.ProcessEnv | undefined>;
+    /**
+     * Ensures that the dependencies required for TensorBoard are installed in Active Environment for the given resource.
+     */
+    ensureDependenciesAreInstalled(resource?: Uri): Promise<boolean>;
+    /**
+     * Whether to allow displaying tensorboard prompt.
+     */
+    isPromptEnabled(): boolean;
+};
+
+type TensorboardExtensionApi = {
+    /**
+     * Registers python extension specific parts with the tensorboard extension
+     */
+    registerPythonApi(interpreterService: PythonApiForTensorboardExtension): void;
+};
+
+@injectable()
+export class TensorboardExtensionIntegration {
+    private tensorboardExtension: Extension<TensorboardExtensionApi> | undefined;
+
+    constructor(
+        @inject(IExtensions) private readonly extensions: IExtensions,
+        @inject(IEnvironmentActivationService) private readonly envActivation: IEnvironmentActivationService,
+        @inject(IWorkspaceService) private workspaceService: IWorkspaceService,
+        @inject(TensorboardDependencyChecker) private readonly dependencyChcker: TensorboardDependencyChecker,
+        @inject(TensorBoardPrompt) private readonly tensorBoardPrompt: TensorBoardPrompt,
+        @inject(IDisposableRegistry) disposables: IDisposableRegistry,
+    ) {
+        this.hideCommands();
+        extensions.onDidChange(this.hideCommands, this, disposables);
+    }
+
+    public registerApi(tensorboardExtensionApi: TensorboardExtensionApi): TensorboardExtensionApi | undefined {
+        this.hideCommands();
+        if (!this.workspaceService.isTrusted) {
+            this.workspaceService.onDidGrantWorkspaceTrust(() => this.registerApi(tensorboardExtensionApi));
+            return undefined;
+        }
+        tensorboardExtensionApi.registerPythonApi({
+            getActivatedEnvironmentVariables: async (resource: Resource) =>
+                this.envActivation.getActivatedEnvironmentVariables(resource, undefined, true),
+            ensureDependenciesAreInstalled: async (resource?: Uri): Promise<boolean> =>
+                this.dependencyChcker.ensureDependenciesAreInstalled(resource),
+            isPromptEnabled: () => this.tensorBoardPrompt.isPromptEnabled(),
+        });
+        return undefined;
+    }
+
+    public hideCommands(): void {
+        if (this.extensions.getExtension(TENSORBOARD_EXTENSION_ID)) {
+            console.error('TensorBoard extension is installed');
+            void commands.executeCommand('setContext', 'python.tensorboardExtInstalled', true);
+        } else {
+            console.error('TensorBoard extension not installed');
+        }
+    }
+
+    public async integrateWithTensorboardExtension(): Promise<void> {
+        const api = await this.getExtensionApi();
+        if (api) {
+            this.registerApi(api);
+        }
+    }
+
+    private async getExtensionApi(): Promise<TensorboardExtensionApi | undefined> {
+        if (!this.tensorboardExtension) {
+            const extension = this.extensions.getExtension<TensorboardExtensionApi>(TENSORBOARD_EXTENSION_ID);
+            if (!extension) {
+                return undefined;
+            }
+            await extension.activate();
+            if (extension.isActive) {
+                this.tensorboardExtension = extension;
+                return this.tensorboardExtension.exports;
+            }
+        } else {
+            return this.tensorboardExtension.exports;
+        }
+        return undefined;
+    }
+}