diff --git a/CHANGELOG.md b/CHANGELOG.md index c8156262e9..737deec0b8 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -63,6 +63,19 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.1.0/), **Operator note**: sessions lost no recorded knowledge, but until now none of it was being shown to the agent. +### Changed + +- **Bare-metal `vllm` default bumped from `0.27.1` to `0.28.0` (still `rocm723`).** + `install_baremetal.sh`'s `VLLM_VERSION` default, `docs/compatibility.rst`, + `docs/install/install.md`, the example `SKILL.md` recipes, and + `assets/slurm/models.tsv` now all name `vllm==0.28.0+rocm723` / + `vllm/vllm-openai-rocm:v0.28.0`. Verified against the real upstream + `v0.28.0` tag that TraceLens' `config_vllm_v0.28.0.patch` applies cleanly + (`git apply --check`), so the TraceLens profiler-config patch path is + unaffected by the bump. `VLLM_ROCM_VARIANT` is unchanged: `wheels.vllm.ai` + only publishes a `rocm723` build for `0.28.0`, same as `0.27.1`. Overridable + via `VLLM_VERSION`/`VLLM_ROCM_VARIANT` as before. + ### Removed - **The `learning/` tuning database, the tracker's scoring layer, and the diff --git a/docs/compatibility.rst b/docs/compatibility.rst index 2940d6ba5e..51190d5bc5 100644 --- a/docs/compatibility.rst +++ b/docs/compatibility.rst @@ -47,7 +47,7 @@ The following table lists the minimum requirements for running Hyperloom. +=====================+========================================================+ | AMD Instinct™ GPU | MI300X, MI325X, MI355X | +---------------------+--------------------------------------------------------+ -| Operating System | Ubuntu 22.04, Ubuntu 24.04 | +| Operating System | Ubuntu 24.04 (recommended); Ubuntu 22.04 (legacy) | +---------------------+--------------------------------------------------------+ | ROCm Version | 7.2.x | +---------------------+--------------------------------------------------------+ @@ -146,7 +146,7 @@ uses a private registry mirror, set the registry prefix accordingly. - MI300X / MI325X * - ``lmsysorg/sglang-rocm:v0.5.18-rocm724-mi35x-20260825`` - MI355X - * - ``vllm/vllm-openai-rocm:v0.27.1`` + * - ``vllm/vllm-openai-rocm:v0.28.0`` - MI300X / MI325X / MI355X The vLLM image entrypoint is ``vllm serve``, so override it (for example @@ -170,6 +170,9 @@ Hyperloom does not install ROCm or torch itself. * - Item - Recommended - Notes + * - Operating System + - Ubuntu 24.04 + - Recommended bare-metal baseline. vLLM 0.28.0+ ROCm wheels require glibc >= 2.39, so Ubuntu 22.04 hosts must downgrade vLLM (for example ``VLLM_VERSION=0.27.1``) or use ``docker`` mode instead. * - ROCm - 7.2.x - The patch level differs per framework and is the same in both setup modes: the vLLM stack uses ROCm 7.2.3 and the SGLang stack uses ROCm 7.2.4 (see the note below). @@ -183,12 +186,12 @@ Hyperloom does not install ROCm or torch itself. - 0.5.18 (rocm724), pinned to commit ``0c7ff19e3b73`` - Installed in ``shared`` mode (reuses the host torch). Uses the ROCm 7.2.4 AMD wheel index (``SGLANG_ROCM_EXTRA=rocm724``), so the SGLang ROCm layer is 7.2.4. ``SGLANG_REF`` is the 0.5.18 pre-release commit the ``lmsysorg/sglang-rocm`` images are built from, not the ``v0.5.18`` tag: upstream removed ``detailed_annotations`` from ``io_struct.py`` between the two, and TraceLens' annotation patches need that field — on the tag three of the ten patches fail to apply, the atomic set rolls back, and kernel-shape profiling is silently unavailable. Note: ``SGLANG_REF`` only pins the version on the source-install branch (non-3.10 Python); on Python 3.10 the AMD wheel index installs ``amd-sglang`` unpinned, which might resolve to a different patch release — and therefore to a build these patches do not fit. * - vLLM - - v0.27.1 (rocm723), isolated venv - - Installs ``vllm==0.27.1+rocm723`` from the wheels.vllm.ai pip index. vLLM's ROCm wheel pins its own torch, so it installs into a dedicated venv (``--framework-env isolated``, the default for vLLM) and never touches the host torch. + - v0.28.0 (rocm723), isolated venv + - Installs ``vllm==0.28.0+rocm723`` from the wheels.vllm.ai pip index on Ubuntu 24.04+. vLLM's ROCm wheel pins its own torch, so it installs into a dedicated venv (``--framework-env isolated``, the default for vLLM) and never touches the host torch. Bare-metal ROCm patch levels differ per framework, and each one matches its container image. The vLLM stack installs the ``rocm723`` variant (ROCm -7.2.3), matching ``vllm/vllm-openai-rocm:v0.27.1``; the SGLang stack +7.2.3), matching ``vllm/vllm-openai-rocm:v0.28.0``; the SGLang stack installs from the ROCm 7.2.4 AMD wheel index, matching the two ``lmsysorg/sglang-rocm:v0.5.18-rocm724-*`` images. ``docker`` mode is still the preferred route for a pre-validated stack, since the images also pin the diff --git a/docs/install/install.md b/docs/install/install.md index 81a2bd8526..1639d1cac0 100644 --- a/docs/install/install.md +++ b/docs/install/install.md @@ -103,6 +103,9 @@ directly. Requirements: +- Ubuntu 24.04 is the recommended host OS for bare-metal setup. vLLM 0.28.0+ + requires Ubuntu 24.04 or newer; on Ubuntu 22.04, downgrade vLLM (for example + ``VLLM_VERSION=0.27.1``) or use Docker mode instead. - ROCm runtime and ROCm torch are already installed. - `git` is available for dependency checkouts. - A serving framework is either already installed, or setup might install one. @@ -393,7 +396,7 @@ It is recommended that you use a ROCm image that already ships the serving framework, so nothing needs to be installed inside the container beyond Hyperloom's runtime deps. The following images are recommended: -- `vllm`: `docker.io/vllm/vllm-openai-rocm:v0.27.1` +- `vllm`: `docker.io/vllm/vllm-openai-rocm:v0.28.0` - `sglang` MI300X: `docker.io/lmsysorg/sglang-rocm:v0.5.18-rocm724-mi30x-20260825` - `sglang` MI355X: `docker.io/lmsysorg/sglang-rocm:v0.5.18-rocm724-mi35x-20260825` @@ -401,7 +404,7 @@ Start a long-running container from the repo root, mounting it at the same path so `.env`, logs, and session artifacts stay valid: ```bash -export HYPERLOOM_IMAGE=docker.io/vllm/vllm-openai-rocm:v0.27.1 +export HYPERLOOM_IMAGE=docker.io/vllm/vllm-openai-rocm:v0.28.0 export REPO_ROOT="$(pwd -P)" docker run -d \ --name "${HYPERLOOM_CONTAINER_NAME:-hyperloom-local}" \ diff --git a/examples/hyperloom-custom-advanced/SKILL.md b/examples/hyperloom-custom-advanced/SKILL.md index 06984c2f4e..4d12ce3449 100644 --- a/examples/hyperloom-custom-advanced/SKILL.md +++ b/examples/hyperloom-custom-advanced/SKILL.md @@ -44,7 +44,7 @@ skip the user-approval step (#1314). Suggested Docker images: -- `vllm`: `docker.io/vllm/vllm-openai-rocm:v0.27.1` +- `vllm`: `docker.io/vllm/vllm-openai-rocm:v0.28.0` - `sglang` MI300X: `docker.io/lmsysorg/sglang-rocm:v0.5.18-rocm724-mi30x-20260825` - `sglang` MI355X: `docker.io/lmsysorg/sglang-rocm:v0.5.18-rocm724-mi35x-20260825` diff --git a/examples/hyperloom-qwen3-14b-fp8-12h-forge/SKILL.md b/examples/hyperloom-qwen3-14b-fp8-12h-forge/SKILL.md index 0a168a543d..1a8def0192 100644 --- a/examples/hyperloom-qwen3-14b-fp8-12h-forge/SKILL.md +++ b/examples/hyperloom-qwen3-14b-fp8-12h-forge/SKILL.md @@ -78,7 +78,7 @@ skip the user-approval step (#1314). Suggested Docker images: -- `vllm`: `docker.io/vllm/vllm-openai-rocm:v0.27.1` +- `vllm`: `docker.io/vllm/vllm-openai-rocm:v0.28.0` - `sglang` MI300X: `docker.io/lmsysorg/sglang-rocm:v0.5.18-rocm724-mi30x-20260825` - `sglang` MI355X: `docker.io/lmsysorg/sglang-rocm:v0.5.18-rocm724-mi35x-20260825` diff --git a/examples/hyperloom-qwen3-14b-fp8-12h/SKILL.md b/examples/hyperloom-qwen3-14b-fp8-12h/SKILL.md index 5f1a924c2a..1cefda57ad 100644 --- a/examples/hyperloom-qwen3-14b-fp8-12h/SKILL.md +++ b/examples/hyperloom-qwen3-14b-fp8-12h/SKILL.md @@ -34,7 +34,7 @@ skip the user-approval step (#1314). Suggested Docker images: -- `vllm`: `docker.io/vllm/vllm-openai-rocm:v0.27.1` +- `vllm`: `docker.io/vllm/vllm-openai-rocm:v0.28.0` - `sglang` MI300X: `docker.io/lmsysorg/sglang-rocm:v0.5.18-rocm724-mi30x-20260825` - `sglang` MI355X: `docker.io/lmsysorg/sglang-rocm:v0.5.18-rocm724-mi35x-20260825` diff --git a/examples/hyperloom-qwen3-8b-3h/SKILL.md b/examples/hyperloom-qwen3-8b-3h/SKILL.md index 93b62f2e24..340b1bceac 100644 --- a/examples/hyperloom-qwen3-8b-3h/SKILL.md +++ b/examples/hyperloom-qwen3-8b-3h/SKILL.md @@ -34,7 +34,7 @@ skip the user-approval step (#1314). Suggested Docker images: -- `vllm`: `docker.io/vllm/vllm-openai-rocm:v0.27.1` +- `vllm`: `docker.io/vllm/vllm-openai-rocm:v0.28.0` - `sglang` MI300X: `docker.io/lmsysorg/sglang-rocm:v0.5.18-rocm724-mi30x-20260825` - `sglang` MI355X: `docker.io/lmsysorg/sglang-rocm:v0.5.18-rocm724-mi35x-20260825` diff --git a/pyproject.toml b/pyproject.toml index ad6d66e843..caa0d21cd2 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -450,6 +450,8 @@ fail_under = 90 # Base trailing argv for ``pytest`` in CI; the workflow appends sharding # (``--splits``/``--group``) and ``-n`` from the values below. pytest_ci_args = [ + "-vv", + "-o", "faulthandler_timeout=180", "-m", "not critic_agent_e2e and not robustness_agent_e2e and not targeted_build_e2e", "--cov=.", "--cov-report=", ] diff --git a/src/hyperloom/inference_optimizer/assets/install_baremetal.sh b/src/hyperloom/inference_optimizer/assets/install_baremetal.sh index e4908d3f51..ea13a0d81c 100755 --- a/src/hyperloom/inference_optimizer/assets/install_baremetal.sh +++ b/src/hyperloom/inference_optimizer/assets/install_baremetal.sh @@ -54,8 +54,8 @@ SGLANG_REPO="${SGLANG_REPO:-https://github.com/sgl-project/sglang.git}" # io_struct.py between the two, and TraceLens' annotation patches need that field. # On the tag, three of the ten patches fail `git apply --check`, the atomic set # rolls back, and kernel-shape profiling is silently unavailable. -# vLLM installs 0.27.1+rocm723 from the wheels.vllm.ai pip index, matching the -# vllm/vllm-openai-rocm:v0.27.1 Docker image. The rocm723 variant puts the +# vLLM installs 0.28.0+rocm723 from the wheels.vllm.ai pip index, matching the +# vllm/vllm-openai-rocm:v0.28.0 Docker image. The rocm723 variant puts the # vLLM ROCm layer at 7.2.3, one patch level above the SGLang stack. AITER_REF # can pin ROCm/aiter to a released tag; when unset, the installer selects the # newest tag compatible with the already-installed ROCm torch/triton stack. @@ -77,7 +77,7 @@ fi SGLANG_ROCM_PYPI_VERSION="${SGLANG_ROCM_PYPI_VERSION:-7.2.4}" AITER_REPO="${AITER_REPO:-https://github.com/ROCm/aiter.git}" AITER_REF="${AITER_REF:-}" -VLLM_VERSION="${VLLM_VERSION:-0.27.1}" +VLLM_VERSION="${VLLM_VERSION:-0.28.0}" VLLM_ROCM_VARIANT="${VLLM_ROCM_VARIANT:-rocm723}" VLLM_ROCM_INDEX="${VLLM_ROCM_INDEX:-https://wheels.vllm.ai/rocm/${VLLM_VERSION}/${VLLM_ROCM_VARIANT}}" _VLLM_VENV_ROOT_WAS_SET="${VLLM_VENV_ROOT+x}" @@ -185,10 +185,9 @@ log() { echo "[install-baremetal] $*"; } warn() { echo "[install-baremetal WARN] $*" >&2; } die() { echo "[install-baremetal ERROR] $*" >&2; exit 1; } -IMAGE_HINT="Provision the ROCm framework base first (run inside an AMD ROCm \ -SGLang/vLLM image such as lmsysorg/sglang-rocm:v0.5.18-rocm724-mi30x|mi35x-* or \ -vllm/vllm-openai-rocm:v0.27.1, or install an equivalent ROCm torch + \ -framework stack), then re-run." +IMAGE_HINT="Provision the ROCm framework base first (SGLang: lmsysorg/sglang-rocm:v0.5.18-rocm724-mi30x|mi35x-*; \ +vLLM bare-metal: Ubuntu 24.04+ host with ROCm torch, or use docker mode with \ +vllm/vllm-openai-rocm:v0.28.0), then re-run." is_interactive() { [ "$ASSUME_YES" -eq 0 ] && [ -t 0 ] && [ -t 1 ]; } @@ -803,6 +802,35 @@ link_vllm_into_shared_bin() { fi } +host_glibc_version() { + getconf GNU_LIBC_VERSION 2>/dev/null | awk '{print $2}' +} + +_version_ge() { + [ "$(printf '%s\n' "$2" "$1" | sort -V | tail -n1)" = "$1" ] +} + +_vllm_semver_base() { + local ver="${VLLM_VERSION%%+*}" + ver="${ver%%-*}" + printf '%s' "$ver" +} + +vllm_version_requires_glibc_239() { + _version_ge "$(_vllm_semver_base)" "0.28.0" +} + +assert_vllm_glibc_compatible() { + vllm_version_requires_glibc_239 || return 0 + local glibc="${1:-$(host_glibc_version)}" + if [ -z "$glibc" ]; then + die "cannot detect host glibc; vLLM ${VLLM_VERSION} requires glibc >= 2.39. Use docker mode or set VLLM_VERSION=0.27.1 on older hosts." + fi + if ! _version_ge "$glibc" "2.39"; then + die "vLLM ${VLLM_VERSION} requires glibc >= 2.39 (host has ${glibc}). Use docker mode or set VLLM_VERSION=0.27.1 before running setup." + fi +} + # Install vLLM from the official ROCm wheel index without replacing ROCm torch. install_vllm_framework() { local py base_py py_mm constraint_file package_spec rocm_torch_ver @@ -837,6 +865,8 @@ PY log "VLLM_ROCM_VARIANT=${VLLM_ROCM_VARIANT}" log "VLLM_ROCM_INDEX=${VLLM_ROCM_INDEX}" + assert_vllm_glibc_compatible + if [ "$CHECK_ONLY" -eq 1 ]; then if [ "$FRAMEWORK_ENV" = "isolated" ] && [ ! -x "$py" ]; then warn "vllm isolated env missing (check-only; would create ${VLLM_VENV_ROOT} and install ${package_spec})" diff --git a/src/hyperloom/inference_optimizer/assets/slurm/models.tsv b/src/hyperloom/inference_optimizer/assets/slurm/models.tsv index 9ad0949219..8127501610 100644 --- a/src/hyperloom/inference_optimizer/assets/slurm/models.tsv +++ b/src/hyperloom/inference_optimizer/assets/slurm/models.tsv @@ -2,5 +2,5 @@ # Example registry consumed by submit.sh. Replace repo_id / image / gpu_type # with your own published models and container images before running. deepseek_r1_sglang deepseek-ai/DeepSeek-R1 sglang MI300X lmsysorg/sglang-rocm:v0.5.18-rocm724-mi30x-20260825 8 fp8 1024 1024 64 moe_mla 6 10 -deepseek_r1_vllm deepseek-ai/DeepSeek-R1 vllm MI300X vllm/vllm-openai-rocm:v0.27.1 8 fp8 1024 1024 64 moe_mla 6 10 -gptoss_vllm openai/gpt-oss-120b vllm MI300X vllm/vllm-openai-rocm:v0.27.1 8 fp4 1024 1024 64 moe_swa 6 10 +deepseek_r1_vllm deepseek-ai/DeepSeek-R1 vllm MI300X vllm/vllm-openai-rocm:v0.28.0 8 fp8 1024 1024 64 moe_mla 6 10 +gptoss_vllm openai/gpt-oss-120b vllm MI300X vllm/vllm-openai-rocm:v0.28.0 8 fp4 1024 1024 64 moe_swa 6 10 diff --git a/src/hyperloom/inference_optimizer/tests/test_baremetal_doc_version_consistency.py b/src/hyperloom/inference_optimizer/tests/test_baremetal_doc_version_consistency.py index a3920936d7..e9bcc7ad5b 100644 --- a/src/hyperloom/inference_optimizer/tests/test_baremetal_doc_version_consistency.py +++ b/src/hyperloom/inference_optimizer/tests/test_baremetal_doc_version_consistency.py @@ -22,12 +22,12 @@ def test_baremetal_defaults_match_compat_doc(): sh = INSTALLER.read_text(encoding="utf-8") doc = COMPAT.read_text(encoding="utf-8") - vllm_version = _default("VLLM_VERSION", sh) # e.g. 0.27.1 + vllm_version = _default("VLLM_VERSION", sh) # e.g. 0.28.0 vllm_variant = _default("VLLM_ROCM_VARIANT", sh) # e.g. rocm723 sglang_ref = _default("SGLANG_REF", sh) # e.g. v0.5.17 sglang_rocm_extra = _default("SGLANG_ROCM_EXTRA", sh) # e.g. rocm724 - # compatibility.rst documents e.g. "v0.27.1 (rocm723)" and the pip spec "vllm==0.27.1+rocm723"; keep both in + # compatibility.rst documents e.g. "v0.28.0 (rocm723)" and the pip spec "vllm==0.28.0+rocm723"; keep both in # lockstep with the script defaults. assert "v%s (%s)" % (vllm_version, vllm_variant) in doc, ( "docs/compatibility.rst must document vLLM 'v%s (%s)' to match " diff --git a/src/hyperloom/inference_optimizer/tests/test_setup_cli.py b/src/hyperloom/inference_optimizer/tests/test_setup_cli.py index 02aaf26fa9..3c8fa19bca 100644 --- a/src/hyperloom/inference_optimizer/tests/test_setup_cli.py +++ b/src/hyperloom/inference_optimizer/tests/test_setup_cli.py @@ -1256,11 +1256,13 @@ def test_baremetal_next_steps_names_the_detected_framework(tmp_path: Path): def test_baremetal_profiler_hotfix_accepts_an_atom_only_host(tmp_path: Path): """The hotfix patches ROCm profiler libs, which torch.profiler uses on any engine.""" + dotenv = tmp_path / ".env" + dotenv.write_text("HYPERLOOM_RUN_MODE=baremetal\n", encoding="utf-8") res = _drive_installer( tmp_path, importable={"atom"}, - dotenv=tmp_path / ".env", - body="rocm_profiler_hotfix_compatible && echo HOTFIX_ELIGIBLE", + dotenv=dotenv, + body=("running_in_container() { return 1; }\nrocm_profiler_hotfix_compatible && echo HOTFIX_ELIGIBLE"), ) assert "HOTFIX_ELIGIBLE" in res.stdout, res.stderr @@ -2006,6 +2008,47 @@ def test_docker_run_mode_applies_the_hotfix_for_an_sglang_image(tmp_path: Path): assert "HOTFIX_ELIGIBLE" in res.stdout, res.stderr +def _drive_vllm_glibc_gate( + tmp_path: Path, + *, + glibc: str, + vllm_version: str = "0.28.0", + check_only: bool = True, +) -> subprocess.CompletedProcess: + return _drive_installer( + tmp_path, + importable=set(), + dotenv=tmp_path / ".env", + body="\n".join( + [ + f'host_glibc_version() {{ echo "{glibc}"; }}', + f'VLLM_VERSION="{vllm_version}"', + f"CHECK_ONLY={1 if check_only else 0}", + "install_vllm_framework", + ] + ), + ) + + +def test_vllm_install_rejects_glibc_235_for_028(tmp_path: Path): + res = _drive_vllm_glibc_gate(tmp_path, glibc="2.35", vllm_version="0.28.0") + + assert res.returncode != 0 + assert "glibc >= 2.39" in res.stderr + + +def test_vllm_install_accepts_glibc_239_for_028(tmp_path: Path): + res = _drive_vllm_glibc_gate(tmp_path, glibc="2.39", vllm_version="0.28.0") + + assert res.returncode == 0, res.stderr + + +def test_vllm_install_allows_pre_028_override_on_glibc_235(tmp_path: Path): + res = _drive_vllm_glibc_gate(tmp_path, glibc="2.35", vllm_version="0.27.1") + + assert res.returncode == 0, res.stderr + + def test_baremetal_run_mode_keeps_the_hotfix_for_vllm(tmp_path: Path): """The framework split is docker-only: a bare-metal vLLM host profiles through the same ROCm libs the overlay fixes, so gating it there would lose profiling.""" res = _drive_hotfix_gate(tmp_path, run_mode="baremetal", importable={"vllm"}) diff --git a/src/hyperloom/skills/hyperloom-setup/SKILL.md b/src/hyperloom/skills/hyperloom-setup/SKILL.md index f12c364bae..26f0230371 100644 --- a/src/hyperloom/skills/hyperloom-setup/SKILL.md +++ b/src/hyperloom/skills/hyperloom-setup/SKILL.md @@ -170,6 +170,12 @@ value. - Do not mark any option as recommended. Present the three options in the exact order above without a default selection. +8. Only when the user chose `baremetal` **and** `vllm (isolated)` in Step 7, + briefly note that the installer enforces the vLLM 0.28.0 glibc floor + (glibc >= 2.39). If setup later fails with that error, explain it in plain + language and point the user to Docker mode or a pre-0.28 override — do not + implement a second version gate here. + ## Step 3: Write `.env` Create or update `.env` in the current directory. @@ -265,6 +271,15 @@ export REPO_ROOT="$(pwd -P)" PYTHONPATH="$REPO_ROOT" python3 -m hyperloom.inference_optimizer.setup -- --install-framework vllm --framework-env isolated --yes ``` +Downgrade path only when the user explicitly chooses a pre-0.28 vLLM on a host +that failed the glibc check: + +```bash +export REPO_ROOT="$(pwd -P)" +export VLLM_VERSION=0.27.1 +PYTHONPATH="$REPO_ROOT" python3 -m hyperloom.inference_optimizer.setup -- --install-framework vllm --framework-env isolated --yes +``` + For `sglang`: ```bash diff --git a/src/kernelforge/loop/task_preparer.py b/src/kernelforge/loop/task_preparer.py index cd171fd1f9..d8d59be736 100644 --- a/src/kernelforge/loop/task_preparer.py +++ b/src/kernelforge/loop/task_preparer.py @@ -127,7 +127,7 @@ def all_failures_are_timeouts(self) -> bool: # timing independently of whatever the driver prints: an eager driver replays zero times, a graph-timed one replays # once per timed iteration. _GRAPH_PROBE_SITECUSTOMIZE = r''' -import atexit, json, os +import atexit, json, os, sys _n = [0] # Collected only when the caller declares a rank count. A single-rank probe @@ -159,17 +159,13 @@ def _ancestor_pids(): _import_pid = os.getpid() -def _install(): +def _install(torch): """Patch CUDAGraph.replay lazily: torch may not be imported yet. - Everything is inside the guard, not only the import. A module named torch - that carries no ``cuda`` is a real shape -- a test stub, a partially - initialized package mid-import -- and reading the attribute outside would - raise out of an import hook that runs in every process on the path. + A module named torch may be a stub or partially initialized. Use the + existing module rather than importing through our own lazy hook again. """ try: - import torch - orig = torch.cuda.CUDAGraph.replay def _replay(self, *a, **k): @@ -182,7 +178,7 @@ def _replay(self, *a, **k): return False -_graph_ready = _install() +_graph_ready = _install(sys.modules.get("torch")) if not _graph_ready: # torch is imported by the driver, not by us. Hook the import so the patch @@ -194,10 +190,10 @@ def _replay(self, *a, **k): def _hooked(name, *a, **k): global _graph_ready mod = _real_import(name, *a, **k) - if name == "torch" or name.startswith("torch."): - _graph_ready = _install() - if _graph_ready: - builtins.__import__ = _real_import + if not _graph_ready and (name == "torch" or name.startswith("torch.")): + _graph_ready = _install(sys.modules.get("torch")) + if _graph_ready and builtins.__import__ is _hooked: + builtins.__import__ = _real_import return mod builtins.__import__ = _hooked diff --git a/src/kernelforge/tests/test_distributed_measurement_contract.py b/src/kernelforge/tests/test_distributed_measurement_contract.py index 98807b0be3..691b58453d 100644 --- a/src/kernelforge/tests/test_distributed_measurement_contract.py +++ b/src/kernelforge/tests/test_distributed_measurement_contract.py @@ -13,7 +13,15 @@ from __future__ import annotations +import atexit +import builtins import json +import subprocess +import sys +import textwrap +import types + +import pytest from kernelforge.loop import task_preparer from kernelforge.loop.task_preparer import ( @@ -178,7 +186,17 @@ async def _fake_create(*args, **kwargs): assert "FORGE_NPROC_PER_NODE" not in captured -def test_the_probe_survives_a_torch_that_carries_no_cuda(tmp_path, monkeypatch): +@pytest.fixture +def probe_runtime(monkeypatch): + original_import = builtins.__import__ + monkeypatch.setattr(builtins, "__import__", original_import) + monkeypatch.setattr(atexit, "register", lambda callback: callback) + stub = types.ModuleType("torch") + monkeypatch.setitem(sys.modules, "torch", stub) + return stub + + +def test_the_probe_survives_a_torch_that_carries_no_cuda(tmp_path, monkeypatch, probe_runtime): """The probe runs in every process on the path, so it may not raise from one. A module named torch with no ``cuda`` is a real shape -- a test stub, or a @@ -186,10 +204,6 @@ def test_the_probe_survives_a_torch_that_carries_no_cuda(tmp_path, monkeypatch): turned the lazy install into an uncaught AttributeError in whatever process happened to import torch next. """ - import sys - import types - - monkeypatch.setitem(sys.modules, "torch", types.ModuleType("torch")) monkeypatch.setenv("GRAPH_PROBE_OUT", str(tmp_path / "probe")) namespace: dict = {} @@ -200,7 +214,7 @@ def test_the_probe_survives_a_torch_that_carries_no_cuda(tmp_path, monkeypatch): assert namespace["_hooked"]("torch") is not None -def test_the_probe_reports_where_a_rank_measured(tmp_path, monkeypatch): +def test_the_probe_reports_where_a_rank_measured(tmp_path, monkeypatch, probe_runtime): """The shard's ``harness`` field is what the module under test actually writes. Asserting the field by hand everywhere else would let the probe and the @@ -219,3 +233,71 @@ class _Harness: monkeypatch.setitem(__import__("sys").modules, "dist_harness", _Harness) assert namespace["_harness_measured"]() is True + + +def test_lazy_probe_imports_are_bounded_and_restore_the_hook(tmp_path): + script = textwrap.dedent( + """ + import atexit + import builtins + import sys + import types + + source = sys.argv[1] + original_import = builtins.__import__ + atexit.register = lambda callback: callback + stub = types.ModuleType("torch") + sys.modules["torch"] = stub + sys.setrecursionlimit(80) + namespaces = [{}, {}] + for namespace in namespaces: + exec(compile(source, "sitecustomize.py", "exec"), namespace) + calls = [0] + def profile(frame, event, arg): + if event == "call" and frame.f_code.co_name == "_install": + calls[0] += 1 + sys.setprofile(profile) + for _ in range(3): + assert builtins.__import__("torch") is stub + sys.setprofile(None) + assert calls[0] == 6, calls + builtins.__import__ = original_import + + namespace = {} + exec(compile(source, "sitecustomize.py", "exec"), namespace) + class Graph: + def replay(self): + return "replayed" + stub.cuda = types.SimpleNamespace(CUDAGraph=Graph) + assert builtins.__import__("torch") is stub + assert namespace["_graph_ready"] is True + assert builtins.__import__ is original_import + replay = Graph.replay + for _ in range(3): + namespace["_hooked"]("torch") + assert Graph.replay is replay + assert Graph().replay() == "replayed" + assert namespace["_n"] == [1] + """ + ) + result = subprocess.run( + [sys.executable, "-c", script, task_preparer._GRAPH_PROBE_SITECUSTOMIZE], + cwd=tmp_path, + capture_output=True, + text=True, + timeout=10, + check=False, + ) + assert result.returncode == 0, result.stdout + result.stderr + + +def test_probe_tests_restore_the_import_hook(tmp_path): + original_import = builtins.__import__ + for test in ( + test_the_probe_survives_a_torch_that_carries_no_cuda, + test_the_probe_reports_where_a_rank_measured, + ): + with pytest.MonkeyPatch.context() as patch: + stub = probe_runtime.__wrapped__(patch) + test(tmp_path, patch, stub) + assert builtins.__import__ is original_import