Skip to content

Commit 87dd3d6

Browse files
IanMatthewHuffIanMatthewHuff
authored
Second refactor for non-jupyter kernel search (#11527)
1 parent c23d799 commit 87dd3d6

7 files changed

Lines changed: 192 additions & 102 deletions

File tree

‎src/client/datascience/jupyter/jupyterExecution.ts‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -145,6 +145,7 @@ export class JupyterExecutionBase implements IJupyterExecution {
145145
traceInfo(`Getting kernel specs for ${options ? options.purpose : 'unknown type of'} server`);
146146
kernelSpecInterpreterPromise = this.kernelSelector.getKernelForLocalConnection(
147147
undefined,
148+
'jupyter',
148149
undefined,
149150
options?.metadata,
150151
!allowUI,

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

Lines changed: 110 additions & 54 deletions
Original file line numberDiff line numberDiff line change
@@ -16,6 +16,7 @@ import { StopWatch } from '../../../common/utils/stopWatch';
1616
import { IInterpreterService, PythonInterpreter } from '../../../interpreter/contracts';
1717
import { IEventNamePropertyMapping, sendTelemetryEvent } from '../../../telemetry';
1818
import { KnownNotebookLanguages, Telemetry } from '../../constants';
19+
import { IKernelFinder } from '../../kernel-launcher/types';
1920
import { reportAction } from '../../progress/decorator';
2021
import { ReportableAction } from '../../progress/types';
2122
import { IJupyterKernelSpec, IJupyterSessionManager, IKernelDependencyService } from '../../types';
@@ -57,7 +58,8 @@ export class KernelSelector {
5758
@inject(IApplicationShell) private readonly applicationShell: IApplicationShell,
5859
@inject(KernelService) private readonly kernelService: KernelService,
5960
@inject(IInterpreterService) private readonly interpreterService: IInterpreterService,
60-
@inject(IKernelDependencyService) private readonly kernelDepdencyService: IKernelDependencyService
61+
@inject(IKernelDependencyService) private readonly kernelDepdencyService: IKernelDependencyService,
62+
@inject(IKernelFinder) private readonly kernelFinder: IKernelFinder
6163
) {}
6264

6365
/**
@@ -160,6 +162,7 @@ export class KernelSelector {
160162
@reportAction(ReportableAction.KernelsGetKernelForLocalConnection)
161163
public async getKernelForLocalConnection(
162164
resource: Resource,
165+
type: 'raw' | 'jupyter' | 'noConnection',
163166
sessionManager?: IJupyterSessionManager,
164167
notebookMetadata?: nbformat.INotebookMetadata,
165168
disableUI?: boolean,
@@ -171,68 +174,28 @@ export class KernelSelector {
171174
interpreterFound: false,
172175
promptedToSelect: false
173176
};
174-
// When this method is called, we know we've started a local jupyter server.
177+
// When this method is called, we know we've started a local jupyter server or are connecting raw
175178
// Lets pre-warm the list of local kernels.
176179
this.selectionProvider
177-
.getKernelSelectionsForLocalSession(resource, 'jupyter', sessionManager, cancelToken)
180+
.getKernelSelectionsForLocalSession(resource, type, sessionManager, cancelToken)
178181
.ignoreErrors();
179182

180183
let selection: KernelSpecInterpreter = {};
181-
if (notebookMetadata?.kernelspec) {
182-
selection.kernelSpec = await this.kernelService.findMatchingKernelSpec(
183-
notebookMetadata?.kernelspec,
184+
185+
if (type === 'jupyter') {
186+
selection = await this.getKernelForLocalJupyterConnection(
187+
resource,
188+
stopWatch,
189+
telemetryProps,
184190
sessionManager,
191+
notebookMetadata,
192+
disableUI,
185193
cancelToken
186194
);
187-
if (selection.kernelSpec) {
188-
selection.interpreter = await this.kernelService.findMatchingInterpreter(
189-
selection.kernelSpec,
190-
cancelToken
191-
);
192-
sendTelemetryEvent(Telemetry.UseExistingKernel);
193-
194-
// Make sure we update the environment in the kernel before using it
195-
await this.kernelService.updateKernelEnvironment(
196-
selection.interpreter,
197-
selection.kernelSpec,
198-
cancelToken
199-
);
200-
} else {
201-
// No kernel info, hence prmopt to use current interpreter as a kernel.
202-
const activeInterpreter = await this.interpreterService.getActiveInterpreter(resource);
203-
if (activeInterpreter) {
204-
selection = await this.useInterpreterAsKernel(
205-
resource,
206-
activeInterpreter,
207-
'jupyter',
208-
notebookMetadata.kernelspec.display_name,
209-
sessionManager,
210-
disableUI,
211-
cancelToken
212-
);
213-
} else {
214-
telemetryProps.promptedToSelect = true;
215-
selection = await this.selectLocalKernel(
216-
resource,
217-
'jupyter',
218-
stopWatch,
219-
sessionManager,
220-
cancelToken
221-
);
222-
}
223-
}
224-
} else {
225-
// No kernel info, hence use current interpreter as a kernel.
226-
const activeInterpreter = await this.interpreterService.getActiveInterpreter(resource);
227-
if (activeInterpreter) {
228-
selection.interpreter = activeInterpreter;
229-
selection.kernelSpec = await this.kernelService.searchAndRegisterKernel(
230-
activeInterpreter,
231-
disableUI,
232-
cancelToken
233-
);
234-
}
195+
} else if (type === 'raw') {
196+
selection = await this.getKernelForLocalRawConnection(resource, notebookMetadata, cancelToken);
235197
}
198+
236199
// If still not found, log an error (this seems possible for some people, so use the default)
237200
if (!selection.kernelSpec) {
238201
traceError('Jupyter Kernel Spec not found for a local connection');
@@ -313,6 +276,99 @@ export class KernelSelector {
313276
interpreter: interpreter
314277
};
315278
}
279+
280+
// Get our kernelspec and matching interpreter for a connection to a local jupyter server
281+
private async getKernelForLocalJupyterConnection(
282+
resource: Resource,
283+
stopWatch: StopWatch,
284+
telemetryProps: IEventNamePropertyMapping[Telemetry.FindKernelForLocalConnection],
285+
sessionManager?: IJupyterSessionManager,
286+
notebookMetadata?: nbformat.INotebookMetadata,
287+
disableUI?: boolean,
288+
cancelToken?: CancellationToken
289+
): Promise<KernelSpecInterpreter> {
290+
let selection: KernelSpecInterpreter = {};
291+
if (notebookMetadata?.kernelspec) {
292+
selection.kernelSpec = await this.kernelService.findMatchingKernelSpec(
293+
notebookMetadata?.kernelspec,
294+
sessionManager,
295+
cancelToken
296+
);
297+
if (selection.kernelSpec) {
298+
selection.interpreter = await this.kernelService.findMatchingInterpreter(
299+
selection.kernelSpec,
300+
cancelToken
301+
);
302+
sendTelemetryEvent(Telemetry.UseExistingKernel);
303+
304+
// Make sure we update the environment in the kernel before using it
305+
await this.kernelService.updateKernelEnvironment(
306+
selection.interpreter,
307+
selection.kernelSpec,
308+
cancelToken
309+
);
310+
} else {
311+
// No kernel info, hence prmopt to use current interpreter as a kernel.
312+
const activeInterpreter = await this.interpreterService.getActiveInterpreter(resource);
313+
if (activeInterpreter) {
314+
selection = await this.useInterpreterAsKernel(
315+
resource,
316+
activeInterpreter,
317+
'jupyter',
318+
notebookMetadata.kernelspec.display_name,
319+
sessionManager,
320+
disableUI,
321+
cancelToken
322+
);
323+
} else {
324+
telemetryProps.promptedToSelect = true;
325+
selection = await this.selectLocalKernel(
326+
resource,
327+
'jupyter',
328+
stopWatch,
329+
sessionManager,
330+
cancelToken
331+
);
332+
}
333+
}
334+
} else {
335+
// No kernel info, hence use current interpreter as a kernel.
336+
const activeInterpreter = await this.interpreterService.getActiveInterpreter(resource);
337+
if (activeInterpreter) {
338+
selection.interpreter = activeInterpreter;
339+
selection.kernelSpec = await this.kernelService.searchAndRegisterKernel(
340+
activeInterpreter,
341+
disableUI,
342+
cancelToken
343+
);
344+
}
345+
}
346+
347+
return selection;
348+
}
349+
350+
// Get our kernelspec and interpreter for a local raw connection
351+
private async getKernelForLocalRawConnection(
352+
resource: Resource,
353+
notebookMetadata?: nbformat.INotebookMetadata,
354+
cancelToken?: CancellationToken
355+
): Promise<KernelSpecInterpreter> {
356+
const selection: KernelSpecInterpreter = {};
357+
358+
// First use our kernel finder to locate a kernelspec on disk
359+
selection.kernelSpec = await this.kernelFinder.findKernelSpec(
360+
resource,
361+
notebookMetadata?.kernelspec?.name,
362+
cancelToken
363+
);
364+
365+
if (selection.kernelSpec) {
366+
// Locate the interpreter that matches our kernelspec
367+
selection.interpreter = await this.kernelService.findMatchingInterpreter(selection.kernelSpec, cancelToken);
368+
}
369+
return selection;
370+
}
371+
316372
private async selectKernel(
317373
resource: Resource,
318374
type: 'raw' | 'jupyter' | 'noConnection',

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

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -289,6 +289,7 @@ export class HostJupyterServer extends LiveShareParticipantHost(JupyterServerBas
289289
const kernelInfo = await (launchInfo.connectionInfo.localLaunch
290290
? this.kernelSelector.getKernelForLocalConnection(
291291
resource,
292+
'jupyter',
292293
sessionManager,
293294
notebookMetadata,
294295
isTestExecution(),

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

Lines changed: 45 additions & 38 deletions
Original file line numberDiff line numberDiff line change
@@ -23,11 +23,10 @@ import * as localize from '../../../common/utils/localize';
2323
import { IServiceContainer } from '../../../ioc/types';
2424
import { Identifiers, LiveShare, Settings } from '../../constants';
2525
import { KernelSelector } from '../../jupyter/kernels/kernelSelector';
26-
import { KernelService } from '../../jupyter/kernels/kernelService';
2726
import { HostJupyterNotebook } from '../../jupyter/liveshare/hostJupyterNotebook';
2827
import { LiveShareParticipantHost } from '../../jupyter/liveshare/liveShareParticipantMixin';
2928
import { IRoleBasedObject } from '../../jupyter/liveshare/roleBasedFactory';
30-
import { IKernelFinder, IKernelLauncher } from '../../kernel-launcher/types';
29+
import { IKernelLauncher } from '../../kernel-launcher/types';
3130
import { ProgressReporter } from '../../progress/progressReporter';
3231
import {
3332
IJupyterKernelSpec,
@@ -57,9 +56,7 @@ export class HostRawNotebookProvider
5756
private fs: IFileSystem,
5857
private serviceContainer: IServiceContainer,
5958
private kernelLauncher: IKernelLauncher,
60-
private kernelFinder: IKernelFinder,
6159
private kernelSelector: KernelSelector,
62-
private kernelService: KernelService,
6360
private progressReporter: ProgressReporter,
6461
private outputChannel: IOutputChannel
6562
) {
@@ -112,46 +109,56 @@ export class HostRawNotebookProvider
112109
try {
113110
const launchTimeout = this.configService.getSettings().datascience.jupyterLaunchTimeout;
114111

115-
// Before we try to connect we need to find a kernel and install ipykernel
116-
const kernelSpec = await this.kernelFinder.findKernelSpec(
112+
// We need to locate kernelspec and possible interpreter for this launch based on resource and notebook metadata
113+
const kernelSpecInterpreter = await this.kernelSelector.getKernelForLocalConnection(
117114
resource,
118-
notebookMetadata?.kernelspec?.name,
115+
'raw',
116+
undefined,
117+
notebookMetadata,
118+
disableUI,
119119
cancelToken
120120
);
121121

122-
// Locate the interpreter that matches our kernelspec
123-
const interpreter = await this.kernelService.findMatchingInterpreter(kernelSpec);
124-
125-
await rawSession.connect(kernelSpec, launchTimeout, interpreter, cancelToken);
126-
127-
// Get the execution info for our notebook
128-
const info = await this.getExecutionInfo(kernelSpec);
129-
130-
if (rawSession.isConnected) {
131-
// Create our notebook
132-
const notebook = new HostJupyterNotebook(
133-
this.liveShare,
134-
rawSession,
135-
this.configService,
136-
this.disposableRegistry,
137-
info,
138-
this.serviceContainer.getAll<INotebookExecutionLogger>(INotebookExecutionLogger),
139-
resource,
140-
identity,
141-
this.getDisposedError.bind(this),
142-
this.workspaceService,
143-
this.appShell,
144-
this.fs
122+
// Interpreter is optional, but we must have a kernel spec for a raw launch
123+
if (!kernelSpecInterpreter.kernelSpec) {
124+
notebookPromise.reject('Failed to find a kernelspec to use for ipykernel launch');
125+
} else {
126+
await rawSession.connect(
127+
kernelSpecInterpreter.kernelSpec,
128+
launchTimeout,
129+
kernelSpecInterpreter.interpreter,
130+
cancelToken
145131
);
146132

147-
// Run initial setup
148-
await notebook.initialize(cancelToken);
149-
150-
traceInfo(`Finished connecting ${this.id}`);
151-
152-
notebookPromise.resolve(notebook);
153-
} else {
154-
notebookPromise.reject(this.getDisposedError());
133+
// Get the execution info for our notebook
134+
const info = await this.getExecutionInfo(kernelSpecInterpreter.kernelSpec);
135+
136+
if (rawSession.isConnected) {
137+
// Create our notebook
138+
const notebook = new HostJupyterNotebook(
139+
this.liveShare,
140+
rawSession,
141+
this.configService,
142+
this.disposableRegistry,
143+
info,
144+
this.serviceContainer.getAll<INotebookExecutionLogger>(INotebookExecutionLogger),
145+
resource,
146+
identity,
147+
this.getDisposedError.bind(this),
148+
this.workspaceService,
149+
this.appShell,
150+
this.fs
151+
);
152+
153+
// Run initial setup
154+
await notebook.initialize(cancelToken);
155+
156+
traceInfo(`Finished connecting ${this.id}`);
157+
158+
notebookPromise.resolve(notebook);
159+
} else {
160+
notebookPromise.reject(this.getDisposedError());
161+
}
155162
}
156163
} catch (ex) {
157164
// Make sure we shut down our session in case we started a process

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

Lines changed: 1 addition & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -19,10 +19,9 @@ import {
1919
import { IServiceContainer } from '../../ioc/types';
2020
import { JUPYTER_OUTPUT_CHANNEL } from '../constants';
2121
import { KernelSelector } from '../jupyter/kernels/kernelSelector';
22-
import { KernelService } from '../jupyter/kernels/kernelService';
2322
import { IRoleBasedObject, RoleBasedFactory } from '../jupyter/liveshare/roleBasedFactory';
2423
import { ILiveShareHasRole } from '../jupyter/liveshare/types';
25-
import { IKernelFinder, IKernelLauncher } from '../kernel-launcher/types';
24+
import { IKernelLauncher } from '../kernel-launcher/types';
2625
import { ProgressReporter } from '../progress/progressReporter';
2726
import { INotebook, IRawConnection, IRawNotebookProvider } from '../types';
2827
import { GuestRawNotebookProvider } from './liveshare/guestRawNotebookProvider';
@@ -42,9 +41,7 @@ type RawNotebookProviderClassType = {
4241
fs: IFileSystem,
4342
serviceContainer: IServiceContainer,
4443
kernelLauncher: IKernelLauncher,
45-
kernelFinder: IKernelFinder,
4644
kernelSelector: KernelSelector,
47-
kernelService: KernelService,
4845
progressReporter: ProgressReporter,
4946
outputChannel: IOutputChannel
5047
): IRawNotebookProviderInterface;
@@ -67,9 +64,7 @@ export class RawNotebookProviderWrapper implements IRawNotebookProvider, ILiveSh
6764
@inject(IFileSystem) fs: IFileSystem,
6865
@inject(IServiceContainer) serviceContainer: IServiceContainer,
6966
@inject(IKernelLauncher) kernelLauncher: IKernelLauncher,
70-
@inject(IKernelFinder) kernelFinder: IKernelFinder,
7167
@inject(KernelSelector) kernelSelector: KernelSelector,
72-
@inject(KernelService) kernelService: KernelService,
7368
@inject(ProgressReporter) progressReporter: ProgressReporter,
7469
@inject(IOutputChannel) @named(JUPYTER_OUTPUT_CHANNEL) outputChannel: IOutputChannel
7570
) {
@@ -88,9 +83,7 @@ export class RawNotebookProviderWrapper implements IRawNotebookProvider, ILiveSh
8883
fs,
8984
serviceContainer,
9085
kernelLauncher,
91-
kernelFinder,
9286
kernelSelector,
93-
kernelService,
9487
progressReporter,
9588
outputChannel
9689
);

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

Lines changed: 8 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -962,7 +962,14 @@ suite('Jupyter Execution', async () => {
962962
env: undefined
963963
};
964964
when(
965-
kernelSelector.getKernelForLocalConnection(anything(), anything(), anything(), anything(), anything())
965+
kernelSelector.getKernelForLocalConnection(
966+
anything(),
967+
anything(),
968+
anything(),
969+
anything(),
970+
anything(),
971+
anything()
972+
)
966973
).thenResolve({
967974
kernelSpec
968975
});

0 commit comments

Comments
 (0)