Add stacked Python snapshot tests (#5206)

Author: hoodmane

build/wd_test.bzl

@@ -113,7 +113,8 @@ SH_RUNTEST_NORMAL = """"$@" -dTEST_TMPDIR=$TEST_TMPDIR"""
 # invoke workerd twice for snapshot tests.
 
 WINDOWS_RUNTEST_SNAPSHOT = """
-powershell -Command \"%* --python-save-snapshot $ENV:PYTHON_SAVE_SNAPSHOT_ARGS\" `-dTEST_TMPDIR=$ENV:TEST_TMPDIR
+$PYTHON_SAVE_SNAPSHOT_OPTIONS = $ENV:PYTHON_SAVE_SNAPSHOT_ARGS -split ' '
+powershell -Command \"%* --python-save-snapshot @PYTHON_SAVE_SNAPSHOT_OPTIONS\" `-dTEST_TMPDIR=$ENV:TEST_TMPDIR
 set TEST_EXIT=!ERRORLEVEL!
 if !TEST_EXIT! EQU 0 (
     powershell -Command \"%* --python-load-snapshot snapshot.bin\" `-dTEST_TMPDIR=$ENV:TEST_TMPDIR

src/pyodide/internal/snapshot.ts

@@ -107,6 +107,8 @@ const CREATED_SNAPSHOT_META: DsoLoadInfo = {
   loadOrder: [],
 };
 if (LOADED_SNAPSHOT_META) {
+  // Make sure we include the soMemoryBases and loadOrder from the baseline snapshot when we are
+  // generating stacked snapshots.
   Object.assign(
     CREATED_SNAPSHOT_META.soMemoryBases,
     LOADED_SNAPSHOT_META.soMemoryBases
@@ -642,16 +644,13 @@ export function isRestoringSnapshot(): boolean {
   return !!LOADED_SNAPSHOT_META;
 }
 
-export function maybeRestoreSnapshot(Module: Module): void {
-  if (!LOADED_SNAPSHOT_META) {
+function checkSnapshotType(snapshotType: string): void {
+  if (SHOULD_SNAPSHOT_TO_DISK) {
     return;
   }
-  const { snapshotSize, snapshotOffset, snapshotReader, settings } =
-    LOADED_SNAPSHOT_META;
-
   if (
     !IS_EW_VALIDATING &&
-    settings.snapshotType === 'dedicated' &&
+    snapshotType === 'dedicated' &&
     !IS_DEDICATED_SNAPSHOT_ENABLED
   ) {
     throw new PythonRuntimeError(
@@ -661,7 +660,7 @@ export function maybeRestoreSnapshot(Module: Module): void {
 
   if (
     !IS_EW_VALIDATING &&
-    settings.snapshotType !== 'dedicated' &&
+    snapshotType !== 'dedicated' &&
     IS_DEDICATED_SNAPSHOT_ENABLED
   ) {
     throw new PythonRuntimeError(
@@ -672,12 +671,21 @@ export function maybeRestoreSnapshot(Module: Module): void {
   // If we have a snapshot in the bundle and the dedicated snapshot flag is enabled, then we
   // should verify that the snapshot in the bundle is a dedicated snapshot. If it is not
   // we should fail with an error.
-  if (settings.snapshotType !== 'dedicated' && IS_SECOND_VALIDATION_PHASE) {
+  if (snapshotType !== 'dedicated' && IS_SECOND_VALIDATION_PHASE) {
     throw new PythonRuntimeError(
       'The second validation phase should receive a dedicated snapshot, got ' +
-        settings.snapshotType
+        snapshotType
     );
   }
+}
+
+export function maybeRestoreSnapshot(Module: Module): void {
+  if (!LOADED_SNAPSHOT_META) {
+    return;
+  }
+  const { snapshotSize, snapshotOffset, snapshotReader, settings } =
+    LOADED_SNAPSHOT_META;
+  checkSnapshotType(settings.snapshotType);
 
   Module.growMemory(snapshotSize);
   snapshotReader.readMemorySnapshot(snapshotOffset, Module.HEAP8);

src/workerd/server/tests/python/BUILD.bazel

@@ -72,3 +72,21 @@ py_wd_test(
     make_snapshot = False,
     use_snapshot = "numpy",
 )
+
+py_wd_test(
+    name = "numpy_stacked_snapshot",
+    directory = "numpy",
+    feature_flags = ["python_dedicated_snapshot"],
+    make_snapshot = True,
+    python_flags = ["0.28.2"],
+    use_snapshot = "baseline",
+)
+
+py_wd_test(
+    name = "fastapi_stacked_snapshot",
+    directory = "fastapi",
+    feature_flags = ["python_dedicated_snapshot"],
+    make_snapshot = True,
+    python_flags = ["0.28.2"],
+    use_snapshot = "baseline",
+)

src/workerd/server/tests/python/py_wd_test.bzl

@@ -17,12 +17,13 @@ def _py_wd_test_helper(
         make_snapshot,
         use_snapshot,
         args,
+        feature_flags,
         data = [],
         **kwargs):
     name_flag = name + "_" + python_flag
     templated_src = name_flag.replace("/", "-") + "@template"
     templated_src = "/".join(src.split("/")[:-1] + [templated_src])
-    flags = _get_enable_flags(python_flag)
+    flags = _get_enable_flags(python_flag) + feature_flags
     feature_flags_txt = ",".join(['"{}"'.format(flag) for flag in flags])
 
     load_snapshot = None
@@ -106,6 +107,7 @@ def py_wd_test(
         name = None,
         python_flags = "all",
         skip_python_flags = [],
+        feature_flags = [],
         args = [],
         size = "enormous",
         tags = [],
@@ -148,6 +150,7 @@ def py_wd_test(
             python_flag,
             make_snapshot = make_snapshot,
             use_snapshot = use_snapshot,
+            feature_flags = feature_flags,
             data = data,
             args = args,
             size = size,