Skip to content

Commit

Permalink
Swap controllers once kernel starts
Browse files Browse the repository at this point in the history
  • Loading branch information
DonJayamanne committed Jul 22, 2022
1 parent 545c0a7 commit 5c0c2ec
Show file tree
Hide file tree
Showing 31 changed files with 832 additions and 351 deletions.
2 changes: 1 addition & 1 deletion .github/workflows/build-test.yml
Original file line number Diff line number Diff line change
Expand Up @@ -281,8 +281,8 @@ jobs:
matrix:
jupyter: [raw, local]
python: [nonConda, conda]
pythonVersion: ['3.9', '3.10']
webOrDesktop: ['']
pythonVersion: ['3.9', '3.10']
# Whether we're using stable (<empty>) or pre-release versions of Python packages.
# When installing pre-release versions, we're only focused on jupyter & related packages.
# Not pre-release versions of pandas, numpy or other such packages that are not core to Jupyter.
Expand Down
8 changes: 4 additions & 4 deletions TELEMETRY.md
Original file line number Diff line number Diff line change
Expand Up @@ -4141,7 +4141,7 @@ No properties for event
const telemetryEvent = isLocalConnection(this.kernelConnection)
? Telemetry.SelectLocalJupyterKernel
: Telemetry.SelectRemoteJupyterKernel;
sendKernelTelemetryEvent(document.uri, telemetryEvent);
sendKernelTelemetryEvent(notebook.uri, telemetryEvent);
this.notebookApi.notebookEditors
```
Expand All @@ -4168,9 +4168,9 @@ No properties for event
const telemetryEvent = isLocalConnection(this.kernelConnection)
? Telemetry.SelectLocalJupyterKernel
: Telemetry.SelectRemoteJupyterKernel;
sendKernelTelemetryEvent(document.uri, telemetryEvent);
sendKernelTelemetryEvent(notebook.uri, telemetryEvent);
this.notebookApi.notebookEditors
.filter((editor) => editor.notebook === document)
.filter((editor) => editor.notebook === notebook)
```
</details>
Expand Down Expand Up @@ -8871,7 +8871,7 @@ No properties for event
default:
// We don't know as its the default kernel on Jupyter server.
}
sendKernelTelemetryEvent(document.uri, Telemetry.SwitchKernel);
sendKernelTelemetryEvent(notebook.uri, Telemetry.SwitchKernel);
// If we have an existing kernel, then we know for a fact the user is changing the kernel.
// Else VSC is just setting a kernel for a notebook after it has opened.
if (existingKernel) {
Expand Down
1 change: 1 addition & 0 deletions news/2 Fixes/10572.md
Original file line number Diff line number Diff line change
@@ -0,0 +1 @@
Ensure the kernel picker points to the live Kernel Connection when starting a new remote kernel.
73 changes: 37 additions & 36 deletions src/interactive-window/interactiveWindow.ts
Original file line number Diff line number Diff line change
Expand Up @@ -32,6 +32,7 @@ import { IConfigurationService, InteractiveWindowMode, IsWebExtension, Resource
import { noop } from '../platform/common/utils/misc';
import {
IKernel,
IKernelProvider,
isLocalConnection,
KernelAction,
KernelConnectionMetadata,
Expand All @@ -41,7 +42,6 @@ import { chainable } from '../platform/common/utils/decorators';
import { InteractiveCellResultError } from '../platform/errors/interactiveCellResultError';
import { DataScience } from '../platform/common/utils/localize';
import { createDeferred, Deferred } from '../platform/common/utils/async';
import { IServiceContainer } from '../platform/ioc/types';
import { SysInfoReason } from '../messageTypes';
import { createOutputWithErrorMessageForDisplay } from '../platform/errors/errorUtils';
import { INotebookExporter } from '../kernels/jupyter/types';
Expand All @@ -54,11 +54,9 @@ import {
IInteractiveWindowDebuggingManager,
InteractiveTab
} from './types';
import { generateInteractiveCode, isInteractiveInputTab } from './helpers';
import { generateInteractiveCode, getInteractiveCellMetadata, isInteractiveInputTab } from './helpers';
import { IControllerSelection, IVSCodeNotebookController } from '../notebooks/controllers/types';
import { DisplayOptions } from '../kernels/displayOptions';
import { getInteractiveCellMetadata } from './helpers';
import { KernelConnector } from '../notebooks/controllers/kernelConnector';
import { getFilePath } from '../platform/common/platform/fs-paths';
import {
ICodeGeneratorFactory,
Expand All @@ -71,6 +69,8 @@ import { updateNotebookMetadata } from '../kernels/execution/helpers';
import { chainWithPendingUpdates } from '../kernels/execution/notebookUpdater';
import { initializeInteractiveOrNotebookTelemetryBasedOnUserAction } from '../kernels/telemetry/helper';
import { generateMarkdownFromCodeLines, parseForComments } from '../platform/common/utils';
import { IServiceContainer } from '../platform/ioc/types';
import { KernelConnector } from '../notebooks/controllers/kernelConnector';

/**
* ViewModel for an interactive window from the Jupyter extension's point of view.
Expand Down Expand Up @@ -105,7 +105,8 @@ export class InteractiveWindow implements IInteractiveWindowLoadable {
private kernelDisposables: Disposable[] = [];
private _insertSysInfoPromise: Promise<NotebookCell> | undefined;
private currentKernelInfo: {
kernel?: Deferred<IKernel>;
kernelStarted?: Deferred<void>;
kernel?: IKernel;
controller?: NotebookController;
metadata?: KernelConnectionMetadata;
} = {};
Expand Down Expand Up @@ -135,6 +136,7 @@ export class InteractiveWindow implements IInteractiveWindowLoadable {
private readonly debuggingManager: IInteractiveWindowDebuggingManager;
private readonly isWebExtension: boolean;
private readonly commandManager: ICommandManager;
private readonly kernelProvider: IKernelProvider;
constructor(
private readonly serviceContainer: IServiceContainer,
private _owner: Resource,
Expand All @@ -156,6 +158,7 @@ export class InteractiveWindow implements IInteractiveWindowLoadable {
this.errorHandler = this.serviceContainer.get<IDataScienceErrorHandler>(IDataScienceErrorHandler);
this.codeGeneratorFactory = this.serviceContainer.get<ICodeGeneratorFactory>(ICodeGeneratorFactory);
this.storageFactory = this.serviceContainer.get<IGeneratedCodeStorageFactory>(IGeneratedCodeStorageFactory);
this.kernelProvider = this.serviceContainer.get<IKernelProvider>(IKernelProvider);
this.debuggingManager = this.serviceContainer.get<IInteractiveWindowDebuggingManager>(
IInteractiveWindowDebuggingManager
);
Expand Down Expand Up @@ -229,31 +232,39 @@ export class InteractiveWindow implements IInteractiveWindowLoadable {
// This cannot happen, but we need to make typescript happy.
throw new Error('Controller not selected');
}
if (this.currentKernelInfo.kernel) {
return this.currentKernelInfo.kernel.promise;
if (this.currentKernelInfo.kernelStarted) {
await this.currentKernelInfo.kernelStarted.promise;
return this.currentKernelInfo.kernel!;
}
const kernelPromise = createDeferred<IKernel>();
kernelPromise.promise.catch(noop);
this.currentKernelInfo = { controller, metadata, kernel: kernelPromise };
const kernelStarted = createDeferred<void>();
kernelStarted.promise.catch(noop);
this.currentKernelInfo = { controller, metadata, kernelStarted };

const sysInfoCell = this.insertSysInfoMessage(metadata, SysInfoReason.Start);
try {
// Try creating a kernel
initializeInteractiveOrNotebookTelemetryBasedOnUserAction(this.owner, metadata);

const onStartKernel = (action: KernelAction, k: IKernel) => {
const onKernelStarted = async (action: KernelAction, k: IKernel) => {
if (action !== 'start' && action !== 'restart') {
return;
}
// Id may be different if the user switched controllers
this.currentKernelInfo.controller = k.controller;
this.currentKernelInfo.metadata = k.kernelConnectionMetadata;
this.currentKernelInfo.kernel = k;
this.updateSysInfoMessage(
this.getSysInfoMessage(k.kernelConnectionMetadata, SysInfoReason.Start),
false,
sysInfoCell
);
};
const onKernelStartCompleted = async (action: KernelAction, _: unknown, k: IKernel) => {
if (action !== 'start' && action !== 'restart') {
return;
}
// Id may be different if the user switched controllers
this.currentKernelInfo.controller = k.controller;
this.currentKernelInfo.metadata = k.kernelConnectionMetadata;
};
// When connecting, we need to update the sys info message
this.updateSysInfoMessage(this.getSysInfoMessage(metadata, SysInfoReason.Start), false, sysInfoCell);
const kernel = await KernelConnector.connectToNotebookKernel(
Expand All @@ -263,7 +274,8 @@ export class InteractiveWindow implements IInteractiveWindowLoadable {
new DisplayOptions(false),
this.internalDisposables,
'jupyterExtension',
onStartKernel
onKernelStarted,
onKernelStartCompleted
);
this.currentKernelInfo.controller = kernel.controller;
this.currentKernelInfo.metadata = kernel.kernelConnectionMetadata;
Expand Down Expand Up @@ -303,10 +315,11 @@ export class InteractiveWindow implements IInteractiveWindowLoadable {
this.fileInKernel = undefined;
await this.runInitialization(kernel, this.owner);
this.finishSysInfoMessage(kernel, sysInfoCell, SysInfoReason.Start);
kernelPromise.resolve(kernel);
kernelStarted.resolve();
return kernel;
} catch (ex) {
kernelPromise.reject(ex);
kernelStarted.reject(ex);
this.currentKernelInfo.kernelStarted = undefined;
this.currentKernelInfo.kernel = undefined;
this.disconnectKernel();
if (this.owner) {
Expand Down Expand Up @@ -422,26 +435,8 @@ export class InteractiveWindow implements IInteractiveWindowLoadable {
message = message.split('\n').join(' \n');
this.updateSysInfoMessage(message, true, cellPromise);
}
private registerControllerChangeListener(controller: IVSCodeNotebookController) {
const controllerChangeListener = controller.controller.onDidChangeSelectedNotebooks(
(selectedEvent: { notebook: NotebookDocument; selected: boolean }) => {
// Controller was deselected for this InteractiveWindow's NotebookDocument
if (selectedEvent.selected === false && selectedEvent.notebook === this.notebookEditor.notebook) {
controllerChangeListener.dispose();
this.disconnectKernel();
}
},
this,
this.internalDisposables
);
}

private listenForControllerSelection() {
const controller = this.notebookControllerSelection.getSelected(this.notebookEditor.notebook);
if (controller !== undefined) {
this.registerControllerChangeListener(controller);
}

// Ensure we hear about any controller changes so we can update our cached promises
this.notebookControllerSelection.onControllerSelected(
(e: { notebook: NotebookDocument; controller: IVSCodeNotebookController }) => {
Expand All @@ -450,8 +445,13 @@ export class InteractiveWindow implements IInteractiveWindowLoadable {
}

// Clear cached kernel when the selected controller for this document changes
this.registerControllerChangeListener(e.controller);
if (e.controller.id !== this.currentKernelInfo.controller?.id) {
const kernel = this.kernelProvider.get(this.notebookDocument.uri);
const isControllerAttachedToTheSameKernel =
kernel && this.currentKernelInfo.kernel && kernel === this.currentKernelInfo.kernel;
if (
e.controller.connection.id !== this.currentKernelInfo.kernel?.kernelConnectionMetadata.id &&
!isControllerAttachedToTheSameKernel
) {
this.disconnectKernel();
this.startKernel(e.controller.controller, e.controller.connection).ignoreErrors();
}
Expand Down Expand Up @@ -614,6 +614,7 @@ export class InteractiveWindow implements IInteractiveWindowLoadable {
private disconnectKernel() {
this.kernelDisposables.forEach((d) => d.dispose());
this.kernelDisposables = [];
this.currentKernelInfo.kernelStarted = undefined;
this.currentKernelInfo.kernel = undefined;
}

Expand Down
72 changes: 67 additions & 5 deletions src/kernels/ipywidgets/commonMessageCoordinator.ts
Original file line number Diff line number Diff line change
Expand Up @@ -4,9 +4,9 @@
'use strict';

import type { KernelMessage } from '@jupyterlab/services';
import { Event, EventEmitter, NotebookDocument } from 'vscode';
import { IApplicationShell, ICommandManager } from '../../platform/common/application/types';
import { STANDARD_OUTPUT_CHANNEL } from '../../platform/common/constants';
import { Event, EventEmitter, NotebookDocument, notebooks } from 'vscode';
import { IApplicationShell, ICommandManager, IVSCodeNotebook } from '../../platform/common/application/types';
import { IPyWidgetRendererId, STANDARD_OUTPUT_CHANNEL } from '../../platform/common/constants';
import { traceVerbose, traceError, traceInfo, traceInfoIfCI } from '../../platform/logging';
import {
IDisposableRegistry,
Expand Down Expand Up @@ -65,7 +65,9 @@ export class CommonMessageCoordinator {
private readonly configService: IConfigurationService;
private webview: IWebviewCommunication | undefined;
private modulesForWhichWeHaveDisplayedWidgetErrorMessage = new Set<string>();

private kernelProvider: IKernelProvider;
private queuedMessages: { type: string; payload: unknown }[] = [];
private readyMessageReceived?: boolean;
public constructor(
private readonly document: NotebookDocument,
private readonly serviceContainer: IServiceContainer
Expand All @@ -75,6 +77,7 @@ export class CommonMessageCoordinator {
this.appShell = this.serviceContainer.get<IApplicationShell>(IApplicationShell, IApplicationShell);
this.commandManager = this.serviceContainer.get<ICommandManager>(ICommandManager);
this.configService = this.serviceContainer.get<IConfigurationService>(IConfigurationService);
this.kernelProvider = this.serviceContainer.get<IKernelProvider>(IKernelProvider);
}

public dispose() {
Expand Down Expand Up @@ -103,6 +106,12 @@ export class CommonMessageCoordinator {
response: webview.asWebviewUri(e.payload)
});
} else {
if (!webview.isReady || !this.readyMessageReceived) {
// Web view is not yet ready to receive messages, hence queue these to be sent later.
this.queuedMessages.push({ type: e.message, payload: e.payload });
return;
}
this.sendPendingWebViewMessages();
webview.postMessage({ type: e.message, payload: e.payload }).then(noop, noop);
}
},
Expand All @@ -113,10 +122,55 @@ export class CommonMessageCoordinator {
(m) => {
traceInfoIfCI(`${ConsoleForegroundColors.Green}Widget Coordinator received ${m.type}`);
this.onMessage(m.type, m.payload);

const kernel = this.kernelProvider.get(this.document.uri);
// Special case the WidgetManager loaded message. It means we're ready
// to use a kernel. (IPyWidget Dispatcher uses this too)
if (m.type === IPyWidgetMessages.IPyWidgets_Ready) {
if (kernel?.kernelConnectionMetadata.kind === 'startUsingRemoteKernelSpec') {
traceInfoIfCI(
'Web view is not ready to receive widget messages (kernel points to remote kernel spec)'
);
const nbEditor = this.serviceContainer
.get<IVSCodeNotebook>(IVSCodeNotebook)
.notebookEditors.find((item) => item.notebook === this.document);
// With remote kernel specs, once the kernel is ready we create a live kernel controller and
// switch to that. At that point the webview also changes, hence
// there's no need to render anything while were in this state.
notebooks
.createRendererMessaging(IPyWidgetRendererId)
.postMessage({ type: IPyWidgetMessages.IPyWidgets_DoNotRenderWidgets }, nbEditor)
.then(noop, noop);
return;
}
if (!webview.isReady) {
traceInfoIfCI('Web view is not ready to receive widget messages');
return;
}
traceInfoIfCI('Web view is ready to receive widget messages');
this.readyMessageReceived = true;
this.sendPendingWebViewMessages();
// At this point, we know the kernels are ready, and the webview is ready to receive messages.
// Its possible the webview was initially unable to render some widgets, but now that everything is ready we
// should be able to render them now.
// E.g assume we're dealing with remote kernel specs,
// Then we start a kernel, next we create a controller that points to the live kernel & change the controller.
// What happens now is the, webview is re-loaded (as theres a change in the controllers).
// However if the user were to run a cell, then the output would get displayed and it could end up,
// being displayed in the old webview & when the webview is re-loaded, it would be displayed in the new webview.
// However its possble the kernel is not yet ready at that point in time.
// We could solve this issue easily by not executing cells, until the webview is ready,
// however that would have significant performance implications.
// Hence for ipywidgets, once the webview has completely been initialized, we can attempt to re-render the widgets.
const nbEditor = this.serviceContainer
.get<IVSCodeNotebook>(IVSCodeNotebook)
.notebookEditors.find((item) => item.notebook === this.document);
if (nbEditor) {
traceInfoIfCI('Re-rendering widgets');
notebooks
.createRendererMessaging(IPyWidgetRendererId)
.postMessage({ type: IPyWidgetMessages.IPyWidgets_ReRenderWidgets }, nbEditor)
.then(noop, noop);
}
promise.resolve();
}
},
Expand Down Expand Up @@ -155,6 +209,14 @@ export class CommonMessageCoordinator {
this.getIPyWidgetScriptSource().onMessage(message, payload);
}

private sendPendingWebViewMessages() {
if (!this.webview || !this.webview.isReady || !this.readyMessageReceived) {
return;
}
while (this.queuedMessages.length) {
this.webview.postMessage(this.queuedMessages.shift()!).then(noop, noop);
}
}
private initialize() {
traceVerbose('initialize CommonMessageCoordinator');
// First hook up the widget script source that will listen to messages even before we start sending messages.
Expand Down
Loading

0 comments on commit 5c0c2ec

Please sign in to comment.