Skip to content

Commit 4cf109b

Browse files
DonJayamanneDonJayamanne
authored
Pre warm raw kernel daemons (microsoft#11487)
1 parent f1bd613 commit 4cf109b

21 files changed

Lines changed: 845 additions & 95 deletions

‎pythonFiles/vscode_datascience_helpers/kernel_launcher_daemon.py‎

Lines changed: 47 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -4,6 +4,7 @@
44
import json
55
import logging
66
import os
7+
import os.path
78
import signal
89
import sys
910
import subprocess
@@ -69,8 +70,43 @@ def m_kill_kernel(self):
6970
self.kernel.kill()
7071
except OSError:
7172
pass
72-
finally:
73-
self.kernel = None
73+
74+
@error_decorator
75+
def m_prewarm_kernel(self):
76+
"""Starts the kernel process with the module
77+
"""
78+
self.log.info("Pre-Warm DS Kernel in DS Kernel Launcher Daemon")
79+
isolated_runner = os.path.join(
80+
os.path.dirname(__file__), "..", "pyvsc-run-isolated.py"
81+
)
82+
kernel_prewarm_starter = os.path.join(
83+
os.path.dirname(__file__), "kernel_prewarm_starter.py"
84+
)
85+
86+
def prewarm_kernel():
87+
cmd = [sys.executable, isolated_runner, kernel_prewarm_starter]
88+
self._start_kernel_observable_in_background(cmd)
89+
self.log.info("Kernel launched, with PID as a daemon %s", self.kernel.pid)
90+
91+
return self._execute_and_capture_output(prewarm_kernel)
92+
93+
@error_decorator
94+
def m_start_prewarmed_kernel(self, args=[]):
95+
"""Starts the pre-warmed kernel process.
96+
"""
97+
self.log.info(
98+
"Start pre-warmed Kernel in DS Kernel Launcher Daemon %s with args %s",
99+
self.kernel.pid,
100+
args,
101+
)
102+
103+
def start_kernel():
104+
self.kernel.stdin.write(
105+
"{0}{1}".format(json.dumps(args), os.linesep).encode("utf-8")
106+
)
107+
self.kernel.stdin.close()
108+
109+
return self._execute_and_capture_output(start_kernel)
74110

75111
@error_decorator
76112
def m_exec_module(self, module_name, args=[], cwd=None, env=None):
@@ -83,7 +119,6 @@ def m_exec_module(self, module_name, args=[], cwd=None, env=None):
83119
)
84120

85121
def start_kernel():
86-
thread_args = (module_name, args, cwd, env)
87122
self._exec_module_observable_in_background(module_name, args, cwd, env)
88123

89124
return self._execute_and_capture_output(start_kernel)
@@ -127,6 +162,13 @@ def _exec_module_observable_in_background(
127162
)
128163
args = [] if args is None else args
129164
cmd = [sys.executable, "-m", module_name] + args
165+
self._start_kernel_observable_in_background(cmd, cwd, env)
166+
self.kernel.stdin.close()
167+
168+
def _start_kernel_observable_in_background(self, cmd, cwd=None, env=None):
169+
self.log.info(
170+
"Exec in DS Kernel Launcher Daemon (observable) %s", cmd,
171+
)
130172
# As the kernel is launched from this same python executable, ensure the kernel variables
131173
# are merged with the variables of this current environment.
132174
new_env_vars = {} if env is None else env
@@ -136,15 +178,12 @@ def _exec_module_observable_in_background(
136178
cmd,
137179
stdout=subprocess.PIPE,
138180
stderr=subprocess.PIPE,
181+
stdin=subprocess.PIPE,
139182
cwd=cwd,
140183
env=env,
141184
independent=False,
142185
)
143-
self.log.info(
144-
"Exec in DS Kernel Launcher Daemon (observable) %s with args %s",
145-
subprocess.PIPE,
146-
subprocess.PIPE,
147-
)
186+
self.log.info("Exec in DS Kernel Launcher Daemon (observable)")
148187

149188
self.kernel = proc
150189
self.log.info("Kernel launched, with PID %s", proc.pid)

‎pythonFiles/vscode_datascience_helpers/kernel_prewarm_starter.py‎

