diff --git a/plans/miles-nccl-broadcast-plan.md b/plans/miles-nccl-broadcast-plan.md new file mode 100644 index 0000000..bd0ab48 --- /dev/null +++ b/plans/miles-nccl-broadcast-plan.md @@ -0,0 +1,731 @@ +# Plan: Enable NCCL broadcast path for miles weight update (rlops/rlix#42) + +Status: DRAFT v8 (2026-08-16) — implementation in progress +v8 (implementation-discovered constraint): miles `assert_rlix_topology` +C1 gate REQUIRES train ⊂ infer (partial overlap is the RLix-mode +contract), so the planned M4 run (b) disjoint-pool topology cannot start +at all, and a fully-disjoint all-broadcast e2e is unreachable without +relaxing the C1 gate (out of scope — rlix_validation untouched per plan +C1). AC7/E9 narrowed: strict `broadcast` mode is verified by (a) E3 unit +matrix (incl. rejection + subset-target success), (b) an e2e REJECTION +check on the overlap harness (mode=broadcast → classification fail-fast, +log captured). The all-broadcast positive e2e is deferred as new O7 +until disjoint topologies are admitted by the M11 validation contract. +Strict mode remains usable today for sync targets that avoid colocate +engines (E3 subset case). +v7: addresses codex round-5 (1 high): the v6 uniform engine-GPU-count +assumption was declared (O6) but not enforced — SGLang server groups can +override `num_gpus_per_engine` per group and the existing C7 checks do not +compare resolved group TP against `rollout_num_gpus_per_engine`. Fixed +with a rlix-side startup uniformity guard (miles still untouched beyond +actor.py): before `register_model_update_resources`, rlix validates the +resolved SGLang config — every server group's `num_gpus_per_engine` equals +`rollout_num_gpus_per_engine` — and fail-fasts otherwise (new E10). +v6 (user redirection): transport is a USER-SELECTED mode, not silent +auto-classification — rlix-side env `RLIX_MILES_UPDATE_TRANSPORT ∈ +{cpu_serialize (default) | broadcast | auto}`; strict `broadcast` never +mixes (fail-fast on colocate targets); `auto` is the only mixing mode and +is what the mandated overlap M4 uses. Minimal-change pass: miles-side +changes shrink to `actor.py` + tests ONLY (dropped the planned +`rollout.py` manager method — rlix derives per-engine GPU counts from its +own `rollout_num_gpus_per_engine` arg); single env replaces the separate +kill-switch. +v5: addresses codex round-3 (1 medium): the harness hardcodes +`--sglang-mem-fraction-static 0.30` (run_smoke_dual.sh:125), contradicting +the mandated 0.8 — resolved by parameterizing the harness mem fraction via +env (default 0.30 preserved; M4 runs 0.8). Also documents the existing +miles C11 transport-flag gate + S2 VRAM gate discovered during the fix. +v4: addresses codex round-2 (2 high): C9 re-aligned to the ACTUAL +run_smoke_dual.sh topology (P1 train [0] / infer [0,1,2]; P2 train [3] / +infer [1,2,3]); explicit timeout hierarchy (sender budgets strictly inside +the service session deadline) + port-claim ownership on abnormal exit +(leak-and-log, never release under a possibly-live TCP store). +v2: M4 e2e re-specified per user requirement — dual-pipeline overlap topology +(not single-pipeline disaggregated), GPU memory limit 0.8. +v3: addresses codex adversarial review (2 high, 2 medium): sender-owned NCCL +teardown + bounded rendezvous timeouts + partial-receiver-failure abort +semantics; pre-committed M4 GPU split + mixed-class startup assertion; +memory preflight + adaptive staging. +Issue: https://github.com/rlops/rlix/issues/42 — "Weight update in miles currently +only support cpu_serialize, we need to support broadcast for distributed setting." + +Branches: +- rlix: `zhenyu/miles-nccl-broadcast` (based on `zhenyu/miles-mvp-e2e` @ b4e0cf6) +- miles: `zhenyu/miles-nccl-broadcast` (based on `zhenyu/m11-mvp-test` @ 6ff8df3) + +--- + +## 1. Problem context + +### Current state + +The M11 RLix-mode weight sync chain is: + +``` +MilesCoordinator (rlix) — sync_base_weights_to_active / _expand_workers + → MilesModelUpdateService (rlix) — builds SyncSessionPlan, atomic timeout, port claim + → MegatronTrainRayActor.run_sync_session (miles, cache_owner) — single composite RPC + → Path A: _dispatch_cpu_serialize_bucket — torch.save bytes → engine tmpfs → HTTP [WORKS] + → Path B: _dispatch_nccl_broadcast — dynamic NCCL group broadcast [BLOCKED] + → SGLangEngine.setup_collective_group / broadcast_parameter / destroy_collective_group [EXIST, dead code] +``` + +Two deliberate fail-fast guards block Path B today: + +1. rlix `miles_model_update_service.py:199` — `sync_selected_workers` raises + `NotImplementedError` when `broadcast_local_ranks` is non-empty. +2. miles `megatron_utils/actor.py:968` — `_dispatch_nccl_broadcast` raises + `NotImplementedError` unconditionally; the receiver-side fan-out code below the + raise (setup / per-bucket broadcast_parameter / destroy) is dead code. + +The missing piece both guards protect against: **sender-side NCCL** on the +cache_owner (join the dynamic group as rank 0, `dist.broadcast` each bucket +tensor, destroy the group). Without it, receivers would block forever inside +SGLang `/init_weights_update_group` waiting for rank 0. + +Additionally, neither rlix call site (`sync_base_weights_to_active`, +`_expand_workers` F40 runtime branch) ever passes `broadcast_local_ranks` — +there is no classification logic deciding which engines need broadcast. + +### Reference implementation (do not modify, reuse pattern only) + +miles standalone (non-RLix) mode already ships a working sender-side NCCL path: +`miles/backends/megatron_utils/update_weight/update_weight_from_distributed/broadcast.py` +(`UpdateWeightFromDistributed`, `connect_rollout_engines_from_distributed`, +`update_weights_from_distributed`). Key proven pattern: + +1. dispatch receiver `init_weights_update_group` `.remote()` refs (do NOT ray.get yet) +2. sender `init_process_group(backend="nccl", init_method=tcp://addr:port, rank=0, world_size)` + — this blocks until all receivers join, which is why step 1 must not block +3. `ray.get(refs)` only after the sender has joined +4. per transfer: dispatch receiver metadata RPC refs (async) → sender + `dist.broadcast(param, src=0, group, async_op=True)` + wait → `ray.get(refs)` +5. teardown: receiver `destroy_weights_update_group` refs + sender + `dist.destroy_process_group` + `ray.get(refs)` + +### Why broadcast matters (issue motivation) + +- cpu_serialize routes weights CPU → tmpfs → HTTP → engine. It requires + sender and receiver on the **same node** (tmpfs path handoff) and pays + serialize/deserialize + disk + HTTP cost per bucket. +- NCCL broadcast is GPU→GPU, required for multi-node ("distributed setting" + in the issue) and substantially faster for disaggregated single-node + topologies (train pool ∩ infer pool = ∅, e.g. M11.2 P2 `[2,3]` vs train `[0,1]`). + +## 2. Assumptions + +- A1: SGLang receiver admin routes behave per the existing standalone path: + `/init_weights_update_group` registers `tp_size` consecutive NCCL ranks + starting at `rank_offset`; `/update_weights_from_distributed` blocks until + the named tensors arrive via NCCL from rank 0. +- A2: `bucket.params` is an insertion-ordered dict (Python ≥3.7), so the + `names`/`dtypes`/`shapes` metadata lists and the sender's broadcast order + agree by construction. +- A3: The cache_owner train actor has a live CUDA context and enough free GPU + memory to stage one bucket (default cap 512 MB, + `miles_model_update_bucket_size_mb`) during transport. +- A4: MVP validation targets a single node, engines with + `nodes_per_engine == 1`; multi-node engines (node_rank > 0 shards) are out + of MVP scope (their `setup_collective_group` returns `{}` today). +- A5: `cluster_device_mappings` (actor_train / actor_infer physical GPU lists) + and `rollout_num_gpus_per_engine` are sufficient to derive engine_index ↔ + physical-GPU sets, as already done in + `MilesPipeline._wait_for_overlap_engines_offloaded`. +- A6: Known reviewer focus (PR rlops/RL#3 history): NCCL rendezvous deadlock + ordering and global-vs-local rank confusion are the two highest-risk bug + classes for this change family. + +## 3. Hard constraints + +- C1 (minimize-upstream-miles, tightened in v6): miles-side changes are + confined to `miles/backends/megatron_utils/actor.py` + (`_dispatch_nccl_broadcast` + in-method helpers — all RLix-port-added + code) and tests. NOTHING else in miles changes: no `rollout.py`, no + `arguments.py`, no `rlix_validation.py` (C11/S2 gates untouched), no + `sglang_engine.py`, and never any pristine upstream line (including + `UpdateWeightFromDistributed`). +- C2 (F04 single composite RPC): no new top-level Ray methods on the train + actor for transport; sender NCCL lives inside `run_sync_session` as + in-method helpers. +- C3 (F26/C16): `master_port != 0`; port claimed via SharedStorage before the + plan is sent (already implemented — keep). +- C4 (F21): exactly one `set_weight_version` publish per sync, by the service + only (unchanged). +- C5 (fail-fast): heterogeneous/invalid plans raise; no silent fallback from + broadcast to cpu_serialize at transport time — and (v6) no silent mode + switching either: the user-selected transport mode is honored exactly or + the run refuses to start (strict `broadcast` + colocate target = + startup rejection, never a quiet downgrade to mixing). +- C6 (atomicity): the whole transport remains inside the service's single + `asyncio.wait_for` timeout. **Enforcement split (codex v3)**: the + service-side `asyncio.wait_for` + `ray.cancel` CANNOT interrupt a sender + blocked inside a native NCCL/HTTP call; the real bound is sender-side + timeouts, so the sender task always unwinds by itself and its own + `finally` performs the teardown. Service-level `inflight_refs` covers + only service-issued refs (top-level RPC, manager calls); it is NOT + claimed to cover the sender's nested receiver refs. + **Deadline hierarchy (codex v4)** — one shared value is NOT a valid + hierarchy; the budgets nest strictly: + - session deadline = plan `timeout_s` (service `asyncio.wait_for`, + default 150 s) — outermost; + - sender internal budgets are fractions of `timeout_s` chosen so the + sender's worst case (rendezvous + all per-bucket collectives + + teardown grace) completes strictly inside the session deadline: + rendezvous ≤ 0.2×, transport total ≤ 0.6×, teardown grace ≤ 0.1×, + safety margin ≥ 0.1× (exact constants are implementation detail; the + invariant "sender worst-case unwind < session deadline" is the + requirement); + - consequence: the service deadline fires only when the sender is truly + wedged past its own budgets (native call refusing to unwind), which is + the fail-fast pipeline-death path — not a normal cleanup path. + **Port-claim ownership on abnormal exit (codex v4, refined v8-r1)**: + the claim is released only after sender teardown is acknowledged — + i.e. after the `run_sync_session` ref resolves (success OR exception; + the sender's `finally` has run by then). Sender-ref resolution is + tracked independently of claim presence, so a cancellation landing in + the post-resolution release window still releases (fire-and-forget) + instead of leaking. Only if the service session deadline expires + while the sender ref is UNRESOLVED AND the session had a non-empty + broadcast set is the claim **intentionally leaked and logged loudly**, + never deleted — the TCP store may still be bound to that port, and the + existing ROLL-backend precedent treats a leaked claim as safer than a + collision. Fresh per-session `get_free_port` picks prevent rendezvous + reuse; leaked claims die with the SharedStorage actor (fail-fast + lifecycle). cpu_serialize-only sessions keep today's release-on-timeout + (no TCP store exists). +- C7 (no-commit): no commit/push without explicit user instruction; codex + review approval required before sign-off. +- C8 (world-size accounting): NCCL group ranks are per-GPU, not per-engine: + `world_size = 1 + Σ gpu_count(broadcast engines)`; engine `rank_offset`s + form a cursor. The current rlix per-engine assignment is only valid for + TP=1 and must be generalized. +- C9 (user-mandated test topology): M4 e2e validation MUST run the + dual-pipeline **overlap** topology (M11.2-style, env-driven + `MILES_DUAL_P*` via `scripts/run_smoke_dual.sh`) with GPU memory limit + 0.8 (SGLang `mem_fraction_static=0.8`). Single-pipeline disaggregated + smoke is NOT an acceptable substitute. **Topology = the harness's actual + checked-in defaults (codex v4 — the v3 draft had invented a split that + contradicted the harness)**, verified against `scripts/run_smoke_dual.sh` + lines 68-71: `MILES_DUAL_P1_TRAIN=0`, `MILES_DUAL_P1_INFER=0,1,2`, + `MILES_DUAL_P2_TRAIN=3`, `MILES_DUAL_P2_INFER=1,2,3` (4 GPUs, TP=1, + 1 GPU/engine). Expected classification: P1 → e0@gpu0 cpu_serialize, + e1@gpu1 + e2@gpu2 broadcast; P2 → e0@gpu1 + e1@gpu2 broadcast, e2@gpu3 + cpu_serialize. Both pipelines are mixed-transport by construction of the + real harness envs — the topology needs no harness change. **Memory + limit needs one (codex v5)**: the harness hardcodes + `--sglang-mem-fraction-static 0.30` (`run_smoke_dual.sh:125`), so the + mandated 0.8 would silently not apply. M4 parameterizes that line via + env — `MILES_SMOKE_MEM_FRACTION` (default `0.30`, preserving existing + harness behavior) — and the M4 invocation exports + `MILES_SMOKE_MEM_FRACTION=0.8`. E4 evidence must include the effective + mem-fraction value echoed from the launch log; a 0.30 run does NOT + satisfy AC6. The topology env values are the source of truth: the + startup assertion (before the first training step) + logs the actual env-derived train/infer mappings AND the actual + scheduler grant used for classification, then verifies each pipeline's + active sync target set contains ≥1 broadcast-classified AND ≥1 + cpu_serialize-classified engine; the smoke is INVALID (fail-fast, not + skip) if the assertion does not hold or if the logged mappings diverge + from the harness envs. + +## 4. Acceptance criteria + +- AC1: A sync session whose plan carries a non-empty `broadcast_local_ranks` + completes end-to-end — no `NotImplementedError` on either side; targeted + engines receive all buckets via NCCL and serve the published weight version. + (Derives from: Goal, A1, A2) +- AC2: All-cpu_serialize behavior is unchanged: existing unit tests and the + colocate smoke pass without modification to expectations. (Derives from: + Goal, C5) +- AC3: NCCL rank/world_size accounting is correct for TP>1 engines + (rank-offset cursor, per-GPU world size), verified at unit level. (Derives + from: C8, A1) +- AC4: Transport is user-selected and its topology inputs are + startup-validated: the v7 uniformity guard rejects resolved SGLang + configs whose server groups diverge from + `rollout_num_gpus_per_engine` before any classification runs. Mode via + rlix-side env + `RLIX_MILES_UPDATE_TRANSPORT ∈ {cpu_serialize | broadcast | auto}`, + default `cpu_serialize` (= today's behavior, zero change until opt-in): + `cpu_serialize` → all engines cpu_serialize; `broadcast` → ALL engines + broadcast, with classification-time fail-fast rejection if any target + engine's GPU intersects the pipeline's train pool (NCCL cannot form a + group with duplicate physical GPUs — sender-colocate targets are + unservable); `auto` → topology classification (engine GPUs ∩ train pool + = ∅ → broadcast, else cpu_serialize; the only mode that mixes). Mode is + logged once per pipeline at startup. (Derives from: Goal, A5, C5) +- AC5: Any abnormal end of a broadcast session — receiver failure during + rendezvous (dies / rejects / hangs), mid-bucket exception, or sender + budget expiry — unwinds within the C6 deadline hierarchy (sender + worst-case < session deadline), executes the sender-owned teardown + (destroy both sides, cancel nested receiver refs), and then releases + the SharedStorage port claim (release strictly after teardown ack). On + the residual wedged-sender path (service session deadline fires first) + the claim is leaked-and-logged, never deleted. No live group, claimed + port under a live TCP store, or unbounded block survives into the next + sync. (Derives from: C6, C3) +- AC6: E2E dual-pipeline overlap smoke on a 4-GPU vast.ai instance (C9 + pre-committed topology, GPU memory limit 0.8, mode `auto` — the overlap + topology requires the mixing mode by NCCL's duplicate-GPU constraint) + exercises **mixed-transport sync sessions** — broadcast for engines + outside each pipeline's granted train GPUs, cpu_serialize for overlap + engines — log-verified per pipeline, with the C9 mixed-class startup + assertion passing, **zero GPU OOM events** across the run, the existing + dual-smoke pass-bar met, and EXIT_CODE=0. (Derives from: Goal, A4, A5, + C9) +- AC7 (narrowed in v8): Strict `broadcast` mode honors its no-mix + contract, verified by (a) the E3 unit matrix — colocate target → + rejection naming the engines; all-disjoint and colocate-avoiding + subset targets → all-broadcast; and (b) an e2e rejection check on the + overlap harness: `RLIX_MILES_UPDATE_TRANSPORT=broadcast` fails fast at + the first classification with the actionable error (log captured). The + all-broadcast POSITIVE e2e is O7-deferred: miles `assert_rlix_topology` + C1 requires train ⊂ infer, so no startable RLix topology can make + every engine broadcast-eligible today. (Derives from: Goal, A4, A5, C5) + +## 5. Milestones + +- M1 — miles sender-side NCCL transport. Implement sender join (bounded + timeout) + per-bucket broadcast with memory preflight + sender-owned + teardown in `_dispatch_nccl_broadcast`; remove the miles-side guard; + unit tests (E1, E7, E8) with mocked engine handles + fake dist. + Delivers: AC1 (miles half), AC5 (miles half). +- M2 — rlix service unlock + rank accounting. Remove the service raise; fix + `comm_ranks`/`world_size` to the per-GPU cursor scheme (uniform stride + from `rollout_num_gpus_per_engine`, injected — no manager RPC, v6); + C6 claim/deadline rules; unit tests. Delivers: AC1 (rlix half), AC3, + AC5 (rlix half). +- M3 — rlix transport-mode wiring. `RLIX_MILES_UPDATE_TRANSPORT` env + (three modes, default cpu_serialize), v7 uniformity startup guard, + strict-broadcast startup rejection, auto classification, wiring at both + call sites; unit tests (E3, E10). Delivers: AC4, AC2. +- M4 — vast.ai e2e validation (v8): (a) dual-pipeline overlap smoke per + C9 (mem 0.8, mode `auto`) — mixed sessions; cpu_serialize legs + + existing dual pass-bar double as the AC2 regression confirmation; + (b') strict-mode e2e REJECTION check on the same overlap harness + (mode `broadcast` → classification fail-fast, log captured; fast run, + no training). Delivers: AC6, AC7, AC2 confirmation. + +## 6. Decision trace + +- D0 (Goal): unblock the NCCL broadcast transport for miles RLix-mode weight + sync so disaggregated/distributed topologies can sync weights without + tmpfs/same-node coupling. AC served: AC1, AC4, AC6. +- D1 (Architecture): keep the existing three-layer shape — service builds the + plan, `run_sync_session` is the single composite RPC (C2), SGLang engines + stay passive receivers. Sender NCCL is added strictly inside the miles + cache_owner; classification is added strictly inside the rlix coordinator. + No protocol keys added or removed from the plan dict (semantics of + `comm_ranks` refined per C8). AC served: AC1, AC2. Depends on: C1, C2. +- D2 (Module boundary): + - miles `megatron_utils/actor.py`: `_dispatch_nccl_broadcast` gains the + sender path; new private helper(s) only (e.g. `_sender_join_group`, + `_broadcast_bucket`). No change to `UpdateWeightFromDistributed`. + - miles `ray/rollout.py`: NOT touched (v6 — the previously planned + `get_engine_gpu_counts` manager method is dropped; rlix already holds + `rollout_num_gpus_per_engine` in miles_args, uniform per pipeline and + validated by the existing C7-engine gate, so per-engine GPU counts are + derived rlix-side with zero new miles surface). + - rlix `miles_model_update_service.py`: remove raise; rank-cursor + `comm_ranks` with uniform per-engine stride from the injected + per-engine GPU count; `world_size = 1 + per_engine × len(broadcast_set)`. + - rlix `miles_coordinator.py`: transport-mode resolution + (`RLIX_MILES_UPDATE_TRANSPORT`, default cpu_serialize) + + `_classify_broadcast_engines(target, mode) -> frozenset[int]` (strict + `broadcast` raises on any colocate target; `auto` splits; logged once + at startup) + wiring at both call sites; topology inputs threaded + through `register_model_update_resources` (train GPU set, infer + mapping, per-engine count). + - rlix `miles_pipeline.py`: pass the topology inputs at registration + (already knows `cluster_device_mappings` + miles_args). **Uniformity + startup guard (v7)**: before `register_model_update_resources` + (phaseB step5), validate the resolved SGLang config that + `_validate_topology` already receives — require every server group's + `num_gpus_per_engine` to equal `miles_args.rollout_num_gpus_per_engine` + (single supported RLix shape); any divergence (per-group TP override, + heterogeneous groups) → fail-fast with an error naming the offending + group. This makes the O6 exclusion enforceable and keeps the + args-derived rank cursor sound without any miles-side query. + - rlix `scripts/run_smoke_dual.sh` (v5): parameterize line 125 as + `--sglang-mem-fraction-static ${MILES_SMOKE_MEM_FRACTION:-0.30}` — + behavior-preserving default; M4 exports 0.8. No other harness change. + AC served: AC1, AC3, AC4. Depends on: C1, C2, A5. +- D3 (Runtime behavior): per-sync dynamic NCCL group, strict ordering to + avoid rendezvous deadlock (A6): (1) dispatch receiver + `setup_collective_group` refs without blocking; (2) sender + `init_process_group(rank=0, timeout=)` + — bounded, so a receiver that dies/rejects/hangs mid-rendezvous cannot + block rank 0 forever; on timeout the sender raises and its `finally` + tears down (codex v3, partial-rendezvous abort); (3) `ray.get` the setup + refs — a receiver setup failure surfaces here as an exception and takes + the same abort path; (4) per bucket: dispatch `broadcast_parameter` refs + → stage bucket tensors on GPU → `dist.broadcast(src=0)` in metadata + order (collectives bounded by the transport budget, C6 hierarchy) → + `ray.get` the bucket refs; (5) sender-owned teardown in the sender's own `finally` + (codex v3): dispatch receiver `destroy_collective_group` refs (tolerating + the existing 400→no-op guard) + sender `dist.destroy_process_group`, + executed on success, exception, AND timeout paths alike — the sender + actor, not the service, owns nested-ref lifecycle. Hard fallback when a + native call refuses to unwind even past its timeout: fail-fast per the + system contract — the pipeline dies and is re-registered (no in-place + recovery attempted). Existing pause/finalize/publish steps in the + service are unchanged. Both transport paths may coexist in one session + (cpu_serialize set and broadcast set are disjoint subsets of the + target). AC served: AC1, AC2, AC5. Depends on: A1, A2, A6, C6. +- D4 (Data/state shape): plan dict keys unchanged (`comm_ranks` becomes the + per-engine rank_offset cursor; `world_size` becomes 1 + per_engine × + |broadcast_set| — consistent with what `setup_collective_group` already + forwards to SGLang as `rank_offset`/`world_size`). Per-engine GPU count + comes from rlix's own `miles_args.rollout_num_gpus_per_engine` (uniform + per pipeline, C7-engine-validated) — no new miles surface (v6). + Transport-mode env `RLIX_MILES_UPDATE_TRANSPORT ∈ {cpu_serialize + (default) | broadcast | auto}` replaces the earlier separate kill-switch + env — one control, user-explicit, no silent mixing outside `auto` (v6). + **Relation to the existing miles `--model-update-transport` flag**: that + flag (choices `cuda_ipc`/`cpu_serialize`, gated by miles + `rlix_validation.py` C11 which forces `cpu_serialize` in RLix mode) + governs only the **colocate-leg** transport inside miles. The rlix env + is a separate rlix-side control; the C11 gate, the flag definition, and + the harness's `--model-update-transport cpu_serialize` argument all + remain untouched (C1). AC served: AC3, AC4. Depends on: C8, C5. +- D5 (Implementation details): + - Staging with memory preflight (codex v3): buckets are CPU tensors + (cache contract). Before staging each bucket, query + `torch.cuda.mem_get_info()`; if free < bucket_size + margin (margin + default 1 GiB, env-overridable), degrade to tensor-by-tensor staging — + `dist.broadcast` is per-tensor anyway, so whole-bucket staging is only + a batching optimization and degradation changes no wire behavior. Free + staged memory before the next bucket → peak GPU overhead ≈ 1 bucket + (healthy) or ≈ 1 tensor (degraded) (A3). + - Warmup (F25 comment in actor.py): after group create, broadcast a + 1-element sentinel tensor before bucket 0 — receivers ignore it? — + `unknown`: whether SGLang's route tolerates an unannounced warmup + collective; resolve during M1 against the reference path (the standalone + path does NOT warm up, so default is NO warmup unless e2e shows NCCL + lazy-init stalls; decision recorded in M1). + - dtype strings already normalized (`str(dtype).replace("torch.", "")`) on + the metadata side; sender broadcasts raw tensors — no cast. + - Cancellation ownership (codex v3, replaces the v1/v2 claim that + service `inflight_refs` suffices): the service's `inflight_refs` + + `_release_port_claim` cover only service-issued refs (top-level + `run_sync_session` ref, manager calls, port claim). The sender's + nested receiver refs are owned by the sender: tracked in a local list + inside `_dispatch_nccl_broadcast`, cancelled (`ray.cancel`) plus + group-destroyed in the sender's `finally` on any exit path. Port-claim + release follows the C6 v4 ownership rule: release after the sender ref + resolves (teardown acknowledged); on session-deadline expiry with a + broadcast session, leak-and-log, never delete. + - miles-side sender must NOT hold any lock beyond the existing + `_cache_lock` (already held for the whole transport in + `run_sync_session`). + AC served: AC1, AC5. Depends on: A2, A3, C6. +- D6 (Evidence): see §7. + +## 7. Evidence and stop conditions + +### Evidence required + +- E1 (M1, AC1/AC5): miles unit tests — sender dispatch ordering (setup refs + before sender join; bucket metadata refs before broadcast; destroy in + finally on success AND on mid-bucket exception), using mocked handles + + monkeypatched `dist`/`init_process_group`. Existing + `tests/test_miles_pipeline.py` style. +- E2 (M2, AC1/AC3): rlix unit tests — service no longer raises for non-empty + broadcast set; plan carries cursor comm_ranks + per-GPU world_size with + uniform per-engine stride (e.g. per_engine=2, broadcast engines {0,1} → + ranks {0:1, 1:3}, world_size 5; per_engine=1 degenerates to the dense + case); count injected from args, no manager RPC. +- E3 (M3, AC4/AC2): rlix unit tests — mode resolution: default/unset → + cpu_serialize (all engines, AC2 regression shape); `broadcast` + any + colocate target → raises at classification time; `broadcast` + all + disjoint → all broadcast; `auto` → overlap engine cpu_serialize, + disjoint engine broadcast (mix); invalid env value → fail-fast error. +- E4 (M4, AC6): dual overlap smoke (C9, mem 0.8) log evidence — effective + `--sglang-mem-fraction-static 0.8` echoed from the launch log (a 0.30 + run does not satisfy AC6); C9 + mixed-class startup assertion passes for both pipelines; for each + pipeline: ≥1 `sync_selected_workers_done ... broadcast=[...]` line with + non-empty broadcast set AND non-empty cpu_serialize set (mixed session); + SGLang `init_weights_update_group` success; ≥1 training step completes + after a broadcast sync; zero CUDA OOM events in all logs; EXIT_CODE=0. +- E5 (M4, AC2): same dual run — existing dual-smoke pass-bar conditions + (`scripts/grep_overlap_log.sh` 7-condition harness) unchanged and green; + cpu_serialize legs of every mixed session complete normally. +- E6 (M1+M2, AC5): unit tests — (a) normal/exception paths: port claim + released only AFTER the sender ref resolves (teardown ack ordering + asserted); miles sender's `finally` cancels its nested receiver refs and + destroys the group on both sides; (b) wedged-sender path (service + session deadline fires with sender ref unresolved, broadcast set + non-empty): claim is NOT deleted — leak-and-log asserted; + cpu_serialize-only session keeps release-on-timeout; (c) budget + hierarchy: sender internal budgets derived from `timeout_s` satisfy + "worst-case unwind < session deadline". +- E7 (M1, AC5): unit test — receiver-fails-during-setup (mocked + `setup_collective_group` raising / hanging past the patched timeout): + sender aborts with a bounded error (no indefinite block), runs teardown, + and surfaces the failure to the service; also covers a mid-bucket + receiver exception taking the same abort path. +- E8 (M1, AC1): unit test — memory-preflight degradation: with + `mem_get_info` patched to report low free memory, staging switches to + tensor-by-tensor and the broadcast sequence/metadata order is unchanged. +- E9 (M4, AC7, narrowed v8): strict-mode e2e rejection evidence — overlap + harness + `RLIX_MILES_UPDATE_TRANSPORT=broadcast` fails fast at the + first classification with the actionable colocate error (log + captured); the strict-mode POSITIVE matrix (all-disjoint, subset + targets) is covered at unit level by E3 (see O7 deferral). +- E10 (M3, AC4): unit tests for the v7 uniformity guard — resolved SGLang + config with a server-group `num_gpus_per_engine` override diverging + from `rollout_num_gpus_per_engine` → startup fail-fast naming the + offending group; uniform config passes; guard runs before + `register_model_update_resources`. + +### Stop conditions (agent must stop and ask) + +- S1: Any change would touch pristine upstream miles lines (C1 conflict). +- S2: E2E requires an SGLang server-side route change (receiver routes turn + out not to match A1) — surface findings first. +- S3: The warmup question (D5) cannot be resolved by observation on the vast + instance within one debugging session. +- S4: Any commit/push/PR action (C7 — needs explicit user instruction). + +## 8. Out of scope + +- O1: Multi-node engines (`nodes_per_engine > 1`) and cross-node tmpfs — MVP + is single-node validation (A4); the rank-cursor scheme is designed + multi-node-ready but not e2e-verified here. +- O2: Performance benchmarking cpu_serialize vs broadcast (issue asks for + capability, not perf numbers). +- O3: Changes to miles standalone (non-RLix) update paths + (`UpdateWeightFromDistributed`, `UpdateWeightP2P`, colocate IPC). +- O4: LoRA/multi-adapter sync over broadcast. +- O5: Dynamic re-classification mid-training beyond what the two existing + call sites already do per-sync. +- O6 (v6, enforcement added v7): heterogeneous per-engine GPU counts + (e.g. prefill TP2 / decode TP4 mixes) are out of scope — and, per codex + round-5, this exclusion is ENFORCED, not merely declared: a rlix-side + startup uniformity guard (see D2) fail-fasts any resolved SGLang config + whose server groups override `num_gpus_per_engine` away from + `rollout_num_gpus_per_engine`. If heterogeneous engines land later, the + documented extension is the miles manager read method + (`engine_gpu_counts` property already exists) — deliberately NOT added + now per minimal-change. +- O7 (v8): fully-disjoint train/infer topologies (and therefore the + all-broadcast positive e2e). miles `assert_rlix_topology` C1 requires + train ⊂ infer as the M11 partial-overlap contract; admitting disjoint + pools is a scheduler/validation design change, not a transport change. + When that lands, run the deferred disjoint dual smoke with + `RLIX_MILES_UPDATE_TRANSPORT=broadcast` as the positive AC7 e2e. + +## 9. Risks + +- R1 (A6): rendezvous deadlock if any `ray.get` lands between receiver + dispatch and sender join → covered by E1 ordering tests. Partial-receiver + failure during rendezvous (the non-happy-path variant, codex v3) → + bounded sender timeout + abort semantics (D3), covered by E7. +- R2 (A6, C8): rank misassignment for TP>1 → covered by E2; e2e M4 runs TP=1 + so unit coverage is the only TP>1 gate (accepted for MVP). +- R3 (A3): GPU OOM while staging a bucket on a busy train GPU. C9 tightens + this risk: mem limit 0.8 on 16 GB GPUs leaves ≈3 GB nominal headroom, + and fragmentation/NCCL buffers/copy temporaries erode it further (codex + v3). Mitigation is now designed-in, not deferred: memory preflight + + tensor-by-tensor degradation (D5, E8) + zero-OOM as an AC6 hard + condition (E4). Defense-in-depth note (v5): miles already ships an S2 + startup gate (`rlix_validation.py`) rejecting configs where bucket_size + + transport scratch ≥ estimated post-wake free VRAM when non-colocate + engines exist — S2 catches misconfiguration at startup; the D5 + preflight catches runtime erosion. Both stay. +- R4: leaked NCCL group after a failed session poisons the next session's + group_name — group names are per-sync (`miles_{pipeline}_{sync_id}`), so a + leak wastes resources but cannot collide; sender-owned + destroy-in-finally + bounded timeouts (D3) + E6/E7 bound the exposure. + Port-claim rule (v4, replaces the v3 note codex flagged): release only + after sender teardown ack; on the wedged-sender path the claim is + leaked-and-logged (existing ROLL-backend precedent: leak is safer than + collision under a possibly-live TCP store). Residual cost of a leak: + one stale SharedStorage key until job restart — accepted. +- R5: `flush_cache`/finalize interaction with paused engines differs under + broadcast — finalize flow is transport-agnostic (unchanged code path); + watched in M4 logs. + +## 10. M4 evidence record (2026-08-16, appended post-run) + +Instance: vast 47858280, 4x A100-SXM4-80GB, fork-baseline image (torch +2.9.1+cu129 / sglang 0.5.10 / ray 2.54.1). Logs on instance: +`/root/smoke_b_reject.log`, `/root/smoke_a_auto08.log`, +`/root/smoke_a_final.log`. + +- E9 (run b'): `RLIX_MILES_UPDATE_TRANSPORT=broadcast` on the overlap + harness → registration logged `transport mode=broadcast ... + train_gpus=[0] infer_gpus=[0,1,2]`; first classification raised the + actionable error ("cannot serve colocate engines [0] ... Use 'auto'"); + SMOKE_EXIT=1. PASS. +- E4 (run a, final): `auto` + `MILES_SMOKE_MEM_FRACTION=0.8` + + `RLIX_ASSERT_MIXED_TRANSPORT=1` → C9 assertion logged+passed for both + pipelines (P1 broadcast=[1,2]/cpu_serialize=[0]; P2 + broadcast=[0,1]/cpu_serialize=[2]); `mem_fraction_static=0.8` + confirmed in engine logs; zero CUDA OOM; SMOKE_EXIT=0. Route-level + mixed-transport attribution (first run): `init_weights_update_group` + 200 ×10, `update_weights_from_distributed` 200 ×11, + `update_weights_from_cpu_bucket` 200 ×6, + `destroy_weights_update_group` 200 ×10; per-pid GPU attribution + matches the predicted classification; P2 same-session mix confirmed by + timestamp adjacency (cpu_bucket 11:39:36-38 → distributed 11:39:39). + PASS. +- E5: `grep_overlap_log.sh` 7-condition harness → RESULT: PASS on both + the first auto run and the final assertion-enabled run. +- Unit evidence: miles 16/16 (E1/E7/E8 + env guard); rlix 32/32 + (E2/E3/E6/E10 + C9 assertion). Regression: rlix suite 91 passed (1 + pre-existing ray-2.54 API failure, verified on pristine base); miles + test_miles_pipeline.py 6/6. +- Codex implementation reviews: r1 (1 high NCCL bound + 1 medium claim + window) fixed; r2 (1 high env setdefault) fixed; r3 approve; r4 + (assertion delta) approve. + +## 11. Post-delivery user-directed expansion: O7 partially delivered (2026-08-16, appended) + +User directive ("我要强制跑nccl" + "不管用户怎么设置gpu拓扑，你都应该支持的吧") +unlocked the dedicated-train (fully-disjoint) topology ahead of the O7 +deferral, expanding the miles diff beyond the original C1 boundary +(actor.py + tests) into `examples/rlix/run_miles_dual.py` — still +RLix-port-added code, per the user's standing minimize-upstream rule. + +Changes (codex delta rounds 5-6, r6: approve): +- `run_miles_dual.py::_overlap_pools_from_env` + `_build_pipeline`: the + per-pipeline invariant is now two-family — subset (overlap, tested) or + fully-disjoint (dedicated train cards, experimental warning); partial + intersections stay rejected with an actionable error (scheduler + shrink/grant accounting unverified for them). +- rlix `miles_pipeline._wait_for_overlap_engines_offloaded`: filters + granted train GPUs to the infer mapping before deriving engine indices + (disjoint grants previously produced nonsense indices like -1, masked + by a broad exception catch). +- Correction to the v8 note: miles C1 (F10) checks the LOCAL count-derived + shape, not physical placement — dual-mode physical mappings were only + blocked by the driver's own asserts, which is why this unlock needed no + `rlix_validation.py` change. The v8 claim "unreachable without relaxing + the C1 gate" was wrong at the physical level. +- Status: strict `broadcast` all-NCCL on the disjoint dual topology + (P1 train [0] / infer [1,2]; P2 train [3] / infer [1,2]) is + user-verification-pending (user runs it themselves); unit + codex + gates green. + +## 12. Disjoint-topology debug: concurrent-sync port collision (2026-08-16, appended) + +First user-run of the dedicated-train topology crashed: sender CUDA +"invalid argument" + receivers losing the rendezvous TCP store mid +ncclUniqueId exchange. Root cause chain (three stacked facts): +miles `RayActor.get_free_port` scans deterministically from 20000, so +both pipelines' cache_owners pick the SAME port; the SharedStorage +port-claim protocol was silently skipped in miles mode (ROLL's actor +does not exist there — the F26/C16 protection was never active); the +overlap topology masked this by scheduling serialization, while +dedicated-train pipelines finish training steps simultaneously and sync +CONCURRENTLY — two NCCL TCP stores cross-wired on port 20000. + +Fix (codex delta rounds 7-8, r8: approve): +- New `_PortClaimStore` detached Ray actor ('rlix:miles_port_claims', + RLIX_NAMESPACE) — rlix-owned fallback backend for the claim protocol; + fallback triggers ONLY on ImportError (no roll) or ray.get_actor + ValueError (actor absent); any other lookup failure fails the sync + fast (no silent skip, no split claim namespace). +- Claim tuple carries the accepting store handle; all three release + paths (awaited, nowait, leak-path unpack) use the carried handle. +- Collision retry now re-picks with get_free_port(start_port=port+1) — + the contested port is claim-reserved but not OS-bound, so the old + plain rescan returned the same port until the retry budget exhausted. +- Evidence bar for the user's rerun: two distinct master_port values in + the log + zero cpu_bucket routes + EXIT 0. rlix suite 35/35. + +## 13. O7 continued: dual-disjoint scheduler gap + single-disjoint unlock (2026-08-16, appended) + +Second user-run finding (after §12's port fix and the tms hook lesson): +dual + disjoint + SHARED infer cards deadlocks at P2 init. Chain: with +dedicated train cards nothing ever shrinks P1's activated engines off the +shared GPUs (in overlap topologies the pipeline's own train reclaims +them, driving the shrink/planned-release rotation); the driver stages +init before training loops, so P1 emits no progress signals; P2's +init-time GENERATION request (lowest priority, progress-driven rotation) +waits forever. **Known limitation, NOT fixed here**: admitting +idle-GEN-yields-to-pending-init is scheduler-regime design work (O7 +follow-up). On 4 GPUs: dual ⇒ use overlap+auto (validated); all-NCCL ⇒ +single-pipeline disjoint. + +Also learned (user runs): (a) tms hook `preload` breaks sender staging +while the train actor is offloaded (stale cached allocator pointers → +cudaErrorInvalidValue at tensor.to); `torch` hook is the validated mode +for broadcast + offload_train — debug_pipeline's "auto" resolves to +preload on A100 and must be overridden. (b) On shared infer cards under +DISJOINT (no time-sharing), co-resident engine pools must divide the +card: mem_fraction ≤ ~0.9/N (dual → 0.4); the 0.8 mandate applies to +overlap (time-shared) topologies where ≤1 pool is awake per card. + +Delivered (codex delta rounds 9-10, r10: approve): +`run_miles_rlix.py::_build_cluster_device_mappings` gains the +`MILES_SINGLE_TRAIN_GPUS` / `MILES_SINGLE_INFER_GPUS` explicit-mapping +override (two-family invariant; env lengths must equal the CLI-derived +actor/rollout counts, fail-fast otherwise) + 9-case test matrix in +miles tests/test_single_mapping_override.py. Minimal all-NCCL topology +on 4 GPUs: single pipeline, train [0] / infer [1,2], mem 0.8, mode +broadcast — user-verification pending. + +## 14. Sync-under-load flush-timeout fix (2026-08-16, appended; codex r11-r16, r16: approve) + +Field failure (user's 20-rollout dual auto run): after_training → +sync_base_weights_to_active → finalize /flush_cache TimeoutError. Root +cause: under the fully-async rollout, generation for the next step is +already live on the target engines and the router keeps dispatching +during the pause/flush window — the queue never drains (2-rollout smokes +had empty queues). Transport-agnostic (cpu path shares finalize). + +Fix — the sync-under-load bracket (rlix miles_coordinator +sync_base_weights_to_active) hardened over six codex rounds: +- quiesce BEFORE the sync, mirroring shrink_engines: per-engine + unregister_from_router (admission close) + manager _abort_engines + (router re-dispatches aborted requests with the NEW weights); +- quiesce lives INSIDE the try so all failures reach cleanup (r11); +- re-admission via the NEW manager-side atomic + RolloutManager.register_router_if_active (state check + idempotent + /add_worker in ONE serialized manager call; manager is + max_concurrency-1, so it cannot interleave with shrink_engines) — + NOT activate_routing (loading-only INIT transition, raises for active + engines, r11) and NOT a read-then-register pair (TOCTOU, r13); the + method filters manually because _resolve_engine_indices raises for + the exact concurrently-shrunk engines it must skip (found via the + r15 contract test); +- concurrently-shrunk engines are never re-admitted (intent ∩ live + state, r12) and are warning-logged; +- a re-registration FAILURE after a successful sync escalates to + RuntimeError (never report healthy-but-unroutable, r14); a failed + sync keeps its original exception as primary; +- registration-time capability probe (empty-list call) fails loud on + new-rlix/old-miles version skew before any sync quiesces engines + (r15); +- abort-idempotency cache reset is an independent best-effort finally + step so future genuine shrinks re-abort. + +miles diff grows by the one manager method (rollout.py RLix section) + +contract test. Tests: rlix AST ordering + behavioral concurrency/ +escalation (105 total green), miles contract 4 cases (29 total green). +User e2e verification of the 20-rollout dual run pending. + +## 15. Final e2e verification — 20-rollout dual overlap run (2026-08-16, user-run) + +Config: debug_pipeline.py, mode=dual, transport=auto, tms hook torch, +mem_fraction 0.8, num_rollout=20, topology P1 [0]/[0,1,2] + P2 +[3]/[1,2,3], 4x A100-SXM4-80GB. Evidence panel from /root/logs/run.log: + +- NCCL group creations: 45 (`backend=nccl`, `world_size=3`, per-sync + group names `miles_{pipeline}_{sync_id}`). +- Mixed transport across the full run: `update_weights_from_distributed` + 200 x48 (NCCL legs) + `update_weights_from_cpu_bucket` 200 x46 + (colocate legs) + `destroy_weights_update_group` 200 x44. +- Port anti-collision IN ACTION: master_port=20000 x27 AND + master_port=20001 x18 — concurrent pipeline syncs claimed distinct + rendezvous ports via the _PortClaimStore fallback (§12 fix verified + under load). +- Zero CUDA OOM, zero "Timeout while flushing cache" (§14 bracket + verified under 20-rollout fully-async load), zero colocate + rejections; 36 /generate requests served across cycles. + +This is the definitive under-load validation on top of §10's 2-rollout +smokes: every fix from §12-§14 exercised in one run. diff --git a/plans/miles-nccl-broadcast-plan.tldr.md b/plans/miles-nccl-broadcast-plan.tldr.md new file mode 100644 index 0000000..602cd40 --- /dev/null +++ b/plans/miles-nccl-broadcast-plan.tldr.md @@ -0,0 +1,301 @@ +# TLDR — Enable NCCL broadcast path for miles weight update (rlops/rlix#42) + +Source plan: `plans/miles-nccl-broadcast-plan.md` (DRAFT v8, 2026-08-16 — implementation in progress) +TLDR mode: `complex` (multi-repo, multi-actor runtime topology, 4 milestones) + +## 0. Audit Dashboard + +- **Goal**: unblock the NCCL broadcast transport for miles RLix-mode weight sync (issue #42), as a **user-selected mode** — `RLIX_MILES_UPDATE_TRANSPORT ∈ {cpu_serialize (default) | broadcast | auto}`; only `auto` ever mixes transports. +- **Blast radius (v6, minimal-change)**: miles = `actor.py` + tests ONLY; rlix = `miles_model_update_service.py` + `miles_coordinator.py` + `miles_pipeline.py` + one env-parameterization line in `run_smoke_dual.sh`. +- **Protocol change**: none on the wire — plan-dict keys unchanged; `comm_ranks`/`world_size` semantics refined per-GPU (C8); per-engine GPU count derived rlix-side from args (no new miles surface). +- **Highest-risk areas**: NCCL rendezvous failure semantics (R1/E7), deadline hierarchy + port-claim ownership (C6), rank accounting (C8/R2), staging OOM under mem 0.8 (R3/E8). +- **Review state**: codex rounds 1-5 (r4: approve on v5; r5 on v6: 1 high — uniformity assumed-not-enforced, fixed in v7 by a rlix-side startup uniformity guard) — pending codex round 6. +- **Milestones**: M1 miles sender → M2 rlix service → M3 rlix mode wiring → M4 e2e (overlap+`auto`+mem 0.8, plus strict-mode rejection check; all-broadcast positive e2e O7-deferred — miles C1 gate requires train ⊂ infer, v8). +- **Open `unknown`s**: 1 — warmup collective after group create (D5). See Appendix G. +- **Stop conditions**: 4 (S1-S4) — see Appendix E. +- **Artifact integrity**: AC-grid PASS (check_tldr_integrity.py exit 0); Mermaid PASS (validate_mermaid.sh exit 0). + +## 1. Context + +**Current**: RLix-mode weight sync chain (`MilesCoordinator` → `MilesModelUpdateService` → cache_owner `run_sync_session` → SGLang engines) fully works for Path A (cpu_serialize tmpfs+HTTP). Path B (NCCL broadcast) is scaffolded on both sides — plan fields, receiver-side fan-out code, SGLang receiver methods all exist — but is dead code behind two deliberate fail-fast guards (rlix `miles_model_update_service.py:199`, miles `actor.py:968`). + +**Gap**: sender-side NCCL on the cache_owner (join dynamic group as rank 0, per-bucket `dist.broadcast`, teardown) was never implemented; without it receivers would hang in SGLang `/init_weights_update_group`. No rlix call site computes `broadcast_local_ranks`; rank accounting is per-engine (wrong for TP>1). + +**v6 design stance (user-directed)**: transport is explicit user choice, not silent auto-classification. `broadcast` mode never mixes — it fail-fast-rejects topologies it cannot serve. The physics forcing a mixing mode to exist at all: NCCL cannot form a group containing the same physical GPU twice, so an engine sharing the cache_owner's GPU can never receive via broadcast — overlap topologies therefore require `auto` (mix) or `cpu_serialize`. + +**Proven reference**: miles standalone `UpdateWeightFromDistributed` (pristine upstream — pattern reuse only, no modification). + +## 2. Assumptions + +| ID | Assumption | +|----|-----------| +| A1 | SGLang receiver routes behave per the standalone path: `/init_weights_update_group` registers `tp_size` consecutive NCCL ranks from `rank_offset`; `/update_weights_from_distributed` blocks until tensors arrive from rank 0 | +| A2 | `bucket.params` dict is insertion-ordered → metadata lists and sender broadcast order agree by construction | +| A3 | cache_owner has a live CUDA context; staging memory is preflight-checked per bucket with tensor-by-tensor degradation; miles S2 startup gate provides the config-level check | +| A4 | MVP validation is single-node, `nodes_per_engine == 1`; multi-node engine shards out of scope | +| A5 | `cluster_device_mappings` + `rollout_num_gpus_per_engine` suffice to derive engine ↔ physical-GPU sets (precedent: `MilesPipeline._wait_for_overlap_engines_offloaded`); the C9 startup assertion verifies derived classification against actual harness envs AND actual scheduler grant | +| A6 | Reviewer-history risk focus: NCCL rendezvous deadlock + global-vs-local rank confusion (PR rlops/RL#3 review classes) | + +## 3. Scope & Constraints + +### 3.1 Out of Scope + +| ID | Excluded | Note | +|----|----------|------| +| O1 | Multi-node engines / cross-node validation | rank-cursor designed multi-node-ready, not e2e-verified here | +| O2 | Performance benchmarking cpu_serialize vs broadcast | issue asks capability, not perf | +| O3 | miles standalone update paths (`UpdateWeightFromDistributed`, `UpdateWeightP2P`, colocate IPC) | pristine upstream (C1) | +| O4 | LoRA / multi-adapter sync over broadcast | — | +| O5 | Mid-training re-classification beyond existing per-sync call sites | — | +| O6 | Heterogeneous per-engine GPU counts (prefill/decode TP mixes) | ENFORCED out of scope (v7): rlix startup uniformity guard fail-fasts server-group `num_gpus_per_engine` overrides diverging from `rollout_num_gpus_per_engine` (E10); future extension = miles `engine_gpu_counts` manager method | +| O7 | Fully-disjoint train/infer topologies + all-broadcast positive e2e (v8) | miles `assert_rlix_topology` C1 requires train ⊂ infer (M11 partial-overlap contract); admitting disjoint pools is a scheduler/validation design change, not a transport change | + +### 3.2 Hard Constraints + +| ID | Constraint | +|----|-----------| +| C1 | v6 tightened: miles changes = `actor.py` (`_dispatch_nccl_broadcast` + in-method helpers, all RLix-port-added) + tests, NOTHING else — no `rollout.py`, `arguments.py`, `rlix_validation.py` (C11/S2 untouched), `sglang_engine.py`, or any pristine upstream line | +| C2 | F04: single composite `run_sync_session` RPC; no new top-level Ray transport methods on train actor | +| C3 | F26/C16: `master_port != 0`, claimed via SharedStorage (existing, keep) | +| C4 | F21: exactly one `set_weight_version` publish per sync, service-only (unchanged) | +| C5 | Fail-fast + v6 no-silent-mode-switching: user-selected transport mode honored exactly or run refuses to start (strict `broadcast` + colocate target = startup rejection, never a quiet downgrade to mixing) | +| C6 | Enforcement split + nested deadline hierarchy (v8-r1): session deadline = `timeout_s` outermost; rendezvous budget (0.15×) doubles as the NCCL pg timeout bounding EACH collective (watchdog); cumulative monotonic deadline (rendezvous+transport=0.65×) checked between collectives; receiver acks get the remaining window; worst-case unwind = 0.9× session. Port claim: release only after sender teardown ack; sender-resolution tracked separately so post-resolution cancellation releases (no false leak); truly-wedged sender + broadcast → leak-and-log; cpu_serialize-only keeps release-on-timeout | +| C7 | No commit/push without explicit user instruction; codex approval before sign-off | +| C8 | NCCL ranks are per-GPU not per-engine: `world_size = 1 + per_engine × len(broadcast_set)`; rank_offset cursor per engine (uniform stride, O6) | +| C9 | M4 run (a) = dual overlap topology from harness defaults (`run_smoke_dual.sh:68-71`): P1 train [0] / infer [0,1,2]; P2 train [3] / infer [1,2,3]; mem fraction via `${MILES_SMOKE_MEM_FRACTION:-0.30}` parameterization, exported `0.8`, effective value echoed (0.30 run ≠ AC6); mode `auto`; startup assertion (env mappings + scheduler grant logged, ≥1 engine of each class per pipeline) — failure/divergence = smoke INVALID | + +### 3.3 Acceptance Criteria + +| ID | Statement | Derives from | Verified by | Milestone | +|----|-----------|--------------|-------------|-----------| +| AC1 | miles `run_sync_session` executes a plan with non-empty `broadcast_local_ranks` end-to-end (sender join + per-bucket broadcast with memory preflight + teardown), no `NotImplementedError` | Goal, A1, A2 | E1, E8 | M1 | +| AC2 | All-cpu_serialize behavior unchanged — default mode IS cpu_serialize (zero change until opt-in); existing unit tests pass unmodified; dual-run cpu_serialize legs + existing dual-smoke pass-bar green | Goal, C5 | E3, E5 | M4 | +| AC3 | Rank/world_size accounting correct for TP>1 (per-GPU cursor, uniform stride from args), unit-verified | C8, A1 | E2 | M2 | +| AC4 | Transport is user-selected via `RLIX_MILES_UPDATE_TRANSPORT` (cpu_serialize default / broadcast strict / auto) with startup-validated topology inputs: the v7 uniformity guard rejects resolved SGLang configs diverging from `rollout_num_gpus_per_engine` before classification; strict mode fail-fast-rejects colocate targets (NCCL duplicate-GPU constraint); `auto` is the only mixing mode; mode logged; invalid value errors | Goal, A5, C5 | E3, E10 | M3 | +| AC5 | Abnormal session end (receiver dies/rejects/hangs, mid-bucket exception, sender budget expiry) unwinds inside the C6 hierarchy with sender-owned teardown; port claim released strictly after teardown ack; wedged-sender residual path leaks-and-logs the claim | C6, C3 | E6, E7 | M2 | +| AC6 | Overlap e2e (C9, mode `auto`, mem 0.8 echoed): mixed-transport sessions per pipeline, startup assertion passes, zero GPU OOM, existing dual pass-bar met, EXIT_CODE=0 | Goal, A4, A5, C9 | E4 | M4 | +| AC7 | (v8 narrowed) Strict `broadcast` honors no-mix: E3 unit matrix (rejection naming colocate engines; all-disjoint + colocate-avoiding subset targets pass) + e2e rejection check on the overlap harness (mode=broadcast → classification fail-fast, log captured); all-broadcast positive e2e O7-deferred (miles C1 gate: train ⊂ infer) | Goal, A4, A5, C5 | E9 | M4 | + +### 3.4 Milestones + +| ID | Scope | Exit criteria | Delivers AC | Depends on | +|----|-------|---------------|-------------|------------| +| M1 | miles sender-side NCCL in `actor.py`: budget-bounded join + per-bucket broadcast with memory preflight + sender-owned teardown + guard removal | E1, E7, E8 green | AC1 | — | +| M2 | rlix service unlock + per-GPU rank cursor (uniform stride from args, no manager RPC) + C6 claim/deadline rules + unit tests | E2, E6 green | AC3, AC5 | M1 | +| M3 | rlix transport-mode wiring: env resolution (3 modes, default cpu_serialize), v7 uniformity startup guard, strict-mode rejection, auto classification, both call sites + unit tests | E3, E10 green | AC4 | M2 | +| M4 | vast.ai e2e (v8): (a) overlap + `auto` + mem 0.8; (b') strict-mode rejection check on the same overlap harness (fast, no training) | E4, E5, E9 green | AC2, AC6, AC7 | M3 | + +### 3.5 Strategy notes + +- v6 control model: ONE rlix-side env selects transport; default `cpu_serialize` = today's behavior, making the merged change a no-op until opt-in (minimal-change). The earlier separate kill-switch env is subsumed. +- Why a mixing mode exists at all: NCCL cannot put the same physical GPU twice in one group, so sender-colocated engines can never receive via broadcast; `broadcast` mode therefore rejects such topologies (C5), and `auto` exists for overlap topologies like the mandated M4 run. +- Rejected alternative: extending miles `--model-update-transport` with a broadcast choice — touches miles `arguments.py` + C11 gate for no functional gain over the rlix env (C1 minimal-change). +- Rejected alternative: modifying/reusing `UpdateWeightFromDistributed` — violates C1; pattern copied, code not. +- Failure-handling stance: bounded budgets + abort + fail-fast pipeline death; leak-don't-release port claim under a possibly-live TCP store. + +## 4. Critical Views + +### 4.1 Architecture Integration View + +*Audit question: where exactly does new code land, and does anything cross the C1/C2 boundaries?* + +```mermaid +flowchart TD + subgraph RLIX [rlix repo] + MP[MilesPipeline MC[MilesCoordinator SVC[MilesModelUpdateService RSS + RSS --> DNB + DNB -- "NCCL broadcast rank 0" --> ENG + DNB -- "Ray RPC setup and metadata" --> ENG +``` + +### 4.2 Runtime / Data Path View + +*Audit question: do the budgets nest so the sender always unwinds before the service deadline, and does teardown precede claim release on every path (C6, AC5)?* + +```mermaid +sequenceDiagram + participant S as Service rlix + participant O as cache_owner miles + participant E as SGLang engines + S->>O: run_sync_session with plan + Note over S: session deadline outermost equals timeout_s + O->>E: setup_collective_group refs dispatched async + Note over O: sender joins as rank 0 within rendezvous budget + O->>O: init_process_group tcp rendezvous bounded + O->>E: ray.get setup refs — receiver failure surfaces here + loop per bucket within transport budget + O->>E: broadcast_parameter refs dispatched async + O->>O: memory preflight then stage bucket or per tensor + O->>E: dist.broadcast each tensor bounded + O->>E: ray.get bucket refs + end + Note over O,E: sender finally within teardown grace — cancel nested refs, destroy both sides + O-->>S: sender ref resolves success or bounded error + S->>S: release port claim only after sender ref resolved + Note over S: wedged sender path — deadline fires, claim leaked and logged, never deleted + S->>E: pause then finalize then continue unchanged + S->>S: publish version via manager unchanged +``` + +### 4.3 Physical Topology View (M4 runs) + +*Audit question: does each M4 run exercise exactly the transport behavior its mode promises?* + +Run (a) — overlap, mode `auto`, mem 0.8 (harness env defaults, unchanged): + +| Pipeline | Train grant | Infer pool (engine@gpu) | cpu_serialize | broadcast | +|----------|------------|------------------------|---------------|-----------| +| P1 | [0] | e0@0, e1@1, e2@2 | e0 | e1, e2 | +| P2 | [3] | e0@1, e1@2, e2@3 | e2 | e0, e1 | + +Run (b) — disjoint pools, mode strict `broadcast`, default mem (env override): + +| Pipeline | Train grant | Infer pool | Expected | +|----------|------------|-----------|----------| +| P1 | [0] | [1] | all broadcast, zero mixing | +| P2 | [2] | [3] | all broadcast, zero mixing | + +Plus rejection check: overlap envs + `broadcast` mode → startup classification error (E9). Startup assertion (run a) logs env mappings + scheduler grant; failure/divergence = smoke INVALID. + +## 5. Decision Map + +*Audit question: does every mechanism trace to a requirement, and are the dependencies acyclic?* + +```mermaid +flowchart LR + D0[D0 Goal] --> D1[D1 Architecture] + D1 --> D2[D2 Module boundary] + D1 --> D3[D3 Runtime behavior] + D2 --> D4[D4 Data shape] + D3 --> D5[D5 Impl details] + D4 --> D5 + D5 --> D6[D6 Evidence] +``` + +| ID | Decision | Depends on | AC served | Notes | +|----|----------|------------|-----------|-------| +| D0 | Unblock NCCL broadcast as a user-selected transport mode | Goal | AC1, AC4, AC6, AC7 | — | +| D1 | Keep 3-layer shape; sender NCCL inside cache_owner only; mode resolution + classification inside rlix coordinator only; no wire-protocol key changes | C1, C2 | AC1, AC2 | `comm_ranks` semantics refined (C8), keys unchanged | +| D2 | v6 module boundary: miles `actor.py` + tests ONLY; rlix service/coordinator/pipeline + one harness line; per-engine GPU count from rlix args (dropped manager method) | C1, C2, A5 | AC1, AC3, AC4 | full file list in Appendix B | +| D3 | Per-sync dynamic group; ordering dispatch→budget-bounded-join→get; per-bucket metadata-then-broadcast within transport budget; sender-owned finally teardown; partial-rendezvous abort; fail-fast pipeline death fallback | A1, A2, A6, C6 | AC1, AC2, AC5 | budgets from the C6 hierarchy | +| D4 | Plan keys unchanged; cursor comm_ranks + `world_size = 1 + per_engine × len(broadcast_set)` (uniform stride from args); mode env `RLIX_MILES_UPDATE_TRANSPORT` (3 values, default cpu_serialize) replaces kill-switch; miles `--model-update-transport` flag + C11 gate untouched | C8, C5 | AC3, AC4 | wire-compatible with `setup_collective_group` | +| D5 | Memory preflight per bucket with tensor-by-tensor degradation (S2 gate = config-level companion); no dtype cast; nested receiver refs sender-owned; port-claim per C6 ownership rule; no new locks; warmup = `unknown` (default NO, decide in M1) | A2, A3, C6 | AC1, AC5 | see Appendix G | +| D6 | Evidence: unit-first (ordering, ranks, mode resolution, claim rules, receiver-failure, preflight), e2e last (two-mode M4 matrix) | — | AC2, AC3, AC5, AC6, AC7 | see §6.1 | + +## 6. Evidence & Stop Conditions + +#### 6.1 Evidence Required + +| ID | Evidence | AC verified | Milestone | Method | +|----|----------|-------------|-----------|--------| +| E1 | miles unit tests: dispatch ordering (setup refs before sender join; metadata refs before broadcast; teardown in finally on success and mid-bucket exception), mocked handles + patched dist | AC1 | M1 | pytest, `test_miles_pipeline.py` style | +| E2 | rlix unit tests: no raise on non-empty broadcast set; cursor comm_ranks + per-GPU world_size with uniform stride (per_engine=2, engines {0,1} → ranks {0:1, 1:3}, ws 5; per_engine=1 degenerates to dense); count injected from args, no manager RPC | AC3 | M2 | pytest | +| E3 | rlix unit tests: mode resolution — default→all cpu_serialize (AC2 shape); `broadcast`+colocate target→classification-time raise; `broadcast`+disjoint→all broadcast; `auto`→mix; invalid env→fail-fast | AC2, AC4 | M3 | pytest | +| E4 | overlap smoke (C9, `auto`, mem 0.8) logs — effective 0.8 echoed (0.30 = fail); startup assertion passes both pipelines; per pipeline ≥1 mixed `sync_selected_workers_done` line; SGLang `init_weights_update_group` success; ≥1 training step after broadcast sync; zero CUDA OOM; EXIT_CODE=0 | AC6 | M4 | vast.ai dual run (a) | +| E5 | same run (a): existing 7-condition dual-smoke pass-bar unchanged and green; cpu_serialize legs of every mixed session complete normally | AC2 | M4 | vast.ai dual run (a) | +| E6 | unit tests: (a) claim released only AFTER sender ref resolves; sender finally cancels nested refs + destroys both sides; (b) wedged-sender → claim NOT deleted, leak-and-log; cpu_serialize-only keeps release-on-timeout; (c) budgets satisfy worst-case-unwind < session deadline | AC5 | M2 | pytest both repos | +| E7 | unit test: receiver-fails-during-setup (raise / hang past patched budget) — sender aborts bounded, runs teardown, surfaces failure; mid-bucket receiver exception same path | AC5 | M1 | pytest miles | +| E8 | unit test: memory-preflight degradation — low `mem_get_info` → tensor-by-tensor staging, broadcast sequence/metadata order unchanged | AC1 | M1 | pytest miles | +| E9 | (v8 narrowed) strict-mode e2e rejection: overlap harness + `broadcast` mode fails fast at first classification with the actionable colocate error (log captured); positive strict matrix covered at unit level by E3 | AC7 | M4 | vast.ai run (b') | +| E10 | unit tests: v7 uniformity guard — server-group `num_gpus_per_engine` override diverging from `rollout_num_gpus_per_engine` → startup fail-fast naming the offending group; uniform config passes; guard runs before `register_model_update_resources` | AC4 | M3 | pytest rlix | + +### 6.2 Stop conditions + +Mirrored in Appendix E (S1-S4): pristine-upstream conflict, SGLang route mismatch, unresolvable warmup question, any commit/push action. + +## 7. Audit Checkpoints + +- [ ] **CHK1** — AC1 is delivered by M1 but its rlix half only unlocks in M2: confirm the M1 exit gate (miles-only, unit-level) is acceptable as "AC1 delivered", or move AC1 to M2. *(auditor decision)* +- [ ] **CHK2** — D5 warmup question is the only `unknown`; plan defaults to NO warmup (standalone path has none). Confirm default. See Appendix G. +- [ ] **CHK3** — R2 accepts unit-only coverage for TP>1 rank accounting (M4 e2e runs TP=1). Confirm acceptable for MVP. +- [ ] **CHK4** — C9's startup assertion covers t=0; scheduler grant drift mid-run would change per-sync classification silently. Confirm t=0 assertion + per-sync mixed-session log evidence (E4) suffices, or require a per-sync assertion (cheap — G4). +- [ ] **CHK5** — `register_model_update_resources` gains topology args; confirm no other callers exist. *(source plan asserts via D2; verify at implementation)* +- [ ] **CHK6** — C6's budget fractions are impl detail with the nesting invariant as the requirement; confirm the invariant formulation is the auditable contract you want. +- [ ] **CHK7** — mem-fraction resolved by one-line harness env parameterization (default 0.30 preserved). Confirm vs forking an M4 wrapper. +- [ ] **CHK8 (v6, revised v7)** — uniformity is now guard-enforced at startup (E10), closing codex round-5's high finding (C7 gates alone do NOT check resolved server-group TP). Residual: the guard reads the resolved SGLang config rlix already holds — confirm at implementation that this config object reflects post-override group values (if not, S1-style stop and reconsider the miles-side query). + +See Appendix G for concrete source-plan patch suggestions. + +## Appendix A: Full Decision Trace + +- D0 → D1: goal + user's explicit-mode directive constrain architecture to existing scaffold with a single rlix-side control surface. +- D1 → D2: C1 (v6 tightened) forces every miles edit into `actor.py`; rlix owns mode + classification because only rlix has `cluster_device_mappings` (A5) and the args-derived per-engine count; harness line is rlix-owned scripting. +- D1 → D3: single-composite-RPC (C2) + C6 enforcement split dictate sender NCCL inside `run_sync_session` under `_cache_lock`, every blocking call budget-bounded, teardown sender-`finally`-scoped. +- D2 → D4: args-derived uniform GPU count (O6) satisfies C8 with zero new miles surface; the mode env replaces the kill-switch as the single user control. +- D3 → D5: deadlock/abort semantics (A6, E7) + invariant-checked staging (A3, E8, S2 companion) + claim-release ordering (C6, E6) are the nontrivial mechanics. +- D4/D5 → D6: unit tests target the two reviewer bug classes + mode resolution + claim/deadline rules; e2e is a two-run matrix proving `auto` mixes correctly under the mandated overlap+0.8 run and `broadcast` never mixes (incl. rejection path). + +## Appendix B: Full Module / File Boundary + +| Repo | File | Change | +|------|------|--------| +| miles | `miles/backends/megatron_utils/actor.py` | `_dispatch_nccl_broadcast`: remove guard; budget-bounded sender join; per-bucket broadcast with memory preflight; sender-owned finally teardown; in-method helpers only | +| miles | `tests/test_miles_pipeline.py` (or sibling) | E1/E7/E8 unit tests | +| miles | everything else (`rollout.py`, `arguments.py`, `rlix_validation.py`, `sglang_engine.py`, all upstream) | NOT touched (C1 v6) | +| rlix | `rlix/pipeline/miles_model_update_service.py` | remove raise; cursor comm_ranks + world_size from injected per-engine count; claim-release ordering + wedged-sender leak-and-log; budget derivation | +| rlix | `rlix/pipeline/miles_coordinator.py` | mode env resolution + `_classify_broadcast_engines(target, mode)` (strict rejection / auto split / default all-cpu_serialize) + wiring both call sites + startup assertion entry | +| rlix | `rlix/pipeline/miles_pipeline.py` | pass topology inputs through `register_model_update_resources`; v7 uniformity guard on the resolved SGLang config (before phaseB step5); startup assertion call | +| rlix | `scripts/run_smoke_dual.sh` | one line (125): `--sglang-mem-fraction-static ${MILES_SMOKE_MEM_FRACTION:-0.30}` | +| rlix | `tests/` | E2/E3/E6 unit tests | + +## Appendix C: Full Risk → Evidence Matrix + +| Risk | Description | Bound by | +|------|-------------|----------| +| R1 | Rendezvous deadlock — happy-path ordering AND partial-receiver failure | E1 + E7; D3 bounded budgets | +| R2 | Rank misassignment for TP>1 (per-engine vs per-GPU) | E2; accepted residual: no TP>1 e2e in MVP (CHK3); uniformity guard-enforced (E10, CHK8) | +| R3 | GPU OOM staging bucket under mem 0.8 | two-level defense: miles S2 startup gate + runtime preflight with degradation (D5, E8); zero-OOM AC6 hard condition (E4) | +| R4 | Leaked NCCL group / port-claim lifecycle | per-sync group names; sender-owned destroy-in-finally + budgets + E6/E7; release-after-ack, leak-and-log on wedged sender | +| R5 | finalize/flush_cache interaction differs under broadcast | transport-agnostic path unchanged; watched in M4 logs (E4) | +| R6 (v6) | User selects `broadcast` on an overlap topology expecting it to work | C5 startup rejection with actionable error naming the colocate engines + suggesting `auto`; verified by E3 (unit) + E9 (e2e) | + +## Appendix D: Implementation Detail Trace + +| Detail | Parent | Note | +|--------|--------|------| +| Mode env parsed once per pipeline, logged, invalid value = error | D4 ← C5 | no silent default on typo | +| Memory preflight per bucket, 1 GiB margin (env-overridable), tensor-by-tensor degradation | D5 ← A3 | S2 gate is the startup-time companion | +| Broadcast order = `bucket.params` insertion order = metadata list order | D5 ← A2 | contract by construction | +| dtype strings normalized on metadata side only; sender broadcasts raw tensors | D5 | no cast | +| Nested receiver refs tracked sender-local; cancelled + destroyed in sender finally | D5 ← C6 | service inflight_refs does NOT cover these | +| Budgets: rendezvous ≤0.2×, transport ≤0.6×, teardown ≤0.1×, margin ≥0.1× of `timeout_s` | D3 ← C6 | invariant is the contract (CHK6) | +| Port claim: release after sender ref resolves; wedged+broadcast → leak-and-log; cpu_serialize-only → release-on-timeout unchanged | D5 ← C6 | ROLL-backend precedent | +| `destroy_collective_group` tolerates 400 missing-group | D3 | receiver side already shipped | +| Sender holds only existing `_cache_lock` | D5 ← C2 | no new locks | +| `world_size = 1` when broadcast set empty stays | D4 | cpu_serialize-only sessions untouched | +| C9 startup assertion: env mappings + scheduler grant logged; classification vs active target set, both pipelines, pre-training | D3 ← C9 | failure/divergence = smoke INVALID | +| Harness mem-fraction env parameterization, default 0.30 | D2 ← C9 | M4 run (a) exports 0.8; E4 echoes | + +## Appendix E: Plan-Mirrored Execution Anchors (auditor view) + +This appendix mirrors execution anchors found in the source plan. **It is not a new instruction set for the implementation agent — agents must read the source plan directly.** + +The plan instructs the agent to stop and ask if: + +- S1: any change would touch pristine upstream miles lines (C1 conflict). +- S2: e2e reveals SGLang receiver routes do not match A1 (server-side change would be required) — surface findings first. +- S3: the warmup question (D5) cannot be resolved by observation within one debugging session on the vast instance. +- S4: any commit/push/PR action is reached (C7 — explicit user instruction required). + +Branch anchors mirrored from the plan header: rlix work on `zhenyu/miles-nccl-broadcast` (from `zhenyu/miles-mvp-e2e` @ b4e0cf6); miles work on `zhenyu/miles-nccl-broadcast` (from `zhenyu/m11-mvp-test` @ 6ff8df3). + +## Appendix G: Plan Patch Suggestions + +| # | Location | Gap | Suggested patch | +|---|----------|-----|-----------------| +| G1 | §6 D5 warmup | `unknown` — whether SGLang route tolerates an unannounced warmup collective | Commit to NO warmup (matching the standalone reference) + explicit M4 fallback note ("if first-bucket broadcast stalls >Ns, add sentinel warmup and re-run") — removes the only `unknown` | +| G2 | §4 AC1 / §5 M1 | AC1 spans both repos but is delivered at M1 (miles-only) | Scope AC1's statement to "miles-side transport" explicitly, or move AC1 to M2 — aligns with CHK1 | +| G3 | §6 D2 | `register_model_update_resources` widening asserted safe without caller inventory | Add one line listing known callers (miles_pipeline phaseB step5 only) as evidence for CHK5 | +| G4 | §3 C9 | Startup assertion covers t=0 only | If CHK4 resolves to "per-sync guarantee required", demand the mixed-class check per sync session (derivable from the E4 log line) | + +--- + +Next suggested artifact: run `/scope-triage @plans/miles-nccl-broadcast-plan.md` after human audit passes. diff --git a/rlix/pipeline/miles_coordinator.py b/rlix/pipeline/miles_coordinator.py index 6df001a..94f1595 100644 --- a/rlix/pipeline/miles_coordinator.py +++ b/rlix/pipeline/miles_coordinator.py @@ -43,6 +43,151 @@ logger = logging.getLogger(__name__) +# --- rlix#42 user-selected weight-update transport mode ----------------- +# One rlix-side control (plan D4): `cpu_serialize` (default — today's +# behavior, no classification), `broadcast` (strict: ALL engines +# broadcast; fail-fast if any target engine shares a GPU with the train +# pool — NCCL cannot form a group containing the same physical GPU +# twice), `auto` (topology split; the only mode that mixes transports). +_TRANSPORT_MODE_ENV = "RLIX_MILES_UPDATE_TRANSPORT" +_VALID_TRANSPORT_MODES = ("cpu_serialize", "broadcast", "auto") +# Smoke-validity guard (plan C9): when set to "1", registration asserts +# that the pipeline's FULL engine set classifies into BOTH transport +# classes under `auto` — an overlap smoke that cannot produce mixed +# sessions is INVALID and must fail fast rather than pass vacuously. +_ASSERT_MIXED_ENV = "RLIX_ASSERT_MIXED_TRANSPORT" + + +def _assert_mixed_transport_startup( + *, + mode: str, + pipeline_id: str, + train_gpu_ids: Optional[List[int]], + infer_gpu_ids: Optional[List[int]], + per_engine: int, +) -> None: + """C9 startup assertion, gated by ``RLIX_ASSERT_MIXED_TRANSPORT=1``. + + Only meaningful for ``auto`` (the mixing mode): classify the full + engine set and require ≥1 broadcast-classified AND ≥1 + cpu_serialize-classified engine; log the resulting split. Any other + mode with the env set is a smoke misconfiguration — fail fast. + """ + if os.environ.get(_ASSERT_MIXED_ENV, "").strip() != "1": + return + if mode != "auto": + raise RuntimeError( + f"{_ASSERT_MIXED_ENV}=1 requires {_TRANSPORT_MODE_ENV}=auto " + f"(the only mixing mode); got mode={mode!r}" + ) + if not infer_gpu_ids: + raise RuntimeError( + f"{_ASSERT_MIXED_ENV}=1 requires registered topology (infer_gpu_ids)" + ) + engine_count = len(infer_gpu_ids) // max(per_engine, 1) + full_set = set(range(engine_count)) + broadcast = _classify_broadcast_engines( + target_engine_indices=full_set, + mode="auto", + train_gpu_ids=train_gpu_ids, + infer_gpu_ids=infer_gpu_ids, + per_engine=per_engine, + ) + cpu_serialize = full_set - set(broadcast) + logger.info( + "[MilesCoordinator] C9 mixed-transport startup assertion pipeline_id=%s " + "engines=%s broadcast=%s cpu_serialize=%s", + pipeline_id, + sorted(full_set), + sorted(broadcast), + sorted(cpu_serialize), + ) + if not broadcast or not cpu_serialize: + raise RuntimeError( + f"C9 mixed-transport startup assertion FAILED pipeline_id={pipeline_id}: " + f"engines={sorted(full_set)} broadcast={sorted(broadcast)} " + f"cpu_serialize={sorted(cpu_serialize)} — the overlap smoke requires " + "≥1 engine of EACH class per pipeline; this topology cannot produce " + "mixed sessions, so the smoke would pass vacuously (INVALID)." + ) + + +def _resolve_transport_mode() -> str: + """Parse the transport-mode env once; invalid value = fail-fast + (plan C5: no silent default on a typo).""" + raw = os.environ.get(_TRANSPORT_MODE_ENV, "").strip().lower() + if not raw: + return "cpu_serialize" + if raw not in _VALID_TRANSPORT_MODES: + raise RuntimeError( + f"{_TRANSPORT_MODE_ENV}={raw!r} is not a valid transport mode; " + f"expected one of {list(_VALID_TRANSPORT_MODES)}" + ) + return raw + + +def _classify_broadcast_engines( + *, + target_engine_indices: Set[int], + mode: str, + train_gpu_ids: Optional[List[int]], + infer_gpu_ids: Optional[List[int]], + per_engine: int, +) -> frozenset: + """Classify a sync's target engines into the broadcast subset (AC4). + + Engine → physical-GPU mapping follows the manager's contiguous + layout (cf. ``MilesPipeline._wait_for_overlap_engines_offloaded``): + local engine ``e`` occupies the ``per_engine`` consecutive GPUs + starting at ``sorted(infer_gpu_ids)[e * per_engine]``. + + - ``cpu_serialize``: empty set (no classification, no topology + required — backward-compatible default). + - ``auto``: engines whose GPU set is disjoint from the train pool → + broadcast; overlapping engines → cpu_serialize (the only mixing + mode). + - ``broadcast``: ALL engines must be broadcast-eligible; any + colocate target = RuntimeError (plan C5: honored exactly or the + run refuses to proceed, never a quiet downgrade). + """ + target = {int(i) for i in target_engine_indices} + if mode == "cpu_serialize" or not target: + return frozenset() + if mode not in _VALID_TRANSPORT_MODES: + raise RuntimeError(f"unknown transport mode {mode!r}") + if train_gpu_ids is None or infer_gpu_ids is None: + raise RuntimeError( + f"{_TRANSPORT_MODE_ENV}={mode} requires topology at " + "register_model_update_resources (train_gpu_ids + infer_gpu_ids); " + "registration did not provide it" + ) + train = {int(g) for g in train_gpu_ids} + infer_sorted = sorted(int(g) for g in infer_gpu_ids) + broadcast: Set[int] = set() + colocate: Set[int] = set() + for engine_index in sorted(target): + lo = engine_index * per_engine + hi = lo + per_engine + engine_gpus = infer_sorted[lo:hi] + if len(engine_gpus) != per_engine: + raise RuntimeError( + f"engine {engine_index} maps outside the infer pool " + f"(pool={infer_sorted}, per_engine={per_engine})" + ) + if train & set(engine_gpus): + colocate.add(engine_index) + else: + broadcast.add(engine_index) + if mode == "broadcast" and colocate: + raise RuntimeError( + f"{_TRANSPORT_MODE_ENV}=broadcast cannot serve colocate engines " + f"{sorted(colocate)} (their GPUs intersect the train pool " + f"{sorted(train)}); NCCL cannot form a group containing the same " + "physical GPU twice. Use 'auto' (mixed transports) or " + "'cpu_serialize'." + ) + return frozenset(broadcast) + _T = TypeVar("_T") # Default max_concurrency for the MILES pipeline actor (F108 Special C1). @@ -313,22 +458,99 @@ def register_model_update_resources( *, cache_owner_actor, rollout_manager, + train_gpu_ids: Optional[List[int]] = None, + infer_gpu_ids: Optional[List[int]] = None, + rollout_num_gpus_per_engine: int = 1, ) -> None: """X2 ctor handle injection — capture handles for lazy :class:`MilesModelUpdateService` construction. Pipeline calls this at init Step 6.6 BEFORE Step 7 INIT, so the service can be lazily built whenever the coordinator first needs to push a sync. + + Topology kwargs (rlix#42): ``train_gpu_ids`` / ``infer_gpu_ids`` + are this pipeline's physical GPU mappings + (``cluster_device_mappings`` actor_train / actor_infer); + ``rollout_num_gpus_per_engine`` is the uniform per-engine GPU + count (v7 uniformity guard runs in MilesPipeline before this + call). They feed transport classification. Optional for + backward compatibility: with the default ``cpu_serialize`` + transport mode no classification runs; any other mode + fail-fasts at sync time if topology was not registered. """ if cache_owner_actor is None or rollout_manager is None: raise ValueError( "register_model_update_resources requires both handles" ) + per_engine = int(rollout_num_gpus_per_engine) + if per_engine < 1: + raise ValueError( + f"rollout_num_gpus_per_engine must be >= 1; got {rollout_num_gpus_per_engine}" + ) + # Resolve the transport mode ONCE per pipeline (invalid value = + # fail-fast here, not at first sync) and log it (plan AC4). + mode = _resolve_transport_mode() + logger.info( + "[MilesCoordinator] transport mode=%s pipeline_id=%s " + "(env %s; train_gpus=%s infer_gpus=%s per_engine=%d)", + mode, + self._pipeline_id, + _TRANSPORT_MODE_ENV, + sorted(int(g) for g in train_gpu_ids) if train_gpu_ids else None, + sorted(int(g) for g in infer_gpu_ids) if infer_gpu_ids else None, + per_engine, + ) + _assert_mixed_transport_startup( + mode=mode, + pipeline_id=self._pipeline_id, + train_gpu_ids=list(train_gpu_ids) if train_gpu_ids is not None else None, + infer_gpu_ids=list(infer_gpu_ids) if infer_gpu_ids is not None else None, + per_engine=per_engine, + ) + # Capability probe (codex impl-r15): the sync-under-load bracket + # depends on the manager-side atomic register_router_if_active + # RPC. On a version-skewed deployment (new rlix + old miles) the + # missing method would only surface INSIDE the first sync's + # finally — after engines were already quiesced. Probe with an + # empty list at registration time so skew fails loud and early. + try: + ray.get(rollout_manager.register_router_if_active.remote([])) + except AttributeError as exc: + raise RuntimeError( + "rollout_manager lacks register_router_if_active — the miles " + "side of this deployment predates the rlix#42 sync-under-load " + "bracket; update miles (miles/ray/rollout.py) before running " + "RLix-mode weight sync" + ) from exc self._model_update_resources = { "cache_owner_actor": cache_owner_actor, "rollout_manager": rollout_manager, + "train_gpu_ids": ( + sorted(int(g) for g in train_gpu_ids) if train_gpu_ids is not None else None + ), + "infer_gpu_ids": ( + sorted(int(g) for g in infer_gpu_ids) if infer_gpu_ids is not None else None + ), + "rollout_num_gpus_per_engine": per_engine, + "transport_mode": mode, } + def _broadcast_set_for(self, target_engine_indices: Set[int]) -> frozenset: + """Transport classification for one sync's target set (plan AC4). + + Reads the mode + topology captured at registration; delegates to + the module-level :func:`_classify_broadcast_engines` (unit-tested + directly). Default ``cpu_serialize`` mode returns an empty set + without touching topology (backward compatible).""" + resources = self._model_update_resources + return _classify_broadcast_engines( + target_engine_indices=target_engine_indices, + mode=resources.get("transport_mode", "cpu_serialize"), + train_gpu_ids=resources.get("train_gpu_ids"), + infer_gpu_ids=resources.get("infer_gpu_ids"), + per_engine=int(resources.get("rollout_num_gpus_per_engine", 1)), + ) + def publish_cache_ready_step(self, step: int) -> int: """Init-bootstrap-only. Sets ``_cache_ready_step`` so any later runtime expand fired before the first ``after_training`` hook @@ -366,6 +588,9 @@ def _ensure_model_update_service(self): pipeline_id=self._pipeline_id, cache_owner_actor=self._model_update_resources["cache_owner_actor"], rollout_manager=self._model_update_resources["rollout_manager"], + rollout_num_gpus_per_engine=int( + self._model_update_resources.get("rollout_num_gpus_per_engine", 1) + ), ) return self._model_update_service @@ -386,19 +611,118 @@ def sync_base_weights_to_active(self, step: int) -> int: # Run the sync OUTSIDE the resize lock so concurrent # report_progress_from_scheduler / clear_progress_stream # callers do not block on weight transport. + broadcast_set = self._broadcast_set_for(set(target)) + # rlix#42 sync-under-load bracket (flush-timeout root cause): + # under the fully-async rollout, generation for the NEXT step is + # already live on these engines, and the router keeps dispatching + # new requests while finalize's /flush_cache waits for the queue + # to drain — a guaranteed 60 s timeout under real load. Mirror + # shrink_engines' proven ordering BEFORE the sync: close router + # admission (unregister raises on non-2xx, aborting the sync + # early rather than hanging in flush), then abort in-flight work + # (the router re-dispatches aborted requests afterwards — same + # contract the donor-shrink path relies on). Re-open admission in + # a finally: serving briefly-stale weights is recoverable, + # unrouted-forever engines are not. + rollout_manager = self._model_update_resources.get("rollout_manager") + handles = ray.get(rollout_manager.get_engine_handles.remote(sorted(target))) + sync_error: Exception | None = None try: + # Quiesce INSIDE the guarded region so any failure — including + # a partial unregister or an abort error — still reaches the + # re-register finally (codex impl-r11 high #1). + ray.get([h.unregister_from_router.remote() for h in handles.values()]) + ray.get(rollout_manager._abort_engines.remote(sorted(target))) return int( ray.get( service.sync_selected_workers.remote( sync_id=None, target_engine_indices=target, version=int(step), + broadcast_local_ranks=broadcast_set, ) ) ) except Exception as exc: + sync_error = exc logger.error("sync_base_weights_to_active failed: %r", exc) raise + finally: + # Re-open router admission via the per-engine + # ``register_with_router`` — idempotent at the router (re-add + # discards from dead_workers), safe even for engines whose + # unregister never happened. NOT ``activate_routing``: that + # is the loading→active INIT transition and raises for these + # still-active engines (codex impl-r11 high #2). + # Concurrent-shrink guard (codex impl-r12/r13): a resize can + # shrink one of these engines while the sync is in flight + # (both flows RPC outside the lock by design, R10-F1). The + # state-check + register must be ATOMIC w.r.t. shrink, so + # both happen inside a single serialized manager call + # (``register_router_if_active`` — the manager actor is + # single-threaded, so no shrink can interleave between its + # state read and the /add_worker). The coordinator-side + # intent intersection under the lock is a cheap pre-filter + # only; the manager call is the authority. + reregister_error: Exception | None = None + try: + with self._resize_sync_lock: + still_intended = sorted( + set(self._active_engine_indices) & set(target) + ) + reregistered = ( + ray.get( + rollout_manager.register_router_if_active.remote( + still_intended + ) + ) + if still_intended + else [] + ) + skipped = sorted(set(int(i) for i in target) - set(reregistered)) + if skipped: + logger.warning( + "sync-under-load bracket: NOT re-registering %s " + "(concurrently shrunk/offloaded or no longer intended " + "active)", + skipped, + ) + except Exception as exc: # noqa: BLE001 + reregister_error = exc + logger.error( + "sync-under-load bracket: failed to re-register router " + "workers for %s: %r", + sorted(target), + exc, + ) + # Independently best-effort: reset the abort-idempotency cache + # so a future genuine shrink re-aborts any in-flights that + # arrive from now on (these engines stay ACTIVE — shrink's + # release-on-offload path never fires for this bracket). + try: + ray.get( + rollout_manager._reset_abort_idempotency_for.remote(sorted(target)) + ) + except Exception as exc: # noqa: BLE001 + logger.error( + "sync-under-load bracket: failed to reset abort " + "idempotency for %s: %r", + sorted(target), + exc, + ) + # codex impl-r14: a re-registration FAILURE (exception — not + # the legitimate skipped-because-shrunk case) must never let a + # successful sync report success: the engines would stay + # intended-active yet unroutable, silently suspending + # generation. Escalate unless an earlier exception is already + # propagating (then keep the original as the primary error). + if reregister_error is not None and sync_error is None: + raise RuntimeError( + "sync_base_weights_to_active: sync succeeded but router " + f"re-registration failed for targets {sorted(target)}; " + "engines may be unroutable — failing fast rather than " + "reporting a healthy sync" + ) from reregister_error # ------------------------------------------------------------------ # F22 / F37 / F40 resize_infer + _expand_workers (iter 23) @@ -541,12 +865,14 @@ def _expand_workers(self, engine_indices: Set[int]) -> None: ) # F40 Runtime branch: wake → service.sync_selected_workers → # activate_routing. + broadcast_set = self._broadcast_set_for(set(engine_indices)) ray.get(rollout_manager.expand_engines.remote(sorted(engine_indices))) ray.get( service.sync_selected_workers.remote( sync_id=None, target_engine_indices=frozenset(engine_indices), version=cached_step, + broadcast_local_ranks=broadcast_set, ) ) ray.get(rollout_manager.activate_routing.remote(sorted(engine_indices))) diff --git a/rlix/pipeline/miles_model_update_service.py b/rlix/pipeline/miles_model_update_service.py index 7d5ccdb..8cece00 100644 --- a/rlix/pipeline/miles_model_update_service.py +++ b/rlix/pipeline/miles_model_update_service.py @@ -28,19 +28,77 @@ logger = logging.getLogger(__name__) +@ray.remote(num_cpus=0) +class _PortClaimStore: + """Minimal KV store implementing the SharedStorage port-claim protocol + (``try_put`` / ``delete``) for MILES mode, where ROLL's SharedStorage + singleton does not exist. + + rlix#42: the port claim is NOT optional once pipelines sync + concurrently (dedicated-train topology): every cache_owner scans free + ports from the same base (miles ``RayActor.get_free_port`` starts at + 20000), so two simultaneous sessions deterministically pick the SAME + rendezvous port, cross-wire their NCCL TCP stores, and crash mid + ``ncclUniqueId`` exchange. Detached + named so every service instance + on the cluster shares one claim space; claims die with the Ray + cluster (consistent with the leak-and-log rule for wedged senders). + """ + + def __init__(self): + self._claims: dict = {} + + def try_put(self, key, value) -> bool: + if key in self._claims: + return False + self._claims[key] = value + return True + + def delete(self, key) -> bool: + self._claims.pop(key, None) + return True + + +def _fallback_port_claim_store() -> Any: + from rlix.protocol.types import RLIX_NAMESPACE + + return _PortClaimStore.options( # type: ignore[attr-defined] + name="rlix:miles_port_claims", + namespace=RLIX_NAMESPACE, + get_if_exists=True, + lifetime="detached", + ).remote() + + def _get_shared_storage_actor() -> Any: + """Resolve the port-claim store: ROLL's SharedStorage when it exists, + else the rlix-owned fallback (see :class:`_PortClaimStore`). + + Fallback triggers ONLY on the two expected miles-mode signals — ROLL + unimportable (ImportError) or the SharedStorage actor absent + (``ray.get_actor`` ValueError). Any OTHER failure propagates: a + transient/unexpected lookup error must fail the sync fast rather + than silently split the claim namespace between the real store and + the fallback (codex impl-r7 medium). + """ try: - from roll.utils.constants import ( # type: ignore[import-not-found] - GLOBAL_STORAGE_NAMESPACE, - STORAGE_NAME, - ) + try: + from roll.utils.constants import ( # type: ignore[import-not-found] + GLOBAL_STORAGE_NAMESPACE, + STORAGE_NAME, + ) + except ImportError: + from roll.distributed.scheduler.storage import ( # type: ignore[import-not-found] + STORAGE_NAME, + ) + from roll.utils.constants import GLOBAL_STORAGE_NAMESPACE # type: ignore[import-not-found] except ImportError: - from roll.distributed.scheduler.storage import ( # type: ignore[import-not-found] - STORAGE_NAME, - ) - from roll.utils.constants import GLOBAL_STORAGE_NAMESPACE # type: ignore[import-not-found] - - return ray.get_actor(STORAGE_NAME, namespace=GLOBAL_STORAGE_NAMESPACE) + return _fallback_port_claim_store() + try: + return ray.get_actor(STORAGE_NAME, namespace=GLOBAL_STORAGE_NAMESPACE) + except ValueError: + # Ray raises ValueError for a named actor that does not exist — + # the normal miles-mode case (no ROLL control plane). + return _fallback_port_claim_store() @dataclasses.dataclass(frozen=True) @@ -119,6 +177,7 @@ def __init__( pipeline_id: str, cache_owner_actor, rollout_manager, + rollout_num_gpus_per_engine: int = 1, ): if not isinstance(pipeline_id, str) or not pipeline_id: raise ValueError("pipeline_id must be a non-empty str") @@ -126,6 +185,18 @@ def __init__( raise ValueError("cache_owner_actor handle is required (F107 / X2 injection)") if rollout_manager is None: raise ValueError("rollout_manager handle is required") + # C8 rank accounting input: NCCL group ranks are per-GPU, so each + # broadcast engine occupies `rollout_num_gpus_per_engine` + # consecutive ranks (its TP workers). Uniform per pipeline in + # RLix mode — enforced upstream by the MilesPipeline uniformity + # startup guard (plan v7); heterogeneous engines are out of + # scope (plan O6). + per_engine = int(rollout_num_gpus_per_engine) + if per_engine < 1: + raise ValueError( + f"rollout_num_gpus_per_engine must be >= 1; got {rollout_num_gpus_per_engine}" + ) + self._rollout_num_gpus_per_engine = per_engine self._pipeline_id = pipeline_id self._cache_owner_actor = cache_owner_actor self._rollout_manager = rollout_manager @@ -189,20 +260,6 @@ async def sync_selected_workers( f"broadcast_local_ranks={sorted(broadcast_set)} must be a subset " f"of target_engine_indices={sorted(target)}" ) - if broadcast_set: - # The cache-owner sender NCCL path (init_process_group + - # per-bucket dist.broadcast + dist.destroy_process_group) - # is not yet implemented inside MILES run_sync_session - # (iter 12 wires only the receiver-side fan-out). Fail - # fast rather than hang the receivers waiting for an - # absent sender. - raise NotImplementedError( - "broadcast_local_ranks transport requires sender-side NCCL " - "(init_process_group + dist.broadcast on the cache_owner). " - "MILES iter 12 only wired the receiver-side fan-out. " - "Until the sender-side path lands, route every target " - "through cpu_serialize." - ) cpu_serialize_set = target - broadcast_set # R06-F1 fix: track every Ray ObjectRef issued by this atomic @@ -216,6 +273,12 @@ async def sync_selected_workers( # Carries the (addr, port) claim out of the atomic unit so the # cancellation handlers below can still release it. port_claim_holder: list = [] + # Set (appended) the moment the run_sync_session ref RESOLVES — + # success or error alike, the sender's finally has run by then + # (teardown ack). Distinguishes a truly wedged sender from a + # cancellation that landed during the post-resolution claim + # release await (codex impl-r1 medium). + sender_resolved_holder: list = [] async def _run() -> int: return await self._run_atomic_unit( @@ -226,6 +289,7 @@ async def _run() -> int: broadcast_set=broadcast_set, inflight_refs=inflight_refs, port_claim_holder=port_claim_holder, + sender_resolved_holder=sender_resolved_holder, ) try: @@ -234,14 +298,24 @@ async def _run() -> int: return await asyncio.wait_for(_run(), timeout=float(self._timeout_s)) except asyncio.TimeoutError: self._cancel_inflight(inflight_refs, reason="wait_for timeout") - self._release_port_claim_nowait(port_claim_holder) + self._finalize_port_claim_on_abort( + port_claim_holder, + has_broadcast=bool(broadcast_set), + sender_resolved=bool(sender_resolved_holder), + reason="wait_for timeout", + ) raise except asyncio.CancelledError: # Outer cancellation (caller cancelled sync_selected_workers # task) — propagate after firing ray.cancel so we don't leak # inflight Ray work either. self._cancel_inflight(inflight_refs, reason="task cancelled") - self._release_port_claim_nowait(port_claim_holder) + self._finalize_port_claim_on_abort( + port_claim_holder, + has_broadcast=bool(broadcast_set), + sender_resolved=bool(sender_resolved_holder), + reason="task cancelled", + ) raise def _cancel_inflight(self, inflight_refs: list, *, reason: str) -> None: @@ -276,6 +350,7 @@ async def _run_atomic_unit( broadcast_set: frozenset[int], inflight_refs: list, port_claim_holder: list, + sender_resolved_holder: list, ) -> int: # Each .remote() is captured into inflight_refs BEFORE we await # so the outer cancellation handler can fire ray.cancel(force=True) @@ -309,13 +384,35 @@ async def _run_atomic_unit( inflight_refs.append(sync_ref) try: await _ray_get(sync_ref) - finally: - # The port is only needed for the NCCL rendezvous during - # transport, so release it however the transport ended. + except asyncio.CancelledError: + # Service session deadline fired while the sender ref was + # UNRESOLVED (wedged sender). Do NOT dispatch a claim delete + # here — the NCCL TCP store behind master_port may still be + # bound; leak-vs-release is decided by the outer handler + # (_finalize_port_claim_on_abort, C6 ownership rule). + raise + except Exception: + # Sender ref RESOLVED with an error: the sender's own finally + # has already run its teardown (teardown ack) — releasing the + # claim is safe. + sender_resolved_holder.append(True) if port_claim is not None: await self._release_port_claim( master_addr=port_claim[0], master_port=port_claim[1], + storage=port_claim[2], + inflight_refs=inflight_refs, + ) + port_claim_holder.clear() + raise + else: + # Success: teardown ack via normal return. + sender_resolved_holder.append(True) + if port_claim is not None: + await self._release_port_claim( + master_addr=port_claim[0], + master_port=port_claim[1], + storage=port_claim[2], inflight_refs=inflight_refs, ) port_claim_holder.clear() @@ -405,17 +502,18 @@ async def _build_plan( # False if the key already exists, which we treat as a # collision and re-pick. Bound the retry budget so a stuck # claim can't hang the sync. - try: - shared_storage = _get_shared_storage_actor() - except Exception as exc: # noqa: BLE001 - logger.warning( - "MilesModelUpdateService: SharedStorage actor unavailable; " - "skipping port claim: %r", - exc, - ) - shared_storage = None - - port_claim: tuple[str, int] | None = None + # Claim store resolution is NOT optional and NOT skippable + # (codex impl-r7): an unexpected lookup failure propagates and + # fails this sync fast instead of running unprotected (the old + # skip path let concurrent sessions collide on the deterministic + # get_free_port scan — the disjoint-topology crash). + shared_storage = _get_shared_storage_actor() + + # The claim tuple carries the ACCEPTING store handle so release + # always deletes from the same store that granted the claim — + # never a re-lookup that might resolve differently (codex + # impl-r7). + port_claim: tuple[str, int, Any] | None = None if shared_storage is not None: for attempt in range(8): claim_key = f"MASTER_ADDR_PORT:{master_addr}:{master_port}" @@ -423,10 +521,17 @@ async def _build_plan( inflight_refs.append(try_put_ref) claimed = bool(await _ray_get(try_put_ref)) if claimed: - port_claim = (master_addr, master_port) + port_claim = (master_addr, master_port, shared_storage) break - # Collision — pick another free port and re-try. - next_port_ref = self._cache_owner_actor.get_free_port.remote() + # Collision — pick another free port and re-try. The + # claimed port is only RESERVED (not yet OS-bound) by the + # winning session, so miles' deterministic free-port scan + # would return the very same value forever; force the + # scan to start PAST the contested port (rlix#42 fix for + # the concurrent-sync rendezvous collision). + next_port_ref = self._cache_owner_actor.get_free_port.remote( + start_port=master_port + 1 + ) inflight_refs.append(next_port_ref) next_port = int(await _ray_get(next_port_ref)) # Always advance master_port to the new pick to avoid @@ -444,11 +549,19 @@ async def _build_plan( "claims before retrying." ) - # Per-engine NCCL rank within the dynamic broadcast group. - # cache_owner is rank 0; receivers are 1..N. + # Per-engine NCCL rank_offset within the dynamic broadcast group + # (C8): ranks are per-GPU, not per-engine. cache_owner is rank 0; + # each broadcast engine occupies `per_engine` consecutive ranks + # (its TP workers — SGLang's init_weights_update_group registers + # tp_size ranks starting at rank_offset), so engine offsets form + # a cursor with uniform stride. per_engine == 1 degenerates to + # the dense 1..N assignment. + per_engine = self._rollout_num_gpus_per_engine comm_ranks: dict[int, int] = {} - for r, idx in enumerate(sorted(broadcast_set), start=1): - comm_ranks[idx] = r + cursor = 1 + for idx in sorted(broadcast_set): + comm_ranks[idx] = cursor + cursor += per_engine # cpu_serialize engines are not in the broadcast group; they # never enter setup_collective_group. Comm-rank assignment is # ignored for them but kept in the dict for plan-shape @@ -456,7 +569,9 @@ async def _build_plan( for idx in cpu_serialize_set: comm_ranks.setdefault(idx, 0) - world_size = 1 + len(broadcast_set) if broadcast_set else 1 + # world_size = 1 (sender) + per-GPU ranks of every broadcast + # engine == the final cursor value. + world_size = cursor if broadcast_set else 1 plan = SyncSessionPlan( sync_id=sync_id, @@ -473,6 +588,46 @@ async def _build_plan( ) return plan.as_wire_dict(), port_claim + def _finalize_port_claim_on_abort( + self, + port_claim_holder: list, + *, + has_broadcast: bool, + sender_resolved: bool, + reason: str, + ) -> None: + """C6 port-claim ownership on the abort path (plan v4/v7). + + cpu_serialize-only sessions keep the historical release-on-abort: + no NCCL TCP store ever bound the port, so deleting the claim is + harmless. A broadcast session whose sender ref already RESOLVED + (teardown ack) but got cancelled during the claim-release await + also releases — fire-and-forget — since the TCP store is already + retired (codex impl-r1 medium). Only the truly wedged sender + (broadcast leg + unresolved ref) intentionally LEAKS the claim: + it may still hold a live TCP store on ``master_port``, and + handing the port to a concurrent session would collide + (ROLL-backend precedent: leak is safer than collision). Fresh + per-session ``get_free_port`` picks keep later sessions safe; the + stale key dies with the SharedStorage actor. + """ + if not port_claim_holder: + return + if not has_broadcast or sender_resolved: + self._release_port_claim_nowait(port_claim_holder) + return + master_addr, master_port, _storage = port_claim_holder.pop() + logger.error( + "[MilesModelUpdateService] WEDGED-SENDER PORT-CLAIM LEAK " + "pipeline_id=%s addr=%s port=%s reason=%s: claim intentionally NOT " + "released (possibly-live NCCL TCP store); stale " + "MASTER_ADDR_PORT key persists until SharedStorage restart.", + self._pipeline_id, + master_addr, + master_port, + reason, + ) + def _release_port_claim_nowait(self, port_claim_holder: list) -> None: """Fire-and-forget release used on the cancellation path. @@ -483,10 +638,9 @@ def _release_port_claim_nowait(self, port_claim_holder: list) -> None: """ if not port_claim_holder: return - master_addr, master_port = port_claim_holder.pop() + master_addr, master_port, storage = port_claim_holder.pop() try: - shared_storage = _get_shared_storage_actor() - shared_storage.delete.remote(f"MASTER_ADDR_PORT:{master_addr}:{int(master_port)}") + storage.delete.remote(f"MASTER_ADDR_PORT:{master_addr}:{int(master_port)}") except Exception as exc: # noqa: BLE001 logger.warning( "MilesModelUpdateService: failed to release port claim " @@ -501,12 +655,15 @@ async def _release_port_claim( *, master_addr: str, master_port: int, + storage: Any, inflight_refs: list, ) -> None: try: - shared_storage = _get_shared_storage_actor() + # Delete from the SAME store that accepted the claim (codex + # impl-r7) — never re-resolve, which could pick a different + # backend under degraded lookup conditions. claim_key = f"MASTER_ADDR_PORT:{master_addr}:{int(master_port)}" - delete_ref = shared_storage.delete.remote(claim_key) + delete_ref = storage.delete.remote(claim_key) inflight_refs.append(delete_ref) await _ray_get(delete_ref) except Exception as exc: # noqa: BLE001 diff --git a/rlix/pipeline/miles_pipeline.py b/rlix/pipeline/miles_pipeline.py index 3591eb9..a9a1ea6 100644 --- a/rlix/pipeline/miles_pipeline.py +++ b/rlix/pipeline/miles_pipeline.py @@ -57,6 +57,43 @@ logger = logging.getLogger(__name__) +def _assert_uniform_engine_gpu_counts(sglang_config: Any, per_engine: int) -> None: + """rlix#42 v7 uniformity startup guard. + + The transport classification + NCCL rank cursor derive per-engine GPU + counts from ``miles_args.rollout_num_gpus_per_engine`` (uniform per + pipeline; heterogeneous engines are plan O6 out-of-scope). SGLang + server groups CAN override ``num_gpus_per_engine`` per group and the + miles C7 gates do not cross-check the resolved group values against + the arg — so this guard fail-fasts any divergence before + ``register_model_update_resources``. ``None`` / attribute-free + configs pass (nothing resolved to diverge). + """ + if sglang_config is None: + return + models = getattr(sglang_config, "models", None) or [] + for model in models: + model_name = getattr(model, "name", "?") + model_val = getattr(model, "num_gpus_per_engine", None) + if model_val is not None and int(model_val) != per_engine: + raise RuntimeError( + f"rlix#42 uniformity guard: sglang_config model {model_name!r} " + f"sets num_gpus_per_engine={model_val}, diverging from " + f"rollout_num_gpus_per_engine={per_engine}; heterogeneous " + "per-engine GPU counts are unsupported in RLix mode (plan O6)" + ) + for group_index, group in enumerate(getattr(model, "server_groups", None) or []): + group_val = getattr(group, "num_gpus_per_engine", None) + if group_val is not None and int(group_val) != per_engine: + raise RuntimeError( + f"rlix#42 uniformity guard: sglang_config model " + f"{model_name!r} server_groups[{group_index}] sets " + f"num_gpus_per_engine={group_val}, diverging from " + f"rollout_num_gpus_per_engine={per_engine}; heterogeneous " + "per-engine GPU counts are unsupported in RLix mode (plan O6)" + ) + + class MilesPipeline: """Per-pipeline actor created by :class:`MilesCoordinator`.""" @@ -401,11 +438,32 @@ def _init_phase_b_infer(self) -> None: # F107 / X2: register handles. F22 (relaxed): in M11.1 single-pipeline # this happens after the manager exists; the dual-pipeline-shell-init # F22 ordering is deferred. + # rlix#42: v7 uniformity guard runs BEFORE registration so the + # args-derived per-engine GPU count the transport classification + # relies on cannot silently diverge from a server-group override. + miles_args = self._pipeline_config.miles_args + per_engine = max(int(getattr(miles_args, "rollout_num_gpus_per_engine", 1) or 1), 1) + _assert_uniform_engine_gpu_counts( + getattr(self._pipeline_config, "sglang_config", None), per_engine + ) + cluster_mappings = ( + getattr(self._pipeline_config, "cluster_device_mappings", None) or {} + ) + train_gpu_ids = [int(g) for g in cluster_mappings.get("actor_train", [])] or None + infer_gpu_ids = [ + int(g) + for g in cluster_mappings.get( + "actor_infer", list(range(int(miles_args.rollout_num_gpus))) + ) + ] logger.info("[MilesPipeline] phaseB step5: register_model_update_resources start") ray.get( self._coordinator_handle.register_model_update_resources.remote( cache_owner_actor=self._cache_owner_actor, rollout_manager=self._rollout_manager, + train_gpu_ids=train_gpu_ids, + infer_gpu_ids=infer_gpu_ids, + rollout_num_gpus_per_engine=per_engine, ) ) logger.info("[MilesPipeline] phaseB step5: register_model_update_resources done") @@ -559,10 +617,18 @@ def _wait_for_overlap_engines_offloaded(self, allocated_train_gpus, *, timeout_s ) ) infer_first = min(infer_mapping) if infer_mapping else 0 + # rlix#42 disjoint-topology guard: only train GPUs that actually + # sit inside the infer pool have overlap engines to wait for. A + # dedicated-train GPU (disjoint mapping) has none — silently + # skipping it is correct, and prevents nonsense engine indices + # like (0 - infer_first) // per_engine == -1. + overlap_train_gpus = [ + int(g) for g in allocated_train_gpus if int(g) in set(infer_mapping) + ] target_indices = sorted( - {(int(g) - infer_first) // per_engine for g in allocated_train_gpus} + {(g - infer_first) // per_engine for g in overlap_train_gpus} ) - target_gpu_ids = sorted(set(int(g) for g in allocated_train_gpus)) + target_gpu_ids = sorted(overlap_train_gpus) if not target_indices: return diff --git a/scripts/run_smoke_dual.sh b/scripts/run_smoke_dual.sh index 63794c7..94909ed 100755 --- a/scripts/run_smoke_dual.sh +++ b/scripts/run_smoke_dual.sh @@ -122,7 +122,7 @@ python /root/miles/examples/rlix/run_miles_dual.py \ --optimizer adam --lr 1e-6 --lr-decay-style constant \ --weight-decay 0.1 --adam-beta1 0.9 --adam-beta2 0.98 \ --use-dynamic-batch-size --max-tokens-per-gpu 512 \ - --sglang-mem-fraction-static 0.30 \ + --sglang-mem-fraction-static "${MILES_SMOKE_MEM_FRACTION:-0.30}" \ --rollout-num-gpus 2 --rollout-num-gpus-per-engine 1 \ --use-miles-router \ --rollout-function-path examples.fully_async.fully_async_rollout.generate_rollout_fully_async \ diff --git a/tests/test_miles_service_broadcast.py b/tests/test_miles_service_broadcast.py new file mode 100644 index 0000000..190f838 --- /dev/null +++ b/tests/test_miles_service_broadcast.py @@ -0,0 +1,337 @@ +"""Unit tests for MilesModelUpdateService broadcast unlock (rlops/rlix#42). + +Covers plan miles-nccl-broadcast v7 evidence items: + +- E2: service no longer raises for a non-empty broadcast set; comm_ranks + is a per-GPU rank_offset cursor with uniform stride from the injected + ``rollout_num_gpus_per_engine``; ``world_size = 1 + per_engine × + len(broadcast_set)``; per_engine == 1 degenerates to the dense case. +- E6(a): the SharedStorage port claim is released only AFTER the + ``run_sync_session`` ref resolves (teardown ack), on both the success + and the sender-raised-exception paths. +- E6(b): the wedged-sender path (service session deadline fires while the + sender ref is unresolved) with a broadcast leg LEAKS the claim + (no delete dispatched); a cpu_serialize-only session keeps the + historical release-on-timeout. + +Ray is imported for real (tests run on the GPU instance) but no cluster +is started: the ``@ray.remote`` wrapper is bypassed via +``__ray_metadata__.modified_class`` and every handle is a local fake with +awaitable refs. +""" + +from __future__ import annotations + +import asyncio +import types +import unittest +from unittest import mock + +import rlix.pipeline.miles_model_update_service as svc_mod + + +def _service_cls(): + cls = svc_mod.MilesModelUpdateService + meta = getattr(cls, "__ray_metadata__", None) + return meta.modified_class if meta is not None else cls + + +class _Ref: + """Awaitable standing in for a Ray ObjectRef.""" + + def __init__(self, value=None, exc=None, event_log=None, event=None): + self._value = value + self._exc = exc + self._event_log = event_log + self._event = event + + def __await__(self): + async def _resolve(): + if self._event_log is not None and self._event is not None: + self._event_log.append(self._event) + if self._exc is not None: + raise self._exc + return self._value + + return _resolve().__await__() + + +class _PendingRef: + """Never-resolving awaitable (wedged sender).""" + + def __await__(self): + async def _forever(): + await asyncio.Future() + + return _forever().__await__() + + +class _RM: + def __init__(self, fn): + self._fn = fn + + def remote(self, *args, **kwargs): + return self._fn(*args, **kwargs) + + +class _FakeEngine: + def __init__(self): + self.pause_generation = _RM(lambda **k: _Ref(None)) + self.finalize_weight_update = _RM(lambda **k: _Ref(None)) + self.continue_generation = _RM(lambda **k: _Ref(None)) + + +class _FakeCacheOwner: + def __init__(self, log, sync_ref_factory): + self.plans = [] + self._log = log + + def _run(plan): + self.plans.append(plan) + return sync_ref_factory() + + # Mirrors miles RayActor.get_free_port: deterministic scan from a + # base, honoring start_port (the collision-retry contract). + self.get_node_ip = _RM(lambda: _Ref("10.0.0.1")) + self.get_free_port = _RM(lambda start_port=29500, **_k: _Ref(int(start_port))) + self.run_sync_session = _RM(_run) + + +class _FakeManager: + def __init__(self, engine_indices): + self._engines = {i: _FakeEngine() for i in engine_indices} + self.get_engine_handles = _RM(lambda idx: _Ref(dict(self._engines))) + self.set_weight_version = _RM(lambda *a, **k: _Ref(None)) + + +class _FakeSharedStorage: + def __init__(self, log, taken_keys=()): + self._log = log + self._taken = set(taken_keys) + self.try_put = _RM(self._try_put) + self.delete = _RM(self._delete) + + def _try_put(self, key, owner): + if key in self._taken: + self._log.append(("claim_collision", key)) + return _Ref(False) + self._log.append(("claim_put", key)) + return _Ref(True) + + def _delete(self, key): + self._log.append(("claim_delete", key)) + return _Ref(True) + + +class _Harness: + def __init__(self, *, engine_indices, per_engine=1, sync_ref_factory=None): + self.log: list = [] + factory = sync_ref_factory or ( + lambda: _Ref(0, event_log=self.log, event=("sender_resolved",)) + ) + self.cache_owner = _FakeCacheOwner(self.log, factory) + self.manager = _FakeManager(engine_indices) + self.storage = _FakeSharedStorage(self.log) + self.cancelled: list = [] + cls = _service_cls() + self.svc = cls( + pipeline_id="pipe-test", + cache_owner_actor=self.cache_owner, + rollout_manager=self.manager, + rollout_num_gpus_per_engine=per_engine, + ) + + def run(self, *, target, broadcast=None, timeout_s=None): + if timeout_s is not None: + self.svc._timeout_s = timeout_s + fake_ray = types.SimpleNamespace( + cancel=lambda ref, force=False: self.cancelled.append(ref) + ) + with mock.patch.object(svc_mod, "ray", fake_ray), mock.patch.object( + svc_mod, "_get_shared_storage_actor", lambda: self.storage + ): + return asyncio.run( + self.svc.sync_selected_workers( + sync_id="sync-test", + target_engine_indices=target, + version=7, + broadcast_local_ranks=broadcast, + ) + ) + + def events(self, name): + return [ev for ev in self.log if ev[0] == name] + + +class TestE2RankCursor(unittest.TestCase): + def test_broadcast_set_no_longer_raises_and_cursor_per_engine_2(self): + h = _Harness(engine_indices={0, 1, 2}, per_engine=2) + version = h.run(target={0, 1, 2}, broadcast={0, 1}) + self.assertEqual(version, 7) + self.assertEqual(len(h.cache_owner.plans), 1) + plan = h.cache_owner.plans[0] + # Cursor with stride 2: engine 0 -> rank_offset 1, engine 1 -> 3. + self.assertEqual(plan["comm_ranks"][0], 1) + self.assertEqual(plan["comm_ranks"][1], 3) + self.assertEqual(plan["comm_ranks"][2], 0) # cpu_serialize placeholder + # world_size = 1 sender + 2 engines x 2 GPUs. + self.assertEqual(plan["world_size"], 5) + self.assertEqual(sorted(plan["broadcast_local_ranks"]), [0, 1]) + self.assertEqual(sorted(plan["cpu_serialize_local_ranks"]), [2]) + + def test_per_engine_1_degenerates_to_dense_ranks(self): + h = _Harness(engine_indices={0, 1}, per_engine=1) + h.run(target={0, 1}, broadcast={0, 1}) + plan = h.cache_owner.plans[0] + self.assertEqual(plan["comm_ranks"], {0: 1, 1: 2}) + self.assertEqual(plan["world_size"], 3) + + def test_empty_broadcast_keeps_world_size_1(self): + h = _Harness(engine_indices={0, 1}, per_engine=4) + h.run(target={0, 1}) + plan = h.cache_owner.plans[0] + self.assertEqual(plan["world_size"], 1) + self.assertEqual(sorted(plan["cpu_serialize_local_ranks"]), [0, 1]) + + def test_broadcast_must_be_subset_of_target(self): + h = _Harness(engine_indices={0, 1}, per_engine=1) + with self.assertRaises(ValueError): + h.run(target={0}, broadcast={0, 5}) + + def test_invalid_per_engine_rejected_at_ctor(self): + with self.assertRaises(ValueError): + _service_cls()( + pipeline_id="p", + cache_owner_actor=object(), + rollout_manager=object(), + rollout_num_gpus_per_engine=0, + ) + + +class TestPortClaimCollision(unittest.TestCase): + def test_contested_port_repicks_past_the_claim(self): + # rlix#42 concurrent-sync fix: a peer pipeline holds the claim on + # the deterministic first pick (e.g. 29500). The retry must ask + # get_free_port to scan PAST the contested port — the port is + # only claim-reserved, not OS-bound, so a plain rescan would + # return the same value forever and exhaust the retry budget. + h = _Harness(engine_indices={0}, per_engine=1) + h.storage = _FakeSharedStorage( + h.log, taken_keys={"MASTER_ADDR_PORT:10.0.0.1:29500"} + ) + h.run(target={0}, broadcast={0}) + plan = h.cache_owner.plans[0] + self.assertEqual(plan["master_port"], 29501) + self.assertEqual(len(h.events("claim_collision")), 1) + self.assertEqual(len(h.events("claim_put")), 1) + + def test_unexpected_storage_lookup_error_fails_fast(self): + # codex impl-r7: a transient/unexpected claim-store lookup error + # must fail the sync, not silently skip claims (unprotected) or + # fall back to a second store (split claim namespace). + h = _Harness(engine_indices={0}, per_engine=1) + + def _boom(): + raise RuntimeError("transient storage lookup failure") + + fake_ray = types.SimpleNamespace(cancel=lambda ref, force=False: None) + with mock.patch.object(svc_mod, "ray", fake_ray), mock.patch.object( + svc_mod, "_get_shared_storage_actor", _boom + ): + with self.assertRaisesRegex(RuntimeError, "transient storage"): + asyncio.run( + h.svc.sync_selected_workers( + sync_id="s", + target_engine_indices={0}, + version=1, + broadcast_local_ranks={0}, + ) + ) + # Nothing reached the sender: fail-fast happened at plan build. + self.assertEqual(h.cache_owner.plans, []) + + def test_port_claim_store_semantics(self): + cls = svc_mod._PortClaimStore + meta = getattr(cls, "__ray_metadata__", None) + store = (meta.modified_class if meta is not None else cls)() + self.assertTrue(store.try_put("k1", "p1")) + self.assertFalse(store.try_put("k1", "p2")) # second claimer loses + self.assertTrue(store.delete("k1")) + self.assertTrue(store.try_put("k1", "p2")) # free after delete + self.assertTrue(store.delete("missing")) # idempotent + + +class TestE6PortClaimOwnership(unittest.TestCase): + def test_claim_released_only_after_sender_resolves_success(self): + h = _Harness(engine_indices={0, 1}, per_engine=1) + h.run(target={0, 1}, broadcast={0}) + resolved = h.log.index(("sender_resolved",)) + deletes = [i for i, ev in enumerate(h.log) if ev[0] == "claim_delete"] + self.assertEqual(len(deletes), 1, f"log: {h.log}") + self.assertLess(resolved, deletes[0]) + + def test_claim_released_when_sender_resolves_with_exception(self): + h = _Harness( + engine_indices={0}, + per_engine=1, + sync_ref_factory=lambda: _Ref(exc=RuntimeError("sender aborted")), + ) + h.log # sync_ref_factory closes over nothing that logs resolution + with self.assertRaisesRegex(RuntimeError, "sender aborted"): + h.run(target={0}, broadcast={0}) + # Sender ref resolved (with an error) => teardown ack => release. + self.assertEqual(len(h.events("claim_delete")), 1) + + def test_wedged_sender_with_broadcast_leaks_claim(self): + h = _Harness( + engine_indices={0}, + per_engine=1, + sync_ref_factory=lambda: _PendingRef(), + ) + with self.assertRaises(asyncio.TimeoutError): + h.run(target={0}, broadcast={0}, timeout_s=0.2) + # C6 v4 rule: no delete dispatched — claim intentionally leaked. + self.assertEqual(h.events("claim_delete"), []) + # Inflight refs were cancelled. + self.assertTrue(h.cancelled) + + def test_wedged_sender_cpu_serialize_only_still_releases(self): + h = _Harness( + engine_indices={0}, + per_engine=1, + sync_ref_factory=lambda: _PendingRef(), + ) + with self.assertRaises(asyncio.TimeoutError): + h.run(target={0}, timeout_s=0.2) + # Historical behavior preserved: no TCP store => release on abort. + self.assertEqual(len(h.events("claim_delete")), 1) + + def test_cancel_during_post_resolution_release_still_releases(self): + # codex impl-r1 medium: sender ref resolves, but the service + # deadline lands during the claim-release await. The claim must + # be released via the fire-and-forget path (the TCP store is + # already retired), NOT leaked as a wedged sender. + h = _Harness(engine_indices={0}, per_engine=1) + # First delete dispatch pends forever (release await never + # completes); the fire-and-forget retry must dispatch a second + # delete after the timeout. + deletes = [] + + def _delete(key): + deletes.append(key) + if len(deletes) == 1: + h.log.append(("claim_delete", key)) + return _PendingRef() + h.log.append(("claim_delete", key)) + return _Ref(True) + + h.storage.delete = _RM(_delete) + with self.assertRaises(asyncio.TimeoutError): + h.run(target={0}, broadcast={0}, timeout_s=0.3) + # Sender resolved before the wedge, so no leak: two delete + # dispatches (the pending awaited one + the nowait retry). + self.assertEqual(len(h.events("claim_delete")), 2) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test_sync_bracket_concurrency.py b/tests/test_sync_bracket_concurrency.py new file mode 100644 index 0000000..0e40fba --- /dev/null +++ b/tests/test_sync_bracket_concurrency.py @@ -0,0 +1,232 @@ +"""Behavioral concurrency test for the sync-under-load bracket +(codex impl-r12): a shrink that commits an engine offloaded while the +sync RPC is in flight must NOT be undone by the bracket's finally — +the concurrently-shrunk engine stays out of router admission; the +still-active engine is re-registered. + +The coordinator instance is built via ``object.__new__`` (bypassing the +Ray-heavy ctor) with only the fields ``sync_base_weights_to_active`` +touches; every remote handle is a local fake whose ``.remote()`` returns +the plain value, and the module's ``ray.get`` is patched to unwrap. +""" + +from __future__ import annotations + +import threading +import types +import unittest +from unittest import mock + +import rlix.pipeline.miles_coordinator as coord_mod + + +def _coordinator_cls(): + cls = coord_mod.MilesCoordinator + meta = getattr(cls, "__ray_metadata__", None) + return meta.modified_class if meta is not None else cls + + +class _RM: + def __init__(self, fn): + self._fn = fn + + def remote(self, *a, **k): + return self._fn(*a, **k) + + +class _FakeEngineHandle: + def __init__(self, log, idx): + self.unregister_from_router = _RM(lambda: log.append(("unregister", idx))) + self.register_with_router = _RM(lambda: log.append(("register", idx))) + + +class TestSyncBracketConcurrentShrink(unittest.TestCase): + def test_finally_skips_concurrently_shrunk_engine(self): + log: list = [] + cls = _coordinator_cls() + coord = object.__new__(cls) + coord._pipeline_id = "pipe-test" + coord._resize_sync_lock = threading.Lock() + coord._active_engine_indices = {0, 1} + coord._cache_ready_step = None + coord._model_update_service = object() # _ensure returns this + + handles = {0: _FakeEngineHandle(log, 0), 1: _FakeEngineHandle(log, 1)} + # Manager: engine 0 gets shrunk mid-sync — states reflect it. + def _register_if_active(idx_list): + # Mirrors the manager-side atomic method: state check and + # register inside one serialized call. + states = {i: ("offloaded" if i == 0 else "active") for i in idx_list} + out = [] + for i in idx_list: + if states[i] == "active": + log.append(("register", i)) + out.append(i) + return out + + manager = types.SimpleNamespace( + get_engine_handles=_RM(lambda idx: dict(handles)), + register_router_if_active=_RM(_register_if_active), + _abort_engines=_RM(lambda idx: log.append(("abort", tuple(idx)))), + _reset_abort_idempotency_for=_RM( + lambda idx: log.append(("reset_idempotency", tuple(idx))) + ), + ) + coord._model_update_resources = { + "cache_owner_actor": object(), + "rollout_manager": manager, + "train_gpu_ids": [0], + "infer_gpu_ids": [0, 1, 2], + "rollout_num_gpus_per_engine": 1, + "transport_mode": "cpu_serialize", + } + + def _sync(**kwargs): + # Concurrent shrink lands while the sync RPC is in flight: + # engine 0 leaves the intended-active set (its manager state + # is already "offloaded" per the fake above). + coord._active_engine_indices = {1} + log.append(("sync", tuple(sorted(kwargs["target_engine_indices"])))) + return 7 + + service = types.SimpleNamespace(sync_selected_workers=_RM(_sync)) + fake_ray = types.SimpleNamespace( + get=lambda refs: refs, # fakes resolve eagerly; unwrap is identity + ) + with mock.patch.object(coord_mod, "ray", fake_ray), mock.patch.object( + cls, "_ensure_model_update_service", lambda self: service + ): + version = coord.sync_base_weights_to_active(7) + + self.assertEqual(version, 7) + registered = [ev for ev in log if ev[0] == "register"] + # Engine 1 (still active + still intended) re-registered; engine 0 + # (concurrently shrunk) must NOT be re-admitted. + self.assertEqual(registered, [("register", 1)]) + # Quiesce ran for both before the sync. + self.assertIn(("unregister", 0), log) + self.assertIn(("unregister", 1), log) + self.assertLess(log.index(("unregister", 0)), log.index(("sync", (0, 1)))) + # Idempotency reset still ran. + self.assertTrue([ev for ev in log if ev[0] == "reset_idempotency"]) + + def test_reregister_failure_after_successful_sync_escalates(self): + # codex impl-r14: a successful sync must NOT report success when + # the authoritative re-register raises — engines would remain + # intended-active yet unroutable. + log: list = [] + cls = _coordinator_cls() + coord = object.__new__(cls) + coord._pipeline_id = "pipe-test" + coord._resize_sync_lock = threading.Lock() + coord._active_engine_indices = {0, 1} + coord._cache_ready_step = None + handles = {0: _FakeEngineHandle(log, 0), 1: _FakeEngineHandle(log, 1)} + + def _register_boom(idx_list): + raise RuntimeError("router add_worker 503") + + manager = types.SimpleNamespace( + get_engine_handles=_RM(lambda idx: dict(handles)), + register_router_if_active=_RM(_register_boom), + _abort_engines=_RM(lambda idx: None), + _reset_abort_idempotency_for=_RM(lambda idx: None), + ) + coord._model_update_resources = { + "cache_owner_actor": object(), + "rollout_manager": manager, + "train_gpu_ids": None, + "infer_gpu_ids": None, + "rollout_num_gpus_per_engine": 1, + "transport_mode": "cpu_serialize", + } + service = types.SimpleNamespace(sync_selected_workers=_RM(lambda **k: 5)) + fake_ray = types.SimpleNamespace(get=lambda refs: refs) + with mock.patch.object(coord_mod, "ray", fake_ray), mock.patch.object( + cls, "_ensure_model_update_service", lambda self: service + ): + with self.assertRaisesRegex(RuntimeError, "re-registration failed"): + coord.sync_base_weights_to_active(5) + + def test_sync_failure_keeps_original_error_over_reregister_failure(self): + # When the sync body itself failed, that original error stays the + # primary exception even if re-register also fails. + log: list = [] + cls = _coordinator_cls() + coord = object.__new__(cls) + coord._pipeline_id = "pipe-test" + coord._resize_sync_lock = threading.Lock() + coord._active_engine_indices = {0} + coord._cache_ready_step = None + handles = {0: _FakeEngineHandle(log, 0)} + manager = types.SimpleNamespace( + get_engine_handles=_RM(lambda idx: dict(handles)), + register_router_if_active=_RM( + lambda idx: (_ for _ in ()).throw(RuntimeError("router down")) + ), + _abort_engines=_RM(lambda idx: None), + _reset_abort_idempotency_for=_RM(lambda idx: None), + ) + coord._model_update_resources = { + "cache_owner_actor": object(), + "rollout_manager": manager, + "train_gpu_ids": None, + "infer_gpu_ids": None, + "rollout_num_gpus_per_engine": 1, + "transport_mode": "cpu_serialize", + } + + def _sync_fail(**k): + raise TimeoutError("flush timeout") + + service = types.SimpleNamespace(sync_selected_workers=_RM(_sync_fail)) + fake_ray = types.SimpleNamespace(get=lambda refs: refs) + with mock.patch.object(coord_mod, "ray", fake_ray), mock.patch.object( + cls, "_ensure_model_update_service", lambda self: service + ): + with self.assertRaisesRegex(TimeoutError, "flush timeout"): + coord.sync_base_weights_to_active(5) + + def test_finally_reregisters_all_on_no_concurrent_change(self): + log: list = [] + cls = _coordinator_cls() + coord = object.__new__(cls) + coord._pipeline_id = "pipe-test" + coord._resize_sync_lock = threading.Lock() + coord._active_engine_indices = {0, 1} + coord._cache_ready_step = None + handles = {0: _FakeEngineHandle(log, 0), 1: _FakeEngineHandle(log, 1)} + def _register_all(idx_list): + for i in idx_list: + log.append(("register", i)) + return list(idx_list) + + manager = types.SimpleNamespace( + get_engine_handles=_RM(lambda idx: dict(handles)), + register_router_if_active=_RM(_register_all), + _abort_engines=_RM(lambda idx: None), + _reset_abort_idempotency_for=_RM(lambda idx: None), + ) + coord._model_update_resources = { + "cache_owner_actor": object(), + "rollout_manager": manager, + "train_gpu_ids": None, + "infer_gpu_ids": None, + "rollout_num_gpus_per_engine": 1, + "transport_mode": "cpu_serialize", + } + service = types.SimpleNamespace( + sync_selected_workers=_RM(lambda **k: 3) + ) + fake_ray = types.SimpleNamespace(get=lambda refs: refs) + with mock.patch.object(coord_mod, "ray", fake_ray), mock.patch.object( + cls, "_ensure_model_update_service", lambda self: service + ): + version = coord.sync_base_weights_to_active(3) + self.assertEqual(version, 3) + registered = sorted(ev[1] for ev in log if ev[0] == "register") + self.assertEqual(registered, [0, 1]) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test_sync_under_load_bracket.py b/tests/test_sync_under_load_bracket.py new file mode 100644 index 0000000..52b068d --- /dev/null +++ b/tests/test_sync_under_load_bracket.py @@ -0,0 +1,108 @@ +"""AST-level ordering assertions for the rlix#42 sync-under-load bracket +in ``MilesCoordinator.sync_base_weights_to_active``. + +The bracket mirrors shrink_engines' proven quiesce ordering so +finalize's /flush_cache cannot race live router dispatch: router +admission closes (unregister_from_router) and in-flight work aborts +(_abort_engines) BEFORE the sync RPC; routing re-activates + the abort +idempotency cache resets in a ``finally``. + +Follows the repo's AST-test pattern (cf. +test_miles_model_update_service_cleanup.py) so it runs without ray. +""" + +from __future__ import annotations + +import ast +from pathlib import Path + +REPO_ROOT = Path(__file__).resolve().parents[1] + + +def _load_sync_base_fn() -> ast.FunctionDef: + source = (REPO_ROOT / "rlix" / "pipeline" / "miles_coordinator.py").read_text( + encoding="utf-8" + ) + tree = ast.parse(source) + for node in ast.walk(tree): + if isinstance(node, ast.FunctionDef) and node.name == "sync_base_weights_to_active": + return node + raise AssertionError("sync_base_weights_to_active not found") + + +def _calls_attr(node: ast.AST, attr: str) -> bool: + return any( + isinstance(child, ast.Attribute) and child.attr == attr + for child in ast.walk(node) + ) + + +def _first_line_calling(fn: ast.FunctionDef, attr: str) -> int: + lines = [ + child.lineno + for child in ast.walk(fn) + if isinstance(child, ast.Attribute) and child.attr == attr + ] + assert lines, f"{attr} not called inside sync_base_weights_to_active" + return min(lines) + + +def test_quiesce_happens_before_the_sync_rpc() -> None: + fn = _load_sync_base_fn() + unregister_line = _first_line_calling(fn, "unregister_from_router") + abort_line = _first_line_calling(fn, "_abort_engines") + sync_line = _first_line_calling(fn, "sync_selected_workers") + assert unregister_line < abort_line < sync_line, ( + "bracket order must be unregister_from_router -> _abort_engines -> " + f"sync_selected_workers (got lines {unregister_line}, {abort_line}, " + f"{sync_line})" + ) + + +def test_quiesce_is_inside_the_guarded_region() -> None: + # codex impl-r11 high #1: a failure in unregister/abort must still + # reach the re-register finally — so BOTH must live in the try body. + fn = _load_sync_base_fn() + try_nodes = [ + node + for node in ast.walk(fn) + if isinstance(node, ast.Try) + and any(_calls_attr(stmt, "sync_selected_workers") for stmt in node.body) + ] + assert try_nodes, "sync RPC must be wrapped in try/finally" + outer = try_nodes[0] + assert any( + _calls_attr(stmt, "unregister_from_router") for stmt in outer.body + ), "unregister_from_router must be inside the try body" + assert any( + _calls_attr(stmt, "_abort_engines") for stmt in outer.body + ), "_abort_engines must be inside the try body" + + +def test_reregister_lives_in_finally_and_is_not_activate_routing() -> None: + fn = _load_sync_base_fn() + try_nodes = [ + node + for node in ast.walk(fn) + if isinstance(node, ast.Try) + and any(_calls_attr(stmt, "sync_selected_workers") for stmt in node.body) + ] + assert try_nodes, "sync RPC must be wrapped in try/finally" + outer = try_nodes[0] + # codex impl-r11 high #2 + impl-r13 TOCTOU: the re-register must be + # the manager-side ATOMIC register_router_if_active (state check and + # /add_worker in one serialized manager call) — never activate_routing + # (loading-only INIT transition) nor a raw per-engine register after a + # separate state read. + assert any( + _calls_attr(stmt, "register_router_if_active") for stmt in outer.finalbody + ), "register_router_if_active must run from the finally block" + assert not any( + _calls_attr(stmt, "activate_routing") for stmt in outer.finalbody + ), "activate_routing must NOT be used for still-active engines" + assert not any( + _calls_attr(stmt, "register_with_router") for stmt in outer.finalbody + ), "raw per-engine register_with_router in the finally is a TOCTOU (impl-r13)" + assert any( + _calls_attr(stmt, "_reset_abort_idempotency_for") for stmt in outer.finalbody + ), "_reset_abort_idempotency_for must run from the finally block" diff --git a/tests/test_transport_mode_classification.py b/tests/test_transport_mode_classification.py new file mode 100644 index 0000000..226dccc --- /dev/null +++ b/tests/test_transport_mode_classification.py @@ -0,0 +1,216 @@ +"""Unit tests for the rlix#42 transport-mode wiring (plan v7, M3). + +- E3: mode resolution (default/unset → cpu_serialize; invalid env → + fail-fast) and classification per mode: `cpu_serialize` never + classifies; `auto` splits overlap→cpu_serialize / disjoint→broadcast + (the only mixing mode); strict `broadcast` raises on any colocate + target and never mixes. +- E10: v7 uniformity startup guard — resolved SGLang config server-group + ``num_gpus_per_engine`` overrides diverging from + ``rollout_num_gpus_per_engine`` fail-fast naming the offender. +""" + +from __future__ import annotations + +import types +import unittest +from unittest import mock + +from rlix.pipeline.miles_coordinator import ( + _ASSERT_MIXED_ENV, + _TRANSPORT_MODE_ENV, + _assert_mixed_transport_startup, + _classify_broadcast_engines, + _resolve_transport_mode, +) +from rlix.pipeline.miles_pipeline import _assert_uniform_engine_gpu_counts + + +class TestE3ModeResolution(unittest.TestCase): + def test_unset_env_defaults_to_cpu_serialize(self): + with mock.patch.dict("os.environ", {}, clear=False): + import os + + os.environ.pop(_TRANSPORT_MODE_ENV, None) + self.assertEqual(_resolve_transport_mode(), "cpu_serialize") + + def test_valid_modes_parse(self): + for mode in ("cpu_serialize", "broadcast", "auto"): + with mock.patch.dict("os.environ", {_TRANSPORT_MODE_ENV: mode}): + self.assertEqual(_resolve_transport_mode(), mode) + + def test_invalid_mode_fails_fast(self): + with mock.patch.dict("os.environ", {_TRANSPORT_MODE_ENV: "brodcast"}): + with self.assertRaisesRegex(RuntimeError, "not a valid transport mode"): + _resolve_transport_mode() + + +class TestE3Classification(unittest.TestCase): + # The M4 run (a) harness topology (run_smoke_dual.sh defaults). + P1 = dict(train_gpu_ids=[0], infer_gpu_ids=[0, 1, 2], per_engine=1) + P2 = dict(train_gpu_ids=[3], infer_gpu_ids=[1, 2, 3], per_engine=1) + + def test_cpu_serialize_mode_never_classifies(self): + out = _classify_broadcast_engines( + target_engine_indices={0, 1, 2}, + mode="cpu_serialize", + train_gpu_ids=None, + infer_gpu_ids=None, + per_engine=1, + ) + self.assertEqual(out, frozenset()) + + def test_auto_mixes_on_overlap_topology_p1(self): + out = _classify_broadcast_engines( + target_engine_indices={0, 1, 2}, mode="auto", **self.P1 + ) + # e0@gpu0 shares the train pool -> cpu_serialize; e1@1, e2@2 -> broadcast. + self.assertEqual(out, frozenset({1, 2})) + + def test_auto_mixes_on_overlap_topology_p2(self): + out = _classify_broadcast_engines( + target_engine_indices={0, 1, 2}, mode="auto", **self.P2 + ) + # infer pool sorted [1,2,3]: e0@1, e1@2 -> broadcast; e2@3 -> colocate. + self.assertEqual(out, frozenset({0, 1})) + + def test_strict_broadcast_rejects_colocate_target(self): + with self.assertRaisesRegex(RuntimeError, r"colocate engines \[0\]"): + _classify_broadcast_engines( + target_engine_indices={0, 1, 2}, mode="broadcast", **self.P1 + ) + + def test_strict_broadcast_all_disjoint_passes(self): + out = _classify_broadcast_engines( + target_engine_indices={0, 1}, + mode="broadcast", + train_gpu_ids=[0], + infer_gpu_ids=[1, 2], + per_engine=1, + ) + self.assertEqual(out, frozenset({0, 1})) + + def test_strict_broadcast_subset_target_avoiding_colocate_passes(self): + # Same overlap topology, but the sync target set excludes the + # colocate engine — strict mode serves it. + out = _classify_broadcast_engines( + target_engine_indices={1, 2}, mode="broadcast", **self.P1 + ) + self.assertEqual(out, frozenset({1, 2})) + + def test_non_default_mode_requires_topology(self): + with self.assertRaisesRegex(RuntimeError, "requires topology"): + _classify_broadcast_engines( + target_engine_indices={0}, + mode="auto", + train_gpu_ids=None, + infer_gpu_ids=None, + per_engine=1, + ) + + def test_engine_outside_pool_rejected(self): + with self.assertRaisesRegex(RuntimeError, "maps outside the infer pool"): + _classify_broadcast_engines( + target_engine_indices={5}, mode="auto", **self.P1 + ) + + def test_per_engine_2_uses_gpu_slices(self): + out = _classify_broadcast_engines( + target_engine_indices={0, 1}, + mode="auto", + train_gpu_ids=[0], + infer_gpu_ids=[0, 1, 2, 3], + per_engine=2, + ) + # e0 -> gpus [0,1] (touches train gpu 0) -> colocate; e1 -> [2,3]. + self.assertEqual(out, frozenset({1})) + + def test_empty_target_returns_empty(self): + out = _classify_broadcast_engines( + target_engine_indices=set(), mode="auto", **self.P1 + ) + self.assertEqual(out, frozenset()) + + +class TestC9MixedStartupAssertion(unittest.TestCase): + P1 = dict(train_gpu_ids=[0], infer_gpu_ids=[0, 1, 2], per_engine=1) + + def test_env_unset_is_noop(self): + with mock.patch.dict("os.environ", {}, clear=False): + import os + + os.environ.pop(_ASSERT_MIXED_ENV, None) + _assert_mixed_transport_startup( + mode="cpu_serialize", + pipeline_id="p", + train_gpu_ids=None, + infer_gpu_ids=None, + per_engine=1, + ) # no raise + + def test_mixed_overlap_topology_passes(self): + with mock.patch.dict("os.environ", {_ASSERT_MIXED_ENV: "1"}): + _assert_mixed_transport_startup( + mode="auto", pipeline_id="p1", **self.P1 + ) # e0 colocate + e1/e2 broadcast → mixed, no raise + + def test_all_colocate_topology_fails_smoke_invalid(self): + with mock.patch.dict("os.environ", {_ASSERT_MIXED_ENV: "1"}): + with self.assertRaisesRegex(RuntimeError, "INVALID"): + _assert_mixed_transport_startup( + mode="auto", + pipeline_id="p", + train_gpu_ids=[0, 1], + infer_gpu_ids=[0, 1], + per_engine=1, + ) + + def test_non_auto_mode_with_assert_env_fails(self): + with mock.patch.dict("os.environ", {_ASSERT_MIXED_ENV: "1"}): + with self.assertRaisesRegex(RuntimeError, "requires"): + _assert_mixed_transport_startup( + mode="cpu_serialize", pipeline_id="p", **self.P1 + ) + + +def _cfg(models): + return types.SimpleNamespace(models=models) + + +def _model(name, num_gpus_per_engine=None, server_groups=None): + return types.SimpleNamespace( + name=name, + num_gpus_per_engine=num_gpus_per_engine, + server_groups=server_groups, + ) + + +def _group(num_gpus_per_engine=None): + return types.SimpleNamespace(num_gpus_per_engine=num_gpus_per_engine) + + +class TestE10UniformityGuard(unittest.TestCase): + def test_none_config_passes(self): + _assert_uniform_engine_gpu_counts(None, 1) + + def test_matching_values_pass(self): + cfg = _cfg([_model("actor", 2, [_group(2), _group(None)])]) + _assert_uniform_engine_gpu_counts(cfg, 2) + + def test_model_level_divergence_fails_naming_model(self): + cfg = _cfg([_model("actor", 4)]) + with self.assertRaisesRegex(RuntimeError, "'actor'.*num_gpus_per_engine=4"): + _assert_uniform_engine_gpu_counts(cfg, 2) + + def test_group_level_divergence_fails_naming_group(self): + cfg = _cfg([_model("actor", 2, [_group(2), _group(4)])]) + with self.assertRaisesRegex(RuntimeError, r"server_groups\[1\].*num_gpus_per_engine=4"): + _assert_uniform_engine_gpu_counts(cfg, 2) + + def test_unset_values_pass(self): + cfg = _cfg([_model("actor", None, [_group(None)])]) + _assert_uniform_engine_gpu_counts(cfg, 3) + + +if __name__ == "__main__": + unittest.main()