From 3e2cd9d9c08eb6c1a64c3cd38401ba6ee2f636c5 Mon Sep 17 00:00:00 2001 From: martinRenou Date: Fri, 7 Jun 2024 16:38:20 +0200 Subject: [PATCH] Implement shared array buffer file access --- .../pyodide-kernel-extension/src/index.ts | 3 ++ packages/pyodide-kernel/src/kernel.ts | 32 ++++++++++++++- packages/pyodide-kernel/src/tokens.ts | 9 ++++ packages/pyodide-kernel/src/worker.ts | 41 +++++++++++++++++-- 4 files changed, 80 insertions(+), 5 deletions(-) diff --git a/packages/pyodide-kernel-extension/src/index.ts b/packages/pyodide-kernel-extension/src/index.ts index 82a7fe39..65ae8a55 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/src/kernel.ts b/packages/pyodide-kernel/src/kernel.ts index 465d6eb4..30776bad 100644 --- a/packages/pyodide-kernel/src/kernel.ts +++ b/packages/pyodide-kernel/src/kernel.ts @@ -3,13 +3,14 @@ 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,9 +26,31 @@ export class PyodideKernel extends BaseKernel implements IKernel { this._worker = this.initWorker(options); this._worker.onmessage = (e) => this._processWorkerMessage(e.data); this._remoteKernel = coincident(this._worker) as IPyodideWorkerKernel; + this._contentsManager = options.contentsManager; + this.setupFilesystemAPIs(); this.initRemote(options); } + 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); + }; + } + /** * Load the worker. * @@ -288,6 +311,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 +359,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; } /** diff --git a/packages/pyodide-kernel/src/worker.ts b/packages/pyodide-kernel/src/worker.ts index e17ba0cf..86bf73c0 100644 --- a/packages/pyodide-kernel/src/worker.ts +++ b/packages/pyodide-kernel/src/worker.ts @@ -3,10 +3,44 @@ import type Pyodide from 'pyodide'; -import type { DriveFS } from '@jupyterlite/contents'; +import { ContentsAPI, DriveFS, ServiceWorkerContentsAPI, TDriveMethod, TDriveRequest, TDriveResponse } from '@jupyterlite/contents'; import type { IPyodideWorkerKernel } from './tokens'; + +/** + * 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 PyodideRemoteKernel { constructor() { this._initialized = new Promise((resolve, reject) => { @@ -132,9 +166,8 @@ export class PyodideRemoteKernel { const mountpoint = '/drive'; const { FS, PATH, ERRNO_CODES } = this._pyodide; const { baseUrl } = options; - const { DriveFS } = await import('@jupyterlite/contents'); - const driveFS = new DriveFS({ + const driveFS = new PyodideDriveFS({ FS, PATH, ERRNO_CODES, @@ -505,5 +538,5 @@ export class PyodideRemoteKernel { protected _stdout_stream: any; protected _stderr_stream: any; protected _resolveInputReply: any; - protected _driveFS: DriveFS | null = null; + protected _driveFS: PyodideDriveFS | null = null; }