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

Author: rchiodo

news/2 Fixes/12530.md

@@ -0,0 +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

@@ -123,7 +123,7 @@ export abstract class InteractiveBase extends WebviewPanelHost<IInteractiveWindo
     protected abstract get notebookMetadata(): INotebookMetadataLive | undefined;
 
     protected abstract get notebookIdentity(): INotebookIdentity;
-
+    protected fileInKernel: string | undefined;
     private unfinishedCells: ICell[] = [];
     private restartingKernel: boolean = false;
     private perceivedJupyterStartupTelemetryCaptured: boolean = false;
@@ -538,6 +538,8 @@ export abstract class InteractiveBase extends WebviewPanelHost<IInteractiveWindo
 
     protected abstract updateNotebookOptions(kernelConnection: KernelConnectionMetadata): Promise<void>;
 
+    protected abstract setFileInKernel(file: string, cancelToken: CancellationToken | undefined): Promise<void>;
+
     protected async clearResult(id: string): Promise<void> {
         await this.ensureConnectionAndNotebook();
         if (this._notebook) {
@@ -631,16 +633,9 @@ export abstract class InteractiveBase extends WebviewPanelHost<IInteractiveWindo
                 }
 
                 // If the file isn't unknown, set the active kernel's __file__ variable to point to that same file.
-                if (file !== Identifiers.EmptyFileName) {
-                    await this._notebook.execute(
-                        `__file__ = '${file.replace(/\\/g, '\\\\')}'`,
-                        file,
-                        line,
-                        uuid(),
-                        cancelToken,
-                        true
-                    );
-                }
+                await this.setFileInKernel(file, cancelToken);
+
+                // Setup telemetry
                 if (stopWatch && !this.perceivedJupyterStartupTelemetryCaptured) {
                     this.perceivedJupyterStartupTelemetryCaptured = true;
                     sendTelemetryEvent(Telemetry.PerceivedJupyterStartupNotebook, stopWatch?.elapsedTime);

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

@@ -579,6 +579,10 @@ export class NativeEditor extends InteractiveBase implements INotebookEditor {
         // Actually don't close, just let the error bubble out
     }
 
+    protected async setFileInKernel(_file: string, _cancelToken: CancellationToken | undefined): Promise<void> {
+        // Native editor doesn't set this as the ipython file should be set for a notebook.
+    }
+
     protected async close(): Promise<void> {
         // Fire our event
         this.closedEvent.fire(this);

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

@@ -2,7 +2,8 @@
 // Licensed under the MIT License.
 import type { nbformat } from '@jupyterlab/coreutils';
 import * as path from 'path';
-import { Event, EventEmitter, Memento, Uri, ViewColumn } from 'vscode';
+import * as uuid from 'uuid';
+import { CancellationToken, Event, EventEmitter, Memento, Uri, ViewColumn } from 'vscode';
 import {
     IApplicationShell,
     ICommandManager,
@@ -420,6 +421,38 @@ export class InteractiveWindow extends InteractiveBase implements IInteractiveWi
     protected async closeBecauseOfFailure(_exc: Error): Promise<void> {
         this.dispose();
     }
+
+    protected async setFileInKernel(file: string, cancelToken: CancellationToken | undefined): Promise<void> {
+        // If in perFile mode, set only once
+        if (this.mode === 'perFile' && !this.fileInKernel && this.notebook && file !== Identifiers.EmptyFileName) {
+            this.fileInKernel = file;
+            await this.notebook.execute(
+                `__file__ = '${file.replace(/\\/g, '\\\\')}'`,
+                file,
+                0,
+                uuid(),
+                cancelToken,
+                true
+            );
+        } else if (
+            (!this.fileInKernel || !this.fs.areLocalPathsSame(this.fileInKernel, file)) &&
+            this.mode !== 'perFile' &&
+            this.notebook &&
+            file !== Identifiers.EmptyFileName
+        ) {
+            // Otherwise we need to reset it every time
+            this.fileInKernel = file;
+            await this.notebook.execute(
+                `__file__ = '${file.replace(/\\/g, '\\\\')}'`,
+                file,
+                0,
+                uuid(),
+                cancelToken,
+                true
+            );
+        }
+    }
+
     protected ensureConnectionAndNotebook(): Promise<void> {
         // Keep track of users who have used interactive window in a worksapce folder.
         // To be used if/when changing workflows related to startup of jupyter.

src/client/datascience/jupyter/jupyterSession.ts

@@ -23,6 +23,7 @@ import { reportAction } from '../progress/decorator';
 import { ReportableAction } from '../progress/types';
 import { IJupyterConnection, ISessionWithSocket } from '../types';
 import { JupyterInvalidKernelError } from './jupyterInvalidKernelError';
+import { JupyterWaitForIdleError } from './jupyterWaitForIdleError';
 import { JupyterWebSockets } from './jupyterWebSocket';
 import { getNameOfKernelConnection } from './kernels/helpers';
 import { KernelConnectionMetadata } from './kernels/types';
@@ -90,9 +91,13 @@ export class JupyterSession extends BaseJupyterSession {
             // Make sure it is idle before we return
             await this.waitForIdleOnSession(newSession, timeoutMS);
         } catch (exc) {
-            traceError('Failed to change kernel', exc);
-            // Throw a new exception indicating we cannot change.
-            throw new JupyterInvalidKernelError(kernelConnection);
+            if (exc instanceof JupyterWaitForIdleError) {
+                throw exc;
+            } else {
+                traceError('Failed to change kernel', exc);
+                // Throw a new exception indicating we cannot change.
+                throw new JupyterInvalidKernelError(kernelConnection);
+            }
         }
 
         return newSession;