Skip to content

Commit

Permalink
Rework how restart works for native notebooks (#6826)
Browse files Browse the repository at this point in the history 
* Rework how restart works for native notebooks

* Put back sys info message for IW

* Wrong index for final cell

* Simplify test

* Fix another failing test

* Linter

* Increase restart timeout
  • Loading branch information
@rchiodo
rchiodo authored Jul 26, 2021
1 parent 2b88aad commit f5c8929
Show file tree
Hide file tree
Showing 4 changed files with 121 additions and 50 deletions.
1 change: 1 addition & 0 deletions news/2 Fixes/5996.md
View file
Original file line number Diff line number Diff line change
@@ -0,0 +1 @@
Run all and restarting does not actually interrupt the rest of the running cells.
26 changes: 10 additions & 16 deletions src/client/datascience/jupyter/kernels/kernel.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -172,18 +172,16 @@ export class Kernel implements IKernel {
if (this.restarting) {
return this.restarting.promise;
}
if (this.notebook) {
this.restarting = createDeferred<void>();
try {
await this.notebook.restartKernel(this.launchTimeout);
await this.initializeAfterStart(SysInfoReason.Restart, notebookDocument);
this.restarting.resolve();
} catch (ex) {
this.restarting.reject(ex);
} finally {
this.restarting = undefined;
}
}
traceInfo(`Restart requested ${notebookDocument.uri}`);
this.startCancellation.cancel();
const restartPromise = this.kernelExecution.restart(notebookDocument, this._notebookPromise);
await restartPromise;

// Interactive window needs a restart sys info
await this.initializeAfterStart(SysInfoReason.Restart, notebookDocument);

// Indicate a restart occurred if it succeeds
this._onRestarted.fire();
}
private async trackNotebookCellPerceivedColdTime(
stopWatch: StopWatch,
Expand Down Expand Up @@ -326,10 +324,6 @@ export class Kernel implements IKernel {
this._notebookPromise = undefined;
this._onDisposed.fire();
});
this.notebook.onKernelRestarted(() => {
traceInfo(`Notebook Kernel restarted ${this.notebook?.identity}`);
this._onRestarted.fire();
});
this.notebook.onSessionStatusChanged(
(e) => {
traceInfo(`Notebook Session status ${this.notebook?.identity} # ${e}`);
Expand Down
79 changes: 56 additions & 23 deletions src/client/datascience/jupyter/kernels/kernelExecution.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -26,8 +26,8 @@ export class KernelExecution implements IDisposable {
private readonly documentExecutions = new WeakMap<NotebookDocument, CellExecutionQueue>();
private readonly executionFactory: CellExecutionFactory;
private readonly disposables: IDisposable[] = [];
private readonly kernelRestartHandlerAdded = new WeakSet<IKernel>();
private _interruptPromise?: Promise<InterruptResult>;
private _restartPromise?: Promise<void>;
constructor(
private readonly kernelProvider: IKernelProvider,
errorHandler: IDataScienceErrorHandler,
Expand All @@ -47,6 +47,12 @@ export class KernelExecution implements IDisposable {
return;
}
sendKernelTelemetryEvent(cell.notebook.uri, Telemetry.ExecuteNativeCell);

// If we're restarting, wait for it to finish
if (this._restartPromise) {
await this._restartPromise;
}

const executionQueue = this.getOrCreateCellExecutionQueue(cell.notebook, notebookPromise);
executionQueue.queueCell(cell);
await executionQueue.waitForCompletion([cell]);
Expand All @@ -56,6 +62,12 @@ export class KernelExecution implements IDisposable {
@captureTelemetry(VSCodeNativeTelemetry.RunAllCells, undefined, true)
public async executeAllCells(notebookPromise: Promise<INotebook>, document: NotebookDocument): Promise<void> {
sendKernelTelemetryEvent(document.uri, Telemetry.ExecuteNativeCell);

// If we're restarting, wait for it to finish
if (this._restartPromise) {
await this._restartPromise;
}

// Only run code cells that are not already running.
const cellsThatWeCanRun = document.getCells().filter((cell) => cell.kind === NotebookCellKind.Code);
if (cellsThatWeCanRun.length === 0) {
Expand Down Expand Up @@ -111,6 +123,40 @@ export class KernelExecution implements IDisposable {

return result;
}
/**
* Restarts the kernel
* If we don't have a kernel (Jupyter Session) available, then just abort all of the cell executions.
*/
public async restart(document: NotebookDocument, notebookPromise?: Promise<INotebook>): Promise<void> {
trackKernelResourceInformation(document.uri, { restartKernel: true });
const executionQueue = this.documentExecutions.get(document);
if (!executionQueue) {
return;
}
// Possible we don't have a notebook.
const notebook = notebookPromise ? await notebookPromise.catch(() => undefined) : undefined;
traceInfo('Restart kernel execution');
// First cancel all the cells & then wait for them to complete.
// Both must happen together, we cannot just wait for cells to complete, as its possible
// that cell1 has started & cell2 has been queued. If Cell1 completes, then Cell2 will start.
// What we want is, if Cell1 completes then Cell2 should not start (it must be cancelled before hand).
const pendingCells = executionQueue.cancel().then(() => executionQueue.waitForCompletion());

if (!notebook) {
traceInfo('No notebook to interrupt');
this._restartPromise = undefined;
await pendingCells;
return;
}

// Restart the active execution
await (this._restartPromise
? this._restartPromise
: (this._restartPromise = this.restartExecution(document, notebook.session)));

// Done restarting, clear restart promise
this._restartPromise = undefined;
}
public dispose() {
this.disposables.forEach((d) => d.dispose());
}
Expand All @@ -122,9 +168,7 @@ export class KernelExecution implements IDisposable {
}

// We need to add the handler to kernel immediately (before we resolve the notebook, else its possible user hits restart or the like and we miss that event).
const wrappedNotebookPromise = this.getKernel(document)
.then((kernel) => this.addKernelRestartHandler(kernel, document))
.then(() => notebookPromise);
const wrappedNotebookPromise = this.getKernel(document).then(() => notebookPromise);

const newCellExecutionQueue = new CellExecutionQueue(
wrappedNotebookPromise,
Expand Down Expand Up @@ -221,27 +265,16 @@ export class KernelExecution implements IDisposable {
return result;
});
}
private addKernelRestartHandler(kernel: IKernel, document: NotebookDocument) {
if (this.kernelRestartHandlerAdded.has(kernel)) {
return;
}
this.kernelRestartHandlerAdded.add(kernel);
traceInfo('Hooked up kernel restart handler');
kernel.onRestarted(
() => {
// We're only interested in restarts of the kernel associated with this document.
const executionQueue = this.documentExecutions.get(document);
if (kernel !== this.kernelProvider.get(document) || !executionQueue) {
return;
}

traceInfo('Cancel all executions as Kernel was restarted');
return executionQueue.cancel(true);
},
this,
this.disposables
);
@captureTelemetry(Telemetry.RestartKernel)
@captureTelemetry(Telemetry.RestartJupyterTime)
private async restartExecution(_document: NotebookDocument, session: IJupyterSession): Promise<void> {
// Just use the internal session. Pending cells should have been canceled by the caller
return session.restart(this.interruptTimeout).catch((exc) => {
traceWarning(`Error during restart: ${exc}`);
});
}

private async getKernel(document: NotebookDocument): Promise<IKernel> {
let kernel = this.kernelProvider.get(document);
if (!kernel) {
Expand Down
65 changes: 54 additions & 11 deletions src/test/datascience/notebook/interruptRestart.vscode.test.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -13,7 +13,6 @@ import { DataScience } from '../../../client/common/utils/localize';
import { noop } from '../../../client/common/utils/misc';
import { Commands } from '../../../client/datascience/constants';
import { IKernelProvider } from '../../../client/datascience/jupyter/kernels/types';
import { traceCellMessage } from '../../../client/datascience/notebook/helpers/helpers';
import { INotebookEditorProvider } from '../../../client/datascience/types';
import { createEventHandler, getOSType, IExtensionTestApi, OSType, waitForCondition } from '../../common';
import { IS_REMOTE_NATIVE_TEST } from '../../constants';
Expand Down Expand Up @@ -156,18 +155,9 @@ suite('DataScience - VSCode Notebook - Restart/Interrupt/Cancel/Errors (slow)',
const waitForKernelToRestart = createEventHandler(kernel, 'onRestarted', disposables);
await commands.executeCommand('jupyter.notebookeditor.restartkernel').then(noop, noop);

await waitForCondition(
async () => {
traceCellMessage(cell, 'Step 8 Cell Status');
return assertVSCCellIsNotRunning(cell);
},
15_000,
'Execution not cancelled first time.'
);

// Wait for kernel to restart before we execute cells again.
traceInfo('Step 9 Wait for restart');
await waitForKernelToRestart.assertFired(15_000);
await waitForKernelToRestart.assertFired(30_000);
traceInfo('Step 10 Restarted');

// Confirm we can execute a cell (using the new kernel session).
Expand All @@ -193,6 +183,59 @@ suite('DataScience - VSCode Notebook - Restart/Interrupt/Cancel/Errors (slow)',
// KERNELPUSH
//await vscEditor.kernel!.interrupt!(vscEditor.document);
});
test('Restarting kernel during run all will skip the rest of the cells', async function () {
traceInfo('Step 1');
await insertCodeCell('print(1)', { index: 0 });
await insertCodeCell('import time\nprint(2)\ntime.sleep(60)', { index: 1 });
await insertCodeCell('print(3)', { index: 2 });
const cell = vscEditor.document.cellAt(1);
// Ensure we click `Yes` when prompted to restart the kernel.
const appShell = api.serviceContainer.get<IApplicationShell>(IApplicationShell);
const showInformationMessage = sinon
.stub(appShell, 'showInformationMessage')
.callsFake(function (message: string) {
traceInfo(`Step 2. ShowInformationMessage ${message}`);
if (message === DataScience.restartKernelMessage()) {
traceInfo(`Step 3. ShowInformationMessage & yes to restart`);
// User clicked ok to restart it.
return DataScience.restartKernelMessageYes();
}
return (appShell.showInformationMessage as any).wrappedMethod.apply(appShell, arguments);
});
disposables.push({ dispose: () => showInformationMessage.restore() });

(editorProvider.activeEditor as any).shouldAskForRestart = () => Promise.resolve(false);
traceInfo(`Step 4. Before execute`);
await runAllCellsInActiveNotebook();
traceInfo(`Step 5. After execute`);

// Wait for cell to get busy.
await waitForCondition(async () => assertVSCCellIsRunning(cell), 15_000, 'Cell not being executed');
traceInfo(`Step 6. Cell is busy`);

// Wait for ?s, and verify cell is still running.
assertVSCCellIsRunning(cell);
// Wait for some output.
await waitForTextOutputInVSCode(cell, '2', 0, false, 15_000); // Wait for 15 seconds for it to start (possibly kernel is still starting).
traceInfo(`Step 7. Cell output`);

// Restart the kernel & use event handler to check if it was restarted successfully.
const kernel = api.serviceContainer.get<IKernelProvider>(IKernelProvider).get(cell.notebook);
if (!kernel) {
throw new Error('Kernel not available');
}
const waitForKernelToRestart = createEventHandler(kernel, 'onRestarted', disposables);
await commands.executeCommand('jupyter.notebookeditor.restartkernel').then(noop, noop);

// Wait for kernel to restart before we execute cells again.
traceInfo('Step 8 Wait for restart');
await waitForKernelToRestart.assertFired(30_000);
traceInfo('Step 9 Restarted');

// Confirm last cell is empty
const lastCell = vscEditor.document.cellAt(2);
assert.equal(lastCell.outputs.length, 0, 'Last cell should not have run');
});
test('Interrupt and running cells again should only run the necessary cells', async function () {
// Interrupts on windows doesn't work well, not as well as on Unix.
// This is how Python works, hence this test is better tested on Unix OS.
Expand Down

0 comments on commit f5c8929

Please sign in to comment.