Lines changed: 31 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,31 @@
1+
"""Starts the python kernel the same way we would if we used the CLI.
2+
IPykernel module must be run in the main thread, hence the blocking mechanism to
3+
ready stdin to figure out when ipykernel needs to be started.
4+
Running the module ipykernel effectively means we're starting the python kernel.
5+
"""
6+
7+
# Copyright (c) Microsoft Corporation. All rights reserved.
8+
# Licensed under the MIT License.
9+
10+
if __name__ != "__main__":
11+
raise Exception("{} cannot be imported".format(__name__))
12+
13+
import json
14+
import runpy
15+
import sys
16+
17+
# Assumption is ipykernel is available.
18+
# Preload module (speed up, so when we really need it, it has already been loaded).
19+
from ipykernel import kernelapp as app
20+
21+
# Block till we read somethign from `stdin`.
22+
# As soon as we get something this is a trigger to start.
23+
# The value passed into stdin would be the arguments we need to place into sys.argv.
24+
input_json = sys.stdin.readline().strip()
25+
sys.argv = json.loads(input_json)
26+
module = sys.argv[0]
27+
28+
# Note, we must launch ipykenel in the main thread for kernel interrupt to work on windows.
29+
30+
# Start kernel in current process.
31+
runpy.run_module(module, run_name="__main__", alter_sys=False)

‎src/client/datascience/activation.ts‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -13,6 +13,7 @@ import { sendTelemetryEvent } from '../telemetry';
1313
import { JupyterDaemonModule, Telemetry } from './constants';
1414
import { ActiveEditorContextService } from './context/activeEditorContext';
1515
import { JupyterInterpreterService } from './jupyter/interpreter/jupyterInterpreterService';
16+
import { KernelDaemonPreWarmer } from './kernel-launcher/kernelDaemonPreWarmer';
1617
import { INotebookAndInteractiveWindowUsageTracker, INotebookEditor, INotebookEditorProvider } from './types';
1718

1819
@injectable()
@@ -24,13 +25,15 @@ export class Activation implements IExtensionSingleActivationService {
2425
@inject(IPythonExecutionFactory) private readonly factory: IPythonExecutionFactory,
2526
@inject(IDisposableRegistry) private readonly disposables: IDisposableRegistry,
2627
@inject(ActiveEditorContextService) private readonly contextService: ActiveEditorContextService,
28+
@inject(KernelDaemonPreWarmer) private readonly daemonPoolPrewarmer: KernelDaemonPreWarmer,
2729
@inject(INotebookAndInteractiveWindowUsageTracker)
2830
private readonly tracker: INotebookAndInteractiveWindowUsageTracker
2931
) {}
3032
public async activate(): Promise<void> {
3133
this.disposables.push(this.notebookEditorProvider.onDidOpenNotebookEditor(this.onDidOpenNotebookEditor, this));
3234
this.disposables.push(this.jupyterInterpreterService.onDidChangeInterpreter(this.onDidChangeInterpreter, this));
3335
this.contextService.activate().ignoreErrors();
36+
this.daemonPoolPrewarmer.activate(undefined).ignoreErrors();
3437
this.tracker.startTracking();
3538
}
3639

‎src/client/datascience/interactive-common/notebookProvider.ts‎

Lines changed: 16 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -5,13 +5,14 @@
55

