Also support creating and using baseline snapshots

src/pyodide/internal/metadata.js

@@ -5,6 +5,7 @@ import { default as ArtifactBundler } from "pyodide-internal:artifacts";
 
 export const IS_WORKERD = MetadataReader.isWorkerd();
 export const IS_TRACING = MetadataReader.isTracing();
+export const IS_CREATING_BASELINE_SNAPSHOT = MetadataReader.isCreatingBaselineSnapshot();
 export const WORKERD_INDEX_URL = PYODIDE_BUCKET.PYODIDE_PACKAGE_BUCKET_URL;
 export const REQUIREMENTS = MetadataReader.getRequirements();
 export const MAIN_MODULE_NAME = MetadataReader.getMainModule();

src/pyodide/internal/python.js

@@ -10,7 +10,8 @@ import {
 } from "pyodide-internal:setupPackages";
 import { default as TarReader } from "pyodide-internal:packages_tar_reader";
 import processScriptImports from "pyodide-internal:process_script_imports.py";
-import { MEMORY_SNAPSHOT_READER } from "pyodide-internal:metadata";
+import { MEMORY_SNAPSHOT_READER, IS_CREATING_BASELINE_SNAPSHOT } from "pyodide-internal:metadata";
+import { reportError } from "pyodide-internal:reportError";
 
 /**
  * This file is a simplified version of the Pyodide loader:
@@ -38,9 +39,7 @@ import pyodideWasmModule from "pyodide-internal:generated/pyodide.asm.wasm";
  */
 import stdlib from "pyodide-internal:generated/python_stdlib.zip";
 
-const SHOULD_UPLOAD_SNAPSHOT =
-  ArtifactBundler.isEnabled() || ArtifactBundler.isEwValidating();
-const DEDICATED_SNAPSHOT = true;
+const SHOULD_UPLOAD_SNAPSHOT = ArtifactBundler.isEnabled() || ArtifactBundler.isEwValidating();
 
 /**
  * Global variable for the memory snapshot. On the first run we stick a copy of
@@ -138,7 +137,7 @@ function prepareFileSystem(Module) {
     );
     Module.FS.mkdirTree(Module.API.config.env.HOME);
   } catch (e) {
-    console.warn(e);
+    reportError(e);
   }
 }
 
@@ -189,9 +188,13 @@ function loadDynlib(Module, path, wasmModuleData) {
  * there.
  */
 function preloadDynamicLibs(Module) {
+  let SO_FILES_TO_LOAD = SITE_PACKAGES_SO_FILES;
+  if (DSO_METADATA?.settings?.baselineSnapshot) {
+    SO_FILES_TO_LOAD = [[ '_lzma.so' ]];
+  }
   try {
     const sitePackages = getSitePackagesPath(Module);
-    for (const soFile of SITE_PACKAGES_SO_FILES) {
+    for (const soFile of SO_FILES_TO_LOAD) {
       let node = SITE_PACKAGES_INFO;
       for (const part of soFile) {
         node = node.children.get(part);
@@ -203,8 +206,8 @@ function preloadDynamicLibs(Module) {
       loadDynlib(Module, path, wasmModuleData);
     }
   } catch (e) {
-    console.log(e);
-    throw e;
+    console.warn("Error in preloadDynamicLibs");
+    reportError(e);
   }
 }
 
@@ -267,9 +270,8 @@ async function instantiateEmscriptenModule(emscriptenSettings) {
     const emscriptenModule = await p;
     return emscriptenModule;
   } catch (e) {
-    // Wish workerd has better error logging...
-    e.stack.split("\n").forEach(console.log.bind(console));
-    throw e;
+    console.warn("Error in instantiateEmscriptenModule");
+    reportError(e);
   }
 }
 
@@ -315,6 +317,10 @@ function recordDsoHandles(Module) {
     }
     dylinkInfo[name].handles.push(handle);
   }
+  dylinkInfo.settings = {};
+  if (IS_CREATING_BASELINE_SNAPSHOT) {
+    dylinkInfo.settings.baselineSnapshot = true;
+  }
   return dylinkInfo;
 }
 
@@ -370,7 +376,7 @@ function memorySnapshotDoImports(Module) {
   simpleRunPython(Module, "sysconfig.get_config_vars()");
   // Delete to avoid polluting globals
   simpleRunPython(Module, `del ${toDelete}`);
-  if (!DEDICATED_SNAPSHOT) {
+  if (IS_CREATING_BASELINE_SNAPSHOT) {
     // We've done all the imports for the baseline snapshot.
     return;
   }
@@ -411,8 +417,7 @@ function setUploadFunction(toUpload) {
       }
     } catch (e) {
       console.warn("Memory snapshot upload failed.");
-      console.warn(e);
-      throw e;
+      reportError(e);
     }
   };
 }
