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/src/libkernelbot/inline_artifacts.py b/src/libkernelbot/inline_artifacts.py new file mode 100644 index 00000000..e7f2ece5 --- /dev/null +++ b/src/libkernelbot/inline_artifacts.py @@ -0,0 +1,193 @@ +"""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 subprocess +import sysconfig +from pathlib import Path + + +def _digest(value: bytes) -> str: + return hashlib.sha256(value).hexdigest() + + +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) + (directory / "manifest.json").write_text( + json.dumps( + { + "module_name": module.__name__, + "sha256": _digest(binary), + "key": directory.name, + "dependencies": _dependencies(Path(module.__file__).parent, work), + }, + sort_keys=True, + ) + ) + + +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" + 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 _dependencies(build: Path, work: Path) -> dict[str, str]: + 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. + 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: dict[str, str] | None, work: Path) -> None: + if dependencies is None: + raise RuntimeError("Artifact has no compiler dependency information") + for name, digest in dependencies.items(): + path = work / name + if _digest(path.read_bytes()) != digest: + raise RuntimeError(f"Build dependency changed: {path}") + + +def install() -> None: + """Install in submission interpreters, including spawned evaluators.""" + import torch + from torch.utils import cpp_extension + + mode = os.environ["KERNELBOT_INLINE_MODE"] + 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) + 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, + } + 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 + record("fallback", reason=str(exc)[:500]) + 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) 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"} + } + 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", + ) + }, + } + # 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 + if mode == "capture": + module = original(*args, **kwargs) + _save_module(directory, module, work) + else: + if key not in loaded: + loaded[key] = _load_module(directory, work) + module = loaded[key] + record("replay") + return module + + cpp_extension.load_inline = load_inline + + +def prepare_environment(work: Path, mode: str, arch: str) -> dict[str, str]: + """Configure the hook in child interpreters without patching the runner.""" + 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_WORKDIR": str(work), + "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", ""), + ], + ) + ), + } diff --git a/src/libkernelbot/launchers/modal.py b/src/libkernelbot/launchers/modal.py index ce7ae06e..a4b073c0 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,29 @@ 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, + ) + except Exception as exc: + # Old deployments and unavailable CPU workers retain the existing GPU path. + packet = { + "artifacts": b"", + "info": { + "status": "fallback", + "duration": time.perf_counter() - started, + "reason": str(exc)[-1000:], + }, + } + 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 +70,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 +86,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..6d2ea1cd --- /dev/null +++ b/src/libkernelbot/python_precompile.py @@ -0,0 +1,172 @@ +"""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, contextmanager, suppress +from pathlib import Path + +from libkernelbot.consts import Timeout +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 + + +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(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 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 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 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") + for name in archive.namelist(): + write_file(root, name, archive.read(name)) + + +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 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. + digits = sm.removesuffix("a") + if not digits.isdigit() or len(digits) < 2: + raise ValueError(f"Invalid CUDA target: {sm}") + return f"{int(digits[:-1])}.{sm[len(digits) - 1 :]}" + + +def _import_submission(work: Path, environment: dict, timeout: float): + with subprocess.Popen( + ["python3", "submission.py"], + cwd=work, + env=environment, + stdout=subprocess.PIPE, + stderr=subprocess.STDOUT, + text=True, + start_new_session=True, + ) as process: + try: + output, _ = process.communicate(timeout=timeout) + except subprocess.TimeoutExpired: + with suppress(ProcessLookupError): + os.killpg(process.pid, signal.SIGKILL) + process.communicate() + raise + if process.returncode: + 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 = {"artifacts": b""} + try: + packet["identity"] = build_identity(config) + with tempfile.TemporaryDirectory(prefix="kernelbot-cpu-") as directory: + work = Path(directory) + 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) + packet["artifacts"] = pack_artifacts(work / "artifacts") + info.artifacts = len(list((work / "artifacts").glob("*/manifest.json"))) + 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)} + + +@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, environment = Path(directory), {} + info = CPUCompileInfo("fallback") + try: + 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:] + previous_environment = dict(os.environ) + try: + with chdir(work if environment else Path.cwd()): + os.environ.update(environment) + yield info + finally: + os.environ.clear() + os.environ.update(previous_environment) + 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/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 new file mode 100644 index 00000000..869c1640 --- /dev/null +++ b/tests/test_inline_artifacts.py @@ -0,0 +1,79 @@ +"""CPU-only adapter tests; Modal integration tests exercise real library loading.""" + +import json +import sys +from types import SimpleNamespace +from unittest.mock import Mock, patch + +import pytest + +from libkernelbot.inline_artifacts import install + + +@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))) + + 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) + + 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() + 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["key"] = "different" + manifest_path.write_text(json.dumps(manifest)) + 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_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..1908f468 --- /dev/null +++ b/tests/test_python_precompile.py @@ -0,0 +1,228 @@ +import asyncio +import json +import os +import subprocess +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, + pack_artifacts, + run_with_cpu_artifacts, + should_precompile, + unpack_artifacts, + write_file, +) +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): + with ( + patch("libkernelbot.python_precompile.build_identity", return_value={"image": "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", ["sources", "image", "version"]) +def test_gpu_rejects_stale_build_without_failing_submission(config, result, changed): + packet = { + "identity": {"image": "image"}, + "artifacts": {"unused.so": b"binary"}, + "info": {"status": "compiled"}, + } + packet["identity"][changed] = "different" + run = MagicMock(return_value=result) + 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, tmp_path): + write_file(tmp_path, "key/extension.so", b"binary") + result.success = success + config["cpu_compile"] = { + "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_WORKDIR"]) / "artifacts" + (root / "events.jsonl").write_text(json.dumps({"mode": "replay"}) + "\n") + return result + + 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(tmp_path): + cpu, gpu = tmp_path / "cpu", tmp_path / "gpu" + artifacts = {"key/extension.so": b"binary" * 1000, "key/manifest.json": b"{}"} + 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(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"): + 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): + 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)