Starting Fewer Jupyter Servers for Notebooks (#7744)

* Start fewer Jupyter servers

* Windows fix for drive casing

* PR Feedback

* Quick fix

* Fixing bug

* Ensure environment variables set 4 session startup

* test fix

* Dummy commit to update comment
This commit is contained in:
Chris LaFreniere
2019-10-20 21:38:58 -07:00
committed by GitHub
parent 0bfb1aab7e
commit e3ae5263c6
9 changed files with 118 additions and 43 deletions

View File

@@ -19,7 +19,7 @@ import { IPrompter, QuestionTypes, IQuestion } from '../prompts/question';
import { AppContext } from '../common/appContext'; import { AppContext } from '../common/appContext';
import { ApiWrapper } from '../common/apiWrapper'; import { ApiWrapper } from '../common/apiWrapper';
import { LocalJupyterServerManager } from './jupyterServerManager'; import { LocalJupyterServerManager, ServerInstanceFactory } from './jupyterServerManager';
import { NotebookCompletionItemProvider } from '../intellisense/completionItemProvider'; import { NotebookCompletionItemProvider } from '../intellisense/completionItemProvider';
import { JupyterNotebookProvider } from './jupyterNotebookProvider'; import { JupyterNotebookProvider } from './jupyterNotebookProvider';
import { ConfigurePythonDialog } from '../dialog/configurePythonDialog'; import { ConfigurePythonDialog } from '../dialog/configurePythonDialog';
@@ -31,6 +31,7 @@ let untitledCounter = 0;
export class JupyterController implements vscode.Disposable { export class JupyterController implements vscode.Disposable {
private _jupyterInstallation: JupyterServerInstallation; private _jupyterInstallation: JupyterServerInstallation;
private _notebookInstances: IServerInstance[] = []; private _notebookInstances: IServerInstance[] = [];
private _serverInstanceFactory: ServerInstanceFactory = new ServerInstanceFactory();
private outputChannel: vscode.OutputChannel; private outputChannel: vscode.OutputChannel;
private prompter: IPrompter; private prompter: IPrompter;
@@ -92,7 +93,8 @@ export class JupyterController implements vscode.Disposable {
documentPath: documentUri.fsPath, documentPath: documentUri.fsPath,
jupyterInstallation: this._jupyterInstallation, jupyterInstallation: this._jupyterInstallation,
extensionContext: this.extensionContext, extensionContext: this.extensionContext,
apiWrapper: this.apiWrapper apiWrapper: this.apiWrapper,
factory: this._serverInstanceFactory
})); }));
azdata.nb.registerNotebookProvider(notebookProvider); azdata.nb.registerNotebookProvider(notebookProvider);
return notebookProvider; return notebookProvider;

View File

