Second refactor for non-jupyter kernel search (#11527)

Author: IanMatthewHuff

src/client/datascience/jupyter/jupyterExecution.ts

@@ -145,6 +145,7 @@ export class JupyterExecutionBase implements IJupyterExecution {
                 traceInfo(`Getting kernel specs for ${options ? options.purpose : 'unknown type of'} server`);
                 kernelSpecInterpreterPromise = this.kernelSelector.getKernelForLocalConnection(
                     undefined,
+                    'jupyter',
                     undefined,
                     options?.metadata,
                     !allowUI,

src/client/datascience/jupyter/kernels/kernelSelector.ts

@@ -16,6 +16,7 @@ import { StopWatch } from '../../../common/utils/stopWatch';
 import { IInterpreterService, PythonInterpreter } from '../../../interpreter/contracts';
 import { IEventNamePropertyMapping, sendTelemetryEvent } from '../../../telemetry';
 import { KnownNotebookLanguages, Telemetry } from '../../constants';
+import { IKernelFinder } from '../../kernel-launcher/types';
 import { reportAction } from '../../progress/decorator';
 import { ReportableAction } from '../../progress/types';
 import { IJupyterKernelSpec, IJupyterSessionManager, IKernelDependencyService } from '../../types';
@@ -57,7 +58,8 @@ export class KernelSelector {
         @inject(IApplicationShell) private readonly applicationShell: IApplicationShell,
         @inject(KernelService) private readonly kernelService: KernelService,
         @inject(IInterpreterService) private readonly interpreterService: IInterpreterService,
-        @inject(IKernelDependencyService) private readonly kernelDepdencyService: IKernelDependencyService
+        @inject(IKernelDependencyService) private readonly kernelDepdencyService: IKernelDependencyService,
+        @inject(IKernelFinder) private readonly kernelFinder: IKernelFinder
     ) {}
 
     /**
@@ -160,6 +162,7 @@ export class KernelSelector {
     @reportAction(ReportableAction.KernelsGetKernelForLocalConnection)
     public async getKernelForLocalConnection(
         resource: Resource,
+        type: 'raw' | 'jupyter' | 'noConnection',
         sessionManager?: IJupyterSessionManager,
         notebookMetadata?: nbformat.INotebookMetadata,
         disableUI?: boolean,
@@ -171,68 +174,28 @@ export class KernelSelector {
             interpreterFound: false,
             promptedToSelect: false
         };
-        // When this method is called, we know we've started a local jupyter server.
+        // When this method is called, we know we've started a local jupyter server or are connecting raw
         // Lets pre-warm the list of local kernels.
         this.selectionProvider
-            .getKernelSelectionsForLocalSession(resource, 'jupyter', sessionManager, cancelToken)
+            .getKernelSelectionsForLocalSession(resource, type, sessionManager, cancelToken)
             .ignoreErrors();
 
         let selection: KernelSpecInterpreter = {};
-        if (notebookMetadata?.kernelspec) {
-            selection.kernelSpec = await this.kernelService.findMatchingKernelSpec(
-                notebookMetadata?.kernelspec,
+
+        if (type === 'jupyter') {
+            selection = await this.getKernelForLocalJupyterConnection(
+                resource,
+                stopWatch,
+                telemetryProps,
                 sessionManager,
+                notebookMetadata,
+                disableUI,
                 cancelToken
             );
-            if (selection.kernelSpec) {
-                selection.interpreter = await this.kernelService.findMatchingInterpreter(
-                    selection.kernelSpec,
-                    cancelToken
-                );
-                sendTelemetryEvent(Telemetry.UseExistingKernel);
-
-                // Make sure we update the environment in the kernel before using it
-                await this.kernelService.updateKernelEnvironment(
-                    selection.interpreter,
-                    selection.kernelSpec,
-                    cancelToken
-                );
-            } else {
-                // No kernel info, hence prmopt to use current interpreter as a kernel.
-                const activeInterpreter = await this.interpreterService.getActiveInterpreter(resource);
-                if (activeInterpreter) {
-                    selection = await this.useInterpreterAsKernel(
-                        resource,
-                        activeInterpreter,
-                        'jupyter',
-                        notebookMetadata.kernelspec.display_name,
-                        sessionManager,
-                        disableUI,
-                        cancelToken
-                    );
-                } else {
-                    telemetryProps.promptedToSelect = true;
-                    selection = await this.selectLocalKernel(
-                        resource,
-                        'jupyter',
-                        stopWatch,
-                        sessionManager,
-                        cancelToken
-                    );
-                }
-            }
-        } else {
-            // No kernel info, hence use current interpreter as a kernel.
-            const activeInterpreter = await this.interpreterService.getActiveInterpreter(resource);
-            if (activeInterpreter) {
-                selection.interpreter = activeInterpreter;
-                selection.kernelSpec = await this.kernelService.searchAndRegisterKernel(
-                    activeInterpreter,
-                    disableUI,
-                    cancelToken
-                );
-            }
+        } else if (type === 'raw') {
+            selection = await this.getKernelForLocalRawConnection(resource, notebookMetadata, cancelToken);
         }
+
         // If still not found, log an error (this seems possible for some people, so use the default)
         if (!selection.kernelSpec) {
             traceError('Jupyter Kernel Spec not found for a local connection');
@@ -313,6 +276,99 @@ export class KernelSelector {
             interpreter: interpreter
         };
     }
+
+    // Get our kernelspec and matching interpreter for a connection to a local jupyter server
+    private async getKernelForLocalJupyterConnection(
+        resource: Resource,
+        stopWatch: StopWatch,
+        telemetryProps: IEventNamePropertyMapping[Telemetry.FindKernelForLocalConnection],
+        sessionManager?: IJupyterSessionManager,
+        notebookMetadata?: nbformat.INotebookMetadata,
+        disableUI?: boolean,
+        cancelToken?: CancellationToken
+    ): Promise<KernelSpecInterpreter> {
+        let selection: KernelSpecInterpreter = {};
+        if (notebookMetadata?.kernelspec) {
+            selection.kernelSpec = await this.kernelService.findMatchingKernelSpec(
+                notebookMetadata?.kernelspec,
+                sessionManager,
+                cancelToken
+            );
+            if (selection.kernelSpec) {
+                selection.interpreter = await this.kernelService.findMatchingInterpreter(
+                    selection.kernelSpec,
+                    cancelToken
+                );
+                sendTelemetryEvent(Telemetry.UseExistingKernel);
+
+                // Make sure we update the environment in the kernel before using it
+                await this.kernelService.updateKernelEnvironment(
+                    selection.interpreter,
+                    selection.kernelSpec,
+                    cancelToken
+                );
+            } else {
+                // No kernel info, hence prmopt to use current interpreter as a kernel.
+                const activeInterpreter = await this.interpreterService.getActiveInterpreter(resource);
+                if (activeInterpreter) {
+                    selection = await this.useInterpreterAsKernel(
+                        resource,
+                        activeInterpreter,
+                        'jupyter',
+                        notebookMetadata.kernelspec.display_name,
+                        sessionManager,
+                        disableUI,
+                        cancelToken
+                    );
+                } else {
+                    telemetryProps.promptedToSelect = true;
+                    selection = await this.selectLocalKernel(
+                        resource,
+                        'jupyter',
+                        stopWatch,
+                        sessionManager,
+                        cancelToken
+                    );
+                }
+            }
+        } else {
+            // No kernel info, hence use current interpreter as a kernel.
+            const activeInterpreter = await this.interpreterService.getActiveInterpreter(resource);
+            if (activeInterpreter) {
+                selection.interpreter = activeInterpreter;
+                selection.kernelSpec = await this.kernelService.searchAndRegisterKernel(
+                    activeInterpreter,
+                    disableUI,
+                    cancelToken
+                );
+            }
+        }
+
+        return selection;
+    }
+
+    // Get our kernelspec and interpreter for a local raw connection
+    private async getKernelForLocalRawConnection(
+        resource: Resource,
+        notebookMetadata?: nbformat.INotebookMetadata,
+        cancelToken?: CancellationToken
+    ): Promise<KernelSpecInterpreter> {
+        const selection: KernelSpecInterpreter = {};
+
+        // First use our kernel finder to locate a kernelspec on disk
+        selection.kernelSpec = await this.kernelFinder.findKernelSpec(
+            resource,
+            notebookMetadata?.kernelspec?.name,
+            cancelToken
+        );
+
+        if (selection.kernelSpec) {
+            // Locate the interpreter that matches our kernelspec
+            selection.interpreter = await this.kernelService.findMatchingInterpreter(selection.kernelSpec, cancelToken);
+        }
+        return selection;
+    }
+
     private async selectKernel(
         resource: Resource,
         type: 'raw' | 'jupyter' | 'noConnection',

src/client/datascience/jupyter/liveshare/hostJupyterServer.ts

@@ -289,6 +289,7 @@ export class HostJupyterServer extends LiveShareParticipantHost(JupyterServerBas
             const kernelInfo = await (launchInfo.connectionInfo.localLaunch
                 ? this.kernelSelector.getKernelForLocalConnection(
                       resource,
+                      'jupyter',
                       sessionManager,
                       notebookMetadata,
                       isTestExecution(),

src/client/datascience/raw-kernel/liveshare/hostRawNotebookProvider.ts

@@ -23,11 +23,10 @@ import * as localize from '../../../common/utils/localize';
 import { IServiceContainer } from '../../../ioc/types';
 import { Identifiers, LiveShare, Settings } from '../../constants';
 import { KernelSelector } from '../../jupyter/kernels/kernelSelector';
-import { KernelService } from '../../jupyter/kernels/kernelService';
 import { HostJupyterNotebook } from '../../jupyter/liveshare/hostJupyterNotebook';
 import { LiveShareParticipantHost } from '../../jupyter/liveshare/liveShareParticipantMixin';
 import { IRoleBasedObject } from '../../jupyter/liveshare/roleBasedFactory';
-import { IKernelFinder, IKernelLauncher } from '../../kernel-launcher/types';
+import { IKernelLauncher } from '../../kernel-launcher/types';
 import { ProgressReporter } from '../../progress/progressReporter';
 import {
     IJupyterKernelSpec,
@@ -57,9 +56,7 @@ export class HostRawNotebookProvider
         private fs: IFileSystem,
         private serviceContainer: IServiceContainer,
         private kernelLauncher: IKernelLauncher,
-        private kernelFinder: IKernelFinder,
         private kernelSelector: KernelSelector,
-        private kernelService: KernelService,
         private progressReporter: ProgressReporter,
         private outputChannel: IOutputChannel
     ) {
@@ -112,46 +109,56 @@ export class HostRawNotebookProvider
         try {
             const launchTimeout = this.configService.getSettings().datascience.jupyterLaunchTimeout;
 
-            // Before we try to connect we need to find a kernel and install ipykernel
-            const kernelSpec = await this.kernelFinder.findKernelSpec(
+            // We need to locate kernelspec and possible interpreter for this launch based on resource and notebook metadata
+            const kernelSpecInterpreter = await this.kernelSelector.getKernelForLocalConnection(
                 resource,
-                notebookMetadata?.kernelspec?.name,
+                'raw',
+                undefined,
+                notebookMetadata,
+                disableUI,
                 cancelToken
             );
 
-            // Locate the interpreter that matches our kernelspec
-            const interpreter = await this.kernelService.findMatchingInterpreter(kernelSpec);
-
-            await rawSession.connect(kernelSpec, launchTimeout, interpreter, cancelToken);
-
-            // Get the execution info for our notebook
-            const info = await this.getExecutionInfo(kernelSpec);
-
-            if (rawSession.isConnected) {
-                // Create our notebook
-                const notebook = new HostJupyterNotebook(
-                    this.liveShare,
-                    rawSession,
-                    this.configService,
-                    this.disposableRegistry,
-                    info,
-                    this.serviceContainer.getAll<INotebookExecutionLogger>(INotebookExecutionLogger),
-                    resource,
-                    identity,
-                    this.getDisposedError.bind(this),
-                    this.workspaceService,
-                    this.appShell,
-                    this.fs
+            // Interpreter is optional, but we must have a kernel spec for a raw launch
+            if (!kernelSpecInterpreter.kernelSpec) {
+                notebookPromise.reject('Failed to find a kernelspec to use for ipykernel launch');
+            } else {
+                await rawSession.connect(
+                    kernelSpecInterpreter.kernelSpec,
+                    launchTimeout,
+                    kernelSpecInterpreter.interpreter,
+                    cancelToken
                 );
 
-                // Run initial setup
-                await notebook.initialize(cancelToken);
-
-                traceInfo(`Finished connecting ${this.id}`);
-
-                notebookPromise.resolve(notebook);
-            } else {
-                notebookPromise.reject(this.getDisposedError());
+                // Get the execution info for our notebook
+                const info = await this.getExecutionInfo(kernelSpecInterpreter.kernelSpec);
+
+                if (rawSession.isConnected) {
+                    // Create our notebook
+                    const notebook = new HostJupyterNotebook(
+                        this.liveShare,
+                        rawSession,
+                        this.configService,
+                        this.disposableRegistry,
+                        info,
+                        this.serviceContainer.getAll<INotebookExecutionLogger>(INotebookExecutionLogger),
+                        resource,
+                        identity,
+                        this.getDisposedError.bind(this),
+                        this.workspaceService,
+                        this.appShell,
+                        this.fs
+                    );
+
+                    // Run initial setup
+                    await notebook.initialize(cancelToken);
+
+                    traceInfo(`Finished connecting ${this.id}`);
+
+                    notebookPromise.resolve(notebook);
+                } else {
+                    notebookPromise.reject(this.getDisposedError());
+                }
             }
         } catch (ex) {
             // Make sure we shut down our session in case we started a process

src/client/datascience/raw-kernel/rawNotebookProviderWrapper.ts

@@ -19,10 +19,9 @@ import {
 import { IServiceContainer } from '../../ioc/types';
 import { JUPYTER_OUTPUT_CHANNEL } from '../constants';
 import { KernelSelector } from '../jupyter/kernels/kernelSelector';
-import { KernelService } from '../jupyter/kernels/kernelService';
 import { IRoleBasedObject, RoleBasedFactory } from '../jupyter/liveshare/roleBasedFactory';
 import { ILiveShareHasRole } from '../jupyter/liveshare/types';
-import { IKernelFinder, IKernelLauncher } from '../kernel-launcher/types';
+import { IKernelLauncher } from '../kernel-launcher/types';
 import { ProgressReporter } from '../progress/progressReporter';
 import { INotebook, IRawConnection, IRawNotebookProvider } from '../types';
 import { GuestRawNotebookProvider } from './liveshare/guestRawNotebookProvider';
@@ -42,9 +41,7 @@ type RawNotebookProviderClassType = {
         fs: IFileSystem,
         serviceContainer: IServiceContainer,
         kernelLauncher: IKernelLauncher,
-        kernelFinder: IKernelFinder,
         kernelSelector: KernelSelector,
-        kernelService: KernelService,
         progressReporter: ProgressReporter,
         outputChannel: IOutputChannel
     ): IRawNotebookProviderInterface;
@@ -67,9 +64,7 @@ export class RawNotebookProviderWrapper implements IRawNotebookProvider, ILiveSh
         @inject(IFileSystem) fs: IFileSystem,
         @inject(IServiceContainer) serviceContainer: IServiceContainer,
         @inject(IKernelLauncher) kernelLauncher: IKernelLauncher,
-        @inject(IKernelFinder) kernelFinder: IKernelFinder,
         @inject(KernelSelector) kernelSelector: KernelSelector,
-        @inject(KernelService) kernelService: KernelService,
         @inject(ProgressReporter) progressReporter: ProgressReporter,
         @inject(IOutputChannel) @named(JUPYTER_OUTPUT_CHANNEL) outputChannel: IOutputChannel
     ) {
@@ -88,9 +83,7 @@ export class RawNotebookProviderWrapper implements IRawNotebookProvider, ILiveSh
             fs,
             serviceContainer,
             kernelLauncher,
-            kernelFinder,
             kernelSelector,
-            kernelService,
             progressReporter,
             outputChannel
         );

src/test/datascience/execution.unit.test.ts

@@ -962,7 +962,14 @@ suite('Jupyter Execution', async () => {
             env: undefined
         };
         when(
-            kernelSelector.getKernelForLocalConnection(anything(), anything(), anything(), anything(), anything())
+            kernelSelector.getKernelForLocalConnection(
+                anything(),
+                anything(),
+                anything(),
+                anything(),
+                anything(),
+                anything()
+            )
         ).thenResolve({
             kernelSpec
         });