From 1b9e71c3097b5687bb2f3fc19a904e7b0e74ab66 Mon Sep 17 00:00:00 2001 From: Mark Saroufim Date: Wed, 9 Sep 2026 21:50:29 -0700 Subject: [PATCH 1/4] Add Modal experiment for CPU-built load_inline extensions --- docs/modal-inline-artifact-prototype.md | 75 +++++++++ scripts/modal_inline_artifact_prototype.py | 185 +++++++++++++++++++++ src/libkernelbot/inline_artifacts.py | 169 +++++++++++++++++++ tests/test_inline_artifacts.py | 130 +++++++++++++++ 4 files changed, 559 insertions(+) create mode 100644 docs/modal-inline-artifact-prototype.md create mode 100644 scripts/modal_inline_artifact_prototype.py create mode 100644 src/libkernelbot/inline_artifacts.py create mode 100644 tests/test_inline_artifacts.py diff --git a/docs/modal-inline-artifact-prototype.md b/docs/modal-inline-artifact-prototype.md new file mode 100644 index 00000000..441d52fb --- /dev/null +++ b/docs/modal-inline-artifact-prototype.md @@ -0,0 +1,75 @@ +# Compile load_inline submissions before allocating a GPU + +This experiment imports the vector-add submission on a CPU-only Modal worker, +copies its compiled `.so` to a T4 worker, and runs KernelBot's correctness harness. +It leaves the submission's `load_inline` call intact and uses the same runner +image for both workers. The production submission path is unchanged. + +```bash +uv sync --extra dev +PYTHONPATH=src:src/runners uv run modal run \ + scripts/modal_inline_artifact_prototype.py \ + --output /tmp/inline-artifact-result.json +``` + +Set `KERNELBOT_PROTOTYPE_GPU` to `L4`, `A100`, `H100`, or `B200` to change the GPU +and compilation target. T4 is the tested default. The app exits after the run. +The compiler wrapper from the runner image is present, but its PCH Volume is +not mounted: both comparison builds start without a precompiled-header cache. + +The default input is `examples/vectoradd_py/submission_cuda_inline.py`. +`--submission path/to/submission.py` accepts another implementation of that same +vector-add task. The harness tests square matrices with dimensions 1, 127, 128, +129, 256, 512 and 1024. + +## How it works + +A `sitecustomize` hook intercepts `load_inline` in the submission process and +the evaluator's spawned processes. On the CPU worker, it calls the compiler and +saves the resulting library and a manifest. On the GPU worker, it checks the +manifest and binary hash, then imports the saved library. Requests are matched +by source, options, target architecture and Python/PyTorch/CUDA ABI information. +Missing artifacts fail; the hook disables PyTorch's Ninja build entrypoint. + +The script transfers library bytes through Modal results and arguments. It +checks that replay creates no build directory, runs the existing +`run_single_evaluation` harness, and verifies that a missing artifact fails. +It then compiles once on the GPU worker to provide a timing baseline. That last +build is diagnostic overhead, not part of the proposed CPU-build path. + +## Results + +The [T4 run](https://modal.com/apps/coreauto/main/ap-kfECmtFhii1jy58mgwXMYY) passed +all seven correctness cases and the missing-artifact check. Replay created no +build directory. Both workers requested 4 CPU cores and 8 GiB RAM. + +| Step | Extension compile/load | Python process including startup | +| --- | ---: | ---: | +| CPU build, zero visible GPUs | 55.684 s | 58.793 s | +| T4 artifact load | 7.728 ms | 2.722 s | +| Fresh build on the T4 worker | 55.413 s | 58.042 s | + +The transferred library and manifest were 1,560,196 bytes. Evaluation took +5.959 s. These are single-run observations, not throughput or billing estimates. + +The JSON output separates extension load/compile time, Python process time, +evaluation time and total experiment time. Compare compilation and replay +separately from scheduling, Python startup and the diagnostic baseline. + +## Limits + +The submission must compile at import time without querying or using a GPU. +Lazy builds, `load()`, non-Python extensions and libraries outside the common +image are unsupported. Errors propagate without a fallback. + +Artifacts are used within one run; this is not a persistent compilation cache. +The manifest does not fingerprint external headers or every compiler input. +Sharing artifacts between submissions would require a complete build identity. +Only use binaries produced by the builder for the same request: the hash detects +corruption, and the Python hook is not a security boundary. + +Run the adapter tests without CUDA: + +```bash +uv run pytest tests/test_inline_artifacts.py +``` diff --git a/scripts/modal_inline_artifact_prototype.py b/scripts/modal_inline_artifact_prototype.py new file mode 100644 index 00000000..5059dfd2 --- /dev/null +++ b/scripts/modal_inline_artifact_prototype.py @@ -0,0 +1,185 @@ +"""Run the CPU-build experiment described in docs/modal-inline-artifact-prototype.md.""" + +import dataclasses +import json +import os +import subprocess +import tempfile +import time +from pathlib import Path + +import modal +from modal_runner import cuda_image + +from libkernelbot.inline_artifacts import ( + pack_artifacts, + prepare_environment, + read_events, + unpack_artifacts, +) + +app = modal.App("kernelbot-inline-artifact-prototype") +GPU = os.environ.get("KERNELBOT_PROTOTYPE_GPU", "T4").upper() +ARCH = {"T4": "7.5", "L4": "8.9", "A100": "8.0", "H100": "9.0a", "B200": "10.0a"}[GPU] + + +def _write_sources(work: Path, sources: dict[str, str]) -> None: + for name, source in sources.items(): + path = work / name + if not path.resolve().is_relative_to(work.resolve()): + raise ValueError(f"Invalid source path: {name}") + path.parent.mkdir(parents=True, exist_ok=True) + path.write_text(source) + + +def _import_submission(work: Path, environment: dict) -> dict: + started = time.perf_counter() + result = subprocess.run( + ["python3", "submission.py"], + cwd=work, + env=environment, + capture_output=True, + text=True, + timeout=240, + ) + if result.returncode: + raise RuntimeError(f"Submission import failed:\n{result.stdout}\n{result.stderr}") + return { + "process_seconds": time.perf_counter() - started, + "events": read_events(work / "artifacts"), + } + + +@app.function( + image=cuda_image, cpu=4, memory=8192, timeout=300, max_containers=1, scaledown_window=2 +) +def compile_cpu(sources: dict[str, str], arch: str) -> dict: + import torch + + if torch.cuda.is_available() or torch.cuda.device_count() != 0: + raise RuntimeError("CPU build stage unexpectedly has a GPU") + with tempfile.TemporaryDirectory(prefix="kb-cpu-build-") as directory: + work = Path(directory) + _write_sources(work, sources) + result = _import_submission(work, prepare_environment(work, "capture", arch)) + artifacts = pack_artifacts(work / "artifacts") + if not artifacts: + raise RuntimeError( + "No import-time load_inline calls captured; " + "lazy/GPU-dependent builds are unsupported" + ) + result.update({"cuda_available": False, "device_count": 0, "artifacts": artifacts}) + return result + + +@app.function( + image=cuda_image, gpu=GPU, cpu=4, memory=8192, timeout=600, max_containers=1, scaledown_window=2 +) +def evaluate_gpu( + sources: dict[str, str], artifacts: dict[str, bytes], arch: str, tests: str +) -> dict: + import torch + + from libkernelbot.run_eval import make_system_info, run_single_evaluation + + expected_capability = tuple(int(x) for x in arch.rstrip("a").split(".")) + if torch.cuda.get_device_capability() != expected_capability: + raise RuntimeError("GPU architecture does not match the CPU compilation target") + started = time.perf_counter() + result = {"gpu": torch.cuda.get_device_name(), "capability": expected_capability} + with tempfile.TemporaryDirectory(prefix="kb-gpu-replay-") as directory: + work = Path(directory) + _write_sources(work, sources) + unpack_artifacts(work / "artifacts", artifacts) + environment = prepare_environment(work, "replay", arch) + result["replay_import"] = _import_submission(work, environment) + + # Keep KernelBot's normal test protocol, including its multiprocessing evaluator. + previous_cwd, previous_environment = Path.cwd(), dict(os.environ) + try: + os.chdir(work) + os.environ.update(environment) + run, _ = run_single_evaluation( + ["python3", "eval.py"], + "test", + system=make_system_info(), + tests=tests, + seed=42, + ) + finally: + os.chdir(previous_cwd) + os.environ.clear() + os.environ.update(previous_environment) + result["evaluation"] = dataclasses.asdict(run) + if not run.success or not run.passed: + raise RuntimeError(f"Artifact evaluation failed: {result['evaluation']}") + result["replay_events"] = read_events(work / "artifacts") + if not result["replay_events"] or any( + e["mode"] != "replay" for e in result["replay_events"] + ): + raise RuntimeError("Expected artifact-only GPU execution") + if (work / "build").exists(): + raise RuntimeError("GPU replay unexpectedly created a build directory") + + # Negative control: missing binaries must fail instead of silently recompiling. + missing = work / "missing" + missing.mkdir() + _write_sources(missing, sources) + probe = subprocess.run( + ["python3", "submission.py"], + cwd=missing, + env=prepare_environment(missing, "replay", arch), + capture_output=True, + text=True, + timeout=60, + ) + result["missing_artifact_rejected"] = ( + probe.returncode != 0 and "No CPU-built artifact" in probe.stderr + ) + if not result["missing_artifact_rejected"]: + raise RuntimeError(f"Missing-artifact negative control failed: {probe.stderr}") + + # Compare against a fresh GPU-side build using the same sources and image. + baseline = work / "baseline" + baseline.mkdir() + _write_sources(baseline, sources) + result["gpu_compile_baseline"] = _import_submission( + baseline, prepare_environment(baseline, "capture", arch) + ) + result["gpu_function_seconds"] = time.perf_counter() - started + return result + + +@app.local_entrypoint() +def main(output: str = "", submission: str = ""): + repo = Path(__file__).resolve().parents[1] + example = repo / "examples" / "vectoradd_py" + sources = { + "submission.py": Path(submission).read_text() + if submission + else (example / "submission_cuda_inline.py").read_text(), + "task.py": (example / "task.py").read_text(), + "reference.py": (example / "reference.py").read_text(), + "utils.py": (repo / "examples" / "utils.py").read_text(), + "eval.py": (repo / "examples" / "eval.py").read_text(), + } + tests = ( + "\n".join(f"size: {size}; seed: 4242" for size in (1, 127, 128, 129, 256, 512, 1024)) + "\n" + ) + started = time.perf_counter() + built = compile_cpu.remote(sources, ARCH) + artifacts = built.pop("artifacts") + print("CPU_BUILD", json.dumps(built), flush=True) + print("ARTIFACT_BYTES", sum(map(len, artifacts.values())), flush=True) + evaluated = evaluate_gpu.remote(sources, artifacts, ARCH, tests) + report = { + "requested_gpu": GPU, + "target_arch": ARCH, + "artifact_bytes": sum(map(len, artifacts.values())), + "cpu_build": built, + "gpu": evaluated, + "client_pipeline_seconds": time.perf_counter() - started, + } + print("RESULT", json.dumps(report, indent=2), flush=True) + if output: + Path(output).write_text(json.dumps(report, indent=2) + "\n") diff --git a/src/libkernelbot/inline_artifacts.py b/src/libkernelbot/inline_artifacts.py new file mode 100644 index 00000000..87a1daa8 --- /dev/null +++ b/src/libkernelbot/inline_artifacts.py @@ -0,0 +1,169 @@ +"""Capture load_inline libraries and reload them in a matching GPU container.""" + +import hashlib +import importlib.util +import inspect +import json +import os +import platform +import sysconfig +import time +from pathlib import Path + + +def _digest(value: bytes) -> str: + return hashlib.sha256(value).hexdigest() + + +def _save_module(directory: Path, module, request: dict, runtime: dict) -> None: + directory.mkdir(exist_ok=True) + binary = Path(module.__file__).read_bytes() + (directory / "extension.so").write_bytes(binary) + (directory / "manifest.json").write_text( + json.dumps( + { + "module_name": module.__name__, + "sha256": _digest(binary), + "runtime": runtime, + "request": request, + }, + sort_keys=True, + ) + ) + + +def _load_module(directory: Path, request: dict, runtime: dict): + manifest_path = directory / "manifest.json" + if not manifest_path.is_file(): + raise RuntimeError( + f"No CPU-built artifact for {request['name']!r}; GPU compilation is disabled" + ) + manifest = json.loads(manifest_path.read_text()) + if manifest["runtime"] != runtime or manifest["request"] != request: + raise RuntimeError("CPU-built extension manifest does not match this request") + binary_path = directory / "extension.so" + if _digest(binary_path.read_bytes()) != manifest["sha256"]: + raise RuntimeError("CPU-built extension failed SHA-256 verification") + spec = importlib.util.spec_from_file_location(manifest["module_name"], binary_path) + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + return module + + +def _refuse_build(*args, **kwargs): + raise RuntimeError("GPU compilation is disabled in artifact replay mode") + + +def install() -> None: + """Called by the prototype's sitecustomize, including in spawned evaluators.""" + import torch + from torch.utils import cpp_extension + + mode = os.environ["KERNELBOT_INLINE_MODE"] + if mode not in {"capture", "replay"}: + raise ValueError(f"Unknown inline artifact mode: {mode}") + root = Path(os.environ["KERNELBOT_INLINE_ARTIFACTS"]) + root.mkdir(parents=True, exist_ok=True) + original = cpp_extension.load_inline + signature = inspect.signature(original) + runtime = { + "torch": str(torch.__version__), + "cuda": torch.version.cuda, + "python_abi": sysconfig.get_config_var("SOABI"), + "machine": platform.machine(), + "cxx11_abi": torch._C._GLIBCXX_USE_CXX11_ABI, + "arch": os.environ["TORCH_CUDA_ARCH_LIST"], + } + loaded = {} + + def load_inline(*args, **kwargs): + bound = signature.bind(*args, **kwargs) + bound.apply_defaults() + if not bound.arguments.get("is_python_module", True): + raise ValueError("Prototype supports Python extension modules only") + # Destination paths and logging do not affect the binary's identity. + request = { + k: v for k, v in bound.arguments.items() if k not in {"build_directory", "verbose"} + } + # PyTorch may mutate source lists while injecting headers/bindings. + request = json.loads(json.dumps(request)) + key = _digest(json.dumps({"request": request, "runtime": runtime}, sort_keys=True).encode()) + directory = root / key + started = time.perf_counter() + if mode == "capture": + module = original(*args, **kwargs) + elapsed = time.perf_counter() - started + _save_module(directory, module, request, runtime) + else: + if key not in loaded: + loaded[key] = _load_module(directory, request, runtime) + module = loaded[key] + elapsed = time.perf_counter() - started + # One append per event also works for the evaluator's spawned processes. + with (root / "events.jsonl").open("a") as log: + log.write( + json.dumps({"mode": mode, "key": key, "seconds": elapsed, "pid": os.getpid()}) + + "\n" + ) + return module + + cpp_extension.load_inline = load_inline + if mode == "replay": + cpp_extension._run_ninja_build = _refuse_build + + +def prepare_environment(work: Path, mode: str, arch: str) -> dict[str, str]: + """Install the opt-in hook only in child interpreters of this experiment.""" + hook = work / "bootstrap" + hook.mkdir(exist_ok=True) + (hook / "sitecustomize.py").write_text( + "import os, traceback\n" + "try:\n" + " from libkernelbot.inline_artifacts import install\n" + " install()\n" + "except BaseException:\n" + " traceback.print_exc()\n" + " os._exit(1)\n" + ) + return { + **os.environ, + "KERNELBOT_INLINE_MODE": mode, + "KERNELBOT_INLINE_ARTIFACTS": str(work / "artifacts"), + "TORCH_EXTENSIONS_DIR": str(work / "build"), + "TORCH_CUDA_ARCH_LIST": arch, + "MAX_JOBS": "2", + "PYTHONPATH": os.pathsep.join( + filter( + None, + [ + str(hook), + str(work), + str(Path(__file__).resolve().parents[1]), + os.environ.get("PYTHONPATH", ""), + ], + ) + ), + } + + +def pack_artifacts(root: Path) -> dict[str, bytes]: + """Transfer only the libraries and manifests, never the Ninja build cache.""" + return { + str(path.relative_to(root)): path.read_bytes() + for pattern in ("*/manifest.json", "*/extension.so") + for path in root.glob(pattern) + } + + +def unpack_artifacts(root: Path, files: dict[str, bytes]) -> None: + for name, data in files.items(): + path = root / name + if not path.resolve().is_relative_to(root.resolve()): + raise ValueError(f"Invalid artifact path: {name}") + path.parent.mkdir(parents=True, exist_ok=True) + path.write_bytes(data) + + +def read_events(root: Path) -> list[dict]: + path = root / "events.jsonl" + return [json.loads(line) for line in path.read_text().splitlines()] if path.exists() else [] diff --git a/tests/test_inline_artifacts.py b/tests/test_inline_artifacts.py new file mode 100644 index 00000000..88222848 --- /dev/null +++ b/tests/test_inline_artifacts.py @@ -0,0 +1,130 @@ +"""CPU-only contract tests; the Modal prototype exercises real compilation/loading.""" + +import json +import os +import sys +import tempfile +import types +import unittest +from pathlib import Path +from unittest.mock import Mock, patch + +from libkernelbot.inline_artifacts import install, pack_artifacts, unpack_artifacts + + +class InlineArtifactTests(unittest.TestCase): + def setUp(self): + self.temporary = tempfile.TemporaryDirectory() + self.addCleanup(self.temporary.cleanup) + self.root = Path(self.temporary.name) + self.binary = self.root / "compiled.so" + self.binary.write_bytes(b"fake compiled extension") + self.builds = 0 + + def load_inline( + name, + cpp_sources, + cuda_sources=None, + extra_cuda_cflags=None, + is_python_module=True, + verbose=False, + build_directory=None, + ): + self.builds += 1 + return types.SimpleNamespace(__name__=name, __file__=str(self.binary)) + + self.original = load_inline + self.extension = types.SimpleNamespace(load_inline=load_inline) + torch = types.SimpleNamespace( + __version__="test", + version=types.SimpleNamespace(cuda="13.3"), + _C=types.SimpleNamespace(_GLIBCXX_USE_CXX11_ABI=True), + ) + modules = patch.dict( + sys.modules, + { + "torch": torch, + "torch.utils": types.SimpleNamespace(cpp_extension=self.extension), + }, + ) + modules.start() + self.addCleanup(modules.stop) + environment = patch.dict( + os.environ, + { + "KERNELBOT_INLINE_MODE": "capture", + "KERNELBOT_INLINE_ARTIFACTS": str(self.root / "artifacts"), + "TORCH_CUDA_ARCH_LIST": "7.5", + }, + ) + environment.start() + self.addCleanup(environment.stop) + install() + self.extension.load_inline("example", "source", cuda_sources="cuda") + + def replay(self): + self.extension.load_inline = self.original + os.environ["KERNELBOT_INLINE_MODE"] = "replay" + install() + + def test_changed_source_and_flags_do_not_reuse_binary(self): + self.replay() + for args in ({"cpp_sources": "changed"}, {"extra_cuda_cflags": ["-O3"]}): + request = {"name": "example", "cpp_sources": "source", "cuda_sources": "cuda", **args} + with self.assertRaisesRegex(RuntimeError, "No CPU-built artifact"): + self.extension.load_inline(**request) + self.assertEqual(self.builds, 1) + + def test_changed_architecture_does_not_reuse_binary(self): + os.environ["TORCH_CUDA_ARCH_LIST"] = "9.0a" + self.replay() + with self.assertRaisesRegex(RuntimeError, "No CPU-built artifact"): + self.extension.load_inline("example", "source", cuda_sources="cuda") + + def test_replay_loads_saved_module_once_without_compiling(self): + self.replay() + module = types.SimpleNamespace() + spec = Mock() + with ( + patch("importlib.util.spec_from_file_location", return_value=spec) as make_spec, + patch("importlib.util.module_from_spec", return_value=module), + ): + for _ in range(2): + self.assertIs( + self.extension.load_inline("example", "source", cuda_sources="cuda"), module + ) + self.assertEqual(make_spec.call_args.args[0], "example") + self.assertEqual(make_spec.call_args.args[1].read_bytes(), self.binary.read_bytes()) + spec.loader.exec_module.assert_called_once_with(module) + self.assertEqual(self.builds, 1) + + def test_wrong_manifest_is_rejected_before_loading(self): + manifest_path = next((self.root / "artifacts").glob("*/manifest.json")) + manifest = json.loads(manifest_path.read_text()) + manifest["runtime"]["cuda"] = "different" + manifest_path.write_text(json.dumps(manifest)) + self.replay() + with self.assertRaisesRegex(RuntimeError, "manifest does not match"): + self.extension.load_inline("example", "source", cuda_sources="cuda") + + def test_corrupted_artifact_is_rejected_before_loading(self): + next((self.root / "artifacts").glob("*/extension.so")).write_bytes(b"corrupted") + self.replay() + with self.assertRaisesRegex(RuntimeError, "SHA-256"): + self.extension.load_inline("example", "source", cuda_sources="cuda") + + def test_ninja_is_disabled_during_replay(self): + self.replay() + with self.assertRaisesRegex(RuntimeError, "GPU compilation is disabled"): + self.extension._run_ninja_build() + + def test_transfer_contains_only_binary_and_manifest(self): + artifacts = pack_artifacts(self.root / "artifacts") + self.assertEqual(len(artifacts), 2) + destination = self.root / "received" + unpack_artifacts(destination, artifacts) + self.assertEqual(pack_artifacts(destination), artifacts) + + def test_transfer_rejects_path_escape(self): + with self.assertRaisesRegex(ValueError, "Invalid artifact path"): + unpack_artifacts(self.root / "received", {"../escape.so": b"binary"}) From 99c09ef63a43efcfbb5d7e2fc8b51a3d455935ae Mon Sep 17 00:00:00 2001 From: Mark Saroufim Date: Wed, 9 Sep 2026 22:25:33 -0700 Subject: [PATCH 2/4] Compile Python load_inline submissions before allocating a Modal GPU --- docs/modal-inline-artifact-prototype.md | 75 ------- docs/python-cpu-compilation.md | 85 ++++++++ scripts/debug_python_precompile.py | 117 +++++++++++ scripts/modal_inline_artifact_prototype.py | 185 ----------------- src/libkernelbot/inline_artifacts.py | 112 +++++++++-- src/libkernelbot/launchers/modal.py | 45 ++++- src/libkernelbot/python_precompile.py | 190 +++++++++++++++++ src/libkernelbot/run_eval.py | 11 + src/runners/modal_runner.py | 19 +- src/runners/modal_runner_archs.py | 3 +- tests/test_inline_artifacts.py | 28 +++ tests/test_modal.py | 7 + tests/test_python_precompile.py | 224 +++++++++++++++++++++ 13 files changed, 816 insertions(+), 285 deletions(-) delete mode 100644 docs/modal-inline-artifact-prototype.md create mode 100644 docs/python-cpu-compilation.md create mode 100644 scripts/debug_python_precompile.py delete mode 100644 scripts/modal_inline_artifact_prototype.py create mode 100644 src/libkernelbot/python_precompile.py create mode 100644 tests/test_python_precompile.py diff --git a/docs/modal-inline-artifact-prototype.md b/docs/modal-inline-artifact-prototype.md deleted file mode 100644 index 441d52fb..00000000 --- a/docs/modal-inline-artifact-prototype.md +++ /dev/null @@ -1,75 +0,0 @@ -# Compile load_inline submissions before allocating a GPU - -This experiment imports the vector-add submission on a CPU-only Modal worker, -copies its compiled `.so` to a T4 worker, and runs KernelBot's correctness harness. -It leaves the submission's `load_inline` call intact and uses the same runner -image for both workers. The production submission path is unchanged. - -```bash -uv sync --extra dev -PYTHONPATH=src:src/runners uv run modal run \ - scripts/modal_inline_artifact_prototype.py \ - --output /tmp/inline-artifact-result.json -``` - -Set `KERNELBOT_PROTOTYPE_GPU` to `L4`, `A100`, `H100`, or `B200` to change the GPU -and compilation target. T4 is the tested default. The app exits after the run. -The compiler wrapper from the runner image is present, but its PCH Volume is -not mounted: both comparison builds start without a precompiled-header cache. - -The default input is `examples/vectoradd_py/submission_cuda_inline.py`. -`--submission path/to/submission.py` accepts another implementation of that same -vector-add task. The harness tests square matrices with dimensions 1, 127, 128, -129, 256, 512 and 1024. - -## How it works - -A `sitecustomize` hook intercepts `load_inline` in the submission process and -the evaluator's spawned processes. On the CPU worker, it calls the compiler and -saves the resulting library and a manifest. On the GPU worker, it checks the -manifest and binary hash, then imports the saved library. Requests are matched -by source, options, target architecture and Python/PyTorch/CUDA ABI information. -Missing artifacts fail; the hook disables PyTorch's Ninja build entrypoint. - -The script transfers library bytes through Modal results and arguments. It -checks that replay creates no build directory, runs the existing -`run_single_evaluation` harness, and verifies that a missing artifact fails. -It then compiles once on the GPU worker to provide a timing baseline. That last -build is diagnostic overhead, not part of the proposed CPU-build path. - -## Results - -The [T4 run](https://modal.com/apps/coreauto/main/ap-kfECmtFhii1jy58mgwXMYY) passed -all seven correctness cases and the missing-artifact check. Replay created no -build directory. Both workers requested 4 CPU cores and 8 GiB RAM. - -| Step | Extension compile/load | Python process including startup | -| --- | ---: | ---: | -| CPU build, zero visible GPUs | 55.684 s | 58.793 s | -| T4 artifact load | 7.728 ms | 2.722 s | -| Fresh build on the T4 worker | 55.413 s | 58.042 s | - -The transferred library and manifest were 1,560,196 bytes. Evaluation took -5.959 s. These are single-run observations, not throughput or billing estimates. - -The JSON output separates extension load/compile time, Python process time, -evaluation time and total experiment time. Compare compilation and replay -separately from scheduling, Python startup and the diagnostic baseline. - -## Limits - -The submission must compile at import time without querying or using a GPU. -Lazy builds, `load()`, non-Python extensions and libraries outside the common -image are unsupported. Errors propagate without a fallback. - -Artifacts are used within one run; this is not a persistent compilation cache. -The manifest does not fingerprint external headers or every compiler input. -Sharing artifacts between submissions would require a complete build identity. -Only use binaries produced by the builder for the same request: the hash detects -corruption, and the Python hook is not a security boundary. - -Run the adapter tests without CUDA: - -```bash -uv run pytest tests/test_inline_artifacts.py -``` diff --git a/docs/python-cpu-compilation.md b/docs/python-cpu-compilation.md new file mode 100644 index 00000000..a6e5bb7b --- /dev/null +++ b/docs/python-cpu-compilation.md @@ -0,0 +1,85 @@ +# CPU compilation for Python submissions + +Python submissions keep their existing `submission.py` interface. Before allocating +a Modal GPU, the launcher attempts a CPU import for submissions containing +`load_inline`. It transfers any compiled Python extensions to the GPU runner, +which continues through the existing test, benchmark and leaderboard evaluator. +Other Python submissions go directly to the GPU runner. + +## Reuse and fallback + +The GPU subprocess hook replaces a matching `load_inline` call with an import +of its compiled library. It checks source strings, compile options, build +environment, Python/PyTorch/CUDA ABI, and the hashes of compiler-reported header +dependencies. A packet also identifies the task sources and runner image. + +| Submission or build condition | Behavior | +| --- | --- | +| Import-time `load_inline`, compatible inputs | Build on CPU and reuse on GPU | +| CUDA tensors or GPU queries during import | CPU import fails; run normally on GPU | +| Compilation inside `custom_kernel` | Compile normally on GPU when called | +| GPU-dependent source or header changes | Reject that artifact and compile on GPU | +| Custom linker flags (`extra_ldflags`) or non-Python libraries | Existing GPU path | +| Triton, CuTe, plain PyTorch, or Python with no detected `load_inline` | Existing GPU path | +| CPU service failure, timeout, missing function, oversized artifact | Existing GPU path | + +The hook falls back per call, so a submission can reuse one extension and compile +another on the GPU. Failing correctness checks are still failures: fallback does +not retry evaluation or alter its result. Raw CUDA submissions are unchanged. + +The CPU function uses the runner image, 4 CPU cores, 8 GiB RAM and a 120-second +import budget. Containers are single-use, Modal access is restricted, and the +PCH Volume is mounted read-only on both CPU and GPU. Artifact bundles are scoped +to the request, compressed, and limited to 1 MiB in transit / 64 MiB unpacked. +Larger results use the normal GPU path; this avoids Modal's large-result blob +upload from a restricted worker. There is no persistent binary cache. + +`FullResult.cpu_compile` records the CPU duration, artifact count, reuse/fallback +counts and reason. It is diagnostic information, not a correctness signal. + +## Operations + +Deploy the runner before restarting the bot with the new launcher: + +```bash +PYTHONPATH=src:src/runners uv run modal deploy src/runners/modal_runner_archs.py +``` + +`KERNELBOT_CPU_COMPILE=0` on the bot disables the CPU attempt. Set +`KERNELBOT_MODAL_APP` on both deployment and bot to use a separate runner app; +the default remains `discord-bot-runner`. A launcher talking to an older runner +falls back if the CPU function is unavailable. + +## Local debug with real Modal workers + +```bash +uv sync --extra dev +PYTHONPATH=src:src/runners uv run python scripts/debug_python_precompile.py --gpu T4 +``` + +The debug command builds ordinary task configs and calls `ModalLauncher`. Its +function lookup uses the registered functions in a temporary Modal app, so it +does not deploy over the live runner or require a local database. CPU and GPU +execution use the configured real Modal account. Results are saved to +`/tmp/kernelbot-native-cpu-compile.json`. + +The inline case uses the unchanged vector-add example and runs test, benchmark, +and leaderboard modes, with a 256-by-256 benchmark. Additional cases exercise +GPU work during import, lazy compilation, Triton, plain PyTorch, and source +generation and included headers that differ between CPU and GPU. Use +`--cases inline` to run only the first case. + +Verified on a real Tesla T4: the unchanged inline submission passed all three +evaluation phases with six artifact loads and no GPU compilation fallback. +The changed-header case rejected the CPU library and passed after GPU compilation +([Modal run](https://modal.com/apps/coreauto/main/ap-b0aEIY861mwk1q0kVHvblV)). +GPU work during import, lazy compilation, Triton, plain PyTorch, and changed source +also passed through their expected paths +([compatibility run](https://modal.com/apps/coreauto/main/ap-JHjO3BkqWXZC39WN7WWKHs)). +The existing Modal integration tests assert reuse on T4 and H100. + +Local tests: + +```bash +uv run pytest tests/test_python_precompile.py tests/test_inline_artifacts.py tests/test_modal.py -m 'not integration' +``` diff --git a/scripts/debug_python_precompile.py b/scripts/debug_python_precompile.py new file mode 100644 index 00000000..13d675e7 --- /dev/null +++ b/scripts/debug_python_precompile.py @@ -0,0 +1,117 @@ +"""Exercise the native Python launcher against an ephemeral, real Modal app. + +PYTHONPATH=src:src/runners uv run python scripts/debug_python_precompile.py +""" + +import argparse +import asyncio +import dataclasses +import json +import os +import textwrap +from pathlib import Path + +import modal + +os.environ.setdefault("KERNELBOT_MODAL_APP", "kernelbot-cpu-compile-debug") + +from modal_runner import compile_python_submission # noqa: E402 +from modal_runner_archs import app, pytorch_functions # noqa: E402 + +from libkernelbot.consts import GPU_TO_SM, SubmissionMode, get_gpu_by_name # noqa: E402 +from libkernelbot.launchers.modal import ModalLauncher # noqa: E402 +from libkernelbot.report import RunProgressReporter # noqa: E402 +from libkernelbot.task import build_task_config, make_task_definition # noqa: E402 + + +class DebugLauncher(ModalLauncher): + def _get_function(self, name): + # Use the same registered functions without deploying over the live app. + if name == "compile_python_submission": + return compile_python_submission + return pytorch_functions[name.removeprefix("run_pytorch_script_")] + + +class ConsoleReporter(RunProgressReporter): + async def push(self, message): + print(message, flush=True) + + async def update(self, message): + print(message, flush=True) + + +def submissions(example: Path) -> dict[str, tuple[str, str | None]]: + inline = (example / "submission_cuda_inline.py").read_text() + begin = inline.index("add_module = load_inline(") + end = inline.index("\n\n\ndef add(", begin) + lazy = ( + inline[:begin] + + "def build_module():\n" + + textwrap.indent(inline[begin:end].replace("add_module =", "return", 1), " ") + + inline[end:] + ).replace("return add_module.add_cuda(A, B)", "return build_module().add_cuda(A, B)") + return { + "inline": (inline, "reused"), + "gpu-import": ("import torch\ntorch.empty(1, device='cuda')\n" + inline, "fallback"), + "lazy": (lazy, "skipped"), + "triton": ((example / "submission_triton.py").read_text(), None), + "torch": ("def custom_kernel(data):\n return data[0] + data[1]\n", None), + "changed-source": ( + inline.replace( + "add_module = load_inline(", + "add_cuda_source += '// GPU' if torch.cuda.is_available() else '// CPU'\n" + "add_module = load_inline(", + ), + "fallback", + ), + "changed-header": ( + "import torch\nfrom pathlib import Path\n" + "Path('/tmp/kernelbot-debug-header.h').write_text(\n" + " '#define OFFSET ' + ('0' if torch.cuda.is_available() else '1'))\n" + + inline.replace( + 'add_cuda_source = """', + 'add_cuda_source = """\n#include "/tmp/kernelbot-debug-header.h"', + ).replace("C[idx] = A[idx] + B[idx];", "C[idx] = A[idx] + B[idx] + OFFSET;"), + "fallback", + ), + } + + +async def run(args): + example = Path(__file__).resolve().parents[1] / "examples" / "vectoradd_py" + task = make_task_definition(example).task + gpu = get_gpu_by_name(args.gpu) + launcher = DebugLauncher([]) + cases = submissions(example) + results = {} + async with app.run.aio(): + for name in args.cases.split(","): + source, expected = cases[name] + config = build_task_config( + task=task, + submission_content=source, + arch=GPU_TO_SM[gpu.name], + mode=SubmissionMode.LEADERBOARD if name == "inline" else SubmissionMode.TEST, + ) + config["benchmarks"] = [{"size": 256, "seed": 54352}] + result = await launcher.run_submission(config, gpu, ConsoleReporter(name)) + info = result.cpu_compile + print(name, json.dumps(dataclasses.asdict(info) if info else None), flush=True) + assert result.success, result.error + assert all(item.run and item.run.passed for item in result.runs.values()), result + assert (info.status if info else None) == expected, result + results[name] = dataclasses.asdict(result) + if args.output: + Path(args.output).write_text(json.dumps(results, default=str, indent=2) + "\n") + return results + + +if __name__ == "__main__": + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--gpu", default="T4") + parser.add_argument( + "--cases", default="inline,gpu-import,lazy,triton,torch,changed-source,changed-header" + ) + parser.add_argument("--output", default="/tmp/kernelbot-native-cpu-compile.json") + with modal.enable_output(): + asyncio.run(run(parser.parse_args())) diff --git a/scripts/modal_inline_artifact_prototype.py b/scripts/modal_inline_artifact_prototype.py deleted file mode 100644 index 5059dfd2..00000000 --- a/scripts/modal_inline_artifact_prototype.py +++ /dev/null @@ -1,185 +0,0 @@ -"""Run the CPU-build experiment described in docs/modal-inline-artifact-prototype.md.""" - -import dataclasses -import json -import os -import subprocess -import tempfile -import time -from pathlib import Path - -import modal -from modal_runner import cuda_image - -from libkernelbot.inline_artifacts import ( - pack_artifacts, - prepare_environment, - read_events, - unpack_artifacts, -) - -app = modal.App("kernelbot-inline-artifact-prototype") -GPU = os.environ.get("KERNELBOT_PROTOTYPE_GPU", "T4").upper() -ARCH = {"T4": "7.5", "L4": "8.9", "A100": "8.0", "H100": "9.0a", "B200": "10.0a"}[GPU] - - -def _write_sources(work: Path, sources: dict[str, str]) -> None: - for name, source in sources.items(): - path = work / name - if not path.resolve().is_relative_to(work.resolve()): - raise ValueError(f"Invalid source path: {name}") - path.parent.mkdir(parents=True, exist_ok=True) - path.write_text(source) - - -def _import_submission(work: Path, environment: dict) -> dict: - started = time.perf_counter() - result = subprocess.run( - ["python3", "submission.py"], - cwd=work, - env=environment, - capture_output=True, - text=True, - timeout=240, - ) - if result.returncode: - raise RuntimeError(f"Submission import failed:\n{result.stdout}\n{result.stderr}") - return { - "process_seconds": time.perf_counter() - started, - "events": read_events(work / "artifacts"), - } - - -@app.function( - image=cuda_image, cpu=4, memory=8192, timeout=300, max_containers=1, scaledown_window=2 -) -def compile_cpu(sources: dict[str, str], arch: str) -> dict: - import torch - - if torch.cuda.is_available() or torch.cuda.device_count() != 0: - raise RuntimeError("CPU build stage unexpectedly has a GPU") - with tempfile.TemporaryDirectory(prefix="kb-cpu-build-") as directory: - work = Path(directory) - _write_sources(work, sources) - result = _import_submission(work, prepare_environment(work, "capture", arch)) - artifacts = pack_artifacts(work / "artifacts") - if not artifacts: - raise RuntimeError( - "No import-time load_inline calls captured; " - "lazy/GPU-dependent builds are unsupported" - ) - result.update({"cuda_available": False, "device_count": 0, "artifacts": artifacts}) - return result - - -@app.function( - image=cuda_image, gpu=GPU, cpu=4, memory=8192, timeout=600, max_containers=1, scaledown_window=2 -) -def evaluate_gpu( - sources: dict[str, str], artifacts: dict[str, bytes], arch: str, tests: str -) -> dict: - import torch - - from libkernelbot.run_eval import make_system_info, run_single_evaluation - - expected_capability = tuple(int(x) for x in arch.rstrip("a").split(".")) - if torch.cuda.get_device_capability() != expected_capability: - raise RuntimeError("GPU architecture does not match the CPU compilation target") - started = time.perf_counter() - result = {"gpu": torch.cuda.get_device_name(), "capability": expected_capability} - with tempfile.TemporaryDirectory(prefix="kb-gpu-replay-") as directory: - work = Path(directory) - _write_sources(work, sources) - unpack_artifacts(work / "artifacts", artifacts) - environment = prepare_environment(work, "replay", arch) - result["replay_import"] = _import_submission(work, environment) - - # Keep KernelBot's normal test protocol, including its multiprocessing evaluator. - previous_cwd, previous_environment = Path.cwd(), dict(os.environ) - try: - os.chdir(work) - os.environ.update(environment) - run, _ = run_single_evaluation( - ["python3", "eval.py"], - "test", - system=make_system_info(), - tests=tests, - seed=42, - ) - finally: - os.chdir(previous_cwd) - os.environ.clear() - os.environ.update(previous_environment) - result["evaluation"] = dataclasses.asdict(run) - if not run.success or not run.passed: - raise RuntimeError(f"Artifact evaluation failed: {result['evaluation']}") - result["replay_events"] = read_events(work / "artifacts") - if not result["replay_events"] or any( - e["mode"] != "replay" for e in result["replay_events"] - ): - raise RuntimeError("Expected artifact-only GPU execution") - if (work / "build").exists(): - raise RuntimeError("GPU replay unexpectedly created a build directory") - - # Negative control: missing binaries must fail instead of silently recompiling. - missing = work / "missing" - missing.mkdir() - _write_sources(missing, sources) - probe = subprocess.run( - ["python3", "submission.py"], - cwd=missing, - env=prepare_environment(missing, "replay", arch), - capture_output=True, - text=True, - timeout=60, - ) - result["missing_artifact_rejected"] = ( - probe.returncode != 0 and "No CPU-built artifact" in probe.stderr - ) - if not result["missing_artifact_rejected"]: - raise RuntimeError(f"Missing-artifact negative control failed: {probe.stderr}") - - # Compare against a fresh GPU-side build using the same sources and image. - baseline = work / "baseline" - baseline.mkdir() - _write_sources(baseline, sources) - result["gpu_compile_baseline"] = _import_submission( - baseline, prepare_environment(baseline, "capture", arch) - ) - result["gpu_function_seconds"] = time.perf_counter() - started - return result - - -@app.local_entrypoint() -def main(output: str = "", submission: str = ""): - repo = Path(__file__).resolve().parents[1] - example = repo / "examples" / "vectoradd_py" - sources = { - "submission.py": Path(submission).read_text() - if submission - else (example / "submission_cuda_inline.py").read_text(), - "task.py": (example / "task.py").read_text(), - "reference.py": (example / "reference.py").read_text(), - "utils.py": (repo / "examples" / "utils.py").read_text(), - "eval.py": (repo / "examples" / "eval.py").read_text(), - } - tests = ( - "\n".join(f"size: {size}; seed: 4242" for size in (1, 127, 128, 129, 256, 512, 1024)) + "\n" - ) - started = time.perf_counter() - built = compile_cpu.remote(sources, ARCH) - artifacts = built.pop("artifacts") - print("CPU_BUILD", json.dumps(built), flush=True) - print("ARTIFACT_BYTES", sum(map(len, artifacts.values())), flush=True) - evaluated = evaluate_gpu.remote(sources, artifacts, ARCH, tests) - report = { - "requested_gpu": GPU, - "target_arch": ARCH, - "artifact_bytes": sum(map(len, artifacts.values())), - "cpu_build": built, - "gpu": evaluated, - "client_pipeline_seconds": time.perf_counter() - started, - } - print("RESULT", json.dumps(report, indent=2), flush=True) - if output: - Path(output).write_text(json.dumps(report, indent=2) + "\n") diff --git a/src/libkernelbot/inline_artifacts.py b/src/libkernelbot/inline_artifacts.py index 87a1daa8..097785ec 100644 --- a/src/libkernelbot/inline_artifacts.py +++ b/src/libkernelbot/inline_artifacts.py @@ -6,6 +6,7 @@ import json import os import platform +import subprocess import sysconfig import time from pathlib import Path @@ -15,7 +16,7 @@ def _digest(value: bytes) -> str: return hashlib.sha256(value).hexdigest() -def _save_module(directory: Path, module, request: dict, runtime: dict) -> None: +def _save_module(directory: Path, module, request: dict, runtime: dict, work: Path) -> None: directory.mkdir(exist_ok=True) binary = Path(module.__file__).read_bytes() (directory / "extension.so").write_bytes(binary) @@ -26,21 +27,21 @@ def _save_module(directory: Path, module, request: dict, runtime: dict) -> None: "sha256": _digest(binary), "runtime": runtime, "request": request, + "dependencies": _dependencies(Path(module.__file__).parent, work), }, sort_keys=True, ) ) -def _load_module(directory: Path, request: dict, runtime: dict): +def _load_module(directory: Path, request: dict, runtime: dict, work: Path): manifest_path = directory / "manifest.json" if not manifest_path.is_file(): - raise RuntimeError( - f"No CPU-built artifact for {request['name']!r}; GPU compilation is disabled" - ) + raise RuntimeError(f"No CPU-built artifact for {request['name']!r}") manifest = json.loads(manifest_path.read_text()) if manifest["runtime"] != runtime or manifest["request"] != request: raise RuntimeError("CPU-built extension manifest does not match this request") + _check_dependencies(manifest.get("dependencies"), work) binary_path = directory / "extension.so" if _digest(binary_path.read_bytes()) != manifest["sha256"]: raise RuntimeError("CPU-built extension failed SHA-256 verification") @@ -54,15 +55,55 @@ def _refuse_build(*args, **kwargs): raise RuntimeError("GPU compilation is disabled in artifact replay mode") +def _dependencies(build: Path, work: Path) -> list[dict]: + result = subprocess.run( + ["ninja", "-C", str(build), "-t", "deps"], + check=True, + capture_output=True, + text=True, + timeout=10, + ) + paths = {line.strip() for line in result.stdout.splitlines() if line.startswith(" ")} + if not paths: + raise RuntimeError("Compiler dependency information is unavailable") + work = work.resolve() + dependencies = [] + for name in sorted(paths): + path = (build / name).resolve() + if path.parent == build.resolve() and path.name in {"main.cpp", "cuda.cu", "sycl.sycl"}: + continue # Generated from the source strings already in the request key. + relative = path.is_relative_to(work) + dependencies.append( + { + "path": str(path.relative_to(work)) if relative else str(path), + "relative": relative, + "sha256": _digest(path.read_bytes()), + } + ) + return dependencies + + +def _check_dependencies(dependencies: list[dict] | None, work: Path) -> None: + if dependencies is None: + raise RuntimeError("Artifact has no compiler dependency information") + for dependency in dependencies: + path = Path(dependency["path"]) + if dependency["relative"]: + path = work / path + if _digest(path.read_bytes()) != dependency["sha256"]: + raise RuntimeError(f"Build dependency changed: {path}") + + def install() -> None: - """Called by the prototype's sitecustomize, including in spawned evaluators.""" + """Install in submission interpreters, including spawned evaluators.""" import torch from torch.utils import cpp_extension mode = os.environ["KERNELBOT_INLINE_MODE"] - if mode not in {"capture", "replay"}: + if mode not in {"capture", "replay", "auto"}: raise ValueError(f"Unknown inline artifact mode: {mode}") root = Path(os.environ["KERNELBOT_INLINE_ARTIFACTS"]) + work = Path(os.environ["KERNELBOT_INLINE_WORKDIR"]) root.mkdir(parents=True, exist_ok=True) original = cpp_extension.load_inline signature = inspect.signature(original) @@ -72,37 +113,77 @@ def install() -> None: "python_abi": sysconfig.get_config_var("SOABI"), "machine": platform.machine(), "cxx11_abi": torch._C._GLIBCXX_USE_CXX11_ABI, - "arch": os.environ["TORCH_CUDA_ARCH_LIST"], } loaded = {} def load_inline(*args, **kwargs): + try: + return dispatch(*args, **kwargs) + except Exception as exc: + if mode != "auto": + raise + with (root / "events.jsonl").open("a") as log: + log.write(json.dumps({"mode": "fallback", "reason": str(exc)[:500]}) + "\n") + return original(*args, **kwargs) + + def dispatch(*args, **kwargs): bound = signature.bind(*args, **kwargs) bound.apply_defaults() - if not bound.arguments.get("is_python_module", True): - raise ValueError("Prototype supports Python extension modules only") + if not bound.arguments.get("is_python_module", True) or bound.arguments.get( + "extra_ldflags" + ): + raise ValueError("Non-Python modules and custom linker inputs require GPU compilation") # Destination paths and logging do not affect the binary's identity. request = { k: v for k, v in bound.arguments.items() if k not in {"build_directory", "verbose"} } # PyTorch may mutate source lists while injecting headers/bindings. request = json.loads(json.dumps(request)) - key = _digest(json.dumps({"request": request, "runtime": runtime}, sort_keys=True).encode()) + current_runtime = { + **runtime, + "cuda_home": getattr(cpp_extension, "CUDA_HOME", None), + "build_environment": { + name: os.environ.get(name) + for name in ( + "TORCH_CUDA_ARCH_LIST", + "CXX", + "CC", + "CUDA_HOME", + "CPATH", + "CPLUS_INCLUDE_PATH", + "C_INCLUDE_PATH", + "NVCC_PREPEND_FLAGS", + "NVCC_APPEND_FLAGS", + "LIBRARY_PATH", + "LD_LIBRARY_PATH", + ) + }, + } + key = _digest( + json.dumps({"request": request, "runtime": current_runtime}, sort_keys=True).encode() + ) directory = root / key started = time.perf_counter() if mode == "capture": module = original(*args, **kwargs) elapsed = time.perf_counter() - started - _save_module(directory, module, request, runtime) + _save_module(directory, module, request, current_runtime, work) else: if key not in loaded: - loaded[key] = _load_module(directory, request, runtime) + loaded[key] = _load_module(directory, request, current_runtime, work) module = loaded[key] elapsed = time.perf_counter() - started # One append per event also works for the evaluator's spawned processes. with (root / "events.jsonl").open("a") as log: log.write( - json.dumps({"mode": mode, "key": key, "seconds": elapsed, "pid": os.getpid()}) + json.dumps( + { + "mode": "replay" if mode == "auto" else mode, + "key": key, + "seconds": elapsed, + "pid": os.getpid(), + } + ) + "\n" ) return module @@ -113,7 +194,7 @@ def load_inline(*args, **kwargs): def prepare_environment(work: Path, mode: str, arch: str) -> dict[str, str]: - """Install the opt-in hook only in child interpreters of this experiment.""" + """Configure the hook in child interpreters without patching the runner.""" hook = work / "bootstrap" hook.mkdir(exist_ok=True) (hook / "sitecustomize.py").write_text( @@ -128,6 +209,7 @@ def prepare_environment(work: Path, mode: str, arch: str) -> dict[str, str]: return { **os.environ, "KERNELBOT_INLINE_MODE": mode, + "KERNELBOT_INLINE_WORKDIR": str(work), "KERNELBOT_INLINE_ARTIFACTS": str(work / "artifacts"), "TORCH_EXTENSIONS_DIR": str(work / "build"), "TORCH_CUDA_ARCH_LIST": arch, diff --git a/src/libkernelbot/launchers/modal.py b/src/libkernelbot/launchers/modal.py index ce7ae06e..0b01d3b0 100644 --- a/src/libkernelbot/launchers/modal.py +++ b/src/libkernelbot/launchers/modal.py @@ -1,8 +1,11 @@ import asyncio +import os +import time import modal -from libkernelbot.consts import GPU, ModalGPU +from libkernelbot.consts import GPU, ModalGPU, Timeout +from libkernelbot.python_precompile import should_precompile from libkernelbot.report import RunProgressReporter from libkernelbot.run_eval import FullResult from libkernelbot.utils import setup_logging @@ -13,9 +16,13 @@ class ModalLauncher(Launcher): - def __init__(self, add_include_dirs: list): + def __init__(self, add_include_dirs: list, *, app_name: str | None = None): super().__init__("Modal", gpus=ModalGPU) self.additional_include_dirs = add_include_dirs + self.app_name = app_name or os.environ.get("KERNELBOT_MODAL_APP", "discord-bot-runner") + + def _get_function(self, name: str): + return modal.Function.from_name(self.app_name, name) async def run_submission( self, config: dict, gpu_type: GPU, status: RunProgressReporter @@ -28,8 +35,34 @@ async def run_submission( await status.push("⏳ Waiting for Modal run to finish...") - function = modal.Function.from_name("discord-bot-runner", func_name) + if os.environ.get("KERNELBOT_CPU_COMPILE", "1") != "0" and should_precompile(config): + started = time.perf_counter() + try: + compiler = self._get_function("compile_python_submission") + packet = await asyncio.wait_for( + compiler.remote.aio(config=config), + timeout=Timeout.COMPILE + 60, + ) + if not isinstance(packet, dict) or not isinstance(packet.get("info"), dict): + raise ValueError("Invalid CPU compilation response") + if "artifacts" not in packet: + raise ValueError("CPU compilation response has no artifact field") + except Exception as exc: + # Old deployments and unavailable CPU workers retain the existing GPU path. + packet = { + "artifacts": {}, + "info": { + "status": "fallback", + "duration": time.perf_counter() - started, + "reason": str(exc)[-1000:], + }, + } + logger.info("Modal CPU compilation: %s", packet["info"]) + config = {**config, "cpu_compile": packet} + function = self._get_function(func_name) result = await function.remote.aio(config=config) + if getattr(result, "cpu_compile", None) is not None: + logger.info("Modal CPU compilation outcome: %s", result.cpu_compile) await status.update("✅ Waiting for modal run to finish... Done") @@ -42,7 +75,7 @@ async def run_validation(self, config: dict, gpu_type: GPU) -> dict: func_name, config.get("version"), ) - function = modal.Function.from_name("discord-bot-runner", func_name) + function = self._get_function(func_name) return await function.remote.aio(config=config) def _function_name(self, config: dict, gpu_type: GPU) -> str: @@ -58,9 +91,7 @@ async def get_queue_status( try: stats = await loop.run_in_executor( None, - lambda: modal.Function.from_name( - "discord-bot-runner", func_name - ).get_current_stats(), + lambda: self._get_function(func_name).get_current_stats(), ) except Exception as e: logger.warning("Could not get Modal queue stats for %s", func_name, exc_info=e) diff --git a/src/libkernelbot/python_precompile.py b/src/libkernelbot/python_precompile.py new file mode 100644 index 00000000..7c2f6e66 --- /dev/null +++ b/src/libkernelbot/python_precompile.py @@ -0,0 +1,190 @@ +"""Best-effort CPU compilation for ordinary Python submissions.""" + +import dataclasses +import hashlib +import io +import json +import os +import signal +import subprocess +import tempfile +import time +import zipfile +from contextlib import chdir +from pathlib import Path + +from libkernelbot.consts import Timeout +from libkernelbot.inline_artifacts import ( + pack_artifacts, + prepare_environment, + read_events, + unpack_artifacts, +) +from libkernelbot.run_eval import CPUCompileInfo + +MAX_ARTIFACT_BYTES = 64 * 1024 * 1024 +MAX_TRANSFER_BYTES = 1024 * 1024 +PROTOCOL_VERSION = 1 + + +def encode_artifacts(artifacts: dict[str, bytes]) -> bytes: + if not artifacts: + return b"" + if sum(map(len, artifacts.values())) > MAX_ARTIFACT_BYTES: + raise ValueError("Compiled artifacts exceed the unpacked size limit") + buffer = io.BytesIO() + with zipfile.ZipFile(buffer, "w", compression=zipfile.ZIP_DEFLATED) as archive: + for name, data in artifacts.items(): + archive.writestr(name, data) + encoded = buffer.getvalue() + # Restricted Modal workers cannot create large-result blobs. Stay below the + # SDK's 2 MiB inline result limit, including metadata and serialization. + if len(encoded) > MAX_TRANSFER_BYTES: + raise ValueError("Compiled artifacts exceed the inline transfer limit") + return encoded + + +def decode_artifacts(encoded: bytes) -> dict[str, bytes]: + if len(encoded) > MAX_TRANSFER_BYTES: + raise ValueError("Artifact transfer exceeds the size limit") + with zipfile.ZipFile(io.BytesIO(encoded)) as archive: + if sum(item.file_size for item in archive.infolist()) > MAX_ARTIFACT_BYTES: + raise ValueError("Unpacked artifacts exceed the size limit") + return {name: archive.read(name) for name in archive.namelist()} + + +def should_precompile(config: dict) -> bool: + return config.get("lang") == "py" and any( + "load_inline" in source + for name, source in config.get("sources", {}).items() + if name.endswith(".py") + ) + + +def source_hash(config: dict) -> str: + return hashlib.sha256(json.dumps(config["sources"], sort_keys=True).encode()).hexdigest() + + +def image_fingerprint() -> str: + return Path("/opt/kernelbot-pch/fingerprint").read_text().strip() + + +def cuda_arch(sm: str) -> str: + # task configs use nvcc's SM spelling, e.g. 90a or 100. + suffix = "a" if sm.endswith("a") else "" + digits = sm.removesuffix("a") + if not digits.isdigit() or len(digits) < 2: + raise ValueError(f"Invalid CUDA target: {sm}") + return f"{int(digits[:-1])}.{digits[-1]}{suffix}" + + +def _import_submission(work: Path, environment: dict, timeout: float): + with subprocess.Popen( + ["python3", "submission.py"], + cwd=work, + env=environment, + stdout=subprocess.PIPE, + stderr=subprocess.PIPE, + text=True, + start_new_session=True, + ) as process: + try: + stdout, stderr = process.communicate(timeout=timeout) + except subprocess.TimeoutExpired: + try: + os.killpg(process.pid, signal.SIGKILL) + except ProcessLookupError: + pass + process.communicate() + raise + if process.returncode: + raise RuntimeError(f"CPU import exited {process.returncode}: {(stderr or stdout)[-1000:]}") + + +def compile_python(config: dict) -> dict: + started = time.perf_counter() + info = CPUCompileInfo(status="skipped") + packet = {"version": PROTOCOL_VERSION, "artifacts": b"", "info": {}} + try: + if not should_precompile(config): + return packet + import torch + + if torch.cuda.is_available() or torch.cuda.device_count(): + raise RuntimeError("CPU compilation worker has a visible GPU") + packet.update(source_hash=source_hash(config), image_fingerprint=image_fingerprint()) + with tempfile.TemporaryDirectory(prefix="kernelbot-cpu-") as directory: + work = Path(directory) + for name, text in config["sources"].items(): + path = work / name + if not path.resolve().is_relative_to(work): + raise ValueError(f"Invalid source path: {name}") + path.parent.mkdir(parents=True, exist_ok=True) + path.write_text(text) + environment = prepare_environment(work, "capture", cuda_arch(config["arch"])) + environment["POPCORN_SEED"] = "1" + _import_submission(work, environment, Timeout.COMPILE) + artifacts = pack_artifacts(work / "artifacts") + packet["artifacts"] = encode_artifacts(artifacts) + info.artifacts = sum(name.endswith("/manifest.json") for name in artifacts) + info.status = "compiled" if info.artifacts else "skipped" + if not info.artifacts: + info.reason = "No import-time load_inline library" + except Exception as exc: + info.status, info.reason = "fallback", str(exc)[-1000:] + packet["artifacts"] = b"" + finally: + info.duration = time.perf_counter() - started + packet["info"] = dataclasses.asdict(info) + return packet + + +def run_with_cpu_artifacts(config: dict, run): + packet = config.get("cpu_compile") + if config.get("lang") != "py" or not packet: + return run(config) + try: + info = CPUCompileInfo(**packet["info"]) + except (KeyError, TypeError): + result = run(config) + result.cpu_compile = CPUCompileInfo("fallback", reason="Invalid CPU build metadata") + return result + if not packet.get("artifacts"): + result = run(config) + result.cpu_compile = info + return result + + with tempfile.TemporaryDirectory(prefix="kernelbot-gpu-") as directory: + work = Path(directory) + try: + if packet["version"] != PROTOCOL_VERSION: + raise ValueError("CPU artifact protocol changed") + if packet["source_hash"] != source_hash(config): + raise ValueError("Submission sources changed after CPU compilation") + if packet["image_fingerprint"] != image_fingerprint(): + raise ValueError("CPU and GPU runner images differ") + unpack_artifacts(work / "artifacts", decode_artifacts(packet["artifacts"])) + environment = prepare_environment(work, "auto", cuda_arch(config["arch"])) + except Exception as exc: + info.status, info.reason = "fallback", str(exc)[-1000:] + result = run(config) + result.cpu_compile = info + return result + previous_environment = dict(os.environ) + try: + with chdir(work): + os.environ.update(environment) + result = run(config) + finally: + os.environ.clear() + os.environ.update(previous_environment) + events = read_events(work / "artifacts") + info.reused = sum(event["mode"] == "replay" for event in events) + info.fallbacks = sum(event["mode"] == "fallback" for event in events) + info.status = "reused" if info.reused else "fallback" + if info.reused and info.fallbacks: + info.status = "partial" + reasons = [event["reason"] for event in events if event["mode"] == "fallback"] + info.reason = reasons[-1] if reasons else "" + result.cpu_compile = info + return result diff --git a/src/libkernelbot/run_eval.py b/src/libkernelbot/run_eval.py index aec59f95..dfabd24e 100644 --- a/src/libkernelbot/run_eval.py +++ b/src/libkernelbot/run_eval.py @@ -82,6 +82,16 @@ class EvalResult: # fmt: on +@dataclasses.dataclass +class CPUCompileInfo: + status: str + duration: float = 0.0 + artifacts: int = 0 + reused: int = 0 + fallbacks: int = 0 + reason: str = "" + + @dataclasses.dataclass class FullResult: # fmt: off @@ -91,6 +101,7 @@ class FullResult: # results of running. There can be multiple runs in one submission, using separate # 'test' and 'benchmark' keys, for example runs: dict[str, EvalResult] = dataclasses.field(default_factory=dict) + cpu_compile: CPUCompileInfo | None = None # fmt: on diff --git a/src/runners/modal_runner.py b/src/runners/modal_runner.py index b177f99d..4a4c01b1 100644 --- a/src/runners/modal_runner.py +++ b/src/runners/modal_runner.py @@ -7,11 +7,13 @@ from modal import App, Image, Volume +from libkernelbot.consts import Timeout +from libkernelbot.python_precompile import compile_python, run_with_cpu_artifacts from libkernelbot.run_eval import FullResult, SystemInfo, run_config # Create a stub for the Modal app # IMPORTANT: This has to stay in separate file or modal breaks -app = App("discord-bot-runner") +app = App(os.environ.get("KERNELBOT_MODAL_APP", "discord-bot-runner")) cuda_version = "13.3.0" flavor = "devel" operating_sys = "ubuntu24.04" @@ -173,7 +175,7 @@ def modal_run_config( # noqa: C901 """Modal version of run_pytorch_script, handling timeouts""" try: with timeout(timeout_seconds): - return run_config(config) + return run_with_cpu_artifacts(config, run_config) except TimeoutException as e: return FullResult( success=False, @@ -191,6 +193,19 @@ def modal_run_config( # noqa: C901 ) +@app.function( + image=cuda_image, + cpu=4, + memory=8192, + timeout=Timeout.COMPILE + 30, + single_use_containers=True, + restrict_modal_access=True, + volumes={PCH_MOUNT: pch_volume.with_mount_options(read_only=True)}, +) +def compile_python_submission(config: dict) -> dict: + return compile_python(config) + + @app.function( image=cuda_image, cpu=4, diff --git a/src/runners/modal_runner_archs.py b/src/runners/modal_runner_archs.py index c3381305..64f327f1 100644 --- a/src/runners/modal_runner_archs.py +++ b/src/runners/modal_runner_archs.py @@ -5,6 +5,7 @@ from libkernelbot.validation_runtime import run_validation_config gpus = ["T4", "L4", "L4:4", "A100-80GB", "H100!", "B200"] +pytorch_functions = {} for gpu in gpus: gpu_slug = gpu.lower().split("-")[0].strip("!").replace(":", "x") app.function( @@ -16,7 +17,7 @@ volumes={PCH_MOUNT: pch_volume.with_mount_options(read_only=True)}, timeout=MODAL_RUN_TIMEOUT_SECONDS, )(modal_run_config) - app.function( + pytorch_functions[gpu_slug] = app.function( gpu=gpu, image=cuda_image, name=f"run_pytorch_script_{gpu_slug}", diff --git a/tests/test_inline_artifacts.py b/tests/test_inline_artifacts.py index 88222848..6eb19e51 100644 --- a/tests/test_inline_artifacts.py +++ b/tests/test_inline_artifacts.py @@ -26,6 +26,7 @@ def load_inline( cpp_sources, cuda_sources=None, extra_cuda_cflags=None, + extra_ldflags=None, is_python_module=True, verbose=False, build_directory=None, @@ -54,11 +55,15 @@ def load_inline( { "KERNELBOT_INLINE_MODE": "capture", "KERNELBOT_INLINE_ARTIFACTS": str(self.root / "artifacts"), + "KERNELBOT_INLINE_WORKDIR": str(self.root), "TORCH_CUDA_ARCH_LIST": "7.5", }, ) environment.start() self.addCleanup(environment.stop) + dependencies = patch("libkernelbot.inline_artifacts._dependencies", return_value=[]) + dependencies.start() + self.addCleanup(dependencies.stop) install() self.extension.load_inline("example", "source", cuda_sources="cuda") @@ -118,6 +123,29 @@ def test_ninja_is_disabled_during_replay(self): with self.assertRaisesRegex(RuntimeError, "GPU compilation is disabled"): self.extension._run_ninja_build() + def test_native_mode_falls_back_to_original_compiler(self): + self.extension.load_inline = self.original + os.environ["KERNELBOT_INLINE_MODE"] = "auto" + install() + self.extension.load_inline("example", "changed source", cuda_sources="cuda") + self.assertEqual(self.builds, 2) + + def test_environment_changes_after_import_do_not_reuse_binary(self): + self.replay() + os.environ["TORCH_CUDA_ARCH_LIST"] = "9.0a" + with self.assertRaisesRegex(RuntimeError, "No CPU-built artifact"): + self.extension.load_inline("example", "source", cuda_sources="cuda") + + def test_custom_linker_inputs_are_left_to_the_gpu(self): + with self.assertRaisesRegex(ValueError, "custom linker inputs"): + self.extension.load_inline("linked", "source", extra_ldflags=["-lcustom"]) + self.assertEqual(self.builds, 1) + self.extension.load_inline = self.original + os.environ["KERNELBOT_INLINE_MODE"] = "auto" + install() + self.extension.load_inline("linked", "source", extra_ldflags=["-lcustom"]) + self.assertEqual(self.builds, 2) + def test_transfer_contains_only_binary_and_manifest(self): artifacts = pack_artifacts(self.root / "artifacts") self.assertEqual(len(artifacts), 2) diff --git a/tests/test_modal.py b/tests/test_modal.py index f3ac827b..241edfb5 100644 --- a/tests/test_modal.py +++ b/tests/test_modal.py @@ -233,6 +233,13 @@ async def test_modal_launcher_python_script( assert result.error == "" assert isinstance(result.runs, dict) + if task[1] == "submission_cuda_inline.py": + assert result.cpu_compile.status == "reused", result.cpu_compile + assert result.cpu_compile.reused >= 2 + assert result.cpu_compile.fallbacks == 0 + else: + assert result.cpu_compile is None + # System info - test actual expected values assert gpu_type.name in result.system.gpu assert "Linux" in result.system.platform diff --git a/tests/test_python_precompile.py b/tests/test_python_precompile.py new file mode 100644 index 00000000..40098d39 --- /dev/null +++ b/tests/test_python_precompile.py @@ -0,0 +1,224 @@ +import asyncio +import json +import os +import subprocess +import sys +from pathlib import Path +from types import SimpleNamespace +from unittest.mock import AsyncMock, MagicMock, patch + +import pytest + +from libkernelbot.consts import get_gpu_by_name +from libkernelbot.inline_artifacts import _check_dependencies, _dependencies +from libkernelbot.launchers.modal import ModalLauncher +from libkernelbot.python_precompile import ( + compile_python, + cuda_arch, + decode_artifacts, + encode_artifacts, + run_with_cpu_artifacts, + should_precompile, + source_hash, +) +from libkernelbot.run_eval import FullResult, SystemInfo + + +@pytest.fixture +def config(): + return { + "lang": "py", + "arch": "75", + "main": "eval.py", + "mode": "test", + "sources": { + "submission.py": "from torch.utils.cpp_extension import load_inline", + "eval.py": "", + }, + } + + +@pytest.fixture +def result(): + return FullResult(success=True, error="", system=SystemInfo()) + + +@pytest.mark.parametrize("sm,expected", [("75", "7.5"), ("90a", "9.0a"), ("100", "10.0")]) +def test_target_architecture(sm, expected): + assert cuda_arch(sm) == expected + + +def test_only_python_load_inline_submissions_are_candidates(config): + assert should_precompile(config) + assert not should_precompile({**config, "lang": "cu"}) + assert not should_precompile({**config, "sources": {"submission.py": "import triton"}}) + + +@pytest.mark.asyncio +async def test_launcher_compiles_before_gpu_and_leaves_input_config_unchanged(config, result): + cpu, gpu = MagicMock(), MagicMock() + packet = {"artifacts": {"module": b"binary"}, "info": {"status": "compiled"}} + cpu.remote.aio = AsyncMock(return_value=packet) + gpu.remote.aio = AsyncMock(return_value=result) + order = [] + + def lookup(app, name): + assert app == "debug-runner" + order.append(name) + if name == "compile_python_submission": + return cpu + cpu.remote.aio.assert_awaited_once_with(config=config) + return gpu + + with patch("modal.Function.from_name", side_effect=lookup): + actual = await ModalLauncher([], app_name="debug-runner").run_submission( + config, + get_gpu_by_name("T4"), + AsyncMock(), + ) + assert actual is result + assert order == ["compile_python_submission", "run_pytorch_script_t4"] + assert "cpu_compile" not in config + gpu.remote.aio.assert_awaited_once_with(config={**config, "cpu_compile": packet}) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("error", [RuntimeError("old deployment"), TimeoutError("CPU timed out")]) +async def test_cpu_service_failure_still_runs_gpu(config, result, error): + cpu, gpu = MagicMock(), MagicMock() + cpu.remote.aio = AsyncMock(side_effect=error) + gpu.remote.aio = AsyncMock(return_value=result) + with patch("modal.Function.from_name", side_effect=[cpu, gpu]): + actual = await ModalLauncher([]).run_submission(config, get_gpu_by_name("T4"), AsyncMock()) + assert actual is result + assert gpu.remote.aio.call_args.kwargs["config"]["cpu_compile"]["info"]["status"] == "fallback" + + +@pytest.mark.asyncio +async def test_cancellation_does_not_launch_gpu(config): + cpu = MagicMock() + cpu.remote.aio = AsyncMock(side_effect=asyncio.CancelledError()) + with patch("modal.Function.from_name", return_value=cpu) as lookup: + with pytest.raises(asyncio.CancelledError): + await ModalLauncher([]).run_submission(config, get_gpu_by_name("T4"), AsyncMock()) + assert lookup.call_count == 1 + + +@pytest.mark.asyncio +@pytest.mark.parametrize("disabled", [True, False]) +async def test_ordinary_python_or_operator_disable_bypasses_cpu( + config, result, monkeypatch, disabled +): + if disabled: + monkeypatch.setenv("KERNELBOT_CPU_COMPILE", "0") + else: + config["sources"]["submission.py"] = "import torch" + gpu = MagicMock() + gpu.remote.aio = AsyncMock(return_value=result) + with patch("modal.Function.from_name", return_value=gpu) as lookup: + await ModalLauncher([]).run_submission(config, get_gpu_by_name("T4"), AsyncMock()) + lookup.assert_called_once_with("discord-bot-runner", "run_pytorch_script_t4") + gpu.remote.aio.assert_awaited_once_with(config=config) + + +@pytest.mark.parametrize( + "error", [RuntimeError("CUDA unavailable"), subprocess.TimeoutExpired("python", 1)] +) +def test_cpu_import_failures_return_fallback(config, monkeypatch, error): + monkeypatch.setitem( + sys.modules, + "torch", + SimpleNamespace( + cuda=SimpleNamespace( + is_available=lambda: False, + device_count=lambda: 0, + ) + ), + ) + with ( + patch("libkernelbot.python_precompile.image_fingerprint", return_value="image"), + patch("libkernelbot.python_precompile._import_submission", side_effect=error), + ): + packet = compile_python(config) + assert packet["info"]["status"] == "fallback" + assert packet["artifacts"] == b"" + + +@pytest.mark.parametrize("changed", ["source_hash", "image_fingerprint", "version"]) +def test_gpu_rejects_stale_build_without_failing_submission(config, result, changed): + packet = { + "version": 1, + "source_hash": source_hash(config), + "image_fingerprint": "image", + "artifacts": {"unused.so": b"binary"}, + "info": {"status": "compiled"}, + } + packet[changed] = "different" + run = MagicMock(return_value=result) + with patch("libkernelbot.python_precompile.image_fingerprint", return_value="image"): + actual = run_with_cpu_artifacts({**config, "cpu_compile": packet}, run) + assert actual is result and actual.cpu_compile.status == "fallback" + run.assert_called_once() + + +@pytest.mark.parametrize("success", [True, False]) +def test_gpu_context_restores_environment_and_records_reuse(config, result, success): + result.success = success + config["cpu_compile"] = { + "version": 1, + "source_hash": source_hash(config), + "image_fingerprint": "image", + "artifacts": encode_artifacts({"extension.so": b"binary"}), + "info": {"status": "compiled"}, + } + previous_cwd, previous_environment = Path.cwd(), dict(os.environ) + + def run(_): + assert os.environ["KERNELBOT_INLINE_MODE"] == "auto" + root = Path(os.environ["KERNELBOT_INLINE_ARTIFACTS"]) + (root / "events.jsonl").write_text(json.dumps({"mode": "replay"}) + "\n") + return result + + with patch("libkernelbot.python_precompile.image_fingerprint", return_value="image"): + actual = run_with_cpu_artifacts(config, run) + assert actual.cpu_compile.status == "reused" and actual.cpu_compile.reused == 1 + assert actual.success == success + assert Path.cwd() == previous_cwd and dict(os.environ) == previous_environment + + +def test_compressed_transport_roundtrip(): + artifacts = {"key/extension.so": b"binary" * 1000, "key/manifest.json": b"{}"} + assert decode_artifacts(encode_artifacts(artifacts)) == artifacts + + +def test_transport_size_limit_uses_fallback_instead_of_large_modal_blobs(monkeypatch): + monkeypatch.setattr("libkernelbot.python_precompile.MAX_TRANSFER_BYTES", 1) + with pytest.raises(ValueError, match="inline transfer limit"): + encode_artifacts({"key/extension.so": b"binary"}) + + +def test_header_dependencies_relocate_and_detect_changes(tmp_path, monkeypatch): + cpu, gpu = tmp_path / "cpu", tmp_path / "gpu" + build = cpu / "build" + build.mkdir(parents=True) + gpu.mkdir() + (cpu / "kernel.h").write_text("#define VALUE 1") + (gpu / "kernel.h").write_text("#define VALUE 1") + monkeypatch.setenv("KERNELBOT_INLINE_WORKDIR", str(cpu)) + output = f"main.o: #deps 2\n {build / 'main.cpp'}\n {cpu / 'kernel.h'}\n" + with patch( + "libkernelbot.inline_artifacts.subprocess.run", return_value=SimpleNamespace(stdout=output) + ): + dependencies = _dependencies(build, cpu) + assert len(dependencies) == 1 + monkeypatch.setenv("KERNELBOT_INLINE_WORKDIR", str(gpu)) + _check_dependencies(dependencies, gpu) + (gpu / "kernel.h").write_text("#define VALUE 2") + with pytest.raises(RuntimeError, match="Build dependency changed"): + _check_dependencies(dependencies, gpu) + + +def test_missing_dependency_information_cannot_be_reused(tmp_path, monkeypatch): + monkeypatch.setenv("KERNELBOT_INLINE_WORKDIR", str(tmp_path)) + with pytest.raises(RuntimeError, match="no compiler dependency"): + _check_dependencies(None, tmp_path) From ffaf9b176a78db1ab82702dcd0d34809e609a8df Mon Sep 17 00:00:00 2001 From: Mark Saroufim Date: Wed, 9 Sep 2026 22:37:46 -0700 Subject: [PATCH 3/4] Simplify CPU artifact transfer and GPU fallback --- src/libkernelbot/inline_artifacts.py | 104 +++---------- src/libkernelbot/launchers/modal.py | 7 +- src/libkernelbot/python_precompile.py | 175 ++++++++++----------- tests/test_inline_artifacts.py | 211 ++++++++------------------ tests/test_python_precompile.py | 68 +++++---- 5 files changed, 204 insertions(+), 361 deletions(-) diff --git a/src/libkernelbot/inline_artifacts.py b/src/libkernelbot/inline_artifacts.py index 097785ec..e7f2ece5 100644 --- a/src/libkernelbot/inline_artifacts.py +++ b/src/libkernelbot/inline_artifacts.py @@ -8,7 +8,6 @@ import platform import subprocess import sysconfig -import time from pathlib import Path @@ -16,7 +15,7 @@ def _digest(value: bytes) -> str: return hashlib.sha256(value).hexdigest() -def _save_module(directory: Path, module, request: dict, runtime: dict, work: Path) -> None: +def _save_module(directory: Path, module, work: Path) -> None: directory.mkdir(exist_ok=True) binary = Path(module.__file__).read_bytes() (directory / "extension.so").write_bytes(binary) @@ -25,8 +24,7 @@ def _save_module(directory: Path, module, request: dict, runtime: dict, work: Pa { "module_name": module.__name__, "sha256": _digest(binary), - "runtime": runtime, - "request": request, + "key": directory.name, "dependencies": _dependencies(Path(module.__file__).parent, work), }, sort_keys=True, @@ -34,12 +32,9 @@ def _save_module(directory: Path, module, request: dict, runtime: dict, work: Pa ) -def _load_module(directory: Path, request: dict, runtime: dict, work: Path): - manifest_path = directory / "manifest.json" - if not manifest_path.is_file(): - raise RuntimeError(f"No CPU-built artifact for {request['name']!r}") - manifest = json.loads(manifest_path.read_text()) - if manifest["runtime"] != runtime or manifest["request"] != request: +def _load_module(directory: Path, work: Path): + manifest = json.loads((directory / "manifest.json").read_text()) + if manifest["key"] != directory.name: raise RuntimeError("CPU-built extension manifest does not match this request") _check_dependencies(manifest.get("dependencies"), work) binary_path = directory / "extension.so" @@ -51,11 +46,7 @@ def _load_module(directory: Path, request: dict, runtime: dict, work: Path): return module -def _refuse_build(*args, **kwargs): - raise RuntimeError("GPU compilation is disabled in artifact replay mode") - - -def _dependencies(build: Path, work: Path) -> list[dict]: +def _dependencies(build: Path, work: Path) -> dict[str, str]: result = subprocess.run( ["ninja", "-C", str(build), "-t", "deps"], check=True, @@ -67,30 +58,22 @@ def _dependencies(build: Path, work: Path) -> list[dict]: if not paths: raise RuntimeError("Compiler dependency information is unavailable") work = work.resolve() - dependencies = [] + dependencies = {} for name in sorted(paths): path = (build / name).resolve() if path.parent == build.resolve() and path.name in {"main.cpp", "cuda.cu", "sycl.sycl"}: continue # Generated from the source strings already in the request key. - relative = path.is_relative_to(work) - dependencies.append( - { - "path": str(path.relative_to(work)) if relative else str(path), - "relative": relative, - "sha256": _digest(path.read_bytes()), - } - ) + name = str(path.relative_to(work)) if path.is_relative_to(work) else str(path) + dependencies[name] = _digest(path.read_bytes()) return dependencies -def _check_dependencies(dependencies: list[dict] | None, work: Path) -> None: +def _check_dependencies(dependencies: dict[str, str] | None, work: Path) -> None: if dependencies is None: raise RuntimeError("Artifact has no compiler dependency information") - for dependency in dependencies: - path = Path(dependency["path"]) - if dependency["relative"]: - path = work / path - if _digest(path.read_bytes()) != dependency["sha256"]: + for name, digest in dependencies.items(): + path = work / name + if _digest(path.read_bytes()) != digest: raise RuntimeError(f"Build dependency changed: {path}") @@ -100,10 +83,8 @@ def install() -> None: from torch.utils import cpp_extension mode = os.environ["KERNELBOT_INLINE_MODE"] - if mode not in {"capture", "replay", "auto"}: - raise ValueError(f"Unknown inline artifact mode: {mode}") - root = Path(os.environ["KERNELBOT_INLINE_ARTIFACTS"]) work = Path(os.environ["KERNELBOT_INLINE_WORKDIR"]) + root = work / "artifacts" root.mkdir(parents=True, exist_ok=True) original = cpp_extension.load_inline signature = inspect.signature(original) @@ -116,14 +97,17 @@ def install() -> None: } loaded = {} + def record(mode, **details): + with (root / "events.jsonl").open("a") as log: + log.write(json.dumps({"mode": mode, **details}) + "\n") + def load_inline(*args, **kwargs): try: return dispatch(*args, **kwargs) except Exception as exc: if mode != "auto": raise - with (root / "events.jsonl").open("a") as log: - log.write(json.dumps({"mode": "fallback", "reason": str(exc)[:500]}) + "\n") + record("fallback", reason=str(exc)[:500]) return original(*args, **kwargs) def dispatch(*args, **kwargs): @@ -137,8 +121,6 @@ def dispatch(*args, **kwargs): request = { k: v for k, v in bound.arguments.items() if k not in {"build_directory", "verbose"} } - # PyTorch may mutate source lists while injecting headers/bindings. - request = json.loads(json.dumps(request)) current_runtime = { **runtime, "cuda_home": getattr(cpp_extension, "CUDA_HOME", None), @@ -159,38 +141,22 @@ def dispatch(*args, **kwargs): ) }, } + # Hash before PyTorch mutates source lists to inject headers/bindings. key = _digest( json.dumps({"request": request, "runtime": current_runtime}, sort_keys=True).encode() ) directory = root / key - started = time.perf_counter() if mode == "capture": module = original(*args, **kwargs) - elapsed = time.perf_counter() - started - _save_module(directory, module, request, current_runtime, work) + _save_module(directory, module, work) else: if key not in loaded: - loaded[key] = _load_module(directory, request, current_runtime, work) + loaded[key] = _load_module(directory, work) module = loaded[key] - elapsed = time.perf_counter() - started - # One append per event also works for the evaluator's spawned processes. - with (root / "events.jsonl").open("a") as log: - log.write( - json.dumps( - { - "mode": "replay" if mode == "auto" else mode, - "key": key, - "seconds": elapsed, - "pid": os.getpid(), - } - ) - + "\n" - ) + record("replay") return module cpp_extension.load_inline = load_inline - if mode == "replay": - cpp_extension._run_ninja_build = _refuse_build def prepare_environment(work: Path, mode: str, arch: str) -> dict[str, str]: @@ -210,7 +176,6 @@ def prepare_environment(work: Path, mode: str, arch: str) -> dict[str, str]: **os.environ, "KERNELBOT_INLINE_MODE": mode, "KERNELBOT_INLINE_WORKDIR": str(work), - "KERNELBOT_INLINE_ARTIFACTS": str(work / "artifacts"), "TORCH_EXTENSIONS_DIR": str(work / "build"), "TORCH_CUDA_ARCH_LIST": arch, "MAX_JOBS": "2", @@ -226,26 +191,3 @@ def prepare_environment(work: Path, mode: str, arch: str) -> dict[str, str]: ) ), } - - -def pack_artifacts(root: Path) -> dict[str, bytes]: - """Transfer only the libraries and manifests, never the Ninja build cache.""" - return { - str(path.relative_to(root)): path.read_bytes() - for pattern in ("*/manifest.json", "*/extension.so") - for path in root.glob(pattern) - } - - -def unpack_artifacts(root: Path, files: dict[str, bytes]) -> None: - for name, data in files.items(): - path = root / name - if not path.resolve().is_relative_to(root.resolve()): - raise ValueError(f"Invalid artifact path: {name}") - path.parent.mkdir(parents=True, exist_ok=True) - path.write_bytes(data) - - -def read_events(root: Path) -> list[dict]: - path = root / "events.jsonl" - return [json.loads(line) for line in path.read_text().splitlines()] if path.exists() else [] diff --git a/src/libkernelbot/launchers/modal.py b/src/libkernelbot/launchers/modal.py index 0b01d3b0..a4b073c0 100644 --- a/src/libkernelbot/launchers/modal.py +++ b/src/libkernelbot/launchers/modal.py @@ -43,21 +43,16 @@ async def run_submission( compiler.remote.aio(config=config), timeout=Timeout.COMPILE + 60, ) - if not isinstance(packet, dict) or not isinstance(packet.get("info"), dict): - raise ValueError("Invalid CPU compilation response") - if "artifacts" not in packet: - raise ValueError("CPU compilation response has no artifact field") except Exception as exc: # Old deployments and unavailable CPU workers retain the existing GPU path. packet = { - "artifacts": {}, + "artifacts": b"", "info": { "status": "fallback", "duration": time.perf_counter() - started, "reason": str(exc)[-1000:], }, } - logger.info("Modal CPU compilation: %s", packet["info"]) config = {**config, "cpu_compile": packet} function = self._get_function(func_name) result = await function.remote.aio(config=config) diff --git a/src/libkernelbot/python_precompile.py b/src/libkernelbot/python_precompile.py index 7c2f6e66..94e44429 100644 --- a/src/libkernelbot/python_precompile.py +++ b/src/libkernelbot/python_precompile.py @@ -10,47 +10,51 @@ import tempfile import time import zipfile -from contextlib import chdir +from contextlib import chdir, contextmanager, suppress from pathlib import Path from libkernelbot.consts import Timeout -from libkernelbot.inline_artifacts import ( - pack_artifacts, - prepare_environment, - read_events, - unpack_artifacts, -) +from libkernelbot.inline_artifacts import prepare_environment from libkernelbot.run_eval import CPUCompileInfo MAX_ARTIFACT_BYTES = 64 * 1024 * 1024 +# Restricted workers cannot upload Modal's large-result blobs. Leave room for metadata. MAX_TRANSFER_BYTES = 1024 * 1024 -PROTOCOL_VERSION = 1 -def encode_artifacts(artifacts: dict[str, bytes]) -> bytes: - if not artifacts: +def pack_artifacts(root: Path) -> bytes: + files = [ + path for pattern in ("*/manifest.json", "*/extension.so") for path in root.glob(pattern) + ] + if not files: return b"" - if sum(map(len, artifacts.values())) > MAX_ARTIFACT_BYTES: + if sum(path.stat().st_size for path in files) > MAX_ARTIFACT_BYTES: raise ValueError("Compiled artifacts exceed the unpacked size limit") buffer = io.BytesIO() with zipfile.ZipFile(buffer, "w", compression=zipfile.ZIP_DEFLATED) as archive: - for name, data in artifacts.items(): - archive.writestr(name, data) - encoded = buffer.getvalue() - # Restricted Modal workers cannot create large-result blobs. Stay below the - # SDK's 2 MiB inline result limit, including metadata and serialization. - if len(encoded) > MAX_TRANSFER_BYTES: + for path in files: + archive.write(path, path.relative_to(root)) + if buffer.tell() > MAX_TRANSFER_BYTES: raise ValueError("Compiled artifacts exceed the inline transfer limit") - return encoded + return buffer.getvalue() + + +def write_file(root: Path, name: str, data: bytes): + path = root / name + if not path.resolve().is_relative_to(root.resolve()): + raise ValueError(f"Invalid artifact path: {name}") + path.parent.mkdir(parents=True, exist_ok=True) + path.write_bytes(data) -def decode_artifacts(encoded: bytes) -> dict[str, bytes]: +def unpack_artifacts(root: Path, encoded: bytes): if len(encoded) > MAX_TRANSFER_BYTES: raise ValueError("Artifact transfer exceeds the size limit") with zipfile.ZipFile(io.BytesIO(encoded)) as archive: if sum(item.file_size for item in archive.infolist()) > MAX_ARTIFACT_BYTES: raise ValueError("Unpacked artifacts exceed the size limit") - return {name: archive.read(name) for name in archive.namelist()} + for name in archive.namelist(): + write_file(root, name, archive.read(name)) def should_precompile(config: dict) -> bool: @@ -61,21 +65,22 @@ def should_precompile(config: dict) -> bool: ) -def source_hash(config: dict) -> str: - return hashlib.sha256(json.dumps(config["sources"], sort_keys=True).encode()).hexdigest() - - -def image_fingerprint() -> str: - return Path("/opt/kernelbot-pch/fingerprint").read_text().strip() +def build_identity(config: dict) -> dict: + return { + "version": 2, + "sources": hashlib.sha256( + json.dumps(config["sources"], sort_keys=True).encode() + ).hexdigest(), + "image": Path("/opt/kernelbot-pch/fingerprint").read_text().strip(), + } def cuda_arch(sm: str) -> str: # task configs use nvcc's SM spelling, e.g. 90a or 100. - suffix = "a" if sm.endswith("a") else "" digits = sm.removesuffix("a") if not digits.isdigit() or len(digits) < 2: raise ValueError(f"Invalid CUDA target: {sm}") - return f"{int(digits[:-1])}.{digits[-1]}{suffix}" + return f"{int(digits[:-1])}.{sm[len(digits) - 1 :]}" def _import_submission(work: Path, environment: dict, timeout: float): @@ -84,107 +89,83 @@ def _import_submission(work: Path, environment: dict, timeout: float): cwd=work, env=environment, stdout=subprocess.PIPE, - stderr=subprocess.PIPE, + stderr=subprocess.STDOUT, text=True, start_new_session=True, ) as process: try: - stdout, stderr = process.communicate(timeout=timeout) + output, _ = process.communicate(timeout=timeout) except subprocess.TimeoutExpired: - try: + with suppress(ProcessLookupError): os.killpg(process.pid, signal.SIGKILL) - except ProcessLookupError: - pass process.communicate() raise if process.returncode: - raise RuntimeError(f"CPU import exited {process.returncode}: {(stderr or stdout)[-1000:]}") + raise RuntimeError(f"CPU import exited {process.returncode}: {output[-1000:]}") def compile_python(config: dict) -> dict: started = time.perf_counter() info = CPUCompileInfo(status="skipped") - packet = {"version": PROTOCOL_VERSION, "artifacts": b"", "info": {}} + packet = {"artifacts": b""} try: - if not should_precompile(config): - return packet - import torch - - if torch.cuda.is_available() or torch.cuda.device_count(): - raise RuntimeError("CPU compilation worker has a visible GPU") - packet.update(source_hash=source_hash(config), image_fingerprint=image_fingerprint()) + packet["identity"] = build_identity(config) with tempfile.TemporaryDirectory(prefix="kernelbot-cpu-") as directory: work = Path(directory) - for name, text in config["sources"].items(): - path = work / name - if not path.resolve().is_relative_to(work): - raise ValueError(f"Invalid source path: {name}") - path.parent.mkdir(parents=True, exist_ok=True) - path.write_text(text) + for name, source in config["sources"].items(): + write_file(work, name, source.encode()) environment = prepare_environment(work, "capture", cuda_arch(config["arch"])) environment["POPCORN_SEED"] = "1" _import_submission(work, environment, Timeout.COMPILE) - artifacts = pack_artifacts(work / "artifacts") - packet["artifacts"] = encode_artifacts(artifacts) - info.artifacts = sum(name.endswith("/manifest.json") for name in artifacts) + packet["artifacts"] = pack_artifacts(work / "artifacts") + info.artifacts = len(list((work / "artifacts").glob("*/manifest.json"))) info.status = "compiled" if info.artifacts else "skipped" - if not info.artifacts: - info.reason = "No import-time load_inline library" except Exception as exc: info.status, info.reason = "fallback", str(exc)[-1000:] - packet["artifacts"] = b"" - finally: - info.duration = time.perf_counter() - started - packet["info"] = dataclasses.asdict(info) - return packet + info.duration = time.perf_counter() - started + return {**packet, "info": dataclasses.asdict(info)} -def run_with_cpu_artifacts(config: dict, run): - packet = config.get("cpu_compile") - if config.get("lang") != "py" or not packet: - return run(config) - try: - info = CPUCompileInfo(**packet["info"]) - except (KeyError, TypeError): - result = run(config) - result.cpu_compile = CPUCompileInfo("fallback", reason="Invalid CPU build metadata") - return result - if not packet.get("artifacts"): - result = run(config) - result.cpu_compile = info - return result - +@contextmanager +def cpu_artifacts(config: dict): + """Set up optional reuse; never catch or retry evaluation failures.""" with tempfile.TemporaryDirectory(prefix="kernelbot-gpu-") as directory: - work = Path(directory) + work, environment = Path(directory), {} + info = CPUCompileInfo("fallback") try: - if packet["version"] != PROTOCOL_VERSION: - raise ValueError("CPU artifact protocol changed") - if packet["source_hash"] != source_hash(config): - raise ValueError("Submission sources changed after CPU compilation") - if packet["image_fingerprint"] != image_fingerprint(): - raise ValueError("CPU and GPU runner images differ") - unpack_artifacts(work / "artifacts", decode_artifacts(packet["artifacts"])) - environment = prepare_environment(work, "auto", cuda_arch(config["arch"])) + packet = config["cpu_compile"] + info = CPUCompileInfo(**packet["info"]) + if packet["artifacts"]: + if packet["identity"] != build_identity(config): + raise ValueError("CPU artifact protocol, sources, or runner image changed") + unpack_artifacts(work / "artifacts", packet["artifacts"]) + environment = prepare_environment(work, "auto", cuda_arch(config["arch"])) except Exception as exc: info.status, info.reason = "fallback", str(exc)[-1000:] - result = run(config) - result.cpu_compile = info - return result previous_environment = dict(os.environ) try: - with chdir(work): + with chdir(work if environment else Path.cwd()): os.environ.update(environment) - result = run(config) + yield info finally: os.environ.clear() os.environ.update(previous_environment) - events = read_events(work / "artifacts") - info.reused = sum(event["mode"] == "replay" for event in events) - info.fallbacks = sum(event["mode"] == "fallback" for event in events) - info.status = "reused" if info.reused else "fallback" - if info.reused and info.fallbacks: - info.status = "partial" - reasons = [event["reason"] for event in events if event["mode"] == "fallback"] - info.reason = reasons[-1] if reasons else "" - result.cpu_compile = info - return result + events = work / "artifacts" / "events.jsonl" + if environment: + for line in events.read_text().splitlines() if events.exists() else []: + event = json.loads(line) + info.reused += event["mode"] == "replay" + info.fallbacks += event["mode"] == "fallback" + info.reason = event.get("reason", info.reason) + info.status = "reused" if info.reused else "fallback" + if info.reused and info.fallbacks: + info.status = "partial" + + +def run_with_cpu_artifacts(config: dict, run): + if config.get("lang") != "py" or not config.get("cpu_compile"): + return run(config) + with cpu_artifacts(config) as info: + result = run(config) + result.cpu_compile = info + return result diff --git a/tests/test_inline_artifacts.py b/tests/test_inline_artifacts.py index 6eb19e51..869c1640 100644 --- a/tests/test_inline_artifacts.py +++ b/tests/test_inline_artifacts.py @@ -1,158 +1,79 @@ -"""CPU-only contract tests; the Modal prototype exercises real compilation/loading.""" +"""CPU-only adapter tests; Modal integration tests exercise real library loading.""" import json -import os import sys -import tempfile -import types -import unittest -from pathlib import Path +from types import SimpleNamespace from unittest.mock import Mock, patch -from libkernelbot.inline_artifacts import install, pack_artifacts, unpack_artifacts +import pytest +from libkernelbot.inline_artifacts import install -class InlineArtifactTests(unittest.TestCase): - def setUp(self): - self.temporary = tempfile.TemporaryDirectory() - self.addCleanup(self.temporary.cleanup) - self.root = Path(self.temporary.name) - self.binary = self.root / "compiled.so" - self.binary.write_bytes(b"fake compiled extension") - self.builds = 0 - def load_inline( - name, - cpp_sources, - cuda_sources=None, - extra_cuda_cflags=None, - extra_ldflags=None, - is_python_module=True, - verbose=False, - build_directory=None, - ): - self.builds += 1 - return types.SimpleNamespace(__name__=name, __file__=str(self.binary)) +@pytest.fixture +def extension(tmp_path, monkeypatch): + binary = tmp_path / "compiled.so" + binary.write_bytes(b"fake compiled extension") + builds = Mock(return_value=SimpleNamespace(__name__="example", __file__=str(binary))) - self.original = load_inline - self.extension = types.SimpleNamespace(load_inline=load_inline) - torch = types.SimpleNamespace( - __version__="test", - version=types.SimpleNamespace(cuda="13.3"), - _C=types.SimpleNamespace(_GLIBCXX_USE_CXX11_ABI=True), - ) - modules = patch.dict( - sys.modules, - { - "torch": torch, - "torch.utils": types.SimpleNamespace(cpp_extension=self.extension), - }, - ) - modules.start() - self.addCleanup(modules.stop) - environment = patch.dict( - os.environ, - { - "KERNELBOT_INLINE_MODE": "capture", - "KERNELBOT_INLINE_ARTIFACTS": str(self.root / "artifacts"), - "KERNELBOT_INLINE_WORKDIR": str(self.root), - "TORCH_CUDA_ARCH_LIST": "7.5", - }, - ) - environment.start() - self.addCleanup(environment.stop) - dependencies = patch("libkernelbot.inline_artifacts._dependencies", return_value=[]) - dependencies.start() - self.addCleanup(dependencies.stop) - install() - self.extension.load_inline("example", "source", cuda_sources="cuda") + def original(name, cpp_sources, cuda_sources=None, extra_cuda_cflags=None, extra_ldflags=None): + return builds(name, cpp_sources, cuda_sources, extra_cuda_cflags, extra_ldflags) - def replay(self): - self.extension.load_inline = self.original - os.environ["KERNELBOT_INLINE_MODE"] = "replay" + extension = SimpleNamespace(load_inline=original, builds=builds) + torch = SimpleNamespace( + __version__="test", + version=SimpleNamespace(cuda="13.3"), + _C=SimpleNamespace(_GLIBCXX_USE_CXX11_ABI=True), + ) + monkeypatch.setitem(sys.modules, "torch", torch) + monkeypatch.setitem(sys.modules, "torch.utils", SimpleNamespace(cpp_extension=extension)) + monkeypatch.setenv("KERNELBOT_INLINE_WORKDIR", str(tmp_path)) + monkeypatch.setenv("TORCH_CUDA_ARCH_LIST", "7.5") + monkeypatch.setenv("KERNELBOT_INLINE_MODE", "capture") + with patch("libkernelbot.inline_artifacts._dependencies", return_value={}): install() - - def test_changed_source_and_flags_do_not_reuse_binary(self): - self.replay() - for args in ({"cpp_sources": "changed"}, {"extra_cuda_cflags": ["-O3"]}): - request = {"name": "example", "cpp_sources": "source", "cuda_sources": "cuda", **args} - with self.assertRaisesRegex(RuntimeError, "No CPU-built artifact"): - self.extension.load_inline(**request) - self.assertEqual(self.builds, 1) - - def test_changed_architecture_does_not_reuse_binary(self): - os.environ["TORCH_CUDA_ARCH_LIST"] = "9.0a" - self.replay() - with self.assertRaisesRegex(RuntimeError, "No CPU-built artifact"): - self.extension.load_inline("example", "source", cuda_sources="cuda") - - def test_replay_loads_saved_module_once_without_compiling(self): - self.replay() - module = types.SimpleNamespace() - spec = Mock() - with ( - patch("importlib.util.spec_from_file_location", return_value=spec) as make_spec, - patch("importlib.util.module_from_spec", return_value=module), - ): - for _ in range(2): - self.assertIs( - self.extension.load_inline("example", "source", cuda_sources="cuda"), module - ) - self.assertEqual(make_spec.call_args.args[0], "example") - self.assertEqual(make_spec.call_args.args[1].read_bytes(), self.binary.read_bytes()) - spec.loader.exec_module.assert_called_once_with(module) - self.assertEqual(self.builds, 1) - - def test_wrong_manifest_is_rejected_before_loading(self): - manifest_path = next((self.root / "artifacts").glob("*/manifest.json")) + extension.load_inline("example", "source", cuda_sources="cuda") + extension.load_inline = original + monkeypatch.setenv("KERNELBOT_INLINE_MODE", "auto") + install() + return extension + + +def test_loads_saved_module_once_without_compiling(extension, tmp_path): + module, spec = SimpleNamespace(), Mock() + with ( + patch("importlib.util.spec_from_file_location", return_value=spec) as make_spec, + patch("importlib.util.module_from_spec", return_value=module), + ): + for _ in range(2): + assert extension.load_inline("example", "source", cuda_sources="cuda") is module + assert make_spec.call_args.args[0] == "example" + assert make_spec.call_args.args[1].read_bytes() == (tmp_path / "compiled.so").read_bytes() + spec.loader.exec_module.assert_called_once_with(module) + assert extension.builds.call_count == 1 + + +@pytest.mark.parametrize("change", ["source", "flags", "arch", "manifest", "binary", "linker"]) +def test_incompatible_artifacts_fall_back_before_loading(extension, tmp_path, monkeypatch, change): + arguments = {"name": "example", "cpp_sources": "source", "cuda_sources": "cuda"} + manifest_path = next((tmp_path / "artifacts").glob("*/manifest.json")) + if change == "source": + arguments["cpp_sources"] = "changed" + elif change == "flags": + arguments["extra_cuda_cflags"] = ["-O3"] + elif change == "arch": + monkeypatch.setenv("TORCH_CUDA_ARCH_LIST", "9.0a") + elif change == "manifest": manifest = json.loads(manifest_path.read_text()) - manifest["runtime"]["cuda"] = "different" + manifest["key"] = "different" manifest_path.write_text(json.dumps(manifest)) - self.replay() - with self.assertRaisesRegex(RuntimeError, "manifest does not match"): - self.extension.load_inline("example", "source", cuda_sources="cuda") - - def test_corrupted_artifact_is_rejected_before_loading(self): - next((self.root / "artifacts").glob("*/extension.so")).write_bytes(b"corrupted") - self.replay() - with self.assertRaisesRegex(RuntimeError, "SHA-256"): - self.extension.load_inline("example", "source", cuda_sources="cuda") - - def test_ninja_is_disabled_during_replay(self): - self.replay() - with self.assertRaisesRegex(RuntimeError, "GPU compilation is disabled"): - self.extension._run_ninja_build() - - def test_native_mode_falls_back_to_original_compiler(self): - self.extension.load_inline = self.original - os.environ["KERNELBOT_INLINE_MODE"] = "auto" - install() - self.extension.load_inline("example", "changed source", cuda_sources="cuda") - self.assertEqual(self.builds, 2) - - def test_environment_changes_after_import_do_not_reuse_binary(self): - self.replay() - os.environ["TORCH_CUDA_ARCH_LIST"] = "9.0a" - with self.assertRaisesRegex(RuntimeError, "No CPU-built artifact"): - self.extension.load_inline("example", "source", cuda_sources="cuda") - - def test_custom_linker_inputs_are_left_to_the_gpu(self): - with self.assertRaisesRegex(ValueError, "custom linker inputs"): - self.extension.load_inline("linked", "source", extra_ldflags=["-lcustom"]) - self.assertEqual(self.builds, 1) - self.extension.load_inline = self.original - os.environ["KERNELBOT_INLINE_MODE"] = "auto" - install() - self.extension.load_inline("linked", "source", extra_ldflags=["-lcustom"]) - self.assertEqual(self.builds, 2) - - def test_transfer_contains_only_binary_and_manifest(self): - artifacts = pack_artifacts(self.root / "artifacts") - self.assertEqual(len(artifacts), 2) - destination = self.root / "received" - unpack_artifacts(destination, artifacts) - self.assertEqual(pack_artifacts(destination), artifacts) - - def test_transfer_rejects_path_escape(self): - with self.assertRaisesRegex(ValueError, "Invalid artifact path"): - unpack_artifacts(self.root / "received", {"../escape.so": b"binary"}) + elif change == "binary": + manifest_path.with_name("extension.so").write_bytes(b"corrupted") + else: + arguments["extra_ldflags"] = ["-lcustom"] + with patch("importlib.util.spec_from_file_location") as load: + extension.load_inline(**arguments) + load.assert_not_called() + assert extension.builds.call_count == 2 + event = json.loads((tmp_path / "artifacts/events.jsonl").read_text()) + assert event["mode"] == "fallback" and event["reason"] diff --git a/tests/test_python_precompile.py b/tests/test_python_precompile.py index 40098d39..1908f468 100644 --- a/tests/test_python_precompile.py +++ b/tests/test_python_precompile.py @@ -2,7 +2,6 @@ import json import os import subprocess -import sys from pathlib import Path from types import SimpleNamespace from unittest.mock import AsyncMock, MagicMock, patch @@ -15,11 +14,11 @@ from libkernelbot.python_precompile import ( compile_python, cuda_arch, - decode_artifacts, - encode_artifacts, + pack_artifacts, run_with_cpu_artifacts, should_precompile, - source_hash, + unpack_artifacts, + write_file, ) from libkernelbot.run_eval import FullResult, SystemInfo @@ -125,18 +124,8 @@ async def test_ordinary_python_or_operator_disable_bypasses_cpu( "error", [RuntimeError("CUDA unavailable"), subprocess.TimeoutExpired("python", 1)] ) def test_cpu_import_failures_return_fallback(config, monkeypatch, error): - monkeypatch.setitem( - sys.modules, - "torch", - SimpleNamespace( - cuda=SimpleNamespace( - is_available=lambda: False, - device_count=lambda: 0, - ) - ), - ) with ( - patch("libkernelbot.python_precompile.image_fingerprint", return_value="image"), + patch("libkernelbot.python_precompile.build_identity", return_value={"image": "image"}), patch("libkernelbot.python_precompile._import_submission", side_effect=error), ): packet = compile_python(config) @@ -144,57 +133,72 @@ def test_cpu_import_failures_return_fallback(config, monkeypatch, error): assert packet["artifacts"] == b"" -@pytest.mark.parametrize("changed", ["source_hash", "image_fingerprint", "version"]) +@pytest.mark.parametrize("changed", ["sources", "image", "version"]) def test_gpu_rejects_stale_build_without_failing_submission(config, result, changed): packet = { - "version": 1, - "source_hash": source_hash(config), - "image_fingerprint": "image", + "identity": {"image": "image"}, "artifacts": {"unused.so": b"binary"}, "info": {"status": "compiled"}, } - packet[changed] = "different" + packet["identity"][changed] = "different" run = MagicMock(return_value=result) - with patch("libkernelbot.python_precompile.image_fingerprint", return_value="image"): + with patch("libkernelbot.python_precompile.build_identity", return_value={"image": "image"}): actual = run_with_cpu_artifacts({**config, "cpu_compile": packet}, run) assert actual is result and actual.cpu_compile.status == "fallback" run.assert_called_once() @pytest.mark.parametrize("success", [True, False]) -def test_gpu_context_restores_environment_and_records_reuse(config, result, success): +def test_gpu_context_restores_environment_and_records_reuse(config, result, success, tmp_path): + write_file(tmp_path, "key/extension.so", b"binary") result.success = success config["cpu_compile"] = { - "version": 1, - "source_hash": source_hash(config), - "image_fingerprint": "image", - "artifacts": encode_artifacts({"extension.so": b"binary"}), + "identity": {"image": "image"}, + "artifacts": pack_artifacts(tmp_path), "info": {"status": "compiled"}, } previous_cwd, previous_environment = Path.cwd(), dict(os.environ) def run(_): assert os.environ["KERNELBOT_INLINE_MODE"] == "auto" - root = Path(os.environ["KERNELBOT_INLINE_ARTIFACTS"]) + root = Path(os.environ["KERNELBOT_INLINE_WORKDIR"]) / "artifacts" (root / "events.jsonl").write_text(json.dumps({"mode": "replay"}) + "\n") return result - with patch("libkernelbot.python_precompile.image_fingerprint", return_value="image"): + with patch("libkernelbot.python_precompile.build_identity", return_value={"image": "image"}): actual = run_with_cpu_artifacts(config, run) assert actual.cpu_compile.status == "reused" and actual.cpu_compile.reused == 1 assert actual.success == success assert Path.cwd() == previous_cwd and dict(os.environ) == previous_environment -def test_compressed_transport_roundtrip(): +def test_compressed_transport_roundtrip(tmp_path): + cpu, gpu = tmp_path / "cpu", tmp_path / "gpu" artifacts = {"key/extension.so": b"binary" * 1000, "key/manifest.json": b"{}"} - assert decode_artifacts(encode_artifacts(artifacts)) == artifacts + for name, data in {**artifacts, "build/main.o": b"exclude"}.items(): + write_file(cpu, name, data) + unpack_artifacts(gpu, pack_artifacts(cpu)) + assert { + str(p.relative_to(gpu)): p.read_bytes() for p in gpu.rglob("*") if p.is_file() + } == artifacts -def test_transport_size_limit_uses_fallback_instead_of_large_modal_blobs(monkeypatch): +def test_transport_size_limit_uses_fallback_instead_of_large_modal_blobs(tmp_path, monkeypatch): monkeypatch.setattr("libkernelbot.python_precompile.MAX_TRANSFER_BYTES", 1) + write_file(tmp_path, "key/extension.so", b"binary") with pytest.raises(ValueError, match="inline transfer limit"): - encode_artifacts({"key/extension.so": b"binary"}) + pack_artifacts(tmp_path) + + +def test_transfer_rejects_path_escape(tmp_path): + import io + import zipfile + + buffer = io.BytesIO() + with zipfile.ZipFile(buffer, "w") as archive: + archive.writestr("../escape.so", b"binary") + with pytest.raises(ValueError, match="Invalid artifact path"): + unpack_artifacts(tmp_path, buffer.getvalue()) def test_header_dependencies_relocate_and_detect_changes(tmp_path, monkeypatch): From 0a1e75c1c96f5f6ad4bf221b6397e1fa64039e6b Mon Sep 17 00:00:00 2001 From: Mark Saroufim Date: Wed, 9 Sep 2026 22:39:33 -0700 Subject: [PATCH 4/4] Discard CPU artifacts if worker cleanup fails --- src/libkernelbot/python_precompile.py | 1 + 1 file changed, 1 insertion(+) diff --git a/src/libkernelbot/python_precompile.py b/src/libkernelbot/python_precompile.py index 94e44429..6d2ea1cd 100644 --- a/src/libkernelbot/python_precompile.py +++ b/src/libkernelbot/python_precompile.py @@ -122,6 +122,7 @@ def compile_python(config: dict) -> dict: info.status = "compiled" if info.artifacts else "skipped" except Exception as exc: info.status, info.reason = "fallback", str(exc)[-1000:] + packet["artifacts"] = b"" info.duration = time.perf_counter() - started return {**packet, "info": dataclasses.asdict(info)}