Change Kernel for raw kernel and change kernel launcher to launch with kernelspec python (#11284)

Author: IanMatthewHuff

src/client/datascience/baseJupyterSession.ts

@@ -136,7 +136,6 @@ export abstract class BaseJupyterSession implements IJupyterSession {
         return this.shutdown();
     }
     // Abstracts for each Session type to implement
-    public abstract async changeKernel(kernel: IJupyterKernelSpec | LiveKernelModel, timeoutMS: number): Promise<void>;
     public abstract async waitForIdle(timeout: number): Promise<void>;
 
     public async shutdown(): Promise<void> {
@@ -179,6 +178,35 @@ export abstract class BaseJupyterSession implements IJupyterSession {
         }
     }
 
+    public async changeKernel(kernel: IJupyterKernelSpec | LiveKernelModel, timeoutMS: number): Promise<void> {
+        let newSession: ISession | undefined;
+
+        // If we are already using this kernel in an active session just return back
+        if (this.kernelSpec?.name === kernel.name && this.session) {
+            return;
+        }
+
+        newSession = await this.createNewKernelSession(kernel, timeoutMS);
+
+        // This is just like doing a restart, kill the old session (and the old restart session), and start new ones
+        if (this.session) {
+            this.shutdownSession(this.session, this.statusHandler).ignoreErrors();
+            this.restartSessionPromise?.then((r) => this.shutdownSession(r, undefined)).ignoreErrors(); // NOSONAR
+        }
+
+        // Update our kernel spec
+        this.kernelSpec = kernel;
+
+        // Save the new session
+        this.session = newSession;
+
+        // Listen for session status changes
+        this.session?.statusChanged.connect(this.statusHandler); // NOSONAR
+
+        // Start the restart session promise too.
+        this.restartSessionPromise = this.createRestartSession(kernel, this.session);
+    }
+
     public async restart(_timeout: number): Promise<void> {
         if (this.session?.isRemoteSession) {
             await this.session.kernel.restart();
@@ -333,6 +361,12 @@ export abstract class BaseJupyterSession implements IJupyterSession {
         cancelToken?: CancellationToken
     ): Promise<ISession>;
 
+    // Sub classes need to implement their own kernel change specific code
+    protected abstract createNewKernelSession(
+        kernel: IJupyterKernelSpec | LiveKernelModel,
+        timeoutMS: number
+    ): Promise<ISession>;
+
     protected async shutdownSession(
         session: ISession | undefined,
         statusHandler: Slot<ISession, Kernel.Status> | undefined

src/client/datascience/jupyter/jupyterSession.ts

@@ -103,14 +103,12 @@ export class JupyterSession extends BaseJupyterSession {
         this.connected = true;
     }
 
-    public async changeKernel(kernel: IJupyterKernelSpec | LiveKernelModel, timeoutMS: number): Promise<void> {
+    public async createNewKernelSession(
+        kernel: IJupyterKernelSpec | LiveKernelModel,
+        timeoutMS: number
+    ): Promise<ISession> {
         let newSession: ISession | undefined;
 
-        // If we are already using this kernel in an active session just return back
-        if (this.kernelSpec?.name === kernel.name && this.session) {
-            return;
-        }
-
         try {
             // Don't immediately assume this kernel is valid. Try creating a session with it first.
             if (kernel.id && this.session && 'session' in kernel) {
@@ -129,23 +127,7 @@ export class JupyterSession extends BaseJupyterSession {
             throw new JupyterInvalidKernelError(kernel);
         }
 
-        // This is just like doing a restart, kill the old session (and the old restart session), and start new ones
-        if (this.session) {
-            this.shutdownSession(this.session, this.statusHandler).ignoreErrors();
-            this.restartSessionPromise?.then((r) => this.shutdownSession(r, undefined)).ignoreErrors();
-        }
-
-        // Update our kernel spec
-        this.kernelSpec = kernel;
-
-        // Save the new session
-        this.session = newSession;
-
-        // Listen for session status changes
-        this.session?.statusChanged.connect(this.statusHandler); // NOSONAR
-
-        // Start the restart session promise too.
-        this.restartSessionPromise = this.createRestartSession(kernel, this.session);
+        return newSession;
     }
 
     protected async createRestartSession(

src/client/datascience/kernel-launcher/kernelLauncher.ts

@@ -14,7 +14,7 @@ import { IFileSystem, TemporaryFile } from '../../common/platform/types';
 import { IPythonExecutionFactory } from '../../common/process/types';
 import { createDeferred, Deferred } from '../../common/utils/async';
 import * as localize from '../../common/utils/localize';
-import { isResource, noop } from '../../common/utils/misc';
+import { noop } from '../../common/utils/misc';
 import { IJupyterKernelSpec } from '../types';
 import { IKernelConnection, IKernelFinder, IKernelLauncher, IKernelProcess } from './types';
 
@@ -50,18 +50,14 @@ class KernelProcess implements IKernelProcess {
         private executionFactory: IPythonExecutionFactory,
         private file: IFileSystem,
         private _connection: IKernelConnection,
-        private _kernelSpec: IJupyterKernelSpec,
-        private _interpreter: InterpreterUri
+        private _kernelSpec: IJupyterKernelSpec
     ) {
         this.readyPromise = createDeferred<void>();
     }
 
     public async launch(): Promise<void> {
         this.connectionFile = await this.file.createTemporaryFile('.json');
 
-        const resource = isResource(this._interpreter) ? this._interpreter : undefined;
-        const pythonPath = isResource(this._interpreter) ? undefined : this._interpreter.path;
-
         const args = [...this._kernelSpec.argv];
         await this.file.writeFile(this.connectionFile.filePath, JSON.stringify(this._connection), {
             encoding: 'utf-8',
@@ -72,7 +68,10 @@ class KernelProcess implements IKernelProcess {
         args[4] = this.connectionFile.filePath;
         args.splice(0, 1);
 
-        const executionService = await this.executionFactory.create({ resource, pythonPath });
+        const executionService = await this.executionFactory.create({
+            resource: undefined,
+            pythonPath: this._kernelSpec.path
+        });
         const exeObs = executionService.execObservable(args, {});
 
         if (exeObs.proc) {
@@ -139,13 +138,7 @@ export class KernelLauncher implements IKernelLauncher {
         }
 
         const connection = await this.getKernelConnection();
-        const kernelProcess = new KernelProcess(
-            this.executionFactory,
-            this.file,
-            connection,
-            kernelSpec,
-            interpreterUri
-        );
+        const kernelProcess = new KernelProcess(this.executionFactory, this.file, connection, kernelSpec);
         await kernelProcess.launch();
         return kernelProcess;
     }

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

@@ -88,8 +88,16 @@ export class RawJupyterSession extends BaseJupyterSession {
         return this.session.process?.kernelSpec;
     }
 
-    public async changeKernel(_kernel: IJupyterKernelSpec | LiveKernelModel, _timeoutMS: number): Promise<void> {
-        throw new Error('Not implemented');
+    public async createNewKernelSession(
+        kernel: IJupyterKernelSpec | LiveKernelModel,
+        _timeoutMS: number
+    ): Promise<ISession> {
+        if (!this.resource || !kernel || 'session' in kernel) {
+            // Don't allow for connecting to a LiveKernelModel
+            throw new Error(localize.DataScience.sessionDisposed());
+        }
+
+        return this.startRawSession(this.resource, kernel);
     }
 
     protected startRestartSession() {

src/test/datascience/raw-kernel/rawJupyterSession.unit.test.ts

@@ -8,7 +8,7 @@ import { EventEmitter } from 'vscode';
 import { KernelSelector } from '../../../client/datascience/jupyter/kernels/kernelSelector';
 import { IKernelLauncher, IKernelProcess } from '../../../client/datascience/kernel-launcher/types';
 import { RawJupyterSession } from '../../../client/datascience/raw-kernel/rawJupyterSession';
-import { IJMPConnection } from '../../../client/datascience/types';
+import { IJMPConnection, IJupyterKernelSpec } from '../../../client/datascience/types';
 import { IServiceContainer } from '../../../client/ioc/types';
 
 // tslint:disable:no-any
@@ -89,6 +89,26 @@ suite('Data Science - RawJupyterSession', () => {
         kernelProcess.verifyAll();
     });
 
+    test('RawJupyterSession - changeKernel', async () => {
+        kernelProcess.setup((kp) => kp.dispose).verifiable(typemoq.Times.exactly(2));
+
+        await rawJupyterSession.connect({} as any, 60_000);
+
+        const newKernel: IJupyterKernelSpec = {
+            argv: [],
+            display_name: 'new kernel',
+            language: 'python',
+            name: 'newkernel',
+            path: 'path'
+        };
+
+        await rawJupyterSession.changeKernel(newKernel, 60_000);
+
+        // Four connects and two processes disposed
+        verify(kernelLauncher.launch(anything(), anything())).times(4);
+        kernelProcess.verifyAll();
+    });
+
     test('RawJupyterSession - Kill process', async () => {
         const shutdown = sinon.stub(rawJupyterSession, 'shutdown');
         shutdown.resolves();