Merge branch 'jupyter13018' of https://github.com/debonte/vscode-jupyter

into jupyter13018
microsoft · Apr 5, 2023 · 3593d53 · 3593d53
2 parents 8bafbc4 + a37b6b3
commit 3593d53
Show file tree

Hide file tree

Showing 21 changed files with 237 additions and 122 deletions.
diff --git a/src/kernels/jupyter/launcher/notebookProvider.ts b/src/kernels/jupyter/launcher/notebookProvider.ts
@@ -15,7 +15,7 @@ import {
 import { Cancellation } from '../../../platform/common/cancellation';
 import { DisplayOptions } from '../../displayOptions';
 import { IRawNotebookProvider } from '../../raw/types';
-import { IJupyterNotebookProvider, IJupyterServerUriStorage } from '../types';
+import { IJupyterNotebookProvider } from '../types';
 import { PythonExtensionNotInstalledError } from '../../../platform/errors/pythonExtNotInstalledError';
 
 /**
@@ -30,8 +30,7 @@ export class NotebookProvider implements INotebookProvider {
         private readonly rawNotebookProvider: IRawNotebookProvider | undefined,
         @inject(IJupyterNotebookProvider)
         private readonly jupyterNotebookProvider: IJupyterNotebookProvider,
-        @inject(IPythonExtensionChecker) private readonly extensionChecker: IPythonExtensionChecker,
-        @inject(IJupyterServerUriStorage) private readonly uriStorage: IJupyterServerUriStorage
+        @inject(IPythonExtensionChecker) private readonly extensionChecker: IPythonExtensionChecker
     ) {}
 
     // Attempt to connect to our server provider, and if we do, return the connection info
@@ -48,11 +47,11 @@ export class NotebookProvider implements INotebookProvider {
         options.ui = this.startupUi;
         if (this.rawNotebookProvider?.isSupported && options.localJupyter) {
             throw new Error('Connect method should not be invoked for local Connections when Raw is supported');
-        } else if (this.extensionChecker.isPythonExtensionInstalled || !this.uriStorage.isLocalLaunch) {
+        } else if (this.extensionChecker.isPythonExtensionInstalled || !options.localJupyter) {
             return this.jupyterNotebookProvider.connect(options).finally(() => handler.dispose());
         } else {
             handler.dispose();
-            if (!this.startupUi.disableUI) {
+            if (!this.startupUi.disableUI && options.localJupyter) {
                 await this.extensionChecker.showPythonExtensionInstallRequiredPrompt();
             }
             throw new PythonExtensionNotInstalledError();

diff --git a/src/kernels/jupyter/launcher/notebookProvider.unit.test.ts b/src/kernels/jupyter/launcher/notebookProvider.unit.test.ts
@@ -8,7 +8,7 @@ import { PythonExtensionChecker } from '../../../platform/api/pythonApi';
 import { IJupyterKernelConnectionSession, KernelConnectionMetadata } from '../../types';
 import { NotebookProvider } from './notebookProvider';
 import { DisplayOptions } from '../../displayOptions';
-import { IJupyterNotebookProvider, IJupyterServerUriStorage } from '../types';
+import { IJupyterNotebookProvider } from '../types';
 import { IRawNotebookProvider } from '../../raw/types';
 import { IDisposable } from '../../../platform/common/types';
 import { disposeAllDisposables } from '../../../platform/common/helpers';
@@ -32,17 +32,13 @@ suite('NotebookProvider', () => {
         when(rawNotebookProvider.isSupported).thenReturn(false);
         const extensionChecker = mock(PythonExtensionChecker);
         when(extensionChecker.isPythonExtensionInstalled).thenReturn(true);
-        const uriStorage = mock<IJupyterServerUriStorage>();
-        when(uriStorage.isLocalLaunch).thenReturn(true);
         const onDidChangeEvent = new vscode.EventEmitter<void>();
         disposables.push(onDidChangeEvent);
-        when(uriStorage.onDidChangeConnectionType).thenReturn(onDidChangeEvent.event);
 
         notebookProvider = new NotebookProvider(
             instance(rawNotebookProvider),
             instance(jupyterNotebookProvider),
-            instance(extensionChecker),
-            instance(uriStorage)
+            instance(extensionChecker)
         );
     });
     teardown(() => disposeAllDisposables(disposables));

diff --git a/src/kernels/kernel.ts b/src/kernels/kernel.ts
@@ -12,7 +12,8 @@ import {
     ColorThemeKind,
     Disposable,
     Uri,
-    NotebookDocument
+    NotebookDocument,
+    Memento
 } from 'vscode';
 import {
     CodeSnippets,
@@ -60,6 +61,7 @@ import { DisplayOptions } from './displayOptions';
 import { SilentExecutionErrorOptions } from './helpers';
 import dedent from 'dedent';
 import { IAnyMessageArgs } from '@jupyterlab/services/lib/kernel/kernel';
+import { cacheKernelInfo, getCacheKernelInfo } from './kernelInfoCache';
 
 const widgetVersionOutPrefix = 'e976ee50-99ed-4aba-9b6b-9dcd5634d07d:IPyWidgets:';
 /**
@@ -169,7 +171,8 @@ abstract class BaseKernel implements IBaseKernel {
         protected readonly kernelSettings: IKernelSettings,
         protected readonly appShell: IApplicationShell,
         protected readonly startupCodeProviders: IStartupCodeProvider[],
-        public readonly _creator: KernelActionSource
+        public readonly _creator: KernelActionSource,
+        private readonly workspaceMemento: Memento
     ) {
         this.disposables.push(this._onStatusChanged);
         this.disposables.push(this._onRestarted);
@@ -738,10 +741,25 @@ abstract class BaseKernel implements IBaseKernel {
                 protocol_version: '',
                 status: 'ok'
             };
-            promises.push(session.requestKernelInfo().then((item) => item?.content));
+            const kernelInfoPromise = session.requestKernelInfo().then((item) => item?.content);
+            promises.push(kernelInfoPromise);
+            kernelInfoPromise
+                .then((content) =>
+                    cacheKernelInfo(
+                        this.workspaceMemento,
+                        this.kernelConnectionMetadata,
+                        content as KernelMessage.IInfoReply | undefined
+                    )
+                )
+                .catch(noop);
             // If this doesn't complete in 5 seconds for remote kernels, assume the kernel is busy & provide some default content.
             if (this.kernelConnectionMetadata.kind === 'connectToLiveRemoteKernel') {
-                promises.push(sleep(5_000).then(() => defaultResponse));
+                const cachedInfo = getCacheKernelInfo(this.workspaceMemento, this.kernelConnectionMetadata);
+                if (cachedInfo) {
+                    promises.push(Promise.resolve(cachedInfo));
+                } else {
+                    promises.push(sleep(5_000).then(() => defaultResponse));
+                }
             }
             const content = await Promise.race(promises);
             if (content === defaultResponse) {
@@ -954,7 +972,8 @@ export class ThirdPartyKernel extends BaseKernel implements IThirdPartyKernel {
         notebookProvider: INotebookProvider,
         appShell: IApplicationShell,
         kernelSettings: IKernelSettings,
-        startupCodeProviders: IStartupCodeProvider[]
+        startupCodeProviders: IStartupCodeProvider[],
+        workspaceMemento: Memento
     ) {
         super(
             uri,
@@ -964,7 +983,8 @@ export class ThirdPartyKernel extends BaseKernel implements IThirdPartyKernel {
             kernelSettings,
             appShell,
             startupCodeProviders,
-            '3rdPartyExtension'
+            '3rdPartyExtension',
+            workspaceMemento
         );
     }
 }
@@ -985,7 +1005,8 @@ export class Kernel extends BaseKernel implements IKernel {
         kernelSettings: IKernelSettings,
         appShell: IApplicationShell,
         public readonly controller: IKernelController,
-        startupCodeProviders: IStartupCodeProvider[]
+        startupCodeProviders: IStartupCodeProvider[],
+        workspaceMemento: Memento
     ) {
         super(
             notebook.uri,
@@ -995,7 +1016,8 @@ export class Kernel extends BaseKernel implements IKernel {
             kernelSettings,
             appShell,
             startupCodeProviders,
-            'jupyterExtension'
+            'jupyterExtension',
+            workspaceMemento
         );
     }
 }

diff --git a/src/kernels/kernelInfoCache.ts b/src/kernels/kernelInfoCache.ts
@@ -0,0 +1,42 @@
+// Copyright (c) Microsoft Corporation.
+// Licensed under the MIT License.
+
+import type { KernelMessage } from '@jupyterlab/services';
+import { KernelConnectionMetadata, isRemoteConnection } from './types';
+import { Memento } from 'vscode';
+
+type CachedKernelInfo = {
+    id: string;
+    info: KernelMessage.IInfoReply;
+    // Time when this was stored in mememto.
+    age: number;
+};
+const KEY = 'KERNEL_INFO';
+const CACHE_EXPIRY_IN_MS = 1000 * 60 * 60 * 24 * 2;
+
+export async function cacheKernelInfo(
+    storage: Memento,
+    kernelConnection: KernelConnectionMetadata,
+    info: KernelMessage.IInfoReply | undefined
+) {
+    if (!info || !isRemoteConnection(kernelConnection)) {
+        return;
+    }
+    const kernelInfos = storage
+        .get<CachedKernelInfo[]>(KEY, [])
+        .filter((item) => Date.now() - item.age < CACHE_EXPIRY_IN_MS)
+        .filter((item) => item.id !== kernelConnection.id);
+    kernelInfos.push({
+        id: kernelConnection.id,
+        age: Date.now(),
+        info
+    });
+
+    await storage.update(KEY, kernelInfos);
+}
+export function getCacheKernelInfo(storage: Memento, kernelConnection: KernelConnectionMetadata) {
+    if (!isRemoteConnection(kernelConnection)) {
+        return;
+    }
+    return storage.get<CachedKernelInfo[]>(KEY, []).find((item) => item.id === kernelConnection.id)?.info;
+}
diff --git a/src/kernels/kernelProvider.node.ts b/src/kernels/kernelProvider.node.ts
@@ -1,14 +1,16 @@
 // Copyright (c) Microsoft Corporation.
 // Licensed under the MIT License.
 
-import { inject, injectable, multiInject } from 'inversify';
-import { NotebookDocument, Uri } from 'vscode';
+import { inject, injectable, multiInject, named } from 'inversify';
+import { Memento, NotebookDocument, Uri } from 'vscode';
 import { IApplicationShell, IVSCodeNotebook } from '../platform/common/application/types';
 import {
     IAsyncDisposableRegistry,
     IConfigurationService,
     IDisposableRegistry,
-    IExtensionContext
+    IExtensionContext,
+    IMemento,
+    WORKSPACE_MEMENTO
 } from '../platform/common/types';
 import { BaseCoreKernelProvider, BaseThirdPartyKernelProvider } from './kernelProvider.base';
 import { InteractiveScheme, InteractiveWindowView, JupyterNotebookView } from '../platform/common/constants';
@@ -42,7 +44,8 @@ export class KernelProvider extends BaseCoreKernelProvider {
         @inject(IJupyterServerUriStorage) jupyterServerUriStorage: IJupyterServerUriStorage,
         @multiInject(ITracebackFormatter)
         private readonly formatters: ITracebackFormatter[],
-        @inject(IStartupCodeProviders) private readonly startupCodeProviders: IStartupCodeProviders
+        @inject(IStartupCodeProviders) private readonly startupCodeProviders: IStartupCodeProviders,
+        @inject(IMemento) @named(WORKSPACE_MEMENTO) private readonly workspaceStorage: Memento
     ) {
         super(asyncDisposables, disposables, notebook);
         disposables.push(jupyterServerUriStorage.onDidRemoveUris(this.handleUriRemoval.bind(this)));
@@ -70,7 +73,8 @@ export class KernelProvider extends BaseCoreKernelProvider {
             settings,
             this.appShell,
             options.controller,
-            this.startupCodeProviders.getProviders(notebookType)
+            this.startupCodeProviders.getProviders(notebookType),
+            this.workspaceStorage
         );
         kernel.onRestarted(() => this._onDidRestartKernel.fire(kernel), this, this.disposables);
         kernel.onDisposed(
@@ -107,7 +111,8 @@ export class ThirdPartyKernelProvider extends BaseThirdPartyKernelProvider {
         @inject(IConfigurationService) private configService: IConfigurationService,
         @inject(IApplicationShell) private readonly appShell: IApplicationShell,
         @inject(IVSCodeNotebook) notebook: IVSCodeNotebook,
-        @inject(IStartupCodeProviders) private readonly startupCodeProviders: IStartupCodeProviders
+        @inject(IStartupCodeProviders) private readonly startupCodeProviders: IStartupCodeProviders,
+        @inject(IMemento) @named(WORKSPACE_MEMENTO) private readonly workspaceStorage: Memento
     ) {
         super(asyncDisposables, disposables, notebook);
     }
@@ -133,7 +138,8 @@ export class ThirdPartyKernelProvider extends BaseThirdPartyKernelProvider {
             this.notebookProvider,
             this.appShell,
             settings,
-            this.startupCodeProviders.getProviders(notebookType)
+            this.startupCodeProviders.getProviders(notebookType),
+            this.workspaceStorage
         );
         kernel.onRestarted(() => this._onDidRestartKernel.fire(kernel), this, this.disposables);
         kernel.onDisposed(

diff --git a/src/kernels/kernelProvider.node.unit.test.ts b/src/kernels/kernelProvider.node.unit.test.ts
@@ -5,7 +5,14 @@
 import { assert } from 'chai';
 import * as sinon from 'sinon';
 import { anything, instance, mock, when } from 'ts-mockito';
-import { EventEmitter, NotebookCellExecutionStateChangeEvent, NotebookController, NotebookDocument, Uri } from 'vscode';
+import {
+    EventEmitter,
+    Memento,
+    NotebookCellExecutionStateChangeEvent,
+    NotebookController,
+    NotebookDocument,
+    Uri
+} from 'vscode';
 import { IApplicationShell, IVSCodeNotebook } from '../platform/common/application/types';
 import {
     IConfigurationService,
@@ -45,6 +52,7 @@ suite('Node Kernel Provider', function () {
     let jupyterServerUriStorage: IJupyterServerUriStorage;
     let metadata: KernelConnectionMetadata;
     let controller: IKernelController;
+    let workspaceMemento: Memento;
     setup(() => {
         notebookProvider = mock<INotebookProvider>();
         configService = mock<IConfigurationService>();
@@ -54,6 +62,11 @@ suite('Node Kernel Provider', function () {
         jupyterServerUriStorage = mock<IJupyterServerUriStorage>();
         metadata = mock<KernelConnectionMetadata>();
         controller = createKernelController();
+        workspaceMemento = mock<Memento>();
+        when(workspaceMemento.update(anything(), anything())).thenResolve();
+        when(workspaceMemento.get(anything(), anything())).thenCall(
+            (_: unknown, defaultValue: unknown) => defaultValue
+        );
     });
     function createKernelProvider() {
         const registry = mock<IStartupCodeProviders>();
@@ -69,7 +82,8 @@ suite('Node Kernel Provider', function () {
             instance(context),
             instance(jupyterServerUriStorage),
             [],
-            instance(registry)
+            instance(registry),
+            instance(workspaceMemento)
         );
     }
     function create3rdPartyKernelProvider() {
@@ -83,7 +97,8 @@ suite('Node Kernel Provider', function () {
             instance(configService),
             instance(appShell),
             instance(vscNotebook),
-            instance(registry)
+            instance(registry),
+            instance(workspaceMemento)
         );
     }
     teardown(async () => {
@@ -211,6 +226,12 @@ suite('KernelProvider Node', () => {
         ]);
         const registry = mock<IStartupCodeProviders>();
         when(registry.getProviders(anything())).thenReturn([]);
+        const workspaceMemento = mock<Memento>();
+        when(workspaceMemento.update(anything(), anything())).thenResolve();
+        when(workspaceMemento.get(anything(), anything())).thenCall(
+            (_: unknown, defaultValue: unknown) => defaultValue
+        );
+
         kernelProvider = new KernelProvider(
             asyncDisposables,
             disposables,
@@ -221,7 +242,8 @@ suite('KernelProvider Node', () => {
             instance(context),
             instance(jupyterServerUriStorage),
             [],
-            instance(registry)
+            instance(registry),
+            instance(workspaceMemento)
         );
         thirdPartyKernelProvider = new ThirdPartyKernelProvider(
             asyncDisposables,
@@ -230,7 +252,8 @@ suite('KernelProvider Node', () => {
             instance(configService),
             instance(appShell),
             instance(vscNotebook),
-            instance(registry)
+            instance(registry),
+            instance(workspaceMemento)
         );
     });
     teardown(async () => {