66
import { inject, injectable } from 'inversify';
77
import { EventEmitter, Uri } from 'vscode';
8+
import { IWorkspaceService } from '../../common/application/types';
89
import { LocalZMQKernel } from '../../common/experimentGroups';
910
import { traceError, traceInfo } from '../../common/logger';
1011
import { IFileSystem } from '../../common/platform/types';
11-
import { IConfigurationService, IDisposableRegistry, IExperimentsManager } from '../../common/types';
12+
import { IConfigurationService, IDisposableRegistry, IExperimentsManager, Resource } from '../../common/types';
1213
import { noop } from '../../common/utils/misc';
1314
import { sendTelemetryEvent } from '../../telemetry';
14-
import { Settings, Telemetry } from '../constants';
15+
import { Identifiers, Settings, Telemetry } from '../constants';
1516
import {
1617
ConnectNotebookProviderOptions,
1718
GetNotebookOptions,
@@ -41,7 +42,8 @@ export class NotebookProvider implements INotebookProvider {
4142
@inject(IRawNotebookProvider) private readonly rawNotebookProvider: IRawNotebookProvider,
4243
@inject(IJupyterNotebookProvider) private readonly jupyterNotebookProvider: IJupyterNotebookProvider,
4344
@inject(IConfigurationService) private readonly configuration: IConfigurationService,
44-
@inject(IExperimentsManager) private readonly experimentsManager: IExperimentsManager
45+
@inject(IExperimentsManager) private readonly experimentsManager: IExperimentsManager,
46+
@inject(IWorkspaceService) private readonly workspaceService: IWorkspaceService
4547
) {
4648
disposables.push(editorProvider.onDidCloseNotebookEditor(this.onDidCloseNotebookEditor, this));
4749
disposables.push(
@@ -96,13 +98,18 @@ export class NotebookProvider implements INotebookProvider {
9698
}
9799

98100
// Finally create if needed
101+
let resource: Resource = options.identity;
102+
if (options.identity.scheme === Identifiers.HistoryPurpose) {
103+
// If we have any workspaces, then use the first available workspace.
104+
// This is required, else using `undefined` as a resource when we have worksapce folders is a different meaning.
105+
// This means interactive window doesn't properly support mult-root workspaces as we pick first workspace.
106+
// Ideally we need to pick the resource of the corresponding Python file.
107+
resource = this.workspaceService.hasWorkspaceFolders
108+
? this.workspaceService.workspaceFolders![0]!.uri
109+
: undefined;
110+
}
99111
const promise = rawKernel
100-
? this.rawNotebookProvider.createNotebook(
101-
options.identity,
102-
options.identity,
103-
options.disableUI,
104-
options.metadata
105-
)
112+
? this.rawNotebookProvider.createNotebook(options.identity, resource, options.disableUI, options.metadata)
106113
: this.jupyterNotebookProvider.createNotebook(options);
107114

108115
this.cacheNotebookPromise(options.identity, promise);

‎src/client/datascience/kernel-launcher/kernelDaemon.ts‎

Lines changed: 74 additions & 35 deletions
Original file line numberDiff line numberDiff line change
@@ -5,18 +5,23 @@
55

66
import { ChildProcess } from 'child_process';
77
import { Subject } from 'rxjs/Subject';
8-
import { MessageConnection, NotificationType, RequestType0 } from 'vscode-jsonrpc';
9-
import { BasePythonDaemon } from '../../common/process/baseDaemon';
8+
import { MessageConnection, NotificationType, RequestType, RequestType0 } from 'vscode-jsonrpc';
9+
import { BasePythonDaemon, ExecResponse } from '../../common/process/baseDaemon';
1010
import {
1111
IPythonExecutionService,
1212
ObservableExecutionResult,
1313
Output,
1414
SpawnOptions,
1515
StdErrError
1616
} from '../../common/process/types';
17+
import { noop } from '../../common/utils/misc';
1718
import { IPythonKernelDaemon, PythonKernelDiedError } from './types';
1819

1920
export class PythonKernelDaemon extends BasePythonDaemon implements IPythonKernelDaemon {
21+
private started?: boolean;
22+
private preWarmed?: boolean;
23+
private outputHooked?: boolean;
24+
private readonly subject = new Subject<Output<string>>();
2025
constructor(
2126
pythonExecutionService: IPythonExecutionService,
2227
pythonPath: string,
@@ -33,21 +38,77 @@ export class PythonKernelDaemon extends BasePythonDaemon implements IPythonKerne
3338
const request = new RequestType0<void, void, void>('kill_kernel');
3439
await this.sendRequestWithoutArgs(request);
3540
}
41+
public dispose() {
42+
this.kill().catch(noop);
43+
super.dispose();
44+
}
45+
public async preWarm() {
46+
if (this.started) {
47+
return;
48+
}
49+
this.preWarmed = true;
50+
this.monitorOutput();
51+
const request = new RequestType0<void, void, void>('prewarm_kernel');
52+
53+
await this.sendRequestWithoutArgs(request);
54+
}
55+
3656
public async start(
3757
moduleName: string,
3858
args: string[],
3959
options: SpawnOptions
4060
): Promise<ObservableExecutionResult<string>> {
41-
const subject = new Subject<Output<string>>();
42-
let stdErr = '';
61+
if (options.throwOnStdErr) {
62+
throw new Error("'throwOnStdErr' not supported in spawnOptions for KernelDaemon.start");
63+
}
64+
if (options.mergeStdOutErr) {
65+
throw new Error("'mergeStdOutErr' not supported in spawnOptions for KernelDaemon.start");
66+
}
67+
if (options.cwd) {
68+
throw new Error("'cwd' not supported in spawnOptions for KernelDaemon.start");
69+
}
70+
if (this.started) {
71+
throw new Error('Kernel has already been started in daemon');
72+
}
73+
this.started = true;
74+
this.monitorOutput();
75+
76+
if (this.preWarmed) {
77+
const request = new RequestType<{ args: string[] }, ExecResponse, void, void>('start_prewarmed_kernel');
78+
await this.sendRequest(request, { args: [moduleName].concat(args) });
79+
} else {
80+
// No need of the output here, we'll tap into the output coming from daemon `this.outputObservale`.
81+
// This is required because execModule will never end.
82+
// We cannot use `execModuleObservable` as that only works where the daemon is busy seeerving on request and we wait for it to finish.
83+
// In this case we're never going to wait for the module to run to end. Cuz when we run `pytohn -m ipykernel`, it never ends.
84+
// It only ends when the kernel dies, meaning the kernel process is dead.
85+
// What we need is to be able to run the module and keep getting a stream of stdout/stderr.
86+
// & also be able to execute other python code. I.e. we need a daemon.
87+
// For this we run the `ipykernel` code in a separate thread.
88+
// This is why when we run `execModule` in the Kernel daemon, it finishes (comes back) quickly.
89+
// However in reality it is running in the background.
90+
// See `m_exec_module_observable` in `kernel_launcher_daemon.py`.
91+
await this.execModule(moduleName, args, options);
92+
}
93+
94+
return {
95+
proc: this.proc,
96+
dispose: () => this.dispose(),
97+
out: this.subject
98+
};
99+
}
100+
private monitorOutput() {
101+
if (this.outputHooked) {
102+
return;
103+
}
104+
this.outputHooked = true;
43105
// Message from daemon when kernel dies.
44106
const KernelDiedNotification = new NotificationType<{ exit_code: string; reason?: string }, void>(
45107
'kernel_died'
46108
);
47109
this.connection.onNotification(KernelDiedNotification, (output) => {
48-
// If we don't have a reason why things failed, then just include the stderr.
49-
subject.error(
50-
new PythonKernelDiedError({ exitCode: parseInt(output.exit_code, 10), reason: output.reason || stdErr })
110+
this.subject.error(
111+
new PythonKernelDiedError({ exitCode: parseInt(output.exit_code, 10), reason: output.reason })
51112
);
52113
});
53114

@@ -56,39 +117,17 @@ export class PythonKernelDaemon extends BasePythonDaemon implements IPythonKerne
56117
// sptting stuff into stdout/stderr.
57118
this.outputObservale.subscribe(
58119
(out) => {
59-
if (out.source === 'stderr' && options.throwOnStdErr) {
60-
stdErr += out.out;
61-
subject.error(new StdErrError(out.out));
62-
} else if (out.source === 'stderr' && options.mergeStdOutErr) {
63-
subject.next({ source: 'stdout', out: out.out });
120+
if (out.source === 'stderr') {
121+
this.subject.error(new StdErrError(out.out));
64122
} else {
65-
subject.next(out);
123+
this.subject.next(out);
66124
}
67125
},
68-
subject.error.bind(subject),
69-
subject.complete.bind(subject)
126+
this.subject.error.bind(this.subject),
127+
this.subject.complete.bind(this.subject)
70128
);
71129

72130
// If the daemon dies, then kernel is also dead.
73-
this.closed.catch((error) => subject.error(new PythonKernelDiedError({ error })));
74-
75-
// No need of the output here, we'll tap into the output coming from daemon `this.outputObservale`.
76-
// This is required because execModule will never end.
77-
// We cannot use `execModuleObservable` as that only works where the daemon is busy seeerving on request and we wait for it to finish.
78-
// In this case we're never going to wait for the module to run to end. Cuz when we run `pytohn -m ipykernel`, it never ends.
79-
// It only ends when the kernel dies, meaning the kernel process is dead.
80-
// What we need is to be able to run the module and keep getting a stream of stdout/stderr.
81-
// & also be able to execute other python code. I.e. we need a daemon.
82-
// For this we run the `ipykernel` code in a separate thread.
83-
// This is why when we run `execModule` in the Kernel daemon, it finishes (comes back) quickly.
84-
// However in reality it is running in the background.
85-
// See `m_exec_module_observable` in `kernel_launcher_daemon.py`.
86-
await this.execModule(moduleName, args, options);
87-
88-
return {
89-
proc: this.proc,
90-
dispose: () => this.dispose(),
91-
out: subject
92-
};
131+
this.closed.catch((error) => this.subject.error(new PythonKernelDiedError({ error })));
93132
}
94133
}

0 commit comments

Comments
 (0)