Add --save-snapshot variant to all wd_py_tests (#3415)

build/wd_test.bzl

@@ -6,6 +6,7 @@ def wd_test(
         name = None,
         args = [],
         ts_deps = [],
+        python_snapshot_test = False,
         **kwargs):
     """Rule to define tests that run `workerd test` with a particular config.
 
@@ -54,6 +55,7 @@ def wd_test(
         name = name,
         data = data,
         args = args,
+        python_snapshot_test = python_snapshot_test,
         **kwargs
     )
 
@@ -69,8 +71,7 @@ if not "%SIDECAR%" == "" (
 )
 
 REM Run the actual test
-powershell -Command \"%*\" `-dTEST_TMPDIR=$ENV:TEST_TMPDIR
-set TEST_EXIT=!ERRORLEVEL!
+$(RUNTEST)
 
 REM Cleanup sidecar if it was started
 if defined SIDECAR_PID (
@@ -98,10 +99,43 @@ if [ ! -z "$(SIDECAR)" ]; then
     sleep 3
 fi
 
-# Run the actual test
+$(RUNTEST)
+"""
+
+WINDOWS_RUNTEST_NORMAL = """
+powershell -Command \"%*\" `-dTEST_TMPDIR=$ENV:TEST_TMPDIR
+set TEST_EXIT=!ERRORLEVEL!
+"""
+
+SH_RUNTEST_NORMAL = """
 "$@" -dTEST_TMPDIR=$TEST_TMPDIR
 """
 
+# We need variants of the RUN_TEST command for Python memory snapshot tests. We have to invoke
+# workerd twice, once with --python-save-snapshot to produce the snapshot and once with
+# --python-load-snapshot to use it.
+#
+# We would like to implement this in py_wd_test and not have to complicate wd_test for it, but
+# unfortunately bazel provides no way for a test to create a file that is used by another test. So
+# we cannot do this with two separate `wd_test` rules. We _could_ use a build step to create the
+# snapshot, but then a failure at this stage would be reported as a build failure when really it
+# should count as a test failure. So the only option left is to make this modification to wd_test to
+# invoke workerd twice for snapshot tests.
+
+WINDOWS_RUNTEST_SNAPSHOT = """
+powershell -Command \"%* --python-save-snapshot\" `-dTEST_TMPDIR=$ENV:TEST_TMPDIR
+set TEST_EXIT=!ERRORLEVEL!
+if !TEST_EXIT! EQ 0 (
+    powershell -Command \"%* --python-load-snapshot\" `-dTEST_TMPDIR=$ENV:TEST_TMPDIR
+    set TEST_EXIT=!ERRORLEVEL!
+)
+"""
+
+SH_RUNTEST_SNAPSHOT = """
+"$@" -dTEST_TMPDIR=$TEST_TMPDIR --python-save-snapshot
+"$@" -dTEST_TMPDIR=$TEST_TMPDIR --python-load-snapshot
+"""
+
 def _wd_test_impl(ctx):
     is_windows = ctx.target_platform_has_constraint(ctx.attr._platforms_os_windows[platform_common.ConstraintValueInfo])
 
@@ -112,9 +146,13 @@ def _wd_test_impl(ctx):
         # Batch script executables must end with ".bat"
         executable = ctx.actions.declare_file("%s_wd_test.bat" % ctx.label.name)
         content = WINDOWS_TEMPLATE.replace("$(SIDECAR)", ctx.file.sidecar.path if ctx.file.sidecar else "")
+        runtest = WINDOWS_RUNTEST_SNAPSHOT if ctx.attr.python_snapshot_test else WINDOWS_RUNTEST_NORMAL
+        content = content.replace("$(RUNTEST)", runtest)
     else:
         executable = ctx.outputs.executable
         content = SH_TEMPLATE.replace("$(SIDECAR)", ctx.file.sidecar.short_path if ctx.file.sidecar else "")
+        runtest = SH_RUNTEST_SNAPSHOT if ctx.attr.python_snapshot_test else SH_RUNTEST_NORMAL
+        content = content.replace("$(RUNTEST)", runtest)
 
     ctx.actions.write(
         output = executable,
@@ -155,6 +193,7 @@ _wd_test = rule(
             executable = True,
             cfg = "exec",
         ),
+        "python_snapshot_test": attr.bool(),
         "_platforms_os_windows": attr.label(default = "@platforms//os:windows"),
     },
 )

src/cloudflare/internal/test/ai/BUILD.bazel

