From 901d29ce266f1b86ee4f9fc07864f23c0700faf2 Mon Sep 17 00:00:00 2001 From: Don Jayamanne Date: Fri, 19 Apr 2024 14:47:09 +1000 Subject: [PATCH 1/2] VS Code API replacement for cell execution events --- src/kernels/execution/cellExecution.ts | 13 ++- src/kernels/execution/cellExecutionQueue.ts | 7 +- .../controllers/controllerRegistration.ts | 6 +- .../notebookCellExecutionStateService.ts | 90 +++++++++++++++++++ .../controllers/serviceRegistry.node.ts | 6 ++ src/notebooks/controllers/types.ts | 39 ++++++++ .../controllers/vscodeNotebookController.ts | 17 +++- .../vscodeNotebookController.unit.test.ts | 8 +- .../notebook/kernelCrashes.vscode.test.ts | 4 +- 9 files changed, 174 insertions(+), 16 deletions(-) create mode 100644 src/notebooks/controllers/notebookCellExecutionStateService.ts diff --git a/src/kernels/execution/cellExecution.ts b/src/kernels/execution/cellExecution.ts index 30b7eb3e766..04237e167a5 100644 --- a/src/kernels/execution/cellExecution.ts +++ b/src/kernels/execution/cellExecution.ts @@ -74,8 +74,11 @@ export class CellExecution implements ICellExecution, IDisposable { public get result(): Promise { return this._result.promise; } - public get preExecute(): Event { - return this._preExecuteEmitter.event; + public get onWillExecute(): Event { + return this._onWillExecute.event; + } + public get onDidExecute(): Event { + return this._onDidExecute.event; } private readonly _result = createDeferred(); @@ -94,7 +97,8 @@ export class CellExecution implements ICellExecution, IDisposable { private disposed?: boolean; private request: Kernel.IShellFuture | undefined; private readonly disposables: IDisposable[] = []; - private _preExecuteEmitter = new EventEmitter(); + private _onWillExecute = new EventEmitter(); + private _onDidExecute = new EventEmitter(); private cellExecutionHandler?: CellExecutionMessageHandler; private session?: IKernelSession; private cancelRequested?: boolean; @@ -159,6 +163,7 @@ export class CellExecution implements ICellExecution, IDisposable { } public async start(session: IKernelSession) { this.session = session; + void this._result.promise.finally(() => this._onDidExecute.fire()); if (this.resumeExecution?.msg_id) { return this.resume(session, this.resumeExecution); } @@ -423,7 +428,7 @@ export class CellExecution implements ICellExecution, IDisposable { const kernelConnection = session.kernel; try { // At this point we're about to ACTUALLY execute some code. Fire an event to indicate that - this._preExecuteEmitter.fire(this.cell); + this._onWillExecute.fire(); traceVerbose(`Cell Index:${this.cell.index} sent to kernel`); // For Jupyter requests, silent === don't output, while store_history === don't update execution count // https://jupyter-client.readthedocs.io/en/stable/api/client.html#jupyter_client.KernelClient.execute diff --git a/src/kernels/execution/cellExecutionQueue.ts b/src/kernels/execution/cellExecutionQueue.ts index 3616527a206..2852df31114 100644 --- a/src/kernels/execution/cellExecutionQueue.ts +++ b/src/kernels/execution/cellExecutionQueue.ts @@ -101,7 +101,8 @@ export class CellExecutionQueue implements Disposable { const cellExecution = this.executionFactory.create(cell, codeOverride, this.metadata); executionItem = cellExecution; this.disposables.push(cellExecution); - cellExecution.preExecute((c) => this._onPreExecute.fire(c), this, this.disposables); + cellExecution.onWillExecute(() => this._onPreExecute.fire(cellExecution.cell), this, this.disposables); + cellExecution.onDidExecute(() => this._onPostExecute.fire(cellExecution.cell), this, this.disposables); this.queueOfItemsToExecute.push(cellExecution); traceCellMessage(cell, 'User queued cell for execution'); @@ -243,10 +244,6 @@ export class CellExecutionQueue implements Disposable { if (index >= 0) { this.queueOfItemsToExecute.splice(index, 1); } - - if (itemToExecute.type === 'cell') { - this._onPostExecute.fire(itemToExecute.cell); - } } let cellErrorsAllowed = false; diff --git a/src/notebooks/controllers/controllerRegistration.ts b/src/notebooks/controllers/controllerRegistration.ts index 2eeed9a47cc..f6f56e572aa 100644 --- a/src/notebooks/controllers/controllerRegistration.ts +++ b/src/notebooks/controllers/controllerRegistration.ts @@ -25,6 +25,7 @@ import { PythonEnvironmentFilter } from '../../platform/interpreter/filter/filte import { IConnectionDisplayDataProvider, IControllerRegistration, + INotebookCellExecutionStateService, InteractiveControllerIdSuffix, IVSCodeNotebookController, IVSCodeNotebookControllerUpdateEvent @@ -360,7 +361,10 @@ export class ControllerRegistration implements IControllerRegistration, IExtensi this.extensionChecker, this.serviceContainer, this.serviceContainer.get(IConnectionDisplayDataProvider), - this.serviceContainer.get(IJupyterVariables, Identifiers.KERNEL_VARIABLES) + this.serviceContainer.get(IJupyterVariables, Identifiers.KERNEL_VARIABLES), + this.serviceContainer.get( + INotebookCellExecutionStateService + ) ); // Hook up to if this NotebookController is selected or de-selected const controllerDisposables: IDisposable[] = []; diff --git a/src/notebooks/controllers/notebookCellExecutionStateService.ts b/src/notebooks/controllers/notebookCellExecutionStateService.ts new file mode 100644 index 00000000000..059b98324cc --- /dev/null +++ b/src/notebooks/controllers/notebookCellExecutionStateService.ts @@ -0,0 +1,90 @@ +// Copyright (c) Microsoft Corporation. +// Licensed under the MIT License. + +import { EventEmitter, type Event, type NotebookCell } from 'vscode'; +import { DisposableBase, DisposableStore } from '../../platform/common/utils/lifecycle'; +import { StopWatch } from '../../platform/common/utils/stopWatch'; +import { inject, injectable } from 'inversify'; +import { IKernelProvider, type IKernel, type INotebookKernelExecution } from '../../kernels/types'; +import { + NotebookCellExecutionState, + type INotebookCellExecutionStateService, + type NotebookCellExecutionStateChangeEvent +} from './types'; +import { IDisposableRegistry } from '../../platform/common/types'; + +@injectable() +export class NotebookCellExecutionStateService extends DisposableBase implements INotebookCellExecutionStateService { + private static cellStates = new WeakMap< + NotebookCell, + { stateTransition: string[]; state: NotebookCellExecutionState; start: StopWatch } + >(); + + /** + * An {@link Event} which fires when the execution state of a cell has changed. + */ + // todo@API this is an event that is fired for a property that cells don't have and that makes me wonder + // how a correct consumer works, e.g the consumer could have been late and missed an event? + private readonly _onDidChangeNotebookCellExecutionState = this._register( + new EventEmitter() + ); + + public static getCellState(cell: NotebookCell): NotebookCellExecutionState | undefined { + return NotebookCellExecutionStateService.cellStates.get(cell)?.state; + } + public static getCellStatus(cell: NotebookCell): string { + return (NotebookCellExecutionStateService.cellStates.get(cell)?.stateTransition || []).join(', ') || ''; + } + public get onDidChangeNotebookCellExecutionState() { + return this._onDidChangeNotebookCellExecutionState.event; + } + private readonly trackedExecutions = new WeakSet(); + private readonly kernelExecutionMaps = new WeakMap(); + private readonly kernelDisposables = new WeakMap(); + constructor( + @inject(IKernelProvider) private readonly kernelProvider: IKernelProvider, + @inject(IDisposableRegistry) disposables: IDisposableRegistry + ) { + super(); + disposables.push(this); + this._register(kernelProvider.onDidCreateKernel(this.monitorKernelExecutionEvents, this)); + this._register(kernelProvider.onDidStartKernel(this.monitorKernelExecutionEvents, this)); + this._register(kernelProvider.onDidDisposeKernel((k) => this.kernelDisposables.get(k)?.dispose(), this)); + } + setPendingState(cell: NotebookCell): void { + this.setCellState(cell, NotebookCellExecutionState.Pending); + } + + private monitorKernelExecutionEvents(kernel: IKernel) { + const execution = this.kernelProvider.getKernelExecution(kernel); + if (this.trackedExecutions.has(execution)) { + return; + } + const disposableStore = this.kernelDisposables.get(kernel) || new DisposableStore(); + this.kernelDisposables.set(kernel, disposableStore); + this._register(disposableStore); + + this.kernelExecutionMaps.set(kernel, execution); + disposableStore.add( + execution.onPreExecute((cell) => this.setCellState(cell, NotebookCellExecutionState.Executing), this) + ); + disposableStore.add( + execution.onPostExecute((cell) => this.setCellState(cell, NotebookCellExecutionState.Idle), this) + ); + } + + private setCellState(cell: NotebookCell, state: NotebookCellExecutionState) { + NotebookCellExecutionStateService.setCellState(cell, state); + this._onDidChangeNotebookCellExecutionState.fire({ cell, state: state }); + } + private static setCellState(cell: NotebookCell, state: NotebookCellExecutionState) { + const stopWatch = NotebookCellExecutionStateService.cellStates.get(cell)?.start || new StopWatch(); + const previousState = NotebookCellExecutionStateService.cellStates.get(cell)?.stateTransition || []; + previousState.push(`${state} ${previousState.length === 0 ? '@ start' : `After ${stopWatch.elapsedTime}ms`}`); + NotebookCellExecutionStateService.cellStates.set(cell, { + stateTransition: previousState, + state, + start: stopWatch + }); + } +} diff --git a/src/notebooks/controllers/serviceRegistry.node.ts b/src/notebooks/controllers/serviceRegistry.node.ts index 9c4bf55c680..55a508095e5 100644 --- a/src/notebooks/controllers/serviceRegistry.node.ts +++ b/src/notebooks/controllers/serviceRegistry.node.ts @@ -10,11 +10,13 @@ import { KernelSourceCommandHandler } from './kernelSource/kernelSourceCommandHa import { LocalNotebookKernelSourceSelector } from './kernelSource/localNotebookKernelSourceSelector.node'; import { LocalPythonEnvNotebookKernelSourceSelector } from './kernelSource/localPythonEnvKernelSourceSelector.node'; import { RemoteNotebookKernelSourceSelector } from './kernelSource/remoteNotebookKernelSourceSelector'; +import { NotebookCellExecutionStateService } from './notebookCellExecutionStateService'; import { IConnectionDisplayDataProvider, IControllerRegistration, ILocalNotebookKernelSourceSelector, ILocalPythonNotebookKernelSourceSelector, + INotebookCellExecutionStateService, IRemoteNotebookKernelSourceSelector } from './types'; @@ -37,6 +39,10 @@ export function registerTypes(serviceManager: IServiceManager, isDevMode: boolea ILocalPythonNotebookKernelSourceSelector, LocalPythonEnvNotebookKernelSourceSelector ); + serviceManager.addSingleton( + INotebookCellExecutionStateService, + NotebookCellExecutionStateService + ); serviceManager.addBinding(ILocalPythonNotebookKernelSourceSelector, IExtensionSyncActivationService); serviceManager.addSingleton( IExtensionSyncActivationService, diff --git a/src/notebooks/controllers/types.ts b/src/notebooks/controllers/types.ts index 8dca9ba44b6..85e882f6995 100644 --- a/src/notebooks/controllers/types.ts +++ b/src/notebooks/controllers/types.ts @@ -141,3 +141,42 @@ export const IConnectionDisplayDataProvider = Symbol('IConnectionDisplayData'); export interface IConnectionDisplayDataProvider { getDisplayData(connection: KernelConnectionMetadata): IConnectionDisplayData; } + +/** + * The execution state of a notebook cell. + */ +export enum NotebookCellExecutionState { + /** + * The cell is idle. + */ + Idle = 1, + /** + * Execution for the cell is pending. + */ + Pending = 2, + /** + * The cell is currently executing. + */ + Executing = 3 +} + +/** + * An event describing a cell execution state change. + */ +export interface NotebookCellExecutionStateChangeEvent { + /** + * The {@link NotebookCell cell} for which the execution state has changed. + */ + readonly cell: vscode.NotebookCell; + + /** + * The new execution state of the cell. + */ + readonly state: NotebookCellExecutionState; +} + +export const INotebookCellExecutionStateService = Symbol('INotebookCellExecutionStateService'); +export interface INotebookCellExecutionStateService { + onDidChangeNotebookCellExecutionState: vscode.Event; + setPendingState(cell: vscode.NotebookCell): void; +} diff --git a/src/notebooks/controllers/vscodeNotebookController.ts b/src/notebooks/controllers/vscodeNotebookController.ts index ef7369f15e3..7ba29d63cd9 100644 --- a/src/notebooks/controllers/vscodeNotebookController.ts +++ b/src/notebooks/controllers/vscodeNotebookController.ts @@ -59,7 +59,12 @@ import { DisplayOptions } from '../../kernels/displayOptions'; import { getNotebookMetadata, isJupyterNotebook, updateNotebookMetadata } from '../../platform/common/utils'; import { ConsoleForegroundColors } from '../../platform/logging/types'; import { KernelConnector } from './kernelConnector'; -import { IConnectionDisplayData, IConnectionDisplayDataProvider, IVSCodeNotebookController } from './types'; +import { + IConnectionDisplayData, + IConnectionDisplayDataProvider, + IVSCodeNotebookController, + type INotebookCellExecutionStateService +} from './types'; import { isCancellationError } from '../../platform/common/cancellation'; import { CellExecutionCreator } from '../../kernels/execution/cellExecutionCreator'; import { @@ -159,7 +164,8 @@ export class VSCodeNotebookController implements Disposable, IVSCodeNotebookCont extensionChecker: IPythonExtensionChecker, serviceContainer: IServiceContainer, displayDataProvider: IConnectionDisplayDataProvider, - jupyterVariables: IJupyterVariables + jupyterVariables: IJupyterVariables, + notebookCellExecutionStateService: INotebookCellExecutionStateService ): IVSCodeNotebookController { return new VSCodeNotebookController( kernelConnection, @@ -173,7 +179,8 @@ export class VSCodeNotebookController implements Disposable, IVSCodeNotebookCont extensionChecker, serviceContainer, displayDataProvider, - jupyterVariables + jupyterVariables, + notebookCellExecutionStateService ); } constructor( @@ -188,7 +195,8 @@ export class VSCodeNotebookController implements Disposable, IVSCodeNotebookCont private readonly extensionChecker: IPythonExtensionChecker, private serviceContainer: IServiceContainer, private readonly displayDataProvider: IConnectionDisplayDataProvider, - jupyterVariables: IJupyterVariables + jupyterVariables: IJupyterVariables, + private readonly notebookCellExecutionStateService: INotebookCellExecutionStateService ) { trackControllerCreation(kernelConnection.id, kernelConnection.interpreter?.id); disposableRegistry.push(this); @@ -392,6 +400,7 @@ export class VSCodeNotebookController implements Disposable, IVSCodeNotebookCont return; } traceInfo(`Handle Execution of Cells ${cells.map((c) => c.index)} for ${getDisplayPath(notebook.uri)}`); + cells.forEach((cell) => this.notebookCellExecutionStateService.setPendingState(cell)); await initializeInteractiveOrNotebookTelemetryBasedOnUserAction(notebook.uri, this.connection); telemetryTracker?.stop(); // Notebook is trusted. Continue to execute cells diff --git a/src/notebooks/controllers/vscodeNotebookController.unit.test.ts b/src/notebooks/controllers/vscodeNotebookController.unit.test.ts index 89ec1a66433..d42325ebd5d 100644 --- a/src/notebooks/controllers/vscodeNotebookController.unit.test.ts +++ b/src/notebooks/controllers/vscodeNotebookController.unit.test.ts @@ -43,6 +43,7 @@ import { mockedVSCodeNamespaces, resetVSCodeMocks } from '../../test/vscode-mock import { IJupyterVariables } from '../../kernels/variables/types'; import { Environment, PythonExtension } from '@vscode/python-extension'; import { crateMockedPythonApi, whenResolveEnvironment } from '../../kernels/helpers.unit.test'; +import { NotebookCellExecutionStateService } from './notebookCellExecutionStateService'; suite(`Notebook Controller`, function () { let controller: NotebookController; @@ -144,6 +145,10 @@ suite(`Notebook Controller`, function () { }); teardown(() => (disposables = dispose(disposables))); function createController(viewType: 'jupyter-notebook' | 'interactive') { + const notebookCellExecutionStateService = new NotebookCellExecutionStateService( + instance(kernelProvider), + disposables + ); new VSCodeNotebookController( instance(kernelConnection), '1', @@ -156,7 +161,8 @@ suite(`Notebook Controller`, function () { instance(extensionChecker), instance(serviceContainer), displayDataProvider, - jupyterVariables + jupyterVariables, + notebookCellExecutionStateService ); notebook = new TestNotebookDocument(undefined, viewType); } diff --git a/src/test/datascience/notebook/kernelCrashes.vscode.test.ts b/src/test/datascience/notebook/kernelCrashes.vscode.test.ts index 11db84c8b2f..49da006b2a8 100644 --- a/src/test/datascience/notebook/kernelCrashes.vscode.test.ts +++ b/src/test/datascience/notebook/kernelCrashes.vscode.test.ts @@ -36,6 +36,7 @@ import { IPlatformService } from '../../../platform/common/platform/types'; import { IInterpreterService } from '../../../platform/interpreter/contracts'; import { ConnectionDisplayDataProvider } from '../../../notebooks/controllers/connectionDisplayData.node'; import { IJupyterVariables } from '../../../kernels/variables/types'; +import { INotebookCellExecutionStateService } from '../../../notebooks/controllers/types'; const codeToKillKernel = dedent` import IPython @@ -125,7 +126,8 @@ suite('VSCode Notebook Kernel Error Handling - @kernelCore', function () { extensionChecker, api.serviceContainer, displayDataProvider, - jupyterVariables + jupyterVariables, + api.serviceContainer.get(INotebookCellExecutionStateService) ); disposables.push(interpreterController); From 5e02126f54edc7b8365d05d6b41cd3735d24608a Mon Sep 17 00:00:00 2001 From: Don Jayamanne Date: Fri, 19 Apr 2024 14:48:10 +1000 Subject: [PATCH 2/2] No state tracking --- .../notebookCellExecutionStateService.ts | 31 +++---------------- 1 file changed, 4 insertions(+), 27 deletions(-) diff --git a/src/notebooks/controllers/notebookCellExecutionStateService.ts b/src/notebooks/controllers/notebookCellExecutionStateService.ts index 059b98324cc..caf84d3ab59 100644 --- a/src/notebooks/controllers/notebookCellExecutionStateService.ts +++ b/src/notebooks/controllers/notebookCellExecutionStateService.ts @@ -3,7 +3,6 @@ import { EventEmitter, type Event, type NotebookCell } from 'vscode'; import { DisposableBase, DisposableStore } from '../../platform/common/utils/lifecycle'; -import { StopWatch } from '../../platform/common/utils/stopWatch'; import { inject, injectable } from 'inversify'; import { IKernelProvider, type IKernel, type INotebookKernelExecution } from '../../kernels/types'; import { @@ -15,11 +14,6 @@ import { IDisposableRegistry } from '../../platform/common/types'; @injectable() export class NotebookCellExecutionStateService extends DisposableBase implements INotebookCellExecutionStateService { - private static cellStates = new WeakMap< - NotebookCell, - { stateTransition: string[]; state: NotebookCellExecutionState; start: StopWatch } - >(); - /** * An {@link Event} which fires when the execution state of a cell has changed. */ @@ -29,12 +23,6 @@ export class NotebookCellExecutionStateService extends DisposableBase implements new EventEmitter() ); - public static getCellState(cell: NotebookCell): NotebookCellExecutionState | undefined { - return NotebookCellExecutionStateService.cellStates.get(cell)?.state; - } - public static getCellStatus(cell: NotebookCell): string { - return (NotebookCellExecutionStateService.cellStates.get(cell)?.stateTransition || []).join(', ') || ''; - } public get onDidChangeNotebookCellExecutionState() { return this._onDidChangeNotebookCellExecutionState.event; } @@ -52,7 +40,7 @@ export class NotebookCellExecutionStateService extends DisposableBase implements this._register(kernelProvider.onDidDisposeKernel((k) => this.kernelDisposables.get(k)?.dispose(), this)); } setPendingState(cell: NotebookCell): void { - this.setCellState(cell, NotebookCellExecutionState.Pending); + this.triggerStateChange(cell, NotebookCellExecutionState.Pending); } private monitorKernelExecutionEvents(kernel: IKernel) { @@ -66,25 +54,14 @@ export class NotebookCellExecutionStateService extends DisposableBase implements this.kernelExecutionMaps.set(kernel, execution); disposableStore.add( - execution.onPreExecute((cell) => this.setCellState(cell, NotebookCellExecutionState.Executing), this) + execution.onPreExecute((cell) => this.triggerStateChange(cell, NotebookCellExecutionState.Executing), this) ); disposableStore.add( - execution.onPostExecute((cell) => this.setCellState(cell, NotebookCellExecutionState.Idle), this) + execution.onPostExecute((cell) => this.triggerStateChange(cell, NotebookCellExecutionState.Idle), this) ); } - private setCellState(cell: NotebookCell, state: NotebookCellExecutionState) { - NotebookCellExecutionStateService.setCellState(cell, state); + private triggerStateChange(cell: NotebookCell, state: NotebookCellExecutionState) { this._onDidChangeNotebookCellExecutionState.fire({ cell, state: state }); } - private static setCellState(cell: NotebookCell, state: NotebookCellExecutionState) { - const stopWatch = NotebookCellExecutionStateService.cellStates.get(cell)?.start || new StopWatch(); - const previousState = NotebookCellExecutionStateService.cellStates.get(cell)?.stateTransition || []; - previousState.push(`${state} ${previousState.length === 0 ? '@ start' : `After ${stopWatch.elapsedTime}ms`}`); - NotebookCellExecutionStateService.cellStates.set(cell, { - stateTransition: previousState, - state, - start: stopWatch - }); - } }