Skip to content

Do not await before queuing cells for execution #15841

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.

Sign up for GitHub

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

Merged
merged 1 commit into from
Jul 29, 2024
Merged

Do not await before queuing cells for execution #15841

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
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
28 changes: 6 additions & 22 deletions src/kernels/kernelExecution.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -21,7 +21,6 @@ import { CellExecutionMessageHandlerService } from './execution/cellExecutionMes
import { CellExecutionQueue } from './execution/cellExecutionQueue';
import { cellOutputToVSCCellOutput, traceCellMessage } from './execution/helpers';
import { executeSilently } from './helpers';
import { initializeInteractiveOrNotebookTelemetryBasedOnUserAction } from './telemetry/helper';
import { sendKernelTelemetryEvent } from './telemetry/sendKernelTelemetryEvent';
import {
IKernel,
Expand Down Expand Up @@ -140,10 +139,6 @@ export class NotebookKernelExecution implements INotebookKernelExecution {
return;
}

await initializeInteractiveOrNotebookTelemetryBasedOnUserAction(
this.kernel.resourceUri,
this.kernel.kernelConnectionMetadata
);
sendKernelTelemetryEvent(this.kernel.resourceUri, Telemetry.ResumeCellExecution);
const sessionPromise = this.kernel.start(new DisplayOptions(false));
const executionQueue = this.getOrCreateCellExecutionQueue(cell.notebook, sessionPromise);
Expand All @@ -166,29 +161,24 @@ export class NotebookKernelExecution implements INotebookKernelExecution {
return;
}

traceCellMessage(cell, `NotebookKernelExecution.executeCell, ${getDisplayPath(cell.notebook.uri)}`);
await initializeInteractiveOrNotebookTelemetryBasedOnUserAction(
this.kernel.resourceUri,
this.kernel.kernelConnectionMetadata
);
const sessionPromise = this.kernel.start(new DisplayOptions(false));

traceCellMessage(cell, `NotebookKernelExecution.executeCell (2), ${getDisplayPath(cell.notebook.uri)}`);
// If we're restarting, wait for it to finish
await this.kernel.restarting;
const sessionPromise = this.kernel.restarting.then(() => this.kernel.start(new DisplayOptions(false)));

traceCellMessage(cell, `NotebookKernelExecution.executeCell (2), ${getDisplayPath(cell.notebook.uri)}`);
traceCellMessage(cell, `NotebookKernelExecution.executeCell (3), ${getDisplayPath(cell.notebook.uri)}`);
const executionQueue = this.getOrCreateCellExecutionQueue(cell.notebook, sessionPromise);
executionQueue.queueCell(cell, codeOverride);
let success = true;
try {
traceCellMessage(cell, `NotebookKernelExecution.executeCell (4), ${getDisplayPath(cell.notebook.uri)}`);
await executionQueue.waitForCompletion(cell);
} catch (ex) {
success = false;
throw ex;
} finally {
traceCellMessage(
cell,
`NotebookKernelExecution.executeCell completed (3), ${getDisplayPath(cell.notebook.uri)}`
`NotebookKernelExecution.executeCell completed (5), ${getDisplayPath(cell.notebook.uri)}`
);
logger.trace(`Cell ${cell.index} executed ${success ? 'successfully' : 'with an error'}`);
sendKernelTelemetryEvent(this.kernel.resourceUri, Telemetry.ExecuteCell, {
Expand All @@ -206,14 +196,8 @@ export class NotebookKernelExecution implements INotebookKernelExecution {
token: CancellationToken
): AsyncGenerator<NotebookCellOutput, void, unknown> {
const stopWatch = new StopWatch();
await initializeInteractiveOrNotebookTelemetryBasedOnUserAction(
this.kernel.resourceUri,
this.kernel.kernelConnectionMetadata
);
const sessionPromise = this.kernel.start(new DisplayOptions(false));

// If we're restarting, wait for it to finish
await this.kernel.restarting;
const sessionPromise = this.kernel.restarting.then(() => this.kernel.start(new DisplayOptions(false)));

const executionQueue = this.getOrCreateCellExecutionQueue(this.notebook, sessionPromise);

Expand Down
88 changes: 71 additions & 17 deletions src/notebooks/controllers/kernelConnector.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -171,7 +171,10 @@ export class KernelConnector {
{
kernel: Deferred<{
kernel: IKernel;
deadKernelAction?: 'deadKernelWasRestarted' | 'deadKernelWasNoRestarted';
deadKernelAction?:
| 'deadKernelNeedsToBeRestarted'
| 'deadKernelWasNoRestarted'
| 'deadKernelWasRestarted';
}>;
options: IDisplayOptions;
}
Expand All @@ -181,7 +184,10 @@ export class KernelConnector {
{
kernel: Deferred<{
kernel: IBaseKernel;
deadKernelAction?: 'deadKernelWasRestarted' | 'deadKernelWasNoRestarted';
deadKernelAction?:
| 'deadKernelNeedsToBeRestarted'
| 'deadKernelWasNoRestarted'
| 'deadKernelWasRestarted';
}>;
options: IDisplayOptions;
}
Expand All @@ -194,17 +200,25 @@ export class KernelConnector {
promise:
| Promise<{
kernel: IBaseKernel;
deadKernelAction?: 'deadKernelWasRestarted' | 'deadKernelWasNoRestarted';
deadKernelAction?:
| 'deadKernelNeedsToBeRestarted'
| 'deadKernelWasNoRestarted'
| 'deadKernelWasRestarted';
}>
| Promise<{
kernel: IKernel;
deadKernelAction?: 'deadKernelWasRestarted' | 'deadKernelWasNoRestarted';
deadKernelAction?:
| 'deadKernelNeedsToBeRestarted'
| 'deadKernelWasNoRestarted'
| 'deadKernelWasRestarted';
}>,
actionSource: KernelActionSource,
onAction: (action: KernelAction, kernel: IBaseKernel) => void,
disposables: IDisposable[]
disposables: IDisposable[],
kernelWasRestartedDueToDeadKernel?: { kernelWasRestartedDueToDeadKernel: Deferred<boolean> }
): Promise<IKernel | IBaseKernel> {
const { kernel, deadKernelAction } = await promise;
const info = await promise;
const { kernel, deadKernelAction } = info;
// Before returning, but without disposing the kernel, double check it's still valid
// If a restart didn't happen, then we can't connect. Throw an error.
// Do this outside of the loop so that subsequent calls will still ask because the kernel isn't disposed
Expand All @@ -215,9 +229,17 @@ export class KernelConnector {
if (deadKernelAction === 'deadKernelWasNoRestarted') {
throw new KernelDeadError(kernel.kernelConnectionMetadata);
} else if (deadKernelAction === 'deadKernelWasRestarted') {
return kernel;
throw new KernelDeadError(kernel.kernelConnectionMetadata);
} else if (kernelWasRestartedDueToDeadKernel?.kernelWasRestartedDueToDeadKernel.value === true) {
throw new KernelDeadError(kernel.kernelConnectionMetadata);
}
info.deadKernelAction = 'deadKernelWasRestarted';
// Kernel is dead and we didn't prompt the user to restart it, hence re-run the code that will prompt the user for a restart.
const deferred = createDeferred<boolean>();
deferred.resolve(true);
kernelWasRestartedDueToDeadKernel = kernelWasRestartedDueToDeadKernel || {
kernelWasRestartedDueToDeadKernel: deferred
};
return KernelConnector.wrapKernelMethod(
kernel.kernelConnectionMetadata,
'start',
Expand All @@ -226,7 +248,8 @@ export class KernelConnector {
notebookResource,
options,
disposables,
onAction
onAction,
kernelWasRestartedDueToDeadKernel
);
}
return kernel;
Expand All @@ -240,7 +263,8 @@ export class KernelConnector {
notebookResource: NotebookResource,
options: IDisplayOptions,
disposables: IDisposable[],
onAction: (action: KernelAction, kernel: IBaseKernel | IKernel) => void = () => noop()
onAction: (action: KernelAction, kernel: IBaseKernel | IKernel) => void = () => noop(),
kernelWasRestartedDueToDeadKernel?: { kernelWasRestartedDueToDeadKernel: Deferred<boolean> }
): Promise<IBaseKernel | IKernel> {
let currentPromise = this.getKernelInfo(notebookResource);
if (!options.disableUI && currentPromise?.options.disableUI) {
Expand Down Expand Up @@ -276,15 +300,22 @@ export class KernelConnector {
)}`
);

kernelWasRestartedDueToDeadKernel = kernelWasRestartedDueToDeadKernel || {
kernelWasRestartedDueToDeadKernel: createDeferred<boolean>()
};
const promise = KernelConnector.wrapKernelMethodImpl(
metadata,
initialContext,
serviceContainer,
notebookResource,
options,
actionSource,
onAction
onAction,
kernelWasRestartedDueToDeadKernel
);
if (kernelWasRestartedDueToDeadKernel?.kernelWasRestartedDueToDeadKernel.resolved) {
kernelWasRestartedDueToDeadKernel.kernelWasRestartedDueToDeadKernel.resolve(true);
}
const deferred = createDeferredFromPromise(promise);
deferred.promise.catch(noop);
// If the kernel gets disposed or we fail to create the kernel, then ensure we remove the cached result.
Expand All @@ -310,7 +341,8 @@ export class KernelConnector {
deferred.promise,
actionSource,
onAction,
disposables
disposables,
kernelWasRestartedDueToDeadKernel
);
}
private static getKernelInfo(notebookResource: NotebookResource) {
Expand All @@ -322,15 +354,23 @@ export class KernelConnector {
notebookResource: NotebookResource,
deferred: Deferred<{
kernel: IBaseKernel;
deadKernelAction?: 'deadKernelWasRestarted' | 'deadKernelWasNoRestarted' | undefined;
deadKernelAction?:
| 'deadKernelNeedsToBeRestarted'
| 'deadKernelWasNoRestarted'
| 'deadKernelWasRestarted'
| undefined;
}>,
options: IDisplayOptions
) {
if ('notebook' in notebookResource) {
KernelConnector.connectionsByNotebook.set(notebookResource.notebook, {
kernel: deferred as Deferred<{
kernel: IKernel;
deadKernelAction?: 'deadKernelWasRestarted' | 'deadKernelWasNoRestarted' | undefined;
deadKernelAction?:
| 'deadKernelNeedsToBeRestarted'
| 'deadKernelWasNoRestarted'
| 'deadKernelWasRestarted'
| undefined;
}>,
options
});
Expand All @@ -342,7 +382,11 @@ export class KernelConnector {
notebookResource: NotebookResource,
matchingKernelPromise?: Promise<{
kernel: IBaseKernel;
deadKernelAction?: 'deadKernelWasRestarted' | 'deadKernelWasNoRestarted' | undefined;
deadKernelAction?:
| 'deadKernelNeedsToBeRestarted'
| 'deadKernelWasNoRestarted'
| 'deadKernelWasRestarted'
| undefined;
}>
) {
if (!matchingKernelPromise) {
Expand Down Expand Up @@ -385,10 +429,11 @@ export class KernelConnector {
notebookResource: NotebookResource,
options: IDisplayOptions,
actionSource: KernelActionSource,
onAction: (action: KernelAction, kernel: IBaseKernel) => void
onAction: (action: KernelAction, kernel: IBaseKernel) => void,
kernelWasRestartedDueToDeadKernel?: { kernelWasRestartedDueToDeadKernel: Deferred<boolean> }
): Promise<{
kernel: IBaseKernel | IKernel;
deadKernelAction?: 'deadKernelWasRestarted' | 'deadKernelWasNoRestarted';
deadKernelAction?: 'deadKernelNeedsToBeRestarted' | 'deadKernelWasNoRestarted';
}> {
const kernelProvider = serviceContainer.get<IKernelProvider>(IKernelProvider);
const thirdPartyKernelProvider = serviceContainer.get<IThirdPartyKernelProvider>(IThirdPartyKernelProvider);
Expand Down Expand Up @@ -416,15 +461,20 @@ export class KernelConnector {
resourceUri: notebookResource.resource
});

let attemptedToRestart = false;
try {
// If the kernel is dead, ask the user if they want to restart.
// We need to perform this check first, as its possible we'd call this method for dead kernels.
// & if the kernel is dead, prompt to restart.
if (initialContext !== 'restart' && isKernelDead(kernel) && !options.disableUI) {
attemptedToRestart = true;
const restarted = await KernelConnector.notifyAndRestartDeadKernel(kernel);
if (restarted && kernelWasRestartedDueToDeadKernel) {
kernelWasRestartedDueToDeadKernel.kernelWasRestartedDueToDeadKernel.resolve(true);
}
return {
kernel,
deadKernelAction: restarted ? 'deadKernelWasRestarted' : 'deadKernelWasNoRestarted'
deadKernelAction: restarted ? 'deadKernelNeedsToBeRestarted' : 'deadKernelWasNoRestarted'
};
} else {
onAction(currentContext, kernel);
Expand All @@ -440,6 +490,10 @@ export class KernelConnector {
}
}
} catch (error) {
if (attemptedToRestart && kernelWasRestartedDueToDeadKernel) {
kernelWasRestartedDueToDeadKernel.kernelWasRestartedDueToDeadKernel.resolve(true);
}

if (!isCancellationError(error)) {
logger.warn(
`Error occurred while trying to ${currentContext} the kernel, options.disableUI=${options.disableUI}`,
Expand Down
3 changes: 2 additions & 1 deletion src/notebooks/controllers/kernelConnector.unit.test.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -29,6 +29,7 @@ import { Disposable } from 'vscode';
import { PythonExtension } from '@vscode/python-extension';
import { setPythonApi } from '../../platform/interpreter/helpers';
import { resolvableInstance } from '../../test/datascience/helpers';
import { noop } from '../../test/core';

suite('Kernel Connector', () => {
const pythonConnection = PythonKernelConnectionMetadata.create({
Expand Down Expand Up @@ -163,7 +164,7 @@ suite('Kernel Connector', () => {
new DisplayOptions(false),
disposables,
'jupyterExtension'
);
).catch(noop);

verify(kernel.restart()).once();
verify(
Expand Down
1 change: 0 additions & 1 deletion src/notebooks/controllers/types.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -38,7 +38,6 @@ export interface IVSCodeNotebookController extends IDisposable {
postMessage(message: any, editor?: vscode.NotebookEditor): Thenable<boolean>;
asWebviewUri(localResource: vscode.Uri): vscode.Uri;
isAssociatedWithDocument(notebook: vscode.NotebookDocument): boolean;
restoreConnection(notebook: vscode.NotebookDocument): Promise<void>;
updateConnection(connection: KernelConnectionMetadata): void;
setPendingCellAddition(notebook: vscode.NotebookDocument, promise: Promise<void>): void;
}
Expand Down
Loading
Loading