diff --git a/.gitignore b/.gitignore index d958933a..31d1f0a4 100644 --- a/.gitignore +++ b/.gitignore @@ -362,4 +362,7 @@ marimo/_lsp/ __marimo__/ # Streamlit -.streamlit/secrets.toml \ No newline at end of file +.streamlit/secrets.toml +tests/deepep_matrix_work/ +deepep_fabric_archive/ +hook_archive/ diff --git a/CMakeLists.txt b/CMakeLists.txt index 923932b2..c40d44b3 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -43,3 +43,43 @@ set_target_properties(cuda_hook PROPERTIES install(TARGETS cuda_hook LIBRARY DESTINATION ${CMAKE_INSTALL_PREFIX} ) + +include(CTest) +if(BUILD_TESTING) + # Pure host-policy tests: no CUDA driver/runtime linkage, matching cuda_hook. + # EXCLUDE_FROM_ALL keeps them out of the default (editable pip) build; they + # are built explicitly (`cmake --build build --target `) and run via + # ctest. + add_executable(vmm_ipc_reservation_test EXCLUDE_FROM_ALL + tests/vmm_ipc_reservation_test.cpp + ) + target_include_directories(vmm_ipc_reservation_test PRIVATE + ${CMAKE_CURRENT_SOURCE_DIR}/include + ${CUDAToolkit_INCLUDE_DIRS} + ) + add_test(NAME vmm_ipc_reservation COMMAND vmm_ipc_reservation_test) + + add_executable(vmm_ipc_ownership_test EXCLUDE_FROM_ALL + tests/vmm_ipc_ownership_test.cpp + ) + target_include_directories(vmm_ipc_ownership_test PRIVATE + ${CMAKE_CURRENT_SOURCE_DIR}/include + ) + add_test(NAME vmm_ipc_ownership COMMAND vmm_ipc_ownership_test) + + add_executable(vmm_ipc_chunk_key_test EXCLUDE_FROM_ALL + tests/vmm_ipc_chunk_key_test.cpp + ) + target_include_directories(vmm_ipc_chunk_key_test PRIVATE + ${CMAKE_CURRENT_SOURCE_DIR}/include + ) + add_test(NAME vmm_ipc_chunk_key COMMAND vmm_ipc_chunk_key_test) + + add_executable(vmm_cleanup_test EXCLUDE_FROM_ALL + tests/vmm_cleanup_test.cpp + ) + target_include_directories(vmm_cleanup_test PRIVATE + ${CMAKE_CURRENT_SOURCE_DIR}/include + ) + add_test(NAME vmm_cleanup COMMAND vmm_cleanup_test) +endif() diff --git a/README.md b/README.md index bec64ee4..1c56c4d1 100644 --- a/README.md +++ b/README.md @@ -81,12 +81,110 @@ Foundry ships engine integrations under `foundry/python/foundry/integration/`. P | Engine | Single GPU | DP | TP | EP | |---|:---:|:---:|:---:|:---:| -| vLLM | ✅ | ✅ | 🚧 | ✅ | -| SGLang | ✅ | ✅ | 🚧 | ✅ | +| vLLM | ✅ | ✅ | ✅ | ✅ | +| SGLang | ✅ | ✅ | ✅ | 🚧 | | TensorRT-LLM | 🚧 | 🚧 | 🚧 | 🚧 | ✅ validated end-to-end (SAVE → LOAD → query)  ·  🚧 not yet +**Provenance.** Actual `foundry-org/foundry` upstream `main` is the +[0.0.2 commit `eef12012`](https://github.com/foundry-org/foundry/commit/eef12012aa0f85ae6079891144797b08c282152d), +which has no official TP support. This `modal-projects` fork is based on that +commit plus its +[agent-guidance commit `21824649`](https://github.com/modal-projects/foundry/commit/21824649e711487a2c6af01cb631b3e41ccc32d8). +Historical TP evidence came from fork draft +[PR #4](https://github.com/modal-projects/foundry/pull/4) and +[PR #2](https://github.com/modal-projects/foundry/pull/2), not upstream `main`; +the pinned SGLang evidence below was produced on this branch. + +TP remains experimental and each check mark is limited to the pinned profile +described below. + +vLLM dense TP=2 is verified only for Foundry +`a8e43442421a80fea7991ee26f97f1003b170e29`, vLLM +`4309c257d3f639e5490d3811293c890c61c76f29` with wheel parent +`6cbe448eed751824d608faf9078ef84724d621c1`, and +`Qwen/Qwen3-8B@b968826d9c46dd6066d109eabc6255188de91218`. The +[six-phase 2×H100 run](https://modal.com/apps/modal-labs/rahul-dev/ap-YfKALdMwnyJQmzQoVvZSMU) +passed all 17 checks. Each rank published 64 JSON/64 binary graph pairs; both +fresh LOADs restored `64/64` graphs per rank at live offset `70724354048`. +SAVE/SAVE2 inventories, offsets, and 68 semantic entries per rank matched; all +265 fingerprinted files remained byte-for-byte unchanged through both LOADs. +Baseline, SAVE2, LOAD, and LOAD2 returned identical sequential and concurrent +text/token strings, every phase exited zero, and cleanup removed both the +authoritative Dict marker and Volume path. + +This vLLM result is a single-host, exactly-two-Hopper, V1, +`FULL_DECODE_ONLY`, PyNCCL P2P/legacy-CUDA-IPC profile. Custom, torch +symmetric, FlashInfer, and fused all-reduce paths plus NCCL CUMEM/NVLS were +disabled. TP>2, V2, multi-host, Ampere/Blackwell, MoE TP, and alternate +collectives remain unsupported or unvalidated. + +The same exact candidate passed the +[native H100 gate](https://modal.com/apps/modal-labs/rahul-dev/ap-0aeSCJNb6eHQhqdMkpyfXk): +10 focused VMM/RNG tests and 823 core tests. It also passed the strict +[2×H100 DeepEP matrix](https://modal.com/apps/modal-labs/rahul-dev/ap-Q8mb8U4EYuE5Jo6Zb0xwrH): +the pinned PyTorch preflight, `fabric`, `nvl_ipc`, `nvl_ipc_prealloc`, and +Modal CLI all returned zero. This is standalone transport/replay evidence, not +general MoE TP or SGLang EP validation. + +SGLang TP=2 is validated only for the pinned single-host symmetric-memory +profile: Foundry +`a8e43442421a80fea7991ee26f97f1003b170e29`, unmodified upstream SGLang +`a25164bda355250015c0ba6d8821f0d4b37abe7d`, and Qwen3-8B revision +`b968826d9c46dd6066d109eabc6255188de91218`. A lightweight +`foundry_sglang_plugin` entry point activates nine named `HookRegistry` targets +with zero SGLang edits. It retains one torch symmetric-memory TP communicator +allocated outside Foundry VMM, performs direct multimem all-gather into that +stable buffer, captures one padded batch-size-256 full-decode graph per rank, +and restores through SGLang's backend-driven capture lifecycle. + +The exact SGLang and model SHAs are validation provenance enforced by the Modal +harness, which checks out the engine SHA and passes the model revision. They +are not inspected by `_validate_profile`, and arbitrary shell-installed SGLang +revisions are not rejected at startup. Shell recipe users must install the +matching SGLang pin; runtime fail-closed checks cover topology, backend, and +settings. + +Plain-NCCL fresh-process SGLang restore is explicitly unsupported. Disproving +app +[`ap-2qSQUKlY24L4RrQkMpUrmc`](https://modal.com/apps/modal-labs/rahul-dev/ap-2qSQUKlY24L4RrQkMpUrmc) +passed structural checks—36 graphs restored per rank, matching graph counts, +offsets, immutable archive bytes, and P2P/IPC topology—but both fresh LOAD +processes produced the same corrupt decode output while baseline and SAVE2 +matched. Those observations reject this plain-NCCL restore path. Archived +process-specific communicator and graph-sequence state is the leading forensic +explanation, not treatment-proven causality; structural restoration alone is +not output correctness. + +The focused symmetric SAVE→LOAD treatment +[`ap-5wO9cFhRZ7ITsLDDv27Xfh`](https://modal.com/apps/modal-labs/rahul-dev/ap-5wO9cFhRZ7ITsLDDv27Xfh) +and final five-phase matrix +[`ap-6fYcSUFLOxocBSjNxakNGZ`](https://modal.com/apps/modal-labs/rahul-dev/ap-6fYcSUFLOxocBSjNxakNGZ) +are `VERIFIED`. Final evidence includes one graph per rank at offset +`68908220416`; 73 symmetric all-reduce plus one symmetric all-gather and zero +NCCL collective nodes per rank; exact baseline/SAVE2/LOAD/LOAD2 sequential and +batch-8 outputs; reproducible SAVE/SAVE2 semantic fingerprints; byte-immutable +LOAD/LOAD2 archives; observed batch-1 and batch-8 replay; clean process exits +and cleanup. Startup times were 217.2 s baseline, 103.0 s SAVE, 103.1 s SAVE2, +and 61.0 s for each LOAD. + +The pinned vLLM profile remains single-host dense Qwen3-8B TP=2 with +full-decode graphs and plain NCCL P2P / legacy CUDA IPC. The SGLang result does +not establish general/upstream TP or plain-NCCL support. SGLang TP>2, multi-host, +custom/NCCL-symmetric/fused collectives, CUMEM, NVLS, EP, DP-attention, +speculative decode, LoRA, pdmux, and unpinned revisions remain unsupported; +Ampere and Blackwell are unvalidated. + +Historical SGLang branch runs from draft +[PR #4](https://github.com/modal-projects/foundry/pull/4) and +[PR #2](https://github.com/modal-projects/foundry/pull/2) are **pre-port fork +evidence only**, not evidence for the zero-edit pinned plugin. +[Upstream Discussion #5](https://github.com/orgs/foundry-org/discussions/5) +explains that NCCL initialization is not generally deterministic and identifies +torch symmetric memory as the intended general backend; official TP support +remains open in [issue #6](https://github.com/foundry-org/foundry/issues/6). + The adapted vLLM / SGLang / TensorRT-LLM forks will be released alongside this repo at `foundry-org/vllm`, `foundry-org/sglang`, `foundry-org/TensorRT-LLM`. ### Performance diff --git a/RELEASE.md b/RELEASE.md index 6586a781..21ddb661 100644 --- a/RELEASE.md +++ b/RELEASE.md @@ -1,10 +1,136 @@ -# Foundry 0.0.2 +# Release Notes + +## Unreleased + +- **Experimental dense tensor parallel profiles.** vLLM dense TP=2 is verified + at Foundry `a8e43442421a80fea7991ee26f97f1003b170e29`, vLLM + `4309c257d3f639e5490d3811293c890c61c76f29`, wheel parent + `6cbe448eed751824d608faf9078ef84724d621c1`, and + `Qwen/Qwen3-8B@b968826d9c46dd6066d109eabc6255188de91218`. The + [final six-phase run](https://modal.com/apps/modal-labs/rahul-dev/ap-YfKALdMwnyJQmzQoVvZSMU) + used exactly two H100 80GB HBM3 GPUs in every phase with CUDA + `13.0.1-cudnn-devel-ubuntu24.04`, driver `580.95.05`, PyTorch + `2.11.0+cu130`, vLLM + `0.1.dev16924+g4309c257d.precompiled`, and NCCL `2.28.9+cuda13.0`. + All 17 checks passed. + + The accepted archive had 64 JSON/64 `.cugraph` files per rank, 68 semantic + entries per rank, and 265 fingerprinted files. SAVE/SAVE2 reproduced + inventory, semantics, and per-rank offset `70724354048`; LOAD and LOAD2 each + restored `64/64` graphs on both ranks, matched that live offset, and left all + 265 fingerprinted files byte-for-byte unchanged. Baseline, SAVE2, LOAD, and + LOAD2 had exact sequential and concurrent text/token-string equality. All six + child processes and the Modal CLI exited zero, then cleanup removed both the + Dict publication marker and Volume directory. + + The vLLM envelope is single-host, exactly two Hopper GPUs, V1, + `FULL_DECODE_ONLY`, and PyNCCL P2P/legacy CUDA IPC. NCCL CUMEM/NVLS and + custom, torch-symmetric, FlashInfer, and fused all-reduce paths are disabled. + TP>2, V2, multi-host, Ampere/Blackwell, MoE TP, and alternate collectives are + unsupported or unvalidated. The exact candidate also passed + [native H100 validation](https://modal.com/apps/modal-labs/rahul-dev/ap-0aeSCJNb6eHQhqdMkpyfXk) + with 10 focused and 823 core tests, after a 320-test local gate plus compile, + pre-commit, diff, and direct-main lineage checks. + + SGLang is distinct: its only validated TP result is + the pinned single-host Qwen3-8B TP=2 profile at Foundry + `a8e43442421a80fea7991ee26f97f1003b170e29`, unmodified upstream SGLang + `a25164bda355250015c0ba6d8821f0d4b37abe7d`, and Qwen3-8B revision + `b968826d9c46dd6066d109eabc6255188de91218`. The SGLang + path uses a retained torch symmetric-memory communicator outside Foundry VMM, + direct stable multimem gather, one padded batch-size-256 full-decode graph, + and backend-driven LOAD. + + The exact engine/model SHAs are validation provenance enforced by the Modal + harness, not by `_validate_profile` or arbitrary shell-install startup + rejection. Shell recipe users must install the matching SGLang pin; runtime + fail-closed checks cover topology, backend, and settings. + + Plain-NCCL fresh-process SGLang restore is rejected. In + [`ap-2qSQUKlY24L4RrQkMpUrmc`](https://modal.com/apps/modal-labs/rahul-dev/ap-2qSQUKlY24L4RrQkMpUrmc), + 36 graphs per rank, matching offsets/inventory, immutable archive bytes, and + P2P/IPC topology all passed, yet both fresh LOAD processes produced the same + corrupt decode while baseline and SAVE2 matched. Archived process-specific + communicator and graph-sequence state is the leading forensic explanation, + not treatment-proven causality. The output observation rejects this + plain-NCCL path; the structural checks alone are not correctness evidence. + + The focused treatment + [`ap-5wO9cFhRZ7ITsLDDv27Xfh`](https://modal.com/apps/modal-labs/rahul-dev/ap-5wO9cFhRZ7ITsLDDv27Xfh) + and final matrix + [`ap-6fYcSUFLOxocBSjNxakNGZ`](https://modal.com/apps/modal-labs/rahul-dev/ap-6fYcSUFLOxocBSjNxakNGZ) + are `VERIFIED`. The matrix recorded one graph per rank, offset `68908220416`, + 73 symmetric all-reduce plus one symmetric all-gather and zero NCCL + collective nodes per rank, exact sequential and batch-8 output equality, + SAVE/SAVE2 semantic reproducibility, LOAD/LOAD2 byte immutability, + batch-1/batch-8 replay, phase times `217.2/103.0/103.1/61.0/61.0` seconds, + clean process exits, and completed cleanup. + + Historical draft PR #4/#2 SGLang records are **pre-port fork evidence only**. + The pinned symmetric result is not general/upstream TP or plain-NCCL support. + General TP remains open in + [upstream issue #6](https://github.com/foundry-org/foundry/issues/6); + [Discussion #5](https://github.com/orgs/foundry-org/discussions/5) identifies + torch symmetric memory as the intended general backend. +- **Zero-edit SGLang plugin activation and pinned-main graph binding.** SGLang + is used unmodified at the pinned commit; Foundry activates through the + dependency-light `foundry_sglang_plugin:activate` entry point and installs + nine `HookRegistry` targets. SAVE and LOAD both traverse the pinned upstream + `DecodeCudaGraphRunner.capture` lifecycle. The backend captures on SAVE and, + on LOAD, performs one pooled `start_graph_builds` followed by ordered + `finish_one_graph_load` calls without native capture or resave. The serialized + `MemoryPoolConfig` is reused by hooking + `KVCacheConfigurator._resolve_memory_pool_config`. This replaces the + monolithic `CudaGraphRunner` fork patch and its ~47-line `sglang/` footprint. + Qwen3.5-122B-A10B remains unsupported because its hybrid + linear-attention/MoE path requires capture-safe warmup, Mamba-state handling, + and a supported fused all-reduce route. + SGLang TP>2, multi-host, custom/NCCL-symmetric/fused collectives, CUMEM, NVLS, + EP, DP-attention, speculative decode, LoRA, pdmux, and unpinned revisions are + outside the validated profile. +- **No-fabric DeepEP IPC.** The VMM-IPC bridge transports shareable file + descriptors with `SCM_RIGHTS`. A strict + [2×H100 matrix](https://modal.com/apps/modal-labs/rahul-dev/ap-Q8mb8U4EYuE5Jo6Zb0xwrH) + passed the pinned PyTorch CPU preflight plus exact SAVE/LOAD output and + sentinel-overwrite checks for `fabric`, `nvl_ipc`, and whole-preallocation + `nvl_ipc_prealloc`. It used Foundry + `a8e43442421a80fea7991ee26f97f1003b170e29`, vLLM + `4309c257d3f639e5490d3811293c890c61c76f29`, DeepEP + `29d31c095796f3c8ece47ee9cdcc167051bbeed9`, NVSHMEM `3.3.24`, + CUDA `13.0.1`, and PyTorch `2.11.0+cu130`. The preflight, all three cases, + and Modal CLI returned zero; the application status was `passed`. + + This establishes the standalone transport/replay paths, not general vLLM + MoE serving, MoE TP, or SGLang EP. The current SGLang HookRegistry EP profile + remains quarantined. The IPC work extends + [upstream PR #3](https://github.com/foundry-org/foundry/pull/3). Historical + branch evidence covers BF16 on 2×H100; GB300/aarch64 FP8 and RDC relinking + reported in [upstream issue #1](https://github.com/foundry-org/foundry/issues/1) + remain outside that scope. +- **FA2 cursor restoration.** When every graph in a load set has zero RNG + increment, LOAD skips unused generator-state restoration instead of + suspending the tracked allocation region. Nonzero-RNG graphs keep their + seed/offset tensors in deterministic memory. This incorporates + [upstream PR #4](https://github.com/foundry-org/foundry/pull/4) and the root + cause from [issue #2](https://github.com/foundry-org/foundry/issues/2), but + does not establish broad Ampere support. +- **Restoration hardening.** Preallocated IPC chunks carry a generation in + their mapping identity so a freed and recreated chunk at the same address + cannot alias stale mappings. The Python graph context always passes the + required pool argument, with the default resolving to `(0, 0)`. CUDA's + zero-alignment sentinel is normalized to Foundry's 2 MiB VMM alignment for + reservation and replay, preserving the deterministic cursor with one + advance. + +--- + +## Foundry 0.0.2 SGLang graduates to a fully validated engine. This release brings the SGLang integration to parity with vLLM across single GPU, data parallel, and expert parallel — with a self-contained recipe and no vLLM build dependency for EP. -## Highlights +### Highlights - **SGLang single GPU / DP / EP all validated end-to-end.** SAVE → LOAD → query verified on single-GPU Qwen3-1.7B / 4B / 14B, @@ -23,7 +149,7 @@ parallel — with a self-contained recipe and no vLLM build dependency for EP. `data_parallel_controller.py`). All save/load logic lives in the integration layer; the edits are inert unless `--foundry-graph-extension-config-path` is set. -## Engine integrations +### Engine integrations - **SGLang** — integration for SGLang v0.5.13. Working configurations: single GPU, data parallel (DP), expert parallel (EP, DeepEP low-latency + DP-attention with fa3). Self-contained recipes under `recipe/sglang/` (shared TOML pair + @@ -45,7 +171,7 @@ parallel — with a self-contained recipe and no vLLM build dependency for EP. on LOAD, with `init_nvshmem_for_loaded_modules` run once on LOAD before any NVSHMEM-kernel graph replays. -## Fixes +### Fixes - **Per-rank VMM device binding (DP/TP/EP).** `set_allocation_region` binds to the current CUDA device, so the integration now calls `set_device(gpu_id)` @@ -63,7 +189,7 @@ parallel — with a self-contained recipe and no vLLM build dependency for EP. capture loop never sets), and the FlashInfer per-bs pre-pass gated off for fa3 while still populating `decode_cuda_graph_metadata` post-load for replay. -## Docs +### Docs - New `docs/sglang/` set (overview, direct-edits, hooks, memory-lifecycle, save-load-workflow, memory-consistency) and a self-contained `recipe/sglang/` @@ -71,10 +197,6 @@ parallel — with a self-contained recipe and no vLLM build dependency for EP. - Top-level README parallelism status table updated to mark SGLang single GPU, DP, and EP as validated. ---- - -## Previous Releases - ## Foundry 0.0.1 First public release of Foundry — a CUDA-graph persistence library that @@ -82,7 +204,7 @@ captures an entire model's CUDA graphs (plus their device context: modules, workspaces, VMM layout) once and replays them at startup, eliminating compile, warmup, and capture from cold-start time. -## Highlights +### Highlights - **Deterministic memory layout — zero patching on graph load.** Foundry indirects memory allocation to the same reserved memory region @@ -99,7 +221,7 @@ warmup, and capture from cold-start time. On load, the template is rebuilt once per group and instances are reconstructed on demand, keeping load fast and asynchronous. -## Engine integrations +### Engine integrations - **vLLM** — compatible with vLLM v0.21. Working configurations: single GPU, data parallel (DP), expert parallel (EP, DeepEP low-latency). End-to-end @@ -111,31 +233,31 @@ warmup, and capture from cold-start time. integration layer (`foundry/integration//`); engine forks contain only minimal hook calls. -## Verified kernel & comm support +### Verified kernel & comm support - **cuBLAS NVJET** kernels (Hopper+). - **torch.compile** modules. - **NVSHMEM / DeepEP** validated. - **DeepGEMM FP8 MoE** validated. -## Dependency +### Dependency - **PyTorch 2.11.0** (compatible with 2.9 – 2.11). - **CUDA 12+**, CMake 4.0+, Boost. -## Documentation +### Documentation - Integration design notes and per-engine recipes under `docs/` and `recipe/`. - vLLM recipe README covers save/load workflow, archive layout, and required env settings. -## Repository hygiene +### Repository hygiene - Open-source pre-commit hooks (ruff, ruff-format, clang-format, markdownlint, actionlint, DCO sign-off). - Smoke tests covering re-export imports and archive round-trip. -## Roadmap +### Roadmap - Adapted vLLM and SGLang forks published alongside the release. - Tensor parallel support. diff --git a/ROADMAP.md b/ROADMAP.md index aa139d65..9ad1e8cd 100644 --- a/ROADMAP.md +++ b/ROADMAP.md @@ -24,14 +24,84 @@ ## Stage 4: vLLM/SGLang Integration -- [x] Sync with latest vLLM release +- [x] Integrate the pinned vLLM fork - [x] EP on vLLM with DeepEP - - [ ] TP on vLLM + - [x] Experimental TP on vLLM + - [x] Strict experimental dense Qwen3-8B TP=2 recipe and CPU contracts + - [x] Final six-phase 2×H100 acceptance + ([`ap-YfKALdMwnyJQmzQoVvZSMU`](https://modal.com/apps/modal-labs/rahul-dev/ap-YfKALdMwnyJQmzQoVvZSMU)) - [x] Quantized MoE with DeepGemm - [x] Drop-in integration layer for CUDA graph persistence in vLLM -- [x] Sync with latest SGLang release - - [ ] EP on SGLang - - [ ] TP on SGLang +- [x] Integrate SGLang through the zero-edit `sglang.srt.plugins` entry point + - [x] EP on SGLang (pre-port fork) + - [ ] Official TP on SGLang + - [x] Strict experimental dense Qwen3-8B TP=2 recipe + - [x] Pinned-main plugin port (pin `a25164bda355250015c0ba6d8821f0d4b37abe7d`, model revision `b968826d9c46dd6066d109eabc6255188de91218`) + - [x] `FullCudaGraphBackend` / `DecodeCudaGraphRunner` / `ShapeKey` SAVE-LOAD binding + `_resolve_memory_pool_config` reuse + - [x] Torch symmetric-memory treatment with retained communicator, direct multimem gather, and one padded batch-size-256 graph + - [x] Focused 2×H100 SAVE→LOAD treatment ([`ap-5wO9cFhRZ7ITsLDDv27Xfh`](https://modal.com/apps/modal-labs/rahul-dev/ap-5wO9cFhRZ7ITsLDDv27Xfh)) + - [x] Final five-phase 2×H100 matrix ([`ap-6fYcSUFLOxocBSjNxakNGZ`](https://modal.com/apps/modal-labs/rahul-dev/ap-6fYcSUFLOxocBSjNxakNGZ)) + +The completed vLLM item is limited to Foundry +`a8e43442421a80fea7991ee26f97f1003b170e29`, vLLM +`4309c257d3f639e5490d3811293c890c61c76f29`, wheel parent +`6cbe448eed751824d608faf9078ef84724d621c1`, and +`Qwen/Qwen3-8B@b968826d9c46dd6066d109eabc6255188de91218` on one host +with exactly two H100s. All 17 checks passed: 64 JSON/64 binary graphs per +rank, `64/64` restoration in both LOADs, offset `70724354048`, exact +baseline/SAVE2/LOAD/LOAD2 sequential and concurrent outputs, 68 semantic +entries per rank, 265 complete-tree fingerprints with byte-for-byte +immutability, all phase/Modal exit codes were zero, and Dict/Volume cleanup +completed. Local validation passed 320 tests plus compile, pre-commit, diff, +and lineage gates; the native H100 run passed 10 focused and 823 core tests +([`ap-0aeSCJNb6eHQhqdMkpyfXk`](https://modal.com/apps/modal-labs/rahul-dev/ap-0aeSCJNb6eHQhqdMkpyfXk)). + +Strict DeepEP transport/replay also passed at the exact final candidate +`a8e43442421a80fea7991ee26f97f1003b170e29` +([`ap-Q8mb8U4EYuE5Jo6Zb0xwrH`](https://modal.com/apps/modal-labs/rahul-dev/ap-Q8mb8U4EYuE5Jo6Zb0xwrH)). +The pinned PyTorch preflight, `fabric`, `nvl_ipc`, `nvl_ipc_prealloc`, and +Modal CLI all returned zero. This does not validate general DeepEP, MoE TP, or +SGLang EP. + +SGLang TP=2 is supported only against the pinned upstream commit +`a25164bda355250015c0ba6d8821f0d4b37abe7d` via the `sglang.srt.plugins` entry +point, Foundry `a8e43442421a80fea7991ee26f97f1003b170e29`, and Qwen3-8B revision +`b968826d9c46dd6066d109eabc6255188de91218`. The zero-edit lightweight plugin +installs nine `HookRegistry` targets. Its validated path retains a torch +symmetric-memory communicator outside Foundry VMM, performs direct stable +multimem gather, captures one padded batch-size-256 graph per rank, and uses +the backend-driven LOAD lifecycle. The final run proved offset `68908220416`, +73 symmetric all-reduce plus one symmetric all-gather and zero NCCL collective +nodes per rank, exact sequential and batch-8 output equality, SAVE/SAVE2 +semantic reproducibility, LOAD/LOAD2 byte immutability, batch-1/batch-8 replay, +and clean shutdown/cleanup. + +The engine/model SHAs are validation provenance enforced by the Modal harness, +not fields inspected by `_validate_profile`. Shell recipe users must install +the matching SGLang pin; runtime fail-closed checks cover topology, backend, +and settings rather than arbitrary installed revisions. + +Plain-NCCL fresh-process SGLang restore is rejected: +[`ap-2qSQUKlY24L4RrQkMpUrmc`](https://modal.com/apps/modal-labs/rahul-dev/ap-2qSQUKlY24L4RrQkMpUrmc) +restored 36 graphs per rank with matching offsets, inventory, archive bytes, +and P2P/IPC topology, yet both LOAD processes emitted the same corrupt decode +while baseline and SAVE2 matched. Archived process-specific communicator and +graph-sequence state is the leading forensic explanation, not +treatment-proven causality. The observed corruption rejects plain-NCCL restore; +the structural pass alone did not establish output correctness. + +This completes only the pinned single-host Qwen3-8B TP=2 symmetric profile. It +is not general/upstream TP or plain-NCCL support. TP>2, multi-host, +custom/NCCL-symmetric/fused collectives, CUMEM, NVLS, EP, DP-attention, +speculative decode, LoRA, pdmux, and unpinned SGLang revisions remain +unsupported. + +Historical SGLang TP=2 branch runs from modal-projects draft +[PR #4](https://github.com/modal-projects/foundry/pull/4) on 2×H100 and +[PR #2](https://github.com/modal-projects/foundry/pull/2) on 2×H200 came from +the **pre-port fork integration** and are supporting evidence only; they do not +broaden the pinned symmetric-profile result above. Official TP remains tracked in +[upstream issue #6](https://github.com/foundry-org/foundry/issues/6). ## Stage 5: Disaggregated and Large-Scale Serving diff --git a/csrc/CUDAGraph.cpp b/csrc/CUDAGraph.cpp index 25df3524..c1135668 100644 --- a/csrc/CUDAGraph.cpp +++ b/csrc/CUDAGraph.cpp @@ -1528,15 +1528,30 @@ GraphLoadResult CUDAGraph::load(const std::string& json_path, MempoolId_t pool) } const json::array& generators_array = root.at("generators").as_array(); + bool has_rng_generator = false; for (const auto& gen_val : generators_array) { const json::object& gen_obj = gen_val.as_object(); - uint64_t state_id = gen_obj.at("id").to_number(); - uint64_t seed = gen_obj.at("seed").to_number(); - uint64_t wholegraph_increment = gen_obj.at("wholegraph_increment").to_number(); + auto increment_it = gen_obj.find("wholegraph_increment"); + // Missing metadata is treated conservatively: registering may expose an + // allocation mismatch, but skipping could silently leave RNG kernel + // arguments pointing at the SAVE process's seed/offset tensors. + if (increment_it == gen_obj.end() || increment_it->value().to_number() != 0) { + has_rng_generator = true; + break; + } + } - auto state = global_generator_state_registry.get_state_from_id(state_id, seed); - state->register_graph(reinterpret_cast(graph.get())); - graph->captured_generator_states_[state] = wholegraph_increment; + if (has_rng_generator) { + for (const auto& gen_val : generators_array) { + const json::object& gen_obj = gen_val.as_object(); + uint64_t state_id = gen_obj.at("id").to_number(); + uint64_t seed = gen_obj.at("seed").to_number(); + uint64_t wholegraph_increment = gen_obj.at("wholegraph_increment").to_number(); + + auto state = global_generator_state_registry.get_state_from_id(state_id, seed); + state->register_graph(reinterpret_cast(graph.get())); + graph->captured_generator_states_[state] = wholegraph_increment; + } } const json::object& allocator_events = root.at("allocator_events").as_object(); diff --git a/csrc/CUDAGraphParallel.cpp b/csrc/CUDAGraphParallel.cpp index 6ebfc0f6..7440b210 100644 --- a/csrc/CUDAGraphParallel.cpp +++ b/csrc/CUDAGraphParallel.cpp @@ -1658,6 +1658,7 @@ std::shared_ptr start_graph_builds_impl( go["seed"] = gens[g].seed; go["wholegraph_increment"] = gens[g].wholegraph_increment; gens_array.push_back(go); + pending->has_rng_generators |= gens[g].wholegraph_increment != 0; } pending->entries[i].generators_meta = boost::json::value(std::move(gens_array)); } @@ -1690,6 +1691,15 @@ std::shared_ptr start_graph_builds_impl( auto gen_it = pre_root.find("generators"); if (gen_it != pre_root.end()) { + const boost::json::array& generators = gen_it->value().as_array(); + for (const auto& gen_val : generators) { + const boost::json::object& generator = gen_val.as_object(); + auto increment_it = generator.find("wholegraph_increment"); + if (increment_it == generator.end() || increment_it->value().to_number() != 0) { + pending->has_rng_generators = true; + break; + } + } pending->entries[i].generators_meta = std::move(gen_it->value()); pre_root.erase(gen_it); } @@ -2030,11 +2040,16 @@ std::shared_ptr start_graph_builds_impl( // SAVE mode. // ============================================================================ -// Per-entry finish. Idempotent on build_complete_ wait (shared_future). +// Per-entry finish. Idempotent on build_complete_ (shared_future). GraphLoadResult finish_one_graph_load_impl(std::shared_ptr pending, size_t index, ReconstructTensorFn reconstruct_fn) { if (pending->build_complete_.valid()) { - pending->build_complete_.wait(); + // get() (not wait()) so a failure the background build thread stored via + // set_exception rethrows HERE, at the per-entry finish boundary, before we + // touch the entry at index (a failed build leaves the entry half-linked). + // shared_future::get() may be called repeatedly, so this stays idempotent + // and the batched path (get() once, then finish-one per index) is unaffected. + pending->build_complete_.get(); } if (index >= pending->entries.size()) { @@ -2045,9 +2060,14 @@ GraphLoadResult finish_one_graph_load_impl(std::shared_ptr pe auto& entry = pending->entries[index]; - // Register deferred generators (before allocator replay to match - // SAVE mode timing where generators are created before graph capture). - if (pending->registry && !entry.generators_meta.is_null()) { + // register_graph lazily creates seed/offset tensors before allocator replay. + // They must stay in the deterministic region when a captured kernel consumes + // Philox state, because their pointers are embedded in that kernel's args. + // When the entire pending set has zero RNG increment, no captured node + // references those tensors; skip registration entirely. This avoids the + // caching-allocator segment growth that caused a 2 MiB A100/FA2 cursor drift + // without replacing deterministic pointers with ordinary CUDA allocations. + if (pending->has_rng_generators && pending->registry && !entry.generators_meta.is_null()) { const boost::json::array& gen_array = entry.generators_meta.as_array(); for (const auto& gen_val : gen_array) { const boost::json::object& gen_obj = gen_val.as_object(); diff --git a/csrc/binding.cpp b/csrc/binding.cpp index 43d1f7ee..6d2622d8 100644 --- a/csrc/binding.cpp +++ b/csrc/binding.cpp @@ -39,8 +39,8 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { "Preallocate memory in the allocation region for fast subsequent allocations"); m.def( - "free_preallocated_region", []() { ::foundry::free_preallocated_region(); }, - "Free the preallocated memory region"); + "free_preallocated_region", []() { return ::foundry::free_preallocated_region(); }, + "Free the preallocated memory region; return whether every cleanup phase succeeded"); m.def( "get_current_alloc_offset", []() { return ::foundry::get_current_alloc_offset(); }, diff --git a/csrc/hook.cpp b/csrc/hook.cpp index d552d08d..a5a65131 100644 --- a/csrc/hook.cpp +++ b/csrc/hook.cpp @@ -14,18 +14,30 @@ #include #include #include +#include #include #include #include #include #include #include +#include +#include +#include +#include +#include +#include +#include #include #include #include #include #include #include "hook.h" +#include "foundry/vmm_cleanup.h" +#include "foundry/vmm_ipc_chunk_key.h" +#include "foundry/vmm_ipc_ownership.h" +#include "foundry/vmm_ipc_reservation.h" namespace fs = boost::filesystem; @@ -213,6 +225,7 @@ struct AllocMetadata { CUmemGenericAllocationHandle handle; CUdeviceptr region_base; bool from_preallocation; // If true, this allocation is carved from preallocated memory + foundry::vmm::CleanupState cleanup; }; struct ThreadLocalStorage { @@ -227,6 +240,7 @@ struct ThreadLocalStorage { size_t preallocated_start_addr; size_t preallocated_end_addr; bool has_preallocation; + foundry::vmm::CleanupState preallocation_cleanup; // Cached values to avoid repeated driver calls CUdevice cached_device; @@ -243,6 +257,7 @@ struct ThreadLocalStorage { preallocated_start_addr(0), preallocated_end_addr(0), has_preallocation(false), + preallocation_cleanup{}, cached_device(0), cached_granularity(0), device_cached(false) {} @@ -253,6 +268,28 @@ static std::once_flag default_allocation_region_flag; static boost::unordered::concurrent_flat_map global_alloc_metadata; static boost::unordered::concurrent_flat_map global_carved_reserve_metadata; +// Defined in the VMM-IPC section below (inside the extern "C" block, hence +// the matching linkage here); closes and forgets the exported fd for a VA +// when its allocation is freed (an exported fd keeps the physical pages +// alive and would otherwise serve stale memory to importers). +extern "C" { +static void vmm_ipc_invalidate_export(CUdeviceptr dptr); +} + +// Process-global mirrors of live LOAD-mode preallocated chunks. The +// authoritative state is thread-local, but cuIpcGetMemHandle may run on a +// different thread. Keep one coherent record per range so concurrent devices +// or regions cannot overwrite each other's handle/base/size tuple. +struct PreallocatedChunk { + size_t size; + CUmemGenericAllocationHandle handle; + uint64_t generation; + foundry::vmm::CleanupState cleanup; +}; +static std::mutex preallocated_chunks_mutex; +static std::map preallocated_chunks; +static std::atomic next_preallocated_chunk_generation{1}; + struct HookAllocationEvent { enum class Type { Alloc, Free, Reserve }; Type type; @@ -2276,7 +2313,7 @@ CUresult cuMemAlloc_v2(CUdeviceptr* dptr, size_t bytesize) { CUdeviceptr region_end_addr = region_base_addr + tls_storage.region.size; // Fast path: check if we can serve from preallocated memory - if (tls_storage.has_preallocation && + if (tls_storage.has_preallocation && !tls_storage.preallocation_cleanup.started && (target_addr + aligned_size) <= tls_storage.preallocated_end_addr) { *dptr = target_addr; @@ -2491,7 +2528,7 @@ CUresult cuMemAllocPitch_v2(CUdeviceptr* dptr, size_t* pPitch, size_t WidthInByt CUdeviceptr region_end_addr = region_base_addr + tls_storage.region.size; // Fast path: check if we can serve from preallocated memory - if (tls_storage.has_preallocation && + if (tls_storage.has_preallocation && !tls_storage.preallocation_cleanup.started && (target_addr + aligned_size) <= tls_storage.preallocated_end_addr) { *dptr = target_addr; *pPitch = pitch; @@ -2543,6 +2580,10 @@ CUresult cuMemAllocPitch_v2(CUdeviceptr* dptr, size_t* pPitch, size_t WidthInByt prop.type = CU_MEM_ALLOCATION_TYPE_PINNED; prop.location.type = CU_MEM_LOCATION_TYPE_DEVICE; prop.location.id = device; + // Tracked pitched allocations are treated as IPC-exportable downstream + // (cuIpcGetMemHandle -> cuMemExportToShareableHandle); request a shareable + // POSIX fd handle here or that export fails late for pitched buffers. + prop.requestedHandleTypes = CU_MEM_HANDLE_TYPE_POSIX_FILE_DESCRIPTOR; result = mem_create_func(&allocHandle, aligned_size, &prop, 0); if (result != CUDA_SUCCESS) { @@ -2629,32 +2670,90 @@ CUresult cuMemFree_v2(CUdeviceptr dptr) { bool found = false; AllocMetadata metadata; + global_alloc_metadata.visit(dptr, [&](std::pair& kv) { + found = true; + metadata = kv.second; + if (!kv.second.from_preallocation) { + // Publish cleanup intent before invalidating the export. Any retry sees + // the retained phase state and skips driver operations that succeeded. + kv.second.cleanup.started = true; + } + }); - size_t visit_count = - global_alloc_metadata.visit(dptr, [&](std::pair& kv) { - found = true; - metadata = kv.second; + if (found) { + if (metadata.from_preallocation) { + // Preallocated carves do not own physical VMM resources. Remove only the + // carve's allocation metadata; free_preallocated_region() owns chunk + // unmap/release and invalidates the chunk-level export. + global_alloc_metadata.erase_if([dptr](const std::pair& kv) { + return kv.first == dptr; }); - if (found) { - // For preallocated memory, we just remove tracking but don't actually free - // The memory will be freed when free_preallocated_region() is called - if (!metadata.from_preallocation) { - typedef CUresult (*cuMemUnmap_t)(CUdeviceptr, size_t); - auto mem_unmap_func = - (cuMemUnmap_t)CUDA_DRIVER_CALL(cuda_driver_entry_table, CUDA_ENTRY_cuMemUnmap); + if (hook_recording_enabled.load()) { + std::lock_guard lock(hook_events_mutex); + HookAllocationEvent event; + event.type = HookAllocationEvent::Type::Free; + event.size = 0; + event.ptr = dptr; + hook_alloc_events.push_back(event); + } - typedef CUresult (*cuMemRelease_t)(CUmemGenericAllocationHandle); - auto mem_release_func = - (cuMemRelease_t)CUDA_DRIVER_CALL(cuda_driver_entry_table, CUDA_ENTRY_cuMemRelease); +#ifdef HOOK_DEBUG + fprintf(stderr, "[HOOK] cuMemFree_v2 ptr=%llu size=%zu (PREALLOCATED carve metadata only)\n", + (unsigned long long)dptr, metadata.size); +#endif + return CUDA_SUCCESS; + } - typedef CUresult (*cuMemAddressFree_t)(CUdeviceptr, size_t); - auto mem_addr_free_func = (cuMemAddressFree_t)CUDA_DRIVER_CALL(cuda_driver_entry_table, - CUDA_ENTRY_cuMemAddressFree); + // Once cleanup starts this allocation must no longer be exportable, even + // if a physical cleanup phase fails and metadata is retained for retry. + vmm_ipc_invalidate_export(dptr); + + typedef CUresult (*cuMemUnmap_t)(CUdeviceptr, size_t); + auto mem_unmap_func = + (cuMemUnmap_t)CUDA_DRIVER_CALL(cuda_driver_entry_table, CUDA_ENTRY_cuMemUnmap); + typedef CUresult (*cuMemRelease_t)(CUmemGenericAllocationHandle); + auto mem_release_func = + (cuMemRelease_t)CUDA_DRIVER_CALL(cuda_driver_entry_table, CUDA_ENTRY_cuMemRelease); + typedef CUresult (*cuMemAddressFree_t)(CUdeviceptr, size_t); + auto mem_addr_free_func = + (cuMemAddressFree_t)CUDA_DRIVER_CALL(cuda_driver_entry_table, CUDA_ENTRY_cuMemAddressFree); - mem_unmap_func(metadata.ptr, metadata.size); - mem_release_func(metadata.handle); - mem_addr_free_func(metadata.ptr, metadata.size); + // Region-owned allocations are subranges of one whole-region VA reservation + // created by set_allocation_region; they own only their mapping and handle. + // IPC imports use an independent reservation (region_base == 0) and own it. + const bool address_free_required = metadata.region_base == 0; + CUresult cleanup_status = CUDA_ERROR_INVALID_VALUE; + bool cleanup_complete = false; + global_alloc_metadata.visit(dptr, [&](std::pair& kv) { + cleanup_status = foundry::vmm::run_cleanup( + kv.second.cleanup, CUDA_SUCCESS, address_free_required, + [&]() { return mem_unmap_func(kv.second.ptr, kv.second.size); }, + [&]() { return mem_release_func(kv.second.handle); }, + [&]() { return mem_addr_free_func(kv.second.ptr, kv.second.size); }, + [&](foundry::vmm::CleanupPhase phase, CUresult status) { + switch (phase) { + case foundry::vmm::CleanupPhase::kUnmap: + fprintf(stderr, "[HOOK] ERROR: cuMemFree_v2: cuMemUnmap failed with error %d\n", + status); + break; + case foundry::vmm::CleanupPhase::kRelease: + fprintf(stderr, "[HOOK] ERROR: cuMemFree_v2: cuMemRelease failed with error %d\n", + status); + break; + case foundry::vmm::CleanupPhase::kAddressFree: + fprintf(stderr, + "[HOOK] ERROR: cuMemFree_v2: cuMemAddressFree failed with error %d\n", + status); + break; + } + }); + cleanup_complete = foundry::vmm::cleanup_complete(kv.second.cleanup, address_free_required); + metadata = kv.second; + }); + + if (!cleanup_complete) { + return cleanup_status; } global_alloc_metadata.erase_if( @@ -2670,11 +2769,11 @@ CUresult cuMemFree_v2(CUdeviceptr dptr) { } #ifdef HOOK_DEBUG - fprintf(stderr, "[HOOK] cuMemFree_v2 ptr=%llu size=%zu (%s)\n", (unsigned long long)dptr, - metadata.size, metadata.from_preallocation ? "PREALLOCATED" : "VMM"); + fprintf(stderr, "[HOOK] cuMemFree_v2 ptr=%llu size=%zu (VMM)\n", (unsigned long long)dptr, + metadata.size); #endif - return CUDA_SUCCESS; + return cleanup_status; } return real_func(dptr); @@ -2702,9 +2801,10 @@ CUresult cuMemAddressReserve(CUdeviceptr* ptr, size_t size, size_t alignment, CU return result; } + const size_t effective_alignment = alignment == 0 ? kAllocAlignment : alignment; constexpr size_t kSmallAllocThreshold = 2ULL << 30; if (size < kSmallAllocThreshold) { - CUdeviceptr aligned_addr = align_to(tls_storage.current_alloc_base_addr, alignment); + CUdeviceptr aligned_addr = align_to(tls_storage.current_alloc_base_addr, effective_alignment); size_t region_base = (size_t)tls_storage.region.base; size_t region_end = region_base + tls_storage.region.size; @@ -2744,7 +2844,7 @@ CUresult cuMemAddressReserve(CUdeviceptr* ptr, size_t size, size_t alignment, CU // constexpr size_t kReserveAlignment = 2ULL << 37; // FIXME(liuxs): This value should be equal to // the first reservation size - CUdeviceptr hint_addr = align_to(tls_storage.current_vmm_reserve_addr, alignment); + CUdeviceptr hint_addr = align_to(tls_storage.current_vmm_reserve_addr, effective_alignment); CUresult result = real_func(ptr, size, alignment, hint_addr, flags); if (result != CUDA_SUCCESS) { fprintf(stderr, "[HOOK] ERROR: cuMemAddressReserve failed with error 0x%x\n", result); @@ -2759,7 +2859,7 @@ CUresult cuMemAddressReserve(CUdeviceptr* ptr, size_t size, size_t alignment, CU (unsigned long long)tls_storage.current_vmm_reserve_addr); } - tls_storage.current_vmm_reserve_addr = align_to(*ptr + size, alignment); + tls_storage.current_vmm_reserve_addr = align_to(*ptr + size, effective_alignment); #ifdef HOOK_DEBUG fprintf(stderr, @@ -2860,11 +2960,409 @@ void* dlsym(void* handle, const char* symbol) { // VMM allocations by translating to cuMemExportToShareableHandle API // ============================================================================= -// Magic marker to identify VMM IPC handles in CUipcMemHandle -static constexpr uint32_t VMM_IPC_MAGIC = 0x564D4D49; // "VMMI" +// Magic marker to identify VMM IPC handles in CUipcMemHandle. +// v2 ("VMM2"): the blob carries {magic, exporter pid, original ptr, aligned size}. +// The POSIX fd itself is transferred via SCM_RIGHTS over a per-process abstract +// unix socket (a raw fd integer is meaningless in another process - the v1 bug +// behind "cuMemImportFromShareableHandle failed with error 999"). +static constexpr uint32_t VMM_IPC_MAGIC = 0x564D4D32; // "VMM2" + +// --------------------------------------------------------------------------- +// VMM-IPC fd transport: each exporting process runs one server thread on an +// abstract unix socket "\0foundry-vmm-ipc.". Importers connect, send the +// 8-byte exporter VA they want, and receive the exported fd via SCM_RIGHTS. +// The server does no CUDA calls: fds are exported on the cuIpcGetMemHandle +// caller's thread and parked in vmm_ipc_exported_fds. +// --------------------------------------------------------------------------- + +// exporter VA -> (allocation handle it was exported from, fd). The handle is +// kept so VA reuse after free/realloc invalidates the cached fd. +static boost::unordered::concurrent_flat_map> + vmm_ipc_exported_fds; +static std::mutex vmm_ipc_server_mutex; +static int vmm_ipc_listen_fd = -1; +// PID that owns this process's VMM-IPC state. Once IPC starts, a fork child +// cannot safely use inherited C++ mutexes/concurrent-map locks or CUDA handles +// (the owning thread does not survive fork), so reject it before touching that +// state. Fork before the first VMM-IPC use remains safe: the owner is still -1. +static std::atomic vmm_ipc_owner_pid{-1}; +// The owner pid is read (and CAS'd) on the async-signal-safe fork-child +// rejection path, which runs in a just-forked child of a multithreaded process +// before exec. A lock-backed atomic could deadlock there (the backing lock may +// be held by a thread that did not survive fork), so require it be lock-free. +static_assert(std::atomic::is_always_lock_free, + "vmm_ipc_owner_pid must be lock-free: it is accessed on the " + "async-signal-safe fork-child rejection path"); +static pid_t vmm_ipc_server_pid = -1; +// Per-process random token carried in the handle blob and verified by the +// server: defends against pid reuse (a recycled pid + the deterministic +// shared-base VAs would otherwise let an importer silently fetch a different +// process's allocation). +static uint64_t vmm_ipc_token = 0; + +// Reject a post-claim fork child using only async-signal-safe operations. This +// can run in a just-forked child of a multithreaded process (before exec), so +// it must not touch stdio/malloc/locks: it writes a fixed compile-time message +// via write(2) and drops the inherited listener fd with close(2) (the fd table +// is per-process, so this does not affect the parent). +static void vmm_ipc_reject_fork_child() { + ssize_t written = write(STDERR_FILENO, foundry::vmm_ipc::kForkChildRejectMessage, + sizeof(foundry::vmm_ipc::kForkChildRejectMessage) - 1); + (void)written; + // Drop the inherited abstract-socket listener so this child does not keep the + // parent's name pinned. We deliberately do NOT walk vmm_ipc_exported_fds: its + // concurrent map cannot be locked safely here. The child therefore inherits + // duplicate references to the parent's parked export fds. Accepted limitation: + // while this rejected child stays alive without exec, those duplicated fds + // keep the exporter's physical allocations pinned even after the parent frees + // them. The parked fds are FD_CLOEXEC, so an exec closes them immediately + // (the supported spawn/exec pattern) and a plain exit closes them too; the + // pin lasts only until this child exits or execs, which it is expected to do + // promptly after this rejection. + int listen_fd = vmm_ipc_listen_fd; + if (listen_fd >= 0) { + close(listen_fd); + vmm_ipc_listen_fd = -1; + } +} + +// Atomically claim process ownership for the first VMM-IPC use. Returns true +// for the owner (including repeat calls) and false for a post-claim fork child, +// which is rejected async-signal-safely and must abort before touching any +// inherited mutex/map/handle. Fork before the first claim is allowed: the child +// observes an unclaimed owner and becomes its own owner. +static bool vmm_ipc_claim_process() { + pid_t pid = getpid(); + pid_t expected = -1; + vmm_ipc_owner_pid.compare_exchange_strong(expected, pid, std::memory_order_acq_rel, + std::memory_order_acquire); + // `expected` now holds the prior owner: -1 if we just claimed, else the + // existing owner (== pid on a repeat call from the owner process). + if (foundry::vmm_ipc::decide_ipc_ownership((long)expected, (long)pid) == + foundry::vmm_ipc::OwnershipDecision::kProceed) { + return true; + } + vmm_ipc_reject_fork_child(); + return false; +} + +// Wire format of a fetch request. +struct VmmIpcRequest { + uint64_t ptr; + uint64_t token; +}; + +static void vmm_ipc_set_timeouts(int fd) { + struct timeval tv = {30, 0}; + setsockopt(fd, SOL_SOCKET, SO_RCVTIMEO, &tv, sizeof(tv)); + setsockopt(fd, SOL_SOCKET, SO_SNDTIMEO, &tv, sizeof(tv)); +} + +// Imported whole-chunk mappings (LOAD-mode exports: a buffer carved from the +// peer's preallocated chunk is shared by exporting the chunk handle plus an +// offset; cuMemMap cannot map a sub-range of an imported handle, so we map +// the entire peer chunk once and hand out interior pointers). The per-process +// token is part of the key so PID reuse cannot return an old exporter's mapping. +// Refcounted by open/close calls. +struct VmmIpcChunkMapping { + CUdeviceptr local_base; + size_t size; + CUmemGenericAllocationHandle handle; + int refcount; + foundry::vmm::CleanupState cleanup; +}; +// Generation-sensitive identity of an imported peer chunk. Defined as a pure, +// CUDA-free value type (foundry/vmm_ipc_chunk_key.h) so its equality/ordering +// contract - crucially that the generation is identity-bearing - is unit-tested +// on CPU (tests/vmm_ipc_chunk_key_test.cpp). +using VmmIpcChunkKey = foundry::vmm_ipc::ChunkKey; +struct VmmIpcChunkSubptr { + VmmIpcChunkKey key; + int open_count; +}; +struct PendingVmmIpcCleanup { + CUdeviceptr local_base; + size_t size; + CUmemGenericAllocationHandle handle; + foundry::vmm::CleanupState cleanup; + const char* context; +}; +static std::mutex vmm_ipc_chunk_mutex; +static std::map vmm_ipc_chunk_mappings; +// Interior pointer handed to a caller -> owning chunk key and number of opens. +static std::map vmm_ipc_chunk_subptrs; +// Failed cleanup of mappings that no caller can close (race losers and +// zero-ref collision mappings). Protected by vmm_ipc_chunk_mutex. A multimap +// preserves every resource even if concurrent VA reuse produces the same key. +static std::multimap vmm_ipc_pending_cleanups; + +static CUresult cleanup_import_mapping(foundry::vmm::CleanupState& cleanup, CUdeviceptr local_base, + size_t size, CUmemGenericAllocationHandle handle, + const char* context) { + typedef CUresult (*cuMemUnmap_t)(CUdeviceptr, size_t); + auto mem_unmap_func = + (cuMemUnmap_t)CUDA_DRIVER_CALL(cuda_driver_entry_table, CUDA_ENTRY_cuMemUnmap); + typedef CUresult (*cuMemRelease_t)(CUmemGenericAllocationHandle); + auto mem_release_func = + (cuMemRelease_t)CUDA_DRIVER_CALL(cuda_driver_entry_table, CUDA_ENTRY_cuMemRelease); + typedef CUresult (*cuMemAddressFree_t)(CUdeviceptr, size_t); + auto address_free_func = + (cuMemAddressFree_t)CUDA_DRIVER_CALL(cuda_driver_entry_table, CUDA_ENTRY_cuMemAddressFree); + + return foundry::vmm::run_cleanup( + cleanup, CUDA_SUCCESS, true, [&]() { return mem_unmap_func(local_base, size); }, + [&]() { return mem_release_func(handle); }, + [&]() { return address_free_func(local_base, size); }, + [&](foundry::vmm::CleanupPhase phase, CUresult status) { + const char* operation = "unknown"; + switch (phase) { + case foundry::vmm::CleanupPhase::kUnmap: + operation = "cuMemUnmap"; + break; + case foundry::vmm::CleanupPhase::kRelease: + operation = "cuMemRelease"; + break; + case foundry::vmm::CleanupPhase::kAddressFree: + operation = "cuMemAddressFree"; + break; + } + fprintf(stderr, "[HOOK] ERROR: VMM-IPC %s cleanup: %s failed with error %d\n", context, + operation, status); + }); +} + +// Requires vmm_ipc_chunk_mutex. +static void queue_pending_import_cleanup(CUdeviceptr local_base, size_t size, + CUmemGenericAllocationHandle handle, + const foundry::vmm::CleanupState& cleanup, + const char* context) { + vmm_ipc_pending_cleanups.emplace( + local_base, PendingVmmIpcCleanup{local_base, size, handle, cleanup, context}); +} + +// Requires vmm_ipc_chunk_mutex. Retry all orphans so one failure cannot starve +// a later entry; retain each incomplete resource and return the first error. +static CUresult drain_pending_import_cleanups_locked() { + return foundry::vmm::drain_pending_cleanups( + vmm_ipc_pending_cleanups, CUDA_SUCCESS, + [](PendingVmmIpcCleanup& pending) { + return cleanup_import_mapping(pending.cleanup, pending.local_base, pending.size, + pending.handle, pending.context); + }, + [](const PendingVmmIpcCleanup& pending) { + return foundry::vmm::cleanup_complete(pending.cleanup, true); + }); +} + +// Dedicated VA zone for relocated peer imports, far below the foundry region +// (default 0x600000000000) and the large-reserve zone right above region end. +// Without this, the driver tends to place a relocated import immediately +// after the region reservation - exactly where the hooked large-reserve hint +// sends the NVSHMEM symmetric heap next, which silently breaks the heap's +// SAVE/LOAD address determinism (observed: heap displaced by a peer-chunk +// import landing at region_end). +static std::atomic vmm_ipc_import_hint{0x300000000000ULL}; + +static socklen_t vmm_ipc_socket_addr(pid_t pid, struct sockaddr_un* addr) { + memset(addr, 0, sizeof(*addr)); + addr->sun_family = AF_UNIX; + // Abstract namespace (sun_path[0] == '\0'): no filesystem entry, vanishes + // with the process. + int n = snprintf(addr->sun_path + 1, sizeof(addr->sun_path) - 2, "foundry-vmm-ipc.%d", (int)pid); + return (socklen_t)(offsetof(struct sockaddr_un, sun_path) + 1 + n); +} + +static void vmm_ipc_server_loop(int listen_fd) { + while (true) { + int conn = accept(listen_fd, nullptr, nullptr); + if (conn < 0) { + if (errno == EBADF || errno == EINVAL) + break; // socket gone + continue; // EINTR/ECONNABORTED/EMFILE/... are transient + } + vmm_ipc_set_timeouts(conn); + // Abstract sockets have no filesystem permissions: only serve same-uid peers. + struct ucred cred; + socklen_t cred_len = sizeof(cred); + if (getsockopt(conn, SOL_SOCKET, SO_PEERCRED, &cred, &cred_len) != 0 || cred.uid != getuid()) { + close(conn); + continue; + } + VmmIpcRequest req = {}; + int fd = -1; + if (recv(conn, &req, sizeof(req), MSG_WAITALL) == (ssize_t)sizeof(req) && + req.token == vmm_ipc_token) { + // Dup under the map lock: the export path may concurrently close and + // replace the parked fd (stale-entry refresh); the dup we send is ours. + vmm_ipc_exported_fds.visit( + (CUdeviceptr)req.ptr, + [&](const std::pair>& + kv) { fd = fcntl(kv.second.second, F_DUPFD_CLOEXEC, 0); }); + } + char status = (fd >= 0) ? 0 : 1; + struct iovec iov = {&status, 1}; + char cbuf[CMSG_SPACE(sizeof(int))] = {}; + struct msghdr msg = {}; + msg.msg_iov = &iov; + msg.msg_iovlen = 1; + if (fd >= 0) { + msg.msg_control = cbuf; + msg.msg_controllen = CMSG_SPACE(sizeof(int)); + struct cmsghdr* c = CMSG_FIRSTHDR(&msg); + c->cmsg_level = SOL_SOCKET; + c->cmsg_type = SCM_RIGHTS; + c->cmsg_len = CMSG_LEN(sizeof(int)); + memcpy(CMSG_DATA(c), &fd, sizeof(int)); + } + // MSG_NOSIGNAL: a peer dying mid-handshake must not SIGPIPE-kill us. + if (sendmsg(conn, &msg, MSG_NOSIGNAL) != 1) { + fprintf(stderr, "[HOOK] WARNING: VMM-IPC fd server sendmsg failed (errno=%d)\n", errno); + } + if (fd >= 0) + close(fd); + close(conn); + } +} + +static bool vmm_ipc_ensure_server() { + // Claim process ownership before locking any VMM-IPC state. A post-claim fork + // child is rejected here so it never touches inherited mutexes/maps/handles. + if (!vmm_ipc_claim_process()) + return false; + std::lock_guard lock(vmm_ipc_server_mutex); + pid_t pid = getpid(); + if (vmm_ipc_server_pid == pid) { + return vmm_ipc_listen_fd >= 0; + } + // First call in this process. VMM-IPC after fork is rejected before this + // mutex (vmm_ipc_claim_process above), so inherited mutex/map state and CUDA + // handles are never locked or mutated by a child. Close any inherited listen + // fd so a parent's abstract name is not pinned alive by us after it exits. + if (vmm_ipc_listen_fd >= 0) { + close(vmm_ipc_listen_fd); + vmm_ipc_listen_fd = -1; + } + int s = socket(AF_UNIX, SOCK_STREAM | SOCK_CLOEXEC, 0); + struct sockaddr_un addr; + socklen_t len = vmm_ipc_socket_addr(pid, &addr); + if (s < 0 || bind(s, (struct sockaddr*)&addr, len) != 0 || listen(s, 64) != 0) { + fprintf(stderr, "[HOOK] ERROR: VMM-IPC fd server setup failed (errno=%d)\n", errno); + if (s >= 0) + close(s); + vmm_ipc_listen_fd = -1; + vmm_ipc_server_pid = pid; // don't retry-spam; exports in this process fail loudly + return false; + } + // Per-process token (re-derived after fork). /dev/urandom, with a clock^pid + // fallback - this is anti-accident (pid reuse), not a security boundary; + // same-uid access control is SO_PEERCRED above. + uint64_t tok = 0; + int ur = open("/dev/urandom", O_RDONLY | O_CLOEXEC); + if (ur >= 0) { + if (read(ur, &tok, sizeof(tok)) != (ssize_t)sizeof(tok)) + tok = 0; + close(ur); + } + if (tok == 0) { + struct timeval tv; + gettimeofday(&tv, nullptr); + tok = ((uint64_t)tv.tv_sec << 32) ^ (uint64_t)tv.tv_usec ^ ((uint64_t)pid << 16) ^ + 0x9e3779b97f4a7c15ULL; + } + vmm_ipc_token = tok; + vmm_ipc_listen_fd = s; + vmm_ipc_server_pid = pid; + std::thread(vmm_ipc_server_loop, s).detach(); + return true; +} + +// Importer side: fetch the fd for `original_ptr` from `exporter_pid`'s server. +// Returns a live fd in THIS process, or -1. +static int vmm_ipc_fetch_fd(pid_t exporter_pid, uint64_t token, CUdeviceptr original_ptr) { + int s = socket(AF_UNIX, SOCK_STREAM | SOCK_CLOEXEC, 0); + if (s < 0) + return -1; + vmm_ipc_set_timeouts(s); + struct sockaddr_un addr; + socklen_t len = vmm_ipc_socket_addr(exporter_pid, &addr); + if (connect(s, (struct sockaddr*)&addr, len) != 0) { + fprintf(stderr, + "[HOOK] ERROR: VMM-IPC could not connect to live exporter server pid %d (errno=%d) - " + "exporter dead or hook version mismatch\n", + (int)exporter_pid, errno); + close(s); + return -1; + } + VmmIpcRequest req = {(uint64_t)original_ptr, token}; + if (send(s, &req, sizeof(req), MSG_NOSIGNAL) != (ssize_t)sizeof(req)) { + close(s); + return -1; + } + char status = 1; + struct iovec iov = {&status, 1}; + char cbuf[CMSG_SPACE(sizeof(int))] = {}; + struct msghdr msg = {}; + msg.msg_iov = &iov; + msg.msg_iovlen = 1; + msg.msg_control = cbuf; + msg.msg_controllen = sizeof(cbuf); + ssize_t r = recvmsg(s, &msg, MSG_CMSG_CLOEXEC); + close(s); + if (r != 1) { + fprintf(stderr, + "[HOOK] ERROR: VMM-IPC could not read a response from live exporter server pid %d " + "(recvmsg=%zd, errno=%d)\n", + (int)exporter_pid, r, errno); + return -1; + } + if (status != 0) { + fprintf(stderr, + "[HOOK] ERROR: VMM-IPC live exporter server rejected stale/unknown export " + "ptr=0x%llx from pid %d (status=%d)\n", + (unsigned long long)original_ptr, (int)exporter_pid, (int)status); + return -1; + } + for (struct cmsghdr* c = CMSG_FIRSTHDR(&msg); c != nullptr; c = CMSG_NXTHDR(&msg, c)) { + if (c->cmsg_level == SOL_SOCKET && c->cmsg_type == SCM_RIGHTS) { + int fd = -1; + memcpy(&fd, CMSG_DATA(c), sizeof(int)); + return fd; + } + } + return -1; +} + +// Close and forget the parked exported fd for a VA (called from the free +// path). Without this, the fd keeps the freed allocation's physical pages +// alive and the server would hand importers a stale mapping. +static void vmm_ipc_invalidate_export(CUdeviceptr dptr) { + // Only the owning process has parked export fds. An unclaimed process (-1) has + // never exported, and a fork child must not lock/mutate the inherited + // concurrent map; both return before any map access. + if (!foundry::vmm_ipc::owns_ipc_state((long)vmm_ipc_owner_pid.load(std::memory_order_acquire), + (long)getpid())) + return; + vmm_ipc_exported_fds.erase_if( + [dptr](const std::pair>& kv) { + if (kv.first != dptr) + return false; + close(kv.second.second); + return true; + }); +} // Hook for cuIpcGetMemHandle - intercept Driver API to support VMM allocations CUresult cuIpcGetMemHandle(CUipcMemHandle* pHandle, CUdeviceptr dptr) { + // Claim/reject as the first executable action, before even the driver symbol + // lookup (CUDA_DRIVER_CALL caches via a lock the fork child must not touch) + // and before any allocation-metadata or server state. Exporting is a VMM-IPC + // use, so the exporter claims ownership; a post-claim fork child is rejected. + // Applied uniformly to plain and VMM handles (conservative all-IPC claim) for + // consistent fork behavior. + if (!vmm_ipc_claim_process()) + return CUDA_ERROR_NOT_SUPPORTED; + typedef CUresult (*cuIpcGetMemHandle_t)(CUipcMemHandle*, CUdeviceptr); auto real_func = (cuIpcGetMemHandle_t)CUDA_DRIVER_CALL(cuda_driver_entry_table, CUDA_ENTRY_cuIpcGetMemHandle); @@ -2878,7 +3376,66 @@ CUresult cuIpcGetMemHandle(CUipcMemHandle* pHandle, CUdeviceptr dptr) { metadata = kv.second; }); - if (found && metadata.handle != 0) { + // Decide what to export: the allocation's own handle (SAVE-mode slow-path + // allocs), or the whole preallocated chunk plus an offset (LOAD-mode fast + // path carves, which have no individual handle). cuMemMap cannot map a + // sub-range of an imported handle, so chunk carves are shared by exporting + // the chunk handle; the importer maps the entire chunk and returns an + // interior pointer. + CUmemGenericAllocationHandle export_handle = 0; + CUdeviceptr export_key = dptr; // fd-registry key == blob lookup key + uint64_t chunk_base = 0; + uint64_t chunk_size = 0; + uint64_t chunk_generation = 0; + bool preallocated_chunk_cleanup_started = false; + std::unique_lock preallocated_chunk_lock; + if (found) { + if (metadata.cleanup.started) { + fprintf(stderr, + "[HOOK] ERROR: cuIpcGetMemHandle rejected ptr=0x%llx because allocation cleanup is " + "pending\n", + (unsigned long long)dptr); + return CUDA_ERROR_INVALID_VALUE; + } + export_handle = metadata.handle; + if (export_handle == 0 && metadata.from_preallocation) { + // Keep the chunk alive and its record stable until the exported fd is + // parked below. free_preallocated_region takes the same mutex. + preallocated_chunk_lock = std::unique_lock(preallocated_chunks_mutex); + auto chunk_it = preallocated_chunks.upper_bound(dptr); + if (chunk_it != preallocated_chunks.begin()) { + --chunk_it; + CUdeviceptr candidate_base = chunk_it->first; + if (dptr >= candidate_base && (uint64_t)(dptr - candidate_base) < chunk_it->second.size) { + if (chunk_it->second.cleanup.started) { + preallocated_chunk_cleanup_started = true; + } else { + chunk_base = (uint64_t)candidate_base; + chunk_size = chunk_it->second.size; + chunk_generation = chunk_it->second.generation; + export_handle = chunk_it->second.handle; + export_key = candidate_base; + } + } + } + if (preallocated_chunk_cleanup_started) { + fprintf(stderr, + "[HOOK] ERROR: cuIpcGetMemHandle rejected carved ptr=0x%llx because " + "preallocated chunk cleanup is pending\n", + (unsigned long long)dptr); + return CUDA_ERROR_INVALID_VALUE; + } + if (export_handle == 0) { + fprintf(stderr, + "[HOOK] ERROR: cuIpcGetMemHandle: carved ptr=0x%llx is outside all live " + "preallocated chunks\n", + (unsigned long long)dptr); + return CUDA_ERROR_INVALID_VALUE; + } + } + } + + if (export_handle != 0) { // This is a VMM allocation - export via shareable handle typedef CUresult (*cuMemExportToShareableHandle_t)( void*, CUmemGenericAllocationHandle, CUmemAllocationHandleType, unsigned long long); @@ -2890,32 +3447,82 @@ CUresult cuIpcGetMemHandle(CUipcMemHandle* pHandle, CUdeviceptr dptr) { return CUDA_ERROR_NOT_SUPPORTED; } - int fd = -1; - CUresult result = - export_func(&fd, metadata.handle, CU_MEM_HANDLE_TYPE_POSIX_FILE_DESCRIPTOR, 0); + // The server also owns the per-process token packed into the blob, so it + // must exist before we hand out any handle. + if (!vmm_ipc_ensure_server()) { + return CUDA_ERROR_NOT_SUPPORTED; + } - if (result != CUDA_SUCCESS) { - fprintf(stderr, - "[HOOK] ERROR: cuMemExportToShareableHandle failed with error %d for ptr=0x%llx\n", - result, (unsigned long long)dptr); - return result; + // Export once per allocation and park the fd for the server thread. + // The cached entry is keyed by VA and validated against the allocation + // handle, so VA reuse after free/realloc re-exports instead of serving a + // stale fd (the free path also eagerly invalidates via + // vmm_ipc_invalidate_export). + int fd = -1; + vmm_ipc_exported_fds.visit( + export_key, + [&](const std::pair>& kv) { + if (kv.second.first == export_handle) + fd = kv.second.second; + }); + if (fd < 0) { + int new_fd = -1; + CUresult result = + export_func(&new_fd, export_handle, CU_MEM_HANDLE_TYPE_POSIX_FILE_DESCRIPTOR, 0); + if (result != CUDA_SUCCESS) { + fprintf(stderr, + "[HOOK] ERROR: cuMemExportToShareableHandle failed with error %d for ptr=0x%llx\n", + result, (unsigned long long)dptr); + return result; + } + // Keep parked fds out of exec'd children. A plain fork child that tries + // to use VMM-IPC is now rejected up front (vmm_ipc_claim_process), so it + // never reaches this export path or serves the parent's inherited fds. + // SCM_RIGHTS transfer is unaffected. + fcntl(new_fd, F_SETFD, FD_CLOEXEC); + vmm_ipc_exported_fds.insert_or_visit( + std::make_pair(export_key, std::make_pair(export_handle, new_fd)), + [&](std::pair>& kv) { + // Entry exists: either a racing thread won (same handle - drop + // ours) or it is stale from a freed allocation (replace). + if (kv.second.first == export_handle) { + close(new_fd); + new_fd = kv.second.second; + } else { + close(kv.second.second); + kv.second = std::make_pair(export_handle, new_fd); + } + }); + fd = new_fd; } - // Pack our custom data into the handle structure - // CUipcMemHandle has 64 reserved bytes - we use them to store our info: - // - Magic marker (4 bytes) to identify VMM handles - // - File descriptor (4 bytes) - // - Original pointer address (8 bytes) - critical for CUDA graph replay! - // - Size (8 bytes) + // Pack our custom data into the handle structure. + // CUipcMemHandle has 64 reserved bytes: + // - Magic marker (4 bytes, "VMM2") + // - Exporter pid (4 bytes) - importer fetches the fd from this process's + // VMM-IPC socket via SCM_RIGHTS (a raw fd integer is not portable) + // - Original pointer address (8 bytes) - fd lookup key + placement hint + // - Aligned size (8 bytes) + // - Per-process token (8 bytes) - server rejects mismatches (pid reuse) + // - Chunk base + chunk size (8+8 bytes) - nonzero iff the pointer is a + // carve from the preallocated chunk; the importer then maps the whole + // chunk and returns base + (ptr - chunk_base) + // - Chunk generation (8 bytes) - distinguishes free/recreate at one base memset(pHandle, 0, sizeof(CUipcMemHandle)); + uint32_t pid_u32 = (uint32_t)getpid(); memcpy(pHandle->reserved, &VMM_IPC_MAGIC, sizeof(uint32_t)); - memcpy(pHandle->reserved + 4, &fd, sizeof(int)); + memcpy(pHandle->reserved + 4, &pid_u32, sizeof(uint32_t)); memcpy(pHandle->reserved + 8, &dptr, sizeof(CUdeviceptr)); memcpy(pHandle->reserved + 16, &metadata.size, sizeof(size_t)); + memcpy(pHandle->reserved + 24, &vmm_ipc_token, sizeof(uint64_t)); + memcpy(pHandle->reserved + 32, &chunk_base, sizeof(uint64_t)); + memcpy(pHandle->reserved + 40, &chunk_size, sizeof(uint64_t)); + memcpy(pHandle->reserved + 48, &chunk_generation, sizeof(uint64_t)); #ifdef HOOK_DEBUG - fprintf(stderr, "[HOOK] cuIpcGetMemHandle: VMM ptr=0x%llx, fd=%d, size=%zu\n", + fprintf(stderr, + "[HOOK] cuIpcGetMemHandle: VMM ptr=0x%llx, fd=%d (served via socket), size=%zu\n", (unsigned long long)dptr, fd, metadata.size); #endif return CUDA_SUCCESS; @@ -2931,6 +3538,23 @@ CUresult cuIpcGetMemHandle(CUipcMemHandle* pHandle, CUdeviceptr dptr) { // Hook for cuIpcOpenMemHandle - import VMM shareable handle and map to original address CUresult cuIpcOpenMemHandle(CUdeviceptr* pdptr, CUipcMemHandle handle, unsigned int Flags) { typedef CUresult (*cuIpcOpenMemHandle_t)(CUdeviceptr*, CUipcMemHandle, unsigned int); + // Claim/reject as the first executable action, before even the driver symbol + // lookup (CUDA_DRIVER_CALL caches via a lock the fork child must not touch), + // before the magic check, and before any chunk-mapping shared state. + // Importing is a VMM-IPC use, so the importer claims ownership; a post-claim + // fork child is rejected. Applied uniformly to plain and VMM handles + // (conservative all-IPC claim) for consistent fork behavior. + if (!vmm_ipc_claim_process()) + return CUDA_ERROR_NOT_SUPPORTED; + + { + std::lock_guard lock(vmm_ipc_chunk_mutex); + CUresult pending_cleanup_status = drain_pending_import_cleanups_locked(); + if (pending_cleanup_status != CUDA_SUCCESS) { + return pending_cleanup_status; + } + } + auto real_func = (cuIpcOpenMemHandle_t)CUDA_DRIVER_CALL(cuda_driver_entry_table, CUDA_ENTRY_cuIpcOpenMemHandle); @@ -2940,19 +3564,83 @@ CUresult cuIpcOpenMemHandle(CUdeviceptr* pdptr, CUipcMemHandle handle, unsigned if (magic == VMM_IPC_MAGIC) { // This is a VMM IPC handle - extract the packed data - int fd; + uint32_t exporter_pid_u32; CUdeviceptr original_ptr; size_t size; + uint64_t token; + uint64_t chunk_base; + uint64_t chunk_size; + uint64_t chunk_generation; - memcpy(&fd, handle.reserved + 4, sizeof(int)); + memcpy(&exporter_pid_u32, handle.reserved + 4, sizeof(uint32_t)); memcpy(&original_ptr, handle.reserved + 8, sizeof(CUdeviceptr)); memcpy(&size, handle.reserved + 16, sizeof(size_t)); + memcpy(&token, handle.reserved + 24, sizeof(uint64_t)); + memcpy(&chunk_base, handle.reserved + 32, sizeof(uint64_t)); + memcpy(&chunk_size, handle.reserved + 40, sizeof(uint64_t)); + memcpy(&chunk_generation, handle.reserved + 48, sizeof(uint64_t)); + pid_t exporter_pid = (pid_t)exporter_pid_u32; + // For chunk carves the fd registry on the exporter is keyed by the chunk + // base, and what we import/map is the whole chunk. + CUdeviceptr fetch_key = (chunk_base != 0) ? (CUdeviceptr)chunk_base : original_ptr; + size_t map_size = (chunk_base != 0) ? (size_t)chunk_size : size; + + if (chunk_base != 0) { + // Fast path: peer chunk already mapped -> hand out an interior pointer. + std::lock_guard lock(vmm_ipc_chunk_mutex); + VmmIpcChunkKey key = + foundry::vmm_ipc::make_chunk_key(exporter_pid, token, chunk_base, chunk_generation); + auto it = vmm_ipc_chunk_mappings.find(key); + if (it != vmm_ipc_chunk_mappings.end()) { + if (it->second.cleanup.started) { + fprintf(stderr, + "[HOOK] ERROR: VMM-IPC chunk open rejected while cleanup is pending for " + "local_base=0x%llx\n", + (unsigned long long)it->second.local_base); + return CUDA_ERROR_INVALID_VALUE; + } + it->second.refcount++; + CUdeviceptr mapped = it->second.local_base + (original_ptr - (CUdeviceptr)chunk_base); + auto sub_it = vmm_ipc_chunk_subptrs.find(mapped); + if (sub_it == vmm_ipc_chunk_subptrs.end()) { + vmm_ipc_chunk_subptrs[mapped] = VmmIpcChunkSubptr{key, 1}; + } else if (sub_it->second.key == key) { + sub_it->second.open_count++; + } else { + it->second.refcount--; + fprintf(stderr, + "[HOOK] ERROR: VMM-IPC interior pointer collision at 0x%llx across chunks\n", + (unsigned long long)mapped); + return CUDA_ERROR_INVALID_VALUE; + } + *pdptr = mapped; + return CUDA_SUCCESS; + } + } #ifdef HOOK_DEBUG - fprintf(stderr, "[HOOK] cuIpcOpenMemHandle: VMM fd=%d, original_ptr=0x%llx, size=%zu\n", fd, - (unsigned long long)original_ptr, size); + fprintf(stderr, + "[HOOK] cuIpcOpenMemHandle: VMM exporter_pid=%d, original_ptr=0x%llx, size=%zu\n", + (int)exporter_pid, (unsigned long long)original_ptr, size); #endif + // Obtain a live fd in THIS process: dup from our own registry for + // same-process opens, SCM_RIGHTS fetch from the exporter otherwise. + int fd = -1; + if (exporter_pid == getpid() && token == vmm_ipc_token) { + vmm_ipc_exported_fds.visit( + fetch_key, + [&](const std::pair>& + kv) { fd = fcntl(kv.second.second, F_DUPFD_CLOEXEC, 0); }); + } else { + fd = vmm_ipc_fetch_fd(exporter_pid, token, fetch_key); + } + if (fd < 0) { + fprintf(stderr, "[HOOK] ERROR: VMM-IPC could not obtain fd for ptr=0x%llx from pid %d\n", + (unsigned long long)fetch_key, (int)exporter_pid); + return CUDA_ERROR_INVALID_VALUE; + } + // Import the allocation handle from file descriptor typedef CUresult (*cuMemImportFromShareableHandle_t)(CUmemGenericAllocationHandle*, void*, CUmemAllocationHandleType); @@ -2961,12 +3649,14 @@ CUresult cuIpcOpenMemHandle(CUdeviceptr* pdptr, CUipcMemHandle handle, unsigned if (import_func == nullptr) { fprintf(stderr, "[HOOK] ERROR: cuMemImportFromShareableHandle not found\n"); + close(fd); return CUDA_ERROR_NOT_SUPPORTED; } CUmemGenericAllocationHandle imported_handle; CUresult result = import_func(&imported_handle, (void*)(intptr_t)fd, CU_MEM_HANDLE_TYPE_POSIX_FILE_DESCRIPTOR); + close(fd); // the driver does not take ownership of our fd copy if (result != CUDA_SUCCESS) { fprintf(stderr, "[HOOK] ERROR: cuMemImportFromShareableHandle failed with error %d\n", @@ -2974,16 +3664,106 @@ CUresult cuIpcOpenMemHandle(CUdeviceptr* pdptr, CUipcMemHandle handle, unsigned return result; } - // Map to the SAME address as original (critical for CUDA graph replay!) + typedef CUresult (*cuMemAddressReserve_t)(CUdeviceptr*, size_t, size_t, CUdeviceptr, + unsigned long long); + auto real_reserve_func = (cuMemAddressReserve_t)CUDA_DRIVER_CALL( + cuda_driver_entry_table, CUDA_ENTRY_cuMemAddressReserve); + typedef CUresult (*cuMemRelease_t)(CUmemGenericAllocationHandle); + auto mem_release_func = + (cuMemRelease_t)CUDA_DRIVER_CALL(cuda_driver_entry_table, CUDA_ENTRY_cuMemRelease); + typedef CUresult (*cuMemAddressFree_t)(CUdeviceptr, size_t); + auto real_address_free_func = + (cuMemAddressFree_t)CUDA_DRIVER_CALL(cuda_driver_entry_table, CUDA_ENTRY_cuMemAddressFree); + + // Reserve our own VA range for the peer mapping (real driver call, NOT the + // hooked wrapper - peer imports must never advance the deterministic + // cursor). The exporter's VA is passed as a hint: with per-rank disjoint + // regions it is honored and the mapping is address-stable; with today's + // shared-base symmetric layouts that range is occupied by our own region + // reservation, the driver picks elsewhere, and the caller gets a valid but + // relocated peer pointer. That is correct for table-indirect consumers + // (DeepEP buffer_ptrs_gpu, custom-allreduce RankData contents): peer VAs + // are never baked into captured graph nodes on those paths. + // First preference: the exporter's own VA (address-stable whenever it is + // locally free, e.g. future per-rank disjoint-base configs). + CUdeviceptr mapped_ptr = 0; + result = real_reserve_func(&mapped_ptr, map_size, 0, fetch_key, 0); + if (result == CUDA_SUCCESS && mapped_ptr != fetch_key) { + // Hint not honored. Do NOT keep the driver's fallback choice - it tends + // to sit immediately after existing reservations, i.e. on the NVSHMEM + // heap hint. Free it, then re-reserve inside the dedicated import zone. + // This runs in the owning importer process. If the free FAILS we cannot + // trust or reclaim the fallback VA, so fail closed immediately: release + // the imported handle (once) and return without attempting the zone + // mapping. Both cleanup statuses are logged. + CUresult fallback_free_status = real_address_free_func(mapped_ptr, map_size); + CUresult fallback_status = foundry::vmm_ipc::require_fallback_reservation_freed( + fallback_free_status, imported_handle, [&](CUmemGenericAllocationHandle h) -> CUresult { + CUresult rc = mem_release_func(h); + fprintf(stderr, + "[HOOK] ERROR: VMM-IPC import: cuMemAddressFree(fallback reservation 0x%llx) " + "failed with error %d; released imported handle (cuMemRelease status %d); " + "failing closed without a dedicated-zone mapping\n", + (unsigned long long)mapped_ptr, fallback_free_status, rc); + return rc; + }); + if (fallback_status != CUDA_SUCCESS) { + return fallback_status; + } + uint64_t zone_hint = + vmm_ipc_import_hint.fetch_add((map_size + ((1ULL << 30) - 1)) & ~((1ULL << 30) - 1)); + mapped_ptr = 0; + result = real_reserve_func(&mapped_ptr, map_size, 0, (CUdeviceptr)zone_hint, 0); + // Strict dedicated zone: a reservation at any other VA can still land on + // the deterministic/NVSHMEM reserve space this zone exists to avoid. Fail + // closed (free the reservation + release the imported handle) unless the + // hint was honored exactly. A driver error is passed through and handled + // by the shared failure path below (which releases the handle). + CUresult zone_status = foundry::vmm_ipc::require_exact_ipc_reservation( + result, mapped_ptr, (CUdeviceptr)zone_hint, map_size, imported_handle, + [&](CUdeviceptr p, size_t s) -> CUresult { + CUresult rc = real_address_free_func(p, s); + if (rc != CUDA_SUCCESS) + fprintf(stderr, + "[HOOK] ERROR: VMM-IPC wrong-VA cleanup: cuMemAddressFree(0x%llx) failed " + "with error %d\n", + (unsigned long long)p, rc); + return rc; + }, + [&](CUmemGenericAllocationHandle h) -> CUresult { + CUresult rc = mem_release_func(h); + if (rc != CUDA_SUCCESS) + fprintf(stderr, + "[HOOK] ERROR: VMM-IPC wrong-VA cleanup: cuMemRelease failed with error %d\n", + rc); + return rc; + }); + if (result == CUDA_SUCCESS && zone_status != CUDA_SUCCESS) { + fprintf(stderr, + "[HOOK] ERROR: VMM-IPC import zone reservation returned 0x%llx instead of " + "required 0x%llx\n", + (unsigned long long)mapped_ptr, (unsigned long long)zone_hint); + return zone_status; + } + } + if (result != CUDA_SUCCESS) { + fprintf(stderr, "[HOOK] ERROR: cuMemAddressReserve for IPC import failed with error %d\n", + result); + mem_release_func(imported_handle); + return result; + } + typedef CUresult (*cuMemMap_t)(CUdeviceptr, size_t, size_t, CUmemGenericAllocationHandle, unsigned long long); auto map_func = (cuMemMap_t)CUDA_DRIVER_CALL(cuda_driver_entry_table, CUDA_ENTRY_cuMemMap); - result = map_func(original_ptr, size, 0, imported_handle, 0); + result = map_func(mapped_ptr, map_size, 0, imported_handle, 0); if (result != CUDA_SUCCESS) { fprintf(stderr, "[HOOK] ERROR: cuMemMap for IPC failed with error %d at addr=0x%llx size=%zu\n", - result, (unsigned long long)original_ptr, size); + result, (unsigned long long)mapped_ptr, map_size); + real_address_free_func(mapped_ptr, map_size); + mem_release_func(imported_handle); return result; } @@ -3003,27 +3783,106 @@ CUresult cuIpcOpenMemHandle(CUdeviceptr* pdptr, CUipcMemHandle handle, unsigned auto set_access_func = (cuMemSetAccess_t)CUDA_DRIVER_CALL(cuda_driver_entry_table, CUDA_ENTRY_cuMemSetAccess); - result = set_access_func(original_ptr, size, &accessDesc, 1); + result = set_access_func(mapped_ptr, map_size, &accessDesc, 1); if (result != CUDA_SUCCESS) { fprintf(stderr, "[HOOK] ERROR: cuMemSetAccess for IPC failed with error %d\n", result); + typedef CUresult (*cuMemUnmap_t)(CUdeviceptr, size_t); + auto mem_unmap_func = + (cuMemUnmap_t)CUDA_DRIVER_CALL(cuda_driver_entry_table, CUDA_ENTRY_cuMemUnmap); + mem_unmap_func(mapped_ptr, map_size); + real_address_free_func(mapped_ptr, map_size); + mem_release_func(imported_handle); return result; } - *pdptr = original_ptr; + if (mapped_ptr != fetch_key) { + fprintf(stderr, + "[HOOK] INFO: VMM-IPC import relocated: exporter VA 0x%llx -> local VA 0x%llx " + "(expected with shared region bases; peer tables must be refreshed by the caller)\n", + (unsigned long long)fetch_key, (unsigned long long)mapped_ptr); + } - // Track this imported allocation + if (chunk_base != 0) { + // Register the whole-chunk mapping and hand out the interior pointer. + CUdeviceptr interior = mapped_ptr + (original_ptr - (CUdeviceptr)chunk_base); + std::lock_guard lock(vmm_ipc_chunk_mutex); + VmmIpcChunkKey key = + foundry::vmm_ipc::make_chunk_key(exporter_pid, token, chunk_base, chunk_generation); + auto it = vmm_ipc_chunk_mappings.find(key); + if (it != vmm_ipc_chunk_mappings.end()) { + // Lost a race to a concurrent open of the same peer chunk. Always tear + // down our mapping with checked cleanup before considering the winner. + // A winner already being cleaned cannot accept another reference. + bool winner_cleanup_started = it->second.cleanup.started; + foundry::vmm::CleanupState loser_cleanup; + CUresult loser_cleanup_status = cleanup_import_mapping(loser_cleanup, mapped_ptr, map_size, + imported_handle, "chunk race loser"); + if (!foundry::vmm::cleanup_complete(loser_cleanup, true)) { + queue_pending_import_cleanup(mapped_ptr, map_size, imported_handle, loser_cleanup, + "pending chunk race loser"); + return loser_cleanup_status; + } + if (winner_cleanup_started) { + fprintf(stderr, + "[HOOK] ERROR: VMM-IPC chunk race winner cleanup is pending for " + "local_base=0x%llx; loser was cleaned and open is rejected\n", + (unsigned long long)it->second.local_base); + return CUDA_ERROR_INVALID_VALUE; + } + it->second.refcount++; + interior = it->second.local_base + (original_ptr - (CUdeviceptr)chunk_base); + } else { + vmm_ipc_chunk_mappings[key] = VmmIpcChunkMapping{mapped_ptr, map_size, imported_handle, 1}; + } + auto sub_it = vmm_ipc_chunk_subptrs.find(interior); + if (sub_it == vmm_ipc_chunk_subptrs.end()) { + vmm_ipc_chunk_subptrs[interior] = VmmIpcChunkSubptr{key, 1}; + } else if (sub_it->second.key == key) { + sub_it->second.open_count++; + } else { + auto mapping_it = vmm_ipc_chunk_mappings.find(key); + if (mapping_it != vmm_ipc_chunk_mappings.end() && --mapping_it->second.refcount == 0) { + CUresult collision_cleanup_status = cleanup_import_mapping( + mapping_it->second.cleanup, mapping_it->second.local_base, mapping_it->second.size, + mapping_it->second.handle, "interior collision"); + if (foundry::vmm::cleanup_complete(mapping_it->second.cleanup, true)) { + vmm_ipc_chunk_mappings.erase(mapping_it); + } else { + queue_pending_import_cleanup(mapping_it->second.local_base, mapping_it->second.size, + mapping_it->second.handle, mapping_it->second.cleanup, + "pending interior collision"); + // This zero-ref mapping has no subptr through which cleanup could + // be retried. Its full resource state now lives in the pending map. + vmm_ipc_chunk_mappings.erase(mapping_it); + return collision_cleanup_status; + } + } + fprintf(stderr, + "[HOOK] ERROR: VMM-IPC interior pointer collision at 0x%llx across chunks\n", + (unsigned long long)interior); + return CUDA_ERROR_INVALID_VALUE; + } + *pdptr = interior; + return CUDA_SUCCESS; + } + + *pdptr = mapped_ptr; + + // Track this imported allocation (close path unmaps, releases, and frees + // the reservation we created above) AllocMetadata alloc_metadata; - alloc_metadata.ptr = original_ptr; + alloc_metadata.ptr = mapped_ptr; alloc_metadata.size = size; alloc_metadata.handle = imported_handle; alloc_metadata.region_base = 0; // region_base == 0 indicates IPC-imported alloc_metadata.from_preallocation = false; - global_alloc_metadata.emplace(original_ptr, alloc_metadata); + global_alloc_metadata.emplace(mapped_ptr, alloc_metadata); #ifdef HOOK_DEBUG fprintf(stderr, - "[HOOK] cuIpcOpenMemHandle: Successfully mapped VMM fd=%d -> ptr=0x%llx, size=%zu\n", - fd, (unsigned long long)*pdptr, size); + "[HOOK] cuIpcOpenMemHandle: Successfully mapped exporter ptr=0x%llx -> ptr=0x%llx, " + "size=%zu\n", + (unsigned long long)original_ptr, (unsigned long long)*pdptr, size); #endif return CUDA_SUCCESS; } @@ -3038,9 +3897,88 @@ CUresult cuIpcOpenMemHandle(CUdeviceptr* pdptr, CUipcMemHandle handle, unsigned // Hook for cuIpcCloseMemHandle - clean up VMM IPC mappings CUresult cuIpcCloseMemHandle(CUdeviceptr dptr) { typedef CUresult (*cuIpcCloseMemHandle_t)(CUdeviceptr); + // Claim/reject as the first executable action, before even the driver symbol + // lookup (CUDA_DRIVER_CALL caches via a lock the fork child must not touch) + // and before touching the chunk-mapping maps or allocation metadata. A + // post-claim fork child is rejected; the owning process (which claimed on + // open) proceeds. Applied uniformly to plain and VMM handles for consistent + // fork behavior. + if (!vmm_ipc_claim_process()) + return CUDA_ERROR_NOT_SUPPORTED; + auto real_func = (cuIpcCloseMemHandle_t)CUDA_DRIVER_CALL(cuda_driver_entry_table, CUDA_ENTRY_cuIpcCloseMemHandle); + // Interior pointer into an imported peer chunk? Refcounted: the chunk + // mapping is torn down when its last interior pointer closes. + { + std::lock_guard lock(vmm_ipc_chunk_mutex); + auto sub_it = vmm_ipc_chunk_subptrs.find(dptr); + if (sub_it != vmm_ipc_chunk_subptrs.end()) { + VmmIpcChunkKey key = sub_it->second.key; + auto it = vmm_ipc_chunk_mappings.find(key); + if (it == vmm_ipc_chunk_mappings.end()) { + vmm_ipc_chunk_subptrs.erase(sub_it); + return CUDA_SUCCESS; + } + + // A failed final close leaves both counts at zero and retains the subptr + // plus mapping. A retry must resume cleanup without decrementing again. + if (sub_it->second.open_count > 0) { + --sub_it->second.open_count; + if (it->second.refcount > 0) { + --it->second.refcount; + } + } + if (it->second.refcount > 0) { + if (sub_it->second.open_count == 0) { + vmm_ipc_chunk_subptrs.erase(sub_it); + } + return CUDA_SUCCESS; + } + + typedef CUresult (*cuMemUnmap_t)(CUdeviceptr, size_t); + auto mem_unmap_func = + (cuMemUnmap_t)CUDA_DRIVER_CALL(cuda_driver_entry_table, CUDA_ENTRY_cuMemUnmap); + typedef CUresult (*cuMemRelease_t)(CUmemGenericAllocationHandle); + auto mem_release_func = + (cuMemRelease_t)CUDA_DRIVER_CALL(cuda_driver_entry_table, CUDA_ENTRY_cuMemRelease); + typedef CUresult (*cuMemAddressFree_t)(CUdeviceptr, size_t); + auto addr_free_func = (cuMemAddressFree_t)CUDA_DRIVER_CALL(cuda_driver_entry_table, + CUDA_ENTRY_cuMemAddressFree); + CUresult cleanup_status = foundry::vmm::run_cleanup( + it->second.cleanup, CUDA_SUCCESS, true, + [&]() { return mem_unmap_func(it->second.local_base, it->second.size); }, + [&]() { return mem_release_func(it->second.handle); }, + [&]() { return addr_free_func(it->second.local_base, it->second.size); }, + [&](foundry::vmm::CleanupPhase phase, CUresult status) { + switch (phase) { + case foundry::vmm::CleanupPhase::kUnmap: + fprintf(stderr, + "[HOOK] ERROR: VMM-IPC chunk close: cuMemUnmap failed with error %d\n", + status); + break; + case foundry::vmm::CleanupPhase::kRelease: + fprintf(stderr, + "[HOOK] ERROR: VMM-IPC chunk close: cuMemRelease failed with error %d\n", + status); + break; + case foundry::vmm::CleanupPhase::kAddressFree: + fprintf( + stderr, + "[HOOK] ERROR: VMM-IPC chunk close: cuMemAddressFree failed with error %d\n", + status); + break; + } + }); + if (foundry::vmm::cleanup_complete(it->second.cleanup, true)) { + vmm_ipc_chunk_subptrs.erase(sub_it); + vmm_ipc_chunk_mappings.erase(it); + } + return cleanup_status; + } + } + AllocMetadata metadata; bool found = false; @@ -3059,17 +3997,55 @@ CUresult cuIpcCloseMemHandle(CUdeviceptr dptr) { auto mem_release_func = (cuMemRelease_t)CUDA_DRIVER_CALL(cuda_driver_entry_table, CUDA_ENTRY_cuMemRelease); - mem_unmap_func(dptr, metadata.size); - mem_release_func(metadata.handle); + typedef CUresult (*cuMemAddressFree_t)(CUdeviceptr, size_t); + auto real_address_free_func = + (cuMemAddressFree_t)CUDA_DRIVER_CALL(cuda_driver_entry_table, CUDA_ENTRY_cuMemAddressFree); - global_alloc_metadata.erase_if( - [dptr](const std::pair& kv) { return kv.first == dptr; }); + CUresult cleanup_status = CUDA_SUCCESS; + bool individual_cleanup_complete = false; + global_alloc_metadata.visit(dptr, [&](std::pair& kv) { + cleanup_status = foundry::vmm::run_cleanup( + kv.second.cleanup, CUDA_SUCCESS, true, + [&]() { return mem_unmap_func(dptr, kv.second.size); }, + [&]() { return mem_release_func(kv.second.handle); }, + // The import path owns a dedicated VA reservation for this mapping. + [&]() { return real_address_free_func(dptr, kv.second.size); }, + [&](foundry::vmm::CleanupPhase phase, CUresult status) { + switch (phase) { + case foundry::vmm::CleanupPhase::kUnmap: + fprintf(stderr, + "[HOOK] ERROR: VMM-IPC individual close: cuMemUnmap failed with error %d\n", + status); + break; + case foundry::vmm::CleanupPhase::kRelease: + fprintf(stderr, + "[HOOK] ERROR: VMM-IPC individual close: cuMemRelease failed with " + "error %d\n", + status); + break; + case foundry::vmm::CleanupPhase::kAddressFree: + fprintf(stderr, + "[HOOK] ERROR: VMM-IPC individual close: cuMemAddressFree failed with " + "error %d\n", + status); + break; + } + }); + individual_cleanup_complete = foundry::vmm::cleanup_complete(kv.second.cleanup, true); + }); + if (individual_cleanup_complete) { + global_alloc_metadata.erase_if([dptr](const std::pair& kv) { + return kv.first == dptr; + }); + } #ifdef HOOK_DEBUG - fprintf(stderr, "[HOOK] cuIpcCloseMemHandle: Closed VMM IPC ptr=0x%llx\n", - (unsigned long long)dptr); + if (individual_cleanup_complete) { + fprintf(stderr, "[HOOK] cuIpcCloseMemHandle: Closed VMM IPC ptr=0x%llx\n", + (unsigned long long)dptr); + } #endif - return CUDA_SUCCESS; + return cleanup_status; } // Not a VMM IPC allocation - fall back to real function @@ -3083,7 +4059,7 @@ CUresult cuIpcCloseMemHandle(CUdeviceptr dptr) { namespace foundry { // Forward declarations -void free_preallocated_region(); +bool free_preallocated_region(); void set_allocation_region(void* base, size_t size) { size_t aligned_base = align_to((size_t)base, kAllocAlignment); @@ -3166,7 +4142,12 @@ bool preallocate_region(size_t size) { if (tls_storage.has_preallocation) { fprintf(stderr, "[HOOK] WARNING: Preallocation already exists, freeing previous allocation\n"); - free_preallocated_region(); + if (!free_preallocated_region()) { + fprintf(stderr, + "[HOOK] ERROR: Cannot replace preallocation because cleanup failed; retry cleanup " + "before preallocating again\n"); + return false; + } } // Get and cache device and granularity for subsequent allocations @@ -3268,15 +4249,21 @@ bool preallocate_region(size_t size) { tls_storage.preallocated_start_addr = target_addr; tls_storage.preallocated_end_addr = target_addr + aligned_size; tls_storage.has_preallocation = true; + tls_storage.preallocation_cleanup = {}; + { + std::lock_guard lock(preallocated_chunks_mutex); + uint64_t generation = next_preallocated_chunk_generation.fetch_add(1); + preallocated_chunks[target_addr] = PreallocatedChunk{aligned_size, allocHandle, generation, {}}; + } // Note: we do NOT advance current_alloc_base_addr here. // The alloc calls will advance it as they consume the preallocated memory. return true; } -void free_preallocated_region() { +bool free_preallocated_region() { if (!tls_storage.has_preallocation) { - return; + return true; } typedef CUresult (*cuMemUnmap_t)(CUdeviceptr, size_t); @@ -3290,13 +4277,53 @@ void free_preallocated_region() { size_t preallocated_size = tls_storage.preallocated_end_addr - tls_storage.preallocated_start_addr; - mem_unmap_func(tls_storage.preallocated_start_addr, preallocated_size); - mem_release_func(tls_storage.preallocated_handle); + CUresult cleanup_status = CUDA_SUCCESS; + { + std::lock_guard lock(preallocated_chunks_mutex); + auto chunk_it = preallocated_chunks.find(tls_storage.preallocated_start_addr); + if (chunk_it != preallocated_chunks.end()) { + chunk_it->second.cleanup.started = true; + } + // Stop the live fd server from handing out this handle as soon as cleanup + // starts. The chunk record remains (marked started) until every required + // driver phase succeeds, so exports are rejected and cleanup is retryable. + vmm_ipc_invalidate_export((CUdeviceptr)tls_storage.preallocated_start_addr); + cleanup_status = foundry::vmm::run_cleanup( + tls_storage.preallocation_cleanup, CUDA_SUCCESS, false, + [&]() { return mem_unmap_func(tls_storage.preallocated_start_addr, preallocated_size); }, + [&]() { return mem_release_func(tls_storage.preallocated_handle); }, + []() { return CUDA_SUCCESS; }, + [](foundry::vmm::CleanupPhase phase, CUresult status) { + switch (phase) { + case foundry::vmm::CleanupPhase::kUnmap: + fprintf(stderr, + "[HOOK] ERROR: free_preallocated_region: cuMemUnmap failed with error %d\n", + status); + break; + case foundry::vmm::CleanupPhase::kRelease: + fprintf(stderr, + "[HOOK] ERROR: free_preallocated_region: cuMemRelease failed with error %d\n", + status); + break; + case foundry::vmm::CleanupPhase::kAddressFree: + break; + } + }); + if (chunk_it != preallocated_chunks.end()) { + chunk_it->second.cleanup = tls_storage.preallocation_cleanup; + } + if (!foundry::vmm::cleanup_complete(tls_storage.preallocation_cleanup, false)) { + return false; + } + preallocated_chunks.erase(tls_storage.preallocated_start_addr); + } tls_storage.preallocated_handle = 0; tls_storage.preallocated_start_addr = 0; tls_storage.preallocated_end_addr = 0; tls_storage.has_preallocation = false; + tls_storage.preallocation_cleanup = {}; + return cleanup_status == CUDA_SUCCESS; } size_t get_current_alloc_offset() { @@ -3464,7 +4491,10 @@ void replay_hook_events_from_json(const boost::json::object& events_obj) { size_t alignment = event_obj.at("alignment").to_number(); uint64_t expected_ptr = event_obj.at("ptr").to_number(); - CUdeviceptr aligned_addr = align_to(tls_storage.current_alloc_base_addr, alignment); + // Archives recorded before alignment normalization stored the caller's + // alignment (including CUDA's default-alignment sentinel, zero). + const size_t effective_alignment = alignment == 0 ? kAllocAlignment : alignment; + CUdeviceptr aligned_addr = align_to(tls_storage.current_alloc_base_addr, effective_alignment); if (aligned_addr != expected_ptr) { fprintf(stderr, "[REPLAY] FATAL: Reserve address mismatch!\n"); diff --git a/docs/sglang/direct-edits.md b/docs/sglang/direct-edits.md index 49c983c9..8decd987 100644 --- a/docs/sglang/direct-edits.md +++ b/docs/sglang/direct-edits.md @@ -1,93 +1,114 @@ # SGLang Direct Edits -Total footprint in the `sglang/` repo: **47 lines across 4 files**. Every other line of integration logic lives in `foundry.integration.sglang`. +## Zero upstream edits -## 1. `sglang/srt/foundry_shim.py` (new, ~27 lines) +The pinned-main integration makes **no edits to the SGLang source tree**. The +previous fork carried ~47 lines across four `sglang/` files +(`srt/foundry_shim.py`, a `ServerArgs` field + `__post_init__` call, and +`run_scheduler_process` / `run_data_parallel_controller_process` re-installs). +All of that is gone. Foundry now activates through SGLang's first-class +plugin entry point and installs its behavior through the upstream +`HookRegistry`, so the engine is used exactly as published at the pinned +commit: -The activation entry point. Re-exports the foundry public API and forces flags that conflict with SAVE/LOAD determinism. - -```python -def apply_server_args(server_args) -> None: - cfg_path = getattr(server_args, "foundry_graph_extension_config_path", None) - if not cfg_path: - return - - server_args.disable_piecewise_cuda_graph = True - server_args.enable_profile_cuda_graph = False - server_args.disable_flashinfer_autotune = True - - from foundry.integration.sglang.hooks import install_hooks - install_hooks(server_args) ``` - -`install_hooks` is idempotent (guarded by a module-level `_INSTALLED` flag) so calling it multiple times across processes is safe. - -## 2. `sglang/srt/server_args.py` (~12 lines) - -Adds the config-path field, the CLI argument, and the `__post_init__` activation: - -```python -class ServerArgs: - ... - foundry_graph_extension_config_path: Optional[str] = None - ... - - def __post_init__(self): - ... - if self.foundry_graph_extension_config_path: - from sglang.srt.foundry_shim import apply_server_args - apply_server_args(self) - ... - - @staticmethod - def add_cli_args(parser): - ... - parser.add_argument( - "--foundry-graph-extension-config-path", - type=str, - default=ServerArgs.foundry_graph_extension_config_path, - help="Path to Foundry CUDA graph extension TOML config.", - ) +foundry @ a8e43442421a80fea7991ee26f97f1003b170e29 +sglang @ a25164bda355250015c0ba6d8821f0d4b37abe7d (upstream main, unmodified) +model @ b968826d9c46dd6066d109eabc6255188de91218 (Qwen/Qwen3-8B) ``` -The field must live on `ServerArgs` (not be monkey-attached) because: +## Activation path -- `argparse` consumes it from the CLI. -- `ServerArgs` is pickled across the `spawn` boundary; only declared fields survive serialization. +1. **Entry point.** Foundry ships a standalone, dependency-light package + `foundry_sglang_plugin` that declares the entry point -## 3. `sglang/srt/managers/scheduler.py` (~4 lines) + ```toml + [project.entry-points."sglang.srt.plugins"] + foundry = "foundry_sglang_plugin:activate" + ``` -```python -def run_scheduler_process(server_args, ...): - load_plugins() - if server_args.foundry_graph_extension_config_path: - from sglang.srt.foundry_shim import apply_server_args - apply_server_args(server_args) - ... -``` + Importing `foundry_sglang_plugin` and calling `activate()` uses only the + Python standard library — no `foundry`, `torch`, or native import. -This is the first thing every spawned scheduler child runs. Without it, foundry's runtime monkey-patches aren't installed in the scheduler process and the model-loading lifecycle takes the upstream path. +2. **Selection.** SGLang's `load_plugins()` (in `sglang.srt.plugins`, called by + `sglang.cli.serve` before model dispatch and again defensively in the + engine) reads the comma-separated `SGLANG_PLUGINS` variable, invokes each + selected plugin as a zero-argument callable, and then applies the + `HookRegistry` hooks those callables registered. The recipe exports: -## 4. `sglang/srt/managers/data_parallel_controller.py` (~4 lines) + ```bash + export SGLANG_PLUGINS=foundry + ``` -```python -def run_data_parallel_controller_process(server_args, ...): - parent_process = psutil.Process().parent() - configure_logger(server_args) - if server_args.foundry_graph_extension_config_path: - from sglang.srt.foundry_shim import apply_server_args - apply_server_args(server_args) - ... -``` +3. **Config gate.** `activate()` reads `FOUNDRY_GRAPH_EXTENSION_CONFIG_PATH`: + + - **Unset (baseline).** `activate()` returns immediately. It imports no + `foundry`, no PyTorch, and no native extension, and registers no hooks — + a true no-op, so `SGLANG_PLUGINS=foundry` alone leaves the engine on its + stock path. + - **Set (SAVE / LOAD).** `activate()` lazily imports Foundry, loads the exact + TOML at that path, then lazily imports and calls `register_hooks()`, which + performs the nine idempotent `HookRegistry` registrations described in + [`hooks.md`](hooks.md). -Same purpose as the scheduler call, but for the DP-controller child. This child later spawns its own scheduler subprocesses, so it must have foundry's spawn-site patches installed before doing so — otherwise its nested `proc.start()` calls won't set `LD_PRELOAD` in their environments. + The recipe selects the mode by pointing the variable at the SAVE or LOAD + TOML (`recipe/experimental/sglang_foundry_tp_{save,load}.toml`) and clears it + for baseline. -## Why these can't be moved into the integration layer +## Why this replaces the fork edits -| Edit | Why direct | +| Fork edit (removed) | Pinned-main replacement | |---|---| -| `foundry_shim.py` | Has to be importable as `sglang.srt.foundry_shim`. It's the only file foundry's hooks can rely on being present in the sglang import path. | -| `ServerArgs.foundry_graph_extension_config_path` | Needs to be a real dataclass field (argparse + pickle). | -| `ServerArgs.__post_init__` call | Earliest activation point in the parent process. | -| `run_scheduler_process` install | `spawn` does not inherit Python state; the child needs an explicit re-install. | -| `run_data_parallel_controller_process` install | Same reason; this child also spawns further children. | +| `sglang/srt/foundry_shim.py` re-export module | `foundry_sglang_plugin:activate` entry point discovered by `load_plugins()` | +| `ServerArgs.foundry_graph_extension_config_path` field + CLI arg | `FOUNDRY_GRAPH_EXTENSION_CONFIG_PATH` environment variable read by `activate()` | +| `ServerArgs.__post_init__` activation call | AFTER-`HookRegistry` hook on `ServerArgs.__post_init__` (fail-closed profile validation) | +| `run_scheduler_process` re-install | `SGLANG_PLUGINS` is inherited across the `spawn` boundary; `load_plugins()` runs in every scheduler child | +| `run_data_parallel_controller_process` re-install | not needed — the dense TP=2 profile has no DP controller; spawn-site `LD_PRELOAD` is handled by the AROUND `Engine._launch_scheduler_processes` hook | + +Because activation is driven entirely by an installed entry point plus two +environment variables, the same image runs baseline, SAVE, and LOAD against the +unmodified pinned SGLang — the only difference between phases is +`FOUNDRY_GRAPH_EXTENSION_CONFIG_PATH`. + +The exact SGLang/model SHAs are validation provenance enforced by the Modal +harness. `_validate_profile` checks runtime topology, backend, and settings; it +does not inspect SHAs or reject an arbitrary shell-installed revision at +startup. Shell recipe users must install the matching SGLang pin. + +The active hooks retain a torch symmetric-memory TP communicator allocated +outside Foundry VMM, route multimem all-gather directly through its stable +buffer, enforce one padded batch-size-256 graph, and let the backend restore +that graph while traversing SGLang's normal LOAD-time capture lifecycle. + +## Evidence and scope + +The zero-edit architecture was verified first by focused SAVE→LOAD app +[`ap-5wO9cFhRZ7ITsLDDv27Xfh`](https://modal.com/apps/modal-labs/rahul-dev/ap-5wO9cFhRZ7ITsLDDv27Xfh) +and then by full matrix app +[`ap-6fYcSUFLOxocBSjNxakNGZ`](https://modal.com/apps/modal-labs/rahul-dev/ap-6fYcSUFLOxocBSjNxakNGZ). +The final run recorded one graph per rank at offset `68908220416`, 73 symmetric +all-reduce plus one symmetric all-gather and zero NCCL collective nodes per +rank, exact sequential and batch-8 outputs, reproducible SAVE/SAVE2 semantics, +byte-immutable LOAD/LOAD2 archives, batch-1/batch-8 replay, clean processes, +and completed cleanup. + +Plain-NCCL fresh-process SGLang restore is rejected by +[`ap-2qSQUKlY24L4RrQkMpUrmc`](https://modal.com/apps/modal-labs/rahul-dev/ap-2qSQUKlY24L4RrQkMpUrmc). +That app passed graph-count, offset, archive-immutability, and P2P/IPC structural +checks but both LOADs produced the same corrupt output. Archived +process-specific communicator and graph-sequence state is the leading forensic +explanation, not treatment-proven causality. The output observation rejects +plain-NCCL restore; the structural checks alone do not prove correctness. + +The positive claim is only for the exact pinned, single-host Qwen3-8B TP=2 +symmetric profile. It does not establish general/upstream TP or plain-NCCL +support. TP>2, multi-host, custom/NCCL-symmetric/fused collectives, CUMEM, NVLS, +EP, DP-attention, speculative decode, LoRA, pdmux, and unpinned revisions are +unsupported. + +## Historical note + +The `foundry-org/sglang` fork, its `foundry_shim.py`, the +`--foundry-graph-extension-config-path` CLI flag, and the "47 lines across 4 +files" footprint are **pre-port fork artifacts**. They are retained here only +to explain the migration; the current integration depends on none of them. diff --git a/docs/sglang/hooks.md b/docs/sglang/hooks.md index cf5b4c82..eb894686 100644 --- a/docs/sglang/hooks.md +++ b/docs/sglang/hooks.md @@ -1,214 +1,265 @@ # SGLang Hook Surface -Every monkey-patch `install_hooks(server_args)` installs. - -## Patches installed by `install_hooks` - -In install order: - -1. `_patch_init_torch_distributed` -2. `_patch_init_memory_pool` -3. `_patch_load_model` -4. `_patch_kernel_warmup` -5. `_patch_cuda_graph_capture` -6. `_patch_spawn_sites` - -> Removed: an earlier `_patch_model_runner_init` wrapped `ModelRunner.__init__` to stash `dp_rank` on `self` (upstream didn't expose it as an attribute). Upstream sglang now does `self.dp_rank = dp_rank` in the constructor, so our wrapper is no longer needed; `_resolve_dp_rank` reads `self.dp_rank` directly. - -### 1. `ModelRunner.init_torch_distributed` - -**Bind the device first (multi-GPU correctness).** Upstream `init_torch_distributed` -calls `torch.get_device_module(self.device).set_device(self.gpu_id)` *inside* -itself. But `setup_graph_extension` reserves the VMM region with -`set_allocation_region`, which binds to **whatever CUDA device is current at call -time** — and we call it *before* upstream. For DP/TP/EP rank > 0 the current -device is still `cuda:0` at that point, so without intervention the region lands -on the wrong GPU and the rank's later allocations fault with an **async illegal -memory access** that surfaces at an unrelated op (e.g. the first -`torch.cuda.Stream()` in `ModelRunner.__init__`). The patch therefore mirrors -upstream's `set_device(self.gpu_id)` up front (guarded on `self.device == "cuda"`) -before reserving the region. Single-GPU and rank 0 are unaffected (`gpu_id == 0`), -which is why this only appeared once DP was exercised. - -Before upstream: `set_device(self.gpu_id)`, then `setup_graph_extension(...)` -reserves the VMM region, loads cached fatbins (LOAD only via -`load_cuda_modules_and_libraries`), and eagerly initializes the cuBLAS handle into -scratch space. - -After upstream: `skip_to_scratch_boundary()` forces the cursor to `cfg.scratch_space_size`. +The pinned-main integration installs its behavior through SGLang's upstream +`HookRegistry` (`sglang.srt.plugins.hook_registry`). `register_hooks()` +(`foundry.integration.sglang.hooks`) registers exactly **nine** idempotent +targets under a module-level `_INSTALLED` guard; no upstream source is patched. + +The registration is unconditional, but each hook body dispatches on the Foundry +extension mode (`NONE` / `SAVE` / `LOAD`). Baseline (`SGLANG_PLUGINS=foundry` +without `FOUNDRY_GRAPH_EXTENSION_CONFIG_PATH`) never reaches `register_hooks()` +at all — `activate()` returns first — so no hook is installed. + +The target strings live in one source of truth, `foundry_sglang_plugin.targets`, +and are reused by both `hooks.py` and the clean-checkout Modal contract. + +## HookRegistry callback conventions + +Pinned `hook_registry.py` documents and implements: + +- `AFTER`: `fn(result, *args, **kwargs)` — runs after the target returns. +- `AROUND`: `fn(original_fn, *args, **kwargs)` — wraps the target; `fn` decides + whether and how to call `original_fn`. +- `REPLACE`: substitutes the target object; class replacements rebind every + `from … import Name` reference (sorted and applied first by `apply_hooks()`). + +## The nine targets + +### 1. `ServerArgs.__post_init__` — AFTER (profile validation) + +Runs on every configured launch (SAVE and LOAD), in the parent, after +`__post_init__` finishes deriving fields. `_validate_profile` aggregates every +violation into one `RuntimeError`, failing closed **before** any scheduler +spawn. It enforces the dense single-host topology, backend, and settings: + +- `nnodes=1`, `tp_size=2`, `pp_size=dp_size=dcp_size=ep_size=1`; +- no DP attention, `moe_a2a_backend="none"`, `disaggregation_mode="null"`, + `enable_pdmux=False`; +- `disable_custom_all_reduce=True`, `enable_torch_symm_mem=True`, + `enable_symm_mem=False`, mscclpp and NCCL-NVLS disabled, + `enforce_disable_flashinfer_allreduce_fusion=True`, and + `flashinfer_allreduce_fusion_backend=None`; +- `disable_flashinfer_autotune=True`; +- effective `flashinfer` prefill and decode attention backends; +- `cuda_graph_config.decode.backend == "full"` and + `cuda_graph_config.decode.max_bs == 256`, + `cuda_graph_config.decode.bs == [256]`, and + `cuda_graph_config.prefill.backend == "disabled"`; +- tokenizer list-batch encode enabled, dynamic tokenizer batching disabled, + CUDA graph padding enabled, and `decode_log_interval=1`; +- environment `NCCL_CUMEM_ENABLE=0`, `NCCL_NVLS_ENABLE=0`, + `PYTORCH_CUDA_ALLOC_CONF=expandable_segments:False`, + `SGLANG_ENABLE_POST_CAPTURE_KV_SIZING=0`. + +This is the "fail-closed through an executed hook" guarantee: unsupported +settings raise before capture. + +`_validate_profile` does not inspect the SGLang or model SHA. Those exact SHAs +are validation provenance enforced by the Modal harness, which checks out the +engine commit and passes the model revision. An arbitrary shell-installed +SGLang revision is not rejected at startup; shell recipe users must install the +matching pin. + +### 2. `Engine._launch_scheduler_processes` — AROUND (spawn LD_PRELOAD) + +Before delegating to the classmethod original, sets up `LD_PRELOAD` +(`libcuda_hook.so`, optionally `libnvshmem_host.so`), `FOUNDRY_MODE`, and a +spawn marker so every scheduler child inherits the Foundry hook library. NONE +mode skips the preload and just delegates. + +### 3. `ModelRunner.init_torch_distributed` — AROUND (rank / VMM setup) + +Reads the rank from the stored parallel-state object +(`model_runner.ps.{tp_rank,pp_rank,dp_rank}`), calls `setup_graph_extension(...)` +to reserve the VMM region (and, on LOAD, restore cached fatbins) **before** +upstream distributed init, then runs the original and forces the cursor to the +scratch boundary. Offsets are logged at each step: ``` -[Foundry] SGLang alloc_offset[after_setup_graph_ext]=… -[Foundry] SGLang alloc_offset[after_init_torch_dist]=… -[Foundry] SGLang alloc_offset[after_scratch_skip]=… +[Foundry] SGLang alloc_offset[after_setup_graph_ext]=… (… MB) +[Foundry] SGLang alloc_offset[after_init_torch_dist]=… (… MB) +[Foundry] SGLang alloc_offset[after_scratch_skip]=… (… MB) ``` -Validated on DP=2 (Qwen3-1.7B, one full replica per rank): both ranks reach an -identical `final_alloc_offset` across both SAVE passes and LOAD replays each -rank's graphs at that offset. Multi-rank runs also export -`NCCL_CUMEM_ENABLE=0` / `NCCL_NVLS_ENABLE=0` from the serve script — the CUMEM -P2P and NVLS multicast fast paths `cuMemMap` with driver-capability flags the -foundry VMM region doesn't carry. - -### 2. `ModelRunnerKVCacheMixin.init_memory_pool` - -- **SAVE**: call upstream `init_memory_pool` unchanged; after it returns, serialize the resolved `MemoryPoolConfig` (via `dataclasses.asdict`) into `warmup_state.json`. -- **LOAD**: skip upstream profiling. Load `MemoryPoolConfig` from `warmup_state.json`. Call `torch.cuda.empty_cache()` to mirror SAVE's `_resolve_memory_pool_config → get_available_gpu_memory(empty_cache=True)` side effect. Then call `_apply_memory_pool_config(config)` directly. - -The `empty_cache()` mirror is load-bearing — see [`memory-consistency.md`](memory-consistency.md) §5. - -### 3. `ModelRunner.load_model` - -Currently a passthrough wrapper. Kept as a future hook point for a LOAD-time `start_graph_builds(all_paths)` overlap with weight IO; the current LOAD path is fast enough (~80 ms total) that overlap isn't needed. +The pinned `ModelRunner.__init__` already binds the device with +`torch.get_device_module(self.device).set_device(ps.gpu_id)` *before* it calls +`init_torch_distributed`, so by the time this AROUND hook front-runs VMM +reservation (which touches the current device via +`torch._C._cuda_getCurrentBlasHandle()`) the correct device is already current. +The hook therefore does **not** call `set_device` itself — the fork's obsolete +second `set_device` shim was removed as redundant. + +### 4. `ModelRunner.alloc_memory_pool` — AROUND (SAVE persists the pool config) + +Logs the offset before/after the upstream allocation. On **SAVE**, rank 0 +serializes the resolved `MemoryPoolConfig` (`dataclasses.asdict`) into the +shared `warmup_state.json` (atomic temp-file + `os.replace`; nonzero ranks do +not write). On **LOAD**, it logs `"[Foundry] SGLang reused saved memory pool +config"` and otherwise delegates — the reuse itself happens in target 5. NONE +delegates untouched. + +### 5. `KVCacheConfigurator._resolve_memory_pool_config` — AROUND (LOAD reuse) + +This is the real memory-pool reuse hook. On **LOAD** it: + +1. loads `warmup_state.json` and validates the persisted dict against the exact + pinned `MemoryPoolConfig` field schema (missing / unknown / out-of-range / + non-finite / unsafely-widened values fail closed with a field-named + `RuntimeError`); +2. reconstructs `MemoryPoolConfig(**saved)`; +3. mirrors the one side effect the bypassed profiling path would have had — + `torch.cuda.empty_cache()` (SAVE's resolver calls + `get_available_gpu_memory(empty_cache=True)`); +4. returns the reconstructed config, so the normal target flow derives pool + sizes and initializes pools from it **without** profiling or calling the + original resolver. + +NONE and SAVE delegate to the original resolver unchanged. + +### 6. `FullCudaGraphBackend` — REPLACE (SAVE writes each graph) + +`build_foundry_full_backend_class()` lazily subclasses the pinned +`FullCudaGraphBackend`. The REPLACE rebinds `resolve_decode_backend`'s import, +so `DecodeCudaGraphRunner.__init__` builds the Foundry subclass as +`self.backend`. `FoundryFullCudaGraphBackend.capture_one(shape_key, forward_fn, +capture_inputs=None, post_warmup_hook=None)`: + +- in NONE mode, delegates to `super().capture_one(...)`; +- on **SAVE**, skips both upstream warmup forwards (their activation + allocations are not reproducible and would drift the VMM cursor) and the + no-op `post_warmup_hook`, captures the single `forward_fn()` into a Foundry + `CUDAGraph` on the shared pool / capture stream, populates + `self._graphs[shape_key]` / `self._outputs[shape_key]`, and calls + `graph_ops.save_graph(graph, output, shape_key)`. +- on **LOAD**, performs no forward or native capture; `capture_session` creates + the shared pool, preallocates, scans the archive, and calls + `start_graph_builds(paths, pool=self._pool, num_threads=4)` once, while each + upstream-requested `capture_one` validates shape/order and calls + `finish_one_graph_load` for the matching archive entry. + +Graphs are keyed by `ShapeKey(size, stream_idx, variant_label)` +(`runner/shape_key.py`). Foundry's dense TP=2 profile is single-stream and +non-LoRA, so `stream_idx` and `variant_label` are always `None`; the archive +filename is `graph_{index}_FULL_t{size}_r{size}_UX_pcN.{json,cugraph}` and the +save/load path fails closed on any populated `stream_idx`/`variant_label`. + +The decode graph output is a `LogitsProcessorOutput`; the dense profile +supports only `next_token_logits`. SAVE fails closed if `hidden_states`, +`full_logits`, or `customized_info` is populated rather than silently dropping +it. + +### 7. `DecodeCudaGraphRunner.capture` — AROUND (SAVE finalize / LOAD restore) + +The SAVE/LOAD branch point: + +- **NONE** — pass through to the upstream capture loop. +- **SAVE** — run the upstream loop (the REPLACE'd `capture_one` writes each + Foundry graph), then finalize the archive: + + ```python + save_graph_manifest() # topology groups for template + on-demand linking + pack_fatbins() # packed kernel fatbins + capture_final_alloc_offset() # per-rank VMM watermark + ``` + +- **LOAD** — delegate to the same pinned upstream `capture()` lifecycle as SAVE. + Upstream still runs warmup, `graph_capture()`, seq-len refill, + `capture_prepare`, `ForwardContext`, and real FlashInfer + `init_forward_metadata_out_graph`. The Foundry backend changes only graph + materialization: + + ```python + # backend.capture_session + pending = FoundryCUDAGraph.start_graph_builds( + paths, + pool=backend._pool, + num_threads=4, + ) + + # backend.capture_one, once per upstream-requested ShapeKey + graph, tensors = FoundryCUDAGraph.finish_one_graph_load(pending, index) + ``` + + This preserves SAVE's allocator trajectory and real FlashInfer metadata + wrappers while avoiding native capture and graph save. The sole pinned shape + is `ShapeKey(size=256)`. + + LOAD logs: + + ``` + [Foundry] Started 1 SGLang graph builds + [Foundry] Loaded SGLang graph graph_0_FULL_t256_r256_UX_pcN.json + [Foundry] SGLang alloc_offset[after_load_all_graphs]=68908220416 + ``` + +### 8. `TorchSymmMemCommunicator.__init__` — AROUND (stable ownership) + +In SAVE and LOAD, this hook raises the torch symmetric-memory signal-pad size, +stops Foundry allocation tracking, invokes the original constructor, and always +resumes tracking in `finally`. The symmetric buffer and rendezvous therefore +remain outside Foundry's VMM region. A successful communicator is retained by +its process-group rank tuple; duplicate construction replaces the same entry +rather than creating a second logical TP communicator. + +### 9. `MultimemAllGatherer.__call__` — AROUND (direct stable gather) + +In SAVE and LOAD, this hook requires exactly one retained communicator and a +contiguous rank-2 `bfloat16` input with supported alignment/capacity. It bypasses +the lazy upstream gather builder and calls +`torch.ops.symm_mem.multimem_all_gather_out` directly with a stable slice of +the retained communicator buffer and its original `group_name`. There is no +second rendezvous and no NCCL fallback. + +## Registration log -### 4. `ModelRunner.kernel_warmup` - -No-op on SAVE/LOAD: - -``` -[Foundry] SGLang kernel_warmup skipped in mode ``` - -The FlashInfer autotune path inside is shut off both by this no-op and by the forced `disable_flashinfer_autotune` flag. - -### 5. `CudaGraphRunner` (the largest patch) - -Four sub-patches: - -#### 5a. `_create_device_graph` - -On SAVE returns `foundry.CUDAGraph()` instead of `torch.cuda.CUDAGraph()`. - -#### 5b. `_capture_graph` - -On SAVE enters `foundry.graph(graph, pool=pool, stream=stream)` instead of `torch.cuda.graph(...)`, captures the `run_once` callable's allocations into foundry's hook event log. - -#### 5c. `capture_one_batch_size` - -On SAVE wraps the `forward` callable with a 3-call counter: - -```python -counter = [0]; real_forward = forward -def warmup_skipping_forward(*args, **kwargs): - counter[0] += 1 - if counter[0] <= 2: - return None - return real_forward(*args, **kwargs) -forward = warmup_skipping_forward -``` - -This suppresses the two pre-capture warmup forwards SGLang does in `for _ in range(2): run_once()`. Only the third invocation — inside `_capture_graph`'s graph capture context — runs the real forward. JIT and autotune allocations that those warmups would normally trigger now happen inside the captured graph and become foundry alloc events that replay verbatim on LOAD. (See `memory-consistency.md` §2.) - -After upstream returns, calls `save_graph(graph, output, key)`. The key matches the inline shape upstream uses for `self.graphs[key]` in `_capture_one_stream`: - -```python -key = bs if stream_idx is None else f"{stream_idx}_{bs}" -``` - -(Sglang removed the `_make_graph_key` / `get_capture_lora_variant` helpers in commit `ce2506e1c`, the same commit that deprecated `record_nolora_graph` dual MoE graph capture.) - -#### 5d. `capture` - -The outermost replacement. - -**SAVE**: - -```python -initialize_all_attention_metadata(self) # pre-pass -attn_backend.forward_metadata = None -real_init = attn_backend.init_forward_metadata_capture_cuda_graph -attn_backend.init_forward_metadata_capture_cuda_graph = reuse_pre_pass_init -try: - result = orig_capture(self, *args, **kwargs) # upstream capture loop -finally: - attn_backend.init_forward_metadata_capture_cuda_graph = real_init -save_graph_manifest() -pack_fatbins() -capture_final_alloc_offset() -``` - -`initialize_all_attention_metadata` walks `reversed(self.capture_bs)` and pre-allocates every per-bs FlashInfer wrapper. The wrappers are stored in `attn_backend.decode_cuda_graph_metadata[bs]`. - -`reuse_pre_pass_init` is a drop-in replacement for the upstream inner init that runs inside `capture_one_batch_size`. For decode mode it: - -1. Looks up the pre-pass wrapper from `decode_cuda_graph_metadata[bs]`. -2. Re-runs `indices_updater_decode.update(...)` with the same buffer slices (idempotent — writes plan info to the same `_int_workspace_buffer`). -3. Sets `attn_backend.forward_metadata = DecodeMetadata(wrappers)` so the captured forward sees the right metadata for this iter. - -No allocation. The captured graph references the pre-pass wrapper's address; LOAD's pre-pass produces a wrapper at the same address. - -**LOAD**: - -```python -if cgr.get_global_graph_memory_pool() is None: - cgr.set_global_graph_memory_pool(self.device_module.graph_pool_handle()) -set_graph_pool_id(cgr.get_global_graph_memory_pool()) -preallocate_for_load_mode() # cuMemCreate+cuMemMap up to final_alloc_offset -initialize_all_attention_metadata(self) # pre-pass (same as SAVE) -load_all_graphs(self) # ONE start_graph_builds + finish_graph_loads -self.graphs = {k: v[0] for k, v in state.loaded_graphs.items()} -self.output_buffers = {k: v[1] for k, v in state.loaded_graphs.items()} -``` - -`load_all_graphs` calls `start_graph_builds(all_paths)` and `finish_graph_loads(pending)` exactly once each. This is required for the manifest's template + on-demand linking to work; per-graph `start_graph_builds([single_path])` calls would leave on-demand graphs without a `shared_exec`, and runtime replay would abort with `Called CUDAGraph::replay without a preceding successful capture or load`. - -### 6. Spawn-site patches - -Two parent-side wrappers: - -```python -# Engine._launch_scheduler_processes -def patched_launch(self, *args, **kwargs): - if get_graph_extension_mode() != CUDAGraphExtensionMode.NONE: - rt.setup_ld_preload_env() - return orig_launch(self, *args, **kwargs) - -# DataParallelController.launch_tensor_parallel_group -def patched_start(self, *args, **kwargs): - if get_graph_extension_mode() != CUDAGraphExtensionMode.NONE: - rt.setup_ld_preload_env() - return orig_start(self, *args, **kwargs) +[Foundry] SGLang hooks registering: mode= workspace= +[Foundry] SGLang hooks registered ``` -`setup_ld_preload_env()` prepends `libcuda_hook.so` (and optionally `libnvshmem_host.so`) to `os.environ["LD_PRELOAD"]`, sets `FOUNDRY_MODE`, and records a wall-clock marker (`FOUNDRY_SPAWN_T0_NS`). All children spawned from these methods inherit the env. - -## Expert parallel (DeepEP) additions - -Active only when `moe_a2a_backend == deepep`. EP runs DP-attention + DeepEP (NCCL-free); -TP attention is unsupported (its NCCL all-reduce is incompatible with the VMM region). - -- **DeepEP buffer pre-capture bootstrap** (`bootstrap_deepep_buffer`, graph_ops). sglang - creates the singleton NVSHMEM `Buffer` lazily on the first MoE dispatch — normally - during the warmup forwards foundry suppresses, which would push creation *inside* the - captured stream (`deep_ep_cpp.Buffer(...)` → "operation not permitted when stream is - capturing"). The hook forces it before the capture loop, unwrapping the - `MaybeTboDeepEPDispatcher._inners` to reach a `DeepEPDispatcher`. Runs on SAVE and LOAD. -- **SAVE-only warmup pass** (`_run_warmup_pass`, `capture()` patch). Reuses the upstream - capture loop with graph capture neutered (run forwards only) to trigger every - pre-capture lazy init — DeepGEMM per-shape JIT (`stream.synchronize()` is illegal in - capture), buffer creation, etc. — outside the captured stream. LOAD doesn't need it - (preallocate + replay place allocations at recorded offsets; bootstrap covers NVSHMEM). -- **`deepep_adapter` mode on LOAD.** LOAD replaces the capture loop, so the adapter's - `_captured_deepep_mode` is never set; replay asserts on it. The hook calls - `deepep_adapter.capture(is_extend_in_batch=False)` after load. -- **FlashInfer pre-pass gated to FlashInfer.** The §5d pre-pass + `reuse_pre_pass_init` - shim handle FlashInfer's per-bs wrappers. fa3 (`FlashAttentionBackend`) uses a single - fixed `init_cuda_graph_state` workspace, so the shim is skipped (detected via absence of - `indices_updater_decode`); but fa3's per-bs `decode_cuda_graph_metadata[bs]` is still - populated post-load for the replay lookup. -- **C++: bind context on graph-build pool workers** (`CUDAGraphParallel.cpp`). EP graphs - carry `NODE_EVENT_RECORD/WAIT` nodes → `cuEventCreate` during on-demand prep, which runs - on `SimpleThreadPool` workers that never `cuCtxSetCurrent(main_ctx)`. Added that call - (mirrors the bg thread). Dense graphs never hit it (no event nodes), which is why it - surfaced only on sglang EP. - -See [`../../recipe/sglang/README.md`](../../recipe/sglang/README.md) for the EP serve -config and required kernel versions (deep_ep `9af0e0d`, sgl-deep-gemm ≥0.1.2, fa3). - -## Patch idiom - -All patches use the `wrap-and-call` idiom — short-circuit on `mode == NONE`, run foundry pre-work, call `orig`, run foundry post-work. The single exception is `capture()` on LOAD, which replaces the upstream method entirely (does not call `orig`). - -## Install order - -`install_hooks` calls the patch helpers in the listed order. Order doesn't matter for correctness here — every patch attaches to a different attribute — but spawn sites go last so the install-completion log line appears after every other patch has registered. +## Historical fork evidence (pre-port) + +The following 2×H100 / 2×H200 numbers came from the **pre-port fork** +integration (the monolithic `CudaGraphRunner` patch + `foundry_shim.py` +activation), not from this pinned-main plugin architecture. They are supporting +context only: + +- modal-projects draft [PR #4](https://github.com/modal-projects/foundry/pull/4) + provides review context; immutable + [commit `21bf734`](https://github.com/modal-projects/foundry/commit/21bf7341847ebce46c534e7f6e86a3929275c410) + recorded the 2×H100 run that wrote and restored 36 graphs per rank at + `final_alloc_offset=69736595456`; four temperature-0 baseline/LOAD responses + matched byte for byte. +- modal-projects draft [PR #2](https://github.com/modal-projects/foundry/pull/2) + provides review context; immutable + [commit `ee7d61a`](https://github.com/modal-projects/foundry/commit/ee7d61a32a4d4732e5f14978975b96ab03bcbede) + recorded the post-#26735 2×H200 run at + `final_alloc_offset=121831948288` with exact SAVE/LOAD equality on 16/16 + prompts. + +Those records do not establish current support. The current evidence is the +focused `VERIFIED` run +[`ap-5wO9cFhRZ7ITsLDDv27Xfh`](https://modal.com/apps/modal-labs/rahul-dev/ap-5wO9cFhRZ7ITsLDDv27Xfh) +and final `VERIFIED` matrix +[`ap-6fYcSUFLOxocBSjNxakNGZ`](https://modal.com/apps/modal-labs/rahul-dev/ap-6fYcSUFLOxocBSjNxakNGZ). +The final matrix recorded one graph per rank at offset `68908220416`; 73 +symmetric all-reduce plus one symmetric all-gather and zero NCCL collective +nodes per rank; exact sequential and batch-8 outputs; SAVE/SAVE2 semantic +reproducibility; LOAD/LOAD2 byte immutability; batch-1/batch-8 replay; times +217.2/103.0/103.1/61.0/61.0 seconds; clean phase processes; and completed +cleanup. + +## Scope + +Only Foundry `a8e43442421a80fea7991ee26f97f1003b170e29`, unmodified SGLang +`a25164bda355250015c0ba6d8821f0d4b37abe7d`, Qwen3-8B revision +`b968826d9c46dd6066d109eabc6255188de91218`, and single-host TP=2 are validated. +Plain-NCCL fresh-process restore is rejected by +[`ap-2qSQUKlY24L4RrQkMpUrmc`](https://modal.com/apps/modal-labs/rahul-dev/ap-2qSQUKlY24L4RrQkMpUrmc): +36 restored graphs per rank and matching structural evidence still produced +corrupt LOAD output. Archived process-specific communicator and graph-sequence +state is the leading forensic explanation, not treatment-proven causality. +The pinned symmetric result is not general/upstream TP or plain-NCCL support. +TP>2, multi-host, custom/NCCL-symmetric/fused collectives, CUMEM, NVLS, EP, +DP-attention, speculative decode, LoRA, pdmux, and unpinned revisions remain +unsupported. diff --git a/docs/sglang/memory-consistency.md b/docs/sglang/memory-consistency.md index d1e21d02..f3e4a870 100644 --- a/docs/sglang/memory-consistency.md +++ b/docs/sglang/memory-consistency.md @@ -1,123 +1,135 @@ # SGLang Memory-Consistency Post-Mortem -The hardest correctness problem in the SGLang integration was making LOAD produce the same output as SAVE. Cursors at every checkpoint matched, foundry's replay reported success, the graph launched cleanly — and the model emitted garbage tokens. - -This doc walks every divergence we hit and the fix that landed. The common cause: anything that allocates on one path but not the other shifts the VMM cursor or changes the address an `_int_workspace_buffer` lands at, and the captured graph kernels read addresses fixed at SAVE time. +The hardest correctness problem in the SGLang integration is making LOAD +produce the same output as SAVE. Cursors at every checkpoint can match, +Foundry's replay can report success, the graph can launch cleanly — and the +model can still emit garbage tokens if any allocation diverges. + +This doc walks each divergence and how the current pinned-main architecture +handles it. The common cause: anything that allocates on one path but not the +other shifts the VMM cursor or changes the address an `_int_workspace_buffer` +lands at, and the captured graph kernels read addresses fixed at SAVE time. + +The same warning applies to non-VMM execution context. Plain-NCCL +fresh-process restore was disproven by +[`ap-2qSQUKlY24L4RrQkMpUrmc`](https://modal.com/apps/modal-labs/rahul-dev/ap-2qSQUKlY24L4RrQkMpUrmc): +baseline and SAVE2 matched, and both LOADs restored 36 graphs per rank with +matching counts, offsets, immutable archive bytes, and P2P/IPC topology, yet +both LOADs produced the same corrupt decode. Each graph retained 73 NCCL +all-reduce and one NCCL all-gather node with process-specific communicator and +graph-sequence state. That state is the leading forensic explanation, not +treatment-proven causality. The output observation rejects this plain-NCCL +path; a structural pass is not an output-correctness pass. ## The general invariant -**SAVE and LOAD must walk an identical VMM cursor trajectory.** Foundry validates `start_base_addr_X` for each captured graph; the validation only catches asymmetries up to the first divergence. Captured-kernel argument addresses (e.g. wrapper `_int_workspace_buffer.data_ptr()`) are **not** validated — they just point at whatever VMM offset SAVE recorded. If LOAD's allocator lands a different buffer at the same offset, kernels read it; if no allocation lands at the offset, kernels read the preallocated zero-initialized memory. - -## Bug 1 — `CUDAGraph::capture_begin/end` toggling the VMM region - -### Symptom - -`replay_hook_events_from_json` errored on iter 2 with `Memory offset mismatch during replay`: current cursor 20 MB ahead of saved `start_base_addr_2`. - -### Root cause - -`foundry/csrc/CUDAGraph.cpp` used to bracket every capture with `resume_allocation_region()` / `stop_allocation_region()`. So on SAVE: - -- iter 0 enters with region ON. bs=24 init allocates inside VMM; cursor advances. -- iter 0 `capture_end` → region OFF. -- iter 1 init runs with region OFF. Allocations bypass VMM, no cursor advance. -- iter 1 `capture_begin` → region ON; `start_base_addr_1` recorded. -- … - -LOAD never calls `_capture_graph`, so nothing flipped the region off. LOAD's per-bs inits all advanced the cursor — past the saved `start_base_addr` values from iter 1 onward. - -### Fix - -Removed both calls in `CUDAGraph::capture_begin/end`. Only the hook recording window (`start_hook_record` / `end_hook_record`) brackets each capture. Region toggling is now each integration's responsibility, and both SGLang and vLLM keep the region ON across the whole capture loop. - -(vLLM initially tried a Python-side `resume/stop_allocation_region` bracket around each capture to suppress the `final_alloc_offset` inflation seen on Qwen3-30B-A3B EP2 — ~194 GB observed vs ~80 GB actual. That was wrong: any tensor a between-capture allocation produced that was later referenced by a captured kernel would have ended up outside the deterministic region, so LOAD's preallocation couldn't recreate it at the right address. The actual root cause was the `empty_cache()` inside `foundry/graph.py:CUDAGraph.__enter__`, which dropped torch caching-allocator segments and forced subsequent inter-capture allocations to take a fresh `cuMemAlloc_v2` path. Removing it brought `final_alloc_offset` back down to ~80 GB. The vLLM integration now also wraps the whole capture loop in `torch.cuda.use_mem_pool(graph_pool)` so SAVE-2 ↔ LOAD caching state stays symmetric for the per-graph `finish_one_graph_load` replay path.) - -## Bug 2 — Pre-capture warmup forwards on SAVE only - -### Symptom - -After bug 1: iter 2 still drifted 20 MB. - -### Root cause - -`CudaGraphRunner.capture_one_batch_size` runs: - -```python -for _ in range(2): - self.device_module.synchronize() - self.model_runner.tp_group.barrier() - run_once() -``` - -before `_capture_graph`. These forwards allocate activation tensors and immediately free them. Torch's caching allocator keeps the freed segments as cached free blocks. Later per-bs init allocations may hit or miss cache depending on segment fragmentation history — and the cache-hit pattern is not reproducible across SAVE and LOAD. - -LOAD never runs these warmups → different cache state → different cache hit/miss decisions → cursor drift. - -### Fix - -In `patched_capture_one_batch_size` wrap `forward` with a 3-call counter; first two invocations return `None`, the third (inside `_capture_graph`) runs the real forward. The `synchronize` and `barrier` calls in the loop are harmless. JIT and autotune that the warmups normally trigger now run inside the captured graph and are recorded as alloc events — foundry replays them verbatim on LOAD. - -## Bug 3 — Per-bs init outside graph capture - -### Symptom - -After bugs 1 and 2: cursors matched everywhere, no foundry-side error, but the model produced garbage tokens. - -### Root cause - -`init_forward_metadata_capture_cuda_graph(bs, ...)` is called by `capture_one_batch_size` **before** `_capture_graph`, so its allocations are not recorded in the captured graph's allocator events. Each call creates a new `BatchDecodeWithPagedKVCacheWrapper`; the wrapper's constructor allocates `_int_workspace_buffer` via `torch.empty(...)`. The captured graph kernels reference `_int_workspace_buffer.data_ptr()` directly. - -For LOAD's wrapper to land at the same VMM address as SAVE's, the cursor at each iter's init time must match. With per-iter init alone, that depends on caching-allocator behavior, which is non-deterministic even after bug 2. - -### Fix (two parts) - -**Both sides**: pre-pass init for every bs in `reversed(capture_bs)` order before the capture / load loop. `initialize_all_attention_metadata` walks the bs list and calls `initialize_attention_metadata_for_bs` for each. Same call sequence on both sides → same cursor trajectory → same wrapper addresses. - -**SAVE only**: the upstream `capture_one_batch_size` still calls the inner `init_forward_metadata_capture_cuda_graph(bs, ...)` per iter. Replaced it with `reuse_pre_pass_init` which: - -- Looks up the existing wrapper from `decode_cuda_graph_metadata[bs]`. -- Re-runs `indices_updater_decode.update(...)` with the same buffer slices — idempotent write to the same `_int_workspace_buffer`. -- Sets `attn_backend.forward_metadata = DecodeMetadata(wrappers)` for the current iter. - -No allocation. Captured graph references the pre-pass wrapper's address; LOAD's identical pre-pass produces a wrapper at the same address. - -(Two earlier approaches didn't work: relying on torch's caching allocator to reuse a freed segment after popping the dict entry, and forcing GC. The caching allocator doesn't deterministically reuse segments when request sizes don't exactly match cached blocks. Explicit reuse via dict lookup is the only reliable path.) - -The `attn_backend.forward_metadata = None` between the pre-pass and the patch install drops the last bs's `DecodeMetadata` reference — defensive, not strictly required by `reuse_pre_pass_init`. - -## Bug 4 — Per-graph `start_graph_builds` broke template/on-demand linking - -### Symptom - -After bugs 1-3: cursors matched perfectly, no foundry replay error — but the runtime crashed at first inference with `Called CUDAGraph::replay without a preceding successful capture or load`. - -### Root cause - -`graph_manifest.json` groups captured graphs by topology. One graph per group is built as a **template** (full cuGraph + cuGraphExec instantiated); the rest are **on-demand** graphs that link to the template's `shared_exec` and patch in per-graph node updates at link time. - -An earlier (interleaved) LOAD called `start_graph_builds([single_path])` once per graph. Every call had 1 input → 1 topology group → graph was its own template. But the manifest marks graphs 2-N as on-demand sharing graph 0's or graph 1's template. Those templates were loaded in earlier `start_graph_builds` calls and no longer in scope — so the on-demand graphs ended up with `shared_exec=null`. Runtime `CUDAGraph::replay` takes the on-demand path when `on_demand_data_` is set; with `shared_exec=null` it falls through to a `capture_ended_ || has_graph_exec_` check that errors. - -### Fix - -`load_all_graphs(self)` calls `start_graph_builds(all_paths)` exactly once, then `finish_graph_loads(pending)` once. All N graphs go through the manifest's template/on-demand linking in a single pass. `finish_graph_loads` then replays each graph's allocator events sequentially, walking the cursor through every `start_base_addr_X` in order. - -## Bug 5 — `_resolve_memory_pool_config`'s hidden `empty_cache` - -### Symptom - -After bugs 1-4: cursors at `after_setup_graph_ext`, `after_init_torch_dist`, `after_scratch_skip`, `before_init_memory_pool`, and `after_init_memory_pool` all matched between SAVE and LOAD. But `save_before_pre_init` (29698 MB on SAVE) was 20 MB ahead of LOAD's `before_preallocate` (29678 MB). - -### Root cause - -SAVE's `init_memory_pool` calls `_resolve_memory_pool_config` → `_profile_available_bytes` → `get_available_gpu_memory(...)` (in `sglang/srt/utils/common.py`), which calls `torch.cuda.empty_cache()`. `empty_cache` releases torch caching-allocator segments back to the driver via `cuMemFree_v2`, which the foundry hook unmaps. - -LOAD's `_patch_init_memory_pool` skipped upstream profiling entirely and went directly to `_apply_memory_pool_config` — no equivalent drain. Caching allocator on LOAD retained segments that SAVE released, so the subsequent attention-backend init (between `init_memory_pool` and `capture()`) took a different `cuMemAlloc` path on each side. 20 MB drift by the time `capture()` was reached. - -### Fix - -One line in `_patch_init_memory_pool`'s LOAD branch: - -```python -torch.cuda.empty_cache() -self._apply_memory_pool_config(self.memory_pool_config) -``` +**SAVE and LOAD must walk an identical VMM cursor trajectory.** Foundry +validates each captured graph's recorded `start_base_addr`; that only catches +asymmetries up to the first divergence. Captured-kernel argument addresses +(e.g. a FlashInfer wrapper's `_int_workspace_buffer.data_ptr()`) are **not** +validated — they point at whatever VMM offset SAVE recorded. If LOAD lands a +different buffer there, kernels read it; if nothing lands there, they read +preallocated zero memory. + +## Divergence 1 — capture toggling the VMM region (resolved in native code) + +An early `CUDAGraph.cpp` bracketed every capture with +`resume_allocation_region()` / `stop_allocation_region()`, so SAVE advanced the +cursor between captures while LOAD (which never captures) did not. The native +fix removed both calls; only the hook recording window brackets each capture, +and both SGLang and vLLM keep the region ON across the whole capture loop. + +## Divergence 2 — pre-capture warmup forwards on SAVE only + +Upstream `capture_one_shape` runs two warmup forwards before capturing. They +allocate activation tensors and free them; torch's caching allocator keeps the +freed segments, and later cache hit/miss decisions are not reproducible across +SAVE and LOAD. + +**Current handling.** `FoundryFullCudaGraphBackend.capture_one` skips **both** +warmups on SAVE (and the no-op `post_warmup_hook`) and captures only the single +`forward_fn()`. JIT/autotune that the warmups would trigger now happens inside +the captured graph and is recorded as Foundry alloc events that replay verbatim +on LOAD. (The pinned FlashInfer backend's `on_after_cuda_graph_warmup` is a +no-op, so there is no warmup-mutated state to undo — the hook is intentionally +not invoked.) + +## Divergence 3 — per-shape wrapper allocation ordering + +Each decode shape's FlashInfer wrapper is built by +`init_forward_metadata_out_graph(fb, in_capture=True)` → +`_create_decode_wrappers`, whose constructor allocates `_int_workspace_buffer`. +SAVE (`capture_one_shape`) allocates that wrapper immediately before capturing +that shape's graph, walking `reversed(capture_bs)`. + +**Current handling.** LOAD delegates to the same pinned upstream `capture()` +lifecycle as SAVE: warmup, `graph_capture`, seq-len refill, `capture_prepare`, +`ForwardContext`, and real +`init_forward_metadata_out_graph(in_capture=True)`. The backend's +`capture_session` issues one +`start_graph_builds(paths, pool=self._pool, num_threads=4)`, and backend +`capture_one` immediately calls `finish_one_graph_load(pending, index)` after +upstream creates that shape's real wrapper. The validated profile has exactly +one padded shape, batch size 256. + +## Divergence 4 — template / on-demand graph linking + +`graph_manifest.json` groups graphs by topology: one template per group (full +cuGraph + exec instantiated), the rest on-demand graphs that link to the +template's `shared_exec`. Issuing `start_graph_builds([single_path])` per graph +leaves on-demand graphs without a template in scope, so replay aborts with +`Called CUDAGraph::replay without a preceding successful capture or load`. + +**Current handling.** Backend `capture_session` issues exactly one +`start_graph_builds(all_paths, pool=self._pool)` so all graphs would remain in +one manifest/template scope, then backend `capture_one` drives per-index +`finish_one_graph_load(pending, i)` over the same index-sorted list; the batched +`finish_graph_loads` is not used. A native fix +(`CUDAGraphParallel.cpp::finish_one_graph_load_impl`) consumes the background +build future with `.get()` (not `.wait()`) so a failed template build surfaces +at the per-entry finish instead of being swallowed. + +## Divergence 5 — the hidden `empty_cache` in memory-pool resolution + +On SAVE, `KVCacheConfigurator._resolve_memory_pool_config` → +`_profile_available_bytes` → `get_available_gpu_memory(...)` calls +`torch.cuda.empty_cache()`, releasing caching-allocator segments to the driver +(which the Foundry hook unmaps). LOAD bypasses that resolver, so without a +mirror the caching state — and every subsequent `cuMemAlloc` decision — drifts. + +**Current handling.** The AROUND `_resolve_memory_pool_config` LOAD hook +performs the same `torch.cuda.empty_cache()` at the point the profiling path +would have, immediately before returning the reconstructed `MemoryPoolConfig`. +See [`memory-lifecycle.md`](memory-lifecycle.md). + +## Verification + +The `tests/modal_sglang_tp.py` acceptance harness converts these invariants +into checks: reproducible per-rank `final_alloc_offset` across SAVE and SAVE2, +LOAD/LOAD2 offsets equal to SAVE2, exact restored graph counts, restored-graph +replay at multiple batch sizes, and token-identical baseline/SAVE2/LOAD/LOAD2 +outputs. + +Focused run +[`ap-5wO9cFhRZ7ITsLDDv27Xfh`](https://modal.com/apps/modal-labs/rahul-dev/ap-5wO9cFhRZ7ITsLDDv27Xfh) +and final run +[`ap-6fYcSUFLOxocBSjNxakNGZ`](https://modal.com/apps/modal-labs/rahul-dev/ap-6fYcSUFLOxocBSjNxakNGZ) +are `VERIFIED`. The final run had one graph per rank at offset `68908220416`; +73 symmetric all-reduce plus one symmetric all-gather and zero NCCL collective +nodes per rank; exact sequential and batch-8 outputs; SAVE/SAVE2 semantic +reproducibility; LOAD/LOAD2 byte immutability; batch-1 and batch-8 replay; times +217.2/103.0/103.1/61.0/61.0 seconds; clean process exits; and completed +cleanup. + +This status is only for Foundry +`a8e43442421a80fea7991ee26f97f1003b170e29`, SGLang +`a25164bda355250015c0ba6d8821f0d4b37abe7d`, Qwen3-8B revision +`b968826d9c46dd6066d109eabc6255188de91218`, single-host TP=2, torch symmetric +memory, direct stable multimem gather, and one batch-size-256 graph. It is not +general/upstream TP or plain-NCCL support; other TP sizes, multi-host, +alternate collectives, EP/DP-attention, speculative decode, LoRA, and pdmux +remain unsupported. + +The exact engine/model SHAs are Modal-harness validation provenance, not fields +checked by `_validate_profile`. Shell recipe users must install the matching +SGLang pin; runtime fail-closed checks cover topology, backend, and settings. diff --git a/docs/sglang/memory-lifecycle.md b/docs/sglang/memory-lifecycle.md index 45b034e2..88abbde2 100644 --- a/docs/sglang/memory-lifecycle.md +++ b/docs/sglang/memory-lifecycle.md @@ -2,7 +2,12 @@ ## The invariant -**SAVE and LOAD must walk an identical VMM cursor trajectory** from `setup_graph_extension` through end-of-capture / end-of-load. Anything that allocates on one path but not the other shifts the cursor by some amount. The captured graph kernels reference VMM addresses fixed at SAVE time; if LOAD's allocations don't land at those exact offsets, the buffers the kernels read are unmapped or contain stale data — **silently**, with no foundry-side error. +**SAVE and LOAD must walk an identical VMM cursor trajectory** from +`setup_graph_extension` through end-of-capture / end-of-load. Anything that +allocates on one path but not the other shifts the cursor. The captured graph +kernels reference VMM addresses fixed at SAVE time; if LOAD's allocations don't +land at those exact offsets, the buffers the kernels read are unmapped or stale +— **silently**, with no Foundry-side error. ## TOML config schema @@ -14,106 +19,191 @@ workspace_root = "foundry_archive_…" # required; on-disk workspace scratch_space_size = "1024MB" # required; pre-deterministic scratch budget hook_library_path = "…/libcuda_hook.so" # optional; auto-discovered -nvshmem_host_path = "…/libnvshmem_host.so" # optional; for EP parity +nvshmem_host_path = "…/libnvshmem_host.so" # optional; EP only (unset for dense TP) ``` -- `base_addr` must be high enough not to collide with model weights or PyTorch's default heap. `0x600000000000` is a safe pick on x86_64. -- `region_size` must exceed `final_alloc_offset` for the model + KV pool + graphs. 256 GiB is generous for a 14 B model. -- `scratch_space_size` should be larger than the sum of NCCL workspace, cuBLAS workspace, and any other one-time allocations that aren't deterministic across runs. -- `workspace_root` is shared across ranks; per-rank artifacts go in `workspace_root/rank_{N}/`. Delete `workspace_root` before any SAVE. +- `base_addr` must not collide with model weights or PyTorch's default heap. + `0x600000000000` is a safe pick on x86_64. +- `region_size` must exceed the model + KV pool + graphs `final_alloc_offset`. +- `scratch_space_size` should exceed the one-time NCCL / cuBLAS workspace and + other non-deterministic allocations. +- `workspace_root` is shared across ranks; per-rank artifacts go in + `workspace_root/rank_{N}/`. Delete `workspace_root` before any SAVE. -## Lifecycle +The retained torch symmetric-memory communicator is deliberately **not** one +of these VMM allocations. Its constructor stops Foundry allocation tracking, +rendezvouses the symmetric buffer under PyTorch's own address contract, and +resumes Foundry tracking in `finally`. Direct multimem gather uses a stable +slice of that retained outside-VMM buffer. + +## Lifecycle (pinned-main HookRegistry targets) ```mermaid flowchart TD - INIT["ModelRunner.__init__
(upstream stashes self.dp_rank)"] - TORCH["ModelRunner.init_torch_distributed"] - TORCH_PRE["setup_graph_extension
(reserve VMM region, eager cuBLAS,
LOAD: load cached fatbins)"] + POST["ServerArgs.__post_init__
AFTER: fail-closed profile validation"] + TORCH["ModelRunner.init_torch_distributed
AROUND"] + TORCH_PRE["setup_graph_extension(ps.tp_rank,…)
reserve VMM region, eager cuBLAS,
LOAD: load cached fatbins"] TORCH_ORIG["<upstream init_torch_distributed>
NCCL warmup + comm bring-up"] TORCH_POST["skip_to_scratch_boundary
cursor → scratch_space_size"] - LOAD_MODEL["ModelRunner.load_model
(passthrough today)"] - POOL["ModelRunnerKVCacheMixin.init_memory_pool"] - POOL_SAVE["SAVE: orig + save resolved
MemoryPoolConfig to warmup_state.json"] - POOL_LOAD["LOAD: load MemoryPoolConfig
+ empty_cache()
+ _apply_memory_pool_config"] - KWARM["ModelRunner.kernel_warmup
(no-op on SAVE/LOAD)"] - GRAPHS["ModelRunner.init_device_graphs
→ CudaGraphRunner(self).capture()"] - SAVE_CAP["SAVE: pre-pass init
+ reuse_pre_pass_init shim
+ capture loop
+ save_graph_manifest
+ pack_fatbins
+ capture_final_alloc_offset"] - LOAD_CAP["LOAD: preallocate_for_load_mode
+ pre-pass init
+ load_all_graphs"] - - INIT --> TORCH + SYMM["TorchSymmMemCommunicator.__init__
AROUND: stop Foundry VMM tracking,
construct + retain communicator, resume"] + ALLOC["ModelRunner.alloc_memory_pool
AROUND"] + ALLOC_SAVE["SAVE: orig + rank 0 serializes resolved
MemoryPoolConfig to warmup_state.json"] + ALLOC_LOAD["LOAD: orig (config reused in resolver below)"] + RESOLVE["KVCacheConfigurator._resolve_memory_pool_config
AROUND"] + RESOLVE_LOAD["LOAD: reconstruct MemoryPoolConfig,
skip profiling + original resolver,
mirror empty_cache()"] + CAP["DecodeCudaGraphRunner.capture
AROUND (backend = FoundryFullCudaGraphBackend, REPLACE)"] + SAVE_CAP["SAVE: upstream loop (capture_one writes graphs)
+ save_graph_manifest + pack_fatbins
+ capture_final_alloc_offset"] + LOAD_CAP["LOAD: same upstream capture lifecycle;
backend preallocates + starts one pooled build,
then finishes one archived ShapeKey(size=256)"] + + POST --> TORCH TORCH --> TORCH_PRE --> TORCH_ORIG --> TORCH_POST - TORCH_POST --> LOAD_MODEL --> POOL - POOL --> POOL_SAVE - POOL --> POOL_LOAD - POOL_SAVE --> KWARM - POOL_LOAD --> KWARM - KWARM --> GRAPHS - GRAPHS --> SAVE_CAP - GRAPHS --> LOAD_CAP + TORCH_POST --> SYMM --> ALLOC + ALLOC --> ALLOC_SAVE + ALLOC --> ALLOC_LOAD + ALLOC_LOAD --> RESOLVE --> RESOLVE_LOAD + ALLOC_SAVE --> CAP + RESOLVE_LOAD --> CAP + CAP --> SAVE_CAP + CAP --> LOAD_CAP ``` +(The `_resolve_memory_pool_config` hook fires *inside* the normal target's +`alloc_memory_pool` on LOAD; the arrows show data flow, not two separate +calls.) + ## Allocation buckets ### Bucket A — pre-deterministic scratch -CUDA context, cuBLAS handle, NCCL warmup, distributed bring-up. These run **inside** the VMM region but are released by `skip_to_scratch_boundary` — the cursor is forced to `cfg.scratch_space_size` regardless of what landed below it. Any non-determinism in this region is invisible to the rest of the lifecycle. +CUDA context, cuBLAS handle, NCCL warmup, distributed bring-up. These run +inside the VMM region but are released by `skip_to_scratch_boundary` — the +cursor is forced to `cfg.scratch_space_size` regardless of what landed below +it. Non-determinism here is invisible to the rest of the lifecycle. + +The retained torch symmetric-memory buffer is a separate allocation class. It +is constructed after the scratch boundary while Foundry's allocation region is +temporarily stopped, so its rendezvous and stable address remain owned by +PyTorch rather than serialized as Foundry VMM state. ### Bucket B — deterministic runtime state -Model weights, KV pools, attention workspace buffers, FlashInfer wrapper `_int_workspace_buffer`s, captured-graph alloc events. These must allocate in identical order on SAVE and LOAD. Their final cursor position is the `final_alloc_offset` persisted to `warmup_state.json`. +Model weights, KV pools, attention workspace buffers, FlashInfer wrapper +`_int_workspace_buffer`s, captured-graph alloc events. These must allocate in +identical order on SAVE and LOAD. Their final cursor position is the +`final_alloc_offset` persisted to `warmup_state.json` and `final_alloc_offset.json`. ### Bucket C — forbidden divergence -Anything that runs on one path but not the other. The fixes we landed identify and align the three known cases: - -1. The two pre-capture warmup forwards in `capture_one_batch_size` — skipped on SAVE so they don't pollute the caching allocator with freed activations LOAD can't reproduce. (See doc 06.) -2. The per-iter inner `init_forward_metadata_capture_cuda_graph(bs)` call — replaced on SAVE with `reuse_pre_pass_init` so it doesn't re-allocate the wrappers the pre-pass already built. (See doc 03 / doc 05.) -3. `_resolve_memory_pool_config` calls `get_available_gpu_memory(empty_cache=True)` on SAVE; LOAD's `_patch_init_memory_pool` mirrors it with an explicit `torch.cuda.empty_cache()` before `_apply_memory_pool_config`. (See below.) +Anything that runs on one path but not the other. The current hooks align the +known cases: + +1. The two pre-capture warmup forwards — `FoundryFullCudaGraphBackend.capture_one` + skips them on SAVE so they don't pollute the caching allocator with freed + activations LOAD can't reproduce. (See `memory-consistency.md`.) +2. FlashInfer wrapper allocation — LOAD runs the same upstream `capture()`, + `capture_prepare`, `ForwardContext`, and + `init_forward_metadata_out_graph(in_capture=True)` lifecycle as SAVE. The + backend's `capture_one` immediately finishes the matching archived graph, so + the real wrapper lands at the SAVE address. There is exactly one padded + shape, batch size 256. +3. `_resolve_memory_pool_config` calls `get_available_gpu_memory(empty_cache=True)` + on SAVE; the LOAD resolver hook mirrors it with an explicit + `torch.cuda.empty_cache()` before returning the reconstructed config. ## The `_resolve_memory_pool_config` mirror -`init_memory_pool` on SAVE calls `_resolve_memory_pool_config` → `_profile_available_bytes` → `get_available_gpu_memory(...)` which calls `torch.cuda.empty_cache()` (`sglang/srt/utils/common.py`). `empty_cache` releases torch caching-allocator segments back to the driver via `cuMemFree_v2`, which the foundry hook unmaps. +On SAVE the normal target resolver calls `_profile_available_bytes` → +`get_available_gpu_memory(...)` which calls `torch.cuda.empty_cache()` +(`sglang/srt/utils/common.py`). `empty_cache` releases torch caching-allocator +segments to the driver via `cuMemFree_v2`, which the Foundry hook unmaps. -LOAD skips upstream `init_memory_pool` and goes straight to `_apply_memory_pool_config`, so it had no equivalent drain. Caching-allocator segments retained on LOAD changed which torch.empty calls hit cache vs. miss during the subsequent attention-backend init — pushing the cursor 20 MB behind SAVE by the time `capture()` was reached. - -The fix is one line: - -```python -torch.cuda.empty_cache() # mirror SAVE's _resolve_memory_pool_config side effect -self._apply_memory_pool_config(self.memory_pool_config) -``` +LOAD bypasses that resolver entirely (it returns the saved config), so it had +no equivalent drain. The LOAD resolver hook therefore performs the same +`torch.cuda.empty_cache()` at the point the profiling path would have, before +returning `MemoryPoolConfig(**saved)`. The bypassed original resolver and +profiler never run. ## Why saving raw memory numbers isn't enough -The pool config is computed from a profile and a configurator: +The pool config is computed from a profile plus a configurator. Persisting the +resolved `MemoryPoolConfig` (via `dataclasses.asdict`) and reconstructing it on +LOAD is necessary and sufficient; persisting raw free-memory numbers would +re-run the configurator on LOAD with stale inputs. -1. `_profile_available_bytes(pre_model_load_memory)` -2. `MemoryPoolConfigurator.calculate_pool_sizes(...)` -3. `_apply_token_constraints(...)` (user caps, page alignment, PP sync) -4. `_resolve_max_num_reqs(...)` -5. `_apply_memory_pool_config(...)` +`WarmupState.memory_pool_config` is exactly this dict. On LOAD the hook +validates it against the exact pinned `MemoryPoolConfig` field schema before +construction: -Persisting the resolved `MemoryPoolConfig` (via `dataclasses.asdict`) and re-applying it on LOAD is necessary and sufficient. Persisting raw free-memory numbers would re-run the configurator on LOAD with stale inputs. +1. `max_total_num_tokens` (required positive int) +2. `max_running_requests` (positive int; the target resolver always fills it) +3. `full_max_total_num_tokens`, `swa_max_total_num_tokens` (optional non-negative) +4. `c4_max_total_num_tokens`, `c128_max_total_num_tokens`, + `c4_state_pool_size`, `c128_state_pool_size` (non-negative int) +5. `mem_fraction_static` (`Optional[float]`; a JSON integer is accepted only + when the exact value is preserved and lies in `(0, 1]`) -`WarmupState.memory_pool_config` is exactly this dict. +Missing / unknown / mistyped / out-of-range / non-finite values fail closed +with a field-named `RuntimeError`. ## VMM region setup -`setup_graph_extension` (in `runtime.py`): +`setup_graph_extension` (in `runtime.py`), driven by the AROUND +`init_torch_distributed` hook using `model_runner.ps.{tp_rank,pp_rank,dp_rank}`: -- Computes the per-rank workspace path (`{workspace_root}/rank_{compute_workspace_rank(...)}`) -- On SAVE: removes the rank workspace if it exists, then creates it fresh -- On LOAD: - - `cge.set_skip_fatbin_processing(True)` - - `cge.load_cuda_modules_and_libraries(workspace_dir)` — restores the fatbins SAVE wrote, into device code memory -- `cge.set_allocation_region(cfg.base_addr, parse_size(cfg.region_size))` — reserves a VMM address range; the cursor starts at `base_addr` -- `torch._C._cuda_getCurrentBlasHandle()` — eager cuBLAS handle so the workspace it wants lands in scratch instead of post-scratch territory -- creates the `CUDAGraphExtensionState` singleton +- computes the per-rank workspace path + (`{workspace_root}/rank_{compute_workspace_rank(...)}`); +- on SAVE: removes and recreates the rank workspace; +- on LOAD: `cge.set_skip_fatbin_processing(True)` then + `cge.load_cuda_modules_and_libraries(workspace_dir)` to restore the fatbins + SAVE wrote; +- `cge.set_allocation_region(cfg.base_addr, parse_size(cfg.region_size))` — + reserves the VMM range; the cursor starts at `base_addr`; +- `torch._C._cuda_getCurrentBlasHandle()` — eager cuBLAS handle into scratch; +- creates the `CUDAGraphExtensionState` singleton. -After upstream `init_torch_distributed` returns, `skip_to_scratch_boundary` forces the cursor to `cfg.scratch_space_size` (default 1 GiB). Allocations below that line are scratch and don't need to be deterministic. +After upstream `init_torch_distributed` returns, `skip_to_scratch_boundary` +forces the cursor to `cfg.scratch_space_size` (default 1 GiB). ## Final watermark -`capture_final_alloc_offset` runs after the SAVE-side capture loop completes (after `save_graph_manifest` and `pack_fatbins`). It writes `final_alloc_offset` to both `rank_{N}/final_alloc_offset.json` and the shared `warmup_state.json`. - -`preallocate_for_load_mode` uses this to call `cge.preallocate_region(final - current)` so the entire deterministic range is mapped to physical memory in one shot. The cursor is **not** advanced — preallocate just pre-maps; the cursor advances naturally as cuMemAllocs land within the preallocated range (fast-path: pointer bump, no driver calls). +`capture_final_alloc_offset` runs after the SAVE capture loop completes (after +`save_graph_manifest` and `pack_fatbins`). It writes `final_alloc_offset` to +both `rank_{N}/final_alloc_offset.json` and (rank 0) the shared +`warmup_state.json`. + +`preallocate_for_load_mode` uses this to `cge.preallocate_region(final - +current)`, mapping the entire deterministic range in one shot. The cursor is +**not** advanced — preallocate just pre-maps; the cursor advances naturally as +allocations land within the preallocated range. + +## Verified boundary + +Focused treatment +[`ap-5wO9cFhRZ7ITsLDDv27Xfh`](https://modal.com/apps/modal-labs/rahul-dev/ap-5wO9cFhRZ7ITsLDDv27Xfh) +and final matrix +[`ap-6fYcSUFLOxocBSjNxakNGZ`](https://modal.com/apps/modal-labs/rahul-dev/ap-6fYcSUFLOxocBSjNxakNGZ) +verified the same per-rank offset `68908220416`, one graph per rank, +73 symmetric all-reduce plus one symmetric all-gather and zero NCCL collective +nodes per rank, exact sequential and batch-8 outputs, SAVE/SAVE2 semantic +reproducibility, and LOAD/LOAD2 byte immutability. Both LOADs replayed batch +sizes 1 and 8; times were 217.2/103.0/103.1/61.0/61.0 seconds; every process +exited cleanly; and cleanup completed. + +This does not validate plain NCCL. Disproving app +[`ap-2qSQUKlY24L4RrQkMpUrmc`](https://modal.com/apps/modal-labs/rahul-dev/ap-2qSQUKlY24L4RrQkMpUrmc) +matched structural graph/offset/archive checks yet produced corrupt fresh-LOAD +decode output. Archived process-specific communicator and graph-sequence state +is the leading forensic explanation, not treatment-proven causality. The +observed corruption rejects this plain-NCCL path without proving that cause. +The positive claim is limited to Foundry +`a8e43442421a80fea7991ee26f97f1003b170e29`, SGLang +`a25164bda355250015c0ba6d8821f0d4b37abe7d`, Qwen3-8B revision +`b968826d9c46dd6066d109eabc6255188de91218`, single-host TP=2, and the exact +torch symmetric-memory profile. Other TP sizes, multi-host, alternate +collectives, EP/DP-attention, speculative decode, LoRA, and pdmux are +unsupported. It is not general/upstream TP or plain-NCCL support. + +The exact engine/model SHAs are validation provenance enforced by the Modal +harness, not fields checked by `_validate_profile`. Shell recipe users must +install the matching SGLang pin; runtime fail-closed checks cover topology, +backend, and settings. diff --git a/docs/sglang/overview.md b/docs/sglang/overview.md index f8ea7cce..78d16e21 100644 --- a/docs/sglang/overview.md +++ b/docs/sglang/overview.md @@ -1,112 +1,226 @@ # SGLang Integration Overview -Foundry persists SGLang's `CudaGraphRunner` graphs to disk on SAVE and restores them on LOAD, skipping graph capture, kernel warmup, and the per-batch-size attention metadata setup costs. +Foundry persists SGLang's decode CUDA graphs to disk on SAVE and restores them +on LOAD. LOAD runs the real upstream lifecycle, including each shape's +`init_forward_metadata_out_graph` planning and FlashInfer wrapper allocation in +SAVE's exact order. The Foundry backend replaces only native graph capture with +archive materialization, so restored kernels find their workspaces at the +recorded addresses. +The current integration targets **pinned upstream SGLang main** and activates +with **zero upstream edits** through SGLang's `sglang.srt.plugins` entry-point +mechanism. -Tested on single-GPU Qwen3-1.7B / 4B / 14B and **data-parallel (DP=2)** Qwen3-1.7B with the FlashInfer attention backend. +``` +sglang @ a25164bda355250015c0ba6d8821f0d4b37abe7d (upstream main, unmodified) +model Qwen/Qwen3-8B @ b968826d9c46dd6066d109eabc6255188de91218 +``` + +## Activation (zero upstream edits) + +Foundry ships a standalone entry point in the `sglang.srt.plugins` group: + +```toml +[project.entry-points."sglang.srt.plugins"] +foundry = "foundry_sglang_plugin:activate" +``` + +- `SGLANG_PLUGINS=foundry` selects it; SGLang's `load_plugins()` invokes + `activate()` and then applies the `HookRegistry` hooks it registers. +- `FOUNDRY_GRAPH_EXTENSION_CONFIG_PATH` gates behavior: unset → `activate()` is + a pure-stdlib no-op (baseline); set → it loads the TOML and calls + `register_hooks()`. + +See [`direct-edits.md`](direct-edits.md) for the full activation path and +[`hooks.md`](hooks.md) for the nine `HookRegistry` targets. ## Parallelism +The pinned plugin integration has one validated profile: + | Mode | Status | Notes | |---|:---:|---| -| Single GPU | ✅ | Qwen3-1.7B / 4B / 14B | -| Data parallel (DP) | ✅ | One full replica per rank; validated DP=2. Requires the per-rank device binding (below) and `NCCL_CUMEM_ENABLE=0` / `NCCL_NVLS_ENABLE=0`. | -| Tensor parallel (TP) | 🚧 | Deterministic NCCL memory layout is under construction. | -| Expert parallel (DeepEP) | ✅ | Validated EP=2 on Qwen3-30B-A3B-FP8 (SAVE/SAVE2/LOAD/query); restored decode graphs match baseline throughput. See **Expert parallel** below. | - -**Expert parallel (DeepEP).** EP runs DP-attention (each rank its own attention — no -NCCL all-reduce) + DeepEP for the MoE all-to-all (NVSHMEM, foundry-compatible). The -serve script is `recipe/sglang/serve_qwen3-30ba3bfp8_ep.sh - [--save|--load]` with: `--enable-dp-attention --moe-a2a-backend deepep ---deepep-mode low_latency --moe-runner-backend deep_gemm --attention-backend fa3 ---disable-custom-all-reduce`. Required kernel builds in the env: `deep_ep` at sglang's -pinned commit (`9af0e0d`, not vLLM's), `sgl-deep-gemm>=0.1.2` (0.1.0 lacks -`m_grouped_bf16_gemm_nt_masked`), `flash-attn-3`. `fa3` is required because the -flashinfer ragged-prefill path has an off-by-one (`q.shape != qo_indptr`) under this -config. DeepEP low-latency caps dispatch at -`SGLANG_DEEPEP_NUM_MAX_DISPATCH_TOKENS_PER_RANK` (default 128); raise it (+ chunk -prefill) for larger batches and keep it identical across SAVE/LOAD. Foundry-specific -EP handling is in [`hooks.md`](hooks.md): a DeepEP buffer pre-capture bootstrap, a -SAVE-only warmup pass (triggers DeepGEMM JIT + buffer creation outside the capture -stream), `deepep_adapter` mode init on LOAD, the FlashInfer pre-pass gated off for -fa3, and a C++ fix binding the CUDA context on the graph-build pool workers. - -**Per-rank device binding (DP/TP/EP).** Foundry's `set_allocation_region` binds the -VMM region to the CUDA device current at call time. Upstream sets the device -*inside* `init_torch_distributed`, which the integration wraps and front-runs, so -the hook explicitly calls `set_device(self.gpu_id)` before reserving the region — -otherwise rank > 0 reserves on `cuda:0` and faults. See [`hooks.md`](hooks.md) §1. -The DP serve script lives at `recipe/sglang/serve_qwen3-1.7b_dp.sh` -(` [--save|--load]`); pick GPUs with `CUDA_VISIBLE_DEVICES`. - -## How to use - -### 1. Install foundry - -```bash -pushd foundry && pip install -e . --no-build-isolation && popd +| Dense tensor parallel (TP=2) | ✅ | Pinned single-host Qwen3-8B, torch symmetric memory, one padded batch-size-256 full-decode graph per rank. | + +The exact SGLang and model SHAs are validation provenance enforced by the Modal +harness: it checks out and verifies the engine SHA and passes the model +revision. `_validate_profile` does not inspect either SHA, so arbitrary +shell-installed SGLang revisions are not rejected at startup. Shell recipe +users must install the matching SGLang pin. + +The `ServerArgs.__post_init__` hook fails closed on the runtime topology, +backend, and settings: TP≠2, PP/DP/DCP/EP≠1, DP-attention, expert parallel, +pdmux, custom/NCCL-symmetric/fused collectives, CUMEM, NVLS, and multi-host +settings are rejected before capture. Speculative decode and LoRA are also +unsupported by the archived `ShapeKey` and output contracts. + +### Historical fork configurations (pre-port) + +Single-GPU (Qwen3-1.7B / 4B / 14B), data-parallel (DP=2 Qwen3-1.7B), and +expert-parallel (EP=2 Qwen3-30B-A3B-FP8) were validated on the **pre-port +`foundry-org/sglang` fork** integration (the monolithic `CudaGraphRunner` +patch + `foundry_shim.py` activation). They predate the pinned-main port and +are **not runnable through the current plugin**, whose profile validation +requires dense TP=2. They are recorded as historical evidence only, not as +current-integration status. + +## Experimental tensor parallel + +The strict entry point is `recipe/experimental/serve_qwen3-8b_sglang_tp.sh 2 +[--save|--load]`; it rejects other TP sizes and extra arguments. Baseline, +SAVE, and LOAD use the same topology and communication settings: + +- `SGLANG_PLUGINS=foundry` in every phase; the SAVE/LOAD TOML is selected via + `FOUNDRY_GRAPH_EXTENSION_CONFIG_PATH` (cleared for baseline); +- `--enable-torch-symm-mem` selects the retained torch symmetric-memory TP + communicator; `--disable-custom-all-reduce` and + `--enforce-disable-flashinfer-allreduce-fusion` prevent unsupported fallback; +- `NCCL_CUMEM_ENABLE=0`, `NCCL_NVLS_ENABLE=0`, + `PYTORCH_CUDA_ALLOC_CONF=expandable_segments:False`, + `SGLANG_ENABLE_POST_CAPTURE_KV_SIZING=0`; +- `--disable-flashinfer-autotune`, FlashInfer attention, and + `--cuda-graph-config '{"decode":{"backend":"full","max_bs":256,"bs":[256]},"prefill":{"backend":"disabled"}}'`; +- a fixed server seed (42 by default); +- `--revision b968826d9c46dd6066d109eabc6255188de91218`. + +The communicator constructor temporarily stops Foundry allocation tracking so +PyTorch allocates and rendezvouses the retained symmetric buffer outside +Foundry VMM. The multimem hook calls +`torch.ops.symm_mem.multimem_all_gather_out` directly into a stable slice of +that buffer and never builds upstream lazy gather state. + +### Why plain NCCL is rejected + +Fresh-process plain-NCCL restore was disproven by +[`ap-2qSQUKlY24L4RrQkMpUrmc`](https://modal.com/apps/modal-labs/rahul-dev/ap-2qSQUKlY24L4RrQkMpUrmc). +Baseline and independent SAVE2 output matched byte for byte. Both fresh LOAD +processes produced the same corrupt decode output even though all 36 graphs per +rank restored, graph inventory/final offsets/archive bytes matched, and NCCL +reported P2P/IPC topology. Every graph retained 73 NCCL all-reduce nodes and +one NCCL all-gather node with SAVE-process communicator and graph-sequence +state. That process-specific state is the leading forensic explanation, not +treatment-proven causality. The observed deterministic corruption rejects this +plain-NCCL restore path; structural restoration is necessary evidence, not +output correctness. + +### Verified evidence + +The focused SAVE→LOAD treatment at Foundry `a8e43442421a80fea7991ee26f97f1003b170e29` +is `VERIFIED`: +[`ap-5wO9cFhRZ7ITsLDDv27Xfh`](https://modal.com/apps/modal-labs/rahul-dev/ap-5wO9cFhRZ7ITsLDDv27Xfh). +The final matrix at Foundry +`a8e43442421a80fea7991ee26f97f1003b170e29` is also `VERIFIED`: +[`ap-6fYcSUFLOxocBSjNxakNGZ`](https://modal.com/apps/modal-labs/rahul-dev/ap-6fYcSUFLOxocBSjNxakNGZ). +Both used SGLang `a25164bda355250015c0ba6d8821f0d4b37abe7d`, +model revision `b968826d9c46dd6066d109eabc6255188de91218`, and two H100s. + +The final run recorded one graph per rank at offset `68908220416`; each rank's +graph had 73 symmetric all-reduce and one symmetric all-gather KernelNode and +zero NCCL collective nodes. SAVE/SAVE2 semantic fingerprints reproduced; +LOAD/LOAD2 left every archive byte unchanged. Both restored processes replayed +batch sizes 1 and 8. All phase processes exited 0, cleanup completed, and no +archive was retained. Time to health was 217.2 s baseline, 103.0 s SAVE, +103.1 s SAVE2, 61.0 s LOAD, and 61.0 s LOAD2. + +The exact sequential outputs for baseline, SAVE2, LOAD, and LOAD2 were: + +```json +[ + " Tensor parallelism is a technique used in distributed computing to \nparallelize the computation of tensors across multiple devices, such as GPUs or \nTPUs, to improve the efficiency and scalability of deep learning models. It \ninvolves splitting the computation of a tensor, which is a multi-dimensional \narray, into smaller sub-tensors that can be processed", + " (5-7-5 syllables)\nSilent power cores, \nCycles of light, heat, and \nmight— \nDigital stars ignite. \n\nLet me know if you'd like a different \nstyle or theme! \n**Explanation of syllables:** \n- Line 1: \"Silent power \ncores\" (", + " Also, list three composite numbers greater than 50. Then, list three \nprime numbers less than 50. Finally, list three composite numbers less than 50. \nAdditionally, list three prime numbers between 50 and 100, and three composite \nnumbers between 50 and 10", + " How do they work? What are the benefits? What are the limitations? \nWhat are the use cases? What are the trade-offs? What are the best practices? \nWhat are the challenges? What are the future directions? What are the key \ntakeaways? What are the common mistakes? What are the performance \nconsiderations? What" +] ``` -An editable install (`pip install -e .`) is the supported path — the recipe serve scripts set no `PYTHONPATH` and rely on `foundry` being importable. To run straight from a source checkout, export `PYTHONPATH=…/foundry/python:…/sglang/python` yourself. +The exact concurrent output in every one of the eight batch-8 slots was +` ?\n\n2 + `. -### 2. Write a TOML config +This evidence is scoped only to the pinned single-host Qwen3-8B TP=2 symmetric +profile. It does not establish general/upstream TP or plain-NCCL support. +TP>2, multi-host, custom/NCCL-symmetric/fused collectives, CUMEM, NVLS, EP, +DP-attention, speculative decode, LoRA, pdmux, unpinned revisions, hybrid MoE, +Ampere, and Blackwell remain unsupported or unvalidated. -```toml -# recipe/sglang/foundry_save.toml -mode = "save" -base_addr = 0x600000000000 -region_size = "256GB" -workspace_root = "foundry_archive" -scratch_space_size = "1024MB" -``` +### Historical branch evidence (pre-port fork) -The matching `foundry_load.toml` just changes `mode = "load"`. See [`memory-lifecycle.md`](memory-lifecycle.md) for what each field controls. +Draft [PR #4](https://github.com/modal-projects/foundry/pull/4) on 2×H100 and +[PR #2](https://github.com/modal-projects/foundry/pull/2) on 2×H200 predate the +zero-edit plugin. They are **pre-port fork evidence only** and do not broaden +the validated profile above. General TP remains open in +[issue #6](https://github.com/foundry-org/foundry/issues/6). -### 3. Run SAVE, then LOAD +## How to use (dense TP=2) + +### 1. Install foundry and pinned SGLang ```bash -# SAVE -rm -rf foundry_archive -bash recipe/sglang/serve_qwen3-mini.sh --save -# Wait for "Application startup complete", then SIGTERM the server. - -# LOAD -bash recipe/sglang/serve_qwen3-mini.sh --load - -# Query -curl -s http://0.0.0.0:12000/v1/completions \ - -H 'Content-Type: application/json' \ - -d '{"model":"Qwen/Qwen3-1.7B","prompt":"The capital of France is", - "max_tokens":12,"temperature":0}' -# → "Paris. The capital of the United States is Washington, D" +git clone https://github.com/sgl-project/sglang.git +git -C sglang checkout --detach a25164bda355250015c0ba6d8821f0d4b37abe7d +cd sglang && pip install -e 'python[all]' --no-build-isolation && cd - +pushd foundry && pip install -e . --no-build-isolation && popd ``` -A single SAVE pass is enough — SGLang doesn't run a profile-forward at startup, so there is no non-determinism that requires a second pass. +Installing `foundry` also installs the `foundry_sglang_plugin` entry point. -## What the integration does +### 2. Run baseline / SAVE / LOAD -SAVE: +```bash +cd recipe/experimental +rm -rf foundry_archive_sglang_tp + +# baseline (no Foundry config → activate() is a no-op) +CUDA_VISIBLE_DEVICES=0,1 bash serve_qwen3-8b_sglang_tp.sh 2 + +# SAVE, then LOAD +CUDA_VISIBLE_DEVICES=0,1 bash serve_qwen3-8b_sglang_tp.sh 2 --save +CUDA_VISIBLE_DEVICES=0,1 bash serve_qwen3-8b_sglang_tp.sh 2 --load -1. `setup_graph_extension(...)` reserves a VMM region at `base_addr` and creates the per-rank workspace. -2. Distributed init / NCCL warmup runs in scratch space; the cursor is then forced to `scratch_space_size`. -3. Model weights, KV pool, and FlashInfer workspace buffers allocate inside the VMM region at byte-deterministic offsets. -4. `kernel_warmup` is a no-op. -5. `CudaGraphRunner.capture` runs a pre-pass that pre-allocates every per-bs FlashInfer wrapper, then enters the upstream capture loop with an idempotent inner-init shim (`reuse_pre_pass_init`) and a wrapper on `forward` that suppresses the two pre-capture warmup forwards. -6. Each captured graph is written to disk; a manifest groups topologically equivalent graphs. -7. The final VMM cursor is recorded as `final_alloc_offset`. +curl -s http://0.0.0.0:12000/generate -H 'Content-Type: application/json' \ + -d '{"text":"The capital of France is", + "sampling_params":{"temperature":0,"max_new_tokens":12}}' +``` + +`--cuda-graph-config` drives the decode `capture_bs` list. A single SAVE pass is +enough — SGLang does not run a profile-forward at startup — but the acceptance +harness runs a second SAVE (`save2`) to prove per-rank offset and archive +reproducibility. -LOAD: +## What the integration does -1. `setup_graph_extension(...)` restores the VMM region and replays captured fatbins into device code memory. -2. Distributed init runs as usual; the cursor advances to the same `scratch_space_size`. -3. Model weights and KV pool re-allocate at the same deterministic offsets. `init_memory_pool` reuses the saved `MemoryPoolConfig` (and calls `torch.cuda.empty_cache()` to mirror SAVE's `_resolve_memory_pool_config` side effect). -4. `CudaGraphRunner.capture` is replaced with: preallocate the entire deterministic range up to `final_alloc_offset`; run the same pre-pass init; call `start_graph_builds(all_paths) + finish_graph_loads(pending)` exactly once. All N graphs are loaded in one shot so the manifest's template/on-demand linking works. -5. `self.graphs` / `self.output_buffers` are populated from `state.loaded_graphs`; the rest of SGLang's serving path runs unchanged. +### SAVE + +1. `ServerArgs.__post_init__` validation hook fails closed on the dense TP=2 + topology/backend/settings envelope; SHA pinning remains Modal provenance. +2. `init_torch_distributed` hook reserves the VMM region on the rank's device + (`ps.tp_rank`), runs upstream NCCL bring-up in scratch, then forces the + cursor to `scratch_space_size`. +3. `alloc_memory_pool` hook lets upstream allocate, then rank 0 serializes the + resolved `MemoryPoolConfig` to `warmup_state.json`. +4. The REPLACE'd `FullCudaGraphBackend.capture_one` captures each decode graph + into a Foundry `CUDAGraph` (skipping the two warmups) and writes it, keyed by + `ShapeKey`. +5. The AROUND `DecodeCudaGraphRunner.capture` runs the upstream loop, then + writes the manifest, packs fatbins, and records `final_alloc_offset`. + +### LOAD + +1–2. Same activation and VMM/rank setup; the region also replays cached fatbins. +3. `_resolve_memory_pool_config` hook reconstructs the saved `MemoryPoolConfig`, + skips profiling and the original resolver, and mirrors + `torch.cuda.empty_cache()`; the normal target flow derives pool sizes from it. +4. The AROUND `DecodeCudaGraphRunner.capture` delegates to the same pinned + upstream lifecycle as SAVE. The backend's `capture_session` establishes the + shared pool, preallocates to `final_alloc_offset`, scans the archive, and + starts one pooled build. Upstream then performs its real warmup, + `graph_capture`, `capture_prepare`, `ForwardContext`, and FlashInfer metadata + setup before backend `capture_one` finishes the matching archived graph. + LOAD performs no native graph capture or save. ## Doc set - [`overview.md`](overview.md) — this file -- [`direct-edits.md`](direct-edits.md) — every line that changed in `sglang/` -- [`hooks.md`](hooks.md) — every monkey-patch foundry installs and what it does on SAVE / LOAD +- [`direct-edits.md`](direct-edits.md) — the zero-upstream-edit plugin activation path +- [`hooks.md`](hooks.md) — the nine `HookRegistry` targets and their SAVE / LOAD behavior - [`memory-lifecycle.md`](memory-lifecycle.md) — VMM region setup, the SAVE↔LOAD allocation parity contract, and the `final_alloc_offset` watermark -- [`save-load-workflow.md`](save-load-workflow.md) — serve scripts, TOML schema, expected logs, validation checks -- [`memory-consistency.md`](memory-consistency.md) — the five known divergences that caused silent LOAD failures and how each was fixed +- [`save-load-workflow.md`](save-load-workflow.md) — serve script, TOML schema, expected logs, validation checks +- [`memory-consistency.md`](memory-consistency.md) — the known divergences that caused silent LOAD failures and how each is handled diff --git a/docs/sglang/save-load-workflow.md b/docs/sglang/save-load-workflow.md index 78be4b49..ff88eec3 100644 --- a/docs/sglang/save-load-workflow.md +++ b/docs/sglang/save-load-workflow.md @@ -1,146 +1,201 @@ # SGLang Save / Load Workflow -## Serve scripts +This describes the pinned-main dense TP=2 recipe +(`recipe/experimental/serve_qwen3-8b_sglang_tp.sh`). Activation is through the +`sglang.srt.plugins` entry point with zero upstream edits; there is no +`--foundry-*` CLI flag. -Key knobs (1.7B example): +## Serve script knobs ```bash -MODEL_NAME="Qwen/Qwen3-1.7B" +MODEL_NAME="Qwen/Qwen3-8B" HOST="0.0.0.0" PORT=12000 -MEM_FRACTION_STATIC=0.6 - -# No LD_PRELOAD / PYTHONPATH here: foundry + the sglang fork are pip-installed, and -# foundry's setup_ld_preload_env auto-detects libcuda_hook.so and LD_PRELOADs it -# into every worker at spawn time. (Running from a source checkout instead? Export -# PYTHONPATH=.../foundry/python:.../sglang/python yourself.) +MEM_FRACTION_STATIC=0.8 +RANDOM_SEED=42 + +# Activation: plugin selector + config gate. No LD_PRELOAD / PYTHONPATH here — +# foundry is pip-installed and its Engine._launch_scheduler_processes hook sets +# LD_PRELOAD for every scheduler child at spawn time. +export SGLANG_PLUGINS=foundry +export SGL_MODEL_REVISION=b968826d9c46dd6066d109eabc6255188de91218 +export NCCL_CUMEM_ENABLE=0 +export NCCL_NVLS_ENABLE=0 +export PYTORCH_CUDA_ALLOC_CONF=expandable_segments:False +export SGLANG_ENABLE_POST_CAPTURE_KV_SIZING=0 + +unset FOUNDRY_GRAPH_EXTENSION_CONFIG_PATH # baseline +# export FOUNDRY_GRAPH_EXTENSION_CONFIG_PATH=…/sglang_foundry_tp_save.toml # SAVE +# export FOUNDRY_GRAPH_EXTENSION_CONFIG_PATH=…/sglang_foundry_tp_load.toml # LOAD sglang serve \ --model-path "$MODEL_NAME" \ + --revision "$SGL_MODEL_REVISION" \ --trust-remote-code \ --host "$HOST" --port "$PORT" \ - --tp-size 1 \ + --tp-size 2 \ + --random-seed "$RANDOM_SEED" \ --mem-fraction-static "$MEM_FRACTION_STATIC" \ --disable-radix-cache \ + --disable-custom-all-reduce \ + --enable-torch-symm-mem \ + --enforce-disable-flashinfer-allreduce-fusion \ + --disable-flashinfer-autotune \ --attention-backend flashinfer \ - --cuda-graph-max-bs 512 \ - --foundry-graph-extension-config-path "$FOUNDRY_TOML" + --decode-log-interval 1 \ + --cuda-graph-config '{"decode":{"backend":"full","max_bs":256,"bs":[256]},"prefill":{"backend":"disabled"}}' ``` -`--cuda-graph-max-bs 512` is the closest analogue to vLLM's `--max-num-seqs 512` — it drives `capture_bs` to span a similar range of decode batch sizes (52 batch sizes from 1 → 512). +The graph config permits exactly one padded decode shape, batch size 256. +Runtime batch sizes 1 and 8 both replay that graph. ## TOML configs -`recipe/sglang/foundry_save.toml`: +`recipe/experimental/sglang_foundry_tp_save.toml`: ```toml mode = "save" base_addr = 0x600000000000 region_size = "256GB" -workspace_root = "foundry_archive" +workspace_root = "foundry_archive_sglang_tp" scratch_space_size = "1024MB" ``` -`foundry_load.toml` is identical except `mode = "load"`. See [`memory-lifecycle.md`](memory-lifecycle.md) for the field semantics. +`sglang_foundry_tp_load.toml` is identical except `mode = "load"`. See +[`memory-lifecycle.md`](memory-lifecycle.md) for the field semantics. Dense TP +needs no NVSHMEM (no DeepEP), so `nvshmem_host_path` is left unset. ## Run sequence ```bash -cd /path/to/foundry/recipe/sglang - -# 1. Clean any prior workspace -rm -rf foundry_archive - -# 2. SAVE -bash serve_qwen3-mini.sh --save -# wait for "Application startup complete", then Ctrl-C / SIGTERM +cd recipe/experimental +rm -rf foundry_archive_sglang_tp -# 3. LOAD -bash serve_qwen3-mini.sh --load -# wait for "Application startup complete" +# baseline (activate() is a no-op without FOUNDRY_GRAPH_EXTENSION_CONFIG_PATH) +CUDA_VISIBLE_DEVICES=0,1 bash serve_qwen3-8b_sglang_tp.sh 2 +# wait for "Application startup complete", query, then Ctrl-C / SIGINT -# 4. Query -curl -s http://0.0.0.0:12000/v1/completions \ - -H 'Content-Type: application/json' \ - -d '{"model":"Qwen/Qwen3-1.7B","prompt":"The capital of France is", - "max_tokens":12,"temperature":0}' -``` +# SAVE +CUDA_VISIBLE_DEVICES=0,1 bash serve_qwen3-8b_sglang_tp.sh 2 --save +# wait for "Application startup complete", then Ctrl-C / SIGINT -For Qwen3-1.7B, expected output: +# LOAD +CUDA_VISIBLE_DEVICES=0,1 bash serve_qwen3-8b_sglang_tp.sh 2 --load +# Query +curl -s http://0.0.0.0:12000/generate -H 'Content-Type: application/json' \ + -d '{"text":"The capital of France is", + "sampling_params":{"temperature":0,"max_new_tokens":12}}' ``` -" Paris. The capital of the United States is Washington, D" -``` - -A single SAVE pass is sufficient. SGLang does not run a profile-forward at startup (only `_profile_available_bytes` samples free GPU memory), so there is no allocation non-determinism that a second SAVE pass would need to mask. ## Per-rank workspace layout ``` -foundry_archive/ - warmup_state.json # shared; MemoryPoolConfig + final_alloc_offset +foundry_archive_sglang_tp/ + warmup_state.json # shared; MemoryPoolConfig + final_alloc_offset (rank 0) rank_0/ - graph_{0..N-1}_FULL_t{bs}_r{bs}_UX_pcN.json # one per captured graph - graph_{0..N-1}_FULL_t{bs}_r{bs}_UX_pcN.cugraph # binary cuGraph blob + graph_0_FULL_t256_r256_UX_pcN.json # sole captured graph + graph_0_FULL_t256_r256_UX_pcN.cugraph # binary cuGraph blob graph_manifest.json # topology groups for template + on-demand linking final_alloc_offset.json # per-rank VMM watermark fatbin_image_packed.img # packed kernel fatbins - fatbin_entrypoint_packed.txt # fatbin entry-point index + rank_1/ + … ``` -Filename scheme (defined in `graph_ops._graph_filename`): +Filename scheme (`graph_ops._graph_filename`), keyed by the dense `ShapeKey`: ``` -graph_{state.capture_index}_FULL_t{bs}_r{bs}_UX_pcN.json +graph_{state.capture_index}_FULL_t{size}_r{size}_UX_pcN.json ``` -`state.capture_index` increments per `save_graph` call so files sort in SAVE-time order. `_GRAPH_FILENAME_RE` in `graph_ops.py` parses them on LOAD. +`state.capture_index` increments per `save_graph` call so files sort in +SAVE-time order; `_GRAPH_FILENAME_RE` parses `index` and `bs` on LOAD to +reconstruct a `ShapeKey(size=bs)`. ## Expected logs SAVE (success): ``` -[Foundry] SGLang hooks installing: mode=save workspace=foundry_archive_qwen_1.7b -[Foundry] SGLang hooks installed -Foundry SGLang integration activated from .../save_qwen_1.7b.toml -[Foundry] SGLang rank=0 workspace_dir=foundry_archive_qwen_1.7b/rank_0 +[Foundry] SGLang hooks registering: mode=save workspace=foundry_archive_sglang_tp +[Foundry] SGLang hooks registered +[Foundry] SGLang rank=0 workspace_dir=foundry_archive_sglang_tp/rank_0 [Foundry] SGLang graph extension setup completed in 0.x s +[Foundry] SGLang alloc_offset[after_setup_graph_ext]=… (… MB) +[Foundry] SGLang alloc_offset[after_init_torch_dist]=… (… MB) [Foundry] SGLang skipped allocator to scratch boundary 1073741824 -[Foundry] SGLang reused saved memory pool config ← only on LOAD -[Foundry] SGLang kernel_warmup skipped in save mode -… -[Foundry] Saved SGLang CUDA graph graph_0_FULL_t512_r512_UX_pcN.json key=512 -… -[Foundry] Saved SGLang CUDA graph graph_51_FULL_t1_r1_UX_pcN.json key=1 -[foundry] Saved graph_manifest.json: 9 topology groups (9 templates, 43 on-demand, 43 stripped) -[Foundry] Saved SGLang warmup state to foundry_archive_qwen_1.7b/warmup_state.json -[Foundry] SGLang final_alloc_offset=22785556480 -… +[Foundry] SGLang alloc_offset[after_scratch_skip]=… (… MB) +[Foundry] Torch symmetric-memory buffer address=0x… +[Foundry] Multimem all-gather buffer address=0x… +[Foundry] SGLang alloc_offset[before_alloc_memory_pool]=… (… MB) +[Foundry] SGLang alloc_offset[after_alloc_memory_pool]=… (… MB) +[Foundry] Saved SGLang CUDA graph graph_0_FULL_t256_r256_UX_pcN.json key=ShapeKey(size=256, stream_idx=None, variant_label=None) +[Foundry] Saved SGLang warmup state to foundry_archive_sglang_tp/warmup_state.json +[Foundry] SGLang final_alloc_offset=68908220416 INFO: Application startup complete. ``` LOAD (success): ``` -[Foundry] SGLang hooks installing: mode=load workspace=foundry_archive_qwen_1.7b -[Foundry] SGLang hooks installed +[Foundry] SGLang hooks registering: mode=load workspace=foundry_archive_sglang_tp +[Foundry] SGLang hooks registered [Foundry] SGLang graph extension setup completed in 0.x s [Foundry] SGLang reused saved memory pool config -[Foundry] SGLang kernel_warmup skipped in load mode +[Foundry] Torch symmetric-memory buffer address=0x… +[Foundry] Multimem all-gather buffer address=0x… [Foundry] SGLang alloc_offset[before_preallocate]=… (… MB) [Foundry] SGLang alloc_offset[after_preallocate]=… (… MB) -[Foundry] SGLang alloc_offset[after_pre_init]=… (… MB) -[CGE] Using graph_manifest.json (9 topology groups) -[CGE] Phase 1: 52 graphs parsed in 0.x ms, 9 topologies, 4 threads, 52 binary + 0 json -[CGE BUILD] Template 0 (...): N nodes, done in X.X ms -… -[CGE] Phase 2: 9 templates + 43 on-demand = 52 graphs built in xx.x ms -[CGE] finish_graph_loads: 52 graphs, xx.x ms -[Foundry] Loaded 52 SGLang graphs in 0.0x s -[Foundry] SGLang alloc_offset[after_load_all_graphs]=22785556480 (… MB) -… +[Foundry] Started 1 SGLang graph builds +[Foundry] Loaded SGLang graph graph_0_FULL_t256_r256_UX_pcN.json +[Foundry] SGLang alloc_offset[after_load_all_graphs]=68908220416 (… MB) INFO: Application startup complete. ``` -The `after_load_all_graphs` value **must** equal SAVE's `final_alloc_offset`. If it doesn't, see [`memory-consistency.md`](memory-consistency.md). +The `after_load_all_graphs` value **must** equal SAVE's `final_alloc_offset`. +LOAD uses the same pinned upstream `capture()` lifecycle as SAVE. Its backend +starts one pooled graph-build batch, then finishes the archived shape when +upstream reaches the corresponding `capture_one`; it does not capture or save. + +## Automated acceptance + +`tests/modal_sglang_tp.py` runs the full baseline → SAVE → SAVE2 → LOAD → +LOAD2 matrix on 2×H100. The final run is `VERIFIED`: +[`ap-6fYcSUFLOxocBSjNxakNGZ`](https://modal.com/apps/modal-labs/rahul-dev/ap-6fYcSUFLOxocBSjNxakNGZ). +It used Foundry `a8e43442421a80fea7991ee26f97f1003b170e29`, SGLang +`a25164bda355250015c0ba6d8821f0d4b37abe7d`, and Qwen3-8B revision +`b968826d9c46dd6066d109eabc6255188de91218`. + +These exact engine/model SHAs are validation provenance enforced by the Modal +harness, not fields checked by `_validate_profile`. Shell recipe users must +install the matching SGLang pin; runtime fail-closed checks cover topology, +backend, and settings. + +Evidence was one graph per rank; offset `68908220416` in SAVE/SAVE2/LOAD/LOAD2; +73 symmetric all-reduce plus one symmetric all-gather and zero NCCL collective +nodes per rank; exact sequential outputs and eight exact ` ?\n\n2 + ` batch-8 +outputs in every query phase; reproducible SAVE/SAVE2 semantic fingerprints; +byte-immutable LOAD/LOAD2; observed batch-1 and batch-8 replay; and process +return code 0 in every phase. Time to health was 217.2, 103.0, 103.1, 61.0, +and 61.0 seconds respectively. Cleanup completed and retained no artifacts. + +The prerequisite focused SAVE→LOAD was separately `VERIFIED` at +[`ap-5wO9cFhRZ7ITsLDDv27Xfh`](https://modal.com/apps/modal-labs/rahul-dev/ap-5wO9cFhRZ7ITsLDDv27Xfh). + +## Rejected plain-NCCL path and support boundary + +Plain-NCCL fresh-process SGLang restore is unsupported. App +[`ap-2qSQUKlY24L4RrQkMpUrmc`](https://modal.com/apps/modal-labs/rahul-dev/ap-2qSQUKlY24L4RrQkMpUrmc) +restored 36 graphs per rank with matching counts, offsets, immutable bytes, and +P2P/IPC topology, but both fresh LOADs emitted the same corrupt decode while +baseline and SAVE2 matched. Archived process-specific communicator and +graph-sequence state is the leading forensic explanation, not +treatment-proven causality. The deterministic corrupt output rejects this +plain-NCCL path; structural pass alone is not output correctness. + +The `VERIFIED` claim is only for the exact pinned single-host Qwen3-8B TP=2 +torch symmetric-memory profile. It is not general/upstream TP or plain-NCCL +support. TP>2, multi-host, custom/NCCL-symmetric/fused collectives, CUMEM, NVLS, +EP, DP-attention, speculative decode, LoRA, pdmux, and unpinned revisions are +unsupported. diff --git a/docs/superpowers/plans/2026-07-23-final-branch-deltas.md b/docs/superpowers/plans/2026-07-23-final-branch-deltas.md new file mode 100644 index 00000000..4f068560 --- /dev/null +++ b/docs/superpowers/plans/2026-07-23-final-branch-deltas.md @@ -0,0 +1,1234 @@ +# Final Branch Deltas Implementation Plan + +> **For agentic workers:** REQUIRED SUB-SKILL: Use +> superpowers:subagent-driven-development or superpowers:executing-plans to +> implement this plan task-by-task. Do not dispatch fast models. Steps use +> checkbox (`- [ ]`) syntax for tracking. + +**Goal:** Selectively implement the remaining audited VMM-IPC, vLLM DeepEP, +DeepEP replay, and vLLM TP acceptance deltas without importing superseded +branch work or widening the verified SGLang profile. + +**Architecture:** Keep the current branch as a linear, merge-free delta from +`origin/main`. Manually re-derive each required behavior against the current +tree, preserving its stronger generation, duplicate-open, current-main SGLang, +and validation infrastructure. Publish archives transactionally, make archived +transport and graph inventory explicit contracts, and accept engine support +only from final-HEAD fresh-process H100 evidence. + +**Tech Stack:** C++17, CUDA Driver VMM/IPC, Python 3.12, PyTorch 2.11/cu130, +pytest, vLLM V1, DeepEP/NVSHMEM, Bash, Modal H100, pre-commit. + +## Global Constraints + +- Start from audited HEAD + `c237554eecc18632ca6938c84571cd0b08953e07` on + `cursor/tp-upstream-consolidation-fe33`. +- Before every implementation task, compare the cited source commit with the + current destination symbols. If either changed, stop that task, document the + discrepancy, and update the task before coding. +- The branch must remain a direct-main lineage. At the audited state, + `merge-base(HEAD, origin/main)` is + `21824649e711487a2c6af01cb631b3e41ccc32d8`, + `origin/main...HEAD` is `0 114`, and `origin/main..HEAD` has zero merge + commits. If `origin/main` advances, reconcile and rebase only with explicit + approval; never merge a feature branch. +- Review with `git diff origin/main...HEAD`. Never use a wholesale branch + cherry-pick, branch merge, or broad final-tree transplant. Source commits are + behavioral references, not integration commands. +- Preserve current VMM alignment normalization from + `0f337a175a36b8204637b3c825c2c70806413dae` and + `f947886f34fd44925fea3084faef53eb13f14929`, generation identity from + `ab22a213ba7d9643a078a076060255247a25295a`, and duplicate interior-pointer + ownership from `eaf3fea5090410f62550e6ace7803d36fb245479`. +- Preserve the verified SGLang Qwen3-8B TP=2 symmetric-memory profile and its + zero-edit HookRegistry architecture. Its evidence anchor is Foundry + `435518206f501b2f6b20080ee53375d0d6a3ff76`; later documentation commits do + not widen the profile. +- Do not port old SGLang `_patch_deepep`, monolithic `CudaGraphRunner`, direct + fork edits, hybrid warmup experiments, graph catalogs, multi-session + backends, TP2/DP2, EP4, or Qwen3.5 claims. +- Quarantine the incompatible SGLang EP recipe. Do not add an SGLang DeepEP + target, archive schema, or support claim in this plan. +- Do not use fast models for implementation, review, or validation. +- Keep optional vLLM/DeepEP imports lazy only where importing them at module + load would break core-only installs; all standard-library and Foundry imports + stay at module scope. +- Run CUDA work only on H100s in Modal environment `rahul-dev`. The local VM + has no CUDA device or toolkit. +- The implementation dependency order is Task 1 → Task 2 → Task 3 → Task 4 → + Task 5 → Task 6 → Task 7 → Task 8 → Task 9. +- Tasks 1-7 require focused commits and reviewer approval. Do not push a GREEN + implementation series until all local contracts through Task 7 pass. A + reviewed test-only RED commit may be pushed when a Modal run is the only way + to capture the failure. +- Task 8 must run from a pushed immutable full SHA. Task 9 documentation and PR + status may claim validation only from that exact SHA. + +## File Responsibility Map + +| Area | Files | Responsibility | +|---|---|---| +| VMM IPC | `csrc/hook.cpp`, `include/foundry/vmm_ipc_reservation.h`, `CMakeLists.txt` | Exportable pitched allocations, process ownership, exact peer-zone reservation, cleanup | +| Native regressions | `tests/test_vmm_alloc.py`, `tests/vmm_ipc_reservation_test.cpp` | Pitched/fork/wrong-VA plus generation, duplicate-open, and stale-export contracts | +| DeepEP policy | `python/foundry/integration/deepep_transport.py` | Shared transport enum, buffer policy, and archive mismatch validation | +| vLLM DeepEP | `python/foundry/integration/vllm/{config,runtime,hooks}.py` | TOML policy, warmup-state schema, early LOAD rejection, first-SAVE metadata, NVSHMEM detection | +| No-IMEX recipes | `recipe/experimental/*ipc*`, `recipe/experimental/foundry_{save,load}*.toml`, `recipe/experimental/README.md` | Archive-defined vLLM no-IMEX launch and explicit SGLang EP quarantine | +| DeepEP replay | `tests/test_deepep_fabric.py`, `tests/run_deepep_matrix.sh`, `tests/modal_deepep.py` | Strict SAVE/LOAD output oracle for fabric, NVL IPC, and preallocation | +| vLLM TP evidence | `tests/tp_validation_helpers.py`, `tests/vllm_tp_validation.py`, `tests/modal_vllm_tp.py` | Exact inventory, immutable transaction, live restore, output, provenance, cleanup | +| CPU contracts | `tests/test_{deepep_transport,deepep_transport_vllm,experimental_recipe,tp_validation_helpers,vllm_tp_recipe,vllm_tp_validation,modal_vllm_tp,modal_deepep}.py` | RED/GREEN contracts without requiring local CUDA | +| Evidence docs | `README.md`, `ROADMAP.md`, `RELEASE.md`, `docs/vllm/{overview,moe-and-deepep}.md`, `recipe/vllm/README.md` | Final-SHA scope and measured H100 evidence only | + +--- + +### Task 1: Harden Native VMM-IPC Ownership and Exact Reservation + +**Files:** + +- Create: `include/foundry/vmm_ipc_reservation.h` +- Create: `tests/vmm_ipc_reservation_test.cpp` +- Modify: `CMakeLists.txt` +- Modify: `csrc/hook.cpp` +- Modify: `tests/test_vmm_alloc.py` + +**Symbols and sources:** + +- Port behavior manually from + `fce0ab2fdcddcd57f0ef868a41d76a50171cdb8d`. +- Port test stabilization from + `13eb6897701f26b449edb6d0a9c8a666644056a9`. +- Modify `cuMemAllocPitch_v2`, `vmm_ipc_ensure_server`, + `vmm_ipc_invalidate_export`, `cuIpcGetMemHandle`, + `cuIpcOpenMemHandle`, and `cuIpcCloseMemHandle`. +- Add `vmm_ipc_owner_pid`, `vmm_ipc_process_matches_owner`, and + `vmm_ipc_claim_process`. +- Preserve `VmmIpcChunkKey` generation and `VmmIpcChunkSubptr::open_count`. + +**Interfaces:** + +- Consumes: current VMM2 handles, SCM_RIGHTS transport, generation-aware + preallocation mappings, and the import-zone cursor. +- Produces: + `require_exact_ipc_reservation(result, actual, expected, size, handle, + free_reservation, release_handle) -> CUresult`, plus fork-safe IPC entry + points returning `CUDA_ERROR_NOT_SUPPORTED`. +- Error contract: a post-claim fork child emits `use spawn/exec` and returns + error 801 before touching inherited mutexes, maps, descriptors, CUDA handles, + or server state. + +- [ ] **Step 1: Reconcile the source commit with the current tree** + +Run: + +```bash +git diff --no-ext-diff fce0ab2fdcddcd57f0ef868a41d76a50171cdb8d^ \ + fce0ab2fdcddcd57f0ef868a41d76a50171cdb8d -- \ + csrc/hook.cpp tests/test_vmm_alloc.py +git blame -L 2570,2578 csrc/hook.cpp +git blame -L 2909,2968 csrc/hook.cpp +git blame -L 3028,3055 csrc/hook.cpp +git blame -L 3418,3440 csrc/hook.cpp +git blame -L 3563,3595 csrc/hook.cpp +``` + +Expected: the pitched path still omits `requestedHandleTypes`; fork handling +still attempts child reinitialization; the second zone hint is not checked; +generation and `open_count` remain current and must not be replaced. + +- [ ] **Step 2: Write RED contracts** + +Add the stabilized `_run_core_pitched_ipc_export` and +`_run_core_rejects_vmm_ipc_after_fork` cases to `tests/test_vmm_alloc.py`. +Both must call `torch.cuda.set_device(0)` and use base +`0x600000000000`. + +Define the pure cleanup policy in +`include/foundry/vmm_ipc_reservation.h` with this interface: + +```cpp +namespace foundry::vmm_ipc { +template +CUresult require_exact_ipc_reservation( + CUresult result, + CUdeviceptr actual, + CUdeviceptr expected, + size_t size, + CUmemGenericAllocationHandle handle, + FreeReservation free_reservation, + ReleaseHandle release_handle); +} +``` + +`tests/vmm_ipc_reservation_test.cpp` must inject three cases: + +1. success at the expected VA: return success, call neither callback; +2. driver error: preserve that error, call neither callback; +3. success at a different VA: call address-free and handle-release exactly + once with the actual reservation and imported handle, then return + `CUDA_ERROR_OUT_OF_MEMORY`. + +Wire the executable under `if(BUILD_TESTING)` and `add_test` in +`CMakeLists.txt`. + +- [ ] **Step 3: Run RED** + +On an H100 checkout containing only the tests: + +```bash +CC=gcc CXX=g++ pip install -e . --no-build-isolation +cmake -S . -B build -G Ninja -DBUILD_TESTING=ON +cmake --build build --target vmm_ipc_reservation_test +python3 -m pytest \ + tests/test_vmm_alloc.py::test_pitched_vmm_allocation_can_be_exported \ + tests/test_vmm_alloc.py::test_vmm_ipc_rejects_use_after_fork -q +``` + +Expected: pitched export fails and post-fork use either times out or violates +the required error. The pure reservation test initially fails to link until +the policy is implemented. + +- [ ] **Step 4: Implement the minimal production behavior** + +In the pitched slow path add: + +```cpp +prop.requestedHandleTypes = CU_MEM_HANDLE_TYPE_POSIX_FILE_DESCRIPTOR; +``` + +Claim one process owner atomically before server initialization and VMM-handle +open. Check that owner before get, close, and export invalidation. Delete the +fork-child cleanup/rebind path: inherited state must never be locked or +mutated. + +Call `require_exact_ipc_reservation` immediately after the dedicated-zone +`cuMemAddressReserve`. Do not call `cuMemMap` unless the reservation exactly +equals `zone_hint`. + +- [ ] **Step 5: Run GREEN** + +On H100: + +```bash +cmake -S . -B build -G Ninja -DBUILD_TESTING=ON +cmake --build build --target vmm_ipc_reservation_test +ctest --test-dir build --output-on-failure -R vmm_ipc_reservation +python3 -m pytest \ + tests/test_vmm_alloc.py::test_pitched_vmm_allocation_can_be_exported \ + tests/test_vmm_alloc.py::test_vmm_ipc_rejects_use_after_fork -q +``` + +Expected: one C++ policy test and both GPU subprocess regressions pass. + +- [ ] **Step 6: Commit and apply the review/push gate** + +```bash +git add CMakeLists.txt csrc/hook.cpp include/foundry/vmm_ipc_reservation.h \ + tests/test_vmm_alloc.py tests/vmm_ipc_reservation_test.cpp +git commit -m "fix: harden VMM IPC process and reservation ownership" +git show --check --stat HEAD +test "$(git merge-base HEAD origin/main)" = "$(git rev-parse origin/main)" +test "$(git rev-list --merges origin/main..HEAD --count)" = 0 +``` + +Reviewer must compare the focused diff with `fce0ab2` and confirm no alignment, +generation, token, or duplicate-open code regressed. Do not push GREEN until +Tasks 1-7 pass locally. + +### Task 2: Add Direct VMM-IPC Generation, Duplicate-Open, and Stale-Export Regressions + +**Files:** + +- Modify: `tests/test_vmm_alloc.py` + +**Symbols and sources:** + +- Exercise existing generation behavior from + `ab22a213ba7d9643a078a076060255247a25295a`. +- Exercise existing duplicate-open behavior from + `eaf3fea5090410f62550e6ace7803d36fb245479`. +- Exercise `vmm_ipc_invalidate_export` and the parked-FD fetch server. +- There is no branch test to cherry-pick; derive these regressions from the + audited current implementation. + +**Interfaces:** + +- Consumes: `CUipcMemHandle`, `cuIpcGetMemHandle`, + `cuIpcOpenMemHandle`, `cuIpcCloseMemHandle`, preallocation APIs, and + subprocess role synchronization. +- Produces: `_run_core_vmm_ipc_generation_reuse`, + `_run_core_vmm_ipc_duplicate_open_close`, and + `_run_core_vmm_ipc_stale_export` subprocess modes. + +- [ ] **Step 1: Write the three H100 regressions** + +Use spawned exporter/importer roles and filesystem control messages; never +fork after IPC ownership is claimed. + +The generation test must: + +1. preallocate, export, open, and close an interior pointer; +2. free the exporter chunk; +3. recreate a chunk at the same `0x600000000000` VA; +4. write a different sentinel value and export it; +5. open through the second handle and assert the new value is observed. + +The duplicate-open test must open the same interior pointer twice, assert both +opens are usable, close once, prove the second reference remains readable, +close again, and prove `cuMemGetAddressRange_v2` no longer recognizes the +mapping. + +The stale-export test must export a normal VMM allocation, free it while the +exporter process remains alive, then make the importer open the stale handle +and assert a non-success CUDA result without a hang. + +- [ ] **Step 2: Prove RED with targeted mutations** + +For each test, start from a clean `csrc/hook.cpp`, make exactly one local +mutation, save that mutation as a patch, run only its matching test on H100, +and reverse the mutation: + +```bash +# Mutation A: omit generation from VmmIpcChunkKey. +git diff -- csrc/hook.cpp > /tmp/vmm-ipc-mutation.patch +python3 -m pytest \ + tests/test_vmm_alloc.py::test_vmm_ipc_generation_changes_after_same_va_recreate -q +git apply -R /tmp/vmm-ipc-mutation.patch +test -z "$(git diff -- csrc/hook.cpp)" + +# Mutation B: erase VmmIpcChunkSubptr on the first close. +git diff -- csrc/hook.cpp > /tmp/vmm-ipc-mutation.patch +python3 -m pytest \ + tests/test_vmm_alloc.py::test_vmm_ipc_duplicate_open_is_refcounted -q +git apply -R /tmp/vmm-ipc-mutation.patch +test -z "$(git diff -- csrc/hook.cpp)" + +# Mutation C: skip vmm_ipc_invalidate_export from cuMemFree_v2. +git diff -- csrc/hook.cpp > /tmp/vmm-ipc-mutation.patch +python3 -m pytest \ + tests/test_vmm_alloc.py::test_vmm_ipc_stale_export_fails -q +git apply -R /tmp/vmm-ipc-mutation.patch +test -z "$(git diff -- csrc/hook.cpp)" +``` + +Expected: each isolated mutation fails its matching test and the source is +restored after every run. + +- [ ] **Step 3: Run GREEN against unmodified production code** + +```bash +python3 -m pytest \ + tests/test_vmm_alloc.py::test_vmm_ipc_generation_changes_after_same_va_recreate \ + tests/test_vmm_alloc.py::test_vmm_ipc_duplicate_open_is_refcounted \ + tests/test_vmm_alloc.py::test_vmm_ipc_stale_export_fails -q +``` + +Expected: all three pass in bounded subprocess time. + +- [ ] **Step 4: Commit and apply the review/push gate** + +```bash +git add tests/test_vmm_alloc.py +git commit -m "test: cover VMM IPC lifecycle ownership" +git show --check --stat HEAD +``` + +Reviewer must confirm all CUDA initialization occurs in subprocess children and +all waits have hard timeouts. No production source belongs in this commit. Do +not push GREEN until Tasks 1-7 pass locally. + +### Task 3: Add the Archive-Defined vLLM DeepEP Transport Contract + +**Files:** + +- Create: `python/foundry/integration/deepep_transport.py` +- Create: `tests/test_deepep_transport.py` +- Create: `tests/test_deepep_transport_vllm.py` +- Modify: `python/foundry/integration/vllm/config.py` +- Modify: `python/foundry/integration/vllm/runtime.py` +- Modify: `python/foundry/integration/vllm/hooks.py` + +**Symbols and sources:** + +- Policy and parsing: + `42875919c64597905cf10257250e7856afa6a7bb`, + `ccda6c5ab431986444646475940608738ace9dc7`. +- Archive validation: + `d74f290a0fff1af3970821da0e60a219e091856e`, + `d6165b75b5d383bba8e3e711240be095f10a5859`. +- vLLM hook application: + `db3555cae006c00e3798d1e88a5af9a26a1e4260`, + `b0234e585f1fdd4f185172cba2270b915ecd213e`. +- First-SAVE preservation: + `69efcd0e51b54c15ba217ea47a0337709bc200fd`, + `60477ff51103d5af679767179855f7237b36bebf`, + `efc8728ad83f94b30955de4545b57c461ca9db30`. +- NVSHMEM detection: + `384017d69c354acb54b02a05190ef47418f9d987`, + `f48ef04401eba7b9e3555a01a38cc91389ac1948`. +- Do not copy the SGLang portions of those commits. + +**Interfaces:** + +- `DeepEPTransport(str, enum.Enum)` has `FABRIC = "fabric"` and + `NVL_IPC = "nvl_ipc"`. +- `parse_deepep_transport(raw: str | None) -> DeepEPTransport` defaults to + fabric and rejects every other value. +- `configure_deepep_buffer(transport: DeepEPTransport, *, + num_nvl_bytes: int, kwargs: dict[str, Any]) -> + tuple[int, dict[str, Any]]` forces + `(use_fabric=True, num_nvl_bytes=0)` for fabric and + `(use_fabric=False, upstream nonzero num_nvl_bytes)` for NVL IPC. +- `validate_transport_match(*, archived: str | None, + configured: DeepEPTransport) -> None` and + `validate_num_nvl_bytes_match(*, archived: int, effective: int, + transport: DeepEPTransport) -> None` fail closed on LOAD mismatch. +- `WarmupState` adds `deepep_transport: str = "fabric"` and + `deepep_num_nvl_bytes: int = 0`. +- LOAD validates transport before module loading/region setup and validates + effective NVL bytes before communicator construction. + +- [ ] **Step 1: Write RED CPU contracts** + +Port only the shared-policy tests from `tests/test_deepep_transport.py`; remove +all SGLang imports and test classes. Port the vLLM tests for: + +- TOML default, `nvl_ipc`, and invalid values; +- explicit path precedence and wheel auto-detection of + `libnvshmem_host.so.3` then `libnvshmem_host.so`; +- legacy warmup-state defaults and new-field round trips; +- rank-zero-only bootstrap/update; +- transport rejection before CUDA operations; +- effective-byte rejection before communicator setup; +- NONE mode pass-through; +- first SAVE profiling despite a transport-only bootstrap file; +- metadata preservation after profiling. + +Run: + +```bash +python3 -m pytest \ + tests/test_deepep_transport.py \ + tests/test_deepep_transport_vllm.py -q +``` + +Expected: collection/import failures because the policy module and schema do +not exist. + +- [ ] **Step 2: Implement policy, config, and early archive validation** + +Add the exact policy interfaces listed above. Extend +`CUDAGraphExtensionConfig.from_toml` and `get_deepep_transport`. In +`setup_graph_extension`, call `validate_transport_match` before +`set_skip_fatbin_processing`, CUDA module loading, or VMM setup. + +`_detect_nvshmem_host_path` must inspect `nvidia.nvshmem` package roots and +return the first existing wheel library. An explicit TOML path always wins. + +- [ ] **Step 3: Implement hook application and first-SAVE preservation** + +Replace the `FOUNDRY_DEEPEP_NVL_IPC` branch in `_patch_deepep` with the parsed +policy. Rank zero stores the selected transport and effective size. A +transport-only bootstrap file must not make `_patch_initialize_kv_caches` +believe KV profiling already completed: + +```python +has_saved_kv_profile = ws is not None and bool(ws.available_gpu_memory) +``` + +When profiling creates a fresh runtime `WarmupState`, copy the two DeepEP +fields from an existing bootstrap state before saving all KV/profile metadata. + +- [ ] **Step 4: Run GREEN** + +```bash +python3 -m pytest \ + tests/test_deepep_transport.py \ + tests/test_deepep_transport_vllm.py -q +python3 -m compileall -q \ + python/foundry/integration/deepep_transport.py \ + python/foundry/integration/vllm +``` + +Expected: all focused tests pass and compilation exits zero. + +- [ ] **Step 5: Commit and apply the review/push gate** + +```bash +git add \ + python/foundry/integration/deepep_transport.py \ + python/foundry/integration/vllm/config.py \ + python/foundry/integration/vllm/runtime.py \ + python/foundry/integration/vllm/hooks.py \ + tests/test_deepep_transport.py \ + tests/test_deepep_transport_vllm.py +git commit -m "feat(vllm): archive DeepEP transport profile" +git show --check --stat HEAD +``` + +Reviewer must verify that no `python/foundry/integration/sglang/**` file changed, +the environment variable is gone from vLLM hooks, and mismatch checks precede +CUDA/communicator side effects. Do not push GREEN until Tasks 1-7 pass locally. + +### Task 4: Convert No-IMEX vLLM Recipes and Quarantine SGLang EP + +**Files:** + +- Modify: `recipe/experimental/foundry_save.toml` +- Modify: `recipe/experimental/foundry_load.toml` +- Modify: `recipe/experimental/foundry_save_fp8.toml` +- Modify: `recipe/experimental/foundry_load_fp8.toml` +- Modify: `recipe/experimental/serve_qwen3-30ba3b_ipc_ep.sh` +- Modify: `recipe/experimental/serve_qwen3-30ba3bfp8_ipc_ep.sh` +- Delete: `recipe/experimental/serve_qwen3-30ba3bfp8_sglang_ipc_ep.sh` +- Delete: `recipe/experimental/sglang_foundry_save.toml` +- Delete: `recipe/experimental/sglang_foundry_load.toml` +- Modify: `recipe/experimental/README.md` +- Modify: `docs/vllm/moe-and-deepep.md` +- Modify: `tests/test_experimental_recipe.py` + +**Symbols and sources:** + +- vLLM recipe semantics: + `f743855c37dc392589858b6df4696a1faebb96d0`. +- Recipe and documentation contracts: + `694f22fb4415822f249882ee78b461473f10d33b`, + `2ec7250390459528511029175284639813e464de`. +- Deliberate divergence: do not port the source branch's SGLang assets or + claims because current `_validate_profile` rejects EP and there is no + HookRegistry DeepEP target. + +**Interfaces:** + +- Every no-IMEX vLLM SAVE/LOAD TOML contains + `deepep_transport = "nvl_ipc"`. +- No no-IMEX serve script exports `FOUNDRY_DEEPEP_NVL_IPC`. +- The dense SGLang TP recipe and verified symmetric profile are unchanged. +- The repository contains no runnable SGLang EP IPC recipe after this task. + +- [ ] **Step 1: Write RED recipe contracts** + +Retain the existing path-portability and missing-NVSHMEM tests. Add assertions +that all four vLLM TOMLs select `nvl_ipc`, both vLLM scripts omit the +environment switch, the three incompatible SGLang EP files are absent, and the +README says the profile is quarantined because current HookRegistry supports +only the verified dense TP envelope. + +Run: + +```bash +python3 -m pytest tests/test_experimental_recipe.py -q +``` + +Expected: failures for missing TOML fields, present environment switches, and +present SGLang EP assets. + +- [ ] **Step 2: Update vLLM recipes and remove the incompatible assets** + +Add the TOML field to all four vLLM files. Remove only +`FOUNDRY_DEEPEP_NVL_IPC`; retain CUMEM/NVLS disablement and V1 selection. + +Delete the three SGLang EP assets. In `recipe/experimental/README.md`, replace +the runnable SGLang EP section with a short quarantine record: + +```markdown +## Quarantined SGLang EP profile + +The old SGLang DeepEP IPC recipe is not runnable on the current HookRegistry +integration. The verified SGLang profile remains dense Qwen3-8B TP=2 with +torch symmetric memory; EP and DeepEP are rejected before launch. Reintroduce +an EP recipe only with a separate current-main profile, HookRegistry target, +archive contract, and fresh SAVE-to-LOAD H100 evidence. +``` + +- [ ] **Step 3: Document the vLLM archive contract** + +In `docs/vllm/moe-and-deepep.md`, document both values, legacy fabric default, +effective NVL bytes, rank-zero archive fields, pre-CUDA mismatch rejection, +first-SAVE metadata preservation, wheel auto-detection, no silent fallback, +and the three strict DeepEP matrix commands: `fabric`, `nvl_ipc`, and +`nvl_ipc_prealloc`. + +- [ ] **Step 4: Run GREEN** + +```bash +python3 -m pytest tests/test_experimental_recipe.py -q +bash -n recipe/experimental/serve_qwen3-30ba3b_ipc_ep.sh +bash -n recipe/experimental/serve_qwen3-30ba3bfp8_ipc_ep.sh +rg 'FOUNDRY_DEEPEP_NVL_IPC' \ + python/foundry/integration/vllm recipe/experimental +``` + +Expected: tests and syntax pass; the final `rg` returns no matches. + +- [ ] **Step 5: Commit and apply the review/push gate** + +```bash +git add docs/vllm/moe-and-deepep.md recipe/experimental \ + tests/test_experimental_recipe.py +git commit -m "docs(vllm): define no-IMEX DeepEP transport" +git show --check --stat HEAD +``` + +Reviewer must confirm the deletion is limited to SGLang EP, not the verified +`serve_qwen3-8b_sglang_tp.sh` profile. Do not push GREEN until Tasks 1-7 pass +locally. + +### Task 5: Make DeepEP Replay Output a Strict SAVE/LOAD Oracle + +**Files:** + +- Modify: `tests/test_deepep_fabric.py` +- Modify: `tests/run_deepep_matrix.sh` +- Create: `tests/modal_deepep.py` +- Create: `tests/test_modal_deepep.py` + +**Symbols and sources:** + +- Strict tuple/output/sentinel behavior: + `9842f2fd76216a7be9cc2e2e14a765fb97131c9d`. +- Add/retain the strict cases `fabric` (`use_fabric=True`, + `num_nvl_bytes=0`), `nvl_ipc` (`use_fabric=False`, 64 MiB), and + `nvl_ipc_prealloc` (the same NVL IPC profile with whole-chunk LOAD + preallocation). Do not substitute `ll_nofabric` for `fabric` and do not + promote NCCL CUMEM/NVLS localization probes. + +**Interfaces:** + +- `_verify_dispatch_result(recv_x, expert_num_tokens, dispatch_handle, + dispatch_event, ...)` validates the modern tuple. +- `_assert_valid_dispatch_equal(actual, expected, expert_counts)` compares only + valid per-expert slices exactly. +- LOAD requires a reconstructed `loaded_recv_x`, fills every valid slice with + NaN, replays, and requires exact warmup equality. +- `tests/modal_deepep.py` runs the three acceptance cases on `H100:2` from one + immutable `FOUNDRY_REF`. + +- [ ] **Step 1: Write and run the RED oracle contract** + +Add a CPU-only +`test_dispatch_oracle_rejects_unwritten_nan_sentinel` to +`tests/test_deepep_fabric.py`. It creates CPU tensors with valid expert counts, +fills the actual valid slices with NaN, and requires +`_assert_valid_dispatch_equal` to raise. Also add source assertions that LOAD +contains no warmup fallback and does contain the NaN fill. + +Run: + +```bash +python3 -m pytest \ + tests/test_deepep_fabric.py::test_dispatch_oracle_rejects_unwritten_nan_sentinel -q +``` + +Expected: failure because `_assert_valid_dispatch_equal` and strict LOAD +sentinel handling do not exist. + +- [ ] **Step 2: Replace the print-only oracle** + +Port the modern tuple checks and deterministic SAVE replay comparison from +`9842f2f`. Delete the unused approximate `graph_num_recv`. On LOAD, remove the +fallback: + +```python +verify_tensor = loaded_recv_x if loaded_recv_x is not None else recv_x +``` + +Require `loaded_recv_x`, sentinel valid slices, replay, synchronize, then +compare exact values. The CPU contract proves a no-op replay leaves NaNs and +fails; the H100 matrix proves real replay overwrites them. + +- [ ] **Step 3: Add a bounded Modal matrix launcher** + +Add this strict non-NVL fabric case to `tests/run_deepep_matrix.sh`: + +```bash +fabric) + export TEST_USE_FABRIC=1 TEST_NVL_BYTES_MB=0 + export NCCL_CUMEM_ENABLE=0 NCCL_NVLS_ENABLE=0 + PORT=29611 + ;; +``` + +The Modal image must pin Foundry by full SHA, build with `CC=gcc CXX=g++`, use +two H100s, install the same DeepEP/NVSHMEM dependencies as the pinned vLLM +toolchain, and invoke: + +```bash +bash tests/run_deepep_matrix.sh fabric both +bash tests/run_deepep_matrix.sh nvl_ipc both +bash tests/run_deepep_matrix.sh nvl_ipc_prealloc both +``` + +Return one JSON result containing command, return code, and output for each +case. Abort on the first nonzero result. CPU contracts must enforce the full +SHA, exact three-case list, `H100:2`, and child-process execution. + +- [ ] **Step 4: Run local GREEN** + +```bash +python3 -m pytest tests/test_modal_deepep.py -q +python3 -m pytest \ + tests/test_deepep_fabric.py::test_dispatch_oracle_rejects_unwritten_nan_sentinel -q +``` + +Expected: CPU contracts pass and the strict helper rejects an unwritten NaN +sentinel. + +- [ ] **Step 5: Commit and apply the review/push gate** + +```bash +git add tests/test_deepep_fabric.py tests/run_deepep_matrix.sh \ + tests/modal_deepep.py tests/test_modal_deepep.py +git commit -m "test: require exact restored DeepEP output" +git show --check --stat HEAD +``` + +Reviewer must confirm all three acceptance profiles run and loaded-output +fallback is impossible. Do not push this commit or run Modal yet: the push gate +opens only after Tasks 6-7 pass their local contracts and review. + +- [ ] **Step 6: Carry the strict matrix into the immutable candidate gate** + +Task 7's aggregate local gate must include `tests/test_modal_deepep.py`. Task 8 +runs `tests/modal_deepep.py` from the pushed full candidate SHA and requires +`fabric`, `nvl_ipc`, and `nvl_ipc_prealloc` to pass SAVE and LOAD with no NaN +sentinel remaining. A failure returns to this task for a new RED contract and +focused follow-up commit; never amend the tested SHA. + +### Task 6: Make the vLLM TP Archive Exact, Transactional, and Restartable + +**Files:** + +- Modify: `tests/tp_validation_helpers.py` +- Modify: `tests/test_tp_validation_helpers.py` +- Modify: `tests/modal_vllm_tp.py` +- Modify: `tests/test_modal_vllm_tp.py` +- Modify: `tests/test_vllm_tp_validation.py` +- Modify: `python/foundry/integration/vllm/hooks.py` +- Modify: `python/foundry/integration/vllm/runtime.py` +- Modify: `tests/vllm_tp_validation.py` + +**Symbols and sources:** + +- Exact count starting point: + `87c38bc43cfdbddc44b06e9c9af41ff0e3863c86`. +- Preserve `LoadRestorationProgress` from + `f79fcb0e1288b76172d3e1a4ac269d1ecd7434dd`. +- Reuse, do not duplicate, `archive_byte_fingerprints` from + `6497383997f5d30ffe9a084cef9cabc9f83299cf`. +- Adapt transaction/LOAD2 patterns from + `0ffb751fcf6585c1ad48f0e4674ce4b4ef39b514` and + `94a4d00bf4ec05b984836f9ee4bb7613a4d110ba`. + +**Interfaces:** + +- `EXPECTED_GRAPH_COUNT = 64`. +- `validate_archive_inventory(root: Path, tp_size: int, + expected_graph_count: int) -> dict[str, dict[str, object]]` requires 64 JSON + and 64 cugraph files per rank, identical stems, indices `0..63`, required + rank files, and root `warmup_state.json`. +- `collect_archive_snapshot(root: Path, tp_size: int, + expected_graph_count: int, foundry_base_addr: int, + foundry_region_end: int) -> dict[str, object]` returns inventory, semantic + fingerprints, full byte fingerprints, graph counts, and offsets. +- `_run_phase_local(phase: str, run_id: str, + expected_snapshot: dict[str, object] | None = None) -> + dict[str, object]` never commits a volume. +- `run_save_transaction(run_id: str) -> dict[str, object]` runs + seed/save/save2 in one remote container and + commits once only after every phase, teardown, log scan, and snapshot check + succeeds. +- `run_load_phase(phase: str, run_id: str, + expected_snapshot: dict[str, object]) -> dict[str, object]` accepts only + `load` or `load2`; two independent remote invocations receive the SAVE2 + snapshot, reread it after teardown, compare all maps, and never commit. + +- [ ] **Step 1: Write RED inventory and transaction contracts** + +Build realistic two-rank fixtures with 64 paired files per rank. Test missing, +duplicate, noncontiguous, mismatched-stem, missing required-file, and missing +root-state failures. Require `fatbin_entrypoint_packed.txt` in semantic +fingerprints; keep unexpected files in the full byte map. + +Static and fake-volume harness tests must require: + +```python +PHASES = ("baseline", "seed", "save", "save2", "load", "load2") +QUERY_PHASES = ("baseline", "save2", "load", "load2") +LOAD_PHASES = ("load", "load2") +``` + +They must also prove failed seed/save/save2 does not call commit and LOAD/LOAD2 +never call commit. + +Run: + +```bash +python3 -m pytest \ + tests/test_tp_validation_helpers.py \ + tests/test_modal_vllm_tp.py \ + tests/test_vllm_tp_validation.py -q +``` + +Expected: failures because exact inventory, transaction publication, and +LOAD2 are absent. + +- [ ] **Step 2: Add exact archive inventory and snapshot collection** + +Use a filename regex that extracts the integer immediately after `graph_`. +Require the same 64 stems for `.json` and `.cugraph`, and require: + +```python +REQUIRED_RANK_FILES = { + "graph_manifest.json", + "fatbin_image_packed.img", + "fatbin_entrypoint_packed.txt", + "final_alloc_offset.json", +} +``` + +The semantic map remains per-rank; the byte map remains every regular file +under the root. + +- [ ] **Step 3: Refactor archive publication into one transaction** + +Move process launch, requests, teardown, log scan, and evidence collection into +`_run_phase_local`. `run_save_transaction` must execute three fresh server +subprocesses in order in the same Modal container: + +```python +results = { + phase: _run_phase_local(phase, run_id) + for phase in ("seed", "save", "save2") +} +snapshot = collect_archive_snapshot(workspace, TP_SIZE, EXPECTED_GRAPH_COUNT) +archive_volume.commit() +return {"results": results, "snapshot": snapshot} +``` + +Wrap commit after all failure checks, not in `finally`. A failed transaction +raises without publishing. + +- [ ] **Step 4: Add LOAD2 and immutable post-LOAD rereads** + +Pass the SAVE2 inventory, semantic map, byte map, counts, and offsets into each +LOAD container. Compare after process teardown and log scanning. Any mismatch +is a phase failure. Neither LOAD calls `archive_volume.commit()`. + +- [ ] **Step 5: Emit and parse live rank offsets** + +Immediately after upstream `GPUModelRunner.capture_model` returns on LOAD, +read the live VMM cursor and emit: + +```text +[foundry] LOAD live final_alloc_offset rank_N: OFFSET +``` + +Derive `N` from the configured `rank_N` workspace. Add a regex to +`scan_log_text`; require exactly one marker per rank in each LOAD and equality +with that rank's SAVE2 `final_alloc_offset.json`. + +Require exactly one `[foundry] LOAD restoration complete rank_N: 64/64 graphs` +per rank per LOAD and zero capture-offset/manifest-save markers. + +- [ ] **Step 6: Run GREEN** + +```bash +python3 -m pytest \ + tests/test_tp_validation_helpers.py \ + tests/test_modal_vllm_tp.py \ + tests/test_vllm_tp_validation.py -q +python3 -m compileall -q \ + python/foundry/integration/vllm \ + tests/modal_vllm_tp.py \ + tests/tp_validation_helpers.py \ + tests/vllm_tp_validation.py +``` + +Expected: exact inventory, transaction, LOAD2, completion, and live-offset +contracts pass. + +- [ ] **Step 7: Commit and apply the review/push gate** + +```bash +git add \ + python/foundry/integration/vllm/hooks.py \ + python/foundry/integration/vllm/runtime.py \ + tests/modal_vllm_tp.py \ + tests/test_modal_vllm_tp.py \ + tests/test_tp_validation_helpers.py \ + tests/test_vllm_tp_validation.py \ + tests/tp_validation_helpers.py \ + tests/vllm_tp_validation.py +git commit -m "test(vllm): make TP archive publication transactional" +git show --check --stat HEAD +``` + +Reviewer must trace every `archive_volume.commit()` and prove there is one +successful save-transaction publication, zero LOAD publications, and no +partial failure path. Do not push GREEN until Task 7 passes locally. + +### Task 7: Complete vLLM TP Output, Pin, Provenance, Error, and Cleanup Gates + +**Files:** + +- Modify: `recipe/vllm/serve_qwen3-8b_tp.sh` +- Modify: `tests/test_vllm_tp_recipe.py` +- Modify: `tests/modal_vllm_tp.py` +- Modify: `tests/test_modal_vllm_tp.py` +- Modify: `tests/vllm_tp_validation.py` +- Modify: `tests/test_vllm_tp_validation.py` + +**Symbols and sources:** + +- Preserve bounded exact completion evidence from + `8bf18e46f87550b6608a3ca1121271838a5ed8f0` and stable process signals from + `e55df61e5caba372cb216c6df0a04bc6f572a543`. +- Preserve atomic vLLM compilation config from + `57243c8dd7ee0420a653d5a4c303a4317e00237f` and NCCL profile from + `3c7cfa3f61582cae63f94cffbd529bc9c473b1f1`. +- Preserve Modal sibling mounts from + `ab041b4281d7ea5eac791a8edf28a3ee3bfd0aec`. +- Adapt current SGLang report/cleanup patterns without importing SGLang + backend behavior. + +**Interfaces:** + +- Four prompts, temperature 0, seed 0, and four-token limit remain fixed. +- Baseline/save2/load/load2 each produce exactly four nonempty sequential and + four nonempty barrier-backed concurrent completion records. +- Every record contains exact `text` and nonempty `token_strings`; all four + phases must equal baseline byte-for-byte at the Python string level. +- Final report records immutable refs, requested/actual versions, graph + evidence, process results, run URL, and cleanup status. + +- [ ] **Step 1: Write RED output and pin contracts** + +Add recipe tests rejecting `main`, short SHAs, non-hex values, and extra +arguments through `VLLM_MODEL_REVISION`. Add harness tests requiring all four +query phases and exact sequential/concurrent equality of both fields. + +Run: + +```bash +python3 -m pytest \ + tests/test_vllm_tp_recipe.py \ + tests/test_modal_vllm_tp.py \ + tests/test_vllm_tp_validation.py -q +``` + +Expected: mutable model override, missing save2/load2 requests, and incomplete +provenance tests fail. + +- [ ] **Step 2: Fail closed on model revision** + +In the shell recipe, reject unless: + +```bash +[[ "$MODEL_REVISION" =~ ^[0-9a-fA-F]{40}$ ]] +``` + +Keep the validated default +`b968826d9c46dd6066d109eabc6255188de91218`. + +- [ ] **Step 3: Enforce the complete output matrix** + +Use `QUERY_PHASES` for both request modes. Reject wrong list length, empty text, +empty token strings, or any phase mismatch. Do not reconstruct token IDs from +the tokenizer; record IDs only if the pinned API returns them directly. + +- [ ] **Step 4: Record complete provenance** + +Add package metadata and `nvidia-smi` helpers that do not initialize CUDA in +the orchestration parent. The final JSON must include: + +- `foundry_ref`, `vllm_ref`, `vllm_wheel_commit`, and model revision; +- CUDA image tag, requested torch `2.11.0`, actual torch version, requested + vLLM SHA, actual vLLM package version, driver, GPU name and count; +- NCCL runtime version parsed from child logs; +- per-rank offsets, inventories, semantic maps, byte maps, completion maps, + process return codes, and phase times; +- `run_url = app.get_dashboard_url()`; +- `cleanup_attempted`, `cleanup_completed`, and `artifacts_retained`. + +- [ ] **Step 5: Preserve error observability and add cleanup** + +Keep bounded health/request/teardown timeouts, early-death rejection, accepted +shutdown codes `{0, -2, 130}`, SIGKILL failure, complete fatal scans, and helper +handle closure. + +Add a remote cleanup function that removes only `RUNS_ROOT/run_id`, verifies it +is gone, then commits that deletion. It is a cleanup transaction, not a LOAD +publication. On failed acceptance, retain artifacts and report that fact. + +- [ ] **Step 6: Run GREEN and the pre-push aggregate gate** + +```bash +python3 -m pytest \ + tests/test_deepep_transport.py \ + tests/test_deepep_transport_vllm.py \ + tests/test_experimental_recipe.py \ + tests/test_modal_deepep.py \ + tests/test_modal_vllm_tp.py \ + tests/test_tp_validation_helpers.py \ + tests/test_vllm_tp_recipe.py \ + tests/test_vllm_tp_validation.py -q +pre-commit run --all-files +git diff --check origin/main...HEAD +``` + +Expected: all focused CPU contracts and repository hooks pass. + +- [ ] **Step 7: Commit, review, and push the immutable candidate** + +```bash +git add recipe/vllm/serve_qwen3-8b_tp.sh \ + tests/modal_vllm_tp.py tests/test_modal_vllm_tp.py \ + tests/test_vllm_tp_recipe.py tests/test_vllm_tp_validation.py \ + tests/vllm_tp_validation.py +git commit -m "test(vllm): complete TP restart acceptance" +git show --check --stat HEAD +test "$(git merge-base HEAD origin/main)" = "$(git rev-parse origin/main)" +test "$(git rev-list --merges origin/main..HEAD --count)" = 0 +git push origin cursor/tp-upstream-consolidation-fe33 +``` + +Push only after a reviewer confirms exact query cardinality, full provenance, +failure precedence, cleanup, and no missing sibling module mount. Record the +pushed full SHA as `CANDIDATE_SHA`; Task 8 must use it unchanged. + +### Task 8: Run Local, Native, DeepEP, and vLLM H100 Validation + +**Files:** + +- Verify: all Task 1-7 files +- Evidence only: Modal JSON/log artifacts for the immutable candidate + +**Sources:** + +- Build/toolchain contract: `AGENTS.md`, `README.md`, and + `CONTRIBUTING.md`. +- Final vLLM command and acceptance criteria: + `.superpowers/sdd/final-vllm-branch-audit.md`. + +**Interfaces:** + +- Consumes: pushed `CANDIDATE_SHA`. +- Produces: local CPU/lint evidence, one native H100 result, one 2×H100 DeepEP + result, and one 2×H100 six-phase vLLM result. + +- [ ] **Step 1: Lock the candidate and carry forward RED evidence** + +```bash +CANDIDATE_SHA="$(git rev-parse HEAD)" +test "${#CANDIDATE_SHA}" = 40 +test "$(git rev-parse origin/cursor/tp-upstream-consolidation-fe33)" = \ + "$CANDIDATE_SHA" +test -z "$(git status --short)" +``` + +Expected: the immutable candidate is pushed and clean. Preserve each owning +task's RED output beside its GREEN result. Any failed command below is a RED +validation result: stop, return to that task, add or strengthen its regression, +and produce a new candidate SHA. + +- [ ] **Step 2: Run the fresh local GREEN gate** + +```bash +python3 -m pytest \ + tests/test_deepep_transport.py \ + tests/test_deepep_transport_vllm.py \ + tests/test_experimental_recipe.py \ + tests/test_modal_deepep.py \ + tests/test_modal_native.py \ + tests/test_modal_vllm_tp.py \ + tests/test_tp_validation_helpers.py \ + tests/test_vllm_tp_recipe.py \ + tests/test_vllm_tp_validation.py -q +python3 -m compileall -q python tests +pre-commit run --all-files +git diff --check origin/main...HEAD +``` + +Expected: every command exits zero. + +- [ ] **Step 3: Run native H100 GREEN** + +```bash +FOUNDRY_REF="$CANDIDATE_SHA" \ + modal run --env rahul-dev tests/modal_native.py --suite all +``` + +Expected: focused VMM/RNG and full core suite pass; no CUDA is initialized in +the local orchestration process. + +- [ ] **Step 4: Run strict DeepEP GREEN** + +```bash +FOUNDRY_REF="$CANDIDATE_SHA" \ + modal run --env rahul-dev tests/modal_deepep.py +``` + +Expected: `fabric`, `nvl_ipc`, and `nvl_ipc_prealloc` each pass SAVE and +LOAD with exact loaded-output sentinel overwrite. + +- [ ] **Step 5: Run final vLLM TP GREEN** + +```bash +FOUNDRY_REF="$CANDIDATE_SHA" \ +VLLM_REF="4309c257d3f639e5490d3811293c890c61c76f29" \ +MODAL_GPU="H100:2" \ + modal run --env rahul-dev tests/modal_vllm_tp.py +``` + +Expected: baseline/seed/save/save2/load/load2 all pass; each rank has exact +64 JSON/cugraph pairs; both LOADs emit one 64/64 marker per rank; live offsets +equal SAVE2 per rank; archive bytes are immutable; sequential and concurrent +text/token strings match exactly; exits and cleanup pass; final JSON includes +the Modal URL. + +- [ ] **Step 6: Apply the review/push gate** + +No code commit is expected in this verification-only task. If any check fails, +do not patch inside Task 8: return to the owning task, add a RED contract, +commit a focused fix, push a new SHA, and rerun all four gates. Do not update +support claims from a failed or superseded SHA. + +### Task 9: Close Branch Disposition and Publish Evidence + +**Files:** + +- Modify only after Task 8 GREEN: `README.md` +- Modify only after Task 8 GREEN: `ROADMAP.md` +- Modify only after Task 8 GREEN: `RELEASE.md` +- Modify only after Task 8 GREEN: `docs/vllm/overview.md` +- Modify only after Task 8 GREEN: `docs/vllm/moe-and-deepep.md` +- Modify only after Task 8 GREEN: `recipe/vllm/README.md` +- Update externally: draft PR #6 body +- Create ignored evidence report under `.superpowers/sdd/` + +**Sources:** + +- `.superpowers/sdd/final-all-branches-audit.md` +- `.superpowers/sdd/final-ipc-branch-audit.md` +- `.superpowers/sdd/final-vllm-branch-audit.md` +- Final Task 8 JSON and Modal URLs only. + +**Interfaces:** + +- Consumes: immutable candidate SHA and complete GREEN evidence. +- Produces: exhaustive all-branch disposition, measured documentation, and an + evidence-backed PR update without broadening SGLang or vLLM scope. + +- [ ] **Step 1: Refresh and classify every branch patch** + +```bash +git fetch origin --prune +for ref in \ + origin/a100-issue \ + origin/ep-ipc \ + origin/devin/1784752573-sglang-nofabric-ep-ipc \ + origin/devin/1784764658-sglang-tensor-parallel \ + origin/cursor/no-imex-deepep-f10b \ + origin/cursor/sglang-v0515-api-port-609f \ + origin/cursor/tp-consolidation-2e2c \ + origin/cursor/vllm-tp-consolidation-5d04 \ + origin/cursor/setup-dev-environment-3aab; do + git cherry -v HEAD "$ref" +done +``` + +For every `+`, record one disposition with commit/tree/blob evidence: +manually represented required behavior, superseded by stronger current code, +documentation-only, diagnostic-only, or conflicting/unvalidated. A `+` is not +automatically missing because manual ports need not share patch IDs. + +- [ ] **Step 2: Verify direct-main lineage and final diff** + +```bash +test "$(git merge-base HEAD origin/main)" = "$(git rev-parse origin/main)" +test "$(git rev-list --merges origin/main..HEAD --count)" = 0 +git rev-list --left-right --count origin/main...HEAD +git diff --check origin/main...HEAD +git diff --stat origin/main...HEAD +``` + +Expected: zero commits on the main-only side, zero merges, and no whitespace +errors. + +- [ ] **Step 3: Confirm documentation is RED for the new candidate** + +```bash +! rg -l "$CANDIDATE_SHA" \ + README.md ROADMAP.md RELEASE.md docs/vllm/overview.md \ + docs/vllm/moe-and-deepep.md recipe/vllm/README.md +rg -n 'git checkout bump|Combined-head 2×H100 acceptance run' \ + README.md ROADMAP.md RELEASE.md docs/vllm recipe/vllm/README.md +``` + +Expected: the final candidate is not yet claimed and stale pending/mutable +wording is identified. If the SHA already appears, verify it came from this +exact Task 8 run before changing anything. + +- [ ] **Step 4: Update measured evidence documentation** + +Record the exact Foundry, vLLM, wheel, model, CUDA, torch, NCCL, GPU, and driver +values; Modal URLs; 64/64 completion maps; per-rank offsets; inventory, +semantic, and byte-map comparisons; exact output comparisons; exits; and +cleanup. + +Restore the vLLM TP invariant section: single host, exactly two Hopper GPUs, +Qwen3-8B, V1, `FULL_DECODE_ONLY`, PyNCCL P2P/legacy IPC, CUMEM/NVLS off, +custom/symmetric/FlashInfer/fused AR off, and one atomic compilation config. + +Replace the mutable `git checkout bump` example with the validated full +Foundry SHA. Keep TP>2, V2, multi-host, Ampere/Blackwell, MoE TP, and alternate +collectives unsupported. Keep the current verified SGLang profile and its +existing evidence unchanged. + +- [ ] **Step 5: Run documentation GREEN** + +```bash +python3 -m pytest \ + tests/test_experimental_recipe.py \ + tests/test_sglang_documentation.py \ + tests/test_vllm_tp_recipe.py -q +pre-commit run markdownlint-cli2 --all-files +pre-commit run --all-files +git diff --check origin/main...HEAD +``` + +Expected: documentation contracts and all quality gates pass. + +- [ ] **Step 6: Commit evidence and apply the final review/push gate** + +```bash +git add README.md ROADMAP.md RELEASE.md \ + docs/vllm/overview.md docs/vllm/moe-and-deepep.md \ + recipe/vllm/README.md +git commit -m "docs: record final vLLM TP validation" +git show --check --stat HEAD +git push origin cursor/tp-upstream-consolidation-fe33 +``` + +Reviewer must match every measured value to Task 8 JSON and confirm no +historical branch run is presented as final-HEAD evidence. + +- [ ] **Step 7: Update draft PR #6** + +Use the PR management tool, not a write-capable `gh` command. Include the final +commit, direct-main proof, all-branch disposition, exact pins, local/native/ +DeepEP/vLLM results, Modal URLs, archive/output/cleanup evidence, and explicit +unsupported scope. Do not merge, enable auto-merge, or mark ready unless the +user separately requests it. + +## Final Completion Criteria + +- All nine tasks are complete in dependency order. +- No wholesale feature-branch cherry-pick or merge occurred. +- The branch remains a merge-free direct delta from current `origin/main`. +- Current generation/open-count/alignment and verified SGLang behavior remain + intact. +- SGLang EP is quarantined rather than silently advertised or widened. +- The immutable final SHA passes local, native H100, strict DeepEP H100:2, and + six-phase vLLM H100:2 validation. +- Documentation and PR status cite only final-SHA evidence. diff --git a/docs/superpowers/plans/2026-07-23-sglang-main-port.md b/docs/superpowers/plans/2026-07-23-sglang-main-port.md new file mode 100644 index 00000000..713be1fd --- /dev/null +++ b/docs/superpowers/plans/2026-07-23-sglang-main-port.md @@ -0,0 +1,1175 @@ +# SGLang Current-Main Port Implementation Plan + +> [!WARNING] +> **SUPERSEDED FOR SGLANG.** The plain-NCCL policy in this historical plan was +> amended after corrupt fresh-LOAD output. Use the +> [symmetric-TP treatment design](../specs/2026-07-23-sglang-symmetric-tp-treatment-design.md) +> and +> [implementation plan](2026-07-23-sglang-symmetric-tp-treatment.md) as the +> current SGLang policy and evidence boundary. + +--- + +> **For agentic workers:** REQUIRED SUB-SKILL: Use superpowers:subagent-driven-development (recommended) or superpowers:executing-plans to implement this plan task-by-task. Steps use checkbox (`- [ ]`) syntax for tracking. + +**Goal:** Retarget Foundry's SGLang TP=2 SAVE-to-LOAD integration onto pinned upstream `sgl-project/sglang` main with zero upstream edits, activated through the `sglang.srt.plugins` entry point and `HookRegistry`, and re-validated on 2x H100. + +**Architecture:** Ship a Foundry-owned general plugin whose `activate()` reads `FOUNDRY_GRAPH_EXTENSION_CONFIG_PATH` and registers `HookRegistry` hooks. SAVE binds to a `FoundryFullCudaGraphBackend(FullCudaGraphBackend)` `capture_one` override plus an AROUND `DecodeCudaGraphRunner.capture`; LOAD REPLACEs `DecodeCudaGraphRunner.capture`; graph keys migrate to `ShapeKey`; LOAD reuses the serialized `MemoryPoolConfig` by hooking `KVCacheConfigurator._resolve_memory_pool_config`. The narrow dense NCCL P2P/IPC profile is enforced fail-closed from an executed `ServerArgs.__post_init__` hook. + +**Tech Stack:** Python 3.12, PyTorch 2.11/cu130, SGLang (pinned main), FlashInfer, NCCL 2.28+, pytest, Bash, Modal, C++17/CUDA Driver API (shared VMM core, unchanged). + +## Global Constraints + +- Work only on branch `cursor/tp-upstream-consolidation-fe33`; do not leave it. The merge base is `origin/main` `21824649e711487a2c6af01cb631b3e41ccc32d8`; the review diff is `origin/main...HEAD`. +- Target immutable SGLang commit `a25164bda355250015c0ba6d8821f0d4b37abe7d`; this is the sole validation target. "Support" means this exact commit, never a moving `main`. +- Immutable model revision Qwen3-8B `b968826d9c46dd6066d109eabc6255188de91218`. +- Zero edits to a clean upstream SGLang checkout and no fork commit. Activation is the `sglang.srt.plugins` entry point `foundry = "foundry_sglang_plugin:activate"`, config via `FOUNDRY_GRAPH_EXTENSION_CONFIG_PATH`, selection via `SGLANG_PLUGINS=foundry`. +- Supported envelope: single-host dense Qwen3-8B, TP=2, `PP=DP=1`, no EP/DP-attention/disaggregation, plain NCCL P2P/legacy CUDA IPC, FlashInfer attention, full decode graphs only, prefill disabled. +- Disable as correctness requirements: SGLang custom all-reduce, NCCL/torch symmetric memory, MSCCL++, FlashInfer autotune and fused all-reduce, NCCL CUMEM/NVLS, post-capture KV sizing (`SGLANG_ENABLE_POST_CAPTURE_KV_SIZING=0`), piecewise/breakable/tc_piecewise decode graphs, pdmux/multi-stream capture, expandable allocator segments. +- Keep engine revision, model revision, seed, graph settings, rank order, and transport identical across baseline, SAVE, and LOAD. +- Put standard-library and Foundry imports at module scope. Two documented no-inline-import exceptions apply to every task below: (1) SGLang imports MUST stay lazy (inside functions) because SGLang is an optional integration and importing it at Foundry import time breaks core-only installs; (2) Foundry imports that transitively load `torch`/`foundry.ops` (the compiled CUDA extension) MUST stay lazy inside `activate()`/`register_hooks()` so an inactive plugin does not import CUDA in every baseline SGLang process at `load_plugins()` time. +- Any discrepancy discovered between this plan and a fresh clean checkout of the pinned SHA MUST be resolved from source evidence in that checkout before writing code. Do not guess an API shape; read the pinned source and update the plan. +- Run all CUDA-native and engine validation only on 2x H100 in Modal `rahul-dev`; do not initialize CUDA in the orchestration parent. The local VM has no NVIDIA device or CUDA toolkit. +- Do not claim support broader than the exact final-head evidence. +- **Process:** No future subagent that implements this plan may use a fast model. Each task is implemented and reviewed by a non-fast model. + +--- + +### Task 1: Plugin Activation and Recipe/Harness Pin Migration + +**Files:** + +- Create: `python/foundry/integration/sglang/plugin.py` +- Modify: `pyproject.toml:12-19` (add `[project.entry-points."sglang.srt.plugins"]`) +- Modify: `python/foundry/integration/sglang/hooks.py:69-100` (add `register_hooks()`; keep `install_hooks` shim) +- Modify: `recipe/experimental/serve_qwen3-8b_sglang_tp.sh:34-70` +- Modify: `tests/modal_sglang_tp.py:65-162` +- Test: `tests/test_sglang_plugin_activation.py` +- Test: `tests/test_sglang_tp_recipe.py:48-101` + +**Interfaces:** + +- Consumes: SGLang `load_plugins()` (calls zero-arg `activate()` in every process), `os.environ["FOUNDRY_GRAPH_EXTENSION_CONFIG_PATH"]`, existing `config.load_graph_extension_config(path: str) -> None`, `config.get_graph_extension_mode() -> CUDAGraphExtensionMode`. +- Produces: `foundry.integration.sglang.plugin.activate() -> None`; `foundry.integration.sglang.hooks.register_hooks() -> None` (registers all `HookRegistry` hooks and is idempotent via the existing `_INSTALLED` flag). + +- [ ] **Step A1: Write the failing activation unit test** + +Create `tests/test_sglang_plugin_activation.py`. Mirror the module-mocking style of `tests/test_sglang_attention_metadata.py` (`_install_module`, `_load_module`) so no real SGLang import is needed. + +```python +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the Foundry project +"""CPU-only contracts for the Foundry SGLang plugin entry point.""" + +from __future__ import annotations + +import importlib.util +from pathlib import Path + +REPO_ROOT = Path(__file__).parents[1] +PLUGIN_PATH = REPO_ROOT / "python" / "foundry" / "integration" / "sglang" / "plugin.py" + + +def _load_plugin(): + spec = importlib.util.spec_from_file_location("foundry_plugin_under_test", PLUGIN_PATH) + assert spec is not None and spec.loader is not None + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + return module + + +def test_activate_is_noop_without_config_env(monkeypatch): + monkeypatch.delenv("FOUNDRY_GRAPH_EXTENSION_CONFIG_PATH", raising=False) + calls = [] + plugin = _load_plugin() + monkeypatch.setattr(plugin, "load_graph_extension_config", lambda path: calls.append(path)) + monkeypatch.setattr(plugin, "_register_hooks", lambda: calls.append("hooks")) + plugin.activate() + assert calls == [] + + +def test_activate_loads_config_and_registers_hooks(monkeypatch, tmp_path): + cfg = tmp_path / "cfg.toml" + cfg.write_text('mode = "save"\n') + monkeypatch.setenv("FOUNDRY_GRAPH_EXTENSION_CONFIG_PATH", str(cfg)) + calls = [] + plugin = _load_plugin() + monkeypatch.setattr(plugin, "load_graph_extension_config", lambda path: calls.append(("cfg", path))) + monkeypatch.setattr(plugin, "_register_hooks", lambda: calls.append("hooks")) + plugin.activate() + assert calls == [("cfg", str(cfg)), "hooks"] +``` + +- [ ] **Step A2: Run the test to verify it fails** + +Run: `python3 -m pytest tests/test_sglang_plugin_activation.py -q` +Expected: FAIL with `ModuleNotFoundError` / file-not-found because `python/foundry/integration/sglang/plugin.py` does not exist. + +- [ ] **Step A3: Create the plugin module** + +Create `python/foundry/integration/sglang/plugin.py`: + +```python +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the Foundry project +"""SGLang plugin entry point for the Foundry CUDA graph extension. + +Invoked by SGLang's ``load_plugins()`` (``sglang.srt.plugins``) in the engine +parent and every scheduler/TP-worker process. Config path is delivered through +``FOUNDRY_GRAPH_EXTENSION_CONFIG_PATH`` (inherited across ``spawn``); when unset, +Foundry is installed but inactive. +""" + +from __future__ import annotations + +import logging +import os + +from foundry.integration.sglang.config import load_graph_extension_config + +logger = logging.getLogger(__name__) + + +def activate() -> None: + cfg_path = os.environ.get("FOUNDRY_GRAPH_EXTENSION_CONFIG_PATH") + if not cfg_path: + return + load_graph_extension_config(cfg_path) + _register_hooks() + + +def _register_hooks() -> None: + # Lazy import (documented exception): register_hooks transitively loads + # torch/foundry.ops, so importing it here — only when a config path is set — + # keeps CUDA out of every baseline SGLang process and preserves core-only + # Foundry installs. + from foundry.integration.sglang.hooks import register_hooks + + register_hooks() + logger.info("[Foundry] SGLang plugin activated") +``` + +- [ ] **Step A4: Run the test to verify it passes** + +Run: `python3 -m pytest tests/test_sglang_plugin_activation.py -q` +Expected: PASS (both tests). `_register_hooks` is patched, so no real `hooks.register_hooks` runs yet. + +- [ ] **Step A5: Add the entry point to `pyproject.toml`** + +After the `[project.optional-dependencies]` block (`pyproject.toml:21-24`), add: + +```toml +[project.entry-points."sglang.srt.plugins"] +foundry = "foundry_sglang_plugin:activate" +``` + +- [ ] **Step A6: Add `register_hooks()` to `hooks.py`** + +In `python/foundry/integration/sglang/hooks.py`, add a `register_hooks()` function that reuses the existing `_INSTALLED` guard and calls the Task 2/3/4 registration helpers. For Task 1 it registers nothing yet beyond the guard so activation is inert until later tasks land: + +```python +def register_hooks() -> None: + """Register all Foundry SGLang hooks through SGLang's HookRegistry. + + Populated incrementally by later tasks (profile validation, spawn/rank/VMM, + memory pool, capture/load). Idempotent via ``_INSTALLED``. + """ + global _INSTALLED + if _INSTALLED: + return + logger.info( + "[Foundry] SGLang hooks registering: mode=%s workspace=%s", + get_graph_extension_mode().value, + get_workspace_root(), + ) + _INSTALLED = True + logger.info("[Foundry] SGLang hooks registered") +``` + +Neutralize the legacy `install_hooks(server_args)` in the same edit: replace its body with an immediate `return` (the `foundry_graph_extension_config_path` field no longer exists upstream, so the plugin path is now the only activation route). Do **not** have it call the legacy `_patch_*` helpers, which Tasks 2–4 delete; leaving them called would create dangling references as those helpers are removed. Task 5 deletes `install_hooks` itself after the clean-checkout contract lands. + +- [ ] **Step A7: Write the failing recipe contract updates** + +In `tests/test_sglang_tp_recipe.py`, update `test_tp_recipe_uses_same_nccl_transport_for_every_phase` (lines 48-59) to assert the new env and flags: + +```python + assert env["PYTORCH_CUDA_ALLOC_CONF"] == "expandable_segments:False" + assert env["NCCL_CUMEM_ENABLE"] == "0" + assert env["NCCL_NVLS_ENABLE"] == "0" + assert env["SGLANG_ENABLE_POST_CAPTURE_KV_SIZING"] == "0" + assert env["SGLANG_PLUGINS"] == "foundry" + assert env["SGL_MODEL_REVISION"] == "b968826d9c46dd6066d109eabc6255188de91218" + cuda_graph_config = args[args.index("--cuda-graph-config") + 1] + assert '"decode"' in cuda_graph_config and '"full"' in cuda_graph_config + assert '"prefill"' in cuda_graph_config and '"disabled"' in cuda_graph_config + assert "--disable-piecewise-cuda-graph" not in args + assert args[args.index("--random-seed") + 1] == "42" +``` + +Replace `test_tp_recipe_selects_matching_foundry_config` (lines 84-101) so `--save`/`--load` set `FOUNDRY_GRAPH_EXTENSION_CONFIG_PATH` in the environment instead of passing `--foundry-graph-extension-config-path`: + +```python +@pytest.mark.parametrize( + ("mode", "config_name"), + [ + ("--save", "sglang_foundry_tp_save.toml"), + ("--load", "sglang_foundry_tp_load.toml"), + ], +) +def test_tp_recipe_activates_foundry_via_env(tmp_path, mode, config_name): + env, args = _run_recipe(tmp_path, mode) + assert Path(env["FOUNDRY_GRAPH_EXTENSION_CONFIG_PATH"]) == RECIPE_DIR / config_name + assert env["SGLANG_PLUGINS"] == "foundry" + assert "--foundry-graph-extension-config-path" not in args + assert args[args.index("--tp-size") + 1] == "2" + assert args[args.index("--revision") + 1] == env["SGL_MODEL_REVISION"] + assert "--disable-custom-all-reduce" in args +``` + +- [ ] **Step A8: Run the recipe contracts and verify RED** + +Run: `python3 -m pytest tests/test_sglang_tp_recipe.py -q` +Expected: FAIL — the current script sets neither `SGLANG_PLUGINS`, `SGLANG_ENABLE_POST_CAPTURE_KV_SIZING`, nor `--cuda-graph-config`, still passes `--foundry-graph-extension-config-path`, and still passes `--disable-piecewise-cuda-graph`. + +- [ ] **Step A9: Migrate the recipe script** + +In `recipe/experimental/serve_qwen3-8b_sglang_tp.sh`, replace the env exports (lines 34-39) and the Foundry-arg block (lines 41-70) so activation is env-driven and the decode backend is explicit. Replace lines 37-70 with: + +```bash +export NCCL_CUMEM_ENABLE=0 +export NCCL_NVLS_ENABLE=0 +export PYTORCH_CUDA_ALLOC_CONF=expandable_segments:False +export SGLANG_ENABLE_POST_CAPTURE_KV_SIZING=0 +export SGLANG_PLUGINS=foundry +export SGL_MODEL_REVISION="${SGL_MODEL_REVISION:-b968826d9c46dd6066d109eabc6255188de91218}" + +if [[ "$MODE" == "--save" ]]; then + export FOUNDRY_GRAPH_EXTENSION_CONFIG_PATH="${SCRIPT_DIR}/sglang_foundry_tp_save.toml" + echo "Using Foundry TP SAVE (NCCL P2P/IPC): ${FOUNDRY_GRAPH_EXTENSION_CONFIG_PATH}" +elif [[ "$MODE" == "--load" ]]; then + export FOUNDRY_GRAPH_EXTENSION_CONFIG_PATH="${SCRIPT_DIR}/sglang_foundry_tp_load.toml" + echo "Using Foundry TP LOAD (NCCL P2P/IPC): ${FOUNDRY_GRAPH_EXTENSION_CONFIG_PATH}" +else + echo "Running without Foundry (baseline SGLang)" +fi + +sglang serve \ + --model-path "$MODEL_NAME" \ + --revision "$SGL_MODEL_REVISION" \ + --trust-remote-code \ + --host "$HOST" --port "$PORT" \ + --tp-size "$TP_SIZE" \ + --random-seed "$RANDOM_SEED" \ + --mem-fraction-static "$MEM_FRACTION_STATIC" \ + --disable-radix-cache \ + --disable-custom-all-reduce \ + --disable-flashinfer-autotune \ + --attention-backend flashinfer \ + --cuda-graph-config '{"decode":{"backend":"full","max_bs":256},"prefill":{"backend":"disabled"}}' +``` + +The `FOUNDRY_ARGS` array and the `--foundry-graph-extension-config-path` argument are removed; baseline (no `MODE`) exports `SGLANG_PLUGINS=foundry` but no config path, so `activate()` is a no-op. + +- [ ] **Step A10: Run the recipe contracts and verify GREEN** + +Run: + +```bash +python3 -m pytest tests/test_sglang_tp_recipe.py -q +bash -n recipe/experimental/serve_qwen3-8b_sglang_tp.sh +``` + +Expected: all recipe contract tests pass and the script parses. If `_run_recipe`'s fake `sglang` shim does not export `SGLANG_PLUGINS` for the baseline case, the baseline `test_tp_recipe_uses_same_nccl_transport_for_every_phase[None]` still passes because the export is unconditional. + +- [ ] **Step A11: Migrate the Modal harness pins** + +In `tests/modal_sglang_tp.py`, replace the fork pin block (lines 65-71) with the upstream pin and drop the cherry-pick: + +```python +# Foundry activates through SGLang's first-class sglang.srt.plugins entry point; +# no fork commit is required. Pin an immutable upstream main commit. +SGLANG_REPO = "https://github.com/sgl-project/sglang.git" +SGLANG_BASE = os.environ.get("SGLANG_BASE", "a25164bda355250015c0ba6d8821f0d4b37abe7d") +MODEL_REVISION = os.environ.get( + "SGL_MODEL_REVISION", "b968826d9c46dd6066d109eabc6255188de91218" +) +``` + +Replace the SGLang build `run_commands` (lines 143-149) with a plain editable install of the pinned checkout (no `git fetch origin foundry`, no cherry-pick): + +```python + .run_commands( + f"git clone {SGLANG_REPO} /sglang", + f"cd /sglang && git checkout {SGLANG_BASE}", + "cd /sglang && pip install -e 'python[all]' --no-build-isolation", + ) +``` + +In `_launch` (lines 250-269) export activation and pin the model revision so the child inherits them: + +```python + env["SGL_MODEL"] = MODEL + env["SGL_MODEL_REVISION"] = MODEL_REVISION + env["SGLANG_PLUGINS"] = "foundry" + env["NCCL_DEBUG"] = "INFO" +``` + +- [ ] **Step A12: Verify harness and imports compile** + +Run: + +```bash +python3 -m compileall -q python/foundry/integration/sglang/plugin.py tests/modal_sglang_tp.py +python3 -c "import ast, pathlib; ast.parse(pathlib.Path('tests/modal_sglang_tp.py').read_text())" +``` + +Expected: exit zero. `modal_sglang_tp.py` still imports `modal`/`tp_validation_helpers` at runtime only under Modal, so `compileall` (byte-compile, no import) is the correct local check. + +- [ ] **Step A13: Commit** + +```bash +git add \ + python/foundry/integration/sglang/plugin.py \ + python/foundry/integration/sglang/hooks.py \ + pyproject.toml \ + recipe/experimental/serve_qwen3-8b_sglang_tp.sh \ + tests/modal_sglang_tp.py \ + tests/test_sglang_plugin_activation.py \ + tests/test_sglang_tp_recipe.py +git commit -m "feat(sglang): activate via plugin entry point and pin upstream main" +``` + +--- + +### Task 2: Current-Main Hook Targets, Profile Validation, and Rank/VMM Setup + +**Files:** + +- Modify: `python/foundry/integration/sglang/hooks.py:69-190,668-700` (registry-based hooks, `self.ps` ranks) +- Modify: `python/foundry/integration/sglang/config.py:38-66` (`_resolve_dp_rank` reads `self.ps`) +- Test: `tests/test_sglang_hook_targets.py` +- Test: `tests/test_sglang_profile_validation.py` + +**Interfaces:** + +- Consumes: `sglang.srt.plugins.hook_registry.HookRegistry.register(target: str, fn, hook_type)` and `HookType.{BEFORE,AFTER,AROUND,REPLACE}`; `ServerArgs.__post_init__(self)`; `Engine._launch_scheduler_processes` (classmethod); `ModelRunner.init_torch_distributed(self)`; `ModelRunner.ps` (`ParallelState` with `tp_rank`, `pp_rank`, `dp_rank`, `gpu_id`). +- Produces: `hooks._register_profile_validation()`, `hooks._register_spawn_hooks()`, `hooks._register_rank_vmm_hook()`, and a module-level `_validate_profile(server_args) -> None` that raises `RuntimeError` outside the supported envelope. + +- [ ] **Step B1: Write the failing profile-validation unit test** + +Create `tests/test_sglang_profile_validation.py`. Build a `SimpleNamespace` `server_args` and assert `_validate_profile` accepts the supported envelope and rejects each unsupported setting. Load `hooks.py` with SGLang modules mocked (reuse `_install_module`/`_load_module` from `test_sglang_attention_metadata.py`, or copy them locally): + +```python +from types import SimpleNamespace +import pytest + + +def _supported_args(): + return SimpleNamespace( + tp_size=2, pp_size=1, dp_size=1, + enable_dp_attention=False, + enable_ep_moe=False, + disable_custom_all_reduce=True, + enable_mscclpp=False, + enable_torch_symm_mem=False, + enable_symm_mem=False, + enable_nccl_nvls=False, + disable_flashinfer_autotune=True, + attention_backend="flashinfer", + ) + + +def test_validate_profile_accepts_supported(hooks_module): + hooks_module._validate_profile(_supported_args()) + + +@pytest.mark.parametrize("field,value", [ + ("tp_size", 3), ("pp_size", 2), ("dp_size", 2), + ("enable_dp_attention", True), ("disable_custom_all_reduce", False), + ("enable_mscclpp", True), ("enable_torch_symm_mem", True), + ("enable_symm_mem", True), ("enable_nccl_nvls", True), + ("disable_flashinfer_autotune", False), +]) +def test_validate_profile_rejects_unsupported(hooks_module, field, value): + args = _supported_args() + setattr(args, field, value) + with pytest.raises(RuntimeError): + hooks_module._validate_profile(args) +``` + +Provide a `hooks_module` fixture that mocks `sglang`, `sglang.srt`, `sglang.srt.plugins`, `sglang.srt.plugins.hook_registry` (with a `HookRegistry` object exposing a `register` classmethod recording calls and a `HookType` enum), then loads `hooks.py` via `_load_module`. + +- [ ] **Step B2: Run and verify RED** + +Run: `python3 -m pytest tests/test_sglang_profile_validation.py -q` +Expected: FAIL — `_validate_profile` does not exist. + +- [ ] **Step B3: Implement `_validate_profile` and validation hook** + +Add to `hooks.py`: + +```python +def _validate_profile(server_args) -> None: + """Fail closed outside the validated dense TP=2 NCCL P2P/IPC envelope. + + Raised from an AFTER hook on ServerArgs.__post_init__ so SGLang's plugin + loader does not swallow it (load_plugins / apply_hooks swallow activation + and apply-time errors, but not errors from a subsequently executed hook). + """ + problems = [] + if getattr(server_args, "tp_size", 1) != 2: + problems.append(f"tp_size={server_args.tp_size} (require 2)") + if getattr(server_args, "pp_size", 1) != 1: + problems.append(f"pp_size={server_args.pp_size} (require 1)") + if getattr(server_args, "dp_size", 1) != 1: + problems.append(f"dp_size={server_args.dp_size} (require 1)") + if getattr(server_args, "enable_dp_attention", False): + problems.append("enable_dp_attention (unsupported)") + if not getattr(server_args, "disable_custom_all_reduce", False): + problems.append("custom all-reduce enabled (require --disable-custom-all-reduce)") + for flag in ("enable_mscclpp", "enable_torch_symm_mem", "enable_symm_mem", "enable_nccl_nvls"): + if getattr(server_args, flag, False): + problems.append(f"{flag} enabled (unsupported)") + if not getattr(server_args, "disable_flashinfer_autotune", False): + problems.append("FlashInfer autotune enabled (require --disable-flashinfer-autotune)") + if problems: + raise RuntimeError( + "[Foundry] SGLang launch outside the validated TP=2 NCCL P2P/IPC " + "profile: " + "; ".join(problems) + ) +``` + +- [ ] **Step B4: Run and verify GREEN** + +Run: `python3 -m pytest tests/test_sglang_profile_validation.py -q` +Expected: PASS. + +- [ ] **Step B5: Write the failing hook-registration test** + +Create `tests/test_sglang_hook_targets.py`. With SGLang mocked and `FOUNDRY_GRAPH_EXTENSION_CONFIG_PATH` set (config mode `save`), call `hooks.register_hooks()` and assert the recorded `HookRegistry.register` targets and hook types include the exact dotted strings: + +```python +EXPECTED = { + ("sglang.srt.server_args.ServerArgs.__post_init__", "AFTER"), + ("sglang.srt.entrypoints.engine.Engine._launch_scheduler_processes", "AROUND"), + ("sglang.srt.model_executor.model_runner.ModelRunner.init_torch_distributed", "AROUND"), +} + + +def test_register_hooks_registers_targets(hooks_module_active): + hooks_module_active.register_hooks() + recorded = set(hooks_module_active._REGISTRY_CALLS) + assert EXPECTED <= recorded +``` + +The `HookRegistry` mock records `(target, hook_type.name)` tuples into a module-visible list. Reset `_INSTALLED` to `False` in the fixture before calling. + +- [ ] **Step B6: Run and verify RED** + +Run: `python3 -m pytest tests/test_sglang_hook_targets.py -q` +Expected: FAIL — `register_hooks()` from Task 1 registers nothing. + +- [ ] **Step B7: Implement registry-based spawn, rank/VMM, and validation registration** + +In `hooks.py`, replace the Task-A stub body of `register_hooks()` so it registers hooks through `HookRegistry`: + +```python +def register_hooks() -> None: + global _INSTALLED + if _INSTALLED: + return + from sglang.srt.plugins.hook_registry import HookRegistry, HookType + + HookRegistry.register( + "sglang.srt.server_args.ServerArgs.__post_init__", + _after_post_init, HookType.AFTER) + HookRegistry.register( + "sglang.srt.entrypoints.engine.Engine._launch_scheduler_processes", + _around_launch_scheduler_processes, HookType.AROUND) + HookRegistry.register( + "sglang.srt.model_executor.model_runner.ModelRunner.init_torch_distributed", + _around_init_torch_distributed, HookType.AROUND) + _register_memory_pool_hooks(HookRegistry, HookType) # Task 3 + _register_capture_hooks(HookRegistry, HookType) # Task 4 + _INSTALLED = True + logger.info("[Foundry] SGLang hooks registered") +``` + +Add the hook bodies (signatures follow `HookRegistry` AROUND convention `fn(original, *args, **kwargs)` where the first positional after `original` is the bound `self`; confirm the exact callback signature against `hook_registry._apply_target` in the pinned checkout before coding and adjust if it differs): + +```python +def _after_post_init(original_result, self): + _validate_profile(self) + return original_result + + +def _around_launch_scheduler_processes(original, *args, **kwargs): + if get_graph_extension_mode() != CUDAGraphExtensionMode.NONE: + rt.setup_ld_preload_env() + return original(*args, **kwargs) + + +def _around_init_torch_distributed(original, self, *args, **kwargs): + mode = get_graph_extension_mode() + if mode == CUDAGraphExtensionMode.NONE: + return original(self, *args, **kwargs) + rt.setup_graph_extension( + self.server_args, + tp_rank=self.ps.tp_rank, + pp_rank=self.ps.pp_rank, + dp_rank=self.ps.dp_rank, + ) + rt.log_alloc_offset("after_setup_graph_ext") + result = original(self, *args, **kwargs) + rt.log_alloc_offset("after_init_torch_dist") + rt.skip_to_scratch_boundary() + rt.log_alloc_offset("after_scratch_skip") + return result +``` + +Delete the redundant explicit `set_device` (device binding now happens in `ModelRunner.__init__` before `init_torch_distributed`). Remove the old `_patch_init_torch_distributed`, `_patch_spawn_sites`, and `_resolve_dp_rank` direct-setattr code; `self.ps.dp_rank` supplies the rank directly. + +- [ ] **Step B8: Simplify `config._resolve_dp_rank` usage** + +`compute_workspace_rank(server_args, tp_rank, pp_rank, dp_rank)` in `config.py:134-142` is unchanged. Remove `hooks._resolve_dp_rank` (lines 38-66); the rank now comes from `self.ps.dp_rank`. Confirm on the pinned checkout that `ParallelState.dp_rank` is `Optional[int]` and pass it through unchanged. + +- [ ] **Step B9: Run and verify GREEN** + +Run: + +```bash +python3 -m pytest tests/test_sglang_hook_targets.py tests/test_sglang_profile_validation.py -q +python3 -m compileall -q python/foundry/integration/sglang +``` + +Expected: PASS and compile clean. + +- [ ] **Step B10: Commit** + +```bash +git add \ + python/foundry/integration/sglang/hooks.py \ + python/foundry/integration/sglang/config.py \ + tests/test_sglang_hook_targets.py \ + tests/test_sglang_profile_validation.py +git commit -m "feat(sglang): register main hooks and enforce TP=2 profile via ParallelState" +``` + +--- + +### Task 3: Memory-Pool Config SAVE/LOAD Reuse + +**Files:** + +- Modify: `python/foundry/integration/sglang/hooks.py` (add `_register_memory_pool_hooks`, remove `_patch_init_memory_pool`) +- Test: `tests/test_sglang_memory_pool_reuse.py` + +**Interfaces:** + +- Consumes: `ModelRunner.alloc_memory_pool(self, memory_pool_config: Optional[MemoryPoolConfig] = None)`; `KVCacheConfigurator._resolve_memory_pool_config(self, pre_model_load_memory)`; `MemoryPoolConfig` (`pool_configurator.py`); `runtime.create_warmup_state(dict) -> WarmupState`, `runtime.save_warmup_state`, `runtime.load_warmup_state() -> WarmupState`. +- Produces: `hooks._register_memory_pool_hooks(HookRegistry, HookType)`, an AROUND `alloc_memory_pool` hook that serializes the resolved config on SAVE, and a LOAD hook on `_resolve_memory_pool_config` returning the saved config. + +- [ ] **Step C1: Write the failing reuse test** + +Create `tests/test_sglang_memory_pool_reuse.py`. With SGLang mocked, drive the SAVE `alloc_memory_pool` hook against a fake `ModelRunner` whose `memory_pool_config` is a fake dataclass, assert `save_warmup_state` received `asdict(...)`; then drive the LOAD `_resolve_memory_pool_config` hook and assert it returns the reconstructed saved config and never calls the (fake) profiler: + +```python +def test_load_resolver_returns_saved_config_without_profiling(mp_hooks, fake_runner): + profiled = [] + fake_runner._profile_available_bytes = lambda *a, **k: profiled.append(1) + mp_hooks.set_mode("load") + resolved = mp_hooks._load_resolve_memory_pool_config( + original=lambda self, m: profiled.append("orig"), + self=fake_runner.kv_cache_configurator, + pre_model_load_memory=123, + ) + assert resolved is fake_runner.kv_cache_configurator._foundry_saved_config + assert profiled == [] +``` + +- [ ] **Step C2: Run and verify RED** + +Run: `python3 -m pytest tests/test_sglang_memory_pool_reuse.py -q` +Expected: FAIL — the memory-pool hooks do not exist. + +- [ ] **Step C3: Implement the memory-pool hooks** + +Add to `hooks.py`: + +```python +def _register_memory_pool_hooks(HookRegistry, HookType) -> None: + HookRegistry.register( + "sglang.srt.model_executor.model_runner.ModelRunner.alloc_memory_pool", + _around_alloc_memory_pool, HookType.AROUND) + HookRegistry.register( + "sglang.srt.mem_cache.kv_cache_configurator." + "KVCacheConfigurator._resolve_memory_pool_config", + _around_resolve_memory_pool_config, HookType.AROUND) + + +def _around_alloc_memory_pool(original, self, *args, **kwargs): + mode = get_graph_extension_mode() + if mode == CUDAGraphExtensionMode.NONE: + return original(self, *args, **kwargs) + if mode == CUDAGraphExtensionMode.LOAD: + import torch + + state = rt.load_warmup_state() + if not state.memory_pool_config: + raise RuntimeError("Foundry LOAD requires a saved memory_pool_config") + # Mirror SAVE's get_available_gpu_memory(empty_cache=True) side effect so + # the caching allocator holds no segments SAVE released; without this the + # VMM cursor drifts from the recorded start_base_addr. + torch.cuda.empty_cache() + rt.log_alloc_offset("before_alloc_memory_pool") + result = original(self, *args, **kwargs) + rt.log_alloc_offset("after_alloc_memory_pool") + logger.info("[Foundry] SGLang reused saved memory pool config") + return result + rt.log_alloc_offset("before_alloc_memory_pool") + result = original(self, *args, **kwargs) + rt.log_alloc_offset("after_alloc_memory_pool") + from dataclasses import asdict + + state = rt.create_warmup_state(asdict(self.memory_pool_config)) + rt.save_warmup_state(state) + return result + + +def _around_resolve_memory_pool_config(original, self, pre_model_load_memory): + if get_graph_extension_mode() != CUDAGraphExtensionMode.LOAD: + return original(self, pre_model_load_memory) + from sglang.srt.model_executor.pool_configurator import MemoryPoolConfig + + saved = rt.load_warmup_state().memory_pool_config + if not saved: + raise RuntimeError("Foundry LOAD requires a saved memory_pool_config") + return MemoryPoolConfig(**saved) +``` + +Keep `dataclasses.asdict` imported at module scope in `hooks.py` (it already is, `hooks.py:12`); the `torch` and `MemoryPoolConfig` imports stay lazy under the documented SGLang/CUDA exception. Remove `_patch_init_memory_pool` (lines 192-233) — its `model_runner_kv_cache_mixin` target no longer exists. + +- [ ] **Step C4: Run and verify GREEN** + +Run: `python3 -m pytest tests/test_sglang_memory_pool_reuse.py -q` +Expected: PASS. + +- [ ] **Step C5: Verify against pinned source before GPU** + +Confirm on the clean pinned checkout that `KVCacheConfigurator.configure` calls `_resolve_memory_pool_config` for the target (non-draft) worker and that `_resolve_memory_pool_config` performs the profiling (`_profile_available_bytes`). If the method name or call shape differs, resolve from source and update this task before proceeding. Then: + +```bash +python3 -m compileall -q python/foundry/integration/sglang +git diff --check origin/main...HEAD +``` + +Expected: compile and whitespace clean. + +- [ ] **Step C6: Commit** + +```bash +git add \ + python/foundry/integration/sglang/hooks.py \ + tests/test_sglang_memory_pool_reuse.py +git commit -m "feat(sglang): reuse saved MemoryPoolConfig on LOAD without profiling" +``` + +--- + +### Task 4: ShapeKey/Full-Backend SAVE and Decode-Runner LOAD + +**Files:** + +- Create: `python/foundry/integration/sglang/backend.py` (`FoundryFullCudaGraphBackend`) +- Modify: `python/foundry/integration/sglang/graph_ops.py` (ShapeKey keying, `captured_req_width`, `backend._graphs`/`_outputs`, pool helpers) +- Modify: `python/foundry/integration/sglang/hooks.py` (add `_register_capture_hooks`, remove `_patch_cuda_graph_capture`/`_patch_load_model`/`_patch_kernel_warmup`) +- Test: `tests/test_sglang_shape_key_archive.py` +- Test: `tests/test_sglang_capture_hooks.py` + +**Interfaces:** + +- Consumes: `FullCudaGraphBackend(cuda_graph_runner, *, enable_memory_saver=False)` with `capture_one(shape_key, forward_fn, capture_inputs=None, post_warmup_hook=None)`, `capture_session(stream)`, `replay(shape_key, static_forward_batch, **kwargs)`, and `_graphs`/`_outputs`/`_pool`/`_device_module`/`_capture_stream`; `ShapeKey(size, stream_idx=None, variant_label=None)`; `DecodeCudaGraphRunner.capture(self)`, `.capture_bs`, `.captured_req_width`; `get_or_create_global_graph_memory_pool`, `set_graph_pool_id`; `init_forward_metadata_out_graph(fb, in_capture=True)`. +- Produces: `backend.FoundryFullCudaGraphBackend`; `graph_ops.save_graph(graph, output, shape_key)`, `graph_ops.load_all_graphs(runner)` populating `backend._graphs`/`_outputs` keyed by `ShapeKey`; `hooks._register_capture_hooks(HookRegistry, HookType)`. + +- [ ] **Step D1: Write the failing ShapeKey archive round-trip test** + +Create `tests/test_sglang_shape_key_archive.py`. With a fake `ShapeKey` dataclass, assert `graph_ops` serializes all three fields and reconstructs them, and that a non-None `stream_idx`/`variant_label` either round-trips or fails closed: + +```python +def test_shape_key_filename_roundtrip(graph_ops): + key = graph_ops.ShapeKeyLike(size=8, stream_idx=None, variant_label=None) + name = graph_ops._graph_filename(3, key) + parsed = graph_ops._shape_key_from_filename(name) + assert parsed.size == 8 and parsed.stream_idx is None and parsed.variant_label is None + + +def test_shape_key_rejects_unsupported_variant(graph_ops): + key = graph_ops.ShapeKeyLike(size=8, stream_idx=1, variant_label="lora") + with pytest.raises(ValueError): + graph_ops._graph_filename(3, key) +``` + +- [ ] **Step D2: Run and verify RED** + +Run: `python3 -m pytest tests/test_sglang_shape_key_archive.py -q` +Expected: FAIL — `_shape_key_from_filename` and the fail-closed guard do not exist. + +- [ ] **Step D3: Migrate `graph_ops.py` to ShapeKey** + +In `graph_ops.py`, replace `_batch_size_from_key` (lines 33-40) and `_graph_filename` (lines 43-45) with `ShapeKey`-aware helpers that assert `stream_idx is None` and `variant_label is None` for the dense profile, and add `_shape_key_from_filename`. Reconstruct `ShapeKey(size=bs)` in `_scan_graph_files` and `load_all_graphs` (lazily import `ShapeKey` from `sglang.srt.model_executor.runner.shape_key`). Update `build_capture_fb_view` (lines 252-275) to read `cuda_graph_runner.captured_req_width` instead of `num_tokens_per_bs`. In `load_all_graphs` (lines 324-363) populate the backend maps instead of runner maps: assign `runner.backend._graphs[shape_key] = graph` and `runner.backend._outputs[shape_key] = _unpack_output(tensors)`. Repoint the graph-pool helpers from `cuda_graph_runner.get_global_graph_memory_pool` to the current location (`model_executor.runner_utils.pool` / `runner` exports) after confirming the exact symbol on the pinned checkout. + +- [ ] **Step D4: Run and verify GREEN** + +Run: `python3 -m pytest tests/test_sglang_shape_key_archive.py -q` +Expected: PASS. + +- [ ] **Step D5: Write the failing capture-hook test** + +Create `tests/test_sglang_capture_hooks.py`. With SGLang mocked, assert `register_hooks()` records the capture targets: + +```python +CAPTURE_EXPECTED = { + ("sglang.srt.model_executor.runner_backend.full_cuda_graph_backend.FullCudaGraphBackend", "REPLACE"), + ("sglang.srt.model_executor.runner.decode_cuda_graph_runner.DecodeCudaGraphRunner.capture", "AROUND"), +} + + +def test_capture_hooks_registered(hooks_module_active): + hooks_module_active.register_hooks() + assert CAPTURE_EXPECTED <= set(hooks_module_active._REGISTRY_CALLS) +``` + +Note: SAVE uses AROUND and LOAD uses REPLACE on `DecodeCudaGraphRunner.capture`; the registered hook body branches on mode internally, so the recorded hook type for the capture target is a single `AROUND` registration whose body performs the LOAD replacement by short-circuiting the original. Assert the `AROUND` registration and document the internal mode branch. + +- [ ] **Step D6: Run and verify RED** + +Run: `python3 -m pytest tests/test_sglang_capture_hooks.py -q` +Expected: FAIL — capture hooks unregistered. + +- [ ] **Step D7: Implement `FoundryFullCudaGraphBackend`** + +Create `python/foundry/integration/sglang/backend.py`: + +```python +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the Foundry project +"""Foundry SAVE override of SGLang's FullCudaGraphBackend.""" + +from __future__ import annotations + +import logging + +from foundry.integration.sglang.config import ( + CUDAGraphExtensionMode, + get_graph_extension_mode, +) + +logger = logging.getLogger(__name__) + + +def build_foundry_full_backend_class(): + """Return a FullCudaGraphBackend subclass that captures into Foundry graphs + on SAVE and delegates otherwise. Built lazily so SGLang stays optional.""" + from sglang.srt.model_executor.runner_backend.full_cuda_graph_backend import ( + FullCudaGraphBackend, + ) + + class FoundryFullCudaGraphBackend(FullCudaGraphBackend): + def capture_one(self, shape_key, forward_fn, capture_inputs=None, + post_warmup_hook=None): + if get_graph_extension_mode() != CUDAGraphExtensionMode.SAVE: + return super().capture_one( + shape_key, forward_fn, capture_inputs, post_warmup_hook) + from foundry.graph import CUDAGraph as FoundryCUDAGraph + from foundry.graph import graph as foundry_graph_ctx + from foundry.integration.sglang.graph_ops import save_graph + + # Skip the two ordinary warmup forwards so JIT/autotune allocations + # land inside the captured graph and replay on LOAD. post_warmup_hook + # semantics are preserved by invoking it once before capture. + if post_warmup_hook is not None: + post_warmup_hook() + graph = FoundryCUDAGraph() + with foundry_graph_ctx(graph, pool=self._pool, stream=self._capture_stream): + output = forward_fn() + self._graphs[shape_key] = graph + self._outputs[shape_key] = output + save_graph(graph, output, shape_key) + return None + + return FoundryFullCudaGraphBackend +``` + +Confirm on the pinned checkout that `capture_one` returns `None` (the runner reads `self._graphs`/`self._outputs`), and that `post_warmup_hook` is invoked once per shape by the runner rather than per warmup forward; adjust if source differs. + +- [ ] **Step D8: Implement the capture hooks in `hooks.py`** + +Add: + +```python +def _register_capture_hooks(HookRegistry, HookType) -> None: + from foundry.integration.sglang.backend import build_foundry_full_backend_class + + HookRegistry.register( + "sglang.srt.model_executor.runner_backend.full_cuda_graph_backend.FullCudaGraphBackend", + build_foundry_full_backend_class(), HookType.REPLACE) + HookRegistry.register( + "sglang.srt.model_executor.runner.decode_cuda_graph_runner.DecodeCudaGraphRunner.capture", + _around_decode_capture, HookType.AROUND) + + +def _around_decode_capture(original, self, *args, **kwargs): + mode = get_graph_extension_mode() + if mode == CUDAGraphExtensionMode.NONE: + return original(self, *args, **kwargs) + if mode == CUDAGraphExtensionMode.SAVE: + result = original(self, *args, **kwargs) + from foundry.integration.sglang.graph_ops import pack_fatbins, save_graph_manifest + + save_graph_manifest() + pack_fatbins() + rt.capture_final_alloc_offset() + return result + # LOAD: replace the capture loop entirely. + from sglang.srt.distributed.device_communicators.pynccl_allocator import ( + set_graph_pool_id, + ) + from foundry.integration.sglang.graph_ops import ( + get_or_create_graph_pool, + initialize_all_attention_metadata, + load_all_graphs, + ) + + backend = self.backend + backend._pool = get_or_create_graph_pool(backend) + set_graph_pool_id(backend._pool) + rt.log_alloc_offset("before_preallocate") + rt.preallocate_for_load_mode() + rt.log_alloc_offset("after_preallocate") + if hasattr(self.attn_backend, "indices_updater_decode"): + initialize_all_attention_metadata(self) + rt.log_alloc_offset("after_pre_init") + load_all_graphs(self) + rt.log_alloc_offset("after_load_all_graphs") + return None +``` + +`get_or_create_graph_pool(backend)` is a small `graph_ops` helper wrapping `get_or_create_global_graph_memory_pool(backend._device_module)`; confirm the exact upstream symbol/location on the pinned checkout. Remove `_patch_cuda_graph_capture` (lines 267-665), `_patch_load_model` (236-246), and `_patch_kernel_warmup` (249-264): their targets are removed upstream and warmup is pinned off by profile. + +- [ ] **Step D9: Assert the narrow output shape in `graph_ops`** + +In `_pack_output` (lines 48-59), keep the `LogitsProcessorOutput.next_token_logits` path but add an explicit assertion that no unsupported speculative/hidden-state fields are populated on the decode output, failing closed if they are. Confirm the decode output type on the pinned checkout is `LogitsProcessorOutput`. + +- [ ] **Step D10: Run and verify GREEN** + +Run: + +```bash +python3 -m pytest \ + tests/test_sglang_shape_key_archive.py \ + tests/test_sglang_capture_hooks.py -q +python3 -m compileall -q python/foundry/integration/sglang +``` + +Expected: PASS and compile clean. + +- [ ] **Step D11: Commit** + +```bash +git add \ + python/foundry/integration/sglang/backend.py \ + python/foundry/integration/sglang/graph_ops.py \ + python/foundry/integration/sglang/hooks.py \ + tests/test_sglang_shape_key_archive.py \ + tests/test_sglang_capture_hooks.py +git commit -m "feat(sglang): bind SAVE/LOAD to FullCudaGraphBackend and ShapeKey" +``` + +--- + +### Task 5: Clean-Checkout Target-Resolution CPU/Modal Contract + +**Files:** + +- Create: `tests/test_sglang_target_resolution.py` +- Create: `tests/modal_sglang_target_resolution.py` +- Modify: `python/foundry/integration/sglang/hooks.py` (remove the legacy `install_hooks` shim) +- Modify: `python/foundry/integration/sglang/__init__.py` (drop `install_hooks` re-export if present) + +**Interfaces:** + +- Consumes: `importlib.metadata.entry_points(group="sglang.srt.plugins")`; `pkgutil.resolve_name("")`; the exact hook-target strings registered in Tasks 2–4. +- Produces: a CPU test asserting the entry point is declared, and a Modal contract that installs the pinned SGLang + editable Foundry and asserts every hook target resolves and that `load_plugins()` with no config env is a no-op. + +- [ ] **Step E1: Write the failing entry-point/target-string CPU test** + +Create `tests/test_sglang_target_resolution.py`. It does not import SGLang; it asserts the entry point is declared in `pyproject.toml` and that the hook-target strings Foundry registers are a fixed, reviewed list (guards against silent target renames in Foundry’s own source): + +```python +import tomllib +from pathlib import Path + +REPO_ROOT = Path(__file__).parents[1] +HOOK_TARGETS = { + "sglang.srt.server_args.ServerArgs.__post_init__", + "sglang.srt.entrypoints.engine.Engine._launch_scheduler_processes", + "sglang.srt.model_executor.model_runner.ModelRunner.init_torch_distributed", + "sglang.srt.model_executor.model_runner.ModelRunner.alloc_memory_pool", + "sglang.srt.mem_cache.kv_cache_configurator.KVCacheConfigurator._resolve_memory_pool_config", + "sglang.srt.model_executor.runner_backend.full_cuda_graph_backend.FullCudaGraphBackend", + "sglang.srt.model_executor.runner.decode_cuda_graph_runner.DecodeCudaGraphRunner.capture", +} + + +def test_entry_point_declared(): + data = tomllib.loads((REPO_ROOT / "pyproject.toml").read_text()) + eps = data["project"]["entry-points"]["sglang.srt.plugins"] + assert eps["foundry"] == "foundry_sglang_plugin:activate" + + +def test_hook_targets_match_registered(monkeypatch, tmp_path): + # Load hooks.py with SGLang mocked and compare the recorded HookRegistry + # targets against HOOK_TARGETS exactly (no missing, no extra). Reuses the + # module loader and module-mocking helpers copied from + # tests/test_sglang_attention_metadata.py (_install_module, _load_module). + import sys + from enum import Enum + from types import ModuleType, SimpleNamespace + + recorded = [] + + class HookType(Enum): + BEFORE = "before" + AFTER = "after" + AROUND = "around" + REPLACE = "replace" + + class HookRegistry: + @classmethod + def register(cls, target, fn, hook_type): + recorded.append(target) + + for name in ("sglang", "sglang.srt", "sglang.srt.plugins", + "sglang.srt.plugins.hook_registry"): + module = ModuleType(name) + monkeypatch.setitem(sys.modules, name, module) + hook_registry_mod = sys.modules["sglang.srt.plugins.hook_registry"] + hook_registry_mod.HookRegistry = HookRegistry + hook_registry_mod.HookType = HookType + + cfg = tmp_path / "cfg.toml" + cfg.write_text('mode = "save"\n') + monkeypatch.setenv("FOUNDRY_GRAPH_EXTENSION_CONFIG_PATH", str(cfg)) + + import importlib.util + + hooks_path = REPO_ROOT / "python" / "foundry" / "integration" / "sglang" / "hooks.py" + spec = importlib.util.spec_from_file_location("foundry_hooks_under_test", hooks_path) + hooks = importlib.util.module_from_spec(spec) + spec.loader.exec_module(hooks) + hooks.load_graph_extension_config(str(cfg)) + hooks._INSTALLED = False + hooks.register_hooks() + + assert set(recorded) == HOOK_TARGETS +``` + +- [ ] **Step E2: Run and verify RED** + +Run: `python3 -m pytest tests/test_sglang_target_resolution.py -q` +Expected: `test_entry_point_declared` PASSES (entry point added in Task 1); `test_hook_targets_match_registered` FAILS only if the registered target set differs from `HOOK_TARGETS` — it should already match given Tasks 2–4, so this test is a regression guard rather than a RED-first test. If it fails, reconcile the target list with the registered hooks. + +- [ ] **Step E3: Confirm no extra or renamed Foundry targets** + +If `test_hook_targets_match_registered` fails, either a hook was added/removed in Tasks 2–4 or a target string drifted. Reconcile `HOOK_TARGETS` and the `HookRegistry.register` calls so they match exactly, then rerun. + +- [ ] **Step E4: Run and verify GREEN** + +Run: `python3 -m pytest tests/test_sglang_target_resolution.py -q` +Expected: PASS. + +- [ ] **Step E5: Remove the legacy `install_hooks` shim** + +Delete `install_hooks(server_args)` (`hooks.py:69-100`) and the `_resolve_dp_rank`/`_patch_*`/`_ep_lazy_init_needed` helpers already superseded. Keep `_patch_deepep`/`bootstrap_deepep_buffer` only if they are naturally isolated behind a DeepEP gate and not referenced by the dense TP path; otherwise remove them in this task and note the removal. Remove any `install_hooks` re-export from `python/foundry/integration/sglang/__init__.py`. + +- [ ] **Step E6: Author the Modal clean-checkout contract** + +Create `tests/modal_sglang_target_resolution.py`: a Modal function reusing the `tests/modal_sglang_tp.py` image (pinned `SGLANG_BASE=a25164bda355250015c0ba6d8821f0d4b37abe7d`, editable Foundry) that, on CPU or a single GPU, asserts: + +```python +from importlib.metadata import entry_points +import pkgutil + +def _check(): + eps = {ep.name: ep.value for ep in entry_points(group="sglang.srt.plugins")} + assert eps.get("foundry") == "foundry_sglang_plugin:activate" + targets = [ + "sglang.srt.server_args.ServerArgs.__post_init__", + "sglang.srt.entrypoints.engine.Engine._launch_scheduler_processes", + "sglang.srt.model_executor.model_runner.ModelRunner.init_torch_distributed", + "sglang.srt.model_executor.model_runner.ModelRunner.alloc_memory_pool", + "sglang.srt.mem_cache.kv_cache_configurator.KVCacheConfigurator._resolve_memory_pool_config", + "sglang.srt.model_executor.runner_backend.full_cuda_graph_backend.FullCudaGraphBackend", + "sglang.srt.model_executor.runner.decode_cuda_graph_runner.DecodeCudaGraphRunner.capture", + ] + for target in targets: + pkgutil.resolve_name(target) # raises on rename + # Inactive plugin must be a no-op with no config env. + import os + os.environ.pop("FOUNDRY_GRAPH_EXTENSION_CONFIG_PATH", None) + from sglang.srt.plugins import load_plugins + load_plugins() +``` + +Any resolution failure means the pinned API drifted; resolve from source and update Tasks 2–4 before proceeding. + +- [ ] **Step E7: Local gates and commit** + +Run: + +```bash +python3 -m pytest tests/test_sglang_target_resolution.py -q +python3 -m compileall -q python/foundry/integration/sglang tests/modal_sglang_target_resolution.py +git diff --check origin/main...HEAD +``` + +Expected: CPU test passes; compile clean; no whitespace errors. Then: + +```bash +git add \ + python/foundry/integration/sglang/hooks.py \ + python/foundry/integration/sglang/__init__.py \ + tests/test_sglang_target_resolution.py \ + tests/modal_sglang_target_resolution.py +git commit -m "test(sglang): assert clean-checkout entry point and hook-target resolution" +``` + +--- + +### Task 6: Documentation and Final 2x H100 Acceptance + +**Files:** + +- Modify: `docs/sglang/overview.md`, `docs/sglang/hooks.md`, `docs/sglang/direct-edits.md`, `docs/sglang/save-load-workflow.md`, `docs/sglang/memory-lifecycle.md`, `docs/sglang/memory-consistency.md` +- Modify: `recipe/experimental/README.md`, `recipe/sglang/README.md` +- Modify: `README.md`, `ROADMAP.md`, `RELEASE.md` +- Test/Evidence: `tests/modal_sglang_tp.py`, PR #6 description + +**Interfaces:** + +- Consumes: the final pushed commit SHA; the pinned SGLang SHA and model revision; all local, clean-checkout, and CPU-contract results. +- Produces: documentation stating zero upstream edits + plugin activation, the exact pinned profile, and separated historical fork evidence; and pinned-commit 2x H100 baseline/SAVE/SAVE2/LOAD/LOAD-restart evidence. + +- [ ] **Step F1: Rewrite the direct-edits/activation docs** + +In `docs/sglang/direct-edits.md` and `docs/sglang/overview.md`, replace the "47 lines across 4 files" fork narrative with: zero upstream edits, activation through the `sglang.srt.plugins` entry point `foundry = "foundry_sglang_plugin:activate"`, config via `FOUNDRY_GRAPH_EXTENSION_CONFIG_PATH`, `SGLANG_PLUGINS=foundry`, and the pinned SHA `a25164bda355250015c0ba6d8821f0d4b37abe7d`. In `docs/sglang/hooks.md`, replace the monolithic `CudaGraphRunner` patch descriptions with the `FullCudaGraphBackend.capture_one` SAVE / `DecodeCudaGraphRunner.capture` LOAD mapping and the `KVCacheConfigurator._resolve_memory_pool_config` reuse hook. Preserve historical fork H200/H100 evidence, clearly labeled as pre-port and separate from pinned-main evidence. + +- [ ] **Step F2: Reconcile scope statements** + +In `README.md`, `ROADMAP.md`, and `RELEASE.md`, keep the TP status scoped to the pinned single-host NCCL P2P/IPC TP=2 recipe and note the new pin/activation model. Use wording equivalent to: + +```markdown +SGLang TP=2 is validated only against the pinned upstream commit +a25164bda355250015c0ba6d8821f0d4b37abe7d via the sglang.srt.plugins entry point. +TP>2, multi-host, custom/symmetric/fused all-reduce, CUMEM, NVLS, EP, +DP-attention, and unpinned SGLang revisions remain unsupported. +``` + +- [ ] **Step F3: Commit documentation** + +```bash +git add README.md ROADMAP.md RELEASE.md docs recipe +git commit -m "docs(sglang): document zero-edit plugin activation and pinned main scope" +``` + +- [ ] **Step F4: Pre-push local gates** + +Run: + +```bash +python3 -m pytest \ + tests/test_sglang_plugin_activation.py \ + tests/test_sglang_profile_validation.py \ + tests/test_sglang_hook_targets.py \ + tests/test_sglang_memory_pool_reuse.py \ + tests/test_sglang_shape_key_archive.py \ + tests/test_sglang_capture_hooks.py \ + tests/test_sglang_target_resolution.py \ + tests/test_sglang_tp_recipe.py -q +python3 -m compileall -q python tests +pre-commit run --all-files +git diff --check origin/main...HEAD +``` + +Expected: all pass. Confirm setup completion first: if `/tmp/cursor/async-install/install-user.status` exists it must read `0`. + +- [ ] **Step F5: Push and update the draft PR** + +Run: + +```bash +git status --short +git push -u origin cursor/tp-upstream-consolidation-fe33 +``` + +Update PR #6 with the pre-test commit and exact local/clean-checkout results. Do not merge or change merge state. + +- [ ] **Step F6: Run the clean-checkout Modal contract** + +Run: + +```bash +FOUNDRY_REF="$(git rev-parse HEAD)" \ + modal run --env rahul-dev tests/modal_sglang_target_resolution.py +``` + +Expected: entry point discoverable, every hook target resolves at the pin, inactive `load_plugins()` is a no-op. Any failure means the pinned API drifted — resolve from source and revisit Tasks 2–4 before the acceptance run. + +- [ ] **Step F7: Run the final-head 2x H100 acceptance** + +Run: + +```bash +FOUNDRY_REF="$(git rev-parse HEAD)" \ + modal run --env rahul-dev tests/modal_sglang_tp.py +``` + +Expected across baseline / SAVE / SAVE2 / LOAD: + +- `EXPECTED_GRAPH_COUNT` graphs per rank, reproducible across SAVE and SAVE2; +- reproducible semantic graph fingerprints and per-rank final offsets; +- LOAD final offset equals SAVE final offset; +- all archived graphs restored and zero recaptured on LOAD; +- restored-graph replay at batch size 1 and at least one larger batch; +- every NCCL channel uses P2P/IPC (`EXPECTED_NCCL_CHANNELS` per rank); +- token-identical baseline/SAVE2/LOAD outputs for sequential and nonempty concurrent prompts; +- no CUDA/NCCL/MMU/Xid/illegal-access/worker-death/scheduler error signature. + +- [ ] **Step F8: Run one additional fresh LOAD restart** + +Rerun LOAD against the same immutable archive (harness restart mode or a second `load` phase). +Expected: identical graph inventory, offsets, and deterministic outputs without stale IPC broker state. + +- [ ] **Step F9: Record evidence and final gates** + +Update the SGLang docs and PR #6 with: exact final commit; pinned SGLang SHA `a25164bda355250015c0ba6d8821f0d4b37abe7d`; model revision `b968826d9c46dd6066d109eabc6255188de91218`; H100 model, CUDA, driver, PyTorch, NCCL, FlashInfer, sgl-kernel versions; run URL; graph counts and offsets; SAVE2/LOAD and restart comparisons; and the explicit unsupported scope. Then: + +```bash +git add docs recipe README.md ROADMAP.md RELEASE.md +git commit -m "docs(sglang): record pinned-main 2xH100 TP validation evidence" +git push origin cursor/tp-upstream-consolidation-fe33 +``` + +If measured values did not change any doc, do not create an empty commit. + +--- + +## Self-Review + +**1. Spec coverage.** Every design decision in `docs/superpowers/specs/2026-07-23-sglang-main-port-design.md` maps to a task: plugin activation + pin migration (A), profile validation + rank/VMM on `self.ps` + spawn LD_PRELOAD (B), memory-pool reuse via `_resolve_memory_pool_config` (C), `FullCudaGraphBackend`/`ShapeKey` SAVE + `DecodeCudaGraphRunner.capture` LOAD + output-shape assertion (D), clean-checkout entry-point/target resolution (E), docs + final 2x H100 acceptance/restart (F). The fail-closed-through-executed-hook requirement is in B3/B7 (AFTER `__post_init__`). + +**2. Placeholder scan.** Every code step shows complete, runnable code, including the E1 target-resolution test with its inline SGLang mock. No "TBD/handle edge cases/similar to Task N", no bare `...` bodies. References to `foundry-org/sglang`, `foundry_shim`, and `foundry_graph_extension_config_path` appear only as the fork artifacts being deleted, never as assumptions the port relies on. + +**3. Type consistency.** `ShapeKey(size, stream_idx, variant_label)`, `MemoryPoolConfig(**saved)`, `capture_one(shape_key, forward_fn, capture_inputs=None, post_warmup_hook=None)`, `replay(shape_key, static_forward_batch, **kwargs)`, `self.ps.{tp_rank,pp_rank,dp_rank}`, `captured_req_width`, and the seven hook-target strings are used identically across Tasks 2–5 and the CPU contract. + +**4. Source-of-truth discipline.** Steps B7, C5, D3, D7, D8, and E6 each require confirming the exact API against a clean pinned checkout before coding and resolving any discrepancy from source, per the Global Constraints. diff --git a/docs/superpowers/plans/2026-07-23-sglang-symmetric-tp-treatment.md b/docs/superpowers/plans/2026-07-23-sglang-symmetric-tp-treatment.md new file mode 100644 index 00000000..5681ef46 --- /dev/null +++ b/docs/superpowers/plans/2026-07-23-sglang-symmetric-tp-treatment.md @@ -0,0 +1,698 @@ +# SGLang Symmetric-TP Treatment Implementation Plan + +> **For agentic workers:** REQUIRED SUB-SKILL: Use +> superpowers:subagent-driven-development or superpowers:executing-plans to +> implement this plan task-by-task. Steps use checkbox (`- [ ]`) syntax for +> tracking. + +**Goal:** Replace the disproven SGLang plain-NCCL restore policy with the proven +single-shape torch symmetric-memory path, verify a focused fresh-process +SAVE→LOAD treatment, then run and document the full five-phase matrix. + +**Architecture:** Two new `HookRegistry` AROUND hooks retain the torch +symmetric communicator allocated outside Foundry VMM and route multimem gather +directly through its stable buffer/group. The current full backend synchronizes +ranks before capture, LOAD receives the shared pool explicitly, and the profile +captures only `ShapeKey(size=256)`. Verification is gated: focused treatment +first; full matrix and support documentation only after it verifies. + +**Tech Stack:** Python 3.12, PyTorch 2.11/cu130 symmetric memory, pinned SGLang, +FlashInfer, Foundry CUDA graphs, pytest, Bash, Modal, two H100s. + +## Global constraints + +- Design authority: + `docs/superpowers/specs/2026-07-23-sglang-symmetric-tp-treatment-design.md`. +- Pin SGLang to `a25164bda355250015c0ba6d8821f0d4b37abe7d` and Qwen3-8B to + `b968826d9c46dd6066d109eabc6255188de91218`. +- Plain-NCCL fresh-process SGLang restore is rejected at this pin. Do not retain + it as a supported fallback. +- Preserve zero-edit activation, atomic rank-zero state writes, strict + `MemoryPoolConfig` validation, complete `ShapeKey` handling, interleaved + FlashInfer metadata/graph finishes, one-start loading, archive/process + hardening, and the two-window fatal scanner. +- Do not port device-group rendezvous, synthesized gather state, warmup + forwards, `create_state` interception, independent per-file loading, or + dependency-strip residue. +- Keep SGLang imports lazy only where the existing optional-integration boundary + requires it, and document that exception at the import site. +- CUDA and engine checks run only on H100s in Modal `rahul-dev`; orchestration + parents must not initialize CUDA. +- Every task requires a non-fast implementer model and an independent non-fast + reviewer model. A fast model may not implement or approve any task. +- Stop after Task 2 unless the focused result is `VERIFIED`. Do not reinterpret + `NOT VERIFIED` or `INCONCLUSIVE` as permission to run Task 3. +- Modal Volume auto-commit may retain partial or mutated unique-run artifacts + after a failed phase. Failures stop later phases, reports mark those artifacts + retained, and the affected run is never verified. + +--- + +### Task 1: Stable Symmetric Hooks, Exact Profile, and Pool Invariants + +**Files:** + +- Modify: `python/foundry_sglang_plugin/targets.py` +- Modify: `python/foundry/integration/sglang/hooks.py` +- Modify: `python/foundry/integration/sglang/backend.py` +- Modify: `python/foundry/integration/sglang/graph_ops.py` +- Modify: `recipe/experimental/serve_qwen3-8b_sglang_tp.sh` +- Test: `tests/test_sglang_hook_targets.py` +- Test: `tests/test_sglang_target_resolution.py` +- Test: `tests/modal_sglang_target_resolution.py` +- Test: `tests/test_sglang_profile_validation.py` +- Test: `tests/test_sglang_capture_hooks.py` +- Test: `tests/test_sglang_graph_load.py` +- Test: `tests/test_sglang_tp_recipe.py` + +**Interfaces:** + +- Add targets `TORCH_SYMM_MEM_INIT` for + `sglang.srt.distributed.device_communicators.torch_symm_mem.TorchSymmMemCommunicator.__init__` + and `MULTIMEM_ALL_GATHER_CALL` for + `sglang.srt.distributed.device_communicators.triton_symm_mem_ag.MultimemAllGatherer.__call__`. +- Produce + `_around_torch_symm_mem_init(original, communicator, *args, **kwargs)` and + `_around_multimem_all_gather(original, gatherer, x)`. +- Retain communicators in + `_torch_symm_mem_communicators: dict[tuple[int, ...], Any]`. +- Preserve + `FoundryFullCudaGraphBackend.capture_one(shape_key, forward_fn, + capture_inputs=None, post_warmup_hook=None)`. +- Call + `FoundryCUDAGraph.start_graph_builds(paths, pool=backend._pool, + num_threads=4)` exactly once. + +- [ ] **Step 1: Write target and symmetric-hook contract tests** + +Extend `tests/test_sglang_hook_targets.py` and +`tests/test_sglang_target_resolution.py` so the exact inventory is nine unique +targets and both additions are AROUND hooks. Add callback tests proving: + +```python +assert events == ["stop", "original", "resume"] +assert retained[tuple(group_ranks)] is communicator +assert op_calls == [(x, communicator.group.group_name, expected_buffer)] +assert output.shape == (tokens, world_size * local_hidden) +``` + +Also require inactive passthrough, `resume_allocation_region()` after an +original-constructor exception, signal-pad minimum +`_MAX_BLOCKS * world_size * 4`, rank-tuple replacement, one observed marker per +gatherer, and failures for zero/multiple communicators, non-rank-2, +non-contiguous/non-bfloat16 input, and insufficient capacity. Update +`tests/modal_sglang_target_resolution.py` to resolve and report all nine +targets at `a25164b`. + +- [ ] **Step 2: Run target tests and verify RED** + +Run: + +```bash +python3 -m pytest -q \ + tests/test_sglang_hook_targets.py \ + tests/test_sglang_target_resolution.py +``` + +Expected: failures show seven targets and missing symmetric callbacks. + +- [ ] **Step 3: Implement the two HookRegistry targets** + +In `targets.py`, add both constants to `HOOK_TARGETS`. In `hooks.py`, register +both through `_register_symmetric_collective_hooks(HookRegistry, HookType)`. +The constructor callback must: + +```python +group = kwargs.get("group", args[0] if args else None) +world_size = torch.distributed.get_world_size(group) +signal_pad_bytes = triton_symm_mem_ag._MAX_BLOCKS * world_size * 4 +symm_mem = torch_symm_mem_module.torch_symm_mem +symm_mem.set_signal_pad_size( + max(symm_mem.get_signal_pad_size(), signal_pad_bytes) +) +cge.stop_allocation_region() +try: + result = original(communicator, *args, **kwargs) +finally: + cge.resume_allocation_region() +``` + +After successful construction, retain a non-null `communicator.buffer` by +`tuple(torch.distributed.get_process_group_ranks(communicator.group))` and log +`[Foundry] Torch symmetric-memory buffer address=0x%x`. + +The gather callback must enforce the design's input/capacity guards, select the +sole retained communicator, and execute: + +```python +comm_buff = communicator.buffer[:required_numel].view( + world_size * token_count, + local_hidden_size, +) +torch.ops.symm_mem.multimem_all_gather_out( + x, + communicator.group.group_name, + comm_buff, +) +return ( + comm_buff.view(world_size, token_count, local_hidden_size) + .movedim(0, 1) + .reshape(token_count, world_size * local_hidden_size) +) +``` + +Log `[Foundry] Multimem all-gather buffer address=0x%x` once per gatherer. +Neither active callback may invoke an upstream lazy gather builder or perform a +second rendezvous. + +- [ ] **Step 4: Write profile and recipe tests** + +Change the supported fixtures in +`tests/test_sglang_profile_validation.py` and +`tests/test_sglang_hook_targets.py` to require: + +```python +enable_torch_symm_mem = True +enable_symm_mem = False +enforce_disable_flashinfer_allreduce_fusion = True +cuda_graph_config.decode.max_bs = 256 +cuda_graph_config.decode.bs = [256] +``` + +Add rejection cases for missing fields, torch symmetric memory disabled, NCCL +symmetric memory enabled, fusion disable false, and empty/multiple/non-256 +decode shapes. In `tests/test_sglang_tp_recipe.py`, parse the JSON and assert +exact equality with: + +```python +{ + "decode": {"backend": "full", "max_bs": 256, "bs": [256]}, + "prefill": {"backend": "disabled"}, +} +``` + +Require `--enable-torch-symm-mem`, +`--enforce-disable-flashinfer-allreduce-fusion`, and +`--disable-custom-all-reduce` in baseline, SAVE, and LOAD. + +- [ ] **Step 5: Run profile/recipe tests and verify RED** + +Run: + +```bash +python3 -m pytest -q \ + tests/test_sglang_profile_validation.py \ + tests/test_sglang_tp_recipe.py +``` + +Expected: the current plain-NCCL/multi-shape expectations fail. + +- [ ] **Step 6: Flip the profile and recipe atomically** + +Update `_validate_profile` to use exact-field checks, require the values above, +and identify the envelope as `torch-symmetric-memory`, not `plain-NCCL`. +Update the serve script to pass the three collective-selection flags and: + +```text +--cuda-graph-config '{"decode":{"backend":"full","max_bs":256,"bs":[256]},"prefill":{"backend":"disabled"}}' +``` + +Keep the immutable model revision, seed 42, FlashInfer attention, existing +environment invariants, and both symmetric SAVE/LOAD TOMLs unchanged. + +- [ ] **Step 7: Write capture/pool ordering tests** + +In `tests/test_sglang_capture_hooks.py`, make `_device_module.synchronize`, +`_tp_group.barrier`, graph-context entry, and `forward_fn` share one event list. +Assert: + +```python +assert events == ["synchronize", "barrier", "enter", "forward", "exit"] +assert forwards == ["forward"] +assert post_warmup_calls == [] +``` + +In `tests/test_sglang_graph_load.py`, give the fake backend +`_pool = "pool-handle"`, accept `pool=` in fake `start_graph_builds`, and assert +the sole start receives `"pool-handle"`. Preserve all metadata-then-finish and +`finish_graph_loads` rejection assertions. + +- [ ] **Step 8: Run capture/pool tests and verify RED** + +Run: + +```bash +python3 -m pytest -q \ + tests/test_sglang_capture_hooks.py \ + tests/test_sglang_graph_load.py +``` + +Expected: missing synchronize/barrier events and missing `pool=` argument. + +- [ ] **Step 9: Implement capture lockstep and explicit pool** + +In `backend.py`, immediately before graph-context construction, call: + +```python +self._device_module.synchronize() +self._tp_group.barrier() +``` + +Keep exactly one captured forward and no warmup/post-warmup calls. In +`graph_ops.load_all_graphs`, reject a missing `backend._pool` and pass it to the +existing one-shot `start_graph_builds`; do not change per-index interleaving. + +- [ ] **Step 10: Run Task 1 gates** + +Run: + +```bash +python3 -m pytest -q \ + tests/test_sglang_hook_targets.py \ + tests/test_sglang_target_resolution.py \ + tests/test_sglang_profile_validation.py \ + tests/test_sglang_capture_hooks.py \ + tests/test_sglang_graph_load.py \ + tests/test_sglang_tp_recipe.py +bash -n recipe/experimental/serve_qwen3-8b_sglang_tp.sh +python3 -m compileall -q \ + python/foundry/integration/sglang \ + python/foundry_sglang_plugin \ + tests/modal_sglang_target_resolution.py +``` + +Expected: all tests pass, shell parses, and Python compilation exits zero. + +- [ ] **Step 11: Commit Task 1** + +```bash +git add \ + python/foundry_sglang_plugin/targets.py \ + python/foundry/integration/sglang/hooks.py \ + python/foundry/integration/sglang/backend.py \ + python/foundry/integration/sglang/graph_ops.py \ + recipe/experimental/serve_qwen3-8b_sglang_tp.sh \ + tests/test_sglang_hook_targets.py \ + tests/test_sglang_target_resolution.py \ + tests/modal_sglang_target_resolution.py \ + tests/test_sglang_profile_validation.py \ + tests/test_sglang_capture_hooks.py \ + tests/test_sglang_graph_load.py \ + tests/test_sglang_tp_recipe.py +git commit -m "fix(sglang): restore graphs with stable symmetric TP collectives" +``` + +- [ ] **Step 12: Push and validate pinned targets on one H100** + +Push the immutable Task 1 commit, then run: + +```bash +git push origin cursor/tp-upstream-consolidation-fe33 +FOUNDRY_REF=<40-hex-task-1-commit> \ +SGLANG_BASE=a25164bda355250015c0ba6d8821f0d4b37abe7d \ +SGL_MODEL_REVISION=b968826d9c46dd6066d109eabc6255188de91218 \ +modal run --env rahul-dev tests/modal_sglang_target_resolution.py +``` + +Expected: nine targets resolve and configured activation reports both new hooks. + +--- + +### Task 2: Focused SAVE-to-LOAD Treatment and Archive Evidence + +**Files:** + +- Modify: `tests/tp_validation_helpers.py` +- Modify: `tests/sglang_tp_validation.py` +- Modify: `tests/modal_sglang_tp.py` +- Test: `tests/test_sglang_tp_validation.py` +- Test: `tests/test_modal_sglang_tp.py` +- Evidence: ignored + `.superpowers/sdd/sglang-symmetric-tp-treatment-runtime-report.md` + +**Interfaces:** + +- Produce + `archive_collective_kernel_counts(workspace: Path, tp_size: int) -> + dict[str, dict[str, int]]`. +- Extend + `scan_log_text(text: str, pre_teardown_text: str | None = None)` with + `torch_symm_mem_ranks: list[str]` and + `multimem_all_gather_ranks: list[str]`. +- Extend `run_phase` with `focused_treatment: bool = False` so focused mode + issues one fixed prompt with temperature 0 and 32 output tokens while + preserving its existing phase and archive-evidence parameters. +- Add a `focused: bool = False` local-entrypoint argument; `--focused` runs only + SAVE then fresh LOAD. + +- [ ] **Step 1: Write scanner and archive-node tests** + +Add JSON fixtures with `KernelNode` records containing +`ncclDevKernel_AllReduce_Sum_bf16_RING_LL`, +`ncclDevKernel_AllGather_RING_LL`, representative +`CUDASymmetricMemoryOps*all_reduce*` / +`CUDASymmetricMemoryOps*all_gather*` function names, and false-positive +noncollective nodes. Assert: + +```python +assert archive_collective_kernel_counts(workspace, 2) == { + "0": { + "all_reduce": 1, + "all_gather": 1, + "symmetric_all_reduce": 1, + "symmetric_all_gather": 1, + }, + "1": { + "all_reduce": 1, + "all_gather": 1, + "symmetric_all_reduce": 1, + "symmetric_all_gather": 1, + }, +} +``` + +Count each kernel node once if any string field under that node begins with +`ncclDevKernel_AllReduce` / `ncclDevKernel_AllGather`, or if one value string +contains both `CUDASymmetricMemoryOps` and the corresponding `all_reduce` / +`all_gather` operation. Do not count arbitrary log text, JSON keys, +non-KernelNodes, or namespace/operation fragments split across values. Add +exact log fixtures for both Foundry markers and retain marker rank/address +diagnostics. + +- [ ] **Step 2: Run evidence tests and verify RED** + +Run: + +```bash +python3 -m pytest -q tests/test_sglang_tp_validation.py +``` + +Expected: missing archive scanner and symmetric marker fields. + +- [ ] **Step 3: Implement pure-Python evidence extraction** + +Implement `archive_collective_kernel_counts` in +`tests/tp_validation_helpers.py` using `json.loads`, rank-local +`graph_*.json`, `node["type"] == "KernelNode"`, and a recursive string walk of +each node. In `tests/sglang_tp_validation.py`, add rank-aware regexes for: + +```text +[Foundry] Torch symmetric-memory buffer address= +[Foundry] Multimem all-gather buffer address= +``` + +Return sorted unique rank lists from `scan_log_text`. Keep NCCL channel/version +parsing as diagnostics and do not weaken either fatal-scan window. + +- [ ] **Step 4: Write focused-mode source contracts** + +In `tests/test_modal_sglang_tp.py`, assert that focused mode: + +- uses exactly `("save", "load")`; +- passes SAVE fingerprints/counts/byte hashes into LOAD; +- queries both phases with one prompt and `max_new_tokens=32`; +- requires graph counts `{"0": 1, "1": 1}`; +- requires the torch symmetric communicator marker on both ranks and at least + one direct multimem gather marker globally in SAVE and LOAD; +- requires positive symmetric all-reduce and all-gather KernelNodes on both + ranks in both phases; +- requires zero archived NCCL all-reduce/all-gather nodes; +- requires equal nonempty output bytes and rank offsets; +- rejects LOAD capture/save, archive mutation, replay absence, bad shutdown, or + scanner errors; and +- never executes baseline, SAVE2, LOAD2, or concurrent queries. + +- [ ] **Step 5: Run focused harness contracts and verify RED** + +Run: + +```bash +python3 -m pytest -q tests/test_modal_sglang_tp.py +``` + +Expected: no focused entrypoint branch or treatment checks exist. + +- [ ] **Step 6: Implement the focused branch** + +Add `focused_treatment=False` to `run_phase`. In focused mode, query one fixed +prompt after readiness, pass `max_new_tokens=32`, omit concurrent queries, and +record archive collective counts. Add `_run_focused_treatment(run_id)` called +from `main(focused: bool = False)` when `focused` is true. + +SAVE must create the empty run-scoped archive and commit it. LOAD must consume +SAVE's inventory/fingerprints without committing. Emit a JSON report with an +explicit `status` of `VERIFIED`, `NOT VERIFIED`, or `INCONCLUSIVE`; any failed +acceptance check raises after printing the report. Environment failure is +`INCONCLUSIVE` only before restored replay and only when unrelated to the +symmetric policy. + +- [ ] **Step 7: Run local Task 2 gates** + +Run: + +```bash +python3 -m pytest -q \ + tests/test_sglang_tp_validation.py \ + tests/test_modal_sglang_tp.py +python3 -m compileall -q \ + tests/tp_validation_helpers.py \ + tests/sglang_tp_validation.py \ + tests/modal_sglang_tp.py +``` + +Expected: all tests pass and compilation exits zero. + +- [ ] **Step 8: Commit and push the focused harness** + +```bash +git add \ + tests/tp_validation_helpers.py \ + tests/sglang_tp_validation.py \ + tests/modal_sglang_tp.py \ + tests/test_sglang_tp_validation.py \ + tests/test_modal_sglang_tp.py +git commit -m "test(sglang): add focused symmetric TP restore treatment" +git push origin cursor/tp-upstream-consolidation-fe33 +``` + +- [ ] **Step 9: Run the focused two-H100 treatment** + +```bash +FOUNDRY_REF=<40-hex-task-2-commit> \ +SGLANG_BASE=a25164bda355250015c0ba6d8821f0d4b37abe7d \ +SGL_MODEL_REVISION=b968826d9c46dd6066d109eabc6255188de91218 \ +modal run --env rahul-dev tests/modal_sglang_tp.py --focused +``` + +Expected: `status=VERIFIED`, exact SAVE/LOAD output, one graph per rank, torch +symmetric communicator markers on both ranks, a global direct multimem marker +in each phase, positive per-rank symmetric all-reduce/all-gather KernelNodes, +zero archived NCCL collective nodes, equal offsets, immutable archive, restored +replay, no LOAD capture/save, and no fatal. + +- [ ] **Step 10: Record and review focused evidence** + +Write the ignored runtime report with the immutable Foundry/SGLang/model SHAs, +Modal app/run URL, GPU/runtime versions, response hash/token IDs, per-rank graph +counts and offsets, symmetric marker ranks, NCCL collective-node counts, +archive hashes, replay evidence, shutdown codes, and final classification. +Have a non-fast reviewer compare raw artifacts to every treatment predicate. +Proceed only on an independent `VERIFIED` review. + +--- + +### Task 3: Full Five-Phase Matrix and Documentation Evidence + +**Files:** + +- Modify: `tests/modal_sglang_tp.py` +- Test: `tests/test_modal_sglang_tp.py` +- Modify: `docs/sglang/overview.md` +- Modify: `docs/sglang/hooks.md` +- Modify: `docs/sglang/save-load-workflow.md` +- Modify: `docs/sglang/memory-lifecycle.md` +- Modify: `docs/sglang/memory-consistency.md` +- Modify: `docs/sglang/direct-edits.md` +- Modify: `recipe/experimental/README.md` +- Modify: `recipe/sglang/README.md` +- Modify: `README.md` +- Modify: `ROADMAP.md` +- Modify: `RELEASE.md` +- Evidence: ignored + `.superpowers/sdd/sglang-symmetric-tp-full-matrix-runtime-report.md` + +**Interfaces:** + +- Preserve `PHASES = + ("baseline", "save", "save2", "load", "load2")`. +- Set `COLLECTIVE_BACKEND = "torch_symmetric_memory"`. +- Replace `nccl_p2p_ipc_every_phase` with symmetric-specific acceptance while + retaining NCCL transport/version fields as diagnostics. + +- [ ] **Step 1: Write full-matrix symmetric acceptance contracts** + +Update `tests/test_modal_sglang_tp.py` to require: + +```python +assert _module_literal("COLLECTIVE_BACKEND") == "torch_symmetric_memory" +assert _module_literal("PHASES") == ( + "baseline", "save", "save2", "load", "load2" +) +``` + +Require one graph per rank in both SAVEs, one restored graph per rank in both +LOADs, torch symmetric communicator markers on ranks 0/1, a global direct +multimem marker, positive per-rank symmetric collective KernelNodes in +SAVE/SAVE2/LOAD/LOAD2, zero archived NCCL collective nodes, and no plain-NCCL +channel acceptance key. +Preserve exact sequential/concurrent baseline comparisons, SAVE/SAVE2 semantic +and byte reproducibility, LOAD/LOAD2 byte immutability, offsets, batch-1 and +batch>1 replay, no recapture, teardown, provenance, and fatal scanning. + +- [ ] **Step 2: Run the full-harness contract and verify RED** + +Run: + +```bash +python3 -m pytest -q tests/test_modal_sglang_tp.py +``` + +Expected: the old `nccl_p2p_ipc_every_phase` and dynamic positive graph-count +acceptance violate the new contract. + +- [ ] **Step 3: Update the full matrix checks** + +Set the collective backend constant, thread archive collective counts through +all archive phases, and require: + +```python +archive_phases = ("save", "save2", "load", "load2") +foundry_phases = ("save", "save2", "load", "load2") +expected_ranks = {"0", "1"} +expected_graph_counts = {"0": 1, "1": 1} +all( + results[phase]["archive_graph_counts"] == expected_graph_counts + for phase in archive_phases +) +all( + set(results[phase]["archive_collective_kernel_counts"]) == expected_ranks + and all( + counts["all_reduce"] == counts["all_gather"] == 0 + and counts["symmetric_all_reduce"] > 0 + and counts["symmetric_all_gather"] > 0 + for counts in results[phase]["archive_collective_kernel_counts"].values() + ) + for phase in archive_phases +) +all( + results[phase]["torch_symm_mem_ranks"] == ["0", "1"] + and any(results[phase]["multimem_all_gather_addresses"].values()) + for phase in foundry_phases +) +``` + +Do not require Foundry markers in baseline because Foundry is inactive there. +Keep baseline output as the correctness oracle. Leave NCCL diagnostics in each +phase result and runtime provenance, but remove them from pass/fail collective +selection. + +- [ ] **Step 4: Run local full-matrix gates** + +Run: + +```bash +python3 -m pytest -q \ + tests/test_sglang_hook_targets.py \ + tests/test_sglang_target_resolution.py \ + tests/test_sglang_profile_validation.py \ + tests/test_sglang_capture_hooks.py \ + tests/test_sglang_graph_load.py \ + tests/test_sglang_tp_recipe.py \ + tests/test_sglang_tp_validation.py \ + tests/test_modal_sglang_tp.py +python3 -m compileall -q python tests +pre-commit run --all-files +git diff --check origin/main...HEAD +``` + +Expected: all tests and hooks pass with no whitespace errors. + +- [ ] **Step 5: Commit and push the full acceptance harness** + +```bash +git add tests/modal_sglang_tp.py tests/test_modal_sglang_tp.py +git commit -m "test(sglang): require symmetric evidence in full TP matrix" +git push origin cursor/tp-upstream-consolidation-fe33 +``` + +- [ ] **Step 6: Run the full five-phase matrix** + +```bash +FOUNDRY_REF=<40-hex-task-3-harness-commit> \ +SGLANG_BASE=a25164bda355250015c0ba6d8821f0d4b37abe7d \ +SGL_MODEL_REVISION=b968826d9c46dd6066d109eabc6255188de91218 \ +modal run --env rahul-dev tests/modal_sglang_tp.py +``` + +Expected: all five phases pass the symmetric checks plus every retained +archive, output, replay, offset, process, provenance, and scanner check. + +- [ ] **Step 7: Record evidence and update documentation** + +Write the ignored full-matrix report first. Then update every documentation path +listed above to: + +- reject plain-NCCL fresh-process SGLang restore and cite the disproving app; +- describe the retained communicator/direct multimem/one-shape policy; +- retain the exact pins and zero-edit `HookRegistry` architecture; +- distinguish the focused treatment from the full matrix; +- record both immutable run URLs, versions, graph counts, offsets, archive-node + counts, and output/replay results; and +- keep TP>2, multi-host, custom/NCCL-symmetric/fused collectives, CUMEM, NVLS, + EP, DP-attention, speculative decode, LoRA, pdmux, and unpinned revisions + unsupported. + +Do not preserve any statement that plain NCCL is supported, that 36 restored +graphs prove correctness, or that pinned-main acceptance is still pending. + +- [ ] **Step 8: Review and commit final evidence docs** + +Run: + +```bash +python3 -m pytest -q tests/test_modal_sglang_tp.py +pre-commit run --all-files +git diff --check origin/main...HEAD +``` + +Have a non-fast reviewer compare the docs to both raw run reports and scan for +plain-NCCL contradictions. Then: + +```bash +git add \ + README.md ROADMAP.md RELEASE.md \ + docs/sglang \ + recipe/experimental/README.md \ + recipe/sglang/README.md +git commit -m "docs(sglang): record symmetric TP treatment and full validation" +git push origin cursor/tp-upstream-consolidation-fe33 +``` + +## Plan self-review + +1. **Spec coverage:** Task 1 owns stable hooks, target inventory, exact profile, + one shape, capture lockstep, and explicit pool. Task 2 owns focused + SAVE→LOAD and zero-NCCL-node evidence. Task 3 is gated on Task 2 and owns the + full matrix plus evidence-backed docs. +2. **Preservation:** No task replaces atomic state, `ShapeKey`, interleaved + finishes, template-aware loading, archive immutability, process cleanup, or + fatal scanning. +3. **Excluded history:** No task contains device-group rendezvous, warmup + forwards, synthesized state, `create_state` interception, or independent + loading. +4. **API consistency:** The same two target constants, callback signatures, + marker strings, one-shape config, pool argument, scanner fields, and archive + count API are used throughout. +5. **Gate consistency:** Task 3 is forbidden unless Task 2 is independently + reviewed as `VERIFIED`; documentation success claims occur only after the + full matrix passes. diff --git a/docs/superpowers/plans/2026-07-23-tp-upstream-consolidation.md b/docs/superpowers/plans/2026-07-23-tp-upstream-consolidation.md new file mode 100644 index 00000000..71625389 --- /dev/null +++ b/docs/superpowers/plans/2026-07-23-tp-upstream-consolidation.md @@ -0,0 +1,643 @@ +# Tensor-Parallel Upstream Consolidation Implementation Plan + +> [!WARNING] +> **SUPERSEDED FOR SGLANG; vLLM portions remain applicable.** The shared +> plain-NCCL policy below is historical for SGLang. Use the +> [symmetric-TP treatment design](../specs/2026-07-23-sglang-symmetric-tp-treatment-design.md) +> and +> [implementation plan](2026-07-23-sglang-symmetric-tp-treatment.md) for current +> SGLang policy and evidence. The vLLM plan remains unchanged. + +--- + +> **For agentic workers:** REQUIRED SUB-SKILL: Use superpowers:subagent-driven-development (recommended) or superpowers:executing-plans to implement this plan task-by-task. Steps use checkbox (`- [ ]`) syntax for tracking. + +**Goal:** Produce one `origin/main`-rooted branch that contains every unique active-branch change and verifiably restores TP=2 CUDA graphs for pinned vLLM and SGLang configurations. + +**Architecture:** Preserve the hardened shared VMM/IPC and graph-restoration history from PR #4, layer the vLLM TP work from PR #5, then add the post-SGLang-#26735 API port from PR #2. Both engines use the same constrained NCCL P2P/IPC profile in baseline, SAVE, and LOAD; CPU contracts enforce the profile and Modal harnesses prove fresh-process restoration on 2x H100. + +**Tech Stack:** C++17/CUDA Driver API, pybind11, Python 3.12, PyTorch 2.11/cu130, pytest, Bash, vLLM, SGLang, NCCL 2.28+, Modal. + +## Global Constraints + +- The final merge base is `origin/main` commit `21824649e711487a2c6af01cb631b3e41ccc32d8` unless `origin/main` advances before integration, in which case rebase the design/plan commits onto that new tip first. +- The review diff is always `origin/main...HEAD`; do not merge another feature branch. +- Initial support is single-host dense Qwen3-8B TP=2 on Hopper using plain NCCL P2P/legacy CUDA IPC. +- Disable custom/fused/symmetric-memory all-reduce, NCCL CUMEM/NVLS, piecewise graphs, SGLang FlashInfer autotuning, vLLM V2, weight offload, sleep mode, and expandable allocator segments. +- Keep engine revisions, model revision, seed, graph settings, rank order, and transport identical across comparable phases. +- Run CUDA and engine validation only on Modal H100s in `rahul-dev`; do not initialize CUDA in the orchestration parent. +- Put imports at module scope unless an optional-engine import must remain lazy; document that exception. +- Do not claim support broader than the exact final-head evidence. + +--- + +### Task 1: Freeze and Integrate the Shared Consolidation + +**Files:** + +- Modify: all files changed by `origin/main..2985288cef719bfa48df18562dd7218abd340d9e` +- Preserve: `docs/superpowers/specs/2026-07-23-tp-upstream-consolidation-design.md` +- Preserve: `docs/superpowers/plans/2026-07-23-tp-upstream-consolidation.md` + +**Interfaces:** + +- Consumes: PR #4 head `2985288cef719bfa48df18562dd7218abd340d9e`. +- Produces: shared VMM2 IPC, graph/RNG restoration, SGLang TP recipe/harness, graph-pool fix, archive comparison, and complete replay evidence on this main-rooted branch. + +- [ ] **Step 1: Re-fetch and verify the frozen source** + +Run: + +```bash +git fetch origin main cursor/tp-consolidation-2e2c +test "$(git merge-base origin/main origin/cursor/tp-consolidation-2e2c)" = \ + "$(git rev-parse origin/main)" +test "$(git rev-parse origin/cursor/tp-consolidation-2e2c)" = \ + 2985288cef719bfa48df18562dd7218abd340d9e +``` + +Expected: all commands exit zero. If the source head moved, inventory its new +commits and update this plan before integration. + +- [ ] **Step 2: Integrate the linear shared range** + +Run: + +```bash +git cherry-pick \ + 4cc4a56484e16aece6311261c9998b7bcc5181e1^..2985288cef719bfa48df18562dd7218abd340d9e +``` + +Expected: the range applies without touching either superpowers document. + +- [ ] **Step 3: Prove source-tree equivalence** + +Run: + +```bash +git diff --exit-code \ + origin/cursor/tp-consolidation-2e2c \ + HEAD \ + -- . ':(exclude)docs/superpowers/**' +git diff --check origin/main...HEAD +``` + +Expected: no functional diff against PR #4 and no whitespace errors. + +- [ ] **Step 4: Verify shared CPU contracts** + +Run: + +```bash +python3 -m pytest tests/test_sglang_tp_recipe.py -q +python3 -m compileall -q python tests/modal_sglang_tp.py +``` + +Expected: six recipe tests pass and compilation exits zero. + +### Task 2: Integrate the vLLM TP Work + +**Files:** + +- Create: `recipe/vllm/serve_qwen3-8b_tp.sh` +- Create: `tests/modal_vllm_tp.py` +- Create: `tests/test_vllm_tp_recipe.py` +- Modify: `csrc/hook.cpp` +- Modify: `tests/test_vmm_alloc.py` + +**Interfaces:** + +- Consumes: shared VMM/IPC implementation from Task 1. +- Produces: vLLM TP recipe and immutable-revision Modal harness; zero-alignment VMM reservations preserve the deterministic cursor. + +- [ ] **Step 1: Integrate only commits unique to PR #5** + +Run: + +```bash +git cherry-pick \ + 08c94a81ddb49fed0d909665e43ec0660299e772 \ + 7b8abf84e5639a160d52c0731bc28bfdb0982cd2 \ + 2baf3181bb87119f8a91626f96c1a3773c66bdea \ + cbeb4c7ffcf64f9ce2b153fe3057847a35ea7ce0 \ + 8ff7050446cc116cb8f03119dbade737c3d8208f \ + 944222dd1e561323850ef366fb54b5ff980ef3e1 \ + ef56a76c080391b3221701ef7aacaaa3943b3543 \ + ea48b5070526bc40ec1df138d21798b5a27b0a30 \ + 160a8c0628d09cc3407ae56a01f328d305d63954 \ + dd76136354a315251edc2d6a7ecbaeb11d1d26a7 \ + d9d757d675e268836e11c4e316d027c04282fb1a +``` + +If `csrc/hook.cpp` conflicts, preserve Task 1's versioned preallocation +generation and lifetime handling, then apply only the zero-alignment behavior: + +```cpp +const size_t effective_alignment = + alignment == 0 ? kDefaultAllocationAlignment : alignment; +``` + +The reservation path must align with `effective_alignment` without advancing +the Foundry cursor twice. + +- [ ] **Step 2: Run the vLLM contracts and zero-alignment test selection** + +Run: + +```bash +python3 -m pytest tests/test_vllm_tp_recipe.py -q +python3 -m pytest tests/test_vmm_alloc.py -q +python3 -m compileall -q tests/modal_vllm_tp.py +bash -n recipe/vllm/serve_qwen3-8b_tp.sh +``` + +Expected: recipe contracts pass; the native test either passes in the CUDA +environment or is deferred with an explicit local no-CUDA reason. + +- [ ] **Step 3: Compare the vLLM-specific tree** + +Run: + +```bash +git diff --check origin/main...HEAD +git diff --name-status \ + origin/cursor/tp-consolidation-2e2c...HEAD +``` + +Expected: the additional files are the vLLM recipe/harness/contracts and the +zero-alignment core regression only. + +### Task 3: Port the SGLang Attention Metadata API + +**Files:** + +- Modify: `python/foundry/integration/sglang/graph_ops.py` +- Modify: `python/foundry/integration/sglang/hooks.py` +- Modify: `recipe/experimental/README.md` +- Test: `tests/modal_sglang_tp.py` + +**Interfaces:** + +- Consumes: pre-#26735 + `init_forward_metadata_capture_cuda_graph(...)` and post-#26735 + `init_forward_metadata_out_graph(forward_batch, in_capture=...)`. +- Produces: one metadata pre-pass and reuse adapter that preserves identical + SAVE/LOAD allocation ordering for either API. + +- [ ] **Step 1: Confirm the recorded pre-port failure** + +Use the preserved PR #2 run and review evidence showing that the pinned +post-#26735 fork fails before capture because the legacy +`init_forward_metadata_capture_cuda_graph` attribute is absent. Do not spend a +new GPU run reproducing an already preserved failure. + +- [ ] **Step 2: Apply the reviewed API-port commit** + +Run: + +```bash +git cherry-pick 427978596e9d1a84245ab5f033dbe0fbe52fbe5f +``` + +Resolve recipe documentation against the newer PR #4 text rather than +replacing final H100 evidence. + +- [ ] **Step 3: Keep standard-library imports at module scope** + +In `python/foundry/integration/sglang/graph_ops.py`, add: + +```python +from types import SimpleNamespace +``` + +at the top and remove the function-local import. Keep SGLang class imports +lazy because SGLang is an optional integration and importing those modules at +Foundry module import time would make CPU/core-only installations fail. + +- [ ] **Step 4: Verify the two review concerns against source** + +Confirm the pinned SGLang revision: + +- exposes `buffers.seq_lens_cpu`; +- ignores the return from `init_forward_metadata_out_graph`; +- allocates `_int_workspace_buffer` only when `in_capture=True`; and +- replans existing wrappers without allocation for `in_capture=False`. + +Run: + +```bash +python3 -m compileall -q python/foundry/integration/sglang +git diff --check origin/main...HEAD +``` + +Expected: compilation and whitespace checks pass. + +- [ ] **Step 5: Commit the import cleanup if it changed the cherry-picked commit** + +Run: + +```bash +git add python/foundry/integration/sglang/graph_ops.py +git commit -m "style: keep SGLang compatibility imports scoped" +``` + +Expected: one focused follow-up commit. + +### Task 4: Enforce the Validated TP=2 Recipe Envelope + +**Files:** + +- Modify: `tests/test_sglang_tp_recipe.py` +- Modify: `tests/test_vllm_tp_recipe.py` +- Modify: `recipe/experimental/serve_qwen3-8b_sglang_tp.sh` +- Modify: `recipe/vllm/serve_qwen3-8b_tp.sh` + +**Interfaces:** + +- Consumes: ` [--save|--load]`. +- Produces: a launch only for TP size exactly 2 with all validated transport + controls set; every unsupported invocation exits nonzero before engine + startup. + +- [ ] **Step 1: Write failing recipe-envelope tests** + +Add to both test files: + +```python +@pytest.mark.parametrize( + "arguments", + [[], ["1"], ["3"], ["2", "--invalid"]], +) +def test_tp_recipe_rejects_unvalidated_profile(arguments: list[str]) -> None: + result = subprocess.run( + ["bash", str(SERVE_SCRIPT), *arguments], + capture_output=True, + text=True, + ) + assert result.returncode != 0 + assert "TP=2" in result.stderr +``` + +Extend each phase-symmetry assertion with: + +```python +assert env["PYTORCH_CUDA_ALLOC_CONF"] == "expandable_segments:False" +``` + +Extend the vLLM assertion with: + +```python +assert env["VLLM_DISABLE_PYNCCL"] == "0" +``` + +- [ ] **Step 2: Run the tests and verify RED** + +Run: + +```bash +python3 -m pytest \ + tests/test_sglang_tp_recipe.py \ + tests/test_vllm_tp_recipe.py -q +``` + +Expected: failures because TP=3 currently starts and the additional +environment controls are absent. + +- [ ] **Step 3: Constrain both recipes** + +At the top of both scripts use: + +```bash +set -euo pipefail + +TP_SIZE="${1:-}" +MODE="${2:-}" +if [[ "$TP_SIZE" != "2" ]] || [[ "$MODE" != "" && "$MODE" != "--save" && "$MODE" != "--load" ]]; then + echo "Usage: $0 2 [--save|--load] (validated profile: TP=2)" >&2 + exit 1 +fi +``` + +Set in both scripts: + +```bash +export PYTORCH_CUDA_ALLOC_CONF=expandable_segments:False +``` + +Set in the vLLM script: + +```bash +export VLLM_DISABLE_PYNCCL=0 +``` + +Use `MODE` instead of direct `$2` references. + +- [ ] **Step 4: Run the tests and verify GREEN** + +Run: + +```bash +python3 -m pytest \ + tests/test_sglang_tp_recipe.py \ + tests/test_vllm_tp_recipe.py -q +bash -n recipe/experimental/serve_qwen3-8b_sglang_tp.sh +bash -n recipe/vllm/serve_qwen3-8b_tp.sh +``` + +Expected: all contract tests and syntax checks pass. + +- [ ] **Step 5: Commit the profile gate** + +Run: + +```bash +git add \ + tests/test_sglang_tp_recipe.py \ + tests/test_vllm_tp_recipe.py \ + recipe/experimental/serve_qwen3-8b_sglang_tp.sh \ + recipe/vllm/serve_qwen3-8b_tp.sh +git commit -m "fix: enforce validated TP communication profile" +``` + +### Task 5: Reconcile Scope and Historical Evidence + +**Files:** + +- Modify: `README.md` +- Modify: `ROADMAP.md` +- Modify: `RELEASE.md` +- Modify: `docs/sglang/overview.md` +- Modify: `docs/sglang/hooks.md` +- Modify: `recipe/experimental/README.md` +- Modify: `recipe/sglang/README.md` +- Modify: `recipe/vllm/README.md` + +**Interfaces:** + +- Consumes: final supported profile and historical run records. +- Produces: documentation that separates final-head evidence, historical H200 + evidence, and unsupported configurations. + +- [ ] **Step 1: Reconcile PR #2's useful documentation** + +Manually carry forward: + +- post-#26735 H200 run evidence; +- 16/16 SAVE/LOAD prompt equality; +- offset `121831948288` from immutable + [commit `ee7d61a`](https://github.com/modal-projects/foundry/commit/ee7d61a32a4d4732e5f14978975b96ab03bcbede); +- explicit Qwen3.5-122B hybrid-MoE non-support; and +- the fact that the API port is complete for the pinned Foundry fork. + +Do not restore the stale statement that only a pre-#26735 revision works. + +- [ ] **Step 2: Keep TP status experimentally scoped** + +Use wording equivalent to: + +```markdown +TP=2 is validated only for the pinned single-host NCCL P2P/IPC recipes. +Custom/symmetric/fused all-reduce, CUMEM, NVLS, TP>2, multi-host, and +unpinned engine revisions remain unsupported. +``` + +Do not mark broad project-wide TP complete in the roadmap. + +- [ ] **Step 3: Correct stale implementation descriptions** + +Ensure release notes describe: + +- zero-RNG graphs skipping unused generator-state restoration; +- versioned preallocation chunks; +- default pool `(0, 0)`; and +- zero-alignment preserving the VMM cursor. + +Remove text that says generator registration is handled by suspending the +allocation region if that no longer matches the final code. + +- [ ] **Step 4: Commit documentation** + +Run: + +```bash +git add README.md ROADMAP.md RELEASE.md docs recipe +git commit -m "docs: define validated tensor parallel scope" +``` + +### Task 6: Local and Native Verification + +**Files:** + +- Test: `tests/test_sglang_tp_recipe.py` +- Test: `tests/test_vllm_tp_recipe.py` +- Test: `tests/test_rng_graph.py` +- Test: `tests/test_vmm_alloc.py` +- Test: core CUDA suite under `tests/` + +**Interfaces:** + +- Consumes: consolidated code. +- Produces: CPU, syntax, lint, and H100-native evidence before engine runs. + +- [ ] **Step 1: Check environment setup completion** + +Run: + +```bash +if [[ -e /tmp/cursor/async-install/install-user.status ]]; then + test "$(< /tmp/cursor/async-install/install-user.status)" = "0" +fi +``` + +Expected: setup complete or no async setup exists. + +- [ ] **Step 2: Run local gates** + +Run: + +```bash +python3 -m pytest \ + tests/test_sglang_tp_recipe.py \ + tests/test_vllm_tp_recipe.py -q +python3 -m compileall -q python tests +pre-commit run --all-files +git diff --check origin/main...HEAD +``` + +Expected: all commands pass. + +- [ ] **Step 3: Commit and push before GPU testing** + +Run: + +```bash +git status --short +git push -u origin cursor/tp-upstream-consolidation-fe33 +``` + +Expected: clean worktree and remote updated. Update draft PR #6 with the +pre-test commit and exact local results. + +- [ ] **Step 4: Run native CUDA and IPC gates on H100** + +Use a Modal command or checked-in harness that builds with: + +```bash +CC=gcc CXX=g++ pip install -e . --no-build-isolation +pytest tests/test_vmm_alloc.py tests/test_rng_graph.py -q +pytest tests/ -q \ + --ignore=tests/test_imports.py \ + --ignore=tests/test_deepep_fabric.py +``` + +Expected: zero-alignment, RNG, default-pool, and core fresh-process tests pass. +Run the DeepEP test separately on two GPUs with its optional dependencies. + +### Task 7: Final-Head SGLang TP Validation + +**Files:** + +- Test: `tests/modal_sglang_tp.py` +- Evidence: PR #6 description and final documentation. + +**Interfaces:** + +- Consumes: final pushed commit SHA. +- Produces: pinned post-#26735 baseline/SAVE/SAVE/LOAD/restart evidence on 2x H100. + +- [ ] **Step 1: Run the final-head harness** + +Run: + +```bash +FOUNDRY_REF="$(git rev-parse HEAD)" \ + modal run --env rahul-dev tests/modal_sglang_tp.py +``` + +Expected: + +- 36 graphs per rank; +- reproducible semantic graph structure and per-rank offsets; +- all archived graphs restored and zero recaptured; +- replay at batch sizes 1 and 8; +- all NCCL channels use P2P/IPC; +- sequential SAVE2/LOAD outputs match exactly; +- concurrent outputs are nonempty; and +- no fatal log signature. + +- [ ] **Step 2: Run one additional fresh LOAD** + +Use the harness's restart mode or rerun LOAD against the same immutable archive. + +Expected: identical graph inventory and deterministic outputs without stale +IPC broker state. + +### Task 8: Final-Head vLLM TP Validation + +**Files:** + +- Test: `tests/modal_vllm_tp.py` +- Evidence: PR #6 description and final documentation. + +**Interfaces:** + +- Consumes: final pushed commit SHA. +- Produces: pinned V1/PyNCCL baseline/seed-SAVE/SAVE/SAVE/LOAD/restart evidence on 2x H100. + +- [ ] **Step 1: Run the final-head harness** + +Run: + +```bash +FOUNDRY_REF="$(git rev-parse HEAD)" \ + modal run --env rahul-dev tests/modal_vllm_tp.py +``` + +Expected: + +- 64 graphs per rank; +- identical per-rank offsets across the deterministic SAVE passes and LOAD + (the PR #5 reference run observed `70724354048`); +- all graphs restored and zero captured on LOAD; +- NCCL P2P/IPC on both ranks; +- bounded deterministic factual completions match for sequential and + concurrent requests; and +- no CUDA, hook, NCCL, worker, or scheduler error. + +- [ ] **Step 2: Run one additional fresh LOAD** + +Expected: same archive inventory, offsets, and outputs. + +### Task 9: Final Review, Evidence Commit, and PR Update + +**Files:** + +- Modify: documentation files only if measured final-head values differ. +- Modify: PR #6 title/body through the PR management tool. + +**Interfaces:** + +- Consumes: all local, native, SGLang, and vLLM results. +- Produces: clean final diff and an evidence-backed draft PR. + +- [ ] **Step 1: Review branch coverage** + +Run: + +```bash +git fetch origin +git cherry -v HEAD origin/a100-issue +git cherry -v HEAD origin/ep-ipc +git cherry -v HEAD origin/devin/1784752573-sglang-nofabric-ep-ipc +git cherry -v HEAD origin/devin/1784764658-sglang-tensor-parallel +git cherry -v HEAD origin/cursor/tp-consolidation-2e2c +git cherry -v HEAD origin/cursor/vllm-tp-consolidation-5d04 +``` + +Expected: no unique functional patch is absent. Documentation-only differences +must be explicitly reconciled. + +- [ ] **Step 2: Run final quality gates** + +Run: + +```bash +git diff --check origin/main...HEAD +python3 -m pytest \ + tests/test_sglang_tp_recipe.py \ + tests/test_vllm_tp_recipe.py -q +pre-commit run --all-files +git status --short --branch +``` + +Expected: all pass and worktree is clean after any evidence update. + +- [ ] **Step 3: Commit measured documentation changes** + +If final-head values changed: + +```bash +git add README.md ROADMAP.md RELEASE.md docs recipe +git commit -m "docs: record final tensor parallel validation" +``` + +If values did not change, do not create an empty commit. + +- [ ] **Step 4: Push and update the draft PR** + +Run: + +```bash +git push -u origin cursor/tp-upstream-consolidation-fe33 +``` + +Update PR #6 with: + +- exact final commit; +- exact engine/model revisions; +- H100 model, CUDA, PyTorch, NCCL versions; +- run URLs; +- graph counts and offsets; +- SAVE2/LOAD and restart comparisons; +- local/native test results; and +- explicit unsupported scope. diff --git a/docs/superpowers/specs/2026-07-23-sglang-main-port-design.md b/docs/superpowers/specs/2026-07-23-sglang-main-port-design.md new file mode 100644 index 00000000..20033daf --- /dev/null +++ b/docs/superpowers/specs/2026-07-23-sglang-main-port-design.md @@ -0,0 +1,472 @@ +# SGLang Current-Main Port Design + +> [!WARNING] +> **SUPERSEDED FOR SGLANG.** The plain-NCCL policy in this historical design was +> amended after corrupt fresh-LOAD output. Use the +> [symmetric-TP treatment design](2026-07-23-sglang-symmetric-tp-treatment-design.md) +> and +> [implementation plan](../plans/2026-07-23-sglang-symmetric-tp-treatment.md) +> as the current SGLang policy and evidence boundary. + +## Goal + +Retarget Foundry's SGLang tensor-parallel SAVE-to-LOAD integration so that it +runs against **current upstream `sgl-project/sglang` main**, pinned at one +immutable commit, with **zero edits to a clean upstream SGLang checkout and no +fork commit**. The integration must: + +- activate through SGLang's first-class `sglang.srt.plugins` entry point and + `HookRegistry`, driven by an inherited environment variable rather than a + bespoke `ServerArgs` field; +- bind SAVE and LOAD to the decomposed CUDA-graph capture path + (`FullCudaGraphBackend` + `DecodeCudaGraphRunner`) introduced by SGLang + PR #23906, keyed by `ShapeKey`; +- reuse the serialized target `MemoryPoolConfig` on LOAD without re-profiling + device memory; +- keep the single-host dense Qwen3-8B TP=2 plain-NCCL P2P/IPC profile as the + only supported envelope; and +- prove fresh-process SAVE-to-LOAD correctness on 2x H100 at the exact pin. + +This is a retarget of the pinned `foundry-org/sglang` fork integration onto a +moving-but-pinned upstream `main`. It supersedes the fork-based activation model +for SGLang; it does not change the vLLM integration or the shared VMM/graph core. + +## Authorization and Provenance + +The user explicitly revised scope to require the current-main port and +authorized every design decision recorded here. The evidence base is two +read-only source-of-truth reports produced this session: + +- `.superpowers/sdd/sglang-main-api-report.md` — upstream API and + lifecycle audit against the pinned checkout. +- `.superpowers/sdd/sglang-main-port-design-report.md` — port + architecture and risk study. + +Both reports inspected the pinned upstream checkout, the current Foundry +integration under `python/foundry/integration/sglang`, and live upstream +PR/issue metadata on 2026-07-23. This spec restates their conclusions as +authorized decisions and is the design of record for the port. + +## Pin and Evidence Boundary + +- Upstream repository: `https://github.com/sgl-project/sglang`. +- Target immutable `main` commit: + `a25164bda355250015c0ba6d8821f0d4b37abe7d`. This SHA was live `main` during + the API audit (observed via `git ls-remote ... refs/heads/main`). "Support" + means this exact commit, never a moving `main`. +- Immutable model revision: Qwen3-8B `b968826d9c46dd6066d109eabc6255188de91218`. +- Any discrepancy discovered between this spec and a fresh clean checkout of the + pinned SHA must be resolved from source evidence in that checkout before + writing code. Do not guess an API shape; read the pinned source. +- All CUDA/native and engine validation runs on 2x H100 in Modal's `rahul-dev` + environment. The local VM has no NVIDIA device or CUDA toolkit. + +## What Changed Upstream + +The two findings that reshape the port: + +1. **A first-class plugin framework now exists.** + `python/sglang/srt/plugins/__init__.py` defines + `GENERAL_PLUGINS_GROUP = "sglang.srt.plugins"` and `load_plugins()`, which + discovers setuptools entry points, invokes each zero-argument callable, then + applies `HookRegistry` hooks. `HookRegistry` (in + `python/sglang/srt/plugins/hook_registry.py`) supports `BEFORE`, `AFTER`, + `AROUND`, and `REPLACE` hooks on any dotted target, preserves `@classmethod` + and `@staticmethod` descriptors, applies class replacements before method + hooks, and propagates class replacements to modules holding stale + `from X import Y` bindings via `_propagate_patch`. `SGLANG_PLUGINS` is a + defined env at `environ.py:1141`. `load_plugins()` runs at the top of + `Engine.__init__`, again in `_launch_subprocesses`, and as the first + statement of `run_scheduler_process` — every process Foundry needs to hook. + +2. **The monolithic capture path was decomposed (PR #23906).** + `CudaGraphRunner` and `model_runner_kv_cache_mixin` no longer exist. Capture + is now phase runners (`DecodeCudaGraphRunner`, `PrefillCudaGraphRunner`, + `EagerRunner`) plus pluggable graph backends (`FullCudaGraphBackend`, + `BreakableCudaGraphBackend`, `TcPiecewiseCudaGraphBackend`) under + `python/sglang/srt/model_executor/runner_backend/`. Graph keys are now + `ShapeKey` dataclasses, not bare ints. + +Two further upstream changes matter: + +- PR #26735 removed the legacy attention capture/replay metadata methods and + established the three-method contract `init_forward_metadata`, + `init_forward_metadata_out_graph(fb, in_capture=False)`, and + `init_forward_metadata_in_graph(fb)`. The removed + `init_forward_metadata_capture_cuda_graph` has no upstream shim. +- PR #28739 moved warmup to `BaseRunner.warmup` / `EagerRunner`; the removed + `ModelRunner.kernel_warmup` must not be recreated. + +## Why the Current Foundry Integration Cannot Run + +On the pinned SHA the current integration is inert or broken: + +1. Activation reads `server_args.foundry_graph_extension_config_path`, a field + absent upstream, so `install_hooks()` returns without installing anything. +2. `_patch_init_memory_pool()` imports the removed + `model_runner_kv_cache_mixin`, raising before later hooks install. +3. `_patch_cuda_graph_capture()` targets the removed monolithic + `CudaGraphRunner` and its removed methods (`_create_device_graph`, + `_capture_graph`, `capture_one_batch_size`, `capture`). +4. `_patch_init_torch_distributed()` reads `self.tp_rank`, `self.pp_rank`, and + `self.dp_rank`; current ranks live on `self.ps` (a frozen `ParallelState`). +5. `_patch_kernel_warmup()` targets the removed `ModelRunner.kernel_warmup`. +6. Graph helpers read runner-owned `graphs`/`output_buffers` and integer keys; + current ownership is `backend._graphs`/`backend._outputs` keyed by `ShapeKey`, + and `build_capture_fb_view` reads `num_tokens_per_bs`, renamed to + `captured_req_width`. + +## Supported Profile + +The only supported envelope is intentionally narrow and identical to the +consolidation profile, restated for the pinned SHA: + +- one host, two Hopper GPUs (validated on 2x H100); +- dense Qwen3-8B at the immutable model revision above; +- TP size exactly 2, with `PP=DP=1`, no EP, no DP-attention, no disaggregation; +- plain NCCL P2P/legacy CUDA IPC collectives; +- FlashInfer attention backend; +- full decode CUDA graphs only, prefill graphs disabled; +- Foundry SAVE and LOAD in fresh processes with identical rank/GPU ordering; +- immutable Foundry, engine, and model revisions. + +The profile disables, as **correctness requirements** grounded in open upstream +defects, not preferences: + +- SGLang custom all-reduce (issue #31117 concurrency deadlock); +- NCCL symmetric memory and torch symmetric-memory all-reduce (PR #29946 captured + symmetric-memory corruption); +- MSCCL++ all-reduce; +- FlashInfer autotune and FlashInfer all-reduce fusion (PR #23317 rank-divergent + tactics); +- NCCL CUMEM and NVLS; +- post-capture KV sizing (must pin `SGLANG_ENABLE_POST_CAPTURE_KV_SIZING=0`); +- piecewise/breakable/tc_piecewise decode graphs; +- pdmux/multi-stream capture; +- expandable CUDA allocator segments. + +A Foundry-enabled launch outside this envelope must fail early with an +actionable message from a hook that SGLang executes (see Failure Handling). It +must not silently select an unvalidated backend. + +This design claims no support for TP greater than 2, multiple hosts, non-Hopper +GPUs, custom/symmetric/fused all-reduce, CUMEM, NVLS, piecewise graphs, EP, +DP-attention, or any unpinned SGLang revision. + +## Activation Architecture + +**Decision: a Foundry-owned general plugin with env-var config; delete the +`foundry-org/sglang` fork dependency from the SGLang design.** + +- Ship a `sglang.srt.plugins` entry point from Foundry's own `pyproject.toml`: + +```toml +[project.entry-points."sglang.srt.plugins"] +foundry = "foundry_sglang_plugin:activate" +``` + +- `activate()` is a zero-argument callable invoked by `load_plugins()` in every + SGLang process. It reads `FOUNDRY_GRAPH_EXTENSION_CONFIG_PATH`; if unset it is + a no-op (Foundry installed but inactive); if set it loads the existing Foundry + TOML through `config.load_graph_extension_config` and registers all hooks + through `HookRegistry`. +- Config-path delivery moves from the removed `--foundry-graph-extension-config-path` + CLI arg to the environment variable, which is inherited by every `spawn` + child and is independent of `ServerArgs` pickling. Recipes may also set + `SGLANG_PLUGINS=foundry` to make selection explicit and exclude other general + plugins during validation. +- This removes all four historical direct SGLang edits (no `foundry_shim.py`, no + `ServerArgs` field, no `apply_server_args` calls in `run_scheduler_process` or + `run_data_parallel_controller_process`). The net upstream direct-edit + footprint is zero. + +Because `load_plugins()` swallows plugin import/execution exceptions +(`plugins/__init__.py`) and `HookRegistry.apply_hooks()` swallows +hook-application exceptions (`hook_registry.py`), Foundry must fail closed from a +hook that is **subsequently executed** — an AROUND/AFTER hook on +`ServerArgs.__post_init__` — not from an exception raised inside `activate()` or +from a missing hook target. + +## Hook Target Mapping + +All hooks are registered via `HookRegistry.register(target, fn, HookType.*)` +using dotted upstream targets. Foundry reuses its existing patch **bodies** where +they still apply, but binds them through the registry so descriptor preservation +and stale-binding propagation are handled by upstream. + +| Concern | Target | Hook type | Foundry action | +|---|---|---|---| +| Fail-closed profile validation | `sglang.srt.server_args.ServerArgs.__post_init__` | AROUND/AFTER | Validate the full profile; raise here (not swallowed). | +| LD_PRELOAD before spawn | `sglang.srt.entrypoints.engine.Engine._launch_scheduler_processes` (classmethod) | AROUND | Call `runtime.setup_ld_preload_env()` before original. | +| Rank/VMM setup | `sglang.srt.model_executor.model_runner.ModelRunner.init_torch_distributed` | AROUND | Reserve VMM region using `self.ps.{tp_rank,pp_rank,dp_rank}` before original; scratch skip after. | +| Memory-pool config | `sglang.srt.model_executor.model_runner.ModelRunner.alloc_memory_pool` | AROUND | SAVE: run original, serialize resulting `MemoryPoolConfig`. LOAD: seed saved config + mirror `empty_cache`. | +| Memory-pool resolver reuse | `sglang.srt.mem_cache.kv_cache_configurator.KVCacheConfigurator._resolve_memory_pool_config` | AROUND/REPLACE | LOAD only: return the saved `MemoryPoolConfig` instead of profiling. | +| Full-backend SAVE | `sglang.srt.model_executor.runner_backend.full_cuda_graph_backend.FullCudaGraphBackend` | REPLACE (class) with `FoundryFullCudaGraphBackend(FullCudaGraphBackend)` | Override `capture_one` for SAVE; delegate otherwise. | +| Decode capture SAVE/LOAD | `sglang.srt.model_executor.runner.decode_cuda_graph_runner.DecodeCudaGraphRunner.capture` | SAVE AROUND / LOAD REPLACE | SAVE: original loop then manifest/fatbins/final offset. LOAD: pool setup, prealloc, metadata prepass, one-shot batched load, populate backend maps. | +| Defensive DP spawn | `sglang.srt.managers.data_parallel_controller.DataParallelController.launch_tensor_parallel_group` | AROUND | Optional; parent engine preload suffices for TP=2. | + +Removed patches that must **not** be recreated: `init_memory_pool`, +`load_model` (no-op wrapper), `kernel_warmup`, `_create_device_graph`, +`_capture_graph`, `capture_one_batch_size`, and the temporary instance patch of +`init_forward_metadata_capture_cuda_graph` (legacy method removed by #26735). + +## Ranks and Distributed Lifecycle + +- Startup rank truth is the frozen `ParallelState` (`parallel_state_wrapper.py`) + stored on `ModelRunner.ps`, exposing `tp_rank`, `tp_size`, `pp_rank`, + `pp_size`, `dp_rank`, `dp_size`, and `gpu_id`. The Foundry rank/VMM hook must + read `self.ps.tp_rank`, `self.ps.pp_rank`, and `self.ps.dp_rank`. +- `config.compute_workspace_rank(server_args, tp_rank, pp_rank, dp_rank)` remains + conceptually correct: DP-attention embeds DP inside TP so workspace rank is + `pp * tp_size + tp`; regular replicated DP needs + `dp * tp_size * pp_size + pp * tp_size + tp`. For TP=2/PP=DP=1 this reduces to + workspace ranks 0 and 1. +- `ModelRunner.__init__` already binds the device via + `torch.get_device_module(self.device).set_device(ps.gpu_id)` before + `init_torch_distributed`. Foundry still reserves the VMM region around + `init_torch_distributed` (before NCCL/process-group allocations), but the + explicit second `set_device` in the current hook is now redundant. +- Distributed init delegates to `bootstrap.init_torch_distributed(...)`, which + calls `_set_all_reduce_flags(server_args=...)` + (`set_custom_all_reduce(not disable_custom_all_reduce)`, + `set_mscclpp_all_reduce(enable_mscclpp)`, + `set_torch_symm_mem_all_reduce(enable_torch_symm_mem)`) and + `_init_parallel_groups(...)`. +- For startup naming use `ModelRunner.ps`; after groups exist use `get_tp_group()` + or `runtime_context.get_parallel()`. + +## Memory-Pool SAVE/LOAD Reuse + +The convenient-looking `alloc_memory_pool(memory_pool_config=...)` path does +**not** reuse a supplied config for a normal target worker. +`KVCacheConfigurator.configure(*, pre_model_load_memory)` calls +`_resolve_memory_pool_config(pre_model_load_memory)` (which profiles device +memory) for target workers; only speculative draft workers consume the supplied +object. Therefore: + +- **SAVE:** AROUND `ModelRunner.alloc_memory_pool`; run the original, then + serialize `dataclasses.asdict(model_runner.memory_pool_config)` into + `warmup_state.json`. +- **LOAD:** before `alloc_memory_pool`, set + `MemoryPoolConfig(**saved_state)` on `model_runner.memory_pool_config`; mirror + the SAVE `empty_cache()` side effect; AROUND/REPLACE + `KVCacheConfigurator._resolve_memory_pool_config` to return that saved config. + Let upstream `configure`, `_derive_pool_sizes`, `_init_pools`, and + `_init_post_memory_pool_components` run unchanged. +- Pin `SGLANG_ENABLE_POST_CAPTURE_KV_SIZING=0` so no post-capture resize sizes + pools from a different free-memory measurement than SAVE. +- `MemoryPoolConfig` is defined in + `python/sglang/srt/model_executor/pool_configurator.py` with + `max_total_num_tokens`, `max_running_requests`, hybrid/DSV4 pool sizes, and + `mem_fraction_static`. + +Retain Foundry's `torch.cuda.empty_cache()` SAVE/LOAD symmetry until an H100 +offset trace proves it unnecessary. + +## SAVE/LOAD Capture Architecture + +### SAVE + +- Class-REPLACE `FullCudaGraphBackend` with + `FoundryFullCudaGraphBackend(FullCudaGraphBackend)`. Override + `capture_one(shape_key, forward_fn, capture_inputs=None, post_warmup_hook=None)` + to skip the two ordinary warmup forwards, create a `foundry.CUDAGraph`, capture + `forward_fn()` under `foundry.graph(cuda_graph=graph, pool=self._pool, + stream=self._capture_stream)`, store into `self._graphs[shape_key]` and + `self._outputs[shape_key]`, and save the graph using the complete `ShapeKey`. + Ordinary (Foundry-inactive) mode delegates to the base implementation + unchanged. +- AROUND `DecodeCudaGraphRunner.capture`: run the upstream loop (which iterates + `reversed(self.capture_bs)`, calls `capture_one_shape`, allocates each + FlashInfer wrapper once via `init_forward_metadata_out_graph(fb, + in_capture=True)`, and calls `backend.capture_one`), then run + `save_graph_manifest()`, `pack_fatbins()`, and `capture_final_alloc_offset()` + once after all shapes. +- Because the upstream runner already allocates each per-bs wrapper exactly once + in `reversed(capture_bs)` order, the SAVE-side FlashInfer pre-pass and reuse + shim from the current fork integration should not be carried over unless an + H100 offset trace shows duplicate allocation recurs. + +### LOAD + +There is no backend restore method in the ABC. LOAD REPLACEs +`DecodeCudaGraphRunner.capture` while leaving the rest of the runner intact: + +1. establish `backend._pool` via `get_or_create_global_graph_memory_pool` and + call `set_graph_pool_id` (mirroring `capture_session`); +2. mirror the runner's capture-time `seq_lens` fill; +3. call `runtime.preallocate_for_load_mode()`; +4. run the FlashInfer metadata prepass for every shape in + `reversed(runner.capture_bs)` order via + `init_forward_metadata_out_graph(fb_view, in_capture=True)`; +5. call `FoundryCUDAGraph.start_graph_builds(all_paths)` once and + `finish_graph_loads(...)` once (the single batched call is load-bearing for + manifest template/on-demand linking — per-graph calls leave on-demand graphs + without `shared_exec` and replay aborts); +6. reconstruct complete `ShapeKey` instances and populate `backend._graphs` and + `backend._outputs`. + +Runtime replay is unchanged: `DecodeCudaGraphRunner.execute` computes the same +`ShapeKey` and calls `backend.replay(shape_key, forward_batch)`; if LOAD +populated the native backend maps, no replay monkey-patch is needed. + +### ShapeKey migration + +`ShapeKey` is a frozen dataclass with `size: int`, `stream_idx: Optional[int] = +None`, `variant_label: Optional[str] = None`. For the target dense profile +`stream_idx=None` and `variant_label=None`, so `size` is batch size. The archive +manifest must serialize all three fields (or fail closed unless the latter two +are `None`). The current `_batch_size_from_key` cannot parse the dataclass +`repr`; the key scheme must migrate to reconstruct `ShapeKey(size=bs)`. + +### Output packing + +Current `_pack_output`/`_unpack_output` preserve only +`LogitsProcessorOutput.next_token_logits`. Current decode outputs may contain +more than `next_token_logits`; the narrow packer must assert the supported +non-speculative, no-hidden-state decode output shape rather than silently +dropping fields. + +## Helper Corrections Required with the Port + +- `build_capture_fb_view` reads `cuda_graph_runner.num_tokens_per_bs`; current + decode runners expose `captured_req_width`. +- graph ownership moves from runner `graphs`/`output_buffers` to + `backend._graphs`/`backend._outputs`. +- graph pool helpers moved to `model_executor.runner_utils.pool` / exports from + `runner`; the old `cuda_graph_runner.get_global_graph_memory_pool` / + `set_global_graph_memory_pool` imports must be repointed. +- DeepEP helpers (`bootstrap_deepep_buffer`, the `deep_ep.Buffer.__init__` + patch) are out of scope for dense TP=2 and stale (`DeepEPBuffer._buffer` moved + to `get_resources().buffers["deepep_ep_state"]`). Retain only where naturally + isolated for validated DeepEP profiles; do not carry into the dense TP port. + +## Recipe and Harness Migration + +- Recipe `recipe/experimental/serve_qwen3-8b_sglang_tp.sh`: drop + `--foundry-graph-extension-config-path`; export + `FOUNDRY_GRAPH_EXTENSION_CONFIG_PATH`, `SGLANG_PLUGINS=foundry`, and + `SGLANG_ENABLE_POST_CAPTURE_KV_SIZING=0`; replace + `--disable-piecewise-cuda-graph` and `--cuda-graph-max-bs 256` with + `--cuda-graph-config '{"decode":{"backend":"full","max_bs":256},"prefill":{"backend":"disabled"}}'`; + keep `--tp-size 2`, `--random-seed 42`, `--mem-fraction-static`, + `--disable-radix-cache`, `--disable-custom-all-reduce`, `--attention-backend + flashinfer`, `NCCL_CUMEM_ENABLE=0`, `NCCL_NVLS_ENABLE=0`, + `PYTORCH_CUDA_ALLOC_CONF=expandable_segments:False`. +- Harness `tests/modal_sglang_tp.py`: set `SGLANG_REPO=sgl-project/sglang`, pin + `SGLANG_BASE=a25164bda355250015c0ba6d8821f0d4b37abe7d`, drop the fork fetch and + cherry-pick, install with `pip install -e 'python[all]'`, pin the immutable + model revision, and export the plugin activation env. + +## Failure Handling + +The supported path fails rather than degrades silently when: + +- `FOUNDRY_GRAPH_EXTENSION_CONFIG_PATH` points at a missing or malformed TOML; +- the plugin entry point is not discoverable in the clean checkout; +- any hook target does not resolve on the pinned SHA; +- the requested profile enables an unsupported topology or collective backend, + a non-full decode backend, CUMEM/NVLS, or post-capture KV sizing; +- TP rank archives are incomplete; +- SAVE passes produce incompatible graph inventories, fingerprints, or final + offsets; +- LOAD restores fewer graphs than archived, saves a new graph, or invokes native + recapture; +- NCCL does not report the intended P2P/IPC transport; or +- server logs contain CUDA, NCCL, MMU, Xid, illegal-access, worker-death, or + scheduler-failure signatures. + +Fail-closed validation runs from an executed `ServerArgs.__post_init__` hook so +the exception is not swallowed by the plugin loader. Harness timeouts are hard +external bounds so collective deadlocks cannot run indefinitely. + +## Verification Design + +### CPU / clean-checkout contract + +- the plugin entry point is discoverable (`entry_points(group="sglang.srt.plugins")` + lists `foundry`); +- `activate()` with the config env unset is a no-op; +- `activate()` with the config env set registers the expected `HookRegistry` + targets; +- every hook target resolves on the pinned SGLang checkout (a rename fails + loudly in CI, not silently on GPU); +- activation/target failure is fatal through an executed hook; +- `ShapeKey` archive round-trip covers all three fields; +- the target memory LOAD path proves `_resolve_memory_pool_config` does not + profile; +- recipes assert the exact NCCL-only env/flags and the `decode.backend=full` + config. + +### 2x H100 acceptance (Modal `rahul-dev`) + +Five phases: baseline ordinary SGLang, Foundry SAVE, second independent SAVE, +fresh-process LOAD, and a second fresh-process LOAD restart. Required evidence: + +- token-identical baseline/SAVE/LOAD output; +- identical per-rank graph counts and allocation offsets across both SAVEs; +- LOAD final offset equals SAVE final offset; +- no graph capture/recapture during LOAD; +- every NCCL channel uses the expected P2P/IPC transport; +- no custom/symmetric/MSCCL++ path in logs; +- no CUDA/NCCL/MMU errors or hangs; +- recorded SGLang SHA, Foundry SHA, model revision, CUDA, driver, PyTorch, NCCL, + FlashInfer, and sgl-kernel versions. + +### Local quality gates + +Focused CPU contract tests, Python compilation, shell syntax checks, +`pre-commit run --all-files`, and `git diff --check origin/main...HEAD`. GPU +work runs remotely because the local VM has neither CUDA nor an NVIDIA device. + +## Risks + +- **R1 main churn (high likelihood, medium impact):** the young target files + drift. Mitigation: pin an immutable SHA; the clean-checkout target-resolution + contract turns renames into loud failures. +- **R2 warmup/JIT determinism in the new `capture_one` (medium/high):** skipping + the two warmup forwards so JIT lands inside capture is the crux. Mitigation: + the SAVE/SAVE2 offset-reproducibility gate and alloc-event log inspection. +- **R3 memory-pool reuse path (medium):** the target resolver re-profiles; + mitigation is the LOAD resolver hook plus pinned sizing inputs and offset + equality. +- **R4 post-capture KV resize (low, mitigated):** env-gated off; recipe pins it + and the fail-closed guard catches a flipped default. +- **R5 decode-backend routing (low):** must land on `FullCudaGraphBackend`; + recipe forces `decode.backend=full` and the plugin asserts the backend type. +- **R6 NCCL non-determinism (inherent, out of scope):** addressed only by the + narrow P2P/IPC profile. + +The only hypothetical scenario forcing an upstream patch is a future main that +removes the `full` decode backend, stops routing dense decode through one +`torch.cuda.CUDAGraph` per shape, or makes `capture`/`capture_one` +un-replaceable. None hold at the pin. + +## Migration Strategy + +- Delete the `foundry-org/sglang` fork dependency from the SGLang design; the + plugin framework replaces its only content (the activation shim). +- Preserve legacy support only where it is naturally isolated and free — the + main architecture governs. Do not maintain a dual monolithic/decomposed runner + adapter; the split is too large to bridge without distorting the design. +- Two-phase landing: (1) fork-free activation + recipe/harness/env migration + + CPU/clean-checkout contracts (all runnable on the no-GPU VM); (2) decomposed + capture/LOAD hooks + memory-pool reuse + H100 acceptance. + +## Acceptance Criteria + +- Foundry activates through the `sglang.srt.plugins` entry point and + `HookRegistry` with zero upstream edits and no fork commit. +- The pinned SHA `a25164bda355250015c0ba6d8821f0d4b37abe7d` and immutable model + revision are the sole validation target. +- SAVE binds to `FoundryFullCudaGraphBackend.capture_one` and the SAVE-side + `DecodeCudaGraphRunner.capture` AROUND hook; LOAD REPLACEs + `DecodeCudaGraphRunner.capture`; keys are `ShapeKey`. +- LOAD reuses the serialized `MemoryPoolConfig` without profiling. +- The recipe enforces the exact NCCL-only, full-decode, post-capture-sizing-off + profile and fails closed outside it. +- The clean-checkout target-resolution contract passes and every hook target + resolves. +- Final-head 2x H100 baseline/SAVE/SAVE2/LOAD/LOAD-restart validation passes with + the direct rank/offset/archive/transport/output/error evidence above. +- Documentation makes no claim broader than the pinned-commit evidence. diff --git a/docs/superpowers/specs/2026-07-23-sglang-symmetric-tp-treatment-design.md b/docs/superpowers/specs/2026-07-23-sglang-symmetric-tp-treatment-design.md new file mode 100644 index 00000000..a7e10fe4 --- /dev/null +++ b/docs/superpowers/specs/2026-07-23-sglang-symmetric-tp-treatment-design.md @@ -0,0 +1,238 @@ +# SGLang Symmetric-TP Treatment Design Amendment + +## Status and authority + +This document amends +[`2026-07-23-sglang-main-port-design.md`](2026-07-23-sglang-main-port-design.md) +for SGLang tensor parallelism. Its collective, graph-shape, verification, and +support decisions supersede the earlier design's plain-NCCL requirements. +The zero-edit plugin architecture, pinned upstream target, memory-pool reuse, +and current SAVE/LOAD hardening remain in force. + +## Evidence boundary + +- SGLang pin: + `a25164bda355250015c0ba6d8821f0d4b37abe7d`. +- Model: Qwen3-8B revision + `b968826d9c46dd6066d109eabc6255188de91218`. +- Disproving H100 app: `ap-2qSQUKlY24L4RrQkMpUrmc`, TP=2 on two H100s. +- In that run, baseline and independent SAVE2 output were byte-identical. + Both fresh LOAD processes produced the same corrupt decode output. +- All 36 graphs per rank restored, graph counts and final offsets matched, the + archive stayed immutable, and NCCL reported the expected P2P/IPC topology. + Corruption began at restored decode replay. +- Every archived graph contained 73 NCCL all-reduce nodes and one NCCL + all-gather node whose argument bytes retained SAVE-process communicator and + graph-sequence state. Foundry does not relocate or reconstruct that state. + +Therefore plain-NCCL fresh-process graph restore is **rejected** for SGLang at +this pin. Equal offsets, archive bytes, and NCCL channels are necessary +lifecycle evidence but do not establish communicator ownership or replay +correctness. + +## Selected policy + +Adopt the proven final PR4 symmetric-memory policy at the current +`HookRegistry` architecture: + +1. Enable SGLang torch symmetric memory and retain one TP communicator. +2. Allocate and rendezvous its symmetric buffer outside Foundry's VMM region. +3. Bypass `MultimemAllGatherer` lazy state and call + `torch.ops.symm_mem.multimem_all_gather_out` directly with a stable slice of + the retained communicator buffer and that communicator's original group + name. +4. Synchronize the device and execute the TP barrier immediately before + entering each Foundry graph-capture context. +5. Pass the backend's shared pool explicitly to the one + `FoundryCUDAGraph.start_graph_builds(...)` call. +6. Capture exactly one padded decode shape, batch size 256. +7. Fail closed on the exact recipe/profile and require symmetric-specific + runtime and archive evidence. + +The relevant SGLang APIs are unchanged between the proven PR4 pin and +`a25164b`; only the Foundry integration mechanism changes from direct +monkeypatching to named `HookRegistry` callbacks. + +## Runtime architecture + +### Stable communicator ownership + +Add AROUND targets for +`TorchSymmMemCommunicator.__init__` and `MultimemAllGatherer.__call__`. +The constructor callback is inactive-mode passthrough. In SAVE or LOAD it: + +- raises the torch symmetric-memory signal-pad size to at least + `_MAX_BLOCKS * world_size * 4`; +- calls `foundry.ops.stop_allocation_region()`, invokes the original + constructor, and always calls `resume_allocation_region()` in `finally`; +- retains successful communicators by + `tuple(torch.distributed.get_process_group_ranks(communicator.group))`; +- replaces a duplicate rank-set entry instead of treating repeated + construction as another TP communicator; and +- logs a rank-visible symmetric-buffer address marker. + +This keeps PyTorch's symmetric allocation and rendezvous outside Foundry's +monotonic VMM allocator while preserving one stable communicator identity. + +### Direct stable multimem gather + +In SAVE and LOAD, the gather callback requires exactly one retained +communicator, contiguous rank-2 `torch.bfloat16` input, and sufficient +communicator-buffer capacity. For input `[tokens, local_hidden]`, it: + +```python +required_numel = communicator.world_size * x.numel() +comm_buff = communicator.buffer[:required_numel].view( + communicator.world_size * x.shape[0], + x.shape[1], +) +torch.ops.symm_mem.multimem_all_gather_out( + x, + communicator.group.group_name, + comm_buff, +) +return ( + comm_buff.view(communicator.world_size, x.shape[0], x.shape[1]) + .movedim(0, 1) + .reshape(x.shape[0], communicator.world_size * x.shape[1]) +) +``` + +The callback logs a multimem-buffer marker and never invokes upstream lazy +`_build`/`create_state` in an active Foundry mode. In inactive mode it delegates +unchanged. + +### Capture and load invariants + +`FoundryFullCudaGraphBackend.capture_one` keeps the existing no-warmup-forward +policy. Immediately before entering +`foundry_graph_ctx(graph, pool=self._pool, stream=self._capture_stream)`, it +calls `self._device_module.synchronize()` and `self._tp_group.barrier()`. + +LOAD keeps one start plus capture-index-ordered per-graph finishes: + +```python +pending = FoundryCUDAGraph.start_graph_builds( + paths, + pool=backend._pool, + num_threads=4, +) +``` + +Each FlashInfer wrapper initialization remains immediately before its matching +`finish_one_graph_load(pending, index)`. The explicit pool fixes a current +binding/API invariant; it does not authorize an independent per-file loader. + +### Strict profile and recipe + +The sole treatment profile is single-host Qwen3-8B, TP=2, +PP/DP/DCP/EP=1, FlashInfer attention, full decode graphs, and disabled prefill +graphs. It requires: + +- `enable_torch_symm_mem=True` and `enable_symm_mem=False`; +- `disable_custom_all_reduce=True`; +- `enforce_disable_flashinfer_allreduce_fusion=True` and + `flashinfer_allreduce_fusion_backend=None`; +- `disable_flashinfer_autotune=True`; +- decode `max_bs == 256` and explicit `bs == [256]`; +- `NCCL_CUMEM_ENABLE=0`, `NCCL_NVLS_ENABLE=0`, + `PYTORCH_CUDA_ALLOC_CONF=expandable_segments:False`, and + `SGLANG_ENABLE_POST_CAPTURE_KV_SIZING=0`. + +Missing resolved fields are violations. The canonical recipe argument is: + +```text +--cuda-graph-config '{"decode":{"backend":"full","max_bs":256,"bs":[256]},"prefill":{"backend":"disabled"}}' +``` + +## Preserved current-tree behavior + +The treatment changes only collective stability, pre-capture lockstep, explicit +pool delivery, and graph-shape/profile policy. It preserves: + +- zero-edit SGLang activation and the named `HookRegistry` target inventory; +- rank-zero-only atomic `warmup_state.json` replacement, exact + `MemoryPoolConfig` schema/value validation, and malformed-state propagation; +- complete dense `ShapeKey(size, stream_idx, variant_label)` validation; +- fail-closed output packing; +- LOAD `seq_lens` reset and FlashInfer wrapper/graph-finish interleaving; +- one-start/template-aware loading and per-entry background-build error + propagation; +- fixed-VMM alignment and explicit graph pool-default handling; +- archive/process-group cleanup, byte and semantic archive immutability, + immutable revisions, and two-window fatal scanning. + +## Explicit non-goals + +Do not port any superseded PR4 experiment: + +- no second/device-group rendezvous; +- no captured rendezvous handle or synthesized `MultimemAllGatherState`; +- no eager or per-shape warmup forwards; +- no `create_state` interception; +- no independent per-file graph loader or `strip_dependencies=False` residue. + +Do not add hybrid Mamba state, DeepEP, speculative decode, LoRA, pdmux, +multi-host, or TP>2 support. + +## Verification sequence + +Modal Volume auto-commit may retain partial or mutated unique-run artifacts +after a failed phase. The harness marks those artifacts retained, failures stop +later phases, and the affected run is never verified. An explicit commit +remains a success-path durability request only. + +### Focused SAVE-to-LOAD treatment + +Before the full matrix, use a new empty archive and run only a fresh SAVE then a +fresh LOAD on two H100s. Use the exact pins above, seed 42, one fixed +temperature-0 prompt, and `max_new_tokens=32`. + +The treatment is `VERIFIED` only when: + +- SAVE and LOAD response bytes are equal, nonempty, and deterministic; +- each rank saves/restores exactly one `ShapeKey(size=256)` graph; +- both phases log the retained symmetric communicator on ranks 0 and 1 and at + least one direct multimem gather marker globally; +- every rank's archived graph JSON contains positive + `CUDASymmetricMemoryOps*all_reduce*` and + `CUDASymmetricMemoryOps*all_gather*` KernelNode function evidence; +- every archived graph JSON contains zero + `ncclDevKernel_AllReduce*` and zero `ncclDevKernel_AllGather*` nodes; +- LOAD performs no native capture/save, replays the restored graph, and leaves + every archive byte unchanged; +- nonzero rank offsets agree across ranks and LOAD equals SAVE; and +- shutdown is clean with no fatal CUDA, NCCL, Foundry, MMU, Python, or + early-process failure. + +Any output mismatch, archived NCCL collective, fallback, archive mutation, +recapture, or offset mismatch is `NOT VERIFIED`. An unrelated setup failure +before graph replay is `INCONCLUSIVE`. Do not run the full matrix after either +result. + +### Full acceptance + +Only after focused verification, run +baseline → SAVE → SAVE2 → LOAD → LOAD2. Retain all existing exact +sequential/concurrent output, archive reproducibility/immutability, two fresh +LOAD, replay, offset, teardown, provenance, and fatal-scan checks. Replace the +plain-NCCL channel acceptance condition with: + +- one graph per rank; +- symmetric communicator markers on every rank and at least one direct + multimem marker globally in every Foundry phase; +- positive symmetric-memory all-reduce/all-gather KernelNode function evidence + in every rank archive; +- zero archived NCCL all-reduce/all-gather nodes; and +- `COLLECTIVE_BACKEND=torch_symmetric_memory`. + +NCCL version and transport logs may remain diagnostic provenance, not a +correctness gate for restored collectives. + +## Acceptance + +This amendment is implemented only when CPU contracts and pinned target +resolution pass, the focused treatment is `VERIFIED`, and the subsequent full +five-phase run passes at one immutable Foundry commit. Documentation must record +both run identifiers and must not claim support beyond the exact pins and +profile above. diff --git a/docs/superpowers/specs/2026-07-23-tp-upstream-consolidation-design.md b/docs/superpowers/specs/2026-07-23-tp-upstream-consolidation-design.md new file mode 100644 index 00000000..a3d4124b --- /dev/null +++ b/docs/superpowers/specs/2026-07-23-tp-upstream-consolidation-design.md @@ -0,0 +1,300 @@ +# Tensor-Parallel Upstream Consolidation Design + +> [!WARNING] +> **SUPERSEDED FOR SGLANG; vLLM portions remain applicable.** The shared +> plain-NCCL policy below is historical for SGLang. Use the +> [symmetric-TP treatment design](2026-07-23-sglang-symmetric-tp-treatment-design.md) +> and +> [implementation plan](../plans/2026-07-23-sglang-symmetric-tp-treatment.md) +> for current SGLang policy and evidence. The vLLM design remains unchanged. + +## Goal + +Produce one reviewable branch whose complete diff is directly against +`origin/main` and that: + +- preserves every unique change on the repository's active branches and draft + pull requests; +- supports and verifies dense tensor parallelism for both vLLM and SGLang; +- defines a narrow, explicit support envelope instead of implying that every + communication backend works with fresh-process CUDA graph restoration; and +- records reproducible evidence from real multi-GPU SAVE-to-LOAD runs. + +The branch starts at `origin/main`, not at another feature branch. Patch +duplicates are represented once, and superseded fixes are replaced by their +latest hardened form. + +## Evidence Considered + +### Foundry repository + +Foundry has issues and GitHub Discussions disabled. The relevant upstream +history is therefore in draft PRs and their review threads: + +- PR #1: VMM-backed CUDA IPC and no-fabric DeepEP. +- PR #2: SGLang dense TP, H200 validation, and the post-SGLang-#26735 attention + metadata API port. +- PR #4: main-rooted consolidation, SGLang TP harness, IPC/RNG hardening, graph + pool fix, archive comparison, and stronger replay checks. +- PR #5: vLLM TP recipe, contract tests, immutable-revision Modal harness, and + compatible CUDA-wheel pinning. + +The standalone `ep-ipc`, `a100-issue`, no-fabric DeepEP, and TP branches are +either represented by patch-equivalent commits in PR #4 or contain newer +commits that must be reconciled explicitly. + +### vLLM upstream + +Upstream reports unresolved lifecycle and correctness concerns for VMM-backed +custom all-reduce, symmetric-memory communication, FlashInfer fused +all-reduce, and CUDA graph capture. Important examples include vLLM issues or +PRs #39788/#39791, #42609/#43923, #46515, #47806, #48071/#48075, and +PRs #49224/#49233. + +The Foundry vLLM fork at `4309c257d3f639e5490d3811293c890c61c76f29` +does not claim TP support and predates relevant upstream fixes. Foundry only +patches its V1 graph runner. + +### SGLang upstream + +SGLang PR #26735 replaced the legacy attention graph-metadata methods with the +three-method `init_forward_metadata` contract. Current upstream has since +replaced the monolithic graph runner as well, so this change supports the +current Foundry fork contract but does not claim compatibility with arbitrary +SGLang `main`. + +Upstream SGLang also reports captured symmetric-memory numerical corruption, +custom-all-reduce VMM incompatibilities, communicator concurrency hangs, and +rank-divergent FlashInfer tuning. Relevant examples include PRs #27593, +PRs #29946, #31049, #31135 and issue #31117. + +These findings make a constrained NCCL path a correctness requirement, not +merely a recipe preference. + +## Supported Tensor-Parallel Profile + +The first supported profile is intentionally narrow: + +- one host; +- two Hopper GPUs, initially validated on 2x H100 and with preserved + historical 2x H200 evidence; +- dense Qwen3-8B; +- TP size 2, with no DP, PP, EP, or disaggregated serving; +- plain NCCL P2P/legacy CUDA IPC collectives; +- Foundry SAVE and LOAD in fresh processes using identical rank/GPU ordering; +- immutable Foundry and engine revisions and an immutable model revision; and +- full decode CUDA graphs only. + +The profile disables: + +- engine custom all-reduce; +- Torch, NCCL, and vLLM symmetric-memory all-reduce; +- FlashInfer fused all-reduce/RMS communication; +- NCCL CUMEM and NVLS; +- piecewise CUDA graphs; +- FlashInfer autotuning during SGLang graph setup; +- vLLM V2; +- weight offloading, sleep mode, and communication/compute fusion; and +- expandable CUDA allocator segments. + +Recipes, validation harnesses, and integration checks use the same transport +and graph settings for baseline, SAVE, and LOAD. A Foundry-enabled TP launch +outside this profile must fail early with an actionable message when the +engine exposes enough configuration to detect it. It must not silently choose +an unvalidated backend. + +This design does not claim support for TP greater than 2, multiple hosts, +Ampere, Blackwell, custom all-reduce, symmetric memory, CUMEM, NVLS, vLLM V2, +piecewise graphs, or current unpinned engine `main` revisions. + +## Consolidation Architecture + +### Main-rooted history + +Create one branch directly from the latest `origin/main`. Reconcile source +branches by patch identity and final behavior rather than merging their branch +tips: + +1. Bring in the hardened VMM/IPC, no-fabric DeepEP, SGLang TP, A100/FA2, RNG, + graph-pool, archive, and replay work from the latest PR #4 head. +2. Bring in the vLLM TP recipe, CPU contract tests, Modal harness, and pinned + extension-wheel setup from PR #5. +3. Bring in PR #2's post-#26735 SGLang attention API compatibility change and + preserve its H200 evidence without restoring stale claims. +4. Re-fetch all remote heads immediately before integration and account for + any new unique patches. +5. Resolve overlaps in source, tests, recipes, and documentation against the + desired final behavior; do not retain duplicate commits solely to preserve + branch ancestry. + +The resulting merge base with `origin/main` must equal the current +`origin/main` commit. The PR base is `main`, and the review diff is +`origin/main...HEAD`. + +### Core VMM and graph restoration + +The consolidated core retains: + +- SCM_RIGHTS transfer for VMM allocation handles instead of process-local raw + file-descriptor values; +- versioned exporter PID/token and preallocated-chunk identities; +- relocated peer mappings outside the deterministic Foundry region; +- whole-chunk export/import with suballocation offsets; +- duplicate-open reference counting and correct close semantics; +- fork-safe descriptor cleanup; +- support for multiple live preallocation ranges; +- deterministic generator-state restoration only when a graph actually uses + RNG state; and +- an explicit default graph pool compatible with the required C++ binding. + +These changes are shared prerequisites for both engines. + +### Engine adapters + +SGLang retains per-TP-rank workspaces, device binding before distributed +initialization, the FlashInfer metadata pre-pass, and fresh NCCL initialization +before restored graph replay. Its attention adapter supports both the legacy +method and the post-#26735 three-method API used by the pinned fork revision. + +vLLM retains per-rank workspaces and its two-pass SAVE lifecycle. LOAD starts +graph construction after weight loading and finalizes graphs in capture order +to reproduce allocator movement. The TP recipe selects only PyNCCL and V1 +`FULL_DECODE_ONLY` graphs. + +Both integrations expose the constrained profile clearly in recipes and +documentation. Validation logs must show the selected communication path. + +## Runtime Flow + +For each rank: + +1. Select the rank's CUDA device. +2. Reserve the Foundry deterministic VMM region. +3. Initialize the fresh NCCL communicator using the constrained transport. +4. Advance beyond the configured scratch region. +5. Allocate model, communication, and graph state deterministically. +6. During SAVE, capture and serialize rank-local graphs, packed modules, + manifests, and final allocation metadata. +7. During LOAD, preallocate the recorded range, load modules, rebuild graph + templates, reproduce engine initialization allocations, and restore graphs + without native recapture. +8. Replay the restored decode graphs and serve real generation requests. + +Each rank writes to `rank_`. SAVE and LOAD must agree on graph +counts and final offsets per rank; equal offsets across different ranks are +observed evidence, not a general requirement. + +## Failure Handling + +The supported path fails rather than degrades silently when: + +- an immutable source or model revision is missing or malformed; +- the expected engine extension is absent or ABI-incompatible; +- TP rank archives are incomplete; +- SAVE passes produce incompatible graph inventories, semantic graph + fingerprints, or final offsets; +- LOAD restores fewer graphs than were archived, saves new graphs, or invokes + native recapture; +- the requested TP profile enables an unsupported communication or graph + backend; +- the engine API is neither of the explicitly supported SGLang metadata + contracts; +- NCCL does not report the intended P2P/IPC transport; or +- server logs contain CUDA, NCCL, MMU, Xid, illegal-access, worker-death, or + scheduler-failure signatures. + +Harness timeouts are hard external bounds so collective deadlocks cannot leave +validation running indefinitely. + +## Verification Design + +### Test-first contract coverage + +CPU tests check recipes and harnesses before GPU implementation changes are +accepted: + +- all required backend-disabling flags apply symmetrically to every phase; +- engine and model revisions are immutable; +- baseline, SAVE, and LOAD commands select identical graph behavior; +- unsupported settings are rejected; +- per-rank archive checks require manifests, graphs, binaries, packed modules, + and offsets; +- LOAD evidence requires complete restoration and zero recapture; and +- harness reports cannot pass without output and transport evidence. + +Native CUDA tests cover: + +- default graph-pool capture; +- RNG-free and Philox graph restoration; +- duplicate IPC opens/closes; +- exporter restart/token changes and fork cleanup; +- multiple preallocation chunks and relocation collisions; and +- fresh-process collective graph SAVE and LOAD where practical. + +### Engine end-to-end gates + +Run each engine on 2x H100 in Modal's `rahul-dev` environment: + +1. constrained ordinary baseline; +2. seed SAVE; +3. deterministic SAVE2 over the intended archive; +4. a second independent SAVE archive when the engine lifecycle permits; +5. fresh LOAD; +6. a second fresh LOAD restart. + +For vLLM, the seed SAVE is required for warmup-state stabilization. SGLang may +not require it architecturally, but uses the same multi-pass validation shape +to prove reproducibility. + +Each run must prove: + +- both ranks become healthy; +- nonzero expected graph counts exist on both ranks; +- SAVE inventories, semantic graph fingerprints, and per-rank final offsets + are reproducible; +- LOAD restores every archived graph, reaches the expected offset, and saves + no graph; +- NCCL uses the intended P2P/IPC path and the engine uses no forbidden + all-reduce backend; +- warmed SAVE2 and fresh LOAD return identical deterministic token IDs and + text for multiple prompts; +- the constrained baseline and LOAD agree where kernel selection is identical; +- sequential requests and concurrent batches exercise restored graphs; +- explicit log evidence shows CUDA graph replay; and +- no fatal signature or timeout occurs. + +Historical H200 results remain documented as supporting evidence, labeled with +their exact older revisions. They do not replace a rerun of the final +consolidated head. + +### Local quality gates + +Run focused tests, the complete CPU-import/contract suite, Python compilation, +shell syntax checks, project formatting/linting, and `git diff --check`. +GPU-native and engine tests run remotely because the local VM has neither CUDA +nor an NVIDIA device. + +## Documentation and Claims + +README, roadmap, release notes, engine docs, and recipe docs state the exact +validated profile. A green TP status means only the profile above. Broader +backend or topology support remains experimental until it has its own +fresh-process SAVE-to-LOAD evidence. + +The final PR description includes exact commit IDs, hardware, software +versions, Modal run links, graph counts, offsets, restart results, and output +comparisons. It also distinguishes historical evidence from runs performed on +the submitted commit. + +## Acceptance Criteria + +- One branch and one draft PR are based directly on `main`. +- Every unique active-branch patch is present or explicitly superseded. +- Patch duplicates appear once. +- SGLang supports the pinned post-#26735 attention API. +- Both engines expose a reproducible TP=2 recipe under the constrained NCCL + profile. +- Unsupported TP communication profiles fail clearly. +- Focused and native tests pass. +- Final-head 2x H100 SAVE-to-LOAD validation passes for vLLM and SGLang. +- Documentation makes no claim broader than the evidence. diff --git a/docs/vllm/direct-edits.md b/docs/vllm/direct-edits.md index a1e63f1f..a18cf1e7 100644 --- a/docs/vllm/direct-edits.md +++ b/docs/vllm/direct-edits.md @@ -55,6 +55,8 @@ class CompilationConfig: This is the earliest activation point in the parent process. As soon as a `CompilationConfig` is fully constructed, foundry's runtime patches are in place. +> Note: vLLM folds `graph_extension_config_path` into `CompilationConfig.compute_hash()` (the torch.compile cache key) even though it never affects codegen. This is harmless **as long as every SAVE/LOAD phase passes the same path string** — run all phases from one canonical CWD (or always an absolute path). If pass 1 and pass 2 see the path via different mount aliases (`/home/...` vs `/data/...`), pass 2 misses pass 1's warm cache and inductor recompiles inside the cuda-graph capture window → `cudaErrorStreamCaptureInvalidated`. We keep vLLM unpatched and rely on consistent invocation instead; the EP recipe scripts are launched from the workspace root. (An optional `ignored_factors` exclusion would harden this at the source but is not applied.) + ## 3. `vllm/v1/engine/core.py` (~6 lines) ```python diff --git a/docs/vllm/moe-and-deepep.md b/docs/vllm/moe-and-deepep.md index e76b48d8..e3b05df2 100644 --- a/docs/vllm/moe-and-deepep.md +++ b/docs/vllm/moe-and-deepep.md @@ -1,58 +1,126 @@ # MoE / DeepEP Support -vLLM-side MoE / expert-parallelism support is what makes the integration interesting beyond dense decode. The remaining MoE-specific concern is that DeepEP all-to-all communicator buffers must be set up in fabric mode for foundry to interpose them, and NVSHMEM module init must fire after the runtime is bootstrapped. +vLLM-side expert parallelism adds DeepEP communication buffers to the graph +archive contract. Foundry therefore treats the DeepEP transport and effective +NVLink allocation size as archived configuration, not launch-time switches. -The dense-only baseline doesn't exercise any of this; turn it on by running an MoE model with EP. +The dense-only profile does not exercise this path. ## MoE weight packing -vLLM's recent versions use modular kernels for both quantized and unquantized MoE — the same `model_loader.load_model()` path produces deterministic packing regardless of foundry mode. Foundry's integration takes the unified preallocate → `start_graph_builds` → weight-load → `prepare_communication_buffer_for_model` path for all MoE variants (and for dense models), so no quant-metadata roundtrip through `WarmupState` is needed. The earlier `collect_moe_quant_metadata` / `inject_moe_quant_metadata` / `reset_moe_quant_config` / `reinit_moe_quant_config` / `split_load_model` ceremony has been removed. +Recent vLLM versions use modular kernels for quantized and unquantized MoE. +The same `model_loader.load_model()` path produces deterministic packing in +every Foundry mode. Foundry uses one preallocate → graph-build → weight-load → +communication-buffer path, so no quantization metadata round trip through +`WarmupState` is required. -## DeepEP fabric mode +## DeepEP transport policy -### Problem +`deepep_transport` accepts exactly two TOML values: -`DeepEPLLAll2AllManager` builds the all-to-all communicator with one of several backends: +```toml +deepep_transport = "fabric" +deepep_transport = "nvl_ipc" +``` + +Omitting the field selects `fabric`. This is also how legacy archives that do +not contain either DeepEP field are interpreted. + +- `fabric` sets `use_fabric=True` and forces the effective + `num_nvl_bytes` to zero. +- `nvl_ipc` sets `use_fabric=False` and preserves vLLM's upstream, exact + positive `num_nvl_bytes`. It is an intranode profile. Foundry translates the + VMM-backed CUDA IPC handles through POSIX file descriptors and `SCM_RIGHTS`, + including buffers carved from a preallocated LOAD chunk. + +Neither profile removes DeepEP's NVSHMEM dependency. The low-latency RDMA +buffer remains on the NVSHMEM symmetric heap. -- NVL (NVLink direct) — uses `cudaIpcGetMemHandle` for buffer sharing. -- Fabric — uses NVSHMEM + symmetric memory. +## Archive contract -NVL's IPC handles are process-local and not persistable; the captured graph can't reference them on LOAD. Fabric mode allocates buffers via NVSHMEM, which foundry's hook can interpose (`init_nvshmem_for_loaded_modules` re-binds them on LOAD). +On SAVE, rank 0 records these fields in the root `warmup_state.json`: -### Solution +```json +{ + "deepep_transport": "nvl_ipc", + "deepep_num_nvl_bytes": 67108864 +} +``` + +The transport-only bootstrap state can be written before vLLM completes memory +profiling. During the first SAVE, the KV profiling path creates the complete +warmup state and explicitly preserves both DeepEP fields. Later SAVE passes +update the same metadata without dropping the archived transport profile. -`_patch_deepep` wraps `DeepEPLLAll2AllManager._make_all2all_kwargs`: +LOAD reads and validates `deepep_transport` while loading the TOML, before CUDA +module loading, VMM region setup, or communicator construction. A mismatch +fails with an error such as: -```python -@functools.wraps(orig) -def patched(self, *args, **kwargs): - kwargs = orig(self, *args, **kwargs) - if get_graph_extension_mode() != CUDAGraphExtensionMode.NONE: - kwargs["use_fabric"] = True - kwargs["num_nvl_bytes"] = 0 # RDMA-only path - return kwargs +```text +Foundry DeepEP transport mismatch: archived=nvl_ipc, configured=fabric ``` -Forces every foundry-mode DeepEP setup onto the fabric path. The patch lives behind a try/except in case the import path moves. +When vLLM builds `DeepEPLLAll2AllManager`, Foundry derives the effective NVLink +size from the current upstream kwargs and validates it before communicator +setup. A size mismatch fails with: -## NVSHMEM initialization order +```text +Foundry DeepEP num_nvl_bytes mismatch for nvl_ipc: archived=67108864, effective=... +``` -NVSHMEM modules must be loaded into device code memory **before** any captured graph that references NVSHMEM kernels is replayed. +There is no silent fallback between `nvl_ipc` and `fabric`. Either mismatch +changes allocation geometry and would invalidate graph addresses, so LOAD +fails closed. -1. `runtime.setup_graph_extension(...)` on LOAD calls `fops.set_skip_fatbin_processing(True)` then `fops.load_cuda_modules_and_libraries(workspace_dir)`. The loader queues each NVSHMEM-using module into `pending_nvshmem_init` but does not flush it (the NVSHMEM runtime isn't up yet). +## NVSHMEM discovery and initialization -2. `_load_model_with_overlap` runs `preallocate_for_load_mode()` and `start_graph_builds()`, then calls `do_original_load()` which eventually invokes `prepare_communication_buffer_for_model(model)`. The orig call constructs the DeepEP `Buffer`, which bootstraps the NVSHMEM runtime. +An explicit `nvshmem_host_path` in the TOML takes precedence. Otherwise Foundry +inspects the optional `nvidia.nvshmem` wheel and selects the first existing +library in this order: -3. `_patch_prepare_comm_buffer`'s post-orig hook then calls `fops.init_nvshmem_for_loaded_modules()` to flush the queue — well before any captured graph that uses NVSHMEM symbols is replayed by `finish_one_graph_load` inside `capture_model`. +1. `libnvshmem_host.so.3` +2. `libnvshmem_host.so` -If NVSHMEM init runs after the first `finish_one_graph_load`, the captured graph that calls into NVSHMEM aborts with `cuLaunchKernel` errors because the kernel handles don't resolve yet. If it runs *before* `prepare_communication_buffer_for_model` (the bug we hit when init was placed inside `start_graph_builds`), `nvshmemx_cumodule_init` silently no-ops because the runtime is not bootstrapped, and replay later sees uninitialized modules — also fatal. +On LOAD, CUDA module loading queues modules that use NVSHMEM. Model loading then +constructs the DeepEP buffer and bootstraps NVSHMEM. Only after that does +Foundry initialize the queued modules, before any restored graph is replayed. -For dense single-GPU models the NVSHMEM count is 0 and the call is a no-op; the post-orig hook is kept on all paths for parity. +## Strict two-GPU matrix -## Validation +The standalone runner's strict SAVE/LOAD cases are: -For a working MoE SAVE→LOAD cycle: +```bash +bash tests/run_deepep_matrix.sh fabric both +bash tests/run_deepep_matrix.sh nvl_ipc both +bash tests/run_deepep_matrix.sh nvl_ipc_prealloc both +``` -- All ranks produce `rank_{N}/warmup_state.json`-equivalent per-rank artifacts. -- All ranks' `final_alloc_offset` are reproducible across SAVE pass 1 → pass 2. -- On LOAD, the first inference call returns coherent tokens. +`fabric` is the zero-NVL fabric baseline (`use_fabric=True`). `nvl_ipc` adds a +64 MiB NVLink IPC buffer with `use_fabric=False`. `nvl_ipc_prealloc` uses the +same NVL IPC transport while LOAD restores from a preallocated Foundry VMM +chunk. All three disable NCCL CUMEM and NVLS. These names match +`tests/run_deepep_matrix.sh`; unknown names fail before execution. + +The strict matrix is verified at Foundry +`a8e43442421a80fea7991ee26f97f1003b170e29` in +[`ap-Q8mb8U4EYuE5Jo6Zb0xwrH`](https://modal.com/apps/modal-labs/rahul-dev/ap-Q8mb8U4EYuE5Jo6Zb0xwrH). +It used two H100s, CUDA `13.0.1-cudnn-devel-ubuntu24.04`, PyTorch +`2.11.0+cu130`, vLLM +`4309c257d3f639e5490d3811293c890c61c76f29`, DeepEP +`29d31c095796f3c8ece47ee9cdcc167051bbeed9`, and NVSHMEM `3.3.24`. +The pinned-PyTorch CPU preflight, all three SAVE/LOAD cases, and Modal CLI +returned zero; the application result was `"status": "passed"`. + +The output oracle was strict rather than a graph-structure proxy. It checked +the modern dispatch tuple, independently gathered top-k route multisets, exact +expert counts and split-field BF16 payloads, and complete metadata. Before +both SAVE and LOAD replay it filled every valid output slice with `NaN` and +every expert-count entry with `-1`; replay had to overwrite both sentinels and +exactly reproduce the expected tensors. `nvl_ipc_prealloc` additionally proved +the same output contract when the imported NVL buffer was carved from a +preallocated Foundry VMM chunk. + +This validates the standalone pinned transport/replay matrix. It does not +establish general DeepEP serving, MoE tensor parallelism, alternate DeepEP +revisions, or SGLang EP. The current SGLang EP profile remains quarantined +pending separate current-HookRegistry SAVE-to-fresh-LOAD evidence. diff --git a/docs/vllm/overview.md b/docs/vllm/overview.md index a96e9e87..67c3023d 100644 --- a/docs/vllm/overview.md +++ b/docs/vllm/overview.md @@ -6,6 +6,7 @@ Validated on: - Single-GPU Qwen3-1.7B / 4B / 14B - DP > 1 (multi-GPU dense) +- Dense Qwen3-8B TP=2 under the pinned profile below - Mixture-of-experts with DeepEP all-to-all (expert parallelism) ## Critical invariants (read this first) @@ -38,6 +39,62 @@ In reality, NVSHMEM's symmetric heap goes through the **large-alloc branch** of All MoE variants (quantized, unquantized) and dense models now share the same path: `preallocate_for_load_mode` → `do_original_load()` (weight load + `prepare_communication_buffer_for_model` + NVSHMEM init post-hook) → `start_graph_builds`. The background template builds were previously kicked off *before* `do_original_load` to overlap with weight IO; that was net-negative due to driver contention, so they now overlap with the cheaper post-load init work instead. +### 5. TP must use one pinned collective and graph profile + +The validated TP path is single-host, dense `Qwen/Qwen3-8B`, exactly two +Hopper GPUs, the vLLM V1 model runner, and `FULL_DECODE_ONLY`. Baseline, both +SAVE passes, and both fresh LOADs all use PyNCCL P2P/legacy CUDA IPC. + +The recipe disables custom all-reduce, torch and NCCL symmetric memory, +standalone FlashInfer all-reduce, the fused all-reduce/RMS compiler pass, NCCL +CUMEM, NCCL NVLS, expandable allocator segments, and vLLM V2. These settings, +including `graph_extension_config_path`, must be passed as one atomic +`--compilation-config` JSON object. Mixing dotted and object forms can replace +earlier nested values and silently select a different graph or collective path. + +## Verified dense TP=2 evidence + +The immutable validation inputs were Foundry +`a8e43442421a80fea7991ee26f97f1003b170e29`, vLLM +`4309c257d3f639e5490d3811293c890c61c76f29`, precompiled-wheel parent +`6cbe448eed751824d608faf9078ef84724d621c1`, and +`Qwen/Qwen3-8B@b968826d9c46dd6066d109eabc6255188de91218`. The +[final Modal run](https://modal.com/apps/modal-labs/rahul-dev/ap-YfKALdMwnyJQmzQoVvZSMU) +used two NVIDIA H100 80GB HBM3 GPUs in every phase, CUDA +`13.0.1-cudnn-devel-ubuntu24.04`, driver `580.95.05`, PyTorch +`2.11.0+cu130`, vLLM `0.1.dev16924+g4309c257d.precompiled`, and NCCL +`2.28.9+cuda13.0`. All 17 checks and all phase/CLI return codes passed. + +SAVE and SAVE2 each produced exactly 64 JSON/64 `.cugraph` files per rank, +with indices `0..63`, four required rank files, and root +`warmup_state.json`. Their inventories, graph counts, per-rank offset +`70724354048`, and 68 semantic fingerprints per rank matched. The published +snapshot had 265 complete-tree byte fingerprints. LOAD and LOAD2 each emitted +one `64/64` completion for ranks 0 and 1, matched the published live offset, +emitted no capture/save marker, and preserved all 265 fingerprinted files +byte-for-byte. Raw SAVE-versus-SAVE2 `.cugraph` bytes are not claimed equal; +immutability is checked from the published SAVE2 snapshot through each LOAD. + +For baseline, SAVE2, LOAD, and LOAD2, both the sequential and concurrent +queries returned these exact text/token-string pairs: + +```text +" Paris. The capital" [" Paris", ".", " The", " capital"] +" ?" [" ?\n\n"] +" H2O." [" H", "2", "O", "."] +" unit. It is" [" unit", ".", " It", " is"] +``` + +Cleanup removed both the authoritative Modal Dict marker and published Volume +path. The candidate also passed +[native H100 validation](https://modal.com/apps/modal-labs/rahul-dev/ap-0aeSCJNb6eHQhqdMkpyfXk) +with 10 focused VMM/RNG and 823 core tests, plus a local 320-test, compile, +pre-commit, diff, and direct-main lineage gate. + +This result does not cover TP>2, V2, multi-host, Ampere/Blackwell, MoE TP, or +alternate collectives. Standalone strict DeepEP evidence is documented +separately and does not widen this dense TP profile. + ## How to use ### 1. Install foundry + vLLM @@ -96,7 +153,9 @@ vLLM's lifecycle splits across the engine-core process and worker processes. Fou 1. Same VMM setup as pass 1. 2. `_initialize_kv_caches` skips the profile forward and uses saved block counts. 3. Same deterministic allocation order ⇒ same `final_alloc_offset`. -4. Captured graphs are byte-identical to pass 1's (already on disk — pass 2 overwrites them). +4. Graph inventory and normalized semantics are reproducible. Raw binary graph + bytes are not a cross-SAVE invariant; pass 2 overwrites the seed archive and + becomes the immutable publication candidate. **LOAD**: @@ -118,7 +177,8 @@ vLLM's lifecycle splits across the engine-core process and worker processes. Fou - [`memory-lifecycle.md`](memory-lifecycle.md) — VMM region setup, the SAVE↔LOAD allocation parity contract, two-pass SAVE rationale - [`save-load-workflow.md`](save-load-workflow.md) — serve scripts, TOML schema, expected logs, validation checks - [`memory-consistency.md`](memory-consistency.md) — what `memory_profiling` does that mandates two-pass SAVE, what's skipped on LOAD, and the silent-mismatch failure modes -- [`moe-and-deepep.md`](moe-and-deepep.md) — DeepEP `use_fabric=True` on LOAD, NVSHMEM init ordering +- [`moe-and-deepep.md`](moe-and-deepep.md) — archived DeepEP transport policy, + strict matrix evidence, and NVSHMEM init ordering ## Scope @@ -126,6 +186,7 @@ vLLM's lifecycle splits across the engine-core process and worker processes. Fou |---|---| | Dense decode (Qwen3, Llama, …) | primary test target | | DP > 1 | multi-GPU dense | +| Dense TP=2 | only the pinned single-host Qwen3-8B Hopper profile above | | MoE with DeepEP | unified path; no split-load | | `torch.compile` (full-graph) | runs on SAVE; `do_not_compile=True` on LOAD (faster startup; graphs replay via `CUDAGraphWrapper`) | diff --git a/include/CUDAGraphInternal.h b/include/CUDAGraphInternal.h index 52c53789..d3e53942 100644 --- a/include/CUDAGraphInternal.h +++ b/include/CUDAGraphInternal.h @@ -21,6 +21,10 @@ struct PendingGraphLoads { MempoolId_t pool; c10::DeviceIndex dev; CUDAGeneratorStateRegistry* registry = nullptr; // for deferred generator registration + // True if any graph in this pending set consumed Philox state. Generator + // state is shared across entries, so this must be decided for the full set + // before the first lazy seed/offset tensor allocation. + bool has_rng_generators = false; // Signaled when background graph building (Phase 2) completes. // finish_graph_loads_impl waits on this before allocator replay. diff --git a/include/foundry/vmm_cleanup.h b/include/foundry/vmm_cleanup.h new file mode 100644 index 00000000..1af76661 --- /dev/null +++ b/include/foundry/vmm_cleanup.h @@ -0,0 +1,93 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright contributors to the Foundry project +#ifndef FOUNDRY_VMM_CLEANUP_H +#define FOUNDRY_VMM_CLEANUP_H + +// Pure, CUDA-free policy for retryable VMM cleanup. Production supplies CUDA +// driver calls; host tests inject deterministic failures. A successful phase is +// never executed twice, while failed phases remain pending for the next retry. + +namespace foundry::vmm { + +enum class CleanupPhase { + kUnmap, + kRelease, + kAddressFree, +}; + +struct CleanupState { + bool started = false; + bool unmapped = false; + bool released = false; + bool address_freed = false; +}; + +inline bool cleanup_complete(const CleanupState& state, bool address_free_required) { + return state.unmapped && state.released && (!address_free_required || state.address_freed); +} + +template +Status run_cleanup(CleanupState& state, Status success, bool address_free_required, Unmap unmap, + Release release, AddressFree address_free, Log log_failure) { + state.started = true; + Status first_failure = success; + + if (!state.unmapped) { + Status status = unmap(); + if (status == success) { + state.unmapped = true; + } else { + log_failure(CleanupPhase::kUnmap, status); + first_failure = status; + } + } + + if (!state.released) { + Status status = release(); + if (status == success) { + state.released = true; + } else { + log_failure(CleanupPhase::kRelease, status); + if (first_failure == success) { + first_failure = status; + } + } + } + + if (address_free_required && !state.address_freed) { + Status status = address_free(); + if (status == success) { + state.address_freed = true; + } else { + log_failure(CleanupPhase::kAddressFree, status); + if (first_failure == success) { + first_failure = status; + } + } + } + + return first_failure; +} + +// Retry every durable orphan cleanup before admitting new import work. Completed +// entries are erased; failed entries retain their resource and phase state. +template +Status drain_pending_cleanups(Container& pending, Status success, Retry retry, Complete complete) { + Status first_failure = success; + for (auto it = pending.begin(); it != pending.end();) { + Status status = retry(it->second); + if (complete(it->second)) { + it = pending.erase(it); + continue; + } + if (first_failure == success) { + first_failure = status; + } + ++it; + } + return first_failure; +} + +} // namespace foundry::vmm + +#endif // FOUNDRY_VMM_CLEANUP_H diff --git a/include/foundry/vmm_ipc_chunk_key.h b/include/foundry/vmm_ipc_chunk_key.h new file mode 100644 index 00000000..ccbc73da --- /dev/null +++ b/include/foundry/vmm_ipc_chunk_key.h @@ -0,0 +1,72 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright contributors to the Foundry project +#ifndef FOUNDRY_VMM_IPC_CHUNK_KEY_H +#define FOUNDRY_VMM_IPC_CHUNK_KEY_H + +#include +#include +#include + +// Identity of an imported peer-chunk mapping. All four fields are +// identity-bearing. In particular the generation distinguishes a chunk that was +// freed and recreated at the SAME base within one exporter process: without it, +// a stale first-generation mapping would be wrongly reused for a recreated +// chunk (reading orphaned physical memory). This is a pure, CUDA-free value +// type so the equality/ordering/hash contract can be unit-tested on a CPU-only +// host; the production hook keys vmm_ipc_chunk_mappings on it (csrc/hook.cpp). + +namespace foundry::vmm_ipc { + +struct ChunkKey { + // pid_t in the hook; widened to long here to stay CUDA/OS-header free. + long exporter_pid; + unsigned long long token; + unsigned long long chunk_base; + unsigned long long generation; + + // Canonical, generation-inclusive field order for equality and the + // strict-weak ordering the std::map uses. + std::tuple as_tuple() const { + return std::make_tuple(exporter_pid, token, chunk_base, generation); + } + + friend bool operator==(const ChunkKey& a, const ChunkKey& b) { + return a.as_tuple() == b.as_tuple(); + } + friend bool operator!=(const ChunkKey& a, const ChunkKey& b) { return !(a == b); } + friend bool operator<(const ChunkKey& a, const ChunkKey& b) { + return a.as_tuple() < b.as_tuple(); + } +}; + +inline ChunkKey make_chunk_key(long exporter_pid, unsigned long long token, + unsigned long long chunk_base, unsigned long long generation) { + return ChunkKey{exporter_pid, token, chunk_base, generation}; +} + +// Hash over all four fields, including the generation, so the equality/hash +// contract is consistent for hashed containers too. +struct ChunkKeyHash { + std::size_t operator()(const ChunkKey& k) const { + std::size_t h = std::hash{}(k.exporter_pid); + auto mix = [&h](unsigned long long v) { + h ^= std::hash{}(v) + 0x9e3779b97f4a7c15ULL + (h << 6) + (h >> 2); + }; + mix(k.token); + mix(k.chunk_base); + mix(k.generation); + return h; + } +}; + +// The identity a generation-LESS key would use. Provided so tests can prove +// that dropping the generation collapses distinct recreated chunks onto one +// key (the exact bug the generation field prevents). Not used in production. +inline std::tuple chunk_identity_without_generation( + const ChunkKey& k) { + return std::make_tuple(k.exporter_pid, k.token, k.chunk_base); +} + +} // namespace foundry::vmm_ipc + +#endif // FOUNDRY_VMM_IPC_CHUNK_KEY_H diff --git a/include/foundry/vmm_ipc_ownership.h b/include/foundry/vmm_ipc_ownership.h new file mode 100644 index 00000000..f8110550 --- /dev/null +++ b/include/foundry/vmm_ipc_ownership.h @@ -0,0 +1,56 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright contributors to the Foundry project +#ifndef FOUNDRY_VMM_IPC_OWNERSHIP_H +#define FOUNDRY_VMM_IPC_OWNERSHIP_H + +// Pure, hardware-independent VMM-IPC process-ownership policy. No CUDA or +// threading dependencies so the decisions can be unit-tested on a CPU-only +// host. The production hook wires these into an atomic owner pid and the IPC +// entry points; see csrc/hook.cpp. + +namespace foundry::vmm_ipc { + +// Fixed, async-signal-safe rejection message for a post-claim fork child. It is +// written verbatim via write(2) from the fork child: between fork() and exec() +// in a multithreaded process only async-signal-safe operations are allowed, so +// fprintf/snprintf/malloc/locks must not be used on that path. The message is a +// compile-time constant (length via sizeof) and includes the "use spawn/exec" +// remediation. +inline constexpr char kForkChildRejectMessage[] = + "[HOOK] ERROR: VMM-IPC is unsupported in this fork child because the parent " + "process already started VMM-IPC; use spawn/exec instead of fork\n"; + +// Sentinel for "no process has claimed VMM-IPC ownership yet". +inline constexpr long kNoOwnerPid = -1; + +enum class OwnershipDecision { + kProceed, // caller may claim/use VMM-IPC in this process + kRejectForkChild // caller is a post-claim fork child; reject before touching + // inherited mutexes/maps/handles +}; + +// Claim-or-match decision given the currently recorded owner pid (kNoOwnerPid +// when unclaimed) and the caller's pid. Unclaimed or same-owner proceeds (the +// first user becomes the owner; the owner may reuse IPC freely). A different +// owner means the caller is a fork child created after the owner claimed, which +// must be rejected. Fork *before* the first claim is therefore always allowed: +// the child observes kNoOwnerPid and proceeds to become its own owner. +inline OwnershipDecision decide_ipc_ownership(long owner_pid, long caller_pid) { + if (owner_pid == kNoOwnerPid || owner_pid == caller_pid) { + return OwnershipDecision::kProceed; + } + return OwnershipDecision::kRejectForkChild; +} + +// Whether the caller may touch owner-only VMM-IPC state (the parked export-fd +// registry). Only the exact owning process may: an unclaimed process +// (kNoOwnerPid) has never exported anything, and a fork child (owner mismatch) +// must not lock or mutate the inherited concurrent map. Both return false so +// the invalidation path returns before any map access. +inline bool owns_ipc_state(long owner_pid, long caller_pid) { + return owner_pid == caller_pid && owner_pid != kNoOwnerPid; +} + +} // namespace foundry::vmm_ipc + +#endif // FOUNDRY_VMM_IPC_OWNERSHIP_H diff --git a/include/foundry/vmm_ipc_reservation.h b/include/foundry/vmm_ipc_reservation.h new file mode 100644 index 00000000..971cdb8c --- /dev/null +++ b/include/foundry/vmm_ipc_reservation.h @@ -0,0 +1,95 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright contributors to the Foundry project +#ifndef FOUNDRY_VMM_IPC_RESERVATION_H +#define FOUNDRY_VMM_IPC_RESERVATION_H + +#include +#include + +namespace foundry::vmm_ipc { + +// Fail-closed policy for the dedicated peer-import VA zone. +// +// After a peer chunk/allocation is imported, the mapping is placed in a +// dedicated VA zone by reserving at `expected` (the zone hint). The reservation +// is only usable if the driver honored the hint exactly: a reservation landing +// at a different VA can collide with a future deterministic/NVSHMEM reserve +// location, silently breaking SAVE/LOAD address determinism. This helper is the +// single authority for that decision so it can be tested without CUDA hardware. +// +// - `result` is the return code of the dedicated-zone `cuMemAddressReserve`. +// - `actual` is the VA it returned; `expected` is the requested zone hint. +// - `size`/`handle` are forwarded to the cleanup callbacks unchanged. +// - `free_reservation(actual, size) -> CUresult` frees the reservation. +// - `release_handle(handle) -> CUresult` releases the imported allocation +// handle. Both callbacks return the driver status so the owner process can +// log a cleanup failure; the values are otherwise not representable in this +// function's single return (the import already failed), so they do not change +// the fail-closed result. +// +// Contract: +// 1. driver error (`result != CUDA_SUCCESS`): return `result`, call neither +// callback (the caller still owns and must release the imported handle); +// 2. success at the requested VA (`actual == expected`): return +// `CUDA_SUCCESS`, call neither callback (the caller proceeds to map); +// 3. success at a different VA (`actual != expected`): free the reservation +// and release the handle exactly once each, in that order, then return +// `CUDA_ERROR_OUT_OF_MEMORY`. A failing address-free never skips the handle +// release (no leak) and neither callback is retried (no double release); +// a failing callback keeps the fail-closed `CUDA_ERROR_OUT_OF_MEMORY`. +// Fallback-VA cleanup policy (runs BEFORE the dedicated-zone reservation). +// +// The peer import first reserves at the exporter's own VA hint. If the driver +// ignores that hint it picks its own fallback VA, which we must free before +// re-reserving inside the dedicated import zone. If that free itself fails we +// cannot safely proceed: the fallback VA may straddle the deterministic/NVSHMEM +// reserve the zone exists to avoid, and the reservation is now dangling. So +// fail closed - release the imported handle exactly once and return +// CUDA_ERROR_OUT_OF_MEMORY - and the caller must return immediately without +// attempting the dedicated-zone mapping. +// +// - `free_status` is the return code of freeing the driver's fallback VA. +// - `release_handle(handle) -> CUresult` releases the imported allocation +// handle; it is invoked only on the failure branch, so the caller must NOT +// release the handle again on this path (no double release). +// +// Contract: +// 1. clean free (`free_status == CUDA_SUCCESS`): return `CUDA_SUCCESS`, do not +// release (the caller continues to the dedicated-zone reservation); +// 2. failed free: release the handle exactly once, return +// `CUDA_ERROR_OUT_OF_MEMORY` (fail closed, no zone mapping). +template +inline CUresult require_fallback_reservation_freed(CUresult free_status, + CUmemGenericAllocationHandle handle, + ReleaseHandle release_handle) { + if (free_status == CUDA_SUCCESS) { + return CUDA_SUCCESS; + } + release_handle(handle); + return CUDA_ERROR_OUT_OF_MEMORY; +} + +template +inline CUresult require_exact_ipc_reservation(CUresult result, CUdeviceptr actual, + CUdeviceptr expected, size_t size, + CUmemGenericAllocationHandle handle, + FreeReservation free_reservation, + ReleaseHandle release_handle) { + if (result != CUDA_SUCCESS) { + return result; + } + if (actual == expected) { + return CUDA_SUCCESS; + } + // Wrong VA: fail closed. Run both cleanups exactly once regardless of the + // first's status so a failed free cannot leak the imported handle. + CUresult free_status = free_reservation(actual, size); + CUresult release_status = release_handle(handle); + (void)free_status; + (void)release_status; + return CUDA_ERROR_OUT_OF_MEMORY; +} + +} // namespace foundry::vmm_ipc + +#endif // FOUNDRY_VMM_IPC_RESERVATION_H diff --git a/include/hook.h b/include/hook.h index 4d6e132b..493f755d 100644 --- a/include/hook.h +++ b/include/hook.h @@ -28,7 +28,7 @@ void set_allocation_region(void* base, size_t size); void stop_allocation_region(); void resume_allocation_region(); bool preallocate_region(size_t size); -void free_preallocated_region(); +bool free_preallocated_region(); size_t get_current_alloc_offset(); void set_current_alloc_offset(size_t offset); void set_pack_fatbins_on_exit(bool enabled); diff --git a/pyproject.toml b/pyproject.toml index b70e5572..08d647c2 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -23,9 +23,13 @@ dev = [ "pytest", ] +[project.entry-points."sglang.srt.plugins"] +foundry = "foundry_sglang_plugin:activate" + [tool.setuptools] packages = [ "foundry", + "foundry_sglang_plugin", "foundry.integration", "foundry.integration.vllm", "foundry.integration.sglang", diff --git a/python/foundry/__init__.py b/python/foundry/__init__.py index bb88e00e..b7f88225 100644 --- a/python/foundry/__init__.py +++ b/python/foundry/__init__.py @@ -1,5 +1,9 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the Foundry project +# ruff: noqa: I001 +# The ops wildcard must come FIRST: .graph's CUDAGraph (and the +# allocation_region helpers) intentionally override the raw pybind names. +from .ops import * from .allocation_region import ( allocation_region, free_preallocated_region, @@ -14,7 +18,6 @@ graph, save_graph_manifest, ) -from .ops import * # Re-exports. Listed here so ruff's --fix doesn't strip them as unused. # (We also configure per-file-ignores for F401 in pyproject.toml as a backstop.) diff --git a/python/foundry/allocation_region.py b/python/foundry/allocation_region.py index 5336c7b7..aff42efb 100644 --- a/python/foundry/allocation_region.py +++ b/python/foundry/allocation_region.py @@ -1,6 +1,7 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the Foundry project import re +import warnings from contextlib import contextmanager from .ops import ( @@ -92,6 +93,7 @@ def allocation_region(base: int | str, size: int | str, prealloc_size: int | str set_allocation_region(base_addr, size_bytes) + body_exception: BaseException | None = None try: if has_prealloc: prealloc_bytes = parse_size(prealloc_size) @@ -99,7 +101,28 @@ def allocation_region(base: int | str, size: int | str, prealloc_size: int | str if not success: raise RuntimeError(f"Failed to preallocate {prealloc_bytes} bytes") yield + except BaseException as exc: + body_exception = exc + raise finally: - if has_prealloc: - free_preallocated_region() - stop_allocation_region() + cleanup_succeeded = True + cleanup_exception: BaseException | None = None + try: + if has_prealloc: + try: + cleanup_succeeded = free_preallocated_region() + except BaseException as exc: + cleanup_succeeded = False + cleanup_exception = exc + finally: + stop_allocation_region() + + if not cleanup_succeeded: + message = "Failed to free preallocated region after stopping allocation region" + if cleanup_exception is not None: + message = f"{message}: {cleanup_exception!r}" + if body_exception is not None: + body_exception.add_note(message) + warnings.warn(message, RuntimeWarning, stacklevel=2) + else: + raise RuntimeError(message) from cleanup_exception diff --git a/python/foundry/graph.py b/python/foundry/graph.py index 427d1599..6a3b4c1c 100644 --- a/python/foundry/graph.py +++ b/python/foundry/graph.py @@ -135,7 +135,10 @@ def __init__( if self.__class__.default_capture_stream is None: self.__class__.default_capture_stream = torch.cuda.Stream() - self.pool = () if pool is None else (pool,) + # The C++ binding accepts a pool tuple or explicit None, but the + # argument itself is required. Keep one positional element so the + # default path passes None instead of omitting `pool` entirely. + self.pool = (pool,) self.capture_stream = ( stream if stream is not None else self.__class__.default_capture_stream ) diff --git a/python/foundry/integration/deepep_transport.py b/python/foundry/integration/deepep_transport.py new file mode 100644 index 00000000..a7145aae --- /dev/null +++ b/python/foundry/integration/deepep_transport.py @@ -0,0 +1,130 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the Foundry project +"""Shared DeepEP transport policy for Foundry engine integrations.""" + +from __future__ import annotations + +import enum +from typing import Any + + +class DeepEPTransport(str, enum.Enum): + FABRIC = "fabric" + NVL_IPC = "nvl_ipc" + + +def parse_deepep_transport(raw: str | None) -> DeepEPTransport: + value = DeepEPTransport.FABRIC.value if raw is None else raw + try: + return DeepEPTransport(value) + except ValueError as exc: + choices = ", ".join(item.value for item in DeepEPTransport) + raise ValueError(f"deepep_transport must be one of: {choices}; got {value!r}") from exc + + +def configure_deepep_buffer( + transport: DeepEPTransport, + *, + num_nvl_bytes: object, + kwargs: dict[str, Any], +) -> tuple[int, dict[str, Any]]: + validated_num_nvl_bytes = validate_num_nvl_bytes_value( + num_nvl_bytes, + transport=transport, + ) + configured = dict(kwargs) + if transport is DeepEPTransport.FABRIC: + configured["use_fabric"] = True + return 0, configured + configured["use_fabric"] = False + return validated_num_nvl_bytes, configured + + +def validate_num_nvl_bytes_value( + value: object, + *, + transport: DeepEPTransport, + field_name: str = "num_nvl_bytes", +) -> int: + if transport is DeepEPTransport.NVL_IPC: + if type(value) is not int or value <= 0: + raise RuntimeError( + f"Foundry DeepEP {field_name} must be an exact positive integer " + f"for {transport.value}; got {value!r}" + ) + elif type(value) is not int or value < 0: + raise RuntimeError( + f"Foundry DeepEP {field_name} must be an exact non-negative integer " + f"for {transport.value}; got {value!r}" + ) + return value + + +def validate_archived_deepep_profile( + data: dict[str, Any], + *, + path: str, +) -> tuple[DeepEPTransport, int]: + if "deepep_transport" in data: + raw_transport = data["deepep_transport"] + if type(raw_transport) is not str: + raise RuntimeError( + f"Foundry invalid warmup state at {path}: field 'deepep_transport' " + f"must be a string; got {raw_transport!r}" + ) + else: + raw_transport = DeepEPTransport.FABRIC.value + + try: + transport = parse_deepep_transport(raw_transport) + except ValueError as exc: + raise RuntimeError( + f"Foundry invalid warmup state at {path}: field 'deepep_transport': {exc}" + ) from exc + + raw_num_nvl_bytes = data.get("deepep_num_nvl_bytes", 0) + try: + num_nvl_bytes = validate_num_nvl_bytes_value( + raw_num_nvl_bytes, + transport=transport, + field_name="deepep_num_nvl_bytes", + ) + except RuntimeError as exc: + raise RuntimeError(f"Foundry invalid warmup state at {path}: {exc}") from exc + return transport, num_nvl_bytes + + +def validate_transport_match( + *, + archived: str | None, + configured: DeepEPTransport, +) -> None: + archived_transport = parse_deepep_transport(archived) + if archived_transport is not configured: + raise RuntimeError( + "Foundry DeepEP transport mismatch: " + f"archived={archived_transport.value}, configured={configured.value}" + ) + + +def validate_num_nvl_bytes_match( + *, + archived: int, + effective: int, + transport: DeepEPTransport, +) -> None: + validate_num_nvl_bytes_value( + archived, + transport=transport, + field_name="archived num_nvl_bytes", + ) + validate_num_nvl_bytes_value( + effective, + transport=transport, + field_name="effective num_nvl_bytes", + ) + if archived != effective: + raise RuntimeError( + "Foundry DeepEP num_nvl_bytes mismatch for " + f"{transport.value}: archived={archived}, effective={effective}" + ) diff --git a/python/foundry/integration/sglang/backend.py b/python/foundry/integration/sglang/backend.py new file mode 100644 index 00000000..64483270 --- /dev/null +++ b/python/foundry/integration/sglang/backend.py @@ -0,0 +1,252 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the Foundry project +"""Foundry SAVE/LOAD override of SGLang's FullCudaGraphBackend. + +The pinned upstream (a25164b) drives every decode shape through the pluggable +``FullCudaGraphBackend`` under the runner's real ``DecodeCudaGraphRunner.capture`` +lifecycle: ``capture()`` warms up, enters ``graph_capture()``, fills the shared +seq-len buffers, and — inside ``backend.capture_session(stream)`` — walks +``reversed(capture_bs)`` calling ``capture_prepare`` / ``ForwardContext`` / +``init_forward_metadata_out_graph(..., in_capture=True)`` and then +``backend.capture_one(shape_key, forward_fn, ...)`` per shape. + +Foundry keeps that exact lifecycle in BOTH SAVE and LOAD (this is the crucial +fidelity point: the allocator trajectory and attention-metadata wrappers must be +built identically on both sides). Only the backend's capture/materialize +mechanics change: + +* SAVE — ``capture_one`` captures the single forward into a Foundry ``CUDAGraph`` + (skipping upstream's two warmup forwards, whose activation allocations are not + LOAD-reproducible) and serializes it. +* LOAD — ``capture_session`` preallocates, scans the archive, and issues exactly + one ``start_graph_builds(paths, pool=self._pool)``; ``capture_one`` then + finishes one archived graph per requested shape (no forward, no native + capture, no warmup), validating shape/order against the archive. + +In every other mode the subclass delegates to the stock upstream implementation. +""" + +from __future__ import annotations + +import logging +import os +from contextlib import contextmanager + +from foundry.integration.sglang import runtime as rt +from foundry.integration.sglang.config import ( + CUDAGraphExtensionMode, + get_graph_extension_mode, +) + +logger = logging.getLogger(__name__) + + +def build_foundry_full_backend_class(): + """Return a ``FullCudaGraphBackend`` subclass that captures into Foundry + graphs on SAVE, restores them on LOAD, and delegates otherwise. + + Built lazily so SGLang (and the pinned ``FullCudaGraphBackend``) stay + optional for core-only Foundry installs. The returned class is registered + as a ``HookType.REPLACE`` on + ``...full_cuda_graph_backend.FullCudaGraphBackend``; the hook registry's + ``_propagate_patch`` rebinds every ``from ... import FullCudaGraphBackend`` + reference (notably ``runner_backend.utils.resolve_decode_backend``) so the + decode runner instantiates this subclass. + """ + from sglang.srt.distributed.device_communicators.pynccl_allocator import ( + set_graph_pool_id, + ) + from sglang.srt.model_executor.runner_backend.full_cuda_graph_backend import ( + FullCudaGraphBackend, + ) + from sglang.srt.model_executor.runner_utils.pool import ( + get_or_create_global_graph_memory_pool, + ) + + from foundry.graph import CUDAGraph as FoundryCUDAGraph + from foundry.graph import graph as foundry_graph_ctx + from foundry.integration.sglang.graph_ops import ( + _require_dense_shape_key, + _scan_graph_files, + _unpack_output, + save_graph, + ) + + class FoundryFullCudaGraphBackend(FullCudaGraphBackend): + def __init__(self, cuda_graph_runner, *, enable_memory_saver: bool = False) -> None: + super().__init__(cuda_graph_runner, enable_memory_saver=enable_memory_saver) + # LOAD bookkeeping on top of the inherited shared-pool state + # (``_pool`` / ``_graphs`` / ``_outputs`` / ``_capture_stream``). + self._pending = None + self._graph_files = [] + self._load_index = 0 + self._closed = False + + @contextmanager + def capture_session(self, stream): + """Establish the shared pool + stream, drive LOAD prepare, then yield. + + Mirrors ``FullCudaGraphBackend.capture_session`` (lazy pool creation + + ``set_graph_pool_id``) so SAVE keeps its exact behaviour, and adds the + proven PR4 LOAD prepare/verify around the same upstream per-shape loop. + SAVE archive finalization stays in the ``capture`` AROUND hook. + """ + if self._closed: + raise RuntimeError("Foundry SGLang backend does not support CUDA graph recapture") + if self._pool is None: + self._pool = get_or_create_global_graph_memory_pool(self._device_module) + set_graph_pool_id(self._pool) + self._capture_stream = stream + + mode = get_graph_extension_mode() + if mode == CUDAGraphExtensionMode.LOAD: + self._prepare_load() + + completed = False + try: + yield + completed = True + finally: + self._capture_stream = None + if completed and mode == CUDAGraphExtensionMode.LOAD: + if self._load_index != len(self._graph_files): + raise RuntimeError( + f"Foundry loaded {self._load_index}/{len(self._graph_files)} " + "SGLang graphs; the upstream capture loop requested a different " + "shape set than the archive holds" + ) + rt.log_alloc_offset("after_load_all_graphs") + logger.info("[Foundry] Loaded %d SGLang graphs", self._load_index) + + def _prepare_load(self) -> None: + """Preallocate, scan the archive, and start the sole build batch. + + SAVE captured into the shared graph memory pool; ``capture_session`` + establishes and publishes it as ``self._pool`` before this runs. + Issue exactly ONE ``start_graph_builds`` (manifest template linking + + background builds); ``capture_one`` finishes each graph per index so + every upstream between-graph wrapper allocation interleaves with the + graph it preceded on SAVE, keeping the LOAD cursor on SAVE's + trajectory. + """ + from foundry.integration.sglang.config import get_config + + cfg = get_config() + if cfg is None or cfg.workspace_dir is None: + raise RuntimeError("Foundry SGLang graph extension is not initialized") + + rt.log_alloc_offset("before_preallocate") + rt.preallocate_for_load_mode() + rt.log_alloc_offset("after_preallocate") + + self._graph_files = _scan_graph_files(cfg.workspace_dir) + if not self._graph_files: + raise RuntimeError(f"No Foundry SGLang graph files found in {cfg.workspace_dir}") + if self._pool is None: + raise RuntimeError("Foundry SGLang graph LOAD requires an initialized pool") + paths = [ + os.path.join(cfg.workspace_dir, filename) + for _index, filename, _meta in self._graph_files + ] + self._pending = FoundryCUDAGraph.start_graph_builds( + paths, pool=self._pool, num_threads=4 + ) + self._load_index = 0 + logger.info("[Foundry] Started %d SGLang graph builds", len(paths)) + + def capture_one( + self, + shape_key, + forward_fn, + capture_inputs=None, + post_warmup_hook=None, + ): + mode = get_graph_extension_mode() + if mode == CUDAGraphExtensionMode.NONE: + return super().capture_one(shape_key, forward_fn, capture_inputs, post_warmup_hook) + if mode == CUDAGraphExtensionMode.LOAD: + self._load_one(shape_key) + return None + + # SAVE. + # Upstream runs two warmup forwards before capture, calling + # ``post_warmup_hook`` after each (``full_cuda_graph_backend.py`` + # capture_one). Foundry deliberately skips BOTH warmups: their + # activation allocations are not reproducible on LOAD and would + # drift the VMM cursor away from each saved ``start_base_addr``. + # JIT / autotune still runs inside the captured graph (the single + # forward below) and is recorded as replayable alloc events. + # + # ``post_warmup_hook`` exists solely to undo state the warmup + # forwards mutate (``base_attn_backend.on_after_cuda_graph_warmup`` + # docstring). With no warmup there is nothing to undo, and the + # pinned FlashInfer backend inherits the base no-op hook anyway, so + # we intentionally do NOT invoke it — invoking it blindly could + # reset state that was never dirtied. + # + # Upstream runs a device synchronize + TP barrier before each warmup + # forward. Since Foundry skips the warmups, restore that lockstep + # immediately before entering the graph context so every rank reaches + # capture together (SGLang's ``tp_group.barrier`` uses its CPU group) + # and the captured symmetric collective sees a settled buffer/group. + self._device_module.synchronize() + self._tp_group.barrier() + graph = FoundryCUDAGraph() + with foundry_graph_ctx(graph, pool=self._pool, stream=self._capture_stream): + output = forward_fn() + + # Mirror upstream's contract: populate the backend maps (the runner + # reads them) and return None. + self._graphs[shape_key] = graph + self._outputs[shape_key] = output + save_graph(graph, output, shape_key) + return None + + def _load_one(self, shape_key) -> None: + """Finish one archived graph for the upstream-requested shape. + + Validates the requested ShapeKey against the archive's capture order + (index ``i`` was captured before ``i+1``, i.e. descending bs), fails + closed on order/extra mismatches, and populates the backend maps the + pluggable runner replays from — no forward, native capture, or warmup. + """ + if self._pending is None: + raise RuntimeError("Foundry SGLang graph builds were not started") + if self._load_index >= len(self._graph_files): + raise RuntimeError("SGLang requested more graph shapes than were saved") + + size = _require_dense_shape_key(shape_key) + _index, filename, meta = self._graph_files[self._load_index] + archived_size = meta["shape_key"].size + if archived_size != size: + raise RuntimeError( + f"SGLang graph order mismatch: requested {size}, archive has " + f"{archived_size} ({filename})" + ) + + graph, tensors = FoundryCUDAGraph.finish_one_graph_load(self._pending, self._load_index) + self._load_index += 1 + self._graphs[shape_key] = graph + self._outputs[shape_key] = _unpack_output(tensors) + logger.info("[Foundry] Loaded SGLang graph %s", filename) + + def cleanup(self) -> None: + # In stock (NONE) operation the subclass is transparent: delegate to + # the upstream cleanup and do NOT latch the recapture guard, so the + # runner's legitimate hidden-mode recapture keeps working. + if get_graph_extension_mode() == CUDAGraphExtensionMode.NONE: + super().cleanup() + return + # Active SAVE/LOAD: a completed capture/load latches ``_closed`` so a + # later capture_session (e.g. hidden-mode recapture) fails closed + # rather than re-capturing Foundry graphs into a torn-down pool. + if self._graphs: + self._closed = True + self._graphs.clear() + self._outputs.clear() + self._pool = None + self._pending = None + self._graph_files = [] + self._load_index = 0 + + return FoundryFullCudaGraphBackend diff --git a/python/foundry/integration/sglang/graph_ops.py b/python/foundry/integration/sglang/graph_ops.py index 18ab1f0e..66a7a227 100644 --- a/python/foundry/integration/sglang/graph_ops.py +++ b/python/foundry/integration/sglang/graph_ops.py @@ -7,41 +7,90 @@ import logging import os import re -import time from typing import Any import torch import foundry as foundry_pkg from foundry import ops as cge -from foundry.graph import CUDAGraph as FoundryCUDAGraph -from foundry.graph import graph as foundry_graph_ctx -from foundry.integration.sglang.config import ( - CUDAGraphExtensionMode, - get_config, - get_graph_extension_mode, -) +from foundry.integration.sglang.config import get_config from foundry.integration.sglang.runtime import get_state logger = logging.getLogger(__name__) -_pending_graph_builds: tuple[Any, list[tuple[int, str, dict[str, Any]]]] | None = None _GRAPH_FILENAME_RE = re.compile(r"^graph_(?P\d+)_FULL_t(?P\d+)_r\d+_UX_pcN\.json$") +# LogitsProcessorOutput fields the dense full-decode archive does not support. +# The dense profile captures only next_token_logits; any populated +# speculative / hidden-state / diffusion field means the runtime replay would +# read a value the archive cannot reconstruct, so SAVE fails closed. +_UNSUPPORTED_DECODE_OUTPUT_FIELDS = ("hidden_states", "full_logits", "customized_info") -def _batch_size_from_key(key: Any) -> int: - if isinstance(key, int): - return key - key_str = str(key) - for part in reversed(key_str.split("_")): - if part.isdigit(): - return int(part) - raise ValueError(f"Cannot derive batch size from SGLang CUDA graph key: {key!r}") + +_SHAPE_KEY_FIELDS = ("size", "stream_idx", "variant_label") + + +def _require_dense_shape_key(key: Any) -> int: + """Return the capture ``size`` for a dense-profile ShapeKey, failing closed. + + The pinned upstream keys graphs by ``ShapeKey(size, stream_idx, + variant_label)`` (``runner/shape_key.py``). Foundry's validated dense TP=2 + full-decode profile is single-stream and non-LoRA. This fails closed unless + the object explicitly exposes every current ``ShapeKey`` field, so a raw int + or an object built against an older/partial schema can never be silently + collapsed into a ``t{size}`` filename that drops a populated dimension: + + * missing any of ``size`` / ``stream_idx`` / ``variant_label`` -> reject + (rejects raw ints and incomplete objects); + * populated ``stream_idx`` (pdmux) or ``variant_label`` (LoRA) -> reject; + * ``size`` must be an exact ``int`` (``bool`` is excluded via ``type(...)``). + """ + missing = [field for field in _SHAPE_KEY_FIELDS if not hasattr(key, field)] + if missing: + raise ValueError( + "Foundry SGLang graph archive requires a ShapeKey exposing " + f"{_SHAPE_KEY_FIELDS}; got {key!r} missing {missing}" + ) + if key.stream_idx is not None or key.variant_label is not None: + raise ValueError( + "Foundry SGLang graph archive supports only the dense single-stream " + "profile; got unsupported " + f"ShapeKey(stream_idx={key.stream_idx!r}, variant_label={key.variant_label!r})" + ) + size = key.size + if type(size) is not int: + raise ValueError(f"Foundry SGLang ShapeKey size must be an int; got {size!r}") + return size def _graph_filename(index: int, key: Any) -> str: - batch_size = _batch_size_from_key(key) - return f"graph_{index}_FULL_t{batch_size}_r{batch_size}_UX_pcN.json" + size = _require_dense_shape_key(key) + return f"graph_{index}_FULL_t{size}_r{size}_UX_pcN.json" + + +def _shape_key_from_filename(name: str): + match = _GRAPH_FILENAME_RE.match(name) + if not match: + raise ValueError(f"Not a Foundry SGLang graph filename: {name!r}") + # Lazy import: SGLang is an optional integration, unavailable in core-only + # Foundry installs; the ShapeKey type lives with the runner it keys. + from sglang.srt.model_executor.runner.shape_key import ShapeKey + + return ShapeKey(size=int(match.group("bs"))) + + +def _assert_dense_decode_output(output: Any) -> None: + populated = [ + field + for field in _UNSUPPORTED_DECODE_OUTPUT_FIELDS + if getattr(output, field, None) is not None + ] + if populated: + raise ValueError( + "Foundry SGLang dense full-decode graph archive supports only " + "next_token_logits; refusing to drop populated unsupported " + f"LogitsProcessorOutput fields: {populated}" + ) def _pack_output(output: Any) -> torch.Tensor: @@ -50,6 +99,7 @@ def _pack_output(output: Any) -> torch.Tensor: if isinstance(output, LogitsProcessorOutput): if output.next_token_logits is None: raise TypeError("SGLang decode CUDA graph output has no next_token_logits") + _assert_dense_decode_output(output) return output.next_token_logits if isinstance(output, torch.Tensor): @@ -76,28 +126,13 @@ def _scan_graph_files(workspace_dir: str) -> list[tuple[int, str, dict[str, Any] continue meta = { "index": int(match.group("index")), - "key": int(match.group("bs")), + "shape_key": _shape_key_from_filename(filename), } graph_files.append((int(meta["index"]), filename, meta)) graph_files.sort(key=lambda x: x[0]) return graph_files -def create_device_graph(): - mode = get_graph_extension_mode() - if mode == CUDAGraphExtensionMode.SAVE: - return FoundryCUDAGraph() - return torch.cuda.CUDAGraph() - - -def capture_graph(graph, pool, stream, run_once_fn): - mode = get_graph_extension_mode() - if mode == CUDAGraphExtensionMode.SAVE: - with foundry_graph_ctx(graph, pool=pool, stream=stream): - return run_once_fn() - return None - - def save_graph(graph, output: Any, key: Any) -> None: cfg = get_config() state = get_state() @@ -126,200 +161,3 @@ def pack_fatbins() -> None: return cge.pack_fatbins_to_folder(cfg.workspace_dir) cge.set_pack_fatbins_on_exit(False) - - -def start_graph_builds() -> None: - global _pending_graph_builds - cfg = get_config() - if cfg is None or cfg.workspace_dir is None or cfg.mode != CUDAGraphExtensionMode.LOAD: - return - - graph_files = _scan_graph_files(cfg.workspace_dir) - if not graph_files: - raise RuntimeError(f"No Foundry SGLang graph files found in {cfg.workspace_dir}") - - paths = [os.path.join(cfg.workspace_dir, filename) for _, filename, _ in graph_files] - t0 = time.perf_counter() - pending = FoundryCUDAGraph.start_graph_builds(paths, num_threads=4) - _pending_graph_builds = (pending, graph_files) - logger.info( - "[Foundry] Started SGLang graph builds for %d graphs in %.3fs", - len(paths), - time.perf_counter() - t0, - ) - - -def preload_all_graphs() -> None: - global _pending_graph_builds - cfg = get_config() - state = get_state() - if cfg is None or state is None or cfg.workspace_dir is None: - raise RuntimeError("Foundry SGLang graph extension is not initialized") - - if _pending_graph_builds is None: - start_graph_builds() - assert _pending_graph_builds is not None - - cge.init_nvshmem_for_loaded_modules() - - pending, graph_files = _pending_graph_builds - _pending_graph_builds = None - - t0 = time.perf_counter() - results = FoundryCUDAGraph.finish_graph_loads(pending) - logger.info( - "[Foundry] Finished SGLang graph loads for %d graphs in %.3fs", - len(results), - time.perf_counter() - t0, - ) - - for i, (_index, _filename, meta) in enumerate(graph_files): - graph, tensors = results[i] - state.loaded_graphs[meta["key"]] = (graph, _unpack_output(tensors)) - - -def bootstrap_deepep_buffer(cuda_graph_runner) -> bool: - """Force the singleton DeepEP ``Buffer`` (NVSHMEM runtime + symmetric heap) - to be created BEFORE the cuda-graph capture loop. - - sglang creates the DeepEP buffer lazily on the first MoE dispatch — normally - during the two pre-capture warmup forwards. Foundry suppresses those warmups - (for allocation determinism), which would push buffer creation into the - captured forward, where ``deep_ep_cpp.Buffer(...)`` aborts with - ``operation not permitted when stream is capturing``. - - Triggering it here (outside any stream capture) creates only the NVSHMEM - runtime + symmetric heap — no model activations — so it stays symmetric - across SAVE and LOAD and lands at the same VMM offset on both. The buffer is - a process-wide singleton (``DeepEPBuffer._buffer``), so one creation per rank - is enough. It is a collective over the EP group, so every rank must reach - this point together — which they do, since ``capture`` runs on all ranks. - - Returns True if a buffer was (or already is) created, False if DeepEP is off. - """ - try: - from sglang.srt.layers.moe.utils import get_moe_a2a_backend - - if not get_moe_a2a_backend().is_deepep(): - return False - except Exception: - return False - - from sglang.srt.layers.moe.token_dispatcher.deepep import ( - DeepEPBuffer, - DeepEPDispatcher, - ) - - if DeepEPBuffer._buffer is not None: - return True - - model = cuda_graph_runner.model_runner.model - for module in model.modules(): - dispatcher = getattr(module, "dispatcher", None) - if dispatcher is None: - continue - # ``module.dispatcher`` is normally a MaybeTboDeepEPDispatcher wrapper - # whose ``_inners`` hold the real DeepEPDispatcher(s); unwrap it. (Also - # handle a bare DeepEPDispatcher for safety.) - candidates = [dispatcher, *getattr(dispatcher, "_inners", [])] - deepep = next((d for d in candidates if isinstance(d, DeepEPDispatcher)), None) - if deepep is None: - continue - # Prefer the low-latency impl (the mode foundry captures); the buffer - # is sized for whichever impls exist, so either bootstraps the shared - # singleton. - impl = getattr(deepep, "_low_latency_dispatcher", None) or getattr( - deepep, "_normal_dispatcher", None - ) - if impl is None: - continue - t0 = time.perf_counter() - impl._get_buffer() - logger.info( - "[Foundry] Bootstrapped DeepEP buffer pre-capture in %.3fs", - time.perf_counter() - t0, - ) - return True - - logger.warning( - "[Foundry] DeepEP backend active but no DeepEPDispatcher found on the " - "model; buffer not bootstrapped (capture may fail inside stream capture)." - ) - return False - - -def initialize_attention_metadata_for_bs(cuda_graph_runner, bs: int) -> None: - """Populate ``decode_cuda_graph_metadata[bs]`` for runtime replay. - - The FlashInfer wrappers and their internal ``_int_workspace_buffer`` - are constructed here, outside the captured graph. The graph's - runtime kernels reference these buffer addresses, so LOAD must - re-run the same call before runtime replay so the wrappers exist - at deterministic VMM addresses. - """ - buffers = cuda_graph_runner.buffers - num_tokens = bs * cuda_graph_runner.num_tokens_per_bs - encoder_lens = buffers.encoder_lens[:bs] if cuda_graph_runner.is_encoder_decoder else None - spec_info = cuda_graph_runner.get_spec_info(num_tokens) - cuda_graph_runner.attn_backend.init_forward_metadata_capture_cuda_graph( - bs, - num_tokens, - buffers.req_pool_indices[:bs], - buffers.seq_lens[:bs], - encoder_lens, - cuda_graph_runner.capture_forward_mode, - spec_info, - ) - - -def initialize_all_attention_metadata(cuda_graph_runner) -> None: - """Pre-pass: populate ``decode_cuda_graph_metadata`` for all bs at once. - - Called on both SAVE and LOAD before the capture/load loop. Walking - ``reversed(self.capture_bs)`` (largest first) matches SAVE's natural - capture order; same order on both sides keeps the VMM cursor - trajectory identical. - """ - for bs in reversed(cuda_graph_runner.capture_bs): - initialize_attention_metadata_for_bs(cuda_graph_runner, bs) - - -def load_all_graphs(cuda_graph_runner) -> None: - """LOAD-time replacement for the upstream capture loop. - - All FlashInfer wrappers are pre-allocated by - ``initialize_all_attention_metadata`` (called by the capture hook - before this function), so the VMM cursor sits where SAVE recorded - ``start_base_addr_0``. Load every graph in one - ``start_graph_builds`` call — this is what enables template + - on-demand linking in the manifest. ``finish_graph_loads`` then - replays each graph's alloc events in sequence, advancing the - cursor exactly the way SAVE did inside its capture loop. - """ - cfg = get_config() - state = get_state() - if cfg is None or state is None or cfg.workspace_dir is None: - raise RuntimeError("Foundry SGLang graph extension is not initialized") - - graph_files = _scan_graph_files(cfg.workspace_dir) - if not graph_files: - raise RuntimeError(f"No Foundry SGLang graph files found in {cfg.workspace_dir}") - - # NVSHMEM init runs once before any graph loads — graphs may reference - # NVSHMEM symbols. Single-GPU dense models have 0 NVSHMEM modules, so - # this is a no-op there but kept for EP parity. - cge.init_nvshmem_for_loaded_modules() - - paths = [os.path.join(cfg.workspace_dir, filename) for _, filename, _ in graph_files] - t0 = time.perf_counter() - pending = FoundryCUDAGraph.start_graph_builds(paths, num_threads=4) - results = FoundryCUDAGraph.finish_graph_loads(pending) - logger.info( - "[Foundry] Loaded %d SGLang graphs in %.3fs", - len(results), - time.perf_counter() - t0, - ) - - for i, (_index, _filename, meta) in enumerate(graph_files): - graph, tensors = results[i] - state.loaded_graphs[meta["key"]] = (graph, _unpack_output(tensors)) diff --git a/python/foundry/integration/sglang/hooks.py b/python/foundry/integration/sglang/hooks.py index 22cd7cfa..ea53ce3e 100644 --- a/python/foundry/integration/sglang/hooks.py +++ b/python/foundry/integration/sglang/hooks.py @@ -4,591 +4,620 @@ from __future__ import annotations -import functools import logging +import math import os -import time -from dataclasses import asdict +from dataclasses import asdict, fields +from typing import Any + +from foundry_sglang_plugin.targets import ( + ALLOC_MEMORY_POOL, + DECODE_CUDA_GRAPH_CAPTURE, + FULL_CUDA_GRAPH_BACKEND, + INIT_TORCH_DISTRIBUTED, + LAUNCH_SCHEDULERS, + MULTIMEM_ALL_GATHER_CALL, + RESOLVE_MEMORY_POOL_CONFIG, + SERVER_ARGS_POST_INIT, + TORCH_SYMM_MEM_INIT, +) from foundry.integration.sglang import runtime as rt from foundry.integration.sglang.config import ( CUDAGraphExtensionMode, get_graph_extension_mode, get_workspace_root, - load_graph_extension_config, ) logger = logging.getLogger(__name__) _INSTALLED = False +_MISSING = object() + +# Torch symmetric-memory TP communicators retained across construction, keyed by +# the tuple of process-group ranks. A repeated construction for the same rank set +# replaces the prior entry so the gather callback always sees exactly one enabled +# communicator rather than an apparent second TP communicator. +_torch_symm_mem_communicators: dict[tuple[int, ...], Any] = {} +_POSITIVE_MEMORY_POOL_INT_FIELDS = { + "max_total_num_tokens": False, + # Optional only while upstream builds an intermediate config. The target + # resolver always fills this before returning the config persisted by SAVE. + "max_running_requests": False, +} +_NONNEGATIVE_MEMORY_POOL_INT_FIELDS = { + "full_max_total_num_tokens": True, + "swa_max_total_num_tokens": True, + "c4_max_total_num_tokens": False, + "c128_max_total_num_tokens": False, + "c4_state_pool_size": False, + "c128_state_pool_size": False, +} + + +def _append_mismatch(problems, obj, field: str, expected) -> None: + actual = getattr(obj, field, _MISSING) + if actual != expected: + rendered = "" if actual is _MISSING else repr(actual) + problems.append(f"{field}={rendered} (require {expected!r})") + + +def _validate_profile(server_args) -> None: + """Reject launches outside the pinned dense TP=2 torch symmetric-memory profile.""" + problems = [] + + for field, expected in ( + ("nnodes", 1), + ("tp_size", 2), + ("pp_size", 1), + ("dp_size", 1), + ("dcp_size", 1), + ("ep_size", 1), + ("enable_dp_attention", False), + ("moe_a2a_backend", "none"), + ("disaggregation_mode", "null"), + ("enable_pdmux", False), + ("disable_custom_all_reduce", True), + ("enable_mscclpp", False), + # Torch symmetric-memory is the proven replay-safe TP collective path. + # ``enable_symm_mem`` is NCCL symmetric memory, a distinct path that must + # stay disabled; ``enforce_disable_flashinfer_allreduce_fusion`` stops an + # SM90 model silently selecting FlashInfer fused all-reduce instead. + ("enable_torch_symm_mem", True), + ("enable_symm_mem", False), + ("enforce_disable_flashinfer_allreduce_fusion", True), + ("enable_nccl_nvls", False), + ("flashinfer_allreduce_fusion_backend", None), + ("disable_flashinfer_autotune", True), + # One native list request must be tokenized and dispatched as one + # BatchTokenizedGenerateReqInput. Dynamic cross-request coalescing is a + # distinct timing-based mode and must stay disabled. + ("enable_tokenizer_batch_encode", True), + ("enable_dynamic_batch_tokenizer", False), + # Runtime batch sizes (e.g. 1, 8) must pad up to the single captured + # ShapeKey 256. With padding disabled SGLang would try to replay a shape + # that was never captured, so require it enabled (missing/True fail). + ("disable_cuda_graph_padding", False), + # decode_log_interval defaults to 40. A short concurrent bs>1 decode + # would then emit no "#running-req ... cuda graph: True" metrics line and + # the acceptance harness could not observe restored bs>1 replay. Require + # the recipe-pinned value 1 so every decode batch is logged; any other + # value or a missing field fails closed. + ("decode_log_interval", 1), + ): + _append_mismatch(problems, server_args, field, expected) + + attention_backend = getattr(server_args, "attention_backend", _MISSING) + for field in ("prefill_attention_backend", "decode_attention_backend"): + phase_backend = getattr(server_args, field, _MISSING) + if phase_backend is _MISSING: + problems.append(f"{field}= (require effective 'flashinfer')") + continue + effective_backend = phase_backend or attention_backend + if effective_backend != "flashinfer": + rendered = "" if attention_backend is _MISSING else repr(attention_backend) + problems.append( + f"{field}={phase_backend!r}, attention_backend={rendered} " + "(require effective 'flashinfer')" + ) - -def _ep_lazy_init_needed() -> bool: - """True when the DeepEP all-to-all backend is active, so pre-capture lazy - init (NVSHMEM buffer, DeepGEMM JIT) must be warmed up outside stream capture.""" - try: - from sglang.srt.layers.moe.utils import get_moe_a2a_backend - - return get_moe_a2a_backend().is_deepep() - except Exception: - return False - - -def _resolve_dp_rank(model_runner) -> int | None: - dp_rank = getattr(model_runner, "dp_rank", None) - if dp_rank is not None: - return dp_rank - - server_args = model_runner.server_args - if getattr(server_args, "enable_dp_attention", False): - from sglang.srt.layers.dp_attention import compute_dp_attention_world_info - - # Returns (attn_tp_rank, attn_tp_size, attn_dp_rank, attn_dp_size) — a - # 4-tuple. Take attn_dp_rank (3rd) and discard the rest. (Unpacking into - # 3 targets would raise ValueError.) - _, _, dp_rank, _ = compute_dp_attention_world_info( - server_args.enable_dp_attention, - model_runner.tp_rank, - server_args.tp_size, - server_args.dp_size, - server_args.attn_cp_size, + graph_config = getattr(server_args, "cuda_graph_config", _MISSING) + for phase, expected in (("decode", "full"), ("prefill", "disabled")): + phase_config = ( + getattr(graph_config, phase, _MISSING) if graph_config is not _MISSING else _MISSING ) - return dp_rank - - if getattr(server_args, "dp_size", 1) > 1: + backend = ( + getattr(phase_config, "backend", _MISSING) if phase_config is not _MISSING else _MISSING + ) + if backend != expected: + rendered = "" if backend is _MISSING else repr(backend) + problems.append(f"cuda_graph_config.{phase}.backend={rendered} (require {expected!r})") + + # Exactly one explicit padded decode shape, 256. SGLang pads smaller runtime + # batches to this single captured graph, so each rank produces one graph. + decode_config = ( + getattr(graph_config, "decode", _MISSING) if graph_config is not _MISSING else _MISSING + ) + for attr, expected in (("max_bs", 256), ("bs", [256])): + actual = ( + getattr(decode_config, attr, _MISSING) if decode_config is not _MISSING else _MISSING + ) + if actual != expected: + rendered = "" if actual is _MISSING else repr(actual) + problems.append(f"cuda_graph_config.decode.{attr}={rendered} (require {expected!r})") + + for name, expected in ( + ("NCCL_CUMEM_ENABLE", "0"), + ("NCCL_NVLS_ENABLE", "0"), + ("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:False"), + ("SGLANG_ENABLE_POST_CAPTURE_KV_SIZING", "0"), + ): + actual = os.environ.get(name) + if actual != expected: + problems.append(f"{name}={actual!r} (require {expected!r})") + + if problems: raise RuntimeError( - "Foundry SGLang integration cannot derive regular DP rank because " - "ModelRunner.dp_rank is absent. Preserve the constructor dp_rank on " - "ModelRunner before initializing torch distributed." + "[Foundry] SGLang launch outside the validated dense TP=2 " + "torch-symmetric-memory full-decode profile: " + "; ".join(problems) ) - return None +def _after_post_init(original_result, server_args): + _validate_profile(server_args) + return original_result -def install_hooks(server_args) -> None: - global _INSTALLED - cfg_path = getattr(server_args, "foundry_graph_extension_config_path", None) - if not cfg_path: - return - if _INSTALLED: - return - t0_ns = os.environ.get("FOUNDRY_SPAWN_T0_NS") - if t0_ns: - logger.info( - "[Foundry] SGLang spawn -> install_hooks: %.1f ms", - (time.perf_counter_ns() - int(t0_ns)) / 1e6, - ) +def _around_launch_scheduler_processes(original, *args, **kwargs): + if get_graph_extension_mode() != CUDAGraphExtensionMode.NONE: + rt.setup_ld_preload_env() + return original(*args, **kwargs) - load_graph_extension_config(cfg_path) - logger.info( - "[Foundry] SGLang hooks installing: mode=%s workspace=%s", - get_graph_extension_mode().value, - get_workspace_root(), + +def _around_init_torch_distributed(original, model_runner, *args, **kwargs): + if get_graph_extension_mode() == CUDAGraphExtensionMode.NONE: + return original(model_runner, *args, **kwargs) + + rt.setup_graph_extension( + model_runner.server_args, + tp_rank=model_runner.ps.tp_rank, + pp_rank=model_runner.ps.pp_rank, + dp_rank=model_runner.ps.dp_rank, + ) + rt.log_alloc_offset("after_setup_graph_ext") + result = original(model_runner, *args, **kwargs) + rt.log_alloc_offset("after_init_torch_dist") + rt.skip_to_scratch_boundary() + rt.log_alloc_offset("after_scratch_skip") + return result + + +def _register_profile_validation(hook_registry, hook_type) -> None: + hook_registry.register( + SERVER_ARGS_POST_INIT, + _after_post_init, + hook_type.AFTER, ) - _patch_init_torch_distributed() - _patch_init_memory_pool() - _patch_load_model() - _patch_kernel_warmup() - _patch_cuda_graph_capture() - _patch_spawn_sites() - _INSTALLED = True - logger.info("[Foundry] SGLang hooks installed") - - -def _patch_init_torch_distributed() -> None: - from sglang.srt.model_executor import model_runner as mr - - cls = mr.ModelRunner - orig = cls.init_torch_distributed - - @functools.wraps(orig) - def patched(self, *args, **kwargs): - mode = get_graph_extension_mode() - if mode == CUDAGraphExtensionMode.NONE: - return orig(self, *args, **kwargs) - - # Bind this rank's CUDA device BEFORE reserving the VMM region. - # init_torch_distributed (orig) calls set_device(self.gpu_id) - # internally, but foundry's set_allocation_region (inside - # setup_graph_extension) reserves the region on the *current* device. - # For DP rank > 0 the current device is still cuda:0 at this point, so - # without setting it first the region lands on the wrong GPU and the - # rank's later allocations fault with an async illegal memory access - # (surfacing at the first Stream()/kernel). Mirrors model_runner's own - # set_device(self.gpu_id). Single-GPU is unaffected (gpu_id == 0). - if self.device == "cuda": - import torch - - torch.get_device_module(self.device).set_device(self.gpu_id) - - rt.setup_graph_extension( - self.server_args, - tp_rank=self.tp_rank, - pp_rank=self.pp_rank, - dp_rank=_resolve_dp_rank(self), - ) - rt.log_alloc_offset("after_setup_graph_ext") - result = orig(self, *args, **kwargs) - rt.log_alloc_offset("after_init_torch_dist") - rt.skip_to_scratch_boundary() - rt.log_alloc_offset("after_scratch_skip") - return result +def _register_spawn_hooks(hook_registry, hook_type) -> None: + hook_registry.register( + LAUNCH_SCHEDULERS, + _around_launch_scheduler_processes, + hook_type.AROUND, + ) - cls.init_torch_distributed = patched +def _register_rank_vmm_hook(hook_registry, hook_type) -> None: + hook_registry.register( + INIT_TORCH_DISTRIBUTED, + _around_init_torch_distributed, + hook_type.AROUND, + ) -def _patch_init_memory_pool() -> None: - from sglang.srt.model_executor import model_runner_kv_cache_mixin as kv_mixin - from sglang.srt.model_executor.pool_configurator import MemoryPoolConfig - cls = kv_mixin.ModelRunnerKVCacheMixin - orig = cls.init_memory_pool - - @functools.wraps(orig) - def patched(self, pre_model_load_memory): - mode = get_graph_extension_mode() - if mode == CUDAGraphExtensionMode.NONE: - return orig(self, pre_model_load_memory) - - if mode == CUDAGraphExtensionMode.LOAD: - import torch - - rt.log_alloc_offset("before_init_memory_pool") - state = rt.load_warmup_state() - if not state.memory_pool_config: - raise RuntimeError("Foundry LOAD requires memory_pool_config") - self.memory_pool_config = MemoryPoolConfig(**state.memory_pool_config) - # Mirror SAVE's ``_resolve_memory_pool_config`` -> - # ``get_available_gpu_memory(empty_cache=True)`` side - # effect. Without this, torch's caching allocator retains - # segments that SAVE released — causing the - # attention-backend init below to take a different - # cuMemAlloc path and drift the VMM cursor away from - # SAVE's recorded ``start_base_addr``. - torch.cuda.empty_cache() - self._apply_memory_pool_config(self.memory_pool_config) - rt.log_alloc_offset("after_init_memory_pool") - logger.info("[Foundry] SGLang reused saved memory pool config") - return None - - rt.log_alloc_offset("before_init_memory_pool") - result = orig(self, pre_model_load_memory) - rt.log_alloc_offset("after_init_memory_pool") - state = rt.create_warmup_state(asdict(self.memory_pool_config)) - rt.save_warmup_state(state) - return result +def _register_memory_pool_hooks(hook_registry, hook_type) -> None: + hook_registry.register( + ALLOC_MEMORY_POOL, + _around_alloc_memory_pool, + hook_type.AROUND, + ) + hook_registry.register( + RESOLVE_MEMORY_POOL_CONFIG, + _around_resolve_memory_pool_config, + hook_type.AROUND, + ) - cls.init_memory_pool = patched - - -def _patch_load_model() -> None: - from sglang.srt.model_executor import model_runner as mr - - cls = mr.ModelRunner - orig = cls.load_model - - @functools.wraps(orig) - def patched(self, *args, **kwargs): - return orig(self, *args, **kwargs) - - cls.load_model = patched - - -def _patch_kernel_warmup() -> None: - from sglang.srt.model_executor import model_runner as mr - - cls = mr.ModelRunner - orig = cls.kernel_warmup - - @functools.wraps(orig) - def patched(self, *args, **kwargs): - mode = get_graph_extension_mode() - if mode == CUDAGraphExtensionMode.NONE: - return orig(self, *args, **kwargs) - # Phase 1 keeps pre-graph model-forward warmups out of SAVE/LOAD. - logger.info("[Foundry] SGLang kernel_warmup skipped in %s mode", mode.value) - return None - - cls.kernel_warmup = patched - - -def _patch_cuda_graph_capture() -> None: - from sglang.srt.model_executor import cuda_graph_runner as cgr - - cls = cgr.CudaGraphRunner - orig_capture = cls.capture - orig_create_device_graph = cls._create_device_graph - orig_capture_graph = cls._capture_graph - orig_capture_one_batch_size = cls.capture_one_batch_size - - # When set, the capture machinery is being reused as a foundry-driven WARMUP - # pass: run a real forward per bs (no graph capture, no save) to trigger all - # of sglang's pre-capture lazy init — DeepEP buffer, DeepGEMM per-shape JIT, - # etc. — that would otherwise fire inside the captured stream and abort with - # "operation not permitted when stream is capturing". See `_run_warmup_pass`. - warmup_active = [False] - - @functools.wraps(orig_create_device_graph) - def patched_create_device_graph(self, *args, **kwargs): - mode = get_graph_extension_mode() - if warmup_active[0]: - # Throwaway graph object; the warmup never captures into it. - return orig_create_device_graph(self, *args, **kwargs) - if mode == CUDAGraphExtensionMode.SAVE: - from foundry.integration.sglang.graph_ops import create_device_graph - - return create_device_graph() - return orig_create_device_graph(self, *args, **kwargs) - - @functools.wraps(orig_capture_graph) - def patched_capture_graph(self, graph, pool, stream, run_once_fn): - mode = get_graph_extension_mode() - if warmup_active[0]: - # Warmup: run the forward eagerly (NO torch.cuda.graph), so DeepGEMM - # JIT compile / NVSHMEM buffer creation happen outside stream capture. - return run_once_fn() - if mode == CUDAGraphExtensionMode.SAVE: - from foundry.integration.sglang.graph_ops import capture_graph - - return capture_graph(graph, pool, stream, run_once_fn) - return orig_capture_graph(self, graph, pool, stream, run_once_fn) - - @functools.wraps(orig_capture_one_batch_size) - def patched_capture_one_batch_size(self, bs, forward, stream_idx=None): - mode = get_graph_extension_mode() - if warmup_active[0]: - # Warmup pass: run upstream unchanged (real warmup forwards + a - # neutered "capture" that is just another real forward). No - # suppression, no save_graph. - return orig_capture_one_batch_size(self, bs, forward, stream_idx) - if mode == CUDAGraphExtensionMode.SAVE: - # Suppress the two pre-capture warmup forwards - # (cuda_graph_runner.py: ``for _ in range(2): run_once()``). - # Their non-deterministic activation allocations would pollute - # the torch caching allocator with freed segments that LOAD - # cannot reproduce — causing the per-bs init's cache-miss vs - # cache-hit asymmetry that drifts the VMM cursor away from - # each saved ``start_base_addr``. JIT / autotune still happens - # inside the graph capture (3rd run_once invocation) and is - # recorded as alloc events, mirroring vLLM doc 04 §2. - counter = [0] - real_forward = forward - - def warmup_skipping_forward(*args, **kwargs): - counter[0] += 1 - if counter[0] <= 2: - return None - return real_forward(*args, **kwargs) - - forward = warmup_skipping_forward - graph, output = orig_capture_one_batch_size(self, bs, forward, stream_idx) - if mode == CUDAGraphExtensionMode.SAVE: - from foundry.integration.sglang.graph_ops import save_graph - - # Mirror the inline key shape upstream uses for self.graphs[key] - # in `_capture_one_stream`. ``_make_graph_key`` and - # ``get_capture_lora_variant`` were removed in sglang - # commit ce2506e1c (record_nolora_graph deprecation). - key = bs if stream_idx is None else f"{stream_idx}_{bs}" - save_graph(graph, output, key) - return graph, output - - def _run_warmup_pass(self): - """Foundry-driven pre-capture warmup for the DeepEP/EP path. - - Reuses the upstream capture loop with graph capture neutered (see - ``warmup_active``) to run one real forward per ``capture_bs`` BEFORE the - real capture/replay work — triggering every pre-capture lazy init sglang - normally does in its (foundry-suppressed) warmup forwards: DeepEP buffer - creation, DeepGEMM per-shape JIT compile, etc. Run on BOTH SAVE and LOAD - so the VMM cursor trajectory stays symmetric. ``self.graphs`` / - ``self.output_buffers`` are saved/cleared/restored so the throwaway - warmup entries don't leak into the real pass. - """ - warmup_active[0] = True - saved_graphs, saved_buffers = self.graphs, self.output_buffers - self.graphs, self.output_buffers = {}, {} - t0 = time.perf_counter() - try: - orig_capture(self) - finally: - warmup_active[0] = False - self.graphs, self.output_buffers = saved_graphs, saved_buffers - logger.info( - "[Foundry] SGLang EP warmup pass (lazy-init) completed in %.3fs", - time.perf_counter() - t0, - ) - @functools.wraps(orig_capture) - def patched(self, *args, **kwargs): - mode = get_graph_extension_mode() - - # DeepEP/EP only (gated so the validated dense / single-GPU / DP paths - # are untouched): - # SAVE: sglang triggers pre-capture lazy init (DeepGEMM per-shape JIT, - # DeepEP buffer, ...) via warmup forwards that foundry suppresses; - # left lazy they fire inside the captured stream and abort - # ("operation not permitted when stream is capturing"). Run a real - # forward per bs up front so all that happens outside capture. - # LOAD: no capture happens — preallocate_for_load_mode reserves the - # whole region up to final_alloc_offset and replay places each alloc - # at its recorded absolute offset, so LOAD need NOT replay the warmup - # cursor trajectory. Running the warmup here is not just unnecessary - # but harmful: its orig_capture re-enters graph_capture() and leaves - # the context in a state that breaks the threaded finish_graph_loads - # ("invalid device context"). So warmup is SAVE-only. - # bootstrap_deepep_buffer runs on both (cheap singleton) to guarantee the - # NVSHMEM runtime is up before replay on LOAD / before capture on SAVE. - if _ep_lazy_init_needed(): - if mode == CUDAGraphExtensionMode.SAVE: - _run_warmup_pass(self) - if mode != CUDAGraphExtensionMode.NONE: - from foundry.integration.sglang.graph_ops import ( - bootstrap_deepep_buffer, - ) +def _raise_invalid_memory_pool_field(field: str, requirement: str, value) -> None: + raise RuntimeError( + "Foundry saved memory_pool_config field " + f"{field!r} must be {requirement}; got {value!r} ({type(value).__name__})" + ) - bootstrap_deepep_buffer(self) - if mode == CUDAGraphExtensionMode.LOAD: - from sglang.srt.distributed.device_communicators.pynccl_allocator import ( - set_graph_pool_id, - ) +def _validate_memory_pool_int_field( + saved_config: dict, + field: str, + *, + optional: bool, + positive: bool, +) -> None: + value = saved_config[field] + if value is None and optional: + return + requirement = ( + f"a {'positive' if positive else 'non-negative'} integer{' or null' if optional else ''}" + ) + if type(value) is not int: + _raise_invalid_memory_pool_field(field, requirement, value) + if (positive and value <= 0) or (not positive and value < 0): + _raise_invalid_memory_pool_field(field, requirement, value) + + +def _validated_memory_pool_config_values(saved_config: dict) -> dict: + normalized = dict(saved_config) + for field, optional in _POSITIVE_MEMORY_POOL_INT_FIELDS.items(): + _validate_memory_pool_int_field( + normalized, + field, + optional=optional, + positive=True, + ) + for field, optional in _NONNEGATIVE_MEMORY_POOL_INT_FIELDS.items(): + _validate_memory_pool_int_field( + normalized, + field, + optional=optional, + positive=False, + ) - from foundry.integration.sglang.graph_ops import ( - initialize_all_attention_metadata, - load_all_graphs, + field = "mem_fraction_static" + value = normalized[field] + if value is None: + return normalized + if type(value) is int: + try: + converted = float(value) + except OverflowError: + _raise_invalid_memory_pool_field( + field, + "an exactly representable finite number or null", + value, ) - - state = rt.get_state() - if state is None: - raise RuntimeError("Foundry SGLang state is not initialized") - # Set up the graph memory pool once — sglang shares one pool - # across all captured graphs, and runtime replay also requires - # it to be set so pynccl knows which pool the graph belongs to. - if cgr.get_global_graph_memory_pool() is None: - cgr.set_global_graph_memory_pool(self.device_module.graph_pool_handle()) - set_graph_pool_id(cgr.get_global_graph_memory_pool()) - rt.log_alloc_offset("before_preallocate") - rt.preallocate_for_load_mode() - rt.log_alloc_offset("after_preallocate") - # Pre-pass: allocate every per-bs FlashInfer wrapper in - # ``reversed(capture_bs)`` order, matching the order SAVE used. - # SAVE's patched ``init_forward_metadata_capture_cuda_graph`` - # is idempotent so the inner per-iter call inside the SAVE - # capture loop does not re-allocate. Same upfront allocation - # sequence on both sides → cursor sits at SAVE's - # ``start_base_addr_0`` when graph load begins. - # FlashInfer-only pre-pass (see SAVE branch). fa3 etc. allocate their - # cuda-graph metadata once in init_cuda_graph_state, so skip it. - if hasattr(self.attn_backend, "indices_updater_decode"): - initialize_all_attention_metadata(self) - rt.log_alloc_offset("after_pre_init") - # Single ``start_graph_builds(all_paths)`` call so templates - # and on-demand graphs link via ``shared_exec`` in the - # manifest. ``finish_graph_loads`` replays alloc events - # graph-by-graph in the same order SAVE captured them. - load_all_graphs(self) - rt.log_alloc_offset("after_load_all_graphs") - self.graphs = {k: v[0] for k, v in state.loaded_graphs.items()} - self.output_buffers = {k: v[1] for k, v in state.loaded_graphs.items()} - # Non-FlashInfer backends (e.g. fa3) populate per-bs decode metadata - # — looked up at replay as attn_backend.decode_cuda_graph_metadata[bs] - # — inside the capture loop, which LOAD replaces. The FlashInfer - # pre-pass above (gated) handled flashinfer; for the rest, populate it - # now. Run AFTER load_all_graphs so the (already-correct) loaded-graph - # VMM offsets are unaffected — fa3's metadata are lightweight views - # over the fixed init_cuda_graph_state workspace, not graph memory. - if not hasattr(self.attn_backend, "indices_updater_decode"): - initialize_all_attention_metadata(self) - # Initialize the DeepEP cuda-graph adapter's captured mode. Upstream - # sets this in deepep_adapter.capture() during the capture loop, - # which LOAD replaces — so without this, runtime replay() asserts - # `_captured_deepep_mode is not None` on the first decode. capture() - # self-gates on the DeepEP backend (no-op otherwise) and sets the - # decode dispatch mode the captured graphs expect. - self.deepep_adapter.capture(is_extend_in_batch=False) - return None - - if mode == CUDAGraphExtensionMode.SAVE: - # FlashInfer allocates a per-bs metadata wrapper (each with its own - # _int_workspace_buffer) on every capture init, so foundry pre-allocates - # them up front and installs a reuse shim that makes the inner init - # reuse them — keeping the VMM cursor deterministic vs LOAD. Backends - # with a single fixed cuda-graph metadata workspace allocated once in - # init_cuda_graph_state (e.g. fa3 / FlashAttentionBackend) don't need - # this and the plain capture is already SAVE/LOAD-deterministic. Detect - # FlashInfer by its per-bs indices_updater_decode. - attn_backend = self.attn_backend - use_fi_prepass = hasattr(attn_backend, "indices_updater_decode") - real_init = attn_backend.init_forward_metadata_capture_cuda_graph - - def reuse_pre_pass_init( - bs, - num_tokens, - req_pool_indices, - seq_lens, - encoder_lens, - forward_mode, - spec_info, - ): - # The pre-pass already allocated a wrapper for this - # bs and stored it in - # ``decode_cuda_graph_metadata`` / - # ``prefill_cuda_graph_metadata``. Reuse it directly - # — no second torch.empty for ``_int_workspace_buffer``. - # Re-run the planner with the same buffer slices the - # capture forward uses, then point - # ``forward_metadata`` at the same wrappers. Same - # plan call on LOAD via the symmetric pre-pass, so - # the captured graph kernels reference VMM addresses - # that LOAD's wrappers actually occupy. - from sglang.srt.layers.attention.flashinfer_backend import ( - DecodeMetadata, - PrefillMetadata, + # JSON has one number type for integers. Accept it for Optional[float] + # only when conversion preserves the exact numeric value. + if not math.isfinite(converted) or int(converted) != value: + _raise_invalid_memory_pool_field( + field, + "an exactly representable finite number or null", + value, + ) + elif type(value) is float: + converted = value + else: + _raise_invalid_memory_pool_field(field, "a finite number or null", value) + if not math.isfinite(converted): + _raise_invalid_memory_pool_field(field, "a finite number or null", value) + if not 0.0 < converted <= 1.0: + _raise_invalid_memory_pool_field( + field, + "a fraction greater than 0 and at most 1, or null", + value, + ) + normalized[field] = converted + return normalized + + +def _around_alloc_memory_pool(original, model_runner, *args, **kwargs): + mode = get_graph_extension_mode() + if mode == CUDAGraphExtensionMode.NONE: + return original(model_runner, *args, **kwargs) + + rt.log_alloc_offset("before_alloc_memory_pool") + result = original(model_runner, *args, **kwargs) + rt.log_alloc_offset("after_alloc_memory_pool") + + if mode == CUDAGraphExtensionMode.SAVE: + extension_state = rt.get_state() + if extension_state is None: + raise RuntimeError("Foundry SGLang state is not initialized before memory-pool SAVE") + # warmup_state.json is shared by all ranks. Only rank zero may write it, + # avoiding concurrent truncation while preserving the existing shared state. + if extension_state.rank == 0: + memory_pool_config = getattr(model_runner, "memory_pool_config", None) + if memory_pool_config is None: + raise RuntimeError( + "Foundry SAVE requires ModelRunner.memory_pool_config after allocation" ) + try: + serialized_config = asdict(memory_pool_config) + except TypeError as exc: + raise RuntimeError( + "Foundry SAVE requires memory_pool_config to be a dataclass" + ) from exc + state = rt.create_warmup_state(serialized_config) + rt.save_warmup_state(state) + else: + logger.info("[Foundry] SGLang reused saved memory pool config") + + return result + + +def _around_resolve_memory_pool_config( + original, + configurator, + pre_model_load_memory, +): + if get_graph_extension_mode() != CUDAGraphExtensionMode.LOAD: + return original(configurator, pre_model_load_memory) + + # Lazy imports: SGLang is optional, and torch's CUDA cache operation is + # required only when LOAD bypasses upstream memory profiling. + from sglang.srt.model_executor.pool_configurator import MemoryPoolConfig - if forward_mode.is_decode_or_idle(): - wrappers = attn_backend.decode_cuda_graph_metadata.get(bs) - if wrappers is None: - return real_init( - bs, - num_tokens, - req_pool_indices, - seq_lens, - encoder_lens, - forward_mode, - spec_info, - ) - seq_lens_sum = seq_lens.sum().item() - attn_backend.indices_updater_decode.update( - req_pool_indices, - seq_lens, - seq_lens.cpu(), - seq_lens_sum, - decode_wrappers=wrappers, - encoder_lens=encoder_lens, - spec_info=spec_info, - fixed_split_size=None, - disable_split_kv=attn_backend.disable_cuda_graph_kv_split, - ) - attn_backend.forward_metadata = DecodeMetadata(wrappers) - return - if ( - forward_mode.is_target_verify() - or forward_mode.is_draft_extend() - or forward_mode.is_dllm_extend() - ): - wrappers = attn_backend.prefill_cuda_graph_metadata.get(bs) - if wrappers is None: - return real_init( - bs, - num_tokens, - req_pool_indices, - seq_lens, - encoder_lens, - forward_mode, - spec_info, - ) - seq_lens_sum = seq_lens.sum().item() - use_ragged = forward_mode.is_dllm_extend() - prefix_lens = ( - seq_lens - attn_backend.dllm_config.block_size - if forward_mode.is_dllm_extend() - else None - ) - spec_info_arg = None if forward_mode.is_dllm_extend() else spec_info - attn_backend.indices_updater_prefill.update( - req_pool_indices, - seq_lens, - seq_lens.cpu(), - seq_lens_sum, - prefix_lens=prefix_lens, - prefill_wrappers=wrappers, - use_ragged=use_ragged, - encoder_lens=encoder_lens, - spec_info=spec_info_arg, - ) - attn_backend.forward_metadata = PrefillMetadata(wrappers, use_ragged, False) - return - # Unknown mode — fall back to real init. - return real_init( - bs, - num_tokens, - req_pool_indices, - seq_lens, - encoder_lens, - forward_mode, - spec_info, - ) + state = rt.load_warmup_state() + saved_config = getattr(state, "memory_pool_config", None) + if not isinstance(saved_config, dict) or not saved_config: + raise RuntimeError("Foundry LOAD requires a saved memory_pool_config") - if use_fi_prepass: - from foundry.integration.sglang.graph_ops import ( - initialize_all_attention_metadata, - ) + expected_fields = {field.name for field in fields(MemoryPoolConfig)} + actual_fields = set(saved_config) + if actual_fields != expected_fields: + missing = sorted(expected_fields - actual_fields) + unknown = sorted(actual_fields - expected_fields) + raise RuntimeError( + "Foundry saved memory_pool_config schema is incompatible with pinned " + f"SGLang MemoryPoolConfig: missing={missing}, unknown={unknown}" + ) - rt.log_alloc_offset("save_before_pre_init") - # Pre-pass: allocate every per-bs FlashInfer wrapper up front, in - # the same ``reversed(capture_bs)`` order LOAD uses. - initialize_all_attention_metadata(self) - rt.log_alloc_offset("save_after_pre_init") - # Drop the pre-pass's last forward_metadata ref so popping the dict - # entry doesn't keep the wrapper alive at refcount 1. - attn_backend.forward_metadata = None - attn_backend.init_forward_metadata_capture_cuda_graph = reuse_pre_pass_init - try: - result = orig_capture(self, *args, **kwargs) - finally: - attn_backend.init_forward_metadata_capture_cuda_graph = real_init + saved_config = _validated_memory_pool_config_values(saved_config) + try: + memory_pool_config = MemoryPoolConfig(**saved_config) + except TypeError as exc: + raise RuntimeError( + "Foundry saved memory_pool_config cannot construct pinned SGLang MemoryPoolConfig" + ) from exc + + import torch + + # SAVE's resolver calls get_available_gpu_memory() with its default + # empty_cache=True. LOAD bypasses that resolver, so mirror only that side effect. + torch.cuda.empty_cache() + return memory_pool_config + + +def _around_decode_capture(original, self, *args, **kwargs): + """AROUND ``DecodeCudaGraphRunner.capture`` — the SAVE/LOAD branch point. + + The registry records a single AROUND registration; the body dispatches on + the Foundry extension mode: + + * NONE — pass through to the upstream capture loop unchanged. + * SAVE — run the upstream capture loop (the REPLACE'd + ``FullCudaGraphBackend.capture_one`` writes each Foundry graph), then + finalize the archive (manifest + fatbins + final alloc offset). + * LOAD — ALSO run the upstream capture loop via ``original``. The pinned + ``capture()`` warms up, enters ``graph_capture()``, fills the shared + seq-len buffers, and — inside ``backend.capture_session`` — runs + ``capture_prepare`` / ``ForwardContext`` / real attention-metadata init + per shape before dispatching to ``backend.capture_one``. The Foundry + backend's LOAD ``capture_session`` (pool + preallocate + one + ``start_graph_builds``) and LOAD ``capture_one`` (per-index + ``finish_one_graph_load``) materialize each archived graph. + + Running the identical upstream lifecycle on both SAVE and LOAD is the point: + the allocator trajectory and the FlashInfer metadata wrappers the captured + kernels reference are built the same way on both sides, so replay is exact. + All LOAD orchestration lives in the backend; the hook adds nothing here. + """ + mode = get_graph_extension_mode() + if mode == CUDAGraphExtensionMode.NONE: + return original(self, *args, **kwargs) + + # Post-construction/runtime fail-closed guard. ``_validate_profile`` checks + # the *requested* ``cuda_graph_config.decode.bs == [256]``, but the runner + # rebuilds ``capture_bs`` from that config filtered/clamped by the request + # pool size and TP/CP alignment (``get_batch_sizes_to_capture``). If pool or + # config resolution shrank or replaced the shape, SAVE would archive — and + # LOAD would restore — a shape other than the single padded ShapeKey 256 + # runtime batches pad to. Refuse before any SAVE/LOAD work. + capture_bs = getattr(self, "capture_bs", None) + if capture_bs != [256]: + raise RuntimeError( + "[Foundry] SGLang decode capture requires exactly one padded shape " + f"capture_bs == [256]; got {capture_bs!r}. Runtime batches pad to this " + "single captured ShapeKey, so a drifted/clamped shape is fatal." + ) - from foundry.integration.sglang.graph_ops import ( - pack_fatbins, - save_graph_manifest, - ) + if mode == CUDAGraphExtensionMode.SAVE: + result = original(self, *args, **kwargs) + from foundry.integration.sglang.graph_ops import ( + pack_fatbins, + save_graph_manifest, + ) - save_graph_manifest() - pack_fatbins() - rt.capture_final_alloc_offset() - return result + save_graph_manifest() + pack_fatbins() + rt.capture_final_alloc_offset() + return result - return orig_capture(self, *args, **kwargs) + # LOAD — delegate to the pinned upstream capture lifecycle; the Foundry + # backend's capture_session / capture_one restore each archived graph. + return original(self, *args, **kwargs) + + +def _around_torch_symm_mem_init(original, communicator, *args, **kwargs): + """AROUND ``TorchSymmMemCommunicator.__init__``. + + PyTorch symmetric memory owns its own cross-rank virtual-address contract, so + its buffer must be rendezvoused outside Foundry's monotonic VMM region. Before + the original constructor runs, raise the symmetric signal-pad size to at least + the multimem gather requirement, ``stop`` Foundry's allocation region, and + always ``resume`` it in ``finally``. A successfully constructed communicator is + retained by its process-group rank tuple for the gather callback. + """ + if get_graph_extension_mode() == CUDAGraphExtensionMode.NONE: + return original(communicator, *args, **kwargs) + + # Lazy imports: torch and the pinned SGLang symmetric-memory modules are + # optional for core-only Foundry installs and only needed once active. + import torch + from sglang.srt.distributed.device_communicators import ( + torch_symm_mem as torch_symm_mem_module, + ) + from sglang.srt.distributed.device_communicators import triton_symm_mem_ag - cls._create_device_graph = patched_create_device_graph - cls._capture_graph = patched_capture_graph - cls.capture_one_batch_size = patched_capture_one_batch_size - cls.capture = patched + from foundry import ops as cge + group = kwargs.get("group", args[0] if args else None) + if group is None: + raise RuntimeError("Foundry could not identify the torch symmetric-memory group") + world_size = torch.distributed.get_world_size(group) + signal_pad_bytes = triton_symm_mem_ag._MAX_BLOCKS * world_size * 4 + symm_mem = torch_symm_mem_module.torch_symm_mem + symm_mem.set_signal_pad_size(max(symm_mem.get_signal_pad_size(), signal_pad_bytes)) -def _patch_spawn_sites() -> None: + cge.stop_allocation_region() try: - from sglang.srt.entrypoints import engine as engine_mod - except Exception: - engine_mod = None + result = original(communicator, *args, **kwargs) + finally: + cge.resume_allocation_region() + + if communicator.buffer is not None: + group_ranks = tuple(torch.distributed.get_process_group_ranks(communicator.group)) + _torch_symm_mem_communicators[group_ranks] = communicator + logger.info( + "[Foundry] Torch symmetric-memory buffer address=0x%x", + communicator.buffer.data_ptr(), + ) + return result - if engine_mod is not None: - orig_launch = engine_mod.Engine._launch_scheduler_processes - @functools.wraps(orig_launch) - def patched_launch(self, *args, **kwargs): - if get_graph_extension_mode() != CUDAGraphExtensionMode.NONE: - rt.setup_ld_preload_env() - return orig_launch(self, *args, **kwargs) +def _around_multimem_all_gather(original, gatherer, x): + """AROUND ``MultimemAllGatherer.__call__``. - engine_mod.Engine._launch_scheduler_processes = patched_launch + In an active Foundry mode this bypasses the upstream lazy gather builder (which + would allocate a process-random symmetric buffer during warmup/capture) and + instead issues ``multimem_all_gather_out`` directly over the sole retained + torch symmetric-memory communicator's stable buffer and rendezvous group name. + It never performs a second rendezvous. + """ + if get_graph_extension_mode() == CUDAGraphExtensionMode.NONE: + return original(gatherer, x) - try: - from sglang.srt.managers import data_parallel_controller as dpc - except Exception: - dpc = None + # Lazy imports: torch and the pinned SGLang gatherer module are optional for + # core-only Foundry installs and only needed once active. + import torch + from sglang.srt.distributed.device_communicators import triton_symm_mem_ag + + if len(_torch_symm_mem_communicators) != 1: + raise RuntimeError( + "Foundry requires exactly one torch symmetric-memory TP communicator, " + f"found {len(_torch_symm_mem_communicators)}" + ) + communicator = next(iter(_torch_symm_mem_communicators.values())) + if x.ndim != 2 or x.dtype != torch.bfloat16 or not x.is_contiguous(): + raise RuntimeError("Foundry multimem all-gather requires contiguous rank-2 bfloat16 input") + world_size = communicator.world_size + token_count, local_hidden_size = x.shape + # Match upstream ``all_gather_inner`` eligibility (fail closed, never fall + # back to NCCL in an active Foundry mode): a positive token count, a 16-byte + # aligned input pointer for the 128-bit ``multimem.st`` store, and a per-rank + # hidden shard that is a multiple of ``_NUMEL_PER_THREAD`` bf16 elements. + if token_count <= 0: + raise RuntimeError( + f"Foundry multimem all-gather requires a positive token count, got {token_count}" + ) + if x.data_ptr() % 16 != 0: + raise RuntimeError( + "Foundry multimem all-gather requires a 16-byte aligned input pointer, " + f"got 0x{x.data_ptr():x}" + ) + numel_per_thread = triton_symm_mem_ag._NUMEL_PER_THREAD + if local_hidden_size % numel_per_thread != 0: + raise RuntimeError( + f"Foundry multimem all-gather requires the per-rank hidden shard " + f"({local_hidden_size}) to be a multiple of {numel_per_thread} bf16 elements" + ) + required_numel = world_size * x.numel() + if required_numel > communicator.buffer.numel(): + raise RuntimeError( + "Foundry torch symmetric-memory buffer is too small for " + f"multimem all-gather: need {required_numel} bf16 elements, " + f"have {communicator.buffer.numel()}" + ) + comm_buff = communicator.buffer[:required_numel].view( + world_size * token_count, + local_hidden_size, + ) + torch.ops.symm_mem.multimem_all_gather_out( + x, + communicator.group.group_name, + comm_buff, + ) + if not getattr(gatherer, "_foundry_observed", False): + logger.info( + "[Foundry] Multimem all-gather buffer address=0x%x", + comm_buff.data_ptr(), + ) + gatherer._foundry_observed = True + return ( + comm_buff.view(world_size, token_count, local_hidden_size) + .movedim(0, 1) + .reshape(token_count, world_size * local_hidden_size) + ) + + +def _register_symmetric_collective_hooks(hook_registry, hook_type) -> None: + hook_registry.register( + TORCH_SYMM_MEM_INIT, + _around_torch_symm_mem_init, + hook_type.AROUND, + ) + hook_registry.register( + MULTIMEM_ALL_GATHER_CALL, + _around_multimem_all_gather, + hook_type.AROUND, + ) + + +def _register_capture_hooks(hook_registry, hook_type) -> None: + # Lazy import: building the subclass touches the pinned + # ``FullCudaGraphBackend``, which is optional for core-only installs. + from foundry.integration.sglang.backend import build_foundry_full_backend_class - if dpc is not None: - orig_start = dpc.DataParallelController.launch_tensor_parallel_group + hook_registry.register( + FULL_CUDA_GRAPH_BACKEND, + build_foundry_full_backend_class(), + hook_type.REPLACE, + ) + hook_registry.register( + DECODE_CUDA_GRAPH_CAPTURE, + _around_decode_capture, + hook_type.AROUND, + ) + + +def register_hooks() -> None: + """Register the current-main profile, rank/VMM, memory-pool, and capture hooks.""" + global _INSTALLED + if _INSTALLED: + return - @functools.wraps(orig_start) - def patched_start(self, *args, **kwargs): - if get_graph_extension_mode() != CUDAGraphExtensionMode.NONE: - rt.setup_ld_preload_env() - return orig_start(self, *args, **kwargs) + # Lazy import: SGLang is an optional integration and is unavailable in + # core-only Foundry installations. + from sglang.srt.plugins.hook_registry import HookRegistry, HookType - dpc.DataParallelController.launch_tensor_parallel_group = patched_start + logger.info( + "[Foundry] SGLang hooks registering: mode=%s workspace=%s", + get_graph_extension_mode().value, + get_workspace_root(), + ) + _register_profile_validation(HookRegistry, HookType) + _register_spawn_hooks(HookRegistry, HookType) + _register_rank_vmm_hook(HookRegistry, HookType) + _register_memory_pool_hooks(HookRegistry, HookType) + _register_capture_hooks(HookRegistry, HookType) + _register_symmetric_collective_hooks(HookRegistry, HookType) + _INSTALLED = True + logger.info("[Foundry] SGLang hooks registered") diff --git a/python/foundry/integration/sglang/runtime.py b/python/foundry/integration/sglang/runtime.py index 1f9032c5..cad04ffa 100644 --- a/python/foundry/integration/sglang/runtime.py +++ b/python/foundry/integration/sglang/runtime.py @@ -8,7 +8,9 @@ import logging import os import shutil +import tempfile import time +from contextlib import suppress from dataclasses import asdict, dataclass, field from datetime import datetime from pathlib import Path @@ -87,12 +89,29 @@ def save_warmup_state(state: WarmupState) -> None: if workspace_root is None: return ext_state = get_state() - path = os.path.join(workspace_root, "warmup_state.json") - if ext_state is not None and ext_state.rank != 0 and os.path.exists(path): + if ext_state is not None and ext_state.rank != 0: return + os.makedirs(workspace_root, exist_ok=True) - with open(path, "w") as f: - json.dump(asdict(state), f, indent=2) + path = os.path.join(workspace_root, "warmup_state.json") + temp_path = None + try: + with tempfile.NamedTemporaryFile( + mode="w", + encoding="utf-8", + dir=workspace_root, + prefix=f".warmup_state.{os.getpid()}.", + suffix=".tmp", + delete=False, + ) as f: + temp_path = f.name + json.dump(asdict(state), f, indent=2) + os.replace(temp_path, path) + temp_path = None + finally: + if temp_path is not None: + with suppress(FileNotFoundError): + os.unlink(temp_path) logger.info("[Foundry] Saved SGLang warmup state to %s", path) @@ -103,8 +122,16 @@ def load_warmup_state() -> WarmupState: path = os.path.join(workspace_root, "warmup_state.json") if not os.path.exists(path): raise RuntimeError(f"Foundry warmup state file not found: {path}") - with open(path) as f: - data = json.load(f) + try: + with open(path) as f: + data = json.load(f) + except json.JSONDecodeError as exc: + raise RuntimeError(f"Foundry warmup state contains invalid JSON at {path}: {exc}") from exc + if not isinstance(data, dict): + raise RuntimeError( + "Foundry incompatible warmup state at " + f"{path}: expected a JSON object, got {type(data).__name__}" + ) valid = set(WarmupState.__dataclass_fields__.keys()) return WarmupState(**{k: v for k, v in data.items() if k in valid}) @@ -170,7 +197,8 @@ def capture_final_alloc_offset() -> int: path = os.path.join(cfg.workspace_dir, "final_alloc_offset.json") with open(path, "w") as f: json.dump({"final_alloc_offset": _final_alloc_offset}, f) - if cfg.workspace_root is not None: + ext_state = get_state() + if cfg.workspace_root is not None and ext_state is not None and ext_state.rank == 0: warmup_state_path = os.path.join(cfg.workspace_root, "warmup_state.json") if os.path.exists(warmup_state_path): state = load_warmup_state() @@ -192,9 +220,26 @@ def preallocate_for_load_mode() -> None: final = json.load(f).get("final_alloc_offset", 0) if final <= 0: final = load_warmup_state().final_alloc_offset - remaining = final - cge.get_current_alloc_offset() - if remaining > 0: - cge.preallocate_region(remaining) + current = cge.get_current_alloc_offset() + if final <= 0: + raise RuntimeError( + "Foundry SGLang LOAD final_alloc_offset must be positive after rank/shared " + f"fallback: final_alloc_offset={final}, current_offset={current}, " + f"workspace={cfg.workspace_dir}" + ) + if current > final: + raise RuntimeError( + "Foundry SGLang LOAD current allocation offset exceeds captured final offset: " + f"current_offset={current}, final_alloc_offset={final}, " + f"workspace={cfg.workspace_dir}" + ) + remaining = final - current + if remaining > 0 and not cge.preallocate_region(remaining): + raise RuntimeError( + "Foundry SGLang LOAD preallocation failed before graph loading: " + f"remaining={remaining}, current_offset={current}, " + f"final_alloc_offset={final}, workspace={cfg.workspace_dir}" + ) def log_alloc_offset(label: str) -> None: diff --git a/python/foundry/integration/vllm/config.py b/python/foundry/integration/vllm/config.py index b03c29e3..0adbf89d 100644 --- a/python/foundry/integration/vllm/config.py +++ b/python/foundry/integration/vllm/config.py @@ -10,6 +10,7 @@ import enum import importlib.util +import json from dataclasses import dataclass from pathlib import Path from typing import TYPE_CHECKING @@ -17,6 +18,13 @@ import tomllib from vllm.logger import init_logger +from foundry.integration.deepep_transport import ( + DeepEPTransport, + parse_deepep_transport, + validate_archived_deepep_profile, + validate_transport_match, +) + if TYPE_CHECKING: from vllm.config import ParallelConfig @@ -35,6 +43,7 @@ class CUDAGraphExtensionConfig: hook_library_path: str | None = None # Path to libnvshmem_host.so for NVSHMEM-based kernels (e.g. DeepEP). nvshmem_host_path: str | None = None + deepep_transport: DeepEPTransport = DeepEPTransport.FABRIC base_addr: int = 0x600000000000 region_size: str = "64GB" workspace_root: str = "foundry_archive" @@ -59,10 +68,15 @@ def from_toml(cls, path: str | Path) -> CUDAGraphExtensionConfig: if hook_library_path is None: hook_library_path = cls._detect_hook_so_path() + nvshmem_host_path = data.get("nvshmem_host_path") + if nvshmem_host_path is None: + nvshmem_host_path = cls._detect_nvshmem_host_path() + return cls( mode=mode, hook_library_path=hook_library_path, - nvshmem_host_path=data.get("nvshmem_host_path"), + nvshmem_host_path=nvshmem_host_path, + deepep_transport=parse_deepep_transport(data.get("deepep_transport")), base_addr=base_addr, region_size=data.get("region_size", cls.region_size), workspace_root=data.get("workspace_root", cls.workspace_root), @@ -79,11 +93,56 @@ def _detect_hook_so_path() -> str | None: return str(hook_so_path) return None + @staticmethod + def _detect_nvshmem_host_path() -> str | None: + """Locate DeepEP's NVSHMEM host library from the optional wheel.""" + try: + spec = importlib.util.find_spec("nvidia.nvshmem") + except (ImportError, ValueError): + return None + if spec is None: + return None + roots = list(spec.submodule_search_locations or []) + if spec.origin: + roots.append(str(Path(spec.origin).parent)) + for root in roots: + for name in ("libnvshmem_host.so.3", "libnvshmem_host.so"): + candidate = Path(root) / "lib" / name + if candidate.exists(): + return str(candidate) + return None + # Module-global config singleton (process-local). _config: CUDAGraphExtensionConfig | None = None +def validate_load_archive_transport(config: CUDAGraphExtensionConfig) -> None: + """Validate LOAD's archived transport without importing CUDA runtime code.""" + if config.mode != CUDAGraphExtensionMode.LOAD: + return + + path = Path(config.workspace_root) / "warmup_state.json" + if not path.exists(): + raise RuntimeError(f"Warmup state file not found: {path}") + try: + with path.open() as f: + data = json.load(f) + except json.JSONDecodeError as exc: + raise RuntimeError(f"Foundry warmup state contains invalid JSON at {path}: {exc}") from exc + if not isinstance(data, dict): + raise RuntimeError( + "Foundry incompatible warmup state at " + f"{path}: expected a JSON object, got {type(data).__name__}" + ) + + archived_transport, _ = validate_archived_deepep_profile(data, path=str(path)) + validate_transport_match( + archived=archived_transport.value, + configured=config.deepep_transport, + ) + + def load_graph_extension_config(path: str) -> None: """Load TOML and set the process-local ``_config`` singleton. @@ -91,6 +150,7 @@ def load_graph_extension_config(path: str) -> None: """ global _config _config = CUDAGraphExtensionConfig.from_toml(path) + validate_load_archive_transport(_config) def get_config() -> CUDAGraphExtensionConfig | None: @@ -127,6 +187,12 @@ def get_nvshmem_host_path() -> str | None: return _config.nvshmem_host_path +def get_deepep_transport() -> DeepEPTransport: + if _config is None: + return DeepEPTransport.FABRIC + return _config.deepep_transport + + def _compute_rank_from_parallel_config( parallel_config: ParallelConfig, ) -> int: diff --git a/python/foundry/integration/vllm/graph_ops.py b/python/foundry/integration/vllm/graph_ops.py index 4cda957f..5c8f9f27 100644 --- a/python/foundry/integration/vllm/graph_ops.py +++ b/python/foundry/integration/vllm/graph_ops.py @@ -40,6 +40,10 @@ get_config, get_graph_extension_mode, ) +from foundry.integration.vllm.restore_progress import ( + LoadRestorationProgress, + rank_from_workspace_dir, +) from foundry.integration.vllm.runtime import get_state if TYPE_CHECKING: @@ -192,6 +196,21 @@ def capture_or_load_graph( if index is None: raise RuntimeError(f"Graph identifier not in pending builds: {identifier}") graph, output = FoundryCUDAGraph.finish_one_graph_load(_pending_graph_builds.pending, index) + restoration = _pending_graph_builds.restoration_progress.record(index) + log.info( + "[foundry] LOAD graph finalized rank_%d index=%d restored=%d/%d", + _pending_graph_builds.restoration_progress.rank, + index, + restoration.count, + _pending_graph_builds.restoration_progress.total, + ) + if restoration.completed_now: + log.info( + "[foundry] LOAD restoration complete rank_%d: %d/%d graphs", + _pending_graph_builds.restoration_progress.rank, + restoration.count, + _pending_graph_builds.restoration_progress.total, + ) # Strong-reference the (graph, output) pair. The upstream # CUDAGraphWrapper stores `entry.output = weak_ref_tensors(output)`, # so without this dict the VMM-backed output tensors get GC'd between @@ -220,6 +239,7 @@ def capture_or_load_graph( class _PendingGraphBuilds: pending: object # opaque PendingGraphLoads handle from C++ identifier_to_index: dict[GraphIdentifier, int] + restoration_progress: LoadRestorationProgress _pending_graph_builds: _PendingGraphBuilds | None = None @@ -301,6 +321,10 @@ def start_graph_builds() -> None: _pending_graph_builds = _PendingGraphBuilds( pending=pending, identifier_to_index=identifier_to_index, + restoration_progress=LoadRestorationProgress( + rank=rank_from_workspace_dir(workspace_dir), + total=len(graph_files), + ), ) diff --git a/python/foundry/integration/vllm/hooks.py b/python/foundry/integration/vllm/hooks.py index f7768630..bb3ab22c 100644 --- a/python/foundry/integration/vllm/hooks.py +++ b/python/foundry/integration/vllm/hooks.py @@ -14,14 +14,21 @@ import os import time from collections.abc import Callable +from contextlib import ExitStack from typing import TYPE_CHECKING +from unittest.mock import patch as mock_patch import torch from vllm.logger import init_logger +from foundry.integration.deepep_transport import ( + configure_deepep_buffer, + validate_num_nvl_bytes_match, +) from foundry.integration.vllm import runtime as rt from foundry.integration.vllm.config import ( CUDAGraphExtensionMode, + get_deepep_transport, get_graph_extension_mode, get_workspace_root, load_graph_extension_config, @@ -30,6 +37,10 @@ if TYPE_CHECKING: from vllm.config import CompilationConfig +# Engine-version-specific vLLM imports intentionally stay lazy. Patch targets +# are resolved only after install_hooks selects the active integration, avoiding +# eager imports of optional/version-dependent engine modules. + log = init_logger("vllm.foundry.hooks") _INSTALLED = False @@ -91,6 +102,10 @@ def install_hooks(compilation_config: CompilationConfig) -> None: if not cfg_path: return + # LOAD archive validation is part of config activation and must run before + # any patch installation or device-capability discovery. + load_graph_extension_config(cfg_path) + # Must run BEFORE anything that might import flashinfer. Cheap + safe to # run in every process: when called in the engine child (after fork), # CUDA init is already OK to happen, the patch just keeps returning the @@ -111,7 +126,6 @@ def install_hooks(compilation_config: CompilationConfig) -> None: spawn_elapsed_ms, ) - load_graph_extension_config(cfg_path) log.info( "[foundry] install_hooks: mode=%s workspace=%s", get_graph_extension_mode().value, @@ -412,6 +426,12 @@ def patched(self): result = orig(self) + # Read the live VMM cursor the instant upstream capture_model returns on + # LOAD — before stop_allocation_region — so the emitted watermark is the + # true post-capture allocator offset the harness compares with SAVE2. + if mode == CUDAGraphExtensionMode.LOAD: + rt.log_load_final_alloc_offset() + if mode == CUDAGraphExtensionMode.SAVE: save_graph_manifest() pack_fatbins() @@ -461,9 +481,6 @@ def patched(self, *args, **kwargs): if get_graph_extension_mode() == CUDAGraphExtensionMode.NONE: return orig(self, *args, **kwargs) - from contextlib import ExitStack - from unittest.mock import patch as mock_patch - from vllm.compilation.counter import compilation_counter from vllm.compilation.cuda_graph import CUDAGraphEntry from vllm.compilation.monitor import ( @@ -534,8 +551,6 @@ def patched(self, *args, **kwargs): def _patch_initialize_kv_caches() -> None: - import os - from vllm.v1.engine import core as core_mod cls = core_mod.EngineCore @@ -547,8 +562,6 @@ def patched(self, vllm_config): if mode == CUDAGraphExtensionMode.NONE: return orig(self, vllm_config) - import time - import vllm.envs as envs from vllm.v1.core.kv_cache_utils import ( generate_scheduler_kv_cache_config, @@ -557,22 +570,30 @@ def patched(self, vllm_config): root = get_workspace_root() ws = None + has_saved_kv_profile = False if root and os.path.exists(os.path.join(root, "warmup_state.json")): ws = rt.load_warmup_state(root) - log.info( - "[foundry] %s mode: warmup_state.json found, skipping KV profiling", - mode.value.upper(), - ) + has_saved_kv_profile = bool(ws.available_gpu_memory) + if has_saved_kv_profile: + log.info( + "[foundry] %s mode: warmup_state.json found, skipping KV profiling", + mode.value.upper(), + ) + else: + log.info( + "[foundry] %s mode: transport bootstrap found; running KV profiling", + mode.value.upper(), + ) specs = self.model_executor.get_kv_cache_specs() has_kv = any(s for s in specs) - if ws is not None: + if has_saved_kv_profile: available_mem = ws.available_gpu_memory num_gpu = ws.num_gpu_blocks num_cpu = ws.num_cpu_blocks elif mode == CUDAGraphExtensionMode.LOAD: - raise RuntimeError("foundry LOAD requires warmup_state.json") + raise RuntimeError("foundry LOAD requires a profiled warmup_state.json") elif has_kv: if envs.VLLM_ELASTIC_EP_SCALE_UP_LAUNCH: assert self.available_gpu_memory_for_kv_cache > 0 @@ -597,7 +618,7 @@ def patched(self, vllm_config): args=(vllm_config.model_config.max_model_len,), ) sched = generate_scheduler_kv_cache_config(cfgs) - if ws is None: + if not has_saved_kv_profile: num_gpu = sched.num_blocks elif sched.num_blocks != num_gpu: log.warning( @@ -621,20 +642,25 @@ def patched(self, vllm_config): dp_shard = getattr(vllm_config.parallel_config, "data_parallel_index", 0) or 0 if mode == CUDAGraphExtensionMode.SAVE and root is not None and dp_shard == 0: os.makedirs(root, exist_ok=True) - if ws is not None: - ws.final_alloc_offset = rt.get_final_alloc_offset() - rt.save_warmup_state(root, ws) + if has_saved_kv_profile: + state = ws + else: + state = rt.create_warmup_state() + if ws is not None: + state.deepep_transport = ws.deepep_transport + state.deepep_num_nvl_bytes = ws.deepep_num_nvl_bytes + state.available_gpu_memory = available_mem + state.num_gpu_blocks = num_gpu + state.num_cpu_blocks = num_cpu + state.gpu_memory_utilization = vllm_config.cache_config.gpu_memory_utilization + state.final_alloc_offset = rt.get_final_alloc_offset() + rt.save_warmup_state(root, state) + if has_saved_kv_profile: log.info( - "[foundry] Updated warmup state: final_alloc_offset=%d", ws.final_alloc_offset + "[foundry] Updated warmup state: final_alloc_offset=%d", + state.final_alloc_offset, ) else: - new = rt.create_warmup_state() - new.available_gpu_memory = available_mem - new.num_gpu_blocks = num_gpu - new.num_cpu_blocks = num_cpu - new.gpu_memory_utilization = vllm_config.cache_config.gpu_memory_utilization - new.final_alloc_offset = rt.get_final_alloc_offset() - rt.save_warmup_state(root, new) log.info("[foundry] Saved warmup state: num_gpu_blocks=%d", num_gpu) return sched @@ -675,12 +701,14 @@ def patched_init(self, *args, **kwargs): # --------------------------------------------------------------------------- -# DeepEP buffer kwargs: use_fabric=True, num_nvl_bytes=0 on any foundry mode. +# DeepEP buffer kwargs: apply the archive-defined transport on SAVE/LOAD. # --------------------------------------------------------------------------- def _patch_deepep() -> None: try: + # vLLM is an optional integration dependency, so discover this + # version-specific module only while installing its patch. from vllm.distributed.device_communicators import all2all except Exception: return @@ -694,13 +722,37 @@ def _patch_deepep() -> None: @functools.wraps(orig) def patched(self, *args, **kwargs): out = orig(self, *args, **kwargs) - if isinstance(out, dict) and get_graph_extension_mode() != CUDAGraphExtensionMode.NONE: - # Force VMM-fabric buffers - # so the foundry cuMem hook tracks them, and disable the - # NVLink buffer (LL-mode uses RDMA) to avoid fabric-cuMemCreate - # collisions with preallocated region. - out["use_fabric"] = True - out["num_nvl_bytes"] = 0 + mode = get_graph_extension_mode() + if not isinstance(out, dict) or mode == CUDAGraphExtensionMode.NONE: + return out + + transport = get_deepep_transport() + upstream_num_nvl_bytes = out.get("num_nvl_bytes", 0) + effective_num_nvl_bytes, configured = configure_deepep_buffer( + transport, + num_nvl_bytes=upstream_num_nvl_bytes, + kwargs=out, + ) + out.clear() + out.update(configured) + out["num_nvl_bytes"] = effective_num_nvl_bytes + + root = get_workspace_root() + if root is None: + raise RuntimeError("Foundry vLLM DeepEP requires workspace_root") + if mode == CUDAGraphExtensionMode.SAVE: + rt.update_warmup_state_deepep_profile( + root, + transport, + effective_num_nvl_bytes, + ) + elif mode == CUDAGraphExtensionMode.LOAD: + state = rt.load_warmup_state(root) + validate_num_nvl_bytes_match( + archived=state.deepep_num_nvl_bytes, + effective=effective_num_nvl_bytes, + transport=transport, + ) return out cls._make_all2all_kwargs = patched diff --git a/python/foundry/integration/vllm/restore_progress.py b/python/foundry/integration/vllm/restore_progress.py new file mode 100644 index 00000000..5aaeb83b --- /dev/null +++ b/python/foundry/integration/vllm/restore_progress.py @@ -0,0 +1,48 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the Foundry project +"""Unique per-rank LOAD graph restoration progress.""" + +from __future__ import annotations + +from pathlib import Path +from typing import NamedTuple + + +class RestorationUpdate(NamedTuple): + count: int + is_new: bool + completed_now: bool + + +def rank_from_workspace_dir(workspace_dir: str) -> int: + rank_dir = Path(workspace_dir).name + prefix = "rank_" + if not rank_dir.startswith(prefix) or not rank_dir.removeprefix(prefix).isdigit(): + raise ValueError(f"LOAD workspace must end in rank_N, got {workspace_dir!r}") + return int(rank_dir.removeprefix(prefix)) + + +class LoadRestorationProgress: + def __init__(self, rank: int, total: int) -> None: + if total <= 0: + raise ValueError("LOAD restoration total must be positive") + self.rank = rank + self.total = total + self._restored_indices: set[int] = set() + self._completion_reported = False + + @property + def restored_indices(self) -> frozenset[int]: + return frozenset(self._restored_indices) + + def record(self, index: int) -> RestorationUpdate: + if not 0 <= index < self.total: + raise IndexError(f"LOAD graph index {index} outside [0, {self.total})") + previous_count = len(self._restored_indices) + self._restored_indices.add(index) + count = len(self._restored_indices) + is_new = count != previous_count + completed_now = is_new and count == self.total and not self._completion_reported + if completed_now: + self._completion_reported = True + return RestorationUpdate(count, is_new, completed_now) diff --git a/python/foundry/integration/vllm/runtime.py b/python/foundry/integration/vllm/runtime.py index c81e64cb..56c36f00 100644 --- a/python/foundry/integration/vllm/runtime.py +++ b/python/foundry/integration/vllm/runtime.py @@ -13,7 +13,9 @@ import json import os import shutil +import tempfile import time +from contextlib import suppress from dataclasses import asdict, dataclass, field from datetime import datetime from pathlib import Path @@ -25,6 +27,12 @@ from foundry import ops as fops from foundry.allocation_region import parse_size +from foundry.integration.deepep_transport import ( + DeepEPTransport, + validate_archived_deepep_profile, + validate_num_nvl_bytes_value, + validate_transport_match, +) from foundry.integration.vllm.config import ( CUDAGraphExtensionMode, _compute_rank_from_parallel_config, @@ -33,6 +41,7 @@ get_hook_library_path, get_nvshmem_host_path, ) +from foundry.integration.vllm.restore_progress import rank_from_workspace_dir if TYPE_CHECKING: from vllm.config import ParallelConfig @@ -57,6 +66,8 @@ class WarmupState: num_cpu_blocks: int = 0 gpu_memory_utilization: float = 0.0 final_alloc_offset: int = 0 + deepep_transport: str = DeepEPTransport.FABRIC.value + deepep_num_nvl_bytes: int = 0 _final_alloc_offset: int = 0 @@ -83,9 +94,26 @@ def create_warmup_state() -> WarmupState: def save_warmup_state(workspace_root: str, state: WarmupState) -> None: + os.makedirs(workspace_root, exist_ok=True) path = os.path.join(workspace_root, "warmup_state.json") - with open(path, "w") as f: - json.dump(asdict(state), f, indent=2) + temp_path = None + try: + with tempfile.NamedTemporaryFile( + mode="w", + encoding="utf-8", + dir=workspace_root, + prefix=f".warmup_state.{os.getpid()}.", + suffix=".tmp", + delete=False, + ) as f: + temp_path = f.name + json.dump(asdict(state), f, indent=2) + os.replace(temp_path, path) + temp_path = None + finally: + if temp_path is not None: + with suppress(FileNotFoundError): + os.unlink(temp_path) log.info("Saved warmup state to %s", path) @@ -93,10 +121,41 @@ def load_warmup_state(workspace_root: str) -> WarmupState: path = os.path.join(workspace_root, "warmup_state.json") if not os.path.exists(path): raise RuntimeError(f"Warmup state file not found: {path}") - with open(path) as f: - data = json.load(f) + try: + with open(path) as f: + data = json.load(f) + except json.JSONDecodeError as exc: + raise RuntimeError(f"Foundry warmup state contains invalid JSON at {path}: {exc}") from exc + if not isinstance(data, dict): + raise RuntimeError( + "Foundry incompatible warmup state at " + f"{path}: expected a JSON object, got {type(data).__name__}" + ) + transport, num_nvl_bytes = validate_archived_deepep_profile(data, path=path) valid = set(WarmupState.__dataclass_fields__.keys()) - return WarmupState(**{k: v for k, v in data.items() if k in valid}) + state_data = {k: v for k, v in data.items() if k in valid} + state_data["deepep_transport"] = transport.value + state_data["deepep_num_nvl_bytes"] = num_nvl_bytes + return WarmupState(**state_data) + + +def update_warmup_state_deepep_profile( + workspace_root: str, + transport: DeepEPTransport, + num_nvl_bytes: int, +) -> None: + validated_num_nvl_bytes = validate_num_nvl_bytes_value( + num_nvl_bytes, + transport=transport, + ) + state = get_state() + if state is None or state.rank != 0: + return + path = os.path.join(workspace_root, "warmup_state.json") + warmup_state = load_warmup_state(workspace_root) if os.path.exists(path) else WarmupState() + warmup_state.deepep_transport = transport.value + warmup_state.deepep_num_nvl_bytes = validated_num_nvl_bytes + save_warmup_state(workspace_root, warmup_state) # --------------------------------------------------------------------------- @@ -107,6 +166,7 @@ def load_warmup_state(workspace_root: str) -> WarmupState: @dataclass class CUDAGraphExtensionState: capture_index: int = 0 + rank: int = 0 # LOAD: strong-reference table keyed by GraphIdentifier. The upstream # CUDAGraphWrapper stores `entry.output = weak_ref_tensors(output)`, so # without a strong ref here the VMM-backed output tensors returned by @@ -147,7 +207,6 @@ def setup_graph_extension(parallel_config: ParallelConfig) -> None: workspace_dir.mkdir(parents=True, exist_ok=True) elif cfg.mode == CUDAGraphExtensionMode.LOAD: - fops.set_skip_fatbin_processing(True) if not workspace_dir.exists(): # rank_0 fallback is only safe for single-rank setups. # Under TP/PP/DP > 1 each rank has its own device addresses, @@ -173,6 +232,12 @@ def setup_graph_extension(parallel_config: ParallelConfig) -> None: f"exist. For multi-rank LOAD, re-run SAVE with the same " f"parallelism topology." ) + warmup_state = load_warmup_state(cfg.workspace_root) + validate_transport_match( + archived=warmup_state.deepep_transport, + configured=cfg.deepep_transport, + ) + fops.set_skip_fatbin_processing(True) log.info("[foundry] LOAD: loading CUDA modules from %s", cfg.workspace_dir) t = time.perf_counter() fops.load_cuda_modules_and_libraries(cfg.workspace_dir) @@ -189,7 +254,7 @@ def setup_graph_extension(parallel_config: ParallelConfig) -> None: # # Bring cuBLAS handle up so its tiny bring-up lands in scratch. # _ = torch._C._cuda_getCurrentBlasHandle() - _state = CUDAGraphExtensionState() + _state = CUDAGraphExtensionState(rank=rank) log.info("[foundry TIMING] setup_graph_extension total: %.3f s", time.perf_counter() - t_total) @@ -227,6 +292,25 @@ def capture_final_alloc_offset() -> int: return offset +def log_load_final_alloc_offset() -> int | None: + """Emit the live per-rank VMM cursor immediately after LOAD capture_model. + + The last per-graph restoration marker can precede post-loop allocator + movement, so the harness needs the true terminal watermark read straight + from the allocator (not a reread of the archived ``final_alloc_offset.json``) + to prove the LOAD trajectory lands exactly on SAVE2's rank offset. The rank + is derived from the configured ``rank_N`` workspace so each LOAD emits + exactly one marker per rank. + """ + cfg = get_config() + if cfg is None or cfg.mode != CUDAGraphExtensionMode.LOAD or cfg.workspace_dir is None: + return None + rank = rank_from_workspace_dir(cfg.workspace_dir) + offset = fops.get_current_alloc_offset() + log.info("[foundry] LOAD live final_alloc_offset rank_%d: %d", rank, offset) + return offset + + # --------------------------------------------------------------------------- # VMM preallocation (LOAD only) # --------------------------------------------------------------------------- diff --git a/python/foundry_sglang_plugin/__init__.py b/python/foundry_sglang_plugin/__init__.py new file mode 100644 index 00000000..3a0c1323 --- /dev/null +++ b/python/foundry_sglang_plugin/__init__.py @@ -0,0 +1,33 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the Foundry project +"""Lightweight SGLang entry point for the Foundry CUDA graph extension. + +SGLang imports this standalone package while discovering plugins in every +process. Module import is stdlib-only so an installed but inactive Foundry +plugin does not import Foundry, PyTorch, or native operations. +""" + +from __future__ import annotations + +import os + + +def activate() -> None: + """Load Foundry only when an explicit graph-extension config is selected.""" + cfg_path = os.environ.get("FOUNDRY_GRAPH_EXTENSION_CONFIG_PATH") + if not cfg_path: + return + + # Lazy import: Foundry loads compiled/optional dependencies and is required + # only for explicitly configured SAVE/LOAD processes. + from foundry.integration.sglang.config import load_graph_extension_config + + load_graph_extension_config(cfg_path) + _register_hooks() + + +def _register_hooks() -> None: + # Lazy import: hooks transitively load torch/foundry.ops after config is set. + from foundry.integration.sglang.hooks import register_hooks + + register_hooks() diff --git a/python/foundry_sglang_plugin/targets.py b/python/foundry_sglang_plugin/targets.py new file mode 100644 index 00000000..6709d5c9 --- /dev/null +++ b/python/foundry_sglang_plugin/targets.py @@ -0,0 +1,50 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the Foundry project +"""Hook targets at SGLang a25164b, shared by registration and validation.""" + +import inspect + +SERVER_ARGS_POST_INIT = "sglang.srt.server_args.ServerArgs.__post_init__" +LAUNCH_SCHEDULERS = "sglang.srt.entrypoints.engine.Engine._launch_scheduler_processes" +INIT_TORCH_DISTRIBUTED = "sglang.srt.model_executor.model_runner.ModelRunner.init_torch_distributed" +ALLOC_MEMORY_POOL = "sglang.srt.model_executor.model_runner.ModelRunner.alloc_memory_pool" +RESOLVE_MEMORY_POOL_CONFIG = ( + "sglang.srt.mem_cache.kv_cache_configurator.KVCacheConfigurator._resolve_memory_pool_config" +) +FULL_CUDA_GRAPH_BACKEND = ( + "sglang.srt.model_executor.runner_backend.full_cuda_graph_backend.FullCudaGraphBackend" +) +DECODE_CUDA_GRAPH_CAPTURE = ( + "sglang.srt.model_executor.runner.decode_cuda_graph_runner.DecodeCudaGraphRunner.capture" +) +TORCH_SYMM_MEM_INIT = ( + "sglang.srt.distributed.device_communicators.torch_symm_mem.TorchSymmMemCommunicator.__init__" +) +MULTIMEM_ALL_GATHER_CALL = ( + "sglang.srt.distributed.device_communicators.triton_symm_mem_ag.MultimemAllGatherer.__call__" +) + +HOOK_TARGETS = ( + SERVER_ARGS_POST_INIT, + LAUNCH_SCHEDULERS, + INIT_TORCH_DISTRIBUTED, + ALLOC_MEMORY_POOL, + RESOLVE_MEMORY_POOL_CONFIG, + FULL_CUDA_GRAPH_BACKEND, + DECODE_CUDA_GRAPH_CAPTURE, + TORCH_SYMM_MEM_INIT, + MULTIMEM_ALL_GATHER_CALL, +) + + +def normalize_resolved_identity(resolved): + """Return stable identity for targets resolved through Python descriptors. + + Pinned ``Engine._launch_scheduler_processes`` is a ``@classmethod``. + Resolving it repeatedly creates distinct bound-method objects whose stable + callable identity is their shared ``__func__``. Classes and ordinary + functions/method descriptors are returned unchanged. + """ + if inspect.ismethod(resolved): + return resolved.__func__ + return resolved diff --git a/recipe/experimental/README.md b/recipe/experimental/README.md new file mode 100644 index 00000000..ad6506d3 --- /dev/null +++ b/recipe/experimental/README.md @@ -0,0 +1,216 @@ +# Foundry recipe — experimental multi-GPU IPC paths + +This directory contains vLLM expert-parallel recipes for hosts without NVIDIA +IMEX and the separately verified SGLang dense tensor-parallel recipe. The vLLM +SAVE and LOAD TOMLs select the archive-defined intranode transport: + +```toml +deepep_transport = "nvl_ipc" +``` + +DeepEP's nonzero NVLink buffer then uses Foundry's VMM-IPC translation: +the hook exports a POSIX file descriptor over `SCM_RIGHTS` and maps it into the +peer rank. DeepEP and NVSHMEM are still required because the RDMA buffer remains +on the NVSHMEM symmetric heap. + +Relevant files: + +```text +recipe/experimental/ +├── foundry_save.toml +├── foundry_load.toml +├── foundry_save_fp8.toml +├── foundry_load_fp8.toml +├── serve_qwen3-30ba3b_ipc_ep.sh +├── serve_qwen3-30ba3bfp8_ipc_ep.sh +├── sglang_foundry_tp_save.toml +├── sglang_foundry_tp_load.toml +└── serve_qwen3-8b_sglang_tp.sh +``` + +Read [`../vllm/README.md`](../vllm/README.md) for installation, the two-pass +SAVE workflow, archive layout, and shared EP flags. Build this Foundry revision +so `libcuda_hook.so` includes the SCM_RIGHTS VMM-IPC transport: + +```bash +CC=gcc CXX=g++ pip install -e . --no-build-isolation +``` + +Foundry auto-detects `libnvshmem_host.so` from the installed NVSHMEM wheel. +Set `nvshmem_host_path` in both matching TOMLs only for a custom build. + +## vLLM workflow + +Run every phase from the same repository path so vLLM sees an identical +`graph_extension_config_path` and reuses its compile cache: + +```bash +rm -rf foundry_archive_ipc +bash recipe/experimental/serve_qwen3-30ba3b_ipc_ep.sh 2 --save +bash recipe/experimental/serve_qwen3-30ba3b_ipc_ep.sh 2 --save +bash recipe/experimental/serve_qwen3-30ba3b_ipc_ep.sh 2 --load +``` + +Use `serve_qwen3-30ba3bfp8_ipc_ep.sh` for the FP8/DeepGEMM model; it uses the +separate `foundry_archive_ipc_fp8` archive. Stop each SAVE server after startup +and leave LOAD running for inference. The scripts preserve vLLM V1 selection +and disable NCCL CUMEM and NVLS in every Foundry phase. The transport itself is +selected only by the SAVE/LOAD TOML and must match the archive. + +## Quarantined SGLang EP profile + +The old SGLang DeepEP IPC recipe is not runnable on the current HookRegistry +integration. The verified SGLang profile remains dense Qwen3-8B TP=2 with +torch symmetric memory; EP and DeepEP are rejected before launch. Reintroduce +an EP recipe only with a separate current-main profile, HookRegistry target, +archive contract, and fresh SAVE-to-LOAD H100 evidence. + +## SGLang dense tensor parallel + +`serve_qwen3-8b_sglang_tp.sh` runs dense Qwen3-8B at exactly TP=2 +(`--tp-size 2`, no DP-attention, no expert parallel) — the first recipe here +that captures a **cross-rank collective into the graph itself**. The script +rejects any other TP size, extra argument, or mode other than `--save` / +`--load`. This profile specifically exercises a per-layer TP all-reduce; +replicated DP and expert-sharded workloads do not. + +```bash +rm -rf foundry_archive_sglang_tp +CUDA_VISIBLE_DEVICES=0,1 bash serve_qwen3-8b_sglang_tp.sh 2 --save +CUDA_VISIBLE_DEVICES=0,1 bash serve_qwen3-8b_sglang_tp.sh 2 --load +curl -s http://0.0.0.0:12000/v1/completions \ + -H 'Content-Type: application/json' \ + -d '{"model":"Qwen/Qwen3-8B","prompt":"The capital of France is","max_tokens":12,"temperature":0}' +``` + +Activation is through SGLang's `sglang.srt.plugins` entry point with zero +upstream edits. The lightweight `foundry_sglang_plugin` entry point installs +nine named `HookRegistry` targets. The script exports `SGLANG_PLUGINS=foundry` +in every phase, selects the SAVE/LOAD TOML through +`FOUNDRY_GRAPH_EXTENSION_CONFIG_PATH` (unset for baseline), and passes model +revision `b968826d9c46dd6066d109eabc6255188de91218` via `--revision`. + +The exact validated pins are: + +```text +Foundry a8e43442421a80fea7991ee26f97f1003b170e29 +SGLang a25164bda355250015c0ba6d8821f0d4b37abe7d +Qwen3-8B b968826d9c46dd6066d109eabc6255188de91218 +``` + +Shell recipe users must install the matching SGLang pin; the serve script does +not install or pin SGLang. The exact engine/model SHAs are validation +provenance enforced by the Modal harness, which clones and verifies the engine +commit and passes the model revision. `_validate_profile` covers topology, +backend, and settings; it does not inspect the SHAs, so arbitrary installed +SGLang revisions are not rejected at startup. + +How the validated collective survives SAVE/LOAD: + +- `--enable-torch-symm-mem` creates the TP communicator. Its constructor runs + while Foundry allocation tracking is stopped, keeping PyTorch's symmetric + allocation/rendezvous outside Foundry VMM, and the communicator is retained. +- The gather hook bypasses upstream lazy state and calls + `torch.ops.symm_mem.multimem_all_gather_out` directly into a stable retained + buffer slice using the communicator's original group name. +- `--disable-custom-all-reduce` and + `--enforce-disable-flashinfer-allreduce-fusion` prevent unsupported fallback. +- The exact graph config captures one padded full-decode shape: + `{"decode":{"backend":"full","max_bs":256,"bs":[256]},"prefill":{"backend":"disabled"}}`. +- LOAD runs SGLang's same upstream capture lifecycle. The backend preallocates, + starts one build with the shared pool, and finishes the archived shape when + upstream reaches `capture_one`; it performs no native capture or resave. +- `NCCL_CUMEM_ENABLE=0` and `NCCL_NVLS_ENABLE=0` remain profile invariants. + NCCL version/transport logs are diagnostic, not correctness evidence for the + symmetric collectives. + +No NVSHMEM is involved (dense, no DeepEP), so the TP TOMLs leave +`nvshmem_host_path` unset. The script pins `--random-seed 42` by default +(`SGL_RANDOM_SEED` overrides it), making SAVE-pass metadata directly +comparable. + +This profile excludes custom, NCCL-symmetric, and fused collectives; NCCL +CUMEM and NVLS; TP>2; multi-host; unpinned engine heads; EP; DP-attention; +speculative decode; LoRA; pdmux; Ampere; Blackwell; and hybrid MoE. Those paths +are unsupported or unvalidated, not expected extensions of this result. + +## Validation status + +Historical branch evidence for the separate Qwen3-30B-A3B EP=2 IPC recipe +covers 2×H200 without IMEX: per-rank `final_alloc_offset` was identical across +passes, LOAD replayed at the saved offset, and queries returned coherent +completions. This is not evidence for dense TP or the current combined head. + +Historical **pre-port fork** evidence for the dense TP recipe (the monolithic +`CudaGraphRunner` patch + `foundry_shim.py` activation, not this pinned-main +plugin architecture): + +- Modal-projects draft [PR #4](https://github.com/modal-projects/foundry/pull/4) + is review context; immutable + [commit `21bf734`](https://github.com/modal-projects/foundry/commit/21bf7341847ebce46c534e7f6e86a3929275c410) + recorded the 2×H100 run where two SAVE passes wrote 36 graphs per rank at + `final_alloc_offset=69736595456`; LOAD restored all graphs at that offset, and + four temperature-0 baseline/LOAD responses matched byte for byte. +- After the SGLang #26735 metadata API port, modal-projects draft + [PR #2](https://github.com/modal-projects/foundry/pull/2) is review context; + immutable + [commit `ee7d61a`](https://github.com/modal-projects/foundry/commit/ee7d61a32a4d4732e5f14978975b96ab03bcbede) + recorded the 2×H200 per-rank `final_alloc_offset=121831948288` and exact + SAVE/LOAD equality on 16/16 prompts. + +These records are pre-port context only. They do not broaden current support. + +Plain-NCCL fresh-process restore was disproven by +[`ap-2qSQUKlY24L4RrQkMpUrmc`](https://modal.com/apps/modal-labs/rahul-dev/ap-2qSQUKlY24L4RrQkMpUrmc). +All 36 graphs per rank restored with matching counts, offsets, immutable bytes, +and P2P/IPC topology, but both fresh LOADs produced the same corrupt decode +while baseline and SAVE2 matched. Archived process-specific communicator and +graph-sequence state is the leading forensic explanation, not +treatment-proven causality. The observation rejects plain-NCCL restore; +structural restoration alone was not output correctness. + +The focused symmetric SAVE→LOAD treatment is `VERIFIED`: +[`ap-5wO9cFhRZ7ITsLDDv27Xfh`](https://modal.com/apps/modal-labs/rahul-dev/ap-5wO9cFhRZ7ITsLDDv27Xfh). +The full baseline → SAVE → SAVE2 → LOAD → LOAD2 matrix is `VERIFIED`: +[`ap-6fYcSUFLOxocBSjNxakNGZ`](https://modal.com/apps/modal-labs/rahul-dev/ap-6fYcSUFLOxocBSjNxakNGZ). + +```bash +modal run --env tests/modal_sglang_tp.py +``` + +Final evidence is one graph per rank at offset `68908220416`; 73 symmetric +all-reduce plus one symmetric all-gather and zero NCCL collective nodes per +rank; exact sequential and eight-slot batch-8 outputs; SAVE/SAVE2 semantic +reproducibility; LOAD/LOAD2 byte immutability; observed batch-1 and batch-8 +replay; time to health 217.2/103.0/103.1/61.0/61.0 seconds; all phase processes +clean; and cleanup complete with no retained archive. + +This verifies only the exact pinned, single-host Qwen3-8B TP=2 symmetric +profile. It is not general/upstream TP or plain-NCCL support. Track official +support in [issue #6](https://github.com/foundry-org/foundry/issues/6). + +## Not supported: Qwen3.5-122B-A10B hybrid MoE + +`Qwen/Qwen3.5-122B-A10B` is not a larger instance of the dense recipe. It +combines Mamba/`Qwen3_5GatedDeltaNet` linear attention, MoE, and multimodal +components. Historical investigation found three independent blockers: + +- lazy torch.compile work during the first captured forward can leave SAVE + with `Graph contains no nodes`; +- FlashInfer fused all-reduce bypassed the **pre-treatment plain-NCCL + experiment**; and +- Mamba SSM state plus KV and graph preallocation exhausted the tested H200 + configuration during LOAD. + +Supporting it requires capture-safe warmup that preserves SAVE/LOAD allocation +identity, explicit Mamba-state handling, and a supported fused-collective path. +It is outside this experimental dense-TP profile. + +## IPC-specific troubleshooting + +| Symptom | Likely cause | +|---|---| +| `[HOOK] ERROR: cuMemImportFromShareableHandle failed with error 999` at `deep_ep.cpp` `Buffer::sync` | Foundry hook not rebuilt with the SCM_RIGHTS transport — it's still packing a raw fd. Rebuild `foundry/` (`uv pip install -e . --no-build-isolation`). | +| `operation failed due to a previous error during capture` on SAVE **pass 2** | vLLM compile-cache over-keying not applied (`graph_extension_config_path` still hashed) → recompile-in-capture. Apply the `compilation.py` `ignored_factors` edit. | +| LOAD `illegal memory access` at replay with an NVL buffer present | Relocated peer import collided with the NVSHMEM heap hint — ensure the hook build includes the dedicated import-VA zone (`0x300000000000`). | +| `error 999` only with `--deepep-mode auto`/`normal` | That's expected for non-LL modes here; this recipe pins `deepep_low_latency`. | diff --git a/recipe/experimental/foundry_load.toml b/recipe/experimental/foundry_load.toml new file mode 100644 index 00000000..1d8bbdbb --- /dev/null +++ b/recipe/experimental/foundry_load.toml @@ -0,0 +1,10 @@ +mode = "load" +base_addr = 0x600000000000 +region_size = "256GB" +workspace_root = "foundry_archive_ipc" +scratch_space_size = "4096MB" +deepep_transport = "nvl_ipc" +# REQUIRED for the EP/IPC recipe — DeepEP low-latency still allocates its RDMA +# buffer on the NVSHMEM symmetric heap (the NVL buffer is the cudaIpc path). +# Point this at the libnvshmem_host.so built by vllm/tools/ep_kernels. +# nvshmem_host_path = "/absolute/path/to/libnvshmem_host.so" diff --git a/recipe/experimental/foundry_load_fp8.toml b/recipe/experimental/foundry_load_fp8.toml new file mode 100644 index 00000000..4edd4e9d --- /dev/null +++ b/recipe/experimental/foundry_load_fp8.toml @@ -0,0 +1,8 @@ +mode = "load" +base_addr = 0x600000000000 +region_size = "256GB" +workspace_root = "foundry_archive_ipc_fp8" +scratch_space_size = "4096MB" +deepep_transport = "nvl_ipc" +# Point this at the libnvshmem_host.so built by vllm/tools/ep_kernels. +# nvshmem_host_path = "/absolute/path/to/libnvshmem_host.so" diff --git a/recipe/experimental/foundry_save.toml b/recipe/experimental/foundry_save.toml new file mode 100644 index 00000000..81c48c4a --- /dev/null +++ b/recipe/experimental/foundry_save.toml @@ -0,0 +1,10 @@ +mode = "save" +base_addr = 0x600000000000 +region_size = "256GB" +workspace_root = "foundry_archive_ipc" +scratch_space_size = "4096MB" +deepep_transport = "nvl_ipc" +# REQUIRED for the EP/IPC recipe — DeepEP low-latency still allocates its RDMA +# buffer on the NVSHMEM symmetric heap (the NVL buffer is the cudaIpc path). +# Point this at the libnvshmem_host.so built by vllm/tools/ep_kernels. +# nvshmem_host_path = "/absolute/path/to/libnvshmem_host.so" diff --git a/recipe/experimental/foundry_save_fp8.toml b/recipe/experimental/foundry_save_fp8.toml new file mode 100644 index 00000000..31b115b3 --- /dev/null +++ b/recipe/experimental/foundry_save_fp8.toml @@ -0,0 +1,8 @@ +mode = "save" +base_addr = 0x600000000000 +region_size = "256GB" +workspace_root = "foundry_archive_ipc_fp8" +scratch_space_size = "4096MB" +deepep_transport = "nvl_ipc" +# Point this at the libnvshmem_host.so built by vllm/tools/ep_kernels. +# nvshmem_host_path = "/absolute/path/to/libnvshmem_host.so" diff --git a/recipe/experimental/serve_qwen3-30ba3b_ipc_ep.sh b/recipe/experimental/serve_qwen3-30ba3b_ipc_ep.sh new file mode 100755 index 00000000..d20997ce --- /dev/null +++ b/recipe/experimental/serve_qwen3-30ba3b_ipc_ep.sh @@ -0,0 +1,80 @@ +#!/bin/bash +# Usage: bash serve_qwen3-30ba3b_ipc_ep.sh [--save|--load] +# +# EXPERIMENTAL: DeepEP expert-parallel with the legacy CUDA-IPC NVLink buffer +# selected by deepep_transport="nvl_ipc" in the matching SAVE/LOAD TOML. +# This exercises foundry's VMM-IPC translation layer (SCM_RIGHTS fd transport +# and whole-chunk peer mapping) on machines without fabric/IMEX. +# +# Build requirements and the canonical-path workflow are documented in +# README.md#vllm-workflow. +SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" + +EP_SIZE=${1:?Usage: $0 [--save|--load]} +MODEL_NAME="Qwen/Qwen3-30B-A3B" +HOST="0.0.0.0" +PORT=12000 +GPU_MEMORY_UTILIZATION=0.8 + +FOUNDRY_ARGS=() +if [[ "$2" == "--save" ]]; then + FOUNDRY_ARGS+=( --compilation-config.graph_extension_config_path "${SCRIPT_DIR}/foundry_save.toml" ) + # Foundry pins NCCL to the IPC/ring transports — the CUMEM P2P and NVLS + # multicast fast paths cuMemMap with driver-capability flags the foundry + # VMM region doesn't carry. + export NCCL_CUMEM_ENABLE=0 + export NCCL_NVLS_ENABLE=0 + # Foundry only patches the V1 model runner; pin V1 explicitly so vLLM + # doesn't quietly route Qwen3ForCausalLM-class architectures to V2. + export VLLM_USE_V2_MODEL_RUNNER=0 + echo "Using foundry SAVE (DeepEP NVL/IPC): ${SCRIPT_DIR}/foundry_save.toml" +elif [[ "$2" == "--load" ]]; then + FOUNDRY_ARGS+=( --compilation-config.graph_extension_config_path "${SCRIPT_DIR}/foundry_load.toml" ) + export NCCL_CUMEM_ENABLE=0 + export NCCL_NVLS_ENABLE=0 + export VLLM_USE_V2_MODEL_RUNNER=0 + echo "Using foundry LOAD (DeepEP NVL/IPC): ${SCRIPT_DIR}/foundry_load.toml" +elif [[ -n "$2" ]]; then + echo "Usage: $0 [--save|--load]" + exit 1 +else + echo "Running without foundry (baseline vLLM)" +fi + +# LD_PRELOAD of libcuda_hook.so / libnvshmem_host.so is set by foundry's +# setup_ld_preload_env at worker spawn time (uses the paths in the TOML +# config). Baseline runs don't need either preloaded by the shell. + +# Foundry only patches the V1 model runner. vLLM defaults certain +# architectures (e.g. Qwen3ForCausalLM) to the V2 runner, which our +# patches don't touch — pin V1 here. +export VLLM_USE_V2_MODEL_RUNNER=0 + +export VLLM_USE_FLASHINFER_SAMPLER=1 +export VLLM_DISABLE_SHARED_EXPERTS_STREAM=1 + +CUDAGRAPH_CAPTURE_SIZES=($(seq 1 256)) + +ARGS=( + --trust-remote-code + --host "$HOST" + --port "$PORT" + --tensor-parallel-size 1 + --data-parallel-size "$EP_SIZE" + --gpu-memory-utilization "$GPU_MEMORY_UTILIZATION" + --distributed-executor-backend uni + --enable-expert-parallel + --all2all-backend deepep_low_latency + --no-enable-prefix-caching + # max-num-batched-tokens == max-num-seqs satisfies vLLM's batched>=seqs check + # and keeps DeepEP's (tokens+1)*2 <= NVSHMEM_QP_DEPTH(=1024) — no QP override. + --max-num-batched-tokens 256 + --max-num-seqs 256 + --attention-config.backend FLASH_ATTN + --compilation-config.cudagraph_mode FULL_DECODE_ONLY + --compilation-config.cudagraph_num_of_warmups 0 + --chat-template-content-format string + --cudagraph-capture-sizes ${CUDAGRAPH_CAPTURE_SIZES[@]} +) + +vllm serve "$MODEL_NAME" "${ARGS[@]}" "${FOUNDRY_ARGS[@]}" diff --git a/recipe/experimental/serve_qwen3-30ba3bfp8_ipc_ep.sh b/recipe/experimental/serve_qwen3-30ba3bfp8_ipc_ep.sh new file mode 100755 index 00000000..bc981a31 --- /dev/null +++ b/recipe/experimental/serve_qwen3-30ba3bfp8_ipc_ep.sh @@ -0,0 +1,68 @@ +#!/bin/bash +# Usage: bash serve_qwen3-30ba3bfp8_ipc_ep.sh [--save|--load] +# +# EXPERIMENTAL: FP8 Qwen3-30B-A3B with DeepGEMM MoE + DeepEP expert-parallel, +# selecting the legacy CUDA-IPC NVLink buffer in the matching SAVE/LOAD TOML. +# Same as serve_qwen3-30ba3b_ipc_ep.sh but: FP8 model + VLLM_USE_DEEP_GEMM=1 +# (blockscale FP8 MoE), and its own archive (foundry_archive_ipc_fp8). +# +# Run every phase from one consistent path (or absolute path) so SAVE pass 1/2 +# and LOAD pass the identical graph_extension_config_path; see +# README.md#vllm-workflow. +SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" + +EP_SIZE=${1:?Usage: $0 [--save|--load]} +MODEL_NAME="Qwen/Qwen3-30B-A3B-FP8" +HOST="0.0.0.0" +PORT=12000 +GPU_MEMORY_UTILIZATION=0.8 + +FOUNDRY_ARGS=() +if [[ "$2" == "--save" ]]; then + FOUNDRY_ARGS+=( --compilation-config.graph_extension_config_path "${SCRIPT_DIR}/foundry_save_fp8.toml" ) + export NCCL_CUMEM_ENABLE=0 + export NCCL_NVLS_ENABLE=0 + export VLLM_USE_V2_MODEL_RUNNER=0 + echo "Using foundry SAVE (FP8/DeepGEMM, DeepEP NVL/IPC): ${SCRIPT_DIR}/foundry_save_fp8.toml" +elif [[ "$2" == "--load" ]]; then + FOUNDRY_ARGS+=( --compilation-config.graph_extension_config_path "${SCRIPT_DIR}/foundry_load_fp8.toml" ) + export NCCL_CUMEM_ENABLE=0 + export NCCL_NVLS_ENABLE=0 + export VLLM_USE_V2_MODEL_RUNNER=0 + echo "Using foundry LOAD (FP8/DeepGEMM, DeepEP NVL/IPC): ${SCRIPT_DIR}/foundry_load_fp8.toml" +elif [[ -n "$2" ]]; then + echo "Usage: $0 [--save|--load]" + exit 1 +else + echo "Running without foundry (baseline vLLM)" +fi + +export VLLM_USE_V2_MODEL_RUNNER=0 +export VLLM_USE_FLASHINFER_SAMPLER=1 +export VLLM_DISABLE_SHARED_EXPERTS_STREAM=1 +# FP8 blockscale MoE via DeepGEMM. +export VLLM_USE_DEEP_GEMM=1 + +CUDAGRAPH_CAPTURE_SIZES=($(seq 1 256)) + +ARGS=( + --trust-remote-code + --host "$HOST" + --port "$PORT" + --tensor-parallel-size 1 + --data-parallel-size "$EP_SIZE" + --gpu-memory-utilization "$GPU_MEMORY_UTILIZATION" + --distributed-executor-backend uni + --enable-expert-parallel + --all2all-backend deepep_low_latency + --no-enable-prefix-caching + --max-num-batched-tokens 256 + --max-num-seqs 256 + --attention-config.backend FLASH_ATTN + --compilation-config.cudagraph_mode FULL_DECODE_ONLY + --compilation-config.cudagraph_num_of_warmups 0 + --chat-template-content-format string + --cudagraph-capture-sizes ${CUDAGRAPH_CAPTURE_SIZES[@]} +) + +vllm serve "$MODEL_NAME" "${ARGS[@]}" "${FOUNDRY_ARGS[@]}" diff --git a/recipe/experimental/serve_qwen3-8b_sglang_tp.sh b/recipe/experimental/serve_qwen3-8b_sglang_tp.sh new file mode 100755 index 00000000..d5a6fca8 --- /dev/null +++ b/recipe/experimental/serve_qwen3-8b_sglang_tp.sh @@ -0,0 +1,78 @@ +#!/bin/bash +# Usage: CUDA_VISIBLE_DEVICES=0,1 bash serve_qwen3-8b_sglang_tp.sh [--save|--load] +# +# EXPERIMENTAL: SGLang dense tensor parallel (tp-size > 1, no DP-attention, no +# expert parallel). Unlike the DP recipe — where each rank is an independent +# replica with no cross-rank collective inside the captured graph — pure TP +# captures a per-layer collective across the TP ranks into every CUDA graph. +# +# That collective is served by torch symmetric memory (--enable-torch-symm-mem): +# PyTorch rendezvouses one persistent symmetric buffer per rank outside Foundry's +# monotonic VMM region, and the captured graph replays a direct +# multimem_all_gather_out over that stable buffer/group. Custom all-reduce and +# FlashInfer fused all-reduce are disabled so neither can win before torch +# symmetric memory. NCCL_CUMEM_ENABLE / NCCL_NVLS_ENABLE stay off so any NCCL +# fallback used only for deterministic initialization keeps a reproducible layout. +set -euo pipefail + +SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" + +TP_SIZE="${1:-}" +MODE="${2:-}" +if (( $# < 1 || $# > 2 )) || [[ "$TP_SIZE" != "2" ]] || [[ "$MODE" != "" && "$MODE" != "--save" && "$MODE" != "--load" ]]; then + echo "Usage: $0 2 [--save|--load] (validated profile: TP=2)" >&2 + exit 1 +fi + +MODEL_NAME="${SGL_MODEL:-Qwen/Qwen3-8B}" +HOST="0.0.0.0" +PORT=12000 +MEM_FRACTION_STATIC=0.8 +RANDOM_SEED="${SGL_RANDOM_SEED:-42}" + +# Keep NCCL initialization identical across baseline, SAVE, and LOAD. CUMEM P2P +# and NVLS multicast use mapping capabilities the Foundry VMM region does not +# carry, so this recipe keeps them off for a reproducible rank-local layout. +export NCCL_CUMEM_ENABLE=0 +export NCCL_NVLS_ENABLE=0 +export PYTORCH_CUDA_ALLOC_CONF=expandable_segments:False +export SGLANG_ENABLE_POST_CAPTURE_KV_SIZING=0 +export SGLANG_PLUGINS=foundry +export SGL_MODEL_REVISION="${SGL_MODEL_REVISION:-b968826d9c46dd6066d109eabc6255188de91218}" + +unset FOUNDRY_GRAPH_EXTENSION_CONFIG_PATH +if [[ "$MODE" == "--save" ]]; then + export FOUNDRY_GRAPH_EXTENSION_CONFIG_PATH="${SCRIPT_DIR}/sglang_foundry_tp_save.toml" + echo "Using Foundry TP SAVE (torch symmetric memory): ${FOUNDRY_GRAPH_EXTENSION_CONFIG_PATH}" +elif [[ "$MODE" == "--load" ]]; then + export FOUNDRY_GRAPH_EXTENSION_CONFIG_PATH="${SCRIPT_DIR}/sglang_foundry_tp_load.toml" + echo "Using Foundry TP LOAD (torch symmetric memory): ${FOUNDRY_GRAPH_EXTENSION_CONFIG_PATH}" +else + echo "Running without Foundry (baseline SGLang)" +fi + +# --disable-custom-all-reduce / --enforce-disable-flashinfer-allreduce-fusion: +# keep SGLang's custom kernel and the FlashInfer fused all-reduce from winning so +# --enable-torch-symm-mem selects the torch symmetric-memory communicator that +# tests/modal_sglang_tp.py validates end-to-end. +# --decode-log-interval 1: decode_log_interval defaults to 40, so a short +# concurrent bs>1 decode would emit no "#running-req ... cuda graph: True" +# metrics line and the acceptance harness could not observe restored bs>1 +# replay. Logging every decode batch keeps that evidence visible in every phase. +sglang serve \ + --model-path "$MODEL_NAME" \ + --revision "$SGL_MODEL_REVISION" \ + --trust-remote-code \ + --host "$HOST" --port "$PORT" \ + --tp-size "$TP_SIZE" \ + --random-seed "$RANDOM_SEED" \ + --mem-fraction-static "$MEM_FRACTION_STATIC" \ + --disable-radix-cache \ + --decode-log-interval 1 \ + --enable-tokenizer-batch-encode \ + --disable-custom-all-reduce \ + --enable-torch-symm-mem \ + --enforce-disable-flashinfer-allreduce-fusion \ + --disable-flashinfer-autotune \ + --attention-backend flashinfer \ + --cuda-graph-config '{"decode":{"backend":"full","max_bs":256,"bs":[256]},"prefill":{"backend":"disabled"}}' diff --git a/recipe/experimental/sglang_foundry_tp_load.toml b/recipe/experimental/sglang_foundry_tp_load.toml new file mode 100644 index 00000000..8ed9f23a --- /dev/null +++ b/recipe/experimental/sglang_foundry_tp_load.toml @@ -0,0 +1,8 @@ +mode = "load" +base_addr = 0x600000000000 +region_size = "256GB" +workspace_root = "foundry_archive_sglang_tp" +scratch_space_size = "1024MB" +# Dense tensor parallel needs no NVSHMEM (no DeepEP). The per-layer TP +# all-reduce is served by NCCL P2P/IPC. NCCL initialization consumes about +# 302 MiB on the validated TP=2 H100 path, leaving headroom before this boundary. diff --git a/recipe/experimental/sglang_foundry_tp_save.toml b/recipe/experimental/sglang_foundry_tp_save.toml new file mode 100644 index 00000000..428ade3d --- /dev/null +++ b/recipe/experimental/sglang_foundry_tp_save.toml @@ -0,0 +1,8 @@ +mode = "save" +base_addr = 0x600000000000 +region_size = "256GB" +workspace_root = "foundry_archive_sglang_tp" +scratch_space_size = "1024MB" +# Dense tensor parallel needs no NVSHMEM (no DeepEP). The per-layer TP +# all-reduce is served by NCCL P2P/IPC. NCCL initialization consumes about +# 302 MiB on the validated TP=2 H100 path, leaving headroom before this boundary. diff --git a/recipe/sglang/README.md b/recipe/sglang/README.md index af4ce528..431a2b51 100644 --- a/recipe/sglang/README.md +++ b/recipe/sglang/README.md @@ -39,15 +39,25 @@ model or topology before SAVE. |---|---|---|---| | Single GPU | `serve_qwen3-mini.sh` | Qwen3-1.7B | FlashInfer backend | | Data parallel | `serve_qwen3-1.7b_dp.sh` | Qwen3-1.7B | one full replica/rank; `NCCL_CUMEM_ENABLE=0`/`NCCL_NVLS_ENABLE=0` | +| Tensor parallel | `../experimental/serve_qwen3-8b_sglang_tp.sh` | Qwen3-8B | validated only for the pinned single-host TP=2 torch symmetric-memory profile | | Expert parallel | `serve_qwen3-30ba3bfp8_ep.sh` | Qwen3-30B-A3B-FP8 | DP-attention + DeepEP; fa3 backend | ## Installation -The recipes assume `foundry` and the SGLang fork are **pip-installed** (editable is -fine) so both import without any `PYTHONPATH`, and foundry's spawn-site patch -auto-detects `libcuda_hook.so` from its install — the scripts set no `LD_PRELOAD` -themselves. The standard workspace layout has `foundry/` (this repo) and `sglang/` -(the foundry-org SGLang fork) as siblings: +> **Historical (pre-port).** The single-GPU / DP / EP recipes in this directory +> were developed against the **pre-port `foundry-org/sglang` fork** (direct +> edits applied to `sglang/`) and have not been re-validated on the pinned-main +> plugin. The current, zero-upstream-edit path is the dense TP=2 recipe under +> [`../experimental/`](../experimental/), which activates through the +> `sglang.srt.plugins` entry point against **unmodified** upstream SGLang; see +> [`../../docs/sglang/`](../../docs/sglang/). The fork layout below is retained +> only for reproducing those historical runs. + +The historical fork recipes assume `foundry` and the SGLang fork are +**pip-installed** (editable is fine) so both import without any `PYTHONPATH`, and +foundry's spawn-site patch auto-detects `libcuda_hook.so` from its install — the +scripts set no `LD_PRELOAD` themselves. The workspace layout had `foundry/` (this +repo) and `sglang/` (the foundry-org SGLang fork) as siblings: ``` / @@ -55,7 +65,7 @@ themselves. The standard workspace layout has `foundry/` (this repo) and `sglang │ ├── python/foundry/ # `pip install -e .` builds libcuda_hook.so here │ ├── recipe/sglang/ # <-- you are here │ └── ... -└── sglang/ # foundry-org/sglang fork (with direct edits applied) +└── sglang/ # foundry-org/sglang fork (pre-port, with direct edits applied) ``` Use a dedicated env, kept separate from the vLLM env so kernel pins don't clash: @@ -92,6 +102,82 @@ CUDA_VISIBLE_DEVICES=0,1 bash serve_qwen3-1.7b_dp.sh 2 --save CUDA_VISIBLE_DEVICES=0,1 bash serve_qwen3-1.7b_dp.sh 2 --load ``` +## Run (experimental tensor parallel) + +Dense TP captures torch symmetric-memory collectives in one padded decode graph +per rank. The entry point accepts only TP=2 and `--save` / `--load`; it rejects +other sizes and surplus arguments. Activation is through the lightweight +`foundry_sglang_plugin` entry point with zero upstream edits +(`SGLANG_PLUGINS=foundry` plus the SAVE/LOAD TOML selected via +`FOUNDRY_GRAPH_EXTENSION_CONFIG_PATH`). It installs nine named `HookRegistry` +targets. + +Shell recipe users must install the matching SGLang pin; the serve script does +not install or pin SGLang, though it passes model revision +`b968826d9c46dd6066d109eabc6255188de91218` via `--revision`. The exact +engine/model SHAs are validation provenance enforced by the Modal acceptance +harness (`tests/modal_sglang_tp.py`), which checks out upstream SGLang commit +`a25164bda355250015c0ba6d8821f0d4b37abe7d` and supplies the model revision. +`_validate_profile` checks topology, backend, and settings; arbitrary installed +SGLang revisions are not rejected at startup. The validated Foundry pin is +`a8e43442421a80fea7991ee26f97f1003b170e29`. + +```bash +cd ../experimental +rm -rf foundry_archive_sglang_tp +CUDA_VISIBLE_DEVICES=0,1 bash serve_qwen3-8b_sglang_tp.sh 2 --save +CUDA_VISIBLE_DEVICES=0,1 bash serve_qwen3-8b_sglang_tp.sh 2 --load +curl -s http://0.0.0.0:12000/v1/completions \ + -H 'Content-Type: application/json' \ + -d '{"model":"Qwen/Qwen3-8B","prompt":"The capital of France is","max_tokens":12,"temperature":0}' +``` + +In SAVE and LOAD, Foundry retains one torch symmetric-memory communicator +allocated/rendezvoused outside Foundry VMM. The gather hook calls +`torch.ops.symm_mem.multimem_all_gather_out` directly with a stable buffer slice +and the original group name. The recipe disables custom and FlashInfer-fused +all-reduce, captures exactly +`{"decode":{"backend":"full","max_bs":256,"bs":[256]},"prefill":{"backend":"disabled"}}`, +and keeps CUMEM/NVLS disabled. LOAD follows SGLang's normal upstream capture +lifecycle while the backend performs one pooled graph build and one ordered +finish; it does not capture or resave. + +Plain-NCCL fresh-process restore is unsupported. In disproving app +[`ap-2qSQUKlY24L4RrQkMpUrmc`](https://modal.com/apps/modal-labs/rahul-dev/ap-2qSQUKlY24L4RrQkMpUrmc), +36 graphs per rank restored with matching inventory, offsets, immutable archive +bytes, and P2P/IPC topology, but both fresh LOAD processes produced the same +corrupt decode while baseline and SAVE2 matched. Archived process-specific +communicator and graph-sequence state is the leading forensic explanation, not +treatment-proven causality. The observed corruption rejects plain-NCCL +restore; the structural pass alone did not prove output correctness. + +The focused run is `VERIFIED`: +[`ap-5wO9cFhRZ7ITsLDDv27Xfh`](https://modal.com/apps/modal-labs/rahul-dev/ap-5wO9cFhRZ7ITsLDDv27Xfh). +The final full matrix is `VERIFIED`: +[`ap-6fYcSUFLOxocBSjNxakNGZ`](https://modal.com/apps/modal-labs/rahul-dev/ap-6fYcSUFLOxocBSjNxakNGZ). +Run the harness from the repository root: + +```bash +modal run --env tests/modal_sglang_tp.py +``` + +The final run recorded one graph per rank at offset `68908220416`; 73 symmetric +all-reduce plus one symmetric all-gather and zero NCCL collective nodes per +rank; exact sequential and eight-slot batch-8 output equality; SAVE/SAVE2 +semantic reproducibility; LOAD/LOAD2 byte immutability; batch-1/batch-8 replay; +times 217.2/103.0/103.1/61.0/61.0 seconds; clean phase processes; and completed +cleanup. + +Do not extrapolate this exact pinned single-host Qwen3-8B TP=2 symmetric result +to general/upstream TP or plain NCCL. Custom/NCCL-symmetric/fused collectives, +CUMEM, NVLS, TP>2, multi-host, EP, DP-attention, speculative decode, LoRA, +pdmux, unpinned revisions, Ampere, Blackwell, and hybrid MoE are unsupported +or unvalidated. Qwen3.5-122B-A10B is explicitly unsupported. General TP remains +open in [issue #6](https://github.com/foundry-org/foundry/issues/6). + +Historical modal-projects draft PR #4/#2 results came from the **pre-port fork +integration** and are context only; they do not broaden this result. + ## Run (expert parallel / DeepEP) EP needs three kernel packages — all SGLang-native, no vLLM involved: @@ -100,14 +186,14 @@ EP needs three kernel packages — all SGLang-native, no vLLM involved: wheel as a dependency (`libnvshmem_host.so.3` under `site-packages/nvidia/nvshmem/lib/`). Foundry auto-detects it from the wheel (just like `libcuda_hook.so`) and the spawn-site patches preload it into each worker — no manual path, no TOML field. -- **DeepEP** @ `9af0e0d` — SGLang's pin. Build via SGLang's own installer - `sglang/scripts/ci/cuda/ci_install_deepep.sh` (it `git checkout`s exactly `9af0e0d` - and builds against the NVSHMEM wheel above). For a single node you can skip the +- **DeepEP** @ `29d31c0` or newer — required for Foundry's fabric/IPC policy + switch. Build via SGLang's own installer, overriding its old `9af0e0d` pin, + and build against the NVSHMEM wheel above. For a single node you can skip the script's gdrcopy/RDMA apt steps and just build the wheel: ```bash git clone https://github.com/deepseek-ai/DeepEP.git && cd DeepEP - git checkout 9af0e0d0e74f3577af1979c9b9e1ac2cad0104ee + git checkout 29d31c095796f3c8ece47ee9cdcc167051bbeed9 TORCH_CUDA_ARCH_LIST="9.0" python setup.py install # Hopper; "9.0;10.0" for Blackwell cd .. ``` diff --git a/recipe/vllm/README.md b/recipe/vllm/README.md index ea55d2f6..c542b618 100644 --- a/recipe/vllm/README.md +++ b/recipe/vllm/README.md @@ -11,6 +11,7 @@ recipe/vllm/ ├── foundry_load.toml # shared LOAD config (same workspace_root) ├── serve_qwen3-mini.sh # Qwen3-1.7B single GPU ├── serve_qwen3-14b_dp.sh # Qwen3-14B data parallel +├── serve_qwen3-8b_tp.sh # Qwen3-8B experimental TP=2 ├── serve_qwen3-30ba3b_ep.sh # Qwen3-30B-A3B (MoE) expert parallel (DeepEP) └── serve_qwen3-30ba3bfp8_ep.sh # Qwen3-30B-A3B FP8 expert parallel (DeepGEMM) ``` @@ -20,10 +21,15 @@ Every script accepts the same trailing `--save` / `--load` flag. Scripts that sc ```bash bash serve_qwen3-mini.sh [--save|--load] bash serve_qwen3-14b_dp.sh [--save|--load] +bash serve_qwen3-8b_tp.sh 2 [--save|--load] bash serve_qwen3-30ba3b_ep.sh [--save|--load] bash serve_qwen3-30ba3bfp8_ep.sh [--save|--load] ``` +The TP entry point is deliberately stricter than the older multi-GPU scripts: +it accepts exactly TP size 2, at most one mode argument, and rejects every +other invocation before starting vLLM. + Because the two TOMLs are shared (single `workspace_root = "foundry_archive"`), one archive is written per host and switching between scripts re-uses or overwrites it — run a fresh `rm -rf foundry_archive` whenever you change models or topology before SAVE pass 1. ## Installation @@ -54,8 +60,11 @@ pip install uv ```bash git clone https://github.com/foundry-org/vllm.git cd vllm -git checkout foundry -VLLM_USE_PRECOMPILED=1 uv pip install --editable . \ +git checkout --detach 4309c257d3f639e5490d3811293c890c61c76f29 +VLLM_USE_PRECOMPILED=1 \ +VLLM_PRECOMPILED_WHEEL_COMMIT=6cbe448eed751824d608faf9078ef84724d621c1 \ +VLLM_PRECOMPILED_WHEEL_VARIANT=cu130 \ +uv pip install --editable . \ --extra-index-url https://wheels.vllm.ai/nightly/cu130 cd .. ``` @@ -65,15 +74,15 @@ Verify the install with `uv pip list`: ``` nvidia-cublas 13.1.0.3 torch 2.11.0(+cu130) -vllm 0.1.dev15646+g040974074.precompiled +vllm 0.1.dev16924+g4309c257d.precompiled ``` ### 3. Install Foundry ```bash -git clone https://github.com/foundry-org/foundry.git +git clone https://github.com/modal-projects/foundry.git cd foundry -git checkout bump +git checkout --detach a8e43442421a80fea7991ee26f97f1003b170e29 uv pip install -e . --no-build-isolation uv pip install pytest pytest tests/ @@ -108,6 +117,14 @@ That will compile and install DeepEP module and put its shared lib in `vllm/tool Then uncomment the `nvshmem_host_path=` in `foundry_save.toml` / `foundry_load.toml` and points to your `libnvshmem_host.so`. +The standalone strict DeepEP transport/replay matrix is verified at this exact +Foundry pin in +[`ap-Q8mb8U4EYuE5Jo6Zb0xwrH`](https://modal.com/apps/modal-labs/rahul-dev/ap-Q8mb8U4EYuE5Jo6Zb0xwrH). +The pinned PyTorch preflight, `fabric`, `nvl_ipc`, `nvl_ipc_prealloc`, and +Modal CLI all returned zero at +`a8e43442421a80fea7991ee26f97f1003b170e29`. This validates the standalone +pinned matrix, not general DeepEP serving, MoE TP, or SGLang EP. + ## Workflow vLLM's `_initialize_kv_caches` does a real-forward memory-profile pass on the first run, so SAVE must run **twice** for a deterministic memory layout: @@ -143,6 +160,65 @@ bash serve_qwen3-30ba3b_ep.sh 2 --load bash ../../../experimental/query.sh 12000 Qwen/Qwen3-30B-A3B ``` +## Workflow (experimental dense TP=2) + +The TP candidate profile is single-host dense `Qwen/Qwen3-8B`, exactly two +Hopper GPUs, vLLM V1, full-decode graphs, and plain NCCL P2P / legacy CUDA IPC. +The script defaults to immutable model revision +`b968826d9c46dd6066d109eabc6255188de91218` and passes it through +`--revision`. The acceptance harness pins the Foundry vLLM fork revision to +`4309c257d3f639e5490d3811293c890c61c76f29`, requires a 40-character Foundry +commit, and records all revisions in its result. + +```bash +rm -rf foundry_archive + +# Ordinary vLLM baseline with the same transport and graph settings. +CUDA_VISIBLE_DEVICES=0,1 bash serve_qwen3-8b_tp.sh 2 + +# vLLM's deterministic two-pass SAVE, then a fresh LOAD. +CUDA_VISIBLE_DEVICES=0,1 bash serve_qwen3-8b_tp.sh 2 --save +CUDA_VISIBLE_DEVICES=0,1 bash serve_qwen3-8b_tp.sh 2 --save +CUDA_VISIBLE_DEVICES=0,1 bash serve_qwen3-8b_tp.sh 2 --load +``` + +Override `VLLM_MODEL_REVISION` only with another immutable model commit and +treat it as a new, unvalidated profile. The recipe forces PyNCCL and disables +custom, torch/NCCL symmetric-memory, and FlashInfer fused all-reduce; NCCL +CUMEM and NVLS; expandable allocator segments; fused all-reduce/RMS passes; and +vLLM V2 in baseline, SAVE, and LOAD. It passes the graph mode, zero warmups, +fused-pass setting, and optional Foundry config path in one atomic +`--compilation-config` object. + +This exact profile is verified by the six-phase +[`ap-YfKALdMwnyJQmzQoVvZSMU`](https://modal.com/apps/modal-labs/rahul-dev/ap-YfKALdMwnyJQmzQoVvZSMU) +run: + +```bash +FOUNDRY_REF=a8e43442421a80fea7991ee26f97f1003b170e29 \ +VLLM_REF=4309c257d3f639e5490d3811293c890c61c76f29 \ +MODAL_GPU=H100:2 \ +modal run --env rahul-dev tests/modal_vllm_tp.py +``` + +Every phase used exactly two NVIDIA H100 80GB HBM3 GPUs. All 17 checks passed: +64 JSON/64 `.cugraph` files per rank; `64/64` restoration for both ranks in +both LOADs; SAVE/SAVE2 and live LOAD/LOAD2 offset `70724354048`; 68 matching +semantic entries per rank; 265 byte-immutable published files; exact +baseline/SAVE2/LOAD/LOAD2 sequential and concurrent text/token strings; all +phase/Modal exit codes were zero; and successful Dict plus Volume cleanup. + +The measured environment was CUDA `13.0.1-cudnn-devel-ubuntu24.04`, driver +`580.95.05`, PyTorch `2.11.0+cu130`, NCCL `2.28.9+cuda13.0`, and vLLM +`0.1.dev16924+g4309c257d.precompiled`. Changing any Foundry, vLLM, wheel, +model, CUDA, or collective pin creates a new unvalidated profile. + +TP>2, multi-host, unpinned engine heads, Ampere, Blackwell, vLLM V2, MoE TP, +and alternate collective backends are unsupported or unvalidated. +Upstream [Discussion #5](https://github.com/orgs/foundry-org/discussions/5) +documents NCCL initialization nondeterminism; official TP remains open in +[issue #6](https://github.com/foundry-org/foundry/issues/6). + ## Archive layout Every SAVE writes into `foundry_archive/` next to the script's working directory: diff --git a/recipe/vllm/serve_qwen3-8b_tp.sh b/recipe/vllm/serve_qwen3-8b_tp.sh new file mode 100644 index 00000000..fc129ee0 --- /dev/null +++ b/recipe/vllm/serve_qwen3-8b_tp.sh @@ -0,0 +1,79 @@ +#!/bin/bash +# Usage: CUDA_VISIBLE_DEVICES=0,1 bash serve_qwen3-8b_tp.sh [--save|--load] +# +# Dense tensor parallelism captures a per-layer all-reduce in every CUDA graph. +# Keep that collective on NCCL's legacy P2P/IPC path in baseline, SAVE, and LOAD +# so all phases exercise the same transport and Foundry can restore its VMM +# mappings at their recorded addresses. +set -euo pipefail + +SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" + +TP_SIZE="${1:-}" +MODE="${2:-}" +if (( $# < 1 || $# > 2 )) || [[ "$TP_SIZE" != "2" ]] || [[ "$MODE" != "" && "$MODE" != "--save" && "$MODE" != "--load" ]]; then + echo "Usage: $0 2 [--save|--load] (validated profile: TP=2)" >&2 + exit 1 +fi + +MODEL_NAME="${VLLM_MODEL:-Qwen/Qwen3-8B}" +MODEL_REVISION="${VLLM_MODEL_REVISION:-b968826d9c46dd6066d109eabc6255188de91218}" +if [[ ! "$MODEL_REVISION" =~ ^[0-9a-fA-F]{40}$ ]]; then + echo "VLLM_MODEL_REVISION must be a full 40-hex commit SHA, got: $MODEL_REVISION" >&2 + exit 1 +fi +HOST="${VLLM_HOST:-0.0.0.0}" +PORT="${VLLM_PORT:-12000}" +GPU_MEMORY_UTILIZATION="${VLLM_GPU_MEMORY_UTILIZATION:-0.8}" + +# Route TP collectives through one reproducible NCCL implementation. vLLM's +# symmetric-memory and FlashInfer paths take precedence over PyNCCL unless they +# are disabled explicitly, even with --disable-custom-all-reduce. +export NCCL_CUMEM_ENABLE=0 +export NCCL_NVLS_ENABLE=0 +export PYTORCH_CUDA_ALLOC_CONF=expandable_segments:False +export VLLM_DISABLE_PYNCCL=0 +export VLLM_ALLREDUCE_USE_SYMM_MEM=0 +export VLLM_ALLREDUCE_USE_FLASHINFER=0 +export VLLM_USE_NCCL_SYMM_MEM=0 + +# Foundry patches the V1 model runner. The pinned vLLM fork accepts spawn (not +# forkserver) for worker processes, and Foundry injects LD_PRELOAD at each spawn. +export VLLM_USE_V2_MODEL_RUNNER=0 +export VLLM_WORKER_MULTIPROC_METHOD=spawn + +if [[ "$MODE" == "--save" ]]; then + FOUNDRY_TOML="${SCRIPT_DIR}/foundry_save.toml" +elif [[ "$MODE" == "--load" ]]; then + FOUNDRY_TOML="${SCRIPT_DIR}/foundry_load.toml" +fi + +if [[ -n "${FOUNDRY_TOML:-}" ]]; then + COMPILATION_CONFIG="{\"cudagraph_mode\":\"FULL_DECODE_ONLY\",\"cudagraph_num_of_warmups\":0,\"pass_config\":{\"fuse_allreduce_rms\":false},\"graph_extension_config_path\":\"${FOUNDRY_TOML}\"}" + echo "Using Foundry TP (NCCL P2P/IPC): ${FOUNDRY_TOML}" +else + COMPILATION_CONFIG='{"cudagraph_mode":"FULL_DECODE_ONLY","cudagraph_num_of_warmups":0,"pass_config":{"fuse_allreduce_rms":false}}' + echo "Running without Foundry (baseline vLLM TP)" +fi + +CUDAGRAPH_CAPTURE_SIZES=($(seq 1 64)) + +ARGS=( + --trust-remote-code + --host "$HOST" + --port "$PORT" + --revision "$MODEL_REVISION" + --tensor-parallel-size "$TP_SIZE" + --distributed-executor-backend mp + --disable-custom-all-reduce + --gpu-memory-utilization "$GPU_MEMORY_UTILIZATION" + --no-enable-prefix-caching + --max-num-batched-tokens 64 + --max-num-seqs 64 + --attention-config.backend FLASH_ATTN + --compilation-config "$COMPILATION_CONFIG" + --chat-template-content-format string + --cudagraph-capture-sizes "${CUDAGRAPH_CAPTURE_SIZES[@]}" +) + +vllm serve "$MODEL_NAME" "${ARGS[@]}" diff --git a/tests/fixtures/sglang_a25164_tokenizer_batch_dispatch.txt b/tests/fixtures/sglang_a25164_tokenizer_batch_dispatch.txt new file mode 100644 index 00000000..9abb917f --- /dev/null +++ b/tests/fixtures/sglang_a25164_tokenizer_batch_dispatch.txt @@ -0,0 +1,83 @@ +# Verbatim source evidence from: +# https://github.com/sgl-project/sglang/blob/a25164bda355250015c0ba6d8821f0d4b37abe7d/python/sglang/srt/managers/tokenizer_manager.py +# commit: a25164bda355250015c0ba6d8821f0d4b37abe7d +# complete upstream file sha256: 4646904f0da00d4810db3e08b90d0dee41d8ecf317c68e99ad8ca727fa96deb8 + + def _should_use_batch_tokenization(self, batch_size, requests) -> bool: + """Return True if we should run the tokenizer in batch mode. + + Current policy: + - Respect explicit server flag `enable_tokenizer_batch_encode`. + - Or, if no request has text or multimodal input (all use pre-tokenized input_ids or input_embeds), batch the requests without tokenization. + - Batch tokenization does not support DP attention yet, and it will make everything goes to the first rank currently + """ + return batch_size > 0 and ( + self.server_args.enable_tokenizer_batch_encode + or ( + (not self.server_args.enable_dp_attention) + and (not self._batch_has_text(batch_size, requests)) + ) + ) + + def _send_batch_request( + self, + tokenized_objs: List[ + Union[TokenizedGenerateReqInput, TokenizedEmbeddingReqInput] + ], + ): + """Send a batch of tokenized requests as a single batched request to the scheduler.""" + set_time_batch(tokenized_objs, "set_api_server_dispatch_time") + time_stats = [tokenized_obj.time_stats for tokenized_obj in tokenized_objs] + for tokenized_obj in tokenized_objs: + tokenized_obj.wrap_pickle_fields() + + if isinstance(tokenized_objs[0], TokenizedGenerateReqInput): + batch_req = BatchTokenizedGenerateReqInput(batch=tokenized_objs) + else: + batch_req = BatchTokenizedEmbeddingReqInput(batch=tokenized_objs) + + self._dispatch_to_scheduler(batch_req) + for tokenized_obj, time_stat in zip(tokenized_objs, time_stats): + tokenized_obj.time_stats = time_stat + set_time_batch(tokenized_objs, "set_api_server_dispatch_finish_time") + + async def _handle_batch_request( + self, + obj: Union[GenerateReqInput, EmbeddingReqInput], + request: Optional[fastapi.Request] = None, + ): + batch_size = obj.batch_size + + generators = [] + rids = [] + if getattr(obj, "parallel_sample_num", 1) == 1: + if self._should_use_batch_tokenization(batch_size, obj): + tokenized_objs = await self._batch_tokenize_and_process(batch_size, obj) + self._send_batch_request(tokenized_objs) + + # Set up generators for each request in the batch + for i in range(batch_size): + tmp_obj = obj[i] + state = self.rid_to_state[tmp_obj.rid] + if tmp_obj.return_prompt_token_ids: + state.prompt_token_ids = list(tokenized_objs[i].input_ids) + generators.append(self._wait_one_response(tmp_obj, request)) + rids.append(tmp_obj.rid) + else: + # Sequential tokenization and processing + with ( + input_blocker_guard_region( + dispatch_to_scheduler=self._dispatch_to_scheduler, + ) + if get_bool_env_var("SGLANG_ENABLE_COLOCATED_BATCH_GEN") + else nullcontext() + ): + for i in range(batch_size): + tmp_obj = obj[i] + tokenized_obj = await self._tokenize_one_request(tmp_obj) + state = self.rid_to_state[tmp_obj.rid] + if tmp_obj.return_prompt_token_ids: + state.prompt_token_ids = list(tokenized_obj.input_ids) + self._send_one_request(tokenized_obj) + generators.append(self._wait_one_response(tmp_obj, request)) + rids.append(tmp_obj.rid) diff --git a/tests/modal_deepep.py b/tests/modal_deepep.py new file mode 100644 index 00000000..57ac6d39 --- /dev/null +++ b/tests/modal_deepep.py @@ -0,0 +1,227 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the Foundry project +"""Run the strict two-GPU DeepEP SAVE/LOAD matrix on Modal. + +The image always checks out an explicit immutable Foundry commit: + + FOUNDRY_REF=<40-hex-commit> \ + modal run --env rahul-dev tests/modal_deepep.py + +The orchestration process never imports a GPU runtime. A pinned-PyTorch CPU +oracle preflight and each native matrix case run in separate child processes; +the matrix runner preloads the Foundry CUDA hook for native cases. +""" + +from __future__ import annotations + +import contextlib +import json +import os +import re +import shlex +import signal +import subprocess + +import modal + +CUDA_TAG = "13.0.1-cudnn-devel-ubuntu24.04" +FOUNDRY_REPO = "https://github.com/modal-projects/foundry.git" +VLLM_REPO = "https://github.com/foundry-org/vllm.git" +VLLM_REF = "4309c257d3f639e5490d3811293c890c61c76f29" +DEEPEP_REF = "29d31c095796f3c8ece47ee9cdcc167051bbeed9" +NVSHMEM_VERSION = "3.3.24" +MATRIX_CASES = ("fabric", "nvl_ipc", "nvl_ipc_prealloc") +PREFLIGHT_COMMAND = ( + "python3", + "-m", + "pytest", + "tests/test_deepep_fabric.py::test_torch_cpu_preflight_covers_split_payload_and_dtypes", + "-q", +) +PREFLIGHT_TIMEOUT_SECONDS = 2 * 60 +CASE_TIMEOUT_SECONDS = 20 * 60 +POST_KILL_TIMEOUT_SECONDS = 10 +FUNCTION_TIMEOUT_SECONDS = 75 * 60 + + +def require_full_commit_sha(value: str | None) -> str: + if value is None or re.fullmatch(r"[0-9a-fA-F]{40}", value) is None: + raise ValueError("FOUNDRY_REF must be a full 40-hex commit SHA") + return value.lower() + + +def commands_for_matrix() -> list[tuple[str, ...]]: + return [("bash", "tests/run_deepep_matrix.sh", case, "both") for case in MATRIX_CASES] + + +def execution_plan() -> list[tuple[str, str, tuple[str, ...], int]]: + return [ + ( + "preflight", + "torch_cpu_oracle", + PREFLIGHT_COMMAND, + PREFLIGHT_TIMEOUT_SECONDS, + ), + *[ + ("case", case, command, CASE_TIMEOUT_SECONDS) + for case, command in zip(MATRIX_CASES, commands_for_matrix(), strict=True) + ], + ] + + +FOUNDRY_REF = require_full_commit_sha(os.environ.get("FOUNDRY_REF")) + +image = ( + modal.Image.from_registry(f"nvidia/cuda:{CUDA_TAG}", add_python="3.12") + .env( + { + "CC": "gcc", + "CXX": "g++", + "FOUNDRY_REF": FOUNDRY_REF, + "UV_SYSTEM_PYTHON": "1", + } + ) + .apt_install( + "git", + "build-essential", + "cmake", + "ninja-build", + "libboost-all-dev", + "libnuma-dev", + "curl", + "xz-utils", + ) + .pip_install( + "cmake>=4.0.0", + "wheel", + "packaging", + "ninja", + "pytest", + "setuptools>=80", + "setuptools-scm", + "uv", + ) + .pip_install( + "torch==2.11.0", + "torchvision==0.26.0", + "torchaudio==2.11.0", + index_url="https://download.pytorch.org/whl/cu130", + ) + .run_commands( + f"git clone {FOUNDRY_REPO} /foundry", + f"git -C /foundry checkout --detach {FOUNDRY_REF}", + f'test "$(git -C /foundry rev-parse HEAD)" = "{FOUNDRY_REF}"', + f"git clone {VLLM_REPO} /vllm", + f"git -C /vllm checkout --detach {VLLM_REF}", + f'test "$(git -C /vllm rev-parse HEAD)" = "{VLLM_REF}"', + ( + "cd /vllm/tools/ep_kernels && " + "TORCH_CUDA_ARCH_LIST=9.0 bash install_python_libraries.sh " + "--workspace /vllm/tools/ep_kernels/ep_kernels_workspace " + f"--deepep-ref {DEEPEP_REF} --nvshmem-ver {NVSHMEM_VERSION}" + ), + "cd /foundry && CC=gcc CXX=g++ pip install -e . --no-build-isolation", + ) +) + +app = modal.App("foundry-deepep-validation") + + +def _decode_output(output: str | bytes | None) -> str: + if output is None: + return "" + if isinstance(output, bytes): + return output.decode(errors="replace") + return output + + +def _run_child_command( + command: tuple[str, ...], + timeout_seconds: int, + popen_factory, + kill_process_group, +) -> tuple[int, str]: + process = popen_factory( + command, + cwd="/foundry", + stdout=subprocess.PIPE, + stderr=subprocess.STDOUT, + text=True, + start_new_session=True, + ) + try: + output, _ = process.communicate(timeout=timeout_seconds) + return process.returncode, _decode_output(output) + except subprocess.TimeoutExpired as timeout_error: + initial_output = _decode_output(timeout_error.output) + with contextlib.suppress(ProcessLookupError): + kill_process_group(process.pid, signal.SIGKILL) + + try: + post_kill_output, _ = process.communicate(timeout=POST_KILL_TIMEOUT_SECONDS) + output = _decode_output(post_kill_output) or initial_output + except subprocess.TimeoutExpired as post_kill_error: + output = _decode_output(post_kill_error.output) or initial_output + + try: + process.wait(timeout=POST_KILL_TIMEOUT_SECONDS) + except subprocess.TimeoutExpired: + with contextlib.suppress(ProcessLookupError): + kill_process_group(process.pid, signal.SIGKILL) + output += ( + f"Process group {process.pid} did not reap within " + f"{POST_KILL_TIMEOUT_SECONDS} seconds\n" + ) + + output += f"Timed out after {timeout_seconds} seconds\n" + return 124, output + + +def execute_deepep_matrix( + popen_factory=subprocess.Popen, + kill_process_group=os.killpg, +) -> dict: + results = [] + for stage, name, command, timeout_seconds in execution_plan(): + returncode, output = _run_child_command( + command, + timeout_seconds, + popen_factory, + kill_process_group, + ) + result = { + "stage": stage, + "name": name, + "command": shlex.join(command), + "output": output, + "returncode": returncode, + } + results.append(result) + print(json.dumps({"event": "deepep_matrix_command", **result}), flush=True) + if returncode != 0: + raise RuntimeError( + json.dumps( + { + "foundry_ref": FOUNDRY_REF, + "results": results, + "status": "failed", + }, + sort_keys=True, + ) + ) + + return { + "foundry_ref": FOUNDRY_REF, + "results": results, + "status": "passed", + } + + +@app.function(image=image, gpu="H100:2", timeout=FUNCTION_TIMEOUT_SECONDS) +def run_deepep_matrix() -> dict: + return execute_deepep_matrix() + + +@app.local_entrypoint() +def main() -> None: + print(json.dumps(run_deepep_matrix.remote(), indent=2, sort_keys=True)) diff --git a/tests/modal_native.py b/tests/modal_native.py new file mode 100644 index 00000000..07a33590 --- /dev/null +++ b/tests/modal_native.py @@ -0,0 +1,143 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the Foundry project +"""Build and run Foundry's single-GPU native CUDA tests on Modal. + +The checkout is always cloned from an explicit immutable commit: + + FOUNDRY_REF=<40-hex-commit> \ + modal run --env rahul-dev tests/modal_native.py --suite all + +The Modal entrypoint only dispatches the GPU function. PyTorch, Foundry, and +CUDA are first loaded by the pytest child processes inside the remote function. +""" + +from __future__ import annotations + +import json +import os +import re +import shlex +import subprocess + +import modal + +CUDA_TAG = "13.0.1-cudnn-devel-ubuntu24.04" +FOUNDRY_REPO = "https://github.com/modal-projects/foundry.git" + +FOCUSED_COMMAND = ( + "python3", + "-m", + "pytest", + "tests/test_vmm_alloc.py", + "tests/test_rng_graph.py", + "-q", +) +CORE_COMMAND = ( + "python3", + "-m", + "pytest", + "tests/", + "-q", + "--ignore=tests/test_imports.py", + "--ignore=tests/test_deepep_fabric.py", +) + + +def require_full_commit_sha(value: str | None) -> str: + if value is None or re.fullmatch(r"[0-9a-fA-F]{40}", value) is None: + raise ValueError("FOUNDRY_REF must be a full 40-hex commit SHA") + return value.lower() + + +def commands_for_suite(suite: str) -> list[tuple[str, ...]]: + if suite == "focused": + return [FOCUSED_COMMAND] + if suite == "core": + return [CORE_COMMAND] + if suite == "all": + return [FOCUSED_COMMAND, CORE_COMMAND] + raise ValueError("suite must be focused, core, or all") + + +FOUNDRY_REF = require_full_commit_sha(os.environ.get("FOUNDRY_REF")) + +image = ( + modal.Image.from_registry(f"nvidia/cuda:{CUDA_TAG}", add_python="3.12") + .env({"CC": "gcc", "CXX": "g++", "FOUNDRY_REF": FOUNDRY_REF}) + .apt_install( + "git", + "build-essential", + "cmake", + "ninja-build", + "libboost-all-dev", + "libnuma-dev", + ) + .pip_install( + "cmake>=4.0.0", + "wheel", + "packaging", + "ninja", + "pytest", + "setuptools>=80", + "setuptools-scm", + ) + .pip_install( + "torch==2.11.0", + "torchvision==0.26.0", + "torchaudio==2.11.0", + index_url="https://download.pytorch.org/whl/cu130", + ) + .run_commands( + f"git clone {FOUNDRY_REPO} /foundry", + f"git -C /foundry checkout --detach {FOUNDRY_REF}", + f'test "$(git -C /foundry rev-parse HEAD)" = "{FOUNDRY_REF}"', + "cd /foundry && CC=gcc CXX=g++ pip install -e . --no-build-isolation", + ) +) + +app = modal.App("foundry-native-validation") + + +def execute_native_suite(suite: str, runner=subprocess.run) -> dict: + results = [] + for command in commands_for_suite(suite): + completed = runner( + command, + cwd="/foundry", + stdout=subprocess.PIPE, + stderr=subprocess.STDOUT, + text=True, + ) + result = { + "command": shlex.join(command), + "output": completed.stdout, + "returncode": completed.returncode, + } + results.append(result) + print(json.dumps({"event": "native_test_command", **result}), flush=True) + if completed.returncode != 0: + failure = { + "foundry_ref": FOUNDRY_REF, + "results": results, + "status": "failed", + "suite": suite, + } + raise RuntimeError(json.dumps(failure)) + + return { + "foundry_ref": FOUNDRY_REF, + "results": results, + "status": "passed", + "suite": suite, + } + + +@app.function(image=image, gpu="H100", timeout=60 * 60) +def run_native_suite(suite: str = "all") -> dict: + return execute_native_suite(suite) + + +@app.local_entrypoint() +def main(suite: str = "all") -> None: + commands_for_suite(suite) + print(json.dumps(run_native_suite.remote(suite), indent=2, sort_keys=True)) diff --git a/tests/modal_sglang_target_resolution.py b/tests/modal_sglang_target_resolution.py new file mode 100644 index 00000000..28c9a8d2 --- /dev/null +++ b/tests/modal_sglang_target_resolution.py @@ -0,0 +1,232 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the Foundry project +"""Clean-checkout contract for Foundry hooks against pinned SGLang main. + +Run only with a pushed, immutable Foundry commit: + + FOUNDRY_REF=<40-hex-commit> \ + modal run --env rahul-dev tests/modal_sglang_target_resolution.py + +The local orchestration process only constructs the existing pinned image and +dispatches remote checks; the local orchestration process never initializes CUDA. +The inactive plugin check is CPU-only, while the configured activation check uses +one H100 so the native Foundry extension can load libcuda. SGLang, PyTorch, and +hook targets are imported only in fresh remote containers. +""" + +from __future__ import annotations + +import json +import os +import pkgutil +import re +import sys +from importlib.metadata import entry_points +from pathlib import Path + +import modal + +GENERAL_PLUGINS_GROUP = "sglang.srt.plugins" +EXPECTED_ENTRY_POINT = "foundry_sglang_plugin:activate" +FULL_COMMIT_SHA_PATTERN = re.compile(r"^[0-9a-fA-F]{40}$") +FOUNDRY_REPO = "https://github.com/modal-projects/foundry.git" +SGLANG_REPO = "https://github.com/sgl-project/sglang.git" +SGLANG_BASE = os.environ.get("SGLANG_BASE", "a25164bda355250015c0ba6d8821f0d4b37abe7d") +MODEL = os.environ.get("SGL_MODEL", "Qwen/Qwen3-8B") +MODEL_REVISION = os.environ.get("SGL_MODEL_REVISION", "b968826d9c46dd6066d109eabc6255188de91218") + +# Validate the clean-checkout revision before importing modal_sglang_tp: that +# module constructs the image while it is imported. +FOUNDRY_REF = os.environ.get("FOUNDRY_REF") +if FOUNDRY_REF is None or FULL_COMMIT_SHA_PATTERN.fullmatch(FOUNDRY_REF) is None: + raise ValueError(f"FOUNDRY_REF must be a full 40-hex commit SHA, got {FOUNDRY_REF!r}") +FOUNDRY_REF = FOUNDRY_REF.lower() +os.environ["FOUNDRY_REF"] = FOUNDRY_REF + +if modal.is_local(): + # Local-only import: Modal remounts this file as a standalone remote module, + # while the sibling harness modules are intentionally unavailable there. + # The remote container already runs in the locally constructed image. + import modal_sglang_tp as shared_harness + + assert ( + FOUNDRY_REF, + FOUNDRY_REPO, + SGLANG_BASE, + SGLANG_REPO, + MODEL, + MODEL_REVISION, + ) == ( + shared_harness.FOUNDRY_REF, + shared_harness.FOUNDRY_REPO, + shared_harness.SGLANG_BASE, + shared_harness.SGLANG_REPO, + shared_harness.MODEL, + shared_harness.MODEL_REVISION, + ) + image = shared_harness.image +else: + # Decorators are re-evaluated when Modal imports the mounted function + # module remotely; the execution image was already selected locally. + image = None + +app = modal.App("foundry-sglang-target-resolution") + + +def _installed_entry_point() -> str: + discovered = [(ep.name, ep.value) for ep in entry_points(group=GENERAL_PLUGINS_GROUP)] + foundry_entry_points = [value for name, value in discovered if name == "foundry"] + if foundry_entry_points != [EXPECTED_ENTRY_POINT]: + raise RuntimeError( + "Installed Foundry SGLang entry point mismatch: " + f"expected one {EXPECTED_ENTRY_POINT!r}, " + f"got {foundry_entry_points!r}; discovered={discovered!r}" + ) + return foundry_entry_points[0] + + +def _installed_hook_targets() -> tuple[str, ...]: + # Remote-image-only import: keep Foundry packages out of the local Modal + # orchestration parent while sharing one target list with hook registration. + from foundry_sglang_plugin.targets import HOOK_TARGETS + + if len(HOOK_TARGETS) != 9 or len(set(HOOK_TARGETS)) != 9: + raise RuntimeError(f"Expected nine unique Foundry hook targets, got {HOOK_TARGETS!r}") + for target in HOOK_TARGETS: + pkgutil.resolve_name(target) + return HOOK_TARGETS + + +@app.function(image=image, timeout=600) +def check_inactive_plugin() -> dict: + os.environ["SGLANG_PLUGINS"] = "foundry" + os.environ.pop("FOUNDRY_GRAPH_EXTENSION_CONFIG_PATH", None) + + entry_point = _installed_entry_point() + + # Remote-image-only imports; importing SGLang locally would defeat the + # no-CUDA-orchestration-parent contract. + from sglang.srt.plugins import load_plugins + from sglang.srt.plugins.hook_registry import HookRegistry + + if HookRegistry._hooks or HookRegistry._patched: + raise RuntimeError("HookRegistry was not empty before inactive plugin discovery") + + load_plugins() + + if HookRegistry._hooks or HookRegistry._patched: + raise RuntimeError( + "Inactive Foundry plugin changed HookRegistry: " + f"registered={sorted(HookRegistry._hooks)}, patched={sorted(HookRegistry._patched)}" + ) + imported_foundry = sorted( + name for name in sys.modules if name == "foundry" or name.startswith("foundry.") + ) + if imported_foundry: + raise RuntimeError(f"Inactive plugin imported Foundry modules: {imported_foundry}") + + return { + "entry_point": entry_point, + "registered_targets": [], + "patched_targets": [], + "imported_foundry_modules": imported_foundry, + } + + +@app.function(image=image, gpu="H100", timeout=600) +def check_configured_plugin() -> dict: + os.environ["SGLANG_PLUGINS"] = "foundry" + config_path = Path("/tmp/foundry-sglang-target-resolution.toml") + config_path.write_text('mode = "save"\n') + os.environ["FOUNDRY_GRAPH_EXTENSION_CONFIG_PATH"] = str(config_path) + + entry_point = _installed_entry_point() + targets = _installed_hook_targets() + expected = set(targets) + + # Remote-image-only imports; HookRegistry intentionally logs and continues + # on target failures, so this contract inspects its private result sets. + from foundry_sglang_plugin.targets import ( + FULL_CUDA_GRAPH_BACKEND, + normalize_resolved_identity, + ) + from sglang.srt.plugins import load_plugins + from sglang.srt.plugins.hook_registry import HookRegistry + + originals = { + target: normalize_resolved_identity(pkgutil.resolve_name(target)) for target in targets + } + + if HookRegistry._hooks or HookRegistry._patched: + raise RuntimeError("HookRegistry was not empty before configured plugin discovery") + + load_plugins() + + registered = set(HookRegistry._hooks) + patched = set(HookRegistry._patched) + if registered != expected: + raise RuntimeError( + "Configured plugin registration mismatch: " + f"missing={sorted(expected - registered)}, extra={sorted(registered - expected)}" + ) + if patched != expected: + raise RuntimeError( + "HookRegistry silently missed configured targets: " + f"missing={sorted(expected - patched)}, extra={sorted(patched - expected)}" + ) + + replacements = { + target: normalize_resolved_identity(pkgutil.resolve_name(target)) for target in targets + } + unchanged = sorted(target for target in targets if replacements[target] is originals[target]) + if unchanged: + raise RuntimeError(f"Applied hooks left target objects unchanged: {unchanged}") + + original_backend = originals[FULL_CUDA_GRAPH_BACKEND] + replacement_backend = replacements[FULL_CUDA_GRAPH_BACKEND] + if not ( + isinstance(original_backend, type) + and isinstance(replacement_backend, type) + and issubclass(replacement_backend, original_backend) + ): + raise RuntimeError( + "FullCudaGraphBackend REPLACE did not resolve to a Foundry subclass: " + f"original={original_backend!r}, replacement={replacement_backend!r}" + ) + + return { + "entry_point": entry_point, + "resolved_targets": list(targets), + "registered_targets": sorted(registered), + "patched_targets": sorted(patched), + "changed_targets": sorted(replacements), + "backend_replacement": ( + f"{replacement_backend.__module__}.{replacement_backend.__qualname__}" + ), + } + + +@app.local_entrypoint() +def main() -> None: + inactive = check_inactive_plugin.remote() + configured = check_configured_plugin.remote() + report = { + "provenance": { + "foundry_repo": FOUNDRY_REPO, + "foundry_ref": FOUNDRY_REF, + "sglang_repo": SGLANG_REPO, + "sglang_base": SGLANG_BASE, + "model": MODEL, + "model_revision": MODEL_REVISION, + }, + "checks": { + "installed_entry_point_discovered": True, + "all_pinned_targets_resolved": True, + "inactive_plugin_is_noop": True, + "configured_targets_registered_and_applied": True, + "class_replacement_resolved": True, + }, + "inactive": inactive, + "configured": configured, + } + print(json.dumps(report, indent=2, sort_keys=True)) diff --git a/tests/modal_sglang_tp.py b/tests/modal_sglang_tp.py new file mode 100644 index 00000000..a85c17a9 --- /dev/null +++ b/tests/modal_sglang_tp.py @@ -0,0 +1,1257 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the Foundry project +"""End-to-end SGLang tensor-parallel validation on Modal. + +The harness runs five phases on two GPUs, each in a fresh server process: + +1. ``baseline`` — a normal SGLang TP run and deterministic queries; the + reference outputs and the reference NCCL transport set. +2. ``save`` — Foundry SAVE into a fresh archive. +3. ``save2`` — a second SAVE that must reproduce the per-rank allocation + offsets and graph inventory, plus deterministic queries. +4. ``load`` — Foundry LOAD in a fresh server process and the same queries. +5. ``load2`` — a second, independent LOAD against the *same* immutable archive + (a genuinely fresh process, no stale IPC broker state) and the same queries. + +Success requires exact bounded text equality of SAVE2 / LOAD / LOAD2 with the +baseline, including one explicit eight-item native SGLang scheduler batch, +reproducible archives / per-rank offsets, exactly one archived and restored +graph per rank, positive per-rank symmetric all-reduce / all-gather KernelNodes +with zero archived NCCL collective nodes, symmetric communicator and +direct-multimem markers in every Foundry phase, restored-graph replay at batch +size 1 and larger, LOAD/LOAD2 final offsets equal to SAVE2, zero +recapture/resave on LOAD, clean process shutdown, and no fatal runtime +signature. NCCL channels and versions remain diagnostic provenance only. + +Every phase uses a unique run directory on a mounted Modal Volume. Modal may +auto-commit that directory when a container exits, including after a failed +phase. Failure reports therefore conservatively mark artifacts retained and +possibly partial or mutated; they stop later phases and never verify the run. + +This harness targets the pinned-main plugin architecture: Foundry activates +through SGLang's ``sglang.srt.plugins`` entry point ``foundry`` with zero +upstream edits, gated by ``FOUNDRY_GRAPH_EXTENSION_CONFIG_PATH``. + +Run the current checkout: + + modal run --env tests/modal_sglang_tp.py + +Run a specific commit: + + FOUNDRY_REF= \ + modal run --env tests/modal_sglang_tp.py + +Use an immutable commit rather than a moving branch name so Modal's image cache +cannot reuse a build from an older branch head. +""" + +from __future__ import annotations + +import importlib.metadata +import json +import os +import re +import shutil +import subprocess +import time +import urllib.request +from pathlib import Path +from typing import IO + +import modal +from sglang_tp_validation import ( + classify_focused_checks, + focused_save_checks, + focused_treatment_checks, + is_allowlisted_startup_environment_failure, + is_modal_infrastructure_error_type, + require_commit_sha, + require_tp_size, + scan_log_text, + stop_process, +) +from tp_validation_helpers import ( + archive_byte_fingerprints, + archive_collective_kernel_counts, + archive_fingerprints, + archive_graph_json_counts, +) + + +def _local_foundry_revision() -> str: + try: + revision = subprocess.check_output( + ["git", "rev-parse", "HEAD"], + cwd=Path(__file__).parents[1], + stderr=subprocess.DEVNULL, + text=True, + ).strip() + except (OSError, subprocess.CalledProcessError) as error: + raise RuntimeError( + "FOUNDRY_REF is required when the local checkout revision cannot be resolved" + ) from error + return require_commit_sha("local Foundry HEAD", revision) + + +CUDA_TAG = "13.0.1-cudnn-devel-ubuntu24.04" + +FOUNDRY_REPO = "https://github.com/modal-projects/foundry.git" +FOUNDRY_REF = require_commit_sha( + "FOUNDRY_REF", + os.environ.get("FOUNDRY_REF") or _local_foundry_revision(), +) + +# Foundry activates through SGLang's first-class sglang.srt.plugins entry point; +# no fork commit is required. Pin an immutable upstream main commit. +SGLANG_REPO = "https://github.com/sgl-project/sglang.git" +SGLANG_BASE = require_commit_sha( + "SGLANG_BASE", + os.environ.get("SGLANG_BASE", "a25164bda355250015c0ba6d8821f0d4b37abe7d"), +) +MODEL_REVISION = require_commit_sha( + "SGL_MODEL_REVISION", + os.environ.get("SGL_MODEL_REVISION", "b968826d9c46dd6066d109eabc6255188de91218"), +) + +MODEL = os.environ.get("SGL_MODEL", "Qwen/Qwen3-8B") +TP_SIZE = require_tp_size(os.environ.get("TP_SIZE", "2")) +MODAL_GPU = os.environ.get("MODAL_GPU", f"H100:{TP_SIZE}") +PORT = 12000 + +ARCHIVE_ROOT = "/data/archive" +RUNS_ROOT = f"{ARCHIVE_ROOT}/runs" +HF_CACHE = "/data/hf" +RECIPE_SCRIPT = "/foundry/recipe/experimental/serve_qwen3-8b_sglang_tp.sh" +FOUNDRY_BASE_ADDR = 0x600000000000 +FOUNDRY_REGION_END = FOUNDRY_BASE_ADDR + 256 * 1024**3 + +# Pinned dense TP=2 full-decode profile knobs (documentation / provenance only; +# the recipe script and the executed ServerArgs.__post_init__ hook enforce them). +ATTENTION_BACKEND = "flashinfer" +COLLECTIVE_BACKEND = "torch_symmetric_memory" +DECODE_GRAPH_BACKEND = "full" +PREFILL_GRAPH_BACKEND = "disabled" + +PROMPTS = [ + "Explain what tensor parallelism is in one paragraph.", + "Write a haiku about GPUs.", + "List three prime numbers greater than 50.", + "Summarize why CUDA graphs speed up inference.", +] +CONCURRENT_PROMPTS = ["Reply with only the result: 2 + 2 ="] * 8 +MAX_NEW_TOKENS = 64 +CONCURRENT_MAX_NEW_TOKENS = 4 +FOCUSED_PHASES = ("save", "load") +FOCUSED_PROMPT = "Explain what tensor parallelism is in one paragraph." +FOCUSED_MAX_NEW_TOKENS = 32 +KEEP_ARTIFACTS = os.environ.get("TP_KEEP_ARTIFACTS", "0") == "1" + +# Phase plan. ``save`` writes the archive; ``save2`` must reproduce it; ``load`` +# and ``load2`` are two independent fresh-process restorations of that archive. +PHASES = ("baseline", "save", "save2", "load", "load2") +LOAD_PHASES = ("load", "load2") +ARCHIVE_PHASES = ("save", "save2", "load", "load2") +FOUNDRY_PHASES = ("save", "save2", "load", "load2") +QUERY_PHASES = ("baseline", "save2", "load", "load2") +PHASE_MODES = { + "baseline": None, + "save": "save", + "save2": "save", + "load": "load", + "load2": "load", +} + +image = ( + modal.Image.from_registry(f"nvidia/cuda:{CUDA_TAG}", add_python="3.12") + .env( + { + "CC": "gcc", + "CXX": "g++", + "FOUNDRY_REF": FOUNDRY_REF, + "SGLANG_BASE": SGLANG_BASE, + "SGL_MODEL_REVISION": MODEL_REVISION, + } + ) + .apt_install( + "git", + "build-essential", + "cmake", + "ninja-build", + "libboost-all-dev", + "libnuma-dev", + ) + .pip_install( + "cmake>=4.0.0", + "wheel", + "packaging", + "ninja", + "setuptools>=80", + "setuptools-scm", + ) + .pip_install( + "torch==2.11.0", + "torchvision==0.26.0", + "torchaudio==2.11.0", + index_url="https://download.pytorch.org/whl/cu130", + ) + .run_commands( + f"git clone {SGLANG_REPO} /sglang", + f"git -C /sglang checkout --detach {SGLANG_BASE}", + f'test "$(git -C /sglang rev-parse HEAD)" = "{SGLANG_BASE}"', + "cd /sglang && pip install -e 'python[all]' --no-build-isolation", + ) + .run_commands( + f"git clone {FOUNDRY_REPO} /foundry", + f"git -C /foundry checkout --detach {FOUNDRY_REF}", + f'test "$(git -C /foundry rev-parse HEAD)" = "{FOUNDRY_REF}"', + "cd /foundry && pip install -e . --no-build-isolation", + ) + .env( + { + "HF_HOME": HF_CACHE, + "HUGGINGFACE_HUB_CACHE": f"{HF_CACHE}/hub", + "PYTHONUNBUFFERED": "1", + } + ) +) + +# Modal auto-mounts only the invoked entrypoint file, not its sibling modules. +# run_phase imports these helpers at module scope, so mount them explicitly onto +# the container PYTHONPATH (runtime mount; copy=False) or hydration fails with +# ModuleNotFoundError before any phase runs. +image = image.add_local_python_source( + "sglang_tp_validation", + "tp_validation_helpers", +) + +# cleanup_run only prunes an archive run directory, so it uses a slim CPU base +# rather than the heavy CUDA image. Modal still re-imports this module (with its +# top-level helper imports) to hydrate cleanup_run, so the slim image must mount +# the same sibling helpers with copy=False or cleanup_run fails to hydrate. +cleanup_image = ( + modal.Image.debian_slim() + .env( + { + "FOUNDRY_REF": FOUNDRY_REF, + "SGLANG_BASE": SGLANG_BASE, + "SGL_MODEL_REVISION": MODEL_REVISION, + "TP_SIZE": str(TP_SIZE), + } + ) + .add_local_python_source( + "sglang_tp_validation", + "tp_validation_helpers", + ) +) + +app = modal.App("foundry-sglang-tp-validation") +archive_volume = modal.Volume.from_name( + os.environ.get( + "TP_ARCHIVE_VOLUME", + "foundry-sglang-tp-validation-archive", + ), + create_if_missing=True, +) +hf_volume = modal.Volume.from_name( + os.environ.get("TP_HF_VOLUME", "foundry-sglang-tp-hf-cache"), + create_if_missing=True, +) + + +def _recipe_command(mode: str | None) -> list[str]: + command = ["bash", RECIPE_SCRIPT, str(TP_SIZE)] + if mode is not None: + command.append(f"--{mode}") + return command + + +def _log_tail(log_path: str, lines: int = 100) -> str: + path = Path(log_path) + if not path.exists(): + return "" + return "\n".join(path.read_text(errors="replace").splitlines()[-lines:]) + + +def _qualified_exception_type(error: BaseException) -> str: + error_type = type(error) + return f"{error_type.__module__}.{error_type.__qualname__}" + + +def _exception_cause_types(error: BaseException | None) -> tuple[str, ...]: + cause_types: list[str] = [] + seen: set[int] = set() + current = error + while current is not None and id(current) not in seen: + seen.add(id(current)) + cause_types.append(_qualified_exception_type(current)) + chained = current.__cause__ or current.__context__ + if chained is None: + reason = getattr(current, "reason", None) + chained = reason if isinstance(reason, BaseException) else None + current = chained + return tuple(cause_types) + + +class FocusedStartupHealthError(RuntimeError): + """Focused startup failure retaining the final health-check cause chain.""" + + def __init__( + self, + message: str, + *, + last_error: BaseException | None, + child_exited: bool, + ) -> None: + super().__init__(message) + self.cause_types = _exception_cause_types(last_error) + self.child_exited = child_exited + + +def _raise_startup_health_error( + message: str, + *, + last_error: BaseException | None, + child_exited: bool, +) -> None: + error = FocusedStartupHealthError( + message, + last_error=last_error, + child_exited=child_exited, + ) + if last_error is None: + raise error + raise error from last_error + + +def _wait_until_healthy( + process: subprocess.Popen, + log_path: str, + timeout: float, +) -> None: + url = f"http://127.0.0.1:{PORT}/health_generate" + deadline = time.time() + timeout + last_error: BaseException | None = None + + while time.time() < deadline: + if process.poll() is not None: + _raise_startup_health_error( + f"SGLang exited with {process.returncode} before becoming healthy\n" + f"{_log_tail(log_path)}", + last_error=last_error, + child_exited=True, + ) + try: + with urllib.request.urlopen(url, timeout=5) as response: + if response.status == 200: + return + except Exception as error: # noqa: BLE001 - preserve the last health error + last_error = error + time.sleep(3) + + last_error_text = str(last_error) if last_error is not None else "" + _raise_startup_health_error( + f"SGLang was not healthy after {timeout}s: {last_error_text}\n{_log_tail(log_path)}", + last_error=last_error, + child_exited=False, + ) + + +def _generate(prompt: str, max_new_tokens: int = MAX_NEW_TOKENS) -> str: + request = urllib.request.Request( + f"http://127.0.0.1:{PORT}/generate", + data=json.dumps( + { + "text": prompt, + "sampling_params": { + "temperature": 0.0, + "max_new_tokens": max_new_tokens, + }, + } + ).encode(), + headers={"Content-Type": "application/json"}, + ) + with urllib.request.urlopen(request, timeout=120) as response: + result = json.loads(response.read()) + return result["text"] if isinstance(result, dict) else result[0]["text"] + + +def _generate_concurrently() -> list[str]: + """Issue one native batch request; keep the historical field/check naming.""" + request = urllib.request.Request( + f"http://127.0.0.1:{PORT}/generate", + data=json.dumps( + { + "text": CONCURRENT_PROMPTS, + "sampling_params": { + "temperature": 0.0, + "max_new_tokens": CONCURRENT_MAX_NEW_TOKENS, + }, + } + ).encode(), + headers={"Content-Type": "application/json"}, + ) + with urllib.request.urlopen(request, timeout=120) as response: + result = json.loads(response.read()) + + if not isinstance(result, list): + raise ValueError("SGLang native batch response must be a list") + if len(result) != len(CONCURRENT_PROMPTS): + raise ValueError( + f"SGLang native batch response must contain exactly {len(CONCURRENT_PROMPTS)} items" + ) + + outputs: list[str] = [] + for index, item in enumerate(result): + if not isinstance(item, dict) or not isinstance(item.get("text"), str): + raise ValueError( + f"SGLang native batch response item {index} must contain a text string" + ) + outputs.append(item["text"]) + return outputs + + +def _launch( + mode: str | None, + log_path: str, + run_root: Path, +) -> tuple[subprocess.Popen, IO[str]]: + env = dict(os.environ) + env["SGL_MODEL"] = MODEL + env["SGL_MODEL_REVISION"] = MODEL_REVISION + env["SGLANG_PLUGINS"] = "foundry" + env["NCCL_DEBUG"] = "INFO" + + # The handle intentionally spans the child lifetime and is closed by stop_process. + log_file = open(log_path, "w") # noqa: SIM115 + process = subprocess.Popen( + _recipe_command(mode), + stdout=log_file, + stderr=subprocess.STDOUT, + env=env, + cwd=run_root, + start_new_session=True, + ) + return process, log_file + + +def _read_rank_offsets(workspace: Path) -> dict[str, int]: + offsets = {} + for rank in range(TP_SIZE): + path = workspace / f"rank_{rank}" / "final_alloc_offset.json" + if not path.exists(): + continue + offsets[str(rank)] = int(json.loads(path.read_text())["final_alloc_offset"]) + return offsets + + +def _archive_graph_counts(workspace: Path) -> dict[str, int]: + counts = {} + for rank in range(TP_SIZE): + rank_dir = workspace / f"rank_{rank}" + counts[str(rank)] = len(list(rank_dir.glob("graph_*.cugraph"))) + return counts + + +def _archive_fingerprints(workspace: Path) -> dict[str, dict[str, str]]: + return archive_fingerprints( + workspace, + TP_SIZE, + FOUNDRY_BASE_ADDR, + FOUNDRY_REGION_END, + ) + + +def _archive_byte_fingerprints(workspace: Path) -> dict[str, str]: + return archive_byte_fingerprints(workspace) + + +def _package_version(*distributions: str) -> str | None: + for distribution in distributions: + try: + return importlib.metadata.version(distribution) + except importlib.metadata.PackageNotFoundError: + continue + return None + + +def _nvidia_smi(query: str) -> str | None: + try: + output = subprocess.check_output( + ["nvidia-smi", f"--query-gpu={query}", "--format=csv,noheader"], + text=True, + stderr=subprocess.DEVNULL, + ) + except (OSError, subprocess.CalledProcessError): + return None + first = output.splitlines()[0].strip() if output.strip() else "" + return first or None + + +def _runtime_versions(nccl_version: str | None) -> dict[str, str | None]: + """Best-effort runtime / package provenance. + + Package versions come from installed distribution metadata (no import, so no + CUDA context is created in this orchestration process), the driver / GPU + come from ``nvidia-smi`` (a subprocess, likewise no CUDA init here), and the + NCCL runtime version is parsed from the child's NCCL_DEBUG banner. + """ + return { + "cuda_toolkit_tag": CUDA_TAG, + "driver": _nvidia_smi("driver_version"), + "gpu_name": _nvidia_smi("name"), + "torch": _package_version("torch"), + "sglang": _package_version("sglang"), + "flashinfer": _package_version("flashinfer-python", "flashinfer"), + "sgl_kernel": _package_version("sgl-kernel"), + "nccl": nccl_version, + } + + +def _candidate_metadata() -> dict[str, str | int]: + return { + "model": MODEL, + "model_revision": MODEL_REVISION, + "tp_size": TP_SIZE, + "attention_backend": ATTENTION_BACKEND, + "collective_backend": COLLECTIVE_BACKEND, + "decode_graph_backend": DECODE_GRAPH_BACKEND, + "prefill_graph_backend": PREFILL_GRAPH_BACKEND, + } + + +@app.function( + image=image, + gpu=MODAL_GPU, + timeout=3600, + volumes={ARCHIVE_ROOT: archive_volume, HF_CACHE: hf_volume}, +) +def run_phase( + phase: str, + run_id: str, + expected_fingerprints: dict | None = None, + expected_graph_counts: dict | None = None, + expected_byte_fingerprints: dict | None = None, + focused_treatment: bool = False, +) -> dict: + mode = PHASE_MODES[phase] + run_root = Path(RUNS_ROOT) / run_id + workspace = run_root / "foundry_archive_sglang_tp" + run_root.mkdir(parents=True, exist_ok=True) + log_path = str(run_root / f"{phase}.log") + print(f"[TP validation] run={run_id} phase={phase} mode={mode} gpu={MODAL_GPU}") + + # Only the first SAVE clears the archive; save2 must overwrite it in place and + # both loads read the same immutable archive. + if phase == "save": + shutil.rmtree(workspace, ignore_errors=True) + + started_at = time.time() + process, log_file = _launch(mode, log_path, run_root) + result = {"phase": phase, "artifact_dir": str(run_root), **_candidate_metadata()} + if focused_treatment: + result.update( + { + "became_healthy": False, + "request_completed": False, + "evidence_errors": [], + "failure_stage": None, + "failure_type": None, + "failure_events": [], + "startup_cause_types": [], + "startup_child_exited": False, + "archive_committed": False, + } + ) + phase_error: Exception | None = None + active_stage = "startup" + try: + _wait_until_healthy(process, log_path, timeout=1800) + result["time_to_healthy_s"] = round(time.time() - started_at, 1) + active_stage = "request" + if focused_treatment: + result["became_healthy"] = True + result["outputs"] = [_generate(FOCUSED_PROMPT, max_new_tokens=FOCUSED_MAX_NEW_TOKENS)] + result["request_completed"] = True + else: + if phase in QUERY_PHASES: + result["outputs"] = [_generate(prompt) for prompt in PROMPTS] + # Compatibility name: this is one explicit native scheduler batch. + result["concurrent_outputs"] = _generate_concurrently() + if process.poll() is not None: + raise RuntimeError( + f"SGLang exited unexpectedly after validation requests\n{_log_tail(log_path)}" + ) + except Exception as error: # noqa: BLE001 - scan complete logs before reporting + phase_error = error + if focused_treatment: + result["failure_stage"] = active_stage + result["failure_type"] = _qualified_exception_type(error) + if active_stage == "startup": + if isinstance(error, FocusedStartupHealthError): + result["startup_cause_types"] = list(error.cause_types) + result["startup_child_exited"] = error.child_exited + else: + result["startup_cause_types"] = list(_exception_cause_types(error)) + result["failure_events"].append( + { + "stage": active_stage, + "type": result["failure_type"], + } + ) + + # Snapshot the log *before* the intentional SIGINT teardown so the fatal scan + # can tell genuine pre-teardown failures apart from the KeyboardInterrupt + # tracebacks / force-kill noise the workers emit when we stop them. The full + # post-teardown log still supplies every other parsed signal. + pre_teardown_text = Path(log_path).read_text(errors="replace") + + teardown_error: Exception | None = None + try: + result["process_return_code"] = stop_process(process, log_file) + except Exception as error: # noqa: BLE001 - include scan evidence with teardown failure + teardown_error = error + result["process_return_code"] = process.returncode + if focused_treatment: + if result["failure_stage"] is None: + result["failure_stage"] = "teardown" + result["failure_type"] = _qualified_exception_type(error) + result["failure_events"].append( + { + "stage": "teardown", + "type": _qualified_exception_type(error), + } + ) + + try: + scan = scan_log_text( + Path(log_path).read_text(errors="replace"), + pre_teardown_text=pre_teardown_text, + ) + except Exception as error: + if not focused_treatment: + raise + result["evidence_errors"].append(f"log scanner: {error}") + if result["failure_stage"] is None: + result["failure_stage"] = "scanner" + result["failure_type"] = _qualified_exception_type(error) + result["failure_events"].append( + { + "stage": "scanner", + "type": _qualified_exception_type(error), + } + ) + scan = { + "errors": [], + "error_categories": ["scanner"], + "python_error_types": [], + "untyped_python_fatal_evidence": [], + "typed_python_traceback_headers": [], + "nccl_version": None, + "load_offsets": {}, + "loaded_graphs": {}, + "torch_symm_mem_ranks": [], + "multimem_all_gather_ranks": [], + "torch_symm_mem_addresses": {}, + "multimem_all_gather_addresses": {}, + "nccl_transport_channels": {}, + "graph_replay_batch_sizes": [], + "native_capture_progress_count": 0, + "saved_graph_log_count": 0, + } + if focused_treatment: + event_categories = { + event["stage"] + for event in result["failure_events"] + if event["stage"] in {"request", "teardown", "scanner", "archive"} + } + scan["error_categories"] = sorted(set(scan.get("error_categories", [])) | event_categories) + result.update(scan) + result["runtime_versions"] = _runtime_versions(scan.get("nccl_version")) + collect_archive_evidence = not focused_treatment or result["became_healthy"] + if collect_archive_evidence: + try: + if phase in {"save", "save2"}: + result["rank_offsets"] = _read_rank_offsets(workspace) + result["archive_graph_counts"] = _archive_graph_counts(workspace) + result["archive_fingerprints"] = _archive_fingerprints(workspace) + result["archive_byte_fingerprints"] = _archive_byte_fingerprints(workspace) + if phase in ARCHIVE_PHASES: + result["archive_graph_json_counts"] = archive_graph_json_counts( + workspace, + TP_SIZE, + ) + result["archive_collective_kernel_counts"] = archive_collective_kernel_counts( + workspace, + TP_SIZE, + ) + except Exception as error: + if not focused_treatment: + raise + result["evidence_errors"].append(f"archive scanner: {error}") + if result["failure_stage"] is None: + result["failure_stage"] = "archive" + result["failure_type"] = _qualified_exception_type(error) + result["failure_events"].append( + { + "stage": "archive", + "type": _qualified_exception_type(error), + } + ) + result["error_categories"] = sorted( + set(result.get("error_categories", [])) | {"archive"} + ) + + # Archive immutability proof: a LOAD must never rewrite the archive it read. + # Re-read the on-disk inventory / semantic fingerprints / full byte map and + # require they still equal the SAVE2 archive that produced them, before + # returning and before any commit. The byte map (every regular file under the + # archive root, including root warmup_state.json and every .cugraph) closes + # the same-container / same-mount mutation gaps the semantic map cannot see. + archive_mutation_error: str | None = None + if phase in LOAD_PHASES and collect_archive_evidence and not result.get("evidence_errors"): + try: + result["archive_graph_counts"] = _archive_graph_counts(workspace) + result["archive_fingerprints"] = _archive_fingerprints(workspace) + result["archive_byte_fingerprints"] = _archive_byte_fingerprints(workspace) + except Exception as error: + if not focused_treatment: + raise + result["evidence_errors"].append(f"archive fingerprint scanner: {error}") + if result["failure_stage"] is None: + result["failure_stage"] = "archive" + result["failure_type"] = _qualified_exception_type(error) + result["failure_events"].append( + { + "stage": "archive", + "type": _qualified_exception_type(error), + } + ) + result["error_categories"] = sorted( + set(result.get("error_categories", [])) | {"archive"} + ) + else: + if ( + expected_graph_counts is not None + and result["archive_graph_counts"] != expected_graph_counts + ): + archive_mutation_error = ( + f"LOAD mutated archive graph counts: " + f"{result['archive_graph_counts']} != {expected_graph_counts}" + ) + elif ( + expected_fingerprints is not None + and result["archive_fingerprints"] != expected_fingerprints + ): + archive_mutation_error = "LOAD mutated archive semantic fingerprints vs SAVE2" + elif ( + expected_byte_fingerprints is not None + and result["archive_byte_fingerprints"] != expected_byte_fingerprints + ): + archive_mutation_error = "LOAD mutated archive bytes vs SAVE2" + + if focused_treatment and archive_mutation_error is not None: + if result["failure_stage"] is None: + result["failure_stage"] = "archive" + result["failure_type"] = "foundry.validation.ArchiveMutationError" + result["failure_events"].append( + { + "stage": "archive", + "type": "foundry.validation.ArchiveMutationError", + } + ) + result["error_categories"] = sorted(set(result.get("error_categories", [])) | {"archive"}) + + failures = [] + if phase_error is not None: + failures.append(f"phase execution: {phase_error}") + if teardown_error is not None: + failures.append(f"process teardown: {teardown_error}") + if archive_mutation_error is not None: + failures.append(f"archive immutability: {archive_mutation_error}") + if result["errors"]: + failures.append(f"fatal log signatures: {result['errors']}") + if focused_treatment and result["evidence_errors"]: + failures.append(f"evidence extraction: {result['evidence_errors']}") + if focused_treatment: + result["phase_failures"] = list(failures) + if phase == "save": + save_checks = focused_save_checks(result, tp_size=TP_SIZE) + failed_save_checks = [ + name for name, passed in save_checks.items() if passed is not True + ] + if failed_save_checks: + failures.append("focused SAVE prerequisites: " + ", ".join(failed_save_checks)) + if result["failure_stage"] is None: + scanner_checks = { + "symmetric_markers_present", + "no_runtime_errors", + "no_scanner_errors", + "no_policy_failure", + } + archive_checks = { + "graph_counts_exact", + "graph_json_counts_exact", + "graph_inventory_matches", + "symmetric_archive_collectives_present", + "no_archived_nccl_collectives", + "offsets_equal_nonzero", + "archive_evidence_complete", + } + if "request_completed" in failed_save_checks: + failure_stage = "request" + elif "expected_shutdown" in failed_save_checks: + failure_stage = "teardown" + elif scanner_checks & set(failed_save_checks): + failure_stage = "scanner" + elif archive_checks & set(failed_save_checks): + failure_stage = "archive" + else: + failure_stage = "scanner" + failure_type = "foundry.validation.FocusedPrerequisiteError" + result["failure_stage"] = failure_stage + result["failure_type"] = failure_type + result["failure_events"].append( + { + "stage": failure_stage, + "type": failure_type, + } + ) + result["error_categories"] = sorted( + set(result.get("error_categories", [])) | {failure_stage} + ) + result["phase_failures"] = failures + + # Request an explicit commit only for a completely successful archive-writing + # phase. Modal auto-commit may retain partial or mutated unique-run artifacts + # after any failed phase. Sequential orchestration stops at that failure and + # must never classify the run VERIFIED. + if phase in {"save", "save2"} and not failures: + archive_volume.commit() + if focused_treatment: + result["archive_committed"] = True + + if failures: + result["status"] = "NOT VERIFIED" + result["artifacts_retained"] = True + result["artifacts_may_be_partial_or_mutated"] = True + print(json.dumps(result, sort_keys=True)) + if failures and not focused_treatment: + raise RuntimeError(f"SGLang phase {phase} failed: {'; '.join(failures)}") + return result + + +@app.function( + image=cleanup_image, + volumes={ARCHIVE_ROOT: archive_volume}, +) +def cleanup_run(run_id: str) -> None: + run_dir = Path(RUNS_ROOT) / run_id + if run_dir.exists(): + shutil.rmtree(run_dir) + if run_dir.exists(): + raise RuntimeError(f"cleanup did not remove archive run directory: {run_dir}") + archive_volume.commit() + + +def _run_focused_treatment(run_id: str) -> None: + results: dict[str, dict] = {} + remote_error: str | None = None + remote_error_type: str | None = None + + def report_and_raise_if_needed( + checks: dict[str, bool | None], + *, + environment_failure: bool, + status_override: str | None = None, + cleanup_attempted: bool = False, + cleanup_completed: bool = False, + cleanup_skipped: bool = True, + artifacts_retained: bool | None = None, + additional_failed_checks: tuple[str, ...] = (), + cleanup_error: str | None = None, + cleanup_error_type: str | None = None, + ) -> str: + restored_replay = checks.get("restored_graph_replay_observed") is True + status = status_override or classify_focused_checks( + checks, environment_failure=environment_failure, restored_replay=restored_replay + ) + failed = [name for name, passed in checks.items() if passed is False] + list( + additional_failed_checks + ) + unavailable = [name for name, passed in checks.items() if passed is None] + if artifacts_retained is None: + # Even without our explicit commit, Modal may auto-commit the unique + # run directory when a failed phase container exits. + artifacts_retained = status != "VERIFIED" + report = { + "run_id": run_id, + "artifacts_retained": artifacts_retained, + "artifacts_may_be_partial_or_mutated": (artifacts_retained and status != "VERIFIED"), + "cleanup_attempted": cleanup_attempted, + "cleanup_completed": cleanup_completed, + "cleanup_skipped": cleanup_skipped, + "cleanup_error": cleanup_error, + "cleanup_error_type": cleanup_error_type, + "foundry_ref": FOUNDRY_REF, + "sglang_base": SGLANG_BASE, + "model_revision": MODEL_REVISION, + "gpu": MODAL_GPU, + **_candidate_metadata(), + "classification": { + "environment_failure": environment_failure, + "known_failure": bool(failed), + "restored_replay": restored_replay, + }, + "remote_error": remote_error, + "remote_error_type": remote_error_type, + "failed_checks": failed, + "unavailable_checks": unavailable, + "checks": checks, + "results": results, + } + report["status"] = status + if status != "VERIFIED": + print(json.dumps(report, indent=2, sort_keys=True)) + detail = remote_error or ", ".join(failed or unavailable) + raise RuntimeError(f"Focused symmetric TP treatment {status}: {detail}") + print(json.dumps(report, indent=2, sort_keys=True)) + return status + + for phase in FOCUSED_PHASES: + try: + if phase == "load": + save_result = results["save"] + results[phase] = run_phase.remote( + phase, + run_id, + expected_fingerprints=save_result.get("archive_fingerprints", {}), + expected_graph_counts=save_result.get("archive_graph_counts", {}), + expected_byte_fingerprints=save_result.get( + "archive_byte_fingerprints", + {}, + ), + focused_treatment=True, + ) + else: + results[phase] = run_phase.remote( + phase, + run_id, + focused_treatment=True, + ) + except Exception as error: # noqa: BLE001 - classify and print before raising + remote_error = f"{phase}: {error}" + remote_error_type = _qualified_exception_type(error) + checks = ( + focused_save_checks( + results.get("save"), + tp_size=TP_SIZE, + require_committed=True, + ) + if phase == "save" + else focused_treatment_checks( + results.get("save"), + results.get("load"), + tp_size=TP_SIZE, + ) + ) + report_and_raise_if_needed( + checks, + environment_failure=is_modal_infrastructure_error_type( + type(error).__module__, + type(error).__name__, + ), + ) + + if phase == "save": + save_result = results["save"] + save_checks = focused_save_checks( + save_result, + tp_size=TP_SIZE, + require_committed=True, + ) + save_environment_failure = is_allowlisted_startup_environment_failure(save_result) + if not all(value is True for value in save_checks.values()): + report_and_raise_if_needed( + save_checks, + environment_failure=save_environment_failure, + ) + + checks = focused_treatment_checks( + results.get("save"), + results.get("load"), + tp_size=TP_SIZE, + ) + environment_failure = is_allowlisted_startup_environment_failure(results.get("load", {})) + status = classify_focused_checks( + checks, + environment_failure=environment_failure, + restored_replay=checks.get("restored_graph_replay_observed") is True, + ) + if status != "VERIFIED": + report_and_raise_if_needed(checks, environment_failure=environment_failure) + if KEEP_ARTIFACTS: + report_and_raise_if_needed( + checks, + environment_failure=environment_failure, + cleanup_skipped=True, + artifacts_retained=True, + ) + return + try: + cleanup_run.remote(run_id) + except Exception as error: # noqa: BLE001 - cleanup is part of verification + remote_error = f"cleanup: {error}" + remote_error_type = _qualified_exception_type(error) + report_and_raise_if_needed( + checks, + environment_failure=False, + status_override="NOT VERIFIED", + cleanup_attempted=True, + cleanup_completed=False, + cleanup_skipped=False, + artifacts_retained=True, + additional_failed_checks=("cleanup_completed",), + cleanup_error=str(error), + cleanup_error_type=remote_error_type, + ) + report_and_raise_if_needed( + checks, + environment_failure=environment_failure, + cleanup_attempted=True, + cleanup_completed=True, + cleanup_skipped=False, + artifacts_retained=False, + ) + + +@app.local_entrypoint() +def main(focused: bool = False) -> None: + safe_ref = re.sub(r"[^A-Za-z0-9_.-]", "-", FOUNDRY_REF)[:24] + run_id = f"{safe_ref}-{time.time_ns()}" + if focused: + _run_focused_treatment(run_id) + return + # Sequential, blocking .remote() calls preserve the baseline -> save -> save2 + # -> load -> load2 ordering the archive lifecycle depends on. Each fresh LOAD + # is handed the SAVE2 archive inventory/fingerprints so run_phase can prove, + # in-container, that it did not mutate the immutable archive. + results: dict[str, dict] = {} + for phase in PHASES: + if phase in LOAD_PHASES: + results[phase] = run_phase.remote( + phase, + run_id, + expected_fingerprints=results["save2"].get("archive_fingerprints", {}), + expected_graph_counts=results["save2"].get("archive_graph_counts", {}), + expected_byte_fingerprints=results["save2"].get("archive_byte_fingerprints", {}), + ) + else: + results[phase] = run_phase.remote(phase, run_id) + + expected_ranks = {str(rank) for rank in range(TP_SIZE)} + expected_graph_counts = {"0": 1, "1": 1} + required_archive_files = { + "graph_manifest.json", + "fatbin_image_packed.img", + "final_alloc_offset.json", + } + + baseline_outputs = results["baseline"].get("outputs", []) + baseline_concurrent = results["baseline"].get("concurrent_outputs", []) + save_offsets = results["save"].get("rank_offsets", {}) + save2_offsets = results["save2"].get("rank_offsets", {}) + save_fingerprints = results["save"].get("archive_fingerprints", {}) + save2_fingerprints = results["save2"].get("archive_fingerprints", {}) + save2_graph_counts = results["save2"].get("archive_graph_counts", {}) + save2_byte_fingerprints = results["save2"].get("archive_byte_fingerprints", {}) + + def _outputs_match_baseline(phase: str) -> bool: + return bool(baseline_outputs) and results[phase].get("outputs", []) == baseline_outputs + + def _concurrent_matches_baseline(phase: str) -> bool: + return ( + bool(baseline_concurrent) + and results[phase].get("concurrent_outputs", []) == baseline_concurrent + ) + + def _rank_fingerprint_ok(rank: str, fingerprint: dict) -> bool: + graph_json_count = sum( + name.startswith("graph_") and name.endswith(".json") and name != "graph_manifest.json" + for name in fingerprint + ) + return ( + required_archive_files <= set(fingerprint) + and graph_json_count == expected_graph_counts[rank] + ) + + def _symmetric_archive_collectives_present(phase: str) -> bool: + counts = results[phase].get("archive_collective_kernel_counts", {}) + return set(counts) == expected_ranks and all( + isinstance(rank_counts, dict) + and type(rank_counts.get("symmetric_all_reduce")) is int + and rank_counts["symmetric_all_reduce"] > 0 + and type(rank_counts.get("symmetric_all_gather")) is int + and rank_counts["symmetric_all_gather"] > 0 + for rank_counts in counts.values() + ) + + def _no_archived_nccl_collectives(phase: str) -> bool: + counts = results[phase].get("archive_collective_kernel_counts", {}) + return set(counts) == expected_ranks and all( + isinstance(rank_counts, dict) + and rank_counts.get("all_reduce") == 0 + and rank_counts.get("all_gather") == 0 + for rank_counts in counts.values() + ) + + def _symmetric_markers_present(phase: str) -> bool: + addresses = results[phase].get("multimem_all_gather_addresses", {}) + return ( + set(results[phase].get("torch_symm_mem_ranks", [])) == expected_ranks + and isinstance(addresses, dict) + and any(addresses.values()) + ) + + archive_graph_counts = { + phase: results[phase].get("archive_graph_counts", {}) for phase in ARCHIVE_PHASES + } + archive_graph_json_counts = { + phase: results[phase].get("archive_graph_json_counts", {}) for phase in ARCHIVE_PHASES + } + load_fingerprints = { + phase: results[phase].get("archive_fingerprints", {}) for phase in LOAD_PHASES + } + load_graph_counts = { + phase: results[phase].get("archive_graph_counts", {}) for phase in LOAD_PHASES + } + load_byte_fingerprints = { + phase: results[phase].get("archive_byte_fingerprints", {}) for phase in LOAD_PHASES + } + + checks = { + "deterministic_outputs_match_baseline": len(baseline_outputs) == len(PROMPTS) + and all(output.strip() for output in baseline_outputs) + and all(_outputs_match_baseline(phase) for phase in ("save2", "load", "load2")), + "concurrent_outputs_match_baseline": len(baseline_concurrent) == len(CONCURRENT_PROMPTS) + and all(output.strip() for output in baseline_concurrent) + and all(_concurrent_matches_baseline(phase) for phase in ("save2", "load", "load2")), + "save_offsets_present": set(save_offsets) == expected_ranks and all(save_offsets.values()), + "save_rank_offsets_symmetric": len(set(save_offsets.values())) == 1, + "save_offsets_reproducible": bool(save_offsets) and save_offsets == save2_offsets, + "load_offsets_match_save2": bool(save2_offsets) + and all(results[phase].get("load_offsets", {}) == save2_offsets for phase in LOAD_PHASES), + "archive_fingerprints_present": set(save_fingerprints) == expected_ranks + and all(_rank_fingerprint_ok(rank, fp) for rank, fp in save_fingerprints.items()), + "archive_fingerprints_reproducible": bool(save_fingerprints) + and save_fingerprints == save2_fingerprints, + "archive_graph_counts_exact": all( + archive_graph_counts[phase] == expected_graph_counts for phase in ARCHIVE_PHASES + ), + "archive_graph_json_counts_exact": all( + archive_graph_json_counts[phase] == expected_graph_counts for phase in ARCHIVE_PHASES + ), + "graphs_restored_each_rank": all( + results[phase].get("loaded_graphs", {}) == expected_graph_counts + for phase in LOAD_PHASES + ), + "symmetric_archive_collectives_each_phase": all( + _symmetric_archive_collectives_present(phase) for phase in ARCHIVE_PHASES + ), + "no_archived_nccl_collectives": all( + _no_archived_nccl_collectives(phase) for phase in ARCHIVE_PHASES + ), + "symmetric_markers_each_phase": all( + _symmetric_markers_present(phase) for phase in FOUNDRY_PHASES + ), + "load_archive_immutable": bool(save2_fingerprints) + and bool(save2_graph_counts) + and all( + load_fingerprints[phase] == save2_fingerprints + and load_graph_counts[phase] == save2_graph_counts + for phase in LOAD_PHASES + ), + "load_archive_bytes_immutable": bool(save2_byte_fingerprints) + and all(load_byte_fingerprints[phase] == save2_byte_fingerprints for phase in LOAD_PHASES), + "restored_graph_replay_observed": all( + 1 in results[phase]["graph_replay_batch_sizes"] + and any(batch_size > 1 for batch_size in results[phase]["graph_replay_batch_sizes"]) + for phase in LOAD_PHASES + ), + # The pinned LOAD lifecycle traverses the real upstream capture() loop, + # so rank0 emits "Capturing batches" progress on LOAD even though the + # Foundry backend performs no CUDA capture. That progress count is + # diagnostic only; no-recapture is proven by positive archived-load + # counts together with zero Foundry graph-save logs. + "load_restored_without_resave": all( + results[phase]["saved_graph_log_count"] == 0 + and sum((results[phase].get("loaded_graphs") or {}).values()) > 0 + for phase in LOAD_PHASES + ), + "no_runtime_errors": not any(results[phase]["errors"] for phase in PHASES), + "expected_process_shutdown": all( + results[phase]["process_return_code"] in {0, -2, 130} for phase in PHASES + ), + } + + report = { + "run_id": run_id, + "foundry_ref": FOUNDRY_REF, + "sglang_base": SGLANG_BASE, + "model_revision": MODEL_REVISION, + "gpu": MODAL_GPU, + **_candidate_metadata(), + "process_return_codes": {phase: results[phase]["process_return_code"] for phase in PHASES}, + "checks": checks, + "results": results, + } + + failed = [name for name, passed in checks.items() if not passed] + if failed: + report.update( + { + "status": "NOT VERIFIED", + "failed_checks": failed, + "artifacts_retained": True, + "artifacts_may_be_partial_or_mutated": True, + "cleanup_attempted": False, + "cleanup_completed": False, + "cleanup_skipped": True, + "cleanup_error": None, + "cleanup_error_type": None, + } + ) + print(json.dumps(report, indent=2, sort_keys=True)) + raise RuntimeError(f"TP validation NOT VERIFIED: {', '.join(failed)}") + try: + cleanup_run.remote(run_id) + except Exception as error: # noqa: BLE001 - cleanup is part of verification + error_type = _qualified_exception_type(error) + report.update( + { + "status": "NOT VERIFIED", + "failed_checks": ["cleanup_completed"], + "artifacts_retained": True, + "artifacts_may_be_partial_or_mutated": True, + "cleanup_attempted": True, + "cleanup_completed": False, + "cleanup_skipped": False, + "cleanup_error": str(error), + "cleanup_error_type": error_type, + } + ) + print(json.dumps(report, indent=2, sort_keys=True)) + raise RuntimeError(f"TP validation NOT VERIFIED: cleanup: {error}") from error + report.update( + { + "status": "VERIFIED", + "failed_checks": [], + "artifacts_retained": False, + "artifacts_may_be_partial_or_mutated": False, + "cleanup_attempted": True, + "cleanup_completed": True, + "cleanup_skipped": False, + "cleanup_error": None, + "cleanup_error_type": None, + } + ) + print(json.dumps(report, indent=2, sort_keys=True)) diff --git a/tests/modal_vllm_tp.py b/tests/modal_vllm_tp.py new file mode 100644 index 00000000..1c3c5875 --- /dev/null +++ b/tests/modal_vllm_tp.py @@ -0,0 +1,2279 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the Foundry project +"""End-to-end vLLM tensor-parallel validation on Modal. + +The harness runs six phases on two GPUs, each in a fresh server process: + +1. ``baseline`` — a normal vLLM TP run and deterministic queries; the reference + completion evidence and the reference NCCL transport set. +2. ``seed`` — a SAVE that records vLLM's memory-profile warmup state. +3. ``save`` — a deterministic SAVE using that warmup state. +4. ``save2`` — a second deterministic SAVE that must reproduce the per-rank + allocation offsets, graph inventory, and semantic fingerprints, and whose + completion outputs must equal the baseline. +5. ``load`` — Foundry LOAD in a fresh server process and the same queries. +6. ``load2`` — a second, independent LOAD against the *same* published archive. + +Storage model (Modal Dict publication authority). +-------------------------------------------------- +Modal automatically commits a mounted Volume when the container exits, so it is +not enough to merely *avoid* an explicit commit — a server that writes into the +mounted archive would publish partial/unvalidated bytes on exit regardless, and +there is no filesystem primitive that makes an on-volume marker reliably the +*last* durable byte. Publication authority therefore lives OUTSIDE the volume, in +a dedicated persistent :data:`publication_index` (a ``modal.Dict``). The volume +holds bytes; the Dict decides which bytes are real. + +* ``baseline`` / ``seed`` / ``save`` / ``save2`` run entirely in ephemeral local + (non-volume) storage under :data:`LOCAL_ROOT`. No server phase ever runs + against the mounted volume. +* Only after the whole transaction is validated is the archive published, keyed + by a unique run id (``--``; the random suffix makes + collisions negligible). Publication: build the bounded ``publication_index`` + record (identity ``{snapshot_digest, published_path, foundry_ref}`` plus + ``run_id`` and the prevalidated transaction payload) and preflight both its + JSON and pickle serialized size against :data:`PUBLICATION_RECORD_MAX_BYTES`; if + an identical record already exists, return idempotently, and if a different one + exists, fail the collision; ``mkdir`` the canonical run dir with + ``exist_ok=False`` as defense-in-depth; copy the validated archive into an + ``_incoming`` child, revalidate, atomically rename ``_incoming`` → ``archive``; + write ``published//manifest.json``; commit the volume with bounded + retries; then, as the single final side effect, atomically + ``publication_index.put(run_id, record, skip_if_exists=True)``. That put is the + authoritative marker: ``False`` (or an ambiguous exception) is reconciled by + fetching the stored record — an exact match is idempotent success, any + difference is a collision. The Dict is never written before a successful commit, + and the only work after a reconciled marker is returning the already-serialized + payload. +* ``load`` / ``load2`` first call ``archive_volume.reload()`` (so a warm container + sees the just-committed bytes and Dict state), then require an exact + ``publication_index`` record identity (matching the expected snapshot digest, + path, and Foundry ref) AND a matching on-volume ``manifest.json`` and + ``archive`` child before copying ``published//archive`` into their own + independent local storage. They execute there and compare the local copy before + vs after; they never write the mounted archive or the Dict. + +Modal semantics (accurate, not aspirational). Modal can auto-commit an +in-progress publication and does not expose an atomic "last byte" on a Volume, so +a reserved-but-unfinished ``published/`` directory can physically persist. +Such a directory is *inert*: it has no ``publication_index`` record, so no +consumer will ever copy it. If the explicit commit fails the function raises +before writing the marker, leaving the (possibly auto-committed) archive +unindexed and inert — no rollback is attempted or needed. If the marker put fails +unambiguously the archive likewise stays inert. Because the marker may land yet +its acknowledgement (or the function's result delivery) be lost, the marker write +reconciles ambiguous outcomes by fetching the record, and the local entrypoint +does the same when ``run_save_transaction.remote`` raises: an exact stored record +means the transaction actually succeeded (recover its payload), otherwise the +original error stands. There is thus no ``ExceptionGroup`` and no on-volume +rollback. Task 7 cleanup deletes both the canonical archive and its Dict entry on +success via :func:`delete_published_run` (the record carries ``run_id`` and +``published_path`` for exactly that). + +Success requires exact bounded completion text/token-string sequences across +baseline / save2 / load / load2, reproducible SAVE/SAVE2 offsets, counts, and +semantic fingerprints (raw ``.cugraph`` byte equality is NOT required), exactly +64 paired JSON/cugraph graph files per rank, per-rank allocation offsets equal to +SAVE2 on both LOADs (from the live post-capture VMM cursor), exactly one 64/64 +restoration per rank per LOAD, no graph recapture on LOAD, byte-immutable local +archives across both LOADs, NCCL P2P/IPC in every phase, clean process shutdown, +and no fatal runtime logs. + +Run the current checkout: + + modal run --env rahul-dev tests/modal_vllm_tp.py + +Run a specific commit: + + FOUNDRY_REF= \ + modal run --env rahul-dev tests/modal_vllm_tp.py + +Both Foundry and vLLM resolve to immutable commits before the remote image is +built so Modal cannot reuse a stale image for a moving branch. +""" + +from __future__ import annotations + +import contextlib +import copy +import hashlib +import importlib.metadata +import json +import os +import pickle +import re +import secrets +import shutil +import subprocess +import threading +import time +import urllib.request +from concurrent.futures import ThreadPoolExecutor +from pathlib import Path +from typing import IO + +import modal +from tp_validation_helpers import ( + archive_byte_fingerprints, + archive_fingerprints, + collect_archive_snapshot, + reject_symlinks, +) +from vllm_tp_validation import ( + completion_evidence, + require_commit_sha, + scan_log_text, + stop_process, +) + + +def _local_foundry_revision() -> str: + try: + revision = subprocess.check_output( + ["git", "rev-parse", "HEAD"], + cwd=Path(__file__).parents[1], + stderr=subprocess.DEVNULL, + text=True, + ).strip() + except (OSError, subprocess.CalledProcessError) as error: + raise RuntimeError( + "FOUNDRY_REF is required when the local checkout revision cannot be resolved" + ) from error + return require_commit_sha("local Foundry HEAD", revision) + + +CUDA_TAG = "13.0.1-cudnn-devel-ubuntu24.04" +REQUESTED_TORCH_VERSION = "2.11.0" + +FOUNDRY_REPO = "https://github.com/modal-projects/foundry.git" +FOUNDRY_REF = require_commit_sha( + "FOUNDRY_REF", + os.environ.get("FOUNDRY_REF") or _local_foundry_revision(), +) + +VLLM_REPO = "https://github.com/foundry-org/vllm.git" +VLLM_REF = require_commit_sha( + "VLLM_REF", + os.environ.get("VLLM_REF", "4309c257d3f639e5490d3811293c890c61c76f29"), +) +# The Foundry fork commit is one Python-only commit on top of this vLLM base. +# Pin its matching CUDA wheel: falling back to the moving nightly wheel can +# silently omit extensions that this older source tree imports as ``vllm._C``. +VLLM_WHEEL_COMMIT = "6cbe448eed751824d608faf9078ef84724d621c1" + +MODEL = "Qwen/Qwen3-8B" +MODEL_REVISION = require_commit_sha( + "VLLM_MODEL_REVISION", + os.environ.get("VLLM_MODEL_REVISION", "b968826d9c46dd6066d109eabc6255188de91218"), +) +VLLM_WHEEL_COMMIT = require_commit_sha( + "VLLM_WHEEL_COMMIT", + os.environ.get("VLLM_WHEEL_COMMIT", VLLM_WHEEL_COMMIT), +) +if os.environ.get("TP_SIZE", "2") != "2": + raise ValueError("TP_SIZE must be exactly 2") +TP_SIZE = 2 +MODAL_GPU = os.environ.get("MODAL_GPU", f"H100:{TP_SIZE}") +PORT = 12000 +CUDAGRAPH_MODE = "FULL_DECODE_ONLY" +DISTRIBUTED_EXECUTOR_BACKEND = "mp" +ATTENTION_BACKEND = "FLASH_ATTN" +COLLECTIVE_BACKEND = "NCCL_P2P_IPC" +# The NCCL channel transport token every rank must report in every phase for +# this single-node TP=2 profile (legacy P2P/IPC, not NET/symmetric-memory). +COLLECTIVE_TRANSPORT = "P2P/IPC" + +# Mounted archive volume layout. Only validated, published runs live under +# PUBLISHED_ROOT; each is a unique run directory gated for consumption by a +# matching publication_index record (the authority) plus its on-volume +# manifest.json. Overridable so CPU contracts can point it at a tmp dir. +ARCHIVE_ROOT = os.environ.get("VLLM_TP_ARCHIVE_ROOT", "/data/archive") +PUBLISHED_ROOT = f"{ARCHIVE_ROOT}/published" +# Ephemeral, NON-volume local storage. Every server phase and every LOAD copy +# lives here so Modal's automatic volume commit on container exit can never +# publish an unvalidated or partially written archive. +LOCAL_ROOT = os.environ.get("VLLM_TP_LOCAL_ROOT", "/tmp/foundry-vllm-tp-runs") +HF_CACHE = "/data/hf" +RECIPE_SCRIPT = "/foundry/recipe/vllm/serve_qwen3-8b_tp.sh" +FOUNDRY_BASE_ADDR = 0x600000000000 +FOUNDRY_REGION_END = FOUNDRY_BASE_ADDR + 256 * 1024**3 + +# Exactly 64 captured graphs per rank are required for archive acceptance. +EXPECTED_GRAPH_COUNT = 64 + +# Upper bound on graph filenames listed in a bounded diagnostic (seed's +# non-canonical inventory and SAVE archive-evidence failures) so a report stays +# legible without dumping the whole archive. +GRAPH_DIAGNOSTIC_SAMPLE_LIMIT = 8 + +# Bounded retries for the durability commit before publication is abandoned. +COMMIT_MAX_ATTEMPTS = 3 + +# The publication-index record must stay small: it holds only bounded primitives — +# the identity fields, the immutable snapshot (which itself carries one semantic +# and one byte digest per archived file, not the file bytes), and a reduced +# transaction payload. Raw logs and the large per-phase validated inventories are +# excluded. Its serialized size (the larger of JSON and pickle) is preflighted +# against this documented limit before any publish so a marker write can never +# fail for size. The limit is generous relative to the observed record size (a +# few tens of KiB for TP=2 x 64 graphs). +PUBLICATION_RECORD_MAX_BYTES = 512 * 1024 + +# The only phase-result fields carried in the stored payload: the primitives that +# ``main`` reads for its checks/report. The large per-phase evidence (the full +# validated archive inventory, the per-phase per-file byte-fingerprint map, +# artifact paths, raw log-derived lists) is excluded from the phase results; the +# immutable snapshot separately carries the one fingerprint map LOAD needs. +# ``seed_diagnostics`` is seed's bounded non-canonical inventory; ``save`` / +# ``save2`` instead carry ``rank_offsets`` / ``archive_graph_counts`` / +# ``archive_fingerprints``. The presence-based reduction naturally keeps these +# per-phase field sets disjoint. +_STORED_PHASE_FIELDS = ( + "phase", + "foundry_ref", + "vllm_ref", + "requested_vllm_ref", + "vllm_wheel_commit", + "model", + "model_revision", + "cuda_image_tag", + "requested_torch_version", + "requested_gpu", + "tp_size", + "cudagraph_mode", + "distributed_executor_backend", + "attention_backend", + "collective_backend", + "process_return_code", + "nccl_transports", + "errors", + "outputs", + "concurrent_outputs", + "rank_offsets", + "archive_graph_counts", + "archive_fingerprints", + "seed_diagnostics", + "runtime_provenance", + "time_to_healthy_s", + "phase_time_s", +) + +# A published run directory holds the validated archive under ``archive/`` plus a +# descriptive ``manifest.json``. During publication the archive first lands under +# an ``_incoming`` child that is atomically renamed to ``archive``. The manifest +# is a convenience/defense record only: publication authority is the separate +# ``publication_index`` Dict entry, not any on-volume file. +MANIFEST_NAME = "manifest.json" +PUBLISHED_ARCHIVE_DIRNAME = "archive" +INCOMING_ARCHIVE_DIRNAME = "_incoming" + + +def _expected_nccl_transports() -> dict[str, list[str]]: + """The strict per-rank transport profile: exactly P2P/IPC for every rank.""" + return {str(rank): [COLLECTIVE_TRANSPORT] for rank in range(TP_SIZE)} + + +# Phase plan. ``seed`` / ``save`` / ``save2`` are one atomic archive-writing +# transaction (all local); ``load`` / ``load2`` are two independent fresh-process +# restorations of the published SAVE2 archive (each into its own local copy). +PHASES = ("baseline", "seed", "save", "save2", "load", "load2") +QUERY_PHASES = ("baseline", "save2", "load", "load2") +LOAD_PHASES = ("load", "load2") +SAVE_TRANSACTION_PHASES = ("seed", "save", "save2") +PHASE_MODES = { + "baseline": None, + "seed": "save", + "save": "save", + "save2": "save", + "load": "load", + "load2": "load", +} + +PROMPTS = [ + "The capital of France is", + "Reply with only the result: 2 + 2 =", + "The chemical formula for water is", + "A GPU is a graphics processing", +] +MAX_TOKENS = 4 +CONCURRENT_BARRIER_TIMEOUT_S = 30.0 + +EARLY_BUILDS_PATTERN = re.compile( + r"\[foundry\] Starting early graph builds \((?P\d+) graphs" +) +NCCL_TRANSPORT_PATTERN = re.compile( + r"\[(?P\d+)\] NCCL INFO Channel \d+/\d+ : .* via (?P\S+)" +) + +image = ( + modal.Image.from_registry(f"nvidia/cuda:{CUDA_TAG}", add_python="3.12") + .env( + { + "CC": "gcc", + "CXX": "g++", + "FOUNDRY_REF": FOUNDRY_REF, + "VLLM_REF": VLLM_REF, + "VLLM_WHEEL_COMMIT": VLLM_WHEEL_COMMIT, + "VLLM_MODEL_REVISION": MODEL_REVISION, + "TP_SIZE": str(TP_SIZE), + "MODAL_GPU": MODAL_GPU, + } + ) + .apt_install( + "git", + "build-essential", + "cmake", + "ninja-build", + "libboost-all-dev", + "libnuma-dev", + ) + .pip_install( + "cmake>=4.0.0", + "wheel", + "packaging", + "ninja", + "setuptools>=80", + "setuptools-scm", + "uv", + ) + .pip_install( + f"torch=={REQUESTED_TORCH_VERSION}", + "torchvision==0.26.0", + "torchaudio==2.11.0", + index_url="https://download.pytorch.org/whl/cu130", + ) + .run_commands( + f"git clone {VLLM_REPO} /vllm", + f"cd /vllm && git checkout {VLLM_REF}", + "cd /vllm && VLLM_USE_PRECOMPILED=1 " + f"VLLM_PRECOMPILED_WHEEL_COMMIT={VLLM_WHEEL_COMMIT} " + "VLLM_PRECOMPILED_WHEEL_VARIANT=cu130 " + "uv pip install --system --editable . " + "--extra-index-url https://wheels.vllm.ai/nightly/cu130", + "test -f /vllm/vllm/_C.abi3.so", + ) + .run_commands( + f"git clone {FOUNDRY_REPO} /foundry", + f"cd /foundry && git checkout {FOUNDRY_REF}", + "cd /foundry && pip install -e . --no-build-isolation", + ) + .env( + { + "HF_HOME": HF_CACHE, + "HUGGINGFACE_HUB_CACHE": f"{HF_CACHE}/hub", + "PYTHONUNBUFFERED": "1", + } + ) +) + +# Modal auto-mounts only the invoked entrypoint file, not its sibling modules. +# The phase functions import these helpers at module scope, so mount them +# explicitly onto the container PYTHONPATH (runtime mount; copy=False) or +# hydration fails with ModuleNotFoundError before any phase runs. +image = image.add_local_python_source( + "vllm_tp_validation", + "tp_validation_helpers", +) +cleanup_image = ( + modal.Image.debian_slim() + .env( + { + "FOUNDRY_REF": FOUNDRY_REF, + "VLLM_REF": VLLM_REF, + "VLLM_WHEEL_COMMIT": VLLM_WHEEL_COMMIT, + "VLLM_MODEL_REVISION": MODEL_REVISION, + "TP_SIZE": str(TP_SIZE), + "MODAL_GPU": MODAL_GPU, + } + ) + .add_local_python_source( + "vllm_tp_validation", + "tp_validation_helpers", + ) +) + +app = modal.App("foundry-vllm-tp-validation") +archive_volume = modal.Volume.from_name( + os.environ.get("VLLM_TP_ARCHIVE_VOLUME", "foundry-vllm-tp-validation-archive"), + create_if_missing=True, +) +hf_volume = modal.Volume.from_name( + os.environ.get("VLLM_TP_HF_VOLUME", "foundry-vllm-tp-hf-cache"), + create_if_missing=True, +) +# The publication authority. A run's archive is "published" only when this Dict +# holds its entry; the mounted volume merely stores bytes. Kept separate from the +# volume precisely because Modal auto-commit cannot make an on-volume marker the +# reliably-last durable write. +publication_index = modal.Dict.from_name( + os.environ.get("VLLM_TP_PUBLICATION_INDEX", "foundry-vllm-tp-validation-publication-index"), + create_if_missing=True, +) + + +class PhaseFailure(RuntimeError): + """A server phase failed; carries the collected result and every failure. + + All secondary (scan / evidence / immutability) failures are appended after + the primary phase/teardown failure so the first fault is never masked and the + full diagnostics survive in ``failures`` and ``result``. + """ + + def __init__(self, phase: str, result: dict, failures: list[str]) -> None: + super().__init__(f"vLLM phase {phase} failed: {'; '.join(failures)}") + self.phase = phase + self.result = result + self.failures = list(failures) + + +class PublicationReconciliationError(RuntimeError): + """A failed remote save transaction could not be reconciled to a marker. + + Raised by the local entrypoint's reconciliation when the publication record is + absent or fails full schema/consistency/digest validation. It is a plain + ``RuntimeError`` subclass (Python 3.10 compatible — no ``ExceptionGroup``); the + original remote error is always preserved as ``__cause__`` via ``raise ... from``. + """ + + +def _txn_local_root(run_id: str) -> Path: + return Path(LOCAL_ROOT) / run_id / "txn" + + +def _baseline_local_root(run_id: str) -> Path: + return Path(LOCAL_ROOT) / run_id / "baseline" + + +def _load_local_root(run_id: str, phase: str) -> Path: + return Path(LOCAL_ROOT) / run_id / phase + + +def _published_dir(run_id: str) -> Path: + return Path(PUBLISHED_ROOT) / run_id + + +def _new_run_id() -> str: + """A unique run id: ``--``. + + The cryptographically random suffix (not the timestamp) is what makes a + canonical-path / index-key collision negligible; the timestamp is retained + only for human-readable ordering. + """ + safe_ref = re.sub(r"[^A-Za-z0-9_.-]", "-", FOUNDRY_REF)[:24] + return f"{safe_ref}-{time.time_ns()}-{secrets.token_hex(8)}" + + +def _recipe_command(mode: str | None) -> list[str]: + command = ["bash", RECIPE_SCRIPT, str(TP_SIZE)] + if mode is not None: + command.append(f"--{mode}") + return command + + +def _log_tail(log_path: str, lines: int = 100) -> str: + path = Path(log_path) + if not path.exists(): + return "" + return "\n".join(path.read_text(errors="replace").splitlines()[-lines:]) + + +def _wait_until_healthy( + process: subprocess.Popen, + log_path: str, + timeout: float, +) -> None: + url = f"http://127.0.0.1:{PORT}/health" + deadline = time.time() + timeout + last_error = "" + + while time.time() < deadline: + if process.poll() is not None: + raise RuntimeError( + f"vLLM exited with {process.returncode} before becoming healthy\n" + f"{_log_tail(log_path)}" + ) + try: + with urllib.request.urlopen(url, timeout=5) as response: + if response.status == 200: + return + except Exception as error: # noqa: BLE001 - preserve the last health error + last_error = str(error) + time.sleep(3) + + raise RuntimeError( + f"vLLM was not healthy after {timeout}s: {last_error}\n{_log_tail(log_path)}" + ) + + +def _generate(prompt: str) -> dict[str, str | list[str]]: + request = urllib.request.Request( + f"http://127.0.0.1:{PORT}/v1/completions", + data=json.dumps( + { + "model": MODEL, + "prompt": prompt, + "max_tokens": MAX_TOKENS, + "temperature": 0.0, + "seed": 0, + "stop": ["\n"], + "logprobs": 1, + } + ).encode(), + headers={"Content-Type": "application/json"}, + ) + with urllib.request.urlopen(request, timeout=180) as response: + result = json.loads(response.read()) + return completion_evidence(result) + + +def _generate_concurrently() -> list[dict[str, str | list[str]]]: + barrier = threading.Barrier(len(PROMPTS)) + + def generate_after_barrier(prompt: str) -> dict[str, str | list[str]]: + barrier.wait(timeout=CONCURRENT_BARRIER_TIMEOUT_S) + return _generate(prompt) + + with ThreadPoolExecutor(max_workers=len(PROMPTS)) as executor: + return list(executor.map(generate_after_barrier, PROMPTS)) + + +def _launch( + mode: str | None, + log_path: str, + run_root: Path, +) -> tuple[subprocess.Popen, IO[str]]: + env = dict(os.environ) + env["CUDA_VISIBLE_DEVICES"] = ",".join(str(rank) for rank in range(TP_SIZE)) + env["VLLM_MODEL"] = MODEL + env["VLLM_MODEL_REVISION"] = MODEL_REVISION + env["VLLM_HOST"] = "127.0.0.1" + env["VLLM_PORT"] = str(PORT) + env["NCCL_DEBUG"] = "INFO" + + # The handle intentionally spans the child lifetime and is closed by stop_process. + log_file = open(log_path, "w") # noqa: SIM115 + try: + process = subprocess.Popen( + _recipe_command(mode), + stdout=log_file, + stderr=subprocess.STDOUT, + env=env, + cwd=run_root, + start_new_session=True, + ) + except BaseException: + log_file.close() + raise + return process, log_file + + +def _read_rank_offsets(workspace: Path) -> dict[str, int]: + offsets = {} + for rank in range(TP_SIZE): + path = workspace / f"rank_{rank}" / "final_alloc_offset.json" + if not path.exists(): + continue + offsets[str(rank)] = int(json.loads(path.read_text())["final_alloc_offset"]) + return offsets + + +def _archive_graph_counts(workspace: Path) -> dict[str, int]: + counts = {} + for rank in range(TP_SIZE): + rank_dir = workspace / f"rank_{rank}" + counts[str(rank)] = len(list(rank_dir.glob("graph_*.cugraph"))) + return counts + + +def _archive_fingerprints(workspace: Path) -> dict[str, dict[str, str]]: + return archive_fingerprints( + workspace, + TP_SIZE, + FOUNDRY_BASE_ADDR, + FOUNDRY_REGION_END, + ) + + +def _archive_byte_fingerprints(workspace: Path) -> dict[str, str]: + return archive_byte_fingerprints(workspace) + + +def _expected_graph_stem(index: int) -> str: + """The canonical descriptor stem for graph ``index`` in the current profile. + + Mirrors :func:`foundry.integration.vllm.graph_ops._build_graph_filename` for + this exact single-node FULL-cudagraph TP profile. The recipe requests capture + sizes ``seq 1 64`` and vLLM captures largest-first, so ``state.capture_index`` + increments in *descending* batch-size order: capture_index ``index`` records + batch size ``EXPECTED_GRAPH_COUNT - index`` (tokens == reqs), uniform (``U``), + no LoRA (``X``), no piecewise (``pcN``). Index 0 is therefore ``t64_r64`` and + index 63 is ``t1_r1``. SAVE / SAVE2 must match this exact descriptor set for + every index — an extra or malformed-but-paired descriptor is rejected even + when the count is exactly 64. + """ + size = EXPECTED_GRAPH_COUNT - index + return f"graph_{index}_FULL_t{size}_r{size}_UX_pcN" + + +def _archive_snapshot(workspace: Path) -> dict: + return collect_archive_snapshot( + workspace, + TP_SIZE, + EXPECTED_GRAPH_COUNT, + FOUNDRY_BASE_ADDR, + FOUNDRY_REGION_END, + expected_stems=_expected_graph_stem, + ) + + +_GRAPH_TOKEN_RE = re.compile(r"^graph_\d+_[A-Z_]+_t(\d+)_") + + +def _graph_file_diagnostics(workspace: Path) -> dict[str, dict[str, object]]: + """Bounded per-rank graph filename/count diagnostics. + + Records the ``.json`` / ``.cugraph`` counts, a capped sample of graph + filenames, and any duplicate batch sizes (``duplicate_sizes`` maps a repeated + ``t`` to its occurrence count) per rank. Duplicate sizes expose seed's + profiling captures (e.g. sizes 64/63 recorded twice) without hiding them. + Used both for seed's non-canonical (diagnostic-only) inventory and to enrich a + SAVE archive-evidence :class:`PhaseFailure` — never validated, published, or + reduced as canonical evidence. + """ + diagnostics: dict[str, dict[str, object]] = {} + for rank in range(TP_SIZE): + rank_dir = workspace / f"rank_{rank}" + names = ( + sorted( + path.name for path in rank_dir.glob("graph_*") if path.name != "graph_manifest.json" + ) + if rank_dir.is_dir() + else [] + ) + json_names = [name for name in names if name.endswith(".json")] + cugraph_names = [name for name in names if name.endswith(".cugraph")] + size_counts: dict[str, int] = {} + for name in json_names: + match = _GRAPH_TOKEN_RE.match(name) + if match is not None: + size = match.group(1) + size_counts[size] = size_counts.get(size, 0) + 1 + diagnostics[str(rank)] = { + "json_count": len(json_names), + "cugraph_count": len(cugraph_names), + "sample_files": names[:GRAPH_DIAGNOSTIC_SAMPLE_LIMIT], + "duplicate_sizes": {size: count for size, count in size_counts.items() if count > 1}, + } + return diagnostics + + +def _scan_log(log_path: str) -> dict: + text = Path(log_path).read_text(errors="replace") + scan = scan_log_text(text) + transports: dict[str, set[str]] = {} + for match in NCCL_TRANSPORT_PATTERN.finditer(text): + transports.setdefault(match.group("rank"), set()).add(match.group("transport")) + scan.update( + { + "early_build_graph_counts": [ + int(match.group("count")) for match in EARLY_BUILDS_PATTERN.finditer(text) + ], + "nccl_transports": { + rank: sorted(rank_transports) for rank, rank_transports in transports.items() + }, + "captured_offset_count": text.count("[foundry] Captured final_alloc_offset:"), + "saved_manifest_count": text.count("[foundry] Saved graph_manifest.json:"), + } + ) + return scan + + +def _empty_scan() -> dict: + return { + "errors": [], + "load_restoration_completions": {}, + "load_final_alloc_offsets": {}, + "nccl_version": None, + "early_build_graph_counts": [], + "nccl_transports": {}, + "captured_offset_count": 0, + "saved_manifest_count": 0, + } + + +def _package_version(*distributions: str) -> str | None: + for distribution in distributions: + try: + return importlib.metadata.version(distribution) + except importlib.metadata.PackageNotFoundError: + continue + return None + + +def _nvidia_smi(query: str) -> list[str]: + """Read all GPU values through nvidia-smi without initializing CUDA.""" + try: + output = subprocess.check_output( + ["nvidia-smi", f"--query-gpu={query}", "--format=csv,noheader"], + text=True, + stderr=subprocess.DEVNULL, + ) + except (OSError, subprocess.CalledProcessError): + return [] + return [line.strip() for line in output.splitlines() if line.strip()] + + +def _runtime_provenance(nccl_version: str | None) -> dict[str, str | int | list[str] | None]: + """Actual package, driver, GPU, and child-parsed NCCL provenance. + + Distribution metadata and ``nvidia-smi`` do not import torch/vLLM or create a + CUDA context in the orchestration process. NCCL comes only from child logs. + """ + drivers = _nvidia_smi("driver_version") + gpu_names = _nvidia_smi("name") + return { + "actual_torch_version": _package_version("torch"), + "actual_vllm_version": _package_version("vllm"), + "driver_version": drivers[0] if drivers else None, + "gpu_count": len(gpu_names), + "gpu_names": gpu_names, + "nccl_runtime_version": nccl_version, + } + + +def _candidate_metadata() -> dict[str, str | int]: + return { + "foundry_ref": FOUNDRY_REF, + "vllm_ref": VLLM_REF, + "requested_vllm_ref": VLLM_REF, + "vllm_wheel_commit": VLLM_WHEEL_COMMIT, + "model": MODEL, + "model_revision": MODEL_REVISION, + "cuda_image_tag": CUDA_TAG, + "requested_torch_version": REQUESTED_TORCH_VERSION, + "requested_gpu": MODAL_GPU, + "tp_size": TP_SIZE, + "cudagraph_mode": CUDAGRAPH_MODE, + "distributed_executor_backend": DISTRIBUTED_EXECUTOR_BACKEND, + "attention_backend": ATTENTION_BACKEND, + "collective_backend": COLLECTIVE_BACKEND, + } + + +def _provenance_failures(result: object) -> list[str]: + """Require exact requested candidate pins and complete actual runtime evidence.""" + if not isinstance(result, dict): + return ["phase result is not an object"] + failures: list[str] = [] + for field, expected in _candidate_metadata().items(): + if result.get(field) != expected: + failures.append(f"{field} {result.get(field)!r} != requested {expected!r}") + + for field in ("phase_time_s", "time_to_healthy_s"): + value = result.get(field) + if not isinstance(value, (int, float)) or isinstance(value, bool) or value <= 0: + failures.append(f"{field} must be positive, got {value!r}") + + runtime = result.get("runtime_provenance") + if not isinstance(runtime, dict): + return [*failures, "runtime_provenance is not an object"] + for field in ( + "actual_torch_version", + "actual_vllm_version", + "driver_version", + "nccl_runtime_version", + ): + value = runtime.get(field) + if not isinstance(value, str) or not value.strip(): + failures.append(f"runtime_provenance.{field} must be nonempty") + gpu_names = runtime.get("gpu_names") + if ( + runtime.get("gpu_count") != TP_SIZE + or not isinstance(gpu_names, list) + or len(gpu_names) != TP_SIZE + or not all(isinstance(name, str) and name.strip() for name in gpu_names) + ): + failures.append(f"runtime GPU inventory must contain exactly {TP_SIZE} nonempty names") + return failures + + +def _completion_set_failures(label: str, records: object) -> list[str]: + failures: list[str] = [] + if not isinstance(records, list): + return [f"{label} completion evidence is not a list"] + if len(records) != len(PROMPTS): + failures.append(f"{label} record count {len(records)} != exact prompt count {len(PROMPTS)}") + for index, record in enumerate(records): + if not isinstance(record, dict): + failures.append(f"{label} record {index} is not an object") + continue + text = record.get("text") + token_strings = record.get("token_strings") + if not isinstance(text, str) or not text: + failures.append(f"{label} record {index} has empty or invalid text") + if ( + not isinstance(token_strings, list) + or not token_strings + or not all(isinstance(token, str) and token for token in token_strings) + ): + failures.append(f"{label} record {index} has empty or invalid token_strings") + return failures + + +def _completion_matrix_failures(results: dict) -> list[str]: + """Require the exact 4x2 output matrix and byte-for-byte baseline equality.""" + failures: list[str] = [] + mode_fields = (("sequential", "outputs"), ("concurrent", "concurrent_outputs")) + for phase in QUERY_PHASES: + phase_result = results.get(phase) + if not isinstance(phase_result, dict): + failures.append(f"{phase} query result is unavailable") + continue + for label, field in mode_fields: + failures.extend(_completion_set_failures(f"{phase} {label}", phase_result.get(field))) + + baseline = results.get("baseline", {}) + for label, field in mode_fields: + expected = baseline.get(field) if isinstance(baseline, dict) else None + for phase in QUERY_PHASES[1:]: + actual = results.get(phase, {}).get(field) + if actual != expected: + failures.append(f"{phase} {label} completions != baseline") + return failures + + +def _run_server_phase( + phase: str, + run_root: Path, + log_path: str, + mode: str | None, + query: bool, +) -> tuple[dict, list[str]]: + """Launch one server, drive it, tear it down, and scan — all fail-closed. + + Returns ``(result, failures)``. ``failures`` is ordered so the primary phase + execution failure comes first, then teardown, then a log-scan failure, then + fatal log signatures. A scan exception is appended without masking an earlier + phase/teardown failure and never raised, so the caller keeps full diagnostics. + This function only runs a server; it never touches the mounted archive volume. + """ + run_root.mkdir(parents=True, exist_ok=True) + result = { + "phase": phase, + "artifact_dir": str(run_root), + **_candidate_metadata(), + } + started_at = time.time() + process, log_file = _launch(mode, log_path, run_root) + + phase_error: Exception | None = None + try: + _wait_until_healthy(process, log_path, timeout=1800) + result["time_to_healthy_s"] = max(round(time.time() - started_at, 3), 0.001) + if query: + result["outputs"] = [_generate(prompt) for prompt in PROMPTS] + result["concurrent_outputs"] = _generate_concurrently() + if process.poll() is not None: + raise RuntimeError( + f"vLLM exited unexpectedly after validation requests\n{_log_tail(log_path)}" + ) + except Exception as error: # noqa: BLE001 - scan complete logs before reporting + phase_error = error + + teardown_error: Exception | None = None + try: + result["process_return_code"] = stop_process(process, log_file) + except Exception as error: # noqa: BLE001 - include scan evidence with teardown failure + teardown_error = error + result["process_return_code"] = getattr(process, "returncode", None) + + scan_error: Exception | None = None + try: + result.update(_scan_log(log_path)) + except Exception as error: # noqa: BLE001 - a malformed log must not mask a real failure + scan_error = error + result.update(_empty_scan()) + provenance_error: Exception | None = None + try: + result["runtime_provenance"] = _runtime_provenance(result.get("nccl_version")) + except Exception as error: # noqa: BLE001 - append after primary/teardown/scan errors + provenance_error = error + result["runtime_provenance"] = {} + result["phase_time_s"] = max(round(time.time() - started_at, 3), 0.001) + + failures: list[str] = [] + if phase_error is not None: + failures.append(f"phase execution: {phase_error}") + if teardown_error is not None: + failures.append(f"process teardown: {teardown_error}") + if scan_error is not None: + failures.append(f"log scan: {scan_error}") + if provenance_error is not None: + failures.append(f"runtime provenance: {provenance_error}") + provenance_failures = _provenance_failures(result) + if provenance_failures: + failures.append(f"runtime provenance: {provenance_failures}") + if result["errors"]: + failures.append(f"fatal log signatures: {result['errors']}") + if query: + failures.extend(_completion_set_failures(f"{phase} sequential", result.get("outputs"))) + failures.extend( + _completion_set_failures( + f"{phase} concurrent", + result.get("concurrent_outputs"), + ) + ) + return result, failures + + +def _collect_save_evidence(workspace: Path, result: dict) -> list[str]: + """Collect canonical per-phase SAVE archive evidence, appending errors. + + Enforces the exact canonical descriptor profile (via :func:`_archive_snapshot`), + so an extra/malformed/wrong-count descriptor is a failure. On failure a + bounded graph filename/count diagnostic is appended after the underlying error + so the resulting :class:`PhaseFailure` carries actionable detail. + """ + failures: list[str] = [] + try: + snapshot = _archive_snapshot(workspace) + result["archive_inventory"] = snapshot["inventory"] + result["rank_offsets"] = snapshot["offsets"] + result["archive_graph_counts"] = snapshot["graph_counts"] + result["archive_fingerprints"] = snapshot["semantic_fingerprints"] + result["archive_byte_fingerprints"] = snapshot["byte_fingerprints"] + except Exception as error: # noqa: BLE001 - malformed archive is a secondary failure + failures.append(f"archive evidence: {error}") + try: + failures.append(f"archive graph diagnostics: {_graph_file_diagnostics(workspace)}") + except Exception as diag_error: # noqa: BLE001 - diagnostics are best-effort + failures.append(f"archive graph diagnostics unavailable: {diag_error}") + return failures + + +def _collect_seed_diagnostics(workspace: Path, result: dict) -> None: + """Record seed's diagnostic (non-canonical) graph inventory without validation. + + Pinned seed profiling legitimately captures duplicate batch sizes (e.g. 64/63 + before the normal 64, yielding 66 pairs). The archive seed writes is + diagnostic only — the next SAVE overwrites it — so it must NOT be held to the + exact-64 canonical descriptor contract. We record bounded per-rank graph + counts and a file sample, but never validate, publish, or reduce them as + canonical SAVE evidence. + """ + result["seed_diagnostics"] = _graph_file_diagnostics(workspace) + + +def _run_transaction_phase(phase: str, txn_root: Path, workspace: Path) -> dict: + """Run one seed/save/save2 phase locally and collect its archive evidence. + + ``seed`` records diagnostic-only graph counts (never the exact-64 validator); + ``save`` / ``save2`` collect the canonical, exactly-validated SAVE evidence. + """ + log_path = str(txn_root / f"{phase}.log") + # The transaction's first phase clears the local archive; ``save`` / ``save2`` + # overwrite it in place. Nothing here is on the mounted volume. + if phase == "seed": + shutil.rmtree(workspace, ignore_errors=True) + result, failures = _run_server_phase( + phase, + txn_root, + log_path, + PHASE_MODES[phase], + phase in QUERY_PHASES, + ) + if phase == "seed": + _collect_seed_diagnostics(workspace, result) + else: + failures.extend(_collect_save_evidence(workspace, result)) + if failures: + raise PhaseFailure(phase, result, failures) + return result + + +def _validate_transaction(results: dict, workspace: Path, baseline_evidence: dict) -> dict: + """Validate the whole transaction and return the immutable SAVE2 snapshot. + + Enforces, before any publication: + + * exact SAVE2 inventory (paired stems/indices/required files) and exactly + :data:`EXPECTED_GRAPH_COUNT` graphs per rank; + * SAVE vs SAVE2 equality of the complete validated inventory, per-rank + offsets, cugraph counts, and semantic fingerprints (raw ``.cugraph`` bytes + may legitimately differ across runs and are deliberately NOT compared); + * every per-rank offset strictly positive and, for this exact single-node + TP=2 profile, observed rank symmetry (all ranks share one offset); + * seed/save/save2 NCCL transport equal to the baseline transport map and no + fatal log signatures; + * SAVE2 sequential + concurrent outputs equal to the passed baseline. + + Raises on any violation so a bad transaction never publishes. + """ + expected_ranks = {str(rank) for rank in range(TP_SIZE)} + expected_counts = {rank: EXPECTED_GRAPH_COUNT for rank in expected_ranks} + save = results["save"] + save2 = results["save2"] + problems: list[str] = [] + + # The exact validated SAVE2 inventory becomes the immutable snapshot. + snapshot = _archive_snapshot(workspace) + if snapshot["graph_counts"] != expected_counts: + problems.append(f"SAVE2 graph counts {snapshot['graph_counts']} != {expected_counts}") + + # SAVE vs SAVE2 must reproduce the complete validated inventory, counts, + # offsets, and semantic fingerprints. + if save.get("archive_inventory") != save2.get("archive_inventory"): + problems.append("SAVE/SAVE2 validated inventory differs") + if save.get("rank_offsets") != save2.get("rank_offsets"): + problems.append( + f"SAVE/SAVE2 offsets differ: {save.get('rank_offsets')} vs {save2.get('rank_offsets')}" + ) + if save.get("archive_graph_counts") != save2.get("archive_graph_counts"): + problems.append("SAVE/SAVE2 graph counts differ") + if save2.get("archive_graph_counts") != expected_counts: + problems.append( + f"SAVE2 graph counts {save2.get('archive_graph_counts')} != {expected_counts}" + ) + if save.get("archive_fingerprints") != save2.get("archive_fingerprints"): + problems.append("SAVE/SAVE2 semantic fingerprints differ") + + # Offsets must be present, non-zero, and (for this profile) rank-symmetric. + save2_offsets = save2.get("rank_offsets") or {} + if set(save2_offsets) != expected_ranks: + problems.append(f"SAVE2 offsets rank set {set(save2_offsets)} != {expected_ranks}") + elif not all(offset > 0 for offset in save2_offsets.values()): + problems.append(f"SAVE2 offsets must all be non-zero: {save2_offsets}") + elif len(set(save2_offsets.values())) != 1: + problems.append(f"SAVE2 offsets are not rank-symmetric for this profile: {save2_offsets}") + + # The baseline transport map must itself be the exact, non-empty P2P/IPC + # profile for every rank (an empty or NET baseline is rejected), and every + # transaction phase must equal that same strict profile. + expected_transports = _expected_nccl_transports() + baseline_transports = baseline_evidence.get("nccl_transports") + if baseline_transports != expected_transports: + problems.append( + f"baseline NCCL transports {baseline_transports} " + f"!= expected P2P/IPC profile {expected_transports}" + ) + for phase in SAVE_TRANSACTION_PHASES: + phase_result = results[phase] + if phase_result.get("nccl_transports") != expected_transports: + problems.append( + f"{phase} NCCL transports {phase_result.get('nccl_transports')} " + f"!= expected P2P/IPC profile {expected_transports}" + ) + if phase_result.get("errors"): + problems.append(f"{phase} has fatal log signatures: {phase_result['errors']}") + + baseline_outputs = baseline_evidence.get("outputs") or [] + baseline_concurrent = baseline_evidence.get("concurrent_outputs") or [] + problems.extend(_completion_set_failures("baseline sequential", baseline_outputs)) + problems.extend(_completion_set_failures("baseline concurrent", baseline_concurrent)) + problems.extend(_completion_set_failures("save2 sequential", save2.get("outputs"))) + problems.extend(_completion_set_failures("save2 concurrent", save2.get("concurrent_outputs"))) + if not baseline_outputs or save2.get("outputs") != baseline_outputs: + problems.append("SAVE2 sequential outputs != baseline") + if not baseline_concurrent or save2.get("concurrent_outputs") != baseline_concurrent: + problems.append("SAVE2 concurrent outputs != baseline") + + if problems: + raise RuntimeError(f"transaction validation failed: {'; '.join(problems)}") + return snapshot + + +def _snapshot_digest(snapshot: dict) -> str: + """A deterministic digest of the validated snapshot from its canonical JSON.""" + canonical = json.dumps(snapshot, sort_keys=True, separators=(",", ":")) + return hashlib.sha256(canonical.encode("utf-8")).hexdigest() + + +def _index_entry(snapshot: dict, published_path: str) -> dict: + """The publication identity for a run: digest, path, and Foundry ref.""" + return { + "snapshot_digest": _snapshot_digest(snapshot), + "published_path": published_path, + "foundry_ref": FOUNDRY_REF, + } + + +def _reduced_phase_result(result: dict) -> dict: + """A bounded, primitive projection of a phase result for the stored payload.""" + return {key: result[key] for key in _STORED_PHASE_FIELDS if key in result} + + +def _publication_record(run_id: str, snapshot: dict, published_path: str, payload: dict) -> dict: + """The full authority record: run id, identity, and the prevalidated payload. + + ``run_id`` and ``published_path`` are the explicit fields Task 7 cleanup uses + to delete both the Dict entry and the on-volume archive. + """ + return { + "run_id": run_id, + **_index_entry(snapshot, published_path), + "payload": payload, + } + + +def _idempotent_record_matches(existing: object, candidate: dict, run_id: str) -> bool: + """Compare deterministic publication identity while ignoring measured timings.""" + if not _record_is_consistent(existing, run_id): + return False + if not _record_is_consistent(candidate, run_id): + return False + + def without_timings(record: dict) -> dict: + normalized = copy.deepcopy(record) + for phase_result in normalized["payload"]["results"].values(): + phase_result.pop("phase_time_s", None) + phase_result.pop("time_to_healthy_s", None) + return normalized + + return without_timings(existing) == without_timings(candidate) + + +def _record_serialized_size(record: dict) -> int: + """The larger of the record's JSON and pickle serialized sizes (bytes).""" + json_size = len(json.dumps(record, sort_keys=True).encode("utf-8")) + pickle_size = len(pickle.dumps(record, protocol=pickle.HIGHEST_PROTOCOL)) + return max(json_size, pickle_size) + + +def _require_record_within_limit(record: dict) -> None: + """Preflight the record's serialized size so a marker write cannot fail on size.""" + size = _record_serialized_size(record) + if size > PUBLICATION_RECORD_MAX_BYTES: + raise RuntimeError( + f"publication record {size} bytes exceeds limit {PUBLICATION_RECORD_MAX_BYTES}" + ) + + +_REQUIRED_PAYLOAD_FIELDS = ( + "snapshot", + "snapshot_digest", + "published_run_id", + "published_path", + "results", +) + +# The reduced per-phase fields common to every transaction phase and the +# container/primitive type each must have. ``bool`` is explicitly rejected where +# an ``int`` is required because ``bool`` is an ``int`` subclass. +_COMMON_REDUCED_TYPES = { + "phase": str, + "foundry_ref": str, + "vllm_ref": str, + "requested_vllm_ref": str, + "vllm_wheel_commit": str, + "model": str, + "model_revision": str, + "cuda_image_tag": str, + "requested_torch_version": str, + "requested_gpu": str, + "tp_size": int, + "cudagraph_mode": str, + "distributed_executor_backend": str, + "attention_backend": str, + "collective_backend": str, + "process_return_code": int, + "nccl_transports": dict, + "errors": list, + "runtime_provenance": dict, + "time_to_healthy_s": float, + "phase_time_s": float, +} +# ``save`` / ``save2`` additionally carry the canonical, exactly-validated SAVE +# evidence main relies on; ``seed`` instead carries only its diagnostic +# (non-canonical) inventory. ``save2`` (the transaction's only QUERY phase) also +# carries the completion outputs. +_SAVE_REDUCED_EXTRA = { + "rank_offsets": dict, + "archive_graph_counts": dict, + "archive_fingerprints": dict, +} +_SEED_REDUCED_EXTRA = {"seed_diagnostics": dict} +_QUERY_REDUCED_EXTRA = {"outputs": list, "concurrent_outputs": list} + + +def _reduced_types_for_phase(phase: str) -> dict: + """The exact reduced field->type schema required for one transaction phase.""" + required = dict(_COMMON_REDUCED_TYPES) + required.update(_SEED_REDUCED_EXTRA if phase == "seed" else _SAVE_REDUCED_EXTRA) + if phase in QUERY_PHASES: + required.update(_QUERY_REDUCED_EXTRA) + return required + + +def _positive_int(value: object) -> bool: + return isinstance(value, int) and not isinstance(value, bool) and value > 0 + + +def _exact_int(value: object, expected: int) -> bool: + return isinstance(value, int) and not isinstance(value, bool) and value == expected + + +def _reduced_results_are_wellformed(results: object) -> bool: + """Whether ``results`` is exactly {seed,save,save2} with typed reduced fields.""" + if not isinstance(results, dict) or set(results) != set(SAVE_TRANSACTION_PHASES): + return False + for phase in SAVE_TRANSACTION_PHASES: + pr = results.get(phase) + if not isinstance(pr, dict): + return False + required = _reduced_types_for_phase(phase) + # Exact per-phase schema: exactly the permitted keys — the common required + # fields plus this phase's specific fields (seed's diagnostic vs save/save2 + # canonical, query extras only on save2) — and nothing else. This rejects + # seed carrying save fields, save/save2 carrying seed_diagnostics, stray + # query extras, and any unknown field. + if set(pr) != set(required): + return False + for field, typ in required.items(): + value = pr.get(field) + if not isinstance(value, typ): + return False + if typ is int and isinstance(value, bool): + return False + if _provenance_failures(pr): + return False + if phase in QUERY_PHASES: + if _completion_set_failures(f"{phase} sequential", pr.get("outputs")): + return False + if _completion_set_failures( + f"{phase} concurrent", + pr.get("concurrent_outputs"), + ): + return False + if pr.get("phase") != phase: + return False + return True + + +def _snapshot_is_wellformed(snapshot: object) -> bool: + """Whether ``snapshot`` has the full immutability schema for this TP profile.""" + if not isinstance(snapshot, dict) or not set(_SNAPSHOT_MATCH_KEYS).issubset(snapshot): + return False + inventory = snapshot.get("inventory") + counts = snapshot.get("graph_counts") + offsets = snapshot.get("offsets") + semantic = snapshot.get("semantic_fingerprints") + byte_map = snapshot.get("byte_fingerprints") + if not all(isinstance(part, dict) for part in (inventory, counts, offsets, semantic, byte_map)): + return False + expected_ranks = {str(rank) for rank in range(TP_SIZE)} + if any(set(part) != expected_ranks for part in (inventory, counts, offsets, semantic)): + return False + if not byte_map: # the flat byte map must cover the archived files + return False + return all( + _exact_int(counts[rank], EXPECTED_GRAPH_COUNT) + and _positive_int(offsets[rank]) + and isinstance(inventory[rank], dict) + and isinstance(semantic[rank], dict) + for rank in expected_ranks + ) + + +def _record_is_consistent(record: object, run_id: str) -> bool: + """Whether a fetched record is a self-consistent publication for ``run_id``. + + The canonical path is derived *only* from ``run_id`` (never trusted from the + record); the record and payload must both name that exact path. Every identity + field, the full required payload schema, the deep ``results`` schema (exact + ``{seed,save,save2}`` with typed reduced fields), the full snapshot schema + (rank keys, exact per-rank counts, positive offsets, inventory / semantic / + byte maps), and the snapshot digest recomputed from the payload snapshot must + all agree. This validator never raises on malformed nested data — it returns + ``False`` — so a caller wrapping it (reconciliation) never masks the original + error. Consumers and Task 7 cleanup share it, rejecting any foreign, tampered, + or malformed record. + """ + if not isinstance(record, dict): + return False + payload = record.get("payload") + if not isinstance(payload, dict): + return False + if any(key not in payload for key in _REQUIRED_PAYLOAD_FIELDS): + return False + snapshot = payload.get("snapshot") + if not _snapshot_is_wellformed(snapshot): + return False + if not _reduced_results_are_wellformed(payload.get("results")): + return False + canonical = str(_published_dir(run_id)) + try: + digest = _snapshot_digest(snapshot) + except (TypeError, ValueError): + return False + return ( + record.get("run_id") == run_id + and record.get("foundry_ref") == FOUNDRY_REF + and record.get("published_path") == canonical + and record.get("snapshot_digest") == digest + and payload.get("published_run_id") == run_id + and payload.get("published_path") == canonical + and payload.get("snapshot_digest") == digest + ) + + +def _archive_manifest(snapshot: dict, published_path: str) -> dict: + """The on-volume manifest: the identity plus the full snapshot.""" + return {**_index_entry(snapshot, published_path), "snapshot": snapshot} + + +def _fetch_record(run_id: str) -> object | None: + """Fetch the stored record for ``run_id`` or ``None`` if absent.""" + try: + return publication_index[run_id] + except KeyError: + return None + + +def _write_publication_marker(run_id: str, record: dict) -> None: + """Atomically write the authority marker, reconciling ambiguity idempotently. + + Uses ``put(..., skip_if_exists=True)``. A ``False`` return means the key + already exists: an identical stored record is idempotent success, anything + else is a collision. An exception from ``put`` is ambiguous (the write may or + may not have landed): fetch the record — an exact match means the marker is + present (success), otherwise re-raise preserving the original error (and any + reconciliation-fetch error). + """ + try: + inserted = publication_index.put(run_id, record, skip_if_exists=True) + except Exception as put_error: # noqa: BLE001 - reconciled below, never swallowed + try: + existing = _fetch_record(run_id) + except Exception as fetch_error: # noqa: BLE001 - both errors are preserved + raise RuntimeError( + f"publication marker for {run_id} failed and reconciliation could not " + f"read the record: {put_error}; {fetch_error}" + ) from put_error + if existing == record: + return + raise RuntimeError( + f"publication marker for {run_id} failed and did not reconcile: {put_error}" + ) from put_error + + if inserted: + return + existing = _fetch_record(run_id) + if existing == record: + return + raise RuntimeError( + f"publication marker collision for {run_id}: a different record already exists" + ) + + +def _reconcile_transaction(run_id: str, original_error: BaseException) -> dict: + """Recover a transaction whose marker landed but whose result delivery failed. + + The local entrypoint cannot inspect the mounted volume, so it validates only + the record's *full* schema/consistency/digest (:func:`_record_is_consistent`) + and returns the stored payload; the on-volume manifest/archive validation is + performed by the first LOAD (:func:`_require_indexed_published_run`) before the + archive is used. If the record is absent or fails validation the reconciliation + itself fails: raise :class:`PublicationReconciliationError` carrying the + reconciliation detail, with the original remote error preserved as ``__cause__`` + so a genuine marker-less transaction failure is never silently swallowed. + """ + try: + record = _fetch_record(run_id) + if record is None: + raise RuntimeError(f"no publication marker exists for {run_id}") + if not _record_is_consistent(record, run_id): + raise RuntimeError( + f"publication record for {run_id} failed full schema/consistency validation" + ) + return record["payload"] + except Exception as reconciliation_error: # noqa: BLE001 - wrapped, cause preserved + raise PublicationReconciliationError( + f"could not reconcile transaction {run_id}: {reconciliation_error}" + ) from original_error + + +_SNAPSHOT_MATCH_KEYS = ( + "inventory", + "graph_counts", + "offsets", + "semantic_fingerprints", + "byte_fingerprints", +) + + +def _require_snapshot_match(actual: dict, expected: dict, context: str) -> None: + for key in _SNAPSHOT_MATCH_KEYS: + if actual.get(key) != expected.get(key): + raise RuntimeError(f"{context} {key} != validated SAVE2 snapshot") + + +def _publish_archive(run_id: str, local_workspace: Path, expected_snapshot: dict) -> str: + """Copy the validated local SAVE2 archive to its unique canonical volume path. + + Does NOT commit and does NOT touch the publication index. The unique + ``mkdir(exist_ok=False)`` is defense-in-depth against an accidental reused run + id; logical exclusivity comes from the random run id plus the Dict authority. + A partial directory left by a failure here is inert (never indexed), so it is + best-effort removed only to avoid clutter, not for correctness. + """ + published = _published_dir(run_id) + Path(PUBLISHED_ROOT).mkdir(parents=True, exist_ok=True) + reject_symlinks(local_workspace) + # Never overwrite an existing canonical path (collision fails closed). + published.mkdir(parents=False, exist_ok=False) + + try: + incoming = published / INCOMING_ARCHIVE_DIRNAME + shutil.copytree(local_workspace, incoming) + reject_symlinks(incoming) + staged = _archive_snapshot(incoming) + _require_snapshot_match(staged, expected_snapshot, "staged archive") + os.rename(incoming, published / PUBLISHED_ARCHIVE_DIRNAME) + manifest = _archive_manifest(expected_snapshot, str(published)) + (published / MANIFEST_NAME).write_text(json.dumps(manifest, sort_keys=True)) + except BaseException: + # Inert either way; remove the partial dir best-effort for tidiness only. + with contextlib.suppress(BaseException): + shutil.rmtree(published, ignore_errors=True) + raise + return str(published) + + +def _commit_with_retries() -> None: + """Commit the archive volume with bounded retries; raise the last error.""" + last_error: BaseException | None = None + for _ in range(COMMIT_MAX_ATTEMPTS): + try: + archive_volume.commit() + return + except Exception as error: # noqa: BLE001 - retried; last error re-raised + last_error = error + raise last_error + + +def _load_validation_failures( + result: dict, + before: dict, + after: dict | None, + expected_snapshot: dict, +) -> list[str]: + """Failures a LOAD must not have. Appended after any phase/teardown failure.""" + problems: list[str] = [] + expected_ranks = {str(rank) for rank in range(TP_SIZE)} + + for key in ( + "inventory", + "graph_counts", + "offsets", + "semantic_fingerprints", + "byte_fingerprints", + ): + if before.get(key) != expected_snapshot.get(key): + problems.append(f"LOAD copy {key} != published SAVE2 snapshot") + + if after is None: + problems.append("LOAD archive reread unavailable") + else: + for key in ( + "inventory", + "graph_counts", + "offsets", + "semantic_fingerprints", + "byte_fingerprints", + ): + if after.get(key) != before.get(key): + problems.append(f"LOAD mutated local {key}") + if after.get(key) != expected_snapshot.get(key): + problems.append(f"LOAD local {key} != published SAVE2 snapshot") + + expected_completions = { + rank: [[EXPECTED_GRAPH_COUNT, EXPECTED_GRAPH_COUNT]] for rank in expected_ranks + } + if result.get("load_restoration_completions") != expected_completions: + problems.append( + f"LOAD restoration completions {result.get('load_restoration_completions')} " + f"!= one {EXPECTED_GRAPH_COUNT}/{EXPECTED_GRAPH_COUNT} per rank" + ) + if result.get("captured_offset_count", 0) != 0 or result.get("saved_manifest_count", 0) != 0: + problems.append("LOAD emitted capture/manifest-save markers") + + expected_live = { + rank: [offset] for rank, offset in expected_snapshot.get("offsets", {}).items() + } + if result.get("load_final_alloc_offsets") != expected_live: + problems.append( + f"LOAD live final_alloc_offsets {result.get('load_final_alloc_offsets')} " + f"!= one-per-rank SAVE2 offsets {expected_live}" + ) + if not result.get("outputs"): + problems.append("LOAD completion outputs unavailable") + return problems + + +def _require_indexed_published_run( + run_id: str, + published_run_dir: str, + expected_snapshot: dict, + *, + verify_bytes: bool = True, +) -> Path: + """Require the publication-index record, on-volume manifest, and archive child. + + Publication authority is the ``publication_index`` Dict. The canonical path is + derived from ``run_id`` and the passed ``published_run_dir`` must equal it (the + record's path is never trusted). A consumer copies a run only when a stored + record passes full schema/consistency validation (:func:`_record_is_consistent`, + shared with Task 7 cleanup) *and* its digest matches the expected snapshot. It + then double-checks the on-volume ``manifest.json`` and the ``archive`` child. A + background archive that was auto-committed but never indexed has no record and + is rejected (inert), as is any record/manifest whose schema, digest, path, or + ref differs. + + When ``verify_bytes`` (the default for consumers and for the idempotent replay + guard), a *fresh* snapshot is collected from the on-volume ``archive/`` bytes + with :func:`collect_archive_snapshot` and compared exactly (inventory, counts, + offsets, semantic and byte fingerprints) to ``expected_snapshot`` — so a + corrupted or drifted published archive is rejected even though its manifest and + directory still exist. Cleanup passes ``verify_bytes=False`` because it only + needs presence to delete. + """ + run_dir = Path(published_run_dir) + canonical = _published_dir(run_id) + if run_dir != canonical: + raise RuntimeError( + f"published dir {run_dir} != canonical {canonical} derived from run id {run_id}" + ) + + record = _fetch_record(run_id) + if record is None: + raise RuntimeError(f"run {run_id} is not in the publication index (unindexed/inert)") + if not _record_is_consistent(record, run_id): + raise RuntimeError( + f"publication record for {run_id} failed full schema/consistency validation" + ) + if record["snapshot_digest"] != _snapshot_digest(expected_snapshot): + raise RuntimeError(f"publication record for {run_id} != expected snapshot digest") + + manifest_path = run_dir / MANIFEST_NAME + if not manifest_path.is_file(): + raise RuntimeError(f"published run {run_dir} has no {MANIFEST_NAME}") + manifest = json.loads(manifest_path.read_text()) + if manifest != _archive_manifest(expected_snapshot, str(run_dir)): + raise RuntimeError(f"published run {run_dir} manifest != expected snapshot/path/ref") + + archive = run_dir / PUBLISHED_ARCHIVE_DIRNAME + if not archive.is_dir(): + raise RuntimeError(f"published run {run_dir} has no {PUBLISHED_ARCHIVE_DIRNAME} child") + if verify_bytes: + # Re-derive the snapshot from the actual on-volume bytes and require an + # exact match: the manifest/dir alone cannot catch byte-level drift. + actual = _archive_snapshot(archive) + _require_snapshot_match(actual, expected_snapshot, f"published archive {run_id}") + return archive + + +def _run_load_local( + phase: str, + run_id: str, + published_run_dir: str, + expected_snapshot: dict, +) -> dict: + """Copy the published archive into independent local storage and LOAD there. + + Requires a matching publication-index record and on-volume manifest, copies + ``/archive`` into local storage, then snapshots the local copy before and + after execution and requires them to be identical, equal to the published + SAVE2 snapshot, with exactly one 64/64 restoration and one live offset per + rank, zero capture/save markers, and available outputs. Never writes the + mounted archive or the Dict; never commits. + """ + # Warm containers keep a stale view of the volume; reload before reading the + # Dict/manifest/bytes so a just-committed publication is visible. + archive_volume.reload() + source = _require_indexed_published_run(run_id, published_run_dir, expected_snapshot) + + run_root = _load_local_root(run_id, phase) + shutil.rmtree(run_root, ignore_errors=True) + run_root.mkdir(parents=True, exist_ok=True) + workspace = run_root / "foundry_archive" + + reject_symlinks(source) + shutil.copytree(source, workspace) + reject_symlinks(workspace) + + before = _archive_snapshot(workspace) + + log_path = str(run_root / f"{phase}.log") + result, failures = _run_server_phase( + phase, + run_root, + log_path, + PHASE_MODES[phase], + query=True, + ) + + after: dict | None = None + try: + after = _archive_snapshot(workspace) + result["archive_graph_counts"] = after["graph_counts"] + result["archive_fingerprints"] = after["semantic_fingerprints"] + result["archive_byte_fingerprints"] = after["byte_fingerprints"] + except Exception as error: # noqa: BLE001 - reread failure is a secondary failure + failures.append(f"archive reread: {error}") + + failures.extend(_load_validation_failures(result, before, after, expected_snapshot)) + if failures: + raise PhaseFailure(phase, result, failures) + return result + + +@app.function( + image=image, + gpu=MODAL_GPU, + timeout=3600, + volumes={ARCHIVE_ROOT: archive_volume, HF_CACHE: hf_volume}, +) +def run_baseline_phase(run_id: str) -> dict: + """Reference (non-Foundry) vLLM TP run, entirely in local storage.""" + run_root = _baseline_local_root(run_id) + log_path = str(run_root / "baseline.log") + result, failures = _run_server_phase("baseline", run_root, log_path, None, query=True) + if failures: + raise PhaseFailure("baseline", result, failures) + print(json.dumps(result, sort_keys=True)) + return result + + +@app.function( + image=image, + gpu=MODAL_GPU, + timeout=7200, + volumes={ARCHIVE_ROOT: archive_volume, HF_CACHE: hf_volume}, +) +def run_save_transaction(run_id: str, baseline_evidence: dict) -> dict: + """Run seed/save/save2 in local storage, validate, then publish atomically. + + All three phases run against ephemeral local storage. After validation the + run is published against the Dict authority: build the bounded record (identity + + run_id + reduced prevalidated payload) and preflight its serialized size; + short-circuit idempotently if an identical record already exists (and fail a + differing one as a collision); copy the validated archive + manifest to the + unique canonical volume path; commit the volume with bounded retries; then, as + the single final side effect, atomically ``put`` the marker (reconciling an + ambiguous put by fetching the record). The Dict is never written before a + successful commit, and the only work after a reconciled marker is returning + the already-serialized payload. A commit or unreconciled marker failure raises + and leaves the archive unindexed and inert (no rollback). + """ + txn_root = _txn_local_root(run_id) + workspace = txn_root / "foundry_archive" + results = { + phase: _run_transaction_phase(phase, txn_root, workspace) + for phase in SAVE_TRANSACTION_PHASES + } + snapshot = _validate_transaction(results, workspace, baseline_evidence) + + published_dir = _published_dir(run_id) + # Bounded, primitive payload: only the reduced per-phase results main needs + # plus the immutable snapshot. Serialize-check before any publication. + payload = { + "results": { + phase: _reduced_phase_result(results[phase]) for phase in SAVE_TRANSACTION_PHASES + }, + "snapshot": snapshot, + "snapshot_digest": _snapshot_digest(snapshot), + "published_run_id": run_id, + "published_path": str(published_dir), + } + json.dumps(payload, sort_keys=True) + record = _publication_record(run_id, snapshot, str(published_dir), payload) + _require_record_within_limit(record) + + # Idempotency / collision on the atomic marker before touching the volume. + existing = _fetch_record(run_id) + if existing is not None: + if _idempotent_record_matches(existing, record, run_id): + # Idempotent replay: confirm the prior publication is actually intact + # on the volume (reload to see another container's commit) before + # declaring success. + archive_volume.reload() + existing_payload = existing["payload"] + _require_indexed_published_run( + run_id, + str(published_dir), + existing_payload["snapshot"], + ) + return existing_payload + raise RuntimeError( + f"publication marker collision for {run_id}: a different record already exists" + ) + if published_dir.exists(): + raise FileExistsError(f"canonical published path already exists: {published_dir}") + + _publish_archive(run_id, workspace, snapshot) + # Durability first; only a successful commit may be followed by the marker. + _commit_with_retries() + # The single final side effect (atomic, self-reconciling). If it raises + # unreconciled, the committed archive stays unindexed and inert. + _write_publication_marker(run_id, record) + return payload + + +@app.function( + image=image, + gpu=MODAL_GPU, + timeout=3600, + volumes={ARCHIVE_ROOT: archive_volume, HF_CACHE: hf_volume}, +) +def run_load_phase( + phase: str, run_id: str, published_run_dir: str, expected_snapshot: dict +) -> dict: + """Independent fresh-process LOAD against the published SAVE2 archive. + + Accepts only ``load`` / ``load2``. Requires the published run's ready manifest + to match the snapshot, copies the canonical archive into its own local + storage, executes there, validates before/after immutability plus the + restoration/offset/marker contract, and never writes the mounted archive or + commits. + """ + if phase not in LOAD_PHASES: + raise ValueError(f"run_load_phase accepts only {LOAD_PHASES}, got {phase!r}") + result = _run_load_local(phase, run_id, published_run_dir, expected_snapshot) + print(json.dumps(result, sort_keys=True)) + return result + + +def delete_published_run(run_id: str) -> None: + """Task 7 cleanup contract: delete a published run's archive and Dict record. + + The canonical path is derived *exclusively* from ``run_id`` — the record's own + ``published_path`` is validated but never used as the deletion target, so a + tampered/foreign record can never redirect ``rmtree`` outside the canonical + tree. Steps, fail-closed: + + 1. ``reload`` the volume first, then derive the canonical path and fetch the + marker. + 2. Marker absent: + * canonical also absent → nothing to do; return idempotent success (covers + a prior pop that committed but whose ack was lost, and any repeated + cleanup of an already-deleted run); + * canonical present → an unindexed orphan; fail closed rather than delete an + unvalidated tree. + 3. Marker present but not self-consistent → refuse. + 4. Marker present, canonical absent → a prior deletion committed the byte + removal but its pop failed; pop the marker and succeed idempotently. + 5. Marker present, canonical present → validate the on-volume ``manifest.json`` + and ``archive`` child (fail closed if a partial deletion left the dir + without them; byte drift must not block deletion), strictly ``rmtree`` (no + ``ignore_errors``) and verify it is gone, commit the deletion with bounded + retries (on commit failure the marker remains, re-cleanable), then pop the + marker last (on pop failure the bytes are gone and committed, so consumers + reject via the missing archive and the next retry finishes at step 4). + + This strict local helper is called by ``cleanup_published_run``; the + orchestrator invokes that remote wrapper as Task 7 after validation. + """ + archive_volume.reload() + canonical = _published_dir(run_id) + record = _fetch_record(run_id) + + if record is None: + # No authority record. If the bytes are gone too, cleanup already completed + # (pop committed but ack lost, or a repeated cleanup) — succeed idempotently. + # If the bytes remain, this is an unindexed orphan; fail closed. + if canonical.exists(): + raise RuntimeError( + f"refusing to delete {run_id}: canonical archive exists but is unindexed (orphan)" + ) + return + + if not _record_is_consistent(record, run_id): + raise RuntimeError( + f"refusing to delete {run_id}: publication record is not self-consistent" + ) + + if not canonical.exists(): + # Marker present but the bytes were already removed+committed by a prior + # attempt whose pop failed: finish by popping the marker. + publication_index.pop(run_id, None) + return + + # The canonical dir exists: its manifest/archive must be present (fail closed + # if a partial deletion left the dir without them). Byte drift must not block + # deletion, so skip the fresh-snapshot comparison here. + _require_indexed_published_run( + run_id, str(canonical), record["payload"]["snapshot"], verify_bytes=False + ) + + shutil.rmtree(canonical) # strict: raises if already missing + if canonical.exists(): + raise RuntimeError(f"canonical archive still present after removal: {canonical}") + + _commit_with_retries() + publication_index.pop(run_id, None) + + +def _cleanup_state( + run_id: str, +) -> tuple[bool | None, bool | None, str | None, str | None]: + """Return marker/path presence, conservatively reporting unreadable state.""" + marker_present: bool | None = None + path_present: bool | None = None + try: + archive_volume.reload() + marker_present = _fetch_record(run_id) is not None + path_present = _published_dir(run_id).exists() + except Exception as error: # noqa: BLE001 - surfaced in the cleanup report + return ( + marker_present, + path_present, + str(error), + f"{type(error).__module__}.{type(error).__name__}", + ) + return marker_present, path_present, None, None + + +@app.function( + image=cleanup_image, + volumes={ARCHIVE_ROOT: archive_volume}, +) +def cleanup_published_run(run_id: str) -> dict: + """Run Task 7's strict cleanup and reconcile ambiguous final-write errors.""" + cleanup_error: BaseException | None = None + try: + delete_published_run(run_id) + except BaseException as error: # cleanup status must be reconciled before reporting + cleanup_error = error + + marker_present, path_present, state_error, state_error_type = _cleanup_state(run_id) + completed = marker_present is False and path_present is False + if completed: + return { + "cleanup_attempted": True, + "cleanup_completed": True, + "artifacts_retained": False, + "marker_present": False, + "published_path_present": False, + "cleanup_error": None, + "cleanup_error_type": None, + "reconciled_error": str(cleanup_error) if cleanup_error is not None else None, + } + + error_text = ( + str(cleanup_error) + if cleanup_error is not None + else state_error or "cleanup returned without removing both marker and canonical path" + ) + error_type = ( + f"{type(cleanup_error).__module__}.{type(cleanup_error).__name__}" + if cleanup_error is not None + else state_error_type + ) + return { + "cleanup_attempted": True, + "cleanup_completed": False, + # Unknown state is retained conservatively; a positive marker or path is + # explicit retained diagnostic evidence. + "artifacts_retained": marker_present is not False or path_present is not False, + "marker_present": marker_present, + "published_path_present": path_present, + "cleanup_error": error_text, + "cleanup_error_type": error_type, + "reconciled_error": None, + } + + +@app.function( + image=cleanup_image, + volumes={ARCHIVE_ROOT: archive_volume}, +) +def inspect_published_run(run_id: str) -> dict: + """Read-only artifact state for final failure reports.""" + archive_volume.reload() + marker_present = _fetch_record(run_id) is not None + published_path_present = _published_dir(run_id).exists() + return { + "marker_present": marker_present, + "published_path_present": published_path_present, + "artifacts_retained": marker_present or published_path_present, + } + + +def _obtain_transaction(run_id: str, baseline_evidence: dict) -> dict: + """Run the save transaction, reconciling a post-marker delivery failure. + + If ``run_save_transaction.remote`` raises — including a failure delivering or + deserializing the result *after* the marker was written — reconcile by + fetching the publication record: a fully schema/consistency-valid record means + the transaction actually succeeded, so recover its payload. Otherwise + :func:`_reconcile_transaction` raises :class:`PublicationReconciliationError` + with the original remote error as ``__cause__`` — a genuine marker-less failure + still fails, but the durable marker is always consulted first. + """ + try: + return run_save_transaction.remote(run_id, baseline_evidence) + except Exception as error: # noqa: BLE001 - reconciled against the durable marker + return _reconcile_transaction(run_id, error) + + +def _authoritative_transaction_payload(run_id: str, returned_payload: object) -> dict: + """Require main's transaction data to come from the durable Dict marker.""" + record = _fetch_record(run_id) + if record is None: + raise RuntimeError(f"transaction {run_id} returned without a publication marker") + if not _record_is_consistent(record, run_id): + raise RuntimeError( + f"publication record for {run_id} failed full schema/consistency validation" + ) + payload = record["payload"] + if returned_payload != payload: + raise RuntimeError(f"transaction {run_id} return payload != authoritative marker payload") + return payload + + +def _acceptance_checks(results: dict, transaction: dict) -> dict[str, bool]: + snapshot = transaction["snapshot"] + expected_ranks = {str(rank) for rank in range(TP_SIZE)} + expected_transports = _expected_nccl_transports() + expected_graph_counts = {rank: EXPECTED_GRAPH_COUNT for rank in expected_ranks} + expected_restoration_completions = { + rank: [[EXPECTED_GRAPH_COUNT, EXPECTED_GRAPH_COUNT]] for rank in expected_ranks + } + + save_offsets = results["save"].get("rank_offsets", {}) + save2_offsets = results["save2"].get("rank_offsets", {}) + save_fingerprints = results["save"].get("archive_fingerprints", {}) + save2_fingerprints = results["save2"].get("archive_fingerprints", {}) + snapshot_offsets = snapshot.get("offsets", {}) + snapshot_graph_counts = snapshot.get("graph_counts", {}) + snapshot_semantic = snapshot.get("semantic_fingerprints", {}) + snapshot_byte = snapshot.get("byte_fingerprints", {}) + expected_live_offsets = {rank: [snapshot_offsets[rank]] for rank in snapshot_offsets} + completion_failures = _completion_matrix_failures(results) + + return { + "exact_bounded_completion_evidence": not any( + "sequential" in problem for problem in completion_failures + ), + "exact_bounded_concurrent_completion_evidence": not any( + "concurrent" in problem for problem in completion_failures + ), + "complete_runtime_provenance_and_timing": all( + not _provenance_failures(results[phase]) for phase in PHASES + ), + "save_offsets_present": set(save_offsets) == expected_ranks and all(save_offsets.values()), + "save_rank_offsets_symmetric": len(set(save_offsets.values())) == 1, + "save_offsets_reproducible": bool(save_offsets) and save_offsets == save2_offsets, + "save_semantic_reproducible": bool(save_fingerprints) + and save_fingerprints == save2_fingerprints, + "archive_inventory_exact": snapshot_graph_counts == expected_graph_counts, + "published_run_available": bool(transaction.get("published_run_id")) + and bool(transaction.get("published_path")), + "graphs_restored_each_rank": all( + results[phase]["load_restoration_completions"] == expected_restoration_completions + for phase in LOAD_PHASES + ), + "load_live_offsets_match_save2": bool(expected_live_offsets) + and all( + results[phase].get("load_final_alloc_offsets", {}) == expected_live_offsets + for phase in LOAD_PHASES + ), + "load_did_not_capture": all( + results[phase]["captured_offset_count"] == 0 + and results[phase]["saved_manifest_count"] == 0 + for phase in LOAD_PHASES + ), + "load_archive_immutable": bool(snapshot_semantic) + and bool(snapshot_graph_counts) + and all( + results[phase].get("archive_fingerprints", {}) == snapshot_semantic + and results[phase].get("archive_graph_counts", {}) == snapshot_graph_counts + for phase in LOAD_PHASES + ), + "load_archive_bytes_immutable": bool(snapshot_byte) + and all( + results[phase].get("archive_byte_fingerprints", {}) == snapshot_byte + for phase in LOAD_PHASES + ), + "nccl_p2p_ipc_every_phase": all( + results[phase]["nccl_transports"] == expected_transports for phase in PHASES + ), + "no_runtime_errors": not any(results[phase]["errors"] for phase in PHASES), + "expected_process_shutdown": all( + results[phase]["process_return_code"] in {0, -2, 130} for phase in PHASES + ), + } + + +def _artifact_state_for_report(run_id: str) -> dict: + try: + state = inspect_published_run.remote(run_id) + except BaseException as error: + return { + "marker_present": None, + "published_path_present": None, + "artifacts_retained": True, + "artifact_state_error": str(error), + "artifact_state_error_type": (f"{type(error).__module__}.{type(error).__name__}"), + } + if ( + not isinstance(state, dict) + or not isinstance(state.get("marker_present"), bool) + or not isinstance(state.get("published_path_present"), bool) + or state.get("artifacts_retained") + is not (state.get("marker_present") or state.get("published_path_present")) + ): + return { + "marker_present": None, + "published_path_present": None, + "artifacts_retained": True, + "artifact_state_error": ("artifact inspection returned an incomplete result"), + "artifact_state_error_type": "artifact_state_result", + } + return { + "marker_present": state.get("marker_present"), + "published_path_present": state.get("published_path_present"), + "artifacts_retained": bool(state.get("artifacts_retained")), + "artifact_state_error": None, + "artifact_state_error_type": None, + } + + +def _report_evidence( + run_id: str, + run_url: str, + results: dict, + transaction: object, + checks: dict[str, bool], +) -> dict: + transaction_data = transaction if isinstance(transaction, dict) else {} + snapshot_value = transaction_data.get("snapshot", {}) + snapshot = snapshot_value if isinstance(snapshot_value, dict) else {} + + def result_value(result: object, field: str, default: object = None) -> object: + return result.get(field, default) if isinstance(result, dict) else default + + def nccl_version(result: object) -> object: + scanned = result_value(result, "nccl_version") + runtime = result_value(result, "runtime_provenance") + runtime_version = runtime.get("nccl_runtime_version") if isinstance(runtime, dict) else None + return scanned or runtime_version + + complete_query_results = all(isinstance(results.get(phase), dict) for phase in QUERY_PHASES) + return { + "run_id": run_id, + "run_url": run_url, + **_candidate_metadata(), + "published_run_id": transaction_data.get("published_run_id"), + "published_path": transaction_data.get("published_path"), + "process_return_codes": { + phase: result_value(result, "process_return_code") for phase, result in results.items() + }, + "phase_times_s": { + phase: result_value(result, "phase_time_s") for phase, result in results.items() + }, + "time_to_healthy_s": { + phase: result_value(result, "time_to_healthy_s") for phase, result in results.items() + }, + "runtime_provenance": { + phase: result_value(result, "runtime_provenance") for phase, result in results.items() + }, + "nccl_runtime_versions": {phase: nccl_version(result) for phase, result in results.items()}, + "completion_maps": { + phase: { + "sequential": result_value(results[phase], "outputs"), + "concurrent": result_value(results[phase], "concurrent_outputs"), + } + for phase in QUERY_PHASES + if phase in results + }, + "per_rank_offsets": { + "published_snapshot": snapshot.get("offsets", {}), + **{ + phase: result_value(result, "load_final_alloc_offsets") + or result_value(result, "rank_offsets") + for phase, result in results.items() + }, + }, + "restoration_completions": { + phase: result_value(results[phase], "load_restoration_completions", {}) + for phase in LOAD_PHASES + if phase in results + }, + "archive_inventory": snapshot.get("inventory", {}), + "archive_graph_counts": snapshot.get("graph_counts", {}), + "archive_semantic_fingerprints": snapshot.get("semantic_fingerprints", {}), + "archive_byte_fingerprints": snapshot.get("byte_fingerprints", {}), + "completion_failures": ( + _completion_matrix_failures(results) if complete_query_results else [] + ), + "checks": checks, + "results": results, + "snapshot": snapshot, + } + + +def _emit_report(report: dict) -> None: + print(json.dumps(report, indent=2, sort_keys=True)) + + +def _normalize_cleanup_result(outcome: object) -> dict: + if isinstance(outcome, dict): + complete = ( + outcome.get("cleanup_attempted") is True + and outcome.get("cleanup_completed") is True + and outcome.get("artifacts_retained") is False + and outcome.get("marker_present") is False + and outcome.get("published_path_present") is False + ) + if complete: + return outcome + marker = outcome.get("marker_present") + path = outcome.get("published_path_present") + error = outcome.get("cleanup_error") + error_type = outcome.get("cleanup_error_type") + else: + marker = None + path = None + error = None + error_type = None + return { + "cleanup_attempted": True, + "cleanup_completed": False, + "artifacts_retained": True, + "marker_present": marker if isinstance(marker, bool) else None, + "published_path_present": path if isinstance(path, bool) else None, + "cleanup_error": error or "cleanup returned an incomplete result", + "cleanup_error_type": error_type or "cleanup_result", + "reconciled_error": None, + } + + +@app.local_entrypoint() +def main() -> None: + run_id = _new_run_id() + run_url = app.get_dashboard_url() + results: dict[str, dict] = {} + transaction: dict | None = None + checks: dict[str, bool] = {} + stage = "baseline" + lifecycle_error: BaseException | None = None + try: + baseline = run_baseline_phase.remote(run_id) + results["baseline"] = baseline + baseline_evidence = { + "outputs": baseline.get("outputs", []), + "concurrent_outputs": baseline.get("concurrent_outputs", []), + "nccl_transports": baseline.get("nccl_transports", {}), + } + + stage = "transaction" + returned_transaction = _obtain_transaction(run_id, baseline_evidence) + transaction = _authoritative_transaction_payload(run_id, returned_transaction) + results.update(transaction["results"]) + snapshot = transaction["snapshot"] + published_path = transaction["published_path"] + + stage = "load" + results["load"] = run_load_phase.remote("load", run_id, published_path, snapshot) + stage = "load2" + results["load2"] = run_load_phase.remote("load2", run_id, published_path, snapshot) + + stage = "aggregate" + checks = _acceptance_checks(results, transaction) + except BaseException as error: + partial_result = getattr(error, "result", None) + partial_phase = getattr(error, "phase", stage) + if isinstance(partial_result, dict) and partial_phase in PHASES: + results.setdefault(partial_phase, partial_result) + if stage == "transaction" and isinstance(error, PublicationReconciliationError): + stage = "reconciliation" + lifecycle_error = error + + report = _report_evidence(run_id, run_url, results, transaction, checks) + if lifecycle_error is not None: + report.update( + { + "status": "NOT VERIFIED", + "failed_stage": stage, + "failed_error": str(lifecycle_error), + "failed_error_type": ( + f"{type(lifecycle_error).__module__}.{type(lifecycle_error).__name__}" + ), + "failed_checks": [], + "cleanup_attempted": False, + "cleanup_completed": False, + "cleanup_error": None, + "cleanup_error_type": None, + **_artifact_state_for_report(run_id), + } + ) + _emit_report(report) + raise RuntimeError( + f"TP validation NOT VERIFIED at {stage}: {lifecycle_error}" + ) from lifecycle_error + + failed = [name for name, passed in checks.items() if not passed] + if failed: + report.update( + { + "status": "NOT VERIFIED", + "failed_stage": "aggregate", + "failed_error": f"failed checks: {', '.join(failed)}", + "failed_error_type": "acceptance_checks", + "failed_checks": failed, + "cleanup_attempted": False, + "cleanup_completed": False, + "cleanup_error": None, + "cleanup_error_type": None, + **_artifact_state_for_report(run_id), + } + ) + _emit_report(report) + raise RuntimeError(f"TP validation NOT VERIFIED: {', '.join(failed)}") + + try: + cleanup = _normalize_cleanup_result(cleanup_published_run.remote(run_id)) + except BaseException as error: + report.update( + { + "status": "NOT VERIFIED", + "failed_stage": "cleanup", + "failed_error": str(error), + "failed_error_type": (f"{type(error).__module__}.{type(error).__name__}"), + "failed_checks": ["cleanup_completed"], + "cleanup_attempted": True, + "cleanup_completed": False, + "cleanup_error": str(error), + "cleanup_error_type": (f"{type(error).__module__}.{type(error).__name__}"), + **_artifact_state_for_report(run_id), + } + ) + _emit_report(report) + raise RuntimeError(f"TP validation NOT VERIFIED at cleanup: {error}") from error + + if not cleanup.get("cleanup_completed"): + report.update( + { + "status": "NOT VERIFIED", + "failed_stage": "cleanup", + "failed_error": cleanup.get("cleanup_error"), + "failed_error_type": cleanup.get("cleanup_error_type") or "cleanup_result", + "failed_checks": ["cleanup_completed"], + **cleanup, + } + ) + _emit_report(report) + raise RuntimeError(f"TP validation NOT VERIFIED: cleanup: {cleanup.get('cleanup_error')}") + + report.update( + { + "status": "VERIFIED", + "failed_stage": None, + "failed_error": None, + "failed_error_type": None, + "failed_checks": [], + **cleanup, + } + ) + _emit_report(report) diff --git a/tests/run_deepep_matrix.sh b/tests/run_deepep_matrix.sh new file mode 100755 index 00000000..02bccda4 --- /dev/null +++ b/tests/run_deepep_matrix.sh @@ -0,0 +1,108 @@ +#!/usr/bin/env bash +# Matrix driver for tests/test_deepep_fabric.py — standalone DeepEP save/load probes +# (no vLLM/sglang init; needs ~4 GB free on each of GPUs 0,1, NOT a fully empty GPU). +# +# fabric strict fabric transport (NVSHMEM heap only, zero NVL bytes), +# NCCL fast paths off; expected PASS +# ll_nofabric pure low-latency (NVSHMEM heap, FD handles), NCCL fast paths off +# -> retained diagnostic; expected PASS +# ll_ncclcumem same, but NCCL_CUMEM_ENABLE=1 -> probe: can NCCL cumem coexist? +# ll_ncclnvls same, but CUMEM=1 + NVLS=1 -> probe: can NCCL NVLS coexist? +# nvl_ipc + 64 MB NVL buffer, use_fabric=0 -> legacy cudaIpc path through the +# hook's VMM-IPC translation (SCM_RIGHTS fd transport; peer mappings +# relocate under shared region bases); expected PASS +# nvl_fabric + 64 MB NVL buffer, use_fabric=1 -> fabric cuMemCreate on a no-IMEX +# machine; documents the fabric dependency; expected FAIL +# +# Usage: run_deepep_matrix.sh +# +# Environment overrides: +# PYTHON_BIN Python with Foundry installed (default: python3) +# VLLM_ROOT vLLM checkout used to build NVSHMEM +# NVSHMEM_SO explicit libnvshmem_host.so path +# DEEPEP_MATRIX_WORK_ROOT per-case work directory root +# DEEPEP_MATRIX_LOG_DIR log output directory +set -euo pipefail + +SCRIPT_DIR=$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" && pwd) +REPO_ROOT=$(cd -- "$SCRIPT_DIR/.." && pwd) +PYTHON_BIN=${PYTHON_BIN:-python3} +VLLM_ROOT=${VLLM_ROOT:-"$REPO_ROOT/../vllm"} +NVSHMEM_SO=${NVSHMEM_SO:-"$VLLM_ROOT/tools/ep_kernels/ep_kernels_workspace/nvshmem/lib/libnvshmem_host.so"} +DEEPEP_MATRIX_WORK_ROOT=${DEEPEP_MATRIX_WORK_ROOT:-"$REPO_ROOT/tests/deepep_matrix_work"} +DEEPEP_MATRIX_LOG_DIR=${DEEPEP_MATRIX_LOG_DIR:-"$REPO_ROOT/logs/deepep_matrix"} +TEST="$REPO_ROOT/tests/test_deepep_fabric.py" + +CASE=${1:?usage: run_deepep_matrix.sh } +PHASE=${2:?usage: run_deepep_matrix.sh } + +export CUDA_VISIBLE_DEVICES=0,1 +# Pin NVSHMEM heap chunks to POSIX FD handles — explicit "no fabric handles anywhere". +# (Auto-detect picks FD on non-MNNVL machines anyway; pinning removes the variable.) +export NVSHMEM_CUMEM_HANDLE_TYPE=FILE_DESCRIPTOR + +case "$CASE" in + fabric) export TEST_USE_FABRIC=1 TEST_NVL_BYTES_MB=0 NCCL_CUMEM_ENABLE=0 NCCL_NVLS_ENABLE=0; PORT=29691 ;; + ll_nofabric) export TEST_USE_FABRIC=0 TEST_NVL_BYTES_MB=0 NCCL_CUMEM_ENABLE=0 NCCL_NVLS_ENABLE=0; PORT=29611 ;; + ll_ncclcumem) export TEST_USE_FABRIC=0 TEST_NVL_BYTES_MB=0 NCCL_CUMEM_ENABLE=1 NCCL_NVLS_ENABLE=0; PORT=29621 ;; + ll_ncclcumem_preinit) export TEST_USE_FABRIC=0 TEST_NVL_BYTES_MB=0 NCCL_CUMEM_ENABLE=1 NCCL_NVLS_ENABLE=0 TEST_NCCL_WARMUP_PRE_REGION=1; PORT=29661 ;; + ll_ncclnvls) export TEST_USE_FABRIC=0 TEST_NVL_BYTES_MB=0 NCCL_CUMEM_ENABLE=1 NCCL_NVLS_ENABLE=1; PORT=29631 ;; + ll_ncclnvls_preinit) export TEST_USE_FABRIC=0 TEST_NVL_BYTES_MB=0 NCCL_CUMEM_ENABLE=1 NCCL_NVLS_ENABLE=1 TEST_NCCL_WARMUP_PRE_REGION=1; PORT=29671 ;; + nvl_ipc) export TEST_USE_FABRIC=0 TEST_NVL_BYTES_MB=64 NCCL_CUMEM_ENABLE=0 NCCL_NVLS_ENABLE=0; PORT=29641 ;; + nvl_fabric) export TEST_USE_FABRIC=1 TEST_NVL_BYTES_MB=64 NCCL_CUMEM_ENABLE=0 NCCL_NVLS_ENABLE=0; PORT=29651 ;; + nvl_ipc_prealloc) export TEST_USE_FABRIC=0 TEST_NVL_BYTES_MB=64 NCCL_CUMEM_ENABLE=0 NCCL_NVLS_ENABLE=0 TEST_LOAD_PREALLOC_MB=1024; PORT=29681 ;; + *) echo "unknown case: $CASE" >&2; exit 2 ;; +esac + +case "$PHASE" in + save|load|both) ;; + *) echo "unknown phase: $PHASE" >&2; exit 2 ;; +esac + +if [[ ! -f "$NVSHMEM_SO" ]]; then + echo "NVSHMEM library not found: $NVSHMEM_SO" >&2 + exit 4 +fi + +HOOK_SO=$( + "$PYTHON_BIN" -c \ + "import foundry.ops, pathlib; print(pathlib.Path(foundry.ops.__file__).parent / 'libcuda_hook.so')" +) +if [[ ! -f "$HOOK_SO" ]]; then + echo "Foundry CUDA hook not found: $HOOK_SO" >&2 + exit 5 +fi + +# GPUs are shared with other users — refuse to run without headroom. +for i in 0 1; do + free=$(nvidia-smi --query-gpu=memory.free --format=csv,noheader,nounits -i "$i") + if [ "$free" -lt 6000 ]; then + echo "GPU $i has only ${free} MiB free — aborting (need ~6 GB headroom)"; exit 3 + fi +done + +# main() overwrites TEST_USE_FABRIC from the CLI flag, so pass the flag too. +FABRIC_ARGS=() +[ "$TEST_USE_FABRIC" = "0" ] && FABRIC_ARGS+=("--no-fabric") + +WORK="$DEEPEP_MATRIX_WORK_ROOT/$CASE" +LOGDIR="$DEEPEP_MATRIX_LOG_DIR" +mkdir -p "$WORK" "$LOGDIR" +cd "$WORK" # deepep_fabric_archive/ is CWD-relative + +export LD_PRELOAD="$NVSHMEM_SO:$HOOK_SO${LD_PRELOAD:+:$LD_PRELOAD}" + +run_phase() { + local p=$1 + # distinct rendezvous port per phase to dodge TIME_WAIT + if [ "$p" = save ]; then export MASTER_PORT=$PORT; else export MASTER_PORT=$((PORT + 1)); fi + echo "=== case=$CASE phase=$p $(date '+%F %T') ===" + echo "env: TEST_USE_FABRIC=$TEST_USE_FABRIC TEST_NVL_BYTES_MB=$TEST_NVL_BYTES_MB" \ + "NCCL_CUMEM_ENABLE=$NCCL_CUMEM_ENABLE NCCL_NVLS_ENABLE=$NCCL_NVLS_ENABLE" \ + "NVSHMEM_CUMEM_HANDLE_TYPE=$NVSHMEM_CUMEM_HANDLE_TYPE MASTER_PORT=$MASTER_PORT" + if [ "$p" = save ]; then rm -rf deepep_fabric_archive; fi + "$PYTHON_BIN" "$TEST" "--$p" "${FABRIC_ARGS[@]}" --num-processes=2 \ + 2>&1 | tee "$LOGDIR/${CASE}_${p}.log" +} + +if [ "$PHASE" = both ]; then run_phase save; run_phase load; else run_phase "$PHASE"; fi diff --git a/tests/sglang_tp_validation.py b/tests/sglang_tp_validation.py new file mode 100644 index 00000000..127c55c0 --- /dev/null +++ b/tests/sglang_tp_validation.py @@ -0,0 +1,1023 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the Foundry project +"""Pure-Python configuration and evidence helpers for the SGLang TP harness. + +This mirrors the reviewed vLLM helper module (``vllm_tp_validation.py``): the +Modal harness stays a thin orchestrator and every string-parsing / process +predicate lives here so it is exercised by CPU contracts before a GPU is +provisioned. + +The scan regexes are pinned to the current architecture: + +* Foundry emits its evidence through ``logging.getLogger(__name__)`` which + propagates to the root logger SGLang configures for each scheduler worker + (``configure_logger(server_args, prefix=" TP{rank}")`` → + ``[%(asctime)s TP{rank}] %(message)s``), so per-rank Foundry lines carry a + ``TP{rank}]`` marker. +* The exact Foundry message formats are: + * ``[Foundry] SGLang alloc_offset[%s]=%d (%.2f MB)`` (``runtime.log_alloc_offset``) + * ``[Foundry] Loaded %d SGLang graphs in %.3fs (...)`` (``graph_ops.load_all_graphs``) + * ``[Foundry] Saved SGLang CUDA graph %s key=%s`` (``graph_ops.save_graph``) +* SGLang's native decode progress bar prints ``Capturing batches (...)`` and + its decode metrics line is ``Decode batch, #running-req: N, ... cuda graph: True``. +""" + +from __future__ import annotations + +import os +import re +import signal +import subprocess +import time +from typing import IO + +FULL_COMMIT_SHA_PATTERN = re.compile(r"^[0-9a-fA-F]{40}$") + +# Two-window fatal scan. +# +# The harness stops each server with SIGINT (see ``stop_process``), and SGLang's +# workers, blocked in ``broadcast_pyobj`` / zmq ``recv`` at the time, print +# ``KeyboardInterrupt`` tracebacks and force-kill (-9) noise as part of that +# *expected* teardown. Scanning the full log with a single broad pattern flagged +# those benign tracebacks (and the benign ``watchdog_timeout=`` ServerArgs field) +# as fatal. The scan is therefore split into two evidence windows that the +# harness feeds separately (a pre-teardown snapshot + the full post-teardown log) +# and merges: +# +# * ``HARD_FATAL_PATTERN`` — hard native / CUDA / NCCL / MMU / Foundry faults and +# anchored worker/scheduler death phrases. These never arise from a clean +# SIGINT teardown, so they are fatal in *either* window (before or after +# teardown). Generic ``RuntimeError`` / ``Exception`` are deliberately *not* +# here; a generic Python failure is handled only in the pre-teardown window by +# ``PYTHON_TRACEBACK_PATTERN`` / ``PYTHON_FAILURE_PATTERN`` plus the +# recovered-retry-block logic below, so the expected SIGINT teardown traceback +# stays ignored post-teardown. The Foundry +# native strings are matched by class so a benign +# ``[HOOK] INFO`` / ``[REPLAY] OK`` / ``[foundry REPLAY]`` / ``[foundry BUILD]`` +# / ``[foundry] WARNING`` line is never a false positive: +# * ``[HOOK] FATAL ERROR:`` / ``[HOOK] ERROR:`` (csrc/hook.cpp) +# * ``[REPLAY] FATAL:`` / ``[REPLAY] ERROR:`` (csrc/hook.cpp) +# * every named ``[foundry ERROR]`` class — ``LOAD`` / ``INSTANTIATE`` +# / ``ON-DEMAND`` and any future class — via ``\[foundry [^\]]*ERROR\]``. +# * ``WATCHDOG_FAILURE_PATTERN`` / ``SUBPROCESS_CRASH_PATTERN`` — anchored real +# watchdog failures (``WatchdogError``, ``watchdog thread crashed``, `` +# watchdog timeout``) and subprocess crashes (``Subprocess (pid=N) +# crashed with exit code -K``), never the benign ``watchdog_timeout=`` / +# ``soft_watchdog_timeout=`` ServerArgs fields. Pre-teardown only (the harness +# force-kills the same workers during teardown, so the identical -9 line is +# expected child-exit noise post-teardown). +# * ``PYTHON_TRACEBACK_PATTERN`` / ``PYTHON_FAILURE_PATTERN`` — generic Python +# failures (a ``Traceback`` header or case-sensitive ``Error:`` match +# anywhere in a line), fatal everywhere in the pre-teardown window *except* +# lines inside an explicitly recognized recovered retry block: a contiguous +# chained traceback ending immediately before SGLang's +# ``retry() failed once (...). Will delay ...s and retry. Error: ...`` marker +# when the server later becomes ready. Unrelated pre-readiness tracebacks stay +# fatal. See :data:`RETRY_RECOVERED_MARKER` / :func:`_recovered_retry_line_indices`. +HARD_FATAL_PATTERN = re.compile( + r"\[HOOK\] (?:FATAL )?ERROR|\[REPLAY\] (?:FATAL|ERROR)|\[foundry [^\]]*ERROR\]|" + r"MMU fault|\bXid\b|CUDA error|" + r"illegal memory access|an illegal memory|NCCL WARN|NCCL error|" + r"Scheduler hit an exception|" + r"(?:TpModelWorkerClient|DataParallelController|Detokenizer).*" + r"(?:hit an exception|died|\bdead\b)|" + r"(?:worker|scheduler|engine)\b.*(?:died|death|crashed)|" + r"CUDA out of memory|segmentation fault|\bsegfault\b|core dumped", + re.IGNORECASE, +) +# Anchored watchdog *failures* — the space in ``watchdog timeout`` / ``watchdog +# thread crashed`` never matches the underscore in the ServerArgs +# ``watchdog_timeout=`` / ``soft_watchdog_timeout=`` fields. +WATCHDOG_FAILURE_PATTERN = re.compile( + r"WatchdogError|watchdog thread crashed|\S+ watchdog timeout", + re.IGNORECASE, +) +# Anchored SGLang subprocess crash, e.g. ``Subprocess detokenizer (pid=50) +# crashed with exit code -9``. Pre-teardown *only*: a worker crash mid-run is a +# genuine failure, but the identical line is expected child-exit noise when the +# harness force-stops the same workers during SIGINT teardown, so HARD (the +# post-teardown window) must keep ignoring it. +SUBPROCESS_CRASH_PATTERN = re.compile( + r"Subprocess \S+ \(pid=\d+\) crashed with exit code -?\d+", + re.IGNORECASE, +) +# SGLang readiness markers. Once any of these appears the HTTP server and every +# scheduler worker are up and serving, so a generic Python traceback / exception +# after this point is a genuine runtime failure. Startup work that emits a +# recovered traceback (e.g. SGLang's ``retry()`` wrapper around a transient HF +# Hub connection reset during model-config resolution) happens *before* readiness +# and must not be treated as fatal. +READINESS_PATTERN = re.compile( + r"The server is fired up and ready to roll!|" + r"Application startup complete\.|" + r"Uvicorn running on " +) +# Generic Python failure markers, applied to the pre-teardown window only (never +# post-teardown, where the expected SIGINT ``Traceback`` / ``KeyboardInterrupt`` +# cascade lives). +# +# * ``Traceback (most recent call last):`` — a Python traceback header. +# * A case-sensitive exception type ``Error:`` / +# ``...Exception:`` anywhere in a line, or an exact timeout class whose name +# does not use those suffixes. Requiring at least one type character before +# the suffix means the lowercase word "exception" in the +# standard chain separator ``The above exception was the direct cause of the +# following exception:`` — and a bare ``Error:`` in retry prose — are never +# matched. This replaces the old whole-log bare ``\bException\b`` / bare +# ``RuntimeError`` matching. +PYTHON_EXCEPTION_TYPE_SOURCE = ( + r"[A-Za-z_][\w.]*(?:Error|Exception|ConnectTimeout|ReadTimeout|PoolTimeout)" +) +PYTHON_TRACEBACK_PATTERN = re.compile(r"Traceback \(most recent call last\):") +PYTHON_FAILURE_PATTERN = re.compile(rf"(?{PYTHON_EXCEPTION_TYPE_SOURCE}):") +# SGLang's exact "recovered retry" warning, e.g. +# ``retry() failed once (0th try, maximum 2 retries). Will delay 0.79s and retry. +# Error: [Errno 104] Connection reset by peer`` (srt/utils/common.py retry()). +# Only a generic Python-failure block *immediately preceding* this marker — and +# only when the server later becomes ready — is treated as recovered. +RETRY_RECOVERED_MARKER = re.compile( + r"retry\(\) failed once \(\d+(?:st|nd|rd|th) try, maximum \d+ retries\)\. " + r"Will delay [\d.]+s and retry\. Error:" +) +# A new timestamped log record ``[YYYY-MM-DD HH:MM:SS ...] ...``. Raw Python +# traceback lines never carry this prefix, so it bounds the contiguous traceback +# block that precedes a retry marker. +LOG_RECORD_PREFIX = re.compile(r"^\[\d{4}-\d{2}-\d{2} \d{2}:\d{2}:\d{2}") + + +def _recovered_retry_line_indices(lines: list[str]) -> set[int]: + """Line indices whose generic Python-failure matches are suppressed because + they belong to an explicitly recognized *recovered retry block*. + + A recovered block is the contiguous chained traceback (one or more + ``Traceback`` headers and chain separators) that ends immediately before a + :data:`RETRY_RECOVERED_MARKER` line, bounded above by the nearest preceding + timestamped log record (:data:`LOG_RECORD_PREFIX`). Suppression additionally + requires a readiness marker *after* the retry marker; without an eventual + readiness the block stays fatal. Blocks separated from the marker by another + timestamped log record are not associated and are never suppressed. + """ + readiness_index = next( + (i for i, line in enumerate(lines) if READINESS_PATTERN.search(line)), None + ) + if readiness_index is None: + return set() + suppressed: set[int] = set() + for i, line in enumerate(lines): + if not RETRY_RECOVERED_MARKER.search(line): + continue + # Require an eventual readiness marker strictly after the retry marker. + if i >= readiness_index: + continue + # Contiguous non-log-record run immediately preceding the marker. + start = i + while start - 1 >= 0 and not LOG_RECORD_PREFIX.match(lines[start - 1]): + start -= 1 + run = range(start, i) + header = next((k for k in run if "Traceback (most recent call last):" in lines[k]), None) + if header is not None: + suppressed.update(range(header, i)) + return suppressed + + +# Foundry per-rank LOAD final offset. ``.*`` bridges the SGLang ``TP{rank}]`` +# prefix and the Foundry message on the same log line. +LOAD_OFFSET_PATTERN = re.compile( + r"TP(?P\d+)\].*alloc_offset\[after_load_all_graphs\]=(?P\d+)" +) +# Foundry per-rank LOAD summary count. The plural "graphs" + count distinguishes +# the summary from the per-graph "Loaded SGLang graph idx=" line. +LOADED_GRAPHS_PATTERN = re.compile(r"TP(?P\d+)\].*Loaded (?P\d+) SGLang graphs") +TORCH_SYMM_MEM_PATTERN = re.compile( + r"(?:^|[\r\n])\[\d{4}-\d{2}-\d{2} \d{2}:\d{2}:\d{2} TP(?P\d+)\] " + r"\[Foundry\] Torch symmetric-memory buffer address=" + r"(?P
0x[0-9a-fA-F]+)$", + re.MULTILINE, +) +MULTIMEM_ALL_GATHER_PATTERN = re.compile( + r"(?:^|[\r\n])\[\d{4}-\d{2}-\d{2} \d{2}:\d{2}:\d{2} TP(?P\d+)\] " + r"\[Foundry\] Multimem all-gather buffer address=" + r"(?P
0x[0-9a-fA-F]+)$", + re.MULTILINE, +) +# NCCL's own connection log carries a [rank] bracket, not SGLang's prefix. +NCCL_TRANSPORT_PATTERN = re.compile( + r"\[(?P\d+)\] NCCL INFO Channel (?P\d+)/\d+ : .* via (?P\S+)" +) +# SGLang decode metrics line; #running-req is the concurrent decode batch width. +GRAPH_REPLAY_PATTERN = re.compile( + r"Decode batch(?: \[\d+\])?, #running-req: (?P\d+).*cuda graph: True" +) +SAVED_GRAPH_MESSAGE = "[Foundry] Saved SGLang CUDA graph" +NATIVE_CAPTURE_MESSAGE = "Capturing batches" +# NCCL prints its own version banner once per process; used only as best-effort +# runtime provenance and never gathered by initializing CUDA in the parent. +NCCL_VERSION_PATTERN = re.compile(r"NCCL version (?P\d+\.\d+(?:\.\d+)?)") + +EXPECTED_SHUTDOWN_RETURNCODES = {0, -2, 130} + + +def require_commit_sha(name: str, value: str | None) -> str: + if value is None or FULL_COMMIT_SHA_PATTERN.fullmatch(value) is None: + raise ValueError(f"{name} must be a full 40-hex commit SHA, got {value!r}") + return value.lower() + + +def require_tp_size(value: str | int) -> int: + try: + tp_size = int(value) + except (TypeError, ValueError) as error: + raise ValueError(f"TP_SIZE must be exactly 2, got {value!r}") from error + if tp_size != 2: + raise ValueError(f"TP_SIZE must be exactly 2, got {tp_size}") + return tp_size + + +def is_symmetric_policy_failure(failure_text: str) -> bool: + """Match only graph-restore/symmetric-policy failures, not benign Foundry phases.""" + normalized = failure_text.lower() + exact_policy_signature = any( + token in normalized + for token in ( + "[foundry load error]", + "[foundry instantiate error]", + "[foundry on-demand error]", + "[hook] fatal error:", + "[hook] error:", + "[replay] fatal:", + "[replay] error:", + "archive immutability", + "fallback to nccl", + "fell back to nccl", + "nccl fallback", + ) + ) + symmetric_failure = any( + policy in normalized + for policy in ("torch symmetric-memory", "symm_mem", "multimem all-gather") + ) and any( + failure in normalized for failure in ("failed", "failure", "error", "fallback", "fell back") + ) + return exact_policy_signature or symmetric_failure + + +ALLOWLISTED_STARTUP_EXCEPTION_TYPES = { + "ConnectionError", + "ConnectionAbortedError", + "ConnectionRefusedError", + "ConnectionResetError", + "ConnectError", + "DownloadError", + "HfHubHTTPError", + "LocalEntryNotFoundError", + "NetworkError", + "ProxyError", + "URLError", + "TimeoutError", + "ConnectTimeout", + "ReadTimeout", + "PoolTimeout", + "TimeoutException", + "builtins.ConnectionError", + "builtins.ConnectionAbortedError", + "builtins.ConnectionRefusedError", + "builtins.ConnectionResetError", + "builtins.TimeoutError", + "urllib.error.URLError", + "httpx.ConnectError", + "httpx.ConnectTimeout", + "httpx.ReadTimeout", + "httpx.PoolTimeout", + "httpx.ProxyError", + "httpx.TimeoutException", + "httpcore.ConnectError", + "httpcore.ConnectTimeout", + "httpcore.ReadTimeout", + "httpcore.PoolTimeout", + "httpcore.ProxyError", + "requests.exceptions.ConnectionError", + "requests.exceptions.ConnectTimeout", + "requests.exceptions.ReadTimeout", + "requests.exceptions.ProxyError", + "requests.exceptions.Timeout", + "huggingface_hub.errors.HfHubHTTPError", + "huggingface_hub.errors.LocalEntryNotFoundError", + "huggingface_hub.utils._errors.HfHubHTTPError", + "huggingface_hub.utils._errors.LocalEntryNotFoundError", +} + + +def is_allowlisted_startup_environment_failure(result: dict) -> bool: + """Classify only structured pre-health network/HF/download failures.""" + if result.get("failure_stage") != "startup" or result.get("archive_committed"): + return False + if result.get("evidence_errors"): + return False + child_exited = result.get("startup_child_exited") is True + event_stages = {event.get("stage") for event in result.get("failure_events", [])} + if event_stages - {"startup", "teardown"}: + return False + if "teardown" in event_stages and not child_exited: + return False + disallowed_categories = { + "hard_native", + "policy", + "harness", + "request", + "scanner", + "archive", + "watchdog", + "subprocess", + } + if disallowed_categories & set(result.get("error_categories", [])): + return False + + untyped_python_evidence = result.get( + "untyped_python_fatal_evidence", + [], + ) + typed_traceback_headers = result.get( + "typed_python_traceback_headers", + [], + ) + if any(evidence not in typed_traceback_headers for evidence in untyped_python_evidence): + return False + + failure_type = str(result.get("failure_type") or "") + startup_envelope = failure_type.rsplit(".", 1)[-1] == "FocusedStartupHealthError" + cause_types = {str(error_type) for error_type in result.get("startup_cause_types", [])} + if not cause_types and not startup_envelope and failure_type: + cause_types.add(failure_type) + python_types = {str(error_type) for error_type in result.get("python_error_types", [])} + available_types = cause_types | python_types + if not available_types or not available_types <= ALLOWLISTED_STARTUP_EXCEPTION_TYPES: + return False + specific_timeout_types = { + "ConnectTimeout", + "ReadTimeout", + "PoolTimeout", + "TimeoutException", + } + cause_type_names = {error_type.rsplit(".", 1)[-1] for error_type in cause_types} + python_type_names = {error_type.rsplit(".", 1)[-1] for error_type in python_types} + return not ( + not startup_envelope + and cause_type_names == {"TimeoutError"} + and not python_type_names & specific_timeout_types + ) + + +def is_modal_infrastructure_error_type(module_name: str, type_name: str) -> bool: + """Allow only exact Modal infrastructure classes, never generic remote errors.""" + return module_name == "modal.exception" and type_name in { + "ClientClosed", + "ConnectionError", + "InternalError", + "InternalFailure", + "ServiceError", + "TimeoutError", + } + + +def _tri_and(*values: bool | None) -> bool | None: + if any(value is False for value in values): + return False + if values and all(value is True for value in values): + return True + return None + + +def _after_request(result: dict | None, value: bool) -> bool | None: + if result is None or not result.get("request_completed"): + return None + return value + + +def _phase_failure_text(result: dict | None) -> str: + if result is None: + return "" + return " ".join( + str(item) + for item in [ + *result.get("phase_failures", []), + *result.get("errors", []), + ] + ) + + +def _has_direct_multimem_gather_marker(result: dict) -> bool: + addresses = result.get("multimem_all_gather_addresses") + if not isinstance(addresses, dict): + return False + return any( + isinstance(rank_addresses, list) + and any(isinstance(address, str) and bool(address) for address in rank_addresses) + for rank_addresses in addresses.values() + ) + + +def _archive_collective_evidence( + result: dict, + expected_ranks: set[str], +) -> tuple[bool, bool]: + counts = result.get("archive_collective_kernel_counts") + if not isinstance(counts, dict) or set(counts) != expected_ranks: + return False, False + + rank_counts = [counts[rank] for rank in expected_ranks] + if not all(isinstance(value, dict) for value in rank_counts): + return False, False + + no_nccl_collectives = all( + type(value.get("all_reduce")) is int + and value["all_reduce"] == 0 + and type(value.get("all_gather")) is int + and value["all_gather"] == 0 + for value in rank_counts + ) + symmetric_collectives_present = all( + type(value.get("symmetric_all_reduce")) is int + and value["symmetric_all_reduce"] > 0 + and type(value.get("symmetric_all_gather")) is int + and value["symmetric_all_gather"] > 0 + for value in rank_counts + ) + return no_nccl_collectives, symmetric_collectives_present + + +def _symmetric_markers_present(result: dict, expected_ranks: set[str]) -> bool: + return set(result.get("torch_symm_mem_ranks", [])) == expected_ranks and ( + _has_direct_multimem_gather_marker(result) + ) + + +def focused_save_checks( + result: dict | None, + *, + tp_size: int, + require_committed: bool = False, +) -> dict[str, bool | None]: + """Evaluate SAVE prerequisites with ``None`` for unavailable pre-health evidence.""" + if result is None: + return { + "request_completed": None, + "graph_counts_exact": None, + "graph_json_counts_exact": None, + "graph_inventory_matches": None, + "symmetric_markers_present": None, + "symmetric_archive_collectives_present": None, + "no_archived_nccl_collectives": None, + "offsets_equal_nonzero": None, + "archive_evidence_complete": None, + "no_scanner_errors": None, + "no_runtime_errors": None, + "expected_shutdown": None, + "no_policy_failure": None, + "post_health_execution_clean": None, + "no_harness_failures": None, + "archive_committed": None, + } + + expected_ranks = {str(rank) for rank in range(tp_size)} + expected_graph_counts = {rank: 1 for rank in expected_ranks} + no_nccl_collectives, symmetric_collectives_present = _archive_collective_evidence( + result, + expected_ranks, + ) + request_completed = ( + True if result.get("request_completed") else False if result.get("became_healthy") else None + ) + graph_counts = result.get("archive_graph_counts", {}) + graph_json_counts = result.get("archive_graph_json_counts", {}) + offsets = result.get("rank_offsets", {}) + phase_failures = result.get("phase_failures", []) + failure_text = _phase_failure_text(result) + startup_environment_failure = is_allowlisted_startup_environment_failure(result) + return { + "request_completed": request_completed, + "graph_counts_exact": _after_request( + result, + graph_counts == expected_graph_counts, + ), + "graph_json_counts_exact": _after_request( + result, + graph_json_counts == expected_graph_counts, + ), + "graph_inventory_matches": _after_request( + result, + graph_counts == graph_json_counts == expected_graph_counts, + ), + "symmetric_markers_present": _after_request( + result, + _symmetric_markers_present(result, expected_ranks), + ), + "symmetric_archive_collectives_present": _after_request( + result, + symmetric_collectives_present, + ), + "no_archived_nccl_collectives": _after_request( + result, + no_nccl_collectives, + ), + "offsets_equal_nonzero": _after_request( + result, + set(offsets) == expected_ranks + and all(offsets.values()) + and len(set(offsets.values())) == 1, + ), + "archive_evidence_complete": _after_request( + result, + set(result.get("archive_fingerprints", {})) == expected_ranks + and bool(result.get("archive_byte_fingerprints")), + ), + "no_scanner_errors": not bool(result.get("evidence_errors")), + "no_runtime_errors": ( + None if startup_environment_failure else not bool(result.get("errors")) + ), + "expected_shutdown": ( + None + if startup_environment_failure + else result["process_return_code"] in EXPECTED_SHUTDOWN_RETURNCODES + if "process_return_code" in result + else None + ), + "no_policy_failure": not is_symmetric_policy_failure(failure_text), + "post_health_execution_clean": ( + not bool(phase_failures) + if result.get("became_healthy") + else None + if phase_failures + else True + ), + "no_harness_failures": ( + True + if startup_environment_failure + else not any( + event.get("stage") in {"request", "teardown", "scanner", "archive"} + for event in result.get("failure_events", []) + ) + ), + "archive_committed": ( + None + if require_committed and not result.get("became_healthy") + else result.get("archive_committed") is True + if require_committed + else True + ), + } + + +def focused_treatment_checks( + save_result: dict | None, + load_result: dict | None, + *, + tp_size: int, +) -> dict[str, bool | None]: + """Evaluate focused SAVE→LOAD evidence without performing I/O.""" + expected_ranks = {str(rank) for rank in range(tp_size)} + expected_graph_counts = {rank: 1 for rank in expected_ranks} + save_checks = focused_save_checks( + save_result, + tp_size=tp_size, + require_committed=True, + ) + load_environment_failure = ( + load_result is not None and is_allowlisted_startup_environment_failure(load_result) + ) + save_ready = save_result is not None and bool(save_result.get("request_completed")) + load_ready = load_result is not None and bool(load_result.get("request_completed")) + both_ready = save_ready and load_ready + save_outputs = save_result.get("outputs", []) if save_result else [] + load_outputs = load_result.get("outputs", []) if load_result else [] + outputs_equal = ( + len(save_outputs) == 1 + and isinstance(save_outputs[0], str) + and bool(save_outputs[0].encode()) + and len(load_outputs) == 1 + and isinstance(load_outputs[0], str) + and save_outputs[0].encode() == load_outputs[0].encode() + ) + save_graph_counts = save_result.get("archive_graph_counts", {}) if save_result else {} + load_graph_counts = load_result.get("archive_graph_counts", {}) if load_result else {} + save_json_counts = save_result.get("archive_graph_json_counts", {}) if save_result else {} + load_json_counts = load_result.get("archive_graph_json_counts", {}) if load_result else {} + save_offsets = save_result.get("rank_offsets", {}) if save_result else {} + load_offsets = load_result.get("load_offsets", {}) if load_result else {} + save_fingerprints = save_result.get("archive_fingerprints", {}) if save_result else {} + load_fingerprints = load_result.get("archive_fingerprints", {}) if load_result else {} + save_byte_fingerprints = save_result.get("archive_byte_fingerprints", {}) if save_result else {} + load_byte_fingerprints = load_result.get("archive_byte_fingerprints", {}) if load_result else {} + collective_evidence = [ + _archive_collective_evidence(result, expected_ranks) + for result in (save_result, load_result) + if result is not None + ] + checks: dict[str, bool | None] = { + **{f"save_{name}": value for name, value in save_checks.items()}, + "outputs_equal_nonempty": outputs_equal if both_ready else None, + "save_graph_counts_exact": save_checks["graph_counts_exact"], + "save_graph_json_counts_exact": save_checks["graph_json_counts_exact"], + "load_graph_counts_exact": _after_request( + load_result, + load_graph_counts == expected_graph_counts, + ), + "load_graph_json_counts_exact": _after_request( + load_result, + load_json_counts == expected_graph_counts, + ), + "graphs_restored_each_rank": _after_request( + load_result, + load_result is not None and load_result.get("loaded_graphs") == expected_graph_counts, + ), + "graph_inventory_matches": ( + save_graph_counts + == save_json_counts + == load_graph_counts + == load_json_counts + == expected_graph_counts + if both_ready + else None + ), + "symmetric_markers_each_phase": ( + all( + _symmetric_markers_present(result, expected_ranks) + for result in (save_result, load_result) + if result is not None + ) + if both_ready + else None + ), + "symmetric_archive_collectives_each_phase": ( + all(symmetric_present for _, symmetric_present in collective_evidence) + if both_ready + else None + ), + "no_archived_nccl_collectives": ( + all(no_nccl for no_nccl, _ in collective_evidence) if both_ready else None + ), + "rank_offsets_equal_nonzero": ( + set(save_offsets) == expected_ranks + and all(save_offsets.values()) + and len(set(save_offsets.values())) == 1 + and load_offsets == save_offsets + if both_ready + else None + ), + "restored_graph_replay_observed": _after_request( + load_result, + load_result is not None and 1 in load_result.get("graph_replay_batch_sizes", []), + ), + # The pinned LOAD lifecycle runs the real upstream capture() traversal, + # so rank0 emits "Capturing batches" progress even though the Foundry + # backend restores archived graphs without any CUDA capture. That + # progress count is diagnostic only; no-recapture is proven by positive + # archived-load counts together with zero Foundry graph-save logs. + "load_restored_without_resave": _after_request( + load_result, + load_result is not None + and load_result.get("saved_graph_log_count") == 0 + and sum((load_result.get("loaded_graphs") or {}).values()) > 0, + ), + "load_archive_immutable": ( + bool(save_fingerprints) + and load_fingerprints == save_fingerprints + and load_graph_counts == save_graph_counts + if both_ready + else None + ), + "load_archive_bytes_immutable": ( + bool(save_byte_fingerprints) and load_byte_fingerprints == save_byte_fingerprints + if both_ready + else None + ), + "no_scanner_errors": _tri_and( + save_checks["no_scanner_errors"], + (not bool(load_result.get("evidence_errors")) if load_result is not None else None), + ), + "no_runtime_errors": _tri_and( + save_checks["no_runtime_errors"], + ( + None + if load_environment_failure + else not bool(load_result.get("errors")) + if load_result is not None + else None + ), + save_checks["post_health_execution_clean"], + ( + not bool(load_result.get("phase_failures")) + if load_result is not None and load_result.get("became_healthy") + else None + ), + ), + "no_policy_failure": not is_symmetric_policy_failure( + " ".join(_phase_failure_text(result) for result in (save_result, load_result)) + ), + "expected_process_shutdown": _tri_and( + save_checks["expected_shutdown"], + ( + None + if load_environment_failure + else load_result["process_return_code"] in EXPECTED_SHUTDOWN_RETURNCODES + if load_result is not None and "process_return_code" in load_result + else None + ), + ), + } + return checks + + +def classify_focused_checks( + checks: dict[str, bool | None], + *, + environment_failure: bool, + restored_replay: bool, +) -> str: + """Classify focused evidence with known failures taking precedence.""" + if checks and all(value is True for value in checks.values()) and not environment_failure: + return "VERIFIED" + known_failure = any(value is False for value in checks.values()) + if environment_failure and not restored_replay and not known_failure: + return "INCONCLUSIVE" + return "NOT VERIFIED" + + +def scan_pre_teardown_errors(text: str) -> list[str]: + """Fatal lines in the pre-teardown window. + + Hard native/CUDA/NCCL/MMU/Foundry faults and anchored watchdog / subprocess + crashes are fatal *anywhere* and are never suppressible. Generic Python + failures (a traceback header or an anchored ``Error:`` line) are fatal + everywhere too, *except* the lines belonging to an explicitly recognized + recovered retry block — a contiguous chained traceback ending immediately + before SGLang's ``retry() failed once (...). Will delay ...s and retry. + Error: ...`` marker when the server later becomes ready (see + :func:`_recovered_retry_line_indices`). Every unrelated pre-readiness + traceback / error stays fatal even if the server later becomes ready, and + when readiness never appears nothing is suppressed. + """ + lines = text.splitlines() + fatal = { + line + for line in lines + if HARD_FATAL_PATTERN.search(line) + or WATCHDOG_FAILURE_PATTERN.search(line) + or SUBPROCESS_CRASH_PATTERN.search(line) + } + suppressed = _recovered_retry_line_indices(lines) + fatal.update( + line + for index, line in enumerate(lines) + if index not in suppressed + and (PYTHON_TRACEBACK_PATTERN.search(line) or PYTHON_FAILURE_PATTERN.search(line)) + ) + return sorted(fatal) + + +def scan_post_teardown_errors(text: str) -> list[str]: + """Fatal lines in the post-teardown window: hard native / CUDA / NCCL / MMU / + Foundry faults only (:data:`HARD_FATAL_PATTERN`). Generic Python failures and + the bare ``Traceback`` / ``KeyboardInterrupt`` emitted by the expected SIGINT + teardown are intentionally not counted here.""" + return sorted({line for line in text.splitlines() if HARD_FATAL_PATTERN.search(line)}) + + +def _python_traceback_evidence( + text: str, +) -> tuple[list[dict[str, int | str]], list[dict[str, int | str]]]: + untyped_headers: list[dict[str, int | str]] = [] + typed_headers: list[dict[str, int | str]] = [] + lines = text.splitlines() + suppressed = _recovered_retry_line_indices(lines) + for index, line in enumerate(lines): + if index in suppressed or PYTHON_TRACEBACK_PATTERN.search(line) is None: + continue + if PYTHON_FAILURE_PATTERN.search(line): + continue + evidence = {"line": index + 1, "text": line} + untyped_headers.append(evidence) + for candidate in lines[index + 1 :]: + # A timestamp starts a separate log record even when that record + # contains an exception type; it cannot terminate this traceback. + if LOG_RECORD_PREFIX.match(candidate): + break + if PYTHON_TRACEBACK_PATTERN.search(candidate): + break + if PYTHON_FAILURE_PATTERN.search(candidate): + typed_headers.append(evidence) + break + return untyped_headers, typed_headers + + +def scan_log_text(text: str, pre_teardown_text: str | None = None) -> dict: + """Extract every acceptance signal from one phase log. + + ``text`` is the full (post-teardown) log; all parsed signals are read from it. + ``errors`` merges the two evidence windows: the pre-teardown snapshot + (``pre_teardown_text``, scanned with :func:`scan_pre_teardown_errors`) and the + full log (scanned with :func:`scan_post_teardown_errors`). When + ``pre_teardown_text`` is ``None`` the pre-teardown window scans ``text`` too, + which keeps single-argument callers (CPU fixtures) broad; the harness always + passes an explicit pre-teardown snapshot so the expected SIGINT teardown noise + is excluded. + + Returns a dict with: + + * ``errors`` — sorted unique fatal lines (positive proof of no crash). + * ``python_error_types`` — exact case-sensitive types parsed from fatal lines. + * ``untyped_python_fatal_evidence`` — fatal traceback headers with no type on + the same line. + * ``typed_python_traceback_headers`` — those headers whose following + traceback block contains a parsed terminal exception type. + * ``load_offsets`` — ``{rank: after_load_all_graphs offset}``. + * ``loaded_graphs`` — ``{rank: restored graph count}``. + * ``torch_symm_mem_ranks`` — ranks logging retained symmetric communicators. + * ``multimem_all_gather_ranks`` — ranks logging direct multimem gather buffers. + * ``nccl_transport_channels`` — ``{rank: {transport: sorted channel ids}}``. + * ``graph_replay_batch_sizes`` — decode ``#running-req`` values under cuda graph. + * ``native_capture_progress_count`` — count of SGLang native capture progress + lines. Diagnostic only: the pinned LOAD lifecycle traverses the real + upstream capture() loop, so rank0 emits these on LOAD while the Foundry + backend performs no CUDA capture. It is NOT an acceptance gate. + * ``saved_graph_log_count`` — count of Foundry graph-save lines. + """ + pre_source = text if pre_teardown_text is None else pre_teardown_text + errors = sorted( + set(scan_pre_teardown_errors(pre_source)) | set(scan_post_teardown_errors(text)) + ) + error_text = "\n".join(errors) + error_categories = set() + if any(HARD_FATAL_PATTERN.search(line) for line in errors): + error_categories.add("hard_native") + if any(WATCHDOG_FAILURE_PATTERN.search(line) for line in errors): + error_categories.add("watchdog") + if any(SUBPROCESS_CRASH_PATTERN.search(line) for line in errors): + error_categories.add("subprocess") + if is_symmetric_policy_failure(error_text): + error_categories.add("policy") + python_error_types = sorted( + {match.group("type") for line in errors for match in PYTHON_FAILURE_PATTERN.finditer(line)} + ) + ( + untyped_python_fatal_evidence, + typed_python_traceback_headers, + ) = _python_traceback_evidence( + pre_source, + ) + load_offsets = { + match.group("rank"): int(match.group("offset")) + for match in LOAD_OFFSET_PATTERN.finditer(text) + } + loaded_graphs = { + match.group("rank"): int(match.group("count")) + for match in LOADED_GRAPHS_PATTERN.finditer(text) + } + transport_channels: dict[str, dict[str, set[int]]] = {} + for match in NCCL_TRANSPORT_PATTERN.finditer(text): + rank_transports = transport_channels.setdefault(match.group("rank"), {}) + rank_transports.setdefault(match.group("transport"), set()).add(int(match.group("channel"))) + graph_replay_batch_sizes = [ + int(match.group("batch_size")) for match in GRAPH_REPLAY_PATTERN.finditer(text) + ] + torch_symm_mem_addresses: dict[str, set[str]] = {} + for match in TORCH_SYMM_MEM_PATTERN.finditer(text): + torch_symm_mem_addresses.setdefault(match.group("rank"), set()).add( + match.group("address").lower() + ) + multimem_all_gather_addresses: dict[str, set[str]] = {} + for match in MULTIMEM_ALL_GATHER_PATTERN.finditer(text): + multimem_all_gather_addresses.setdefault(match.group("rank"), set()).add( + match.group("address").lower() + ) + nccl_version_match = NCCL_VERSION_PATTERN.search(text) + return { + "errors": errors, + "error_categories": sorted(error_categories), + "python_error_types": python_error_types, + "untyped_python_fatal_evidence": untyped_python_fatal_evidence, + "typed_python_traceback_headers": typed_python_traceback_headers, + "nccl_version": nccl_version_match.group("version") if nccl_version_match else None, + "load_offsets": load_offsets, + "loaded_graphs": loaded_graphs, + "torch_symm_mem_ranks": sorted(torch_symm_mem_addresses, key=int), + "multimem_all_gather_ranks": sorted(multimem_all_gather_addresses, key=int), + "torch_symm_mem_addresses": { + rank: sorted(addresses) for rank, addresses in sorted(torch_symm_mem_addresses.items()) + }, + "multimem_all_gather_addresses": { + rank: sorted(addresses) + for rank, addresses in sorted(multimem_all_gather_addresses.items()) + }, + "nccl_transport_channels": { + rank: { + transport: sorted(channels) + for transport, channels in sorted(rank_transports.items()) + } + for rank, rank_transports in sorted(transport_channels.items()) + }, + "graph_replay_batch_sizes": graph_replay_batch_sizes, + "native_capture_progress_count": text.count(NATIVE_CAPTURE_MESSAGE), + "saved_graph_log_count": text.count(SAVED_GRAPH_MESSAGE), + } + + +def channels_are_p2p_ipc(channels: dict, expected_ranks: set[str]) -> bool: + """Every rank present, every transport P2P/IPC-prefixed, every channel set nonempty. + + Accepts NCCL P2P/IPC submodes such as ``P2P/IPC/read`` (the transport token + is whatever NCCL logged after ``via``), and fails closed on a missing rank, + an empty channel list, or any non-P2P/IPC transport on any rank. + """ + if set(channels) != set(expected_ranks): + return False + for rank_channels in channels.values(): + if not rank_channels: + return False + for transport, channel_ids in rank_channels.items(): + if not channel_ids or not transport.startswith("P2P/IPC"): + return False + return True + + +def validate_shutdown_returncode(returncode: int, died_before_teardown: bool) -> None: + if died_before_teardown: + raise RuntimeError(f"SGLang process died before intentional teardown (code {returncode})") + if returncode not in EXPECTED_SHUTDOWN_RETURNCODES: + raise RuntimeError(f"SGLang process exited with unexpected code {returncode}") + + +def _process_group_has_members(pgid: int) -> bool: + """True if the process group still has at least one live member.""" + try: + os.killpg(pgid, 0) + except ProcessLookupError: + return False + except PermissionError: + # A member exists but is owned by another uid; treat as still present. + return True + return True + + +def _wait_for_group_exit(pgid: int, timeout: float) -> bool: + """Poll until the whole process group has drained or the timeout elapses.""" + deadline = time.monotonic() + timeout + while time.monotonic() < deadline: + if not _process_group_has_members(pgid): + return True + time.sleep(0.05) + return not _process_group_has_members(pgid) + + +def stop_process( + process: subprocess.Popen, + log_file: IO[str], + timeout: float = 120, + group_timeout: float = 30, +) -> int: + """SIGINT the child process group, drain it, close the log, and validate exit. + + Mirrors the reviewed vLLM helper and hardens teardown for the SGLang launcher + tree (a ``bash`` launcher that spawns a scheduler which spawns per-rank + workers, all in one process group): a child that already exited before the + intentional teardown is a failure, only an explicit clean/SIGINT code + (0, -2, 130) is accepted afterwards, and the *entire* process group must drain + after the launcher exits. The process-group id is captured before signalling + (it survives the leader's exit while any member is alive); if any group member + is still alive after ``group_timeout`` it is force-killed and teardown fails. + """ + died_before_teardown = process.poll() is not None + pgid: int | None = None + if not died_before_teardown: + try: + pgid = os.getpgid(process.pid) + except ProcessLookupError: + pgid = None + + lingering_group_children = False + try: + if not died_before_teardown and pgid is not None: + os.killpg(pgid, signal.SIGINT) + try: + process.wait(timeout=timeout) + except subprocess.TimeoutExpired: + os.killpg(pgid, signal.SIGKILL) + process.wait(timeout=30) + # The launcher has exited; ensure every worker in the group also did. + if not _wait_for_group_exit(pgid, group_timeout): + lingering_group_children = True + try: + os.killpg(pgid, signal.SIGKILL) + except ProcessLookupError: + lingering_group_children = False + else: + _wait_for_group_exit(pgid, 30) + finally: + log_file.close() + if process.returncode is None: + raise RuntimeError("SGLang process did not produce a return code after teardown") + validate_shutdown_returncode(process.returncode, died_before_teardown) + if lingering_group_children: + raise RuntimeError( + "SGLang process group had lingering group children after SIGINT; force-killed" + ) + return process.returncode diff --git a/tests/test_allocation_region_cleanup.py b/tests/test_allocation_region_cleanup.py new file mode 100644 index 00000000..06a437a8 --- /dev/null +++ b/tests/test_allocation_region_cleanup.py @@ -0,0 +1,74 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the Foundry project +"""CPU-only contracts for allocation-region teardown failures.""" + +from __future__ import annotations + +import importlib.util +import sys +from pathlib import Path +from types import ModuleType, SimpleNamespace + +import pytest + +REPO_ROOT = Path(__file__).parents[1] +ALLOCATION_REGION_PATH = REPO_ROOT / "python" / "foundry" / "allocation_region.py" + + +def _load_allocation_region(monkeypatch, cleanup_result: bool): + events = [] + foundry = ModuleType("foundry") + foundry.__path__ = [] + monkeypatch.setitem(sys.modules, "foundry", foundry) + + ops = ModuleType("foundry.ops") + ops.free_preallocated_region = lambda: events.append("free") or cleanup_result + ops.get_current_alloc_offset = lambda: 0 + ops.preallocate_region = lambda size: events.append(("preallocate", size)) or True + ops.resume_allocation_region = lambda: None + ops.set_allocation_region = lambda base, size: events.append(("set", base, size)) + ops.set_current_alloc_offset = lambda _offset: None + ops.stop_allocation_region = lambda: events.append("stop") + monkeypatch.setitem(sys.modules, "foundry.ops", ops) + + spec = importlib.util.spec_from_file_location( + "foundry.allocation_region_contract", + ALLOCATION_REGION_PATH, + ) + assert spec is not None and spec.loader is not None + module = importlib.util.module_from_spec(spec) + monkeypatch.setitem(sys.modules, spec.name, module) + spec.loader.exec_module(module) + return SimpleNamespace(module=module, events=events) + + +def test_cleanup_false_raises_only_after_region_is_stopped(monkeypatch) -> None: + context = _load_allocation_region(monkeypatch, cleanup_result=False) + + with ( + pytest.raises(RuntimeError, match="Failed to free preallocated region"), + context.module.allocation_region(0x500000000000, 4096, prealloc_size=1024), + ): + context.events.append("body") + + assert context.events == [ + ("set", 0x500000000000, 4096), + ("preallocate", 1024), + "body", + "free", + "stop", + ] + + +def test_cleanup_false_preserves_body_exception_and_warns(monkeypatch) -> None: + context = _load_allocation_region(monkeypatch, cleanup_result=False) + + with ( + pytest.warns(RuntimeWarning, match="Failed to free preallocated region"), + pytest.raises(ValueError, match="body failed") as exc_info, + context.module.allocation_region(0x500000000000, 4096, prealloc_size=1024), + ): + raise ValueError("body failed") + + assert any("Failed to free preallocated region" in note for note in exc_info.value.__notes__) + assert context.events[-2:] == ["free", "stop"] diff --git a/tests/test_cuda_graph_parallel_finish_propagates.py b/tests/test_cuda_graph_parallel_finish_propagates.py new file mode 100644 index 00000000..8d31eec8 --- /dev/null +++ b/tests/test_cuda_graph_parallel_finish_propagates.py @@ -0,0 +1,76 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the Foundry project +"""Source contract: per-entry graph finish must propagate build failures. + +The background template build stores any failure in the ``build_complete_`` +promise via ``set_exception`` (``CUDAGraphParallel.cpp`` bg thread). A +``shared_future::wait()`` blocks until ready but SUPPRESSES that stored +exception, so ``finish_one_graph_load_impl`` would then read a half-built +``pending->entries[index]`` (unlinked graph / missing template) and crash with +an opaque error instead of surfacing the real build failure. + +``shared_future::get()`` rethrows the stored exception (and may be called +repeatedly), so the per-entry finish must consume the future via ``get()`` +BEFORE touching ``pending->entries``. The batched ``finish_graph_loads_impl`` +stays valid: it gets once, then calls finish-one repeatedly (repeated get on a +shared_future is well-defined). + +This is a local, CUDA-free regression guard: building/exercising the native +path needs an H100 (see AGENTS.md), so we assert the fix at the source level. +""" + +from __future__ import annotations + +from pathlib import Path + +import pytest + +SRC = Path(__file__).parents[1] / "csrc" / "CUDAGraphParallel.cpp" + + +def _extract_body(text: str, signature_anchor: str) -> str: + start = text.index(signature_anchor) + brace = text.index("{", start) + depth = 0 + for i in range(brace, len(text)): + char = text[i] + if char == "{": + depth += 1 + elif char == "}": + depth -= 1 + if depth == 0: + return text[brace : i + 1] + raise AssertionError(f"unbalanced braces after {signature_anchor!r}") + + +@pytest.fixture +def finish_one_body() -> str: + return _extract_body( + SRC.read_text(), + "GraphLoadResult finish_one_graph_load_impl(", + ) + + +def test_finish_one_consumes_future_with_get_not_wait(finish_one_body: str) -> None: + assert "build_complete_.get()" in finish_one_body + assert "build_complete_.wait()" not in finish_one_body + + +def test_finish_one_propagates_before_accessing_entries(finish_one_body: str) -> None: + get_pos = finish_one_body.index("build_complete_.get()") + entries_pos = finish_one_body.index("pending->entries") + assert get_pos < entries_pos, ( + "build_complete_.get() must run before any pending->entries access so a " + "failed background build is rethrown, not read as a half-built entry" + ) + + +def test_batched_finish_still_gets_the_future() -> None: + # The batched path may get() once up front then call finish-one repeatedly; + # repeated shared_future::get() is valid. Guard that it still consumes the + # future (so a build failure surfaces even before the per-entry loop). + batched = _extract_body( + SRC.read_text(), + "std::vector finish_graph_loads_impl(", + ) + assert "build_complete_.get()" in batched diff --git a/tests/test_deepep_fabric.py b/tests/test_deepep_fabric.py index be6be723..93ae73b6 100644 --- a/tests/test_deepep_fabric.py +++ b/tests/test_deepep_fabric.py @@ -13,15 +13,22 @@ - Multiple GPUs for distributed testing """ +import ast import os import shutil import subprocess import sys +from collections import Counter from pathlib import Path import pytest -import torch -import torch.distributed as dist + +try: + import torch + import torch.distributed as dist +except ModuleNotFoundError: + torch = None + dist = None # Use lower base address (matching working deepep_args example) # Note: 0x60000000000 (6TB) NOT 0x600000000000 (96TB) - the extra zero matters! @@ -29,6 +36,8 @@ REGION_SIZE_STR = "512GB" ARCHIVE_DIR = "deepep_fabric_archive" HOOK_ARCHIVE_DIR = "hook_archive" +SOURCE_RANK_OFFSET = 128 +LAYOUT_COUNT_MASK = (1 << 32) - 1 def _get_hook_so_path(): @@ -77,20 +86,48 @@ def _init_dist(local_rank: int, num_local_ranks: int): ) +def _maybe_warm_nccl_pre_region(group): + """Diagnostic knob (TEST_NCCL_WARMUP_PRE_REGION=1): force NCCL transport + setup BEFORE the foundry region exists, so NCCL's cuMem allocations + (NCCL_CUMEM_ENABLE=1) pass through untracked. Localizes whether the + NCCL-CUMEM incompatibility is "NCCL cuMem while region enabled" only. + The warmup tensor's caching-allocator segment is drained afterwards so + post-region allocations don't reuse an untracked (non-deterministic) + segment. + """ + if os.environ.get("TEST_NCCL_WARMUP_PRE_REGION", "0") != "1": + return + t = torch.ones(1024, 1024, device="cuda") + dist.all_reduce(t) # default group: ring/p2p transport setup + dist.all_reduce(t, group=group) # subgroup used by the DeepEP Buffer + dist.all_gather_object([None] * dist.get_world_size(), 0) # object-coll path + torch.cuda.synchronize() + del t + torch.cuda.empty_cache() # legal here: region not set yet, no captures + print("[TEST] NCCL pre-region warmup done (transports up before region)", flush=True) + + def _create_deterministic_inputs( rank: int, num_tokens: int, hidden: int, num_experts: int, num_topk: int ): """Create deterministic input tensors with traceable patterns. - Pattern for x: each token i has value (rank * 1000 + i) in all hidden dims + The first hidden-128 values encode source rank and the final 128 encode + token id. Both fields are exactly representable in BF16 for this test. Pattern for topk_idx: deterministic expert selection based on token id """ import deep_ep - # x[token_i, :] = rank * 1000 + token_i (easy to trace which rank/token) - x = torch.zeros((num_tokens, hidden), dtype=torch.bfloat16, device="cuda") - for i in range(num_tokens): - x[i, :] = float(rank * 1000 + i) + assert hidden > 128 + assert num_tokens <= 257 + assert rank - SOURCE_RANK_OFFSET <= 256 + x = torch.full( + (num_tokens, hidden), + float(rank - SOURCE_RANK_OFFSET), + dtype=torch.bfloat16, + device="cuda", + ) + x[:, -128:] = torch.arange(num_tokens, device="cuda").to(torch.bfloat16).view(-1, 1) # Deterministic topk_idx: token i selects experts [i % num_experts, (i+1) % num_experts, ...] topk_idx = torch.zeros((num_tokens, num_topk), dtype=deep_ep.topk_idx_t, device="cuda") @@ -101,6 +138,16 @@ def _create_deterministic_inputs( return x, topk_idx +def _gather_source_topk(topk_idx, num_ranks: int, num_tokens: int, num_topk: int, group): + gathered = torch.empty( + (num_ranks, num_tokens, num_topk), + dtype=topk_idx.dtype, + device=topk_idx.device, + ) + dist.all_gather_into_tensor(gathered, topk_idx, group=group) + return gathered.cpu().tolist() + + def _print_buffer_info(buffer, rank: int, prefix: str): """Print DeepEP buffer memory addresses and info.""" print(f"[Rank {rank}] {prefix}: DeepEP Buffer Info:", flush=True) @@ -146,87 +193,344 @@ def _print_buffer_info(buffer, rank: int, prefix: str): ) -def _verify_dispatch_result( +def _tensor_contract(tensor_contract=None): + if tensor_contract is not None: + return tensor_contract + assert torch is not None, "PyTorch is required for the real tensor contract" + return { + "tensor_type": torch.Tensor, + "bfloat16_dtype": torch.bfloat16, + "int32_dtype": torch.int32, + "int64_dtype": torch.int64, + } + + +def _assert_tensor_contract(value, *, shape, dtype, name: str, tensor_contract=None): + contract = _tensor_contract(tensor_contract) + assert isinstance(value, contract["tensor_type"]), f"{name} is not a tensor" + assert tuple(value.shape) == shape, f"{name} shape {tuple(value.shape)} != {shape}" + assert value.dtype == contract[dtype], f"{name} dtype {value.dtype} != {contract[dtype]}" + + +def _assert_dispatch_tensor_contract( recv_x, - recv_topk_idx, - recv_src_idx, - num_recv, + recv_count, + recv_src_info, + recv_layout_range, + *, + num_local_experts: int, + num_ranks: int, + max_tokens: int, + hidden: int, + tensor_contract=None, +): + capacity = num_ranks * max_tokens + _assert_tensor_contract( + recv_x, + shape=(num_local_experts, capacity, hidden), + dtype="bfloat16_dtype", + name="recv_x", + tensor_contract=tensor_contract, + ) + _assert_tensor_contract( + recv_count, + shape=(num_local_experts,), + dtype="int32_dtype", + name="recv_count", + tensor_contract=tensor_contract, + ) + _assert_tensor_contract( + recv_src_info, + shape=(num_local_experts, capacity), + dtype="int32_dtype", + name="recv_src_info", + tensor_contract=tensor_contract, + ) + _assert_tensor_contract( + recv_layout_range, + shape=(num_local_experts, num_ranks), + dtype="int64_dtype", + name="recv_layout_range", + tensor_contract=tensor_contract, + ) + + +def _unpack_dispatch_result( + result, + *, + num_ranks: int, + max_tokens: int, + hidden: int, + num_experts: int, + tensor_contract=None, +): + assert type(result) is tuple, "low_latency_dispatch must return a tuple" + assert len(result) == 5, f"low_latency_dispatch returned {len(result)} fields instead of 5" + recv_x, recv_count, handle, dispatch_event, recv_hook = result + assert type(handle) is tuple, "dispatch handle must be a tuple" + assert len(handle) == 5, f"dispatch handle has {len(handle)} fields instead of 5" + recv_src_info, recv_layout_range, handle_max_tokens, handle_hidden, handle_experts = handle + assert type(handle_max_tokens) is int and handle_max_tokens == max_tokens + assert type(handle_hidden) is int and handle_hidden == hidden + assert type(handle_experts) is int and handle_experts == num_experts + assert dispatch_event is not None, "missing dispatch event" + assert recv_hook is None, "recv_hook must be None when return_recv_hook=False" + assert num_experts % num_ranks == 0 + _assert_dispatch_tensor_contract( + recv_x, + recv_count, + recv_src_info, + recv_layout_range, + num_local_experts=num_experts // num_ranks, + num_ranks=num_ranks, + max_tokens=max_tokens, + hidden=hidden, + tensor_contract=tensor_contract, + ) + return recv_x, recv_count, recv_src_info, recv_layout_range, dispatch_event + + +def _unpack_saved_dispatch_outputs( + output_tensors, + *, + num_ranks: int, + max_tokens: int, + hidden: int, + num_experts: int, +): + assert isinstance(output_tensors, (list, tuple)), "loaded graph outputs must be a sequence" + assert len(output_tensors) == 4, ( + f"loaded graph restored {len(output_tensors)} tensors instead of 4" + ) + recv_x, recv_count, recv_src_info, recv_layout_range = output_tensors + _assert_dispatch_tensor_contract( + recv_x, + recv_count, + recv_src_info, + recv_layout_range, + num_local_experts=num_experts // num_ranks, + num_ranks=num_ranks, + max_tokens=max_tokens, + hidden=hidden, + ) + return recv_x, recv_count, recv_src_info, recv_layout_range + + +def _expected_route_counters( + source_topk, + *, rank: int, num_ranks: int, num_tokens: int, - hidden: int, num_experts: int, - prefix: str, + num_topk: int, ): - """Verify dispatch results are correct based on our deterministic pattern.""" + assert num_experts % num_ranks == 0 + assert len(source_topk) == num_ranks num_local_experts = num_experts // num_ranks local_expert_start = rank * num_local_experts - local_expert_end = local_expert_start + num_local_experts - - # num_recv might be an EventOverlap object (async), need to sync and get value - actual_num_recv = num_recv - if hasattr(num_recv, "wait"): - # It's an async object, wait for it - num_recv.wait() - if hasattr(num_recv, "value"): - actual_num_recv = num_recv.value - elif hasattr(num_recv, "item"): - actual_num_recv = num_recv.item() - elif not isinstance(num_recv, int): - # Try to get from recv_x shape - actual_num_recv = ( - recv_x.shape[0] * recv_x.shape[1] - if recv_x is not None and len(recv_x.shape) >= 2 - else 0 - ) - print( - f"[Rank {rank}] {prefix}: num_recv is {type(num_recv)}, using recv_x shape to estimate", - flush=True, - ) + expected = [Counter() for _ in range(num_local_experts)] + for source_rank, source_routes in enumerate(source_topk): + assert len(source_routes) == num_tokens + for token, selected_experts in enumerate(source_routes): + assert len(selected_experts) == num_topk + for selected_expert in selected_experts: + assert isinstance(selected_expert, int) and not isinstance(selected_expert, bool) + local_expert = selected_expert - local_expert_start + if 0 <= local_expert < num_local_experts: + expected[local_expert][(source_rank, token)] += 1 + return expected - print(f"[Rank {rank}] {prefix}: Verification - num_recv = {actual_num_recv}", flush=True) - print( - f"[Rank {rank}] {prefix}: Verification - local experts range = [{local_expert_start}, {local_expert_end})", - flush=True, - ) - # recv_x shape is [num_local_experts, max_tokens_per_expert, hidden] - # Print shape info for debugging - if recv_x is not None: - print(f"[Rank {rank}] {prefix}: recv_x shape = {recv_x.shape}", flush=True) +def _payload_row_values(recv_x, expert: int, position: int): + row = recv_x[expert, position] + return row.tolist() if hasattr(row, "tolist") else row - # Check some values in recv_x to verify the pattern - if recv_x is not None and recv_x.numel() > 0: - # Check first expert's first few tokens - num_experts_in_recv = recv_x.shape[0] - tokens_per_expert = recv_x.shape[1] - check_experts = min(2, num_experts_in_recv) - check_tokens = min(3, tokens_per_expert) - print( - f"[Rank {rank}] {prefix}: Checking first {check_experts} experts, {check_tokens} tokens each:", - flush=True, +def _verify_dispatch_tensors( + recv_x, + recv_count, + recv_src_info, + recv_layout_range, + *, + rank: int, + num_ranks: int, + num_tokens: int, + hidden: int, + num_experts: int, + num_topk: int, + source_topk, + prefix: str, + tensor_contract=None, +): + assert hidden > 128 + assert num_experts % num_ranks == 0 + num_local_experts = num_experts // num_ranks + _assert_dispatch_tensor_contract( + recv_x, + recv_count, + recv_src_info, + recv_layout_range, + num_local_experts=num_local_experts, + num_ranks=num_ranks, + max_tokens=num_tokens, + hidden=hidden, + tensor_contract=tensor_contract, + ) + expected_by_expert = _expected_route_counters( + source_topk, + rank=rank, + num_ranks=num_ranks, + num_tokens=num_tokens, + num_experts=num_experts, + num_topk=num_topk, + ) + recv_x_cpu = recv_x.cpu() + counts = recv_count.cpu().tolist() + src_info = recv_src_info.cpu().tolist() + layout_ranges = recv_layout_range.cpu().tolist() + + for local_expert, expected in enumerate(expected_by_expert): + global_expert = rank * num_local_experts + local_expert + actual_count = counts[local_expert] + assert isinstance(actual_count, int) and not isinstance(actual_count, bool) + assert actual_count == sum(expected.values()), ( + f"{prefix}: expert {global_expert} count {actual_count} != " + f"expected {sum(expected.values())}" ) + intervals = [] + actual = Counter() + for source_rank, encoded_range in enumerate(layout_ranges[local_expert]): + assert isinstance(encoded_range, int) and not isinstance(encoded_range, bool) + assert encoded_range >= 0, ( + f"{prefix}: expert {global_expert} source {source_rank} " + f"has negative layout range {encoded_range}" + ) + begin = encoded_range >> 32 + source_count = encoded_range & LAYOUT_COUNT_MASK + end = begin + source_count + expected_source_count = sum( + count + for (expected_source, _token), count in expected.items() + if expected_source == source_rank + ) + assert source_count == expected_source_count, ( + f"{prefix}: expert {global_expert} source {source_rank} count " + f"{source_count} != expected {expected_source_count}" + ) + assert 0 <= begin <= actual_count and begin <= end <= actual_count, ( + f"{prefix}: expert {global_expert} source {source_rank} interval " + f"[{begin}, {end}) outside [0, {actual_count})" + ) + if source_count: + intervals.append((begin, end, source_rank)) + for position in range(begin, end): + token = src_info[local_expert][position] + assert isinstance(token, int) and not isinstance(token, bool) + assert 0 <= token < num_tokens, ( + f"{prefix}: expert {global_expert} position {position} " + f"has invalid source token {token}" + ) + route = (source_rank, token) + actual[route] += 1 + expected_row = [float(source_rank - SOURCE_RANK_OFFSET)] * (hidden - 128) + [ + float(token) + ] * 128 + actual_row = _payload_row_values( + recv_x_cpu, + local_expert, + position, + ) + assert actual_row == expected_row, ( + f"{prefix}: expert {global_expert} position {position} " + f"payload differs for route {route}" + ) - for expert_i in range(check_experts): - expert_global_idx = local_expert_start + expert_i - for token_i in range(check_tokens): - recv_val = recv_x[expert_i, token_i, 0].item() + cursor = 0 + for begin, end, source_rank in sorted(intervals): + assert begin == cursor, ( + f"{prefix}: expert {global_expert} source {source_rank} interval " + f"starts at {begin}, expected {cursor}" + ) + cursor = end + assert cursor == actual_count, ( + f"{prefix}: expert {global_expert} layout covers [0, {cursor}), " + f"expected [0, {actual_count})" + ) + assert actual == expected, ( + f"{prefix}: expert {global_expert} routes {actual} != expected {expected}" + ) - # Decode: src_rank = recv_val // 1000, token_id = recv_val % 1000 - decoded_src_rank = int(recv_val) // 1000 - decoded_token_id = int(recv_val) % 1000 + print(f"[Rank {rank}] {prefix}: verified expert counts {counts}", flush=True) - print( - f"[Rank {rank}] {prefix}: expert[{expert_i}] (global {expert_global_idx}), " - f"token[{token_i}]: value={recv_val:.1f} " - f"(decoded: src_rank={decoded_src_rank}, token_id={decoded_token_id})", - flush=True, - ) - else: - print(f"[Rank {rank}] {prefix}: recv_x is empty or None", flush=True) - return True +def _verify_dispatch_result( + result, + *, + rank: int, + num_ranks: int, + num_tokens: int, + hidden: int, + num_experts: int, + num_topk: int, + source_topk, + prefix: str, + tensor_contract=None, +): + recv_x, recv_count, recv_src_info, recv_layout_range, dispatch_event = _unpack_dispatch_result( + result, + num_ranks=num_ranks, + max_tokens=num_tokens, + hidden=hidden, + num_experts=num_experts, + tensor_contract=tensor_contract, + ) + if hasattr(dispatch_event, "wait"): + dispatch_event.wait() + _verify_dispatch_tensors( + recv_x, + recv_count, + recv_src_info, + recv_layout_range, + rank=rank, + num_ranks=num_ranks, + num_tokens=num_tokens, + hidden=hidden, + num_experts=num_experts, + num_topk=num_topk, + source_topk=source_topk, + prefix=prefix, + tensor_contract=tensor_contract, + ) + return recv_x, recv_count, recv_src_info, recv_layout_range + + +def _sentinel_dispatch_outputs( + recv_x, + recv_count, + recv_src_info, + recv_layout_range, + *, + rank: int, + num_ranks: int, + num_experts: int, + source_topk, +): + num_tokens = len(source_topk[0]) + num_topk = len(source_topk[0][0]) if num_tokens else 0 + expected_by_expert = _expected_route_counters( + source_topk, + rank=rank, + num_ranks=num_ranks, + num_tokens=num_tokens, + num_experts=num_experts, + num_topk=num_topk, + ) + for local_expert, expected in enumerate(expected_by_expert): + recv_x[local_expert, : sum(expected.values())].fill_(float("nan")) + recv_count.fill_(-1) + recv_src_info.fill_(-1) + recv_layout_range.fill_(-1) def _run_save(local_rank: int, num_processes: int): @@ -263,6 +567,8 @@ def _run_save(local_rank: int, num_processes: int): num_experts = num_ranks * 4 # 4 experts per rank num_topk = 2 + _maybe_warm_nccl_pre_region(group) + # Setup allocation region region_size = fdry.parse_size(REGION_SIZE_STR) print(f"[Rank {rank}] SAVE: Setting up allocation region at {hex(BASE_ADDR)}", flush=True) @@ -300,15 +606,17 @@ def _run_save(local_rank: int, num_processes: int): # Create deterministic input tensors x, topk_idx = _create_deterministic_inputs(rank, num_tokens, hidden, num_experts, num_topk) cumulative_stats = torch.zeros((num_local_experts,), dtype=torch.int, device="cuda") + source_topk = _gather_source_topk(topk_idx, num_ranks, num_tokens, num_topk, group) print(f"[Rank {rank}] SAVE: Input tensors:", flush=True) print(f"[Rank {rank}] SAVE: x address = {hex(x.data_ptr())}, shape = {x.shape}", flush=True) print( - f"[Rank {rank}] SAVE: x[0, :5] = {x[0, :5].tolist()} (expect {rank * 1000 + 0})", + f"[Rank {rank}] SAVE: x[0, :5] = {x[0, :5].tolist()} " + f"(expect {rank - SOURCE_RANK_OFFSET})", flush=True, ) print( - f"[Rank {rank}] SAVE: x[1, :5] = {x[1, :5].tolist()} (expect {rank * 1000 + 1})", + f"[Rank {rank}] SAVE: x[1, -5:] = {x[1, -5:].tolist()} (expect 1)", flush=True, ) print(f"[Rank {rank}] SAVE: topk_idx address = {hex(topk_idx.data_ptr())}", flush=True) @@ -330,8 +638,17 @@ def _run_save(local_rank: int, num_processes: int): ) torch.cuda.synchronize() - # Unpack dispatch result - recv_x, recv_topk_idx, recv_src_idx, num_recv, *rest = result + recv_x, recv_count, recv_src_info, recv_layout_range = _verify_dispatch_result( + result, + rank=rank, + num_ranks=num_ranks, + num_tokens=num_tokens, + hidden=hidden, + num_experts=num_experts, + num_topk=num_topk, + source_topk=source_topk, + prefix="SAVE-warmup", + ) print(f"[Rank {rank}] SAVE: Warmup dispatch result:", flush=True) print( f"[Rank {rank}] SAVE: recv_x address = {hex(recv_x.data_ptr()) if recv_x is not None else 'None'}", @@ -341,20 +658,11 @@ def _run_save(local_rank: int, num_processes: int): f"[Rank {rank}] SAVE: recv_x shape = {recv_x.shape if recv_x is not None else 'None'}", flush=True, ) - print(f"[Rank {rank}] SAVE: num_recv = {num_recv}", flush=True) - - # Verify results - _verify_dispatch_result( - recv_x, - recv_topk_idx, - recv_src_idx, - num_recv, - rank, - num_ranks, - num_tokens, - hidden, - num_experts, - "SAVE-warmup", + print(f"[Rank {rank}] SAVE: recv_count = {recv_count}", flush=True) + print(f"[Rank {rank}] SAVE: recv_src_info shape = {recv_src_info.shape}", flush=True) + print( + f"[Rank {rank}] SAVE: recv_layout_range shape = {recv_layout_range.shape}", + flush=True, ) dist.barrier(group) @@ -380,45 +688,62 @@ def _run_save(local_rank: int, num_processes: int): print(f"[Rank {rank}] SAVE: Graph capture completed", flush=True) - # Print graph output tensor addresses - # graph_result structure varies - extract only actual tensors - print(f"[Rank {rank}] SAVE: Graph output result type: {type(graph_result)}", flush=True) - print(f"[Rank {rank}] SAVE: Graph output result length: {len(graph_result)}", flush=True) - - # Collect only tensor objects from the result - output_tensors_to_save = [] - for i, item in enumerate(graph_result): - item_type = type(item).__name__ - if isinstance(item, torch.Tensor): - print( - f"[Rank {rank}] SAVE: result[{i}] = Tensor, address = {hex(item.data_ptr())}, shape = {item.shape}", - flush=True, - ) - output_tensors_to_save.append(item) - elif isinstance(item, tuple): - print(f"[Rank {rank}] SAVE: result[{i}] = tuple of length {len(item)}", flush=True) - # Check if tuple contains tensors - for j, sub_item in enumerate(item): - if isinstance(sub_item, torch.Tensor): - print( - f"[Rank {rank}] SAVE: result[{i}][{j}] = Tensor, address = {hex(sub_item.data_ptr())}, shape = {sub_item.shape}", - flush=True, - ) - output_tensors_to_save.append(sub_item) - else: - print(f"[Rank {rank}] SAVE: result[{i}] = {item_type}", flush=True) + graph_recv_x, graph_recv_count, graph_recv_src_info, graph_recv_layout_range, _event = ( + _unpack_dispatch_result( + graph_result, + num_ranks=num_ranks, + max_tokens=num_tokens, + hidden=hidden, + num_experts=num_experts, + ) + ) + output_tensors_to_save = [ + graph_recv_x, + graph_recv_count, + graph_recv_src_info, + graph_recv_layout_range, + ] + for name, tensor in zip( + ("recv_x", "recv_count", "recv_src_info", "recv_layout_range"), + output_tensors_to_save, + strict=True, + ): + print( + f"[Rank {rank}] SAVE: {name} address = {hex(tensor.data_ptr())}, " + f"shape = {tensor.shape}", + flush=True, + ) # Replay and verify + _sentinel_dispatch_outputs( + graph_recv_x, + graph_recv_count, + graph_recv_src_info, + graph_recv_layout_range, + rank=rank, + num_ranks=num_ranks, + num_experts=num_experts, + source_topk=source_topk, + ) + torch.cuda.synchronize() print(f"[Rank {rank}] SAVE: Replaying graph", flush=True) graph.replay() torch.cuda.synchronize() - # Verify graph replay result - graph_num_recv = cumulative_stats.sum().item() # Approximate + _verify_dispatch_result( + graph_result, + rank=rank, + num_ranks=num_ranks, + num_tokens=num_tokens, + hidden=hidden, + num_experts=num_experts, + num_topk=num_topk, + source_topk=source_topk, + prefix="SAVE-replay", + ) print(f"[Rank {rank}] SAVE: Graph replay completed", flush=True) # Save graph and fatbins to per-rank archive - # output_tensors_to_save was already built above by filtering tensor objects rank_archive = os.path.join(ARCHIVE_DIR, f"rank_{rank}") os.makedirs(rank_archive, exist_ok=True) @@ -487,11 +812,22 @@ def _run_load(local_rank: int, num_processes: int): print(f"[Rank {rank}] LOAD: Loading CUDA modules from {rank_archive}", flush=True) fdry.load_cuda_modules_and_libraries(rank_archive) + _maybe_warm_nccl_pre_region(group) + # Step 2: Setup allocation region region_size = fdry.parse_size(REGION_SIZE_STR) print(f"[Rank {rank}] LOAD: Setting up allocation region at {hex(BASE_ADDR)}", flush=True) fdry.set_allocation_region(BASE_ADDR, region_size) + # Optional (TEST_LOAD_PREALLOC_MB=N): reproduce the production LOAD shape - + # one upfront-preallocated chunk that subsequent allocations (including the + # DeepEP NVL buffer) carve from with NO individual cuMem handle. This + # exercises the hook's whole-chunk VMM-IPC export path. + prealloc_mb = int(os.environ.get("TEST_LOAD_PREALLOC_MB", "0")) + if prealloc_mb > 0: + assert fdry.preallocate_region(prealloc_mb * 1024 * 1024), "preallocate_region failed" + print(f"[Rank {rank}] LOAD: preallocated {prealloc_mb} MB chunk", flush=True) + # Step 3: Create DeepEP buffer (initializes NVSHMEM runtime) num_local_experts = num_experts // num_ranks num_rdma_bytes = deep_ep.Buffer.get_low_latency_rdma_size_hint( @@ -499,13 +835,17 @@ def _run_load(local_rank: int, num_processes: int): ) use_fabric = os.environ.get("TEST_USE_FABRIC", "1") == "1" + # Must match SAVE: an NVL buffer allocated on SAVE but not LOAD (or vice versa) + # diverges the allocation trajectory and the peer-pointer setup. + num_nvl_bytes = int(os.environ.get("TEST_NVL_BYTES_MB", "0")) * 1024 * 1024 print( - f"[Rank {rank}] LOAD: Creating DeepEP buffer with use_fabric={use_fabric}, size={num_rdma_bytes / 1e6:.2f} MB", + f"[Rank {rank}] LOAD: Creating DeepEP buffer with use_fabric={use_fabric}, num_nvl_bytes={num_nvl_bytes / 1e6:.2f} MB, num_rdma_bytes={num_rdma_bytes / 1e6:.2f} MB", flush=True, ) buffer = deep_ep.Buffer( group, + num_nvl_bytes=num_nvl_bytes, num_rdma_bytes=num_rdma_bytes, low_latency_mode=True, num_qps_per_rank=num_local_experts, @@ -526,15 +866,17 @@ def _run_load(local_rank: int, num_processes: int): # Create deterministic input tensors (same as SAVE) x, topk_idx = _create_deterministic_inputs(rank, num_tokens, hidden, num_experts, num_topk) cumulative_stats = torch.zeros((num_local_experts,), dtype=torch.int, device="cuda") + source_topk = _gather_source_topk(topk_idx, num_ranks, num_tokens, num_topk, group) print(f"[Rank {rank}] LOAD: Input tensors:", flush=True) print(f"[Rank {rank}] LOAD: x address = {hex(x.data_ptr())}, shape = {x.shape}", flush=True) print( - f"[Rank {rank}] LOAD: x[0, :5] = {x[0, :5].tolist()} (expect {rank * 1000 + 0})", + f"[Rank {rank}] LOAD: x[0, :5] = {x[0, :5].tolist()} " + f"(expect {rank - SOURCE_RANK_OFFSET})", flush=True, ) print( - f"[Rank {rank}] LOAD: x[1, :5] = {x[1, :5].tolist()} (expect {rank * 1000 + 1})", + f"[Rank {rank}] LOAD: x[1, -5:] = {x[1, -5:].tolist()} (expect 1)", flush=True, ) print(f"[Rank {rank}] LOAD: topk_idx address = {hex(topk_idx.data_ptr())}", flush=True) @@ -556,8 +898,17 @@ def _run_load(local_rank: int, num_processes: int): ) torch.cuda.synchronize() - # Unpack and verify warmup result - recv_x, recv_topk_idx, recv_src_idx, num_recv, *rest = warmup_result + recv_x, recv_count, recv_src_info, recv_layout_range = _verify_dispatch_result( + warmup_result, + rank=rank, + num_ranks=num_ranks, + num_tokens=num_tokens, + hidden=hidden, + num_experts=num_experts, + num_topk=num_topk, + source_topk=source_topk, + prefix="LOAD-warmup", + ) print(f"[Rank {rank}] LOAD: Warmup dispatch result:", flush=True) print( f"[Rank {rank}] LOAD: recv_x address = {hex(recv_x.data_ptr()) if recv_x is not None else 'None'}", @@ -567,20 +918,11 @@ def _run_load(local_rank: int, num_processes: int): f"[Rank {rank}] LOAD: recv_x shape = {recv_x.shape if recv_x is not None else 'None'}", flush=True, ) - print(f"[Rank {rank}] LOAD: num_recv = {num_recv}", flush=True) - - # Verify warmup results - _verify_dispatch_result( - recv_x, - recv_topk_idx, - recv_src_idx, - num_recv, - rank, - num_ranks, - num_tokens, - hidden, - num_experts, - "LOAD-warmup", + print(f"[Rank {rank}] LOAD: recv_count = {recv_count}", flush=True) + print(f"[Rank {rank}] LOAD: recv_src_info shape = {recv_src_info.shape}", flush=True) + print( + f"[Rank {rank}] LOAD: recv_layout_range shape = {recv_layout_range.shape}", + flush=True, ) dist.barrier(group) @@ -593,60 +935,80 @@ def _run_load(local_rank: int, num_processes: int): flush=True, ) - # Load graph from per-rank archive + # Load graph from per-rank archive. + # Default: the production path (start_graph_builds + finish_graph_loads), + # same machinery the vLLM/sglang integrations use. TEST_LOAD_API=single + # selects the legacy CUDAGraph.load path — KNOWN BROKEN for DeepEP graphs: + # it does not merge root-level common_kernel_node_attrs, so factored-out + # cooperative/clusterDim attrs are dropped and the dispatch kernel faults + # with "unspecified launch failure" at replay. graph_json = os.path.join(rank_archive, "deepep_dispatch_graph.json") - print(f"[Rank {rank}] LOAD: Loading graph from {graph_json}", flush=True) - graph, output_tensors = fdry.CUDAGraph.load(graph_json) + load_api = os.environ.get("TEST_LOAD_API", "parallel") + print(f"[Rank {rank}] LOAD: Loading graph from {graph_json} (api={load_api})", flush=True) + if load_api == "single": + graph, output_tensors = fdry.CUDAGraph.load(graph_json) + else: + pending = fdry.CUDAGraph.start_graph_builds([graph_json]) + ((graph, output_tensors),) = fdry.CUDAGraph.finish_graph_loads(pending) print(f"[Rank {rank}] LOAD: Graph loaded successfully", flush=True) - # Print loaded output tensor info - loaded_recv_x = None - if output_tensors is not None: - if isinstance(output_tensors, (list, tuple)): - for i, t in enumerate(output_tensors): - if hasattr(t, "data_ptr"): - print( - f"[Rank {rank}] LOAD: Loaded output[{i}] address = {hex(t.data_ptr())}, shape = {t.shape}", - flush=True, - ) - # First tensor should be recv_x - if i == 0: - loaded_recv_x = t - elif hasattr(output_tensors, "data_ptr"): - print( - f"[Rank {rank}] LOAD: Loaded output address = {hex(output_tensors.data_ptr())}, shape = {output_tensors.shape}", - flush=True, - ) - loaded_recv_x = output_tensors + loaded_recv_x, loaded_recv_count, loaded_recv_src_info, loaded_recv_layout_range = ( + _unpack_saved_dispatch_outputs( + output_tensors, + num_ranks=num_ranks, + max_tokens=num_tokens, + hidden=hidden, + num_experts=num_experts, + ) + ) + for name, tensor in zip( + ("recv_x", "recv_count", "recv_src_info", "recv_layout_range"), + ( + loaded_recv_x, + loaded_recv_count, + loaded_recv_src_info, + loaded_recv_layout_range, + ), + strict=True, + ): + print( + f"[Rank {rank}] LOAD: Loaded {name} address = {hex(tensor.data_ptr())}, " + f"shape = {tensor.shape}", + flush=True, + ) + + _sentinel_dispatch_outputs( + loaded_recv_x, + loaded_recv_count, + loaded_recv_src_info, + loaded_recv_layout_range, + rank=rank, + num_ranks=num_ranks, + num_experts=num_experts, + source_topk=source_topk, + ) + torch.cuda.synchronize() # Replay loaded graph print(f"[Rank {rank}] LOAD: Replaying loaded graph", flush=True) graph.replay() torch.cuda.synchronize() - # Verify after replay - check the LOADED output tensors (not warmup recv_x) - # recv_x shape is [num_local_experts, max_tokens_per_expert, hidden] print(f"[Rank {rank}] LOAD: After graph replay:", flush=True) - verify_tensor = loaded_recv_x if loaded_recv_x is not None else recv_x - if verify_tensor is not None and verify_tensor.numel() > 0: - print(f"[Rank {rank}] LOAD: verify_tensor shape = {verify_tensor.shape}", flush=True) - num_local_experts = num_experts // num_ranks - local_expert_start = rank * num_local_experts - - # Check first expert's first few tokens - check_experts = min(2, verify_tensor.shape[0]) - check_tokens = min(3, verify_tensor.shape[1]) - - for expert_i in range(check_experts): - for token_i in range(check_tokens): - val = verify_tensor[expert_i, token_i, 0].item() - decoded_src_rank = int(val) // 1000 - decoded_token_id = int(val) % 1000 - print( - f"[Rank {rank}] LOAD: expert[{expert_i}], token[{token_i}]: " - f"value={val:.1f} (rank={decoded_src_rank}, token={decoded_token_id})", - flush=True, - ) + _verify_dispatch_tensors( + loaded_recv_x, + loaded_recv_count, + loaded_recv_src_info, + loaded_recv_layout_range, + rank=rank, + num_ranks=num_ranks, + num_tokens=num_tokens, + hidden=hidden, + num_experts=num_experts, + num_topk=num_topk, + source_topk=source_topk, + prefix="LOAD-replay", + ) print(f"[Rank {rank}] LOAD: Graph replay successful", flush=True) @@ -706,7 +1068,383 @@ def _check_deepep_available(): def _check_multi_gpu(): """Check if multiple GPUs are available.""" - return torch.cuda.device_count() >= 2 + return torch is not None and torch.cuda.device_count() >= 2 + + +class _FakeDispatchTensor: + def __init__(self, values, shape, dtype): + self.values = values + self.shape = shape + self.dtype = dtype + + def cpu(self): + return self + + def tolist(self): + return self.values + + def __getitem__(self, key): + value = self.values + for part in key if isinstance(key, tuple) else (key,): + value = value[part] + return value + + +_FAKE_TENSOR_CONTRACT = { + "tensor_type": _FakeDispatchTensor, + "bfloat16_dtype": "bfloat16", + "int32_dtype": "int32", + "int64_dtype": "int64", +} + + +def _fake_source_topk( + num_ranks: int = 2, + num_tokens: int = 8, + num_experts: int = 8, + num_topk: int = 2, +): + return [ + [ + [(token + topk_slot) % num_experts for topk_slot in range(num_topk)] + for token in range(num_tokens) + ] + for _ in range(num_ranks) + ] + + +def _fake_payload_row(source_rank: int, token: int, hidden: int): + return [float(source_rank - 128)] * (hidden - 128) + [float(token)] * 128 + + +def _fake_dispatch_result( + *, + rank: int = 1, + num_ranks: int = 2, + num_tokens: int = 8, + hidden: int = 256, + num_experts: int = 8, + num_topk: int = 2, + source_order=(0, 1), + reverse_tokens: bool = False, + source_topk=None, +): + source_topk = source_topk or _fake_source_topk(num_ranks, num_tokens, num_experts, num_topk) + num_local_experts = num_experts // num_ranks + max_tokens = num_tokens + capacity = num_ranks * max_tokens + recv_x_values = [ + [[float("nan")] * hidden for _ in range(capacity)] for _ in range(num_local_experts) + ] + recv_count_values = [0] * num_local_experts + recv_src_info_values = [[-99] * capacity for _ in range(num_local_experts)] + recv_layout_range_values = [[0] * num_ranks for _ in range(num_local_experts)] + + for local_expert in range(num_local_experts): + global_expert = rank * num_local_experts + local_expert + packed_routes = [] + for source_rank in source_order: + source_tokens = [ + token + for token, selected_experts in enumerate(source_topk[source_rank]) + for selected_expert in selected_experts + if selected_expert == global_expert + ] + if reverse_tokens: + source_tokens.reverse() + begin = len(packed_routes) + recv_layout_range_values[local_expert][source_rank] = (begin << 32) | len(source_tokens) + packed_routes.extend((source_rank, token) for token in source_tokens) + + recv_count_values[local_expert] = len(packed_routes) + for position, (source_rank, token) in enumerate(packed_routes): + recv_x_values[local_expert][position] = _fake_payload_row(source_rank, token, hidden) + recv_src_info_values[local_expert][position] = token + + recv_x = _FakeDispatchTensor( + recv_x_values, + (num_local_experts, capacity, hidden), + "bfloat16", + ) + recv_count = _FakeDispatchTensor( + recv_count_values, + (num_local_experts,), + "int32", + ) + recv_src_info = _FakeDispatchTensor( + recv_src_info_values, + (num_local_experts, capacity), + "int32", + ) + recv_layout_range = _FakeDispatchTensor( + recv_layout_range_values, + (num_local_experts, num_ranks), + "int64", + ) + handle = ( + recv_src_info, + recv_layout_range, + max_tokens, + hidden, + num_experts, + ) + return (recv_x, recv_count, handle, object(), None), source_topk + + +def _verify_fake_dispatch(result, source_topk, **overrides): + parameters = { + "rank": 1, + "num_ranks": 2, + "num_tokens": 8, + "hidden": 256, + "num_experts": 8, + "num_topk": 2, + "source_topk": source_topk, + "prefix": "fake", + "tensor_contract": _FAKE_TENSOR_CONTRACT, + } + parameters.update(overrides) + _verify_dispatch_result(result, **parameters) + + +def test_oracle_accepts_valid_rank1_global_expert6_payload(): + result, source_topk = _fake_dispatch_result() + + _verify_fake_dispatch(result, source_topk) + + +def test_oracle_accepts_permuted_valid_tokens(): + result, source_topk = _fake_dispatch_result(reverse_tokens=True) + + _verify_fake_dispatch(result, source_topk) + + +def test_oracle_accepts_permuted_source_segments(): + result, source_topk = _fake_dispatch_result(source_order=(1, 0)) + + _verify_fake_dispatch(result, source_topk) + + +def test_oracle_accepts_exact_zero_counts(): + source_topk = [[[0, 0] for _ in range(8)] for _ in range(2)] + result, source_topk = _fake_dispatch_result(source_topk=source_topk) + + _verify_fake_dispatch(result, source_topk) + + +def test_oracle_rejects_bad_fourth_token(): + result, source_topk = _fake_dispatch_result() + recv_x = result[0] + recv_x.values[0][3][-1] = 99.0 + + with pytest.raises(AssertionError): + _verify_fake_dispatch(result, source_topk) + + +def test_oracle_rejects_missing_route_with_plausible_count(): + result, source_topk = _fake_dispatch_result() + recv_count = result[1] + recv_layout_range = result[2][1] + recv_count.values[0] -= 1 + source_one_range = recv_layout_range.values[0][1] + recv_layout_range.values[0][1] = (source_one_range & ~0xFFFFFFFF) | ( + (source_one_range & 0xFFFFFFFF) - 1 + ) + + with pytest.raises(AssertionError): + _verify_fake_dispatch(result, source_topk) + + +@pytest.mark.parametrize( + "malformation", + ( + "result_not_tuple", + "missing_result_field", + "extra_result_field", + "handle_not_tuple", + "missing_handle_field", + "extra_handle_field", + "wrong_max_tokens", + "wrong_hidden", + "wrong_num_experts", + "recv_hook", + "wrong_src_dtype", + "wrong_layout_shape", + ), +) +def test_oracle_requires_exact_low_latency_tuple_and_handle(malformation): + result, source_topk = _fake_dispatch_result() + recv_x, recv_count, handle, event, recv_hook = result + recv_src_info, recv_layout_range, max_tokens, hidden, num_experts = handle + + if malformation == "result_not_tuple": + malformed = list(result) + elif malformation == "missing_result_field": + malformed = result[:-1] + elif malformation == "extra_result_field": + malformed = (*result, None) + elif malformation == "handle_not_tuple": + malformed = (recv_x, recv_count, list(handle), event, recv_hook) + elif malformation == "missing_handle_field": + malformed = (recv_x, recv_count, handle[:-1], event, recv_hook) + elif malformation == "extra_handle_field": + malformed = (recv_x, recv_count, (*handle, 0), event, recv_hook) + elif malformation == "wrong_max_tokens": + malformed = ( + recv_x, + recv_count, + (recv_src_info, recv_layout_range, max_tokens + 1, hidden, num_experts), + event, + recv_hook, + ) + elif malformation == "wrong_hidden": + malformed = ( + recv_x, + recv_count, + (recv_src_info, recv_layout_range, max_tokens, hidden + 1, num_experts), + event, + recv_hook, + ) + elif malformation == "wrong_num_experts": + malformed = ( + recv_x, + recv_count, + (recv_src_info, recv_layout_range, max_tokens, hidden, num_experts + 1), + event, + recv_hook, + ) + elif malformation == "recv_hook": + malformed = (recv_x, recv_count, handle, event, object()) + elif malformation == "wrong_src_dtype": + recv_src_info.dtype = "int64" + malformed = result + elif malformation == "wrong_layout_shape": + recv_layout_range.shape = (4, 3) + malformed = result + else: + raise AssertionError(f"unhandled malformation {malformation}") + + with pytest.raises(AssertionError): + _verify_fake_dispatch(malformed, source_topk) + + +def test_replay_oracle_rejects_unwritten_metadata_and_payload(): + result, source_topk = _fake_dispatch_result() + recv_x, recv_count, handle, _event, _recv_hook = result + recv_src_info, recv_layout_range, *_scalars = handle + for expert, count in enumerate(recv_count.values): + for position in range(count): + recv_x.values[expert][position] = [float("nan")] * recv_x.shape[2] + recv_count.values[:] = [-1] * len(recv_count.values) + for row in recv_src_info.values: + row[:] = [-1] * len(row) + for row in recv_layout_range.values: + row[:] = [-1] * len(row) + + with pytest.raises(AssertionError): + _verify_fake_dispatch(result, source_topk) + + +@pytest.mark.skipif(torch is None, reason="PyTorch not installed") +def test_torch_cpu_preflight_covers_split_payload_and_dtypes(): + scalar_values = torch.arange(1000, 1009, dtype=torch.float32).to(torch.bfloat16) + assert len(set(scalar_values.tolist())) < scalar_values.numel() + assert scalar_values[1].item() == 1000.0 + + fake_result, source_topk = _fake_dispatch_result( + source_order=(1, 0), + reverse_tokens=True, + ) + fake_recv_x, fake_recv_count, fake_handle, event, recv_hook = fake_result + fake_recv_src_info, fake_recv_layout_range, *handle_scalars = fake_handle + recv_x = torch.tensor(fake_recv_x.values, dtype=torch.bfloat16, device="cpu") + recv_count = torch.tensor(fake_recv_count.values, dtype=torch.int32, device="cpu") + recv_src_info = torch.tensor( + fake_recv_src_info.values, + dtype=torch.int32, + device="cpu", + ) + recv_layout_range = torch.tensor( + fake_recv_layout_range.values, + dtype=torch.int64, + device="cpu", + ) + result = ( + recv_x, + recv_count, + (recv_src_info, recv_layout_range, *handle_scalars), + event, + recv_hook, + ) + + _verify_dispatch_result( + result, + rank=1, + num_ranks=2, + num_tokens=8, + hidden=256, + num_experts=8, + num_topk=2, + source_topk=source_topk, + prefix="torch-cpu", + ) + _sentinel_dispatch_outputs( + recv_x, + recv_count, + recv_src_info, + recv_layout_range, + rank=1, + num_ranks=2, + num_experts=8, + source_topk=source_topk, + ) + assert all( + torch.isnan(recv_x[expert, :count]).all() + for expert, count in enumerate(fake_recv_count.values) + ) + assert torch.equal(recv_count, torch.full_like(recv_count, -1)) + assert torch.equal(recv_src_info, torch.full_like(recv_src_info, -1)) + assert torch.equal(recv_layout_range, torch.full_like(recv_layout_range, -1)) + with pytest.raises(AssertionError): + _verify_dispatch_result( + result, + rank=1, + num_ranks=2, + num_tokens=8, + hidden=256, + num_experts=8, + num_topk=2, + source_topk=source_topk, + prefix="torch-cpu-noop-replay", + ) + + +def test_save_and_load_replay_use_exact_four_output_metadata_sentinels(): + source = Path(__file__).read_text() + tree = ast.parse(source) + save_node = next( + node for node in tree.body if isinstance(node, ast.FunctionDef) and node.name == "_run_save" + ) + load_node = next( + node for node in tree.body if isinstance(node, ast.FunctionDef) and node.name == "_run_load" + ) + save_source = ast.get_source_segment(source, save_node) + load_source = ast.get_source_segment(source, load_node) + assert save_source is not None + assert load_source is not None + + output_names = ( + "recv_x", + "recv_count", + "recv_src_info", + "recv_layout_range", + ) + assert all(name in save_source for name in output_names) + assert all(name in load_source for name in output_names) + assert save_source.count("_sentinel_dispatch_outputs(") == 1 + assert load_source.count("_sentinel_dispatch_outputs(") == 1 + assert "_assert_valid_dispatch_equal" not in save_source + assert "_assert_valid_dispatch_equal" not in load_source @pytest.mark.skipif(not _check_deepep_available(), reason="DeepEP not installed") diff --git a/tests/test_deepep_transport.py b/tests/test_deepep_transport.py new file mode 100644 index 00000000..666d6f63 --- /dev/null +++ b/tests/test_deepep_transport.py @@ -0,0 +1,101 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the Foundry project +from __future__ import annotations + +import importlib.util +from pathlib import Path + +import pytest + +POLICY_PATH = ( + Path(__file__).parents[1] / "python" / "foundry" / "integration" / "deepep_transport.py" +) +POLICY_SPEC = importlib.util.spec_from_file_location("deepep_transport_under_test", POLICY_PATH) +assert POLICY_SPEC is not None and POLICY_SPEC.loader is not None +POLICY = importlib.util.module_from_spec(POLICY_SPEC) +POLICY_SPEC.loader.exec_module(POLICY) + +DeepEPTransport = POLICY.DeepEPTransport +configure_deepep_buffer = POLICY.configure_deepep_buffer +parse_deepep_transport = POLICY.parse_deepep_transport +validate_num_nvl_bytes_match = POLICY.validate_num_nvl_bytes_match +validate_transport_match = POLICY.validate_transport_match + + +@pytest.mark.parametrize( + ("raw", "expected"), + [ + (None, DeepEPTransport.FABRIC), + ("fabric", DeepEPTransport.FABRIC), + ("nvl_ipc", DeepEPTransport.NVL_IPC), + ], +) +def test_parse_deepep_transport(raw, expected): + assert parse_deepep_transport(raw) is expected + + +def test_parse_deepep_transport_rejects_unknown_value(): + with pytest.raises(ValueError, match="deepep_transport"): + parse_deepep_transport("rdma") + + +@pytest.mark.parametrize( + ("transport", "expected_bytes", "expected_fabric"), + [ + (DeepEPTransport.FABRIC, 0, True), + (DeepEPTransport.NVL_IPC, 128, False), + ], +) +def test_configure_deepep_buffer(transport, expected_bytes, expected_fabric): + num_nvl_bytes, kwargs = configure_deepep_buffer( + transport, + num_nvl_bytes=128, + kwargs={"other": "kept"}, + ) + + assert num_nvl_bytes == expected_bytes + assert kwargs == {"other": "kept", "use_fabric": expected_fabric} + + +@pytest.mark.parametrize("num_nvl_bytes", [0, -1, True, False, 1.0, "128", None]) +def test_nvl_ipc_rejects_nonpositive_or_non_exact_int_bytes(num_nvl_bytes): + with pytest.raises(RuntimeError, match="num_nvl_bytes.*positive integer"): + configure_deepep_buffer( + DeepEPTransport.NVL_IPC, + num_nvl_bytes=num_nvl_bytes, + kwargs={}, + ) + + +@pytest.mark.parametrize("num_nvl_bytes", [-1, True, False, 1.0, "128", None]) +def test_fabric_validates_upstream_bytes_before_forcing_zero(num_nvl_bytes): + with pytest.raises(RuntimeError, match="num_nvl_bytes.*non-negative integer"): + configure_deepep_buffer( + DeepEPTransport.FABRIC, + num_nvl_bytes=num_nvl_bytes, + kwargs={}, + ) + + +def test_transport_mismatch_names_both_values(): + with pytest.raises(RuntimeError, match="archived=nvl_ipc.*configured=fabric"): + validate_transport_match( + archived="nvl_ipc", + configured=DeepEPTransport.FABRIC, + ) + + +def test_legacy_transport_defaults_to_fabric_during_validation(): + validate_transport_match( + archived=None, + configured=DeepEPTransport.FABRIC, + ) + + +def test_nvl_bytes_mismatch_names_both_values(): + with pytest.raises(RuntimeError, match="archived=64.*effective=128"): + validate_num_nvl_bytes_match( + archived=64, + effective=128, + transport=DeepEPTransport.NVL_IPC, + ) diff --git a/tests/test_deepep_transport_vllm.py b/tests/test_deepep_transport_vllm.py new file mode 100644 index 00000000..68575c18 --- /dev/null +++ b/tests/test_deepep_transport_vllm.py @@ -0,0 +1,926 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the Foundry project +from __future__ import annotations + +import importlib.machinery +import importlib.util +import json +import logging +import os +import sys +import threading +import types +from dataclasses import asdict +from pathlib import Path + +import pytest + +REPO_ROOT = Path(__file__).parents[1] +POLICY_PATH = REPO_ROOT / "python" / "foundry" / "integration" / "deepep_transport.py" +VLLM_CONFIG_PATH = REPO_ROOT / "python" / "foundry" / "integration" / "vllm" / "config.py" +VLLM_RESTORE_PROGRESS_PATH = ( + REPO_ROOT / "python" / "foundry" / "integration" / "vllm" / "restore_progress.py" +) +VLLM_RUNTIME_PATH = REPO_ROOT / "python" / "foundry" / "integration" / "vllm" / "runtime.py" +VLLM_HOOKS_PATH = REPO_ROOT / "python" / "foundry" / "integration" / "vllm" / "hooks.py" + +POLICY_SPEC = importlib.util.spec_from_file_location( + "deepep_transport_vllm_under_test", POLICY_PATH +) +assert POLICY_SPEC is not None and POLICY_SPEC.loader is not None +POLICY = importlib.util.module_from_spec(POLICY_SPEC) +POLICY_SPEC.loader.exec_module(POLICY) +DeepEPTransport = POLICY.DeepEPTransport + + +def _install_module(monkeypatch, name: str) -> types.ModuleType: + module = types.ModuleType(name) + module.__spec__ = importlib.machinery.ModuleSpec(name, loader=None) + monkeypatch.setitem(sys.modules, name, module) + if "." in name: + parent_name, child_name = name.rsplit(".", 1) + parent = sys.modules.get(parent_name) + if parent is not None: + setattr(parent, child_name, module) + return module + + +def _load_module(monkeypatch, name: str, path: Path) -> types.ModuleType: + spec = importlib.util.spec_from_file_location(name, path) + assert spec is not None and spec.loader is not None + module = importlib.util.module_from_spec(spec) + monkeypatch.setitem(sys.modules, name, module) + spec.loader.exec_module(module) + parent_name, child_name = name.rsplit(".", 1) + setattr(sys.modules[parent_name], child_name, module) + return module + + +@pytest.fixture +def vllm_modules(monkeypatch): + torch = _install_module(monkeypatch, "torch") + torch.cuda = types.SimpleNamespace() + torch.version = types.SimpleNamespace(cuda=None) + torch.nn = types.SimpleNamespace(Module=object) + + vllm = _install_module(monkeypatch, "vllm") + vllm.__path__ = [] + vllm.__version__ = "test" + logger_module = _install_module(monkeypatch, "vllm.logger") + logger_module.init_logger = logging.getLogger + + foundry = _install_module(monkeypatch, "foundry") + foundry.__path__ = [] + ops = _install_module(monkeypatch, "foundry.ops") + ops.__spec__.origin = str(REPO_ROOT / "python" / "foundry" / "ops.so") + ops.set_skip_fatbin_processing = lambda _value: None + ops.load_cuda_modules_and_libraries = lambda _path: None + ops.set_allocation_region = lambda _base, _size: None + allocation_region = _install_module(monkeypatch, "foundry.allocation_region") + allocation_region.parse_size = lambda _value: 1024 + integration = _install_module(monkeypatch, "foundry.integration") + integration.__path__ = [] + monkeypatch.setitem(sys.modules, "foundry.integration.deepep_transport", POLICY) + integration.deepep_transport = POLICY + vllm_integration = _install_module(monkeypatch, "foundry.integration.vllm") + vllm_integration.__path__ = [] + + config = _load_module( + monkeypatch, + "foundry.integration.vllm.config", + VLLM_CONFIG_PATH, + ) + _load_module( + monkeypatch, + "foundry.integration.vllm.restore_progress", + VLLM_RESTORE_PROGRESS_PATH, + ) + runtime = _load_module( + monkeypatch, + "foundry.integration.vllm.runtime", + VLLM_RUNTIME_PATH, + ) + hooks = _load_module( + monkeypatch, + "foundry.integration.vllm.hooks", + VLLM_HOOKS_PATH, + ) + return types.SimpleNamespace( + config=config, + hooks=hooks, + ops=ops, + runtime=runtime, + ) + + +def _install_fake_all2all_module(monkeypatch, manager_cls): + distributed_module = _install_module(monkeypatch, "vllm.distributed") + distributed_module.__path__ = [] + communicators_module = _install_module( + monkeypatch, + "vllm.distributed.device_communicators", + ) + communicators_module.__path__ = [] + all2all_module = _install_module( + monkeypatch, + "vllm.distributed.device_communicators.all2all", + ) + all2all_module.DeepEPLLAll2AllManager = manager_cls + + +def _install_fake_kv_cache_modules(monkeypatch, engine_core_cls, scheduler_config): + v1_module = _install_module(monkeypatch, "vllm.v1") + v1_module.__path__ = [] + engine_module = _install_module(monkeypatch, "vllm.v1.engine") + engine_module.__path__ = [] + core_module = _install_module(monkeypatch, "vllm.v1.engine.core") + core_module.EngineCore = engine_core_cls + _install_module(monkeypatch, "vllm.v1.core") + core_utils_module = _install_module(monkeypatch, "vllm.v1.core.kv_cache_utils") + core_utils_module.get_kv_cache_configs = lambda _config, specs, available_memory: [ + types.SimpleNamespace( + specs=specs, + available_memory=available_memory, + ) + ] + core_utils_module.generate_scheduler_kv_cache_config = lambda _configs: scheduler_config + envs_module = _install_module(monkeypatch, "vllm.envs") + envs_module.VLLM_ELASTIC_EP_SCALE_UP_LAUNCH = False + + +@pytest.mark.parametrize( + ("contents", "expected"), + [ + ("", DeepEPTransport.FABRIC), + ('deepep_transport = "nvl_ipc"\n', DeepEPTransport.NVL_IPC), + ], +) +def test_vllm_config_parses_deepep_transport(tmp_path, vllm_modules, contents, expected): + config_path = tmp_path / "config.toml" + config_path.write_text(contents) + + config = vllm_modules.config.CUDAGraphExtensionConfig.from_toml(config_path) + + assert config.deepep_transport is expected + + +def test_vllm_config_rejects_unknown_deepep_transport(tmp_path, vllm_modules): + config_path = tmp_path / "config.toml" + config_path.write_text('deepep_transport = "rdma"\n') + + with pytest.raises(ValueError, match="deepep_transport"): + vllm_modules.config.CUDAGraphExtensionConfig.from_toml(config_path) + + +def test_vllm_get_deepep_transport_uses_loaded_config(tmp_path, vllm_modules): + config_path = tmp_path / "config.toml" + config_path.write_text('deepep_transport = "nvl_ipc"\n') + + assert vllm_modules.config.get_deepep_transport() is DeepEPTransport.FABRIC + vllm_modules.config.load_graph_extension_config(str(config_path)) + assert vllm_modules.config.get_deepep_transport() is DeepEPTransport.NVL_IPC + + +def test_vllm_config_prefers_versioned_nvshmem_wheel_library( + tmp_path, + monkeypatch, + vllm_modules, +): + package_root = tmp_path / "nvidia" / "nvshmem" + lib_dir = package_root / "lib" + lib_dir.mkdir(parents=True) + versioned = lib_dir / "libnvshmem_host.so.3" + unversioned = lib_dir / "libnvshmem_host.so" + versioned.touch() + unversioned.touch() + real_find_spec = vllm_modules.config.importlib.util.find_spec + + def fake_find_spec(name): + if name == "nvidia.nvshmem": + return types.SimpleNamespace( + origin=None, + submodule_search_locations=[str(package_root)], + ) + return real_find_spec(name) + + monkeypatch.setattr(vllm_modules.config.importlib.util, "find_spec", fake_find_spec) + config_path = tmp_path / "config.toml" + config_path.write_text("") + + config = vllm_modules.config.CUDAGraphExtensionConfig.from_toml(config_path) + + assert config.nvshmem_host_path == str(versioned) + + +def test_vllm_config_falls_back_to_unversioned_nvshmem_wheel_library( + tmp_path, + monkeypatch, + vllm_modules, +): + package_root = tmp_path / "nvidia" / "nvshmem" + lib_dir = package_root / "lib" + lib_dir.mkdir(parents=True) + unversioned = lib_dir / "libnvshmem_host.so" + unversioned.touch() + real_find_spec = vllm_modules.config.importlib.util.find_spec + + def fake_find_spec(name): + if name == "nvidia.nvshmem": + return types.SimpleNamespace( + origin=None, + submodule_search_locations=[str(package_root)], + ) + return real_find_spec(name) + + monkeypatch.setattr(vllm_modules.config.importlib.util, "find_spec", fake_find_spec) + config_path = tmp_path / "config.toml" + config_path.write_text("") + + config = vllm_modules.config.CUDAGraphExtensionConfig.from_toml(config_path) + + assert config.nvshmem_host_path == str(unversioned) + + +def test_vllm_explicit_nvshmem_path_bypasses_wheel_detection( + tmp_path, + monkeypatch, + vllm_modules, +): + explicit = tmp_path / "custom" / "libnvshmem_host.so" + config_path = tmp_path / "config.toml" + config_path.write_text(f'nvshmem_host_path = "{explicit}"\n') + monkeypatch.setattr( + vllm_modules.config.CUDAGraphExtensionConfig, + "_detect_nvshmem_host_path", + staticmethod(lambda: pytest.fail("explicit path must bypass wheel detection")), + raising=False, + ) + + config = vllm_modules.config.CUDAGraphExtensionConfig.from_toml(config_path) + + assert config.nvshmem_host_path == str(explicit) + + +def test_vllm_runtime_preloads_nvshmem_before_hook(monkeypatch, vllm_modules): + monkeypatch.setattr(vllm_modules.runtime, "get_nvshmem_host_path", lambda: "/wheel/nvshmem.so") + monkeypatch.setattr(vllm_modules.runtime, "get_hook_library_path", lambda: "/foundry/hook.so") + # Seed a tracked value so pytest restores the process environment even when + # LD_PRELOAD was initially absent and setup_ld_preload_env writes it directly. + monkeypatch.setenv("LD_PRELOAD", "") + + vllm_modules.runtime.setup_ld_preload_env() + + assert os.environ["LD_PRELOAD"].split(":")[:2] == [ + "/wheel/nvshmem.so", + "/foundry/hook.so", + ] + + +def test_vllm_warmup_state_round_trips_deepep_profile(tmp_path, vllm_modules): + state = vllm_modules.runtime.WarmupState( + deepep_transport="nvl_ipc", + deepep_num_nvl_bytes=64 * 1024 * 1024, + ) + + vllm_modules.runtime.save_warmup_state(str(tmp_path), state) + + contents = json.loads((tmp_path / "warmup_state.json").read_text()) + assert contents["deepep_transport"] == "nvl_ipc" + assert contents["deepep_num_nvl_bytes"] == 64 * 1024 * 1024 + loaded = vllm_modules.runtime.load_warmup_state(str(tmp_path)) + assert loaded.deepep_transport == "nvl_ipc" + assert loaded.deepep_num_nvl_bytes == 64 * 1024 * 1024 + + +def test_vllm_legacy_warmup_state_uses_fabric_defaults(tmp_path, vllm_modules): + (tmp_path / "warmup_state.json").write_text('{"vllm_version": "legacy"}') + + loaded = vllm_modules.runtime.load_warmup_state(str(tmp_path)) + + assert loaded.deepep_transport == "fabric" + assert loaded.deepep_num_nvl_bytes == 0 + + +@pytest.mark.parametrize( + ("profile", "field"), + [ + ({"deepep_transport": None}, "deepep_transport"), + ({"deepep_transport": 1}, "deepep_transport"), + ({"deepep_transport": "rdma"}, "deepep_transport"), + ({"deepep_num_nvl_bytes": True}, "deepep_num_nvl_bytes"), + ({"deepep_num_nvl_bytes": "64"}, "deepep_num_nvl_bytes"), + ({"deepep_num_nvl_bytes": -1}, "deepep_num_nvl_bytes"), + ( + { + "deepep_transport": "nvl_ipc", + "deepep_num_nvl_bytes": 0, + }, + "deepep_num_nvl_bytes", + ), + ({"deepep_transport": "nvl_ipc"}, "deepep_num_nvl_bytes"), + ], +) +def test_vllm_warmup_state_rejects_invalid_deepep_fields( + tmp_path, + vllm_modules, + profile, + field, +): + path = tmp_path / "warmup_state.json" + path.write_text(json.dumps(profile)) + + with pytest.raises(RuntimeError, match=field) as exc_info: + vllm_modules.runtime.load_warmup_state(str(tmp_path)) + + assert str(path) in str(exc_info.value) + + +def test_vllm_atomic_state_write_hides_partial_temp_content( + tmp_path, + monkeypatch, + vllm_modules, +): + runtime = vllm_modules.runtime + old_state = runtime.WarmupState(vllm_version="old") + new_state = runtime.WarmupState(vllm_version="new") + runtime.save_warmup_state(str(tmp_path), old_state) + shared_path = tmp_path / "warmup_state.json" + write_started = threading.Event() + finish_write = threading.Event() + errors = [] + real_dump = runtime.json.dump + + def paused_dump(value, stream, *args, **kwargs): + stream.write("{") + stream.flush() + write_started.set() + if not finish_write.wait(timeout=2): + raise TimeoutError("reader did not release paused writer") + stream.seek(0) + stream.truncate() + return real_dump(value, stream, *args, **kwargs) + + monkeypatch.setattr(runtime.json, "dump", paused_dump) + + def write_state(): + try: + runtime.save_warmup_state(str(tmp_path), new_state) + except Exception as exc: + errors.append(exc) + + writer = threading.Thread(target=write_state) + writer.start() + assert write_started.wait(timeout=2) + try: + observed = json.loads(shared_path.read_text()) + finally: + finish_write.set() + writer.join(timeout=2) + + assert not writer.is_alive() + assert errors == [] + assert observed == asdict(old_state) + assert json.loads(shared_path.read_text()) == asdict(new_state) + + +def test_vllm_load_rejects_malformed_or_non_object_state(tmp_path, vllm_modules): + path = tmp_path / "warmup_state.json" + path.write_text('{"vllm_version":') + with pytest.raises(RuntimeError, match="invalid JSON") as malformed: + vllm_modules.runtime.load_warmup_state(str(tmp_path)) + assert str(path) in str(malformed.value) + + path.write_text("[]") + with pytest.raises(RuntimeError, match="incompatible warmup state") as non_object: + vllm_modules.runtime.load_warmup_state(str(tmp_path)) + assert str(path) in str(non_object.value) + + +def test_vllm_profile_bootstrap_and_update_are_rank_zero_only( + tmp_path, + vllm_modules, +): + runtime = vllm_modules.runtime + runtime._state = runtime.CUDAGraphExtensionState(rank=1) + runtime.update_warmup_state_deepep_profile( + str(tmp_path), + DeepEPTransport.NVL_IPC, + 64 * 1024 * 1024, + ) + assert not (tmp_path / "warmup_state.json").exists() + + runtime._state.rank = 0 + runtime.update_warmup_state_deepep_profile( + str(tmp_path), + DeepEPTransport.NVL_IPC, + 64 * 1024 * 1024, + ) + + saved = runtime.load_warmup_state(str(tmp_path)) + assert saved.deepep_transport == "nvl_ipc" + assert saved.deepep_num_nvl_bytes == 64 * 1024 * 1024 + + +@pytest.mark.parametrize("num_nvl_bytes", [0, -1, True, 1.0, "67108864"]) +def test_vllm_profile_update_rejects_invalid_nvl_bytes_before_archive_write( + tmp_path, + vllm_modules, + num_nvl_bytes, +): + runtime = vllm_modules.runtime + runtime._state = runtime.CUDAGraphExtensionState(rank=0) + + with pytest.raises(RuntimeError, match="num_nvl_bytes.*positive integer"): + runtime.update_warmup_state_deepep_profile( + str(tmp_path), + DeepEPTransport.NVL_IPC, + num_nvl_bytes, + ) + + assert not (tmp_path / "warmup_state.json").exists() + + +def test_vllm_profile_update_preserves_existing_metadata(tmp_path, vllm_modules): + runtime = vllm_modules.runtime + original = runtime.WarmupState( + available_gpu_memory=[123], + num_gpu_blocks=7, + vllm_version="preserved", + ) + runtime.save_warmup_state(str(tmp_path), original) + runtime._state = runtime.CUDAGraphExtensionState(rank=0) + + runtime.update_warmup_state_deepep_profile( + str(tmp_path), + DeepEPTransport.NVL_IPC, + 64 * 1024 * 1024, + ) + + saved = runtime.load_warmup_state(str(tmp_path)) + assert saved.available_gpu_memory == [123] + assert saved.num_gpu_blocks == 7 + assert saved.vllm_version == "preserved" + + +def test_vllm_load_rejects_transport_before_foundry_ops( + tmp_path, + monkeypatch, + vllm_modules, +): + (tmp_path / "rank_0").mkdir() + (tmp_path / "warmup_state.json").write_text( + json.dumps( + { + "deepep_transport": "nvl_ipc", + "deepep_num_nvl_bytes": 64 * 1024 * 1024, + } + ) + ) + config = types.SimpleNamespace( + mode=vllm_modules.config.CUDAGraphExtensionMode.LOAD, + deepep_transport=DeepEPTransport.FABRIC, + workspace_root=str(tmp_path), + workspace_dir=None, + base_addr=0x600000000000, + region_size="1GB", + ) + monkeypatch.setattr(vllm_modules.runtime, "get_config", lambda: config) + operations = [] + monkeypatch.setattr( + vllm_modules.ops, + "set_skip_fatbin_processing", + lambda _value: operations.append("set_skip_fatbin_processing"), + ) + monkeypatch.setattr( + vllm_modules.ops, + "load_cuda_modules_and_libraries", + lambda _path: operations.append("load_cuda_modules_and_libraries"), + ) + monkeypatch.setattr( + vllm_modules.ops, + "set_allocation_region", + lambda _base, _size: operations.append("set_allocation_region"), + ) + parallel_config = types.SimpleNamespace( + rank=0, + world_size=1, + data_parallel_index=0, + tensor_parallel_size=1, + pipeline_parallel_size=1, + data_parallel_size=1, + nnodes=1, + ) + + with pytest.raises(RuntimeError, match="archived=nvl_ipc.*configured=fabric"): + vllm_modules.runtime.setup_graph_extension(parallel_config) + + assert operations == [] + + +def test_vllm_worker_load_rejects_transport_before_cuda_lifecycle( + tmp_path, + monkeypatch, + vllm_modules, +): + config_path = tmp_path / "config.toml" + config_path.write_text( + "\n".join( + [ + 'mode = "load"', + f'workspace_root = "{tmp_path}"', + 'deepep_transport = "fabric"', + ] + ) + ) + (tmp_path / "warmup_state.json").write_text( + json.dumps( + { + "deepep_transport": "nvl_ipc", + "deepep_num_nvl_bytes": 64 * 1024 * 1024, + } + ) + ) + hooks = vllm_modules.hooks + monkeypatch.setattr(hooks, "_INSTALLED", False) + lifecycle_events = [] + monkeypatch.setattr( + hooks, + "_quarantine_get_device_capability", + lambda: lifecycle_events.append("quarantine"), + ) + for patch_name in ( + "_patch_init_worker_distributed_environment", + "_patch_kernel_warmup", + "_patch_dummy_runs", + "_patch_load_model", + "_patch_prepare_comm_buffer", + "_patch_capture_model", + "_patch_cuda_graph_wrapper_call", + "_patch_initialize_kv_caches", + "_patch_subprocess_spawn_sites", + "_patch_deepep", + ): + monkeypatch.setattr( + hooks, + patch_name, + lambda name=patch_name: lifecycle_events.append(name), + ) + compilation_config = types.SimpleNamespace( + graph_extension_config_path=str(config_path), + cudagraph_num_of_warmups=1, + ) + + class Worker: + def __init__(self): + hooks.install_hooks(compilation_config) + + def init_device(self): + lifecycle_events.append("torch.accelerator.set_device_index") + + def start_worker(): + worker = Worker() + worker.init_device() + + with pytest.raises(RuntimeError, match="archived=nvl_ipc.*configured=fabric"): + start_worker() + + assert lifecycle_events == [] + + +def test_vllm_setup_records_computed_multiprocess_rank( + tmp_path, + monkeypatch, + vllm_modules, +): + config = types.SimpleNamespace( + mode=vllm_modules.config.CUDAGraphExtensionMode.SAVE, + workspace_root=str(tmp_path), + workspace_dir=None, + base_addr=0x600000000000, + region_size="1GB", + ) + monkeypatch.setattr(vllm_modules.runtime, "get_config", lambda: config) + parallel_config = types.SimpleNamespace( + rank=1, + world_size=2, + data_parallel_index=1, + ) + + vllm_modules.runtime.setup_graph_extension(parallel_config) + + assert config.workspace_dir == str(tmp_path / "rank_3") + assert vllm_modules.runtime.get_state().rank == 3 + + +@pytest.mark.parametrize( + ("transport", "expected"), + [ + ( + DeepEPTransport.FABRIC, + {"num_nvl_bytes": 0, "use_fabric": True}, + ), + ( + DeepEPTransport.NVL_IPC, + {"num_nvl_bytes": 64 * 1024 * 1024, "use_fabric": False}, + ), + ], +) +def test_vllm_hook_applies_archive_transport_on_save( + tmp_path, + monkeypatch, + vllm_modules, + transport, + expected, +): + original_calls = [] + + class Manager: + def _make_all2all_kwargs( + self, + max_num_tokens_per_dp_rank, + token_hidden_size, + num_ep_ranks, + num_global_experts, + num_local_experts, + ): + original_calls.append( + ( + max_num_tokens_per_dp_rank, + token_hidden_size, + num_ep_ranks, + num_global_experts, + num_local_experts, + ) + ) + return { + "group": "cpu-group", + "num_nvl_bytes": 64 * 1024 * 1024, + "num_rdma_bytes": 4096, + } + + _install_fake_all2all_module(monkeypatch, Manager) + updates = [] + monkeypatch.setattr( + vllm_modules.hooks, + "get_graph_extension_mode", + lambda: vllm_modules.config.CUDAGraphExtensionMode.SAVE, + ) + monkeypatch.setattr(vllm_modules.hooks, "get_deepep_transport", lambda: transport) + monkeypatch.setattr(vllm_modules.hooks, "get_workspace_root", lambda: str(tmp_path)) + monkeypatch.setattr( + vllm_modules.hooks.rt, + "update_warmup_state_deepep_profile", + lambda root, selected, num_nvl_bytes: updates.append((root, selected, num_nvl_bytes)), + ) + monkeypatch.setenv("FOUNDRY_DEEPEP_NVL_IPC", "1") + + vllm_modules.hooks._patch_deepep() + result = Manager()._make_all2all_kwargs( + 128, + 7168, + 8, + num_global_experts=128, + num_local_experts=16, + ) + + assert result == { + **expected, + "group": "cpu-group", + "num_rdma_bytes": 4096, + } + assert original_calls == [(128, 7168, 8, 128, 16)] + assert updates == [(str(tmp_path), transport, expected["num_nvl_bytes"])] + + +def test_vllm_hook_leaves_none_mode_unchanged(monkeypatch, vllm_modules): + upstream = { + "num_nvl_bytes": 64 * 1024 * 1024, + "use_fabric": "upstream", + } + + class Manager: + def _make_all2all_kwargs( + self, + max_num_tokens_per_dp_rank, + token_hidden_size, + num_ep_ranks, + num_global_experts, + num_local_experts, + ): + return dict(upstream) + + _install_fake_all2all_module(monkeypatch, Manager) + monkeypatch.setattr( + vllm_modules.hooks, + "get_graph_extension_mode", + lambda: vllm_modules.config.CUDAGraphExtensionMode.NONE, + ) + monkeypatch.setattr( + vllm_modules.hooks.rt, + "update_warmup_state_deepep_profile", + lambda *_args: pytest.fail("NONE mode must not update the archive"), + raising=False, + ) + + vllm_modules.hooks._patch_deepep() + + assert Manager()._make_all2all_kwargs(128, 7168, 8, 128, 16) == upstream + + +@pytest.mark.parametrize( + ("transport", "num_nvl_bytes"), + [ + (DeepEPTransport.NVL_IPC, 0), + (DeepEPTransport.NVL_IPC, -1), + (DeepEPTransport.NVL_IPC, True), + (DeepEPTransport.NVL_IPC, 1.0), + (DeepEPTransport.NVL_IPC, "67108864"), + (DeepEPTransport.FABRIC, True), + ], +) +@pytest.mark.parametrize( + "mode_name", + ["SAVE", "LOAD"], +) +def test_vllm_hook_rejects_invalid_upstream_bytes_before_archive_access( + tmp_path, + monkeypatch, + vllm_modules, + transport, + num_nvl_bytes, + mode_name, +): + class Manager: + def _make_all2all_kwargs( + self, + max_num_tokens_per_dp_rank, + token_hidden_size, + num_ep_ranks, + num_global_experts, + num_local_experts, + ): + return {"num_nvl_bytes": num_nvl_bytes} + + _install_fake_all2all_module(monkeypatch, Manager) + mode = getattr(vllm_modules.config.CUDAGraphExtensionMode, mode_name) + monkeypatch.setattr(vllm_modules.hooks, "get_graph_extension_mode", lambda: mode) + monkeypatch.setattr(vllm_modules.hooks, "get_deepep_transport", lambda: transport) + monkeypatch.setattr(vllm_modules.hooks, "get_workspace_root", lambda: str(tmp_path)) + monkeypatch.setattr( + vllm_modules.hooks.rt, + "update_warmup_state_deepep_profile", + lambda *_args: pytest.fail("invalid bytes must not update the archive"), + ) + monkeypatch.setattr( + vllm_modules.hooks.rt, + "load_warmup_state", + lambda *_args: pytest.fail("invalid bytes must not read the archive"), + ) + vllm_modules.hooks._patch_deepep() + + with pytest.raises(RuntimeError, match="num_nvl_bytes"): + Manager()._make_all2all_kwargs(128, 7168, 8, 128, 16) + + +def test_vllm_load_rejects_effective_bytes_before_communicator_setup( + tmp_path, + monkeypatch, + vllm_modules, +): + original_calls = [] + + class Manager: + def _make_all2all_kwargs( + self, + max_num_tokens_per_dp_rank, + token_hidden_size, + num_ep_ranks, + num_global_experts, + num_local_experts, + ): + original_calls.append( + ( + max_num_tokens_per_dp_rank, + token_hidden_size, + num_ep_ranks, + num_global_experts, + num_local_experts, + ) + ) + return {"num_nvl_bytes": 128 * 1024 * 1024} + + _install_fake_all2all_module(monkeypatch, Manager) + communicator_calls = [] + monkeypatch.setattr( + vllm_modules.hooks, + "get_graph_extension_mode", + lambda: vllm_modules.config.CUDAGraphExtensionMode.LOAD, + ) + monkeypatch.setattr( + vllm_modules.hooks, + "get_deepep_transport", + lambda: DeepEPTransport.NVL_IPC, + ) + monkeypatch.setattr(vllm_modules.hooks, "get_workspace_root", lambda: str(tmp_path)) + monkeypatch.setattr( + vllm_modules.hooks.rt, + "load_warmup_state", + lambda _root: types.SimpleNamespace(deepep_num_nvl_bytes=64 * 1024 * 1024), + ) + vllm_modules.hooks._patch_deepep() + + def construct_communicator(): + kwargs = Manager()._make_all2all_kwargs( + 128, + 7168, + 8, + num_global_experts=128, + num_local_experts=16, + ) + communicator_calls.append(kwargs) + + with pytest.raises(RuntimeError, match="archived=67108864.*effective=134217728"): + construct_communicator() + + assert communicator_calls == [] + assert original_calls == [(128, 7168, 8, 128, 16)] + + +def test_vllm_transport_bootstrap_does_not_skip_first_save_profile( + tmp_path, + monkeypatch, + vllm_modules, +): + hooks = vllm_modules.hooks + profile_calls = [] + initialized_configs = [] + + class ModelExecutor: + def get_kv_cache_specs(self): + return [object()] + + def determine_available_memory(self): + profile_calls.append("determine_available_memory") + return [123456] + + def initialize_from_config(self, configs): + initialized_configs.append(configs) + + class EngineCore: + def __init__(self): + self.model_executor = ModelExecutor() + self.available_gpu_memory_for_kv_cache = 0 + + def _initialize_kv_caches(self, _vllm_config): + pytest.fail("SAVE mode must use Foundry KV-cache initialization") + + scheduler_config = types.SimpleNamespace(num_blocks=7, kv_cache_groups=[]) + _install_fake_kv_cache_modules(monkeypatch, EngineCore, scheduler_config) + bootstrap = hooks.rt.WarmupState( + deepep_transport="nvl_ipc", + deepep_num_nvl_bytes=64 * 1024 * 1024, + ) + hooks.rt.save_warmup_state(str(tmp_path), bootstrap) + monkeypatch.setattr( + hooks, + "get_graph_extension_mode", + lambda: vllm_modules.config.CUDAGraphExtensionMode.SAVE, + ) + monkeypatch.setattr(hooks, "get_workspace_root", lambda: str(tmp_path)) + runtime_metadata = { + "vllm_version": "test-vllm", + "timestamp": "2026-07-23T10:00:00", + "cuda_version": "13.0", + "gpu_name": "H100", + "gpu_total_memory": 80 * 1024**3, + } + monkeypatch.setattr( + hooks.rt, + "create_warmup_state", + lambda: hooks.rt.WarmupState(**runtime_metadata), + ) + monkeypatch.setattr(hooks.rt, "get_final_alloc_offset", lambda: 0) + hooks._patch_initialize_kv_caches() + vllm_config = types.SimpleNamespace( + model_config=types.SimpleNamespace(max_model_len=1024), + parallel_config=types.SimpleNamespace(data_parallel_index=0, rank=0), + cache_config=types.SimpleNamespace( + num_gpu_blocks=0, + block_size=0, + gpu_memory_utilization=0.9, + ), + validate_block_size=lambda: None, + ) + + result = EngineCore()._initialize_kv_caches(vllm_config) + + assert result is scheduler_config + assert profile_calls == ["determine_available_memory"] + assert len(initialized_configs) == 1 + saved = hooks.rt.load_warmup_state(str(tmp_path)) + assert saved.available_gpu_memory == [123456] + assert saved.num_gpu_blocks == 7 + for field, expected in runtime_metadata.items(): + assert getattr(saved, field) == expected + assert saved.deepep_transport == "nvl_ipc" + assert saved.deepep_num_nvl_bytes == 64 * 1024 * 1024 diff --git a/tests/test_experimental_recipe.py b/tests/test_experimental_recipe.py new file mode 100644 index 00000000..888fc7c7 --- /dev/null +++ b/tests/test_experimental_recipe.py @@ -0,0 +1,232 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the Foundry project +"""Portable configuration contracts for the experimental recipes.""" + +import os +import re +import subprocess +from pathlib import Path + +import tomllib + +ROOT = Path(__file__).resolve().parents[1] +EXPERIMENTAL_DIR = ROOT / "recipe" / "experimental" +VLLM_CONFIGS = ( + "foundry_save.toml", + "foundry_load.toml", + "foundry_save_fp8.toml", + "foundry_load_fp8.toml", +) +VLLM_SCRIPTS = ( + "serve_qwen3-30ba3b_ipc_ep.sh", + "serve_qwen3-30ba3bfp8_ipc_ep.sh", +) +INCOMPATIBLE_SGLANG_EP_ASSETS = ( + "serve_qwen3-30ba3bfp8_sglang_ipc_ep.sh", + "sglang_foundry_save.toml", + "sglang_foundry_load.toml", +) +DEEPEP_MATRIX_COMMANDS = ( + "bash tests/run_deepep_matrix.sh fabric both", + "bash tests/run_deepep_matrix.sh nvl_ipc both", + "bash tests/run_deepep_matrix.sh nvl_ipc_prealloc both", +) +DEEPEP_EVIDENCE_DOCS = ( + ROOT / "README.md", + ROOT / "ROADMAP.md", + ROOT / "RELEASE.md", + ROOT / "docs" / "vllm" / "moe-and-deepep.md", + ROOT / "recipe" / "vllm" / "README.md", +) +VLLM_EVIDENCE_DOCS = ( + *DEEPEP_EVIDENCE_DOCS, + ROOT / "docs" / "vllm" / "overview.md", +) +FINAL_FOUNDRY_REF = "a8e43442421a80fea7991ee26f97f1003b170e29" +FINAL_NATIVE_APP = "ap-0aeSCJNb6eHQhqdMkpyfXk" +FINAL_DEEPEP_APP = "ap-Q8mb8U4EYuE5Jo6Zb0xwrH" +FINAL_VLLM_APP = "ap-YfKALdMwnyJQmzQoVvZSMU" +STALE_FINAL_EVIDENCE = ( + "6ee2216ba619d26b4c9e26ae9fbb59c707ee4315", + "de79cb813d28e1149f0dba8316731d33495dc6d8", + "ap-xgSEhpwHIgrfcVjESRYZz8", + "ap-8mg9uXqgyBQHCpCauEcu1S", + "ap-Cgzmx0YbjvxdFKvpwuLBdZ", +) +MATRIX_SCRIPT = ROOT / "tests" / "run_deepep_matrix.sh" + + +def test_vllm_configs_do_not_enable_author_specific_nvshmem_path() -> None: + for filename in VLLM_CONFIGS: + text = (EXPERIMENTAL_DIR / filename).read_text() + config = tomllib.loads(text) + assert "nvshmem_host_path" not in config, filename + assert "/data/liuxs" not in text, filename + + +def test_no_imex_vllm_configs_select_archived_nvl_ipc_transport() -> None: + for filename in VLLM_CONFIGS: + config = tomllib.loads((EXPERIMENTAL_DIR / filename).read_text()) + assert config.get("deepep_transport") == "nvl_ipc", filename + + +def test_no_imex_vllm_scripts_keep_profile_flags_without_env_transport_switch() -> None: + for filename in VLLM_SCRIPTS: + text = (EXPERIMENTAL_DIR / filename).read_text() + assert "FOUNDRY_DEEPEP_NVL_IPC" not in text, filename + assert "README.md#vllm-workflow" in text, filename + assert 'README "Required code"' not in text, filename + for required_export in ( + "export NCCL_CUMEM_ENABLE=0", + "export NCCL_NVLS_ENABLE=0", + "export VLLM_USE_V2_MODEL_RUNNER=0", + ): + assert required_export in text, f"{filename}: missing {required_export}" + + +def test_incompatible_sglang_ep_assets_are_quarantined_without_dense_tp_deletion() -> None: + for filename in INCOMPATIBLE_SGLANG_EP_ASSETS: + assert not (EXPERIMENTAL_DIR / filename).exists(), filename + + for filename in ( + "serve_qwen3-8b_sglang_tp.sh", + "sglang_foundry_tp_save.toml", + "sglang_foundry_tp_load.toml", + ): + assert (EXPERIMENTAL_DIR / filename).is_file(), filename + + readme = (EXPERIMENTAL_DIR / "README.md").read_text() + for term in ( + "Quarantined SGLang EP profile", + "current HookRegistry", + "dense Qwen3-8B TP=2", + "torch symmetric memory", + "EP and DeepEP are rejected before launch", + "fresh SAVE-to-LOAD H100 evidence", + ): + assert term in readme + assert "the EP recipe puts" not in readme + + +def test_vllm_docs_define_archive_transport_contract_and_actual_matrix_cases() -> None: + text = (ROOT / "docs" / "vllm" / "moe-and-deepep.md").read_text() + + for term in ( + 'deepep_transport = "fabric"', + 'deepep_transport = "nvl_ipc"', + "legacy archives", + "deepep_num_nvl_bytes", + "rank 0", + "before CUDA", + "first SAVE", + "libnvshmem_host.so.3", + "libnvshmem_host.so", + "silent fallback", + "Foundry DeepEP transport mismatch: archived=nvl_ipc, configured=fabric", + "Foundry DeepEP num_nvl_bytes mismatch", + ): + assert term in text + for command in DEEPEP_MATRIX_COMMANDS: + assert command in text + assert "run_deepep_matrix.sh ll_nofabric both" not in text + + +def test_deepep_evidence_docs_pin_final_candidate_and_verified_run() -> None: + for path in DEEPEP_EVIDENCE_DOCS: + text = path.read_text() + assert FINAL_FOUNDRY_REF in text, path + assert FINAL_DEEPEP_APP in text, path + + +def test_vllm_evidence_docs_use_only_final_runtime_evidence() -> None: + for path in VLLM_EVIDENCE_DOCS: + text = path.read_text() + assert FINAL_FOUNDRY_REF in text, path + stale = [value for value in STALE_FINAL_EVIDENCE if value in text] + assert not stale, f"{path}: stale final evidence {stale}" + + for path in ( + ROOT / "README.md", + ROOT / "ROADMAP.md", + ROOT / "RELEASE.md", + ROOT / "docs" / "vllm" / "overview.md", + ): + text = path.read_text() + assert FINAL_NATIVE_APP in text, path + assert "823 core" in text, path + + for path in ( + ROOT / "README.md", + ROOT / "ROADMAP.md", + ROOT / "RELEASE.md", + ROOT / "docs" / "vllm" / "overview.md", + ROOT / "recipe" / "vllm" / "README.md", + ): + assert FINAL_VLLM_APP in path.read_text(), path + + +def test_vllm_evidence_docs_do_not_describe_fingerprint_count_as_bytes() -> None: + inaccurate_count = re.compile(r"\b265\b[^\n.]{0,40}\bbytes\b", re.IGNORECASE) + inaccurate_exit_phrases = ( + "zero phase exits", + "zero phase and modal exits", + ) + + for path in VLLM_EVIDENCE_DOCS: + text = path.read_text() + assert inaccurate_count.search(text) is None, path + for phrase in inaccurate_exit_phrases: + assert phrase not in text.lower(), path + + +def test_deepep_matrix_fabric_case_is_strict_and_has_unique_port() -> None: + text = MATRIX_SCRIPT.read_text() + match = re.search(r"^\s*fabric\)(.*?)\s*;;", text, flags=re.MULTILINE) + assert match is not None + case_body = match.group(1) + + for assignment in ( + "TEST_USE_FABRIC=1", + "TEST_NVL_BYTES_MB=0", + "NCCL_CUMEM_ENABLE=0", + "NCCL_NVLS_ENABLE=0", + ): + assert assignment in case_body + + fabric_port = re.search(r"\bPORT=(\d+)\b", case_body) + assert fabric_port is not None + all_ports = re.findall(r"\bPORT=(\d+)\b", text) + assert all_ports.count(fabric_port.group(1)) == 1 + + +def test_troubleshooting_only_documents_real_compile_cache_remedy() -> None: + readme = (EXPERIMENTAL_DIR / "README.md").read_text() + assert "FOUNDRY_DISABLE_COMBO_KERNELS" not in readme + assert "experimental/expert-parallel" not in readme + + +def test_deepep_matrix_script_rejects_missing_nvshmem(tmp_path: Path) -> None: + script = MATRIX_SCRIPT + text = script.read_text() + for override in ( + "PYTHON_BIN", + "VLLM_ROOT", + "NVSHMEM_SO", + "DEEPEP_MATRIX_WORK_ROOT", + "DEEPEP_MATRIX_LOG_DIR", + ): + assert override in text + assert "/data/liuxs" not in text + + missing_nvshmem = tmp_path / "libnvshmem_host.so" + env = os.environ.copy() + env["NVSHMEM_SO"] = str(missing_nvshmem) + result = subprocess.run( + ["bash", str(script), "ll_nofabric", "save"], + check=False, + capture_output=True, + env=env, + text=True, + ) + assert result.returncode == 4 + assert result.stderr.strip() == f"NVSHMEM library not found: {missing_nvshmem}" diff --git a/tests/test_generation_recreate_contract.py b/tests/test_generation_recreate_contract.py new file mode 100644 index 00000000..d8d7a0ec --- /dev/null +++ b/tests/test_generation_recreate_contract.py @@ -0,0 +1,82 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the Foundry project +"""Source contract for the native generation-recreate test. + +The generation-recreate regression must exercise BOTH generations inside a +single live allocation region: it may free/recreate only the preallocation +between generations, never stop/re-set the region. This is a pure text contract +(no CUDA/torch/foundry import) so it runs on a CPU-only host and guards against +a stop/set sneaking back into the test. +""" + +from __future__ import annotations + +from pathlib import Path + +TEST_FILE = Path(__file__).parent / "test_vmm_alloc.py" + + +def _function_source(source: str, def_line: str) -> str: + """Return the source of a top-level function (from its ``def`` line to the + next top-level ``def``/``if __name__`` at column 0).""" + start = source.index(def_line) + rest = source[start + len(def_line) :] + end = len(rest) + for marker in ("\ndef ", "\nif __name__"): + idx = rest.find(marker) + if idx != -1: + end = min(end, idx) + return def_line + rest[:end] + + +def _gen_recreate_body() -> str: + return _function_source(TEST_FILE.read_text(), "def _run_core_generation_recreate():") + + +def test_generation_recreate_keeps_one_region_for_both_generations() -> None: + body = _gen_recreate_body() + # Exactly one region lifecycle spanning both generations. + assert body.count("set_allocation_region(") == 1, ( + "generation test must set the allocation region exactly once (one region " + "for both generations)" + ) + assert body.count("stop_allocation_region(") == 1, ( + "generation test must stop the allocation region exactly once" + ) + + +def test_generation_recreate_recreates_only_the_preallocation() -> None: + body = _gen_recreate_body() + # The preallocation is freed between generations (recreate the preallocation, + # not the region). + assert "free_preallocated_region()" in body, ( + "generation test must free the preallocation between generations" + ) + # A second generation is genuinely created (the helper is invoked twice). + assert body.count("preallocate_export_import_read(") >= 3, ( + "generation test must define and invoke the export/import helper for both generations" + ) + + +def test_generation_recreate_drains_caching_allocator_before_free() -> None: + body = _gen_recreate_body() + del_idx = body.find("del tensor") + empty_idx = body.find("torch.cuda.empty_cache()") + free_idx = body.find("free_preallocated_region()") + assert del_idx != -1, "must delete the generation tensor" + assert empty_idx != -1, "must drain the caching allocator via empty_cache()" + assert free_idx != -1, "must free the preallocation" + # del tensor -> empty_cache -> free_preallocated_region so the caching + # allocator cannot retain a pointer into unmapped generation backing. + assert del_idx < empty_idx < free_idx, ( + "must order del tensor -> empty_cache -> free_preallocated_region" + ) + + +def test_generation_recreate_makes_no_same_va_claim() -> None: + body = _gen_recreate_body() + # Distinct chunk base is asserted; identical source VA / chunk base is not. + assert "cb2 != cb1" in body, "must assert generation 2 lands at a different chunk base" + assert "== ptr1" not in body and "== cb1" not in body, ( + "must not claim an identical VA / chunk base across generations" + ) diff --git a/tests/test_imports.py b/tests/test_imports.py index 2de75a6e..1b03e9b2 100644 --- a/tests/test_imports.py +++ b/tests/test_imports.py @@ -13,6 +13,11 @@ """ import importlib +import os +import subprocess +import sys +import textwrap +from pathlib import Path PUBLIC_MODULES = [ # Top-level @@ -69,6 +74,60 @@ def test_re_exports_present(): ) +def test_public_cudagraph_is_python_wrapper(): + """The top-level API must not be overwritten by the raw pybind export.""" + source = textwrap.dedent( + """ + import importlib.abc + import importlib.util + import sys + import types + + torch = types.ModuleType("torch") + + class Tensor: + pass + + torch.Tensor = Tensor + sys.modules["torch"] = torch + + class RawCUDAGraph: + pass + + class OpsLoader(importlib.abc.Loader): + def exec_module(self, module): + module.CUDAGraph = RawCUDAGraph + for name in ( + "free_preallocated_region", + "get_current_alloc_offset", + "preallocate_region", + "resume_allocation_region", + "set_allocation_region", + "set_current_alloc_offset", + "stop_allocation_region", + ): + setattr(module, name, lambda *args, **kwargs: None) + + class OpsFinder(importlib.abc.MetaPathFinder): + def find_spec(self, fullname, path, target=None): + if fullname == "foundry.ops": + return importlib.util.spec_from_loader(fullname, OpsLoader()) + return None + + sys.meta_path.insert(0, OpsFinder()) + + import foundry + from foundry.graph import CUDAGraph as PythonCUDAGraph + + assert foundry.CUDAGraph is PythonCUDAGraph + assert foundry.ops.CUDAGraph is RawCUDAGraph + """ + ) + env = os.environ.copy() + env["PYTHONPATH"] = str(Path(__file__).resolve().parents[1] / "python") + subprocess.run([sys.executable, "-c", source], check=True, env=env) + + def test_vllm_integration_public_api(): mod = importlib.import_module("foundry.integration.vllm") for name in ("install_hooks", "CUDAGraphExtensionMode", "get_graph_extension_mode"): diff --git a/tests/test_modal_deepep.py b/tests/test_modal_deepep.py new file mode 100644 index 00000000..e4a142fb --- /dev/null +++ b/tests/test_modal_deepep.py @@ -0,0 +1,422 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the Foundry project +"""CPU-only contracts for the strict DeepEP Modal matrix.""" + +from __future__ import annotations + +import ast +import importlib.util +import json +import signal +import subprocess +import sys +import types +from pathlib import Path + +import pytest + +REPO_ROOT = Path(__file__).parents[1] +HARNESS_PATH = REPO_ROOT / "tests" / "modal_deepep.py" +FULL_SHA = "0123456789abcdef0123456789abcdef01234567" +EXPECTED_CASES = ("fabric", "nvl_ipc", "nvl_ipc_prealloc") +EXPECTED_PREFLIGHT_COMMAND = ( + "python3", + "-m", + "pytest", + "tests/test_deepep_fabric.py::test_torch_cpu_preflight_covers_split_payload_and_dtypes", + "-q", +) +EXPECTED_CASE_COMMANDS = [ + ("bash", "tests/run_deepep_matrix.sh", case, "both") for case in EXPECTED_CASES +] + + +class _FakeImage: + def __init__(self): + self.environment = {} + self.environment_before_run_commands = [] + + @classmethod + def from_registry(cls, *args, **kwargs): + return cls() + + def env(self, values): + self.environment.update(values) + return self + + def run_commands(self, *args, **kwargs): + self.environment_before_run_commands.append(dict(self.environment)) + return self + + def __getattr__(self, _name): + return lambda *args, **kwargs: self + + +class _FakeApp: + def __init__(self, _name): + pass + + def function(self, **_kwargs): + return lambda function: function + + def local_entrypoint(self): + return lambda function: function + + +class _CompletedProcess: + next_pid = 4000 + + def __init__(self, command, returncode: int, output: str): + self.command = command + self.returncode = returncode + self.output = output + self.pid = self.next_pid + type(self).next_pid += 1 + self.communicate_timeouts = [] + + def communicate(self, timeout): + self.communicate_timeouts.append(timeout) + return self.output, None + + +class _TimeoutProcess: + pid = 9123 + + def __init__(self, command): + self.command = command + self.returncode = None + self.group_alive = True + self.communicate_timeouts = [] + self.wait_timeouts = [] + + def communicate(self, timeout): + self.communicate_timeouts.append(timeout) + if len(self.communicate_timeouts) == 1: + raise subprocess.TimeoutExpired( + self.command, + timeout, + output="partial output\n", + ) + assert not self.group_alive + self.returncode = -signal.SIGKILL + return "partial output\n", None + + def wait(self, timeout): + self.wait_timeouts.append(timeout) + assert not self.group_alive + return self.returncode + + +def _load_harness(monkeypatch: pytest.MonkeyPatch): + fake_modal = types.ModuleType("modal") + fake_modal.App = _FakeApp + fake_modal.Image = _FakeImage + monkeypatch.setitem(sys.modules, "modal", fake_modal) + monkeypatch.setenv("FOUNDRY_REF", FULL_SHA) + spec = importlib.util.spec_from_file_location("modal_deepep_contract", HARNESS_PATH) + assert spec is not None and spec.loader is not None + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + return module + + +def _assignment_source(name: str) -> str: + source = HARNESS_PATH.read_text() + tree = ast.parse(source) + for node in tree.body: + if isinstance(node, ast.Assign) and any( + isinstance(target, ast.Name) and target.id == name for target in node.targets + ): + segment = ast.get_source_segment(source, node) + assert segment is not None + return segment + raise AssertionError(f"{name} assignment not found") + + +def _assert_popen_kwargs(kwargs) -> None: + assert kwargs == { + "cwd": "/foundry", + "stdout": subprocess.PIPE, + "stderr": subprocess.STDOUT, + "text": True, + "start_new_session": True, + } + + +def test_matrix_plan_has_torch_preflight_then_exactly_three_cases( + monkeypatch: pytest.MonkeyPatch, +) -> None: + harness = _load_harness(monkeypatch) + + assert harness.MATRIX_CASES == EXPECTED_CASES + assert harness.PREFLIGHT_COMMAND == EXPECTED_PREFLIGHT_COMMAND + assert harness.commands_for_matrix() == EXPECTED_CASE_COMMANDS + assert harness.execution_plan() == [ + ( + "preflight", + "torch_cpu_oracle", + EXPECTED_PREFLIGHT_COMMAND, + harness.PREFLIGHT_TIMEOUT_SECONDS, + ), + *[ + ("case", case, command, harness.CASE_TIMEOUT_SECONDS) + for case, command in zip(EXPECTED_CASES, EXPECTED_CASE_COMMANDS, strict=True) + ], + ] + + +def test_matrix_is_bounded_h100_pair_with_outer_timeout_buffer( + monkeypatch: pytest.MonkeyPatch, +) -> None: + harness = _load_harness(monkeypatch) + assert harness.PREFLIGHT_TIMEOUT_SECONDS == 2 * 60 + assert harness.CASE_TIMEOUT_SECONDS == 20 * 60 + assert harness.POST_KILL_TIMEOUT_SECONDS == 10 + assert harness.FUNCTION_TIMEOUT_SECONDS == 75 * 60 + required_child_time = ( + harness.PREFLIGHT_TIMEOUT_SECONDS + len(EXPECTED_CASES) * harness.CASE_TIMEOUT_SECONDS + ) + + assert required_child_time + 5 * 60 <= harness.FUNCTION_TIMEOUT_SECONDS + source = HARNESS_PATH.read_text() + assert 'gpu="H100:2"' in source + assert "timeout=FUNCTION_TIMEOUT_SECONDS" in source + + +@pytest.mark.parametrize("foundry_ref", [None, "", "main", "abc123", "g" * 40]) +def test_harness_requires_full_immutable_foundry_ref( + monkeypatch: pytest.MonkeyPatch, + foundry_ref: str | None, +) -> None: + harness = _load_harness(monkeypatch) + + with pytest.raises(ValueError, match="FOUNDRY_REF must be a full 40-hex commit SHA"): + harness.require_full_commit_sha(foundry_ref) + + +def test_image_clones_candidate_repository_and_pins_dependency_toolchain() -> None: + source = HARNESS_PATH.read_text() + image_source = _assignment_source("image") + + for pin in ( + 'FOUNDRY_REPO = "https://github.com/modal-projects/foundry.git"', + 'CUDA_TAG = "13.0.1-cudnn-devel-ubuntu24.04"', + 'VLLM_REF = "4309c257d3f639e5490d3811293c890c61c76f29"', + 'DEEPEP_REF = "29d31c095796f3c8ece47ee9cdcc167051bbeed9"', + 'NVSHMEM_VERSION = "3.3.24"', + ): + assert pin in source + assert 'add_python="3.12"' in image_source + assert '"torch==2.11.0"' in image_source + assert 'index_url="https://download.pytorch.org/whl/cu130"' in image_source + assert "--deepep-ref {DEEPEP_REF}" in image_source + assert "--nvshmem-ver {NVSHMEM_VERSION}" in image_source + assert "TORCH_CUDA_ARCH_LIST=9.0" in image_source + + +def test_image_exports_uv_system_python_before_commands_without_changing_pins( + monkeypatch: pytest.MonkeyPatch, +) -> None: + harness = _load_harness(monkeypatch) + image_source = _assignment_source("image") + + assert harness.image.environment_before_run_commands + assert all( + environment.get("UV_SYSTEM_PYTHON") == "1" + for environment in harness.image.environment_before_run_commands + ) + assert image_source.index('"UV_SYSTEM_PYTHON": "1"') < image_source.index(".run_commands(") + assert { + "FOUNDRY_REPO": harness.FOUNDRY_REPO, + "VLLM_REPO": harness.VLLM_REPO, + "VLLM_REF": harness.VLLM_REF, + "DEEPEP_REF": harness.DEEPEP_REF, + "NVSHMEM_VERSION": harness.NVSHMEM_VERSION, + } == { + "FOUNDRY_REPO": "https://github.com/modal-projects/foundry.git", + "VLLM_REPO": "https://github.com/foundry-org/vllm.git", + "VLLM_REF": "4309c257d3f639e5490d3811293c890c61c76f29", + "DEEPEP_REF": "29d31c095796f3c8ece47ee9cdcc167051bbeed9", + "NVSHMEM_VERSION": "3.3.24", + } + + +def test_image_checks_out_full_foundry_ref_before_gcc_build( + monkeypatch: pytest.MonkeyPatch, +) -> None: + harness = _load_harness(monkeypatch) + source = HARNESS_PATH.read_text() + image_source = _assignment_source("image") + ref_assignment = 'FOUNDRY_REF = require_full_commit_sha(os.environ.get("FOUNDRY_REF"))' + checkout = "git -C /foundry checkout --detach {FOUNDRY_REF}" + verify = 'test "$(git -C /foundry rev-parse HEAD)" = "{FOUNDRY_REF}"' + build = "CC=gcc CXX=g++ pip install -e . --no-build-isolation" + + assert source.index(ref_assignment) < source.index("image =") + assert checkout in image_source + assert verify in image_source + assert build in image_source + assert image_source.index(checkout) < image_source.index(verify) < image_source.index(build) + assert {key: harness.image.environment[key] for key in ("CC", "CXX", "FOUNDRY_REF")} == { + "CC": "gcc", + "CXX": "g++", + "FOUNDRY_REF": FULL_SHA, + } + + +def test_preflight_failure_is_structured_and_aborts_before_cases( + monkeypatch: pytest.MonkeyPatch, + capsys: pytest.CaptureFixture[str], +) -> None: + harness = _load_harness(monkeypatch) + calls = [] + + def popen_factory(command, **kwargs): + calls.append(command) + _assert_popen_kwargs(kwargs) + return _CompletedProcess(command, returncode=7, output="oracle failed\n") + + with pytest.raises(RuntimeError) as error: + harness.execute_deepep_matrix(popen_factory=popen_factory) + + assert calls == [EXPECTED_PREFLIGHT_COMMAND] + event = json.loads(capsys.readouterr().out) + assert event == { + "command": "python3 -m pytest " + "tests/test_deepep_fabric.py::test_torch_cpu_preflight_covers_split_payload_and_dtypes " + "-q", + "event": "deepep_matrix_command", + "name": "torch_cpu_oracle", + "output": "oracle failed\n", + "returncode": 7, + "stage": "preflight", + } + failure = json.loads(str(error.value)) + assert failure == { + "foundry_ref": FULL_SHA, + "results": [ + { + "command": event["command"], + "name": "torch_cpu_oracle", + "output": "oracle failed\n", + "returncode": 7, + "stage": "preflight", + } + ], + "status": "failed", + } + + +def test_case_failure_aborts_after_successful_preflight( + monkeypatch: pytest.MonkeyPatch, +) -> None: + harness = _load_harness(monkeypatch) + calls = [] + outcomes = iter(((0, "oracle passed\n"), (9, "fabric failed\n"))) + + def popen_factory(command, **kwargs): + calls.append(command) + _assert_popen_kwargs(kwargs) + returncode, output = next(outcomes) + return _CompletedProcess(command, returncode, output) + + with pytest.raises(RuntimeError) as error: + harness.execute_deepep_matrix(popen_factory=popen_factory) + + assert calls == [EXPECTED_PREFLIGHT_COMMAND, EXPECTED_CASE_COMMANDS[0]] + failure = json.loads(str(error.value)) + assert [result["stage"] for result in failure["results"]] == ["preflight", "case"] + assert [result["name"] for result in failure["results"]] == [ + "torch_cpu_oracle", + "fabric", + ] + + +def test_timeout_kills_process_group_and_uses_bounded_post_kill_reap( + monkeypatch: pytest.MonkeyPatch, + capsys: pytest.CaptureFixture[str], +) -> None: + harness = _load_harness(monkeypatch) + process = _TimeoutProcess(EXPECTED_PREFLIGHT_COMMAND) + killed_groups = [] + + def popen_factory(command, **kwargs): + assert command == EXPECTED_PREFLIGHT_COMMAND + _assert_popen_kwargs(kwargs) + return process + + def kill_process_group(process_group, sig): + killed_groups.append((process_group, sig)) + assert process_group == process.pid + assert sig == signal.SIGKILL + process.group_alive = False + + with pytest.raises(RuntimeError) as error: + harness.execute_deepep_matrix( + popen_factory=popen_factory, + kill_process_group=kill_process_group, + ) + + assert killed_groups == [(process.pid, signal.SIGKILL)] + assert not process.group_alive + assert process.communicate_timeouts == [ + harness.PREFLIGHT_TIMEOUT_SECONDS, + harness.POST_KILL_TIMEOUT_SECONDS, + ] + assert process.wait_timeouts == [harness.POST_KILL_TIMEOUT_SECONDS] + event = json.loads(capsys.readouterr().out) + assert event["stage"] == "preflight" + assert event["returncode"] == 124 + assert event["output"] == ( + f"partial output\nTimed out after {harness.PREFLIGHT_TIMEOUT_SECONDS} seconds\n" + ) + failure = json.loads(str(error.value)) + assert failure["results"][0]["returncode"] == 124 + + +def test_success_returns_preflight_and_all_three_case_results( + monkeypatch: pytest.MonkeyPatch, + capsys: pytest.CaptureFixture[str], +) -> None: + harness = _load_harness(monkeypatch) + calls = [] + + def popen_factory(command, **kwargs): + calls.append(command) + _assert_popen_kwargs(kwargs) + name = "torch_cpu_oracle" if command == EXPECTED_PREFLIGHT_COMMAND else command[2] + return _CompletedProcess(command, returncode=0, output=f"{name} passed\n") + + result = harness.execute_deepep_matrix(popen_factory=popen_factory) + + assert calls == [EXPECTED_PREFLIGHT_COMMAND, *EXPECTED_CASE_COMMANDS] + assert result["foundry_ref"] == FULL_SHA + assert result["status"] == "passed" + assert [(item["stage"], item["name"]) for item in result["results"]] == [ + ("preflight", "torch_cpu_oracle"), + *[("case", case) for case in EXPECTED_CASES], + ] + assert [item["returncode"] for item in result["results"]] == [0, 0, 0, 0] + events = [json.loads(line) for line in capsys.readouterr().out.splitlines()] + assert [(event["stage"], event["name"]) for event in events] == [ + ("preflight", "torch_cpu_oracle"), + *[("case", case) for case in EXPECTED_CASES], + ] + + +def test_orchestration_uses_process_groups_without_importing_gpu_runtime() -> None: + source = HARNESS_PATH.read_text() + tree = ast.parse(source) + imported_roots = set() + for node in ast.walk(tree): + if isinstance(node, ast.Import): + imported_roots.update(alias.name.split(".", 1)[0] for alias in node.names) + elif isinstance(node, ast.ImportFrom) and node.module is not None: + imported_roots.add(node.module.split(".", 1)[0]) + + assert imported_roots.isdisjoint({"torch", "foundry", "deep_ep"}) + assert "subprocess.Popen" in source + assert "subprocess.run" not in source + assert "start_new_session=True" in source + assert "os.killpg" in source + assert "signal.SIGKILL" in source diff --git a/tests/test_modal_native.py b/tests/test_modal_native.py new file mode 100644 index 00000000..401b0d27 --- /dev/null +++ b/tests/test_modal_native.py @@ -0,0 +1,178 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the Foundry project +"""CPU-only contracts for the native CUDA Modal harness.""" + +from __future__ import annotations + +import ast +import importlib.util +import json +import os +import subprocess +import sys +from pathlib import Path + +import pytest + +REPO_ROOT = Path(__file__).parents[1] +HARNESS_PATH = REPO_ROOT / "tests" / "modal_native.py" +FULL_SHA = "0123456789abcdef0123456789abcdef01234567" + + +def _load_harness(monkeypatch: pytest.MonkeyPatch): + monkeypatch.setenv("FOUNDRY_REF", FULL_SHA) + spec = importlib.util.spec_from_file_location("modal_native_contract", HARNESS_PATH) + assert spec is not None and spec.loader is not None + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + return module + + +def _image_definition_source() -> str: + source = HARNESS_PATH.read_text() + tree = ast.parse(source) + for node in tree.body: + if isinstance(node, ast.Assign) and any( + isinstance(target, ast.Name) and target.id == "image" for target in node.targets + ): + image_source = ast.get_source_segment(source, node) + assert image_source is not None + return image_source + raise AssertionError("image assignment not found") + + +@pytest.mark.parametrize("foundry_ref", [None, "", "main", "abc123", "g" * 40]) +def test_harness_requires_explicit_full_immutable_ref(foundry_ref: str | None) -> None: + env = os.environ.copy() + if foundry_ref is None: + env.pop("FOUNDRY_REF", None) + else: + env["FOUNDRY_REF"] = foundry_ref + + result = subprocess.run( + [sys.executable, str(HARNESS_PATH)], + cwd=REPO_ROOT, + env=env, + capture_output=True, + text=True, + ) + + assert result.returncode != 0 + assert "FOUNDRY_REF must be a full 40-hex commit SHA" in result.stderr + + +def test_harness_selects_focused_and_core_commands(monkeypatch: pytest.MonkeyPatch) -> None: + harness = _load_harness(monkeypatch) + + focused = ( + "python3", + "-m", + "pytest", + "tests/test_vmm_alloc.py", + "tests/test_rng_graph.py", + "-q", + ) + core = ( + "python3", + "-m", + "pytest", + "tests/", + "-q", + "--ignore=tests/test_imports.py", + "--ignore=tests/test_deepep_fabric.py", + ) + + assert harness.commands_for_suite("focused") == [focused] + assert harness.commands_for_suite("core") == [core] + assert harness.commands_for_suite("all") == [focused, core] + with pytest.raises(ValueError, match="focused, core, or all"): + harness.commands_for_suite("unknown") + + +def test_image_pins_cuda_python_torch_and_provisions_latest_pytest() -> None: + image_source = _image_definition_source() + + assert 'CUDA_TAG = "13.0.1-cudnn-devel-ubuntu24.04"' in HARNESS_PATH.read_text() + assert 'add_python="3.12"' in image_source + assert '"torch==2.11.0"' in image_source + assert 'index_url="https://download.pytorch.org/whl/cu130"' in image_source + assert '"pytest"' in image_source + assert '"pytest=="' not in image_source + assert image_source.index('"pytest"') < image_source.index(".run_commands(") + + +def test_image_checks_out_exact_detached_ref_and_builds_with_gcc() -> None: + image_source = _image_definition_source() + checkout = "git -C /foundry checkout --detach {FOUNDRY_REF}" + verify = 'test "$(git -C /foundry rev-parse HEAD)" = "{FOUNDRY_REF}"' + build = "CC=gcc CXX=g++ pip install -e . --no-build-isolation" + + assert checkout in image_source + assert verify in image_source + assert build in image_source + assert image_source.index(checkout) < image_source.index(verify) < image_source.index(build) + assert '.env({"CC": "gcc", "CXX": "g++", "FOUNDRY_REF": FOUNDRY_REF})' in image_source + + +def test_command_failure_emits_json_and_raises_nonzero( + monkeypatch: pytest.MonkeyPatch, + capsys: pytest.CaptureFixture[str], +) -> None: + harness = _load_harness(monkeypatch) + + def fail(command, **kwargs): + assert kwargs == { + "cwd": "/foundry", + "stdout": subprocess.PIPE, + "stderr": subprocess.STDOUT, + "text": True, + } + return subprocess.CompletedProcess(command, returncode=7, stdout="native failure\n") + + with pytest.raises(RuntimeError) as error: + harness.execute_native_suite("focused", runner=fail) + + event = json.loads(capsys.readouterr().out) + assert event == { + "command": ("python3 -m pytest tests/test_vmm_alloc.py tests/test_rng_graph.py -q"), + "event": "native_test_command", + "output": "native failure\n", + "returncode": 7, + } + failure = json.loads(str(error.value)) + assert failure["status"] == "failed" + assert failure["foundry_ref"] == FULL_SHA + assert failure["suite"] == "focused" + assert failure["results"] == [ + { + "command": event["command"], + "output": "native failure\n", + "returncode": 7, + } + ] + + +def test_orchestration_module_and_entrypoints_do_not_initialize_cuda() -> None: + source = HARNESS_PATH.read_text() + tree = ast.parse(source) + imported_roots = set() + for node in ast.walk(tree): + if isinstance(node, ast.Import): + imported_roots.update(alias.name.split(".", 1)[0] for alias in node.names) + elif isinstance(node, ast.ImportFrom) and node.module is not None: + imported_roots.add(node.module.split(".", 1)[0]) + + assert imported_roots.isdisjoint({"torch", "foundry"}) + functions = { + node.name: ast.get_source_segment(source, node) + for node in tree.body + if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef)) + } + for name in ("run_native_suite", "main"): + entrypoint = functions[name] + assert entrypoint is not None + assert "import torch" not in entrypoint + assert "import foundry" not in entrypoint + assert "cuda.init" not in entrypoint.lower() + assert "torch.cuda" not in entrypoint.lower() + assert "run_native_suite.remote(" in functions["main"] diff --git a/tests/test_modal_sglang_tp.py b/tests/test_modal_sglang_tp.py new file mode 100644 index 00000000..f50f0874 --- /dev/null +++ b/tests/test_modal_sglang_tp.py @@ -0,0 +1,1780 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the Foundry project +"""CPU-only contracts for the SGLang TP Modal harness.""" + +from __future__ import annotations + +import ast +import importlib.util +import json +import os +import subprocess +import sys +import urllib.error +from pathlib import Path + +import pytest + +REPO_ROOT = Path(__file__).parents[1] +HARNESS_PATH = REPO_ROOT / "tests" / "modal_sglang_tp.py" +VALIDATION_PATH = REPO_ROOT / "tests" / "sglang_tp_validation.py" +TP_RECIPE_PATH = REPO_ROOT / "recipe" / "experimental" / "serve_qwen3-8b_sglang_tp.sh" +PINNED_BATCH_DISPATCH_FIXTURE = ( + REPO_ROOT / "tests" / "fixtures" / "sglang_a25164_tokenizer_batch_dispatch.txt" +) +FULL_SHA = "0123456789abcdef0123456789abcdef01234567" + + +def _load_validation(): + assert VALIDATION_PATH.exists(), f"required validation module is missing: {VALIDATION_PATH}" + spec = importlib.util.spec_from_file_location("sglang_tp_validation_contract", VALIDATION_PATH) + assert spec is not None and spec.loader is not None + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + return module + + +def _run_harness(**updates: str) -> subprocess.CompletedProcess[str]: + env = os.environ.copy() + env.update( + { + "FOUNDRY_REF": FULL_SHA, + "SGLANG_BASE": FULL_SHA, + "SGL_MODEL_REVISION": FULL_SHA, + "TP_SIZE": "2", + } + ) + env.update(updates) + code = """ +import runpy +import sys +from types import ModuleType + +class BlockModalProvisioning(ModuleType): + def __getattr__(self, name): + raise AssertionError("Modal image construction reached") + +sys.modules["modal"] = BlockModalProvisioning("modal") +sys.path.insert(0, sys.argv[1]) +runpy.run_path(sys.argv[2], run_name="__main__") +""" + return subprocess.run( + [sys.executable, "-c", code, str(HARNESS_PATH.parent), str(HARNESS_PATH)], + cwd=REPO_ROOT, + env=env, + capture_output=True, + text=True, + ) + + +def _image_definition_source() -> str: + source = HARNESS_PATH.read_text() + tree = ast.parse(source) + for node in tree.body: + if isinstance(node, ast.Assign) and any( + isinstance(target, ast.Name) and target.id == "image" for target in node.targets + ): + image_source = ast.get_source_segment(source, node) + assert image_source is not None + return image_source + raise AssertionError("image assignment not found") + + +def _module_literal(name: str): + """Return the ast.literal_eval of a module-level ``name = `` assign.""" + source = HARNESS_PATH.read_text() + tree = ast.parse(source) + for node in tree.body: + if isinstance(node, ast.Assign) and any( + isinstance(target, ast.Name) and target.id == name for target in node.targets + ): + return ast.literal_eval(node.value) + raise AssertionError(f"module-level literal {name!r} not found") + + +def _function_source(name: str) -> str: + source = HARNESS_PATH.read_text() + tree = ast.parse(source) + for node in ast.walk(tree): + if isinstance(node, ast.FunctionDef) and node.name == name: + segment = ast.get_source_segment(source, node) + assert segment is not None + return segment + raise AssertionError(f"function {name!r} not found") + + +class _ImageRecorder: + """A fake Modal image that records hydration-relevant builder calls. + + Every unhandled builder method is chainable and returns the same recorder, + so the fully constructed image accumulates the same configuration. + """ + + def __init__(self, base: str) -> None: + self.base = base + self.environment: dict[str, str] = {} + self.local_python_sources: list[tuple[tuple[str, ...], bool]] = [] + + def env(self, values: dict[str, str]): + self.environment.update(values) + return self + + def add_local_python_source(self, *modules: str, copy: bool = False, **_kwargs): + self.local_python_sources.append((tuple(modules), copy)) + return self + + def __getattr__(self, _name: str): + def _chain(*_args, **_kwargs): + return self + + return _chain + + +def _fake_modal(): + modal_module = type(sys)("modal") + + class _ImageFactory: + @staticmethod + def from_registry(*_args, **_kwargs): + return _ImageRecorder("registry") + + @staticmethod + def debian_slim(*_args, **_kwargs): + return _ImageRecorder("debian_slim") + + class _App: + def __init__(self, *_args, **_kwargs): + pass + + def function(self, *_args, **_kwargs): + return lambda fn: fn + + def local_entrypoint(self, *_args, **_kwargs): + return lambda fn: fn + + class _Volume: + @staticmethod + def from_name(*_args, **_kwargs): + return object() + + modal_module.Image = _ImageFactory + modal_module.App = _App + modal_module.Volume = _Volume + modal_module.is_local = lambda: True + return modal_module + + +def _recorded_harness_namespace(path: Path, **env_updates: str) -> dict: + """Execute ``path`` at import time under a recording fake ``modal`` and + return its module namespace (image recorders included).""" + import runpy + + saved_modules = { + name: sys.modules.get(name) + for name in ( + "modal", + "modal_sglang_tp", + "modal_vllm_tp", + "sglang_tp_validation", + "vllm_tp_validation", + "tp_validation_helpers", + ) + } + saved_path = list(sys.path) + saved_env = os.environ.copy() + try: + sys.modules["modal"] = _fake_modal() + sys.path.insert(0, str(path.parent)) + os.environ.update( + { + "FOUNDRY_REF": FULL_SHA, + "SGLANG_BASE": FULL_SHA, + "SGL_MODEL_REVISION": FULL_SHA, + "TP_SIZE": "2", + } + ) + os.environ.update(env_updates) + return runpy.run_path(str(path), run_name="_image_recorder_contract") + finally: + sys.path[:] = saved_path + os.environ.clear() + os.environ.update(saved_env) + for name, module in saved_modules.items(): + if module is None: + sys.modules.pop(name, None) + else: + sys.modules[name] = module + + +def _attached_modules(image: _ImageRecorder) -> set[str]: + return {module for modules, _copy in image.local_python_sources for module in modules} + + +def _copy_flags(image: _ImageRecorder) -> set[bool]: + return {copy for _modules, copy in image.local_python_sources} + + +def test_images_include_remote_hydration_dependencies() -> None: + namespace = _recorded_harness_namespace(HARNESS_PATH) + + image = namespace["image"] + cleanup_image = namespace["cleanup_image"] + assert isinstance(image, _ImageRecorder), type(image) + assert isinstance(cleanup_image, _ImageRecorder), type(cleanup_image) + + for recorder in (image, cleanup_image): + attached = _attached_modules(recorder) + assert "sglang_tp_validation" in attached + assert "tp_validation_helpers" in attached + # Runtime mount (copy=False) so the entrypoint's sibling modules are on + # the container PYTHONPATH without baking a new image layer. + assert _copy_flags(recorder) == {False} + + assert cleanup_image.base == "debian_slim" + assert cleanup_image.environment == { + "FOUNDRY_REF": FULL_SHA, + "SGLANG_BASE": FULL_SHA, + "SGL_MODEL_REVISION": FULL_SHA, + "TP_SIZE": "2", + } + + cleanup_source = _function_source("cleanup_run") + assert "modal.Image.debian_slim()" not in cleanup_source + assert "gpu=" not in cleanup_source + + +def test_cleanup_image_is_lightweight_debian_slim() -> None: + source = HARNESS_PATH.read_text() + + # The cleanup image is a slim base (no CUDA/from_registry, no GPU) with the + # sibling helpers mounted so cleanup_run can hydrate. + assert "modal.Image.debian_slim()" in source + cleanup_decorator = _function_source("cleanup_run") + assert "image=cleanup_image" in source + assert cleanup_decorator + + +def test_cleanup_remote_module_hydration_does_not_require_git( + monkeypatch: pytest.MonkeyPatch, +) -> None: + def reject_git(*_args, **_kwargs): + raise AssertionError("remote cleanup hydration must not invoke git") + + monkeypatch.setattr(subprocess, "check_output", reject_git) + + namespace = _recorded_harness_namespace(HARNESS_PATH) + + assert namespace["FOUNDRY_REF"] == FULL_SHA + assert namespace["TP_SIZE"] == 2 + + +def test_commit_refs_must_be_full_immutable_shas() -> None: + validation = _load_validation() + + assert validation.require_commit_sha("FOUNDRY_REF", FULL_SHA.upper()) == FULL_SHA + for invalid in (None, "", "main", "abc123", "g" * 40): + with pytest.raises(ValueError, match="40-hex commit SHA"): + validation.require_commit_sha("FOUNDRY_REF", invalid) + + +def test_tp_size_must_be_exactly_two() -> None: + validation = _load_validation() + + assert validation.require_tp_size("2") == 2 + for invalid in (1, 3, "two"): + with pytest.raises(ValueError, match="TP_SIZE must be exactly 2"): + validation.require_tp_size(invalid) + + +@pytest.mark.parametrize("name", ["FOUNDRY_REF", "SGLANG_BASE", "SGL_MODEL_REVISION"]) +@pytest.mark.parametrize("value", ["main", "abc123", "g" * 40]) +def test_harness_rejects_mutable_or_partial_commit_overrides(name: str, value: str) -> None: + result = _run_harness(**{name: value}) + + assert result.returncode != 0 + assert f"{name} must be a full 40-hex commit SHA" in result.stderr + + +@pytest.mark.parametrize("tp_size", [1, 3]) +def test_harness_rejects_unvalidated_tp_sizes_before_provisioning(tp_size: int) -> None: + result = _run_harness(TP_SIZE=str(tp_size)) + + assert result.returncode != 0 + assert f"TP_SIZE must be exactly 2, got {tp_size}" in result.stderr + + +def test_image_cache_inputs_include_every_immutable_revision() -> None: + source = HARNESS_PATH.read_text() + image_source = _image_definition_source() + + assert 'return "main"' not in source + assert 'return require_commit_sha("local Foundry HEAD", revision)' in source + assert '"FOUNDRY_REF": FOUNDRY_REF' in image_source + assert '"SGLANG_BASE": SGLANG_BASE' in image_source + assert '"SGL_MODEL_REVISION": MODEL_REVISION' in image_source + assert "git -C /sglang checkout --detach {SGLANG_BASE}" in image_source + assert "git -C /foundry checkout --detach {FOUNDRY_REF}" in image_source + + +def test_phase_sequence_and_mode_mapping_add_second_fresh_load() -> None: + assert _module_literal("PHASES") == ("baseline", "save", "save2", "load", "load2") + assert _module_literal("LOAD_PHASES") == ("load", "load2") + assert _module_literal("QUERY_PHASES") == ("baseline", "save2", "load", "load2") + assert _module_literal("PHASE_MODES") == { + "baseline": None, + "save": "save", + "save2": "save", + "load": "load", + "load2": "load", + } + + +def test_modal_volume_failure_semantics_are_conservative() -> None: + source = HARNESS_PATH.read_text() + run_phase_source = _function_source("run_phase") + focused_source = _function_source("_run_focused_treatment") + main_source = _function_source("main") + + assert "partial/failed SAVE remains private" not in source + assert "auto-commit may retain partial or mutated unique-run artifacts" in run_phase_source + assert '"artifacts_may_be_partial_or_mutated"' in focused_source + assert '"artifacts_may_be_partial_or_mutated"' in main_source + + +def test_full_harness_requires_exactly_one_graph_per_rank() -> None: + source = HARNESS_PATH.read_text() + main_source = _function_source("main") + + assert _module_literal("ARCHIVE_PHASES") == ("save", "save2", "load", "load2") + assert _module_literal("FOUNDRY_PHASES") == ("save", "save2", "load", "load2") + assert "EXPECTED_NCCL_CHANNELS" not in source + assert 'expected_graph_counts = {"0": 1, "1": 1}' in main_source + assert "archive_graph_counts_exact" in main_source + assert "archive_graph_json_counts_exact" in main_source + assert "graphs_restored_each_rank" in main_source + # No hardcoded channel count or historical 36-graph expectation remains. + assert "range(24)" not in source + assert "== 36" not in source + + +def test_current_harness_reports_torch_symmetric_memory_collective_backend() -> None: + assert _module_literal("COLLECTIVE_BACKEND") == "torch_symmetric_memory" + assert '"collective_backend": COLLECTIVE_BACKEND' in _function_source("_candidate_metadata") + assert "**_candidate_metadata()" in _function_source("_run_focused_treatment") + + +def test_full_checks_keep_baseline_oracle_without_plain_nccl_acceptance() -> None: + source = HARNESS_PATH.read_text() + run_phase_source = _function_source("run_phase") + main_source = _function_source("main") + + # Exact output and concurrent-output acceptance remains baseline-derived. + assert "baseline_outputs" in main_source + assert "baseline_concurrent" in main_source + assert "deterministic_outputs_match_baseline" in main_source + assert "concurrent_outputs_match_baseline" in main_source + assert "graphs_restored_each_rank" in main_source + assert "load_offsets_match_save2" in main_source + # NCCL version/channel maps remain phase diagnostics, never pass/fail keys. + assert "nccl_transport_channels" in run_phase_source + assert "nccl_version" in run_phase_source + assert "nccl_p2p_ipc_every_phase" not in main_source + assert "channels_are_p2p_ipc" not in source + # Both fresh LOADs are checked. + assert "load2" in main_source + + +def test_full_archive_phases_collect_symmetric_and_graph_json_evidence() -> None: + run_phase_source = _function_source("run_phase") + main_source = _function_source("main") + + assert "if phase in ARCHIVE_PHASES:" in run_phase_source + assert "archive_graph_json_counts(" in run_phase_source + assert "archive_collective_kernel_counts(" in run_phase_source + assert "symmetric_archive_collectives_each_phase" in main_source + assert "no_archived_nccl_collectives" in main_source + + +def test_full_foundry_phases_require_symmetric_runtime_markers() -> None: + main_source = _function_source("main") + + assert "FOUNDRY_PHASES" in main_source + assert "symmetric_markers_each_phase" in main_source + assert "torch_symm_mem_ranks" in main_source + assert "multimem_all_gather_addresses" in main_source + + +def test_uses_shared_scan_and_stop_helpers() -> None: + source = HARNESS_PATH.read_text() + + assert "from sglang_tp_validation import" in source + assert "scan_log_text" in source + assert "stop_process" in source + assert "expected_process_shutdown" in source + + +def test_provenance_reports_immutable_revisions_and_versions() -> None: + source = HARNESS_PATH.read_text() + metadata_source = _function_source("_candidate_metadata") + + for field in ('"model"', '"model_revision"', '"tp_size"', '"attention_backend"'): + assert field in metadata_source + # Runtime / package versions gathered where available (no CUDA init in parent). + version_source = _function_source("_runtime_versions") + for token in ("driver", "torch", "sglang", "flashinfer", "sgl", "nccl"): + assert token in version_source + # The report still carries the immutable revisions. + assert '"foundry_ref": FOUNDRY_REF' in source + assert '"sglang_base": SGLANG_BASE' in source + assert "MODEL_REVISION" in source + + +def test_query_phases_produce_outputs_for_baseline_save2_and_both_loads() -> None: + run_phase_source = _function_source("run_phase") + + assert "QUERY_PHASES" in run_phase_source + assert "outputs" in run_phase_source + assert "concurrent_outputs" in run_phase_source + + +def test_load_phases_receive_and_verify_save2_archive_and_never_commit() -> None: + run_phase_source = _function_source("run_phase") + main_source = _function_source("main") + + # run_phase accepts the SAVE2 expected inventory/fingerprints for verification. + signature = run_phase_source.splitlines()[0] + assert "expected_fingerprints" in run_phase_source + assert "expected_graph_counts" in run_phase_source + # LOAD phases re-read the on-disk archive and prove it is byte-for-byte SAVE2. + assert "archive_fingerprints" in run_phase_source + assert "archive_graph_counts" in run_phase_source + # The archive volume must never be committed for a LOAD phase (a mutating or + # failed LOAD cannot persist into LOAD2). + assert 'if phase in {"save", "save2"} and not failures:' in run_phase_source + assert run_phase_source.count("archive_volume.commit()") == 1 + assert run_phase_source.index("failures = []") < run_phase_source.index( + "archive_volume.commit()" + ) + # main threads SAVE2 expectations into each fresh LOAD remote call. + assert "expected_fingerprints=" in main_source + assert "expected_graph_counts=" in main_source + assert signature # signature captured for clarity; presence asserted above + + +def test_final_checks_require_both_loads_preserve_save2_archive() -> None: + main_source = _function_source("main") + + assert "load_archive_immutable" in main_source + # The immutability check compares against the SAVE2 archive, for both loads. + assert "save2_fingerprints" in main_source + assert "save2_graph_counts" in main_source + + +def test_byte_level_archive_fingerprints_are_recorded_and_verified() -> None: + source = HARNESS_PATH.read_text() + run_phase_source = _function_source("run_phase") + main_source = _function_source("main") + + # A byte-level recursive fingerprint helper backs the harness. + assert "archive_byte_fingerprints" in source + # run_phase records byte fingerprints and verifies LOAD/LOAD2 against SAVE2. + assert "expected_byte_fingerprints" in run_phase_source + assert "archive_byte_fingerprints" in run_phase_source + # main threads SAVE2's byte fingerprints into each fresh LOAD. + assert "expected_byte_fingerprints=" in main_source + # Final check requires both loads' full byte maps to equal SAVE2. + assert "load_archive_bytes_immutable" in main_source + assert "save2_byte_fingerprints" in main_source + + +def test_concurrent_outputs_compared_exactly_across_phases() -> None: + main_source = _function_source("main") + + # bs>1 concurrent evidence is compared for exact equality with the baseline, + # not merely required to be nonempty. + assert "concurrent_outputs_match_baseline" in main_source + assert "concurrent_outputs_nonempty" not in main_source + # Restored-graph replay at bs>1 is still required. + assert "restored_graph_replay_observed" in main_source + + +def test_concurrent_oracle_uses_eight_identical_factual_prompts_and_tight_bound() -> None: + namespace = _recorded_harness_namespace(HARNESS_PATH) + factual_prompt = "Reply with only the result: 2 + 2 =" + + assert namespace["CONCURRENT_PROMPTS"] == [factual_prompt] * 8 + assert len(namespace["CONCURRENT_PROMPTS"]) == 8 + assert len(set(namespace["CONCURRENT_PROMPTS"])) == 1 + assert namespace["CONCURRENT_MAX_NEW_TOKENS"] == 4 + # Sequential coverage and its longer response bound stay unchanged. + assert namespace["PROMPTS"] == [ + "Explain what tensor parallelism is in one paragraph.", + "Write a haiku about GPUs.", + "List three prime numbers greater than 50.", + "Summarize why CUDA graphs speed up inference.", + ] + assert namespace["MAX_NEW_TOKENS"] == 64 + + +def test_concurrent_generation_uses_one_native_batch_with_exact_shared_payload( + monkeypatch: pytest.MonkeyPatch, +) -> None: + namespace = _recorded_harness_namespace(HARNESS_PATH) + globals_dict = namespace["_generate_concurrently"].__globals__ + expected_prompts = namespace["CONCURRENT_PROMPTS"] + expected_outputs = [f"batch-output-{index}" for index in range(8)] + requests: list[tuple[str, dict, float]] = [] + + assert len(expected_prompts) == 8 + assert expected_prompts == [expected_prompts[0]] * 8 + + class Response: + def __enter__(self): + return self + + def __exit__(self, *_args): + return None + + def read(self) -> bytes: + return json.dumps([{"text": output} for output in expected_outputs]).encode() + + def urlopen(request, timeout: float): + requests.append((request.full_url, json.loads(request.data), timeout)) + return Response() + + monkeypatch.setattr(globals_dict["urllib"].request, "urlopen", urlopen) + + assert namespace["_generate_concurrently"]() == expected_outputs + assert requests == [ + ( + f"http://127.0.0.1:{namespace['PORT']}/generate", + { + "text": expected_prompts, + "sampling_params": { + "temperature": 0.0, + "max_new_tokens": 4, + }, + }, + 120, + ) + ] + + +@pytest.mark.parametrize( + ("response_body", "message"), + [ + ({"text": ["output"] * 8}, "must be a list"), + ([{"text": "output"}] * 7, "exactly 8"), + ([{"text": "output"}] * 7 + [{"text": 7}], "text string"), + ], +) +def test_concurrent_native_batch_rejects_malformed_or_wrong_length_responses( + monkeypatch: pytest.MonkeyPatch, + response_body, + message: str, +) -> None: + namespace = _recorded_harness_namespace(HARNESS_PATH) + globals_dict = namespace["_generate_concurrently"].__globals__ + + class Response: + def __enter__(self): + return self + + def __exit__(self, *_args): + return None + + def read(self) -> bytes: + return json.dumps(response_body).encode() + + monkeypatch.setattr( + globals_dict["urllib"].request, + "urlopen", + lambda *_args, **_kwargs: Response(), + ) + + with pytest.raises(ValueError, match=message): + namespace["_generate_concurrently"]() + + +def test_native_batch_removes_timing_based_concurrency_helpers() -> None: + source = HARNESS_PATH.read_text() + batch_source = _function_source("_generate_concurrently") + + assert "import threading" not in source + assert "ThreadPoolExecutor" not in source + assert "Barrier" not in batch_source + + +def test_native_list_request_and_recipe_reach_pinned_batch_dispatcher() -> None: + harness_source = HARNESS_PATH.read_text() + batch_source = _function_source("_generate_concurrently") + recipe_source = TP_RECIPE_PATH.read_text() + pinned_source = PINNED_BATCH_DISPATCH_FIXTURE.read_text() + + assert "a25164bda355250015c0ba6d8821f0d4b37abe7d" in harness_source + assert '"text": CONCURRENT_PROMPTS' in batch_source + assert len(_recorded_harness_namespace(HARNESS_PATH)["CONCURRENT_PROMPTS"]) == 8 + assert "--enable-tokenizer-batch-encode" in recipe_source + assert "--enable-dynamic-batch-tokenizer" not in recipe_source + + assert "commit: a25164bda355250015c0ba6d8821f0d4b37abe7d" in pinned_source + assert "self.server_args.enable_tokenizer_batch_encode" in pinned_source + assert "if self._should_use_batch_tokenization(batch_size, obj):" in pinned_source + assert "self._send_batch_request(tokenized_objs)" in pinned_source + assert "batch_req = BatchTokenizedGenerateReqInput(batch=tokenized_objs)" in pinned_source + assert "self._dispatch_to_scheduler(batch_req)" in pinned_source + assert "self._send_one_request(tokenized_obj)" in pinned_source + + +def test_nccl_channels_and_versions_are_diagnostic_only() -> None: + source = HARNESS_PATH.read_text() + main_source = _function_source("main") + + assert "channels_are_p2p_ipc" not in source + assert "nccl_p2p_ipc_every_phase" not in main_source + assert '"nccl_transport_channels"' in source + assert '"nccl_version"' in source + + +def test_run_phase_uses_group_aware_teardown() -> None: + source = HARNESS_PATH.read_text() + + # The harness relies on the shared group-draining stop_process helper. + assert "stop_process" in source + + +def test_run_phase_two_window_scan_before_and_after_teardown() -> None: + run_phase_source = _function_source("run_phase") + + # A pre-teardown snapshot is captured before the intentional SIGINT teardown, + # so the SIGINT KeyboardInterrupt tracebacks are excluded from the pre window. + assert "pre_teardown_text" in run_phase_source + snapshot_index = run_phase_source.index("pre_teardown_text") + stop_index = run_phase_source.index("stop_process(") + assert snapshot_index < stop_index, "pre-teardown snapshot must precede stop_process" + + # The full log is scanned after teardown and both windows are merged via the + # shared helper (all other parsed signals still come from the full log). + scan_index = run_phase_source.index("scan_log_text(") + assert stop_index < scan_index, "full-log scan must run after teardown" + assert "pre_teardown_text=pre_teardown_text" in run_phase_source + + +def test_focused_mode_has_exact_two_phase_query_contract() -> None: + assert _module_literal("FOCUSED_PHASES") == ("save", "load") + assert _module_literal("FOCUSED_MAX_NEW_TOKENS") == 32 + assert isinstance(_module_literal("FOCUSED_PROMPT"), str) + + generate_source = _function_source("_generate") + run_phase_source = _function_source("run_phase") + + assert "max_new_tokens: int = MAX_NEW_TOKENS" in generate_source + assert '"max_new_tokens": max_new_tokens' in generate_source + assert "focused_treatment: bool = False" in run_phase_source + assert "_generate(FOCUSED_PROMPT, max_new_tokens=FOCUSED_MAX_NEW_TOKENS)" in run_phase_source + assert "if focused_treatment:" in run_phase_source + assert "else:" in run_phase_source + assert "_generate_concurrently()" in run_phase_source + + +def test_focused_treatment_threads_save_archive_evidence_into_fresh_load() -> None: + focused_source = _function_source("_run_focused_treatment") + + assert "for phase in FOCUSED_PHASES" in focused_source + assert "focused_treatment=True" in focused_source + assert "expected_fingerprints=save_result" in focused_source + assert "expected_graph_counts=save_result" in focused_source + assert "expected_byte_fingerprints=save_result" in focused_source + for forbidden in ("baseline", "save2", "load2", "_generate_concurrently"): + assert forbidden not in focused_source + + +def test_focused_treatment_requires_every_planned_acceptance_signal() -> None: + source = HARNESS_PATH.read_text() + validation_source = VALIDATION_PATH.read_text() + run_phase_source = _function_source("run_phase") + focused_source = _function_source("_run_focused_treatment") + + assert "archive_collective_kernel_counts" in source + assert "archive_graph_json_counts" in source + assert "archive_collective_kernel_counts" in run_phase_source + assert "archive_graph_json_counts" in run_phase_source + assert "focused_save_checks" in source + assert "focused_treatment_checks" in source + assert "classify_focused_checks" in source + assert "focused_save_checks(" in focused_source + assert "focused_treatment_checks(" in focused_source + assert "classify_focused_checks(" in focused_source + assert "symmetric_archive_collectives_present" in validation_source + assert "symmetric_archive_collectives_each_phase" in validation_source + assert "multimem_all_gather_addresses" in validation_source + + +@pytest.mark.parametrize( + ("failure_text", "expected"), + [ + ("/foundry/python/module.py: ConnectionError: setup unavailable", False), + ("CUDA out of memory while loading model weights", False), + ("gpu [0] NCCL WARN transient bootstrap failure", False), + ("[foundry BUILD] Template 0 (t256): building...", False), + ("[foundry] WARNING: optional binary graph write failed", False), + ("[foundry LOAD ERROR] graph node restoration failed", True), + ("[HOOK] FATAL ERROR: fixed-address allocation failed", True), + ("[REPLAY] ERROR: graph launch failed", True), + ("Torch symmetric-memory communicator failed", True), + ("Multimem all-gather used a fallback", True), + ("archive immutability: LOAD mutated archive bytes", True), + ], +) +def test_symmetric_policy_failure_detection_ignores_unrelated_setup_paths( + failure_text: str, + expected: bool, +) -> None: + validation = _load_validation() + + assert validation.is_symmetric_policy_failure(failure_text) is expected + + +def test_focused_report_prints_explicit_status_before_failed_checks_raise() -> None: + focused_source = _function_source("_run_focused_treatment") + main_source = _function_source("main") + + assert 'report["status"] = status' in focused_source + assert "print(json.dumps(report" in focused_source + assert focused_source.index("print(json.dumps(report") < focused_source.rindex( + "raise RuntimeError" + ) + assert "def main(focused: bool = False)" in main_source + assert "if focused:" in main_source + assert "_run_focused_treatment(run_id)" in main_source + + +class _FakeProcess: + def __init__(self) -> None: + self.returncode = 0 + + def poll(self): + return None + + +class _EarlyExitProcess: + def __init__(self, returncode: int = 17) -> None: + self.returncode = returncode + self._poll_results = iter((None, returncode)) + + def poll(self): + return next(self._poll_results, self.returncode) + + +class _RecordingVolume: + def __init__(self, events: list[str]) -> None: + self.events = events + + def commit(self) -> None: + self.events.append("commit") + + +def _focused_scan(phase: str) -> dict: + return { + "errors": [], + "error_categories": [], + "python_error_types": [], + "untyped_python_fatal_evidence": [], + "typed_python_traceback_headers": [], + "nccl_version": "2.28.3", + "load_offsets": {"0": 4096, "1": 4096} if phase == "load" else {}, + "loaded_graphs": {"0": 1, "1": 1} if phase == "load" else {}, + "torch_symm_mem_ranks": ["0", "1"], + "multimem_all_gather_ranks": ["0", "1"], + "torch_symm_mem_addresses": {"0": ["0x1000"], "1": ["0x2000"]}, + "multimem_all_gather_addresses": {"0": ["0x3000"], "1": ["0x4000"]}, + "nccl_transport_channels": {}, + "graph_replay_batch_sizes": [1] if phase == "load" else [], + "native_capture_progress_count": 0, + "saved_graph_log_count": 0, + } + + +def _install_run_phase_fakes( + namespace: dict, + tmp_path: Path, + events: list[str], + *, + phase: str, + startup_error: Exception | None = None, + phase_error: bool = False, + teardown_error: bool = False, + fatal_error: bool = False, + archive_error: bool = False, + invalid_inventory: bool = False, + missing_marker: bool = False, + missing_symmetric_collective: bool = False, + patch_archive: bool = True, +) -> _RecordingVolume: + globals_dict = namespace["run_phase"].__globals__ + volume = _RecordingVolume(events) + globals_dict["RUNS_ROOT"] = str(tmp_path) + globals_dict["archive_volume"] = volume + + def wait_until_healthy(*_args, **_kwargs): + if startup_error is not None: + raise startup_error + + globals_dict["_wait_until_healthy"] = wait_until_healthy + + def generate(*_args, **_kwargs): + if phase_error: + raise RuntimeError("request failed after health") + return "deterministic output" + + def launch(_mode, log_path, _run_root): + log_file = Path(log_path).open("w") # noqa: SIM115 - spans fake process lifetime + log_file.write("healthy\n") + log_file.flush() + return _FakeProcess(), log_file + + def stop_process(_process, log_file): + log_file.close() + if teardown_error: + raise RuntimeError("teardown failed") + return 0 + + def scan_log_text(*_args, **_kwargs): + scan = _focused_scan(phase) + if fatal_error: + scan["errors"] = ["[foundry LOAD ERROR] graph restore failed"] + scan["error_categories"] = ["hard_native", "policy"] + if startup_error is not None: + error_type = type(startup_error).__name__ + scan["errors"] = [f"{error_type}: startup failed"] + scan["python_error_types"] = [error_type] + if missing_marker: + scan["torch_symm_mem_ranks"] = ["0"] + return scan + + def collective_counts(*_args, **_kwargs): + events.append("collective_evidence") + if archive_error: + raise ValueError("malformed graph JSON") + return { + "0": { + "all_reduce": 0, + "all_gather": 0, + "symmetric_all_reduce": 0 if missing_symmetric_collective else 3, + "symmetric_all_gather": 1, + }, + "1": { + "all_reduce": 0, + "all_gather": 0, + "symmetric_all_reduce": 3, + "symmetric_all_gather": 1, + }, + } + + globals_dict["_generate"] = generate + globals_dict["_launch"] = launch + globals_dict["stop_process"] = stop_process + globals_dict["scan_log_text"] = scan_log_text + globals_dict["_runtime_versions"] = lambda _version: {} + if patch_archive: + globals_dict["_read_rank_offsets"] = lambda _workspace: { + "0": 4096, + "1": 4096, + } + globals_dict["_archive_graph_counts"] = lambda _workspace: { + "0": 1, + "1": 1, + } + globals_dict["archive_graph_json_counts"] = lambda *_args: ( + events.append("json_inventory") + or ({"0": 0, "1": 1} if invalid_inventory else {"0": 1, "1": 1}) + ) + globals_dict["_archive_fingerprints"] = lambda _workspace: { + "0": {"graph_0.json": "a"}, + "1": {"graph_0.json": "b"}, + } + globals_dict["_archive_byte_fingerprints"] = lambda _workspace: ( + events.append("archive_evidence") + or {"rank_0/graph_0.json": "a", "rank_1/graph_0.json": "b"} + ) + globals_dict["archive_collective_kernel_counts"] = collective_counts + return volume + + +def test_production_health_wait_preserves_last_exception_on_early_exit( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + namespace = _recorded_harness_namespace(HARNESS_PATH) + process = _EarlyExitProcess() + log_path = tmp_path / "startup.log" + log_path.write_text("ConnectionError: model download failed\n") + connection_error = ConnectionError("health connection unavailable") + health_error = urllib.error.URLError(connection_error) + + def fail_health_request(*_args, **_kwargs): + raise health_error + + monkeypatch.setattr(namespace["urllib"].request, "urlopen", fail_health_request) + monkeypatch.setattr(namespace["time"], "sleep", lambda _seconds: None) + + with pytest.raises(namespace["FocusedStartupHealthError"]) as raised: + namespace["_wait_until_healthy"]( + process, + str(log_path), + timeout=30, + ) + + assert raised.value.__cause__ is health_error + assert raised.value.cause_types == ( + "urllib.error.URLError", + "builtins.ConnectionError", + ) + assert raised.value.child_exited is True + + +def test_actual_health_wrapper_allows_proven_network_early_exit_only( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, + capsys: pytest.CaptureFixture[str], +) -> None: + namespace = _recorded_harness_namespace(HARNESS_PATH) + production_wait = namespace["_wait_until_healthy"] + production_scan = namespace["scan_log_text"] + events: list[str] = [] + _install_run_phase_fakes( + namespace, + tmp_path, + events, + phase="save", + teardown_error=True, + patch_archive=False, + ) + globals_dict = namespace["run_phase"].__globals__ + globals_dict["_wait_until_healthy"] = production_wait + globals_dict["scan_log_text"] = production_scan + + def launch(_mode, log_path, _run_root): + log_file = Path(log_path).open("w") # noqa: SIM115 - spans fake process lifetime + log_file.write("ConnectionError: model download failed\n") + log_file.flush() + return _EarlyExitProcess(), log_file + + def fail_health_request(*_args, **_kwargs): + raise urllib.error.URLError(ConnectionError("health connection unavailable")) + + globals_dict["_launch"] = launch + monkeypatch.setattr(globals_dict["urllib"].request, "urlopen", fail_health_request) + monkeypatch.setattr(globals_dict["time"], "sleep", lambda _seconds: None) + + result = namespace["run_phase"]( + "save", + "production-startup-failure", + focused_treatment=True, + ) + + assert result["failure_stage"] == "startup" + assert result["failure_type"].endswith("FocusedStartupHealthError") + assert result["startup_cause_types"] == [ + "urllib.error.URLError", + "builtins.ConnectionError", + ] + assert result["startup_child_exited"] is True + assert result["python_error_types"] == ["ConnectionError"] + assert result["evidence_errors"] == [] + assert result["archive_committed"] is False + assert "commit" not in events + capsys.readouterr() + + remote = _SaveOnlyRemoteFunction(result) + globals_dict["run_phase"] = remote + globals_dict["KEEP_ARTIFACTS"] = False + + with pytest.raises(RuntimeError, match="INCONCLUSIVE"): + namespace["_run_focused_treatment"]("production-startup-failure") + + report = json.loads(capsys.readouterr().out) + assert remote.calls == ["save"] + assert report["status"] == "INCONCLUSIVE" + assert report["artifacts_retained"] is True + assert report["artifacts_may_be_partial_or_mutated"] is True + + +@pytest.mark.parametrize( + ("failure_kind", "expected_stage"), + [ + ("phase", "request"), + ("teardown", "teardown"), + ("fatal", "scanner"), + ("archive", "archive"), + ("inventory", "archive"), + ("marker", "scanner"), + ("symmetric", "archive"), + ], +) +def test_focused_save_never_commits_partial_archive( + tmp_path: Path, + failure_kind: str, + expected_stage: str, +) -> None: + namespace = _recorded_harness_namespace(HARNESS_PATH) + events: list[str] = [] + _install_run_phase_fakes( + namespace, + tmp_path, + events, + phase="save", + phase_error=failure_kind == "phase", + teardown_error=failure_kind == "teardown", + fatal_error=failure_kind == "fatal", + archive_error=failure_kind == "archive", + invalid_inventory=failure_kind == "inventory", + missing_marker=failure_kind == "marker", + missing_symmetric_collective=failure_kind == "symmetric", + ) + + result = namespace["run_phase"]( + "save", + "transactional-save", + focused_treatment=True, + ) + + assert result["phase_failures"] or result["evidence_errors"] + assert result["failure_stage"] == expected_stage + assert result["archive_committed"] is False + assert "commit" not in events + + +def test_focused_save_commits_only_after_all_evidence_succeeds(tmp_path: Path) -> None: + namespace = _recorded_harness_namespace(HARNESS_PATH) + events: list[str] = [] + _install_run_phase_fakes(namespace, tmp_path, events, phase="save") + + result = namespace["run_phase"]( + "save", + "complete-save", + focused_treatment=True, + ) + + assert result["phase_failures"] == [] + assert result["archive_committed"] is True + assert result["archive_graph_json_counts"] == {"0": 1, "1": 1} + assert events.count("commit") == 1 + assert events.index("json_inventory") < events.index("commit") + assert events.index("collective_evidence") < events.index("commit") + assert events.index("archive_evidence") < events.index("commit") + + +def test_focused_load_never_commits_archive(tmp_path: Path) -> None: + namespace = _recorded_harness_namespace(HARNESS_PATH) + events: list[str] = [] + _install_run_phase_fakes(namespace, tmp_path, events, phase="load") + + result = namespace["run_phase"]( + "load", + "immutable-load", + expected_fingerprints={ + "0": {"graph_0.json": "a"}, + "1": {"graph_0.json": "b"}, + }, + expected_graph_counts={"0": 1, "1": 1}, + expected_byte_fingerprints={ + "rank_0/graph_0.json": "a", + "rank_1/graph_0.json": "b", + }, + focused_treatment=True, + ) + + assert result["archive_committed"] is False + assert "commit" not in events + + +def _complete_focused_results() -> tuple[dict, dict]: + archive_fingerprints = { + "0": {"graph_0.json": "a"}, + "1": {"graph_0.json": "b"}, + } + byte_fingerprints = { + "rank_0/graph_0.json": "a", + "rank_1/graph_0.json": "b", + } + common = { + "became_healthy": True, + "request_completed": True, + "outputs": ["deterministic output"], + "archive_graph_counts": {"0": 1, "1": 1}, + "archive_graph_json_counts": {"0": 1, "1": 1}, + "archive_collective_kernel_counts": { + "0": { + "all_reduce": 0, + "all_gather": 0, + "symmetric_all_reduce": 3, + "symmetric_all_gather": 1, + }, + "1": { + "all_reduce": 0, + "all_gather": 0, + "symmetric_all_reduce": 3, + "symmetric_all_gather": 1, + }, + }, + "archive_fingerprints": archive_fingerprints, + "archive_byte_fingerprints": byte_fingerprints, + "torch_symm_mem_ranks": ["0", "1"], + "multimem_all_gather_ranks": ["1"], + "torch_symm_mem_addresses": {"0": ["0x1000"], "1": ["0x2000"]}, + "multimem_all_gather_addresses": {"1": ["0x3000"]}, + "errors": [], + "evidence_errors": [], + "phase_failures": [], + "process_return_code": 0, + "archive_committed": False, + "native_capture_progress_count": 0, + "saved_graph_log_count": 0, + } + save = { + **common, + "phase": "save", + "rank_offsets": {"0": 4096, "1": 4096}, + "archive_committed": True, + } + load = { + **common, + "phase": "load", + "load_offsets": {"0": 4096, "1": 4096}, + "loaded_graphs": {"0": 1, "1": 1}, + "graph_replay_batch_sizes": [1], + } + return save, load + + +def _complete_full_results() -> dict[str, dict]: + outputs = [f"output-{index}" for index in range(4)] + concurrent_outputs = [f"concurrent-{index}" for index in range(8)] + channels = { + "0": {"P2P/IPC": [0, 1]}, + "1": {"P2P/IPC": [0, 1]}, + } + fingerprints = { + rank: { + "graph_manifest.json": "manifest", + "fatbin_image_packed.img": f"fatbin-{rank}", + "final_alloc_offset.json": "offset", + "graph_0.json": f"graph-{rank}", + } + for rank in ("0", "1") + } + byte_fingerprints = { + "rank_0/graph_0.json": "graph-0", + "rank_1/graph_0.json": "graph-1", + } + collective_counts = { + rank: { + "all_reduce": 0, + "all_gather": 0, + "symmetric_all_reduce": 3, + "symmetric_all_gather": 1, + } + for rank in ("0", "1") + } + common = { + "outputs": outputs, + "concurrent_outputs": concurrent_outputs, + "nccl_transport_channels": channels, + "nccl_version": "2.28.9", + "errors": [], + "process_return_code": 0, + "graph_replay_batch_sizes": [], + "saved_graph_log_count": 0, + "loaded_graphs": {}, + "load_offsets": {}, + "archive_fingerprints": {}, + "archive_graph_counts": {}, + "archive_graph_json_counts": {}, + "archive_collective_kernel_counts": {}, + "archive_byte_fingerprints": {}, + "rank_offsets": {}, + "torch_symm_mem_ranks": [], + "multimem_all_gather_addresses": {}, + } + results = { + phase: {**common, "phase": phase} + for phase in ("baseline", "save", "save2", "load", "load2") + } + for phase in ("save", "save2", "load", "load2"): + results[phase]["archive_fingerprints"] = fingerprints + results[phase]["archive_graph_counts"] = {"0": 1, "1": 1} + results[phase]["archive_graph_json_counts"] = {"0": 1, "1": 1} + results[phase]["archive_collective_kernel_counts"] = collective_counts + results[phase]["archive_byte_fingerprints"] = byte_fingerprints + results[phase]["torch_symm_mem_ranks"] = ["0", "1"] + results[phase]["multimem_all_gather_addresses"] = {"1": ["0x3000"]} + for phase in ("save", "save2"): + results[phase]["rank_offsets"] = {"0": 4096, "1": 4096} + for phase in ("load", "load2"): + results[phase]["load_offsets"] = {"0": 4096, "1": 4096} + results[phase]["loaded_graphs"] = {"0": 1, "1": 1} + results[phase]["graph_replay_batch_sizes"] = [1, 2] + return results + + +class _FakeRemoteFunction: + def __init__(self, save_result: dict, load_result: dict) -> None: + self.save_result = save_result + self.load_result = load_result + self.calls: list[str] = [] + + def remote(self, phase: str, *_args, **_kwargs) -> dict: + self.calls.append(phase) + return self.save_result if phase == "save" else self.load_result + + +class _FullRemoteFunction: + def __init__(self, results: dict[str, dict]) -> None: + self.results = results + self.calls: list[str] = [] + + def remote(self, phase: str, *_args, **_kwargs) -> dict: + self.calls.append(phase) + return self.results[phase] + + +class _RaisingRemoteFunction: + def __init__(self, error: Exception) -> None: + self.error = error + + def remote(self, *_args, **_kwargs): + raise self.error + + +class _CleanupRemoteFunction: + def __init__( + self, + capsys: pytest.CaptureFixture[str], + error: Exception | None = None, + ) -> None: + self.capsys = capsys + self.error = error + self.calls: list[str] = [] + + def remote(self, run_id: str) -> None: + assert self.capsys.readouterr().out == "" + self.calls.append(run_id) + if self.error is not None: + raise self.error + + +class _RemoteAdapter: + def __init__(self, function) -> None: + self.function = function + + def remote(self, *args, **kwargs): + return self.function(*args, **kwargs) + + +class ServiceError(Exception): + pass + + +ServiceError.__module__ = "modal.exception" + + +@pytest.mark.parametrize( + ("error", "expected_status"), + [ + (ServiceError("service unavailable"), "INCONCLUSIVE"), + (AttributeError("harness bug"), "NOT VERIFIED"), + ], +) +def test_focused_remote_errors_distinguish_environment_from_harness_failure( + capsys: pytest.CaptureFixture[str], + error: Exception, + expected_status: str, +) -> None: + namespace = _recorded_harness_namespace(HARNESS_PATH) + globals_dict = namespace["_run_focused_treatment"].__globals__ + globals_dict["run_phase"] = _RaisingRemoteFunction(error) + globals_dict["KEEP_ARTIFACTS"] = False + + with pytest.raises(RuntimeError, match=expected_status): + namespace["_run_focused_treatment"]("remote-error") + + report = json.loads(capsys.readouterr().out) + assert report["status"] == expected_status + assert report["artifacts_retained"] is True + assert report["artifacts_may_be_partial_or_mutated"] is True + + +class _SaveOnlyRemoteFunction: + def __init__(self, save_result: dict) -> None: + self.save_result = save_result + self.calls: list[str] = [] + + def remote(self, phase: str, *_args, **_kwargs) -> dict: + self.calls.append(phase) + if phase != "save": + raise AssertionError("LOAD must not run after failed startup") + return self.save_result + + +@pytest.mark.parametrize( + ("startup_error", "expected_status"), + [ + (ConnectionError("network unavailable"), "INCONCLUSIVE"), + (AttributeError("startup harness bug"), "NOT VERIFIED"), + ], +) +def test_actual_startup_failure_skips_archive_and_stops_before_load( + tmp_path: Path, + capsys: pytest.CaptureFixture[str], + startup_error: Exception, + expected_status: str, +) -> None: + namespace = _recorded_harness_namespace(HARNESS_PATH) + events: list[str] = [] + _install_run_phase_fakes( + namespace, + tmp_path, + events, + phase="save", + startup_error=startup_error, + patch_archive=False, + ) + actual_run_phase = namespace["run_phase"] + + result = actual_run_phase( + "save", + "startup-failure", + focused_treatment=True, + ) + + assert result["failure_stage"] == "startup" + assert result["failure_type"].endswith(type(startup_error).__name__) + assert result["evidence_errors"] == [] + assert result["archive_committed"] is False + assert result["python_error_types"] == [type(startup_error).__name__] + assert "commit" not in events + capsys.readouterr() + + remote = _SaveOnlyRemoteFunction(result) + globals_dict = namespace["_run_focused_treatment"].__globals__ + globals_dict["run_phase"] = remote + globals_dict["KEEP_ARTIFACTS"] = False + + with pytest.raises(RuntimeError, match=expected_status): + namespace["_run_focused_treatment"]("startup-failure") + + report = json.loads(capsys.readouterr().out) + assert remote.calls == ["save"] + assert report["status"] == expected_status + assert report["artifacts_retained"] is True + assert report["artifacts_may_be_partial_or_mutated"] is True + + +@pytest.mark.parametrize("keep_artifacts", [False, True]) +def test_focused_orchestration_rejects_uncommitted_save_with_possible_auto_committed_artifacts( + capsys: pytest.CaptureFixture[str], + keep_artifacts: bool, +) -> None: + namespace = _recorded_harness_namespace(HARNESS_PATH) + save, load = _complete_focused_results() + save["archive_graph_json_counts"] = {"0": 0, "1": 1} + save["archive_committed"] = False + remote = _FakeRemoteFunction(save, load) + cleanup_calls: list[str] = [] + globals_dict = namespace["_run_focused_treatment"].__globals__ + globals_dict["run_phase"] = remote + globals_dict["cleanup_run"] = type( + "Cleanup", + (), + {"remote": staticmethod(cleanup_calls.append)}, + )() + globals_dict["KEEP_ARTIFACTS"] = keep_artifacts + + with pytest.raises(RuntimeError, match="NOT VERIFIED"): + namespace["_run_focused_treatment"]("bad-save") + + report = json.loads(capsys.readouterr().out) + assert remote.calls == ["save"] + assert report["status"] == "NOT VERIFIED" + assert report["artifacts_retained"] is True + assert report["artifacts_may_be_partial_or_mutated"] is True + assert cleanup_calls == [] + + +def test_focused_load_failure_reports_committed_save_artifact_retained( + capsys: pytest.CaptureFixture[str], +) -> None: + namespace = _recorded_harness_namespace(HARNESS_PATH) + save, load = _complete_focused_results() + load["outputs"] = ["different output"] + remote = _FakeRemoteFunction(save, load) + cleanup_calls: list[str] = [] + globals_dict = namespace["_run_focused_treatment"].__globals__ + globals_dict["run_phase"] = remote + globals_dict["cleanup_run"] = type( + "Cleanup", + (), + {"remote": staticmethod(cleanup_calls.append)}, + )() + globals_dict["KEEP_ARTIFACTS"] = False + + with pytest.raises(RuntimeError, match="NOT VERIFIED"): + namespace["_run_focused_treatment"]("load-failure") + + report = json.loads(capsys.readouterr().out) + assert remote.calls == ["save", "load"] + assert report["artifacts_retained"] is True + assert cleanup_calls == [] + + +@pytest.mark.parametrize( + ("keep_artifacts", "expected_retained", "cleanup_count"), + [(False, False, 1), (True, True, 0)], +) +def test_focused_success_reports_actual_artifact_retention( + capsys: pytest.CaptureFixture[str], + keep_artifacts: bool, + expected_retained: bool, + cleanup_count: int, +) -> None: + namespace = _recorded_harness_namespace(HARNESS_PATH) + save, load = _complete_focused_results() + remote = _FakeRemoteFunction(save, load) + cleanup = _CleanupRemoteFunction(capsys) + globals_dict = namespace["_run_focused_treatment"].__globals__ + globals_dict["run_phase"] = remote + globals_dict["cleanup_run"] = cleanup + globals_dict["KEEP_ARTIFACTS"] = keep_artifacts + + namespace["_run_focused_treatment"]("good-run") + + report = json.loads(capsys.readouterr().out) + assert remote.calls == ["save", "load"] + assert report["status"] == "VERIFIED" + assert report["artifacts_retained"] is expected_retained + assert report["artifacts_may_be_partial_or_mutated"] is False + assert report["cleanup_attempted"] is (not keep_artifacts) + assert report["cleanup_completed"] is (not keep_artifacts) + assert report["cleanup_skipped"] is keep_artifacts + assert len(cleanup.calls) == cleanup_count + + +def test_focused_cleanup_failure_emits_only_final_not_verified_report( + capsys: pytest.CaptureFixture[str], +) -> None: + namespace = _recorded_harness_namespace(HARNESS_PATH) + save, load = _complete_focused_results() + cleanup = _CleanupRemoteFunction(capsys, RuntimeError("cleanup image failed")) + globals_dict = namespace["_run_focused_treatment"].__globals__ + globals_dict["run_phase"] = _FakeRemoteFunction(save, load) + globals_dict["cleanup_run"] = cleanup + globals_dict["KEEP_ARTIFACTS"] = False + + with pytest.raises(RuntimeError, match="Focused symmetric TP treatment NOT VERIFIED: cleanup"): + namespace["_run_focused_treatment"]("cleanup-failure") + + report = json.loads(capsys.readouterr().out) + assert report["status"] == "NOT VERIFIED" + assert report["failed_checks"] == ["cleanup_completed"] + assert report["cleanup_attempted"] is True + assert report["cleanup_completed"] is False + assert report["cleanup_skipped"] is False + assert report["artifacts_retained"] is True + assert report["remote_error"] == "cleanup: cleanup image failed" + assert report["remote_error_type"] == "builtins.RuntimeError" + assert cleanup.calls == ["cleanup-failure"] + + +@pytest.mark.parametrize("deletion_failure", ["exception", "silent_noop"]) +def test_actual_cleanup_failure_does_not_commit_and_is_not_verified( + tmp_path: Path, + capsys: pytest.CaptureFixture[str], + monkeypatch: pytest.MonkeyPatch, + deletion_failure: str, +) -> None: + namespace = _recorded_harness_namespace(HARNESS_PATH) + actual_cleanup = namespace["cleanup_run"] + globals_dict = actual_cleanup.__globals__ + run_id = f"cleanup-{deletion_failure}" + run_dir = tmp_path / run_id + run_dir.mkdir() + (run_dir / "archive.bin").write_bytes(b"archive") + events: list[str] = [] + rmtree_kwargs: list[dict] = [] + + def fail_rmtree(_path, **kwargs): + rmtree_kwargs.append(kwargs) + if deletion_failure == "exception": + raise PermissionError("deletion denied") + + monkeypatch.setattr(globals_dict["shutil"], "rmtree", fail_rmtree) + globals_dict["RUNS_ROOT"] = str(tmp_path) + globals_dict["archive_volume"] = _RecordingVolume(events) + save, load = _complete_focused_results() + globals_dict["run_phase"] = _FakeRemoteFunction(save, load) + globals_dict["cleanup_run"] = _RemoteAdapter(actual_cleanup) + globals_dict["KEEP_ARTIFACTS"] = False + + with pytest.raises(RuntimeError, match="Focused symmetric TP treatment NOT VERIFIED: cleanup"): + namespace["_run_focused_treatment"](run_id) + + report = json.loads(capsys.readouterr().out) + assert report["status"] == "NOT VERIFIED" + assert report["failed_checks"] == ["cleanup_completed"] + assert report["cleanup_attempted"] is True + assert report["cleanup_completed"] is False + assert report["artifacts_retained"] is True + assert report["cleanup_error_type"] in { + "builtins.PermissionError", + "builtins.RuntimeError", + } + assert rmtree_kwargs == [{}] + assert run_dir.exists() + assert events == [] + + +def test_cleanup_absent_run_directory_is_idempotent( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + namespace = _recorded_harness_namespace(HARNESS_PATH) + cleanup = namespace["cleanup_run"] + globals_dict = cleanup.__globals__ + events: list[str] = [] + + def reject_rmtree(*_args, **_kwargs): + raise AssertionError("rmtree must not run for an absent directory") + + monkeypatch.setattr(globals_dict["shutil"], "rmtree", reject_rmtree) + globals_dict["RUNS_ROOT"] = str(tmp_path) + globals_dict["archive_volume"] = _RecordingVolume(events) + + cleanup("already-absent") + + assert events == ["commit"] + + +def _assert_full_check_failure( + namespace: dict, + results: dict[str, dict], + capsys: pytest.CaptureFixture[str], + expected_check: str, +) -> dict: + cleanup = _CleanupRemoteFunction(capsys) + globals_dict = namespace["main"].__globals__ + globals_dict["run_phase"] = _FullRemoteFunction(results) + globals_dict["cleanup_run"] = cleanup + globals_dict["KEEP_ARTIFACTS"] = False + + with pytest.raises(RuntimeError, match="TP validation NOT VERIFIED"): + namespace["main"](focused=False) + + report = json.loads(capsys.readouterr().out) + assert expected_check in report["failed_checks"] + assert report["status"] == "NOT VERIFIED" + assert report["cleanup_attempted"] is False + assert report["cleanup_completed"] is False + assert report["artifacts_retained"] is True + assert report["artifacts_may_be_partial_or_mutated"] is True + assert cleanup.calls == [] + return report + + +@pytest.mark.parametrize("phase", ["save", "save2", "load", "load2"]) +@pytest.mark.parametrize( + ("field", "bad_counts", "expected_check"), + [ + ("archive_graph_counts", {"0": 2, "1": 1}, "archive_graph_counts_exact"), + ("archive_graph_json_counts", {"0": 1, "1": 0}, "archive_graph_json_counts_exact"), + ], +) +def test_full_acceptance_requires_one_archived_graph_per_rank_in_every_archive_phase( + capsys: pytest.CaptureFixture[str], + phase: str, + field: str, + bad_counts: dict[str, int], + expected_check: str, +) -> None: + namespace = _recorded_harness_namespace(HARNESS_PATH) + results = _complete_full_results() + results[phase][field] = bad_counts + + _assert_full_check_failure(namespace, results, capsys, expected_check) + + +@pytest.mark.parametrize("phase", ["load", "load2"]) +def test_full_acceptance_requires_one_restored_graph_per_rank_in_both_loads( + capsys: pytest.CaptureFixture[str], + phase: str, +) -> None: + namespace = _recorded_harness_namespace(HARNESS_PATH) + results = _complete_full_results() + results[phase]["loaded_graphs"] = {"0": 0, "1": 1} + + _assert_full_check_failure(namespace, results, capsys, "graphs_restored_each_rank") + + +@pytest.mark.parametrize("phase", ["save", "save2", "load", "load2"]) +@pytest.mark.parametrize( + ("field", "bad_value", "expected_check"), + [ + ("symmetric_all_reduce", 0, "symmetric_archive_collectives_each_phase"), + ("symmetric_all_gather", 0, "symmetric_archive_collectives_each_phase"), + ("all_reduce", 1, "no_archived_nccl_collectives"), + ("all_gather", 1, "no_archived_nccl_collectives"), + ], +) +def test_full_acceptance_requires_symmetric_and_zero_nccl_archive_nodes_per_rank( + capsys: pytest.CaptureFixture[str], + phase: str, + field: str, + bad_value: int, + expected_check: str, +) -> None: + namespace = _recorded_harness_namespace(HARNESS_PATH) + results = _complete_full_results() + counts = { + rank: dict(rank_counts) + for rank, rank_counts in results[phase]["archive_collective_kernel_counts"].items() + } + counts["0"][field] = bad_value + results[phase]["archive_collective_kernel_counts"] = counts + + _assert_full_check_failure(namespace, results, capsys, expected_check) + + +@pytest.mark.parametrize("phase", ["save", "save2", "load", "load2"]) +@pytest.mark.parametrize("missing_evidence", ["rank", "multimem"]) +def test_full_acceptance_requires_symmetric_runtime_markers_in_every_foundry_phase( + capsys: pytest.CaptureFixture[str], + phase: str, + missing_evidence: str, +) -> None: + namespace = _recorded_harness_namespace(HARNESS_PATH) + results = _complete_full_results() + if missing_evidence == "rank": + results[phase]["torch_symm_mem_ranks"] = ["0"] + else: + results[phase]["multimem_all_gather_addresses"] = {} + + _assert_full_check_failure(namespace, results, capsys, "symmetric_markers_each_phase") + + +def test_full_acceptance_keeps_nccl_channel_and_version_maps_diagnostic( + capsys: pytest.CaptureFixture[str], +) -> None: + namespace = _recorded_harness_namespace(HARNESS_PATH) + results = _complete_full_results() + results["baseline"]["nccl_transport_channels"] = {"0": {"NET/IB": [7]}} + results["save"]["nccl_transport_channels"] = {} + results["save2"]["nccl_version"] = "diagnostic-version" + cleanup = _CleanupRemoteFunction(capsys) + globals_dict = namespace["main"].__globals__ + globals_dict["run_phase"] = _FullRemoteFunction(results) + globals_dict["cleanup_run"] = cleanup + globals_dict["KEEP_ARTIFACTS"] = False + + namespace["main"](focused=False) + + report = json.loads(capsys.readouterr().out) + assert report["status"] == "VERIFIED" + assert "nccl_p2p_ipc_every_phase" not in report["checks"] + assert report["results"]["baseline"]["nccl_transport_channels"] == {"0": {"NET/IB": [7]}} + assert report["results"]["save2"]["nccl_version"] == "diagnostic-version" + assert cleanup.calls + + +@pytest.mark.parametrize("keep_artifacts", [False, True]) +def test_full_success_reports_only_after_verified_cleanup( + capsys: pytest.CaptureFixture[str], + keep_artifacts: bool, +) -> None: + namespace = _recorded_harness_namespace(HARNESS_PATH) + remote = _FullRemoteFunction(_complete_full_results()) + cleanup = _CleanupRemoteFunction(capsys) + globals_dict = namespace["main"].__globals__ + globals_dict["run_phase"] = remote + globals_dict["cleanup_run"] = cleanup + globals_dict["KEEP_ARTIFACTS"] = keep_artifacts + + namespace["main"](focused=False) + + report = json.loads(capsys.readouterr().out) + assert remote.calls == ["baseline", "save", "save2", "load", "load2"] + assert report["status"] == "VERIFIED" + assert report["failed_checks"] == [] + assert report["artifacts_retained"] is False + assert report["cleanup_attempted"] is True + assert report["cleanup_completed"] is True + assert report["cleanup_skipped"] is False + assert cleanup.calls + + +def test_full_cleanup_failure_emits_only_final_not_verified_report( + capsys: pytest.CaptureFixture[str], +) -> None: + namespace = _recorded_harness_namespace(HARNESS_PATH) + cleanup = _CleanupRemoteFunction(capsys, RuntimeError("cleanup image failed")) + globals_dict = namespace["main"].__globals__ + globals_dict["run_phase"] = _FullRemoteFunction(_complete_full_results()) + globals_dict["cleanup_run"] = cleanup + globals_dict["KEEP_ARTIFACTS"] = False + + with pytest.raises(RuntimeError, match="TP validation NOT VERIFIED: cleanup"): + namespace["main"](focused=False) + + report = json.loads(capsys.readouterr().out) + assert report["status"] == "NOT VERIFIED" + assert report["failed_checks"] == ["cleanup_completed"] + assert report["cleanup_attempted"] is True + assert report["cleanup_completed"] is False + assert report["cleanup_skipped"] is False + assert report["artifacts_retained"] is True + assert report["cleanup_error"] == "cleanup image failed" + assert report["cleanup_error_type"] == "builtins.RuntimeError" + assert len(cleanup.calls) == 1 + + +def test_full_check_failure_prints_not_verified_before_raise_and_skips_cleanup( + capsys: pytest.CaptureFixture[str], +) -> None: + namespace = _recorded_harness_namespace(HARNESS_PATH) + results = _complete_full_results() + results["load"]["outputs"] = ["wrong"] + cleanup = _CleanupRemoteFunction(capsys) + globals_dict = namespace["main"].__globals__ + globals_dict["run_phase"] = _FullRemoteFunction(results) + globals_dict["cleanup_run"] = cleanup + globals_dict["KEEP_ARTIFACTS"] = False + + with pytest.raises(RuntimeError, match="TP validation NOT VERIFIED"): + namespace["main"](focused=False) + + report = json.loads(capsys.readouterr().out) + assert report["status"] == "NOT VERIFIED" + assert "deterministic_outputs_match_baseline" in report["failed_checks"] + assert report["cleanup_attempted"] is False + assert report["cleanup_completed"] is False + assert report["cleanup_skipped"] is True + assert report["artifacts_retained"] is True + assert cleanup.calls == [] + + +def test_full_concurrent_output_mismatch_remains_fail_closed( + capsys: pytest.CaptureFixture[str], +) -> None: + namespace = _recorded_harness_namespace(HARNESS_PATH) + results = _complete_full_results() + load2_outputs = list(results["load2"]["concurrent_outputs"]) + load2_outputs[3] = "different exact bytes" + results["load2"]["concurrent_outputs"] = load2_outputs + + report = _assert_full_check_failure( + namespace, + results, + capsys, + "concurrent_outputs_match_baseline", + ) + + assert report["checks"]["concurrent_outputs_match_baseline"] is False diff --git a/tests/test_modal_vllm_tp.py b/tests/test_modal_vllm_tp.py new file mode 100644 index 00000000..fe197106 --- /dev/null +++ b/tests/test_modal_vllm_tp.py @@ -0,0 +1,2878 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the Foundry project +"""CPU-only contracts for the vLLM TP Modal harness.""" + +from __future__ import annotations + +import ast +import concurrent.futures +import json +import os +import runpy +import shutil +import subprocess +import sys +import threading +from pathlib import Path + +import pytest + +REPO_ROOT = Path(__file__).parents[1] +HARNESS_PATH = REPO_ROOT / "tests" / "modal_vllm_tp.py" +FULL_SHA = "0123456789abcdef0123456789abcdef01234567" + +# The strict transport profile every phase (baseline + save phases) must report. +_P2P_TRANSPORTS = {"0": ["P2P/IPC"], "1": ["P2P/IPC"]} +_NCCL_P2P_LOG = ( + "[0] NCCL INFO Channel 00/02 : 0 1 via P2P/IPC\n[1] NCCL INFO Channel 00/02 : 1 0 via P2P/IPC\n" +) + + +class _ImageRecorder: + """A fake Modal image that records ``add_local_python_source`` calls. + + Every other builder method is chainable and returns the same recorder, so + the fully constructed image accumulates the sibling-module mounts exactly + as the real Modal chain would. + """ + + def __init__(self, base: str = "registry") -> None: + self.base = base + self.environment: dict[str, str] = {} + self.local_python_sources: list[tuple[tuple[str, ...], bool]] = [] + + def env(self, values: dict[str, str]): + self.environment.update(values) + return self + + def add_local_python_source(self, *modules: str, copy: bool = False, **_kwargs): + self.local_python_sources.append((tuple(modules), copy)) + return self + + def __getattr__(self, _name: str): + def _chain(*_args, **_kwargs): + return self + + return _chain + + +def _fake_modal(): + modal_module = type(sys)("modal") + + class _ImageFactory: + @staticmethod + def from_registry(*_args, **_kwargs): + return _ImageRecorder("registry") + + @staticmethod + def debian_slim(*_args, **_kwargs): + return _ImageRecorder("debian_slim") + + class _App: + def __init__(self, *_args, **_kwargs): + pass + + def get_dashboard_url(self): + return "https://modal.test/apps/vllm" + + def function(self, *_args, **_kwargs): + return lambda fn: fn + + def local_entrypoint(self, *_args, **_kwargs): + return lambda fn: fn + + class _Volume: + @staticmethod + def from_name(*_args, **_kwargs): + return object() + + class _Dict: + @staticmethod + def from_name(*_args, **_kwargs): + return {} + + modal_module.Image = _ImageFactory + modal_module.App = _App + modal_module.Volume = _Volume + modal_module.Dict = _Dict + modal_module.is_local = lambda: True + return modal_module + + +def _construct_image_with_recorder(path: Path, **env_updates: str) -> _ImageRecorder: + saved_modules = { + name: sys.modules.get(name) + for name in ("modal", "modal_vllm_tp", "vllm_tp_validation", "tp_validation_helpers") + } + saved_path = list(sys.path) + saved_env = os.environ.copy() + try: + sys.modules["modal"] = _fake_modal() + sys.path.insert(0, str(path.parent)) + os.environ.update({"FOUNDRY_REF": FULL_SHA, "VLLM_REF": FULL_SHA}) + os.environ.update(env_updates) + namespace = runpy.run_path(str(path), run_name="_image_recorder_contract") + image = namespace["image"] + assert isinstance(image, _ImageRecorder), type(image) + return image + finally: + sys.path[:] = saved_path + os.environ.clear() + os.environ.update(saved_env) + for name, module in saved_modules.items(): + if module is None: + sys.modules.pop(name, None) + else: + sys.modules[name] = module + + +def test_image_mounts_sibling_helper_modules_for_remote_hydration() -> None: + image = _construct_image_with_recorder(HARNESS_PATH) + + attached = {module for modules, _copy in image.local_python_sources for module in modules} + copy_flags = {copy for _modules, copy in image.local_python_sources} + + assert "vllm_tp_validation" in attached + assert "tp_validation_helpers" in attached + # Runtime mount (copy=False): sibling modules on the container PYTHONPATH + # without baking a new image layer. + assert copy_flags == {False} + + +def test_primary_and_cleanup_images_pin_remote_hydration_environment() -> None: + recorded = _load_harness_globals() + primary = recorded["image"] + cleanup = recorded["cleanup_image"] + expected = { + "FOUNDRY_REF": FULL_SHA, + "VLLM_REF": FULL_SHA, + "VLLM_WHEEL_COMMIT": recorded["VLLM_WHEEL_COMMIT"], + "VLLM_MODEL_REVISION": recorded["MODEL_REVISION"], + "TP_SIZE": "2", + "MODAL_GPU": recorded["MODAL_GPU"], + } + for candidate in (primary, cleanup): + assert isinstance(candidate, _ImageRecorder) + assert {key: candidate.environment[key] for key in expected} == expected + attached = { + module for modules, _copy in candidate.local_python_sources for module in modules + } + assert {"vllm_tp_validation", "tp_validation_helpers"} <= attached + assert cleanup.base == "debian_slim" + + +def test_remote_module_hydration_in_subprocess_requires_no_git_and_imports_helpers() -> None: + env = os.environ.copy() + env.update( + { + "FOUNDRY_REF": FULL_SHA, + "VLLM_REF": FULL_SHA, + "VLLM_WHEEL_COMMIT": FULL_SHA, + "VLLM_MODEL_REVISION": FULL_SHA, + "TP_SIZE": "2", + } + ) + code = r""" +import runpy +import subprocess +import sys +from types import ModuleType + +class Image: + def env(self, *_a, **_k): return self + def add_local_python_source(self, *modules, **_k): + assert {"vllm_tp_validation", "tp_validation_helpers"} <= set(modules) + return self + def __getattr__(self, _name): + return lambda *_a, **_k: self + +class ImageFactory: + @staticmethod + def from_registry(*_a, **_k): return Image() + @staticmethod + def debian_slim(*_a, **_k): return Image() + +class App: + def __init__(self, *_a, **_k): pass + def function(self, *_a, **_k): return lambda fn: fn + def local_entrypoint(self, *_a, **_k): return lambda fn: fn + +class Named: + @staticmethod + def from_name(*_a, **_k): return {} + +modal = ModuleType("modal") +modal.Image = ImageFactory +modal.App = App +modal.Volume = Named +modal.Dict = Named +sys.modules["modal"] = modal +sys.path.insert(0, sys.argv[1]) +subprocess.check_output = lambda *_a, **_k: (_ for _ in ()).throw( + AssertionError("remote hydration invoked git") +) +namespace = runpy.run_path(sys.argv[2], run_name="_remote_hydration_contract") +assert namespace["FOUNDRY_REF"] == sys.argv[3] +assert namespace["VLLM_REF"] == sys.argv[3] +assert namespace["TP_SIZE"] == 2 +print("hydrated") +""" + result = subprocess.run( + [ + sys.executable, + "-c", + code, + str(HARNESS_PATH.parent), + str(HARNESS_PATH), + FULL_SHA, + ], + cwd=REPO_ROOT, + env=env, + capture_output=True, + text=True, + ) + + assert result.returncode == 0, result.stderr + assert result.stdout.strip() == "hydrated" + + +# --------------------------------------------------------------------------- +# Transaction / LOAD2 contracts against the live harness module. +# +# ``runpy.run_path`` returns a *copy* of the module globals, so its top-level +# functions keep their real (shared) ``__globals__`` dict. Patching that dict +# swaps the server primitives (``_launch`` / ``_wait_until_healthy`` / +# ``_generate`` / ``stop_process``) for fakes and lets us drive the transaction +# and LOADs on CPU without Modal or a GPU. The fake ``modal.Volume`` records +# ``commit()`` so every publication path is traceable. +# --------------------------------------------------------------------------- + + +class _RecordingVolume: + def __init__(self, events: list[str]) -> None: + self.events = events + self.commit_count = 0 + self.reload_count = 0 + # Number of upcoming commit() calls to fail before succeeding. Set very + # high to make every bounded-retry attempt fail. + self.fail_commits = 0 + self.commit_error = RuntimeError("commit boom") + + def commit(self) -> None: + self.commit_count += 1 + if self.fail_commits > 0: + self.fail_commits -= 1 + raise self.commit_error + self.events.append("commit") + + def reload(self) -> None: + self.reload_count += 1 + self.events.append("reload") + + +class _RecordingDict: + """A fake ``modal.Dict`` with atomic ``put(skip_if_exists=...)`` semantics.""" + + def __init__(self, events: list[str]) -> None: + self.events = events + self.store: dict[str, object] = {} + self.put_count = 0 + self.pop_count = 0 + # Number of upcoming put()/pop() calls to raise before behaving normally. + self.fail_puts = 0 + self.fail_pops = 0 + # Number of upcoming pop() calls that remove the key AND then raise, to + # model a write that committed but whose acknowledgement was lost. + self.fail_pops_after_write = 0 + self.put_error = RuntimeError("index boom") + self.pop_error = RuntimeError("pop boom") + + def put(self, key: str, value: object, skip_if_exists: bool = False) -> bool: + self.put_count += 1 + if self.fail_puts > 0: + self.fail_puts -= 1 + raise self.put_error + if skip_if_exists and key in self.store: + return False + self.store[key] = value + self.events.append("index") + return True + + def __setitem__(self, key: str, value: object) -> None: + self.store[key] = value + self.events.append("index") + + def __getitem__(self, key: str) -> object: + value = self.store[key] # a miss raises KeyError before recording a read + self.events.append("read") + return value + + def __contains__(self, key: str) -> bool: + return key in self.store + + def get(self, key: str, default: object = None) -> object: + return self.store.get(key, default) + + def pop(self, key: str, default: object = None) -> object: + self.pop_count += 1 + if self.fail_pops > 0: + self.fail_pops -= 1 + raise self.pop_error + if self.fail_pops_after_write > 0: + self.fail_pops_after_write -= 1 + self.store.pop(key, default) # the removal commits, then the ack is lost + raise self.pop_error + self.events.append("pop") + return self.store.pop(key, default) + + +def _fake_modal_with_recording_volumes(): + modal_module = type(sys)("modal") + + class _ImageFactory: + @staticmethod + def from_registry(*_args, **_kwargs): + return _ImageRecorder("registry") + + @staticmethod + def debian_slim(*_args, **_kwargs): + return _ImageRecorder("debian_slim") + + class _App: + def __init__(self, *_args, **_kwargs): + pass + + def get_dashboard_url(self): + return "https://modal.test/apps/vllm" + + def function(self, *_args, **_kwargs): + return lambda fn: fn + + def local_entrypoint(self, *_args, **_kwargs): + return lambda fn: fn + + # A single shared event log orders volume commits vs Dict index writes so a + # test can assert the index write is the last durable side effect. + events: list[str] = [] + volumes: dict[str, _RecordingVolume] = {} + dicts: dict[str, _RecordingDict] = {} + + class _Volume: + @staticmethod + def from_name(name, *_args, **_kwargs): + return volumes.setdefault(name, _RecordingVolume(events)) + + class _Dict: + @staticmethod + def from_name(name, *_args, **_kwargs): + return dicts.setdefault(name, _RecordingDict(events)) + + modal_module.Image = _ImageFactory + modal_module.App = _App + modal_module.Volume = _Volume + modal_module.Dict = _Dict + modal_module.is_local = lambda: True + return modal_module + + +def _load_harness_globals(**env_updates: str) -> dict: + saved_modules = { + name: sys.modules.get(name) + for name in ("modal", "modal_vllm_tp", "vllm_tp_validation", "tp_validation_helpers") + } + saved_path = list(sys.path) + saved_env = os.environ.copy() + try: + sys.modules["modal"] = _fake_modal_with_recording_volumes() + sys.path.insert(0, str(HARNESS_PATH.parent)) + os.environ.update({"FOUNDRY_REF": FULL_SHA, "VLLM_REF": FULL_SHA}) + os.environ.update(env_updates) + namespace = runpy.run_path(str(HARNESS_PATH), run_name="_transaction_contract") + return namespace["run_save_transaction"].__globals__ + finally: + sys.path[:] = saved_path + os.environ.clear() + os.environ.update(saved_env) + for name, module in saved_modules.items(): + if module is None: + sys.modules.pop(name, None) + else: + sys.modules[name] = module + + +def _build_exact_archive(workspace: Path, tp_size: int = 2, graph_count: int = 64) -> dict: + workspace.mkdir(parents=True, exist_ok=True) + (workspace / "warmup_state.json").write_text( + '{"memory_pool_config": {"max_total_num_tokens": 1}}' + ) + stems: dict[str, dict[int, str]] = {} + for rank in range(tp_size): + rank_dir = workspace / f"rank_{rank}" + rank_dir.mkdir(parents=True) + rank_stems: dict[int, str] = {} + for index in range(graph_count): + # Descending capture order: capture_index i -> batch size count-i, so + # index 0 is the largest size (t{count}) and index count-1 is t1. + size = graph_count - index + stem = f"graph_{index}_FULL_t{size}_r{size}_UX_pcN" + rank_stems[index] = stem + (rank_dir / f"{stem}.json").write_text('{"nodes": []}') + (rank_dir / f"{stem}.cugraph").write_bytes(b"\x00cugraph" + bytes([index % 256])) + stems[str(rank)] = rank_stems + (rank_dir / "graph_manifest.json").write_text('{"groups": []}') + (rank_dir / "fatbin_image_packed.img").write_bytes(b"FATBIN\x00") + (rank_dir / "fatbin_entrypoint_packed.txt").write_text("entrypoint\n") + (rank_dir / "final_alloc_offset.json").write_text('{"final_alloc_offset": 69736595456}') + return stems + + +def _commit_calls_by_function() -> dict[str, int]: + tree = ast.parse(HARNESS_PATH.read_text()) + counts: dict[str, int] = {} + for node in ast.walk(tree): + if not isinstance(node, ast.FunctionDef): + continue + count = 0 + for sub in ast.walk(node): + if ( + isinstance(sub, ast.Call) + and isinstance(sub.func, ast.Attribute) + and sub.func.attr == "commit" + and isinstance(sub.func.value, ast.Name) + and sub.func.value.id == "archive_volume" + ): + count += 1 + counts[node.name] = count + return counts + + +def test_phase_plan_constants_are_exact() -> None: + g = _load_harness_globals() + + assert g["EXPECTED_GRAPH_COUNT"] == 64 + assert g["PHASES"] == ("baseline", "seed", "save", "save2", "load", "load2") + assert g["QUERY_PHASES"] == ("baseline", "save2", "load", "load2") + assert g["LOAD_PHASES"] == ("load", "load2") + + +def test_commit_is_isolated_to_the_commit_helper() -> None: + counts = _commit_calls_by_function() + + # ``archive_volume.commit()`` appears textually in exactly one place: the + # bounded-retry helper. Every other function reaches the volume only by + # calling that helper, never by committing directly. + assert counts.get("_commit_with_retries") == 1 + assert counts.get("run_save_transaction", 0) == 0 + assert counts.get("_publish_archive", 0) == 0 + assert counts.get("_run_transaction_phase", 0) == 0 + assert counts.get("_run_load_local", 0) == 0 + assert counts.get("run_load_phase", 0) == 0 + assert counts.get("run_baseline_phase", 0) == 0 + assert counts.get("main", 0) == 0 + assert sum(counts.values()) == 1 + + +def _dict_puts_by_function() -> dict[str, int]: + tree = ast.parse(HARNESS_PATH.read_text()) + counts: dict[str, int] = {} + for node in ast.walk(tree): + if not isinstance(node, ast.FunctionDef): + continue + count = 0 + for sub in ast.walk(node): + if ( + isinstance(sub, ast.Call) + and isinstance(sub.func, ast.Attribute) + and sub.func.attr == "put" + and isinstance(sub.func.value, ast.Name) + and sub.func.value.id == "publication_index" + ): + count += 1 + counts[node.name] = count + return counts + + +def test_marker_put_is_isolated_to_the_marker_helper() -> None: + counts = _dict_puts_by_function() + + # ``publication_index.put(...)`` appears in exactly one place: the atomic + # marker helper. No other function mutates the Dict via put. + assert counts.get("_write_publication_marker") == 1 + assert sum(counts.values()) == 1 + + # The Dict is never mutated via subscript assignment anywhere in the harness. + tree = ast.parse(HARNESS_PATH.read_text()) + subscript_writes = [ + target + for node in ast.walk(tree) + if isinstance(node, ast.Assign) + for target in node.targets + if isinstance(target, ast.Subscript) + and isinstance(target.value, ast.Name) + and target.value.id == "publication_index" + ] + assert subscript_writes == [] + + +def test_marker_write_is_the_single_final_side_effect() -> None: + # In run_save_transaction the marker write is the last statement before the + # sole return, and comes after the commit — no fallible op may follow it. + tree = ast.parse(HARNESS_PATH.read_text()) + fn = next( + node + for node in ast.walk(tree) + if isinstance(node, ast.FunctionDef) and node.name == "run_save_transaction" + ) + body = fn.body + + def _call_pos(name: str) -> int: + return next( + i + for i, stmt in enumerate(body) + if isinstance(stmt, ast.Expr) + and isinstance(stmt.value, ast.Call) + and isinstance(stmt.value.func, ast.Name) + and stmt.value.func.id == name + ) + + commit_pos = _call_pos("_commit_with_retries") + marker_pos = _call_pos("_write_publication_marker") + assert commit_pos < marker_pos + assert marker_pos == len(body) - 2 # marker write, then return + assert isinstance(body[-1], ast.Return) + + +def test_source_is_python310_compatible_no_exception_group() -> None: + tree = ast.parse(HARNESS_PATH.read_text()) + + # ``ExceptionGroup`` is a 3.11+ builtin; the source must not use it (as a name + # or a raise) so it runs unchanged on the repo's declared 3.10 minimum. + names = {node.id for node in ast.walk(tree) if isinstance(node, ast.Name)} + assert "ExceptionGroup" not in names + + # The removed publication complexity must be gone as well. + source = HARNESS_PATH.read_text() + assert "renameat2" not in source + assert "_commit_or_rollback" not in source + + +class _FakeProcess: + def __init__(self) -> None: + self.returncode = 0 + self.pid = 4321 + + def poll(self): + return None + + +def _byte_map(root: Path) -> dict[str, bytes]: + return { + path.relative_to(root).as_posix(): path.read_bytes() + for path in root.rglob("*") + if path.is_file() + } + + +def _load_log(offsets: dict[str, int]) -> str: + lines = [ + f"[foundry] LOAD live final_alloc_offset rank_{rank}: {offset}" + for rank, offset in sorted(offsets.items()) + ] + lines += [ + f"[foundry] LOAD restoration complete rank_{rank}: 64/64 graphs" for rank in sorted(offsets) + ] + return "\n".join(lines) + "\n" + + +def _install_fake_server( + g: dict, + *, + archive_writer=lambda mode, workspace: None, + log_body_for=lambda mode: _NCCL_P2P_LOG, + stop_hook=lambda mode, workspace: None, +) -> None: + """Drive server phases on CPU: write a fake local archive and fake logs. + + ``archive_writer`` populates ``run_root/foundry_archive`` (the local, never + mounted, workspace); ``log_body_for`` returns the log text; ``stop_hook`` + runs inside the faked teardown so tests can simulate a mutation that happens + after the server exits but before the post-run archive reread. + """ + + def fake_launch(mode, log_path, run_root): + workspace = Path(run_root) / "foundry_archive" + archive_writer(mode, workspace) + Path(log_path).write_text(log_body_for(mode)) + handle = Path(log_path).open("a") # noqa: SIM115 - closed by the faked stop_process + handle._foundry_mode = mode # type: ignore[attr-defined] + handle._foundry_workspace = workspace # type: ignore[attr-defined] + return _FakeProcess(), handle + + def fake_stop_process(process, log_file, *args, **kwargs): + stop_hook( + getattr(log_file, "_foundry_mode", None), + getattr(log_file, "_foundry_workspace", None), + ) + log_file.close() + return 0 + + g["_launch"] = fake_launch + g["_wait_until_healthy"] = lambda *a, **k: None + g["_generate"] = lambda prompt: {"text": f"o:{prompt}", "token_strings": [prompt]} + g["_generate_concurrently"] = lambda: [ + {"text": f"c:{index}", "token_strings": ["c", str(index)]} for index in range(4) + ] + g["_runtime_provenance"] = lambda nccl_version: { + "actual_torch_version": "2.11.0+cu130", + "actual_vllm_version": "0.15.0", + "driver_version": "580.95.05", + "gpu_count": 2, + "gpu_names": ["NVIDIA H100", "NVIDIA H100"], + "nccl_runtime_version": nccl_version or "2.28.9", + } + g["stop_process"] = fake_stop_process + + +def test_launch_closes_log_handle_when_popen_raises( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + g, *_ = _harness_with_storage(tmp_path) + + class Handle: + closed = False + + def close(self) -> None: + self.closed = True + + handle = Handle() + g["open"] = lambda *_a, **_k: handle + monkeypatch.setattr( + g["subprocess"], + "Popen", + lambda *_a, **_k: (_ for _ in ()).throw(OSError("spawn failed")), + ) + + with pytest.raises(OSError, match="spawn failed"): + g["_launch"](None, str(tmp_path / "launch.log"), tmp_path) + + assert handle.closed is True + + +def test_real_concurrent_generation_is_barrier_backed_and_preserves_prompt_order( + tmp_path: Path, +) -> None: + g, *_ = _harness_with_storage(tmp_path) + prompts = ["p3", "p1", "p4", "p2"] + g["PROMPTS"] = prompts + entered_together = threading.Barrier(len(prompts), timeout=1) + + def generate(prompt: str) -> dict[str, object]: + entered_together.wait() + return {"text": prompt, "token_strings": [prompt]} + + g["_generate"] = generate + + assert g["_generate_concurrently"]() == [ + {"text": prompt, "token_strings": [prompt]} for prompt in prompts + ] + + +def test_real_concurrent_generation_uses_bounded_barrier_timeout( + tmp_path: Path, +) -> None: + g, *_ = _harness_with_storage(tmp_path) + g["CONCURRENT_BARRIER_TIMEOUT_S"] = 0.05 + real_executor = concurrent.futures.ThreadPoolExecutor + g["ThreadPoolExecutor"] = lambda max_workers: real_executor(max_workers=max_workers - 1) + g["_generate"] = lambda prompt: {"text": prompt, "token_strings": [prompt]} + + with pytest.raises(threading.BrokenBarrierError): + g["_generate_concurrently"]() + + +def test_real_concurrent_generation_propagates_request_error_after_executor_cleanup( + tmp_path: Path, +) -> None: + g, *_ = _harness_with_storage(tmp_path) + started: list[str] = [] + lock = threading.Lock() + + def generate(prompt: str) -> dict[str, object]: + with lock: + started.append(prompt) + if prompt == g["PROMPTS"][1]: + raise RuntimeError("request failed") + return {"text": prompt, "token_strings": [prompt]} + + g["_generate"] = generate + + with pytest.raises(RuntimeError, match="request failed"): + g["_generate_concurrently"]() + + assert sorted(started) == sorted(g["PROMPTS"]) + + +def _write_save_archive(mode: str | None, workspace: Path) -> None: + if mode != "save": + return + shutil.rmtree(workspace, ignore_errors=True) + _build_exact_archive(workspace) + + +def _harness_with_storage(tmp_path: Path): + archive_root = tmp_path / "archive" + local_root = tmp_path / "local" + g = _load_harness_globals( + VLLM_TP_ARCHIVE_ROOT=str(archive_root), + VLLM_TP_LOCAL_ROOT=str(local_root), + ) + return g, archive_root, local_root + + +def _snapshot_of(g: dict, root: Path) -> dict: + return g["collect_archive_snapshot"]( + root, 2, 64, g["FOUNDRY_BASE_ADDR"], g["FOUNDRY_REGION_END"] + ) + + +# --------------------------------------------------------------------------- +# Save transaction: local run, canonical copy, commit, Dict index write. +# --------------------------------------------------------------------------- + + +def test_save_transaction_runs_locally_commits_then_indexes(tmp_path: Path) -> None: + g, archive_root, local_root = _harness_with_storage(tmp_path) + volume = g["archive_volume"] + index = g["publication_index"] + _install_fake_server(g, archive_writer=_write_save_archive) + + run_id = "run-publish" + baseline = g["run_baseline_phase"](run_id) + evidence = { + "outputs": baseline["outputs"], + "concurrent_outputs": baseline["concurrent_outputs"], + "nccl_transports": baseline["nccl_transports"], + } + out = g["run_save_transaction"](run_id, evidence) + + # Exactly one commit and one atomic marker put; the marker is the single last + # durable side effect (commit strictly precedes it). + assert volume.commit_count == 1 + assert index.put_count == 1 + assert volume.events == ["commit", "index"] + + # The transaction ran against ephemeral local storage, never the volume. + assert (local_root / run_id / "txn" / "foundry_archive" / "rank_0").is_dir() + + published = Path(out["published_path"]) + assert published == archive_root / "published" / run_id + assert out["published_run_id"] == run_id + assert published.is_dir() + archive = published / "archive" + assert len(list((archive / "rank_0").glob("graph_*.cugraph"))) == 64 + assert not (published / "_incoming").exists() + + # The stored record carries identity, run_id (for cleanup), and the payload. + digest = g["_snapshot_digest"](out["snapshot"]) + assert out["snapshot_digest"] == digest + record = index.store[run_id] + assert record["run_id"] == run_id + assert record["snapshot_digest"] == digest + assert record["published_path"] == str(published) + assert record["foundry_ref"] == g["FOUNDRY_REF"] + assert record["payload"] == out + + # The on-volume manifest mirrors the identity plus the full snapshot. + manifest = json.loads((published / "manifest.json").read_text()) + assert manifest == { + "snapshot_digest": digest, + "published_path": str(published), + "foundry_ref": g["FOUNDRY_REF"], + "snapshot": out["snapshot"], + } + assert out["snapshot"] == _snapshot_of(g, archive) + + +def test_stored_payload_is_bounded_and_excludes_large_evidence(tmp_path: Path) -> None: + g, archive_root, _ = _harness_with_storage(tmp_path) + _install_fake_server(g, archive_writer=_write_save_archive) + + run_id = "run-bounded" + baseline = g["run_baseline_phase"](run_id) + evidence = { + "outputs": baseline["outputs"], + "concurrent_outputs": baseline["concurrent_outputs"], + "nccl_transports": baseline["nccl_transports"], + } + out = g["run_save_transaction"](run_id, evidence) + + # Reduced per-phase results carry only bounded primitives, not the large + # validated inventory or per-file byte fingerprints. + for phase in ("seed", "save", "save2"): + pr = out["results"][phase] + assert "archive_inventory" not in pr + assert "archive_byte_fingerprints" not in pr + assert set(pr).issubset(set(g["_STORED_PHASE_FIELDS"])) + # save/save2 keep exactly the fields main needs for its checks. + assert "rank_offsets" in out["results"]["save"] + assert "archive_fingerprints" in out["results"]["save2"] + + +def _build_archive_with_count(workspace: Path, count: int) -> None: + """Build a paired, contiguous archive of ``count`` canonical-descriptor graphs.""" + workspace.mkdir(parents=True, exist_ok=True) + (workspace / "warmup_state.json").write_text( + '{"memory_pool_config": {"max_total_num_tokens": 1}}' + ) + for rank in range(2): + rank_dir = workspace / f"rank_{rank}" + rank_dir.mkdir(parents=True, exist_ok=True) + for index in range(count): + size = count - index # descending capture order + stem = f"graph_{index}_FULL_t{size}_r{size}_UX_pcN" + (rank_dir / f"{stem}.json").write_text('{"nodes": []}') + (rank_dir / f"{stem}.cugraph").write_bytes(b"\x00cg" + bytes([index % 256])) + (rank_dir / "graph_manifest.json").write_text('{"groups": []}') + (rank_dir / "fatbin_image_packed.img").write_bytes(b"FATBIN\x00") + (rank_dir / "fatbin_entrypoint_packed.txt").write_text("entrypoint\n") + (rank_dir / "final_alloc_offset.json").write_text('{"final_alloc_offset": 69736595456}') + + +def _build_seed_capture_archive(workspace: Path) -> None: + """Model the actual pinned SEED capture order (66 graphs with duplicate sizes). + + Profiling captures batch sizes 64 then 63 at indices 0/1, then the normal + descending capture loop records sizes 64..1 at indices 2..65. Sizes 64 and 63 + therefore each appear twice (across distinct indices) — exactly the seed + duplication the exact-64 contract must NOT reject. + """ + workspace.mkdir(parents=True, exist_ok=True) + (workspace / "warmup_state.json").write_text( + '{"memory_pool_config": {"max_total_num_tokens": 1}}' + ) + sizes = [64, 63, *range(64, 0, -1)] # profiling 64/63, then descending 64..1 + assert len(sizes) == 66 + for rank in range(2): + rank_dir = workspace / f"rank_{rank}" + rank_dir.mkdir(parents=True, exist_ok=True) + for index, size in enumerate(sizes): + stem = f"graph_{index}_FULL_t{size}_r{size}_UX_pcN" + (rank_dir / f"{stem}.json").write_text('{"nodes": []}') + (rank_dir / f"{stem}.cugraph").write_bytes(b"\x00cg" + bytes([index % 256])) + (rank_dir / "graph_manifest.json").write_text('{"groups": []}') + (rank_dir / "fatbin_image_packed.img").write_bytes(b"FATBIN\x00") + (rank_dir / "fatbin_entrypoint_packed.txt").write_text("entrypoint\n") + (rank_dir / "final_alloc_offset.json").write_text('{"final_alloc_offset": 69736595456}') + + +def test_expected_graph_stem_matches_source_profile(tmp_path: Path) -> None: + g, *_ = _harness_with_storage(tmp_path) + + # Source-confirmed profile (recipe seq 1..64, vLLM captures largest-first): + # capture_index i -> batch size 64-i, so index 0 is t64 and index 63 is t1. + assert g["_expected_graph_stem"](0) == "graph_0_FULL_t64_r64_UX_pcN" + assert g["_expected_graph_stem"](63) == "graph_63_FULL_t1_r1_UX_pcN" + + +def test_seed_phase_accepts_diagnostic_graph_count(tmp_path: Path) -> None: + g, *_ = _harness_with_storage(tmp_path) + + def seed66(mode, workspace): + # Pinned seed profiling captures 66 pairs with duplicate sizes 64/63. + _build_seed_capture_archive(Path(workspace)) + + _install_fake_server(g, archive_writer=seed66) + + txn_root = tmp_path / "seedtxn" + workspace = txn_root / "foundry_archive" + result = g["_run_transaction_phase"]("seed", txn_root, workspace) + + # Seed's diagnostic archive is recorded but NOT held to the exact-64 canonical + # contract: 66 pairs must not fail the phase. + diag = result["seed_diagnostics"] + assert diag["0"]["json_count"] == 66 + assert diag["0"]["cugraph_count"] == 66 + assert len(diag["0"]["sample_files"]) <= g["GRAPH_DIAGNOSTIC_SAMPLE_LIMIT"] + # The diagnostic exposes the profiling duplicates (sizes 64 and 63 twice). + assert diag["0"]["duplicate_sizes"] == {"64": 2, "63": 2} + # No canonical SAVE evidence is attached to seed. + assert "archive_inventory" not in result + assert "rank_offsets" not in result + + +def test_save_phase_rejects_noncanonical_graph_count(tmp_path: Path) -> None: + g, *_ = _harness_with_storage(tmp_path) + + def save66(mode, workspace): + _build_seed_capture_archive(Path(workspace)) + + _install_fake_server(g, archive_writer=save66) + + txn_root = tmp_path / "savetxn" + workspace = txn_root / "foundry_archive" + with pytest.raises(g["PhaseFailure"]) as excinfo: + g["_run_transaction_phase"]("save", txn_root, workspace) + + message = str(excinfo.value) + # SAVE's 66 is rejected by the exact-64 contract, and the failure carries + # bounded graph filename/count diagnostics for future debugging. + assert "66" in message + assert "graph_0_FULL_t64_r64_UX_pcN" in message + + +def test_save_phase_rejects_wrong_descriptor_with_exact_count(tmp_path: Path) -> None: + g, *_ = _harness_with_storage(tmp_path) + + def wrong_descriptor(mode, workspace): + workspace = Path(workspace) + _build_exact_archive(workspace) + # Exactly 64 paired, contiguous graphs, but index 5 (canonically t59) has + # a wrong descriptor. + for suffix in (".json", ".cugraph"): + (workspace / "rank_0" / f"graph_5_FULL_t59_r59_UX_pcN{suffix}").rename( + workspace / "rank_0" / f"graph_5_FULL_t777_r777_UX_pcN{suffix}" + ) + + _install_fake_server(g, archive_writer=wrong_descriptor) + + txn_root = tmp_path / "descrtxn" + workspace = txn_root / "foundry_archive" + with pytest.raises(g["PhaseFailure"], match="(?i)descriptor|expected"): + g["_run_transaction_phase"]("save", txn_root, workspace) + + +def test_reduced_results_enforce_exact_per_phase_schema(tmp_path: Path) -> None: + g, *_ = _harness_with_storage(tmp_path) + + # The canonical reduced results are well-formed. + assert g["_reduced_results_are_wellformed"](_stored_results(g)) is True + + # seed carrying a save-only canonical field is rejected. + bad = _stored_results(g) + bad["seed"]["rank_offsets"] = {"0": 1, "1": 1} + assert g["_reduced_results_are_wellformed"](bad) is False + + # save carrying seed's diagnostic field is rejected. + bad = _stored_results(g) + bad["save"]["seed_diagnostics"] = {} + assert g["_reduced_results_are_wellformed"](bad) is False + + # save2 carrying seed's diagnostic field is rejected. + bad = _stored_results(g) + bad["save2"]["seed_diagnostics"] = {} + assert g["_reduced_results_are_wellformed"](bad) is False + + # An unknown extra field on any phase is rejected. + bad = _stored_results(g) + bad["save"]["mystery"] = 1 + assert g["_reduced_results_are_wellformed"](bad) is False + + # Query extras are permitted only on save2: save carrying them is rejected. + bad = _stored_results(g) + bad["save"]["outputs"] = bad["save2"]["outputs"] + bad["save"]["concurrent_outputs"] = bad["save2"]["concurrent_outputs"] + assert g["_reduced_results_are_wellformed"](bad) is False + + # seed missing its required diagnostic is rejected. + bad = _stored_results(g) + del bad["seed"]["seed_diagnostics"] + assert g["_reduced_results_are_wellformed"](bad) is False + + # save2 missing its required query extras is rejected. + bad = _stored_results(g) + del bad["save2"]["outputs"] + assert g["_reduced_results_are_wellformed"](bad) is False + + +def test_stored_payload_seed_carries_diagnostic_shape(tmp_path: Path) -> None: + g, archive_root, _ = _harness_with_storage(tmp_path) + _install_fake_server(g, archive_writer=_write_save_archive) + + run_id = "run-seed-shape" + baseline = g["run_baseline_phase"](run_id) + evidence = { + "outputs": baseline["outputs"], + "concurrent_outputs": baseline["concurrent_outputs"], + "nccl_transports": baseline["nccl_transports"], + } + out = g["run_save_transaction"](run_id, evidence) + + seed = out["results"]["seed"] + # Seed's stored shape is diagnostic-only: no canonical SAVE reduction. + assert "seed_diagnostics" in seed + assert "rank_offsets" not in seed + assert "archive_fingerprints" not in seed + # save/save2 keep the canonical reduced fields main relies on. + assert "rank_offsets" in out["results"]["save"] + assert "archive_fingerprints" in out["results"]["save2"] + + +def test_transaction_rejects_oversized_record(tmp_path: Path) -> None: + g, archive_root, _ = _harness_with_storage(tmp_path) + volume = g["archive_volume"] + index = g["publication_index"] + g["PUBLICATION_RECORD_MAX_BYTES"] = 1 # force the size preflight to fail + _install_fake_server(g, archive_writer=_write_save_archive) + + run_id = "run-toobig" + baseline = g["run_baseline_phase"](run_id) + evidence = { + "outputs": baseline["outputs"], + "concurrent_outputs": baseline["concurrent_outputs"], + "nccl_transports": baseline["nccl_transports"], + } + + with pytest.raises(RuntimeError, match="(?i)exceeds limit"): + g["run_save_transaction"](run_id, evidence) + assert volume.commit_count == 0 + assert index.put_count == 0 + assert not (archive_root / "published" / run_id).exists() + + +def test_failed_transaction_publishes_no_canonical_and_never_indexes(tmp_path: Path) -> None: + g, archive_root, _ = _harness_with_storage(tmp_path) + volume = g["archive_volume"] + index = g["publication_index"] + + def broken_save2(mode, workspace): + _write_save_archive(mode, workspace) + # SAVE2 drops one graph: exact-inventory validation must reject it. + if mode == "save" and getattr(broken_save2, "calls", 0) == 2: + (workspace / "rank_1" / "graph_63_FULL_t1_r1_UX_pcN.cugraph").unlink() + broken_save2.calls = getattr(broken_save2, "calls", 0) + 1 + + _install_fake_server(g, archive_writer=broken_save2) + + run_id = "run-broken" + evidence = { + "outputs": [{"text": "o", "token_strings": ["o"]}], + "concurrent_outputs": [{"text": "c", "token_strings": ["c"]}], + } + + with pytest.raises((RuntimeError, ValueError)): + g["run_save_transaction"](run_id, evidence) + + assert volume.commit_count == 0 + assert index.put_count == 0 + assert run_id not in index + assert not (archive_root / "published" / run_id).exists() + + +def test_transaction_requires_save2_outputs_equal_baseline(tmp_path: Path) -> None: + g, archive_root, _ = _harness_with_storage(tmp_path) + volume = g["archive_volume"] + _install_fake_server(g, archive_writer=_write_save_archive) + + run_id = "run-outputs" + # Baseline evidence that SAVE2's deterministic queries cannot reproduce. + evidence = { + "outputs": [{"text": "DIVERGENT", "token_strings": ["x"]}], + "concurrent_outputs": [{"text": "c", "token_strings": ["c"]}], + "nccl_transports": _P2P_TRANSPORTS, + } + + with pytest.raises(RuntimeError, match="(?i)baseline"): + g["run_save_transaction"](run_id, evidence) + + assert volume.commit_count == 0 + assert not (archive_root / "published" / run_id).exists() + + +def test_transaction_requires_reproducible_save_offsets(tmp_path: Path) -> None: + g, archive_root, _ = _harness_with_storage(tmp_path) + volume = g["archive_volume"] + + def drifting_offset(mode, workspace): + _write_save_archive(mode, workspace) + drifting_offset.calls = getattr(drifting_offset, "calls", 0) + 1 + # SAVE2 (third SAVE) records a different per-rank offset than SAVE. + if mode == "save" and drifting_offset.calls == 3: + for rank in range(2): + (workspace / f"rank_{rank}" / "final_alloc_offset.json").write_text( + '{"final_alloc_offset": 111}' + ) + + _install_fake_server(g, archive_writer=drifting_offset) + + run_id = "run-drift" + baseline = { + "outputs": [{"text": "o:" + p, "token_strings": [p]} for p in g["PROMPTS"]], + "concurrent_outputs": [{"text": "c", "token_strings": ["c"]}], + "nccl_transports": _P2P_TRANSPORTS, + } + + with pytest.raises(RuntimeError, match="(?i)offset"): + g["run_save_transaction"](run_id, baseline) + + assert volume.commit_count == 0 + assert not (archive_root / "published" / run_id).exists() + + +def test_transaction_tolerates_cugraph_byte_drift(tmp_path: Path) -> None: + g, archive_root, _ = _harness_with_storage(tmp_path) + volume = g["archive_volume"] + + def byte_drift(mode, workspace): + _write_save_archive(mode, workspace) + byte_drift.calls = getattr(byte_drift, "calls", 0) + 1 + # SAVE2 writes different raw cugraph bytes, but identical graph JSON, + # offsets, and counts. Byte equality of SAVE vs SAVE2 is NOT required. + if mode == "save" and byte_drift.calls == 3: + for rank in range(2): + for path in (workspace / f"rank_{rank}").glob("graph_*.cugraph"): + path.write_bytes(path.read_bytes() + b"DRIFT") + + _install_fake_server(g, archive_writer=byte_drift) + + run_id = "run-bytedrift" + baseline = g["run_baseline_phase"](run_id) + evidence = { + "outputs": baseline["outputs"], + "concurrent_outputs": baseline["concurrent_outputs"], + "nccl_transports": baseline["nccl_transports"], + } + out = g["run_save_transaction"](run_id, evidence) + + assert volume.commit_count == 1 + published = Path(out["published_path"]) + assert published.is_dir() + # SAVE2's (drifted) bytes are the published immutable baseline. + archive = published / "archive" + assert out["snapshot"]["byte_fingerprints"] == _snapshot_of(g, archive)["byte_fingerprints"] + + +def test_publish_rejects_symlinked_local_archive(tmp_path: Path) -> None: + g, archive_root, _ = _harness_with_storage(tmp_path) + local_ws = tmp_path / "local_ws" + _build_exact_archive(local_ws) + (local_ws / "rank_0" / "alias.bin").symlink_to(local_ws / "rank_0" / "fatbin_image_packed.img") + snapshot = { + "inventory": {}, + "graph_counts": {}, + "offsets": {}, + "semantic_fingerprints": {}, + "byte_fingerprints": {}, + } + + with pytest.raises(ValueError, match="(?i)symlink"): + g["_publish_archive"]("run-symlink", local_ws, snapshot) + + assert g["archive_volume"].commit_count == 0 + assert not (archive_root / "published" / "run-symlink").exists() + + +def test_transaction_requires_nonzero_offsets(tmp_path: Path) -> None: + g, archive_root, _ = _harness_with_storage(tmp_path) + volume = g["archive_volume"] + + def zero_offset(mode, workspace): + _write_save_archive(mode, workspace) + if mode == "save": + for rank in range(2): + (workspace / f"rank_{rank}" / "final_alloc_offset.json").write_text( + '{"final_alloc_offset": 0}' + ) + + _install_fake_server(g, archive_writer=zero_offset) + run_id = "run-zero" + baseline = g["run_baseline_phase"](run_id) + evidence = { + "outputs": baseline["outputs"], + "concurrent_outputs": baseline["concurrent_outputs"], + "nccl_transports": baseline["nccl_transports"], + } + + with pytest.raises(RuntimeError, match="(?i)non-zero"): + g["run_save_transaction"](run_id, evidence) + assert volume.commit_count == 0 + assert not (archive_root / "published" / run_id).exists() + + +def test_transaction_requires_rank_symmetric_offsets(tmp_path: Path) -> None: + g, archive_root, _ = _harness_with_storage(tmp_path) + volume = g["archive_volume"] + + def asymmetric_offset(mode, workspace): + _write_save_archive(mode, workspace) + if mode == "save": + (workspace / "rank_0" / "final_alloc_offset.json").write_text( + '{"final_alloc_offset": 100}' + ) + (workspace / "rank_1" / "final_alloc_offset.json").write_text( + '{"final_alloc_offset": 200}' + ) + + _install_fake_server(g, archive_writer=asymmetric_offset) + run_id = "run-asym" + baseline = g["run_baseline_phase"](run_id) + evidence = { + "outputs": baseline["outputs"], + "concurrent_outputs": baseline["concurrent_outputs"], + "nccl_transports": baseline["nccl_transports"], + } + + with pytest.raises(RuntimeError, match="(?i)symmetric"): + g["run_save_transaction"](run_id, evidence) + assert volume.commit_count == 0 + assert not (archive_root / "published" / run_id).exists() + + +def test_transaction_requires_phase_nccl_is_exact_p2p_ipc(tmp_path: Path) -> None: + g, archive_root, _ = _harness_with_storage(tmp_path) + volume = g["archive_volume"] + + def save_nccl_log(mode): + # SAVE reports NET/Socket instead of the required P2P/IPC transport. + if mode == "save": + return ( + "[0] NCCL INFO Channel 00/02 : 0 1 via NET/Socket\n" + "[1] NCCL INFO Channel 00/02 : 1 0 via NET/Socket\n" + ) + return _NCCL_P2P_LOG + + _install_fake_server(g, archive_writer=_write_save_archive, log_body_for=save_nccl_log) + run_id = "run-nccl" + baseline = g["run_baseline_phase"](run_id) + evidence = { + "outputs": baseline["outputs"], + "concurrent_outputs": baseline["concurrent_outputs"], + "nccl_transports": baseline["nccl_transports"], + } + + with pytest.raises(RuntimeError, match="(?i)p2p/ipc"): + g["run_save_transaction"](run_id, evidence) + assert volume.commit_count == 0 + assert not (archive_root / "published" / run_id).exists() + + +@pytest.mark.parametrize( + "bad_baseline", + [ + pytest.param({}, id="empty"), + pytest.param({"0": ["NET/Socket"], "1": ["NET/Socket"]}, id="net"), + pytest.param({"0": ["P2P/IPC"]}, id="missing-rank"), + ], +) +def test_transaction_rejects_non_p2p_ipc_baseline(tmp_path: Path, bad_baseline: dict) -> None: + g, archive_root, _ = _harness_with_storage(tmp_path) + volume = g["archive_volume"] + _install_fake_server(g, archive_writer=_write_save_archive) + + run_id = "run-badbaseline" + # Reproduce SAVE2's deterministic outputs so only the baseline transport + # profile is at fault. + evidence = { + "outputs": [{"text": "o:" + p, "token_strings": [p]} for p in g["PROMPTS"]], + "concurrent_outputs": [{"text": "c", "token_strings": ["c"]}], + "nccl_transports": bad_baseline, + } + + with pytest.raises(RuntimeError, match="(?i)baseline nccl"): + g["run_save_transaction"](run_id, evidence) + assert volume.commit_count == 0 + assert not (archive_root / "published" / run_id).exists() + + +def test_commit_retries_then_succeeds_and_indexes(tmp_path: Path) -> None: + g, archive_root, _ = _harness_with_storage(tmp_path) + volume = g["archive_volume"] + index = g["publication_index"] + # First two commit attempts fail; the third (within bounded retries) succeeds. + volume.fail_commits = 2 + _install_fake_server(g, archive_writer=_write_save_archive) + + run_id = "run-commitretry" + baseline = g["run_baseline_phase"](run_id) + evidence = { + "outputs": baseline["outputs"], + "concurrent_outputs": baseline["concurrent_outputs"], + "nccl_transports": baseline["nccl_transports"], + } + + out = g["run_save_transaction"](run_id, evidence) + + published = Path(out["published_path"]) + assert published.is_dir() + assert (published / "manifest.json").is_file() + assert volume.commit_count == 3 # two failures + one success + # The marker is written only after the eventual successful commit. + assert index.put_count == 1 + assert run_id in index + + +def test_publish_is_idempotent_on_identical_rerun(tmp_path: Path) -> None: + g, archive_root, _ = _harness_with_storage(tmp_path) + volume = g["archive_volume"] + index = g["publication_index"] + _install_fake_server(g, archive_writer=_write_save_archive) + + run_id = "run-idem" + baseline = g["run_baseline_phase"](run_id) + evidence = { + "outputs": baseline["outputs"], + "concurrent_outputs": baseline["concurrent_outputs"], + "nccl_transports": baseline["nccl_transports"], + } + + first = g["run_save_transaction"](run_id, evidence) + assert volume.commit_count == 1 + assert index.put_count == 1 + + # Re-running the identical deterministic transaction reconciles idempotently: + # the existing record matches, so no re-publish, commit, or marker write. + second = g["run_save_transaction"](run_id, evidence) + assert second == first + assert volume.commit_count == 1 + assert index.put_count == 1 + + +def test_idempotent_replay_rejects_byte_corrupted_published_archive(tmp_path: Path) -> None: + g, archive_root, _ = _harness_with_storage(tmp_path) + volume = g["archive_volume"] + _install_fake_server(g, archive_writer=_write_save_archive) + + run_id = "run-idem-corrupt" + baseline = g["run_baseline_phase"](run_id) + evidence = { + "outputs": baseline["outputs"], + "concurrent_outputs": baseline["concurrent_outputs"], + "nccl_transports": baseline["nccl_transports"], + } + out = g["run_save_transaction"](run_id, evidence) + assert volume.commit_count == 1 + + # Corrupt a byte of the published archive; the identical replay rebuilds the + # same record but the fresh on-volume snapshot no longer matches, so the + # idempotency guard rejects it rather than silently declaring success. + victim = next((Path(out["published_path"]) / "archive" / "rank_0").glob("graph_*.cugraph")) + victim.write_bytes(victim.read_bytes() + b"CORRUPT") + + with pytest.raises(RuntimeError, match="(?i)published archive"): + g["run_save_transaction"](run_id, evidence) + assert volume.commit_count == 1 # no second commit + + +def test_commit_failure_leaves_no_index(tmp_path: Path) -> None: + g, archive_root, _ = _harness_with_storage(tmp_path) + volume = g["archive_volume"] + index = g["publication_index"] + # Every bounded-retry commit attempt fails. + volume.fail_commits = 10_000 + _install_fake_server(g, archive_writer=_write_save_archive) + + run_id = "run-commitfail" + baseline = g["run_baseline_phase"](run_id) + evidence = { + "outputs": baseline["outputs"], + "concurrent_outputs": baseline["concurrent_outputs"], + "nccl_transports": baseline["nccl_transports"], + } + + # No ExceptionGroup: the plain commit error propagates. + with pytest.raises(RuntimeError, match="commit boom"): + g["run_save_transaction"](run_id, evidence) + + # The run was never marked, so any auto-committed archive bytes are inert. + assert index.put_count == 0 + assert run_id not in index + # Bounded retries only (no rollback re-commit). + assert volume.commit_count == g["COMMIT_MAX_ATTEMPTS"] + + +def test_marker_put_failure_leaves_inert_unindexed_archive(tmp_path: Path) -> None: + g, archive_root, _ = _harness_with_storage(tmp_path) + volume = g["archive_volume"] + index = g["publication_index"] + # The commit succeeds but the marker put fails and cannot be reconciled + # (nothing was actually written). + index.fail_puts = 1 + _install_fake_server(g, archive_writer=_write_save_archive) + + run_id = "run-markerfail" + baseline = g["run_baseline_phase"](run_id) + evidence = { + "outputs": baseline["outputs"], + "concurrent_outputs": baseline["concurrent_outputs"], + "nccl_transports": baseline["nccl_transports"], + } + + with pytest.raises(RuntimeError, match="(?i)did not reconcile"): + g["run_save_transaction"](run_id, evidence) + + # The archive bytes were committed but the run is unindexed and therefore + # inert: a consumer must reject it (proven by the LOAD-unindexed test). + assert volume.commit_count == 1 + assert run_id not in index + assert (archive_root / "published" / run_id / "archive" / "rank_0").is_dir() + + +def test_publish_rejects_conflicting_existing_record(tmp_path: Path) -> None: + g, archive_root, _ = _harness_with_storage(tmp_path) + volume = g["archive_volume"] + index = g["publication_index"] + _install_fake_server(g, archive_writer=_write_save_archive) + + run_id = "run-collide-record" + baseline = g["run_baseline_phase"](run_id) + evidence = { + "outputs": baseline["outputs"], + "concurrent_outputs": baseline["concurrent_outputs"], + "nccl_transports": baseline["nccl_transports"], + } + # A *different* record for this run id already exists: fail the collision. + index.store[run_id] = {"run_id": run_id, "snapshot_digest": "x", "payload": {}} + + with pytest.raises(RuntimeError, match="(?i)collision"): + g["run_save_transaction"](run_id, evidence) + assert volume.commit_count == 0 + assert not (archive_root / "published" / run_id).exists() + + +def test_publish_rejects_existing_canonical_path(tmp_path: Path) -> None: + g, archive_root, _ = _harness_with_storage(tmp_path) + volume = g["archive_volume"] + _install_fake_server(g, archive_writer=_write_save_archive) + + run_id = "run-collide-path" + baseline = g["run_baseline_phase"](run_id) + evidence = { + "outputs": baseline["outputs"], + "concurrent_outputs": baseline["concurrent_outputs"], + "nccl_transports": baseline["nccl_transports"], + } + # A canonical directory already exists on the volume (never overwrite). + existing = archive_root / "published" / run_id + existing.mkdir(parents=True) + (existing / "sentinel").write_text("keep") + + with pytest.raises(FileExistsError): + g["run_save_transaction"](run_id, evidence) + assert volume.commit_count == 0 + assert (existing / "sentinel").read_text() == "keep" + + +def test_publish_archive_refuses_second_publication_of_same_run(tmp_path: Path) -> None: + g, archive_root, _ = _harness_with_storage(tmp_path) + local_ws = tmp_path / "local_ws" + _build_exact_archive(local_ws) + snapshot = _snapshot_of(g, local_ws) + + run_id = "run-twice" + first = g["_publish_archive"](run_id, local_ws, snapshot) + first_bytes = _byte_map(Path(first) / "archive") + + # The canonical run dir already exists: the exclusive mkdir refuses to + # overwrite it (never clobbers a concurrent/prior publication). + with pytest.raises(FileExistsError): + g["_publish_archive"](run_id, local_ws, snapshot) + + assert _byte_map(Path(first) / "archive") == first_bytes + + +# --------------------------------------------------------------------------- +# LOAD: copy published archive locally, execute, compare local before/after. +# --------------------------------------------------------------------------- + + +def _index_entry_for(g: dict, run_dir: Path, snapshot: dict) -> dict: + return { + "snapshot_digest": g["_snapshot_digest"](snapshot), + "published_path": str(run_dir), + "foundry_ref": g["FOUNDRY_REF"], + } + + +def _stored_results(g: dict) -> dict: + """Valid reduced per-phase results (seed/save/save2) for a stored payload.""" + transports = {"0": ["P2P/IPC"], "1": ["P2P/IPC"]} + + def phase_result(phase: str) -> dict: + result = { + "phase": phase, + **g["_candidate_metadata"](), + "process_return_code": 0, + "nccl_transports": transports, + "errors": [], + "runtime_provenance": _valid_runtime_provenance(), + "time_to_healthy_s": 0.5, + "phase_time_s": 1.0, + } + if phase == "seed": + # Seed's stored shape is diagnostic-only (never the canonical SAVE + # reduction): bounded per-rank graph counts and a file sample. + result["seed_diagnostics"] = { + "0": {"json_count": 66, "cugraph_count": 66, "sample_files": ["graph_0.json"]}, + "1": {"json_count": 66, "cugraph_count": 66, "sample_files": ["graph_0.json"]}, + } + else: + result["rank_offsets"] = {"0": 69736595456, "1": 69736595456} + result["archive_graph_counts"] = {"0": 64, "1": 64} + result["archive_fingerprints"] = {"0": {"rank_0/graph_0.json": "aa"}, "1": {}} + if phase == "save2": + result["outputs"] = [ + {"text": f"o:{index}", "token_strings": ["o", str(index)]} for index in range(4) + ] + result["concurrent_outputs"] = [ + {"text": f"c:{index}", "token_strings": ["c", str(index)]} for index in range(4) + ] + return result + + return {phase: phase_result(phase) for phase in ("seed", "save", "save2")} + + +def _record_for(g: dict, run_id: str, run_dir: Path, snapshot: dict) -> dict: + entry = _index_entry_for(g, run_dir, snapshot) + payload = { + "snapshot": snapshot, + "snapshot_digest": entry["snapshot_digest"], + "published_path": str(run_dir), + "published_run_id": run_id, + "results": _stored_results(g), + } + return {"run_id": run_id, **entry, "payload": payload} + + +def _publish_fixture(g: dict, archive_root: Path, run_id: str) -> tuple[Path, dict]: + """Build an indexed published run (archive/ + manifest.json + Dict record).""" + run_dir = archive_root / "published" / run_id + archive = run_dir / "archive" + _build_exact_archive(archive) + snapshot = _snapshot_of(g, archive) + entry = _index_entry_for(g, run_dir, snapshot) + (run_dir / "manifest.json").write_text( + json.dumps({**entry, "snapshot": snapshot}, sort_keys=True) + ) + g["publication_index"].put( + run_id, _record_for(g, run_id, run_dir, snapshot), skip_if_exists=True + ) + return run_dir, snapshot + + +def test_load_copies_published_archive_locally_and_never_writes_volume(tmp_path: Path) -> None: + g, archive_root, local_root = _harness_with_storage(tmp_path) + volume = g["archive_volume"] + run_id = "run-load" + published, snapshot = _publish_fixture(g, archive_root, run_id) + published_before = _byte_map(published) + + _install_fake_server(g, log_body_for=lambda mode: _load_log(snapshot["offsets"])) + + result = g["run_load_phase"]("load", run_id, str(published), snapshot) + + assert volume.commit_count == 0 + # The mounted, published archive was never written by the LOAD. + assert _byte_map(published) == published_before + # The LOAD executed against an independent local copy. + assert (local_root / run_id / "load" / "foundry_archive" / "rank_0").is_dir() + assert result["load_final_alloc_offsets"] == {"0": [69736595456], "1": [69736595456]} + assert result["load_restoration_completions"] == {"0": [[64, 64]], "1": [[64, 64]]} + assert result["captured_offset_count"] == 0 + assert result["saved_manifest_count"] == 0 + + +def test_load_rejects_byte_corrupted_published_archive(tmp_path: Path) -> None: + g, archive_root, local_root = _harness_with_storage(tmp_path) + volume = g["archive_volume"] + run_id = "run-load-corrupt" + published, snapshot = _publish_fixture(g, archive_root, run_id) + _install_fake_server(g, log_body_for=lambda mode: _load_log(snapshot["offsets"])) + + # Corrupt a byte of the on-volume archive without changing its structure: the + # manifest/dir still exist, so only the fresh-snapshot byte comparison catches + # it. The consumer must reject before copying anything locally. + victim = next((published / "archive" / "rank_0").glob("graph_*.cugraph")) + victim.write_bytes(victim.read_bytes() + b"CORRUPT") + + with pytest.raises(RuntimeError, match="(?i)published archive"): + g["run_load_phase"]("load", run_id, str(published), snapshot) + assert volume.commit_count == 0 + assert not (local_root / run_id / "load").exists() + + +def test_load_detects_mutation_during_run(tmp_path: Path) -> None: + g, archive_root, _ = _harness_with_storage(tmp_path) + volume = g["archive_volume"] + run_id = "run-load-mut" + published, snapshot = _publish_fixture(g, archive_root, run_id) + published_before = _byte_map(published) + + def mutating_writer(mode, workspace): + # The server mutates its local copy mid-run (before teardown). + (workspace / "warmup_state.json").write_text('{"memory_pool_config": {"x": 1}}') + + _install_fake_server( + g, + archive_writer=mutating_writer, + log_body_for=lambda mode: _load_log(snapshot["offsets"]), + ) + + with pytest.raises(RuntimeError, match="(?i)mutat"): + g["run_load_phase"]("load", run_id, str(published), snapshot) + + assert volume.commit_count == 0 + assert _byte_map(published) == published_before + + +def test_load_detects_mutation_after_teardown(tmp_path: Path) -> None: + g, archive_root, _ = _harness_with_storage(tmp_path) + volume = g["archive_volume"] + run_id = "run-load-post" + published, snapshot = _publish_fixture(g, archive_root, run_id) + + def post_teardown_mutation(mode, workspace): + # A mutation that lands after the server exits but before the reread. + # It keeps the archive structurally valid, so only the byte map diverges. + (workspace / "warmup_state.json").write_text( + '{"memory_pool_config": {"max_total_num_tokens": 777}}' + ) + + _install_fake_server( + g, + log_body_for=lambda mode: _load_log(snapshot["offsets"]), + stop_hook=post_teardown_mutation, + ) + + with pytest.raises(RuntimeError, match="(?i)mutat"): + g["run_load_phase"]("load", run_id, str(published), snapshot) + + assert volume.commit_count == 0 + + +def test_load_fails_on_live_offset_divergence(tmp_path: Path) -> None: + g, archive_root, _ = _harness_with_storage(tmp_path) + volume = g["archive_volume"] + run_id = "run-load-off" + published, snapshot = _publish_fixture(g, archive_root, run_id) + + diverged = dict(snapshot["offsets"]) + diverged["1"] = 42 + _install_fake_server(g, log_body_for=lambda mode: _load_log(diverged)) + + with pytest.raises(RuntimeError, match="(?i)offset"): + g["run_load_phase"]("load", run_id, str(published), snapshot) + + assert volume.commit_count == 0 + + +def test_server_phase_appends_secondary_errors_without_masking(tmp_path: Path) -> None: + g, _, _ = _harness_with_storage(tmp_path) + _install_fake_server(g) + + def boom(*_a, **_k): + raise RuntimeError("healthcheck boom") + + def kaboom(*_a, **_k): + raise ValueError("scan kaboom") + + g["_wait_until_healthy"] = boom + g["_scan_log"] = kaboom + + run_root = tmp_path / "phase" + result, failures = g["_run_server_phase"]( + "baseline", run_root, str(run_root / "p.log"), None, False + ) + + # The primary phase failure is preserved and ordered before the secondary + # log-scan failure; neither masks the other. + assert failures[0].startswith("phase execution:") + assert "healthcheck boom" in failures[0] + scan_failures = [f for f in failures if f.startswith("log scan:")] + assert scan_failures and "scan kaboom" in scan_failures[0] + assert failures.index(scan_failures[0]) > 0 + # Full diagnostics survive: the empty-scan defaults are still populated. + assert result["errors"] == [] + assert result["load_final_alloc_offsets"] == {} + + +def test_server_phase_provenance_failure_does_not_mask_primary_or_teardown( + tmp_path: Path, +) -> None: + g, _, _ = _harness_with_storage(tmp_path) + _install_fake_server(g) + g["_wait_until_healthy"] = lambda *_a, **_k: (_ for _ in ()).throw(RuntimeError("primary boom")) + g["_runtime_provenance"] = lambda *_a, **_k: (_ for _ in ()).throw( + ValueError("provenance boom") + ) + + run_root = tmp_path / "phase-provenance" + result, failures = g["_run_server_phase"]( + "baseline", + run_root, + str(run_root / "p.log"), + None, + False, + ) + + assert failures[0].startswith("phase execution:") + assert "primary boom" in failures[0] + assert any("runtime provenance: provenance boom" in failure for failure in failures[1:]) + assert result["runtime_provenance"] == {} + assert isinstance(result["phase_time_s"], float) + + +def test_load_rejects_unindexed_background_archive(tmp_path: Path) -> None: + g, archive_root, local_root = _harness_with_storage(tmp_path) + volume = g["archive_volume"] + index = g["publication_index"] + run_id = "run-unindexed" + published, snapshot = _publish_fixture(g, archive_root, run_id) + _install_fake_server(g, log_body_for=lambda mode: _load_log(snapshot["offsets"])) + + # Simulate an auto-committed-but-never-indexed background archive: the bytes + # and manifest are on the volume, but there is no publication-index entry. + del index.store[run_id] + with pytest.raises(RuntimeError, match="(?i)publication index"): + g["run_load_phase"]("load", run_id, str(published), snapshot) + assert not (local_root / run_id / "load").exists() + assert volume.commit_count == 0 + + +@pytest.mark.parametrize("field", ["snapshot_digest", "published_path", "foundry_ref"]) +def test_load_rejects_wrong_index_entry(tmp_path: Path, field: str) -> None: + g, archive_root, local_root = _harness_with_storage(tmp_path) + index = g["publication_index"] + run_id = f"run-badindex-{field}" + published, snapshot = _publish_fixture(g, archive_root, run_id) + _install_fake_server(g, log_body_for=lambda mode: _load_log(snapshot["offsets"])) + + # Corrupt exactly one authoritative field of the record identity. + entry = dict(index[run_id]) + entry[field] = entry[field] + "-tampered" if isinstance(entry[field], str) else "tampered" + index.store[run_id] = entry + + with pytest.raises(RuntimeError, match="(?i)schema/consistency|expected snapshot digest"): + g["run_load_phase"]("load", run_id, str(published), snapshot) + assert not (local_root / run_id / "load").exists() + + +def test_load_rejects_missing_or_mismatched_manifest(tmp_path: Path) -> None: + g, archive_root, local_root = _harness_with_storage(tmp_path) + volume = g["archive_volume"] + run_id = "run-manifest" + published, snapshot = _publish_fixture(g, archive_root, run_id) + _install_fake_server(g, log_body_for=lambda mode: _load_log(snapshot["offsets"])) + + # (a) The index entry is present but the on-volume manifest is missing. + (published / "manifest.json").unlink() + with pytest.raises(RuntimeError, match="(?i)manifest"): + g["run_load_phase"]("load", run_id, str(published), snapshot) + assert not (local_root / run_id / "load").exists() + + # (b) A manifest that does not match the expected snapshot/path/ref. + (published / "manifest.json").write_text(json.dumps({"snapshot": {"offsets": {"0": 1}}})) + with pytest.raises(RuntimeError, match="(?i)manifest"): + g["run_load_phase"]("load", run_id, str(published), snapshot) + + assert volume.commit_count == 0 + + +def test_load_rejects_run_without_archive_child(tmp_path: Path) -> None: + g, archive_root, local_root = _harness_with_storage(tmp_path) + volume = g["archive_volume"] + run_id = "run-noarchive" + published, snapshot = _publish_fixture(g, archive_root, run_id) + _install_fake_server(g, log_body_for=lambda mode: _load_log(snapshot["offsets"])) + + # Index + manifest are present but the archive child is gone: reject. + shutil.rmtree(published / "archive") + with pytest.raises(RuntimeError, match="(?i)archive"): + g["run_load_phase"]("load", run_id, str(published), snapshot) + + assert not (local_root / run_id / "load").exists() + assert volume.commit_count == 0 + + +def test_load_rejects_non_load_phase(tmp_path: Path) -> None: + g, archive_root, _ = _harness_with_storage(tmp_path) + published, snapshot = _publish_fixture(g, archive_root, "run-load-bad") + + for bad in ("baseline", "seed", "save", "save2"): + with pytest.raises(ValueError): + g["run_load_phase"](bad, "run-load-bad", str(published), snapshot) + assert g["archive_volume"].commit_count == 0 + + +def test_load_reloads_volume_before_reading_index(tmp_path: Path) -> None: + g, archive_root, _ = _harness_with_storage(tmp_path) + volume = g["archive_volume"] + run_id = "run-reload" + published, snapshot = _publish_fixture(g, archive_root, run_id) + _install_fake_server(g, log_body_for=lambda mode: _load_log(snapshot["offsets"])) + + volume.events.clear() + g["run_load_phase"]("load", run_id, str(published), snapshot) + + # The warm container reloads the volume before it reads the Dict record. + assert volume.reload_count == 1 + assert volume.events[:2] == ["reload", "read"] + + +def test_load_reload_reveals_stale_index_record(tmp_path: Path) -> None: + g, archive_root, _ = _harness_with_storage(tmp_path) + volume = g["archive_volume"] + index = g["publication_index"] + run_id = "run-stale" + run_dir = archive_root / "published" / run_id + archive = run_dir / "archive" + _build_exact_archive(archive) + snapshot = _snapshot_of(g, archive) + entry = _index_entry_for(g, run_dir, snapshot) + (run_dir / "manifest.json").write_text( + json.dumps({**entry, "snapshot": snapshot}, sort_keys=True) + ) + record = _record_for(g, run_id, run_dir, snapshot) + + # Warm-container stale view: the record is only visible after reload(). + def reveal() -> None: + volume.reload_count += 1 + index.store[run_id] = record + + volume.reload = reveal + _install_fake_server(g, log_body_for=lambda mode: _load_log(snapshot["offsets"])) + + # The LOAD succeeds only because reload runs first and reveals the record. + result = g["run_load_phase"]("load", run_id, str(run_dir), snapshot) + assert volume.reload_count == 1 + assert run_id in index.store + assert result["load_restoration_completions"] == {"0": [[64, 64]], "1": [[64, 64]]} + + +# --------------------------------------------------------------------------- +# Atomic marker + reconciliation (put skip_if_exists, ambiguous put, main). +# --------------------------------------------------------------------------- + + +class _FakeRemote: + """Stand-in for a Modal function whose ``.remote`` call is controllable.""" + + def __init__(self, fn) -> None: + self._fn = fn + + def remote(self, *args, **kwargs): + return self._fn(*args, **kwargs) + + +def _wellformed_snapshot() -> dict: + """A hand-built snapshot passing the full schema (rank keys, counts, offsets).""" + return { + "inventory": { + "0": {"graph_count": 64, "indices": [0], "stems": {"0": "graph_0"}}, + "1": {"graph_count": 64, "indices": [0], "stems": {"0": "graph_0"}}, + }, + "graph_counts": {"0": 64, "1": 64}, + "offsets": {"0": 69736595456, "1": 69736595456}, + "semantic_fingerprints": { + "0": {"rank_0/graph_0.json": "aa"}, + "1": {"rank_1/graph_0.json": "bb"}, + }, + "byte_fingerprints": {"rank_0/graph_0.cugraph": "cc", "warmup_state.json": "dd"}, + } + + +def _consistent_record(g: dict, run_id: str, published_path: str | None = None) -> dict: + # The canonical path is derived from run_id; a consistent record must name it. + if published_path is None: + published_path = str(g["_published_dir"](run_id)) + snapshot = _wellformed_snapshot() + digest = g["_snapshot_digest"](snapshot) + payload = { + "snapshot": snapshot, + "snapshot_digest": digest, + "published_path": published_path, + "published_run_id": run_id, + "results": _stored_results(g), + } + return { + "run_id": run_id, + "snapshot_digest": digest, + "published_path": published_path, + "foundry_ref": g["FOUNDRY_REF"], + "payload": payload, + } + + +def test_marker_put_false_identical_is_idempotent(tmp_path: Path) -> None: + g, *_ = _harness_with_storage(tmp_path) + index = g["publication_index"] + record = _consistent_record(g, "rid") + assert index.put("rid", record, skip_if_exists=True) is True + + # A second write of the identical record reconciles idempotently (put False). + g["_write_publication_marker"]("rid", record) + assert index.store["rid"] == record + + +def test_marker_put_false_different_is_collision(tmp_path: Path) -> None: + g, *_ = _harness_with_storage(tmp_path) + index = g["publication_index"] + index.store["rid"] = _consistent_record(g, "rid", published_path="/vol/published/OTHER") + + with pytest.raises(RuntimeError, match="(?i)collision"): + g["_write_publication_marker"]("rid", _consistent_record(g, "rid")) + + +def test_marker_ambiguous_put_reconciles_when_record_present(tmp_path: Path) -> None: + g, *_ = _harness_with_storage(tmp_path) + index = g["publication_index"] + record = _consistent_record(g, "rid") + # The record is already present, but put raises an ambiguous error. + index.store["rid"] = record + index.fail_puts = 1 + + # No raise: reconciliation finds the exact record. + g["_write_publication_marker"]("rid", record) + + +def test_marker_ambiguous_put_preserves_error_when_absent(tmp_path: Path) -> None: + g, *_ = _harness_with_storage(tmp_path) + index = g["publication_index"] + index.fail_puts = 1 # put raises and nothing was written + + with pytest.raises(RuntimeError, match="(?i)did not reconcile"): + g["_write_publication_marker"]("rid", _consistent_record(g, "rid")) + + +def test_reconcile_returns_payload_for_consistent_record(tmp_path: Path) -> None: + g, *_ = _harness_with_storage(tmp_path) + index = g["publication_index"] + record = _consistent_record(g, "rid") + index.store["rid"] = record + + out = g["_reconcile_transaction"]("rid", RuntimeError("delivery boom")) + assert out == record["payload"] + + +def test_reconcile_wraps_when_no_marker_preserving_cause(tmp_path: Path) -> None: + g, *_ = _harness_with_storage(tmp_path) + original = RuntimeError("genuine boom") + + with pytest.raises(g["PublicationReconciliationError"]) as excinfo: + g["_reconcile_transaction"]("missing", original) + + # The original remote error is preserved as __cause__; detail names the run. + assert excinfo.value.__cause__ is original + assert "missing" in str(excinfo.value) + + +def test_reconcile_wraps_on_inconsistent_record_preserving_cause(tmp_path: Path) -> None: + g, *_ = _harness_with_storage(tmp_path) + index = g["publication_index"] + record = _consistent_record(g, "rid") + record["snapshot_digest"] = "wrong-digest" + index.store["rid"] = record + original = RuntimeError("delivery boom") + + with pytest.raises(g["PublicationReconciliationError"]) as excinfo: + g["_reconcile_transaction"]("rid", original) + assert excinfo.value.__cause__ is original + + +def test_reconcile_wraps_on_missing_payload_fields(tmp_path: Path) -> None: + g, *_ = _harness_with_storage(tmp_path) + index = g["publication_index"] + record = _consistent_record(g, "rid") + del record["payload"]["results"] # incomplete payload schema + index.store["rid"] = record + original = RuntimeError("delivery boom") + + with pytest.raises(g["PublicationReconciliationError"]) as excinfo: + g["_reconcile_transaction"]("rid", original) + assert excinfo.value.__cause__ is original + + +@pytest.mark.parametrize( + "mutate", + [ + lambda r: r["payload"]["results"].pop("save2"), + lambda r: r["payload"]["results"].__setitem__("extra", {}), + lambda r: r["payload"]["results"]["save"].pop("rank_offsets"), + lambda r: r["payload"]["results"]["save"].__setitem__("errors", "boom"), + lambda r: r["payload"]["results"]["save2"].pop("outputs"), + lambda r: r["payload"]["results"]["save2"].__setitem__( + "outputs", [{"text": "only-one", "token_strings": ["only-one"]}] + ), + lambda r: r["payload"]["results"]["save2"].__setitem__( + "concurrent_outputs", + [ + {"text": f"c:{index}", "token_strings": ([] if index == 2 else ["c"])} + for index in range(4) + ], + ), + lambda r: r["payload"]["results"]["save"].pop("foundry_ref"), + lambda r: r["payload"]["results"]["save"].__setitem__("vllm_ref", "f" * 40), + lambda r: r["payload"]["results"]["save"].__setitem__("phase_time_s", 0.0), + lambda r: r["payload"]["results"]["save"]["runtime_provenance"].__setitem__( + "nccl_runtime_version", None + ), + lambda r: r["payload"]["results"]["save"]["runtime_provenance"].__setitem__( + "gpu_names", ["NVIDIA H100"] + ), + lambda r: r["payload"]["results"]["seed"].__setitem__("process_return_code", True), + lambda r: r["payload"]["results"]["seed"].__setitem__("phase", "nope"), + lambda r: r["payload"].__setitem__("results", "not-a-dict"), + ], +) +def test_record_is_consistent_rejects_bad_results(tmp_path: Path, mutate) -> None: + g, *_ = _harness_with_storage(tmp_path) + record = _consistent_record(g, "rid") + assert g["_record_is_consistent"](record, "rid") is True + + mutate(record) + assert g["_record_is_consistent"](record, "rid") is False + + +@pytest.mark.parametrize( + "mutate", + [ + lambda s: s.pop("byte_fingerprints"), + lambda s: s.__setitem__("byte_fingerprints", {}), + lambda s: s["graph_counts"].__setitem__("0", 63), + lambda s: s["offsets"].__setitem__("0", 0), + lambda s: s["offsets"].__setitem__("0", -5), + lambda s: s["offsets"].__setitem__("0", True), + lambda s: s["graph_counts"].pop("1"), + lambda s: s.__setitem__("inventory", ["not", "a", "dict"]), + lambda s: s["semantic_fingerprints"].__setitem__("0", "not-a-dict"), + ], +) +def test_snapshot_wellformed_rejects_bad_snapshot(tmp_path: Path, mutate) -> None: + g, *_ = _harness_with_storage(tmp_path) + snapshot = _wellformed_snapshot() + assert g["_snapshot_is_wellformed"](snapshot) is True + + mutate(snapshot) + assert g["_snapshot_is_wellformed"](snapshot) is False + + +def test_reconcile_does_not_mask_original_on_malformed_nested_data(tmp_path: Path) -> None: + g, *_ = _harness_with_storage(tmp_path) + index = g["publication_index"] + record = _consistent_record(g, "rid") + # Malformed nested data that passes the shallow schema but is not + # JSON-serializable: the digest recomputation must be swallowed (return False), + # not raised, so reconciliation still wraps and preserves the original error. + record["payload"]["snapshot"]["inventory"]["0"] = {"stems": {1, 2, 3}} + index.store["rid"] = record + original = RuntimeError("delivery boom") + + with pytest.raises(g["PublicationReconciliationError"]) as excinfo: + g["_reconcile_transaction"]("rid", original) + assert excinfo.value.__cause__ is original + + +def test_main_recovers_when_marker_landed_but_remote_raised(tmp_path: Path) -> None: + g, *_ = _harness_with_storage(tmp_path) + index = g["publication_index"] + record = _consistent_record(g, "rid") + # The transaction wrote its marker server-side, then result delivery failed. + index.store["rid"] = record + + def deliver_fail(run_id, evidence): + raise RuntimeError("result delivery boom") + + g["run_save_transaction"] = _FakeRemote(deliver_fail) + + out = g["_obtain_transaction"]("rid", {}) + assert out == record["payload"] + + +def test_main_wraps_genuine_failure_without_marker_preserving_cause(tmp_path: Path) -> None: + g, *_ = _harness_with_storage(tmp_path) + + def hard_fail(run_id, evidence): + raise RuntimeError("genuine transaction boom") + + g["run_save_transaction"] = _FakeRemote(hard_fail) + + with pytest.raises(g["PublicationReconciliationError"]) as excinfo: + g["_obtain_transaction"]("missing", {}) + assert isinstance(excinfo.value.__cause__, RuntimeError) + assert "genuine transaction boom" in str(excinfo.value.__cause__) + + +@pytest.mark.parametrize( + ("mutate", "why"), + [ + (lambda r: r.__setitem__("run_id", "other"), "record run_id"), + (lambda r: r.__setitem__("foundry_ref", "not-the-ref"), "record ref"), + (lambda r: r.__setitem__("published_path", "/tmp/evil"), "record path"), + (lambda r: r.__setitem__("snapshot_digest", "deadbeef"), "record digest"), + (lambda r: r["payload"].__setitem__("published_run_id", "other"), "payload run_id"), + (lambda r: r["payload"].__setitem__("published_path", "/tmp/evil"), "payload path"), + (lambda r: r["payload"].__setitem__("snapshot_digest", "deadbeef"), "payload digest"), + (lambda r: r["payload"].pop("results"), "missing payload field"), + (lambda r: r.pop("payload"), "missing payload"), + ], +) +def test_record_is_consistent_rejects_mismatches(tmp_path: Path, mutate, why: str) -> None: + g, *_ = _harness_with_storage(tmp_path) + record = _consistent_record(g, "rid") + assert g["_record_is_consistent"](record, "rid") is True + + mutate(record) + assert g["_record_is_consistent"](record, "rid") is False, why + + +def test_delete_published_run_hook_removes_bytes_then_pops_marker(tmp_path: Path) -> None: + g, archive_root, _ = _harness_with_storage(tmp_path) + volume = g["archive_volume"] + index = g["publication_index"] + run_id = "run-delete" + published, _snapshot = _publish_fixture(g, archive_root, run_id) + assert published.exists() + assert run_id in index.store + volume.events.clear() + + # Task 7 cleanup contract: delete both the archive bytes and the Dict record. + g["delete_published_run"](run_id) + + assert not published.exists() + assert run_id not in index.store + assert volume.reload_count == 1 + assert volume.commit_count == 1 + # Order: reload the volume, commit the deletion, then pop the marker last. + durable = [e for e in volume.events if e in ("reload", "commit", "pop")] + assert durable == ["reload", "commit", "pop"] + + +def test_delete_published_run_refuses_unindexed_orphan(tmp_path: Path) -> None: + g, archive_root, _ = _harness_with_storage(tmp_path) + volume = g["archive_volume"] + run_id = "run-del-orphan" + published, _snapshot = _publish_fixture(g, archive_root, run_id) + # Drop the marker while the bytes remain: an unindexed orphan. Cleanup must + # fail closed rather than delete an unvalidated tree. + del g["publication_index"].store[run_id] + + with pytest.raises(RuntimeError, match="(?i)orphan"): + g["delete_published_run"](run_id) + assert published.exists() + assert volume.commit_count == 0 + + +def test_delete_fully_absent_is_idempotent_success(tmp_path: Path) -> None: + g, archive_root, _ = _harness_with_storage(tmp_path) + volume = g["archive_volume"] + index = g["publication_index"] + run_id = "run-never-published" + + # Neither a marker nor canonical bytes exist: cleanup is a no-op success, and + # repeating it stays successful. + g["delete_published_run"](run_id) + g["delete_published_run"](run_id) + + assert run_id not in index.store + assert not (archive_root / "published" / run_id).exists() + assert volume.commit_count == 0 + + +def test_delete_pop_write_then_ack_lost_then_retry(tmp_path: Path) -> None: + g, archive_root, _ = _harness_with_storage(tmp_path) + volume = g["archive_volume"] + index = g["publication_index"] + run_id = "run-del-acklost" + published, _snapshot = _publish_fixture(g, archive_root, run_id) + + # First attempt: rmtree + commit succeed, and the pop removes the marker but + # its acknowledgement is lost (the call raises after the write committed). + index.fail_pops_after_write = 1 + with pytest.raises(RuntimeError, match="pop boom"): + g["delete_published_run"](run_id) + assert not published.exists() + assert run_id not in index.store # the pop actually committed + assert volume.commit_count == 1 + + # Retry: both the marker and the bytes are absent, so cleanup succeeds + # idempotently without another commit or pop. + g["delete_published_run"](run_id) + assert run_id not in index.store + assert volume.commit_count == 1 + + +def test_delete_published_run_never_trusts_record_path(tmp_path: Path) -> None: + g, archive_root, _ = _harness_with_storage(tmp_path) + index = g["publication_index"] + run_id = "run-malicious" + # A hostile external directory the record tries to redirect deletion at. + victim = archive_root.parent / "victim" + victim.mkdir(parents=True) + (victim / "keep.txt").write_text("precious") + + # A record whose paths point at the victim, keyed under a run whose canonical + # dir does not even exist. Cleanup derives the target from run_id only. + record = _consistent_record(g, run_id, published_path=str(victim)) + index.store[run_id] = record + + with pytest.raises(RuntimeError, match="(?i)refusing to delete"): + g["delete_published_run"](run_id) + # The victim tree is untouched: the record path was never used as a target. + assert (victim / "keep.txt").read_text() == "precious" + + +def test_delete_published_run_rejects_stale_record_missing_archive(tmp_path: Path) -> None: + g, archive_root, _ = _harness_with_storage(tmp_path) + volume = g["archive_volume"] + index = g["publication_index"] + run_id = "run-del-stale" + published, _snapshot = _publish_fixture(g, archive_root, run_id) + # The record is present and consistent, but the on-volume archive is gone. + shutil.rmtree(published / "archive") + + with pytest.raises(RuntimeError, match="(?i)archive"): + g["delete_published_run"](run_id) + # Nothing was committed or popped: the marker remains for a later retry. + assert volume.commit_count == 0 + assert run_id in index.store + + +def test_delete_published_run_commit_failure_keeps_marker(tmp_path: Path) -> None: + g, archive_root, _ = _harness_with_storage(tmp_path) + volume = g["archive_volume"] + index = g["publication_index"] + run_id = "run-del-commitfail" + published, _snapshot = _publish_fixture(g, archive_root, run_id) + volume.fail_commits = 10_000 # every bounded-retry commit fails + + with pytest.raises(RuntimeError, match="commit boom"): + g["delete_published_run"](run_id) + + # The bytes were removed but the deletion could not be committed durably, so + # the marker MUST remain (the run stays discoverable/re-cleanable). + assert not published.exists() + assert run_id in index.store + assert index.pop_count == 0 + + +def test_delete_published_run_pop_failure_leaves_bytes_gone(tmp_path: Path) -> None: + g, archive_root, _ = _harness_with_storage(tmp_path) + volume = g["archive_volume"] + index = g["publication_index"] + run_id = "run-del-popfail" + published, _snapshot = _publish_fixture(g, archive_root, run_id) + index.fail_pops = 1 # bytes removed + committed, but the final pop fails + + with pytest.raises(RuntimeError, match="pop boom"): + g["delete_published_run"](run_id) + + # Bytes are gone and committed; the marker lingers but consumers reject the + # run because the archive is missing (proven by the missing-archive test). + assert not published.exists() + assert volume.commit_count == 1 + assert run_id in index.store + + +def test_delete_published_run_retry_after_pop_failure_succeeds(tmp_path: Path) -> None: + g, archive_root, _ = _harness_with_storage(tmp_path) + volume = g["archive_volume"] + index = g["publication_index"] + run_id = "run-del-retry" + published, _snapshot = _publish_fixture(g, archive_root, run_id) + + # First attempt: bytes removed + committed, but the marker pop fails. + index.fail_pops = 1 + with pytest.raises(RuntimeError, match="pop boom"): + g["delete_published_run"](run_id) + assert not published.exists() + assert run_id in index.store + assert volume.commit_count == 1 + + # Retry: the canonical dir is already absent and the marker is still + # consistent, so cleanup pops it idempotently and succeeds without a second + # commit or an rmtree. + g["delete_published_run"](run_id) + assert run_id not in index.store + assert volume.commit_count == 1 + + +def _completion_records(prefix: str) -> list[dict[str, object]]: + return [ + {"text": f"{prefix}-{index}", "token_strings": [prefix, str(index)]} for index in range(4) + ] + + +def _complete_query_results() -> dict[str, dict[str, object]]: + sequential = _completion_records("sequential") + concurrent = _completion_records("concurrent") + return { + phase: { + "outputs": [dict(record) for record in sequential], + "concurrent_outputs": [dict(record) for record in concurrent], + } + for phase in ("baseline", "save2", "load", "load2") + } + + +def test_completion_matrix_requires_exact_four_nonempty_records_in_every_query_phase( + tmp_path: Path, +) -> None: + g, *_ = _harness_with_storage(tmp_path) + results = _complete_query_results() + + assert g["_completion_matrix_failures"](results) == [] + + results["load"]["outputs"] = results["load"]["outputs"][:3] + results["load2"]["concurrent_outputs"][2]["token_strings"] = [] + failures = " ".join(g["_completion_matrix_failures"](results)) + assert "load sequential record count" in failures + assert "load2 concurrent record 2" in failures + + +@pytest.mark.parametrize( + ("phase", "mode", "field", "bad_value"), + [ + ("baseline", "outputs", "text", ""), + ("save2", "outputs", "token_strings", []), + ("load", "concurrent_outputs", "text", ""), + ("load2", "concurrent_outputs", "token_strings", [""]), + ], +) +def test_completion_matrix_rejects_empty_evidence_and_phase_divergence( + tmp_path: Path, + phase: str, + mode: str, + field: str, + bad_value: object, +) -> None: + g, *_ = _harness_with_storage(tmp_path) + results = _complete_query_results() + results[phase][mode][1][field] = bad_value + + failures = g["_completion_matrix_failures"](results) + assert failures + assert any(phase in failure for failure in failures) + + +def test_runtime_provenance_uses_metadata_and_nvidia_smi_without_importing_cuda( + tmp_path: Path, +) -> None: + g, *_ = _harness_with_storage(tmp_path) + package_calls: list[tuple[str, ...]] = [] + smi_calls: list[str] = [] + + def package_version(*names: str) -> str | None: + package_calls.append(names) + return {"torch": "2.11.0+cu130", "vllm": "0.15.0"}[names[0]] + + def nvidia_smi(query: str) -> list[str]: + smi_calls.append(query) + return { + "driver_version": ["580.95.05", "580.95.05"], + "name": ["NVIDIA H100 80GB HBM3", "NVIDIA H100 80GB HBM3"], + }[query] + + g["_package_version"] = package_version + g["_nvidia_smi"] = nvidia_smi + provenance = g["_runtime_provenance"]("2.28.9") + + assert provenance == { + "actual_torch_version": "2.11.0+cu130", + "actual_vllm_version": "0.15.0", + "driver_version": "580.95.05", + "gpu_count": 2, + "gpu_names": ["NVIDIA H100 80GB HBM3", "NVIDIA H100 80GB HBM3"], + "nccl_runtime_version": "2.28.9", + } + assert package_calls == [("torch",), ("vllm",)] + assert smi_calls == ["driver_version", "name"] + + +def test_candidate_metadata_contains_every_immutable_requested_pin(tmp_path: Path) -> None: + g, *_ = _harness_with_storage(tmp_path) + + metadata = g["_candidate_metadata"]() + + assert metadata["foundry_ref"] == g["FOUNDRY_REF"] + assert metadata["vllm_ref"] == g["VLLM_REF"] + assert metadata["vllm_wheel_commit"] == g["VLLM_WHEEL_COMMIT"] + assert metadata["model_revision"] == g["MODEL_REVISION"] + assert metadata["cuda_image_tag"] == g["CUDA_TAG"] + assert metadata["requested_torch_version"] == "2.11.0" + assert metadata["requested_gpu"] == g["MODAL_GPU"] + + +def _valid_runtime_provenance() -> dict[str, object]: + return { + "actual_torch_version": "2.11.0+cu130", + "actual_vllm_version": "0.15.0", + "driver_version": "580.95.05", + "gpu_count": 2, + "gpu_names": ["NVIDIA H100", "NVIDIA H100"], + "nccl_runtime_version": "2.28.9", + } + + +def _valid_phase_provenance(g: dict, phase: str = "baseline") -> dict[str, object]: + return { + "phase": phase, + **g["_candidate_metadata"](), + "runtime_provenance": _valid_runtime_provenance(), + "phase_time_s": 1.0, + "time_to_healthy_s": 0.5, + } + + +@pytest.mark.parametrize( + "mutate", + [ + lambda result: result.pop("foundry_ref"), + lambda result: result.__setitem__("vllm_ref", "f" * 40), + lambda result: result.__setitem__("vllm_wheel_commit", "short"), + lambda result: result.__setitem__("model_revision", "main"), + lambda result: result.__setitem__("tp_size", 1), + lambda result: result.__setitem__("cuda_image_tag", ""), + lambda result: result.__setitem__("requested_torch_version", ""), + lambda result: result.__setitem__("phase_time_s", 0.0), + lambda result: result["runtime_provenance"].__setitem__("actual_torch_version", None), + lambda result: result["runtime_provenance"].__setitem__("actual_vllm_version", ""), + lambda result: result["runtime_provenance"].__setitem__("nccl_runtime_version", None), + lambda result: result["runtime_provenance"].__setitem__("driver_version", ""), + lambda result: result["runtime_provenance"].__setitem__("gpu_count", 1), + lambda result: result["runtime_provenance"].__setitem__("gpu_names", ["NVIDIA H100", ""]), + ], +) +def test_phase_provenance_requires_exact_candidate_and_complete_actual_runtime( + tmp_path: Path, + mutate, +) -> None: + g, *_ = _harness_with_storage(tmp_path) + result = _valid_phase_provenance(g) + assert g["_provenance_failures"](result) == [] + + mutate(result) + + assert g["_provenance_failures"](result) + + +def test_incomplete_runtime_provenance_prevents_transaction_publication( + tmp_path: Path, +) -> None: + g, archive_root, _ = _harness_with_storage(tmp_path) + _install_fake_server(g, archive_writer=_write_save_archive) + g["_runtime_provenance"] = lambda _nccl: { + **_valid_runtime_provenance(), + "actual_vllm_version": None, + } + baseline = { + "outputs": [{"text": f"o:{prompt}", "token_strings": [prompt]} for prompt in g["PROMPTS"]], + "concurrent_outputs": [ + {"text": f"c:{index}", "token_strings": ["c", str(index)]} for index in range(4) + ], + "nccl_transports": _P2P_TRANSPORTS, + } + + with pytest.raises(g["PhaseFailure"], match="runtime provenance"): + g["run_save_transaction"]("run-incomplete-runtime", baseline) + + assert g["archive_volume"].commit_count == 0 + assert g["publication_index"].put_count == 0 + assert not (archive_root / "published" / "run-incomplete-runtime").exists() + + +def test_cleanup_remote_reconciles_lost_pop_ack_as_completed(tmp_path: Path) -> None: + g, archive_root, _ = _harness_with_storage(tmp_path) + run_id = "run-cleanup-ack-lost" + published, _snapshot = _publish_fixture(g, archive_root, run_id) + index = g["publication_index"] + index.fail_pops_after_write = 1 + + outcome = g["cleanup_published_run"](run_id) + + assert outcome["cleanup_attempted"] is True + assert outcome["cleanup_completed"] is True + assert outcome["artifacts_retained"] is False + assert outcome["marker_present"] is False + assert outcome["published_path_present"] is False + assert outcome["reconciled_error"] == "pop boom" + assert not published.exists() + + +def test_cleanup_remote_reports_retained_indexed_artifacts_on_refusal(tmp_path: Path) -> None: + g, archive_root, _ = _harness_with_storage(tmp_path) + run_id = "run-cleanup-refused" + published, _snapshot = _publish_fixture(g, archive_root, run_id) + index = g["publication_index"] + index.store[run_id]["foundry_ref"] = "tampered" + + outcome = g["cleanup_published_run"](run_id) + + assert outcome["cleanup_attempted"] is True + assert outcome["cleanup_completed"] is False + assert outcome["artifacts_retained"] is True + assert outcome["marker_present"] is True + assert outcome["published_path_present"] is True + assert "not self-consistent" in outcome["cleanup_error"] + assert published.exists() + + +def test_cleanup_remote_reports_reconciliation_state_error_type(tmp_path: Path) -> None: + g, *_ = _harness_with_storage(tmp_path) + g["delete_published_run"] = lambda _run_id: None + g["archive_volume"].reload = lambda: (_ for _ in ()).throw(PermissionError("state unreadable")) + + outcome = g["cleanup_published_run"]("run-state-error") + + assert outcome["cleanup_completed"] is False + assert outcome["artifacts_retained"] is True + assert outcome["cleanup_error"] == "state unreadable" + assert outcome["cleanup_error_type"] == "builtins.PermissionError" + assert outcome["marker_present"] is None + assert outcome["published_path_present"] is None + + +def test_main_uses_dashboard_url_authoritative_payload_and_final_cleanup_stage() -> None: + tree = ast.parse(HARNESS_PATH.read_text()) + main = next( + node for node in ast.walk(tree) if isinstance(node, ast.FunctionDef) and node.name == "main" + ) + source = ast.get_source_segment(HARNESS_PATH.read_text(), main) + assert source is not None + + assert "app.get_dashboard_url()" in source + assert "_authoritative_transaction_payload(" in source + assert "cleanup_published_run.remote(" in source + assert source.index("if failed:") < source.index("cleanup_published_run.remote(") + assert source.index("cleanup_published_run.remote(") < source.index('"status": "VERIFIED"') + + +def test_cleanup_function_has_lightweight_helper_mount_and_archive_volume() -> None: + source = HARNESS_PATH.read_text() + + assert "modal.Image.debian_slim()" in source + assert '"vllm_tp_validation"' in source + assert '"tp_validation_helpers"' in source + assert "@app.function(" in source + assert "image=cleanup_image" in source + assert "volumes={ARCHIVE_ROOT: archive_volume}" in source + + +class _LifecycleRemote: + def __init__(self, fn) -> None: + self.fn = fn + self.calls: list[tuple[object, ...]] = [] + + def remote(self, *args): + self.calls.append(args) + return self.fn(*args) + + +def _lifecycle_transaction(run_id: str) -> dict[str, object]: + return { + "snapshot": { + "inventory": {}, + "graph_counts": {}, + "offsets": {}, + "semantic_fingerprints": {}, + "byte_fingerprints": {}, + }, + "snapshot_digest": "digest", + "published_run_id": run_id, + "published_path": f"/published/{run_id}", + "results": {"seed": {}, "save": {}, "save2": {}}, + } + + +def _install_lifecycle_fakes( + g: dict, + *, + fail_stage: str | None = None, + checks: dict[str, bool] | None = None, + cleanup_error: Exception | None = None, + inspection_error: Exception | None = None, +) -> tuple[_LifecycleRemote, _LifecycleRemote]: + run_id = "lifecycle-run" + baseline = { + "outputs": _completion_records("sequential"), + "concurrent_outputs": _completion_records("concurrent"), + } + transaction = _lifecycle_transaction(run_id) + + def baseline_call(_run_id: str): + if fail_stage == "baseline": + raise RuntimeError("baseline failed") + return baseline + + def obtain(_run_id: str, _evidence: dict): + if fail_stage == "transaction": + raise RuntimeError("transaction failed") + return transaction + + def load_call(phase: str, *_args): + if fail_stage == phase: + raise RuntimeError(f"{phase} failed") + return {"phase": phase} + + def cleanup_call(_run_id: str): + if cleanup_error is not None: + raise cleanup_error + return { + "cleanup_attempted": True, + "cleanup_completed": True, + "artifacts_retained": False, + "marker_present": False, + "published_path_present": False, + "cleanup_error": None, + "cleanup_error_type": None, + "reconciled_error": None, + } + + def inspect_call(_run_id: str): + if inspection_error is not None: + raise inspection_error + return { + "marker_present": fail_stage != "baseline", + "published_path_present": fail_stage != "baseline", + "artifacts_retained": fail_stage != "baseline", + } + + cleanup = _LifecycleRemote(cleanup_call) + inspector = _LifecycleRemote(inspect_call) + g["_new_run_id"] = lambda: run_id + g["run_baseline_phase"] = _LifecycleRemote(baseline_call) + g["_obtain_transaction"] = obtain + g["_authoritative_transaction_payload"] = lambda _rid, payload: payload + g["run_load_phase"] = _LifecycleRemote(load_call) + g["_acceptance_checks"] = lambda _results, _transaction: checks or {"accepted": True} + g["cleanup_published_run"] = cleanup + g["inspect_published_run"] = inspector + return cleanup, inspector + + +@pytest.mark.parametrize("failed_stage", ["baseline", "transaction", "load", "load2"]) +def test_main_phase_error_emits_exactly_one_structured_not_verified_report( + tmp_path: Path, + capsys: pytest.CaptureFixture[str], + failed_stage: str, +) -> None: + g, *_ = _harness_with_storage(tmp_path) + cleanup, inspector = _install_lifecycle_fakes(g, fail_stage=failed_stage) + + with pytest.raises(RuntimeError, match="TP validation NOT VERIFIED"): + g["main"]() + + report = json.loads(capsys.readouterr().out) + assert report["status"] == "NOT VERIFIED" + assert report["run_url"] == "https://modal.test/apps/vllm" + assert report["failed_stage"] == failed_stage + assert report["failed_error_type"] == "builtins.RuntimeError" + assert report["cleanup_attempted"] is False + assert report["cleanup_completed"] is False + assert report["artifacts_retained"] is (failed_stage != "baseline") + assert ( + set(report["results"]) + == { + "baseline": set(), + "transaction": {"baseline"}, + "load": {"baseline", "seed", "save", "save2"}, + "load2": {"baseline", "seed", "save", "save2", "load"}, + }[failed_stage] + ) + assert cleanup.calls == [] + assert inspector.calls == [("lifecycle-run",)] + + +def test_main_labels_publication_reconciliation_failure_and_preserves_cause( + tmp_path: Path, + capsys: pytest.CaptureFixture[str], +) -> None: + g, *_ = _harness_with_storage(tmp_path) + _cleanup, _inspector = _install_lifecycle_fakes(g) + original = RuntimeError("save transaction failed") + reconciliation = g["PublicationReconciliationError"]("marker absent") + reconciliation.__cause__ = original + g["_obtain_transaction"] = lambda *_a: (_ for _ in ()).throw(reconciliation) + + with pytest.raises(RuntimeError, match="TP validation NOT VERIFIED") as excinfo: + g["main"]() + + report = json.loads(capsys.readouterr().out) + assert report["failed_stage"] == "reconciliation" + assert report["failed_error_type"].endswith("PublicationReconciliationError") + assert excinfo.value.__cause__ is reconciliation + assert reconciliation.__cause__ is original + + +def test_main_includes_phase_failure_partial_result_in_final_report( + tmp_path: Path, + capsys: pytest.CaptureFixture[str], +) -> None: + g, *_ = _harness_with_storage(tmp_path) + _cleanup, _inspector = _install_lifecycle_fakes(g) + partial = {"phase": "load", "evidence_errors": ["request failed"]} + + def load_with_partial(phase: str, *_args): + if phase == "load": + raise g["PhaseFailure"]("load", partial, ["request failed"]) + return {"phase": phase} + + g["run_load_phase"] = _LifecycleRemote(load_with_partial) + + with pytest.raises(RuntimeError, match="TP validation NOT VERIFIED"): + g["main"]() + + report = json.loads(capsys.readouterr().out) + assert report["failed_stage"] == "load" + assert report["results"]["load"] == partial + assert report["cleanup_attempted"] is False + + +def test_main_aggregate_failure_reports_and_never_cleans( + tmp_path: Path, + capsys: pytest.CaptureFixture[str], +) -> None: + g, *_ = _harness_with_storage(tmp_path) + cleanup, _inspector = _install_lifecycle_fakes( + g, + checks={"exact_outputs": False, "provenance_complete": True}, + ) + + with pytest.raises(RuntimeError, match="TP validation NOT VERIFIED"): + g["main"]() + + report = json.loads(capsys.readouterr().out) + assert report["status"] == "NOT VERIFIED" + assert report["failed_stage"] == "aggregate" + assert report["failed_error_type"] == "acceptance_checks" + assert report["failed_checks"] == ["exact_outputs"] + assert report["cleanup_attempted"] is False + assert report["artifacts_retained"] is True + assert cleanup.calls == [] + + +def test_main_aggregate_exception_emits_one_report_and_never_cleans( + tmp_path: Path, + capsys: pytest.CaptureFixture[str], +) -> None: + g, *_ = _harness_with_storage(tmp_path) + cleanup, _inspector = _install_lifecycle_fakes(g) + g["_acceptance_checks"] = lambda *_a: (_ for _ in ()).throw(ValueError("aggregate crashed")) + + with pytest.raises(RuntimeError, match="TP validation NOT VERIFIED"): + g["main"]() + + report = json.loads(capsys.readouterr().out) + assert report["status"] == "NOT VERIFIED" + assert report["failed_stage"] == "aggregate" + assert report["failed_error_type"] == "builtins.ValueError" + assert report["cleanup_attempted"] is False + assert cleanup.calls == [] + + +def test_partial_report_tolerates_malformed_remote_result(tmp_path: Path) -> None: + g, *_ = _harness_with_storage(tmp_path) + + report = g["_report_evidence"]( + "rid", + "https://modal.test/apps/vllm", + {"baseline": "malformed"}, + None, + {}, + ) + + assert report["results"] == {"baseline": "malformed"} + assert report["process_return_codes"] == {"baseline": None} + assert report["runtime_provenance"] == {"baseline": None} + + +def test_partial_report_tolerates_malformed_transaction_payload(tmp_path: Path) -> None: + g, *_ = _harness_with_storage(tmp_path) + + report = g["_report_evidence"]( + "rid", + "https://modal.test/apps/vllm", + {}, + "malformed", + {}, + ) + + assert report["published_run_id"] is None + assert report["published_path"] is None + assert report["snapshot"] == {} + + +def test_main_cleanup_transport_error_reports_conservative_retained_state( + tmp_path: Path, + capsys: pytest.CaptureFixture[str], +) -> None: + g, *_ = _harness_with_storage(tmp_path) + cleanup, inspector = _install_lifecycle_fakes( + g, + cleanup_error=RuntimeError("cleanup transport failed"), + inspection_error=RuntimeError("inspector hydration failed"), + ) + + with pytest.raises(RuntimeError, match="TP validation NOT VERIFIED"): + g["main"]() + + report = json.loads(capsys.readouterr().out) + assert report["status"] == "NOT VERIFIED" + assert report["failed_stage"] == "cleanup" + assert report["failed_error_type"] == "builtins.RuntimeError" + assert report["cleanup_attempted"] is True + assert report["cleanup_completed"] is False + assert report["artifacts_retained"] is True + assert report["marker_present"] is None + assert report["published_path_present"] is None + assert report["artifact_state_error"] == "inspector hydration failed" + assert len(cleanup.calls) == 1 + assert len(inspector.calls) == 1 + + +def test_artifact_inspection_malformed_result_is_conservatively_retained( + tmp_path: Path, +) -> None: + g, *_ = _harness_with_storage(tmp_path) + g["inspect_published_run"] = _LifecycleRemote(lambda _run_id: "malformed") + + state = g["_artifact_state_for_report"]("rid") + + assert state["marker_present"] is None + assert state["published_path_present"] is None + assert state["artifacts_retained"] is True + assert state["artifact_state_error"] == "artifact inspection returned an incomplete result" + assert state["artifact_state_error_type"] == "artifact_state_result" + + +def test_main_malformed_cleanup_result_is_not_verified_and_conservatively_retained( + tmp_path: Path, + capsys: pytest.CaptureFixture[str], +) -> None: + g, *_ = _harness_with_storage(tmp_path) + _cleanup, _inspector = _install_lifecycle_fakes(g) + malformed = _LifecycleRemote(lambda _run_id: {"cleanup_completed": "yes"}) + g["cleanup_published_run"] = malformed + + with pytest.raises(RuntimeError, match="TP validation NOT VERIFIED"): + g["main"]() + + report = json.loads(capsys.readouterr().out) + assert report["status"] == "NOT VERIFIED" + assert report["failed_stage"] == "cleanup" + assert report["cleanup_attempted"] is True + assert report["cleanup_completed"] is False + assert report["artifacts_retained"] is True + assert report["marker_present"] is None + assert report["published_path_present"] is None + assert report["cleanup_error"] == "cleanup returned an incomplete result" + assert len(malformed.calls) == 1 + + +def test_main_success_requires_cleanup_completed_report( + tmp_path: Path, + capsys: pytest.CaptureFixture[str], +) -> None: + g, *_ = _harness_with_storage(tmp_path) + cleanup, _inspector = _install_lifecycle_fakes(g) + + g["main"]() + + report = json.loads(capsys.readouterr().out) + assert report["status"] == "VERIFIED" + assert report["failed_stage"] is None + assert report["cleanup_attempted"] is True + assert report["cleanup_completed"] is True + assert report["artifacts_retained"] is False + assert "nccl_runtime_versions" in report + assert "completion_failures" in report + assert report["per_rank_offsets"]["published_snapshot"] == {} + assert len(cleanup.calls) == 1 diff --git a/tests/test_no_output_graph_load_contract.py b/tests/test_no_output_graph_load_contract.py new file mode 100644 index 00000000..8dd70aec --- /dev/null +++ b/tests/test_no_output_graph_load_contract.py @@ -0,0 +1,125 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the Foundry project +"""CPU contracts for native tests that save graphs without output tensors.""" + +from __future__ import annotations + +import ast +import os +import subprocess +import sys +import textwrap +from pathlib import Path + +import pytest + +TESTS_DIR = Path(__file__).parent +NO_OUTPUT_GRAPH_TESTS = ( + TESTS_DIR / "test_load_cublas_ws.py", + TESTS_DIR / "test_nvjet_graph.py", +) + + +def _calls_named(tree: ast.AST, name: str) -> list[ast.Call]: + return [ + node + for node in ast.walk(tree) + if isinstance(node, ast.Call) + and isinstance(node.func, ast.Attribute) + and node.func.attr == name + ] + + +@pytest.mark.parametrize("test_file", NO_OUTPUT_GRAPH_TESTS, ids=lambda path: path.stem) +def test_native_no_output_save_loads_one_graph(test_file: Path) -> None: + tree = ast.parse(test_file.read_text()) + save_calls = _calls_named(tree, "save") + load_calls = _calls_named(tree, "load") + + assert len(save_calls) == 1 + assert len(save_calls[0].args) == 1 + assert not save_calls[0].keywords + assert len(load_calls) == 1 + + assignments = [ + node + for node in ast.walk(tree) + if isinstance(node, ast.Assign) and node.value is load_calls[0] + ] + assert len(assignments) == 1 + assert len(assignments[0].targets) == 1 + assert isinstance(assignments[0].targets[0], ast.Name), ( + f"{test_file.name} saves no output tensors, so CUDAGraph.load returns one graph " + "and must not be tuple-unpacked" + ) + + +def test_public_wrapper_returns_one_graph_when_raw_outputs_are_none() -> None: + source = textwrap.dedent( + """ + import contextlib + import importlib.abc + import importlib.util + import sys + import types + + torch = types.ModuleType("torch") + + class Tensor: + pass + + class Stream: + pass + + torch.Tensor = Tensor + torch.cuda = types.SimpleNamespace( + Stream=Stream, + stream=lambda stream: contextlib.nullcontext(), + synchronize=lambda: None, + ) + sys.modules["torch"] = torch + + loaded_graph = object() + + class RawCUDAGraph: + @staticmethod + def load(json_path, pool): + assert json_path == "graph.json" + assert pool is None + return loaded_graph, None + + class OpsLoader(importlib.abc.Loader): + def exec_module(self, module): + module.CUDAGraph = RawCUDAGraph + for name in ( + "free_preallocated_region", + "get_current_alloc_offset", + "preallocate_region", + "resume_allocation_region", + "set_allocation_region", + "set_current_alloc_offset", + "stop_allocation_region", + ): + setattr(module, name, lambda *args, **kwargs: None) + + class OpsFinder(importlib.abc.MetaPathFinder): + def find_spec(self, fullname, path, target=None): + if fullname == "foundry.ops": + return importlib.util.spec_from_loader(fullname, OpsLoader()) + return None + + sys.meta_path.insert(0, OpsFinder()) + + import foundry + from foundry.graph import CUDAGraph as PythonCUDAGraph + + result = foundry.CUDAGraph.load("graph.json") + + assert foundry.CUDAGraph is PythonCUDAGraph + assert result is loaded_graph + assert not isinstance(result, tuple) + """ + ) + env = os.environ.copy() + env["PYTHONPATH"] = str(TESTS_DIR.parent / "python") + subprocess.run([sys.executable, "-c", source], check=True, env=env) diff --git a/tests/test_rng_graph.py b/tests/test_rng_graph.py new file mode 100644 index 00000000..099e9d77 --- /dev/null +++ b/tests/test_rng_graph.py @@ -0,0 +1,116 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the Foundry project +"""Fresh-process replay coverage for CUDA graphs that consume Philox state.""" + +from __future__ import annotations + +import importlib.util +import os +import shutil +import subprocess +import sys +from pathlib import Path + +import foundry as fdry +import pytest +import torch + +BASE_ADDR = 0x650000000000 +REGION_SIZE = "2GB" +ARCHIVE_DIR = "rng_graph_archive" +HOOK_ARCHIVE_DIR = "hook_archive" +GRAPH_PATH = f"{ARCHIVE_DIR}/graph.json" +ORACLE_PATH = f"{ARCHIVE_DIR}/replay_oracle.pt" + + +def _hook_path() -> str: + spec = importlib.util.find_spec("foundry.ops") + if not spec or not spec.origin: + raise RuntimeError("foundry.ops is not installed") + path = Path(spec.origin).resolve().parent / "libcuda_hook.so" + if not path.exists(): + raise RuntimeError(f"CUDA hook not found at {path}") + return str(path) + + +def _replay_twice( + graph: fdry.CUDAGraph, + output: torch.Tensor, +) -> tuple[torch.Tensor, torch.Tensor]: + graph.replay() + torch.cuda.synchronize() + first = output.detach().cpu().clone() + + graph.replay() + torch.cuda.synchronize() + second = output.detach().cpu().clone() + return first, second + + +def _save() -> None: + torch.cuda.init() + torch.cuda.set_device(0) + torch.cuda.manual_seed_all(1234) + Path(ARCHIVE_DIR).mkdir(exist_ok=True) + + fdry.set_allocation_region(BASE_ADDR, fdry.parse_size(REGION_SIZE)) + graph = fdry.CUDAGraph() + with fdry.graph(graph, pool=(0, 0)): + output = torch.rand(4096, device="cuda") + graph.save(GRAPH_PATH, output_tensors=output) + torch.save(_replay_twice(graph, output), ORACLE_PATH) + fdry.stop_allocation_region() + + +def _load() -> None: + torch.cuda.init() + torch.cuda.set_device(0) + torch.cuda.manual_seed_all(1234) + + fdry.load_cuda_modules_and_libraries(HOOK_ARCHIVE_DIR) + fdry.set_allocation_region(BASE_ADDR, fdry.parse_size(REGION_SIZE)) + graph, output = fdry.CUDAGraph.load(GRAPH_PATH, pool=(0, 0)) + expected_first, expected_second = torch.load( + ORACLE_PATH, + map_location="cpu", + weights_only=True, + ) + first, second = _replay_twice(graph, output) + + assert torch.equal(first, expected_first), "fresh LOAD replay 1 differs from SAVE" + assert torch.equal(second, expected_second), "fresh LOAD replay 2 differs from SAVE" + fdry.stop_allocation_region() + + +def _spawn(mode: str) -> None: + env = os.environ.copy() + hook = _hook_path() + env["LD_PRELOAD"] = f"{hook}:{env['LD_PRELOAD']}" if env.get("LD_PRELOAD") else hook + subprocess.check_call( + [sys.executable, str(Path(__file__).resolve()), f"--{mode}"], + env=env, + ) + + +def _cleanup() -> None: + shutil.rmtree(ARCHIVE_DIR, ignore_errors=True) + shutil.rmtree(HOOK_ARCHIVE_DIR, ignore_errors=True) + + +@pytest.mark.filterwarnings("ignore:TORCH_CUDA_ARCH_LIST is not set") +def test_rng_graph_replays_in_fresh_process() -> None: + _cleanup() + try: + _spawn("save") + _spawn("load") + finally: + _cleanup() + + +if __name__ == "__main__": + if "--save" in sys.argv: + _save() + elif "--load" in sys.argv: + _load() + else: + test_rng_graph_replays_in_fresh_process() diff --git a/tests/test_sglang_capture_hooks.py b/tests/test_sglang_capture_hooks.py new file mode 100644 index 00000000..991b4a24 --- /dev/null +++ b/tests/test_sglang_capture_hooks.py @@ -0,0 +1,656 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the Foundry project +"""CPU-only contracts for the current-main SGLang capture SAVE/LOAD hooks. + +Foundry binds graph SAVE/LOAD to the pinned upstream (a25164b) pluggable-backend +architecture, behaviour-faithful to the proven PR4 design: + + * ``FullCudaGraphBackend`` is REPLACE'd with a Foundry subclass whose + ``capture_one`` captures into a Foundry graph on SAVE and, on LOAD, finishes + one archived graph per requested shape (no forward / warmup / native + capture). Its ``capture_session`` override drives the LOAD prepare + (preallocate, scan, one ``start_graph_builds``) and fail-closed exit. + * ``DecodeCudaGraphRunner.capture`` is wrapped AROUND. SAVE delegates to the + original then finalizes the archive; LOAD ALSO delegates to the original so + the pinned upstream lifecycle (warmup, ``graph_capture`` context, seq fill, + ``capture_prepare`` / ``ForwardContext`` / real attention metadata, per-shape + dispatch) runs — the backend materializes each archived graph. The hook does + no manual LOAD orchestration. + +The backend LOAD lifecycle itself is covered in ``test_sglang_graph_load.py``. +These contracts exercise the hook bodies + backend SAVE/NONE with SGLang fully +mocked (the local VM has no CUDA). +""" + +from __future__ import annotations + +import importlib.util +import sys +from dataclasses import dataclass +from enum import Enum +from pathlib import Path +from types import ModuleType, SimpleNamespace + +import pytest + +REPO_ROOT = Path(__file__).parents[1] +SGLANG_DIR = REPO_ROOT / "python" / "foundry" / "integration" / "sglang" +HOOKS_PATH = SGLANG_DIR / "hooks.py" +BACKEND_PATH = SGLANG_DIR / "backend.py" + +CAPTURE_EXPECTED = { + ( + "sglang.srt.model_executor.runner_backend.full_cuda_graph_backend.FullCudaGraphBackend", + "REPLACE", + ), + ( + "sglang.srt.model_executor.runner.decode_cuda_graph_runner.DecodeCudaGraphRunner.capture", + "AROUND", + ), +} + + +@dataclass(frozen=True) +class ShapeKey: + size: int + stream_idx: int | None = None + variant_label: str | None = None + + +class ExtensionMode(Enum): + NONE = "none" + SAVE = "save" + LOAD = "load" + + +class HookType(Enum): + BEFORE = "before" + AFTER = "after" + AROUND = "around" + REPLACE = "replace" + + +def _install_module(monkeypatch, name: str) -> ModuleType: + module = ModuleType(name) + monkeypatch.setitem(sys.modules, name, module) + if "." in name: + parent_name, child_name = name.rsplit(".", 1) + parent = sys.modules.get(parent_name) + if parent is not None: + setattr(parent, child_name, module) + return module + + +def _load_module(path: Path, name: str): + spec = importlib.util.spec_from_file_location(name, path) + assert spec is not None and spec.loader is not None + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + return module + + +class _Recorder: + """Records graph_ops / runtime calls made by the hook bodies.""" + + def __init__(self): + self.events = [] + + def record(self, name): + def _fn(*args, **kwargs): + self.events.append((name, args, kwargs)) + + return _fn + + +def _install_common_sglang(monkeypatch): + _install_module(monkeypatch, "sglang").__path__ = [] + _install_module(monkeypatch, "sglang.srt").__path__ = [] + _install_module(monkeypatch, "sglang.srt.model_executor").__path__ = [] + _install_module(monkeypatch, "sglang.srt.model_executor.runner").__path__ = [] + shape_key = _install_module(monkeypatch, "sglang.srt.model_executor.runner.shape_key") + shape_key.ShapeKey = ShapeKey + + +def _build_config_module(monkeypatch, mode): + config = _install_module(monkeypatch, "foundry.integration.sglang.config") + config.CUDAGraphExtensionMode = ExtensionMode + current = [mode] + config.get_graph_extension_mode = lambda: current[0] + config.get_workspace_root = lambda: None + config.get_config = lambda: SimpleNamespace(workspace_dir="/ws") + config._current = current + return config + + +def _load_hooks(monkeypatch, mode, recorder): + monkeypatch.syspath_prepend(str(REPO_ROOT / "python")) + foundry = _install_module(monkeypatch, "foundry") + foundry.__path__ = [] + _install_module(monkeypatch, "foundry.integration").__path__ = [] + _install_module(monkeypatch, "foundry.integration.sglang").__path__ = [] + + _build_config_module(monkeypatch, mode) + + runtime = _install_module(monkeypatch, "foundry.integration.sglang.runtime") + for name in ( + "setup_ld_preload_env", + "setup_graph_extension", + "log_alloc_offset", + "skip_to_scratch_boundary", + "capture_final_alloc_offset", + "preallocate_for_load_mode", + ): + setattr(runtime, name, recorder.record(name)) + + graph_ops = _install_module(monkeypatch, "foundry.integration.sglang.graph_ops") + for name in ("save_graph_manifest", "pack_fatbins"): + setattr(graph_ops, name, recorder.record(name)) + + backend_mod = _install_module(monkeypatch, "foundry.integration.sglang.backend") + backend_mod.build_foundry_full_backend_class = lambda: FoundryBackendSentinel + + _install_common_sglang(monkeypatch) + + hooks = _load_module(HOOKS_PATH, "sglang_capture_hooks_contract") + return hooks, graph_ops, runtime + + +class FoundryBackendSentinel: + """Stand-in for the REPLACE class produced by backend.build_*.""" + + +class RegistrySpy: + def __init__(self): + self.calls = [] + + def register(self, target, hook, hook_type=HookType.AFTER): + self.calls.append((target, hook, hook_type)) + + +@pytest.fixture +def hooks_module_active(monkeypatch): + recorder = _Recorder() + hooks, _graph_ops, _runtime = _load_hooks(monkeypatch, ExtensionMode.SAVE, recorder) + calls = [] + + class TrackingRegistry: + def register(self, target, hook, hook_type=HookType.AFTER): + calls.append((target, hook_type.name)) + + registry = TrackingRegistry() + + def _register_hooks(): + hooks._register_capture_hooks(registry, HookType) + + hooks.register_hooks = _register_hooks + hooks._REGISTRY_CALLS = calls + return hooks + + +def test_capture_hooks_registered(hooks_module_active): + hooks_module_active.register_hooks() + assert set(hooks_module_active._REGISTRY_CALLS) >= CAPTURE_EXPECTED + + +def test_register_capture_hooks_uses_replace_and_around(monkeypatch): + recorder = _Recorder() + hooks, _graph_ops, _runtime = _load_hooks(monkeypatch, ExtensionMode.SAVE, recorder) + registry = RegistrySpy() + hooks._register_capture_hooks(registry, HookType) + + by_target = {target: (hook, ht) for target, hook, ht in registry.calls} + backend_target = ( + "sglang.srt.model_executor.runner_backend.full_cuda_graph_backend.FullCudaGraphBackend" + ) + capture_target = ( + "sglang.srt.model_executor.runner.decode_cuda_graph_runner.DecodeCudaGraphRunner.capture" + ) + assert by_target[backend_target][0] is FoundryBackendSentinel + assert by_target[backend_target][1] is HookType.REPLACE + assert by_target[capture_target][1] is HookType.AROUND + + +def _get_capture_around(hooks, registry): + capture_target = ( + "sglang.srt.model_executor.runner.decode_cuda_graph_runner.DecodeCudaGraphRunner.capture" + ) + return next(hook for target, hook, _ht in registry.calls if target == capture_target) + + +def test_capture_around_none_mode_is_passthrough(monkeypatch): + recorder = _Recorder() + hooks, _graph_ops, _runtime = _load_hooks(monkeypatch, ExtensionMode.NONE, recorder) + registry = RegistrySpy() + hooks._register_capture_hooks(registry, HookType) + around = _get_capture_around(hooks, registry) + + calls = [] + + def original(self, *args, **kwargs): + calls.append((self, args, kwargs)) + return "orig" + + runner = SimpleNamespace() + assert around(original, runner) == "orig" + assert calls == [(runner, (), {})] + assert recorder.events == [] + + +def test_capture_around_save_delegates_then_finalizes(monkeypatch): + recorder = _Recorder() + hooks, _graph_ops, _runtime = _load_hooks(monkeypatch, ExtensionMode.SAVE, recorder) + registry = RegistrySpy() + hooks._register_capture_hooks(registry, HookType) + around = _get_capture_around(hooks, registry) + + order = [] + + def original(self, *args, **kwargs): + order.append("original") + return "captured" + + runner = SimpleNamespace(capture_bs=[256]) + result = around(original, runner) + + assert result == "captured" + event_names = [name for name, _a, _k in recorder.events] + # SAVE must run the upstream capture first, then finalize the archive. + assert order == ["original"] + assert event_names == ["save_graph_manifest", "pack_fatbins", "capture_final_alloc_offset"] + + +def test_capture_around_load_calls_original_lifecycle(monkeypatch): + # The strongest PR4 divergence fix: LOAD delegates to the pinned upstream + # capture() so graph_capture context, seq fills, capture_prepare, + # ForwardContext, real attention metadata and per-shape dispatch all run. + # The backend materializes each archived graph; the hook does NO manual + # LOAD orchestration. + recorder = _Recorder() + hooks, graph_ops, _runtime = _load_hooks(monkeypatch, ExtensionMode.LOAD, recorder) + registry = RegistrySpy() + hooks._register_capture_hooks(registry, HookType) + around = _get_capture_around(hooks, registry) + + calls = [] + + def original(self, *args, **kwargs): + calls.append((self, args, kwargs)) + return "loaded-by-upstream" + + runner = SimpleNamespace(capture_bs=[256], backend=SimpleNamespace(_pool=None)) + result = around(original, runner) + + assert result == "loaded-by-upstream" + assert calls == [(runner, (), {})] + # No manual pool publish, preallocate, seq fill, or load_all_graphs. + assert runner.backend._pool is None + assert recorder.events == [] + assert not hasattr(graph_ops, "load_all_graphs") + assert not hasattr(graph_ops, "get_or_create_graph_pool") + + +@pytest.mark.parametrize("capture_bs", [[128], [256, 128], [], [512], [256, 256]]) +def test_capture_around_save_rejects_capture_bs_drift(monkeypatch, capture_bs): + # The profile validates cuda_graph_config.decode.bs==[256], but the runtime + # runner rebuilds capture_bs (request-pool/alignment clamping can shrink or + # replace it). SAVE must fail closed before any capture if it is not [256]. + recorder = _Recorder() + hooks, _graph_ops, _runtime = _load_hooks(monkeypatch, ExtensionMode.SAVE, recorder) + registry = RegistrySpy() + hooks._register_capture_hooks(registry, HookType) + around = _get_capture_around(hooks, registry) + + calls = [] + + def original(self, *args, **kwargs): + calls.append("original") + return "captured" + + runner = SimpleNamespace(capture_bs=capture_bs) + with pytest.raises(RuntimeError, match="capture_bs"): + around(original, runner) + + # Fail closed BEFORE any capture or archive finalization. + assert calls == [] + assert recorder.events == [] + + +@pytest.mark.parametrize("capture_bs", [[128], [256, 128], [], [512], [256, 256]]) +def test_capture_around_load_rejects_capture_bs_drift(monkeypatch, capture_bs): + recorder = _Recorder() + hooks, _graph_ops, _runtime = _load_hooks(monkeypatch, ExtensionMode.LOAD, recorder) + registry = RegistrySpy() + hooks._register_capture_hooks(registry, HookType) + around = _get_capture_around(hooks, registry) + + calls = [] + + def original(self, *args, **kwargs): + calls.append("original") + return None + + runner = SimpleNamespace(capture_bs=capture_bs, backend=SimpleNamespace(_pool=None)) + with pytest.raises(RuntimeError, match="capture_bs"): + around(original, runner) + + # Fail closed BEFORE delegating to the upstream capture lifecycle. + assert calls == [] + assert recorder.events == [] + assert runner.backend._pool is None + + +# --- backend subclass contracts ------------------------------------------- + + +def _load_backend(monkeypatch, mode): + foundry = _install_module(monkeypatch, "foundry") + foundry.__path__ = [] + _install_module(monkeypatch, "foundry.integration").__path__ = [] + _install_module(monkeypatch, "foundry.integration.sglang").__path__ = [] + _build_config_module(monkeypatch, mode) + + captured = {"events": []} + + runtime = _install_module(monkeypatch, "foundry.integration.sglang.runtime") + runtime.preallocate_for_load_mode = lambda: captured["events"].append("preallocate") + runtime.log_alloc_offset = lambda label: captured["events"].append(("offset", label)) + + graph_module = _install_module(monkeypatch, "foundry.graph") + + class FoundryCUDAGraph: + instances = [] + pending_sentinel = object() + + def __init__(self): + FoundryCUDAGraph.instances.append(self) + + @staticmethod + def start_graph_builds(paths, pool=None, num_threads=4): + captured["events"].append(("start_graph_builds", tuple(paths), pool)) + return FoundryCUDAGraph.pending_sentinel + + @staticmethod + def finish_one_graph_load(pending, index): + captured["events"].append(("finish_one", index)) + return (f"graph-{index}", ("tensors", index)) + + FoundryCUDAGraph.instances = [] + + class _GraphCtx: + def __init__(self, cuda_graph, pool=None, stream=None): + captured["ctx"] = {"graph": cuda_graph, "pool": pool, "stream": stream} + + def __enter__(self): + captured["entered"] = True + captured["events"].append("enter") + return self + + def __exit__(self, *args): + captured["exited"] = True + captured["events"].append("exit") + return False + + graph_module.CUDAGraph = FoundryCUDAGraph + graph_module.graph = _GraphCtx + + graph_ops = _install_module(monkeypatch, "foundry.integration.sglang.graph_ops") + saved = [] + graph_ops.save_graph = lambda graph, output, key: saved.append((graph, output, key)) + graph_ops._scan_graph_files = lambda workspace_dir: [] + graph_ops._unpack_output = lambda tensors: ("unpacked", tensors) + graph_ops._require_dense_shape_key = lambda key: int(key.size) + + # Base FullCudaGraphBackend + pool/pynccl the subclass capture_session uses. + _install_module(monkeypatch, "sglang").__path__ = [] + _install_module(monkeypatch, "sglang.srt").__path__ = [] + _install_module(monkeypatch, "sglang.srt.distributed").__path__ = [] + _install_module(monkeypatch, "sglang.srt.distributed.device_communicators").__path__ = [] + pynccl = _install_module( + monkeypatch, "sglang.srt.distributed.device_communicators.pynccl_allocator" + ) + pynccl.set_graph_pool_id = lambda pool: captured["events"].append(("set_graph_pool_id", pool)) + + _install_module(monkeypatch, "sglang.srt.model_executor").__path__ = [] + _install_module(monkeypatch, "sglang.srt.model_executor.runner_backend").__path__ = [] + _install_module(monkeypatch, "sglang.srt.model_executor.runner_utils").__path__ = [] + pool_mod = _install_module(monkeypatch, "sglang.srt.model_executor.runner_utils.pool") + pool_mod.get_or_create_global_graph_memory_pool = lambda device_module: "pool-handle" + + fcgb = _install_module( + monkeypatch, + "sglang.srt.model_executor.runner_backend.full_cuda_graph_backend", + ) + + class FullCudaGraphBackend: + def __init__(self, cuda_graph_runner, *, enable_memory_saver=False): + self.cuda_graph_runner = cuda_graph_runner + self.enable_memory_saver = enable_memory_saver + self._graphs = {} + self._outputs = {} + self._pool = None + self._capture_stream = None + self._device_module = SimpleNamespace( + synchronize=lambda: captured["events"].append("synchronize") + ) + self._tp_group = SimpleNamespace(barrier=lambda: captured["events"].append("barrier")) + self.super_calls = [] + + def capture_one(self, shape_key, forward_fn, capture_inputs=None, post_warmup_hook=None): + self.super_calls.append((shape_key, capture_inputs, post_warmup_hook)) + return "super-capture" + + def cleanup(self): + captured["events"].append("super_cleanup") + self._graphs.clear() + self._outputs.clear() + self._pool = None + + fcgb.FullCudaGraphBackend = FullCudaGraphBackend + + backend = _load_module(BACKEND_PATH, "sglang_capture_backend_contract") + return backend, FullCudaGraphBackend, FoundryCUDAGraph, saved, captured + + +def test_backend_subclass_init_sets_load_state(monkeypatch): + backend_mod, base, _foundry_graph, _saved, _captured = _load_backend( + monkeypatch, ExtensionMode.LOAD + ) + cls = backend_mod.build_foundry_full_backend_class() + assert issubclass(cls, base) + inst = cls(SimpleNamespace()) + # LOAD bookkeeping initialized on top of the inherited shared-pool state. + assert inst._pending is None + assert inst._graph_files == [] + assert inst._load_index == 0 + assert inst._closed is False + assert inst._pool is None + + +def test_backend_subclass_delegates_when_not_save(monkeypatch): + backend_mod, base, foundry_graph, saved, _captured = _load_backend( + monkeypatch, ExtensionMode.NONE + ) + cls = backend_mod.build_foundry_full_backend_class() + assert issubclass(cls, base) + inst = cls(SimpleNamespace()) + + result = inst.capture_one(ShapeKey(size=4), lambda: "out", post_warmup_hook=lambda: None) + assert result == "super-capture" + assert len(inst.super_calls) == 1 + assert saved == [] + assert foundry_graph.instances == [] + + +def test_backend_subclass_saves_on_save_without_warmup(monkeypatch): + backend_mod, _base, foundry_graph, saved, captured = _load_backend( + monkeypatch, ExtensionMode.SAVE + ) + cls = backend_mod.build_foundry_full_backend_class() + inst = cls(SimpleNamespace()) + + forwards = [] + hook_calls = [] + + def forward_fn(): + forwards.append("forward") + return "decode-output" + + key = ShapeKey(size=8) + # SAVE runs through the backend capture_session (which establishes the pool + # + stream) then the upstream loop calls capture_one. + with inst.capture_session("stream-handle"): + result = inst.capture_one( + key, forward_fn, post_warmup_hook=lambda: hook_calls.append("hook") + ) + + assert result is None + # Exactly one forward — the captured one; the two warmup forwards are skipped. + assert forwards == ["forward"] + # post_warmup_hook must NOT be invoked (no warmup ran -> no mutated state). + assert hook_calls == [] + # Captured into a fresh Foundry graph, using the backend's pool + stream. + assert len(foundry_graph.instances) == 1 + assert captured["ctx"]["graph"] is foundry_graph.instances[0] + assert captured["ctx"]["pool"] == "pool-handle" + assert captured["ctx"]["stream"] == "stream-handle" + assert captured["entered"] and captured["exited"] + # backend maps populated and archive written. + assert inst._graphs[key] is foundry_graph.instances[0] + assert inst._outputs[key] == "decode-output" + assert saved == [(foundry_graph.instances[0], "decode-output", key)] + + +def test_backend_subclass_synchronizes_and_barriers_before_capture(monkeypatch): + backend_mod, _base, _foundry_graph, _saved, captured = _load_backend( + monkeypatch, ExtensionMode.SAVE + ) + cls = backend_mod.build_foundry_full_backend_class() + inst = cls(SimpleNamespace()) + + forwards = [] + post_warmup_calls = [] + + def forward_fn(): + forwards.append("forward") + captured["events"].append("forward") + return "decode-output" + + with inst.capture_session("stream-handle"): + inst.capture_one( + ShapeKey(size=256), + forward_fn, + post_warmup_hook=lambda: post_warmup_calls.append("hook"), + ) + + # Device synchronize + TP barrier must lockstep every rank immediately + # before entering the Foundry graph context; exactly one captured forward + # runs inside it, and no warmup/post-warmup work happens. capture_session + # publishes the pool first (no synchronize/barrier of its own). + assert captured["events"][-5:] == ["synchronize", "barrier", "enter", "forward", "exit"] + assert forwards == ["forward"] + assert post_warmup_calls == [] + + +def test_backend_cleanup_none_mode_delegates_to_super_and_allows_new_session(monkeypatch): + backend_mod, _base, _foundry_graph, _saved, captured = _load_backend( + monkeypatch, ExtensionMode.NONE + ) + cls = backend_mod.build_foundry_full_backend_class() + inst = cls(SimpleNamespace()) + inst._graphs[ShapeKey(size=256)] = "graph" + inst._outputs[ShapeKey(size=256)] = "out" + + inst.cleanup() + + # In stock (NONE) operation the subclass is transparent: cleanup delegates + # to super and does NOT latch the fail-closed recapture guard. + assert "super_cleanup" in captured["events"] + assert inst._graphs == {} + assert inst._outputs == {} + assert inst._closed is False + + # A subsequent stock capture_session is therefore allowed (e.g. the + # hidden-mode recapture the upstream runner performs). + with inst.capture_session("stream-handle"): + pass + + +def test_backend_cleanup_active_mode_latches_recapture(monkeypatch): + backend_mod, _base, _foundry_graph, _saved, captured = _load_backend( + monkeypatch, ExtensionMode.SAVE + ) + cls = backend_mod.build_foundry_full_backend_class() + inst = cls(SimpleNamespace()) + inst._graphs[ShapeKey(size=256)] = "graph" + + inst.cleanup() + + # Active SAVE/LOAD retains the fail-closed recapture latch (Foundry graphs + # cannot be safely recaptured into a torn-down pool) and does its own + # extended cleanup rather than the stock one. + assert inst._closed is True + assert "super_cleanup" not in captured["events"] + assert inst._graphs == {} + with pytest.raises(RuntimeError, match="recapture"), inst.capture_session("stream-handle"): + pass + + +# --- output narrowing contracts (D9) -------------------------------------- + + +def _load_graph_ops_for_output(monkeypatch): + torch = _install_module(monkeypatch, "torch") + + class _Tensor: + pass + + torch.Tensor = _Tensor + + foundry = _install_module(monkeypatch, "foundry") + foundry.__path__ = [] + _install_module(monkeypatch, "foundry.ops") + graph = _install_module(monkeypatch, "foundry.graph") + graph.CUDAGraph = object + graph.graph = lambda *a, **k: None + _install_module(monkeypatch, "foundry.integration").__path__ = [] + _install_module(monkeypatch, "foundry.integration.sglang").__path__ = [] + config = _install_module(monkeypatch, "foundry.integration.sglang.config") + config.get_config = lambda: None + runtime = _install_module(monkeypatch, "foundry.integration.sglang.runtime") + runtime.get_state = lambda: None + + _install_module(monkeypatch, "sglang").__path__ = [] + _install_module(monkeypatch, "sglang.srt").__path__ = [] + _install_module(monkeypatch, "sglang.srt.layers").__path__ = [] + lp = _install_module(monkeypatch, "sglang.srt.layers.logits_processor") + + @dataclass + class LogitsProcessorOutput: + next_token_logits: object = None + hidden_states: object = None + full_logits: object = None + customized_info: object = None + + lp.LogitsProcessorOutput = LogitsProcessorOutput + + graph_ops_path = SGLANG_DIR / "graph_ops.py" + return ( + _load_module(graph_ops_path, "sglang_graph_ops_output_contract"), + _Tensor, + LogitsProcessorOutput, + ) + + +def test_pack_output_accepts_dense_next_token_logits(monkeypatch): + graph_ops, tensor_cls, lp = _load_graph_ops_for_output(monkeypatch) + logits = tensor_cls() + out = lp(next_token_logits=logits) + assert graph_ops._pack_output(out) is logits + + +@pytest.mark.parametrize("field", ["hidden_states", "full_logits", "customized_info"]) +def test_pack_output_rejects_unsupported_fields(monkeypatch, field): + graph_ops, tensor_cls, lp = _load_graph_ops_for_output(monkeypatch) + out = lp(next_token_logits=tensor_cls(), **{field: object()}) + with pytest.raises(ValueError): + graph_ops._pack_output(out) diff --git a/tests/test_sglang_documentation.py b/tests/test_sglang_documentation.py new file mode 100644 index 00000000..d4ab3d7a --- /dev/null +++ b/tests/test_sglang_documentation.py @@ -0,0 +1,234 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the Foundry project +"""Contracts for the pinned SGLang symmetric-TP evidence documentation.""" + +import re +from pathlib import Path + +REPO_ROOT = Path(__file__).resolve().parents[1] +EVIDENCE_DOCS = ( + REPO_ROOT / "README.md", + REPO_ROOT / "ROADMAP.md", + REPO_ROOT / "RELEASE.md", + REPO_ROOT / "docs" / "sglang" / "direct-edits.md", + REPO_ROOT / "docs" / "sglang" / "hooks.md", + REPO_ROOT / "docs" / "sglang" / "memory-consistency.md", + REPO_ROOT / "docs" / "sglang" / "memory-lifecycle.md", + REPO_ROOT / "docs" / "sglang" / "overview.md", + REPO_ROOT / "docs" / "sglang" / "save-load-workflow.md", + REPO_ROOT / "recipe" / "experimental" / "README.md", + REPO_ROOT / "recipe" / "sglang" / "README.md", +) +REQUIRED_EVIDENCE = ( + "a8e43442421a80fea7991ee26f97f1003b170e29", + "a25164bda355250015c0ba6d8821f0d4b37abe7d", + "b968826d9c46dd6066d109eabc6255188de91218", + "ap-2qSQUKlY24L4RrQkMpUrmc", + "ap-5wO9cFhRZ7ITsLDDv27Xfh", + "ap-6fYcSUFLOxocBSjNxakNGZ", + "68908220416", +) +STALE_FINAL_EVIDENCE = ( + "435518206f501b2f6b20080ee53375d0d6a3ff76", + "733f76678a73171a1534a9d61baabc8ac1fc20e3", + "ap-SwjRFbh7zqLfmKlGGkDuTM", + "ap-RHFcD96CAsGSqO7Z236yzF", +) +STALE_CLAIMS = ( + "acceptance pending", + "acceptance is pending", + "acceptance run is still pending", + "acceptance **pending", + "pinned-main head is not yet validated", + "pinned-main architecture has not yet", + "targets nccl p2p/ipc all-reduce", + "selects plain nccl", + "only supported profile. custom / symmetric-memory", + "load short-circuits that capture", + "short-circuit the upstream loop", + "exactly **seven**", + "seven `hookregistry` targets", +) +SUPERSEDED_SGLANG_DOCS = { + REPO_ROOT / "docs" / "superpowers" / "plans" / "2026-07-23-sglang-main-port.md": False, + REPO_ROOT / "docs" / "superpowers" / "plans" / "2026-07-23-tp-upstream-consolidation.md": True, + REPO_ROOT / "docs" / "superpowers" / "specs" / "2026-07-23-sglang-main-port-design.md": False, + REPO_ROOT + / "docs" + / "superpowers" + / "specs" + / "2026-07-23-tp-upstream-consolidation-design.md": True, +} +HISTORICAL_EVIDENCE_COMMITS = { + "69736595456": ( + "https://github.com/modal-projects/foundry/commit/21bf7341847ebce46c534e7f6e86a3929275c410" + ), + "121831948288": ( + "https://github.com/modal-projects/foundry/commit/ee7d61a32a4d4732e5f14978975b96ab03bcbede" + ), +} +DEFINITIVE_NCCL_CAUSE_PHRASES = ( + "because save-process communicator", + "did not reconstruct save-process communicator", + "does not reconstruct archived save-process", + "do not reconstruct archived save-process", + "foundry does not reconstruct that state", +) +CURRENT_ACTIVATION_TARGET = "foundry_sglang_plugin:activate" +SUPERSEDED_PORT_DOCS = ( + REPO_ROOT / "docs" / "superpowers" / "plans" / "2026-07-23-sglang-main-port.md", + REPO_ROOT / "docs" / "superpowers" / "specs" / "2026-07-23-sglang-main-port-design.md", +) +TREATMENT_DOCS = ( + REPO_ROOT / "docs" / "superpowers" / "plans" / "2026-07-23-sglang-symmetric-tp-treatment.md", + REPO_ROOT + / "docs" + / "superpowers" + / "specs" + / "2026-07-23-sglang-symmetric-tp-treatment-design.md", +) + + +def _normalized(path: Path) -> str: + return " ".join(path.read_text().lower().split()) + + +def test_sglang_evidence_docs_record_exact_pins_runs_and_offset() -> None: + assert len(EVIDENCE_DOCS) == 11 + for path in EVIDENCE_DOCS: + text = path.read_text() + missing = [item for item in REQUIRED_EVIDENCE if item not in text] + assert not missing, f"{path.relative_to(REPO_ROOT)} missing {missing}" + stale = [item for item in STALE_FINAL_EVIDENCE if item in text] + assert not stale, f"{path.relative_to(REPO_ROOT)} contains stale {stale}" + + +def test_sglang_evidence_docs_reject_stale_plain_nccl_and_pending_claims() -> None: + for path in EVIDENCE_DOCS: + normalized = _normalized(path) + stale = [claim for claim in STALE_CLAIMS if claim in normalized] + assert not stale, f"{path.relative_to(REPO_ROOT)} contains stale claims {stale}" + + +def test_pins_are_modal_provenance_not_runtime_profile_validation() -> None: + hooks = _normalized(REPO_ROOT / "docs" / "sglang" / "hooks.md") + hooks_plain = hooks.replace("`", "") + overview = _normalized(REPO_ROOT / "docs" / "sglang" / "overview.md") + assert "_validate_profile does not inspect the sglang or model sha" in hooks_plain + assert "topology, backend, and settings" in hooks + assert "validation provenance enforced by the modal harness" in overview + + for relative_path in ( + "recipe/experimental/README.md", + "recipe/sglang/README.md", + ): + recipe = _normalized(REPO_ROOT / relative_path) + assert "shell recipe users must install the matching sglang pin" in recipe + assert "arbitrary installed sglang revisions are not rejected at startup" in recipe + + harness = (REPO_ROOT / "tests" / "modal_sglang_tp.py").read_text() + assert 'f"git -C /sglang checkout --detach {SGLANG_BASE}"' in harness + assert 'f\'test "$(git -C /sglang rev-parse HEAD)" = "{SGLANG_BASE}"\'' in harness + assert 'env["SGL_MODEL_REVISION"] = MODEL_REVISION' in harness + + +def test_overview_load_runs_real_attention_metadata_planning() -> None: + overview = _normalized(REPO_ROOT / "docs" / "sglang" / "overview.md") + assert "skipping per-batch-size attention-metadata" not in overview + assert "runs the real upstream lifecycle" in overview + assert "`init_forward_metadata_out_graph` planning" in overview + assert "replaces only native graph capture with archive materialization" in overview + + +def test_plain_nccl_rejection_does_not_claim_treatment_proved_the_cause() -> None: + for path in EVIDENCE_DOCS: + normalized = _normalized(path) + assert "leading forensic explanation" in normalized + definitive = [phrase for phrase in DEFINITIVE_NCCL_CAUSE_PHRASES if phrase in normalized] + assert not definitive, ( + f"{path.relative_to(REPO_ROOT)} states unproven NCCL causality: {definitive}" + ) + + +def test_final_evidence_values_and_pinned_scope_are_documented() -> None: + overview = _normalized(REPO_ROOT / "docs" / "sglang" / "overview.md") + for evidence in ( + "one graph per rank", + "offset `68908220416`", + "73 symmetric all-reduce", + "one symmetric all-gather", + "zero nccl collective nodes", + "217.2 s baseline", + "103.0 s save", + "103.1 s save2", + "61.0 s load", + "61.0 s load2", + "ap-5wo9cfhrz7itslddv27xfh", + "ap-6fycsufloxocbsjnxakngz", + ): + assert evidence in overview + + for path in EVIDENCE_DOCS: + normalized = _normalized(path) + for scope_token in ("pinned", "single-host", "qwen3-8b", "tp=2"): + assert scope_token in normalized, ( + f"{path.relative_to(REPO_ROOT)} missing pinned scope token {scope_token}" + ) + assert "general/upstream tp" in normalized + assert "plain-nccl support" in normalized or "plain nccl" in normalized + if "modal-projects/foundry/pull/2" in normalized or "pull/4" in normalized: + assert "pre-port" in normalized + + +def test_older_plain_nccl_designs_are_prominently_superseded_for_sglang() -> None: + superpowers_root = REPO_ROOT / "docs" / "superpowers" + policy_patterns = ( + r"both engines use.{0,240}nccl p2p", + r"initial support is.{0,240}plain nccl", + r"only supported envelope.{0,240}plain-nccl", + r"supported envelope:.{0,240}plain nccl", + r"plain nccl p2p/legacy cuda ipc collectives", + ) + detected = { + path + for path in superpowers_root.glob("**/*.md") + if "sglang" in _normalized(path) + and any(re.search(pattern, _normalized(path)) for pattern in policy_patterns) + } + assert detected == set(SUPERSEDED_SGLANG_DOCS) + + for path, preserves_vllm in SUPERSEDED_SGLANG_DOCS.items(): + banner = _normalized(path)[:1200] + assert "superseded for sglang" in banner + assert "2026-07-23-sglang-symmetric-tp-treatment-design.md" in banner + assert "2026-07-23-sglang-symmetric-tp-treatment.md" in banner + if preserves_vllm: + assert "vllm portions remain applicable" in banner + + +def test_historical_numeric_evidence_links_to_immutable_commits() -> None: + for path in (*EVIDENCE_DOCS, *SUPERSEDED_SGLANG_DOCS): + text = path.read_text() + for value, commit_url in HISTORICAL_EVIDENCE_COMMITS.items(): + if value in text: + assert commit_url in text, ( + f"{path.relative_to(REPO_ROOT)} attributes {value} without {commit_url}" + ) + + +def test_superseded_port_docs_use_current_activation_target_and_relative_evidence_paths() -> None: + for path in SUPERSEDED_PORT_DOCS: + text = path.read_text() + assert CURRENT_ACTIVATION_TARGET in text + assert "foundry.integration.sglang.plugin:activate" not in text + assert "/workspace" not in text + + +def test_treatment_docs_describe_modal_auto_commit_failure_retention() -> None: + for path in TREATMENT_DOCS: + normalized = _normalized(path) + assert "modal volume auto-commit" in normalized + assert "partial or mutated unique-run artifacts" in normalized + assert "stop later phases" in normalized + assert "never verified" in normalized + assert "private" not in normalized diff --git a/tests/test_sglang_graph_load.py b/tests/test_sglang_graph_load.py new file mode 100644 index 00000000..598ea9a0 --- /dev/null +++ b/tests/test_sglang_graph_load.py @@ -0,0 +1,350 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the Foundry project +"""CPU-only contract for the pinned-upstream-faithful SGLang graph LOAD. + +The proven PR4 backend does NOT short-circuit ``DecodeCudaGraphRunner.capture`` +on LOAD. Upstream ``capture()`` runs its real lifecycle (warmup, ``graph_capture`` +context, seq-len fill, ``capture_prepare`` / ``ForwardContext`` / attention +metadata, then per-shape dispatch), and the pluggable backend materializes each +archived graph: + + * ``capture_session`` (LOAD): establish the inherited shared pool + stream, + preallocate, scan every archived graph path, issue exactly ONE + ``start_graph_builds(paths, pool=self._pool)``, then yield to the upstream + per-shape loop. On exit it requires every archived file consumed and logs + the final offset/count. + * ``capture_one`` (LOAD): validate the requested ShapeKey against the archive + order, ``finish_one_graph_load(pending, index)``, and populate + ``self._graphs`` / ``self._outputs`` — no forward, no native capture, no + warmup. + +These contracts drive the backend module through a fake upstream ``capture()`` +loop with SGLang fully mocked (the local VM has no CUDA). +""" + +from __future__ import annotations + +import importlib.util +import sys +from dataclasses import dataclass +from pathlib import Path +from types import ModuleType, SimpleNamespace + +import pytest + +REPO_ROOT = Path(__file__).parents[1] +BACKEND_PATH = REPO_ROOT / "python" / "foundry" / "integration" / "sglang" / "backend.py" + + +@dataclass(frozen=True) +class ShapeKey: + size: int + stream_idx: int | None = None + variant_label: str | None = None + + +class ExtensionMode: + class _M: + def __init__(self, value): + self.value = value + + NONE = _M("none") + SAVE = _M("save") + LOAD = _M("load") + + +def _install_module(monkeypatch, name: str) -> ModuleType: + module = ModuleType(name) + monkeypatch.setitem(sys.modules, name, module) + if "." in name: + parent_name, child_name = name.rsplit(".", 1) + parent = sys.modules.get(parent_name) + if parent is not None: + setattr(parent, child_name, module) + return module + + +def _load_backend(monkeypatch, mode, *, graph_files): + """Install the fake SGLang/Foundry surface and load ``backend.py``. + + ``graph_files`` is the archive ``_scan_graph_files`` returns: a list of + ``(index, filename, {"shape_key": ShapeKey(...)})`` in capture (descending-bs) + order. + """ + events: list = [] + + foundry = _install_module(monkeypatch, "foundry") + foundry.__path__ = [] + _install_module(monkeypatch, "foundry.integration").__path__ = [] + _install_module(monkeypatch, "foundry.integration.sglang").__path__ = [] + + config = _install_module(monkeypatch, "foundry.integration.sglang.config") + config.CUDAGraphExtensionMode = ExtensionMode + config.get_graph_extension_mode = lambda: mode + config.get_config = lambda: SimpleNamespace(workspace_dir="/ws") + + runtime = _install_module(monkeypatch, "foundry.integration.sglang.runtime") + runtime.preallocate_for_load_mode = lambda: events.append(("preallocate",)) + runtime.log_alloc_offset = lambda label: events.append(("offset", label)) + + # foundry.graph — Foundry graph type + capture context + build primitives. + graph_module = _install_module(monkeypatch, "foundry.graph") + + class FoundryCUDAGraph: + instances: list = [] + pending_sentinel = object() + + def __init__(self): + FoundryCUDAGraph.instances.append(self) + + @staticmethod + def start_graph_builds(paths, pool=None, num_threads=4): + events.append(("start_graph_builds", tuple(paths), pool, num_threads)) + return FoundryCUDAGraph.pending_sentinel + + @staticmethod + def finish_one_graph_load(pending, index): + assert pending is FoundryCUDAGraph.pending_sentinel + events.append(("finish_one", index)) + return (f"graph-{index}", ("tensors", index)) + + @staticmethod + def finish_graph_loads(pending): + raise AssertionError("interleaved LOAD must not call finish_graph_loads") + + class _GraphCtx: + def __init__(self, cuda_graph, pool=None, stream=None): + events.append(("graph_ctx", pool, stream)) + + def __enter__(self): + events.append(("ctx_enter",)) + return self + + def __exit__(self, *args): + events.append(("ctx_exit",)) + return False + + FoundryCUDAGraph.instances = [] + graph_module.CUDAGraph = FoundryCUDAGraph + graph_module.graph = _GraphCtx + + graph_ops = _install_module(monkeypatch, "foundry.integration.sglang.graph_ops") + graph_ops.save_graph = lambda graph, output, key: events.append(("save_graph", key)) + graph_ops._scan_graph_files = lambda workspace_dir: list(graph_files) + graph_ops._unpack_output = lambda tensors: ("unpacked", tensors) + graph_ops._require_dense_shape_key = lambda key: int(key.size) + + # Pinned SGLang surface. + _install_module(monkeypatch, "sglang").__path__ = [] + _install_module(monkeypatch, "sglang.srt").__path__ = [] + _install_module(monkeypatch, "sglang.srt.distributed").__path__ = [] + _install_module(monkeypatch, "sglang.srt.distributed.device_communicators").__path__ = [] + pynccl = _install_module( + monkeypatch, "sglang.srt.distributed.device_communicators.pynccl_allocator" + ) + pynccl.set_graph_pool_id = lambda pool: events.append(("set_graph_pool_id", pool)) + + _install_module(monkeypatch, "sglang.srt.model_executor").__path__ = [] + _install_module(monkeypatch, "sglang.srt.model_executor.runner_backend").__path__ = [] + _install_module(monkeypatch, "sglang.srt.model_executor.runner_utils").__path__ = [] + pool_mod = _install_module(monkeypatch, "sglang.srt.model_executor.runner_utils.pool") + pool_mod.get_or_create_global_graph_memory_pool = lambda device_module: "pool-handle" + + fcgb = _install_module( + monkeypatch, "sglang.srt.model_executor.runner_backend.full_cuda_graph_backend" + ) + + class FullCudaGraphBackend: + def __init__(self, cuda_graph_runner, *, enable_memory_saver=False): + self._graphs = {} + self._outputs = {} + self._pool = None + self._capture_stream = None + self._device_module = SimpleNamespace(synchronize=lambda: events.append(("sync",))) + self._tp_group = SimpleNamespace(barrier=lambda: events.append(("barrier",))) + self.super_calls = [] + + def capture_one(self, shape_key, forward_fn, capture_inputs=None, post_warmup_hook=None): + self.super_calls.append((shape_key, capture_inputs, post_warmup_hook)) + return "super-capture" + + fcgb.FullCudaGraphBackend = FullCudaGraphBackend + + module = _load_module_from(BACKEND_PATH, "sglang_backend_load_contract") + return module, FoundryCUDAGraph, events + + +def _load_module_from(path: Path, name: str): + spec = importlib.util.spec_from_file_location(name, path) + assert spec is not None and spec.loader is not None + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + return module + + +def _archive(*sizes: int): + return [ + (i, f"graph_{i}_FULL_t{bs}_r{bs}_UX_pcN.json", {"shape_key": ShapeKey(size=bs)}) + for i, bs in enumerate(sizes) + ] + + +def _drive_upstream_load(inst, shape_keys, events, *, complete=True): + """Simulate the pinned ``DecodeCudaGraphRunner.capture`` LOAD lifecycle. + + Upstream fills seq_lens, enters ``backend.capture_session`` and, per shape, + runs ``capture_prepare`` + ``init_forward_metadata_out_graph`` (recorded as a + ``metadata`` event) then calls ``backend.capture_one`` inside a + ``ForwardContext``. + """ + forwards: list = [] + + def _forward(): + forwards.append("forward") + return "forward-out" + + with inst.capture_session("load-stream"): + for sk in shape_keys: + events.append(("metadata", sk.size)) + inst.capture_one(sk, _forward, capture_inputs=None, post_warmup_hook=None) + if not complete: + break + return forwards + + +def _build(monkeypatch, mode, *sizes): + backend_mod, graph_cls, events = _load_backend(monkeypatch, mode, graph_files=_archive(*sizes)) + cls = backend_mod.build_foundry_full_backend_class() + inst = cls(SimpleNamespace()) + return inst, graph_cls, events + + +def test_load_capture_session_establishes_pool_and_starts_builds_once(monkeypatch): + inst, _graph_cls, events = _build(monkeypatch, ExtensionMode.LOAD, 8, 4, 2) + keys = [ShapeKey(size=8), ShapeKey(size=4), ShapeKey(size=2)] + + forwards = _drive_upstream_load(inst, keys, events) + + # No forward, warmup, or native capture on LOAD. + assert forwards == [] + assert _graph_cls.instances == [] + + # Exactly one start_graph_builds over all archived paths, into the inherited + # shared pool, before any finish. + starts = [e for e in events if e[0] == "start_graph_builds"] + assert len(starts) == 1 + assert [Path(p).name for p in starts[0][1]] == [ + "graph_0_FULL_t8_r8_UX_pcN.json", + "graph_1_FULL_t4_r4_UX_pcN.json", + "graph_2_FULL_t2_r2_UX_pcN.json", + ] + assert starts[0][2] == "pool-handle" + assert starts[0][3] == 4 + assert inst._pool == "pool-handle" + assert ("set_graph_pool_id", "pool-handle") in events + + # start precedes every finish. + start_idx = next(i for i, e in enumerate(events) if e[0] == "start_graph_builds") + finish_idxs = [i for i, e in enumerate(events) if e[0] == "finish_one"] + assert finish_idxs and all(start_idx < i for i in finish_idxs) + + +def test_load_interleaves_upstream_metadata_then_per_index_finish(monkeypatch): + inst, _graph_cls, events = _build(monkeypatch, ExtensionMode.LOAD, 8, 4, 2) + keys = [ShapeKey(size=8), ShapeKey(size=4), ShapeKey(size=2)] + + _drive_upstream_load(inst, keys, events) + + interleaved = [e for e in events if e[0] in ("metadata", "finish_one")] + assert interleaved == [ + ("metadata", 8), + ("finish_one", 0), + ("metadata", 4), + ("finish_one", 1), + ("metadata", 2), + ("finish_one", 2), + ] + + +def test_load_populates_backend_maps_by_shape_key(monkeypatch): + inst, _graph_cls, events = _build(monkeypatch, ExtensionMode.LOAD, 8, 4) + keys = [ShapeKey(size=8), ShapeKey(size=4)] + + _drive_upstream_load(inst, keys, events) + + assert inst._graphs[ShapeKey(size=8)] == "graph-0" + assert inst._graphs[ShapeKey(size=4)] == "graph-1" + assert inst._outputs[ShapeKey(size=8)] == ("unpacked", ("tensors", 0)) + assert inst._outputs[ShapeKey(size=4)] == ("unpacked", ("tensors", 1)) + + +def test_load_preallocates_before_starting_builds(monkeypatch): + inst, _graph_cls, events = _build(monkeypatch, ExtensionMode.LOAD, 8) + _drive_upstream_load(inst, [ShapeKey(size=8)], events) + + names = [e[0] for e in events] + assert names.index("preallocate") < names.index("start_graph_builds") + offset_labels = [e[1] for e in events if e[0] == "offset"] + assert "before_preallocate" in offset_labels + assert "after_preallocate" in offset_labels + + +def test_load_final_offset_label_and_summary(monkeypatch, caplog): + inst, _graph_cls, events = _build(monkeypatch, ExtensionMode.LOAD, 8, 4) + with caplog.at_level("INFO"): + _drive_upstream_load(inst, [ShapeKey(size=8), ShapeKey(size=4)], events) + + offset_labels = [e[1] for e in events if e[0] == "offset"] + # Final offset log uses the label the runtime harness scanner matches. + assert offset_labels[-1] == "after_load_all_graphs" + summary = [r for r in caplog.records if "Loaded 2 SGLang graphs" in r.getMessage()] + assert len(summary) == 1 + + +def test_load_never_calls_finish_graph_loads(monkeypatch): + inst, _graph_cls, events = _build(monkeypatch, ExtensionMode.LOAD, 8) + # finish_graph_loads raises if touched; a clean per-index finish proves the + # batched variant is unused. + _drive_upstream_load(inst, [ShapeKey(size=8)], events) + assert [e for e in events if e[0] == "finish_one"] == [("finish_one", 0)] + + +def test_load_requires_started_builds_before_capture_one(monkeypatch): + inst, _graph_cls, _events = _build(monkeypatch, ExtensionMode.LOAD, 8) + # capture_one outside capture_session (no start_graph_builds) must fail closed. + with pytest.raises(RuntimeError, match="not started"): + inst.capture_one(ShapeKey(size=8), lambda: None) + + +def test_load_rejects_shape_order_mismatch(monkeypatch): + inst, _graph_cls, events = _build(monkeypatch, ExtensionMode.LOAD, 8, 4) + # Archive order is [8, 4]; requesting 4 first is a fatal order mismatch. + with pytest.raises(RuntimeError, match="order mismatch"): + _drive_upstream_load(inst, [ShapeKey(size=4)], events) + + +def test_load_rejects_extra_shape(monkeypatch): + inst, _graph_cls, events = _build(monkeypatch, ExtensionMode.LOAD, 8) + # Archive has one graph; a second requested shape is fatal. + with pytest.raises(RuntimeError, match="more graph shapes"): + _drive_upstream_load(inst, [ShapeKey(size=8), ShapeKey(size=4)], events) + + +def test_load_rejects_missing_shape_on_session_exit(monkeypatch): + inst, _graph_cls, events = _build(monkeypatch, ExtensionMode.LOAD, 8, 4) + # Only one of two archived graphs consumed: capture_session exit fails closed. + with pytest.raises(RuntimeError, match="1/2"): + _drive_upstream_load(inst, [ShapeKey(size=8)], events, complete=False) + + +def test_load_then_cleanup_blocks_recapture(monkeypatch): + inst, _graph_cls, events = _build(monkeypatch, ExtensionMode.LOAD, 8) + _drive_upstream_load(inst, [ShapeKey(size=8)], events) + assert inst._graphs # populated + + inst.cleanup() + assert inst._graphs == {} + assert inst._pool is None + + with pytest.raises(RuntimeError, match="recapture"), inst.capture_session("load-stream"): + pass diff --git a/tests/test_sglang_hook_targets.py b/tests/test_sglang_hook_targets.py new file mode 100644 index 00000000..e0f20e68 --- /dev/null +++ b/tests/test_sglang_hook_targets.py @@ -0,0 +1,833 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the Foundry project +"""CPU-only contracts for current-main SGLang hook targets and callbacks.""" + +from __future__ import annotations + +import importlib.util +import sys +from enum import Enum +from pathlib import Path +from types import ModuleType, SimpleNamespace + +import pytest + +REPO_ROOT = Path(__file__).parents[1] +HOOKS_PATH = REPO_ROOT / "python" / "foundry" / "integration" / "sglang" / "hooks.py" + +EXPECTED_HOOKS = { + ( + "sglang.srt.server_args.ServerArgs.__post_init__", + "AFTER", + ), + ( + "sglang.srt.entrypoints.engine.Engine._launch_scheduler_processes", + "AROUND", + ), + ( + "sglang.srt.model_executor.model_runner.ModelRunner.init_torch_distributed", + "AROUND", + ), + ( + "sglang.srt.model_executor.model_runner.ModelRunner.alloc_memory_pool", + "AROUND", + ), + ( + "sglang.srt.mem_cache.kv_cache_configurator." + "KVCacheConfigurator._resolve_memory_pool_config", + "AROUND", + ), + ( + "sglang.srt.model_executor.runner_backend.full_cuda_graph_backend.FullCudaGraphBackend", + "REPLACE", + ), + ( + "sglang.srt.model_executor.runner.decode_cuda_graph_runner.DecodeCudaGraphRunner.capture", + "AROUND", + ), + ( + "sglang.srt.distributed.device_communicators.torch_symm_mem." + "TorchSymmMemCommunicator.__init__", + "AROUND", + ), + ( + "sglang.srt.distributed.device_communicators.triton_symm_mem_ag." + "MultimemAllGatherer.__call__", + "AROUND", + ), +} + + +def _install_module(monkeypatch, name: str) -> ModuleType: + module = ModuleType(name) + monkeypatch.setitem(sys.modules, name, module) + if "." in name: + parent_name, child_name = name.rsplit(".", 1) + parent = sys.modules.get(parent_name) + if parent is not None: + setattr(parent, child_name, module) + return module + + +def _load_hooks(monkeypatch): + monkeypatch.syspath_prepend(str(REPO_ROOT / "python")) + foundry = _install_module(monkeypatch, "foundry") + foundry.__path__ = [] + _install_module(monkeypatch, "foundry.integration").__path__ = [] + _install_module(monkeypatch, "foundry.integration.sglang").__path__ = [] + + # The capture-hook registration builds the REPLACE class lazily from the + # backend module; stub it so this contract stays SGLang-free. + backend = _install_module(monkeypatch, "foundry.integration.sglang.backend") + backend.build_foundry_full_backend_class = lambda: type("FoundryFullCudaGraphBackend", (), {}) + + events = [] + runtime = _install_module(monkeypatch, "foundry.integration.sglang.runtime") + runtime.setup_ld_preload_env = lambda: events.append(("preload",)) + runtime.setup_graph_extension = lambda server_args, **ranks: events.append( + ("setup", server_args, ranks) + ) + runtime.log_alloc_offset = lambda label: events.append(("log", label)) + runtime.skip_to_scratch_boundary = lambda: events.append(("scratch",)) + + class ExtensionMode(Enum): + NONE = "none" + SAVE = "save" + LOAD = "load" + + current_mode = [ExtensionMode.SAVE] + config = _install_module(monkeypatch, "foundry.integration.sglang.config") + config.CUDAGraphExtensionMode = ExtensionMode + config.get_graph_extension_mode = lambda: current_mode[0] + config.get_workspace_root = lambda: None + + sglang = _install_module(monkeypatch, "sglang") + sglang.__path__ = [] + srt = _install_module(monkeypatch, "sglang.srt") + srt.__path__ = [] + plugins = _install_module(monkeypatch, "sglang.srt.plugins") + plugins.__path__ = [] + registry_module = _install_module(monkeypatch, "sglang.srt.plugins.hook_registry") + + class HookType(Enum): + BEFORE = "before" + AFTER = "after" + AROUND = "around" + REPLACE = "replace" + + class HookRegistry: + calls = [] + + @classmethod + def register(cls, target, hook, hook_type=HookType.AFTER): + cls.calls.append((target, hook, hook_type)) + + registry_module.HookRegistry = HookRegistry + registry_module.HookType = HookType + + spec = importlib.util.spec_from_file_location("sglang_target_hooks_contract", HOOKS_PATH) + assert spec is not None and spec.loader is not None + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + return module, HookRegistry, events, current_mode, ExtensionMode + + +@pytest.fixture +def hook_context(monkeypatch): + for name, value in ( + ("NCCL_CUMEM_ENABLE", "0"), + ("NCCL_NVLS_ENABLE", "0"), + ("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:False"), + ("SGLANG_ENABLE_POST_CAPTURE_KV_SIZING", "0"), + ): + monkeypatch.setenv(name, value) + return _load_hooks(monkeypatch) + + +def _registered_callback(registry, target: str): + return next(hook for path, hook, _hook_type in registry.calls if path == target) + + +def _supported_args(): + return SimpleNamespace( + nnodes=1, + tp_size=2, + pp_size=1, + dp_size=1, + dcp_size=1, + ep_size=1, + enable_dp_attention=False, + moe_a2a_backend="none", + disaggregation_mode="null", + enable_pdmux=False, + disable_custom_all_reduce=True, + enable_mscclpp=False, + enable_torch_symm_mem=True, + enable_symm_mem=False, + enforce_disable_flashinfer_allreduce_fusion=True, + enable_nccl_nvls=False, + flashinfer_allreduce_fusion_backend=None, + disable_flashinfer_autotune=True, + enable_tokenizer_batch_encode=True, + enable_dynamic_batch_tokenizer=False, + disable_cuda_graph_padding=False, + decode_log_interval=1, + attention_backend="flashinfer", + prefill_attention_backend=None, + decode_attention_backend=None, + cuda_graph_config=SimpleNamespace( + decode=SimpleNamespace(backend="full", max_bs=256, bs=[256]), + prefill=SimpleNamespace(backend="disabled"), + ), + ) + + +def test_register_hooks_preserves_task_2_and_adds_memory_pool_targets( + hook_context, +) -> None: + hooks, registry, _events, _current_mode, _extension_mode = hook_context + + hooks.register_hooks() + + recorded = {(target, hook_type.name) for target, _callback, hook_type in registry.calls} + assert recorded == EXPECTED_HOOKS + + +def test_register_hooks_is_idempotent(hook_context) -> None: + hooks, registry, _events, _current_mode, _extension_mode = hook_context + + hooks.register_hooks() + hooks.register_hooks() + + assert len(registry.calls) == len(EXPECTED_HOOKS) + + +def test_after_post_init_uses_pinned_after_calling_convention(hook_context) -> None: + hooks, registry, _events, _current_mode, _extension_mode = hook_context + hooks.register_hooks() + callback = _registered_callback(registry, "sglang.srt.server_args.ServerArgs.__post_init__") + result = object() + + assert callback(result, _supported_args()) is result + + invalid = _supported_args() + invalid.tp_size = 4 + with pytest.raises(RuntimeError, match="tp_size"): + callback(result, invalid) + + +def test_spawn_hook_preserves_classmethod_descriptor_arguments(hook_context) -> None: + hooks, registry, events, _current_mode, _extension_mode = hook_context + hooks.register_hooks() + callback = _registered_callback( + registry, + "sglang.srt.entrypoints.engine.Engine._launch_scheduler_processes", + ) + + class Engine: + pass + + def original(cls, server_args, port_args, run_scheduler_process_func): + events.append( + ( + "original_spawn", + cls, + server_args, + port_args, + run_scheduler_process_func, + ) + ) + return "spawned" + + run_scheduler = object() + result = callback(original, Engine, "server", "ports", run_scheduler) + + assert result == "spawned" + assert events == [ + ("preload",), + ("original_spawn", Engine, "server", "ports", run_scheduler), + ] + + +def test_rank_hook_uses_parallel_state_and_preserves_setup_order(hook_context) -> None: + hooks, registry, events, _current_mode, _extension_mode = hook_context + hooks.register_hooks() + callback = _registered_callback( + registry, + "sglang.srt.model_executor.model_runner.ModelRunner.init_torch_distributed", + ) + server_args = object() + runner = SimpleNamespace( + server_args=server_args, + ps=SimpleNamespace(tp_rank=1, pp_rank=0, dp_rank=None, gpu_id=1), + ) + + def original(self, *, timeout): + events.append(("original_dist", self, timeout)) + return "distributed" + + result = callback(original, runner, timeout=30) + + assert result == "distributed" + assert events == [ + ( + "setup", + server_args, + {"tp_rank": 1, "pp_rank": 0, "dp_rank": None}, + ), + ("log", "after_setup_graph_ext"), + ("original_dist", runner, 30), + ("log", "after_init_torch_dist"), + ("scratch",), + ("log", "after_scratch_skip"), + ] + + +def test_rank_hook_is_passthrough_when_extension_is_inactive(hook_context) -> None: + hooks, registry, events, current_mode, extension_mode = hook_context + hooks.register_hooks() + callback = _registered_callback( + registry, + "sglang.srt.model_executor.model_runner.ModelRunner.init_torch_distributed", + ) + current_mode[0] = extension_mode.NONE + runner = SimpleNamespace(server_args=object(), ps=SimpleNamespace()) + + def original(self): + events.append(("original_dist", self)) + return "distributed" + + assert callback(original, runner) == "distributed" + assert events == [("original_dist", runner)] + + +# --- symmetric-memory collective hook contracts ---------------------------- + +TORCH_SYMM_MEM_INIT = ( + "sglang.srt.distributed.device_communicators.torch_symm_mem.TorchSymmMemCommunicator.__init__" +) +MULTIMEM_ALL_GATHER_CALL = ( + "sglang.srt.distributed.device_communicators.triton_symm_mem_ag.MultimemAllGatherer.__call__" +) + +_BF16 = object() + + +def _prod(shape) -> int: + result = 1 + for dim in shape: + result *= dim + return result + + +class _SymTensor: + """Minimal tensor stand-in tracking shape/dtype/contiguity/numel for the + symmetric-memory gather callback (the local VM has no torch/CUDA). + + ``_storage`` models the underlying allocation: a slice or ``view`` shares its + parent's ``_storage`` object, so a test can prove the tensor handed to + ``multimem_all_gather_out`` is a view of the retained communicator buffer's + storage rather than a coincidentally same-shaped tensor. ``movedim`` returns + a non-contiguous permuted view that still shares storage, but ``reshape`` on + a non-contiguous tensor materializes a fresh contiguous allocation (new + ``_storage``) — mirroring real torch, where the ``movedim(0,1).reshape(...)`` + that builds the gather's *return value* legitimately does not alias the + buffer. ``_lineage`` records the exact op chain for the same reason.""" + + def __init__( + self, shape, *, dtype=_BF16, contiguous=True, numel=None, ptr=0, storage=None, lineage=None + ): + self.shape = tuple(shape) + self.dtype = dtype + self._contiguous = contiguous + self._numel = numel if numel is not None else _prod(self.shape) + self._ptr = ptr + self._storage = storage if storage is not None else object() + self._lineage = list(lineage) if lineage is not None else [] + + @property + def ndim(self) -> int: + return len(self.shape) + + def is_contiguous(self) -> bool: + return self._contiguous + + def numel(self) -> int: + return self._numel + + def data_ptr(self) -> int: + return self._ptr + + def _derive(self, shape, *, numel, op, storage=None, contiguous=True): + return _SymTensor( + shape, + dtype=self.dtype, + contiguous=contiguous, + numel=numel, + ptr=self._ptr, + storage=storage if storage is not None else self._storage, + lineage=self._lineage + [op], + ) + + def __getitem__(self, key): + if isinstance(key, slice): + # A leading contiguous slice stays a contiguous view of the storage. + stop = key.stop + return self._derive((stop,), numel=stop, op=("slice", stop)) + raise KeyError(key) + + def view(self, *shape): + # ``view`` requires contiguous input and shares storage with its parent. + if len(shape) == 1 and isinstance(shape[0], (tuple, list)): + shape = tuple(shape[0]) + return self._derive(shape, numel=_prod(shape), op=("view", tuple(shape))) + + def movedim(self, src, dst): + # A permuted view: still aliases storage, but is now non-contiguous. + dims = list(self.shape) + moved = dims.pop(src) + dims.insert(dst, moved) + return self._derive( + tuple(dims), numel=self._numel, op=("movedim", src, dst), contiguous=False + ) + + def reshape(self, *shape): + # Contiguous reshape is a view; non-contiguous reshape (e.g. after + # movedim) copies into a fresh contiguous allocation with new storage. + if len(shape) == 1 and isinstance(shape[0], (tuple, list)): + shape = tuple(shape[0]) + storage = self._storage if self._contiguous else object() + return self._derive( + shape, numel=_prod(shape), op=("reshape", tuple(shape)), storage=storage + ) + + +def _install_symmetric_env( + monkeypatch, + *, + world_size=2, + group_ranks=(0, 1), + signal_pad_current=0, + max_blocks=32, + numel_per_thread=8, +): + order: list[str] = [] + op_calls: list[tuple] = [] + set_calls: list[int] = [] + + torch = _install_module(monkeypatch, "torch") + torch.bfloat16 = _BF16 + + class _Dist: + @staticmethod + def get_world_size(group): + return world_size + + @staticmethod + def get_process_group_ranks(group): + return list(group_ranks) + + torch.distributed = _Dist() + + def _multimem_all_gather_out(x, group_name, out): + op_calls.append((x, group_name, out)) + + torch.ops = SimpleNamespace( + symm_mem=SimpleNamespace(multimem_all_gather_out=_multimem_all_gather_out) + ) + + _install_module(monkeypatch, "sglang.srt.distributed").__path__ = [] + _install_module(monkeypatch, "sglang.srt.distributed.device_communicators").__path__ = [] + torch_symm_mem = _install_module( + monkeypatch, "sglang.srt.distributed.device_communicators.torch_symm_mem" + ) + + class _SymmMemModule: + def __init__(self): + self.size = signal_pad_current + + def get_signal_pad_size(self): + return self.size + + def set_signal_pad_size(self, value): + self.size = value + set_calls.append(value) + + torch_symm_mem.torch_symm_mem = _SymmMemModule() + triton = _install_module( + monkeypatch, "sglang.srt.distributed.device_communicators.triton_symm_mem_ag" + ) + triton._MAX_BLOCKS = max_blocks + triton._NUMEL_PER_THREAD = numel_per_thread + + cge = _install_module(monkeypatch, "foundry.ops") + cge.stop_allocation_region = lambda: order.append("stop") + cge.resume_allocation_region = lambda: order.append("resume") + + return SimpleNamespace( + order=order, + op_calls=op_calls, + set_calls=set_calls, + symm_mem=torch_symm_mem.torch_symm_mem, + max_blocks=max_blocks, + world_size=world_size, + group_ranks=tuple(group_ranks), + ) + + +def _make_communicator(*, buffer, group_name, world_size): + return SimpleNamespace( + buffer=buffer, + group=SimpleNamespace(group_name=group_name), + world_size=world_size, + ) + + +def test_symmetric_hooks_are_registered_as_around(hook_context) -> None: + hooks, registry, _events, _current_mode, _extension_mode = hook_context + hooks.register_hooks() + + recorded = {(target, hook_type.name) for target, _callback, hook_type in registry.calls} + assert (TORCH_SYMM_MEM_INIT, "AROUND") in recorded + assert (MULTIMEM_ALL_GATHER_CALL, "AROUND") in recorded + assert len(registry.calls) == 9 + assert len({target for target, _hook, _ht in registry.calls}) == 9 + + +def test_symm_mem_init_stops_and_resumes_around_original(hook_context, monkeypatch) -> None: + hooks, registry, _events, _current_mode, _extension_mode = hook_context + hooks.register_hooks() + env = _install_symmetric_env(monkeypatch, world_size=2, group_ranks=(0, 1)) + callback = _registered_callback(registry, TORCH_SYMM_MEM_INIT) + + buffer = _SymTensor((4096,), numel=4096, ptr=0xABCD000) + communicator = _make_communicator(buffer=None, group_name="tp:0", world_size=2) + + def original(comm, group, device): + env.order.append("original") + comm.buffer = buffer + return None + + callback(original, communicator, "group", "cuda:0") + + assert env.order == ["stop", "original", "resume"] + group_ranks = env.group_ranks + assert hooks._torch_symm_mem_communicators[group_ranks] is communicator + + +def test_symm_mem_init_signal_pad_minimum(hook_context, monkeypatch) -> None: + hooks, registry, _events, _current_mode, _extension_mode = hook_context + hooks.register_hooks() + env = _install_symmetric_env( + monkeypatch, world_size=2, group_ranks=(0, 1), signal_pad_current=8, max_blocks=32 + ) + callback = _registered_callback(registry, TORCH_SYMM_MEM_INIT) + + buffer = _SymTensor((4096,), numel=4096) + communicator = _make_communicator(buffer=None, group_name="tp:0", world_size=2) + + def original(comm, *args, **kwargs): + env.order.append("original") + comm.buffer = buffer + + callback(original, communicator, "group", "cuda:0") + + expected_minimum = env.max_blocks * env.world_size * 4 + assert env.set_calls == [max(8, expected_minimum)] + assert env.symm_mem.size >= expected_minimum + + +def test_symm_mem_init_resumes_region_on_original_exception(hook_context, monkeypatch) -> None: + hooks, registry, _events, _current_mode, _extension_mode = hook_context + hooks.register_hooks() + env = _install_symmetric_env(monkeypatch) + callback = _registered_callback(registry, TORCH_SYMM_MEM_INIT) + + communicator = _make_communicator(buffer=None, group_name="tp:0", world_size=2) + + def original(comm, *args, **kwargs): + env.order.append("original") + raise RuntimeError("constructor blew up") + + with pytest.raises(RuntimeError, match="constructor blew up"): + callback(original, communicator, "group", "cuda:0") + + # resume_allocation_region must run in finally even on failure. + assert env.order == ["stop", "original", "resume"] + assert hooks._torch_symm_mem_communicators == {} + + +def test_symm_mem_init_replaces_communicator_for_same_rank_tuple(hook_context, monkeypatch) -> None: + hooks, registry, _events, _current_mode, _extension_mode = hook_context + hooks.register_hooks() + env = _install_symmetric_env(monkeypatch, world_size=2, group_ranks=(0, 1)) + callback = _registered_callback(registry, TORCH_SYMM_MEM_INIT) + + def make_original(buffer): + def original(comm, *args, **kwargs): + comm.buffer = buffer + + return original + + first = _make_communicator(buffer=None, group_name="tp:0", world_size=2) + second = _make_communicator(buffer=None, group_name="tp:0", world_size=2) + callback(make_original(_SymTensor((4096,), numel=4096)), first, "group", "cuda:0") + callback(make_original(_SymTensor((4096,), numel=4096)), second, "group", "cuda:0") + + assert len(hooks._torch_symm_mem_communicators) == 1 + assert hooks._torch_symm_mem_communicators[env.group_ranks] is second + + +def test_symm_mem_init_inactive_passthrough(hook_context) -> None: + hooks, registry, _events, current_mode, extension_mode = hook_context + hooks.register_hooks() + current_mode[0] = extension_mode.NONE + callback = _registered_callback(registry, TORCH_SYMM_MEM_INIT) + + called = [] + communicator = _make_communicator(buffer=None, group_name="tp:0", world_size=2) + + def original(comm, *args, **kwargs): + called.append((comm, args, kwargs)) + return "constructed" + + assert callback(original, communicator, "group", "cuda:0") == "constructed" + assert called == [(communicator, ("group", "cuda:0"), {})] + assert hooks._torch_symm_mem_communicators == {} + + +def _run_gather(hooks, registry, env, *, world_size, tokens, local_hidden, capacity=None): + callback = _registered_callback(registry, MULTIMEM_ALL_GATHER_CALL) + if capacity is None: + capacity = world_size * tokens * local_hidden + buffer = _SymTensor((capacity,), numel=capacity, ptr=0xBEEF000) + communicator = _make_communicator(buffer=buffer, group_name="tp:grp", world_size=world_size) + hooks._torch_symm_mem_communicators[env.group_ranks] = communicator + gatherer = SimpleNamespace() + x = _SymTensor((tokens, local_hidden), dtype=_BF16, numel=tokens * local_hidden, ptr=0x1000) + + def original(g, tensor): + raise AssertionError("active gather must not call the upstream lazy builder") + + output = callback(original, gatherer, x) + return output, communicator, x, gatherer + + +def test_multimem_all_gather_uses_retained_communicator(hook_context, monkeypatch) -> None: + hooks, registry, _events, _current_mode, _extension_mode = hook_context + hooks.register_hooks() + env = _install_symmetric_env(monkeypatch, world_size=2) + + tokens, local_hidden = 8, 16 + output, communicator, x, _gatherer = _run_gather( + hooks, registry, env, world_size=2, tokens=tokens, local_hidden=local_hidden + ) + + assert len(env.op_calls) == 1 + op_x, op_group_name, op_out = env.op_calls[0] + assert op_x is x + assert op_group_name == communicator.group.group_name + assert op_out.shape == (2 * tokens, local_hidden) + assert output.shape == (tokens, 2 * local_hidden) + + # The tensor handed to multimem_all_gather_out must be an actual view of the + # RETAINED communicator buffer's storage (not merely a same-shaped tensor), + # reached by the exact buffer[:required_numel].view(world_size*tokens, + # local_hidden) chain. This is the critical aliasing contract: the op writes + # gathered data directly into the retained symmetric buffer. + required_numel = 2 * tokens * local_hidden + assert op_out._storage is communicator.buffer._storage + assert op_out._storage is not x._storage + assert op_out._lineage == [ + ("slice", required_numel), + ("view", (2 * tokens, local_hidden)), + ] + # The RETURNED tensor is a movedim+reshape rearrangement of the op buffer. + # reshape after movedim operates on a non-contiguous view, which in real + # torch materializes a fresh contiguous allocation; asserting the return + # value aliases the buffer storage would be production-inaccurate. We only + # require the correct shape and that it is the rearranged, materialized + # result — distinct from the op output's buffer-aliased storage. + assert output.shape == (tokens, 2 * local_hidden) + assert output._storage is not communicator.buffer._storage + assert output._lineage[-2:] == [ + ("movedim", 0, 1), + ("reshape", (tokens, 2 * local_hidden)), + ] + + +def test_multimem_all_gather_marks_gatherer_once(hook_context, monkeypatch, caplog) -> None: + hooks, registry, _events, _current_mode, _extension_mode = hook_context + hooks.register_hooks() + env = _install_symmetric_env(monkeypatch, world_size=2) + + callback = _registered_callback(registry, MULTIMEM_ALL_GATHER_CALL) + capacity = 2 * 8 * 16 + buffer = _SymTensor((capacity,), numel=capacity) + communicator = _make_communicator(buffer=buffer, group_name="tp:grp", world_size=2) + hooks._torch_symm_mem_communicators[env.group_ranks] = communicator + gatherer = SimpleNamespace() + + def original(g, tensor): + raise AssertionError("active gather must not call the upstream lazy builder") + + with caplog.at_level("INFO"): + for _ in range(2): + x = _SymTensor((8, 16), dtype=_BF16, numel=8 * 16) + callback(original, gatherer, x) + + assert gatherer._foundry_observed is True + markers = [r for r in caplog.records if "Multimem all-gather buffer address" in r.getMessage()] + assert len(markers) == 1 + assert len(env.op_calls) == 2 + + +def test_multimem_all_gather_requires_exactly_one_communicator(hook_context, monkeypatch) -> None: + hooks, registry, _events, _current_mode, _extension_mode = hook_context + hooks.register_hooks() + env = _install_symmetric_env(monkeypatch, world_size=2) + callback = _registered_callback(registry, MULTIMEM_ALL_GATHER_CALL) + + gatherer = SimpleNamespace() + x = _SymTensor((8, 16), dtype=_BF16, numel=8 * 16) + + def original(g, tensor): + raise AssertionError("active gather must not call the upstream lazy builder") + + # Zero retained communicators. + with pytest.raises(RuntimeError, match="exactly one"): + callback(original, gatherer, x) + + # Two retained communicators (distinct rank tuples). + buf = _SymTensor((2 * 8 * 16,), numel=2 * 8 * 16) + hooks._torch_symm_mem_communicators[(0, 1)] = _make_communicator( + buffer=buf, group_name="a", world_size=2 + ) + hooks._torch_symm_mem_communicators[(2, 3)] = _make_communicator( + buffer=buf, group_name="b", world_size=2 + ) + with pytest.raises(RuntimeError, match="exactly one"): + callback(original, gatherer, x) + + +@pytest.mark.parametrize( + "x", + [ + _SymTensor((2, 4, 8), dtype=_BF16, numel=64), + _SymTensor((8, 16), dtype=object(), numel=128), + _SymTensor((8, 16), dtype=_BF16, contiguous=False, numel=128), + ], +) +def test_multimem_all_gather_rejects_bad_input(hook_context, monkeypatch, x) -> None: + hooks, registry, _events, _current_mode, _extension_mode = hook_context + hooks.register_hooks() + env = _install_symmetric_env(monkeypatch, world_size=2) + callback = _registered_callback(registry, MULTIMEM_ALL_GATHER_CALL) + + capacity = 2 * 8 * 16 + buffer = _SymTensor((capacity,), numel=capacity) + hooks._torch_symm_mem_communicators[env.group_ranks] = _make_communicator( + buffer=buffer, group_name="tp:grp", world_size=2 + ) + gatherer = SimpleNamespace() + + def original(g, tensor): + raise AssertionError("active gather must not call the upstream lazy builder") + + with pytest.raises(RuntimeError, match="contiguous rank-2 bfloat16"): + callback(original, gatherer, x) + assert env.op_calls == [] + + +def test_multimem_all_gather_rejects_insufficient_capacity(hook_context, monkeypatch) -> None: + hooks, registry, _events, _current_mode, _extension_mode = hook_context + hooks.register_hooks() + env = _install_symmetric_env(monkeypatch, world_size=2) + callback = _registered_callback(registry, MULTIMEM_ALL_GATHER_CALL) + + tokens, local_hidden = 8, 16 + # Buffer one element short of world_size * x.numel(). + capacity = 2 * tokens * local_hidden - 1 + buffer = _SymTensor((capacity,), numel=capacity) + hooks._torch_symm_mem_communicators[env.group_ranks] = _make_communicator( + buffer=buffer, group_name="tp:grp", world_size=2 + ) + gatherer = SimpleNamespace() + x = _SymTensor((tokens, local_hidden), dtype=_BF16, numel=tokens * local_hidden) + + def original(g, tensor): + raise AssertionError("active gather must not call the upstream lazy builder") + + with pytest.raises(RuntimeError, match="too small"): + callback(original, gatherer, x) + assert env.op_calls == [] + + +def _retain_one(hooks, env, *, world_size=2, capacity=2 * 8 * 16): + buffer = _SymTensor((capacity,), numel=capacity, ptr=0xBEEF000) + hooks._torch_symm_mem_communicators[env.group_ranks] = _make_communicator( + buffer=buffer, group_name="tp:grp", world_size=world_size + ) + + +def _no_fallback_original(g, tensor): + raise AssertionError("active gather must not call the upstream lazy builder") + + +def test_multimem_all_gather_rejects_zero_token_count(hook_context, monkeypatch) -> None: + hooks, registry, _events, _current_mode, _extension_mode = hook_context + hooks.register_hooks() + env = _install_symmetric_env(monkeypatch, world_size=2) + callback = _registered_callback(registry, MULTIMEM_ALL_GATHER_CALL) + _retain_one(hooks, env) + + x = _SymTensor((0, 16), dtype=_BF16, numel=0, ptr=0x1000) + with pytest.raises(RuntimeError, match="positive token count"): + callback(_no_fallback_original, SimpleNamespace(), x) + assert env.op_calls == [] + + +def test_multimem_all_gather_rejects_unaligned_input_pointer(hook_context, monkeypatch) -> None: + hooks, registry, _events, _current_mode, _extension_mode = hook_context + hooks.register_hooks() + env = _install_symmetric_env(monkeypatch, world_size=2) + callback = _registered_callback(registry, MULTIMEM_ALL_GATHER_CALL) + _retain_one(hooks, env) + + # 0x1008 % 16 == 8, so the 128-bit multimem.st store would be misaligned. + x = _SymTensor((8, 16), dtype=_BF16, numel=8 * 16, ptr=0x1008) + with pytest.raises(RuntimeError, match="16-byte aligned"): + callback(_no_fallback_original, SimpleNamespace(), x) + assert env.op_calls == [] + + +def test_multimem_all_gather_rejects_unaligned_hidden(hook_context, monkeypatch) -> None: + hooks, registry, _events, _current_mode, _extension_mode = hook_context + hooks.register_hooks() + env = _install_symmetric_env(monkeypatch, world_size=2, numel_per_thread=8) + callback = _registered_callback(registry, MULTIMEM_ALL_GATHER_CALL) + _retain_one(hooks, env, capacity=2 * 8 * 12) + + # local hidden 12 is not a multiple of _NUMEL_PER_THREAD (8). + x = _SymTensor((8, 12), dtype=_BF16, numel=8 * 12, ptr=0x1000) + with pytest.raises(RuntimeError, match="multiple of"): + callback(_no_fallback_original, SimpleNamespace(), x) + assert env.op_calls == [] + + +def test_multimem_all_gather_inactive_passthrough(hook_context) -> None: + hooks, registry, _events, current_mode, extension_mode = hook_context + hooks.register_hooks() + current_mode[0] = extension_mode.NONE + callback = _registered_callback(registry, MULTIMEM_ALL_GATHER_CALL) + + called = [] + + def original(g, tensor): + called.append((g, tensor)) + return "gathered" + + gatherer = SimpleNamespace() + sentinel = object() + assert callback(original, gatherer, sentinel) == "gathered" + assert called == [(gatherer, sentinel)] diff --git a/tests/test_sglang_memory_pool_reuse.py b/tests/test_sglang_memory_pool_reuse.py new file mode 100644 index 00000000..844ad443 --- /dev/null +++ b/tests/test_sglang_memory_pool_reuse.py @@ -0,0 +1,369 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the Foundry project +"""CPU-only contracts for current-main SGLang memory-pool reuse hooks.""" + +from __future__ import annotations + +import importlib.util +import sys +from dataclasses import asdict, dataclass +from enum import Enum +from pathlib import Path +from types import ModuleType, SimpleNamespace + +import pytest + +REPO_ROOT = Path(__file__).parents[1] +HOOKS_PATH = REPO_ROOT / "python" / "foundry" / "integration" / "sglang" / "hooks.py" + + +@dataclass +class FakeMemoryPoolConfig: + max_total_num_tokens: int + max_running_requests: int | None = None + full_max_total_num_tokens: int | None = None + swa_max_total_num_tokens: int | None = None + c4_max_total_num_tokens: int = 0 + c128_max_total_num_tokens: int = 0 + c4_state_pool_size: int = 0 + c128_state_pool_size: int = 0 + mem_fraction_static: float | None = None + + +def _install_module(monkeypatch, name: str) -> ModuleType: + module = ModuleType(name) + monkeypatch.setitem(sys.modules, name, module) + if "." in name: + parent_name, child_name = name.rsplit(".", 1) + parent = sys.modules.get(parent_name) + if parent is not None: + setattr(parent, child_name, module) + return module + + +def _load_hooks(monkeypatch): + monkeypatch.syspath_prepend(str(REPO_ROOT / "python")) + foundry = _install_module(monkeypatch, "foundry") + foundry.__path__ = [] + _install_module(monkeypatch, "foundry.integration").__path__ = [] + _install_module(monkeypatch, "foundry.integration.sglang").__path__ = [] + + events = [] + saved_state = [None] + rank = [0] + runtime = _install_module(monkeypatch, "foundry.integration.sglang.runtime") + runtime.get_state = lambda: SimpleNamespace(rank=rank[0]) + runtime.log_alloc_offset = lambda label: events.append(("log", label)) + + def create_warmup_state(memory_pool_config): + events.append(("create", memory_pool_config)) + return SimpleNamespace(memory_pool_config=memory_pool_config) + + def save_warmup_state(state): + events.append(("save", state)) + saved_state[0] = state + + def load_warmup_state(): + events.append(("load",)) + if saved_state[0] is None: + raise RuntimeError("Foundry warmup state file not found") + return saved_state[0] + + runtime.create_warmup_state = create_warmup_state + runtime.save_warmup_state = save_warmup_state + runtime.load_warmup_state = load_warmup_state + + class ExtensionMode(Enum): + NONE = "none" + SAVE = "save" + LOAD = "load" + + current_mode = [ExtensionMode.SAVE] + config = _install_module(monkeypatch, "foundry.integration.sglang.config") + config.CUDAGraphExtensionMode = ExtensionMode + config.get_graph_extension_mode = lambda: current_mode[0] + config.get_workspace_root = lambda: None + + torch = _install_module(monkeypatch, "torch") + torch.cuda = SimpleNamespace( + empty_cache=lambda: events.append(("empty_cache",)), + ) + + sglang = _install_module(monkeypatch, "sglang") + sglang.__path__ = [] + srt = _install_module(monkeypatch, "sglang.srt") + srt.__path__ = [] + model_executor = _install_module(monkeypatch, "sglang.srt.model_executor") + model_executor.__path__ = [] + pool_configurator = _install_module( + monkeypatch, + "sglang.srt.model_executor.pool_configurator", + ) + pool_configurator.MemoryPoolConfig = FakeMemoryPoolConfig + + spec = importlib.util.spec_from_file_location("sglang_memory_pool_hooks", HOOKS_PATH) + assert spec is not None and spec.loader is not None + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + return SimpleNamespace( + events=events, + hooks=module, + mode=current_mode, + modes=ExtensionMode, + rank=rank, + saved_state=saved_state, + ) + + +@pytest.fixture +def hook_context(monkeypatch): + return _load_hooks(monkeypatch) + + +def _pool_config() -> FakeMemoryPoolConfig: + return FakeMemoryPoolConfig( + max_total_num_tokens=8192, + max_running_requests=128, + full_max_total_num_tokens=4096, + swa_max_total_num_tokens=8192, + c4_max_total_num_tokens=32, + c128_max_total_num_tokens=16, + c4_state_pool_size=8, + c128_state_pool_size=4, + mem_fraction_static=0.8, + ) + + +def test_save_alloc_serializes_resolved_config_after_upstream_allocation(hook_context) -> None: + runner = SimpleNamespace(memory_pool_config=None) + resolved = _pool_config() + + def original(self, memory_pool_config=None): + assert memory_pool_config is None + hook_context.events.append(("original_alloc",)) + self.memory_pool_config = resolved + return "allocated" + + result = hook_context.hooks._around_alloc_memory_pool(original, runner) + + assert result == "allocated" + assert hook_context.saved_state[0].memory_pool_config == asdict(resolved) + assert hook_context.events == [ + ("log", "before_alloc_memory_pool"), + ("original_alloc",), + ("log", "after_alloc_memory_pool"), + ("create", asdict(resolved)), + ("save", hook_context.saved_state[0]), + ] + + +def test_nonzero_rank_does_not_write_shared_warmup_state(hook_context) -> None: + hook_context.rank[0] = 1 + runner = SimpleNamespace(memory_pool_config=None) + + def original(self): + self.memory_pool_config = _pool_config() + + hook_context.hooks._around_alloc_memory_pool(original, runner) + + assert hook_context.saved_state[0] is None + assert all(event[0] not in {"create", "save"} for event in hook_context.events) + + +def test_load_resolver_reconstructs_exact_config_without_profiling(hook_context) -> None: + config = _pool_config() + hook_context.saved_state[0] = SimpleNamespace(memory_pool_config=asdict(config)) + hook_context.mode[0] = hook_context.modes.LOAD + profiled = [] + + def original(self, pre_model_load_memory): + profiled.append((self, pre_model_load_memory)) + + resolver = SimpleNamespace( + _profile_available_bytes=lambda *_args, **_kwargs: profiled.append("profile"), + ) + resolved = hook_context.hooks._around_resolve_memory_pool_config( + original, + resolver, + 123, + ) + + assert resolved == config + assert resolved is not config + assert profiled == [] + assert hook_context.events == [("load",), ("empty_cache",)] + + +def test_load_config_reaches_normal_target_configure_path(hook_context) -> None: + saved = _pool_config() + hook_context.saved_state[0] = SimpleNamespace(memory_pool_config=asdict(saved)) + hook_context.mode[0] = hook_context.modes.LOAD + profiled = [] + configured = [] + + def original_resolver(self, pre_model_load_memory): + profiled.append(pre_model_load_memory) + return FakeMemoryPoolConfig(max_total_num_tokens=1) + + def wrapped_resolver(self, pre_model_load_memory): + return hook_context.hooks._around_resolve_memory_pool_config( + original_resolver, + self, + pre_model_load_memory, + ) + + class FakeKVCacheConfigurator: + def _resolve_memory_pool_config(self, pre_model_load_memory): + return wrapped_resolver(self, pre_model_load_memory) + + def configure(self, *, pre_model_load_memory): + config = self._resolve_memory_pool_config(pre_model_load_memory) + configured.append(config) + return SimpleNamespace(memory_pool_config=config) + + runner = SimpleNamespace( + memory_pool_config=FakeMemoryPoolConfig(max_total_num_tokens=2), + kv_cache_configurator=FakeKVCacheConfigurator(), + pre_model_load_memory=123, + ) + + def original_alloc(self, memory_pool_config=None): + if memory_pool_config is not None: + self.memory_pool_config = memory_pool_config + result = self.kv_cache_configurator.configure( + pre_model_load_memory=self.pre_model_load_memory + ) + self.memory_pool_config = result.memory_pool_config + return "allocated" + + result = hook_context.hooks._around_alloc_memory_pool(original_alloc, runner) + + assert result == "allocated" + assert configured == [runner.memory_pool_config] + assert configured[0] is runner.memory_pool_config + assert runner.memory_pool_config == saved + assert profiled == [] + assert hook_context.events == [ + ("log", "before_alloc_memory_pool"), + ("load",), + ("empty_cache",), + ("log", "after_alloc_memory_pool"), + ] + + +@pytest.mark.parametrize( + ("saved_config", "message"), + [ + ({}, "saved memory_pool_config"), + ({"max_total_num_tokens": 8192}, "schema"), + ({**asdict(_pool_config()), "unknown_field": 1}, "schema"), + ], +) +def test_load_fails_clearly_for_missing_or_incompatible_config( + hook_context, + saved_config, + message, +) -> None: + hook_context.saved_state[0] = SimpleNamespace(memory_pool_config=saved_config) + hook_context.mode[0] = hook_context.modes.LOAD + + with pytest.raises(RuntimeError, match=message): + hook_context.hooks._around_resolve_memory_pool_config( + lambda *_args, **_kwargs: pytest.fail("original resolver called"), + object(), + 123, + ) + + assert ("empty_cache",) not in hook_context.events + + +@pytest.mark.parametrize( + ("field", "value"), + [ + ("max_total_num_tokens", True), + ("max_total_num_tokens", 0), + ("max_running_requests", "128"), + ("max_running_requests", 0), + ("max_running_requests", None), + ("full_max_total_num_tokens", -1), + ("swa_max_total_num_tokens", []), + ("c4_max_total_num_tokens", -1), + ("c128_max_total_num_tokens", True), + ("c4_state_pool_size", {}), + ("c128_state_pool_size", "4"), + ("mem_fraction_static", True), + ("mem_fraction_static", "0.8"), + ("mem_fraction_static", []), + ("mem_fraction_static", 0), + ("mem_fraction_static", 1.1), + ("mem_fraction_static", float("nan")), + ("mem_fraction_static", 2**53 + 1), + ("mem_fraction_static", 10**1000), + ], +) +def test_load_rejects_invalid_field_values_before_original_or_profile( + hook_context, + field, + value, +) -> None: + saved_config = asdict(_pool_config()) + saved_config[field] = value + hook_context.saved_state[0] = SimpleNamespace(memory_pool_config=saved_config) + hook_context.mode[0] = hook_context.modes.LOAD + calls = [] + resolver = SimpleNamespace( + _profile_available_bytes=lambda *_args, **_kwargs: calls.append("profile"), + ) + + with pytest.raises(RuntimeError, match=field): + hook_context.hooks._around_resolve_memory_pool_config( + lambda *_args, **_kwargs: calls.append("original"), + resolver, + 123, + ) + + assert calls == [] + assert ("empty_cache",) not in hook_context.events + + +def test_load_normalizes_exact_json_integer_for_optional_float(hook_context) -> None: + saved_config = asdict(_pool_config()) + saved_config["mem_fraction_static"] = 1 + hook_context.saved_state[0] = SimpleNamespace(memory_pool_config=saved_config) + hook_context.mode[0] = hook_context.modes.LOAD + + resolved = hook_context.hooks._around_resolve_memory_pool_config( + lambda *_args, **_kwargs: pytest.fail("original resolver called"), + object(), + 123, + ) + + assert resolved.mem_fraction_static == 1.0 + assert type(resolved.mem_fraction_static) is float + + +def test_inactive_memory_hooks_delegate_without_side_effects(hook_context) -> None: + hook_context.mode[0] = hook_context.modes.NONE + runner = object() + resolver = object() + calls = [] + + alloc_result = hook_context.hooks._around_alloc_memory_pool( + lambda self, *args, **kwargs: calls.append(("alloc", self, args, kwargs)) or "allocated", + runner, + "arg", + option=True, + ) + resolve_result = hook_context.hooks._around_resolve_memory_pool_config( + lambda self, value: calls.append(("resolve", self, value)) or "resolved", + resolver, + 123, + ) + + assert alloc_result == "allocated" + assert resolve_result == "resolved" + assert calls == [ + ("alloc", runner, ("arg",), {"option": True}), + ("resolve", resolver, 123), + ] + assert hook_context.events == [] diff --git a/tests/test_sglang_plugin_activation.py b/tests/test_sglang_plugin_activation.py new file mode 100644 index 00000000..a7f78f1a --- /dev/null +++ b/tests/test_sglang_plugin_activation.py @@ -0,0 +1,108 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the Foundry project +"""CPU-only contracts for the Foundry SGLang plugin entry point.""" + +from __future__ import annotations + +import importlib.util +import os +import subprocess +import sys +from pathlib import Path +from types import ModuleType + +import tomllib + +REPO_ROOT = Path(__file__).parents[1] +PYTHON_ROOT = REPO_ROOT / "python" +PLUGIN_PATH = PYTHON_ROOT / "foundry_sglang_plugin" / "__init__.py" +PYPROJECT_PATH = REPO_ROOT / "pyproject.toml" + + +def _load_plugin(): + spec = importlib.util.spec_from_file_location("foundry_plugin_under_test", PLUGIN_PATH) + assert spec is not None and spec.loader is not None + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + return module + + +def test_activate_is_noop_without_config_env(monkeypatch): + monkeypatch.delenv("FOUNDRY_GRAPH_EXTENSION_CONFIG_PATH", raising=False) + plugin = _load_plugin() + plugin.activate() + + +def test_activate_loads_config_and_registers_hooks(monkeypatch, tmp_path): + cfg = tmp_path / "cfg.toml" + cfg.write_text('mode = "save"\n') + monkeypatch.setenv("FOUNDRY_GRAPH_EXTENSION_CONFIG_PATH", str(cfg)) + calls = [] + foundry = ModuleType("foundry") + foundry.__path__ = [] + integration = ModuleType("foundry.integration") + integration.__path__ = [] + sglang = ModuleType("foundry.integration.sglang") + sglang.__path__ = [] + config = ModuleType("foundry.integration.sglang.config") + config.load_graph_extension_config = lambda path: calls.append(("cfg", path)) + hooks = ModuleType("foundry.integration.sglang.hooks") + hooks.register_hooks = lambda: calls.append("hooks") + monkeypatch.setitem(sys.modules, "foundry", foundry) + monkeypatch.setitem(sys.modules, "foundry.integration", integration) + monkeypatch.setitem(sys.modules, "foundry.integration.sglang", sglang) + monkeypatch.setitem(sys.modules, "foundry.integration.sglang.config", config) + monkeypatch.setitem(sys.modules, "foundry.integration.sglang.hooks", hooks) + + plugin = _load_plugin() + plugin.activate() + + assert calls == [("cfg", str(cfg)), "hooks"] + + +def test_inactive_entry_point_import_avoids_foundry_torch_and_native_ops() -> None: + code = """ +import importlib.abc +import sys + +class BlockForbiddenImports(importlib.abc.MetaPathFinder): + def find_spec(self, fullname, path=None, target=None): + if ( + fullname == "foundry" + or fullname.startswith("foundry.") + or fullname == "torch" + or fullname.startswith("torch.") + ): + raise ImportError(f"forbidden import: {fullname}") + return None + +sys.meta_path.insert(0, BlockForbiddenImports()) +import foundry_sglang_plugin +foundry_sglang_plugin.activate() +for forbidden in ("foundry", "foundry.ops", "foundry._C", "torch"): + assert forbidden not in sys.modules, forbidden +""" + env = os.environ.copy() + env.pop("FOUNDRY_GRAPH_EXTENSION_CONFIG_PATH", None) + env["PYTHONPATH"] = str(PYTHON_ROOT) + + result = subprocess.run( + [sys.executable, "-c", code], + cwd=REPO_ROOT, + env=env, + capture_output=True, + text=True, + ) + + assert result.returncode == 0, result.stderr + + +def test_entry_point_uses_explicit_lightweight_package() -> None: + project = tomllib.loads(PYPROJECT_PATH.read_text()) + + assert ( + project["project"]["entry-points"]["sglang.srt.plugins"]["foundry"] + == "foundry_sglang_plugin:activate" + ) + assert "foundry_sglang_plugin" in project["tool"]["setuptools"]["packages"] + assert not (PYTHON_ROOT / "foundry" / "integration" / "sglang" / "plugin.py").exists() diff --git a/tests/test_sglang_profile_validation.py b/tests/test_sglang_profile_validation.py new file mode 100644 index 00000000..48502323 --- /dev/null +++ b/tests/test_sglang_profile_validation.py @@ -0,0 +1,293 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the Foundry project +"""CPU-only contracts for the Foundry SGLang supported-profile guard.""" + +from __future__ import annotations + +import importlib.util +import sys +from enum import Enum +from pathlib import Path +from types import ModuleType, SimpleNamespace + +import pytest + +REPO_ROOT = Path(__file__).parents[1] +HOOKS_PATH = REPO_ROOT / "python" / "foundry" / "integration" / "sglang" / "hooks.py" + +SUPPORTED_ENV = { + "NCCL_CUMEM_ENABLE": "0", + "NCCL_NVLS_ENABLE": "0", + "PYTORCH_CUDA_ALLOC_CONF": "expandable_segments:False", + "SGLANG_ENABLE_POST_CAPTURE_KV_SIZING": "0", +} + + +def _install_module(monkeypatch, name: str) -> ModuleType: + module = ModuleType(name) + monkeypatch.setitem(sys.modules, name, module) + if "." in name: + parent_name, child_name = name.rsplit(".", 1) + parent = sys.modules.get(parent_name) + if parent is not None: + setattr(parent, child_name, module) + return module + + +def _load_hooks(monkeypatch): + monkeypatch.syspath_prepend(str(REPO_ROOT / "python")) + foundry = _install_module(monkeypatch, "foundry") + foundry.__path__ = [] + _install_module(monkeypatch, "foundry.integration").__path__ = [] + _install_module(monkeypatch, "foundry.integration.sglang").__path__ = [] + _install_module(monkeypatch, "foundry.integration.sglang.runtime") + + class ExtensionMode(Enum): + NONE = "none" + SAVE = "save" + LOAD = "load" + + config = _install_module(monkeypatch, "foundry.integration.sglang.config") + config.CUDAGraphExtensionMode = ExtensionMode + config.get_graph_extension_mode = lambda: ExtensionMode.SAVE + config.get_workspace_root = lambda: None + + spec = importlib.util.spec_from_file_location("sglang_profile_hooks_contract", HOOKS_PATH) + assert spec is not None and spec.loader is not None + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + return module + + +@pytest.fixture +def hooks_module(monkeypatch): + return _load_hooks(monkeypatch) + + +@pytest.fixture(autouse=True) +def supported_environment(monkeypatch): + for name, value in SUPPORTED_ENV.items(): + monkeypatch.setenv(name, value) + + +def _supported_args(): + return SimpleNamespace( + nnodes=1, + tp_size=2, + pp_size=1, + dp_size=1, + dcp_size=1, + ep_size=1, + enable_dp_attention=False, + moe_a2a_backend="none", + disaggregation_mode="null", + enable_pdmux=False, + disable_custom_all_reduce=True, + enable_mscclpp=False, + enable_torch_symm_mem=True, + enable_symm_mem=False, + enforce_disable_flashinfer_allreduce_fusion=True, + enable_nccl_nvls=False, + flashinfer_allreduce_fusion_backend=None, + disable_flashinfer_autotune=True, + enable_tokenizer_batch_encode=True, + enable_dynamic_batch_tokenizer=False, + disable_cuda_graph_padding=False, + decode_log_interval=1, + attention_backend="flashinfer", + prefill_attention_backend=None, + decode_attention_backend=None, + cuda_graph_config=SimpleNamespace( + decode=SimpleNamespace(backend="full", max_bs=256, bs=[256]), + prefill=SimpleNamespace(backend="disabled"), + ), + ) + + +def test_validate_profile_accepts_supported(hooks_module) -> None: + hooks_module._validate_profile(_supported_args()) + + +@pytest.mark.parametrize( + ("field", "value"), + [ + ("nnodes", 2), + ("tp_size", 3), + ("pp_size", 2), + ("dp_size", 2), + ("dcp_size", 2), + ("ep_size", 2), + ("enable_dp_attention", True), + ("moe_a2a_backend", "deepep"), + ("disaggregation_mode", "decode"), + ("enable_pdmux", True), + ("disable_custom_all_reduce", False), + ("enable_mscclpp", True), + ("enable_torch_symm_mem", False), + ("enable_symm_mem", True), + ("enforce_disable_flashinfer_allreduce_fusion", False), + ("enable_nccl_nvls", True), + ("flashinfer_allreduce_fusion_backend", "auto"), + ("disable_flashinfer_autotune", False), + ("enable_tokenizer_batch_encode", False), + ("enable_dynamic_batch_tokenizer", True), + # Runtime bs1/bs8 must pad to the single ShapeKey 256; padding disabled + # would replay a mismatched shape. + ("disable_cuda_graph_padding", True), + # decode_log_interval must resolve to 1 so every decode batch (including + # the short concurrent bs>1 request) emits a "#running-req" metrics line + # the acceptance harness parses for restored bs>1 replay evidence. The + # SGLang default of 40 would drop that evidence. + ("decode_log_interval", 40), + ("decode_log_interval", 2), + ("attention_backend", "fa3"), + ("prefill_attention_backend", "fa3"), + ("decode_attention_backend", "fa3"), + ], +) +def test_validate_profile_rejects_unsupported_field(hooks_module, field: str, value) -> None: + args = _supported_args() + setattr(args, field, value) + + with pytest.raises(RuntimeError, match=field): + hooks_module._validate_profile(args) + + +@pytest.mark.parametrize( + ("phase", "backend"), + [("decode", "breakable"), ("decode", "disabled"), ("prefill", "full")], +) +def test_validate_profile_rejects_unsupported_cuda_graph_backend( + hooks_module, phase: str, backend: str +) -> None: + args = _supported_args() + getattr(args.cuda_graph_config, phase).backend = backend + + with pytest.raises(RuntimeError, match=rf"cuda_graph_config\.{phase}\.backend"): + hooks_module._validate_profile(args) + + +@pytest.mark.parametrize( + ("attr", "value"), + [ + ("max_bs", 128), + ("max_bs", 512), + ("bs", []), + ("bs", [128]), + ("bs", [128, 256]), + ("bs", [256, 256]), + ], +) +def test_validate_profile_rejects_non_single_256_decode_shape( + hooks_module, attr: str, value +) -> None: + args = _supported_args() + setattr(args.cuda_graph_config.decode, attr, value) + + with pytest.raises(RuntimeError, match=rf"cuda_graph_config\.decode\.{attr}"): + hooks_module._validate_profile(args) + + +@pytest.mark.parametrize("attr", ["max_bs", "bs"]) +def test_validate_profile_fails_closed_when_decode_shape_field_missing( + hooks_module, attr: str +) -> None: + args = _supported_args() + delattr(args.cuda_graph_config.decode, attr) + + with pytest.raises(RuntimeError, match=rf"cuda_graph_config\.decode\.{attr}"): + hooks_module._validate_profile(args) + + +def test_validate_profile_fails_closed_when_padding_flag_missing(hooks_module) -> None: + args = _supported_args() + del args.disable_cuda_graph_padding + + with pytest.raises(RuntimeError, match="disable_cuda_graph_padding"): + hooks_module._validate_profile(args) + + +def test_validate_profile_fails_closed_when_decode_log_interval_missing( + hooks_module, +) -> None: + args = _supported_args() + del args.decode_log_interval + + with pytest.raises(RuntimeError, match="decode_log_interval"): + hooks_module._validate_profile(args) + + +@pytest.mark.parametrize( + "field", + ["enable_tokenizer_batch_encode", "enable_dynamic_batch_tokenizer"], +) +def test_validate_profile_fails_closed_when_tokenizer_batch_mode_missing( + hooks_module, + field: str, +) -> None: + args = _supported_args() + delattr(args, field) + + with pytest.raises(RuntimeError, match=field): + hooks_module._validate_profile(args) + + +def test_validate_profile_error_identifies_torch_symmetric_memory(hooks_module) -> None: + args = _supported_args() + args.enable_torch_symm_mem = False + + with pytest.raises(RuntimeError, match="torch-symmetric-memory") as exc_info: + hooks_module._validate_profile(args) + assert "plain-NCCL" not in str(exc_info.value) + + +@pytest.mark.parametrize( + ("name", "value"), + [ + ("NCCL_CUMEM_ENABLE", "1"), + ("NCCL_NVLS_ENABLE", "1"), + ("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True"), + ("SGLANG_ENABLE_POST_CAPTURE_KV_SIZING", "1"), + ], +) +def test_validate_profile_rejects_unsupported_environment( + hooks_module, monkeypatch, name: str, value: str +) -> None: + monkeypatch.setenv(name, value) + + with pytest.raises(RuntimeError, match=name): + hooks_module._validate_profile(_supported_args()) + + +def test_validate_profile_reports_every_violation(hooks_module, monkeypatch) -> None: + args = _supported_args() + args.tp_size = 4 + args.enable_dp_attention = True + args.cuda_graph_config.decode.backend = "breakable" + args.cuda_graph_config.prefill.backend = "full" + monkeypatch.setenv("NCCL_CUMEM_ENABLE", "1") + monkeypatch.setenv("SGLANG_ENABLE_POST_CAPTURE_KV_SIZING", "1") + + with pytest.raises(RuntimeError) as exc_info: + hooks_module._validate_profile(args) + + message = str(exc_info.value) + for name in ( + "tp_size", + "enable_dp_attention", + "cuda_graph_config.decode.backend", + "cuda_graph_config.prefill.backend", + "NCCL_CUMEM_ENABLE", + "SGLANG_ENABLE_POST_CAPTURE_KV_SIZING", + ): + assert name in message + + +def test_validate_profile_fails_closed_when_resolved_config_is_missing( + hooks_module, +) -> None: + args = _supported_args() + del args.cuda_graph_config + + with pytest.raises(RuntimeError, match="cuda_graph_config"): + hooks_module._validate_profile(args) diff --git a/tests/test_sglang_runtime_state.py b/tests/test_sglang_runtime_state.py new file mode 100644 index 00000000..0511004b --- /dev/null +++ b/tests/test_sglang_runtime_state.py @@ -0,0 +1,353 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the Foundry project +"""CPU-only contracts for SGLang shared warmup-state persistence.""" + +from __future__ import annotations + +import builtins +import importlib.util +import json +import sys +import threading +from dataclasses import asdict +from enum import Enum +from pathlib import Path +from types import ModuleType, SimpleNamespace + +import pytest + +REPO_ROOT = Path(__file__).parents[1] +RUNTIME_PATH = REPO_ROOT / "python" / "foundry" / "integration" / "sglang" / "runtime.py" + + +def _install_module(monkeypatch, name: str) -> ModuleType: + module = ModuleType(name) + monkeypatch.setitem(sys.modules, name, module) + if "." in name: + parent_name, child_name = name.rsplit(".", 1) + parent = sys.modules.get(parent_name) + if parent is not None: + setattr(parent, child_name, module) + return module + + +def _load_runtime(monkeypatch, tmp_path): + foundry = _install_module(monkeypatch, "foundry") + foundry.__path__ = [] + + current_offset = [4096] + ops = _install_module(monkeypatch, "foundry.ops") + ops.get_current_alloc_offset = lambda: current_offset[0] + ops.preallocate_region = lambda _size: True + + allocation_region = _install_module(monkeypatch, "foundry.allocation_region") + allocation_region.parse_size = lambda value: value + + _install_module(monkeypatch, "foundry.integration").__path__ = [] + _install_module(monkeypatch, "foundry.integration.sglang").__path__ = [] + + class ExtensionMode(Enum): + NONE = "none" + SAVE = "save" + LOAD = "load" + + workspace_dir = tmp_path / "rank_0" + workspace_dir.mkdir() + cfg = SimpleNamespace( + mode=ExtensionMode.SAVE, + workspace_dir=str(workspace_dir), + workspace_root=str(tmp_path), + ) + config = _install_module(monkeypatch, "foundry.integration.sglang.config") + config.CUDAGraphExtensionMode = ExtensionMode + config.compute_workspace_rank = lambda *_args, **_kwargs: 0 + config.get_config = lambda: cfg + config.get_graph_extension_mode = lambda: cfg.mode + config.get_hook_library_path = lambda: None + config.get_nvshmem_host_path = lambda: None + + torch = _install_module(monkeypatch, "torch") + torch.cuda = SimpleNamespace() + + spec = importlib.util.spec_from_file_location("sglang_runtime_state_contract", RUNTIME_PATH) + assert spec is not None and spec.loader is not None + runtime = importlib.util.module_from_spec(spec) + monkeypatch.setitem(sys.modules, spec.name, runtime) + spec.loader.exec_module(runtime) + return SimpleNamespace( + cfg=cfg, + current_offset=current_offset, + modes=ExtensionMode, + runtime=runtime, + ) + + +@pytest.fixture +def runtime_context(monkeypatch, tmp_path): + return _load_runtime(monkeypatch, tmp_path) + + +def test_initialized_nonzero_rank_never_opens_shared_warmup_state( + runtime_context, + monkeypatch, +) -> None: + runtime = runtime_context.runtime + runtime._state = runtime.CUDAGraphExtensionState(rank=1) + shared_path = Path(runtime_context.cfg.workspace_root) / "warmup_state.json" + opened_shared_paths = [] + real_open = builtins.open + + def guarded_open(path, *args, **kwargs): + if Path(path) == shared_path: + opened_shared_paths.append(Path(path)) + raise AssertionError("nonzero rank opened shared warmup state") + return real_open(path, *args, **kwargs) + + monkeypatch.setattr(builtins, "open", guarded_open) + + runtime.save_warmup_state(runtime.WarmupState(memory_pool_config={"rank": 1})) + + assert opened_shared_paths == [] + assert not shared_path.exists() + + +def test_nonzero_capture_writes_rank_offset_without_reading_shared_state( + runtime_context, + monkeypatch, +) -> None: + runtime = runtime_context.runtime + runtime._state = runtime.CUDAGraphExtensionState(rank=1) + rank_dir = Path(runtime_context.cfg.workspace_root) / "rank_1" + rank_dir.mkdir() + runtime_context.cfg.workspace_dir = str(rank_dir) + shared_path = Path(runtime_context.cfg.workspace_root) / "warmup_state.json" + shared_path.write_text(json.dumps(asdict(runtime.WarmupState(final_alloc_offset=1)))) + + monkeypatch.setattr( + runtime, + "load_warmup_state", + lambda: pytest.fail("nonzero rank read shared warmup state"), + ) + monkeypatch.setattr( + runtime, + "save_warmup_state", + lambda _state: pytest.fail("nonzero rank updated shared warmup state"), + ) + + assert runtime.capture_final_alloc_offset() == 4096 + assert json.loads((rank_dir / "final_alloc_offset.json").read_text()) == { + "final_alloc_offset": 4096 + } + + +def test_rank_zero_atomic_write_hides_partial_temp_content_from_readers( + runtime_context, + monkeypatch, +) -> None: + runtime = runtime_context.runtime + runtime._state = runtime.CUDAGraphExtensionState(rank=0) + shared_path = Path(runtime_context.cfg.workspace_root) / "warmup_state.json" + old_state = runtime.WarmupState(memory_pool_config={"generation": "old"}) + new_state = runtime.WarmupState(memory_pool_config={"generation": "new"}) + shared_path.write_text(json.dumps(asdict(old_state))) + + write_started = threading.Event() + finish_write = threading.Event() + writer_errors = [] + real_dump = runtime.json.dump + + def paused_dump(value, stream, *args, **kwargs): + stream.write("{") + stream.flush() + write_started.set() + if not finish_write.wait(timeout=2): + raise TimeoutError("reader did not release paused writer") + stream.seek(0) + stream.truncate() + return real_dump(value, stream, *args, **kwargs) + + monkeypatch.setattr(runtime.json, "dump", paused_dump) + + def write_state(): + try: + runtime.save_warmup_state(new_state) + except Exception as exc: + writer_errors.append(exc) + + writer = threading.Thread(target=write_state) + writer.start() + assert write_started.wait(timeout=2) + try: + observed_during_write = json.loads(shared_path.read_text()) + finally: + finish_write.set() + writer.join(timeout=2) + + assert not writer.is_alive() + assert writer_errors == [] + assert observed_during_write == asdict(old_state) + assert json.loads(shared_path.read_text()) == asdict(new_state) + + +def test_rank_zero_atomic_write_cleans_temp_and_preserves_state_on_error( + runtime_context, + monkeypatch, +) -> None: + runtime = runtime_context.runtime + runtime._state = runtime.CUDAGraphExtensionState(rank=0) + workspace_root = Path(runtime_context.cfg.workspace_root) + shared_path = workspace_root / "warmup_state.json" + old_state = runtime.WarmupState(memory_pool_config={"generation": "old"}) + shared_path.write_text(json.dumps(asdict(old_state))) + + def fail_dump(*_args, **_kwargs): + raise OSError("simulated write failure") + + monkeypatch.setattr(runtime.json, "dump", fail_dump) + + with pytest.raises(OSError, match="simulated write failure"): + runtime.save_warmup_state(runtime.WarmupState(memory_pool_config={"generation": "new"})) + + assert json.loads(shared_path.read_text()) == asdict(old_state) + assert list(workspace_root.glob(".warmup_state.*.tmp")) == [] + + +@pytest.mark.parametrize("json_root", [[], "state", 1, None, True]) +def test_load_rejects_non_object_json_root(runtime_context, json_root) -> None: + runtime = runtime_context.runtime + shared_path = Path(runtime_context.cfg.workspace_root) / "warmup_state.json" + shared_path.write_text(json.dumps(json_root)) + + with pytest.raises(RuntimeError, match="incompatible warmup state") as exc_info: + runtime.load_warmup_state() + + assert str(shared_path) in str(exc_info.value) + + +def test_load_wraps_malformed_json_with_state_path(runtime_context) -> None: + runtime = runtime_context.runtime + shared_path = Path(runtime_context.cfg.workspace_root) / "warmup_state.json" + shared_path.write_text('{"memory_pool_config":') + + with pytest.raises(RuntimeError, match="invalid JSON") as exc_info: + runtime.load_warmup_state() + + assert str(shared_path) in str(exc_info.value) + + +def test_load_preserves_missing_file_error(runtime_context) -> None: + runtime = runtime_context.runtime + shared_path = Path(runtime_context.cfg.workspace_root) / "warmup_state.json" + + with pytest.raises(RuntimeError, match="warmup state file not found") as exc_info: + runtime.load_warmup_state() + + assert str(shared_path) in str(exc_info.value) + assert "invalid JSON" not in str(exc_info.value) + + +def test_load_preallocation_failure_raises_with_offsets( + runtime_context, + monkeypatch, +) -> None: + runtime_context.cfg.mode = runtime_context.modes.LOAD + final_offset = 16384 + current_offset = runtime_context.current_offset[0] + state_path = Path(runtime_context.cfg.workspace_dir) / "final_alloc_offset.json" + state_path.write_text(json.dumps({"final_alloc_offset": final_offset})) + attempted = [] + + def fail_preallocation(size: int) -> bool: + attempted.append(size) + return False + + monkeypatch.setattr(runtime_context.runtime.cge, "preallocate_region", fail_preallocation) + + with pytest.raises(RuntimeError, match="SGLang LOAD preallocation failed") as exc_info: + runtime_context.runtime.preallocate_for_load_mode() + + assert attempted == [final_offset - current_offset] + message = str(exc_info.value) + assert f"remaining={final_offset - current_offset}" in message + assert f"current_offset={current_offset}" in message + assert f"final_alloc_offset={final_offset}" in message + + +def test_load_preallocation_success_is_accepted(runtime_context, monkeypatch) -> None: + runtime_context.cfg.mode = runtime_context.modes.LOAD + final_offset = 8192 + state_path = Path(runtime_context.cfg.workspace_dir) / "final_alloc_offset.json" + state_path.write_text(json.dumps({"final_alloc_offset": final_offset})) + attempted = [] + monkeypatch.setattr( + runtime_context.runtime.cge, + "preallocate_region", + lambda size: attempted.append(size) or True, + ) + + runtime_context.runtime.preallocate_for_load_mode() + + assert attempted == [final_offset - runtime_context.current_offset[0]] + + +def test_load_preallocation_requires_positive_final_after_shared_fallback( + runtime_context, + monkeypatch, +) -> None: + runtime_context.cfg.mode = runtime_context.modes.LOAD + monkeypatch.setattr( + runtime_context.runtime, + "load_warmup_state", + lambda: runtime_context.runtime.WarmupState(final_alloc_offset=0), + ) + + with pytest.raises(RuntimeError, match="final_alloc_offset must be positive") as exc_info: + runtime_context.runtime.preallocate_for_load_mode() + + message = str(exc_info.value) + assert "final_alloc_offset=0" in message + assert f"current_offset={runtime_context.current_offset[0]}" in message + + +def test_load_preallocation_rejects_current_offset_beyond_final( + runtime_context, + monkeypatch, +) -> None: + runtime_context.cfg.mode = runtime_context.modes.LOAD + final_offset = runtime_context.current_offset[0] - 1 + state_path = Path(runtime_context.cfg.workspace_dir) / "final_alloc_offset.json" + state_path.write_text(json.dumps({"final_alloc_offset": final_offset})) + attempted = [] + monkeypatch.setattr( + runtime_context.runtime.cge, + "preallocate_region", + lambda size: attempted.append(size) or True, + ) + + with pytest.raises(RuntimeError, match="current allocation offset exceeds") as exc_info: + runtime_context.runtime.preallocate_for_load_mode() + + assert attempted == [] + message = str(exc_info.value) + assert f"current_offset={runtime_context.current_offset[0]}" in message + assert f"final_alloc_offset={final_offset}" in message + + +def test_load_preallocation_equal_offsets_is_valid_noop( + runtime_context, + monkeypatch, +) -> None: + runtime_context.cfg.mode = runtime_context.modes.LOAD + final_offset = runtime_context.current_offset[0] + state_path = Path(runtime_context.cfg.workspace_dir) / "final_alloc_offset.json" + state_path.write_text(json.dumps({"final_alloc_offset": final_offset})) + attempted = [] + monkeypatch.setattr( + runtime_context.runtime.cge, + "preallocate_region", + lambda size: attempted.append(size) or True, + ) + + runtime_context.runtime.preallocate_for_load_mode() + + assert attempted == [] diff --git a/tests/test_sglang_shape_key_archive.py b/tests/test_sglang_shape_key_archive.py new file mode 100644 index 00000000..d78a722f --- /dev/null +++ b/tests/test_sglang_shape_key_archive.py @@ -0,0 +1,156 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the Foundry project +"""CPU-only contracts for ShapeKey-based SGLang graph-archive filenames. + +The pinned upstream (a25164b) keys every captured CUDA graph by +``sglang.srt.model_executor.runner.shape_key.ShapeKey(size, stream_idx, +variant_label)``. Foundry's dense TP=2 full-decode profile only supports the +single-stream, no-LoRA variant, so the archive round-trips ``size`` and fails +closed on any populated ``stream_idx`` / ``variant_label``. +""" + +from __future__ import annotations + +import importlib.util +import sys +from dataclasses import dataclass +from pathlib import Path +from types import ModuleType, SimpleNamespace + +import pytest + +REPO_ROOT = Path(__file__).parents[1] +GRAPH_OPS_PATH = REPO_ROOT / "python" / "foundry" / "integration" / "sglang" / "graph_ops.py" + + +@dataclass(frozen=True) +class ShapeKeyLike: + """Mirror of the pinned upstream ShapeKey dataclass.""" + + size: int + stream_idx: int | None = None + variant_label: str | None = None + + +def _install_module(monkeypatch, name: str) -> ModuleType: + module = ModuleType(name) + monkeypatch.setitem(sys.modules, name, module) + if "." in name: + parent_name, child_name = name.rsplit(".", 1) + parent = sys.modules.get(parent_name) + if parent is not None: + setattr(parent, child_name, module) + return module + + +def _load_graph_ops(monkeypatch): + torch = _install_module(monkeypatch, "torch") + torch.Tensor = object + + foundry = _install_module(monkeypatch, "foundry") + foundry.__path__ = [] + _install_module(monkeypatch, "foundry.ops") + graph = _install_module(monkeypatch, "foundry.graph") + graph.CUDAGraph = object + graph.graph = lambda *args, **kwargs: None + + _install_module(monkeypatch, "foundry.integration").__path__ = [] + _install_module(monkeypatch, "foundry.integration.sglang").__path__ = [] + config = _install_module(monkeypatch, "foundry.integration.sglang.config") + config.CUDAGraphExtensionMode = SimpleNamespace(SAVE="save", LOAD="load") + config.get_config = lambda: None + config.get_graph_extension_mode = lambda: None + runtime = _install_module(monkeypatch, "foundry.integration.sglang.runtime") + runtime.get_state = lambda: None + + # Lazy import target for _shape_key_from_filename. + _install_module(monkeypatch, "sglang").__path__ = [] + _install_module(monkeypatch, "sglang.srt").__path__ = [] + _install_module(monkeypatch, "sglang.srt.model_executor").__path__ = [] + _install_module(monkeypatch, "sglang.srt.model_executor.runner").__path__ = [] + shape_key = _install_module(monkeypatch, "sglang.srt.model_executor.runner.shape_key") + shape_key.ShapeKey = ShapeKeyLike + + spec = importlib.util.spec_from_file_location( + "sglang_shape_key_archive_contract", GRAPH_OPS_PATH + ) + assert spec is not None and spec.loader is not None + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + module.ShapeKeyLike = ShapeKeyLike + return module + + +@pytest.fixture +def graph_ops(monkeypatch): + return _load_graph_ops(monkeypatch) + + +def test_shape_key_filename_roundtrip(graph_ops): + key = graph_ops.ShapeKeyLike(size=8, stream_idx=None, variant_label=None) + name = graph_ops._graph_filename(3, key) + parsed = graph_ops._shape_key_from_filename(name) + assert parsed.size == 8 and parsed.stream_idx is None and parsed.variant_label is None + + +def test_graph_filename_encodes_index_and_size(graph_ops): + key = graph_ops.ShapeKeyLike(size=16) + assert graph_ops._graph_filename(5, key) == "graph_5_FULL_t16_r16_UX_pcN.json" + + +def test_shape_key_rejects_unsupported_variant(graph_ops): + key = graph_ops.ShapeKeyLike(size=8, stream_idx=1, variant_label="lora") + with pytest.raises(ValueError): + graph_ops._graph_filename(3, key) + + +def test_shape_key_rejects_stream_idx_only(graph_ops): + key = graph_ops.ShapeKeyLike(size=8, stream_idx=0) + with pytest.raises(ValueError): + graph_ops._graph_filename(3, key) + + +def test_shape_key_rejects_variant_label_only(graph_ops): + key = graph_ops.ShapeKeyLike(size=8, variant_label="nolora") + with pytest.raises(ValueError): + graph_ops._graph_filename(3, key) + + +def test_shape_key_from_filename_rejects_non_matching(graph_ops): + with pytest.raises(ValueError): + graph_ops._shape_key_from_filename("not_a_graph_file.json") + + +# --- Finding 3: strict full-field validation -------------------------------- + + +def test_require_dense_shape_key_rejects_raw_int(graph_ops): + # A bare int is not a ShapeKey: it exposes none of the current fields, so it + # must fail closed rather than being treated as a size. + with pytest.raises(ValueError): + graph_ops._graph_filename(3, 8) + + +@dataclass(frozen=True) +class _MissingVariantLabel: + size: int + stream_idx: int | None = None + + +@dataclass(frozen=True) +class _MissingStreamIdx: + size: int + variant_label: str | None = None + + +def test_require_dense_shape_key_rejects_incomplete_object(graph_ops): + for incomplete in (_MissingVariantLabel(size=8), _MissingStreamIdx(size=8)): + with pytest.raises(ValueError): + graph_ops._graph_filename(3, incomplete) + + +def test_require_dense_shape_key_rejects_bool_size(graph_ops): + # bool is an int subclass; a boolean size is not a valid capture size. + key = graph_ops.ShapeKeyLike(size=True, stream_idx=None, variant_label=None) + with pytest.raises(ValueError): + graph_ops._graph_filename(3, key) diff --git a/tests/test_sglang_target_resolution.py b/tests/test_sglang_target_resolution.py new file mode 100644 index 00000000..a10e3f62 --- /dev/null +++ b/tests/test_sglang_target_resolution.py @@ -0,0 +1,361 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the Foundry project +"""CPU contracts for installed SGLang discovery and pinned-main hook targets.""" + +from __future__ import annotations + +import ast +import functools +import importlib.util +import os +import subprocess +import sys +from enum import Enum +from pathlib import Path +from types import ModuleType + +import tomllib + +REPO_ROOT = Path(__file__).parents[1] +PYTHON_ROOT = REPO_ROOT / "python" +HOOKS_PATH = PYTHON_ROOT / "foundry" / "integration" / "sglang" / "hooks.py" +INIT_PATH = PYTHON_ROOT / "foundry" / "integration" / "sglang" / "__init__.py" +TARGETS_PATH = PYTHON_ROOT / "foundry_sglang_plugin" / "targets.py" +MODAL_CONTRACT_PATH = REPO_ROOT / "tests" / "modal_sglang_target_resolution.py" + +EXPECTED_ENTRY_POINT = "foundry_sglang_plugin:activate" +EXPECTED_HOOK_TARGETS = { + "sglang.srt.server_args.ServerArgs.__post_init__", + "sglang.srt.entrypoints.engine.Engine._launch_scheduler_processes", + "sglang.srt.model_executor.model_runner.ModelRunner.init_torch_distributed", + "sglang.srt.model_executor.model_runner.ModelRunner.alloc_memory_pool", + ("sglang.srt.mem_cache.kv_cache_configurator.KVCacheConfigurator._resolve_memory_pool_config"), + "sglang.srt.model_executor.runner_backend.full_cuda_graph_backend.FullCudaGraphBackend", + ("sglang.srt.model_executor.runner.decode_cuda_graph_runner.DecodeCudaGraphRunner.capture"), + ( + "sglang.srt.distributed.device_communicators.torch_symm_mem." + "TorchSymmMemCommunicator.__init__" + ), + ("sglang.srt.distributed.device_communicators.triton_symm_mem_ag.MultimemAllGatherer.__call__"), +} + + +def _install_module(monkeypatch, name: str) -> ModuleType: + module = ModuleType(name) + monkeypatch.setitem(sys.modules, name, module) + if "." in name: + parent_name, child_name = name.rsplit(".", 1) + parent = sys.modules.get(parent_name) + if parent is not None: + setattr(parent, child_name, module) + return module + + +def _load_module(path: Path, name: str): + spec = importlib.util.spec_from_file_location(name, path) + assert spec is not None and spec.loader is not None + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + return module + + +def _load_hooks(monkeypatch): + targets_package = _install_module(monkeypatch, "foundry_sglang_plugin") + targets_package.__path__ = [] + targets_module = _load_module(TARGETS_PATH, "foundry_sglang_plugin.targets") + monkeypatch.setitem(sys.modules, "foundry_sglang_plugin.targets", targets_module) + targets_package.targets = targets_module + + foundry = _install_module(monkeypatch, "foundry") + foundry.__path__ = [] + _install_module(monkeypatch, "foundry.integration").__path__ = [] + _install_module(monkeypatch, "foundry.integration.sglang").__path__ = [] + + backend = _install_module(monkeypatch, "foundry.integration.sglang.backend") + backend.build_foundry_full_backend_class = lambda: type("FoundryFullCudaGraphBackend", (), {}) + + runtime = _install_module(monkeypatch, "foundry.integration.sglang.runtime") + runtime.setup_ld_preload_env = lambda: None + runtime.setup_graph_extension = lambda *args, **kwargs: None + runtime.log_alloc_offset = lambda *args: None + runtime.skip_to_scratch_boundary = lambda: None + + class ExtensionMode(Enum): + NONE = "none" + SAVE = "save" + LOAD = "load" + + config = _install_module(monkeypatch, "foundry.integration.sglang.config") + config.CUDAGraphExtensionMode = ExtensionMode + config.get_graph_extension_mode = lambda: ExtensionMode.SAVE + config.get_workspace_root = lambda: None + + _install_module(monkeypatch, "sglang").__path__ = [] + _install_module(monkeypatch, "sglang.srt").__path__ = [] + _install_module(monkeypatch, "sglang.srt.plugins").__path__ = [] + registry_module = _install_module(monkeypatch, "sglang.srt.plugins.hook_registry") + + class HookType(Enum): + BEFORE = "before" + AFTER = "after" + AROUND = "around" + REPLACE = "replace" + + class HookRegistry: + calls = [] + + @classmethod + def register(cls, target, hook, hook_type=HookType.AFTER): + cls.calls.append((target, hook_type)) + + registry_module.HookRegistry = HookRegistry + registry_module.HookType = HookType + return _load_module(HOOKS_PATH, "sglang_target_resolution_hooks"), HookRegistry + + +def test_entry_point_declaration_is_exact() -> None: + project = tomllib.loads((REPO_ROOT / "pyproject.toml").read_text()) + + entry_points = project["project"]["entry-points"]["sglang.srt.plugins"] + assert entry_points == {"foundry": EXPECTED_ENTRY_POINT} + + +def test_registered_targets_match_reviewed_single_source(monkeypatch) -> None: + targets = _load_module(TARGETS_PATH, "foundry_sglang_targets_contract") + hooks, registry = _load_hooks(monkeypatch) + + hooks.register_hooks() + + registered = {target for target, _hook_type in registry.calls} + assert set(targets.HOOK_TARGETS) == EXPECTED_HOOK_TARGETS + assert registered == EXPECTED_HOOK_TARGETS + assert len(registry.calls) == len(EXPECTED_HOOK_TARGETS) + + +def test_resolved_identity_normalizes_classmethods_and_detects_wrappers() -> None: + targets = _load_module(TARGETS_PATH, "foundry_sglang_identity_contract") + + class Engine: + @classmethod + def launch(cls): + return cls + + def ordinary_method(self): + return self + + first = Engine.launch + second = Engine.launch + assert first is not second + assert targets.normalize_resolved_identity(first) is targets.normalize_resolved_identity(second) + assert targets.normalize_resolved_identity(Engine) is Engine + assert targets.normalize_resolved_identity(Engine.ordinary_method) is Engine.ordinary_method + + original = targets.normalize_resolved_identity(first) + + @functools.wraps(original) + def wrapped(*args, **kwargs): + return original(*args, **kwargs) + + Engine.launch = classmethod(wrapped) + replacement = targets.normalize_resolved_identity(Engine.launch) + + assert replacement is not original + assert replacement.__wrapped__ is original + + +def test_legacy_install_shim_and_removed_targets_are_absent() -> None: + integration_source = HOOKS_PATH.read_text() + INIT_PATH.read_text() + + assert "install_hooks" not in integration_source + assert "model_runner_kv_cache_mixin" not in integration_source + assert "sglang.srt.model_executor.cuda_graph_runner.CudaGraphRunner" not in integration_source + + +def test_modal_contract_reuses_pinned_clean_checkout_image() -> None: + source = MODAL_CONTRACT_PATH.read_text() + + for required in ( + "import modal_sglang_tp as shared_harness", + "FOUNDRY_REF", + "SGLANG_BASE", + "MODEL", + "MODEL_REVISION", + "entry_points(group=GENERAL_PLUGINS_GROUP)", + "pkgutil.resolve_name", + "normalize_resolved_identity(pkgutil.resolve_name(target))", + "HookRegistry._hooks", + "HookRegistry._patched", + '"foundry_ref": FOUNDRY_REF', + '"sglang_base": SGLANG_BASE', + '"model": MODEL', + '"model_revision": MODEL_REVISION', + ): + assert required in source + + documentation = ast.get_docstring(ast.parse(source)) + assert documentation is not None + normalized_documentation = " ".join(documentation.split()) + for required in ( + "inactive plugin check is CPU-only", + "configured activation check uses one H100", + "local orchestration process never initializes CUDA", + ): + assert required in normalized_documentation + + +def test_modal_remote_checks_split_cuda_driver_resources() -> None: + source = MODAL_CONTRACT_PATH.read_text() + + assert "@app.function(image=image, timeout=600)\ndef check_inactive_plugin() -> dict:" in source + assert ( + '@app.function(image=image, gpu="H100", timeout=600)\n' + "def check_configured_plugin() -> dict:" in source + ) + assert source.count('gpu="H100"') == 1 + + +class _ImageRecorder: + """Fake Modal image recording ``add_local_python_source`` calls.""" + + def __init__(self) -> None: + self.local_python_sources: list[tuple[tuple[str, ...], bool]] = [] + + def add_local_python_source(self, *modules: str, copy: bool = False, **_kwargs): + self.local_python_sources.append((tuple(modules), copy)) + return self + + def __getattr__(self, _name: str): + def _chain(*_args, **_kwargs): + return self + + return _chain + + +def _fake_modal(): + modal_module = ModuleType("modal") + + class _ImageFactory: + @staticmethod + def from_registry(*_args, **_kwargs): + return _ImageRecorder() + + @staticmethod + def debian_slim(*_args, **_kwargs): + return _ImageRecorder() + + class _App: + def __init__(self, *_args, **_kwargs): + pass + + def function(self, *_args, **_kwargs): + return lambda fn: fn + + def local_entrypoint(self, *_args, **_kwargs): + return lambda fn: fn + + class _Volume: + @staticmethod + def from_name(*_args, **_kwargs): + return object() + + modal_module.Image = _ImageFactory + modal_module.App = _App + modal_module.Volume = _Volume + modal_module.is_local = lambda: True + return modal_module + + +def test_reused_image_inherits_sibling_module_mounts() -> None: + import runpy + + full_sha = "0123456789abcdef0123456789abcdef01234567" + saved_modules = { + name: sys.modules.get(name) + for name in ( + "modal", + "modal_sglang_tp", + "modal_sglang_target_resolution", + "sglang_tp_validation", + "tp_validation_helpers", + ) + } + saved_path = list(sys.path) + saved_env = os.environ.copy() + try: + sys.modules["modal"] = _fake_modal() + sys.path.insert(0, str(MODAL_CONTRACT_PATH.parent)) + os.environ.update( + { + "FOUNDRY_REF": full_sha, + "SGLANG_BASE": full_sha, + "SGL_MODEL_REVISION": full_sha, + "TP_SIZE": "2", + } + ) + namespace = runpy.run_path( + str(MODAL_CONTRACT_PATH), run_name="_target_resolution_image_contract" + ) + image = namespace["image"] + assert isinstance(image, _ImageRecorder), type(image) + attached = {module for modules, _copy in image.local_python_sources for module in modules} + assert "sglang_tp_validation" in attached + assert "tp_validation_helpers" in attached + finally: + sys.path[:] = saved_path + os.environ.clear() + os.environ.update(saved_env) + for name, module in saved_modules.items(): + if module is None: + sys.modules.pop(name, None) + else: + sys.modules[name] = module + + +def test_modal_remote_import_does_not_require_local_harness_modules() -> None: + code = """ +import importlib.abc +import runpy +import sys +from types import ModuleType + +class BlockLocalHarnessImports(importlib.abc.MetaPathFinder): + def find_spec(self, fullname, path=None, target=None): + if fullname in {"modal_sglang_tp", "sglang_tp_validation"}: + raise ImportError(f"local harness module is unavailable remotely: {fullname}") + return None + +class FakeApp: + def __init__(self, *args, **kwargs): + pass + + def function(self, **kwargs): + return lambda fn: fn + + def local_entrypoint(self): + return lambda fn: fn + +modal = ModuleType("modal") +modal.App = FakeApp +modal.is_local = lambda: False +sys.modules["modal"] = modal +sys.meta_path.insert(0, BlockLocalHarnessImports()) +runpy.run_path(sys.argv[1], run_name="modal_remote_import_contract") +""" + env = os.environ.copy() + env.update( + { + "FOUNDRY_REF": "0123456789abcdef0123456789abcdef01234567", + "SGLANG_BASE": "123456789abcdef0123456789abcdef012345678", + "SGL_MODEL_REVISION": "23456789abcdef0123456789abcdef0123456789", + } + ) + + result = subprocess.run( + [sys.executable, "-c", code, str(MODAL_CONTRACT_PATH)], + cwd=REPO_ROOT, + env=env, + capture_output=True, + text=True, + ) + + assert result.returncode == 0, result.stderr diff --git a/tests/test_sglang_tp_recipe.py b/tests/test_sglang_tp_recipe.py new file mode 100644 index 00000000..5288a17b --- /dev/null +++ b/tests/test_sglang_tp_recipe.py @@ -0,0 +1,148 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the Foundry project +"""CPU-only contract tests for the SGLang tensor-parallel recipe.""" + +from __future__ import annotations + +import json +import os +import subprocess +from pathlib import Path + +import pytest +import tomllib + +EXPECTED_CUDA_GRAPH_CONFIG = { + "decode": {"backend": "full", "max_bs": 256, "bs": [256]}, + "prefill": {"backend": "disabled"}, +} + +REPO_ROOT = Path(__file__).parents[1] +RECIPE_DIR = REPO_ROOT / "recipe" / "experimental" +SERVE_SCRIPT = RECIPE_DIR / "serve_qwen3-8b_sglang_tp.sh" + + +def _run_recipe(tmp_path: Path, mode: str | None) -> tuple[dict[str, str], list[str]]: + bin_dir = tmp_path / "bin" + bin_dir.mkdir() + fake_sglang = bin_dir / "sglang" + fake_sglang.write_text( + '#!/usr/bin/env bash\nenv > "$CAPTURE_ENV"\nprintf "%s\\n" "$@" > "$CAPTURE_ARGS"\n' + ) + fake_sglang.chmod(0o755) + + env_path = tmp_path / "env" + args_path = tmp_path / "args" + env = os.environ.copy() + env.update( + { + "CAPTURE_ENV": str(env_path), + "CAPTURE_ARGS": str(args_path), + "FOUNDRY_GRAPH_EXTENSION_CONFIG_PATH": "/stale/foundry-config.toml", + "PATH": f"{bin_dir}:{env['PATH']}", + } + ) + + command = ["bash", str(SERVE_SCRIPT), "2"] + if mode is not None: + command.append(mode) + subprocess.run(command, check=True, env=env, capture_output=True, text=True) + + captured_env = dict(line.split("=", 1) for line in env_path.read_text().splitlines()) + return captured_env, args_path.read_text().splitlines() + + +@pytest.mark.parametrize("mode", [None, "--save", "--load"]) +def test_tp_recipe_uses_same_nccl_transport_for_every_phase( + tmp_path: Path, + mode: str | None, +) -> None: + env, args = _run_recipe(tmp_path, mode) + + assert env["PYTORCH_CUDA_ALLOC_CONF"] == "expandable_segments:False" + assert env["NCCL_CUMEM_ENABLE"] == "0" + assert env["NCCL_NVLS_ENABLE"] == "0" + assert env["SGLANG_ENABLE_POST_CAPTURE_KV_SIZING"] == "0" + assert env["SGLANG_PLUGINS"] == "foundry" + assert env["SGL_MODEL_REVISION"] == "b968826d9c46dd6066d109eabc6255188de91218" + cuda_graph_config = args[args.index("--cuda-graph-config") + 1] + assert json.loads(cuda_graph_config) == EXPECTED_CUDA_GRAPH_CONFIG + assert "--disable-piecewise-cuda-graph" not in args + # Collective-selection flags must be present in baseline, SAVE, and LOAD so + # torch symmetric memory wins over custom/FlashInfer-fused all-reduce. + assert "--disable-custom-all-reduce" in args + assert "--enable-torch-symm-mem" in args + assert "--enforce-disable-flashinfer-allreduce-fusion" in args + assert args.count("--enable-tokenizer-batch-encode") == 1 + assert "--enable-dynamic-batch-tokenizer" not in args + assert args[args.index("--random-seed") + 1] == "42" + if mode is None: + assert "FOUNDRY_GRAPH_EXTENSION_CONFIG_PATH" not in env + + +@pytest.mark.parametrize("mode", [None, "--save", "--load"]) +def test_tp_recipe_logs_every_decode_batch_in_every_phase( + tmp_path: Path, + mode: str | None, +) -> None: + # SGLang's decode_log_interval defaults to 40, so a short concurrent bs>1 + # decode would emit no "#running-req" metrics line and the acceptance + # bs>1-replay evidence would silently disappear. Pin it to 1 in every phase + # so baseline, SAVE, and LOAD log every decode batch. + _, args = _run_recipe(tmp_path, mode) + + assert args.count("--decode-log-interval") == 1 + assert args[args.index("--decode-log-interval") + 1] == "1" + + +@pytest.mark.parametrize( + "arguments", + [ + [], + ["1"], + ["3"], + ["2", "--invalid"], + ["2", "--save", "unexpected"], + ["2", "--load", "unexpected"], + ], +) +def test_tp_recipe_rejects_unvalidated_profile(arguments: list[str]) -> None: + result = subprocess.run( + ["bash", str(SERVE_SCRIPT), *arguments], + capture_output=True, + text=True, + ) + + assert result.returncode != 0 + assert "TP=2" in result.stderr + + +@pytest.mark.parametrize( + ("mode", "config_name"), + [ + ("--save", "sglang_foundry_tp_save.toml"), + ("--load", "sglang_foundry_tp_load.toml"), + ], +) +def test_tp_recipe_activates_foundry_via_env( + tmp_path: Path, + mode: str, + config_name: str, +) -> None: + env, args = _run_recipe(tmp_path, mode) + + assert Path(env["FOUNDRY_GRAPH_EXTENSION_CONFIG_PATH"]) == RECIPE_DIR / config_name + assert env["SGLANG_PLUGINS"] == "foundry" + assert "--foundry-graph-extension-config-path" not in args + assert args[args.index("--tp-size") + 1] == "2" + assert args[args.index("--revision") + 1] == env["SGL_MODEL_REVISION"] + assert "--disable-custom-all-reduce" in args + + +def test_tp_save_and_load_configs_are_symmetric() -> None: + save = tomllib.loads((RECIPE_DIR / "sglang_foundry_tp_save.toml").read_text()) + load = tomllib.loads((RECIPE_DIR / "sglang_foundry_tp_load.toml").read_text()) + + assert save.pop("mode") == "save" + assert load.pop("mode") == "load" + assert save == load diff --git a/tests/test_sglang_tp_validation.py b/tests/test_sglang_tp_validation.py new file mode 100644 index 00000000..a610a309 --- /dev/null +++ b/tests/test_sglang_tp_validation.py @@ -0,0 +1,1756 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the Foundry project +"""CPU-only contracts for the SGLang TP validation evidence helpers. + +Every fixture line below is built from the *exact* Foundry / SGLang log +message it stands in for, so a drift in either the Foundry message strings +(``python/foundry/integration/sglang/{graph_ops,runtime,hooks}.py``) or the +pinned SGLang logger prefix (` TP{rank}` from +``sglang.srt.managers.scheduler.run_scheduler_process``) breaks a contract +here before a GPU is ever provisioned. +""" + +from __future__ import annotations + +import importlib.util +import json +import os +import subprocess +import sys +import time +from pathlib import Path + +import pytest + +REPO_ROOT = Path(__file__).parents[1] +VALIDATION_PATH = REPO_ROOT / "tests" / "sglang_tp_validation.py" +HELPERS_PATH = REPO_ROOT / "tests" / "tp_validation_helpers.py" +FULL_SHA = "0123456789abcdef0123456789abcdef01234567" + + +def _load_validation(name: str = "sglang_tp_validation_contract"): + assert VALIDATION_PATH.exists(), f"required validation module is missing: {VALIDATION_PATH}" + spec = importlib.util.spec_from_file_location(name, VALIDATION_PATH) + assert spec is not None and spec.loader is not None + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + return module + + +def _load_helpers(name: str = "tp_validation_helpers_contract"): + assert HELPERS_PATH.exists(), f"required helper module is missing: {HELPERS_PATH}" + spec = importlib.util.spec_from_file_location(name, HELPERS_PATH) + assert spec is not None and spec.loader is not None + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + return module + + +# Log lines assembled from the pinned ` TP{rank}` scheduler prefix +# (``[ TP0] ``) and the exact Foundry message format strings. +def _sglang_line(rank: int, message: str) -> str: + return f"[2026-07-23 09:54:00 TP{rank}] {message}" + + +def _load_offset_line(rank: int, offset: int) -> str: + # runtime.log_alloc_offset: "[Foundry] SGLang alloc_offset[%s]=%d (%.2f MB)" + return _sglang_line( + rank, f"[Foundry] SGLang alloc_offset[after_load_all_graphs]={offset} (0.00 MB)" + ) + + +def _loaded_graphs_line(rank: int, count: int) -> str: + # graph_ops.load_all_graphs summary: + # "[Foundry] Loaded %d SGLang graphs in %.3fs (alloc_offset %d -> %d, +%d bytes)" + return _sglang_line( + rank, + f"[Foundry] Loaded {count} SGLang graphs in 0.500s (alloc_offset 100 -> 200, +100 bytes)", + ) + + +def _per_graph_loaded_line(rank: int, index: int, bs: int) -> str: + # graph_ops.load_all_graphs per-graph: + # "[Foundry] Loaded SGLang graph idx=%d bs=%d alloc_offset=%d" + return _sglang_line( + rank, f"[Foundry] Loaded SGLang graph idx={index} bs={bs} alloc_offset=12345" + ) + + +def _saved_graph_line(rank: int, index: int, bs: int) -> str: + # graph_ops.save_graph: "[Foundry] Saved SGLang CUDA graph %s key=%s" + return _sglang_line( + rank, + f"[Foundry] Saved SGLang CUDA graph graph_{index}_FULL_t{bs}_r{bs}_UX_pcN.json " + f"key=ShapeKey(size={bs}, stream_idx=None, variant_label=None)", + ) + + +def _torch_symm_mem_line(rank: int, address: int) -> str: + # hooks._around_torch_symm_mem_init: + # "[Foundry] Torch symmetric-memory buffer address=0x%x" + return _sglang_line( + rank, + f"[Foundry] Torch symmetric-memory buffer address=0x{address:x}", + ) + + +def _multimem_all_gather_line(rank: int, address: int) -> str: + # hooks._around_multimem_all_gather: + # "[Foundry] Multimem all-gather buffer address=0x%x" + return _sglang_line( + rank, + f"[Foundry] Multimem all-gather buffer address=0x{address:x}", + ) + + +def _nccl_channel_line(rank: int, channel: int, transport: str = "P2P/IPC") -> str: + # NCCL's own connection log; rank is NCCL's [rank] bracket, not SGLang's prefix. + return ( + f"gpu-host:1234:1250 [{rank}] NCCL INFO Channel {channel:02d}/24 : {rank} 1 via {transport}" + ) + + +def _decode_replay_line(rank: int, running_req: int) -> str: + # scheduler_components.metrics_reporter decode log; label is "cuda graph" on CUDA. + return _sglang_line( + rank, + f"Decode batch, #running-req: {running_req}, #token: 128, token usage: 0.01, " + f"cuda graph: True, gen throughput (token/s): 42.00, #queue-req: 0", + ) + + +def _native_capture_line(bs: int) -> str: + # decode_cuda_graph_runner capture progress: "Capturing batches ({bs=} {avail_mem=:.2f} GB)" + return _sglang_line(0, f"Capturing batches (bs={bs} avail_mem=50.00 GB)") + + +def _server_args_line() -> str: + # The exact benign ServerArgs repr fields that previously false-positived on + # the bare ``Watchdog`` token (run ap-4glkv7uGUyl28MutKmVGhi baseline, line 8). + return ( + "[2026-07-23 11:37:54] server_args=ServerArgs(model_path='Qwen/Qwen3-8B', " + "watchdog_timeout=300, soft_watchdog_timeout=None, custom_sigquit_handler=None, " + "schedule_policy='fcfs', scheduler_recv_interval=1, attention_backend='flashinfer')" + ) + + +# The exact benign SIGINT teardown noise from run ap-4glkv7uGUyl28MutKmVGhi +# (baseline.log): the scheduler / detokenizer workers were blocked in +# broadcast_pyobj / zmq recv when the harness sent SIGINT, so PyErr_CheckSignals +# raised KeyboardInterrupt, and the detokenizer was force-killed (-9) during the +# expected shutdown cascade. None of these are genuine failures. +_TEARDOWN_KEYBOARD_INTERRUPT = [ + "Process Process-1:", + "Process Process-2:", + "Traceback (most recent call last):", + ' File "/usr/local/lib/python3.12/multiprocessing/process.py", line 314, in _bootstrap', + " self.run()", + ' File "/sglang/python/sglang/srt/managers/scheduler.py", line 1573, in event_loop_overlap', + " recv_reqs = self.request_receiver.recv_requests()", + ' File "/sglang/python/sglang/srt/utils/common.py", line 2190, in broadcast_pyobj', + " dist.broadcast(tensor_size, src=src, group=dist_group)", + "KeyboardInterrupt", + "[2026-07-23 11:40:55] Subprocess detokenizer (pid=50) crashed with exit code -9. " + "Triggering SIGQUIT for cleanup...", +] + + +# The exact recovered startup-retry traceback from run ap-XO6ZBjyZ9xXBpYmPngh8IZ +# (baseline.log lines 22-100): a transient HF Hub connection reset during model +# config resolution, wrapped by SGLang's own retry() helper, which logged the +# full chained traceback and then retried successfully. The server went on to +# full readiness and served every request (return code 0, 55 graphs). None of +# these lines are a genuine failure. +_HF_RETRY_TRACEBACK = [ + "Traceback (most recent call last):", + ' File "/usr/local/lib/python3.12/site-packages/httpx/_transports/default.py", ' + "line 101, in map_httpcore_exceptions", + " yield", + "httpcore.ConnectError: [Errno 104] Connection reset by peer", + "", + "The above exception was the direct cause of the following exception:", + "", + "Traceback (most recent call last):", + ' File "/sglang/python/sglang/srt/utils/common.py", line 3177, in retry', + " return fn()", + ' File "/sglang/python/sglang/srt/configs/model_config.py", line 1143, in ', + " lambda: hf_api.file_exists(", + "httpx.ConnectError: [Errno 104] Connection reset by peer", + "[2026-07-23 12:15:12 TP0] retry() failed once (0th try, maximum 2 retries). " + "Will delay 0.79s and retry. Error: [Errno 104] Connection reset by peer", +] + + +def _readiness_line() -> str: + # SGLang's canonical readiness banner (baseline.log line 420). + return _sglang_line(0, "The server is fired up and ready to roll!") + + +def test_commit_and_tp_helpers_are_strict() -> None: + validation = _load_validation("sglang_tp_validation_helpers") + + assert validation.require_commit_sha("FOUNDRY_REF", FULL_SHA.upper()) == FULL_SHA + for invalid in (None, "", "main", "abc123", "g" * 40): + with pytest.raises(ValueError, match="40-hex commit SHA"): + validation.require_commit_sha("FOUNDRY_REF", invalid) + + assert validation.require_tp_size("2") == 2 + for invalid in (1, 3, "two"): + with pytest.raises(ValueError, match="TP_SIZE must be exactly 2"): + validation.require_tp_size(invalid) + + +def test_scan_parses_load_evidence_from_exact_foundry_messages() -> None: + validation = _load_validation("sglang_tp_validation_load_scan") + text = "\n".join( + [ + _load_offset_line(0, 69736595456), + _load_offset_line(1, 69736595456), + _per_graph_loaded_line(0, 0, 256), + _per_graph_loaded_line(1, 0, 256), + _loaded_graphs_line(0, 36), + _loaded_graphs_line(1, 36), + ] + ) + + scan = validation.scan_log_text(text) + + assert scan["load_offsets"] == {"0": 69736595456, "1": 69736595456} + assert scan["loaded_graphs"] == {"0": 36, "1": 36} + # The per-graph "Loaded SGLang graph idx=" line must not be miscounted as a summary. + assert scan["errors"] == [] + + +def test_scan_parses_symmetric_collective_markers_per_rank() -> None: + validation = _load_validation("sglang_tp_validation_symmetric_markers") + text = "\n".join( + [ + _torch_symm_mem_line(1, 0x7F1000), + _multimem_all_gather_line(0, 0x7F2000), + _torch_symm_mem_line(0, 0x7F3000), + _multimem_all_gather_line(1, 0x7F4000), + _torch_symm_mem_line(0, 0x7F3000), + ] + ) + + scan = validation.scan_log_text(text) + + assert scan["torch_symm_mem_ranks"] == ["0", "1"] + assert scan["multimem_all_gather_ranks"] == ["0", "1"] + + +def test_scan_parses_symmetric_markers_after_carriage_return_progress() -> None: + validation = _load_validation("sglang_tp_validation_cr_symmetric_markers") + text = ( + f"capture progress 50%\r{_torch_symm_mem_line(0, 0x7F1000)}\n" + f"{_torch_symm_mem_line(1, 0x7F2000)}\n" + f"progress...\r{_multimem_all_gather_line(0, 0x7F3000)}\n" + f"{_multimem_all_gather_line(1, 0x7F4000)}\n" + f'progress...\r"{_torch_symm_mem_line(2, 0x7F5000)}"\n' + f"progress...\rnested {_multimem_all_gather_line(3, 0x7F6000)}" + ) + + scan = validation.scan_log_text(text) + + assert scan["torch_symm_mem_ranks"] == ["0", "1"] + assert scan["torch_symm_mem_addresses"] == { + "0": ["0x7f1000"], + "1": ["0x7f2000"], + } + assert scan["multimem_all_gather_ranks"] == ["0", "1"] + assert scan["multimem_all_gather_addresses"] == { + "0": ["0x7f3000"], + "1": ["0x7f4000"], + } + + +def test_scan_keeps_missing_symmetric_marker_ranks_visible() -> None: + validation = _load_validation("sglang_tp_validation_missing_symmetric_marker") + text = "\n".join( + [ + _torch_symm_mem_line(0, 0x7F1000), + _multimem_all_gather_line(1, 0x7F2000), + ] + ) + + scan = validation.scan_log_text(text) + + assert scan["torch_symm_mem_ranks"] == ["0"] + assert scan["multimem_all_gather_ranks"] == ["1"] + + +def test_scan_accepts_only_exact_anchored_symmetric_marker_lines() -> None: + validation = _load_validation("sglang_tp_validation_anchored_symmetric_markers") + valid_torch = _torch_symm_mem_line(0, 0x7F1000) + valid_gather = _multimem_all_gather_line(1, 0x7F2000) + adversarial = [ + _sglang_line(2, "[Foundry] Torch symmetric-memory buffer address=0x"), + f'"{_torch_symm_mem_line(3, 0x7F3000)}"', + _sglang_line( + 4, + f"nested {_multimem_all_gather_line(5, 0x7F4000)}", + ), + "[2026-07-23 09:54:00 TPwrong] [Foundry] Torch symmetric-memory buffer address=0x7f5000", + f"prefix {_multimem_all_gather_line(6, 0x7F6000)}", + f"{_torch_symm_mem_line(7, 0x7F7000)} trailing", + ] + + scan = validation.scan_log_text("\n".join([valid_torch, valid_gather, *adversarial])) + + assert scan["torch_symm_mem_ranks"] == ["0"] + assert scan["multimem_all_gather_ranks"] == ["1"] + assert scan["torch_symm_mem_addresses"] == {"0": ["0x7f1000"]} + assert scan["multimem_all_gather_addresses"] == {"1": ["0x7f2000"]} + + +def test_archive_collective_kernel_counts_recurses_kernel_node_strings(tmp_path: Path) -> None: + helpers = _load_helpers("tp_validation_helpers_collective_nodes") + symmetric_all_reduce_name = ( + "void c10d::cuda::CUDASymmetricMemoryOps::" + "one_shot_all_reduce_kernel(at::BFloat16 const*)" + ) + symmetric_all_gather_name = ( + "void c10d::cuda::CUDASymmetricMemoryOps::" + "multimem_all_gather_kernel(at::BFloat16 const*)" + ) + kernel_nodes = [ + { + "type": "KernelNode", + "params": { + "name": "ncclDevKernel_AllReduce_Sum_bf16_RING_LL", + "nested": ["ncclDevKernel_AllReduce_duplicate", {"label": "ordinary_kernel"}], + }, + }, + { + "type": "KernelNode", + "params": {"nested": {"symbol": "ncclDevKernel_AllGather_RING_LL"}}, + }, + { + "type": "KernelNode", + "params": { + "function_name": symmetric_all_reduce_name, + "nested_duplicate": [symmetric_all_reduce_name], + }, + }, + { + "type": "KernelNode", + "params": {"nested": {"function_name": symmetric_all_gather_name}}, + }, + { + "type": "KernelNode", + "params": {"name": "flashinfer_attention_kernel"}, + }, + { + "type": "KernelNode", + "params": {"ncclDevKernel_AllReduce_log_field": "ordinary text"}, + }, + { + "type": "MemcpyNode", + "params": {"log": "ncclDevKernel_AllReduce_not_a_kernel_node"}, + }, + { + "type": "KernelNode", + "params": { + "function_name": "CUDASymmetricMemoryOps::copy_kernel", + "description": "all_reduce appears only in another value", + }, + }, + { + "type": "KernelNode", + "params": {"function_name": "unrelated_all_gather_kernel"}, + }, + { + "type": "KernelNode", + "params": {"CUDASymmetricMemoryOps::all_reduce": "ordinary value"}, + }, + { + "type": "MemcpyNode", + "params": {"function_name": symmetric_all_gather_name}, + }, + ] + for rank in range(2): + rank_dir = tmp_path / f"rank_{rank}" + rank_dir.mkdir() + (rank_dir / "graph_0.json").write_text( + json.dumps( + { + "nodes": kernel_nodes, + "log": "ncclDevKernel_AllGather_not_a_node", + } + ) + ) + + assert helpers.archive_collective_kernel_counts(tmp_path, 2) == { + "0": { + "all_reduce": 1, + "all_gather": 1, + "symmetric_all_reduce": 1, + "symmetric_all_gather": 1, + }, + "1": { + "all_reduce": 1, + "all_gather": 1, + "symmetric_all_reduce": 1, + "symmetric_all_gather": 1, + }, + } + assert helpers.archive_graph_json_counts(tmp_path, 2) == {"0": 1, "1": 1} + + +def test_archive_collective_scanner_rejects_missing_rank_directory(tmp_path: Path) -> None: + helpers = _load_helpers("tp_validation_helpers_missing_rank") + (tmp_path / "rank_0").mkdir() + (tmp_path / "rank_0" / "graph_0.json").write_text(json.dumps({"nodes": []})) + + with pytest.raises(ValueError, match=r"rank_1.*does not exist"): + helpers.archive_collective_kernel_counts(tmp_path, 2) + + +def test_archive_collective_scanner_rejects_rank_without_graph_json(tmp_path: Path) -> None: + helpers = _load_helpers("tp_validation_helpers_no_graph_json") + for rank in range(2): + rank_dir = tmp_path / f"rank_{rank}" + rank_dir.mkdir() + (rank_dir / "graph_manifest.json").write_text(json.dumps({"graphs": []})) + + with pytest.raises(ValueError, match=r"rank_0.*no graph JSON"): + helpers.archive_collective_kernel_counts(tmp_path, 2) + + +def test_archive_collective_scanner_rejects_malformed_graph_json(tmp_path: Path) -> None: + helpers = _load_helpers("tp_validation_helpers_bad_graph_json") + for rank in range(2): + rank_dir = tmp_path / f"rank_{rank}" + rank_dir.mkdir() + (rank_dir / "graph_0.json").write_text( + "not-json" if rank == 0 else json.dumps({"nodes": []}) + ) + + with pytest.raises(ValueError, match=r"rank_0/graph_0\.json.*malformed JSON"): + helpers.archive_collective_kernel_counts(tmp_path, 2) + + +@pytest.mark.parametrize("nodes", [None, {}, "KernelNode"]) +def test_archive_collective_scanner_rejects_non_list_nodes( + tmp_path: Path, + nodes: object, +) -> None: + helpers = _load_helpers(f"tp_validation_helpers_bad_nodes_{type(nodes).__name__}") + for rank in range(2): + rank_dir = tmp_path / f"rank_{rank}" + rank_dir.mkdir() + (rank_dir / "graph_0.json").write_text(json.dumps({"nodes": nodes})) + + with pytest.raises(ValueError, match=r"rank_0/graph_0\.json.*nodes.*list"): + helpers.archive_collective_kernel_counts(tmp_path, 2) + + +def _focused_save_result() -> dict: + return { + "phase": "save", + "became_healthy": True, + "request_completed": True, + "outputs": ["deterministic output"], + "archive_graph_counts": {"0": 1, "1": 1}, + "archive_graph_json_counts": {"0": 1, "1": 1}, + "archive_collective_kernel_counts": { + "0": { + "all_reduce": 0, + "all_gather": 0, + "symmetric_all_reduce": 3, + "symmetric_all_gather": 1, + }, + "1": { + "all_reduce": 0, + "all_gather": 0, + "symmetric_all_reduce": 3, + "symmetric_all_gather": 1, + }, + }, + "rank_offsets": {"0": 4096, "1": 4096}, + "torch_symm_mem_ranks": ["0", "1"], + "multimem_all_gather_ranks": ["1"], + "torch_symm_mem_addresses": {"0": ["0x1000"], "1": ["0x2000"]}, + "multimem_all_gather_addresses": {"1": ["0x3000"]}, + "archive_fingerprints": {"0": {"graph_0.json": "a"}, "1": {"graph_0.json": "b"}}, + "archive_byte_fingerprints": {"rank_0/graph_0.json": "a", "rank_1/graph_0.json": "b"}, + "errors": [], + "error_categories": [], + "python_error_types": [], + "evidence_errors": [], + "failure_events": [], + "phase_failures": [], + "process_return_code": 0, + "archive_committed": True, + } + + +def _focused_load_result() -> dict: + save = _focused_save_result() + return { + **save, + "phase": "load", + "loaded_graphs": {"0": 1, "1": 1}, + "load_offsets": {"0": 4096, "1": 4096}, + "graph_replay_batch_sizes": [1], + # The pinned LOAD lifecycle runs the real upstream capture() traversal, + # so rank0 emits "Capturing batches" progress even though the Foundry + # backend performs no CUDA capture. It is diagnostic, not an acceptance + # gate — a healthy LOAD has a positive count. + "native_capture_progress_count": 1, + "saved_graph_log_count": 0, + "archive_committed": False, + } + + +@pytest.mark.parametrize( + ("updates", "failed_check"), + [ + ({"request_completed": False}, "request_completed"), + ({"archive_graph_counts": {"0": 1}}, "graph_counts_exact"), + ({"archive_graph_json_counts": {"0": 0, "1": 1}}, "graph_json_counts_exact"), + ({"torch_symm_mem_ranks": ["0"]}, "symmetric_markers_present"), + ({"multimem_all_gather_addresses": {}}, "symmetric_markers_present"), + ( + { + "archive_collective_kernel_counts": { + "0": { + "all_reduce": 1, + "all_gather": 0, + "symmetric_all_reduce": 1, + "symmetric_all_gather": 1, + }, + "1": { + "all_reduce": 0, + "all_gather": 0, + "symmetric_all_reduce": 1, + "symmetric_all_gather": 1, + }, + } + }, + "no_archived_nccl_collectives", + ), + ( + { + "archive_collective_kernel_counts": { + "0": { + "all_reduce": 0, + "all_gather": 0, + "symmetric_all_reduce": 0, + "symmetric_all_gather": 1, + }, + "1": { + "all_reduce": 0, + "all_gather": 0, + "symmetric_all_reduce": 1, + "symmetric_all_gather": 1, + }, + } + }, + "symmetric_archive_collectives_present", + ), + ( + { + "archive_collective_kernel_counts": { + "0": { + "all_reduce": 0, + "all_gather": 0, + "symmetric_all_reduce": 1, + "symmetric_all_gather": 1, + }, + "1": { + "all_reduce": 0, + "all_gather": 0, + "symmetric_all_reduce": 1, + "symmetric_all_gather": 0, + }, + } + }, + "symmetric_archive_collectives_present", + ), + ({"rank_offsets": {"0": 4096, "1": 8192}}, "offsets_equal_nonzero"), + ({"evidence_errors": ["archive scanner failed"]}, "no_scanner_errors"), + ({"process_return_code": 1}, "expected_shutdown"), + ( + {"phase_failures": ["[foundry LOAD ERROR] restored graph failed"]}, + "no_policy_failure", + ), + ], +) +def test_focused_save_checks_expose_known_violations( + updates: dict, + failed_check: str, +) -> None: + validation = _load_validation(f"sglang_tp_validation_save_{failed_check}") + result = _focused_save_result() + result.update(updates) + + checks = validation.focused_save_checks(result, tp_size=2) + + assert checks[failed_check] is False + + +def test_focused_save_checks_leave_pre_health_environment_evidence_unknown() -> None: + validation = _load_validation("sglang_tp_validation_save_unknown") + result = { + "phase": "save", + "became_healthy": False, + "request_completed": False, + "phase_failures": ["phase execution: network setup unavailable"], + "evidence_errors": [], + "errors": [], + "process_return_code": 0, + } + + checks = validation.focused_save_checks(result, tp_size=2) + + assert checks["request_completed"] is None + assert checks["graph_counts_exact"] is None + assert checks["post_health_execution_clean"] is None + assert not any(value is False for value in checks.values()) + + +@pytest.mark.parametrize( + ("mutate_phase", "field", "value", "failed_check"), + [ + ("load", "outputs", ["different output"], "outputs_equal_nonempty"), + ("load", "loaded_graphs", {"0": 0, "1": 1}, "graphs_restored_each_rank"), + ("load", "graph_replay_batch_sizes", [], "restored_graph_replay_observed"), + # A Foundry graph-save log on LOAD means a graph was re-captured/re-saved. + ("load", "saved_graph_log_count", 1, "load_restored_without_resave"), + # Zero restored graphs on LOAD (with no save logs) still fails the + # positive-archived-load requirement of the no-resave gate. + ("load", "loaded_graphs", {"0": 0, "1": 0}, "load_restored_without_resave"), + ("load", "load_offsets", {"0": 4096, "1": 8192}, "rank_offsets_equal_nonzero"), + ("load", "archive_byte_fingerprints", {"mutated": "x"}, "load_archive_bytes_immutable"), + ("load", "evidence_errors", ["malformed graph JSON"], "no_scanner_errors"), + ("load", "multimem_all_gather_addresses", {}, "symmetric_markers_each_phase"), + ( + "load", + "archive_collective_kernel_counts", + { + "0": { + "all_reduce": 0, + "all_gather": 0, + "symmetric_all_reduce": 0, + "symmetric_all_gather": 1, + }, + "1": { + "all_reduce": 0, + "all_gather": 0, + "symmetric_all_reduce": 1, + "symmetric_all_gather": 1, + }, + }, + "symmetric_archive_collectives_each_phase", + ), + ], +) +def test_focused_treatment_checks_expose_functional_evidence( + mutate_phase: str, + field: str, + value: object, + failed_check: str, +) -> None: + validation = _load_validation(f"sglang_tp_validation_treatment_{failed_check}") + save = _focused_save_result() + load = _focused_load_result() + target = save if mutate_phase == "save" else load + target[field] = value + + checks = validation.focused_treatment_checks(save, load, tp_size=2) + + assert checks[failed_check] is False + + +def test_focused_treatment_accepts_global_gather_marker_with_per_rank_archive_proof() -> None: + validation = _load_validation("sglang_tp_validation_pr4_focused_acceptance") + checks = validation.focused_treatment_checks( + _focused_save_result(), + _focused_load_result(), + tp_size=2, + ) + + assert checks["symmetric_markers_each_phase"] is True + assert checks["symmetric_archive_collectives_each_phase"] is True + assert checks["no_archived_nccl_collectives"] is True + assert ( + validation.classify_focused_checks( + checks, + environment_failure=False, + restored_replay=True, + ) + == "VERIFIED" + ) + + +def test_focused_treatment_accepts_upstream_load_capture_progress() -> None: + # The pinned LOAD lifecycle traverses the real upstream capture() loop, so + # rank0 prints many "Capturing batches" progress lines while the backend + # only restores archived graphs. That progress count must NOT gate + # acceptance: a LOAD with a high native_capture_progress_count but zero + # graph-save logs and positive restored counts still VERIFIES. + validation = _load_validation("sglang_tp_validation_load_progress_diagnostic") + save = _focused_save_result() + load = _focused_load_result() + load["native_capture_progress_count"] = 256 + + checks = validation.focused_treatment_checks(save, load, tp_size=2) + + assert checks["load_restored_without_resave"] is True + assert ( + validation.classify_focused_checks( + checks, + environment_failure=False, + restored_replay=True, + ) + == "VERIFIED" + ) + + +@pytest.mark.parametrize( + ("checks", "environment_failure", "restored_replay", "expected"), + [ + ({"complete": True}, False, True, "VERIFIED"), + ({"unavailable": None}, True, False, "INCONCLUSIVE"), + ({"functional": False, "unavailable": None}, True, False, "NOT VERIFIED"), + ({"unavailable": None}, True, True, "NOT VERIFIED"), + ({"harness": False}, False, False, "NOT VERIFIED"), + ({"incomplete": None}, False, False, "NOT VERIFIED"), + ], +) +def test_focused_classification_is_evidence_first( + checks: dict[str, bool | None], + environment_failure: bool, + restored_replay: bool, + expected: str, +) -> None: + validation = _load_validation(f"sglang_tp_validation_classify_{expected.replace(' ', '_')}") + + assert ( + validation.classify_focused_checks( + checks, + environment_failure=environment_failure, + restored_replay=restored_replay, + ) + == expected + ) + + +def test_later_environment_failure_cannot_override_output_mismatch() -> None: + validation = _load_validation("sglang_tp_validation_output_precedence") + save = _focused_save_result() + load = _focused_load_result() + load["outputs"] = ["different output"] + load["phase_failures"] = ["phase execution: unrelated network failure"] + checks = validation.focused_treatment_checks(save, load, tp_size=2) + + assert checks["outputs_equal_nonempty"] is False + assert ( + validation.classify_focused_checks( + checks, + environment_failure=True, + restored_replay=False, + ) + == "NOT VERIFIED" + ) + + +def test_allowlisted_load_startup_failure_is_inconclusive() -> None: + validation = _load_validation("sglang_tp_validation_load_startup_environment") + save = _focused_save_result() + load = { + "phase": "load", + "became_healthy": False, + "request_completed": False, + "failure_stage": "startup", + "failure_type": "builtins.ConnectionError", + "failure_events": [ + { + "stage": "startup", + "type": "builtins.ConnectionError", + } + ], + "errors": ["ConnectionError: network unavailable"], + "error_categories": [], + "python_error_types": ["ConnectionError"], + "evidence_errors": [], + "phase_failures": ["phase execution: network unavailable"], + "process_return_code": 0, + "archive_committed": False, + } + + checks = validation.focused_treatment_checks(save, load, tp_size=2) + environment_failure = validation.is_allowlisted_startup_environment_failure(load) + + assert environment_failure is True + assert not any(value is False for value in checks.values()) + assert ( + validation.classify_focused_checks( + checks, + environment_failure=environment_failure, + restored_replay=False, + ) + == "INCONCLUSIVE" + ) + + +def test_scan_reports_structured_python_types_and_error_categories() -> None: + validation = _load_validation("sglang_tp_validation_structured_errors") + text = "\n".join( + [ + "httpx.ConnectError: connection reset by peer", + "httpx.ConnectTimeout: model download timed out", + "AttributeError: harness_result has no attribute checks", + "[HOOK] FATAL ERROR: fixed-address allocation failed", + ] + ) + + scan = validation.scan_log_text(text) + + assert scan["python_error_types"] == [ + "AttributeError", + "httpx.ConnectError", + "httpx.ConnectTimeout", + ] + assert scan["error_categories"] == ["hard_native", "policy"] + assert scan["untyped_python_fatal_evidence"] == [] + assert scan["typed_python_traceback_headers"] == [] + + +def test_scan_extracts_python_types_from_timestamped_and_nested_fatal_lines() -> None: + validation = _load_validation("sglang_tp_validation_nested_python_types") + text = "\n".join( + [ + "requests.exceptions.ConnectionError: model download failed", + "[2026-07-23 12:00:00 TP0] AttributeError: nested log prose", + 'wrapper="httpx.ConnectTimeout: quoted failure"', + " RuntimeError: indented failure", + ] + ) + + scan = validation.scan_log_text(text) + + assert scan["python_error_types"] == [ + "AttributeError", + "RuntimeError", + "httpx.ConnectTimeout", + "requests.exceptions.ConnectionError", + ] + assert scan["untyped_python_fatal_evidence"] == [] + + +@pytest.mark.parametrize( + ("text", "typed_headers"), + [ + ( + "\n".join( + [ + "Traceback (most recent call last):", + ' File "/tmp/model.py", line 1, in load', + "ConnectionError: model download failed", + ] + ), + ["Traceback (most recent call last):"], + ), + ( + "\n".join( + [ + "Traceback (most recent call last):", + ' File "/tmp/model.py", line 1, in load', + ] + ), + [], + ), + ( + "\n".join( + [ + "Traceback (most recent call last):", + ' File "/tmp/model.py", line 1, in load', + "[2026-07-23 12:00:00 TP0] ConnectionError: separate log record", + ] + ), + [], + ), + ], +) +def test_scan_returns_untyped_traceback_headers_with_block_typing( + text: str, + typed_headers: list[str], +) -> None: + validation = _load_validation("sglang_tp_validation_traceback_blocks") + + scan = validation.scan_log_text(text) + + assert scan["untyped_python_fatal_evidence"] == [ + { + "line": 1, + "text": "Traceback (most recent call last):", + } + ] + assert scan["typed_python_traceback_headers"] == [ + { + "line": 1, + "text": header, + } + for header in typed_headers + ] + + +def _startup_result_with_scan(scan: dict) -> dict: + return { + "failure_stage": "startup", + "failure_type": "foundry.validation.FocusedStartupHealthError", + "startup_cause_types": ["builtins.ConnectionError"], + "startup_child_exited": False, + "python_error_types": scan["python_error_types"], + "untyped_python_fatal_evidence": scan["untyped_python_fatal_evidence"], + "typed_python_traceback_headers": scan["typed_python_traceback_headers"], + "error_categories": scan["error_categories"], + "evidence_errors": [], + "archive_committed": False, + "failure_events": [ + { + "stage": "startup", + "type": "foundry.validation.FocusedStartupHealthError", + } + ], + } + + +@pytest.mark.parametrize( + ("text", "expected"), + [ + ( + "[2026-07-23 12:00:00 TP0] AttributeError: startup bug", + False, + ), + ( + 'health wrapper="requests.exceptions.ConnectionError: unavailable"', + True, + ), + ( + "\n".join( + [ + "nested ConnectionError: unavailable", + "[2026-07-23 12:00:00 TP0] TypeError: startup bug", + ] + ), + False, + ), + ( + "\n".join( + [ + "Traceback (most recent call last):", + ' File "/tmp/model.py", line 1, in load', + "ConnectionError: model download failed", + ] + ), + True, + ), + ( + "\n".join( + [ + "Traceback (most recent call last):", + ' File "/tmp/model.py", line 1, in load', + "[2026-07-23 12:00:00 TP0] ConnectionError: separate log record", + ] + ), + False, + ), + ( + "\n".join( + [ + "Traceback (most recent call last):", + ' File "/tmp/model.py", line 1, in load', + "[2026-07-23 12:00:00 TP0] AttributeError: separate startup bug", + ] + ), + False, + ), + ("Traceback (most recent call last):", False), + ( + "\n".join( + [ + "Traceback (most recent call last):", + "ConnectionError: first block is allowlisted", + "Traceback (most recent call last):", + ] + ), + False, + ), + ( + "\n".join( + [ + "Traceback (most recent call last):", + ' File "/tmp/model.py", line 1, in load', + "AttributeError: startup bug", + ] + ), + False, + ), + ], +) +def test_startup_environment_classification_uses_all_nested_python_evidence( + text: str, + expected: bool, +) -> None: + validation = _load_validation("sglang_tp_validation_nested_startup_classification") + scan = validation.scan_log_text(text) + result = _startup_result_with_scan(scan) + + assert validation.is_allowlisted_startup_environment_failure(result) is expected + + +@pytest.mark.parametrize( + ("updates", "expected"), + [ + ({}, True), + ( + { + "python_error_types": [ + "ConnectionError", + "AttributeError", + ] + }, + False, + ), + ( + { + "startup_cause_types": [], + "python_error_types": [], + }, + False, + ), + ( + { + "startup_cause_types": [ + "builtins.ConnectionError", + "builtins.RuntimeError", + ] + }, + False, + ), + ( + { + "startup_cause_types": ["project.ConnectionError"], + "python_error_types": [], + }, + False, + ), + ( + { + "failure_events": [ + { + "stage": "startup", + "type": "foundry.validation.FocusedStartupHealthError", + }, + { + "stage": "scanner", + "type": "builtins.ValueError", + }, + ] + }, + False, + ), + ({"startup_child_exited": False}, False), + ], +) +def test_startup_environment_requires_all_causes_and_tolerates_only_exit_teardown( + updates: dict, + expected: bool, +) -> None: + validation = _load_validation("sglang_tp_validation_startup_cause_chain") + result = { + "failure_stage": "startup", + "failure_type": "foundry.validation.FocusedStartupHealthError", + "startup_cause_types": ["builtins.ConnectionError"], + "startup_child_exited": True, + "python_error_types": ["ConnectionError"], + "error_categories": ["teardown"], + "evidence_errors": [], + "archive_committed": False, + "failure_events": [ + { + "stage": "startup", + "type": "foundry.validation.FocusedStartupHealthError", + }, + { + "stage": "teardown", + "type": "builtins.RuntimeError", + }, + ], + **updates, + } + + assert validation.is_allowlisted_startup_environment_failure(result) is expected + + +@pytest.mark.parametrize( + ("updates", "expected"), + [ + ( + { + "failure_stage": "startup", + "failure_type": "builtins.ConnectionError", + "python_error_types": ["ConnectionError"], + }, + True, + ), + ( + { + "failure_stage": "startup", + "failure_type": "httpx.ConnectError", + "python_error_types": ["httpx.ConnectError"], + }, + True, + ), + ( + { + "failure_stage": "startup", + "failure_type": "builtins.TimeoutError", + "python_error_types": [], + }, + False, + ), + ( + { + "failure_stage": "startup", + "failure_type": "builtins.TimeoutError", + "python_error_types": ["httpx.ConnectTimeout"], + }, + True, + ), + ( + { + "failure_stage": "startup", + "failure_type": "builtins.AttributeError", + "python_error_types": ["AttributeError"], + }, + False, + ), + ( + { + "failure_stage": "request", + "failure_type": "builtins.ConnectionError", + "python_error_types": ["ConnectionError"], + }, + False, + ), + ( + { + "failure_stage": "startup", + "failure_type": "builtins.ConnectionError", + "python_error_types": ["ConnectionError"], + "error_categories": ["hard_native"], + }, + False, + ), + ( + { + "failure_stage": "scanner", + "failure_type": "json.JSONDecodeError", + "evidence_errors": ["scanner failed"], + }, + False, + ), + ( + { + "failure_stage": "startup", + "failure_type": "builtins.RuntimeError", + "python_error_types": ["RuntimeError"], + "error_categories": ["policy"], + }, + False, + ), + ], +) +def test_startup_environment_classification_uses_structured_evidence( + updates: dict, + expected: bool, +) -> None: + validation = _load_validation("sglang_tp_validation_structured_environment") + result = { + "failure_stage": None, + "failure_type": None, + "python_error_types": [], + "error_categories": [], + "evidence_errors": [], + "archive_committed": False, + **updates, + } + + assert validation.is_allowlisted_startup_environment_failure(result) is expected + + +@pytest.mark.parametrize( + ("module_name", "type_name", "expected"), + [ + ("modal.exception", "InternalFailure", True), + ("modal.exception", "ServiceError", True), + ("modal.exception", "ConnectionError", True), + ("modal.exception", "RemoteError", False), + ("builtins", "RuntimeError", False), + ("project.modal_adapter", "InfrastructureError", False), + ], +) +def test_modal_infrastructure_type_classification_is_exact( + module_name: str, + type_name: str, + expected: bool, +) -> None: + validation = _load_validation("sglang_tp_validation_modal_environment") + + assert validation.is_modal_infrastructure_error_type(module_name, type_name) is expected + + +def test_scan_reads_nccl_version_banner() -> None: + validation = _load_validation("sglang_tp_validation_nccl_version") + + assert ( + validation.scan_log_text("gpu:1:2 [0] NCCL INFO NCCL version 2.28.3+cuda13.0")[ + "nccl_version" + ] + == "2.28.3" + ) + assert validation.scan_log_text("no banner here")["nccl_version"] is None + + +def test_scan_parses_nccl_transport_channels_per_rank() -> None: + validation = _load_validation("sglang_tp_validation_nccl") + text = "\n".join( + [_nccl_channel_line(0, ch) for ch in range(3)] + + [_nccl_channel_line(1, ch) for ch in range(3)] + ) + + scan = validation.scan_log_text(text) + + assert scan["nccl_transport_channels"] == { + "0": {"P2P/IPC": [0, 1, 2]}, + "1": {"P2P/IPC": [0, 1, 2]}, + } + + +def test_scan_detects_non_p2p_ipc_transport() -> None: + validation = _load_validation("sglang_tp_validation_nccl_net") + text = "\n".join([_nccl_channel_line(0, 0, transport="NET/IB/0")]) + + scan = validation.scan_log_text(text) + + assert scan["nccl_transport_channels"] == {"0": {"NET/IB/0": [0]}} + + +def test_scan_parses_p2p_ipc_submode_transport() -> None: + validation = _load_validation("sglang_tp_validation_nccl_submode") + text = "\n".join( + [_nccl_channel_line(0, ch, transport="P2P/IPC/read") for ch in range(2)] + + [_nccl_channel_line(1, ch, transport="P2P/IPC/read") for ch in range(2)] + ) + + scan = validation.scan_log_text(text) + + assert scan["nccl_transport_channels"] == { + "0": {"P2P/IPC/read": [0, 1]}, + "1": {"P2P/IPC/read": [0, 1]}, + } + + +def test_channels_are_p2p_ipc_accepts_submodes_and_requires_every_rank() -> None: + validation = _load_validation("sglang_tp_validation_channels_ok") + expected_ranks = {"0", "1"} + + # Plain P2P/IPC and the P2P/IPC/read submode are both accepted. + assert validation.channels_are_p2p_ipc( + {"0": {"P2P/IPC": [0, 1, 2]}, "1": {"P2P/IPC/read": [0, 1, 2]}}, + expected_ranks, + ) + + +def test_channels_are_p2p_ipc_rejects_non_p2p_missing_or_empty() -> None: + validation = _load_validation("sglang_tp_validation_channels_bad") + expected_ranks = {"0", "1"} + + # A non-P2P/IPC transport on any rank fails closed. + assert not validation.channels_are_p2p_ipc( + {"0": {"P2P/IPC": [0, 1]}, "1": {"NET/IB/0": [0, 1]}}, + expected_ranks, + ) + # A missing rank fails closed. + assert not validation.channels_are_p2p_ipc({"0": {"P2P/IPC": [0, 1]}}, expected_ranks) + # Empty channel evidence for a rank fails closed. + assert not validation.channels_are_p2p_ipc( + {"0": {"P2P/IPC": [0, 1]}, "1": {"P2P/IPC": []}}, + expected_ranks, + ) + # No channels at all for a rank fails closed. + assert not validation.channels_are_p2p_ipc({"0": {"P2P/IPC": [0, 1]}, "1": {}}, expected_ranks) + + +def test_fatal_scan_covers_native_and_python_signatures() -> None: + validation = _load_validation("sglang_tp_validation_native_fatal") + # Exact native strings from csrc/hook.cpp, csrc/CUDAGraphParallel.cpp, and + # csrc/CUDAGraph.cpp; every named [foundry ERROR] class is fatal. + fatal_lines = [ + "[HOOK] FATAL ERROR: cuMemGetAllocationGranularity failed with error 999", + "[REPLAY] FATAL: Allocation address mismatch!", + "[foundry LOAD ERROR] cuGraphAddKernelNode FAILED for node 3 with error 400", + "[foundry INSTANTIATE ERROR] cudaGraphInstantiate FAILED with error 1: bad", + "[foundry ON-DEMAND ERROR] cuGraphExecUpdate FAILED for graph 2 (t256): 900 (err)", + # Anchored Python exception-type lines (Error/Exception:), + # not bare case-insensitive prose. With no readiness marker these are + # scanned over the whole pre-teardown text. + _sglang_line(0, "RuntimeError: Foundry LOAD requires a saved memory_pool_config"), + _sglang_line(1, "torch.cuda.OutOfMemoryError: CUDA out of memory. Tried to allocate"), + ] + + scan = validation.scan_log_text("\n".join(fatal_lines)) + + for signature in ( + "[HOOK] FATAL ERROR", + "[REPLAY] FATAL", + "[foundry LOAD ERROR]", + "[foundry INSTANTIATE ERROR]", + "[foundry ON-DEMAND ERROR]", + "RuntimeError", + "OutOfMemoryError", + ): + assert any(signature in line for line in scan["errors"]), signature + assert len(scan["errors"]) == len(fatal_lines) + + +def test_fatal_scan_does_not_flag_benign_native_or_prose_lines() -> None: + validation = _load_validation("sglang_tp_validation_benign") + # Benign [foundry ...] lines that carry no "ERROR]" class must not trip the + # generic native matcher. + benign_lines = [ + "[HOOK] INFO: VMM-IPC import remapped base to 0x600000000000", + "[HOOK] DEBUG: Stored 5 fatbin segments for device linking", + "[REPLAY] OK: Allocated 1024 bytes at expected address 0x600000000000", + "[REPLAY] HINT: Try a different base_addr in your TOML config", + "[foundry REPLAY] graph 2 (t256): updated 4 nodes in 1.0 us (prev graph 1)", + "[foundry ON-DEMAND] graph 2 (t256): updated 4 nodes in 1.0 us", + "[foundry BUILD] Template 0 (t256): building...", + "[foundry] WARNING: Failed to write binary graph: /tmp/x", + # The benign ServerArgs repr with watchdog_timeout=/soft_watchdog_timeout= + # must never be flagged (the proven baseline false positive). + _server_args_line(), + _decode_replay_line(0, 2), + _loaded_graphs_line(0, 36), + _load_offset_line(0, 69736595456), + ] + + scan = validation.scan_log_text("\n".join(benign_lines)) + + assert scan["errors"] == [] + + +def test_scan_counts_capture_and_save_and_replay_evidence() -> None: + validation = _load_validation("sglang_tp_validation_counts") + text = "\n".join( + [ + _saved_graph_line(0, 0, 256), + _saved_graph_line(0, 1, 128), + _native_capture_line(256), + _native_capture_line(128), + _decode_replay_line(0, 1), + _decode_replay_line(0, 3), + ] + ) + + scan = validation.scan_log_text(text) + + assert scan["saved_graph_log_count"] == 2 + assert scan["native_capture_progress_count"] == 2 + assert scan["graph_replay_batch_sizes"] == [1, 3] + + +def test_fatal_scan_is_line_based_and_broad() -> None: + validation = _load_validation("sglang_tp_validation_fatal") + text = "\n".join( + [ + "[2026-07-23 09:54:00 TP0] Scheduler hit an exception: RuntimeError", + # Anchored real watchdog failure (not the benign watchdog_timeout= field). + "[2026-07-23 09:54:00 TP1] Scheduler watchdog timeout after 300s, killing the server", + "gpu-host:1234:1250 [0] NCCL WARN Cuda failure", + "CUDA error: an illegal memory access was encountered", + "Traceback (most recent call last):", + "TpModelWorkerClient hit an exception", + "Xid 74 detected", + "segmentation fault", + "[2026-07-23 09:54:00 TP0] [Foundry] Loaded 36 SGLang graphs in 0.5s (alloc_offset 100 -> 200, +100 bytes)", + ] + ) + + scan = validation.scan_log_text(text) + joined = " ".join(scan["errors"]).lower() + + for signature in ( + "scheduler hit an exception", + "watchdog timeout", + "nccl warn", + "illegal memory access", + "traceback", + "tpmodelworkerclient hit an exception", + "xid", + "segmentation fault", + ): + assert signature in joined + # A healthy Foundry LOAD summary must not be flagged as fatal. + assert not any("Loaded 36 SGLang graphs" in line for line in scan["errors"]) + + +def test_two_window_scan_ignores_expected_sigint_teardown_noise() -> None: + validation = _load_validation("sglang_tp_validation_teardown") + healthy = [ + _server_args_line(), + _sglang_line(0, "The server is fired up and ready to roll!"), + _decode_replay_line(0, 1), + _decode_replay_line(0, 8), + ] + # run_phase snapshots the log before SIGINT (no teardown tracebacks yet), then + # scans the full log (with teardown noise) after stop_process. + pre_teardown_text = "\n".join(healthy) + full_text = "\n".join(healthy + _TEARDOWN_KEYBOARD_INTERRUPT) + + scan = validation.scan_log_text(full_text, pre_teardown_text=pre_teardown_text) + + assert scan["errors"] == [] + # All other signals are still parsed from the full log. + assert scan["graph_replay_batch_sizes"] == [1, 8] + + +def test_post_teardown_window_excludes_bare_traceback_and_keyboardinterrupt() -> None: + validation = _load_validation("sglang_tp_validation_post_excl") + + # The post-teardown window (full log) never counts the SIGINT teardown noise. + assert validation.scan_post_teardown_errors("\n".join(_TEARDOWN_KEYBOARD_INTERRUPT)) == [] + + +def test_pre_teardown_window_flags_genuine_python_traceback() -> None: + validation = _load_validation("sglang_tp_validation_pre_trace") + pre_teardown_text = "\n".join( + [ + _sglang_line(0, "Traceback (most recent call last):"), + ' File "/sglang/python/sglang/srt/model_executor/model_runner.py", line 1, in init', + _sglang_line(0, "RuntimeError: Foundry LOAD requires a saved memory_pool_config"), + ] + ) + + # A real failure before teardown is fatal (Traceback + RuntimeError). + pre_errors = validation.scan_pre_teardown_errors(pre_teardown_text) + assert any("Traceback" in line for line in pre_errors) + assert any("RuntimeError" in line for line in pre_errors) + + scan = validation.scan_log_text(pre_teardown_text, pre_teardown_text=pre_teardown_text) + assert scan["errors"] + + +def test_pre_teardown_watchdog_failures_fatal_but_serverargs_field_benign() -> None: + validation = _load_validation("sglang_tp_validation_watchdog") + + for fatal in ( + _sglang_line(0, "WatchdogError: scheduler stalled"), + _sglang_line(0, "Scheduler watchdog timeout after 300s"), + _sglang_line(0, "watchdog thread crashed"), + ): + assert validation.scan_pre_teardown_errors(fatal), fatal + + # The benign ServerArgs fields must never be flagged in either window. + assert validation.scan_pre_teardown_errors(_server_args_line()) == [] + assert validation.scan_post_teardown_errors(_server_args_line()) == [] + + +def test_subprocess_crash_is_fatal_pre_teardown_and_ignored_post_teardown() -> None: + validation = _load_validation("sglang_tp_validation_subproc") + + # The exact captured line from run ap-4glkv7uGUyl28MutKmVGhi (baseline.log). + captured = _TEARDOWN_KEYBOARD_INTERRUPT[-1] + assert "Subprocess detokenizer (pid=50) crashed with exit code -9" in captured + + # Pre-teardown: a real detokenizer/tokenizer subprocess crash is fatal. + assert validation.scan_pre_teardown_errors(captured) == [captured] + tokenizer_crash = ( + "[2026-07-23 11:40:55] Subprocess tokenizer (pid=51) crashed with exit code -6. " + "Triggering SIGQUIT for cleanup..." + ) + assert validation.scan_pre_teardown_errors(tokenizer_crash) == [tokenizer_crash] + + # Post-teardown: the same line is expected child-exit noise during SIGINT. + assert validation.scan_post_teardown_errors(captured) == [] + assert validation.scan_post_teardown_errors(tokenizer_crash) == [] + + # A crash that surfaces in the pre-teardown snapshot (genuine mid-run) is fatal + # through the merged two-window scan; when it only appears post-teardown it is + # ignored. + assert validation.scan_log_text(captured, pre_teardown_text=captured)["errors"] == [captured] + healthy = _sglang_line(0, "The server is fired up and ready to roll!") + scan = validation.scan_log_text("\n".join([healthy, captured]), pre_teardown_text=healthy) + assert scan["errors"] == [] + + # Benign ServerArgs config is unaffected in either window. + assert validation.scan_pre_teardown_errors(_server_args_line()) == [] + + +def test_recovered_startup_retry_traceback_before_readiness_is_not_fatal() -> None: + validation = _load_validation("sglang_tp_validation_retry") + + # 1. Exact run-9 fixture: chained httpx.ConnectError + retry() log BEFORE + # readiness, followed by readiness and healthy lines => no fatal. + healthy_tail = [ + _readiness_line(), + _decode_replay_line(0, 1), + _sglang_line(0, 'INFO: 127.0.0.1:17331 - "POST /generate HTTP/1.1" 200 OK'), + ] + recovered = "\n".join(_HF_RETRY_TRACEBACK + healthy_tail) + assert validation.scan_pre_teardown_errors(recovered) == [] + assert validation.scan_log_text(recovered, pre_teardown_text=recovered)["errors"] == [] + + # 2. The SAME traceback AFTER readiness is a genuine runtime failure => fatal. + after_readiness = "\n".join([_readiness_line(), *_HF_RETRY_TRACEBACK]) + post_ready_errors = validation.scan_pre_teardown_errors(after_readiness) + assert any("Traceback" in line for line in post_ready_errors) + assert any("httpx.ConnectError:" in line for line in post_ready_errors) + + # 3. Unrecovered / no readiness ever => generic failures scanned over the whole + # pre-teardown text => fatal (health logic would fail too). + never_ready = "\n".join(_HF_RETRY_TRACEBACK) + assert validation.scan_pre_teardown_errors(never_ready) + + # 4. The lowercase chain-separator prose "The above exception was the direct + # cause of the following exception:" is never fatal on its own (no bare + # case-insensitive \bException\b match). + separator = "The above exception was the direct cause of the following exception:" + assert validation.scan_pre_teardown_errors("\n".join([_readiness_line(), separator])) == [] + assert validation.scan_post_teardown_errors(separator) == [] + + # 5. Explicit hard failures are fatal anywhere, even before readiness. + hard_before_ready = "\n".join( + [ + "CUDA error: an illegal memory access was encountered", + _readiness_line(), + ] + ) + assert validation.scan_pre_teardown_errors(hard_before_ready) + assert validation.scan_post_teardown_errors(hard_before_ready) + + +def test_only_recovered_retry_blocks_are_suppressed_not_all_pre_readiness() -> None: + validation = _load_validation("sglang_tp_validation_retryblock") + ready = _readiness_line() + + # (1) The exact HF chained traceback + adjacent retry marker + a later + # readiness marker is recovered => ignored. + recovered = "\n".join([*_HF_RETRY_TRACEBACK, ready]) + assert validation.scan_pre_teardown_errors(recovered) == [] + + # (2) The SAME chained traceback WITHOUT the retry marker is fatal, even + # though readiness follows (nothing recovers it). + no_marker = "\n".join([*_HF_RETRY_TRACEBACK[:-1], ready]) + no_marker_errors = validation.scan_pre_teardown_errors(no_marker) + assert any("Traceback" in line for line in no_marker_errors) + assert any("httpx.ConnectError:" in line for line in no_marker_errors) + + # (3) An unrelated pre-readiness traceback (not adjacent to any retry marker) + # stays fatal even though the server later becomes ready. + unrelated = "\n".join( + [ + "Traceback (most recent call last):", + ' File "/sglang/python/sglang/srt/model_executor/model_runner.py", line 1, in init', + "ValueError: unrelated startup failure", + ready, + ] + ) + unrelated_errors = validation.scan_pre_teardown_errors(unrelated) + assert any("ValueError: unrelated startup failure" in line for line in unrelated_errors) + assert any("Traceback" in line for line in unrelated_errors) + + # (4) A traceback separated from a later retry marker by another timestamped + # log record is NOT part of the retry block => fatal. + separated = "\n".join( + [ + "Traceback (most recent call last):", + ' File "/x.py", line 1, in f', + "ValueError: earlier unrelated failure", + _sglang_line(0, "Init torch distributed begin."), + _HF_RETRY_TRACEBACK[-1], # retry marker, now not adjacent to the traceback + ready, + ] + ) + separated_errors = validation.scan_pre_teardown_errors(separated) + assert any("ValueError: earlier unrelated failure" in line for line in separated_errors) + + # (5) A hard native error inside a recovered retry block is never suppressed, + # while the generic ConnectError in the same block is. + with_hard = "\n".join( + [ + "Traceback (most recent call last):", + ' File "/sglang/python/sglang/srt/configs/model_config.py", line 1143, in ', + "CUDA error: an illegal memory access was encountered", + "httpx.ConnectError: [Errno 104] Connection reset by peer", + _HF_RETRY_TRACEBACK[-1], # retry marker adjacent + ready, + ] + ) + with_hard_errors = validation.scan_pre_teardown_errors(with_hard) + assert any("CUDA error" in line for line in with_hard_errors) + assert not any("httpx.ConnectError" in line for line in with_hard_errors) + + # (6) No readiness ever => nothing is suppressed (the marked block is fatal). + never_ready = "\n".join(_HF_RETRY_TRACEBACK) + assert validation.scan_pre_teardown_errors(never_ready) + + +def test_post_teardown_window_keeps_hard_errors_fatal_after_teardown() -> None: + validation = _load_validation("sglang_tp_validation_post_hard") + hard = [ + "[foundry LOAD ERROR] cuGraphAddKernelNode FAILED for node 3 with error 400", + "CUDA error: an illegal memory access was encountered", + "gpu-host:1234:1250 [0] NCCL WARN Cuda failure", + ] + healthy = [_readiness_line()] + # Hard native faults that only surface in the full (post-teardown) log stay + # fatal even though the pre-teardown snapshot was clean. + full_text = "\n".join(healthy + hard) + pre_teardown_text = "\n".join(healthy) + + scan = validation.scan_log_text(full_text, pre_teardown_text=pre_teardown_text) + + for signature in ("[foundry LOAD ERROR]", "CUDA error", "NCCL WARN"): + assert any(signature in line for line in scan["errors"]), signature + assert len(scan["errors"]) == len(hard) + + # Post-teardown is hard-native-only: a generic Python RuntimeError that only + # surfaces after teardown is NOT counted (the SIGINT teardown emits chained + # tracebacks, and generic Python failures are gated to the pre-teardown text). + with_generic_post = "\n".join( + healthy + hard + [_sglang_line(0, "RuntimeError: kaboom during shutdown")] + ) + scan_generic = validation.scan_log_text(with_generic_post, pre_teardown_text=pre_teardown_text) + assert not any("RuntimeError" in line for line in scan_generic["errors"]) + assert ( + validation.scan_post_teardown_errors( + _sglang_line(0, "RuntimeError: kaboom during shutdown") + ) + == [] + ) + + +@pytest.mark.parametrize("returncode", [0, -2, 130]) +def test_intentional_shutdown_codes_accepted(returncode: int) -> None: + validation = _load_validation(f"sglang_tp_validation_exit_{returncode}") + + validation.validate_shutdown_returncode(returncode, died_before_teardown=False) + + +@pytest.mark.parametrize( + ("returncode", "died_before_teardown"), + [(1, False), (-9, False), (0, True), (-2, True)], +) +def test_unexpected_or_early_exit_rejected(returncode: int, died_before_teardown: bool) -> None: + validation = _load_validation(f"sglang_tp_validation_bad_{returncode}_{died_before_teardown}") + + with pytest.raises(RuntimeError, match="SGLang process"): + validation.validate_shutdown_returncode(returncode, died_before_teardown) + + +def test_stop_process_returns_and_validates_child_status(tmp_path: Path) -> None: + validation = _load_validation("sglang_tp_validation_stop") + log_path = tmp_path / "child.log" + log_file = log_path.open("w") + process = subprocess.Popen( + [ + sys.executable, + "-c", + ( + "import signal,time\n" + "signal.signal(signal.SIGINT, lambda *_: exit(0))\n" + "print('ready', flush=True)\n" + "time.sleep(30)\n" + ), + ], + stdout=log_file, + stderr=subprocess.STDOUT, + start_new_session=True, + ) + deadline = time.monotonic() + 5 + while time.monotonic() < deadline and "ready" not in log_path.read_text(): + time.sleep(0.01) + + assert validation.stop_process(process, log_file, timeout=5) == 0 + assert log_file.closed + + +def test_stop_process_rejects_child_that_died_before_teardown(tmp_path: Path) -> None: + validation = _load_validation("sglang_tp_validation_early_death") + log_file = (tmp_path / "child.log").open("w") + process = subprocess.Popen( + [sys.executable, "-c", "pass"], + stdout=log_file, + stderr=subprocess.STDOUT, + start_new_session=True, + ) + process.wait(timeout=5) + + with pytest.raises(RuntimeError, match="died before intentional teardown"): + validation.stop_process(process, log_file, timeout=1) + assert log_file.closed + + +def _spawn_launcher_with_grandchild( + tmp_path: Path, + grandchild_ignores_sigint: bool, +) -> tuple[subprocess.Popen, object]: + """Launch a group-leader launcher that spawns a grandchild in the same group. + + The launcher always exits cleanly on SIGINT. The grandchild either ignores + SIGINT (a lingering worker that must be force-killed) or exits cleanly. + """ + log_path = tmp_path / "group.log" + log_file = log_path.open("w") + grandchild_handler = "signal.SIG_IGN" if grandchild_ignores_sigint else "lambda *_: sys.exit(0)" + grandchild_body = ( + "import signal,sys,time;" + f"signal.signal(signal.SIGINT, {grandchild_handler});" + "print('gc-ready', flush=True);" + "time.sleep(60)" + ) + launcher_body = ( + "import signal, subprocess, sys, time\n" + f"subprocess.Popen([sys.executable, '-c', {grandchild_body!r}])\n" + "signal.signal(signal.SIGINT, lambda *_: sys.exit(0))\n" + "print('ready', flush=True)\n" + "time.sleep(60)\n" + ) + process = subprocess.Popen( + [sys.executable, "-c", launcher_body], + stdout=log_file, + stderr=subprocess.STDOUT, + start_new_session=True, + ) + deadline = time.monotonic() + 5 + while time.monotonic() < deadline and log_path.read_text().count("ready") < 2: + time.sleep(0.01) + return process, log_file + + +def test_stop_process_waits_for_all_group_children_to_exit(tmp_path: Path) -> None: + validation = _load_validation("sglang_tp_validation_group_clean") + process, log_file = _spawn_launcher_with_grandchild(tmp_path, grandchild_ignores_sigint=False) + + assert validation.stop_process(process, log_file, timeout=5, group_timeout=5) == 0 + assert log_file.closed + + +def test_stop_process_kills_and_fails_on_lingering_group_children(tmp_path: Path) -> None: + validation = _load_validation("sglang_tp_validation_group_linger") + process, log_file = _spawn_launcher_with_grandchild(tmp_path, grandchild_ignores_sigint=True) + pgid = os.getpgid(process.pid) + + with pytest.raises(RuntimeError, match="lingering|group child"): + validation.stop_process(process, log_file, timeout=3, group_timeout=2) + assert log_file.closed + # The force-kill must have drained the group so nothing leaks past teardown. + with pytest.raises(ProcessLookupError): + for _ in range(200): + os.killpg(pgid, 0) + time.sleep(0.02) diff --git a/tests/test_tp_validation_helpers.py b/tests/test_tp_validation_helpers.py new file mode 100644 index 00000000..9547c274 --- /dev/null +++ b/tests/test_tp_validation_helpers.py @@ -0,0 +1,434 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the Foundry project +"""CPU-only contracts for the shared TP archive evidence helpers.""" + +from __future__ import annotations + +import importlib.util +from pathlib import Path + +import pytest + +REPO_ROOT = Path(__file__).parents[1] +HELPERS_PATH = REPO_ROOT / "tests" / "tp_validation_helpers.py" + +FOUNDRY_BASE_ADDR = 0x600000000000 +FOUNDRY_REGION_END = FOUNDRY_BASE_ADDR + 256 * 1024**3 + + +def _load_helpers(name: str = "tp_validation_helpers_contract"): + assert HELPERS_PATH.exists(), f"required helpers module is missing: {HELPERS_PATH}" + spec = importlib.util.spec_from_file_location(name, HELPERS_PATH) + assert spec is not None and spec.loader is not None + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + return module + + +def _build_archive(root: Path) -> None: + """A realistic two-rank archive: root warmup_state, graphs, packed modules, stray file.""" + (root / "warmup_state.json").write_text('{"memory_pool_config": {"max_total_num_tokens": 1}}') + for rank in range(2): + rank_dir = root / f"rank_{rank}" + rank_dir.mkdir(parents=True) + (rank_dir / "graph_0_FULL_t256_r256_UX_pcN.json").write_text('{"nodes": []}') + (rank_dir / "graph_0_FULL_t256_r256_UX_pcN.cugraph").write_bytes(b"\x00\x01\x02cugraph") + (rank_dir / "graph_manifest.json").write_text('{"groups": []}') + (rank_dir / "fatbin_image_packed.img").write_bytes(b"FATBIN\x00") + (rank_dir / "final_alloc_offset.json").write_text('{"final_alloc_offset": 69736595456}') + # An unexpected file that byte-level fingerprints must still catch. + (root / "rank_0" / "unexpected_scratch.bin").write_bytes(b"stray") + + +def test_archive_byte_fingerprints_hash_every_regular_file(tmp_path: Path) -> None: + helpers = _load_helpers("tp_helpers_byte_all") + _build_archive(tmp_path) + + fingerprints = helpers.archive_byte_fingerprints(tmp_path) + + expected_relpaths = { + "warmup_state.json", + "rank_0/graph_0_FULL_t256_r256_UX_pcN.json", + "rank_0/graph_0_FULL_t256_r256_UX_pcN.cugraph", + "rank_0/graph_manifest.json", + "rank_0/fatbin_image_packed.img", + "rank_0/final_alloc_offset.json", + "rank_0/unexpected_scratch.bin", + "rank_1/graph_0_FULL_t256_r256_UX_pcN.json", + "rank_1/graph_0_FULL_t256_r256_UX_pcN.cugraph", + "rank_1/graph_manifest.json", + "rank_1/fatbin_image_packed.img", + "rank_1/final_alloc_offset.json", + } + assert set(fingerprints) == expected_relpaths + # Keys are POSIX relative paths; values are 64-hex SHA-256 digests. + assert all(len(digest) == 64 for digest in fingerprints.values()) + + +def test_archive_byte_fingerprints_detect_cugraph_and_root_mutation(tmp_path: Path) -> None: + helpers = _load_helpers("tp_helpers_byte_mut") + _build_archive(tmp_path) + baseline = helpers.archive_byte_fingerprints(tmp_path) + + # A single byte flipped in a .cugraph binary must change the map — the + # semantic graph fingerprint never hashes .cugraph, so only the byte map + # closes this same-mount reuse mutation gap. + (tmp_path / "rank_1" / "graph_0_FULL_t256_r256_UX_pcN.cugraph").write_bytes( + b"\x00\x01\x03cugraph" + ) + after_cugraph = helpers.archive_byte_fingerprints(tmp_path) + assert after_cugraph != baseline + assert ( + after_cugraph["rank_1/graph_0_FULL_t256_r256_UX_pcN.cugraph"] + != baseline["rank_1/graph_0_FULL_t256_r256_UX_pcN.cugraph"] + ) + + # A mutated root warmup_state.json (outside every rank dir) is also caught. + (tmp_path / "warmup_state.json").write_text( + '{"memory_pool_config": {"max_total_num_tokens": 2}}' + ) + after_root = helpers.archive_byte_fingerprints(tmp_path) + assert after_root["warmup_state.json"] != baseline["warmup_state.json"] + + +def test_semantic_and_byte_fingerprints_are_separate(tmp_path: Path) -> None: + helpers = _load_helpers("tp_helpers_separate") + _build_archive(tmp_path) + + semantic = helpers.archive_fingerprints(tmp_path, 2, FOUNDRY_BASE_ADDR, FOUNDRY_REGION_END) + byte_map = helpers.archive_byte_fingerprints(tmp_path) + + # Semantic fingerprints are per-rank and do NOT cover the root + # warmup_state.json, the .cugraph binaries, or unexpected files. + assert set(semantic) == {"0", "1"} + assert not any(name.endswith(".cugraph") for name in semantic["0"]) + # The byte map is a flat relpath->digest map that does cover them. + assert "warmup_state.json" in byte_map + assert "rank_0/graph_0_FULL_t256_r256_UX_pcN.cugraph" in byte_map + assert "rank_0/unexpected_scratch.bin" in byte_map + + +EXPECTED_GRAPH_COUNT = 64 + + +def _build_exact_archive(root: Path, tp_size: int = 2, graph_count: int = EXPECTED_GRAPH_COUNT): + """Realistic two-rank archive: ``graph_count`` paired JSON/cugraph per rank, + contiguous indices ``0..graph_count-1``, all required rank files, root + ``warmup_state.json``, and one unexpected file the byte map must still see. + + Returns the ``{rank: {index: stem}}`` map so tests can mutate a single + known stem/index. + """ + (root / "warmup_state.json").write_text('{"memory_pool_config": {"max_total_num_tokens": 1}}') + stems: dict[str, dict[int, str]] = {} + for rank in range(tp_size): + rank_dir = root / f"rank_{rank}" + rank_dir.mkdir(parents=True) + rank_stems: dict[int, str] = {} + for index in range(graph_count): + # Descending capture order: capture_index i -> batch size count-i. + size = graph_count - index + stem = f"graph_{index}_FULL_t{size}_r{size}_UX_pcN" + rank_stems[index] = stem + (rank_dir / f"{stem}.json").write_text('{"nodes": []}') + (rank_dir / f"{stem}.cugraph").write_bytes(b"\x00cugraph" + bytes([index % 256])) + stems[str(rank)] = rank_stems + (rank_dir / "graph_manifest.json").write_text('{"groups": []}') + (rank_dir / "fatbin_image_packed.img").write_bytes(b"FATBIN\x00") + (rank_dir / "fatbin_entrypoint_packed.txt").write_text("entrypoint\n") + (rank_dir / "final_alloc_offset.json").write_text('{"final_alloc_offset": 69736595456}') + (root / "rank_0" / "unexpected_scratch.bin").write_bytes(b"stray") + return stems + + +def test_validate_archive_inventory_accepts_exact_paired_archive(tmp_path: Path) -> None: + helpers = _load_helpers("tp_helpers_inv_ok") + _build_exact_archive(tmp_path) + + inventory = helpers.validate_archive_inventory(tmp_path, 2, EXPECTED_GRAPH_COUNT) + + assert set(inventory) == {"0", "1"} + for rank in ("0", "1"): + assert inventory[rank]["graph_count"] == EXPECTED_GRAPH_COUNT + assert sorted(inventory[rank]["indices"]) == list(range(EXPECTED_GRAPH_COUNT)) + + +def test_validate_archive_inventory_rejects_missing_graph(tmp_path: Path) -> None: + helpers = _load_helpers("tp_helpers_inv_missing") + stems = _build_exact_archive(tmp_path) + (tmp_path / "rank_1" / f"{stems['1'][63]}.cugraph").unlink() + + with pytest.raises(ValueError, match="rank_1"): + helpers.validate_archive_inventory(tmp_path, 2, EXPECTED_GRAPH_COUNT) + + +def test_validate_archive_inventory_rejects_duplicate_index(tmp_path: Path) -> None: + helpers = _load_helpers("tp_helpers_inv_dup") + _build_exact_archive(tmp_path) + (tmp_path / "rank_0" / "graph_0_FULL_t9_r9_UX_pcN.json").write_text('{"nodes": []}') + (tmp_path / "rank_0" / "graph_0_FULL_t9_r9_UX_pcN.cugraph").write_bytes(b"dup") + + with pytest.raises(ValueError, match="(?i)duplicate|index 0"): + helpers.validate_archive_inventory(tmp_path, 2, EXPECTED_GRAPH_COUNT) + + +def test_validate_archive_inventory_rejects_noncontiguous_indices(tmp_path: Path) -> None: + helpers = _load_helpers("tp_helpers_inv_gap") + stems = _build_exact_archive(tmp_path) + for suffix in (".json", ".cugraph"): + original = tmp_path / "rank_0" / f"{stems['0'][63]}{suffix}" + original.rename(tmp_path / "rank_0" / f"graph_64_FULL_t99_r99_UX_pcN{suffix}") + + with pytest.raises(ValueError, match="(?i)contiguous|index|0\\.\\.63|range"): + helpers.validate_archive_inventory(tmp_path, 2, EXPECTED_GRAPH_COUNT) + + +def test_validate_archive_inventory_rejects_mismatched_stems(tmp_path: Path) -> None: + helpers = _load_helpers("tp_helpers_inv_stem") + stems = _build_exact_archive(tmp_path) + mismatch = tmp_path / "rank_1" / f"{stems['1'][5]}.cugraph" + mismatch.rename(tmp_path / "rank_1" / "graph_5_FULL_t999_r999_UX_pcN.cugraph") + + with pytest.raises(ValueError, match="(?i)stem|mismatch"): + helpers.validate_archive_inventory(tmp_path, 2, EXPECTED_GRAPH_COUNT) + + +def test_validate_archive_inventory_rejects_missing_required_file(tmp_path: Path) -> None: + helpers = _load_helpers("tp_helpers_inv_reqfile") + _build_exact_archive(tmp_path) + (tmp_path / "rank_0" / "fatbin_entrypoint_packed.txt").unlink() + + with pytest.raises(ValueError, match="fatbin_entrypoint_packed.txt"): + helpers.validate_archive_inventory(tmp_path, 2, EXPECTED_GRAPH_COUNT) + + +def test_validate_archive_inventory_rejects_missing_root_state(tmp_path: Path) -> None: + helpers = _load_helpers("tp_helpers_inv_rootstate") + _build_exact_archive(tmp_path) + (tmp_path / "warmup_state.json").unlink() + + with pytest.raises(ValueError, match="warmup_state.json"): + helpers.validate_archive_inventory(tmp_path, 2, EXPECTED_GRAPH_COUNT) + + +def test_collect_archive_snapshot_returns_every_map(tmp_path: Path) -> None: + helpers = _load_helpers("tp_helpers_snapshot") + _build_exact_archive(tmp_path) + + snapshot = helpers.collect_archive_snapshot( + tmp_path, + 2, + EXPECTED_GRAPH_COUNT, + FOUNDRY_BASE_ADDR, + FOUNDRY_REGION_END, + ) + + assert set(snapshot) == { + "inventory", + "semantic_fingerprints", + "byte_fingerprints", + "graph_counts", + "offsets", + } + assert snapshot["graph_counts"] == {"0": EXPECTED_GRAPH_COUNT, "1": EXPECTED_GRAPH_COUNT} + assert snapshot["offsets"] == {"0": 69736595456, "1": 69736595456} + # Semantic map is per-rank and must include the packed entrypoint. + assert set(snapshot["semantic_fingerprints"]) == {"0", "1"} + assert "fatbin_entrypoint_packed.txt" in snapshot["semantic_fingerprints"]["0"] + # Byte map is a flat relpath map covering the root state and unexpected files. + assert "warmup_state.json" in snapshot["byte_fingerprints"] + assert "rank_0/unexpected_scratch.bin" in snapshot["byte_fingerprints"] + + +def test_collect_archive_snapshot_rejects_incomplete_archive(tmp_path: Path) -> None: + helpers = _load_helpers("tp_helpers_snapshot_bad") + stems = _build_exact_archive(tmp_path) + (tmp_path / "rank_0" / f"{stems['0'][0]}.json").unlink() + + with pytest.raises(ValueError): + helpers.collect_archive_snapshot( + tmp_path, + 2, + EXPECTED_GRAPH_COUNT, + FOUNDRY_BASE_ADDR, + FOUNDRY_REGION_END, + ) + + +def test_semantic_fingerprints_include_packed_entrypoint(tmp_path: Path) -> None: + helpers = _load_helpers("tp_helpers_entrypoint") + _build_exact_archive(tmp_path) + + semantic = helpers.archive_fingerprints(tmp_path, 2, FOUNDRY_BASE_ADDR, FOUNDRY_REGION_END) + + for rank in ("0", "1"): + assert "fatbin_entrypoint_packed.txt" in semantic[rank] + + +def test_validate_archive_inventory_rejects_file_symlink(tmp_path: Path) -> None: + helpers = _load_helpers("tp_helpers_symlink_file") + _build_exact_archive(tmp_path) + # An otherwise-complete archive with a single stray symlink whose name does + # not collide with the graph inventory: only symlink rejection can catch it. + (tmp_path / "rank_0" / "alias_blob.bin").symlink_to( + tmp_path / "rank_0" / "fatbin_image_packed.img" + ) + + with pytest.raises(ValueError, match="(?i)symlink"): + helpers.validate_archive_inventory(tmp_path, 2, EXPECTED_GRAPH_COUNT) + + +def test_validate_archive_inventory_rejects_directory_symlink(tmp_path: Path) -> None: + helpers = _load_helpers("tp_helpers_symlink_dir") + _build_exact_archive(tmp_path) + (tmp_path / "rank_0_alias").symlink_to(tmp_path / "rank_0", target_is_directory=True) + + with pytest.raises(ValueError, match="(?i)symlink"): + helpers.validate_archive_inventory(tmp_path, 2, EXPECTED_GRAPH_COUNT) + + +def test_validate_archive_inventory_rejects_symlinked_root(tmp_path: Path) -> None: + helpers = _load_helpers("tp_helpers_symlink_root") + real_root = tmp_path / "real_archive" + real_root.mkdir() + _build_exact_archive(real_root) + linked_root = tmp_path / "linked_archive" + linked_root.symlink_to(real_root, target_is_directory=True) + + with pytest.raises(ValueError, match="(?i)symlink"): + helpers.validate_archive_inventory(linked_root, 2, EXPECTED_GRAPH_COUNT) + + +def test_reject_symlinks_rejects_symlinked_root(tmp_path: Path) -> None: + helpers = _load_helpers("tp_helpers_reject_root") + real_root = tmp_path / "real" + real_root.mkdir() + linked_root = tmp_path / "linked" + linked_root.symlink_to(real_root, target_is_directory=True) + + with pytest.raises(ValueError, match="(?i)symlink"): + helpers.reject_symlinks(linked_root) + + +def test_collect_archive_snapshot_rejects_symlink(tmp_path: Path) -> None: + helpers = _load_helpers("tp_helpers_snapshot_symlink") + _build_exact_archive(tmp_path) + (tmp_path / "rank_1" / "sneaky_link.bin").symlink_to( + tmp_path / "rank_1" / "fatbin_image_packed.img" + ) + + with pytest.raises(ValueError, match="(?i)symlink"): + helpers.collect_archive_snapshot( + tmp_path, + 2, + EXPECTED_GRAPH_COUNT, + FOUNDRY_BASE_ADDR, + FOUNDRY_REGION_END, + ) + + +def test_inventory_and_byte_map_cover_same_regular_file_universe(tmp_path: Path) -> None: + helpers = _load_helpers("tp_helpers_universe") + _build_exact_archive(tmp_path) + + # A symlink-free archive validates, and its byte map is exactly every regular + # file on disk (no symlinks skipped, so no hidden divergence). + helpers.validate_archive_inventory(tmp_path, 2, EXPECTED_GRAPH_COUNT) + byte_map = helpers.archive_byte_fingerprints(tmp_path) + + on_disk_regular_files = { + path.relative_to(tmp_path).as_posix() + for path in tmp_path.rglob("*") + if path.is_file() and not path.is_symlink() + } + assert set(byte_map) == on_disk_regular_files + + +def _canonical_stem(index: int) -> str: + """The source-confirmed current profile descriptor for a graph index. + + Descending capture order: capture_index ``index`` -> batch size + ``EXPECTED_GRAPH_COUNT - index`` (index 0 is the largest, t64). + """ + size = EXPECTED_GRAPH_COUNT - index + return f"graph_{index}_FULL_t{size}_r{size}_UX_pcN" + + +def test_validate_archive_inventory_accepts_matching_expected_stems(tmp_path: Path) -> None: + helpers = _load_helpers("tp_helpers_stems_ok") + _build_exact_archive(tmp_path) + + # The exact profile fixture matches the canonical descriptor for every index, + # so passing the profile validator still accepts it. + inventory = helpers.validate_archive_inventory( + tmp_path, 2, EXPECTED_GRAPH_COUNT, expected_stems=_canonical_stem + ) + assert set(inventory) == {"0", "1"} + + +def test_validate_archive_inventory_rejects_wrong_descriptor_with_exact_count( + tmp_path: Path, +) -> None: + helpers = _load_helpers("tp_helpers_stems_bad") + stems = _build_exact_archive(tmp_path) + # Rename index 7's JSON+cugraph to a paired but wrong descriptor (still 64 + # contiguous paired graphs); only the profile validator can reject it. + wrong = "graph_7_FULL_t999_r999_UX_pcN" + for suffix in (".json", ".cugraph"): + (tmp_path / "rank_0" / f"{stems['0'][7]}{suffix}").rename( + tmp_path / "rank_0" / f"{wrong}{suffix}" + ) + + # Without the profile validator the archive still passes (count/index/pairing + # are intact) — proving the descriptor check is exactly what closes the gap. + helpers.validate_archive_inventory(tmp_path, 2, EXPECTED_GRAPH_COUNT) + + with pytest.raises(ValueError, match="(?i)descriptor|expected"): + helpers.validate_archive_inventory( + tmp_path, 2, EXPECTED_GRAPH_COUNT, expected_stems=_canonical_stem + ) + + +def test_validate_archive_inventory_without_expected_stems_is_descriptor_agnostic( + tmp_path: Path, +) -> None: + helpers = _load_helpers("tp_helpers_stems_none") + # A non-canonical but paired/contiguous archive (mirrors SGLang's differing + # profile). Default use — no ``expected_stems`` — must accept it unchanged. + (tmp_path / "warmup_state.json").write_text( + '{"memory_pool_config": {"max_total_num_tokens": 1}}' + ) + for rank in range(2): + rank_dir = tmp_path / f"rank_{rank}" + rank_dir.mkdir(parents=True) + for index in range(EXPECTED_GRAPH_COUNT): + stem = f"graph_{index}_SOMETHING_ELSE_t{index}" + (rank_dir / f"{stem}.json").write_text('{"nodes": []}') + (rank_dir / f"{stem}.cugraph").write_bytes(b"cg") + (rank_dir / "graph_manifest.json").write_text('{"groups": []}') + (rank_dir / "fatbin_image_packed.img").write_bytes(b"FATBIN\x00") + (rank_dir / "fatbin_entrypoint_packed.txt").write_text("entrypoint\n") + (rank_dir / "final_alloc_offset.json").write_text('{"final_alloc_offset": 1}') + + inventory = helpers.validate_archive_inventory(tmp_path, 2, EXPECTED_GRAPH_COUNT) + assert inventory["0"]["stems"]["3"] == "graph_3_SOMETHING_ELSE_t3" + + +def test_collect_archive_snapshot_forwards_expected_stems(tmp_path: Path) -> None: + helpers = _load_helpers("tp_helpers_snapshot_stems") + stems = _build_exact_archive(tmp_path) + wrong = "graph_2_FULL_t42_r42_UX_pcN" + for suffix in (".json", ".cugraph"): + (tmp_path / "rank_1" / f"{stems['1'][2]}{suffix}").rename( + tmp_path / "rank_1" / f"{wrong}{suffix}" + ) + + with pytest.raises(ValueError, match="(?i)descriptor|expected"): + helpers.collect_archive_snapshot( + tmp_path, + 2, + EXPECTED_GRAPH_COUNT, + FOUNDRY_BASE_ADDR, + FOUNDRY_REGION_END, + expected_stems=_canonical_stem, + ) diff --git a/tests/test_vllm_source_contract.py b/tests/test_vllm_source_contract.py new file mode 100644 index 00000000..68d25a30 --- /dev/null +++ b/tests/test_vllm_source_contract.py @@ -0,0 +1,82 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the Foundry project +"""CPU-only source contracts for vLLM import and cleanup documentation.""" + +from __future__ import annotations + +import ast +from pathlib import Path + +REPO_ROOT = Path(__file__).parents[1] +HOOKS_PATH = REPO_ROOT / "python" / "foundry" / "integration" / "vllm" / "hooks.py" +HARNESS_PATH = REPO_ROOT / "tests" / "modal_vllm_tp.py" + + +def _function_node(tree: ast.Module, name: str) -> ast.FunctionDef: + for node in tree.body: + if isinstance(node, ast.FunctionDef) and node.name == name: + return node + raise AssertionError(f"function not found: {name}") + + +def test_standard_library_hook_imports_are_module_scoped() -> None: + tree = ast.parse(HOOKS_PATH.read_text()) + top_level_import_lines = { + node.lineno for node in tree.body if isinstance(node, (ast.Import, ast.ImportFrom)) + } + misplaced = [] + for node in ast.walk(tree): + if not isinstance(node, (ast.Import, ast.ImportFrom)): + continue + imports_target = ( + isinstance(node, ast.Import) + and any(alias.name in {"os", "time"} for alias in node.names) + ) or ( + isinstance(node, ast.ImportFrom) + and ( + ( + node.module == "contextlib" + and any(alias.name == "ExitStack" for alias in node.names) + ) + or ( + node.module == "unittest.mock" + and any( + alias.name == "patch" and alias.asname == "mock_patch" + for alias in node.names + ) + ) + ) + ) + if imports_target and node.lineno not in top_level_import_lines: + misplaced.append(node.lineno) + assert misplaced == [] + + +def test_engine_version_imports_remain_lazy_for_hook_installation() -> None: + source = HOOKS_PATH.read_text() + tree = ast.parse(source) + function = _function_node(tree, "_patch_init_worker_distributed_environment") + assert any( + isinstance(node, ast.ImportFrom) and node.module == "vllm.v1.worker" + for node in ast.walk(function) + ) + assert "Engine-version-specific vLLM imports intentionally stay lazy" in source + assert "only after install_hooks selects the active integration" in source + + +def test_vllm_cleanup_remote_docstring_references_task_seven() -> None: + tree = ast.parse(HARNESS_PATH.read_text()) + function = _function_node(tree, "cleanup_published_run") + docstring = ast.get_docstring(function) + assert docstring is not None + assert "Task 7" in docstring + assert "Task 6" not in docstring + + +def test_vllm_delete_docstring_matches_orchestrated_cleanup() -> None: + tree = ast.parse(HARNESS_PATH.read_text()) + function = _function_node(tree, "delete_published_run") + docstring = ast.get_docstring(function) + assert docstring is not None + assert "called by ``cleanup_published_run``" in docstring + assert "NOT called by this harness" not in docstring diff --git a/tests/test_vllm_tp_recipe.py b/tests/test_vllm_tp_recipe.py new file mode 100644 index 00000000..71a8994e --- /dev/null +++ b/tests/test_vllm_tp_recipe.py @@ -0,0 +1,163 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the Foundry project +"""CPU-only contract tests for the vLLM tensor-parallel recipe.""" + +from __future__ import annotations + +import json +import os +import subprocess +from pathlib import Path + +import pytest +import tomllib + +REPO_ROOT = Path(__file__).parents[1] +RECIPE_DIR = REPO_ROOT / "recipe" / "vllm" +SERVE_SCRIPT = RECIPE_DIR / "serve_qwen3-8b_tp.sh" + + +def _run_recipe( + tmp_path: Path, + mode: str | None, + *, + model_revision: str = "b968826d9c46dd6066d109eabc6255188de91218", +) -> tuple[dict[str, str], list[str]]: + bin_dir = tmp_path / "bin" + bin_dir.mkdir() + fake_vllm = bin_dir / "vllm" + fake_vllm.write_text( + '#!/usr/bin/env bash\nenv > "$CAPTURE_ENV"\nprintf "%s\\n" "$@" > "$CAPTURE_ARGS"\n' + ) + fake_vllm.chmod(0o755) + + env_path = tmp_path / "env" + args_path = tmp_path / "args" + env = os.environ.copy() + env.update( + { + "CAPTURE_ENV": str(env_path), + "CAPTURE_ARGS": str(args_path), + "PATH": f"{bin_dir}:{env['PATH']}", + "VLLM_MODEL_REVISION": model_revision, + } + ) + + command = ["bash", str(SERVE_SCRIPT), "2"] + if mode is not None: + command.append(mode) + subprocess.run(command, check=True, env=env, capture_output=True, text=True) + + captured_env = dict(line.split("=", 1) for line in env_path.read_text().splitlines()) + return captured_env, args_path.read_text().splitlines() + + +def _compilation_config(args: list[str]) -> dict: + assert args.count("--compilation-config") == 1 + config_index = args.index("--compilation-config") + 1 + return json.loads(args[config_index]) + + +@pytest.mark.parametrize("mode", [None, "--save", "--load"]) +def test_tp_recipe_uses_same_collective_transport_for_every_phase( + tmp_path: Path, + mode: str | None, +) -> None: + env, args = _run_recipe(tmp_path, mode) + + assert env["PYTORCH_CUDA_ALLOC_CONF"] == "expandable_segments:False" + assert env["NCCL_CUMEM_ENABLE"] == "0" + assert env["NCCL_NVLS_ENABLE"] == "0" + assert env["VLLM_DISABLE_PYNCCL"] == "0" + assert env["VLLM_ALLREDUCE_USE_SYMM_MEM"] == "0" + assert env["VLLM_ALLREDUCE_USE_FLASHINFER"] == "0" + assert env["VLLM_USE_NCCL_SYMM_MEM"] == "0" + assert env["VLLM_USE_V2_MODEL_RUNNER"] == "0" + assert env["VLLM_WORKER_MULTIPROC_METHOD"] == "spawn" + assert "--disable-custom-all-reduce" in args + config = _compilation_config(args) + assert config["cudagraph_mode"] == "FULL_DECODE_ONLY" + assert config["cudagraph_num_of_warmups"] == 0 + assert config["pass_config"]["fuse_allreduce_rms"] is False + + +@pytest.mark.parametrize( + ("mode", "config_name"), + [ + ("--save", "foundry_save.toml"), + ("--load", "foundry_load.toml"), + ], +) +def test_tp_recipe_selects_matching_foundry_config( + tmp_path: Path, + mode: str, + config_name: str, +) -> None: + _env, args = _run_recipe(tmp_path, mode) + + config = _compilation_config(args) + assert Path(config["graph_extension_config_path"]) == RECIPE_DIR / config_name + assert args[0] == "serve" + assert args[args.index("--tensor-parallel-size") + 1] == "2" + assert args[args.index("--distributed-executor-backend") + 1] == "mp" + + +def test_tp_baseline_does_not_enable_foundry(tmp_path: Path) -> None: + _env, args = _run_recipe(tmp_path, None) + + assert "graph_extension_config_path" not in _compilation_config(args) + + +def test_tp_recipe_pins_model_revision(tmp_path: Path) -> None: + _env, args = _run_recipe(tmp_path, None) + + assert args.count("--revision") == 1 + assert args[args.index("--revision") + 1] == "b968826d9c46dd6066d109eabc6255188de91218" + + +@pytest.mark.parametrize( + "revision", + [ + "main", + "b968826", + "g" * 40, + "b968826d9c46dd6066d109eabc6255188de91218 --served-model-name injected", + ], +) +def test_tp_recipe_rejects_non_sha_model_revision(tmp_path: Path, revision: str) -> None: + with pytest.raises(subprocess.CalledProcessError) as excinfo: + _run_recipe(tmp_path, None, model_revision=revision) + + assert "full 40-hex commit SHA" in excinfo.value.stderr + + +def test_tp_save_and_load_configs_are_symmetric() -> None: + save = tomllib.loads((RECIPE_DIR / "foundry_save.toml").read_text()) + load = tomllib.loads((RECIPE_DIR / "foundry_load.toml").read_text()) + + assert save.pop("mode") == "save" + assert load.pop("mode") == "load" + assert save == load + assert save["scratch_space_size"] == "4096MB" + + +@pytest.mark.parametrize( + "arguments", + [ + [], + ["1"], + ["3"], + ["2", "--invalid"], + ["2", "--save", "unexpected"], + ["2", "--load", "unexpected"], + ], +) +def test_tp_recipe_rejects_unvalidated_profile(arguments: list[str]) -> None: + result = subprocess.run( + ["bash", str(SERVE_SCRIPT), *arguments], + capture_output=True, + text=True, + ) + + assert result.returncode != 0 + assert "TP=2" in result.stderr diff --git a/tests/test_vllm_tp_validation.py b/tests/test_vllm_tp_validation.py new file mode 100644 index 00000000..5b24bb48 --- /dev/null +++ b/tests/test_vllm_tp_validation.py @@ -0,0 +1,432 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the Foundry project +"""CPU-only contracts for the vLLM TP validation evidence.""" + +from __future__ import annotations + +import ast +import importlib.util +import json +import subprocess +import sys +import time +from pathlib import Path + +import pytest + +REPO_ROOT = Path(__file__).parents[1] +HARNESS_PATH = REPO_ROOT / "tests" / "modal_vllm_tp.py" +HELPERS_PATH = REPO_ROOT / "tests" / "tp_validation_helpers.py" +VALIDATION_PATH = REPO_ROOT / "tests" / "vllm_tp_validation.py" +RESTORE_PROGRESS_PATH = ( + REPO_ROOT / "python" / "foundry" / "integration" / "vllm" / "restore_progress.py" +) +HOOK_PATH = REPO_ROOT / "csrc" / "hook.cpp" +GRAPH_OPS_PATH = REPO_ROOT / "python" / "foundry" / "integration" / "vllm" / "graph_ops.py" +HOOKS_PATH = REPO_ROOT / "python" / "foundry" / "integration" / "vllm" / "hooks.py" +RUNTIME_PATH = REPO_ROOT / "python" / "foundry" / "integration" / "vllm" / "runtime.py" +MODEL_REVISION = "b968826d9c46dd6066d109eabc6255188de91218" + + +def _load_module(path: Path, name: str): + assert path.exists(), f"required validation module is missing: {path}" + spec = importlib.util.spec_from_file_location(name, path) + assert spec is not None and spec.loader is not None + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + return module + + +def test_commit_refs_must_be_full_immutable_shas() -> None: + validation = _load_module(VALIDATION_PATH, "vllm_tp_validation_contract") + sha = "0123456789abcdef0123456789abcdef01234567" + + assert validation.require_commit_sha("FOUNDRY_REF", sha) == sha + for invalid in (None, "", "main", "abc123", "g" * 40): + with pytest.raises(ValueError, match="40-hex commit SHA"): + validation.require_commit_sha("FOUNDRY_REF", invalid) + + +def test_error_scan_and_rank_restoration_are_positive_proof() -> None: + validation = _load_module(VALIDATION_PATH, "vllm_tp_validation_scan") + text = "\n".join( + [ + "[foundry] LOAD restoration complete rank_0: 3/3 graphs", + "[foundry] LOAD restoration complete rank_1: 3/3 graphs", + "ERROR worker process died", + "RuntimeError: CUDA error", + "Traceback (most recent call last):", + "Exception: NCCL WARN MMU fault Xid segmentation fault", + ] + ) + + scan = validation.scan_log_text(text) + + assert scan["load_restoration_completions"] == { + "0": [[3, 3]], + "1": [[3, 3]], + } + # No live LOAD offset markers in this fixture. + assert scan["load_final_alloc_offsets"] == {} + errors = " ".join(scan["errors"]) + for signature in ( + "ERROR", + "RuntimeError", + "Traceback", + "Exception", + "NCCL WARN", + "MMU fault", + "Xid", + "segmentation fault", + ): + assert signature.lower() in errors.lower() + + +def test_scan_parses_live_load_final_alloc_offsets_per_rank() -> None: + validation = _load_module(VALIDATION_PATH, "vllm_tp_validation_live_offset") + text = "\n".join( + [ + "[foundry] LOAD live final_alloc_offset rank_0: 69736595456", + "[foundry] LOAD live final_alloc_offset rank_1: 69736595456", + "[foundry] LOAD restoration complete rank_0: 64/64 graphs", + "[foundry] LOAD restoration complete rank_1: 64/64 graphs", + ] + ) + + scan = validation.scan_log_text(text) + + # Each rank's markers are collected as a list so the harness can require + # exactly one per rank and compare it with SAVE2's rank offset. + assert scan["load_final_alloc_offsets"] == { + "0": [69736595456], + "1": [69736595456], + } + assert scan["load_restoration_completions"] == { + "0": [[64, 64]], + "1": [[64, 64]], + } + + +def test_scan_collects_duplicate_live_offset_markers_per_rank() -> None: + validation = _load_module(VALIDATION_PATH, "vllm_tp_validation_dup_offset") + text = "\n".join( + [ + "[foundry] LOAD live final_alloc_offset rank_0: 111", + "[foundry] LOAD live final_alloc_offset rank_0: 222", + ] + ) + + scan = validation.scan_log_text(text) + + # Duplicates are retained so the harness can reject "not exactly one". + assert scan["load_final_alloc_offsets"] == {"0": [111, 222]} + + +def test_scan_parses_nccl_runtime_version_from_child_log() -> None: + validation = _load_module(VALIDATION_PATH, "vllm_tp_validation_nccl_version") + + scan = validation.scan_log_text("[0] NCCL INFO NCCL version 2.28.9+cuda13.0\n") + + assert scan["nccl_version"] == "2.28.9+cuda13.0" + + +@pytest.mark.parametrize("returncode", [0, -2, 130]) +def test_expected_intentional_shutdown_codes_are_accepted(returncode: int) -> None: + validation = _load_module(VALIDATION_PATH, f"vllm_tp_validation_exit_{returncode}") + + validation.validate_shutdown_returncode(returncode, died_before_teardown=False) + + +@pytest.mark.parametrize( + ("returncode", "died_before_teardown"), + [(1, False), (-9, False), (0, True), (-2, True)], +) +def test_unexpected_or_early_process_exit_is_rejected( + returncode: int, + died_before_teardown: bool, +) -> None: + validation = _load_module( + VALIDATION_PATH, + f"vllm_tp_validation_bad_exit_{returncode}_{died_before_teardown}", + ) + + with pytest.raises(RuntimeError, match="vLLM process"): + validation.validate_shutdown_returncode(returncode, died_before_teardown) + + +def test_stop_process_returns_and_validates_child_status(tmp_path: Path) -> None: + validation = _load_module(VALIDATION_PATH, "vllm_tp_validation_stop") + log_path = tmp_path / "child.log" + log_file = log_path.open("w") + process = subprocess.Popen( + [ + sys.executable, + "-c", + ( + "import signal,time\n" + "signal.signal(signal.SIGINT, lambda *_: exit(0))\n" + "print('ready', flush=True)\n" + "time.sleep(30)\n" + ), + ], + stdout=log_file, + stderr=subprocess.STDOUT, + start_new_session=True, + ) + deadline = time.monotonic() + 5 + while time.monotonic() < deadline and "ready" not in log_path.read_text(): + time.sleep(0.01) + + assert validation.stop_process(process, log_file, timeout=1) == 0 + assert log_file.closed + + +def test_stop_process_rejects_child_that_died_before_teardown(tmp_path: Path) -> None: + validation = _load_module(VALIDATION_PATH, "vllm_tp_validation_early_death") + log_file = (tmp_path / "child.log").open("w") + process = subprocess.Popen( + [sys.executable, "-c", "pass"], + stdout=log_file, + stderr=subprocess.STDOUT, + start_new_session=True, + ) + process.wait(timeout=5) + + with pytest.raises(RuntimeError, match="died before intentional teardown"): + validation.stop_process(process, log_file, timeout=1) + assert log_file.closed + + +def test_completion_evidence_uses_exact_text_and_token_strings() -> None: + validation = _load_module(VALIDATION_PATH, "vllm_tp_validation_output") + response = { + "choices": [ + { + "text": " Paris", + "logprobs": { + "tokens": [" Paris"], + "token_logprobs": [-0.01], + }, + } + ] + } + + assert validation.completion_evidence(response) == { + "text": " Paris", + "token_strings": [" Paris"], + } + + +@pytest.mark.parametrize( + "text,tokens", + [ + ("", [" Paris"]), + (" Paris", []), + (" Paris", [""]), + ], +) +def test_completion_evidence_rejects_empty_text_or_token_strings( + text: str, + tokens: list[str], +) -> None: + validation = _load_module(VALIDATION_PATH, "vllm_tp_validation_nonempty_output") + response = { + "choices": [ + { + "text": text, + "logprobs": {"tokens": tokens}, + } + ] + } + + with pytest.raises(RuntimeError, match="nonempty"): + validation.completion_evidence(response) + + +def test_archive_fingerprints_use_shared_semantic_normalization(tmp_path: Path) -> None: + helpers = _load_module(HELPERS_PATH, "tp_validation_helpers_contract") + base = 0x600000000000 + region_end = base + 256 * 1024**3 + rank_dir = tmp_path / "rank_0" + rank_dir.mkdir() + + def write_graph(path: Path, *, generator_id: int, seed: int, extra: int, pointer: int) -> None: + path.write_text( + json.dumps( + { + "generators": [{"id": generator_id, "seed": seed, "state": "stable"}], + "nodes": [ + { + "id": 7, + "type": "KernelNode", + "params": { + "kernelParams": [ + { + "index": 0, + "value_hex": pointer.to_bytes(8, "little").hex(), + } + ], + "extra_argBuffer_hex": "", + "extra": [extra], + }, + } + ], + } + ) + ) + + graph = rank_dir / "graph_0_FULL.json" + same_semantics = tmp_path / "same.json" + changed_pointer = tmp_path / "changed.json" + write_graph(graph, generator_id=1, seed=2, extra=3, pointer=base + 0x1000) + write_graph(same_semantics, generator_id=9, seed=10, extra=11, pointer=base + 0x1000) + write_graph(changed_pointer, generator_id=9, seed=10, extra=11, pointer=base + 0x2000) + (rank_dir / "graph_manifest.json").write_text('{"graphs": ["graph_0_FULL.json"]}') + (rank_dir / "fatbin_image_packed.img").write_bytes(b"packed module") + (rank_dir / "final_alloc_offset.json").write_text('{"final_alloc_offset": 1}') + + assert helpers.semantic_graph_sha256(graph, base, region_end) == helpers.semantic_graph_sha256( + same_semantics, + base, + region_end, + ) + assert helpers.semantic_graph_sha256( + graph, + base, + region_end, + ) != helpers.semantic_graph_sha256(changed_pointer, base, region_end) + fingerprints = helpers.archive_fingerprints(tmp_path, 1, base, region_end) + assert set(fingerprints["0"]) == { + "fatbin_image_packed.img", + "final_alloc_offset.json", + "graph_0_FULL.json", + "graph_manifest.json", + } + + +def test_restoration_progress_counts_unique_indices_only() -> None: + progress_module = _load_module(RESTORE_PROGRESS_PATH, "vllm_restore_progress_contract") + progress = progress_module.LoadRestorationProgress(rank=1, total=2) + + first = progress.record(0) + duplicate = progress.record(0) + final = progress.record(1) + + assert (first.count, first.is_new, first.completed_now) == (1, True, False) + assert (duplicate.count, duplicate.is_new, duplicate.completed_now) == (1, False, False) + assert (final.count, final.is_new, final.completed_now) == (2, True, True) + assert progress.restored_indices == frozenset({0, 1}) + + +def test_restoration_rank_comes_from_rank_workspace() -> None: + progress_module = _load_module(RESTORE_PROGRESS_PATH, "vllm_restore_progress_rank") + + assert progress_module.rank_from_workspace_dir("/archive/run/rank_7") == 7 + with pytest.raises(ValueError, match="rank_N"): + progress_module.rank_from_workspace_dir("/archive/run/shared") + + +def test_hook_and_graph_ops_expose_required_native_evidence_contracts() -> None: + hook = HOOK_PATH.read_text() + graph_ops = GRAPH_OPS_PATH.read_text() + + assert "alignment == 0 ? kAllocAlignment : alignment" in hook + assert hook.count("align_to(tls_storage.current_alloc_base_addr, effective_alignment)") >= 2 + assert "align_to(tls_storage.current_vmm_reserve_addr, effective_alignment)" in hook + assert "align_to(*ptr + size, effective_alignment)" in hook + assert "event.alignment = alignment;" in hook + assert "LOAD graph finalized rank_%d index=%d restored=%d/%d" in graph_ops + assert "LOAD restoration complete rank_%d: %d/%d graphs" in graph_ops + + +def _capture_model_patched(tree: ast.Module) -> ast.FunctionDef: + for node in ast.walk(tree): + if isinstance(node, ast.FunctionDef) and node.name == "_patch_capture_model": + for sub in ast.walk(node): + if isinstance(sub, ast.FunctionDef) and sub.name == "patched": + return sub + raise AssertionError("hooks.py has no _patch_capture_model.patched function") + + +def _is_load_mode_test(test: ast.expr) -> bool: + return ( + isinstance(test, ast.Compare) + and isinstance(test.left, ast.Name) + and test.left.id == "mode" + and len(test.ops) == 1 + and isinstance(test.ops[0], ast.Eq) + and isinstance(test.comparators[0], ast.Attribute) + and test.comparators[0].attr == "LOAD" + ) + + +def _offset_calls(node: ast.AST) -> list[ast.Call]: + return [ + sub + for sub in ast.walk(node) + if isinstance(sub, ast.Call) + and isinstance(sub.func, ast.Attribute) + and sub.func.attr == "log_load_final_alloc_offset" + ] + + +def test_load_emits_live_final_alloc_offset_after_capture() -> None: + runtime = RUNTIME_PATH.read_text() + hooks_source = HOOKS_PATH.read_text() + + # The live post-capture LOAD marker is emitted from runtime with the exact + # format the harness scanner parses, keyed by the rank derived from the + # configured rank_N workspace. + assert "[foundry] LOAD live final_alloc_offset rank_%d: %d" in runtime + assert "def log_load_final_alloc_offset(" in runtime + assert "rank_from_workspace_dir(" in runtime + + # AST contract: the emitter fires exactly once in the whole hooks module, + # only inside capture_model's ``if mode == CUDAGraphExtensionMode.LOAD`` + # branch, and strictly after the wrapped ``result = orig(self)`` capture. + hooks_tree = ast.parse(hooks_source) + assert len(_offset_calls(hooks_tree)) == 1 + + patched = _capture_model_patched(hooks_tree) + calls = _offset_calls(patched) + assert len(calls) == 1 + call = calls[0] + + enclosing_load_if = None + for node in ast.walk(patched): + if isinstance(node, ast.If) and _is_load_mode_test(node.test): + for stmt in node.body: + if isinstance(stmt, ast.Expr) and stmt.value is call: + enclosing_load_if = node + assert enclosing_load_if is not None, "emitter is not gated on a LOAD-only branch" + + orig_linenos = [ + node.lineno + for node in ast.walk(patched) + if isinstance(node, ast.Assign) + and any(isinstance(t, ast.Name) and t.id == "result" for t in node.targets) + and isinstance(node.value, ast.Call) + and isinstance(node.value.func, ast.Name) + and node.value.func.id == "orig" + ] + assert orig_linenos, "capture_model does not assign result = orig(self)" + assert call.lineno > min(orig_linenos) + + +def test_harness_fixes_candidate_and_reports_complete_identity() -> None: + source = HARNESS_PATH.read_text() + + assert "TP_SIZE = 2" in source + assert "TP_SIZE = int(os.environ.get" not in source + assert MODEL_REVISION in source + assert 'return "main"' not in source + assert source.index("require_commit_sha(") < source.index("image =") + for field in ( + '"model"', + '"model_revision"', + '"tp_size"', + '"cudagraph_mode"', + '"distributed_executor_backend"', + '"process_return_code"', + ): + assert field in source diff --git a/tests/test_vmm_alloc.py b/tests/test_vmm_alloc.py index 5df909f0..60c8bde9 100644 --- a/tests/test_vmm_alloc.py +++ b/tests/test_vmm_alloc.py @@ -1,12 +1,46 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the Foundry project +import ctypes import os +import signal import subprocess import sys +import tempfile +import time from pathlib import Path +import foundry as fdry import pytest import torch +from vmm_driver_bindings import load_driver_bindings + +# Hard upper bound for any lifecycle subprocess / protocol wait. A hung child +# (deadlocked IPC, missing signal) fails loudly with captured output instead of +# stalling the suite forever. +_LIFECYCLE_TIMEOUT_S = 120.0 +# Secondary, shorter bound for post-kill collection so even the cleanup path +# cannot block indefinitely on an unkillable/zombie child. +_KILL_TIMEOUT_S = 30.0 + + +class CUipcMemHandle(ctypes.Structure): + _fields_ = [("reserved", ctypes.c_byte * 64)] + + +def _bind_driver(): + """Bind the CUDA driver entry points used by the IPC lifecycle tests.""" + return load_driver_bindings(CUipcMemHandle) + + +def _handle_field(handle, lo, hi): + """Little-endian integer packed into handle.reserved[lo:hi].""" + return int.from_bytes(bytes(bytearray(handle.reserved[lo:hi])), "little") + + +def _read_device_bytes(drv, dptr, n): + buf = (ctypes.c_ubyte * n)() + assert drv.memcpy_dtoh(buf, dptr, n) == 0 + return bytes(buf) def _get_hook_so_path(): @@ -26,8 +60,6 @@ def _get_hook_so_path(): def _run_core(): - import foundry as fdry - torch.cuda.init() device = torch.device("cuda:0") @@ -104,8 +136,6 @@ def _run_core_without_region(): def _run_core_size_parsing(): - import foundry as fdry - torch.cuda.init() device = torch.device("cuda:0") @@ -146,16 +176,434 @@ def _run_core_size_parsing(): print("[TEST] test_size_parsing PASSED") -def _spawn_with_preload(test_mode): +def _run_core_zero_alignment_reserve(): + torch.cuda.init() + + base_addr = 0x7F4000000000 + region_size = 1024 * 1024 * 1024 + default_alignment = 2 * 1024 * 1024 + deliberately_unaligned_offset = default_alignment + 12345 + reserve_size = 2 * 1024 * 1024 + + driver = ctypes.CDLL(None) + reserve = driver.cuMemAddressReserve + reserve.argtypes = [ + ctypes.POINTER(ctypes.c_uint64), + ctypes.c_size_t, + ctypes.c_size_t, + ctypes.c_uint64, + ctypes.c_ulonglong, + ] + reserve.restype = ctypes.c_int + address_free = driver.cuMemAddressFree + address_free.argtypes = [ctypes.c_uint64, ctypes.c_size_t] + address_free.restype = ctypes.c_int + + ptr = ctypes.c_uint64() + with fdry.allocation_region(base_addr, region_size): + fdry.set_current_alloc_offset(deliberately_unaligned_offset) + result = reserve(ctypes.byref(ptr), reserve_size, 0, 0, 0) + assert result == 0 + expected_ptr = (base_addr + deliberately_unaligned_offset + default_alignment - 1) & ~( + default_alignment - 1 + ) + assert ptr.value == expected_ptr, ( + f"zero-alignment reserve returned {hex(ptr.value)}, expected " + f"default-aligned {hex(expected_ptr)}" + ) + expected_offset = expected_ptr - base_addr + reserve_size + assert fdry.get_current_alloc_offset() == expected_offset + assert address_free(ptr.value, reserve_size) == 0 + + +def _run_core_pitched_ipc_export(): + torch.cuda.init() + torch.cuda.set_device(0) + + driver = ctypes.CDLL(None) + alloc_pitch = driver.cuMemAllocPitch_v2 + alloc_pitch.argtypes = [ + ctypes.POINTER(ctypes.c_uint64), + ctypes.POINTER(ctypes.c_size_t), + ctypes.c_size_t, + ctypes.c_size_t, + ctypes.c_uint, + ] + alloc_pitch.restype = ctypes.c_int + ipc_get_handle = driver.cuIpcGetMemHandle + ipc_get_handle.argtypes = [ctypes.POINTER(CUipcMemHandle), ctypes.c_uint64] + ipc_get_handle.restype = ctypes.c_int + mem_free = driver.cuMemFree_v2 + mem_free.argtypes = [ctypes.c_uint64] + mem_free.restype = ctypes.c_int + + ptr = ctypes.c_uint64() + pitch = ctypes.c_size_t() + handle = CUipcMemHandle() + with fdry.allocation_region(0x600000000000, "1GB"): + assert alloc_pitch(ctypes.byref(ptr), ctypes.byref(pitch), 1024, 1024, 4) == 0 + assert ipc_get_handle(ctypes.byref(handle), ptr.value) == 0 + assert mem_free(ptr.value) == 0 + + +def _run_core_rejects_vmm_ipc_after_fork(): + torch.cuda.init() + torch.cuda.set_device(0) + + driver = ctypes.CDLL(None) + mem_alloc = driver.cuMemAlloc_v2 + mem_alloc.argtypes = [ctypes.POINTER(ctypes.c_uint64), ctypes.c_size_t] + mem_alloc.restype = ctypes.c_int + ipc_get_handle = driver.cuIpcGetMemHandle + ipc_get_handle.argtypes = [ctypes.POINTER(CUipcMemHandle), ctypes.c_uint64] + ipc_get_handle.restype = ctypes.c_int + ipc_open_handle = driver.cuIpcOpenMemHandle + ipc_open_handle.argtypes = [ctypes.POINTER(ctypes.c_uint64), CUipcMemHandle, ctypes.c_uint] + ipc_open_handle.restype = ctypes.c_int + ipc_close_handle = driver.cuIpcCloseMemHandle + ipc_close_handle.argtypes = [ctypes.c_uint64] + ipc_close_handle.restype = ctypes.c_int + mem_free = driver.cuMemFree_v2 + mem_free.argtypes = [ctypes.c_uint64] + mem_free.restype = ctypes.c_int + + not_supported = 801 # CUDA_ERROR_NOT_SUPPORTED + + ptr = ctypes.c_uint64() + parent_handle = CUipcMemHandle() + with fdry.allocation_region(0x600000000000, "1GB"): + assert mem_alloc(ctypes.byref(ptr), 2 * 1024 * 1024) == 0 + # The parent claims VMM-IPC ownership by exporting a handle. + assert ipc_get_handle(ctypes.byref(parent_handle), ptr.value) == 0 + + child_pid = os.fork() + if child_pid == 0: + # Every VMM-IPC entry point must fail fast in the post-claim child, + # before touching any inherited lock/map/handle. Exercise get, open, + # and close; each must return CUDA_ERROR_NOT_SUPPORTED. + signal.alarm(5) + child_handle = CUipcMemHandle() + r_get = ipc_get_handle(ctypes.byref(child_handle), ptr.value) + opened = ctypes.c_uint64() + r_open = ipc_open_handle(ctypes.byref(opened), parent_handle, 1) + r_close = ipc_close_handle(ptr.value) + ok = r_get == not_supported and r_open == not_supported and r_close == not_supported + os._exit(0 if ok else 1) + + _, status = os.waitpid(child_pid, 0) + assert os.waitstatus_to_exitcode(status) == 0 + # The parent (owner) is unaffected and can still use and free the buffer. + assert mem_free(ptr.value) == 0 + + +def _run_core_export_server(work_dir, sentinel): + """Exporter process: allocate + fill a VMM buffer, export it (which starts + the abstract-socket server thread), publish the serialized 64-byte handle + blob, then stay alive until signaled so a live importer can fetch the fd via + SCM_RIGHTS. On the free signal, free the allocation but keep the socket/token + server alive until a stale importer is rejected; exit only on a later signal. + Both protocol waits are self-bounded so they cannot hang the suite.""" + torch.cuda.init() + torch.cuda.set_device(0) + drv = _bind_driver() + work = Path(work_dir) + size = 2 * 1024 * 1024 + + fdry.set_allocation_region(0x610000000000, 1 * 1024 * 1024 * 1024) + try: + ptr = ctypes.c_uint64() + assert drv.mem_alloc(ctypes.byref(ptr), size) == 0 + assert drv.memset_d8(ptr.value, sentinel, size) == 0 + torch.cuda.synchronize() + + handle = CUipcMemHandle() + assert drv.ipc_get(ctypes.byref(handle), ptr.value) == 0 # starts server thread + + # Publish the handle blob atomically (write-then-rename), then signal + # readiness with our pid. + blob = bytes(bytearray(handle.reserved)) + assert len(blob) == 64 + tmp = work / "handle.bin.tmp" + tmp.write_bytes(blob) + tmp.rename(work / "handle.bin") + (work / "ready").write_text(str(os.getpid())) + + # Stay alive until the orchestrator asks us to invalidate the exported + # allocation. The server must remain alive after this point so the stale + # importer exercises a negative server response, not a dead socket. + deadline = time.monotonic() + 60.0 + while not (work / "free").exists(): + if time.monotonic() > deadline: + sys.stderr.write("[TEST] exporter timed out waiting for free signal\n") + os._exit(3) + time.sleep(0.02) + + assert drv.mem_free(ptr.value) == 0 + (work / "freed").write_text("1") + + deadline = time.monotonic() + 60.0 + while not (work / "exit").exists(): + if time.monotonic() > deadline: + sys.stderr.write("[TEST] exporter timed out waiting for exit signal\n") + os._exit(4) + time.sleep(0.02) + finally: + fdry.stop_allocation_region() + + +def _load_serialized_handle(work_dir): + blob = (Path(work_dir) / "handle.bin").read_bytes() + assert len(blob) == 64, f"serialized handle must be 64 bytes, got {len(blob)}" + handle = CUipcMemHandle() + ctypes.memmove(handle.reserved, blob, 64) + return handle + + +def _run_core_import_read(work_dir, sentinel): + """Importer process (exporter alive): open the serialized handle - this + connects to the exporter's abstract socket, receives the fd via SCM_RIGHTS + (token verified server-side) - and reads back the exported sentinel.""" + torch.cuda.init() + torch.cuda.set_device(0) + drv = _bind_driver() + n = 64 + + handle = _load_serialized_handle(work_dir) + opened = ctypes.c_uint64() + rc = drv.ipc_open(ctypes.byref(opened), handle, 1) + assert rc == 0, f"live cross-process import failed rc={rc}" + observed = set(_read_device_bytes(drv, opened.value, n)) + assert observed == {sentinel}, f"cross-process import read {observed}, expected {{{sentinel}}}" + assert drv.ipc_close(opened.value) == 0 + + +def _run_core_import_stale(work_dir): + """Importer process (exporter alive, allocation freed): opening the same + serialized handle must fail promptly because the live token server no longer + has a parked fd for the invalidated export.""" + torch.cuda.init() + torch.cuda.set_device(0) + drv = _bind_driver() + + handle = _load_serialized_handle(work_dir) + opened = ctypes.c_uint64() + rc = drv.ipc_open(ctypes.byref(opened), handle, 1) + assert rc != 0, f"stale cross-process import unexpectedly succeeded (rc={rc})" + + +def _run_core_duplicate_interior_open(): + """Opening the same imported interior pointer twice is refcounted: usable + after the first close, torn down only after the second. After the second + close the whole-chunk mapping's VA range must be unreserved (proven by + reserving the exact former range with a hint that must be honored), and a + fresh open must remap and read current data.""" + torch.cuda.init() + torch.cuda.set_device(0) + drv = _bind_driver() + + n = 64 + with fdry.allocation_region(0x600000000000, "1GB", prealloc_size="128MB"): + tensor = torch.empty(1024, 1024, device="cuda") # fast-path chunk carve + ptr = tensor.data_ptr() + assert drv.memset_d8(ptr, 0x11, n) == 0 + torch.cuda.synchronize() + + handle = CUipcMemHandle() + assert drv.ipc_get(ctypes.byref(handle), ptr) == 0 + chunk_base = _handle_field(handle, 32, 40) + chunk_size = _handle_field(handle, 40, 48) + assert chunk_base != 0, "carve from a preallocated chunk must export a chunk base" + interior_offset = ptr - chunk_base + + p1 = ctypes.c_uint64() + p2 = ctypes.c_uint64() + assert drv.ipc_open(ctypes.byref(p1), handle, 1) == 0 + assert drv.ipc_open(ctypes.byref(p2), handle, 1) == 0 + assert p1.value == p2.value, "duplicate interior open must return the same VA" + assert set(_read_device_bytes(drv, p1.value, n)) == {0x11} + + # The whole chunk is mapped at local_base; the interior pointer is + # local_base + (ptr - chunk_base). + local_base = p1.value - interior_offset + + # First close: open_count 2 -> 1, mapping survives and stays usable. + assert drv.ipc_close(p1.value) == 0 + assert set(_read_device_bytes(drv, p2.value, n)) == {0x11} + + # A write through the source is visible through the surviving mapping. + assert drv.memset_d8(ptr, 0x22, n) == 0 + torch.cuda.synchronize() + assert set(_read_device_bytes(drv, p2.value, n)) == {0x22} + + # Second close tears the mapping down. Prove the former mapped range is + # unreserved: reserve exactly [local_base, local_base+chunk_size) with + # local_base as the hint. A still-reserved range would relocate the + # reservation to another VA. No dereference of the freed range. + assert drv.ipc_close(p2.value) == 0 + probe = ctypes.c_uint64() + rc = drv.addr_reserve(ctypes.byref(probe), chunk_size, 0, local_base, 0) + assert rc == 0, f"probe reservation failed rc={rc}" + assert probe.value == local_base, ( + f"former mapped range still reserved after 2nd close: probe landed at " + f"{hex(probe.value)} != {hex(local_base)}" + ) + assert drv.addr_free(probe.value, chunk_size) == 0 + + # A fresh open must remap and read the current source data. + p3 = ctypes.c_uint64() + assert drv.ipc_open(ctypes.byref(p3), handle, 1) == 0 + assert set(_read_device_bytes(drv, p3.value, n)) == {0x22} + assert drv.ipc_close(p3.value) == 0 + + del tensor + + +def _run_core_generation_recreate(): + """Within ONE allocation region (never stopped/re-set between generations), + preallocate + export + import generation 1, free the preallocation, then + preallocate generation 2 at the current monotonic cursor. Because the region + cursor advances and is never rewound, generation 2 gets a DIFFERENT chunk + base; the global generation counter also increments. Assert the OBSERVABLE + behavior: the exported chunk generation increments and a fresh import reads + the recreated (fresh) bytes, never stale generation-1 data. + + We deliberately do NOT stop/set the region between generations and make no + claim about identical VAs. The generation-sensitivity of the mapping key + itself (same pid/token/base + different generation -> distinct key; dropping + the generation collapses them) is proven hermetically by the pure CPU test + tests/vmm_ipc_chunk_key_test.cpp; this native test complements it with a real + device round trip across a free/recreate within a single live region.""" + torch.cuda.init() + torch.cuda.set_device(0) + drv = _bind_driver() + + base = 0x600000000000 + region_size = 1 * 1024 * 1024 * 1024 + prealloc = 128 * 1024 * 1024 + n = 64 + + def preallocate_export_import_read(sentinel): + assert fdry.preallocate_region(prealloc) + tensor = torch.empty(1024, 1024, device="cuda") + ptr = tensor.data_ptr() + assert drv.memset_d8(ptr, sentinel, n) == 0 + torch.cuda.synchronize() + handle = CUipcMemHandle() + assert drv.ipc_get(ctypes.byref(handle), ptr) == 0 + opened = ctypes.c_uint64() + assert drv.ipc_open(ctypes.byref(opened), handle, 1) == 0 + observed = set(_read_device_bytes(drv, opened.value, n)) + assert drv.ipc_close(opened.value) == 0 + chunk_base = _handle_field(handle, 32, 40) + gen = _handle_field(handle, 48, 56) + # Importer is closed above; now drop the tensor and drain the caching + # allocator BEFORE the caller frees the preallocation, so no cached block + # retains a pointer into this generation's soon-to-be-unmapped backing. + # empty_cache() releases blocks via cuMemFree_v2, which does NOT rewind + # the region cursor, so the monotonic cursor (and the distinct gen-2 + # base) is preserved. + del tensor + torch.cuda.synchronize() + torch.cuda.empty_cache() + return chunk_base, gen, observed + + # One region for BOTH generations; only the preallocation is freed/recreated. + fdry.set_allocation_region(base, region_size) + try: + cb1, gen1, observed1 = preallocate_export_import_read(0xA1) + assert observed1 == {0xA1}, "gen-1 import must read its own bytes" + # Free only the preallocation; the region stays live and the cursor keeps + # its advanced position (freeing does not rewind it). + fdry.free_preallocated_region() + + cb2, gen2, observed2 = preallocate_export_import_read(0xB2) + assert observed2 == {0xB2}, "recreated import must read FRESH bytes, never stale gen-1 data" + fdry.free_preallocated_region() + finally: + fdry.stop_allocation_region() + + assert cb2 != cb1, ( + f"gen-2 must land at a different chunk base on the monotonic cursor " + f"(cb1={hex(cb1)}, cb2={hex(cb2)})" + ) + assert gen2 > gen1, f"chunk generation must increment on recreate (gen1={gen1}, gen2={gen2})" + + +def _preload_env(): so_path = _get_hook_so_path() env = os.environ.copy() if env.get("LD_PRELOAD"): env["LD_PRELOAD"] = f"{so_path}:{env['LD_PRELOAD']}" else: env["LD_PRELOAD"] = so_path + return env + + +def _kill_and_collect(proc, label, timeout=_KILL_TIMEOUT_S): + """Kill proc if still running and collect its output under a bounded + secondary wait. Never blocks indefinitely: if the process will not die + within `timeout` after the kill, raise with diagnostics rather than wait + forever. Returns (out, err).""" + if proc.poll() is None: + proc.kill() + try: + return proc.communicate(timeout=timeout) + except subprocess.TimeoutExpired as exc: + raise AssertionError( + f"{label} pid={proc.pid} did not terminate within {timeout}s after kill; " + f"abandoning to avoid an indefinite wait" + ) from exc + + +def _run_capture(cmd, env, label, timeout=_LIFECYCLE_TIMEOUT_S): + """Run a subprocess to completion under a hard timeout, asserting rc==0 and + surfacing captured output. On timeout the child is killed and collected + under a bounded secondary wait (no indefinite blocking).""" + proc = subprocess.Popen(cmd, env=env, stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True) + try: + out, err = proc.communicate(timeout=timeout) + except subprocess.TimeoutExpired as exc: + out, err = _kill_and_collect(proc, label) + raise AssertionError( + f"{label} timed out after {timeout}s\nSTDOUT:\n{out}\nSTDERR:\n{err}" + ) from exc + assert proc.returncode == 0, ( + f"{label} failed rc={proc.returncode}\nSTDOUT:\n{out}\nSTDERR:\n{err}" + ) + return out, err + + +def _spawn_with_preload(test_mode, timeout=_LIFECYCLE_TIMEOUT_S): + cmd = [sys.executable, str(Path(__file__).resolve()), f"--{test_mode}"] + _run_capture(cmd, _preload_env(), test_mode, timeout=timeout) + +def _spawn_with_preload_capture(test_mode, timeout=_LIFECYCLE_TIMEOUT_S): + """Run a mode under LD_PRELOAD, asserting success and returning (out, err).""" cmd = [sys.executable, str(Path(__file__).resolve()), f"--{test_mode}"] - subprocess.check_call(cmd, env=env) + return _run_capture(cmd, _preload_env(), test_mode, timeout=timeout) + + +def _wait_for_file(path, proc, label, timeout=_LIFECYCLE_TIMEOUT_S): + """Poll for a protocol file, failing fast (with captured output) if the + producer process dies early or the deadline elapses. Every output + collection is bounded via _kill_and_collect.""" + deadline = time.monotonic() + timeout + while not path.exists(): + rc = proc.poll() + if rc is not None: + out, err = _kill_and_collect(proc, label) + raise AssertionError( + f"{label} exited early rc={rc} before publishing {path.name}\n" + f"STDOUT:\n{out}\nSTDERR:\n{err}" + ) + if time.monotonic() > deadline: + out, err = _kill_and_collect(proc, label) + raise AssertionError( + f"timed out after {timeout}s waiting for {path.name} from {label}\n" + f"STDOUT:\n{out}\nSTDERR:\n{err}" + ) + time.sleep(0.02) @pytest.mark.filterwarnings("ignore:TORCH_CUDA_ARCH_LIST is not set") @@ -176,6 +624,110 @@ def test_size_parsing(): _spawn_with_preload("size-parsing") +@pytest.mark.filterwarnings("ignore:TORCH_CUDA_ARCH_LIST is not set") +def test_zero_alignment_reserve_stays_in_region(): + """CUDA's alignment=0 convention uses Foundry's default 2 MiB alignment.""" + _spawn_with_preload("zero-alignment-reserve") + + +@pytest.mark.filterwarnings("ignore:TORCH_CUDA_ARCH_LIST is not set") +def test_pitched_vmm_allocation_can_be_exported(): + """Tracked pitched allocations must request a shareable POSIX handle.""" + _spawn_with_preload("pitched-ipc-export") + + +@pytest.mark.filterwarnings("ignore:TORCH_CUDA_ARCH_LIST is not set") +def test_vmm_ipc_rejects_use_after_fork(): + """A fork child must fail promptly (get/open/close) instead of touching + inherited IPC locks, and emit the async-signal-safe 'use spawn/exec' + rejection message.""" + _, stderr = _spawn_with_preload_capture("reject-vmm-ipc-after-fork") + assert "use spawn/exec" in stderr + + +@pytest.mark.filterwarnings("ignore:TORCH_CUDA_ARCH_LIST is not set") +def test_vmm_ipc_stale_exporter_cross_process(tmp_path): + """Cross-process stale exporter over the abstract socket + SCM_RIGHTS + token + path: while the exporter is alive a fresh importer opens the serialized + handle and reads the sentinel; the exporter then frees the allocation but + remains alive, and a fresh importer using the same handle fails promptly + through the still-live server. Every wait is bounded and hangs fail with + diagnostics.""" + sentinel = 0xCD + work = tmp_path / "ipc" + work.mkdir() + env = _preload_env() + self_path = str(Path(__file__).resolve()) + + exporter = subprocess.Popen( + [sys.executable, self_path, "--export-server", str(work), str(sentinel)], + env=env, + stdout=subprocess.PIPE, + stderr=subprocess.PIPE, + text=True, + ) + try: + # 1. Exporter publishes the handle + readiness. + _wait_for_file(work / "ready", exporter, "exporter") + + # 2. Live importer fetches the fd cross-process and reads the sentinel. + _run_capture( + [sys.executable, self_path, "--import-read", str(work), str(sentinel)], + env, + "live cross-process import", + ) + + # 3. Signal only the allocation free, wait for acknowledgement, and + # prove the exporter/socket server remains alive. + (work / "free").write_text("1") + _wait_for_file(work / "freed", exporter, "exporter") + assert exporter.poll() is None, "exporter exited before stale import retry" + + # 4. Fresh importer with the same serialized handle must fail promptly + # through the live socket/token server's negative response. + stale_stdout, stale_stderr = _run_capture( + [sys.executable, self_path, "--import-stale", str(work)], + env, + "stale cross-process import", + ) + assert "live exporter server rejected stale/unknown export" in stale_stderr, ( + "stale import did not prove a live-server negative response\n" + f"STDOUT:\n{stale_stdout}\nSTDERR:\n{stale_stderr}" + ) + assert "could not connect to live exporter server" not in stale_stderr + assert exporter.poll() is None, "exporter exited during stale import retry" + + # 5. Only after the stale rejection may the exporter exit. Both the + # normal wait and the post-kill collection are bounded. + (work / "exit").write_text("1") + try: + out, err = exporter.communicate(timeout=_LIFECYCLE_TIMEOUT_S) + except subprocess.TimeoutExpired as exc: + out, err = _kill_and_collect(exporter, "exporter") + raise AssertionError( + f"exporter did not exit after signal\nSTDOUT:\n{out}\nSTDERR:\n{err}" + ) from exc + assert exporter.returncode == 0, ( + f"exporter rc={exporter.returncode}\nSTDOUT:\n{out}\nSTDERR:\n{err}" + ) + finally: + # Bounded cleanup: never wait indefinitely even on the teardown path. + if exporter.poll() is None: + _kill_and_collect(exporter, "exporter (cleanup)") + + +@pytest.mark.filterwarnings("ignore:TORCH_CUDA_ARCH_LIST is not set") +def test_vmm_ipc_duplicate_interior_open_refcounts(): + """Duplicate interior opens are refcounted across open/close/close.""" + _spawn_with_preload("duplicate-interior-open") + + +@pytest.mark.filterwarnings("ignore:TORCH_CUDA_ARCH_LIST is not set") +def test_vmm_ipc_preallocation_generation_recreate(): + """Recreating a preallocation bumps the generation and maps fresh data.""" + _spawn_with_preload("generation-recreate") + + if __name__ == "__main__": if "--core" in sys.argv: _run_core() @@ -183,7 +735,33 @@ def test_size_parsing(): _run_core_without_region() elif "--size-parsing" in sys.argv: _run_core_size_parsing() + elif "--zero-alignment-reserve" in sys.argv: + _run_core_zero_alignment_reserve() + elif "--pitched-ipc-export" in sys.argv: + _run_core_pitched_ipc_export() + elif "--reject-vmm-ipc-after-fork" in sys.argv: + _run_core_rejects_vmm_ipc_after_fork() + elif "--export-server" in sys.argv: + idx = sys.argv.index("--export-server") + _run_core_export_server(sys.argv[idx + 1], int(sys.argv[idx + 2])) + elif "--import-read" in sys.argv: + idx = sys.argv.index("--import-read") + _run_core_import_read(sys.argv[idx + 1], int(sys.argv[idx + 2])) + elif "--import-stale" in sys.argv: + idx = sys.argv.index("--import-stale") + _run_core_import_stale(sys.argv[idx + 1]) + elif "--duplicate-interior-open" in sys.argv: + _run_core_duplicate_interior_open() + elif "--generation-recreate" in sys.argv: + _run_core_generation_recreate() else: test_vmm_allocation() test_vmm_allocation_without_region() test_size_parsing() + test_zero_alignment_reserve_stays_in_region() + test_pitched_vmm_allocation_can_be_exported() + test_vmm_ipc_rejects_use_after_fork() + with tempfile.TemporaryDirectory() as _tmp: + test_vmm_ipc_stale_exporter_cross_process(Path(_tmp)) + test_vmm_ipc_duplicate_interior_open_refcounts() + test_vmm_ipc_preallocation_generation_recreate() diff --git a/tests/test_vmm_cleanup_contract.py b/tests/test_vmm_cleanup_contract.py new file mode 100644 index 00000000..c2d48980 --- /dev/null +++ b/tests/test_vmm_cleanup_contract.py @@ -0,0 +1,241 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the Foundry project +"""CUDA-free source contracts for retryable native VMM cleanup.""" + +from __future__ import annotations + +import ast +from pathlib import Path + +REPO_ROOT = Path(__file__).parents[1] +HOOK_CPP = REPO_ROOT / "csrc" / "hook.cpp" +HOOK_H = REPO_ROOT / "include" / "hook.h" +BINDING_CPP = REPO_ROOT / "csrc" / "binding.cpp" +VMM_ALLOC_TEST = REPO_ROOT / "tests" / "test_vmm_alloc.py" + + +def _function_body(source: str, signature: str) -> str: + start = 0 + while True: + idx = source.find(signature, start) + if idx == -1: + raise AssertionError(f"signature not found: {signature!r}") + brace = source.find("{", idx) + semi = source.find(";", idx) + if brace != -1 and (semi == -1 or brace < semi): + break + start = idx + len(signature) + depth = 0 + for pos in range(brace, len(source)): + if source[pos] == "{": + depth += 1 + elif source[pos] == "}": + depth -= 1 + if depth == 0: + return source[brace : pos + 1] + raise AssertionError(f"unbalanced braces for {signature!r}") + + +def _python_function_source(source: str, name: str) -> str: + tree = ast.parse(source) + for node in tree.body: + if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef)) and node.name == name: + segment = ast.get_source_segment(source, node) + assert segment is not None + return segment + raise AssertionError(f"Python function not found: {name!r}") + + +def test_free_preallocated_region_returns_bool_through_native_binding() -> None: + assert "bool free_preallocated_region();" in HOOK_H.read_text() + binding = BINDING_CPP.read_text() + assert ( + '"free_preallocated_region", []() { return ::foundry::free_preallocated_region(); }' + in binding + ) + + +def test_preallocation_replacement_refuses_failed_cleanup() -> None: + body = _function_body(HOOK_CPP.read_text(), "bool preallocate_region(size_t size)") + free_call = body.find("if (!free_preallocated_region())") + create_call = body.find("mem_create_func(") + assert free_call != -1 + assert create_call != -1 + assert free_call < create_call + assert "Cannot replace preallocation because cleanup failed" in body + + +def test_free_preallocation_retains_phase_state_until_full_success() -> None: + source = HOOK_CPP.read_text() + assert "foundry::vmm::CleanupState preallocation_cleanup" in source + body = _function_body(source, "bool free_preallocated_region()") + assert "foundry::vmm::run_cleanup(" in body + assert "cleanup_complete(" in body + assert "return false;" in body + assert body.find("run_cleanup(") < body.find("preallocated_chunks.erase(") + assert body.find("cleanup_complete(") < body.find("has_preallocation = false") + for phase in ("cuMemUnmap", "cuMemRelease"): + assert f"free_preallocated_region: {phase} failed" in body + + +def test_chunk_and_individual_close_use_retryable_three_phase_cleanup() -> None: + source = HOOK_CPP.read_text() + assert "foundry::vmm::CleanupState cleanup;" in source + body = _function_body(source, "CUresult cuIpcCloseMemHandle(CUdeviceptr dptr)") + assert body.count("foundry::vmm::run_cleanup(") == 2 + assert body.count("foundry::vmm::cleanup_complete(") >= 2 + for ownership in ("chunk", "individual"): + for phase in ("cuMemUnmap", "cuMemRelease", "cuMemAddressFree"): + assert f"VMM-IPC {ownership} close: {phase} failed" in body + + +def test_ipc_close_erases_tracking_only_after_cleanup_completes() -> None: + body = _function_body(HOOK_CPP.read_text(), "CUresult cuIpcCloseMemHandle(CUdeviceptr dptr)") + chunk_complete = body.find("cleanup_complete(it->second.cleanup") + chunk_erase = body.find("vmm_ipc_chunk_mappings.erase(it)") + individual_complete = body.find("individual_cleanup_complete") + individual_erase = body.find("global_alloc_metadata.erase_if") + assert -1 not in (chunk_complete, chunk_erase, individual_complete, individual_erase) + assert chunk_complete < chunk_erase + assert individual_complete < individual_erase + + +def test_vmm_free_cleanup_distinguishes_region_owned_and_imported_addresses() -> None: + body = _function_body(HOOK_CPP.read_text(), "CUresult cuMemFree_v2(CUdeviceptr dptr)") + + mark_started = body.find("kv.second.cleanup.started = true") + invalidate = body.find("vmm_ipc_invalidate_export(dptr)") + ownership = body.find("bool address_free_required = metadata.region_base == 0") + cleanup = body.find("foundry::vmm::run_cleanup(") + complete = body.find("foundry::vmm::cleanup_complete(") + erase = body.find("global_alloc_metadata.erase_if", complete) + record = body.find("HookAllocationEvent::Type::Free", erase) + assert -1 not in (mark_started, invalidate, ownership, cleanup, complete, erase, record) + assert mark_started < invalidate < ownership < cleanup < complete < erase < record + assert "Region-owned allocations are subranges of one whole-region VA reservation" in body + compact_body = " ".join(body.split()) + assert "kv.second.cleanup, CUDA_SUCCESS, address_free_required" in compact_body + assert "cleanup_complete(kv.second.cleanup, address_free_required)" in compact_body + assert "return cleanup_status;" in body + for phase in ("cuMemUnmap", "cuMemRelease", "cuMemAddressFree"): + assert f"cuMemFree_v2: {phase} failed" in body + + +def test_preallocated_carve_free_is_explicitly_nonphysical() -> None: + body = _function_body(HOOK_CPP.read_text(), "CUresult cuMemFree_v2(CUdeviceptr dptr)") + carve_branch = body.find("if (metadata.from_preallocation)") + cleanup = body.find("foundry::vmm::run_cleanup(") + assert -1 not in (carve_branch, cleanup) + assert carve_branch < cleanup + assert "Preallocated carves do not own physical VMM resources" in body + + +def test_preallocated_carve_export_fails_while_chunk_cleanup_is_pending() -> None: + body = _function_body(HOOK_CPP.read_text(), "CUresult cuIpcGetMemHandle") + assert "preallocated_chunk_cleanup_started" in body + assert "preallocated chunk cleanup is pending" in body + pending_check = body.find("if (preallocated_chunk_cleanup_started)") + real_fallback = body.find("return real_func(pHandle, dptr)") + assert -1 not in (pending_check, real_fallback) + assert pending_check < real_fallback + + +def test_chunk_open_race_and_collision_use_checked_cleanup_policy() -> None: + source = HOOK_CPP.read_text() + helper = _function_body(source, "static CUresult cleanup_import_mapping") + for operation in ("cuMemUnmap", "cuMemRelease", "cuMemAddressFree"): + assert f'operation = "{operation}"' in helper + assert "failed with error %d" in helper + + body = _function_body(source, "CUresult cuIpcOpenMemHandle") + race = body[body.find("// Register the whole-chunk mapping") :] + + assert '"chunk race loser"' in race + assert '"interior collision"' in race + assert race.count("cleanup_import_mapping(") == 2 + assert "winner_cleanup_started" in race + loser_cleanup = race.find("CUresult loser_cleanup_status = cleanup_import_mapping(") + pending_reject = race.find("if (winner_cleanup_started)") + winner_increment = race.find("it->second.refcount++") + assert -1 not in (loser_cleanup, pending_reject, winner_increment) + assert loser_cleanup < pending_reject < winner_increment + + collision_cleanup = race.find("CUresult collision_cleanup_status = cleanup_import_mapping(") + collision_complete = race.find("cleanup_complete(mapping_it->second.cleanup", collision_cleanup) + collision_erase = race.find("vmm_ipc_chunk_mappings.erase(mapping_it)", collision_cleanup) + assert -1 not in (collision_cleanup, collision_complete, collision_erase) + assert collision_cleanup < collision_complete < collision_erase + + +def test_pending_import_cleanup_is_drained_before_any_new_open() -> None: + source = HOOK_CPP.read_text() + assert "struct PendingVmmIpcCleanup" in source + assert "std::multimap" in source + assert "vmm_ipc_pending_cleanups" in source + + body = _function_body(source, "CUresult cuIpcOpenMemHandle") + claim = body.find("if (!vmm_ipc_claim_process())") + drain = body.find("drain_pending_import_cleanups_locked()") + real_lookup = body.find("CUDA_ENTRY_cuIpcOpenMemHandle") + magic = body.find("if (magic == VMM_IPC_MAGIC)") + assert -1 not in (claim, drain, real_lookup, magic) + assert claim < drain < real_lookup < magic + assert "if (pending_cleanup_status != CUDA_SUCCESS)" in body + assert body.find("return pending_cleanup_status") < real_lookup + + drain_body = _function_body(source, "static CUresult drain_pending_import_cleanups_locked()") + assert "foundry::vmm::drain_pending_cleanups(" in drain_body + assert "cleanup_import_mapping(" in drain_body + assert "cleanup_complete(pending.cleanup, true)" in drain_body + + +def test_failed_race_and_collision_cleanup_moves_to_durable_pending_storage() -> None: + body = _function_body(HOOK_CPP.read_text(), "CUresult cuIpcOpenMemHandle") + race = body[body.find("// Register the whole-chunk mapping") :] + + loser_cleanup = race.find("CUresult loser_cleanup_status = cleanup_import_mapping(") + loser_pending = race.find("queue_pending_import_cleanup(", loser_cleanup) + loser_return = race.find("return loser_cleanup_status", loser_cleanup) + winner_increment = race.find("it->second.refcount++", loser_cleanup) + assert -1 not in (loser_cleanup, loser_pending, loser_return, winner_increment) + assert loser_cleanup < loser_pending < loser_return < winner_increment + + collision_cleanup = race.find("CUresult collision_cleanup_status = cleanup_import_mapping(") + collision_pending = race.find("queue_pending_import_cleanup(", collision_cleanup) + collision_erase = race.find("vmm_ipc_chunk_mappings.erase(mapping_it)", collision_pending) + collision_return = race.find("return collision_cleanup_status", collision_pending) + pointer_return = race.find("*pdptr = interior", collision_cleanup) + assert -1 not in ( + collision_cleanup, + collision_pending, + collision_erase, + collision_return, + pointer_return, + ) + assert ( + collision_cleanup < collision_pending < collision_erase < collision_return < pointer_return + ) + + +def test_stale_handle_is_retried_while_export_server_remains_alive() -> None: + source = VMM_ALLOC_TEST.read_text() + exporter = _python_function_source(source, "_run_core_export_server") + lifecycle = _python_function_source(source, "test_vmm_ipc_stale_exporter_cross_process") + + free_call = exporter.find("drv.mem_free(ptr.value)") + freed_signal = exporter.find('(work / "freed").write_text') + exit_wait = exporter.find('while not (work / "exit").exists()') + assert -1 not in (free_call, freed_signal, exit_wait) + assert free_call < freed_signal < exit_wait + + signal_free = lifecycle.find('(work / "free").write_text') + wait_freed = lifecycle.find('_wait_for_file(work / "freed"') + alive_check = lifecycle.find("assert exporter.poll() is None") + stale_import = lifecycle.find('"--import-stale"') + signal_exit = lifecycle.find('(work / "exit").write_text') + assert -1 not in (signal_free, wait_freed, alive_check, stale_import, signal_exit) + assert signal_free < wait_freed < alive_check < stale_import < signal_exit + assert "live exporter server rejected stale/unknown export" in HOOK_CPP.read_text() + assert "stale_stdout, stale_stderr = _run_capture(" in lifecycle + assert '"live exporter server rejected stale/unknown export" in stale_stderr' in lifecycle + assert '"could not connect to live exporter server" not in stale_stderr' in lifecycle diff --git a/tests/test_vmm_driver_bindings.py b/tests/test_vmm_driver_bindings.py new file mode 100644 index 00000000..de5fcbe7 --- /dev/null +++ b/tests/test_vmm_driver_bindings.py @@ -0,0 +1,179 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the Foundry project +"""CPU-only contracts for VMM lifecycle CUDA driver bindings.""" + +from __future__ import annotations + +import ctypes +import importlib.util +from pathlib import Path + +import pytest + +HELPER_PATH = Path(__file__).parent / "vmm_driver_bindings.py" + + +class FakeFunction: + def __init__(self, name: str) -> None: + self.name = name + self.argtypes = None + self.restype = None + + +class FakeIpcHandle(ctypes.Structure): + _fields_ = [("reserved", ctypes.c_byte * 64)] + + +class FakeLibrary: + def __init__(self, *symbols: str) -> None: + self.symbols = {name: FakeFunction(name) for name in symbols} + self.lookups: list[str] = [] + + def __getattr__(self, name: str) -> FakeFunction: + self.lookups.append(name) + try: + return self.symbols[name] + except KeyError as exc: + raise AttributeError(name) from exc + + +def _load_helper(): + assert HELPER_PATH.exists(), f"required helper module is missing: {HELPER_PATH}" + spec = importlib.util.spec_from_file_location("vmm_driver_bindings_contract", HELPER_PATH) + assert spec is not None and spec.loader is not None + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + return module + + +def _process_library() -> FakeLibrary: + return FakeLibrary( + "cuMemAlloc_v2", + "cuMemFree_v2", + "cuIpcGetMemHandle", + "cuIpcOpenMemHandle", + "cuIpcCloseMemHandle", + "cuMemAddressReserve", + "cuMemAddressFree", + ) + + +def _cuda_library(*utility_symbols: str) -> FakeLibrary: + return FakeLibrary(*utility_symbols) + + +def test_loads_process_global_and_libcuda_for_distinct_operation_classes() -> None: + helper = _load_helper() + process = _process_library() + cuda = _cuda_library("cuMemsetD8_v2", "cuMemcpyDtoH_v2") + load_requests = [] + + def load_library(name): + load_requests.append(name) + return process if name is None else cuda + + bindings = helper.load_driver_bindings(FakeIpcHandle, load_library=load_library) + + assert load_requests == [None, "libcuda.so.1"] + assert bindings.mem_alloc is process.symbols["cuMemAlloc_v2"] + assert bindings.mem_free is process.symbols["cuMemFree_v2"] + assert bindings.ipc_get is process.symbols["cuIpcGetMemHandle"] + assert bindings.ipc_open is process.symbols["cuIpcOpenMemHandle"] + assert bindings.ipc_close is process.symbols["cuIpcCloseMemHandle"] + assert bindings.addr_reserve is process.symbols["cuMemAddressReserve"] + assert bindings.addr_free is process.symbols["cuMemAddressFree"] + assert bindings.memset_d8 is cuda.symbols["cuMemsetD8_v2"] + assert bindings.memcpy_dtoh is cuda.symbols["cuMemcpyDtoH_v2"] + + +@pytest.mark.parametrize( + "symbol", + [ + "cuMemAlloc_v2", + "cuMemFree_v2", + "cuIpcGetMemHandle", + "cuIpcOpenMemHandle", + "cuIpcCloseMemHandle", + "cuMemAddressReserve", + "cuMemAddressFree", + ], +) +def test_each_intercepted_operation_uses_only_process_global(symbol: str) -> None: + helper = _load_helper() + process = _process_library() + del process.symbols[symbol] + cuda = _cuda_library( + "cuMemAlloc_v2", + "cuMemFree_v2", + "cuIpcGetMemHandle", + "cuIpcOpenMemHandle", + "cuIpcCloseMemHandle", + "cuMemAddressReserve", + "cuMemAddressFree", + "cuMemsetD8_v2", + "cuMemcpyDtoH_v2", + ) + + with pytest.raises(AttributeError, match=symbol): + helper.bind_driver_libraries(process, cuda, FakeIpcHandle) + + assert cuda.lookups == [] + + +@pytest.mark.parametrize( + "symbol", + [ + "cuMemsetD8_v2", + "cuMemcpyDtoH_v2", + ], +) +def test_each_utility_uses_only_exact_libcuda_v2_symbol(symbol: str) -> None: + helper = _load_helper() + process = _process_library() + process.symbols[symbol] = FakeFunction(symbol) + cuda = _cuda_library( + "cuMemsetD8_v2", + "cuMemcpyDtoH_v2", + ) + del cuda.symbols[symbol] + + with pytest.raises( + AttributeError, + match=rf"required CUDA driver symbol {symbol} is unavailable from libcuda\.so\.1", + ): + helper.bind_driver_libraries(process, cuda, FakeIpcHandle) + + assert process.lookups.count(symbol) == 0 + assert cuda.lookups[-1:] == [symbol] + assert symbol.removesuffix("_v2") not in cuda.lookups + + +def test_all_bindings_have_exact_ctypes_signatures() -> None: + helper = _load_helper() + process = _process_library() + cuda = _cuda_library("cuMemsetD8_v2", "cuMemcpyDtoH_v2") + + bindings = helper.bind_driver_libraries(process, cuda, FakeIpcHandle) + + u64 = ctypes.c_uint64 + expected_argtypes = { + "mem_alloc": [ctypes.POINTER(u64), ctypes.c_size_t], + "mem_free": [u64], + "ipc_get": [ctypes.POINTER(FakeIpcHandle), u64], + "ipc_open": [ctypes.POINTER(u64), FakeIpcHandle, ctypes.c_uint], + "ipc_close": [u64], + "addr_reserve": [ + ctypes.POINTER(u64), + ctypes.c_size_t, + ctypes.c_size_t, + u64, + ctypes.c_ulonglong, + ], + "addr_free": [u64, ctypes.c_size_t], + "memset_d8": [u64, ctypes.c_ubyte, ctypes.c_size_t], + "memcpy_dtoh": [ctypes.c_void_p, u64, ctypes.c_size_t], + } + for binding_name, argtypes in expected_argtypes.items(): + function = getattr(bindings, binding_name) + assert function.argtypes == argtypes, binding_name + assert function.restype is ctypes.c_int, binding_name diff --git a/tests/test_vmm_ipc_ownership_contract.py b/tests/test_vmm_ipc_ownership_contract.py new file mode 100644 index 00000000..3e2b77a4 --- /dev/null +++ b/tests/test_vmm_ipc_ownership_contract.py @@ -0,0 +1,210 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the Foundry project +"""Source contracts for VMM-IPC process-ownership guards. + +These assert placement/ordering invariants in ``csrc/hook.cpp`` that cannot be +observed natively: a post-claim fork child cannot be driven to introspect the +server-init or export-invalidation paths, and the async-signal-safety of the +rejection path is a property of *which* calls appear on it, not of a return +code. They are pure text/AST-free contracts and need no CUDA device. +""" + +from __future__ import annotations + +from pathlib import Path + +import pytest + +HOOK_CPP = Path(__file__).parents[1] / "csrc" / "hook.cpp" + + +def _source() -> str: + return HOOK_CPP.read_text() + + +def _function_body(source: str, signature: str) -> str: + """Return the brace-matched body of the function whose signature contains + ``signature`` (the first occurrence that is a definition, i.e. followed by + a ``{`` before any ``;``).""" + start = 0 + while True: + idx = source.find(signature, start) + if idx == -1: + raise AssertionError(f"signature not found: {signature!r}") + brace = source.find("{", idx) + semi = source.find(";", idx) + if brace != -1 and (semi == -1 or brace < semi): + break + start = idx + len(signature) + depth = 0 + for pos in range(brace, len(source)): + ch = source[pos] + if ch == "{": + depth += 1 + elif ch == "}": + depth -= 1 + if depth == 0: + return source[brace : pos + 1] + raise AssertionError(f"unbalanced braces for {signature!r}") + + +def test_owner_pid_defaults_to_unclaimed_so_fork_before_claim_is_allowed() -> None: + # Owner starts unclaimed; a fork before the first claim observes -1 and is + # allowed to become its own owner (pre-claim owner semantics). + assert "std::atomic vmm_ipc_owner_pid{-1}" in _source() + + +def test_owner_pid_atomic_is_asserted_lock_free() -> None: + # The owner pid is CAS'd on the async-signal-safe fork-child reject path; a + # lock-backed atomic could deadlock there, so it must be statically checked. + src = _source() + assert "static_assert(std::atomic::is_always_lock_free" in src + + +def test_server_init_claims_ownership_before_locking_server_mutex() -> None: + body = _function_body(_source(), "static bool vmm_ipc_ensure_server()") + claim = body.find("vmm_ipc_claim_process()") + lock = body.find("std::lock_guard lock(vmm_ipc_server_mutex)") + assert claim != -1, "ensure_server must claim ownership" + assert lock != -1, "ensure_server must lock the server mutex" + assert claim < lock, "claim must happen before locking inherited mutex" + + +def test_invalidation_guards_before_touching_the_export_map() -> None: + body = _function_body(_source(), "static void vmm_ipc_invalidate_export(CUdeviceptr dptr)") + guard = body.find("owns_ipc_state(") + erase = body.find("vmm_ipc_exported_fds.erase_if") + assert guard != -1, "invalidation must consult owns_ipc_state" + assert erase != -1, "invalidation must erase from the export map" + assert guard < erase, "owner guard must return before any map access" + # The guard must early-return (no map access on the not-owner path). + assert body[guard : guard + 200].count("return") >= 1 + + +@pytest.mark.parametrize( + ("signature", "shared_state"), + [ + ( + "CUresult cuIpcGetMemHandle(CUipcMemHandle* pHandle, CUdeviceptr dptr)", + "global_alloc_metadata.visit", + ), + ( + "CUresult cuIpcOpenMemHandle(CUdeviceptr* pdptr, CUipcMemHandle handle, unsigned int Flags)", + "memcpy(&magic, handle.reserved", + ), + ("CUresult cuIpcCloseMemHandle(CUdeviceptr dptr)", "vmm_ipc_chunk_mutex"), + ], +) +def test_entry_points_claim_ownership_before_shared_state( + signature: str, shared_state: str +) -> None: + body = _function_body(_source(), signature) + claim = body.find("vmm_ipc_claim_process()") + shared = body.find(shared_state) + assert claim != -1, f"{signature} must claim ownership" + assert shared != -1, f"{signature} must reach its shared state {shared_state!r}" + assert claim < shared, f"{signature} must claim before touching {shared_state!r}" + + +@pytest.mark.parametrize( + "signature", + [ + "CUresult cuIpcGetMemHandle(CUipcMemHandle* pHandle, CUdeviceptr dptr)", + "CUresult cuIpcOpenMemHandle(CUdeviceptr* pdptr, CUipcMemHandle handle, unsigned int Flags)", + "CUresult cuIpcCloseMemHandle(CUdeviceptr dptr)", + ], +) +def test_entry_points_claim_before_driver_symbol_lookup(signature: str) -> None: + # The claim must be the first executable action, ahead of the cached driver + # symbol lookup (CUDA_DRIVER_CALL takes a lock a fork child must not touch). + # A leading typedef is not executable, so ordering is measured against the + # CUDA_DRIVER_CALL that resolves the real entry point. + body = _function_body(_source(), signature) + claim = body.find("vmm_ipc_claim_process()") + lookup = body.find("CUDA_DRIVER_CALL(") + assert claim != -1, f"{signature} must claim ownership" + assert lookup != -1, f"{signature} must resolve its real entry via CUDA_DRIVER_CALL" + assert claim < lookup, f"{signature} must claim before the driver symbol lookup" + + +def test_import_fallback_free_failure_fails_closed_before_zone_mapping() -> None: + # A failed free of the driver's fallback reservation must fail closed: route + # the decision through the pure helper, log both cleanup statuses, and return + # immediately BEFORE the dedicated-zone reservation - never continue to map. + body = _function_body( + _source(), + "CUresult cuIpcOpenMemHandle(CUdeviceptr* pdptr, CUipcMemHandle handle, unsigned int Flags)", + ) + helper = body.find("require_fallback_reservation_freed(") + fail_return = body.find("return fallback_status;") + zone_reserve = body.find("real_reserve_func(&mapped_ptr, map_size, 0, (CUdeviceptr)zone_hint") + assert helper != -1, "fallback free must route through require_fallback_reservation_freed" + assert fail_return != -1, "fallback failure must return immediately" + assert zone_reserve != -1, "the dedicated-zone reservation must exist" + assert helper < fail_return < zone_reserve, ( + "must fail closed (return) before attempting the dedicated-zone reservation" + ) + # Both cleanup statuses are logged (the free status and the release status). + assert "cuMemAddressFree(fallback reservation" in body + assert "cuMemRelease status" in body + + +def test_chunk_key_uses_pure_generation_sensitive_type() -> None: + # The map key must be the pure, unit-tested ChunkKey built via make_chunk_key + # (which always includes the generation), not an ad-hoc tuple that could drop + # the generation. + src = _source() + assert "using VmmIpcChunkKey = foundry::vmm_ipc::ChunkKey;" in src + assert 'foundry/vmm_ipc_chunk_key.h"' in src + # Every chunk-mapping key construction goes through the generation-inclusive + # factory; no bare make_tuple key with the generation appended. + assert ( + "foundry::vmm_ipc::make_chunk_key(exporter_pid, token, chunk_base, chunk_generation)" + in (src) + ) + assert "std::make_tuple(exporter_pid, token, chunk_base, chunk_generation)" not in src + + +def test_reject_fork_child_closes_listener_and_skips_export_map() -> None: + # The rejected child must fail fast: close the inherited listener fd but NOT + # iterate the inherited concurrent export-fd map (which cannot be locked + # safely here). This keeps the parked-fd limitation explicit, not "fixed". + body = _function_body(_source(), "static void vmm_ipc_reject_fork_child()") + assert "close(listen_fd)" in body, "must close the inherited listener fd" + assert "vmm_ipc_listen_fd = -1" in body, "must clear the listener fd after closing" + # No member access on the concurrent export-fd map (a bare mention in a + # comment explaining WHY it is skipped is fine; a `.method(` call is not). + for access in ("vmm_ipc_exported_fds.", "vmm_ipc_exported_fds["): + assert access not in body, ( + "must not access the inherited concurrent export-fd map from a fork child" + ) + + +def test_reject_fork_child_documents_parked_fd_pinning_limitation() -> None: + # The rejected child cannot lock the export-fd map, so it keeps inherited + # duplicate fds. The limitation (physical memory stays pinned until the + # child exits/execs) must be documented accurately, not glossed over. + body = _function_body(_source(), "static void vmm_ipc_reject_fork_child()") + assert "physical allocations pinned" in body, ( + "must document that inherited duplicate fds keep physical memory pinned" + ) + assert "until this child exits or execs" in body, ( + "must document the pin lasts until the child exits/execs" + ) + assert "FD_CLOEXEC" in body, "must note FD_CLOEXEC closes parked fds on exec" + + +def test_fork_child_rejection_is_async_signal_safe() -> None: + body = _function_body(_source(), "static void vmm_ipc_reject_fork_child()") + # Must emit via write(2) and must not use non-async-signal-safe calls. + assert "write(STDERR_FILENO" in body + assert "kForkChildRejectMessage" in body + for forbidden in ("fprintf", "snprintf", "sprintf", "malloc", "std::lock", "lock_guard"): + assert forbidden not in body, f"reject path must not call {forbidden}" + + +def test_claim_helper_itself_avoids_stdio_and_locks() -> None: + body = _function_body(_source(), "static bool vmm_ipc_claim_process()") + # The claim helper runs on the fork-child reject path; keep it stdio/lock free. + for forbidden in ("fprintf", "snprintf", "sprintf", "malloc", "lock_guard"): + assert forbidden not in body, f"claim helper must not call {forbidden}" diff --git a/tests/tp_validation_helpers.py b/tests/tp_validation_helpers.py new file mode 100644 index 00000000..203547af --- /dev/null +++ b/tests/tp_validation_helpers.py @@ -0,0 +1,371 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the Foundry project +"""Pure-Python archive evidence shared by TP validation harnesses.""" + +from __future__ import annotations + +import hashlib +import json +import os +import re +from collections.abc import Callable +from pathlib import Path + +# Every rank archive directory must contain exactly these non-graph files. +REQUIRED_RANK_FILES = { + "graph_manifest.json", + "fatbin_image_packed.img", + "fatbin_entrypoint_packed.txt", + "final_alloc_offset.json", +} + +# The integer index is whatever immediately follows the ``graph_`` prefix. +_GRAPH_INDEX_RE = re.compile(r"graph_(\d+)") + + +def sha256(path: Path) -> str: + digest = hashlib.sha256() + with path.open("rb") as file: + for chunk in iter(lambda: file.read(1024 * 1024), b""): + digest.update(chunk) + return digest.hexdigest() + + +def semantic_graph_sha256( + path: Path, + foundry_base_addr: int, + foundry_region_end: int, +) -> str: + """Hash stable graph semantics while retaining Foundry VMM pointers.""" + graph = json.loads(path.read_text()) + for generator in graph.get("generators", []): + generator.pop("id", None) + generator.pop("seed", None) + + vmm_kernel_pointers = [] + for node in graph.get("nodes", []): + if node.get("type") != "KernelNode": + continue + params = node["params"] + encoded_params = params.pop("kernelParams", []) + encoded_arg_buffer = params.pop("extra_argBuffer_hex", "") + params.pop("extra", None) + encoded_values = [ + (f"param:{param['index']}", param.get("value_hex", "")) for param in encoded_params + ] + encoded_values.append(("arg_buffer", encoded_arg_buffer)) + for source, encoded in encoded_values: + raw = bytes.fromhex(encoded) + for offset in range(0, max(0, len(raw) - 7), 8): + value = int.from_bytes(raw[offset : offset + 8], "little") + if foundry_base_addr <= value < foundry_region_end: + vmm_kernel_pointers.append([node["id"], source, offset, value]) + + graph["vmm_kernel_pointers"] = vmm_kernel_pointers + canonical = json.dumps(graph, sort_keys=True, separators=(",", ":")).encode() + return hashlib.sha256(canonical).hexdigest() + + +def archive_byte_fingerprints(root: Path) -> dict[str, str]: + """Byte-for-byte SHA-256 of every regular file under ``root``. + + Recurses the whole archive root and hashes each regular file (skipping + symlinks) keyed by its POSIX path relative to ``root``. Unlike the semantic + per-rank fingerprints, this covers the root ``warmup_state.json``, every + ``graph_*.cugraph`` binary, every JSON, the packed module image, and any + unexpected file — so a same-container / same-mount LOAD mutation of any of + them is detected byte-for-byte. + """ + fingerprints: dict[str, str] = {} + for path in sorted(root.rglob("*")): + if path.is_symlink() or not path.is_file(): + continue + fingerprints[path.relative_to(root).as_posix()] = sha256(path) + return fingerprints + + +def _recursive_strings(value: object): + if isinstance(value, str): + yield value + elif isinstance(value, dict): + for item in value.values(): + yield from _recursive_strings(item) + elif isinstance(value, (list, tuple)): + for item in value: + yield from _recursive_strings(item) + + +def _archive_graph_json_paths( + workspace: Path, + tp_size: int, +) -> dict[str, list[Path]]: + paths_by_rank: dict[str, list[Path]] = {} + for rank in range(tp_size): + rank_dir = workspace / f"rank_{rank}" + if not rank_dir.is_dir(): + raise ValueError(f"archive rank directory {rank_dir} does not exist") + paths = sorted( + path for path in rank_dir.glob("graph_*.json") if path.name != "graph_manifest.json" + ) + if not paths: + raise ValueError(f"archive rank directory {rank_dir} has no graph JSON files") + paths_by_rank[str(rank)] = paths + return paths_by_rank + + +def archive_graph_json_counts(workspace: Path, tp_size: int) -> dict[str, int]: + """Count concrete archived graph JSON files after validating rank inventory.""" + return { + rank: len(paths) for rank, paths in _archive_graph_json_paths(workspace, tp_size).items() + } + + +def archive_collective_kernel_counts( + workspace: Path, + tp_size: int, +) -> dict[str, dict[str, int]]: + """Count archived NCCL and symmetric collective KernelNodes per rank.""" + counts: dict[str, dict[str, int]] = {} + for rank, paths in _archive_graph_json_paths(workspace, tp_size).items(): + rank_counts = { + "all_reduce": 0, + "all_gather": 0, + "symmetric_all_reduce": 0, + "symmetric_all_gather": 0, + } + for path in paths: + try: + graph = json.loads(path.read_text()) + except json.JSONDecodeError as error: + raise ValueError( + f"archive graph {path} contains malformed JSON: {error}" + ) from error + if not isinstance(graph, dict): + raise ValueError(f"archive graph {path} root must be a JSON object") + nodes = graph.get("nodes") + if not isinstance(nodes, list): + raise ValueError(f"archive graph {path} field 'nodes' must be a list") + for index, node in enumerate(nodes): + if not isinstance(node, dict): + raise ValueError(f"archive graph {path} node {index} must be a JSON object") + if node.get("type") != "KernelNode": + continue + strings = tuple(_recursive_strings(node)) + rank_counts["all_reduce"] += any( + value.startswith("ncclDevKernel_AllReduce") for value in strings + ) + rank_counts["all_gather"] += any( + value.startswith("ncclDevKernel_AllGather") for value in strings + ) + rank_counts["symmetric_all_reduce"] += any( + "CUDASymmetricMemoryOps" in value and "all_reduce" in value for value in strings + ) + rank_counts["symmetric_all_gather"] += any( + "CUDASymmetricMemoryOps" in value and "all_gather" in value for value in strings + ) + counts[rank] = rank_counts + return counts + + +def archive_fingerprints( + workspace: Path, + tp_size: int, + foundry_base_addr: int, + foundry_region_end: int, +) -> dict[str, dict[str, str]]: + """Hash graph JSON semantically and stable archive files byte-for-byte.""" + fingerprints = {} + for rank in range(tp_size): + rank_dir = workspace / f"rank_{rank}" + files = [*rank_dir.glob("graph_*.json")] + files.extend( + path + for path in ( + rank_dir / "graph_manifest.json", + rank_dir / "fatbin_image_packed.img", + rank_dir / "fatbin_entrypoint_packed.txt", + rank_dir / "final_alloc_offset.json", + ) + if path.exists() + ) + fingerprints[str(rank)] = { + path.relative_to(rank_dir).as_posix(): ( + semantic_graph_sha256(path, foundry_base_addr, foundry_region_end) + if path.name.startswith("graph_") and path.name != "graph_manifest.json" + else sha256(path) + ) + for path in sorted(files) + } + return fingerprints + + +def _rank_graph_indices(rank_dir: Path, suffix: str) -> dict[int, str]: + """Map ``{index: stem}`` for every ``graph__*`` file. + + ``stem`` is the filename with ``suffix`` stripped so the JSON and cugraph + stems can be compared one-to-one. Raises on an unparsable name or a + repeated index (two files claiming the same ``graph_``). + """ + by_index: dict[int, str] = {} + for path in rank_dir.glob(f"graph_*{suffix}"): + if path.name == "graph_manifest.json": + continue + match = _GRAPH_INDEX_RE.match(path.name) + if match is None: + raise ValueError(f"archive graph file {path} has no parsable graph_ prefix") + index = int(match.group(1)) + stem = path.name[: -len(suffix)] + if index in by_index: + raise ValueError( + f"archive rank directory {rank_dir} has duplicate {suffix} files for " + f"index {index}: {by_index[index]}{suffix} and {stem}{suffix}" + ) + by_index[index] = stem + return by_index + + +def reject_symlinks(root: Path) -> None: + """Fail closed if any entry anywhere under ``root`` is a symlink. + + ``os.walk(followlinks=False)`` does not descend into symlinked directories, + but still lists them in ``dirnames`` so a symlink to a directory is caught + without traversing it. Rejecting every symlink keeps the archive's + regular-file universe identical to what :func:`archive_byte_fingerprints` + hashes (which skips symlinks), so a symlink can never smuggle bytes past the + byte map or point the inventory at a file outside the archive. + """ + root = Path(root) + # ``os.walk`` starts at ``root`` and follows it even with followlinks=False, + # so a symlinked root would be traversed silently. Reject it explicitly. + if root.is_symlink(): + raise ValueError(f"archive root is a symlink (rejected): {root}") + for dirpath, dirnames, filenames in os.walk(root, followlinks=False): + for name in (*dirnames, *filenames): + candidate = Path(dirpath) / name + if candidate.is_symlink(): + raise ValueError(f"archive contains a symlink (rejected): {candidate}") + + +def validate_archive_inventory( + root: Path, + tp_size: int, + expected_graph_count: int, + expected_stems: Callable[[int], str] | None = None, +) -> dict[str, dict[str, object]]: + """Require an exact, paired, contiguous per-rank archive. + + For each rank ``0..tp_size-1`` this enforces exactly ``expected_graph_count`` + ``.json`` graphs and exactly ``expected_graph_count`` ``.cugraph`` graphs, + identical stems for each index, contiguous indices ``0..expected-1``, and the + presence of every file in :data:`REQUIRED_RANK_FILES`. The archive root must + also carry ``warmup_state.json`` and contain no symlinks. Any violation + raises ``ValueError``. + + When ``expected_stems`` is given it is called with each index ``0..expected-1`` + and must return the exact descriptor stem that index's paired files must have; + a divergent stem is rejected even when the count, indices, and pairing are all + correct (an extra/malformed-but-paired descriptor). It is optional so callers + with a different profile (e.g. SGLang) keep the count/index/pairing contract + unchanged by simply not passing it. + """ + root = Path(root) + reject_symlinks(root) + warmup_state = root / "warmup_state.json" + if not warmup_state.is_file(): + raise ValueError(f"archive root is missing warmup_state.json: {warmup_state}") + + inventory: dict[str, dict[str, object]] = {} + expected_indices = set(range(expected_graph_count)) + for rank in range(tp_size): + rank_dir = root / f"rank_{rank}" + if not rank_dir.is_dir(): + raise ValueError(f"archive rank directory {rank_dir} does not exist") + + json_by_index = _rank_graph_indices(rank_dir, ".json") + cugraph_by_index = _rank_graph_indices(rank_dir, ".cugraph") + + if len(json_by_index) != expected_graph_count: + raise ValueError( + f"archive rank_{rank} has {len(json_by_index)} graph JSON files, " + f"expected exactly {expected_graph_count}" + ) + if len(cugraph_by_index) != expected_graph_count: + raise ValueError( + f"archive rank_{rank} has {len(cugraph_by_index)} graph cugraph files, " + f"expected exactly {expected_graph_count}" + ) + if set(json_by_index) != expected_indices: + raise ValueError( + f"archive rank_{rank} graph JSON indices are not the contiguous range " + f"0..{expected_graph_count - 1}: {sorted(json_by_index)}" + ) + if set(cugraph_by_index) != expected_indices: + raise ValueError( + f"archive rank_{rank} graph cugraph indices are not the contiguous range " + f"0..{expected_graph_count - 1}: {sorted(cugraph_by_index)}" + ) + for index in expected_indices: + if json_by_index[index] != cugraph_by_index[index]: + raise ValueError( + f"archive rank_{rank} index {index} has mismatched stems: " + f"{json_by_index[index]}.json vs {cugraph_by_index[index]}.cugraph" + ) + if expected_stems is not None: + for index in sorted(expected_indices): + expected = expected_stems(index) + if json_by_index[index] != expected: + raise ValueError( + f"archive rank_{rank} index {index} descriptor " + f"{json_by_index[index]!r} != expected {expected!r}" + ) + for required in sorted(REQUIRED_RANK_FILES): + if not (rank_dir / required).is_file(): + raise ValueError(f"archive rank_{rank} is missing required file {required}") + + inventory[str(rank)] = { + "graph_count": expected_graph_count, + "indices": sorted(expected_indices), + "stems": {str(index): json_by_index[index] for index in sorted(expected_indices)}, + } + return inventory + + +def _read_rank_offsets(root: Path, tp_size: int) -> dict[str, int]: + offsets: dict[str, int] = {} + for rank in range(tp_size): + path = root / f"rank_{rank}" / "final_alloc_offset.json" + offsets[str(rank)] = int(json.loads(path.read_text())["final_alloc_offset"]) + return offsets + + +def collect_archive_snapshot( + root: Path, + tp_size: int, + expected_graph_count: int, + foundry_base_addr: int, + foundry_region_end: int, + expected_stems: Callable[[int], str] | None = None, +) -> dict[str, object]: + """Validate the exact inventory and return every immutability map. + + Returns the strict per-rank ``inventory``, the per-rank ``semantic_fingerprints`` + (which include the packed entrypoint), the flat ``byte_fingerprints`` covering + every regular file under ``root``, the per-rank ``graph_counts``, and the + per-rank final-allocation ``offsets``. Raises ``ValueError`` if the inventory + is not exact. ``expected_stems`` is forwarded to + :func:`validate_archive_inventory` for exact descriptor-profile enforcement. + """ + root = Path(root) + inventory = validate_archive_inventory(root, tp_size, expected_graph_count, expected_stems) + return { + "inventory": inventory, + "semantic_fingerprints": archive_fingerprints( + root, + tp_size, + foundry_base_addr, + foundry_region_end, + ), + "byte_fingerprints": archive_byte_fingerprints(root), + "graph_counts": {rank: info["graph_count"] for rank, info in inventory.items()}, + "offsets": _read_rank_offsets(root, tp_size), + } diff --git a/tests/vllm_tp_validation.py b/tests/vllm_tp_validation.py new file mode 100644 index 00000000..b957d475 --- /dev/null +++ b/tests/vllm_tp_validation.py @@ -0,0 +1,115 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the Foundry project +"""Pure-Python evidence helpers for the vLLM TP Modal harness.""" + +from __future__ import annotations + +import os +import re +import signal +import subprocess +from typing import IO + +FULL_COMMIT_SHA_PATTERN = re.compile(r"^[0-9a-fA-F]{40}$") +FATAL_LOG_PATTERN = re.compile( + r"\bERROR\b|RuntimeError|Traceback|(?:worker|process).*(?:died|death)|" + r"Engine core proc died|\bException\b|CUDA error|illegal memory access|" + r"an illegal memory|NCCL WARN|NCCL error|MMU fault|\bXid\b|" + r"segmentation fault|\bsegfault\b", + re.IGNORECASE, +) +LOAD_RESTORATION_COMPLETE_PATTERN = re.compile( + r"\[foundry\] LOAD restoration complete rank_(?P\d+): " + r"(?P\d+)/(?P\d+) graphs" +) +# Live per-rank VMM cursor emitted immediately after ``capture_model`` returns +# on LOAD. This is the real allocator watermark, not a reread of the archived +# ``final_alloc_offset.json``, so the harness can prove the live LOAD trajectory +# lands exactly on SAVE2's offset. +LOAD_LIVE_FINAL_ALLOC_OFFSET_PATTERN = re.compile( + r"\[foundry\] LOAD live final_alloc_offset rank_(?P\d+): (?P\d+)" +) +NCCL_VERSION_PATTERN = re.compile(r"\bNCCL version (?P\S+)") +EXPECTED_SHUTDOWN_RETURNCODES = {0, -2, 130} + + +def require_commit_sha(name: str, value: str | None) -> str: + if value is None or FULL_COMMIT_SHA_PATTERN.fullmatch(value) is None: + raise ValueError(f"{name} must be a full 40-hex commit SHA, got {value!r}") + return value.lower() + + +def scan_log_text(text: str) -> dict: + errors = sorted({line for line in text.splitlines() if FATAL_LOG_PATTERN.search(line)}) + completions: dict[str, list[list[int]]] = {} + for match in LOAD_RESTORATION_COMPLETE_PATTERN.finditer(text): + completions.setdefault(match.group("rank"), []).append( + [int(match.group("count")), int(match.group("total"))] + ) + live_offsets: dict[str, list[int]] = {} + for match in LOAD_LIVE_FINAL_ALLOC_OFFSET_PATTERN.finditer(text): + live_offsets.setdefault(match.group("rank"), []).append(int(match.group("offset"))) + nccl_versions = {match.group("version") for match in NCCL_VERSION_PATTERN.finditer(text)} + if len(nccl_versions) > 1: + raise RuntimeError(f"vLLM child logs reported conflicting NCCL versions: {nccl_versions}") + return { + "errors": errors, + "load_restoration_completions": completions, + "load_final_alloc_offsets": live_offsets, + "nccl_version": next(iter(nccl_versions), None), + } + + +def validate_shutdown_returncode(returncode: int, died_before_teardown: bool) -> None: + if died_before_teardown: + raise RuntimeError(f"vLLM process died before intentional teardown (code {returncode})") + if returncode not in EXPECTED_SHUTDOWN_RETURNCODES: + raise RuntimeError(f"vLLM process exited with unexpected code {returncode}") + + +def stop_process( + process: subprocess.Popen, + log_file: IO[str], + timeout: float = 120, +) -> int: + died_before_teardown = process.poll() is not None + try: + if not died_before_teardown: + os.killpg(os.getpgid(process.pid), signal.SIGINT) + try: + process.wait(timeout=timeout) + except subprocess.TimeoutExpired: + os.killpg(os.getpgid(process.pid), signal.SIGKILL) + process.wait(timeout=30) + finally: + log_file.close() + if process.returncode is None: + raise RuntimeError("vLLM process did not produce a return code after teardown") + validate_shutdown_returncode(process.returncode, died_before_teardown) + return process.returncode + + +def completion_evidence(response: dict) -> dict[str, str | list[str]]: + try: + choice = response["choices"][0] + text = choice["text"] + token_strings = choice["logprobs"]["tokens"] + except (IndexError, KeyError, TypeError) as error: + raise RuntimeError( + "vLLM response did not include completion text and token logprobs" + ) from error + if ( + not isinstance(text, str) + or not text + or not isinstance(token_strings, list) + or not token_strings + or not all(isinstance(token, str) for token in token_strings) + or not all(token for token in token_strings) + ): + raise RuntimeError( + "vLLM response completion evidence must contain nonempty text and token strings" + ) + return { + "text": text, + "token_strings": token_strings, + } diff --git a/tests/vmm_cleanup_test.cpp b/tests/vmm_cleanup_test.cpp new file mode 100644 index 00000000..61482847 --- /dev/null +++ b/tests/vmm_cleanup_test.cpp @@ -0,0 +1,228 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright contributors to the Foundry project +// +// Pure host-policy regression for retryable VMM cleanup. Driver operations are +// injected lambdas, so failures and retries are exercised without CUDA. + +#include +#include +#include +#include +#include "foundry/vmm_cleanup.h" + +namespace { + +int failures = 0; + +struct CallCounts { + int unmap = 0; + int release = 0; + int address_free = 0; +}; + +void expect(bool cond, const char* what) { + if (!cond) { + fprintf(stderr, "[FAIL] %s\n", what); + ++failures; + } +} + +struct Faults { + std::vector unmap; + std::vector release; + std::vector address_free; + size_t unmap_index = 0; + size_t release_index = 0; + size_t address_free_index = 0; + int unmap_calls = 0; + int release_calls = 0; + int address_free_calls = 0; + CallCounts* observed = nullptr; + + int next_unmap() { + ++unmap_calls; + if (observed != nullptr) + ++observed->unmap; + return unmap.at(unmap_index++); + } + + int next_release() { + ++release_calls; + if (observed != nullptr) + ++observed->release; + return release.at(release_index++); + } + + int next_address_free() { + ++address_free_calls; + if (observed != nullptr) + ++observed->address_free; + return address_free.at(address_free_index++); + } +}; + +} // namespace + +int main() { + using foundry::vmm::cleanup_complete; + using foundry::vmm::CleanupPhase; + using foundry::vmm::CleanupState; + using foundry::vmm::drain_pending_cleanups; + using foundry::vmm::run_cleanup; + + constexpr int kSuccess = 0; + constexpr int kUnmapFailure = 11; + constexpr int kReleaseFailure = 22; + constexpr int kAddressFreeFailure = 33; + constexpr bool kIpcImportAddressFreeRequired = true; + constexpr bool kRegionOwnedAddressFreeRequired = false; + + { + CleanupState state; + Faults faults{ + {kUnmapFailure, kSuccess}, {kReleaseFailure, kSuccess}, {kAddressFreeFailure, kSuccess}}; + std::vector> logged; + int status = run_cleanup( + state, kSuccess, kIpcImportAddressFreeRequired, [&] { return faults.next_unmap(); }, + [&] { return faults.next_release(); }, [&] { return faults.next_address_free(); }, + [&](CleanupPhase phase, int result) { logged.emplace_back(phase, result); }); + + expect(status == kUnmapFailure, "first attempt returns the first failure"); + expect(state.started, "cleanup records that an attempt started"); + expect(!cleanup_complete(state, kIpcImportAddressFreeRequired), + "all-failed IPC import cleanup remains incomplete"); + expect(faults.unmap_calls == 1 && faults.release_calls == 1 && faults.address_free_calls == 1, + "all required phases are attempted despite earlier failures"); + expect(logged.size() == 3, "every failed phase is reported"); + + status = run_cleanup( + state, kSuccess, kIpcImportAddressFreeRequired, [&] { return faults.next_unmap(); }, + [&] { return faults.next_release(); }, [&] { return faults.next_address_free(); }, + [&](CleanupPhase phase, int result) { logged.emplace_back(phase, result); }); + expect(status == kSuccess, "retry succeeds after injected failures clear"); + expect(cleanup_complete(state, kIpcImportAddressFreeRequired), + "successful IPC import retry completes all three owned phases"); + expect(faults.unmap_calls == 2 && faults.release_calls == 2 && faults.address_free_calls == 2, + "failed phases are retried exactly once"); + } + + { + CleanupState state; + Faults faults{{kSuccess}, {kReleaseFailure, kSuccess}, {kSuccess}}; + int status = run_cleanup( + state, kSuccess, kIpcImportAddressFreeRequired, [&] { return faults.next_unmap(); }, + [&] { return faults.next_release(); }, [&] { return faults.next_address_free(); }, + [](CleanupPhase, int) {}); + expect(status == kReleaseFailure, "middle-phase failure is returned"); + expect(state.unmapped && !state.released && state.address_freed, + "successful phases are retained around a failed phase"); + + status = run_cleanup( + state, kSuccess, kIpcImportAddressFreeRequired, [&] { return faults.next_unmap(); }, + [&] { return faults.next_release(); }, [&] { return faults.next_address_free(); }, + [](CleanupPhase, int) {}); + expect(status == kSuccess, "partial cleanup retry succeeds"); + expect(faults.unmap_calls == 1 && faults.release_calls == 2 && faults.address_free_calls == 1, + "retry skips completed phases without double cleanup"); + } + + { + CleanupState state; + Faults faults{{kSuccess}, {kReleaseFailure, kSuccess}, {}}; + int status = run_cleanup( + state, kSuccess, kRegionOwnedAddressFreeRequired, [&] { return faults.next_unmap(); }, + [&] { return faults.next_release(); }, + [&] { + ++faults.address_free_calls; + return kAddressFreeFailure; + }, + [](CleanupPhase, int) {}); + expect(status == kReleaseFailure, "region-owned cleanup reports release failure"); + expect(!cleanup_complete(state, kRegionOwnedAddressFreeRequired), + "failed region-owned two-phase cleanup remains incomplete"); + expect(faults.address_free_calls == 0, + "region-owned allocation never frees a subrange of the whole region VA reservation"); + + status = run_cleanup( + state, kSuccess, kRegionOwnedAddressFreeRequired, [&] { return faults.next_unmap(); }, + [&] { return faults.next_release(); }, + [&] { + ++faults.address_free_calls; + return kAddressFreeFailure; + }, + [](CleanupPhase, int) {}); + expect(status == kSuccess && cleanup_complete(state, kRegionOwnedAddressFreeRequired), + "region-owned retry completes only unmap and release"); + expect(faults.unmap_calls == 1 && faults.release_calls == 2, + "region-owned retry skips its completed unmap"); + } + + { + struct PendingResource { + CleanupState cleanup; + Faults faults; + }; + + CallCounts loser_calls; + PendingResource loser{ + {}, + {{kUnmapFailure, kUnmapFailure, kSuccess}, + {kReleaseFailure, kReleaseFailure, kSuccess}, + {kAddressFreeFailure, kAddressFreeFailure, kSuccess}, + 0, + 0, + 0, + 0, + 0, + 0, + &loser_calls}, + }; + CallCounts collision_calls; + PendingResource collision{ + {}, + {{kSuccess}, {kReleaseFailure, kSuccess}, {kSuccess}, 0, 0, 0, 0, 0, 0, &collision_calls}, + }; + + auto retry = [&](PendingResource& pending) { + return run_cleanup( + pending.cleanup, kSuccess, kIpcImportAddressFreeRequired, + [&] { return pending.faults.next_unmap(); }, + [&] { return pending.faults.next_release(); }, + [&] { return pending.faults.next_address_free(); }, [](CleanupPhase, int) {}); + }; + auto complete = [](const PendingResource& pending) { + return cleanup_complete(pending.cleanup, kIpcImportAddressFreeRequired); + }; + + int status = retry(loser); + expect(status == kUnmapFailure, "race loser initial cleanup returns its first failure"); + status = retry(collision); + expect(status == kReleaseFailure, "collision initial cleanup reports release failure"); + + std::multimap pending; + pending.emplace(0x1000, std::move(loser)); + pending.emplace(0x2000, std::move(collision)); + status = drain_pending_cleanups(pending, kSuccess, retry, complete); + expect(status == kUnmapFailure, "pending drain returns first retained cleanup failure"); + expect(pending.size() == 1 && pending.begin()->first == 0x1000, + "completed collision cleanup is erased while failed race loser persists"); + expect(loser_calls.unmap == 2 && loser_calls.release == 2 && loser_calls.address_free == 2, + "failed loser phases are retried and retained"); + expect(collision_calls.unmap == 1 && collision_calls.release == 2 && + collision_calls.address_free == 1, + "collision retry skips successful unmap and address-free phases"); + + status = drain_pending_cleanups(pending, kSuccess, retry, complete); + expect(status == kSuccess && pending.empty(), + "later open drains the retained loser after failures clear"); + expect(loser_calls.unmap == 3 && loser_calls.release == 3 && loser_calls.address_free == 3, + "final loser drain retries each failed phase exactly once"); + } + + if (failures == 0) { + printf("[PASS] VMM cleanup retry policy contract\n"); + return 0; + } + fprintf(stderr, "[FAIL] %d assertion(s) failed\n", failures); + return 1; +} diff --git a/tests/vmm_driver_bindings.py b/tests/vmm_driver_bindings.py new file mode 100644 index 00000000..2075049e --- /dev/null +++ b/tests/vmm_driver_bindings.py @@ -0,0 +1,92 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the Foundry project +"""Torch-free CUDA driver bindings for the native VMM lifecycle tests.""" + +from __future__ import annotations + +import ctypes +import types + + +def _bind(library, name, argtypes, source, restype=ctypes.c_int): + try: + function = getattr(library, name) + except AttributeError as exc: + raise AttributeError( + f"required CUDA driver symbol {name} is unavailable from {source}" + ) from exc + function.argtypes = argtypes + function.restype = restype + return function + + +def bind_driver_libraries(process_global, cuda_driver, ipc_handle_type): + """Bind intercepted calls globally and unhooked utilities from libcuda.""" + u64 = ctypes.c_uint64 + + return types.SimpleNamespace( + # These calls exercise Foundry behavior and must resolve through the + # LD_PRELOAD interposer in the process-global symbol table. + mem_alloc=_bind( + process_global, + "cuMemAlloc_v2", + [ctypes.POINTER(u64), ctypes.c_size_t], + "process global symbols", + ), + mem_free=_bind( + process_global, + "cuMemFree_v2", + [u64], + "process global symbols", + ), + ipc_get=_bind( + process_global, + "cuIpcGetMemHandle", + [ctypes.POINTER(ipc_handle_type), u64], + "process global symbols", + ), + ipc_open=_bind( + process_global, + "cuIpcOpenMemHandle", + [ctypes.POINTER(u64), ipc_handle_type, ctypes.c_uint], + "process global symbols", + ), + ipc_close=_bind( + process_global, + "cuIpcCloseMemHandle", + [u64], + "process global symbols", + ), + addr_reserve=_bind( + process_global, + "cuMemAddressReserve", + [ctypes.POINTER(u64), ctypes.c_size_t, ctypes.c_size_t, u64, ctypes.c_ulonglong], + "process global symbols", + ), + addr_free=_bind( + process_global, + "cuMemAddressFree", + [u64, ctypes.c_size_t], + "process global symbols", + ), + # CUDA 13 exposes the size_t utility ABIs under exact v2 symbols. + memset_d8=_bind( + cuda_driver, + "cuMemsetD8_v2", + [u64, ctypes.c_ubyte, ctypes.c_size_t], + "libcuda.so.1", + ), + memcpy_dtoh=_bind( + cuda_driver, + "cuMemcpyDtoH_v2", + [ctypes.c_void_p, u64, ctypes.c_size_t], + "libcuda.so.1", + ), + ) + + +def load_driver_bindings(ipc_handle_type, load_library=ctypes.CDLL): + """Load both symbol scopes required by the native lifecycle tests.""" + process_global = load_library(None) + cuda_driver = load_library("libcuda.so.1") + return bind_driver_libraries(process_global, cuda_driver, ipc_handle_type) diff --git a/tests/vmm_ipc_chunk_key_test.cpp b/tests/vmm_ipc_chunk_key_test.cpp new file mode 100644 index 00000000..76ddde3c --- /dev/null +++ b/tests/vmm_ipc_chunk_key_test.cpp @@ -0,0 +1,90 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright contributors to the Foundry project +// +// Pure CPU contract for the generation-sensitive VMM-IPC chunk key. No CUDA: +// the key is a plain value type so its equality/ordering/hash identity can be +// proven on a CPU-only host. The crucial property is that the generation is an +// identity-bearing field, so a chunk freed and recreated at the same base is a +// DISTINCT map key; dropping the generation would collapse the two onto one key +// and reuse a stale mapping. + +#include +#include +#include + +#include "foundry/vmm_ipc_chunk_key.h" + +static int g_failures = 0; + +static void expect(bool cond, const char* what) { + if (!cond) { + fprintf(stderr, "[FAIL] %s\n", what); + g_failures++; + } +} + +int main() { + using foundry::vmm_ipc::chunk_identity_without_generation; + using foundry::vmm_ipc::ChunkKey; + using foundry::vmm_ipc::ChunkKeyHash; + using foundry::vmm_ipc::make_chunk_key; + + const long pid = 4321; + const unsigned long long token = 0xDEADBEEFCAFEULL; + const unsigned long long base = 0x600000000000ULL; + + ChunkKey g1 = make_chunk_key(pid, token, base, 1); + ChunkKey g2 = make_chunk_key(pid, token, base, 2); + ChunkKey g1_again = make_chunk_key(pid, token, base, 1); + + // Same pid/token/base but different generation -> distinct keys. + expect(g1 != g2, "same pid/token/base, different generation -> not equal"); + expect(!(g1 == g2), "same pid/token/base, different generation -> operator== false"); + expect(g1 == g1_again, "identical fields -> equal"); + + // Strict-weak ordering: exactly one of a m; + m[g1] = 10; + m[g2] = 20; + expect(m.size() == 2, "map keeps distinct entries per generation"); + expect(m[g1] == 10 && m[g2] == 20, "map does not alias generations"); + // Recreating g1 addresses the same slot (usability), not g2's. + m[g1_again] = 11; + expect(m.size() == 2 && m[g1] == 11, "recreated identical key hits the same slot"); + } + + // Hash is generation-inclusive too (consistent with equality). + { + std::unordered_map m; + m[g1] = 1; + m[g2] = 2; + expect(m.size() == 2, "hashed map keeps distinct entries per generation"); + } + + // Negative contract: REMOVING the generation must fail. A generation-less + // identity aliases the two distinct chunks onto one key, which is exactly the + // stale-mapping-reuse bug the generation field prevents. + expect(chunk_identity_without_generation(g1) == chunk_identity_without_generation(g2), + "dropping generation collapses distinct recreated chunks onto one key (bug)"); + { + std::map, int> without_gen; + without_gen[chunk_identity_without_generation(g1)] = 10; + without_gen[chunk_identity_without_generation(g2)] = 20; + expect(without_gen.size() == 1, + "generation-less key wrongly collapses two generations into one slot"); + } + + if (g_failures == 0) { + printf("[PASS] vmm_ipc chunk-key generation identity contract\n"); + return 0; + } + fprintf(stderr, "[FAIL] %d chunk-key contract checks failed\n", g_failures); + return 1; +} diff --git a/tests/vmm_ipc_ownership_test.cpp b/tests/vmm_ipc_ownership_test.cpp new file mode 100644 index 00000000..919b5720 --- /dev/null +++ b/tests/vmm_ipc_ownership_test.cpp @@ -0,0 +1,65 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright contributors to the Foundry project +// +// Pure-policy regression for VMM-IPC process ownership. No CUDA, no threads: +// exercises the claim/match decision and the owner-only-state predicate that +// the production hook wires into an atomic owner pid, plus the fixed +// async-signal-safe fork-child rejection message. + +#include +#include +#include "foundry/vmm_ipc_ownership.h" + +namespace { + +int failures = 0; + +void expect(bool cond, const char* what) { + if (!cond) { + fprintf(stderr, "[FAIL] %s\n", what); + ++failures; + } +} + +} // namespace + +int main() { + using foundry::vmm_ipc::decide_ipc_ownership; + using foundry::vmm_ipc::kForkChildRejectMessage; + using foundry::vmm_ipc::kNoOwnerPid; + using foundry::vmm_ipc::OwnershipDecision; + using foundry::vmm_ipc::owns_ipc_state; + + const long owner = 4321; + const long child = 8765; + + // Claim/match decision. + expect(decide_ipc_ownership(kNoOwnerPid, child) == OwnershipDecision::kProceed, + "unclaimed owner -> proceed (pre-claim fork / first use allowed)"); + expect(decide_ipc_ownership(owner, owner) == OwnershipDecision::kProceed, + "same owner -> proceed (owner may reuse IPC)"); + expect(decide_ipc_ownership(owner, child) == OwnershipDecision::kRejectForkChild, + "different owner -> reject post-claim fork child"); + + // Owner-only state (export-fd registry / invalidation guard). + expect(!owns_ipc_state(kNoOwnerPid, child), + "unclaimed process owns no IPC state (invalidation returns before map)"); + expect(owns_ipc_state(owner, owner), "owner owns its IPC state"); + expect(!owns_ipc_state(owner, child), "fork child does not own inherited IPC state"); + + // Fixed async-signal-safe rejection message contract. + expect(strstr(kForkChildRejectMessage, "use spawn/exec") != nullptr, + "reject message includes the 'use spawn/exec' remediation"); + expect(strncmp(kForkChildRejectMessage, "[HOOK] ERROR:", 13) == 0, + "reject message uses the standard [HOOK] ERROR: prefix"); + const size_t len = sizeof(kForkChildRejectMessage) - 1; + expect(len > 0 && kForkChildRejectMessage[len - 1] == '\n', + "reject message is newline-terminated"); + + if (failures == 0) { + printf("[PASS] vmm_ipc ownership policy contract\n"); + return 0; + } + fprintf(stderr, "[FAIL] %d assertion(s) failed\n", failures); + return 1; +} diff --git a/tests/vmm_ipc_reservation_test.cpp b/tests/vmm_ipc_reservation_test.cpp new file mode 100644 index 00000000..a3e9772b --- /dev/null +++ b/tests/vmm_ipc_reservation_test.cpp @@ -0,0 +1,148 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright contributors to the Foundry project +// +// Pure-policy regression for the dedicated peer-import VA zone. No CUDA calls: +// the reservation/release side effects are injected as callbacks so the +// fail-closed decision can be exercised on a CPU-only host. + +#include +#include "foundry/vmm_ipc_reservation.h" + +namespace { + +struct Recorder { + int free_calls = 0; + int release_calls = 0; + CUdeviceptr freed_ptr = 0; + size_t freed_size = 0; + CUmemGenericAllocationHandle released_handle = 0; + // Statuses the injected cleanup callbacks return (default success). + CUresult free_status = CUDA_SUCCESS; + CUresult release_status = CUDA_SUCCESS; + + CUresult on_free(CUdeviceptr p, size_t s) { + ++free_calls; + freed_ptr = p; + freed_size = s; + return free_status; + } + CUresult on_release(CUmemGenericAllocationHandle h) { + ++release_calls; + released_handle = h; + return release_status; + } +}; + +int failures = 0; + +void expect(bool cond, const char* what) { + if (!cond) { + fprintf(stderr, "[FAIL] %s\n", what); + ++failures; + } +} + +CUresult run(Recorder& rec, CUresult result, CUdeviceptr actual, CUdeviceptr expected, size_t size, + CUmemGenericAllocationHandle handle) { + return foundry::vmm_ipc::require_exact_ipc_reservation( + result, actual, expected, size, handle, + [&](CUdeviceptr p, size_t s) { return rec.on_free(p, s); }, + [&](CUmemGenericAllocationHandle h) { return rec.on_release(h); }); +} + +} // namespace + +int main() { + const CUdeviceptr expected = 0x300000000000ULL; + const size_t size = 2ULL * 1024 * 1024; + const CUmemGenericAllocationHandle handle = 0xABCDEF12ULL; + const CUdeviceptr wrong = 0x7F0000000000ULL; + + // Case 1: success exactly at the requested VA -> no cleanup, return success. + { + Recorder rec; + CUresult r = run(rec, CUDA_SUCCESS, expected, expected, size, handle); + expect(r == CUDA_SUCCESS, "case1: returns CUDA_SUCCESS"); + expect(rec.free_calls == 0, "case1: address-free not called"); + expect(rec.release_calls == 0, "case1: handle-release not called"); + } + + // Case 2: driver error -> preserve the error, no cleanup (caller owns handle). + { + Recorder rec; + CUresult r = run(rec, CUDA_ERROR_OUT_OF_MEMORY, 0, expected, size, handle); + expect(r == CUDA_ERROR_OUT_OF_MEMORY, "case2: preserves driver error"); + expect(rec.free_calls == 0, "case2: address-free not called"); + expect(rec.release_calls == 0, "case2: handle-release not called"); + } + + // Case 3: success at the wrong VA -> free reservation + release handle once, + // then fail closed. + { + Recorder rec; + CUresult r = run(rec, CUDA_SUCCESS, wrong, expected, size, handle); + expect(r == CUDA_ERROR_OUT_OF_MEMORY, "case3: fails closed with OOM"); + expect(rec.free_calls == 1, "case3: address-free called exactly once"); + expect(rec.release_calls == 1, "case3: handle-release called exactly once"); + expect(rec.freed_ptr == wrong, "case3: freed the actual reservation VA"); + expect(rec.freed_size == size, "case3: freed the reservation size"); + expect(rec.released_handle == handle, "case3: released the imported handle"); + } + + // Case 4: wrong VA with a failing address-free -> the handle release still + // runs exactly once (no leak), and the result stays fail-closed. + { + Recorder rec; + rec.free_status = CUDA_ERROR_INVALID_VALUE; + CUresult r = run(rec, CUDA_SUCCESS, wrong, expected, size, handle); + expect(r == CUDA_ERROR_OUT_OF_MEMORY, "case4: still fails closed with OOM"); + expect(rec.free_calls == 1, "case4: address-free called exactly once"); + expect(rec.release_calls == 1, "case4: handle-release still called exactly once"); + expect(rec.released_handle == handle, "case4: released the imported handle"); + } + + // Case 5: wrong VA with both cleanups failing -> each invoked exactly once + // (no double release), result stays fail-closed. + { + Recorder rec; + rec.free_status = CUDA_ERROR_INVALID_VALUE; + rec.release_status = CUDA_ERROR_INVALID_VALUE; + CUresult r = run(rec, CUDA_SUCCESS, wrong, expected, size, handle); + expect(r == CUDA_ERROR_OUT_OF_MEMORY, "case5: still fails closed with OOM"); + expect(rec.free_calls == 1, "case5: address-free called exactly once"); + expect(rec.release_calls == 1, "case5: handle-release called exactly once (no double release)"); + } + + // --- require_fallback_reservation_freed (runs before the zone reservation) --- + auto run_fallback = [](Recorder& rec, CUresult free_status, + CUmemGenericAllocationHandle h) -> CUresult { + rec.free_status = CUDA_SUCCESS; // unused here; release status drives rec + return foundry::vmm_ipc::require_fallback_reservation_freed( + free_status, h, [&](CUmemGenericAllocationHandle hh) { return rec.on_release(hh); }); + }; + + // Case 6: clean fallback free -> proceed (SUCCESS), no handle release. + { + Recorder rec; + CUresult r = run_fallback(rec, CUDA_SUCCESS, handle); + expect(r == CUDA_SUCCESS, "case6: clean fallback free -> proceed to zone"); + expect(rec.release_calls == 0, "case6: handle not released on clean free"); + } + + // Case 7: failed fallback free -> release handle once, fail closed with OOM + // (caller must not continue to the dedicated-zone mapping). + { + Recorder rec; + CUresult r = run_fallback(rec, CUDA_ERROR_INVALID_VALUE, handle); + expect(r == CUDA_ERROR_OUT_OF_MEMORY, "case7: failed fallback free -> fail closed OOM"); + expect(rec.release_calls == 1, "case7: handle released exactly once (no double release)"); + expect(rec.released_handle == handle, "case7: released the imported handle"); + } + + if (failures == 0) { + printf("[PASS] require_exact_ipc_reservation policy contract\n"); + return 0; + } + fprintf(stderr, "[FAIL] %d assertion(s) failed\n", failures); + return 1; +}