Skip to content

Commit dba5e37

Browse files
rchiodorchiodo
authored
Fix multiprocessing problems with setting __file__ (microsoft#14376)
* Fix multiprocessing problems with setting __file__ * Update news entry * Problem with wait for idle not propagating outwards
1 parent 0a1bbbd commit dba5e37

5 files changed

Lines changed: 53 additions & 15 deletions

File tree

‎news/2 Fixes/12530.md‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1 @@
1+
Make sure not to set ```__file__``` unless necessary as this can mess up some modules (like multiprocessing)

‎src/client/datascience/interactive-common/interactiveBase.ts‎

Lines changed: 6 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -123,7 +123,7 @@ export abstract class InteractiveBase extends WebviewPanelHost<IInteractiveWindo
123123
protected abstract get notebookMetadata(): INotebookMetadataLive | undefined;
124124

125125
protected abstract get notebookIdentity(): INotebookIdentity;
126-
126+
protected fileInKernel: string | undefined;
127127
private unfinishedCells: ICell[] = [];
128128
private restartingKernel: boolean = false;
129129
private perceivedJupyterStartupTelemetryCaptured: boolean = false;
@@ -538,6 +538,8 @@ export abstract class InteractiveBase extends WebviewPanelHost<IInteractiveWindo
538538

539539
protected abstract updateNotebookOptions(kernelConnection: KernelConnectionMetadata): Promise<void>;
540540

541+
protected abstract setFileInKernel(file: string, cancelToken: CancellationToken | undefined): Promise<void>;
542+
541543
protected async clearResult(id: string): Promise<void> {
542544
await this.ensureConnectionAndNotebook();
543545
if (this._notebook) {
@@ -631,16 +633,9 @@ export abstract class InteractiveBase extends WebviewPanelHost<IInteractiveWindo
631633
}
632634

633635
// If the file isn't unknown, set the active kernel's __file__ variable to point to that same file.
634-
if (file !== Identifiers.EmptyFileName) {
635-
await this._notebook.execute(
636-
`__file__ = '${file.replace(/\\/g, '\\\\')}'`,
637-
file,
638-
line,
639-
uuid(),
640-
cancelToken,
641-
true
642-
);
643-
}
636+
await this.setFileInKernel(file, cancelToken);
637+
638+
// Setup telemetry
644639
if (stopWatch && !this.perceivedJupyterStartupTelemetryCaptured) {
645640
this.perceivedJupyterStartupTelemetryCaptured = true;
646641
sendTelemetryEvent(Telemetry.PerceivedJupyterStartupNotebook, stopWatch?.elapsedTime);

‎src/client/datascience/interactive-ipynb/nativeEditor.ts‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -579,6 +579,10 @@ export class NativeEditor extends InteractiveBase implements INotebookEditor {
579579
// Actually don't close, just let the error bubble out
580580
}
581581

582+
protected async setFileInKernel(_file: string, _cancelToken: CancellationToken | undefined): Promise<void> {
583+
// Native editor doesn't set this as the ipython file should be set for a notebook.
584+
}
585+
582586
protected async close(): Promise<void> {
583587
// Fire our event
584588
this.closedEvent.fire(this);

‎src/client/datascience/interactive-window/interactiveWindow.ts‎

Lines changed: 34 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -2,7 +2,8 @@
22
// Licensed under the MIT License.
33
import type { nbformat } from '@jupyterlab/coreutils';
44
import * as path from 'path';
5-
import { Event, EventEmitter, Memento, Uri, ViewColumn } from 'vscode';
5+
import * as uuid from 'uuid';
6+
import { CancellationToken, Event, EventEmitter, Memento, Uri, ViewColumn } from 'vscode';
67
import {
78
IApplicationShell,
89
ICommandManager,
@@ -420,6 +421,38 @@ export class InteractiveWindow extends InteractiveBase implements IInteractiveWi
420421
protected async closeBecauseOfFailure(_exc: Error): Promise<void> {
421422
this.dispose();
422423
}
424+
425+
protected async setFileInKernel(file: string, cancelToken: CancellationToken | undefined): Promise<void> {
426+
// If in perFile mode, set only once
427+
if (this.mode === 'perFile' && !this.fileInKernel && this.notebook && file !== Identifiers.EmptyFileName) {
428+
this.fileInKernel = file;
429+
await this.notebook.execute(
430+
`__file__ = '${file.replace(/\\/g, '\\\\')}'`,
431+
file,
432+
0,
433+
uuid(),
434+
cancelToken,
435+
true
436+
);
437+
} else if (
438+
(!this.fileInKernel || !this.fs.areLocalPathsSame(this.fileInKernel, file)) &&
439+
this.mode !== 'perFile' &&
440+
this.notebook &&
441+
file !== Identifiers.EmptyFileName
442+
) {
443+
// Otherwise we need to reset it every time
444+
this.fileInKernel = file;
445+
await this.notebook.execute(
446+
`__file__ = '${file.replace(/\\/g, '\\\\')}'`,
447+
file,
448+
0,
449+
uuid(),
450+
cancelToken,
451+
true
452+
);
453+
}
454+
}
455+
423456
protected ensureConnectionAndNotebook(): Promise<void> {
424457
// Keep track of users who have used interactive window in a worksapce folder.
425458
// To be used if/when changing workflows related to startup of jupyter.

‎src/client/datascience/jupyter/jupyterSession.ts‎

Lines changed: 8 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -23,6 +23,7 @@ import { reportAction } from '../progress/decorator';
2323
import { ReportableAction } from '../progress/types';
2424
import { IJupyterConnection, ISessionWithSocket } from '../types';
2525
import { JupyterInvalidKernelError } from './jupyterInvalidKernelError';
26+
import { JupyterWaitForIdleError } from './jupyterWaitForIdleError';
2627
import { JupyterWebSockets } from './jupyterWebSocket';
2728
import { getNameOfKernelConnection } from './kernels/helpers';
2829
import { KernelConnectionMetadata } from './kernels/types';
@@ -90,9 +91,13 @@ export class JupyterSession extends BaseJupyterSession {
9091
// Make sure it is idle before we return
9192
await this.waitForIdleOnSession(newSession, timeoutMS);
9293
} catch (exc) {
93-
traceError('Failed to change kernel', exc);
94-
// Throw a new exception indicating we cannot change.
95-
throw new JupyterInvalidKernelError(kernelConnection);
94+
if (exc instanceof JupyterWaitForIdleError) {
95+
throw exc;
96+
} else {
97+
traceError('Failed to change kernel', exc);
98+
// Throw a new exception indicating we cannot change.
99+
throw new JupyterInvalidKernelError(kernelConnection);
100+
}
96101
}
97102

98103
return newSession;

0 commit comments

Comments
 (0)