Skip to content

Commit 6dc177c

Browse files
DonJayamanneDonJayamanne
committed
Parallelize
1 parent ebcdc65 commit 6dc177c

File tree

1 file changed

+102
-51
lines changed

1 file changed

+102
-51
lines changed

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

Lines changed: 102 additions & 51 deletions
Original file line numberDiff line numberDiff line change
@@ -3,7 +3,7 @@
33

44
'use strict';
55

6-
import { execSync } from 'child_process';
6+
import * as cp from 'child_process';
77
import * as os from 'os';
88
import * as path from 'path';
99
import * as uuid from 'uuid/v4';
@@ -25,7 +25,6 @@ import { JupyterCommandFinder } from './jupyterCommandFinder';
2525
import { JupyterConnection, JupyterServerInfo } from './jupyterConnection';
2626
import { KernelService } from './kernelService';
2727

28-
2928
/**
3029
* Responsible for starting a notebook.
3130
* Separate class as theres quite a lot of work involved in starting a notebook.
@@ -63,57 +62,11 @@ export class NotebookStarter implements Disposable {
6362
let exitCode: number | null = 0;
6463
try {
6564
// Generate a temp dir with a unique GUID, both to match up our started server and to easily clean up after
66-
const tempDir = await this.generateTempDir();
67-
this.disposables.push(tempDir);
68-
69-
// In the temp dir, create an empty config python file. This is the same
70-
// as starting jupyter with all of the defaults.
71-
const configFile = useDefaultConfig ? path.join(tempDir.path, 'jupyter_notebook_config.py') : undefined;
72-
if (configFile) {
73-
await this.fileSystem.writeFile(configFile, '');
74-
traceInfo(`Generating custom default config at ${configFile}`);
75-
}
76-
77-
// Create extra args based on if we have a config or not
78-
const extraArgs: string[] = [];
79-
if (useDefaultConfig) {
80-
extraArgs.push(`--config=${configFile}`);
81-
}
82-
// Check for the debug environment variable being set. Setting this
83-
// causes Jupyter to output a lot more information about what it's doing
84-
// under the covers and can be used to investigate problems with Jupyter.
85-
if (process.env && process.env.VSCODE_PYTHON_DEBUG_JUPYTER) {
86-
extraArgs.push('--debug');
87-
}
88-
89-
// Modify the data rate limit if starting locally. The default prevents large dataframes from being returned.
90-
extraArgs.push('--NotebookApp.iopub_data_rate_limit=10000000000.0');
91-
92-
// Check for a docker situation.
93-
try {
94-
if (await this.fileSystem.fileExists('/proc/self/cgroup')) {
95-
const cgroup = await this.fileSystem.readFile('/proc/self/cgroup');
96-
if (cgroup.includes('docker')) {
97-
// We definitely need an ip address.
98-
extraArgs.push('--ip');
99-
extraArgs.push('127.0.0.1');
100-
101-
// Now see if we need --allow-root.
102-
const idResults = execSync('id', { encoding: 'utf-8' });
103-
if (idResults.includes('(root)')) {
104-
extraArgs.push('--allow-root');
105-
}
106-
}
107-
}
108-
} catch {
109-
noop();
110-
}
111-
112-
// Use this temp file and config file to generate a list of args for our command
113-
const args: string[] = [...['--no-browser', `--notebook-dir=${tempDir.path}`], ...extraArgs];
65+
const tempDirPromise = this.generateTempDir();
66+
tempDirPromise.then(dir => this.disposables.push(dir)).ignoreErrors();
11467

11568
// Before starting the notebook process, make sure we generate a kernel spec
116-
const kernelSpec = await this.kernelService.getMatchingKernelSpec(undefined, cancelToken);
69+
const [args, kernelSpec] = await Promise.all([this.generateArguments(useDefaultConfig, tempDirPromise), this.kernelService.getMatchingKernelSpec(undefined, cancelToken)]);
11770

11871
// Make sure we haven't canceled already.
11972
if (cancelToken && cancelToken.isCancellationRequested) {
@@ -137,6 +90,7 @@ export class NotebookStarter implements Disposable {
13790
}
13891

13992
// Wait for the connection information on this result
93+
const tempDir = await tempDirPromise;
14094
const connection = await JupyterConnection.waitForConnection(tempDir.path, this.getJupyterServerInfo, launchResult, this.serviceContainer, cancelToken);
14195

14296
// Fire off telemetry for the process being talkable
@@ -159,6 +113,103 @@ export class NotebookStarter implements Disposable {
159113
}
160114
}
161115
}
116+
117+
private async generateArguments(useDefaultConfig: boolean, tempDirPromise: Promise<TemporaryDirectory>): Promise<string[]>{
118+
const mainArgs: string [] = ['--no-browser'];
119+
120+
// Create extra args based on if we have a config or not
121+
const extraArgs: string[] = [];
122+
123+
// Check for the debug environment variable being set. Setting this
124+
// causes Jupyter to output a lot more information about what it's doing
125+
// under the covers and can be used to investigate problems with Jupyter.
126+
if (process.env && process.env.VSCODE_PYTHON_DEBUG_JUPYTER) {
127+
extraArgs.push('--debug');
128+
}
129+
130+
// Modify the data rate limit if starting locally. The default prevents large dataframes from being returned.
131+
extraArgs.push('--NotebookApp.iopub_data_rate_limit=10000000000.0');
132+
133+
// Parallelize as much as possible.
134+
const promises = [this.addNotebookDirArgument(mainArgs, tempDirPromise), this.addDockerArguments(extraArgs)];
135+
if (useDefaultConfig) {
136+
promises.push(this.addConfigArgument(extraArgs, tempDirPromise));
137+
}
138+
139+
await Promise.all(promises);
140+
141+
// Use this temp file and config file to generate a list of args for our command
142+
return [...mainArgs, ...extraArgs];
143+
}
144+
145+
/**
146+
* Adds the `--notebook-dir` argument.
147+
*
148+
* @private
149+
* @param {string[]} args
150+
* @param {Promise<TemporaryDirectory>} tempDirectory
151+
* @returns {Promise<void>}
152+
* @memberof NotebookStarter
153+
*/
154+
private async addNotebookDirArgument(args: string[], tempDirectory: Promise<TemporaryDirectory>): Promise<void> {
155+
const tempDir = await tempDirectory;
156+
args.push(`--notebook-dir=${tempDir.path}`);
157+
}
158+
159+
/**
160+
* Adds the `--config` argument.
161+
*
162+
* @private
163+
* @param {string[]} args
164+
* @param {Promise<TemporaryDirectory>} tempDirectory
165+
* @returns {Promise<void>}
166+
* @memberof NotebookStarter
167+
*/
168+
private async addConfigArgument(args: string[], tempDirectory: Promise<TemporaryDirectory>): Promise<void> {
169+
const tempDir = await tempDirectory;
170+
// In the temp dir, create an empty config python file. This is the same
171+
// as starting jupyter with all of the defaults.
172+
const configFile = path.join(tempDir.path, 'jupyter_notebook_config.py');
173+
await this.fileSystem.writeFile(configFile, '');
174+
traceInfo(`Generating custom default config at ${configFile}`);
175+
176+
// Create extra args based on if we have a config or not
177+
args.push(`--config=${configFile}`);
178+
}
179+
180+
/**
181+
* Adds the `--ip` and `--alow-root` arguments when in docker.
182+
*
183+
* @private
184+
* @param {string[]} args
185+
* @param {Promise<TemporaryDirectory>} tempDirectory
186+
* @returns {Promise<void>}
187+
* @memberof NotebookStarter
188+
*/
189+
private async addDockerArguments(args: string[]): Promise<void> {
190+
// Check for a docker situation.
191+
try {
192+
const cgroup = await this.fileSystem.readFile('/proc/self/cgroup').catch(() => '');
193+
if (!cgroup.includes('docker')) {
194+
return;
195+
}
196+
// We definitely need an ip address.
197+
args.push('--ip');
198+
args.push('127.0.0.1');
199+
200+
// Now see if we need --allow-root.
201+
await new Promise(resolve => {
202+
cp.exec('id', { encoding: 'utf-8' }, (_, stdout: string | Buffer) => {
203+
if (stdout && stdout.toString().includes('(root)')) {
204+
args.push('--allow-root');
205+
}
206+
resolve();
207+
});
208+
});
209+
} catch {
210+
noop();
211+
}
212+
}
162213
private async generateTempDir(): Promise<TemporaryDirectory> {
163214
const resultDir = path.join(os.tmpdir(), uuid());
164215
await this.fileSystem.createDirectory(resultDir);

0 commit comments

Comments
 (0)