Add run notebook in dedicated extension host command (#9492)

Author: rebornix

package.json

@@ -68,6 +68,7 @@
         "onCommand:jupyter.selectNativeJupyterUriFromToolBar",
         "onCommand:jupyter.clearSavedJupyterUris",
         "onCommand:jupyter.filterKernels",
+        "onCommand:jupyter.runInDedicatedExtensionHost",
         "onWalkthrough:jupyterWelcome",
         "onNotebook:jupyter-notebook",
         "onNotebook:interactive"
@@ -812,6 +813,11 @@
                 "title": "%DataScience.interactiveCopyCell%",
                 "icon": "$(copy)",
                 "category": "Jupyter"
+            },
+            {
+                "command": "jupyter.runInDedicatedExtensionHost",
+                "title": "%DataScience.runInDedicatedExtensionHost%",
+                "category": "Jupyter"
             }
         ],
         "menus": {

package.nls.json

@@ -171,6 +171,7 @@
     "DataScience.expandVariableExplorerTooltip": "Show variables active in jupyter kernel",
     "DataScience.close": "Close",
     "DataScience.interactiveCopyCell": "Copy Cell",
+    "DataScience.runInDedicatedExtensionHost": "Run in Dedicated Extension Host",
     "DataScience.variableLoadingValue": "Loading...",
     "DataScience.importDialogTitle": "Import Jupyter Notebook",
     "DataScience.importDialogFilter": "Jupyter Notebooks",
@@ -572,4 +573,4 @@
     "Installer.couldNotInstallLibrary": "Could not install {0}. If pip is not available, please use the package manager of your choice to manually install this library into your Python environment.",
     "Installer.dataScienceInstallPrompt": "Data Science library {0} is not installed. Install?",
     "Products.installingModule": "Installing {0}"
-}
+}

src/platform/common/application/commands.ts

@@ -86,6 +86,7 @@ export interface ICommandNameArgumentTypeMapping extends ICommandNameWithoutArgu
     ['extension.open']: [string];
     ['setContext']: [string, boolean] | ['jupyter.vscode.channel', Channel];
     ['jupyter.reloadVSCode']: [string];
+    ['jupyter.runInDedicatedExtensionHost']: [string];
     ['revealLine']: [{ lineNumber: number; at: 'top' | 'center' | 'bottom' }];
     ['python._loadLanguageServerExtension']: {}[];
     ['python.SelectAndInsertDebugConfiguration']: [TextDocument, Position, CancellationToken];

src/platform/common/application/commands/runInDedicatedExtensionHost.ts

@@ -0,0 +1,62 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
+'use strict';
+
+import { inject, injectable } from 'inversify';
+import { ConfigurationTarget, extensions } from 'vscode';
+import { IExtensionSingleActivationService } from '../../../activation/types';
+import { PythonExtension, PylanceExtension } from '../../constants';
+import { noop } from '../../utils/misc';
+import { ICommandManager, IWorkspaceService } from '../types';
+
+/**
+ * Prompts user to reload VS Code with a custom message, and reloads if necessary.
+ */
+@injectable()
+export class RunInDedicatedExtensionHostCommandHandler implements IExtensionSingleActivationService {
+    constructor(
+        @inject(ICommandManager) private readonly commandManager: ICommandManager,
+        @inject(IWorkspaceService) private readonly workspaceService: IWorkspaceService
+    ) {}
+    public async activate(): Promise<void> {
+        this.commandManager.registerCommand('jupyter.runInDedicatedExtensionHost', this.updateAffinity, this);
+    }
+    private async updateAffinity() {
+        const affinity = this.workspaceService.getConfiguration('extensions').get('experimental.affinity') as
+            | { [key: string]: number }
+            | undefined;
+        let maxAffinity = 0;
+        if (affinity) {
+            Object.values(affinity).forEach((value) => {
+                maxAffinity = Math.max(maxAffinity, value);
+            });
+        }
+
+        const targetAffinity = maxAffinity + 1;
+
+        let update: { [key: string]: number } = {
+            'ms-toolsai.jupyter': targetAffinity,
+            'ms-toolsai.jupyter-renderers': targetAffinity
+        };
+
+        if (extensions.getExtension(PythonExtension)) {
+            update[PythonExtension] = targetAffinity;
+        }
+
+        if (extensions.getExtension(PylanceExtension)) {
+            update[PylanceExtension] = targetAffinity;
+        }
+
+        await this.workspaceService.getConfiguration('extensions').update(
+            'experimental.affinity',
+            {
+                ...(affinity ?? {}),
+                ...update
+            },
+            ConfigurationTarget.Global
+        );
+
+        this.commandManager.executeCommand('workbench.action.reloadWindow').then(noop, noop);
+    }
+}

src/platform/common/serviceRegistry.ts

@@ -10,6 +10,7 @@ import { ActiveResourceService } from './application/activeResource';
 import { ApplicationEnvironment } from './application/applicationEnvironment';
 import { ClipboardService } from './application/clipboard';
 import { ReloadVSCodeCommandHandler } from './application/commands/reloadCommand';
+import { RunInDedicatedExtensionHostCommandHandler } from './application/commands/runInDedicatedExtensionHost';
 import { DebugService } from './application/debugService';
 import { DocumentManager } from './application/documentManager';
 import { EncryptedStorage } from './application/encryptedStorage';
@@ -91,6 +92,10 @@ export function registerTypes(serviceManager: IServiceManager) {
         IExtensionSingleActivationService,
         ReloadVSCodeCommandHandler
     );
+    serviceManager.addSingleton<IExtensionSingleActivationService>(
+        IExtensionSingleActivationService,
+        RunInDedicatedExtensionHostCommandHandler
+    );
     serviceManager.addSingleton<IExtensionSingleActivationService>(
         IExtensionSingleActivationService,
         AmlComputeContext