@@ -446,13 +451,21 @@ function encodeSnapshot(heap, dsoJSON) {
  * Decode heap and dsoJSON from the memory snapshot artifact we downloaded
  */
 function decodeSnapshot() {
-  const buf = new Uint32Array(2);
-  MEMORY_SNAPSHOT_READER.readMemorySnapshot(0, buf);
+  let buf = new Uint32Array(2);
+  let offset = 0;
+  MEMORY_SNAPSHOT_READER.readMemorySnapshot(offset, buf);
+  offset += 8;
+  let snapshotVersion = 0;
+  if (buf[0] == SNAPSHOT_MAGIC) {
+    snapshotVersion = buf[1];
+    buf = MEMORY_SNAPSHOT_READER.readMemorySnapshot(offset, buf);
+    offset += 8;
+  }
   const snapshotOffset = buf[0];
   SNAPSHOT_SIZE = MEMORY_SNAPSHOT_READER.getMemorySnapshotSize() - snapshotOffset;
   const jsonLength = buf[1];
   const jsonBuf = new Uint8Array(jsonLength);
-  MEMORY_SNAPSHOT_READER.readMemorySnapshot(8, jsonBuf);
+  MEMORY_SNAPSHOT_READER.readMemorySnapshot(offset, jsonBuf);
   DSO_METADATA = JSON.parse(new TextDecoder().decode(jsonBuf));
   READ_MEMORY = function(Module) {
     // restore memory from snapshot
@@ -503,23 +516,28 @@ function simpleRunPython(emscriptenModule, code) {
 
 let TEST_SNAPSHOT = undefined;
 (function () {
-  // Lookup memory snapshot from artifact store.
-  if (!MEMORY_SNAPSHOT_READER) {
-    // snapshots are disabled or there isn't one yet
-    return;
-  }
+  try {
+    // Lookup memory snapshot from artifact store.
+    if (!MEMORY_SNAPSHOT_READER) {
+      // snapshots are disabled or there isn't one yet
+      return;
+    }
 
-  // Simple sanity check to ensure this snapshot isn't corrupted.
-  //
-  // TODO(later): we need better detection when this is corrupted. Right now the isolate will
-  // just die.
-  const snapshotSize = MEMORY_SNAPSHOT_READER.getMemorySnapshotSize();
-  if (snapshotSize <= 100) {
-    TEST_SNAPSHOT = new Uint8Array(snapshotSize);
-    MEMORY_SNAPSHOT_READER.readMemorySnapshot(0, TEST_SNAPSHOT);
-    return;
+    // Simple sanity check to ensure this snapshot isn't corrupted.
+    //
+    // TODO(later): we need better detection when this is corrupted. Right now the isolate will
+    // just die.
+    const snapshotSize = MEMORY_SNAPSHOT_READER.getMemorySnapshotSize();
+    if (snapshotSize <= 100) {
+      TEST_SNAPSHOT = new Uint8Array(snapshotSize);
+      MEMORY_SNAPSHOT_READER.readMemorySnapshot(0, TEST_SNAPSHOT);
+      return;
+    }
+    decodeSnapshot();
+  } catch (e) {
+    console.warn("Error in top level of python.js");
+    reportError(e);
   }
-  decodeSnapshot();
 })();
 
 export async function loadPyodide(lockfile, indexURL) {

src/pyodide/internal/reportError.js

@@ -0,0 +1,5 @@
+export function reportError(e) {
+  e.stack.split("\n").forEach((s) => console.warn(s));
+  throw e;
+}
+

src/pyodide/python-entrypoint-helper.js

@@ -21,6 +21,7 @@ import {
   WORKERD_INDEX_URL,
 } from "pyodide-internal:metadata";
 import { default as ArtifactBundler } from "pyodide-internal:artifacts";
+import { reportError } from "pyodide-internal:reportError";
 
 function pyimportMainModule(pyodide) {
   if (!MAIN_MODULE_NAME.endsWith(".py")) {
@@ -155,8 +156,8 @@ function makeHandler(pyHandlerName) {
         return mainModule[pyHandlerName].callRelaxed(...args);
       });
     } catch (e) {
-      console.warn(e.stack);
-      throw e;
+      console.warn("Error in makeHandler");
+      reportError(e);
     } finally {
       args[2].waitUntil(uploadArtifacts());
     }
@@ -205,8 +206,8 @@ try {
     ArtifactBundler.storeMemorySnapshot(getMemoryToUpload());
   }
 } catch (e) {
-  console.warn(e);
-  throw e;
+  console.warn("Error in top level in python-entrypoint-helper.js");
+  reportError(e);
 }
 
 export default handlers;

src/workerd/api/pyodide/pyodide.c++

@@ -117,7 +117,7 @@ jsg::Ref<PyodideMetadataReader> makePyodideMetadataReader(Worker::Reader conf) {
   }
   return jsg::alloc<PyodideMetadataReader>(kj::mv(mainModule), names.finish(), contents.finish(),
                                            requirements.finish(), true /* isWorkerd */,
-                                           false /* isTracing */, kj::none /* memorySnapshot */);
+                                           false /* isTracing */, false /* createBaselineSnapshot */, kj::none /* memorySnapshot */);
 }
 
 } // namespace workerd::api::pyodide

src/workerd/api/pyodide/pyodide.h

@@ -36,15 +36,18 @@ class PyodideMetadataReader : public jsg::Object {
   kj::Array<kj::String> requirements;
   bool isWorkerdFlag;
   bool isTracingFlag;
+  bool createBaselineSnapshot;
   kj::Maybe<kj::Array<kj::byte>> memorySnapshot;
 
 public:
   PyodideMetadataReader(kj::String mainModule, kj::Array<kj::String> names,
                         kj::Array<kj::Array<kj::byte>> contents, kj::Array<kj::String> requirements,
                         bool isWorkerd, bool isTracing,
+                        bool createBaselineSnapshot,
                         kj::Maybe<kj::Array<kj::byte>> memorySnapshot)
       : mainModule(kj::mv(mainModule)), names(kj::mv(names)), contents(kj::mv(contents)),
         requirements(kj::mv(requirements)), isWorkerdFlag(isWorkerd), isTracingFlag(isTracing),
+        createBaselineSnapshot(createBaselineSnapshot),
         memorySnapshot(kj::mv(memorySnapshot)) {}
 
   bool isWorkerd() {
@@ -55,6 +58,10 @@ class PyodideMetadataReader : public jsg::Object {
     return this->isTracingFlag;
   }
 
+  bool isCreatingBaselineSnapshot() {
+    return createBaselineSnapshot;
+  }
+
   kj::String getMainModule() {
     return kj::str(this->mainModule);
   }
@@ -94,6 +101,7 @@ class PyodideMetadataReader : public jsg::Object {
     JSG_METHOD(getMemorySnapshotSize);
     JSG_METHOD(readMemorySnapshot);
     JSG_METHOD(disposeMemorySnapshot);
+    JSG_METHOD(isCreatingBaselineSnapshot);
   }
 
   void visitForMemoryInfo(jsg::MemoryTracker& tracker) const {
@@ -123,22 +131,24 @@ class ArtifactBundler : public jsg::Object {
         existingSnapshot(kj::mv(existingSnapshot)),
         uploadMemorySnapshotCb(kj::mv(uploadMemorySnapshotCb)),
         hasUploaded(false),
-        isValidating(false)
-        {};
+        isValidating(false),
+        createBaselineSnapshot(false) {};
 
   ArtifactBundler(kj::Maybe<kj::Array<kj::byte>> existingSnapshot)
       : storedSnapshot(kj::none),
         existingSnapshot(kj::mv(existingSnapshot)),
         uploadMemorySnapshotCb(kj::none),
         hasUploaded(false),
-        isValidating(false) {};
+        isValidating(false),
+        createBaselineSnapshot(false) {};
 
   ArtifactBundler(bool isValidating = false)
       : storedSnapshot(kj::none),
         existingSnapshot(kj::none),
         uploadMemorySnapshotCb(kj::none),
         hasUploaded(false),
-        isValidating(isValidating) {};
+        isValidating(isValidating),
+        createBaselineSnapshot(false) {};
 
   jsg::Promise<bool> uploadMemorySnapshot(jsg::Lock& js, kj::Array<kj::byte> snapshot) {
     // Prevent multiple uploads.
@@ -218,6 +228,7 @@ class ArtifactBundler : public jsg::Object {
   kj::Maybe<kj::Function<kj::Promise<bool>(kj::Array<kj::byte> snapshot)>> uploadMemorySnapshotCb;
   bool hasUploaded;
   bool isValidating;
+  bool createBaselineSnapshot;
 };