compiler: remove the implicit last-build state (B2)

generate_config/compile_kernel/load_module also recorded the last build
directory and PTX in ContextVars so that bare compile_kernel() and
load_module() calls worked: a second, implicit interface beside the
explicit build_dir/ptx_path one used by every owner. compile_kernel now
requires build_dir= and load_module requires the PTX path.

Build contract tests use explicit paths; they still check per-build
header/output isolation, interleaved generation, path escapes, refusal of a
removed snapshot, and that a failed compile keeps the previous PTX and
leaves no partial output. No kernel, numerical or launch change.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

src/CelerisLab/cuda/compiler_v2.py

@@ -19,7 +19,6 @@ import shutil
 import tempfile
 import hashlib
 import json
-from contextvars import ContextVar
 from dataclasses import asdict
 from pathlib import Path
 from typing import Optional
@@ -36,9 +35,6 @@ _KERNEL_DIR = os.path.join(
 )
 
 
-_last_build_dir = ContextVar("celeris_last_build_dir", default=None)
-_last_ptx_path = ContextVar("celeris_last_ptx_path", default=None)
-
 
 def kernel_path(filename: str) -> str:
     return os.path.join(_KERNEL_DIR, filename)
@@ -60,9 +56,7 @@ def generate_config(cfg: LBMConfig, n_objects: int = 0,
                     *, build_root: Optional[str] = None, block_root_bc: int = 0) -> str:
     """Snapshot kernels and render headers in a fresh private directory.
 
-    Returns its path. Pass it explicitly to compile_kernel for interleaved
+    Returns its path; pass it to compile_kernel(build_dir=...). The caller owns cleanup_build_dir() once the loaded module is no longer
-    builds. Legacy generate_config(); compile_kernel() remains context-local.
-    The caller owns cleanup_build_dir() once the loaded module is no longer
     needed for recompilation/debugging. No package files are changed.
     block_root_bc specializes a blocked root module; fine/dense builds use 0.
     """
@@ -91,8 +85,6 @@ def generate_config(cfg: LBMConfig, n_objects: int = 0,
     except BaseException:
         shutil.rmtree(build_dir)
         raise
-    _last_build_dir.set(build_dir)
-    _last_ptx_path.set(None)
     return build_dir
 
 
@@ -104,9 +96,6 @@ def cleanup_build_dir(build_dir: Optional[str]) -> None:
     if path == Path(_KERNEL_DIR).resolve() or not (path / ".celeris-build").is_file():
         raise ValueError("Refusing to remove a directory not owned by the compiler")
     shutil.rmtree(path)
-    if _last_build_dir.get() == str(build_dir):
-        _last_build_dir.set(None)
-        _last_ptx_path.set(None)
 
 
 def _render_config(m: dict, n_objects: int, cdir: str, *, block_root_bc: int = 0):
@@ -253,12 +242,8 @@ def _render_config(m: dict, n_objects: int, cdir: str, *, block_root_bc: int = 0
 # Compilation
 # ---------------------------------------------------------------------------
 def compile_kernel(arch: str = "sm_70", source: str = "kernel_v2.cu",
-                   output: Optional[str] = None, *,
+                   output: Optional[str] = None, *, build_dir: str) -> str:
-                   build_dir: Optional[str] = None) -> str:
+    """Compile the snapshot in build_dir (from generate_config); never package headers."""
-    """Compile a generated private source snapshot; never use package headers."""
-    # A failed attempt must not make implicit load_module() reuse older code.
-    _last_ptx_path.set(None)
-    build_dir = build_dir or _last_build_dir.get()
     if not build_dir or not (Path(build_dir) / ".celeris-build").is_file():
         raise RuntimeError("Call generate_config() and pass its build_dir before compiling")
     root = Path(build_dir).resolve()
@@ -283,15 +268,11 @@ def compile_kernel(arch: str = "sm_70", source: str = "kernel_v2.cu",
     finally:
         if os.path.exists(tmp):
             os.unlink(tmp)
-    _last_ptx_path.set(str(dst))
     return str(dst)
 
 
-def load_module(ptx_path: Optional[str] = None):
+def load_module(ptx_path: str):
-    """Load explicit PTX or the last PTX compiled in this Python context."""
+    """Load one compiled PTX file as a new PyCUDA module."""
-    ptx_path = ptx_path or _last_ptx_path.get()
-    if not ptx_path:
-        raise RuntimeError("No compiled PTX; call compile_kernel() first")
     import pycuda.driver as cuda
     return cuda.module_from_file(ptx_path)
 

tests/unit/test_build_contracts.py

@@ -126,8 +126,7 @@ class BuildContracts(unittest.TestCase):
             def work(nx, count):
                 build = compiler.generate_config(LBMConfig(nx=nx), count, build_root=root)
                 barrier.wait()
-                # Compatibility form is context-local even without explicit path.
+                output = compiler.compile_kernel(build_dir=build)
-                output = compiler.compile_kernel()
                 data = json.loads(Path(output).read_text())
                 self.assertEqual(data, {"nx": nx, "objects": count})
                 metadata = json.loads(Path(output + ".json").read_text())
@@ -156,7 +155,7 @@ class BuildContracts(unittest.TestCase):
             compiler.cleanup_build_dir(a)
             compiler.cleanup_build_dir(b)
             with self.assertRaises(RuntimeError):
-                compiler.compile_kernel()
+                compiler.compile_kernel(build_dir=a)  # removed snapshot is never compiled
 
     def test_failed_compile_cannot_implicitly_load_previous_artifact(self):
         with tempfile.TemporaryDirectory() as root:
@@ -169,8 +168,7 @@ class BuildContracts(unittest.TestCase):
                 with self.assertRaises(subprocess.CalledProcessError):
                     compiler.compile_kernel(build_dir=build)
             self.assertEqual(Path(previous).read_bytes(), old_bytes)
-            with self.assertRaisesRegex(RuntimeError, "No compiled PTX"):
+            self.assertEqual(sorted(Path(build).glob("kernel_v2-*.ptx")), [])  # no partial output
-                compiler.load_module()
             compiler.cleanup_build_dir(build)