Skip to content

Commit 7ce36cd

Browse files
DonJayamanneDonJayamanne
committed
WIP
1 parent bcfc5b1 commit 7ce36cd

12 files changed

+197
-59
lines changed

‎src/kernels/execution/cellExecution.ts

Lines changed: 40 additions & 17 deletions
Original file line numberDiff line numberDiff line change
@@ -9,7 +9,8 @@ import {
99
NotebookCellOutput,
1010
NotebookCellExecutionState,
1111
Event,
12-
EventEmitter
12+
EventEmitter,
13+
CancellationToken
1314
} from 'vscode';
1415

1516
import { Kernel } from '@jupyterlab/services';
@@ -20,7 +21,6 @@ import { disposeAllDisposables } from '../../platform/common/helpers';
2021
import { traceError, traceInfoIfCI, traceVerbose, traceWarning } from '../../platform/logging';
2122
import { IDisposable } from '../../platform/common/types';
2223
import { createDeferred } from '../../platform/common/utils/async';
23-
import { StopWatch } from '../../platform/common/utils/stopWatch';
2424
import { noop } from '../../platform/common/utils/misc';
2525
import { getDisplayNameOrNameOfKernelConnection } from '../../kernels/helpers';
2626
import { isCancellationError } from '../../platform/common/cancellation';
@@ -49,7 +49,10 @@ export class CellExecutionFactory {
4949
code: string | undefined,
5050
metadata: Readonly<KernelConnectionMetadata>,
5151
resumeExecutionMsgId?: string,
52-
restoreOutput?: boolean
52+
restoreOutput?: boolean,
53+
token?: CancellationToken,
54+
startTime?: number,
55+
executionCount?: number
5356
) {
5457
// eslint-disable-next-line @typescript-eslint/no-use-before-define
5558
return CellExecution.fromCell(
@@ -59,7 +62,10 @@ export class CellExecutionFactory {
5962
this.controller,
6063
this.requestListener,
6164
resumeExecutionMsgId,
62-
restoreOutput
65+
restoreOutput,
66+
token,
67+
startTime,
68+
executionCount
6369
);
6470
}
6571
}
@@ -81,8 +87,6 @@ export class CellExecution implements IDisposable {
8187
return this._preExecuteEmitter.event;
8288
}
8389

84-
private stopWatch = new StopWatch();
85-
8690
private readonly _result = createDeferred<NotebookCellRunState>();
8791

