From dbcdb8708129aaeb7bedaa7924dfd6eac530c528 Mon Sep 17 00:00:00 2001 From: Claude Date: Mon, 26 Oct 2026 14:05:58 +0000 Subject: [PATCH] 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 --- src/CelerisLab/cuda/compiler_v2.py | 29 +++++------------------------ tests/unit/test_build_contracts.py | 8 +++----- 2 files changed, 8 insertions(+), 29 deletions(-) diff --git a/src/CelerisLab/cuda/compiler_v2.py b/src/CelerisLab/cuda/compiler_v2.py index 0a687fa..ec22b02 100644 --- a/src/CelerisLab/cuda/compiler_v2.py +++ b/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 - builds. Legacy generate_config(); compile_kernel() remains context-local. - The caller owns cleanup_build_dir() once the loaded module is no longer + Returns its path; pass it to compile_kernel(build_dir=...). 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, *, - build_dir: Optional[str] = None) -> str: - """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() + output: Optional[str] = None, *, build_dir: str) -> str: + """Compile the snapshot in build_dir (from generate_config); never package headers.""" 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): - """Load explicit PTX or the last PTX compiled in this Python context.""" - ptx_path = ptx_path or _last_ptx_path.get() - if not ptx_path: - raise RuntimeError("No compiled PTX; call compile_kernel() first") +def load_module(ptx_path: str): + """Load one compiled PTX file as a new PyCUDA module.""" import pycuda.driver as cuda return cuda.module_from_file(ptx_path) diff --git a/tests/unit/test_build_contracts.py b/tests/unit/test_build_contracts.py index 4d8f11a..6b5fa16 100644 --- a/tests/unit/test_build_contracts.py +++ b/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() + output = compiler.compile_kernel(build_dir=build) 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"): - compiler.load_module() + self.assertEqual(sorted(Path(build).glob("kernel_v2-*.ptx")), []) # no partial output compiler.cleanup_build_dir(build)