Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Non-blocking background Python kernel execution #15230

Merged
merged 1 commit into from
Feb 26, 2024
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
20 changes: 20 additions & 0 deletions src/api.proposed.displayUpdate.d.ts
Original file line number Diff line number Diff line change
@@ -0,0 +1,20 @@
// Copyright (c) Microsoft Corporation.
// Licensed under the MIT License.

import { Event, NotebookCellOutput } from 'vscode';

declare module './api' {
/**
* Represents a Jupyter Kernel.
*/
export interface Kernel {
/**
* Event that fires when a display_update message is sent, one that
* belongs to a display output from a previous code execution call.
* This allows 3rd party extensions to send messages from the kernel back to the extension host
* even after the code execution has completed.
* E.g. via background threads.
*/
onDidRecieveDisplayUpdate: Event<NotebookCellOutput>;
}
}
3 changes: 2 additions & 1 deletion src/kernels/execution/cellExecutionMessageHandler.ts
Original file line number Diff line number Diff line change
Expand Up @@ -1101,7 +1101,8 @@ export class CellExecutionMessageHandler implements IDisposable {
}
const outputToBeUpdated = CellOutputDisplayIdTracker.getMappedOutput(this.cell.notebook, displayId);
if (!outputToBeUpdated) {
traceWarning('Update display data message received, but no output found to update', msg.content);
// Possible this is a display Id that was created by code executed by an extension.
traceVerbose('Update display data message received, but no output found to update', msg.content);
return;
}
if (outputToBeUpdated.cell.document.isClosed) {
Expand Down
38 changes: 25 additions & 13 deletions src/kernels/execution/codeExecution.ts
Original file line number Diff line number Diff line change
Expand Up @@ -15,8 +15,12 @@ import { ICodeExecution } from './types';
import { executeSilentlyAndEmitOutput } from '../helpers';
import { EventEmitter, NotebookCellOutput } from 'vscode';
import { KernelError } from '../errors/kernelError';
import { JVSC_EXTENSION_ID } from '../../platform/common/constants';

function traceExecMessage(executionId: string, message: string) {
function traceExecMessage(extensionId: string, executionId: string, message: string) {
if (extensionId === JVSC_EXTENSION_ID) {
return;
}
traceVerbose(`Execution Id:${executionId}. ${message}.`);
}

Expand Down Expand Up @@ -72,18 +76,22 @@ export class CodeExecution implements ICodeExecution, IDisposable {
public async start(session: IKernelSession) {
this.session = session;
if (this.cancelHandled) {
traceExecMessage(this.executionId, 'Not starting as it was cancelled');
traceExecMessage(this.extensionId, this.executionId, 'Not starting as it was cancelled');
return;
}
traceExecMessage(this.executionId, 'Start Code execution');
traceExecMessage(this.extensionId, this.executionId, 'Start Code execution');
traceInfoIfCI(`Code Exec contents ${this.code.substring(0, 50)}...`);
if (!session.kernel || session.kernel.isDisposed || session.isDisposed) {
this._done.reject(new SessionDisposedError());
return;
}

if (this.started) {
traceExecMessage(this.executionId, 'Code has already been started yet CodeExecution.Start invoked again');
traceExecMessage(
this.extensionId,
this.executionId,
'Code has already been started yet CodeExecution.Start invoked again'
);
traceError(`Code has already been started yet CodeExecution.Start invoked again ${this.executionId}`);
// TODO: Send telemetry this should never happen, if it does we have problems.
return this.done;
Expand All @@ -107,11 +115,17 @@ export class CodeExecution implements ICodeExecution, IDisposable {
// At this point the cell execution can only be stopped from kernel by interrupting it.
// stop handling execution results & the like from the kernel.
traceExecMessage(
this.extensionId,
this.executionId,
'Code is already running, interrupting and waiting for it to finish or kernel to start'
);
const kernel = this.session?.kernel;
if (kernel) {
// When Jupyter extnsion runs code internall using this,
// we do not want to interrupt the kernel as we're more in control of the kernel.
// I.e. only when we need to cancel code thats executed by 3rd party extension code
// should we interrupt the kernel (for internal code we're more careful and do not want to interrupt).
// User is always notified when 3rd party code runs, not the same for our internal code.
if (kernel && this.extensionId !== JVSC_EXTENSION_ID) {
await kernel.interrupt().catch(noop);
}
// This is the only time we cancel the request.
Expand All @@ -122,7 +136,7 @@ export class CodeExecution implements ICodeExecution, IDisposable {
if (this.cancelHandled || this._completed) {
return;
}
traceExecMessage(this.executionId, 'Execution cancelled');
traceExecMessage(this.extensionId, this.executionId, 'Execution cancelled');
this.cancelHandled = true;
this._done.resolve();
this.dispose();
Expand All @@ -137,7 +151,7 @@ export class CodeExecution implements ICodeExecution, IDisposable {
}
this.disposed = true;
if (!this._completed) {
traceExecMessage(this.executionId, 'Execution disposed');
traceExecMessage(this.extensionId, this.executionId, 'Execution disposed');
}
dispose(this.disposables);
}
Expand All @@ -148,13 +162,13 @@ export class CodeExecution implements ICodeExecution, IDisposable {
this._done.resolve();
throw error;
}
traceExecMessage(this.executionId, 'Send code for execution');
traceExecMessage(this.extensionId, this.executionId, 'Send code for execution');

const kernelConnection = session.kernel;
try {
this.started = true;
this._onRequestSent.fire();
traceExecMessage(this.executionId, `Execution Request Sent to Kernel`);
traceExecMessage(this.extensionId, this.executionId, `Execution Request 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
this.request = executeSilentlyAndEmitOutput(
Expand Down Expand Up @@ -183,10 +197,10 @@ export class CodeExecution implements ICodeExecution, IDisposable {
const response = await this.request!.done;
this._completed = true;
if (response.content.status === 'error') {
traceExecMessage(this.executionId, 'Executed with errors');
traceExecMessage(this.extensionId, this.executionId, 'Executed with errors');
this._done.reject(new KernelError(response.content));
} else {
traceExecMessage(this.executionId, 'Executed successfully');
traceExecMessage(this.extensionId, this.executionId, 'Executed successfully');
this._done.resolve();
}
} catch (ex) {
Expand All @@ -198,8 +212,6 @@ export class CodeExecution implements ICodeExecution, IDisposable {
// @jupyterlab/services throws a `Canceled` error when the kernel is interrupted.
// Or even when the kernel dies when running a cell with the code `os.kill(os.getpid(), 9)`
traceError(`Error in waiting for code ${this.executionId} to complete`, ex);
} else {
traceError(`Some other execution error for exec ${this.executionId}`, ex);
}
this._done.reject(ex);
}
Expand Down
Loading
Loading