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

extensions/notebook/src/jupyter/jupyterController.ts

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

extensions/notebook/src/jupyter/jupyterNotebookManager.ts

@@ -53,4 +53,4 @@ export class JupyterNotebookManager implements nb.NotebookManager, vscode.Dispos
 			this._serverManager.stopServer().then(success => undefined, error => this.apiWrapper.showErrorMessage(error));
 		}
 	}
-}
+}

extensions/notebook/src/jupyter/jupyterNotebookProvider.ts

@@ -3,16 +3,17 @@
  *  Licensed under the Source EULA. See License.txt in the project root for license information.
  *--------------------------------------------------------------------------------------------*/
 
-'use strict';
-
 import { nb } from 'azdata';
 import * as vscode from 'vscode';
 import * as nls from 'vscode-nls';
+import * as path from 'path';
 const localize = nls.loadMessageBundle();
 
 import * as constants from '../common/constants';
+import * as utils from '../common/utils';
 import { JupyterNotebookManager } from './jupyterNotebookManager';
 import { LocalJupyterServerManager } from './jupyterServerManager';
+import { JupyterSessionManager } from './jupyterSessionManager';
 
 export type ServerManagerFactory = (documentUri: vscode.Uri) => LocalJupyterServerManager;
 
@@ -31,12 +32,16 @@ export class JupyterNotebookProvider implements nb.NotebookProvider {
 	}
 
 	private doGetNotebookManager(notebookUri: vscode.Uri): nb.NotebookManager {
-		let uriString = notebookUri.toString();
-		let manager = this.managerTracker.get(uriString);
+		let baseFolder = this.transformToBaseFolder(notebookUri.fsPath.toString());
+		let manager = this.managerTracker.get(baseFolder);
 		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);
-			this.managerTracker.set(uriString, manager);
+			this.managerTracker.set(baseFolder, 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
 			return;
 		}
-		let uriString = notebookUri.toString();
-		let manager = this.managerTracker.get(uriString);
+		let baseFolder = this.transformToBaseFolder(notebookUri.fsPath.toString());
+		let manager = this.managerTracker.get(baseFolder);
 		if (manager) {
-			this.managerTracker.delete(uriString);
-			manager.dispose();
+			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();
+					}
+				}, FiveMinutesInMs);
+			}
 		}
 	}
 
 	public get standardKernels(): nb.IStandardKernel[] {
 		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;
+			}
+		}
+	}
 }

extensions/notebook/src/jupyter/jupyterServerManager.ts

@@ -16,7 +16,7 @@ import { ApiWrapper } from '../common/apiWrapper';
 import { JupyterServerInstallation } from './jupyterServerInstallation';
 import * as utils from '../common/utils';
 import { IServerInstance } from './common';
