From ae20dc725fccc990c87253e659ed6aa1eba5969b Mon Sep 17 00:00:00 2001 From: martinRenou Date: Fri, 21 Jun 2024 08:23:55 +0100 Subject: [PATCH] FileSystem calls over Atomics.wait instead of service worker when available (#114) * Implement shared array buffer file access * Linter * Iterate * Reduce diff * Debug --- .../pyodide-kernel-extension/src/index.ts | 3 + packages/pyodide-kernel/package.json | 2 +- .../pyodide-kernel/src/coincident.worker.ts | 74 ++++++++++++++++++- packages/pyodide-kernel/src/kernel.ts | 36 ++++++++- packages/pyodide-kernel/src/tokens.ts | 9 +++ 5 files changed, 120 insertions(+), 4 deletions(-) diff --git a/packages/pyodide-kernel-extension/src/index.ts b/packages/pyodide-kernel-extension/src/index.ts index 630dffa2..1909c897 100644 --- a/packages/pyodide-kernel-extension/src/index.ts +++ b/packages/pyodide-kernel-extension/src/index.ts @@ -41,6 +41,8 @@ const kernel: JupyterLiteServerPlugin = { serviceWorker?: IServiceWorkerManager, broadcastChannel?: IBroadcastChannelWrapper, ) => { + const contentsManager = app.serviceManager.contents; + const config = JSON.parse(PageConfig.getOption('litePluginSettings') || '{}')[PLUGIN_ID] || {}; @@ -93,6 +95,7 @@ const kernel: JupyterLiteServerPlugin = { disablePyPIFallback, mountDrive, loadPyodideOptions, + contentsManager, }); }, }); diff --git a/packages/pyodide-kernel/package.json b/packages/pyodide-kernel/package.json index fe6d4528..3d6dcf1d 100644 --- a/packages/pyodide-kernel/package.json +++ b/packages/pyodide-kernel/package.json @@ -37,7 +37,7 @@ "build:lib": "tsc -b", "build:prod": "jlpm build", "build:py": "python scripts/generate-wheels-js.py", - "build:worker": "esbuild --bundle --minify --sourcemap --target=es2019 --format=esm --outfile=lib/worker.js src/worker.ts", + "build:worker": "esbuild --bundle --minify --sourcemap --target=es2019 --format=esm --outfile=lib/coincident.worker.js src/coincident.worker.ts", "dist": "cd ../../dist && npm pack ../packages/pyodide-kernel", "clean": "jlpm clean:lib && jlpm clean:py", "clean:all": "jlpm clean", diff --git a/packages/pyodide-kernel/src/coincident.worker.ts b/packages/pyodide-kernel/src/coincident.worker.ts index 35d3321a..78bc10ca 100644 --- a/packages/pyodide-kernel/src/coincident.worker.ts +++ b/packages/pyodide-kernel/src/coincident.worker.ts @@ -6,13 +6,83 @@ */ import coincident from 'coincident'; +import { + ContentsAPI, + DriveFS, + ServiceWorkerContentsAPI, + TDriveMethod, + TDriveRequest, + TDriveResponse, +} from '@jupyterlite/contents'; + import { PyodideRemoteKernel } from './worker'; import { IPyodideWorkerKernel } from './tokens'; -const worker = new PyodideRemoteKernel(); - const workerAPI: IPyodideWorkerKernel = coincident(self) as IPyodideWorkerKernel; +/** + * An Emscripten-compatible synchronous Contents API using shared array buffers. + */ +export class SharedBufferContentsAPI extends ContentsAPI { + request(data: TDriveRequest): TDriveResponse { + return workerAPI.processDriveRequest(data); + } +} + +/** + * A custom drive implementation which uses shared array buffers if available, service worker otherwise + */ +class PyodideDriveFS extends DriveFS { + createAPI(options: DriveFS.IOptions): ContentsAPI { + if (crossOriginIsolated) { + return new SharedBufferContentsAPI( + options.driveName, + options.mountpoint, + options.FS, + options.ERRNO_CODES, + ); + } else { + return new ServiceWorkerContentsAPI( + options.baseUrl, + options.driveName, + options.mountpoint, + options.FS, + options.ERRNO_CODES, + ); + } + } +} + +export class PyodideCoincidentKernel extends PyodideRemoteKernel { + /** + * Setup custom Emscripten FileSystem + */ + protected async initFilesystem( + options: IPyodideWorkerKernel.IOptions, + ): Promise { + if (options.mountDrive) { + const mountpoint = '/drive'; + const { FS, PATH, ERRNO_CODES } = this._pyodide; + const { baseUrl } = options; + + const driveFS = new PyodideDriveFS({ + FS, + PATH, + ERRNO_CODES, + baseUrl, + driveName: this._driveName, + mountpoint, + }); + FS.mkdir(mountpoint); + FS.mount(driveFS, {}, mountpoint); + FS.chdir(mountpoint); + this._driveFS = driveFS; + } + } +} + +const worker = new PyodideCoincidentKernel(); + workerAPI.initialize = worker.initialize.bind(worker); workerAPI.execute = worker.execute.bind(worker); workerAPI.complete = worker.complete.bind(worker); diff --git a/packages/pyodide-kernel/src/kernel.ts b/packages/pyodide-kernel/src/kernel.ts index ec6ca5af..5c14cb89 100644 --- a/packages/pyodide-kernel/src/kernel.ts +++ b/packages/pyodide-kernel/src/kernel.ts @@ -3,13 +3,18 @@ import coincident from 'coincident'; import { PromiseDelegate } from '@lumino/coreutils'; import { PageConfig } from '@jupyterlab/coreutils'; -import { KernelMessage } from '@jupyterlab/services'; +import { Contents, KernelMessage } from '@jupyterlab/services'; import { BaseKernel, IKernel } from '@jupyterlite/kernel'; import { IPyodideWorkerKernel, IRemotePyodideWorkerKernel } from './tokens'; import { allJSONUrl, pipliteWheelUrl } from './_pypi'; +import { + DriveContentsProcessor, + TDriveMethod, + TDriveRequest, +} from '@jupyterlite/contents'; /** * A kernel that executes Python code with Pyodide. @@ -25,6 +30,28 @@ export class PyodideKernel extends BaseKernel implements IKernel { this._worker = this.initWorker(options); this._worker.onmessage = (e) => this._processWorkerMessage(e.data); this._remoteKernel = this.initRemote(options); + this._contentsManager = options.contentsManager; + this.setupFilesystemAPIs(); + } + + private setupFilesystemAPIs() { + (this._remoteKernel.processDriveRequest as any) = async ( + data: TDriveRequest, + ) => { + if (!DriveContentsProcessor) { + throw new Error( + 'File system calls over Atomics.wait is only supported with jupyterlite>=0.4.0a3', + ); + } + + if (this._contentsProcessor === undefined) { + this._contentsProcessor = new DriveContentsProcessor({ + contentsManager: this._contentsManager, + }); + } + + return await this._contentsProcessor.processDriveRequest(data); + }; } /** @@ -288,6 +315,8 @@ export class PyodideKernel extends BaseKernel implements IKernel { return await this._remoteKernel.inputReply(content, this.parent); } + private _contentsManager: Contents.IManager; + private _contentsProcessor: DriveContentsProcessor | undefined; private _worker: Worker; private _remoteKernel: IRemotePyodideWorkerKernel; private _ready = new PromiseDelegate(); @@ -334,5 +363,10 @@ export namespace PyodideKernel { lockFileURL: string; packages: string[]; }; + + /** + * The Jupyterlite content manager + */ + contentsManager: Contents.IManager; } } diff --git a/packages/pyodide-kernel/src/tokens.ts b/packages/pyodide-kernel/src/tokens.ts index 90ae6bdf..3449f5b2 100644 --- a/packages/pyodide-kernel/src/tokens.ts +++ b/packages/pyodide-kernel/src/tokens.ts @@ -5,6 +5,7 @@ * Definitions for the Pyodide kernel. */ +import { TDriveMethod, TDriveRequest, TDriveResponse } from '@jupyterlite/contents'; import { IWorkerKernel } from '@jupyterlite/kernel'; /** @@ -20,6 +21,14 @@ export interface IPyodideWorkerKernel extends IWorkerKernel { * Handle any lazy initialization activities. */ initialize(options: IPyodideWorkerKernel.IOptions): Promise; + + /** + * Process drive request + * @param data + */ + processDriveRequest( + data: TDriveRequest, + ): TDriveResponse; } /**