@@ -15,6 +15,8 @@ py_wd_test(
         "*.js",
         "*.py",
     ]),
+    # Works but times out frequently
+    make_snapshot = False,
     tags = [
         # TODO(someday): Fix asan failure for this, see https://github.com/cloudflare/workerd/pull/3140#discussion_r1858273318
         "no-asan",

src/cloudflare/internal/test/aig/BUILD.bazel

@@ -15,6 +15,8 @@ py_wd_test(
         "*.js",
         "*.py",
     ]),
+    # Works but times out
+    make_snapshot = False,
     tags = [
         # TODO(someday): Fix asan failure for this, see https://github.com/cloudflare/workerd/pull/3140#discussion_r1858273318
         "no-asan",

src/cloudflare/internal/test/br/BUILD.bazel

@@ -15,6 +15,8 @@ py_wd_test(
         "*.js",
         "*.py",
     ]),
+    # Works but times out
+    make_snapshot = False,
     tags = [
         # TODO(someday): Fix asan failure for this, see https://github.com/cloudflare/workerd/pull/3140#discussion_r1858273318
         "no-asan",

src/cloudflare/internal/test/d1/BUILD.bazel

@@ -24,6 +24,7 @@ py_wd_test(
         "*.py",
         "*.js",
     ]),
+    make_snapshot = False,
     tags = [
         # TODO(someday): Fix asan failure for this, see https://github.com/cloudflare/workerd/pull/3140#discussion_r1858273318
         "no-asan",

src/cloudflare/internal/test/vectorize/BUILD.bazel

@@ -13,6 +13,8 @@ py_wd_test(
         "*.py",
         "*.js",
     ]),
+    # Works but times out
+    make_snapshot = False,
     tags = [
         # TODO(someday): Fix asan failure for this, see https://github.com/cloudflare/workerd/pull/3140#discussion_r1858273318
         "no-asan",

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

@@ -445,9 +445,21 @@ jsg::Ref<PyodideMetadataReader> makePyodideMetadataReader(Worker::Reader conf,
     }
     names.add(kj::str(module.getName()));
   }
-  bool createSnapshot = pythonConfig.createSnapshot;
-  bool createBaselineSnapshot = pythonConfig.createBaselineSnapshot;
-  bool snapshotToDisk = createSnapshot || createBaselineSnapshot;
+  bool snapshotToDisk = pythonConfig.createSnapshot || pythonConfig.createBaselineSnapshot;
+  if (pythonConfig.loadSnapshotFromDisk && snapshotToDisk) {
+    KJ_FAIL_ASSERT(
+        "Doesn't make sense to pass both --python-save-snapshot and --python-load-snapshot");
+  }
+  kj::Maybe<kj::Array<kj::byte>> memorySnapshot = kj::none;
+  if (pythonConfig.loadSnapshotFromDisk) {
+    auto& root = KJ_REQUIRE_NONNULL(pythonConfig.packageDiskCacheRoot);
+    kj::Path path("snapshot.bin");
+    auto maybeFile = root->tryOpenFile(path);
+    if (maybeFile == kj::none) {
+      KJ_FAIL_REQUIRE("Expected to find snapshot.bin in the package cache directory");
+    }
+    memorySnapshot = KJ_REQUIRE_NONNULL(maybeFile)->readAllBytes();
+  }
   auto lock = KJ_ASSERT_NONNULL(getPyodideLock(pythonRelease),
       kj::str("No lock file defined for Python packages release ", pythonRelease.getPackages()));
 
@@ -462,9 +474,9 @@ jsg::Ref<PyodideMetadataReader> makePyodideMetadataReader(Worker::Reader conf,
     true      /* isWorkerd */,
     false     /* isTracing */,
     snapshotToDisk,
-    createBaselineSnapshot,
+    pythonConfig.createBaselineSnapshot,
     false,    /* usePackagesInArtifactBundler */
-    kj::none  /* memorySnapshot */
+    kj::mv(memorySnapshot)
   );
   // clang-format on
 }

src/workerd/api/pyodide/pyodide.h

@@ -49,6 +49,7 @@ struct PythonConfig {
   const PyodideBundleManager pyodideBundleManager;
   bool createSnapshot;
   bool createBaselineSnapshot;
+  bool loadSnapshotFromDisk;
 };
 
 // A function to read a segment of the tar file into a buffer

src/workerd/server/server.h

