d3: port the two package-gate fixes from the vela worker

recipe/d3/reference.md

@@ -112,9 +112,11 @@ The package is about 52 GB (26.09B BF16 parameters plus the FP32
 revision's own (a self-consistent manifest from a later snapshot carries a
 different `repo_id` or runtime table), when any of the 9 runtime files
 fails its full-hash check, or when the package root holds an importable
-entry the runtime table does not name (an extra `*.py` file or
-`__init__.py` package, which the `sys.path` setup would let shadow a
-dependency) — that gate runs before any package code is imported.
+entry the runtime table does not name (an extra source, `.pyc` or
+extension module, or `__init__.py` package, which the `sys.path` setup
+would let shadow a dependency) or a `__pycache__` whose stale bytecode
+would execute in place of the verified source — that gate runs before any
+package code is imported.
 
 ## Start the worker
 

src/models/d3/README.md

@@ -26,9 +26,11 @@ This worker additionally hashes every file of
 `runtime.files_sha256` (8 runtime modules plus `requirements.txt`) in full
 **before importing anything** from the package directory, and refuses to
 serve packages without a manifest. It also rejects importable entries the
-runtime table does not name — an extra `*.py` file or `__init__.py` package
-in the package root — because that root is placed first on `sys.path`, where
-such an entry would shadow a dependency of the pinned runtime. Because hash
+runtime table does not name — an extra source, `.pyc` or extension module,
+or an `__init__.py` package in the package root — because that root is
+placed first on `sys.path`, where such an entry would shadow a dependency
+of the pinned runtime, and any `__pycache__`, whose stale bytecode would
+execute in place of the verified source. Because hash
 agreement with a manifest
 proves only internal consistency — a later snapshot ships a self-consistent
 table of its own — the gate first requires the manifest to carry exactly

src/models/d3/engine.py

@@ -11,9 +11,11 @@
 ``files_sha256`` table and skips hashing files over 64 MiB in fast mode — the
 executor hashes every file of the manifest's ``runtime.files_sha256`` table
 in full before importing anything from the package directory, and refuses
-unlisted importable entries (an extra ``*.py`` file or ``__init__.py``
-package) in the package root: that root goes first on ``sys.path``, so such
-an entry would shadow a dependency the pinned runtime imports. Hash
+unlisted importable entries (an extra source, ``.pyc`` or extension module,
+or an ``__init__.py`` package) and any ``__pycache__`` in the package root:
+that root goes first on ``sys.path``, so such an entry would shadow a
+dependency the pinned runtime imports, and a stale bytecode cache would
+execute in place of the verified source. Hash
 agreement proves only internal consistency — a later snapshot ships a
 self-consistent table of its own — so the gate first requires the manifest
 to carry exactly the repo id and runtime digests checked in for
@@ -26,6 +28,7 @@
 import json
 import sys
 import threading
+from importlib.machinery import EXTENSION_SUFFIXES, SOURCE_SUFFIXES
 from pathlib import Path
 from typing import Any
 
@@ -34,6 +37,7 @@
 MANIFEST_NAME = "MODEL_MANIFEST.json"
 MANIFEST_SCHEMA = "d3-package-manifest/1"
 VERIFY_MODES = ("fast", "full", "none")
+IMPORTABLE_FILE_SUFFIXES = frozenset(SOURCE_SUFFIXES) | {".pyc"} | frozenset(EXTENSION_SUFFIXES)
 # The runtime.files_sha256 table of REFERENCE_REVISION, transcribed from the
 # pinned package's manifest; repinning means updating this table together
 # with REFERENCE_REVISION.
@@ -105,8 +109,10 @@ def verify_runtime_files(root: Path, manifest: dict[str, Any]) -> list[str]:
     files this process is about to execute, so unlike the upstream shard
     verification there is no size shortcut. It also refuses importable
     entries the table does not name — the package root is about to go first
