From b86183ed438f2f01ef1a8cb5553c335448c1bd89 Mon Sep 17 00:00:00 2001 From: Jason Mancuso <7891333+jvmncs@users.noreply.github.com> Date: Mon, 29 Jun 2026 17:44:58 -0400 Subject: [PATCH 01/18] feat(miles_disagg): per-node host-RAM monitor for the kimi K2.6 trainer The trainer can OOM-kill on host RAM (the publish/update_weights full-model gather is the peak consumer at ~1.69 TiB on the ~1.95 TiB B200:8 node), and Megatron only reports GPU memory. helpers.start_host_mem_monitor() logs a per-node [hostmem] trace from a daemon thread (called in Trainer.start_ray on every node) so 'modal app logs -f' shows the phase/peak that blows the budget. --- cookbook/miles_disagg/helpers.py | 52 ++++++++++++++++++++++++++++ cookbook/miles_disagg/modal_train.py | 5 +++ 2 files changed, 57 insertions(+) diff --git a/cookbook/miles_disagg/helpers.py b/cookbook/miles_disagg/helpers.py index 9d50018..c14ceeb 100644 --- a/cookbook/miles_disagg/helpers.py +++ b/cookbook/miles_disagg/helpers.py @@ -221,3 +221,55 @@ def _post_json(url: str, payload: dict, *, timeout: float) -> dict: ) with urllib.request.urlopen(request, timeout=timeout) as resp: return json.load(resp) + + +def start_host_mem_monitor(interval_s: int = 20) -> None: + """Log this node's host-RAM trajectory to stdout from a daemon thread. + + The trainer can OOM-kill on host-RAM exhaustion (the publish/update_weights + full-model gather is the peak consumer), but Megatron only reports GPU memory + and the kill leaves no durable peak behind. This logs MemTotal/MemAvailable + + the container cgroup usage every ``interval_s`` so a live ``modal app logs -f`` + shows exactly which phase blows the ~1.95 TiB B200:8 node and how high it peaks. + Runs on EVERY node (called from the SPMD enter()), so whichever rank OOMs has + its own trace. Best-effort: never raises.""" + import threading + + host = socket.gethostname() + + def _meminfo() -> tuple[float, float]: + total = avail = 0.0 + try: + with open("/proc/meminfo") as f: + for line in f: + if line.startswith("MemTotal:"): + total = int(line.split()[1]) / 1024 / 1024 # GiB + elif line.startswith("MemAvailable:"): + avail = int(line.split()[1]) / 1024 / 1024 + except Exception: # noqa: BLE001 + pass + return total, avail + + def _cgroup_used_gib() -> float: + for path in ("/sys/fs/cgroup/memory.current", # cgroup v2 + "/sys/fs/cgroup/memory/memory.usage_in_bytes"): # v1 + try: + with open(path) as f: + return int(f.read().strip()) / 1024**3 + except Exception: # noqa: BLE001 + continue + return -1.0 + + def _loop() -> None: + while True: + total, avail = _meminfo() + used = total - avail + cg = _cgroup_used_gib() + print( + f"[hostmem] {host} used={used:.0f}GiB avail={avail:.0f}GiB " + f"total={total:.0f}GiB cgroup_used={cg:.0f}GiB", + flush=True, + ) + time.sleep(interval_s) + + threading.Thread(target=_loop, daemon=True, name="host-mem-monitor").start() diff --git a/cookbook/miles_disagg/modal_train.py b/cookbook/miles_disagg/modal_train.py index b3a9686..8811748 100644 --- a/cookbook/miles_disagg/modal_train.py +++ b/cookbook/miles_disagg/modal_train.py @@ -411,6 +411,11 @@ class Trainer: def start_ray(self) -> None: rank, master_addr, my_ip = helpers.get_modal_cluster_context(N_TRAIN_NODES) self.rank = rank + # Per-node host-RAM trace: the trainer can OOM-kill at host-RAM exhaustion + # (the publish/update_weights gather is the peak), and Modal's kill leaves no + # peak behind. This makes `modal app logs -f` show which node/phase blows the + # ~1.95 TiB budget. Runs on every node (SPMD enter()). + helpers.start_host_mem_monitor() os.environ.update( { "MILES_HOST_IP": my_ip, From 6f143fbc1d42e5418e7481ca15cd91c33dde053b Mon Sep 17 00:00:00 2001 From: Jason Mancuso <7891333+jvmncs@users.noreply.github.com> Date: Mon, 29 Jun 2026 19:05:17 -0700 Subject: [PATCH 02/18] docs: remove overlap-drain-fix references from in_place commit mode (#2) The overlap-drain race was never reproduced and may have been theoretical; stop claiming in_place requires a special SGLang build with that fix. --- cookbook/slime_disagg/README.md | 2 +- cookbook/slime_disagg/configs/moonlight_disagg.py | 4 +--- cookbook/slime_disagg/configs/qwen3_4b_delta_flash.py | 3 +-- cookbook/slime_disagg/sidecar.py | 5 ++--- src/stitch/sync.py | 6 ++---- 5 files changed, 7 insertions(+), 13 deletions(-) diff --git a/cookbook/slime_disagg/README.md b/cookbook/slime_disagg/README.md index 211d067..fbd6a4e 100644 --- a/cookbook/slime_disagg/README.md +++ b/cookbook/slime_disagg/README.md @@ -100,4 +100,4 @@ Bounding the per-run replay with periodic recovery anchors is left for later. Sidecars default to `quiesce` commit mode, which drains in-flight requests before applying a delta. The `in_place` mode (set `SIDECAR_COMMIT_MODE` in the config module) applies without draining and relies on version-namespaced -KV keys. It needs an SGLang build with the overlap-drain fix. +KV keys. diff --git a/cookbook/slime_disagg/configs/moonlight_disagg.py b/cookbook/slime_disagg/configs/moonlight_disagg.py index 0c747a0..37a2327 100644 --- a/cookbook/slime_disagg/configs/moonlight_disagg.py +++ b/cookbook/slime_disagg/configs/moonlight_disagg.py @@ -38,9 +38,7 @@ # M3: in_place commit applies weights without draining in-flight rollouts. Stale # KV is isolated per weight version by the sidecar's extra_key stamping (so old # requests keep decoding on their version's KV and it drains as they finish); -# min-version pins cross commits freely, only exact pins are quiesced. (The -# overlap-drain race the docs warned about is a tiny, un-reproduced window in the -# overlap-spec path, so we don't gate on the patched build.) +# min-version pins cross commits freely, only exact pins are quiesced. SIDECAR_COMMIT_MODE = "in_place" SIDECAR_DEBUG_REQUESTS = True diff --git a/cookbook/slime_disagg/configs/qwen3_4b_delta_flash.py b/cookbook/slime_disagg/configs/qwen3_4b_delta_flash.py index eae2097..82ab900 100644 --- a/cookbook/slime_disagg/configs/qwen3_4b_delta_flash.py +++ b/cookbook/slime_disagg/configs/qwen3_4b_delta_flash.py @@ -13,8 +13,7 @@ # the engine, applies, and resumes — in-flight requests keep decoding on stale # KV and cross-version isolation comes from extra_key stamping, so commits stop # blocking behind over-generation/eval stragglers and skip the full-tree flush. -# It relies on the targeted SGLang build's overlap-drain fix. "quiesce" is the -# safe fallback that drains in-flight requests before applying. +# "quiesce" is the safe fallback that drains in-flight requests before applying. SIDECAR_COMMIT_MODE = "in_place" # Log every versioned sidecar proxy request (start/end + injected rid) at INFO, diff --git a/cookbook/slime_disagg/sidecar.py b/cookbook/slime_disagg/sidecar.py index 4a17d3a..ba62fb1 100644 --- a/cookbook/slime_disagg/sidecar.py +++ b/cookbook/slime_disagg/sidecar.py @@ -81,9 +81,8 @@ def main() -> None: help=( "in_place (default): pause/apply/continue without flushing; " "in-flight requests keep decoding on stale KV and version isolation " - "comes from extra_key stamping. Relies on the engine's overlap-drain " - "fix. quiesce: wait out active requests and flush before applying " - "(safe on any build)." + "comes from extra_key stamping. quiesce: wait out active requests " + "and flush before applying." ), ) parser.add_argument( diff --git a/src/stitch/sync.py b/src/stitch/sync.py index 7b39e99..590d620 100644 --- a/src/stitch/sync.py +++ b/src/stitch/sync.py @@ -204,10 +204,8 @@ class WeightSyncManager(RolloutAdmissionGate): - ``in_place``: pause the engine in place, apply without flushing, and continue — in-flight requests resume decoding on their existing KV. Cross-version KV isolation comes from the sidecar stamping a composed, - version-namespaced ``extra_key`` onto every proxied request. Requires - an engine build with the overlap-drain fix; without it a forward in - flight at pause time can race the weight mutation. Only exact-version - requests are quiesced/gated. + version-namespaced ``extra_key`` onto every proxied request. Only + exact-version requests are quiesced/gated. """ def __init__( From 72a5e47e296f6297497df8d82b402f5ac449a1e1 Mon Sep 17 00:00:00 2001 From: Jason Mancuso <7891333+jvmncs@users.noreply.github.com> Date: Mon, 29 Jun 2026 19:57:18 -0700 Subject: [PATCH 03/18] test: collect cookbook tests + de-flake/decouple sync gate tests (#3) Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- pyproject.toml | 2 +- src/stitch/sync_test.py | 45 ++++++++++++++++++++++++++++++++--------- 2 files changed, 37 insertions(+), 10 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index 71ecd54..46af5e7 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -22,5 +22,5 @@ packages = ["src/stitch"] exclude = ["*_test.py"] [tool.pytest.ini_options] -testpaths = ["src"] +testpaths = ["src", "cookbook"] python_files = ["*_test.py"] diff --git a/src/stitch/sync_test.py b/src/stitch/sync_test.py index 8b65fea..c8ceeff 100644 --- a/src/stitch/sync_test.py +++ b/src/stitch/sync_test.py @@ -11,6 +11,14 @@ from stitch.sync import PolicyViolation, RolloutAdmissionGate, WeightSyncManager +async def _settle() -> None: + """Drain currently-ready callbacks so a task that is about to park on an + asyncio primitive reaches and parks at its await point. Deterministic + replacement for a wall-clock ``sleep`` when asserting a task is *blocked*.""" + for _ in range(10): + await asyncio.sleep(0) + + def _write_slime_version(base: Path, version: int, prev: int) -> None: vdir = base / f"weight_v{version:06d}" vdir.mkdir(parents=True) @@ -197,7 +205,7 @@ async def admit_exact_zero() -> str: return exc.error["error"]["type"] admit = asyncio.create_task(admit_exact_zero()) - await asyncio.sleep(0.05) + await _settle() # The admission must be gated, not validated against version 0. self.assertFalse(admit.done()) @@ -275,7 +283,7 @@ async def run() -> None: exact_ctx = manager.request_context(WeightVersionPolicy(exact_version=0)) await exact_ctx.__aenter__() sync = asyncio.create_task(manager.sync_to(1)) - await asyncio.sleep(0.05) + await _settle() self.assertFalse(engine.apply_started.is_set()) # Releasing the exact request lets the commit proceed. @@ -295,7 +303,7 @@ async def admit_exact() -> str: return exc.error["error"]["type"] gated = asyncio.create_task(admit_exact()) - await asyncio.sleep(0.05) + await _settle() self.assertFalse(gated.done()) engine.apply_gate.set() @@ -351,10 +359,19 @@ async def resume() -> None: resume=resume, ) - # Pause, apply, advance the served version, then resume — and the - # gate is reopened afterward. + # Pause, apply, advance the served version, then resume. self.assertEqual(events, ["pause", "apply", "applied", "resume"]) - self.assertFalse(gate._committing) + # The gate reopened afterward: a second commit drives cleanly through. + await gate.commit_version( + apply=apply, + on_applied=lambda: events.append("applied"), + pause=pause, + resume=resume, + ) + self.assertEqual( + events, + ["pause", "apply", "applied", "resume", "pause", "apply", "applied", "resume"], + ) asyncio.run(run()) @@ -381,10 +398,20 @@ async def resume() -> None: resume=resume, ) - # quiesce never pauses; a failed apply advances nothing and the gate - # is cleared so admissions resume. + # quiesce never pauses; a failed apply advances nothing. self.assertEqual(events, ["apply"]) - self.assertFalse(gate._committing) + + async def ok_apply() -> None: + events.append("apply") + + # The gate cleared despite the failure: a subsequent commit proceeds. + await gate.commit_version( + apply=ok_apply, + on_applied=lambda: events.append("applied"), + pause=pause, + resume=resume, + ) + self.assertEqual(events, ["apply", "apply", "applied"]) asyncio.run(run()) From 51961553e19905407cbd62c68973e6bff7a51a95 Mon Sep 17 00:00:00 2001 From: Jason Mancuso <7891333+jvmncs@users.noreply.github.com> Date: Mon, 29 Jun 2026 22:30:30 -0700 Subject: [PATCH 04/18] perf(sync): compose multi-version catch-up into one engine reload (#4) * perf(sync): compose multi-version catch-up into one engine reload _sync_once applied + reloaded the engine once per intermediate version on a multi-version catch-up (N pauses, N full update_weights_from_disk). The host-side apply (apply_deltas) already chain-replays the whole tail in a single call, so the per-version loop only multiplied the expensive engine reload. Verify chain contiguity up front, then drive a single commit that applies the target manifest once and advances current_version to latest -- one pause, one host-apply, one reload regardless of tail depth. Co-Authored-By: jason.mancuso@modal.com * docs(protocol): document apply_manifest chain-replay requirement The sync manager now composes a multi-version catch-up into a single apply_manifest call, so the EngineAdapter contract requires it to replay every intermediate version from the applied one to the target, not apply manifest as a single delta off the previous version. Make that explicit. --- src/stitch/protocol.py | 10 ++++++- src/stitch/sync.py | 59 +++++++++++++++++++++++++---------------- src/stitch/sync_test.py | 10 ++++--- 3 files changed, 52 insertions(+), 27 deletions(-) diff --git a/src/stitch/protocol.py b/src/stitch/protocol.py index 458a0eb..5644006 100644 --- a/src/stitch/protocol.py +++ b/src/stitch/protocol.py @@ -539,7 +539,15 @@ async def flush_cache(self) -> None: ... async def apply_manifest(self, manifest: VersionManifest, version_path: str) -> None: - """Apply one published weight version to the engine.""" + """Bring the engine to ``manifest.version``, replaying every intermediate + version from the currently applied one up to the target. + + The sync manager composes a multi-version catch-up into a single call + (it verifies chain contiguity, then invokes this once for the whole + tail), so an adapter MUST advance across all skipped versions, not just + apply ``manifest`` as a single delta off the previous one. The canonical + adapter satisfies this because slime's ``apply_deltas`` chain-replays + from the applied version to ``manifest.version``.""" ... async def pause_generation(self) -> None: diff --git a/src/stitch/sync.py b/src/stitch/sync.py index 590d620..35e0c49 100644 --- a/src/stitch/sync.py +++ b/src/stitch/sync.py @@ -353,34 +353,47 @@ async def _sync_once(self) -> bool: self.sync_state = SyncState.PREFETCHING self.last_sync_error = None + + # Verify the whole tail is a contiguous chain before touching the + # engine, so a broken chain fails before any pause/apply. apply_deltas + # re-verifies each step's base_version internally as it replays. + expected_base = self.current_version + target_manifest: VersionManifest | None = None for version in range(self.current_version + 1, latest + 1): manifest = self.board.read_manifest(self.current_run_id, version) - if manifest.base_version != self.current_version: + if manifest.base_version != expected_base: raise RuntimeError( f"cannot apply version {version}: manifest base " - f"{manifest.base_version} != current {self.current_version}" + f"{manifest.base_version} != expected {expected_base}" ) - version_path = str(self.board.version_dir(self.current_run_id, version)) - self.sync_state = SyncState.PREPARING - - async def apply(manifest: VersionManifest = manifest, version_path: str = version_path) -> None: - # quiesce flushes before applying; in_place skips the flush - # (the gate paused the engine and stale KV resumes as-is). - if self.commit_mode != "in_place": - await self.engine.flush_cache() - self.sync_state = SyncState.COMMITTING - await self.engine.apply_manifest(manifest, version_path) - - # on_applied bumps current_version under the gate (before - # continue_generation in in_place mode), so a request can never - # be admitted observing the stale version on mutated weights. - await self.commit_version( - apply=apply, - on_applied=lambda version=version: setattr(self, "current_version", version), - pause=self.engine.pause_generation, - resume=self.engine.continue_generation, - ) - self.sync_state = SyncState.PREFETCHING + expected_base = version + target_manifest = manifest + assert target_manifest is not None # latest > current_version, so the loop ran + target_path = str(self.board.version_dir(self.current_run_id, latest)) + + # Compose the tail and reload once: apply_manifest(target) replays + # every delta from the applied version up to `latest` host-side, then + # does a single engine reload — not one reload per intermediate version. + self.sync_state = SyncState.PREPARING + + async def apply(manifest: VersionManifest = target_manifest, version_path: str = target_path) -> None: + # quiesce flushes before applying; in_place skips the flush + # (the gate paused the engine and stale KV resumes as-is). + if self.commit_mode != "in_place": + await self.engine.flush_cache() + self.sync_state = SyncState.COMMITTING + await self.engine.apply_manifest(manifest, version_path) + + # on_applied bumps current_version under the gate (before + # continue_generation in in_place mode), so a request can never be + # admitted observing the stale version on mutated weights. + await self.commit_version( + apply=apply, + on_applied=lambda: setattr(self, "current_version", latest), + pause=self.engine.pause_generation, + resume=self.engine.continue_generation, + ) + self.sync_state = SyncState.PREFETCHING # Reached the board's latest for this run as captured at pass start; a # version published mid-pass is picked up by the next reconcile tick. diff --git a/src/stitch/sync_test.py b/src/stitch/sync_test.py index c8ceeff..1d654f9 100644 --- a/src/stitch/sync_test.py +++ b/src/stitch/sync_test.py @@ -65,7 +65,7 @@ async def continue_generation(self) -> None: class SyncManagerTest(unittest.TestCase): - def test_startup_sync_applies_transition_chain(self) -> None: + def test_startup_sync_composes_tail_into_one_apply(self) -> None: async def run() -> None: with tempfile.TemporaryDirectory() as tmp: board = FilesystemBulletinBoard(tmp) @@ -78,7 +78,10 @@ async def run() -> None: self.assertEqual(manager.current_version, 2) self.assertEqual(manager.sync_state, SyncState.IDLE) - self.assertEqual([v for v, _ in engine.applies], [1, 2]) + # A multi-version catch-up composes the tail host-side and reloads + # once: a single apply at the target version, not one per version. + self.assertEqual([v for v, _ in engine.applies], [2]) + self.assertEqual(engine.flushes, 1) asyncio.run(run()) @@ -96,7 +99,8 @@ async def run() -> None: await manager.startup_sync() self.assertEqual(manager.current_run_id, "run-a") self.assertEqual(manager.current_version, 2) - self.assertEqual([v for v, _ in engine.applies], [1, 2]) + # The two-version tail is composed into one apply at the target. + self.assertEqual([v for v, _ in engine.applies], [2]) # A new run forks at base with its version space restarting at 1. _write_slime_version(root / "run-b", 1, 0) From 78f88c5823bb1196813966ae8688bb77902ad5d4 Mon Sep 17 00:00:00 2001 From: Jason Mancuso <7891333+jvmncs@users.noreply.github.com> Date: Mon, 29 Jun 2026 23:38:39 -0700 Subject: [PATCH 05/18] chore: remove validation-only probes and probe-stapled scaffolding (#7) Delete the throwaway modal_probes/ directory and cookbook/miles_disagg/ modal_probes.py (validation/profiling scratch), and root out everything stapled into the examples that existed only to serve them: - modal_train.py: drop the env-gated R3_PROBE all-to-all logging injection (the functional R3 dispatch + GLM4Moe bias bakes are unchanged). - kimi_k2_6_nvfp4_disagg: drop ANCHOR_MODEL. Its only real consumers were the deleted probes; prepare_checkpoints merely snapshot_download'd it and never diffed anything, so the 'design anchor / diff against our base' lifecycle comments described validation that never ran. Remove the dead fetch and rewrite the docstring/README/last-layer-carveout comment to stand on their own. - slime moonlight_disagg: drop the 'verified by modal_probes/verify_mla_routing' doc reference. No change to any train/serve/sync code path. Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- cookbook/miles_disagg/README.md | 5 +- .../configs/kimi_k2_6_nvfp4_disagg.py | 15 +- cookbook/miles_disagg/modal_probes.py | 350 ------------------ cookbook/miles_disagg/modal_train.py | 28 +- .../slime_disagg/configs/moonlight_disagg.py | 2 +- modal_probes/__init__.py | 7 - modal_probes/confirm_router_bias_fp32.py | 59 --- modal_probes/inspect_dapo_dataset.py | 33 -- modal_probes/inspect_moonlight_weights.py | 125 ------- modal_probes/inspect_s3_transport.py | 83 ----- modal_probes/verify_mla_routing.py | 234 ------------ 11 files changed, 9 insertions(+), 932 deletions(-) delete mode 100644 cookbook/miles_disagg/modal_probes.py delete mode 100644 modal_probes/__init__.py delete mode 100644 modal_probes/confirm_router_bias_fp32.py delete mode 100644 modal_probes/inspect_dapo_dataset.py delete mode 100644 modal_probes/inspect_moonlight_weights.py delete mode 100644 modal_probes/inspect_s3_transport.py delete mode 100644 modal_probes/verify_mla_routing.py diff --git a/cookbook/miles_disagg/README.md b/cookbook/miles_disagg/README.md index 7567667..7b51275 100644 --- a/cookbook/miles_disagg/README.md +++ b/cookbook/miles_disagg/README.md @@ -42,8 +42,7 @@ prepared BF16 Hugging Face checkpoint directly. 2. **Served NVFP4 base** (`--hf-checkpoint`): produced from the masters with miles' own `tools/convert_hf_to_nvfp4.py`, so the served packing equals the trainer's export packing **by construction** (smallest deltas, no byte-exact - risk). For Kimi, `nvidia/Kimi-K2.6-NVFP4` is the design anchor / validation - target — diff it against the prepared base — not the literal served bytes. + risk). 3. **Megatron torch_dist** (`--load`/`--save`): trainer-internal rollout ckpts. The trainer reads the NVFP4 base for both the export quant config and the diff @@ -135,7 +134,7 @@ m run -m cookbook.miles_disagg.modal_train::smoke_flash_pool --weight-version 3 ``` The full `kimi_k2_6_nvfp4_disagg` recipe is a 32×8 B200 trainer footprint that -exceeds the probe budget — run the Moonlight de-risk first to validate the +exceeds the de-risk budget — run the Moonlight de-risk first to validate the QAT → NVFP4-export → XOR-delta → SGLang-reload loop, then scale. ### Fork dependencies diff --git a/cookbook/miles_disagg/configs/kimi_k2_6_nvfp4_disagg.py b/cookbook/miles_disagg/configs/kimi_k2_6_nvfp4_disagg.py index ee6b05b..59ac9cc 100644 --- a/cookbook/miles_disagg/configs/kimi_k2_6_nvfp4_disagg.py +++ b/cookbook/miles_disagg/configs/kimi_k2_6_nvfp4_disagg.py @@ -23,15 +23,13 @@ CHECKPOINT LIFECYCLE (three roles; see modal_train.prepare_checkpoints) ----------------------------------------------------------------------- Neither published K2.6 checkpoint is bf16: ``moonshotai/Kimi-K2.6`` ships as -compressed-tensors INT4, and ``nvidia/Kimi-K2.6-NVFP4`` is NVFP4. So: +compressed-tensors INT4. So: * BF16 masters (``ref_load``): dequantize the moonshotai INT4 checkpoint with ``tools/convert_kimi_int4_to_bf16.py``. * Served NVFP4 base (``hf_checkpoint``): produce with miles' ``tools/convert_hf_to_nvfp4.py`` from the bf16 masters, so the served packing == the trainer's export packing BY CONSTRUCTION (smallest deltas, no byte- - exact risk). ``nvidia/Kimi-K2.6-NVFP4`` is the DESIGN ANCHOR / validation - target: prepare_checkpoints fetches it and diffs it against our base to - confirm SGLang serves the layout — it is not the literal served bytes. + exact risk). * Megatron torch_dist (``load``/``save``): trainer-internal rollout ckpts. The trainer reads ``hf_checkpoint`` (NVFP4 base) for the export quant config AND @@ -72,9 +70,7 @@ # Source checkpoints the prepare step consumes (see modal_train.prepare_checkpoints): # SOURCE_MODEL -> dequantize INT4 -> bf16 masters -> convert -> served NVFP4 base -# ANCHOR_MODEL -> fetched for validation only (diff against our base) SOURCE_MODEL = "moonshotai/Kimi-K2.6" -ANCHOR_MODEL = "nvidia/Kimi-K2.6-NVFP4" MODEL_TAG = "kimi-k2-6-nvfp4" SIDECAR_COMMIT_MODE = "in_place" @@ -260,10 +256,9 @@ class _Miles(MilesConfig): # NOTE: the END carve-out is 0, NOT 1. A bf16 LAST MoE layer (layer 60) made the served # base's last-layer experts bf16 while SGLang's fused-MoE update_weights_from_disk RELOAD # loader allocates NVFP4 for every expert layer -> it can't load a bf16 [7168,2048] expert - # into the NVFP4-packed buffer (cold-load honors the carve-out, reload does not). The nvidia - # anchor quantizes all routed-expert layers incl. the last and reloads cleanly, so all-NVFP4 - # routed experts is a proven, reload-safe scheme for this model. (If QAT later needs the - # last-layer bf16 carve-out, the alternative is patching SGLang's fused-MoE reload loader.) + # into the NVFP4-packed buffer (cold-load honors the carve-out, reload does not). So the + # reload-safe scheme is to keep all routed experts NVFP4; if QAT later needs the last-layer + # bf16 carve-out, the alternative is patching SGLang's fused-MoE reload loader. num_layers_at_start_in_bf16 = 1 num_layers_at_end_in_bf16 = 0 diff --git a/cookbook/miles_disagg/modal_probes.py b/cookbook/miles_disagg/modal_probes.py deleted file mode 100644 index d842ba4..0000000 --- a/cookbook/miles_disagg/modal_probes.py +++ /dev/null @@ -1,350 +0,0 @@ -"""TEMP profiling probe: why is the SGLang weight reload so slow? - -Brings up ONE rollout engine identical to the miles_disagg Server cls (same -serving image, same SGLangEndpoint(model_path=NVFP4 base, tp, SGLANG_SERVER_ARGS), -same volumes + ephemeral disk), then times each phase of the weight-update path -and isolates disk-read throughput from SGLang's load processing: - - 1. raw read throughput of the base shards on the /prep Volume, - 2. SGLang cold startup load (reads the base from /prep), - 3. parallel materialize /prep base -> /local-checkpoint (ephemeral NVMe), - 4. raw read throughput of /local-checkpoint shards (ephemeral, cold), - 5. CLEAN reload of the base from /local-checkpoint via update_weights_from_disk - (no delta -> won't hit the merged-column crash; pure load-perf number), - 6. apply the existing delta (18ec126bbabc/weight_v000001) + reload (the real - scenario; may crash at the merged-column param — caught + timed). - -Run (its own app, no collision with the main deploy): - EXPERIMENT_CONFIG=kimi_k2_6_nvfp4_disagg MODAL_PROFILE=modal-labs \ - uv run --extra modal modal run --detach -m cookbook.miles_disagg.modal_probes::profile_reload -Read results: modal app logs -e jason-dev miles-kimi-probe -""" - -from __future__ import annotations - -import glob -import os -import subprocess -import time - -import modal - -# Reuse the EXACT image / volumes / constants / engine args the Server cls uses. -from cookbook.miles_disagg import modal_train as mt - -probe_app = modal.App("miles-kimi-probe") - -DELTA_RUN_ID = "18ec126bbabc" -DELTA_VERSION = 1 - - -def _dd_read_gbps(path: str, label: str) -> float: - """Cold-read a file with dd (drop page cache first); return GB/s.""" - try: - with open("/proc/sys/vm/drop_caches", "w") as f: - f.write("3") - except Exception as e: # noqa: BLE001 - print(f"[diskread] (could not drop page cache: {e})") - size = os.path.getsize(path) - t0 = time.perf_counter() - subprocess.run(["dd", f"if={path}", "of=/dev/null", "bs=16M"], check=True, - stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL) - dt = time.perf_counter() - t0 - gbps = size / 1e9 / dt - print(f"[diskread] {label}: {size/1e9:.1f} GB in {dt:.1f}s = {gbps:.2f} GB/s ({path})") - return gbps - - -@probe_app.function( - image=mt.server_image, - gpu=f"{mt.modal_cfg.gpu}:{mt.miles_cfg.rollout_num_gpus_per_engine}", - cloud=mt.modal_cfg.cloud, - region=mt.modal_cfg.region, - volumes={ - str(mt.HF_CACHE_PATH): mt.hf_cache_volume, - str(mt.PREP_PATH): mt.prep_volume, - mt.SGLANG_CACHE_PATH: mt.sglang_cache_volume, - mt.exp.DELTA_BULLETIN_ROOT: mt.delta_volume, - }, - ephemeral_disk=mt.modal_cfg.rollout_ephemeral_disk_mib, - timeout=60 * mt.MINUTES, -) -def profile_reload(delta_run_id: str = DELTA_RUN_ID, version: int = DELTA_VERSION) -> None: - import asyncio - - import httpx - from autoinference_utils.endpoint import SGLangEndpoint - - from miles.utils.disk_delta import apply_deltas - from stitch.engines.sglang import SGLangDiskDeltaAdapter - from stitch.protocol import VersionManifest - from cookbook.miles_disagg.sidecar import _parallel_init_local_checkpoint - - base = mt.MODEL_NAME # /prep/kimi-k2-6-nvfp4/nvfp4 - local = mt.LOCAL_CHECKPOINT_PATH # /local-checkpoint (ephemeral) - sglang_url = f"http://127.0.0.1:{mt.SGLANG_PORT}" - R: dict[str, object] = {} # durable results, printed once at the end - - def base_shards(): - return sorted(glob.glob(f"{base}/*.safetensors")) - - def alive() -> bool: - try: - return httpx.get(f"{sglang_url}/health", timeout=10.0).status_code == 200 - except Exception: # noqa: BLE001 - return False - - async def _reload(model_path: str) -> float: - t = time.perf_counter() - async with httpx.AsyncClient(timeout=None, trust_env=False) as c: - r = await c.post(f"{sglang_url}/update_weights_from_disk", - json={"model_path": model_path, "weight_version": "probe", "flush_cache": False}) - r.raise_for_status() - if r.json().get("success") is False: - raise RuntimeError(f"reject: {r.json()}") - return time.perf_counter() - t - - print(f"=== PROBE START — base={base} tp={mt.miles_cfg.rollout_num_gpus_per_engine} ===") - R["base_shards"] = len(base_shards()) - R["base_total_gb"] = round(sum(os.path.getsize(p) for p in base_shards()) / 1e9, 1) - - # 1. raw read throughput of a base shard (Modal Volume /prep) - if base_shards(): - R["dd_prep_gbps"] = round(_dd_read_gbps(base_shards()[0], "prep-volume base shard (cold)"), 2) - - # 2. SGLang cold startup load (reads base from /prep Volume) — the cold-load path - print("[phase] SGLang startup (cold load of base from /prep)...") - t0 = time.perf_counter() - endpoint = SGLangEndpoint( - model_path=base, worker_port=mt.SGLANG_PORT, - tp=mt.miles_cfg.rollout_num_gpus_per_engine, extra_server_args=mt.SGLANG_SERVER_ARGS, - health_timeout=mt.SERVER_STARTUP_TIMEOUT, health_poll_interval=10.0, - ) - endpoint.start() - R["startup_coldload_from_prep_s"] = round(time.perf_counter() - t0, 1) - print(f"[phase] startup cold load: {R['startup_coldload_from_prep_s']}s") - - # 3. KEY TEST 1 — reload /prep DIRECTLY via update_weights_from_disk (the reload path, - # complete base, no materialize). Crash here => SGLang reload-path bug (cold-load vs - # reload divergence). Success => the crash is materialize/local-checkpoint specific. - print("[phase] KEY TEST 1: reload /prep directly via update_weights_from_disk...") - if alive(): - try: - R["reload_prep_s"] = round(asyncio.run(_reload(base)), 1) - R["reload_prep"] = "OK" - print(f"[phase] reload /prep: OK {R['reload_prep_s']}s") - except Exception as e: # noqa: BLE001 - R["reload_prep"] = f"CRASH: {str(e)[:120]}" - print(f"[phase] reload /prep CRASH: {e}") - else: - R["reload_prep"] = "SKIP (engine not alive)" - - # 4. materialize + KEY TEST 2 — only if /prep reload didn't kill the engine - if alive(): - print("[phase] materialize base -> /local-checkpoint (32 workers)...") - t0 = time.perf_counter() - _parallel_init_local_checkpoint()(local, base) - R["materialize_s"] = round(time.perf_counter() - t0, 1) - print(f"[phase] materialize: {R['materialize_s']}s") - local_shards = sorted(glob.glob(f"{local}/*.safetensors")) - if local_shards: - R["dd_local_gbps"] = round(_dd_read_gbps(local_shards[0], "ephemeral /local-checkpoint shard (cold)"), 2) - print("[phase] KEY TEST 2: reload /local-checkpoint via update_weights_from_disk...") - if alive(): - try: - R["reload_local_s"] = round(asyncio.run(_reload(local)), 1) - R["reload_local"] = "OK" - except Exception as e: # noqa: BLE001 - R["reload_local"] = f"CRASH: {str(e)[:120]}" - print(f"[phase] reload /local CRASH: {e}") - else: - R["materialize_s"] = R["reload_local"] = "SKIP (engine died on /prep reload)" - - print("=== PROBE SUMMARY ===") - for k, v in R.items(): - print(f" {k}: {v}") - try: - endpoint.stop() - except Exception: # noqa: BLE001 - pass - - -@probe_app.function( - image=mt.server_image, - volumes={ - str(mt.PREP_PATH): mt.prep_volume, - mt.exp.DELTA_BULLETIN_ROOT: mt.delta_volume, - }, - timeout=15 * mt.MINUTES, -) -def rca_checksum(run_id: str = "ece3aa5c17c0") -> None: - """CPU RCA: for a few of the checksum-mismatched tensors, reconstruct base XOR delta - against the CURRENT /prep base and compare to the delta's stored `want` checksum. - match => base is consistent with the delta => the SIDECAR's base_local was different - (volume skew / stale materialize) -> the pool diffed against a different base. - mismatch => the CURRENT /prep base != the base the trainer snapshotted (trainer-side skew). - """ - import glob, io, json, struct - import numpy as np - import zstandard - from miles.utils.disk_delta import checksum, make_tensor_reader - - base = mt.MODEL_NAME - vdir = f"{mt.exp.DELTA_BULLETIN_ROOT}/{run_id}/weight_v000001" - reader = make_tensor_reader(base) - meta = json.load(open(f"{vdir}/model.safetensors.index.json"))["metadata"] - algo = meta["checksum_format"] - print(f"=== RCA checksum: base={base} delta={vdir} algo={algo} enc={meta['delta_encoding']} ===") - - probe = ["language_model.lm_head.weight", "language_model.model.embed_tokens.weight", - "language_model.model.layers.3.mlp.experts.0.gate_proj.weight", - "language_model.model.layers.3.self_attn.o_proj.weight"] - checked = 0 - for df in sorted(glob.glob(f"{vdir}/*.safetensors")): - blob = open(df, "rb").read() - (hlen,) = struct.unpack(" None: - """CPU: compare NVFP4 tensor SHAPES (weight + scales) of a routed expert and a - merged-column param between OUR /prep base and the nvidia anchor — to pinpoint - the layout divergence the reload's fused-MoE loader rejects.""" - import json - from huggingface_hub import snapshot_download - from safetensors import safe_open - - ours = mt.MODEL_NAME - anchor = snapshot_download(mt.exp.ANCHOR_MODEL, local_files_only=True) - - def shapes_for(root: str, names: list[str]) -> dict: - idx = json.load(open(f"{root}/model.safetensors.index.json")).get("weight_map", {}) - out = {} - for n in names: - f = idx.get(n) - if not f: - out[n] = "ABSENT"; continue - with safe_open(f"{root}/{f}", framework="numpy") as sf: - sl = sf.get_slice(n) - out[n] = f"{sl.get_shape()} {sl.get_dtype()}" - return out - - probe_names = [] - for base_mod in ["language_model.model.layers.5.mlp.experts.0.gate_proj", - "language_model.model.layers.5.mlp.experts.0.down_proj"]: - for suf in [".weight", ".weight_scale", ".weight_scale_2", ".input_scale"]: - probe_names.append(base_mod + suf) - o = shapes_for(ours, probe_names) - a = shapes_for(anchor, probe_names) - print("=== ROUTED-EXPERT NVFP4 SHAPE DIFF (ours vs anchor) ===") - for n in probe_names: - mark = "" if o[n] == a[n] else " <<< DIFFERS" - print(f" {n.split('experts.0.')[-1]:28s} ours={o[n]:22s} anchor={a[n]:22s}{mark}") - - -@probe_app.function( - image=mt.server_image, - gpu=f"{mt.modal_cfg.gpu}:{mt.miles_cfg.rollout_num_gpus_per_engine}", - cloud=mt.modal_cfg.cloud, - region=mt.modal_cfg.region, - volumes={ - str(mt.HF_CACHE_PATH): mt.hf_cache_volume, - str(mt.PREP_PATH): mt.prep_volume, - mt.SGLANG_CACHE_PATH: mt.sglang_cache_volume, - }, - ephemeral_disk=mt.modal_cfg.rollout_ephemeral_disk_mib, - timeout=60 * mt.MINUTES, -) -def profile_anchor_reload() -> None: - """DECISIVE test: is the crash SGLang's reload path (any checkpoint) or OUR - prepare_checkpoints NVFP4 output specifically? - - Cold-load the KNOWN-GOOD reference nvidia/Kimi-K2.6-NVFP4 (the design anchor), - then reload IT via update_weights_from_disk: - anchor reload CRASH -> SGLang reload-path bug (prepare_checkpoints innocent) - anchor reload OK -> our /prep convert_hf_to_nvfp4 output is the culprit - (probe v1/v2 showed our /prep reload crashes); the - bf16 carve-out (layers 0 & 60, no scales) is the suspect. - """ - import asyncio - - import httpx - from huggingface_hub import snapshot_download - from autoinference_utils.endpoint import SGLangEndpoint - - anchor = snapshot_download(mt.exp.ANCHOR_MODEL, local_files_only=True) - sglang_url = f"http://127.0.0.1:{mt.SGLANG_PORT}" - R: dict[str, object] = {"anchor_path": anchor} - print(f"=== ANCHOR PROBE — {mt.exp.ANCHOR_MODEL} @ {anchor} ===") - - def alive() -> bool: - try: - return httpx.get(f"{sglang_url}/health", timeout=10.0).status_code == 200 - except Exception: # noqa: BLE001 - return False - - async def _reload(model_path: str) -> float: - t = time.perf_counter() - async with httpx.AsyncClient(timeout=None, trust_env=False) as c: - r = await c.post(f"{sglang_url}/update_weights_from_disk", - json={"model_path": model_path, "weight_version": "probe", "flush_cache": False}) - r.raise_for_status() - if r.json().get("success") is False: - raise RuntimeError(f"reject: {r.json()}") - return time.perf_counter() - t - - print("[phase] cold-load anchor...") - t0 = time.perf_counter() - endpoint = SGLangEndpoint( - model_path=anchor, worker_port=mt.SGLANG_PORT, - tp=mt.miles_cfg.rollout_num_gpus_per_engine, extra_server_args=mt.SGLANG_SERVER_ARGS, - health_timeout=mt.SERVER_STARTUP_TIMEOUT, health_poll_interval=10.0, - ) - endpoint.start() - R["anchor_coldload_s"] = round(time.perf_counter() - t0, 1) - print(f"[phase] anchor cold-load: {R['anchor_coldload_s']}s") - - print("[phase] KEY: reload anchor via update_weights_from_disk...") - if alive(): - try: - R["anchor_reload_s"] = round(asyncio.run(_reload(anchor)), 1) - R["anchor_reload"] = "OK" - except Exception as e: # noqa: BLE001 - R["anchor_reload"] = f"CRASH: {str(e)[:120]}" - print(f"[phase] anchor reload CRASH: {e}") - else: - R["anchor_reload"] = "SKIP (engine not alive after cold-load)" - - print("=== ANCHOR PROBE SUMMARY ===") - for k, v in R.items(): - print(f" {k}: {v}") - try: - endpoint.stop() - except Exception: # noqa: BLE001 - pass diff --git a/cookbook/miles_disagg/modal_train.py b/cookbook/miles_disagg/modal_train.py index 8811748..bbed7d4 100644 --- a/cookbook/miles_disagg/modal_train.py +++ b/cookbook/miles_disagg/modal_train.py @@ -185,28 +185,6 @@ if getattr(exp, "PATCH_GLM4MOE_EXPERT_BIAS_FP32", False): image = image.run_commands(f"python3 -c {shlex.quote(_GLM4MOE_EXPERT_BIAS_BAKE_PY)}") -# R3 debug probe (env-gated, removable): log the actual EP all-to-all split sizes -# right before token_dispatch's all_to_all, to diagnose the K2.6 routing-replay -# "Split sizes doesn't match total dim 0 size" crash. Enable with R3_PROBE=1 at -# deploy. Prepends a logging line before `global_input_tokens = all_to_all(`. -if os.environ.get("R3_PROBE"): - _R3_PROBE_TARGET = "global_input_tokens = all_to_all(" - _R3_PROBE_LOG = ( - 'import logging as _r3; _r3.getLogger("R3PROBE").warning(' - 'f"[R3PROBE] perm={permutated_local_input_tokens.size(0)} ' - "in_sum={int(self.input_splits.sum())} out_sum={int(self.output_splits.sum())} " - 'nout={int(self.num_out_tokens)} ep={self.ep_size} tp={self.tp_size}"); ' - ) - _R3_PROBE_PY = ( - "import pathlib;" - f"p=pathlib.Path({_R3_DISPATCH_TARGET!r});" - "s=p.read_text();" - f"s=s.replace({_R3_PROBE_TARGET!r}, {_R3_PROBE_LOG + _R3_PROBE_TARGET!r}, 1);" - "p.write_text(s);" - "print('[R3 probe] injected all_to_all size logging')" - ) - image = image.run_commands(f"python3 -c {shlex.quote(_R3_PROBE_PY)}") - # Local source mounted at container start (no rebuild on code edits). Modal puts # /root on PYTHONPATH, so both packages import from subprocesses (the sidecar, Ray # workers). MUST come after any .run_commands() above (Modal forbids run_commands @@ -528,7 +506,6 @@ def prepare_checkpoints() -> None: downloaded checkpoint IS the masters. 2. served NVFP4 base: tools/convert_hf_to_nvfp4.py over the masters, so the served packing == the trainer's export packing by construction. - 3. (optional) fetch ANCHOR_MODEL (nvidia NVFP4) for validation only. """ if getattr(exp, "DISABLE_HF_XET", False): os.environ["HF_HUB_DISABLE_XET"] = "1" @@ -580,7 +557,7 @@ def _build_bf16(out: str) -> None: # (for Kimi VLMs it's nested under text_config). Strip it: the masters are # bf16, and convert_hf_to_nvfp4 inherits text_config verbatim — leaving it # would make the served NVFP4 base claim compressed-tensors INT4 there and - # diverge from the clean nvidia anchor (transformers/SGLang could then pick + # diverge from the clean served base (transformers/SGLang could then pick # the wrong quant loader for the language model). _strip_stale_quant_config(os.path.join(out, "config.json")) @@ -606,9 +583,6 @@ def _build_bf16(out: str) -> None: check=True, ), ) - # 3. validation anchor (diff against nvfp4_dir to confirm SGLang-compatibility) - if anchor := getattr(exp, "ANCHOR_MODEL", None): - snapshot_download(anchor) prep_volume.commit() print(f"Prepared masters={bf16_dir} served_base={nvfp4_dir}") diff --git a/cookbook/slime_disagg/configs/moonlight_disagg.py b/cookbook/slime_disagg/configs/moonlight_disagg.py index 37a2327..fa863ab 100644 --- a/cookbook/slime_disagg/configs/moonlight_disagg.py +++ b/cookbook/slime_disagg/configs/moonlight_disagg.py @@ -20,7 +20,7 @@ 1. The training image's slime fork ref (modal_train.SLIME_REPO_REF) must include scripts/models/moonlight.sh, the deepseekv3 megatron_to_hf export, and the routing-replay code. The base image's SGLang (used by the rollout pool) is - independent and already verified by modal_probes/verify_mla_routing. + independent. 2. Bridge-mode HF load is UNVERIFIED for Moonlight/DeepSeek-V3 (it is proven for Qwen here). If load fails, the fallback is the proven torch_dist path (tools/convert_hf_to_torch_dist.py + ref_load pointed at the converted dir). diff --git a/modal_probes/__init__.py b/modal_probes/__init__.py deleted file mode 100644 index c71798d..0000000 --- a/modal_probes/__init__.py +++ /dev/null @@ -1,7 +0,0 @@ -"""One-off Modal probes and smoke checks. - -Throwaway verification harnesses that exercise infrastructure or engine behavior -(e.g. "does the pinned SGLang build emit routed experts for an MLA MoE?"). These -are intentionally kept out of ``cookbook/``, which holds pedagogical end-to-end -examples. -""" diff --git a/modal_probes/confirm_router_bias_fp32.py b/modal_probes/confirm_router_bias_fp32.py deleted file mode 100644 index 450192c..0000000 --- a/modal_probes/confirm_router_bias_fp32.py +++ /dev/null @@ -1,59 +0,0 @@ -"""Probe: confirm Megatron keeps the MoE router's expert_bias buffer in fp32, and -show WHY (so we know whether exporting it as bf16 for the rollout is correct). - -The M1 disk-delta export crashed XORing a 256-byte tensor against a 128-byte base -for `e_score_correction_bias`. The base is uniformly bf16 (per inspect_moonlight_weights), -and `64 experts x 4 bytes (fp32) = 256` vs `64 x 2 (bf16) = 128`. moonlight.sh sets -`--moe-router-dtype fp32`, so the hypothesis is: Megatron keeps the router `expert_bias` -in fp32 and slime's deepseekv3 converter exported it without casting to the bf16 -checkpoint dtype. - -This is CPU-only source inspection of megatron-core's router -- it (1) confirms the -fp32 path without a GPU build, and (2) dumps `_maintain_float32_expert_bias` + the -expert_bias buffer registration + its update, so we can judge whether the fp32 is a -training-side accumulation concern (=> bf16 export to the rollout is the correct HF -representation) or something the inference path also depends on. - - m run -m modal_probes.confirm_router_bias_fp32::confirm -""" - -from __future__ import annotations - -import modal - -SLIME_IMAGE_TAG = "slimerl/slime:nightly-dev-20260527a" - -image = modal.Image.from_registry(SLIME_IMAGE_TAG).entrypoint([]) -app = modal.App("confirm-router-bias-fp32") - - -@app.function(image=image, timeout=10 * 60) -def confirm() -> None: - import importlib - import inspect - - router_mod = importlib.import_module("megatron.core.transformer.moe.router") - print(f"router source: {inspect.getsourcefile(router_mod)}\n") - - # 1) Dump _maintain_float32_expert_bias from whatever class defines it. - printed = set() - for cls_name in dir(router_mod): - cls = getattr(router_mod, cls_name) - if not isinstance(cls, type): - continue - method = getattr(cls, "_maintain_float32_expert_bias", None) - if method is None or method in printed: - continue - printed.add(method) - print(f"===== {cls_name}._maintain_float32_expert_bias =====") - try: - print(inspect.getsource(method)) - except (OSError, TypeError) as exc: - print(f"(could not getsource: {exc})") - - # 2) Show every line in the module that touches expert_bias creation / update / dtype. - print("===== router.py lines: expert_bias / register_buffer / bias update / dtype =====") - for i, line in enumerate(inspect.getsource(router_mod).splitlines(), 1): - s = line.strip() - if ("expert_bias" in s or "register_buffer" in s) and ("#" != s[:1] if s else False): - print(f"{i:5}: {s}") diff --git a/modal_probes/inspect_dapo_dataset.py b/modal_probes/inspect_dapo_dataset.py deleted file mode 100644 index a38569b..0000000 --- a/modal_probes/inspect_dapo_dataset.py +++ /dev/null @@ -1,33 +0,0 @@ -"""Probe: inspect DAPO-Math-17k columns so the Moonlight config's prepare_data -writes a jsonl with the fields slime's --input-key prompt / --label-key label + -rm_type=math expect. The slime recipes consume a *preprocessed* jsonl; the raw HF -dataset may put the answer under e.g. reward_model.ground_truth, not `label`. - - m run -m modal_probes.inspect_dapo_dataset::inspect -""" - -from __future__ import annotations - -import modal - -SLIME_IMAGE_TAG = "slimerl/slime:nightly-dev-20260527a" - -image = modal.Image.from_registry(SLIME_IMAGE_TAG).entrypoint([]).pip_install("datasets") -app = modal.App("inspect-dapo-dataset") - - -@app.function(image=image, timeout=15 * 60, secrets=[modal.Secret.from_name("huggingface-secret")]) -def inspect() -> None: - import json - - from datasets import load_dataset - - ds = load_dataset("BytedTsinghua-SIA/DAPO-Math-17k", split="train") - print(f"num_rows={ds.num_rows}") - print(f"column_names={ds.column_names}") - print(f"features={ds.features}") - ex = ds[0] - print("\n=== example[0] (values truncated to 600 chars) ===") - for k, v in ex.items(): - s = json.dumps(v, default=str) - print(f" {k}: {s[:600]}{' …' if len(s) > 600 else ''}") diff --git a/modal_probes/inspect_moonlight_weights.py b/modal_probes/inspect_moonlight_weights.py deleted file mode 100644 index d0e4ec2..0000000 --- a/modal_probes/inspect_moonlight_weights.py +++ /dev/null @@ -1,125 +0,0 @@ -"""Probe: dump Moonlight's base HF tensor dtypes/shapes to debug the M1 -disk-delta export shape mismatch. - -The M1 trainer crashed in slime's disk-delta export: - - update_weight_from_disk_delta._encode_delta -> diff_and_compress: - diff = new ^ old - ValueError: operands could not be broadcast together with shapes (256,) (128,) - -The disk-delta XORs the Megatron->HF *exported* tensor (``new``) against the base -HF checkpoint tensor (``old``) AS FLAT uint8 BYTES (it does -``tensor.view(torch.uint8).reshape(-1)`` first). So this is 256 bytes vs 128 -bytes -- the export produced exactly DOUBLE the bytes of the base for one key. -Two candidate causes: - - (a) dtype mismatch -- the base checkpoint stores that tensor in a 1-byte dtype - (fp8) while the export is 2-byte bf16 (same element count, 2x bytes); or - (b) element-count doubling in slime's convert_deepseekv3_to_hf for Moonlight's - MLA variant (no q-LoRA, --qk-layernorm) -- that converter is written for - DeepSeek-V3.2 (it has the DSA "indexer" wq_b/wk/k_norm branches) and one of - its hard-coded 128/64 reshapes does not match Moonlight's layout. - -This reads the safetensors headers only (dtype + shape + byte span per tensor), -so it is CPU-only and instant -- no weights are materialized. It surfaces the -base side: the dtype histogram tells (a) apart from (b), and the per-tensor byte -sizes pinpoint which base tensor is 128 bytes (the ``old`` operand). - - alias m="uv run --extra modal modal" - m run -m modal_probes.inspect_moonlight_weights::inspect -""" - -from __future__ import annotations - -import json -import os -import struct -from collections import Counter -from pathlib import Path - -import modal - -SLIME_IMAGE_TAG = "slimerl/slime:nightly-dev-20260527a" -HF_CACHE_PATH = "/root/.cache/huggingface" -MODEL_NAME = os.environ.get("VERIFY_MODEL", "moonshotai/Moonlight-16B-A3B-Instruct") - -# safetensors dtype -> bytes/element, for spotting fp8 (1B) vs bf16 (2B) bases. -_DTYPE_BYTES = { - "F64": 8, "I64": 8, "F32": 4, "I32": 4, "F16": 2, "BF16": 2, - "I16": 2, "F8_E4M3": 1, "F8_E5M2": 1, "I8": 1, "U8": 1, "BOOL": 1, -} - -image = ( - modal.Image.from_registry(SLIME_IMAGE_TAG) - .entrypoint([]) - # The base image bakes an HF cache here; remove it so the volume can mount. - .run_commands(f"rm -rf {HF_CACHE_PATH}") - .env({"HF_HUB_ENABLE_HF_TRANSFER": "1"}) -) -app = modal.App("inspect-moonlight-weights") -hf_cache_volume = modal.Volume.from_name("huggingface-cache", create_if_missing=True) - - -def _read_safetensors_header(path: Path) -> dict: - """Parse just the JSON header of a .safetensors file (no tensor data).""" - with open(path, "rb") as f: - (header_len,) = struct.unpack(" None: - from huggingface_hub import snapshot_download - - model_dir = Path(snapshot_download(MODEL_NAME)) # cached from the M1 download - config = json.loads((model_dir / "config.json").read_text()) - mla_keys = [ - "q_lora_rank", "kv_lora_rank", "qk_nope_head_dim", "qk_rope_head_dim", - "v_head_dim", "qk_head_dim", "num_attention_heads", "num_key_value_heads", - "hidden_size", "head_dim", "torch_dtype", "quantization_config", - ] - print(f"=== {MODEL_NAME} config (MLA-relevant) ===") - for k in mla_keys: - if k in config: - print(f" {k} = {config[k]}") - - tensors: dict[str, dict] = {} - for shard in sorted(model_dir.glob("*.safetensors")): - tensors.update(_read_safetensors_header(shard)) - print(f"\n=== {len(tensors)} tensors across {len(list(model_dir.glob('*.safetensors')))} shard(s) ===") - - dtypes = Counter(meta["dtype"] for meta in tensors.values()) - print(f"dtype histogram: {dict(dtypes)}") - if len(dtypes) > 1: - print(" -> MIXED dtypes: a base tensor in a 1-byte dtype vs a bf16 export") - print(" would XOR-mismatch by exactly 2x. This is hypothesis (a).") - - def nbytes(meta: dict) -> int: - off = meta.get("data_offsets") - if off: - return off[1] - off[0] - n = 1 - for d in meta["shape"]: - n *= d - return n * _DTYPE_BYTES.get(meta["dtype"], 0) - - # The `old` operand was 128 bytes. List every base tensor at that byte size. - print("\n=== base tensors that are exactly 128 bytes (the `old` operand) ===") - for name, meta in sorted(tensors.items()): - if nbytes(meta) == 128: - print(f" {name}: shape={meta['shape']} dtype={meta['dtype']} nbytes=128") - - # Full attention/norm layout for layer 0 (the converter branches act here). - print("\n=== layer-0 attention + norm tensors ===") - for name, meta in sorted(tensors.items()): - if (".layers.0." in name or "layers.0." in name) and ( - "self_attn" in name or "norm" in name.lower() - ): - print(f" {name}: shape={meta['shape']} dtype={meta['dtype']} nbytes={nbytes(meta)}") diff --git a/modal_probes/inspect_s3_transport.py b/modal_probes/inspect_s3_transport.py deleted file mode 100644 index 6452284..0000000 --- a/modal_probes/inspect_s3_transport.py +++ /dev/null @@ -1,83 +0,0 @@ -"""Probe: inspect (and, with --confirm, wipe) the standalone provider's S3 -transport state. - -The run-id design partitions each run's delta chain under ``/weight_v{N}/`` -and a single self-identifying ``latest`` pointer (``/weight_vNNNNNN``). -Pre-run-id deployments left a flat ``weight_v*`` chain + a bare ``latest`` -(``"NNNNNN"``); that layout is incompatible and must be wiped once before the -new code runs (a cold-starting replica would otherwise read the stale bare -pointer and choke on the orphaned flat chain until the first publish). - - m run -m modal_probes.inspect_s3_transport::inspect - m run -m modal_probes.inspect_s3_transport::wipe --confirm -""" - -from __future__ import annotations - -import json -import shutil -from pathlib import Path - -import modal - -BUCKET = "modal-stitch-s3-transport" -PREFIX = "standalone-rollouts/moonlight" -OIDC_ROLE = "arn:aws:iam::459781239556:role/modal-buckets/stitch-s3-transport-role" - - -def _mount(read_only: bool) -> modal.CloudBucketMount: - return modal.CloudBucketMount( - bucket_name=BUCKET, - key_prefix=f"{PREFIX}/", - oidc_auth_role_arn=OIDC_ROLE, - read_only=read_only, - ) - - -image = modal.Image.debian_slim() -app = modal.App("inspect-s3-transport") - - -def _report(root: Path) -> None: - latest = root / "latest" - print(f"latest pointer = {latest.read_text().strip() if latest.exists() else 'MISSING'}") - print(f"top-level entries: {sorted(p.name for p in root.iterdir())}") - # Run partitions: any top-level dir holding its own weight_v* chain. - for d in sorted(p for p in root.iterdir() if p.is_dir() and not p.name.startswith("weight_v")): - wv = sorted(x.name for x in d.glob("weight_v*")) - if wv: - print(f"[run dir] {d.name}: versions={wv}") - flat = sorted(root.glob("weight_v*")) - if flat: - print(f"!! legacy flat (pre-run-id) version dirs present: {[d.name for d in flat]}") - for d in flat: - idx = d / "model.safetensors.index.json" - meta = json.loads(idx.read_text()).get("metadata", {}) if idx.exists() else {} - print(f" {d.name}.index.metadata = {meta}") - - -@app.function(image=image, volumes={"/mnt/t": _mount(read_only=True)}, timeout=10 * 60, region="us") -def inspect() -> None: - _report(Path("/mnt/t")) - - -@app.function(image=image, volumes={"/mnt/t": _mount(read_only=False)}, timeout=10 * 60, region="us") -def wipe(confirm: bool = False) -> None: - """Delete everything under the prefix (pointer + all run/flat version dirs). - - The state is migration-incompatible (and, for the diagnosed app, corrupt), so - this is a clean reset: replicas then cold-start at (None, 0) on base, and the - first training run writes a fresh ``/`` partition + pointer. - """ - root = Path("/mnt/t") - if not confirm: - print("Refusing to wipe without --confirm. Current state:") - _report(root) - return - for entry in sorted(root.iterdir()): - if entry.is_dir(): - shutil.rmtree(entry, ignore_errors=True) - else: - entry.unlink(missing_ok=True) - print("Wiped. Post-wipe state:") - _report(root) diff --git a/modal_probes/verify_mla_routing.py b/modal_probes/verify_mla_routing.py deleted file mode 100644 index 5fbefb6..0000000 --- a/modal_probes/verify_mla_routing.py +++ /dev/null @@ -1,234 +0,0 @@ -"""M0 verify gate for the async-MoE disaggregated demo. - -The cheap question this answers before any training config is written: does the -SGLang build pinned in this repo's slime image serve a DeepSeek-V3-architecture -(multi-head latent attention) MoE *and* return per-token routed-expert ids? - -Why this is the one gate worth running first. Kimi K2.6 is a DeepSeek-V3-arch -model -- MLA (``q_a/q_b`` + ``kv_a_proj_with_mqa``/``kv_b`` projections) plus -DeepSeek-MoE (sigmoid router, shared experts, grouped top-k). The disk-delta -weight sync and routing replay (``--use-rollout-routing-replay``) are already -validated on the GLM-4.x MoEs in bf16 and fp8 -- but GLM-4.x is GQA, not MLA. -MLA is the single axis those runs never exercised. Moonlight-16B-A3B is -Moonshot's small same-family proxy (MLA + DeepSeek-MoE, *real* HF weights), so -if SGLang here serves it and emits ``meta_info['routed_experts']``, the -routing-replay rollout path works for the Kimi proxy too. - -Scope: this checks the *serving* axis only. It deliberately needs real HF -weights, which is why it serves Moonlight rather than the deepseek-v3-5layer -arch wrapper (that wrapper is a Megatron arch config with no HF checkpoint to -serve). The *trainer-side* MLA round-trip -- bridge-mode HF load into Megatron -and disk-delta export producing HF keys that match the served safetensors -- is -exercised by the first M1 bring-up run, where deepseek-v3-5layer is the fast -iteration vehicle. - - alias m="uv run --extra modal modal" - m run -m modal_probes.verify_mla_routing::download_model - m run -m modal_probes.verify_mla_routing::verify - -PASS prints the SGLang version and a routed-experts tensor that reshapes to -[tokens, num_layers, moe_router_topk] with expert ids inside [0, num_experts). -""" - -from __future__ import annotations - -import base64 -import json -import os -import signal -import subprocess -import time -import urllib.error -import urllib.request -from pathlib import Path - -import modal - -# The pinned slime image carries the SGLang build under test; serving from it is -# what makes this check meaningful. We add nothing to the sglang wheel. -SLIME_IMAGE_TAG = "slimerl/slime:nightly-dev-20260527a" -HF_CACHE_PATH = "/root/.cache/huggingface" - -# Moonlight-16B-A3B: small DeepSeek-V3-arch (MLA) MoE with real weights. Override -# VERIFY_MODEL to point at another DeepSeek/Kimi-family HF checkpoint. -MODEL_NAME = os.environ.get("VERIFY_MODEL", "moonshotai/Moonlight-16B-A3B-Instruct") - -MINUTES = 60 -SGLANG_PORT = 8001 -STARTUP_TIMEOUT = 25 * MINUTES - -image = ( - modal.Image.from_registry(SLIME_IMAGE_TAG) - .entrypoint([]) - # The base image bakes in an HF cache; remove it so it cannot shadow the - # cache volume mounted at the same path. - .run_commands(f"rm -rf {HF_CACHE_PATH}") - .env({"HF_HUB_ENABLE_HF_TRANSFER": "1", "HF_XET_HIGH_PERFORMANCE": "1"}) -) - -app = modal.App("verify-mla-routing") -hf_cache_volume = modal.Volume.from_name("huggingface-cache", create_if_missing=True) - - -@app.function( - image=image, - volumes={HF_CACHE_PATH: hf_cache_volume}, - timeout=60 * MINUTES, - secrets=[modal.Secret.from_name("huggingface-secret")], -) -def download_model() -> None: - from huggingface_hub import snapshot_download - - snapshot_download(repo_id=MODEL_NAME) - hf_cache_volume.commit() - print(f"Downloaded {MODEL_NAME} into the huggingface-cache volume.") - - -@app.function( - image=image, - gpu="H200:1", - volumes={HF_CACHE_PATH: hf_cache_volume}, - timeout=30 * MINUTES, -) -def verify() -> None: - """Serve the proxy model on the pinned SGLang and assert routed-expert output.""" - import numpy as np - from huggingface_hub import snapshot_download - - model_dir = snapshot_download(MODEL_NAME, local_files_only=True) - config = json.loads((Path(model_dir) / "config.json").read_text()) - num_layers = int(config["num_hidden_layers"]) - topk = int(config["num_experts_per_tok"]) - num_experts = int(config.get("n_routed_experts") or config.get("num_experts") or 0) - print( - f"Serving {MODEL_NAME}: num_hidden_layers={num_layers}, " - f"num_experts_per_tok={topk}, n_routed_experts={num_experts or '?'}" - ) - - proc = subprocess.Popen( - [ - "python3", - "-m", - "sglang.launch_server", - "--model-path", - model_dir, - "--port", - str(SGLANG_PORT), - "--trust-remote-code", - "--enable-return-routed-experts", # the capability under test - "--mem-fraction-static", - "0.85", - "--cuda-graph-max-bs", - "8", - ], - start_new_session=True, - ) - try: - info = _wait_ready(proc, STARTUP_TIMEOUT) - print( - f"SGLang up: version={info.get('version')!r} " - f"attention_backend={info.get('attention_backend')!r}" - ) - - payload = { - "text": "Give me a one-sentence fun fact about the Moon.", - "sampling_params": {"temperature": 0.0, "max_new_tokens": 24}, - "return_logprob": True, - "return_routed_experts": True, # what the rollout path sends - } - out = _post_json( - f"http://127.0.0.1:{SGLANG_PORT}/generate", payload, timeout=180 - ) - meta = out["meta_info"] - - if "routed_experts" not in meta: - raise SystemExit( - "FAIL: meta_info has no 'routed_experts'. The pinned SGLang build does not emit " - "routed experts for this MLA model -- routing replay for the Kimi proxy is blocked " - "at the engine. (Confirm --enable-return-routed-experts is supported by this build.)" - ) - - # Mirror the slime rollout decode exactly: base64 -> int32, reshape to - # [tokens, num_layers, moe_router_topk] (see slime sglang_rollout.generate). - flat = np.frombuffer( - base64.b64decode(meta["routed_experts"].encode("ascii")), dtype=np.int32 - ) - per_layer_topk = num_layers * topk - if per_layer_topk == 0 or flat.size % per_layer_topk != 0: - raise SystemExit( - f"FAIL: routed_experts size {flat.size} is not divisible by " - f"num_layers*topk ({num_layers}*{topk}); the rollout reshape contract is broken." - ) - tokens = flat.size // per_layer_topk - matrix = flat.reshape(tokens, num_layers, topk) - - completion = meta.get("completion_tokens") - prompt = meta.get("prompt_tokens") - print( - f"routed_experts: {flat.size} int32 -> [tokens={tokens}, layers={num_layers}, topk={topk}] " - f"(prompt_tokens={prompt}, completion_tokens={completion})" - ) - lo, hi = int(matrix.min()), int(matrix.max()) - print( - f" expert id range [{lo}, {hi}]; sample (token 0, layer 0): {matrix[0, 0].tolist()}" - ) - if num_experts and not (0 <= lo and hi < num_experts): - raise SystemExit( - f"FAIL: expert ids [{lo}, {hi}] fall outside [0, {num_experts}); " - "the routed-experts payload is malformed for this arch." - ) - print( - "PASS: pinned SGLang serves the MLA MoE and emits well-formed routed_experts." - ) - finally: - _terminate(proc) - - -def _wait_ready(proc: subprocess.Popen, timeout: int) -> dict: - """Poll until SGLang answers server_info; return it (carries the version).""" - deadline = time.time() + timeout - last_error = "" - while time.time() < deadline: - if proc.poll() is not None: - raise SystemExit( - f"FAIL: sglang.launch_server exited early (code={proc.returncode})." - ) - for endpoint in ("/get_server_info", "/server_info"): - try: - return _get_json( - f"http://127.0.0.1:{SGLANG_PORT}{endpoint}", timeout=10 - ) - except (urllib.error.URLError, TimeoutError, OSError, ValueError) as exc: - last_error = f"{endpoint}: {type(exc).__name__}: {exc}" - time.sleep(5) - raise SystemExit( - f"FAIL: SGLang did not become ready in {timeout}s; last error: {last_error}" - ) - - -def _get_json(url: str, *, timeout: float) -> dict: - with urllib.request.urlopen(url, timeout=timeout) as resp: - return json.load(resp) - - -def _post_json(url: str, payload: dict, *, timeout: float) -> dict: - req = urllib.request.Request( - url, - data=json.dumps(payload).encode(), - headers={"Content-Type": "application/json"}, - ) - with urllib.request.urlopen(req, timeout=timeout) as resp: - return json.load(resp) - - -def _terminate(proc: subprocess.Popen) -> None: - if proc.poll() is not None: - return - try: - os.killpg(proc.pid, signal.SIGTERM) - proc.wait(timeout=20) - except Exception: # noqa: BLE001 - try: - os.killpg(proc.pid, signal.SIGKILL) - except Exception: # noqa: BLE001 - pass From 21674a5739464c7a55cd0e1ad00229bbac2d9a6a Mon Sep 17 00:00:00 2001 From: Jason Mancuso <7891333+jvmncs@users.noreply.github.com> Date: Tue, 30 Jun 2026 01:25:51 -0700 Subject: [PATCH 06/18] Explicit pool claim/reset: one trainer call owns one pool epoch (#5) * Explicit pool claim/reset: one trainer call owns one pool epoch * Align reconcile tests with #4 composed-reload semantics After rebasing onto main (which includes #4, composing a multi-version catch-up into a single engine reload), two reconcile tests still asserted the pre-#4 per-version apply accounting: - local_disagg/harness_test.py (added by this PR): a cold latecomer catching up v1..v2 now records a single composed apply [2], not [1, 2]. - standalone_rollouts/provider_test.py (pre-existing; #4 changed the behavior but did not update this cookbook test): same one-apply tail. The served version / run_id invariants are unchanged; only the per-apply accounting moved. No production code change. --- cookbook/bulletin_hooks.py | 71 +++++-- cookbook/local_disagg/README.md | 50 +++++ cookbook/local_disagg/__init__.py | 9 + cookbook/local_disagg/harness.py | 150 ++++++++++++++ cookbook/local_disagg/harness_test.py | 183 ++++++++++++++++++ cookbook/miles_disagg/hooks.py | 16 +- cookbook/miles_disagg/modal_train.py | 11 ++ cookbook/slime_disagg/hooks.py | 16 +- cookbook/slime_disagg/modal_train.py | 11 ++ cookbook/standalone_rollouts/frontdoor.py | 27 +-- .../standalone_rollouts/frontdoor_test.py | 8 + cookbook/standalone_rollouts/modal_serve.py | 9 +- cookbook/standalone_rollouts/provider_test.py | 6 +- src/stitch/bulletin.py | 40 ++++ src/stitch/bulletin_test.py | 41 ++++ src/stitch/protocol.py | 72 +++++++ src/stitch/protocol_test.py | 38 ++++ 17 files changed, 721 insertions(+), 37 deletions(-) create mode 100644 cookbook/local_disagg/README.md create mode 100644 cookbook/local_disagg/__init__.py create mode 100644 cookbook/local_disagg/harness.py create mode 100644 cookbook/local_disagg/harness_test.py diff --git a/cookbook/bulletin_hooks.py b/cookbook/bulletin_hooks.py index d2c0e74..1e742f6 100644 --- a/cookbook/bulletin_hooks.py +++ b/cookbook/bulletin_hooks.py @@ -23,7 +23,7 @@ from typing import Any from stitch.bulletin import FilesystemBulletinBoard -from stitch.protocol import parse_weight_identity +from stitch.protocol import BASE_VERSION, PointerRewind, parse_weight_identity from stitch.providers.modal import commit_volume, discover_flash_targets, volume_reloader, wake_targets @@ -56,23 +56,56 @@ def commit_and_wake( version = parse_weight_identity(Path(version_dir).name) rank = distributed_rank() - # Rank 0 owns the `latest` pointer and writes it before committing, so the + # Rank 0 owns the `latest` pointer and advances it before committing, so the # committed bulletin is self-consistent for the poll/startup path. The pointer # lives at the transport root (the Volume mount) and is self-identifying — # `/weight_vN` — while the trainer wrote the version dir under the run - # partition (update_weight_disk_dir = /), so a new run is a - # forward move of the pointer, never a colliding rewind. + # partition (update_weight_disk_dir = /). `advance` enforces the + # monotonic-within-run rule (the new run's first publish forks at base); a + # same-run rewind (e.g. a republish) is dropped rather than serving stale + # weights — never silently overwritten. if version is not None and rank in (None, 0): - FilesystemBulletinBoard(_transport_root(args), layout="slime").write_latest( - _run_id(args), version - ) + board = FilesystemBulletinBoard(_transport_root(args), layout="slime") + try: + board.advance(_run_id(args), version) + except PointerRewind: + logger.warning( + "publish of version %s would rewind latest; dropping (run %r)", + version, + _run_id(args), + exc_info=True, + ) + return commit_volume(_volume_name(args)) if version is None or rank not in (None, 0): return - # Waking warm containers is a best-effort latency optimization: a transient - # Modal control-plane error must not kill the training step — `latest` is - # already committed and sidecars self-sync on their next poll. + _best_effort_wake(args, version, app_name_env=app_name_env, cls_name_env=cls_name_env) + + +def claim_pool(args: Any, *, app_name_env: str, cls_name_env: str) -> None: + """Trainer launch hook (rank 0): claim the rollout pool for this run. + + Write the empty pointer ``/weight_v000000``, commit the Volume, and + wake the pool — so every replica (cold or already-warm on a finished run) + resets to base *before* the first delta publishes, instead of inferring the + reset from the first publish's run mismatch. ``run_id`` must be fresh per + launch (the run's epoch/fence token); claiming a run already at the pointer + raises :class:`PointerRewind`, which fails the launch loudly rather than + leaving the pool pinned to a dead incarnation's high-water mark. + """ + if distributed_rank() not in (None, 0): + return + board = FilesystemBulletinBoard(_transport_root(args), layout="slime") + board.claim(_run_id(args)) + commit_volume(_volume_name(args)) + _best_effort_wake(args, BASE_VERSION, app_name_env=app_name_env, cls_name_env=cls_name_env) + + +def _best_effort_wake(args: Any, version: int, *, app_name_env: str, cls_name_env: str) -> None: + """Nudge warm Flash containers to reconcile now. Best-effort: a transient + Modal control-plane error must not kill the training step — `latest` is + already committed and sidecars self-sync on their next poll/startup.""" try: app_name = getattr(args, "rollout_modal_flash_app_name", None) or os.environ[app_name_env] cls_name = getattr(args, "rollout_modal_flash_server_cls_name", None) or os.getenv( @@ -195,9 +228,21 @@ def _transport_root(args: Any) -> str: def _run_id(args: Any) -> str: - """The run partition (chain identity). Passed explicitly via custom_config, - falling back to the basename of the per-run write dir.""" - return str(getattr(args, "run_id", None) or Path(bulletin_root(args)).name) + """The run partition (chain identity), passed explicitly via custom_config. + + Required — never derived from the write-dir basename. A fresh run_id per + launch is the epoch/fence token that makes restart safe (a restart is just a + new epoch that claims and resets the pool); deriving it from a per-run dir + that a crash-restart could reuse is exactly the reuse hazard this design + removes, so a missing run_id is a launch misconfiguration, not a fallback. + """ + run_id = getattr(args, "run_id", None) + if not run_id: + raise ValueError( + "run_id is required (pass it via custom_config_path); the bulletin " + "hooks no longer derive it from the write-dir basename" + ) + return str(run_id) def _gate_board(args: Any) -> FilesystemBulletinBoard: diff --git a/cookbook/local_disagg/README.md b/cookbook/local_disagg/README.md new file mode 100644 index 0000000..9948802 --- /dev/null +++ b/cookbook/local_disagg/README.md @@ -0,0 +1,50 @@ +# local_disagg — minimal pool-claim harness + +The smallest possible disaggregated-rollout setup: **no Modal, no slime/miles, +no GPUs.** A local-filesystem bulletin board, an in-memory rollout pool, and a +trainer-as-writer, wired with the *same* claim / advance / reconcile primitives +the production cookbooks use. It exists so the pool-ownership invariants can be +exercised and iterated on in milliseconds. + +## The model + +One **trainer** call owns one **run**, which owns one **pool epoch**: + +``` +trainer.claim() -> board.claim(run_id) # write /weight_v000000 (empty) +trainer.publish() -> board.advance(run_id, N) # write /weight_vN, monotonic +replica.reconcile() -> WeightSyncManager.sync_to() # converge to latest (reset on run switch) +``` + +- The **bulletin board** is the single source of truth: a self-identifying + `latest` pointer `/weight_v{N}`. +- The **trainer** is the single writer. `claim` resets the pool to base for a + fresh run; `publish` advances monotonically within the run. Both go through + the board's guarded writers (`stitch.protocol.decide_pointer_move`), so a + reused `run_id` or a non-monotonic publish raises `PointerRewind` instead of + serving stale weights. +- Each **replica** is a pure reconciler — it reads `latest` and converges + (replay the chain forward, or reset-to-base then replay on a run switch). + +`run_id` is a per-launch epoch/fence token (default: a fresh `uuid4`). A restart +is just a new epoch that claims and resets the pool — there is no special restart +path, and a crash-restart can never reuse a `run_id` to resurrect a dead pointer. + +## Invariants (see `harness_test.py`) + +- A claim resets every replica to base (v0) under the new run's id. +- Within a run the pool converges to each published version, in order. +- A new run forks at base: the pool resets even from a higher prior version. +- Re-claiming a run already at the pointer (a reused `run_id`) is rejected as a + rewind; the correct restart mints a fresh `run_id`. +- A non-monotonic publish within a run is rejected. +- A late-joining (cold / scaled-up) replica reconciles to the *current* run. + +## Run it + +```bash +uv run python -m pytest cookbook/local_disagg/harness_test.py -q +``` + +To iterate by hand, build a board + trainer + pool and step through claim / +publish / reconcile; see `harness.py` for the (tiny) surface. diff --git a/cookbook/local_disagg/__init__.py b/cookbook/local_disagg/__init__.py new file mode 100644 index 0000000..f587cbb --- /dev/null +++ b/cookbook/local_disagg/__init__.py @@ -0,0 +1,9 @@ +"""Minimal, dependency-free disaggregated-rollout harness. + +No Modal, no slime/miles, no GPUs: a filesystem bulletin board, an in-memory +rollout pool, and a trainer-as-writer, wired with the *same* claim/advance/ +reconcile primitives the real cookbooks use. It exists to exercise — and pin +down with tests — the pool-claim invariants quickly and locally. + +See :mod:`cookbook.local_disagg.harness`. +""" diff --git a/cookbook/local_disagg/harness.py b/cookbook/local_disagg/harness.py new file mode 100644 index 0000000..cc814f6 --- /dev/null +++ b/cookbook/local_disagg/harness.py @@ -0,0 +1,150 @@ +"""A minimal in-memory disaggregated-rollout harness. + +Three pieces mirror the real cookbooks with none of the infrastructure: + +- :class:`MemoryEngine` — an in-memory rollout server. It "serves" exactly one + weight version; ``apply_manifest`` steps it forward one delta and ``reset`` + drops it back to base (v0). Stands in for the SGLang adapter. +- :class:`LocalReplica` — one ``WeightSyncManager`` driving one ``MemoryEngine`` + against a shared bulletin board. The pool is a list of these; each is a pure + reconciler that converges to ``latest`` on ``reconcile()``. +- :class:`LocalTrainer` — the single writer for one run. It ``claim``s the pool + (resets every replica to base) at launch, then ``publish``es monotonic delta + versions. One trainer ↔ one run ↔ one pool epoch. + +The trainer writes the *same* slime-layout chain the real disk-delta publisher +writes (``//weight_v{N}/model.safetensors.index.json``), so the +replicas reconcile through the production ``WeightSyncManager`` path — only the +engine and the transport (a local dir instead of a Modal Volume / S3) are fakes. +""" + +from __future__ import annotations + +import json +import uuid +from pathlib import Path + +from stitch.bulletin import FilesystemBulletinBoard +from stitch.protocol import BASE_VERSION, PointerMove, VersionManifest + + +class MemoryEngine: + """In-memory rollout engine: tracks the single weight version it serves. + + base is v0; each ``apply_manifest`` advances exactly one delta and ``reset`` + re-materializes base. ``applied`` / ``resets`` are recorded so tests can + assert the reconcile path (replay the chain, reset on a run switch). + """ + + backend = "memory" + + def __init__(self) -> None: + self.version = 0 + self.applied: list[int] = [] + self.resets = 0 + + async def prepare(self) -> None: + pass + + async def flush_cache(self) -> None: + pass + + async def apply_manifest(self, manifest: VersionManifest, version_path: str) -> None: + self.version = manifest.version + self.applied.append(manifest.version) + + async def reset(self) -> None: + self.version = BASE_VERSION + self.resets += 1 + + async def pause_generation(self) -> None: + pass + + async def continue_generation(self) -> None: + pass + + +class LocalReplica: + """One rollout-pool replica: a ``WeightSyncManager`` + its ``MemoryEngine``. + + Constructed lazily (the sync manager needs a running event loop), so the + pool can be sized before any reconcile. ``reconcile`` converges the replica + to the board's current ``(run_id, version)`` — following a run switch + (reset → replay) exactly like the production sidecar. + """ + + def __init__(self, board: FilesystemBulletinBoard) -> None: + from stitch.sync import WeightSyncManager + + self.engine = MemoryEngine() + self.manager = WeightSyncManager(board=board, engine=self.engine, commit_mode="in_place") + + async def reconcile(self) -> None: + await self.manager.sync_to() + + @property + def served_version(self) -> int: + return self.engine.version + + @property + def served_run_id(self) -> str | None: + return self.manager.current_run_id + + +class LocalTrainer: + """The single writer for one run: claim the pool, then publish deltas. + + ``run_id`` defaults to a fresh per-launch token (the epoch/fence id that + makes restart a clean new-run reset, never a colliding rewind). ``claim`` + writes the empty pointer; ``publish`` writes the next version dir and + advances the pointer. Both go through the board's guarded writers, so a + reused run_id or a non-monotonic publish raises rather than serving stale + weights. + """ + + def __init__(self, board: FilesystemBulletinBoard, run_id: str | None = None) -> None: + self.board = board + self.run_id = run_id or uuid.uuid4().hex[:12] + self.version = BASE_VERSION + + def claim(self) -> PointerMove: + move = self.board.claim(self.run_id) + self.version = BASE_VERSION + return move + + def publish(self) -> PointerMove: + """Materialize the next delta version and advance ``latest`` to it.""" + nxt = self.version + 1 + _write_version_dir(self.board.version_dir(self.run_id, nxt), version=nxt, base=self.version) + move = self.board.advance(self.run_id, nxt) + self.version = nxt + return move + + +def make_pool(board: FilesystemBulletinBoard, size: int) -> list[LocalReplica]: + return [LocalReplica(board) for _ in range(size)] + + +async def reconcile_pool(pool: list[LocalReplica]) -> None: + for replica in pool: + await replica.reconcile() + + +def open_board(root: str | Path) -> FilesystemBulletinBoard: + """The slime-layout board both trainer and pool share (run-scoped chains).""" + return FilesystemBulletinBoard(str(root), layout="slime") + + +def _write_version_dir(version_dir: Path, *, version: int, base: int) -> None: + """Write the slime disk-delta version dir the real publisher would write: + a canonical HF index whose metadata carries the version lineage.""" + version_dir.mkdir(parents=True, exist_ok=True) + (version_dir / "model.safetensors.index.json").write_text( + json.dumps( + { + "metadata": {"version": f"{version:06d}", "base_version": f"{base:06d}"}, + "weight_map": {"w": "model-00001-of-00001.safetensors"}, + } + ), + encoding="utf-8", + ) diff --git a/cookbook/local_disagg/harness_test.py b/cookbook/local_disagg/harness_test.py new file mode 100644 index 0000000..9617d97 --- /dev/null +++ b/cookbook/local_disagg/harness_test.py @@ -0,0 +1,183 @@ +"""Pool-claim invariants, exercised on the minimal in-memory harness. + +These are the guarantees the explicit claim/advance design is meant to provide, +phrased against the same primitives the real cookbooks use (board.claim / +board.advance + a WeightSyncManager pool). Each test reads as one invariant. +""" + +from __future__ import annotations + +import asyncio +import tempfile +import unittest + +from cookbook.local_disagg.harness import ( + LocalReplica, + LocalTrainer, + make_pool, + open_board, + reconcile_pool, +) +from stitch.protocol import PointerRewind + + +class PoolClaimInvariantTest(unittest.TestCase): + def test_claim_resets_pool_to_base(self) -> None: + """A claim drops every replica to base (v0) under the new run's id — + the explicit 'empty' starting state, written before any delta.""" + + async def run() -> None: + with tempfile.TemporaryDirectory() as tmp: + board = open_board(tmp) + trainer = LocalTrainer(board) + pool = make_pool(board, size=3) + + move = trainer.claim() + self.assertTrue(move.reset) + self.assertEqual(move.version, 0) + await reconcile_pool(pool) + + self.assertTrue(all(r.served_version == 0 for r in pool)) + self.assertTrue(all(r.served_run_id == trainer.run_id for r in pool)) + + asyncio.run(run()) + + def test_publish_advances_pool_monotonically(self) -> None: + """Within a run the pool converges to each published version in order.""" + + async def run() -> None: + with tempfile.TemporaryDirectory() as tmp: + board = open_board(tmp) + trainer = LocalTrainer(board) + pool = make_pool(board, size=2) + trainer.claim() + + for expected in (1, 2, 3): + self.assertEqual(trainer.publish().version, expected) + await reconcile_pool(pool) + self.assertTrue(all(r.served_version == expected for r in pool)) + + # The chain was replayed delta-by-delta, never skipped. + self.assertEqual(pool[0].engine.applied, [1, 2, 3]) + + asyncio.run(run()) + + def test_new_run_resets_pool_even_from_higher_version(self) -> None: + """A fresh run forks at base: the pool re-materializes base and replays + the new chain, even though the finished run reached a higher version.""" + + async def run() -> None: + with tempfile.TemporaryDirectory() as tmp: + board = open_board(tmp) + pool = make_pool(board, size=2) + + old = LocalTrainer(board) + old.claim() + for _ in range(5): + old.publish() + await reconcile_pool(pool) + self.assertTrue(all(r.served_version == 5 for r in pool)) + + new = LocalTrainer(board) + self.assertNotEqual(new.run_id, old.run_id) + new.claim() + new.publish() # the new run is only at v1 + await reconcile_pool(pool) + + self.assertTrue(all(r.served_version == 1 for r in pool)) + self.assertTrue(all(r.served_run_id == new.run_id for r in pool)) + # The drop from v5 to the new run's base went through an engine reset. + self.assertTrue(all(r.engine.resets >= 1 for r in pool)) + + asyncio.run(run()) + + def test_restart_with_reused_run_id_is_rejected_as_rewind(self) -> None: + """The restart hazard, made impossible: re-claiming a run already at the + pointer is a rewind, not a silent stale-pointer reuse. A restart must + mint a fresh run_id (a new epoch), which claims cleanly.""" + + async def run() -> None: + with tempfile.TemporaryDirectory() as tmp: + board = open_board(tmp) + crashed = LocalTrainer(board, run_id="run-fixed") + crashed.claim() + for _ in range(3): + crashed.publish() + + # Crash-restart that reused the same run_id would rewind latest + # (v3 -> v0) onto stale weights — rejected. + restarted_same = LocalTrainer(board, run_id="run-fixed") + with self.assertRaises(PointerRewind): + restarted_same.claim() + + # The correct restart is a new epoch: a fresh run_id claims clean. + pool = make_pool(board, size=2) + restarted_fresh = LocalTrainer(board) + restarted_fresh.claim() + await reconcile_pool(pool) + self.assertTrue(all(r.served_version == 0 for r in pool)) + self.assertTrue(all(r.served_run_id == restarted_fresh.run_id for r in pool)) + + asyncio.run(run()) + + def test_non_monotonic_publish_within_run_is_rejected(self) -> None: + """Within a run the pointer only moves forward; re-advancing to an + already-published version is a rewind.""" + + async def run() -> None: + with tempfile.TemporaryDirectory() as tmp: + board = open_board(tmp) + trainer = LocalTrainer(board) + trainer.claim() + trainer.publish() + trainer.publish() # at v2 + + with self.assertRaises(PointerRewind): + board.advance(trainer.run_id, 2) + with self.assertRaises(PointerRewind): + board.advance(trainer.run_id, 1) + + asyncio.run(run()) + + def test_late_joining_replica_reconciles_to_current_run(self) -> None: + """A replica that joins after the claim (a scaled-up / cold container) + converges to the *current* run's chain, never a finished run's pointer.""" + + async def run() -> None: + with tempfile.TemporaryDirectory() as tmp: + board = open_board(tmp) + + old = LocalTrainer(board) + old.claim() + for _ in range(4): + old.publish() + + new = LocalTrainer(board) + new.claim() + new.publish() + new.publish() # current run at v2 + + # Replica created only now, with no prior state. + latecomer = LocalReplica(board) + await latecomer.reconcile() + + self.assertEqual(latecomer.served_version, 2) + self.assertEqual(latecomer.served_run_id, new.run_id) + # The cold catch-up over the v1..v2 tail composes into a single + # engine apply at the target version (not one apply per + # intermediate version) — see WeightSyncManager._sync_once. + self.assertEqual(latecomer.engine.applied, [2]) + + asyncio.run(run()) + + def test_claim_requires_run_id(self) -> None: + """A claim must name its run (the per-launch epoch token); an empty run + id is a launch misconfiguration, not a usable claim.""" + with tempfile.TemporaryDirectory() as tmp: + board = open_board(tmp) + with self.assertRaises(ValueError): + board.claim("") + + +if __name__ == "__main__": + unittest.main() diff --git a/cookbook/miles_disagg/hooks.py b/cookbook/miles_disagg/hooks.py index 67b8f3e..9f75160 100644 --- a/cookbook/miles_disagg/hooks.py +++ b/cookbook/miles_disagg/hooks.py @@ -9,21 +9,31 @@ from typing import Any from cookbook.bulletin_hooks import ( + claim_pool as _claim_pool, commit_and_wake as _commit_and_wake, gated_rollout_request_hook, ) +_APP_NAME_ENV = "MILES_DELTA_APP_NAME" +_CLS_NAME_ENV = "MILES_DELTA_SERVER_CLS_NAME" + + +def claim_pool(args: Any) -> None: + """Claim the rollout pool for this run at launch (resets every replica to base).""" + _claim_pool(args, app_name_env=_APP_NAME_ENV, cls_name_env=_CLS_NAME_ENV) + + def commit_and_wake(args: Any, version_dir: str, rollout_engines: list[Any]) -> None: """miles ``custom_delta_pre_push_path`` hook (publish-only, bulletin board).""" _commit_and_wake( args, version_dir, rollout_engines, - app_name_env="MILES_DELTA_APP_NAME", - cls_name_env="MILES_DELTA_SERVER_CLS_NAME", + app_name_env=_APP_NAME_ENV, + cls_name_env=_CLS_NAME_ENV, ) # Re-export for the trainer's custom_rollout_request_hook_path. -__all__ = ["commit_and_wake", "gated_rollout_request_hook"] +__all__ = ["claim_pool", "commit_and_wake", "gated_rollout_request_hook"] diff --git a/cookbook/miles_disagg/modal_train.py b/cookbook/miles_disagg/modal_train.py index bbed7d4..c01aa7c 100644 --- a/cookbook/miles_disagg/modal_train.py +++ b/cookbook/miles_disagg/modal_train.py @@ -21,6 +21,7 @@ import tempfile import uuid from pathlib import Path +from types import SimpleNamespace import modal import modal.experimental @@ -465,6 +466,16 @@ def train(self, experiment: str, payload: dict) -> None: helpers.prepare_miles_config(cfg, tempfile.mkdtemp()) cmd = helpers.build_train_cmd(cfg, MILES_ROOT) + # Claim the pool for this run *before* miles starts publishing: write the + # empty pointer (/weight_v000000) and wake the pool so every + # replica resets to base now, closing the window where a replica could + # reconcile to a finished run's stale high-water version. + from cookbook.miles_disagg import hooks + + hooks.claim_pool( + SimpleNamespace(update_weight_disk_dir=cfg.update_weight_disk_dir, **cfg.custom_config_path) + ) + print(f"Training {experiment}: nodes={N_TRAIN_NODES}, rollout_endpoint={cfg.rollout_endpoint_url}") print(f"Command: {cmd}") # Tee the full training output to a committed Volume file so failures are diff --git a/cookbook/slime_disagg/hooks.py b/cookbook/slime_disagg/hooks.py index 93ea67f..d4b1814 100644 --- a/cookbook/slime_disagg/hooks.py +++ b/cookbook/slime_disagg/hooks.py @@ -9,21 +9,31 @@ from typing import Any from cookbook.bulletin_hooks import ( + claim_pool as _claim_pool, commit_and_wake as _commit_and_wake, gated_rollout_request_hook, ) +_APP_NAME_ENV = "SLIME_DELTA_APP_NAME" +_CLS_NAME_ENV = "SLIME_DELTA_SERVER_CLS_NAME" + + +def claim_pool(args: Any) -> None: + """Claim the rollout pool for this run at launch (resets every replica to base).""" + _claim_pool(args, app_name_env=_APP_NAME_ENV, cls_name_env=_CLS_NAME_ENV) + + def commit_and_wake(args: Any, version_dir: str, rollout_engines: list[Any]) -> None: """SLIME ``custom_delta_pre_push_path`` hook (publish-only, bulletin board).""" _commit_and_wake( args, version_dir, rollout_engines, - app_name_env="SLIME_DELTA_APP_NAME", - cls_name_env="SLIME_DELTA_SERVER_CLS_NAME", + app_name_env=_APP_NAME_ENV, + cls_name_env=_CLS_NAME_ENV, ) # Re-export for the trainer's custom_rollout_request_hook_path. -__all__ = ["commit_and_wake", "gated_rollout_request_hook"] +__all__ = ["claim_pool", "commit_and_wake", "gated_rollout_request_hook"] diff --git a/cookbook/slime_disagg/modal_train.py b/cookbook/slime_disagg/modal_train.py index b4db9b5..6ecd046 100644 --- a/cookbook/slime_disagg/modal_train.py +++ b/cookbook/slime_disagg/modal_train.py @@ -17,6 +17,7 @@ import tempfile import uuid from pathlib import Path +from types import SimpleNamespace import modal import modal.experimental @@ -369,6 +370,16 @@ def train(self, experiment: str, payload: dict) -> None: helpers.prepare_slime_config(cfg, tempfile.mkdtemp()) cmd = helpers.build_train_cmd(cfg, SLIME_ROOT) + # Claim the pool for this run *before* slime starts publishing: write the + # empty pointer (/weight_v000000) and wake the pool so every + # replica resets to base now, closing the window where a replica could + # reconcile to a finished run's stale high-water version. + from cookbook.slime_disagg import hooks + + hooks.claim_pool( + SimpleNamespace(update_weight_disk_dir=cfg.update_weight_disk_dir, **cfg.custom_config_path) + ) + print( f"Training {experiment}: nodes={N_TRAIN_NODES}, rollout_endpoint={cfg.rollout_endpoint_url}" ) diff --git a/cookbook/standalone_rollouts/frontdoor.py b/cookbook/standalone_rollouts/frontdoor.py index 299180d..09ac7e8 100644 --- a/cookbook/standalone_rollouts/frontdoor.py +++ b/cookbook/standalone_rollouts/frontdoor.py @@ -27,8 +27,10 @@ from typing import Any from stitch.protocol import ( + PointerRewind, RolloutPoolState, RolloutReplicaState, + decide_pointer_move, format_snapshot_identity, parse_weight_identity, ) @@ -54,6 +56,11 @@ def advance_latest_decision( Returns ``{"run_id": str|None, "version": int, "reset": bool}`` to accept, or ``{"error": {...}}`` to reject (``InvalidIdentity`` / ``WeightRewindRejected``). + The accept/reset/rewind call is :func:`stitch.protocol.decide_pointer_move`, + the same rule the bulletin-board publish path advances through; this wrapper + only parses the wire identity and shapes the error dict. An explicit claim is + just a signal at ``weight_v000000`` with a fresh ``run_id`` (a cross-run move, + so ``reset=True``). """ version = parse_weight_identity(identity) if version is None: @@ -63,22 +70,20 @@ def advance_latest_decision( "message": f"identity {identity!r} is not weight_v", } } - if request_run_id != current_run_id: - # New run: its version space restarts, so accepting it is not a rewind. - return {"run_id": request_run_id, "version": int(version), "reset": True} - if version <= current_version: + try: + move = decide_pointer_move( + current_run_id, current_version, run_id=request_run_id, version=version + ) + except PointerRewind as rewind: return { "error": { "type": "WeightRewindRejected", - "message": ( - f"latest is at version {current_version} (run {current_run_id!r}); " - f"refusing to rewind to {version}" - ), - "current_version": int(current_version), - "requested_version": int(version), + "message": str(rewind), + "current_version": rewind.current_version, + "requested_version": rewind.requested_version, } } - return {"run_id": request_run_id, "version": int(version), "reset": False} + return {"run_id": move.run_id, "version": move.version, "reset": move.reset} def pool_state_from_server_infos(infos: list[dict[str, Any]]) -> RolloutPoolState: diff --git a/cookbook/standalone_rollouts/frontdoor_test.py b/cookbook/standalone_rollouts/frontdoor_test.py index 439199f..7d7523f 100644 --- a/cookbook/standalone_rollouts/frontdoor_test.py +++ b/cookbook/standalone_rollouts/frontdoor_test.py @@ -37,6 +37,14 @@ def test_runless_layout_stays_monotonic(self) -> None: rewind = advance_latest_decision(None, 5, "weight_v000005", None) self.assertEqual(rewind["error"]["type"], "WeightRewindRejected") + def test_claim_is_a_base_version_signal_for_a_fresh_run(self) -> None: + # An explicit claim is just weight_v000000 with a fresh run id: a + # cross-run move, so it resets the pool to base before any delta. + self.assertEqual( + advance_latest_decision("run-a", 5, "weight_v000000", "run-b"), + {"run_id": "run-b", "version": 0, "reset": True}, + ) + def test_rejects_unparseable_identity(self) -> None: self.assertEqual( advance_latest_decision(None, 0, "base", None)["error"]["type"], "InvalidIdentity" diff --git a/cookbook/standalone_rollouts/modal_serve.py b/cookbook/standalone_rollouts/modal_serve.py index 9adec36..f06a963 100644 --- a/cookbook/standalone_rollouts/modal_serve.py +++ b/cookbook/standalone_rollouts/modal_serve.py @@ -24,7 +24,6 @@ from cookbook.slime_disagg import helpers from cookbook.standalone_rollouts import frontdoor as frontdoor_mod from stitch.bulletin import FilesystemBulletinBoard -from stitch.protocol import format_snapshot_identity from stitch.providers.modal import ( discover_flash_targets, resolve_flash_gateway_url, @@ -482,10 +481,10 @@ async def advance_to(run_id: str | None, version: int) -> None: # Singleton writer: a single small write is one atomic S3 PutObject, # so no rename dance is needed. The pointer is self-identifying # (`/weight_vN`), so a new run is a forward move (not a rewind) - # and there is no separate run pointer to flip. - (S3_TRANSPORT_MOUNT_PATH / "latest").write_text( - format_snapshot_identity(run_id, version), encoding="utf-8" - ) + # and there is no separate run pointer to flip. The monotonic/reset + # decision already ran in advance_latest_decision, so this writes the + # decided move through the same board the pool reconciles against. + board.write_latest(run_id, version) async def list_server_infos() -> list[dict]: targets = await asyncio.to_thread( diff --git a/cookbook/standalone_rollouts/provider_test.py b/cookbook/standalone_rollouts/provider_test.py index d94f097..4af57fe 100644 --- a/cookbook/standalone_rollouts/provider_test.py +++ b/cookbook/standalone_rollouts/provider_test.py @@ -89,9 +89,11 @@ async def run() -> None: ) await manager.startup_sync() - # Reconciled to the `latest` pointer, applying the chain in order. + # Reconciled to the `latest` pointer: the v1..v2 tail composes + # into a single engine apply at the target version (not one apply + # per intermediate version) — see WeightSyncManager._sync_once. self.assertEqual(manager.current_version, 2) - self.assertEqual(engine.applies, [1, 2]) + self.assertEqual(engine.applies, [2]) asyncio.run(run()) diff --git a/src/stitch/bulletin.py b/src/stitch/bulletin.py index 473a7ec..b91a703 100644 --- a/src/stitch/bulletin.py +++ b/src/stitch/bulletin.py @@ -9,8 +9,11 @@ from typing import Any, Protocol from stitch.protocol import ( + BASE_VERSION, + PointerMove, VersionManifest, atomic_write_text, + decide_pointer_move, format_snapshot_identity, parse_snapshot_identity, read_latest, @@ -29,6 +32,10 @@ def read_latest(self) -> tuple[str | None, int]: ... def write_latest(self, run_id: str | None, version: int) -> None: ... + def advance(self, run_id: str | None, version: int) -> PointerMove: ... + + def claim(self, run_id: str) -> PointerMove: ... + def version_dir(self, run_id: str | None, version: int) -> Path: ... def read_manifest(self, run_id: str | None, version: int) -> VersionManifest: ... @@ -99,11 +106,44 @@ def read_latest(self) -> tuple[str | None, int]: return (None, read_latest(self.root)) def write_latest(self, run_id: str | None, version: int) -> None: + """Overwrite the pointer unconditionally. Prefer :meth:`advance` / + :meth:`claim`, which enforce the monotonic-within-run rule; this raw + write is for callers that have already made the move decision.""" if self.layout == "slime": atomic_write_text(self.root / "latest", format_snapshot_identity(run_id, version)) else: write_latest(self.root, version) + def advance(self, run_id: str | None, version: int) -> PointerMove: + """Move ``latest`` to ``(run_id, version)`` under the shared monotonic + rule (see :func:`decide_pointer_move`), then write it. + + Returns the :class:`PointerMove` (``reset=True`` when the move crosses + runs); raises :class:`PointerRewind` on a same-run rewind. This is the + single guarded writer both cookbook patterns publish through, so neither + can silently rewind the pointer onto stale weights. + """ + current_run_id, current_version = self.read_latest() + move = decide_pointer_move( + current_run_id, current_version, run_id=run_id, version=version + ) + self.write_latest(move.run_id, move.version) + return move + + def claim(self, run_id: str) -> PointerMove: + """Claim the pool for a fresh run: advance to the empty pointer + ``/weight_v000000``, resetting every replica to base before any + delta is published. + + ``run_id`` is the run's epoch/fence token and must be unique per launch; + claiming a run already at the pointer is a rewind (a reused run_id after + a restart), which surfaces as :class:`PointerRewind` rather than leaving + the pool pinned to the dead incarnation's high-water mark. + """ + if not run_id: + raise ValueError("claim requires a run_id (the run's per-launch epoch token)") + return self.advance(run_id, BASE_VERSION) + def version_dir(self, run_id: str | None, version: int) -> Path: if self.layout == "slime": base = self.root / run_id if run_id else self.root diff --git a/src/stitch/bulletin_test.py b/src/stitch/bulletin_test.py index 7380ee6..c011a1b 100644 --- a/src/stitch/bulletin_test.py +++ b/src/stitch/bulletin_test.py @@ -7,6 +7,7 @@ from stitch.bulletin import FilesystemBulletinBoard from stitch.protocol import ( + PointerRewind, VersionManifest, format_snapshot_identity, parse_snapshot_identity, @@ -106,6 +107,46 @@ def test_publish_manifest_advances_run_pointer(self) -> None: ) self.assertEqual(board.read_latest(), ("run-a", 1)) + def test_advance_is_monotonic_within_run_and_resets_across_runs(self) -> None: + with tempfile.TemporaryDirectory() as tmp: + board = FilesystemBulletinBoard(Path(tmp), layout="slime") + + first = board.advance("run-a", 1) + self.assertEqual((first.run_id, first.version, first.reset), ("run-a", 1, True)) + self.assertFalse(board.advance("run-a", 2).reset) + self.assertEqual(board.read_latest(), ("run-a", 2)) + + with self.assertRaises(PointerRewind): + board.advance("run-a", 2) + # A rejected advance leaves the pointer untouched. + self.assertEqual(board.read_latest(), ("run-a", 2)) + + crossed = board.advance("run-b", 1) + self.assertTrue(crossed.reset) + self.assertEqual(board.read_latest(), ("run-b", 1)) + + def test_claim_writes_empty_pointer_and_rejects_reuse(self) -> None: + with tempfile.TemporaryDirectory() as tmp: + board = FilesystemBulletinBoard(Path(tmp), layout="slime") + + move = board.claim("run-a") + self.assertTrue(move.reset) + self.assertEqual(board.read_latest(), ("run-a", 0)) + + board.advance("run-a", 1) + # Re-claiming the same run (a restart that reused its run_id) rewinds. + with self.assertRaises(PointerRewind): + board.claim("run-a") + # A fresh run id claims cleanly. + self.assertTrue(board.claim("run-b").reset) + self.assertEqual(board.read_latest(), ("run-b", 0)) + + def test_claim_requires_run_id(self) -> None: + with tempfile.TemporaryDirectory() as tmp: + board = FilesystemBulletinBoard(Path(tmp), layout="slime") + with self.assertRaises(ValueError): + board.claim("") + def test_unknown_layout_rejected(self) -> None: with self.assertRaises(ValueError): FilesystemBulletinBoard("/tmp", layout="bogus") diff --git a/src/stitch/protocol.py b/src/stitch/protocol.py index 5644006..195199b 100644 --- a/src/stitch/protocol.py +++ b/src/stitch/protocol.py @@ -14,6 +14,11 @@ PROTOCOL_VERSION = 1 LATEST_FILE = "latest.json" +# The "empty" pointer a launch claims before publishing any delta: version 0 is +# the run's base (every run forks at base, so a chain starts at v1 on top of v0). +# A pool reconciling to ``/weight_v000000`` resets to base weights. +BASE_VERSION = 0 + class SyncState(str, Enum): IDLE = "IDLE" @@ -403,6 +408,73 @@ def parse_snapshot_identity(text: str) -> tuple[str | None, int]: return (run_id, version) +class PointerRewind(Exception): + """A pointer move would rewind ``latest`` within the same run. + + The only writer of ``latest`` advances it monotonically *within a run*; a + move to an equal-or-lower version on the same run (e.g. a restarted trainer + that reused its run_id, or a duplicate claim) is rejected as a rewind rather + than silently serving stale weights. Crossing to a *different* run is not a + rewind — it forks at base (see :func:`decide_pointer_move`). + """ + + def __init__( + self, *, run_id: str | None, current_version: int, requested_version: int + ) -> None: + super().__init__( + f"latest is at version {current_version} (run {run_id!r}); " + f"refusing to rewind to {requested_version}" + ) + self.run_id = run_id + self.current_version = int(current_version) + self.requested_version = int(requested_version) + + +@dataclass(frozen=True) +class PointerMove: + """An accepted move of the ``latest`` pointer. + + ``reset`` is True when the move crosses to a different run, i.e. the pool + must re-materialize base and restart the version space (a claim, or a new + run's first publish); False for an ordinary monotonic advance within a run. + """ + + run_id: str | None + version: int + reset: bool + + +def decide_pointer_move( + current_run_id: str | None, + current_version: int, + *, + run_id: str | None, + version: int, +) -> PointerMove: + """Decide whether the single ``latest`` writer may move to ``(run_id, version)``. + + The one rule both the trainer-as-writer (bulletin board) and the + frontdoor-as-writer (hot-load API) paths share, so they can't bake in + divergent semantics: + + - A *different* run forks at base, so its version space restarts; accepting + it (even at a lower number, including the ``BASE_VERSION`` claim) is a + reset, not a rewind. + - Within the same run (and the run-less layout where both ids are None) the + move must be strictly newer; otherwise :class:`PointerRewind` is raised. + """ + version = int(version) + if run_id != current_run_id: + return PointerMove(run_id=run_id, version=version, reset=True) + if version <= int(current_version): + raise PointerRewind( + run_id=current_run_id, + current_version=current_version, + requested_version=version, + ) + return PointerMove(run_id=run_id, version=version, reset=False) + + def version_dir(root: str | Path, version: int) -> Path: return Path(root) / "versions" / weight_identity(version) diff --git a/src/stitch/protocol_test.py b/src/stitch/protocol_test.py index 37abe90..97acf76 100644 --- a/src/stitch/protocol_test.py +++ b/src/stitch/protocol_test.py @@ -6,11 +6,14 @@ from stitch.bulletin import FilesystemBulletinBoard from stitch.protocol import ( + BASE_VERSION, Artifact, + PointerRewind, RolloutPoolState, RolloutReplicaState, VersionManifest, WeightVersionPolicy, + decide_pointer_move, evaluate_version_policy, parse_weight_identity, read_latest, @@ -18,6 +21,41 @@ ) +class DecidePointerMoveTest(unittest.TestCase): + """The single accept/reset/rewind rule both the bulletin-board publish path + and the frontdoor hot-load path share (so they can't diverge).""" + + def test_forward_within_run_is_a_non_reset_advance(self) -> None: + move = decide_pointer_move("run-a", 4, run_id="run-a", version=5) + self.assertEqual((move.run_id, move.version, move.reset), ("run-a", 5, False)) + + def test_same_or_lower_version_within_run_rewinds(self) -> None: + with self.assertRaises(PointerRewind): + decide_pointer_move("run-a", 5, run_id="run-a", version=5) + with self.assertRaises(PointerRewind) as cm: + decide_pointer_move("run-a", 5, run_id="run-a", version=3) + self.assertEqual(cm.exception.current_version, 5) + self.assertEqual(cm.exception.requested_version, 3) + + def test_different_run_forks_at_base_as_a_reset(self) -> None: + # A new run is accepted even at a lower version (its space restarts) ... + move = decide_pointer_move("run-a", 5, run_id="run-b", version=1) + self.assertEqual((move.run_id, move.version, move.reset), ("run-b", 1, True)) + # ... including the empty BASE_VERSION claim. + claim = decide_pointer_move("run-a", 5, run_id="run-b", version=BASE_VERSION) + self.assertEqual((claim.run_id, claim.version, claim.reset), ("run-b", 0, True)) + + def test_first_claim_against_empty_pointer_is_a_reset(self) -> None: + move = decide_pointer_move(None, 0, run_id="run-a", version=BASE_VERSION) + self.assertEqual((move.run_id, move.version, move.reset), ("run-a", 0, True)) + + def test_runless_layout_keeps_monotonic_cas(self) -> None: + move = decide_pointer_move(None, 2, run_id=None, version=3) + self.assertFalse(move.reset) + with self.assertRaises(PointerRewind): + decide_pointer_move(None, 3, run_id=None, version=3) + + class ProtocolTest(unittest.TestCase): def test_manifest_round_trips_extended_and_legacy_fields(self) -> None: with tempfile.TemporaryDirectory() as tmp: From 390804f01f46876bef59e96c46e347be3eb3b14f Mon Sep 17 00:00:00 2001 From: Jason Mancuso <7891333+jvmncs@users.noreply.github.com> Date: Tue, 30 Jun 2026 08:42:13 -0700 Subject: [PATCH 07/18] Consolidate cookbook trainers onto a shared sidecar/serving/launch spine (#9) * Consolidate cookbook trainers onto a shared sidecar/serving/launch spine Collapse the accidentally-forked cookbook scaffolding onto one trainer-agnostic spine, leaving each example to encode only its intended axes (slime vs miles quant; disagg trainer-writer vs standalone frontdoor-writer): - cookbook/sidecar.py: unified sidecar (arg parsing, bulletin board + Volume reloader, concurrent stale-aware base materialization). slime/miles sidecars become thin adapters naming only their host-side delta decoder. - cookbook/serving.py: one B200 SGLang image builder; build_int4/build_nvfp4 remain as thin wrappers. - cookbook/trainer_helpers.py: shared config-prep, train-command, host-RAM monitor, and a smoke check with a wake_on_demand flag. slime/miles helpers re-export with their per-trainer params. Fixes a latent missing socket import in the miles host-mem monitor. - cookbook/rollout_control.py: shared trainer-side primitives (distributed_rank, read_setting, session affinity) for both control planes; collapses the standalone slime hook fork onto them. - standalone announce_claim: claim-at-launch parity with bulletin claim_pool, wired into the standalone trainer launch. - Tests for the new shared modules + announce_claim. Co-Authored-By: jason.mancuso@modal.com * Fix consolidation runtime gaps flagged in review - serving image: mount the whole cookbook package at /root/cookbook (not just the per-trainer subdir) so the sidecar subprocess (python3 -m cookbook..sidecar), which is never imported at deploy time and so is invisible to Modal's import-time automounting, can resolve the shared cookbook.sidecar spine it delegates to. - modal_probes: import parallel_init_local_checkpoint from cookbook.sidecar (the function moved there from cookbook.miles_disagg.sidecar and gained a required disk_delta_module arg). Co-Authored-By: jason.mancuso@modal.com * Mount the whole cookbook package in the trainer images The consolidation moved the shared spine (helpers/hooks/sidecar/ trainer_helpers/ray_cluster/sidecar_process) up to the cookbook package root, so cookbook..helpers now imports cookbook.* root modules and the per-trainer sidecar is a thin adapter over cookbook.sidecar. The trainer images still mounted only their per-trainer subdir (/root/cookbook/) with include_source=False, so those root modules were absent at runtime: the trainer's helpers import and the 'python3 -m cookbook..sidecar' subprocess (the server image for the small trainer-image-reuse experiments, e.g. qwen) would ImportError. Mount the whole cookbook package, matching cookbook/serving.py and the standalone modal_train/modal_serve images (which were already widened). Co-Authored-By: jason.mancuso@modal.com --------- Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- cookbook/bulletin_hooks.py | 21 +- cookbook/miles_disagg/helpers.py | 228 ++------------- cookbook/miles_disagg/modal_train.py | 9 +- cookbook/miles_disagg/serving.py | 133 ++------- cookbook/miles_disagg/sidecar.py | 221 +------------- cookbook/rollout_control.py | 80 ++++++ cookbook/rollout_control_test.py | 59 ++++ cookbook/serving.py | 141 +++++++++ cookbook/sidecar.py | 252 ++++++++++++++++ cookbook/sidecar_test.py | 109 +++++++ cookbook/slime_disagg/helpers.py | 134 ++------- cookbook/slime_disagg/modal_train.py | 9 +- cookbook/slime_disagg/serving.py | 115 +------- cookbook/slime_disagg/sidecar.py | 132 +-------- cookbook/standalone_rollouts/slime/hooks.py | 81 +++--- .../standalone_rollouts/slime/hooks_test.py | 40 +++ .../standalone_rollouts/slime/modal_train.py | 6 + cookbook/trainer_helpers.py | 271 ++++++++++++++++++ cookbook/trainer_helpers_test.py | 77 +++++ 19 files changed, 1204 insertions(+), 914 deletions(-) create mode 100644 cookbook/rollout_control.py create mode 100644 cookbook/rollout_control_test.py create mode 100644 cookbook/serving.py create mode 100644 cookbook/sidecar.py create mode 100644 cookbook/sidecar_test.py create mode 100644 cookbook/trainer_helpers.py create mode 100644 cookbook/trainer_helpers_test.py diff --git a/cookbook/bulletin_hooks.py b/cookbook/bulletin_hooks.py index 1e742f6..3493fac 100644 --- a/cookbook/bulletin_hooks.py +++ b/cookbook/bulletin_hooks.py @@ -22,6 +22,7 @@ from pathlib import Path from typing import Any +from cookbook.rollout_control import apply_session_affinity, distributed_rank from stitch.bulletin import FilesystemBulletinBoard from stitch.protocol import BASE_VERSION, PointerRewind, parse_weight_identity from stitch.providers.modal import commit_volume, discover_flash_targets, volume_reloader, wake_targets @@ -186,29 +187,13 @@ async def gated_rollout_request_hook(args: Any, sample: Any, request: dict[str, request["max_retries"] = int(getattr(args, "rollout_request_retry_attempts", request.get("max_retries", 60))) request["retry_sleep"] = float(getattr(args, "rollout_request_retry_sleep", request.get("retry_sleep", 1.0))) - session_id = getattr(sample, "session_id", None) - if session_id: - header = str(getattr(args, "rollout_session_affinity_header", "x-session-affinity")) - headers = dict(request.get("headers") or {}) - headers.setdefault(header, session_id) - request["headers"] = headers + header = str(getattr(args, "rollout_session_affinity_header", "x-session-affinity")) + apply_session_affinity(request, getattr(sample, "session_id", None), header) # ── Shared helpers ──────────────────────────────────────────────────────────── -def distributed_rank() -> int | None: - """Return the torch distributed rank, or None if not initialized.""" - try: - import torch.distributed as dist - - if dist.is_available() and dist.is_initialized(): - return int(dist.get_rank()) - except Exception: # noqa: BLE001 - return None - return None - - def _volume_name(args: Any) -> str: return str(getattr(args, "update_weight_delta_volume_name", None) or os.environ["DELTA_VOLUME_NAME"]) diff --git a/cookbook/miles_disagg/helpers.py b/cookbook/miles_disagg/helpers.py index c14ceeb..48dfd66 100644 --- a/cookbook/miles_disagg/helpers.py +++ b/cookbook/miles_disagg/helpers.py @@ -1,23 +1,20 @@ """Trainer-specific helpers for the miles_disagg example. -Ray cluster, sidecar, and process helpers are shared across trainers via -:mod:`cookbook.ray_cluster` and :mod:`cookbook.sidecar_process`. This module -provides miles-specific config preparation, train command building, and the -Flash pool smoke check. +Thin wrappers over the shared launch spine: Ray-cluster/sidecar/process helpers +come from :mod:`cookbook.ray_cluster` / :mod:`cookbook.sidecar_process`, and +config-prep / train-command / smoke-check / host-RAM monitor come from +:mod:`cookbook.trainer_helpers`. This module only supplies the miles-specific +axes: the sidecar module path, the config-field tuple to materialize, the +model-script attribute, and that the rollout pool scales from zero (wake on +demand). """ from __future__ import annotations -import json -import os -import shlex import subprocess -import time -import urllib.error -import urllib.request from typing import Any -from stitch.providers.modal import discover_flash_targets, resolve_flash_gateway_url +from cookbook.miles_disagg.configs.base import YAML_CONFIG_FIELDS # Re-export shared helpers so existing callers (modal_train.py) don't break. from cookbook.ray_cluster import ( # noqa: F401 @@ -32,6 +29,14 @@ terminate_process, wait_http, ) +from cookbook.trainer_helpers import ( # noqa: F401 + VersionAheadError, + build_train_cmd as _build_train_cmd, + materialize_node_local_yaml, + prepare_config, + smoke_flash_pool as _smoke_flash_pool, + start_host_mem_monitor, +) SIDECAR_MODULE = "cookbook.miles_disagg.sidecar" @@ -63,9 +68,6 @@ def start_sglang_sidecar( ) -# ── miles launch ────────────────────────────────────────────────────────────── - - def prepare_miles_config(miles_cfg: Any, tmpdir: str) -> None: """Resolve HF repo IDs to local paths and materialize inline YAML configs. @@ -73,70 +75,12 @@ def prepare_miles_config(miles_cfg: Any, tmpdir: str) -> None: NVFP4 base and bf16 masters), so the ``startswith("/")`` guard skips them; a repo-id-shaped value (if any) is snapshot-downloaded from the HF cache. """ - from huggingface_hub import snapshot_download - import yaml - - from cookbook.miles_disagg.configs.base import YAML_CONFIG_FIELDS - - for attr in ("hf_checkpoint", "load", "ref_load", "critic_load"): - if (val := getattr(miles_cfg, attr, None)) and not str(val).startswith("/"): - setattr(miles_cfg, attr, snapshot_download(val, local_files_only=True)) - - for field in YAML_CONFIG_FIELDS: - if isinstance(val := getattr(miles_cfg, field, None), dict): - path = os.path.join(tmpdir, f"{field}.yaml") - with open(path, "w") as f: - yaml.dump(val, f) - setattr(miles_cfg, field, path) - - -def materialize_node_local_yaml(miles_cfg: Any, field: str, dest_dir: str = "/root/.miles_node_yaml") -> None: - """Materialize a per-actor-read YAML config to a deterministic node-local path. - - Some config files (notably ``te_precision_config_file``, which - ``load_quantization_recipe`` re-reads on every Ray actor during model build) - are read independently on each trainer node — not just parsed once on the head. - ``prepare_miles_config`` writes them under ``tempfile.mkdtemp()`` on the head - only, so on a multi-node cluster the other containers can't see that path. - - Call this on EVERY node (SPMD train()), before the rank-0 gate: each node - writes identical content (from the shared payload) to the same fixed path, so - the path the head embeds in the args resolves locally on all actors. No volume - commit/reload race — Ray actors are long-lived and wouldn't see post-start - volume writes anyway. - """ - import yaml - - if isinstance(val := getattr(miles_cfg, field, None), dict): - os.makedirs(dest_dir, exist_ok=True) - path = os.path.join(dest_dir, f"{field}.yaml") - with open(path, "w") as f: - yaml.dump(val, f) - setattr(miles_cfg, field, path) + prepare_config(miles_cfg, tmpdir, YAML_CONFIG_FIELDS) def build_train_cmd(miles_cfg: Any, miles_root: str) -> str: - """Build the training command, sourcing model arch args if needed. - - miles' train.py / train_async.py live at the repo root and consume the - ``MODEL_ARGS`` bash array defined by the sourced model script, exactly like - slime's launcher. - """ - train_script = f"{miles_root}/{'train_async.py' if miles_cfg.async_mode else 'train.py'}" - if miles_cfg.miles_model_script: - inner = ( - f"source {miles_root}/{miles_cfg.miles_model_script} && " - f"python3 {train_script} ${{MODEL_ARGS[@]}} {shlex.join(miles_cfg.cli_args())}" - ) - return f"bash -c {shlex.quote(inner)}" - return f"python3 {train_script} {shlex.join(miles_cfg.cli_args())}" - - -# ── Flash pool smoke check ──────────────────────────────────────────────────── - - -class VersionAheadError(RuntimeError): - """Raised when a monotonic rollout pool has already advanced past a smoke version.""" + """Build the training command, sourcing miles' model arch args if needed.""" + return _build_train_cmd(miles_cfg, miles_root, model_script_attr="miles_model_script") def smoke_flash_pool( @@ -148,128 +92,14 @@ def smoke_flash_pool( expect_min_containers: int, timeout_seconds: int, ) -> None: - """Wake the elastic pool on demand and confirm it serves at the expected - weight version. - - The pool has no warm floor (min_containers=0), so a completion sent to the - Flash gateway is what scales it 0->1; Flash holds the request during the - container's cold start (model load + FP4 kernel tuning), so the warmup uses a - generous timeout. ``expect_min_containers`` is advisory only — a value > 0 - just means "also confirm at least one direct container reports the version." - """ - deadline = time.time() + timeout_seconds - last_error: str | None = None - while True: - try: - _check_flash_pool_once(app_name, cls_name, model_name, weight_version) - return - except VersionAheadError: - raise - except Exception as exc: # noqa: BLE001 - last_error = f"{type(exc).__name__}: {exc}" - if time.time() >= deadline: - raise TimeoutError(f"Flash pool smoke did not pass before timeout: {last_error}") - print(f"Waiting for Flash pool to wake/serve: {last_error}") - time.sleep(10) - - -def _check_flash_pool_once(app_name: str, cls_name: str, model_name: str, expected: int) -> None: - gateway = resolve_flash_gateway_url(app_name, cls_name) - print(f"Gateway URL: {gateway}") - - # Wake the (scaled-to-zero) pool and wait for a container to serve. Flash - # holds the request through the cold start, so the timeout must exceed it. - payload = { - "model": model_name, - "messages": [{"role": "user", "content": "Reply with exactly OK."}], - "max_tokens": 8, - "temperature": 0, - "weight_version": {"exact_version": expected}, - "chat_template_kwargs": {"enable_thinking": False}, - } - data = _post_json(f"{gateway}/v1/chat/completions", payload, timeout=900) - print(f"Gateway completion: {data}") - start, end = int(data.get("weight_version_start", -1)), int(data.get("weight_version_end", -1)) - if start > expected or end > expected: - raise VersionAheadError(f"gateway served version {start}->{end}, already past expected {expected}") - if start != expected or end != expected: - raise RuntimeError(f"unexpected gateway weight metadata: {data}") - - # The pool is warm now; confirm each live container reports the version. - targets = discover_flash_targets(app_name, cls_name) - print(f"Direct container URLs ({len(targets)}):") - for target in [gateway, *targets]: - info = _get_json(f"{target}/server_info", timeout=30) - print(f"{target} server_info={info}") - current = int(info["current_version"]) - if current > expected: - raise VersionAheadError(f"{target} current_version={current} already passed expected {expected}") - if current != expected: - raise RuntimeError(f"{target} current_version={current} expected {expected}") - - -def _get_json(url: str, *, timeout: float) -> dict: - with urllib.request.urlopen(url, timeout=timeout) as resp: - return json.load(resp) - - -def _post_json(url: str, payload: dict, *, timeout: float) -> dict: - request = urllib.request.Request( - url, - data=json.dumps(payload).encode(), - headers={"Content-Type": "application/json"}, + """Smoke the scale-from-zero miles rollout pool (min_containers=0): the + completion wakes it, then each warmed container is confirmed at the version.""" + _smoke_flash_pool( + app_name=app_name, + cls_name=cls_name, + model_name=model_name, + weight_version=weight_version, + expect_min_containers=expect_min_containers, + timeout_seconds=timeout_seconds, + wake_on_demand=True, ) - with urllib.request.urlopen(request, timeout=timeout) as resp: - return json.load(resp) - - -def start_host_mem_monitor(interval_s: int = 20) -> None: - """Log this node's host-RAM trajectory to stdout from a daemon thread. - - The trainer can OOM-kill on host-RAM exhaustion (the publish/update_weights - full-model gather is the peak consumer), but Megatron only reports GPU memory - and the kill leaves no durable peak behind. This logs MemTotal/MemAvailable + - the container cgroup usage every ``interval_s`` so a live ``modal app logs -f`` - shows exactly which phase blows the ~1.95 TiB B200:8 node and how high it peaks. - Runs on EVERY node (called from the SPMD enter()), so whichever rank OOMs has - its own trace. Best-effort: never raises.""" - import threading - - host = socket.gethostname() - - def _meminfo() -> tuple[float, float]: - total = avail = 0.0 - try: - with open("/proc/meminfo") as f: - for line in f: - if line.startswith("MemTotal:"): - total = int(line.split()[1]) / 1024 / 1024 # GiB - elif line.startswith("MemAvailable:"): - avail = int(line.split()[1]) / 1024 / 1024 - except Exception: # noqa: BLE001 - pass - return total, avail - - def _cgroup_used_gib() -> float: - for path in ("/sys/fs/cgroup/memory.current", # cgroup v2 - "/sys/fs/cgroup/memory/memory.usage_in_bytes"): # v1 - try: - with open(path) as f: - return int(f.read().strip()) / 1024**3 - except Exception: # noqa: BLE001 - continue - return -1.0 - - def _loop() -> None: - while True: - total, avail = _meminfo() - used = total - avail - cg = _cgroup_used_gib() - print( - f"[hostmem] {host} used={used:.0f}GiB avail={avail:.0f}GiB " - f"total={total:.0f}GiB cgroup_used={cg:.0f}GiB", - flush=True, - ) - time.sleep(interval_s) - - threading.Thread(target=_loop, daemon=True, name="host-mem-monitor").start() diff --git a/cookbook/miles_disagg/modal_train.py b/cookbook/miles_disagg/modal_train.py index c01aa7c..76857f4 100644 --- a/cookbook/miles_disagg/modal_train.py +++ b/cookbook/miles_disagg/modal_train.py @@ -189,10 +189,13 @@ # Local source mounted at container start (no rebuild on code edits). Modal puts # /root on PYTHONPATH, so both packages import from subprocesses (the sidecar, Ray # workers). MUST come after any .run_commands() above (Modal forbids run_commands -# after a non-copy local mount). +# after a non-copy local mount). The whole cookbook package is mounted (not just +# the per-trainer subdir) so the trainer and the `python3 -m +# cookbook.miles_disagg.sidecar` subprocess can import the shared cookbook spine +# (helpers/hooks/sidecar) the thin adapters delegate to. image = image.add_local_python_source("stitch").add_local_dir( - Path(__file__).parent, - remote_path="/root/cookbook/miles_disagg", + Path(__file__).parent.parent, + remote_path="/root/cookbook", ignore=["**/__pycache__"], ) diff --git a/cookbook/miles_disagg/serving.py b/cookbook/miles_disagg/serving.py index 1bd2635..3e7a43c 100644 --- a/cookbook/miles_disagg/serving.py +++ b/cookbook/miles_disagg/serving.py @@ -1,26 +1,15 @@ -"""Dedicated B200 NVFP4 SGLang serving image for the rollout pool. - -The trainer half of this example runs on the miles/Megatron image -(``modal_train.image``). The rollout half that serves the NVFP4 checkpoint needs -a different stack: a Blackwell SGLang build that loads the model's NVFP4 weights -and runs MLA attention with the tuned hierarchical KV cache. This module builds -that image — the miles twin of cookbook/slime_disagg/serving.py. - -Two deliberate choices keep it lean: - - * **SGLang comes from the modal-projects/sglang fork** (Blackwell fa4 / - cutlass-dsl prerelease kernels + the tokenspeed MLA attention backend) — the - same build the slime example uses, proven for NVFP4. No ``--quantization`` - flag is baked in; the served checkpoint's own NVFP4 quant config drives - weight loading (see the config module's ``SGLANG_SERVER_ARGS``). - * **miles is cloned ``--no-deps`` for one module.** The sidecar only imports - ``miles.utils.disk_delta`` (stdlib + numpy + zstandard; xxhash/blake3 lazy), - so Megatron is intentionally absent from the rollout pool. NOTE: this assumes - ``miles.utils.disk_delta`` is import-light (no heavy package __init__ chain); - verify on a warm container during bring-up. - -Pin the SAME miles ref the trainer image uses so the host-side delta decoder -matches the trainer's delta encoder. +"""Dedicated B200 NVFP4 SGLang serving image for the miles rollout pool. + +A thin wrapper over :func:`cookbook.serving.build_b200_serving_image` — the miles +twin of cookbook/slime_disagg/serving.py. The image itself is trainer-agnostic +(see that module): NVFP4 vs INT4 is driven by the served checkpoint's own quant +config, not by this builder. This wrapper pins miles as the ``--no-deps`` decoder +package, does a full clone, and clears the SGLang kernel cache as the final +filesystem step (modal_train mounts a kernel-cache volume at /root/.cache/sglang, +which can't mount over a non-empty path). + +NOTE: this assumes ``miles.utils.disk_delta`` is import-light (no heavy package +__init__ chain); verify on a warm container during bring-up. """ from __future__ import annotations @@ -29,26 +18,7 @@ import modal -# Pinned Blackwell SGLang fork — matches the slime 4xB200 Kimi serve recipe, -# proven for NVFP4. Only the python sources are checked out over the prebuilt -# base image's kernels. -SGLANG_IMAGE_TAG = "lmsysorg/sglang:v0.5.12" -SGLANG_FORK_REPO = "https://github.com/modal-projects/sglang.git" -SGLANG_FORK_BRANCH = "timmy/dflash-fa4-fp8" -SGLANG_FORK_COMMIT = "dafb2b325b40298c5097564811463c585b7e9814" - -# SGLang runtime tunables carried over from the B200 deployment. -SERVING_IMAGE_ENV = { - "HF_XET_HIGH_PERFORMANCE": "1", - "HF_HUB_ENABLE_HF_TRANSFER": "1", - "SGLANG_ALLOW_OVERWRITE_LONGER_CONTEXT_LEN": "1", - "SGLANG_DISABLE_CUDNN_CHECK": "1", - "SGLANG_ENABLE_OVERLAP_PLAN_STREAM": "1", - "SGLANG_TIMEOUT_KEEP_ALIVE": "300", -} -# NOTE: kernel-cache persistence is handled by mounting SGLang's cache dir -# (/root/.cache/sglang, which nests flashinfer/DeepGemm) as a volume in -# modal_train — no env override here, so SGLang's default placement is kept. +from cookbook.serving import build_b200_serving_image def build_nvfp4_b200_serving_image( @@ -59,72 +29,13 @@ def build_nvfp4_b200_serving_image( hf_cache_path: str, experiment: str, ) -> modal.Image: - """Build the rollout-pool serving image (see module docstring). - - The miles fork ref / root and the HF cache path are passed in by - ``modal_train`` so the pool and the trainer pin the identical miles commit - (the sidecar's ``disk_delta`` must match the trainer's delta encoder). - """ - # serving.py lives in cookbook/miles_disagg, mounted to /root/cookbook/miles_disagg - # exactly as the trainer image mounts it, so `cookbook.miles_disagg.sidecar` - # imports identically in either container. - miles_disagg_dir = Path(__file__).parent - return ( - modal.Image.from_registry(SGLANG_IMAGE_TAG) - .run_commands( - f"cd /sgl-workspace/sglang && git remote add modal-fork {SGLANG_FORK_REPO}" - f" && git fetch modal-fork {SGLANG_FORK_BRANCH}" - f" && git checkout {SGLANG_FORK_COMMIT} -- python/", - ) - # Pre-release CUDA wheels (cutlass-dsl / sglang-kernel / flash-attn-4) — - # keep the deployment's known-good pip resolution. - .run_commands( - "pip install nvidia-cutlass-dsl==4.5.1 sglang-kernel==0.4.3 'flash-attn-4>=4.0.0b10'" - ) - # flash-attn-4 checks for the deprecated MmaFP8Op but cutlass-dsl 4.5.1 now - # generates MmaF8F6F4Op instead. Patch the isinstance check to handle both. - .run_commands( - "sed -i 's/isinstance(op, tcgen05.mma.MmaFP8Op)/isinstance(op, (tcgen05.mma.MmaFP8Op, tcgen05.mma.MmaF8F6F4Op))/' " - "/usr/local/lib/python3.12/dist-packages/flash_attn/cute/blackwell_helpers.py" - ) - # The base image bakes in an HF cache; remove it so it cannot shadow the - # cache volume mounted at the same path. - .run_commands(f"rm -rf {hf_cache_path}") - # miles --no-deps gives the sidecar `miles.utils.disk_delta` (host-side - # delta apply). Megatron is NOT installed — the pool never trains. Pin the - # SAME ref the trainer image uses so the delta encoder/decoder match. - .run_commands( - f"git clone {miles_repo_url} {miles_root}" - f" && cd {miles_root}" - f" && git fetch origin {miles_repo_ref}" - f" && git checkout FETCH_HEAD" - f" && python3 -m pip install --no-deps -e {miles_root}" - ) - .pip_install( - "autoinference-utils==0.2.0", # SGLang server lifecycle for the rollout pool - "fastapi", # stitch sidecar - "httpx", # stitch sidecar - "uvicorn", # stitch sidecar - # disk_delta host-side apply: zstd decompress + xxhash (xxh3-128 - # default) / blake3 checksums. miles is installed --no-deps. - "zstandard", - "xxhash", - "blake3", - ) - # MUST be the last filesystem step: modal_train mounts the - # miles-sglang-cache volume at /root/.cache/sglang, and a volume can't - # mount over a non-empty path. The sglang-kernel/flashinfer and miles - # installs above populate this dir, so clear it AFTER them (the volume - # repopulates the JIT/autotuner cache on first boot). - .run_commands("rm -rf /root/.cache/sglang") - .env({"EXPERIMENT_CONFIG": experiment, **SERVING_IMAGE_ENV}) - # Mounted at container start (not copied into the image) so code edits to - # stitch / the sidecar never rebuild the image. Modal puts /root on - # PYTHONPATH for subprocesses (the sidecar). - .add_local_python_source("stitch") - .add_local_dir( - miles_disagg_dir, - remote_path="/root/cookbook/miles_disagg", - ignore=["**/__pycache__"], - ) + return build_b200_serving_image( + trainer_repo_url=miles_repo_url, + trainer_repo_ref=miles_repo_ref, + trainer_root=miles_root, + cookbook_dir=Path(__file__).parent, + hf_cache_path=hf_cache_path, + experiment=experiment, + shallow_clone=False, + clear_sglang_cache_at_end=True, ) diff --git a/cookbook/miles_disagg/sidecar.py b/cookbook/miles_disagg/sidecar.py index 912bc23..d629fb0 100644 --- a/cookbook/miles_disagg/sidecar.py +++ b/cookbook/miles_disagg/sidecar.py @@ -1,223 +1,26 @@ -"""SGLang weight-sync sidecar launcher for the miles_disagg example. +"""SGLang weight-sync sidecar entry point for the miles_disagg example. -The reusable versioned-proxy library lives in ``stitch.servers.sglang``; this -module wires the concrete bulletin-board + Modal Volume + SGLang-disk-delta -realization and runs it as a process (``helpers.start_sglang_sidecar`` launches -``python3 -m cookbook.miles_disagg.sidecar``). +A thin adapter over :mod:`cookbook.sidecar`: the shared spine owns every knob; +this only names miles' host-side delta decoder and asks for it to be injected. +``helpers.start_sglang_sidecar`` launches it via +``python3 -m cookbook.miles_disagg.sidecar``. -The ONLY functional difference from cookbook/slime_disagg/sidecar.py is the -host-side delta decoder: this injects ``miles.utils.disk_delta`` into the -adapter instead of letting it default to ``slime.utils.disk_delta``. The two are -byte-identical (same XOR/overwrite + zstd + xxh3/blake3/adler32 wire format), but -the rollout pool image ships miles (``--no-deps``), not slime, so the miles -functions must be injected explicitly. +miles' ``disk_delta`` is byte-identical to slime's (same XOR/overwrite + zstd + +xxh3/blake3/adler32 wire format), but the rollout pool image ships miles +``--no-deps``, not slime, so the decoder must be injected explicitly rather than +relying on the engine's slime default. """ from __future__ import annotations -import argparse -import logging -import os +from cookbook.sidecar import run_sidecar -from stitch.bulletin import FilesystemBulletinBoard -from stitch.engines.sglang import SGLangDiskDeltaAdapter -from stitch.servers.sglang import create_app -from stitch.sync import CommitMode, WeightSyncManager -logger = logging.getLogger(__name__) - - -def _parallel_init_local_checkpoint(workers: int = 32): - """Return an ``init_local_checkpoint(local, base)`` that copies the base - shards concurrently. miles' default copies them one at a time with - ``shutil.copy2`` — fine for a 16 B Moonlight base, far too slow for K2.6's - ~591 GB. We reuse miles' apply-lock + applied-version bookkeeping so - ``apply_deltas`` stays compatible, and fall back to the miles default if those - internals move under us.""" - import concurrent.futures - import hashlib - import shutil - - from miles.utils import disk_delta as dd - - def _base_fingerprint(base_dir: str) -> str: - """Identity of the base's bytes: (filename, size, mtime_ns) over every shard. Re-prep - rewrites the shards (new mtime/size), so this changes even when the tensor index/shapes - don't — which is exactly the case that bit us (NVFP4 values changed, names didn't).""" - h = hashlib.sha256() - for e in sorted(os.scandir(base_dir), key=lambda e: e.name): - if e.is_file(): - st = e.stat() - h.update(f"{e.name}:{st.st_size}:{st.st_mtime_ns}\n".encode()) - return h.hexdigest() - - def _init(local_ckpt_dir: str, base_dir: str) -> None: - try: - apply_lock = dd._apply_lock - read_version = dd._read_applied_version - write_version = dd._write_applied_version - drop_page_cache = dd.drop_page_cache - except AttributeError: # miles internals changed — use its own copier - dd.init_local_checkpoint(local_ckpt_dir, base_dir) - return - - fp_path = os.path.join(local_ckpt_dir, ".base_fingerprint") - base_fp = _base_fingerprint(base_dir) - with apply_lock(local_ckpt_dir): - if read_version(local_ckpt_dir) is not None: - try: - cur = open(fp_path).read().strip() - except FileNotFoundError: - cur = None - if cur == base_fp: - return # current base already materialized - # STALE: /local-checkpoint holds a different (older) base — e.g. after a re-prep. - # Reusing it makes the host-side delta (diffed against the NEW base) XOR onto the - # WRONG bytes -> base_local XOR delta != export -> checksum mismatch on every tensor. - logger.warning( - "stale /local-checkpoint (base changed: %s != %s) — wiping + re-materializing", - cur, base_fp, - ) - shutil.rmtree(local_ckpt_dir, ignore_errors=True) - logger.info("Materializing base %s -> %s (%d copy workers)", base_dir, local_ckpt_dir, workers) - os.makedirs(local_ckpt_dir, exist_ok=True) - files = [e for e in os.scandir(base_dir) if e.is_file()] - - def _copy(entry: os.DirEntry) -> None: - import shutil - - shutil.copy2(entry.path, os.path.join(local_ckpt_dir, entry.name)) - drop_page_cache(entry.path) # don't evict the local copy we keep resident - - with concurrent.futures.ThreadPoolExecutor(max_workers=min(workers, max(1, len(files)))) as ex: - for _ in ex.map(_copy, files): # re-raises the first copy failure - pass - write_version(local_ckpt_dir, "000000") - with open(fp_path, "w") as f: - f.write(base_fp) - - return _init - - -def build_manager( - *, - upstream_url: str, - bulletin_root: str, - local_checkpoint_dir: str, - base_checkpoint_dir: str, - volume_name: str = "", - run_id: str | None = None, - commit_mode: CommitMode = "in_place", - debug_requests: bool = False, -) -> WeightSyncManager: - # The miles host-side delta decoder. miles is installed --no-deps in the - # serving image for exactly this module (Megatron is absent — the pool never - # trains). Pinned to the same miles ref as the trainer so encoder == decoder. - from miles.utils.disk_delta import apply_deltas - - refresh = None - if volume_name: - from stitch.providers.modal import volume_reloader - - refresh = volume_reloader(volume_name) - # miles publish-only writes the flat layout (weight_v{N}/ + - # model.safetensors.index.json + a raw `latest` pointer) to the Volume — the - # same layout slime uses, so layout="slime" reads it unchanged. - board = FilesystemBulletinBoard(bulletin_root, refresh=refresh, layout="slime") - engine = SGLangDiskDeltaAdapter( - upstream_url=upstream_url, - local_checkpoint_dir=local_checkpoint_dir, - base_checkpoint_dir=base_checkpoint_dir, - apply_deltas=apply_deltas, - init_local_checkpoint=_parallel_init_local_checkpoint(), - ) - return WeightSyncManager( - board=board, - engine=engine, - run_id=run_id, - commit_mode=commit_mode, - debug_requests=debug_requests, - ) +DISK_DELTA_MODULE = "miles.utils.disk_delta" def main() -> None: - parser = argparse.ArgumentParser() - parser.add_argument("--host", default="0.0.0.0") - parser.add_argument("--port", type=int, default=8000) - parser.add_argument("--upstream-url", required=True) - parser.add_argument( - "--bulletin-root", - default=os.environ.get("DELTA_BULLETIN_ROOT", "/delta-bulletin"), - ) - parser.add_argument("--volume-name", default=os.environ.get("DELTA_VOLUME_NAME", "")) - parser.add_argument( - "--local-checkpoint-dir", - default=os.environ.get("STITCH_LOCAL_CHECKPOINT_DIR", "/local-checkpoint"), - help="Writable host-local full HF checkpoint patched in place by each delta.", - ) - parser.add_argument( - "--base-checkpoint-dir", - default=os.environ.get("STITCH_BASE_CHECKPOINT_DIR"), - help="Base HF checkpoint the local copy is seeded from (deltas build on it).", - ) - parser.add_argument("--run-id", default=os.environ.get("DISAGG_RUN_ID")) - parser.add_argument( - "--commit-mode", - choices=("quiesce", "in_place"), - default=os.environ.get("SIDECAR_COMMIT_MODE", "in_place"), - help=( - "in_place (default): pause/apply/continue without flushing; in-flight " - "requests keep decoding on stale KV and version isolation comes from " - "extra_key stamping. quiesce: wait out active requests and flush " - "before applying (safe on any build)." - ), - ) - parser.add_argument( - "--debug-requests", - action="store_true", - default=os.environ.get("SIDECAR_DEBUG_REQUESTS", "").lower() in {"1", "true", "yes"}, - help="Log every versioned sidecar proxy request at INFO level.", - ) - parser.add_argument( - "--upstream-timeout", - type=float, - default=float(os.environ.get("SIDECAR_UPSTREAM_TIMEOUT", "3600")), - help=( - "Seconds to wait for an upstream SGLang response before failing the " - "request (5xx) instead of holding it open forever. Must exceed the " - "slowest legitimate generation." - ), - ) - args = parser.parse_args() - if not args.base_checkpoint_dir: - raise SystemExit( - "--base-checkpoint-dir/STITCH_BASE_CHECKPOINT_DIR is required: deltas are" - " applied host-side on top of a copy of this base HF checkpoint." - ) - - logging.basicConfig(level=logging.INFO) - import uvicorn - - manager = build_manager( - upstream_url=args.upstream_url, - bulletin_root=args.bulletin_root, - local_checkpoint_dir=args.local_checkpoint_dir, - base_checkpoint_dir=args.base_checkpoint_dir, - volume_name=args.volume_name, - run_id=args.run_id, - commit_mode=args.commit_mode, - debug_requests=args.debug_requests, - ) - uvicorn.run( - create_app( - manager, - upstream_url=args.upstream_url, - upstream_timeout=args.upstream_timeout, - ), - host=args.host, - port=args.port, - log_level="info", - ) + run_sidecar(disk_delta_module=DISK_DELTA_MODULE, inject_apply_deltas=True) if __name__ == "__main__": diff --git a/cookbook/rollout_control.py b/cookbook/rollout_control.py new file mode 100644 index 0000000..de601dd --- /dev/null +++ b/cookbook/rollout_control.py @@ -0,0 +1,80 @@ +"""Trainer-side primitives shared by both rollout-control planes. + +The cookbook has two *intended* control planes for advancing the rollout pool's +``latest`` pointer: + +- **trainer-as-writer** (slime_disagg / miles_disagg): the trainer's rank-0 + publish hook writes the Modal-Volume bulletin board directly + (:mod:`cookbook.bulletin_hooks`). +- **frontdoor-as-writer** (standalone_rollouts): the trainer POSTs a hot-load + signal to an external front door that owns the pointer + (:mod:`cookbook.standalone_rollouts.slime.hooks`). + +Those two planes are the real "disagg vs standalone" axis and stay distinct. But +the trainer-side *plumbing* around them was copy-forked: an identical torch rank +probe, an identical args-or-env settings reader, and identical session-affinity +header stamping. This module owns that plumbing once so the two hook stacks only +encode their genuine difference (who writes the pointer), not boilerplate. +""" + +from __future__ import annotations + +import os +from typing import Any + + +def distributed_rank() -> int | None: + """Return this process's torch-distributed rank, or ``None`` if torch + distributed isn't initialized (single-process / pre-init). + + Used to gate rank-0-only side effects (pointer writes, transport copies, + pool wakes) so only one writer acts per publish. + """ + try: + import torch.distributed as dist + + if dist.is_available() and dist.is_initialized(): + return int(dist.get_rank()) + except Exception: # noqa: BLE001 + return None + return None + + +def read_setting( + args: Any | None, + attr: str, + env: str, + *, + default: str | None = None, + required: bool = False, +) -> str: + """Read a string setting from the trainer ``args`` namespace, falling back to + an environment variable, then a default. + + The trainer's ``--custom-config-path`` setattr's every config key onto + ``args``; an env var is the deploy-time fallback for values that can't ride a + per-launch arg. Returns ``""`` for an absent, non-required setting. + """ + value = getattr(args, attr, None) if args is not None else None + if value is None: + value = os.environ.get(env) + if value is None: + value = default + if value is None and required: + raise RuntimeError( + f"Missing required setting {attr!r} or environment variable {env}" + ) + return str(value) if value is not None else "" + + +def apply_session_affinity(request: dict[str, Any], session_id: Any, header: str) -> None: + """Stamp ``header: session_id`` onto a rollout request's headers (idempotent). + + Sticky-session routing: a sample carrying a ``session_id`` should land on the + same replica across turns. ``setdefault`` so an explicit caller header wins. + """ + if not session_id: + return + headers = dict(request.get("headers") or {}) + headers.setdefault(header, str(session_id)) + request["headers"] = headers diff --git a/cookbook/rollout_control_test.py b/cookbook/rollout_control_test.py new file mode 100644 index 0000000..cc3a643 --- /dev/null +++ b/cookbook/rollout_control_test.py @@ -0,0 +1,59 @@ +from __future__ import annotations + +import unittest +from argparse import Namespace +from unittest import mock + +from cookbook import rollout_control + + +class DistributedRankTest(unittest.TestCase): + def test_none_when_torch_distributed_unavailable(self) -> None: + # No torch dist initialized in tests (torch may be absent entirely): the + # probe must degrade to None, which the hooks treat as "single writer". + self.assertIsNone(rollout_control.distributed_rank()) + + +class ReadSettingTest(unittest.TestCase): + def test_args_beats_env_beats_default(self) -> None: + args = Namespace(api_key="from-args") + with mock.patch.dict("os.environ", {"API_KEY": "from-env"}, clear=True): + self.assertEqual( + rollout_control.read_setting(args, "api_key", "API_KEY", default="d"), "from-args" + ) + with mock.patch.dict("os.environ", {"API_KEY": "from-env"}, clear=True): + self.assertEqual( + rollout_control.read_setting(Namespace(), "api_key", "API_KEY", default="d"), + "from-env", + ) + with mock.patch.dict("os.environ", {}, clear=True): + self.assertEqual( + rollout_control.read_setting(Namespace(), "api_key", "API_KEY", default="d"), "d" + ) + + def test_empty_string_for_absent_optional_and_raises_when_required(self) -> None: + with mock.patch.dict("os.environ", {}, clear=True): + self.assertEqual(rollout_control.read_setting(None, "x", "X"), "") + with self.assertRaises(RuntimeError): + rollout_control.read_setting(None, "x", "X", required=True) + + +class SessionAffinityTest(unittest.TestCase): + def test_stamps_header_without_overriding_explicit_value(self) -> None: + request: dict = {"headers": {"x-aff": "explicit"}} + rollout_control.apply_session_affinity(request, "sess-1", "x-aff") + self.assertEqual(request["headers"]["x-aff"], "explicit") # setdefault, not override + + request = {} + rollout_control.apply_session_affinity(request, "sess-1", "x-aff") + self.assertEqual(request["headers"]["x-aff"], "sess-1") + + def test_noop_for_falsy_session_id(self) -> None: + request: dict = {"headers": {"a": "b"}} + rollout_control.apply_session_affinity(request, None, "x-aff") + self.assertEqual(request["headers"], {"a": "b"}) + self.assertNotIn("x-aff", request["headers"]) + + +if __name__ == "__main__": + unittest.main() diff --git a/cookbook/serving.py b/cookbook/serving.py new file mode 100644 index 0000000..42b622c --- /dev/null +++ b/cookbook/serving.py @@ -0,0 +1,141 @@ +"""Shared B200 SGLang serving-image builder for the disagg rollout pool. + +The trainer half of each disagg example runs on its own slime/miles Megatron +image. The rollout half serves the model on a Blackwell SGLang build (fa4 / +cutlass-dsl prerelease kernels + the tokenspeed MLA attention backend) that +loads the served checkpoint's *own* quant config — there is no ``--quantization`` +flag baked in, so INT4 (slime/Kimi K2.6) vs NVFP4 (miles) is a property of the +served checkpoint, not of this image. The per-trainer ``serving.py`` modules are +thin wrappers that pass their trainer package + repo pin here. + +Two deliberate choices keep the image lean (identical for every trainer): + + * **SGLang comes from the modal-projects/sglang fork** pinned below — the same + build the standalone 4xB200 Kimi deployment uses (proven for NVFP4; the INT4 + MLA-MoE path is the one un-de-risked axis — verify on a warm container). + * **The trainer package is cloned ``--no-deps`` for one module.** The sidecar + only imports ``.utils.disk_delta`` (stdlib + numpy + zstandard; + xxhash/blake3 lazy), so Megatron is intentionally absent — the pool never + trains. Pin the SAME ref the trainer image uses so the host-side delta + decoder matches the trainer's delta encoder. +""" + +from __future__ import annotations + +from pathlib import Path + +import modal + +# Pinned Blackwell SGLang fork. Only the python sources are checked out over the +# prebuilt base image's kernels. +SGLANG_IMAGE_TAG = "lmsysorg/sglang:v0.5.12" +SGLANG_FORK_REPO = "https://github.com/modal-projects/sglang.git" +SGLANG_FORK_BRANCH = "timmy/dflash-fa4-fp8" +SGLANG_FORK_COMMIT = "dafb2b325b40298c5097564811463c585b7e9814" + +# SGLang runtime tunables carried over from the standalone B200 deployment. +SERVING_IMAGE_ENV = { + "HF_XET_HIGH_PERFORMANCE": "1", + "HF_HUB_ENABLE_HF_TRANSFER": "1", + "SGLANG_ALLOW_OVERWRITE_LONGER_CONTEXT_LEN": "1", + "SGLANG_DISABLE_CUDNN_CHECK": "1", + "SGLANG_ENABLE_OVERLAP_PLAN_STREAM": "1", + "SGLANG_TIMEOUT_KEEP_ALIVE": "300", +} + + +def build_b200_serving_image( + *, + trainer_repo_url: str, + trainer_repo_ref: str, + trainer_root: str, + cookbook_dir: Path, + hf_cache_path: str, + experiment: str, + shallow_clone: bool = True, + clear_sglang_cache_at_end: bool = False, +) -> modal.Image: + """Build the rollout-pool serving image (see module docstring). + + ``trainer_repo_url`` / ``trainer_repo_ref`` / ``trainer_root`` pin the + ``--no-deps`` trainer checkout (so the pool's ``disk_delta`` matches the + trainer's encoder). ``cookbook_dir`` is the per-trainer cookbook package dir; + its parent (the whole ``cookbook`` package) is mounted at ``/root/cookbook``, + so the sidecar subprocess can import both ``cookbook..sidecar`` and the + shared ``cookbook.sidecar`` spine it delegates to. Mounting the package (not + just the subdir) is required because the sidecar is launched as + ``python3 -m cookbook..sidecar`` and is never imported at deploy time, + so Modal's import-time automounting never sees the shared module. + + ``shallow_clone`` does a ``--depth 1`` clone+fetch (fine when the ref is a + branch tip or recent commit). ``clear_sglang_cache_at_end`` removes + ``/root/.cache/sglang`` as the final filesystem step, required when + ``modal_train`` mounts a kernel-cache volume there (a volume can't mount over + a non-empty path). + """ + depth = "--depth 1 " if shallow_clone else "" + fetch_depth = "--depth 1 " if shallow_clone else "" + image = ( + modal.Image.from_registry(SGLANG_IMAGE_TAG) + .run_commands( + f"cd /sgl-workspace/sglang && git remote add modal-fork {SGLANG_FORK_REPO}" + f" && git fetch modal-fork {SGLANG_FORK_BRANCH}" + f" && git checkout {SGLANG_FORK_COMMIT} -- python/", + ) + # Pre-release CUDA wheels (cutlass-dsl / sglang-kernel / flash-attn-4) — + # keep the deployment's known-good pip resolution. + .run_commands( + "pip install nvidia-cutlass-dsl==4.5.1 sglang-kernel==0.4.3 'flash-attn-4>=4.0.0b10'" + ) + # flash-attn-4 checks for the deprecated MmaFP8Op but cutlass-dsl 4.5.1 now + # generates MmaF8F6F4Op instead. Patch the isinstance check to handle both. + .run_commands( + "sed -i 's/isinstance(op, tcgen05.mma.MmaFP8Op)/isinstance(op, (tcgen05.mma.MmaFP8Op, tcgen05.mma.MmaF8F6F4Op))/' " + "/usr/local/lib/python3.12/dist-packages/flash_attn/cute/blackwell_helpers.py" + ) + # The base image bakes in an HF cache; remove it so it cannot shadow the + # cache volume mounted at the same path. + .run_commands(f"rm -rf {hf_cache_path}") + # trainer pkg --no-deps gives the sidecar `.utils.disk_delta` + # (host-side delta apply). Megatron is NOT installed — the pool never + # trains. Pin the SAME ref the trainer image uses so encoder == decoder. + .run_commands( + f"git clone {depth}{trainer_repo_url} {trainer_root}" + f" && cd {trainer_root}" + f" && git fetch {fetch_depth}origin {trainer_repo_ref}" + f" && git checkout FETCH_HEAD" + f" && python3 -m pip install --no-deps -e {trainer_root}" + ) + .pip_install( + "autoinference-utils==0.2.0", # SGLang server lifecycle for the rollout pool + "fastapi", # stitch sidecar + "httpx", # stitch sidecar + "uvicorn", # stitch sidecar + # disk_delta host-side apply: zstd decompress + xxhash (xxh3-128 + # default) / blake3 checksums. trainer pkg is installed --no-deps. + "zstandard", + "xxhash", + "blake3", + ) + ) + if clear_sglang_cache_at_end: + # MUST be the last filesystem step: modal_train mounts a kernel-cache + # volume at /root/.cache/sglang, and a volume can't mount over a non-empty + # path. The sglang-kernel/flashinfer and trainer installs above populate + # this dir, so clear it AFTER them (the volume repopulates the JIT/ + # autotuner cache on first boot). + image = image.run_commands("rm -rf /root/.cache/sglang") + return ( + image.env({"EXPERIMENT_CONFIG": experiment, **SERVING_IMAGE_ENV}) + # Mounted at container start (not copied into the image) so code edits to + # stitch / the sidecar never rebuild the image. Modal puts /root on + # PYTHONPATH for subprocesses (the sidecar). The whole cookbook package is + # mounted (not just the per-trainer subdir) so the sidecar can import the + # shared cookbook.sidecar spine it is a thin adapter over. + .add_local_python_source("stitch") + .add_local_dir( + cookbook_dir.parent, + remote_path="/root/cookbook", + ignore=["**/__pycache__"], + ) + ) diff --git a/cookbook/sidecar.py b/cookbook/sidecar.py new file mode 100644 index 0000000..72cda04 --- /dev/null +++ b/cookbook/sidecar.py @@ -0,0 +1,252 @@ +"""Shared SGLang weight-sync sidecar spine for the disagg cookbook trainers. + +The reusable versioned-proxy library lives in ``stitch.servers.sglang``; this +module wires the concrete bulletin-board + Modal Volume + SGLang-disk-delta +realization once, and the per-trainer ``cookbook//sidecar.py`` modules +are thin adapters that select the host-side delta decoder. + +The ONLY axis that varies between trainers is the host-side ``disk_delta`` +module the rollout pool ships ``--no-deps`` (``slime.utils.disk_delta`` vs +``miles.utils.disk_delta``). Their XOR/overwrite + zstd + xxh3/blake3 wire +formats are byte-identical, but the pool image installs exactly one of them, so +the adapter names which to import. Everything else — arg parsing, the bulletin +board + Volume reloader, the base-materialization strategy — lives here. +""" + +from __future__ import annotations + +import argparse +import concurrent.futures +import hashlib +import logging +import os +import shutil +from typing import Callable + +from stitch.bulletin import FilesystemBulletinBoard +from stitch.engines.sglang import SGLangDiskDeltaAdapter +from stitch.servers.sglang import create_app +from stitch.sync import CommitMode, WeightSyncManager + + +logger = logging.getLogger(__name__) + + +def parallel_init_local_checkpoint(disk_delta_module: str, workers: int = 32) -> Callable[[str, str], None]: + """Return an ``init_local_checkpoint(local, base)`` that copies the base + shards concurrently and skips the copy when the materialized base is current. + + ``disk_delta``'s default seeds the local checkpoint one shard at a time with + ``shutil.copy2`` — fine for a 16 B Moonlight base, far too slow for K2.6's + ~591 GB. We reuse the decoder's apply-lock + applied-version bookkeeping so + ``apply_deltas`` stays compatible, and fall back to the decoder's own copier + if those internals move under us. ``disk_delta_module`` is the trainer's + decoder (``slime.utils.disk_delta`` / ``miles.utils.disk_delta``); both expose + the same internals (miles is forked from slime), so one implementation serves + both. Imported lazily so the module imports without slime/miles installed.""" + import importlib + + def _base_fingerprint(base_dir: str) -> str: + """Identity of the base's bytes: (filename, size, mtime_ns) over every shard. Re-prep + rewrites the shards (new mtime/size), so this changes even when the tensor index/shapes + don't — which is exactly the case that bit K2.6 (NVFP4 values changed, names didn't).""" + h = hashlib.sha256() + for e in sorted(os.scandir(base_dir), key=lambda e: e.name): + if e.is_file(): + st = e.stat() + h.update(f"{e.name}:{st.st_size}:{st.st_mtime_ns}\n".encode()) + return h.hexdigest() + + def _init(local_ckpt_dir: str, base_dir: str) -> None: + dd = importlib.import_module(disk_delta_module) + try: + apply_lock = dd._apply_lock + read_version = dd._read_applied_version + write_version = dd._write_applied_version + drop_page_cache = dd.drop_page_cache + except AttributeError: # decoder internals changed — use its own copier + dd.init_local_checkpoint(local_ckpt_dir, base_dir) + return + + fp_path = os.path.join(local_ckpt_dir, ".base_fingerprint") + base_fp = _base_fingerprint(base_dir) + with apply_lock(local_ckpt_dir): + if read_version(local_ckpt_dir) is not None: + try: + cur = open(fp_path).read().strip() + except FileNotFoundError: + cur = None + if cur == base_fp: + return # current base already materialized + # STALE: /local-checkpoint holds a different (older) base — e.g. after a re-prep. + # Reusing it makes the host-side delta (diffed against the NEW base) XOR onto the + # WRONG bytes -> base_local XOR delta != export -> checksum mismatch on every tensor. + logger.warning( + "stale /local-checkpoint (base changed: %s != %s) — wiping + re-materializing", + cur, base_fp, + ) + shutil.rmtree(local_ckpt_dir, ignore_errors=True) + logger.info("Materializing base %s -> %s (%d copy workers)", base_dir, local_ckpt_dir, workers) + os.makedirs(local_ckpt_dir, exist_ok=True) + files = [e for e in os.scandir(base_dir) if e.is_file()] + + def _copy(entry: os.DirEntry) -> None: + shutil.copy2(entry.path, os.path.join(local_ckpt_dir, entry.name)) + drop_page_cache(entry.path) # don't evict the local copy we keep resident + + with concurrent.futures.ThreadPoolExecutor(max_workers=min(workers, max(1, len(files)))) as ex: + for _ in ex.map(_copy, files): # re-raises the first copy failure + pass + write_version(local_ckpt_dir, "000000") + with open(fp_path, "w") as f: + f.write(base_fp) + + return _init + + +def build_manager( + *, + upstream_url: str, + bulletin_root: str, + local_checkpoint_dir: str, + base_checkpoint_dir: str, + disk_delta_module: str, + inject_apply_deltas: bool, + base_copy_workers: int = 32, + volume_name: str = "", + run_id: str | None = None, + commit_mode: CommitMode = "in_place", + debug_requests: bool = False, +) -> WeightSyncManager: + """Build the reconciling :class:`WeightSyncManager` for one rollout replica. + + ``disk_delta_module`` is the trainer's host-side decoder. ``inject_apply_deltas`` + passes that module's ``apply_deltas`` to the engine explicitly — needed when the + decoder is not slime's default (miles), and harmless when it is. The base copy is + always the shared concurrent/stale-aware materializer. + """ + import importlib + + refresh = None + if volume_name: + from stitch.providers.modal import volume_reloader + + refresh = volume_reloader(volume_name) + # Publish-only writes the flat slime-native layout (weight_v{N}/ + + # model.safetensors.index.json + a raw `latest` pointer) to the transport; + # miles writes the same layout, so layout="slime" reads either unchanged. + board = FilesystemBulletinBoard(bulletin_root, refresh=refresh, layout="slime") + apply_deltas = None + if inject_apply_deltas: + # The decoder is installed --no-deps in the serving image for exactly this + # module (Megatron is absent — the pool never trains). Pinned to the same + # trainer ref as the encoder so encoder == decoder. + apply_deltas = importlib.import_module(disk_delta_module).apply_deltas + engine = SGLangDiskDeltaAdapter( + upstream_url=upstream_url, + local_checkpoint_dir=local_checkpoint_dir, + base_checkpoint_dir=base_checkpoint_dir, + apply_deltas=apply_deltas, + init_local_checkpoint=parallel_init_local_checkpoint(disk_delta_module, base_copy_workers), + ) + return WeightSyncManager( + board=board, + engine=engine, + run_id=run_id, + commit_mode=commit_mode, + debug_requests=debug_requests, + ) + + +def run_sidecar(*, disk_delta_module: str, inject_apply_deltas: bool) -> None: + """Parse args and run the sidecar uvicorn server. + + The per-trainer adapter calls this with its decoder module; ``run_sidecar`` + owns every other knob (host/port/transport/commit-mode/timeout/...). + """ + parser = argparse.ArgumentParser() + parser.add_argument("--host", default="0.0.0.0") + parser.add_argument("--port", type=int, default=8000) + parser.add_argument("--upstream-url", required=True) + parser.add_argument( + "--bulletin-root", + default=os.environ.get("DELTA_BULLETIN_ROOT", "/delta-bulletin"), + ) + parser.add_argument("--volume-name", default=os.environ.get("DELTA_VOLUME_NAME", "")) + parser.add_argument( + "--local-checkpoint-dir", + default=os.environ.get("STITCH_LOCAL_CHECKPOINT_DIR", "/local-checkpoint"), + help="Writable host-local full HF checkpoint patched in place by each delta.", + ) + parser.add_argument( + "--base-checkpoint-dir", + default=os.environ.get("STITCH_BASE_CHECKPOINT_DIR"), + help="Base HF checkpoint the local copy is seeded from (deltas build on it).", + ) + parser.add_argument("--run-id", default=os.environ.get("DISAGG_RUN_ID")) + parser.add_argument( + "--base-copy-workers", + type=int, + default=int(os.environ.get("SIDECAR_BASE_COPY_WORKERS", "32")), + help="Threads used to materialize the base checkpoint (concurrent shard copy).", + ) + parser.add_argument( + "--commit-mode", + choices=("quiesce", "in_place"), + default=os.environ.get("SIDECAR_COMMIT_MODE", "in_place"), + help=( + "in_place (default): pause/apply/continue without flushing; in-flight " + "requests keep decoding on stale KV and version isolation comes from " + "extra_key stamping. quiesce: wait out active requests and flush " + "before applying (safe on any build)." + ), + ) + parser.add_argument( + "--debug-requests", + action="store_true", + default=os.environ.get("SIDECAR_DEBUG_REQUESTS", "").lower() in {"1", "true", "yes"}, + help="Log every versioned sidecar proxy request at INFO level.", + ) + parser.add_argument( + "--upstream-timeout", + type=float, + default=float(os.environ.get("SIDECAR_UPSTREAM_TIMEOUT", "3600")), + help=( + "Seconds to wait for an upstream SGLang response before failing the " + "request (5xx) instead of holding it open forever. Must exceed the " + "slowest legitimate generation." + ), + ) + args = parser.parse_args() + if not args.base_checkpoint_dir: + raise SystemExit( + "--base-checkpoint-dir/STITCH_BASE_CHECKPOINT_DIR is required: deltas are" + " applied host-side on top of a copy of this base HF checkpoint." + ) + + logging.basicConfig(level=logging.INFO) + import uvicorn + + manager = build_manager( + upstream_url=args.upstream_url, + bulletin_root=args.bulletin_root, + local_checkpoint_dir=args.local_checkpoint_dir, + base_checkpoint_dir=args.base_checkpoint_dir, + disk_delta_module=disk_delta_module, + inject_apply_deltas=inject_apply_deltas, + base_copy_workers=args.base_copy_workers, + volume_name=args.volume_name, + run_id=args.run_id, + commit_mode=args.commit_mode, + debug_requests=args.debug_requests, + ) + uvicorn.run( + create_app( + manager, + upstream_url=args.upstream_url, + upstream_timeout=args.upstream_timeout, + ), + host=args.host, + port=args.port, + log_level="info", + ) diff --git a/cookbook/sidecar_test.py b/cookbook/sidecar_test.py new file mode 100644 index 0000000..edc8001 --- /dev/null +++ b/cookbook/sidecar_test.py @@ -0,0 +1,109 @@ +from __future__ import annotations + +import contextlib +import os +import sys +import tempfile +import types +import unittest +from pathlib import Path + +from cookbook import sidecar + + +def _install_fake_decoder(name: str, *, with_internals: bool) -> types.ModuleType: + """Register a fake ``disk_delta`` module in sys.modules. + + ``with_internals=True`` exposes the apply-lock + applied-version bookkeeping + the concurrent copier reuses; False omits them so the copier falls back to the + decoder's own ``init_local_checkpoint``. + """ + mod = types.ModuleType(name) + + def init_local_checkpoint(local: str, base: str) -> None: + os.makedirs(local, exist_ok=True) + Path(local, ".fallback_used").write_text("1") + + mod.init_local_checkpoint = init_local_checkpoint # type: ignore[attr-defined] + + if with_internals: + @contextlib.contextmanager + def _apply_lock(local: str): + yield + + def _read_applied_version(local: str) -> str | None: + p = Path(local, ".applied_version") + return p.read_text().strip() if p.exists() else None + + def _write_applied_version(local: str, version: str) -> None: + Path(local).mkdir(parents=True, exist_ok=True) + Path(local, ".applied_version").write_text(version) + + mod._apply_lock = _apply_lock # type: ignore[attr-defined] + mod._read_applied_version = _read_applied_version # type: ignore[attr-defined] + mod._write_applied_version = _write_applied_version # type: ignore[attr-defined] + mod.drop_page_cache = lambda path: None # type: ignore[attr-defined] + + sys.modules[name] = mod + return mod + + +class ParallelInitLocalCheckpointTest(unittest.TestCase): + def setUp(self) -> None: + self._tmp = tempfile.TemporaryDirectory() + self.root = Path(self._tmp.name) + self.base = self.root / "base" + self.local = self.root / "local" + self.base.mkdir() + (self.base / "shard-0.safetensors").write_bytes(b"AAAA") + (self.base / "shard-1.safetensors").write_bytes(b"BBBB") + self._names: list[str] = [] + + def tearDown(self) -> None: + for n in self._names: + sys.modules.pop(n, None) + self._tmp.cleanup() + + def _decoder(self, name: str, *, with_internals: bool = True) -> str: + self._names.append(name) + _install_fake_decoder(name, with_internals=with_internals) + return name + + def test_copies_all_base_shards_and_records_fingerprint(self) -> None: + init = sidecar.parallel_init_local_checkpoint(self._decoder("fake_dd_copy")) + init(str(self.local), str(self.base)) + self.assertEqual((self.local / "shard-0.safetensors").read_bytes(), b"AAAA") + self.assertEqual((self.local / "shard-1.safetensors").read_bytes(), b"BBBB") + self.assertTrue((self.local / ".base_fingerprint").exists()) + self.assertEqual((self.local / ".applied_version").read_text(), "000000") + + def test_skips_recopy_when_base_unchanged(self) -> None: + init = sidecar.parallel_init_local_checkpoint(self._decoder("fake_dd_skip")) + init(str(self.local), str(self.base)) + # A delta would have patched the local copy in place; an unnecessary + # re-materialize would clobber it, so prove the second call is a no-op. + (self.local / "shard-0.safetensors").write_bytes(b"PATCHED") + init(str(self.local), str(self.base)) + self.assertEqual((self.local / "shard-0.safetensors").read_bytes(), b"PATCHED") + + def test_wipes_and_rematerializes_when_base_changed(self) -> None: + init = sidecar.parallel_init_local_checkpoint(self._decoder("fake_dd_stale")) + init(str(self.local), str(self.base)) + (self.local / "stale-marker").write_text("x") + # Re-prep rewrites a base shard (new size/mtime) — fingerprint changes, + # so the stale local copy must be wiped and rebuilt against the new base. + (self.base / "shard-0.safetensors").write_bytes(b"CCCCCCCC") + init(str(self.local), str(self.base)) + self.assertEqual((self.local / "shard-0.safetensors").read_bytes(), b"CCCCCCCC") + self.assertFalse((self.local / "stale-marker").exists()) + + def test_falls_back_to_decoder_copier_when_internals_missing(self) -> None: + init = sidecar.parallel_init_local_checkpoint( + self._decoder("fake_dd_fallback", with_internals=False) + ) + init(str(self.local), str(self.base)) + self.assertTrue((self.local / ".fallback_used").exists()) + + +if __name__ == "__main__": + unittest.main() diff --git a/cookbook/slime_disagg/helpers.py b/cookbook/slime_disagg/helpers.py index 6295ad7..031afe4 100644 --- a/cookbook/slime_disagg/helpers.py +++ b/cookbook/slime_disagg/helpers.py @@ -1,23 +1,19 @@ """Trainer-specific helpers for the slime_disagg example. -Ray cluster, sidecar, and process helpers are shared across trainers via -:mod:`cookbook.ray_cluster` and :mod:`cookbook.sidecar_process`. This module -provides slime-specific config preparation, train command building, and the -Flash pool smoke check. +Thin wrappers over the shared launch spine: Ray-cluster/sidecar/process helpers +come from :mod:`cookbook.ray_cluster` / :mod:`cookbook.sidecar_process`, and +config-prep / train-command / smoke-check come from :mod:`cookbook.trainer_helpers`. +This module only supplies the slime-specific axes: the sidecar module path, the +config-field tuple to materialize, the model-script attribute, and that the +rollout pool runs with a warm floor (not scale-from-zero). """ from __future__ import annotations -import json -import os -import shlex import subprocess -import time -import urllib.error -import urllib.request from typing import Any -from stitch.providers.modal import discover_flash_targets, resolve_flash_gateway_url +from cookbook.slime_disagg.configs.base import YAML_CONFIG_FIELDS # Re-export shared helpers so existing callers (modal_train.py) don't break. from cookbook.ray_cluster import ( # noqa: F401 @@ -32,6 +28,12 @@ terminate_process, wait_http, ) +from cookbook.trainer_helpers import ( # noqa: F401 + VersionAheadError, + build_train_cmd as _build_train_cmd, + prepare_config, + smoke_flash_pool as _smoke_flash_pool, +) SIDECAR_MODULE = "cookbook.slime_disagg.sidecar" @@ -63,45 +65,14 @@ def start_sglang_sidecar( ) -# ── SLIME launch ────────────────────────────────────────────────────────────── - - def prepare_slime_config(slime_cfg: Any, tmpdir: str) -> None: """Resolve HF repo IDs to local paths and materialize inline YAML configs.""" - from huggingface_hub import snapshot_download - import yaml - - from cookbook.slime_disagg.configs.base import YAML_CONFIG_FIELDS - - for attr in ("hf_checkpoint", "load", "ref_load", "critic_load"): - if (val := getattr(slime_cfg, attr, None)) and not str(val).startswith("/"): - setattr(slime_cfg, attr, snapshot_download(val, local_files_only=True)) - - for field in YAML_CONFIG_FIELDS: - if isinstance(val := getattr(slime_cfg, field, None), dict): - path = os.path.join(tmpdir, f"{field}.yaml") - with open(path, "w") as f: - yaml.dump(val, f) - setattr(slime_cfg, field, path) + prepare_config(slime_cfg, tmpdir, YAML_CONFIG_FIELDS) def build_train_cmd(slime_cfg: Any, slime_root: str) -> str: - """Build the training command, sourcing model arch args if needed.""" - train_script = f"{slime_root}/{'train_async.py' if slime_cfg.async_mode else 'train.py'}" - if slime_cfg.slime_model_script: - inner = ( - f"source {slime_root}/{slime_cfg.slime_model_script} && " - f"python3 {train_script} ${{MODEL_ARGS[@]}} {shlex.join(slime_cfg.cli_args())}" - ) - return f"bash -c {shlex.quote(inner)}" - return f"python3 {train_script} {shlex.join(slime_cfg.cli_args())}" - - -# ── Flash pool smoke check ──────────────────────────────────────────────────── - - -class VersionAheadError(RuntimeError): - """Raised when a monotonic rollout pool has already advanced past a smoke version.""" + """Build the training command, sourcing slime's model arch args if needed.""" + return _build_train_cmd(slime_cfg, slime_root, model_script_attr="slime_model_script") def smoke_flash_pool( @@ -113,68 +84,13 @@ def smoke_flash_pool( expect_min_containers: int, timeout_seconds: int, ) -> None: - """Poll the Flash gateway and direct container URLs until the pool serves - completions at the expected weight version.""" - deadline = time.time() + timeout_seconds - last_error: str | None = None - while True: - gateway = resolve_flash_gateway_url(app_name, cls_name) - targets = discover_flash_targets(app_name, cls_name) - if len(targets) < expect_min_containers: - last_error = f"expected at least {expect_min_containers} containers, found {len(targets)}: {targets}" - else: - try: - _check_flash_pool_once(gateway, targets, model_name, weight_version) - return - except VersionAheadError: - raise - except Exception as exc: # noqa: BLE001 - last_error = f"{type(exc).__name__}: {exc}" - if time.time() >= deadline: - raise TimeoutError(f"Flash pool smoke did not pass before timeout: {last_error}") - print(f"Waiting for Flash pool readiness: {last_error}") - time.sleep(10) - - -def _check_flash_pool_once(gateway: str, targets: list[str], model_name: str, expected: int) -> None: - print(f"Gateway URL: {gateway}") - print(f"Direct container URLs ({len(targets)}):") - for target in targets: - print(f" {target}") - - for target in [gateway, *targets]: - info = _get_json(f"{target}/server_info", timeout=30) - print(f"{target} server_info={info}") - current = int(info["current_version"]) - if current > expected: - raise VersionAheadError(f"{target} current_version={current} already passed expected {expected}") - if current != expected: - raise RuntimeError(f"{target} current_version={current} expected {expected}") - - payload = { - "model": model_name, - "messages": [{"role": "user", "content": "Reply with exactly OK."}], - "max_tokens": 8, - "temperature": 0, - "weight_version": {"exact_version": expected}, - "chat_template_kwargs": {"enable_thinking": False}, - } - data = _post_json(f"{gateway}/v1/chat/completions", payload, timeout=180) - print(f"Gateway completion: {data}") - if int(data.get("weight_version_start", -1)) != expected or int(data.get("weight_version_end", -1)) != expected: - raise RuntimeError(f"unexpected gateway weight metadata: {data}") - - -def _get_json(url: str, *, timeout: float) -> dict: - with urllib.request.urlopen(url, timeout=timeout) as resp: - return json.load(resp) - - -def _post_json(url: str, payload: dict, *, timeout: float) -> dict: - request = urllib.request.Request( - url, - data=json.dumps(payload).encode(), - headers={"Content-Type": "application/json"}, + """Smoke the warm-floor slime rollout pool (min_containers > 0).""" + _smoke_flash_pool( + app_name=app_name, + cls_name=cls_name, + model_name=model_name, + weight_version=weight_version, + expect_min_containers=expect_min_containers, + timeout_seconds=timeout_seconds, + wake_on_demand=False, ) - with urllib.request.urlopen(request, timeout=timeout) as resp: - return json.load(resp) diff --git a/cookbook/slime_disagg/modal_train.py b/cookbook/slime_disagg/modal_train.py index 6ecd046..da5f95c 100644 --- a/cookbook/slime_disagg/modal_train.py +++ b/cookbook/slime_disagg/modal_train.py @@ -113,11 +113,14 @@ # Local source is mounted when containers start rather than copied into the # image, so code changes never trigger an image rebuild. Modal puts /root # on PYTHONPATH, which also makes both packages importable from - # subprocesses (the sidecar, Ray workers). + # subprocesses (the sidecar, Ray workers). The whole cookbook package is + # mounted (not just the per-trainer subdir) so the trainer and the + # `python3 -m cookbook.slime_disagg.sidecar` subprocess can import the shared + # cookbook spine (helpers/hooks/sidecar) the thin adapters delegate to. .add_local_python_source("stitch") .add_local_dir( - Path(__file__).parent, - remote_path="/root/cookbook/slime_disagg", + Path(__file__).parent.parent, + remote_path="/root/cookbook", ignore=["**/__pycache__"], ) ) diff --git a/cookbook/slime_disagg/serving.py b/cookbook/slime_disagg/serving.py index 4607c94..d1956e4 100644 --- a/cookbook/slime_disagg/serving.py +++ b/cookbook/slime_disagg/serving.py @@ -1,26 +1,10 @@ -"""Dedicated B200 native-INT4 SGLang serving image for the rollout pool. +"""Dedicated B200 native-INT4 SGLang serving image for the slime rollout pool. -The trainer half of this example runs on the slime/Megatron image -(``modal_train.image``). The rollout half that serves Kimi K2.6 needs a -different stack: a Blackwell SGLang build that loads the model's native -compressed-tensors INT4 (W4A16) checkpoint and runs MLA attention with the -tuned hierarchical KV cache. This module builds that image. - -Two deliberate choices keep it lean: - - * **SGLang comes from the modal-projects/sglang fork** (Blackwell fa4 / - cutlass-dsl prerelease kernels + the tokenspeed MLA attention backend) — the - same build the standalone 4xB200 Kimi deployment uses. The FP4-specific - ``--quantization modelopt_fp4`` is *not* baked in here; the served - checkpoint's own ``compressed-tensors`` config drives INT4 weight loading - (see the config module's ``SGLANG_SERVER_ARGS``). - * **slime is cloned ``--no-deps`` for one module.** The sidecar only imports - ``slime.utils.disk_delta`` (stdlib + numpy + zstandard; xxhash/blake3 lazy), - so Megatron is intentionally absent from the rollout pool. - -The single un-de-risked axis is whether this fork serves native-INT4 MLA MoE on -Blackwell as cleanly as it serves NVFP4 (it is proven for NVFP4). Verify on a -warm container before a long run; the FP4 path is the known-good fallback. +A thin wrapper over :func:`cookbook.serving.build_b200_serving_image`. The image +itself is trainer-agnostic (see that module): INT4 vs NVFP4 is driven by the +served checkpoint's own ``compressed-tensors`` config, not by this builder. This +wrapper only pins slime as the ``--no-deps`` decoder package and does a shallow +clone (the slime ref is a branch-tip commit). """ from __future__ import annotations @@ -29,22 +13,7 @@ import modal -# Pinned Blackwell SGLang fork — matches the standalone 4xB200 Kimi serve recipe. -# Only the python sources are checked out over the prebuilt base image's kernels. -SGLANG_IMAGE_TAG = "lmsysorg/sglang:v0.5.12" -SGLANG_FORK_REPO = "https://github.com/modal-projects/sglang.git" -SGLANG_FORK_BRANCH = "timmy/dflash-fa4-fp8" -SGLANG_FORK_COMMIT = "dafb2b325b40298c5097564811463c585b7e9814" - -# SGLang runtime tunables carried over from the standalone B200 deployment. -SERVING_IMAGE_ENV = { - "HF_XET_HIGH_PERFORMANCE": "1", - "HF_HUB_ENABLE_HF_TRANSFER": "1", - "SGLANG_ALLOW_OVERWRITE_LONGER_CONTEXT_LEN": "1", - "SGLANG_DISABLE_CUDNN_CHECK": "1", - "SGLANG_ENABLE_OVERLAP_PLAN_STREAM": "1", - "SGLANG_TIMEOUT_KEEP_ALIVE": "300", -} +from cookbook.serving import build_b200_serving_image def build_int4_b200_serving_image( @@ -55,66 +24,12 @@ def build_int4_b200_serving_image( hf_cache_path: str, experiment: str, ) -> modal.Image: - """Build the rollout-pool serving image (see module docstring). - - The slime fork ref / root and the HF cache path are passed in by - ``modal_train`` so the pool and the trainer pin the identical slime commit - (the sidecar's ``disk_delta`` must match the trainer's delta encoder). - """ - # serving.py lives in cookbook/slime_disagg, mounted to /root/cookbook/slime_disagg - # exactly as the trainer image mounts it, so `cookbook.slime_disagg.sidecar` - # imports identically in either container. - slime_disagg_dir = Path(__file__).parent - return ( - modal.Image.from_registry(SGLANG_IMAGE_TAG) - .run_commands( - f"cd /sgl-workspace/sglang && git remote add modal-fork {SGLANG_FORK_REPO}" - f" && git fetch modal-fork {SGLANG_FORK_BRANCH}" - f" && git checkout {SGLANG_FORK_COMMIT} -- python/", - ) - # Pre-release CUDA wheels (cutlass-dsl / sglang-kernel / flash-attn-4) — - # keep the deployment's known-good pip resolution. - .run_commands( - "pip install nvidia-cutlass-dsl==4.5.1 sglang-kernel==0.4.3 'flash-attn-4>=4.0.0b10'" - ) - # flash-attn-4 checks for the deprecated MmaFP8Op but cutlass-dsl 4.5.1 now - # generates MmaF8F6F4Op instead. Patch the isinstance check to handle both. - .run_commands( - "sed -i 's/isinstance(op, tcgen05.mma.MmaFP8Op)/isinstance(op, (tcgen05.mma.MmaFP8Op, tcgen05.mma.MmaF8F6F4Op))/' " - "/usr/local/lib/python3.12/dist-packages/flash_attn/cute/blackwell_helpers.py" - ) - # The base image bakes in an HF cache; remove it so it cannot shadow the - # cache volume mounted at the same path. - .run_commands(f"rm -rf {hf_cache_path}") - # slime --no-deps gives the sidecar `slime.utils.disk_delta` (host-side - # delta apply). Megatron is NOT installed — the pool never trains. Pin the - # SAME ref the trainer image uses so the delta encoder/decoder match. - .run_commands( - f"git clone --depth 1 {slime_repo_url} {slime_root}" - f" && cd {slime_root}" - f" && git fetch --depth 1 origin {slime_repo_ref}" - f" && git checkout FETCH_HEAD" - f" && python3 -m pip install --no-deps -e {slime_root}" - ) - .pip_install( - "autoinference-utils==0.2.0", # SGLang server lifecycle for the rollout pool - "fastapi", # stitch sidecar - "httpx", # stitch sidecar - "uvicorn", # stitch sidecar - # disk_delta host-side apply: zstd decompress + xxhash (xxh3-128 - # default) / blake3 checksums. slime is installed --no-deps. - "zstandard", - "xxhash", - "blake3", - ) - .env({"EXPERIMENT_CONFIG": experiment, **SERVING_IMAGE_ENV}) - # Mounted at container start (not copied into the image) so code edits to - # stitch / the sidecar never rebuild the image. Modal puts /root on - # PYTHONPATH for subprocesses (the sidecar). - .add_local_python_source("stitch") - .add_local_dir( - slime_disagg_dir, - remote_path="/root/cookbook/slime_disagg", - ignore=["**/__pycache__"], - ) + return build_b200_serving_image( + trainer_repo_url=slime_repo_url, + trainer_repo_ref=slime_repo_ref, + trainer_root=slime_root, + cookbook_dir=Path(__file__).parent, + hf_cache_path=hf_cache_path, + experiment=experiment, + shallow_clone=True, ) diff --git a/cookbook/slime_disagg/sidecar.py b/cookbook/slime_disagg/sidecar.py index ba62fb1..2e21ed1 100644 --- a/cookbook/slime_disagg/sidecar.py +++ b/cookbook/slime_disagg/sidecar.py @@ -1,137 +1,21 @@ -"""SGLang weight-sync sidecar launcher for the slime_disagg example. +"""SGLang weight-sync sidecar entry point for the slime_disagg example. -The reusable versioned-proxy library lives in ``stitch.servers.sglang``; this -module wires the concrete bulletin-board + Modal Volume + SGLang-disk-delta -realization and runs it as a process (``helpers.start_sglang_sidecar`` launches -``python3 -m cookbook.slime_disagg.sidecar``). +A thin adapter over :mod:`cookbook.sidecar`: the shared spine owns every knob; +this only names slime's host-side delta decoder. ``helpers.start_sglang_sidecar`` +launches it via ``python3 -m cookbook.slime_disagg.sidecar``. """ from __future__ import annotations -import argparse -import logging -import os +from cookbook.sidecar import run_sidecar -from stitch.bulletin import FilesystemBulletinBoard -from stitch.engines.sglang import SGLangDiskDeltaAdapter -from stitch.servers.sglang import create_app -from stitch.sync import CommitMode, WeightSyncManager - -def build_manager( - *, - upstream_url: str, - bulletin_root: str, - local_checkpoint_dir: str, - base_checkpoint_dir: str, - volume_name: str = "", - run_id: str | None = None, - commit_mode: CommitMode = "in_place", - debug_requests: bool = False, -) -> WeightSyncManager: - refresh = None - if volume_name: - from stitch.providers.modal import volume_reloader - - refresh = volume_reloader(volume_name) - # slime publish-only writes the flat slime-native layout (weight_v{N}/ + - # model.safetensors.index.json + a raw `latest` pointer) to the Volume. - board = FilesystemBulletinBoard(bulletin_root, refresh=refresh, layout="slime") - engine = SGLangDiskDeltaAdapter( - upstream_url=upstream_url, - local_checkpoint_dir=local_checkpoint_dir, - base_checkpoint_dir=base_checkpoint_dir, - ) - return WeightSyncManager( - board=board, - engine=engine, - run_id=run_id, - commit_mode=commit_mode, - debug_requests=debug_requests, - ) +# slime's decoder is the engine's lazy default, so it need not be injected. +DISK_DELTA_MODULE = "slime.utils.disk_delta" def main() -> None: - parser = argparse.ArgumentParser() - parser.add_argument("--host", default="0.0.0.0") - parser.add_argument("--port", type=int, default=8000) - parser.add_argument("--upstream-url", required=True) - parser.add_argument( - "--bulletin-root", - default=os.environ.get("DELTA_BULLETIN_ROOT", "/delta-bulletin"), - ) - parser.add_argument( - "--volume-name", default=os.environ.get("DELTA_VOLUME_NAME", "") - ) - parser.add_argument( - "--local-checkpoint-dir", - default=os.environ.get("STITCH_LOCAL_CHECKPOINT_DIR", "/local-checkpoint"), - help="Writable host-local full HF checkpoint patched in place by each delta.", - ) - parser.add_argument( - "--base-checkpoint-dir", - default=os.environ.get("STITCH_BASE_CHECKPOINT_DIR"), - help="Base HF checkpoint the local copy is seeded from (deltas build on it).", - ) - parser.add_argument("--run-id", default=os.environ.get("DISAGG_RUN_ID")) - parser.add_argument( - "--commit-mode", - choices=("quiesce", "in_place"), - default=os.environ.get("SIDECAR_COMMIT_MODE", "in_place"), - help=( - "in_place (default): pause/apply/continue without flushing; " - "in-flight requests keep decoding on stale KV and version isolation " - "comes from extra_key stamping. quiesce: wait out active requests " - "and flush before applying." - ), - ) - parser.add_argument( - "--debug-requests", - action="store_true", - default=os.environ.get("SIDECAR_DEBUG_REQUESTS", "").lower() - in {"1", "true", "yes"}, - help="Log every versioned sidecar proxy request at INFO level.", - ) - parser.add_argument( - "--upstream-timeout", - type=float, - default=float(os.environ.get("SIDECAR_UPSTREAM_TIMEOUT", "3600")), - help=( - "Seconds to wait for an upstream SGLang response before failing the " - "request (5xx) instead of holding it open forever. Must exceed the " - "slowest legitimate generation." - ), - ) - args = parser.parse_args() - if not args.base_checkpoint_dir: - raise SystemExit( - "--base-checkpoint-dir/STITCH_BASE_CHECKPOINT_DIR is required: deltas are" - " applied host-side on top of a copy of this base HF checkpoint." - ) - - logging.basicConfig(level=logging.INFO) - import uvicorn - - manager = build_manager( - upstream_url=args.upstream_url, - bulletin_root=args.bulletin_root, - local_checkpoint_dir=args.local_checkpoint_dir, - base_checkpoint_dir=args.base_checkpoint_dir, - volume_name=args.volume_name, - run_id=args.run_id, - commit_mode=args.commit_mode, - debug_requests=args.debug_requests, - ) - uvicorn.run( - create_app( - manager, - upstream_url=args.upstream_url, - upstream_timeout=args.upstream_timeout, - ), - host=args.host, - port=args.port, - log_level="info", - ) + run_sidecar(disk_delta_module=DISK_DELTA_MODULE, inject_apply_deltas=False) if __name__ == "__main__": diff --git a/cookbook/standalone_rollouts/slime/hooks.py b/cookbook/standalone_rollouts/slime/hooks.py index 8001312..f6b484a 100644 --- a/cookbook/standalone_rollouts/slime/hooks.py +++ b/cookbook/standalone_rollouts/slime/hooks.py @@ -24,6 +24,9 @@ from pathlib import Path from typing import Any +from cookbook.rollout_control import apply_session_affinity +from cookbook.rollout_control import distributed_rank as _distributed_rank +from cookbook.rollout_control import read_setting as _setting from stitch.protocol import RolloutPoolState, parse_weight_identity, weight_identity @@ -198,6 +201,37 @@ def announce_and_wait(args: Any, version_dir: str, rollout_engines: list[Any]) - ) +def announce_claim(args: Any, *, run_id: str) -> None: + """Trainer launch hook (rank 0): claim the rollout pool for this run via the + front door, the standalone counterpart to :func:`cookbook.bulletin_hooks.claim_pool`. + + POST the empty pointer (``weight_v000000``) under a fresh ``run_id`` so the + front door's cross-run :func:`decide_pointer_move` resets ``latest`` to base + *before* the first delta — every replica (cold, or warm on a finished run) + reconciles to base up front instead of inferring the reset from the first + publish's run mismatch. Best-effort and non-blocking: the pool may be scaled + to zero at launch, and the first :func:`announce_and_wait` already blocks on + readiness, so a transient claim failure costs only the early reset, not + correctness. ``run_id`` must be fresh per launch (the epoch/fence token). + """ + if _distributed_rank() not in (None, 0): + return + cfg = replace(ShimConfig.from_env(args), run_id=run_id) + try: + _post_hot_load( + cfg, + identity=weight_identity(0), + previous_identity=cfg.base_snapshot_identity, + ) + except Exception: # noqa: BLE001 + logger.warning( + "Best-effort pool claim failed for run %r; the first publish's " + "cross-run reset will still establish base", + run_id, + exc_info=True, + ) + + def rollout_request_weight_version_hook( args: Any, sample: Any, request: dict[str, Any] ) -> None: @@ -270,17 +304,16 @@ def rollout_request_weight_version_hook( headers["Provider-Model"] = cfg.provider_model if cfg.provider_deployment: headers["Provider-Deployment"] = cfg.provider_deployment - if getattr(sample, "session_id", None): - # Neutral by default. The front-door relabel proxy maps this to - # Modal-Session-ID before the gateway. - affinity_header = _setting( - args, - "api_shim_session_affinity_header", - "STITCH_SHIM_SESSION_AFFINITY_HEADER", - default="x-session-affinity", - ) - headers.setdefault(affinity_header, sample.session_id) request["headers"] = headers + # Neutral by default. The front-door relabel proxy maps this to + # Modal-Session-ID before the gateway. + affinity_header = _setting( + args, + "api_shim_session_affinity_header", + "STITCH_SHIM_SESSION_AFFINITY_HEADER", + default="x-session-affinity", + ) + apply_session_affinity(request, getattr(sample, "session_id", None), affinity_header) def _rollout_request_target_version(args: Any, rollout_id: int, evaluation: bool) -> int: @@ -413,32 +446,4 @@ def _copy_version_to_transport(version_dir: Path, destination: Path) -> None: shutil.copyfileobj(src, dst) -def _distributed_rank() -> int | None: - try: - import torch.distributed as dist - - if dist.is_available() and dist.is_initialized(): - return int(dist.get_rank()) - except Exception: # noqa: BLE001 - return None - return None - -def _setting( - args: Any | None, - attr: str, - env: str, - *, - default: str | None = None, - required: bool = False, -) -> str: - value = getattr(args, attr, None) if args is not None else None - if value is None: - value = os.environ.get(env) - if value is None: - value = default - if value is None and required: - raise RuntimeError( - f"Missing required setting {attr!r} or environment variable {env}" - ) - return str(value) if value is not None else "" diff --git a/cookbook/standalone_rollouts/slime/hooks_test.py b/cookbook/standalone_rollouts/slime/hooks_test.py index b280a45..e15841b 100644 --- a/cookbook/standalone_rollouts/slime/hooks_test.py +++ b/cookbook/standalone_rollouts/slime/hooks_test.py @@ -89,6 +89,46 @@ def test_skips_baseline_non_version_dir(self) -> None: self.assertEqual(posted, []) +class AnnounceClaimTest(unittest.TestCase): + def test_posts_base_pointer_under_fresh_run_on_rank_zero(self) -> None: + # Claim-at-launch parity with bulletin_hooks.claim_pool: POST the empty + # pointer (weight_v000000) under the fresh run id so the front door resets + # latest to base before the first delta. + posted: list[tuple[str | None, str, str]] = [] + + def fake_post(cfg, *, identity, previous_identity) -> None: + posted.append((cfg.run_id, identity, previous_identity)) + + args = Namespace(api_shim_base_url="http://provider") + with mock.patch.object(hooks, "_distributed_rank", return_value=0), mock.patch.object( + hooks, "_post_hot_load", fake_post + ): + hooks.announce_claim(args, run_id="run-a") + self.assertEqual(posted, [("run-a", "weight_v000000", "base")]) + + def test_is_noop_off_rank_zero(self) -> None: + posted: list[int] = [] + args = Namespace(api_shim_base_url="http://provider") + with mock.patch.object(hooks, "_distributed_rank", return_value=1), mock.patch.object( + hooks, "_post_hot_load", lambda *a, **k: posted.append(1) + ): + hooks.announce_claim(args, run_id="run-a") + self.assertEqual(posted, []) + + def test_swallows_claim_failure(self) -> None: + # Best-effort: a transient claim POST failure must not kill the launch + # (the first publish's cross-run reset still establishes base). + args = Namespace(api_shim_base_url="http://provider") + + def boom(*a, **k) -> None: + raise RuntimeError("front door down") + + with mock.patch.object(hooks, "_distributed_rank", return_value=0), mock.patch.object( + hooks, "_post_hot_load", boom + ): + hooks.announce_claim(args, run_id="run-a") # must not raise + + class RolloutRequestHookTest(unittest.TestCase): def test_skips_pin_without_rollout_id_but_sets_affinity(self) -> None: # PR #5's request carries no rollout_id: the hook must not crash, must diff --git a/cookbook/standalone_rollouts/slime/modal_train.py b/cookbook/standalone_rollouts/slime/modal_train.py index 673410e..7e97744 100644 --- a/cookbook/standalone_rollouts/slime/modal_train.py +++ b/cookbook/standalone_rollouts/slime/modal_train.py @@ -18,6 +18,7 @@ import modal.experimental from cookbook.standalone_rollouts import modal_serve as provider_app +from cookbook.standalone_rollouts.slime import hooks from cookbook.slime_disagg import helpers from cookbook.slime_disagg.configs.base import ( CHECKPOINTS_PATH, @@ -199,6 +200,11 @@ def train(self, experiment: str, payload: dict) -> None: cfg.custom_config_path = _custom_config( cfg, rollout_num_engines=getattr(run, "ROLLOUT_NUM_ENGINES", 1) ) + # Claim the pool for this fresh run before any delta: reset `latest` to + # base via the front door so replicas reconcile to base up front instead + # of inferring the reset from the first publish (the explicit-claim model + # PR #5 established for the bulletin path, now mirrored here). + hooks.announce_claim(cfg, run_id=run_id) helpers.prepare_slime_config(cfg, tempfile.mkdtemp()) cmd = helpers.build_train_cmd(cfg, SLIME_ROOT) diff --git a/cookbook/trainer_helpers.py b/cookbook/trainer_helpers.py new file mode 100644 index 0000000..68468f1 --- /dev/null +++ b/cookbook/trainer_helpers.py @@ -0,0 +1,271 @@ +"""Shared trainer-side launch helpers for the disagg cookbook trainers. + +slime and miles drive the same launch spine: resolve HF repo ids + materialize +inline YAML configs, build the ``train.py`` command (optionally sourcing a model +arch script), monitor host RAM, and smoke the deployed Flash rollout pool. The +per-trainer ``helpers.py`` modules are thin wrappers that supply the only axes +that actually differ: which config-field tuple to materialize, the model-script +attribute name, and whether the rollout pool has a warm floor or scales from zero. +""" + +from __future__ import annotations + +import json +import os +import shlex +import socket +import threading +import time +import urllib.error +import urllib.request +from typing import Any, Iterable + +from stitch.providers.modal import discover_flash_targets, resolve_flash_gateway_url + + +# ── Config preparation ──────────────────────────────────────────────────────── + + +def prepare_config(cfg: Any, tmpdir: str, yaml_config_fields: Iterable[str]) -> None: + """Resolve HF repo IDs to local paths and materialize inline YAML configs. + + Repo-id-shaped checkpoint fields are snapshot-downloaded from the HF cache; + absolute paths (already-prepared checkpoints) are left untouched. Inline dict + configs in ``yaml_config_fields`` are written to ``tmpdir`` and the field is + repointed at the file the trainer reads. + """ + from huggingface_hub import snapshot_download + import yaml + + for attr in ("hf_checkpoint", "load", "ref_load", "critic_load"): + if (val := getattr(cfg, attr, None)) and not str(val).startswith("/"): + setattr(cfg, attr, snapshot_download(val, local_files_only=True)) + + for field in yaml_config_fields: + if isinstance(val := getattr(cfg, field, None), dict): + path = os.path.join(tmpdir, f"{field}.yaml") + with open(path, "w") as f: + yaml.dump(val, f) + setattr(cfg, field, path) + + +def materialize_node_local_yaml(cfg: Any, field: str, dest_dir: str = "/root/.miles_node_yaml") -> None: + """Materialize a per-actor-read YAML config to a deterministic node-local path. + + Some config files (notably miles' ``te_precision_config_file``, which + ``load_quantization_recipe`` re-reads on every Ray actor during model build) + are read independently on each trainer node — not just parsed once on the head. + ``prepare_config`` writes them under ``tempfile.mkdtemp()`` on the head only, + so on a multi-node cluster the other containers can't see that path. + + Call this on EVERY node (SPMD train()), before the rank-0 gate: each node + writes identical content (from the shared payload) to the same fixed path, so + the path the head embeds in the args resolves locally on all actors. No volume + commit/reload race — Ray actors are long-lived and wouldn't see post-start + volume writes anyway. + """ + import yaml + + if isinstance(val := getattr(cfg, field, None), dict): + os.makedirs(dest_dir, exist_ok=True) + path = os.path.join(dest_dir, f"{field}.yaml") + with open(path, "w") as f: + yaml.dump(val, f) + setattr(cfg, field, path) + + +def build_train_cmd(cfg: Any, trainer_root: str, *, model_script_attr: str) -> str: + """Build the training command, sourcing model arch args if needed. + + slime/miles ``train.py`` / ``train_async.py`` live at the repo root and consume + the ``MODEL_ARGS`` bash array defined by the sourced model script. + ``model_script_attr`` is the config attribute naming that script + (``slime_model_script`` / ``miles_model_script``). + """ + train_script = f"{trainer_root}/{'train_async.py' if cfg.async_mode else 'train.py'}" + model_script = getattr(cfg, model_script_attr) + if model_script: + inner = ( + f"source {trainer_root}/{model_script} && " + f"python3 {train_script} ${{MODEL_ARGS[@]}} {shlex.join(cfg.cli_args())}" + ) + return f"bash -c {shlex.quote(inner)}" + return f"python3 {train_script} {shlex.join(cfg.cli_args())}" + + +# ── Flash pool smoke check ──────────────────────────────────────────────────── + + +class VersionAheadError(RuntimeError): + """Raised when a monotonic rollout pool has already advanced past a smoke version.""" + + +def smoke_flash_pool( + *, + app_name: str, + cls_name: str, + model_name: str, + weight_version: int, + expect_min_containers: int, + timeout_seconds: int, + wake_on_demand: bool, +) -> None: + """Poll the Flash gateway until the pool serves completions at the expected + weight version, via the gateway and each container directly. + + ``wake_on_demand`` chooses the readiness model. False (warm floor): require + ``expect_min_containers`` live containers, then check each, then complete via + the gateway. True (scale-from-zero): a completion sent to the gateway is what + scales the pool 0->1 (Flash holds the request through the cold start, so the + timeout must be generous); ``expect_min_containers`` is advisory and the + direct-container check confirms the warmed pool afterward. + """ + deadline = time.time() + timeout_seconds + last_error: str | None = None + while True: + try: + if wake_on_demand: + _smoke_wake_on_demand(app_name, cls_name, model_name, weight_version) + else: + _smoke_warm_floor(app_name, cls_name, model_name, weight_version, expect_min_containers) + return + except VersionAheadError: + raise + except Exception as exc: # noqa: BLE001 + last_error = f"{type(exc).__name__}: {exc}" + if time.time() >= deadline: + raise TimeoutError(f"Flash pool smoke did not pass before timeout: {last_error}") + print(f"Waiting for Flash pool readiness: {last_error}") + time.sleep(10) + + +def _smoke_warm_floor( + app_name: str, cls_name: str, model_name: str, expected: int, expect_min_containers: int +) -> None: + gateway = resolve_flash_gateway_url(app_name, cls_name) + targets = discover_flash_targets(app_name, cls_name) + if len(targets) < expect_min_containers: + raise RuntimeError( + f"expected at least {expect_min_containers} containers, found {len(targets)}: {targets}" + ) + print(f"Gateway URL: {gateway}") + print(f"Direct container URLs ({len(targets)}):") + for target in targets: + print(f" {target}") + _assert_containers_at_version([gateway, *targets], expected) + _assert_gateway_completion_exact(gateway, model_name, expected) + + +def _smoke_wake_on_demand(app_name: str, cls_name: str, model_name: str, expected: int) -> None: + gateway = resolve_flash_gateway_url(app_name, cls_name) + print(f"Gateway URL: {gateway}") + # Wake the (scaled-to-zero) pool and wait for a container to serve. Flash + # holds the request through the cold start, so the timeout must exceed it. + data = _post_json(f"{gateway}/v1/chat/completions", _completion_payload(model_name, expected), timeout=900) + print(f"Gateway completion: {data}") + start, end = int(data.get("weight_version_start", -1)), int(data.get("weight_version_end", -1)) + if start > expected or end > expected: + raise VersionAheadError(f"gateway served version {start}->{end}, already past expected {expected}") + if start != expected or end != expected: + raise RuntimeError(f"unexpected gateway weight metadata: {data}") + # The pool is warm now; confirm each live container reports the version. + targets = discover_flash_targets(app_name, cls_name) + print(f"Direct container URLs ({len(targets)}):") + _assert_containers_at_version([gateway, *targets], expected) + + +def _assert_containers_at_version(targets: list[str], expected: int) -> None: + for target in targets: + info = _get_json(f"{target}/server_info", timeout=30) + print(f"{target} server_info={info}") + current = int(info["current_version"]) + if current > expected: + raise VersionAheadError(f"{target} current_version={current} already passed expected {expected}") + if current != expected: + raise RuntimeError(f"{target} current_version={current} expected {expected}") + + +def _assert_gateway_completion_exact(gateway: str, model_name: str, expected: int) -> None: + data = _post_json(f"{gateway}/v1/chat/completions", _completion_payload(model_name, expected), timeout=180) + print(f"Gateway completion: {data}") + if int(data.get("weight_version_start", -1)) != expected or int(data.get("weight_version_end", -1)) != expected: + raise RuntimeError(f"unexpected gateway weight metadata: {data}") + + +def _completion_payload(model_name: str, expected: int) -> dict: + return { + "model": model_name, + "messages": [{"role": "user", "content": "Reply with exactly OK."}], + "max_tokens": 8, + "temperature": 0, + "weight_version": {"exact_version": expected}, + "chat_template_kwargs": {"enable_thinking": False}, + } + + +def _get_json(url: str, *, timeout: float) -> dict: + with urllib.request.urlopen(url, timeout=timeout) as resp: + return json.load(resp) + + +def _post_json(url: str, payload: dict, *, timeout: float) -> dict: + request = urllib.request.Request( + url, + data=json.dumps(payload).encode(), + headers={"Content-Type": "application/json"}, + ) + with urllib.request.urlopen(request, timeout=timeout) as resp: + return json.load(resp) + + +# ── Host-RAM monitor ────────────────────────────────────────────────────────── + + +def start_host_mem_monitor(interval_s: int = 20) -> None: + """Log this node's host-RAM trajectory to stdout from a daemon thread. + + The trainer can OOM-kill on host-RAM exhaustion (the publish/update_weights + full-model gather is the peak consumer), but Megatron only reports GPU memory + and the kill leaves no durable peak behind. This logs MemTotal/MemAvailable + + the container cgroup usage every ``interval_s`` so a live ``modal app logs -f`` + shows exactly which phase blows a big node and how high it peaks. Runs on EVERY + node (called from the SPMD enter()), so whichever rank OOMs has its own trace. + Best-effort: never raises.""" + host = socket.gethostname() + + def _meminfo() -> tuple[float, float]: + total = avail = 0.0 + try: + with open("/proc/meminfo") as f: + for line in f: + if line.startswith("MemTotal:"): + total = int(line.split()[1]) / 1024 / 1024 # GiB + elif line.startswith("MemAvailable:"): + avail = int(line.split()[1]) / 1024 / 1024 + except Exception: # noqa: BLE001 + pass + return total, avail + + def _cgroup_used_gib() -> float: + for path in ("/sys/fs/cgroup/memory.current", # cgroup v2 + "/sys/fs/cgroup/memory/memory.usage_in_bytes"): # v1 + try: + with open(path) as f: + return int(f.read().strip()) / 1024**3 + except Exception: # noqa: BLE001 + continue + return -1.0 + + def _loop() -> None: + while True: + total, avail = _meminfo() + used = total - avail + cg = _cgroup_used_gib() + print( + f"[hostmem] {host} used={used:.0f}GiB avail={avail:.0f}GiB " + f"total={total:.0f}GiB cgroup_used={cg:.0f}GiB", + flush=True, + ) + time.sleep(interval_s) + + threading.Thread(target=_loop, daemon=True, name="host-mem-monitor").start() diff --git a/cookbook/trainer_helpers_test.py b/cookbook/trainer_helpers_test.py new file mode 100644 index 0000000..207f969 --- /dev/null +++ b/cookbook/trainer_helpers_test.py @@ -0,0 +1,77 @@ +from __future__ import annotations + +import unittest +from argparse import Namespace +from unittest import mock + +from cookbook import trainer_helpers + + +def _cfg(*, async_mode: bool, model_script: str | None) -> Namespace: + cfg = Namespace(async_mode=async_mode, trainer_model_script=model_script) + cfg.cli_args = lambda: ["--foo", "bar baz"] # a value with a space to prove quoting + return cfg + + +class BuildTrainCmdTest(unittest.TestCase): + def test_plain_python_invocation_without_model_script(self) -> None: + cmd = trainer_helpers.build_train_cmd( + _cfg(async_mode=False, model_script=None), "/root/slime", model_script_attr="trainer_model_script" + ) + self.assertEqual(cmd, "python3 /root/slime/train.py --foo 'bar baz'") + + def test_async_mode_selects_train_async(self) -> None: + cmd = trainer_helpers.build_train_cmd( + _cfg(async_mode=True, model_script=None), "/root/slime", model_script_attr="trainer_model_script" + ) + self.assertTrue(cmd.startswith("python3 /root/slime/train_async.py ")) + + def test_sources_model_script_and_passes_model_args(self) -> None: + cmd = trainer_helpers.build_train_cmd( + _cfg(async_mode=False, model_script="scripts/qwen.sh"), + "/root/slime", + model_script_attr="trainer_model_script", + ) + # bash -c wrapping a `source