Compile Python load_inline submissions on CPU before Modal GPU evaluation

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'
+```

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()))

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", ""),
+                ],
+            )
+        ),
+    }