-import { PerNotebookServerInstance, IInstanceOptions } from './serverInstance';
+import { PerFolderServerInstance, IInstanceOptions } from './serverInstance';
 import { CommandContext } from '../common/constants';
 
 export interface IServerManagerOptions {
@@ -30,11 +30,11 @@ export class LocalJupyterServerManager implements nb.ServerManager, vscode.Dispo
 	private _serverSettings: Partial<ServerConnection.ISettings>;
 	private _onServerStarted = new vscode.EventEmitter<void>();
 	private _instanceOptions: IInstanceOptions;
-	private apiWrapper: ApiWrapper;
-	private jupyterServer: IServerInstance;
+	private _apiWrapper: ApiWrapper;
+	private _jupyterServer: IServerInstance;
 	factory: ServerInstanceFactory;
 	constructor(private options: IServerManagerOptions) {
-		this.apiWrapper = options.apiWrapper || new ApiWrapper();
+		this._apiWrapper = options.apiWrapper || new ApiWrapper();
 		this.factory = options.factory || new ServerInstanceFactory();
 	}
 
@@ -43,7 +43,7 @@ export class LocalJupyterServerManager implements nb.ServerManager, vscode.Dispo
 	}
 
 	public get isStarted(): boolean {
-		return !!this.jupyterServer;
+		return !!this._jupyterServer;
 	}
 
 	public get instanceOptions(): IInstanceOptions {
@@ -56,11 +56,13 @@ export class LocalJupyterServerManager implements nb.ServerManager, vscode.Dispo
 
 	public async startServer(): Promise<void> {
 		try {
-			this.jupyterServer = await this.doStartServer();
-			this.options.extensionContext.subscriptions.push(this);
-			let partialSettings = LocalJupyterServerManager.getLocalConnectionSettings(this.jupyterServer.uri);
-			this._serverSettings = partialSettings;
-			this._onServerStarted.fire();
+			if (!this._jupyterServer) {
+				this._jupyterServer = await this.doStartServer();
+				this.options.extensionContext.subscriptions.push(this);
+				let partialSettings = LocalJupyterServerManager.getLocalConnectionSettings(this._jupyterServer.uri);
+				this._serverSettings = partialSettings;
+				this._onServerStarted.fire();
+			}
 
 		} catch (error) {
 			// 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 {
 		this.stopServer().catch(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> {
-		if (this.jupyterServer) {
-			await this.jupyterServer.stop();
+		if (this._jupyterServer) {
+			await this._jupyterServer.stop();
 		}
 	}
 
@@ -106,7 +108,7 @@ export class LocalJupyterServerManager implements nb.ServerManager, vscode.Dispo
 		let installation = this.options.jupyterInstallation;
 		await installation.promptForPythonInstall();
 		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
 		// 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/nb3.ipynb
 		// ... 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) {
-			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 the user is using a system version of python, then
 				// '.' 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 {
 
 	createInstance(options: IInstanceOptions): IServerInstance {
-		return new PerNotebookServerInstance(options);
+		return new PerFolderServerInstance(options);
 	}
 }
 

extensions/notebook/src/jupyter/jupyterSessionManager.ts

@@ -117,13 +117,14 @@ export class JupyterSessionManager implements nb.SessionManager {
 		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) {
 			// no-op
 			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 jupyterSession = new JupyterSession(sessionImpl);
+		let jupyterSession = new JupyterSession(sessionImpl, skipSettingEnvironmentVars);
+		await jupyterSession.messagesComplete;
 		let index = JupyterSessionManager._sessions.findIndex(session => session.path === options.path);
 		if (index > -1) {
 			JupyterSessionManager._sessions.splice(index);
@@ -167,8 +168,10 @@ export class JupyterSessionManager implements nb.SessionManager {
 
 export class JupyterSession implements nb.ISession {
 	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 {
@@ -205,6 +208,11 @@ export class JupyterSession implements nb.ISession {
 		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> {
 		// For now, Jupyter implementation handles disposal etc. so we can just
 		// 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());
 	}
+
+	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 {

extensions/notebook/src/jupyter/serverInstance.ts

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

extensions/notebook/src/test/model/serverInstance.test.ts

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

extensions/notebook/src/test/model/sessionManager.test.ts

@@ -85,7 +85,7 @@ describe('Jupyter Session Manager', function (): void {
 		mockJupyterManager.setup(m => m.startNew(TypeMoq.It.isAny())).returns(() => Promise.resolve(expectedSessionInfo));
 
 		// 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
 		should(session.path).equal(sessionOptions.path);
 		should(session.canChangeKernels).be.true();
@@ -147,7 +147,7 @@ describe('Jupyter Session', function (): void {
 		kernel = session.kernel;
 		// Then I expect it to have the ID, and only be called once
 		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> {

src/sql/workbench/api/common/extHostNotebook.ts

@@ -34,7 +34,7 @@ export class ExtHostNotebook implements ExtHostNotebookShape {
 		let adapter = this.findManagerForUri(uriString);
 		if (!adapter) {
 			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;
 	}
 
-	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 uriString = notebookUri.toString();
 		let adapter = new NotebookManagerAdapter(provider, manager, uriString);