@@ -3,16 +3,17 @@
* Licensed under the Source EULA. See License.txt in the project root for license information. * Licensed under the Source EULA. See License.txt in the project root for license information.
*--------------------------------------------------------------------------------------------*/ *--------------------------------------------------------------------------------------------*/
'use strict';
import { nb } from 'azdata'; import { nb } from 'azdata';
import * as vscode from 'vscode'; import * as vscode from 'vscode';
import * as nls from 'vscode-nls'; import * as nls from 'vscode-nls';
import * as path from 'path';
const localize = nls.loadMessageBundle(); const localize = nls.loadMessageBundle();
import * as constants from '../common/constants'; import * as constants from '../common/constants';
import * as utils from '../common/utils';
import { JupyterNotebookManager } from './jupyterNotebookManager'; import { JupyterNotebookManager } from './jupyterNotebookManager';
import { LocalJupyterServerManager } from './jupyterServerManager'; import { LocalJupyterServerManager } from './jupyterServerManager';
import { JupyterSessionManager } from './jupyterSessionManager';
export type ServerManagerFactory = (documentUri: vscode.Uri) => LocalJupyterServerManager; export type ServerManagerFactory = (documentUri: vscode.Uri) => LocalJupyterServerManager;
@@ -31,12 +32,16 @@ export class JupyterNotebookProvider implements nb.NotebookProvider {
} }
private doGetNotebookManager(notebookUri: vscode.Uri): nb.NotebookManager { private doGetNotebookManager(notebookUri: vscode.Uri): nb.NotebookManager {
let uriString = notebookUri.toString(); let baseFolder = this.transformToBaseFolder(notebookUri.fsPath.toString());
let manager = this.managerTracker.get(uriString); let manager = this.managerTracker.get(baseFolder);
if (!manager) { if (!manager) {
let serverManager = this.createServerManager(notebookUri); let baseFolderUri = vscode.Uri.file(baseFolder);
if (!baseFolderUri) {
baseFolderUri = notebookUri;
}
let serverManager = this.createServerManager(baseFolderUri);
manager = new JupyterNotebookManager(serverManager); manager = new JupyterNotebookManager(serverManager);
this.managerTracker.set(uriString, manager); this.managerTracker.set(baseFolder, manager);
} }
return manager; return manager;
} }
@@ -46,15 +51,52 @@ export class JupyterNotebookProvider implements nb.NotebookProvider {
// As this is a notification method, will skip throwing an error here // As this is a notification method, will skip throwing an error here
return; return;
} }
let uriString = notebookUri.toString(); let baseFolder = this.transformToBaseFolder(notebookUri.fsPath.toString());
let manager = this.managerTracker.get(uriString); let manager = this.managerTracker.get(baseFolder);
if (manager) { if (manager) {
this.managerTracker.delete(uriString); let sessionManager = (manager.sessionManager as JupyterSessionManager);
let session = sessionManager.listRunning().find(e => e.path === notebookUri.fsPath);
if (session) {
manager.sessionManager.shutdown(session.id);
}
if (sessionManager.listRunning().length === 0) {
const FiveMinutesInMs = 5 * 60 * 1000;
setTimeout(() => {
if (sessionManager.listRunning().length === 0) {
this.managerTracker.delete(baseFolder);
manager.dispose(); manager.dispose();
} }
}, FiveMinutesInMs);
}
}
} }
public get standardKernels(): nb.IStandardKernel[] { public get standardKernels(): nb.IStandardKernel[] {
return []; return [];
} }
private transformToBaseFolder(notebookPath: string): string {
let parsedPath = path.parse(notebookPath);
let userHome = utils.getUserHome();
let relativePathStrUserHome = path.relative(notebookPath, userHome);
if (notebookPath === '.' || relativePathStrUserHome.endsWith('..') || relativePathStrUserHome === '') {
// If you don't match the notebookPath's casing for drive letters,
// a 404 will result when trying to create a new session on Windows
if (notebookPath && userHome && notebookPath[0].toLowerCase() === userHome[0].toLowerCase()) {
userHome = notebookPath[0] + userHome.substr(1);
}
// If the user is using a system version of python, then
// '.' will try to create a notebook in a system directory.
// Since this will fail due to permissions issues, use the user's
// home folder instead.
return userHome;
} else {
let splitDirName: string[] = path.dirname(notebookPath).split(path.sep);
if (splitDirName.length > 1) {
return path.join(parsedPath.root, splitDirName[1]);
} else {
return userHome;
}
}
}
} }

View File

