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>
This commit is contained in:
Claude
2026-10-26 14:05:58 +00:00
co-authored by Claude Opus 5.5
parent 0d62c3cfb8
commit dbcdb87081
2 changed files with 8 additions and 29 deletions
+5 -24
View File
@@ -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)
+3 -5
View File
@@ -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)