Update Configure Python dialog to allow packages to be installed for specific kernels. (#10286)

extensions/notebook/src/jupyter/jupyterController.ts

@@ -22,7 +22,7 @@ import { ApiWrapper } from '../common/apiWrapper';
 import { LocalJupyterServerManager, ServerInstanceFactory } from './jupyterServerManager';
 import { NotebookCompletionItemProvider } from '../intellisense/completionItemProvider';
 import { JupyterNotebookProvider } from './jupyterNotebookProvider';
-import { ConfigurePythonDialog } from '../dialog/configurePythonDialog';
+import { ConfigurePythonWizard } from '../dialog/configurePython/configurePythonWizard';
 import CodeAdapter from '../prompts/adapter';
 import { ManagePackagesDialog } from '../dialog/managePackages/managePackagesDialog';
 import { IPackageManageProvider } from '../types';
@@ -30,6 +30,7 @@ import { LocalPipPackageManageProvider } from './localPipPackageManageProvider';
 import { LocalCondaPackageManageProvider } from './localCondaPackageManageProvider';
 import { ManagePackagesDialogModel, ManagePackageDialogOptions } from '../dialog/managePackages/managePackagesDialogModel';
 import { PiPyClient } from './pipyClient';
+import { ConfigurePythonDialog } from '../dialog/configurePython/configurePythonDialog';
 
 let untitledCounter = 0;
 
@@ -250,10 +251,21 @@ export class JupyterController implements vscode.Disposable {
 	}
 
 	public doConfigurePython(jupyterInstaller: JupyterServerInstallation): void {
-		let pythonDialog = new ConfigurePythonDialog(this.apiWrapper, jupyterInstaller);
-		pythonDialog.showDialog().catch((err: any) => {
-			this.apiWrapper.showErrorMessage(utils.getErrorMessage(err));
-		});
+		let enablePreviewFeatures = this.apiWrapper.getConfiguration('workbench').get('enablePreviewFeatures');
+		if (enablePreviewFeatures) {
+			let pythonWizard = new ConfigurePythonWizard(this.apiWrapper, jupyterInstaller);
+			pythonWizard.start().catch((err: any) => {
+				this.apiWrapper.showErrorMessage(utils.getErrorMessage(err));
+			});
+			pythonWizard.setupComplete.catch((err: any) => {
+				this.apiWrapper.showErrorMessage(utils.getErrorMessage(err));
+			});
+		} else {
+			let pythonDialog = new ConfigurePythonDialog(this.apiWrapper, jupyterInstaller);
+			pythonDialog.showDialog().catch((err: any) => {
+				this.apiWrapper.showErrorMessage(utils.getErrorMessage(err));
+			});
+		}
 	}
 
 	public get jupyterInstallation() {

extensions/notebook/src/jupyter/jupyterNotebookManager.ts

@@ -20,8 +20,8 @@ export class JupyterNotebookManager implements nb.NotebookManager, vscode.Dispos
 		this._sessionManager = sessionManager || new JupyterSessionManager(pythonEnvVarPath);
 		this._serverManager.onServerStarted(() => {
 			this.setServerSettings(this._serverManager.serverSettings);
+			this._sessionManager.installation = this._serverManager.instanceOptions.install;
 		});
-
 	}
 	public get contentManager(): nb.ContentManager {
 		return undefined;

extensions/notebook/src/jupyter/jupyterServerInstallation.ts

@@ -17,9 +17,10 @@ import * as constants from '../common/constants';
 import * as utils from '../common/utils';
 import { OutputChannel, ConfigurationTarget, window } from 'vscode';
 import { Deferred } from '../common/promise';
-import { ConfigurePythonDialog } from '../dialog/configurePythonDialog';
+import { ConfigurePythonWizard } from '../dialog/configurePython/configurePythonWizard';
 import { IPrompter, IQuestion, QuestionTypes } from '../prompts/question';
 import CodeAdapter from '../prompts/adapter';
+import { ConfigurePythonDialog } from '../dialog/configurePython/configurePythonDialog';
 
 const localize = nls.loadMessageBundle();
 const msgInstallPkgProgress = localize('msgInstallPkgProgress', "Notebook dependencies installation is in progress");
@@ -39,10 +40,15 @@ function msgDependenciesInstallationFailed(errorMessage: string): string { retur
 function msgDownloadPython(platform: string, pythonDownloadUrl: string): string { return localize('msgDownloadPython', "Downloading local python for platform: {0} to {1}", platform, pythonDownloadUrl); }
 function msgPackageRetrievalFailed(errorMessage: string): string { return localize('msgPackageRetrievalFailed', "Encountered an error when trying to retrieve list of installed packages: {0}", errorMessage); }
 
+export interface PythonInstallSettings {
+	installPath: string;
+	existingPython: boolean;
+	specificPackages?: PythonPkgDetails[];
+}
 export interface IJupyterServerInstallation {
 	installCondaPackages(packages: PythonPkgDetails[], useMinVersion: boolean): Promise<void>;
 	configurePackagePaths(): Promise<void>;
-	startInstallProcess(forceInstall: boolean, installSettings?: { installPath: string, existingPython: boolean }): Promise<void>;
+	startInstallProcess(forceInstall: boolean, installSettings?: PythonInstallSettings): Promise<void>;
 	getInstalledPipPackages(): Promise<PythonPkgDetails[]>;
 	getInstalledCondaPackages(): Promise<PythonPkgDetails[]>;
 	uninstallCondaPackages(packages: PythonPkgDetails[]): Promise<void>;
@@ -66,7 +72,6 @@ export class JupyterServerInstallation implements IJupyterServerInstallation {
 
 	private _pythonInstallationPath: string;
 	private _pythonExecutable: string;
-	private _pythonPackageDir: string;
 	private _usingExistingPython: boolean;
 	private _usingConda: boolean;
 
@@ -104,6 +109,8 @@ export class JupyterServerInstallation implements IJupyterServerInstallation {
 	private readonly _expectedCondaPipPackages = this._commonPipPackages;
 	private readonly _expectedCondaPackages: PythonPkgDetails[];
 
+	private _kernelSetupCache: Map<string, boolean>;
+
 	constructor(extensionPath: string, outputChannel: OutputChannel, apiWrapper: ApiWrapper, pythonInstallationPath?: string) {
 		this.extensionPath = extensionPath;
 		this.outputChannel = outputChannel;
@@ -120,9 +127,11 @@ export class JupyterServerInstallation implements IJupyterServerInstallation {
 		} else {
 			this._expectedCondaPackages = this._commonPackages;
 		}
+
+		this._kernelSetupCache = new Map<string, boolean>();
 	}
 
-	private async installDependencies(backgroundOperation: azdata.BackgroundOperation, forceInstall: boolean): Promise<void> {
+	private async installDependencies(backgroundOperation: azdata.BackgroundOperation, forceInstall: boolean, specificPackages?: PythonPkgDetails[]): Promise<void> {
 		if (!(await utils.exists(this._pythonExecutable)) || forceInstall || this._usingExistingPython) {
 			window.showInformationMessage(msgInstallPkgStart);
 
@@ -132,12 +141,7 @@ export class JupyterServerInstallation implements IJupyterServerInstallation {
 
 			try {
 				await this.installPythonPackage(backgroundOperation, this._usingExistingPython, this._pythonInstallationPath, this.outputChannel);
-
-				if (this._usingExistingPython) {
-					await this.upgradePythonPackages(false, forceInstall);
-				} else {
-					await this.installOfflinePipDependencies();
-				}
+				await this.upgradePythonPackages(false, forceInstall, specificPackages);
 			} catch (err) {
 				this.outputChannel.appendLine(msgDependenciesInstallationFailed(utils.getErrorMessage(err)));
 				throw err;
@@ -282,19 +286,13 @@ export class JupyterServerInstallation implements IJupyterServerInstallation {
 			? this._pythonInstallationPath
 			: path.join(this._pythonInstallationPath, constants.pythonBundleVersion);
 
-		if (this._usingExistingPython) {
-			this._pythonPackageDir = undefined;
-		} else {
-			this._pythonPackageDir = path.join(pythonSourcePath, 'offlinePackages');
-		}
-
 		// Update python paths and properties to reference user's local python.
 		let pythonBinPathSuffix = process.platform === constants.winPlatform ? '' : 'bin';
 
 		this._pythonExecutable = JupyterServerInstallation.getPythonExePath(this._pythonInstallationPath, this._usingExistingPython);
 		this.pythonBinPath = path.join(pythonSourcePath, pythonBinPathSuffix);
 
-		this._usingConda = this.checkCondaExists();
+		this._usingConda = this.isCondaInstalled();
 
 		// Store paths to python libraries required to run jupyter.
 		this.pythonEnvVarPath = process.env['PATH'];
@@ -360,7 +358,7 @@ export class JupyterServerInstallation implements IJupyterServerInstallation {
 	 * @param installSettings Optional parameter that specifies where to install python, and whether the install targets an existing python install.
 	 * The previous python path (or the default) is used if a new path is not specified.
 	 */
-	public async startInstallProcess(forceInstall: boolean, installSettings?: { installPath: string, existingPython: boolean }): Promise<void> {
+	public async startInstallProcess(forceInstall: boolean, installSettings?: PythonInstallSettings): Promise<void> {
 		let isPythonRunning: boolean;
 		if (installSettings) {
 			isPythonRunning = await this.isPythonRunning(installSettings.installPath, installSettings.existingPython);
@@ -399,7 +397,7 @@ export class JupyterServerInstallation implements IJupyterServerInstallation {
 				description: msgTaskName,
 				isCancelable: false,
 				operation: op => {
-					this.installDependencies(op, forceInstall)
+					this.installDependencies(op, forceInstall, installSettings?.specificPackages)
 						.then(async () => {
 							await updateConfig();
 							this._installCompletion.resolve();
@@ -427,28 +425,45 @@ export class JupyterServerInstallation implements IJupyterServerInstallation {
 	/**
 	 * Opens a dialog for configuring the installation path for the Notebook Python dependencies.
 	 */
-	public async promptForPythonInstall(): Promise<void> {
+	public async promptForPythonInstall(kernelDisplayName: string): Promise<void> {
 		if (!JupyterServerInstallation.isPythonInstalled(this.apiWrapper)) {
-			let pythonDialog = new ConfigurePythonDialog(this.apiWrapper, this);
-			return pythonDialog.showDialog(true);
+			let enablePreviewFeatures = this.apiWrapper.getConfiguration('workbench').get('enablePreviewFeatures');
+			if (enablePreviewFeatures) {
+				let pythonWizard = new ConfigurePythonWizard(this.apiWrapper, this);
+				await pythonWizard.start(kernelDisplayName, true);
+				return pythonWizard.setupComplete;
+			} else {
+				let pythonDialog = new ConfigurePythonDialog(this.apiWrapper, this);
+				return pythonDialog.showDialog(true);
+			}
 		}
 	}
 
 	/**
 	 * Prompts user to upgrade certain python packages if they're below the minimum expected version.
 	 */
-	public async promptForPackageUpgrade(): Promise<void> {
+	public async promptForPackageUpgrade(kernelName: string): Promise<void> {
 		if (this._installInProgress) {
 			this.apiWrapper.showInfoMessage(msgWaitingForInstall);
 			return this._installCompletion.promise;
 		}
 
+		let requiredPackages: PythonPkgDetails[];
+		let enablePreviewFeatures = this.apiWrapper.getConfiguration('workbench').get('enablePreviewFeatures');
+		if (enablePreviewFeatures) {
+			if (this._kernelSetupCache.get(kernelName)) {
+				return;
+			}
+			requiredPackages = JupyterServerInstallation.getRequiredPackagesForKernel(kernelName);
+		}
+
 		this._installInProgress = true;
 		this._installCompletion = new Deferred<void>();
-		this.upgradePythonPackages(true, false)
+		this.upgradePythonPackages(true, false, requiredPackages)
 			.then(() => {
 				this._installCompletion.resolve();
 				this._installInProgress = false;
+				this._kernelSetupCache.set(kernelName, true);
 			})
 			.catch(err => {
 				let errorMsg = msgDependenciesInstallationFailed(utils.getErrorMessage(err));
@@ -458,10 +473,14 @@ export class JupyterServerInstallation implements IJupyterServerInstallation {
 		return this._installCompletion.promise;
 	}
 
-	private async upgradePythonPackages(promptForUpgrade: boolean, forceInstall: boolean): Promise<void> {
+	private async upgradePythonPackages(promptForUpgrade: boolean, forceInstall: boolean, specificPackages?: PythonPkgDetails[]): Promise<void> {
 		let expectedCondaPackages: PythonPkgDetails[];
 		let expectedPipPackages: PythonPkgDetails[];
-		if (this._usingConda) {
+		if (specificPackages) {
+			// Always install generic packages with pip, since conda may not have them.
+			expectedCondaPackages = [];
+			expectedPipPackages = specificPackages;
+		} else if (this._usingConda) {
 			expectedCondaPackages = this._expectedCondaPackages;
 			expectedPipPackages = this._expectedCondaPipPackages;
 		} else {
@@ -510,9 +529,12 @@ export class JupyterServerInstallation implements IJupyterServerInstallation {
 			if (promptForUpgrade) {
 				doUpgrade = await this._prompter.promptSingle<boolean>(<IQuestion>{
 					type: QuestionTypes.confirm,
-					message: localize('confirmPackageUpgrade', "Some installed python packages need to be upgraded. Would you like to upgrade them now?"),
+					message: localize('confirmPackageUpgrade', "Some required python packages need to be installed. Would you like to install them now?"),
 					default: true
 				});
+				if (!doUpgrade) {
+					throw new Error(localize('configurePython.packageInstallDeclined', "Package installation was declined."));
+				}
 			} else {
 				doUpgrade = true;
 			}
@@ -563,9 +585,17 @@ export class JupyterServerInstallation implements IJupyterServerInstallation {
 		}
 	}
 
-	public async getInstalledPipPackages(): Promise<PythonPkgDetails[]> {
+	public async getInstalledPipPackages(pythonExePath?: string): Promise<PythonPkgDetails[]> {
 		try {
-			let cmd = `"${this.pythonExecutable}" -m pip list --format=json`;
+			if (pythonExePath) {
+				if (!fs.existsSync(pythonExePath)) {
+					return [];
+				}
+			} else if (!JupyterServerInstallation.isPythonInstalled(this.apiWrapper)) {
+				return [];
+			}
+
+			let cmd = `"${pythonExePath ?? this.pythonExecutable}" -m pip list --format=json`;
 			let packagesInfo = await this.executeBufferedCommand(cmd);
 			let packagesResult: PythonPkgDetails[] = [];
 			if (packagesInfo) {
@@ -598,6 +628,10 @@ export class JupyterServerInstallation implements IJupyterServerInstallation {
 
 	public async getInstalledCondaPackages(): Promise<PythonPkgDetails[]> {
 		try {
+			if (!this.isCondaInstalled()) {
+				return [];
+			}
+
 			let condaExe = this.getCondaExePath();
 			let cmd = `"${condaExe}" list --json`;
 			let packagesInfo = await this.executeBufferedCommand(cmd);
@@ -637,32 +671,6 @@ export class JupyterServerInstallation implements IJupyterServerInstallation {
 		return this.executeStreamedCommand(cmd);
 	}
 
-	private async installOfflinePipDependencies(): Promise<void> {
-		// Skip this step if using existing python, since this is for our provided package
-		if (!this._usingExistingPython && process.platform === constants.winPlatform) {
-			this.outputChannel.show(true);
-			this.outputChannel.appendLine(localize('msgInstallStart', "Installing required packages to run Notebooks..."));
-
-			let requirements = path.join(this._pythonPackageDir, 'requirements.txt');
-			let installJupyterCommand = `"${this._pythonExecutable}" -m pip install --no-index -r "${requirements}" --find-links "${this._pythonPackageDir}" --no-warn-script-location`;
-			await this.executeStreamedCommand(installJupyterCommand);
-
-			// Force reinstall pip to update shebangs in pip*.exe files
-			installJupyterCommand = `"${this._pythonExecutable}" -m pip install --force-reinstall --no-index pip --find-links "${this._pythonPackageDir}" --no-warn-script-location`;
-			await this.executeStreamedCommand(installJupyterCommand);
-
-			fs.remove(this._pythonPackageDir, (err: Error) => {
-				if (err) {
-					this.outputChannel.appendLine(err.message);
-				}
-			});
-
-			this.outputChannel.appendLine(localize('msgJupyterInstallDone', "... Jupyter installation complete."));
-		} else {
-			return Promise.resolve();
-		}
-	}
-
 	public async executeStreamedCommand(command: string): Promise<void> {
 		await utils.executeStreamedCommand(command, { env: this.execOptions.env }, this.outputChannel);
 	}
@@ -691,7 +699,7 @@ export class JupyterServerInstallation implements IJupyterServerInstallation {
 		return this._usingConda;
 	}
 
-	private checkCondaExists(): boolean {
+	private isCondaInstalled(): boolean {
 		if (!this._usingExistingPython) {
 			return false;
 		}
@@ -800,6 +808,55 @@ export class JupyterServerInstallation implements IJupyterServerInstallation {
 
 		return undefined;
 	}
+
+	public static getRequiredPackagesForKernel(kernelName: string): PythonPkgDetails[] {
+		let packages = [{
+			name: 'jupyter',
+			version: '1.0.0'
+		}];
+		switch (kernelName) {
+			case constants.python3DisplayName:
+				break;
+			case constants.pysparkDisplayName:
+			case constants.sparkScalaDisplayName:
+			case constants.sparkRDisplayName:
+				packages.push({
+					name: 'sparkmagic',
+					version: '0.12.9'
+				}, {
+					name: 'pandas',
+					version: '0.24.2'
+				}, {
+					name: 'prose-codeaccelerator',
+					version: '1.3.0'
+				});
+				break;
+			case constants.powershellDisplayName:
+				packages.push({
+					name: 'powershell-kernel',
+					version: '0.1.3'
+				});
+				break;
+			case constants.allKernelsName:
+				packages.push({
+					name: 'sparkmagic',
+					version: '0.12.9'
+				}, {
+					name: 'pandas',
+					version: '0.24.2'
+				}, {
+					name: 'prose-codeaccelerator',
+					version: '1.3.0'
+				}, {
+					name: 'powershell-kernel',
+					version: '0.1.3'
+				});
+				break;
+			default:
+				return undefined;
+		}
+		return packages;
+	}
 }
 
 export interface PythonPkgDetails {

extensions/notebook/src/jupyter/jupyterServerManager.ts

@@ -56,16 +56,15 @@ export class LocalJupyterServerManager implements nb.ServerManager, vscode.Dispo
 		return this.options && this.options.jupyterInstallation;
 	}
 
-	public async startServer(): Promise<void> {
+	public async startServer(kernelSpec: nb.IKernelSpec): Promise<void> {
 		try {
 			if (!this._jupyterServer) {
-				this._jupyterServer = await this.doStartServer();
+				this._jupyterServer = await this.doStartServer(kernelSpec);
 				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
 			throw error;
@@ -107,10 +106,10 @@ export class LocalJupyterServerManager implements nb.ServerManager, vscode.Dispo
 		return this.options.documentPath;
 	}
 
-	private async doStartServer(): Promise<IServerInstance> { // We can't find or create servers until the installation is complete
+	private async doStartServer(kernelSpec: nb.IKernelSpec): Promise<IServerInstance> { // We can't find or create servers until the installation is complete
 		let installation = this.options.jupyterInstallation;
-		await installation.promptForPythonInstall();
-		await installation.promptForPackageUpgrade();
+		await installation.promptForPythonInstall(kernelSpec.display_name);
+		await installation.promptForPackageUpgrade(kernelSpec.display_name);
 		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

extensions/notebook/src/jupyter/jupyterSessionManager.ts

@@ -15,6 +15,7 @@ const localize = nls.loadMessageBundle();
 
 import { JupyterKernel } from './jupyterKernel';
 import { Deferred } from '../common/promise';
+import { JupyterServerInstallation } from './jupyterServerInstallation';
 
 const configBase = {
 	'kernel_python_credentials': {
@@ -66,6 +67,7 @@ export class JupyterSessionManager implements nb.SessionManager {
 	private _isReady: boolean;
 	private _sessionManager: Session.IManager;
 	private static _sessions: JupyterSession[] = [];
+	private _installation: JupyterServerInstallation;
 
 	constructor(private _pythonEnvVarPath?: string) {
 		this._isReady = false;
@@ -84,6 +86,12 @@ export class JupyterSessionManager implements nb.SessionManager {
 			});
 	}
 
+	public set installation(installation: JupyterServerInstallation) {
+		this._installation = installation;
+		JupyterSessionManager._sessions.forEach(session => {
+			session.installation = installation;
+		});
+	}
 	public get isReady(): boolean {
 		return this._isReady;
 	}
@@ -126,7 +134,7 @@ export class JupyterSessionManager implements nb.SessionManager {
 			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, skipSettingEnvironmentVars, this._pythonEnvVarPath);
+		let jupyterSession = new JupyterSession(sessionImpl, this._installation, skipSettingEnvironmentVars, this._pythonEnvVarPath);
 		await jupyterSession.messagesComplete;
 		let index = JupyterSessionManager._sessions.findIndex(session => session.path === options.path);
 		if (index > -1) {
@@ -173,7 +181,7 @@ export class JupyterSession implements nb.ISession {
 	private _kernel: nb.IKernel;
 	private _messagesComplete: Deferred<void> = new Deferred<void>();
 
-	constructor(private sessionImpl: Session.ISession, skipSettingEnvironmentVars?: boolean, private _pythonEnvVarPath?: string) {
+	constructor(private sessionImpl: Session.ISession, private _installation: JupyterServerInstallation, skipSettingEnvironmentVars?: boolean, private _pythonEnvVarPath?: string) {
 		this.setEnvironmentVars(skipSettingEnvironmentVars).catch(error => {
 			console.error(`Unexpected exception setting Jupyter Session variables : ${error}`);
 			// We don't want callers to hang forever waiting - it's better to continue on even if we weren't
@@ -221,7 +229,20 @@ export class JupyterSession implements nb.ISession {
 		return this._messagesComplete.promise;
 	}
 
+	public set installation(installation: JupyterServerInstallation) {
+		this._installation = installation;
+	}
+
 	public async changeKernel(kernelInfo: nb.IKernelSpec): Promise<nb.IKernel> {
+		if (this._installation) {
+			try {
+				await this._installation.promptForPackageUpgrade(kernelInfo.display_name);
+			} catch (err) {
+				// Have to swallow the error here to prevent hangs when changing back to the old kernel.
+				console.error(err.toString());
+				return this._kernel;
+			}
+		}
 		// For now, Jupyter implementation handles disposal etc. so we can just
 		// null out our kernel and let the changeKernel call handle this
 		this._kernel = undefined;