diff --git a/examples/rlix/run_miles_dual.py b/examples/rlix/run_miles_dual.py index f637de0944a..c65a82bc8f6 100644 --- a/examples/rlix/run_miles_dual.py +++ b/examples/rlix/run_miles_dual.py @@ -140,15 +140,40 @@ def _overlap_pools_from_env(num_gpus_per_node: int) -> ( ) if len(set(mapping)) != len(mapping): raise ValueError(f"{label}={mapping} has duplicate GPU ids") - if not set(p1_train).issubset(set(p1_infer)): - raise ValueError( - f"p1_train={p1_train} not ⊆ p1_infer={p1_infer} " - f"(per-pipeline partial-overlap invariant)" - ) - if not set(p2_train).issubset(set(p2_infer)): + # Per-pipeline topology families (rlix#42): + # subset (train ⊆ infer) — classic M11 partial overlap, fully tested; + # disjoint (train ∩ infer == ∅) — dedicated train card(s); no engine + # shares the sender's GPU, so all-broadcast NCCL transport + # is possible. Newly admitted; logged as experimental. + # A PARTIAL intersection (crossing but neither subset nor disjoint) + # stays rejected: the scheduler's donor-shrink/grant accounting for a + # train pool that is half-inside / half-outside the infer pool has no + # test coverage yet — fail fast with the reason rather than wedge a + # run mid-training. + for label_t, train, label_i, infer in ( + ("p1_train", p1_train, "p1_infer", p1_infer), + ("p2_train", p2_train, "p2_infer", p2_infer), + ): + train_set, infer_set = set(train), set(infer) + if train_set.issubset(infer_set): + continue + if not (train_set & infer_set): + import logging as _logging + + _logging.getLogger("run_miles_dual").warning( + "%s=%s is fully DISJOINT from %s=%s — dedicated-train " + "topology (rlix#42): no colocate engines; every engine is " + "NCCL-broadcast-eligible. Newer than the overlap contract; " + "watch the first run.", + label_t, train, label_i, infer, + ) + continue raise ValueError( - f"p2_train={p2_train} not ⊆ p2_infer={p2_infer} " - f"(per-pipeline partial-overlap invariant)" + f"{label_t}={train} partially intersects {label_i}={infer}: " + "each pipeline's train pool must be either a subset of its " + "infer pool (overlap/time-sharing) or fully disjoint from it " + "(dedicated train cards). Mixed shapes are unsupported — the " + "scheduler's shrink/grant accounting for them is unverified." ) return (p1_train, p1_infer), (p2_train, p2_infer) @@ -226,10 +251,16 @@ def _build_pipeline( train_size = len(train_mapping) infer_size = len(infer_mapping) - if not set(train_mapping).issubset(set(infer_mapping)): + # Same two-family rule as _overlap_pools_from_env (rlix#42): subset + # (overlap/time-sharing) or fully-disjoint (dedicated train cards) + # are valid; partial intersections are rejected there before this + # point, so only re-assert the invariant pair here. + _train_set, _infer_set = set(train_mapping), set(infer_mapping) + if not _train_set.issubset(_infer_set) and (_train_set & _infer_set): raise ValueError( - f"mp{pipeline_index}: train_mapping={train_mapping} not ⊆ " - f"infer_mapping={infer_mapping} (per-pipeline partial-overlap)" + f"mp{pipeline_index}: train_mapping={train_mapping} partially " + f"intersects infer_mapping={infer_mapping}; must be a subset " + "(overlap) or fully disjoint (dedicated train cards)" ) args = _per_pipeline_args( base_args, diff --git a/examples/rlix/run_miles_rlix.py b/examples/rlix/run_miles_rlix.py index 681c11f24fb..d14906edb9c 100644 --- a/examples/rlix/run_miles_rlix.py +++ b/examples/rlix/run_miles_rlix.py @@ -37,7 +37,62 @@ def _build_cluster_device_mappings(args) -> dict[str, list[int]]: convention) so train can be a strict subset of infer (partial overlap topology). No new device_mapping CLI args are introduced (Layer 1 forbidden). + + rlix#42 explicit-mapping override: ``MILES_SINGLE_TRAIN_GPUS`` / + ``MILES_SINGLE_INFER_GPUS`` (comma lists of physical GPU ids) replace + the range derivation, admitting the dedicated-train (fully-disjoint) + topology — the minimal all-NCCL-broadcast shape (e.g. train "0", + infer "1,2"). Same two-family invariant as the dual driver: subset + (overlap) or fully disjoint; partial intersections rejected. """ + import os + + train_env = os.environ.get("MILES_SINGLE_TRAIN_GPUS", "").strip() + infer_env = os.environ.get("MILES_SINGLE_INFER_GPUS", "").strip() + if train_env or infer_env: + if not (train_env and infer_env): + raise ValueError( + "MILES_SINGLE_TRAIN_GPUS and MILES_SINGLE_INFER_GPUS must be " + "set together (comma lists of physical GPU ids)" + ) + train = [int(g) for g in train_env.split(",") if g.strip() != ""] + infer = [int(g) for g in infer_env.split(",") if g.strip() != ""] + for label, mapping in (("train", train), ("infer", infer)): + if len(set(mapping)) != len(mapping): + raise ValueError(f"MILES_SINGLE_{label.upper()}_GPUS has duplicates: {mapping}") + train_set, infer_set = set(train), set(infer) + if not train_set.issubset(infer_set) and (train_set & infer_set): + raise ValueError( + f"train={train} partially intersects infer={infer}: must be a " + "subset (overlap) or fully disjoint (dedicated train cards)" + ) + # The mapping lengths MUST match the CLI counts that still size + # the actual actors (RayTrainGroup uses actor_num_nodes × + # actor_num_gpus_per_node; Phase B uses rollout_num_gpus) — a + # divergence would fail late in scheduler/placement work or + # silently allocate an unintended topology (codex impl-r9). + expected_train = int(args.actor_num_nodes) * int(args.actor_num_gpus_per_node) + expected_infer = int(args.rollout_num_gpus) + if len(train) != expected_train or len(infer) != expected_infer: + raise ValueError( + f"MILES_SINGLE_TRAIN_GPUS={train} / MILES_SINGLE_INFER_GPUS={infer} " + f"lengths must match the CLI-derived worker counts: expected " + f"len(train)=={expected_train} (actor_num_nodes × " + f"actor_num_gpus_per_node) and len(infer)=={expected_infer} " + f"(rollout_num_gpus); got {len(train)} / {len(infer)}" + ) + if not (train_set & infer_set): + import logging as _logging + + _logging.getLogger("run_miles_rlix").warning( + "train=%s is fully DISJOINT from infer=%s — dedicated-train " + "topology (rlix#42): no colocate engines; every engine is " + "NCCL-broadcast-eligible.", + train, + infer, + ) + return {"actor_train": train, "actor_infer": infer} + actor_count = int(args.actor_num_nodes) * int(args.actor_num_gpus_per_node) rollout_count = int(args.rollout_num_gpus) return { diff --git a/miles/backends/megatron_utils/actor.py b/miles/backends/megatron_utils/actor.py index 93e6e330f56..03ce11dfef3 100644 --- a/miles/backends/megatron_utils/actor.py +++ b/miles/backends/megatron_utils/actor.py @@ -898,6 +898,7 @@ def run_sync_session(self, plan) -> int: master_port=master_port, comm_ranks=comm_ranks, world_size=world_size, + timeout_s=float(plan["timeout_s"]), ) return version @@ -949,80 +950,28 @@ def _dispatch_nccl_broadcast( master_port: int, comm_ranks: dict[int, int], world_size: int, + timeout_s: float = 0.0, ) -> None: """In-method helper: dynamic NCCL broadcast of bucket payloads. - **MILES-side self-guard (cross-cutting review P1-8)**: the sender- - side ``init_process_group`` + per-bucket ``dist.broadcast`` + - ``dist.destroy_process_group`` is NOT yet wired here. The - receiver-side fan-out below (``setup_collective_group`` + - ``broadcast_parameter`` + ``destroy_collective_group``) would - block forever on ``init_weights_update_group`` waiting for an - absent rank-0 sender. Until the sender path lands, refuse to - even attempt the receiver-side setup so MILES does not depend - on the RLix-side service guard for safety. The MILES contract - is "zero RLix import dependency"; MILES must self-guard. + Sender-side NCCL (rlops/rlix#42): the cache_owner joins the dynamic + group as rank 0 and broadcasts each bucket tensor to the SGLang + receiver engines. Delegates to the module-level + :func:`_run_sender_broadcast` so the transport logic is + unit-testable without constructing a full Megatron actor; the + function is a private helper, NOT a Ray RPC (scope F04). """ - if not target_handles: - return - raise NotImplementedError( - "broadcast transport requires sender-side NCCL " - "(init_process_group + dist.broadcast on the cache_owner). " - "Until the sender-side path lands, plans must route every " - "target through cpu_serialize. MilesModelUpdateService " - "already raises in iter 19/20; this MILES-side guard makes " - "the same invariant explicit at the receiver fan-out." - ) - if world_size <= 0: - raise ValueError( - f"_dispatch_nccl_broadcast requires world_size > 0; got {world_size}" - ) - # Receiver-side group create with the same world_size value the - # sender uses (cache_owner == rank 0, plus one entry per - # receiver engine that participates in the broadcast). - ray.get( - [ - handle.setup_collective_group.remote( - group_name=group_name, - master_addr=master_addr, - master_port=master_port, - rank=comm_ranks[engine_index], - world_size=int(world_size), - ) - for engine_index, handle in target_handles.items() - ] + _run_sender_broadcast( + sync_id=sync_id, + buckets=buckets, + target_handles=target_handles, + group_name=group_name, + master_addr=master_addr, + master_port=master_port, + comm_ranks=comm_ranks, + world_size=world_size, + timeout_s=timeout_s, ) - try: - for bucket in buckets: - # Per-bucket metadata for SGLang's - # update_weights_from_distributed admin route. - names: list[str] = [] - dtypes: list[str] = [] - shapes: list[list[int]] = [] - for name, tensor in bucket.params.items(): - names.append(name) - dtypes.append(str(tensor.dtype).replace("torch.", "")) - shapes.append(list(tensor.shape)) - ray.get( - [ - handle.broadcast_parameter.remote( - sync_id=sync_id, - bucket_index=int(bucket.bucket_index), - group_name=group_name, - names=names, - dtypes=dtypes, - shapes=shapes, - ) - for handle in target_handles.values() - ] - ) - finally: - ray.get( - [ - handle.destroy_collective_group.remote(group_name=group_name) - for handle in target_handles.values() - ] - ) def load_other_checkpoint(self, model_tag: str, path: str) -> None: old_args = self.args.load, self.args.no_load_optim, self.args.no_load_rng, self.args.finetune @@ -1072,3 +1021,325 @@ def connect_actor_critic( rank=0 if self.role == "actor" else 1, group_name=group_name, ) + + +# ---------------------------------------------------------------------- +# F4 Path B sender-side NCCL broadcast helpers (rlops/rlix#42). +# +# Module-level so the transport logic is unit-testable without a full +# MegatronTrainRayActor (megatron init). These are private helpers of +# run_sync_session's dispatch, NOT Ray RPCs (scope F04). The nested +# receiver refs issued here are OWNED here: the finally block cancels +# them and destroys the group on both sides on every exit path — the +# RLix service's inflight_refs never covers them (plan C6 enforcement +# split). +# ---------------------------------------------------------------------- + +_BROADCAST_STAGING_MARGIN_ENV = "MILES_BROADCAST_STAGING_MARGIN_GB" +# Fallback budgets when the plan carries no positive timeout_s (the RLix +# service's asyncio.wait_for is disabled): rendezvous / transport / +# teardown, in seconds. +_BROADCAST_FALLBACK_BUDGETS_S = (30.0, 90.0, 15.0) + + +def _broadcast_budgets(timeout_s: float) -> tuple[float, float, float]: + """C6 deadline hierarchy: (rendezvous, transport, teardown_grace). + + Fractions of the session deadline chosen so the sender's worst case + completes strictly inside the RLix service's + ``asyncio.wait_for(timeout_s)`` — the service deadline then only + fires for a truly wedged sender (fail-fast path), never as a normal + cleanup path. + + Enforcement layers for the native NCCL work (codex impl-r1 high): + the rendezvous budget doubles as the process-group timeout, so the + NCCL watchdog bounds EACH individual collective; a cumulative + monotonic deadline (rendezvous + transport) is checked between + tensors/buckets in :func:`_run_sender_broadcast`. Worst-case unwind + is therefore rendezvous + transport + one in-flight collective + + teardown = (0.15 + 0.5 + 0.15 + 0.1) x timeout_s = 0.9 x timeout_s + < the session deadline, with 0.1 x margin. + """ + if timeout_s and timeout_s > 0: + return 0.15 * timeout_s, 0.5 * timeout_s, 0.1 * timeout_s + return _BROADCAST_FALLBACK_BUDGETS_S + + +def _ensure_nccl_async_error_handling() -> None: + """The per-collective hard bound (C6) is enforced by the NCCL + watchdog, which async error handling enables. An explicitly DISABLED + setting would silently reintroduce unbounded native collectives + (codex impl-r2 high), so fail fast on it instead of overwriting the + operator's choice; when unset, pin the torch default explicitly. + Checks the legacy ``NCCL_ASYNC_ERROR_HANDLING`` alias too (older + torch reads it as a fallback).""" + import os as _os + + for var in ("TORCH_NCCL_ASYNC_ERROR_HANDLING", "NCCL_ASYNC_ERROR_HANDLING"): + raw = _os.environ.get(var) + if raw is not None and raw.strip() == "0": + raise RuntimeError( + f"{var}=0 disables the NCCL watchdog, which the broadcast " + "transport's per-collective hard bound (plan C6) depends on; " + "a blocked dist.broadcast would then hang the cache_owner " + "past the session deadline. Unset it or set it to a non-zero " + "handling mode before enabling broadcast transport." + ) + _os.environ.setdefault("TORCH_NCCL_ASYNC_ERROR_HANDLING", "1") + + +def _monotonic() -> float: + """Seam for unit tests to drive the transport deadline clock.""" + import time as _time + + return _time.monotonic() + + +def _check_transport_deadline(deadline: float, sync_id: str, bucket_index: int) -> None: + """Cumulative sender-side transport deadline check, run between + collectives (a blocked in-flight collective is bounded separately by + the process-group timeout / NCCL watchdog).""" + if _monotonic() > deadline: + raise TimeoutError( + f"_run_sender_broadcast sync_id={sync_id} bucket={bucket_index}: " + "cumulative transport deadline exceeded; aborting before the next " + "collective (C6 sender-side bound)" + ) + + +def _staging_margin_bytes() -> int: + """Free-VRAM safety margin required on top of a bucket before + whole-bucket staging (default 1 GiB, env-overridable).""" + import os as _os + + raw = _os.environ.get(_BROADCAST_STAGING_MARGIN_ENV, "") + if not raw: + return 1024**3 + try: + margin_gb = float(raw) + except ValueError as exc: + raise ValueError( + f"{_BROADCAST_STAGING_MARGIN_ENV} must be a float, got {raw!r}" + ) from exc + if margin_gb < 0: + raise ValueError( + f"{_BROADCAST_STAGING_MARGIN_ENV} must be >= 0, got {margin_gb}" + ) + return int(margin_gb * 1024**3) + + +def _plan_bucket_staging(bucket_size_bytes: int, free_bytes: int, margin_bytes: int) -> str: + """Memory preflight decision: ``"bucket"`` = stage the whole bucket on + GPU at once (batching optimization); ``"tensor"`` = degrade to + tensor-by-tensor staging. Degradation changes batching only — the + broadcast sequence and metadata order are identical either way. + """ + if free_bytes >= bucket_size_bytes + margin_bytes: + return "bucket" + return "tensor" + + +def _resolve_staging_device(): + """The cache_owner's CUDA device used to stage CPU bucket tensors for + NCCL. Separated out so unit tests can patch it to a CPU device.""" + return torch.device("cuda", torch.cuda.current_device()) + + +def _query_free_bytes(device) -> int: + """Free bytes on ``device`` per the CUDA allocator. Separated out for + unit-test patching.""" + free_bytes, _total = torch.cuda.mem_get_info(device) + return int(free_bytes) + + +def _run_sender_broadcast( + *, + sync_id: str, + buckets, + target_handles: dict[int, Any], + group_name: str, + master_addr: str, + master_port: int, + comm_ranks: dict[int, int], + world_size: int, + timeout_s: float, +) -> None: + """Sender-side dynamic NCCL broadcast (F4 Path B). + + Deadlock-safe ordering (mirrors the proven standalone + ``UpdateWeightFromDistributed`` pattern; plan R1/A6): + + 1. dispatch receiver ``setup_collective_group`` refs async — + no ``ray.get`` yet (receivers block in + ``/init_weights_update_group`` until rank 0 joins); + 2. sender joins as rank 0 with a bounded rendezvous timeout, so a + receiver that dies / rejects / hangs mid-rendezvous cannot + wedge rank 0 (E7); + 3. ``ray.get`` the setup refs — a receiver-side setup failure + surfaces here and takes the abort path; + 4. per bucket: dispatch ``broadcast_parameter`` metadata refs + async, memory-preflight the staging (E8), ``dist.broadcast`` + each tensor in metadata order, then ``ray.get`` the refs; + 5. ``finally``: sender-owned teardown on success, exception, and + budget expiry alike — cancel outstanding nested refs, receiver + ``destroy_collective_group`` fan-out (400-tolerant on the + engine side), sender ``destroy_process_group``. + """ + if not target_handles: + return + if world_size <= 0: + raise ValueError( + f"_run_sender_broadcast requires world_size > 0; got {world_size}" + ) + missing = sorted(i for i in target_handles if i not in comm_ranks) + if missing: + raise KeyError( + f"_run_sender_broadcast: comm_ranks missing engine indices {missing}" + ) + + import datetime as _datetime + + rendezvous_budget_s, transport_budget_s, teardown_grace_s = _broadcast_budgets( + float(timeout_s) + ) + # The process-group timeout below is the per-collective hard bound; + # it is enforced by the NCCL watchdog. Fail fast if async error + # handling has been explicitly disabled (must run before the group + # is constructed). + _ensure_nccl_async_error_handling() + transport_deadline = _monotonic() + rendezvous_budget_s + transport_budget_s + nested_refs: list = [] + sender_group = None + try: + setup_refs = [ + handle.setup_collective_group.remote( + group_name=group_name, + master_addr=master_addr, + master_port=int(master_port), + rank=int(comm_ranks[engine_index]), + world_size=int(world_size), + ) + for engine_index, handle in sorted(target_handles.items()) + ] + nested_refs.extend(setup_refs) + # Sender join BEFORE getting the setup refs: init_process_group + # blocks until all ranks join and the receiver HTTP routes block + # until rank 0 shows up — a ray.get here first would deadlock. + # The timeout doubles as the group's per-collective bound (see + # _broadcast_budgets). + sender_group = init_process_group( + backend="nccl", + init_method=f"tcp://{master_addr}:{int(master_port)}", + world_size=int(world_size), + rank=0, + group_name=group_name, + timeout=_datetime.timedelta(seconds=rendezvous_budget_s), + ) + ray.get(setup_refs, timeout=rendezvous_budget_s) + + device = _resolve_staging_device() + margin_bytes = _staging_margin_bytes() + for bucket in buckets: + names: list[str] = [] + dtypes: list[str] = [] + shapes: list[list[int]] = [] + for name, tensor in bucket.params.items(): + names.append(name) + dtypes.append(str(tensor.dtype).replace("torch.", "")) + shapes.append(list(tensor.shape)) + bucket_refs = [ + handle.broadcast_parameter.remote( + sync_id=sync_id, + bucket_index=int(bucket.bucket_index), + group_name=group_name, + names=names, + dtypes=dtypes, + shapes=shapes, + ) + for _engine_index, handle in sorted(target_handles.items()) + ] + nested_refs.extend(bucket_refs) + + free_bytes = _query_free_bytes(device) + staging_mode = _plan_bucket_staging( + int(bucket.size_bytes), free_bytes, margin_bytes + ) + if staging_mode == "bucket": + staged = { + name: tensor.to(device) for name, tensor in bucket.params.items() + } + for name in names: + _check_transport_deadline( + transport_deadline, sync_id, int(bucket.bucket_index) + ) + dist.broadcast(staged[name], 0, group=sender_group) + staged.clear() + else: + logger.warning( + "_run_sender_broadcast sync_id=%s bucket=%s: free VRAM %d B < " + "bucket %d B + margin %d B; degrading to tensor-by-tensor staging", + sync_id, + bucket.bucket_index, + free_bytes, + bucket.size_bytes, + margin_bytes, + ) + for name, tensor in bucket.params.items(): + _check_transport_deadline( + transport_deadline, sync_id, int(bucket.bucket_index) + ) + gpu_tensor = tensor.to(device) + dist.broadcast(gpu_tensor, 0, group=sender_group) + del gpu_tensor + # Cumulative budget: the receivers get whatever remains of the + # transport window, not a fresh per-bucket allowance. + remaining_s = transport_deadline - _monotonic() + if remaining_s <= 0: + raise TimeoutError( + f"_run_sender_broadcast sync_id={sync_id} " + f"bucket={bucket.bucket_index}: transport deadline exhausted " + "before receiver bucket acks" + ) + ray.get(bucket_refs, timeout=remaining_s) + finally: + # Sender-owned teardown (plan C6): cancel outstanding nested refs + # first so receivers blocked in HTTP routes cannot hold the group + # alive, then destroy both sides. Completed refs cancel as no-ops. + for ref in nested_refs: + try: + ray.cancel(ref) + except Exception as exc: # noqa: BLE001 + logger.debug("_run_sender_broadcast: ray.cancel failed: %r", exc) + destroy_refs = [] + for _engine_index, handle in sorted(target_handles.items()): + try: + destroy_refs.append( + handle.destroy_collective_group.remote(group_name=group_name) + ) + except Exception as exc: # noqa: BLE001 + logger.warning( + "_run_sender_broadcast sync_id=%s: destroy dispatch failed: %r", + sync_id, + exc, + ) + if destroy_refs: + try: + ray.get(destroy_refs, timeout=teardown_grace_s) + except Exception as exc: # noqa: BLE001 + logger.warning( + "_run_sender_broadcast sync_id=%s: receiver group destroy " + "incomplete within %.1fs: %r", + sync_id, + teardown_grace_s, + exc, + ) + if sender_group is not None: + try: + dist.destroy_process_group(sender_group) + except Exception as exc: # noqa: BLE001 + logger.warning( + "_run_sender_broadcast sync_id=%s: sender destroy_process_group " + "failed: %r", + sync_id, + exc, + ) diff --git a/miles/ray/rollout.py b/miles/ray/rollout.py index 19fb0dbf7b3..843b7b63aaa 100644 --- a/miles/ray/rollout.py +++ b/miles/ray/rollout.py @@ -906,6 +906,33 @@ def _reset_abort_idempotency_for(self, engine_indices: Iterable[int]) -> None: # so callers document intent. self._release_abort_idempotency_for(engine_indices) + def register_router_if_active(self, engine_indices: Iterable[int]) -> list[int]: + """rlix#42 sync-under-load bracket: atomically re-open router + admission for engines that are STILL ``active`` at this exact + moment, skipping every other state. + + Runs on the manager's single-threaded execution (plain + ``@ray.remote`` — max_concurrency 1), so it is serialized against + ``shrink_engines``: the TOCTOU between a coordinator-side state + read and a separate per-engine ``/add_worker`` cannot be + interleaved by a shrink that starts in between (codex impl-r13). + ``register_with_router`` is idempotent at the router and raises + on non-2xx. Returns the indices actually re-registered. + + Deliberately does NOT use ``_resolve_engine_indices``: that + helper raises for non-alive engines, but skipping a + concurrently-shrunk (offloaded/disabling) engine is exactly this + method's contract — silent skip, reported via the return value. + """ + registered: list[int] = [] + for idx in sorted(set(int(i) for i in engine_indices)): + info = self._engines.get(idx) + if info is None or info.state != "active": + continue + ray.get(info.handle.register_with_router.remote()) + registered.append(idx) + return registered + def get_router_enabled_workers(self) -> list[str]: """M11.2 Option β 3e: snapshot the router's ``enabled_workers`` set. diff --git a/tests/test_nccl_broadcast_dispatch.py b/tests/test_nccl_broadcast_dispatch.py new file mode 100644 index 00000000000..8a106f4d49d --- /dev/null +++ b/tests/test_nccl_broadcast_dispatch.py @@ -0,0 +1,381 @@ +"""Unit tests for the sender-side NCCL broadcast dispatch (rlops/rlix#42). + +Covers plan miles-nccl-broadcast v7 evidence items: + +- E1: dispatch ordering — receiver setup refs BEFORE sender join; metadata + refs BEFORE tensor broadcasts; sender-owned teardown in ``finally`` on + success AND on mid-bucket exception. +- E7: receiver-fails-during-setup — sender aborts bounded (setup ray.get + failure or rendezvous timeout), still runs teardown, surfaces the error. +- E8: memory-preflight degradation — low free VRAM switches staging to + tensor-by-tensor with the broadcast sequence/metadata order unchanged. + +All Ray / torch.distributed / CUDA interactions are patched at the +``actor`` module namespace; ``_run_sender_broadcast`` is exercised +directly (module-level helper — no MegatronTrainRayActor construction). +""" + +import types +import unittest +from unittest import mock + +import torch + +import miles.backends.megatron_utils.actor as actor_mod +from miles.backends.megatron_utils.update_weight.cpu_bucket_cache import BucketEntry + + +def _make_bucket(idx: int = 0, n_tensors: int = 2, numel: int = 4) -> BucketEntry: + params = {f"w{idx}_{i}": torch.zeros(numel) for i in range(n_tensors)} + size = sum(t.element_size() * t.numel() for t in params.values()) + return BucketEntry( + bucket_index=idx, + params=params, + size_bytes=size, + element_count=n_tensors * numel, + ) + + +class _FakeRef: + def __init__(self, kind: str, engine_index: int, kwargs: dict): + self.kind = kind + self.engine_index = engine_index + self.kwargs = kwargs + + +class _FakeRemoteMethod: + def __init__(self, log: list, kind: str, engine_index: int): + self._log = log + self._kind = kind + self._engine_index = engine_index + + def remote(self, **kwargs): + self._log.append((f"{self._kind}.remote", self._engine_index)) + return _FakeRef(self._kind, self._engine_index, kwargs) + + +class _FakeHandle: + def __init__(self, log: list, engine_index: int): + self.setup_collective_group = _FakeRemoteMethod(log, "setup", engine_index) + self.broadcast_parameter = _FakeRemoteMethod(log, "metadata", engine_index) + self.destroy_collective_group = _FakeRemoteMethod(log, "destroy", engine_index) + + +class _Harness: + """Patches ray/dist/init_process_group/staging seams in the actor + module namespace and records a global event log.""" + + def __init__(self, fail_get_kinds: set | None = None, join_error: Exception | None = None): + self.log: list = [] + self.join_calls: list = [] + self.get_refs: list = [] + self._fail_get_kinds = fail_get_kinds or set() + self._join_error = join_error + self.sender_group = object() + + # --- fakes ------------------------------------------------------- + def ray_get(self, refs, timeout=None): + kinds = tuple(sorted({r.kind for r in refs})) + self.log.append(("ray.get", kinds, timeout)) + self.get_refs.extend(refs) + for kind in kinds: + if kind in self._fail_get_kinds: + raise RuntimeError(f"injected {kind} failure") + return [None for _ in refs] + + def ray_cancel(self, ref): + self.log.append(("ray.cancel", ref.kind, ref.engine_index)) + + def dist_broadcast(self, tensor, src, group=None): + assert src == 0 + assert group is self.sender_group + self.log.append(("broadcast", int(tensor.numel()))) + + def dist_destroy(self, group): + assert group is self.sender_group + self.log.append(("sender_destroy",)) + + def join(self, **kwargs): + self.log.append(("sender_join",)) + self.join_calls.append(kwargs) + if self._join_error is not None: + raise self._join_error + return self.sender_group + + # --- helpers ----------------------------------------------------- + def patches(self, free_bytes: int = 10**12): + fake_ray = types.SimpleNamespace(get=self.ray_get, cancel=self.ray_cancel) + fake_dist = types.SimpleNamespace( + broadcast=self.dist_broadcast, destroy_process_group=self.dist_destroy + ) + return [ + mock.patch.object(actor_mod, "ray", fake_ray), + mock.patch.object(actor_mod, "dist", fake_dist), + mock.patch.object(actor_mod, "init_process_group", self.join), + mock.patch.object( + actor_mod, "_resolve_staging_device", lambda: torch.device("cpu") + ), + mock.patch.object(actor_mod, "_query_free_bytes", lambda device: free_bytes), + ] + + def run(self, *, buckets, handles, free_bytes: int = 10**12, timeout_s: float = 150.0): + patches = self.patches(free_bytes=free_bytes) + for p in patches: + p.start() + try: + actor_mod._run_sender_broadcast( + sync_id="sync-test", + buckets=buckets, + target_handles=handles, + group_name="grp", + master_addr="10.0.0.1", + master_port=29500, + comm_ranks={i: i + 1 for i in handles}, + world_size=1 + len(handles), + timeout_s=timeout_s, + ) + finally: + for p in patches: + p.stop() + + # --- log queries ------------------------------------------------- + def indices(self, predicate): + return [i for i, ev in enumerate(self.log) if predicate(ev)] + + def first(self, predicate): + idx = self.indices(predicate) + assert idx, f"no matching event in log: {self.log}" + return idx[0] + + def last(self, predicate): + idx = self.indices(predicate) + assert idx, f"no matching event in log: {self.log}" + return idx[-1] + + +class TestE1DispatchOrdering(unittest.TestCase): + def test_happy_path_ordering(self): + h = _Harness() + handles = {0: _FakeHandle(h.log, 0), 1: _FakeHandle(h.log, 1)} + h.run(buckets=[_make_bucket(0), _make_bucket(1)], handles=handles) + + setup_last = h.last(lambda ev: ev[0] == "setup.remote") + join_idx = h.first(lambda ev: ev[0] == "sender_join") + setup_get = h.first(lambda ev: ev[0] == "ray.get" and ev[1] == ("setup",)) + # (1) every receiver setup dispatched before the sender joins; + # (2) sender joins before the setup refs are awaited. + self.assertLess(setup_last, join_idx) + self.assertLess(join_idx, setup_get) + + # (3) per bucket: metadata refs dispatched before the first tensor + # broadcast of that bucket; bucket refs awaited after broadcasts. + first_meta = h.first(lambda ev: ev[0] == "metadata.remote") + first_bcast = h.first(lambda ev: ev[0] == "broadcast") + self.assertLess(first_meta, first_bcast) + # 2 buckets x 2 tensors = 4 tensor broadcasts, once each (to the + # group — not per receiver). + self.assertEqual(len(h.indices(lambda ev: ev[0] == "broadcast")), 4) + + # (4) teardown ran at the end: receiver destroy fan-out + sender + # destroy after the last broadcast. + last_bcast = h.last(lambda ev: ev[0] == "broadcast") + destroy_first = h.first(lambda ev: ev[0] == "destroy.remote") + sender_destroy = h.first(lambda ev: ev[0] == "sender_destroy") + self.assertLess(last_bcast, destroy_first) + self.assertLess(destroy_first, sender_destroy) + + def test_teardown_on_mid_bucket_receiver_exception(self): + h = _Harness(fail_get_kinds={"metadata"}) + handles = {0: _FakeHandle(h.log, 0)} + with self.assertRaisesRegex(RuntimeError, "injected metadata failure"): + h.run(buckets=[_make_bucket(0)], handles=handles) + # Sender-owned finally still cancelled refs and destroyed both sides. + self.assertTrue(h.indices(lambda ev: ev[0] == "ray.cancel")) + self.assertTrue(h.indices(lambda ev: ev[0] == "destroy.remote")) + self.assertTrue(h.indices(lambda ev: ev[0] == "sender_destroy")) + + def test_empty_target_handles_is_noop(self): + h = _Harness() + h.run(buckets=[_make_bucket(0)], handles={}) + self.assertEqual(h.log, []) + + def test_invalid_world_size_raises(self): + h = _Harness() + with self.assertRaises(ValueError): + with mock.patch.object(actor_mod, "ray"), mock.patch.object(actor_mod, "dist"): + actor_mod._run_sender_broadcast( + sync_id="s", + buckets=[], + target_handles={0: _FakeHandle(h.log, 0)}, + group_name="g", + master_addr="a", + master_port=1, + comm_ranks={0: 1}, + world_size=0, + timeout_s=1.0, + ) + + def test_missing_comm_rank_raises(self): + h = _Harness() + with self.assertRaises(KeyError): + actor_mod._run_sender_broadcast( + sync_id="s", + buckets=[], + target_handles={0: _FakeHandle(h.log, 0), 7: _FakeHandle(h.log, 7)}, + group_name="g", + master_addr="a", + master_port=1, + comm_ranks={0: 1}, + world_size=3, + timeout_s=1.0, + ) + + +class TestE7ReceiverFailureAbort(unittest.TestCase): + def test_setup_failure_aborts_and_tears_down(self): + h = _Harness(fail_get_kinds={"setup"}) + handles = {0: _FakeHandle(h.log, 0), 1: _FakeHandle(h.log, 1)} + with self.assertRaisesRegex(RuntimeError, "injected setup failure"): + h.run(buckets=[_make_bucket(0)], handles=handles) + # No tensor was broadcast; teardown still ran on both sides. + self.assertEqual(h.indices(lambda ev: ev[0] == "broadcast"), []) + self.assertTrue(h.indices(lambda ev: ev[0] == "destroy.remote")) + self.assertTrue(h.indices(lambda ev: ev[0] == "sender_destroy")) + + def test_sender_join_timeout_aborts_bounded(self): + h = _Harness(join_error=RuntimeError("rendezvous timed out")) + handles = {0: _FakeHandle(h.log, 0)} + with self.assertRaisesRegex(RuntimeError, "rendezvous timed out"): + h.run(buckets=[_make_bucket(0)], handles=handles) + # Receiver destroy fan-out still dispatched; the sender group was + # never created so sender-side destroy is correctly absent. + self.assertTrue(h.indices(lambda ev: ev[0] == "destroy.remote")) + self.assertEqual(h.indices(lambda ev: ev[0] == "sender_destroy"), []) + + def test_join_timeout_uses_rendezvous_budget(self): + h = _Harness() + handles = {0: _FakeHandle(h.log, 0)} + h.run(buckets=[], handles=handles, timeout_s=150.0) + self.assertEqual(len(h.join_calls), 1) + timeout = h.join_calls[0]["timeout"] + self.assertAlmostEqual(timeout.total_seconds(), 0.15 * 150.0) + + def test_budget_hierarchy_nests_inside_session_deadline(self): + rendezvous, transport, teardown = actor_mod._broadcast_budgets(150.0) + self.assertAlmostEqual(rendezvous, 0.15 * 150.0) + self.assertAlmostEqual(transport, 0.5 * 150.0) + self.assertAlmostEqual(teardown, 0.1 * 150.0) + # C6 invariant: sender worst-case unwind — rendezvous + transport + # + one in-flight collective (bounded by the pg timeout == + # rendezvous budget, NCCL watchdog) + teardown — < session + # deadline. + self.assertLess(rendezvous + transport + rendezvous + teardown, 150.0) + # Fallback budgets when the service timeout is disabled. + self.assertEqual( + actor_mod._broadcast_budgets(0.0), actor_mod._BROADCAST_FALLBACK_BUDGETS_S + ) + + def test_disabled_async_error_handling_refuses_to_run(self): + # codex impl-r2 high: an explicitly disabled NCCL watchdog would + # reintroduce unbounded native collectives — the sender must + # refuse to start rather than silently overwrite the env. + for var in ("TORCH_NCCL_ASYNC_ERROR_HANDLING", "NCCL_ASYNC_ERROR_HANDLING"): + h = _Harness() + handles = {0: _FakeHandle(h.log, 0)} + with mock.patch.dict("os.environ", {var: "0"}): + with self.assertRaisesRegex(RuntimeError, "NCCL watchdog"): + h.run(buckets=[_make_bucket(0)], handles=handles) + # Refused before any rendezvous: no join, no receiver setup get. + self.assertEqual(h.indices(lambda ev: ev[0] == "sender_join"), []) + + def test_unset_async_error_handling_is_pinned_on(self): + h = _Harness() + handles = {0: _FakeHandle(h.log, 0)} + import os + + with mock.patch.dict("os.environ", {}, clear=False): + os.environ.pop("TORCH_NCCL_ASYNC_ERROR_HANDLING", None) + os.environ.pop("NCCL_ASYNC_ERROR_HANDLING", None) + h.run(buckets=[], handles=handles) + self.assertEqual(os.environ.get("TORCH_NCCL_ASYNC_ERROR_HANDLING"), "1") + + def test_nonzero_async_error_handling_values_accepted(self): + h = _Harness() + handles = {0: _FakeHandle(h.log, 0)} + with mock.patch.dict("os.environ", {"TORCH_NCCL_ASYNC_ERROR_HANDLING": "2"}): + h.run(buckets=[], handles=handles) + self.assertTrue(h.indices(lambda ev: ev[0] == "sender_join")) + + def test_stalled_transport_hits_cumulative_deadline_and_tears_down(self): + # Receiver setup succeeds but the transport clock burns past the + # cumulative deadline (e.g. a stalled peer eating watchdog-bounded + # collectives): the sender must abort between collectives and + # still run its finally teardown (codex impl-r1 high). + h = _Harness() + handles = {0: _FakeHandle(h.log, 0)} + clock = {"now": 0.0} + + def _fake_monotonic(): + # Every observation advances the clock far beyond the budget. + clock["now"] += 1000.0 + return clock["now"] + + with mock.patch.object(actor_mod, "_monotonic", _fake_monotonic): + with self.assertRaisesRegex(TimeoutError, "deadline"): + h.run( + buckets=[_make_bucket(0, n_tensors=2)], + handles=handles, + timeout_s=150.0, + ) + self.assertTrue(h.indices(lambda ev: ev[0] == "destroy.remote")) + self.assertTrue(h.indices(lambda ev: ev[0] == "sender_destroy")) + + +class TestE8MemoryPreflight(unittest.TestCase): + def test_staging_plan_pure_decision(self): + self.assertEqual(actor_mod._plan_bucket_staging(100, 200, 50), "bucket") + self.assertEqual(actor_mod._plan_bucket_staging(100, 149, 50), "tensor") + self.assertEqual(actor_mod._plan_bucket_staging(100, 150, 50), "bucket") + + def test_low_memory_degrades_without_changing_order(self): + rich = _Harness() + handles_rich = {0: _FakeHandle(rich.log, 0)} + rich.run(buckets=[_make_bucket(0, n_tensors=3)], handles=handles_rich) + + poor = _Harness() + handles_poor = {0: _FakeHandle(poor.log, 0)} + poor.run( + buckets=[_make_bucket(0, n_tensors=3)], handles=handles_poor, free_bytes=0 + ) + + # Same broadcast sequence either way (3 tensors, same sizes) — + # degradation changes batching only, not wire behavior. + rich_bcasts = [ev for ev in rich.log if ev[0] == "broadcast"] + poor_bcasts = [ev for ev in poor.log if ev[0] == "broadcast"] + self.assertEqual(rich_bcasts, poor_bcasts) + # Metadata names preserve bucket insertion order in both runs. + rich_meta = [r for r in rich.get_refs if r.kind == "metadata"] + poor_meta = [r for r in poor.get_refs if r.kind == "metadata"] + self.assertEqual(rich_meta[0].kwargs["names"], poor_meta[0].kwargs["names"]) + self.assertEqual(rich_meta[0].kwargs["names"], ["w0_0", "w0_1", "w0_2"]) + + def test_staging_margin_env_override(self): + with mock.patch.dict("os.environ", {actor_mod._BROADCAST_STAGING_MARGIN_ENV: "2.5"}): + self.assertEqual(actor_mod._staging_margin_bytes(), int(2.5 * 1024**3)) + with mock.patch.dict("os.environ", {}, clear=False): + import os + + os.environ.pop(actor_mod._BROADCAST_STAGING_MARGIN_ENV, None) + self.assertEqual(actor_mod._staging_margin_bytes(), 1024**3) + with mock.patch.dict( + "os.environ", {actor_mod._BROADCAST_STAGING_MARGIN_ENV: "not-a-float"} + ): + with self.assertRaises(ValueError): + actor_mod._staging_margin_bytes() + with mock.patch.dict("os.environ", {actor_mod._BROADCAST_STAGING_MARGIN_ENV: "-1"}): + with self.assertRaises(ValueError): + actor_mod._staging_margin_bytes() + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test_register_router_if_active.py b/tests/test_register_router_if_active.py new file mode 100644 index 00000000000..b3557103918 --- /dev/null +++ b/tests/test_register_router_if_active.py @@ -0,0 +1,82 @@ +"""Contract test for the REAL RolloutManager.register_router_if_active +(rlix#42 sync-under-load bracket; codex impl-r15 asked for a test against +the production class, not coordinator-side fakes). + +The manager instance is built via ``object.__new__`` with a hand-rolled +``_engines`` table; ``ray.get`` in the rollout module is patched to +unwrap the fakes' eager returns. +""" + +import types +import unittest +from unittest import mock + +import miles.ray.rollout as rollout_mod + + +def _manager_cls(): + cls = rollout_mod.RolloutManager + meta = getattr(cls, "__ray_metadata__", None) + return meta.modified_class if meta is not None else cls + + +class _FakeHandle: + def __init__(self, log, idx, fail=False): + def _register(): + if fail: + raise RuntimeError(f"router add_worker failed for {idx}") + log.append(idx) + return None + + self.register_with_router = types.SimpleNamespace(remote=lambda: _register()) + + +def _info(state, handle): + return types.SimpleNamespace(state=state, handle=handle) + + +class TestRegisterRouterIfActive(unittest.TestCase): + def _manager(self, engines): + mgr = object.__new__(_manager_cls()) + mgr._engines = engines + return mgr + + def test_registers_only_active_engines(self): + log: list = [] + mgr = self._manager( + { + 0: _info("offloaded", _FakeHandle(log, 0)), + 1: _info("active", _FakeHandle(log, 1)), + 2: _info("disabling", _FakeHandle(log, 2)), + 3: _info("active", _FakeHandle(log, 3)), + } + ) + with mock.patch.object(rollout_mod, "ray", types.SimpleNamespace(get=lambda r: r)): + out = mgr.register_router_if_active([0, 1, 2, 3]) + self.assertEqual(out, [1, 3]) + self.assertEqual(sorted(log), [1, 3]) + + def test_unknown_index_is_skipped_not_raised(self): + log: list = [] + mgr = self._manager({1: _info("active", _FakeHandle(log, 1))}) + with mock.patch.object(rollout_mod, "ray", types.SimpleNamespace(get=lambda r: r)): + out = mgr.register_router_if_active([1, 99]) + self.assertEqual(out, [1]) + + def test_empty_probe_returns_empty(self): + # The rlix coordinator uses an empty-list call as a version-skew + # capability probe at registration time — must be a cheap no-op. + mgr = self._manager({}) + out = mgr.register_router_if_active([]) + self.assertEqual(out, []) + + def test_register_failure_propagates(self): + log: list = [] + mgr = self._manager({1: _info("active", _FakeHandle(log, 1, fail=True))}) + with mock.patch.object(rollout_mod, "ray", types.SimpleNamespace(get=lambda r: r)): + with self.assertRaisesRegex(RuntimeError, "add_worker failed"): + mgr.register_router_if_active([1]) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test_single_mapping_override.py b/tests/test_single_mapping_override.py new file mode 100644 index 00000000000..a6940448662 --- /dev/null +++ b/tests/test_single_mapping_override.py @@ -0,0 +1,87 @@ +"""Tests for the rlix#42 single-mode explicit-mapping override +(`MILES_SINGLE_TRAIN_GPUS` / `MILES_SINGLE_INFER_GPUS` in +examples/rlix/run_miles_rlix.py::_build_cluster_device_mappings). + +Covers codex impl-r9: matching, undersized, oversized, duplicate, +partial-intersection, subset, disjoint, half-set, and unset-legacy cases. +""" + +import os +import types +import unittest +from unittest import mock + +os.environ.setdefault("RLIX_CONTROL_PLANE", "rlix") # module import guard + +from examples.rlix.run_miles_rlix import _build_cluster_device_mappings # noqa: E402 + + +def _args(actor_gpus=1, rollout_gpus=2): + return types.SimpleNamespace( + actor_num_nodes=1, + actor_num_gpus_per_node=actor_gpus, + rollout_num_gpus=rollout_gpus, + ) + + +def _env(train=None, infer=None): + env = {} + if train is not None: + env["MILES_SINGLE_TRAIN_GPUS"] = train + if infer is not None: + env["MILES_SINGLE_INFER_GPUS"] = infer + return mock.patch.dict("os.environ", env, clear=False) + + +class TestSingleMappingOverride(unittest.TestCase): + def setUp(self): + for var in ("MILES_SINGLE_TRAIN_GPUS", "MILES_SINGLE_INFER_GPUS"): + os.environ.pop(var, None) + + def test_unset_envs_keep_legacy_range_derivation(self): + out = _build_cluster_device_mappings(_args(1, 2)) + self.assertEqual(out, {"actor_train": [0], "actor_infer": [0, 1]}) + + def test_disjoint_override_accepted(self): + with _env("0", "1,2"): + out = _build_cluster_device_mappings(_args(1, 2)) + self.assertEqual(out, {"actor_train": [0], "actor_infer": [1, 2]}) + + def test_subset_override_accepted(self): + with _env("0", "0,1"): + out = _build_cluster_device_mappings(_args(1, 2)) + self.assertEqual(out, {"actor_train": [0], "actor_infer": [0, 1]}) + + def test_partial_intersection_rejected(self): + with _env("0,1", "1,2"): + with self.assertRaisesRegex(ValueError, "partially intersects"): + _build_cluster_device_mappings(_args(2, 2)) + + def test_undersized_infer_rejected(self): + with _env("0", "1"): + with self.assertRaisesRegex(ValueError, "lengths must match"): + _build_cluster_device_mappings(_args(1, 2)) + + def test_oversized_infer_rejected(self): + with _env("0", "1,2,3"): + with self.assertRaisesRegex(ValueError, "lengths must match"): + _build_cluster_device_mappings(_args(1, 2)) + + def test_train_length_mismatch_rejected(self): + with _env("0,3", "1,2"): + with self.assertRaisesRegex(ValueError, "lengths must match"): + _build_cluster_device_mappings(_args(1, 2)) + + def test_duplicate_ids_rejected(self): + with _env("0", "1,1"): + with self.assertRaisesRegex(ValueError, "duplicates"): + _build_cluster_device_mappings(_args(1, 2)) + + def test_half_set_envs_rejected(self): + with _env(train="0"): + with self.assertRaisesRegex(ValueError, "must be .*set together"): + _build_cluster_device_mappings(_args(1, 2)) + + +if __name__ == "__main__": + unittest.main()