8892
private started?: boolean;
@@ -97,14 +101,18 @@ export class CellExecution implements IDisposable {
97101
private readonly disposables: IDisposable[] = [];
98102
private _preExecuteEmitter = new EventEmitter<NotebookCell>();
99103
private cellExecutionHandler?: CellExecutionMessageHandler;
104+
private cancelRequested?: boolean;
100105
private constructor(
101106
public readonly cell: NotebookCell,
102107
private readonly codeOverride: string | undefined,
103108
private readonly kernelConnection: Readonly<KernelConnectionMetadata>,
104109
private readonly controller: IKernelController,
105110
private readonly requestListener: CellExecutionMessageHandlerService,
106111
private readonly resumeExecutionMsgId?: string,
107-
private readonly restoreOutput?: boolean
112+
private readonly restoreOutput?: boolean,
113+
private readonly token?: CancellationToken,
114+
private readonly _startTime?: number,
115+
private readonly executionCount?: number
108116
) {
109117
workspace.onDidCloseTextDocument(
110118
(e) => {
@@ -152,7 +160,10 @@ export class CellExecution implements IDisposable {
152160
controller: IKernelController,
153161
requestListener: CellExecutionMessageHandlerService,
154162
resumeExecutionMsgId?: string,
155-
restoreOutput?: boolean
163+
restoreOutput?: boolean,
164+
token?: CancellationToken,
165+
startTime?: number,
166+
executionCount?: number
156167
) {
157168
return new CellExecution(
158169
cell,
@@ -161,11 +172,14 @@ export class CellExecution implements IDisposable {
161172
controller,
162173
requestListener,
163174
resumeExecutionMsgId,
164-
restoreOutput
175+
restoreOutput,
176+
token,
177+
startTime,
178+
executionCount
165179
);
166180
}
167181
public async resume(session: IKernelConnectionSession) {
168-
if (this.cancelHandled) {
182+
if (this.cancelHandled || this.token?.isCancellationRequested) {
169183
traceCellMessage(this.cell, 'Not resuming as it was cancelled');
170184
return;
171185
}
@@ -186,11 +200,13 @@ export class CellExecution implements IDisposable {
186200
}
187201
this.started = true;
188202

189-
this.startTime = new Date().getTime();
203+
this.startTime = this._startTime || new Date().getTime();
190204
activeNotebookCellExecution.set(this.cell.notebook, this.execution);
191205
this.execution?.start(this.startTime);
206+
if (this.executionCount && this.execution) {
207+
this.execution.executionOrder = this.executionCount;
208+
}
192209
NotebookCellStateTracker.setCellState(this.cell, NotebookCellExecutionState.Executing);
193-
this.stopWatch.reset();
194210

195211
this.cellExecutionHandler = this.requestListener.registerListenerForResumingExecution(this.cell, {
196212
kernel: session.kernel!,
@@ -237,7 +253,6 @@ export class CellExecution implements IDisposable {
237253
// Else when running cells with existing outputs, the outputs don't get cleared & it doesn't look like its running.
238254
// Ideally we shouldn't have any awaits, but here we want the UI to get updated.
239255
await this.execution?.clearOutput();
240-
this.stopWatch.reset();
241256

242257
// Begin the request that will modify our cell.
243258
this.execute(this.codeOverride || this.cell.document.getText().replace(/\r\n/g, '\n'), session)
@@ -258,6 +273,7 @@ export class CellExecution implements IDisposable {
258273
if (this.cancelHandled) {
259274
return;
260275
}
276+
this.cancelRequested = true;
261277
if (this.started && !forced) {
262278
// At this point the cell execution can only be stopped from kernel & we should not
263279
// stop handling execution results & the like from the kernel.
@@ -288,7 +304,11 @@ export class CellExecution implements IDisposable {
288304
disposeAllDisposables(this.disposables);
289305
}
290306
private completedWithErrors(error: Partial<Error>) {
291-
traceWarning(`Cell completed with errors`, error);
307+
if (!this.disposed && !this.cancelRequested) {
308+
traceWarning(`Cell completed with errors`, error);
309+
} else {
310+
traceWarning(`Cell completed with errors (${this.disposed ? 'disposed' : 'cancelled'})`);
311+
}
292312
traceCellMessage(this.cell, 'Completed with errors');
293313

294314
traceCellMessage(this.cell, 'Update with error state & output');
@@ -420,6 +440,7 @@ export class CellExecution implements IDisposable {
420440
this.cellExecutionHandler = this.requestListener.registerListenerForExecution(this.cell, {
421441
kernel: session.kernel!,
422442
cellExecution: this.execution!,
443+
startTime: this.startTime!,
423444
request: this.request,
424445
onErrorHandlingExecuteRequestIOPubMessage: (error) => {
425446
traceError(`Cell (index = ${this.cell.index}) execution completed with errors (2).`, error);
@@ -442,9 +463,11 @@ export class CellExecution implements IDisposable {
442463
this.completedSuccessfully();
443464
traceCellMessage(this.cell, 'Executed successfully in executeCell');
444465
} catch (ex) {
445-
// @jupyterlab/services throws a `Canceled` error when the kernel is interrupted.
446-
// Or even when the kernel dies when running a cell with the code `os.kill(os.getpid(), 9)`
447-
traceError('Error in waiting for cell to complete', ex);
466+
if (!this.disposed && !this.cancelRequested) {
467+
// @jupyterlab/services throws a `Canceled` error when the kernel is interrupted.
468+
// Or even when the kernel dies when running a cell with the code `os.kill(os.getpid(), 9)`
469+
traceError('Error in waiting for cell to complete', ex);
470+
}
448471
traceCellMessage(this.cell, 'Some other execution error');
449472
if (ex && ex instanceof Error && isCancellationError(ex, true)) {
450473
// No point displaying the error stack trace from Jupyter npm package.

‎src/kernels/execution/cellExecutionMessageHandlerService.ts

Lines changed: 31 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -50,6 +50,7 @@ export class CellExecutionMessageHandlerService {
5050
kernel: Kernel.IKernelConnection;
5151
request: Kernel.IShellFuture<KernelMessage.IExecuteRequestMsg, KernelMessage.IExecuteReplyMsg>;
5252
cellExecution: NotebookCellExecution;
53+
startTime: number;
5354
onErrorHandlingExecuteRequestIOPubMessage: (error: Error) => void;
5455
}
5556
): CellExecutionMessageHandler {
@@ -59,9 +60,37 @@ export class CellExecutionMessageHandlerService {
5960
this.workspaceStorage
6061
.update(`LAST_EXECUTED_CELL_${cell.notebook.uri.toString()}`, {
6162
index: cell.index,
62-
msg_id: options.request?.msg.header.msg_id
63+
msg_id: options.request?.msg.header.msg_id,
64+
startTime: options.startTime
6365
})
6466
.then(noop, noop);
67+
const iopubMessageHandler = (_: unknown, msg: KernelMessage.IIOPubMessage<KernelMessage.IOPubMessageType>) => {
68+
if (
69+
'execution_count' in msg.content &&
70+
typeof msg.content.execution_count === 'number' &&
71+
'msg_id' in msg.parent_header &&
72+
msg.parent_header.msg_id === options.request.msg.header.msg_id
73+
) {
74+
const currentInfo = this.workspaceStorage.get<
75+
| {
76+
index: number;
77+
msg_id: string;
78+
startTime: number;
79+
execution_count: number;
80+
}
81+
| undefined
82+
>(`LAST_EXECUTED_CELL_${cell.notebook.uri.toString()}`, undefined);
83+
if (currentInfo?.msg_id === options.request.msg.header.msg_id) {
84+
this.workspaceStorage
85+
.update(`LAST_EXECUTED_CELL_${cell.notebook.uri.toString()}`, {
86+
...currentInfo,
87+
execution_count: msg.content.execution_count
88+
})
89+
.then(noop, noop);
90+
}
91+
}
92+
};
93+
options.kernel.iopubMessage.connect(iopubMessageHandler);
6594
const handler = new CellExecutionMessageHandler(
6695
cell,
6796
this.appShell,
@@ -76,6 +105,7 @@ export class CellExecutionMessageHandlerService {
76105
// This object must be kept in memory has it monitors the kernel messages.
77106
this.messageHandlers.set(cell, handler);
78107
handler.completed.finally(() => {
108+
options.kernel.iopubMessage.disconnect(iopubMessageHandler);
79109
const info = this.workspaceStorage.get<
80110
| {
81111
index: number;

‎src/kernels/execution/cellExecutionQueue.ts

Lines changed: 18 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,7 @@
11
// Copyright (c) Microsoft Corporation.
22
// Licensed under the MIT License.
33

4-
import { Disposable, EventEmitter, NotebookCell } from 'vscode';
4+
import { CancellationToken, Disposable, EventEmitter, NotebookCell } from 'vscode';
55
import { traceError, traceVerbose, traceWarning } from '../../platform/logging';
66
import { noop } from '../../platform/common/utils/misc';
77
import { traceCellMessage } from './helpers';
@@ -81,13 +81,28 @@ export class CellExecutionQueue implements Disposable {
8181
/**
8282
* Queue the cell for execution & start processing it immediately.
8383
*/
84-
public resumeCell(cell: NotebookCell, msg_id: string): void {
84+
public resumeCell(
85+
cell: NotebookCell,
86+
msg_id: string,
87+
token: CancellationToken,
88+
startTime: number,
89+
executionCount: number
90+
): void {
8591
const existingCellExecution = this.queueOfCellsToExecute.find((item) => item.cell === cell);
8692
if (existingCellExecution) {
8793
traceCellMessage(cell, 'Use existing cell execution');
8894
return;
8995
}
90-
const cellExecution = this.executionFactory.create(cell, '', this.metadata, msg_id);
96+
const cellExecution = this.executionFactory.create(
97+
cell,
98+
'',
99+
this.metadata,
100+
msg_id,
101+
false,
102+
token,
103+
startTime,
104+
executionCount
105+
);
91106
this.disposables.push(cellExecution);
92107
this.queueOfCellsToExecute.push(cellExecution);
93108

‎src/kernels/jupyter/launcher/notebookProvider.ts

Lines changed: 4 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -15,7 +15,7 @@ import {
1515
import { Cancellation } from '../../../platform/common/cancellation';
1616
import { DisplayOptions } from '../../displayOptions';
1717
import { IRawNotebookProvider } from '../../raw/types';
18-
import { IJupyterNotebookProvider, IJupyterServerUriStorage } from '../types';
18+
import { IJupyterNotebookProvider } from '../types';
1919
import { PythonExtensionNotInstalledError } from '../../../platform/errors/pythonExtNotInstalledError';
2020

2121
/**
@@ -30,8 +30,7 @@ export class NotebookProvider implements INotebookProvider {
3030
private readonly rawNotebookProvider: IRawNotebookProvider | undefined,
3131
@inject(IJupyterNotebookProvider)
3232
private readonly jupyterNotebookProvider: IJupyterNotebookProvider,
33-
@inject(IPythonExtensionChecker) private readonly extensionChecker: IPythonExtensionChecker,
34-
@inject(IJupyterServerUriStorage) private readonly uriStorage: IJupyterServerUriStorage
33+
@inject(IPythonExtensionChecker) private readonly extensionChecker: IPythonExtensionChecker
3534
) {}
3635

3736
// Attempt to connect to our server provider, and if we do, return the connection info
@@ -48,11 +47,11 @@ export class NotebookProvider implements INotebookProvider {
4847
options.ui = this.startupUi;
4948
if (this.rawNotebookProvider?.isSupported && options.localJupyter) {
5049
throw new Error('Connect method should not be invoked for local Connections when Raw is supported');
51-
} else if (this.extensionChecker.isPythonExtensionInstalled || !this.uriStorage.isLocalLaunch) {
50+
} else if (this.extensionChecker.isPythonExtensionInstalled || !options.localJupyter) {
5251
return this.jupyterNotebookProvider.connect(options).finally(() => handler.dispose());
5352
} else {
5453
handler.dispose();
55-
if (!this.startupUi.disableUI) {
54+
if (!this.startupUi.disableUI && options.localJupyter) {
5655
await this.extensionChecker.showPythonExtensionInstallRequiredPrompt();
5756
}
5857
throw new PythonExtensionNotInstalledError();

‎src/kernels/jupyter/launcher/notebookProvider.unit.test.ts

Lines changed: 1 addition & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -8,7 +8,7 @@ import { PythonExtensionChecker } from '../../../platform/api/pythonApi';
88
import { IJupyterKernelConnectionSession, KernelConnectionMetadata } from '../../types';
99
import { NotebookProvider } from './notebookProvider';
1010
import { DisplayOptions } from '../../displayOptions';
11-
import { IJupyterNotebookProvider, IJupyterServerUriStorage } from '../types';
11+
import { IJupyterNotebookProvider } from '../types';
1212
import { IRawNotebookProvider } from '../../raw/types';
1313
import { IDisposable } from '../../../platform/common/types';
1414
import { disposeAllDisposables } from '../../../platform/common/helpers';
@@ -32,17 +32,13 @@ suite('NotebookProvider', () => {
3232
when(rawNotebookProvider.isSupported).thenReturn(false);
3333
const extensionChecker = mock(PythonExtensionChecker);
3434
when(extensionChecker.isPythonExtensionInstalled).thenReturn(true);
35-
const uriStorage = mock<IJupyterServerUriStorage>();
36-
when(uriStorage.isLocalLaunch).thenReturn(true);
3735
const onDidChangeEvent = new vscode.EventEmitter<void>();
3836
disposables.push(onDidChangeEvent);
39-
when(uriStorage.onDidChangeConnectionType).thenReturn(onDidChangeEvent.event);
4037

4138
notebookProvider = new NotebookProvider(
4239
instance(rawNotebookProvider),
4340
instance(jupyterNotebookProvider),
4441
instance(extensionChecker),
45-
instance(uriStorage)
4642
);
4743
});
4844
teardown(() => disposeAllDisposables(disposables));

0 commit comments

Comments
 (0)