diff --git a/AGENTS.md b/AGENTS.md index badfbf9..f636c18 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -2,4 +2,19 @@ Use Google-style testing, i.e. test `src/stitch/foo.py` with `src/stitch/foo_tes If jj is available, use that instead of git for version control. -Code legibility is highly valued, both for humans and future agents. Throughout the session, ensure that you are maintaining status quo or improving the legibility of the codebase. \ No newline at end of file +Code legibility is highly valued, both for humans and future agents. Throughout the session, ensure that you are maintaining status quo or improving the legibility of the codebase. + +## Session protocol (hard-won; do not skip) + +Grounding: +- Reproduce a defect (temp Modal probe or failing test) before fixing it. Never commit a fix whose premise is only code-reading or a subagent audit claim. Grounding first tends to produce *simpler* fixes. +- Label every finding by confidence: confirmed-by-probe / code-read inference / hypothesis. Subagent audit output is a lead to verify, not a fact. + +Modal: +- Preflight before any Modal work: confirm auth (`modal profile current`) and the target environment; pass `-e ` explicitly on every modal command, monitors included. +- Before an expensive GPU launch, run a cheap fail-fast check (CPU-only arg-parse/dry-run inside the image). Bundle image-build sanity checks into one probe; don't discover gotchas one rebuild at a time. +- Watch loops must assert positive progress (app exists, log lines grow) and must not redirect stderr to /dev/null — empty output is a failure signal, not "still booting". Smoke-test the monitor command manually once before trusting it. Long-running probes should print a machine-readable verdict line. + +Verification: +- `uv run pytest` before every commit. +- A Modal-touching change is not done until exercised in Modal. \ No newline at end of file diff --git a/CLAUDE.md b/CLAUDE.md deleted file mode 100644 index 3a19d3e..0000000 --- a/CLAUDE.md +++ /dev/null @@ -1,3 +0,0 @@ -Always start the session by immediately loading up [AGENTS.md](./AGENTS.md). - -NOTE: This is not a frontier pretraining codebase. That should be obvious, just look around. Any pipeline parallelism is embarassingly small. There is no hot-swapping of GPUs, no heterogenous hardware support, no TPUs. The training this repository enables is convenient, small finetuning for niche use cases, not serious economically valuable work that frontier models can do. Please keep this in mind as you work. diff --git a/README.md b/README.md index 6165fef..938504b 100644 --- a/README.md +++ b/README.md @@ -1,65 +1,85 @@ # stitch -A framework-agnostic protocol for disaggregated reinforcement learning of -LLMs. - -When training and rollout generation run on separate machines, the rollout -servers need a way to pick up new weights as the trainer produces them, and -each rollout request needs to know which weight version it was served with. -`stitch` provides the protocol and glue for that. Trainers publish immutable, -versioned weight artifacts to a shared "bulletin board" directory, rollout -servers sync to a requested version, and completion requests declare which -versions they will accept. - -`stitch` is unopinionated about algorithm/training framework but is strongly -opinionated about supporting workloads that are: -- async-first -- agentic-first -- elastic rollout compute +A framework-agnostic protocol for **disaggregated reinforcement learning** of LLMs. + +When training and rollout generation run on separate machines, the rollout servers need +new weights as the trainer produces them, and every rollout sample needs to know which +weight version produced it. stitch is the protocol and glue for that: trainers publish +immutable, versioned weight artifacts to a shared store; rollout replicas sync themselves +to a requested version; completion requests declare which versions they accept, and +responses report which version served them. + +It is unopinionated about the algorithm and training framework, and opinionated about the +workloads it targets: +- **async-first** — the trainer never blocks on rollout, nor rollout on the trainer +- **agentic-first** — a sample can be a long multi-turn episode spanning weight versions +- **elastic** — rollout replicas join and leave mid-run + +## What it does that a raw endpoint can't + +1. **Weight sync** — get each published version into an elastic, multi-replica pool and + reloaded into the live engine, in place, without stopping generation. +2. **Version correctness** — every sample is attributable to the weight version(s) that + produced it, and the trainer can bound staleness (serve only within N versions of + `latest`). +3. **Elastic sync** — a replica added mid-run boots, catches up over the delta chain, and + joins the rotation on its own; one still behind rejects requests it cannot serve yet + (retryable `409`) rather than emit a stale-version generation. +4. **Coordination** — full-vs-delta apply, session affinity, and MoE router replay, agreed + across publish and serve. ## How it works -1. After an optimizer step, the trainer writes weight artifacts (e.g. sparse - deltas) for version `v` to the bulletin board, publishes a - `manifest.json` describing them, then advances `latest.json`. -2. A sidecar in front of each rollout server watches the board. When a request - arrives pinned to version `v` (via a `weight_version` constraint in the - request body), the sidecar applies versions in order until it reaches `v`, - then proxies the request to the engine. -3. Responses carry the version they were served with, and a server that hasn't - caught up returns `409 WeightVersionNotReady` so the trainer can retry. +1. After an optimizer step, the trainer publishes version `v`'s weight artifacts (a full + anchor, or a sparse delta against an earlier version) under the store, then advances the + `latest` pointer. The version's HF `model.safetensors.index.json` *is* its manifest. +2. A sidecar in front of each rollout replica reconciles to `latest`. A request pinned to + version `v` waits until the replica has applied the chain up to `v`, then proxies to the + engine; the engine reloads new versions in place while it keeps serving. +3. Responses carry the version that served them; a replica behind the pinned version + returns a retryable `409`, so the caller waits or reroutes and never gets a stale + generation. + +## Elastic rollout — spin up engines mid-run + +The pool is a set of independent, self-syncing replicas: each reads the authoritative +`latest` pointer and converges itself, so adding one needs no coordination. Scale the Modal +Flash pool up and the new containers boot, base-seed, replay the delta chain to the current +version, and join the rotation on their own: + +```bash +# bump the floor from 2 -> 4; Flash boots 2 more containers that self-sync +python -c "from stitch.pools.modal_flash import ModalFlashPool; \ + ModalFlashPool('', 'Server').scale(min=4, max=4)" +``` + +A joiner pays a one-time catch-up (materialize the base, replay from the newest anchor), +which periodic full anchors bound. While it is behind, version-pinned requests it cannot +serve yet get a retryable `409` and route to caught-up replicas. + +## The sglang fork + +Rollout engines run a patched sglang: the disaggregated `/pull_weights`, correct quantized +reloads, and the O(delta) partial reload are not upstream yet. The pin and its rationale +live in **[cookbook/common/SGLANG_FORK.md](cookbook/common/SGLANG_FORK.md)** — the full +patch stack, the upstreaming PRs, and how to re-port onto a newer sglang release. ## Layout -- `src/stitch/protocol.py`: Wire protocol. Version manifests, artifacts, - request policies, the `latest.json` pointer, version-namespaced cache keys. -- `src/stitch/bulletin.py`: Bulletin board storage (filesystem-backed, with - a pluggable refresh hook for remote volumes). -- `src/stitch/sync.py`: Sync state machine that drives a server from its - current version to a target, with `quiesce` (drain, then apply) and - `in_place` (pause/apply/continue) commit modes. -- `src/stitch/servers/sglang.py`: HTTP sidecar that adds version semantics - to an SGLang server. -- `src/stitch/engines/sglang.py`: SGLang prepare/commit adapter using - `/update_weights_from_disk`. -- `src/stitch/trainers/slime.py`: Hooks for - [slime](https://github.com/THUDM/slime) that publish sparse-delta versions - from training ranks. -- `src/stitch/providers/modal.py`: Modal helpers for Volume commit/reload and - Flash container discovery. -- `cookbook/`: End-to-end examples. - - `slime_disagg/`: SLIME plus a stitch-managed Modal Flash/SGLang pool. - - `standalone_rollouts/`: standalone Modal/SGLang rollout provider with a hot-load API shim. - -The core package has no required dependencies; extras pull in what each -adapter needs (`modal`, `sglang`, `slime`). - -## Adding adapters - -Trainer adapters should publish canonical Hugging Face tensor names so engine -adapters stay trainer-agnostic. Engine adapters implement the same -prepare/commit contract as the SGLang one; the request protocol doesn't -change. +- `src/stitch/` — the library: the domain vocabulary and the sync / serve / publish logic, + plus the three ports (**Store**, **Engine**, **Pool**) and their instances. Framework-, + engine-, and provider-agnostic; no required dependencies (extras pull in `modal` / + `sglang` / `boto3`). +- `cookbook/` — deployments, not core: `common/` (shared image builds, the sidecar, the + publish/claim/request hooks, launch helpers) and the `miles_disagg/` / `slime_disagg/` + recipes with per-experiment `configs/`. + +## Adding to it + +- A different store / engine / pool is a new subclass behind the port — zero core edits. +- A different deployment or model is a new `cookbook/` recipe — core never changes. +- Trainer adapters publish canonical Hugging Face tensor names, so engine adapters stay + trainer-agnostic. ## Development diff --git a/cookbook/bulletin_hooks.py b/cookbook/bulletin_hooks.py deleted file mode 100644 index d2c0e74..0000000 --- a/cookbook/bulletin_hooks.py +++ /dev/null @@ -1,206 +0,0 @@ -"""Shared Modal Volume bulletin-board hooks for publish + rollout gating. - -Two hooks that any trainer (slime, miles, ...) plugs into via its -``custom_delta_pre_push_path`` and ``custom_rollout_request_hook_path``: - -- :func:`commit_and_wake` — advance the ``latest`` pointer, commit the Volume, - and best-effort wake the Flash rollout pool. -- :func:`gated_rollout_request_hook` — pin each rollout request to a bounded- - staleness weight version so unusable (too-stale) rollouts are never generated. - -Both hooks read their config off the trainer's ``args`` namespace (the trainer's -``--custom-config-path`` setattr's every key onto ``args``). The only -trainer-specific axis is the env-var fallback for the Flash app / class name; -callers pass those as ``app_name_env`` / ``cls_name_env``. -""" - -from __future__ import annotations - -import logging -import os -import time -from pathlib import Path -from typing import Any - -from stitch.bulletin import FilesystemBulletinBoard -from stitch.protocol import parse_weight_identity -from stitch.providers.modal import commit_volume, discover_flash_targets, volume_reloader, wake_targets - - -logger = logging.getLogger(__name__) - - -# ── Publish hook ────────────────────────────────────────────────────────────── - - -def commit_and_wake( - args: Any, - version_dir: str, - rollout_engines: list[Any], - *, - app_name_env: str, - cls_name_env: str, -) -> None: - """Trainer ``custom_delta_pre_push_path`` hook (publish-only, bulletin board). - - The trainer has written ``weight_v{N}/`` to the Modal Volume. Advance the - committed ``latest`` pointer, commit the Volume so the rollout pool's - ``reload`` sees the new version, then best-effort wake the Flash pool. The - sidecars self-sync (wake RPC, periodic poll, startup), so a missed wake only - costs latency. - - ``app_name_env`` / ``cls_name_env`` are the env-var names the trainer uses - for the Flash app and server class (e.g. ``"SLIME_DELTA_APP_NAME"``). - """ - del rollout_engines - version = parse_weight_identity(Path(version_dir).name) - rank = distributed_rank() - - # Rank 0 owns the `latest` pointer and writes 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. - if version is not None and rank in (None, 0): - FilesystemBulletinBoard(_transport_root(args), layout="slime").write_latest( - _run_id(args), version - ) - 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. - 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( - cls_name_env, "Server" - ) - wake_targets(discover_flash_targets(app_name=app_name, cls_name=cls_name), version) - except Exception: # noqa: BLE001 - logger.warning( - "Best-effort rollout wake failed for version %s; sidecars will self-sync", - version, - exc_info=True, - ) - - -# ── Staleness-gated rollout requests ────────────────────────────────────────── - - -class CachedLatestPointer: - """TTL-cached ``(run_id, version)`` from the bulletin board's ``latest`` pointer. - - The per-request hook gets no rollout_id, so the staleness floor is derived - out-of-band from the published ``latest`` pointer (the publish hook already - advanced + committed it). TTL-cached with a Volume reload so a (possibly - cross-node) rollout actor sees rank-0's committed pointer without a Volume - reload per request. - """ - - def __init__(self) -> None: - self.version: int = 0 - self.run_id: str | None = None - self._refreshed_at: float = -1e9 - self._board: FilesystemBulletinBoard | None = None - - async def get(self, args: Any, ttl: float = 2.0) -> int: - now = time.monotonic() - if self._board is None: - self._board = _gate_board(args) - if now - self._refreshed_at >= ttl: - self._refreshed_at = now - try: - await self._board.refresh() - run_id, version = self._board.read_latest() - # Staleness floor is per-run (version restarts at 1 each run); on a - # run change adopt the new run's pointer immediately so the floor - # isn't pinned to a finished run's higher version number. - self.run_id = run_id - self.version = int(version) - except Exception: # noqa: BLE001 - logger.warning( - "gate: could not read latest published version; using cached %s", - self.version, - exc_info=True, - ) - return self.version - - -_latest_cache = CachedLatestPointer() - - -async def gated_rollout_request_hook(args: Any, sample: Any, request: dict[str, Any]) -> None: - """Trainer ``custom_rollout_request_hook_path``: gate each rollout on - ``weight_version - k`` so unusable (too-stale) rollouts are never generated. - - A request pinned to ``min_required_version = latest - lag`` is admitted only - by a replica within ``lag`` versions of the newest weights; a lagging replica - returns a retryable 409 (which also nudges it to sync forward), so the trainer - never spends rollout compute on weights staler than its bound. ``min`` mode - (not ``exact``) lets the request cross in_place commits without being quiesced. - """ - mode = str(getattr(args, "rollout_request_weight_version_mode", "min")) - if mode != "none": - latest = await _latest_cache.get(args) - lag = int(getattr(args, "rollout_request_weight_version_lag", 0)) - target = max(0, latest - lag) - key = "exact_version" if mode == "exact" else "min_required_version" - request["payload"]["weight_version"] = {key: target} - - 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 - - -# ── 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"]) - - -def bulletin_root(args: Any) -> str: - """Where the trainer writes version dirs: ``/``.""" - return str( - getattr(args, "update_weight_disk_dir", None) - or os.environ.get("DELTA_BULLETIN_ROOT", "/delta-bulletin") - ) - - -def _transport_root(args: Any) -> str: - """The Volume mount root that holds the canonical ``latest`` pointer and that - the sidecar boards are rooted at — the parent of the per-run write dir.""" - return str(Path(bulletin_root(args)).parent) - - -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) - - -def _gate_board(args: Any) -> FilesystemBulletinBoard: - vol = getattr(args, "update_weight_delta_volume_name", None) or os.environ.get("DELTA_VOLUME_NAME") - refresh = volume_reloader(vol) if vol else None - return FilesystemBulletinBoard(_transport_root(args), refresh=refresh, layout="slime") diff --git a/cookbook/common/SGLANG_FORK.md b/cookbook/common/SGLANG_FORK.md new file mode 100644 index 0000000..6830755 --- /dev/null +++ b/cookbook/common/SGLANG_FORK.md @@ -0,0 +1,127 @@ +# The sglang fork stitch pins + +stitch's rollout engines run a **patched sglang**, not upstream: the disaggregated +weight-sync path (engine-side `/pull_weights`, correct quantized reloads, and the +O(delta) partial reload) is not in upstream sglang yet. This file is the source of +truth for *what* we patch, *why*, and *how to move the patch stack onto a newer +sglang release*. + +## The pin + +`cookbook/common/serving_image.py` builds the serving image by overlaying the fork's +`python/` onto a stock sglang container: + +```python +SGLANG_IMAGE_TAG = "lmsysorg/sglang:v0.5.15" # base kernels/CUDA +SGLANG_FORK_BRANCH = "stitch-sglang-v0.5.15" # modal-projects/sglang +SGLANG_FORK_COMMIT = "13479b59cccd77459fb003d2f2e138e4cca8ed17" +``` + +The branch is **`v0.5.15` + the patch stack below, nothing else** — every commit +`git log v0.5.15..stitch-sglang-v0.5.15` is one of ours. Because only `python/` is +overlaid, `SGLANG_IMAGE_TAG` **must** be the same sglang version the branch is based +on (`v0.5.15`), or the baked C++/CUDA ops will be ABI-mismatched with the python. + +## The patch stack + +Two tiers. The **base case** makes disaggregated weight sync work at all; the +**optimization** tier makes a reload O(delta) instead of O(full checkpoint). Read the +commit bodies for the details — this is the map. + +### Base case — weight sync works + +1. **`[RL] /pull_weights: engine-side pull ... into a host-local checkpoint`** + The disaggregated receiver: `POST /pull_weights` walks the published + `weight_v{N}/` chain from the nearest full anchor, replays deltas (xor + zstd, + xxh3 checksum) into a host-local checkpoint that `/update_weights_from_disk` then + reloads, while the engine keeps serving. Hardened for an eventually-consistent + volume mount (whole-file in-memory read + size-verify before the xor, one reload + per host, reseed from the pristine boot checkpoint). + *Upstreaming:* https://github.com/sgl-project/sglang/pull/30367 + +2. **`[RL] update_weights_from_disk: load quantized checkpoints like initial loading`** + Makes an in-place reload of a quantized checkpoint reproduce `init(checkpoint)` + (fp8 blockwise + compressed-tensors block/channel): restore latched quant scale + state before loading, fix the rollback path, keep weights refillable, and fix the + UE8M0 scale inverse for row counts not divisible by 128. Without this, fp8 + reloads silently diverge from the served kernel format. + *Upstreaming:* https://github.com/sgl-project/sglang/pull/30761 + +### Optimization — O(delta) reload + +3. **`[RL] reload: record/replay load plans for repeated reloads`** + Record the model's first-reload weight dispatch once, replay it directly after + (skip the per-tensor routing scan, parallelize the load). Opt in per model + (`supports_load_plan_replay`); gated by `SGLANG_ENABLE_RELOAD_LOAD_PLAN=1`. Falls + back to the legacy loader on any failure. + +4. **`[RL] reload: O(delta) partial reload via touched checkpoint names`** + Given the touched tensor names (`weight_names`, from a delta apply), reload only + those tensors + their fused/expert closures and re-post-process only the touched + modules. Pre-flights every touched module's quant method for incremental support; + any gap falls back to a full reload. + +5. **`[RL] modelopt fp4: incremental post-loading for partial reloads`** + The NVFP4 model-side enabler for (4): `process_weights_after_partial_loading` + re-derives kernel state for only the touched experts — CUTLASS per-expert + re-swizzle and TRT-LLM per-expert re-alignment — declining marlin/cutedsl and any + padded/whole-layer layout to a safe full reload. Also restores NVFP4 raw bit-exact + compare in the weight checker (v0.5.15 replaced it with a NotImplementedError) so a + partial reload is byte-verifiable against a full one via `/weights_checker`. Without + this, NVFP4 partial reload declines and pays a full reload. + *Upstreaming:* not yet filed. + +6. **`[RL] load plan: record during the initial load so reloads start already-replaying`** + Record the load plan during the model's initial boot load, so the first + `update_weights_from_disk` already replays / goes O(delta) partial instead of paying + a full record-reload — the full reload is eliminated from steady state (matters most + for elastic joiners that boot then immediately catch up via deltas). Gated on the + same flag; drops the plan and falls back to a plain load on any failure. + +7. **`[RL] fastsafetensors: env-gated nogds for sandboxed (no-GDS) hosts`** + The fast full-reload path for a DENSE checkpoint, where the O(delta) partial reload + (3–6) doesn't apply (e.g. fp8 — the delta touches ~all weight tensors). Upstream + `--load-format fastsafetensors` splits the safetensors files across TP ranks (each + rank reads only its 1/N) but defaults to GDS/cuFile, which needs the `nvidia-fs` + driver — absent under gVisor, so it fails at open (`"Error opening file"`). Read + `SGLANG_FASTSAFETENSORS_NOGDS` to force the nogds path (O_DIRECT + host bounce + buffer); default preserves upstream GDS behavior. The win is *fewer bytes*, not the + disk: under gVisor the read is bound by the Sentry's single-thread byte-copy + (~5 GB/s/node), so cutting the per-node read ~408→~102 GB takes GLM-4.5-Air-fp8's + reload ~130 s → ~23 s. Enabled per recipe via `--load-format fastsafetensors`. + *Upstreaming:* not yet filed. + +`SGLANG_ENABLE_RELOAD_LOAD_PLAN` is opt-in per recipe (off unless set): the NVFP4 configs +enable it via `SGLANG_ENV` — their native load is single-threaded, so replay is a large win. +A dense config where partial reload can't apply (fp8) instead uses `--load-format +fastsafetensors` (patch 7) for a fast *full* reload and leaves the load plan off (see +`cookbook/miles_disagg/configs/glm45_air_fp8.py`). + +## Re-porting to a newer sglang release (`stitch-sglang-vX`) + +When bumping the base (e.g. to `v0.5.16`): + +1. Branch `stitch-sglang-vX` off the new tag on `modal-projects/sglang`. +2. Re-apply the seven commits **in the order above** (cherry-pick from + `stitch-sglang-v0.5.15`, or from the source branches `weight-sync-miles` / + `fp8-reload-main` / `weight-sync-upstream`). Squash-preserving is fine; keep the + two tiers legible. Patch 7 (fastsafetensors nogds) is a small standalone one-liner. +3. **Audit before trusting a clean cherry-pick** — a clean apply does not prove the + patch is still needed or correct. In particular: + - `[RL] modelopt fp4 ...` (5) is tightly coupled to sglang's NVFP4 MoE post-load + (`ModelOptNvFp4FusedMoEMethod.process_weights_after_loading`, the swizzle/alias + helpers, and the MoE runner backends). This function is restructured often; the + partial pass usually needs a genuine **rewrite**, not a port. The v0.5.15 + rewrite scopes the fast path to CUTLASS and declines other backends to a safe + full reload. + - Check whether the new base already upstreamed any patch (then drop it). +4. **Validate** a quantized partial reload is byte-identical to a full reload of the + same bytes via `/weights_checker` (per-tensor checksums) before shipping — this is + how the NVFP4 partial pass is proven correct. +5. Update the three constants in `cookbook/common/serving_image.py` and this file. + +## Upstreaming + +We want these in upstream sglang so the fork shrinks to zero. Open PRs: +`/pull_weights` → sgl-project/sglang#30367; quantized reload → sgl-project/sglang#30761. +The load-plan / partial-reload tier is not yet filed. diff --git a/cookbook/common/__init__.py b/cookbook/common/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/cookbook/common/config.py b/cookbook/common/config.py new file mode 100644 index 0000000..c0000db --- /dev/null +++ b/cookbook/common/config.py @@ -0,0 +1,38 @@ +"""``ModalConfig`` — the Modal-infrastructure half of an experiment config, shared by +every recipe (GPU model, region, rollout-pool sizing, prep topology). + +The framework *training* config (``MilesConfig`` / ``SlimeConfig``) is self-contained in +each framework's subdir; only this infra grouping is common. +""" + +from __future__ import annotations + +from typing import Any, Literal + +GPUType = Literal["H100", "H200", "B200", "B300", "A100"] + + +class ModalConfig: + """Modal infrastructure: GPU model, region, rollout-pool sizing, prep topology.""" + + gpu: GPUType = "B200" + memory: tuple[int, int] | None = None + cloud: str | None = None + region: str | None = None + rollout_min_containers: int = 2 + rollout_max_containers: int | None = None + # Flash autoscaler target: keep well below the sglang engine concurrency so Flash + # adds containers instead of packing requests until KV saturates and they stall. + rollout_target_inputs: int | None = None + proxy_regions: list[str] = ["us-west"] + rollout_ephemeral_disk_mib: int | None = None + rollout_memory_mib: int | None = None + torch_dist_prep_nodes: int = 2 + torch_dist_prep_gpus_per_node: int = 8 + torch_dist_convert_extra_args: str = "" + torch_dist_prep_ephemeral_disk_mib: int | None = None + trainer_ephemeral_disk_mib: int | None = None + + def __init__(self, **kwargs: Any) -> None: + for k, v in kwargs.items(): + setattr(self, k, v) diff --git a/cookbook/common/constants.py b/cookbook/common/constants.py new file mode 100644 index 0000000..7925748 --- /dev/null +++ b/cookbook/common/constants.py @@ -0,0 +1,24 @@ +"""Shared deployment constants — container mount points, ports, and timeouts, general +across every recipe. The per-experiment names/values (APP_NAME, DELTA_VOLUME_NAME, the +disk layout) live in the config modules. +""" + +from __future__ import annotations + +from pathlib import Path + +# Container mount points (Modal Volumes attach here). +HF_CACHE_PATH = Path("/root/.cache/huggingface") +DATA_PATH = Path("/data") +CHECKPOINTS_PATH = Path("/checkpoints") +PREP_PATH = Path("/prep") # //{bf16 masters, served base, torch_dist ref_load} +SGLANG_CACHE_PATH = "/root/.cache/sglang" # sglang kernel/JIT cache; survives cold starts + +# Ports. +SIDECAR_PORT = 8000 # the versioned rollout proxy — the container's public port +SGLANG_PORT = 8001 # the private sglang server behind the sidecar +RAY_PORT = 6379 + +# Timeouts. +MINUTES = 60 +SERVER_STARTUP_TIMEOUT = 35 * MINUTES diff --git a/cookbook/common/hooks.py b/cookbook/common/hooks.py new file mode 100644 index 0000000..203f8df --- /dev/null +++ b/cookbook/common/hooks.py @@ -0,0 +1,147 @@ +"""Shared framework-hook logic (miles and slime both write the same HF-delta layout to +a Modal Volume and wake a Modal Flash pool, so the logic is one place). + +Each framework re-exports these from its own ``hooks.py`` under the dotted paths its run +config points at; the only per-framework difference is that thin re-export. Every shim +reads the run's coordinates off the trainer's ``args`` namespace (the framework +``setattr``s its ``custom_config_path`` onto it) and calls the stitch core against a +ModalVolumeStore + ModalFlashPool. +""" + +from __future__ import annotations + +import asyncio +import logging +import os +import time +from pathlib import Path +from typing import Any + +from stitch.pools.modal_flash import ModalFlashPool +from stitch.publish import claim_run, constrain_request, publish_version +from stitch.stores.modal_volume import ModalVolumeStore +from stitch.versions import PointerRewind + +logger = logging.getLogger(__name__) + + +# ── publish ──────────────────────────────────────────────────────────────────── +def commit_and_wake(args: Any, published_dir: str, rollout_engines: Any = None) -> None: + """Bridge the framework's disk-delta publish to the stitch store. The framework fires + this at each shared-store durability boundary (a dir to flush to the volume): an actual + version dir (``weight_vNNNNNN``, holding the HF index) and — at baseline / pointer + commit — the run dir. Every rank flushes its own writes; rank 0 publishes (manifest + + pointer advance + pool wake) ONLY when handed a version dir. Keying on the dir the + framework names, not on reading an index, is what makes the run-dir calls (no index yet) + a clean no-op instead of a missing-file crash.""" + del rollout_engines + store = _store(args) + store.commit() # flush this rank's local writes (shards / index / pointer) to the volume + if _rank() not in (None, 0) or not Path(published_dir).name.startswith("weight_v"): + return + try: + publish_version(store, _pool(args), published_dir, run_id=_run_id(args)) + except PointerRewind: + # A same-run republish (e.g. a retried step) — drop it rather than serve stale. + logger.warning("publish of %s would rewind latest; dropping", published_dir, exc_info=True) + + +def claim_pool(args: Any) -> None: + """Launch hook (rank 0): reset every replica to base before the first publish, so a + cold or finished-run-warm pool starts this run clean.""" + if _rank() not in (None, 0): + return + claim_run(_store(args), _pool(args), _run_id(args)) # raises PointerRewind on a reused run_id + + +# ── staleness-gated rollout requests ──────────────────────────────────────────── +async def gated_rollout_request_hook(args: Any, sample: Any, request: dict[str, Any]) -> None: + """Pin each request to a bounded-staleness version, so a too-stale replica returns a + retryable 409 (nudging it to sync) instead of the trainer spending rollout compute on + weights beyond its lag bound.""" + payload, headers = request["payload"], dict(request.get("headers") or {}) + mode = str(getattr(args, "rollout_request_weight_version_mode", "min")) + affinity = str(getattr(args, "rollout_session_affinity_header", "x-session-affinity")) + session_id = getattr(sample, "session_id", None) + + latest = exact = None + lag = 0 + if mode != "none": + floor = await _latest.get(args) + lag = int(getattr(args, "rollout_request_weight_version_lag", 0)) + if mode == "exact": + exact = max(0, floor - lag) + else: + latest = floor + constrain_request( + payload, headers, latest=latest, lag=lag, exact=exact, + session_id=session_id, affinity_header=affinity, + ) + request["headers"] = headers + 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))) + + +class _CachedPointer: + """TTL-cached ``latest`` version. The per-request hook gets no rollout id, so the + staleness floor comes from the published pointer (already advanced by the publish + hook), cached with a Volume reload so it isn't reloaded once per request.""" + + def __init__(self) -> None: + self._version = 0 + self._at = -1e9 + self._store: ModalVolumeStore | None = None + + async def get(self, args: Any, ttl: float = 2.0) -> int: + store = self._store + if store is None: + store = self._store = _store(args) + now = time.monotonic() + if now - self._at >= ttl: + self._at = now + try: + await asyncio.to_thread(store.refresh) # reload is blocking; keep the loop free + pointer = store.read_pointer() + self._version = pointer.version if pointer else 0 + except Exception: # noqa: BLE001 + logger.warning("gate: could not read latest; using cached %s", self._version, exc_info=True) + return self._version + + +_latest = _CachedPointer() + + +# ── args → run coordinates ─────────────────────────────────────────────────────── +def _store(args: Any) -> ModalVolumeStore: + volume = getattr(args, "update_weight_delta_volume_name", None) or os.environ.get("DELTA_VOLUME_NAME") + return ModalVolumeStore(_transport_root(args), volume_name=volume or None) + + +def _pool(args: Any) -> ModalFlashPool: + app = getattr(args, "rollout_modal_flash_app_name", None) or os.environ["DELTA_APP_NAME"] + cls = getattr(args, "rollout_modal_flash_server_cls_name", None) or os.environ.get("DELTA_SERVER_CLS_NAME", "Server") + return ModalFlashPool(app, cls) + + +def _transport_root(args: Any) -> str: + # The trainer writes version dirs under /; the Store is rooted at . + write_dir = getattr(args, "update_weight_disk_dir", None) or os.environ.get("DELTA_BULLETIN_ROOT", "/delta-bulletin") + return str(Path(write_dir).parent) + + +def _run_id(args: Any) -> str: + run_id = getattr(args, "run_id", None) + if not run_id: + raise ValueError("run_id is required (pass it via custom_config_path) — it is the run's fence token") + return str(run_id) + + +def _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 diff --git a/cookbook/common/hooks_test.py b/cookbook/common/hooks_test.py new file mode 100644 index 0000000..efbdaad --- /dev/null +++ b/cookbook/common/hooks_test.py @@ -0,0 +1,119 @@ +"""Harness for the shared hook shims. + +Runs without Modal/torch: the real ``_store`` builds a local dir (volume_name=None from +the temp root) and only ``_pool`` is faked. Run directly: + PYTHONPATH=src:. python cookbook/common/hooks_test.py +""" + +from __future__ import annotations + +import asyncio +import json +import tempfile +from pathlib import Path +from types import SimpleNamespace + +from cookbook.common import hooks +from stitch.stores.modal_volume import ModalVolumeStore +from stitch.versions import VersionRef + + +class _FakePool: + def __init__(self) -> None: + self.woke: list = [] + + def discover_replicas(self): + return ["http://r1"] + + def wake(self, replicas, ref): + self.woke.append(ref) + + +def _args(root: str, run_id: str = "run-abc", **extra): + # transport root = parent of update_weight_disk_dir, so the Store roots at `root`. + return SimpleNamespace(update_weight_disk_dir=f"{root}/{run_id}", run_id=run_id, **extra) + + +def _write_version(root: Path, ref: VersionRef) -> str: + d = root / ref.identity + d.mkdir(parents=True) + (d / "model.safetensors.index.json").write_text( + json.dumps({"metadata": {"version": ref.version}, "weight_map": {"w": "model-00001.safetensors"}}) + ) + return str(d) + + +def test_commit_and_wake_publishes() -> None: + with tempfile.TemporaryDirectory() as tmp: + root = Path(tmp) + pool = _FakePool() + hooks._pool = lambda args: pool # rank is None in tests -> treated as writer + vdir = _write_version(root, VersionRef("run-abc", 1)) + hooks.commit_and_wake(_args(str(root)), vdir) + assert ModalVolumeStore(root).read_pointer() == VersionRef("run-abc", 1) + assert pool.woke == [VersionRef("run-abc", 1)] + + +def test_commit_and_wake_baseline_is_noop() -> None: + # The framework fires the hook at baseline/pointer-commit with the RUN dir (name is the + # run_id, no version written). Keying on the dir name means we flush the volume but read + # no index and publish nothing — the regression was reading an index that wasn't there. + with tempfile.TemporaryDirectory() as tmp: + root = Path(tmp) + pool = _FakePool() + hooks._pool = lambda args: pool + run_dir = root / "run-abc" + run_dir.mkdir(parents=True) + hooks.commit_and_wake(_args(str(root)), str(run_dir)) + assert ModalVolumeStore(root).read_pointer() is None + assert pool.woke == [] + + +def test_claim_pool_resets_to_base() -> None: + with tempfile.TemporaryDirectory() as tmp: + pool = _FakePool() + hooks._pool = lambda args: pool + hooks.claim_pool(_args(tmp)) + assert ModalVolumeStore(Path(tmp)).read_pointer() == VersionRef("run-abc", 0) + assert pool.woke == [VersionRef("run-abc", 0)] + + +def test_request_hook_min_lag() -> None: + with tempfile.TemporaryDirectory() as tmp: + root = Path(tmp) + ModalVolumeStore(root).advance_pointer(VersionRef("run-abc", 10)) # published latest + hooks._latest = hooks._CachedPointer() # fresh cache reading this store + args = _args(str(root), rollout_request_weight_version_lag=2, + rollout_request_retry_attempts=900, rollout_session_affinity_header="Modal-Session-ID") + request = {"payload": {}} + asyncio.run(hooks.gated_rollout_request_hook(args, SimpleNamespace(session_id="grp-1"), request)) + assert request["payload"]["weight_version"] == {"min_version": 8, "exact_version": None} + assert request["headers"]["Modal-Session-ID"] == "grp-1" + assert request["max_retries"] == 900 + + +def test_request_hook_exact_and_none() -> None: + with tempfile.TemporaryDirectory() as tmp: + root = Path(tmp) + ModalVolumeStore(root).advance_pointer(VersionRef("run-abc", 10)) + hooks._latest = hooks._CachedPointer() + exact_req = {"payload": {}} + asyncio.run(hooks.gated_rollout_request_hook( + _args(str(root), rollout_request_weight_version_mode="exact", rollout_request_weight_version_lag=1), + SimpleNamespace(session_id=None), exact_req)) + assert exact_req["payload"]["weight_version"] == {"min_version": None, "exact_version": 9} + + hooks._latest = hooks._CachedPointer() + none_req = {"payload": {}} + asyncio.run(hooks.gated_rollout_request_hook( + _args(str(root), rollout_request_weight_version_mode="none"), + SimpleNamespace(session_id=None), none_req)) + assert none_req["payload"]["weight_version"] == {"min_version": None, "exact_version": None} + + +if __name__ == "__main__": + tests = [v for k, v in sorted(globals().items()) if k.startswith("test_")] + for t in tests: + t() + print(f" ok {t.__name__}") + print(f"common hooks harness: {len(tests)} PASS") diff --git a/cookbook/common/launch.py b/cookbook/common/launch.py new file mode 100644 index 0000000..f483593 --- /dev/null +++ b/cookbook/common/launch.py @@ -0,0 +1,41 @@ +"""Trainer-launch helpers shared by every recipe: build the train command (sourcing the +model-arch MODEL_ARGS script) and resolve/materialize a config before launch. Both are +framework-agnostic given the framework's root + model-script attribute + YAML fields. +""" + +from __future__ import annotations + +import os +import shlex +from typing import Any + + +def build_train_cmd(cfg: Any, root: str, model_script_attr: str) -> str: + """The train command. ``train_async.py`` / ``train.py`` live at the framework root + and consume the ``MODEL_ARGS`` bash array defined by the sourced model script.""" + train_script = f"{root}/{'train_async.py' if cfg.async_mode else 'train.py'}" + model_script = getattr(cfg, model_script_attr, "") + if model_script: + inner = ( + f"source {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())}" + + +def resolve_config(cfg: Any, tmpdir: str, yaml_fields: tuple[str, ...]) -> None: + """Resolve HF repo-id checkpoint fields to local paths and materialize inline YAML + config dicts to files the trainer reads. Absolute paths are left untouched.""" + 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_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) diff --git a/cookbook/common/process.py b/cookbook/common/process.py new file mode 100644 index 0000000..e551e6a --- /dev/null +++ b/cookbook/common/process.py @@ -0,0 +1,113 @@ +"""Subprocess + runtime helpers shared by every recipe: launch the sidecar beside +sglang, wait on HTTP liveness, terminate cleanly, monitor host RAM, apply git patches. +""" + +from __future__ import annotations + +import os +import signal +import socket +import subprocess +import threading +import time +import urllib.error +import urllib.request + +SIDECAR_MODULE = "cookbook.common.sidecar" # the one shared serve() entrypoint + + +def start_sidecar( + *, sidecar_port: int, sglang_port: int, bulletin_root: str, local_checkpoint_dir: str, + volume_name: str, commit_mode: str, debug_requests: bool = False, +) -> subprocess.Popen: + """Launch the versioned rollout proxy (the shared sidecar) beside sglang.""" + cmd = [ + "python3", "-m", SIDECAR_MODULE, + "--host", "0.0.0.0", "--port", str(sidecar_port), + "--upstream", f"http://127.0.0.1:{sglang_port}", + "--bulletin-root", bulletin_root, + "--local-checkpoint-dir", local_checkpoint_dir, + "--volume-name", volume_name, + "--commit-mode", commit_mode, + ] + if debug_requests: + cmd.append("--debug-requests") + print("Starting sidecar:", " ".join(cmd)) + return subprocess.Popen(cmd, start_new_session=True) + + +def wait_http(url: str, process: subprocess.Popen | None, timeout: int) -> None: + deadline = time.time() + timeout + last_error: str | None = None + while time.time() < deadline: + if process is not None and process.poll() is not None: + raise RuntimeError(f"process exited while waiting for {url}: code={process.returncode}") + try: + with urllib.request.urlopen(url, timeout=5) as resp: + if 200 <= resp.status < 500: + return + except (urllib.error.URLError, TimeoutError, OSError) as exc: + last_error = f"{type(exc).__name__}: {exc}" + time.sleep(2) + raise TimeoutError(f"timed out waiting for {url}; last error: {last_error}") + + +def terminate_process(process: subprocess.Popen | None) -> None: + if process is None or process.poll() is not None: + return + try: + os.killpg(process.pid, signal.SIGTERM) + process.wait(timeout=20) + except Exception: # noqa: BLE001 + try: + os.killpg(process.pid, signal.SIGKILL) + except Exception: # noqa: BLE001 + pass + + +def apply_git_patches(patch_paths: list[str], repo_dir: str, label: str) -> None: + """Apply git patches to a runtime checkout, tolerating an already-applied patch + (idempotent across container restarts).""" + for patch_path in patch_paths: + if not os.path.exists(patch_path): + raise FileNotFoundError(f"{label} not found: {patch_path}") + check = subprocess.run(["git", "-C", repo_dir, "apply", "--check", patch_path], capture_output=True, text=True) + if check.returncode == 0: + subprocess.run(["git", "-C", repo_dir, "apply", patch_path], check=True) + print(f"[{label}] applied {patch_path}", flush=True) + continue + reverse = subprocess.run( + ["git", "-C", repo_dir, "apply", "--reverse", "--check", patch_path], capture_output=True, text=True + ) + if reverse.returncode == 0: + print(f"[{label}] already applied {patch_path}", flush=True) + continue + raise RuntimeError(f"cannot apply {label} {patch_path}\ncheck: {check.stderr}\nreverse: {reverse.stderr}") + + +def start_host_mem_monitor(interval_s: int = 20) -> None: + """Trace this node's host-RAM from a daemon thread. Host-RAM exhaustion OOM-kills the + trainer (the peak is the publish weight-gather) and Modal exposes no host-RAM metric, so + this log line is the only signal. Best-effort.""" + 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 + elif line.startswith("MemAvailable:"): + avail = int(line.split()[1]) / 1024 / 1024 + except Exception: # noqa: BLE001 + pass + return total, avail + + def _loop() -> None: + while True: + total, avail = _meminfo() + print(f"[hostmem] {host} used={total - avail:.0f}GiB avail={avail:.0f}GiB total={total:.0f}GiB", flush=True) + time.sleep(interval_s) + + threading.Thread(target=_loop, daemon=True, name="host-mem-monitor").start() diff --git a/cookbook/common/ray_cluster.py b/cookbook/common/ray_cluster.py new file mode 100644 index 0000000..04204d3 --- /dev/null +++ b/cookbook/common/ray_cluster.py @@ -0,0 +1,97 @@ +"""Ray cluster bring-up for multi-node Modal training — framework-agnostic.""" + +from __future__ import annotations + +import socket +import subprocess +import time +from pathlib import Path + +RAY_START_TIMEOUT = 240 +RAY_WORKER_JOIN_TIMEOUT = 180 + + +def get_modal_cluster_context(n_nodes: int) -> tuple[int, str, str]: + """(rank, master_addr, my_ip) for the current Modal cluster (single-node safe).""" + import modal.experimental + + try: + info = modal.experimental.get_cluster_info() + except Exception: # noqa: BLE001 + if n_nodes == 1: + ip = _local_ip() + return 0, ip, ip + raise + actual = len(info.container_ipv4_ips) + if actual == 0 and n_nodes == 1: + ip = _local_ip() + return 0, ip, ip + if actual != n_nodes: + raise RuntimeError(f"cluster size mismatch: expected {n_nodes} node(s), got {actual}") + return info.rank, info.container_ipv4_ips[0], info.container_ipv4_ips[info.rank] + + +def _local_ip() -> str: + with socket.socket(socket.AF_INET, socket.SOCK_DGRAM) as sock: + try: + sock.connect(("8.8.8.8", 80)) + return sock.getsockname()[0] + except OSError: + return socket.gethostbyname(socket.gethostname()) + + +def start_ray_head(my_ip: str, n_nodes: int, *, ray_port: int) -> None: + import ray + + try: + subprocess.run( + ["ray", "start", "--head", f"--node-ip-address={my_ip}", f"--port={ray_port}", + "--disable-usage-stats", "--include-dashboard=false"], + check=True, timeout=RAY_START_TIMEOUT, + ) + except (subprocess.TimeoutExpired, subprocess.CalledProcessError) as exc: + _print_ray_logs() + raise RuntimeError(f"Ray head failed to start: {exc}") from exc + + last_error = "" + for _ in range(RAY_START_TIMEOUT): + try: + ray.init(address=f"{my_ip}:{ray_port}") + break + except Exception as exc: # noqa: BLE001 + last_error = f"{type(exc).__name__}: {exc}" + time.sleep(1) + else: + _print_ray_logs() + raise RuntimeError(f"Ray head failed to start before timeout: {last_error}") + + for _ in range(RAY_WORKER_JOIN_TIMEOUT): + alive = [n for n in ray.nodes() if n["Alive"]] + print(f"Waiting for workers: {len(alive)}/{n_nodes} alive") + if len(alive) == n_nodes: + return + time.sleep(1) + _print_ray_logs() + raise RuntimeError(f"Timed out waiting for all {n_nodes} Ray nodes to join") + + +def start_ray_worker(my_ip: str, master_addr: str, *, ray_port: int) -> None: + subprocess.run( + ["ray", "start", f"--node-ip-address={my_ip}", "--address", f"{master_addr}:{ray_port}", + "--disable-usage-stats"], + check=True, timeout=RAY_START_TIMEOUT, + ) + + +def _print_ray_logs() -> None: + log_dir = Path("/tmp/ray/session_latest/logs") + for name in ("gcs_server.out", "gcs_server.err", "raylet.out", "raylet.err", "monitor.err"): + path = log_dir / name + if not path.exists(): + continue + print(f"===== {path} =====") + try: + for line in path.read_text(errors="replace").splitlines()[-80:]: + print(line) + except OSError as exc: + print(f"could not read {path}: {exc}") diff --git a/cookbook/common/server.py b/cookbook/common/server.py new file mode 100644 index 0000000..f4b10fe --- /dev/null +++ b/cookbook/common/server.py @@ -0,0 +1,67 @@ +"""Shared rollout-Server logic: boot sglang + the stitch sidecar, and tear them down. + +Modal requires an ``@app.cls`` class to live at module (global) scope, so each framework's +``app.py`` defines a thin ``Server`` class there; its ``@enter``/``@exit`` delegate to these +functions. The behavior is shared here; only the class skeleton + its per-app decorators +(image / gpu / volumes) stay local. The public container port is the sidecar (it fronts the +private sglang on SGLANG_PORT). +""" + +from __future__ import annotations + +import os +from typing import Any + +from . import process +from .constants import SGLANG_PORT, SIDECAR_PORT + + +def serve_startup( + replica: Any, + *, + model_name: str, + sglang_args: dict, + tp: int, + concurrency: int, + bulletin_root: str, + local_checkpoint_dir: str, + volume_name: str, + commit_mode: str, + startup_timeout: int, + sglang_env: dict[str, str] | None = None, +) -> None: + """Start sglang + the versioned-proxy sidecar on a Server replica (from ``@modal.enter``). + The engine serves ``model_name`` and materializes each version into + ``local_checkpoint_dir`` itself via /pull_weights; the sidecar drives the sync. + ``sglang_env`` is a per-config override of the sglang process env (over the image's + baked defaults) — set before launch so the engine subprocess inherits it.""" + from autoinference_utils.endpoint import SGLangEndpoint, warmup_chat_completions + + if sglang_env: + os.environ.update(sglang_env) + replica.endpoint = SGLangEndpoint( + model_path=model_name, worker_port=SGLANG_PORT, tp=tp, + extra_server_args=sglang_args, health_timeout=startup_timeout, health_poll_interval=10.0, + ) + replica.endpoint.start() + warmup = { + "model": model_name, + "messages": [{"role": "user", "content": "Reply with exactly OK."}], + "max_tokens": 8, "temperature": 0, "chat_template_kwargs": {"enable_thinking": False}, + } + warmup_chat_completions(port=SGLANG_PORT, payload=warmup, successful_requests=2, + request_timeout=120.0, max_attempts_per_request=3) + replica.sidecar = process.start_sidecar( + sidecar_port=SIDECAR_PORT, sglang_port=SGLANG_PORT, bulletin_root=bulletin_root, + local_checkpoint_dir=local_checkpoint_dir, volume_name=volume_name, commit_mode=commit_mode, + ) + process.wait_http(f"http://127.0.0.1:{SIDECAR_PORT}/health", replica.sidecar, startup_timeout) + print(f"Rollout server ready: model={model_name}, target_inputs={concurrency}") + + +def serve_stop(replica: Any) -> None: + """Tear down the sidecar + sglang (from ``@modal.exit``).""" + process.terminate_process(getattr(replica, "sidecar", None)) + endpoint = getattr(replica, "endpoint", None) + if endpoint is not None: + endpoint.stop() diff --git a/cookbook/common/serving_image.py b/cookbook/common/serving_image.py new file mode 100644 index 0000000..552f27d --- /dev/null +++ b/cookbook/common/serving_image.py @@ -0,0 +1,82 @@ +"""The weight-sync sglang SERVING image — shared by every recipe. + +Trainer-agnostic: no trainer package is installed (the delta apply lives in the engine +behind ``/pull_weights``), so miles and slime serve on the identical image; precision +comes from the served checkpoint, not a ``--quantization`` flag. The base image supplies +kernels/CUDA; the fork pin (branch ``stitch-sglang-v0.5.15`` over base tag ``v0.5.15``) +carries ``/pull_weights``, the hardened local_checkpoint receiver, the quantized-reload +restore protocol (reload == init), and the O(delta) partial-reload load plan. See +``SGLANG_FORK.md`` next to this file for the full patch stack, the upstreaming PRs, and +how to re-port these patches onto a newer sglang release. +""" + +from __future__ import annotations + +from pathlib import Path + +import modal + +# Fork = base sglang tag + the stitch weight-sync patch stack (see SGLANG_FORK.md). +# The base tag MUST match the branch's base tag: the fork overlays python/ only, so the +# baked kernels/CUDA come from this image and must be ABI-compatible with that python/. +SGLANG_IMAGE_TAG = "lmsysorg/sglang:v0.5.15" +SGLANG_FORK_REPO = "https://github.com/modal-projects/sglang.git" +SGLANG_FORK_BRANCH = "stitch-sglang-v0.5.15" +SGLANG_FORK_COMMIT = "13479b59cccd77459fb003d2f2e138e4cca8ed17" + +_COOKBOOK_DIR = Path(__file__).resolve().parent.parent # .../cookbook + +_SERVING_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", + # GDS/cuFile needs nvidia-fs, absent under gVisor → force fastsafetensors onto its + # nogds path (O_DIRECT + host bounce). Only read when --load-format fastsafetensors. + "SGLANG_FASTSAFETENSORS_NOGDS": "1", +} + + +def build_serving_image( + *, + hf_cache_path: str, + delta_volume_name: str, + experiment: str, + extra_commands: tuple[str, ...] = (), + extra_env: dict[str, str] | None = None, +) -> modal.Image: + """The rollout-pool image. ``DELTA_VOLUME_NAME`` is read by the engine's pre-read hook + and the sidecar's Store; ``EXPERIMENT_CONFIG`` so the container's re-import resolves the + same experiment as the deploy; stitch + the whole cookbook package are mounted so the + sidecar (``python -m cookbook.common.sidecar``) and the framework hooks resolve. + + ``extra_commands`` / ``extra_env`` let an experiment extend the shared image (e.g. the + NVFP4 4/6 recipe's flashinfer lockstep upgrade and its FLASHINFER_* quantizer + contract) without forking it; commands run after the fork checkout.""" + 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} && git checkout {SGLANG_FORK_COMMIT} -- python/" + ) + ) + for cmd in extra_commands: + image = image.run_commands(cmd) + return ( + image + .run_commands(f"rm -rf {hf_cache_path}") # baked HF cache must not shadow the mounted volume + .pip_install( + "autoinference-utils==0.2.0", # sglang server lifecycle + "fastapi", "httpx", "uvicorn", # the stitch sidecar + "zstandard", "xxhash", "blake3", # engine-side /pull_weights receiver's codecs + "fastsafetensors", # --load-format fastsafetensors: per-rank read (nogds, see env below) + ) + .env({**_SERVING_ENV, **(extra_env or {}), "DELTA_VOLUME_NAME": delta_volume_name, "EXPERIMENT_CONFIG": experiment}) + # The kernel-cache volume mounts at /root/.cache/sglang, which can't mount over a + # non-empty path — clear it as the final filesystem step (repopulated on boot). + .run_commands("rm -rf /root/.cache/sglang") + .add_local_python_source("stitch") + .add_local_dir(str(_COOKBOOK_DIR), remote_path="/root/cookbook", ignore=["**/__pycache__"]) + ) diff --git a/cookbook/common/sidecar.py b/cookbook/common/sidecar.py new file mode 100644 index 0000000..649fd4a --- /dev/null +++ b/cookbook/common/sidecar.py @@ -0,0 +1,48 @@ +"""The shared Server-side entrypoint: build the Store + Engine and run the versioned +rollout proxy in front of the local sglang. One entrypoint for every recipe. + +The Server container launches this as a subprocess (see common/process.py): + python -m cookbook.common.sidecar --port 8000 --upstream http://127.0.0.1:8001 ... +Defaults come from the serving container's env, so the flags are optional. +""" + +from __future__ import annotations + +import argparse +import os + +from stitch.engines.sglang import SGLangEngine +from stitch.service import serve +from stitch.stores.modal_volume import ModalVolumeStore + + +def main() -> None: + p = argparse.ArgumentParser() + p.add_argument("--host", default="0.0.0.0") + p.add_argument("--port", type=int, default=8000) + p.add_argument("--upstream", default="http://127.0.0.1:8001") + p.add_argument("--bulletin-root", default=os.environ.get("DELTA_BULLETIN_ROOT", "/delta-bulletin")) + p.add_argument("--volume-name", default=os.environ.get("DELTA_VOLUME_NAME", "")) + p.add_argument("--local-checkpoint-dir", default=os.environ.get("STITCH_LOCAL_CHECKPOINT_DIR", "/local-checkpoint")) + p.add_argument("--commit-mode", choices=["quiesce", "in_place"], default=os.environ.get("SIDECAR_COMMIT_MODE", "quiesce")) + p.add_argument("--run-id", default=os.environ.get("DISAGG_RUN_ID") or None) + p.add_argument("--debug-requests", action="store_true") + # Background convergence backstop: re-check latest every N seconds so a replica that + # missed its wake (cold start racing the last publish, or a lost best-effort wake) still + # catches up. 0 disables it (wake + 409-self-heal only). + p.add_argument("--reconcile-interval", type=float, + default=float(os.environ.get("STITCH_RECONCILE_INTERVAL", "5.0"))) + args = p.parse_args() + + store = ModalVolumeStore(args.bulletin_root, volume_name=args.volume_name or None) + engine = SGLangEngine(args.upstream, args.local_checkpoint_dir) + serve( + store, engine, + run_id=args.run_id, commit_mode=args.commit_mode, + host=args.host, port=args.port, debug_requests=args.debug_requests, + reconcile_interval=args.reconcile_interval, + ) + + +if __name__ == "__main__": + main() diff --git a/cookbook/common/smoke.py b/cookbook/common/smoke.py new file mode 100644 index 0000000..19e3818 --- /dev/null +++ b/cookbook/common/smoke.py @@ -0,0 +1,102 @@ +"""Flash-pool smoke check — shared by every recipe. Confirms the pool serves +completions at an expected weight version, through the gateway and each replica.""" + +from __future__ import annotations + +import json +import time +import urllib.request + +from stitch.pools.modal_flash import ModalFlashPool +from stitch.versions import VersionRef + + +class VersionAheadError(RuntimeError): + """A monotonic pool has already advanced past the smoke's expected version.""" + + +def smoke_flash_pool(*, app_name: str, cls_name: str, model_name: str, weight_version: int, timeout_seconds: int) -> None: + """Poll until the pool serves completions at ``weight_version`` — through the gateway + (which also wakes a scaled-down pool; Flash holds the request through the cold start) + and then each live replica's ``/server_info``.""" + pool = ModalFlashPool(app_name, cls_name) + deadline = time.time() + timeout_seconds + last_error: str | None = None + while True: + try: + gateway = pool.gateway_url() + print(f"Gateway URL: {gateway}") + # A fresh pool serves the base but has no claimed run, so version 0 (run-scoped) + # is unpinnable — an exact-version request would 409. Pre-claim, gate on plain + # serving; the version check only makes sense once a run has claimed the pool. + if _get_json(f"{gateway}/server_info", timeout=60).get("run_id") is None: + data = _post_json(f"{gateway}/v1/chat/completions", _completion(model_name), timeout=900) + _check_serves(data) + print(f"Pool serves base (unclaimed): {data.get('choices')}") + return + data = _post_json(f"{gateway}/v1/chat/completions", _completion(model_name, weight_version), timeout=900) + print(f"Gateway completion: {data}") + _check_completion(data, weight_version) + for target in pool.discover_replicas(): + info = _get_json(f"{target}/server_info", timeout=30) + print(f"{target} server_info={info}") + _check_version(_applied_version(info), weight_version, target) + 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 _applied_version(info: dict) -> int: + applied = info.get("applied") + return VersionRef.parse(applied).version if applied else -1 + + +def _check_version(current: int, expected: int, target: str) -> None: + if current > expected: + raise VersionAheadError(f"{target} applied={current} already past expected {expected}") + if current != expected: + raise RuntimeError(f"{target} applied={current}, expected {expected}") + + +def _check_completion(data: dict, expected: int) -> None: + 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 {start}->{end}, already past expected {expected}") + if start != expected or end != expected: + raise RuntimeError(f"unexpected gateway weight metadata: {data}") + + +def _check_serves(data: dict) -> None: + choices = data.get("choices") or [] + if not choices or not ((choices[0].get("message") or {}).get("content")): + raise RuntimeError(f"pool did not return a completion: {data}") + + +def _completion(model_name: str, expected: int | None = None) -> dict: + payload = { + "model": model_name, + "messages": [{"role": "user", "content": "Reply with exactly OK."}], + "max_tokens": 8, + "temperature": 0, + "chat_template_kwargs": {"enable_thinking": False}, + } + if expected is not None: # pin the version only against a claimed pool + payload["weight_version"] = {"exact_version": expected} + return payload + + +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) diff --git a/cookbook/miles_disagg/README.md b/cookbook/miles_disagg/README.md deleted file mode 100644 index 7567667..0000000 --- a/cookbook/miles_disagg/README.md +++ /dev/null @@ -1,192 +0,0 @@ -# Disaggregated miles on Modal - -Train Kimi K2.6 or Moonlight with GRPO under **native NVFP4 QAT**, or run -GLM-4.5-Air as a BF16 H200 experiment, while rollouts run on a separate, -elastic pool of SGLang servers. The miles twin of -[`../slime_disagg`](../slime_disagg): same two-half architecture and the same -`stitch` bulletin-board / sidecar machinery, but the trainer is **miles**. - -The app has two halves. - -- **`Trainer`** is a clustered miles/Ray job. For NVFP4 configs it QAT-trains - with `--fp4-format e2m1` / `NVFP4BlockScaling`; for GLM-4.5-Air it trains and - exports BF16. After each step it writes a sparse XOR weight delta to a Modal - Volume (the "bulletin board") and publishes a new weight version. -- **`Server`** is a Modal Flash pool of SGLang servers on the configured GPU - type: B200 for NVFP4 configs, H200 for GLM-4.5-Air BF16. A sidecar in each - container watches the bulletin board, applies deltas in order, and serves - rollouts pinned to an exact weight version. Requests for a version a container - hasn't reached get `409` and miles retries. - -Weights flow through the Volume; rollout traffic flows through the Flash -gateway; either side scales or restarts on its own. - -## NVFP4 is all-Blackwell - -NVFP4 QAT is native Megatron FP4 training (`NVFP4BlockScaling`, TransformerEngine -≥ 2.7.0.dev0). FP4 GEMM requires Blackwell, so **both** the trainer and the -rollout pool run on B200 — unlike the INT4 recipe, which fake-quantizes on H200. -There is no simulated/non-Blackwell NVFP4 weight-QAT path. - -`glm45_air_bf16_disagg` is the exception: it is a BF16 experiment on H200. It -does not use NVFP4 QAT, does not run `convert_hf_to_nvfp4.py`, and serves the -prepared BF16 Hugging Face checkpoint directly. - -## Checkpoint lifecycle (three roles) - -`prepare_checkpoints` builds them on a GPU (see `modal_train.py`): - -1. **BF16 masters** (`--ref-load`): the trainable parameters. Moonlight ships - bf16 (masters = the download); Kimi K2.6 ships INT4, so masters are - dequantized with `tools/convert_kimi_int4_to_bf16.py`. -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. -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 -baseline, so applying delta_vN reproduces export_vN byte-for-byte — the served -weights become the trainer's NVFP4 export. - -## GLM-4.5-Air BF16 on H200 - -Use `glm45_air_bf16_disagg` for `zai-org/GLM-4.5-Air`. - -This config has two prepared checkpoint paths: - -- `/prep/glm45-air-bf16/bf16`: the prepared Hugging Face BF16 checkpoint. This - is both the SGLang served base (`--hf-checkpoint`) and the disk-delta baseline. -- `/prep/glm45-air-bf16/torch_dist`: the Megatron raw-mode checkpoint loaded by - the trainer (`--ref-load`). - -The weight-sync loop is still the same disk-delta loop: miles exports BF16 HF -tensors after each update, XORs the new bytes against the previous bytes, writes -the delta to the Modal Volume bulletin board, and the stitch sidecar applies the -delta onto each rollout container's local BF16 checkpoint copy. - -The GLM path has two Modal-specific details: - -- `prepare_checkpoints` disables Xet and `hf_transfer`; the standard Hugging - Face downloader was the path that finished reliably for this large checkpoint. -- `prepare_torch_dist` uses a small wrapper around miles' - `convert_hf_to_torch_dist.py` so multi-node Modal Volume commits merge all - `iter_0000001` shard files instead of only rank 0's renamed `release` dir. - -Run from the repo root: - -```bash -alias m="uv run --extra modal modal" -export EXPERIMENT_CONFIG=glm45_air_bf16_disagg - -# Long-running one-time prep. Keep the rollout pool down until the served base -# exists, otherwise warm containers crash-loop on the missing model path. If you -# use --detach, wait for each prep job to finish before starting the next one. -POOL_MIN_CONTAINERS=0 m run --detach -m cookbook.miles_disagg.modal_train::prepare_checkpoints -POOL_MIN_CONTAINERS=0 m run --detach -m cookbook.miles_disagg.modal_train::prepare_torch_dist -POOL_MIN_CONTAINERS=0 m run -m cookbook.miles_disagg.modal_train::prepare_dataset - -# Deploy the H200 rollout pool and trainer app. -m deploy --strategy recreate -m cookbook.miles_disagg.modal_train - -# Verify SGLang serves the prepared BF16 base, then launch training. -m run -m cookbook.miles_disagg.modal_train::smoke_flash_pool -m run -m cookbook.miles_disagg.modal_train::launch_train - -# Optional: check a later synced weight version. -m run -m cookbook.miles_disagg.modal_train::smoke_flash_pool --weight-version 1 -``` - -Expected prepared outputs in the `miles-prep-checkpoints` Volume: - -```text -glm45-air-bf16/bf16/ -glm45-air-bf16/torch_dist/latest_checkpointed_iteration.txt -glm45-air-bf16/torch_dist/iter_0000001/ -``` - -## Run it - -You need a Modal account and a `huggingface-secret` Modal secret. Work from the -repo root, with this alias: - -```bash -alias m="uv run --extra modal modal" - -# Start with the small, runnable de-risk (single Blackwell node). -export EXPERIMENT_CONFIG=moonlight_nvfp4_disagg - -# One-time setup: fetch + convert checkpoints (GPU), and the dataset. -m run -m cookbook.miles_disagg.modal_train::prepare_checkpoints -m run -m cookbook.miles_disagg.modal_train::prepare_dataset - -# Deploy the rollout pool + trainer. -m deploy -m cookbook.miles_disagg.modal_train - -# Wait for the pool to come up and answer at version 0. -m run -m cookbook.miles_disagg.modal_train::smoke_flash_pool - -# Train. Returns immediately; the run continues on Modal. -m run -m cookbook.miles_disagg.modal_train::launch_train - -# Smoke-check the pool at a given version (the chain advances one per rollout). -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 -QAT → NVFP4-export → XOR-delta → SGLang-reload loop, then scale. - -### Fork dependencies - -The image pins the miles fork branch `nvfp4-disagg-fixes` (`MILES_REPO_REF`), -which carries the disaggregated-rollout features plus the publish-only / NVFP4 -fixes this cookbook needs: NVFP4 export dispatch -(`megatron_to_hf/processors/__init__.py`), the publish-only rollout semaphore -and HTTP client, the 0-dim NVFP4-scale delta encode, and the `encoding_dsv4` -import guard. Push that branch before deploying. - -The megatron routing-replay (R3) fix lives in `radixark/Megatron-LM` and is -**baked into the trainer image at build time** (a `.run_commands` step in -`modal_train.py`; source diff in `megatron_r3_num_out_tokens.patch`). The bake -is idempotent — it becomes a no-op once the fork itself ships the fix. - -Dev iteration: overlay a local miles checkout at deploy time (no rebuild, no -push). This only overlays miles; the megatron R3 fix still comes from the bake. - -```bash -MILES_LOCAL_DIR=/path/to/miles EXPERIMENT_CONFIG=moonlight_nvfp4_disagg \ - m deploy --strategy recreate -m cookbook.miles_disagg.modal_train -``` - -## Configuration - -Each experiment is a module in `configs/` holding a `ModalConfig` (GPU type, -pool size, regions), a `MilesConfig` (every miles/Megatron CLI arg), and the -names of the Modal resources it owns. `launch_train` ships the training args as -plain data, so editing or adding a `MilesConfig` needs no redeploy; only -infrastructure (GPU, nodes, pool size, Volume names, serving flags) does. - -```bash -EXPERIMENT_CONFIG= m deploy --strategy recreate -m cookbook.miles_disagg.modal_train -``` - -Useful log searches (the app name is the config's `APP_NAME`): - -```bash -m app logs --since 4h --search "passrate " -m app logs --since 4h --search "weight_v" -``` - -## Bring-up checklist (flagged, validated by the Moonlight run) - -- The miles image's TransformerEngine is ≥ 2.7.0.dev0 and the trainer runs on - Blackwell (NVFP4 BlockScaling). -- SGLang serves the prepared NVFP4 base on Blackwell (the `serving.py` fork is - proven for NVFP4) — verify on a warm container. -- The `convert_hf_to_nvfp4.py` quantization scope (which tensors get NVFP4 + - exclude rules) matches the export processor's scope, so the XOR delta aligns. - A scope mismatch fails loud on the first delta apply (checksum/shape) — which - is exactly what the Moonlight run catches cheaply. -- `miles.utils.disk_delta` is import-light in the `--no-deps` serving image. diff --git a/cookbook/miles_disagg/_spawn_into_deployed.py b/cookbook/miles_disagg/_spawn_into_deployed.py deleted file mode 100644 index 4f31e78..0000000 --- a/cookbook/miles_disagg/_spawn_into_deployed.py +++ /dev/null @@ -1,24 +0,0 @@ -"""Spawn Trainer.train on the already-deployed app via the Modal SDK. - -This is the non-client-tied equivalent of the launch_train local_entrypoint: -it does a fire-and-forget `Cls.from_name(APP_NAME, "Trainer").train.spawn(...)` -against the DEPLOYED app, so it never creates an ephemeral app context (which is -what collides with / stops the deployed serving app). Run with plain `python`, -NOT `modal run`. -""" - -import importlib -import sys - -import modal - - -def main(experiment: str = "kimi_k2_6_nvfp4_disagg") -> None: - run = importlib.import_module(f"cookbook.miles_disagg.configs.{experiment}") - trainer = modal.Cls.from_name(run.APP_NAME, "Trainer")() - call = trainer.train.spawn(experiment, run.miles.to_payload()) - print(f"Spawned train({experiment!r}) on {run.APP_NAME}: {call.object_id}") - - -if __name__ == "__main__": - main(*sys.argv[1:]) diff --git a/cookbook/miles_disagg/app.py b/cookbook/miles_disagg/app.py new file mode 100644 index 0000000..842d8f1 --- /dev/null +++ b/cookbook/miles_disagg/app.py @@ -0,0 +1,319 @@ +"""Disaggregated miles training on Modal, assembled on the stitch core. + +``EXPERIMENT_CONFIG`` selects a config module under ``cookbook.miles_disagg.configs``. The +Server (sglang + stitch sidecar) is the shared common one; the Trainer runs miles on Ray +and publishes XOR deltas through a Modal Volume the pool syncs from. + + EXPERIMENT_CONFIG=glm45_air_fp8 uv run --extra modal modal deploy -m cookbook.miles_disagg.app + +Config access is uniform: the experiment module ``exp`` is the single source of truth — +its ``exp.modal`` (infra), ``exp.miles`` (training), and ``exp.`` are read directly; +shared deployment constants come from ``common.constants``. ``ROLLOUT_CONCURRENCY`` is the +one resolved value (the experiment's Flash target, else the engine's concurrency). +""" + +from __future__ import annotations + +import importlib +import os +import subprocess +import tempfile +import uuid +from types import SimpleNamespace +from typing import Any + +import modal +import modal.experimental + +from stitch.pools.modal_flash import ModalFlashPool + +from cookbook.common import launch, ray_cluster, serving_image, server, smoke +from cookbook.common.constants import ( + CHECKPOINTS_PATH, DATA_PATH, HF_CACHE_PATH, MINUTES, PREP_PATH, RAY_PORT, + SERVER_STARTUP_TIMEOUT, SGLANG_CACHE_PATH, SIDECAR_PORT, +) +from cookbook.miles_disagg import prep, trainer_image +from cookbook.miles_disagg.config import MilesConfig, YAML_CONFIG_FIELDS +from cookbook.miles_disagg.trainer_image import MEGATRON_PATH, MILES_ROOT + +EXPERIMENT = os.environ["EXPERIMENT_CONFIG"] # required; a default would silently serve the wrong experiment +MILES_LOCAL_DIR = os.environ.get("MILES_LOCAL_DIR") # optional dev overlay of a local miles checkout + +exp = importlib.import_module(f"cookbook.miles_disagg.configs.{EXPERIMENT}") +modal_cfg = exp.modal +miles_cfg = exp.miles + +# The Flash autoscaler target (and sglang concurrency cap): the experiment's explicit +# target_inputs, else the engine's configured concurrency. +ROLLOUT_CONCURRENCY = modal_cfg.rollout_target_inputs or miles_cfg.sglang_server_concurrency + +# EXPERIMENT_CONFIG is baked into both images (inside build_*_image) so the container's +# re-import resolves the same experiment as the deploy, not the default. An experiment +# may override the trainer stack (image tag / miles ref / setup) and extend the shared +# serving image (extra commands / env) — see the 4/6 recipe config. +image = trainer_image.build_trainer_image( + hf_cache_path=str(HF_CACHE_PATH), experiment=EXPERIMENT, miles_local=MILES_LOCAL_DIR, + image_tag=getattr(exp, "MILES_IMAGE_TAG", None), + miles_repo_ref=getattr(exp, "MILES_REPO_REF", None), + setup_commands=tuple(getattr(exp, "TRAINER_IMAGE_SETUP", ())), +) +server_image = serving_image.build_serving_image( + hf_cache_path=str(HF_CACHE_PATH), delta_volume_name=exp.DELTA_VOLUME_NAME, experiment=EXPERIMENT, + extra_commands=tuple(getattr(exp, "SERVING_IMAGE_EXTRA_COMMANDS", ())), + extra_env=getattr(exp, "SERVING_IMAGE_ENV", None), +) +if MILES_LOCAL_DIR: + server_image = server_image.add_local_dir(MILES_LOCAL_DIR, remote_path=MILES_ROOT, ignore=[".git", "**/__pycache__", "**/*.pyc"]) + +def _volume(name: str) -> modal.Volume: + """Volume v2 by default, tolerating environments where the name already exists as v1. + + ``from_name`` is lazy, so the version-mismatch error would otherwise only surface at + deploy/run; hydrating here resolves it eagerly and falls back to the volume's actual + version. Fresh environments always get v2.""" + vol = modal.Volume.from_name(name, create_if_missing=True, version=2) + try: + vol.hydrate() + except modal.exception.InvalidError: + vol = modal.Volume.from_name(name, create_if_missing=True) + vol.hydrate() + return vol + + +hf_cache_volume = _volume("huggingface-cache") +data_volume = _volume("miles-data") +checkpoints_volume = _volume("miles-checkpoints") +prep_volume = _volume("miles-prep-checkpoints") +sglang_cache_volume = _volume("miles-sglang-cache") # survives cold starts +delta_volume = _volume(exp.DELTA_VOLUME_NAME) + +train_volumes = { + str(HF_CACHE_PATH): hf_cache_volume, + str(DATA_PATH): data_volume, + str(CHECKPOINTS_PATH): checkpoints_volume, + str(PREP_PATH): prep_volume, + exp.DELTA_BULLETIN_ROOT: delta_volume, +} + +app = modal.App(exp.APP_NAME) + +SGLANG_SERVER_ARGS = { + "--served-model-name": miles_cfg.hf_checkpoint, + "--cuda-graph-max-bs-decode": str(ROLLOUT_CONCURRENCY), + "--max-running-requests": str(ROLLOUT_CONCURRENCY), + "--trust-remote-code": "", + # The engine reads published versions off the Modal Volume; this hook reloads it + # first (object stores lack cross-host read-after-write consistency). + "--custom-pull-weights-pre-read-hook": "stitch.stores.modal_volume.pull_weights_pre_read_hook", + **exp.SGLANG_SERVER_ARGS, +} + + +# The rollout Server: a thin module-level class (Modal requires @app.cls at global scope) +# whose enter/exit delegate to the shared common.server logic. sglang + the stitch sidecar. +@app.cls( + image=server_image, + gpu=f"{modal_cfg.gpu}:{miles_cfg.rollout_num_gpus_per_engine}", + cloud=modal_cfg.cloud, region=modal_cfg.region, + volumes={ + str(HF_CACHE_PATH): hf_cache_volume, + str(PREP_PATH): prep_volume, + SGLANG_CACHE_PATH: sglang_cache_volume, + exp.DELTA_BULLETIN_ROOT: delta_volume, + }, + min_containers=modal_cfg.rollout_min_containers, max_containers=modal_cfg.rollout_max_containers, + timeout=40 * MINUTES, scaledown_window=15 * MINUTES, + ephemeral_disk=modal_cfg.rollout_ephemeral_disk_mib, memory=modal_cfg.rollout_memory_mib, + include_source=False, +) +@modal.experimental.http_server( + port=SIDECAR_PORT, proxy_regions=modal_cfg.proxy_regions, + exit_grace_period=25, startup_timeout=SERVER_STARTUP_TIMEOUT, +) +@modal.concurrent(target_inputs=ROLLOUT_CONCURRENCY) +class Server: + @modal.enter() + def startup(self) -> None: + server.serve_startup( + self, model_name=miles_cfg.hf_checkpoint, sglang_args=SGLANG_SERVER_ARGS, + tp=miles_cfg.rollout_num_gpus_per_engine, concurrency=ROLLOUT_CONCURRENCY, + bulletin_root=exp.DELTA_BULLETIN_ROOT, local_checkpoint_dir=exp.LOCAL_CHECKPOINT_PATH, + volume_name=exp.DELTA_VOLUME_NAME, commit_mode=exp.SIDECAR_COMMIT_MODE, + startup_timeout=SERVER_STARTUP_TIMEOUT, sglang_env=getattr(exp, "SGLANG_ENV", {}), + ) + + @modal.exit() + def stop(self) -> None: + server.serve_stop(self) + + +# ── Trainer (miles on Ray) ──────────────────────────────────────────────────── +# Multi-node needs an RDMA gang (clustered) over the EFA fabric; single-node takes +# neither. Both are inline on the decorator so there's one declaration, not a rebind. +_MULTINODE = miles_cfg.n_train_nodes > 1 + + +@app.cls( + image=image, + gpu=f"{modal_cfg.gpu}:{miles_cfg.actor_num_gpus_per_node}", + memory=modal_cfg.memory, + cloud=modal_cfg.cloud, region=modal_cfg.region, + volumes=train_volumes, + ephemeral_disk=modal_cfg.trainer_ephemeral_disk_mib, + timeout=24 * 60 * MINUTES, startup_timeout=20 * MINUTES, scaledown_window=30 * MINUTES, + include_source=False, + **({"experimental_options": {"efa_enabled": True}} if _MULTINODE else {}), +) +@(modal.experimental.clustered(miles_cfg.n_train_nodes, rdma=True) if _MULTINODE else lambda c: c) +class Trainer: + """miles actor cluster. Ray comes up once per container in enter(), so back-to-back + runs reuse it.""" + + @modal.enter() + def start_ray(self) -> None: + from cookbook.common import process + + rank, master_addr, my_ip = ray_cluster.get_modal_cluster_context(miles_cfg.n_train_nodes) + process.apply_git_patches(list(getattr(exp, "MEGATRON_RUNTIME_PATCHES", [])), MEGATRON_PATH, "Megatron patch") + self.rank = rank + process.start_host_mem_monitor() # per-node host-RAM trace (publish gather is the OOM peak) + os.environ.update({ + "MILES_HOST_IP": my_ip, "SGLANG_HOST_IP": my_ip, "HOST_IP": my_ip, + "MASTER_ADDR": master_addr, "RAY_ADDRESS": f"{master_addr}:{RAY_PORT}", + "no_proxy": f"127.0.0.1,{master_addr},{my_ip}", "NO_PROXY": f"127.0.0.1,{master_addr},{my_ip}", + "PYTHONPATH": f"{MEGATRON_PATH}:{os.environ.get('PYTHONPATH', '')}", # source-only megatron.training + **miles_cfg.environment, + }) + if rank == 0: + ray_cluster.start_ray_head(my_ip, miles_cfg.n_train_nodes, ray_port=RAY_PORT) + else: + ray_cluster.start_ray_worker(my_ip, master_addr, ray_port=RAY_PORT) + + @modal.method() + def train(self, payload: dict) -> None: + """Run one training job from a MilesConfig payload (see MilesConfig.to_payload).""" + for volume in train_volumes.values(): + volume.reload() + + cfg = MilesConfig.from_payload(payload) + _materialize_node_local_yaml(cfg, "te_precision_config_file") # re-read per Ray actor; needs identical local path + if self.rank != 0: + return + + cfg.rollout_endpoint_url = ModalFlashPool(exp.APP_NAME, "Server").gateway_url() + run_id = uuid.uuid4().hex[:12] # per-launch fence token; forks a fresh chain + cfg.update_weight_disk_dir = f"{exp.DELTA_BULLETIN_ROOT}/{run_id}" + if getattr(cfg, "save_interval", None) is None: + cfg.load = cfg.save = cfg.save_hf = None + else: + cfg.load = cfg.save = f"{CHECKPOINTS_PATH}/{run_id}/checkpoints" + # Merge the run's bulletin identity into custom_config_path (already carrying the + # request-gating knobs); miles setattr's every key onto args for the hooks. + custom_config = { + **(cfg.custom_config_path or {}), + "update_weight_delta_volume_name": exp.DELTA_VOLUME_NAME, + "rollout_modal_flash_app_name": exp.APP_NAME, + "rollout_modal_flash_server_cls_name": "Server", + "run_id": run_id, + } + cfg.custom_config_path = custom_config + launch.resolve_config(cfg, tempfile.mkdtemp(), YAML_CONFIG_FIELDS) + cmd = launch.build_train_cmd(cfg, MILES_ROOT, "miles_model_script") + + # Claim the pool before miles publishes: reset every replica to base for this run. + from cookbook.common import hooks + + hooks.claim_pool(SimpleNamespace(update_weight_disk_dir=cfg.update_weight_disk_dir, **custom_config)) + + print(f"Training {EXPERIMENT}: nodes={miles_cfg.n_train_nodes}, rollout_endpoint={cfg.rollout_endpoint_url}") + print(f"Command: {cmd}") + log_path = f"{CHECKPOINTS_PATH}/{run_id}/train.log" + os.makedirs(os.path.dirname(log_path), exist_ok=True) + teed = f"set -o pipefail; ({cmd}) 2>&1 | tee {log_path}" # tee to a committed log; survives the app-logs buffer + try: + subprocess.run(["bash", "-lc", teed], check=True) + finally: + try: + checkpoints_volume.commit() + print(f"Train log committed to miles-checkpoints at {run_id}/train.log") + except Exception as exc: # noqa: BLE001 + print(f"WARNING: could not commit train log: {exc}") + + +# ── miles launch helper (te_precision_config_file is miles-only) ────────────────── +def _materialize_node_local_yaml(cfg: Any, field: str, dest_dir: str = "/root/.miles_node_yaml") -> None: + """Write an inline YAML config to a deterministic node-local path on EVERY node. + Fields like te_precision_config_file are re-read on each Ray actor, so they must + resolve to identical content at an identical path on all nodes.""" + 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) + + +# ── Preparation + entrypoints ────────────────────────────────────────────────── +@app.function( + image=image, gpu=f"{modal_cfg.gpu}:1", + volumes={str(HF_CACHE_PATH): hf_cache_volume, str(PREP_PATH): prep_volume}, + timeout=6 * 60 * MINUTES, secrets=[modal.Secret.from_name("huggingface-secret")], include_source=False, +) +def prepare_checkpoints() -> None: + prep.prepare_checkpoints(exp, prep_volume) + + +# torch_dist conversion is clustered across nodes (a large MoE won't fit an 8-way split). +_TORCH_DIST_MULTINODE = modal_cfg.torch_dist_prep_nodes > 1 + + +@app.function( + image=image, gpu=f"{modal_cfg.gpu}:{modal_cfg.torch_dist_prep_gpus_per_node}", + volumes={str(HF_CACHE_PATH): hf_cache_volume, str(PREP_PATH): prep_volume}, + timeout=6 * 60 * MINUTES, + ephemeral_disk=(modal_cfg.torch_dist_prep_ephemeral_disk_mib or modal_cfg.rollout_ephemeral_disk_mib), + secrets=[modal.Secret.from_name("huggingface-secret")], include_source=False, + **({"experimental_options": {"efa_enabled": True}} if _TORCH_DIST_MULTINODE else {}), +) +@(modal.experimental.clustered(modal_cfg.torch_dist_prep_nodes, rdma=True) if _TORCH_DIST_MULTINODE else lambda fn: fn) +def prepare_torch_dist() -> None: + rank, master_addr, _ = ray_cluster.get_modal_cluster_context(modal_cfg.torch_dist_prep_nodes) + prep.prepare_torch_dist(exp, prep_volume, rank=rank, master_addr=master_addr) + + +@app.function( + image=image, volumes={str(DATA_PATH): data_volume}, + timeout=2 * 60 * MINUTES, secrets=[modal.Secret.from_name("huggingface-secret")], include_source=False, +) +def prepare_dataset() -> None: + data_volume.reload() + miles_cfg.prepare_data() + data_volume.commit() + + +@app.local_entrypoint() +def launch_train() -> None: + """Spawn training on the deployed app. Config ships as data, so edits run without a + redeploy; infrastructure changes still require one.""" + from modal.exception import NotFoundError + + try: + trainer = modal.Cls.from_name(exp.APP_NAME, "Trainer")() + call = trainer.train.spawn(miles_cfg.to_payload()) + except NotFoundError: + raise SystemExit( + f"App {exp.APP_NAME!r} is not deployed. Run:\n" + f" EXPERIMENT_CONFIG={EXPERIMENT} uv run --extra modal modal deploy -m cookbook.miles_disagg.app" + ) + print(f"Spawned train on {exp.APP_NAME}: {call.object_id}") + + +@app.local_entrypoint() +def smoke_flash_pool(weight_version: int = 0, timeout_seconds: int = 30 * MINUTES) -> None: + """Check the deployed Flash pool serves completions at the expected weight version.""" + smoke.smoke_flash_pool( + app_name=exp.APP_NAME, cls_name="Server", model_name=miles_cfg.hf_checkpoint, + weight_version=weight_version, timeout_seconds=timeout_seconds, + ) diff --git a/cookbook/miles_disagg/config.py b/cookbook/miles_disagg/config.py new file mode 100644 index 0000000..0bac52f --- /dev/null +++ b/cookbook/miles_disagg/config.py @@ -0,0 +1,95 @@ +"""``MilesConfig`` — miles training arguments as a reflected config (self-contained). + +Every public, non-callable attribute becomes a miles CLI arg via ``cli_args`` (miles +wraps Megatron's parser, so Megatron args pass straight through); ``environment`` / +``async_mode`` / ``miles_model_script`` are launcher instructions, not CLI args. The +Modal-infra half of an experiment is ``common.config.ModalConfig``. +""" + +from __future__ import annotations + +from typing import Any + +# Fields that are launcher instructions, not miles CLI args. +_MILES_SKIP = {"environment", "async_mode", "miles_model_script"} +# Fields miles reads as YAML files; inline dicts are materialized before launch. +# (te_precision_config_file is handled separately in app.py — it needs an identical +# node-local path on every Ray actor, not a per-launch tmpdir.) +YAML_CONFIG_FIELDS = ("eval_config", "custom_config_path", "sglang_config") + + +class MilesConfig: + """Subclass and set class attributes; all public, non-callable, non-skip attributes + become miles CLI args via ``cli_args``.""" + + environment: dict = {} + async_mode: bool = False # True -> train_async.py + miles_model_script: str = "" # shell script (relative to the miles root) defining MODEL_ARGS + + def __init__(self, **kwargs: Any) -> None: + self.environment = dict(type(self).environment) # fresh per instance; never mutate the class default + for k, v in kwargs.items(): + setattr(self, k, v) + + @property + def n_train_nodes(self) -> int: + """Trainer node count: actor nodes, plus critic nodes for PPO/critic setups.""" + nodes = int(getattr(self, "actor_num_nodes", 1)) + if getattr(self, "use_critic", False) or getattr(self, "advantage_estimator", None) == "ppo": + nodes += int(getattr(self, "critic_num_nodes", nodes)) + return nodes + + def _fields(self) -> dict[str, Any]: + """Merged fields across the class hierarchy; instance attrs win.""" + fields: dict[str, Any] = {} + for cls in reversed(type(self).__mro__): + if cls is object: + continue + fields.update( + { + k: v + for k, v in vars(cls).items() + if not k.startswith("_") + and not callable(v) + and not isinstance(v, (classmethod, staticmethod, property)) + } + ) + fields.update(vars(self)) + return {k: v for k, v in fields.items() if k not in _MILES_SKIP} + + def cli_args(self) -> list[str]: + """miles CLI args: field_name -> --field-name; True -> bare flag; False/None -> + omitted; list -> --flag v1 v2; else --flag value.""" + out: list[str] = [] + for key, val in self._fields().items(): + if val is None or val is False: + continue + flag = f"--{key.replace('_', '-')}" + if val is True: + out.append(flag) + elif isinstance(val, list): + out += [flag] + [str(v) for v in val] + else: + out += [flag, str(val)] + return out + + def prepare_data(self) -> None: + raise NotImplementedError(f"{type(self).__name__} has no prepare_data()") + + def to_payload(self) -> dict[str, Any]: + """Flatten to plain data so the launcher can ship a config to the deployed Trainer + — new or edited experiments run without a redeploy.""" + return { + "fields": self._fields(), + "environment": dict(self.environment), + "async_mode": self.async_mode, + "miles_model_script": self.miles_model_script, + } + + @classmethod + def from_payload(cls, payload: dict[str, Any]) -> "MilesConfig": + cfg = cls(**payload["fields"]) + cfg.environment = dict(payload["environment"]) + cfg.async_mode = payload["async_mode"] + cfg.miles_model_script = payload["miles_model_script"] + return cfg diff --git a/cookbook/miles_disagg/configs/base.py b/cookbook/miles_disagg/configs/base.py deleted file mode 100644 index d19882f..0000000 --- a/cookbook/miles_disagg/configs/base.py +++ /dev/null @@ -1,189 +0,0 @@ -"""Base configuration classes and volume mount paths for the miles example. - -Two separate concerns: - - ModalConfig — Modal infrastructure (GPU model, regions, rollout pool size) - MilesConfig — miles training arguments - -Each experiment module defines one instance of each, plus a handful of -module-level constants (APP_NAME, DELTA_VOLUME_NAME, ...) that name the -Modal resources the experiment owns. All non-private, non-callable -attributes on a MilesConfig subclass become miles CLI args automatically -via cli_args() (miles wraps Megatron's parser, so Megatron args like -``--fp4-format`` pass straight through). The 'environment' field is the only -exception — it is exported into the Ray runtime environment, not passed to -miles directly. - -This is the miles twin of cookbook/slime_disagg/configs/base.py. The -reflection mechanism is identical; only the names and the skip set differ -(miles uses ``miles_model_script``). -""" - -from pathlib import Path -from typing import Any, Literal - -# ── Volume mount paths ──────────────────────────────────────────────────────── - -HF_CACHE_PATH = Path("/root/.cache/huggingface") -DATA_PATH = Path("/data") -CHECKPOINTS_PATH = Path("/checkpoints") -# Derived checkpoints the prepare step builds on a GPU (see modal_train): -# /-bf16 bf16 masters (dequantized from the published INT4) -# /-nvfp4 served NVFP4 base (miles convert_hf_to_nvfp4 of the masters) -PREP_PATH = Path("/prep") - -# ── Types ───────────────────────────────────────────────────────────────────── - -GPUType = Literal["H100", "H200", "B200", "B300", "A100"] - -# Fields on MilesConfig that are NOT miles CLI args. -_MILES_SKIP = {"environment", "async_mode", "miles_model_script"} - -# MilesConfig fields that miles reads as YAML files at runtime. Experiments may -# set these as inline dicts; the launcher materializes them to temp YAML files -# before building the CLI command. miles' --custom-config-path setattr's every -# key in the YAML onto the args namespace (arguments.py), which is how the -# bulletin/gating knobs below reach the hooks. -YAML_CONFIG_FIELDS = ("eval_config", "custom_config_path", "sglang_config", "te_precision_config_file") - - -class ModalConfig: - """Modal infrastructure configuration.""" - - gpu: GPUType = "B200" # NVFP4 QAT + NVFP4 serving are both Blackwell-only - memory: tuple[int, int] | None = None # per-container memory in MiB (request, limit) - cloud: str | None = None # e.g. "aws", "gcp" - region: str | None = None # e.g. "us-east-2" - rollout_min_containers: int = 2 # warm Flash rollout containers - # Flash autoscaler target: concurrent inputs (requests) per container before it - # scales OUT. None = use sglang_server_concurrency (legacy). Set it well below the - # SGLang engine concurrency so Flash adds containers instead of packing requests - # onto a few until their KV cache saturates (which 502'd / stalled the rollout). - rollout_target_inputs: int | None = None - proxy_regions: list[str] = ["us-west"] # Flash gateway proxy regions - # Ephemeral disk (MiB) for the rollout Server. The sidecar materializes a - # writable local copy of the served base (~full model size) onto ephemeral - # disk; large models exceed Modal's default. None = Modal default (fine for - # small bases like Moonlight). - rollout_ephemeral_disk_mib: int | None = None - # Nodes x GPUs/node for the prepare_torch_dist conversion. The full 1T K2.6 - # needs >=2 nodes (8-way OOMs); a small proxy (e.g. 2-layer) fits 1 GPU (the - # convert auto-derives pp=world_size and asserts pp <= num_layers). - torch_dist_prep_nodes: int = 2 - torch_dist_prep_gpus_per_node: int = 8 - # Extra parallelism args for the prepare_torch_dist conversion. Large MoE needs - # explicit EP sharding so experts fit (else convert's auto-pp=world_size/EP1 OOMs). - torch_dist_convert_extra_args: str = "" - # Ephemeral disk (MiB) for the prepare_torch_dist conversion containers. Each node - # buffers its whole distcp shard set in the Volume's local write cache until commit; - # the 1T K2.6 save is ~700 GB/node, so the served-pool default (~800 GB) leaves no - # headroom and rank 0's commit (shards + .metadata + common.pt) hits ENOSPC. None = - # fall back to rollout_ephemeral_disk_mib. - torch_dist_prep_ephemeral_disk_mib: int | None = None - # Ephemeral disk (MiB) for the Trainer nodes. miles' ray.init(address="auto") uses - # default spill+log dirs under /tmp/ray, so over a multi-hour run Ray logs + object - # spill (rollout batches + per-publish full-model gathers) accumulate on the node's - # disk; Modal's default is far too small and progressively ENOSPC'd (`No space left - # on device` writing /tmp/ray/.../logs). None = Modal default. - trainer_ephemeral_disk_mib: int | None = None - - def __init__(self, **kwargs: Any) -> None: - for k, v in kwargs.items(): - setattr(self, k, v) - - -class MilesConfig: - """Base miles training configuration. - - Subclass and set class attributes to configure an experiment. All - attributes (except those in _MILES_SKIP) are forwarded to miles as CLI - args. Each experiment must be fully self-contained — no inherited defaults - beyond this base class. - - Fields in _MILES_SKIP are launcher instructions, not miles CLI args: - environment — exported into the Ray runtime environment - async_mode — selects train_async.py vs train.py - miles_model_script — path relative to the miles root to a shell script - that defines MODEL_ARGS for model architecture; - sourced before running the train command - """ - - # Launcher instructions — not passed to miles CLI (see _MILES_SKIP). - environment: dict = { - "CUDA_DEVICE_MAX_CONNECTIONS": "1", - "NCCL_NVLS_ENABLE": "1", - } - async_mode: bool = False # True → use train_async.py - miles_model_script: str = "" # shell script path relative to the miles root - - def __init__(self, **kwargs: Any) -> None: - # Fresh environment dict per instance — never mutate the class-level default. - self.environment = dict(type(self).environment) - for k, v in kwargs.items(): - setattr(self, k, v) - - def _fields(self) -> dict[str, Any]: - """Merged field dict from the class hierarchy; instance attrs win.""" - fields: dict[str, Any] = {} - for cls in reversed(type(self).__mro__): - if cls is object: - continue - fields.update( - { - k: v - for k, v in vars(cls).items() - if not k.startswith("_") - and not callable(v) - and not isinstance(v, (classmethod, staticmethod, property)) - } - ) - fields.update(vars(self)) - return {k: v for k, v in fields.items() if k not in _MILES_SKIP} - - def cli_args(self) -> list[str]: - """miles CLI arguments derived from this config. - - Conversion rules: - field_name → --field-name (underscore to hyphen) - True → --flag (no value) - False/None → omitted - list → --flag v1 v2 ... - other → --flag value - """ - out: list[str] = [] - for key, val in self._fields().items(): - if val is None or val is False: - continue - flag = f"--{key.replace('_', '-')}" - if val is True: - out.append(flag) - elif isinstance(val, list): - out += [flag] + [str(v) for v in val] - else: - out += [flag, str(val)] - return out - - def prepare_data(self) -> None: - raise NotImplementedError(f"{type(self).__name__} has no prepare_data()") - - def to_payload(self) -> dict[str, Any]: - """Flatten to plain data for sending to a deployed Trainer. - - launch_train resolves config modules locally and ships the result, so - the deployed app never imports them — new or edited experiments run - without a redeploy. - """ - return { - "fields": self._fields(), - "environment": dict(self.environment), - "async_mode": self.async_mode, - "miles_model_script": self.miles_model_script, - } - - @classmethod - def from_payload(cls, payload: dict[str, Any]) -> "MilesConfig": - cfg = cls(**payload["fields"]) - cfg.environment = dict(payload["environment"]) - cfg.async_mode = payload["async_mode"] - cfg.miles_model_script = payload["miles_model_script"] - return cfg diff --git a/cookbook/miles_disagg/configs/glm45_air_bf16_disagg.py b/cookbook/miles_disagg/configs/glm45_air_bf16.py similarity index 84% rename from cookbook/miles_disagg/configs/glm45_air_bf16_disagg.py rename to cookbook/miles_disagg/configs/glm45_air_bf16.py index 0013720..a123d97 100644 --- a/cookbook/miles_disagg/configs/glm45_air_bf16_disagg.py +++ b/cookbook/miles_disagg/configs/glm45_air_bf16.py @@ -2,40 +2,37 @@ from __future__ import annotations -from cookbook.miles_disagg.configs.base import DATA_PATH, PREP_PATH, ModalConfig, MilesConfig +from cookbook.common.config import ModalConfig +from cookbook.common.constants import DATA_PATH, PREP_PATH +from cookbook.miles_disagg.config import MilesConfig -APP_NAME = "miles-glm45-air-bf16-disagg" -DELTA_VOLUME_NAME = "miles-delta-bulletin-glm45-air-bf16" +APP_NAME = "stitch-glm45-air-bf16" +DELTA_VOLUME_NAME = "stitch-delta-glm45-air-bf16" DELTA_BULLETIN_ROOT = "/delta-bulletin" +LOCAL_CHECKPOINT_PATH = "/local-checkpoint" SOURCE_MODEL = "zai-org/GLM-4.5-Air" MODEL_TAG = "glm45-air-bf16" SERVED_CHECKPOINT_FORMAT = "bf16" -PATCH_GLM4MOE_EXPERT_BIAS_FP32 = True USE_MODAL_TORCH_DIST_WRAPPER = True -# The standard HF downloader was the path that finished reliably for this model. -DISABLE_HF_XET = True +DISABLE_HF_XET = True # the plain HF downloader is the path that finished reliably here DISABLE_HF_TRANSFER = True SIDECAR_COMMIT_MODE = "in_place" SIDECAR_DEBUG_REQUESTS = True - - -def build_serving_image(**kwargs): - from cookbook.miles_disagg.serving import build_nvfp4_b200_serving_image - - return build_nvfp4_b200_serving_image(**kwargs) +MEGATRON_RUNTIME_PATCHES = ["/root/cookbook/miles_disagg/patches/megatron-r3-dispatch.patch"] SGLANG_SERVER_ARGS = { + "--weight-loader-prefetch-checkpoints": "", + "--weight-loader-prefetch-num-threads": "8", "--reasoning-parser": "glm45", "--tool-call-parser": "glm45", "--dist-timeout": "3600", "--context-length": "32768", "--mem-fraction-static": "0.8", "--chunked-prefill-size": "16384", - "--model-loader-extra-config": '{"enable_multithread_load":true,"num_threads":8}', "--skip-server-warmup": "", } @@ -77,7 +74,7 @@ class _Miles(MilesConfig): rollout_endpoint_url = None use_miles_router = True - custom_rollout_request_hook_path = "cookbook.miles_disagg.hooks.gated_rollout_request_hook" + custom_rollout_request_hook_path = "cookbook.common.hooks.gated_rollout_request_hook" custom_config_path = { "rollout_request_weight_version_mode": "min", "rollout_request_weight_version_lag": 1, @@ -90,12 +87,11 @@ class _Miles(MilesConfig): async_mode = True update_weights_interval = 1 - update_weight_transfer_mode = "disk" - update_weight_disk_delta = True + update_weight_transfer_mode = "disk-delta" update_weight_delta_encoding = "xor" update_weight_delta_checksum = "xxh3-128" update_weight_disk_dir = DELTA_BULLETIN_ROOT - custom_delta_pre_push_path = "cookbook.miles_disagg.hooks.commit_and_wake" + custom_update_weight_post_write_path = "cookbook.common.hooks.commit_and_wake" prompt_data = f"{DATA_PATH}/dapo-math-17k/dapo-math-17k.jsonl" input_key = "prompt" diff --git a/cookbook/miles_disagg/configs/glm45_air_fp8.py b/cookbook/miles_disagg/configs/glm45_air_fp8.py new file mode 100644 index 0000000..2881f97 --- /dev/null +++ b/cookbook/miles_disagg/configs/glm45_air_fp8.py @@ -0,0 +1,174 @@ +"""GLM-4.5-Air: bf16 trainer, served in native HF FP8 through the disaggregated pool.""" + +from __future__ import annotations + +from cookbook.common.config import ModalConfig +from cookbook.common.constants import DATA_PATH, PREP_PATH +from cookbook.miles_disagg.config import MilesConfig + +APP_NAME = "stitch-glm45-air-fp8" +DELTA_VOLUME_NAME = "stitch-delta-glm45-air-fp8" +DELTA_BULLETIN_ROOT = "/delta-bulletin" +LOCAL_CHECKPOINT_PATH = "/local-checkpoint" +SIDECAR_COMMIT_MODE = "quiesce" # fp8 reload is exact; draining buys clean per-version attribution + +MODEL_TAG = "glm45-air-bf16" +SOURCE_MODEL = "zai-org/GLM-4.5-Air" +ROLLOUT_SOURCE_MODEL = "zai-org/GLM-4.5-Air-FP8" +SERVED_CHECKPOINT_FORMAT = "fp8" +USE_MODAL_TORCH_DIST_WRAPPER = True +DISABLE_HF_XET = True # the plain HF downloader is the path that finished reliably here +DISABLE_HF_TRANSFER = True + +MEGATRON_RUNTIME_PATCHES = ["/root/cookbook/miles_disagg/patches/megatron-r3-dispatch.patch"] + +SGLANG_SERVER_ARGS = { + # Dense FP8 delta => no partial reload => every step reloads the full checkpoint. Under + # gVisor the read is bound by the Sentry's single-thread byte-copy (~5 GB/s), so the win + # is fewer BYTES: fastsafetensors splits files across TP ranks (each reads 1/N), ~408 -> + # ~102 GB/node => reload ~130s -> ~23s. nogds is forced via SGLANG_FASTSAFETENSORS_NOGDS + # in the serving image (GDS/nvidia-fs is absent under gVisor). + "--load-format": "fastsafetensors", + "--dtype": "auto", + "--reasoning-parser": "glm45", + "--tool-call-parser": "glm45", + "--dist-timeout": "3600", + "--context-length": "32768", + "--mem-fraction-static": "0.7", + "--chunked-prefill-size": "8192", + "--max-prefill-tokens": "16384", + "--cuda-graph-max-bs-prefill": "2048", # avoid H200 cold-start graph-compile hangs + "--skip-server-warmup": "", +} + + +class _Miles(MilesConfig): + miles_model_script = "scripts/models/glm4.5-106B-A12B.sh" + + hf_checkpoint = f"{PREP_PATH}/{MODEL_TAG}/fp8" + ref_load = f"{PREP_PATH}/{MODEL_TAG}/torch_dist" + megatron_to_hf_mode = "raw" + model_name = "glm4moe" + + actor_num_nodes = 4 + actor_num_gpus_per_node = 8 + num_gpus_per_node = 8 + colocate = False + rollout_num_gpus = 0 # external rollout: framework runs no local engines + rollout_num_gpus_per_engine = 4 + rollout_endpoint_url = None # filled at launch from the pool gateway + use_miles_router = True + + custom_rollout_request_hook_path = "cookbook.common.hooks.gated_rollout_request_hook" + custom_update_weight_post_write_path = "cookbook.common.hooks.commit_and_wake" + custom_config_path = { + "rollout_request_weight_version_mode": "min", + "rollout_request_weight_version_lag": 1, + "rollout_request_retry_attempts": 900, + "rollout_request_retry_sleep": 1.0, + "rollout_session_affinity_header": "Modal-Session-ID", + "rollout_request_timeout_secs": 300, + } + + async_mode = True + update_weights_interval = 1 + update_weight_transfer_mode = "disk-delta" + update_weight_delta_encoding = "xor" + update_weight_delta_checksum = "xxh3-128" + update_weight_disk_dir = DELTA_BULLETIN_ROOT # run-scoped at launch to / + + prompt_data = f"{DATA_PATH}/dapo-math-17k/dapo-math-17k.jsonl" + input_key = "prompt" + label_key = "label" + apply_chat_template = True + rollout_shuffle = True + balance_data = True + rm_type = "deepscaler" + eval_interval = None + + num_rollout = 10 + save_interval = None # miles forces a final save when set, regardless of interval; leave off + rollout_batch_size = 16 + rollout_max_response_len = 4096 + rollout_temperature = 0.8 + n_samples_per_prompt = 4 + global_batch_size = 64 + use_dynamic_global_batch_size = True + sglang_server_concurrency = 128 + use_rollout_routing_replay = False + + tensor_model_parallel_size = 1 + sequence_parallel = True + pipeline_model_parallel_size = 4 + context_parallel_size = 1 + expert_model_parallel_size = 8 + expert_tensor_parallel_size = 1 + decoder_last_pipeline_num_layers = 10 + use_dynamic_batch_size = True + max_tokens_per_gpu = 8192 + recompute_granularity = "full" + recompute_method = "uniform" + recompute_num_layers = 1 + attention_dropout = 0.0 + hidden_dropout = 0.0 + accumulate_allreduce_grads_in_fp32 = True + attention_softmax_in_fp32 = True + no_check_for_nan_in_loss_and_grad = True + + optimizer = "adam" + lr = 1e-6 + lr_decay_style = "constant" + weight_decay = 0.1 + adam_beta1 = 0.9 + adam_beta2 = 0.98 + optimizer_cpu_offload = True + overlap_cpu_optimizer_d2h_h2d = True + use_precision_aware_optimizer = True + + advantage_estimator = "gspo" + eps_clip = 4e-4 + eps_clip_high = None + use_kl_loss = False + kl_loss_coef = 0.0 + kl_loss_type = "low_var_kl" + entropy_coef = 0.0 + use_tis = True + + environment = { + "CUDA_DEVICE_MAX_CONNECTIONS": "1", + "NCCL_NVLS_ENABLE": "1", + "NVSHMEM_DISABLE_NCCL": "1", + "NCCL_TIMEOUT_MS": "360000000", + } + + def prepare_data(self) -> None: + from datasets import load_dataset + + ds = load_dataset("BytedTsinghua-SIA/DAPO-Math-17k", split="train") + ds = ds.shuffle(seed=42).select(range(min(50000, ds.num_rows))) + ds = ds.map(lambda ex: {"label": ex["reward_model"]["ground_truth"]}) + ds = ds.select_columns(["prompt", "label"]) + ds.to_json(f"{DATA_PATH}/dapo-math-17k/dapo-math-17k.jsonl") + + +modal = ModalConfig( + gpu="H200", + memory=1_048_576, + rollout_min_containers=2, + rollout_max_containers=4, # start at 2; scale to 4 mid-run to exercise elastic join + rollout_target_inputs=32, + proxy_regions=["us-west"], + rollout_ephemeral_disk_mib=819_200, + torch_dist_prep_nodes=4, + torch_dist_prep_gpus_per_node=8, + torch_dist_convert_extra_args=( + "--tensor-model-parallel-size 1 " + "--pipeline-model-parallel-size 4 " + "--expert-model-parallel-size 8 " + "--expert-tensor-parallel-size 1 " + "--decoder-last-pipeline-num-layers 10" + ), + torch_dist_prep_ephemeral_disk_mib=819_200, +) + +miles = _Miles() diff --git a/cookbook/miles_disagg/configs/glm5_2_nvfp4.py b/cookbook/miles_disagg/configs/glm5_2_nvfp4.py new file mode 100644 index 0000000..9997745 --- /dev/null +++ b/cookbook/miles_disagg/configs/glm5_2_nvfp4.py @@ -0,0 +1,247 @@ +"""GLM-5.2 GRPO on Modal, disaggregated, native-NVFP4 end to end. + +Deploy: EXPERIMENT_CONFIG=glm5_2_nvfp4 uv run --extra modal modal deploy --strategy recreate -m cookbook.miles_disagg.app +""" + +from __future__ import annotations + +from cookbook.common.config import ModalConfig +from cookbook.common.constants import DATA_PATH, PREP_PATH +from cookbook.miles_disagg.config import MilesConfig + + +APP_NAME = "stitch-glm5-2-nvfp4" +DELTA_VOLUME_NAME = "stitch-delta-glm5-2-nvfp4" +DELTA_BULLETIN_ROOT = "/delta-bulletin" +LOCAL_CHECKPOINT_PATH = "/local-checkpoint" + +# GLM-5.2 ships bf16 (unlike Kimi's INT4): the bf16 masters ARE the source, no dequant. +SOURCE_MODEL = "zai-org/GLM-5.2" +MODEL_TAG = "glm5-2-nvfp4" + +SIDECAR_COMMIT_MODE = "in_place" +SIDECAR_DEBUG_REQUESTS = True +# R3 routing-replay needs the dropless Megatron dispatch fix at startup. +MEGATRON_RUNTIME_PATCHES = [ + "/root/cookbook/miles_disagg/patches/megatron-r3-dispatch.patch", +] + + +SGLANG_SERVER_ARGS = { + "--weight-loader-prefetch-checkpoints": "", + "--weight-loader-prefetch-num-threads": "8", + # v0.5.15 has no glm5.2 parser; glm45/glm47 are closest. TODO(glm5.2): confirm 5.2's format. + "--reasoning-parser": "glm45", + "--tool-call-parser": "glm47", + "--dist-timeout": "3600", + "--kv-cache-dtype": "fp8_e4m3", + "--attention-backend": "tokenspeed_mla", + "--context-length": "32768", + "--mem-fraction-static": "0.85", + "--chunked-prefill-size": "16384", + "--schedule-conservativeness": "0.5", + "--schedule-policy": "lpm", + "--enable-hierarchical-cache": "", + "--hicache-ratio": "2", + "--hicache-io-backend": "kernel", + "--hicache-mem-layout": "page_first", + "--hicache-write-policy": "write_through", + "--skip-server-warmup": "", + "--enable-return-routed-experts": "", # routing replay (DeepSeek-V3-arch MoE) +} + +SGLANG_ENV = {"SGLANG_ENABLE_RELOAD_LOAD_PLAN": "1"} # NVFP4: load-plan replay + O(delta) partial reload + +modal = ModalConfig( + gpu="B200", + region="us", + # TODO(glm5.2): size to the actual model. These are Kimi's (~1T) numbers — scale + # down for a smaller GLM 5.2 (memory, ephemeral disk, node count). + memory=1_650_688, + rollout_min_containers=8, # warm floor; Flash scales above under load + rollout_target_inputs=32, + proxy_regions=["us-west"], + rollout_ephemeral_disk_mib=819_200, # NVFP4 base copy + in-place delta headroom + trainer_ephemeral_disk_mib=2_097_152, + # TODO(glm5.2): torch_dist conversion parallelism — match the trainer EP/TP below. + torch_dist_prep_nodes=4, + torch_dist_prep_gpus_per_node=8, + torch_dist_convert_extra_args=( + "--tensor-model-parallel-size 1 --pipeline-model-parallel-size 1 --expert-model-parallel-size 32" + ), + torch_dist_prep_ephemeral_disk_mib=2_097_152, +) + + +class _Miles(MilesConfig): + miles_model_script = "scripts/models/glm5.2-744B-A40B.sh" + + hf_checkpoint = f"{PREP_PATH}/{MODEL_TAG}/nvfp4" + ref_load = f"{PREP_PATH}/{MODEL_TAG}/torch_dist" + megatron_to_hf_mode = "raw" + model_name = "glm_moe_dsa" + + # TODO(glm5.2): size actor_num_nodes / parallelism to the model (Kimi's 16x8=128). + actor_num_nodes = 16 + actor_num_gpus_per_node = 8 + num_gpus_per_node = 8 + colocate = False + rollout_num_gpus = 0 + rollout_num_gpus_per_engine = 4 # TODO(glm5.2): B200s per rollout engine (fit the base) + rollout_endpoint_url = None + use_miles_router = True + + custom_rollout_request_hook_path = ( + "cookbook.common.hooks.gated_rollout_request_hook" + ) + custom_config_path = { + "rollout_request_weight_version_mode": "min", + "rollout_request_weight_version_lag": 1, + "rollout_request_retry_attempts": 1200, # outlast a full cold pool load + "rollout_request_retry_sleep": 1.0, + "rollout_session_affinity_header": "Modal-Session-ID", + "rollout_request_timeout_secs": 300, + } + + async_mode = True + update_weights_interval = 1 + + # NVFP4 QAT — miles' canonical NVFP4 RL recipe (shared with Kimi). + fp4_format = "e2m1" + fp4_recipe = "nvfp4" + fp4_param_gather = False + # NVFP4 only on the routed expert GEMMs, everything else bf16 — matches the served base. + te_precision_config_file = { + "configs": { + "nvfp4": { + "transformer_engine_config_type": "TEQuantizationParams", + "training_recipe": {"fp4_quantization_recipe": "nvfp4"}, + }, + "bf16": { + "transformer_engine_config_type": "TEQuantizationParams", + "training_recipe": {}, + }, + }, + "matchers": { + "routed_experts_fc1_nvfp4": { + "type": "glob", + "enabled": True, + "pattern": "*.mlp.experts.linear_fc1", + "config": "nvfp4", + }, + "routed_experts_fc2_nvfp4": { + "type": "glob", + "enabled": True, + "pattern": "*.mlp.experts.linear_fc2", + "config": "nvfp4", + }, + "default_bf16": { + "type": "glob", + "enabled": True, + "pattern": "*", + "config": "bf16", + }, + }, + } + # Keep ALL routed experts NVFP4 (reload-safe: sglang's fused-MoE reload loader + # allocates NVFP4 for every expert layer; a bf16 carve-out can't reload). + num_layers_at_start_in_bf16 = 3 + num_layers_at_end_in_bf16 = 0 + + update_weight_transfer_mode = "disk-delta" + update_weight_delta_encoding = "xor" + update_weight_delta_checksum = "xxh3-128" + update_weight_disk_dir = DELTA_BULLETIN_ROOT + custom_update_weight_post_write_path = "cookbook.common.hooks.commit_and_wake" + + prompt_data = f"{DATA_PATH}/dapo-math-17k/dapo-math-17k.jsonl" + input_key = "prompt" + label_key = "label" + apply_chat_template = True + rollout_shuffle = True + balance_data = True + rm_type = "deepscaler" + eval_interval = None + + num_rollout = 3 # bring-up smoke length; scale up for a real run + save_interval = 20 + rollout_batch_size = 32 + rollout_max_response_len = 4096 + rollout_temperature = 0.8 + n_samples_per_prompt = 8 + global_batch_size = 256 + use_dynamic_global_batch_size = True + sglang_server_concurrency = 256 + + use_rollout_routing_replay = True + + # TODO(glm5.2): trainer parallelism — size to the model + actor_num_nodes above + # (mirrors Kimi's 128-GPU layout; adjust TP/PP/CP/EP to GLM 5.2's layer/expert counts). + tensor_model_parallel_size = 8 + sequence_parallel = True + pipeline_model_parallel_size = 8 + context_parallel_size = 2 + expert_model_parallel_size = 16 + expert_tensor_parallel_size = 1 + decoder_last_pipeline_num_layers = 5 + use_dynamic_batch_size = True + max_tokens_per_gpu = 16384 + recompute_granularity = "full" + recompute_method = "uniform" + recompute_num_layers = 1 + attention_dropout = 0.0 + hidden_dropout = 0.0 + accumulate_allreduce_grads_in_fp32 = True + attention_softmax_in_fp32 = True + no_check_for_nan_in_loss_and_grad = True + + optimizer = "adam" + lr = 1e-6 + lr_decay_style = "constant" + weight_decay = 0.1 + adam_beta1 = 0.9 + adam_beta2 = 0.98 + optimizer_cpu_offload = True + overlap_cpu_optimizer_d2h_h2d = True + use_precision_aware_optimizer = True + + advantage_estimator = "grpo" + eps_clip = 0.2 + eps_clip_high = 0.28 + use_kl_loss = True + kl_loss_coef = 0.0 + kl_loss_type = "low_var_kl" + entropy_coef = 0.0 + use_tis = True + + environment = { + "CUDA_DEVICE_MAX_CONNECTIONS": "1", + "NCCL_NVLS_ENABLE": "1", + "NVSHMEM_DISABLE_NCCL": "1", + "NCCL_TIMEOUT_MS": "360000000", + # NVFP4 numerics (shared with Kimi; required for correct NVFP4 QAT). + "NVTE_NVFP4_DISABLE_2D_QUANTIZATION": "1", + "NVTE_NVFP4_DISABLE_RHT": "1", + "NVTE_NVFP4_DISABLE_STOCHASTIC_ROUNDING": "1", + "NVTE_NVFP4_ROW_SCALED_ACTIVATION": "1", + "NVTE_BACKWARD_OVERRIDE": "high_precision", + "NVTE_USE_FAST_MATH": "0", + } + + def prepare_data(self) -> None: + from datasets import load_dataset + + ds = load_dataset("BytedTsinghua-SIA/DAPO-Math-17k", split="train") + ds = ds.shuffle(seed=42).select(range(min(50000, ds.num_rows))) + ds = ds.map(lambda ex: {"label": ex["reward_model"]["ground_truth"]}) + ds = ds.select_columns(["prompt", "label"]) + ds.to_json(f"{DATA_PATH}/dapo-math-17k/dapo-math-17k.jsonl") + + +# TODO(glm5.2) — still to confirm before a real run: +# 1. Parsers: glm45 vs a GLM-5.2-specific reasoning/tool parser in sglang v0.5.15. +# 2. Size the trainer to 744B-A40B / 78 layers: actor_num_nodes + TP/PP/CP/EP + +# decoder_last_pipeline_num_layers + memory / ephemeral disk (values above are +# Kimi's ~1T/128-GPU layout as a starting point). + +miles = _Miles() diff --git a/cookbook/miles_disagg/configs/kimi_k25_2layer_nvfp4_disagg.py b/cookbook/miles_disagg/configs/kimi_k25_2layer_nvfp4.py similarity index 74% rename from cookbook/miles_disagg/configs/kimi_k25_2layer_nvfp4_disagg.py rename to cookbook/miles_disagg/configs/kimi_k25_2layer_nvfp4.py index 62ad400..13007d6 100644 --- a/cookbook/miles_disagg/configs/kimi_k25_2layer_nvfp4_disagg.py +++ b/cookbook/miles_disagg/configs/kimi_k25_2layer_nvfp4.py @@ -1,46 +1,35 @@ -"""Kimi-K2.5 2-layer NVFP4 disaggregated — the FAST iteration proxy for K2.6. - -CharyZeng/Kimi-K2.5-2layer is the identical arch class as K2.6 -(KimiK25ForConditionalGeneration, model_type "kimi_k25": MLA + DeepSeek-MoE, 384 -experts, topk 8, first_k_dense_replace 1, compressed-tensors INT4) but only 2 -layers. So it exercises the EXACT K2.6-specific path — INT4->bf16 dequant, NVFP4 -convert, the KimiK25 mbridge import converter, raw-mode build, torch_dist load, -te-precision NVFP4 QAT, disk-delta publish, NVFP4 serving — on a SINGLE node in -minutes, instead of 16x8 B200 + ~40 min init. Shake out the init/converter path -here, then run the full kimi_k2_6_nvfp4_disagg once it's clean. - - EXPERIMENT_CONFIG=kimi_k25_2layer_nvfp4_disagg \ - uv run --extra modal modal run --detach -m cookbook.miles_disagg.modal_train::prepare_checkpoints - EXPERIMENT_CONFIG=kimi_k25_2layer_nvfp4_disagg \ - uv run --extra modal modal run --detach -m cookbook.miles_disagg.modal_train::prepare_torch_dist - EXPERIMENT_CONFIG=kimi_k25_2layer_nvfp4_disagg \ - uv run --extra modal modal deploy --strategy recreate -m cookbook.miles_disagg.modal_train - # then ::launch_train +"""Kimi-K2.5 2-layer NVFP4 disaggregated — the FAST single-node iteration proxy for K2.6. + +Deploy (::prepare_checkpoints, ::prepare_torch_dist, deploy, then ::launch_train): + EXPERIMENT_CONFIG=kimi_k25_2layer_nvfp4 uv run --extra modal modal deploy --strategy recreate -m cookbook.miles_disagg.app """ from __future__ import annotations -from cookbook.miles_disagg.configs.base import DATA_PATH, PREP_PATH, ModalConfig, MilesConfig +from cookbook.common.config import ModalConfig +from cookbook.common.constants import DATA_PATH, PREP_PATH +from cookbook.miles_disagg.config import MilesConfig -APP_NAME = "miles-kimi-k25-2layer-nvfp4-disagg" -DELTA_VOLUME_NAME = "miles-delta-bulletin-kimi-k25-2layer" +APP_NAME = "stitch-kimi-k25-2layer-nvfp4" +DELTA_VOLUME_NAME = "stitch-delta-kimi-k25-2layer-nvfp4" DELTA_BULLETIN_ROOT = "/delta-bulletin" +LOCAL_CHECKPOINT_PATH = "/local-checkpoint" SOURCE_MODEL = "CharyZeng/Kimi-K2.5-2layer" # INT4, KimiK25 arch, 2 layers MODEL_TAG = "kimi-k25-2layer-nvfp4" SIDECAR_COMMIT_MODE = "in_place" SIDECAR_DEBUG_REQUESTS = True - - -def build_serving_image(**kwargs): - from cookbook.miles_disagg.serving import build_nvfp4_b200_serving_image - - return build_nvfp4_b200_serving_image(**kwargs) +# R3 routing-replay needs the dropless Megatron dispatch fix at startup. +MEGATRON_RUNTIME_PATCHES = [ + "/root/cookbook/miles_disagg/patches/megatron-r3-dispatch.patch", +] SGLANG_SERVER_ARGS = { + "--weight-loader-prefetch-checkpoints": "", + "--weight-loader-prefetch-num-threads": "8", "--tool-call-parser": "kimi_k2", "--reasoning-parser": "kimi_k2", "--dist-timeout": "3600", @@ -52,6 +41,8 @@ def build_serving_image(**kwargs): "--enable-return-routed-experts": "", } +SGLANG_ENV = {"SGLANG_ENABLE_RELOAD_LOAD_PLAN": "1"} # NVFP4: load-plan replay + O(delta) partial reload + modal = ModalConfig( gpu="B200", region="us", @@ -71,7 +62,6 @@ class _Miles(MilesConfig): megatron_to_hf_mode = "raw" model_name = "kimi_k25" # KimiK25 mbridge import + convert_kimi_k25_to_hf export - # Disaggregated publish-only rollout on a single node (B200:4 trainer). actor_num_nodes = 1 actor_num_gpus_per_node = 4 num_gpus_per_node = 4 @@ -81,7 +71,7 @@ class _Miles(MilesConfig): rollout_endpoint_url = None use_miles_router = True - custom_rollout_request_hook_path = "cookbook.miles_disagg.hooks.gated_rollout_request_hook" + custom_rollout_request_hook_path = "cookbook.common.hooks.gated_rollout_request_hook" custom_config_path = { "rollout_request_weight_version_mode": "min", "rollout_request_weight_version_lag": 1, @@ -93,7 +83,7 @@ class _Miles(MilesConfig): async_mode = True update_weights_interval = 1 - # NVFP4 QAT — canonical recipe (radixark/miles#1261), same as K2.6. + # NVFP4 QAT — same canonical recipe as K2.6. fp4_format = "e2m1" fp4_recipe = "nvfp4" fp4_param_gather = False @@ -126,12 +116,11 @@ class _Miles(MilesConfig): num_layers_at_start_in_bf16 = 1 num_layers_at_end_in_bf16 = 0 - update_weight_transfer_mode = "disk" - update_weight_disk_delta = True + update_weight_transfer_mode = "disk-delta" update_weight_delta_encoding = "xor" update_weight_delta_checksum = "xxh3-128" update_weight_disk_dir = DELTA_BULLETIN_ROOT - custom_delta_pre_push_path = "cookbook.miles_disagg.hooks.commit_and_wake" + custom_update_weight_post_write_path = "cookbook.common.hooks.commit_and_wake" prompt_data = f"{DATA_PATH}/dapo-math-17k/dapo-math-17k.jsonl" input_key = "prompt" diff --git a/cookbook/miles_disagg/configs/kimi_k2_6_nvfp4.py b/cookbook/miles_disagg/configs/kimi_k2_6_nvfp4.py new file mode 100644 index 0000000..60bf77f --- /dev/null +++ b/cookbook/miles_disagg/configs/kimi_k2_6_nvfp4.py @@ -0,0 +1,254 @@ +"""Kimi K2.6 GRPO on Modal, disaggregated, native-NVFP4 end to end. + +Deploy: EXPERIMENT_CONFIG=kimi_k2_6_nvfp4 uv run --extra modal modal deploy --strategy recreate -m cookbook.miles_disagg.app +""" + +from __future__ import annotations + +from cookbook.common.config import ModalConfig +from cookbook.common.constants import DATA_PATH, PREP_PATH +from cookbook.miles_disagg.config import MilesConfig + + +APP_NAME = "stitch-kimi-k2-6-nvfp4" +DELTA_VOLUME_NAME = "stitch-delta-kimi-k2-6-nvfp4" +DELTA_BULLETIN_ROOT = "/delta-bulletin" +LOCAL_CHECKPOINT_PATH = "/local-checkpoint" + +# SOURCE_MODEL (INT4) -> dequant -> bf16 masters -> convert -> served NVFP4 base. +SOURCE_MODEL = "moonshotai/Kimi-K2.6" +MODEL_TAG = "kimi-k2-6-nvfp4" + +SIDECAR_COMMIT_MODE = "in_place" +SIDECAR_DEBUG_REQUESTS = True +# R3 routing-replay needs the dropless Megatron dispatch fix at startup. +MEGATRON_RUNTIME_PATCHES = [ + "/root/cookbook/miles_disagg/patches/megatron-r3-dispatch.patch", +] + + +# mem-fraction / context-length are STARTING POINTS — measure on a warm B200:4 and adjust. +SGLANG_SERVER_ARGS = { + "--weight-loader-prefetch-checkpoints": "", + "--weight-loader-prefetch-num-threads": "8", + "--tool-call-parser": "kimi_k2", + "--reasoning-parser": "kimi_k2", + "--dist-timeout": "3600", + "--kv-cache-dtype": "fp8_e4m3", + "--attention-backend": "tokenspeed_mla", + "--context-length": "32768", + "--mem-fraction-static": "0.85", + "--chunked-prefill-size": "16384", + "--schedule-conservativeness": "0.5", + "--schedule-policy": "lpm", + "--enable-hierarchical-cache": "", + "--hicache-ratio": "2", + "--hicache-io-backend": "kernel", + "--hicache-mem-layout": "page_first", + "--hicache-write-policy": "write_through", + "--skip-server-warmup": "", + "--enable-return-routed-experts": "", +} + +SGLANG_ENV = {"SGLANG_ENABLE_RELOAD_LOAD_PLAN": "1"} # NVFP4: load-plan replay + O(delta) partial reload + +modal = ModalConfig( + gpu="B200", + # Request Modal's max (nodes are 1.79 TiB): the CPU-offloaded 1T optimizer + publish + # snapshot peak ~1.2 TiB, and a smaller *request* got OOM-killed (excess is best-effort). + memory=1_650_688, + # Keep the rollout pool small (2x B200:4 = 8 GPUs) so the 64-GPU trainer gang can + # co-schedule; min==max pins it, else it autoscales up and re-competes for B200. + rollout_min_containers=2, + rollout_max_containers=2, + rollout_target_inputs=32, + proxy_regions=["us-west"], + rollout_ephemeral_disk_mib=819_200, # ~591 GB base copy + in-place delta-apply headroom + # Ray logs + object spill fill the default disk (ENOSPC'd); 2 TiB gives headroom. + trainer_ephemeral_disk_mib=2_097_152, + # 4x8=32-way: EP32 shards 384 experts to ~90 GB/rank so distcp writes finish inside the + # 900s heartbeat (EP16 raced it); torch_dist reshards on load into the EP16 trainer. + torch_dist_prep_nodes=4, + torch_dist_prep_gpus_per_node=8, + torch_dist_convert_extra_args=( + "--tensor-model-parallel-size 1 --pipeline-model-parallel-size 1 --expert-model-parallel-size 32" + ), + # ~700 GB of distcp shards buffer before commit; 2 TB gives headroom. + torch_dist_prep_ephemeral_disk_mib=2_097_152, +) + + +class _Miles(MilesConfig): + # Arch comes from the model script (shared with Kimi-K2-Thinking; K2.6 matches). + miles_model_script = "scripts/models/kimi-k2-thinking.sh" + + hf_checkpoint = f"{PREP_PATH}/{MODEL_TAG}/nvfp4" + ref_load = f"{PREP_PATH}/{MODEL_TAG}/torch_dist" + # "raw" (not "bridge"): K2.6's HF arch is the KimiK25 VLM wrapper, which AutoBridge + # can't build. Export still routes via model_name below. + megatron_to_hf_mode = "raw" + model_name = ( + "kimi_k25" # megatron_to_hf export dispatch (convert_kimi_k25_to_hf + NVFP4) + ) + + actor_num_nodes = 8 # 8x8 B200 = 64 GPUs (trainer only; pool is elastic on top) + actor_num_gpus_per_node = 8 + num_gpus_per_node = 8 + colocate = False + rollout_num_gpus = 0 + rollout_num_gpus_per_engine = 4 # B200:4 per rollout container (K2.6 NVFP4 fits) + rollout_endpoint_url = None + use_miles_router = True + + # Staleness gate; knobs ride in custom_config_path (read by the hook, not miles core). + custom_rollout_request_hook_path = ( + "cookbook.common.hooks.gated_rollout_request_hook" + ) + custom_config_path = { + "rollout_request_weight_version_mode": "min", + "rollout_request_weight_version_lag": 1, + # 1200×1s=20 min outlasts a full ~16 min cold-load (240s was too short, rollout 503'd). + "rollout_request_retry_attempts": 1200, + "rollout_request_retry_sleep": 1.0, + "rollout_session_affinity_header": "Modal-Session-ID", + # finite read timeout; without it a request to a scaled-down Flash container hangs forever. + "rollout_request_timeout_secs": 300, + } + + async_mode = True + update_weights_interval = 1 + + # NVFP4 QAT — miles' canonical NVFP4 RL recipe. + fp4_format = "e2m1" + fp4_recipe = "nvfp4" + # False: with it True, TE NVFP4Tensor params crash Megatron DDP's param-buffer repoint. + fp4_param_gather = False + # NVFP4 only on the routed expert GEMMs, everything else bf16 — matches the served base. + te_precision_config_file = { + "configs": { + "nvfp4": { + "transformer_engine_config_type": "TEQuantizationParams", + "training_recipe": {"fp4_quantization_recipe": "nvfp4"}, + }, + "bf16": { + "transformer_engine_config_type": "TEQuantizationParams", + "training_recipe": {}, + }, + }, + "matchers": { + "routed_experts_fc1_nvfp4": { + "type": "glob", + "enabled": True, + "pattern": "*.mlp.experts.linear_fc1", + "config": "nvfp4", + }, + "routed_experts_fc2_nvfp4": { + "type": "glob", + "enabled": True, + "pattern": "*.mlp.experts.linear_fc2", + "config": "nvfp4", + }, + "default_bf16": { + "type": "glob", + "enabled": True, + "pattern": "*", + "config": "bf16", + }, + }, + } + # bf16 carve-out for the dense first layer; END stays 0 (reload-safe): SGLang's fused-MoE + # reload loader allocates NVFP4 for every expert layer, so a bf16 last layer can't reload. + num_layers_at_start_in_bf16 = 1 + num_layers_at_end_in_bf16 = 0 + + update_weight_transfer_mode = "disk-delta" + update_weight_delta_encoding = "xor" + update_weight_delta_checksum = "xxh3-128" + update_weight_disk_dir = DELTA_BULLETIN_ROOT + custom_update_weight_post_write_path = "cookbook.common.hooks.commit_and_wake" + + prompt_data = f"{DATA_PATH}/dapo-math-17k/dapo-math-17k.jsonl" + input_key = "prompt" + label_key = "label" + apply_chat_template = True + rollout_shuffle = True + balance_data = True + rm_type = "deepscaler" + eval_interval = None + + num_rollout = 10 # 10 GRPO steps; each publishes a delta to inspect for tensor-sparsity + save_interval = 20 # megatron requires it; > num_rollout so this run skips saves + rollout_batch_size = 32 + rollout_max_response_len = 4096 + rollout_temperature = 0.8 + n_samples_per_prompt = 8 + global_batch_size = 256 + use_dynamic_global_batch_size = True + sglang_server_concurrency = 256 + + use_rollout_routing_replay = True + + # Trainer parallelism for 8x8=64 GPUs: TP8*PP2*CP4=64, EP32 over the experts — the + # proven kimi_k25_lora_8nodes topology. EP32 keeps expert weights+grads ~63GB/rank. + tensor_model_parallel_size = 8 + sequence_parallel = True + pipeline_model_parallel_size = 2 + context_parallel_size = 4 + expert_model_parallel_size = 32 + expert_tensor_parallel_size = 1 + decoder_last_pipeline_num_layers = 30 + use_dynamic_batch_size = True + max_tokens_per_gpu = 16384 + recompute_granularity = "full" + recompute_method = "uniform" + recompute_num_layers = 1 + attention_dropout = 0.0 + hidden_dropout = 0.0 + accumulate_allreduce_grads_in_fp32 = True + attention_softmax_in_fp32 = True + no_check_for_nan_in_loss_and_grad = True + + optimizer = "adam" + lr = 1e-6 + lr_decay_style = "constant" + weight_decay = 0.1 + adam_beta1 = 0.9 + adam_beta2 = 0.98 + optimizer_cpu_offload = True + overlap_cpu_optimizer_d2h_h2d = True + use_precision_aware_optimizer = True + + advantage_estimator = "grpo" + eps_clip = 0.2 + eps_clip_high = 0.28 + use_kl_loss = True + kl_loss_coef = 0.0 + kl_loss_type = "low_var_kl" + entropy_coef = 0.0 + use_tis = True + + environment = { + "CUDA_DEVICE_MAX_CONNECTIONS": "1", + "NCCL_NVLS_ENABLE": "1", + "NVSHMEM_DISABLE_NCCL": "1", + "NCCL_TIMEOUT_MS": "360000000", + # NVFP4 numerics: required for correct NVFP4 QAT. + "NVTE_NVFP4_DISABLE_2D_QUANTIZATION": "1", + "NVTE_NVFP4_DISABLE_RHT": "1", + "NVTE_NVFP4_DISABLE_STOCHASTIC_ROUNDING": "1", + "NVTE_NVFP4_ROW_SCALED_ACTIVATION": "1", + "NVTE_BACKWARD_OVERRIDE": "high_precision", + "NVTE_USE_FAST_MATH": "0", + } + + def prepare_data(self) -> None: + from datasets import load_dataset + + ds = load_dataset("BytedTsinghua-SIA/DAPO-Math-17k", split="train") + ds = ds.shuffle(seed=42).select(range(min(50000, ds.num_rows))) + ds = ds.map(lambda ex: {"label": ex["reward_model"]["ground_truth"]}) + ds = ds.select_columns(["prompt", "label"]) + ds.to_json(f"{DATA_PATH}/dapo-math-17k/dapo-math-17k.jsonl") + + +miles = _Miles() diff --git a/cookbook/miles_disagg/configs/kimi_k2_6_nvfp4_disagg.py b/cookbook/miles_disagg/configs/kimi_k2_6_nvfp4_disagg.py deleted file mode 100644 index ee6b05b..0000000 --- a/cookbook/miles_disagg/configs/kimi_k2_6_nvfp4_disagg.py +++ /dev/null @@ -1,374 +0,0 @@ -"""Kimi K2.6 GRPO on Modal, disaggregated, native-NVFP4 end to end. - -Kimi K2.6 is a ~1T-parameter DeepSeek-V3-architecture MoE (MLA + DeepSeek-MoE: -sigmoid router, shared experts, grouped top-k). The rollout pool serves the -model's **NVFP4 (W4A4, experts-only)** checkpoint on Blackwell; the trainer -QAT-trains in NVFP4 and publishes byte-exact XOR deltas against that served base. - -This is the full-scale recipe — a 16x8 B200 (128 GPU) trainer footprint, with an -elastic Flash B200 rollout pool on top. Run ``moonlight_nvfp4_disagg`` first: it -exercises the exact same QAT → NVFP4-export → XOR-delta → SGLang-reload machinery -on a single Blackwell node and de-risks the one axis this recipe can only verify -on a warm container (export packing == served packing). The Kimi-specific paths -this adds over Moonlight are the INT4→bf16 dequant (``convert_kimi_int4_to_bf16``) -and the ``kimi_k25`` megatron_to_hf export converter (``convert_kimi_k25_to_hf``). - -NVFP4 IS ALL-BLACKWELL ----------------------- -NVFP4 QAT = native Megatron ``--fp4-format e2m1`` → ``NVFP4BlockScaling`` (TE >= -2.7.0.dev0). FP4 GEMM requires Blackwell, so the TRAINER is B200 too (not H200 + -fake-QAT like the INT4 recipe). The serving pool is the existing Blackwell SGLang -fork (serving.py), proven for NVFP4. - -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: - * 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. - * Megatron torch_dist (``load``/``save``): trainer-internal rollout ckpts. - -The trainer reads ``hf_checkpoint`` (NVFP4 base) for the export quant config AND -the diff baseline (``_capture_baseline`` seeds the snapshot from those exact -bytes), so applying delta_vN reproduces export_vN byte-for-byte. - -INVARIANT: NVFP4 group size is fixed at 16 and the served base's -quantization_config (quant_algo == "NVFP4", experts-only, FP8 KV) drives both -the engine load and the miles export dispatch (added in -megatron_to_hf/processors/__init__.py: route on quant_algo == "NVFP4"). - -Deploy as its own app: - EXPERIMENT_CONFIG=kimi_k2_6_nvfp4_disagg \ - uv run --extra modal modal deploy --strategy recreate -m cookbook.miles_disagg.modal_train - -Prerequisites the bring-up depends on (flagged, validated by the moonlight run): - 1. The miles trainer image's TransformerEngine must be >= 2.7.0.dev0 (NVFP4 - BlockScaling), and the trainer must run on Blackwell. Verify on a warm B200. - 2. SGLang must serve NVFP4 MLA MoE on Blackwell (serving.py fork, proven for - NVFP4) — verify the prepared base loads on a warm container. - 3. The exact ``--fp4-format`` companion args for K2.6 (e.g. high-precision - first/last layers) may need tuning; confirm against the Megatron fork. -""" - -from __future__ import annotations - -from cookbook.miles_disagg.configs.base import ( - DATA_PATH, - PREP_PATH, - ModalConfig, - MilesConfig, -) - - -APP_NAME = "miles-kimi-k2-6-nvfp4-disagg" -DELTA_VOLUME_NAME = "miles-delta-bulletin-kimi-k2-6-nvfp4" -DELTA_BULLETIN_ROOT = "/delta-bulletin" - -# 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" -SIDECAR_DEBUG_REQUESTS = True - - -def build_serving_image(**kwargs): - from cookbook.miles_disagg.serving import build_nvfp4_b200_serving_image - - return build_nvfp4_b200_serving_image(**kwargs) - - -# Kimi K2.6 on the B200 pool serving the NVFP4 base. The trainer image injects -# --served-model-name / --dtype / --cuda-graph-max-bs / --max-running-requests / -# --trust-remote-code; only the Kimi/MLA/cache extras live here. mem-fraction and -# context-length are STARTING POINTS — measure on a warm B200:4 and adjust. -SGLANG_SERVER_ARGS = { - "--tool-call-parser": "kimi_k2", - "--reasoning-parser": "kimi_k2", - "--dist-timeout": "3600", - "--kv-cache-dtype": "fp8_e4m3", - "--attention-backend": "tokenspeed_mla", - "--context-length": "32768", - "--mem-fraction-static": "0.85", - "--chunked-prefill-size": "16384", - "--schedule-conservativeness": "0.5", - "--schedule-policy": "lpm", - "--enable-hierarchical-cache": "", - "--hicache-ratio": "2", - "--hicache-io-backend": "kernel", - "--hicache-mem-layout": "page_first", - "--hicache-write-policy": "write_through", - "--skip-server-warmup": "", - "--enable-return-routed-experts": "", -} - -modal = ModalConfig( - gpu="B200", - region="us", - # Host RAM (request, limit MiB). The 1T optimizer is CPU-offloaded (optimizer_cpu_offload — - # required; the largest PP-stage ranks' fp32 Adam+master ~104 GB wouldn't fit GPU alongside - # the publish gather), so each rank holds ~145 GB host RAM => ~1.15 TB/node steady, and the - # disk-delta publish (baseline snapshot of ~24k tensors) spikes higher. A node OOM-killed at - # ~1.2 TiB because the *request* (768 GiB) — not the limit — was the binding guarantee (memory - # above the request is best-effort, reclaimed under node pressure). AWS B200:8 nodes have - # 1.79 TiB RAM; Modal caps the request at 1,650,688 MiB (1.574 TiB), so request that max — - # guaranteed past the ~1.2 TiB publish peak, with OS headroom on the 1.79 TiB node. - memory=1_650_688, - # Warm floor so the pool is up before the trainer sends rollouts (see the - # moonlight config). Flash scales above this under load. - rollout_min_containers=8, # warm floor: skip the cold 2->N ramp that 502'd v5's rollout - # Scale OUT at ~32 concurrent/container (vs the 256 engine concurrency) so Flash - # spreads the rollout across containers instead of saturating a few KV caches. - rollout_target_inputs=32, - proxy_regions=["us-west"], - # The served NVFP4 base is ~591 GB; the sidecar copies it to /local-checkpoint - # on ephemeral disk. 800 GiB covers the copy plus in-place delta-apply headroom. - rollout_ephemeral_disk_mib=819_200, - # Trainer nodes: Ray logs + object spill (rollout batches + per-publish full-model - # gathers) accumulate under /tmp/ray over the run; the default disk progressively - # ENOSPC'd. 2 TiB of the B200:8 local NVMe gives ample headroom. - trainer_ephemeral_disk_mib=2_097_152, - # torch_dist conversion: 4x8 B200 = 32-way. EP32 shards the 384 experts to - # ~90 GB/rank so each rank's distcp write finishes well inside the 900s cluster - # heartbeat window (EP16's ~150 GB/rank raced it and missed .metadata). TP1/PP1 - # keeps the (replicated) non-expert small; PP1 is intentional (the convert skips - # auto-pp=world_size when EP>1). torch_dist reshards on load, so EP32-saved loads - # fine into the EP16 trainer. - torch_dist_prep_nodes=4, - torch_dist_prep_gpus_per_node=8, - torch_dist_convert_extra_args=( - "--tensor-model-parallel-size 1 --pipeline-model-parallel-size 1 --expert-model-parallel-size 32" - ), - # Each node buffers ~700 GB of distcp shards in the Volume write cache before commit; - # 2 TB gives rank 0 (shards + .metadata + common.pt) headroom over the 800 GB pool default. - torch_dist_prep_ephemeral_disk_mib=2_097_152, -) - - -class _Miles(MilesConfig): - # Architecture is sourced from the model script (MLA + the full DeepSeek-MoE - # arg set). Shared with Kimi-K2-Thinking, whose arch K2.6 matches. - miles_model_script = "scripts/models/kimi-k2-thinking.sh" - - # Checkpoints (absolute -> launcher skips HF download). hf_checkpoint is the - # served NVFP4 base; ref_load is the Megatron torch_dist checkpoint (built from - # the bf16 masters by prepare_torch_dist). Raw mode can ONLY load torch_dist - # (miles' HF load is bridge-only); the conversion uses the KimiK25 mbridge. - hf_checkpoint = f"{PREP_PATH}/{MODEL_TAG}/nvfp4" - ref_load = f"{PREP_PATH}/{MODEL_TAG}/torch_dist" - # "raw" builds the GPTModel from the model script's MODEL_ARGS (MLA + DeepSeek- - # MoE), NOT via megatron-bridge AutoBridge. K2.6's HF arch is the VLM wrapper - # KimiK25ForConditionalGeneration, which AutoBridge does not support (it knows - # DeepseekV3ForCausalLM etc.) — so "bridge" crashes in init. This is how slime - # builds DeepSeek-V3/MLA MoE models. Export still routes via model_name below. - megatron_to_hf_mode = "raw" - model_name = ( - "kimi_k25" # megatron_to_hf export dispatch (convert_kimi_k25_to_hf + NVFP4) - ) - - # Disaggregated publish-only rollout; modal_train fills rollout_endpoint_url. - actor_num_nodes = 16 # 16x8 B200 = 128 GPUs (trainer only; pool is elastic on top) - actor_num_gpus_per_node = 8 - num_gpus_per_node = 8 - colocate = False - rollout_num_gpus = 0 - rollout_num_gpus_per_engine = 4 # B200:4 per rollout container (K2.6 NVFP4 fits) - rollout_endpoint_url = None - use_miles_router = True - - # Staleness gate. The hook path is a real miles CLI arg; the gate KNOBS ride - # in custom_config_path (consumed by the hook, not miles core). modal_train - # merges the dynamic bulletin identity (run_id, volume, app) into this dict. - custom_rollout_request_hook_path = ( - "cookbook.miles_disagg.hooks.gated_rollout_request_hook" - ) - custom_config_path = { - "rollout_request_weight_version_mode": "min", - "rollout_request_weight_version_lag": 1, - # 240×1s=240s was too short: after a (re)deploy the pool cold-loads the ~591 GB - # base for ~16 min, and the trainer can reach iter-1 rollouts first -> the rollout - # exhausted retries on 503s and failed. 1200×1s=20 min outlasts a full cold-load, - # so rollouts ride out a cold/restarting pool instead of dying on the race. - "rollout_request_retry_attempts": 1200, - "rollout_request_retry_sleep": 1.0, - "rollout_session_affinity_header": "Modal-Session-ID", - # Finite per-request read timeout for the disagg /generate client (init_http_client). - # Without it the client uses httpx.Timeout(None) and a request to a Flash container - # that scaled down mid-flight hangs forever (rollout stalls at ~N%, engines idle). - # 300s comfortably exceeds a full 4096-token generation (~130s at ~30 tok/s) but - # recovers a lost request fast: it errors -> the retry above reroutes it. - "rollout_request_timeout_secs": 300, - } - - async_mode = True - update_weights_interval = 1 - - # NVFP4 QAT — canonical miles recipe per radixark/miles#1261 (NVFP4 RL). - fp4_format = "e2m1" - fp4_recipe = "nvfp4" - # fp4_param_gather=False is canonical (#1261 never sets --fp4-param-gather): - # keeps NVFP4 GEMM compute QAT (config.fp4 in raw mode) with bf16 master params. - # With it True, params are TE NVFP4Tensor and Megatron DDP's param-buffer repoint - # (modify_underlying_storage -> TE replace_raw_data) crashes (TE: FP8 yes, NVFP4 no). - fp4_param_gather = False - # Per-module TE precision config (#1261's mechanism): NVFP4 ONLY on the routed - # expert GEMMs, everything else bf16 — matches the experts-only served base. - # Materialized to a temp YAML and passed as --te-precision-config-file. - te_precision_config_file = { - "configs": { - "nvfp4": { - "transformer_engine_config_type": "TEQuantizationParams", - "training_recipe": {"fp4_quantization_recipe": "nvfp4"}, - }, - "bf16": { - "transformer_engine_config_type": "TEQuantizationParams", - "training_recipe": {}, - }, - }, - "matchers": { - "routed_experts_fc1_nvfp4": { - "type": "glob", - "enabled": True, - "pattern": "*.mlp.experts.linear_fc1", - "config": "nvfp4", - }, - "routed_experts_fc2_nvfp4": { - "type": "glob", - "enabled": True, - "pattern": "*.mlp.experts.linear_fc2", - "config": "nvfp4", - }, - "default_bf16": { - "type": "glob", - "enabled": True, - "pattern": "*", - "config": "bf16", - }, - }, - } - # bf16 carve-out for the dense first layer (FIRST_K_DENSE_REPLACE=1). Passed to BOTH - # convert_hf_to_nvfp4 (served base) and the trainer so they agree. - # 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.) - num_layers_at_start_in_bf16 = 1 - num_layers_at_end_in_bf16 = 0 - - # Disk-delta publish-only over the Modal Volume bulletin board. - update_weight_transfer_mode = "disk" - update_weight_disk_delta = True - update_weight_delta_encoding = "xor" - update_weight_delta_checksum = "xxh3-128" - update_weight_disk_dir = DELTA_BULLETIN_ROOT - custom_delta_pre_push_path = "cookbook.miles_disagg.hooks.commit_and_wake" - - # Data: dapo-math-17k. - prompt_data = f"{DATA_PATH}/dapo-math-17k/dapo-math-17k.jsonl" - input_key = "prompt" - label_key = "label" - apply_chat_template = True - rollout_shuffle = True - balance_data = True - rm_type = "deepscaler" - eval_interval = None - - # Rollout (bring-up smoke length; scale num_rollout for a real run). - # Sized for FAST loop closure (proving the publish), not a real run: a small - # batch and short responses keep the first rollout to minutes, not the ~1.6 h - # the full 128x8x16384 took (per-sequence decode of 16 K-token responses, not - # pool size, is the bound). Restore rollout_batch_size=128 / - # rollout_max_response_len=16384 once v1 publishes and we scale num_rollout. - num_rollout = 3 # trimmed e2e loop-closure smoke (scale up for a real run) - save_interval = 20 # megatron requires it; > num_rollout so the smoke skips saves - rollout_batch_size = 32 - rollout_max_response_len = 4096 - rollout_temperature = 0.8 - n_samples_per_prompt = 8 - global_batch_size = 256 - use_dynamic_global_batch_size = True - sglang_server_concurrency = 256 - - use_rollout_routing_replay = True - - # Trainer parallelism, sized for 16x8=128 GPUs. Derived from the proven Kimi - # 32x8 layout (TP8/PP8/CP4/EP32) by halving the non-PP width: TP8*PP8*CP2 = - # 128 (DP1), and for the experts ETP1*EP16 = 16 = TP8*CP2*DP1. Keeping TP8/PP8 - # preserves the per-stage layer split (and decoder_last_pipeline_num_layers). - tensor_model_parallel_size = 8 - sequence_parallel = True - pipeline_model_parallel_size = 8 - context_parallel_size = 2 - expert_model_parallel_size = 16 - expert_tensor_parallel_size = 1 - decoder_last_pipeline_num_layers = 5 - use_dynamic_batch_size = True - max_tokens_per_gpu = 16384 - recompute_granularity = "full" - recompute_method = "uniform" - recompute_num_layers = 1 - attention_dropout = 0.0 - hidden_dropout = 0.0 - accumulate_allreduce_grads_in_fp32 = True - attention_softmax_in_fp32 = True - no_check_for_nan_in_loss_and_grad = True - - # Optimizer. - optimizer = "adam" - lr = 1e-6 - lr_decay_style = "constant" - weight_decay = 0.1 - adam_beta1 = 0.9 - adam_beta2 = 0.98 - optimizer_cpu_offload = True - overlap_cpu_optimizer_d2h_h2d = True - use_precision_aware_optimizer = True - - # Algorithm (GRPO + truncated importance sampling). - advantage_estimator = "grpo" - eps_clip = 0.2 - eps_clip_high = 0.28 - use_kl_loss = True - kl_loss_coef = 0.0 - kl_loss_type = "low_var_kl" - entropy_coef = 0.0 - use_tis = True - - environment = { - "CUDA_DEVICE_MAX_CONNECTIONS": "1", - "NCCL_NVLS_ENABLE": "1", - "NVSHMEM_DISABLE_NCCL": "1", - "NCCL_TIMEOUT_MS": "360000000", - # NVFP4 numerics (radixark/miles#1261 NVFP4 train env). Without these the - # NVFP4 QAT is mis-configured even once the build/DDP/load gaps are cleared. - "NVTE_NVFP4_DISABLE_2D_QUANTIZATION": "1", - "NVTE_NVFP4_DISABLE_RHT": "1", - "NVTE_NVFP4_DISABLE_STOCHASTIC_ROUNDING": "1", - "NVTE_NVFP4_ROW_SCALED_ACTIVATION": "1", - "NVTE_BACKWARD_OVERRIDE": "high_precision", - "NVTE_USE_FAST_MATH": "0", - } - - def prepare_data(self) -> None: - from datasets import load_dataset - - ds = load_dataset("BytedTsinghua-SIA/DAPO-Math-17k", split="train") - ds = ds.shuffle(seed=42).select(range(min(50000, ds.num_rows))) - ds = ds.map(lambda ex: {"label": ex["reward_model"]["ground_truth"]}) - ds = ds.select_columns(["prompt", "label"]) - ds.to_json(f"{DATA_PATH}/dapo-math-17k/dapo-math-17k.jsonl") - - -miles = _Miles() diff --git a/cookbook/miles_disagg/configs/moonlight_nvfp4.py b/cookbook/miles_disagg/configs/moonlight_nvfp4.py new file mode 100644 index 0000000..20abd71 --- /dev/null +++ b/cookbook/miles_disagg/configs/moonlight_nvfp4.py @@ -0,0 +1,177 @@ +"""Moonlight-16B-A3B GRPO on Modal, disaggregated, native-NVFP4 end to end — the small runnable de-risk for kimi_k2_6_nvfp4. + +Deploy: EXPERIMENT_CONFIG=moonlight_nvfp4 uv run --extra modal modal deploy --strategy recreate -m cookbook.miles_disagg.app +""" + +from __future__ import annotations + +from cookbook.common.config import ModalConfig +from cookbook.common.constants import DATA_PATH, PREP_PATH +from cookbook.miles_disagg.config import MilesConfig + + +APP_NAME = "stitch-moonlight-nvfp4" +DELTA_VOLUME_NAME = "stitch-delta-moonlight-nvfp4" +DELTA_BULLETIN_ROOT = "/delta-bulletin" +LOCAL_CHECKPOINT_PATH = "/local-checkpoint" + +SOURCE_MODEL = "moonshotai/Moonlight-16B-A3B-Instruct" +MODEL_TAG = "moonlight-16b-nvfp4" # names the prepared dirs under PREP_PATH + +# in_place applies weights without draining in-flight rollouts; stale KV isolated per version. +SIDECAR_COMMIT_MODE = "in_place" +SIDECAR_DEBUG_REQUESTS = True +# R3 routing-replay needs the dropless Megatron dispatch fix at startup. +MEGATRON_RUNTIME_PATCHES = [ + "/root/cookbook/miles_disagg/patches/megatron-r3-dispatch.patch", +] + + +# No --quantization flag — NVFP4 comes from the served checkpoint's quant config. +# mem-fraction / context-length are STARTING POINTS; measure. +SGLANG_SERVER_ARGS = { + "--weight-loader-prefetch-checkpoints": "", + "--weight-loader-prefetch-num-threads": "8", + "--attention-backend": "tokenspeed_mla", + "--kv-cache-dtype": "fp8_e4m3", # tokenspeed_mla requires this + "--context-length": "8192", # Moonlight's max_position_embeddings + "--mem-fraction-static": "0.8", + "--chunked-prefill-size": "4096", + "--skip-server-warmup": "", + # routing replay: pool emits per-token routed experts for the trainer to replay. + "--enable-return-routed-experts": "", +} + +SGLANG_ENV = {"SGLANG_ENABLE_RELOAD_LOAD_PLAN": "1"} # NVFP4: load-plan replay + O(delta) partial reload + +modal = ModalConfig( + gpu="B200", + region="us", + # warm floor of 1: the pool must be UP before the trainer sends rollouts (scale-from-0 + # raced the retry budget). Flash still scales ABOVE this under load. + rollout_min_containers=1, + proxy_regions=["us-west"], +) + + +class _Miles(MilesConfig): + # Arch comes from the model script; do NOT inline arch attrs here. + miles_model_script = "scripts/models/moonlight.sh" + + hf_checkpoint = f"{PREP_PATH}/{MODEL_TAG}/nvfp4" + ref_load = f"{PREP_PATH}/{MODEL_TAG}/bf16" + megatron_to_hf_mode = "bridge" + model_name = "deepseekv3" # bridge dispatch: Moonlight is DeepSeek-V3 arch + + actor_num_nodes = 1 + actor_num_gpus_per_node = 4 # 1 node x 4 B200 trainer (matches the proven moonlight recipe) + num_gpus_per_node = 4 + colocate = False # disk-delta is incompatible with --colocate + rollout_num_gpus = 0 # publish-only forces this; set explicitly for clarity + rollout_num_gpus_per_engine = 1 # B200:1 per rollout container (Moonlight NVFP4 is tiny) + rollout_endpoint_url = None + use_miles_router = True + + # Staleness gate; the knobs ride in custom_config_path (read by the hook, not miles core). + custom_rollout_request_hook_path = "cookbook.common.hooks.gated_rollout_request_hook" + custom_config_path = { + "rollout_request_weight_version_mode": "min", + "rollout_request_weight_version_lag": 1, # bounded staleness window + "rollout_request_retry_attempts": 240, + "rollout_request_retry_sleep": 1.0, + "rollout_session_affinity_header": "Modal-Session-ID", + } + + # async-first: one-step off-policy; publish weights every step. + async_mode = True + update_weights_interval = 1 + + # NVFP4 QAT (native Megatron FP4; Blackwell + TE >= 2.7.0.dev0). + # fp4_param_gather stays False: with it True, TE NVFP4Tensor params crash Megatron DDP. + fp4_format = "e2m1" + fp4_param_gather = False + + update_weight_transfer_mode = "disk-delta" + update_weight_delta_encoding = "xor" + update_weight_delta_checksum = "xxh3-128" + update_weight_disk_dir = DELTA_BULLETIN_ROOT # modal_train run-scopes this + custom_update_weight_post_write_path = "cookbook.common.hooks.commit_and_wake" + + prompt_data = f"{DATA_PATH}/dapo-math-17k/dapo-math-17k.jsonl" + input_key = "prompt" + label_key = "label" + apply_chat_template = True + rollout_shuffle = True + balance_data = True + rm_type = "deepscaler" + eval_interval = None + + num_rollout = 20 + save_interval = 1000 # megatron requires it; > num_rollout so the smoke skips megatron saves + rollout_batch_size = 32 + rollout_max_response_len = 4096 # fits within the 8192 context (prompt + response) + rollout_temperature = 0.8 + n_samples_per_prompt = 8 + global_batch_size = 128 + use_dynamic_global_batch_size = True + sglang_server_concurrency = 64 + + # R3: replay sglang's routed experts in the train/log-prob forward. + use_rollout_routing_replay = True + + # Trainer parallelism (the proven moonlight setting: TP2/SP/PP1/CP1/EP4). + tensor_model_parallel_size = 2 + sequence_parallel = True + pipeline_model_parallel_size = 1 + context_parallel_size = 1 + expert_model_parallel_size = 4 + expert_tensor_parallel_size = 1 + use_dynamic_batch_size = True + max_tokens_per_gpu = 8192 + recompute_granularity = "full" + recompute_method = "uniform" + recompute_num_layers = 1 + attention_dropout = 0.0 + hidden_dropout = 0.0 + accumulate_allreduce_grads_in_fp32 = True + attention_softmax_in_fp32 = True + no_check_for_nan_in_loss_and_grad = True + + # Optimizer (CPU offload keeps GPU state tiny for ~3B active). + optimizer = "adam" + lr = 1e-6 + lr_decay_style = "constant" + weight_decay = 0.1 + adam_beta1 = 0.9 + adam_beta2 = 0.98 + optimizer_cpu_offload = True + overlap_cpu_optimizer_d2h_h2d = True + use_precision_aware_optimizer = True + + advantage_estimator = "grpo" + eps_clip = 0.2 + eps_clip_high = 0.28 + use_kl_loss = True + kl_loss_coef = 0.0 + kl_loss_type = "low_var_kl" + entropy_coef = 0.0 + use_tis = True + + environment = { + "CUDA_DEVICE_MAX_CONNECTIONS": "1", + "NCCL_NVLS_ENABLE": "1", + "NVSHMEM_DISABLE_NCCL": "1", + "NCCL_TIMEOUT_MS": "360000000", + } + + def prepare_data(self) -> None: + from datasets import load_dataset + + ds = load_dataset("BytedTsinghua-SIA/DAPO-Math-17k", split="train") + ds = ds.shuffle(seed=42).select(range(min(50000, ds.num_rows))) + ds = ds.map(lambda ex: {"label": ex["reward_model"]["ground_truth"]}) + ds = ds.select_columns(["prompt", "label"]) + ds.to_json(f"{DATA_PATH}/dapo-math-17k/dapo-math-17k.jsonl") + + +miles = _Miles() diff --git a/cookbook/miles_disagg/configs/moonlight_nvfp4_disagg.py b/cookbook/miles_disagg/configs/moonlight_nvfp4_disagg.py deleted file mode 100644 index 2000599..0000000 --- a/cookbook/miles_disagg/configs/moonlight_nvfp4_disagg.py +++ /dev/null @@ -1,236 +0,0 @@ -"""Moonlight-16B-A3B GRPO on Modal, disaggregated, native-NVFP4 end to end. - -This is the small, *runnable* de-risk for the Kimi-K2.6-NVFP4 recipe -(``kimi_k2_6_nvfp4_disagg``). Moonlight is the same DeepSeek-V3 architecture -family (MLA + DeepSeek-MoE: sigmoid router, shared experts, grouped top-k) at a -size that fits a single Blackwell node, so it exercises the full -QAT → NVFP4 export → XOR-delta → SGLang-reload loop cheaply. - -WHY NVFP4 IS ALL-BLACKWELL --------------------------- -NVFP4 QAT in miles is native Megatron ``--fp4-format e2m1`` → ``NVFP4BlockScaling`` -(TransformerEngine ≥ 2.7.0.dev0). FP4 GEMM requires Blackwell, so BOTH the -trainer and the rollout pool run on B200 — unlike the INT4 recipe, which fake- -quantizes (straight-through) on H200. There is no simulated/non-Blackwell NVFP4 -weight-QAT path. - -CHECKPOINT LIFECYCLE (three roles; see modal_train.prepare_checkpoints) ------------------------------------------------------------------------ - * BF16 masters (``ref_load``): Moonlight ships bf16, so the masters are the - downloaded checkpoint directly (no dequant needed — unlike Kimi, which is - published as INT4 and must be dequantized). - * Served NVFP4 base (``hf_checkpoint``): produced by miles' own - ``tools/convert_hf_to_nvfp4.py`` from the bf16 masters. Because the SAME - TE reference quantizer produces the base and the trainer's export, the - served packing == the export packing BY CONSTRUCTION, the step-0 export - equals the base, and the first XOR delta is ~empty. This is exactly what - de-risks the byte-exact axis that the Kimi recipe must verify on a warm - container. - * Megatron torch_dist (``load``/``save``): the trainer's own rollout - checkpoints — never seen by the rollout pool. - -The trainer reads ``hf_checkpoint`` (the NVFP4 base) for both the export -quantization_config and the diff baseline; ``_capture_baseline`` seeds the -snapshot from those exact bytes, so applying delta_vN reproduces export_vN -byte-for-byte and the served weights become the trainer's NVFP4 export. - -Deploy as its own app: - EXPERIMENT_CONFIG=moonlight_nvfp4_disagg \ - uv run --extra modal modal deploy --strategy recreate -m cookbook.miles_disagg.modal_train -""" - -from __future__ import annotations - -from cookbook.miles_disagg.configs.base import DATA_PATH, PREP_PATH, ModalConfig, MilesConfig - - -APP_NAME = "miles-moonlight-nvfp4-disagg" -DELTA_VOLUME_NAME = "miles-delta-bulletin-moonlight-nvfp4" -DELTA_BULLETIN_ROOT = "/delta-bulletin" - -# Source HF repo (bf16) the prepare step turns into masters + NVFP4 base. -SOURCE_MODEL = "moonshotai/Moonlight-16B-A3B-Instruct" -MODEL_TAG = "moonlight-16b-nvfp4" # names the prepared dirs under PREP_PATH - -# Async one-step off-policy: in_place applies weights without draining in-flight -# rollouts; stale-version KV is isolated per weight version by the sidecar. -SIDECAR_COMMIT_MODE = "in_place" -SIDECAR_DEBUG_REQUESTS = True - - -def build_serving_image(**kwargs): - """Per-experiment rollout-pool image (modal_train picks this up if present).""" - from cookbook.miles_disagg.serving import build_nvfp4_b200_serving_image - - return build_nvfp4_b200_serving_image(**kwargs) - - -# Moonlight NVFP4 on a B200 pool. The trainer image injects --served-model-name / -# --dtype / --trust-remote-code; only the MLA/MoE/cache extras live here. No -# --quantization flag — NVFP4 is driven by the served checkpoint's own -# quant config. mem-fraction / context-length are STARTING POINTS; measure. -SGLANG_SERVER_ARGS = { - "--attention-backend": "tokenspeed_mla", - "--kv-cache-dtype": "fp8_e4m3", # tokenspeed_mla requires this - "--context-length": "8192", # Moonlight's max_position_embeddings - "--mem-fraction-static": "0.8", - "--chunked-prefill-size": "4096", - "--skip-server-warmup": "", - # Routing replay: the pool emits per-token routed experts so the trainer can - # replay them (DeepSeek-V3-arch MoE supports it). - "--enable-return-routed-experts": "", -} - -modal = ModalConfig( - gpu="B200", - region="us", - # Warm floor of 1: the pool must be UP before the trainer sends rollouts. - # With min=0 the pool only scales on rollout traffic, and scale-from-0 either - # didn't trigger or lost the race against the rollout retry budget, stalling - # the trainer in model-load. Flash still scales ABOVE this floor under load. - rollout_min_containers=1, - proxy_regions=["us-west"], -) - - -class _Miles(MilesConfig): - # Architecture comes from the model script (MLA + the full DeepSeek-MoE arg - # set); do NOT inline arch attrs here. - miles_model_script = "scripts/models/moonlight.sh" - - # Checkpoints (absolute paths -> the launcher skips HF download for them). - # hf_checkpoint is the served NVFP4 base; ref_load is the bf16 masters. - hf_checkpoint = f"{PREP_PATH}/{MODEL_TAG}/nvfp4" - ref_load = f"{PREP_PATH}/{MODEL_TAG}/bf16" - megatron_to_hf_mode = "bridge" - model_name = "deepseekv3" # bridge dispatch: Moonlight is DeepSeek-V3 arch - - # Disaggregated publish-only rollout; modal_train fills rollout_endpoint_url. - actor_num_nodes = 1 - actor_num_gpus_per_node = 4 # 1 node x 4 B200 trainer (matches the proven moonlight recipe) - num_gpus_per_node = 4 - colocate = False # disk-delta is incompatible with --colocate - rollout_num_gpus = 0 # publish-only forces this; set explicitly for clarity - rollout_num_gpus_per_engine = 1 # B200:1 per rollout container (Moonlight NVFP4 is tiny) - rollout_endpoint_url = None - use_miles_router = True - - # Staleness gate. custom_rollout_request_hook_path is a real miles CLI arg; - # the gate KNOBS are consumed only by the hook (not miles core), so they ride - # in custom_config_path — miles setattr's every YAML key onto the args - # namespace, and the hook reads them via getattr(args, ...). modal_train - # merges the dynamic bulletin identity (run_id, volume, app) into this dict. - custom_rollout_request_hook_path = "cookbook.miles_disagg.hooks.gated_rollout_request_hook" - custom_config_path = { - "rollout_request_weight_version_mode": "min", - "rollout_request_weight_version_lag": 1, # bounded staleness window - "rollout_request_retry_attempts": 240, - "rollout_request_retry_sleep": 1.0, - "rollout_session_affinity_header": "Modal-Session-ID", - } - - # Async-first: train_async pipelines generate(N+1) with train(N). - async_mode = True - update_weights_interval = 1 - - # ── NVFP4 QAT (native Megatron FP4 training; Blackwell + TE >= 2.7.0.dev0) ── - # --fp4-format takes the element format 'e2m1'; the recipe (fp4_recipe) already - # defaults to 'nvfp4' (NVFP4BlockScaling). --fp4-param-gather keeps params in - # fp4 to save memory. NVFP4 group size is fixed at 16. - fp4_format = "e2m1" - fp4_param_gather = True - - # ── Disk-delta publish-only over the Modal Volume bulletin board ── - update_weight_transfer_mode = "disk" - update_weight_disk_delta = True - update_weight_delta_encoding = "xor" - update_weight_delta_checksum = "xxh3-128" - update_weight_disk_dir = DELTA_BULLETIN_ROOT # modal_train run-scopes this - custom_delta_pre_push_path = "cookbook.miles_disagg.hooks.commit_and_wake" - - # Data: dapo-math-17k (the math-reasoning set the Kimi recipe uses). - prompt_data = f"{DATA_PATH}/dapo-math-17k/dapo-math-17k.jsonl" - input_key = "prompt" - label_key = "label" - apply_chat_template = True - rollout_shuffle = True - balance_data = True - rm_type = "deepscaler" - eval_interval = None - - # Rollout (bring-up smoke length; scale num_rollout for a real run). - num_rollout = 20 - save_interval = 1000 # megatron requires it; > num_rollout so the smoke skips megatron saves - rollout_batch_size = 32 - rollout_max_response_len = 4096 # fits within the 8192 context (prompt + response) - rollout_temperature = 0.8 - n_samples_per_prompt = 8 - global_batch_size = 128 - use_dynamic_global_batch_size = True - sglang_server_concurrency = 64 - - # Routing replay (R3): replay sglang's routed experts in the train/log-prob - # forward. Being root-caused — instrumented run to capture the MoE token-count - # mismatch behind "Split sizes doesn't match total dim 0 size". - use_rollout_routing_replay = True - - # Trainer parallelism (the proven moonlight setting: TP2/SP/PP1/CP1/EP4). - tensor_model_parallel_size = 2 - sequence_parallel = True - pipeline_model_parallel_size = 1 - context_parallel_size = 1 - expert_model_parallel_size = 4 - expert_tensor_parallel_size = 1 - use_dynamic_batch_size = True - max_tokens_per_gpu = 8192 - recompute_granularity = "full" - recompute_method = "uniform" - recompute_num_layers = 1 - attention_dropout = 0.0 - hidden_dropout = 0.0 - accumulate_allreduce_grads_in_fp32 = True - attention_softmax_in_fp32 = True - no_check_for_nan_in_loss_and_grad = True - - # Optimizer (CPU offload keeps GPU state tiny for ~3B active). - optimizer = "adam" - lr = 1e-6 - lr_decay_style = "constant" - weight_decay = 0.1 - adam_beta1 = 0.9 - adam_beta2 = 0.98 - optimizer_cpu_offload = True - overlap_cpu_optimizer_d2h_h2d = True - use_precision_aware_optimizer = True - - # Algorithm (GRPO + truncated importance sampling). - advantage_estimator = "grpo" - eps_clip = 0.2 - eps_clip_high = 0.28 - use_kl_loss = True - kl_loss_coef = 0.0 - kl_loss_type = "low_var_kl" - entropy_coef = 0.0 - use_tis = True - - # Ray runtime environment. - environment = { - "CUDA_DEVICE_MAX_CONNECTIONS": "1", - "NCCL_NVLS_ENABLE": "1", - "NVSHMEM_DISABLE_NCCL": "1", - "NCCL_TIMEOUT_MS": "360000000", - } - - def prepare_data(self) -> None: - from datasets import load_dataset - - # DAPO-Math-17k columns are `prompt` (chat list) and the gold answer under - # `reward_model.ground_truth` — lift the answer into a `label` column. - ds = load_dataset("BytedTsinghua-SIA/DAPO-Math-17k", split="train") - ds = ds.shuffle(seed=42).select(range(min(50000, ds.num_rows))) - ds = ds.map(lambda ex: {"label": ex["reward_model"]["ground_truth"]}) - ds = ds.select_columns(["prompt", "label"]) - ds.to_json(f"{DATA_PATH}/dapo-math-17k/dapo-math-17k.jsonl") - - -miles = _Miles() diff --git a/cookbook/miles_disagg/configs/qwen3_30b_a3b_nvfp4_46.py b/cookbook/miles_disagg/configs/qwen3_30b_a3b_nvfp4_46.py new file mode 100644 index 0000000..0c6577a --- /dev/null +++ b/cookbook/miles_disagg/configs/qwen3_30b_a3b_nvfp4_46.py @@ -0,0 +1,342 @@ +"""Qwen3-30B-A3B GRPO on Modal, disaggregated — the humans& NVFP4 "4-bitter" recipe. + +The full recipe from https://humansand.ai/blog/nvfp4-rl (miles PR #1261), on the blog's +own model/dataset (Qwen3-30B-A3B, DAPO-Math-17k), end to end on B200: + + 1. Per-token row-scaled NVFP4 forward on the MoE experts + (NVTE_NVFP4_ROW_SCALED_ACTIVATION=1, RHT/stochastic-rounding/2D off). + 2. Dequantized backward — the BF16 backward differentiates DQ(Q(w)), the same quantized + function the forward evaluated (NVTE_BACKWARD_OVERRIDE=dequantized; TE 2.17 + the + TE#3141 save-original-input fix, applied in TRAINER_IMAGE_SETUP). + 3. Four-over-six adaptive block scaling on weights AND activations + (NVTE_NVFP4_4OVER6=all / FLASHINFER_NVFP4_4OVER6=1), the trainer and sampler sides of + the env contract pinned to identical conventions: MAE error mode, e4m3-max 256, fast + math off. + 4. Selective precision: the last 15% of layers (7 of 48) stay BF16; Qwen3 MoE has no + shared expert, so the layer carve-out is the whole of this axis. + +WHY THE WEIGHT BYTES ARE EXACT BY CONSTRUCTION +---------------------------------------------- +miles' export and tools/convert_hf_to_nvfp4.py both use the same TE-direct NVFP4Quantizer +(PR #1261), so served packing == export packing and the delta baseline aligns — the +FlashInfer weight quantizer is never on the served path. PREP_ENV pins the conversion to +the trainer's exact NVTE_* settings; a drift there fails loud on the first delta apply. +FlashInfer's runtime role is activation quantization only, where the FLASHINFER_* mirror +keeps rollout logprobs near the trainer's (train_rollout_logprob_abs_diff ~= 0.044 in the +validated smoke; PR #1261 reports 0.031). + +STACK (deltas over the shared images, via the per-experiment hooks) +------------------------------------------------------------------- +Trainer: the shared dated base + TE 2.17.0 cu13 (dequantized backward, 4/6, row-scaled) ++ miles branch ``nvfp4-46-recipe-v2`` (the shared pin merged with radixark PR #1261). +Serving: the shared stitch-sglang-v0.5.15 fork image + flashinfer 0.6.13 +(python/cubin/jit-cache in lockstep — 0.6.13 carries the 4/6 error-domain fix, FI#3448) ++ the FLASHINFER_* side of the quantizer contract. + +GPU footprint: trainer 1x8 B200 + rollout pool 1-3x1 B200 (<= 11 concurrent). + +Deploy: EXPERIMENT_CONFIG=qwen3_30b_a3b_nvfp4_46 uv run --extra modal modal deploy --strategy recreate -m cookbook.miles_disagg.app +""" + +from __future__ import annotations + +from cookbook.common.config import ModalConfig +from cookbook.common.constants import DATA_PATH, PREP_PATH +from cookbook.miles_disagg.config import MilesConfig + + +APP_NAME = "stitch-qwen3-30b-nvfp4-46" +DELTA_VOLUME_NAME = "stitch-delta-qwen3-30b-nvfp4-46" +DELTA_BULLETIN_ROOT = "/delta-bulletin" +LOCAL_CHECKPOINT_PATH = "/local-checkpoint" + +SOURCE_MODEL = "Qwen/Qwen3-30B-A3B" # bf16 -> masters ARE the download +MODEL_TAG = "qwen3-30b-a3b-nvfp4-46" + +SIDECAR_COMMIT_MODE = "in_place" +SIDECAR_DEBUG_REQUESTS = True +# R3 routing-replay needs the dropless Megatron dispatch fix at startup. +MEGATRON_RUNTIME_PATCHES = [ + "/root/cookbook/miles_disagg/patches/megatron-r3-dispatch.patch", +] + +# ── The 4/6 recipe's miles + TE stack (overrides the shared trainer image) ──── +# The shared miles pin merged with radixark/miles PR #1261 (the humans& recipe: +# TE-direct NVFP4 export, 4/6 support, env forwarding). +MILES_REPO_REF = "a80810a9b835f33aa9b099324df0f8ec5c2558ce" # branch nvfp4-46-recipe-v2 + +# The TE#3141 fix TE 2.17 is missing: without it the dequantized backward override still +# saves the pre-quantization input, so the wgrad is computed against different operands +# than the forward used. Applied as a self-verifying edit — the upstream patch file's +# hunks were cut against TE main and reject on the v2.17 layout. Drop at TE >= 2.18. +_TE_DEQUANT_BWD_FIX = """python - <<'EOF' +import importlib.util, pathlib +te = pathlib.Path(importlib.util.find_spec("transformer_engine").submodule_search_locations[0]) +edits = { + "pytorch/module/linear.py": ( + 'if backward_override == "high_precision":\\n save_original_input = True\\n', + 'if backward_override == "high_precision":\\n save_original_input = True\\n' + ' elif backward_override == "dequantized":\\n save_original_input = False\\n', + ), + "pytorch/module/grouped_linear.py": ( + 'if backward_override == "high_precision":\\n save_original_input = True\\n', + 'if backward_override == "high_precision":\\n save_original_input = True\\n' + ' elif backward_override == "dequantized":\\n save_original_input = False\\n', + ), +} +for rel, (old, new) in edits.items(): + p = te / rel + src = p.read_text() + assert src.count(old) == 1, f"{rel}: expected exactly 1 match, got {src.count(old)}" + p.write_text(src.replace(old, new)) + print(f"patched {rel}") +EOF""" + +TRAINER_IMAGE_SETUP = ( + # TE 2.17 on the cu13 base (matches PR #1261's own Dockerfile). The torch extension + # builds from sdist against the image's torch 2.11. + "pip uninstall -y transformer_engine transformer_engine_cu12 transformer_engine_cu13 transformer_engine_torch" + " ; pip install --no-deps transformer_engine==2.17.0" + " && pip install transformer_engine_cu13==2.17.0 nvidia-mathdx==25.6.0", + "pip -v install --no-build-isolation transformer_engine_torch==2.17.0", + _TE_DEQUANT_BWD_FIX, +) + +# ── Trainer/sampler 4/6 + NVFP4 quantizer contract ──────────────────────────── +# Both sides must agree on: 4/6 on, MAE candidate-error mode, e4m3-max 256 under 4/6, +# fast math off. The NVTE_* side drives TE (training, export, and the served-base +# conversion via PREP_ENV); the FLASHINFER_* side drives the pool's activation +# quantization kernels. +_NVTE_QUANT_ENV = { + "NVTE_NVFP4_DISABLE_2D_QUANTIZATION": "1", + "NVTE_NVFP4_DISABLE_RHT": "1", + "NVTE_NVFP4_DISABLE_STOCHASTIC_ROUNDING": "1", + "NVTE_NVFP4_ROW_SCALED_ACTIVATION": "1", + "NVTE_NVFP4_4OVER6": "all", + "NVTE_NVFP4_4OVER6_ERR_MODE": "MAE", + "NVTE_NVFP4_4OVER6_E4M3_USE_256": "all", + "NVTE_BACKWARD_OVERRIDE": "dequantized", + "NVTE_USE_FAST_MATH": "0", + "TRTLLM_DISABLE_FP4_QUANT_FAST_MATH": "1", +} +_FLASHINFER_QUANT_ENV = { + "FLASHINFER_NVFP4_4OVER6": "1", + "FLASHINFER_NVFP4_4OVER6_ERR_MODE": "MAE", + "FLASHINFER_NVFP4_4OVER6_E4M3_USE_256": "1", + "FLASHINFER_NVFP4_4OVER6_ERR_USE_FAST_MATH": "0", + "FLASHINFER_DISABLE_FP4_QUANT_FAST_MATH": "1", + "TRTLLM_DISABLE_FP4_QUANT_FAST_MATH": "1", + "SGLANG_FLASHINFER_NVFP4_PER_TOKEN_ACTIVATION": "1", +} + +# prepare_checkpoints must quantize the served base under the trainer's exact settings. +PREP_ENV = dict(_NVTE_QUANT_ENV) + +# flashinfer 0.6.13: the 4/6 error-domain fix (FI#3448). The three packages move in +# lockstep — sglang hard-fails on a version mismatch, correctly, since a stale jit-cache +# would silently serve kernels from the old version. +SERVING_IMAGE_EXTRA_COMMANDS = ( + "pip install 'flashinfer_python[cu13]==0.6.13' 'flashinfer-cubin==0.6.13'" + " && pip install flashinfer-jit-cache==0.6.13 --index-url https://flashinfer.ai/whl/cu130", +) +SERVING_IMAGE_ENV = dict(_FLASHINFER_QUANT_ENV) + +SGLANG_ENV = {"SGLANG_ENABLE_RELOAD_LOAD_PLAN": "1"} # NVFP4: load-plan replay + O(delta) partial reload + +# Qwen3 MoE is GQA (no MLA): trtllm_mha attention + the routed FlashInfer TRTLLM MoE +# runner (emits per-token routed experts for R3 replay). NVFP4 comes from the served +# checkpoint's own quant config — no --quantization flag. +SGLANG_SERVER_ARGS = { + "--weight-loader-prefetch-checkpoints": "", + "--weight-loader-prefetch-num-threads": "8", + "--attention-backend": "trtllm_mha", + "--moe-runner-backend": "flashinfer_trtllm_routed", + "--kv-cache-dtype": "bfloat16", + "--context-length": "16384", + "--mem-fraction-static": "0.8", + "--chunked-prefill-size": "4096", + "--skip-server-warmup": "", + "--enable-return-routed-experts": "", +} + +modal = ModalConfig( + gpu="B200", + region="us", + # Warm floor of 1 (the pool must be UP before the trainer sends rollouts); cap at 3 + # to bound the footprint at trainer 8 + pool 3 = 11 concurrent B200. + rollout_min_containers=1, + rollout_max_containers=3, + # Per-container autoscaler target, well below the trainer's client concurrency: a + # rollout wave (rollout_batch_size x n_samples_per_prompt = 256) must register as + # queue pressure so Flash scales OUT to the container cap instead of one engine + # absorbing the whole wave at its concurrency ceiling. + rollout_target_inputs=24, + proxy_regions=["us-west"], +) + + +class _Miles(MilesConfig): + miles_model_script = "scripts/models/qwen3-30B-A3B.sh" + + # Bridge mode: ref_load is the bf16 HF masters directly (no torch_dist prep). + hf_checkpoint = f"{PREP_PATH}/{MODEL_TAG}/nvfp4" + ref_load = f"{PREP_PATH}/{MODEL_TAG}/bf16" + megatron_to_hf_mode = "bridge" + model_name = "qwen3moe" # megatron_to_hf export dispatch + + actor_num_nodes = 1 + actor_num_gpus_per_node = 8 + num_gpus_per_node = 8 + colocate = False # disk-delta is incompatible with --colocate + rollout_num_gpus = 0 + rollout_num_gpus_per_engine = 1 # 30B NVFP4 is ~17 GB packed; 1 B200 per engine + rollout_endpoint_url = None + use_miles_router = True + + custom_rollout_request_hook_path = "cookbook.common.hooks.gated_rollout_request_hook" + custom_config_path = { + "rollout_request_weight_version_mode": "min", + "rollout_request_weight_version_lag": 1, + "rollout_request_retry_attempts": 240, + "rollout_request_retry_sleep": 1.0, + "rollout_session_affinity_header": "Modal-Session-ID", + } + + async_mode = True + update_weights_interval = 1 + + # ── NVFP4 QAT: the recipe's trainer precision config ── + fp4_format = "e2m1" + fp4_recipe = "nvfp4" + fp4_param_gather = False + te_precision_config_file = { + "configs": { + "nvfp4": { + "transformer_engine_config_type": "TEQuantizationParams", + "training_recipe": {"fp4_quantization_recipe": "nvfp4"}, + }, + "bf16": {"transformer_engine_config_type": "TEQuantizationParams", "training_recipe": {}}, + }, + "matchers": { + "routed_experts_fc1_nvfp4": { + "type": "glob", + "enabled": True, + "pattern": "*.mlp.experts.linear_fc1", + "config": "nvfp4", + }, + "routed_experts_fc2_nvfp4": { + "type": "glob", + "enabled": True, + "pattern": "*.mlp.experts.linear_fc2", + "config": "nvfp4", + }, + "default_bf16": {"type": "glob", "enabled": True, "pattern": "*", "config": "bf16"}, + }, + } + # Selective precision: last 15% of the 48 layers (7) stay BF16 (the blog's "spend + # precision where it matters"; Qwen3 MoE has no shared expert). All layers are MoE + # (FIRST_K_DENSE_REPLACE=0), so no start carve-out. + num_layers_at_start_in_bf16 = 0 + num_layers_at_end_in_bf16 = 7 + + update_weight_transfer_mode = "disk-delta" + update_weight_delta_encoding = "xor" + update_weight_delta_checksum = "xxh3-128" + update_weight_disk_dir = DELTA_BULLETIN_ROOT # app.py run-scopes this + custom_update_weight_post_write_path = "cookbook.common.hooks.commit_and_wake" + + # Data: DAPO-Math-17k — the blog's dataset. + prompt_data = f"{DATA_PATH}/dapo-math-17k/dapo-math-17k.jsonl" + input_key = "prompt" + label_key = "label" + apply_chat_template = True + rollout_shuffle = True + balance_data = True + rm_type = "deepscaler" + eval_interval = None + + # Rollout (bring-up smoke length; scale num_rollout for a real run). + num_rollout = 20 + # None (not just > num_rollout): miles still writes a final megatron save at the last + # rollout, and the ~120 GB torch_dist save blows the trainer's default ephemeral disk + # (the volume write cache lives there). None makes the app null out --load/--save so + # no save path exists. Set a real save_interval (plus trainer_ephemeral_disk_mib) for + # long runs. + save_interval = None + rollout_batch_size = 32 + # Long enough that math reasoning traces mostly terminate rather than truncate + # (4096 clipped most responses, starving the reward signal). + rollout_max_response_len = 12288 + rollout_temperature = 0.8 + n_samples_per_prompt = 8 + global_batch_size = 128 + use_dynamic_global_batch_size = True + # Trainer-side client concurrency to the pool gateway: high enough to drive the + # scaled-out pool (3 x 24 targets), not just one engine. + sglang_server_concurrency = 128 + + use_rollout_routing_replay = True + + # TE 2.17's cuDNN fused-attention graph is incompatible with the base image's loaded + # libcudnn (9.16: CUDNN_STATUS_BAD_PARAM in the bwd reshape). Attention is BF16 in + # this recipe (only MoE experts are NVFP4), so the backend sits outside the + # quantization contract. Drop when the base ships cuDNN >= 9.19. + attention_backend = "flash" + # Trainer parallelism: the upstream Qwen3-30B recipe (TP4/SP/PP1/CP1) with EP = the + # full node under no_colocate, as run_qwen3_30b_a3b.py does. + tensor_model_parallel_size = 4 + sequence_parallel = True + pipeline_model_parallel_size = 1 + context_parallel_size = 1 + expert_model_parallel_size = 8 + expert_tensor_parallel_size = 1 + use_dynamic_batch_size = True + max_tokens_per_gpu = 8192 + recompute_granularity = "full" + recompute_method = "uniform" + recompute_num_layers = 1 + attention_dropout = 0.0 + hidden_dropout = 0.0 + accumulate_allreduce_grads_in_fp32 = True + attention_softmax_in_fp32 = True + no_check_for_nan_in_loss_and_grad = True + + optimizer = "adam" + lr = 1e-6 + lr_decay_style = "constant" + weight_decay = 0.1 + adam_beta1 = 0.9 + adam_beta2 = 0.98 + optimizer_cpu_offload = True + overlap_cpu_optimizer_d2h_h2d = True + use_precision_aware_optimizer = True + + # Algorithm (GRPO + truncated importance sampling, per the blog's baseline). + advantage_estimator = "grpo" + eps_clip = 0.2 + eps_clip_high = 0.28 + use_kl_loss = True + kl_loss_coef = 0.0 + kl_loss_type = "low_var_kl" + entropy_coef = 0.0 + use_tis = True + + environment = { + "CUDA_DEVICE_MAX_CONNECTIONS": "1", + "NCCL_NVLS_ENABLE": "1", + "NVSHMEM_DISABLE_NCCL": "1", + "NCCL_TIMEOUT_MS": "360000000", + **_NVTE_QUANT_ENV, + } + + def prepare_data(self) -> None: + from datasets import load_dataset + + ds = load_dataset("BytedTsinghua-SIA/DAPO-Math-17k", split="train") + ds = ds.shuffle(seed=42).select(range(min(50000, ds.num_rows))) + ds = ds.map(lambda ex: {"label": ex["reward_model"]["ground_truth"]}) + ds = ds.select_columns(["prompt", "label"]) + ds.to_json(f"{DATA_PATH}/dapo-math-17k/dapo-math-17k.jsonl") + + +miles = _Miles() diff --git a/cookbook/miles_disagg/helpers.py b/cookbook/miles_disagg/helpers.py deleted file mode 100644 index 9d50018..0000000 --- a/cookbook/miles_disagg/helpers.py +++ /dev/null @@ -1,223 +0,0 @@ -"""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. -""" - -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 - -# Re-export shared helpers so existing callers (modal_train.py) don't break. -from cookbook.ray_cluster import ( # noqa: F401 - RAY_START_TIMEOUT, - RAY_WORKER_JOIN_TIMEOUT, - get_modal_cluster_context, - start_ray_head, - start_ray_worker, - training_nodes, -) -from cookbook.sidecar_process import ( # noqa: F401 - terminate_process, - wait_http, -) - - -SIDECAR_MODULE = "cookbook.miles_disagg.sidecar" - - -def start_sglang_sidecar( - *, - sidecar_port: int, - sglang_port: int, - bulletin_root: str, - local_checkpoint_dir: str, - base_checkpoint_dir: str, - volume_name: str, - commit_mode: str, - debug_requests: bool = False, -) -> subprocess.Popen: - from cookbook.sidecar_process import start_sglang_sidecar as _start - - return _start( - sidecar_module=SIDECAR_MODULE, - sidecar_port=sidecar_port, - sglang_port=sglang_port, - bulletin_root=bulletin_root, - local_checkpoint_dir=local_checkpoint_dir, - base_checkpoint_dir=base_checkpoint_dir, - volume_name=volume_name, - commit_mode=commit_mode, - debug_requests=debug_requests, - ) - - -# ── miles launch ────────────────────────────────────────────────────────────── - - -def prepare_miles_config(miles_cfg: Any, tmpdir: str) -> None: - """Resolve HF repo IDs to local paths and materialize inline YAML configs. - - hf_checkpoint / ref_load already point at prepared absolute paths (the served - 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) - - -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.""" - - -def smoke_flash_pool( - *, - app_name: str, - cls_name: str, - model_name: str, - weight_version: int, - 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"}, - ) - with urllib.request.urlopen(request, timeout=timeout) as resp: - return json.load(resp) diff --git a/cookbook/miles_disagg/hooks.py b/cookbook/miles_disagg/hooks.py deleted file mode 100644 index 67b8f3e..0000000 --- a/cookbook/miles_disagg/hooks.py +++ /dev/null @@ -1,29 +0,0 @@ -"""Modal publish + rollout-gating hooks for the miles_disagg example. - -Thin wrappers around :mod:`cookbook.bulletin_hooks` with the miles-specific -env-var fallbacks for the Flash app / server class name. -""" - -from __future__ import annotations - -from typing import Any - -from cookbook.bulletin_hooks import ( - commit_and_wake as _commit_and_wake, - gated_rollout_request_hook, -) - - -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", - ) - - -# Re-export for the trainer's custom_rollout_request_hook_path. -__all__ = ["commit_and_wake", "gated_rollout_request_hook"] diff --git a/cookbook/miles_disagg/megatron_r3_num_out_tokens.patch b/cookbook/miles_disagg/megatron_r3_num_out_tokens.patch deleted file mode 100644 index 407134f..0000000 --- a/cookbook/miles_disagg/megatron_r3_num_out_tokens.patch +++ /dev/null @@ -1,43 +0,0 @@ -megatron/core/transformer/moe/token_dispatcher.py --- Python -521 self.num_out_tokens = n 521 self.num_out_tokens = n -... um_local_tokens_per_expert.sum() ... um_local_tokens_per_expert.sum() -522 self._maybe_update_cuda 522 self._maybe_update_cuda -... _sync_point("before_permutation_1") ... _sync_point("before_permutation_1") -523 else: 523 else: -524 # For dropless training 524 # For dropless training -... , output size is static (num_tokens ... , output size is normally num_token -... * topk) ... s * topk. -525 # No explicit sync need 525 # Under routing replay -... ed ... (R3), the replayed routing_map can -... ... have fewer -... 526 # than topk experts per -... ... token (duplicate / -1 entries coll -... ... apse in the -... 527 # boolean map), so num_ -... ... tokens*topk overcounts and disagree -... ... s with -... 528 # input_splits (= routi -... ... ng_map True-count), causing "Split -... ... sizes doesn't -... 529 # match total dim 0 siz -... ... e" in the EP all-to-all. Derive it -... ... from the -... 530 # actual routing_map to -... ... stay consistent (equals num_tokens -... ... *topk in the -... 531 # normal dense case; ne -... ... eds the GPU->CPU sync like the drop -... ... ping path). -526 self.num_out_tokens = r 532 self.num_out_tokens = n -... outing_map.size(0) * self.config.mo ... um_local_tokens_per_expert.sum() -... e_router_topk ... -... 533 self._maybe_update_cuda -... ... _sync_point("before_permutation_1") -527 if self.ep_size > 1 or self 534 if self.ep_size > 1 or self -... .tp_size > 1: ... .tp_size > 1: -528 # ===================== 535 # ===================== -... ============================== ... ============================== -529 # Calculate input_split 536 # Calculate input_split -... s, output_splits for alltoall/allga ... s, output_splits for alltoall/allga -... ther in variable size. ... ther in variable size. - 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 deleted file mode 100644 index b3a9686..0000000 --- a/cookbook/miles_disagg/modal_train.py +++ /dev/null @@ -1,759 +0,0 @@ -"""Disaggregated miles training on Modal (NVFP4, Blackwell end to end). - -A Modal Flash pool of SGLang servers handles rollouts; a clustered Trainer runs -miles on Ray and publishes XOR weight deltas through a Modal Volume bulletin -board that the rollout servers sync from. The miles twin of -cookbook/slime_disagg/modal_train.py — same two-half architecture, same stitch -sidecar/bulletin machinery; the trainer is miles instead of slime and the -precision is NVFP4 (so BOTH halves are Blackwell — see the config docstrings). - -Run all commands as modules from the repo root, e.g.: - - uv run --extra modal modal deploy -m cookbook.miles_disagg.modal_train -""" - -from __future__ import annotations - -import importlib -import os -import shlex -import subprocess -import tempfile -import uuid -from pathlib import Path - -import modal -import modal.experimental - -from cookbook.miles_disagg import helpers -from cookbook.miles_disagg.configs.base import ( - CHECKPOINTS_PATH, - DATA_PATH, - HF_CACHE_PATH, - PREP_PATH, - MilesConfig, -) -from stitch.providers.modal import resolve_flash_gateway_url - -# The one deploy-time knob: which experiment config this app is built around. -EXPERIMENT = os.environ.get("EXPERIMENT_CONFIG", "moonlight_nvfp4_disagg") -exp = importlib.import_module(f"cookbook.miles_disagg.configs.{EXPERIMENT}") - -modal_cfg = exp.modal -miles_cfg = exp.miles - -# The rollout Server pool warm-boots `min_containers` replicas the moment this app -# is materialized — by a deploy OR by `modal run` of ANY function in it. Those -# replicas serve the prepared NVFP4 base (miles_cfg.hf_checkpoint), so before that -# base exists they crash-loop on a missing --model-path. prepare_checkpoints builds -# the base, so it must run with the pool down: POOL_MIN_CONTAINERS=0. -POOL_MIN_CONTAINERS = int(os.environ.get("POOL_MIN_CONTAINERS", modal_cfg.rollout_min_containers)) - -APP_NAME = exp.APP_NAME -MODEL_NAME = miles_cfg.hf_checkpoint # the served NVFP4 base (a prepared local path) -ROLLOUT_CONCURRENCY = modal_cfg.rollout_target_inputs or miles_cfg.sglang_server_concurrency -N_TRAIN_NODES = helpers.training_nodes(miles_cfg) - -MINUTES = 60 -SIDECAR_PORT = 8000 -SGLANG_PORT = 8001 -RAY_PORT = 6379 -SERVER_STARTUP_TIMEOUT = 35 * MINUTES -LOCAL_CHECKPOINT_PATH = "/local-checkpoint" - -# radixark/miles bakes Megatron-LM (miles-main: native --fp4-format NVFP4 -# BlockScaling) + TransformerEngine. NVFP4 QAT requires TE >= 2.7.0.dev0 and -# Blackwell; verify the image's TE on a warm B200 during bring-up. -MILES_IMAGE_TAG = "radixark/miles:latest" -MILES_ROOT = "/root/miles" -# megatron.core is pip-installed in the base image, but megatron.training lives -# only in this source tree; miles' own launcher puts it on PYTHONPATH, and so -# must we (we run train_async.py directly rather than via execute_train). -MEGATRON_PATH = "/root/Megatron-LM" -# Fork branch with the disaggregated-rollout features (opaque HTTP endpoint, -# publish-only disk-delta, request hook) AND the NVFP4 / publish-only fixes -# (NVFP4 export dispatch, publish-only semaphore + http client, 0-dim delta -# encode, encoding_dsv4 import guard). See branch nvfp4-disagg-fixes. -MILES_REPO_URL = "https://github.com/modal-projects/miles.git" -# nvfp4-disagg-v2 @ 9e98062af = merge(feat/disaggregated-rollout @121035147, -# radixark/miles#1261 NVFP4 RL @f95c2c495) + our publish-only/0-dim/encoding_dsv4 -# fixes + KimiK25 mbridge import converter (miles_plugins/mbridge/kimi.py). PUSH -# this branch to modal-projects/miles before deploying. (Built in /tmp/miles-merge.) -MILES_REPO_REF = "e9ad52dbbe09b6113b4fa4dccfb5ace35341540e" - -# Build-time bake of the megatron R3 dispatch fix (see the .run_commands call -# below). Kept as a string so the build step has no host-file dependency; the -# replace targets a single, stable line and reports how many sites it hit. -_R3_DISPATCH_TARGET = "/root/Megatron-LM/megatron/core/transformer/moe/token_dispatcher.py" -_R3_DISPATCH_OLD = "self.num_out_tokens = routing_map.size(0) * self.config.moe_router_topk" -_R3_DISPATCH_NEW = ( - "self.num_out_tokens = num_local_tokens_per_expert.sum()\n" - ' self._maybe_update_cuda_sync_point("before_permutation_1")' -) -_R3_DISPATCH_BAKE_PY = ( - "import pathlib;" - f"p=pathlib.Path({_R3_DISPATCH_TARGET!r});" - "s=p.read_text();" - f"n=s.count({_R3_DISPATCH_OLD!r});" - f"p.write_text(s.replace({_R3_DISPATCH_OLD!r}, {_R3_DISPATCH_NEW!r}));" - "print(f'[R3 bake] patched {n} dispatch site(s) in token_dispatcher.py')" -) - -# Keep GLM4Moe expert-bias exports fp32 for BF16 disk-delta byte-shape parity. -_GLM4MOE_EXPERT_BIAS_TARGET = f"{MILES_ROOT}/miles/backends/megatron_utils/megatron_to_hf/glm4moe.py" -_GLM4MOE_EXPERT_BIAS_OLD = ( - ' elif rest == "mlp.router.expert_bias":\n' - ' return [(f"model.layers.{layer_idx}.mlp.gate.e_score_correction_bias", to_model_dtype(args, param))]' -) -_GLM4MOE_EXPERT_BIAS_NEW = ( - ' elif rest == "mlp.router.expert_bias":\n' - " return [(f\"model.layers.{layer_idx}.mlp.gate.e_score_correction_bias\", param)]" -) -_GLM4MOE_EXPERT_BIAS_BAKE_PY = ( - "import pathlib;" - f"p=pathlib.Path({_GLM4MOE_EXPERT_BIAS_TARGET!r});" - "s=p.read_text();" - f"n=s.count({_GLM4MOE_EXPERT_BIAS_OLD!r});" - "assert n == 1, f'expected exactly one GLM4Moe expert-bias export site, found {n}';" - f"p.write_text(s.replace({_GLM4MOE_EXPERT_BIAS_OLD!r}, {_GLM4MOE_EXPERT_BIAS_NEW!r}));" - "print(f'[GLM4Moe bake] patched {n} expert-bias dtype export site(s) in glm4moe.py')" -) - -TORCH_DIST_CONVERT_WRAPPER = "/root/convert_hf_to_torch_dist_modal.py" - -image = ( - modal.Image.from_registry(MILES_IMAGE_TAG) - .entrypoint([]) - # RDMA/EFA userspace stack. The miles Dockerfile installs only nvtop/rsync/etc; - # without the full ibverbs stack NCCL silently falls back to TCP sockets on the - # multi-node trainer (catastrophic throughput) and emits libibverbs driver - # warnings. These let the AWS-OFI/libfabric path NCCL uses under rdma=True bind EFA. - .apt_install( - "libibverbs-dev", - "libibverbs1", - "libhwloc-dev", - "libnl-route-3-200", - ) - # 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}") - # Replace the bundled miles with the fork branch (keeps the baked Megatron-LM - # + TE, which provide the NVFP4 training path). - .run_commands( - f"rm -rf {MILES_ROOT}" - 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}" - ) - # R3 routing-replay fix for the baked radixark/Megatron-LM. The dropless MoE - # token dispatcher computes num_out_tokens = routing_map.size(0) * topk, which - # assumes topk distinct experts/token — false under routing replay (duplicate - # / -1 experts collapse in the boolean map), causing the EP all-to-all to - # crash with "Split sizes doesn't match total dim 0 size". Derive it from the - # actual per-expert counts instead (identical to size*topk in the dense - # non-replay case). Idempotent: a no-op once the fork itself ships the fix. - # Source diff: cookbook/miles_disagg/megatron_r3_num_out_tokens.patch. - .run_commands(f"python3 -c {shlex.quote(_R3_DISPATCH_BAKE_PY)}") - .pip_install( - "fastapi", # stitch sidecar (rollout pool reuses this image only if no serving image) - "httpx", - "uvicorn", - # miles applies disk deltas host-side via miles.utils.disk_delta; ensure - # its codec deps are present even when miles is reinstalled --no-deps. - "zstandard", - "xxhash", - "blake3", - ) - .env( - { - "EXPERIMENT_CONFIG": EXPERIMENT, - "HF_XET_HIGH_PERFORMANCE": "1", - "HF_HUB_ENABLE_HF_TRANSFER": "1", - } - ) -) - -if getattr(exp, "USE_MODAL_TORCH_DIST_WRAPPER", False): - image = image.add_local_file( - "cookbook/miles_disagg/convert_hf_to_torch_dist_modal.py", - TORCH_DIST_CONVERT_WRAPPER, - copy=True, - ) - -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 -# after a non-copy local mount). -image = image.add_local_python_source("stitch").add_local_dir( - Path(__file__).parent, - remote_path="/root/cookbook/miles_disagg", - ignore=["**/__pycache__"], -) - -# Dev iteration: MILES_LOCAL_DIR overlays a local miles checkout onto the image's -# cloned fork, so fork edits (e.g. the NVFP4 export-dispatch fix) take effect on -# container start with no image rebuild or push. Unset by default. -if miles_local := os.environ.get("MILES_LOCAL_DIR"): - image = image.add_local_dir( - miles_local, - remote_path=MILES_ROOT, - ignore=[".git", "**/__pycache__", "**/*.pyc"], - ) - - -def _select_server_image() -> modal.Image: - """The rollout pool needs a Blackwell SGLang build that serves NVFP4; an - experiment opts in via build_serving_image(...). Either way the pool pins the - trainer's exact miles ref so the sidecar's disk_delta matches the encoder.""" - builder = getattr(exp, "build_serving_image", None) - if builder is None: - return image - return builder( - miles_repo_url=MILES_REPO_URL, - miles_repo_ref=MILES_REPO_REF, - miles_root=MILES_ROOT, - hf_cache_path=str(HF_CACHE_PATH), - experiment=EXPERIMENT, - ) - - -server_image = _select_server_image() -if miles_local and server_image is not image: - server_image = server_image.add_local_dir( - miles_local, - remote_path=MILES_ROOT, - ignore=[".git", "**/__pycache__", "**/*.pyc"], - ) - -with server_image.imports(): - from autoinference_utils.endpoint import SGLangEndpoint, warmup_chat_completions - - -hf_cache_volume = modal.Volume.from_name("huggingface-cache", create_if_missing=True) -data_volume = modal.Volume.from_name("miles-data", create_if_missing=True) -checkpoints_volume = modal.Volume.from_name("miles-checkpoints", create_if_missing=True) -prep_volume = modal.Volume.from_name("miles-prep-checkpoints", create_if_missing=True) -# Persists SGLang's kernel caches across cold starts. SGLang nests flashinfer's -# FP4 MoE autotuner + JIT cache (and DeepGemm etc.) under its own cache dir, so we -# mount the whole tree. Without it every elastic scale-up re-runs `[AutoTuner] -# Tuning trtllm_fp4_block_scale_moe` and recompiles kernels (minutes per Server). -sglang_cache_volume = modal.Volume.from_name("miles-sglang-cache", create_if_missing=True) -delta_volume = modal.Volume.from_name(exp.DELTA_VOLUME_NAME, create_if_missing=True, version=2) - -# SGLang's default cache dir ($HOME/.cache/sglang), which holds the nested -# flashinfer/DeepGemm caches; mounted so they survive container teardown. -SGLANG_CACHE_PATH = "/root/.cache/sglang" - -train_volumes = { - str(HF_CACHE_PATH): hf_cache_volume, - str(DATA_PATH): data_volume, - str(CHECKPOINTS_PATH): checkpoints_volume, - str(PREP_PATH): prep_volume, - exp.DELTA_BULLETIN_ROOT: delta_volume, -} - -app = modal.App(APP_NAME) - -SGLANG_SERVER_ARGS = { - "--served-model-name": MODEL_NAME, - "--dtype": "bfloat16", - "--cuda-graph-max-bs": str(ROLLOUT_CONCURRENCY), - "--max-running-requests": str(ROLLOUT_CONCURRENCY), - "--trust-remote-code": "", - # NVFP4 is driven by the served checkpoint's own quant config — no - # --quantization flag. MLA/cache/routing extras come from the experiment. - **exp.SGLANG_SERVER_ARGS, -} - -WARMUP_PAYLOAD = { - "model": MODEL_NAME, - "messages": [{"role": "user", "content": "Reply with exactly OK."}], - "max_tokens": 8, - "temperature": 0, - "chat_template_kwargs": {"enable_thinking": False}, -} - - -@app.cls( - image=server_image, - gpu=f"{modal_cfg.gpu}:{miles_cfg.rollout_num_gpus_per_engine}", - cloud=modal_cfg.cloud, - region=modal_cfg.region, - volumes={ - str(HF_CACHE_PATH): hf_cache_volume, - str(PREP_PATH): prep_volume, - SGLANG_CACHE_PATH: sglang_cache_volume, - exp.DELTA_BULLETIN_ROOT: delta_volume, - }, - min_containers=POOL_MIN_CONTAINERS, - timeout=40 * MINUTES, - scaledown_window=15 * MINUTES, - # The sidecar copies the full served base (~591 GB for K2.6) to - # /local-checkpoint on ephemeral disk; Modal's default is too small. - ephemeral_disk=modal_cfg.rollout_ephemeral_disk_mib, - include_source=False, -) -@modal.experimental.http_server( - port=SIDECAR_PORT, - proxy_regions=modal_cfg.proxy_regions, - exit_grace_period=25, - startup_timeout=SERVER_STARTUP_TIMEOUT, -) -@modal.concurrent(target_inputs=ROLLOUT_CONCURRENCY) -class Server: - """One SGLang rollout server plus the stitch weight-sync sidecar. - - The sidecar proxies rollout traffic, reloads the delta Volume, and applies - published weight versions so requests pinned to a version are served by - matching weights. - """ - - @modal.enter() - def startup(self) -> None: - self.endpoint = SGLangEndpoint( - model_path=MODEL_NAME, - worker_port=SGLANG_PORT, - tp=miles_cfg.rollout_num_gpus_per_engine, - extra_server_args=SGLANG_SERVER_ARGS, - health_timeout=SERVER_STARTUP_TIMEOUT, - health_poll_interval=10.0, - ) - self.endpoint.start() - warmup_chat_completions( - port=SGLANG_PORT, - payload=WARMUP_PAYLOAD, - successful_requests=2, - request_timeout=120.0, - max_attempts_per_request=3, - ) - # The served NVFP4 base is a prepared directory on the prep Volume - # (MODEL_NAME is its absolute path); deltas are applied host-side onto a - # copy of it. - self.sidecar = helpers.start_sglang_sidecar( - sidecar_port=SIDECAR_PORT, - sglang_port=SGLANG_PORT, - bulletin_root=exp.DELTA_BULLETIN_ROOT, - local_checkpoint_dir=LOCAL_CHECKPOINT_PATH, - base_checkpoint_dir=MODEL_NAME, - volume_name=exp.DELTA_VOLUME_NAME, - commit_mode=exp.SIDECAR_COMMIT_MODE, - debug_requests=getattr(exp, "SIDECAR_DEBUG_REQUESTS", False), - ) - helpers.wait_http( - f"http://127.0.0.1:{SIDECAR_PORT}/health", - self.sidecar, - SERVER_STARTUP_TIMEOUT, - ) - print(f"Rollout server ready: model={MODEL_NAME}, target_inputs={ROLLOUT_CONCURRENCY}") - - @modal.exit() - def stop(self) -> None: - helpers.terminate_process(getattr(self, "sidecar", None)) - if hasattr(self, "endpoint"): - self.endpoint.stop() - - -# Multi-node clustering (RDMA + EFA) is applied only when N_TRAIN_NODES > 1: -# Modal requires clustered B200 functions to use all 8 GPUs per node, so a -# single-node trainer (e.g. the Moonlight de-risk on B200:4) runs as a plain cls. -_TRAINER_KWARGS = dict( - image=image, - gpu=f"{modal_cfg.gpu}:{miles_cfg.actor_num_gpus_per_node}", - memory=modal_cfg.memory, - cloud=modal_cfg.cloud, - region=modal_cfg.region, - volumes=train_volumes, - # Ray (address="auto") spills objects + writes logs under /tmp/ray on the node's - # ephemeral disk; Modal's default is far too small for a multi-hour 128-GPU run and - # progressively ENOSPC'd. Give the B200:8 nodes' local NVMe room. - ephemeral_disk=modal_cfg.trainer_ephemeral_disk_mib, - timeout=24 * 60 * MINUTES, - startup_timeout=20 * MINUTES, - scaledown_window=30 * MINUTES, - include_source=False, -) -if N_TRAIN_NODES > 1: - _TRAINER_KWARGS["experimental_options"] = {"efa_enabled": True} - - -class Trainer: - """miles actor cluster. The Ray cluster comes up once per container in - enter(), so back-to-back training runs reuse it instead of rebuilding it.""" - - @modal.enter() - def start_ray(self) -> None: - rank, master_addr, my_ip = helpers.get_modal_cluster_context(N_TRAIN_NODES) - self.rank = rank - os.environ.update( - { - "MILES_HOST_IP": my_ip, - "SGLANG_HOST_IP": my_ip, - "HOST_IP": my_ip, - "MASTER_ADDR": master_addr, - "RAY_ADDRESS": f"{master_addr}:{RAY_PORT}", - "no_proxy": f"127.0.0.1,{master_addr},{my_ip}", - "NO_PROXY": f"127.0.0.1,{master_addr},{my_ip}", - # Expose megatron.training (source-only) before `ray start`, so both - # the launch subprocess and the Ray actors inherit it. Prepend to - # preserve Modal's existing PYTHONPATH (/pkg, /root for stitch). - "PYTHONPATH": f"{MEGATRON_PATH}:{os.environ.get('PYTHONPATH', '')}", - **miles_cfg.environment, - } - ) - if rank == 0: - helpers.start_ray_head(my_ip, N_TRAIN_NODES, ray_port=RAY_PORT) - else: - helpers.start_ray_worker(my_ip, master_addr, ray_port=RAY_PORT) - - @modal.method() - def train(self, experiment: str, payload: dict) -> None: - """Run one training job from a MilesConfig payload (see to_payload()).""" - for volume in train_volumes.values(): - volume.reload() - - cfg = MilesConfig.from_payload(payload) - # te_precision_config_file is re-read on every Ray actor during model build, - # so it must exist at an identical local path on all 16 nodes. Materialize it - # here (SPMD: train() runs on every container) before the rank-0 gate; node 0 - # then embeds this deterministic path in the args. (prepare_miles_config below - # skips it once it's a path string, not a dict.) - helpers.materialize_node_local_yaml(cfg, "te_precision_config_file") - if self.rank != 0: - return - - if helpers.training_nodes(cfg) != N_TRAIN_NODES: - raise ValueError( - f"experiment {experiment!r} needs {helpers.training_nodes(cfg)} node(s) but this app " - f"was deployed with {N_TRAIN_NODES}; deploy it as its own app with EXPERIMENT_CONFIG={experiment}" - ) - if cfg.environment != miles_cfg.environment: - print( - f"WARNING: experiment {experiment!r} changes `environment`, which only " - f"takes effect after a redeploy restarts the Ray cluster." - ) - - cfg.rollout_endpoint_url = resolve_flash_gateway_url(APP_NAME, Server.__name__) - # Fresh run id per launch: miles writes this run's chain under a partition - # (//weight_v{N}/), while the canonical pointer at - # /latest is self-identifying (/weight_vN), so a new - # run is a forward pointer move, never a colliding rewind. - run_id = uuid.uuid4().hex[:12] - cfg.update_weight_disk_dir = f"{exp.DELTA_BULLETIN_ROOT}/{run_id}" - cfg.load = f"{CHECKPOINTS_PATH}/{run_id}/checkpoints" - cfg.save = f"{CHECKPOINTS_PATH}/{run_id}/checkpoints" - # Merge the dynamic bulletin identity into custom_config_path (which already - # carries the request-gating knobs). miles setattr's every key onto args, - # so the publish + request hooks read them via getattr(args, ...). - existing = dict(getattr(cfg, "custom_config_path", {}) or {}) - cfg.custom_config_path = { - **existing, - "update_weight_delta_volume_name": exp.DELTA_VOLUME_NAME, - "rollout_modal_flash_app_name": APP_NAME, - "rollout_modal_flash_server_cls_name": Server.__name__, - "run_id": run_id, - } - helpers.prepare_miles_config(cfg, tempfile.mkdtemp()) - cmd = helpers.build_train_cmd(cfg, MILES_ROOT) - - 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 - # observable after the fact — the app-logs buffer scrolls past tracebacks. - # set -o pipefail preserves the trainer's exit code through the tee. - log_path = f"{CHECKPOINTS_PATH}/{run_id}/train.log" - os.makedirs(os.path.dirname(log_path), exist_ok=True) - teed = f"set -o pipefail; ({cmd}) 2>&1 | tee {log_path}" - try: - subprocess.run(["bash", "-lc", teed], check=True) - finally: - try: - checkpoints_volume.commit() - print(f"Train log committed to volume miles-checkpoints at {run_id}/train.log") - except Exception as exc: # noqa: BLE001 - print(f"WARNING: could not commit train log: {exc}") - - -# Apply clustering only for genuine multi-node trainers (see _TRAINER_KWARGS). -if N_TRAIN_NODES > 1: - Trainer = modal.experimental.clustered(N_TRAIN_NODES, rdma=True)(Trainer) -Trainer = app.cls(**_TRAINER_KWARGS)(Trainer) - - -@app.function( - image=image, - gpu=f"{modal_cfg.gpu}:1", - volumes={str(HF_CACHE_PATH): hf_cache_volume, str(PREP_PATH): prep_volume}, - timeout=6 * 60 * MINUTES, - secrets=[modal.Secret.from_name("huggingface-secret")], - include_source=False, -) -def prepare_checkpoints() -> None: - """Build the bf16 masters + served NVFP4 base on a GPU (NVFP4 quant needs CUDA). - - Lifecycle (see the config docstrings): - 1. masters (bf16): if the source ships quantized (Kimi INT4), dequantize - with tools/convert_kimi_int4_to_bf16.py; if bf16 (Moonlight), the - 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" - os.environ.pop("HF_XET_HIGH_PERFORMANCE", None) - if getattr(exp, "DISABLE_HF_TRANSFER", False): - os.environ.pop("HF_HUB_ENABLE_HF_TRANSFER", None) - - from huggingface_hub import snapshot_download - - prep_volume.reload() - tag = exp.MODEL_TAG - bf16_dir = f"{PREP_PATH}/{tag}/bf16" - nvfp4_dir = f"{PREP_PATH}/{tag}/nvfp4" - served_checkpoint_format = getattr(exp, "SERVED_CHECKPOINT_FORMAT", "nvfp4") - if served_checkpoint_format not in {"bf16", "nvfp4"}: - raise SystemExit(f"Unsupported SERVED_CHECKPOINT_FORMAT={served_checkpoint_format!r}") - tools = f"{MILES_ROOT}/tools" - - def _staged(final_dir: str, build) -> None: - """Build into a sibling .partial dir and atomically rename on success, so - an interrupted multi-hour step never leaves a half-built dir that the - `already done?` check below would mistake for complete (this run's first - attempt died mid-dequant with 7/64 files written).""" - if os.path.isdir(final_dir) and os.listdir(final_dir): - print(f"reusing existing {final_dir}") - return - partial = f"{final_dir}.partial" - subprocess.run(["rm", "-rf", partial], check=True) - os.makedirs(partial, exist_ok=True) - build(partial) - os.rename(partial, final_dir) - - src = snapshot_download(exp.SOURCE_MODEL) - is_int4 = "int4" in exp.SOURCE_MODEL.lower() or _is_int4(src) - - # 1. masters (bf16) - def _build_bf16(out: str) -> None: - if is_int4: - subprocess.run( - ["python", f"{tools}/convert_kimi_int4_to_bf16.py", "--model-dir", src, "--output-dir", out], - check=True, - ) - else: - # bf16 source IS the masters; dereference (-L) so the prep dir holds - # real files, not symlinks into the (separate) HF cache volume. - subprocess.run(f"cp -aL {src}/. {out}/", shell=True, check=True) - # convert_kimi_int4_to_bf16 copies the source config verbatim, so the - # dequantized masters still carry the source's INT4 quantization_config - # (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 - # the wrong quant loader for the language model). - _strip_stale_quant_config(os.path.join(out, "config.json")) - - _staged(bf16_dir, _build_bf16) - - if served_checkpoint_format == "bf16": - prep_volume.commit() - print(f"Prepared masters={bf16_dir} served_base={bf16_dir}") - return - - # 2. served NVFP4 base (TE-direct quantizer per radixark/miles#1261). bf16 - # carve-outs for the dense first / last layers must match the trainer's - # --num-layers-at-start/end-in-bf16 so the served base == the export layout. - _nvfp4_carveouts = [] - if (n := getattr(miles_cfg, "num_layers_at_start_in_bf16", None)) is not None: - _nvfp4_carveouts += ["--num-layers-at-start-in-bf16", str(n)] - if (n := getattr(miles_cfg, "num_layers_at_end_in_bf16", None)) is not None: - _nvfp4_carveouts += ["--num-layers-at-end-in-bf16", str(n)] - _staged( - nvfp4_dir, - lambda out: subprocess.run( - ["python", f"{tools}/convert_hf_to_nvfp4.py", "--model-dir", bf16_dir, "--save-dir", out, *_nvfp4_carveouts], - 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}") - - -def _strip_stale_quant_config(config_path: str) -> None: - """Remove any `quantization_config` (top-level AND text_config-nested) from an - HF config.json. Used to clean dequantized bf16 masters whose config still - claims the source quant scheme.""" - import json - - if not os.path.exists(config_path): - return - with open(config_path) as f: - cfg = json.load(f) - removed = bool(cfg.pop("quantization_config", None)) - if isinstance(cfg.get("text_config"), dict): - removed = bool(cfg["text_config"].pop("quantization_config", None)) or removed - if removed: - with open(config_path, "w") as f: - json.dump(cfg, f, indent=2) - print(f"stripped stale quantization_config from {config_path}") - - -def _is_int4(model_dir: str) -> bool: - import json - - cfg_path = os.path.join(model_dir, "config.json") - if not os.path.exists(cfg_path): - return False - with open(cfg_path) as f: - cfg = json.load(f) or {} - # VLMs (Kimi K2.x are KimiK25ForConditionalGeneration) nest the quant config - # under text_config; convert_kimi_int4_to_bf16 reads it the same way. - qc = (cfg.get("text_config") or {}).get("quantization_config") or cfg.get("quantization_config") or {} - return qc.get("quant_method") == "compressed-tensors" - - -# Raw-mode training loads a Megatron torch_dist checkpoint, not HF (miles' HF load -# is bridge-only). Convert the bf16 masters -> torch_dist with convert_hf_to_torch_dist -# (raw build from MODEL_ARGS + the KimiK25 mbridge weight loader). The full 1T K2.6 -# won't fit an 8-way split (~250 GB/GPU), so this runs clustered across -# TORCH_DIST_PREP_NODES x 8 B200 via torchrun (convert auto-derives pp from world size). -TORCH_DIST_PREP_NODES = modal_cfg.torch_dist_prep_nodes -_TORCH_DIST_GPUS_PER_NODE = modal_cfg.torch_dist_prep_gpus_per_node - - -def prepare_torch_dist() -> None: - """Build {tag}/torch_dist (the raw-mode ref_load) from the {tag}/bf16 masters.""" - rank, master_addr, _my_ip = helpers.get_modal_cluster_context(TORCH_DIST_PREP_NODES) - prep_volume.reload() - tag = exp.MODEL_TAG - bf16_dir = f"{PREP_PATH}/{tag}/bf16" - torch_dist_dir = f"{PREP_PATH}/{tag}/torch_dist" - if os.path.exists(os.path.join(torch_dist_dir, "latest_checkpointed_iteration.txt")): - print(f"reusing existing torch_dist {torch_dist_dir}") - return - if not miles_cfg.miles_model_script: - raise SystemExit("prepare_torch_dist requires miles_model_script (MODEL_ARGS)") - use_modal_wrapper = TORCH_DIST_PREP_NODES > 1 and getattr(exp, "USE_MODAL_TORCH_DIST_WRAPPER", False) - convert_script = TORCH_DIST_CONVERT_WRAPPER if use_modal_wrapper else f"{MILES_ROOT}/tools/convert_hf_to_torch_dist.py" - # source the model script for ${MODEL_ARGS[@]}, then torchrun the conversion. - inner = ( - f"source {MILES_ROOT}/{miles_cfg.miles_model_script} && " - f"PYTHONPATH={MEGATRON_PATH} torchrun" - f" --nnodes {TORCH_DIST_PREP_NODES} --node-rank {rank}" - f" --master-addr {master_addr} --master-port 29500" - f" --nproc-per-node {_TORCH_DIST_GPUS_PER_NODE}" - f" {convert_script} ${{MODEL_ARGS[@]}}" - f" --hf-checkpoint {bf16_dir} --save {torch_dist_dir} --megatron-to-hf-mode raw" - f" {modal_cfg.torch_dist_convert_extra_args}" - ) - env = os.environ.copy() - if use_modal_wrapper: - env["SKIP_RELEASE_RENAME"] = "1" - subprocess.run(["bash", "-c", inner], check=True, env=env) - # Each node wrote its own distcp shards to its own Volume mount, so EVERY node must - # commit — a rank-0-only commit drops nodes 1..N-1's shards and the checkpoint loads - # short a shard (FileNotFoundError on ___.distcp). Rank 0 additionally - # committed .metadata/common.pt and the iteration tracker. Disjoint files across nodes - # merge cleanly on the Volume. - prep_volume.commit() - if rank == 0: - print(f"Prepared torch_dist={torch_dist_dir}") - - -_torch_dist_fn_kwargs = dict( - image=image, - gpu=f"{modal_cfg.gpu}:{_TORCH_DIST_GPUS_PER_NODE}", - volumes={str(HF_CACHE_PATH): hf_cache_volume, str(PREP_PATH): prep_volume}, - timeout=6 * 60 * MINUTES, - # Headroom for the Modal Volume's local write cache while each node buffers its full - # distcp shard set (~700 GB for the 1T save) until commit. The served-pool default - # leaves no room and rank 0's commit hits ENOSPC, so allow a dedicated larger value. - ephemeral_disk=(modal_cfg.torch_dist_prep_ephemeral_disk_mib or modal_cfg.rollout_ephemeral_disk_mib), - secrets=[modal.Secret.from_name("huggingface-secret")], - include_source=False, -) -if TORCH_DIST_PREP_NODES > 1: - prepare_torch_dist = modal.experimental.clustered(TORCH_DIST_PREP_NODES, rdma=True)(prepare_torch_dist) - _torch_dist_fn_kwargs["experimental_options"] = {"efa_enabled": True} -prepare_torch_dist = app.function(**_torch_dist_fn_kwargs)(prepare_torch_dist) - - -@app.function( - image=image, - volumes={str(DATA_PATH): data_volume}, - timeout=2 * 60 * MINUTES, - secrets=[modal.Secret.from_name("huggingface-secret")], - include_source=False, -) -def prepare_dataset() -> None: - data_volume.reload() - miles_cfg.prepare_data() - data_volume.commit() - - -@app.local_entrypoint() -def launch_train(experiment: str = EXPERIMENT) -> None: - """Resolve an experiment from the local working tree and spawn it on the - deployed app. Training args ship as data, so new or edited configs run - without a redeploy; infrastructure changes still require one.""" - from modal.exception import NotFoundError - - run = importlib.import_module(f"cookbook.miles_disagg.configs.{experiment}") - if run.DELTA_VOLUME_NAME != exp.DELTA_VOLUME_NAME: - raise SystemExit( - f"Experiment {experiment!r} owns Volume {run.DELTA_VOLUME_NAME!r}, but app {APP_NAME!r} " - f"mounts {exp.DELTA_VOLUME_NAME!r}. Deploy it as its own app with EXPERIMENT_CONFIG={experiment}." - ) - - try: - trainer = modal.Cls.from_name(APP_NAME, Trainer.__name__)() - call = trainer.train.spawn(experiment, run.miles.to_payload()) - except NotFoundError: - raise SystemExit( - f"App {APP_NAME!r} is not deployed. Run:\n" - f" uv run --extra modal modal deploy -m cookbook.miles_disagg.modal_train" - ) - print(f"Spawned train({experiment!r}) on {APP_NAME}: {call.object_id}") - - -@app.local_entrypoint() -def smoke_flash_pool(weight_version: int = 0, timeout_seconds: int = 30 * MINUTES) -> None: - """Check that the deployed Flash pool serves completions at the expected - weight version, via the gateway and each container directly.""" - helpers.smoke_flash_pool( - app_name=APP_NAME, - cls_name=Server.__name__, - model_name=MODEL_NAME, - weight_version=weight_version, - expect_min_containers=modal_cfg.rollout_min_containers, - timeout_seconds=timeout_seconds, - ) diff --git a/cookbook/miles_disagg/patches/megatron-r3-dispatch.patch b/cookbook/miles_disagg/patches/megatron-r3-dispatch.patch new file mode 100644 index 0000000..f10d7a6 --- /dev/null +++ b/cookbook/miles_disagg/patches/megatron-r3-dispatch.patch @@ -0,0 +1,21 @@ +# R3 routing-replay fix for the baked radixark/Megatron-LM. The dropless MoE +# token dispatcher computes num_out_tokens = routing_map.size(0) * topk, which +# assumes topk distinct experts/token -- false under routing replay (duplicate +# / -1 experts collapse in the boolean map), causing the EP all-to-all to crash +# with "Split sizes doesn't match total dim 0 size". Derive it from the actual +# per-expert counts instead (identical to size*topk in the dense non-replay +# case). Idempotent: a no-op once the fork itself ships the fix. +diff --git a/megatron/core/transformer/moe/token_dispatcher.py b/megatron/core/transformer/moe/token_dispatcher.py +index 327dbc8a3..afeca2603 100644 +--- a/megatron/core/transformer/moe/token_dispatcher.py ++++ b/megatron/core/transformer/moe/token_dispatcher.py +@@ -523,7 +523,8 @@ class MoEAlltoAllTokenDispatcher(MoETokenDispatcher): + else: + # For dropless training, output size is static (num_tokens * topk) + # No explicit sync needed +- self.num_out_tokens = routing_map.size(0) * self.config.moe_router_topk ++ self.num_out_tokens = num_local_tokens_per_expert.sum() ++ self._maybe_update_cuda_sync_point("before_permutation_1") + if self.ep_size > 1 or self.tp_size > 1: + # =================================================== + # Calculate input_splits, output_splits for alltoall/allgather in variable size. diff --git a/cookbook/miles_disagg/prep.py b/cookbook/miles_disagg/prep.py new file mode 100644 index 0000000..b5b5d7b --- /dev/null +++ b/cookbook/miles_disagg/prep.py @@ -0,0 +1,152 @@ +"""miles checkpoint preparation, ported to plain functions the app registers as Modal +functions: build the bf16 masters + the served base (bf16 / fp8 / nvfp4) and the raw-mode +torch_dist ref_load. All read their experiment constants off the selected config module. +""" + +from __future__ import annotations + +import json +import os +import subprocess + +from cookbook.common.constants import PREP_PATH +from cookbook.miles_disagg.trainer_image import MEGATRON_PATH, MILES_ROOT, TORCH_DIST_CONVERT_WRAPPER + + +def prepare_checkpoints(exp, prep_volume) -> None: + """Build the bf16 masters (trainer arch source) + the served base on a GPU. + + masters (bf16): a quantized source (Kimi INT4) is dequantized; a bf16 source IS the + masters. served base: bf16 = masters; fp8 = the published ROLLOUT_SOURCE_MODEL; nvfp4 = + miles' TE-direct quantizer over the masters (packing == the trainer's export packing). + """ + if getattr(exp, "DISABLE_HF_XET", False): + os.environ["HF_HUB_DISABLE_XET"] = "1" + os.environ.pop("HF_XET_HIGH_PERFORMANCE", None) + if getattr(exp, "DISABLE_HF_TRANSFER", False): + os.environ.pop("HF_HUB_ENABLE_HF_TRANSFER", None) + # Quantizer envs the conversion must share with the trainer (e.g. the 4/6 recipe's + # NVTE_NVFP4_4OVER6 family): the served base and the trainer export must be produced + # under identical settings or the delta baseline diverges from the export layout. + os.environ.update(getattr(exp, "PREP_ENV", {})) + from huggingface_hub import snapshot_download + + prep_volume.reload() + tag = exp.MODEL_TAG + bf16_dir, fp8_dir, nvfp4_dir = f"{PREP_PATH}/{tag}/bf16", f"{PREP_PATH}/{tag}/fp8", f"{PREP_PATH}/{tag}/nvfp4" + served_format = getattr(exp, "SERVED_CHECKPOINT_FORMAT", "nvfp4") + if served_format not in {"bf16", "fp8", "nvfp4"}: + raise SystemExit(f"unsupported SERVED_CHECKPOINT_FORMAT={served_format!r}") + tools = f"{MILES_ROOT}/tools" + + src = snapshot_download(exp.SOURCE_MODEL) + is_int4 = _is_int4(src) # read the source's quant scheme, not its repo name + + def _build_bf16(out: str) -> None: + if is_int4: + subprocess.run(["python", f"{tools}/convert_kimi_int4_to_bf16.py", "--model-dir", src, "--output-dir", out], check=True) + else: + subprocess.run(f"cp -aL {src}/. {out}/", shell=True, check=True) # -L: real files, not cache symlinks + _strip_stale_quant_config(os.path.join(out, "config.json")) + + _staged(bf16_dir, _build_bf16) + + if served_format == "bf16": + prep_volume.commit() + print(f"Prepared masters={bf16_dir} served_base={bf16_dir}") + return + + if served_format == "fp8": + fp8_source = getattr(exp, "ROLLOUT_SOURCE_MODEL", None) + if not fp8_source: + raise SystemExit("SERVED_CHECKPOINT_FORMAT='fp8' requires ROLLOUT_SOURCE_MODEL") + _staged(fp8_dir, lambda out: subprocess.run(f"cp -aL {snapshot_download(fp8_source)}/. {out}/", shell=True, check=True)) + prep_volume.commit() + print(f"Prepared masters={bf16_dir} served_base={fp8_dir}") + return + + # nvfp4: miles' TE-direct quantizer. bf16 carve-outs must match the trainer's + # --num-layers-at-start/end-in-bf16 so the served base == the export layout. + carveouts: list[str] = [] + if (n := getattr(exp.miles, "num_layers_at_start_in_bf16", None)) is not None: + carveouts += ["--num-layers-at-start-in-bf16", str(n)] + if (n := getattr(exp.miles, "num_layers_at_end_in_bf16", None)) is not None: + carveouts += ["--num-layers-at-end-in-bf16", str(n)] + _staged(nvfp4_dir, lambda out: subprocess.run( + ["python", f"{tools}/convert_hf_to_nvfp4.py", "--model-dir", bf16_dir, "--save-dir", out, *carveouts], check=True)) + prep_volume.commit() + print(f"Prepared masters={bf16_dir} served_base={nvfp4_dir}") + + +def prepare_torch_dist(exp, prep_volume, *, rank: int, master_addr: str) -> None: + """Build {tag}/torch_dist (the raw-mode ref_load) from the {tag}/bf16 masters via a + clustered torchrun conversion (large MoE won't fit an 8-way split).""" + prep_volume.reload() + tag = exp.MODEL_TAG + bf16_dir, torch_dist_dir = f"{PREP_PATH}/{tag}/bf16", f"{PREP_PATH}/{tag}/torch_dist" + if os.path.exists(os.path.join(torch_dist_dir, "latest_checkpointed_iteration.txt")): + print(f"reusing existing torch_dist {torch_dist_dir}") + return + if not exp.miles.miles_model_script: + raise SystemExit("prepare_torch_dist requires miles_model_script (MODEL_ARGS)") + nodes = exp.modal.torch_dist_prep_nodes + use_wrapper = nodes > 1 and getattr(exp, "USE_MODAL_TORCH_DIST_WRAPPER", False) + convert = TORCH_DIST_CONVERT_WRAPPER if use_wrapper else f"{MILES_ROOT}/tools/convert_hf_to_torch_dist.py" + inner = ( + f"source {MILES_ROOT}/{exp.miles.miles_model_script} && " + f"PYTHONPATH={MEGATRON_PATH} torchrun" + f" --nnodes {nodes} --node-rank {rank} --master-addr {master_addr} --master-port 29500" + f" --nproc-per-node {exp.modal.torch_dist_prep_gpus_per_node}" + f" {convert} ${{MODEL_ARGS[@]}}" + f" --hf-checkpoint {bf16_dir} --save {torch_dist_dir} --megatron-to-hf-mode raw" + f" {exp.modal.torch_dist_convert_extra_args}" + ) + env = {**os.environ} + if use_wrapper: + env["SKIP_RELEASE_RENAME"] = "1" + subprocess.run(["bash", "-c", inner], check=True, env=env) + # Every node commits its own distcp shards (disjoint files merge on the Volume); + # a rank-0-only commit would drop the other nodes' shards. + prep_volume.commit() + if rank == 0: + print(f"Prepared torch_dist={torch_dist_dir}") + + +def _staged(final_dir: str, build) -> None: + """Build into a .partial sibling and atomically rename, so an interrupted step never + leaves a half-built dir the reuse check mistakes for complete.""" + if os.path.isdir(final_dir) and os.listdir(final_dir): + print(f"reusing existing {final_dir}") + return + partial = f"{final_dir}.partial" + subprocess.run(["rm", "-rf", partial], check=True) + os.makedirs(partial, exist_ok=True) + build(partial) + os.rename(partial, final_dir) + + +def _strip_stale_quant_config(config_path: str) -> None: + """Drop any quantization_config (top-level and text_config-nested) from an HF config, + so the bf16 masters don't claim the source's quant scheme.""" + if not os.path.exists(config_path): + return + with open(config_path) as f: + cfg = json.load(f) + removed = bool(cfg.pop("quantization_config", None)) + if isinstance(cfg.get("text_config"), dict): + removed = bool(cfg["text_config"].pop("quantization_config", None)) or removed + if removed: + with open(config_path, "w") as f: + json.dump(cfg, f, indent=2) + print(f"stripped stale quantization_config from {config_path}") + + +def _is_int4(model_dir: str) -> bool: + cfg_path = os.path.join(model_dir, "config.json") + if not os.path.exists(cfg_path): + return False + with open(cfg_path) as f: + cfg = json.load(f) or {} + # VLMs (Kimi K2.x) nest the quant config under text_config. + qc = (cfg.get("text_config") or {}).get("quantization_config") or cfg.get("quantization_config") or {} + return qc.get("quant_method") == "compressed-tensors" diff --git a/cookbook/miles_disagg/serving.py b/cookbook/miles_disagg/serving.py deleted file mode 100644 index 1bd2635..0000000 --- a/cookbook/miles_disagg/serving.py +++ /dev/null @@ -1,130 +0,0 @@ -"""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. -""" - -from __future__ import annotations - -from pathlib import Path - -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. - - -def build_nvfp4_b200_serving_image( - *, - miles_repo_url: str, - miles_repo_ref: str, - miles_root: str, - 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__"], - ) - ) diff --git a/cookbook/miles_disagg/sidecar.py b/cookbook/miles_disagg/sidecar.py deleted file mode 100644 index 912bc23..0000000 --- a/cookbook/miles_disagg/sidecar.py +++ /dev/null @@ -1,224 +0,0 @@ -"""SGLang weight-sync sidecar launcher 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``). - -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. -""" - -from __future__ import annotations - -import argparse -import logging -import os - -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, - ) - - -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", - ) - - -if __name__ == "__main__": - main() diff --git a/cookbook/miles_disagg/trainer_image.py b/cookbook/miles_disagg/trainer_image.py new file mode 100644 index 0000000..82ab09a --- /dev/null +++ b/cookbook/miles_disagg/trainer_image.py @@ -0,0 +1,70 @@ +"""The miles trainer image + the versions pinned to launch a miles run. + +The base image bakes Megatron-LM (native --fp4-format NVFP4) + TransformerEngine; the +miles fork is cloned over it at a pinned commit. The serving half is separate and shared +(common/serving_image.py) — the pool installs no trainer package. +""" + +from __future__ import annotations + +from pathlib import Path + +import modal + +# Dated tag, never `latest`: Modal caches from_registry per tag string and won't re-pull +# a moved mutable tag, so `latest` silently serves whatever was first pulled. +MILES_IMAGE_TAG = "radixark/miles:dev-202607090055" +MILES_REPO_URL = "https://github.com/modal-projects/miles.git" +MILES_REPO_REF = "bdaf8d04bbc26981ff4bc95b8cfe679c3cd29013" + +MILES_ROOT = "/root/miles" +MEGATRON_PATH = "/root/Megatron-LM" # source-only megatron.training must be on PYTHONPATH +TORCH_DIST_CONVERT_WRAPPER = "/root/convert_hf_to_torch_dist_modal.py" + +_COOKBOOK_DIR = Path(__file__).resolve().parent.parent # .../cookbook +_TORCH_DIST_WRAPPER_SRC = Path(__file__).resolve().parent / "convert_hf_to_torch_dist_modal.py" + + +def build_trainer_image( + *, + hf_cache_path: str, + experiment: str, + miles_local: str | None = None, + image_tag: str | None = None, + miles_repo_ref: str | None = None, + setup_commands: tuple[str, ...] = (), +) -> modal.Image: + """The miles trainer image: RDMA/EFA userspace + the pinned miles fork + the + trainer-side delta encoder's codecs. stitch + the cookbook package are mounted so the + trainer, Ray actors, and the sidecar subprocess resolve their imports. + + ``image_tag`` / ``miles_repo_ref`` / ``setup_commands`` let an experiment carry a + different stack (e.g. the 4/6 recipe's TE 2.17 upgrade) without forking the app; + ``setup_commands`` run after the cache cleanup and before the miles clone, so a setup + change never invalidates the miles layer cache.""" + image = ( + modal.Image.from_registry(image_tag or MILES_IMAGE_TAG) + .entrypoint([]) + # RDMA/EFA userspace so multi-node NCCL binds EFA under rdma=True instead of TCP. + .apt_install("libibverbs-dev", "libibverbs1", "libhwloc-dev", "libnl-route-3-200") + .run_commands(f"rm -rf {hf_cache_path}") # baked HF cache must not shadow the mounted volume + ) + for cmd in setup_commands: + image = image.run_commands(cmd) + image = ( + image.run_commands( + f"rm -rf {MILES_ROOT}" + f" && git clone {MILES_REPO_URL} {MILES_ROOT}" + f" && cd {MILES_ROOT} && git fetch origin {miles_repo_ref or MILES_REPO_REF} && git checkout FETCH_HEAD" + f" && python3 -m pip install --no-deps -e {MILES_ROOT}" + ) + # The trainer-side delta ENCODER (miles delta.py) needs the codecs even under --no-deps. + .pip_install("fastapi", "httpx", "uvicorn", "zstandard", "xxhash", "blake3") + .env({"HF_XET_HIGH_PERFORMANCE": "1", "HF_HUB_ENABLE_HF_TRANSFER": "1", "EXPERIMENT_CONFIG": experiment}) + .add_local_file(str(_TORCH_DIST_WRAPPER_SRC), TORCH_DIST_CONVERT_WRAPPER, copy=True) + .add_local_python_source("stitch") + .add_local_dir(str(_COOKBOOK_DIR), remote_path="/root/cookbook", ignore=["**/__pycache__"]) + ) + if miles_local: # dev overlay: replace the cloned fork with a local checkout (no rebuild) + image = image.add_local_dir(miles_local, remote_path=MILES_ROOT, ignore=[".git", "**/__pycache__", "**/*.pyc"]) + return image diff --git a/cookbook/ray_cluster.py b/cookbook/ray_cluster.py deleted file mode 100644 index 86fee30..0000000 --- a/cookbook/ray_cluster.py +++ /dev/null @@ -1,139 +0,0 @@ -"""Ray cluster bring-up helpers for multi-node Modal training. - -Shared across cookbook trainers (slime, miles, etc.). Originally adapted from -https://github.com/modal-projects/multinode-training-guide. -""" - -from __future__ import annotations - -import socket -import subprocess -import time -from pathlib import Path -from typing import Any - - -RAY_START_TIMEOUT = 240 -RAY_WORKER_JOIN_TIMEOUT = 180 - - -def training_nodes(cfg: Any) -> int: - nodes = int(getattr(cfg, "actor_num_nodes", 1)) - if getattr(cfg, "use_critic", False) or getattr(cfg, "advantage_estimator", None) == "ppo": - nodes += int(getattr(cfg, "critic_num_nodes", nodes)) - return nodes - - -def get_modal_cluster_context(n_nodes: int) -> tuple[int, str, str]: - """Return (rank, master_addr, my_ip) for the current Modal cluster.""" - import modal.experimental - - try: - info = modal.experimental.get_cluster_info() - except Exception: # noqa: BLE001 - if n_nodes == 1: - ip = _get_local_ip() - return 0, ip, ip - raise - actual_nodes = len(info.container_ipv4_ips) - if actual_nodes == 0 and n_nodes == 1: - ip = _get_local_ip() - return 0, ip, ip - if actual_nodes != n_nodes: - raise RuntimeError(f"cluster size mismatch: expected {n_nodes} node(s), got {actual_nodes}") - return info.rank, info.container_ipv4_ips[0], info.container_ipv4_ips[info.rank] - - -def _get_local_ip() -> str: - with socket.socket(socket.AF_INET, socket.SOCK_DGRAM) as sock: - try: - sock.connect(("8.8.8.8", 80)) - return sock.getsockname()[0] - except OSError: - return socket.gethostbyname(socket.gethostname()) - - -def start_ray_head(my_ip: str, n_nodes: int, *, ray_port: int) -> None: - """Start the Ray head node and wait for all workers to join.""" - import ray - - start_cmd = [ - "ray", - "start", - "--head", - f"--node-ip-address={my_ip}", - f"--port={ray_port}", - "--disable-usage-stats", - "--include-dashboard=false", - ] - try: - subprocess.run(start_cmd, check=True, timeout=RAY_START_TIMEOUT) - except subprocess.TimeoutExpired as exc: - _print_ray_logs() - raise RuntimeError(f"Ray head node did not finish startup within {RAY_START_TIMEOUT}s") from exc - except subprocess.CalledProcessError as exc: - _print_ray_logs() - raise RuntimeError(f"Ray head node failed to start with exit code {exc.returncode}") from exc - - last_error = "" - for _ in range(RAY_START_TIMEOUT): - try: - ray.init(address=f"{my_ip}:{ray_port}") - break - except Exception as exc: # noqa: BLE001 - last_error = f"{type(exc).__name__}: {exc}" - time.sleep(1) - else: - _print_ray_logs() - raise RuntimeError(f"Ray head node failed to start before timeout: {last_error}") - - for _ in range(RAY_WORKER_JOIN_TIMEOUT): - alive = [n for n in ray.nodes() if n["Alive"]] - print(f"Waiting for workers: {len(alive)}/{n_nodes} alive") - if len(alive) == n_nodes: - break - time.sleep(1) - else: - _print_ray_logs() - raise RuntimeError(f"Timed out waiting for all {n_nodes} Ray nodes to join") - - -def start_ray_worker(my_ip: str, master_addr: str, *, ray_port: int) -> None: - """Join this container to the Ray cluster as a worker node.""" - subprocess.run( - [ - "ray", - "start", - f"--node-ip-address={my_ip}", - "--address", - f"{master_addr}:{ray_port}", - "--disable-usage-stats", - ], - check=True, - timeout=RAY_START_TIMEOUT, - ) - - -def _print_ray_logs() -> None: - log_dir = Path("/tmp/ray/session_latest/logs") - for name in ( - "dashboard.log", - "dashboard.err", - "gcs_server.out", - "gcs_server.err", - "raylet.out", - "raylet.err", - "monitor.out", - "monitor.err", - ): - path = log_dir / name - if not path.exists(): - continue - print(f"===== {path} =====") - try: - lines = path.read_text(errors="replace").splitlines() - except OSError as exc: - print(f"could not read {path}: {exc}") - continue - for line in lines[-80:]: - print(line) diff --git a/cookbook/sidecar_process.py b/cookbook/sidecar_process.py deleted file mode 100644 index ad76d3f..0000000 --- a/cookbook/sidecar_process.py +++ /dev/null @@ -1,87 +0,0 @@ -"""Sidecar process launch + HTTP liveness helpers. - -Shared across cookbook trainers (slime, miles, etc.). The sidecar module path is -parameterized so each trainer can wire its own sidecar entry point. -""" - -from __future__ import annotations - -import os -import signal -import subprocess -import time -import urllib.error -import urllib.request - - -def start_sglang_sidecar( - *, - sidecar_module: str, - sidecar_port: int, - sglang_port: int, - bulletin_root: str, - local_checkpoint_dir: str, - base_checkpoint_dir: str, - volume_name: str, - commit_mode: str, - debug_requests: bool = False, -) -> subprocess.Popen: - """Launch the weight-sync sidecar as a subprocess. - - ``sidecar_module`` is the Python module path (e.g. - ``"cookbook.slime_disagg.sidecar"``). - """ - cmd = [ - "python3", - "-m", - sidecar_module, - "--host", - "0.0.0.0", - "--port", - str(sidecar_port), - "--upstream-url", - f"http://127.0.0.1:{sglang_port}", - "--bulletin-root", - bulletin_root, - "--local-checkpoint-dir", - local_checkpoint_dir, - "--base-checkpoint-dir", - base_checkpoint_dir, - "--volume-name", - volume_name, - "--commit-mode", - commit_mode, - ] - if debug_requests: - cmd.append("--debug-requests") - print("Starting sidecar:", " ".join(cmd)) - return subprocess.Popen(cmd, start_new_session=True) - - -def wait_http(url: str, process: subprocess.Popen | None, timeout: int) -> None: - deadline = time.time() + timeout - last_error: str | None = None - while time.time() < deadline: - if process is not None and process.poll() is not None: - raise RuntimeError(f"process exited while waiting for {url}: code={process.returncode}") - try: - with urllib.request.urlopen(url, timeout=5) as resp: - if 200 <= resp.status < 500: - return - except (urllib.error.URLError, TimeoutError, OSError) as exc: - last_error = f"{type(exc).__name__}: {exc}" - time.sleep(2) - raise TimeoutError(f"Timed out waiting for {url}; last error: {last_error}") - - -def terminate_process(process: subprocess.Popen | None) -> None: - if process is None or process.poll() is not None: - return - try: - os.killpg(process.pid, signal.SIGTERM) - process.wait(timeout=20) - except Exception: - try: - os.killpg(process.pid, signal.SIGKILL) - except Exception: - pass diff --git a/cookbook/slime_disagg/README.md b/cookbook/slime_disagg/README.md deleted file mode 100644 index 211d067..0000000 --- a/cookbook/slime_disagg/README.md +++ /dev/null @@ -1,103 +0,0 @@ -# Disaggregated SLIME on Modal - -Train a model with GRPO while rollouts run on a separate, elastic pool of -SGLang servers. - -The app has two halves. - -- **`Trainer`** is a clustered SLIME/Ray job. After each optimizer step it - writes a sparse weight delta to a Modal Volume (the "bulletin board") and - publishes a new weight version. -- **`Server`** is a Modal Flash pool of SGLang servers. A sidecar in each - container watches the bulletin board, applies deltas in order, and serves - rollout requests pinned to an exact weight version. Requests for a version - the container hasn't reached yet get `409` and SLIME retries. - -The two halves never talk directly. Weights flow through the Volume, rollout -traffic flows through the Flash gateway, and either side can scale or restart -on its own. - -## Layout - -`modal_train.py` defines the Modal app (the `Server` pool, the `Trainer` -cluster, and the setup/launch entrypoints); `configs/` holds one module per -experiment. The `stitch` package (this repo) provides the bulletin-board -protocol, the SGLang sidecar, and the SLIME hooks. Both `stitch` and this -example are mounted into containers at startup, so code edits never rebuild -the image. - -## Run it - -You need a Modal account and a `huggingface-secret` Modal secret. Work from -the repo root, with this alias to keep the commands short: - -```bash -alias m="uv run --extra modal modal" - -# One-time setup. Fetch the model and dataset onto Volumes. -m run -m cookbook.slime_disagg.modal_train::download_model -m run -m cookbook.slime_disagg.modal_train::prepare_dataset - -# Deploy the rollout pool + trainer. -m deploy -m cookbook.slime_disagg.modal_train - -# Wait for the pool to come up and answer at version 0. -m run -m cookbook.slime_disagg.modal_train::smoke_flash_pool - -# Train. Returns immediately; the run continues on Modal. -m run -m cookbook.slime_disagg.modal_train::launch_train - -# Smoke-check the pool at a given version (the chain advances one per rollout). -m run -m cookbook.slime_disagg.modal_train::smoke_flash_pool --weight-version 3 -``` - -To train again, just run `launch_train` again. Each launch gets a fresh run id -and writes its delta chain under its own `/` partition, so sequential -runs never collide — no bulletin-board reset between runs. The warm `Trainer` -cluster (Ray started once per container in `@modal.enter()`) goes straight to -training, and the rollout pool re-materializes to the new run's base on its own -when it sees the pointer move to the new ``. - -## Configuration - -Each experiment is a module in `configs/` holding a `ModalConfig` (GPU type, -pool size, regions), a `SlimeConfig` (every SLIME CLI arg), and the names of -the Modal resources it owns. - -`launch_train` imports the experiment from your local working tree and ships -the training arguments to the deployed `Trainer` as plain data. Editing or -adding a `SlimeConfig` therefore needs no redeploy; just `launch_train` it. -Infrastructure binds at deploy time, so changes to GPU type, node count, -pool size, Volume names, or SGLang server flags still need a deploy. - -The only environment variable is `EXPERIMENT_CONFIG`, which picks the config -module at deploy time. Each experiment becomes its own Modal app: - -```bash -EXPERIMENT_CONFIG= m deploy --strategy recreate -m cookbook.slime_disagg.modal_train -``` - -A config with a real acceptance signal (rather than the smoke-test protocol -check) runs the same transport while the reward metrics climb. The -`passrate/pass@1` and `passrate/pass@8` metrics should trend up. Useful log -searches (the app name is the config's `APP_NAME`): - -```bash -m app logs --since 4h --search "passrate " -m app logs --since 4h --search "Published sparse delta" -``` - -## Protocol notes - -Each launch gets a fresh `run_id`. The trainer writes that run's delta chain -under `/delta-bulletin//weight_v{N}/`, and a single `/delta-bulletin/latest` -pointer names the active snapshot (`/weight_v{N}`). Sidecars apply -versions in order from their current version; when the pointer moves to a new -run they re-materialize the base and replay that run's chain from the start. -Each run is isolated under its own `run_id`, so sequential runs never collide. -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. diff --git a/cookbook/slime_disagg/__init__.py b/cookbook/slime_disagg/__init__.py index 9bfb841..e69de29 100644 --- a/cookbook/slime_disagg/__init__.py +++ b/cookbook/slime_disagg/__init__.py @@ -1,2 +0,0 @@ -"""Modal Flash sparse-delta SLIME example.""" - diff --git a/cookbook/slime_disagg/app.py b/cookbook/slime_disagg/app.py new file mode 100644 index 0000000..7849ffc --- /dev/null +++ b/cookbook/slime_disagg/app.py @@ -0,0 +1,232 @@ +"""Disaggregated slime training on Modal, assembled on the stitch core. + +``EXPERIMENT_CONFIG`` selects a config module under ``cookbook.slime_disagg.configs``. The +Server (sglang + stitch sidecar) is the shared common one; the Trainer runs slime on Ray +and publishes XOR deltas through a Modal Volume the pool syncs from. + + EXPERIMENT_CONFIG=kimi_k2_6_int4 uv run --extra modal modal deploy -m cookbook.slime_disagg.app + +Config access is uniform: the experiment module ``exp`` is the single source of truth — +its ``exp.modal`` (infra), ``exp.slime`` (training), and ``exp.`` are read directly; +shared deployment constants come from ``common.constants``. ``ROLLOUT_CONCURRENCY`` is the +one resolved value (the experiment's Flash target, else the engine's concurrency). +""" + +from __future__ import annotations + +import importlib +import os +import subprocess +import tempfile +import uuid +from types import SimpleNamespace + +import modal +import modal.experimental + +from stitch.pools.modal_flash import ModalFlashPool + +from cookbook.common import launch, ray_cluster, serving_image, server, smoke +from cookbook.common.constants import ( + CHECKPOINTS_PATH, DATA_PATH, HF_CACHE_PATH, MINUTES, RAY_PORT, SERVER_STARTUP_TIMEOUT, SIDECAR_PORT, +) +from cookbook.slime_disagg import trainer_image +from cookbook.slime_disagg.config import SlimeConfig, YAML_CONFIG_FIELDS +from cookbook.slime_disagg.trainer_image import SLIME_ROOT + +# Deploy-time environment (selection + dev overlay, NOT experiment config). +EXPERIMENT = os.environ["EXPERIMENT_CONFIG"] # required; a default would silently serve the wrong experiment +SLIME_LOCAL_DIR = os.environ.get("SLIME_LOCAL_DIR") # optional dev overlay of a local slime checkout + +exp = importlib.import_module(f"cookbook.slime_disagg.configs.{EXPERIMENT}") +modal_cfg = exp.modal +slime_cfg = exp.slime + +# The Flash autoscaler target (and sglang concurrency cap): the experiment's explicit +# target_inputs, else the engine's configured concurrency. +ROLLOUT_CONCURRENCY = modal_cfg.rollout_target_inputs or slime_cfg.sglang_server_concurrency + +# EXPERIMENT_CONFIG is baked into both images (inside build_*_image) so the container's +# re-import resolves the same experiment as the deploy, not the default. +image = trainer_image.build_trainer_image(hf_cache_path=str(HF_CACHE_PATH), experiment=EXPERIMENT, slime_local=SLIME_LOCAL_DIR) +server_image = serving_image.build_serving_image(hf_cache_path=str(HF_CACHE_PATH), delta_volume_name=exp.DELTA_VOLUME_NAME, experiment=EXPERIMENT) +if SLIME_LOCAL_DIR: + server_image = server_image.add_local_dir(SLIME_LOCAL_DIR, remote_path=SLIME_ROOT, ignore=[".git", "**/__pycache__", "**/*.pyc"]) + +hf_cache_volume = modal.Volume.from_name("huggingface-cache", create_if_missing=True, version=2) +data_volume = modal.Volume.from_name("slime-data", create_if_missing=True, version=2) +checkpoints_volume = modal.Volume.from_name("slime-checkpoints", create_if_missing=True, version=2) +delta_volume = modal.Volume.from_name(exp.DELTA_VOLUME_NAME, create_if_missing=True, version=2) + +train_volumes = { + str(HF_CACHE_PATH): hf_cache_volume, + str(DATA_PATH): data_volume, + str(CHECKPOINTS_PATH): checkpoints_volume, + exp.DELTA_BULLETIN_ROOT: delta_volume, +} + +app = modal.App(exp.APP_NAME) + +SGLANG_SERVER_ARGS = { + "--served-model-name": slime_cfg.hf_checkpoint, + "--dtype": "bfloat16", + "--cuda-graph-max-bs-decode": str(ROLLOUT_CONCURRENCY), + "--max-running-requests": str(ROLLOUT_CONCURRENCY), + "--trust-remote-code": "", + "--custom-pull-weights-pre-read-hook": "stitch.stores.modal_volume.pull_weights_pre_read_hook", + **exp.SGLANG_SERVER_ARGS, +} + + +# The rollout Server: a thin module-level class (Modal requires @app.cls at global scope) +# whose enter/exit delegate to the shared common.server logic. sglang + the stitch sidecar. +@app.cls( + image=server_image, + gpu=f"{modal_cfg.gpu}:{slime_cfg.rollout_num_gpus_per_engine}", + cloud=modal_cfg.cloud, region=modal_cfg.region, + volumes={str(HF_CACHE_PATH): hf_cache_volume, exp.DELTA_BULLETIN_ROOT: delta_volume}, + min_containers=modal_cfg.rollout_min_containers, max_containers=modal_cfg.rollout_max_containers, + timeout=40 * MINUTES, scaledown_window=15 * MINUTES, + ephemeral_disk=modal_cfg.rollout_ephemeral_disk_mib, memory=modal_cfg.rollout_memory_mib, + include_source=False, +) +@modal.experimental.http_server( + port=SIDECAR_PORT, proxy_regions=modal_cfg.proxy_regions, + exit_grace_period=25, startup_timeout=SERVER_STARTUP_TIMEOUT, +) +@modal.concurrent(target_inputs=ROLLOUT_CONCURRENCY) +class Server: + @modal.enter() + def startup(self) -> None: + server.serve_startup( + self, model_name=slime_cfg.hf_checkpoint, sglang_args=SGLANG_SERVER_ARGS, + tp=slime_cfg.rollout_num_gpus_per_engine, concurrency=ROLLOUT_CONCURRENCY, + bulletin_root=exp.DELTA_BULLETIN_ROOT, local_checkpoint_dir=exp.LOCAL_CHECKPOINT_PATH, + volume_name=exp.DELTA_VOLUME_NAME, commit_mode=exp.SIDECAR_COMMIT_MODE, + startup_timeout=SERVER_STARTUP_TIMEOUT, sglang_env=getattr(exp, "SGLANG_ENV", {}), + ) + + @modal.exit() + def stop(self) -> None: + server.serve_stop(self) + + +# ── Trainer (slime on Ray) ──────────────────────────────────────────────────── +# Multi-node needs an RDMA gang (clustered) over the EFA fabric; single-node takes +# neither. Both are inline on the decorator so there's one declaration, not a rebind. +_MULTINODE = slime_cfg.n_train_nodes > 1 + + +@app.cls( + image=image, + gpu=f"{modal_cfg.gpu}:{slime_cfg.actor_num_gpus_per_node}", + memory=modal_cfg.memory, + cloud=modal_cfg.cloud, region=modal_cfg.region, + volumes=train_volumes, + ephemeral_disk=modal_cfg.trainer_ephemeral_disk_mib, + timeout=24 * 60 * MINUTES, startup_timeout=20 * MINUTES, scaledown_window=30 * MINUTES, + include_source=False, + **({"experimental_options": {"efa_enabled": True}} if _MULTINODE else {}), +) +@(modal.experimental.clustered(slime_cfg.n_train_nodes, rdma=True) if _MULTINODE else lambda c: c) +class Trainer: + """slime actor cluster. Ray comes up once per container in enter(), so back-to-back + runs reuse it.""" + + @modal.enter() + def start_ray(self) -> None: + from cookbook.common import process + + rank, master_addr, my_ip = ray_cluster.get_modal_cluster_context(slime_cfg.n_train_nodes) + self.rank = rank + process.start_host_mem_monitor() # per-node host-RAM trace (publish gather is the OOM peak) + os.environ.update({ + "SLIME_HOST_IP": my_ip, "SGLANG_HOST_IP": my_ip, "HOST_IP": my_ip, + "MASTER_ADDR": master_addr, "RAY_ADDRESS": f"{master_addr}:{RAY_PORT}", + "no_proxy": f"127.0.0.1,{master_addr},{my_ip}", "NO_PROXY": f"127.0.0.1,{master_addr},{my_ip}", + **slime_cfg.environment, + }) + if rank == 0: + ray_cluster.start_ray_head(my_ip, slime_cfg.n_train_nodes, ray_port=RAY_PORT) + else: + ray_cluster.start_ray_worker(my_ip, master_addr, ray_port=RAY_PORT) + + @modal.method() + def train(self, payload: dict) -> None: + """Run one training job from a SlimeConfig payload (see SlimeConfig.to_payload).""" + for volume in train_volumes.values(): + volume.reload() + if self.rank != 0: # rank 0 drives; other ranks only need their Ray workers (started in enter) + return + + cfg = SlimeConfig.from_payload(payload) + cfg.rollout_endpoint_url = ModalFlashPool(exp.APP_NAME, "Server").gateway_url() + run_id = uuid.uuid4().hex[:12] # per-launch fence token; forks a fresh chain + cfg.update_weight_disk_dir = f"{exp.DELTA_BULLETIN_ROOT}/{run_id}" + hook_knobs = { + "update_weight_delta_volume_name": exp.DELTA_VOLUME_NAME, + "rollout_modal_flash_app_name": exp.APP_NAME, + "rollout_modal_flash_server_cls_name": "Server", + "run_id": run_id, + } + cfg.custom_config_path = hook_knobs # materialized to a YAML path below; keep the mapping for claim + launch.resolve_config(cfg, tempfile.mkdtemp(), YAML_CONFIG_FIELDS) + cmd = launch.build_train_cmd(cfg, SLIME_ROOT, "slime_model_script") + + # Claim the pool before slime publishes: reset every replica to base for this run. + from cookbook.common import hooks + + hooks.claim_pool(SimpleNamespace(update_weight_disk_dir=cfg.update_weight_disk_dir, **hook_knobs)) + + print(f"Training {EXPERIMENT}: nodes={slime_cfg.n_train_nodes}, rollout_endpoint={cfg.rollout_endpoint_url}") + print(f"Command: {cmd}") + subprocess.run(["bash", "-lc", cmd], check=True) + + +# ── Preparation + entrypoints ────────────────────────────────────────────────── +@app.function( + image=image, volumes={str(HF_CACHE_PATH): hf_cache_volume}, + timeout=2 * 60 * MINUTES, secrets=[modal.Secret.from_name("huggingface-secret")], include_source=False, +) +def download_model() -> None: + """Snapshot the served model into the HF cache (sglang serves it by repo id).""" + from huggingface_hub import snapshot_download + + snapshot_download(repo_id=slime_cfg.hf_checkpoint) + hf_cache_volume.commit() + + +@app.function( + image=image, volumes={str(DATA_PATH): data_volume}, + timeout=2 * 60 * MINUTES, secrets=[modal.Secret.from_name("huggingface-secret")], include_source=False, +) +def prepare_dataset() -> None: + data_volume.reload() + slime_cfg.prepare_data() + data_volume.commit() + + +@app.local_entrypoint() +def launch_train() -> None: + """Spawn training on the deployed app. Config ships as data, so edits run without a + redeploy; infrastructure changes still require one.""" + from modal.exception import NotFoundError + + try: + trainer = modal.Cls.from_name(exp.APP_NAME, "Trainer")() + call = trainer.train.spawn(slime_cfg.to_payload()) + except NotFoundError: + raise SystemExit( + f"App {exp.APP_NAME!r} is not deployed. Run:\n" + f" EXPERIMENT_CONFIG={EXPERIMENT} uv run --extra modal modal deploy -m cookbook.slime_disagg.app" + ) + print(f"Spawned train on {exp.APP_NAME}: {call.object_id}") + + +@app.local_entrypoint() +def smoke_flash_pool(weight_version: int = 0, timeout_seconds: int = 30 * MINUTES) -> None: + """Check the deployed Flash pool serves completions at the expected weight version.""" + smoke.smoke_flash_pool( + app_name=exp.APP_NAME, cls_name="Server", model_name=slime_cfg.hf_checkpoint, + weight_version=weight_version, timeout_seconds=timeout_seconds, + ) diff --git a/cookbook/slime_disagg/config.py b/cookbook/slime_disagg/config.py new file mode 100644 index 0000000..e274cc0 --- /dev/null +++ b/cookbook/slime_disagg/config.py @@ -0,0 +1,92 @@ +"""``SlimeConfig`` — slime training arguments as a reflected config (self-contained). + +Every public, non-callable attribute becomes a slime CLI arg via ``cli_args``; +``environment`` / ``async_mode`` / ``slime_model_script`` are launcher instructions, not +CLI args. The Modal-infra half of an experiment is ``common.config.ModalConfig``. +""" + +from __future__ import annotations + +from typing import Any + +# Fields that are launcher instructions, not slime CLI args. +_SLIME_SKIP = {"environment", "async_mode", "slime_model_script"} +# Fields slime reads as YAML files; inline dicts are materialized before launch. +YAML_CONFIG_FIELDS = ("eval_config", "custom_config_path", "sglang_config") + + +class SlimeConfig: + """Subclass and set class attributes; all public, non-callable, non-skip attributes + become slime CLI args via ``cli_args``.""" + + environment: dict = {} + async_mode: bool = False # True -> train_async.py + slime_model_script: str = "" # shell script (relative to the slime root) defining MODEL_ARGS + + def __init__(self, **kwargs: Any) -> None: + self.environment = dict(type(self).environment) # fresh per instance; never mutate the class default + for k, v in kwargs.items(): + setattr(self, k, v) + + @property + def n_train_nodes(self) -> int: + """Trainer node count: actor nodes, plus critic nodes for PPO/critic setups.""" + nodes = int(getattr(self, "actor_num_nodes", 1)) + if getattr(self, "use_critic", False) or getattr(self, "advantage_estimator", None) == "ppo": + nodes += int(getattr(self, "critic_num_nodes", nodes)) + return nodes + + def _fields(self) -> dict[str, Any]: + """Merged fields across the class hierarchy; instance attrs win.""" + fields: dict[str, Any] = {} + for cls in reversed(type(self).__mro__): + if cls is object: + continue + fields.update( + { + k: v + for k, v in vars(cls).items() + if not k.startswith("_") + and not callable(v) + and not isinstance(v, (classmethod, staticmethod, property)) + } + ) + fields.update(vars(self)) + return {k: v for k, v in fields.items() if k not in _SLIME_SKIP} + + def cli_args(self) -> list[str]: + """slime CLI args: field_name -> --field-name; True -> bare flag; False/None -> + omitted; list -> --flag v1 v2; else --flag value.""" + out: list[str] = [] + for key, val in self._fields().items(): + if val is None or val is False: + continue + flag = f"--{key.replace('_', '-')}" + if val is True: + out.append(flag) + elif isinstance(val, list): + out += [flag] + [str(v) for v in val] + else: + out += [flag, str(val)] + return out + + def prepare_data(self) -> None: + raise NotImplementedError(f"{type(self).__name__} has no prepare_data()") + + def to_payload(self) -> dict[str, Any]: + """Flatten to plain data so the launcher can ship a config to the deployed Trainer + — new or edited experiments run without a redeploy.""" + return { + "fields": self._fields(), + "environment": dict(self.environment), + "async_mode": self.async_mode, + "slime_model_script": self.slime_model_script, + } + + @classmethod + def from_payload(cls, payload: dict[str, Any]) -> "SlimeConfig": + cfg = cls(**payload["fields"]) + cfg.environment = dict(payload["environment"]) + cfg.async_mode = payload["async_mode"] + cfg.slime_model_script = payload["slime_model_script"] + return cfg diff --git a/cookbook/slime_disagg/configs/__init__.py b/cookbook/slime_disagg/configs/__init__.py index 08d8d4c..e69de29 100644 --- a/cookbook/slime_disagg/configs/__init__.py +++ b/cookbook/slime_disagg/configs/__init__.py @@ -1,2 +0,0 @@ -"""Example SLIME configs for Modal Flash sparse-delta training.""" - diff --git a/cookbook/slime_disagg/configs/base.py b/cookbook/slime_disagg/configs/base.py deleted file mode 100644 index a325432..0000000 --- a/cookbook/slime_disagg/configs/base.py +++ /dev/null @@ -1,149 +0,0 @@ -"""Base configuration classes and volume mount paths. - -Two separate concerns: - - ModalConfig — Modal infrastructure (GPU model, regions, rollout pool size) - SlimeConfig — SLIME training arguments - -Each experiment module defines one instance of each, plus a handful of -module-level constants (APP_NAME, DELTA_VOLUME_NAME, ...) that name the -Modal resources the experiment owns. All non-private, non-callable -attributes on a SlimeConfig subclass become SLIME CLI args automatically -via cli_args(). The 'environment' field is the only exception — it is -exported into the Ray runtime environment, not passed to SLIME directly. -""" - -from pathlib import Path -from typing import Any, Literal - -# ── Volume mount paths ──────────────────────────────────────────────────────── - -HF_CACHE_PATH = Path("/root/.cache/huggingface") -DATA_PATH = Path("/data") -CHECKPOINTS_PATH = Path("/checkpoints") - -# ── Types ───────────────────────────────────────────────────────────────────── - -GPUType = Literal["H100", "H200", "B200", "B300", "A100"] - -# Fields on SlimeConfig that are NOT SLIME CLI args. -_SLIME_SKIP = {"environment", "async_mode", "slime_model_script"} - -# SlimeConfig fields that SLIME reads as YAML files at runtime. -# Experiments may set these as inline dicts; the launcher materializes -# them to temp YAML files before building the CLI command. -YAML_CONFIG_FIELDS = ("eval_config", "custom_config_path", "sglang_config") - - -class ModalConfig: - """Modal infrastructure configuration.""" - - gpu: GPUType = "H200" - memory: tuple[int, int] | None = ( - None # per-container memory in MiB (request, limit) - ) - cloud: str | None = None # e.g. "aws", "gcp" - region: str | None = None # e.g. "us-east-2" - rollout_min_containers: int = 4 # warm Flash rollout containers - proxy_regions: list[str] = ["us-east"] # Flash gateway proxy regions - - def __init__(self, **kwargs: Any) -> None: - for k, v in kwargs.items(): - setattr(self, k, v) - - -class SlimeConfig: - """Base SLIME training configuration. - - Subclass and set class attributes to configure an experiment. - All attributes (except those in _SLIME_SKIP) are forwarded to SLIME as - CLI args. Each experiment must be fully self-contained — no inherited - defaults beyond this base class. - - Fields in _SLIME_SKIP are launcher instructions, not SLIME CLI args: - environment — exported into the Ray runtime environment - async_mode — selects train_async.py vs train.py - slime_model_script — path relative to /root/slime to a shell script that - defines MODEL_ARGS for model architecture; sourced - before running the train command - """ - - # Launcher instructions — not passed to SLIME CLI (see _SLIME_SKIP). - environment: dict = { - "CUDA_DEVICE_MAX_CONNECTIONS": "1", - "NCCL_NVLS_ENABLE": "1", - } - async_mode: bool = False # True → use train_async.py - slime_model_script: str = "" # shell script path relative to /root/slime - - def __init__(self, **kwargs: Any) -> None: - # Fresh environment dict per instance — never mutate the class-level default. - self.environment = dict(type(self).environment) - for k, v in kwargs.items(): - setattr(self, k, v) - - def _fields(self) -> dict[str, Any]: - """Merged field dict from the class hierarchy; instance attrs win.""" - fields: dict[str, Any] = {} - for cls in reversed(type(self).__mro__): - if cls is object: - continue - fields.update( - { - k: v - for k, v in vars(cls).items() - if not k.startswith("_") - and not callable(v) - and not isinstance(v, (classmethod, staticmethod, property)) - } - ) - fields.update(vars(self)) - return {k: v for k, v in fields.items() if k not in _SLIME_SKIP} - - def cli_args(self) -> list[str]: - """SLIME CLI arguments derived from this config. - - Conversion rules: - field_name → --field-name (underscore to hyphen) - True → --flag (no value) - False/None → omitted - list → --flag v1 v2 ... - other → --flag value - """ - out: list[str] = [] - for key, val in self._fields().items(): - if val is None or val is False: - continue - flag = f"--{key.replace('_', '-')}" - if val is True: - out.append(flag) - elif isinstance(val, list): - out += [flag] + [str(v) for v in val] - else: - out += [flag, str(val)] - return out - - def prepare_data(self) -> None: - raise NotImplementedError(f"{type(self).__name__} has no prepare_data()") - - def to_payload(self) -> dict[str, Any]: - """Flatten to plain data for sending to a deployed Trainer. - - launch_train resolves config modules locally and ships the result, so - the deployed app never imports them — new or edited experiments run - without a redeploy. - """ - return { - "fields": self._fields(), - "environment": dict(self.environment), - "async_mode": self.async_mode, - "slime_model_script": self.slime_model_script, - } - - @classmethod - def from_payload(cls, payload: dict[str, Any]) -> "SlimeConfig": - cfg = cls(**payload["fields"]) - cfg.environment = dict(payload["environment"]) - cfg.async_mode = payload["async_mode"] - cfg.slime_model_script = payload["slime_model_script"] - return cfg diff --git a/cookbook/slime_disagg/configs/kimi_k2_6_int4.py b/cookbook/slime_disagg/configs/kimi_k2_6_int4.py new file mode 100644 index 0000000..d34b97e --- /dev/null +++ b/cookbook/slime_disagg/configs/kimi_k2_6_int4.py @@ -0,0 +1,168 @@ +"""Kimi K2.6 GRPO on Modal, disaggregated, native-INT4 end to end. + +Deploy: EXPERIMENT_CONFIG=kimi_k2_6_int4 uv run --extra modal modal deploy -m cookbook.slime_disagg.app +""" + +from __future__ import annotations + +from cookbook.common.config import ModalConfig +from cookbook.common.constants import DATA_PATH +from cookbook.slime_disagg.config import SlimeConfig + +APP_NAME = "stitch-kimi-k2-6-int4" +DELTA_VOLUME_NAME = "stitch-delta-kimi-k2-6-int4" +DELTA_BULLETIN_ROOT = "/delta-bulletin" +LOCAL_CHECKPOINT_PATH = "/local-checkpoint" + +# QAT grouping; MUST match the served INT4 checkpoint's compressed-tensors group_size. +INT4_GROUP_SIZE = "32" + +# in_place applies weights without draining in-flight rollouts; stale KV isolated per version. +SIDECAR_COMMIT_MODE = "in_place" + +SGLANG_SERVER_ARGS = { + "--weight-loader-prefetch-checkpoints": "", + "--weight-loader-prefetch-num-threads": "8", + "--tool-call-parser": "kimi_k2", + "--reasoning-parser": "kimi_k2", + "--dist-timeout": "3600", + "--kv-cache-dtype": "fp8_e4m3", + "--attention-backend": "tokenspeed_mla", + "--context-length": "32768", + "--mem-fraction-static": "0.85", + "--chunked-prefill-size": "16384", + "--schedule-conservativeness": "0.5", + "--schedule-policy": "lpm", + "--enable-hierarchical-cache": "", + "--hicache-ratio": "2", + "--hicache-io-backend": "kernel", + "--hicache-mem-layout": "page_first", + "--hicache-write-policy": "write_through", + "--skip-server-warmup": "", + # Routing replay: the pool emits per-token routed experts so the trainer can replay them. + "--enable-return-routed-experts": "", +} + +modal = ModalConfig( + gpu="B200", + region="us", + rollout_min_containers=2, + rollout_target_inputs=256, + proxy_regions=["us-west"], +) + + +class _Slime(SlimeConfig): + # Arch comes from the model script. + slime_model_script = "scripts/models/kimi-k2-thinking.sh" + + # The native-INT4 base is the served model, the QAT init, and the disk-delta base. + hf_checkpoint = "moonshotai/Kimi-K2.6" + ref_load = hf_checkpoint + megatron_to_hf_mode = "bridge" + + actor_num_nodes = 32 # 32x8 = 256 GPUs + actor_num_gpus_per_node = 8 + colocate = False + rollout_num_gpus = 0 # external rollout: the framework runs no local engines + rollout_num_gpus_per_engine = 4 # B200:4 per rollout container (native INT4 fits) + rollout_endpoint_url = None # filled at launch from the pool gateway + + # The three plug points stitch fills (slime's publish hook key is custom_delta_pre_push_path): + custom_rollout_request_hook_path = "cookbook.common.hooks.gated_rollout_request_hook" + custom_delta_pre_push_path = "cookbook.common.hooks.commit_and_wake" + rollout_request_weight_version_mode = "min" + rollout_request_weight_version_lag = 1 # bounded staleness window + rollout_request_retry_attempts = 240 + rollout_request_retry_sleep = 1.0 + rollout_session_affinity_header = "Modal-Session-ID" + + async_mode = True + update_weights_interval = 1 + update_weight_mode = "delta" + update_weight_transport = "disk" + update_weight_delta_encoding = "xor" + update_weight_delta_checksum = "xxh3-128" + update_weight_disk_dir = DELTA_BULLETIN_ROOT # run-scoped at launch to / + + prompt_data = f"{DATA_PATH}/dapo-math-17k/dapo-math-17k.jsonl" + input_key = "prompt" + label_key = "label" + apply_chat_template = True + rollout_shuffle = True + rm_type = "math" + eval_interval = None + + rollout_function_path = "slime.rollout.sglang_rollout.generate_rollout" + num_rollout = 20 + rollout_batch_size = 128 + rollout_max_response_len = 16384 + rollout_temperature = 0.8 + rollout_top_p = 1.0 + n_samples_per_prompt = 8 + over_sampling_batch_size = 256 + dynamic_sampling_filter_path = "slime.rollout.filter_hub.dynamic_sampling_filters.check_reward_nonzero_std" + num_steps_per_rollout = 4 + balance_data = True + sglang_server_concurrency = 256 + use_fault_tolerance = False + + tensor_model_parallel_size = 8 + pipeline_model_parallel_size = 8 + context_parallel_size = 4 + expert_model_parallel_size = 32 + expert_tensor_parallel_size = 1 + decoder_last_pipeline_num_layers = 5 + sequence_parallel = True + use_dynamic_batch_size = True + max_tokens_per_gpu = 16384 + recompute_granularity = "full" + recompute_method = "uniform" + recompute_num_layers = 1 + attention_dropout = 0.0 + hidden_dropout = 0.0 + accumulate_allreduce_grads_in_fp32 = True + attention_softmax_in_fp32 = True + no_check_for_nan_in_loss_and_grad = True + + optimizer = "adam" + lr = 1e-6 + lr_decay_style = "constant" + weight_decay = 0.1 + adam_beta1 = 0.9 + adam_beta2 = 0.98 + optimizer_cpu_offload = True + overlap_cpu_optimizer_d2h_h2d = True + use_precision_aware_optimizer = True + + advantage_estimator = "grpo" + eps_clip = 0.2 + eps_clip_high = 0.28 + use_kl_loss = True + kl_loss_coef = 0.0 + kl_loss_type = "low_var_kl" + entropy_coef = 0.0 + use_tis = True + + use_rollout_routing_replay = True # needs the pool's --enable-return-routed-experts + + environment = { + "CUDA_DEVICE_MAX_CONNECTIONS": "1", + "NCCL_NVLS_ENABLE": "1", + "NVSHMEM_DISABLE_NCCL": "1", + "NCCL_TIMEOUT_MS": "360000000", + "OPEN_TRAINING_INT4_FAKE_QAT_FLAG": "1", + "OPEN_TRAINING_INT4_GROUP_SIZE": INT4_GROUP_SIZE, + } + + def prepare_data(self) -> None: + from datasets import load_dataset + + ds = load_dataset("BytedTsinghua-SIA/DAPO-Math-17k", split="train") + ds = ds.shuffle(seed=42).select(range(min(50000, ds.num_rows))) + ds = ds.map(lambda ex: {"label": ex["reward_model"]["ground_truth"]}) + ds = ds.select_columns(["prompt", "label"]) + ds.to_json(f"{DATA_PATH}/dapo-math-17k/dapo-math-17k.jsonl") + + +slime = _Slime() diff --git a/cookbook/slime_disagg/configs/kimi_k2_6_int4_disagg.py b/cookbook/slime_disagg/configs/kimi_k2_6_int4_disagg.py deleted file mode 100644 index e2ea7a4..0000000 --- a/cookbook/slime_disagg/configs/kimi_k2_6_int4_disagg.py +++ /dev/null @@ -1,284 +0,0 @@ -"""Kimi K2.6 GRPO on Modal, disaggregated, native-INT4 end to end. - -Kimi K2.6 is a ~1T-parameter DeepSeek-V3-architecture MoE (MLA + DeepSeek-MoE: -sigmoid router, shared experts, grouped top-k). It is far too large to serve in -bf16/fp8 on 4xB200 (~2TB / ~1TB vs 768GB), so the rollout pool serves the -model's **native compressed-tensors INT4 (W4A16)** checkpoint (~0.5TB, fits). - -That single INT4 format is what makes the disaggregated weight-sync work here. -The bulletin-board path XORs the trainer's exported HF tensors byte-for-byte -against the served base checkpoint, so train and serve must agree on the on-disk -format. We get that with INT4 QAT: - - * the trainer fake-quantizes weights in-loop (OPEN_TRAINING_INT4_FAKE_QAT_FLAG) - so it learns INT4-robust weights; - * its megatron->hf export emits native compressed-tensors W4A16 tensors - (.weight_packed / .weight_scale / .weight_shape, see - megatron_to_hf/processors/quantizer_compressed_tensors.py); - * the sidecar XORs those onto a copy of the native-INT4 base and the pool - reloads it. One format, end to end — no bf16<->FP4 re-quantization step. - -INVARIANT: OPEN_TRAINING_INT4_GROUP_SIZE (the QAT simulation grouping) MUST equal -the served checkpoint's compressed-tensors group_size. The export reads the group -size from the base checkpoint's quantization_config; QAT must simulate the same -grouping or the exported codes won't match what the engine loads. - -Reshape provenance (colocated scripts/low_precision/run-kimi-k2-Thinking-int4.sh --> this disagg config): - ARCHITECTURE -> scripts/models/kimi-k2-thinking.sh via `slime_model_script` - (NO arch attrs here — the script is the single source). - PARALLELISM -> this config, from the recipe's 32x8 TP8/PP8/CP4/EP32. - INFERENCE -> SGLANG_SERVER_ARGS + the dedicated B200 serving image - (serving.py); rollout_num_gpus_per_engine=4 (B200:4). - ALGO/DATA -> this config (GRPO + TIS, dapo-math-17k). - -Deploy as its own app: - EXPERIMENT_CONFIG=kimi_k2_6_int4_disagg m deploy --strategy recreate -m cookbook.slime_disagg.modal_train - -Prerequisites the bring-up depends on (flagged, not yet automated): - 1. hf_checkpoint must resolve to a native compressed-tensors INT4 (W4A16) Kimi - K2.6 checkpoint whose group_size equals INT4_GROUP_SIZE below. Verify the - repo id and its config.json (quant_method == "compressed-tensors"). If no - native-INT4 K2.6 exists yet, fall back by setting hf_checkpoint to - moonshotai/Kimi-K2-Thinking: it is a drop-in swap here — native W4A16 at - group_size 32 (matches INT4_GROUP_SIZE) with the same kimi-k2-thinking.sh - arch and a published slime int4 recipe. - 2. kimi-k2-thinking.sh must describe K2.6's arch (rope scaling, norm eps). If - K2.6 diverges, add scripts/models/kimi-k2.6.sh to the slime fork and point - slime_model_script at it. - 3. The serving image's SGLang fork must serve native-INT4 MLA MoE on Blackwell - (proven for NVFP4; verify on a warm container — see serving.py). - -The smaller `moonlight_int4_disagg` config is the same machinery at a size that -fits a couple of GPUs — run it first to de-risk the INT4-QAT/disk-delta loop. -""" - -from __future__ import annotations - -from cookbook.slime_disagg.configs.base import DATA_PATH, ModalConfig, SlimeConfig - - -APP_NAME = "slime-kimi-k2-6-int4-disagg" -DELTA_VOLUME_NAME = "slime-delta-bulletin-kimi-k2-6-int4" -DELTA_BULLETIN_ROOT = "/delta-bulletin" - -# QAT grouping; MUST match the served INT4 checkpoint's compressed-tensors -# group_size (see the INVARIANT in the module docstring). -INT4_GROUP_SIZE = "32" - -# Async one-step off-policy: in_place applies weights without draining in-flight -# rollouts; stale-version KV is isolated per weight version by the sidecar's -# extra_key stamping and drains as those requests finish. min-version pins cross -# commits freely; only exact pins are quiesced. -SIDECAR_COMMIT_MODE = "in_place" -SIDECAR_DEBUG_REQUESTS = True - - -def build_serving_image(**kwargs): - """Per-experiment rollout-pool image (modal_train picks this up if present). - - Lazy import keeps this config module importable without the modal SDK (it is - resolved locally by launch_train and in unit tests). modal_train passes the - slime ref / root / cache path so the pool pins the trainer's exact slime. - """ - from cookbook.slime_disagg.serving import build_int4_b200_serving_image - - return build_int4_b200_serving_image(**kwargs) - - -# Kimi K2.6 on the elastic B200:4 pool, serving the native compressed-tensors -# INT4 checkpoint. The trainer image injects --served-model-name / --dtype / -# --cuda-graph-max-bs / --max-running-requests / --trust-remote-code, so only the -# Kimi/MLA/cache extras live here. NOTE: no --quantization flag — INT4 is driven -# by the checkpoint's own compressed-tensors config. mem-fraction-static and -# context-length are STARTING POINTS; measure on a warm B200:4 and adjust. -SGLANG_SERVER_ARGS = { - "--tool-call-parser": "kimi_k2", - "--reasoning-parser": "kimi_k2", - "--dist-timeout": "3600", - "--kv-cache-dtype": "fp8_e4m3", - "--attention-backend": "tokenspeed_mla", - "--context-length": "32768", - "--mem-fraction-static": "0.85", - "--chunked-prefill-size": "16384", - "--schedule-conservativeness": "0.5", - "--schedule-policy": "lpm", - # Hierarchical KV cache (tuned in the standalone B200 deployment). - "--enable-hierarchical-cache": "", - "--hicache-ratio": "2", - "--hicache-io-backend": "kernel", - "--hicache-mem-layout": "page_first", - "--hicache-write-policy": "write_through", - "--skip-server-warmup": "", - # Routing replay: the pool must emit per-token routed experts so the trainer - # can replay them. slime launches no engine in publish-only mode, so this is - # set here (DeepSeek-V3-arch MoE supports it). - "--enable-return-routed-experts": "", -} - -# B200 in us-west (Blackwell availability), matching the standalone serve recipe. -# rollout_min_containers is a warm floor only; Modal scales replicas up to meet -# target_inputs (= sglang_server_concurrency), converging on the count that -# saturates the trainer. Raise the floor once that count is known. -modal = ModalConfig( - gpu="B200", - region="us", - rollout_min_containers=2, - proxy_regions=["us-west"], -) - - -class _Slime(SlimeConfig): - # Architecture is sourced from the model script; do NOT inline arch attrs - # (the script carries MLA + the full DeepSeek-MoE arg set). MLA models must - # not set --attention-backend flash, so it is omitted (see the recipe note). - slime_model_script = "scripts/models/kimi-k2-thinking.sh" - - # Model + checkpoint: the native-INT4 base is the served model, the QAT init, - # and the disk-delta base — all the same checkpoint (see prerequisite 1). - hf_checkpoint = "moonshotai/Kimi-K2.6" - ref_load = hf_checkpoint - megatron_to_hf_mode = "bridge" - - # Disaggregated publish-only rollout through slime's opaque HTTP endpoint; - # modal_train fills rollout_endpoint_url from the Flash gateway at launch. - actor_num_nodes = 32 # 32x8 H200 = 256 GPUs (the recipe's trainer footprint) - actor_num_gpus_per_node = 8 - colocate = False - rollout_num_gpus = 0 - rollout_num_gpus_per_engine = 4 # B200:4 per rollout container (native INT4 fits) - rollout_endpoint_url = None - # Staleness gate: each rollout is pinned to min_required_version = latest - lag, - # so a replica more than `lag` versions behind 409s (-> retried, and nudged to - # sync forward) rather than generating too-stale rollouts. - custom_rollout_request_hook_path = "cookbook.slime_disagg.hooks.gated_rollout_request_hook" - rollout_request_weight_version_mode = "min" - rollout_request_weight_version_lag = 1 # bounded staleness window (raise if 409 retries bubble) - rollout_request_retry_attempts = 240 - rollout_request_retry_sleep = 1.0 - rollout_session_affinity_header = "Modal-Session-ID" - - # Async-first: train_async pipelines generate(N+1) with train(N); publish - # every step. For a model this large, raise update_weights_interval if the - # per-step export+delta+reload cost dominates. - async_mode = True - update_weights_interval = 1 - - # Disk-delta publish-only over the Modal Volume bulletin board. The export - # emits native compressed-tensors INT4, so the XOR delta is byte-exact - # against the native-INT4 base. - update_weight_mode = "delta" - update_weight_transport = "disk" - update_weight_delta_encoding = "xor" - update_weight_delta_checksum = "xxh3-128" - update_weight_disk_dir = DELTA_BULLETIN_ROOT - custom_delta_pre_push_path = "cookbook.slime_disagg.hooks.commit_and_wake" - - # Data: dapo-math-17k, the hard math-reasoning set the proven Kimi recipe uses. - prompt_data = f"{DATA_PATH}/dapo-math-17k/dapo-math-17k.jsonl" - input_key = "prompt" - label_key = "label" - apply_chat_template = True - rollout_shuffle = True - rm_type = "math" - eval_interval = None # skip eval during bring-up (see the aime note below) - - # Rollout. Production-leaning throughput knobs from the recipe; num_rollout is - # the bring-up smoke length — scale it up for a real run. - rollout_function_path = "slime.rollout.sglang_rollout.generate_rollout" - num_rollout = 20 - rollout_batch_size = 128 - rollout_max_response_len = 16384 - rollout_temperature = 0.8 - rollout_top_p = 1.0 - n_samples_per_prompt = 8 - over_sampling_batch_size = 256 - dynamic_sampling_filter_path = ( - "slime.rollout.filter_hub.dynamic_sampling_filters.check_reward_nonzero_std" - ) - num_steps_per_rollout = 4 - balance_data = True - # sglang_server_concurrency is THE knob to size for B200:4 saturation: it sets - # the rollout pool's target_inputs (Modal autoscaling) and SGLang's - # cuda-graph-max-bs / max-running-requests. Tune for max throughput per replica. - sglang_server_concurrency = 256 - use_fault_tolerance = False - - # Trainer parallelism from the recipe (world = TP8 * PP8 * CP4 = 256 = 32x8; - # EP32 over the expert region). decoder_last_pipeline_num_layers balances the - # 61-layer model over 8 pipeline stages. Arch flags come from the model script. - tensor_model_parallel_size = 8 - pipeline_model_parallel_size = 8 - context_parallel_size = 4 - expert_model_parallel_size = 32 - expert_tensor_parallel_size = 1 - decoder_last_pipeline_num_layers = 5 - sequence_parallel = True - use_dynamic_batch_size = True - max_tokens_per_gpu = 16384 - recompute_granularity = "full" - recompute_method = "uniform" - recompute_num_layers = 1 - attention_dropout = 0.0 - hidden_dropout = 0.0 - accumulate_allreduce_grads_in_fp32 = True - attention_softmax_in_fp32 = True - no_check_for_nan_in_loss_and_grad = True - - # Optimizer (from the recipe; CPU offload keeps GPU state tiny for ~32B active). - optimizer = "adam" - lr = 1e-6 - lr_decay_style = "constant" - weight_decay = 0.1 - adam_beta1 = 0.9 - adam_beta2 = 0.98 - optimizer_cpu_offload = True - overlap_cpu_optimizer_d2h_h2d = True - use_precision_aware_optimizer = True - - # Algorithm (GRPO + truncated importance sampling, from the recipe). - advantage_estimator = "grpo" - eps_clip = 0.2 - eps_clip_high = 0.28 - use_kl_loss = True - kl_loss_coef = 0.0 - kl_loss_type = "low_var_kl" - entropy_coef = 0.0 - use_tis = True - - # Routing replay: replay the rollout engine's per-token expert routing during - # the training forward/backward to cut train-inference divergence. Needs the - # pool's --enable-return-routed-experts (in SGLANG_SERVER_ARGS above). - use_rollout_routing_replay = True - - # Ray runtime environment: base flags + INT4 QAT. NCCL timeout is generous for - # the large all-to-all expert traffic. - environment = { - "CUDA_DEVICE_MAX_CONNECTIONS": "1", - "NCCL_NVLS_ENABLE": "1", - "NVSHMEM_DISABLE_NCCL": "1", - "NCCL_TIMEOUT_MS": "360000000", - "OPEN_TRAINING_INT4_FAKE_QAT_FLAG": "1", - "OPEN_TRAINING_INT4_GROUP_SIZE": INT4_GROUP_SIZE, - } - - def prepare_data(self) -> None: - from datasets import load_dataset - - # Raw DAPO-Math-17k columns are `prompt` (a chat list) and the gold answer - # under `reward_model.ground_truth` — there is no `label`. slime reads - # --input-key prompt (+ apply_chat_template) and --label-key label, so lift - # the answer into a `label` column. A shuffled subset is ample for a run. - ds = load_dataset("BytedTsinghua-SIA/DAPO-Math-17k", split="train") - ds = ds.shuffle(seed=42).select(range(min(50000, ds.num_rows))) - ds = ds.map(lambda ex: {"label": ex["reward_model"]["ground_truth"]}) - ds = ds.select_columns(["prompt", "label"]) - ds.to_json(f"{DATA_PATH}/dapo-math-17k/dapo-math-17k.jsonl") - - # To add the recipe's aime eval: fetch aime-2024 into {prompt, label} jsonl in - # prepare_data, then set eval_interval (e.g. 20), eval_prompt_data = - # ["aime", f"{DATA_PATH}/aime-2024.jsonl"], n_samples_per_eval_prompt = 16, - # eval_max_response_len = 16384, eval_top_p = 0.7. - - -slime = _Slime() diff --git a/cookbook/slime_disagg/configs/kimi_k2_6_int4_disagg_test.py b/cookbook/slime_disagg/configs/kimi_k2_6_int4_disagg_test.py deleted file mode 100644 index 8541f72..0000000 --- a/cookbook/slime_disagg/configs/kimi_k2_6_int4_disagg_test.py +++ /dev/null @@ -1,79 +0,0 @@ -"""Contract tests for the Kimi K2.6 native-INT4 disaggregated config. - -These assert the invariants that make the INT4-QAT + disk-delta loop coherent — -they are pure data checks and import no modal/slime, so they run anywhere. -""" - -from __future__ import annotations - -import unittest - -from cookbook.slime_disagg.configs import kimi_k2_6_int4_disagg as cfg -from cookbook.slime_disagg.configs.base import SlimeConfig - - -class KimiK26Int4ConfigTest(unittest.TestCase): - def setUp(self) -> None: - self.slime = cfg.slime - self.args = cfg.slime.cli_args() - - def test_is_slime_config_and_cli_args_build(self) -> None: - self.assertIsInstance(self.slime, SlimeConfig) - self.assertTrue(all(isinstance(a, str) for a in self.args)) - - def test_int4_qat_wired(self) -> None: - env = self.slime.environment - self.assertEqual(env.get("OPEN_TRAINING_INT4_FAKE_QAT_FLAG"), "1") - # The QAT simulation grouping MUST match the served checkpoint's - # compressed-tensors group_size (the load-bearing INT4 invariant). - self.assertEqual(env.get("OPEN_TRAINING_INT4_GROUP_SIZE"), cfg.INT4_GROUP_SIZE) - - def test_disaggregated_disk_delta_contract(self) -> None: - self.assertFalse(self.slime.colocate) - self.assertEqual(self.slime.rollout_num_gpus, 0) - self.assertEqual(self.slime.rollout_num_gpus_per_engine, 4) # B200:4 - self.assertTrue(self.slime.async_mode) - self.assertEqual(self.slime.update_weight_mode, "delta") - self.assertEqual(self.slime.update_weight_transport, "disk") - self.assertEqual(self.slime.update_weight_delta_encoding, "xor") - - def test_trainer_footprint_matches_recipe(self) -> None: - # world = TP8 * PP8 * CP4 = 256 = actor_num_nodes(32) * gpus_per_node(8). - self.assertEqual(self.slime.actor_num_nodes, 32) - self.assertEqual(self.slime.actor_num_gpus_per_node, 8) - world = ( - self.slime.tensor_model_parallel_size - * self.slime.pipeline_model_parallel_size - * self.slime.context_parallel_size - ) - self.assertEqual(world, self.slime.actor_num_nodes * self.slime.actor_num_gpus_per_node) - - def test_no_fp4_leakage_and_routing_replay(self) -> None: - keys = cfg.SGLANG_SERVER_ARGS - # INT4 is driven by the checkpoint's compressed-tensors config, never a - # quantization flag — an FP4 leak would mismatch the disk-delta bytes. - self.assertNotIn("--quantization", keys) - self.assertNotIn("modelopt_fp4", " ".join(keys)) - # Routing replay needs the pool to emit per-token routed experts. - self.assertIn("--enable-return-routed-experts", keys) - self.assertTrue(self.slime.use_rollout_routing_replay) - - def test_mla_does_not_force_flash_attention(self) -> None: - # MLA models must not set --attention-backend flash on the trainer. - self.assertNotIn("--attention-backend", self.args) - - def test_launcher_only_fields_excluded_from_cli(self) -> None: - # environment / async_mode / slime_model_script are launcher instructions, - # not SLIME CLI args. - for flag in ("--environment", "--async-mode", "--slime-model-script"): - self.assertNotIn(flag, self.args) - self.assertEqual(self.slime.slime_model_script, "scripts/models/kimi-k2-thinking.sh") - - def test_serving_image_builder_is_lazy_and_callable(self) -> None: - # Present so modal_train deploys the dedicated B200 image; lazy so importing - # this config needs no modal SDK (verified by this test importing cleanly). - self.assertTrue(callable(cfg.build_serving_image)) - - -if __name__ == "__main__": - unittest.main() diff --git a/cookbook/slime_disagg/configs/moonlight.py b/cookbook/slime_disagg/configs/moonlight.py new file mode 100644 index 0000000..4d13fc2 --- /dev/null +++ b/cookbook/slime_disagg/configs/moonlight.py @@ -0,0 +1,136 @@ +"""Moonlight-16B-A3B GRPO on Modal Flash, disaggregated — the cheap rung of the K2.6 ladder. + +Deploy: EXPERIMENT_CONFIG=moonlight m deploy --strategy recreate -m cookbook.slime_disagg.app +""" + +from __future__ import annotations + +from cookbook.common.config import ModalConfig +from cookbook.common.constants import DATA_PATH +from cookbook.slime_disagg.config import SlimeConfig + + +APP_NAME = "stitch-moonlight" +DELTA_VOLUME_NAME = "stitch-delta-moonlight" +DELTA_BULLETIN_ROOT = "/delta-bulletin" +LOCAL_CHECKPOINT_PATH = "/local-checkpoint" + +# in_place applies weights without draining in-flight rollouts; stale KV isolated per version. +SIDECAR_COMMIT_MODE = "in_place" +SIDECAR_DEBUG_REQUESTS = True + +# mem-fraction-static is a STARTING POINT -- measure on a warm container and adjust. +SGLANG_SERVER_ARGS = { + "--weight-loader-prefetch-checkpoints": "", + "--weight-loader-prefetch-num-threads": "8", + "--context-length": "8192", + "--mem-fraction-static": "0.85", + # routing replay: set here since slime launches no engine in publish-only mode. + "--enable-return-routed-experts": "", +} + +modal = ModalConfig(gpu="H200", region="us") + + +class _Slime(SlimeConfig): + # Arch comes from the model script; do NOT inline arch attrs here. + slime_model_script = "scripts/models/moonlight.sh" + + hf_checkpoint = "moonshotai/Moonlight-16B-A3B-Instruct" + ref_load = hf_checkpoint + megatron_to_hf_mode = "bridge" + + actor_num_nodes = 2 # 2x8 H200 -> exercises multinode RDMA + expert parallel + actor_num_gpus_per_node = 8 + colocate = False + rollout_num_gpus = 0 + rollout_num_gpus_per_engine = 1 # 1xH200 per rollout container (MLA -> cheap KV) + rollout_endpoint_url = None + # Staleness gate: pin each request to latest_published - lag; over-stale replicas 409 -> retry. + custom_rollout_request_hook_path = "cookbook.common.hooks.gated_rollout_request_hook" + rollout_request_weight_version_mode = "min" + rollout_request_weight_version_lag = 1 # k: bounded staleness window (tune up if 409 retries bubble) + rollout_request_retry_attempts = 240 + rollout_request_retry_sleep = 1.0 + rollout_session_affinity_header = "Modal-Session-ID" + + # async-first: one-step off-policy; publish weights every step. + async_mode = True + update_weights_interval = 1 + + # disk-delta publish-only (export uses convert_deepseekv3_to_hf for this arch). + update_weight_mode = "delta" + update_weight_transport = "disk" + update_weight_delta_encoding = "xor" + update_weight_delta_checksum = "xxh3-128" + update_weight_disk_dir = DELTA_BULLETIN_ROOT + custom_delta_pre_push_path = "cookbook.common.hooks.commit_and_wake" + + prompt_data = f"{DATA_PATH}/dapo-math-17k/dapo-math-17k.jsonl" + input_key = "prompt" + label_key = "label" + apply_chat_template = True + rollout_shuffle = True + rm_type = "math" + eval_interval = None # skip eval during bring-up + + rollout_function_path = "slime.rollout.sglang_rollout.generate_rollout" + num_rollout = 5 + rollout_batch_size = 64 + rollout_max_response_len = 4096 + rollout_temperature = 1.0 + rollout_top_p = 1.0 + n_samples_per_prompt = 8 + global_batch_size = 128 + sglang_server_concurrency = 64 + use_fault_tolerance = False + + # Trainer parallelism scaled from the recipe's TP4/EP8 to 2x8 (world = TP4*DP4 = 16). + tensor_model_parallel_size = 4 + expert_model_parallel_size = 8 + expert_tensor_parallel_size = 1 + pipeline_model_parallel_size = 1 + context_parallel_size = 1 + sequence_parallel = True + use_dynamic_batch_size = True + max_tokens_per_gpu = 8192 + recompute_granularity = "full" + recompute_method = "uniform" + recompute_num_layers = 1 + attention_dropout = 0.0 + hidden_dropout = 0.0 + accumulate_allreduce_grads_in_fp32 = True + + # Optimizer (from the recipe; CPU offload keeps GPU state tiny for 3B active). + optimizer = "adam" + lr = 1e-6 + lr_decay_style = "constant" + weight_decay = 0.1 + adam_beta1 = 0.9 + adam_beta2 = 0.98 + optimizer_cpu_offload = True + overlap_cpu_optimizer_d2h_h2d = True + use_precision_aware_optimizer = True + + advantage_estimator = "grpo" + eps_clip = 0.2 + eps_clip_high = 0.28 + use_kl_loss = True + kl_loss_coef = 0.0 + kl_loss_type = "low_var_kl" + entropy_coef = 0.0 + + # R3 (arxiv 2510.11370): replay the rollout engine's expert routing in the train forward. + use_rollout_routing_replay = True + + def prepare_data(self) -> None: + from datasets import load_dataset + + ds = load_dataset("BytedTsinghua-SIA/DAPO-Math-17k", split="train") + ds = ds.shuffle(seed=42).select(range(min(50000, ds.num_rows))) + ds = ds.map(lambda ex: {"label": ex["reward_model"]["ground_truth"]}) + ds = ds.select_columns(["prompt", "label"]) + ds.to_json(f"{DATA_PATH}/dapo-math-17k/dapo-math-17k.jsonl") + + +slime = _Slime() diff --git a/cookbook/slime_disagg/configs/moonlight_disagg.py b/cookbook/slime_disagg/configs/moonlight_disagg.py deleted file mode 100644 index 0c747a0..0000000 --- a/cookbook/slime_disagg/configs/moonlight_disagg.py +++ /dev/null @@ -1,200 +0,0 @@ -"""Moonlight-16B-A3B GRPO on Modal Flash, disaggregated (M1: synchronous bring-up). - -Moonlight-16B-A3B is Moonshot's small DeepSeek-V3-architecture MoE -- MLA plus -DeepSeek-MoE (sigmoid router, shared experts, grouped top-k) -- i.e. the Kimi -K2.6 family at a size that fits a single H200 rollout container. This config is -the synchronous bring-up rung; routing replay and async-first layer on top (see -the M2/M3 notes at the bottom). - -Reshape provenance rule (colocated `run-moonlight-16B-A3B.sh` -> disagg config): - ARCHITECTURE -> sourced from scripts/models/moonlight.sh via `slime_model_script` - (NO arch attrs are set here -- the script is the single source). - PARALLELISM -> this config, scaled from the recipe's single-node TP4/EP8. - INFERENCE -> SGLANG_SERVER_ARGS / rollout_num_gpus_per_engine (Modal layer). - ALGO/DATA -> this config. - -Deploy as its own app: - EXPERIMENT_CONFIG=moonlight_disagg m deploy --strategy recreate -m cookbook.slime_disagg.modal_train - -Prerequisites the bring-up depends on (flagged, not yet automated): - 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. - 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). -""" - -from __future__ import annotations - -from cookbook.slime_disagg.configs.base import DATA_PATH, ModalConfig, SlimeConfig - - -APP_NAME = "slime-moonlight-disagg" -DELTA_VOLUME_NAME = "slime-delta-bulletin-moonlight" -DELTA_BULLETIN_ROOT = "/delta-bulletin" - -# 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.) -SIDECAR_COMMIT_MODE = "in_place" -SIDECAR_DEBUG_REQUESTS = True - -# Moonlight serving on the elastic pool. MLA's compressed KV is tiny, so one -# H200 holds the ~32 GB bf16 weights plus a large KV pool. mem-fraction-static -# is a STARTING POINT -- measure it on a warm container and adjust (the only -# datapoint from the colocated recipe is 0.7 at TP8, a different topology). -SGLANG_SERVER_ARGS = { - "--context-length": "8192", - "--mem-fraction-static": "0.85", - # M2 routing replay: the pool must emit per-token routed experts. slime - # launches no engine in publish-only mode, so this is set here (not by - # sglang_engine.py). num_layers/moe_router_topk come from moonlight.sh, so - # the rollout's [tokens, 27, 6] reshape matches the served model (M0-verified). - "--enable-return-routed-experts": "", -} - -modal = ModalConfig(gpu="H200", region="us") - - -class _Slime(SlimeConfig): - # Architecture is sourced from the model script; do NOT inline arch attrs - # (the script carries MLA + the full DeepSeek-MoE arg set). - slime_model_script = "scripts/models/moonlight.sh" - - # Model + checkpoint. Bridge-mode HF load, same as the dense disagg example - # (see prerequisite 2 in the module docstring -- this is the main M1 unknown). - hf_checkpoint = "moonshotai/Moonlight-16B-A3B-Instruct" - ref_load = hf_checkpoint - megatron_to_hf_mode = "bridge" - - # Disaggregated publish-only rollout through slime's opaque HTTP endpoint; - # modal_train fills rollout_endpoint_url from the Flash gateway at launch. - actor_num_nodes = 2 # 2x8 H200 -> exercises multinode RDMA + expert parallel - actor_num_gpus_per_node = 8 - colocate = False - rollout_num_gpus = 0 - rollout_num_gpus_per_engine = 1 # 1xH200 per rollout container (MLA -> cheap KV) - rollout_endpoint_url = None - # M3 staleness gate: each rollout request is pinned to - # min_required_version = latest_published - lag (derived out-of-band from the - # bulletin `latest`, since the per-request hook gets no rollout_id). A replica - # more than `lag` versions behind 409s -> retried, so no too-stale rollouts. - custom_rollout_request_hook_path = "cookbook.slime_disagg.hooks.gated_rollout_request_hook" - rollout_request_weight_version_mode = "min" - rollout_request_weight_version_lag = 1 # k: bounded staleness window (tune up if 409 retries bubble) - rollout_request_retry_attempts = 240 - rollout_request_retry_sleep = 1.0 - rollout_session_affinity_header = "Modal-Session-ID" - - # M3 async-first: one-step off-policy (train_async pipelines generate(N+1) with - # train(N)); publish weights every step. - async_mode = True - update_weights_interval = 1 - - # Disk-delta publish-only over the Modal Volume bulletin board (export uses - # convert_deepseekv3_to_hf for this arch). - update_weight_mode = "delta" - update_weight_transport = "disk" - update_weight_delta_encoding = "xor" - update_weight_delta_checksum = "xxh3-128" - update_weight_disk_dir = DELTA_BULLETIN_ROOT - custom_delta_pre_push_path = "cookbook.slime_disagg.hooks.commit_and_wake" - - # Data: dapo-math-17k, the hard math-reasoning set the proven Moonlight recipe - # uses (MoE-worthy, unlike GSM8K). The convincing demo adds aime eval. - prompt_data = f"{DATA_PATH}/dapo-math-17k/dapo-math-17k.jsonl" - input_key = "prompt" - label_key = "label" - apply_chat_template = True - rollout_shuffle = True - rm_type = "math" - eval_interval = None # skip eval during bring-up - - # Rollout. Small num_rollout for the bring-up smoke; scale up for the demo. - rollout_function_path = "slime.rollout.sglang_rollout.generate_rollout" - num_rollout = 5 - rollout_batch_size = 64 - rollout_max_response_len = 4096 - rollout_temperature = 1.0 - rollout_top_p = 1.0 - n_samples_per_prompt = 8 - global_batch_size = 128 - sglang_server_concurrency = 64 - use_fault_tolerance = False - - # Trainer parallelism: scaled from the recipe's single-node TP4/EP8 to 2x8 - # (world = TP4 * DP4 = 16; EP8 over the expert region). MLA / MoE arch flags - # come from the model script; the bring-up keeps the recipe's alltoall - # dispatcher (the recipe's flex + --moe-enable-deepep is a later perf lever - # and needs DeepEP in the image). - tensor_model_parallel_size = 4 - expert_model_parallel_size = 8 - expert_tensor_parallel_size = 1 - pipeline_model_parallel_size = 1 - context_parallel_size = 1 - sequence_parallel = True - use_dynamic_batch_size = True - max_tokens_per_gpu = 8192 - recompute_granularity = "full" - recompute_method = "uniform" - recompute_num_layers = 1 - attention_dropout = 0.0 - hidden_dropout = 0.0 - accumulate_allreduce_grads_in_fp32 = True - - # Optimizer (from the recipe; CPU offload keeps GPU state tiny for 3B active). - optimizer = "adam" - lr = 1e-6 - lr_decay_style = "constant" - weight_decay = 0.1 - adam_beta1 = 0.9 - adam_beta2 = 0.98 - optimizer_cpu_offload = True - overlap_cpu_optimizer_d2h_h2d = True - use_precision_aware_optimizer = True - - # Algorithm (GRPO). - advantage_estimator = "grpo" - eps_clip = 0.2 - eps_clip_high = 0.28 - use_kl_loss = True - kl_loss_coef = 0.0 - kl_loss_type = "low_var_kl" - entropy_coef = 0.0 - - # Routing replay (M2): replay the rollout engine's per-token expert routing - # during the training forward/backward to cut train-inference divergence - # (R3, arxiv 2510.11370). Auto-implies use_routing_replay; needs the pool's - # --enable-return-routed-experts (in SGLANG_SERVER_ARGS above). - use_rollout_routing_replay = True - - def prepare_data(self) -> None: - from datasets import load_dataset - - # Raw DAPO-Math-17k columns are `prompt` (a chat list) and the gold answer - # under `reward_model.ground_truth` — there is no `label`. slime reads - # --input-key prompt (+ apply_chat_template) and --label-key label, so lift - # the answer into a `label` column. The HF artifact is ~1.8M rows; a shuffled - # subset is ample for bring-up and a short reward-climb. - ds = load_dataset("BytedTsinghua-SIA/DAPO-Math-17k", split="train") - ds = ds.shuffle(seed=42).select(range(min(50000, ds.num_rows))) - ds = ds.map(lambda ex: {"label": ex["reward_model"]["ground_truth"]}) - ds = ds.select_columns(["prompt", "label"]) - ds.to_json(f"{DATA_PATH}/dapo-math-17k/dapo-math-17k.jsonl") - - -slime = _Slime() - -# Milestone progression (kept here so it's legible): -# M1 (sync MoE bring-up): DONE. -# M2 (routing replay): DONE — use_rollout_routing_replay=True + -# --enable-return-routed-experts; num_layers(27)/moe_router_topk(6) from moonlight.sh. -# M3 (async-first): DONE above — async_mode=True (train_async one-step off-policy), -# SIDECAR_COMMIT_MODE="in_place" (in-flight updates + extra_key KV namespacing), -# min-version gate at latest-lag via gated_rollout_request_hook. Tune -# rollout_request_weight_version_lag for the staleness/throughput trade-off. diff --git a/cookbook/slime_disagg/configs/moonlight_int4_disagg.py b/cookbook/slime_disagg/configs/moonlight_int4.py similarity index 63% rename from cookbook/slime_disagg/configs/moonlight_int4_disagg.py rename to cookbook/slime_disagg/configs/moonlight_int4.py index aa148f8..c2ba96e 100644 --- a/cookbook/slime_disagg/configs/moonlight_int4_disagg.py +++ b/cookbook/slime_disagg/configs/moonlight_int4.py @@ -1,42 +1,31 @@ -"""Moonlight-16B-A3B GRPO on Modal, disaggregated, native-INT4 end to end. - -The cheap de-risk rung for `kimi_k2_6_int4_disagg`. Moonlight-16B-A3B is the same -DeepSeek-V3 architecture (MLA + DeepSeek-MoE) and the same native compressed- -tensors INT4 (W4A16) serving path as Kimi K2.6, at a size that fits a single -H200 — so it exercises the *exact* INT4-QAT + disk-delta + native-INT4-serve loop -without the B200 / Blackwell-SGLang-fork variable (it reuses the trainer image's -SGLang, which already serves INT4 in the colocated recipe). Get this green first; -then `kimi_k2_6_int4_disagg` changes only scale and the serving image. - -INVARIANT (as in the Kimi config): OPEN_TRAINING_INT4_GROUP_SIZE MUST equal the -served checkpoint's compressed-tensors group_size — 128 here, per the colocated -scripts/low_precision/run-moonlight-16B-A3B-int4.sh recipe. - -Deploy as its own app: - EXPERIMENT_CONFIG=moonlight_int4_disagg m deploy --strategy recreate -m cookbook.slime_disagg.modal_train - -Prerequisite: hf_checkpoint must resolve to a native compressed-tensors INT4 -(W4A16) Moonlight checkpoint with group_size 128 (the recipe quantizes -Moonlight-16B-A3B-Instruct to INT4 offline). Verify the repo id / config.json. +"""Moonlight-16B-A3B GRPO on Modal, disaggregated, native-INT4 end to end — the cheap de-risk for kimi_k2_6_int4. + +INVARIANT: OPEN_TRAINING_INT4_GROUP_SIZE MUST equal the served checkpoint's compressed-tensors group_size (128 here). + +Deploy: EXPERIMENT_CONFIG=moonlight_int4 m deploy --strategy recreate -m cookbook.slime_disagg.app """ from __future__ import annotations -from cookbook.slime_disagg.configs.base import DATA_PATH, ModalConfig, SlimeConfig +from cookbook.common.config import ModalConfig +from cookbook.common.constants import DATA_PATH +from cookbook.slime_disagg.config import SlimeConfig -APP_NAME = "slime-moonlight-int4-disagg" -DELTA_VOLUME_NAME = "slime-delta-bulletin-moonlight-int4" +APP_NAME = "stitch-moonlight-int4" +DELTA_VOLUME_NAME = "stitch-delta-moonlight-int4" DELTA_BULLETIN_ROOT = "/delta-bulletin" +LOCAL_CHECKPOINT_PATH = "/local-checkpoint" INT4_GROUP_SIZE = "128" SIDECAR_COMMIT_MODE = "in_place" SIDECAR_DEBUG_REQUESTS = True -# No build_serving_image: the pool reuses the trainer image, whose SGLang serves -# this native-INT4 checkpoint on a single H200 (no Blackwell fork needed). +# The pool reuses the trainer image (its SGLang serves native INT4; no Blackwell fork). SGLANG_SERVER_ARGS = { + "--weight-loader-prefetch-checkpoints": "", + "--weight-loader-prefetch-num-threads": "8", "--context-length": "16384", "--mem-fraction-static": "0.85", "--enable-return-routed-experts": "", # routing replay @@ -46,24 +35,21 @@ class _Slime(SlimeConfig): - # Architecture is sourced from the model script (MLA + DeepSeek-MoE). MLA - # models must not set --attention-backend flash, so it is omitted. + # Arch comes from the model script. MLA models must not set --attention-backend flash. slime_model_script = "scripts/models/moonlight.sh" - # The native-INT4 base is the served model, the QAT init, and the disk-delta - # base — all the same checkpoint (see prerequisite). + # The native-INT4 base is the served model, the QAT init, and the disk-delta base. hf_checkpoint = "moonshotai/Moonlight-16B-A3B-Instruct-INT4" ref_load = hf_checkpoint megatron_to_hf_mode = "bridge" - # Disaggregated publish-only rollout; modal_train fills rollout_endpoint_url. actor_num_nodes = 1 # 1x8 H200 (the recipe trains on 1x4; 1 node here) actor_num_gpus_per_node = 8 colocate = False rollout_num_gpus = 0 rollout_num_gpus_per_engine = 1 # 1xH200 per rollout container (16B INT4 fits) rollout_endpoint_url = None - custom_rollout_request_hook_path = "cookbook.slime_disagg.hooks.gated_rollout_request_hook" + custom_rollout_request_hook_path = "cookbook.common.hooks.gated_rollout_request_hook" rollout_request_weight_version_mode = "min" rollout_request_weight_version_lag = 1 rollout_request_retry_attempts = 240 @@ -73,14 +59,13 @@ class _Slime(SlimeConfig): async_mode = True update_weights_interval = 1 - # Disk-delta publish-only: the export emits native compressed-tensors INT4, so - # the XOR delta is byte-exact against the native-INT4 base. + # disk-delta: export emits native INT4, so the XOR delta is byte-exact against the base. update_weight_mode = "delta" update_weight_transport = "disk" update_weight_delta_encoding = "xor" update_weight_delta_checksum = "xxh3-128" update_weight_disk_dir = DELTA_BULLETIN_ROOT - custom_delta_pre_push_path = "cookbook.slime_disagg.hooks.commit_and_wake" + custom_delta_pre_push_path = "cookbook.common.hooks.commit_and_wake" prompt_data = f"{DATA_PATH}/dapo-math-17k/dapo-math-17k.jsonl" input_key = "prompt" @@ -156,8 +141,6 @@ class _Slime(SlimeConfig): def prepare_data(self) -> None: from datasets import load_dataset - # DAPO-Math-17k: lift reward_model.ground_truth into a `label` column (see - # the Kimi config for the rationale). A shuffled subset is ample. ds = load_dataset("BytedTsinghua-SIA/DAPO-Math-17k", split="train") ds = ds.shuffle(seed=42).select(range(min(50000, ds.num_rows))) ds = ds.map(lambda ex: {"label": ex["reward_model"]["ground_truth"]}) diff --git a/cookbook/slime_disagg/configs/qwen3_4b_delta_flash.py b/cookbook/slime_disagg/configs/qwen3_4b_delta_flash.py index eae2097..536fde5 100644 --- a/cookbook/slime_disagg/configs/qwen3_4b_delta_flash.py +++ b/cookbook/slime_disagg/configs/qwen3_4b_delta_flash.py @@ -2,27 +2,25 @@ from __future__ import annotations -from cookbook.slime_disagg.configs.base import DATA_PATH, ModalConfig, SlimeConfig +from cookbook.common.config import ModalConfig +from cookbook.common.constants import DATA_PATH +from cookbook.slime_disagg.config import SlimeConfig -APP_NAME = "slime-qwen3-4b-delta-flash" -DELTA_VOLUME_NAME = "slime-delta-bulletin-qwen3-4b" +APP_NAME = "stitch-qwen3-4b" +DELTA_VOLUME_NAME = "stitch-delta-qwen3-4b" DELTA_BULLETIN_ROOT = "/delta-bulletin" +LOCAL_CHECKPOINT_PATH = "/local-checkpoint" -# How the rollout sidecar applies published weight versions. "in_place" pauses -# 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. +# in_place applies weights without draining; stale KV isolated per version via extra_key. SIDECAR_COMMIT_MODE = "in_place" -# Log every versioned sidecar proxy request (start/end + injected rid) at INFO, -# so a stuck rollout can be traced hop-by-hop: slime rid -> sidecar -> SGLang. +# Log versioned sidecar proxy requests at INFO to trace a stuck rollout hop-by-hop. SIDECAR_DEBUG_REQUESTS = True -# SGLang server tuning, merged over the structural args set in modal_train.py. SGLANG_SERVER_ARGS = { + "--weight-loader-prefetch-checkpoints": "", + "--weight-loader-prefetch-num-threads": "8", "--reasoning-parser": "qwen3", "--context-length": "16384", "--mem-fraction-static": "0.84", @@ -34,42 +32,32 @@ class _Slime(SlimeConfig): - # Model hf_checkpoint = "Qwen/Qwen3-4B" ref_load = hf_checkpoint megatron_to_hf_mode = "bridge" - # Modal Flash disaggregated rollout through slime's generic HTTP endpoint mode. actor_num_nodes = 1 actor_num_gpus_per_node = 8 colocate = False rollout_num_gpus = 0 rollout_num_gpus_per_engine = 1 - # Publish-only: slime launches no engines and routes /generate to the Modal - # Flash gateway (filled in at launch); rollouts pull weights from the pool. - rollout_endpoint_url = None - custom_rollout_request_hook_path = "stitch.trainers.slime.rollout_request_weight_version_hook" + rollout_endpoint_url = None # publish-only: slime routes /generate to the Flash gateway + custom_rollout_request_hook_path = "cookbook.common.hooks.gated_rollout_request_hook" rollout_request_weight_version_mode = "exact" rollout_request_weight_version_lag = 0 rollout_request_retry_attempts = 240 rollout_request_retry_sleep = 1.0 - # The trainer hits the Modal Flash gateway directly, which routes session - # affinity on Modal-Session-ID; emit that so GRPO siblings co-locate. + # session affinity so GRPO siblings co-locate on one Flash replica. rollout_session_affinity_header = "Modal-Session-ID" - # Disk-delta publish-only over the Modal Volume bulletin board: slime writes - # weight_v{N}/ + a `latest` pointer to update_weight_disk_dir (the Volume), - # the pre-push hook commits the Volume + wakes the pool, and each sidecar - # applies the delta host-side via slime's disk_delta. Publish-only is implied - # by rollout_endpoint_url, so no local-checkpoint dir is required here. + # disk-delta publish-only: slime writes weight_v{N}/ + `latest`; the hook commits and wakes the pool. update_weight_mode = "delta" update_weight_transport = "disk" update_weight_delta_encoding = "xor" update_weight_delta_checksum = "xxh3-128" update_weight_disk_dir = DELTA_BULLETIN_ROOT - custom_delta_pre_push_path = "cookbook.slime_disagg.hooks.commit_and_wake" + custom_delta_pre_push_path = "cookbook.common.hooks.commit_and_wake" - # Data prompt_data = f"{DATA_PATH}/gsm8k/train.parquet" eval_prompt_data = ["gsm8k", f"{DATA_PATH}/gsm8k/test.parquet"] input_key = "messages" @@ -78,7 +66,6 @@ class _Slime(SlimeConfig): rollout_shuffle = True rm_type = "math" - # Rollout rollout_function_path = "slime.rollout.sglang_rollout.generate_rollout" num_rollout = 3 rollout_batch_size = 64 @@ -90,13 +77,11 @@ class _Slime(SlimeConfig): sglang_server_concurrency = 64 use_fault_tolerance = False - # Eval eval_interval = None n_samples_per_eval_prompt = 4 eval_max_response_len = 8192 eval_top_p = 1.0 - # Training tensor_model_parallel_size = 1 sequence_parallel = False use_dynamic_batch_size = True @@ -109,7 +94,6 @@ class _Slime(SlimeConfig): accumulate_allreduce_grads_in_fp32 = True attention_softmax_in_fp32 = True - # Optimizer optimizer = "adam" lr = 1e-6 lr_decay_style = "constant" @@ -117,7 +101,6 @@ class _Slime(SlimeConfig): adam_beta1 = 0.9 adam_beta2 = 0.98 - # Algorithm advantage_estimator = "grpo" eps_clip = 0.2 eps_clip_high = 0.28 @@ -126,7 +109,6 @@ class _Slime(SlimeConfig): kl_loss_type = "low_var_kl" entropy_coef = 0.0 - # Qwen3-4B architecture num_layers = 36 hidden_size = 2560 ffn_hidden_size = 9728 diff --git a/cookbook/slime_disagg/configs/qwen3_4b_delta_flash_hillclimb.py b/cookbook/slime_disagg/configs/qwen3_4b_delta_flash_hillclimb.py index 0780e1a..33c4ba6 100644 --- a/cookbook/slime_disagg/configs/qwen3_4b_delta_flash_hillclimb.py +++ b/cookbook/slime_disagg/configs/qwen3_4b_delta_flash_hillclimb.py @@ -5,9 +5,10 @@ from cookbook.slime_disagg.configs import qwen3_4b_delta_flash as _base -APP_NAME = "slime-qwen3-4b-delta-flash-hillclimb" -DELTA_VOLUME_NAME = "slime-delta-bulletin-qwen3-4b-hillclimb" +APP_NAME = "stitch-qwen3-4b-hillclimb" +DELTA_VOLUME_NAME = "stitch-delta-qwen3-4b-hillclimb" DELTA_BULLETIN_ROOT = _base.DELTA_BULLETIN_ROOT +LOCAL_CHECKPOINT_PATH = "/local-checkpoint" SIDECAR_COMMIT_MODE = _base.SIDECAR_COMMIT_MODE SGLANG_SERVER_ARGS = _base.SGLANG_SERVER_ARGS @@ -15,13 +16,11 @@ class _Slime(_base._Slime): - # Longer than the protocol smoke, but still bounded enough to use as a - # repeatable disaggregated reward-hillclimb check. + # Longer than the protocol smoke but still bounded — a repeatable reward-hillclimb check. num_rollout = 120 eval_interval = 20 log_passrate = True - # update_weight_disk_dir is inherited from the base config (same Volume mount - # path); only DELTA_VOLUME_NAME differs, so this app owns its own bulletin. + # update_weight_disk_dir is inherited from the base; only DELTA_VOLUME_NAME differs. slime = _Slime() diff --git a/cookbook/slime_disagg/helpers.py b/cookbook/slime_disagg/helpers.py deleted file mode 100644 index 6295ad7..0000000 --- a/cookbook/slime_disagg/helpers.py +++ /dev/null @@ -1,180 +0,0 @@ -"""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. -""" - -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 - -# Re-export shared helpers so existing callers (modal_train.py) don't break. -from cookbook.ray_cluster import ( # noqa: F401 - RAY_START_TIMEOUT, - RAY_WORKER_JOIN_TIMEOUT, - get_modal_cluster_context, - start_ray_head, - start_ray_worker, - training_nodes, -) -from cookbook.sidecar_process import ( # noqa: F401 - terminate_process, - wait_http, -) - - -SIDECAR_MODULE = "cookbook.slime_disagg.sidecar" - - -def start_sglang_sidecar( - *, - sidecar_port: int, - sglang_port: int, - bulletin_root: str, - local_checkpoint_dir: str, - base_checkpoint_dir: str, - volume_name: str, - commit_mode: str, - debug_requests: bool = False, -) -> subprocess.Popen: - from cookbook.sidecar_process import start_sglang_sidecar as _start - - return _start( - sidecar_module=SIDECAR_MODULE, - sidecar_port=sidecar_port, - sglang_port=sglang_port, - bulletin_root=bulletin_root, - local_checkpoint_dir=local_checkpoint_dir, - base_checkpoint_dir=base_checkpoint_dir, - volume_name=volume_name, - commit_mode=commit_mode, - debug_requests=debug_requests, - ) - - -# ── 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) - - -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.""" - - -def smoke_flash_pool( - *, - app_name: str, - cls_name: str, - model_name: str, - weight_version: int, - 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"}, - ) - with urllib.request.urlopen(request, timeout=timeout) as resp: - return json.load(resp) diff --git a/cookbook/slime_disagg/hooks.py b/cookbook/slime_disagg/hooks.py deleted file mode 100644 index 93ea67f..0000000 --- a/cookbook/slime_disagg/hooks.py +++ /dev/null @@ -1,29 +0,0 @@ -"""Modal publish + rollout-gating hooks for the slime_disagg example. - -Thin wrappers around :mod:`cookbook.bulletin_hooks` with the slime-specific -env-var fallbacks for the Flash app / server class name. -""" - -from __future__ import annotations - -from typing import Any - -from cookbook.bulletin_hooks import ( - commit_and_wake as _commit_and_wake, - gated_rollout_request_hook, -) - - -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", - ) - - -# Re-export for the trainer's custom_rollout_request_hook_path. -__all__ = ["commit_and_wake", "gated_rollout_request_hook"] diff --git a/cookbook/slime_disagg/hooks_test.py b/cookbook/slime_disagg/hooks_test.py deleted file mode 100644 index b0d2f81..0000000 --- a/cookbook/slime_disagg/hooks_test.py +++ /dev/null @@ -1,103 +0,0 @@ -from __future__ import annotations - -import asyncio -import os -import tempfile -import unittest -from argparse import Namespace -from pathlib import Path -from unittest import mock - -import cookbook.bulletin_hooks as bulletin_hooks -from cookbook.slime_disagg import hooks - - -class CommitAndWakeTest(unittest.TestCase): - def test_advances_run_pointer_commits_and_wakes(self) -> None: - with tempfile.TemporaryDirectory() as tmp: - root = Path(tmp) # transport root (Volume mount) - disk_dir = root / "run-a" # update_weight_disk_dir = / - version_dir = disk_dir / "weight_v000001" - version_dir.mkdir(parents=True) - args = Namespace( - update_weight_disk_dir=str(disk_dir), - run_id="run-a", - rollout_modal_flash_app_name="app", - rollout_modal_flash_server_cls_name="Server", - ) - - with mock.patch.dict(os.environ, {"DELTA_VOLUME_NAME": "delta-volume"}), mock.patch.object( - bulletin_hooks, "distributed_rank", return_value=0 - ), mock.patch.object(bulletin_hooks, "commit_volume") as commit_volume, mock.patch.object( - bulletin_hooks, "discover_flash_targets", return_value=["https://c"] - ), mock.patch.object(bulletin_hooks, "wake_targets") as wake_targets: - hooks.commit_and_wake(args, str(version_dir), []) - - # The canonical pointer lives at the transport root and is self-identifying. - self.assertEqual((root / "latest").read_text(encoding="utf-8"), "run-a/weight_v000001") - commit_volume.assert_called_once_with("delta-volume") - wake_targets.assert_called_once_with(["https://c"], 1) - - def test_baseline_commits_but_does_not_advance_or_wake(self) -> None: - with tempfile.TemporaryDirectory() as tmp: - root = Path(tmp) - disk_dir = root / "run-a" - disk_dir.mkdir(parents=True) - # The baseline call passes the disk-dir root, not a weight_v{N} dir. - args = Namespace(update_weight_disk_dir=str(disk_dir), run_id="run-a") - with mock.patch.dict(os.environ, {"DELTA_VOLUME_NAME": "delta-volume"}), mock.patch.object( - bulletin_hooks, "distributed_rank", return_value=0 - ), mock.patch.object(bulletin_hooks, "commit_volume") as commit_volume, mock.patch.object( - bulletin_hooks, "wake_targets" - ) as wake_targets: - hooks.commit_and_wake(args, str(disk_dir), []) - - self.assertFalse((root / "latest").exists()) - commit_volume.assert_called_once_with("delta-volume") - wake_targets.assert_not_called() - - def test_non_rank_zero_commits_only(self) -> None: - with tempfile.TemporaryDirectory() as tmp: - root = Path(tmp) - disk_dir = root / "run-a" - (disk_dir / "weight_v000002").mkdir(parents=True) - args = Namespace(update_weight_disk_dir=str(disk_dir), run_id="run-a") - with mock.patch.dict(os.environ, {"DELTA_VOLUME_NAME": "delta-volume"}), mock.patch.object( - bulletin_hooks, "distributed_rank", return_value=3 - ), mock.patch.object(bulletin_hooks, "commit_volume") as commit_volume, mock.patch.object( - bulletin_hooks, "wake_targets" - ) as wake_targets: - hooks.commit_and_wake(args, str(disk_dir / "weight_v000002"), []) - - self.assertFalse((root / "latest").exists()) # only rank 0 writes latest - commit_volume.assert_called_once_with("delta-volume") # every rank commits shards - wake_targets.assert_not_called() - - -class GateTest(unittest.TestCase): - def test_gate_reads_version_from_run_pointer(self) -> None: - async def run() -> None: - with tempfile.TemporaryDirectory() as tmp: - root = Path(tmp) - disk_dir = root / "run-a" - disk_dir.mkdir(parents=True) - # Canonical pointer at the transport root carries the run identity; - # the staleness gate cares only about the version within the run. - (root / "latest").write_text("run-a/weight_v000005", encoding="utf-8") - args = Namespace(update_weight_disk_dir=str(disk_dir), run_id="run-a") - # Reset the module-level cache before each test. - cache = bulletin_hooks._latest_cache - cache.version = 0 - cache.run_id = None - cache._refreshed_at = -1e9 - cache._board = None - with mock.patch.dict(os.environ, {}, clear=True): - version = await cache.get(args) - self.assertEqual(version, 5) - self.assertEqual(cache.run_id, "run-a") - - asyncio.run(run()) - - -if __name__ == "__main__": - unittest.main() diff --git a/cookbook/slime_disagg/modal_train.py b/cookbook/slime_disagg/modal_train.py deleted file mode 100644 index b4db9b5..0000000 --- a/cookbook/slime_disagg/modal_train.py +++ /dev/null @@ -1,445 +0,0 @@ -"""Disaggregated SLIME training on Modal. - -A Modal Flash pool of SGLang servers handles rollouts; a clustered Trainer -runs SLIME on Ray and publishes sparse weight deltas through a Modal Volume -bulletin board that the rollout servers sync from. - -Run all commands as modules from the repo root, e.g.: - - uv run --extra modal modal deploy -m cookbook.slime_disagg.modal_train -""" - -from __future__ import annotations - -import importlib -import os -import subprocess -import tempfile -import uuid -from pathlib import Path - -import modal -import modal.experimental - -from cookbook.slime_disagg import helpers -from cookbook.slime_disagg.configs.base import ( - CHECKPOINTS_PATH, - DATA_PATH, - HF_CACHE_PATH, - SlimeConfig, -) -from stitch.providers.modal import resolve_flash_gateway_url - -# The one deploy-time knob: which experiment config this app is built around. -# `modal deploy` takes no function arguments, so the selection has to come from -# the environment; the image records the same value so containers reconstruct -# the identical app. Everything else is configured in the experiment modules. -EXPERIMENT = os.environ.get("EXPERIMENT_CONFIG", "qwen3_4b_delta_flash") -exp = importlib.import_module(f"cookbook.slime_disagg.configs.{EXPERIMENT}") - -modal_cfg = exp.modal -slime_cfg = exp.slime - -APP_NAME = exp.APP_NAME -MODEL_NAME = slime_cfg.hf_checkpoint -ROLLOUT_CONCURRENCY = slime_cfg.sglang_server_concurrency -N_TRAIN_NODES = helpers.training_nodes(slime_cfg) - -MINUTES = 60 -SIDECAR_PORT = 8000 -SGLANG_PORT = 8001 -RAY_PORT = 6379 -SERVER_STARTUP_TIMEOUT = 35 * MINUTES -# Ephemeral host-local full HF checkpoint the sidecar patches in place per delta -# (seeded once from the base; rebuilt from base on a cold container). -LOCAL_CHECKPOINT_PATH = "/local-checkpoint" - -SLIME_IMAGE_TAG = "slimerl/slime:nightly-dev-20260527a" -SLIME_ROOT = "/root/slime" -# Fork branch with the generic HTTP rollout endpoint and publish-only -# disk-delta hooks that this example drives. -SLIME_REPO_URL = "https://github.com/modal-projects/slime.git" -# Pin to an exact commit, not the branch tip: the build's `git fetch ... && -# checkout` is a cached image layer, so a moving branch tip silently leaves the -# container on a stale slime. This is PR #5 head (disaggregated-rollout): disk- -# delta publish-only + rollout_endpoint_url + custom_rollout_request_hook_path. -# Bump this SHA to roll slime forward. -SLIME_REPO_REF = "ebfe153949b1a69c39e92f947ed5d475166dd724" # incl. deepseekv3 router-dtype export fix + per-request rollout hook - -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}") - # Replace the bundled slime with the fork branch. - .run_commands( - f"rm -rf {SLIME_ROOT}" - 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}" - ) - # The base image installs megatron-core as a PEP 660 *strict* editable that - # exposes only `megatron.core`, hiding `megatron.training` (which slime's - # megatron backend imports) even though the full tree exists on disk at - # /root/Megatron-LM. Reinstall in compat editable mode so a .pth puts the - # whole source tree on the path and `megatron.training` is importable. - .run_commands( - "cd /root/Megatron-LM" - " && python3 -m pip install --no-deps -e . --config-settings editable_mode=compat" - ) - .pip_install( - "autoinference-utils==0.2.0", # SGLang server lifecycle for the rollout pool - "fastapi", # stitch sidecar - "httpx", # stitch sidecar - "uvicorn", # stitch sidecar - # The sidecar applies disk deltas host-side via slime.utils.disk_delta, - # which decompresses with zstandard and checksums with xxhash (xxh3-128 - # default) or blake3. slime is installed --no-deps, so add them here. - "zstandard", - "xxhash", - "blake3", - ) - .env( - { - "EXPERIMENT_CONFIG": EXPERIMENT, - "HF_XET_HIGH_PERFORMANCE": "1", - "HF_HUB_ENABLE_HF_TRANSFER": "1", - } - ) - # 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). - .add_local_python_source("stitch") - .add_local_dir( - Path(__file__).parent, - remote_path="/root/cookbook/slime_disagg", - ignore=["**/__pycache__"], - ) -) - -# Dev iteration: SLIME_LOCAL_DIR overlays a local slime checkout onto the image's -# cloned fork (installed editable at /root/slime), so fork edits take effect on -# container start with no image rebuild or push. Unset by default, so the committed -# example always builds from the pinned SLIME_REPO_REF. -if slime_local := os.environ.get("SLIME_LOCAL_DIR"): - image = image.add_local_dir( - slime_local, - remote_path=SLIME_ROOT, - ignore=[".git", "**/__pycache__", "**/*.pyc"], - ) - - -# The rollout pool may need a different serving stack than the trainer — e.g. a -# Blackwell SGLang build that serves native-INT4 Kimi K2.6, which the slime -# trainer image does not provide. An experiment opts in by defining -# build_serving_image(...); otherwise the pool reuses the trainer image (the -# Qwen/Moonlight bf16 and Moonlight-INT4 experiments do). Either way the pool -# pins the trainer's exact slime ref, so the sidecar's disk_delta decoder matches -# the trainer's delta encoder. -def _select_server_image() -> modal.Image: - builder = getattr(exp, "build_serving_image", None) - if builder is None: - return image - return builder( - slime_repo_url=SLIME_REPO_URL, - slime_repo_ref=SLIME_REPO_REF, - slime_root=SLIME_ROOT, - hf_cache_path=str(HF_CACHE_PATH), - experiment=EXPERIMENT, - ) - - -server_image = _select_server_image() -if slime_local and server_image is not image: - # Mirror the trainer's dev-iteration overlay onto a dedicated serving image. - server_image = server_image.add_local_dir( - slime_local, - remote_path=SLIME_ROOT, - ignore=[".git", "**/__pycache__", "**/*.pyc"], - ) - -with server_image.imports(): - from autoinference_utils.endpoint import SGLangEndpoint, warmup_chat_completions - - -hf_cache_volume = modal.Volume.from_name("huggingface-cache", create_if_missing=True) -data_volume = modal.Volume.from_name("slime-data", create_if_missing=True) -checkpoints_volume = modal.Volume.from_name("slime-checkpoints", create_if_missing=True) -delta_volume = modal.Volume.from_name( - exp.DELTA_VOLUME_NAME, create_if_missing=True, version=2 -) - -train_volumes = { - str(HF_CACHE_PATH): hf_cache_volume, - str(DATA_PATH): data_volume, - str(CHECKPOINTS_PATH): checkpoints_volume, - exp.DELTA_BULLETIN_ROOT: delta_volume, -} - -app = modal.App(APP_NAME) - -SGLANG_SERVER_ARGS = { - "--served-model-name": MODEL_NAME, - "--dtype": "bfloat16", - "--cuda-graph-max-bs": str(ROLLOUT_CONCURRENCY), - "--max-running-requests": str(ROLLOUT_CONCURRENCY), - "--trust-remote-code": "", - # The disk-delta branch applies deltas host-side and reloads via the ordinary - # update_weights_from_disk path, so the old engine-side delta server args - # (--update-weight-delta-chunk-bytes/--update-weight-delta-read-workers) no - # longer exist and must not be passed. - **exp.SGLANG_SERVER_ARGS, -} - -WARMUP_PAYLOAD = { - "model": MODEL_NAME, - "messages": [{"role": "user", "content": "Reply with exactly OK."}], - "max_tokens": 8, - "temperature": 0, - "chat_template_kwargs": {"enable_thinking": False}, -} - - -@app.cls( - image=server_image, - gpu=f"{modal_cfg.gpu}:{slime_cfg.rollout_num_gpus_per_engine}", - cloud=modal_cfg.cloud, - region=modal_cfg.region, - volumes={ - str(HF_CACHE_PATH): hf_cache_volume, - exp.DELTA_BULLETIN_ROOT: delta_volume, - }, - min_containers=modal_cfg.rollout_min_containers, - timeout=40 * MINUTES, - scaledown_window=15 * MINUTES, - include_source=False, -) -@modal.experimental.http_server( - port=SIDECAR_PORT, - proxy_regions=modal_cfg.proxy_regions, - exit_grace_period=25, - startup_timeout=SERVER_STARTUP_TIMEOUT, -) -@modal.concurrent(target_inputs=ROLLOUT_CONCURRENCY) -class Server: - """One SGLang rollout server plus the stitch weight-sync sidecar. - - The sidecar proxies rollout traffic, reloads the delta Volume, and applies - published weight versions so requests pinned to a version are served by - matching weights. - """ - - @modal.enter() - def startup(self) -> None: - self.endpoint = SGLangEndpoint( - model_path=MODEL_NAME, - worker_port=SGLANG_PORT, - tp=slime_cfg.rollout_num_gpus_per_engine, - extra_server_args=SGLANG_SERVER_ARGS, - health_timeout=SERVER_STARTUP_TIMEOUT, - health_poll_interval=10.0, - ) - self.endpoint.start() - warmup_chat_completions( - port=SGLANG_PORT, - payload=WARMUP_PAYLOAD, - successful_requests=2, - request_timeout=120.0, - max_attempts_per_request=3, - ) - # Deltas are applied host-side onto a copy of the base checkpoint; the - # base resolves to the same HF cache snapshot the SGLang server loaded. - from huggingface_hub import snapshot_download - - base_checkpoint_dir = snapshot_download(MODEL_NAME, local_files_only=True) - self.sidecar = helpers.start_sglang_sidecar( - sidecar_port=SIDECAR_PORT, - sglang_port=SGLANG_PORT, - bulletin_root=exp.DELTA_BULLETIN_ROOT, - local_checkpoint_dir=LOCAL_CHECKPOINT_PATH, - base_checkpoint_dir=base_checkpoint_dir, - volume_name=exp.DELTA_VOLUME_NAME, - commit_mode=exp.SIDECAR_COMMIT_MODE, - debug_requests=getattr(exp, "SIDECAR_DEBUG_REQUESTS", False), - ) - helpers.wait_http( - f"http://127.0.0.1:{SIDECAR_PORT}/health", - self.sidecar, - SERVER_STARTUP_TIMEOUT, - ) - print( - f"Rollout server ready: model={MODEL_NAME}, target_inputs={ROLLOUT_CONCURRENCY}" - ) - - @modal.exit() - def stop(self) -> None: - helpers.terminate_process(getattr(self, "sidecar", None)) - if hasattr(self, "endpoint"): - self.endpoint.stop() - - -@app.cls( - image=image, - gpu=f"{modal_cfg.gpu}:{slime_cfg.actor_num_gpus_per_node}", - memory=modal_cfg.memory, - cloud=modal_cfg.cloud, - region=modal_cfg.region, - volumes=train_volumes, - timeout=24 * 60 * MINUTES, - startup_timeout=20 * MINUTES, - scaledown_window=30 * MINUTES, - experimental_options={"efa_enabled": True}, - include_source=False, -) -@modal.experimental.clustered(N_TRAIN_NODES, rdma=True) -class Trainer: - """SLIME actor cluster. The Ray cluster comes up once per container in - enter(), so back-to-back training runs reuse it instead of rebuilding it.""" - - @modal.enter() - def start_ray(self) -> None: - rank, master_addr, my_ip = helpers.get_modal_cluster_context(N_TRAIN_NODES) - self.rank = rank - # Ray actors inherit the raylet's environment, so everything the - # training processes need must be exported before `ray start`. - os.environ.update( - { - "SLIME_HOST_IP": my_ip, - "SGLANG_HOST_IP": my_ip, - "HOST_IP": my_ip, - "MASTER_ADDR": master_addr, - "RAY_ADDRESS": f"{master_addr}:{RAY_PORT}", - "no_proxy": f"127.0.0.1,{master_addr},{my_ip}", - "NO_PROXY": f"127.0.0.1,{master_addr},{my_ip}", - **slime_cfg.environment, - } - ) - if rank == 0: - helpers.start_ray_head(my_ip, N_TRAIN_NODES, ray_port=RAY_PORT) - else: - helpers.start_ray_worker(my_ip, master_addr, ray_port=RAY_PORT) - - @modal.method() - def train(self, experiment: str, payload: dict) -> None: - """Run one training job from a SlimeConfig payload (see to_payload()). - - The config arrives as data instead of a module name, so launch_train - can resolve experiments from the local working tree and new or edited - configs run without a redeploy. - """ - for volume in train_volumes.values(): - volume.reload() - # Rank 0 drives the run; the cluster stays up until its call returns, - # so the other ranks only need their Ray workers, started in enter(). - if self.rank != 0: - return - - cfg = SlimeConfig.from_payload(payload) - if helpers.training_nodes(cfg) != N_TRAIN_NODES: - raise ValueError( - f"experiment {experiment!r} needs {helpers.training_nodes(cfg)} node(s) but this app " - f"was deployed with {N_TRAIN_NODES}; deploy it as its own app with EXPERIMENT_CONFIG={experiment}" - ) - if cfg.environment != slime_cfg.environment: - # Ray inherited the deploy-time environment when enter() started it. - print( - f"WARNING: experiment {experiment!r} changes `environment`, which only " - f"takes effect after a redeploy restarts the Ray cluster." - ) - - cfg.rollout_endpoint_url = resolve_flash_gateway_url(APP_NAME, Server.__name__) - # Fresh run id per launch: slime writes this run's chain under a partition - # (//weight_v{N}/), while the canonical pointer at - # /latest is self-identifying (/weight_vN). So a new - # run never overwrites a finished run's version dirs and its pointer move is - # a forward step, not a colliding rewind — no manual bulletin reset needed. - run_id = uuid.uuid4().hex[:12] - cfg.update_weight_disk_dir = f"{exp.DELTA_BULLETIN_ROOT}/{run_id}" - # stitch's publish hooks read these off the slime args namespace. - cfg.custom_config_path = { - "update_weight_delta_volume_name": exp.DELTA_VOLUME_NAME, - "rollout_modal_flash_app_name": APP_NAME, - "rollout_modal_flash_server_cls_name": Server.__name__, - "run_id": run_id, - } - helpers.prepare_slime_config(cfg, tempfile.mkdtemp()) - cmd = helpers.build_train_cmd(cfg, SLIME_ROOT) - - print( - f"Training {experiment}: nodes={N_TRAIN_NODES}, rollout_endpoint={cfg.rollout_endpoint_url}" - ) - print(f"Command: {cmd}") - subprocess.run(["bash", "-lc", cmd], check=True) - - -@app.function( - image=image, - volumes={str(HF_CACHE_PATH): hf_cache_volume}, - timeout=2 * 60 * MINUTES, - secrets=[modal.Secret.from_name("huggingface-secret")], - include_source=False, -) -def download_model() -> None: - from huggingface_hub import snapshot_download - - snapshot_download(repo_id=MODEL_NAME) - hf_cache_volume.commit() - - -@app.function( - image=image, - volumes={str(DATA_PATH): data_volume}, - timeout=2 * 60 * MINUTES, - secrets=[modal.Secret.from_name("huggingface-secret")], - include_source=False, -) -def prepare_dataset() -> None: - data_volume.reload() - slime_cfg.prepare_data() - data_volume.commit() - - -@app.local_entrypoint() -def launch_train(experiment: str = EXPERIMENT) -> None: - """Resolve an experiment from the local working tree and spawn it on the - deployed app. Training args ship as data, so new or edited configs run - without a redeploy; infrastructure changes (GPU, nodes, pool size, - Volume names) still require one.""" - from modal.exception import NotFoundError - - run = importlib.import_module(f"cookbook.slime_disagg.configs.{experiment}") - if run.DELTA_VOLUME_NAME != exp.DELTA_VOLUME_NAME: - raise SystemExit( - f"Experiment {experiment!r} owns Volume {run.DELTA_VOLUME_NAME!r}, but app {APP_NAME!r} " - f"mounts {exp.DELTA_VOLUME_NAME!r}. Deploy it as its own app with EXPERIMENT_CONFIG={experiment}." - ) - - try: - trainer = modal.Cls.from_name(APP_NAME, Trainer.__name__)() - call = trainer.train.spawn(experiment, run.slime.to_payload()) - except NotFoundError: - raise SystemExit( - f"App {APP_NAME!r} is not deployed. Run:\n" - f" uv run --extra modal modal deploy -m cookbook.slime_disagg.modal_train" - ) - print(f"Spawned train({experiment!r}) on {APP_NAME}: {call.object_id}") - - -@app.local_entrypoint() -def smoke_flash_pool( - weight_version: int = 0, timeout_seconds: int = 30 * MINUTES -) -> None: - """Check that the deployed Flash pool serves completions at the expected - weight version, via the gateway and each container directly.""" - helpers.smoke_flash_pool( - app_name=APP_NAME, - cls_name=Server.__name__, - model_name=MODEL_NAME, - weight_version=weight_version, - expect_min_containers=modal_cfg.rollout_min_containers, - timeout_seconds=timeout_seconds, - ) diff --git a/cookbook/slime_disagg/serving.py b/cookbook/slime_disagg/serving.py deleted file mode 100644 index 4607c94..0000000 --- a/cookbook/slime_disagg/serving.py +++ /dev/null @@ -1,120 +0,0 @@ -"""Dedicated B200 native-INT4 SGLang serving image for the 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. -""" - -from __future__ import annotations - -from pathlib import Path - -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", -} - - -def build_int4_b200_serving_image( - *, - slime_repo_url: str, - slime_repo_ref: str, - slime_root: str, - 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__"], - ) - ) diff --git a/cookbook/slime_disagg/sidecar.py b/cookbook/slime_disagg/sidecar.py deleted file mode 100644 index 4a17d3a..0000000 --- a/cookbook/slime_disagg/sidecar.py +++ /dev/null @@ -1,139 +0,0 @@ -"""SGLang weight-sync sidecar launcher 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``). -""" - -from __future__ import annotations - -import argparse -import logging -import os - -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, - ) - - -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. Relies on the engine's overlap-drain " - "fix. 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", - ) - - -if __name__ == "__main__": - main() diff --git a/cookbook/slime_disagg/trainer_image.py b/cookbook/slime_disagg/trainer_image.py new file mode 100644 index 0000000..a9dbc7c --- /dev/null +++ b/cookbook/slime_disagg/trainer_image.py @@ -0,0 +1,47 @@ +"""The slime trainer image + the versions pinned to launch a slime run. + +The serving half is separate and shared (common/serving_image.py) — the pool installs no +trainer package, so slime and miles serve on the identical weight-sync sglang image. +""" + +from __future__ import annotations + +from pathlib import Path + +import modal + +SLIME_IMAGE_TAG = "slimerl/slime:nightly-dev-20260527a" +SLIME_REPO_URL = "https://github.com/modal-projects/slime.git" +# Pin to an exact commit, not the branch tip (the fetch+checkout is a cached layer). +SLIME_REPO_REF = "11bb0fa48aa37d5c54fe297143c6bc1d40f311bf" +SLIME_ROOT = "/root/slime" + +_COOKBOOK_DIR = Path(__file__).resolve().parent.parent # .../cookbook + + +def build_trainer_image(*, hf_cache_path: str, experiment: str, slime_local: str | None = None) -> modal.Image: + """The slime trainer image: the pinned slime fork over the slime base, Megatron-LM + reinstalled so ``megatron.training`` is importable, + the delta encoder's codecs. + stitch + the cookbook package are mounted for the trainer, Ray actors, and sidecar.""" + image = ( + modal.Image.from_registry(SLIME_IMAGE_TAG) + .entrypoint([]) + .run_commands(f"rm -rf {hf_cache_path}") # baked HF cache must not shadow the mounted volume + .run_commands( + f"rm -rf {SLIME_ROOT}" + f" && git clone --depth 1 {SLIME_REPO_URL} {SLIME_ROOT}" + f" && cd {SLIME_ROOT} && git fetch --depth 1 origin {SLIME_REPO_REF} && git checkout FETCH_HEAD" + f" && python3 -m pip install --no-deps -e {SLIME_ROOT}" + ) + # The base installs megatron-core as a strict editable that hides + # megatron.training; reinstall in compat mode so the whole source tree is importable. + .run_commands("cd /root/Megatron-LM && python3 -m pip install --no-deps -e . --config-settings editable_mode=compat") + # The trainer-side delta ENCODER (slime.utils.disk_delta) needs the codecs even under --no-deps. + .pip_install("fastapi", "httpx", "uvicorn", "zstandard", "xxhash", "blake3") + .env({"HF_XET_HIGH_PERFORMANCE": "1", "HF_HUB_ENABLE_HF_TRANSFER": "1", "EXPERIMENT_CONFIG": experiment}) + .add_local_python_source("stitch") + .add_local_dir(str(_COOKBOOK_DIR), remote_path="/root/cookbook", ignore=["**/__pycache__"]) + ) + if slime_local: # dev overlay: replace the cloned fork with a local checkout (no rebuild) + image = image.add_local_dir(slime_local, remote_path=SLIME_ROOT, ignore=[".git", "**/__pycache__", "**/*.pyc"]) + return image diff --git a/cookbook/standalone_rollouts/README.md b/cookbook/standalone_rollouts/README.md deleted file mode 100644 index c805a6b..0000000 --- a/cookbook/standalone_rollouts/README.md +++ /dev/null @@ -1,218 +0,0 @@ -# Standalone SGLang Rollout Provider - -This cookbook deploys a standalone Modal Flash pool of SGLang rollout servers -that implements the customer hot-load API. External trainers upload checkpoints -or deltas to S3, call the provider hot-load endpoint, poll readiness, then send -rollout traffic to the same provider URL. - -It is a **log-as-truth** design: a durable, monotonic `latest` pointer in the S3 -transport is the source of truth, and the elastic pool reconciles to it by pull. - -1. Modal starts one SGLang server + a stitch weight-sync sidecar per warm - container. Each sidecar is a `WeightSyncManager` that reconciles its engine - to `latest` (on startup, on a wake, and on a periodic poll). -2. The **front door** — a singleton `App.server` (`min_containers=1`, pinned to - the same routing region as the pool) and the only writer of `latest` — serves - the customer API. `POST /hot_load/...` advances - `latest` (monotonic CAS; a rewind is rejected) and best-effort wakes the - pool. The pool pulls the new `weight_v{N}/`, applies the disk delta host-side - (slime `disk_delta`: chain-replay + per-tensor checksum), and reloads SGLang. -3. `GET /hot_load/...` reports readiness by enumerating the **live** containers - and querying each `/server_info` — no self-reported replica state, so a - scaled-down replica can't haunt the readiness fraction. -4. Inference (`/generate`, `/v1/chat/completions`, `/v1/completions`, …) is - proxied to the SGLang gateway. - -## Layout - -| File | What it is | -|---|---| -| `modal_serve.py` | Standalone Modal rollout-provider app; owns the Modal `App`, the Server pool, and the singleton front door | -| `frontdoor.py` | Front-door hot-load adapter logic (advance `latest`, live-readiness, proxy) — injected I/O, unit-tested | -| `provider.py` | Per-container sidecar: a `WeightSyncManager` over a slime-layout board on the transport | -| `configs/moonlight_hot_load.py` | Moonlight-16B-A3B provider config (DeepSeek-V3-arch MoE; routing replay) | -| `slime/` | Optional SLIME integration-test harness for this provider | - -## Compatibility Notes - -The provider targets slime's `disk-delta-weight-sync` branch + PR #5. Each -version is a canonical HF/SafeTensors directory `weight_v{N}/` with a -`model.safetensors.index.json`; the engine applies deltas **host-side** (slime -`disk_delta`) onto a local full checkpoint and then reloads through the ordinary -`update_weights_from_disk` path — there is no engine-side `load_format="delta"` -receiver. The delta format is XOR (or `overwrite`) encoding, zstd compression, -and xxh3-128 (or blake3/adler32) per-tensor checksums; the version's -`index.json` carries `delta_encoding`/`compression_format`/`checksum_format`. - -A customer-produced delta (XOR + adler32 + zstd) is directly applicable by this -applier. The only adapter work is metadata *location*: the customer sends -`compression_format`/`checksum_format`/`previous_snapshot_identity` in the POST -body and ships a weight-map-only `index.json`, whereas the applier reads those -from the index's `metadata` block — so a customer-facing front door normalizes -the POST metadata into the index before advancing `latest`. - -Session affinity is delegated to Modal's Flash gateway. External clients send -the neutral `x-session-affinity` header to the **front door** (the advertised -provider URL, an `App.server` in front of the pool); the front door rewrites it to -`Modal-Session-ID` *before* the gateway, which then consistently routes related -requests to the same replica. The rewrite must happen pre-gateway, so it lives -in the front door rather than the per-container sidecar. Requests without the -header are routed normally. - -## Deploy the Provider - -Work from the repo root. The provider expects: - -- `huggingface-secret` for downloading the base model into the HF cache Volume. -- `stitch-api-shim-provider` for optional hot-load API auth. - -The default config mounts the S3 bucket `modal-stitch-s3-transport` at -`/mnt/stitch-s3-transport` with this Modal OIDC role: - -```text -arn:aws:iam::459781239556:role/modal-buckets/stitch-s3-transport-role -``` - -The mounted bucket prefix is `standalone-rollouts/moonlight/`, so the external -S3 location for uploaded snapshots is: - -```text -s3://modal-stitch-s3-transport/standalone-rollouts/moonlight// -``` - -Override `STITCH_SHIM_S3_BUCKET_NAME`, `STITCH_SHIM_S3_KEY_PREFIX`, -`STITCH_SHIM_S3_REGION`, or `STITCH_SHIM_S3_OIDC_AUTH_ROLE_ARN` at deploy time -if you create another bucket or prefix. `STITCH_SHIM_S3_REGION` is optional, -but useful when Modal's S3 Mountpoint cannot auto-detect the bucket region. - -Create the secret: - -```bash -uv run --extra modal modal secret create stitch-api-shim-provider \ - STITCH_SHIM_API_KEY=... \ - STITCH_SHIM_PROVIDER_MODEL=moonlight \ - STITCH_SHIM_PROVIDER_DEPLOYMENT=rollout-prod -``` - -The S3 OIDC role must allow `s3:PutObject` on the prefix: the singleton front -door writes the `latest` pointer there. - -Deploy: - -```bash -alias m="uv run --extra modal modal" - -m run -m cookbook.standalone_rollouts.modal_serve::download_model -m deploy -m cookbook.standalone_rollouts.modal_serve -m run -m cookbook.standalone_rollouts.modal_serve::print_url -# Authenticated smoke from inside Modal (reads the provider secret; the API key -# never leaves Modal): polls GET /hot_load readiness + a base completion. -m run -m cookbook.standalone_rollouts.modal_serve::check -``` - -The default app is `stitch-moonlight-api-shim`. To create a separate deployment, -add a config under `configs/` and deploy with: - -```bash -PROVIDER_CONFIG=my_provider_config m deploy -m cookbook.standalone_rollouts.modal_serve -``` - -## External Trainer Contract - -Upload each snapshot to: - -```text -s3://modal-stitch-s3-transport/standalone-rollouts/moonlight// -``` - -Signal it: - -```bash -curl -X POST "$GATEWAY/hot_load/v1/models/hot_load" \ - -H "Authorization: Bearer $STITCH_SHIM_API_KEY" \ - -H "Provider-Model: moonlight" \ - -H "Provider-Deployment: rollout-prod" \ - -H "Content-Type: application/json" \ - -d '{"identity":"weight_v000001"}' -``` - -For a compatible SGLang/SLIME delta: - -```bash -curl -X POST "$GATEWAY/hot_load/v1/models/hot_load" \ - -H "Authorization: Bearer $STITCH_SHIM_API_KEY" \ - -H "Provider-Model: moonlight" \ - -H "Provider-Deployment: rollout-prod" \ - -H "Content-Type: application/json" \ - -d '{ - "identity": "weight_v000002", - "incremental_snapshot_metadata": { - "previous_snapshot_identity": "weight_v000001", - "compression_format": "zstd", - "checksum_format": "xxh3-128" - }, - "reset_prompt_cache": "new_session" - }' -``` - -Poll readiness: - -```bash -curl "$GATEWAY/hot_load/v1/models/hot_load" \ - -H "Authorization: Bearer $STITCH_SHIM_API_KEY" \ - -H "Provider-Model: moonlight" \ - -H "Provider-Deployment: rollout-prod" -``` - -A replica counts as ready only when `readiness` is true and -`current_snapshot_identity` matches the requested identity. - -## Optional SLIME Test Harness - -The provider app above is standalone and can be used by any external trainer -that follows the same S3 + hot-load contract. For end-to-end testing, this -cookbook also includes `slime/`, a Modal-hosted SLIME trainer harness that -copies sparse deltas into the mounted S3 transport and calls the deployed -provider. - -The provider owns the Modal app. Deploying `modal_serve.py` publishes only the -SGLang rollout provider. Deploying `slime/modal_train.py` imports that same -app and adds the trainer functions, so the resulting Modal app contains both -the standalone provider and the SLIME integration-test trainer. - -The trainer reuses `stitch-api-shim-provider` for optional shim API auth. It -derives the provider URL from the deployed Flash `Server` class, and the S3 -transport uses the same Modal `CloudBucketMount` OIDC role as the provider. - -Then redeploy the same Modal app with the trainer functions included and launch -the trainer: - -```bash -m run -m cookbook.standalone_rollouts.modal_serve::download_model -m run -m cookbook.standalone_rollouts.slime.modal_train::prepare_dataset -m deploy -m cookbook.standalone_rollouts.slime.modal_train -m run -m cookbook.standalone_rollouts.slime.modal_train::launch_train -``` - -The trainer (`slime/configs/moonlight_slime_trainer.py`) runs **async-first** -(`train_async`, one-step off-policy) with **rollout routing replay** enabled, in -SLIME publish-only mode: it launches no rollout engines, routes `/generate` to -the provider front door, and writes each `weight_v{N}/` to a local disk dir. - -Staleness is gated by a **publish-hook readiness barrier**, not a per-request -pin (this is the key difference from `slime_disagg`, which uses a min-version -request gate). The `custom_delta_pre_push_path` hook (`announce_and_wait`) copies -the new version to the S3 transport, POSTs the customer hot-load API, and blocks -until the front door reports **every live replica** ready on the new version -(`readiness_threshold` 1.0) — so the next rollouts always run on current weights. -The request hook therefore leaves the version gate off -(`api_shim_rollout_request_weight_version_mode="none"`) and only carries auth -headers, retries, and session affinity. - -Routing replay requires the pool to return per-token routed experts: the provider -config sets `--enable-return-routed-experts` and the trainer sets -`use_rollout_routing_replay=True` (the layer count / router top-k come from the -sourced `scripts/models/moonlight.sh`). - -The first weight update only seeds SLIME's baseline snapshot and publishes -nothing; subsequent updates publish `weight_v000001`, `weight_v000002`, … . diff --git a/cookbook/standalone_rollouts/__init__.py b/cookbook/standalone_rollouts/__init__.py deleted file mode 100644 index e99b1ed..0000000 --- a/cookbook/standalone_rollouts/__init__.py +++ /dev/null @@ -1 +0,0 @@ -"""Standalone rollout-provider cookbook.""" diff --git a/cookbook/standalone_rollouts/configs/__init__.py b/cookbook/standalone_rollouts/configs/__init__.py deleted file mode 100644 index 652ec63..0000000 --- a/cookbook/standalone_rollouts/configs/__init__.py +++ /dev/null @@ -1 +0,0 @@ -"""Provider configs for the standalone rollout cookbook.""" diff --git a/cookbook/standalone_rollouts/configs/moonlight_hot_load.py b/cookbook/standalone_rollouts/configs/moonlight_hot_load.py deleted file mode 100644 index 2faf7e4..0000000 --- a/cookbook/standalone_rollouts/configs/moonlight_hot_load.py +++ /dev/null @@ -1,56 +0,0 @@ -"""Moonlight-16B-A3B rollout provider config for the hot-load API shim. - -Moonlight is a small DeepSeek-V3-architecture (MLA + DeepSeek-MoE) model — the -Kimi K2.6 family at a size that fits one H200. The pool serves rollouts and emits -per-token routed experts so the trainer can replay them (routing replay). -""" - -from __future__ import annotations - -from pathlib import Path - - -APP_NAME = "stitch-moonlight-api-shim" -MODEL_NAME = "moonshotai/Moonlight-16B-A3B-Instruct" - -HF_SECRET_NAME = "huggingface-secret" -SHIM_SECRET_NAME = "stitch-api-shim-provider" -HF_CACHE_VOLUME_NAME = "huggingface-cache" - -HF_CACHE_PATH = Path("/root/.cache/huggingface") -# Ephemeral host-local full HF checkpoint the sidecar patches in place per delta -# (seeded from the base; rebuilt on a cold container). -LOCAL_CHECKPOINT_PATH = "/local-checkpoint" -# in_place pauses/applies/continues without flushing; stale-version KV is isolated -# by the sidecar's extra_key stamping and drains as its in-flight requests finish. -COMMIT_MODE = "in_place" -S3_TRANSPORT_BUCKET_NAME = "modal-stitch-s3-transport" -S3_TRANSPORT_KEY_PREFIX = "standalone-rollouts/moonlight" -S3_TRANSPORT_MOUNT_PATH = Path("/mnt/stitch-s3-transport") -S3_TRANSPORT_REGION = None -S3_TRANSPORT_OIDC_AUTH_ROLE_ARN = ( - "arn:aws:iam::459781239556:role/modal-buckets/stitch-s3-transport-role" -) - -GPU = "H200" -CLOUD = None -# Pin the rollout pool (and the front door) to the US so they stay co-located. -REGION = "us" -# Region inputs are routed through. The front door's `routing_region` and the -# rollout pool's Flash `proxy_regions` are kept identical so customer traffic and -# the pool share one entry region. -ROUTING_REGION = "us-east" -PROXY_REGIONS = [ROUTING_REGION] -ROLLOUT_MIN_CONTAINERS = 4 -ROLLOUT_NUM_GPUS_PER_ENGINE = 1 # Moonlight (~32 GB bf16, tiny MLA KV) fits 1xH200 -ROLLOUT_CONCURRENCY = 32 - -# Moonlight is DeepSeek-V3-arch (no qwen reasoning parser). Routing replay needs -# the pool to emit per-token routed experts: the trainer launches no engine in -# publish-only mode, so --enable-return-routed-experts must be set here, not by -# slime's sglang_engine. mem-fraction-static is a starting point — measure it. -SGLANG_SERVER_ARGS = { - "--context-length": "8192", - "--mem-fraction-static": "0.85", - "--enable-return-routed-experts": "", -} diff --git a/cookbook/standalone_rollouts/frontdoor.py b/cookbook/standalone_rollouts/frontdoor.py deleted file mode 100644 index 299180d..0000000 --- a/cookbook/standalone_rollouts/frontdoor.py +++ /dev/null @@ -1,198 +0,0 @@ -"""Front-door hot-load adapter for the standalone rollout provider. - -The front door is the single public entry and the single writer of the -bulletin board's monotonic ``latest`` pointer. It implements the customer's -hot-load API as a thin projection of the canonical log-as-truth design: - -- ``POST /hot_load/...`` advances ``latest`` to the signalled checkpoint - (monotonic CAS — a rewind is rejected for now; rolling the fleet back from a - recovery anchor is the future story), then best-effort wakes the pool. The - elastic rollout pool reconciles to ``latest`` on its own. -- ``GET /hot_load/...`` reports pool readiness by enumerating the *live* - containers and querying each ``/server_info`` — no self-reported replica - state, so a scaled-down container can't haunt the readiness fraction. -- Everything else is proxied to the rollout gateway. - -All I/O (reading/writing ``latest``, enumerating replicas, proxying, auth) is -injected so the adapter logic is testable without Modal; ``modal_serve.py`` -supplies the real implementations. - -No ``from __future__ import annotations`` here: the route handlers' -``request: Request`` annotation must evaluate eagerly against the factory-local -fastapi import, or FastAPI mistakes ``request`` for a query parameter. -""" - -import asyncio -from collections.abc import Awaitable, Callable -from typing import Any - -from stitch.protocol import ( - RolloutPoolState, - RolloutReplicaState, - format_snapshot_identity, - parse_weight_identity, -) - - -HOT_LOAD_PATH = "/hot_load/v1/models/hot_load" - - -def advance_latest_decision( - current_run_id: str | None, - current_version: int, - identity: str, - request_run_id: str | None, -) -> dict[str, Any]: - """Decide whether a hot-load signal should advance the ``latest`` pointer. - - Run-id partitioning makes this idempotent across runs. A signal from a - *different* run (``request_run_id != current_run_id``) is a fresh chain whose - version space restarts at 1, so it is accepted and resets the version space - (``"reset": True``) — accepting v1 after a finished run reached v5 is the - intended begin-a-new-run path, not a rewind. Within the same run (and the - run-less customer layout where both are None) the monotonic CAS still applies. - - Returns ``{"run_id": str|None, "version": int, "reset": bool}`` to accept, or - ``{"error": {...}}`` to reject (``InvalidIdentity`` / ``WeightRewindRejected``). - """ - version = parse_weight_identity(identity) - if version is None: - return { - "error": { - "type": "InvalidIdentity", - "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: - 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), - } - } - return {"run_id": request_run_id, "version": int(version), "reset": False} - - -def pool_state_from_server_infos(infos: list[dict[str, Any]]) -> RolloutPoolState: - """Build a pool-readiness report from live ``/server_info`` responses. - - A replica is ready when it is reachable and idle (not mid-sync, no sticky - sync error). The trainer separately matches ``current_snapshot_identity`` - against its target, so an idle replica still on an old version is observable - but correctly not counted toward the target. - """ - replicas: list[RolloutReplicaState] = [] - for info in infos: - current_version = info.get("current_version") - current_run_id = info.get("current_run_id") - sync_state = info.get("sync_state") - last_error = info.get("last_sync_error") - ready = sync_state == "IDLE" and not last_error - # The snapshot identity carries the run, so a replica still serving a - # finished run's same-numbered version isn't counted ready for a new run. - identity = ( - format_snapshot_identity(current_run_id, current_version) - if isinstance(current_version, int) and current_version >= 0 - else None - ) - replicas.append( - RolloutReplicaState( - readiness=ready, - current_version=current_version if isinstance(current_version, int) else None, - current_snapshot_identity=identity, - replica_id=info.get("run_id") or info.get("replica_id"), - sync_state=sync_state, - readiness_reason=None if ready else (last_error or sync_state or "unreachable"), - ) - ) - return RolloutPoolState(replicas=replicas) - - -def create_frontdoor_app( - *, - read_current_pointer: Callable[[], Awaitable[tuple[str | None, int]]], - advance_to: Callable[[str | None, int], Awaitable[None]], - list_server_infos: Callable[[], Awaitable[list[dict[str, Any]]]], - proxy: Callable[..., Awaitable[Any]], - authorize: Callable[[Any], Any] | None = None, - wake: Callable[[int], Awaitable[None]] | None = None, -): - """Build the front-door FastAPI app from injected I/O. - - ``read_current_pointer`` returns the active ``(run_id, version)``; - ``advance_to(run_id, version)`` atomically writes the single self-identifying - ``latest`` pointer. ``list_server_infos`` enumerates live replicas; ``proxy`` - forwards non-hot-load requests; ``authorize`` returns a rejection Response or - ``None``; ``wake`` is a best-effort post-advance nudge. - - The read-decide-advance sequence is serialized so the singleton front door - never races itself into a transient rewind (two POSTs both reading the same - current version and advancing out of order). - """ - from fastapi import FastAPI, Request - from fastapi.responses import JSONResponse, Response - - app = FastAPI() - advance_lock = asyncio.Lock() - - def _auth(request: Request): - return authorize(request.headers) if authorize is not None else None - - @app.post(HOT_LOAD_PATH, response_model=None) - async def post_hot_load(request: Request) -> Response: - rejected = _auth(request) - if rejected is not None: - return rejected - payload = await request.json() - if not isinstance(payload, dict) or not payload.get("identity"): - return JSONResponse({"error": "body.identity is required"}, status_code=400) - identity = str(payload["identity"]) - request_run_id = payload.get("run_id") - async with advance_lock: - current_run_id, current_version = await read_current_pointer() - decision = advance_latest_decision( - current_run_id, current_version, identity, request_run_id - ) - if "error" in decision: - status = 400 if decision["error"]["type"] == "InvalidIdentity" else 409 - return JSONResponse(decision, status_code=status) - await advance_to(decision["run_id"], decision["version"]) - snapshot_identity = format_snapshot_identity(decision["run_id"], decision["version"]) - if wake is not None: - try: - await wake(decision["version"]) - except Exception: # noqa: BLE001 — wake is a latency optimization only - pass - return JSONResponse( - { - "accepted": True, - "identity": identity, - "current_snapshot_identity": snapshot_identity, - "run_id": decision["run_id"], - } - ) - - @app.get(HOT_LOAD_PATH, response_model=None) - async def get_hot_load(request: Request) -> Response: - rejected = _auth(request) - if rejected is not None: - return rejected - infos = await list_server_infos() - return JSONResponse(pool_state_from_server_infos(infos).to_dict()) - - @app.api_route("/{path:path}", methods=["GET", "POST", "PUT", "DELETE", "PATCH"]) - async def catch_all(path: str, request: Request) -> Response: - rejected = _auth(request) - if rejected is not None: - return rejected - return await proxy(request, path) - - return app diff --git a/cookbook/standalone_rollouts/frontdoor_test.py b/cookbook/standalone_rollouts/frontdoor_test.py deleted file mode 100644 index 439199f..0000000 --- a/cookbook/standalone_rollouts/frontdoor_test.py +++ /dev/null @@ -1,176 +0,0 @@ -from __future__ import annotations - -import unittest - -from cookbook.standalone_rollouts.frontdoor import ( - HOT_LOAD_PATH, - advance_latest_decision, - create_frontdoor_app, - pool_state_from_server_infos, -) - - -class AdvanceDecisionTest(unittest.TestCase): - def test_accepts_strictly_newer_same_run(self) -> None: - self.assertEqual( - advance_latest_decision("run-a", 4, "weight_v000005", "run-a"), - {"run_id": "run-a", "version": 5, "reset": False}, - ) - - def test_rejects_rewind_same_run(self) -> None: - for identity in ("weight_v000004", "weight_v000003"): - decision = advance_latest_decision("run-a", 4, identity, "run-a") - self.assertEqual(decision["error"]["type"], "WeightRewindRejected") - self.assertEqual(decision["error"]["current_version"], 4) - - def test_new_run_resets_even_if_version_lower(self) -> None: - # A new run restarts at v1; accepting it after the prior run reached v5 is - # the begin-a-new-run path, not a rewind — the idempotency fix. - self.assertEqual( - advance_latest_decision("run-a", 5, "weight_v000001", "run-b"), - {"run_id": "run-b", "version": 1, "reset": True}, - ) - - def test_runless_layout_stays_monotonic(self) -> None: - ok = advance_latest_decision(None, 4, "weight_v000005", None) - self.assertEqual(ok, {"run_id": None, "version": 5, "reset": False}) - rewind = advance_latest_decision(None, 5, "weight_v000005", None) - self.assertEqual(rewind["error"]["type"], "WeightRewindRejected") - - def test_rejects_unparseable_identity(self) -> None: - self.assertEqual( - advance_latest_decision(None, 0, "base", None)["error"]["type"], "InvalidIdentity" - ) - - -class PoolStateTest(unittest.TestCase): - def test_ready_only_when_idle_and_no_error(self) -> None: - state = pool_state_from_server_infos( - [ - {"run_id": "a", "current_run_id": "run-x", "current_version": 5, "sync_state": "IDLE", "last_sync_error": None}, - {"run_id": "b", "current_run_id": "run-x", "current_version": 4, "sync_state": "PREFETCHING", "last_sync_error": None}, - {"run_id": "c", "current_run_id": "run-x", "current_version": 5, "sync_state": "ERROR", "last_sync_error": "boom"}, - ] - ) - by_id = {r.replica_id: r for r in state.replicas} - self.assertTrue(by_id["a"].readiness) - self.assertEqual(by_id["a"].current_snapshot_identity, "run-x/weight_v000005") - self.assertFalse(by_id["b"].readiness) - self.assertEqual(by_id["b"].readiness_reason, "PREFETCHING") - self.assertFalse(by_id["c"].readiness) - self.assertEqual(by_id["c"].readiness_reason, "boom") - self.assertEqual(state.ready_count(target_snapshot_identity="run-x/weight_v000005"), 1) - - def test_identity_carries_run_id(self) -> None: - # A replica on run-b advertises a run-scoped identity, so it is not - # miscounted ready for a different run's same-numbered version. - state = pool_state_from_server_infos( - [{"run_id": "a", "current_run_id": "run-b", "current_version": 1, "sync_state": "IDLE"}] - ) - self.assertEqual(state.replicas[0].current_snapshot_identity, "run-b/weight_v000001") - self.assertEqual(state.ready_count(target_snapshot_identity="run-b/weight_v000001"), 1) - self.assertEqual(state.ready_count(target_snapshot_identity="run-a/weight_v000001"), 0) - - -class FrontdoorAppTest(unittest.TestCase): - def _client(self, *, run_id="run-a", version=5, authorize=None): - from fastapi.responses import JSONResponse - from fastapi.testclient import TestClient - - state = {"run_id": run_id, "version": version} - calls: dict[str, list] = {"advanced": [], "woke": [], "proxied": []} - - async def read_current_pointer(): - return (state["run_id"], state["version"]) - - async def advance_to(rid, v: int) -> None: - calls["advanced"].append((rid, v)) - state["run_id"], state["version"] = rid, v - - async def list_server_infos(): - return [ - { - "run_id": "a", - "current_run_id": state["run_id"], - "current_version": state["version"], - "sync_state": "IDLE", - } - ] - - async def wake(v: int) -> None: - calls["woke"].append(v) - - async def proxy(request, path): - calls["proxied"].append(path) - return JSONResponse({"proxied": path}) - - app = create_frontdoor_app( - read_current_pointer=read_current_pointer, - advance_to=advance_to, - list_server_infos=list_server_infos, - proxy=proxy, - authorize=authorize, - wake=wake, - ) - return TestClient(app), calls - - def test_post_advances_on_newer_identity_and_wakes(self) -> None: - client, calls = self._client(run_id="run-a", version=4) - resp = client.post(HOT_LOAD_PATH, json={"identity": "weight_v000005", "run_id": "run-a"}) - self.assertEqual(resp.status_code, 200) - self.assertTrue(resp.json()["accepted"]) - self.assertEqual(resp.json()["current_snapshot_identity"], "run-a/weight_v000005") - self.assertEqual(calls["advanced"], [("run-a", 5)]) - self.assertEqual(calls["woke"], [5]) - - def test_post_new_run_accepted_as_reset(self) -> None: - client, calls = self._client(run_id="run-a", version=5) - resp = client.post(HOT_LOAD_PATH, json={"identity": "weight_v000001", "run_id": "run-b"}) - self.assertEqual(resp.status_code, 200) - self.assertEqual(calls["advanced"], [("run-b", 1)]) - self.assertEqual(resp.json()["current_snapshot_identity"], "run-b/weight_v000001") - - def test_post_rejects_rewind_without_advancing(self) -> None: - client, calls = self._client(run_id="run-a", version=5) - resp = client.post(HOT_LOAD_PATH, json={"identity": "weight_v000005", "run_id": "run-a"}) - self.assertEqual(resp.status_code, 409) - self.assertEqual(resp.json()["error"]["type"], "WeightRewindRejected") - self.assertEqual(calls["advanced"], []) - self.assertEqual(calls["woke"], []) - - def test_post_requires_identity(self) -> None: - client, _ = self._client() - self.assertEqual(client.post(HOT_LOAD_PATH, json={}).status_code, 400) - - def test_get_reports_pool_readiness(self) -> None: - client, _ = self._client(run_id="run-a", version=5) - body = client.get(HOT_LOAD_PATH).json() - self.assertEqual(len(body["replicas"]), 1) - self.assertEqual(body["replicas"][0]["current_snapshot_identity"], "run-a/weight_v000005") - self.assertTrue(body["replicas"][0]["readiness"]) - - def test_catch_all_proxies(self) -> None: - client, calls = self._client() - resp = client.post("/v1/chat/completions", json={"model": "m"}) - self.assertEqual(resp.json(), {"proxied": "v1/chat/completions"}) - self.assertEqual(calls["proxied"], ["v1/chat/completions"]) - - def test_auth_rejection_blocks_every_route(self) -> None: - from fastapi.responses import JSONResponse - - def deny(_headers): - return JSONResponse({"error": "unauthorized"}, status_code=401) - - client, calls = self._client(authorize=deny) - self.assertEqual( - client.post(HOT_LOAD_PATH, json={"identity": "weight_v000009", "run_id": "run-a"}).status_code, - 401, - ) - self.assertEqual(client.get(HOT_LOAD_PATH).status_code, 401) - self.assertEqual(client.post("/v1/chat/completions", json={}).status_code, 401) - self.assertEqual(calls["advanced"], []) - self.assertEqual(calls["proxied"], []) - - -if __name__ == "__main__": - unittest.main() diff --git a/cookbook/standalone_rollouts/modal_serve.py b/cookbook/standalone_rollouts/modal_serve.py deleted file mode 100644 index 9adec36..0000000 --- a/cookbook/standalone_rollouts/modal_serve.py +++ /dev/null @@ -1,599 +0,0 @@ -"""Standalone Modal rollout provider implementing the customer hot-load API. - -Log-as-truth design: the front door is a singleton that owns the monotonic -``latest`` pointer in the S3 transport and implements the customer's -``POST/GET /hot_load`` API. An elastic Flash pool of -SGLang servers + stitch sidecars reconciles to ``latest`` on its own and serves -inference. There is no central desired-state mailbox: the pool pulls, and the -front door derives readiness by enumerating the live containers. - -Run commands from the repo root, for example: - - uv run --extra modal modal deploy -m cookbook.standalone_rollouts.modal_serve -""" - -import asyncio -import importlib -import os -import subprocess -from pathlib import Path - -import modal -import modal.experimental - -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, - wake_targets, -) - - -PROVIDER_CONFIG = os.environ.get("PROVIDER_CONFIG", "moonlight_hot_load") -exp = importlib.import_module(f"cookbook.standalone_rollouts.configs.{PROVIDER_CONFIG}") - -APP_NAME = exp.APP_NAME -MODEL_NAME = exp.MODEL_NAME -ROLLOUT_CONCURRENCY = exp.ROLLOUT_CONCURRENCY -SIDECAR_PORT = 8000 -SGLANG_PORT = 8001 -MINUTES = 60 -SERVER_STARTUP_TIMEOUT = 35 * MINUTES -LOCAL_CHECKPOINT_PATH = exp.LOCAL_CHECKPOINT_PATH -S3_TRANSPORT_BUCKET_NAME = os.environ.get( - "STITCH_SHIM_S3_BUCKET_NAME", exp.S3_TRANSPORT_BUCKET_NAME -) -S3_TRANSPORT_KEY_PREFIX = os.environ.get( - "STITCH_SHIM_S3_KEY_PREFIX", exp.S3_TRANSPORT_KEY_PREFIX -) -S3_TRANSPORT_MOUNT_PATH = exp.S3_TRANSPORT_MOUNT_PATH -S3_TRANSPORT_REGION = os.environ.get("STITCH_SHIM_S3_REGION", exp.S3_TRANSPORT_REGION) -S3_TRANSPORT_OIDC_AUTH_ROLE_ARN = os.environ.get( - "STITCH_SHIM_S3_OIDC_AUTH_ROLE_ARN", exp.S3_TRANSPORT_OIDC_AUTH_ROLE_ARN -) - -SLIME_IMAGE_TAG = "slimerl/slime:nightly-dev-20260527a" -SLIME_ROOT = "/root/slime" -SLIME_REPO_URL = "https://github.com/modal-projects/slime.git" -# PR #5 head (disaggregated-rollout, stacked on disk-delta-weight-sync). The -# provider sidecar applies disk deltas host-side via slime.utils.disk_delta, so -# the image must carry that branch's slime plus its checksum/compression deps. -# Pin a SHA, not the branch tip: the clone is a cached image layer. -SLIME_REPO_REF = "ebfe153949b1a69c39e92f947ed5d475166dd724" # incl. deepseekv3 router-dtype export fix + per-request rollout hook - -image = ( - modal.Image.from_registry(SLIME_IMAGE_TAG) - .entrypoint([]) - .run_commands(f"rm -rf {exp.HF_CACHE_PATH}") - # Replace the bundled slime with the disk-delta branch so the sidecar can - # import slime.utils.disk_delta (host-side apply). - .run_commands( - f"rm -rf {SLIME_ROOT}" - 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 - "boto3", - "fastapi", - "httpx", - "uvicorn", - # slime.utils.disk_delta host-side apply: zstd decompress + xxhash - # (xxh3-128 default) / blake3 checksums. slime is installed --no-deps. - "zstandard", - "xxhash", - "blake3", - ) - .env( - { - "EXPERIMENT_CONFIG": PROVIDER_CONFIG, - "PROVIDER_CONFIG": PROVIDER_CONFIG, - "HF_XET_HIGH_PERFORMANCE": "1", - "HF_HUB_ENABLE_HF_TRANSFER": "1", - "STITCH_SHIM_MODAL_APP_NAME": APP_NAME, - "STITCH_SHIM_MODAL_CLS_NAME": "Server", - "STITCH_SHIM_TRANSPORT_ROOT": str(S3_TRANSPORT_MOUNT_PATH), - "STITCH_LOCAL_CHECKPOINT_DIR": LOCAL_CHECKPOINT_PATH, - } - ) - .add_local_python_source("stitch") - .add_local_dir( - Path(__file__).parents[1], - remote_path="/root/cookbook", - ignore=["**/__pycache__"], - ) -) - -# Dev iteration: SLIME_LOCAL_DIR overlays a local slime checkout onto the image's -# cloned fork (installed editable at /root/slime), so fork edits take effect on -# container start with no image rebuild. Unset by default. -if slime_local := os.environ.get("SLIME_LOCAL_DIR"): - image = image.add_local_dir( - slime_local, - remote_path=SLIME_ROOT, - ignore=[".git", "**/__pycache__", "**/*.pyc"], - ) - -with image.imports(): - from autoinference_utils.endpoint import SGLangEndpoint, warmup_chat_completions - - -def _key_prefix_for_mount(prefix: str) -> str | None: - prefix = prefix.strip("/") - if not prefix: - return None - return f"{prefix}/" - - -hf_cache_volume = modal.Volume.from_name( - exp.HF_CACHE_VOLUME_NAME, create_if_missing=True -) -# read_only=False: the front door writes the `latest` pointer here. -s3_transport_mount = modal.CloudBucketMount( - bucket_name=S3_TRANSPORT_BUCKET_NAME, - key_prefix=_key_prefix_for_mount(S3_TRANSPORT_KEY_PREFIX), - secret=modal.Secret.from_dict({"AWS_REGION": S3_TRANSPORT_REGION}) - if S3_TRANSPORT_REGION - else None, - oidc_auth_role_arn=S3_TRANSPORT_OIDC_AUTH_ROLE_ARN, - read_only=False, -) -app = modal.App(APP_NAME) - -SGLANG_SERVER_ARGS = { - "--served-model-name": MODEL_NAME, - "--dtype": "bfloat16", - "--cuda-graph-max-bs": str(ROLLOUT_CONCURRENCY), - "--max-running-requests": str(ROLLOUT_CONCURRENCY), - "--trust-remote-code": "", - **exp.SGLANG_SERVER_ARGS, -} - -WARMUP_PAYLOAD = { - "model": MODEL_NAME, - "messages": [{"role": "user", "content": "Reply with exactly OK."}], - "max_tokens": 8, - "temperature": 0, - "chat_template_kwargs": {"enable_thinking": False}, -} - - -@app.cls( - image=image, - gpu=f"{exp.GPU}:{exp.ROLLOUT_NUM_GPUS_PER_ENGINE}", - cloud=exp.CLOUD, - region=exp.REGION, - volumes={ - str(exp.HF_CACHE_PATH): hf_cache_volume, - str(S3_TRANSPORT_MOUNT_PATH): s3_transport_mount, - }, - secrets=[modal.Secret.from_name(exp.SHIM_SECRET_NAME)], - min_containers=exp.ROLLOUT_MIN_CONTAINERS, - timeout=40 * MINUTES, - scaledown_window=15 * MINUTES, - include_source=False, -) -@modal.experimental.http_server( - port=SIDECAR_PORT, - proxy_regions=exp.PROXY_REGIONS, - exit_grace_period=25, - startup_timeout=SERVER_STARTUP_TIMEOUT, -) -@modal.concurrent(target_inputs=ROLLOUT_CONCURRENCY) -class Server: - """One SGLang server plus the stitch weight-sync sidecar, reconciling to the - `latest` pointer the front door advances.""" - - @modal.enter() - def startup(self) -> None: - self.endpoint = SGLangEndpoint( - model_path=MODEL_NAME, - worker_port=SGLANG_PORT, - tp=exp.ROLLOUT_NUM_GPUS_PER_ENGINE, - extra_server_args=SGLANG_SERVER_ARGS, - health_timeout=SERVER_STARTUP_TIMEOUT, - health_poll_interval=10.0, - ) - self.endpoint.start() - warmup_chat_completions( - port=SGLANG_PORT, - payload=WARMUP_PAYLOAD, - successful_requests=2, - request_timeout=120.0, - max_attempts_per_request=3, - ) - # Deltas are applied host-side onto a copy of the base checkpoint; the - # base resolves to the same HF cache snapshot the SGLang server loaded. - from huggingface_hub import snapshot_download - - base_checkpoint_dir = snapshot_download(MODEL_NAME, local_files_only=True) - self.sidecar = _start_provider_sidecar(base_checkpoint_dir=base_checkpoint_dir) - helpers.wait_http( - f"http://127.0.0.1:{SIDECAR_PORT}/health", - self.sidecar, - SERVER_STARTUP_TIMEOUT, - ) - print( - f"Rollout server ready: model={MODEL_NAME}, target_inputs={ROLLOUT_CONCURRENCY}" - ) - - @modal.exit() - def stop(self) -> None: - helpers.terminate_process(getattr(self, "sidecar", None)) - if hasattr(self, "endpoint"): - self.endpoint.stop() - - -@app.function( - image=image, - volumes={str(exp.HF_CACHE_PATH): hf_cache_volume}, - timeout=2 * 60 * MINUTES, - secrets=[modal.Secret.from_name(exp.HF_SECRET_NAME)], - include_source=False, -) -def download_model() -> None: - from huggingface_hub import snapshot_download - - snapshot_download(repo_id=MODEL_NAME) - hf_cache_volume.commit() - - -@app.function( - image=image, - secrets=[modal.Secret.from_name(exp.SHIM_SECRET_NAME)], - timeout=35 * MINUTES, - include_source=False, -) -def check(timeout_seconds: int = 20 * MINUTES) -> None: - """Authenticated smoke from inside Modal: reads the provider secret (so the - API key never leaves Modal), polls GET /hot_load readiness through the front - door, then serves a base completion through it. Run with: - - uv run --extra modal modal run -m cookbook.standalone_rollouts.modal_serve::check - """ - import json - import time - import urllib.request - - gateway = modal.Server.from_name(APP_NAME, "FrontDoor").get_url().rstrip("/") - headers = _shim_headers() - deadline = time.time() + timeout_seconds - last = "" - while True: - try: - req = urllib.request.Request( - f"{gateway}/hot_load/v1/models/hot_load", headers=headers - ) - with urllib.request.urlopen(req, timeout=30) as resp: - pool = json.load(resp) - ready = [r for r in pool.get("replicas", []) if r.get("readiness")] - print( - f"pool: {len(pool.get('replicas', []))} replicas, {len(ready)} ready :: {pool}" - ) - if len(ready) >= exp.ROLLOUT_MIN_CONTAINERS: - break - last = f"only {len(ready)} ready" - except Exception as exc: # noqa: BLE001 - last = f"{type(exc).__name__}: {exc}" - if time.time() > deadline: - raise TimeoutError(f"front-door readiness timed out: {last}") - time.sleep(10) - - payload = { - "model": MODEL_NAME, - "messages": [{"role": "user", "content": "Reply with exactly OK."}], - "max_tokens": 8, - "temperature": 0, - "chat_template_kwargs": {"enable_thinking": False}, - } - req = urllib.request.Request( - f"{gateway}/v1/chat/completions", - data=json.dumps(payload).encode("utf-8"), - headers={"Content-Type": "application/json", **headers}, - ) - with urllib.request.urlopen(req, timeout=180) as resp: - print("completion:", json.dumps(json.load(resp))[:1200]) - - -@app.local_entrypoint() -def print_url() -> None: - print(frontdoor_url()) - - -@app.local_entrypoint() -def print_secret_template() -> None: - print( - "\n".join( - [ - f"modal secret create {exp.SHIM_SECRET_NAME} \\", - " STITCH_SHIM_API_KEY=... \\", - " STITCH_SHIM_PROVIDER_MODEL=moonlight \\", - " STITCH_SHIM_PROVIDER_DEPLOYMENT=rollout-prod", - ] - ) - ) - - -@app.local_entrypoint() -def smoke( - timeout_seconds: int = 30 * MINUTES, - api_key: str = "", - provider_model: str = "", - provider_deployment: str = "", -) -> None: - """Check the deployed gateway can report pool state and serve a completion.""" - import json - import time - import urllib.request - - gateway = frontdoor_url() - deadline = time.time() + timeout_seconds - headers = _shim_headers( - api_key=api_key, - provider_model=provider_model, - provider_deployment=provider_deployment, - ) - last_error = "" - while True: - try: - req = urllib.request.Request( - f"{gateway}/hot_load/v1/models/hot_load", headers=headers - ) - with urllib.request.urlopen(req, timeout=30) as resp: - pool = json.load(resp) - if len(pool.get("replicas", [])) >= exp.ROLLOUT_MIN_CONTAINERS: - break - last_error = f"expected {exp.ROLLOUT_MIN_CONTAINERS} replicas, got {pool}" - except Exception as exc: # noqa: BLE001 - last_error = f"{type(exc).__name__}: {exc}" - if time.time() > deadline: - raise TimeoutError(f"Provider smoke failed: {last_error}") - print(f"Waiting for provider pool: {last_error}") - time.sleep(10) - - payload = { - "model": MODEL_NAME, - "messages": [{"role": "user", "content": "Reply with exactly OK."}], - "max_tokens": 8, - "temperature": 0, - "chat_template_kwargs": {"enable_thinking": False}, - } - req = urllib.request.Request( - f"{gateway}/v1/chat/completions", - data=json.dumps(payload).encode("utf-8"), - headers={"Content-Type": "application/json", **headers}, - ) - with urllib.request.urlopen(req, timeout=180) as resp: - print(json.dumps(json.load(resp), indent=2)[:2000]) - - -def provider_gateway_url() -> str: - urls = Server._experimental_get_flash_urls() - if not urls: - return resolve_flash_gateway_url(APP_NAME, Server.__name__) - return str(urls[0]).rstrip("/") - - -def _frontdoor_headers(raw: dict[str, str]) -> dict[str, str]: - """Map the protocol-neutral ``x-session-affinity`` onto Modal's gateway - session header so Flash co-locates related requests at routing time. This - must run *before* the gateway (i.e. outside the Server Cls), since the - gateway picks the container from the header.""" - headers: dict[str, str] = {} - affinity: str | None = None - for key, value in raw.items(): - lower = key.lower() - if lower in {"host", "content-length"}: - continue - if lower == "x-session-affinity": - affinity = value - continue - headers[key] = value - if affinity: - headers["Modal-Session-ID"] = affinity - return headers - - -def _auth_error(headers): - """Validate the customer auth headers against the provider secret. Returns a - JSONResponse to reject, or None to allow.""" - from fastapi.responses import JSONResponse - - api_key = os.environ.get("STITCH_SHIM_API_KEY") - if api_key and headers.get("authorization") != f"Bearer {api_key}": - return JSONResponse({"error": "unauthorized"}, status_code=401) - provider_model = os.environ.get("STITCH_SHIM_PROVIDER_MODEL") - if provider_model and headers.get("provider-model") != provider_model: - return JSONResponse( - {"error": "Provider-Model header does not match this deployment"}, - status_code=400, - ) - provider_deployment = os.environ.get("STITCH_SHIM_PROVIDER_DEPLOYMENT") - if ( - provider_deployment - and headers.get("provider-deployment") != provider_deployment - ): - return JSONResponse( - {"error": "Provider-Deployment header does not match this deployment"}, - status_code=400, - ) - return None - - -FRONTDOOR_PORT = 8000 - - -@app.server( - image=image, - volumes={str(S3_TRANSPORT_MOUNT_PATH): s3_transport_mount}, - secrets=[modal.Secret.from_name(exp.SHIM_SECRET_NAME)], - min_containers=1, - max_containers=1, # singleton: exactly one writer of the `latest` pointer - nonpreemptible=True, # keep the sole writer up; a preemption blips the API - scaledown_window=2, - region=exp.REGION, # co-locate the front door with the rollout pool (us) - routing_region=exp.ROUTING_REGION, # share the pool's Flash proxy region - port=FRONTDOOR_PORT, - unauthenticated=True, # public customer endpoint; auth is enforced in-app - target_concurrency=1000, # one container, many concurrent inputs - startup_timeout=5 * MINUTES, - exit_grace_period=25, - include_source=False, -) -class FrontDoor: - """Public front door: the single writer of `latest` and the customer - hot-load API, plus an affinity-relabeling proxy to the rollout gateway. - - Serves the front-door FastAPI app under uvicorn (App.server). No - `from __future__ import annotations` interplay here — frontdoor_mod's - create_frontdoor_app resolves its handler annotations against its own - eager fastapi import. - """ - - @modal.enter() - def start(self) -> None: - import threading - - import httpx - import uvicorn - from fastapi.responses import Response - - board = FilesystemBulletinBoard(str(S3_TRANSPORT_MOUNT_PATH), layout="slime") - gateway: dict[str, str | None] = {"url": None} - clients: dict[str, httpx.AsyncClient] = {} - - def _proxy_client() -> httpx.AsyncClient: - client = clients.get("client") - if client is None: - client = httpx.AsyncClient(timeout=None, trust_env=False) - clients["client"] = client - return client - - async def read_current_pointer() -> tuple[str | None, int]: - return board.read_latest() - - 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" - ) - - async def list_server_infos() -> list[dict]: - targets = await asyncio.to_thread( - discover_flash_targets, APP_NAME, Server.__name__ - ) - infos: list[dict] = [] - async with httpx.AsyncClient(timeout=10.0, trust_env=False) as client: - for target in targets: - try: - resp = await client.get(f"{target}/server_info") - infos.append(resp.json()) - except Exception: # noqa: BLE001 — unreachable replica reported as not-ready - infos.append( - {"sync_state": None, "last_sync_error": "unreachable"} - ) - return infos - - async def wake(version: int) -> None: - targets = await asyncio.to_thread( - discover_flash_targets, APP_NAME, Server.__name__ - ) - await asyncio.to_thread(wake_targets, targets, version) - - async def proxy(request, path: str) -> Response: - if gateway["url"] is None: - gateway["url"] = provider_gateway_url() - headers = _frontdoor_headers(dict(request.headers)) - body = await request.body() - resp = await _proxy_client().request( - request.method, - f"{gateway['url']}/{path}", - headers=headers, - content=body, - params=request.query_params, - ) - return Response( - content=resp.content, - status_code=resp.status_code, - media_type=resp.headers.get("content-type") or None, - ) - - asgi_app = frontdoor_mod.create_frontdoor_app( - read_current_pointer=read_current_pointer, - advance_to=advance_to, - list_server_infos=list_server_infos, - proxy=proxy, - authorize=_auth_error, - wake=wake, - ) - config = uvicorn.Config( - asgi_app, host="0.0.0.0", port=FRONTDOOR_PORT, log_level="info" - ) - self._server = uvicorn.Server(config) - self._thread = threading.Thread(target=self._server.run, daemon=True) - self._thread.start() - - @modal.exit() - def stop(self) -> None: - server = getattr(self, "_server", None) - if server is not None: - server.should_exit = True - thread = getattr(self, "_thread", None) - if thread is not None: - thread.join(timeout=25) - - -def frontdoor_url() -> str: - """Advertised provider URL: external clients hit the front door, which owns - `latest` and relabels affinity pre-gateway.""" - return FrontDoor.get_url().rstrip("/") - - -def _start_provider_sidecar(*, base_checkpoint_dir: str) -> subprocess.Popen: - cmd = [ - "python3", - "-m", - "cookbook.standalone_rollouts.provider", - "--host", - "0.0.0.0", - "--port", - str(SIDECAR_PORT), - "--upstream-url", - f"http://127.0.0.1:{SGLANG_PORT}", - "--transport-root", - str(S3_TRANSPORT_MOUNT_PATH), - "--local-checkpoint-dir", - LOCAL_CHECKPOINT_PATH, - "--base-checkpoint-dir", - base_checkpoint_dir, - "--commit-mode", - exp.COMMIT_MODE, - ] - print("Starting provider sidecar:", " ".join(cmd)) - return subprocess.Popen(cmd, start_new_session=True) - - -def _shim_headers( - *, api_key: str = "", provider_model: str = "", provider_deployment: str = "" -) -> dict[str, str]: - headers: dict[str, str] = {} - if api_key := (api_key or os.environ.get("STITCH_SHIM_API_KEY")): - headers["Authorization"] = f"Bearer {api_key}" - if provider_model := ( - provider_model or os.environ.get("STITCH_SHIM_PROVIDER_MODEL") - ): - headers["Provider-Model"] = provider_model - if provider_deployment := ( - provider_deployment or os.environ.get("STITCH_SHIM_PROVIDER_DEPLOYMENT") - ): - headers["Provider-Deployment"] = provider_deployment - return headers diff --git a/cookbook/standalone_rollouts/provider.py b/cookbook/standalone_rollouts/provider.py deleted file mode 100644 index b2c11d1..0000000 --- a/cookbook/standalone_rollouts/provider.py +++ /dev/null @@ -1,143 +0,0 @@ -"""Per-container rollout sidecar for the standalone hot-load provider. - -A pool member of the log-as-truth design: a :class:`WeightSyncManager` that -reconciles its local SGLang server to the shared ``latest`` pointer. The trainer -uploads ``weight_v{N}/`` version dirs to object storage (mounted here as the -transport root) in the flat slime/customer layout, and the front door advances -``latest`` on ``POST /hot_load``. Each replica pulls from that durable monotonic -log — on startup, on a best-effort wake, and on a periodic reconcile — applies -the disk delta host-side, and reloads the engine. - -There is no desired-mailbox and no self-reported replica state: a scaled-up -container catches up by reading ``latest`` with no push, and the front door -derives pool readiness by enumerating the live containers and querying -``/server_info``. Authentication and the customer hot-load API live in the front -door (``modal_serve.py``), not here. -""" - -from __future__ import annotations - -import argparse -import logging -import os -import uuid - -from stitch.bulletin import FilesystemBulletinBoard -from stitch.engines.sglang import SGLangDiskDeltaAdapter -from stitch.servers.sglang import create_app as create_sglang_app -from stitch.sync import CommitMode, WeightSyncManager - - -logger = logging.getLogger(__name__) -VERSIONED_ROUTES = frozenset({"generate", "v1/chat/completions", "v1/completions"}) - - -def build_manager( - *, - upstream_url: str, - transport_root: str, - local_checkpoint_dir: str, - base_checkpoint_dir: str, - commit_mode: CommitMode = "in_place", - run_id: str | None = None, - debug_requests: bool = False, -) -> WeightSyncManager: - # The transport root is the object-store mount holding the flat - # weight_v{N}/ version dirs and the raw `latest` pointer (slime/customer - # layout). Deltas are read straight from the mount during apply. - board = FilesystemBulletinBoard(transport_root, 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, - ) - - -def create_app(manager: WeightSyncManager, *, upstream_url: str, poll_interval: float = 5.0): - # include_sync_routes=True keeps /rpc_sync_from_bulletin_board so the front - # door can wake replicas the moment it advances `latest`; the periodic - # reconcile is the fallback that converges anything that missed the wake. - return create_sglang_app( - manager, - upstream_url=upstream_url, - versioned_routes=VERSIONED_ROUTES, - include_sync_routes=True, - background_sync_interval=poll_interval, - ) - - -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( - "--transport-root", - default=os.environ.get("STITCH_SHIM_TRANSPORT_ROOT"), - help="Object-store mount holding weight_v{N}/ dirs and the raw `latest` pointer.", - ) - 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( - "--commit-mode", - choices=("quiesce", "in_place"), - default=os.environ.get("STITCH_SHIM_COMMIT_MODE", "in_place"), - ) - parser.add_argument( - "--poll-interval", - type=float, - default=float(os.environ.get("STITCH_SHIM_POLL_INTERVAL", "5.0")), - help="Seconds between background reconciles against the `latest` pointer.", - ) - parser.add_argument("--run-id", default=os.environ.get("MODAL_TASK_ID") or uuid.uuid4().hex) - parser.add_argument( - "--debug-requests", - action="store_true", - default=os.environ.get("STITCH_SHIM_DEBUG_REQUESTS", "").lower() in {"1", "true", "yes"}, - ) - args = parser.parse_args() - if not args.transport_root: - raise SystemExit("--transport-root/STITCH_SHIM_TRANSPORT_ROOT is required") - 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, - transport_root=args.transport_root, - local_checkpoint_dir=args.local_checkpoint_dir, - base_checkpoint_dir=args.base_checkpoint_dir, - commit_mode=args.commit_mode, - run_id=args.run_id, - debug_requests=args.debug_requests, - ) - uvicorn.run( - create_app(manager, upstream_url=args.upstream_url, poll_interval=args.poll_interval), - host=args.host, - port=args.port, - log_level="info", - ) - - -if __name__ == "__main__": - main() diff --git a/cookbook/standalone_rollouts/provider_test.py b/cookbook/standalone_rollouts/provider_test.py deleted file mode 100644 index d94f097..0000000 --- a/cookbook/standalone_rollouts/provider_test.py +++ /dev/null @@ -1,150 +0,0 @@ -from __future__ import annotations - -import asyncio -import json -import tempfile -import unittest -from pathlib import Path -from unittest import mock - -from cookbook.standalone_rollouts.provider import build_manager, create_app -from stitch.bulletin import FilesystemBulletinBoard -from stitch.engines.sglang import SGLangDiskDeltaAdapter -from stitch.sync import WeightSyncManager - - -class _FakeEngine: - backend = "fake" - - def __init__(self) -> None: - self.applies: list[int] = [] - self.events: list[str] = [] - - async def flush_cache(self) -> None: - self.events.append("flush") - - async def apply_manifest(self, manifest, version_path) -> None: - self.applies.append(manifest.version) - self.events.append("apply") - - async def pause_generation(self) -> None: - self.events.append("pause") - - async def continue_generation(self) -> None: - self.events.append("continue") - - -def _write_slime_version(root: Path, version: int, base: int) -> None: - vdir = root / f"weight_v{version:06d}" - vdir.mkdir(parents=True) - (vdir / "model.safetensors.index.json").write_text( - json.dumps( - { - "metadata": { - "version": f"{version:06d}", - "base_version": f"{base:06d}", - "delta_encoding": "xor", - "compression_format": "zstd", - "checksum_format": "xxh3-128", - }, - "weight_map": {"w": "model-00001-of-00001.safetensors"}, - } - ), - encoding="utf-8", - ) - - -class BuildManagerTest(unittest.TestCase): - def test_wires_slime_board_and_disk_delta_adapter(self) -> None: - with tempfile.TemporaryDirectory() as tmp: - manager = build_manager( - upstream_url="http://127.0.0.1:30000/", - transport_root=tmp, - local_checkpoint_dir="/local", - base_checkpoint_dir="/base", - ) - - self.assertIsInstance(manager, WeightSyncManager) - self.assertEqual(manager.board.layout, "slime") - self.assertEqual(str(manager.board.root), tmp) - self.assertIsInstance(manager.engine, SGLangDiskDeltaAdapter) - self.assertEqual(manager.engine.local_checkpoint_dir, "/local") - self.assertEqual(manager.engine.base_checkpoint_dir, "/base") - # The provider serves the customer; in_place is the perf default. - self.assertEqual(manager.commit_mode, "in_place") - - -class ProviderSyncTest(unittest.TestCase): - def test_startup_sync_pulls_chain_from_slime_latest(self) -> None: - async def run() -> None: - with tempfile.TemporaryDirectory() as tmp: - root = Path(tmp) - _write_slime_version(root, 1, 0) - _write_slime_version(root, 2, 1) - (root / "latest").write_text("000002", encoding="utf-8") - - engine = _FakeEngine() - manager = WeightSyncManager( - board=FilesystemBulletinBoard(root, layout="slime"), engine=engine - ) - await manager.startup_sync() - - # Reconciled to the `latest` pointer, applying the chain in order. - self.assertEqual(manager.current_version, 2) - self.assertEqual(engine.applies, [1, 2]) - - asyncio.run(run()) - - -class ProviderAppTest(unittest.TestCase): - def test_create_app_serves_versioned_proxy_and_server_info(self) -> None: - from fastapi.testclient import TestClient - - with tempfile.TemporaryDirectory() as tmp: - manager = WeightSyncManager( - board=FilesystemBulletinBoard(tmp, layout="slime"), engine=_FakeEngine() - ) - # poll_interval=0 disables the background reconcile task for the test. - app = create_app(manager, upstream_url="http://127.0.0.1:9", poll_interval=0) - - with TestClient(app) as client, mock.patch("httpx.AsyncClient", _RecordingUpstream): - _RecordingUpstream.last_json = None - resp = client.post("/generate", json={"text": "hi"}) - - self.assertEqual(resp.status_code, 200) - # Versioned route: stamped with the composed extra_key namespace. - self.assertEqual(_RecordingUpstream.last_json["extra_key"], "wv0;") - self.assertEqual(resp.json()["meta_info"]["weight_version_start"], 0) - self.assertEqual(client.get("/server_info").json()["current_version"], 0) - - -class _RecordingUpstream: - last_json: dict | None = None - last_url: str | None = None - - def __init__(self, *args, **kwargs) -> None: - pass - - async def __aenter__(self) -> "_RecordingUpstream": - return self - - async def __aexit__(self, *exc) -> bool: - return False - - async def aclose(self) -> None: - return None - - async def request(self, method, url, **kwargs): - import httpx - - type(self).last_url = url - type(self).last_json = kwargs.get("json") - return httpx.Response( - 200, - json={"text": "ok", "meta_info": {"finish_reason": {"type": "length"}}}, - request=httpx.Request(method, url), - ) - - -if __name__ == "__main__": - unittest.main() diff --git a/cookbook/standalone_rollouts/slime/__init__.py b/cookbook/standalone_rollouts/slime/__init__.py deleted file mode 100644 index 55c8feb..0000000 --- a/cookbook/standalone_rollouts/slime/__init__.py +++ /dev/null @@ -1 +0,0 @@ -"""Optional SLIME integration harness for standalone rollout providers.""" diff --git a/cookbook/standalone_rollouts/slime/configs/__init__.py b/cookbook/standalone_rollouts/slime/configs/__init__.py deleted file mode 100644 index 8ecba1d..0000000 --- a/cookbook/standalone_rollouts/slime/configs/__init__.py +++ /dev/null @@ -1 +0,0 @@ -"""SLIME trainer configs for the standalone rollout cookbook.""" diff --git a/cookbook/standalone_rollouts/slime/configs/moonlight_slime_trainer.py b/cookbook/standalone_rollouts/slime/configs/moonlight_slime_trainer.py deleted file mode 100644 index 1148d6e..0000000 --- a/cookbook/standalone_rollouts/slime/configs/moonlight_slime_trainer.py +++ /dev/null @@ -1,157 +0,0 @@ -"""Moonlight-16B-A3B SLIME trainer config for the API-shim provider. - -The async + routing-replay analog of cookbook/slime_disagg/configs/moonlight_disagg.py, -but publishing through the standalone provider's customer hot-load API (S3 -transport + singleton front door) instead of the in-cluster Modal Volume -bulletin board. - -Key difference from slime_disagg's M3: staleness is gated by a PUBLISH-HOOK -readiness BARRIER, not a per-request version pin. `announce_and_wait` copies the -new version to S3, POSTs the hot-load API, and blocks until the front door -reports every live replica ready on the new version (readiness_threshold defaults -to 1.0), so the next rollouts always run on current weights. The request hook -therefore leaves the version gate OFF and only carries auth headers, retries, and -session affinity. -""" - -from __future__ import annotations - -from cookbook.slime_disagg.configs.base import DATA_PATH, ModalConfig, SlimeConfig - - -HF_SECRET_NAME = "huggingface-secret" -# The trainer calls the provider shim, so it needs the same optional auth values. -SHIM_SECRET_NAME = "stitch-api-shim-provider" -# The mounted S3 transport the provider pool pulls from (flat customer layout). -TRANSPORT_ROOT = "/mnt/stitch-s3-transport" -# slime's disk-delta writer uses atomic rename (ENOSYS on the S3 mount), so it -# publishes weight_v{N}/ to local disk; announce_and_wait copies each version to -# the transport (PutObject) before signalling the hot-load API. -LOCAL_DELTA_DIR = "/tmp/slime-api-shim-deltas" -ROLLOUT_NUM_ENGINES = 1 - -modal = ModalConfig(gpu="H200", region="us") - - -class _Slime(SlimeConfig): - # Architecture from the slime model script (MLA + DeepSeek-MoE); no inline arch. - slime_model_script = "scripts/models/moonlight.sh" - - # Model + checkpoint. Bridge-mode HF load (no torch_dist conversion); matches - # the model the provider pool serves. - hf_checkpoint = "moonshotai/Moonlight-16B-A3B-Instruct" - ref_load = hf_checkpoint - megatron_to_hf_mode = "bridge" - - # Publish-only rollout through the provider front door (URL filled at launch). - actor_num_nodes = 2 # 2x8 H200 -> multinode RDMA + expert parallel - actor_num_gpus_per_node = 8 - colocate = False - rollout_num_gpus = 0 - rollout_num_gpus_per_engine = 1 - rollout_endpoint_url = None - - # M3 async-first: one-step off-policy (train_async); publish every step. - async_mode = True - update_weights_interval = 1 - - # Routing replay: replay the rollout engine's per-token expert routing during - # training. Needs the provider pool's --enable-return-routed-experts - # (configs/moonlight_hot_load.py); num_layers(27)/moe_router_topk(6) come from - # the model script so the rollout reshape stays consistent. - use_rollout_routing_replay = True - - # Staleness is gated by the announce_and_wait readiness BARRIER (the publish - # hook below), so the request hook leaves the per-request version pin OFF and - # only carries auth headers, retries, and affinity. - custom_rollout_request_hook_path = ( - "cookbook.standalone_rollouts.slime.hooks.rollout_request_weight_version_hook" - ) - api_shim_rollout_request_weight_version_mode = "none" - api_shim_rollout_request_retry_attempts = 240 - api_shim_rollout_request_retry_sleep = 1.0 - - # Disk-delta publish-only: write weight_v{N}/ to local disk; the pre-push hook - # copies to the S3 transport, POSTs the hot-load API, and blocks until all live - # replicas report the new version (readiness_threshold 1.0 by default). - update_weight_mode = "delta" - update_weight_transport = "disk" - update_weight_delta_encoding = "xor" - update_weight_delta_checksum = "xxh3-128" - update_weight_disk_dir = LOCAL_DELTA_DIR - api_shim_transport_root = TRANSPORT_ROOT - custom_delta_pre_push_path = ( - "cookbook.standalone_rollouts.slime.hooks.announce_and_wait" - ) - - # Data: dapo-math-17k (MoE-worthy), matching the slime_disagg Moonlight config. - prompt_data = f"{DATA_PATH}/dapo-math-17k/dapo-math-17k.jsonl" - input_key = "prompt" - label_key = "label" - apply_chat_template = True - rollout_shuffle = True - rm_type = "math" - eval_interval = None - - # Rollout - rollout_function_path = "slime.rollout.sglang_rollout.generate_rollout" - num_rollout = 5 - rollout_batch_size = 64 - rollout_max_response_len = 4096 - rollout_temperature = 1.0 - rollout_top_p = 1.0 - n_samples_per_prompt = 8 - global_batch_size = 128 - sglang_server_concurrency = 64 - use_fault_tolerance = False - - # Trainer parallelism (TP4/EP8 over 2x8; MLA/MoE arch flags from the script). - tensor_model_parallel_size = 4 - expert_model_parallel_size = 8 - expert_tensor_parallel_size = 1 - pipeline_model_parallel_size = 1 - context_parallel_size = 1 - sequence_parallel = True - use_dynamic_batch_size = True - max_tokens_per_gpu = 8192 - recompute_granularity = "full" - recompute_method = "uniform" - recompute_num_layers = 1 - attention_dropout = 0.0 - hidden_dropout = 0.0 - accumulate_allreduce_grads_in_fp32 = True - - # Optimizer (CPU offload keeps GPU state tiny for 3B active). - optimizer = "adam" - lr = 1e-6 - lr_decay_style = "constant" - weight_decay = 0.1 - adam_beta1 = 0.9 - adam_beta2 = 0.98 - optimizer_cpu_offload = True - overlap_cpu_optimizer_d2h_h2d = True - use_precision_aware_optimizer = True - - # Algorithm (GRPO) - advantage_estimator = "grpo" - eps_clip = 0.2 - eps_clip_high = 0.28 - use_kl_loss = True - kl_loss_coef = 0.0 - kl_loss_type = "low_var_kl" - entropy_coef = 0.0 - - def prepare_data(self) -> None: - from datasets import load_dataset - - # Raw DAPO-Math-17k has `prompt` (chat list) + reward_model.ground_truth; - # lift the answer into `label` and subset the ~1.8M-row artifact (see - # cookbook/slime_disagg/configs/moonlight_disagg.py). - ds = load_dataset("BytedTsinghua-SIA/DAPO-Math-17k", split="train") - ds = ds.shuffle(seed=42).select(range(min(50000, ds.num_rows))) - ds = ds.map(lambda ex: {"label": ex["reward_model"]["ground_truth"]}) - ds = ds.select_columns(["prompt", "label"]) - ds.to_json(f"{DATA_PATH}/dapo-math-17k/dapo-math-17k.jsonl") - - -slime = _Slime() diff --git a/cookbook/standalone_rollouts/slime/hooks.py b/cookbook/standalone_rollouts/slime/hooks.py deleted file mode 100644 index 8001312..0000000 --- a/cookbook/standalone_rollouts/slime/hooks.py +++ /dev/null @@ -1,444 +0,0 @@ -"""Optional SLIME trainer-side hooks for a provider hot-load API shim. - -These hooks run in SLIME, not in the rollout provider: - -1. Copy the completed SLIME disk-delta version directory to the mounted S3 transport. -2. Announce the new snapshot identity to the provider. -3. Poll the provider's pool readiness endpoint until enough replicas report - the target identity. - -They assume the provider can consume the files written by SLIME, or that a -provider-specific materialization step has been inserted before upload. -""" - -from __future__ import annotations - -import json -import logging -import os -import shutil -import time -import urllib.error -import urllib.request -from dataclasses import dataclass, replace -from pathlib import Path -from typing import Any - -from stitch.protocol import RolloutPoolState, parse_weight_identity, weight_identity - - -logger = logging.getLogger(__name__) - - -@dataclass(frozen=True) -class ShimConfig: - api_base_url: str - transport_root: Path | None = None - api_key: str | None = None - provider_model: str | None = None - provider_deployment: str | None = None - base_snapshot_identity: str = "base" - compression_format: str = "zstd" - checksum_format: str = "xxh3-128" - reset_prompt_cache: str = "new_session" - readiness_threshold: float = 1.0 - poll_timeout_seconds: float = 30 * 60 - poll_interval_seconds: float = 5.0 - # Run id: partitions the transport (`/weight_v{N}/`) and is folded - # into the single self-identifying pointer, so a fresh run's chain is isolated - # and a finished run's pointer can't fast-forward a new run's cold start. - # None = the run-less customer flat layout. - run_id: str | None = None - - @classmethod - def from_env(cls, args: Any | None = None) -> "ShimConfig": - transport = _setting( - args, "api_shim_transport_root", "STITCH_SHIM_TRANSPORT_ROOT", default="" - ) - return cls( - api_base_url=_setting( - args, - "api_shim_base_url", - "STITCH_SHIM_API_BASE_URL", - default=getattr(args, "rollout_endpoint_url", None), - required=True, - ).rstrip("/"), - transport_root=Path(transport) if transport else None, - api_key=_setting(args, "api_shim_api_key", "STITCH_SHIM_API_KEY"), - provider_model=_setting( - args, "api_shim_provider_model", "STITCH_SHIM_PROVIDER_MODEL" - ), - provider_deployment=_setting( - args, - "api_shim_provider_deployment", - "STITCH_SHIM_PROVIDER_DEPLOYMENT", - ), - base_snapshot_identity=_setting( - args, - "api_shim_base_snapshot_identity", - "STITCH_SHIM_BASE_SNAPSHOT_IDENTITY", - default="base", - ), - compression_format=_setting( - args, - "api_shim_compression_format", - "STITCH_SHIM_COMPRESSION_FORMAT", - default="zstd", - ), - checksum_format=_setting( - args, - "api_shim_checksum_format", - "STITCH_SHIM_CHECKSUM_FORMAT", - default=str(getattr(args, "update_weight_delta_checksum", "xxh3-128")), - ), - reset_prompt_cache=_setting( - args, - "api_shim_reset_prompt_cache", - "STITCH_SHIM_RESET_PROMPT_CACHE", - default="new_session", - ), - readiness_threshold=float( - _setting( - args, - "api_shim_readiness_threshold", - "STITCH_SHIM_READINESS_THRESHOLD", - default="1.0", - ) - ), - poll_timeout_seconds=float( - _setting( - args, - "api_shim_poll_timeout_seconds", - "STITCH_SHIM_POLL_TIMEOUT_SECONDS", - default="1800", - ) - ), - poll_interval_seconds=float( - _setting( - args, - "api_shim_poll_interval_seconds", - "STITCH_SHIM_POLL_INTERVAL_SECONDS", - default="5", - ) - ), - # run_id is NOT read here: a per-launch value can't ride an --api-shim-* - # arg (dropped by slime's parser) or a late env var (doesn't reach the - # Ray actor this hook runs in). announce_and_wait derives it from the - # version dir slime wrote (update_weight_disk_dir = /). - ) - - def identity_for_version(self, version: int) -> str: - return weight_identity(version) - - def previous_identity_for_version(self, version: int) -> str: - if int(version) <= 1: - return self.base_snapshot_identity - return self.identity_for_version(int(version) - 1) - - def transport_path_for_identity(self, identity: str) -> Path: - if self.transport_root is None: - raise RuntimeError("transport_root is not configured (api_shim_transport_root)") - base = self.transport_root / self.run_id if self.run_id else self.transport_root - return base / identity - - def readiness_identity(self, identity: str) -> str: - """The run-scoped snapshot identity the pool reports for readiness matching - (``/weight_vN``), so a replica on a finished run's same-numbered - version isn't miscounted as ready for this run.""" - return f"{self.run_id}/{identity}" if self.run_id else identity - - -def announce_and_wait(args: Any, version_dir: str, rollout_engines: list[Any]) -> None: - """SLIME ``custom_delta_pre_push_path`` hook (publish-only mode). - - slime publishes ``weight_v{N}/`` to a LOCAL disk dir (its disk-delta writer - uses atomic rename, which the S3 CloudBucketMount does not support). Rank 0 - copies that version dir to the shared transport the provider pool pulls from, - then signals the provider's customer hot-load API and blocks until the - elastic pool reports the version ready — so the next rollout only runs once - enough replicas serve the new weights. The POST drives the front door, which - owns the canonical ``latest`` pointer. - """ - - del rollout_engines - if _distributed_rank() not in (None, 0): - return - identity = Path(version_dir).name - version = parse_weight_identity(identity) - if version is None: - # _capture_baseline (the first update_weights) calls the pre-push hook - # with the disk-dir root, not a published version dir — nothing to - # announce yet. Only weight_v{N} dirs are hot-loaded. - return - cfg = ShimConfig.from_env(args) - # The run id is the partition dir slime wrote into - # (update_weight_disk_dir = /): derive it from the version dir - # rather than an env var or --api-shim-* arg, neither of which crosses into the - # Ray training actor this hook runs in. update_weight_disk_dir is a known slime - # arg, so the partition reliably reaches here. - run_id = Path(version_dir).parent.name or None - if run_id: - cfg = replace(cfg, run_id=run_id) - # slime's atomic-rename writer can't target the S3 mount (ENOSYS), so the - # version dir lives on local disk; copy it to the transport (PutObject) the - # provider pool reads before signalling the hot-load. - _copy_version_to_transport(Path(version_dir), cfg.transport_path_for_identity(identity)) - _post_hot_load( - cfg, - identity=identity, - previous_identity=cfg.previous_identity_for_version(version), - ) - target = cfg.readiness_identity(identity) - state = wait_until_ready(cfg, identity) - logger.info( - "Provider hot-load ready for %s: %s/%s replicas", - target, - state.ready_count(target_snapshot_identity=target), - len(state.replicas), - ) - - -def rollout_request_weight_version_hook( - args: Any, sample: Any, request: dict[str, Any] -) -> None: - """SLIME ``custom_rollout_request_hook_path`` hook. - - The publish hook has already decided which versions are usable by polling the - provider pool. This request hook turns that trainer decision into provider - admission control, so requests routed to lagging replicas fail with a - retryable 409 instead of producing unusable samples. - """ - - mode = str( - getattr( - args, - "api_shim_rollout_request_weight_version_mode", - os.environ.get("STITCH_SHIM_ROLLOUT_REQUEST_WEIGHT_VERSION_MODE", "exact"), - ) - ) - # PR #5's apply_rollout_request_hook builds request={url,payload,headers, - # max_retries,retry_sleep} with NO rollout_id — that is per-rollout context, - # not per-request (sample.index is the batch position, not the trainer step). - # So the step-based version pin only runs if a caller supplies a rollout_id. - # announce_and_wait already blocks the next rollout until the pool serves the - # target, so skipping the pin is safe; it is belt-and-suspenders admission - # control for lagging replicas. TODO: re-derive a per-request target under - # the PR #5 contract (e.g. the latest published version) if pinning is wanted. - rollout_id = request.get("rollout_id") - if mode != "none" and rollout_id is not None: - target_version = _rollout_request_target_version( - args, int(rollout_id), bool(request.get("evaluation", False)) - ) - if mode == "exact": - request["payload"]["weight_version"] = {"exact_version": target_version} - elif mode == "min": - request["payload"]["weight_version"] = {"min_required_version": target_version} - else: - raise ValueError( - f"Unsupported api_shim_rollout_request_weight_version_mode: {mode!r}" - ) - - # Generous retries on every rollout request so a cold/scaling/lagging replica - # (a transient 503, or a 409 weight-version reject) is retried rather than - # failing the rollout — the point of the elastic-pool readiness model. This - # also rides out a pool that is still warming when rollouts begin. - request["max_retries"] = int( - _setting( - args, - "api_shim_rollout_request_retry_attempts", - "STITCH_SHIM_ROLLOUT_REQUEST_RETRY_ATTEMPTS", - default="60", - ) - ) - request["retry_sleep"] = float( - _setting( - args, - "api_shim_rollout_request_retry_sleep", - "STITCH_SHIM_ROLLOUT_REQUEST_RETRY_SLEEP", - default="1.0", - ) - ) - - # The customer authenticates every request, and the provider front door - # enforces it on all routes (inference included), so the trainer must attach - # the same auth headers it sends to the hot-load API. - cfg = ShimConfig.from_env(args) - headers = dict(request.get("headers") or {}) - if cfg.api_key: - headers["Authorization"] = f"Bearer {cfg.api_key}" - if cfg.provider_model: - 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 - - -def _rollout_request_target_version(args: Any, rollout_id: int, evaluation: bool) -> int: - lag = int( - _setting( - args, - "api_shim_rollout_request_version_lag", - "STITCH_SHIM_ROLLOUT_REQUEST_VERSION_LAG", - default="0", - ) - ) - if evaluation: - lag = 0 - return max(0, int(rollout_id) - lag) - - -def wait_until_ready(cfg: ShimConfig, identity: str) -> RolloutPoolState: - # The pool reports current_snapshot_identity run-scoped (`/weight_vN`), - # so match against the same composite — otherwise a replica still serving a - # finished run's same-numbered version would falsely count as ready. - target = cfg.readiness_identity(identity) - deadline = time.monotonic() + cfg.poll_timeout_seconds - while True: - state = _get_hot_load_state(cfg) - if state.is_ready( - threshold=cfg.readiness_threshold, - target_snapshot_identity=target, - ): - return state - if time.monotonic() >= deadline: - ready = state.ready_count(target_snapshot_identity=target) - raise TimeoutError( - f"Timed out waiting for {target}: {ready}/{len(state.replicas)} " - f"replicas ready at threshold {cfg.readiness_threshold}; last_state={state.to_dict()}" - ) - logger.info( - "Waiting for %s readiness: %.3f < %.3f", - target, - state.readiness_fraction(target_snapshot_identity=target), - cfg.readiness_threshold, - ) - time.sleep(cfg.poll_interval_seconds) - - -def _post_hot_load(cfg: ShimConfig, *, identity: str, previous_identity: str) -> None: - payload = { - "identity": identity, - # The run id tells the front door which (possibly new) run this belongs to; - # a new run restarts the version space instead of being a rewind. - "run_id": cfg.run_id, - "incremental_snapshot_metadata": { - "previous_snapshot_identity": previous_identity, - "compression_format": cfg.compression_format, - "checksum_format": cfg.checksum_format, - }, - "reset_prompt_cache": cfg.reset_prompt_cache, - } - _request_json( - f"{cfg.api_base_url}/hot_load/v1/models/hot_load", - method="POST", - headers=_headers(cfg), - payload=payload, - ) - - -def _get_hot_load_state(cfg: ShimConfig) -> RolloutPoolState: - payload = _request_json( - f"{cfg.api_base_url}/hot_load/v1/models/hot_load", - method="GET", - headers=_headers(cfg), - ) - if not isinstance(payload, dict): - raise TypeError( - f"hot-load readiness response must be a JSON object, got {type(payload).__name__}" - ) - return RolloutPoolState.from_dict(payload) - - -def _request_json( - url: str, - *, - method: str, - headers: dict[str, str], - payload: dict[str, Any] | None = None, -) -> Any: - body = None if payload is None else json.dumps(payload).encode("utf-8") - request = urllib.request.Request(url, data=body, headers=headers, method=method) - try: - with urllib.request.urlopen(request, timeout=60) as response: - content = response.read() - except urllib.error.HTTPError as exc: - detail = exc.read().decode("utf-8", errors="replace") - raise RuntimeError( - f"{method} {url} failed with HTTP {exc.code}: {detail}" - ) from exc - if not content: - return {} - return json.loads(content.decode("utf-8")) - - -def _headers(cfg: ShimConfig) -> dict[str, str]: - headers = {"Content-Type": "application/json"} - if cfg.api_key: - headers["Authorization"] = f"Bearer {cfg.api_key}" - if cfg.provider_model: - headers["Provider-Model"] = cfg.provider_model - if cfg.provider_deployment: - headers["Provider-Deployment"] = cfg.provider_deployment - return headers - - -def _copy_version_to_transport(version_dir: Path, destination: Path) -> None: - """Copy a locally-published version dir to the (S3-mounted) transport. - - Uses plain writes, not rename: PutObject works on the CloudBucketMount while - rename does not. On a single-node trainer all ranks share the local FS, so - rank 0 copies the whole dir (every rank's shard + the index). - """ - files = sorted( - path - for path in version_dir.rglob("*") - if path.is_file() and not path.name.endswith(".tmp") - ) - for path in files: - target = destination / path.relative_to(version_dir) - target.parent.mkdir(parents=True, exist_ok=True) - if target.exists(): - target.unlink() # the transport may reject overwrites - with path.open("rb") as src, target.open("wb") as dst: - 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 deleted file mode 100644 index b280a45..0000000 --- a/cookbook/standalone_rollouts/slime/hooks_test.py +++ /dev/null @@ -1,143 +0,0 @@ -from __future__ import annotations - -import tempfile -import unittest -from argparse import Namespace -from pathlib import Path -from unittest import mock - -from cookbook.standalone_rollouts.slime import hooks -from stitch.protocol import RolloutPoolState, RolloutReplicaState - - -class ShimConfigTest(unittest.TestCase): - def test_uses_rollout_endpoint_as_base_url_fallback_and_zstd_defaults(self) -> None: - args = Namespace(rollout_endpoint_url="http://provider/") - - with mock.patch.dict("os.environ", {}, clear=True): - cfg = hooks.ShimConfig.from_env(args) - - self.assertEqual(cfg.api_base_url, "http://provider") - # The disk-delta branch ships zstd compression + xxh3-128 checksums. - self.assertEqual(cfg.compression_format, "zstd") - self.assertEqual(cfg.checksum_format, "xxh3-128") - self.assertIsNone(cfg.run_id) # run-less by default - - def test_run_id_scopes_transport_path_and_readiness(self) -> None: - cfg = hooks.ShimConfig(api_base_url="http://p", transport_root=Path("/t"), run_id="run-a") - self.assertEqual(cfg.transport_path_for_identity("weight_v000003"), Path("/t/run-a/weight_v000003")) - self.assertEqual(cfg.readiness_identity("weight_v000003"), "run-a/weight_v000003") - runless = hooks.ShimConfig(api_base_url="http://p", transport_root=Path("/t")) - self.assertEqual(runless.transport_path_for_identity("weight_v000003"), Path("/t/weight_v000003")) - self.assertEqual(runless.readiness_identity("weight_v000003"), "weight_v000003") - - -class AnnounceAndWaitTest(unittest.TestCase): - def test_copies_to_run_partition_then_announces_on_rank_zero(self) -> None: - with tempfile.TemporaryDirectory() as tmp: - root = Path(tmp) - # slime writes to update_weight_disk_dir = /, so the run - # id is the version dir's parent — announce_and_wait derives it there. - version_dir = root / "local" / "run-a" / "weight_v000003" - version_dir.mkdir(parents=True) - (version_dir / "model-00000-of-00001.safetensors").write_bytes(b"delta") - (version_dir / "model.safetensors.index.json").write_text("{}", encoding="utf-8") - transport = root / "transport" - - posted: list[tuple[str | None, str, str]] = [] - - def fake_post(cfg, *, identity, previous_identity) -> None: - posted.append((cfg.run_id, identity, previous_identity)) - - # The pool reports the run-scoped identity; readiness must match it. - ready = RolloutPoolState( - replicas=[ - RolloutReplicaState(readiness=True, current_snapshot_identity="run-a/weight_v000003") - ] - ) - args = Namespace(api_shim_base_url="http://provider", api_shim_transport_root=str(transport)) - with mock.patch.object(hooks, "_distributed_rank", return_value=0), mock.patch.object( - hooks, "_post_hot_load", fake_post - ), mock.patch.object(hooks, "_get_hot_load_state", return_value=ready): - hooks.announce_and_wait(args, str(version_dir), []) - - # Copied under the run partition (not the flat root), then announced - # with the derived run id (v2 as predecessor). - self.assertTrue((transport / "run-a" / "weight_v000003" / "model.safetensors.index.json").exists()) - self.assertFalse((transport / "weight_v000003").exists()) - self.assertEqual(posted, [("run-a", "weight_v000003", "weight_v000002")]) - - 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_and_wait(args, "/work/weight_v000003", []) - - self.assertEqual(posted, []) - - def test_skips_baseline_non_version_dir(self) -> None: - posted: list[int] = [] - 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", lambda *a, **k: posted.append(1) - ): - # _capture_baseline calls the hook with the disk-dir root, not a - # weight_v{N} dir — it must be a no-op, not an error. - hooks.announce_and_wait(args, "/mnt/stitch-s3-transport", []) - self.assertEqual(posted, []) - - -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 - # skip the version pin, and still apply session affinity. - args = Namespace( - api_shim_rollout_request_weight_version_mode="exact", - rollout_endpoint_url="http://provider", - ) - sample = Namespace(session_id="grp-1") - request = {"url": "u", "payload": {}, "headers": None, "max_retries": 60, "retry_sleep": 1.0} - with mock.patch.dict("os.environ", {}, clear=True): - hooks.rollout_request_weight_version_hook(args, sample, request) - self.assertNotIn("weight_version", request["payload"]) - self.assertEqual(request["headers"]["x-session-affinity"], "grp-1") - # Retries are applied even without a version pin (cold/scaling pool). - self.assertEqual(request["max_retries"], 60) - - def test_attaches_auth_headers(self) -> None: - # The front door enforces auth on inference too, so every rollout request - # must carry the provider auth headers. - args = Namespace( - api_shim_rollout_request_weight_version_mode="none", - rollout_endpoint_url="http://provider", - ) - sample = Namespace(session_id=None) - request = {"payload": {}, "headers": None} - env = { - "STITCH_SHIM_API_KEY": "k", - "STITCH_SHIM_PROVIDER_MODEL": "moonlight", - "STITCH_SHIM_PROVIDER_DEPLOYMENT": "prod", - } - with mock.patch.dict("os.environ", env, clear=True): - hooks.rollout_request_weight_version_hook(args, sample, request) - self.assertEqual(request["headers"]["Authorization"], "Bearer k") - self.assertEqual(request["headers"]["Provider-Model"], "moonlight") - self.assertEqual(request["headers"]["Provider-Deployment"], "prod") - - def test_pins_exact_when_rollout_id_supplied(self) -> None: - args = Namespace( - api_shim_rollout_request_weight_version_mode="exact", - api_shim_rollout_request_version_lag=0, - rollout_endpoint_url="http://provider", - ) - sample = Namespace(session_id=None) - request = {"payload": {}, "headers": None, "rollout_id": 3} - with mock.patch.dict("os.environ", {}, clear=True): - hooks.rollout_request_weight_version_hook(args, sample, request) - self.assertEqual(request["payload"]["weight_version"], {"exact_version": 3}) - - -if __name__ == "__main__": - unittest.main() diff --git a/cookbook/standalone_rollouts/slime/modal_train.py b/cookbook/standalone_rollouts/slime/modal_train.py deleted file mode 100644 index 673410e..0000000 --- a/cookbook/standalone_rollouts/slime/modal_train.py +++ /dev/null @@ -1,291 +0,0 @@ -"""Optional SLIME trainer for testing the standalone API-shim provider. - -This module imports the provider app and adds trainer-only functions to it. A -deploy of `cookbook.standalone_rollouts.modal_serve` publishes only the rollout -provider; a deploy of this module publishes both the provider and the trainer. -""" - -from __future__ import annotations - -import importlib -import os -import subprocess -import tempfile -import uuid -from pathlib import Path - -import modal -import modal.experimental - -from cookbook.standalone_rollouts import modal_serve as provider_app -from cookbook.slime_disagg import helpers -from cookbook.slime_disagg.configs.base import ( - CHECKPOINTS_PATH, - DATA_PATH, - HF_CACHE_PATH, - SlimeConfig, -) - - -TRAINER_CONFIG = os.environ.get( - "TRAINER_CONFIG", os.environ.get("EXPERIMENT_CONFIG", "moonlight_slime_trainer") -) -exp = importlib.import_module( - f"cookbook.standalone_rollouts.slime.configs.{TRAINER_CONFIG}" -) - -modal_cfg = exp.modal -slime_cfg = exp.slime - -APP_NAME = provider_app.APP_NAME -MODEL_NAME = slime_cfg.hf_checkpoint -N_TRAIN_NODES = helpers.training_nodes(slime_cfg) - -MINUTES = 60 -RAY_PORT = 6379 - -SLIME_IMAGE_TAG = "slimerl/slime:nightly-dev-20260527a" -SLIME_ROOT = "/root/slime" -SLIME_REPO_URL = "https://github.com/modal-projects/slime.git" -# Pin to an exact commit (see cookbook/slime_disagg/modal_train.py): the cached -# clone layer otherwise leaves the container on a stale slime. This is PR #5 -# head (disaggregated-rollout): disk-delta publish-only + rollout_endpoint_url + -# custom_rollout_request_hook_path. -SLIME_REPO_REF = "ebfe153949b1a69c39e92f947ed5d475166dd724" # incl. deepseekv3 router-dtype export fix + per-request rollout hook - -trainer_image = ( - modal.Image.from_registry(SLIME_IMAGE_TAG) - .entrypoint([]) - .run_commands(f"rm -rf {HF_CACHE_PATH}") - .run_commands( - f"rm -rf {SLIME_ROOT}" - 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}" - ) - # The base image installs megatron-core as a PEP 660 *strict* editable that - # hides megatron.training (which slime's megatron backend imports). Reinstall - # in compat mode so the whole source tree is importable (mirrors slime_disagg). - .run_commands( - "cd /root/Megatron-LM" - " && python3 -m pip install --no-deps -e . --config-settings editable_mode=compat" - ) - .pip_install( - # Disk-delta encode side (slime.utils.disk_delta) compresses with - # zstandard and checksums with xxhash (xxh3-128 default) / blake3. - "zstandard", - "xxhash", - "blake3", - ) - .env( - { - "EXPERIMENT_CONFIG": TRAINER_CONFIG, - "TRAINER_CONFIG": TRAINER_CONFIG, - "HF_XET_HIGH_PERFORMANCE": "1", - "HF_HUB_ENABLE_HF_TRANSFER": "1", - "STITCH_SHIM_TRANSPORT_ROOT": str(provider_app.S3_TRANSPORT_MOUNT_PATH), - } - ) - .add_local_python_source("stitch") - .add_local_dir( - Path(__file__).parents[2], - remote_path="/root/cookbook", - ignore=["**/__pycache__"], - ) -) - -# Dev iteration: SLIME_LOCAL_DIR overlays a local slime checkout onto the image's -# cloned fork (installed editable at /root/slime), so fork edits take effect on -# container start with no image rebuild. Unset by default. -if slime_local := os.environ.get("SLIME_LOCAL_DIR"): - trainer_image = trainer_image.add_local_dir( - slime_local, - remote_path=SLIME_ROOT, - ignore=[".git", "**/__pycache__", "**/*.pyc"], - ) - -hf_cache_volume = provider_app.hf_cache_volume -data_volume = modal.Volume.from_name("slime-data", create_if_missing=True) -checkpoints_volume = modal.Volume.from_name("slime-checkpoints", create_if_missing=True) - -train_volumes = { - str(HF_CACHE_PATH): hf_cache_volume, - str(DATA_PATH): data_volume, - str(CHECKPOINTS_PATH): checkpoints_volume, - str(provider_app.S3_TRANSPORT_MOUNT_PATH): provider_app.s3_transport_mount, -} -reloadable_train_volumes = (hf_cache_volume, data_volume, checkpoints_volume) - -app = provider_app.app - - -@app.cls( - image=trainer_image, - gpu=f"{modal_cfg.gpu}:{slime_cfg.actor_num_gpus_per_node}", - memory=modal_cfg.memory, - cloud=modal_cfg.cloud, - region=modal_cfg.region, - volumes=train_volumes, - secrets=[modal.Secret.from_name(exp.SHIM_SECRET_NAME)], - timeout=24 * 60 * MINUTES, - startup_timeout=20 * MINUTES, - scaledown_window=30 * MINUTES, - experimental_options={"efa_enabled": True}, - include_source=False, -) -@modal.experimental.clustered(N_TRAIN_NODES, rdma=True) -class Trainer: - @modal.enter() - def start_ray(self) -> None: - rank, master_addr, my_ip = helpers.get_modal_cluster_context(N_TRAIN_NODES) - self.rank = rank - os.environ.update( - { - "SLIME_HOST_IP": my_ip, - "SGLANG_HOST_IP": my_ip, - "HOST_IP": my_ip, - "MASTER_ADDR": master_addr, - "RAY_ADDRESS": f"{master_addr}:{RAY_PORT}", - "no_proxy": f"127.0.0.1,{master_addr},{my_ip}", - "NO_PROXY": f"127.0.0.1,{master_addr},{my_ip}", - **_trainer_environment(), - } - ) - if rank == 0: - helpers.start_ray_head(my_ip, N_TRAIN_NODES, ray_port=RAY_PORT) - else: - helpers.start_ray_worker(my_ip, master_addr, ray_port=RAY_PORT) - - @modal.method() - def train(self, experiment: str, payload: dict) -> None: - for volume in reloadable_train_volumes: - volume.reload() - if self.rank != 0: - return - - _validate_trainer_environment() - provider_url = ( - os.environ.get("STITCH_SHIM_API_BASE_URL") or provider_app.frontdoor_url() - ).rstrip("/") - os.environ["STITCH_SHIM_API_BASE_URL"] = provider_url - cfg = SlimeConfig.from_payload(payload) - run = importlib.import_module( - f"cookbook.standalone_rollouts.slime.configs.{experiment}" - ) - if helpers.training_nodes(cfg) != N_TRAIN_NODES: - raise ValueError( - f"experiment {experiment!r} needs {helpers.training_nodes(cfg)} node(s) but this app " - f"was deployed with {N_TRAIN_NODES}; redeploy with TRAINER_CONFIG={experiment}" - ) - if cfg.environment != slime_cfg.environment: - print( - f"WARNING: experiment {experiment!r} changes `environment`, which only " - f"takes effect after a redeploy restarts the Ray cluster." - ) - - cfg.rollout_endpoint_url = provider_url - cfg.api_shim_base_url = provider_url - # Fresh run id per launch, carried on update_weight_disk_dir — a *known* - # slime arg, so it reaches the publish hook running in the Ray training - # actor. (A per-launch env var doesn't: Ray workers start in @modal.enter() - # before train() runs, so they don't inherit this env; and an unknown - # --api-shim-run-id arg is dropped by slime's parser.) slime writes this - # run's chain to //weight_v{N}/, and announce_and_wait - # derives the run id from that dir to scope the S3 transport + pointer. - run_id = uuid.uuid4().hex[:12] - cfg.update_weight_disk_dir = f"{cfg.update_weight_disk_dir}/{run_id}" - cfg.custom_config_path = _custom_config( - cfg, rollout_num_engines=getattr(run, "ROLLOUT_NUM_ENGINES", 1) - ) - helpers.prepare_slime_config(cfg, tempfile.mkdtemp()) - cmd = helpers.build_train_cmd(cfg, SLIME_ROOT) - - print( - f"Training {experiment}: nodes={N_TRAIN_NODES}, rollout_endpoint={cfg.rollout_endpoint_url}" - ) - print(f"Command: {cmd}") - subprocess.run(["bash", "-lc", cmd], check=True) - - -@app.function( - image=trainer_image, - volumes={str(DATA_PATH): data_volume}, - timeout=2 * 60 * MINUTES, - secrets=[modal.Secret.from_name(exp.HF_SECRET_NAME)], - include_source=False, -) -def prepare_dataset() -> None: - data_volume.reload() - slime_cfg.prepare_data() - data_volume.commit() - - -@app.local_entrypoint() -def launch_train(experiment: str = TRAINER_CONFIG) -> None: - from modal.exception import NotFoundError - - run = importlib.import_module( - f"cookbook.standalone_rollouts.slime.configs.{experiment}" - ) - trainer = modal.Cls.from_name(APP_NAME, Trainer.__name__)() - try: - call = trainer.train.spawn(experiment, run.slime.to_payload()) - except NotFoundError: - raise SystemExit( - f"App {APP_NAME!r} is not deployed. Run:\n" - f" uv run --extra modal modal deploy -m cookbook.standalone_rollouts.slime.modal_train" - ) - print(f"Spawned train({experiment!r}) on {APP_NAME}: {call.object_id}") - - -@app.local_entrypoint() -def print_trainer_secret_template() -> None: - print( - "\n".join( - [ - f"modal secret create {exp.SHIM_SECRET_NAME} \\", - " STITCH_SHIM_API_KEY=... \\", - " STITCH_SHIM_PROVIDER_MODEL=moonlight \\", - " STITCH_SHIM_PROVIDER_DEPLOYMENT=rollout-prod \\", - " STITCH_SHIM_BASE_SNAPSHOT_IDENTITY=base", - ] - ) - ) - - -def _validate_trainer_environment() -> None: - required = ("STITCH_SHIM_TRANSPORT_ROOT",) - missing = [name for name in required if not os.environ.get(name)] - if missing: - raise RuntimeError( - f"Missing required runtime setting(s): {', '.join(missing)}. " - f"Create/update the {exp.SHIM_SECRET_NAME!r} Modal secret or redeploy the app." - ) - transport_root = Path(os.environ["STITCH_SHIM_TRANSPORT_ROOT"]) - if not transport_root.exists(): - raise RuntimeError(f"Transport root is not mounted: {transport_root}") - - -def _trainer_environment() -> dict[str, str]: - env = {str(key): str(value) for key, value in slime_cfg.environment.items()} - existing = env.get("PYTHONPATH") or os.environ.get("PYTHONPATH", "") - paths = ["/root/Megatron-LM", "/root"] - paths.extend(path for path in existing.split(os.pathsep) if path) - env["PYTHONPATH"] = os.pathsep.join(dict.fromkeys(paths)) - return env - - -def _custom_config(cfg: SlimeConfig, *, rollout_num_engines: int) -> dict: - current = getattr(cfg, "custom_config_path", None) - if isinstance(current, dict): - data = dict(current) - elif current: - raise ValueError( - "This trainer app expects custom_config_path to be unset or a dict before materialization." - ) - else: - data = {} - data["rollout_num_engines"] = int(rollout_num_engines) - return data 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 diff --git a/pyproject.toml b/pyproject.toml index 71ecd54..5746b2b 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -6,9 +6,12 @@ requires-python = ">=3.11" dependencies = [] [project.optional-dependencies] -# providers/modal.py uses the Modal SDK and httpx (rollout-pool wake). +# The Modal instances (stores/modal_volume.py, pools/modal_flash.py) use the Modal +# SDK and httpx (the pool wake). modal = ["modal>=1.5.1.dev9", "httpx"] -sglang = ["fastapi", "httpx", "uvicorn"] +# The sidecar runs the versioned proxy under uvicorn (Modal images install it +# explicitly); the packaged service code needs only fastapi + httpx. +sglang = ["fastapi", "httpx"] [dependency-groups] dev = ["pytest"] @@ -22,5 +25,6 @@ packages = ["src/stitch"] exclude = ["*_test.py"] [tool.pytest.ini_options] -testpaths = ["src"] +testpaths = ["src", "cookbook"] # google-style: foo_test.py sits beside foo.py python_files = ["*_test.py"] +pythonpath = ["."] # so the co-located cookbook hook test imports the top-level cookbook/ package diff --git a/src/stitch/__init__.py b/src/stitch/__init__.py index ed24ac1..cdcc687 100644 --- a/src/stitch/__init__.py +++ b/src/stitch/__init__.py @@ -1,27 +1 @@ """Disaggregated rollout protocol and integration helpers.""" - -from stitch.protocol import ( - Artifact, - EngineAdapter, - RolloutPoolState, - RolloutReplicaState, - SyncState, - VersionManifest, - WeightVersionPolicy, - read_latest, - version_dir, - write_latest, -) - -__all__ = [ - "Artifact", - "EngineAdapter", - "RolloutPoolState", - "RolloutReplicaState", - "SyncState", - "VersionManifest", - "WeightVersionPolicy", - "read_latest", - "version_dir", - "write_latest", -] diff --git a/src/stitch/bulletin.py b/src/stitch/bulletin.py deleted file mode 100644 index 473a7ec..0000000 --- a/src/stitch/bulletin.py +++ /dev/null @@ -1,132 +0,0 @@ -"""Bulletin board storage interfaces.""" - -from __future__ import annotations - -import asyncio -import inspect -from collections.abc import Callable -from pathlib import Path -from typing import Any, Protocol - -from stitch.protocol import ( - VersionManifest, - atomic_write_text, - format_snapshot_identity, - parse_snapshot_identity, - read_latest, - version_dir, - weight_identity, - write_latest, -) - - -class BulletinBoard(Protocol): - root: Path - - async def refresh(self) -> None: ... - - def read_latest(self) -> tuple[str | None, int]: ... - - def write_latest(self, run_id: str | None, version: int) -> None: ... - - def version_dir(self, run_id: str | None, version: int) -> Path: ... - - def read_manifest(self, run_id: str | None, version: int) -> VersionManifest: ... - - def publish_manifest( - self, - manifest: VersionManifest, - version_path: str | Path | None = None, - *, - run_id: str | None = None, - ) -> None: ... - - -class FilesystemBulletinBoard: - """Filesystem-backed bulletin board. - - A refresh callback can be provided by storage providers such as Modal Volume. - - Two on-disk layouts are supported: - - - ``"stitch"`` (default): the engine-neutral protocol — ``versions/`` -nested - version dirs, a JSON ``latest.json`` pointer, and a ``manifest.json`` per - version. - - ``"slime"``: slime's native disk-delta publish output (and the customer's - object-store layout). Each run's chain lives under ``/weight_v{N:06d}/`` - and the single ``latest`` pointer holds the self-identifying snapshot identity - ``/weight_v{N:06d}`` (a bare ``weight_v{N:06d}`` for the degenerate - run-less layout). The manifest is read from each version's - ``model.safetensors.index.json``. Run-id partitioning is what makes sequential - runs collision-free: a new run writes a fresh ``/`` chain and the - pointer moves to it (a new run is not a rewind), so a finished run's chain can - never be overwritten or fast-forward a cold start. The pool reconciles against - ``latest``; the front door (or the publish hook) advances it. - """ - - def __init__( - self, - root: str | Path, - refresh: Callable[[], Any] | None = None, - *, - layout: str = "stitch", - ) -> None: - if layout not in ("stitch", "slime"): - raise ValueError(f"unknown bulletin layout {layout!r}") - self.root = Path(root) - self._refresh = refresh - self.layout = layout - - async def refresh(self) -> None: - if self._refresh is None: - return - result = await asyncio.to_thread(self._refresh) - if inspect.isawaitable(result): - await result - - def read_latest(self) -> tuple[str | None, int]: - """The active snapshot pointer as ``(run_id, version)``. - - slime: parse ``/weight_v{N}`` (or a bare/legacy pointer) from the - ``latest`` file; missing/empty/unparseable -> ``(None, 0)``. stitch: the - JSON pointer is run-less, so ``(None, )``. - """ - if self.layout == "slime": - path = self.root / "latest" - if not path.exists(): - return (None, 0) - return parse_snapshot_identity(path.read_text(encoding="utf-8")) - return (None, read_latest(self.root)) - - def write_latest(self, run_id: str | None, version: int) -> None: - if self.layout == "slime": - atomic_write_text(self.root / "latest", format_snapshot_identity(run_id, version)) - else: - write_latest(self.root, 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 - return base / weight_identity(version) - return version_dir(self.root, version) - - def read_manifest(self, run_id: str | None, version: int) -> VersionManifest: - if self.layout == "slime": - return VersionManifest.from_slime_index(self.version_dir(run_id, version)) - return VersionManifest.read(self.version_dir(run_id, version) / "manifest.json") - - def publish_manifest( - self, - manifest: VersionManifest, - version_path: str | Path | None = None, - *, - run_id: str | None = None, - ) -> None: - if self.layout == "slime": - # The version dir + index.json already exist (written by slime/the - # uploader under /); publishing is only advancing the pointer. - self.write_latest(run_id, manifest.version) - return - target = Path(version_path) if version_path is not None else self.version_dir(run_id, manifest.version) - manifest.write(target / "manifest.json") - self.write_latest(run_id, manifest.version) diff --git a/src/stitch/bulletin_test.py b/src/stitch/bulletin_test.py deleted file mode 100644 index 7380ee6..0000000 --- a/src/stitch/bulletin_test.py +++ /dev/null @@ -1,129 +0,0 @@ -from __future__ import annotations - -import json -import tempfile -import unittest -from pathlib import Path - -from stitch.bulletin import FilesystemBulletinBoard -from stitch.protocol import ( - VersionManifest, - format_snapshot_identity, - parse_snapshot_identity, -) - - -class SnapshotIdentityTest(unittest.TestCase): - def test_format_round_trips_with_and_without_run_id(self) -> None: - self.assertEqual(format_snapshot_identity("run-a", 5), "run-a/weight_v000005") - self.assertEqual(format_snapshot_identity(None, 5), "weight_v000005") - self.assertEqual(parse_snapshot_identity("run-a/weight_v000005"), ("run-a", 5)) - - def test_parse_tolerates_bare_legacy_and_garbage(self) -> None: - self.assertEqual(parse_snapshot_identity("weight_v000005"), (None, 5)) # bare - self.assertEqual(parse_snapshot_identity("000005"), (None, 5)) # legacy flat pointer - self.assertEqual(parse_snapshot_identity(""), (None, 0)) # missing -> not-ready - self.assertEqual(parse_snapshot_identity("garbage"), (None, 0)) # unparseable -> not-ready - - -class SlimeLayoutBulletinTest(unittest.TestCase): - """The slime/customer layout: per-run ``/weight_v{N}/`` chains and a - single self-identifying ``latest`` pointer ``/weight_v{N}``.""" - - def _write_version(self, base: Path, version: int, prev: int) -> Path: - vdir = base / f"weight_v{version:06d}" - vdir.mkdir(parents=True) - (vdir / "model.safetensors.index.json").write_text( - json.dumps( - { - "metadata": { - "version": f"{version:06d}", - "base_version": f"{prev:06d}", - "delta_encoding": "xor", - "compression_format": "zstd", - "checksum_format": "xxh3-128", - }, - "weight_map": {"w": "model-00001-of-00001.safetensors"}, - } - ), - encoding="utf-8", - ) - return vdir - - def test_run_id_layout_reads_chain(self) -> None: - with tempfile.TemporaryDirectory() as tmp: - root = Path(tmp) - self._write_version(root / "run-a", 1, 0) - self._write_version(root / "run-a", 2, 1) - (root / "latest").write_text("run-a/weight_v000002", encoding="utf-8") - - board = FilesystemBulletinBoard(root, layout="slime") - - self.assertEqual(board.read_latest(), ("run-a", 2)) - self.assertEqual(board.version_dir("run-a", 2), root / "run-a" / "weight_v000002") - manifest = board.read_manifest("run-a", 2) - self.assertEqual(manifest.version, 2) - self.assertEqual(manifest.base_version, 1) - self.assertEqual(manifest.compression_format, "zstd") - self.assertEqual(manifest.checksum_format, "xxh3-128") - - def test_read_latest_absent_is_none_zero(self) -> None: - with tempfile.TemporaryDirectory() as tmp: - self.assertEqual(FilesystemBulletinBoard(Path(tmp), layout="slime").read_latest(), (None, 0)) - - def test_write_latest_run_id_round_trips(self) -> None: - with tempfile.TemporaryDirectory() as tmp: - root = Path(tmp) - board = FilesystemBulletinBoard(root, layout="slime") - board.write_latest("run-a", 7) - self.assertEqual((root / "latest").read_text(encoding="utf-8"), "run-a/weight_v000007") - self.assertEqual(board.read_latest(), ("run-a", 7)) - - def test_write_latest_none_run_id_writes_bare_identity(self) -> None: - with tempfile.TemporaryDirectory() as tmp: - root = Path(tmp) - board = FilesystemBulletinBoard(root, layout="slime") - board.write_latest(None, 7) - self.assertEqual((root / "latest").read_text(encoding="utf-8"), "weight_v000007") - self.assertEqual(board.read_latest(), (None, 7)) - - def test_legacy_bare_pointer_parses_runless(self) -> None: - # A pre-run-id deployment left `latest` = "000005"; it must parse as a - # run-less pointer (None, 5), never be misread as a phantom run. - with tempfile.TemporaryDirectory() as tmp: - root = Path(tmp) - (root / "latest").write_text("000005", encoding="utf-8") - self.assertEqual(FilesystemBulletinBoard(root, layout="slime").read_latest(), (None, 5)) - - def test_publish_manifest_advances_run_pointer(self) -> None: - with tempfile.TemporaryDirectory() as tmp: - root = Path(tmp) - self._write_version(root / "run-a", 1, 0) - board = FilesystemBulletinBoard(root, layout="slime") - board.publish_manifest( - VersionManifest(version=1, base_version=0, backend="disk_delta", load_format="auto"), - run_id="run-a", - ) - self.assertEqual(board.read_latest(), ("run-a", 1)) - - def test_unknown_layout_rejected(self) -> None: - with self.assertRaises(ValueError): - FilesystemBulletinBoard("/tmp", layout="bogus") - - -class StitchLayoutBulletinTest(unittest.TestCase): - """The engine-neutral protocol layout is run-less: run_id is always None.""" - - def test_publish_and_read_round_trip_run_id_none(self) -> None: - with tempfile.TemporaryDirectory() as tmp: - root = Path(tmp) - board = FilesystemBulletinBoard(root) # default layout="stitch" - manifest = VersionManifest(version=3, base_version=2, backend="fake", load_format="noop") - board.publish_manifest(manifest) - self.assertEqual(board.read_latest(), (None, 3)) - self.assertEqual(board.version_dir(None, 3), root / "versions" / "weight_v000003") - self.assertEqual(board.read_manifest(None, 3).version, 3) - - -if __name__ == "__main__": - unittest.main() diff --git a/src/stitch/engines/base.py b/src/stitch/engines/base.py new file mode 100644 index 0000000..b611b8c --- /dev/null +++ b/src/stitch/engines/base.py @@ -0,0 +1,63 @@ +"""The ``Engine`` port — a client to one inference engine. + +Instances subclass this base and override every method: ``engines/sglang.py`` (sglang). +Add vllm as a new subclass. +""" + +from __future__ import annotations + +from typing import Any + +from stitch.versions import VersionManifest, VersionRef + + +class Engine: + """Drives one engine and translates the version protocol; the heavy weight + apply runs inside the engine, not here. Subclasses override every method.""" + + async def stage(self, manifest: VersionManifest, source_dir: str) -> None: + """Bring the local checkpoint to ``manifest.ref``: seed from the nearest FULL + anchor, then replay deltas forward. May run while the engine serves.""" + raise NotImplementedError + + async def commit(self, ref: VersionRef) -> None: + """Reload the staged checkpoint into the serving weights — the gate covers only this.""" + raise NotImplementedError + + async def flush(self) -> None: + """Evict cached state (KV / radix tree). Called before commit in quiesce mode.""" + raise NotImplementedError + + async def pause(self) -> None: + """Pause the scheduler in place (in_place commit); in-flight requests stay resident.""" + raise NotImplementedError + + async def resume(self) -> None: + """Resume the scheduler after a pause.""" + raise NotImplementedError + + async def reset(self) -> None: + """Reseed the local checkpoint to the engine's boot base.""" + raise NotImplementedError + + async def prefetch(self) -> None: + """Optional: seed the host-local checkpoint from the engine's base ahead of the first + stage(), so stage only applies the delta instead of copying the full base off the + critical path. Default no-op — an engine with no host-local checkpoint needs nothing.""" + return + + def stamp_request(self, request: dict[str, Any], served: VersionRef) -> None: + """Namespace a request to the version it's served on so requests from different + versions can't share KV prefixes (engine-specific, e.g. sglang's extra_key). + Mutates ``request`` in place.""" + raise NotImplementedError + + def stamp_response(self, response: dict[str, Any], served: VersionRef, current: VersionRef) -> None: + """Record which version generated a response, in the engine's response shape + (e.g. sglang's meta_info vs OpenAI top-level). Mutates ``response`` in place.""" + raise NotImplementedError + + def base_url(self) -> str: + """The engine's base HTTP URL — the proxy forwards to it, and the engine's own + stage/commit calls target it.""" + raise NotImplementedError diff --git a/src/stitch/engines/sglang.py b/src/stitch/engines/sglang.py index 6486dd3..12425ac 100644 --- a/src/stitch/engines/sglang.py +++ b/src/stitch/engines/sglang.py @@ -1,170 +1,150 @@ -"""SGLang rollout engine adapters.""" +"""``SGLangEngine`` — the ``Engine`` instance for a single sglang server. + +The weight apply lives inside sglang (``weight_sync/local_checkpoint``): ``stage`` +POSTs ``/pull_weights``, which chain-replays deltas from the applied checkpoint with +per-tensor checksum verification (reseeding from base on corruption), and ``commit`` +reloads the materialized checkpoint via ``/update_weights_from_disk``. This client +only drives those endpoints and translates the version protocol — it holds no +host-side decoder and imports no trainer package. +""" from __future__ import annotations import asyncio -import logging -import time -from collections.abc import Callable -from dataclasses import dataclass +import shutil from pathlib import Path - -from stitch.protocol import VersionManifest - - -logger = logging.getLogger(__name__) -EXTRA_KEY_DELIMITER = ";" - - -def compose_extra_key( - version: int, user_extra_key: str | None = None, run_id: str | None = None -) -> str: - """Compose a weight-version-namespaced SGLang ``extra_key``. - - The version segment sits at a fixed position (the prefix) and is - delimiter-terminated, so it parses unambiguously regardless of the user - key's content. sglang appends ``lora_id`` to ``extra_key`` with no - delimiter, so the version must never be parsed from the right. - Examples: ``wv7;`` (no user key), ``wv7;my-key``. This is an SGLang - radix-cache namespacing concern, not part of the engine-neutral protocol. - - ``run_id`` is folded into the key *content* (after the delimiter), so two - runs that both restart version numbering at 1 get distinct radix namespaces - (``wv1;run-a/`` vs ``wv1;run-b/``) and a stale cross-run request can never - reuse another run's same-numbered KV. ``run_id=None`` keeps the bare form. - """ - run_segment = f"{run_id}/" if run_id else "" - return f"wv{int(version)}{EXTRA_KEY_DELIMITER}{run_segment}{user_extra_key or ''}" - - -def parse_extra_key_version(extra_key: str) -> int | None: - """Inverse of :func:`compose_extra_key`. None for non-composed keys.""" - if not extra_key.startswith("wv"): - return None - head, delim, _rest = extra_key.partition(EXTRA_KEY_DELIMITER) - if not delim or not head[2:].isdigit(): - return None - return int(head[2:]) - - -@dataclass -class SGLangDiskDeltaAdapter: - """Applies disk-delta weight versions to one local SGLang server. - - The delta is applied *host-side*: slime's ``disk_delta`` patches a local - full HF checkpoint in place (chain-replayed from the base, per-tensor - checksum-verified, with a base-version precondition), and the engine then - reloads that checkpoint through the ordinary ``update_weights_from_disk`` - path. The engine carries no delta receiver, so no ``load_format`` / ``files`` - delta payload is sent — that is all the new disk-delta slime branch needs. - - ``apply_deltas`` / ``init_local_checkpoint`` are injectable so the adapter is - testable without slime/numpy installed; by default they bind lazily to - ``slime.utils.disk_delta``. - """ - - upstream_url: str - local_checkpoint_dir: str - base_checkpoint_dir: str - backend: str = "disk_delta" - apply_deltas: Callable[[str, str, int], None] | None = None - init_local_checkpoint: Callable[[str, str], None] | None = None - - def __post_init__(self) -> None: - self.upstream_url = self.upstream_url.rstrip("/") - - def _apply_deltas(self) -> Callable[[str, str, int], None]: - if self.apply_deltas is not None: - return self.apply_deltas - from slime.utils.disk_delta import apply_deltas - - return apply_deltas - - def _init_local_checkpoint(self) -> Callable[[str, str], None]: - if self.init_local_checkpoint is not None: - return self.init_local_checkpoint - from slime.utils.disk_delta import init_local_checkpoint - - return init_local_checkpoint - - async def prepare(self) -> None: - """Materialize the host-local full checkpoint from the base once - (idempotent) so later deltas apply on top of it in place. Run at startup; - blocking copy, so it is offloaded to a thread.""" - await asyncio.to_thread( - self._init_local_checkpoint(), self.local_checkpoint_dir, self.base_checkpoint_dir +from typing import Any + +from stitch.engines.base import Engine +from stitch.versions import VersionManifest, VersionRef + + +class SGLangEngine(Engine): + def __init__( + self, + base_url: str, + local_checkpoint_dir: str, + *, + control_timeout: float = 120.0, + flush_on_commit: bool = False, + ) -> None: + self._base_url = base_url.rstrip("/") + self.local_checkpoint_dir = local_checkpoint_dir + self._control_timeout = control_timeout + self._flush_on_commit = flush_on_commit + + def base_url(self) -> str: + return self._base_url + + async def stage(self, manifest: VersionManifest, source_dir: str) -> None: + # source_dir is the target version's dir; its parent is the root of weight_v* + # dirs the pull walks. Disk-only, so it runs while the engine serves. + await self._post( + "/pull_weights", + { + "local_checkpoint_dir": self.local_checkpoint_dir, + "source_dir": str(Path(source_dir).parent), + "target_version": manifest.ref.version, + }, + timeout=None, + action="weight pull", ) - async def reset(self) -> None: - """Discard the local checkpoint and re-materialize the base from scratch. - - Used on a run switch: the new run's chain forks at base, so the locally - patched checkpoint must be thrown away before replaying the new chain. The - sync manager invokes this under the commit gate (engine paused), so no - request decodes across the wipe; ``prepare`` then rebuilds a complete base - (a partial wipe from a crash is re-seeded by ``init_local_checkpoint``). - """ - import shutil + async def prefetch(self) -> None: + # Seed the host-local checkpoint from the served base now (target_version=0): the + # receiver copies its own model_path into local_checkpoint_dir and applies no deltas, + # so the first real stage() only applies the delta rather than paying the full base + # copy. Disk-only (serves throughout) and flock-serialized + idempotent with a + # concurrent stage. source_dir is unused for the base seed (a placeholder here). + await self._post( + "/pull_weights", + {"local_checkpoint_dir": self.local_checkpoint_dir, "source_dir": self.local_checkpoint_dir, "target_version": 0}, + timeout=None, + action="base prefetch", + ) - await asyncio.to_thread(shutil.rmtree, self.local_checkpoint_dir, ignore_errors=True) - await self.prepare() + async def commit(self, ref: VersionRef) -> None: + # flush_cache defaults off: the reconciler owns flushing — it calls flush() + # itself before a quiesce reload, and in_place deliberately keeps in-flight KV. + # flush_on_commit is the escape hatch for a setup that wants the reload to evict. + await self._post( + "/update_weights_from_disk", + {"model_path": self.local_checkpoint_dir, "weight_version": str(ref.version), + "flush_cache": self._flush_on_commit}, + timeout=None, + action="weight update", + ) - async def flush_cache(self) -> None: - import httpx + async def flush(self) -> None: + await self._get("/flush_cache", ok=(200, 404)) - async with httpx.AsyncClient(timeout=120.0, trust_env=False) as client: - resp = await client.get(f"{self.upstream_url}/flush_cache") - if resp.status_code not in (200, 404): - resp.raise_for_status() + async def pause(self) -> None: + await self._post("/pause_generation", {"mode": "in_place"}, timeout=self._control_timeout) - async def pause_generation(self) -> None: - """Pause the scheduler loop in place: in-flight requests stay resident - and resume decoding on their existing KV after continue_generation.""" - import httpx + async def resume(self) -> None: + await self._post("/continue_generation", {}, timeout=self._control_timeout) - async with httpx.AsyncClient(timeout=120.0, trust_env=False) as client: - resp = await client.post(f"{self.upstream_url}/pause_generation", json={"mode": "in_place"}) - resp.raise_for_status() + async def reset(self) -> None: + # Re-materialize base and reload it into the engine, so a run switch that lands at + # v0 serves base on the GPU -- not the previous run's weights under the new + # (run, 0) identity. Runs under the commit gate (engine paused). Ref: stitch#32. + await asyncio.to_thread(shutil.rmtree, self.local_checkpoint_dir, ignore_errors=True) + await self.prefetch() # reseed base (target_version=0) into the wiped checkpoint + await self._post( + "/update_weights_from_disk", + {"model_path": self.local_checkpoint_dir, "weight_version": "0", + "flush_cache": self._flush_on_commit}, + timeout=None, + action="reset reload to base", + ) - async def continue_generation(self) -> None: + def stamp_request(self, request: dict[str, Any], served: VersionRef) -> None: + user = request.get("extra_key") + if isinstance(user, list): + request["extra_key"] = [self._extra_key(served, k) for k in user] + else: + request["extra_key"] = self._extra_key(served, user) + + def stamp_response(self, response: dict[str, Any], served: VersionRef, current: VersionRef) -> None: + meta = response.get("meta_info") + if isinstance(meta, dict): # sglang /generate carries attribution in meta_info + meta["weight_version"] = str(served.version) + meta["weight_version_start"] = served.version + meta["weight_version_end"] = current.version + else: # OpenAI-style routes at the top level + response["weight_version_start"] = served.version + response["weight_version_end"] = current.version + + def _extra_key(self, served: VersionRef, user: str | None) -> str: + # Namespace the KV cache by version (and run, so two runs that both restart at + # v1 stay distinct): requests on different versions can't share radix prefixes. + run = f"{served.run_id}/" if served.run_id else "" + return f"wv{served.version};{run}{user or ''}" + + async def _post(self, path: str, payload: dict[str, Any], *, timeout: float | None, action: str | None = None) -> None: import httpx - async with httpx.AsyncClient(timeout=120.0, trust_env=False) as client: - resp = await client.post(f"{self.upstream_url}/continue_generation", json={}) - resp.raise_for_status() + async with httpx.AsyncClient(timeout=timeout, trust_env=False) as client: + resp = await client.post(f"{self._base_url}{path}", json=payload) + _raise_for_engine(resp, action or path) - async def apply_manifest(self, manifest: VersionManifest, version_path: str) -> None: + async def _get(self, path: str, *, ok: tuple[int, ...] = (200,)) -> None: import httpx - # Bring the local checkpoint up to this version host-side (apply_deltas - # chain-replays from whatever is applied, verifying each base_version), - # then reload the full local checkpoint. version_path is the published - # version dir; its parent is the root of weight_v* dirs apply_deltas walks. - delta_root = str(Path(version_path).parent) - _t_apply = time.perf_counter() - await asyncio.to_thread( - self._apply_deltas(), self.local_checkpoint_dir, delta_root, int(manifest.version) - ) - logger.info( - "[apply timing] v=%s host_delta_apply=%.2fs", manifest.version, time.perf_counter() - _t_apply - ) - - _t_reload = time.perf_counter() - payload = { - "model_path": self.local_checkpoint_dir, - "weight_version": str(manifest.version), - # The sync manager flushes via GET /flush_cache while quiesced. - # The engine-side post-apply flush hard-asserts on failure - # (killing the scheduler process) if any request slipped in, so - # it must stay disabled here. - "flush_cache": False, - } - async with httpx.AsyncClient(timeout=None, trust_env=False) as client: - resp = await client.post(f"{self.upstream_url}/update_weights_from_disk", json=payload) - resp.raise_for_status() - data = resp.json() - if data.get("success") is False: - raise RuntimeError(f"SGLang rejected weight update: {data}") - logger.info( - "[apply timing] v=%s engine_reload=%.2fs", manifest.version, time.perf_counter() - _t_reload - ) + async with httpx.AsyncClient(timeout=self._control_timeout, trust_env=False) as client: + resp = await client.get(f"{self._base_url}{path}") + if resp.status_code not in ok: + _raise_for_engine(resp, path) + + +def _raise_for_engine(resp: Any, action: str) -> None: + # A failed control call comes back as HTTP 4xx with the engine's traceback in the + # JSON body — read the body before the status so the real error isn't lost. + try: + data = resp.json() + if not isinstance(data, dict): + data = {"message": data} + except ValueError: + data = {"message": resp.text} + if resp.status_code != 200 or data.get("success") is False: + raise RuntimeError(f"sglang rejected {action} (HTTP {resp.status_code}): {data.get('message', data)}") diff --git a/src/stitch/engines/sglang_test.py b/src/stitch/engines/sglang_test.py index 3385597..fc0cd49 100644 --- a/src/stitch/engines/sglang_test.py +++ b/src/stitch/engines/sglang_test.py @@ -1,88 +1,43 @@ -from __future__ import annotations - -import asyncio -import unittest -from unittest import mock - -from stitch.engines.sglang import ( - SGLangDiskDeltaAdapter, - compose_extra_key, - parse_extra_key_version, -) -from stitch.protocol import VersionManifest - - -class _RecordingPost: - """Stand-in for httpx.AsyncClient that records the reload POST.""" - - last_url: str | None = None - last_json: dict | None = None +"""SGLangEngine harness: version stamping (pure dict mutation, no HTTP). - def __init__(self, *args, **kwargs) -> None: - pass - - async def __aenter__(self) -> "_RecordingPost": - return self - - async def __aexit__(self, *exc) -> bool: - return False - - async def post(self, url, json=None): - import httpx - - type(self).last_url = url - type(self).last_json = json - return httpx.Response(200, json={"success": True}, request=httpx.Request("POST", url)) +The /pull_weights + /update_weights_from_disk control paths hit a real engine, so they +are validated e2e; the request/response stamping is the provable-without-sglang part.""" +from __future__ import annotations -class SGLangDiskDeltaAdapterTest(unittest.TestCase): - def test_apply_manifest_applies_host_side_then_plain_reload(self) -> None: - async def run() -> None: - calls: dict[str, list] = {"apply": [], "init": []} +from stitch.engines.base import Engine +from stitch.engines.sglang import SGLangEngine +from stitch.versions import VersionRef - adapter = SGLangDiskDeltaAdapter( - upstream_url="http://up/", - local_checkpoint_dir="/local", - base_checkpoint_dir="/base", - apply_deltas=lambda local, root, version: calls["apply"].append((local, root, version)), - init_local_checkpoint=lambda local, base: calls["init"].append((local, base)), - ) - await adapter.prepare() - manifest = VersionManifest( - version=5, base_version=4, backend="disk_delta", load_format="auto" - ) - with mock.patch("httpx.AsyncClient", _RecordingPost): - _RecordingPost.last_json = None - await adapter.apply_manifest(manifest, "/bulletin/versions/weight_v000005") +def test_satisfies_engine_port() -> None: + assert isinstance(SGLangEngine("http://engine", "/ckpt"), Engine) - # Base materialized once; delta chain applied against the version - # dir's parent (the root of weight_v* dirs) up to version 5. - self.assertEqual(calls["init"], [("/local", "/base")]) - self.assertEqual(calls["apply"], [("/local", "/bulletin/versions", 5)]) - # Engine reload is the plain disk path: local checkpoint, no - # load_format / files delta payload. - self.assertEqual(_RecordingPost.last_url, "http://up/update_weights_from_disk") - self.assertEqual( - _RecordingPost.last_json, - {"model_path": "/local", "weight_version": "5", "flush_cache": False}, - ) - asyncio.run(run()) +def test_stamp_request_namespaces_by_version() -> None: + engine = SGLangEngine("http://engine", "/ckpt") + req: dict = {"text": "hi"} + engine.stamp_request(req, VersionRef("r1", 7)) + assert req["extra_key"] == "wv7;r1/" # version + run namespace, no user key + listed: dict = {"extra_key": ["a", "b"]} + engine.stamp_request(listed, VersionRef(None, 3)) + assert listed["extra_key"] == ["wv3;a", "wv3;b"] # run-less, per-element -class ExtraKeyTest(unittest.TestCase): - def test_compose_extra_key_round_trips_and_is_position_fixed(self) -> None: - self.assertEqual(compose_extra_key(0), "wv0;") - self.assertEqual(compose_extra_key(7, "my-key"), "wv7;my-key") - self.assertEqual(parse_extra_key_version(compose_extra_key(12, None)), 12) - self.assertEqual(parse_extra_key_version(compose_extra_key(3, "wv9;decoy")), 3) - # The user key cannot shift or forge the version segment. - self.assertEqual(parse_extra_key_version("wv1;anything;else"), 1) - self.assertIsNone(parse_extra_key_version("plain-user-key")) - self.assertIsNone(parse_extra_key_version("wv12")) # no terminator - self.assertIsNone(parse_extra_key_version("wvx;k")) +def test_stamp_response_generate_vs_openai() -> None: + engine = SGLangEngine("http://engine", "/ckpt") + gen: dict = {"text": "x", "meta_info": {}} + engine.stamp_response(gen, VersionRef("r1", 4), VersionRef("r1", 5)) + assert gen["meta_info"] == {"weight_version": "4", "weight_version_start": 4, "weight_version_end": 5} + openai: dict = {"choices": []} + engine.stamp_response(openai, VersionRef("r1", 4), VersionRef("r1", 4)) + assert openai["weight_version_start"] == 4 and openai["weight_version_end"] == 4 + assert "meta_info" not in openai and "weight_version" not in openai if __name__ == "__main__": - unittest.main() + tests = [v for k, v in sorted(globals().items()) if k.startswith("test_")] + for t in tests: + t() + print(f" ok {t.__name__}") + print(f"sglang engine harness: {len(tests)} PASS") diff --git a/src/stitch/engines/vllm.py b/src/stitch/engines/vllm.py new file mode 100644 index 0000000..a526723 --- /dev/null +++ b/src/stitch/engines/vllm.py @@ -0,0 +1,72 @@ +"""``VLLMEngine`` — the ``Engine`` instance for a vLLM server. **TODO: not implemented.** + +Scaffold only. ``engines/sglang.py`` is the reference implementation; a vLLM engine slots +in behind the same ``Engine`` port with zero core changes. The work is mapping stage/commit +onto vLLM's weight-update surface, which differs from sglang's ``/pull_weights`` + +``/update_weights_from_disk`` — vLLM has no built-in disk-delta receiver, so the host-side +apply (walk to the nearest anchor, replay xor/overwrite deltas, checksum-verify) has to live +either in a worker extension or a small sidecar step before the reload. Each method below +notes what its vLLM equivalent should do; fill them in when we bring a vLLM pool up. +""" + +from __future__ import annotations + +from typing import Any + +from stitch.engines.base import Engine +from stitch.versions import VersionManifest, VersionRef + + +class VLLMEngine(Engine): + def __init__( + self, + base_url: str, + local_checkpoint_dir: str, + *, + control_timeout: float = 120.0, + ) -> None: + self._base_url = base_url.rstrip("/") + self.local_checkpoint_dir = local_checkpoint_dir + self._control_timeout = control_timeout + + def base_url(self) -> str: + return self._base_url + + async def stage(self, manifest: VersionManifest, source_dir: str) -> None: + # TODO: bring the host-local checkpoint to manifest.ref — seed from the nearest FULL + # anchor, replay deltas forward, checksum-verify. vLLM lacks sglang's /pull_weights, so + # this apply must be provided (worker extension or a sidecar copy+decode) before commit. + raise NotImplementedError("VLLMEngine.stage: TODO") + + async def commit(self, ref: VersionRef) -> None: + # TODO: reload the staged checkpoint into the serving weights (the gate covers only this) + # — e.g. vLLM collective_rpc into a WorkerExtension that runs load_weights, or vLLM's + # runtime weight-update API. Must reproduce an initial load for the served quant format. + raise NotImplementedError("VLLMEngine.commit: TODO") + + async def flush(self) -> None: + # TODO: evict KV / prefix cache before a quiesce reload. + raise NotImplementedError("VLLMEngine.flush: TODO") + + async def pause(self) -> None: + # TODO: pause the scheduler in place (in_place commit), keeping in-flight requests resident. + raise NotImplementedError("VLLMEngine.pause: TODO") + + async def resume(self) -> None: + # TODO: resume the scheduler after a pause. + raise NotImplementedError("VLLMEngine.resume: TODO") + + async def reset(self) -> None: + # TODO: wipe local_checkpoint_dir so the next stage reseeds from the engine's boot base. + raise NotImplementedError("VLLMEngine.reset: TODO") + + # prefetch() intentionally inherits the Engine no-op default until stage() exists. + + def stamp_request(self, request: dict[str, Any], served: VersionRef) -> None: + # TODO: namespace the request to the served version for KV isolation (vLLM equivalent of + # sglang's extra_key — e.g. a prefix-cache salt / request tag). Mutates request in place. + raise NotImplementedError("VLLMEngine.stamp_request: TODO") + + def stamp_response(self, response: dict[str, Any], served: VersionRef, current: VersionRef) -> None: + # TODO: record which version generated the response, in vLLM's response shape. + raise NotImplementedError("VLLMEngine.stamp_response: TODO") diff --git a/src/stitch/pools/__init__.py b/src/stitch/pools/__init__.py new file mode 100644 index 0000000..4694569 --- /dev/null +++ b/src/stitch/pools/__init__.py @@ -0,0 +1 @@ +"""Pool instances (see base.py for the port).""" diff --git a/src/stitch/pools/base.py b/src/stitch/pools/base.py new file mode 100644 index 0000000..5342f62 --- /dev/null +++ b/src/stitch/pools/base.py @@ -0,0 +1,29 @@ +"""The ``Pool`` port — a client to the running elastic replica pool. + +Instances subclass this base: ``pools/modal_flash.py`` (Modal Flash). Add k8s as a new +subclass. Override ``gateway_url`` and ``discover_replicas``; ``wake`` and ``scale`` are +optional (their no-op defaults fall back to the replicas' own polling / load-autoscale). +""" + +from __future__ import annotations + +from stitch.versions import VersionRef + + +class Pool: + """Reach, enumerate, and (optionally) nudge/scale the serving replicas — a client + to a running pool, not its deployment.""" + + def gateway_url(self) -> str: + """The single URL rollout traffic is sent to (the pool's front door).""" + raise NotImplementedError + + def discover_replicas(self) -> list[str]: + """Base URLs of the currently-live replicas.""" + raise NotImplementedError + + def wake(self, replicas: list[str], ref: VersionRef) -> None: + """Nudge replicas to reconcile now. Optional — the default relies on their polling.""" + + def scale(self, *, min: int | None = None, max: int | None = None) -> None: + """Adjust the replica floor/cap. Optional — the default relies on load-autoscale.""" diff --git a/src/stitch/pools/modal_flash.py b/src/stitch/pools/modal_flash.py new file mode 100644 index 0000000..936715a --- /dev/null +++ b/src/stitch/pools/modal_flash.py @@ -0,0 +1,83 @@ +"""``ModalFlashPool`` — the ``Pool`` instance for a Modal Flash service. + +Replicas are the Flash containers; the gateway is the Flash URL. This is a *client* +to a running pool — reach, enumerate, wake, scale — not the pool's deployment (that +is an example). Every Modal call is import-lazy, so the module loads without Modal. +""" + +from __future__ import annotations + +import logging +from concurrent.futures import ThreadPoolExecutor + +from stitch.pools.base import Pool +from stitch.versions import VersionRef + +logger = logging.getLogger(__name__) + + +class ModalFlashPool(Pool): + def __init__(self, app_name: str, cls_name: str) -> None: + self.app_name = app_name + self.cls_name = cls_name + + def gateway_url(self) -> str: + import modal + + cls = modal.Cls.from_name(self.app_name, self.cls_name) + urls = cls._experimental_get_flash_urls() + if not urls: + raise RuntimeError( + f"no Flash gateway URL for {self.app_name}.{self.cls_name} — deploy the app first" + ) + return str(urls[0]).rstrip("/") + + def discover_replicas(self) -> list[str]: + import modal + import modal.experimental + + modal.Cls.from_name(self.app_name, self.cls_name) # client-side resolve side effect + containers = modal.experimental.flash_get_containers(self.app_name, self.cls_name) + return [_normalize_url(h) for c in containers if (h := _host(c))] + + def wake(self, replicas: list[str], ref: VersionRef) -> None: + # Kick each replica to reconcile now; it re-reads the authoritative pointer, + # so no target version travels in the body. Runs in the trainer's publish hot + # path, so fan out over one shared client instead of a serial round-trip each. + if not replicas: + return + import httpx + + with httpx.Client(timeout=5.0, trust_env=False) as client: + + def wake_one(url: str) -> None: + try: + client.post(f"{url}/wake").raise_for_status() + except Exception as exc: # noqa: BLE001 + logger.warning("failed to wake %s for %s: %s", url, ref.identity, exc) + + with ThreadPoolExecutor(max_workers=min(16, len(replicas))) as pool: + list(pool.map(wake_one, replicas)) + + def scale(self, *, min: int | None = None, max: int | None = None) -> None: + import modal + + fn = modal.Cls.from_name(self.app_name, self.cls_name)._get_class_service_function() + kwargs: dict[str, int] = {} + if min is not None: + kwargs["min_containers"] = min + if max is not None: + kwargs["max_containers"] = max + if kwargs: + fn.update_autoscaler(**kwargs) + + +def _host(container) -> str | None: + if isinstance(container, dict): + return container.get("host") + return getattr(container, "host", None) + + +def _normalize_url(host: str) -> str: + host = str(host).rstrip("/") + return host if host.startswith(("http://", "https://")) else f"https://{host}" diff --git a/src/stitch/pools/modal_flash_test.py b/src/stitch/pools/modal_flash_test.py new file mode 100644 index 0000000..8468000 --- /dev/null +++ b/src/stitch/pools/modal_flash_test.py @@ -0,0 +1,20 @@ +"""ModalFlashPool harness. Every method is a lazy Modal call (discover / wake / scale / +gateway), so the provable-without-Modal surface is port conformance; the live calls are +validated e2e.""" + +from __future__ import annotations + +from stitch.pools.base import Pool +from stitch.pools.modal_flash import ModalFlashPool + + +def test_satisfies_pool_port() -> None: + assert isinstance(ModalFlashPool("app", "Server"), Pool) + + +if __name__ == "__main__": + tests = [v for k, v in sorted(globals().items()) if k.startswith("test_")] + for t in tests: + t() + print(f" ok {t.__name__}") + print(f"modal_flash harness: {len(tests)} PASS") diff --git a/src/stitch/protocol.py b/src/stitch/protocol.py deleted file mode 100644 index 458a0eb..0000000 --- a/src/stitch/protocol.py +++ /dev/null @@ -1,552 +0,0 @@ -"""Core wire protocol helpers for disaggregated rollout weight sync.""" - -from __future__ import annotations - -import json -import os -import time -from dataclasses import dataclass, field -from enum import Enum -from pathlib import Path -from typing import Any, Protocol, runtime_checkable - - -PROTOCOL_VERSION = 1 -LATEST_FILE = "latest.json" - - -class SyncState(str, Enum): - IDLE = "IDLE" - QUEUED = "QUEUED" - PREFETCHING = "PREFETCHING" - PREPARING = "PREPARING" - COMMITTING = "COMMITTING" - ERROR = "ERROR" - - -@dataclass(frozen=True) -class RolloutReplicaState: - """Readiness report for one rollout server replica. - - Providers may identify weights by an integer stitch version, an external - snapshot identity, or both. Readiness is separate from identity matching: - a healthy replica on an old snapshot is observable, but it is not ready for - requests that require the new target. - """ - - readiness: bool - current_version: int | None = None - current_snapshot_identity: str | None = None - replica_id: str | None = None - sync_state: str | None = None - readiness_reason: str | None = None - metadata: dict[str, Any] = field(default_factory=dict) - - @classmethod - def from_dict(cls, data: dict[str, Any]) -> "RolloutReplicaState": - known = { - "readiness", - "current_version", - "current_snapshot_identity", - "replica_id", - "sync_state", - "readiness_reason", - "metadata", - } - metadata = {k: v for k, v in data.items() if k not in known} - metadata.update(dict(data.get("metadata") or {})) - return cls( - readiness=_bool(data.get("readiness", False)), - current_version=_optional_int(data.get("current_version")), - current_snapshot_identity=_optional_str(data.get("current_snapshot_identity")), - replica_id=_optional_str(data.get("replica_id")), - sync_state=_optional_str(data.get("sync_state")), - readiness_reason=_optional_str(data.get("readiness_reason")), - metadata=metadata, - ) - - def to_dict(self) -> dict[str, Any]: - data: dict[str, Any] = {"readiness": self.readiness} - if self.current_version is not None: - data["current_version"] = self.current_version - if self.current_snapshot_identity is not None: - data["current_snapshot_identity"] = self.current_snapshot_identity - if self.replica_id is not None: - data["replica_id"] = self.replica_id - if self.sync_state is not None: - data["sync_state"] = self.sync_state - if self.readiness_reason is not None: - data["readiness_reason"] = self.readiness_reason - if self.metadata: - data["metadata"] = self.metadata - return data - - def matches_target( - self, - *, - target_version: int | None = None, - target_snapshot_identity: str | None = None, - ) -> bool: - if not self.readiness: - return False - if target_version is not None and self.current_version != int(target_version): - return False - if target_snapshot_identity is not None and self.current_snapshot_identity != str(target_snapshot_identity): - return False - return True - - -@dataclass(frozen=True) -class RolloutPoolState: - """Readiness report for a rollout server pool.""" - - replicas: list[RolloutReplicaState] = field(default_factory=list) - protocol_version: int = PROTOCOL_VERSION - metadata: dict[str, Any] = field(default_factory=dict) - - @classmethod - def from_dict(cls, data: dict[str, Any]) -> "RolloutPoolState": - known = {"protocol_version", "replicas", "metadata"} - metadata = {k: v for k, v in data.items() if k not in known} - metadata.update(dict(data.get("metadata") or {})) - return cls( - protocol_version=int(data.get("protocol_version", PROTOCOL_VERSION)), - replicas=[RolloutReplicaState.from_dict(x) for x in data.get("replicas", [])], - metadata=metadata, - ) - - def to_dict(self) -> dict[str, Any]: - data: dict[str, Any] = { - "protocol_version": int(self.protocol_version), - "replicas": [replica.to_dict() for replica in self.replicas], - } - if self.metadata: - data["metadata"] = self.metadata - return data - - def ready_count( - self, - *, - target_version: int | None = None, - target_snapshot_identity: str | None = None, - ) -> int: - return sum( - replica.matches_target( - target_version=target_version, - target_snapshot_identity=target_snapshot_identity, - ) - for replica in self.replicas - ) - - def readiness_fraction( - self, - *, - target_version: int | None = None, - target_snapshot_identity: str | None = None, - ) -> float: - if not self.replicas: - return 0.0 - return self.ready_count( - target_version=target_version, - target_snapshot_identity=target_snapshot_identity, - ) / len(self.replicas) - - def is_ready( - self, - *, - threshold: float = 1.0, - target_version: int | None = None, - target_snapshot_identity: str | None = None, - ) -> bool: - return bool(self.replicas) and self.readiness_fraction( - target_version=target_version, - target_snapshot_identity=target_snapshot_identity, - ) >= float(threshold) - - -@dataclass(frozen=True) -class WeightVersionPolicy: - """Request-level policy for acceptable rollout server weights.""" - - min_required_version: int | None = None - exact_version: int | None = None - - @classmethod - def from_payload(cls, payload: dict[str, Any] | None) -> "WeightVersionPolicy": - raw = (payload or {}).get("weight_version") or {} - if not isinstance(raw, dict): - raw = {} - return cls( - min_required_version=_optional_int(raw.get("min_required_version")), - exact_version=_optional_int(raw.get("exact_version")), - ) - - def to_payload(self) -> dict[str, int | None]: - return { - "min_required_version": self.min_required_version, - "exact_version": self.exact_version, - } - - -@dataclass(frozen=True) -class Artifact: - """An immutable artifact referenced by a version manifest.""" - - kind: str - path: str - size_bytes: int | None = None - checksum: str | None = None - metadata: dict[str, Any] = field(default_factory=dict) - - @classmethod - def from_dict(cls, data: dict[str, Any]) -> "Artifact": - return cls( - kind=str(data["kind"]), - path=str(data["path"]), - size_bytes=_optional_int(data.get("size_bytes")), - checksum=None if data.get("checksum") is None else str(data["checksum"]), - metadata=dict(data.get("metadata") or {}), - ) - - def to_dict(self) -> dict[str, Any]: - data: dict[str, Any] = { - "kind": self.kind, - "path": self.path, - } - if self.size_bytes is not None: - data["size_bytes"] = self.size_bytes - if self.checksum is not None: - data["checksum"] = self.checksum - if self.metadata: - data["metadata"] = self.metadata - return data - - -@dataclass(frozen=True) -class VersionManifest: - """Published description of one trainer-to-rollout weight version.""" - - version: int - base_version: int - backend: str - load_format: str - # How the transition artifacts are encoded/compressed/checksummed. Mirrors - # slime's disk-delta index.json metadata so the engine-neutral manifest is - # self-describing; None on full snapshots or pre-delta manifests. - delta_encoding: str | None = None - compression_format: str | None = None - checksum_format: str | None = None - transition_files: list[str] = field(default_factory=list) - created_at: float = field(default_factory=time.time) - protocol_version: int = PROTOCOL_VERSION - artifacts: list[Artifact] = field(default_factory=list) - run_id: str | None = None - base_model: str | None = None - recovery: dict[str, Any] | None = None - metadata: dict[str, Any] = field(default_factory=dict) - - @classmethod - def read(cls, path: str | Path) -> "VersionManifest": - with Path(path).open("r", encoding="utf-8") as f: - data = json.load(f) - protocol_version = int(data.get("protocol_version", PROTOCOL_VERSION)) - if protocol_version != PROTOCOL_VERSION: - raise ValueError( - f"manifest at {path} declares protocol_version {protocol_version}, " - f"but this build supports {PROTOCOL_VERSION}; refusing to read it" - ) - artifacts = [Artifact.from_dict(x) for x in data.get("artifacts", [])] - transition_files = [str(x) for x in data.get("transition_files", [])] - if not transition_files: - transition_files = [a.path for a in artifacts if a.kind == "transition"] - return cls( - version=int(data["version"]), - base_version=int(data["base_version"]), - backend=str(data.get("backend", "sparse_delta")), - load_format=str(data.get("load_format", "delta")), - delta_encoding=_optional_str(data.get("delta_encoding")), - compression_format=_optional_str(data.get("compression_format")), - checksum_format=_optional_str(data.get("checksum_format")), - transition_files=transition_files, - created_at=float(data.get("created_at", 0.0)), - protocol_version=protocol_version, - artifacts=artifacts, - run_id=None if data.get("run_id") is None else str(data["run_id"]), - base_model=None if data.get("base_model") is None else str(data["base_model"]), - recovery=data.get("recovery"), - metadata=dict(data.get("metadata") or {}), - ) - - @classmethod - def from_slime_index( - cls, - version_dir: str | Path, - *, - run_id: str | None = None, - base_model: str | None = None, - metadata: dict[str, Any] | None = None, - ) -> "VersionManifest": - """Lift a slime disk-delta version directory's - ``model.safetensors.index.json`` into the engine-neutral manifest. - - slime's disk-delta publisher writes a canonical HF index whose - ``metadata`` block carries the version lineage and the delta - encoding/compression/checksum. The delta is applied host-side and the - engine then reloads the full local checkpoint, so ``load_format`` is the - plain ``auto`` path, not a delta receiver. - """ - index_path = Path(version_dir) / "model.safetensors.index.json" - with index_path.open("r", encoding="utf-8") as f: - index = json.load(f) - meta = index.get("metadata") or {} - files = sorted({str(name) for name in (index.get("weight_map") or {}).values()}) - return cls( - version=int(meta["version"]), - base_version=int(meta["base_version"]), - backend="disk_delta", - load_format="auto", - delta_encoding=_optional_str(meta.get("delta_encoding")), - compression_format=_optional_str(meta.get("compression_format")), - checksum_format=_optional_str(meta.get("checksum_format")), - transition_files=files, - artifacts=[Artifact(kind="transition", path=path) for path in files], - run_id=run_id, - base_model=base_model, - metadata=metadata if metadata is not None else {"trainer": "slime", "transport": "disk"}, - ) - - def write(self, path: str | Path) -> None: - atomic_write_json(path, self.to_dict()) - - def to_dict(self) -> dict[str, Any]: - artifacts = self.artifacts or [Artifact(kind="transition", path=p) for p in self.transition_files] - data: dict[str, Any] = { - "protocol_version": int(self.protocol_version), - "version": int(self.version), - "base_version": int(self.base_version), - "backend": self.backend, - "load_format": self.load_format, - "transition_files": list(self.transition_files), - "artifacts": [a.to_dict() for a in artifacts], - "created_at": float(self.created_at), - } - if self.delta_encoding is not None: - data["delta_encoding"] = self.delta_encoding - if self.compression_format is not None: - data["compression_format"] = self.compression_format - if self.checksum_format is not None: - data["checksum_format"] = self.checksum_format - if self.run_id is not None: - data["run_id"] = self.run_id - if self.base_model is not None: - data["base_model"] = self.base_model - if self.recovery is not None: - data["recovery"] = self.recovery - if self.metadata: - data["metadata"] = self.metadata - return data - - def transition_artifact_paths(self) -> list[str]: - if self.transition_files: - return list(self.transition_files) - return [a.path for a in self.artifacts if a.kind == "transition"] - - -def weight_identity(version: int) -> str: - """Canonical snapshot-identity string for an integer weight version.""" - return f"weight_v{int(version):06d}" - - -def parse_weight_identity(identity: str) -> int | None: - """Inverse of :func:`weight_identity`; None if not ``weight_v``.""" - prefix = "weight_v" - if not identity.startswith(prefix): - return None - digits = identity[len(prefix):] - return int(digits) if digits.isdigit() else None - - -def format_snapshot_identity(run_id: str | None, version: int) -> str: - """The canonical pointer/snapshot identity for a (run_id, version). - - ``/weight_v`` when a run is named, else the bare - ``weight_v`` (the degenerate single-run / customer flat layout). This - is the single self-identifying value written to the slime-layout ``latest`` - pointer: a run-scoped chain can never be mistaken for a different run's, and - an old bare pointer parses back to ``run_id=None`` rather than a phantom run. - """ - identity = weight_identity(version) - return f"{run_id}/{identity}" if run_id else identity - - -def parse_snapshot_identity(text: str) -> tuple[str | None, int]: - """Inverse of :func:`format_snapshot_identity`, tolerant of legacy pointers. - - ``/weight_v`` -> ``(run_id, version)``; a bare - ``weight_v`` -> ``(None, version)``; a legacy raw ```` -> - ``(None, version)``; empty / unparseable -> ``(None, 0)`` (treated as - no-valid-pointer, i.e. not-ready rather than a misparse). - """ - text = (text or "").strip() - if not text: - return (None, 0) - run_id: str | None = None - tail = text - if "/" in text: - run_id, tail = text.rsplit("/", 1) - run_id = run_id or None - version = parse_weight_identity(tail) - if version is None: - version = int(tail) if tail.isdigit() else None - if version is None: - return (None, 0) - return (run_id, version) - - -def version_dir(root: str | Path, version: int) -> Path: - return Path(root) / "versions" / weight_identity(version) - - -def latest_path(root: str | Path) -> Path: - return Path(root) / LATEST_FILE - - -def read_latest(root: str | Path) -> int: - path = latest_path(root) - if not path.exists(): - return 0 - with path.open("r", encoding="utf-8") as f: - data = json.load(f) - return int(data.get("version", 0)) - - -def write_latest(root: str | Path, version: int) -> None: - atomic_write_json( - latest_path(root), - { - "protocol_version": PROTOCOL_VERSION, - "version": int(version), - "updated_at": time.time(), - }, - ) - - -def version_not_ready_error(current: int, target: int) -> dict[str, Any]: - return { - "error": { - "type": "WeightVersionNotReady", - "message": f"server is at version {current}, target {target} is not ready", - "current_version": int(current), - "target_version": int(target), - } - } - - -def version_too_old_error(current: int, target: int) -> dict[str, Any]: - return { - "error": { - "type": "WeightVersionTooOld", - "message": f"server is at version {current}, cannot roll back to {target}", - "current_version": int(current), - "target_version": int(target), - } - } - - -def evaluate_version_policy( - current_version: int, policy: WeightVersionPolicy -) -> dict[str, Any] | None: - """Shared exact/min admission check. Returns a typed error dict or None. - - Callers decide how to react to a not-ready error (pull toward the target vs - reject): the bulletin-board manager queues a sync, the hot-load shim rejects. - """ - if policy.exact_version is not None: - target = int(policy.exact_version) - if current_version < target: - return version_not_ready_error(current_version, target) - if current_version > target: - return version_too_old_error(current_version, target) - return None - if policy.min_required_version is not None and current_version < int( - policy.min_required_version - ): - return version_not_ready_error(current_version, int(policy.min_required_version)) - return None - - -def atomic_write_json(path: str | Path, payload: dict[str, Any]) -> None: - path = Path(path) - path.parent.mkdir(parents=True, exist_ok=True) - tmp = path.with_name(path.name + ".tmp") - with tmp.open("w", encoding="utf-8") as f: - json.dump(payload, f, sort_keys=True) - f.write("\n") - f.flush() - os.fsync(f.fileno()) - os.replace(tmp, path) - - -def atomic_write_text(path: str | Path, text: str) -> None: - path = Path(path) - path.parent.mkdir(parents=True, exist_ok=True) - tmp = path.with_name(path.name + ".tmp") - with tmp.open("w", encoding="utf-8") as f: - f.write(text) - f.flush() - os.fsync(f.fileno()) - os.replace(tmp, path) - - -def _optional_int(value: Any) -> int | None: - if value is None: - return None - return int(value) - - -def _optional_str(value: Any) -> str | None: - if value is None: - return None - return str(value) - - -def _bool(value: Any) -> bool: - if isinstance(value, bool): - return value - if isinstance(value, str): - return value.lower() in {"1", "true", "yes"} - return bool(value) - - -@runtime_checkable -class EngineAdapter(Protocol): - """Contract for rollout engine adapters driven by :class:`WeightSyncManager`. - - An adapter bridges the sync manager to one inference engine instance. - Required methods are called during every version commit; optional methods - (``prepare``, ``reset``) are probed with ``getattr`` at startup and on run - switches, so adapters may omit them. - - See :class:`stitch.engines.sglang.SGLangDiskDeltaAdapter` for the canonical - implementation. - """ - - backend: str - - async def flush_cache(self) -> None: - """Evict all cached state (KV, radix tree). Called before ``apply_manifest`` - in quiesce mode; skipped in in_place mode.""" - ... - - async def apply_manifest(self, manifest: VersionManifest, version_path: str) -> None: - """Apply one published weight version to the engine.""" - ... - - async def pause_generation(self) -> None: - """Pause the engine's scheduler in place. Required for ``commit_mode="in_place"``; - in-flight requests stay resident and resume after ``continue_generation``.""" - ... - - async def continue_generation(self) -> None: - """Resume the engine's scheduler after a pause.""" - ... diff --git a/src/stitch/protocol_test.py b/src/stitch/protocol_test.py deleted file mode 100644 index 37abe90..0000000 --- a/src/stitch/protocol_test.py +++ /dev/null @@ -1,181 +0,0 @@ -from __future__ import annotations - -import tempfile -import unittest -from pathlib import Path - -from stitch.bulletin import FilesystemBulletinBoard -from stitch.protocol import ( - Artifact, - RolloutPoolState, - RolloutReplicaState, - VersionManifest, - WeightVersionPolicy, - evaluate_version_policy, - parse_weight_identity, - read_latest, - weight_identity, -) - - -class ProtocolTest(unittest.TestCase): - def test_manifest_round_trips_extended_and_legacy_fields(self) -> None: - with tempfile.TemporaryDirectory() as tmp: - root = Path(tmp) - board = FilesystemBulletinBoard(root) - manifest = VersionManifest( - version=3, - base_version=2, - backend="disk_delta", - load_format="auto", - delta_encoding="xor", - compression_format="zstd", - checksum_format="xxh3-128", - transition_files=["rank0000_flush000000.safetensors"], - artifacts=[ - Artifact( - kind="transition", - path="rank0000_flush000000.safetensors", - checksum="sha256:abc", - ) - ], - run_id="run-1", - base_model="Qwen/Qwen3-4B", - ) - - board.publish_manifest(manifest) - loaded = board.read_manifest(None, 3) - - self.assertEqual(read_latest(root), 3) - self.assertEqual(loaded.version, 3) - self.assertEqual(loaded.base_version, 2) - self.assertEqual(loaded.transition_artifact_paths(), ["rank0000_flush000000.safetensors"]) - self.assertEqual(loaded.artifacts[0].checksum, "sha256:abc") - self.assertEqual(loaded.run_id, "run-1") - self.assertEqual(loaded.delta_encoding, "xor") - self.assertEqual(loaded.compression_format, "zstd") - self.assertEqual(loaded.checksum_format, "xxh3-128") - - def test_manifest_from_slime_index(self) -> None: - import json - - with tempfile.TemporaryDirectory() as tmp: - version_dir = Path(tmp) / "weight_v000007" - version_dir.mkdir(parents=True) - (version_dir / "model.safetensors.index.json").write_text( - json.dumps( - { - "metadata": { - "version": "000007", - "base_version": "000006", - "delta_encoding": "xor", - "compression_format": "zstd", - "checksum_format": "xxh3-128", - }, - "weight_map": { - "model.layers.0.weight": "model-00001-of-00002.safetensors", - "model.layers.1.weight": "model-00002-of-00002.safetensors", - }, - } - ), - encoding="utf-8", - ) - - manifest = VersionManifest.from_slime_index(version_dir, run_id="run-9") - - self.assertEqual(manifest.version, 7) - self.assertEqual(manifest.base_version, 6) - self.assertEqual(manifest.backend, "disk_delta") - self.assertEqual(manifest.load_format, "auto") - self.assertEqual(manifest.delta_encoding, "xor") - self.assertEqual(manifest.compression_format, "zstd") - self.assertEqual(manifest.checksum_format, "xxh3-128") - self.assertEqual( - manifest.transition_files, - ["model-00001-of-00002.safetensors", "model-00002-of-00002.safetensors"], - ) - self.assertEqual(manifest.run_id, "run-9") - - def test_weight_identity_round_trips(self) -> None: - self.assertEqual(weight_identity(0), "weight_v000000") - self.assertEqual(weight_identity(123), "weight_v000123") - self.assertEqual(parse_weight_identity("weight_v000123"), 123) - self.assertIsNone(parse_weight_identity("base")) - self.assertIsNone(parse_weight_identity("weight_vxyz")) - - def test_evaluate_version_policy(self) -> None: - self.assertIsNone(evaluate_version_policy(5, WeightVersionPolicy())) - self.assertIsNone(evaluate_version_policy(5, WeightVersionPolicy(exact_version=5))) - self.assertEqual( - evaluate_version_policy(4, WeightVersionPolicy(exact_version=5))["error"]["type"], - "WeightVersionNotReady", - ) - self.assertEqual( - evaluate_version_policy(6, WeightVersionPolicy(exact_version=5))["error"]["type"], - "WeightVersionTooOld", - ) - self.assertEqual( - evaluate_version_policy(4, WeightVersionPolicy(min_required_version=5))["error"]["type"], - "WeightVersionNotReady", - ) - self.assertIsNone(evaluate_version_policy(5, WeightVersionPolicy(min_required_version=5))) - - def test_weight_version_policy_ignores_malformed_payload(self) -> None: - self.assertEqual(WeightVersionPolicy.from_payload({}), WeightVersionPolicy()) - self.assertEqual(WeightVersionPolicy.from_payload({"weight_version": 7}), WeightVersionPolicy()) - self.assertEqual( - WeightVersionPolicy.from_payload({"weight_version": {"min_required_version": "5", "exact_version": 6}}), - WeightVersionPolicy(min_required_version=5, exact_version=6), - ) - - def test_rollout_pool_state_parses_snapshot_readiness(self) -> None: - state = RolloutPoolState.from_dict( - { - "replicas": [ - { - "replica_id": "a", - "readiness": True, - "current_snapshot_identity": "ckpt-2", - "zone": "us-east-1a", - }, - { - "replica_id": "b", - "readiness": True, - "current_snapshot_identity": "ckpt-2", - }, - { - "replica_id": "c", - "readiness": False, - "current_snapshot_identity": "ckpt-1", - "readiness_reason": "downloading weights", - }, - ] - } - ) - - self.assertEqual(state.ready_count(target_snapshot_identity="ckpt-2"), 2) - self.assertEqual(state.readiness_fraction(target_snapshot_identity="ckpt-2"), 2 / 3) - self.assertTrue(state.is_ready(target_snapshot_identity="ckpt-2", threshold=0.5)) - self.assertFalse(state.is_ready(target_snapshot_identity="ckpt-2", threshold=1.0)) - self.assertEqual(state.replicas[0].metadata["zone"], "us-east-1a") - self.assertEqual(state.replicas[2].readiness_reason, "downloading weights") - - def test_rollout_pool_state_matches_integer_versions(self) -> None: - state = RolloutPoolState( - replicas=[ - RolloutReplicaState(readiness=True, current_version=7), - RolloutReplicaState(readiness=True, current_version=6), - ] - ) - - self.assertEqual(state.ready_count(target_version=7), 1) - self.assertEqual(state.ready_count(), 2) - self.assertFalse(RolloutPoolState().is_ready()) - self.assertEqual( - RolloutPoolState.from_dict(state.to_dict()), - state, - ) - - -if __name__ == "__main__": - unittest.main() diff --git a/src/stitch/providers/__init__.py b/src/stitch/providers/__init__.py deleted file mode 100644 index 447597b..0000000 --- a/src/stitch/providers/__init__.py +++ /dev/null @@ -1 +0,0 @@ -"""Infrastructure provider helpers.""" diff --git a/src/stitch/providers/modal.py b/src/stitch/providers/modal.py deleted file mode 100644 index a7ab880..0000000 --- a/src/stitch/providers/modal.py +++ /dev/null @@ -1,95 +0,0 @@ -"""Modal provider helpers for disaggregated rollout.""" - -from __future__ import annotations - -from collections.abc import Callable - - -def commit_volume(volume_name: str) -> None: - import modal - - modal.Volume.from_name(volume_name, version=2, create_if_missing=True).commit() - - -def reload_volume(volume_name: str) -> None: - import modal - - modal.Volume.from_name(volume_name, version=2, create_if_missing=True).reload() - - -def volume_reloader(volume_name: str) -> Callable[[], None]: - return lambda: reload_volume(volume_name) - - -def discover_flash_targets(app_name: str, cls_name: str) -> list[str]: - import modal - import modal.experimental - - # Resolve the Cls by name for its Modal-client-side side effect; the returned - # handle is intentionally unused (flash_get_containers queries by name). - modal.Cls.from_name(app_name, cls_name) - containers = modal.experimental.flash_get_containers(app_name, cls_name) - return _flash_targets_from_containers(containers) - - -def _flash_targets_from_containers(containers) -> list[str]: - targets: list[str] = [] - for container in containers: - host = ( - container.get("host") - if isinstance(container, dict) - else getattr(container, "host", None) - ) - if host: - targets.append(normalize_base_url(str(host))) - return targets - - -def resolve_flash_gateway_url(app_name: str, cls_name: str) -> str: - import modal - - cls = modal.Cls.from_name(app_name, cls_name) - urls = cls._experimental_get_flash_urls() - if not urls: - raise RuntimeError( - f"No Flash gateway URL found for {app_name}.{cls_name}. " - "Deploy the app first so Modal starts the Flash pool." - ) - return str(urls[0]).rstrip("/") - - -def wake_targets(targets: list[str], version: int, *, timeout: float = 5.0) -> None: - import logging - from concurrent.futures import ThreadPoolExecutor - - import httpx - - logger = logging.getLogger(__name__) - - if not targets: - return - - # Wakes run in the trainer publish hot path; fan out instead of paying one - # round-trip per container serially. One httpx.Client is shared across the - # pool (Clients are thread-safe) so connections and keep-alive are reused. - # trust_env=False keeps proxy env vars from rerouting localhost/gateway hops. - with httpx.Client(timeout=timeout, trust_env=False) as client: - - def wake_one(target: str) -> None: - url = f"{target}/rpc_sync_from_bulletin_board" - try: - resp = client.post(url, json={"target_version": int(version)}) - resp.raise_for_status() - logger.info("Wake sync accepted by %s: %s", target, resp.text[:200]) - except Exception as exc: # noqa: BLE001 - logger.warning("Failed to wake %s for version %s: %s", target, version, exc) - - with ThreadPoolExecutor(max_workers=min(16, len(targets))) as pool: - list(pool.map(wake_one, targets)) - - -def normalize_base_url(host: str) -> str: - host = host.rstrip("/") - if host.startswith(("http://", "https://")): - return host - return f"https://{host}" diff --git a/src/stitch/publish.py b/src/stitch/publish.py new file mode 100644 index 0000000..bbf26f8 --- /dev/null +++ b/src/stitch/publish.py @@ -0,0 +1,81 @@ +"""Trainer-side helpers: publish a version, claim a run, constrain a rollout request. + +These are what a training framework wires into its publish and request hooks. They +compose the Store and Pool ports, so they work with any backend — no Modal here. +""" + +from __future__ import annotations + +import logging +from typing import Any + +from stitch.pools.base import Pool +from stitch.stores.base import Store +from stitch.versions import VersionConstraint, VersionManifest, VersionRef, decide_pointer_move + +logger = logging.getLogger(__name__) + + +def publish_version( + store: Store, + pool: Pool | None, + version_dir: str, + *, + run_id: str, + base_model: str | None = None, +) -> VersionRef: + """Publish one version from a framework-written directory (full or delta): derive the + manifest from its HF index, write it durably, advance ``latest`` (rejecting a rewind), + then wake the pool. Files land before the pointer moves, so a replica never sees a + pointer to incomplete bytes.""" + manifest = VersionManifest.from_hf_index(version_dir, run_id=run_id, base_model=base_model) + decide_pointer_move(store.read_pointer(), manifest.ref) # raises PointerRewind on a rewind + store.publish(manifest, version_dir) + store.advance_pointer(manifest.ref) + _wake(pool, manifest.ref) + return manifest.ref + + +def claim_run(store: Store, pool: Pool | None, run_id: str) -> None: + """Start a run at base before its first publish: write the base pointer and wake the + pool, so every replica (cold or warm on a finished run) resets to base up front. A + reused ``run_id`` (the run's per-launch fence token) is a rewind — rejected here so a + restart can't leave the pool pinned to a dead incarnation's high-water mark.""" + decide_pointer_move(store.read_pointer(), VersionRef(run_id, 0)) # raises PointerRewind on a reused run_id + store.claim(run_id) + _wake(pool, VersionRef(run_id, 0)) + + +def _wake(pool: Pool | None, ref: VersionRef) -> None: + """Best-effort pool wake: the pointer is already durable, so a transient control-plane + error just costs latency (replicas self-sync on their next poll/startup).""" + if pool is None: + return + try: + pool.wake(pool.discover_replicas(), ref) + except Exception: # noqa: BLE001 + logger.warning("pool wake failed for %s; replicas will self-sync", ref.identity, exc_info=True) + + +def constrain_request( + payload: dict[str, Any], + headers: dict[str, str], + *, + latest: int | None = None, + lag: int = 0, + exact: int | None = None, + session_id: Any = None, + affinity_header: str | None = None, +) -> None: + """Set the version constraint (on ``payload``) and session affinity (on ``headers``) + for one outgoing rollout request. ``exact`` pins a single version; otherwise a + bounded-lag request floors the version at ``latest - lag``. Mutates in place.""" + if exact is not None: + constraint = VersionConstraint(exact_version=int(exact)) + elif latest is not None: + constraint = VersionConstraint(min_version=max(0, int(latest) - int(lag))) + else: + constraint = VersionConstraint() + payload["weight_version"] = constraint.to_payload() + if affinity_header and session_id is not None: + headers[affinity_header] = str(session_id) diff --git a/src/stitch/publish_test.py b/src/stitch/publish_test.py new file mode 100644 index 0000000..b0d27a5 --- /dev/null +++ b/src/stitch/publish_test.py @@ -0,0 +1,119 @@ +"""Publish-side harness: publish_version / claim_run / constrain_request against fakes.""" + +from __future__ import annotations + +import json +import tempfile +from pathlib import Path + +from stitch.pools.base import Pool +from stitch.publish import claim_run, constrain_request, publish_version +from stitch.stores.base import Store +from stitch.versions import PointerRewind, VersionRef + + +class FakeStore(Store): + def __init__(self, pointer: VersionRef | None = None) -> None: + self._pointer = pointer + self.published: list = [] + + def read_pointer(self): + return self._pointer + + def publish(self, manifest, files_dir): + self.published.append(manifest) + + def advance_pointer(self, ref): + self._pointer = ref + + def claim(self, run_id): + self._pointer = VersionRef(run_id, 0) + + +class FakePool(Pool): + def __init__(self) -> None: + self.woke: list = [] + + def discover_replicas(self): + return ["http://r1"] + + def wake(self, replicas, ref): + self.woke.append(ref) + + +def _version_dir(tmp: str, *, version: int, base: int | None = None, diff: str | None = None) -> str: + d = Path(tmp) / f"weight_v{version:06d}" + d.mkdir() + meta: dict = {"version": version} + if diff: + meta.update({"diff": diff, "base_version": base, "compression": "zstd", "checksum": "xxh3-128"}) + index = {"metadata": meta, "weight_map": {"w": "model-00001.safetensors"}} + (d / "model.safetensors.index.json").write_text(json.dumps(index)) + return str(d) + + +def test_publish_full() -> None: + with tempfile.TemporaryDirectory() as tmp: + store, pool = FakeStore(), FakePool() + ref = publish_version(store, pool, _version_dir(tmp, version=1), run_id="r1") + assert ref == VersionRef("r1", 1) + assert store._pointer == VersionRef("r1", 1) + man = store.published[0] + assert man.kind.value == "full" and man.base_version is None + assert pool.woke == [VersionRef("r1", 1)] + + +def test_publish_delta() -> None: + with tempfile.TemporaryDirectory() as tmp: + store = FakeStore(VersionRef("r1", 1)) + publish_version(store, None, _version_dir(tmp, version=2, base=1, diff="xor"), run_id="r1") + man = store.published[0] + assert man.kind.value == "delta" and man.base_version == 1 and man.delta_encoding == "xor" + assert store._pointer == VersionRef("r1", 2) + + +def test_publish_rewind_rejected() -> None: + with tempfile.TemporaryDirectory() as tmp: + store = FakeStore(VersionRef("r1", 3)) + try: + publish_version(store, None, _version_dir(tmp, version=1), run_id="r1") + raise AssertionError("expected PointerRewind") + except PointerRewind: + pass + + +def test_claim_run() -> None: + store, pool = FakeStore(), FakePool() + claim_run(store, pool, "r2") + assert store._pointer == VersionRef("r2", 0) + assert pool.woke == [VersionRef("r2", 0)] + + +def test_claim_rewind_rejected() -> None: + store = FakeStore(VersionRef("r2", 3)) + try: + claim_run(store, None, "r2") # reused run_id already at v3 -> rewind to base + raise AssertionError("expected PointerRewind") + except PointerRewind: + pass + + +def test_constrain_lag() -> None: + payload, headers = {}, {} + constrain_request(payload, headers, latest=10, lag=2, session_id="g1", affinity_header="X-Session") + assert payload["weight_version"] == {"min_version": 8, "exact_version": None} + assert headers["X-Session"] == "g1" + + +def test_constrain_exact() -> None: + payload, headers = {}, {} + constrain_request(payload, headers, exact=5) + assert payload["weight_version"] == {"min_version": None, "exact_version": 5} + + +if __name__ == "__main__": + tests = [v for k, v in sorted(globals().items()) if k.startswith("test_")] + for t in tests: + t() + print(f" ok {t.__name__}") + print(f"publish harness: {len(tests)} PASS") diff --git a/src/stitch/servers/__init__.py b/src/stitch/servers/__init__.py deleted file mode 100644 index e4ceec7..0000000 --- a/src/stitch/servers/__init__.py +++ /dev/null @@ -1 +0,0 @@ -"""HTTP server adapters.""" diff --git a/src/stitch/servers/sglang.py b/src/stitch/servers/sglang.py deleted file mode 100644 index 338ecae..0000000 --- a/src/stitch/servers/sglang.py +++ /dev/null @@ -1,356 +0,0 @@ -"""HTTP sidecar that adds weight-version protocol semantics to SGLang.""" - -import asyncio -import contextlib -import inspect -import logging -import uuid -from collections.abc import Callable, Iterable -from contextlib import asynccontextmanager -from typing import Any - -from stitch.engines.sglang import compose_extra_key -from stitch.protocol import WeightVersionPolicy -from stitch.sync import PolicyViolation, RolloutSyncManager - - -logger = logging.getLogger(__name__) - -# Engine control routes that must not be reachable through the body-blind -# gateway: a stray call would mutate engine state behind the sync manager's -# back and silently break its version bookkeeping. The sidecar itself calls -# these on the upstream directly, not through the proxy. -BLOCKED_ROUTES = frozenset( - { - "update_weights_from_disk", - "update_weights_from_distributed", - "update_weights_from_tensor", - "update_weight_version", - "init_weights_update_group", - "destroy_weights_update_group", - "flush_cache", - "pause_generation", - "continue_generation", - "abort_request", - "release_memory_occupation", - "resume_memory_occupation", - "post_process_weights", - } -) - - -def create_app( - manager: RolloutSyncManager, - *, - upstream_url: str, - versioned_routes: Iterable[str] = ("generate", "v1/chat/completions"), - register_routes: Callable[[Any], None] | None = None, - include_sync_routes: bool = True, - upstream_timeout: float | None = 3600.0, - background_sync_interval: float | None = None, -): - from fastapi import FastAPI, Request - from fastapi.responses import JSONResponse, Response - import httpx - - upstream_url = upstream_url.rstrip("/") - # Bound the wait for an upstream (SGLang) response. A generation that - # finishes but never delivers its HTTP body would otherwise hang this proxy - # forever (timeout=None), holding the request open and wedging the client - # awaiting it — exactly the failure mode that stalled a rollout for hours. - # On timeout the upstream call raises, surfacing as a 5xx the client can - # retry, instead of an infinite hold. connect stays short (upstream is - # localhost); pass None to opt out. - upstream_request_timeout = httpx.Timeout(upstream_timeout, connect=10.0) - versioned_route_set = {route.strip("/") for route in versioned_routes} - - # One pooled upstream client for the whole process. A rollout proxy that - # reconnects to the engine on every request pays a TCP/pool setup per hop; - # reusing the client keeps connections warm across the sidecar->engine hop - # that carries every rollout. Created lazily on the first proxied request - # (so request-time test patching of httpx.AsyncClient still wins) and closed - # on shutdown. - pooled: dict[str, Any] = {} - - def upstream_client() -> Any: - client = pooled.get("client") - if client is None: - client = httpx.AsyncClient(timeout=upstream_request_timeout, trust_env=False) - pooled["client"] = client - return client - - async def _reconcile_loop(interval: float) -> None: - # Pull-based reconcile against the bulletin board's `latest` pointer. - # In the log-as-truth deployment the front door advances `latest` and - # the pool catches up here, so a replica that missed a wake (or scaled - # up after one) still converges without any request-version pin. - board = getattr(manager, "board", None) - queue_sync = getattr(manager, "queue_sync", None) - if board is None or queue_sync is None: - return - while True: - await asyncio.sleep(interval) - try: - await board.refresh() - queue_sync() - except Exception: # noqa: BLE001 - logger.warning("background reconcile failed", exc_info=True) - - reconcile: dict[str, Any] = {} - - @asynccontextmanager - async def lifespan(_app: FastAPI): - await manager.startup_sync() - if background_sync_interval and background_sync_interval > 0: - reconcile["task"] = asyncio.ensure_future(_reconcile_loop(background_sync_interval)) - try: - yield - finally: - task = reconcile.pop("task", None) - if task is not None: - task.cancel() - with contextlib.suppress(BaseException): - await task - client = pooled.pop("client", None) - if client is not None: - await client.aclose() - shutdown_sync = getattr(manager, "shutdown_sync", None) - if shutdown_sync is not None: - result = shutdown_sync() - if inspect.isawaitable(result): - await result - - app = FastAPI(lifespan=lifespan) - - @app.get("/health") - async def health() -> dict[str, Any]: - return {"ok": True, "current_version": manager.current_version} - - @app.get("/server_info") - async def server_info() -> dict[str, Any]: - return await manager.server_info() - - @app.get("/get_weight_version") - async def get_weight_version() -> dict[str, str]: - return {"weight_version": str(manager.current_version)} - - if include_sync_routes: - - @app.post("/rpc_sync_from_bulletin_board") - async def rpc_sync_from_bulletin_board(request: Request) -> dict[str, Any]: - payload = await request.json() - target = payload.get("target_version") - manager.queue_sync(int(target) if target is not None else None) - return { - "accepted": True, - "current_version": manager.current_version, - "queued_target_version": getattr( - manager, "queued_target_version", None - ), - "sync_state": _sync_state_value(getattr(manager, "sync_state", None)), - } - - if register_routes is not None: - register_routes(app) - - async def _watch_disconnect(request: Request) -> None: - while True: - message = await request.receive() - if message["type"] == "http.disconnect": - return - - async def _abort_upstream(rid: str) -> None: - try: - await upstream_client().request( - "POST", f"{upstream_url}/abort_request", json={"rid": rid}, timeout=10.0 - ) - except Exception: # noqa: BLE001 - logger.warning( - "sidecar_proxy failed to abort upstream rid=%s", rid, exc_info=True - ) - - @app.api_route("/{path:path}", methods=["GET", "POST", "PUT", "DELETE", "PATCH"]) - async def proxy(path: str, request: Request) -> Response: - route = path.strip("/") - if route in BLOCKED_ROUTES: - return JSONResponse( - { - "error": { - "type": "RouteBlocked", - "message": f"/{route} is managed by the weight-sync sidecar and is not proxied", - } - }, - status_code=403, - ) - - body = await request.body() - payload: dict[str, Any] | None = None - if body and request.headers.get("content-type", "").startswith( - "application/json" - ): - parsed = await request.json() - if isinstance(parsed, dict): - payload = parsed - - versioned_route = route in versioned_route_set - policy = ( - WeightVersionPolicy.from_payload(payload) - if versioned_route - else WeightVersionPolicy() - ) - request_id = request.headers.get("x-slime-request-id", "-") - - forward_headers = _forward_headers(request.headers) - - rid: str | None = None - if versioned_route and payload is not None: - payload.pop("weight_version", None) - # Inject a request id so the upstream generation can be aborted if - # the client disconnects; otherwise abandoned requests keep - # generating, holding the commit quiesce point for their full - # remaining length. - rid = payload.get("rid") - if rid is None: - rid = uuid.uuid4().hex - payload["rid"] = rid - - try: - ctx = manager.request_context(policy if versioned_route else None) - async with ctx as start_version: - if versioned_route and payload is not None: - # Stamp the serving version into the engine's KV cache - # namespace: requests admitted under different versions - # structurally cannot share radix-tree prefixes. - run_id = getattr(manager, "current_run_id", None) - user_key = payload.get("extra_key") - if isinstance(user_key, list): - payload["extra_key"] = [ - compose_extra_key(start_version, k, run_id) for k in user_key - ] - else: - payload["extra_key"] = compose_extra_key( - start_version, user_key, run_id - ) - started = asyncio.get_running_loop().time() - if versioned_route and manager.debug_requests: - logger.info( - "sidecar_proxy start request_id=%s path=%s exact=%s current=%s active=%s", - request_id, - path, - policy.exact_version, - start_version, - manager.active_requests, - ) - - async def _upstream_call() -> Any: - request_kwargs: dict[str, Any] = { - "params": request.query_params, - "headers": forward_headers, - } - if payload is not None: - request_kwargs["json"] = payload - else: - request_kwargs["content"] = body - return await upstream_client().request( - request.method, - f"{upstream_url}/{path}", - **request_kwargs, - ) - - upstream_task = asyncio.ensure_future(_upstream_call()) - disconnect_task = asyncio.ensure_future(_watch_disconnect(request)) - try: - await asyncio.wait( - {upstream_task, disconnect_task}, - return_when=asyncio.FIRST_COMPLETED, - ) - if not upstream_task.done(): - upstream_task.cancel() - with contextlib.suppress(BaseException): - await upstream_task - if rid is not None: - await _abort_upstream(rid) - logger.info( - "sidecar_proxy client_disconnect request_id=%s path=%s rid=%s elapsed=%.2fs", - request_id, - path, - rid, - asyncio.get_running_loop().time() - started, - ) - return Response(status_code=499) - finally: - disconnect_task.cancel() - with contextlib.suppress(BaseException): - await disconnect_task - - try: - resp = upstream_task.result() - content_type = resp.headers.get("content-type", "") - if "application/json" not in content_type: - return Response( - content=resp.content, - status_code=resp.status_code, - media_type=content_type or None, - ) - data = resp.json() - except Exception: - if versioned_route: - logger.exception( - "sidecar_proxy error request_id=%s path=%s elapsed=%.2fs current=%s active=%s", - request_id, - path, - asyncio.get_running_loop().time() - started, - manager.current_version, - manager.active_requests, - ) - raise - if versioned_route and manager.debug_requests: - logger.info( - "sidecar_proxy end request_id=%s path=%s status=%s elapsed=%.2fs current=%s active=%s", - request_id, - path, - resp.status_code, - asyncio.get_running_loop().time() - started, - manager.current_version, - manager.active_requests, - ) - # Captured while the request is still pinned, so a commit - # cannot advance the version between serving and reporting. - end_version = manager.current_version - except PolicyViolation as exc: - logger.info( - "sidecar_proxy reject request_id=%s path=%s current=%s error=%s", - request_id, - path, - manager.current_version, - exc.error, - ) - return JSONResponse(exc.error, status_code=409) - - if versioned_route and isinstance(data, dict): - # Driven by the same `versioned_route` flag that gated and stamped the - # request, so injection can't diverge from gating (previously a fixed - # path list here meant /v1/completions got version metadata while - # going ungated, and a custom versioned route got gated but no - # metadata). /generate carries it in meta_info; OpenAI-style routes - # at the top level. - if route == "generate": - meta = data.setdefault("meta_info", {}) - meta["weight_version"] = str(start_version) - meta["weight_version_start"] = start_version - meta["weight_version_end"] = end_version - else: - data["weight_version_start"] = start_version - data["weight_version_end"] = end_version - return JSONResponse(data, status_code=resp.status_code) - - return app - - -def _forward_headers(headers: Any) -> dict[str, str]: - blocked = {"host", "content-length", "connection"} - return {k: v for k, v in headers.items() if k.lower() not in blocked} - - -def _sync_state_value(sync_state: Any) -> Any: - return getattr(sync_state, "value", sync_state) diff --git a/src/stitch/servers/sglang_test.py b/src/stitch/servers/sglang_test.py deleted file mode 100644 index ebad0df..0000000 --- a/src/stitch/servers/sglang_test.py +++ /dev/null @@ -1,323 +0,0 @@ -from __future__ import annotations - -import asyncio -import tempfile -import unittest -from unittest import mock - -from stitch.bulletin import FilesystemBulletinBoard -from stitch.sync import WeightSyncManager - - -class FakeEngine: - backend = "fake" - - def __init__(self) -> None: - self.events: list[str] = [] - - async def flush_cache(self) -> None: - self.events.append("flush") - - async def apply_manifest(self, manifest, version_path) -> None: - self.events.append("apply") - - async def pause_generation(self) -> None: - self.events.append("pause") - - async def continue_generation(self) -> None: - self.events.append("continue") - - -class SidecarProxyTest(unittest.TestCase): - def _client(self, tmp: str): - from fastapi.testclient import TestClient - - from stitch.servers.sglang import create_app - - board = FilesystemBulletinBoard(tmp) - manager = WeightSyncManager(board=board, engine=FakeEngine()) - app = create_app(manager, upstream_url="http://127.0.0.1:9") - return manager, TestClient(app) - - def test_engine_control_routes_are_blocked(self) -> None: - with tempfile.TemporaryDirectory() as tmp: - _manager, client = self._client(tmp) - with client: - for route in ( - "update_weights_from_disk", - "flush_cache", - "pause_generation", - "abort_request", - ): - resp = client.post(f"/{route}", json={}) - self.assertEqual(resp.status_code, 403, route) - self.assertEqual(resp.json()["error"]["type"], "RouteBlocked") - - def test_health_and_server_info(self) -> None: - with tempfile.TemporaryDirectory() as tmp: - _manager, client = self._client(tmp) - with client: - self.assertEqual( - client.get("/health").json(), {"ok": True, "current_version": 0} - ) - info = client.get("/server_info").json() - self.assertEqual(info["current_version"], 0) - self.assertEqual(info["sync_state"], "IDLE") - - def test_version_policy_rejections_are_409(self) -> None: - with tempfile.TemporaryDirectory() as tmp: - manager, client = self._client(tmp) - with client: - manager.current_version = 2 - resp = client.post( - "/generate", - json={"text": "hi", "weight_version": {"exact_version": 1}}, - ) - self.assertEqual(resp.status_code, 409) - self.assertEqual(resp.json()["error"]["type"], "WeightVersionTooOld") - - -class _RecordingUpstream: - """Stand-in for httpx.AsyncClient that records the forwarded payload. - - Fully synchronous (no suspension points) so the proxy's upstream task - completes before its disconnect watcher can observe the test client's - immediate ASGI disconnect. - """ - - last_json: dict | None = None - last_url: str | None = None - last_headers: dict | None = None - - def __init__(self, *args, **kwargs) -> None: - pass - - async def __aenter__(self) -> "_RecordingUpstream": - return self - - async def __aexit__(self, *exc) -> bool: - return False - - async def aclose(self) -> None: - return None - - async def request(self, method, url, **kwargs): - import httpx - - type(self).last_url = url - type(self).last_json = kwargs.get("json") - type(self).last_headers = kwargs.get("headers") - return httpx.Response( - 200, - json={"text": "ok", "meta_info": {"finish_reason": {"type": "length"}}}, - request=httpx.Request(method, url), - ) - - -class _BlockingUpstream: - """Fake upstream whose 'slow' generations block until released, recording - every forwarded payload in order.""" - - calls: list[dict] = [] - release: "asyncio.Event" - - def __init__(self, *args, **kwargs) -> None: - pass - - async def __aenter__(self) -> "_BlockingUpstream": - return self - - async def __aexit__(self, *exc) -> bool: - return False - - async def aclose(self) -> None: - return None - - async def request(self, method, url, **kwargs): - import httpx - - payload = kwargs.get("json") or {} - type(self).calls.append(payload) - if payload.get("text") == "slow": - await type(self).release.wait() - return httpx.Response( - 200, - json={"text": "ok", "meta_info": {}}, - request=httpx.Request(method, url), - ) - - -class SidecarInPlaceCommitTest(unittest.TestCase): - def test_request_crossing_in_place_commit_is_stamped_start_end(self) -> None: - """End-to-end through the sidecar app: an in-flight request crosses an - in-place bulletin-board commit without being drained, finishes with - (start=0, end=1) metadata, and subsequent requests land in the new - extra_key namespace.""" - - async def run() -> None: - import httpx - from httpx import ASGITransport - - from stitch.protocol import VersionManifest - from stitch.servers.sglang import create_app - - with tempfile.TemporaryDirectory() as tmp: - board = FilesystemBulletinBoard(tmp) - engine = FakeEngine() - manager = WeightSyncManager( - board=board, engine=engine, commit_mode="in_place" - ) - app = create_app(manager, upstream_url="http://127.0.0.1:9") - - _BlockingUpstream.calls = [] - _BlockingUpstream.release = asyncio.Event() - - async def wait_for(predicate, timeout: float = 5.0) -> None: - deadline = asyncio.get_running_loop().time() + timeout - while not predicate(): - if asyncio.get_running_loop().time() > deadline: - raise AssertionError("timed out waiting for condition") - await asyncio.sleep(0.01) - - driver = httpx.AsyncClient( - transport=ASGITransport(app=app), base_url="http://sidecar" - ) - async with driver: - with mock.patch("httpx.AsyncClient", _BlockingUpstream): - slow = asyncio.create_task( - driver.post("/generate", json={"text": "slow"}) - ) - await wait_for(lambda: manager.active_requests == 1) - self.assertEqual( - _BlockingUpstream.calls[0]["extra_key"], "wv0;" - ) - - board.publish_manifest( - VersionManifest( - version=1, - base_version=0, - backend="fake", - load_format="noop", - ) - ) - rpc = await driver.post( - "/rpc_sync_from_bulletin_board", json={"target_version": 1} - ) - self.assertTrue(rpc.json()["accepted"]) - - # The commit lands while the request is still in flight: - # in_place mode must not drain non-strict traffic. - await wait_for(lambda: manager.current_version == 1) - self.assertEqual(manager.active_requests, 1) - self.assertEqual(engine.events, ["pause", "apply", "continue"]) - - _BlockingUpstream.release.set() - meta = (await slow).json()["meta_info"] - self.assertEqual(meta["weight_version_start"], 0) - self.assertEqual(meta["weight_version_end"], 1) - - # New admissions are stamped with the new namespace. - await driver.post("/generate", json={"text": "hi"}) - self.assertEqual( - _BlockingUpstream.calls[-1]["extra_key"], "wv1;" - ) - - asyncio.run(run()) - - -class SidecarStampingTest(unittest.TestCase): - def _client(self, tmp: str): - from fastapi.testclient import TestClient - - from stitch.servers.sglang import create_app - - board = FilesystemBulletinBoard(tmp) - manager = WeightSyncManager(board=board, engine=FakeEngine()) - app = create_app(manager, upstream_url="http://127.0.0.1:9") - return manager, TestClient(app) - - def test_generate_payload_is_stamped_with_composed_extra_key(self) -> None: - with tempfile.TemporaryDirectory() as tmp: - _RecordingUpstream.last_json = None - manager, client = self._client(tmp) - with client, mock.patch("httpx.AsyncClient", _RecordingUpstream): - resp = client.post("/generate", json={"text": "hi"}) - self.assertEqual(resp.status_code, 200) - forwarded = _RecordingUpstream.last_json - self.assertEqual(forwarded["extra_key"], "wv0;") - self.assertIn("rid", forwarded) - meta = resp.json()["meta_info"] - self.assertEqual(meta["weight_version_start"], 0) - self.assertEqual(meta["weight_version_end"], 0) - - def test_stamping_composes_user_extra_keys_and_lists(self) -> None: - with tempfile.TemporaryDirectory() as tmp: - manager, client = self._client(tmp) - manager.current_version = 3 - with client, mock.patch("httpx.AsyncClient", _RecordingUpstream): - client.post("/generate", json={"text": "hi", "extra_key": "user-key"}) - self.assertEqual( - _RecordingUpstream.last_json["extra_key"], "wv3;user-key" - ) - - client.post( - "/generate", json={"text": ["a", "b"], "extra_key": ["k1", "k2"]} - ) - self.assertEqual( - _RecordingUpstream.last_json["extra_key"], ["wv3;k1", "wv3;k2"] - ) - - client.post( - "/v1/chat/completions", - json={ - "model": "m", - "messages": [{"role": "user", "content": "hi"}], - }, - ) - self.assertEqual(_RecordingUpstream.last_json["extra_key"], "wv3;") - - def test_unversioned_routes_are_not_stamped(self) -> None: - with tempfile.TemporaryDirectory() as tmp: - manager, client = self._client(tmp) - with client, mock.patch("httpx.AsyncClient", _RecordingUpstream): - _RecordingUpstream.last_json = None - # /v1/completions is not in the default versioned route set, so - # it must be neither stamped (request) nor version-annotated - # (response). Previously the response-metadata block hardcoded - # this path and injected version fields onto an ungated request. - resp = client.post("/v1/completions", json={"model": "m", "prompt": "hi"}) - self.assertNotIn("extra_key", _RecordingUpstream.last_json or {}) - self.assertNotIn("weight_version_start", resp.json()) - self.assertNotIn("weight_version_end", resp.json()) - - -class _CountingUpstream(_RecordingUpstream): - instances = 0 - - def __init__(self, *args, **kwargs) -> None: - super().__init__(*args, **kwargs) - type(self).instances += 1 - - -class SidecarClientPoolTest(unittest.TestCase): - def test_upstream_client_is_pooled_across_requests(self) -> None: - with tempfile.TemporaryDirectory() as tmp: - from fastapi.testclient import TestClient - - from stitch.servers.sglang import create_app - - board = FilesystemBulletinBoard(tmp) - manager = WeightSyncManager(board=board, engine=FakeEngine()) - app = create_app(manager, upstream_url="http://127.0.0.1:9") - - _CountingUpstream.instances = 0 - with TestClient(app) as client, mock.patch("httpx.AsyncClient", _CountingUpstream): - for _ in range(3): - self.assertEqual(client.post("/generate", json={"text": "hi"}).status_code, 200) - # One pooled client served all three requests, plus it was closed - # on shutdown (no per-request construct/teardown). - self.assertEqual(_CountingUpstream.instances, 1) - - -if __name__ == "__main__": - unittest.main() diff --git a/src/stitch/service.py b/src/stitch/service.py new file mode 100644 index 0000000..7aefce4 --- /dev/null +++ b/src/stitch/service.py @@ -0,0 +1,218 @@ +"""The rollout-service runtime: the versioned proxy (``create_app``), the sidecar +entrypoint (``serve``), and cross-replica readiness aggregation (``readiness``). + +Engine- and provider-agnostic: request/response version stamping is delegated to the +Engine, and the proxy forwards everything else to the engine's own HTTP surface. + +No ``from __future__ import annotations`` here: the FastAPI route handlers below are +introspected at runtime, and their ``Request`` type is a create_app-local import — under +stringized annotations FastAPI can't resolve it (it looks only in module globals) and +demotes ``request`` to a required query param, 422-ing every call. +""" + +import asyncio +import contextlib +import logging +import uuid +from collections.abc import Iterable +from contextlib import asynccontextmanager +from typing import Any + +from stitch.engines.base import Engine +from stitch.pools.base import Pool +from stitch.stores.base import Store +from stitch.sync import CommitMode, ConstraintUnmet, Reconciler +from stitch.versions import PoolState, ReplicaState, VersionConstraint + +logger = logging.getLogger(__name__) + +# Engine control routes the body-blind gateway must never reach: a stray call would +# mutate engine state behind the reconciler's back and corrupt its version bookkeeping. +# The sidecar drives these on the engine directly, not through the proxy. +BLOCKED_ROUTES = frozenset( + { + "update_weights_from_disk", + "update_weights_from_distributed", + "update_weights_from_tensor", + "pull_weights", + "flush_cache", + "pause_generation", + "continue_generation", + "abort_request", + } +) + +VERSIONED_ROUTES = ("generate", "v1/chat/completions", "v1/completions") + +# Hop-by-hop / rewritten headers the proxy never forwards upstream. +_DROP_HEADERS = {"host", "content-length", "connection"} + + +def create_app( + reconciler: Reconciler, + engine: Engine, + *, + versioned_routes: Iterable[str] = VERSIONED_ROUTES, + upstream_timeout: float | None = 3600.0, +): + """The versioned rollout proxy. Versioned routes are admitted through the gate + (constraint enforced, serving version captured), stamped by the engine, forwarded, + and the response stamped with the served version. A rejected constraint returns a + retryable 409; a client disconnect aborts the upstream generation.""" + from fastapi import FastAPI, Request + from fastapi.responses import JSONResponse, Response + import httpx + + engine_url = engine.base_url().rstrip("/") + timeout = httpx.Timeout(upstream_timeout, connect=10.0) + versioned = {r.strip("/") for r in versioned_routes} + pooled: dict[str, Any] = {} + + def client() -> Any: + c = pooled.get("client") + if c is None: + c = httpx.AsyncClient(timeout=timeout, trust_env=False) + pooled["client"] = c + return c + + @asynccontextmanager + async def lifespan(_app: FastAPI): + await reconciler.startup() + try: + yield + finally: + await reconciler.shutdown() + c = pooled.pop("client", None) + if c is not None: + await c.aclose() + + app = FastAPI(lifespan=lifespan) + + @app.get("/health") + async def health() -> dict[str, Any]: + return {"ok": True} + + @app.get("/server_info") + async def server_info() -> dict[str, Any]: + return reconciler.server_info() + + @app.post("/wake") + async def wake() -> dict[str, Any]: + reconciler.wake() + return reconciler.server_info() + + async def _watch_disconnect(request: Request) -> None: + while True: + if (await request.receive())["type"] == "http.disconnect": + return + + @app.api_route("/{path:path}", methods=["GET", "POST", "PUT", "DELETE", "PATCH"]) + async def proxy(path: str, request: Request) -> Response: + route = path.strip("/") + if route in BLOCKED_ROUTES: + return JSONResponse( + {"error": {"type": "RouteBlocked", "message": f"/{route} is managed by the sidecar"}}, + status_code=403, + ) + + body = await request.body() + payload: dict[str, Any] | None = None + if body and request.headers.get("content-type", "").startswith("application/json"): + parsed = await request.json() + payload = parsed if isinstance(parsed, dict) else None + + is_versioned = route in versioned + constraint = VersionConstraint.from_payload(payload) if is_versioned else VersionConstraint() + + # A request id lets us abort the upstream generation on client disconnect — + # otherwise an abandoned request keeps generating, holding the quiesce point. + rid = None + if is_versioned and payload is not None: + payload.pop("weight_version", None) + rid = payload.setdefault("rid", uuid.uuid4().hex) + + headers = {k: v for k, v in request.headers.items() if k.lower() not in _DROP_HEADERS} + + try: + async with reconciler.admit(constraint if is_versioned else None) as served: + if is_versioned and payload is not None and served is not None: + engine.stamp_request(payload, served) + kwargs: dict[str, Any] = {"params": request.query_params, "headers": headers} + kwargs["json" if payload is not None else "content"] = payload if payload is not None else body + + upstream_task = asyncio.ensure_future(client().request(request.method, f"{engine_url}/{path}", **kwargs)) + disconnect_task = asyncio.ensure_future(_watch_disconnect(request)) + try: + await asyncio.wait({upstream_task, disconnect_task}, return_when=asyncio.FIRST_COMPLETED) + if not upstream_task.done(): + upstream_task.cancel() + with contextlib.suppress(BaseException): + await upstream_task + if rid is not None: + await _abort(client(), engine_url, rid) + return Response(status_code=499) + finally: + disconnect_task.cancel() + with contextlib.suppress(BaseException): + await disconnect_task + + resp = upstream_task.result() + if "application/json" not in resp.headers.get("content-type", ""): + return Response(content=resp.content, status_code=resp.status_code, + media_type=resp.headers.get("content-type") or None) + data = resp.json() + current = reconciler.applied # captured while still pinned, before any commit advances it + except ConstraintUnmet as exc: + return JSONResponse(exc.error, status_code=409) + + if is_versioned and isinstance(data, dict) and served is not None and current is not None: + engine.stamp_response(data, served, current) + return JSONResponse(data, status_code=resp.status_code) + + return app + + +async def _abort(client: Any, engine_url: str, rid: str) -> None: + try: + await client.request("POST", f"{engine_url}/abort_request", json={"rid": rid}, timeout=10.0) + except Exception: # noqa: BLE001 + logger.warning("failed to abort upstream rid=%s", rid, exc_info=True) + + +def serve( + store: Store, + engine: Engine, + *, + run_id: str | None = None, + commit_mode: CommitMode = "quiesce", + host: str = "0.0.0.0", + port: int = 8000, + debug_requests: bool = False, + reconcile_interval: float = 5.0, +) -> None: + """Run one replica's sidecar: build the Reconciler over the given store+engine + and serve the versioned proxy. The deployment supplies the concrete instances.""" + import uvicorn + + reconciler = Reconciler( + store=store, engine=engine, run_id=run_id, commit_mode=commit_mode, + debug_requests=debug_requests, reconcile_interval=reconcile_interval, + ) + uvicorn.run(create_app(reconciler, engine), host=host, port=port, log_level="info") + + +async def readiness(pool: Pool, *, timeout: float = 15.0) -> PoolState: + """Aggregate every replica's ``/server_info`` into a PoolState (drives the readiness + poll and the smoke check). A replica that fails to answer counts as not ready.""" + import httpx + + async def probe(c: Any, url: str) -> ReplicaState: + try: + resp = await c.get(f"{url.rstrip('/')}/server_info", timeout=timeout) + return ReplicaState.from_dict(resp.json()) + except Exception as exc: # noqa: BLE001 + return ReplicaState(ready=False, reason=str(exc)[:80]) + + async with httpx.AsyncClient(trust_env=False) as c: + states = await asyncio.gather(*(probe(c, url) for url in pool.discover_replicas())) + return PoolState(list(states)) diff --git a/src/stitch/stores/__init__.py b/src/stitch/stores/__init__.py new file mode 100644 index 0000000..63b8284 --- /dev/null +++ b/src/stitch/stores/__init__.py @@ -0,0 +1 @@ +"""Store instances (see base.py for the port).""" diff --git a/src/stitch/stores/base.py b/src/stitch/stores/base.py new file mode 100644 index 0000000..a4bbbeb --- /dev/null +++ b/src/stitch/stores/base.py @@ -0,0 +1,43 @@ +"""The ``Store`` port — where published versions and the ``latest`` pointer live. + +Instances subclass this base and override every method: ``stores/modal_volume.py`` +(Modal Volume). Add S3 / NFS as new subclasses. +""" + +from __future__ import annotations + +from stitch.versions import VersionManifest, VersionRef + + +class Store: + """A versioned checkpoint holder: the version bytes plus the monotonic + pointer / run-epoch coordination. Subclasses override every method.""" + + def refresh(self) -> None: + """Make other hosts' writes visible (Volume reload; no-op if strongly consistent).""" + raise NotImplementedError + + def read_pointer(self) -> VersionRef | None: + """The current ``latest`` pointer, or None if no run has been claimed.""" + raise NotImplementedError + + def advance_pointer(self, ref: VersionRef) -> None: + """Move ``latest`` to ``ref`` — the caller has already run ``decide_pointer_move``.""" + raise NotImplementedError + + def claim(self, run_id: str) -> None: + """Start a new run epoch at base, forking the version space.""" + raise NotImplementedError + + def read_manifest(self, ref: VersionRef) -> VersionManifest: + """The manifest for ``ref``, derived from its on-disk HF index.""" + raise NotImplementedError + + def publish(self, manifest: VersionManifest, files_dir: str) -> None: + """Durably write a version's files; must be visible before ``advance_pointer``.""" + raise NotImplementedError + + def materialize(self, ref: VersionRef) -> str: + """Ensure the version's files are locally readable and return their directory + (hides mount vs download).""" + raise NotImplementedError diff --git a/src/stitch/stores/modal_volume.py b/src/stitch/stores/modal_volume.py new file mode 100644 index 0000000..c95336e --- /dev/null +++ b/src/stitch/stores/modal_volume.py @@ -0,0 +1,118 @@ +"""``ModalVolumeStore`` — the ``Store`` instance backed by a Modal Volume. + +Each run's chain lives under ``//weight_vNNNNNN/`` (run-less: +``/weight_vNNNNNN/``) as HF-safetensors + delta metadata, and a ``latest`` +text file holds the self-identifying pointer identity. Durability is an explicit +Volume commit; cross-host visibility is a reload. With ``volume_name=None`` it is a +plain local directory, so the class is exercisable without Modal. +""" + +from __future__ import annotations + +import contextlib +import os +import tempfile +from pathlib import Path + +from stitch.stores.base import Store +from stitch.versions import VersionManifest, VersionRef + +_POINTER = "latest" + + +class ModalVolumeStore(Store): + def __init__(self, root: str | Path, *, volume_name: str | None = None) -> None: + self.root = Path(root) + self.volume_name = volume_name + + def refresh(self) -> None: + if self.volume_name: + _volume(self.volume_name).reload() + + def read_pointer(self) -> VersionRef | None: + path = self.root / _POINTER + if not path.exists(): + return None + text = path.read_text(encoding="utf-8").strip() + return VersionRef.parse(text) if text else None + + def advance_pointer(self, ref: VersionRef) -> None: + # The caller has already run decide_pointer_move; this is the durable write. + self.root.mkdir(parents=True, exist_ok=True) + _atomic_write(self.root / _POINTER, ref.identity) + if self.volume_name: + _volume(self.volume_name).commit() + + def claim(self, run_id: str) -> None: + if not run_id: + raise ValueError("claim requires a run_id (the run's per-launch epoch token)") + self.advance_pointer(VersionRef(run_id, 0)) + + def read_manifest(self, ref: VersionRef) -> VersionManifest: + return VersionManifest.from_hf_index(self._version_dir(ref), run_id=ref.run_id) + + def publish(self, manifest: VersionManifest, files_dir: str) -> None: + # The framework usually writes straight into the mounted volume, so files_dir + # is already the version dir; copy only when it isn't. Commit so the bytes are + # durable and visible on every host before the pointer moves to them. + target = self._version_dir(manifest.ref) + source = Path(files_dir) + if source.resolve() != target.resolve(): + import shutil + + shutil.copytree(source, target, dirs_exist_ok=True) + if self.volume_name: + _volume(self.volume_name).commit() + + def materialize(self, ref: VersionRef) -> str: + self.refresh() # materialize this host's view before the files are read + return str(self._version_dir(ref)) + + def commit(self) -> None: + """Durably flush pending writes on this host (e.g. one trainer rank's shard of a + version's files). Not part of the Store port — a Modal-Volume affordance the + publish hook uses on non-writer ranks; a no-op without a backing volume.""" + if self.volume_name: + _volume(self.volume_name).commit() + + def _version_dir(self, ref: VersionRef) -> Path: + # ref.identity is exactly the run-partitioned dir name (/weight_vNNNNNN), + # so the pointer spelling and the on-disk layout share one source of truth. + return self.root / ref.identity + + +def pull_weights_pre_read_hook(source_dir: str, target_version: int) -> None: + """Engine-side ``--custom-pull-weights-pre-read-hook``: reload the delta Volume onto + THIS host exactly once so the engine's pull can read the published version. + + One reload, not a loop: the sidecar only pulls after seeing ``target_version`` at the + pointer, and a reload loop thrashes the Modal-v2 mount (turned ~3s delta pulls into + 100-500s ones and tripped the engine's 300s watchdog). Completeness is handled + downstream — the engine size-verifies each staged delta, and a not-yet-materialized + blob fails fast so the sidecar retries with a fresh single reload. The volume name + travels via ``DELTA_VOLUME_NAME`` (set on the serving container).""" + del source_dir # the engine passes it; the reload is by volume name, not path + volume_name = os.environ.get("DELTA_VOLUME_NAME", "") + if not volume_name or target_version <= 0: + return + _volume(volume_name).reload() + + +def _volume(name: str): + import modal + + return modal.Volume.from_name(name, version=2, create_if_missing=True) + + +def _atomic_write(path: Path, text: str) -> None: + fd, tmp = tempfile.mkstemp(dir=str(path.parent), prefix=".tmp-") + try: + with os.fdopen(fd, "w", encoding="utf-8") as f: + f.write(text) + f.flush() + os.fsync(f.fileno()) # durable before the rename (stitch#30) + os.replace(tmp, path) + except BaseException: + with contextlib.suppress(FileNotFoundError): + os.unlink(tmp) + raise diff --git a/src/stitch/stores/modal_volume_test.py b/src/stitch/stores/modal_volume_test.py new file mode 100644 index 0000000..9026133 --- /dev/null +++ b/src/stitch/stores/modal_volume_test.py @@ -0,0 +1,78 @@ +"""ModalVolumeStore harness, driven through the real ``publish_version`` flow. + +Covers everything provable without Modal (volume_name=None → a local dir): the pointer +round-trip, the claim → delta-chain path, and the external-staging copy. The volume-backed +path is validated e2e.""" + +from __future__ import annotations + +import json +import tempfile +from pathlib import Path + +from stitch.publish import publish_version +from stitch.stores.base import Store +from stitch.stores.modal_volume import ModalVolumeStore +from stitch.versions import VersionKind, VersionRef + + +def _write_version(root: Path, ref: VersionRef, *, base: int | None = None, diff: str | None = None) -> str: + d = root / ref.identity + d.mkdir(parents=True) + meta: dict = {"version": ref.version} + if diff: + meta.update({"diff": diff, "base_version": base, "compression": "zstd", "checksum": "xxh3-128"}) + (d / "model.safetensors.index.json").write_text( + json.dumps({"metadata": meta, "weight_map": {"w": "model-00001.safetensors"}}) + ) + (d / "model-00001.safetensors").write_bytes(b"\x00") + return str(d) + + +def test_satisfies_store_port() -> None: + assert isinstance(ModalVolumeStore("/tmp/store"), Store) + + +def test_publish_full_roundtrip() -> None: + with tempfile.TemporaryDirectory() as tmp: + root = Path(tmp) + store = ModalVolumeStore(root) + assert store.read_pointer() is None + vdir = _write_version(root, VersionRef("r1", 1)) # framework wrote it in place + ref = publish_version(store, None, vdir, run_id="r1") + assert ref == VersionRef("r1", 1) + assert store.read_pointer() == VersionRef("r1", 1) # pointer parses back to the ref + man = store.read_manifest(ref) + assert man.kind is VersionKind.FULL and man.base_version is None + assert (Path(store.materialize(ref)) / "model.safetensors.index.json").exists() + + +def test_claim_then_delta_chain() -> None: + with tempfile.TemporaryDirectory() as tmp: + root = Path(tmp) + store = ModalVolumeStore(root) + store.claim("r1") + assert store.read_pointer() == VersionRef("r1", 0) # base before any publish + publish_version(store, None, _write_version(root, VersionRef("r1", 1)), run_id="r1") + publish_version(store, None, _write_version(root, VersionRef("r1", 2), base=1, diff="xor"), run_id="r1") + assert store.read_pointer() == VersionRef("r1", 2) + man = store.read_manifest(VersionRef("r1", 2)) + assert man.kind is VersionKind.DELTA and man.base_version == 1 and man.delta_encoding == "xor" + + +def test_copies_when_files_dir_is_external() -> None: + with tempfile.TemporaryDirectory() as tmp: + root, staging = Path(tmp) / "store", Path(tmp) / "staging" + root.mkdir() + store = ModalVolumeStore(root) + src = _write_version(staging, VersionRef("r1", 1)) # a staging dir, not the store layout + publish_version(store, None, src, run_id="r1") + assert (root / "r1" / "weight_v000001" / "model.safetensors.index.json").exists() + + +if __name__ == "__main__": + tests = [v for k, v in sorted(globals().items()) if k.startswith("test_")] + for t in tests: + t() + print(f" ok {t.__name__}") + print(f"modal_volume harness: {len(tests)} PASS") diff --git a/src/stitch/stores/s3.py b/src/stitch/stores/s3.py new file mode 100644 index 0000000..b09a936 --- /dev/null +++ b/src/stitch/stores/s3.py @@ -0,0 +1,133 @@ +"""``S3Store`` — the ``Store`` instance backed by an S3 (or S3-compatible) bucket. + +Each run's chain lives under ``//weight_vNNNNNN/`` as HF-safetensors + +delta metadata, and a ``/latest`` object holds the self-identifying pointer +identity. S3 is strongly read-after-write consistent, so ``refresh`` is a no-op and a +single ``latest`` PUT is the durable, immediately-visible pointer move (no rename). Unlike +a mounted Volume there is no shared filesystem, so ``materialize`` downloads a version's +chain into ``cache_dir`` — the local path the engine's pull then reads. + +``boto3`` is imported lazily (only when the store is used), so the module loads without it; +credentials/region come from the standard AWS resolution chain, and ``endpoint_url`` targets +S3-compatible stores (MinIO, R2, ...). +""" + +from __future__ import annotations + +import os +from pathlib import Path +from urllib.parse import urlparse + +from stitch.stores.base import Store +from stitch.versions import VersionManifest, VersionRef + +_POINTER = "latest" +_INDEX = "model.safetensors.index.json" + + +class S3Store(Store): + def __init__(self, root: str, *, cache_dir: str | Path, endpoint_url: str | None = None) -> None: + # root is an ``s3://bucket/prefix`` URI (bare ``bucket/prefix`` is accepted too); + # cache_dir is the local directory materialize downloads into (what the engine reads). + parsed = urlparse(root if "://" in root else f"s3://{root}") + self.bucket = parsed.netloc + self.prefix = parsed.path.strip("/") + self.cache_dir = Path(cache_dir) + self.endpoint_url = endpoint_url + self._client = None + + def refresh(self) -> None: + pass # S3 is strongly read-after-write consistent; nothing to make visible + + def read_pointer(self) -> VersionRef | None: + text = self._read_text(self._key(_POINTER)) + return VersionRef.parse(text) if text else None + + def advance_pointer(self, ref: VersionRef) -> None: + # The caller has already run decide_pointer_move; a single-object PUT is the atomic, + # immediately-visible durable write. Monotonicity is the single-writer caller's job. + self._s3().put_object(Bucket=self.bucket, Key=self._key(_POINTER), Body=ref.identity.encode("utf-8")) + + def claim(self, run_id: str) -> None: + if not run_id: + raise ValueError("claim requires a run_id (the run's per-launch epoch token)") + self.advance_pointer(VersionRef(run_id, 0)) + + def read_manifest(self, ref: VersionRef) -> VersionManifest: + # from_hf_index reads a local dir, so land just the small index where materialize + # would put the version and parse it (the heavy shards download in materialize). + index_dir = self.cache_dir / ref.identity + self._download(self._key(ref.identity, _INDEX), index_dir / _INDEX) + return VersionManifest.from_hf_index(index_dir, run_id=ref.run_id) + + def publish(self, manifest: VersionManifest, files_dir: str) -> None: + # Upload every file the trainer wrote under files_dir to the version's key prefix. + # Objects land before the caller moves the pointer; read-after-write makes them visible. + src = Path(files_dir) + for path in sorted(p for p in src.rglob("*") if p.is_file()): + key = self._key(manifest.ref.identity, path.relative_to(src).as_posix()) + self._s3().upload_file(str(path), self.bucket, key) + + def materialize(self, ref: VersionRef) -> str: + # No shared mount: sync the run's published versions (the delta chain back to the + # nearest anchor) into the local cache so the engine's pull can read them, then return + # the version dir. Its parent is the run dir the pull walks back over. + self._sync(self._key(ref.run_id), self.cache_dir / ref.run_id) + return str(self.cache_dir / ref.identity) + + def commit(self) -> None: + """No-op: S3 writes are immediately durable. Present only so the shared publish hook's + ``store.commit()`` (a Modal-Volume affordance for non-writer ranks) is safe here too.""" + + # ── S3 helpers ─────────────────────────────────────────────────────────────── + def _key(self, *parts: str) -> str: + return "/".join(p.strip("/") for p in (self.prefix, *parts) if p.strip("/")) + + def _s3(self): + if self._client is None: + import boto3 + + self._client = boto3.client("s3", endpoint_url=self.endpoint_url) + return self._client + + def _read_text(self, key: str) -> str | None: + s3 = self._s3() + try: + return s3.get_object(Bucket=self.bucket, Key=key)["Body"].read().decode("utf-8").strip() + except s3.exceptions.NoSuchKey: + return None + + def _download(self, key: str, dest: Path) -> None: + dest.parent.mkdir(parents=True, exist_ok=True) + self._s3().download_file(self.bucket, key, str(dest)) + + def _sync(self, key_prefix: str, dest_root: Path) -> None: + """Download every object under ``key_prefix`` into ``dest_root``, skipping any file + already present at the same size — so re-materializing a chain only fetches new versions.""" + s3 = self._s3() + for page in s3.get_paginator("list_objects_v2").paginate(Bucket=self.bucket, Prefix=key_prefix + "/"): + for obj in page.get("Contents", []): + rel = obj["Key"][len(key_prefix) + 1:] + if not rel: # the prefix "directory" placeholder, if any + continue + dest = dest_root / rel + if dest.exists() and dest.stat().st_size == obj["Size"]: + continue + dest.parent.mkdir(parents=True, exist_ok=True) + s3.download_file(self.bucket, obj["Key"], str(dest)) + + +def pull_weights_pre_read_hook(source_dir: str, target_version: int) -> None: + """Engine-side ``--custom-pull-weights-pre-read-hook``: download the run's published + versions onto THIS host's ``source_dir`` so the engine's pull can read them — the S3 + analogue of the Modal-Volume reload, for engines that span hosts not sharing the cache. + + ``source_dir`` is the run dir (``/``); its basename is the run. The bucket/ + prefix travel via ``DELTA_S3_URI`` (and optional ``DELTA_S3_ENDPOINT_URL``) on the serving + container. Guarded on ``target_version > 0`` — version 0 is the engine's own served base.""" + uri = os.environ.get("DELTA_S3_URI", "") + if not uri or target_version <= 0: + return + run_id = Path(source_dir).name + store = S3Store(uri, cache_dir=Path(source_dir).parent, endpoint_url=os.environ.get("DELTA_S3_ENDPOINT_URL") or None) + store.materialize(VersionRef(run_id, target_version)) diff --git a/src/stitch/stores/s3_test.py b/src/stitch/stores/s3_test.py new file mode 100644 index 0000000..89e0ff8 --- /dev/null +++ b/src/stitch/stores/s3_test.py @@ -0,0 +1,108 @@ +"""Harness for S3Store over an in-memory fake S3 (no boto3 / no network). + +Exercises the pointer round-trip and the publish -> read_manifest -> materialize path the +reconciler drives, so the store's S3 key layout + download-to-cache logic is verified without +a bucket. Run directly: PYTHONPATH=src python src/stitch/stores/s3_test.py +""" + +from __future__ import annotations + +import io +import json +import tempfile +from pathlib import Path + +from stitch.stores.s3 import S3Store +from stitch.versions import VersionKind, VersionManifest, VersionRef + + +class _FakeS3: + """Dict-backed stand-in for a boto3 S3 client (only the calls S3Store makes).""" + + class exceptions: # noqa: N801 — mirrors boto3 client.exceptions.NoSuchKey + class NoSuchKey(Exception): + pass + + def __init__(self) -> None: + self.objs: dict[tuple[str, str], bytes] = {} + + def put_object(self, Bucket, Key, Body): + self.objs[(Bucket, Key)] = Body + + def get_object(self, Bucket, Key): + if (Bucket, Key) not in self.objs: + raise self.exceptions.NoSuchKey() + return {"Body": io.BytesIO(self.objs[(Bucket, Key)])} + + def upload_file(self, Filename, Bucket, Key): + self.objs[(Bucket, Key)] = Path(Filename).read_bytes() + + def download_file(self, Bucket, Key, Filename): + Path(Filename).parent.mkdir(parents=True, exist_ok=True) + Path(Filename).write_bytes(self.objs[(Bucket, Key)]) + + def get_paginator(self, _name): + objs = self.objs + + class _Paginator: + def paginate(self, Bucket, Prefix): + contents = [{"Key": k, "Size": len(v)} for (b, k), v in objs.items() if b == Bucket and k.startswith(Prefix)] + return [{"Contents": contents}] + + return _Paginator() + + +def _store(tmp: str) -> S3Store: + store = S3Store("s3://bkt/deltas", cache_dir=Path(tmp) / "cache") + store._client = _FakeS3() # inject the fake so _s3() never imports boto3 + return store + + +def _write_local_version(root: Path, ref: VersionRef) -> str: + """A version dir as the trainer would write it locally (delta index + one shard).""" + d = root / ref.identity + d.mkdir(parents=True) + (d / "model.safetensors.index.json").write_text( + json.dumps({ + "metadata": {"version": ref.version, "base_version": ref.version - 1, "diff": "xor", + "compression": "zstd", "checksum": "xxh3-128"}, + "weight_map": {"w": "model-00001-of-00001.safetensors"}, + }) + ) + (d / "model-00001-of-00001.safetensors").write_bytes(b"\x00\x11\x22\x33") + return str(d) + + +def test_s3_pointer_roundtrip() -> None: + with tempfile.TemporaryDirectory() as tmp: + store = _store(tmp) + assert store.read_pointer() is None + store.claim("run-x") + assert store.read_pointer() == VersionRef("run-x", 0) + store.advance_pointer(VersionRef("run-x", 3)) + assert store.read_pointer() == VersionRef("run-x", 3) + + +def test_s3_publish_manifest_materialize() -> None: + with tempfile.TemporaryDirectory() as tmp: + store = _store(tmp) + ref = VersionRef("run-x", 1) + local = _write_local_version(Path(tmp) / "trainer", ref) # trainer's local write + store.publish(VersionManifest.from_hf_index(local, run_id="run-x"), local) + + manifest = store.read_manifest(ref) # reads the index back from S3 + assert manifest.ref == ref + assert manifest.kind is VersionKind.DELTA + assert manifest.delta_encoding == "xor" and manifest.base_version == 0 + + version_dir = Path(store.materialize(ref)) # downloads the chain into the cache + assert version_dir == store.cache_dir / ref.identity + assert (version_dir / "model-00001-of-00001.safetensors").read_bytes() == b"\x00\x11\x22\x33" + + +if __name__ == "__main__": + tests = [v for k, v in sorted(globals().items()) if k.startswith("test_")] + for t in tests: + t() + print(f" ok {t.__name__}") + print(f"s3 store harness: {len(tests)} PASS") diff --git a/src/stitch/sync.py b/src/stitch/sync.py index 7b39e99..0a330bb 100644 --- a/src/stitch/sync.py +++ b/src/stitch/sync.py @@ -1,183 +1,144 @@ -"""Generic rollout server sync manager.""" +"""The per-replica sync brain: an ``AdmissionGate`` (gate rollout requests on the +served version) and a ``Reconciler`` (converge the replica to the store's pointer). + +They share one lock and one ``applied`` version so that reporting stays correct: +a request's constraint is checked and its serving version captured under the same +lock the committer holds across the weight apply *and* the version flip. +""" from __future__ import annotations import asyncio import logging +import time from collections import defaultdict -from collections.abc import Awaitable, Callable, Mapping -from contextlib import asynccontextmanager -from typing import Any, Literal, Protocol - -from stitch.bulletin import BulletinBoard -from stitch.protocol import ( - EngineAdapter, - SyncState, - VersionManifest, - WeightVersionPolicy, - evaluate_version_policy, -) +from collections.abc import Awaitable, Callable +from contextlib import asynccontextmanager, suppress +from typing import Any, Literal +from stitch.engines.base import Engine +from stitch.stores.base import Store +from stitch.versions import SyncState, VersionConstraint, VersionKind, VersionRef logger = logging.getLogger(__name__) + CommitMode = Literal["quiesce", "in_place"] -class PolicyViolation(Exception): - """A request's weight-version policy cannot be satisfied by this server.""" +class ConstraintUnmet(Exception): + """A request's version constraint cannot be met by this replica (a retryable 409).""" - def __init__(self, error: Mapping[str, Any]) -> None: - super().__init__(error["error"]["message"]) + def __init__(self, error: dict[str, Any]) -> None: + super().__init__(error["message"]) self.error = error -class RolloutSyncManager(Protocol): - """The surface stitch.servers.sglang.create_app drives. WeightSyncManager - implements it, and both the bulletin-board sidecar and the hot-load provider - sidecar build one — they differ only in board backend and front-door surface. +class AdmissionGate: + """Request admission + the commit gate. - create_app also probes optional members defensively (shutdown_sync, and the - sync-route trio queue_sync/queued_target_version/sync_state); those are not - part of this required contract. - """ - - debug_requests: bool - current_version: int - active_requests: int - - async def startup_sync(self) -> None: ... - - async def server_info(self) -> dict[str, Any]: ... - - def request_context(self, policy: WeightVersionPolicy | None = None) -> Any: ... - - -class RolloutAdmissionGate: - """Shared request-admission + commit gate for rollout sync managers. - - Owns the active-request accounting and the ``_committing`` gate, and - centralizes the lock discipline that keeps version reporting correct: a - request's policy is checked and its serving version captured under the same - ``_active_cond`` acquisition the committer uses, and commits hold the gate - across the engine apply *and* the version advance (cleared only after). - Subclasses provide ``current_version`` and override the hooks for their - policy / admission / exact-pin specifics. ``WeightSyncManager`` composes it - (so the bulletin-board and hot-load provider sidecars, both WeightSyncManager, - share the gate semantics and the P0.1 commit-window fix in one place). + ``quiesce`` drains all in-flight requests and flushes before applying; ``in_place`` + pauses the engine and lets non-exact requests keep decoding on stale KV (only exact + pins are drained) — but only across *compatible* commits: an incompatible transition + (a run switch's base reset) commits with ``drain_all=True``, which drains and gates + everything regardless of mode. Either way the committer holds the lock across the + apply and the version flip, so no request is ever admitted seeing the stale version + on new weights. """ def __init__(self, *, commit_mode: CommitMode = "quiesce") -> None: self.commit_mode = commit_mode - self._active_cond = asyncio.Condition() - self._active_requests = 0 + self.applied: VersionRef | None = None + self._cond = asyncio.Condition() + self._active = 0 self._committing = False - # Exact-version pins in flight, summed across versions, so a commit can - # wait for strict traffic to drain. Tracked on the gate (not per - # subclass) so the bulletin-board manager and the hot-load shim share - # the in_place commit semantics. + self._drain_all = False self._exact_inflight: dict[int, int] = defaultdict(int) @property def active_requests(self) -> int: - return self._active_requests - - @property - def inflight_exact_versions(self) -> dict[str, int]: - return {str(version): count for version, count in sorted(self._exact_inflight.items()) if count} + return self._active - def _admission_gated(self, policy: WeightVersionPolicy) -> bool: + def _gated(self, c: VersionConstraint) -> bool: if not self._committing: return False - if self.commit_mode == "in_place": - # Non-strict requests cross commits freely: stamped with the version - # current at admission, a mislabel around the commit is old-era - # impurity only. Exact pins must not cross. - return policy.exact_version is not None - return True - - def _policy_error(self, policy: WeightVersionPolicy) -> dict[str, Any] | None: - return evaluate_version_policy(self.current_version, policy) - - def _on_admit(self, policy: WeightVersionPolicy) -> None: - if policy.exact_version is not None: - self._exact_inflight[int(policy.exact_version)] += 1 - - def _on_release(self, policy: WeightVersionPolicy) -> None: - if policy.exact_version is not None: - key = int(policy.exact_version) - self._exact_inflight[key] -= 1 - if not self._exact_inflight[key]: - del self._exact_inflight[key] - - def _on_policy_violation(self, error: dict[str, Any]) -> None: - """Hook run under the lock when admission is rejected.""" + # in_place lets non-exact requests cross a compatible commit — they were stamped + # with the admission version, so a mislabel is only old-era impurity. Exact pins + # can't cross, and nothing crosses an incompatible (drain_all) transition. + if self._drain_all or self.commit_mode != "in_place": + return True + return c.exact_version is not None + + def _rejection(self, c: VersionConstraint) -> dict[str, Any] | None: + applied = self.applied.version if self.applied else None + if c.satisfied_by(applied): + return None + target = c.exact_version if c.exact_version is not None else c.min_version + return { + "type": "WeightVersionNotReady", + "target_version": target, + "applied": applied, + "message": f"served version {applied} does not satisfy {c}", + } + + def _on_reject(self, error: dict[str, Any]) -> None: + """Hook, run under the lock, when a request is rejected.""" def _commit_ready(self) -> bool: - """The predicate a commit waits on before closing the admission gate: - in_place drains only exact pins (so at most one exact version is ever - live); quiesce drains all in-flight proxied requests.""" - if self.commit_mode == "in_place": + if self.commit_mode == "in_place" and not self._drain_all: return not any(self._exact_inflight.values()) - return self._active_requests == 0 + return self._active == 0 @asynccontextmanager - async def request_context(self, policy: WeightVersionPolicy | None = None): - """Admit one request: gate on in-progress commits, enforce the policy, - and capture the serving version — all under one ``_active_cond`` - acquisition, so the yielded version is exactly what the engine serves - the request on. Raises :class:`PolicyViolation` when the policy fails. - """ - policy = policy or WeightVersionPolicy() - async with self._active_cond: - await self._active_cond.wait_for(lambda: not self._admission_gated(policy)) - error = self._policy_error(policy) + async def admit(self, constraint: VersionConstraint | None = None): + """Admit one request under a single lock acquisition, yielding the version it + is served on. Raises :class:`ConstraintUnmet` if the constraint can't be met.""" + c = constraint or VersionConstraint() + async with self._cond: + await self._cond.wait_for(lambda: not self._gated(c)) + error = self._rejection(c) if error is not None: - self._on_policy_violation(error) - raise PolicyViolation(error) - start_version = self.current_version - self._active_requests += 1 - self._on_admit(policy) + self._on_reject(error) + raise ConstraintUnmet(error) + served = self.applied + self._active += 1 + if c.exact_version is not None: + self._exact_inflight[c.exact_version] += 1 try: - yield start_version + yield served finally: - async with self._active_cond: - self._active_requests -= 1 - self._on_release(policy) - self._active_cond.notify_all() - - async def _begin_commit(self, ready: Callable[[], bool]) -> None: - """Wait for the quiesce predicate, then close the admission gate under - the same lock acquisition (so no request slips in between).""" - async with self._active_cond: - await self._active_cond.wait_for(ready) - self._committing = True - - async def _end_commit(self) -> None: - async with self._active_cond: - self._committing = False - self._active_cond.notify_all() - - async def commit_version( + async with self._cond: + self._active -= 1 + if c.exact_version is not None: + self._exact_inflight[c.exact_version] -= 1 + if not self._exact_inflight[c.exact_version]: + del self._exact_inflight[c.exact_version] + self._cond.notify_all() + + async def commit( self, *, apply: Callable[[], Awaitable[None]], on_applied: Callable[[], None], pause: Callable[[], Awaitable[None]] | None = None, resume: Callable[[], Awaitable[None]] | None = None, + drain_all: bool = False, ) -> None: - """Drive one commit through the gate: wait for the commit point and - close the admission gate, apply the new weights, advance the served - version (``on_applied``) while the gate is still held, then reopen. - - In ``in_place`` mode the engine is paused around the apply and resumed in - a finally, with the version advanced before resume so new admissions see - the new namespace. On failure the gate (and pause) are unwound and the - served version is left unchanged — ``on_applied`` runs only after a - successful apply. Both the bulletin-board manager and the hot-load shim - commit through here, so the gate sequencing lives in one place. - """ - await self._begin_commit(self._commit_ready) + """Wait for the commit point, close the gate, apply, flip the served version + (``on_applied``) while the gate is held, then reopen. ``on_applied`` runs only + after a successful apply; in ``in_place`` the flip happens before ``resume``. + ``drain_all`` marks an incompatible transition (a base reset): drain and gate + every request regardless of mode — rolling requests may cross a compatible + weight update, never a change of lineage (stitch#32).""" + # Close admission BEFORE draining (stitch#32): with the gate still open during + # the drain, new requests keep being admitted and an in_place non-exact request + # can straddle a base reset. Set _committing first, then wait for the drain. + async with self._cond: + self._committing = True + self._drain_all = drain_all + self._cond.notify_all() try: + async with self._cond: + await self._cond.wait_for(self._commit_ready) if self.commit_mode == "in_place" and pause is not None: await pause() try: @@ -190,200 +151,234 @@ async def commit_version( await apply() on_applied() finally: - await self._end_commit() - + async with self._cond: + self._committing = False + self._drain_all = False + self._cond.notify_all() -class WeightSyncManager(RolloutAdmissionGate): - """Local rollout server sync manager. - Commit modes: - - - ``quiesce`` (default): wait for active proxied requests per - ``commit_wait_policy``, flush the engine cache, then apply. Safe on any - engine build. - - ``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. - """ +class Reconciler(AdmissionGate): + """Converges one replica to the store's ``latest`` pointer: stage the chain, + reload once, flip the served version. A run change resets to base (the engine + reseeds), so a run's chain is never mistaken for another's.""" def __init__( self, *, - board: BulletinBoard, - engine: EngineAdapter, + store: Store, + engine: Engine, run_id: str | None = None, commit_mode: CommitMode = "quiesce", debug_requests: bool = False, + reconcile_interval: float = 5.0, ) -> None: super().__init__(commit_mode=commit_mode) - self.board = board + self.store = store self.engine = engine - self.run_id = run_id + self.run_id = run_id # static replica label for server_info, not the active chain self.debug_requests = debug_requests - # The active *chain* identity, parsed from the pointer (`/weight_vN`). - # Distinct from `self.run_id` (the static replica label). None for the - # run-less stitch layout, where the manager never switches. - self.current_run_id: str | None = None - self.current_version = 0 - self.latest_seen_version = 0 - self.queued_target_version: int | None = None + self.reconcile_interval = reconcile_interval self.sync_state = SyncState.IDLE - self.last_sync_error: str | None = None - self._sync_task: asyncio.Task[None] | None = None - self._sync_lock = asyncio.Lock() - - async def startup_sync(self) -> None: - prepare = getattr(self.engine, "prepare", None) - if prepare is not None: - await prepare() - # Converge to the board's current (run_id, latest): _sync_once follows a - # run switch (re-materialize + reset to v0) and replays the chain. - await self.sync_to() - - async def server_info(self) -> dict[str, Any]: + self.last_error: str | None = None + self.metrics: dict[str, Any] = {} + self._task: asyncio.Task[None] | None = None + self._prefetch_task: asyncio.Task[None] | None = None + self._prefetch_done = False + self._prefetch_error: str | None = None + self._periodic_task: asyncio.Task[None] | None = None + self._lock = asyncio.Lock() + + async def startup(self) -> None: + # Seed the base into the host-local checkpoint in the background so the first real + # stage() is delta-only, not a full base copy on the trainer's critical path. The + # replica serves throughout (this is disk-only); a stage() waits for this to finish + # before writing the same dir (see _reconcile_once_measured), so the two never race. + self._prefetch_task = asyncio.create_task(self._prefetch_base()) + await self.reconcile() + if self.reconcile_interval > 0: + self._periodic_task = asyncio.create_task(self._periodic_reconcile()) + + async def shutdown(self) -> None: + for task in (self._periodic_task, self._prefetch_task): + if task is not None: + task.cancel() + with suppress(BaseException): + await task + + async def _periodic_reconcile(self) -> None: + # Convergence backstop: re-check the pointer every interval so a replica that never + # got a wake still catches up — a cold start that raced the last publish/wake, or a + # best-effort wake that was lost, would otherwise sit stale until the NEXT publish. + # wake() is idempotent and non-cancelling, so an already-caught-up tick is a cheap + # pointer read and an in-flight reconcile is left to finish. + while True: + await asyncio.sleep(self.reconcile_interval) + self.wake() + + async def _prefetch_base(self) -> None: + try: + await self.engine.prefetch() + self._prefetch_done = True + except Exception as exc: # noqa: BLE001 — best-effort; the first real pull falls back to a full copy + self._prefetch_error = str(exc) + logger.exception("base prefetch failed; first sync will pay the full base copy") + + def server_info(self) -> dict[str, Any]: + # Superset of ReplicaState's wire keys (ready/applied/sync_state/reason) plus + # diagnostics; the readiness poll reads the first four. Two-level readiness: `ready` + # means serveable (weights on the GPU, version applied) and does NOT wait on the base + # prefetch — the replica serves while the prefetch runs. `prefetch_done`/`prefetch_error` + # expose whether the O(delta) fast path is primed, so a slow/failed seed is observable. return { + "ready": self.applied is not None and self.sync_state is not SyncState.ERROR, + "applied": self.applied.identity if self.applied else None, + "sync_state": self.sync_state.value, + "reason": self.last_error, "run_id": self.run_id, - "backend": self.engine.backend, "commit_mode": self.commit_mode, - "current_run_id": self.current_run_id, - "current_version": self.current_version, - "latest_seen_version": self.latest_seen_version, - "queued_target_version": self.queued_target_version, - "sync_state": self.sync_state.value, - "last_sync_error": self.last_sync_error, - "sync_task_active": self._sync_task is not None and not self._sync_task.done(), - "active_requests": self._active_requests, - "inflight_exact_versions": self.inflight_exact_versions, + "active_requests": self._active, + "prefetch_done": self._prefetch_done, + "prefetch_error": self._prefetch_error, + "metrics": self.metrics, } - def _on_policy_violation(self, error: dict[str, Any]) -> None: - if error["error"]["type"] == "WeightVersionNotReady": - self.queue_sync(error["error"]["target_version"]) - - async def validate_policy(self, policy: WeightVersionPolicy) -> tuple[bool, int, Mapping[str, Any] | None]: - """Advisory pre-check. The authoritative check is in request_context.""" - error = self._policy_error(policy) - if error is not None and error["error"]["type"] == "WeightVersionNotReady": - self.queue_sync(error["error"]["target_version"]) - return error is None, self.current_version, error - - def queue_sync(self, target_version: int | None = None) -> None: - run_id, latest = self.board.read_latest() - self.latest_seen_version = max(self.latest_seen_version, latest) - hint = int(target_version) if target_version is not None else 0 - # A run change is always work (the pointer moved to a fresh chain, even if - # its version number is lower than what we serve); within a run it's work - # only when the target exceeds what we've applied. - needs_switch = run_id != self.current_run_id - if not needs_switch and max(latest, hint) <= self.current_version: - return - self.queued_target_version = max(latest, hint, self.queued_target_version or 0) - if self.sync_state is SyncState.IDLE: - self.sync_state = SyncState.QUEUED - if self._sync_task is None or self._sync_task.done(): - self._sync_task = asyncio.get_running_loop().create_task(self.sync_to()) - - async def sync_to(self, target_version: int | None = None) -> None: - # target_version is an optional lower-bound hint from a waker; the - # authoritative target is always the board's current (run_id, latest), so - # we converge to it — and follow a run switch — regardless of the hint. - if target_version is not None: - self.queued_target_version = max(int(target_version), self.queued_target_version or 0) + def _on_reject(self, error: dict[str, Any]) -> None: + self.wake() # a version-not-ready 409 is our cue to catch up + def wake(self) -> None: + """Nudge a reconcile now (a publish wake or a 409). Non-cancelling: starts a + task only if none is running; the running loop re-reads the authoritative pointer.""" + if self._task is None or self._task.done(): + self._task = asyncio.get_running_loop().create_task(self.reconcile()) + + async def reconcile(self) -> None: + """Loop until caught up to the store's (run, latest); on error, record it and + stop — a later wake or poll retries.""" while True: try: - reached = await self._sync_once() + caught_up = await self._reconcile_once() + if caught_up: + # A publish can land mid-commit; re-check before idling. + await asyncio.to_thread(self.store.refresh) + if self._behind(self.store.read_pointer()): + caught_up = False except Exception as exc: # noqa: BLE001 - self.last_sync_error = str(exc) + self.last_error = str(exc) self.sync_state = SyncState.ERROR - logger.exception("Weight sync failed") + logger.exception("reconcile failed") return - - if reached: - self.queued_target_version = None + if caught_up: self.sync_state = SyncState.IDLE return await asyncio.sleep(1.0) - async def _switch_run(self, new_run_id: str | None) -> None: - """Rebase onto a new run's chain: re-materialize base and reset to v0. - - A new run forks at base (slime sets ``base_version=000000`` for v1 of every - run), so switching chains means discarding the current weights and replaying - the new chain. The re-materialize runs through the commit gate (pause → - reset → resume) exactly like a version commit, so no in-flight request - decodes across the weight wipe. - """ - logger.info( - "run change %r -> %r: re-materializing base, resetting to v0", - self.current_run_id, - new_run_id, - ) - reset = getattr(self.engine, "reset", None) - was_patched = self.current_version > 0 - - async def _apply() -> None: - if reset is not None and was_patched: - await reset() - - def _applied() -> None: - self.current_run_id = new_run_id - self.current_version = 0 - self.latest_seen_version = 0 - self.last_sync_error = None - - await self.commit_version( - apply=_apply, - on_applied=_applied, - pause=self.engine.pause_generation, - resume=self.engine.continue_generation, + def _behind(self, pointer: VersionRef | None) -> bool: + if pointer is None: + return False + if self.applied is None or pointer.run_id != self.applied.run_id: + return True + return pointer.version > self.applied.version + + async def _reconcile_once(self) -> bool: + async with self._lock: + m: dict[str, Any] = {} + try: + return await self._reconcile_once_measured(m) + except Exception as exc: + m["error"] = str(exc) + raise + finally: + if len(m) > 1 or "error" in m: # a no-work pass leaves the last breakdown alone + m["at"] = time.time() + self.metrics = m + + async def _reconcile_once_measured(self, m: dict[str, Any]) -> bool: + # refresh() may block on I/O; offload it so it never stalls the serving loop. + await asyncio.to_thread(self.store.refresh) + pointer = self.store.read_pointer() + if pointer is None: + return True + if self.applied is None or pointer.run_id != self.applied.run_id: + await self._switch_run(pointer.run_id) + if not self._behind(pointer): + return True + + self.sync_state = SyncState.PREFETCHING + self.last_error = None + target = self.store.read_manifest(pointer) + source_dir = await asyncio.to_thread(self.store.materialize, pointer) + m["target_version"] = pointer.version + + # Serialize the base prefetch with this stage: both write the host-local checkpoint, + # so wait for the prefetch here rather than rely on the engine's file lock — prefetch + # then stage is strictly ordered. If the prefetch failed, this stage seeds the full + # base itself (recorded so the slow fallback is visible, not silent). + if self._prefetch_task is not None and not self._prefetch_task.done(): + await self._prefetch_task + if self._prefetch_error is not None: + m["paid_base_copy"] = True + + # Stage (host-side apply — the engine walks back to the nearest anchor and + # replays the whole tail) runs while serving; the gate covers only the reload. + self.sync_state = SyncState.PREPARING + t = time.perf_counter() + await self.engine.stage(target, source_dir) + m["stage_s"] = round(time.perf_counter() - t, 3) + + def on_applied() -> None: + self.applied = pointer + + if target.kind is VersionKind.DELTA and not target.files: + # An empty delta advances the version but needs no reload — stale KV stays + # valid on byte-identical weights. + await self.commit(apply=self._commit_noop, on_applied=on_applied) + m["skipped_reload"] = True + else: + async def apply() -> None: + if self.commit_mode != "in_place": + await self.engine.flush() + self.sync_state = SyncState.COMMITTING + t2 = time.perf_counter() + await self.engine.commit(pointer) + m["commit_s"] = round(time.perf_counter() - t2, 3) + + await self.commit( + apply=apply, + on_applied=on_applied, + pause=self.engine.pause, + resume=self.engine.resume, + ) + self.sync_state = SyncState.PREFETCHING + return not self._behind(self.store.read_pointer()) # a mid-pass publish is next tick's work + + async def _commit_noop(self) -> None: + self.sync_state = SyncState.COMMITTING + + async def _switch_run(self, new_run: str | None) -> None: + """Rebase onto a new run: reset to base (the engine reseeds on the next stage). + Commits with ``drain_all`` — a base reset is an incompatible transition, so even + in ``in_place`` mode no rolling request is admitted during, or decodes across, + the wipe.""" + old_run = self.applied.run_id if self.applied else None + logger.info("run change %r -> %r: resetting to base", old_run, new_run) + # Reset when the engine holds patched weights (v>0), or a prior pass errored and + # may have left the local checkpoint dirty even at v0 (stitch#32). + was_patched = self.applied is not None and ( + self.applied.version > 0 or self.last_error is not None ) - async def _sync_once(self) -> bool: - async with self._sync_lock: - await self.board.refresh() - run_id, latest = self.board.read_latest() - if run_id != self.current_run_id: - await self._switch_run(run_id) - self.latest_seen_version = max(self.latest_seen_version, latest) - if latest <= self.current_version: - return True - - self.sync_state = SyncState.PREFETCHING - self.last_sync_error = 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: - raise RuntimeError( - f"cannot apply version {version}: manifest base " - f"{manifest.base_version} != current {self.current_version}" - ) - 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 - - # 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. - return self.current_version >= latest + async def apply() -> None: + if was_patched: + await self.engine.reset() + + def on_applied() -> None: + self.applied = VersionRef(new_run, 0) + self.last_error = None + + await self.commit( + apply=apply, + on_applied=on_applied, + pause=self.engine.pause, + resume=self.engine.resume, + drain_all=True, + ) diff --git a/src/stitch/sync_test.py b/src/stitch/sync_test.py index 8b65fea..58486fa 100644 --- a/src/stitch/sync_test.py +++ b/src/stitch/sync_test.py @@ -1,393 +1,443 @@ +"""In-memory core harness (the Phase-1 gate): the real Reconciler + AdmissionGate +against fake Store / Engine — no Modal, sglang, or GPU. Runnable directly +(``python src/stitch/sync_test.py``) or under pytest.""" + from __future__ import annotations import asyncio -import json -import tempfile -import unittest -from pathlib import Path - -from stitch.bulletin import FilesystemBulletinBoard -from stitch.protocol import SyncState, VersionManifest, WeightVersionPolicy -from stitch.sync import PolicyViolation, RolloutAdmissionGate, WeightSyncManager - - -def _write_slime_version(base: Path, version: int, prev: int) -> None: - vdir = base / f"weight_v{version:06d}" - vdir.mkdir(parents=True) - (vdir / "model.safetensors.index.json").write_text( - json.dumps( - { - "metadata": {"version": f"{version:06d}", "base_version": f"{prev:06d}"}, - "weight_map": {"w": "model-00001-of-00001.safetensors"}, - } - ), - encoding="utf-8", - ) +import queue +import threading + +from stitch.engines.base import Engine +from stitch.stores.base import Store +from stitch.sync import ConstraintUnmet, Reconciler +from stitch.versions import ( + VersionConstraint, + VersionKind, + VersionManifest, + VersionRef, + SyncState, +) + + +class FakeStore(Store): + def __init__(self, pointer: VersionRef | None = None, *manifests: VersionManifest) -> None: + self._pointer = pointer + self._manifests = {(m.ref.run_id, m.ref.version): m for m in manifests} + self.refreshed = 0 + + def refresh(self) -> None: + self.refreshed += 1 + + def read_pointer(self) -> VersionRef | None: + return self._pointer + + def read_manifest(self, ref: VersionRef) -> VersionManifest: + return self._manifests[(ref.run_id, ref.version)] + + def materialize(self, ref: VersionRef) -> str: + return f"/fake/{ref.identity}" + + def advance_pointer(self, ref: VersionRef) -> None: + self._pointer = ref + + def claim(self, run_id: str) -> None: + self._pointer = VersionRef(run_id, 0) + def publish(self, manifest: VersionManifest, files_dir: str) -> None: + self._manifests[(manifest.ref.run_id, manifest.ref.version)] = manifest -class FakeEngine: - backend = "fake" +class FakeEngine(Engine): def __init__(self) -> None: - self.flushes = 0 - self.applies: list[tuple[int, str]] = [] - self.events: list[str] = [] - self.apply_gate: asyncio.Event | None = None - self.apply_started: asyncio.Event = asyncio.Event() - - async def flush_cache(self) -> None: - self.flushes += 1 - self.events.append("flush") - - async def apply_manifest(self, manifest: VersionManifest, version_path: str) -> None: - self.apply_started.set() - if self.apply_gate is not None: - await self.apply_gate.wait() - self.applies.append((manifest.version, version_path)) - self.events.append("apply") + self.calls: list[str] = [] # ordered pause/flush/commit/reset/resume + stage + self.staged: list[VersionRef] = [] + self.committed: list[VersionRef] = [] + + async def stage(self, manifest: VersionManifest, source_dir: str) -> None: + self.staged.append(manifest.ref) + self.calls.append(f"stage:{manifest.ref.version}") + + async def commit(self, ref: VersionRef) -> None: + self.committed.append(ref) + self.calls.append(f"commit:{ref.version}") + + async def flush(self) -> None: + self.calls.append("flush") + + async def pause(self) -> None: + self.calls.append("pause") + + async def resume(self) -> None: + self.calls.append("resume") async def reset(self) -> None: - self.events.append("reset") - - async def pause_generation(self) -> None: - self.events.append("pause") - - async def continue_generation(self) -> None: - self.events.append("continue") - - -class SyncManagerTest(unittest.TestCase): - def test_startup_sync_applies_transition_chain(self) -> None: - async def run() -> None: - with tempfile.TemporaryDirectory() as tmp: - board = FilesystemBulletinBoard(tmp) - board.publish_manifest(VersionManifest(version=1, base_version=0, backend="fake", load_format="noop")) - board.publish_manifest(VersionManifest(version=2, base_version=1, backend="fake", load_format="noop")) - engine = FakeEngine() - manager = WeightSyncManager(board=board, engine=engine) - - await manager.startup_sync() - - self.assertEqual(manager.current_version, 2) - self.assertEqual(manager.sync_state, SyncState.IDLE) - self.assertEqual([v for v, _ in engine.applies], [1, 2]) - - asyncio.run(run()) - - def test_run_change_rematerializes_and_resets_under_gate(self) -> None: - async def run() -> None: - with tempfile.TemporaryDirectory() as tmp: - root = Path(tmp) - board = FilesystemBulletinBoard(root, layout="slime") - _write_slime_version(root / "run-a", 1, 0) - _write_slime_version(root / "run-a", 2, 1) - board.write_latest("run-a", 2) - engine = FakeEngine() - manager = WeightSyncManager(board=board, engine=engine, commit_mode="in_place") - - 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]) - - # A new run forks at base with its version space restarting at 1. - _write_slime_version(root / "run-b", 1, 0) - board.write_latest("run-b", 1) - await manager.sync_to() - - self.assertEqual(manager.current_run_id, "run-b") - self.assertEqual(manager.current_version, 1) - self.assertIn("run-b", engine.applies[-1][1]) # applied the new run's chain - # The re-materialize ran under the commit gate: the reset is - # bracketed by pause/continue, so no request decodes across it. - i = engine.events.index("reset") - self.assertEqual(engine.events[i - 1], "pause") - self.assertEqual(engine.events[i + 1], "continue") - - asyncio.run(run()) - - def test_stitch_layout_never_switches_run(self) -> None: - async def run() -> None: - with tempfile.TemporaryDirectory() as tmp: - board = FilesystemBulletinBoard(tmp) # stitch layout: run-less - board.publish_manifest(VersionManifest(version=1, base_version=0, backend="fake", load_format="noop")) - board.publish_manifest(VersionManifest(version=2, base_version=1, backend="fake", load_format="noop")) - engine = FakeEngine() - manager = WeightSyncManager(board=board, engine=engine) - - await manager.startup_sync() - - self.assertIsNone(manager.current_run_id) - self.assertEqual(manager.current_version, 2) - self.assertNotIn("reset", engine.events) - - asyncio.run(run()) - - def test_exact_and_min_policy_errors_are_retryable(self) -> None: - async def run() -> None: - with tempfile.TemporaryDirectory() as tmp: - board = FilesystemBulletinBoard(tmp) - board.publish_manifest(VersionManifest(version=1, base_version=0, backend="fake", load_format="noop")) - manager = WeightSyncManager(board=board, engine=FakeEngine()) - - ok, current, error = await manager.validate_policy(WeightVersionPolicy(exact_version=1)) - self.assertFalse(ok) - self.assertEqual(current, 0) - self.assertEqual(error["error"]["type"], "WeightVersionNotReady") - - await manager.sync_to(1) - ok, current, error = await manager.validate_policy(WeightVersionPolicy(min_required_version=1)) - self.assertTrue(ok) - self.assertEqual(current, 1) - self.assertIsNone(error) - - ok, _current, error = await manager.validate_policy(WeightVersionPolicy(exact_version=0)) - self.assertFalse(ok) - self.assertEqual(error["error"]["type"], "WeightVersionTooOld") - - asyncio.run(run()) - - def test_request_context_pins_and_reports_serving_version(self) -> None: - async def run() -> None: - with tempfile.TemporaryDirectory() as tmp: - board = FilesystemBulletinBoard(tmp) - board.publish_manifest(VersionManifest(version=1, base_version=0, backend="fake", load_format="noop")) - manager = WeightSyncManager(board=board, engine=FakeEngine()) - await manager.sync_to(1) - - async with manager.request_context(WeightVersionPolicy(exact_version=1)) as version: - self.assertEqual(version, 1) - self.assertEqual(manager.active_requests, 1) - info = await manager.server_info() - self.assertEqual(info["inflight_exact_versions"], {"1": 1}) - self.assertEqual(manager.active_requests, 0) - info = await manager.server_info() - self.assertEqual(info["inflight_exact_versions"], {}) - - with self.assertRaises(PolicyViolation) as cm: - async with manager.request_context(WeightVersionPolicy(exact_version=0)): - pass - self.assertEqual(cm.exception.error["error"]["type"], "WeightVersionTooOld") - - asyncio.run(run()) - - def test_commit_window_gates_admissions(self) -> None: - """A request arriving while the engine apply is in flight must not be - validated against the stale pre-commit version (it would otherwise be - served on the new weights while reporting the old version).""" - - async def run() -> None: - with tempfile.TemporaryDirectory() as tmp: - board = FilesystemBulletinBoard(tmp) - board.publish_manifest(VersionManifest(version=1, base_version=0, backend="fake", load_format="noop")) - engine = FakeEngine() - engine.apply_gate = asyncio.Event() - manager = WeightSyncManager(board=board, engine=engine) - - sync = asyncio.create_task(manager.sync_to(1)) - await engine.apply_started.wait() - # Engine apply for v1 is now in flight; current_version still 0. - self.assertEqual(manager.current_version, 0) - - async def admit_exact_zero() -> str: - try: - async with manager.request_context(WeightVersionPolicy(exact_version=0)): - return "served" - except PolicyViolation as exc: - return exc.error["error"]["type"] - - admit = asyncio.create_task(admit_exact_zero()) - await asyncio.sleep(0.05) - # The admission must be gated, not validated against version 0. - self.assertFalse(admit.done()) - - engine.apply_gate.set() - await sync - # After the commit lands, exact_version=0 is no longer - # satisfiable and must be rejected, not silently served on v1. - self.assertEqual(await admit, "WeightVersionTooOld") - self.assertEqual(manager.current_version, 1) - - async with manager.request_context(WeightVersionPolicy(min_required_version=1)) as version: - self.assertEqual(version, 1) - - asyncio.run(run()) - - def test_in_place_commit_pauses_applies_continues_without_flush(self) -> None: - async def run() -> None: - with tempfile.TemporaryDirectory() as tmp: - board = FilesystemBulletinBoard(tmp) - board.publish_manifest(VersionManifest(version=1, base_version=0, backend="fake", load_format="noop")) - - versions_at_continue: list[int] = [] - - class InPlaceEngine(FakeEngine): - async def continue_generation(self) -> None: - versions_at_continue.append(manager.current_version) - await super().continue_generation() - - engine = InPlaceEngine() - manager = WeightSyncManager(board=board, engine=engine, commit_mode="in_place") - await manager.sync_to(1) - - self.assertEqual(manager.current_version, 1) - self.assertEqual(engine.events, ["pause", "apply", "continue"]) - self.assertEqual(engine.flushes, 0) - # New admissions must see the new namespace before the engine - # resumes, so the bump happens before continue_generation. - self.assertEqual(versions_at_continue, [1]) - - asyncio.run(run()) - - def test_in_place_commit_continues_engine_on_apply_failure(self) -> None: - async def run() -> None: - with tempfile.TemporaryDirectory() as tmp: - board = FilesystemBulletinBoard(tmp) - board.publish_manifest(VersionManifest(version=1, base_version=0, backend="fake", load_format="noop")) - - class FailingEngine(FakeEngine): - async def apply_manifest(self, manifest: VersionManifest, version_path: str) -> None: - raise RuntimeError("apply blew up") - - engine = FailingEngine() - manager = WeightSyncManager(board=board, engine=engine, commit_mode="in_place") - await manager.sync_to(1) - - self.assertEqual(manager.sync_state, SyncState.ERROR) - self.assertEqual(manager.current_version, 0) - # The engine must not be left paused after a failed apply. - self.assertEqual(engine.events, ["pause", "continue"]) - async with manager.request_context() as version: - self.assertEqual(version, 0) - - asyncio.run(run()) - - def test_in_place_commit_gates_exact_but_admits_nonstrict(self) -> None: - async def run() -> None: - with tempfile.TemporaryDirectory() as tmp: - board = FilesystemBulletinBoard(tmp) - board.publish_manifest(VersionManifest(version=1, base_version=0, backend="fake", load_format="noop")) - engine = FakeEngine() - engine.apply_gate = asyncio.Event() - manager = WeightSyncManager(board=board, engine=engine, commit_mode="in_place") - - # An exact-version request in flight blocks the commit point. - 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) - self.assertFalse(engine.apply_started.is_set()) - - # Releasing the exact request lets the commit proceed. - await exact_ctx.__aexit__(None, None, None) - await engine.apply_started.wait() - - # Mid-commit: non-strict requests are admitted (and stamped - # with the pre-commit version); exact requests are gated. - async with manager.request_context() as version: - self.assertEqual(version, 0) - - async def admit_exact() -> str: - try: - async with manager.request_context(WeightVersionPolicy(exact_version=0)): - return "served" - except PolicyViolation as exc: - return exc.error["error"]["type"] - - gated = asyncio.create_task(admit_exact()) - await asyncio.sleep(0.05) - self.assertFalse(gated.done()) - - engine.apply_gate.set() - await sync - self.assertEqual(manager.current_version, 1) - self.assertEqual(await gated, "WeightVersionTooOld") - - asyncio.run(run()) - - def test_commit_gate_clears_on_apply_failure(self) -> None: - async def run() -> None: - with tempfile.TemporaryDirectory() as tmp: - board = FilesystemBulletinBoard(tmp) - board.publish_manifest(VersionManifest(version=1, base_version=0, backend="fake", load_format="noop")) - - class FailingEngine(FakeEngine): - async def apply_manifest(self, manifest: VersionManifest, version_path: str) -> None: - raise RuntimeError("apply blew up") - - manager = WeightSyncManager(board=board, engine=FailingEngine()) - await manager.sync_to(1) - self.assertEqual(manager.sync_state, SyncState.ERROR) - - # The gate must not be left set after a failed commit. - async with manager.request_context() as version: - self.assertEqual(version, 0) - - asyncio.run(run()) - - -class AdmissionGateCommitDriverTest(unittest.TestCase): - """The shared commit_version driver both WeightSyncManager and the hot-load - ProviderShim commit through.""" - - def test_in_place_commit_advances_version_before_resume(self) -> None: - async def run() -> None: - gate = RolloutAdmissionGate(commit_mode="in_place") - events: list[str] = [] - - async def pause() -> None: - events.append("pause") - - async def apply() -> None: - events.append("apply") - - async def resume() -> None: - events.append("resume") - - await gate.commit_version( - apply=apply, - on_applied=lambda: events.append("applied"), - pause=pause, - resume=resume, - ) - - # Pause, apply, advance the served version, then resume — and the - # gate is reopened afterward. - self.assertEqual(events, ["pause", "apply", "applied", "resume"]) - self.assertFalse(gate._committing) - - asyncio.run(run()) - - def test_quiesce_commit_skips_pause_and_unwinds_on_failure(self) -> None: - async def run() -> None: - gate = RolloutAdmissionGate(commit_mode="quiesce") - events: list[str] = [] - - async def pause() -> None: - events.append("pause") - - async def apply() -> None: - events.append("apply") - raise RuntimeError("boom") + self.calls.append("reset") + + async def prefetch(self) -> None: + self.calls.append("prefetch") + + def stamp_request(self, request, served) -> None: + pass + + def stamp_response(self, response, served, current) -> None: + pass + + def base_url(self) -> str: + return "http://engine" + + +def _full(run: str, version: int) -> VersionManifest: + return VersionManifest(VersionRef(run, version), VersionKind.FULL, ["model.safetensors"]) + + +def _delta(run: str, version: int, *, files: list[str]) -> VersionManifest: + return VersionManifest( + VersionRef(run, version), VersionKind.DELTA, files, base_version=version - 1, delta_encoding="xor" + ) + + +def _run(coro) -> None: + asyncio.run(coro) + + +# ── reconcile ──────────────────────────────────────────────────────────────── +def test_fresh_reconcile() -> None: + async def go() -> None: + engine = FakeEngine() + r = Reconciler(store=FakeStore(VersionRef("r1", 3), _full("r1", 3)), engine=engine) + await r.startup() + assert r.applied == VersionRef("r1", 3) + assert engine.staged[-1] == VersionRef("r1", 3) + assert VersionRef("r1", 3) in engine.committed + assert r.sync_state is SyncState.IDLE + assert engine.calls.index("flush") < engine.calls.index("commit:3") # quiesce flushes first + + _run(go()) + + +def test_startup_prefetches_base() -> None: + async def go() -> None: + engine = FakeEngine() + r = Reconciler(store=FakeStore(), engine=engine) # unclaimed pool: reconcile is a no-op + await r.startup() + await r._prefetch_task # let the background base-seed finish + assert "prefetch" in engine.calls + + _run(go()) + + +def test_catch_up() -> None: + async def go() -> None: + engine = FakeEngine() + r = Reconciler(store=FakeStore(VersionRef("r1", 5), _full("r1", 5)), engine=engine) + r.applied = VersionRef("r1", 3) + await r.reconcile() + assert r.applied == VersionRef("r1", 5) + assert engine.committed == [VersionRef("r1", 5)] # one composed stage+reload, not per-version + + _run(go()) + + +def test_run_switch_resets_in_place() -> None: + async def go() -> None: + engine = FakeEngine() + r = Reconciler(store=FakeStore(VersionRef("r2", 2), _full("r2", 2)), engine=engine, commit_mode="in_place") + r.applied = VersionRef("r1", 5) + await r.reconcile() + assert r.applied == VersionRef("r2", 2) + assert "reset" in engine.calls # was patched -> reseed base for the new run + assert engine.calls.index("pause") < engine.calls.index("reset") < engine.calls.index("resume") + assert "flush" not in engine.calls # in_place never flushes + + _run(go()) + + +def test_run_switch_drains_rolling_requests() -> None: + # A base reset is an incompatible transition: even in in_place mode no rolling request + # is admitted during, or decodes across, the wipe (drain_all; stitch#32 / review P0-B). + async def go() -> None: + engine = FakeEngine() + r = Reconciler(store=FakeStore(VersionRef("r2", 1), _full("r2", 1)), engine=engine, commit_mode="in_place") + r.applied = VersionRef("r1", 5) + release = asyncio.Event() + late_served: list[VersionRef | None] = [] + + async def rolling() -> None: + async with r.admit(None): + await release.wait() + + async def late() -> None: + async with r.admit(None) as served: + late_served.append(served) + + req = asyncio.create_task(rolling()) + await asyncio.sleep(0) # admitted before the switch begins + sync = asyncio.create_task(r.reconcile()) + for _ in range(1000): # bounded: without drain_all the switch completes without draining + if r._committing: + break + await asyncio.sleep(0.001) + late_task = asyncio.create_task(late()) + await asyncio.sleep(0.05) + assert "reset" not in engine.calls # the wipe waits for the rolling request + assert not late_served # and nothing is admitted while draining + release.set() + await asyncio.gather(req, sync, late_task) + assert "reset" in engine.calls + assert late_served[0] is not None and late_served[0].run_id == "r2" # admitted post-wipe + + _run(go()) + + +def test_rolling_requests_cross_in_place_commit() -> None: + # The counterpart: a *compatible* in_place commit never waits on rolling traffic — + # the update applies while the request keeps decoding; only a base reset drains. + async def go() -> None: + engine = FakeEngine() + r = Reconciler(store=FakeStore(VersionRef("r1", 4), _full("r1", 4)), engine=engine, commit_mode="in_place") + r.applied = VersionRef("r1", 3) + release = asyncio.Event() + + async def rolling() -> None: + async with r.admit(None): + await release.wait() + + req = asyncio.create_task(rolling()) + await asyncio.sleep(0) + await r.reconcile() # completes while the rolling request is still in flight + assert r.applied == VersionRef("r1", 4) + release.set() + await req + + _run(go()) + + +def test_empty_delta_skips_reload() -> None: + async def go() -> None: + engine = FakeEngine() + r = Reconciler(store=FakeStore(VersionRef("r1", 4), _delta("r1", 4, files=[])), engine=engine) + r.applied = VersionRef("r1", 3) + await r.reconcile() + assert r.applied == VersionRef("r1", 4) + assert engine.staged == [VersionRef("r1", 4)] + assert engine.committed == [] # no reload for a zero-file delta + assert r.metrics.get("skipped_reload") is True + + _run(go()) + + +def test_periodic_reconcile_recovers_missed_wake() -> None: + async def go() -> None: + engine = FakeEngine() + store = FakeStore(VersionRef("r1", 3), _full("r1", 3), _full("r1", 5)) + r = Reconciler(store=store, engine=engine, reconcile_interval=0.02) + await r.startup() + assert r.applied == VersionRef("r1", 3) + # A publish advances latest, but its wake never lands (cold start raced it, or the + # best-effort wake was lost). No wake() call, no request — only the background loop. + store.advance_pointer(VersionRef("r1", 5)) + await asyncio.sleep(0.1) + assert r.applied == VersionRef("r1", 5) # the backstop caught up on its own + await r.shutdown() + + _run(go()) + + +def test_reconcile_interval_zero_disables_backstop() -> None: + async def go() -> None: + store = FakeStore(VersionRef("r1", 3), _full("r1", 3), _full("r1", 5)) + r = Reconciler(store=store, engine=FakeEngine(), reconcile_interval=0.0) + await r.startup() + store.advance_pointer(VersionRef("r1", 5)) + await asyncio.sleep(0.1) + assert r.applied == VersionRef("r1", 3) # no backstop: stays until a wake/409 + assert r._periodic_task is None + await r.shutdown() + + _run(go()) + + +def test_stage_waits_for_prefetch() -> None: + # The base prefetch and a stage both write /local; the stage must wait for the prefetch + # (serialized in the reconciler, not left to the engine's file lock) so they never race. + async def go() -> None: + engine = FakeEngine() + release = asyncio.Event() + base_prefetch = engine.prefetch + + async def slow_prefetch() -> None: + await release.wait() + await base_prefetch() + + engine.prefetch = slow_prefetch # type: ignore[method-assign] + r = Reconciler(store=FakeStore(VersionRef("r1", 2), _full("r1", 2)), engine=engine, reconcile_interval=0.0) + r.applied = VersionRef("r1", 0) # same run, behind -> stage v2 (no run switch) + task = asyncio.create_task(r.startup()) + await asyncio.sleep(0.05) # startup fires the prefetch; reconcile reaches the await + assert "stage:2" not in engine.calls # blocked on the (still-running) prefetch + release.set() + await task + assert engine.calls.index("prefetch") < engine.calls.index("stage:2") # prefetch, then stage + await r.shutdown() + + _run(go()) + + +# ── convergence liveness ───────────────────────────────────────────────────── +# The backstop self-heals what wake-driven-only cannot (the red tests behind +# stitch#45): each scenario converges with the backstop and never with interval=0. + + +async def _converged(r: Reconciler, target: VersionRef, timeout: float = 1.0) -> bool: + deadline = asyncio.get_running_loop().time() + timeout + while asyncio.get_running_loop().time() < deadline: + if r.applied == target and r.sync_state is SyncState.IDLE: + return True + await asyncio.sleep(0.02) + return False + + +class FlakyStore(FakeStore): + """read_manifest fails once, then heals — a transient store-side error.""" + + failures = 1 + + def read_manifest(self, ref: VersionRef) -> VersionManifest: + if self.failures: + self.failures -= 1 + raise RuntimeError("transient store error") + return super().read_manifest(ref) + + +async def _heals_transient_error(interval: float) -> bool: + r = Reconciler(store=FlakyStore(VersionRef("r1", 1), _full("r1", 1)), engine=FakeEngine(), + reconcile_interval=interval) + await r.startup() # the pass hits the error -> ERROR; the store is healed from here on + async with r.admit(None): + pass # unconstrained traffic never 409s, so it nudges nothing + ok = await _converged(r, VersionRef("r1", 1)) + await r.shutdown() + return ok + + +def test_transient_error_recovery_needs_backstop() -> None: + assert asyncio.run(_heals_transient_error(0.05)) + assert not asyncio.run(_heals_transient_error(0)) # ERROR retries only on external wake + + +class HostViewStore(FakeStore): + """advance_pointer lands on the durable *remote*; read_pointer sees it only after + refresh() snapshots remote -> local (Volume reload semantics). refresh_gate lets a + test hold a pass open on a pre-publish snapshot.""" + + refresh_gate: queue.Queue[threading.Event] | None = None + + def __init__(self, pointer: VersionRef | None = None, *manifests: VersionManifest) -> None: + super().__init__(None, *manifests) + self.remote_pointer = pointer + + def refresh(self) -> None: + self._pointer = self.remote_pointer + if self.refresh_gate is not None: + self.refresh_gate.put(release := threading.Event()) + release.wait(timeout=10) + + def advance_pointer(self, ref: VersionRef) -> None: + self.remote_pointer = ref + + +async def _heals_dropped_wake(interval: float) -> bool: + """A wake IS delivered, but mid-pass: wake() no-ops against the running task, whose + caught-up recheck already snapshotted pre-publish state — the wake is lost.""" + store = HostViewStore(VersionRef("r1", 1), _full("r1", 1), _full("r1", 2)) + r = Reconciler(store=store, engine=FakeEngine(), reconcile_interval=interval) + await r.startup() # converges to v1, ungated + + gate = store.refresh_gate = queue.Queue() + r.wake() # start an idle pass + (await asyncio.to_thread(gate.get, True, 10)).set() # release its pass-start refresh + recheck = await asyncio.to_thread(gate.get, True, 10) # its recheck: snapshotted v1, held + store.advance_pointer(VersionRef("r1", 2)) + r.wake() # v2's wake: delivered mid-pass -> dropped; the pass idles on its v1 snapshot + recheck.set() + store.refresh_gate = None + ok = await _converged(r, VersionRef("r1", 2)) + await r.shutdown() + return ok + + +def test_dropped_wake_recovery_needs_backstop() -> None: + assert asyncio.run(_heals_dropped_wake(0.05)) + assert not asyncio.run(_heals_dropped_wake(0)) + + +def test_constrained_409_recovers_without_backstop() -> None: + """The event-driven channel: a min_version 409 self-wakes a stale ERROR replica.""" + + async def go() -> None: + r = Reconciler(store=FlakyStore(VersionRef("r1", 1), _full("r1", 1)), engine=FakeEngine(), + reconcile_interval=0) + await r.startup() + assert r.sync_state is SyncState.ERROR + try: + async with r.admit(VersionConstraint(min_version=1)): + raise AssertionError("should have rejected") + except ConstraintUnmet: + pass + assert await _converged(r, VersionRef("r1", 1)) + + _run(go()) + + +# ── admission gate ─────────────────────────────────────────────────────────── +def test_admit_satisfied() -> None: + async def go() -> None: + r = Reconciler(store=FakeStore(), engine=FakeEngine()) + r.applied = VersionRef("r1", 5) + async with r.admit(VersionConstraint(min_version=3)) as served: + assert served == VersionRef("r1", 5) + + _run(go()) + + +def test_admit_rejected_triggers_wake() -> None: + async def go() -> None: + r = Reconciler(store=FakeStore(VersionRef("r1", 5), _full("r1", 5)), engine=FakeEngine()) + r.applied = VersionRef("r1", 2) + try: + async with r.admit(VersionConstraint(min_version=5)): + raise AssertionError("should have rejected") + except ConstraintUnmet as e: + assert e.error["type"] == "WeightVersionNotReady" + assert e.error["target_version"] == 5 + assert r._task is not None # the 409 kicked off a catch-up reconcile - async def resume() -> None: - events.append("resume") + _run(go()) - with self.assertRaises(RuntimeError): - await gate.commit_version( - apply=apply, - on_applied=lambda: events.append("applied"), - pause=pause, - resume=resume, - ) - # quiesce never pauses; a failed apply advances nothing and the gate - # is cleared so admissions resume. - self.assertEqual(events, ["apply"]) - self.assertFalse(gate._committing) +def test_version_flips_before_resume() -> None: + async def go() -> None: + engine = FakeEngine() + r = Reconciler(store=FakeStore(VersionRef("r1", 4), _full("r1", 4)), engine=engine, commit_mode="in_place") + r.applied = VersionRef("r1", 3) + seen: dict[str, VersionRef | None] = {} + base_resume = engine.resume - asyncio.run(run()) + async def resume_spy() -> None: + seen["applied"] = r.applied + await base_resume() + + engine.resume = resume_spy # type: ignore[method-assign] + await r.reconcile() + assert seen["applied"] == VersionRef("r1", 4) # flipped under the gate, before resume + + _run(go()) if __name__ == "__main__": - unittest.main() + tests = [v for k, v in sorted(globals().items()) if k.startswith("test_")] + for t in tests: + t() + print(f" ok {t.__name__}") + print(f"reconcile harness: {len(tests)} PASS") diff --git a/src/stitch/trainers/__init__.py b/src/stitch/trainers/__init__.py deleted file mode 100644 index e274eaa..0000000 --- a/src/stitch/trainers/__init__.py +++ /dev/null @@ -1 +0,0 @@ -"""Trainer framework adapters.""" diff --git a/src/stitch/trainers/slime.py b/src/stitch/trainers/slime.py deleted file mode 100644 index b19c61e..0000000 --- a/src/stitch/trainers/slime.py +++ /dev/null @@ -1,165 +0,0 @@ -"""Slime trainer adapter for disaggregated rollout weight sync.""" - -from __future__ import annotations - -import logging -import time -from argparse import Namespace -from pathlib import Path -from typing import Any - -from stitch.bulletin import FilesystemBulletinBoard -from stitch.protocol import Artifact, VersionManifest - - -logger = logging.getLogger(__name__) - - -def publish_delta_version( - args: Any, - version_dir: str, - files: list[str], - weight_version: str | int, - rollout_engines: list[Any], -) -> list[Any]: - """Slime ``custom_delta_publish_path`` hook: write the version manifest and - advance ``latest.json`` on the bulletin board. - - Provider-agnostic — no provider import. Durability (e.g. committing a Modal - Volume) and best-effort rollout-pool wake are layered on by the consuming - example's hook (see ``cookbook/slime_disagg/hooks.py``). - """ - del rollout_engines - version = int(weight_version) - root = Path(_bulletin_root(args)) - version_path = Path(version_dir) - - # Disk-delta slime writes a canonical model.safetensors.index.json carrying - # the delta encoding/compression/checksum; lift it instead of hardcoding the - # format. Fall back to the pre-index layout when the engine didn't write one. - if (version_path / "model.safetensors.index.json").exists(): - manifest = VersionManifest.from_slime_index( - version_path, - run_id=getattr(args, "run_id", None), - base_model=getattr(args, "hf_checkpoint", None), - ) - else: - sorted_files = sorted(files) - manifest = VersionManifest( - version=version, - base_version=version - 1, - backend="sparse_delta", - load_format="delta", - transition_files=sorted_files, - artifacts=[Artifact(kind="transition", path=path) for path in sorted_files], - created_at=time.time(), - run_id=getattr(args, "run_id", None), - base_model=getattr(args, "hf_checkpoint", None), - metadata={"trainer": "slime", "transport": "disk"}, - ) - FilesystemBulletinBoard(root).publish_manifest(manifest, version_path=version_path) - logger.info("Published delta version %s with %d file(s)", manifest.version, len(manifest.transition_files)) - return [] - - -def rollout_request_weight_version_hook(args: Namespace, sample: Any, request: dict[str, Any]) -> None: - """Attach provider admission constraints to one SLIME rollout request. - - The hook is request-level control, not the trainer's staleness policy: it - prevents an opaque rollout router from spending compute on a replica that - cannot serve a version the trainer has already decided is usable. - """ - - mode = str(getattr(args, "rollout_request_weight_version_mode", "exact")) - # PR #5's per-request hook receives no rollout_id (it is per-rollout context, - # not per-request — sample.index is the batch position), so the trainer-step - # pin only runs when one is supplied. Otherwise the pool serves the latest - # hot-loaded version and a lagging replica returns a retryable 409. TODO: - # re-derive the per-request target (e.g. the latest published version) under - # the PR #5 contract if strict pinning is needed. - rollout_id = request.get("rollout_id") - if mode != "none" and rollout_id is not None: - target_version = rollout_target_weight_version( - args, - int(rollout_id), - evaluation=bool(request.get("evaluation", False)), - ) - if not bool(request.get("evaluation", False)): - target_version = max(0, target_version - int(getattr(args, "rollout_request_weight_version_lag", 0))) - if mode == "exact": - request["payload"]["weight_version"] = {"exact_version": target_version} - elif mode == "min": - request["payload"]["weight_version"] = {"min_required_version": target_version} - else: - raise ValueError(f"Unsupported rollout_request_weight_version_mode: {mode!r}") - - # Generous retries on every request so a lagging/scaling replica (a 409 - # weight-version reject or a transient error) is retried, not failed. - 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))) - if getattr(sample, "session_id", None): - # Provider-neutral by default. Modal's Flash gateway routes session - # affinity on the Modal-Session-ID header, so the Modal configs set this - # to that name and affinity is honored at the gateway (one hop) rather - # than re-routed inside the rollout container. - affinity_header = str( - getattr(args, "rollout_session_affinity_header", "x-session-affinity") - ) - headers = dict(request.get("headers") or {}) - headers.setdefault(affinity_header, sample.session_id) - request["headers"] = headers - - -def generate_rollout( - args: Namespace, - rollout_id: int, - data_source: Any, - evaluation: bool = False, -): - """Run SLIME's default SGLang rollout. - - Kept as a compatibility wrapper for older configs. New configs should use - ``slime.rollout.sglang_rollout.generate_rollout`` directly plus - ``custom_rollout_request_hook_path`` when they need request constraints. - """ - from slime.rollout import sglang_rollout as upstream_rollout - - assert args.rollout_global_dataset - logger.info( - "Disaggregated %s rollout_id=%s", - "eval" if evaluation else "train", - rollout_id, - ) - - with upstream_rollout.rollout_request_context(args, rollout_id, evaluation=evaluation): - if evaluation: - output, _ = upstream_rollout.run(upstream_rollout.eval_rollout(args, rollout_id)) - return output - - output, aborted_samples = upstream_rollout.run( - upstream_rollout.generate_rollout_async(args, rollout_id, data_source.get_samples) - ) - if aborted_samples: - data_source.add_samples(aborted_samples) - return output - - -def rollout_target_weight_version(args: Namespace, rollout_id: int, evaluation: bool = False) -> int: - if not evaluation: - return int(rollout_id) - # Eval pins to the latest published version (the slime-native `latest`). - try: - _, version = FilesystemBulletinBoard(_bulletin_root(args), layout="slime").read_latest() - return version - except Exception: # noqa: BLE001 - return int(rollout_id) - - -def _bulletin_root(args: Any) -> str: - root = ( - getattr(args, "update_weight_disk_dir", None) - or getattr(args, "update_weight_delta_root", None) - ) - if root: - return str(root) - return str(Path(args.update_weight_delta_dir).parent) diff --git a/src/stitch/trainers/slime_test.py b/src/stitch/trainers/slime_test.py deleted file mode 100644 index 2592fdb..0000000 --- a/src/stitch/trainers/slime_test.py +++ /dev/null @@ -1,155 +0,0 @@ -from __future__ import annotations - -import tempfile -import unittest -from argparse import Namespace -from pathlib import Path - -from stitch.protocol import VersionManifest, read_latest -from stitch.trainers.slime import publish_delta_version, rollout_request_weight_version_hook - - -class SlimeHooksTest(unittest.TestCase): - def test_rollout_request_hook_adds_exact_policy_retry_and_affinity(self) -> None: - args = Namespace( - rollout_request_weight_version_mode="exact", - rollout_request_weight_version_lag=1, - rollout_request_retry_attempts=240, - rollout_request_retry_sleep=0.25, - ) - sample = Namespace(session_id="session-1") - request = { - "payload": {}, - "headers": None, - "max_retries": 60, - "retry_sleep": 1.0, - "rollout_id": 3, - "evaluation": False, - } - - rollout_request_weight_version_hook(args, sample, request) - - self.assertEqual(request["payload"]["weight_version"], {"exact_version": 2}) - self.assertEqual(request["max_retries"], 240) - self.assertEqual(request["retry_sleep"], 0.25) - self.assertEqual(request["headers"]["x-session-affinity"], "session-1") - - def test_rollout_request_hook_uses_configured_affinity_header(self) -> None: - args = Namespace( - rollout_request_weight_version_mode="exact", - rollout_session_affinity_header="Modal-Session-ID", - ) - sample = Namespace(session_id="group-7") - request = { - "payload": {}, - "headers": None, - "max_retries": 60, - "retry_sleep": 1.0, - "rollout_id": 3, - "evaluation": False, - } - - rollout_request_weight_version_hook(args, sample, request) - - self.assertEqual(request["headers"]["Modal-Session-ID"], "group-7") - self.assertNotIn("x-session-affinity", request["headers"]) - - def test_rollout_request_hook_can_add_min_policy(self) -> None: - args = Namespace(rollout_request_weight_version_mode="min") - sample = Namespace(session_id=None) - request = { - "payload": {}, - "headers": None, - "max_retries": 60, - "retry_sleep": 1.0, - "rollout_id": 3, - "evaluation": False, - } - - rollout_request_weight_version_hook(args, sample, request) - - self.assertEqual(request["payload"]["weight_version"], {"min_required_version": 3}) - self.assertIsNone(request["headers"]) - - def test_rollout_request_hook_degrades_without_rollout_id(self) -> None: - # PR #5's request carries no rollout_id: skip the pin (no crash), still - # apply the retry budget and session affinity. - args = Namespace( - rollout_request_weight_version_mode="exact", - rollout_request_retry_attempts=240, - ) - sample = Namespace(session_id="grp-9") - request = {"payload": {}, "headers": None, "max_retries": 60, "retry_sleep": 1.0} - - rollout_request_weight_version_hook(args, sample, request) - - self.assertNotIn("weight_version", request["payload"]) - self.assertEqual(request["max_retries"], 240) - self.assertEqual(request["headers"]["x-session-affinity"], "grp-9") - - def test_publish_delta_version_writes_manifest_and_latest(self) -> None: - with tempfile.TemporaryDirectory() as tmp: - root = Path(tmp) - version_dir = root / "versions" / "weight_v000001" - version_dir.mkdir(parents=True) - args = Namespace( - update_weight_delta_root=str(root), - update_weight_delta_dir=str(root / "versions"), - hf_checkpoint="Qwen/Qwen3-4B", - run_id="run-1", - ) - - refs = publish_delta_version(args, str(version_dir), ["b.safetensors", "a.safetensors"], 1, []) - - manifest = VersionManifest.read(version_dir / "manifest.json") - self.assertEqual(refs, []) - self.assertEqual(read_latest(root), 1) - self.assertEqual(manifest.version, 1) - self.assertEqual(manifest.base_version, 0) - self.assertEqual(manifest.transition_files, ["a.safetensors", "b.safetensors"]) - self.assertEqual(manifest.artifacts[0].kind, "transition") - self.assertEqual(manifest.metadata["trainer"], "slime") - - def test_publish_delta_version_lifts_slime_index(self) -> None: - import json - - with tempfile.TemporaryDirectory() as tmp: - root = Path(tmp) - version_dir = root / "versions" / "weight_v000002" - version_dir.mkdir(parents=True) - (version_dir / "model.safetensors.index.json").write_text( - json.dumps( - { - "metadata": { - "version": "000002", - "base_version": "000001", - "delta_encoding": "xor", - "compression_format": "zstd", - "checksum_format": "xxh3-128", - }, - "weight_map": {"w": "model-00001-of-00001.safetensors"}, - } - ), - encoding="utf-8", - ) - args = Namespace( - update_weight_delta_root=str(root), - update_weight_delta_dir=str(root / "versions"), - hf_checkpoint="Qwen/Qwen3-4B", - run_id="run-2", - ) - - # files arg is ignored when the engine wrote a canonical index. - publish_delta_version(args, str(version_dir), [], 2, []) - - manifest = VersionManifest.read(version_dir / "manifest.json") - self.assertEqual(read_latest(root), 2) - self.assertEqual(manifest.base_version, 1) - self.assertEqual(manifest.backend, "disk_delta") - self.assertEqual(manifest.compression_format, "zstd") - self.assertEqual(manifest.checksum_format, "xxh3-128") - self.assertEqual(manifest.transition_files, ["model-00001-of-00001.safetensors"]) - - -if __name__ == "__main__": - unittest.main() diff --git a/src/stitch/versions.py b/src/stitch/versions.py new file mode 100644 index 0000000..6f3ab2d --- /dev/null +++ b/src/stitch/versions.py @@ -0,0 +1,214 @@ +"""stitch domain vocabulary: version types, the request constraint, replica/pool +readiness, and the pure pointer rules. + +No I/O, no Modal, no engine, no framework. Every other module depends on this, +and this depends on nothing else in stitch. The Store/Engine/Pool ports live +with their instances (``stores/base.py``, ``engines/base.py``, ``pools/base.py``). +""" + +from __future__ import annotations + +import json +import time +from dataclasses import dataclass +from enum import Enum +from pathlib import Path +from typing import Any + +_WEIGHT_PREFIX = "weight_v" + + +@dataclass(frozen=True) +class VersionRef: + """A published policy version: a run epoch (``run_id``) and a monotonic number. + + Its string form — ``/weight_v000007``, or bare ``weight_v000007`` + run-less — is self-identifying, so one run's chain can't be mistaken for + another's; external specs call this the checkpoint identity. + """ + + run_id: str | None + version: int + + @property + def identity(self) -> str: + name = f"{_WEIGHT_PREFIX}{self.version:06d}" + return f"{self.run_id}/{name}" if self.run_id else name + + @classmethod + def parse(cls, text: str) -> "VersionRef": + # ``[/]weight_vNNNNNN`` (or a legacy bare ``NNNNNN``). A missing pointer + # is handled by callers (they pass "" -> None before this); non-empty content + # that doesn't parse is corrupt control state -> raise, never a silent fall back + # to base (which would serve the wrong weights). Ref: stitch#31. + text = (text or "").strip() + run_id, _, tail = text.rpartition("/") + digits = tail[len(_WEIGHT_PREFIX):] if tail.startswith(_WEIGHT_PREFIX) else tail + if not digits.isdigit(): + raise ValueError(f"unparseable snapshot pointer: {text!r}") + return cls(run_id or None, int(digits)) + + +class VersionKind(str, Enum): + FULL = "full" # an anchor: the version's files are the weights + DELTA = "delta" # a diff against base_version (same run), chaining back to an anchor + + +@dataclass(frozen=True) +class VersionManifest: + """One published version, derived from its directory's HF index (never stored + separately). ``kind`` alone decides how a replica applies it: FULL seeds from + ``files``; DELTA decodes (``delta_encoding`` / ``compression`` / ``checksum``) against + ``base_version``, which chains back to the nearest FULL anchor. + """ + + ref: VersionRef + kind: VersionKind + files: list[str] + base_version: int | None = None # required iff DELTA; always the same run as ref + delta_encoding: str | None = None + compression: str | None = None + checksum: str | None = None + base_model: str | None = None + created_at: float = 0.0 + + @classmethod + def from_hf_index( + cls, + version_dir: str | Path, + *, + run_id: str | None = None, + base_model: str | None = None, + ) -> "VersionManifest": + # The dir's model.safetensors.index.json carries the lineage + codec under + # `metadata` and the files as `weight_map` values. A `diff` key => DELTA. + index = json.loads((Path(version_dir) / "model.safetensors.index.json").read_text()) + meta = index.get("metadata") or {} + weight_map = index.get("weight_map") or {} + diff = meta.get("diff") or None + base = meta.get("base_version") + return cls( + ref=VersionRef(run_id, int(meta["version"])), + kind=VersionKind.DELTA if diff else VersionKind.FULL, + files=sorted({str(f) for f in weight_map.values()}), + base_version=int(base) if diff and base is not None else None, + delta_encoding=diff, + compression=meta.get("compression") or None, + checksum=meta.get("checksum") or None, + base_model=base_model, + created_at=float(meta.get("created_at", time.time())), + ) + + +@dataclass(frozen=True) +class VersionConstraint: + """What a rollout request requires of the serving version. ``exact_version`` + pins one version; ``min_version`` is a floor (a bounded-lag request sets it to + ``latest - lag``). Both None means no constraint. + """ + + min_version: int | None = None + exact_version: int | None = None + + @classmethod + def from_payload(cls, payload: dict[str, Any] | None) -> "VersionConstraint": + raw = (payload or {}).get("weight_version") + raw = raw if isinstance(raw, dict) else {} + mn, ex = raw.get("min_version"), raw.get("exact_version") + return cls(int(mn) if mn is not None else None, int(ex) if ex is not None else None) + + def to_payload(self) -> dict[str, int | None]: + return {"min_version": self.min_version, "exact_version": self.exact_version} + + def satisfied_by(self, applied: int | None) -> bool: + if applied is None: + return self.min_version is None and self.exact_version is None + if self.exact_version is not None: + return applied == self.exact_version + return self.min_version is None or applied >= self.min_version + + +class SyncState(str, Enum): + IDLE = "IDLE" + QUEUED = "QUEUED" + PREFETCHING = "PREFETCHING" + PREPARING = "PREPARING" + COMMITTING = "COMMITTING" + ERROR = "ERROR" + + +_SYNC_STATES = {s.value for s in SyncState} + + +@dataclass(frozen=True) +class ReplicaState: + """One replica's readiness report — the body of its ``server_info``.""" + + ready: bool + applied: VersionRef | None = None + sync_state: SyncState | None = None + reason: str | None = None # why not ready, surfaced to the readiness poll + + @classmethod + def from_dict(cls, data: dict[str, Any]) -> "ReplicaState": + applied, state = data.get("applied"), data.get("sync_state") + return cls( + ready=bool(data.get("ready", False)), + applied=VersionRef.parse(applied) if applied else None, + sync_state=SyncState(state) if state in _SYNC_STATES else None, + reason=data.get("reason"), + ) + + def to_dict(self) -> dict[str, Any]: + data: dict[str, Any] = {"ready": self.ready} + if self.applied is not None: + data["applied"] = self.applied.identity + if self.sync_state is not None: + data["sync_state"] = self.sync_state.value + if self.reason is not None: + data["reason"] = self.reason + return data + + +@dataclass(frozen=True) +class PoolState: + """Aggregate readiness across a pool's replicas; drives the readiness poll.""" + + replicas: list[ReplicaState] + + def ready_fraction(self, target: VersionRef) -> float: + if not self.replicas: + return 0.0 + return sum(r.ready and r.applied == target for r in self.replicas) / len(self.replicas) + + def is_ready(self, target: VersionRef, *, threshold: float = 1.0) -> bool: + return bool(self.replicas) and self.ready_fraction(target) >= threshold + + +class PointerRewind(Exception): + """A move that would rewind ``latest`` within the same run. The single writer + advances monotonically per run; crossing to a *different* run forks at base + and is not a rewind.""" + + def __init__(self, current: VersionRef, proposed: VersionRef) -> None: + super().__init__( + f"latest is {current.identity!r}; refusing to rewind to {proposed.identity!r}" + ) + self.current, self.proposed = current, proposed + + +@dataclass(frozen=True) +class PointerMove: + ref: VersionRef + reset: bool # crossed to a new run -> re-materialize base and restart the version space + + +def decide_pointer_move(current: VersionRef | None, proposed: VersionRef) -> PointerMove: + """A different run forks at base (a reset, even at a lower number); within the + same run the move must be strictly newer, else :class:`PointerRewind`.""" + current_run = current.run_id if current is not None else None + if proposed.run_id != current_run: + return PointerMove(proposed, reset=True) + if current is not None and proposed.version <= current.version: + raise PointerRewind(current, proposed) + return PointerMove(proposed, reset=False) diff --git a/tools/probes/README.md b/tools/probes/README.md new file mode 100644 index 0000000..3d83a96 --- /dev/null +++ b/tools/probes/README.md @@ -0,0 +1,71 @@ +# probes — pressure-test harness (dev-only, never pedagogical) + +Disposable instrumentation for certifying the pool + weight-sync protocol under +realistic and adversarial conditions. Unlike `cookbook/` (best-practice recipes), +nothing here demonstrates how stitch *should* be used — probes exist to find where +it breaks and to measure perf under pressure. Expect rough edges; harden a probe +only when it graduates into the standing regression harness. + +## Design + +A pool certification doesn't need a trainer. The serving side needs **publishes** +and **traffic**, so the harness decouples them: + +- **`replay_publisher.py`** — replays a *recorded* delta chain (past runs persist on + the delta Volume) through the real `publish_version()` path at a controlled + cadence, under a fresh `run_id`. Real delta densities, sizes, and chain structure; + zero trainer GPUs. The target pool must serve the same base model. +- **`traffic.py`** — reward-free load shapes against the pool gateway: + `long_decode`, `long_prefill`, `agentic` (multi-turn sessions with growing + context, synthetic tool-result injections, session affinity), `mixed`. Responses + carry `weight_version_start/end`, so the generator doubles as the + straddle-attribution collector. +- **`poller.py`** — scrapes every replica's `/server_info` into JSONL and + summarizes: applied-version timelines, per-version convergence lag, + stage/commit timings, not-ready windows. +- **`app.py`** — Modal wrapper to run the above from containers (the replay + publisher needs the delta Volume mounted). + +## Conventions + +- **Environment:** everything probe-related lives in the `stitch-dev` Modal + environment (`modal environment create stitch-dev`, once). The target pool must + be deployed in the *same* environment — `ModalFlashPool` resolves names in the + caller's environment — so deploy the cookbook recipe with `-e stitch-dev` for + probe runs. +- **Results:** JSONL baselines land on the `stitch-probe-results` Volume, one + directory per tagged run. Baselines are *recorded and human-judged*, not CI + gates — do not wire thresholds into CI. + +## Quickstart + +```bash +# once +uv run --extra modal modal environment create stitch-dev + +# deploy the target pool (example: glm45_air_fp8) into stitch-dev +EXPERIMENT_CONFIG=glm45_air_fp8 uv run --extra modal modal deploy -m cookbook.miles_disagg.app -e stitch-dev + +# deploy the probe app, bound to that recipe's delta volume +PROBE_DELTA_VOLUME=stitch-delta-glm45-air-fp8 \ + uv run --extra modal modal deploy -m tools.probes.app -e stitch-dev + +# start the poller (background), then traffic, then the replay +uv run --extra modal modal run -e stitch-dev -m tools.probes.app::poll --pool-app stitch-glm45-air-fp8 --tag demo & +uv run --extra modal modal run -e stitch-dev -m tools.probes.app::traffic \ + --pool-app stitch-glm45-air-fp8 --shape agentic --concurrency 32 --duration 1800 --tag demo & +uv run --extra modal modal run -e stitch-dev -m tools.probes.app::replay \ + --pool-app stitch-glm45-air-fp8 --source-run --cadence-s 60 --tag demo +``` + +## Known limitations (skeleton) + +- **No TTFT.** The versioned proxy buffers responses (no streaming), so only + end-to-end latency is measurable. Revisit when the streaming-proxy decision lands. +- **Synthetic filler text.** Prompt content is pseudo-text at controlled lengths; + swap in real document corpora when content realism starts to matter. +- **Token counts are approximated** from word counts (~0.75 words/token). +- The replay publisher **copies** each version dir under the new run prefix + (that's the real publish path) — delta volumes grow per replay and nothing GCs. +- Version floor tracking polls the gateway's `/server_info`, which answers from an + arbitrary replica — probe-grade, not exact. diff --git a/tools/probes/__init__.py b/tools/probes/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/tools/probes/app.py b/tools/probes/app.py new file mode 100644 index 0000000..4baa482 --- /dev/null +++ b/tools/probes/app.py @@ -0,0 +1,84 @@ +"""Modal wrapper for the probes — everything runs in the ``stitch-dev`` environment. + +The replay publisher needs the delta Volume mounted, so which volume to mount is fixed +at deploy time via ``PROBE_DELTA_VOLUME`` (one probe deploy per target recipe): + + PROBE_DELTA_VOLUME=stitch-delta-glm45-air-fp8 \\ + uv run --extra modal modal deploy -m tools.probes.app -e stitch-dev + +The target pool must be deployed in the same environment (ModalFlashPool resolves names +in the caller's environment). Results land on the ``stitch-probe-results`` Volume under +``//``; baselines are recorded and human-judged, never CI gates. +""" + +from __future__ import annotations + +import asyncio +import json +import os + +import modal + +DELTA_ROOT = "/delta-bulletin" +RESULTS_ROOT = "/probe-results" +MINUTES = 60 + +app = modal.App("stitch-probes") +delta_volume_name = os.environ.get("PROBE_DELTA_VOLUME", "stitch-probe-scratch") +delta_volume = modal.Volume.from_name(delta_volume_name, version=2, create_if_missing=True) +results_volume = modal.Volume.from_name("stitch-probe-results", version=2, create_if_missing=True) + +image = ( + modal.Image.debian_slim(python_version="3.12") + .pip_install("httpx") + # Bake the deploy-time volume choice: containers re-import this module, where the + # shell's PROBE_DELTA_VOLUME doesn't exist — without this the store would resolve + # (and commit) a different volume than the one mounted. + .env({"PROBE_DELTA_VOLUME": delta_volume_name}) + .add_local_python_source("stitch", "tools") +) + + +@app.function(image=image, volumes={RESULTS_ROOT: results_volume}, timeout=120 * MINUTES) +def poll(pool_app: str, pool_cls: str = "Server", interval: float = 2.0, duration: float = 3600.0, tag: str = "run") -> None: + from tools.probes import poller + + out = f"{RESULTS_ROOT}/{tag}/server_info.jsonl" + asyncio.run(poller.poll(pool_app, pool_cls, interval=interval, duration=duration, out_path=out)) + results_volume.commit() + print(json.dumps(poller.summarize(out), indent=2)) + + +@app.function(image=image, volumes={RESULTS_ROOT: results_volume}, timeout=120 * MINUTES) +def traffic( + pool_app: str, + pool_cls: str = "Server", + model: str = "default", + shape: str = "mixed", + concurrency: int = 16, + duration: float = 600.0, + lag: int | None = None, + tag: str = "run", +) -> None: + from stitch.pools.modal_flash import ModalFlashPool + from tools.probes import traffic as traffic_mod + + gateway = ModalFlashPool(pool_app, pool_cls).gateway_url() + out = f"{RESULTS_ROOT}/{tag}/traffic-{shape}.jsonl" + summary = asyncio.run(traffic_mod.run( + gateway, model, shape=shape, concurrency=concurrency, duration=duration, lag=lag, out_path=out, + )) + results_volume.commit() + print(json.dumps(summary, indent=2)) + + +@app.function(image=image, volumes={DELTA_ROOT: delta_volume, RESULTS_ROOT: results_volume}, timeout=240 * MINUTES) +def replay(pool_app: str, source_run: str, pool_cls: str = "Server", cadence_s: float = 30.0, limit: int | None = None, tag: str = "run") -> None: + from tools.probes.replay_publisher import replay as replay_chain + + delta_volume.reload() + run_id = replay_chain( + root=DELTA_ROOT, source_run=source_run, app_name=pool_app, cls_name=pool_cls, + volume_name=delta_volume_name, cadence_s=cadence_s, limit=limit, + ) + print(f"replay complete: run_id={run_id} tag={tag}") diff --git a/tools/probes/poller.py b/tools/probes/poller.py new file mode 100644 index 0000000..4b3d789 --- /dev/null +++ b/tools/probes/poller.py @@ -0,0 +1,133 @@ +"""Scrape every replica's ``/server_info`` into JSONL, and summarize. + +The certification runs' observability spine: ``/server_info`` is point-in-time, so this +turns it into a timeseries — applied-version timelines per replica, per-version +convergence lag (first replica at vN -> last replica at vN), stage/commit timings, and +not-ready windows. +""" + +from __future__ import annotations + +import asyncio +import json +import time +from collections import defaultdict +from pathlib import Path +from typing import Any + +REDISCOVER_EVERY = 30.0 # seconds between replica-list refreshes (containers come and go) + + +async def poll( + app_name: str, + cls_name: str = "Server", + *, + interval: float = 2.0, + duration: float | None = None, + out_path: str, + replicas: list[str] | None = None, # static URL list: skip Modal discovery +) -> None: + import httpx + + fixed = list(replicas) if replicas else None + pool = None + if fixed is None: + from stitch.pools.modal_flash import ModalFlashPool + + pool = ModalFlashPool(app_name, cls_name) + out = Path(out_path) + out.parent.mkdir(parents=True, exist_ok=True) + replicas = fixed or [] + last_discover = 0.0 + deadline = time.time() + duration if duration else None + async with httpx.AsyncClient(timeout=10.0, trust_env=False) as client: + with out.open("a") as f: + while deadline is None or time.time() < deadline: + if pool is not None and (not replicas or time.time() - last_discover >= REDISCOVER_EVERY): + replicas = await asyncio.to_thread(pool.discover_replicas) + last_discover = time.time() + for row in await asyncio.gather(*(_probe(client, url) for url in replicas)): + f.write(json.dumps(row) + "\n") + f.flush() + await asyncio.sleep(interval) + + +async def _probe(client: Any, url: str) -> dict[str, Any]: + row: dict[str, Any] = {"t": time.time(), "replica": url} + try: + row["info"] = (await client.get(f"{url.rstrip('/')}/server_info")).json() + except Exception as exc: # noqa: BLE001 — an unreachable replica is itself a data point + row["error"] = str(exc)[:200] + return row + + +def summarize(path: str) -> dict[str, Any]: + """Reduce a poll JSONL to the certification quantities.""" + from stitch.versions import VersionRef + + first_seen: dict[str, dict[int, float]] = defaultdict(dict) # replica -> version -> t + timings: dict[tuple[str, float], dict[str, Any]] = {} # (replica, metrics.at) -> metrics + unready: dict[str, float] = defaultdict(float) # replica -> seconds observed not ready + prev_t: dict[str, float] = {} + errors = 0 + + for line in Path(path).read_text().splitlines(): + row = json.loads(line) + replica, t = row["replica"], row["t"] + info = row.get("info") + if info is None: + errors += 1 + continue + if info.get("applied"): + v = VersionRef.parse(info["applied"]).version + first_seen[replica].setdefault(v, t) + if not info.get("ready") and replica in prev_t: + unready[replica] += t - prev_t[replica] + if (m := info.get("metrics")) and "at" in m: + timings[(replica, m["at"])] = m + prev_t[replica] = t + + versions = sorted({v for per in first_seen.values() for v in per}) + convergence = { + v: round(max(per[v] for per in first_seen.values() if v in per) + - min(per[v] for per in first_seen.values() if v in per), 3) + for v in versions + if sum(v in per for per in first_seen.values()) > 1 + } + stage = sorted(m["stage_s"] for m in timings.values() if "stage_s" in m) + commit = sorted(m["commit_s"] for m in timings.values() if "commit_s" in m) + return { + "replicas": len(first_seen), + "versions_seen": versions, + "convergence_lag_s": convergence, # first-replica-at-vN -> last-replica-at-vN + "stage_s": _dist(stage), + "commit_s": _dist(commit), + "unready_s": {k: round(v, 1) for k, v in unready.items() if v}, + "probe_errors": errors, + } + + +def _dist(xs: list[float]) -> dict[str, float] | None: + if not xs: + return None + return {"n": len(xs), "p50": xs[len(xs) // 2], "p95": xs[int(len(xs) * 0.95)], "max": xs[-1]} + + +if __name__ == "__main__": + import argparse + + ap = argparse.ArgumentParser() + sub = ap.add_subparsers(dest="cmd", required=True) + p = sub.add_parser("poll") + p.add_argument("--app", required=True) + p.add_argument("--cls", default="Server") + p.add_argument("--interval", type=float, default=2.0) + p.add_argument("--duration", type=float, default=None) + p.add_argument("--out", required=True) + s = sub.add_parser("summarize") + s.add_argument("path") + args = ap.parse_args() + if args.cmd == "poll": + asyncio.run(poll(args.app, args.cls, interval=args.interval, duration=args.duration, out_path=args.out)) + else: + print(json.dumps(summarize(args.path), indent=2)) diff --git a/tools/probes/replay_publisher.py b/tools/probes/replay_publisher.py new file mode 100644 index 0000000..74a1ba9 --- /dev/null +++ b/tools/probes/replay_publisher.py @@ -0,0 +1,78 @@ +"""Replay a recorded delta chain through the real publish path at a controlled cadence. + +A pool certification doesn't need a trainer: past runs' chains persist on the delta +Volume as ``//weight_vNNNNNN/`` dirs, and ``publish_version()`` only +needs those dirs. Replaying re-publishes them under a fresh ``run_id`` against a live +pool — real delta densities, sizes, and chain structure, zero trainer GPUs. + +Caveats: the pool must serve the same base model the chain was recorded against, and +each publish *copies* the version dir under the new run prefix (that's the real path: +files land before the pointer moves), so the volume grows per replay and nothing GCs. +""" + +from __future__ import annotations + +import time +import uuid +from pathlib import Path + + +def replay( + *, + root: str, + source_run: str, + app_name: str | None = None, + cls_name: str = "Server", + pool: object = None, # any stitch Pool; overrides app_name + volume_name: str | None = None, + cadence_s: float = 30.0, + limit: int | None = None, + run_id: str | None = None, +) -> str: + """Claim the pool under a fresh run, then publish the recorded chain v1..vN with + ``cadence_s`` between publishes. Returns the run_id used. With neither ``pool`` + nor ``app_name`` there are no wakes — replicas converge via their backstop.""" + from stitch.publish import claim_run, publish_version + from stitch.stores.modal_volume import ModalVolumeStore + + store = ModalVolumeStore(root, volume_name=volume_name) + if pool is None and app_name is not None: + from stitch.pools.modal_flash import ModalFlashPool + + pool = ModalFlashPool(app_name, cls_name) + run_id = run_id or f"replay-{uuid.uuid4().hex[:8]}" + + src = Path(root) / source_run + dirs = sorted(d for d in src.iterdir() if d.is_dir() and d.name.startswith("weight_v")) + dirs = [d for d in dirs if int(d.name.removeprefix("weight_v")) > 0] # v0 is the pool's own base + if limit is not None: + dirs = dirs[:limit] + if not dirs: + raise SystemExit(f"no weight_v* dirs under {src}") + + print(f"replaying {len(dirs)} versions from {source_run!r} as run {run_id!r} at {cadence_s}s cadence") + claim_run(store, pool, run_id) + for d in dirs: + t = time.time() + ref = publish_version(store, pool, str(d), run_id=run_id) + print(f"published {ref.identity} in {time.time() - t:.1f}s") + time.sleep(cadence_s) + return run_id + + +if __name__ == "__main__": + import argparse + + ap = argparse.ArgumentParser() + ap.add_argument("--root", required=True, help="delta bulletin root (the mounted Volume path)") + ap.add_argument("--source-run", required=True) + ap.add_argument("--app", required=True, help="target pool app name") + ap.add_argument("--cls", default="Server") + ap.add_argument("--volume-name", default=None) + ap.add_argument("--cadence-s", type=float, default=30.0) + ap.add_argument("--limit", type=int, default=None) + args = ap.parse_args() + replay( + root=args.root, source_run=args.source_run, app_name=args.app, cls_name=args.cls, + volume_name=args.volume_name, cadence_s=args.cadence_s, limit=args.limit, + ) diff --git a/tools/probes/traffic.py b/tools/probes/traffic.py new file mode 100644 index 0000000..9b82aea --- /dev/null +++ b/tools/probes/traffic.py @@ -0,0 +1,201 @@ +"""Reward-free traffic shapes against a pool gateway. + +Shapes stress the serving/sync protocol, not the model — what matters is prompt and +decode length, concurrency, and session structure. ``agentic`` is the headline shape +(the top real use case): multi-turn sessions whose context grows every turn with a +synthetic tool-result blob, pinned to one replica via the session-affinity header — +which makes per-publish KV-namespace rotation (``extra_key``) directly measurable as +turn-latency inflation right after a version flip. + +Responses carry ``weight_version_start``/``weight_version_end`` (top-level on OpenAI +routes), so the generator doubles as the straddle-attribution collector. +""" + +from __future__ import annotations + +import asyncio +import json +import random +import time +from dataclasses import dataclass +from pathlib import Path +from typing import Any + +AFFINITY_HEADER = "Modal-Session-ID" # what the cookbook configs use +RETRY_409_SLEEP = 1.0 +RETRY_409_LIMIT = 120 +_WORDS = ("the model weights version pool replica delta chain anchor policy rollout " + "context token prefill decode publish commit stage pointer session request").split() + + +@dataclass(frozen=True) +class Shape: + prompt_tokens: tuple[int, int] + max_tokens: tuple[int, int] + turns: tuple[int, int] = (1, 1) + tool_tokens: tuple[int, int] = (0, 0) # per-turn context growth (agentic) + + +SHAPES: dict[str, Shape] = { + "long_decode": Shape(prompt_tokens=(200, 800), max_tokens=(4096, 12288)), + "long_prefill": Shape(prompt_tokens=(8_000, 24_000), max_tokens=(256, 1024)), + "agentic": Shape(prompt_tokens=(1_000, 3_000), max_tokens=(256, 1536), turns=(4, 12), tool_tokens=(500, 4_000)), +} +MIXED_WEIGHTS = {"long_decode": 0.4, "long_prefill": 0.2, "agentic": 0.4} + + +def _filler(rng: random.Random, tokens: int) -> str: + return " ".join(rng.choices(_WORDS, k=max(1, int(tokens * 0.75)))) # ~0.75 words/token + + +def _estimate_tokens(messages: list[dict[str, str]]) -> int: + return int(sum(len(m["content"].split()) for m in messages) / 0.75) + + +async def run( + gateway: str, + model: str, + *, + shape: str = "mixed", + concurrency: int = 16, + duration: float = 600.0, + lag: int | None = None, # floor requests at (gateway-observed version - lag); None = unconstrained + context_limit: int = 16384, # engine --context-length; sessions stop growing before it + out_path: str | None = None, + seed: int = 0, +) -> dict[str, Any]: + import httpx + + rows: list[dict[str, Any]] = [] + floor = _VersionFloor(gateway) if lag is not None else None + deadline = time.time() + duration + async with httpx.AsyncClient(timeout=3600.0, trust_env=False) as client: + if floor: + await floor.start(client) + workers = [ + asyncio.create_task( + _worker(client, gateway, model, shape, deadline, rows, floor, lag, context_limit, i, random.Random(seed + i)) + ) + for i in range(concurrency) + ] + await asyncio.gather(*workers) + if floor: + floor.stop() + if out_path: + p = Path(out_path) + p.parent.mkdir(parents=True, exist_ok=True) + p.write_text("".join(json.dumps(r) + "\n" for r in rows)) + return summarize(rows) + + +async def _worker(client, gateway, model, shape_name, deadline, rows, floor, lag, context_limit, worker_id, rng) -> None: # noqa: ANN001 + n = 0 # one worker = sequential sessions + while time.time() < deadline: + name = shape_name if shape_name != "mixed" else rng.choices(*zip(*MIXED_WEIGHTS.items()))[0] + await _session( + client, gateway, model, name, SHAPES[name], rng, rows, floor, lag, context_limit, + session_id=f"w{worker_id}-{n}", + ) + n += 1 + + +async def _session(client, gateway, model, name, spec, rng, rows, floor, lag, context_limit, session_id) -> None: # noqa: ANN001 + prompt_tokens = min(rng.randint(*spec.prompt_tokens), context_limit - max(spec.max_tokens) - 256) + messages = [{"role": "user", "content": _filler(rng, prompt_tokens)}] + headers = {AFFINITY_HEADER: session_id} + for turn in range(rng.randint(*spec.turns)): + max_tokens = rng.randint(*spec.max_tokens) + if _estimate_tokens(messages) + max_tokens + 256 > context_limit: + break # the session has outgrown the engine's context window + payload: dict[str, Any] = { + "model": model, + "messages": messages, + "max_tokens": max_tokens, + "temperature": 0.8, + } + if floor is not None and floor.version is not None: + payload["weight_version"] = {"min_version": max(0, floor.version - (lag or 0))} + row = {"t": time.time(), "shape": name, "session": session_id, "turn": turn, "retries_409": 0} + data = await _post_with_retry(client, f"{gateway}/v1/chat/completions", payload, headers, row) + rows.append(row) + if data is None: + return # session dies with its failed request + row.update( + wv_start=data.get("weight_version_start"), + wv_end=data.get("weight_version_end"), + straddled=data.get("weight_version_start") != data.get("weight_version_end"), + ) + content = (data.get("choices") or [{}])[0].get("message", {}).get("content") or "" + messages.append({"role": "assistant", "content": content}) + if spec.tool_tokens[1]: # agentic: the "tool result" grows the context every turn + messages.append({"role": "user", "content": f"tool result:\n{_filler(rng, rng.randint(*spec.tool_tokens))}\ncontinue."}) + else: + break + + +async def _post_with_retry(client, url, payload, headers, row) -> dict[str, Any] | None: # noqa: ANN001 + start = time.time() + for _ in range(RETRY_409_LIMIT): + try: + resp = await client.post(url, json=payload, headers=headers) + except Exception as exc: # noqa: BLE001 + row.update(latency=time.time() - start, error=str(exc)[:200]) + return None + if resp.status_code == 409: # version not ready: the retryable staleness signal + row["retries_409"] += 1 + await asyncio.sleep(RETRY_409_SLEEP) + continue + row.update(latency=time.time() - start, status=resp.status_code) + if resp.status_code != 200: + row["error"] = resp.text[:200] + return None + return resp.json() + row.update(latency=time.time() - start, error="409 retry budget exhausted") + return None + + +class _VersionFloor: + """Track the pool's applied version via the gateway's /server_info (answers from an + arbitrary replica — probe-grade, not exact).""" + + def __init__(self, gateway: str) -> None: + self.gateway = gateway + self.version: int | None = None + self._task: asyncio.Task[None] | None = None + + async def start(self, client) -> None: # noqa: ANN001 + from stitch.versions import VersionRef + + async def loop() -> None: + while True: + try: + info = (await client.get(f"{self.gateway}/server_info")).json() + if info.get("applied"): + self.version = VersionRef.parse(info["applied"]).version + except Exception: # noqa: BLE001 + pass + await asyncio.sleep(2.0) + + self._task = asyncio.create_task(loop()) + + def stop(self) -> None: + if self._task is not None: + self._task.cancel() + + +def summarize(rows: list[dict[str, Any]]) -> dict[str, Any]: + by_shape: dict[str, list[dict[str, Any]]] = {} + for r in rows: + by_shape.setdefault(r["shape"], []).append(r) + out: dict[str, Any] = {"requests": len(rows)} + for name, rs in sorted(by_shape.items()): + lat = sorted(r["latency"] for r in rs if "latency" in r) + out[name] = { + "n": len(rs), + "latency_p50": lat[len(lat) // 2] if lat else None, + "latency_p95": lat[int(len(lat) * 0.95)] if lat else None, + "straddled": sum(bool(r.get("straddled")) for r in rs), + "retries_409": sum(r.get("retries_409", 0) for r in rs), + "errors": sum("error" in r for r in rs), + } + return out diff --git a/tools/profiling/__init__.py b/tools/profiling/__init__.py new file mode 100644 index 0000000..620fdfb --- /dev/null +++ b/tools/profiling/__init__.py @@ -0,0 +1 @@ +"""Modal measurement probes (weight-sync reload profiling, disk surveys).""" diff --git a/tools/profiling/delta_audit.py b/tools/profiling/delta_audit.py new file mode 100644 index 0000000..23d584f --- /dev/null +++ b/tools/profiling/delta_audit.py @@ -0,0 +1,156 @@ +"""Audit an experiment's published deltas: touched names, reload closure, density. + +The O(delta) partial-reload cost scales with each version's *closure* (the +touched checkpoint tensors plus everything they drag in), so this audit is the +cheap leading indicator to run against a bulletin before spending GPU-hours: +it recomputes, from each version's ``model.safetensors.index.json`` alone, the +same closure the engine's load plan applies: + + - expert closure: one touched tensor of an expert pulls every tensor of that + expert (the per-expert post-loading transform consumes weights AND scales + together, and destroys the raw forms in place), and + - fused-sibling closure: a touched input of a load-time fusion pulls its + sibling (the model re-fuses only from a complete part set). + +It also fingerprints WHICH (layer, expert) slots each version touches and +compares them across runs — identical sets across independent runs indicate a +structural effect (e.g. padding tokens routed to expert 0), not learning. The +2026-06 K2.6 bulletin audit found exactly that: every published delta touched +only expert 0 per layer, so its 0.44% density reflects smoke-test-scale +training plus NVFP4 quantization masking sub-quantum updates, NOT a property +to extrapolate to production-scale runs. + +CPU-only, no GPUs, seconds per bulletin. Run: + EXPERIMENT_CONFIG=kimi_k2_6_nvfp4_disagg \ + uv run --extra modal modal run -m profiling.delta_audit::audit +Optionally filter with --run-id . +""" + +from __future__ import annotations + +import json +import os +import re +from collections import Counter +from datetime import datetime, timezone + +import modal + +# Local: resolve the experiment's volumes/paths from EXPERIMENT_CONFIG, and +# bake them into the image env so the in-container module re-import resolves +# the SAME objects without needing the cookbook package. +if modal.is_local(): + from cookbook.miles_disagg import modal_train as mt + + _DELTA_VOLUME_NAME = mt.exp.DELTA_VOLUME_NAME + _BULLETIN_ROOT = mt.exp.DELTA_BULLETIN_ROOT + # e.g. "kimi-k2-6-nvfp4/nvfp4" under the prep volume + _BASE_REL = os.path.relpath(mt.MODEL_NAME, str(mt.PREP_PATH)) +else: + _DELTA_VOLUME_NAME = os.environ["AUDIT_DELTA_VOLUME"] + _BULLETIN_ROOT = os.environ["AUDIT_BULLETIN_ROOT"] + _BASE_REL = os.environ["AUDIT_BASE_REL"] + +# Keep in sync with cookbook/miles_disagg/modal_train.py. +_PREP_VOLUME_NAME = "miles-prep-checkpoints" + +app = modal.App("weight-sync-delta-audit") +image = modal.Image.debian_slim(python_version="3.12").env( + { + "AUDIT_DELTA_VOLUME": _DELTA_VOLUME_NAME, + "AUDIT_BULLETIN_ROOT": _BULLETIN_ROOT, + "AUDIT_BASE_REL": _BASE_REL, + } +) +delta_volume = modal.Volume.from_name(_DELTA_VOLUME_NAME) +prep_volume = modal.Volume.from_name(_PREP_VOLUME_NAME) + +EXPERT_RE = re.compile(r"\.layers\.(\d+)\..*?\.experts\.(\d+)\.") +# Load-time fusions whose siblings must reload together (DSv3/Kimi attention). +FUSED_SIBLINGS = [("self_attn.q_a_proj", "self_attn.kv_a_proj_with_mqa")] + + +def _closure(touched: set[str], base_names: list[str]) -> set[str]: + expert_index: dict[str, list[str]] = {} + for name in base_names: + m = EXPERT_RE.search(name) + if m: + expert_index.setdefault(name[: m.end() - 1], []).append(name) + + out = set(touched) + for name in touched: + m = EXPERT_RE.search(name) + if m: + out.update(expert_index.get(name[: m.end() - 1], ())) + for a, b in FUSED_SIBLINGS: + for x, y in ((a, b), (b, a)): + if f".{x}." in name: + prefix = name.split(f".{x}.")[0] + out.update(n for n in base_names if n.startswith(f"{prefix}.{y}.")) + return out + + +@app.function( + image=image, + volumes={_BULLETIN_ROOT: delta_volume, "/prep": prep_volume}, + timeout=1800, +) +def audit(run_id: str = "") -> None: + base_index = os.path.join("/prep", os.environ["AUDIT_BASE_REL"], "model.safetensors.index.json") + with open(base_index) as f: + base_names = list(json.load(f)["weight_map"]) + expert_slots = {m.group(1, 2) for n in base_names if (m := EXPERT_RE.search(n))} + root = os.environ["AUDIT_BULLETIN_ROOT"] + print(f"base: {len(base_names)} tensors, {len(expert_slots)} expert slots ({base_index})") + print(f"{'run':>14} {'ver':>4} {'mtime':>12} {'touched':>8} {'experts':>8} {'closure':>8} {'density':>8} payload") + + per_version_experts: dict[str, dict[str, set]] = {} + for run in sorted(os.listdir(root)): + if run_id and run != run_id: + continue + run_dir = os.path.join(root, run) + if not os.path.isdir(run_dir): + continue + for ver in sorted(d for d in os.listdir(run_dir) if d.startswith("weight_v")): + vdir = os.path.join(run_dir, ver) + files = sorted(os.listdir(vdir)) if os.path.isdir(vdir) else [] + idx_path = os.path.join(vdir, "model.safetensors.index.json") + if not os.path.isfile(idx_path): + if files: + print(f"{run:>14} {ver[-4:]:>4} (no index; files: {files[:4]})") + continue + with open(idx_path) as f: + weight_map = json.load(f)["weight_map"] + touched = set(weight_map) + payload = sum( + os.path.getsize(os.path.join(vdir, f)) + for f in files + if f.endswith(".safetensors") + ) + missing_shards = set(weight_map.values()) - set(files) + mtime = max(os.path.getmtime(os.path.join(vdir, f)) for f in files) + stamp = datetime.fromtimestamp(mtime, tz=timezone.utc).strftime("%m-%d %H:%M") + experts = {m.group(1, 2) for n in touched if (m := EXPERT_RE.search(n))} + per_version_experts.setdefault(ver, {})[run] = experts + clo = _closure(touched, base_names) + print( + f"{run:>14} {ver[-4:]:>4} {stamp:>12} {len(touched):8d} {len(experts):8d} " + f"{len(clo):8d} {len(clo) / len(base_names):8.2%} {payload / 1e6:8.1f} MB" + + (f" MISSING SHARDS: {sorted(missing_shards)}" if missing_shards else "") + ) + kinds = Counter( + "expert" if EXPERT_RE.search(n) else n.rsplit(".", 1)[-1] for n in touched + ) + print(f"{'':>14} breakdown: {dict(kinds.most_common(5))}") + if experts and len(experts) <= 6: + print(f"{'':>14} touched experts (layer, id): {sorted(experts, key=lambda t: (int(t[0]), int(t[1])))}") + + print("\n===== cross-run expert-slot comparison (identical sets => structural, not learning) =====") + for ver, by_run in sorted(per_version_experts.items()): + sets = {r: e for r, e in by_run.items() if e} + if len(sets) < 2: + continue + inter = set.intersection(*sets.values()) + union = set.union(*sets.values()) + print(f"{ver}: runs={len(sets)} intersection={len(inter)} union={len(union)}") + print("\nAUDIT DONE") diff --git a/tools/profiling/modal_probes.py b/tools/profiling/modal_probes.py new file mode 100644 index 0000000..1604550 --- /dev/null +++ b/tools/profiling/modal_probes.py @@ -0,0 +1,414 @@ +"""Weight-sync measurement probes: where does the SGLang reload time go? + +Brings up ONE rollout engine identical to the miles_disagg Server cls (same +serving image, same SGLangEndpoint(model_path, tp, SGLANG_SERVER_ARGS), same +volumes + ephemeral disk) and times each phase of the weight-update path, +isolating disk-read throughput from SGLang's load processing: + + 1. raw cold-read throughput of a base shard 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 cold+warm read throughput of a /local-checkpoint shard — the disk the + steady-state reload actually reads, + 5. reload from /local-checkpoint with a COLD page cache, + 6. immediate second reload (WARM page cache; doubles as the two-consecutive- + reloads stability gate for the quant postprocess), + 7. optionally, host-side delta apply + reload (pass --delta-run-id/--version). + +With the instrumented serving pin, phases 5-7 also emit the engine's +"[reload timing] iter_wait=..s load=..s postprocess=..s total=..s" line, which +splits disk/materialize waits from weight_loader dispatch and the quant +postprocess — read those (and "[disk delta apply] ..." for phase 7) from the +same app logs as this probe's summary. + +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 profiling.modal_probes::profile_reload +Read results: modal app logs -e weight-sync-probes +""" + +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("weight-sync-probes") + +# The reference NVFP4 checkpoint the anchor probe cold-loads/reloads (must +# already be in the HF cache volume; the probe never downloads it). +DEFAULT_ANCHOR_MODEL = "nvidia/Kimi-K2.6-NVFP4" + + +def _drop_page_caches() -> bool: + try: + with open("/proc/sys/vm/drop_caches", "w") as f: + f.write("3") + return True + except Exception as e: # noqa: BLE001 + print(f"[diskread] (could not drop page cache: {e})") + return False + + +def _dd_read_gbps(path: str, label: str, *, cold: bool = True) -> float: + """Read a file with dd (optionally dropping the page cache first); return GB/s.""" + if cold: + _drop_page_caches() + 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, + memory=mt.modal_cfg.rollout_memory_mib, + timeout=120 * mt.MINUTES, +) +def profile_reload(delta_run_id: str = "", delta_version: int = 0, load_plan: int = 0) -> None: + import asyncio + + import httpx + from autoinference_utils.endpoint import SGLangEndpoint + + from cookbook.miles_disagg import helpers + from cookbook.sidecar import parallel_init_local_checkpoint + + # Same runtime sglang patches the Server cls applies (e.g. the FP8 + # reload-attrs patch) — the probe must serve the identical code. + helpers.apply_sglang_runtime_patches(list(getattr(mt.exp, "SGLANG_RUNTIME_PATCHES", []))) + + if load_plan: + # Reload record/replay fast path (fork's model_loader/load_plan.py): + # reload #1 records the dispatch plan, #2/#3 replay it. Inherited by + # the SGLang server subprocess. + os.environ["SGLANG_ENABLE_RELOAD_LOAD_PLAN"] = "1" + + base = mt.MODEL_NAME # the served base on /prep + 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() + body = r.json() + if body.get("success") is False: + raise RuntimeError(f"reject: {body}") + print(f"[reload message] {body.get('message')}") + return time.perf_counter() - t + + def _weights_checksum(label: str) -> str: + """Per-GPU weight checksums via /weights_checker — identical bytes in + must yield identical checksums out, which is the record-vs-replay + correctness gate for the load plan.""" + try: + r = httpx.post(f"{sglang_url}/weights_checker", + json={"action": "checksum"}, timeout=1800.0) + r.raise_for_status() + digest = str(r.json()) + print(f"[checksum] {label}: {digest[:160]}") + return digest + except Exception as e: # noqa: BLE001 + print(f"[checksum] {label} failed: {e}") + return f"ERROR: {e}" + + 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. materialize base -> /local-checkpoint, the disk every steady-state + # reload actually reads + print("[phase] materialize base -> /local-checkpoint (32 workers)...") + t0 = time.perf_counter() + parallel_init_local_checkpoint("miles.utils.disk_delta")(local, base) + R["materialize_s"] = round(time.perf_counter() - t0, 1) + print(f"[phase] materialize: {R['materialize_s']}s") + + # 4. raw ephemeral-NVMe read rate, cold and warm — the disk-bound-vs- + # loader-bound reference points for the reload numbers below + local_shards = sorted(glob.glob(f"{local}/*.safetensors")) + if local_shards: + R["dd_local_cold_gbps"] = round(_dd_read_gbps(local_shards[0], "ephemeral /local-checkpoint shard (cold)"), 2) + R["dd_local_warm_gbps"] = round( + _dd_read_gbps(local_shards[0], "ephemeral /local-checkpoint shard (warm)", cold=False), 2 + ) + + # 5. reload from /local-checkpoint with a cold page cache (the steady-state + # worst case: the checkpoint got evicted during the rollout). With + # load_plan=1 this reload RECORDS the dispatch plan. + if alive(): + print("[phase] reload /local-checkpoint (cold page cache)...") + R["page_cache_dropped"] = _drop_page_caches() + try: + R["reload_local_cold_s"] = round(asyncio.run(_reload(local)), 1) + except Exception as e: # noqa: BLE001 + R["reload_local_cold"] = f"CRASH: {str(e)[:120]}" + print(f"[phase] cold reload CRASH: {e}") + if load_plan: + R["checksum_after_record"] = _weights_checksum("after record reload") + + # 6. immediate second reload (warm page cache) — isolates disk reads from + # dispatch+postprocess, and gates on the quant postprocess surviving + # consecutive reloads in one process. With load_plan=1 this is the first + # REPLAY of the plan recorded by the previous reload. + if alive(): + print("[phase] reload /local-checkpoint again (warm page cache)...") + try: + R["reload_local_warm_s"] = round(asyncio.run(_reload(local)), 1) + R["consecutive_reload"] = "OK" + except Exception as e: # noqa: BLE001 + R["consecutive_reload"] = f"CRASH: {str(e)[:120]}" + print(f"[phase] warm reload CRASH: {e}") + if load_plan: + R["checksum_after_replay"] = _weights_checksum("after replay reload") + R["plan_checksum_match"] = R.get("checksum_after_record") == R["checksum_after_replay"] + + # 6b. with the load plan enabled: a second replay — the FAST pass executing + # compiled copy programs — with its own checksum gate, since it is a + # different code path than the observed/compile pass before it + if load_plan and alive(): + print("[phase] reload /local-checkpoint a third time (plan fast pass)...") + try: + R["reload_local_replay2_s"] = round(asyncio.run(_reload(local)), 1) + R["checksum_after_fast"] = _weights_checksum("after fast reload") + R["plan_checksum_match_fast"] = R.get("checksum_after_record") == R["checksum_after_fast"] + except Exception as e: # noqa: BLE001 + R["reload_local_replay2"] = f"CRASH: {str(e)[:120]}" + print(f"[phase] fast reload CRASH: {e}") + + # 7. optional: the full steady-state O(delta) cycle against REAL published + # deltas — host-apply the chain, PARTIAL reload (touched names only), + # then prove it byte-equivalent to a FULL reload of the same bytes. + if delta_run_id and alive(): + import json as _json + + from miles.utils.disk_delta import apply_deltas + + delta_root = f"{mt.exp.DELTA_BULLETIN_ROOT}/{delta_run_id}" + print(f"[phase] apply delta chain {delta_root} up to v{delta_version}...") + t0 = time.perf_counter() + R["delta_apply_stats"] = apply_deltas(local, delta_root, delta_version) + R["delta_apply_s"] = round(time.perf_counter() - t0, 1) + + touched: set[str] = set() + for v in range(1, delta_version + 1): + index = f"{delta_root}/weight_v{v:06d}/model.safetensors.index.json" + with open(index) as f: + touched.update(_json.load(f)["weight_map"]) + R["delta_touched_names"] = len(touched) + + async def _partial_reload() -> 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": local, "weight_version": "probe-partial", + "flush_cache": False, "weight_names": sorted(touched)}) + r.raise_for_status() + body = r.json() + if body.get("success") is False: + raise RuntimeError(f"reject: {body}") + print(f"[reload message] {body.get('message')}") + return time.perf_counter() - t + + try: + print(f"[phase] PARTIAL reload ({len(touched)} touched names)...") + R["partial_reload_s"] = round(asyncio.run(_partial_reload()), 1) + R["checksum_after_partial"] = _weights_checksum("after partial reload") + print("[phase] FULL reload of the same bytes (equivalence reference)...") + R["full_reload_after_delta_s"] = round(asyncio.run(_reload(local)), 1) + R["checksum_after_full"] = _weights_checksum("after full reload") + R["partial_checksum_match"] = ( + R["checksum_after_partial"] == R["checksum_after_full"] + ) + # steady-state partial timing, engine state already at the target + R["partial_reload2_s"] = round(asyncio.run(_partial_reload()), 1) + except Exception as e: # noqa: BLE001 + R["partial_reload"] = f"CRASH: {str(e)[:120]}" + print(f"[phase] partial reload CRASH: {e}") + + 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, + 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(anchor_model: str = DEFAULT_ANCHOR_MODEL) -> None: + """Cold-load a reference checkpoint of the same model (e.g. the nvidia NVFP4 + anchor), then reload IT via update_weights_from_disk — the ours-vs-anchor + comparison that separates checkpoint-layout effects from engine reload-path + behavior. Record the resolved moe_runner_backend from the server logs with + every number: anchor and ours may resolve differently.""" + import asyncio + + import httpx + from autoinference_utils.endpoint import SGLangEndpoint + from huggingface_hub import snapshot_download + + anchor = snapshot_download(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 — {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() + body = r.json() + if body.get("success") is False: + raise RuntimeError(f"reject: {body}") + print(f"[reload message] {body.get('message')}") + 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] 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 + + +@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, + ephemeral_disk=mt.modal_cfg.rollout_ephemeral_disk_mib, + timeout=20 * mt.MINUTES, + single_use_containers=True, # one sample per container: the whole point is fresh disk draws +) +def disk_survey(sample: int = 0, write_gb: int = 16) -> dict: + """One ephemeral-disk sample on a fresh container of the rollout pool's + exact shape: direct-IO write and read rates (bypassing the page cache, + which containers cannot drop) plus buffered/cached re-read, with placement + metadata. Evidence base for the observed 0.2-1.1 GB/s cold-read spread.""" + import json + + root = mt.LOCAL_CHECKPOINT_PATH + os.makedirs(root, exist_ok=True) + path = f"{root}/_disk_survey.bin" + blocks = max(1, (write_gb << 30) // (16 << 20)) + result: dict[str, object] = { + "sample": sample, + "placement": {k: v for k, v in os.environ.items() if k.startswith("MODAL_") and len(v) < 80}, + } + + def timed_dd(label: str, args: list[str]) -> float: + t0 = time.perf_counter() + subprocess.run(args, check=True, stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL) + gbps = (blocks * (16 << 20)) / 1e9 / (time.perf_counter() - t0) + result[label] = round(gbps, 2) + print(f"[disk survey {sample}] {label}: {gbps:.2f} GB/s") + return gbps + + timed_dd("write_direct_gbps", + ["dd", "if=/dev/zero", f"of={path}", "bs=16M", f"count={blocks}", "oflag=direct", "conv=fsync"]) + timed_dd("read_direct_gbps", + ["dd", f"if={path}", "of=/dev/null", "bs=16M", "iflag=direct"]) + timed_dd("read_buffered_gbps", ["dd", f"if={path}", "of=/dev/null", "bs=16M"]) + timed_dd("read_cached_gbps", ["dd", f"if={path}", "of=/dev/null", "bs=16M"]) + os.remove(path) + print(f"[disk survey {sample}] {json.dumps(result)}") + return result + + +@probe_app.local_entrypoint() +def survey(samples: int = 4) -> None: + """Fan the disk survey across `samples` fresh containers and print each.""" + for result in disk_survey.map(range(samples)): + print(result) diff --git a/uv.lock b/uv.lock index ba7f14f..0dd7104 100644 --- a/uv.lock +++ b/uv.lock @@ -971,7 +971,6 @@ modal = [ sglang = [ { name = "fastapi" }, { name = "httpx" }, - { name = "uvicorn" }, ] [package.dev-dependencies] @@ -985,7 +984,6 @@ requires-dist = [ { name = "httpx", marker = "extra == 'modal'" }, { name = "httpx", marker = "extra == 'sglang'" }, { name = "modal", marker = "extra == 'modal'", specifier = ">=1.5.1.dev9" }, - { name = "uvicorn", marker = "extra == 'sglang'" }, ] provides-extras = ["modal", "sglang"] @@ -1052,19 +1050,6 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/dc/9b/47798a6c91d8bdb567fe2698fe81e0c6b7cb7ef4d13da4114b41d239f65d/typing_inspection-0.4.2-py3-none-any.whl", hash = "sha256:4ed1cacbdc298c220f1bd249ed5287caa16f34d44ef4e9c3d0cbad5b521545e7", size = 14611, upload-time = "2025-10-01T02:14:40.154Z" }, ] -[[package]] -name = "uvicorn" -version = "0.49.0" -source = { registry = "https://pypi.org/simple" } -dependencies = [ - { name = "click" }, - { name = "h11" }, -] -sdist = { url = "https://files.pythonhosted.org/packages/c4/1f/fa18009dea8469069cca78a4e877a008ab78f08b064bfc9ab891579077ff/uvicorn-0.49.0.tar.gz", hash = "sha256:ebf4271aa580d9de97f93192d4595176df6e91f9aae919ca73e4fc07df1e66a3", size = 91284, upload-time = "2026-06-03T22:01:30.448Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/88/fa/e1388bbcf24ef3274f45c0c1c7b501fd14971037c1b6ee23610553307497/uvicorn-0.49.0-py3-none-any.whl", hash = "sha256:ba3d14c3ee7e41c6c654c46c9eb489d33213cdd30aa1696eab1374337c13f68f", size = 71376, upload-time = "2026-06-03T22:01:29.037Z" }, -] - [[package]] name = "watchfiles" version = "1.2.0"