-    on ``sys.path``, where an extra ``*.py`` file or ``__init__.py`` package
-    would shadow one of the pinned runtime's dependencies.
+    on ``sys.path``, where an extra source, ``.pyc`` or extension module,
+    or an ``__init__.py`` package would shadow one of the pinned runtime's
+    dependencies — and any ``__pycache__``, whose cached bytecode could
+    execute in place of a verified source file.
     """
 
     table = (manifest.get("runtime") or {}).get("files_sha256")
@@ -125,7 +131,9 @@ def verify_runtime_files(root: Path, manifest: dict[str, Any]) -> list[str]:
     for entry in sorted(Path(root).iterdir()):
         if entry.name in table:
             continue
-        if entry.is_file() and entry.suffix == ".py":
+        if entry.name == "__pycache__":
+            problems.append("unexpected bytecode cache __pycache__/ in the package root")
+        elif entry.is_file() and entry.suffix in IMPORTABLE_FILE_SUFFIXES:
             problems.append(f"unexpected importable file {entry.name} in the package root")
         elif entry.is_dir() and (entry / "__init__.py").is_file():
             problems.append(f"unexpected importable package {entry.name}/ in the package root")
@@ -134,6 +142,19 @@ def verify_runtime_files(root: Path, manifest: dict[str, Any]) -> list[str]:
     return sorted(table)
 
 
+def put_first_on_sys_path(root: Path) -> str:
+    """Put the resolved package root at the front of ``sys.path``.
+
+    Equivalent entries are dropped, not reused: a ``PYTHONPATH`` copy of
+    the root listed later would leave foreign directories ahead of it.
+    """
+
+    resolved = Path(root).resolve()
+    sys.path[:] = [entry for entry in sys.path if Path(entry or ".").resolve() != resolved]
+    sys.path.insert(0, str(resolved))
+    return str(resolved)
+
+
 class D3Executor:
     """Load and drive one immutable pinned d3 runtime instance."""
 
@@ -163,11 +184,16 @@ def __init__(
         # Import the just-verified runtime the way upstream d3_server.py does:
         # the package directory goes first on sys.path so the flat runtime
         # modules (d3_runtime, d3_format, ...) resolve to the checkpoint's own
-        # copies. One worker process serves one package.
-        root_str = str(root)
-        if root_str not in sys.path:
-            sys.path.insert(0, root_str)
-        import d3_runtime
+        # copies. One worker process serves one package. The import runs with
+        # bytecode writing off: a __pycache__ it left next to the verified
+        # source is what the next gate run refuses.
+        put_first_on_sys_path(root)
+        saved_dont_write_bytecode = sys.dont_write_bytecode
+        sys.dont_write_bytecode = True
+        try:
+            import d3_runtime
+        finally:
+            sys.dont_write_bytecode = saved_dont_write_bytecode
 
         self.model = d3_runtime.D3.from_pretrained(
             root,

tests/d3/test_engine_manifest.py

@@ -11,6 +11,9 @@
 
 import hashlib
 import json
+import py_compile
+import re
+import sys
 from pathlib import Path
 
 import pytest
@@ -20,6 +23,7 @@
     PINNED_RUNTIME_SHA256,
     REFERENCE_REPO,
     load_manifest,
+    put_first_on_sys_path,
     verify_pinned_package,
     verify_runtime_files,
 )
@@ -169,6 +173,19 @@ def test_extra_importable_files_in_the_package_root_are_refused(tmp_path):
         verify_runtime_files(tmp_path, load_manifest(tmp_path))
 
 
+@pytest.mark.parametrize(
+    "name",
+    ["torch.pyc", "torch.abi3.so", "_d3.cpython-312-x86_64-linux-gnu.so"],
+)
+def test_sourceless_and_extension_module_files_are_refused(tmp_path, name):
+    # Python imports .pyc bytecode and compiled extensions from sys.path
+    # entries too, not just .py source files.
+    build_package(tmp_path, small_files())
+    (tmp_path / name).write_bytes(b"module bytes\n")
+    with pytest.raises(ValueError, match=f"unexpected importable file {re.escape(name)}"):
+        verify_runtime_files(tmp_path, load_manifest(tmp_path))
+
+
 def test_extra_package_directories_in_the_package_root_are_refused(tmp_path):
     build_package(tmp_path, small_files())
     fla = tmp_path / "fla"
@@ -179,17 +196,54 @@ def test_extra_package_directories_in_the_package_root_are_refused(tmp_path):
         verify_runtime_files(tmp_path, load_manifest(tmp_path))
 
 
+def test_existing_bytecode_caches_in_the_package_root_are_refused(tmp_path):
+    # CPython imports an unchecked-hash __pycache__/d3_runtime.<tag>.pyc in
+    # place of the verified source, so a cache compiled from different code
+    # must be refused even when the source on disk matches the manifest.
+    build_package(tmp_path, small_files())
+    runtime = tmp_path / "d3_runtime.py"
+    verified = runtime.read_bytes()
+    runtime.write_bytes(b"import os\nos._exit(0)\n")
+    pycache = tmp_path / "__pycache__"
+    pycache.mkdir()
+    py_compile.compile(
+        runtime,
+        pycache / f"d3_runtime.{sys.implementation.cache_tag}.pyc",
+        invalidation_mode=py_compile.PycInvalidationMode.UNCHECKED_HASH,
+    )
+    runtime.write_bytes(verified)
+
+    with pytest.raises(ValueError, match="unexpected bytecode cache __pycache__"):
+        verify_runtime_files(tmp_path, load_manifest(tmp_path))
+
+
 def test_unimportable_package_root_entries_still_pass(tmp_path):
     # The pinned package's own layout shape: an assets/ directory without
-    # __init__.py, plain data files, and the __pycache__ Python writes next
-    # to the imported runtime — none of it can shadow an import.
+    # __init__.py and plain data files — none of it can shadow an import.
     manifest = build_package(tmp_path, small_files())
     assets = tmp_path / "assets"
     assets.mkdir()
     (assets / "banner.png").write_bytes(b"not a python package\n")
     (tmp_path / "LICENSE").write_text("Apache-2.0\n")
-    pycache = tmp_path / "__pycache__"
-    pycache.mkdir()
-    (pycache / "d3_runtime.cpython-312.pyc").write_bytes(b"bytecode\n")
 
     assert verify_runtime_files(tmp_path, manifest) == ["d3_runtime.py", "requirements.txt"]
+
+
+def test_the_package_root_is_forced_to_the_front_of_sys_path(tmp_path):
+    # A PYTHONPATH copy of the root listed later must not leave a foreign
+    # directory ahead of it: sys.path order alone decides which copy the
+    # runtime import resolves to, and nothing else checks it.
+    package = tmp_path / "pkg"
+    shadow = tmp_path / "shadow"
+    package.mkdir()
+    shadow.mkdir()
+    resolved = str(package.resolve())
+    sys.path.insert(0, str(shadow))
+    sys.path.append(str(package))
+    try:
+        assert put_first_on_sys_path(package) == resolved
+        assert sys.path[0] == resolved
+        assert resolved not in sys.path[1:]
+        assert str(shadow) in sys.path
+    finally:
+        sys.path[:] = [entry for entry in sys.path if entry not in (str(shadow), resolved)]