Skip to content

Commit 96496cc

Browse files
DonJayamanneDonJayamanne
committed
Parallelize
1 parent 072ccea commit 96496cc

File tree

1 file changed

+102
-50
lines changed

1 file changed

+102
-50
lines changed

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

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

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

11771
// Make sure we haven't canceled already.
11872
if (cancelToken && cancelToken.isCancellationRequested) {
@@ -136,6 +90,7 @@ export class NotebookStarter implements Disposable {
13690
}
13791

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

14196
// Fire off telemetry for the process being talkable
@@ -158,6 +113,103 @@ export class NotebookStarter implements Disposable {
158113
}
159114
}
160115
}
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+
}
161213
private async generateTempDir(): Promise<TemporaryDirectory> {
162214
const resultDir = path.join(os.tmpdir(), uuid());
163215
await this.fileSystem.createDirectory(resultDir);

0 commit comments

Comments
 (0)