@@ -79,6 +79,9 @@ class Server final: private kj::TaskSet::ErrorHandler {
   void setPythonCreateBaselineSnapshot() {
     pythonConfig.createBaselineSnapshot = true;
   }
+  void setPythonLoadSnapshot() {
+    pythonConfig.loadSnapshotFromDisk = true;
+  }
 
   // Runs the server using the given config.
   kj::Promise<void> run(jsg::V8System& v8System,
@@ -118,7 +121,8 @@ class Server final: private kj::TaskSet::ErrorHandler {
   PythonConfig pythonConfig = PythonConfig{.packageDiskCacheRoot = kj::none,
     .pyodideDiskCacheRoot = kj::none,
     .createSnapshot = false,
-    .createBaselineSnapshot = false};
+    .createBaselineSnapshot = false,
+    .loadSnapshotFromDisk = false};
 
   bool experimental = false;
 

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

@@ -20,6 +20,8 @@ py_wd_test("env-param")
 
 py_wd_test(
     "asgi",
+    # TODO: investigate failure!
+    make_snapshot = False,
     skip_python_flags = ["0.27.1"],
 )
 

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

@@ -53,6 +53,7 @@ def gen_import_tests(to_test, pkg_skip_versions = {}):
             directory = lib,
             src = wd_test_fname,
             skip_python_flags = pkg_skip_versions.get(lib, []),
+            make_snapshot = False,
             args = ["--experimental", "--pyodide-package-disk-cache-dir", "../all_pyodide_wheels"],
             data = [worker_py_fname, "@all_pyodide_wheels//:whls"],
             size = "enormous",

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

@@ -11,12 +11,16 @@ def _py_wd_test_helper(
         name,
         src,
         python_flag,
+        snapshot,
+        *,
+        args = [],
         **kwargs):
     name_flag = name + "_" + python_flag
     templated_src = name_flag.replace("/", "-") + "@template"
     templated_src = "/".join(src.split("/")[:-1] + [templated_src])
     flags = FEATURE_FLAGS[python_flag] + ["python_workers"]
     feature_flags_txt = ",".join(['"{}"'.format(flag) for flag in flags])
+
     expand_template(
         name = name_flag + "@rule",
         out = templated_src,
@@ -27,6 +31,8 @@ def _py_wd_test_helper(
     wd_test(
         src = templated_src,
         name = name_flag + "@",
+        args = args,
+        python_snapshot_test = snapshot,
         **kwargs
     )
 
@@ -41,6 +47,7 @@ def py_wd_test(
         args = [],
         size = "enormous",
         tags = [],
+        make_snapshot = True,
         **kwargs):
     if python_flags == "all":
         python_flags = FEATURE_FLAGS.keys()
@@ -59,13 +66,21 @@ def py_wd_test(
     elif name == None:
         name = src.removesuffix(".wd-test")
     data += ["//src/workerd/server/tests/python:pyodide_dev.capnp.bin@rule"]
-    args = args + ["--pyodide-bundle-disk-cache-dir", "$(location //src/workerd/server/tests/python:pyodide_dev.capnp.bin@rule)/..", "--experimental"]
+    args = args + [
+        "--pyodide-bundle-disk-cache-dir",
+        "$(location //src/workerd/server/tests/python:pyodide_dev.capnp.bin@rule)/..",
+        "--experimental",
+        "--pyodide-package-disk-cache-dir",
+        ".",
+    ]
     tags = tags + ["py_wd_test"]
+
     for python_flag in python_flags:
         _py_wd_test_helper(
             name,
             src,
             python_flag,
+            snapshot = make_snapshot,
             data = data,
             args = args,
             size = size,

src/workerd/server/workerd.c++

@@ -779,10 +779,15 @@ class CliMain final: public SchemaFileImpl::ErrorReporter {
       server->setPythonCreateSnapshot();
       return true;
     }, "Save a dedicated snapshot to the disk cache")
-        .addOption({"python-save-baseline-snapshot"}, [this]() {
+        .addOption({"python-save-baseline-snapshot"},
+            [this]() {
       server->setPythonCreateBaselineSnapshot();
       return true;
-    }, "Save a baseline snapshot to the disk cache");
+    }, "Save a baseline snapshot to the disk cache")
+        .addOption({"python-load-snapshot"}, [this]() {
+      server->setPythonLoadSnapshot();
+      return true;
+    }, "Load a snapshot from the package disk cache");
   }
 
   kj::MainFunc addServeOptions(kj::MainBuilder& builder) {