Skip to content

Commit ef6d780

Browse files
DonJayamanneDonJayamanne
authored
Improvements to startup of notebook (microsoft#8332)
* Parallelize * Fix tests * Fix typo * Fixed order of arguments to ensure tests pass * Linter issues
1 parent 58a6e24 commit ef6d780

2 files changed

Lines changed: 99 additions & 56 deletions

File tree

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

Lines changed: 94 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';
@@ -14,7 +14,6 @@ import { IFileSystem, TemporaryDirectory } from '../../common/platform/types';
1414
import { IPythonExecutionFactory, SpawnOptions } from '../../common/process/types';
1515
import { IDisposable } from '../../common/types';
1616
import * as localize from '../../common/utils/localize';
17-
import { noop } from '../../common/utils/misc';
1817
import { StopWatch } from '../../common/utils/stopWatch';
1918
import { EXTENSION_ROOT_DIR } from '../../constants';
2019
import { IServiceContainer } from '../../ioc/types';
@@ -62,57 +61,11 @@ export class NotebookStarter implements Disposable {
6261
let exitCode: number | null = 0;
6362
try {
6463
// 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];
64+
const tempDirPromise = this.generateTempDir();
65+
tempDirPromise.then(dir => this.disposables.push(dir)).ignoreErrors();
11366

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

11770
// Make sure we haven't canceled already.
11871
if (cancelToken && cancelToken.isCancellationRequested) {
@@ -136,6 +89,7 @@ export class NotebookStarter implements Disposable {
13689
}
13790

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

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

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

Lines changed: 5 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -460,7 +460,7 @@ suite('Jupyter Execution', async () => {
460460
const getServerInfoPath = path.join(EXTENSION_ROOT_DIR, 'pythonFiles', 'datascience', 'getServerInfo.py');
461461
setupPythonService(service, undefined, [getServerInfoPath], Promise.resolve({ stdout: 'failure to get server infos' }));
462462
setupPythonServiceExecObservable(service, 'jupyter', ['kernelspec', 'list'], [], []);
463-
setupPythonServiceExecObservable(service, 'jupyter', ['notebook', '--no-browser', /--notebook-dir=.*/, /.*/, '--NotebookApp.iopub_data_rate_limit=10000000000.0'], [], notebookStdErr ? notebookStdErr : ['http://localhost:8888/?token=198']);
463+
setupPythonServiceExecObservable(service, 'jupyter', ['notebook', '--no-browser', /--notebook-dir=.*/, /--config=.*/, '--NotebookApp.iopub_data_rate_limit=10000000000.0'], [], notebookStdErr ? notebookStdErr : ['http://localhost:8888/?token=198']);
464464

465465
}
466466

@@ -472,7 +472,7 @@ suite('Jupyter Execution', async () => {
472472
const getServerInfoPath = path.join(EXTENSION_ROOT_DIR, 'pythonFiles', 'datascience', 'getServerInfo.py');
473473
setupPythonService(service, undefined, [getServerInfoPath], Promise.resolve({ stdout: 'failure to get server infos' }));
474474
setupPythonServiceExecObservable(service, 'jupyter', ['kernelspec', 'list'], [], []);
475-
setupPythonServiceExecObservable(service, 'jupyter', ['notebook', '--no-browser', /--notebook-dir=.*/, /.*/, '--NotebookApp.iopub_data_rate_limit=10000000000.0'], [], notebookStdErr ? notebookStdErr : ['http://localhost:8888/?token=198']);
475+
setupPythonServiceExecObservable(service, 'jupyter', ['notebook', '--no-browser', /--notebook-dir=.*/, /--config=.*/, '--NotebookApp.iopub_data_rate_limit=10000000000.0'], [], notebookStdErr ? notebookStdErr : ['http://localhost:8888/?token=198']);
476476
}
477477

478478
function setupMissingNotebookPythonService(service: TypeMoq.IMock<IPythonExecutionService>) {
@@ -500,15 +500,15 @@ suite('Jupyter Execution', async () => {
500500
const getServerInfoPath = path.join(EXTENSION_ROOT_DIR, 'pythonFiles', 'datascience', 'getServerInfo.py');
501501
setupProcessServiceExec(service, workingPython.path, [getServerInfoPath], Promise.resolve({ stdout: 'failure to get server infos' }));
502502
setupProcessServiceExecObservable(service, workingPython.path, ['-m', 'jupyter', 'kernelspec', 'list'], [], []);
503-
setupProcessServiceExecObservable(service, workingPython.path, ['-m', 'jupyter', 'notebook', '--no-browser', /--notebook-dir=.*/, /.*/, '--NotebookApp.iopub_data_rate_limit=10000000000.0'], [], notebookStdErr ? notebookStdErr : ['http://localhost:8888/?token=198']);
503+
setupProcessServiceExecObservable(service, workingPython.path, ['-m', 'jupyter', 'notebook', '--no-browser', /--notebook-dir=.*/, /--config=.*/, '--NotebookApp.iopub_data_rate_limit=10000000000.0'], [], notebookStdErr ? notebookStdErr : ['http://localhost:8888/?token=198']);
504504
}
505505

506506
function setupMissingKernelProcessService(service: TypeMoq.IMock<IProcessService>, notebookStdErr?: string[]) {
507507
setupProcessServiceExec(service, missingKernelPython.path, ['-m', 'jupyter', 'kernelspec', 'list'], Promise.resolve({ stdout: `working ${path.dirname(workingKernelSpec)}` }));
508508
const getServerInfoPath = path.join(EXTENSION_ROOT_DIR, 'pythonFiles', 'datascience', 'getServerInfo.py');
509509
setupProcessServiceExec(service, missingKernelPython.path, [getServerInfoPath], Promise.resolve({ stdout: 'failure to get server infos' }));
510510
setupProcessServiceExecObservable(service, missingKernelPython.path, ['-m', 'jupyter', 'kernelspec', 'list'], [], []);
511-
setupProcessServiceExecObservable(service, missingKernelPython.path, ['-m', 'jupyter', 'notebook', '--no-browser', /--notebook-dir=.*/, /.*/, '--NotebookApp.iopub_data_rate_limit=10000000000.0'], [], notebookStdErr ? notebookStdErr : ['http://localhost:8888/?token=198']);
511+
setupProcessServiceExecObservable(service, missingKernelPython.path, ['-m', 'jupyter', 'notebook', '--no-browser', /--notebook-dir=.*/, /--config=.*/, '--NotebookApp.iopub_data_rate_limit=10000000000.0'], [], notebookStdErr ? notebookStdErr : ['http://localhost:8888/?token=198']);
512512
}
513513

514514
function setupPathProcessService(jupyterPath: string, service: TypeMoq.IMock<IProcessService>, notebookStdErr?: string[]) {
@@ -517,7 +517,7 @@ suite('Jupyter Execution', async () => {
517517
setupProcessServiceExec(service, jupyterPath, ['--version'], Promise.resolve({ stdout: '1.1.1.1' }));
518518
setupProcessServiceExec(service, jupyterPath, ['notebook', '--version'], Promise.resolve({ stdout: '1.1.1.1' }));
519519
setupProcessServiceExec(service, jupyterPath, ['kernelspec', '--version'], Promise.resolve({ stdout: '1.1.1.1' }));
520-
setupProcessServiceExecObservable(service, jupyterPath, ['notebook', '--no-browser', /--notebook-dir=.*/, /.*/, '--NotebookApp.iopub_data_rate_limit=10000000000.0'], [], notebookStdErr ? notebookStdErr : ['http://localhost:8888/?token=198']);
520+
setupProcessServiceExecObservable(service, jupyterPath, ['notebook', '--no-browser', /--notebook-dir=.*/, /--config=.*/, '--NotebookApp.iopub_data_rate_limit=10000000000.0'], [], notebookStdErr ? notebookStdErr : ['http://localhost:8888/?token=198']);
521521

522522
// WE also check for existence with just the key jupyter
523523
setupProcessServiceExec(service, 'jupyter', ['--version'], Promise.resolve({ stdout: '1.1.1.1' }));

0 commit comments

Comments
 (0)