@@ -16,7 +16,7 @@ import { ApiWrapper } from '../common/apiWrapper';
import { JupyterServerInstallation } from './jupyterServerInstallation'; import { JupyterServerInstallation } from './jupyterServerInstallation';
import * as utils from '../common/utils'; import * as utils from '../common/utils';
import { IServerInstance } from './common'; import { IServerInstance } from './common';
import { PerNotebookServerInstance, IInstanceOptions } from './serverInstance'; import { PerFolderServerInstance, IInstanceOptions } from './serverInstance';
import { CommandContext } from '../common/constants'; import { CommandContext } from '../common/constants';
export interface IServerManagerOptions { export interface IServerManagerOptions {
@@ -30,11 +30,11 @@ export class LocalJupyterServerManager implements nb.ServerManager, vscode.Dispo
private _serverSettings: Partial<ServerConnection.ISettings>; private _serverSettings: Partial<ServerConnection.ISettings>;
private _onServerStarted = new vscode.EventEmitter<void>(); private _onServerStarted = new vscode.EventEmitter<void>();
private _instanceOptions: IInstanceOptions; private _instanceOptions: IInstanceOptions;
private apiWrapper: ApiWrapper; private _apiWrapper: ApiWrapper;
private jupyterServer: IServerInstance; private _jupyterServer: IServerInstance;
factory: ServerInstanceFactory; factory: ServerInstanceFactory;
constructor(private options: IServerManagerOptions) { constructor(private options: IServerManagerOptions) {
this.apiWrapper = options.apiWrapper || new ApiWrapper(); this._apiWrapper = options.apiWrapper || new ApiWrapper();
this.factory = options.factory || new ServerInstanceFactory(); this.factory = options.factory || new ServerInstanceFactory();
} }
@@ -43,7 +43,7 @@ export class LocalJupyterServerManager implements nb.ServerManager, vscode.Dispo
} }
public get isStarted(): boolean { public get isStarted(): boolean {
return !!this.jupyterServer; return !!this._jupyterServer;
} }
public get instanceOptions(): IInstanceOptions { public get instanceOptions(): IInstanceOptions {
@@ -56,11 +56,13 @@ export class LocalJupyterServerManager implements nb.ServerManager, vscode.Dispo
public async startServer(): Promise<void> { public async startServer(): Promise<void> {
try { try {
this.jupyterServer = await this.doStartServer(); if (!this._jupyterServer) {
this._jupyterServer = await this.doStartServer();
this.options.extensionContext.subscriptions.push(this); this.options.extensionContext.subscriptions.push(this);
let partialSettings = LocalJupyterServerManager.getLocalConnectionSettings(this.jupyterServer.uri); let partialSettings = LocalJupyterServerManager.getLocalConnectionSettings(this._jupyterServer.uri);
this._serverSettings = partialSettings; this._serverSettings = partialSettings;
this._onServerStarted.fire(); this._onServerStarted.fire();
}
} catch (error) { } catch (error) {
// this is caught and notified up the stack, no longer showing a message here // this is caught and notified up the stack, no longer showing a message here
@@ -71,13 +73,13 @@ export class LocalJupyterServerManager implements nb.ServerManager, vscode.Dispo
public dispose(): void { public dispose(): void {
this.stopServer().catch(err => { this.stopServer().catch(err => {
let msg = utils.getErrorMessage(err); let msg = utils.getErrorMessage(err);
this.apiWrapper.showErrorMessage(localize('shutdownError', 'Shutdown of Notebook server failed: {0}', msg)); this._apiWrapper.showErrorMessage(localize('shutdownError', 'Shutdown of Notebook server failed: {0}', msg));
}); });
} }
public async stopServer(): Promise<void> { public async stopServer(): Promise<void> {
if (this.jupyterServer) { if (this._jupyterServer) {
await this.jupyterServer.stop(); await this._jupyterServer.stop();
} }
} }
@@ -106,7 +108,7 @@ export class LocalJupyterServerManager implements nb.ServerManager, vscode.Dispo
let installation = this.options.jupyterInstallation; let installation = this.options.jupyterInstallation;
await installation.promptForPythonInstall(); await installation.promptForPythonInstall();
await installation.promptForPackageUpgrade(); await installation.promptForPackageUpgrade();
this.apiWrapper.setCommandContext(CommandContext.NotebookPythonInstalled, true); this._apiWrapper.setCommandContext(CommandContext.NotebookPythonInstalled, true);
// Calculate the path to use as the notebook-dir for Jupyter based on the path of the uri of the // Calculate the path to use as the notebook-dir for Jupyter based on the path of the uri of the
// notebook to open. This will be the workspace folder if the notebook uri is inside a workspace // notebook to open. This will be the workspace folder if the notebook uri is inside a workspace
@@ -118,9 +120,11 @@ export class LocalJupyterServerManager implements nb.ServerManager, vscode.Dispo
// /path2/nb2.ipynb // /path2/nb2.ipynb
// /path2/nb3.ipynb // /path2/nb3.ipynb
// ... will result in 2 notebook servers being started, one for /path1/ and one for /path2/ // ... will result in 2 notebook servers being started, one for /path1/ and one for /path2/
let notebookDir = this.apiWrapper.getWorkspacePathFromUri(vscode.Uri.file(this.documentPath)); let notebookDir = this._apiWrapper.getWorkspacePathFromUri(vscode.Uri.file(this.documentPath));
if (!notebookDir) { if (!notebookDir) {
let docDir = path.dirname(this.documentPath); let docDir;
// If a folder is passed in for documentPath, use the folder instead of calling dirname
docDir = path.extname(this.documentPath) ? path.dirname(this.documentPath) : this.documentPath;
if (docDir === '.') { if (docDir === '.') {
// If the user is using a system version of python, then // If the user is using a system version of python, then
// '.' will try to create a notebook in a system directory. // '.' will try to create a notebook in a system directory.
@@ -156,7 +160,7 @@ export class LocalJupyterServerManager implements nb.ServerManager, vscode.Dispo
export class ServerInstanceFactory { export class ServerInstanceFactory {
createInstance(options: IInstanceOptions): IServerInstance { createInstance(options: IInstanceOptions): IServerInstance {
return new PerNotebookServerInstance(options); return new PerFolderServerInstance(options);
} }
} }

View File

@@ -117,13 +117,14 @@ export class JupyterSessionManager implements nb.SessionManager {
return allKernels; return allKernels;
} }
public async startNew(options: nb.ISessionOptions): Promise<nb.ISession> { public async startNew(options: nb.ISessionOptions, skipSettingEnvironmentVars?: boolean): Promise<nb.ISession> {
if (!this._isReady) { if (!this._isReady) {
// no-op // no-op
return Promise.reject(new Error(localize('errorStartBeforeReady', 'Cannot start a session, the manager is not yet initialized'))); return Promise.reject(new Error(localize('errorStartBeforeReady', 'Cannot start a session, the manager is not yet initialized')));
} }
let sessionImpl = await this._sessionManager.startNew(options); let sessionImpl = await this._sessionManager.startNew(options);
let jupyterSession = new JupyterSession(sessionImpl); let jupyterSession = new JupyterSession(sessionImpl, skipSettingEnvironmentVars);
await jupyterSession.messagesComplete;
let index = JupyterSessionManager._sessions.findIndex(session => session.path === options.path); let index = JupyterSessionManager._sessions.findIndex(session => session.path === options.path);
if (index > -1) { if (index > -1) {
JupyterSessionManager._sessions.splice(index); JupyterSessionManager._sessions.splice(index);
@@ -167,8 +168,10 @@ export class JupyterSessionManager implements nb.SessionManager {
export class JupyterSession implements nb.ISession { export class JupyterSession implements nb.ISession {
private _kernel: nb.IKernel; private _kernel: nb.IKernel;
private _messagesComplete: Deferred<void> = new Deferred<void>();
constructor(private sessionImpl: Session.ISession) { constructor(private sessionImpl: Session.ISession, skipSettingEnvironmentVars?: boolean) {
this.setEnvironmentVars(skipSettingEnvironmentVars);
} }
public get canChangeKernels(): boolean { public get canChangeKernels(): boolean {
@@ -205,6 +208,11 @@ export class JupyterSession implements nb.ISession {
return this._kernel; return this._kernel;
} }
// Sent when startup messages have been sent
public get messagesComplete(): Promise<void> {
return this._messagesComplete.promise;
}
public async changeKernel(kernelInfo: nb.IKernelSpec): Promise<nb.IKernel> { public async changeKernel(kernelInfo: nb.IKernelSpec): Promise<nb.IKernel> {
// For now, Jupyter implementation handles disposal etc. so we can just // For now, Jupyter implementation handles disposal etc. so we can just
// null out our kernel and let the changeKernel call handle this // null out our kernel and let the changeKernel call handle this
@@ -321,6 +329,22 @@ export class JupyterSession implements nb.ISession {
} }
return endpoints.find(ep => ep.serviceName.toLowerCase() === serviceName.toLowerCase()); return endpoints.find(ep => ep.serviceName.toLowerCase() === serviceName.toLowerCase());
} }
private async setEnvironmentVars(skip: boolean = false): Promise<void> {
if (!skip && this.sessionImpl) {
let allCode: string = '';
for (let i = 0; i < Object.keys(process.env).length; i++) {
let key = Object.keys(process.env)[i];
// Jupyter doesn't seem to alow for setting multiple variables at once, so doing it with multiple commands
allCode += `%set_env ${key}=${process.env[key]}${EOL}`;
}
let future = this.sessionImpl.kernel.requestExecute({
code: allCode
}, true);
await future.done;
}
this._messagesComplete.resolve();
}
} }
interface ICredentials { interface ICredentials {

View File

@@ -99,7 +99,7 @@ export class ServerInstanceUtils {
} }
} }
export class PerNotebookServerInstance implements IServerInstance { export class PerFolderServerInstance implements IServerInstance {
/** /**
* Root of the jupyter directory structure. Config and data roots will be * Root of the jupyter directory structure. Config and data roots will be
@@ -123,6 +123,7 @@ export class PerNotebookServerInstance implements IServerInstance {
private _port: string; private _port: string;
private _uri: vscode.Uri; private _uri: vscode.Uri;
private _isStarted: boolean = false; private _isStarted: boolean = false;
private _isStopping: boolean = false;
private utils: ServerInstanceUtils; private utils: ServerInstanceUtils;
private childProcess: ChildProcess; private childProcess: ChildProcess;
private errorHandler: ErrorHandler = new ErrorHandler(); private errorHandler: ErrorHandler = new ErrorHandler();
@@ -132,7 +133,7 @@ export class PerNotebookServerInstance implements IServerInstance {
} }
public get isStarted(): boolean { public get isStarted(): boolean {
return this._isStarted; return this._isStarted && !this._isStopping;
} }
public get port(): string { public get port(): string {
@@ -153,13 +154,14 @@ export class PerNotebookServerInstance implements IServerInstance {
public async stop(): Promise<void> { public async stop(): Promise<void> {
try { try {
this._isStopping = true;
if (this.baseDir) { if (this.baseDir) {
let exists = await this.utils.pathExists(this.baseDir); let exists = await this.utils.pathExists(this.baseDir);
if (exists) { if (exists) {
await this.utils.removeDir(this.baseDir); await this.utils.removeDir(this.baseDir);
} }
} }
if (this.isStarted) { if (this._isStarted) {
let install = this.options.install; let install = this.options.install;
let stopCommand = `"${install.pythonExecutable}" -m jupyter notebook stop ${this._port}`; let stopCommand = `"${install.pythonExecutable}" -m jupyter notebook stop ${this._port}`;
await this.utils.executeBufferedCommand(stopCommand, install.execOptions, install.outputChannel); await this.utils.executeBufferedCommand(stopCommand, install.execOptions, install.outputChannel);
@@ -171,6 +173,7 @@ export class PerNotebookServerInstance implements IServerInstance {
this._isStarted = false; this._isStarted = false;
this.utils.ensureProcessEnded(this.childProcess); this.utils.ensureProcessEnded(this.childProcess);
this.handleConnectionClosed(); this.handleConnectionClosed();
} }
} }

View File

@@ -11,7 +11,7 @@ import 'mocha';
import { JupyterServerInstallation } from '../../jupyter/jupyterServerInstallation'; import { JupyterServerInstallation } from '../../jupyter/jupyterServerInstallation';
import { ApiWrapper } from '../..//common/apiWrapper'; import { ApiWrapper } from '../..//common/apiWrapper';
import { PerNotebookServerInstance, ServerInstanceUtils } from '../../jupyter/serverInstance'; import { PerFolderServerInstance, ServerInstanceUtils } from '../../jupyter/serverInstance';
import { MockOutputChannel } from '../common/stubs'; import { MockOutputChannel } from '../common/stubs';
import * as testUtils from '../common/testUtils'; import * as testUtils from '../common/testUtils';
import { LocalJupyterServerManager } from '../../jupyter/jupyterServerManager'; import { LocalJupyterServerManager } from '../../jupyter/jupyterServerManager';
@@ -27,7 +27,7 @@ describe('Jupyter server instance', function (): void {
let mockOutputChannel: TypeMoq.IMock<MockOutputChannel>; let mockOutputChannel: TypeMoq.IMock<MockOutputChannel>;
let mockApiWrapper: TypeMoq.IMock<ApiWrapper>; let mockApiWrapper: TypeMoq.IMock<ApiWrapper>;
let mockUtils: TypeMoq.IMock<ServerInstanceUtils>; let mockUtils: TypeMoq.IMock<ServerInstanceUtils>;
let serverInstance: PerNotebookServerInstance; let serverInstance: PerFolderServerInstance;
beforeEach(() => { beforeEach(() => {
mockApiWrapper = TypeMoq.Mock.ofType(ApiWrapper); mockApiWrapper = TypeMoq.Mock.ofType(ApiWrapper);
@@ -40,7 +40,7 @@ describe('Jupyter server instance', function (): void {
mockInstall.object.execOptions = { env: Object.assign({}, process.env) }; mockInstall.object.execOptions = { env: Object.assign({}, process.env) };
mockUtils = TypeMoq.Mock.ofType(ServerInstanceUtils); mockUtils = TypeMoq.Mock.ofType(ServerInstanceUtils);
mockUtils.setup(u => u.ensureProcessEnded(TypeMoq.It.isAny())).returns(() => undefined); mockUtils.setup(u => u.ensureProcessEnded(TypeMoq.It.isAny())).returns(() => undefined);
serverInstance = new PerNotebookServerInstance({ serverInstance = new PerFolderServerInstance({
documentPath: expectedPath, documentPath: expectedPath,
install: mockInstall.object install: mockInstall.object
}, mockUtils.object); }, mockUtils.object);

View File

@@ -85,7 +85,7 @@ describe('Jupyter Session Manager', function (): void {
mockJupyterManager.setup(m => m.startNew(TypeMoq.It.isAny())).returns(() => Promise.resolve(expectedSessionInfo)); mockJupyterManager.setup(m => m.startNew(TypeMoq.It.isAny())).returns(() => Promise.resolve(expectedSessionInfo));
// When I call startSession // When I call startSession
let session = await sessionManager.startNew(sessionOptions); let session = await sessionManager.startNew(sessionOptions, true);
// Then I expect the parameters passed to be correct // Then I expect the parameters passed to be correct
should(session.path).equal(sessionOptions.path); should(session.path).equal(sessionOptions.path);
should(session.canChangeKernels).be.true(); should(session.canChangeKernels).be.true();
@@ -147,7 +147,7 @@ describe('Jupyter Session', function (): void {
kernel = session.kernel; kernel = session.kernel;
// Then I expect it to have the ID, and only be called once // Then I expect it to have the ID, and only be called once
should(kernel.id).equal('id'); should(kernel.id).equal('id');
mockJupyterSession.verify(s => s.kernel, TypeMoq.Times.once()); mockJupyterSession.verify(s => s.kernel, TypeMoq.Times.exactly(2));
}); });
it('should send name in changeKernel request', async function (): Promise<void> { it('should send name in changeKernel request', async function (): Promise<void> {

View File

@@ -34,7 +34,7 @@ export class ExtHostNotebook implements ExtHostNotebookShape {
let adapter = this.findManagerForUri(uriString); let adapter = this.findManagerForUri(uriString);
if (!adapter) { if (!adapter) {
adapter = await this._withProvider(providerHandle, (provider) => { adapter = await this._withProvider(providerHandle, (provider) => {
return this.createManager(provider, uri); return this.getOrCreateManager(provider, uri);
}); });
} }
@@ -242,7 +242,7 @@ export class ExtHostNotebook implements ExtHostNotebookShape {
return undefined; return undefined;
} }
private async createManager(provider: azdata.nb.NotebookProvider, notebookUri: URI): Promise<NotebookManagerAdapter> { private async getOrCreateManager(provider: azdata.nb.NotebookProvider, notebookUri: URI): Promise<NotebookManagerAdapter> {
let manager = await provider.getNotebookManager(notebookUri); let manager = await provider.getNotebookManager(notebookUri);
let uriString = notebookUri.toString(); let uriString = notebookUri.toString();
let adapter = new NotebookManagerAdapter(provider, manager, uriString); let adapter = new NotebookManagerAdapter(provider, manager, uriString);