From cca0ffaf02be6dce467573146ab96ab3bc97745b Mon Sep 17 00:00:00 2001 From: cytsoi <154662279+tcyjohn@users.noreply.github.com> Date: Thu, 22 Oct 2026 01:33:43 +0800 Subject: [PATCH 1/9] fix: publish weights to current physical rollout services Drain admitted requests before installation and reopen only after physical TP worker receipts and prefix-cache reset. Give each transfer a fresh snapshot identity and actual optimizer revision, while renewing NCCL sessions only when receiver membership changes. Keep training ranks in a common export failure protocol, fence NCCL stream dependencies, and ignore stale startup commands. Omit the unused KV reservation and migration APIs. Validation: 59 distinct targeted CPU cases passed across publication, serving, NCCL lifecycle, and two-rank collective export tests. --- engine/megatron_core_train_engine.py | 2 + infra/sync/collective_export.py | 98 ++++++ infra/sync/nccl_weight_sync.py | 101 +++++- infra/sync/prefix_cache_reset.py | 40 +++ infra/sync/publication_session.py | 37 +++ pyproject.toml | 1 + requirements-ci.txt | 2 + scripts/restartable_vllm_server.py | 46 ++- scripts/serving_protocol.py | 157 +++++++++ scripts/vllm_hot_reload_api_server.py | 87 +++-- tests/test_checked_prefix_cache_reset.py | 52 +++ tests/test_collective_weight_export.py | 128 ++++++++ tests/test_nccl_store_lifecycle.py | 147 +++++++++ tests/test_nccl_store_timeout.py | 54 ++++ tests/test_nccl_weight_sync.py | 132 ++++++++ tests/test_publication_sessions.py | 89 ++++++ tests/test_restartable_vllm_server.py | 82 +++++ tests/test_serving_protocol.py | 386 +++++++++++++++++++++++ trainer/async_rl_trainer.py | 146 +++++---- vllm_hot_reload.py | 34 +- 20 files changed, 1720 insertions(+), 101 deletions(-) create mode 100644 infra/sync/collective_export.py create mode 100644 infra/sync/prefix_cache_reset.py create mode 100644 infra/sync/publication_session.py create mode 100644 scripts/serving_protocol.py create mode 100644 tests/test_checked_prefix_cache_reset.py create mode 100644 tests/test_collective_weight_export.py create mode 100644 tests/test_nccl_store_lifecycle.py create mode 100644 tests/test_nccl_store_timeout.py create mode 100644 tests/test_publication_sessions.py create mode 100644 tests/test_serving_protocol.py diff --git a/engine/megatron_core_train_engine.py b/engine/megatron_core_train_engine.py index 2b83444..444ca9a 100644 --- a/engine/megatron_core_train_engine.py +++ b/engine/megatron_core_train_engine.py @@ -104,6 +104,7 @@ def __init__( self.world_size = int(os.getenv("WORLD_SIZE", "1")) self.is_main_process = self.rank == 0 self.current_version = 0 + self.optimizer_revision = 0 self.bridge = None self.provider = None @@ -533,6 +534,7 @@ def grpo_update( update_successful, grad_norm, _ = self.optimizer.step() if not update_successful: raise FloatingPointError("Megatron optimizer rejected the update") + self.optimizer_revision += 1 merged = self._merge_stats(local_stats) if isinstance(grad_norm, torch.Tensor): merged["grad_norm"] = float(grad_norm.detach().item()) diff --git a/infra/sync/collective_export.py b/infra/sync/collective_export.py new file mode 100644 index 0000000..ad11273 --- /dev/null +++ b/infra/sync/collective_export.py @@ -0,0 +1,98 @@ +"""Keep every training rank in the same weight-export iterator protocol.""" +from __future__ import annotations + +import torch.distributed as dist + + +def _gather(value, group): + if not dist.is_available() or not dist.is_initialized(): + return [value] + values = [None] * dist.get_world_size(group) + dist.all_gather_object(values, value, group=group) + return values + + +class _ExportIterator: + """One control checkpoint per yielded tensor, before the next TP export. + + An iterator exception aborts this publication on every rank. No rank may + advance to the next tensor after a peer failed. An exception *inside* a + tensor's distributed collective remains a process-group failure and is + subject to that group's timeout/abort; it cannot be repaired by draining. + """ + + def __init__(self, factory, group, prepare): + self.group, self.prepare, self.finished = group, prepare, False + self.iterator = None + error = "" + try: + self.iterator = iter(factory()) + except BaseException as exc: + error = f"{type(exc).__name__}: {exc}" + errors = _gather(error, group) + if any(errors): + self.finished = True + raise RuntimeError(f"collective weight export initialization failed: {errors}") + + def __iter__(self): + return self + + def __next__(self): + if self.finished: + raise StopIteration + item, status = None, {"done": False, "error": "", "name": None} + try: + item = next(self.iterator) + if self.prepare is not None: + item = self.prepare(item) + status["name"] = str(item[0]) + except StopIteration: + status["done"] = True + except BaseException as exc: + status["error"] = f"{type(exc).__name__}: {exc}" + statuses = _gather(status, self.group) + errors = [state["error"] for state in statuses if state["error"]] + if errors: + self.finished = True + raise RuntimeError(f"collective weight export iterator aborted: {errors}") + if any(state["done"] for state in statuses): + self.finished = True + if not all(state["done"] for state in statuses): + raise RuntimeError("collective weight export ended at different tensor positions") + raise StopIteration + if len({state["name"] for state in statuses}) != 1: + self.finished = True + raise RuntimeError("collective weight export tensor identities differ across ranks") + return item + + +def stream_export_with_consensus(factory, *, is_sender, send, group=None, prepare=None): + """Drain training collectives after a transport failure, then fail together. + + ``send`` runs only on the sender and must consume the supplied iterator. + A transport error leaves that iterator alive; the sender drains the same + stream as its peers without sending any further weights. Initialization, + preparation, iteration and final transport status all have common control + positions, so an iterator error never silently becomes successful EOF. + """ + weights = _ExportIterator(factory, group, prepare) + error = "" + try: + if is_sender: + send(weights) + else: + for _name, tensor in weights: + del tensor + except BaseException as exc: + error = f"{type(exc).__name__}: {exc}" + # Also handle a sender which returned without consuming the entire stream. + # Draining is safe after transport errors, and a failed iterator marks + # itself finished so it cannot enter another TP collective. + try: + for _name, tensor in weights: + del tensor + except BaseException as exc: + error += f"; export drain failed: {type(exc).__name__}: {exc}" + errors = _gather(error, group) + if any(errors): + raise RuntimeError(f"collective weight publication failed: {errors}") diff --git a/infra/sync/nccl_weight_sync.py b/infra/sync/nccl_weight_sync.py index 1844ce3..0802077 100644 --- a/infra/sync/nccl_weight_sync.py +++ b/infra/sync/nccl_weight_sync.py @@ -3,18 +3,27 @@ from __future__ import annotations import os +import ctypes import ipaddress +import math import threading import time from dataclasses import dataclass +from datetime import timedelta from typing import Callable, Iterable, Iterator import torch -_COMMUNICATOR_CACHE: dict[tuple[str, int, int, int, str], object] = {} +_COMMUNICATOR_CACHE: dict[tuple[str, int, int, int, str, str], object] = {} _COMMUNICATOR_CACHE_LOCK = threading.Lock() +_BASE_STORE_CACHE: dict[tuple[str, int, bool], object] = {} +_BASE_STORE_CACHE_LOCK = threading.Lock() _PIPELINE_DEPTH = 2 +# vLLM constructs datetime.timedelta(seconds=store_timeout), so None/inf and +# timedelta.max are unsafe. A century represents the application's unbounded +# wait while staying below signed int64 even after conversion to nanoseconds. +_UNBOUNDED_STORE_TIMEOUT_S = 100 * 365 * 24 * 60 * 60 @dataclass(frozen=True) @@ -29,6 +38,7 @@ class NcclReloadSpec: timeout_s: float = 1200.0 chunk_bytes: int = 256 * 1024 * 1024 rate_limit_gbps: float = 0.0 + session_id: str = "" def __post_init__(self): if not self.host.strip(): @@ -44,13 +54,47 @@ def __post_init__(self): def _group(spec: NcclReloadSpec): from vllm.distributed.utils import StatelessProcessGroup - return StatelessProcessGroup.create( + timeout_s = float(spec.timeout_s) + if not math.isfinite(timeout_s) or timeout_s < 0: + raise ValueError("NCCL store timeout must be finite non-negative seconds") + timeout_s = timeout_s if timeout_s > 0 else _UNBOUNDED_STORE_TIMEOUT_S + if spec.session_id: + # Keep one TCP listener across member replacement. A new namespace must + # reset metadata counters/NCCL IDs without rebinding the occupied port. + key = (str(spec.host), int(spec.port), int(spec.rank) == 0) + timeout = timedelta(seconds=timeout_s) + with _BASE_STORE_CACHE_LOCK: + store = _BASE_STORE_CACHE.get(key) + if store is None: + store = torch.distributed.TCPStore( + host_name=spec.host, port=int(spec.port), world_size=None, + is_master=key[2], timeout=timeout, wait_for_workers=False, + use_libuv=False, + ) + _BASE_STORE_CACHE[key] = store + store.set_timeout(timeout) + namespaced = torch.distributed.PrefixStore(spec.session_id, store) + namespaced.set_timeout(timeout) + return StatelessProcessGroup( + rank=int(spec.rank), world_size=int(spec.world_size), + store=namespaced, socket=None, + ) + group = StatelessProcessGroup.create( host=spec.host, port=int(spec.port), rank=int(spec.rank), world_size=int(spec.world_size), - store_timeout=max(30, int(spec.timeout_s)), + store_timeout=timeout_s, ) + # vLLM 0.9.2 keeps a Python socket for the master_listen_fd handed to + # TCPStore. TCPStore owns and closes that descriptor. Letting the Python + # socket also close it can abort the store's daemon during cache eviction + # or interpreter shutdown. Transfer ownership while both objects are alive; + # receiver groups have no listening socket. + listen_socket = getattr(group, "socket", None) + if spec.rank == 0 and listen_socket is not None: + listen_socket.detach() + return group def _communicator(spec: NcclReloadSpec): @@ -82,6 +126,7 @@ def all_reduce(self, in_tensor, op=None, stream=None): int(spec.world_size), int(spec.rank), device_key, + spec.session_id, ) with _COMMUNICATOR_CACHE_LOCK: communicator = _COMMUNICATOR_CACHE.get(key) @@ -100,13 +145,51 @@ def all_reduce(self, in_tensor, op=None, stream=None): communicator = _BroadcastCommunicator(_group(spec), device=device) communicator._constructor_warmup_output = None _COMMUNICATOR_CACHE[key] = communicator + # The new group has completed NCCL initialization/warm-up. World size + # and local rank describe membership, not route ownership: repartition + # can change a surviving worker's rank without changing its CUDA device. + # Owners serialize each host/port/device route and drain old transfers + # first, so superseded sessions can be aborted without departed peers. + if spec.session_id: + for previous in list(_COMMUNICATOR_CACHE): + if (previous != key and previous[:2] == key[:2] + and previous[4] == key[4]): + _abort_communicator(_COMMUNICATOR_CACHE[previous]) + del _COMMUNICATOR_CACHE[previous] return communicator +def _abort_communicator(communicator) -> None: + """Release a local NCCL handle without waiting for an exited member.""" + handle = getattr(communicator, "comm", None) + if handle is not None: + nccl = communicator.nccl + # vLLM 0.9.2 does not wrap abort. Bind the library symbol locally with + # the same ncclResult_t(ncclComm_t) ABI as its ncclCommDestroy wrapper. + abort = nccl.lib.ncclCommAbort + abort.restype = ctypes.c_int + abort.argtypes = [ctypes.c_void_p] + device = getattr(communicator, "device", None) + if isinstance(device, torch.device) and device.type == "cuda" and torch.cuda.is_available(): + with torch.cuda.device(device): + result = abort(handle) + else: + result = abort(handle) + nccl.NCCL_CHECK(result) + communicator.comm = None + communicator.disabled = True + communicator.available = False + communicator.group = None + + def clear_communicator_cache() -> None: - """Drop cached communicators, typically during orderly process shutdown.""" + """Release owned NCCL groups before dropping their shared TCPStore listener.""" with _COMMUNICATOR_CACHE_LOCK: - _COMMUNICATOR_CACHE.clear() + for key in list(_COMMUNICATOR_CACHE): + _abort_communicator(_COMMUNICATOR_CACHE[key]) + del _COMMUNICATOR_CACHE[key] + with _BASE_STORE_CACHE_LOCK: + _BASE_STORE_CACHE.clear() def _dtype_name(dtype: torch.dtype) -> str: @@ -167,6 +250,10 @@ def drain_one() -> None: if tensor.device.type != "cuda": tensor = tensor.to(device=spec.device, non_blocking=True) tensor = tensor.contiguous() + if stream is not None: + # Bridge export, nonblocking copies and contiguous() produce on + # the current stream; PyNcclCommunicator does not fence that work. + stream.wait_stream(torch.cuda.current_stream(device=torch.device(spec.device))) chunk_numel = max(1, int(spec.chunk_bytes) // tensor.element_size()) metadata = { "done": False, @@ -222,6 +309,10 @@ def receive_weights(spec: NcclReloadSpec) -> Iterator[tuple[str, torch.Tensor]]: dtype=_dtype_from_name(str(metadata["dtype"])), device=spec.device, ) + if stream is not None: + # The allocator may reuse memory still owned by work on its + # allocation stream, even though the NCCL receive only writes it. + stream.wait_stream(torch.cuda.current_stream(device=torch.device(spec.device))) flat = tensor.view(-1) chunk_numel = max(1, int(metadata["chunk_numel"])) for start in range(0, flat.numel(), chunk_numel): diff --git a/infra/sync/prefix_cache_reset.py b/infra/sync/prefix_cache_reset.py new file mode 100644 index 0000000..b7b80c8 --- /dev/null +++ b/infra/sync/prefix_cache_reset.py @@ -0,0 +1,40 @@ +"""Opt-in V1 utility preserving the scheduler's prefix-reset success signal.""" + + +def install_checked_prefix_cache_reset(core_class=None): + """Install in each spawned core process, never in installed source files. + + vLLM 0.9.2's EngineCore and AsyncLLM reset wrappers discard the scheduler's + boolean return. A distinct utility makes missing child-process installation + an explicit RPC error instead of silently accepting a discarded failure. + """ + if core_class is None: + import vllm + if vllm.__version__.split('+', 1)[0] != '0.9.2': + raise RuntimeError('Checked prefix reset requires vLLM 0.9.2') + from vllm.v1.engine.core import EngineCore + core_class = EngineCore + if getattr(core_class, '_libra_checked_prefix_reset', False): + return + + def reset(self): + if self.scheduler.reset_prefix_cache() is not True: + raise RuntimeError('Prefix cache reset failed: active KV blocks remain') + return True + + core_class.libra_reset_prefix_cache = reset + core_class._libra_checked_prefix_reset = True + + +async def reset_prefix_cache_verified(client, *, require_verified=False): + if require_verified: + core_client = getattr(client, 'engine_core', None) + if core_client is None or not hasattr(core_client, 'call_utility_async'): + raise RuntimeError('Checked prefix reset requires the V1 utility client') + if await core_client.call_utility_async('libra_reset_prefix_cache') is not True: + raise RuntimeError('Prefix cache reset utility did not attest success') + return True + result = await client.reset_prefix_cache() + if result is False: + raise RuntimeError('Prefix cache reset failed') + return True if result is True else None diff --git a/infra/sync/publication_session.py b/infra/sync/publication_session.py new file mode 100644 index 0000000..0c21b93 --- /dev/null +++ b/infra/sync/publication_session.py @@ -0,0 +1,37 @@ +"""Keep global NCCL membership aligned with the live rollout processes.""" +from __future__ import annotations + +import json +import urllib.request +import uuid + + +def prepare_nccl_session(owner, offsets): + """Reuse the group while the current receivers keep their physical UUIDs.""" + config = owner.config + hetero = getattr(config, 'heterogeneous_rollout', None) + instances = getattr(hetero, 'instances', ()) if hetero else () + if instances: + receivers = {} + for index, instance in enumerate(instances): + host = instance.host or hetero.vllm_host + host = '127.0.0.1' if host == '0.0.0.0' else host + port = int(instance.port or hetero.vllm_base_port + index) + identity = str(getattr(instance, 'instance_id', f'instance_{index}')) + receivers[identity] = (f'http://{host}:{port}', int(instance.tp)) + else: + # The existing homogeneous setup already resolves VLLM_ENDPOINTS and + # host/base-port/count defaults. Use those actual clients unchanged. + receivers = {f'instance_{index}': (engine.base_url, int(config.vllm_tp_size)) + for index, engine in enumerate(owner.rollout_engine.engines)} + members = [] + for instance_id, offset in offsets.items(): + endpoint, tp = receivers[instance_id] + with urllib.request.urlopen(endpoint.rstrip('/') + '/publication/status', timeout=130) as response: + state = json.load(response) + members.append((instance_id, endpoint, tp, offset, state['server_instance_id'])) + membership = tuple(members) + if getattr(owner, '_nccl_session_members', None) != membership: + owner._nccl_session_id = uuid.uuid4().hex + owner._nccl_session_members = membership + return dict(nccl_session_id=owner._nccl_session_id) diff --git a/pyproject.toml b/pyproject.toml index 5b6eab7..1389452 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -52,6 +52,7 @@ packages = [ "RL_Framework.infra.scheduling", "RL_Framework.infra.sync", "RL_Framework.launcher", + "RL_Framework.scripts", "RL_Framework.trainer", "RL_Framework.workflow", ] diff --git a/requirements-ci.txt b/requirements-ci.txt index 36bfd40..f804ad6 100644 --- a/requirements-ci.txt +++ b/requirements-ci.txt @@ -4,3 +4,5 @@ pytest-asyncio==0.26.0 numpy==1.26.4 PyYAML==6.0.2 httpx==0.28.1 +fastapi # Serving lifecycle endpoints share the vLLM runtime API. +transformers==4.56.2 diff --git a/scripts/restartable_vllm_server.py b/scripts/restartable_vllm_server.py index f41af8d..efc964e 100644 --- a/scripts/restartable_vllm_server.py +++ b/scripts/restartable_vllm_server.py @@ -44,8 +44,8 @@ def parse_args() -> argparse.Namespace: def wait_until_ready(process: subprocess.Popen, url: str, timeout: float) -> None: - deadline = time.monotonic() + timeout - while time.monotonic() < deadline: + deadline = time.monotonic() + timeout if timeout > 0 else None + while deadline is None or time.monotonic() < deadline: if process.poll() is not None: raise RuntimeError(f"vLLM exited with status {process.returncode}") try: @@ -70,7 +70,7 @@ def reload_inplace(health_url: str, checkpoint_path: str, timeout: float) -> dic headers={"Content-Type": "application/json"}, method="POST", ) - with urllib.request.urlopen(request, timeout=timeout + 10) as response: + with urllib.request.urlopen(request, timeout=timeout + 10 if timeout > 0 else None) as response: if response.status != 200: raise RuntimeError(f"vLLM reload returned HTTP {response.status}") return json.loads(response.read().decode("utf-8")) @@ -89,7 +89,7 @@ def reload_nccl( headers={"Content-Type": "application/json"}, method="POST", ) - with urllib.request.urlopen(request, timeout=timeout + 10) as response: + with urllib.request.urlopen(request, timeout=timeout + 10 if timeout > 0 else None) as response: if response.status != 200: raise RuntimeError(f"vLLM NCCL reload returned HTTP {response.status}") return json.loads(response.read().decode("utf-8")) @@ -127,12 +127,29 @@ def write_json_atomic(path: Path, payload: dict) -> None: tmp_path.replace(path) +def _request_identity(stat) -> tuple[int, int, int, int]: + return stat.st_dev, stat.st_ino, stat.st_mtime_ns, stat.st_size + + +def read_reload_request(path: Path) -> tuple[dict, tuple[int, int, int, int]]: + """Read contents and identity from one open file, across atomic replacements.""" + with path.open(encoding="utf-8") as stream: + request = json.load(stream) + stat = os.fstat(stream.fileno()) + return request, _request_identity(stat) + + def main() -> int: args = parse_args() control_dir = Path(args.control_dir) control_dir.mkdir(parents=True, exist_ok=True) request_path = control_dir / "reload_request.json" + try: + startup_request = _request_identity(request_path.stat()) + except FileNotFoundError: + startup_request = None current_version = 0 + current_snapshot = None process: subprocess.Popen | None = None def request_shutdown(signum, _frame): @@ -157,7 +174,7 @@ def launch(model_path: str) -> subprocess.Popen: raise RuntimeError(f"vLLM exited unexpectedly: {process.returncode}") if request_path.exists(): try: - request = json.loads(request_path.read_text(encoding="utf-8")) + request, request_identity = read_reload_request(request_path) requested_version = int(request["version"]) reload_mode = str(request.get("mode", "checkpoint")).lower() checkpoint_path = str(request.get("checkpoint_path", "")) @@ -178,7 +195,17 @@ def launch(model_path: str) -> subprocess.Popen: except (OSError, ValueError, KeyError, json.JSONDecodeError): time.sleep(max(0.05, args.poll_interval)) continue - if requested_version > current_version: + # A fresh service cannot join the old NCCL sender. Wait for a + # new request file, even if it names the same revision/snapshot. + # Capture happened before launch, so requests written while + # waiting for the child's health are still consumed normally. + if reload_mode == "nccl" and request_identity == startup_request: + time.sleep(max(0.05, args.poll_interval)) + continue + requested_snapshot = request.get('snapshot_id') + if (requested_version > current_version or + requested_version == current_version and requested_snapshot + and requested_snapshot != current_snapshot): if reload_mode == "checkpoint" and not Path(checkpoint_path).is_dir(): raise FileNotFoundError(checkpoint_path) current_version = requested_version @@ -205,6 +232,10 @@ def launch(model_path: str) -> subprocess.Popen: ), "timeout_seconds": float(args.ready_timeout), "chunk_bytes": int(request["nccl_chunk_bytes"]), + **{key: request[key] for key in ('snapshot_id', 'optimizer_revision', 'expected_tensor_count') + if key in request}, + **(dict(session_id=request['nccl_session_id']) + if request.get('nccl_session_id') else {}), }, args.ready_timeout, ) @@ -226,6 +257,7 @@ def launch(model_path: str) -> subprocess.Popen: { "instance_id": args.instance_id, "version": current_version, + **{key: request[key] for key in ('snapshot_id', 'optimizer_revision', 'topology_epoch') if key in request}, "checkpoint_path": checkpoint_path, "reload_strategy": reload_strategy, "reload_method": reload_method, @@ -241,6 +273,7 @@ def launch(model_path: str) -> subprocess.Popen: { "instance_id": args.instance_id, "version": current_version, + **{key: request[key] for key in ('snapshot_id', 'optimizer_revision', 'topology_epoch') if key in request}, "checkpoint_path": checkpoint_path, "reload_strategy": reload_strategy, "reload_method": reload_method, @@ -256,6 +289,7 @@ def launch(model_path: str) -> subprocess.Popen: "server_response": response, }, ) + current_snapshot = requested_snapshot print( f"[vllm-reloader] ready version={current_version} " f"method={reload_method} strategy={reload_strategy} " diff --git a/scripts/serving_protocol.py b/scripts/serving_protocol.py new file mode 100644 index 0000000..bfb7e69 --- /dev/null +++ b/scripts/serving_protocol.py @@ -0,0 +1,157 @@ +"""Keep request lifetimes and weight installation on one serving boundary.""" +from __future__ import annotations + +import asyncio +import hashlib +import json +import time +import uuid + + +class ServingState: + def __init__(self, tp_size, model_path=''): + if int(tp_size) < 1: + raise ValueError('Expected positive physical TP worker count') + self.tp_size = int(tp_size) + self.server_instance_id = uuid.uuid4().hex + self.server_generation = 0 + self.snapshot_id = 'initial:' + hashlib.sha256(str(model_path).encode()).hexdigest() + self.optimizer_revision = 0 + self.workers = [] + self.active_requests = 0 + self.closed = False + self.failed = None + self.condition = asyncio.Condition() + + def status(self): + return dict(server_instance_id=self.server_instance_id, server_generation=self.server_generation, + snapshot_id=self.snapshot_id, optimizer_revision=self.optimizer_revision, + tp_size=self.tp_size, workers=self.workers, active_requests=self.active_requests, + closed=self.closed, failed=self.failed, + server_generation_verified=self.failed is None and not self.closed, + server_monotonic=time.monotonic()) + + async def enter(self): + async with self.condition: + await self.condition.wait_for(lambda: not self.closed or self.failed) + if self.failed: + raise RuntimeError('Rollout service is quarantined: ' + self.failed) + self.active_requests += 1 + return self.server_generation + + async def leave(self): + async with self.condition: + self.active_requests -= 1 + if self.active_requests < 0: + raise RuntimeError('Unbalanced serving request lifetime') + self.condition.notify_all() + + async def begin_reload(self, timeout): + async with self.condition: + if self.closed or self.failed: + raise RuntimeError('Rollout service is unavailable for reload') + self.closed = True + wait = self.condition.wait_for(lambda: self.active_requests == 0) + try: + if timeout > 0: + await asyncio.wait_for(wait, timeout) + else: + await wait + except BaseException: + # No mutation occurred before the drain; retry is safe. + self.closed = False + self.condition.notify_all() + raise + + def install(self, *, snapshot_id, optimizer_revision, workers, rank_offset=0, expected_tensor_count=None): + if not self.closed: + raise RuntimeError('Reload has no exclusive drained service ownership') + ranks = [worker.get('rank') for worker in workers] + if (len(workers) != self.tp_size or sorted(ranks) != list(range(self.tp_size)) + or any(worker.get('service_generation') != self.server_generation + 1 for worker in workers)): + raise RuntimeError('All TP workers must attest the next service generation') + if snapshot_id is not None: + counts = {worker.get('received_tensors') for worker in workers} + count = expected_tensor_count + if count is None and len(counts) == 1: + count = next(iter(counts)) + if (not snapshot_id or type(optimizer_revision) is not int or optimizer_revision < self.optimizer_revision + or type(count) is not int or count < 1 + or sorted(worker.get('nccl_rank', -1) for worker in workers) != list(range(1+rank_offset, 1+rank_offset+self.tp_size)) + or any(worker.get('receive_complete') is not True or worker.get('received_tensors') != count + or worker.get('loaded_parameters', 0) < 1 for worker in workers)): + raise RuntimeError('Missing complete immutable-snapshot tensor receipt on a TP worker') + self.server_generation += 1 + self.snapshot_id, self.optimizer_revision = snapshot_id, optimizer_revision if snapshot_id else 0 + self.workers = [dict(worker) for worker in workers] + + async def reopen(self): + async with self.condition: + if self.failed: + raise RuntimeError('A failed service cannot be reopened') + self.closed = False + self.condition.notify_all() + + async def fail(self, error): + async with self.condition: + self.failed = str(error) + self.closed = True + self.condition.notify_all() + + +def generation_headers(state, request_id, generation, status, streaming=False): + verified = bool(state.failed is None and generation == state.server_generation + and 200 <= status < 300 and not streaming) + values = { + 'x-libra-request-id': request_id, + 'x-libra-service-instance': state.server_instance_id, + 'x-libra-service-generation': str(state.server_generation), + 'x-libra-service-generation-start': str(generation), + 'x-libra-service-generation-end': str(state.server_generation), + 'x-libra-service-generation-verified': '1' if verified else '0', + 'x-libra-reload-overlap': '0' if generation == state.server_generation else '1', + } + return [(key.encode(), value.encode()) for key, value in values.items()] + + +class GenerationMiddleware: + """Admit request lifetimes through the same boundary used by weight reload.""" + def __init__(self, app, state): + self.app, self.state = app, state + + async def __call__(self, scope, receive, send): + if (scope['type'] != 'http' or scope.get('method') != 'POST' + or scope.get('path') not in {'/v1/completions', '/v1/chat/completions'}): + return await self.app(scope, receive, send) + headers = dict(scope.get('headers', [])) + request_id = headers.get(b'x-libra-request-id', uuid.uuid4().hex.encode()).decode() + try: + generation = await self.state.enter() + except RuntimeError as exc: + body = json.dumps(dict(detail=str(exc))).encode() + await send(dict(type='http.response.start', status=409, headers=[(b'content-type', b'application/json')])) + return await send(dict(type='http.response.body', body=body)) + start, bodies, streaming = None, [], False + async def observe(message): + nonlocal start, streaming + if message['type'] == 'http.response.start': + start = dict(message) + streaming = b'text/event-stream' in dict(message.get('headers', [])).get(b'content-type', b'') + if streaming: + start['headers'] = list(start.get('headers', [])) + generation_headers( + self.state, request_id, generation, start['status'], streaming=True) + await send(start) + return + if message['type'] != 'http.response.body' or streaming: + return await send(message) + bodies.append(message) + if not message.get('more_body', False): + start['headers'] = list(start.get('headers', [])) + generation_headers( + self.state, request_id, generation, start['status']) + await send(start) + for part in bodies: + await send(part) + try: + await self.app(scope, receive, observe) + finally: + await self.state.leave() diff --git a/scripts/vllm_hot_reload_api_server.py b/scripts/vllm_hot_reload_api_server.py index 6f5ba7d..1017419 100644 --- a/scripts/vllm_hot_reload_api_server.py +++ b/scripts/vllm_hot_reload_api_server.py @@ -18,6 +18,13 @@ from vllm.entrypoints.openai.cli_args import make_arg_parser from vllm.entrypoints.openai.cli_args import validate_parsed_serve_args from vllm.utils import FlexibleArgumentParser +from RL_Framework.infra.sync.prefix_cache_reset import ( + install_checked_prefix_cache_reset, reset_prefix_cache_verified, +) +from RL_Framework.scripts.serving_protocol import ServingState, GenerationMiddleware + +# EngineCore subprocesses import this project entrypoint before utility RPCs. +install_checked_prefix_cache_reset() class ReloadWeightsRequest(BaseModel): @@ -36,12 +43,21 @@ class ReloadWeightsNcclRequest(BaseModel): rank_offset: int timeout_seconds: float = 1200.0 chunk_bytes: int = 268435456 + snapshot_id: str | None = None + optimizer_revision: int | None = None + expected_tensor_count: int | None = None + session_id: str = '' _reload_lock = asyncio.Lock() _custom_router = APIRouter() +@_custom_router.get('/publication/status') +async def publication_status(request: Request): + return request.app.state.serving_state.status() + + def _kv_directory(key: str) -> Path: if not re.fullmatch(r"[a-f0-9]{32}", key): raise HTTPException(status_code=400, detail="invalid KV key") @@ -88,22 +104,24 @@ async def reload_weights(payload: ReloadWeightsRequest, request: Request): async with _reload_lock: started = time.perf_counter() + state = request.app.state.serving_state + await state.begin_reload(payload.timeout_seconds) client = api_server.engine_client(request) try: results = await client.collective_rpc( - "reload_weights", - timeout=payload.timeout_seconds, + "reload_weights", timeout=payload.timeout_seconds if payload.timeout_seconds > 0 else None, args=(str(checkpoint),), ) - # Prefix-cache entries encode outputs from the previous parameters. - await client.reset_prefix_cache() - except Exception as exc: - raise HTTPException(status_code=500, detail=str(exc)) from exc - return { - "checkpoint_path": str(checkpoint), - "total_seconds": time.perf_counter() - started, - "workers": results, - } + cleared = await reset_prefix_cache_verified(client, + require_verified=os.environ.get('LIBRA_CHECK_PREFIX_CACHE') == '1') + state.install(snapshot_id=None, optimizer_revision=None, workers=results) + await state.reopen() + except BaseException as exc: + await state.fail(repr(exc)) + raise + return dict(checkpoint_path=str(checkpoint), total_seconds=time.perf_counter()-started, + workers=results, prefix_cache_cleared=cleared is True, + service_generation=state.server_generation, service_instance_id=state.server_instance_id) @_custom_router.post("/reload_weights_nccl") @@ -114,27 +132,33 @@ async def reload_weights_nccl( """Receive the next model directly from the Megatron GPU stream.""" async with _reload_lock: started = time.perf_counter() + state = request.app.state.serving_state + if ((payload.snapshot_id is None) != (payload.optimizer_revision is None) + or payload.optimizer_revision is not None and payload.optimizer_revision < state.optimizer_revision): + raise HTTPException(status_code=409, detail='Missing or regressing snapshot identity') + await state.begin_reload(payload.timeout_seconds) client = api_server.engine_client(request) try: results = await client.collective_rpc( - "reload_weights_nccl", - timeout=payload.timeout_seconds, - args=( - payload.host, - payload.port, - payload.world_size, - payload.rank_offset, - payload.timeout_seconds, - payload.chunk_bytes, - ), + "reload_weights_nccl", timeout=payload.timeout_seconds if payload.timeout_seconds > 0 else None, + args=(payload.host, payload.port, payload.world_size, payload.rank_offset, + payload.timeout_seconds, payload.chunk_bytes) + + ((payload.expected_tensor_count, payload.session_id) if payload.session_id else + (payload.expected_tensor_count,) if payload.expected_tensor_count is not None else ()), ) - await client.reset_prefix_cache() - except Exception as exc: - raise HTTPException(status_code=500, detail=str(exc)) from exc - return { - "total_seconds": time.perf_counter() - started, - "workers": results, - } + cleared = await reset_prefix_cache_verified(client, + require_verified=bool(payload.snapshot_id) or os.environ.get('LIBRA_CHECK_PREFIX_CACHE') == '1') + state.install(snapshot_id=payload.snapshot_id, optimizer_revision=payload.optimizer_revision, + workers=results, rank_offset=payload.rank_offset, + expected_tensor_count=payload.expected_tensor_count) + await state.reopen() + except BaseException as exc: + await state.fail(repr(exc)) + raise + return dict(total_seconds=time.perf_counter()-started, workers=results, + prefix_cache_cleared=cleared is True, service_generation=state.server_generation, + service_instance_id=state.server_instance_id, snapshot_id=state.snapshot_id, + optimizer_revision=state.optimizer_revision) # vLLM 0.9.x constructs a fresh FastAPI app in build_app() and includes its @@ -144,9 +168,16 @@ async def reload_weights_nccl( _build_app = api_server.build_app +def make_serving_state(args): + """Build the service state shared by generation and weight reload.""" + return ServingState(args.tensor_parallel_size, args.model) + + def build_app_with_reload_routes(args): app = _build_app(args) + app.state.serving_state = make_serving_state(args) app.include_router(_custom_router) + app.add_middleware(GenerationMiddleware, state=app.state.serving_state) return app diff --git a/tests/test_checked_prefix_cache_reset.py b/tests/test_checked_prefix_cache_reset.py new file mode 100644 index 0000000..5ca224e --- /dev/null +++ b/tests/test_checked_prefix_cache_reset.py @@ -0,0 +1,52 @@ +"""A failed or missing V1 utility cannot produce a cache-clear ACK.""" + +import asyncio +from types import SimpleNamespace + +import pytest + +from RL_Framework.infra.sync.prefix_cache_reset import ( + install_checked_prefix_cache_reset, reset_prefix_cache_verified, +) + + +@pytest.mark.parametrize('success', [True, False]) +def test_checked_core_utility_preserves_scheduler_outcome(success): + class Core: + def reset_prefix_cache(self): + # Mirrors vLLM 0.9.2's original wrapper, which discards the bool. + self.scheduler.reset_prefix_cache() + + original = Core.reset_prefix_cache + install_checked_prefix_cache_reset(Core) + core = Core() + core.scheduler = SimpleNamespace(reset_prefix_cache=lambda: success) + assert Core.reset_prefix_cache is original + if success: + assert core.libra_reset_prefix_cache() is True + else: + with pytest.raises(RuntimeError, match='active KV blocks'): + core.libra_reset_prefix_cache() + + +@pytest.mark.parametrize('outcome', [True, False, 'missing']) +def test_verified_reset_accepts_only_the_real_utility_success(outcome): + async def utility(name): + assert name == 'libra_reset_prefix_cache' + if outcome == 'missing': + raise AttributeError('Unknown utility method') + return outcome + + client = SimpleNamespace(engine_core=SimpleNamespace(call_utility_async=utility)) + if outcome is True: + assert asyncio.run(reset_prefix_cache_verified(client, require_verified=True)) is True + else: + with pytest.raises((RuntimeError, AttributeError)): + asyncio.run(reset_prefix_cache_verified(client, require_verified=True)) + + +def test_legacy_none_return_is_preserved_without_claiming_verified_reset(): + async def reset(): + return None + + assert asyncio.run(reset_prefix_cache_verified(SimpleNamespace(reset_prefix_cache=reset))) is None diff --git a/tests/test_collective_weight_export.py b/tests/test_collective_weight_export.py new file mode 100644 index 0000000..bbc2caa --- /dev/null +++ b/tests/test_collective_weight_export.py @@ -0,0 +1,128 @@ +"""Transport failures must not abandon peers inside the next TP export.""" + +import io +import json +from datetime import timedelta +from pathlib import Path +from types import SimpleNamespace +from unittest.mock import patch + +import torch +import torch.distributed as dist +import torch.multiprocessing as mp + + +def _worker(rank, rendezvous, directory): + torch.set_num_threads(1) + dist.init_process_group("gloo", init_method=f"file://{rendezvous}", rank=rank, world_size=2, + timeout=timedelta(seconds=12)) + control = dist.new_group(backend="gloo", timeout=timedelta(seconds=12)) + try: + from RL_Framework.infra.sync.collective_export import stream_export_with_consensus + from RL_Framework.infra.sync import nccl_weight_sync, publication_session + from RL_Framework.trainer.async_rl_trainer import AsyncRLTrainer + for entrypoint in ("helper", "native"): + for failure in ("before-first", "after-first", "iterator", "prepare", "request"): + if failure == "prepare" and entrypoint != "helper": + continue + if failure == "request" and entrypoint == "helper": + continue + exported, sent = [], [] + + def weights(): + for index in range(3): + # Real inter-rank compute/export synchronization. + dist.barrier() + if failure == "iterator" and rank == 1 and index == 1: + raise RuntimeError("iterator failed on rank 1") + exported.append(index) + yield f"weight-{index}", torch.ones(2) * index + + def prepare(item): + if failure == "prepare" and rank == 0 and item[0] == "weight-1": + raise RuntimeError("tensor preparation failed") + return item + + def send(items, *args, **kwargs): + if failure == "before-first": + raise RuntimeError("transport bootstrap failed") + for item in items: + sent.append(item[0]) + if failure == "after-first": + raise RuntimeError("transport failed after first tensor") + + folder = Path(directory) / entrypoint / failure + config = SimpleNamespace( + rollout_weight_sync_control_dir=str(folder), rollout_nccl_host="127.0.0.1", + rollout_nccl_port=29620, rollout_nccl_chunk_mb=1, + rollout_nccl_rate_limit_gbps=0, rollout_weight_sync_timeout_s=2, + rollout_weight_sync_poll_interval_s=0.001, + heterogeneous_rollout=SimpleNamespace(instances=[ + SimpleNamespace(instance_id="engine", host="127.0.0.1", port=8000, tp=2), + ]), + ) + + def command(value): + packet = [value] + dist.broadcast_object_list(packet, src=0, group=control) + if packet[0].get("error"): + raise RuntimeError(packet[0]["error"]) + return packet[0] + + def write_request(path, value): + if failure == "request": + raise OSError("control request write failed") + path.write_text(json.dumps(value)) + + trainer = SimpleNamespace( + config=config, is_main_process=rank == 0, control=control, + train_engine=SimpleNamespace(train_backend="megatron_core", optimizer_revision=23, + stream_rollout_weights=weights), + optimizer_revision=23, service_generations={}, + events=SimpleNamespace(emit=lambda *args, **kwargs: None), command=command, + _nccl_rollout_topology=lambda: (3, {"engine": 0}), + _write_json_atomic=write_request, + ) + caught = "" + with patch.object(nccl_weight_sync, "send_weights", send), \ + patch.object(publication_session.urllib.request, "urlopen", + lambda *args, **kwargs: io.StringIO('{"server_instance_id":"live-service"}')), \ + patch.object(torch.cuda, "current_device", lambda: "cpu"), \ + patch.object(torch.cuda, "current_stream", lambda: SimpleNamespace(synchronize=lambda: None)): + try: + if entrypoint == "helper": + stream_export_with_consensus(weights, is_sender=rank == 0, send=send, + group=control, prepare=prepare) + elif entrypoint == "native": + AsyncRLTrainer._sync_weights_nccl(trainer, 23) + except RuntimeError as exc: + caught = str(exc) + if failure == "request": + assert "control request write failed" in caught, caught + assert exported == [] + else: + assert "collective weight publication failed" in caught, caught + errors = [None, None] + dist.all_gather_object(errors, caught, group=control) + assert all(errors) + if failure in ("before-first", "after-first"): + assert exported == [0, 1, 2] + assert sent == (["weight-0"] if rank == 0 and failure == "after-first" else []) + elif failure in ("iterator", "prepare"): + assert 2 not in exported # nobody enters another export collective + # A matching barrier proves the failed publication left no rank + # waiting in an orphaned export iteration. + dist.barrier() + + captured = [] + stream_export_with_consensus(lambda: iter([("ok", torch.ones(1))]), + is_sender=rank == 0, + send=lambda items: captured.extend(items), group=control) + assert len(captured) == (1 if rank == 0 else 0) + finally: + dist.destroy_process_group() + + +def test_two_rank_stream_failures_drain_then_fail_all_entrypoints(tmp_path): + mp.start_processes(_worker, args=(str(tmp_path / "rendezvous"), str(tmp_path)), + nprocs=2, join=True, start_method="spawn") diff --git a/tests/test_nccl_store_lifecycle.py b/tests/test_nccl_store_lifecycle.py new file mode 100644 index 0000000..1116c7e --- /dev/null +++ b/tests/test_nccl_store_lifecycle.py @@ -0,0 +1,147 @@ +"""The stateless rendezvous transfers its listening FD to the real TCPStore.""" +import os +from pathlib import Path +import subprocess +import sys +import textwrap +from types import ModuleType, SimpleNamespace + +import pytest + +from RL_Framework.infra.sync import nccl_weight_sync +from RL_Framework.infra.sync.nccl_weight_sync import NcclReloadSpec + + +@pytest.mark.parametrize("rank", [0, 1]) +def test_only_master_transfers_listening_socket_ownership(monkeypatch, rank): + detached = [] + group = SimpleNamespace(socket=SimpleNamespace(detach=lambda: detached.append(True))) + utils = ModuleType("vllm.distributed.utils") + utils.StatelessProcessGroup = SimpleNamespace(create=lambda **_kwargs: group) + monkeypatch.setitem(sys.modules, "vllm.distributed.utils", utils) + + assert nccl_weight_sync._group(NcclReloadSpec("127.0.0.1", 12345, 2, rank, "cpu")) is group + assert detached == ([True] if rank == 0 else []) + + +def test_real_tcp_store_survives_python_socket_cleanup_and_cache_eviction(): + # A subprocess is essential: closing a socket that still owns TCPStore's FD + # triggers SIGABRT in a C++ daemon, bypassing Python exception handling. Use + # real CPU PyTorch but reproduce vLLM's small create() without importing GPU + # dependencies. Multiple groups also exercise independent cache lifetimes. + script = textwrap.dedent(""" + import gc + import socket + import sys + from datetime import timedelta + from types import ModuleType, SimpleNamespace + import torch.distributed as dist + from RL_Framework.infra.sync import nccl_weight_sync + + def create(host, port, rank, world_size, store_timeout): + sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM) + sock.bind((host, port)) + sock.listen() + port = sock.getsockname()[1] + store = dist.TCPStore(host, port, world_size, True, + timeout=timedelta(seconds=store_timeout), use_libuv=False, + master_listen_fd=sock.fileno()) + return SimpleNamespace(socket=sock, store=store) + + utils = ModuleType('vllm.distributed.utils') + utils.StatelessProcessGroup = SimpleNamespace(create=create) + sys.modules['vllm.distributed.utils'] = utils + for index in range(3): + spec = nccl_weight_sync.NcclReloadSpec('127.0.0.1', 0, 1, 0, 'cpu', timeout_s=2) + group = nccl_weight_sync._group(spec) + assert group.socket.fileno() == -1 + group.socket.close() + group.store.set('working', b'yes') + assert group.store.get('working') == b'yes' + nccl_weight_sync._COMMUNICATOR_CACHE[('test', index)] = SimpleNamespace(group=group) + del group + nccl_weight_sync.clear_communicator_cache() + gc.collect() + print('ORDERLY_TCPSTORE_EXIT', flush=True) + """) + root = Path(nccl_weight_sync.__file__).resolve().parents[2] + # Preserve the package alias provided by the ordinary project test runner. + environment = dict(os.environ) + completed = subprocess.run([sys.executable, "-B", "-c", script], cwd=root, + env=environment, capture_output=True, text=True, timeout=30) + assert completed.returncode == 0, completed.stdout + completed.stderr + assert "ORDERLY_TCPSTORE_EXIT" in completed.stdout + assert "TCPStoreMasterDaemon::run() failed" not in completed.stderr + + +def test_real_namespaced_store_reuses_listener_and_reconnects_new_member(): + # Keep TCPStore lifetime errors isolated in a real subprocess, as above. + # Only vLLM's dataclass container is stubbed; all stores/connections are real. + script = textwrap.dedent(''' + import gc + import os + import subprocess + import sys + import textwrap + from types import ModuleType, SimpleNamespace + from RL_Framework.infra.sync import nccl_weight_sync as transport + + utils = ModuleType('vllm.distributed.utils') + utils.StatelessProcessGroup = lambda **kwargs: SimpleNamespace(**kwargs) + sys.modules[utils.__name__] = utils + sender = transport.NcclReloadSpec('127.0.0.1', 0, 2, 0, 'cpu', timeout_s=2, session_id='old') + old_sender = transport._group(sender) + port = transport._BASE_STORE_CACHE[('127.0.0.1', 0, True)].port + receiver = transport.NcclReloadSpec('127.0.0.1', port, 2, 1, 'cpu', timeout_s=2, session_id='old') + old_receiver = transport._group(receiver) + old_sender.store.set('broadcast_from/0/0', b'old-nccl-id') + assert old_receiver.store.get('broadcast_from/0/0') == b'old-nccl-id' + + from dataclasses import replace + new_sender = transport._group(replace(sender, session_id='new')) + new_receiver = transport._group(replace(receiver, session_id='new')) + assert len(transport._BASE_STORE_CACHE) == 2 + assert transport._BASE_STORE_CACHE[('127.0.0.1', 0, True)].port == port + assert not new_receiver.store.check(['broadcast_from/0/0']) + new_sender.store.set('broadcast_from/0/0', b'new-nccl-id') + assert new_receiver.store.get('broadcast_from/0/0') == b'new-nccl-id' + assert old_receiver.store.get('broadcast_from/0/0') == b'old-nccl-id' + assert all(group.socket is None for group in (old_sender, old_receiver, new_sender, new_receiver)) + + client = textwrap.dedent(""" + import sys + from types import ModuleType, SimpleNamespace + from RL_Framework.infra.sync import nccl_weight_sync as transport + utils = ModuleType('vllm.distributed.utils') + utils.StatelessProcessGroup = lambda **kwargs: SimpleNamespace(**kwargs) + sys.modules[utils.__name__] = utils + group = transport._group(transport.NcclReloadSpec('127.0.0.1', int(sys.argv[1]), + 2, 1, 'cpu', timeout_s=2, session_id='new')) + assert group.store.get('broadcast_from/0/0') == b'new-nccl-id' + group.store.set('reconnected', b'yes') + del group + transport.clear_communicator_cache() + """) + result = subprocess.run([sys.executable, '-B', '-c', client, str(port)], + env=dict(os.environ), capture_output=True, text=True, timeout=15) + assert result.returncode == 0, result.stdout + result.stderr + assert new_sender.store.get('reconnected') == b'yes' + + del old_sender, old_receiver, new_sender, new_receiver + transport.clear_communicator_cache() + gc.collect() + assert not transport._BASE_STORE_CACHE + # After whole cleanup, the actual port can be bound by a fresh master. + fresh = transport._group(replace(sender, port=port, session_id='fresh')) + fresh.store.set('healthy', b'yes') + assert fresh.store.get('healthy') == b'yes' + del fresh + transport.clear_communicator_cache() + gc.collect() + print('NAMESPACED_STORE_RECONNECT_AND_CLEANUP_OK', flush=True) + ''') + root = Path(nccl_weight_sync.__file__).resolve().parents[2] + completed = subprocess.run([sys.executable, '-B', '-c', script], cwd=root, + env=dict(os.environ), capture_output=True, text=True, timeout=40) + assert completed.returncode == 0, completed.stdout + completed.stderr + assert 'NAMESPACED_STORE_RECONNECT_AND_CLEANUP_OK' in completed.stdout diff --git a/tests/test_nccl_store_timeout.py b/tests/test_nccl_store_timeout.py new file mode 100644 index 0000000..3e7cc32 --- /dev/null +++ b/tests/test_nccl_store_timeout.py @@ -0,0 +1,54 @@ +"""Shared rendezvous waits preserve application timeout=0 and subsecond limits.""" +from datetime import timedelta +import sys +from types import ModuleType, SimpleNamespace + +import pytest +import torch.distributed as dist + +from RL_Framework.infra.sync import nccl_weight_sync +from RL_Framework.infra.sync.nccl_weight_sync import NcclReloadSpec + + +@pytest.fixture +def create_group(monkeypatch): + calls = [] + utils = ModuleType('vllm.distributed.utils') + utils.StatelessProcessGroup = SimpleNamespace(create=lambda **kwargs: calls.append(kwargs) or kwargs) + for name in ('vllm', 'vllm.distributed'): + package = ModuleType(name) + package.__path__ = [] + monkeypatch.setitem(sys.modules, name, package) + monkeypatch.setitem(sys.modules, 'vllm.distributed.utils', utils) + return calls + + +@pytest.mark.parametrize('timeout', [0.125, 1, 29.5, 1200]) +def test_positive_nccl_store_timeout_preserves_seconds_without_minimum_or_truncation(create_group, timeout): + nccl_weight_sync._group(NcclReloadSpec('127.0.0.1', 12345, 2, 0, 'cpu', timeout_s=timeout)) + assert create_group[0]['store_timeout'] == timeout + + +def test_zero_nccl_store_timeout_uses_safe_unbounded_sentinel(create_group): + nccl_weight_sync._group(NcclReloadSpec('127.0.0.1', 12345, 2, 0, 'cpu', timeout_s=0)) + timeout = create_group[0]['store_timeout'] + assert timeout > 365 * 24 * 60 * 60 + assert timeout * 1_000_000_000 < 2**63 + # Exercise the same timedelta conversion performed by vLLM. + assert timedelta(seconds=timeout).total_seconds() == timeout + + +@pytest.mark.parametrize('timeout', [-1, float('nan'), float('inf')]) +def test_invalid_nccl_store_timeout_fails_before_group_creation(create_group, timeout): + with pytest.raises(ValueError, match='non-negative seconds'): + nccl_weight_sync._group(NcclReloadSpec('127.0.0.1', 12345, 2, 0, 'cpu', timeout_s=timeout)) + assert not create_group + + +def test_unbounded_sentinel_is_accepted_by_real_torch_tcp_store(): + # No peer or wait is needed: this verifies PyTorch's C++ duration conversion + # accepts the sentinel and the resulting store remains functional. + timeout = timedelta(seconds=nccl_weight_sync._UNBOUNDED_STORE_TIMEOUT_S) + store = dist.TCPStore('127.0.0.1', 0, 1, True, timeout=timeout) + store.set('probe', b'ready') + assert store.get('probe') == b'ready' diff --git a/tests/test_nccl_weight_sync.py b/tests/test_nccl_weight_sync.py index 85790f4..7480a2f 100644 --- a/tests/test_nccl_weight_sync.py +++ b/tests/test_nccl_weight_sync.py @@ -1,8 +1,11 @@ import sys +import ctypes +from dataclasses import replace from types import SimpleNamespace from types import ModuleType import torch +import pytest from RL_Framework.infra.sync import nccl_weight_sync from RL_Framework.infra.sync.nccl_weight_sync import NcclReloadSpec @@ -163,3 +166,132 @@ def all_reduce(self, _tensor, op=None, stream=None): assert first is second assert first.__class__.init_count == 1 nccl_weight_sync.clear_communicator_cache() + + +@pytest.mark.parametrize('direction', ['send', 'receive']) +def test_nccl_side_stream_waits_for_each_tensor_producer_before_broadcast(monkeypatch, direction): + events = [] + + class Stream: + def wait_stream(self, other): + assert self is transfer and other is producer + events.append('wait_producer') + + def wait_event(self, event): + assert self is producer + events.append('wait_receive') + + class Event: + def record(self, stream): + assert stream is transfer + events.append('record_event') + + def synchronize(self): + events.append('finish_send') + + producer, transfer = Stream(), Stream() + monkeypatch.setattr(nccl_weight_sync, '_cuda_stream', lambda device: transfer) + monkeypatch.setattr(torch.cuda, 'current_stream', lambda **kwargs: producer) + monkeypatch.setattr(torch.cuda, 'Event', Event) + monkeypatch.setattr(torch.cuda, 'synchronize', + lambda *args, **kwargs: pytest.fail('Transport dependencies must not synchronize the device')) + + def broadcast(tensor, src, stream): + assert src == 0 and stream is transfer + events.append('broadcast') + + if direction == 'send': + communicator = SimpleNamespace(group=FakeGroup(0, []), broadcast=broadcast) + contiguous = torch.Tensor.contiguous + + def copy_to_device(tensor, **kwargs): + assert kwargs['non_blocking'] is True + events.append('copy') + return tensor + + def prepare_contiguous(tensor): + result = contiguous(tensor) + events.append('contiguous') + return result + + monkeypatch.setattr(torch.Tensor, 'to', copy_to_device) + monkeypatch.setattr(torch.Tensor, 'contiguous', prepare_contiguous) + monkeypatch.setattr(nccl_weight_sync, '_communicator', lambda spec: communicator) + source = [(str(i), torch.arange(4).reshape(2, 2).t()) for i in range(2)] + spec = NcclReloadSpec('host', 29620, 2, 0, 'cuda:0', chunk_bytes=16) + assert nccl_weight_sync.send_weights(source, spec) == 2 + per_tensor = ['copy', 'contiguous', 'wait_producer', 'broadcast', 'broadcast', 'record_event'] + assert events == per_tensor * 2 + ['finish_send', 'finish_send'] + else: + metadata = [dict(name=str(i), shape=[4], dtype='float32', chunk_numel=2, done=False) + for i in range(2)] + [dict(done=True)] + communicator = SimpleNamespace(group=FakeGroup(1, metadata), broadcast=broadcast) + allocate = torch.empty + + def allocate_on_current_stream(shape, **kwargs): + events.append('allocate') + return allocate(shape, dtype=kwargs['dtype']) + + def record_consumer_stream(tensor, stream): + assert stream is producer + events.append('record_consumer') + + monkeypatch.setattr(torch, 'empty', allocate_on_current_stream) + monkeypatch.setattr(torch.Tensor, 'record_stream', record_consumer_stream) + monkeypatch.setattr(nccl_weight_sync, '_communicator', lambda spec: communicator) + spec = NcclReloadSpec('host', 29620, 2, 1, 'cuda:0') + assert len(list(nccl_weight_sync.receive_weights(spec))) == 2 + per_tensor = ['allocate', 'wait_producer', 'broadcast', 'broadcast', + 'record_event', 'wait_receive', 'record_consumer'] + assert events == per_tensor * 2 + + +@pytest.fixture +def session_transport(monkeypatch): + events = [] + + class FakePyNcclCommunicator: + def __init__(self, group, device): + session = group.session_id + events.append(('build', session)) + self.group, self.device = group, device + self.comm = ctypes.c_void_p(len(events)) + + def abort(handle): + assert handle is self.comm + events.append(('abort', session)) + return 0 + + def check(result): + assert result == 0 + + self.nccl = SimpleNamespace(lib=SimpleNamespace(ncclCommAbort=abort), NCCL_CHECK=check) + self.all_reduce('warmup') + events.append(('ready', session)) + + def all_reduce(self, _tensor, op=None, stream=None): + return object() + + module = ModuleType('vllm.distributed.device_communicators.pynccl') + module.PyNcclCommunicator = FakePyNcclCommunicator + monkeypatch.setitem(sys.modules, module.__name__, module) + monkeypatch.setattr(nccl_weight_sync, '_group', lambda spec: SimpleNamespace(session_id=spec.session_id)) + nccl_weight_sync.clear_communicator_cache() + yield SimpleNamespace(events=events, + spec=NcclReloadSpec('127.0.0.1', 29620, 3, 0, 'cpu', session_id='a')) + nccl_weight_sync.clear_communicator_cache() + + +def test_new_physical_session_rebuilds_then_releases_old_group(session_transport): + fixture = session_transport + first = nccl_weight_sync._communicator(fixture.spec) + assert nccl_weight_sync._communicator(fixture.spec) is first + other = nccl_weight_sync._communicator(replace(fixture.spec, port=29621, session_id='other')) + second = nccl_weight_sync._communicator(replace(fixture.spec, session_id='new-members')) + assert second is not first + assert fixture.events.index(('ready', 'new-members')) < fixture.events.index(('abort', 'a')) + assert first.comm is None and first.group is None + assert set(nccl_weight_sync._COMMUNICATOR_CACHE.values()) == {second, other} + assert ('abort', 'other') not in fixture.events + nccl_weight_sync.clear_communicator_cache() + assert second.comm is other.comm is None diff --git a/tests/test_publication_sessions.py b/tests/test_publication_sessions.py new file mode 100644 index 0000000..07b7962 --- /dev/null +++ b/tests/test_publication_sessions.py @@ -0,0 +1,89 @@ +"""The native sender reuses only the current physical receiver membership.""" +import io +import json +from types import SimpleNamespace as NS + +import pytest +import torch + +from RL_Framework.infra.sync import nccl_weight_sync, publication_session +from RL_Framework.trainer.async_rl_trainer import AsyncRLTrainer + + +@pytest.mark.parametrize('endpoints,environment,expected', [ + ('', '', ['http://127.0.0.1:8000', 'http://127.0.0.1:8001', 'http://127.0.0.1:8002']), + ('', 'remote-a:8100,remote-b:8200', ['http://remote-a:8100', 'http://remote-b:8200']), + ('explicit:8300', 'ignored:8400', ['http://explicit:8300']), +]) +def test_session_reuses_the_existing_homogeneous_endpoint_setup( + monkeypatch, endpoints, environment, expected): + config = NS(rollout_backend='vllm', model_path='/model', vllm_endpoints=endpoints, + vllm_host='127.0.0.1', vllm_port=8000, vllm_num_instances=0, + rollout_gpus=3, vllm_tp_size=1) + owner = NS(config=config, is_main_process=False) + monkeypatch.setenv('VLLM_ENDPOINTS', environment) + AsyncRLTrainer._setup_homogeneous_engine(owner) + assert [engine.base_url for engine in owner.rollout_engine.engines] == expected + observed = [] + + def status(url, timeout): + observed.append(url) + return io.StringIO(json.dumps(dict(server_instance_id=url))) + + monkeypatch.setattr(publication_session.urllib.request, 'urlopen', status) + offsets = {f'instance_{index}': index for index in range(owner.rollout_engine.num_instances)} + first = publication_session.prepare_nccl_session(owner, offsets) + assert publication_session.prepare_nccl_session(owner, offsets) == first + assert observed == [endpoint + '/publication/status' for endpoint in expected] * 2 + + +def test_native_publication_reuses_membership_and_sends_a_fresh_same_revision_transfer(tmp_path, monkeypatch): + states = dict(a='service-a', b='service-b') + config = NS(heterogeneous_rollout=NS(instances=[ + NS(instance_id=name, host=name, port=8000, tp=1) for name in states]), + rollout_weight_sync_control_dir=str(tmp_path), rollout_nccl_host='core', + rollout_nccl_port=29620, rollout_nccl_chunk_mb=1, rollout_nccl_rate_limit_gbps=0, + rollout_weight_sync_timeout_s=2, rollout_weight_sync_poll_interval_s=.001) + trainer = NS(config=config, is_main_process=True, control=None, + train_engine=NS(train_backend='megatron_core', optimizer_revision=7, + stream_rollout_weights=lambda: iter([('weight', torch.ones(2))])), + _nccl_rollout_topology=lambda: (3, {'a': 0, 'b': 1}), + _write_json_atomic=lambda path, value: path.write_text(json.dumps(value)), + _record_rollout_weight_version=lambda **kwargs: None, + _record_weight_sync_phase=lambda phase: None) + + def status(url, timeout): + host = url.split('://', 1)[1].split(':', 1)[0] + return io.StringIO(json.dumps(dict(server_instance_id=states[host]))) + + monkeypatch.setattr(publication_session.urllib.request, 'urlopen', status) + monkeypatch.setattr(torch.cuda, 'current_device', lambda: 'cpu') + requests, generations = [], {} + + def send(weights, spec, **kwargs): + tensors = list(weights) + for name, tensor in tensors: + kwargs['on_tensor'](name, tensor.numel() * tensor.element_size()) + request = json.loads((tmp_path / 'reload_request.json').read_text()) + assert spec.session_id == request['nccl_session_id'] + assert request['optimizer_revision'] == trainer.train_engine.optimizer_revision + requests.append(request) + for identity, offset in request['nccl_rank_offsets'].items(): + service = states[identity] + generations[service] = generations.get(service, 0) + 1 + worker = dict(rank=0, nccl_rank=offset+1, received_tensors=len(tensors), + receive_complete=True, loaded_parameters=1, service_generation=generations[service]) + record = dict(version=request['version'], snapshot_id=request['snapshot_id'], + optimizer_revision=request['optimizer_revision'], server_response=dict( + service_instance_id=service, service_generation=generations[service], + prefix_cache_cleared=True, workers=[worker])) + (tmp_path / f'ack_{identity}_{request["version"]}.json').write_text(json.dumps(record)) + + monkeypatch.setattr(nccl_weight_sync, 'send_weights', send) + AsyncRLTrainer._sync_weights_nccl(trainer, 1) + AsyncRLTrainer._sync_weights_nccl(trainer, 1) + assert requests[0]['snapshot_id'] != requests[1]['snapshot_id'] + assert requests[0]['nccl_session_id'] == requests[1]['nccl_session_id'] + states['a'] = 'returned-service-a' + AsyncRLTrainer._sync_weights_nccl(trainer, 1) + assert requests[2]['nccl_session_id'] != requests[1]['nccl_session_id'] diff --git a/tests/test_restartable_vllm_server.py b/tests/test_restartable_vllm_server.py index 1017a8f..ff721dc 100644 --- a/tests/test_restartable_vllm_server.py +++ b/tests/test_restartable_vllm_server.py @@ -1,6 +1,9 @@ import importlib.util import json from pathlib import Path +from types import SimpleNamespace + +import pytest SCRIPT_PATH = Path(__file__).parents[1] / "scripts" / "restartable_vllm_server.py" @@ -55,3 +58,82 @@ def test_parallel_is_default_reload_strategy(monkeypatch): args = module.parse_args() assert args.reload_method == "restart" assert args.reload_strategy == "parallel" + + +@pytest.fixture +def supervisor(tmp_path, monkeypatch): + module = _load_module() + args = SimpleNamespace(instance_id='returned', control_dir=str(tmp_path), + health_url='http://service/health', initial_model='/model', ready_timeout=1, + reload_strategy='parallel', reload_method='inplace', poll_interval=.01, + command=['server']) + monkeypatch.setattr(module, 'parse_args', lambda: args) + monkeypatch.setattr(module.subprocess, 'Popen', lambda *args, **kwargs: SimpleNamespace(poll=lambda: None)) + monkeypatch.setattr(module, 'wait_until_ready', lambda *args: None) + monkeypatch.setattr(module, 'stop_process', lambda *args: None) + monkeypatch.setattr(module.signal, 'signal', lambda *args: None) + calls = [] + + def reload(url, payload, timeout): + calls.append(payload) + return dict(service_instance_id='new-service', service_generation=len(calls)) + + def stop(seconds): + raise KeyboardInterrupt() + + monkeypatch.setattr(module, 'reload_nccl', reload) + monkeypatch.setattr(module.time, 'sleep', stop) + request = dict(mode='nccl', version=0, snapshot_id='same-snapshot', optimizer_revision=0, + nccl_host='train', nccl_port=29620, nccl_world_size=2, + nccl_rank_offsets={'returned': 0}, nccl_chunk_bytes=1024) + return SimpleNamespace(module=module, args=args, calls=calls, request=request, + request_path=tmp_path / 'reload_request.json', ack=tmp_path / 'ack_returned_0.json') + + +def test_startup_ignores_old_nccl_until_a_new_transfer_arrives(supervisor, monkeypatch): + s = supervisor + s.module.write_json_atomic(s.request_path, s.request) + polls = 0 + + def poll(seconds): + nonlocal polls + polls += 1 + if polls == 1: + assert not s.calls and not s.ack.exists() + s.request['snapshot_id'] = 'returned-transfer' + s.module.write_json_atomic(s.request_path, s.request) + else: + raise KeyboardInterrupt() + + monkeypatch.setattr(s.module.time, 'sleep', poll) + assert s.module.main() == 130 + assert len(s.calls) == 1 + ack = json.loads(s.ack.read_text()) + assert ack['snapshot_id'] == 'returned-transfer' and ack['optimizer_revision'] == 0 + assert ack['server_response'] == dict(service_instance_id='new-service', service_generation=1) + + +@pytest.mark.parametrize('old_request', [False, True]) +def test_new_nccl_request_during_launch_is_not_missed(supervisor, monkeypatch, old_request): + s = supervisor + if old_request: + s.module.write_json_atomic(s.request_path, s.request) + monkeypatch.setattr(s.module, 'wait_until_ready', + lambda *args: s.module.write_json_atomic(s.request_path, s.request)) + assert s.module.main() == 130 + assert len(s.calls) == 1 + assert json.loads(s.ack.read_text())['server_response']['service_generation'] == 1 + + +def test_existing_disk_request_keeps_its_original_reload_behavior(supervisor, tmp_path, monkeypatch): + s = supervisor + checkpoint = tmp_path / 'checkpoint' + checkpoint.mkdir() + s.module.write_json_atomic(s.request_path, + dict(mode='checkpoint', version=1, checkpoint_path=str(checkpoint))) + disk_calls = [] + monkeypatch.setattr(s.module, 'reload_inplace', lambda *args: disk_calls.append(args) or {}) + assert s.module.main() == 130 + assert disk_calls == [(s.args.health_url, str(checkpoint), s.args.ready_timeout)] + assert not s.calls + assert json.loads((tmp_path / 'ack_returned_1.json').read_text())['version'] == 1 diff --git a/tests/test_serving_protocol.py b/tests/test_serving_protocol.py new file mode 100644 index 0000000..eaa805d --- /dev/null +++ b/tests/test_serving_protocol.py @@ -0,0 +1,386 @@ +"""Shared serving correctness across physical TP workers and request lifetimes.""" + +import asyncio +import importlib.util +import json +from pathlib import Path +from types import SimpleNamespace as NS +import sys + +import pytest +import torch + +from RL_Framework.scripts.serving_protocol import ServingState, GenerationMiddleware +from RL_Framework.vllm_hot_reload import InplaceReloadWorkerExtension + + +def workers(tp, generation=1, tensors=7, offset=0): + return [dict(rank=rank, nccl_rank=1+offset+rank, service_generation=generation, + received_tensors=tensors, loaded_parameters=tensors, receive_complete=True) for rank in range(tp)] + + +@pytest.mark.parametrize('tp', [1, 2, 4]) +def test_initial_and_all_tp_publication_identity(tp): + async def run(): + state = ServingState(tp, '/models/test') + assert state.status()['server_generation_verified'] + assert state.optimizer_revision == state.server_generation == 0 + await state.begin_reload(0) + state.install(snapshot_id='revision100', optimizer_revision=100, workers=workers(tp, offset=4), rank_offset=4) + await state.reopen() + status = state.status() + assert status['server_generation'] == 1 and status['optimizer_revision'] == 100 + assert status['workers'] == workers(tp, offset=4) + assert status['server_generation_verified'] + asyncio.run(run()) + + +@pytest.mark.parametrize('mutate', [ + lambda rows: rows.pop(), + lambda rows: rows[1].update(rank=0), + lambda rows: rows[1].update(service_generation=0), + lambda rows: rows[1].update(receive_complete=False), + lambda rows: rows[1].update(received_tensors=6), + lambda rows: rows[1].update(nccl_rank=20), +]) +def test_incomplete_stale_or_inconsistent_tp_receipts_cannot_advance(mutate): + async def run(): + state = ServingState(2) + await state.begin_reload(0) + rows = workers(2) + mutate(rows) + with pytest.raises(RuntimeError): + state.install(snapshot_id='new', optimizer_revision=1, workers=rows) + await state.fail('invalid receipt') + assert not state.status()['server_generation_verified'] + assert state.server_generation == 0 + with pytest.raises(RuntimeError): + await state.reopen() + with pytest.raises(RuntimeError): + await state.enter() + asyncio.run(run()) + + +def test_reload_drains_inflight_requests_before_install(): + async def run(): + state = ServingState(2) + assert await state.enter() == 0 + installing = asyncio.create_task(state.begin_reload(0)) + await asyncio.sleep(0) + assert state.closed and not installing.done() + waiting = asyncio.create_task(state.enter()) + await asyncio.sleep(0) + assert not waiting.done() + await state.leave() + await installing + state.install(snapshot_id='new', optimizer_revision=1, workers=workers(2)) + await state.reopen() + assert await waiting == 1 + await state.leave() + assert state.active_requests == 0 + asyncio.run(run()) + + +def test_positive_drain_timeout_reopens_before_any_mutation(): + async def run(): + state = ServingState(1) + await state.enter() + with pytest.raises(TimeoutError): + await state.begin_reload(.001) + assert not state.closed and state.failed is None and state.server_generation == 0 + await state.leave() + asyncio.run(run()) + + +def test_generation_headers_attest_completion_without_changing_request_body(): + async def run(): + state, sent = ServingState(1), [] + received_payload = [] + async def app(scope, receive, send): + received_payload.append(json.loads((await receive())['body'])) + await send(dict(type='http.response.start', status=200, headers=[(b'content-type', b'application/json')])) + await send(dict(type='http.response.body', body=b'{"choices":[]}')) + middleware = GenerationMiddleware(app, state) + payload = dict(prompt=[1, 2, 3], seed=42, temperature=.8, top_p=.95, custom_field='unchanged') + async def receive(): + return dict(type='http.request', body=json.dumps(payload).encode(), more_body=False) + async def send(message): + sent.append(message) + scope = dict(type='http', method='POST', path='/v1/completions', + headers=[(b'x-libra-request-id', b'actual-request')]) + await middleware(scope, receive, send) + headers = dict(sent[0]['headers']) + assert headers[b'x-libra-request-id'] == b'actual-request' + assert headers[b'x-libra-service-generation'] == b'0' + assert headers[b'x-libra-service-generation-verified'] == b'1' + assert headers[b'x-libra-reload-overlap'] == b'0' + assert state.active_requests == 0 + assert received_payload == [payload] + asyncio.run(run()) + + +@pytest.mark.parametrize('consume_all', [True, False]) +def test_worker_attests_actual_source_tensor_count_and_done_marker(monkeypatch, consume_all): + from RL_Framework.infra.sync import nccl_weight_sync + monkeypatch.setattr(torch.cuda, 'current_device', lambda: 0) + monkeypatch.setattr(torch.cuda, 'synchronize', lambda: None) + def source(spec): + for i in range(3): + yield f'weight{i}', torch.ones(1) + monkeypatch.setattr(nccl_weight_sync, 'receive_weights', source) + worker = InplaceReloadWorkerExtension() + worker.rank = 1 + def load(weights): + if consume_all: + return {name for name, tensor in weights} + return {next(weights)[0]} + worker.model_runner = NS(model=NS(load_weights=load)) + if not consume_all: + with pytest.raises(RuntimeError, match='done marker'): + worker.reload_weights_nccl('host', 1, 3, 0) + assert getattr(worker, '_libra_service_generation', 0) == 0 + else: + result = worker.reload_weights_nccl('host', 1, 3, 0, expected_tensor_count=3) + assert result['receive_complete'] and result['received_tensors'] == 3 + assert result['service_generation'] == 1 and result['nccl_rank'] == 2 + + +def test_supervisor_echoes_snapshot_and_revision_into_physical_ack(tmp_path, monkeypatch): + path = Path(__file__).resolve().parents[1] / 'scripts/restartable_vllm_server.py' + spec = importlib.util.spec_from_file_location('serving_supervisor_test', path) + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + request = dict(mode='nccl', version=100, snapshot_id='final', optimizer_revision=100, + topology_epoch=8, nccl_host='train', nccl_port=29620, nccl_world_size=5, + nccl_rank_offsets={'current': 0}, nccl_chunk_bytes=1024) + args = NS(instance_id='current', control_dir=str(tmp_path), health_url='http://service/health', + initial_model='/models/test', ready_timeout=0, reload_strategy='parallel', reload_method='inplace', + poll_interval=.01, command=['server']) + monkeypatch.setattr(module, 'parse_args', lambda: args) + child = NS(poll=lambda: None) + monkeypatch.setattr(module.subprocess, 'Popen', lambda *args, **kwargs: child) + monkeypatch.setattr(module, 'wait_until_ready', + lambda *args: module.write_json_atomic(tmp_path / 'reload_request.json', request)) + monkeypatch.setattr(module, 'stop_process', lambda *args: None) + monkeypatch.setattr(module.signal, 'signal', lambda *args: None) + forwarded = [] + def reload(url, payload, timeout): + forwarded.append(payload) + return dict(workers=workers(4), service_instance_id='actual', service_generation=1) + monkeypatch.setattr(module, 'reload_nccl', reload) + calls = 0 + def sleep(seconds): + nonlocal calls + calls += 1 + if calls == 1: + request['snapshot_id'] = 'same-params-new-topology' + request['topology_epoch'] = 9 + (tmp_path / 'reload_request.json').write_text(json.dumps(request)) + else: + raise KeyboardInterrupt() + monkeypatch.setattr(module.time, 'sleep', sleep) + assert module.main() == 130 + assert [payload['snapshot_id'] for payload in forwarded] == ['final', 'same-params-new-topology'] + ack = json.loads((tmp_path / 'ack_current_100.json').read_text()) + assert ack['snapshot_id'] == 'same-params-new-topology' and ack['optimizer_revision'] == 100 + assert ack['topology_epoch'] == 9 and ack['server_response']['service_instance_id'] == 'actual' + + + +def load_serving_app(monkeypatch, tp=4, invalid_ack=False, request_started=None, request_release=None): + """Import the actual FastAPI routes with only the unavailable vLLM runtime stubbed.""" + import argparse + import types + from fastapi import FastAPI, Request + from fastapi.responses import JSONResponse + modules = {} + for name in ('vllm', 'vllm.entrypoints', 'vllm.entrypoints.openai', + 'vllm.entrypoints.openai.api_server', 'vllm.entrypoints.openai.cli_args', + 'vllm.utils', 'vllm.v1', 'vllm.v1.engine', 'vllm.v1.engine.core'): + modules[name] = types.ModuleType(name) + monkeypatch.setitem(sys.modules, name, modules[name]) + modules['vllm'].__version__ = '0.9.2' + modules['vllm.entrypoints.openai'].api_server = modules['vllm.entrypoints.openai.api_server'] + modules['vllm.entrypoints.openai.cli_args'].make_arg_parser = lambda parser: parser + modules['vllm.entrypoints.openai.cli_args'].validate_parsed_serve_args = lambda args: None + modules['vllm.utils'].FlexibleArgumentParser = argparse.ArgumentParser + modules['vllm.v1.engine.core'].EngineCore = type('Core', (), {}) + monkeypatch.setitem(sys.modules, 'uvloop', NS(run=lambda coroutine: None)) + def build(args): + app = FastAPI() + @app.post('/v1/completions') + async def completion(request: Request): + if request_started is not None: + request_started.set() + await request_release.wait() + return JSONResponse(dict(choices=[dict(text='ok')], received=await request.json())) + return app + async def utility(name): + assert name == 'libra_reset_prefix_cache' + return True + async def rpc(name, timeout, args=()): + assert timeout is None or timeout == 120 + if name == 'reload_weights_nccl': + rows = workers(tp, offset=args[3]) + return rows[:-1] if invalid_ack else rows + raise AssertionError(name) + client = NS(collective_rpc=rpc, engine_core=NS(call_utility_async=utility)) + api = modules['vllm.entrypoints.openai.api_server'] + api.build_app = build + api.engine_client = lambda request: request.app.state.client + path = Path(__file__).resolve().parents[1] / 'scripts/vllm_hot_reload_api_server.py' + spec = importlib.util.spec_from_file_location('libra_serving_routes_test', path) + module = importlib.util.module_from_spec(spec) + monkeypatch.setitem(sys.modules, spec.name, module) + spec.loader.exec_module(module) + app = module.build_app_with_reload_routes(NS(tensor_parallel_size=tp, model='/models/test')) + app.state.client = client + return app + + +def test_real_http_routes_drain_requests_and_attest_physical_workers(monkeypatch): + import httpx + async def run(): + started, release = asyncio.Event(), asyncio.Event() + app = load_serving_app(monkeypatch, request_started=started, request_release=release) + async with httpx.AsyncClient(transport=httpx.ASGITransport(app=app), base_url='http://service') as client: + initial = (await client.get('/publication/status')).json() + assert initial['server_generation_verified'] and initial['tp_size'] == 4 + generating = asyncio.create_task(client.post('/v1/completions', + headers={'X-Libra-Request-Id': 'old-request'}, json=dict(prompt=[1]))) + await started.wait() + publishing = asyncio.create_task(client.post('/reload_weights_nccl', json=dict( + host='training', port=29620, world_size=5, rank_offset=0, + snapshot_id='final100', optimizer_revision=100, timeout_seconds=0))) + await asyncio.sleep(0) + assert not publishing.done() + waiting = asyncio.create_task(client.post('/v1/completions', + headers={'X-Libra-Request-Id': 'new-request'}, json=dict(prompt=[1, 2], seed=42))) + await asyncio.sleep(0) + assert not waiting.done() + release.set() + old = await generating + assert old.headers['x-libra-service-generation'] == '0' + reply = await publishing + assert reply.status_code == 200 + assert reply.json()['prefix_cache_cleared'] + assert len(reply.json()['workers']) == 4 + live = (await client.get('/publication/status')).json() + assert live['optimizer_revision'] == 100 and live['snapshot_id'] == 'final100' + assert live['server_generation'] == 1 and live['server_generation_verified'] + assert live['workers'] == reply.json()['workers'] + response = await waiting + assert response.headers['x-libra-service-generation-verified'] == '1' + assert response.headers['x-libra-service-generation'] == '1' + assert response.json()['received'] == dict(prompt=[1, 2], seed=42) + asyncio.run(run()) + + +def test_real_http_partial_tp_failure_quarantines_service(monkeypatch): + import httpx + app = load_serving_app(monkeypatch, invalid_ack=True) + async def run(): + async with httpx.AsyncClient(transport=httpx.ASGITransport(app=app, raise_app_exceptions=False), base_url='http://service') as client: + response = await client.post('/reload_weights_nccl', json=dict(host='training', port=29620, + world_size=5, rank_offset=0, snapshot_id='broken', optimizer_revision=1, timeout_seconds=0)) + assert response.status_code == 500 + state = (await client.get('/publication/status')).json() + assert state['failed'] and not state['server_generation_verified'] + assert state['server_generation'] == 0 + assert (await client.post('/v1/completions', json=dict(prompt=[1]))).status_code == 409 + asyncio.run(run()) + + +@pytest.mark.parametrize('session_id', ['', 'run-global-new-members']) +@pytest.mark.parametrize('tensor_count', [None, 3]) +def test_nccl_session_wire_from_request_file_through_http_rpc_and_worker_spec( + tmp_path, monkeypatch, session_id, tensor_count): + import httpx + from RL_Framework.infra.sync import nccl_weight_sync + app = load_serving_app(monkeypatch, tp=1) + service_id = app.state.serving_state.server_instance_id + specs, payloads, rpc_args = [], [], [] + worker = InplaceReloadWorkerExtension() + worker.rank = 0 + worker.model_runner = NS(model=NS(load_weights=lambda source: {name for name, tensor in source})) + monkeypatch.setattr(torch.cuda, 'current_device', lambda: 0) + monkeypatch.setattr(torch.cuda, 'synchronize', lambda: None) + + def source(spec): + specs.append(spec) + yield from ((f'weight{i}', torch.ones(1)) for i in range(3)) + + monkeypatch.setattr(nccl_weight_sync, 'receive_weights', source) + + async def rpc(name, timeout, args=()): + assert name == 'reload_weights_nccl' and timeout is None + rpc_args.append(args) + return [worker.reload_weights_nccl(*args)] + + app.state.client.collective_rpc = rpc + path = Path(__file__).resolve().parents[1] / 'scripts/restartable_vllm_server.py' + spec = importlib.util.spec_from_file_location('session_wire_supervisor_test', path) + supervisor = importlib.util.module_from_spec(spec) + spec.loader.exec_module(supervisor) + request = dict(mode='nccl', version=2, snapshot_id='after-return', optimizer_revision=2, + nccl_host='train', nccl_port=29620, nccl_world_size=2, + nccl_rank_offsets={'returned': 0}, nccl_chunk_bytes=1024) + if tensor_count is not None: + request['expected_tensor_count'] = tensor_count + if session_id: + request.update(nccl_session_id=session_id) + args = NS(instance_id='returned', control_dir=str(tmp_path), health_url='http://service/health', + initial_model='/model', ready_timeout=0, reload_strategy='parallel', reload_method='inplace', + poll_interval=.01, command=['server']) + monkeypatch.setattr(supervisor, 'parse_args', lambda: args) + monkeypatch.setattr(supervisor.subprocess, 'Popen', lambda *args, **kwargs: NS(poll=lambda: None)) + monkeypatch.setattr(supervisor, 'wait_until_ready', lambda *args: + supervisor.write_json_atomic(tmp_path / 'reload_request.json', request)) + monkeypatch.setattr(supervisor, 'stop_process', lambda *args: None) + monkeypatch.setattr(supervisor.signal, 'signal', lambda *args: None) + + def urlopen(message, timeout): + payload = json.loads(message.data) + payloads.append(payload) + assert message.full_url == 'http://service/reload_weights_nccl' + + async def post(): + async with httpx.AsyncClient(transport=httpx.ASGITransport(app=app), + base_url='http://service') as client: + response = await client.post('/reload_weights_nccl', json=payload) + assert response.status_code == 200, response.text + return response.content + + body = asyncio.run(post()) + class Response: + status = 200 + def __enter__(self): + return self + def __exit__(self, *args): + return False + def read(self): + return body + return Response() + + def stop(seconds): + raise KeyboardInterrupt() + + monkeypatch.setattr(supervisor.urllib.request, 'urlopen', urlopen) + monkeypatch.setattr(supervisor.time, 'sleep', stop) + assert supervisor.main() == 130 + assert len(payloads) == len(rpc_args) == len(specs) == 1 + if session_id: + assert payloads[0]['session_id'] == session_id + assert rpc_args[0][-2:] == (tensor_count, session_id) + else: + assert 'session_id' not in payloads[0] + assert len(rpc_args[0]) == (6 if tensor_count is None else 7) + assert specs[0].session_id == session_id + assert specs[0].rank == 1 and specs[0].world_size == 2 + ack = json.loads((tmp_path / 'ack_returned_2.json').read_text()) + reply = ack['server_response'] + assert reply['service_instance_id'] == service_id + assert reply['prefix_cache_cleared'] and reply['workers'][0]['receive_complete'] + state = app.state.serving_state.status() + assert state['server_generation_verified'] and state['optimizer_revision'] == 2 + assert state['server_generation'] == 1 and not state['closed'] diff --git a/trainer/async_rl_trainer.py b/trainer/async_rl_trainer.py index 48d89a8..f3efcdd 100644 --- a/trainer/async_rl_trainer.py +++ b/trainer/async_rl_trainer.py @@ -1319,8 +1319,6 @@ def _sync_weights_nccl(self, target_version: int) -> None: if not control_dir_value: raise ValueError("NCCL weight synchronization requires rollout_weight_sync_control_dir") control_dir = Path(control_dir_value).resolve() - if self.is_main_process: - control_dir.mkdir(parents=True, exist_ok=True) import socket @@ -1351,30 +1349,47 @@ def _sync_weights_nccl(self, target_version: int) -> None: ), "created_at": time.time(), } + preparation_error = None if self.is_main_process: - for instance_id in offsets: - (control_dir / f"ack_{instance_id}_{target_version}.json").unlink( - missing_ok=True - ) - (control_dir / f"error_{instance_id}_{target_version}.json").unlink( - missing_ok=True + try: + from RL_Framework.infra.sync.publication_session import prepare_nccl_session + + request.update( + prepare_nccl_session(self, offsets), + snapshot_id=uuid.uuid4().hex, + optimizer_revision=int(self.train_engine.optimizer_revision), ) - self._write_json_atomic(control_dir / "reload_request.json", request) + control_dir.mkdir(parents=True, exist_ok=True) + for instance_id in offsets: + (control_dir / f"ack_{instance_id}_{target_version}.json").unlink( + missing_ok=True + ) + (control_dir / f"error_{instance_id}_{target_version}.json").unlink( + missing_ok=True + ) + self._write_json_atomic(control_dir / "reload_request.json", request) + except Exception as exc: + preparation_error = f"{type(exc).__name__}: {exc}" + preparation = [preparation_error] + if dist.is_initialized(): + dist.broadcast_object_list(preparation, src=0, group=getattr(self, "control", None)) + if preparation[0]: + raise RuntimeError(f"Rollout publication preparation failed: {preparation[0]}") from RL_Framework.infra.sync.nccl_weight_sync import ( NcclReloadSpec, send_weights, ) - weights = self.train_engine.stream_rollout_weights() started = time.perf_counter() transferred_bytes = 0 transport_keepalive: list[object] = [] - if self.is_main_process: - def record_tensor(_name: str, size_bytes: int) -> None: - nonlocal transferred_bytes - transferred_bytes += size_bytes + def record_tensor(_name: str, size_bytes: int) -> None: + nonlocal transferred_bytes + transferred_bytes += size_bytes + + def send_export(weights): send_weights( weights, NcclReloadSpec( @@ -1387,6 +1402,7 @@ def record_tensor(_name: str, size_bytes: int) -> None: getattr(self.config, "rollout_weight_sync_timeout_s", 1200.0) ), chunk_bytes=request["nccl_chunk_bytes"], + session_id=request["nccl_session_id"], rate_limit_gbps=float( getattr(self.config, "rollout_nccl_rate_limit_gbps", 0.0) ), @@ -1394,55 +1410,67 @@ def record_tensor(_name: str, size_bytes: int) -> None: on_tensor=record_tensor, keepalive=transport_keepalive, ) - else: - # The export itself contains Megatron TP/EP collectives. - for _name, tensor in weights: - del tensor + from RL_Framework.infra.sync.collective_export import stream_export_with_consensus + stream_export_with_consensus( + self.train_engine.stream_rollout_weights, + is_sender=self.is_main_process, send=send_export, + group=getattr(self, "control", None), + ) + + failure = None if self.is_main_process: - timeout = float(getattr(self.config, "rollout_weight_sync_timeout_s", 1200.0)) - deadline = time.monotonic() + timeout - pending = set(offsets) - while pending and time.monotonic() < deadline: - failures = { - instance_id: (control_dir / f"error_{instance_id}_{target_version}.json").read_text( - encoding="utf-8" - ) - for instance_id in pending - if (control_dir / f"error_{instance_id}_{target_version}.json").exists() - } - if failures: - raise RuntimeError(f"vLLM NCCL reload failed: {failures}") - pending = { - instance_id - for instance_id in pending - if not (control_dir / f"ack_{instance_id}_{target_version}.json").exists() - } - if pending: - time.sleep( - float( - getattr( - self.config, - "rollout_weight_sync_poll_interval_s", - 0.05, + try: + timeout = float(getattr(self.config, "rollout_weight_sync_timeout_s", 1200.0)) + deadline = time.monotonic() + timeout if timeout > 0 else float("inf") + pending = set(offsets) + while pending and time.monotonic() < deadline: + failures = { + instance_id: (control_dir / f"error_{instance_id}_{target_version}.json").read_text( + encoding="utf-8" + ) + for instance_id in pending + if (control_dir / f"error_{instance_id}_{target_version}.json").exists() + } + if failures: + raise RuntimeError(f"vLLM NCCL reload failed: {failures}") + pending = { + instance_id + for instance_id in pending + if not (control_dir / f"ack_{instance_id}_{target_version}.json").exists() + } + if pending: + time.sleep( + float( + getattr( + self.config, + "rollout_weight_sync_poll_interval_s", + 0.05, + ) ) ) - ) - if pending: - raise TimeoutError(f"Timed out waiting for NCCL reload: {sorted(pending)}") - self._record_rollout_weight_version( - target_version=target_version, - checkpoint_path="", - instance_ids=list(offsets), - ) - print( - f"Direct NCCL rollout refresh complete: version={target_version}, " - f"world_size={world_size}, model_bytes={transferred_bytes}, " - f"fanout_bytes={transferred_bytes * (world_size - 1)}, " - f"seconds={time.perf_counter() - started:.3f}", - flush=True, - ) - transport_keepalive.clear() + if pending: + raise TimeoutError(f"Timed out waiting for NCCL reload: {sorted(pending)}") + self._record_rollout_weight_version( + target_version=target_version, + checkpoint_path="", + instance_ids=list(offsets), + ) + print( + f"Direct NCCL rollout refresh complete: version={target_version}, " + f"world_size={world_size}, model_bytes={transferred_bytes}, " + f"fanout_bytes={transferred_bytes * (world_size - 1)}, " + f"seconds={time.perf_counter() - started:.3f}", + flush=True, + ) + transport_keepalive.clear() + except Exception as exc: + failure = f"{type(exc).__name__}: {exc}" + outcome = [failure] + if dist.is_initialized(): + dist.broadcast_object_list(outcome, src=0, group=getattr(self, "control", None)) + if outcome[0]: + raise RuntimeError(f"Rollout publication failed: {outcome[0]}") if dist.is_initialized(): dist.barrier() self._record_weight_sync_phase("sync_complete") diff --git a/vllm_hot_reload.py b/vllm_hot_reload.py index 5cc55f2..012afd7 100644 --- a/vllm_hot_reload.py +++ b/vllm_hot_reload.py @@ -48,7 +48,7 @@ def reload_weights(self, checkpoint_path: str) -> dict[str, Any]: self.vllm_config.model_config.model = previous_path raise finished_at = time.time() - return { + result = { "rank": int(getattr(self, "rank", 0)), "checkpoint_path": str(checkpoint), "previous_path": previous_path, @@ -56,6 +56,7 @@ def reload_weights(self, checkpoint_path: str) -> dict[str, Any]: "finished_at": finished_at, "load_seconds": time.perf_counter() - started, } + return self._attest_reload(result) def reload_weights_nccl( self, @@ -65,6 +66,8 @@ def reload_weights_nccl( rank_offset: int, timeout_seconds: float = 1200.0, chunk_bytes: int = 256 * 1024 * 1024, + expected_tensor_count: int | None = None, + session_id: str = '', ) -> dict[str, Any]: """Load full HF-format weights from a Megatron NCCL broadcast.""" import torch @@ -83,13 +86,38 @@ def reload_weights_nccl( device=torch.cuda.current_device(), timeout_s=float(timeout_seconds), chunk_bytes=int(chunk_bytes), + session_id=session_id, ) started = time.perf_counter() - loaded = self.model_runner.model.load_weights(receive_weights(spec)) + count, complete = 0, False + source = receive_weights(spec) + def observed(): + nonlocal count, complete + for name, tensor in source: + count += 1 + yield name, tensor + complete = True + weights = observed() + try: + loaded = self.model_runner.model.load_weights(weights) + if not complete: + raise RuntimeError('Model loader returned before the NCCL done marker') + if expected_tensor_count is not None and count != expected_tensor_count: + raise RuntimeError('Received source tensor count differs from the snapshot') + finally: + weights.close() + source.close() torch.cuda.synchronize() - return { + result = { "rank": local_rank, "nccl_rank": spec.rank, "loaded_parameters": len(loaded) if loaded is not None else 0, + "received_tensors": count, + "receive_complete": complete, "load_seconds": time.perf_counter() - started, } + return self._attest_reload(result) + + def _attest_reload(self, result): + self._libra_service_generation = getattr(self, '_libra_service_generation', 0) + 1 + return dict(result, service_generation=self._libra_service_generation) From 8e38c5c0717eecb35ebcc6a176d7ba727a9c328f Mon Sep 17 00:00:00 2001 From: cytsoi <154662279+tcyjohn@users.noreply.github.com> Date: Thu, 22 Oct 2026 01:34:53 +0800 Subject: [PATCH 2/9] fix: retain accepted rollout work across publication pauses --- infra/execution/async_runner.py | 102 +++++++++++++++------------- infra/execution/batch_dispatcher.py | 25 +++---- scripts/test_cpu_ci.sh | 8 +++ tests/test_async_task_runner.py | 55 ++++++++++++++- tests/test_sync_drain_prefetch.py | 10 ++- trainer/async_rl_trainer.py | 4 +- 6 files changed, 140 insertions(+), 64 deletions(-) diff --git a/infra/execution/async_runner.py b/infra/execution/async_runner.py index c362eb3..781485a 100644 --- a/infra/execution/async_runner.py +++ b/infra/execution/async_runner.py @@ -36,6 +36,8 @@ def __init__( self.exiting = threading.Event() self.paused = threading.Event() + self._draining = threading.Event() + self._admission_lock = threading.Lock() self.input_queue: queue.Queue = queue.Queue(maxsize=max_queue_size) @@ -62,6 +64,7 @@ def initialize(self): """Initialize.""" self.exiting.clear() self.paused.clear() + self._draining.clear() self._loop_ready = threading.Event() self.thread = threading.Thread( @@ -183,11 +186,9 @@ async def _run_async_loop( try: while not exit_event.is_set(): - # Pausing blocks new dispatcher submissions. Tasks already - # accepted by this runner must still be started and reaped so - # callers can drain the rollout pipeline before restarting - # serving processes for a weight update. - if not pause_event.is_set(): + # Explicit drain starts accepted inputs. Ordinary pause keeps + # them available for cancel_queued(). + if not pause_event.is_set() or self._draining.is_set(): self._drain_pending_inputs(running_tasks, input_queue) @@ -316,33 +317,33 @@ def submit( ) -> int: """Submit.""" - with self._thread_exception_lock: - if self._thread_exception is not None: - raise RuntimeError(f"AsyncTaskRunner is unusable: {self._thread_exception}") - - - with self._active_task_ids_lock: - if task_id in self._active_task_ids: - raise ValueError(f"duplicate task ID: {task_id}") - self._active_task_ids.add(task_id) - - - fn_kwargs = {k: v for k, v in kwargs.items() if k != 'task_id'} - - task_input = { - "async_fn": async_fn, - "args": args, - "kwargs": fn_kwargs, - "task_id": task_id, - "create_time": time.monotonic_ns(), - } + with self._admission_lock: + if self.paused.is_set(): + raise queue.Full("runner admission is paused") + with self._thread_exception_lock: + if self._thread_exception is not None: + raise RuntimeError(f"AsyncTaskRunner is unusable: {self._thread_exception}") - try: - self.input_queue.put_nowait(task_input) - except queue.Full: with self._active_task_ids_lock: - self._active_task_ids.discard(task_id) - raise queue.Full("input queue is full") + if task_id in self._active_task_ids: + raise ValueError(f"duplicate task ID: {task_id}") + self._active_task_ids.add(task_id) + + fn_kwargs = {k: v for k, v in kwargs.items() if k != 'task_id'} + task_input = { + "async_fn": async_fn, + "args": args, + "kwargs": fn_kwargs, + "task_id": task_id, + "create_time": time.monotonic_ns(), + } + + try: + self.input_queue.put_nowait(task_input) + except queue.Full: + with self._active_task_ids_lock: + self._active_task_ids.discard(task_id) + raise queue.Full("input queue is full") self._signal_new_input() return task_id @@ -387,11 +388,13 @@ def wait( def pause(self): """Pause.""" - self.paused.set() + with self._admission_lock: + self.paused.set() def resume(self): """Resume.""" - self.paused.clear() + with self._admission_lock: + self.paused.clear() self._signal_new_input() def cancel_queued(self, timeout: float = 2.0) -> list[int]: @@ -426,22 +429,27 @@ def drain_input_queue() -> None: def wait_until_idle(self, timeout: float = 3600.0) -> None: """Wait until every task accepted before pause has completed.""" deadline = time.monotonic() + timeout - while True: - with self._thread_exception_lock: - if self._thread_exception is not None: - raise RuntimeError( - f"AsyncTaskRunner error: {self._thread_exception}" + self._draining.set() + self._signal_new_input() + try: + while True: + with self._thread_exception_lock: + if self._thread_exception is not None: + raise RuntimeError( + f"AsyncTaskRunner error: {self._thread_exception}" + ) + with self._active_task_ids_lock: + active_count = len(self._active_task_ids) + if active_count == 0 and self.input_queue.empty(): + return + if time.monotonic() >= deadline: + raise TimeoutError( + "Timed out draining rollout tasks before weight reload: " + f"active={active_count}, queued={self.input_queue.qsize()}" ) - with self._active_task_ids_lock: - active_count = len(self._active_task_ids) - if active_count == 0 and self.input_queue.empty(): - return - if time.monotonic() >= deadline: - raise TimeoutError( - "Timed out draining rollout tasks before weight reload: " - f"active={active_count}, queued={self.input_queue.qsize()}" - ) - time.sleep(0.05) + time.sleep(0.05) + finally: + self._draining.clear() def get_queue_sizes(self) -> tuple[int, int]: """Get queue sizes.""" diff --git a/infra/execution/batch_dispatcher.py b/infra/execution/batch_dispatcher.py index 600d662..cbf99c3 100644 --- a/infra/execution/batch_dispatcher.py +++ b/infra/execution/batch_dispatcher.py @@ -147,16 +147,17 @@ def _commit_loop(self) -> None: with self._reconfigure_lock: - # pause() can race the small window between dequeue and - # submit. Recheck under the same lock used by - # cancel_queued() so an unsent request never crosses the - # weight-sync boundary. - if self.runner.paused.is_set(): - self.staleness_manager.on_rollouts_cancelled(enqueued=1) - with self._result_cv: - self._active_task_ids.discard(task_input.task_id) - self._result_cv.notify_all() - continue + # Keep the selected task owned by the pending queue until + # commit, so cancel_queued() can still return its identity + # if pause wins the race for this lock. + with self._input_cv: + if ( + not self._pending_inputs + or self._pending_inputs[0] is not task_input + or self.runner.paused.is_set() + ): + continue + self._pending_inputs.popleft() task_fn = self.task_factory(task_input) try: @@ -191,7 +192,7 @@ def _commit_loop(self) -> None: break def _get_next_task_for_submission(self) -> TInput | None: - """Get next task for submission.""" + """Wait for a candidate, retaining queue ownership until commit.""" with self._input_cv: while not self._shutdown_event.is_set(): self._check_thread_exception() @@ -202,7 +203,7 @@ def _get_next_task_for_submission(self) -> TInput | None: and self.staleness_manager.get_capacity() > 0 and self._pending_inputs ): - return self._pending_inputs.popleft() + return self._pending_inputs[0] self._input_cv.wait(timeout=0.5) diff --git a/scripts/test_cpu_ci.sh b/scripts/test_cpu_ci.sh index ee6430f..8b90c37 100644 --- a/scripts/test_cpu_ci.sh +++ b/scripts/test_cpu_ci.sh @@ -6,8 +6,16 @@ python -m pytest -q --strict-markers --junitxml=reports/cpu-tests.xml \ tests/test_qwen3_config.py \ tests/test_planner_global_batch_workload.py \ tests/test_async_task_runner.py \ + tests/test_sync_drain_prefetch.py \ tests/test_phase_tracer.py \ tests/test_restartable_vllm_server.py \ + tests/test_checked_prefix_cache_reset.py \ + tests/test_nccl_store_lifecycle.py \ + tests/test_nccl_store_timeout.py \ + tests/test_nccl_weight_sync.py \ + tests/test_publication_sessions.py \ + tests/test_serving_protocol.py \ + tests/test_collective_weight_export.py \ tests/test_cmlfq_scheduler.py \ tests/test_cmlfq_cost_scheduler.py \ tests/test_cost_model_units.py \ diff --git a/tests/test_async_task_runner.py b/tests/test_async_task_runner.py index 8e32aed..b0c6207 100644 --- a/tests/test_async_task_runner.py +++ b/tests/test_async_task_runner.py @@ -1,6 +1,10 @@ import asyncio +from contextlib import contextmanager +import queue import threading +import pytest + from RL_Framework.infra.execution.async_runner import AsyncTaskRunner @@ -108,8 +112,9 @@ async def mark_started(): runner = AsyncTaskRunner() runner.initialize() try: - runner.pause() - runner.submit(mark_started, task_id=10) + with _hold_event_loop(runner): + runner.submit(mark_started, task_id=10) + runner.pause() cancelled = runner.cancel_queued() @@ -119,3 +124,49 @@ async def mark_started(): runner.wait_until_idle(timeout=1.0) finally: runner.destroy() + + +@contextmanager +def _hold_event_loop(runner): + entered, release = threading.Event(), threading.Event() + + def block(): + entered.set() + release.wait(5.) + + runner._loop.call_soon_threadsafe(block) + assert entered.wait(1.) + try: + yield + finally: + release.set() + + + +def test_explicit_drain_runs_inputs_accepted_before_pause(): + runner = AsyncTaskRunner() + runner.initialize() + try: + # Freeze the real event loop before submission so these inputs are + # accepted but deterministically have not started when pause occurs. + with _hold_event_loop(runner): + runner.submit(_return_value, 'first', task_id=7) + runner.submit(_return_value, 'second', task_id=8) + runner.pause() + with pytest.raises(queue.Full, match='paused'): + runner.submit(_return_value, 'rejected', task_id=9) + assert runner.input_queue.qsize() == 2 + assert runner._active_task_ids == {7, 8} + + runner.wait_until_idle(timeout=1.) + assert {row.data for row in runner.wait(count=2, timeout=1.)} == {'first', 'second'} + assert runner.paused.is_set() + assert not runner._draining.is_set() + assert not runner._active_task_ids and runner.input_queue.empty() + with pytest.raises(queue.Full, match='paused'): + runner.submit(_return_value, 'still rejected', task_id=9) + runner.resume() + runner.submit(_return_value, 'after resume', task_id=9) + assert runner.wait(count=1, timeout=1.)[0].data == 'after resume' + finally: + runner.destroy() diff --git a/tests/test_sync_drain_prefetch.py b/tests/test_sync_drain_prefetch.py index 4fe1458..2135234 100644 --- a/tests/test_sync_drain_prefetch.py +++ b/tests/test_sync_drain_prefetch.py @@ -71,7 +71,7 @@ def test_sync_drain_prefetch_limit_uses_two_then_one_batch(): assert trainer._sync_aware_rollout_limit(step=2, batch_size=4) is None assert trainer._sync_aware_rollout_limit(step=3, batch_size=4) == 8 assert trainer._sync_aware_rollout_limit(step=4, batch_size=4) == 4 - assert trainer._sync_aware_rollout_limit(step=5, batch_size=4) == 4 + assert trainer._sync_aware_rollout_limit(step=5, batch_size=4) is None assert trainer._sync_aware_rollout_limit(step=6, batch_size=4) is None @@ -87,6 +87,9 @@ def __init__(self): def submit(self, _fn, *, task_id): self.submitted.append(task_id) + def cancel_queued(self): + return [] + runner = _Runner() manager = StalenessManager( _VersionProvider(), @@ -113,5 +116,10 @@ def dequeue_at_pause_boundary(): dispatcher._commit_loop() assert runner.submitted == [] + assert manager.get_stats().enqueued == 1 + assert manager.get_stats().cancelled == 0 + assert dispatcher.cancel_queued()['task_ids'] == [7] assert manager.get_stats().enqueued == 0 assert manager.get_stats().cancelled == 1 + assert dispatcher.cancel_queued()['task_ids'] == [] + assert manager.get_stats().cancelled == 1 diff --git a/trainer/async_rl_trainer.py b/trainer/async_rl_trainer.py index f3efcdd..87e3684 100644 --- a/trainer/async_rl_trainer.py +++ b/trainer/async_rl_trainer.py @@ -951,7 +951,7 @@ def train(self, workflow, dataset=None): weight_sync_time = 0.0 rollout_sync_drain = None - if self.config.sync_interval > 0 and step > 0 and step % self.config.sync_interval == 0: + if self.config.sync_interval > 0 and (step + 1) % self.config.sync_interval == 0: self._trace_train_phase(step, "weight_sync_start") if self.dispatcher is not None: drain_started = time.time() @@ -1115,7 +1115,7 @@ def _sync_aware_rollout_limit( if interval <= 0 or lead_steps <= 0: return None remainder = step % interval - steps_until_sync = 0 if step > 0 and remainder == 0 else interval - remainder + steps_until_sync = interval - remainder if steps_until_sync > lead_steps: return None prefetch_batches = max(1, steps_until_sync) From 9f4b49f31c9e0543e5ba83ed99c40d2779b7e1a1 Mon Sep 17 00:00:00 2001 From: cytsoi <154662279+tcyjohn@users.noreply.github.com> Date: Thu, 22 Oct 2026 01:34:53 +0800 Subject: [PATCH 3/9] fix: consume rollout credit only for accepted results --- scripts/test_cpu_ci.sh | 1 + tests/test_rollout_credit.py | 50 ++++++++++++++++++++++++++++++++++++ trainer/async_rl_trainer.py | 6 ++++- 3 files changed, 56 insertions(+), 1 deletion(-) create mode 100644 tests/test_rollout_credit.py diff --git a/scripts/test_cpu_ci.sh b/scripts/test_cpu_ci.sh index 8b90c37..d73bb53 100644 --- a/scripts/test_cpu_ci.sh +++ b/scripts/test_cpu_ci.sh @@ -7,6 +7,7 @@ python -m pytest -q --strict-markers --junitxml=reports/cpu-tests.xml \ tests/test_planner_global_batch_workload.py \ tests/test_async_task_runner.py \ tests/test_sync_drain_prefetch.py \ + tests/test_rollout_credit.py \ tests/test_phase_tracer.py \ tests/test_restartable_vllm_server.py \ tests/test_checked_prefix_cache_reset.py \ diff --git a/tests/test_rollout_credit.py b/tests/test_rollout_credit.py new file mode 100644 index 0000000..32dc826 --- /dev/null +++ b/tests/test_rollout_credit.py @@ -0,0 +1,50 @@ +"""A timed-out task releases running credit, never accepted credit.""" +import asyncio +from types import SimpleNamespace + +from RL_Framework.infra.execution.batch_dispatcher import TaskInput +from RL_Framework.infra.sync.staleness import StalenessManager +from RL_Framework.trainer.async_rl_trainer import AsyncRLTrainer + + +def test_timeout_notification_does_not_consume_successful_result_credit(monkeypatch): + trainer = AsyncRLTrainer.__new__(AsyncRLTrainer) + trainer.config = SimpleNamespace(n_samples=4, max_head_offpolicyness=3, + enable_rollout_tracing=False) + trainer.is_main_process = True + trainer.train_engine = SimpleNamespace(get_version=lambda: 0) + trainer.rollout_engine = object() + trainer._refresh_rollout_engine_from_manifest_for_task = lambda: None + trainer._pending_grpo_groups = {} + manager = trainer.staleness_manager = StalenessManager( + trainer.train_engine, max_concurrent_rollouts=8, + consumer_batch_size=4, max_staleness=3) + + async def episode(engine, data, **kwargs): + if data['timeout']: + await asyncio.Future() + return {} + + trainer.workflow = SimpleNamespace(run_episode=episode) + monkeypatch.setenv('ROLLOUT_TASK_TIMEOUT_S', '.001') + + async def produce(): + packets = [] + for index in range(5): + manager.on_rollout_enqueued() + manager.on_rollout_submitted() + task = TaskInput(index, {'timeout': index == 0}, version=0, + group_id='failed' if index == 0 else 'complete', + rollout_index=max(0, index - 1)) + packets.append(await trainer._create_rollout_task(task)()) + return packets + + packets = asyncio.run(produce()) + assert (manager.get_stats().accepted, manager.get_stats().rejected) == (4, 1) + chunks = [packets[:4], packets[4:]] + trainer.dispatcher = SimpleNamespace(active_submit_and_wait=lambda **kw: chunks.pop(0)) + batch = trainer._collect_complete_grpo_batch(iter(()), batch_size=4) + assert [row['grpo_rollout_index'] for row in batch] == [0, 1, 2, 3] + assert (manager.get_stats().accepted, manager.get_stats().running, + manager.get_stats().rejected) == (0, 0, 1) + assert manager.get_capacity() == 8 diff --git a/trainer/async_rl_trainer.py b/trainer/async_rl_trainer.py index 87e3684..45f267f 100644 --- a/trainer/async_rl_trainer.py +++ b/trainer/async_rl_trainer.py @@ -569,7 +569,11 @@ def _collect_complete_grpo_batch(self, data_gen, batch_size: int): return selected continue if self.staleness_manager is not None: - self.staleness_manager.on_batch_consumed(len(raw_batch)) + accepted_count = sum( + isinstance(trajectory, dict) and trajectory.get("status") != "failed" + for trajectory in raw_batch + ) + self.staleness_manager.on_batch_consumed(accepted_count) for trajectory in raw_batch: if trajectory.get("status") == "failed": continue From b21cf139fef67f2eaa320d4cfd5f4e22c46fe7a4 Mon Sep 17 00:00:00 2001 From: cytsoi <154662279+tcyjohn@users.noreply.github.com> Date: Thu, 22 Oct 2026 01:28:50 +0800 Subject: [PATCH 4/9] Bound native Adam gradient preparation and update temporaries --- engine/megatron_core_train_engine.py | 9 + engine/optimizer_memory.py | 265 +++++++++++++++++ tests/test_torch_optimizer_memory.py | 407 +++++++++++++++++++++++++++ 3 files changed, 681 insertions(+) create mode 100644 engine/optimizer_memory.py create mode 100644 tests/test_torch_optimizer_memory.py diff --git a/engine/megatron_core_train_engine.py b/engine/megatron_core_train_engine.py index 444ca9a..971593b 100644 --- a/engine/megatron_core_train_engine.py +++ b/engine/megatron_core_train_engine.py @@ -17,6 +17,7 @@ from RL_Framework.engine.megatron_core_checkpointing import ( MegatronDistributedCheckpointManager, ) +from RL_Framework.engine.optimizer_memory import bound_default_torch_adam_memory from RL_Framework.engine.device_utils import set_device from RL_Framework.infra.elastic import ( GradientPayload, @@ -260,6 +261,10 @@ def init_phase(name: str) -> None: self.model, use_gloo_process_groups=False, ) + bound_default_torch_adam_memory(self.optimizer, + release_main_gradients=(getattr(self, 'use_distributed_optimizer', False) + and not getattr(self, 'grad_reduce_in_fp32', False) + and not getattr(self, 'optimizer_cpu_offload', False))) if self.is_main_process: print("MCORE_INIT_PHASE optimizer_build_complete", flush=True) if self.is_main_process: @@ -575,6 +580,10 @@ def load_weights(self, path: str, version: int): optimizer=self.optimizer, version=version, ) + bound_default_torch_adam_memory(self.optimizer, + release_main_gradients=(getattr(self, 'use_distributed_optimizer', False) + and not getattr(self, 'grad_reduce_in_fp32', False) + and not getattr(self, 'optimizer_cpu_offload', False))) self.current_version = version def get_parallel_state(self) -> dict[str, Any]: diff --git a/engine/optimizer_memory.py b/engine/optimizer_memory.py new file mode 100644 index 0000000..0d2e083 --- /dev/null +++ b/engine/optimizer_memory.py @@ -0,0 +1,265 @@ +"""Bound temporary storage in the native PyTorch Adam fallback.""" + +from functools import wraps + +import torch +from torch.optim import optimizer as optimizer_module + + +def _native_streaming_eligible(optimizer): + seen = set() + if (getattr(optimizer, '_optimizer_step_pre_hooks', {}) + or getattr(optimizer, '_optimizer_step_post_hooks', {}) + or getattr(optimizer_module, '_global_optimizer_pre_hooks', {}) + or getattr(optimizer_module, '_global_optimizer_post_hooks', {})): + return False + for group in optimizer.param_groups: + if (group.get('foreach') is not False or group.get('fused') + or group.get('capturable') or group.get('differentiable')): + return False + for parameter in group['params']: + if id(parameter) in seen or (parameter.grad is not None and parameter.grad.is_sparse): + return False + seen.add(id(parameter)) + return True + + +def _materialize_gradient(parameter, state): + parameter.grad = state['sources'][parameter].float() + coefficient = state['coefficient'] + if coefficient is not None and coefficient < 1.0: + # Preserve local_multi_tensor_scale's src * scale followed by copy_. + state['scale'](None, [[parameter.grad], [parameter.grad]], coefficient) + + +def _stream_main_gradient_step(optimizer): + """Materialize deferred gradients and release each copy after its Adam update.""" + if getattr(optimizer, '_libra_streamed_main_gradient_step', False): + return + original_step = optimizer.step + + @wraps(original_step) + def step(*args, **kwargs): + # Closures, graph capture and custom hooks can observe the entire + # optimizer call. Preserve their native behavior instead of splitting it. + groups = optimizer.param_groups + state = getattr(optimizer, '_libra_bf16_gradient_state', None) + streaming = state is not None and state['active'] + has_gradient = streaming or any(p.grad is not None for g in groups for p in g['params']) + eligible = not args and not kwargs and _native_streaming_eligible(optimizer) + if not eligible or not has_gradient: + if streaming: + for parameter in state['sources']: + _materialize_gradient(parameter, state) + state['active'] = False + state['sources'].clear() + return original_step(*args, **kwargs) + + try: + for group in groups: + for parameter in group['params']: + if streaming and parameter in state['sources']: + _materialize_gradient(parameter, state) + if parameter.grad is None: + continue + # Keep the original parameter/state objects and every group + # hyperparameter. MCore's global norm determines clipping; + # deferred gradients receive that FP32 clip just above. + optimizer.param_groups = [dict(group, params=[parameter])] + original_step() + parameter.grad = None + if streaming: + state['active'] = False + state['sources'].clear() + finally: + # MCore's group-range maps and checkpoint layout refer to these + # original group objects, including when the native update fails. + optimizer.param_groups = groups + + optimizer.step = step + optimizer._libra_streamed_main_gradient_step = True + + +def _install_mcore_bf16_streaming(root, wrappers): + """Defer MCore's FP32 copies, only with its exact local norm/scale backend.""" + try: + from megatron.core import utils + from megatron.core.optimizer import clip_grads + from megatron.core.optimizer import optimizer as mcore_optimizer + from megatron.core.optimizer.distrib_optimizer import DistributedOptimizer + except ImportError: + return + if not ( + clip_grads.multi_tensor_applier is utils.local_multi_tensor_applier + and clip_grads.l2_norm_impl is utils.local_multi_tensor_l2_norm + and clip_grads.multi_tensor_scale_impl is utils.local_multi_tensor_scale + ): + return + + def eligible(wrapper): + config = getattr(wrapper, 'config', None) + return ( + type(wrapper) is DistributedOptimizer + and type(wrapper.optimizer) in (torch.optim.Adam, torch.optim.AdamW) + and _native_streaming_eligible(wrapper.optimizer) + and bool(getattr(config, 'bf16', False)) + and getattr(wrapper, 'grad_scaler', None) is None + and not getattr(config, 'log_num_zeros_in_grad', False) + and not getattr(getattr(root, 'config', None), 'log_num_zeros_in_grad', False) + and not getattr(config, 'use_precision_aware_optimizer', False) + and not getattr(config, 'use_precision_aware_optimizer_no_fp8_or_ds_fp8', False) + and not getattr(config, 'optimizer_cpu_offload', False) + and not getattr(getattr(wrapper, 'ddp_config', None), 'use_megatron_fsdp', False) + and hasattr(wrapper, 'model_float16_groups') + ) + + for wrapper in wrappers: + if eligible(wrapper) and not getattr(wrapper, '_libra_streamed_bf16_main_gradients', False): + _install_mcore_bf16_child(wrapper, eligible, clip_grads, mcore_optimizer) + # ChainedOptimizer computes a common norm before clipping any child. Keep + # that original call and its collectives, then share its exact coefficient. + for wrapper in wrappers: + if (any(getattr(child, '_libra_streamed_bf16_main_gradients', False) for child in wrappers) + and getattr(wrapper, 'chained_optimizers', None) is not None + and not getattr(wrapper, '_libra_streamed_bf16_norm', False)): + _record_gradient_norm(wrapper) + + +def _record_gradient_norm(wrapper): + original_norm = wrapper.get_grad_norm + + @wraps(original_norm) + def norm(): + value = original_norm() + pending = [wrapper] + while pending: + child = pending.pop() + children = getattr(child, 'chained_optimizers', None) + if children is not None: + pending.extend(children) + continue + state = getattr(getattr(child, 'optimizer', None), '_libra_bf16_gradient_state', None) + if state is not None and state['active']: + state['coefficient'] = ( + child.config.clip_grad / (value + 1.0e-6) + if child.config.clip_grad > 0 else None + ) + return value + + wrapper.get_grad_norm = norm + wrapper._libra_streamed_bf16_norm = True + + +def _install_mcore_bf16_child(wrapper, eligible, clip_grads, mcore_optimizer): + original_copy = wrapper._copy_model_grads_to_main_grads + original_norm_grads = wrapper.get_main_grads_for_grad_norm + original_clip = wrapper.clip_grad_norm + state = dict(active=False, sources={}, coefficient=None, + scale=lambda flag, tensors, coefficient: clip_grads.multi_tensor_applier( + clip_grads.multi_tensor_scale_impl, flag, tensors, coefficient)) + wrapper.optimizer._libra_bf16_gradient_state = state + + @torch.no_grad() + def copy_views(): + state.update(active=False, sources={}, coefficient=None) + if not eligible(wrapper) or wrapper.is_stub_optimizer: + return original_copy() + sources = {} + for model_groups, main_groups in ( + (wrapper.model_float16_groups, wrapper.shard_fp32_from_float16_groups), + (wrapper.model_fp32_groups, wrapper.shard_fp32_groups), + ): + for models, masters in zip(model_groups, main_groups): + for model, master in zip(models, masters): + interval = wrapper._get_model_param_range_map(model)['param'] + assert interval.size == master.nelement() + source = model.main_grad.view(-1)[interval.start:interval.end] + if source.dtype not in (torch.bfloat16, torch.float32): + return original_copy() + sources[master] = source + for master in sources: + master.grad = None + state.update(active=True, sources=sources) + + @torch.no_grad() + def norm_grads(): + if not state['active']: + return original_norm_grads() + # The exact local fallback first norms every FP32 tensor, then norms + # torch.tensor(the scalar norms). Norming a nonnegative scalar again + # is exact, allowing the original aggregate/collective to remain intact. + return [torch.norm(state['sources'][parameter].float()) + for parameter in wrapper.get_parameters() + if parameter in state['sources'] + and mcore_optimizer.param_is_not_shared(parameter) + and mcore_optimizer.tensor_parallel.param_is_not_tensor_parallel_duplicate(parameter)] + + @torch.no_grad() + def clip(max_norm): + if not state['active']: + return original_clip(max_norm) + value = wrapper.get_grad_norm() + state['coefficient'] = max_norm / (value + 1.0e-6) + return value + + wrapper._copy_model_grads_to_main_grads = copy_views + wrapper.get_main_grads_for_grad_norm = norm_grads + wrapper.clip_grad_norm = clip + wrapper._libra_streamed_bf16_main_gradients = True + _record_gradient_norm(wrapper) + + +def bound_default_torch_adam_memory(optimizer, *, release_main_gradients=False) -> int: + """Disable automatic foreach selection in native Adam/AdamW leaves. + + CUDA's automatic foreach Adam path creates a full parameter-sized list of + temporaries (for example, sqrt of every second moment at once). MCore's + non-fused fallback can therefore exceed memory even when model, master + parameters, gradients and Adam moments fit. The per-tensor implementation + retains the same precision and bounds those temporaries to one parameter. + + Explicit foreach choices and fused implementations are preserved. This + walks the MCore chained/single optimizer wrappers without accessing a + ChainedOptimizer's single-child-only ``optimizer`` property. Return the + number of parameter groups changed; repeated calls are idempotent. + + With release_main_gradients=True, each FP32 copy is released after its + native Adam update. For BF16 DistributedOptimizer instances using MCore's + exact local norm/scale fallback (without grad scaling, zero counting, + precision-aware optimization, FSDP or CPU offload), prepare_grads also + defers conversion: retain BF16 shard views, compute scalar FP32 norms one + at a time, and leave the original aggregate norm/collectives in place. + Each Adam update then materializes and clips only its own FP32 gradient. + BF16 accumulation, FP32 clipping/Adam arithmetic and checkpoint state stay + unchanged; existing moments no longer overlap a full FP32 gradient copy + on the second or subsequent update. Unsupported paths retain eager copy. + """ + pending = [optimizer] + visited = set() + wrappers = [] + changed = 0 + while pending: + current = pending.pop() + if current is None or id(current) in visited: + continue + visited.add(id(current)) + if type(current) in (torch.optim.Adam, torch.optim.AdamW): + for group in current.param_groups: + if group.get("foreach") is None and group.get("fused") is not True: + group["foreach"] = False + changed += 1 + if (current.defaults.get("foreach") is None + and current.defaults.get("fused") is not True): + current.defaults["foreach"] = False + if release_main_gradients: + _stream_main_gradient_step(current) + continue + children = getattr(current, "chained_optimizers", None) + wrappers.append(current) + if children is not None: + pending.extend(children) + else: + pending.append(getattr(current, "optimizer", None)) + if release_main_gradients and wrappers: + _install_mcore_bf16_streaming(optimizer, wrappers) + return changed diff --git a/tests/test_torch_optimizer_memory.py b/tests/test_torch_optimizer_memory.py new file mode 100644 index 0000000..d50fcb4 --- /dev/null +++ b/tests/test_torch_optimizer_memory.py @@ -0,0 +1,407 @@ +"""The fallback optimizer must fit without full-model foreach temporaries.""" +import copy +import importlib +import weakref +from types import SimpleNamespace + +import pytest +import torch + +from RL_Framework.engine.optimizer_memory import bound_default_torch_adam_memory + + +def _parameters(): + return [torch.nn.Parameter(torch.linspace(-1., 1., size)) for size in (19, 31)] + + +def test_native_adamw_automatic_foreach_allocates_all_moments_but_bounded_path_does_not(monkeypatch): + # CPU uses another automatic choice; force CUDA's foreach decision while + # retaining the real PyTorch optimizer step and tensor implementations. + adam_module = importlib.import_module('torch.optim.adam') + monkeypatch.setattr(adam_module, '_default_to_fused_or_foreach', + lambda *args, **kwargs: (False, True)) + original_sqrt = torch._foreach_sqrt + allocations = [] + + def measured_sqrt(tensors): + allocations.append(sum(t.numel() * t.element_size() for t in tensors)) + return original_sqrt(tensors) + + monkeypatch.setattr(torch, '_foreach_sqrt', measured_sqrt) + automatic_parameters = _parameters() + automatic = torch.optim.AdamW(automatic_parameters, lr=.01) + for parameter in automatic_parameters: + parameter.grad = torch.ones_like(parameter) + automatic.step() + assert allocations == [sum(p.numel() * p.element_size() for p in automatic_parameters)] + + masters, reference_masters = _parameters(), _parameters() + bounded = torch.optim.AdamW(masters, lr=.01, betas=(.9, .95)) + reference = torch.optim.AdamW(reference_masters, lr=.01, betas=(.9, .95), foreach=False) + assert bound_default_torch_adam_memory(bounded) == 1 + for step in range(3): + for parameter, expected in zip(masters, reference_masters): + gradient = torch.linspace(-.3, .7, parameter.numel()) * (step + 1) + parameter.grad = gradient.clone() + expected.grad = gradient.clone() + bounded.step() + reference.step() + for parameter, expected in zip(masters, reference_masters): + assert parameter.grad is not None # The default policy does not consume gradients. + assert torch.equal(parameter, expected) + # MCore's BF16 model copy and FP32 master/moments are unchanged. + assert torch.equal(parameter.to(torch.bfloat16), expected.to(torch.bfloat16)) + for key in ('exp_avg', 'exp_avg_sq', 'step'): + assert torch.equal(bounded.state[parameter][key], reference.state[expected][key]) + assert len(allocations) == 1 + + +@pytest.mark.parametrize('optimizer_type', [torch.optim.Adam, torch.optim.AdamW]) +def test_nested_mcore_chained_and_distributed_wrappers_are_configured_once(optimizer_type): + ChainedOptimizer = pytest.importorskip("megatron.core.optimizer.optimizer").ChainedOptimizer + DistributedOptimizer = pytest.importorskip("megatron.core.optimizer.distrib_optimizer").DistributedOptimizer + + first, second = optimizer_type(_parameters()), optimizer_type(_parameters()) + distributed = DistributedOptimizer.__new__(DistributedOptimizer) + distributed.optimizer = first + inner = ChainedOptimizer.__new__(ChainedOptimizer) + inner.chained_optimizers = [distributed, SimpleNamespace(optimizer=second)] + outer = ChainedOptimizer.__new__(ChainedOptimizer) + # This also proves we never access the asserting .optimizer property of a + # multi-child chain, and shared references are not counted twice. + outer.chained_optimizers = [inner, distributed] + assert bound_default_torch_adam_memory(outer) == 2 + assert all(group['foreach'] is False for opt in (first, second) for group in opt.param_groups) + assert bound_default_torch_adam_memory(outer) == 0 + + +def test_explicit_foreach_and_fused_choices_and_non_native_optimizers_are_preserved(): + explicit = torch.optim.AdamW([ + {'params': [_parameters()[0]], 'foreach': True}, + {'params': [_parameters()[0]], 'foreach': False}, + {'params': [_parameters()[0]], 'foreach': None}, + ]) + fused = torch.optim.AdamW(_parameters(), fused=True) + + class OtherAdam(torch.optim.AdamW): + pass + + other = OtherAdam(_parameters()) + sgd = torch.optim.SGD(_parameters(), lr=.1) + tree = SimpleNamespace(chained_optimizers=[explicit, fused, other, sgd, None]) + assert bound_default_torch_adam_memory(tree) == 1 + assert [group['foreach'] for group in explicit.param_groups] == [True, False, False] + assert fused.param_groups[0]['foreach'] is None + assert fused.param_groups[0]['fused'] is True + assert other.param_groups[0]['foreach'] is None + assert sgd.param_groups[0]['foreach'] is None + + +def test_old_optimizer_state_restoration_reapplies_only_automatic_policy(): + optimizer = torch.optim.AdamW(_parameters()) + for parameter in optimizer.param_groups[0]['params']: + parameter.grad = torch.ones_like(parameter) + optimizer.step() + saved = copy.deepcopy(optimizer.state_dict()) + assert saved['param_groups'][0]['foreach'] is None + assert bound_default_torch_adam_memory(optimizer) == 1 + optimizer.load_state_dict(saved) + assert optimizer.param_groups[0]['foreach'] is None + assert bound_default_torch_adam_memory(optimizer) == 1 + actual = optimizer.state_dict() + for key, values in saved['state'].items(): + for field, tensor in values.items(): + assert torch.equal(actual['state'][key][field], tensor) + assert optimizer.defaults['foreach'] is False + optimizer.add_param_group({'params': [_parameters()[0]]}) + assert optimizer.param_groups[-1]['foreach'] is False + + +def _master_groups(parameters): + return [dict(params=parameters[:2], lr=.003, weight_decay=.01, name='body'), + dict(params=parameters[2:], lr=.007, weight_decay=0., amsgrad=True, name='head')] + + +def _assert_adam_equal(actual, expected, parameters, reference): + for parameter, target in zip(parameters, reference): + assert torch.equal(parameter, target) + assert set(actual.state[parameter]) == set(expected.state[target]) + for field, value in actual.state[parameter].items(): + assert torch.equal(value, expected.state[target][field]) + + +@pytest.mark.parametrize('optimizer_type', [torch.optim.Adam, torch.optim.AdamW]) +def test_streamed_main_gradients_preserve_multigroup_updates_and_reloadable_state(optimizer_type): + parameters = _parameters() + [torch.nn.Parameter(torch.linspace(-.4, .8, 7))] + reference = [torch.nn.Parameter(p.detach().clone()) for p in parameters] + actual = optimizer_type(_master_groups(parameters), betas=(.9, .95)) + expected = optimizer_type(_master_groups(reference), betas=(.9, .95), foreach=False) + original_groups = actual.param_groups + group_objects = list(original_groups) + bound_default_torch_adam_memory(actual, release_main_gradients=True) + streamed_step = actual.step + bound_default_torch_adam_memory(actual, release_main_gradients=True) + assert actual.step is streamed_step + for step in range(4): + for index, (parameter, target) in enumerate(zip(parameters, reference)): + gradient = torch.linspace(-.3, .7, parameter.numel()) * (step + 1) + parameter.grad = gradient.clone() if (step, index) != (2, 1) else None + target.grad = gradient.clone() if (step, index) != (2, 1) else None + # MCore computes and applies the common global clipping before Adam. + torch.nn.utils.clip_grad_norm_(parameters, 1.) + torch.nn.utils.clip_grad_norm_(reference, 1.) + actual.step() + expected.step() + assert all(parameter.grad is None for parameter in parameters) + assert actual.param_groups is original_groups + assert all(group is original for group, original in zip(actual.param_groups, group_objects)) + _assert_adam_equal(actual, expected, parameters, reference) + + saved = copy.deepcopy(actual.state_dict()) + assert [len(group['params']) for group in saved['param_groups']] == [2, 1] + assert [group['name'] for group in saved['param_groups']] == ['body', 'head'] + reloaded_parameters = [torch.nn.Parameter(p.detach().clone()) for p in parameters] + reloaded = optimizer_type(_master_groups(reloaded_parameters), betas=(.9, .95), foreach=False) + reloaded.load_state_dict(saved) + bound_default_torch_adam_memory(reloaded, release_main_gradients=True) + for parameter, target in zip(reloaded_parameters, reference): + parameter.grad = torch.ones_like(parameter) + target.grad = torch.ones_like(target) + reloaded.step() + expected.step() + _assert_adam_equal(reloaded, expected, reloaded_parameters, reference) + + +def test_streamed_state_initialization_releases_each_main_gradient_before_the_next(monkeypatch): + parameters = _parameters() + [torch.nn.Parameter(torch.zeros(7))] + optimizer = torch.optim.AdamW(parameters, foreach=False) + for parameter in parameters: + parameter.grad = torch.ones_like(parameter) + bound_default_torch_adam_memory(optimizer, release_main_gradients=True) + allocate = torch.zeros_like + resident_gradients = [] + + def measured_zeros(tensor, *args, **kwargs): + if any(tensor is parameter for parameter in parameters): + resident_gradients.append(sum(p.grad.numel() for p in parameters if p.grad is not None)) + return allocate(tensor, *args, **kwargs) + + monkeypatch.setattr(torch, 'zeros_like', measured_zeros) + optimizer.step() + # Each parameter initializes two Adam moments. Later moments must not + # overlap the FP32 main gradients of parameters already updated. + assert resident_gradients == [57, 57, 38, 38, 7, 7] + assert all(parameter.grad is None for parameter in parameters) + + +def test_streamed_step_restores_original_groups_and_propagates_allocation_failure(monkeypatch): + parameters = _parameters() + [torch.nn.Parameter(torch.zeros(7))] + optimizer = torch.optim.AdamW(_master_groups(parameters), foreach=False) + original_groups, group_objects = optimizer.param_groups, list(optimizer.param_groups) + gradients = [torch.ones_like(parameter) for parameter in parameters] + for parameter, gradient in zip(parameters, gradients): + parameter.grad = gradient + bound_default_torch_adam_memory(optimizer, release_main_gradients=True) + allocate = torch.zeros_like + + def failing_zeros(tensor, *args, **kwargs): + if tensor is parameters[1]: + raise MemoryError('injected moment allocation') + return allocate(tensor, *args, **kwargs) + + monkeypatch.setattr(torch, 'zeros_like', failing_zeros) + with pytest.raises(MemoryError, match='injected moment allocation'): + optimizer.step() + assert optimizer.param_groups is original_groups + assert all(group is original for group, original in zip(optimizer.param_groups, group_objects)) + assert parameters[0].grad is None + assert parameters[1].grad is gradients[1] and parameters[2].grad is gradients[2] + + +def test_closure_uses_one_native_step_and_keeps_main_gradients(): + parameters = _parameters() + optimizer = torch.optim.AdamW(parameters, foreach=False) + gradients = [torch.ones_like(parameter) for parameter in parameters] + for parameter, gradient in zip(parameters, gradients): + parameter.grad = gradient + bound_default_torch_adam_memory(optimizer, release_main_gradients=True) + calls = [] + + def closure(): + calls.append(True) + return torch.tensor(4.) + + assert optimizer.step(closure).item() == 4. + assert calls == [True] + assert all(p.grad is gradient for p, gradient in zip(parameters, gradients)) + + +@pytest.mark.parametrize('unsupported', ['foreach', 'fused', 'differentiable', 'capturable']) +def test_unsupported_streaming_configuration_delegates_without_clearing_gradients(unsupported): + parameters = _parameters() + optimizer = torch.optim.AdamW([ + dict(params=parameters[:1], foreach=False), + dict(params=parameters[1:], **{unsupported: True}), + ]) + gradients = [torch.ones_like(parameter) for parameter in parameters] + for parameter, gradient in zip(parameters, gradients): + parameter.grad = gradient + original_groups = optimizer.param_groups + calls, sentinel = [], object() + + def native_step(*args, **kwargs): + # Spy at the native boundary so unsupported CUDA-only modes can be + # checked on CPU without silently exercising another implementation. + calls.append(([len(group['params']) for group in optimizer.param_groups], args, kwargs)) + return sentinel + + optimizer.step = native_step + bound_default_torch_adam_memory(optimizer, release_main_gradients=True) + assert optimizer.step() is sentinel + assert calls == [([1, 1], (), {})] + assert optimizer.param_groups is original_groups + assert all(p.grad is gradient for p, gradient in zip(parameters, gradients)) + + +@pytest.fixture +def mcore_cpu(monkeypatch): + """Run MCore's real optimizer flow/math, redirecting CUDA-only test plumbing.""" + module = pytest.importorskip('megatron.core.optimizer.optimizer') + distributed = pytest.importorskip('megatron.core.optimizer.distrib_optimizer') + layers = pytest.importorskip('megatron.core.tensor_parallel.layers') + tensor, tensor_type = torch.tensor, torch.Tensor.type + reductions = [] + + def cpu_tensor(*args, **kwargs): + if kwargs.get('device') == 'cuda': + kwargs['device'] = 'cpu' + return tensor(*args, **kwargs) + + def cuda_type_for_clip(value, *args, **kwargs): + if not args and not kwargs and value.dtype == torch.float32: + return 'torch.cuda.FloatTensor' + return tensor_type(value, *args, **kwargs) + + monkeypatch.setattr(torch, 'tensor', cpu_tensor) + monkeypatch.setattr(torch.Tensor, 'type', cuda_type_for_clip) + monkeypatch.setattr(layers, 'get_tensor_model_parallel_rank', lambda: 1) + monkeypatch.setattr(torch.distributed, 'all_reduce', + lambda value, **kwargs: reductions.append((value.clone(), kwargs['group']))) + + def make(children=1): + wrappers = [] + group = object() + for child_index in range(children): + masters = _parameters() + [torch.nn.Parameter(torch.linspace(-.4, .8, 7))] + models = [] + for index, master in enumerate(masters): + master.tensor_model_parallel = index != 2 + master.shared = index == 1 + model = torch.nn.Parameter(torch.zeros(master.numel() + 4, dtype=torch.bfloat16)) + model.main_grad = torch.zeros_like(model) + models.append(model) + wrapper = distributed.DistributedOptimizer.__new__(distributed.DistributedOptimizer) + wrapper.optimizer = torch.optim.AdamW(_master_groups(masters), betas=(.9, .95), foreach=False) + wrapper.config = SimpleNamespace(bf16=True, use_precision_aware_optimizer=False, + use_precision_aware_optimizer_no_fp8_or_ds_fp8=False, optimizer_cpu_offload=False, + log_num_zeros_in_grad=False, clip_grad=.7 / (child_index + 1), timers=None, + reuse_grad_buf_for_mxfp8_param_ag=False, overlap_param_gather_with_optimizer_step=False) + wrapper.grad_scaler, wrapper.is_stub_optimizer = None, False + wrapper.ddp_config = SimpleNamespace(use_megatron_fsdp=False, overlap_param_gather=False) + wrapper.grad_stats_parallel_group = group + wrapper.model_float16_groups = [models[:2], models[2:]] + wrapper.shard_fp32_from_float16_groups = [masters[:2], masters[2:]] + wrapper.model_fp32_groups = wrapper.shard_fp32_groups = [[], []] + wrapper._get_model_param_range_map = lambda model: { + 'param': SimpleNamespace(start=2, end=model.numel()-2, size=model.numel()-4)} + + @torch.no_grad() + def copy_models(models=models, masters=masters): + for model, master in zip(models, masters): + model[2:-2].copy_(master) + + wrapper._copy_main_params_to_model_params = copy_models + wrapper.model_chunks = [SimpleNamespace(start_param_sync=lambda: None)] + wrappers.append(wrapper) + chain = module.ChainedOptimizer.__new__(module.ChainedOptimizer) + chain.chained_optimizers = wrappers + chain.config = wrappers[0].config + return SimpleNamespace(root=chain, children=wrappers, + parameters=[p for child in wrappers for group in child.optimizer.param_groups for p in group['params']], + models=[p for child in wrappers for group in child.model_float16_groups for p in group]) + + return SimpleNamespace(make=make, reductions=reductions) + + +@pytest.mark.parametrize('children', [1, 2]) +def test_bf16_prepare_streams_second_step_with_exact_global_norm_clip_and_state(mcore_cpu, monkeypatch, children): + actual, expected = mcore_cpu.make(children), mcore_cpu.make(children) + bound_default_torch_adam_memory(actual.root, release_main_gradients=True) + assert all(child._libra_streamed_bf16_main_gradients for child in actual.children) + copy_method = actual.children[0]._copy_model_grads_to_main_grads + bound_default_torch_adam_memory(actual.root, release_main_gradients=True) + assert actual.children[0]._copy_model_grads_to_main_grads is copy_method + for step in range(3): + # The second pass starts with all Adam moments already resident. + for index, (model, reference) in enumerate(zip(actual.models, expected.models)): + gradient = torch.linspace(-3., 7., model.numel()).to(torch.bfloat16) * (step + index + 1) + model.main_grad.copy_(gradient) + reference.main_grad.copy_(gradient) + source_copies = [model.main_grad.clone() for model in actual.models] + if step == 2: + for child in actual.children: + child.optimizer.load_state_dict(copy.deepcopy(child.optimizer.state_dict())) + bound_default_torch_adam_memory(actual.root, release_main_gradients=True) + expected_result = expected.root.step() + references, live_counts = [], [] + original_float = torch.Tensor.float + + def counted_float(value, *args, **kwargs): + result = original_float(value, *args, **kwargs) + if value.dtype == torch.bfloat16: + references.append(weakref.ref(result)) + live_counts.append(sum(ref() is not None for ref in references)) + return result + + with monkeypatch.context() as patch: + patch.setattr(torch.Tensor, 'float', counted_float) + actual_result = actual.root.step() + assert actual_result == expected_result + assert actual_result[1] > 1 # Clipping is exercised, across all children. + assert max(live_counts) == 1 + assert not any(ref() is not None for ref in references) + assert all(parameter.grad is None for parameter in actual.parameters) + for model, source in zip(actual.models, source_copies): + assert torch.equal(model.main_grad, source) # Never clip the BF16 source. + for child, reference in zip(actual.children, expected.children): + parameters = [p for group in child.optimizer.param_groups for p in group['params']] + targets = [p for group in reference.optimizer.param_groups for p in group['params']] + _assert_adam_equal(child.optimizer, reference.optimizer, parameters, targets) + assert all(torch.equal(a, b) for a, b in zip(actual.models, expected.models)) + # Original MCore performs one global SUM for each chain, using its own group. + reductions = mcore_cpu.reductions[-2:] + assert len(reductions) == 2 + assert torch.equal(reductions[0][0], reductions[1][0]) + assert reductions[0][1] is expected.root.get_grad_stats_parallel_group() + assert reductions[1][1] is actual.root.get_grad_stats_parallel_group() + + +@pytest.mark.parametrize('unsupported', ['grad_scaler', 'zeros', 'precision_aware', 'fsdp', 'backend']) +def test_bf16_streaming_leaves_unsupported_mcore_preparation_untouched(mcore_cpu, monkeypatch, unsupported): + actual = mcore_cpu.make() + child = actual.children[0] + original = child._copy_model_grads_to_main_grads + if unsupported == 'grad_scaler': + child.grad_scaler = object() + elif unsupported == 'zeros': + child.config.log_num_zeros_in_grad = True + elif unsupported == 'precision_aware': + child.config.use_precision_aware_optimizer = True + elif unsupported == 'fsdp': + child.ddp_config.use_megatron_fsdp = True + else: + from megatron.core.optimizer import clip_grads + monkeypatch.setattr(clip_grads, 'l2_norm_impl', lambda *args: None) + bound_default_torch_adam_memory(actual.root, release_main_gradients=True) + assert not getattr(child, '_libra_streamed_bf16_main_gradients', False) + assert child._copy_model_grads_to_main_grads == original From 5ec582f8274065fd160d05105474701e8e57ae1d Mon Sep 17 00:00:00 2001 From: cytsoi <154662279+tcyjohn@users.noreply.github.com> Date: Thu, 22 Oct 2026 01:29:59 +0800 Subject: [PATCH 5/9] Restore native master parameters and independent Adam state for Hybrid lanes --- engine/megatron_core_checkpointing.py | 4 +- engine/megatron_core_train_engine.py | 104 +++++++++++++++++++- engine/optimizer_state.py | 57 +++++++++++ tests/test_ehp_dp_resharding.py | 117 +++++++++++++++++++++++ tests/test_elastic_optimizer_state.py | 132 ++++++++++++++++++++++++++ tests/test_optimizer_state.py | 52 ++++++++++ 6 files changed, 463 insertions(+), 3 deletions(-) create mode 100644 engine/optimizer_state.py create mode 100644 tests/test_ehp_dp_resharding.py create mode 100644 tests/test_elastic_optimizer_state.py create mode 100644 tests/test_optimizer_state.py diff --git a/engine/megatron_core_checkpointing.py b/engine/megatron_core_checkpointing.py index 5982299..ccb15e8 100644 --- a/engine/megatron_core_checkpointing.py +++ b/engine/megatron_core_checkpointing.py @@ -168,7 +168,9 @@ def load( ) self._load_model_state(self._unwrap_model(model), loaded) if optimizer is not None and "optimizer" in loaded: - optimizer.load_state_dict(loaded["optimizer"]) + from RL_Framework.engine.optimizer_state import optimizer_state_load_context + with optimizer_state_load_context(optimizer): + optimizer.load_state_dict(loaded["optimizer"]) self._barrier() def _build_sharded_state_dict( diff --git a/engine/megatron_core_train_engine.py b/engine/megatron_core_train_engine.py index 971593b..32a36c6 100644 --- a/engine/megatron_core_train_engine.py +++ b/engine/megatron_core_train_engine.py @@ -3,6 +3,9 @@ from __future__ import annotations import os +import copy +import json +import uuid import threading import time from pathlib import Path @@ -259,7 +262,7 @@ def init_phase(name: str) -> None: self.optimizer = get_megatron_optimizer( optimizer_config, self.model, - use_gloo_process_groups=False, + use_gloo_process_groups=True, ) bound_default_torch_adam_memory(self.optimizer, release_main_gradients=(getattr(self, 'use_distributed_optimizer', False) @@ -798,7 +801,104 @@ def apply_elastic_gradient_update(self, tensors: tuple[torch.Tensor, ...], *, st ok, _norm, _zeros = self.optimizer.step() if not ok: raise FloatingPointError("Hybrid Megatron optimizer rejected the update") - self.current_version = int(state_version) + self.optimizer_revision = int(state_version) + + @staticmethod + def _clone_cpu_state(value): + if isinstance(value, torch.Tensor): + return value.detach().to(device="cpu", copy=True) + if isinstance(value, dict): + return {key: MegatronCoreTrainEngine._clone_cpu_state(item) for key, item in value.items()} + if isinstance(value, (tuple, list)): + items = [MegatronCoreTrainEngine._clone_cpu_state(item) for item in value] + return tuple(items) if isinstance(value, tuple) else items + return copy.deepcopy(value) + + def capture_elastic_state(self): + """Clone current parameters, optimizer master parameters and Adam state.""" + self._synchronize_optimizer_parameters() + optimizer = self.optimizer.state_dict() + # MCore DistributedOptimizer deliberately omits DP-zero parameter and + # Adam shards from state_dict; include its actual DP-zero state too. + distributed_state = self._capture_optimizer_parameter_state(self.optimizer) + return { + "format_version": 1, + "snapshot_id": uuid.uuid4().hex, + "model": self._clone_cpu_state([chunk.state_dict() for chunk in self.model]), + "optimizer": self._clone_cpu_state(optimizer), + "distributed_optimizer": self._clone_cpu_state(distributed_state), + "has_distributed_optimizer": distributed_state is not None, + "optimizer_revision": int(self.optimizer_revision), + "version": int(self.current_version), + "parallel_state": self.get_parallel_state(), + } + + @classmethod + def _capture_optimizer_parameter_state(cls, optimizer): + chained = getattr(optimizer, "chained_optimizers", None) + if chained is not None: + return {"kind": "chain", "children": [cls._capture_optimizer_parameter_state(child) for child in chained]} + if hasattr(optimizer, "get_parameter_state_dp_zero"): + if getattr(optimizer, "data_parallel_group_gloo", None) is None: + raise RuntimeError("live EHP optimizer state requires its initialized DP Gloo group") + return {"kind": "dp_zero", "state": optimizer.get_parameter_state_dp_zero()} + return None + + @classmethod + def _restore_optimizer_parameter_state(cls, optimizer, payload): + if payload is None: + return + if payload["kind"] == "chain": + children = getattr(optimizer, "chained_optimizers", ()) + if len(children) != len(payload["children"]): + raise ValueError("live EHP optimizer chain count mismatch") + for child, state in zip(children, payload["children"]): + cls._restore_optimizer_parameter_state(child, state) + elif payload["kind"] == "dp_zero": + if not hasattr(optimizer, "load_parameter_state_from_dp_zero"): + raise RuntimeError("optimizer cannot restore DP-zero Adam/master state") + group = getattr(optimizer, "data_parallel_group_gloo", None) + if payload["state"] is None and (group is None or group.rank() == 0): + raise ValueError("live EHP snapshot lacks the master parameters and Adam state on its receiving DP root") + optimizer.load_parameter_state_from_dp_zero(payload["state"]) + else: + raise ValueError("unknown live EHP optimizer parameter state format") + + def _synchronize_optimizer_parameters(self): + if self.optimizer is not None and hasattr(self.optimizer, "finish_param_sync"): + for index in range(len(self.model)): + self.optimizer.finish_param_sync(index) + if torch.cuda.is_available(): + torch.cuda.synchronize(self._device()) + + def restore_elastic_state(self, payload): + if payload.get("format_version") != 1 or "optimizer" not in payload: + raise ValueError("live EHP snapshot is missing optimizer state") + if len(payload["model"]) != len(self.model): + raise ValueError("live EHP model chunk count mismatch") + for chunk, state in zip(self.model, payload["model"]): + chunk.load_state_dict(state, strict=True) + from RL_Framework.engine.optimizer_state import optimizer_state_load_context + # DP-zero restoration reaches the same MCore 0.14 low-level tensor + # setter as model-space restoration. Keep the real step supplied by + # common metadata visible while it installs master/Adam tensors. + with optimizer_state_load_context(self.optimizer): + self.optimizer.load_state_dict(payload["optimizer"]) + parameter_state = payload.get("distributed_optimizer") + if payload.get("has_distributed_optimizer", parameter_state is not None): + self._restore_optimizer_parameter_state(self.optimizer, parameter_state) + bound_default_torch_adam_memory(self.optimizer, + release_main_gradients=(getattr(self, 'use_distributed_optimizer', False) + and not getattr(self, 'grad_reduce_in_fp32', False) + and not getattr(self, 'optimizer_cpu_offload', False))) + self.optimizer_revision = int(payload["optimizer_revision"]) + self.current_version = int(payload["version"]) + self._synchronize_optimizer_parameters() + receipt = {"acknowledged": True, "snapshot_id": payload["snapshot_id"], + "state_version": self.optimizer_revision, + "optimizer_revision": self.optimizer_revision} + self._elastic_state_receipt = receipt + return receipt def capture_elastic_state_snapshot( self, diff --git a/engine/optimizer_state.py b/engine/optimizer_state.py new file mode 100644 index 0000000..86d602c --- /dev/null +++ b/engine/optimizer_state.py @@ -0,0 +1,57 @@ +"""Strict native MCore optimizer metadata restoration for fixed-topology EHP.""" + +from contextlib import contextmanager + +import torch + + +@contextmanager +def optimizer_state_load_context(optimizer): + """Preserve native metadata's Torch Adam step during MCore tensor loads. + + MCore 0.14's model-space loader deliberately excludes ``step`` from tensor + state, but its low-level setter iterates every Torch Adam state key. Supply + the step already loaded from common metadata to that setter only. Other + missing master/moment tensors still fail through the native strict copy. + This temporary instance adapter never edits installed MCore dependencies. + """ + parts = getattr(optimizer, "chained_optimizers", (optimizer,)) + restores = [] + absent = object() + completed = False + try: + for part in parts: + inner = getattr(part, "optimizer", None) + native = getattr(part, "_set_main_param_and_optimizer_states", None) + if not isinstance(inner, (torch.optim.Adam, torch.optim.AdamW)) or not callable(native): + continue + name = "_set_main_param_and_optimizer_states" + previous = vars(part).get(name, absent) + + def copy_with_metadata_step(model_param, tensors, *, part=part, inner=inner, native=native): + group, order = part.model_param_group_index_map[model_param] + main_param = inner.param_groups[group]["params"][order] + step = inner.state[main_param].get("step") + if "step" not in tensors and step is not None: + tensors = dict(tensors, step=step) + return native(model_param, tensors) + + setattr(part, name, copy_with_metadata_step) + restores.append((part, name, previous)) + yield + completed = True + finally: + for part, name, previous in reversed(restores): + if previous is absent: + delattr(part, name) + else: + setattr(part, name, previous) + # Native MCore's Torch fallback assigns one shared scalar tensor + # to every Adam parameter during metadata loading. Torch Adam + # increments each parameter's step in-place; shared references + # would advance it once per parameter and change the next update. + if completed: + for state in part.optimizer.state.values(): + step = state.get("step") + if isinstance(step, torch.Tensor): + state["step"] = step.detach().clone() diff --git a/tests/test_ehp_dp_resharding.py b/tests/test_ehp_dp_resharding.py new file mode 100644 index 0000000..c728d93 --- /dev/null +++ b/tests/test_ehp_dp_resharding.py @@ -0,0 +1,117 @@ +"""Native MCore Adam resharding from two Core DP replicas to one Hybrid. + +Each test isolates one model-parallel lane: real Gloo gathers the Core shards, +then a separately initialized DP1 group restores the native DP-zero payload. +The small optimizer fixture follows test_ehp_correctness.py and needs no CUDA. +""" +from pathlib import Path +from types import SimpleNamespace +from datetime import timedelta +import torch +import torch.distributed as dist +import torch.multiprocessing as mp +import pytest + +from RL_Framework.engine.megatron_core_train_engine import MegatronCoreTrainEngine + + +def _engine(module, rank, world, values, padded_extra=0): + full = torch.tensor(values, dtype=torch.float32) + padded_numel = len(values) + padded_extra + width = padded_numel // world + start, end = rank * width, min((rank + 1) * width, len(values)) + master = torch.nn.Parameter(full[start:end].clone(), requires_grad=False) + model = torch.nn.Parameter(full.bfloat16().clone()) + inner = torch.optim.AdamW([{'params': [master], 'wd_mult': 1., 'lr_mult': 1., + 'is_expert_parallel': False, 'is_decoupled_lr': False}], + lr=.01, betas=(.9, .95)) + optimizer = module.DistributedOptimizer.__new__(module.DistributedOptimizer) + optimizer.optimizer = inner + optimizer.ddp_config = SimpleNamespace(use_megatron_fsdp=False) + optimizer.config = SimpleNamespace(exp_avg_dtype=torch.float32, exp_avg_sq_dtype=torch.float32, + use_precision_aware_optimizer_no_fp8_or_ds_fp8=False, fp16=False) + optimizer.grad_scaler = None + optimizer.data_parallel_group_gloo = dist.group.WORLD + optimizer.model_param_group_index_map = {model: (0, 0)} + dtype = (torch.bfloat16, torch.bfloat16) + optimizer.gbuf_ranges = [{dtype: [{'param_map': {model: { + 'gbuf_world': module.Range(start, end), 'param': module.Range(start, end), + 'gbuf_local': module.Range(0, end - start)}}}]}] + optimizer.buffers = [SimpleNamespace(numel_unpadded=len(values), params=[model], + buckets=[SimpleNamespace(grad_data=torch.zeros(padded_numel), numel_unpadded=len(values))])] + chunk = torch.nn.Module() + chunk.register_parameter('weight', model) + engine = MegatronCoreTrainEngine.__new__(MegatronCoreTrainEngine) + engine.model, engine.optimizer = [chunk], optimizer + engine.optimizer_revision, engine.current_version = 3, 0 + engine.get_parallel_state = lambda: dict(train_tp=1, train_dp=world) + return engine, master, model + + +def _capture_core_shards(rank, rendezvous, directory, padded_extra): + from megatron.core.optimizer import distrib_optimizer as module + module.HAVE_APEX_OR_TE = False + torch.cuda.current_device = lambda: 'cpu' + dist.init_process_group('gloo', init_method=f'file://{rendezvous}', rank=rank, world_size=2, + timeout=timedelta(seconds=30)) + try: + engine, master, model = _engine(module, rank, 2, [1., 2., 3., 4.], padded_extra) + gradient = torch.tensor([.1, -.2, .3, -.4]) + width = (4 + padded_extra) // 2 + for _ in range(3): + master.grad = gradient[rank * width:rank * width + master.numel()].clone() + engine.optimizer.optimizer.step() + local = torch.zeros(width) + local[:master.numel()].copy_(master.detach()) + shards = [torch.empty_like(local) for _ in range(2)] + dist.all_gather(shards, local) + with torch.no_grad(): + model.copy_(torch.cat(shards)[:4]) + snapshot = engine.capture_elastic_state() + if rank == 0: + torch.save(snapshot, Path(directory) / 'core-dp0.pt') + else: + # MCore returns empty per-bucket metadata on followers, not full + # Adam/master tensors. Production canonicalizes all targets to dp0. + state = snapshot['distributed_optimizer']['state'] + assert not state[0][(torch.bfloat16, torch.bfloat16)] + finally: + dist.destroy_process_group() + + +@pytest.mark.parametrize('source_padding', [0, 2]) +def test_native_mcore_dp2_to_dp1_restores_full_master_moments_step_and_next_update(tmp_path, monkeypatch, source_padding): + module = pytest.importorskip( + 'megatron.core.optimizer.distrib_optimizer', + reason='native DP-zero resharding requires the pinned Megatron-Core source', + ) + monkeypatch.setattr(module, 'HAVE_APEX_OR_TE', False) + monkeypatch.setattr(torch.cuda, 'current_device', lambda: 'cpu') + mp.spawn(_capture_core_shards, args=(str(tmp_path / 'source-rendezvous'), str(tmp_path), source_padding), nprocs=2, join=True) + snapshot = torch.load(tmp_path / 'core-dp0.pt', weights_only=False) + dist.init_process_group('gloo', init_method=f'file://{tmp_path / "target-rendezvous"}', + rank=0, world_size=1, timeout=timedelta(seconds=30)) + try: + target, master, model = _engine(module, 0, 1, [-9.] * 4) + receipt = target.restore_elastic_state(snapshot) + assert receipt['acknowledged'] and receipt['optimizer_revision'] == 3 + reference = torch.nn.Parameter(torch.tensor([1., 2., 3., 4.])) + adam = torch.optim.AdamW([reference], lr=.01, betas=(.9, .95)) + for _ in range(3): + reference.grad = torch.tensor([.1, -.2, .3, -.4]) + adam.step() + torch.testing.assert_close(master, reference, rtol=0, atol=0) + torch.testing.assert_close(model, reference.bfloat16(), rtol=0, atol=0) + for key in ('step', 'exp_avg', 'exp_avg_sq'): + torch.testing.assert_close(target.optimizer.optimizer.state[master][key], + adam.state[reference][key], rtol=0, atol=0) + master.grad = torch.tensor([-.7, .4, .2, -.1]) + reference.grad = master.grad.clone() + target.optimizer.optimizer.step() + adam.step() + torch.testing.assert_close(master, reference, rtol=0, atol=0) + for key in ('step', 'exp_avg', 'exp_avg_sq'): + torch.testing.assert_close(target.optimizer.optimizer.state[master][key], + adam.state[reference][key], rtol=0, atol=0) + finally: + dist.destroy_process_group() diff --git a/tests/test_elastic_optimizer_state.py b/tests/test_elastic_optimizer_state.py new file mode 100644 index 0000000..8c7a689 --- /dev/null +++ b/tests/test_elastic_optimizer_state.py @@ -0,0 +1,132 @@ +"""Native master/moment/step restoration at an elastic boundary.""" +import pytest +import torch +import torch.distributed as dist +from RL_Framework.engine.megatron_core_train_engine import MegatronCoreTrainEngine + + +def test_chained_distributed_optimizer_snapshot_includes_actual_master_and_adam_tensors(): + from types import SimpleNamespace + class DistributedChild: + data_parallel_group_gloo = object() + def __init__(self, value): + self.state = {'param': torch.tensor([value]), 'exp_avg': torch.tensor([value + 1]), + 'exp_avg_sq': torch.tensor([value + 2]), 'step': torch.tensor(9)} + def get_parameter_state_dp_zero(self): + return self.state + def load_parameter_state_from_dp_zero(self, state): + self.state = state + source = SimpleNamespace(chained_optimizers=[DistributedChild(1.), DistributedChild(3.)]) + target = SimpleNamespace(chained_optimizers=[DistributedChild(0.), DistributedChild(0.)]) + payload = MegatronCoreTrainEngine._clone_cpu_state(MegatronCoreTrainEngine._capture_optimizer_parameter_state(source)) + source.chained_optimizers[0].state['param'].add_(100) + MegatronCoreTrainEngine._restore_optimizer_parameter_state(target, payload) + assert target.chained_optimizers[0].state['param'].item() == 1. + assert target.chained_optimizers[1].state['exp_avg_sq'].item() == 5. + assert target.chained_optimizers[1].state['step'].item() == 9 + + + +def test_receiving_dp_root_rejects_nonroot_source_missing_adam_and_accepts_canonical_dp0(): + from types import SimpleNamespace + loaded = [] + target = SimpleNamespace(data_parallel_group_gloo=SimpleNamespace(rank=lambda: 0), + load_parameter_state_from_dp_zero=loaded.append) + with pytest.raises(ValueError, match='master parameters and Adam'): + MegatronCoreTrainEngine._restore_optimizer_parameter_state(target, {'kind': 'dp_zero', 'state': None}) + assert not loaded + canonical = {'param': torch.tensor([2.]), 'exp_avg': torch.tensor([.2]), 'exp_avg_sq': torch.tensor([.02])} + MegatronCoreTrainEngine._restore_optimizer_parameter_state(target, {'kind': 'dp_zero', 'state': canonical}) + assert loaded == [canonical] + + + +def test_actual_mcore_dp_zero_restore_preserves_step_master_moments_and_removes_adapter(tmp_path, monkeypatch): + from types import SimpleNamespace + module = pytest.importorskip('megatron.core.optimizer.distrib_optimizer') + monkeypatch.setattr(module, 'HAVE_APEX_OR_TE', False) + monkeypatch.setattr(torch.cuda, 'current_device', lambda: 'cpu') + # The production initialization may already have installed its global + # ordinary-checkpoint compatibility wrapper. Isolate the real native setter + # for the negative reproduction, then restore that wrapper before the EHP + # installation to verify that the scoped context can coexist with it. + existing_setter = module.DistributedOptimizer._set_main_param_and_optimizer_states + native_setter = existing_setter + if getattr(native_setter, '_rl_framework_step_compat', False): + closure = dict(zip(native_setter.__code__.co_freevars, + (cell.cell_contents for cell in native_setter.__closure__))) + native_setter = closure['original_set_state'] + monkeypatch.setattr(module.DistributedOptimizer, '_set_main_param_and_optimizer_states', native_setter) + assert not dist.is_initialized() + dist.init_process_group('gloo', init_method=f'file://{tmp_path / "dp-zero-rendezvous"}', rank=0, world_size=1) + try: + def engine(values): + master = torch.nn.Parameter(torch.tensor(values, dtype=torch.float32), requires_grad=False) + model = torch.nn.Parameter(master.detach().bfloat16().clone()) + inner = torch.optim.AdamW([{ + 'params': [master], 'wd_mult': 1., 'lr_mult': 1., + 'is_expert_parallel': False, 'is_decoupled_lr': False, + }], lr=.01, betas=(.9, .95)) + optimizer = module.DistributedOptimizer.__new__(module.DistributedOptimizer) + optimizer.optimizer = inner + optimizer.ddp_config = SimpleNamespace(use_megatron_fsdp=False) + optimizer.config = SimpleNamespace( + exp_avg_dtype=torch.float32, exp_avg_sq_dtype=torch.float32, + use_precision_aware_optimizer_no_fp8_or_ds_fp8=False, fp16=False, + ) + optimizer.grad_scaler = None + optimizer.data_parallel_group_gloo = dist.group.WORLD + optimizer.model_param_group_index_map = {model: (0, 0)} + dtype = (torch.bfloat16, torch.float32) + optimizer.gbuf_ranges = [{dtype: [{'param_map': { + model: {'gbuf_world': module.Range(0, 2), 'param': module.Range(0, 2), + 'gbuf_local': module.Range(0, 2)}, + }}]}] + optimizer.buffers = [SimpleNamespace( + numel_unpadded=2, params=[model], + buckets=[SimpleNamespace(grad_data=torch.zeros(2), numel_unpadded=2)], + )] + chunk = torch.nn.Module() + chunk.register_parameter('weight', model) + result = MegatronCoreTrainEngine.__new__(MegatronCoreTrainEngine) + result.model, result.optimizer = [chunk], optimizer + result.optimizer_revision, result.current_version = 3, 0 + result.get_parallel_state = lambda: dict(train_tp=1, train_dp=1) + return result, master, model + + source, source_master, source_model = engine([1., 2.]) + for _ in range(3): + source_master.grad = torch.tensor([.1, -.2]) + source.optimizer.optimizer.step() + with torch.no_grad(): + source_model.copy_(source_master) + snapshot = source.capture_elastic_state() + target, target_master, target_model = engine([-9., -9.]) + # Actual pinned DP-zero gather/scatter excludes scalar step from tensor + # transfers, while its native setter iterates the destination step key. + target.optimizer.load_state_dict(snapshot['optimizer']) + with pytest.raises(KeyError, match='step'): + target.optimizer.load_parameter_state_from_dp_zero(snapshot['distributed_optimizer']['state']) + monkeypatch.setattr(module.DistributedOptimizer, '_set_main_param_and_optimizer_states', existing_setter) + receipt = target.restore_elastic_state(snapshot) + assert receipt['acknowledged'] and receipt['optimizer_revision'] == 3 + assert '_set_main_param_and_optimizer_states' not in vars(target.optimizer) + assert module.DistributedOptimizer._set_main_param_and_optimizer_states is existing_setter + torch.testing.assert_close(target_master, source_master, rtol=0, atol=0) + torch.testing.assert_close(target_model, source_model, rtol=0, atol=0) + for key in ('step', 'exp_avg', 'exp_avg_sq'): + torch.testing.assert_close(target.optimizer.optimizer.state[target_master][key], + source.optimizer.optimizer.state[source_master][key], rtol=0, atol=0) + for adapter, parameter in ((source, source_master), (target, target_master)): + parameter.grad = torch.tensor([-.3, .4]) + adapter.optimizer.optimizer.step() + torch.testing.assert_close(target_master, source_master, rtol=0, atol=0) + incomplete = MegatronCoreTrainEngine._clone_cpu_state(snapshot) + dtype = (torch.bfloat16, torch.float32) + del incomplete['distributed_optimizer']['state'][0][dtype]['exp_avg'] + with pytest.raises(KeyError, match='exp_avg'): + target.restore_elastic_state(incomplete) + assert '_set_main_param_and_optimizer_states' not in vars(target.optimizer) + assert module.DistributedOptimizer._set_main_param_and_optimizer_states is existing_setter + finally: + dist.destroy_process_group() diff --git a/tests/test_optimizer_state.py b/tests/test_optimizer_state.py new file mode 100644 index 0000000..9b7c9f2 --- /dev/null +++ b/tests/test_optimizer_state.py @@ -0,0 +1,52 @@ +"""Strict Adam metadata restoration must also preserve the next optimizer step.""" +from types import SimpleNamespace + +import pytest +import torch + +from RL_Framework.engine.optimizer_state import optimizer_state_load_context + + +def test_restored_parameter_steps_do_not_share_mutable_metadata_tensor(): + restored = [torch.nn.Parameter(torch.tensor([1.])), torch.nn.Parameter(torch.tensor([2.]))] + reference = [torch.nn.Parameter(parameter.detach().clone()) for parameter in restored] + actual = torch.optim.AdamW(restored, lr=.01) + expected = torch.optim.AdamW(reference, lr=.01) + shared = torch.tensor(3.) + for optimizer, parameters in [(actual, restored), (expected, reference)]: + for parameter in parameters: + optimizer.state[parameter] = { + 'step': shared if optimizer is actual else shared.clone(), + 'exp_avg': torch.tensor([.1]), 'exp_avg_sq': torch.tensor([.2]), + } + wrapper = SimpleNamespace(optimizer=actual, model_param_group_index_map={ + parameter: (0, index) for index, parameter in enumerate(restored)}) + received = [] + wrapper._set_main_param_and_optimizer_states = lambda parameter, tensors: received.append(tensors) + original = wrapper._set_main_param_and_optimizer_states + with optimizer_state_load_context(wrapper): + wrapper._set_main_param_and_optimizer_states(restored[0], {'exp_avg': torch.tensor([.1])}) + assert received[0]['step'].item() == 3 + assert wrapper._set_main_param_and_optimizer_states is original + assert actual.state[restored[0]]['step'] is not actual.state[restored[1]]['step'] + for optimizer, parameters in [(actual, restored), (expected, reference)]: + for parameter in parameters: + parameter.grad = torch.tensor([.3]) + optimizer.step() + for actual_parameter, expected_parameter in zip(restored, reference): + torch.testing.assert_close(actual_parameter, expected_parameter, rtol=0, atol=0) + assert actual.state[actual_parameter]['step'].item() == 4 + + +def test_missing_master_or_moment_state_is_not_invented(): + parameter = torch.nn.Parameter(torch.tensor([1.])) + optimizer = torch.optim.AdamW([parameter]) + optimizer.state[parameter]['step'] = torch.tensor(3.) + def strict_setter(parameter, tensors): + return tensors['exp_avg'] + wrapper = SimpleNamespace(optimizer=optimizer, model_param_group_index_map={parameter: (0, 0)}, + _set_main_param_and_optimizer_states=strict_setter) + with pytest.raises(KeyError, match='exp_avg'): + with optimizer_state_load_context(wrapper): + wrapper._set_main_param_and_optimizer_states(parameter, {}) + assert wrapper._set_main_param_and_optimizer_states is strict_setter From 5b06370ece826444de8cfd81cd631cf67df0d65a Mon Sep 17 00:00:00 2001 From: cytsoi <154662279+tcyjohn@users.noreply.github.com> Date: Thu, 22 Oct 2026 13:52:51 +0800 Subject: [PATCH 6/9] Keep GRPO gradients and statistics independent of replica partitioning Partition one fixed global batch into complete GRPO groups, carry actual sample counts through Hybrid tasks and gradient packets, and weight Core and Hybrid contributions by those counts. Preserve empty Core lanes in the collective and give asynchronous CPU gradient packets independent storage. Use the mean of each trajectory's valid-token loss, followed by a sample mean, for the shared Core/Hybrid objective and PG/KL statistics. Zero-mask trajectories contribute zero while remaining in the batch. This intentionally changes the previous objective for variable-length microbatches larger than one; clipping, the KL coefficient and Adam remain unchanged. Include the partition, nonzero-gradient, sample-mean, worker-task and transport regressions with the implementation. --- engine/megatron_core_train_engine.py | 140 ++++++-- infra/elastic/batch_partition.py | 68 ++++ infra/elastic/gradient_ipc.py | 10 +- infra/elastic/hybrid_pool.py | 133 ++++++-- infra/elastic/native_rdma.py | 6 + scripts/elastic_hybrid_worker.py | 4 +- scripts/test_cpu_ci.sh | 13 +- tests/hybrid_statistics_helpers.py | 77 +++++ tests/test_ehp_batch_partition.py | 210 ++++++++++++ tests/test_elastic_gradient_idle_path.py | 50 ++- tests/test_fixed_batch_gradients.py | 269 +++++++++++++++ tests/test_grpo_sample_mean.py | 95 ++++++ tests/test_hybrid_sample_statistics_data.py | 53 +++ tests/test_hybrid_worker_batch_protocol.py | 61 ++++ tests/test_megatron_core_train_engine.py | 8 +- trainer/async_rl_trainer.py | 355 +++++++++++++------- 16 files changed, 1366 insertions(+), 186 deletions(-) create mode 100644 infra/elastic/batch_partition.py create mode 100644 tests/hybrid_statistics_helpers.py create mode 100644 tests/test_ehp_batch_partition.py create mode 100644 tests/test_fixed_batch_gradients.py create mode 100644 tests/test_grpo_sample_mean.py create mode 100644 tests/test_hybrid_sample_statistics_data.py create mode 100644 tests/test_hybrid_worker_batch_protocol.py diff --git a/engine/megatron_core_train_engine.py b/engine/megatron_core_train_engine.py index 32a36c6..e411149 100644 --- a/engine/megatron_core_train_engine.py +++ b/engine/megatron_core_train_engine.py @@ -109,6 +109,7 @@ def __init__( self.is_main_process = self.rank == 0 self.current_version = 0 self.optimizer_revision = 0 + self._elastic_batch_partition = None self.bridge = None self.provider = None @@ -532,7 +533,8 @@ def grpo_update( logits = self._forward(micro_batch) loss, stats = self._loss(micro_batch, logits) scaled_loss = self.optimizer.scale_loss( - loss / max(len(micro_batches), 1) + loss * (len(micro_batch["trajectories"]) / max(len(trajectories), 1)) + * self._elastic_core_loss_scale() ) scaled_loss.backward() local_stats.append(stats) @@ -757,6 +759,23 @@ def set_elastic_training_step(self, step: int, active_hybrid_workers=None) -> No if active_hybrid_workers is not None: self._elastic_step_hybrid_workers = tuple(str(w) for w in active_hybrid_workers) + def set_elastic_batch_partition(self, local_sample_count, global_sample_count, hybrid_sample_counts): + if local_sample_count < 0 or global_sample_count <= 0 or local_sample_count > global_sample_count: + raise ValueError("invalid fixed global batch partition") + self._elastic_batch_partition = (int(local_sample_count), int(global_sample_count), dict(hybrid_sample_counts)) + + + def _elastic_core_loss_scale(self): + partition = self._elastic_batch_partition + if partition is None: + return 1.0 + local_count, total_count, _ = partition + # MCore finalizes a MEAN across core DP first. Weight before that + # operation, then remove its duplicate copies in the elastic SUM. + world = max(1, self.get_data_parallel_world_size()) + return local_count * world / total_count + + def set_elastic_active_gradient_timeout(self, timeout_s: float) -> None: self._elastic_active_gradient_timeout_s = max(0.0, float(timeout_s)) @@ -775,14 +794,17 @@ def _model_main_gradients(self) -> list[torch.Tensor]: def compute_elastic_gradient_payload( self, trajectories: list[dict[str, Any]], *, worker_id: str, target_core_id: str, step: int, state_version: int, membership_epoch: int, + global_sample_count: int = 0, ) -> GradientPayload: from megatron.core.distributed import finalize_model_grads micro_batches = list(self._iter_micro_batches(trajectories)) self._zero_grad() self._set_train_mode(True) + local_stats = [] for micro_batch in micro_batches: - loss, _ = self._loss(micro_batch, self._forward(micro_batch)) - self.optimizer.scale_loss(loss / max(len(micro_batches), 1)).backward() + loss, stats = self._loss(micro_batch, self._forward(micro_batch)) + self.optimizer.scale_loss(loss * (len(micro_batch["trajectories"]) / max(len(trajectories), 1))).backward() + local_stats.append(stats) finalize_model_grads(self.model) return GradientPayload( replica_id=worker_id, target_core_id=target_core_id, @@ -790,6 +812,8 @@ def compute_elastic_gradient_payload( replica_rank=int(os.environ.get("RANK", "0")), replica_world_size=int(os.environ.get("WORLD_SIZE", "1")), step=int(step), state_version=int(state_version), membership_epoch=int(membership_epoch), + sample_count=len(trajectories), global_sample_count=int(global_sample_count), + stats=self._merge_stats(local_stats), ) def apply_elastic_gradient_update(self, tensors: tuple[torch.Tensor, ...], *, state_version: int) -> None: @@ -801,7 +825,7 @@ def apply_elastic_gradient_update(self, tensors: tuple[torch.Tensor, ...], *, st ok, _norm, _zeros = self.optimizer.step() if not ok: raise FloatingPointError("Hybrid Megatron optimizer rejected the update") - self.optimizer_revision = int(state_version) + self.current_version = int(state_version) @staticmethod def _clone_cpu_state(value): @@ -1081,11 +1105,13 @@ def _loss( -advantages.unsqueeze(1) * ratio, -advantages.unsqueeze(1) * clipped_ratio, ) - count = loss_mask.sum().clamp(min=1.0) - pg_loss = (pg_loss * loss_mask).sum() / count + # A trajectory is one sample of the fixed global batch. Its weight + # must not depend on its token count or on the elastic microbatch split. + count = loss_mask.sum(dim=-1).clamp(min=1.0) + pg_loss = ((pg_loss * loss_mask).sum(dim=-1) / count).mean() # Positive second-order approximation of KL(old_policy || new_policy). kl_per_token = (ratio - 1.0) - log_ratio - kl = (kl_per_token * loss_mask).sum() / count + kl = ((kl_per_token * loss_mask).sum(dim=-1) / count).mean() loss = pg_loss + self.kl_coef * kl return loss, { "loss": float(loss.detach()), @@ -1098,11 +1124,12 @@ def _loss( def _reduce_stats(self, stats: dict[str, float]) -> dict[str, float]: updates = float(stats["n_updates"]) + samples = float(stats["n_samples"]) values = torch.tensor( [ - stats["loss"] * updates, - stats["pg_loss"] * updates, - stats["kl"] * updates, + stats["loss"] * samples, + stats["pg_loss"] * samples, + stats["kl"] * samples, stats["reward_sum"], stats["n_samples"], updates, @@ -1118,10 +1145,11 @@ def _reduce_stats(self, stats: dict[str, float]) -> dict[str, float]: ) dist.all_reduce(values, op=dist.ReduceOp.SUM, group=group) global_updates = values[5].clamp(min=1.0) + global_samples = values[4].clamp(min=1.0) return { - "loss": float(values[0] / global_updates), - "pg_loss": float(values[1] / global_updates), - "kl": float(values[2] / global_updates), + "loss": float(values[0] / global_samples), + "pg_loss": float(values[1] / global_samples), + "kl": float(values[2] / global_samples), "reward_sum": float(values[3]), "n_samples": float(values[4]), "n_updates": float(values[5]), @@ -1146,17 +1174,17 @@ def _merge_stats(stats: list[dict[str, float]]) -> dict[str, float]: reward_sum = sum(float(item.get("reward_sum", 0.0)) for item in stats) return { "loss": sum( - float(item["loss"]) * float(item.get("n_updates", 1.0)) + float(item["loss"]) * float(item.get("n_samples", 0.0)) for item in stats - ) / max(updates, 1.0), + ) / max(samples, 1.0), "pg_loss": sum( - float(item["pg_loss"]) * float(item.get("n_updates", 1.0)) + float(item["pg_loss"]) * float(item.get("n_samples", 0.0)) for item in stats - ) / max(updates, 1.0), + ) / max(samples, 1.0), "kl": sum( - float(item["kl"]) * float(item.get("n_updates", 1.0)) + float(item["kl"]) * float(item.get("n_samples", 0.0)) for item in stats - ) / max(updates, 1.0), + ) / max(samples, 1.0), "reward_sum": reward_sum, "reward_mean": reward_sum / max(samples, 1.0), "n_samples": samples, @@ -1177,16 +1205,49 @@ def _set_train_mode(self, enabled: bool) -> None: for model_chunk in self.model: model_chunk.train(enabled) + def _materialize_elastic_core_gradients(self): + """Restore full bucket gradients after MCore distributed-optimizer RS. + + MCore 0.14 only synchronizes the local shard of each grad bucket when + using DistributedOptimizer. Hybrid exchange is per model lane, so its + payload must see all finalized parameter gradients, not stale remote + shards that remain in the bucket after reduce-scatter. + """ + if not getattr(self, "use_distributed_optimizer", False): + return + if not dist.is_available() or not dist.is_initialized(): + return + for model in self.model: + groups = list(getattr(model, "bucket_groups", ())) + groups += list(getattr(model, "expert_parallel_bucket_groups", ())) + for bucket_group in groups: + group = bucket_group.intra_distributed_optimizer_instance_group + world = dist.get_world_size(group) + if world <= 1: + continue + rank = dist.get_rank(group) + for bucket in bucket_group.buckets: + flat = bucket.grad_data.view(-1) + if flat.numel() % world: + raise ValueError("distributed optimizer grad bucket is not evenly sharded") + width = flat.numel() // world + # The output aliases the original RS input; clone only the + # finalized local shard before gathering back into it. + local = flat[rank * width:(rank + 1) * width].clone() + dist.all_gather_into_tensor(flat, local, group=group) + + def _apply_elastic_inter_replica_gradients(self) -> None: pending = getattr(self, "_pending_hybrid_gradients", []) expected = set(getattr(self, "_elastic_step_hybrid_workers", ())) - if not pending and not expected: + if self.elastic_gradient_domain is None and not pending and not expected: return if self.elastic_gradient_domain is None: with self._hybrid_gradient_condition: self._pending_hybrid_gradients.clear() return - deadline = time.monotonic() + self._elastic_active_gradient_timeout_s + deadline = (time.monotonic() + self._elastic_active_gradient_timeout_s + if self._elastic_active_gradient_timeout_s > 0 else None) with self._hybrid_gradient_condition: while expected: present = { @@ -1197,14 +1258,15 @@ def _apply_elastic_inter_replica_gradients(self) -> None: missing = expected - present if not missing: break - remaining = deadline - time.monotonic() - if remaining <= 0: + remaining = deadline - time.monotonic() if deadline is not None else None + if remaining is not None and remaining <= 0: raise TimeoutError( "timed out waiting for active hybrid gradients: " f"step={self._elastic_training_step}, missing={sorted(missing)}" ) - self._hybrid_gradient_condition.wait(timeout=min(remaining, 0.1)) + self._hybrid_gradient_condition.wait(timeout=min(remaining, 0.1) if remaining is not None else 0.1) pending = list(self._pending_hybrid_gradients) + self._materialize_elastic_core_gradients() params_and_grads = [(None, grad) for grad in self._model_main_gradients()] if not params_and_grads: with self._hybrid_gradient_condition: @@ -1212,11 +1274,21 @@ def _apply_elastic_inter_replica_gradients(self) -> None: return core_id = f"dp{self.get_data_parallel_rank()}" + partition = self._elastic_batch_partition + weighted = partition is not None + world = len(self.elastic_gradient_domain.core_replica_ids) + # finalize_model_grads has already averaged the weighted core sum + # across DP. Each lane holds the same core contribution, so divide by + # DP before the second SUM, while adding each hybrid exactly once. + local = tuple(grad.detach() for _, grad in params_and_grads) + if weighted and world > 1: + for grad in local: + grad.div_(world) reduced = self.elastic_gradient_domain.reduce_core_gradients( - core_gradients={ - core_id: tuple(grad.detach() for _, grad in params_and_grads), - }, - hybrid_payloads=pending, + core_gradients={core_id: local}, hybrid_payloads=pending, + core_sample_counts={core_id: partition[1]} if weighted else None, + global_sample_count=partition[1] if weighted else 0, + in_place=True, ) for (_, grad), reduced_grad in zip( params_and_grads, @@ -1226,12 +1298,22 @@ def _apply_elastic_inter_replica_gradients(self) -> None: reduced_grad.to(device=grad.device, dtype=grad.dtype) ) if self._elastic_gradient_update_callback is not None: + # The callback can queue this update for asynchronous serialization. + # Optimizer clipping and the next zero_grad may mutate main_grad as + # soon as it returns. Preserve an owned CPU snapshot, rather than + # restoring a full-model FP32 duplicate on the GPU. + update_tensors = None for payload in pending: if payload.replica_id in expected: + if update_tensors is None: + update_tensors = tuple( + tensor.detach().to(device="cpu", copy=True) + for tensor in reduced[core_id] + ) self._elastic_gradient_update_callback( GradientUpdate( replica_id=payload.replica_id, - tensors=tuple(reduced[core_id]), + tensors=update_tensors, step=payload.step, state_version=self.current_version, membership_epoch=payload.membership_epoch, diff --git a/infra/elastic/batch_partition.py b/infra/elastic/batch_partition.py new file mode 100644 index 0000000..b4577ea --- /dev/null +++ b/infra/elastic/batch_partition.py @@ -0,0 +1,68 @@ +"""Fixed global-batch allocation in whole GRPO prompt groups.""" +from __future__ import annotations + +from collections import OrderedDict +from typing import Any, Sequence + + +def complete_group_sample_counts(global_batch_size: int, n_samples: int, replicas: int) -> list[int]: + """Balance complete groups, allowing replicas with no assigned group.""" + if n_samples < 1 or replicas < 1 or global_batch_size < 1: + raise ValueError("batch size, group size and replica count must be positive") + if global_batch_size % n_samples: + raise ValueError("global batch size must contain complete GRPO groups") + groups, remainder = divmod(global_batch_size // n_samples, replicas) + return [(groups + (index < remainder)) * n_samples for index in range(replicas)] + + +def partition_complete_grpo_groups( + trajectories: list[dict[str, Any]], n_samples: int, participant_ids: Sequence[str], +) -> dict[str, list[dict[str, Any]]]: + """Partition one global batch without expanding, duplicating or splitting it.""" + if len(set(participant_ids)) != len(participant_ids): + raise ValueError("training participant IDs must be unique") + counts = complete_group_sample_counts(len(trajectories), n_samples, len(participant_ids)) + groups: dict[str, list[dict[str, Any]]] = OrderedDict() + for trajectory in trajectories: + groups.setdefault(str(trajectory["grpo_group_id"]), []).append(trajectory) + if any(len(members) != n_samples for members in groups.values()): + raise ValueError("global batch contains incomplete or duplicate GRPO groups") + for members in groups.values(): + indices = [member.get("rollout_index") for member in members] + if all(index is not None for index in indices) and len(set(indices)) != n_samples: + raise ValueError("global batch contains duplicate GRPO members") + ordered = [member for members in groups.values() for member in members] + result: dict[str, list[dict[str, Any]]] = {} + offset = 0 + for participant, count in zip(participant_ids, counts): + result[participant] = ordered[offset:offset + count] + offset += count + return result + + +def unique_group_members(members, n_samples=None): + """Keep the first successful occurrence of each explicit member identity. + + Older workflows without member IDs retain their positional semantics; their + distinct samples cannot safely be identified by text or tensor equality. + """ + selected, seen = [], set() + for member in members: + if not isinstance(member, dict) or member.get("status") == "failed": + continue + index = member.get("grpo_rollout_index", member.get("rollout_index")) + if index is not None: + try: + index = int(index) + except (ValueError, TypeError): + continue + if n_samples is not None and not 0 <= index < n_samples: + continue + keys = [("member", index)] if index is not None else [] + if member.get("task_id") is not None: + keys.append(("task", member["task_id"])) + if any(key in seen for key in keys): + continue + seen.update(keys) + selected.append(member) + return selected diff --git a/infra/elastic/gradient_ipc.py b/infra/elastic/gradient_ipc.py index e1d3b1c..4455832 100644 --- a/infra/elastic/gradient_ipc.py +++ b/infra/elastic/gradient_ipc.py @@ -70,7 +70,7 @@ def _serialize_metadata(self, payload: GradientPayload, *, expect_update: bool) buffer = io.BytesIO() torch.save( { - "protocol_version": 3, + "protocol_version": 4, "authkey": self.endpoint.authkey, "replica_id": payload.replica_id, "target_core_id": payload.target_core_id, @@ -82,6 +82,9 @@ def _serialize_metadata(self, payload: GradientPayload, *, expect_update: bool) "state_version": payload.state_version, "membership_epoch": payload.membership_epoch, "expect_update": bool(expect_update), + "sample_count": payload.sample_count, + "global_sample_count": payload.global_sample_count, + "stats": payload.stats, }, buffer, ) @@ -279,7 +282,7 @@ def _recv_payload(self, conn: socket.socket) -> tuple[GradientPayload, bool]: data = torch.load(io.BytesIO(self._recv_exact(conn, length)), map_location="cpu", weights_only=False) if self.authkey and data.get("authkey") != self.authkey: raise PermissionError("invalid elastic gradient authkey") - if int(data.get("protocol_version", 0)) not in {2, 3}: + if int(data.get("protocol_version", 0)) not in {2, 3, 4}: raise ValueError("unsupported elastic gradient protocol") tensors = [] for _ in range(int(data["tensor_count"])): @@ -295,6 +298,9 @@ def _recv_payload(self, conn: socket.socket) -> tuple[GradientPayload, bool]: step=int(data.get("step", -1)), state_version=int(data.get("state_version", -1)), membership_epoch=int(data.get("membership_epoch", 0)), + sample_count=int(data.get("sample_count", 1)), + global_sample_count=int(data.get("global_sample_count", 0)), + stats=data.get("stats", {}), ), bool(data.get("expect_update", False)) @staticmethod diff --git a/infra/elastic/hybrid_pool.py b/infra/elastic/hybrid_pool.py index d704ae4..d574e3c 100644 --- a/infra/elastic/hybrid_pool.py +++ b/infra/elastic/hybrid_pool.py @@ -71,6 +71,9 @@ class GradientPayload: step: int = -1 state_version: int = -1 membership_epoch: int = 0 + sample_count: int = 1 + global_sample_count: int = 0 + stats: dict[str, float] = field(default_factory=dict) @staticmethod def zeros_like( @@ -510,12 +513,22 @@ def reduce_core_gradients( *, core_gradients: Mapping[str, Iterable[torch.Tensor]], hybrid_payloads: Iterable[GradientPayload] = (), + core_sample_counts: Mapping[str, int] | None = None, + global_sample_count: int = 0, + step: int = -1, + state_version: int = -1, + in_place: bool = False, ) -> dict[str, tuple[torch.Tensor, ...]]: - """Return synchronized gradients for each core replica. - - ``core_gradients`` must contain all fixed core replicas. Hybrid payloads - whose workers are still joining are treated as zero placeholders even if - a caller accidentally supplies non-zero tensors. + """Reduce local sample means into the fixed global batch gradient. + + With sample counts, each local mean contributes ``n_local / B`` and the + core collective sums these contributions. Joining replicas have zero + weight. Every core lane participates, including lanes with no Hybrid. + Unweighted calls retain the legacy backend-neutral simulation API. + ``in_place`` explicitly transfers ownership of one local core lane's + gradient buffers to this call. It bounds temporary storage to transport + chunks, rather than duplicating every FP32 model gradient. The default + retains independent output buffers and leaves its inputs untouched. """ distributed = ( dist is not None @@ -532,66 +545,126 @@ def reduce_core_gradients( raise ValueError("no local core gradients were provided") effective = { - rid: tuple(t.clone() for t in core_gradients[rid]) + rid: tuple(core_gradients[rid]) for rid in provided_core_ids } self._validate_same_structure(effective.values()) - + if in_place and len(provided_core_ids) != 1: + raise ValueError("in-place reduction requires exactly one local core lane") + + weighted = global_sample_count > 0 + if weighted: + if core_sample_counts is None: + raise ValueError("sample counts are required for a fixed global batch") + scales = {rid: int(core_sample_counts[rid]) / global_sample_count + for rid in provided_core_ids} + else: + scales = {rid: 1.0 for rid in provided_core_ids} with self._lock: targets = dict(self._hybrid_targets) joining = set(self._joining) - + epochs = dict(self._membership_epochs) + accepted = [] for payload in hybrid_payloads: - if payload.replica_id not in targets: + target = targets.get(payload.replica_id) + if target not in effective or payload.replica_id in joining: continue - target = targets[payload.replica_id] - if target not in effective: + if payload.zero_placeholder: continue if target != payload.target_core_id: - raise ValueError( - f"payload target mismatch for {payload.replica_id}: " - f"{payload.target_core_id} != {target}" - ) - if payload.replica_id in joining or payload.zero_placeholder: - continue + raise ValueError(f"payload target mismatch for {payload.replica_id}") if payload.replica_world_size != self.replica_world_size: - raise ValueError( - f"payload replica world size mismatch for {payload.replica_id}: " - f"{payload.replica_world_size} != {self.replica_world_size}" - ) + raise ValueError(f"payload replica world size mismatch for {payload.replica_id}") if payload.replica_rank != self.local_replica_rank: - # Each fixed TP/PP/CP lane consumes only the matching rank from - # the complete hybrid DP replica. continue if len(payload.tensors) != len(effective[target]): raise ValueError("hybrid payload tensor count mismatch") - effective[target] = tuple( - base + extra.to(device=base.device, dtype=base.dtype) - for base, extra in zip(effective[target], payload.tensors) - ) - + if any(a.shape != b.shape for a, b in zip(effective[target], payload.tensors)): + raise ValueError("hybrid payload tensor shape mismatch") + if weighted and payload.global_sample_count != global_sample_count: + raise ValueError("hybrid global batch size mismatch") + if payload.sample_count < 0: + raise ValueError("negative hybrid sample count") + scale = payload.sample_count / global_sample_count if weighted else 1.0 + accepted.append((target, payload, scale)) + # Validate the complete payload set before mutating an owned buffer or + # marking a payload consumed, including failures in a later payload. + if in_place: + all_tensors = [t for tensors in effective.values() for t in tensors] + all_tensors += [t for _, payload, _ in accepted for t in payload.tensors] + if any(not t.is_contiguous() for t in all_tensors): + raise ValueError("in-place reduction requires contiguous gradients") + with torch.no_grad(): + for rid, tensors in effective.items(): + if scales[rid] != 1.0: + for tensor in tensors: + tensor.mul_(scales[rid]) + for target, payload, scale in accepted: + for base, extra in zip(effective[target], payload.tensors): + width = max(DEFAULT_CHUNK_BYTES // max(base.element_size(), 1), 1) + destination, source = base.view(-1), extra.view(-1) + for start in range(0, base.numel(), width): + stop = start + width + # Preserve multiply-then-add rounding; add_(alpha=) + # can fuse them and change the numerical update. + contribution = source[start:stop].to( + device=base.device, dtype=base.dtype, copy=True) + contribution.mul_(scale) + destination[start:stop].add_(contribution) + del contribution + else: + effective = {rid: tuple(t.clone() for t in tensors) + for rid, tensors in effective.items()} + if weighted: + effective = {rid: tuple(t * scales[rid] for t in tensors) + for rid, tensors in effective.items()} + for target, payload, scale in accepted: + effective[target] = tuple( + base + extra.to(device=base.device, dtype=base.dtype) * scale + for base, extra in zip(effective[target], payload.tensors) + ) # In a real distributed launch each rank owns one core replica lane. The # transport call performs the RDMA/NCCL all-reduce. For local tests we # compute the same reduction explicitly across the in-memory core map. if distributed and self.process_group is not None: + if in_place: + with torch.no_grad(): + for tensors in effective.values(): + for tensor in tensors: + width = max(DEFAULT_CHUNK_BYTES // max(tensor.element_size(), 1), 1) + flat = tensor.view(-1) + for start in range(0, flat.numel(), width): + chunk = flat[start:start + width] + reduced = self.transport.exchange( + chunk, group=self.process_group, + average=self.average_core_replicas and not weighted, + ) + chunk.copy_(reduced) + del reduced + return effective return { rid: tuple( self.transport.exchange( tensor, group=self.process_group, - average=self.average_core_replicas, + average=self.average_core_replicas and not weighted, ) for tensor in tensors ) for rid, tensors in effective.items() } + if in_place: + # There is exactly one lane here. Avoid stack/sum/clone of that + # entire lane when no collective is needed (notably Core DP1). + return effective + reduced = [] for tensors_at_idx in zip(*(effective[rid] for rid in provided_core_ids)): total = torch.stack( [t.to(tensors_at_idx[0].device) for t in tensors_at_idx] ).sum(dim=0) - if self.average_core_replicas: + if self.average_core_replicas and not weighted: total = total / len(provided_core_ids) reduced.append(total) return { diff --git a/infra/elastic/native_rdma.py b/infra/elastic/native_rdma.py index b80308a..d8aedd2 100644 --- a/infra/elastic/native_rdma.py +++ b/infra/elastic/native_rdma.py @@ -259,6 +259,9 @@ def send(self, payload: GradientPayload): "step": payload.step, "state_version": payload.state_version, "membership_epoch": payload.membership_epoch, + "sample_count": payload.sample_count, + "global_sample_count": payload.global_sample_count, + "stats": payload.stats, }, path, ) @@ -358,6 +361,9 @@ def _serve(self): step=int(data.get("step", -1)), state_version=int(data.get("state_version", -1)), membership_epoch=int(data.get("membership_epoch", 0)), + sample_count=int(data.get("sample_count", 1)), + global_sample_count=int(data.get("global_sample_count", 0)), + stats=data.get("stats", {}), ) self.on_payload(payload) self._received += 1 diff --git a/scripts/elastic_hybrid_worker.py b/scripts/elastic_hybrid_worker.py index 0384578..992d237 100644 --- a/scripts/elastic_hybrid_worker.py +++ b/scripts/elastic_hybrid_worker.py @@ -59,6 +59,8 @@ def _load_payload_file(path: Path, *, worker_id: str, target_core_id: str, args) step=int(data.get("step", -1)) if isinstance(data, dict) else -1, state_version=int(data.get("state_version", -1)) if isinstance(data, dict) else -1, membership_epoch=int(data.get("membership_epoch", 0)) if isinstance(data, dict) else 0, + sample_count=int(data.get("sample_count", 1)) if isinstance(data, dict) else 1, + global_sample_count=int(data.get("global_sample_count", 0)) if isinstance(data, dict) else 0, ) @@ -179,7 +181,7 @@ def main(): engine.load_elastic_state_snapshot(task_snapshot) snapshot = Path(task_snapshot) engine.recompute_logprobs(data["trajectories"]) - payload = engine.compute_elastic_gradient_payload(data["trajectories"], worker_id=args.replica_id or args.worker_id, target_core_id=args.target_core_id, step=int(data["step"]), state_version=int(data["state_version"]), membership_epoch=int(data["membership_epoch"])) + payload = engine.compute_elastic_gradient_payload(data["trajectories"], worker_id=args.replica_id or args.worker_id, target_core_id=args.target_core_id, step=int(data["step"]), state_version=int(data["state_version"]), membership_epoch=int(data["membership_epoch"]), global_sample_count=int(data.get("global_sample_count", 0))) else: payload = _load_payload_file(task_path, worker_id=args.worker_id, target_core_id=args.target_core_id, args=args) if payload is not None: diff --git a/scripts/test_cpu_ci.sh b/scripts/test_cpu_ci.sh index d73bb53..8bebfe7 100644 --- a/scripts/test_cpu_ci.sh +++ b/scripts/test_cpu_ci.sh @@ -25,4 +25,15 @@ python -m pytest -q --strict-markers --junitxml=reports/cpu-tests.xml \ tests/test_rollout_engine.py \ tests/test_cpu_offload_backend.py \ tests/test_r2e_patch_harness.py \ - tests/test_r2e_gym_workflow.py + tests/test_r2e_gym_workflow.py \ + tests/test_optimizer_state.py \ + tests/test_elastic_optimizer_state.py \ + tests/test_ehp_dp_resharding.py \ + tests/test_ehp_batch_partition.py \ + tests/test_fixed_batch_gradients.py \ + tests/test_grpo_sample_mean.py \ + tests/test_hybrid_sample_statistics_data.py \ + tests/test_elastic_gradient_idle_path.py \ + tests/test_torch_optimizer_memory.py \ + tests/test_megatron_core_train_engine.py \ + tests/test_hybrid_worker_batch_protocol.py diff --git a/tests/hybrid_statistics_helpers.py b/tests/hybrid_statistics_helpers.py new file mode 100644 index 0000000..19fed15 --- /dev/null +++ b/tests/hybrid_statistics_helpers.py @@ -0,0 +1,77 @@ +"""Small CPU model fixture shared by payload-data and activation-integration tests.""" +import sys +import threading +from types import SimpleNamespace as NS + +import torch + +from RL_Framework.engine import megatron_core_train_engine as module +from RL_Framework.engine.megatron_core_train_engine import MegatronCoreTrainEngine as Engine +from RL_Framework.infra.elastic.hybrid_pool import InterReplicaGradientDomain + + +def rows(start, count, zero_mask=False): + return [dict(input_ids=torch.zeros(1, 2+i%3, dtype=torch.long), + logprobs=torch.full((1,2+i%3), -.7), + loss_mask=torch.zeros(1,2+i%3) if zero_mask else torch.ones(1,2+i%3), + rewards=torch.tensor([float(i)]), advantages=torch.tensor([1.+i/5]), + feature=float(i)/3) + for i in range(start,start+count)] + + +def engine(monkeypatch, lane=0, width=1): + # Only MCore finalization is replaced: the loss, microbatching, payload + # compute, enqueue, gradient reduction and statistics methods are real. + monkeypatch.setitem(sys.modules, 'megatron.core.distributed', + NS(finalize_model_grads=lambda model: None)) + ex=Engine.__new__(Engine) + parameter=torch.tensor(.1,requires_grad=True) + ex.train_tp_size=1 + ex.micro_batch_size=3 + ex.tokenizer=NS(pad_token_id=0) + ex.clip_epsilon=.2 + ex.kl_coef=.1 + ex.optimizer_revision=2 + ex.current_version=2 + ex.model=[] + ex._device=lambda:'cpu' + ex._zero_grad=lambda:setattr(parameter,'grad',torch.zeros_like(parameter)) + ex._set_train_mode=lambda enabled:None + def forward(batch): + features=torch.tensor([r['feature'] for r in batch['trajectories']])[:,None] + values=(parameter+features).expand_as(batch['input_ids']) + return torch.stack((values,torch.zeros_like(values)),dim=-1) + ex._forward=forward + ex._model_main_gradients=lambda:[parameter.grad] + ex.use_distributed_optimizer=False + ex.get_data_parallel_rank=lambda:0 + ex.get_data_parallel_world_size=lambda:1 + ex._parallel_state=lambda:NS(get_data_parallel_group=lambda **kw:'dp-lane-'+str(lane)) + monkeypatch.setattr(module.dist,'all_reduce',lambda values,**kw:None) + steps=[] + ex.optimizer=NS(scale_loss=lambda loss:loss, + step=lambda:(steps.append(True) or True,7.,0)) + ex._pending_hybrid_gradients=[] + ex._hybrid_gradient_condition=threading.Condition() + ex._elastic_seen_payloads=set() + ex._elastic_active_gradient_timeout_s=0 + ex._elastic_training_step=3 + ex._elastic_expected_state_version=2 + ex._elastic_gradient_update_callback=None + domain=InterReplicaGradientDomain(core_replica_ids=['dp0'], + replica_world_size=width,local_replica_rank=lane) + domain.request_join('hybrid','dp0',membership_epoch=4) + domain.mark_active('hybrid',membership_epoch=4) + ex.elastic_gradient_domain=domain + ex._elastic_step_hybrid_workers=('hybrid',) + return ex,steps + + +def compute(monkeypatch, subset, lane=0, width=1): + worker,_=engine(monkeypatch,lane,width) + monkeypatch.setenv('RANK',str(lane)) + monkeypatch.setenv('WORLD_SIZE',str(width)) + payload=worker.compute_elastic_gradient_payload(subset,worker_id='hybrid', + target_core_id='dp0',step=3,state_version=2,membership_epoch=4, + global_sample_count=12) + return worker,payload diff --git a/tests/test_ehp_batch_partition.py b/tests/test_ehp_batch_partition.py new file mode 100644 index 0000000..86b5978 --- /dev/null +++ b/tests/test_ehp_batch_partition.py @@ -0,0 +1,210 @@ +"""Fixed-B tests through production trainer collection, dispatch and train paths.""" +import json +from datetime import timedelta +from pathlib import Path +from types import SimpleNamespace + +import pytest +import torch +import torch.distributed as dist +import torch.multiprocessing as mp + +from RL_Framework.infra.elastic.batch_partition import ( + complete_group_sample_counts, partition_complete_grpo_groups, +) +from RL_Framework.trainer.async_rl_trainer import AsyncRLTrainer + + +def trajectories(prefix, count, n_samples=4): + return [dict(grpo_group_id=f"{prefix}-g{i // n_samples}", rollout_index=i % n_samples, + rewards=torch.tensor([float(i % n_samples)]), input_ids=torch.ones(1, 2, dtype=torch.long)) + for i in range(count)] + + +@pytest.mark.parametrize("batch,replicas,expected", [(16, 3, [8, 4, 4]), (8, 4, [4, 4, 0, 0])]) +def test_balances_whole_groups_including_zero_assignments(batch, replicas, expected): + assert complete_group_sample_counts(batch, 4, replicas) == expected + original = trajectories("global", batch) + assignment = partition_complete_grpo_groups(original, 4, [str(i) for i in range(replicas)]) + assert [len(part) for part in assignment.values()] == expected + assert [id(x) for part in assignment.values() for x in part] == [id(x) for x in original] + assert all(len({x["grpo_group_id"] for x in part}) * 4 == len(part) for part in assignment.values()) + + +def test_rejects_duplicate_member_and_incomplete_group(): + batch = trajectories("global", 8) + batch[1]["rollout_index"] = 0 + with pytest.raises(ValueError, match="duplicate GRPO members"): + partition_complete_grpo_groups(batch, 4, ["dp0"]) + with pytest.raises(ValueError, match="complete GRPO groups"): + partition_complete_grpo_groups(batch[:-1], 4, ["dp0"]) + + +class Engine: + def __init__(self, dp_rank=0, dp_size=1, source=True): + self.dp_rank, self.dp_size, self.source = dp_rank, dp_size, source + self.assignments, self.updated = [], [] + self.optimizer_revision = 0 + + def get_data_parallel_world_size(self): + return self.dp_size + + def get_data_parallel_rank(self): + return self.dp_rank + + def is_batch_source(self): + return self.source + + def get_version(self): + return 0 + + def set_elastic_batch_partition(self, local_count, global_count, hybrid_counts): + self.assignments.append((local_count, global_count, hybrid_counts)) + + def grpo_update(self, batch, ppo_epochs): + self.updated.append([x["grpo_group_id"] for x in batch]) + self.optimizer_revision += 1 + return {"loss": 0., "pg_loss": 0., "kl": 0.} + + +def make_trainer(tmp_path, rank=0, dp_size=1, dp_rank=0, source=True, hybrids=()): + trainer = object.__new__(AsyncRLTrainer) + trainer.rank, trainer.local_rank = rank, rank + trainer.world_size, trainer.is_main_process = dp_size, rank == 0 + trainer.config = SimpleNamespace( + batch_size=16, n_samples=4, total_steps=2, micro_batch_size=1, + ppo_epochs=1, recompute_logprobs=False, sync_interval=0, save_interval=0, + log_dir=str(tmp_path), enable_rollout_tracing=False, + global_resource_planner=SimpleNamespace(hybrid_worker_task_dir=str(tmp_path), + hybrid_lockstep_gradient_sync=True), + ) + trainer.train_engine = Engine(dp_rank, dp_size, source) + trainer.start_step = 0 + trainer.runtime_elastic_executor = None + trainer._use_heterogeneous = False + trainer.staleness_manager = None + trainer._pending_hybrid_training_batches = {} + trainer._elastic_membership_run_id = "fixed-b-test" + trainer.setup = lambda workflow: None + trainer._create_data_generator = lambda dataset, workflow: iter(()) + trainer._prepare_elastic_training_step = lambda step: list(hybrids) + for name in ("_trace_train_phase", "_configure_sync_aware_rollout_prefetch", + "_runtime_reconfiguration_boundary", "_wait_for_all_ranks_after_rollout", + "_log_to_wandb", "_record_history_step", "_run_global_resource_planner_at_boundary", + "_run_periodic_evaluation", "_cleanup"): + setattr(trainer, name, lambda *args, **kwargs: None) + return trainer + + +@pytest.mark.parametrize("hybrid_count,expected", [(0, 16), (1, 8), (5, 4)]) +def test_actual_train_collects_fixed_batch_and_dispatches_complete_groups(tmp_path, hybrid_count, expected): + hybrids = [f"hybrid{i}" for i in range(hybrid_count)] + trainer = make_trainer(tmp_path, hybrids=hybrids) + (tmp_path / "membership").mkdir() + for worker in hybrids: + (tmp_path / "membership" / f"{worker}.json").write_text(json.dumps( + dict(run_id="fixed-b-test", role="hybrid_training", membership_epoch=7))) + collected, histories = [], [] + + def collect(data_gen, batch_size): + collected.append(batch_size) + return trajectories(f"step{len(collected)}", batch_size) + + trainer._collect_complete_grpo_batch = collect + trainer._record_history_step = lambda step, batch, stats, *args: histories.append((len(batch), stats.copy())) + trainer.train(None, dataset=[1]) + assert collected == [16, 16] + assert [len(batch) for batch in trainer.train_engine.updated] == [expected, expected] + assert [row[0] for row in histories] == [16, 16] + assert [row[1]["n_trajectories"] for row in histories] == [16, 16] + for step in range(2): + tasks = [torch.load(tmp_path / f"{worker}.step_{step}.pt", weights_only=False) for worker in hybrids] + assert sum(len(task["trajectories"]) for task in tasks) + expected == 16 + assert all(task["global_sample_count"] == 16 for task in tasks) + assert all(task["sample_count"] == len(task["trajectories"]) for task in tasks) + assert all(task["state_version"] == 0 for task in tasks) + all_groups = trainer.train_engine.updated[step] + [x["grpo_group_id"] for task in tasks for x in task["trajectories"]] + assert len(all_groups) == 16 + assert len(set(all_groups)) == 4 + assert all(all_groups.count(group) == 4 for group in set(all_groups)) + + +def _gloo_collection_worker(rank, rendezvous, directory): + dist.init_process_group("gloo", rank=rank, world_size=4, + init_method=f"file://{rendezvous}", timeout=timedelta(seconds=30)) + try: + path = Path(directory) + results = [] + # Each case is a separately initialized fixed Core layout. These do + # not exercise or claim support for runtime Core topology changes. + for global_count, compute_dp_size in ((16, 2), (16, 4), (16, 1), (8, 4)): + width = 4 // compute_dp_size + trainer = make_trainer(path, rank, compute_dp_size, rank // width, rank % width == 0) + trainer.world_size = 4 + trainer.config.batch_size = global_count + trainer._initialize_rollout_source_layout() + trainer._collect_complete_grpo_batch = lambda generator, count: trajectories( + f"source{trainer.train_engine.dp_rank}", count) + batch, participants = trainer._collect_fixed_global_batch(None, trainer._source_batch_sample_count(), []) + own = trainer._assign_fixed_global_batch(0, batch, participants) + results.append({"global_groups": [x["grpo_group_id"] for x in batch], + "own_groups": [x["grpo_group_id"] for x in own]}) + torch.save(results, path / f"rank{rank}.pt") + trainer._collect_complete_grpo_batch = lambda generator, count: trajectories("partial", 0 if trainer.train_engine.dp_rank == 1 else count) + with pytest.raises(RuntimeError, match="incomplete fixed batch"): + trainer._collect_fixed_global_batch(None, trainer._source_batch_sample_count(), []) + finally: + dist.destroy_process_group() + + +def test_real_gloo_dp_sources_survive_tp_followers(tmp_path): + mp.spawn(_gloo_collection_worker, args=(str(tmp_path / "rendezvous"), str(tmp_path)), nprocs=4, join=True) + results = [torch.load(tmp_path / f"rank{rank}.pt", weights_only=False) for rank in range(4)] + for topology_index, (global_count, dp_size) in enumerate(((16, 2), (16, 4), (16, 1), (8, 4))): + topology = [result[topology_index] for result in results] + assert all(len(set(item["global_groups"])) == global_count // 4 for item in topology) + representatives = topology[::4 // dp_size] + assigned = [group for item in representatives for group in item["own_groups"]] + assert len(assigned) == global_count + assert set(assigned) == set(topology[0]["global_groups"]) + assert all(assigned.count(group) == 4 for group in set(assigned)) + assert results[0][0] == results[1][0] + assert results[2][0] == results[3][0] + assert results[0][2] == results[3][2] + assert results[2][3]["own_groups"] == results[3][3]["own_groups"] == [] + + +def test_incomplete_source_batch_fails_before_optimizer_update(tmp_path): + trainer = make_trainer(tmp_path) + trainer._collect_complete_grpo_batch = lambda generator, count: trajectories("partial", 12) + with pytest.raises(RuntimeError, match="incomplete fixed batch"): + trainer.train(None, dataset=[1]) + assert not trainer.train_engine.updated + + +def test_logical_sampler_cursor_survives_rollout_client_rebind(tmp_path): + trainer = make_trainer(tmp_path, dp_size=2) + trainer._rollout_source_world_size = 2 + trainer._rollout_source_rank = 0 + trainer._task_counter = 0 + dataset = [{"id": index} for index in range(12)] + stream = AsyncRLTrainer._create_data_generator(trainer, dataset, None) + first = [next(stream) for _ in range(4)] + trainer._reset_rollout_pipeline_after_reconfigure() + second = [next(stream) for _ in range(4)] + assert [task.task_id for task in first + second] == list(range(8)) + assert {task.group_id for task in first} == {"dp0:epoch0:sample0"} + assert {task.group_id for task in second} == {"dp0:epoch0:sample1"} + assert first[0].data != second[0].data + assert trainer._is_rollout_source() + assert trainer._source_batch_sample_count() == 8 + + +def test_rebind_preserves_admitted_results_and_incomplete_group_members(tmp_path): + trainer = make_trainer(tmp_path) + pending = {"group": trajectories("part", 3)} + trainer._pending_grpo_groups = pending + trainer.dispatcher = SimpleNamespace(reset_after_reconfigure=lambda: pytest.fail("valid work discarded")) + trainer._reset_rollout_pipeline_after_reconfigure() + assert trainer._pending_grpo_groups is pending + assert len(pending["group"]) == 3 diff --git a/tests/test_elastic_gradient_idle_path.py b/tests/test_elastic_gradient_idle_path.py index 417bac9..b1f6e25 100644 --- a/tests/test_elastic_gradient_idle_path.py +++ b/tests/test_elastic_gradient_idle_path.py @@ -1,8 +1,12 @@ from types import SimpleNamespace +import threading + +import torch from RL_Framework.engine.megatron_core_train_engine import MegatronCoreTrainEngine from RL_Framework.engine.megatron_train_engine import Megatron3DTrainEngine from RL_Framework.engine.train_engine import FSDPTrainEngine +from RL_Framework.infra.elastic.hybrid_pool import GradientPayload, InterReplicaGradientDomain class _UnexpectedReduce: @@ -10,14 +14,56 @@ def reduce_core_gradients(self, **_kwargs): raise AssertionError("elastic reduce must not run without a hybrid payload") -def test_mcore_skips_elastic_collective_without_hybrid_payload(): +def test_mcore_skips_elastic_collective_when_domain_is_not_configured(): engine = object.__new__(MegatronCoreTrainEngine) - engine.elastic_gradient_domain = _UnexpectedReduce() + engine.elastic_gradient_domain = None engine._pending_hybrid_gradients = [] engine._apply_elastic_inter_replica_gradients() +def test_mcore_dp1_reuses_main_gradient_for_weighted_hybrid_update(): + domain = InterReplicaGradientDomain(core_replica_ids=['dp0']) + domain.request_join('hybrid', 'dp0', membership_epoch=1) + domain.mark_active('hybrid', membership_epoch=1) + engine = object.__new__(MegatronCoreTrainEngine) + engine.elastic_gradient_domain = domain + engine._pending_hybrid_gradients = [GradientPayload( + 'hybrid', 'dp0', (torch.full((73,), 7.),), + step=3, state_version=2, membership_epoch=1, + sample_count=2, global_sample_count=4, + )] + engine._elastic_step_hybrid_workers = ('hybrid',) + engine._elastic_active_gradient_timeout_s = 0 + engine._hybrid_gradient_condition = threading.Condition() + engine._elastic_batch_partition = (2, 4, {'hybrid': 2}) + engine._elastic_training_step = 3 + engine.current_version = 2 + # This is the already weighted, finalized Core contribution. Passing it + # through an out-of-place divide-by-one would waste a full model buffer. + grad = torch.full((73,), 3.) + engine._model_main_gradients = lambda: [grad] + engine.get_data_parallel_rank = lambda: 0 + reduce = domain.reduce_core_gradients + + def checked_reduce(**kwargs): + assert kwargs['in_place'] is True + assert kwargs['core_gradients']['dp0'][0].data_ptr() == grad.data_ptr() + return reduce(**kwargs) + + domain.reduce_core_gradients = checked_reduce + acknowledgements = [] + engine._elastic_gradient_update_callback = acknowledgements.append + engine._apply_elastic_inter_replica_gradients() + assert torch.equal(grad, torch.full((73,), 6.5)) + assert not engine._pending_hybrid_gradients + assert len(acknowledgements) == 1 + assert acknowledgements[0].tensors[0].data_ptr() != grad.data_ptr() + assert acknowledgements[0].state_version == 2 + grad.zero_() + assert torch.equal(acknowledgements[0].tensors[0], torch.full((73,), 6.5)) + + def test_megatron3d_skips_elastic_collective_without_hybrid_payload(): engine = object.__new__(Megatron3DTrainEngine) engine.elastic_gradient_domain = _UnexpectedReduce() diff --git a/tests/test_fixed_batch_gradients.py b/tests/test_fixed_batch_gradients.py new file mode 100644 index 0000000..0eb3246 --- /dev/null +++ b/tests/test_fixed_batch_gradients.py @@ -0,0 +1,269 @@ +"""Fixed-batch gradient weighting, delivery identity and CPU reduction.""" +from dataclasses import replace +import threading +import pytest +import torch +import torch.distributed as dist +import torch.multiprocessing as mp +from RL_Framework.engine.megatron_core_train_engine import MegatronCoreTrainEngine +from RL_Framework.infra.elastic.hybrid_pool import GradientPayload, InterReplicaGradientDomain, ReplicaRole +from RL_Framework.infra.elastic.gradient_ipc import ElasticGradientClient, ElasticGradientServer + +def _domain(): + domain = InterReplicaGradientDomain(core_replica_ids=['dp0', 'dp1']) + domain.request_join('hybrid', 'dp0', membership_epoch=7) + domain.mark_active('hybrid', membership_epoch=7) + return domain + + + +def _payload(**changes): + return replace(GradientPayload('hybrid', 'dp0', (torch.tensor([9.]),), + step=3, state_version=2, membership_epoch=7, + sample_count=2, global_sample_count=8), **changes) + + + +def test_uneven_partitions_match_full_batch_gradient_and_adam_update(): + # Core means 1 (four samples) and 5 (two); Hybrid mean 9 (two). + result = _domain().reduce_core_gradients( + core_gradients={'dp0': (torch.tensor([1.]),), 'dp1': (torch.tensor([5.]),)}, + core_sample_counts={'dp0': 4, 'dp1': 2}, global_sample_count=8, + hybrid_payloads=[_payload()], step=3, state_version=2, + ) + expected = torch.tensor([(4 + 10 + 18) / 8]) + for gradients in result.values(): + torch.testing.assert_close(gradients[0], expected) + full = torch.nn.Parameter(torch.tensor([2.])) + split = torch.nn.Parameter(full.detach().clone()) + reference = torch.optim.Adam([full], lr=.1) + elastic = torch.optim.Adam([split], lr=.1) + full.grad, split.grad = expected.clone(), result['dp0'][0].clone() + reference.step() + elastic.step() + torch.testing.assert_close(full, split) + torch.testing.assert_close(reference.state[full]['exp_avg'], elastic.state[split]['exp_avg']) + + + + + + + + + +def _gloo_worker(rank, rendezvous, output, in_place=False): + dist.init_process_group('gloo', init_method=f'file://{rendezvous}', rank=rank, world_size=2) + isolated = dist.new_group([0, 1], backend='gloo') + domain = _domain() + domain.bind_hybrid_process_group(isolated) + result = domain.reduce_core_gradients( + core_gradients={f'dp{rank}': (torch.tensor([1. if rank == 0 else 5.]),)}, + core_sample_counts={f'dp{rank}': 4 if rank == 0 else 2}, + global_sample_count=8, hybrid_payloads=[_payload()] if rank == 0 else [], + step=3, state_version=2, + in_place=in_place, + ) + output[rank] = float(result[f'dp{rank}'][0].item()) + dist.destroy_process_group() + + + +@pytest.mark.parametrize('in_place', [False, True]) +def test_real_gloo_lane_without_hybrid_still_participates(tmp_path, in_place): + context = mp.get_context('spawn') + with context.Manager() as manager: + results = manager.dict() + processes = [context.Process(target=_gloo_worker, args=(rank, str(tmp_path / 'rendezvous'), results, in_place)) for rank in range(2)] + for process in processes: + process.start() + for process in processes: + process.join(30) + try: + assert [process.exitcode for process in processes] == [0, 0] + assert dict(results) == {0: 4., 1: 4.} + finally: + for process in processes: + if process.is_alive(): + process.terminate() + process.join() + + + + + + + + + + + + + + + +def test_tcp_preserves_batch_weighting_metadata(): + received = [] + server = ElasticGradientServer(on_payload=received.append) + endpoint = server.start() + try: + ElasticGradientClient(endpoint).send(_payload()) + assert received[0].sample_count == 2 + assert received[0].global_sample_count == 8 + assert received[0].state_version == 2 + finally: + server.close() + + + + + + +def test_empty_core_lane_enters_elastic_reduce_without_pending_or_expected_hybrids(): + engine = MegatronCoreTrainEngine.__new__(MegatronCoreTrainEngine) + engine.elastic_gradient_domain = _domain() + engine._pending_hybrid_gradients = [] + engine._elastic_step_hybrid_workers = () + engine._elastic_active_gradient_timeout_s = 0 + engine._hybrid_gradient_condition = threading.Condition() + engine._elastic_batch_partition = (0, 8, {'hybrid': 4}) + engine._elastic_training_step = 3 + engine._elastic_expected_state_version = 2 + engine._elastic_gradient_update_callback = None + grad = torch.zeros(1) + engine._model_main_gradients = lambda: [grad] + engine.get_data_parallel_rank = lambda: 1 + entered = [] + def reduce(**kwargs): + entered.append(kwargs) + return {'dp1': (torch.tensor([7.]),)} + engine.elastic_gradient_domain.reduce_core_gradients = reduce + engine._apply_elastic_inter_replica_gradients() + assert len(entered) == 1 + assert entered[0]['in_place'] is True + assert entered[0]['core_gradients']['dp1'][0].data_ptr() == grad.data_ptr() + torch.testing.assert_close(grad, torch.tensor([7.])) + + + +def test_native_dp_without_hybrid_domain_keeps_loss_scale_one(): + engine = MegatronCoreTrainEngine.__new__(MegatronCoreTrainEngine) + engine.elastic_gradient_domain = None + engine.get_data_parallel_world_size = lambda: 2 + engine.set_elastic_batch_partition(8, 16, {}) + assert engine._elastic_core_loss_scale() == 1. + + + +class _GradientModel(torch.nn.Module): + def __init__(self): + super().__init__() + self.weight = torch.nn.Parameter(torch.tensor([2.])) + self.weight.main_grad = torch.zeros(1) + self.weight.register_hook(self._accumulate) + + def _accumulate(self, gradient): + self.weight.main_grad.add_(gradient) + return gradient + + def zero_grad_buffer(self): + self.weight.main_grad.zero_() + + + +class _AdamAdapter: + def __init__(self, model): + self.model = model + self.inner = torch.optim.Adam(model.parameters(), lr=.01) + + def zero_grad(self): + self.inner.zero_grad() + + def scale_loss(self, loss): + return loss + + def step(self): + self.model.weight.grad = self.model.weight.main_grad.clone() + self.inner.step() + return True, self.model.weight.main_grad.norm(), 0 + + + +def _gloo_engine_worker(rank, rendezvous, output): + import sys + import types + dist.init_process_group('gloo', init_method=f'file://{rendezvous}', rank=rank, world_size=2) + isolated = dist.new_group([0, 1], backend='gloo') + def finalize(models): + for model in models: + dist.all_reduce(model.weight.main_grad) + model.weight.main_grad.div_(2) + sys.modules['megatron.core.distributed'] = types.SimpleNamespace(finalize_model_grads=finalize) + engine = MegatronCoreTrainEngine.__new__(MegatronCoreTrainEngine) + model = _GradientModel() + engine.model = [model] + engine.optimizer = _AdamAdapter(model) + engine.optimizer_revision = 0 + engine.current_version = 0 + engine.elastic_gradient_domain = _domain() + engine.elastic_gradient_domain.bind_hybrid_process_group(isolated) + engine.get_data_parallel_world_size = lambda: 2 + engine.get_data_parallel_rank = lambda: rank + engine._hybrid_gradient_condition = threading.Condition() + engine._elastic_active_gradient_timeout_s = 10 + engine._elastic_gradient_update_callback = None + engine._elastic_seen_payloads = set() + engine._pending_hybrid_gradients = [] + engine.set_elastic_training_step(0, ['hybrid'] if rank == 0 else []) + engine.set_elastic_batch_partition(4 if rank == 0 else 0, 8, {'hybrid': 4}) + if rank == 0: + engine.enqueue_hybrid_gradient_payload(_payload(tensors=(torch.tensor([-7.]),), step=0, + state_version=0, sample_count=4)) + engine._iter_micro_batches = lambda rows: ({'trajectories': rows[start:start + 3]} for start in range(0, len(rows), 3)) + engine._forward = lambda batch: (model.weight - torch.tensor(batch['trajectories'])).square().mean() + engine._loss = lambda batch, loss: (loss, {'loss': float(loss.detach()), 'pg_loss': float(loss.detach()), + 'kl': 0., 'reward_sum': 0., 'n_samples': len(batch['trajectories']), 'n_updates': 1}) + engine._reduce_stats = lambda stats: stats + engine.grpo_update([0., 1., 2., 3.] if rank == 0 else []) + optimizer_state = engine.optimizer.inner.state[model.weight] + # Simulate a real RS buffer: only this rank's shard contains finalized + # gradients; the other shard still has stale local values. + bucket_data = torch.tensor([1., 2., 999., 999.]) if rank == 0 else torch.tensor([-999., -999., 3., 4.]) + model.bucket_groups = [types.SimpleNamespace( + intra_distributed_optimizer_instance_group=dist.group.WORLD, + buckets=[types.SimpleNamespace(grad_data=bucket_data)], + )] + engine.use_distributed_optimizer = True + engine._materialize_elastic_core_gradients() + output[rank] = (float(model.weight.detach()), float(optimizer_state['exp_avg']), + float(optimizer_state['exp_avg_sq']), engine.optimizer_revision, bucket_data.tolist()) + dist.destroy_process_group() + + + +def test_full_engine_gloo_prescale_finalize_elastic_sum_matches_full_batch_with_empty_core(tmp_path): + context = mp.get_context('spawn') + with context.Manager() as manager: + results = manager.dict() + processes = [context.Process(target=_gloo_engine_worker, args=(rank, str(tmp_path / 'rendezvous'), results)) for rank in range(2)] + for process in processes: + process.start() + for process in processes: + process.join(30) + try: + assert [process.exitcode for process in processes] == [0, 0] + expected = torch.nn.Parameter(torch.tensor([2.])) + optimizer = torch.optim.Adam([expected], lr=.01) + (expected - torch.arange(8)).square().mean().backward() + optimizer.step() + for result in results.values(): + assert result[0] == pytest.approx(float(expected)) + assert result[1] == pytest.approx(float(optimizer.state[expected]['exp_avg'])) + assert result[2] == pytest.approx(float(optimizer.state[expected]['exp_avg_sq'])) + assert result[3] == 1 + assert result[4] == [1., 2., 3., 4.] + finally: + for process in processes: + if process.is_alive(): + process.terminate() + process.join() diff --git a/tests/test_grpo_sample_mean.py b/tests/test_grpo_sample_mean.py new file mode 100644 index 0000000..537ea81 --- /dev/null +++ b/tests/test_grpo_sample_mean.py @@ -0,0 +1,95 @@ +"""Trajectory weights remain fixed across uneven elastic microbatches.""" +from types import SimpleNamespace + +import pytest +import torch + +from RL_Framework.engine import megatron_core_train_engine as engine_module +from RL_Framework.engine.megatron_core_train_engine import MegatronCoreTrainEngine +from RL_Framework.infra.elastic.batch_partition import partition_complete_grpo_groups +from RL_Framework.trainer.async_rl_trainer import AsyncRLTrainer + + +def _engine(): + engine = MegatronCoreTrainEngine.__new__(MegatronCoreTrainEngine) + engine.train_tp_size = 1 + engine.clip_epsilon, engine.kl_coef = .2, .1 + return engine + + +def test_b12_n4_mb3_gradient_matches_each_trajectory_reference_with_or_without_hybrid(): + rows = [dict(grpo_group_id=str(i // 4), rollout_index=i % 4, index=i, + rewards=torch.tensor([float(i == 8)])) for i in range(12)] + AsyncRLTrainer.__new__(AsyncRLTrainer)._compute_advantages(rows) + engine = _engine() + + def gradient(participants, microbatch_size): + parameter = torch.tensor(.1, requires_grad=True) + assignments = partition_complete_grpo_groups(rows, 4, participants) + for subset in assignments.values(): + for start in range(0, len(subset), microbatch_size): + part = subset[start:start + microbatch_size] + count = len(part) + mask = torch.zeros(count, 11) + mask[:, 1] = 1 + for index, row in enumerate(part): + if row['index'] in (6, 8): + mask[index, 1:] = 1 + # Distinct row features keep the reference gradient nonzero. + features = torch.tensor([1 + row['index'] / 20 for row in part])[:, None] + logits = torch.stack([(parameter * features).expand(count, 11), + torch.zeros(count, 11)], -1) + batch = dict(input_ids=torch.zeros(count, 11, dtype=torch.long), + old_logprobs=torch.full((count, 11), -torch.log(torch.tensor(2.))), + loss_mask=mask, advantages=torch.cat([row['advantages'] for row in part]), + rewards=torch.cat([row['rewards'] for row in part])) + loss, _ = engine._loss(batch, logits) + (loss * count / len(rows)).backward() + return parameter.grad + + reference = gradient(['dp0'], 1) + assert reference.abs() > 1e-4 + torch.testing.assert_close(gradient(['dp0'], 3), reference) + torch.testing.assert_close(gradient(['dp0', 'hybrid'], 3), reference) + + +def test_masked_sample_losses_and_statistics_match_single_trajectory_reference(monkeypatch): + engine = _engine() + values = torch.tensor([.1, -.3, .8], requires_grad=True) + logits = torch.stack([values[:, None].expand(3, 4), torch.zeros(3, 4)], -1) + batch = dict(input_ids=torch.zeros(3, 4, dtype=torch.long), + old_logprobs=torch.full((3, 4), -.7), + loss_mask=torch.tensor([[0., 1., 0., 0.], [0., 1., 1., 1.], [0., 0., 0., 0.]]), + advantages=torch.tensor([.7, -.4, 1.]), rewards=torch.tensor([1., 2., 3.])) + + def subset(start, stop): + return {key: value[start:stop] for key, value in batch.items()} + + singles = [engine._loss(subset(i, i + 1), logits[i:i + 1]) for i in range(3)] + reference_loss = torch.stack([loss for loss, _ in singles]).mean() + actual_loss, actual_stats = engine._loss(batch, logits) + torch.testing.assert_close(actual_loss, reference_loss) + actual_gradient, = torch.autograd.grad(actual_loss, values, retain_graph=True) + reference_gradient, = torch.autograd.grad(reference_loss, values, retain_graph=True) + torch.testing.assert_close(actual_gradient, reference_gradient) + assert actual_gradient[2] == 0 + assert singles[2][0] == 0 and singles[2][1]['n_samples'] == 1 + + left = engine._loss(subset(0, 2), logits[:2])[1] + right = singles[2][1] + merged = engine._merge_stats([left, right]) + for key in ('loss', 'pg_loss', 'kl'): + assert merged[key] == pytest.approx(actual_stats[key]) + assert merged['n_samples'] == 3 and merged['n_updates'] == 2 + + # Exercise the real DP-stat encoder and decoder with a captured peer vector. + engine._device = lambda: 'cpu' + engine._parallel_state = lambda: SimpleNamespace(get_data_parallel_group=lambda **_: None) + peer = [] + monkeypatch.setattr(engine_module.dist, 'all_reduce', lambda values, **_: peer.append(values.clone())) + engine._reduce_stats(right) + monkeypatch.setattr(engine_module.dist, 'all_reduce', lambda values, **_: values.add_(peer[0])) + reduced = engine._reduce_stats(left) + for key in ('loss', 'pg_loss', 'kl'): + assert reduced[key] == pytest.approx(actual_stats[key]) + assert reduced['n_samples'] == 3 and reduced['n_updates'] == 2 diff --git a/tests/test_hybrid_sample_statistics_data.py b/tests/test_hybrid_sample_statistics_data.py new file mode 100644 index 0000000..a222739 --- /dev/null +++ b/tests/test_hybrid_sample_statistics_data.py @@ -0,0 +1,53 @@ +"""Hybrid loss statistics are sample means carried by the existing transports.""" +import shutil + +import pytest + +from hybrid_statistics_helpers import compute, rows +from RL_Framework.infra.elastic.gradient_ipc import ElasticGradientClient, ElasticGradientServer +from RL_Framework.infra.elastic.native_rdma import NativeRDMAGradientClient, NativeRDMAGradientServer + + +@pytest.mark.parametrize('zero_mask',[False,True]) +def test_payload_statistics_match_single_sample_means(monkeypatch,zero_mask): + samples=rows(8,4,zero_mask=zero_mask) + worker,payload=compute(monkeypatch,samples) + per_sample=[] + for row in samples: + batch=worker._pad_micro_batch([row]) + per_sample.append(worker._loss(batch,worker._forward(batch))[1]) + reference=worker._merge_stats(per_sample) + for key in ('loss','pg_loss','kl','reward_sum','reward_mean','n_samples'): + assert payload.stats[key]==pytest.approx(reference[key],abs=1e-6) + assert payload.stats['n_samples']==4 and payload.stats['n_updates']==2 + + +def test_tcp_transports_statistics_from_real_hybrid_compute(monkeypatch): + _,payload=compute(monkeypatch,rows(8,4)) + assert hasattr(payload,'stats'), 'Hybrid compute discarded the loss statistics' + received=[] + server=ElasticGradientServer(on_payload=received.append) + endpoint=server.start() + try: + ElasticGradientClient(endpoint).send(payload) + assert received[0].stats==payload.stats + finally: + server.close() + + +def test_native_metadata_transports_statistics_without_requiring_rdma_hardware(tmp_path,monkeypatch): + _,payload=compute(monkeypatch,rows(8,4)) + assert hasattr(payload,'stats'), 'Hybrid compute discarded the loss statistics' + received=[] + server=NativeRDMAGradientServer(on_payload=received.append,work_dir=tmp_path/'server') + staged=tmp_path/'wire.pt' + client=NativeRDMAGradientClient(server.endpoint) + client.transfer.send_file_once=lambda input_path,**kw:shutil.copyfile(input_path,staged) + client.send(payload) + def transfer(output_path,**kw): + shutil.copyfile(staged,output_path) + server._stop.set() + server.transfer.receive_file_once=transfer + server._serve() + assert server.last_error is None + assert received[0].stats==payload.stats diff --git a/tests/test_hybrid_worker_batch_protocol.py b/tests/test_hybrid_worker_batch_protocol.py new file mode 100644 index 0000000..b07ad6c --- /dev/null +++ b/tests/test_hybrid_worker_batch_protocol.py @@ -0,0 +1,61 @@ +"""The worker forwards fixed-batch counts through its actual task loop.""" +import importlib.util +from pathlib import Path +from types import SimpleNamespace + +import torch + +from RL_Framework.infra.elastic.hybrid_pool import GradientPayload + + +def _worker(): + path = Path(__file__).parents[1] / 'scripts' / 'elastic_hybrid_worker.py' + spec = importlib.util.spec_from_file_location('hybrid_batch_worker', path) + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + return module + + +def test_payload_file_retains_local_and_global_sample_counts(tmp_path): + worker = _worker() + path = tmp_path / 'payload.pt' + torch.save(dict(tensors=[torch.ones(1)], sample_count=4, global_sample_count=12), path) + args = SimpleNamespace(replica_id='h', replica_rank=0, replica_world_size=1) + payload = worker._load_payload_file(path, worker_id='h', target_core_id='dp0', args=args) + assert payload.sample_count == 4 and payload.global_sample_count == 12 + + +def test_training_task_passes_global_count_to_engine_and_gradient_client(tmp_path, monkeypatch): + worker = _worker() + snapshot = tmp_path / 'snapshot.pt' + snapshot.touch() + args = SimpleNamespace(snapshot_path=str(snapshot), initialize_only=False, + worker_mode='megatron_core', gradient_host='localhost', gradient_port=1, + authkey='', gradient_transport='tcp', payload_file='', idle=False, + task_dir=str(tmp_path), worker_id='h', replica_id='h', target_core_id='dp0') + received, sent = [], [] + + def compute(trajectories, **metadata): + received.append(metadata) + return GradientPayload('h', 'dp0', (torch.ones(1),), + step=metadata['step'], state_version=metadata['state_version'], + membership_epoch=metadata['membership_epoch'], sample_count=len(trajectories), + global_sample_count=metadata['global_sample_count']) + + engine = SimpleNamespace(recompute_logprobs=lambda _: None, + compute_elastic_gradient_payload=compute) + monkeypatch.setattr(worker, 'parse_args', lambda: args) + monkeypatch.setattr(worker, '_initialize_megatron_worker', lambda _: engine) + monkeypatch.setattr(worker, '_endpoint_for_lane', lambda *_: {}) + monkeypatch.setattr(worker, '_gradient_client', lambda *_: + SimpleNamespace(send=lambda payload, **_: sent.append(payload))) + monkeypatch.setattr(worker.torch.distributed, 'is_initialized', lambda: False) + monkeypatch.setenv('RANK', '0') + torch.save(dict(trajectories=[{}] * 4, step=3, state_version=2, + membership_epoch=7, sample_count=4, global_sample_count=12), tmp_path / 'h.0.pt') + torch.save(dict(shutdown=True), tmp_path / 'h.1.pt') + worker.main() + assert received[0]['global_sample_count'] == 12 + assert received[0]['state_version'] == 2 + assert len(sent) == 1 and sent[0].sample_count == 4 and sent[0].global_sample_count == 12 + assert (tmp_path / 'h.0.pt.done').exists() diff --git a/tests/test_megatron_core_train_engine.py b/tests/test_megatron_core_train_engine.py index ef164f8..50fbf1d 100644 --- a/tests/test_megatron_core_train_engine.py +++ b/tests/test_megatron_core_train_engine.py @@ -8,7 +8,7 @@ ) -def test_merge_stats_uses_update_weighted_averages(): +def test_merge_stats_weights_losses_by_samples_and_grad_norm_by_updates(): merged = MegatronCoreTrainEngine._merge_stats( [ { @@ -32,9 +32,9 @@ def test_merge_stats_uses_update_weighted_averages(): ] ) - assert merged["loss"] == pytest.approx(2.5) - assert merged["pg_loss"] == pytest.approx(3.5) - assert merged["kl"] == pytest.approx(0.25) + assert merged["loss"] == pytest.approx(2.0) + assert merged["pg_loss"] == pytest.approx(3.0) + assert merged["kl"] == pytest.approx(0.2) assert merged["grad_norm"] == pytest.approx(4.5) assert merged["reward_mean"] == pytest.approx(1.5) diff --git a/trainer/async_rl_trainer.py b/trainer/async_rl_trainer.py index 45f267f..351bd11 100644 --- a/trainer/async_rl_trainer.py +++ b/trainer/async_rl_trainer.py @@ -26,6 +26,9 @@ from transformers import AutoTokenizer from RL_Framework.config import AsyncRLConfig, HeterogeneousInstanceConfig +from RL_Framework.infra.elastic.batch_partition import ( + complete_group_sample_counts, partition_complete_grpo_groups, +) from RL_Framework.engine.rollout_engine import MockRolloutEngine, MultiInstanceRolloutEngine from RL_Framework.engine.heterogeneous_engine import HeterogeneousRolloutEngine from RL_Framework.engine.train_engine import TrainEngine @@ -253,6 +256,7 @@ def setup(self, workflow): self.train_engine = create_train_engine(self.config) self.train_engine.initialize(max_seq_length=self.config.max_seq_length) self._run_grp_memory_preflight() + self._initialize_rollout_source_layout() if self.start_step > 0 or self.initial_model_version > 0: self.train_engine.set_version(self.initial_model_version) if self.is_main_process: @@ -480,8 +484,8 @@ def _create_data_generator(self, dataset, workflow): sampler = DistributedSampler( dataset, - num_replicas=self.train_engine.get_data_parallel_world_size(), - rank=self.train_engine.get_data_parallel_rank(), + num_replicas=int(getattr(self, "_rollout_source_world_size", self.train_engine.get_data_parallel_world_size())), + rank=int(getattr(self, "_rollout_source_rank", self.train_engine.get_data_parallel_rank()) or 0), shuffle=True, ) @@ -501,7 +505,7 @@ def _create_data_generator(self, dataset, workflow): for sample_index, batch in enumerate(dataloader): for item in batch: group_id = ( - f"dp{self.train_engine.get_data_parallel_rank()}:" + f"dp{getattr(self, '_rollout_source_rank', self.train_engine.get_data_parallel_rank())}:" f"epoch{epoch}:sample{sample_index}" ) for rollout_index in range(n_samples): @@ -520,81 +524,96 @@ def _create_data_generator(self, dataset, workflow): ) def _collect_complete_grpo_batch(self, data_gen, batch_size: int): - """Collect complete prompt groups even when individual rollouts are rejected.""" + """Consume ready complete groups before requesting missing members.""" + from RL_Framework.infra.elastic.batch_partition import unique_group_members + n_samples = int(self.config.n_samples) + if batch_size < 0 or n_samples < 1 or batch_size % n_samples: + raise ValueError("requested batch must contain complete GRPO groups") selected: list[dict[str, Any]] = [] + consumed = getattr(self, "_consumed_grpo_group_ids", None) + if consumed is None: + consumed = self._consumed_grpo_group_ids = set() planner_cfg = getattr(self.config, "global_resource_planner", None) - collection_timeout = float( - getattr(planner_cfg, "runtime_batch_collection_timeout_s", 0.0) - if planner_cfg is not None - else 0.0 - ) - max_retries = int( - getattr(planner_cfg, "runtime_batch_collection_max_retries", 0) - if planner_cfg is not None - else 0 - ) + collection_timeout = float(getattr(planner_cfg, "runtime_batch_collection_timeout_s", 0.0)) + max_retries = int(getattr(planner_cfg, "runtime_batch_collection_max_retries", 0)) retries = 0 is_main_process = getattr(self, "is_main_process", True) while len(selected) < batch_size: + # This must precede every blocking dispatcher call: after EHP + # borrows all engines a complete group may already be assembled. + for group_id in list(self._pending_grpo_groups): + if str(group_id) in consumed: + del self._pending_grpo_groups[group_id] + continue + members = unique_group_members(self._pending_grpo_groups[group_id], n_samples) + self._pending_grpo_groups[group_id] = members + if len(members) < n_samples: + continue + selected.extend(members[:n_samples]) + consumed.add(str(group_id)) + del self._pending_grpo_groups[group_id] + if len(selected) == batch_size: + return selected if not is_main_process: self._follow_applied_runtime_reconfiguration_if_available( self._runtime_reconfiguration_coord_dir() ) try: - raw_batch = self.dispatcher.active_submit_and_wait( - input_generator=data_gen, - batch_size=batch_size, - timeout=collection_timeout if collection_timeout > 0 else None, - ) + peek = getattr(self.dispatcher, "peek_ready_data", None) + ready = peek() if callable(peek) else [] + paused = callable(getattr(self.dispatcher, "is_paused", None)) and self.dispatcher.is_paused() + if ready or paused: + # Existing results can be consumed even while production is + # paused. Never enqueue new work while every engine is lent. + count = min(len(ready), batch_size - len(selected)) if ready else 1 + results = self.dispatcher.wait_results( + count=count, + timeout=collection_timeout if collection_timeout > 0 else None, + ) + raw_batch = [result.data for result in results if result.data is not None] + else: + pending_members = sum(len(unique_group_members(members, n_samples)) + for members in self._pending_grpo_groups.values()) + needed = max(1, batch_size - len(selected) - pending_members) + raw_batch = self.dispatcher.active_submit_and_wait( + input_generator=data_gen, batch_size=needed, + timeout=collection_timeout if collection_timeout > 0 else None, + ) except TimeoutError as exc: retries += 1 if is_main_process: - print( - "[BatchCollection] timeout " - f"retry={retries}/{max_retries} " - f"error={exc}" - ) + print(f"[BatchCollection] timeout retry={retries}/{max_retries} error={exc}") self._reset_rollout_pipeline_after_reconfigure() if retries > max_retries: - # Do not strand non-source ranks at the next collective - # when long generations leave the queue short. Return - # complete groups collected so far; the caller will - # distribute the same partial batch to every rank. - if is_main_process: - print( - "[BatchCollection] returning partial batch after " - f"{retries} timeout(s): {len(selected)} trajectories" - ) + # The fixed-B caller propagates this shortfall to all ranks + # and rejects an optimizer update with a partial batch. return selected continue if self.staleness_manager is not None: + # Timeout notifications already released running credit through + # rejection. Every successful result releases accepted credit, + # including members later discarded by group filtering. accepted_count = sum( isinstance(trajectory, dict) and trajectory.get("status") != "failed" for trajectory in raw_batch ) self.staleness_manager.on_batch_consumed(accepted_count) for trajectory in raw_batch: - if trajectory.get("status") == "failed": + if not isinstance(trajectory, dict) or trajectory.get("status") == "failed": continue group_id = str(trajectory["grpo_group_id"]) - self._pending_grpo_groups.setdefault(group_id, []).append(trajectory) - - for group_id in list(self._pending_grpo_groups): - members = self._pending_grpo_groups[group_id] - if len(members) < n_samples: - continue - remaining = batch_size - len(selected) - if remaining < n_samples: - break - selected.extend(members[:n_samples]) - del self._pending_grpo_groups[group_id] - - # Rejected tasks can leave incomplete groups forever. Bound their memory. + if group_id not in consumed: + self._pending_grpo_groups.setdefault(group_id, []).append(trajectory) + # Complete admitted groups remain available for future fixed batches. + # Bound abandoned incomplete groups only, preserving their order. max_pending_groups = max(64, batch_size * 8) - while len(self._pending_grpo_groups) > max_pending_groups: - oldest_group = next(iter(self._pending_grpo_groups)) - del self._pending_grpo_groups[oldest_group] + if len(self._pending_grpo_groups) > max_pending_groups: + for group_id in list(self._pending_grpo_groups): + if len(self._pending_grpo_groups) <= max_pending_groups: + break + if len(unique_group_members(self._pending_grpo_groups[group_id], n_samples)) < n_samples: + del self._pending_grpo_groups[group_id] return selected def _compute_advantages(self, trajectories): @@ -796,12 +815,7 @@ def train(self, workflow, dataset=None): data_gen = self._create_data_generator(dataset, workflow) - local_batch_size = self.train_engine.get_local_batch_size(self.config.batch_size) - if local_batch_size % int(self.config.n_samples) != 0: - raise ValueError( - "Each data-parallel replica's local batch size must be divisible by n_samples: " - f"local_batch_size={local_batch_size}, n_samples={self.config.n_samples}" - ) + local_batch_size = self._source_batch_sample_count() if self.is_main_process: print("\nStarting asynchronous pipeline training") @@ -827,56 +841,24 @@ def train(self, workflow, dataset=None): # opening the next step to peers. if self.runtime_elastic_executor is not None: self.runtime_elastic_executor.hybrid_runtime_state() - active_hybrid_ids = self._prepare_elastic_training_step(step) self._runtime_reconfiguration_boundary(step, "pre") else: self._trace_train_phase(step, "follow_reconfig_start") self._runtime_reconfiguration_boundary(step, "pre") self._trace_train_phase(step, "follow_reconfig_done") - active_hybrid_ids = self._prepare_elastic_training_step(step) + active_hybrid_ids = self._prepare_elastic_training_step(step) + # Reconfiguration can change the number of DP data sources. + local_batch_size = self._source_batch_sample_count() if self._use_heterogeneous and hasattr(self.rollout_engine, "notify_epoch_start"): self.rollout_engine.notify_epoch_start(epoch=step) - batch = None - if self.train_engine.is_batch_source(): - self._trace_train_phase(step, "collect_batch_start") - batch = self._collect_complete_grpo_batch( - data_gen=data_gen, - batch_size=local_batch_size * (1 + len(active_hybrid_ids)), - ) - self._trace_train_phase( - step, - "collect_batch_done", - batch_size=len(batch), - ) - self._trace_train_phase(step, "distribute_start") - # The Megatron helper broadcasts only within each tensor-parallel - # group. With data parallelism, the other TP group's source rank - # would receive None and enter the next collective early. Use one - # explicit world-wide batch broadcast so every rank observes the - # same batch decision and payload before any training collective. - if torch.distributed.is_available() and torch.distributed.is_initialized(): - payload = [batch if self.rank == 0 else None] - torch.distributed.broadcast_object_list(payload, src=0) - batch = payload[0] or [] - else: - batch = self.train_engine.distribute_trajectories(batch) - self._trace_train_phase( - step, - "distribute_done", - batch_size=len(batch) if batch is not None else 0, + self._trace_train_phase(step, "collect_batch_start") + batch, participants = self._collect_fixed_global_batch( + data_gen, local_batch_size, active_hybrid_ids, ) - self._trace_train_phase(step, "rollout_rank_wait_start") - self._wait_for_all_ranks_after_rollout(step, batch) - self._trace_train_phase(step, "rollout_rank_wait_done") - - if not batch: - if self.is_main_process: - print(f"Step {step}: received no valid trajectories; skipping") - self._runtime_reconfiguration_boundary(step, "post") - continue + self._trace_train_phase(step, "collect_batch_done", batch_size=len(batch)) align_trajectories = getattr( self.train_engine, @@ -894,15 +876,8 @@ def train(self, workflow, dataset=None): adv_start = time.time() self._trace_train_phase(step, "advantage_start") batch = self._compute_advantages(batch) - if active_hybrid_ids and self.is_main_process: - core_count = local_batch_size - for index, worker_id in enumerate(active_hybrid_ids): - start = core_count + index * local_batch_size - self._pending_hybrid_training_batches[worker_id] = list( - batch[start : start + local_batch_size] - ) - batch = batch[:core_count] - self._dispatch_hybrid_training_batches(step) + global_batch = batch + batch = self._assign_fixed_global_batch(step, global_batch, participants) advantage_time = time.time() - adv_start self._trace_train_phase(step, "advantage_done") @@ -929,7 +904,7 @@ def train(self, workflow, dataset=None): reward_start = time.time() self._trace_train_phase(step, "reward_compute_start") reward_values = torch.cat( - [traj["rewards"].reshape(-1).float() for traj in batch] + [traj["rewards"].reshape(-1).float() for traj in global_batch] ) stats.update( { @@ -1026,7 +1001,8 @@ def train(self, workflow, dataset=None): step_total_time = time.time() - step_start - stats["n_trajectories"] = len(batch) + stats["n_trajectories"] = len(global_batch) + stats["local_n_trajectories"] = len(batch) stats["rollout_time"] = rollout_time stats["train_time"] = train_time stats["weight_sync_time"] = weight_sync_time @@ -1044,12 +1020,12 @@ def train(self, workflow, dataset=None): if self.is_main_process: - self._record_history_step(step, batch, stats, + self._record_history_step(step, global_batch, stats, rollout_time, train_time, weight_sync_time, advantage_time, recompute_time, step_total_time) self._trace_train_phase(step, "resource_planner_start") - self._run_global_resource_planner_at_boundary(step, batch, stats) + self._run_global_resource_planner_at_boundary(step, global_batch, stats) self._trace_train_phase(step, "resource_planner_done") else: self._trace_train_phase(step, "follow_reconfig_after_step_start") @@ -1080,7 +1056,7 @@ def train(self, workflow, dataset=None): print(f" PG Loss: {stats.get('pg_loss', 0):.4f}") print(f" KL: {stats.get('kl', 0):.4f}") print(f" Reward: {stats.get('reward_mean', 0):.4f}") - print(f" Trajectories: {len(batch)}") + print(f" Trajectories: {len(global_batch)}") print(f" Version: {self.train_engine.get_version()}") print( " Pipeline: " @@ -2801,11 +2777,11 @@ def _dispatch_hybrid_training_batches(self, step: int) -> None: for worker_id, trajectories in self._pending_hybrid_training_batches.items(): membership_path = task_dir / "membership" / f"{worker_id}.json" if not membership_path.exists(): - continue + raise RuntimeError(f"active hybrid has no membership record: {worker_id}") state = json.loads(membership_path.read_text(encoding="utf-8")) if (state.get("run_id") != self._elastic_membership_run_id or str(state.get("role")) != "hybrid_training"): - continue + raise RuntimeError(f"active hybrid membership changed before dispatch: {worker_id}") task_path = task_dir / f"{worker_id}.step_{step}.pt" tmp_path = task_path.with_suffix(task_path.suffix + ".tmp") torch.save( @@ -2814,6 +2790,8 @@ def _dispatch_hybrid_training_batches(self, step: int) -> None: "step": int(step), "state_version": int(self.train_engine.get_version()), "membership_epoch": int(state.get("membership_epoch", 0)), + "sample_count": len(trajectories), + "global_sample_count": int(self.config.batch_size), "snapshot_path": "" if lockstep else str(self.train_engine.get_elastic_state_snapshot_path(int(self.train_engine.get_version()))), }, tmp_path, @@ -3156,12 +3134,14 @@ def _apply_global_resource_planner_decision(self, step: int, decision, stats: di ) def _reset_rollout_pipeline_after_reconfigure(self) -> None: - if self.dispatcher is not None and hasattr( - self.dispatcher, - "reset_after_reconfigure", - ): - self.dispatcher.reset_after_reconfigure() - self._pending_grpo_groups.clear() + """Preserve admitted work when a client or compute topology is rebound. + + Pending input tasks resolve the rollout client at execution time. Existing + result queues and partially assembled complete groups contain valid + sampled trajectories, so clearing them would silently drop work and + restart the logical input stream unevenly between configurations. + """ + return def _pause_rollout_pipeline(self) -> None: if self.dispatcher is not None and hasattr(self.dispatcher, "pause"): @@ -3356,3 +3336,144 @@ def _record_history_step( resource_config=self._resource_config_snapshot, control_plane=control_plane, ) + + def _initialize_rollout_source_layout(self) -> None: + """Register one stable data producer per initialized Core DP replica. + + Keep source identities and admitted groups across rollout client + rebindings and elastic Hybrid membership changes. + """ + source = None + if self.train_engine.is_batch_source(): + source = (int(self.train_engine.get_data_parallel_rank()), int(self.rank)) + if dist.is_available() and dist.is_initialized(): + records = [None] * dist.get_world_size() + dist.all_gather_object(records, source) + else: + records = [source] + sources = sorted(item for item in records if item is not None) + expected = int(self.train_engine.get_data_parallel_world_size()) + if [item[0] for item in sources] != list(range(expected)): + raise RuntimeError("initial topology requires one logical source per DP replica") + self._rollout_source_ranks = tuple(item[1] for item in sources) + self._rollout_source_world_size = len(sources) + self._rollout_source_rank = ( + self._rollout_source_ranks.index(self.rank) + if self.rank in self._rollout_source_ranks else None + ) + + + def _is_rollout_source(self) -> bool: + if hasattr(self, "_rollout_source_rank"): + return self._rollout_source_rank is not None + return self.train_engine.is_batch_source() + + + def _source_batch_sample_count(self) -> int: + source_size = int(getattr(self, "_rollout_source_world_size", self.train_engine.get_data_parallel_world_size())) + source_rank = getattr(self, "_rollout_source_rank", self.train_engine.get_data_parallel_rank()) + if source_rank is None: + return 0 + counts = complete_group_sample_counts( + int(self.config.batch_size), int(self.config.n_samples), source_size, + ) + return counts[int(source_rank)] + + + def _collect_fixed_global_batch(self, data_gen, source_count, active_hybrid_ids): + """Gather distinct DP sources before assigning the fixed batch to compute. + + TP/PP followers participate in the same collective but cannot overwrite + another DP source's complete groups. + """ + record = None + if self._is_rollout_source(): + try: + trajectories = self._collect_complete_grpo_batch(data_gen, source_count) + error = "" if len(trajectories) == source_count else ( + f"incomplete fixed batch: expected {source_count}, got {len(trajectories)}" + ) + except Exception as exc: + trajectories, error = [], f"{type(exc).__name__}: {exc}" + record = { + "source_rank": int(getattr(self, "_rollout_source_rank", self.train_engine.get_data_parallel_rank())), + "trajectories": trajectories, "error": error, + } + # Rollout may run longer than the NCCL process-group timeout. Wait + # through the existing out-of-band readiness protocol before entering + # the data collective, including on a source that collected an error. + self._wait_for_all_ranks_after_rollout( + int(getattr(self, "global_step", 0)), + record["trajectories"] if record is not None else [], + ) + if dist.is_available() and dist.is_initialized(): + records = [None] * dist.get_world_size() + dist.all_gather_object(records, { + "source": record, + "core_id": int(self.train_engine.get_data_parallel_rank()), + "hybrids": sorted(active_hybrid_ids), + }) + else: + records = [{"source": record, "core_id": int(self.train_engine.get_data_parallel_rank()), + "hybrids": sorted(active_hybrid_ids)}] + sources, hybrid_targets = {}, {} + for item in records: + for worker_id in item["hybrids"]: + if worker_id in hybrid_targets and hybrid_targets[worker_id] != item["core_id"]: + raise RuntimeError(f"hybrid assigned to multiple core replicas: {worker_id}") + hybrid_targets[worker_id] = item["core_id"] + source = item["source"] + if source is None: + continue + if source["error"]: + raise RuntimeError(f"source {source['source_rank']} batch collection failed: {source['error']}") + if source["source_rank"] in sources: + raise RuntimeError(f"multiple logical batch sources for {source['source_rank']}") + sources[source["source_rank"]] = source + source_size = int(getattr(self, "_rollout_source_world_size", self.train_engine.get_data_parallel_world_size())) + if set(sources) != set(range(source_size)): + raise RuntimeError("missing logical batch source") + dp_size = int(self.train_engine.get_data_parallel_world_size()) + global_batch = [] + for source_rank in range(source_size): + global_batch.extend(sources[source_rank]["trajectories"]) + hybrids = sorted(hybrid_targets) + if len(global_batch) != int(self.config.batch_size): + raise RuntimeError("collected batch does not match fixed global batch size") + participants = [f"dp{i}" for i in range(dp_size)] + sorted(hybrids) + partition_complete_grpo_groups(global_batch, int(self.config.n_samples), participants) + return global_batch, participants + + + def _assign_fixed_global_batch(self, step, global_batch, participants): + assignments = partition_complete_grpo_groups( + global_batch, int(self.config.n_samples), participants, + ) + core_id = f"dp{self.train_engine.get_data_parallel_rank()}" + core_batch = assignments[core_id] + core_ids = {f"dp{i}" for i in range(self.train_engine.get_data_parallel_world_size())} + hybrid_counts = { + worker_id: len(assignments[worker_id]) + for worker_id in participants if worker_id not in core_ids + } + setter = getattr(self.train_engine, "set_elastic_batch_partition", None) + if callable(setter): + setter(len(core_batch), len(global_batch), hybrid_counts) + elif hybrid_counts: + raise RuntimeError("train engine cannot normalize an elastic fixed batch") + dispatch_error = "" + if self.is_main_process and hybrid_counts: + self._pending_hybrid_training_batches = { + worker_id: assignments[worker_id] for worker_id in hybrid_counts + } + try: + self._dispatch_hybrid_training_batches(step) + except Exception as exc: + dispatch_error = f"{type(exc).__name__}: {exc}" + if hybrid_counts and dist.is_available() and dist.is_initialized(): + errors = [None] * dist.get_world_size() + dist.all_gather_object(errors, dispatch_error) + dispatch_error = next((error for error in errors if error), "") + if dispatch_error: + raise RuntimeError(f"fixed-batch hybrid dispatch failed: {dispatch_error}") + return core_batch From 02cc3bff65feba18bf1242033689f9fed357ab3f Mon Sep 17 00:00:00 2001 From: cytsoi <154662279+tcyjohn@users.noreply.github.com> Date: Thu, 22 Oct 2026 13:54:22 +0800 Subject: [PATCH 7/9] Release transient CUDA storage before allocation peaks Release logits after recompute copies or loss construction, and release loss/scaled_loss after backward so unused references do not overlap the next allocation peak. Include lifetime and numerical regression coverage. Release unused CUDA cache once before the first Hybrid optimizer update, retaining the mitigation introduced after an observed NCCL allocation failure. Register its regression test with the file that implements it. --- engine/megatron_core_train_engine.py | 15 +++ scripts/test_cpu_ci.sh | 1 + tests/test_hybrid_optimizer_cuda_cache.py | 125 ++++++++++++++++++++++ tests/test_megatron_core_train_engine.py | 84 +++++++++++++++ 4 files changed, 225 insertions(+) create mode 100644 tests/test_hybrid_optimizer_cuda_cache.py diff --git a/engine/megatron_core_train_engine.py b/engine/megatron_core_train_engine.py index e411149..3b4d7ff 100644 --- a/engine/megatron_core_train_engine.py +++ b/engine/megatron_core_train_engine.py @@ -126,6 +126,7 @@ def __init__( self._elastic_training_step = -1 self._elastic_step_hybrid_workers: tuple[str, ...] = () self._elastic_gradient_update_callback = None + self._hybrid_optimizer_cuda_cache_released = False self.hybrid_lockstep_gradient_sync = True self._elastic_active_gradient_timeout_s = 300.0 self.checkpoints = MegatronDistributedCheckpointManager( @@ -506,6 +507,7 @@ def recompute_logprobs(self, trajectories: list[dict[str, Any]]): logits, micro_batch["input_ids"], ).cpu() + del logits for row, trajectory in zip( action_log_probs, micro_batch["trajectories"], @@ -532,11 +534,13 @@ def grpo_update( for micro_batch in micro_batches: logits = self._forward(micro_batch) loss, stats = self._loss(micro_batch, logits) + del logits scaled_loss = self.optimizer.scale_loss( loss * (len(micro_batch["trajectories"]) / max(len(trajectories), 1)) * self._elastic_core_loss_scale() ) scaled_loss.backward() + del loss, scaled_loss local_stats.append(stats) finalize_model_grads(self.model) @@ -822,9 +826,20 @@ def apply_elastic_gradient_update(self, tensors: tuple[torch.Tensor, ...], *, st raise ValueError(f"Hybrid update tensor count mismatch: {len(tensors)} != {len(local)}") for dst, src in zip(local, tensors): dst.copy_(src.to(device=dst.device, dtype=dst.dtype)) + cache_released = False + if not getattr(self, "_hybrid_optimizer_cuda_cache_released", False) and torch.cuda.is_available(): + device = next((grad.device for grad in local if grad.device.type == "cuda"), None) + if device is not None: + # Return unused cached blocks to external CUDA allocators + # before this Hybrid's first optimizer update. + with torch.cuda.device(device): + torch.cuda.empty_cache() + cache_released = True ok, _norm, _zeros = self.optimizer.step() if not ok: raise FloatingPointError("Hybrid Megatron optimizer rejected the update") + if cache_released: + self._hybrid_optimizer_cuda_cache_released = True self.current_version = int(state_version) @staticmethod diff --git a/scripts/test_cpu_ci.sh b/scripts/test_cpu_ci.sh index 8bebfe7..7df1cfb 100644 --- a/scripts/test_cpu_ci.sh +++ b/scripts/test_cpu_ci.sh @@ -34,6 +34,7 @@ python -m pytest -q --strict-markers --junitxml=reports/cpu-tests.xml \ tests/test_grpo_sample_mean.py \ tests/test_hybrid_sample_statistics_data.py \ tests/test_elastic_gradient_idle_path.py \ + tests/test_hybrid_optimizer_cuda_cache.py \ tests/test_torch_optimizer_memory.py \ tests/test_megatron_core_train_engine.py \ tests/test_hybrid_worker_batch_protocol.py diff --git a/tests/test_hybrid_optimizer_cuda_cache.py b/tests/test_hybrid_optimizer_cuda_cache.py new file mode 100644 index 0000000..34a3d7c --- /dev/null +++ b/tests/test_hybrid_optimizer_cuda_cache.py @@ -0,0 +1,125 @@ +"""Hybrid cache release happens once, without changing gradient/update semantics.""" + +from contextlib import contextmanager +from types import SimpleNamespace + +import pytest +import torch + +from RL_Framework.engine.megatron_core_train_engine import MegatronCoreTrainEngine + + +class DeviceGradient: + def __init__(self, calls): + self.calls = calls + self.device = torch.device('cuda:1') + self.dtype = torch.bfloat16 + self.tensor = torch.zeros(2, dtype=self.dtype) + + def copy_(self, value): + self.calls.append('copy') + self.tensor.copy_(value) + + +class Source: + def __init__(self, value): + self.tensor = torch.tensor(value, dtype=torch.float32) + + def to(self, *, device, dtype): + assert device == torch.device('cuda:1') + return self.tensor.to(dtype=dtype) + + +def cuda_engine(monkeypatch): + calls = [] + gradients = [DeviceGradient(calls), DeviceGradient(calls)] + engine = MegatronCoreTrainEngine.__new__(MegatronCoreTrainEngine) + engine.rank, engine.current_version = 0, 1 + engine._model_main_gradients = lambda: gradients + engine.optimizer = SimpleNamespace(step=lambda: (calls.append('step') or (True, 0., 0))) + monkeypatch.setattr(torch.cuda, 'is_available', lambda: True) + + @contextmanager + def device_context(device): + assert device == torch.device('cuda:1') + calls.append('enter') + yield + calls.append('exit') + + monkeypatch.setattr(torch.cuda, 'device', device_context) + monkeypatch.setattr(torch.cuda, 'empty_cache', lambda: calls.append('empty')) + return engine, gradients, calls + + +def test_first_cuda_hybrid_update_releases_after_copy_and_before_step(monkeypatch, capsys): + engine, gradients, calls = cuda_engine(monkeypatch) + sources = (Source([1.25, 2.5]), Source([3., 4.])) + engine.apply_elastic_gradient_update(sources, state_version=2) + assert calls == ['copy', 'copy', 'enter', 'empty', 'exit', 'step'] + assert engine.current_version == 2 and engine._hybrid_optimizer_cuda_cache_released + for gradient, source in zip(gradients, sources): + assert torch.equal(gradient.tensor, source.tensor.bfloat16()) + assert not capsys.readouterr().out + calls.clear() + engine.apply_elastic_gradient_update(sources, state_version=3) + assert calls == ['copy', 'copy', 'step'] and engine.current_version == 3 + assert not capsys.readouterr().out + + +def test_optimizer_exception_is_preserved_and_cache_release_can_retry(monkeypatch): + engine, _gradients, calls = cuda_engine(monkeypatch) + failure = RuntimeError('original NCCL failure') + + def failed_step(): + calls.append('step') + raise failure + + engine.optimizer.step = failed_step + with pytest.raises(RuntimeError) as error: + engine.apply_elastic_gradient_update((Source([1., 2.]), Source([3., 4.])), state_version=2) + assert error.value is failure and engine.current_version == 1 + assert not getattr(engine, '_hybrid_optimizer_cuda_cache_released', False) + engine.optimizer.step = lambda: (True, 0., 0) + engine.apply_elastic_gradient_update((Source([1., 2.]), Source([3., 4.])), state_version=2) + assert calls.count('empty') == 2 and engine.current_version == 2 + + +def test_rejected_optimizer_update_does_not_mark_cache_or_revision_complete(monkeypatch): + engine, _gradients, calls = cuda_engine(monkeypatch) + engine.optimizer.step = lambda: (False, 0., 0) + with pytest.raises(FloatingPointError, match='rejected the update'): + engine.apply_elastic_gradient_update((Source([1., 2.]), Source([3., 4.])), state_version=2) + assert calls.count('empty') == 1 and engine.current_version == 1 + assert not getattr(engine, '_hybrid_optimizer_cuda_cache_released', False) + + +@pytest.mark.parametrize('cuda_available', [False, True]) +def test_cpu_gradients_keep_original_copy_dtype_and_optimizer_update(monkeypatch, cuda_available): + engine = MegatronCoreTrainEngine.__new__(MegatronCoreTrainEngine) + gradient = torch.zeros(2, dtype=torch.bfloat16) + source = torch.tensor([1.25, -2.5], dtype=torch.float32) + source_before = source.clone() + weight = torch.tensor([10., 20.]) + engine.current_version = 1 + engine._model_main_gradients = lambda: [gradient] + monkeypatch.setattr(torch.cuda, 'is_available', lambda: cuda_available) + monkeypatch.setattr(torch.cuda, 'empty_cache', lambda: pytest.fail('CPU path released CUDA cache')) + + def step(): + weight.sub_(gradient.float(), alpha=.1) + return True, 0., 0 + + engine.optimizer = SimpleNamespace(step=step) + engine.apply_elastic_gradient_update((source,), state_version=2) + assert torch.equal(source, source_before) + assert torch.equal(gradient, source.bfloat16()) + assert torch.equal(weight, torch.tensor([10., 20.]) - source.bfloat16().float() * .1) + assert engine.current_version == 2 + assert not getattr(engine, '_hybrid_optimizer_cuda_cache_released', False) + + +def test_tensor_count_failure_happens_before_copy_or_cache_release(monkeypatch): + engine, _gradients, calls = cuda_engine(monkeypatch) + with pytest.raises(ValueError, match='tensor count mismatch'): + engine.apply_elastic_gradient_update((Source([1., 2.]),), state_version=2) + assert calls == [] and engine.current_version == 1 diff --git a/tests/test_megatron_core_train_engine.py b/tests/test_megatron_core_train_engine.py index 50fbf1d..0a99807 100644 --- a/tests/test_megatron_core_train_engine.py +++ b/tests/test_megatron_core_train_engine.py @@ -1,4 +1,6 @@ from types import SimpleNamespace +import sys +import weakref import pytest import torch @@ -148,3 +150,85 @@ def new_group(**kwargs): } ] assert engine._elastic_gradient_group_ranks == (0, 8, 16, 24) + + +def test_recompute_releases_logits_before_next_forward_and_preserves_logprobs(monkeypatch): + engine = MegatronCoreTrainEngine.__new__(MegatronCoreTrainEngine) + engine.train_tp_size = 1 + weights = torch.arange(49, dtype=torch.float32).reshape(7, 7) / 10 + rows = [dict(input_ids=torch.tensor([[0, 2, 3, 1]]), logprobs=torch.zeros(1, 4)), + dict(input_ids=torch.tensor([[4, 1, 5, 6]]), logprobs=torch.zeros(1, 4))] + outputs, modes = [], [] + def forward(batch): + assert all(output() is None for output in outputs) + logits = weights[batch['input_ids']] + outputs.append(weakref.ref(logits)) + return logits + monkeypatch.setattr(engine, '_forward', forward) + monkeypatch.setattr(engine, '_set_train_mode', modes.append) + monkeypatch.setattr(engine, '_iter_recompute_micro_batches', + lambda trajectories: (dict(input_ids=row['input_ids'], trajectories=[row]) for row in trajectories)) + engine.recompute_logprobs(rows) + assert all(output() is None for output in outputs) + assert modes == [False, True] + for row in rows: + expected = torch.log_softmax(weights[row['input_ids']][:, :-1], dim=-1) + expected = expected.gather(-1, row['input_ids'][:, 1:, None]).squeeze(-1) + torch.testing.assert_close(row['logprobs'][:, 1:], expected) + assert row['logprobs'][0, 0] == 0 + + + +def test_grpo_releases_backward_outputs_before_next_microbatch_and_optimizer(monkeypatch): + monkeypatch.setitem(sys.modules, 'megatron.core.distributed', + SimpleNamespace(finalize_model_grads=lambda _: None)) + engine = MegatronCoreTrainEngine.__new__(MegatronCoreTrainEngine) + engine.train_tp_size = 1 + engine.clip_epsilon, engine.kl_coef = .2, .1 + engine.current_version, engine.optimizer_revision = 3, 0 + initial = torch.arange(49, dtype=torch.float32).reshape(7, 7) / 10 + parameter = torch.nn.Parameter(initial.clone()) + reference = torch.nn.Parameter(initial.clone()) + optimizer = torch.optim.SGD([parameter], lr=.01) + reference_optimizer = torch.optim.SGD([reference], lr=.01) + rows = [dict(input_ids=torch.tensor([[0, 2, 3, 1]])), + dict(input_ids=torch.tensor([[4, 1, 5, 6]]))] + def batches(trajectories): + for row in trajectories: + yield dict(input_ids=row['input_ids'], old_logprobs=torch.full((1, 4), -2.), + loss_mask=torch.tensor([[0., 1., 1., 1.]]), advantages=torch.tensor([.7]), + rewards=torch.tensor([1.]), trajectories=[row]) + references = [] + def forward(batch): + assert all(output() is None for output in references) + logits = parameter[batch['input_ids']] + references.append(weakref.ref(logits)) + return logits + gradients = [] + def step(): + assert all(output() is None for output in references) + gradients.append(parameter.grad.clone()) + optimizer.step() + return True, parameter.grad.norm(), 0 + def scale_loss(value): + assert all(output() is None for output in references) + return value + engine.model = [] + engine.optimizer = SimpleNamespace(scale_loss=scale_loss, step=step) + monkeypatch.setattr(engine, '_iter_micro_batches', batches) + monkeypatch.setattr(engine, '_forward', forward) + monkeypatch.setattr(engine, '_zero_grad', lambda: optimizer.zero_grad(set_to_none=True)) + monkeypatch.setattr(engine, '_elastic_core_loss_scale', lambda: 1.) + monkeypatch.setattr(engine, '_apply_elastic_inter_replica_gradients', lambda: None) + monkeypatch.setattr(engine, '_reduce_stats', lambda stats: stats) + engine.grpo_update(rows, ppo_epochs=2) + for gradient in gradients: + reference_optimizer.zero_grad(set_to_none=True) + for batch in batches(rows): + loss, _ = engine._loss(batch, reference[batch['input_ids']]) + (loss / len(rows)).backward() + torch.testing.assert_close(gradient, reference.grad) + reference_optimizer.step() + torch.testing.assert_close(parameter, reference) + assert engine.optimizer_revision == 2 + assert all(output() is None for output in references) From 6fd100abf7b3e12f0b87437d41c801068486a3a3 Mon Sep 17 00:00:00 2001 From: cytsoi <154662279+tcyjohn@users.noreply.github.com> Date: Thu, 22 Oct 2026 13:54:44 +0800 Subject: [PATCH 8/9] Keep Hybrid ownership coherent across lending, activation and publication Publish to the current serving receivers while preserving borrowed Hybrid replicas. Return resources through planner or sampling-supply transitions, confirm old process exit, and install current weights before reopening. Lend complete same-host replicas using the upstream TP/EP/PP/CP width and physical rollout reserve. Expand existing loans without changing their slots or generation; roll back only the new allocation after partial failure. Hand off cold startup after the first installation ACK. Ready joins wait for a Core boundary without a separate queue deadline, while worker startup and warm installation retain their existing timeout budgets. Reconcile activation and cancellation under the existing lifecycle lock. Release unused failed-join snapshots at a shared Core boundary after all loaders finish. Wait for assigned tasks during normal return, retire stale tasks after confirmed abnormal exit, and aggregate only consumed Hybrid gradient statistics. Include the lifecycle and boundary regression tests. --- engine/megatron_core_train_engine.py | 167 ++-- infra/cost_model/global_resource_planner.py | 121 ++- infra/elastic/core_reconfiguration.py | 129 +++ infra/elastic/hybrid_pool.py | 81 +- infra/elastic/process_lifecycle.py | 62 ++ infra/elastic/readiness.py | 82 ++ infra/elastic/rollout_lending.py | 231 +++++ infra/elastic/rollout_lifecycle.py | 126 +++ infra/elastic/rollout_process.py | 410 ++++++++ infra/elastic/rollout_supply.py | 25 + infra/elastic/runtime_executor.py | 873 +++++++++++++----- infra/elastic/state_transfer.py | 148 +++ infra/execution/batch_dispatcher.py | 6 + scripts/elastic_hybrid_controller.py | 84 +- scripts/elastic_hybrid_worker.py | 59 +- scripts/test_cpu_ci.sh | 27 +- tests/test_core_reconfiguration_protocol.py | 188 ++++ tests/test_core_subset_publication.py | 68 ++ tests/test_ehp_batch_partition.py | 16 +- tests/test_ehp_membership_release.py | 25 +- tests/test_ehp_readiness_stop.py | 231 +++++ tests/test_elastic_gradient_idle_path.py | 4 +- tests/test_elastic_hybrid_pool.py | 44 + tests/test_fixed_batch_gradients.py | 43 + tests/test_grp_decision_gates.py | 103 +++ tests/test_grpo_grouping.py | 74 ++ tests/test_hybrid_boundary_activation.py | 197 ++++ tests/test_hybrid_join_readiness_budget.py | 221 +++++ tests/test_hybrid_optimizer_cuda_cache.py | 18 +- tests/test_hybrid_sample_statistics.py | 44 + tests/test_hybrid_task_lifecycle.py | 124 +++ tests/test_publication_preserves_hybrids.py | 179 ++++ tests/test_rollout_lending.py | 221 +++++ tests/test_rollout_lifecycle.py | 333 +++++++ tests/test_rollout_process.py | 200 ++++ tests/test_rollout_supply.py | 132 +++ .../test_runtime_blue_green_gpu_ownership.py | 122 +++ tests/test_runtime_boundary_integration.py | 21 +- tests/test_runtime_core_capabilities.py | 218 +++++ tests/test_runtime_elastic_executor.py | 216 ++++- tests/test_runtime_executor_shutdown.py | 31 +- tests/test_runtime_failed_snapshot_cleanup.py | 248 +++++ tests/test_runtime_repartition_eligibility.py | 64 ++ tests/test_runtime_rollout_ownership.py | 59 ++ trainer/async_rl_trainer.py | 530 ++++++----- 45 files changed, 5985 insertions(+), 620 deletions(-) create mode 100644 infra/elastic/core_reconfiguration.py create mode 100644 infra/elastic/process_lifecycle.py create mode 100644 infra/elastic/readiness.py create mode 100644 infra/elastic/rollout_lending.py create mode 100644 infra/elastic/rollout_lifecycle.py create mode 100644 infra/elastic/rollout_process.py create mode 100644 infra/elastic/rollout_supply.py create mode 100644 infra/elastic/state_transfer.py create mode 100644 tests/test_core_reconfiguration_protocol.py create mode 100644 tests/test_core_subset_publication.py create mode 100644 tests/test_ehp_readiness_stop.py create mode 100644 tests/test_grp_decision_gates.py create mode 100644 tests/test_hybrid_boundary_activation.py create mode 100644 tests/test_hybrid_join_readiness_budget.py create mode 100644 tests/test_hybrid_sample_statistics.py create mode 100644 tests/test_hybrid_task_lifecycle.py create mode 100644 tests/test_publication_preserves_hybrids.py create mode 100644 tests/test_rollout_lending.py create mode 100644 tests/test_rollout_lifecycle.py create mode 100644 tests/test_rollout_process.py create mode 100644 tests/test_rollout_supply.py create mode 100644 tests/test_runtime_blue_green_gpu_ownership.py create mode 100644 tests/test_runtime_core_capabilities.py create mode 100644 tests/test_runtime_failed_snapshot_cleanup.py create mode 100644 tests/test_runtime_repartition_eligibility.py create mode 100644 tests/test_runtime_rollout_ownership.py diff --git a/engine/megatron_core_train_engine.py b/engine/megatron_core_train_engine.py index 3b4d7ff..67a8704 100644 --- a/engine/megatron_core_train_engine.py +++ b/engine/megatron_core_train_engine.py @@ -109,6 +109,10 @@ def __init__( self.is_main_process = self.rank == 0 self.current_version = 0 self.optimizer_revision = 0 + self._elastic_state_server = None + self._elastic_snapshot_urls = {} + self._elastic_expected_state_version = 0 + self._elastic_seen_payloads = set() self._elastic_batch_partition = None self.bridge = None @@ -544,12 +548,15 @@ def grpo_update( local_stats.append(stats) finalize_model_grads(self.model) - self._apply_elastic_inter_replica_gradients() + hybrid_stats = self._apply_elastic_inter_replica_gradients() update_successful, grad_norm, _ = self.optimizer.step() if not update_successful: raise FloatingPointError("Megatron optimizer rejected the update") self.optimizer_revision += 1 - merged = self._merge_stats(local_stats) + merged = self._merge_stats(local_stats + (hybrid_stats or [])) + # Hybrid contributes samples, not extra Core optimizer updates or + # additional weight in the existing DP gradient-norm statistic. + merged["n_updates"] = sum(float(item.get("n_updates", 1.0)) for item in local_stats) if isinstance(grad_norm, torch.Tensor): merged["grad_norm"] = float(grad_norm.detach().item()) else: @@ -740,6 +747,10 @@ def _elastic_communication_domains(self): ) def close_elastic_communication_domain(self) -> None: + if self._elastic_state_server is not None: + self._elastic_state_server.close() + self._elastic_state_server = None + self._elastic_snapshot_urls.clear() group = self._elastic_gradient_process_group self._elastic_gradient_process_group = None self._elastic_gradient_group_ranks = () @@ -754,14 +765,45 @@ def get_elastic_gradient_process_group(self): return self._elastic_gradient_process_group def enqueue_hybrid_gradient_payload(self, payload: GradientPayload) -> None: + domain = self.elastic_gradient_domain + if domain is None: + raise RuntimeError("hybrid gradient arrived without an elastic domain") + membership = domain.membership_state() + epochs = membership["membership_epochs"] + if payload.step != self._elastic_training_step: + raise ValueError("hybrid gradient training step does not match") + if payload.state_version != self._elastic_expected_state_version: + raise ValueError("hybrid gradient optimizer revision does not match") + if payload.membership_epoch != epochs.get(payload.replica_id): + raise ValueError("hybrid gradient membership epoch does not match") + if payload.replica_id not in self._elastic_step_hybrid_workers: + raise ValueError("unexpected hybrid gradient replica") + if payload.zero_placeholder: + raise ValueError("active hybrid replicas cannot replace assigned data with placeholders") + if payload.replica_rank != domain.local_replica_rank or payload.replica_world_size != domain.replica_world_size: + raise ValueError("hybrid gradient model lane does not match") + partition = self._elastic_batch_partition + if partition is not None: + _, total, counts = partition + if payload.sample_count != counts.get(payload.replica_id) or payload.global_sample_count != total: + raise ValueError("hybrid gradient sample count does not match partition") + identity = (payload.replica_id, payload.replica_rank, payload.step, + payload.state_version, payload.membership_epoch) with self._hybrid_gradient_condition: + if identity in self._elastic_seen_payloads: + raise ValueError("duplicate hybrid gradient payload") + self._elastic_seen_payloads.add(identity) self._pending_hybrid_gradients.append(payload) self._hybrid_gradient_condition.notify_all() - def set_elastic_training_step(self, step: int, active_hybrid_workers=None) -> None: + def set_elastic_training_step(self, step: int, active_hybrid_workers=None, *, state_version=None) -> None: self._elastic_training_step = int(step) + self._elastic_expected_state_version = int( + self.optimizer_revision if state_version is None else state_version + ) if active_hybrid_workers is not None: self._elastic_step_hybrid_workers = tuple(str(w) for w in active_hybrid_workers) + self._elastic_seen_payloads = {key for key in self._elastic_seen_payloads if key[2] >= step - 1} def set_elastic_batch_partition(self, local_sample_count, global_sample_count, hybrid_sample_counts): if local_sample_count < 0 or global_sample_count <= 0 or local_sample_count > global_sample_count: @@ -801,6 +843,8 @@ def compute_elastic_gradient_payload( global_sample_count: int = 0, ) -> GradientPayload: from megatron.core.distributed import finalize_model_grads + if int(state_version) != self.optimizer_revision: + raise ValueError("hybrid task optimizer revision does not match installed state") micro_batches = list(self._iter_micro_batches(trajectories)) self._zero_grad() self._set_train_mode(True) @@ -840,7 +884,7 @@ def apply_elastic_gradient_update(self, tensors: tuple[torch.Tensor, ...], *, st raise FloatingPointError("Hybrid Megatron optimizer rejected the update") if cache_released: self._hybrid_optimizer_cuda_cache_released = True - self.current_version = int(state_version) + self.optimizer_revision = int(state_version) @staticmethod def _clone_cpu_state(value): @@ -939,50 +983,55 @@ def restore_elastic_state(self, payload): self._elastic_state_receipt = receipt return receipt - def capture_elastic_state_snapshot( - self, - worker_id: str, - target_core_id: str, - ) -> int: - """Write metadata for the latest collective distributed checkpoint.""" + def capture_elastic_state_snapshot(self, worker_id, target_core_id): del worker_id, target_core_id - snapshot_path = Path( - self.get_elastic_state_snapshot_path(self.current_version) - ) - snapshot_path.parent.mkdir(parents=True, exist_ok=True) - manifest_path = ( - self.checkpoints.sync_path - / f"v{self.current_version}" - / "megatron_checkpoint_manifest.json" - ) - torch.save( - { - "backend": self.train_backend, - "version": self.current_version, - "manifest_path": str(manifest_path), - "parallel_state": self.get_parallel_state(), - }, - snapshot_path, - ) - return self.current_version - - def get_elastic_state_snapshot_path(self, version: int) -> str: - state_dir = Path( - os.getenv( - "ELASTIC_TRAINING_STATE_DIR", - "./logs/elastic_training_state", - ) - ) - return str(state_dir / f"v{version}" / f"rank_{self.rank}.pt") - - def load_elastic_state_snapshot(self, snapshot_path: str) -> None: - payload = torch.load( - snapshot_path, - map_location="cpu", - weights_only=False, - ) - version = int(payload["version"]) - self.load_weights(str(self.checkpoints.sync_path), version) + from RL_Framework.infra.elastic.state_transfer import LiveStateServer + if self._elastic_state_server is None: + self._elastic_state_server = LiveStateServer() + revision = int(self.optimizer_revision) + signature = tuple(sorted((key, repr(value)) for key, value in self.get_parallel_state().items())) + if getattr(self, "_elastic_snapshot_topology", signature) != signature: + for previous in list(self._elastic_snapshot_urls): + self.release_elastic_state_snapshot(previous) + self._elastic_snapshot_topology = signature + if revision not in self._elastic_snapshot_urls: + self._elastic_snapshot_urls[revision] = self._elastic_state_server.publish(self.capture_elastic_state()) + return revision + + def get_elastic_state_snapshot_path(self, version): + try: + return self._elastic_snapshot_urls[int(version)] + except KeyError as exc: + raise RuntimeError(f"live EHP state revision {version} has not been captured") from exc + + def release_elastic_state_snapshot(self, version): + uri = self._elastic_snapshot_urls.pop(int(version), None) + if uri is not None: + self._elastic_state_server.release(uri) + + + def load_elastic_state_snapshot(self, snapshot_path, *, target_core_id=None, replica_rank=None): + from RL_Framework.infra.elastic.state_transfer import receive_live_state + uri = snapshot_path + manifest = None + if not str(snapshot_path).startswith("ehp://"): + manifest = json.loads(Path(snapshot_path).read_text(encoding="utf-8")) + if manifest.get("format") != "libra-ehp-lanes-v1": + raise ValueError("invalid live EHP lane manifest") + if target_core_id is None: + raise ValueError("target core identity is required for a live EHP lane manifest") + lane = int(os.environ.get("RANK", "0")) if replica_rank is None else int(replica_rank) + uri = manifest["lanes"][str(target_core_id)][str(lane)] + payload = receive_live_state(uri) + if manifest is not None and int(payload["optimizer_revision"]) != int(manifest["optimizer_revision"]): + raise ValueError("live EHP manifest revision does not match transferred state") + receipt = self.restore_elastic_state(payload) + if manifest is not None: + receipt["lane_snapshot_id"] = receipt["snapshot_id"] + receipt["snapshot_id"] = manifest["snapshot_id"] + receipt["target_core_id"] = str(target_core_id) + receipt["source_uri"] = str(uri) + return receipt def _iter_micro_batches(self, trajectories: list[dict[str, Any]]): for start in range(0, len(trajectories), self.micro_batch_size): @@ -1252,15 +1301,15 @@ def _materialize_elastic_core_gradients(self): dist.all_gather_into_tensor(flat, local, group=group) - def _apply_elastic_inter_replica_gradients(self) -> None: + def _apply_elastic_inter_replica_gradients(self) -> list[dict[str, float]]: pending = getattr(self, "_pending_hybrid_gradients", []) expected = set(getattr(self, "_elastic_step_hybrid_workers", ())) if self.elastic_gradient_domain is None and not pending and not expected: - return + return [] if self.elastic_gradient_domain is None: with self._hybrid_gradient_condition: self._pending_hybrid_gradients.clear() - return + return [] deadline = (time.monotonic() + self._elastic_active_gradient_timeout_s if self._elastic_active_gradient_timeout_s > 0 else None) with self._hybrid_gradient_condition: @@ -1268,7 +1317,8 @@ def _apply_elastic_inter_replica_gradients(self) -> None: present = { p.replica_id for p in self._pending_hybrid_gradients if p.replica_id in expected - and (p.step < 0 or p.step == self._elastic_training_step) + and p.step == self._elastic_training_step + and p.state_version == self._elastic_expected_state_version } missing = expected - present if not missing: @@ -1286,7 +1336,7 @@ def _apply_elastic_inter_replica_gradients(self) -> None: if not params_and_grads: with self._hybrid_gradient_condition: self._pending_hybrid_gradients.clear() - return + return [] core_id = f"dp{self.get_data_parallel_rank()}" partition = self._elastic_batch_partition @@ -1299,12 +1349,22 @@ def _apply_elastic_inter_replica_gradients(self) -> None: if weighted and world > 1: for grad in local: grad.div_(world) + consumed_before = set(self.elastic_gradient_domain._consumed_payloads) reduced = self.elastic_gradient_domain.reduce_core_gradients( core_gradients={core_id: local}, hybrid_payloads=pending, core_sample_counts={core_id: partition[1]} if weighted else None, global_sample_count=partition[1] if weighted else 0, + step=self._elastic_training_step, + state_version=self._elastic_expected_state_version, in_place=True, ) + # The reducer owns acceptance and lane routing. Only identities newly + # consumed by this successful reduction may contribute sample metrics. + consumed = self.elastic_gradient_domain._consumed_payloads - consumed_before + hybrid_stats = [payload.stats for payload in pending + if (payload.replica_id, payload.replica_rank, payload.step, + payload.state_version, payload.membership_epoch) in consumed + and payload.stats] for (_, grad), reduced_grad in zip( params_and_grads, reduced[core_id], @@ -1330,12 +1390,13 @@ def _apply_elastic_inter_replica_gradients(self) -> None: replica_id=payload.replica_id, tensors=update_tensors, step=payload.step, - state_version=self.current_version, + state_version=self._elastic_expected_state_version + 1, membership_epoch=payload.membership_epoch, ) ) with self._hybrid_gradient_condition: self._pending_hybrid_gradients.clear() + return hybrid_stats @staticmethod def _parallel_state(): diff --git a/infra/cost_model/global_resource_planner.py b/infra/cost_model/global_resource_planner.py index 3b5be43..5fc5a70 100644 --- a/infra/cost_model/global_resource_planner.py +++ b/infra/cost_model/global_resource_planner.py @@ -26,6 +26,51 @@ logger = logging.getLogger(__name__) +@dataclass(frozen=True) +class CoreRuntimeCapabilities: + """Fixed Core layout on the already initialized physical ranks. + + Hybrid replicas use a separate process group and do not enlarge this rank + set. No capability is supplied to startup GRP, which may choose its size. + """ + world_size: int + tp: int + pp: int + cp: int + ep: int + dp: int + rebuild_tp_dp: bool = field(default=False, init=False) + + @classmethod + def from_engine(cls, engine): + world = getattr(engine, "world_size", None) + if not isinstance(world, int) or world < 1: + return None + getter = getattr(engine, "get_parallel_state", None) + state = getter() if callable(getter) else {} + dimensions = {axis: int(state.get(f"train_{axis}", + getattr(engine, f"train_{axis}_size", 1))) for axis in ("tp", "pp", "cp", "ep", "dp")} + return cls(world_size=world, **dimensions) + + def incompatibility(self, plan): + topology = plan.train_config + if int(plan.train_gpus) != self.world_size: + return "physical_core_rank_set_changed" + for axis in ("pp", "cp", "ep"): + if int(getattr(topology, axis, 1)) != getattr(self, axis): + return f"core_{axis}_rebuild_unsupported" + if not self.rebuild_tp_dp and (int(topology.tp) != self.tp + or int(topology.dp) * int(getattr(topology, "ep", 1)) != self.dp): + return "core_tp_dp_rebuild_unsupported" + return "" + + def to_dict(self): + return {"physical_rank_count": self.world_size, "tp": self.tp, "pp": self.pp, + "cp": self.cp, "ep": self.ep, "dp": self.dp, + "rebuild_tp_dp": self.rebuild_tp_dp, + "physical_rank_resize": False, "source": "initialized_training_engine"} + + @dataclass class GlobalResourcePlan: """Concrete cross-stage plan selected by the planner.""" @@ -189,11 +234,19 @@ class PlannerDecision: trigger: str = "" runtime_metrics: RuntimeResourceMetrics | None = None elastic_hybrid_signal: ElasticHybridSignal | None = None + # None preserves legacy hand-constructed decisions. Planner-produced + # decisions explicitly distinguish full repartition from EHP adjustment. + resource_should_reconfigure: bool | None = None def to_dict(self) -> dict[str, Any]: return { "step": self.step, "should_reconfigure": self.should_reconfigure, + "resource_should_reconfigure": self.resource_should_reconfigure, + "elastic_should_adjust": bool( + self.elastic_hybrid_signal is not None + and self.elastic_hybrid_signal.should_adjust + ), "reason": self.reason, "trigger": self.trigger, "num_requests": self.num_requests, @@ -317,6 +370,7 @@ def __init__( 1, int(elastic_hybrid_signal_ttl_steps) ) self.verbose = verbose + self.runtime_capabilities: CoreRuntimeCapabilities | None = None self.optimizer = TwoLevelNestedOptimizer( cost_model=evaluator, @@ -524,6 +578,7 @@ def plan_if_needed( num_requests=len(self._history), trigger="warmup", runtime_metrics=runtime_metrics, + resource_should_reconfigure=False, ) ) @@ -538,6 +593,7 @@ def plan_if_needed( num_requests=len(self._history), trigger=trigger, runtime_metrics=runtime_metrics, + resource_should_reconfigure=False, ) ) if trigger == "interval_skip": @@ -550,6 +606,7 @@ def plan_if_needed( num_requests=len(self._history), trigger=trigger, runtime_metrics=runtime_metrics, + resource_should_reconfigure=False, ) ) if len(self._history) < self.min_history_size: @@ -562,6 +619,7 @@ def plan_if_needed( num_requests=len(self._history), trigger=trigger, runtime_metrics=runtime_metrics, + resource_should_reconfigure=False, ) ) @@ -614,6 +672,7 @@ def plan_if_needed( num_requests=len(self._history), trigger=trigger, runtime_metrics=runtime_metrics, + resource_should_reconfigure=False, ) ) @@ -671,12 +730,42 @@ def plan_if_needed( }, ) + # Keep the unconstrained capacity recommendation as EHP demand, while + # selecting Core/rollout layouts that the live backend can execute. + # Startup has no live capabilities and retains the full 4/8/... search. + elastic_candidate = candidate + capabilities = self.runtime_capabilities + rejected = capabilities.incompatibility(candidate) if capabilities is not None else "" + if rejected: + unconstrained = candidate.to_dict() + if not forced_candidate: + feasible = self.optimizer.optimize( + n_total_gpus=self.n_total_gpus, requests=workload, + B_global=config.batch_size, + allowed_train_tp=(self.allowed_train_tp if capabilities.rebuild_tp_dp else [capabilities.tp]), + allowed_train_ep=[capabilities.ep], allowed_train_pp=[capabilities.pp], + micro_batch_sizes=self.micro_batch_sizes, + fixed_train_gpus=capabilities.world_size, + allocation_granularity_gpus=self.allocation_granularity_gpus, + min_train_gpus=self.min_train_gpus, min_rollout_gpus=self.min_rollout_gpus, + rollout_node_tp_pattern=list(getattr(config.global_resource_planner, + "rollout_node_tp_pattern", []) or []), + train_cp_size=capabilities.cp, + dp_batch_group_size=max(1, int(getattr(config, "n_samples", 1) or 1))) + candidate = self._plan_from_result(feasible, config) or current + result = feasible + else: + # Forced test commands also cannot authorize inventing ranks. + candidate = current + candidate.metadata["runtime_candidate_filter"] = { + "reason": rejected, "capabilities": capabilities.to_dict(), + "unconstrained_candidate": unconstrained, + "elastic_demand_preserved": True} + if capabilities is not None: + candidate.metadata["runtime_core_capabilities"] = capabilities.to_dict() candidate.metadata["workload"] = { - "global_batch_size": int(config.batch_size), - "history_size": len(self._history), - "representative_requests": len(workload), - "selection": "deterministic_length_quantiles", - } + "global_batch_size": int(config.batch_size), "history_size": len(self._history), + "representative_requests": len(workload), "selection": "deterministic_length_quantiles"} raw_gain = current.t_global - candidate.t_global net_gain = raw_gain - self.reconfiguration_cost_s gain_ratio = raw_gain / max(current.t_global, 1e-6) @@ -689,7 +778,7 @@ def plan_if_needed( step=step, config=config, current=current, - candidate=candidate, + candidate=elastic_candidate, runtime_metrics=runtime_metrics, ) elastic_adjust = bool( @@ -700,6 +789,7 @@ def plan_if_needed( ) resource_should = ( (not same_plan or force_runtime_reconfigure) + and not (rejected and forced_candidate) and ( forced_candidate or force_runtime_reconfigure @@ -733,12 +823,16 @@ def plan_if_needed( trigger=trigger, runtime_metrics=runtime_metrics, elastic_hybrid_signal=elastic_signal, + resource_should_reconfigure=resource_should, ) - if should: + if resource_should: self._active_plan = candidate - self._last_replan_step = int(step) - elif self._active_plan is None: + else: + # An EHP join/release does not apply the rejected candidate's + # training topology or rollout TP layout. self._active_plan = current + if should: + self._last_replan_step = int(step) self._mark_rejected_trigger_observed(trigger, runtime_metrics) return self._remember(decision) @@ -1029,6 +1123,8 @@ def apply_plan_to_config( plan: GlobalResourcePlan, config: AsyncRLConfig, ) -> None: + preserve_slots = (config.train_gpus == plan.train_gpus + and config.rollout_gpus == plan.rollout_gpus) config.n_total_gpus = plan.n_total_gpus config.train_gpus = plan.train_gpus config.rollout_gpus = plan.rollout_gpus @@ -1047,6 +1143,13 @@ def apply_plan_to_config( hetero.enabled = True hetero.total_gpus = plan.rollout_gpus host_gpu_pools = self._rollout_host_gpu_pools(config, plan.rollout_gpus) + if preserve_slots and hetero.instances: + observed: dict[str, set[int]] = {} + for instance in hetero.instances: + observed.setdefault(instance.host or hetero.vllm_host, set()).update(instance.gpus) + # Repartition the actual rollout pool, excluding the fixed Core GPUs. + if sum(map(len, observed.values())) == plan.rollout_gpus: + host_gpu_pools = [(host, sorted(gpus)) for host, gpus in observed.items()] if host_gpu_pools: hetero.available_gpus = sorted( {gpu for _, gpus in host_gpu_pools for gpu in gpus} diff --git a/infra/elastic/core_reconfiguration.py b/infra/elastic/core_reconfiguration.py new file mode 100644 index 0000000..05894bc --- /dev/null +++ b/infra/elastic/core_reconfiguration.py @@ -0,0 +1,129 @@ +"""Collective core-state operations inside the existing training boundary. + +Only small commands and receipts use the shared filesystem. Model/Adam tensors +stay in memory and move through the live-state transport. +""" +from __future__ import annotations + +import time +import uuid +from pathlib import Path + +from RL_Framework.infra.elastic.runtime_boundary import read_record, write_record + + +def _run_operation(engine, command, trainer=None): + operation = command["operation"] + if operation == "capture": + revision = int(engine.capture_elastic_state_snapshot("boundary", "all")) + return { + "optimizer_revision": revision, + "snapshot_uri": engine.get_elastic_state_snapshot_path(revision), + "lane": engine.get_elastic_lane_state(), + } + if operation == "release": + for revision in command.get("revisions", []): + engine.release_elastic_state_snapshot(int(revision)) + return {} + if operation == "publish": + if trainer is None: + raise RuntimeError("collective publication requires the live trainer") + target = dict(command["target"]) + ids = [instance["instance_id"] for instance in target["instances"]] + if target.get("scope") == "returned_subset": + # An extension rollback only republishes its new incarnations. + # Existing Hybrid loans and the full trainer configuration survive. + missing = object() + previous_targets = getattr(trainer, "_publication_target_ids", missing) + previous_pending = getattr(trainer, "_restored_rollouts_pending_publication", missing) + trainer._publication_target_ids = set(ids) + trainer._restored_rollouts_pending_publication = set(ids) + try: + from .rollout_lifecycle import publish_with_consensus + revision = int(engine.optimizer_revision) + publish_with_consensus(trainer, revision) + finally: + for name, previous in ( + ("_publication_target_ids", previous_targets), + ("_restored_rollouts_pending_publication", previous_pending), + ): + if previous is missing: + delattr(trainer, name) + else: + setattr(trainer, name, previous) + return {"optimizer_revision": revision, "acknowledged": True} + target["unavailable_rollout_ids"] = ids + trainer._restored_rollouts_pending_publication = set(ids) + trainer._apply_runtime_reconfiguration_state(target) + revision = int(engine.optimizer_revision) + from .rollout_lifecycle import publish_with_consensus + publish_with_consensus(trainer, revision) + engine.set_version(revision) + trainer._reopen_restored_rollouts() + trainer._rebind_rollout_engine_from_config(reason="core_boundary_publication") + return {"optimizer_revision": revision, "acknowledged": True} + raise ValueError(f"unknown collective core operation: {operation}") + + +def execute_core_command(engine, command, directory: Path, trainer=None): + rank = int(getattr(engine, "rank", 0)) + receipt = { + "command_id": command["command_id"], "coord_id": command["coord_id"], + "run_id": command["run_id"], "rank": rank, + } + try: + receipt.update(_run_operation(engine, command, trainer)) + receipt["status"] = "success" + except Exception as exc: + receipt.update(status="failed", error=f"{type(exc).__name__}: {exc}") + write_record(directory / "core_receipts" / f"rank_{rank}.json", receipt) + return receipt + + +def follow_core_command(trainer, directory: Path, coord_id: str): + command = read_record(directory / "core_command.json") + if (command.get("coord_id") != coord_id + or not trainer._runtime_state_is_current(command)): + return + handled = getattr(trainer, "_handled_core_commands", set()) + command_id = command.get("command_id") + if not command_id or command_id in handled: + return + receipt = execute_core_command(trainer.train_engine, command, directory, trainer) + handled.add(command_id) + trainer._handled_core_commands = handled + if receipt["status"] != "success": + raise RuntimeError(receipt["error"]) + + +def dispatch_core_command(executor, operation: str, *, target=None, revisions=()): + """Run on every core rank already parked in this transaction boundary.""" + engine = executor.train_engine + directory = executor._coordination_dir() + command = { + "operation": operation, "target": target or {}, + "revisions": list(revisions), "command_id": uuid.uuid4().hex, + "coord_id": executor._active_runtime_coord_id, + "run_id": executor.runtime_run_id, + } + write_record(directory / "core_command.json", command) + execute_core_command(engine, command, directory, getattr(executor, "trainer", None)) + timeout = float(getattr(executor.config.global_resource_planner, "runtime_drain_timeout_s", 3600)) + deadline = time.monotonic() + timeout if timeout > 0 else None + size = int(getattr(engine, "world_size", 1)) + while True: + receipts = [read_record(directory / "core_receipts" / f"rank_{r}.json") for r in range(size)] + matches = [r for rank, r in enumerate(receipts) + if r.get("command_id") == command["command_id"] + and r.get("run_id") == command["run_id"] + and r.get("coord_id") == command["coord_id"] + and r.get("rank") == rank] + errors = [r for r in matches if r.get("status") == "failed"] + if errors: + raise RuntimeError(f"collective core {operation} failed: {errors}") + if len(matches) == size: + return matches + if deadline is not None and time.monotonic() >= deadline: + raise TimeoutError(f"core {operation} missing rank receipts: {len(matches)}/{size}") + if executor._closing.wait(0.02): + raise RuntimeError("runtime executor closed during core state operation") diff --git a/infra/elastic/hybrid_pool.py b/infra/elastic/hybrid_pool.py index d574e3c..25bf407 100644 --- a/infra/elastic/hybrid_pool.py +++ b/infra/elastic/hybrid_pool.py @@ -5,8 +5,8 @@ * the core TP/PP training topology stays fixed; * hybrid workers join/leave only in the inter-replica gradient domain; -* joining workers first contribute zero-gradient placeholders, so core training - can keep stepping while their state is fetched and aligned; +* joining workers contribute no samples until a state-installation ACK proves + that their parameters and optimizer are caught up; * gradient payloads are exchanged through a chunked RDMA-style data plane, with a torch.distributed implementation for CUDA/NCCL or CPU/Gloo tests. """ @@ -292,6 +292,7 @@ def __init__( self._membership_epochs: dict[str, int] = {} self._joining: set[str] = set() self._membership_epoch = 0 + self._consumed_payloads: set[tuple[str, int, int, int, int]] = set() self._lock = threading.RLock() @property @@ -564,12 +565,15 @@ def reduce_core_gradients( targets = dict(self._hybrid_targets) joining = set(self._joining) epochs = dict(self._membership_epochs) + seen = set() accepted = [] for payload in hybrid_payloads: target = targets.get(payload.replica_id) if target not in effective or payload.replica_id in joining: continue if payload.zero_placeholder: + if step >= 0: + raise ValueError("active hybrid gradient cannot be a zero placeholder") continue if target != payload.target_core_id: raise ValueError(f"payload target mismatch for {payload.replica_id}") @@ -577,6 +581,17 @@ def reduce_core_gradients( raise ValueError(f"payload replica world size mismatch for {payload.replica_id}") if payload.replica_rank != self.local_replica_rank: continue + if (step >= 0 or weighted) and payload.membership_epoch != epochs.get(payload.replica_id, 0): + raise ValueError("stale hybrid membership epoch") + if step >= 0 and payload.step != step: + raise ValueError("stale hybrid training step") + if state_version >= 0 and payload.state_version != state_version: + raise ValueError("stale hybrid optimizer revision") + key = (payload.replica_id, payload.replica_rank, payload.step, + payload.state_version, payload.membership_epoch) + if key in seen or (payload.step >= 0 and key in self._consumed_payloads): + raise ValueError("duplicate hybrid gradient payload") + seen.add(key) if len(payload.tensors) != len(effective[target]): raise ValueError("hybrid payload tensor count mismatch") if any(a.shape != b.shape for a, b in zip(effective[target], payload.tensors)): @@ -623,6 +638,11 @@ def reduce_core_gradients( base + extra.to(device=base.device, dtype=base.dtype) * scale for base, extra in zip(effective[target], payload.tensors) ) + with self._lock: + self._consumed_payloads.update(k for k in seen if k[2] >= 0) + if step >= 0: + self._consumed_payloads = {k for k in self._consumed_payloads if k[2] >= step - 1} + # In a real distributed launch each rank owns one core replica lane. The # transport call performs the RDMA/NCCL all-reduce. For local tests we # compute the same reduction explicitly across the in-memory core map. @@ -704,6 +724,7 @@ def __init__( snapshot_fetcher: Callable[[str, str], int] | None = None, zero_sync_steps: int = 1, max_background_workers: int = 0, + require_state_ack: bool = True, ): if zero_sync_steps < 0: raise ValueError("zero_sync_steps cannot be negative") @@ -711,6 +732,7 @@ def __init__( core_rollout_workers = tuple(core_rollout_workers) hybrid_workers = tuple(hybrid_workers) self.zero_sync_steps = zero_sync_steps + self.require_state_ack = bool(require_state_ack) self.snapshot_fetcher = snapshot_fetcher or self._default_snapshot_fetcher self.gradient_domain = gradient_domain or InterReplicaGradientDomain( core_replica_ids=core_train_workers @@ -989,10 +1011,30 @@ def _invoke_activation(self, callback, args, *, identity, generation, replica, d previous = getattr(self._activation_context, "join", None) self._activation_context.join = (identity, generation, replica, deadline) try: - callback(*args) + receipt = callback(*args) + deadline = self._activation_context.join[3] + check = self._assert_replica_join_current if replica else self._assert_join_current + check(identity, generation, deadline) + if self.require_state_ack: + if not isinstance(receipt, Mapping) or receipt.get("acknowledged") is not True: + raise RuntimeError("Hybrid activation requires an actual state installation ACK") + if int(receipt.get("state_version", -1)) < int(args[-1]): + raise RuntimeError("Hybrid activation ACK is older than the fetched state") + if int(receipt.get("optimizer_revision", -1)) != int(receipt["state_version"]): + raise RuntimeError("Hybrid activation optimizer revision ACK mismatch") + return receipt, deadline finally: self._activation_context.join = previous + def handoff_activation_to_boundary(self) -> None: + """Release the cold deadline once the runtime takes over activation. + + Called under activation_guard so the cold deadline and generation are + checked before the request enters the boundary queue. + """ + identity, generation, replica, _ = self._activation_context.join + self._activation_context.join = (identity, generation, replica, None) + @contextmanager def activation_guard(self, handle=None): """Serialize a launch commit with cancellation, never readiness waits. @@ -1024,7 +1066,7 @@ def _finish_join( ) -> ElasticWorker: deadline = ( time.monotonic() + max(float(timeout_s), 0.0) - if timeout_s is not None + if timeout_s is not None and float(timeout_s) > 0 else None ) try: @@ -1050,10 +1092,12 @@ def _finish_join( state_version=version, ) - for _ in range(self.zero_sync_steps): - time.sleep(0) + # State installation is acknowledged by the activation callback; + # yielding the host thread does not synchronize optimizer updates. self._assert_join_current(worker_id, generation, deadline) + if activation_barrier is None and self.require_state_ack: + raise RuntimeError("Hybrid activation requires a state installation barrier") if activation_barrier is not None: with self._lock: self._assert_join_current(worker_id, generation, deadline) @@ -1065,10 +1109,12 @@ def _finish_join( state_version=version, transition_generation=generation, ) - self._invoke_activation( + receipt, deadline = self._invoke_activation( activation_barrier, (worker_id, target_core_id, version), identity=worker_id, generation=generation, replica=False, deadline=deadline, ) + if isinstance(receipt, Mapping): + version = int(receipt["state_version"]) self._assert_join_current(worker_id, generation, deadline) with self._lock: @@ -1125,7 +1171,7 @@ def _finish_replica_join( ) -> tuple[ElasticWorker, ...]: deadline = ( time.monotonic() + max(float(timeout_s), 0.0) - if timeout_s is not None + if timeout_s is not None and float(timeout_s) > 0 else None ) try: @@ -1157,10 +1203,12 @@ def _finish_replica_join( join_state=JoinState.ZERO_SYNC, state_version=version, ) - for _ in range(self.zero_sync_steps): - time.sleep(0) + # State installation is acknowledged by the activation callback; + # yielding the host thread does not synchronize optimizer updates. self._assert_replica_join_current(replica_id, generation, deadline) + if replica_activation_barrier is None and activation_barrier is None and self.require_state_ack: + raise RuntimeError("Hybrid activation requires a state installation barrier") if replica_activation_barrier is not None: with self._lock: self._assert_replica_join_current(replica_id, generation, deadline) @@ -1172,10 +1220,12 @@ def _finish_replica_join( state_version=version, transition_generation=generation, ) - self._invoke_activation( + receipt, deadline = self._invoke_activation( replica_activation_barrier, (replica_id, members, target_core_id, version), identity=replica_id, generation=generation, replica=True, deadline=deadline, ) + if isinstance(receipt, Mapping): + version = int(receipt["state_version"]) self._assert_replica_join_current(replica_id, generation, deadline) elif activation_barrier is not None: with self._lock: @@ -1188,12 +1238,19 @@ def _finish_replica_join( state_version=version, transition_generation=generation, ) + installed_versions = [] for worker_id in members: - self._invoke_activation( + receipt, deadline = self._invoke_activation( activation_barrier, (worker_id, target_core_id, version), identity=replica_id, generation=generation, replica=True, deadline=deadline, ) + if isinstance(receipt, Mapping): + installed_versions.append(int(receipt["state_version"])) self._assert_replica_join_current(replica_id, generation, deadline) + if installed_versions: + if len(set(installed_versions)) != 1: + raise RuntimeError("Hybrid replica members acknowledged different optimizer revisions") + version = installed_versions[0] with self._lock: self._assert_replica_join_current(replica_id, generation, deadline) diff --git a/infra/elastic/process_lifecycle.py b/infra/elastic/process_lifecycle.py new file mode 100644 index 0000000..1b2fe8b --- /dev/null +++ b/infra/elastic/process_lifecycle.py @@ -0,0 +1,62 @@ +"""Stop only owned worker sessions and verify their live members exited.""" +from __future__ import annotations + +import os +from pathlib import Path +import signal +import subprocess +import time + + +def _live_group_members(group): + members = [] + for entry in Path('/proc').iterdir(): + if not entry.name.isdigit(): + continue + try: + fields = (entry / 'stat').read_text().rsplit(')', 1)[1].split() + if int(fields[2]) == group and fields[0] != 'Z': + members.append(int(entry.name)) + except (FileNotFoundError, ProcessLookupError, PermissionError): + continue + return members + + +def stop_local_worker(proc, timeout=5.): + """Use a dedicated process group when owned, retaining legacy mock support.""" + real_process = any(base.__module__ == "subprocess" and base.__name__ == "Popen" + for base in type(proc).__mro__) + try: + group = os.getpgid(proc.pid) if real_process else None + except (ProcessLookupError, AttributeError): + group = None + # All actual runtime launches use start_new_session=True. Never signal a + # shared shell/trainer process group for a legacy subprocess or test double. + owned_group = group == proc.pid or (real_process and group is None and bool(_live_group_members(proc.pid))) + if group is not None and not owned_group and proc.poll() is None: + raise RuntimeError('Hybrid launcher lacks a dedicated owned process group; exit cannot be proved safely') + if owned_group: + group = proc.pid + for sig in (signal.SIGTERM, signal.SIGKILL): + try: + os.killpg(group, sig) + except ProcessLookupError: + break + deadline = time.monotonic() + timeout + while _live_group_members(group) and time.monotonic() < deadline: + time.sleep(.02) + if not _live_group_members(group): + break + if _live_group_members(group): + raise RuntimeError(f'Hybrid process group remains live: {group}') + proc.wait(timeout=timeout) + else: + if proc.poll() is None: + proc.terminate() + try: + proc.wait(timeout=timeout) + except subprocess.TimeoutExpired: + proc.kill() + proc.wait(timeout=timeout) + if proc.poll() is None: + raise RuntimeError('Hybrid launcher has not exited') diff --git a/infra/elastic/readiness.py b/infra/elastic/readiness.py new file mode 100644 index 0000000..c4cd3cd --- /dev/null +++ b/infra/elastic/readiness.py @@ -0,0 +1,82 @@ +"""Validate actual per-model-lane EHP installation receipts.""" +from __future__ import annotations + +import json +from pathlib import Path + + +def snapshot_contract(snapshot_path, *, state_version=None, target_core_id=None): + if not snapshot_path: + raise ValueError('Hybrid readiness requires a captured state identity') + if str(snapshot_path).startswith('ehp://'): + if state_version is None: + raise ValueError('direct live-state readiness requires its optimizer revision') + return {'snapshot_path': str(snapshot_path), 'snapshot_id': None, + 'optimizer_revision': int(state_version), 'target_core_id': target_core_id} + manifest = json.loads(Path(snapshot_path).read_text(encoding='utf-8')) + if manifest.get('format') != 'libra-ehp-lanes-v1' or not manifest.get('snapshot_id'): + raise ValueError('Hybrid readiness requires a live lane manifest') + revision = int(manifest['optimizer_revision']) + if state_version is not None and revision != int(state_version): + raise ValueError('Hybrid readiness revision differs from its snapshot manifest') + if target_core_id is not None and target_core_id not in manifest['lanes']: + raise ValueError('Hybrid target core is missing from its state manifest') + return {'snapshot_path': str(snapshot_path), 'snapshot_id': str(manifest['snapshot_id']), + 'optimizer_revision': revision, 'target_core_id': target_core_id, + 'lanes': manifest['lanes']} + + +def read_ready_receipts(task_dir, worker_id, world_size, contract): + """Return a receipt only when every actual TP worker matches the contract. + + Old and partially written files are pending observations, never readiness. + The process owner removes same-snapshot files before each new launch. + """ + receipts = [] + revision = contract['optimizer_revision'] + for rank in range(int(world_size)): + path = Path(task_dir) / f'{worker_id}.rank_{rank}.ready' + try: + rows = path.read_text(encoding='utf-8').splitlines() + receipt = dict(row.split('=', 1) for row in rows if '=' in row) + if receipt.get('acknowledged') != 'True': + return None + if int(receipt.get('rank', -1)) != rank or int(receipt.get('pid', -1)) <= 0: + return None + if receipt.get('snapshot') != contract['snapshot_path'] or not receipt.get('snapshot_id'): + return None + expected_identity = contract.get('snapshot_id') + if expected_identity is not None and receipt['snapshot_id'] != expected_identity: + return None + if int(receipt.get('optimizer_revision', -1)) != revision or int(receipt.get('state_version', -1)) != revision: + return None + target = contract.get('target_core_id') + if target is not None and receipt.get('target_core_id') != target: + return None + if expected_identity is not None: + if not receipt.get('lane_snapshot_id'): + return None + lane_core = target or receipt.get('target_core_id') + expected_uri = contract['lanes'].get(lane_core, {}).get(str(rank)) + if not expected_uri or receipt.get('source_uri') != expected_uri: + return None + except (OSError, ValueError, TypeError): + return None + receipts.append(receipt) + if not receipts or len({receipt['pid'] for receipt in receipts}) != len(receipts): + return None + if len({receipt['snapshot_id'] for receipt in receipts}) != 1: + return None + targets = {receipt.get('target_core_id') for receipt in receipts} + if len(targets) != 1: + return None + return {'acknowledged': True, 'snapshot_id': receipts[0]['snapshot_id'], + 'state_version': revision, 'optimizer_revision': revision, + 'target_core_id': receipts[0].get('target_core_id'), 'rank_receipts': receipts} + + +def clear_ready_receipts(task_dir, worker_id): + directory = Path(task_dir) + (directory / f'{worker_id}.ready').unlink(missing_ok=True) + for path in directory.glob(f'{worker_id}.rank_*.ready'): + path.unlink(missing_ok=True) diff --git a/infra/elastic/rollout_lending.py b/infra/elastic/rollout_lending.py new file mode 100644 index 0000000..a4bdded --- /dev/null +++ b/infra/elastic/rollout_lending.py @@ -0,0 +1,231 @@ +"""Lend complete live rollout engines to EHP without applying a GRP layout.""" +from __future__ import annotations + +from contextlib import nullcontext +from copy import deepcopy +from dataclasses import replace +from itertools import combinations +import json +from pathlib import Path +import time +import uuid + + +def _set_unavailable(executor, instance_ids): + scheduler = getattr(getattr(executor, "rollout_engine", None), "scheduler", None) + if scheduler is None: + return + with getattr(scheduler, "_lock", nullcontext()): + for handle in getattr(scheduler, "_instances", []): + if handle.instance_id in instance_ids: + handle.is_ready = False + + +def _capacity(physical, width): + by_host = {} + for host, gpu in physical: + by_host.setdefault(host, set()).add(gpu) + return sum(len(gpus) // width * width for gpus in by_host.values()) + + +def _choose_engines(executor, required): + loan = getattr(executor, "_lent_rollout", None) or {} + old_ids = {meta.instance_id for meta in loan.get("engines", [])} + metas = sorted((meta for meta in executor._process_meta.values() + if meta.instance_id not in old_ids), key=lambda meta: meta.instance_id) + width = executor._elastic_replica_size_gpus() + reserve = max(0, int(getattr(executor.config.global_resource_planner, + "elastic_hybrid_min_rollout_gpus", 0) or 0)) + serving = {(meta.host, int(gpu)) for meta in metas for gpu in meta.gpus} + existing = {(slot["host"], int(gpu)) for slot in loan.get("slots", {}).values() + for gpu in slot["gpus"]} + candidates = [] + for count in range(1, len(metas) + 1): + for subset in combinations(metas, count): + physical = {(meta.host, int(gpu)) for meta in subset for gpu in meta.gpus} + if len(serving - physical) < reserve: + continue + if _capacity(existing | physical, width) >= required: + candidates.append((len(physical), count, subset)) + return min(candidates, key=lambda item: item[:2])[2] if candidates else () + + +def _restart_engines(executor, saved): + restored = saved["restored"] + for meta in saved["engines"]: + if meta.instance_id not in saved["stopped"]: + continue + started = next((item for item in restored if item.instance_id == meta.instance_id), None) + if started is None: + command = executor._command_for_instance_id(meta.instance_id, meta.port) + target = replace(meta, command=command, pid=0, adopted=False) + started = executor._start_one_rollout_process(target) + # A readiness failure must not lose ownership of an already started + # process or launch another process on its GPU when recovery retries. + restored.append(started) + started.incarnation_id = uuid.uuid4().hex + started.service_generation = None + executor._wait_rollout_process_start_ready(started) + return list(restored) + + +def _rollback_expansion(executor): + pending = getattr(executor, "_rollout_loan_rollback", None) + if pending is None: + return + # A failed stop may have signalled a process without proving its exit. + # Resolve that same ownership before launching replacement CUDA workers. + for identity in pending["uncertain"][:]: + stopped = executor._stop_rollout_process(identity) + if not stopped: + raise RuntimeError(f"rollout stop did not confirm ownership: {identity}") + pending["stopped"].append(identity) + pending["uncertain"].remove(identity) + restored = _restart_engines(executor, pending) + if restored: + executor._publish_current_rollout_weights( + instance_ids=[meta.instance_id for meta in restored]) + scheduler = getattr(getattr(executor, "rollout_engine", None), "scheduler", None) + # The subset publisher only installs weights. Routing is restored here, + # after its current-revision ACK; untouched engines keep prior readiness. + for handle in getattr(scheduler, "_instances", []): + if handle.instance_id in pending["readiness"]: + handle.is_ready = pending["readiness"][handle.instance_id] + executor._rollout_loan_rollback = None + + +def lend_rollout(executor, required_gpu_count): + """Grow the loan to a total GPU target without replacing active Hybrids. + + Only newly selected whole engines are stopped. Their physical slots become + eligible for training together, after every new stop is acknowledged. + """ + required = int(required_gpu_count) + previous = getattr(executor, "_lent_rollout", None) + existing = previous["slots"] if previous else {} + physical = {(slot["host"], int(gpu)) for slot in existing.values() for gpu in slot["gpus"]} + if required < 1 or _capacity(physical, executor._elastic_replica_size_gpus()) >= required: + return deepcopy(existing) + if getattr(executor, "_rollout_loan_rollback", None): + raise RuntimeError("an incomplete rollout loan expansion must be recovered before retry") + executor._adopt_existing_rollout_processes() + selected = _choose_engines(executor, required) + if not selected: + return deepcopy(existing) + # Freeze the namespace against the complete physical rollout allocation at + # the first loan. Removing an engine must not renumber later worker slots. + namespace = previous["physical_slots"] if previous else { + address: f"rollout{index}" for index, address in enumerate(sorted({ + (meta.host, int(gpu)) for meta in executor._process_meta.values() for gpu in meta.gpus}))} + additional = {(meta.host, int(gpu)) for meta in selected for gpu in meta.gpus} + slots = {namespace[address]: {"host": address[0], "gpus": [address[1]]} + for address in sorted(additional)} + scheduler = getattr(getattr(executor, "rollout_engine", None), "scheduler", None) + selected_ids = {meta.instance_id for meta in selected} + pending = {"engines": [deepcopy(meta) for meta in selected], "stopped": [], + "uncertain": [], "restored": [], + "readiness": {handle.instance_id: handle.is_ready + for handle in getattr(scheduler, "_instances", []) + if handle.instance_id in selected_ids}} + executor._rollout_loan_rollback = pending + dispatcher = getattr(executor, "dispatcher", None) + paused_here = dispatcher is not None and not dispatcher.is_paused() + if paused_here: + dispatcher.pause() + timeout = float(getattr(executor.config.global_resource_planner, "runtime_drain_timeout_s", 3600.)) + try: + if dispatcher is not None: + dispatcher.wait_until_idle(timeout=timeout) + engine = getattr(executor, "rollout_engine", None) + if engine is not None and hasattr(engine, "wait_until_idle"): + engine.wait_until_idle(timeout=timeout) + _set_unavailable(executor, {meta.instance_id for meta in selected}) + for meta in selected: + pending["uncertain"].append(meta.instance_id) + stopped = executor._stop_rollout_process(meta.instance_id) + if not stopped: + raise RuntimeError(f"rollout stop did not confirm ownership: {meta.instance_id}") + pending["stopped"].append(meta.instance_id) + pending["uncertain"].remove(meta.instance_id) + saved = previous if previous is not None else { + "engines": [], "slots": {}, "stopped": [], "restored": [], + "generation": uuid.uuid4().hex, "physical_slots": namespace} + saved["engines"].extend(pending["engines"]) + saved["stopped"].extend(pending["stopped"]) + saved["slots"].update(slots) + executor._lent_rollout = saved + executor._rollout_loan_rollback = None + return deepcopy(saved["slots"]) + except Exception as error: + try: + _rollback_expansion(executor) + except Exception as rollback_error: + raise RuntimeError(f"rollout loan rollback incomplete: {rollback_error}") from error + raise + finally: + if paused_here: + scheduler = getattr(getattr(executor, "rollout_engine", None), "scheduler", None) + if any(handle.is_ready for handle in getattr(scheduler, "_instances", [])): + dispatcher.resume() + else: + saved = getattr(executor, "_lent_rollout", None) + if saved is not None: + saved["resume_dispatcher_after_publication"] = True + + +def wait_hybrid_stopped(executor, worker_id, request_id=None): + """Await the controller's proof that the recorded remote process exited.""" + directory = Path(executor._hybrid_worker_control_dir()) + timeout = float(getattr(executor.config.global_resource_planner, "hybrid_worker_stop_timeout_s", 60.)) + deadline = time.monotonic() + timeout if timeout > 0 else float("inf") + while True: + done = directory / f".stop_{worker_id}.done" + error = directory / f".stop_{worker_id}.error" + if error.exists(): + record = json.loads(error.read_text()) + if request_id is None or record.get("request_id") == request_id: + raise RuntimeError(f"hybrid stop failed: {record.get('error')}") + if done.exists(): + record = json.loads(done.read_text()) + if ((request_id is None or record.get("request_id") == request_id) + and record.get("worker_id") == worker_id and record.get("exited") is True): + return record + if time.monotonic() >= deadline: + raise TimeoutError(f"hybrid exit not acknowledged: {worker_id}") + time.sleep(.05) + + +def return_rollout(executor): + """Restart borrowed allocations only after Hybrid stop acknowledgements. + + Returned engines remain unavailable until the caller publishes the current + revision to the new service generation and explicitly reopens admission. + """ + _rollback_expansion(executor) + saved = getattr(executor, "_lent_rollout", None) + if not saved: + return [] + physical = {(meta.host, int(gpu)) for meta in saved["engines"] for gpu in meta.gpus} + for worker_id, meta in list(getattr(executor, "_hybrid_worker_meta", {}).items()): + if not any((meta.host, int(gpu)) in physical for gpu in meta.gpus): + continue + # The executor retains ownership until a matching remote exit ACK (or + # a local process wait) completes. It also withdraws replica aliases. + if worker_id in executor._hybrid_worker_meta: + executor.wait_for_hybrid_training_tasks(worker_id) + executor._stop_hybrid_worker_process(worker_id) + # A preceding pool release may already have withdrawn the local metadata. + directory = Path(executor._hybrid_worker_control_dir()) + for worker_id in saved["slots"]: + for suffix in ("json", "running"): + pending = directory / f".stop_{worker_id}.{suffix}" + try: + request = json.loads(pending.read_text()) + except FileNotFoundError: + continue + wait_hybrid_stopped(executor, worker_id, request.get("request_id")) + _set_unavailable(executor, {meta.instance_id for meta in saved["engines"]}) + restored = _restart_engines(executor, saved) + executor._rollout_lending_resume_pending = bool(saved.get("resume_dispatcher_after_publication", False)) + executor._lent_rollout = None + return restored diff --git a/infra/elastic/rollout_lifecycle.py b/infra/elastic/rollout_lifecycle.py new file mode 100644 index 0000000..0865170 --- /dev/null +++ b/infra/elastic/rollout_lifecycle.py @@ -0,0 +1,126 @@ +"""Training-boundary coordination for rollout GPUs lent to EHP.""" +from __future__ import annotations + +import torch.distributed as dist + + +def _broadcast(trainer, value): + packet = [value] + if dist.is_available() and dist.is_initialized(): + dist.broadcast_object_list(packet, src=0, group=getattr(trainer, "control", None)) + if packet[0].get("error"): + raise RuntimeError(packet[0]["error"]) + return packet[0] + + +def unavailable_rollouts(executor): + loan = getattr(executor, "_lent_rollout", None) + pending = getattr(executor, "_rollout_loan_rollback", None) + return list(dict.fromkeys(meta.instance_id for record in (loan, pending) if record + for meta in record["engines"])) + + +def apply_availability(trainer, unavailable): + previous = set(getattr(trainer, "_rollout_unavailable_ids", ())) + trainer._rollout_unavailable_ids = set(unavailable) + scheduler = getattr(trainer.rollout_engine, "scheduler", None) + for handle in getattr(scheduler, "_instances", []): + if handle.instance_id in trainer._rollout_unavailable_ids: + handle.is_ready = False + elif handle.instance_id in previous: + handle.is_ready = True + + +def restore_for_publication(trainer): + """Return physical GPUs on rank zero; every rank withdraws old membership. + + The caller has closed admission on all logical sources. Restarted services + remain unroutable until publication has validated their new worker receipts. + """ + outcome = None + if trainer.is_main_process: + executor = trainer.runtime_elastic_executor + try: + ids = unavailable_rollouts(executor) + if ids: + from .rollout_lending import return_rollout + executor._release_all_hybrid_workers(reason="restore_rollout_supply") + return_rollout(executor) + outcome = {"restored": ids} + except Exception as exc: + outcome = {"error": f"failed to restore rollout GPUs: {exc}"} + outcome = _broadcast(trainer, outcome) + if outcome["restored"]: + apply_availability(trainer, outcome["restored"]) + trainer._restored_rollouts_pending_publication = set(outcome["restored"]) + trainer._refresh_nonblocking_elastic_membership() + trainer.train_engine._elastic_step_hybrid_workers = () + return bool(outcome["restored"]) + + +def reopen_after_publication(trainer): + # A subset rollback has a separate owner: its publication must neither + # withdraw an existing loan nor open routing before rank consensus. + if getattr(trainer, "_publication_target_ids", None) is not None: + return + if not getattr(trainer, "_restored_rollouts_pending_publication", None): + return + apply_availability(trainer, ()) + trainer._restored_rollouts_pending_publication = set() + executor = getattr(trainer, "runtime_elastic_executor", None) + if executor is not None: + executor._rollout_lending_resume_pending = False + + +def publish_with_consensus(trainer, revision): + error = "" + try: + trainer._sync_weights(target_version=revision) + except Exception as exc: + error = f"{type(exc).__name__}: {exc}" + errors = [error] + if dist.is_available() and dist.is_initialized(): + errors = [None] * dist.get_world_size() + dist.all_gather_object(errors, error, group=getattr(trainer, "control", None)) + if any(errors): + raise RuntimeError(f"Rollout publication failed across core ranks: {errors}") + + +def ensure_batch_supply(trainer): + """Keep the loan while ready groups suffice; return it before a supply wait.""" + cfg = getattr(trainer.config, "global_resource_planner", None) + if not (getattr(cfg, "enabled", False) + and getattr(cfg, "runtime_reconfigure_training", False) + and getattr(cfg, "hybrid_worker_launch_enabled", False) + and not getattr(cfg, "runtime_training_pool_plan_only", True)): + return + outcome = None + if trainer.is_main_process: + executor = trainer.runtime_elastic_executor + ids = unavailable_rollouts(executor) + instances = getattr(trainer.config.heterogeneous_rollout, "instances", ()) + outcome = {"all_lent": bool(ids) and len(ids) == len(instances)} + outcome = _broadcast(trainer, outcome) + if not outcome["all_lent"]: + return + from .rollout_supply import ready_complete_samples + sufficient = (not trainer._is_rollout_source() or + ready_complete_samples(trainer) >= trainer._source_batch_sample_count()) + records = [sufficient] + if dist.is_available() and dist.is_initialized(): + records = [None] * dist.get_world_size() + dist.all_gather_object(records, sufficient, group=getattr(trainer, "control", None)) + if all(records): + return + trainer._pause_rollout_pipeline() + # Publication implementations must invalidate old service receipts before + # a loan is returned, even when the replacement keeps the same endpoints. + restore = getattr(trainer, "_restore_lent_rollouts_for_publication", None) + restored = restore() if callable(restore) else restore_for_publication(trainer) + if restored: + revision = int(trainer.train_engine.optimizer_revision) + publish_with_consensus(trainer, revision) + trainer.train_engine.set_version(revision) + reopen_after_publication(trainer) + trainer._rebind_rollout_engine_from_config(reason="ehp_supply_restore") + trainer._resume_rollout_pipeline() diff --git a/infra/elastic/rollout_process.py b/infra/elastic/rollout_process.py new file mode 100644 index 0000000..4f52256 --- /dev/null +++ b/infra/elastic/rollout_process.py @@ -0,0 +1,410 @@ +"""Place rollout restarts inside the current allocation and stop only owned work.""" +from __future__ import annotations + +import argparse +import ipaddress +import json +import os +from pathlib import Path +import shlex +import signal +import socket +import subprocess +import sys +import time +import uuid + + +def is_local_host(host, env=None): + env = os.environ if env is None else env + host = str(host or '') + try: + address = ipaddress.ip_address(host) + if address.is_loopback or address.is_unspecified: + return True + except ValueError: + pass + names = {'', 'localhost', socket.gethostname(), socket.gethostname().split('.', 1)[0], + str(env.get('SLURMD_NODENAME', ''))} + return host in names or host.split('.', 1)[0] in {name.split('.', 1)[0] for name in names if name} + + +def allocation_nodes(env): + expression = env.get('SLURM_JOB_NODELIST') or env.get('SLURM_NODELIST') + if not env.get('SLURM_JOB_ID') or not expression: + raise RuntimeError('Remote rollout launch requires the current Slurm allocation and node list') + result = subprocess.run(['scontrol', 'show', 'hostnames', expression], text=True, + capture_output=True, check=True, env=env) + nodes = result.stdout.split() + if not nodes: + raise RuntimeError('The current Slurm allocation contains no nodes') + return nodes + + +def target_node(host, env, nodes=None): + nodes = allocation_nodes(env) if nodes is None else nodes + exact = [node for node in nodes if node == host or node.split('.', 1)[0] == host.split('.', 1)[0]] + if not exact: + try: + address = str(ipaddress.ip_address(host)) + exact = [node for node in nodes if address in { + result[4][0] for result in socket.getaddrinfo(node, None)}] + except (ValueError, OSError): + pass + if len(exact) != 1: + raise RuntimeError(f'Rollout host {host!r} is not one unambiguous node in the current Slurm allocation') + return exact[0] + + +def slurm_command(command, host, gpus, env, *, step_name=None): + """Use allocated node GPUs; never request an arbitrary replacement subset.""" + node = target_node(host, env) + ids = ','.join(str(gpu) for gpu in gpus) + if not gpus or len(set(gpus)) != len(gpus) or min(gpus) < 0: + raise ValueError('Expected distinct assigned node GPU IDs') + # With --overlap the step inherits the job's existing GPU allocation on + # this node. Asking for --gpus=N could select a different physical subset. + argv = ['srun', '--overlap', '--jobid', env['SLURM_JOB_ID'], '--nodes=1', '--ntasks=1', + '--nodelist', node, '--cpus-per-task=4', '--cpu-bind=none'] + if step_name: + argv += ['--job-name', step_name] + argv += ['env', 'CUDA_VISIBLE_DEVICES=' + ids, 'PLANNED_CUDA_VISIBLE_DEVICES=' + ids, + 'bash', '-c', command] + return shlex.join(argv) + + +def placement_command(command, host, gpus, env): + return command if is_local_host(host, env) else slurm_command(command, host, gpus, env) + + +def _write(path, value): + path = Path(path) + path.parent.mkdir(parents=True, exist_ok=True) + temporary = path.with_suffix('.tmp.' + str(os.getpid())) + temporary.write_text(json.dumps(value)) + temporary.replace(path) + + +def _process_record(pid): + try: + fields = (Path('/proc') / str(pid) / 'stat').read_text().rsplit(')', 1)[1].split() + return dict(pid=int(pid), parent=int(fields[1]), group=int(fields[2]), + state=fields[0], start=fields[19]) + except (FileNotFoundError, ProcessLookupError, PermissionError): + return None + + +def _owned_descendants(pid): + records = {int(path.name): _process_record(int(path.name)) for path in Path('/proc').iterdir() + if path.name.isdigit()} + owned, frontier = {}, [pid] + while frontier: + current = frontier.pop() + record = records.get(current) + if record is not None and current not in owned: + owned[current] = record + frontier.extend(key for key, child in records.items() if child and child['parent'] == current) + return owned + + +def _still_owned(record): + current = _process_record(record['pid']) + return bool(current and current['start'] == record['start'] and current['state'] != 'Z') + + +def stop_owned_session(proc, timeout=30.): + """Include owned descendant sessions; a supervisor may start a new session.""" + real = any(base.__module__ == 'subprocess' and base.__name__ == 'Popen' for base in type(proc).__mro__) + if not real: + from .process_lifecycle import stop_local_worker + return stop_local_worker(proc, timeout=timeout) + records = _owned_descendants(proc.pid) + root = records.get(proc.pid) + if root is not None and root['group'] != proc.pid and proc.poll() is None: + raise RuntimeError('Rollout launcher does not own a dedicated process session') + for sig in (signal.SIGTERM, signal.SIGKILL): + # Capture all descendant sessions before removing the supervisor that + # otherwise supplies their parent relationship. PID birth times fence + # unrelated processes if an identifier is reused during cleanup. + records.update(_owned_descendants(proc.pid)) + for record in reversed(list(records.values())): + if _still_owned(record): + try: + os.kill(record['pid'], sig) + except ProcessLookupError: + pass + deadline = time.monotonic() + timeout if timeout > 0 else None + while any(_still_owned(record) for record in records.values()) and (deadline is None or time.monotonic() < deadline): + time.sleep(.02) + if not any(_still_owned(record) for record in records.values()): + proc.wait(timeout=max(timeout, .1) if timeout > 0 else None) + return + raise RuntimeError('Owned rollout descendant processes remain live after termination') + + +def launch_rollout_process(meta, env, stdout=None, control_dir=None): + remote = not is_local_host(meta.host, env) + token = uuid.uuid4().hex + owner = dict(token=token, instance_id=meta.instance_id, host=meta.host, gpus=list(meta.gpus), remote=remote, + job_id=env.get('SLURM_JOB_ID'), step_name='libroll-' + token[:16]) + command = meta.command + if remote: + root = Path(control_dir or env.get('LIBRA_ROLLOUT_PROCESS_DIR') or '') + if not str(root) or str(root) == '.': + raise ValueError('Remote rollout ownership requires a shared project control directory') + path = root / 'process-owners' / (token + '.json') + owner['state_path'] = str(path) + argv = [sys.executable, '-m', 'RL_Framework.infra.elastic.rollout_process', + '--state', str(path), '--token', token, '--', 'bash', '-c', command] + command = slurm_command(shlex.join(argv), meta.host, meta.gpus, env, step_name=owner['step_name']) + proc = subprocess.Popen(command, shell=True, env=env, stdout=stdout, + stderr=subprocess.STDOUT if stdout is not None else None, start_new_session=True) + proc._libra_rollout_owner = owner + return proc + + +def _owned_remote_step(owner, env): + path = Path(owner['state_path']) + if path.exists(): + record = json.loads(path.read_text()) + if record.get('token') != owner['token'] or str(record.get('job_id')) != str(owner['job_id']): + raise RuntimeError('Remote rollout ownership receipt does not match this launcher') + step = record.get('step_id') + if not str(step).isdigit(): + raise RuntimeError('Remote rollout receipt lacks an actual Slurm step ID') + return str(owner['job_id']) + '.' + str(step), record + # srun may still be waiting to start. Its unique step name lets cancellation + # remain scoped even before the remote command publishes its first receipt. + result = subprocess.run(['squeue', '--steps', '--jobs', str(owner['job_id']), '--noheader', + '--format=%i|%j'], check=True, capture_output=True, text=True, env=env) + matches = [line.split('|', 1)[0].strip() for line in result.stdout.splitlines() + if '|' in line and line.split('|', 1)[1].strip() == owner['step_name']] + if len(matches) > 1: + raise RuntimeError('Remote rollout owner maps to more than one Slurm step') + return (matches[0] if matches else None), None + + +def stop_rollout_process(proc, meta=None, timeout=30., control_dir=None): + owner = getattr(proc, '_libra_rollout_owner', None) + if not owner or not owner['remote']: + return stop_owned_session(proc, timeout) + env = os.environ.copy() + if str(env.get('SLURM_JOB_ID')) != str(owner['job_id']): + raise RuntimeError('Cannot stop a rollout from another Slurm allocation') + step, record = _owned_remote_step(owner, env) + if record and record.get('status') == 'exited' and record.get('owned_processes_exited') is True: + return stop_owned_session(proc, timeout) + if step: + subprocess.run(['scancel', '--signal=TERM', step], check=True, env=env, capture_output=True, text=True) + else: + # Cancel a pending srun before it can create a remote step. + stop_owned_session(proc, timeout) + deadline = time.monotonic() + timeout if timeout > 0 else None + while deadline is None or time.monotonic() < deadline: + _, record = _owned_remote_step(owner, env) + if record and record.get('status') == 'exited' and record.get('owned_processes_exited') is True: + stop_owned_session(proc, timeout) + return + query = subprocess.run(['squeue', '--steps', '--jobs', str(owner['job_id']), '--noheader', + '--format=%i|%j'], check=True, capture_output=True, text=True, env=env) + if not any('|' in line and line.split('|', 1)[1].strip() == owner['step_name'] + for line in query.stdout.splitlines()) and proc.poll() is not None: + return + time.sleep(.05) + if step: + subprocess.run(['scancel', '--signal=KILL', step], check=True, env=env, capture_output=True, text=True) + raise RuntimeError('Remote rollout Slurm step exit was not verified; retaining ownership') + + + +def _process_command(pid): + root = Path('/proc') / str(pid) + try: + if root.stat().st_uid != os.getuid(): + raise RuntimeError('Adopted rollout PID belongs to another user') + argv = [part.decode() for part in (root / 'cmdline').read_bytes().split(b'\0') if part] + environment = dict(part.decode().split('=', 1) for part in (root / 'environ').read_bytes().split(b'\0') + if b'=' in part) + except (OSError, UnicodeError) as exc: + raise RuntimeError('Cannot verify adopted rollout PID identity') from exc + return argv, environment + + +def _argument_value(argv, flag): + for index, argument in enumerate(argv): + if argument == flag and index + 1 < len(argv): + return argv[index + 1] + if argument.startswith(flag + '='): + return argument.split('=', 1)[1] + return None + + +def _flatten_shell_arguments(argv): + values = list(argv) + for argument in argv: + if '--instance-id' in argument or '--port' in argument: + try: + values.extend(shlex.split(argument)) + except ValueError: + pass + return values + + +def _verify_adopted_identity(pid, contract, *, launcher=False): + record = _process_record(pid) + if record is None or record['state'] == 'Z': + raise RuntimeError('Adopted rollout PID is absent; its ownership cannot be proved') + argv, environment = _process_command(pid) + if str(environment.get('SLURM_JOB_ID') or environment.get('SLURM_JOBID')) != str(contract['job_id']): + raise RuntimeError('Adopted rollout PID belongs to a different Slurm allocation') + flattened = _flatten_shell_arguments(argv) + supervisor = any(value in {'RL_Framework.scripts.restartable_vllm_server'} or + value.endswith('/restartable_vllm_server.py') for value in argv) + server = any(value in {'RL_Framework.scripts.vllm_hot_reload_api_server', + 'vllm.entrypoints.openai.api_server'} + or value.endswith('/vllm_hot_reload_api_server.py') for value in argv) + if launcher: + if not argv or Path(argv[0]).name != 'srun': + raise RuntimeError('PID is not the expected local Slurm rollout launcher') + node = _argument_value(flattened, '--nodelist') or _argument_value(flattened, '-w') + if not node or node.split('.', 1)[0] != str(contract['host']).split('.', 1)[0]: + raise RuntimeError('Adopted srun launcher targets a different rollout node') + elif not supervisor and not server: + raise RuntimeError('Adopted PID is not an expected vLLM service or supervisor') + if (supervisor or launcher) and _argument_value(flattened, '--instance-id') != contract['instance_id']: + raise RuntimeError('Adopted rollout instance identity does not match') + port = _argument_value(flattened, '--port') + health = _argument_value(flattened, '--health-url') + if str(port) != str(contract['port']) and not (health and health.endswith(':' + str(contract['port']) + '/health')): + raise RuntimeError('Adopted rollout port does not match') + control = contract.get('control_dir') + if control and supervisor and _argument_value(flattened, '--control-dir') != control: + raise RuntimeError('Adopted rollout belongs to a different run control directory') + return record + + +def _stop_verified_pid_tree(pid, expected_start, timeout): + records = _owned_descendants(pid) + root = records.get(pid) + if root is None or root['start'] != expected_start: + raise RuntimeError('Adopted rollout PID changed before termination') + for sig in (signal.SIGTERM, signal.SIGKILL): + current = _process_record(pid) + if current and current['start'] == expected_start: + records.update(_owned_descendants(pid)) + for record in reversed(list(records.values())): + if _still_owned(record): + try: + os.kill(record['pid'], sig) + except ProcessLookupError: + pass + deadline = time.monotonic() + timeout if timeout > 0 else None + while any(_still_owned(record) for record in records.values()) and (deadline is None or time.monotonic() < deadline): + time.sleep(.02) + if not any(_still_owned(record) for record in records.values()): + return dict(pid=pid, process_start=expected_start, owned_processes_exited=True) + raise RuntimeError('Adopted rollout process tree remains live after termination') + + +def _stop_adopted_on_host(contract, timeout): + """Run only on the host owning the service; never interpret its PID locally.""" + if str(os.environ.get('SLURM_JOB_ID') or os.environ.get('SLURM_JOBID')) != str(contract['job_id']): + raise RuntimeError('Adopted rollout stop is outside its Slurm allocation') + if contract.get('discover_service'): + candidates = [] + for path in Path('/proc').iterdir(): + if not path.name.isdigit(): + continue + try: + record = _verify_adopted_identity(int(path.name), contract) + except RuntimeError: + continue + candidates.append(record) + ids = {record['pid'] for record in candidates} + roots = [record for record in candidates if record['parent'] not in ids] + # A server is usually a supervisor's child. Intermediate shell parents + # still leave one root supervisor; multiple independent roots are unsafe. + if len(roots) != 1: + raise RuntimeError('Known launcher does not identify one unique owned rollout service tree') + record = roots[0] + else: + record = _verify_adopted_identity(int(contract['pid']), contract) + result = _stop_verified_pid_tree(record['pid'], record['start'], timeout) + return dict(result, instance_id=contract['instance_id'], port=contract['port'], job_id=contract['job_id']) + + +def stop_adopted_rollout(meta, env=None, timeout=30., control_dir=None): + """Validate host, role, run and job ownership before stopping an adopted PID. + + Old registries sometimes store the controller's local srun PID beside the + remote service host. Only a verified launcher permits lookup of the unique + matching service on that remote node; a bare PID never authorizes a scan. + """ + env = dict(os.environ if env is None else env) + job = env.get('SLURM_JOB_ID') or env.get('SLURM_JOBID') + if not job or int(meta.pid) <= 0: + raise RuntimeError('Adopted rollout stop requires a known PID and current Slurm allocation') + contract = dict(pid=int(meta.pid), host=meta.host, instance_id=meta.instance_id, + port=int(meta.port), job_id=str(job), control_dir=str(control_dir) if control_dir else None) + if is_local_host(meta.host, env): + return _stop_adopted_on_host(contract, timeout) + target_node(meta.host, env) # Prove allocation ownership before touching any PID. + local_launcher = None + try: + local_launcher = _verify_adopted_identity(int(meta.pid), contract, launcher=True) + except RuntimeError: + pass + if local_launcher is not None: + contract['discover_service'] = True + argv = [sys.executable, '-m', 'RL_Framework.infra.elastic.rollout_process', + '--stop-adopted', json.dumps(contract, separators=(',', ':')), '--timeout', str(timeout)] + command = slurm_command(shlex.join(argv), meta.host, meta.gpus, env) + result = subprocess.run(shlex.split(command), env=env, check=True, capture_output=True, text=True, + timeout=3*timeout+30 if timeout > 0 else None) + try: + receipt = json.loads(result.stdout.strip().splitlines()[-1]) + except (ValueError, IndexError) as exc: + raise RuntimeError('Remote adopted rollout stop did not provide an exit receipt') from exc + if (receipt.get('owned_processes_exited') is not True or receipt.get('job_id') != str(job) + or receipt.get('instance_id') != meta.instance_id or receipt.get('port') != int(meta.port)): + raise RuntimeError('Remote adopted rollout exit receipt differs from the requested owner') + if local_launcher is not None and _still_owned(local_launcher): + _stop_verified_pid_tree(local_launcher['pid'], local_launcher['start'], timeout) + return receipt + +def main(argv=None): + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument('--state', type=Path) + parser.add_argument('--token') + parser.add_argument('--stop-adopted', type=json.loads) + parser.add_argument('--timeout', type=float, default=30.) + parser.add_argument('command', nargs=argparse.REMAINDER) + args = parser.parse_args(argv) + if args.stop_adopted is not None: + result = _stop_adopted_on_host(args.stop_adopted, args.timeout) + print(json.dumps(result), flush=True) + return 0 + if args.command and args.command[0] == '--': + args.command = args.command[1:] + if not args.state or not args.token or not args.command or not os.environ.get('SLURM_JOB_ID') or not os.environ.get('SLURM_STEP_ID', '').isdigit(): + raise RuntimeError('Owned remote rollout runner requires a real Slurm job step') + owner = dict(token=args.token, job_id=os.environ['SLURM_JOB_ID'], step_id=os.environ['SLURM_STEP_ID'], + host=socket.gethostname(), status='starting') + _write(args.state, owner) + def terminate(signum, frame): + raise KeyboardInterrupt(signum) + signal.signal(signal.SIGTERM, terminate) + signal.signal(signal.SIGINT, terminate) + proc = None + try: + proc = subprocess.Popen(args.command, start_new_session=True) + _write(args.state, dict(owner, status='running', child_pid=proc.pid)) + return proc.wait() + finally: + if proc is not None: + stop_owned_session(proc) + _write(args.state, dict(owner, status='exited', owned_processes_exited=True)) + + +if __name__ == '__main__': + sys.exit(main()) diff --git a/infra/elastic/rollout_supply.py b/infra/elastic/rollout_supply.py new file mode 100644 index 0000000..5e3e4eb --- /dev/null +++ b/infra/elastic/rollout_supply.py @@ -0,0 +1,25 @@ +"""Inspect already generated GRPO supply without consuming or reordering it.""" +from __future__ import annotations + + +from .batch_partition import unique_group_members + +def ready_complete_samples(trainer) -> int: + """Count whole groups immediately available across both consumer queues.""" + group_size = int(trainer.config.n_samples) + if group_size < 1: + raise ValueError("n_samples must be positive") + consumed = getattr(trainer, "_consumed_grpo_group_ids", set()) + grouped = {str(group): list(members) for group, members in trainer._pending_grpo_groups.items() + if str(group) not in consumed} + dispatcher = getattr(trainer, "dispatcher", None) + peek = getattr(dispatcher, "peek_ready_data", None) + for member in peek() if callable(peek) else (): + if not isinstance(member, dict) or member.get("status") == "failed": + continue + group = member.get("grpo_group_id") + if group is None or str(group) in consumed: + continue + grouped.setdefault(str(group), []).append(member) + return sum(group_size for members in grouped.values() + if len(unique_group_members(members, group_size)) >= group_size) diff --git a/infra/elastic/runtime_executor.py b/infra/elastic/runtime_executor.py index 6628046..3321305 100644 --- a/infra/elastic/runtime_executor.py +++ b/infra/elastic/runtime_executor.py @@ -19,6 +19,7 @@ from RL_Framework.config import AsyncRLConfig from RL_Framework.infra.cost_model.global_resource_planner import ( + CoreRuntimeCapabilities, GlobalResourcePlan, GlobalResourcePlanner, PlannerDecision, @@ -103,6 +104,7 @@ class PreparedRuntimeTransaction: should_pause_dispatcher: bool drain_required: bool coord_id: str + resource_reconfiguration: bool = True class RuntimeElasticExecutor: @@ -171,6 +173,13 @@ def __init__( self.runtime_run_id = "" self._peer_drain_request_coord_id = "" self._peer_drain_completed_coord_id = "" + self._pending_activation_lock = threading.Lock() + self._pending_activations: dict[str, dict[str, Any]] = {} + self._snapshot_manifests: dict[int, str] = {} + self._snapshot_cleanup_requested = False + self._retired_hybrid_joins: list[Any] = [] + self._snapshot_stop_pending: set[str] = set() + self._leased_rollout_processes: dict[str, ManagedRolloutProcess] = {} @contextmanager def _hybrid_resource_guard(self, handle=None): @@ -192,9 +201,37 @@ def prepare_transaction(self, decision: PlannerDecision) -> PreparedRuntimeTrans """Freeze runtime eligibility without launching, stopping, or waiting.""" planner_cfg = self.config.global_resource_planner plan = decision.candidate_plan + resource_permission = getattr(decision, "resource_should_reconfigure", None) + resource_reconfiguration = ( + bool(decision.should_reconfigure) + if resource_permission is None else bool(resource_permission) + ) + if not resource_reconfiguration: + # EHP changes membership within the existing allocation. Its demand + # signal does not authorize installing a rejected GRP candidate. + plan = getattr(decision, "current_plan", None) should_execute = bool(decision.should_reconfigure and plan is not None) reason = str(decision.reason) current_train_gpus = int(self.config.train_gpus) + capabilities = CoreRuntimeCapabilities.from_engine(self.train_engine) + if should_execute and resource_reconfiguration and capabilities is not None: + rejection = capabilities.incompatibility(plan) + if rejection: + # Fence stale/manual decisions before publishing a peer drain, + # pausing generation or closing live gradient servers. + signal = getattr(decision, "elastic_hybrid_signal", None) + plan = getattr(decision, "current_plan", None) + resource_reconfiguration = False + should_execute = bool(plan is not None and signal is not None and signal.should_adjust) + reason = "runtime_candidate_rejected:" + rejection + if self.has_pending_boundary_work() and not should_execute: + # Cleanup uses the same all-Core transaction as activation, even + # when the planner has no resource change to apply. + plan = getattr(decision, "current_plan", None) + if plan is not None: + should_execute = True + resource_reconfiguration = False + reason = "pending_hybrid_boundary_work" rollout_runtime_enabled = bool( getattr(planner_cfg, "apply_to_runtime", True) or getattr(planner_cfg, "runtime_manage_rollout_processes", False) @@ -222,6 +259,7 @@ def prepare_transaction(self, decision: PlannerDecision) -> PreparedRuntimeTrans if ( should_execute + and resource_reconfiguration and plan is not None and getattr(planner_cfg, "runtime_manage_rollout_processes", False) ): @@ -259,6 +297,7 @@ def prepare_transaction(self, decision: PlannerDecision) -> PreparedRuntimeTrans should_pause_dispatcher=should_pause_dispatcher, drain_required=drain_required, coord_id=coord_id, + resource_reconfiguration=resource_reconfiguration, ) def execute( @@ -323,11 +362,19 @@ def _execute_prepared( prepared.effective_planner_signal, result=result, ) + if prepared.should_execute: + self._release_unused_core_snapshots_at_boundary() if not prepared.should_execute or prepared.candidate_plan is None: return result plan = prepared.candidate_plan planner_cfg = self.config.global_resource_planner + if not prepared.resource_reconfiguration: + return self._execute_membership_only(prepared) + if getattr(self, "_lent_rollout", None): + self._release_all_hybrid_workers(reason="resource_repartition", result=result) + from RL_Framework.infra.elastic.rollout_lending import return_rollout + result.started_processes.extend(return_rollout(self)) rollout_runtime_enabled = prepared.rollout_runtime_enabled current_train_gpus = prepared.current_train_gpus strategy = prepared.strategy @@ -484,6 +531,7 @@ def _execute_prepared( if ( not cluster_swap_applied and getattr(planner_cfg, "runtime_reconfigure_training", False) + and not self._hybrid_worker_launch_enabled ): if getattr(planner_cfg, "runtime_training_pool_plan_only", True): result.training_actions.extend( @@ -556,6 +604,23 @@ def _execute_prepared( if getattr(planner_cfg, "runtime_manage_rollout_processes", False): result.actions.append("wait_rollout_ready") + if getattr(self, "trainer", None) is not None and rollout_runtime_enabled: + # Recreated services start from the HF checkpoint. Publish the + # current optimizer state before any client can route work. + self._publish_current_rollout_weights() + if (not cluster_swap_applied and self._hybrid_worker_launch_enabled + and getattr(planner_cfg, "runtime_reconfigure_training", False) + and not getattr(planner_cfg, "runtime_training_pool_plan_only", True)): + from RL_Framework.infra.elastic.rollout_lending import lend_rollout + from RL_Framework.infra.elastic.rollout_lifecycle import apply_availability, unavailable_rollouts + self._ensure_elastic_pool() + desired = self._desired_hybrid_workers(int(self.config.train_gpus), plan) + if desired: + self._cluster_swap_training_slots = lend_rollout(self, desired) + result.training_actions.extend(self._reconfigure_training_pool(int(self.config.train_gpus), plan, result)) + if getattr(self, "trainer", None) is not None: + apply_availability(self.trainer, unavailable_rollouts(self)) + self._release_unused_core_snapshots_at_boundary() self._write_rollout_manifest(plan, phase="applied") self._write_peer_rank_reconfiguration_state(plan, phase="applied") @@ -565,10 +630,212 @@ def _execute_prepared( result.errors.append(str(exc)) raise finally: - if paused and self.dispatcher is not None and hasattr(self.dispatcher, "resume"): - self.dispatcher.resume() + if paused and result.applied and self.dispatcher is not None and hasattr(self.dispatcher, "resume"): + if getattr(self, "trainer", None) is not None: + self.trainer._resume_rollout_pipeline() + else: + self.dispatcher.resume() result.actions.append("dispatcher_resume") + def _capture_core_snapshot(self) -> str: + from RL_Framework.infra.elastic.core_reconfiguration import dispatch_core_command + revision = int(getattr(self.train_engine, "optimizer_revision", 0)) + if revision in self._snapshot_manifests: + return self._snapshot_manifests[revision] + receipts = dispatch_core_command(self, "capture") + lanes: dict[str, dict[str, str]] = {} + for receipt in receipts: + if int(receipt["optimizer_revision"]) != revision: + raise RuntimeError("core snapshot ranks disagree on optimizer revision") + lane = receipt["lane"] + core_id = f"dp{int(lane['data_parallel_rank'])}" + index = (int(lane.get("tensor_parallel_rank", 0)) + + int(lane.get("context_parallel_rank", 0)) * int(self.config.train_tp_size) + + int(lane.get("pipeline_parallel_rank", 0)) * int(self.config.train_tp_size) + * int(getattr(self.config, "train_cp_size", 1))) + lanes.setdefault(core_id, {})[str(index)] = receipt["snapshot_uri"] + width = self._elastic_replica_size_gpus() + if (set(lanes) != {f"dp{i}" for i in range(int(self.config.train_dp_size))} + or any(set(value) != {str(i) for i in range(width)} for value in lanes.values())): + raise RuntimeError("core snapshot is missing model-parallel lanes") + # DistributedOptimizer gathers each TP lane's complete Adam/master + # state on its DP-zero rank. DP followers have identical parameters but + # no complete optimizer payload; every joining replica reads DP zero. + if int(getattr(self.config, "train_ep_size", 1)) == 1: + lanes = {core: dict(lanes["dp0"]) for core in lanes} + manifest = { + "format": "libra-ehp-lanes-v1", "optimizer_revision": revision, + "snapshot_id": uuid.uuid4().hex, "lanes": lanes, + } + path = self._coordination_dir() / f"ehp_state_{revision}_{manifest['snapshot_id']}.json" + self._write_json_atomic(path, manifest) + self._snapshot_manifests[revision] = str(path) + self._snapshot_cleanup_requested = True + return str(path) + + def _queue_activation(self, replica_id: str, receipt: dict[str, Any]) -> dict[str, Any]: + """Wait for Core boundary completion or cancellation after initial ACK.""" + request = {"receipt": receipt, "event": threading.Event()} + with self._hybrid_resource_guard(): + with self._pending_activation_lock: + self.elastic_pool.handoff_activation_to_boundary() + self._pending_activations[replica_id] = request + while not request["event"].wait(0.05): + if self._closing.is_set(): + raise RuntimeError("EHP activation cancelled during cleanup") + if "error" in request: + raise RuntimeError(request["error"]) + return request["receipt"] + + def has_pending_activations(self) -> bool: + with self._pending_activation_lock: + return bool(self._pending_activations) + + def has_pending_boundary_work(self) -> bool: + """Request a common Core boundary; never issue collectives here.""" + if self.has_pending_activations(): + return True + with self._lifecycle_lock: + if not self._snapshot_cleanup_requested: + return False + with self._pending_hybrid_join_lock: + if any(not handle.done() for handle in self._pending_hybrid_joins.values()): + return False + if any(not handle.done() for handle in self._retired_hybrid_joins): + return False + return bool(self._snapshot_manifests or self._snapshot_stop_pending) + + def _release_unused_core_snapshots_at_boundary(self) -> None: + """Release collectively only after every possible loader is fenced. + + Called exclusively by execute_prepared while all Core ranks follow the + transaction. Reconciliation and loader threads only request cleanup. + A cancelled Future may still be fetching, so removing its membership + does not by itself prove that its snapshot is no longer in use. + """ + with self._lifecycle_lock: + if not self._snapshot_cleanup_requested: + return + self._retired_hybrid_joins = [ + handle for handle in self._retired_hybrid_joins if not handle.done() + ] + with self._pending_hybrid_join_lock: + if self._pending_hybrid_joins or self._retired_hybrid_joins: + return + with self._pending_activation_lock: + if self._pending_activations or self._snapshot_stop_pending: + return + if self._snapshot_manifests: + from RL_Framework.infra.elastic.core_reconfiguration import dispatch_core_command + revisions = dict(self._snapshot_manifests) + dispatch_core_command(self, "release", revisions=revisions) + # Keep manifests and the request on any rank/transport error. + # Repeating a successful local release is idempotent. + for revision, manifest in revisions.items(): + Path(manifest).unlink(missing_ok=True) + self._snapshot_manifests.pop(revision, None) + self._snapshot_cleanup_requested = False + + def _stop_snapshot_loader(self, worker_id, result=None): + self._snapshot_stop_pending.add(worker_id) + self._stop_hybrid_worker_process(worker_id, result) + self._snapshot_stop_pending.discard(worker_id) + + def _finish_pending_activations(self) -> None: + import torch + with self._pending_activation_lock: + pending = dict(self._pending_activations) + if not pending: + return + revision = int(self.train_engine.optimizer_revision) + path = self._capture_core_snapshot() + task_dir = Path(self.config.global_resource_planner.hybrid_worker_task_dir) + try: + for replica_id, request in pending.items(): + if int(request["receipt"]["optimizer_revision"]) != revision: + task = task_dir / f"{replica_id}.install_{uuid.uuid4().hex}.pt" + temporary = task.with_suffix(".tmp") + # This file contains only an installation command and URI. + torch.save({"install_state": True, "snapshot_path": path}, temporary) + temporary.replace(task) + receipt = self._wait_hybrid_worker_ready( + replica_id, state_version=revision, + replica_world_size=self._elastic_replica_size_gpus(), snapshot_path=path, + ) + if int(self.train_engine.optimizer_revision) != revision: + raise RuntimeError("optimizer advanced during the activation boundary") + request["receipt"] = receipt + request["event"].set() + handle = self._pending_hybrid_joins.get(replica_id) + if handle is not None: + timeout = float(getattr(self.config.global_resource_planner, "elastic_hybrid_join_timeout_s", 180)) + handle.result(timeout=timeout if timeout > 0 else None) + self._reconcile_pending_hybrid_joins() + except Exception as exc: + for request in pending.values(): + request["error"] = str(exc) + request["event"].set() + raise + finally: + with self._pending_activation_lock: + for replica_id in pending: + self._pending_activations.pop(replica_id, None) + self._snapshot_cleanup_requested = True + + def _execute_membership_only(self, prepared): + from RL_Framework.infra.elastic.rollout_lending import lend_rollout, return_rollout + plan, result = prepared.candidate_plan, prepared.result + cfg = self.config.global_resource_planner + paused = self.dispatcher is not None + if paused: + self.dispatcher.pause() + self.dispatcher.wait_until_idle(timeout=float(getattr(cfg, "runtime_drain_timeout_s", 3600))) + try: + if getattr(cfg, "runtime_reconfigure_training", False): + if getattr(cfg, "runtime_training_pool_plan_only", True): + result.training_actions.extend(self._plan_training_pool_reconfiguration(prepared.current_train_gpus, plan)) + else: + self._ensure_elastic_pool() + desired = self._desired_hybrid_workers(prepared.current_train_gpus, plan) + if self._hybrid_worker_launch_enabled and desired: + self._cluster_swap_training_slots = lend_rollout(self, desired) + result.training_actions.extend(self._reconfigure_training_pool(prepared.current_train_gpus, plan, result)) + self._finish_pending_activations() + if self._hybrid_worker_launch_enabled and not desired: + restored = return_rollout(self) + if restored: + self._publish_current_rollout_weights() + from RL_Framework.infra.elastic.rollout_lifecycle import apply_availability, unavailable_rollouts + trainer = getattr(self, "trainer", None) + if trainer is not None: + apply_availability(trainer, unavailable_rollouts(self)) + self._release_unused_core_snapshots_at_boundary() + self._write_peer_rank_reconfiguration_state(plan, phase="applied") + result.actions.append("elastic_membership_only") + result.applied = True + return result + finally: + if paused and result.applied: + if getattr(self, "trainer", None) is not None: + self.trainer._resume_rollout_pipeline() + elif not getattr(self, "_lent_rollout", None): + self.dispatcher.resume() + + def _publish_current_rollout_weights(self, instance_ids=None): + """Publish while all core ranks are parked and admission stays closed.""" + from RL_Framework.infra.elastic.core_reconfiguration import dispatch_core_command + targets = [{"instance_id": inst.instance_id, "tp": int(inst.tp), + "gpus": list(inst.gpus), "host": inst.host, "port": inst.port} + for inst in self.config.heterogeneous_rollout.instances] + target = {"instances": targets} + if instance_ids is not None: + wanted = set(instance_ids) + selected = [inst for inst in targets if inst["instance_id"] in wanted] + if {inst["instance_id"] for inst in selected} != wanted: + raise ValueError("returned rollout publication contains unknown instances") + target = {"instances": selected, "scope": "returned_subset"} + dispatch_core_command(self, "publish", target=target) + def _reconfigure_training_pool( self, current_train_gpus: int, @@ -613,7 +880,31 @@ def _reconfigure_training_pool( delta = desired_replicas - active_replicas if delta > 0: - groups = self._candidate_replica_groups(candidates, replica_size) + if self._hybrid_worker_launch_enabled: + # Only confirmed loans/cluster swaps transfer physical GPU + # ownership. A failed loan must not fall back to live rollout + # slots or start a snapshot/join with an empty allocation. + by_host: dict[str, list[str]] = {} + physical = set() + for worker_id in candidates: + slot = self._cluster_swap_training_slot(worker_id) + gpus = slot.get("gpus", []) + if len(gpus) != 1: + continue + host = str(slot.get("host", "")) + address = (host, int(gpus[0])) + if address in physical: + continue + physical.add(address) + by_host.setdefault(host, []).append(worker_id) + groups = [ + group for workers in by_host.values() + for group in self._candidate_replica_groups(workers, replica_size) + ] + if groups: + self._capture_core_snapshot() + else: + groups = self._candidate_replica_groups(candidates, replica_size) for offset, (replica_id, members) in enumerate(groups[:delta]): target_core = core_ids[(active_replicas + offset) % len(core_ids)] member_ranks = {worker_id: rank for rank, worker_id in enumerate(members)} @@ -623,7 +914,7 @@ def activate_replica( _members: tuple[str, ...], target_core_id: str, state_version: int, - ) -> None: + ) -> dict[str, Any]: slots = [self._cluster_swap_training_slot(wid) for wid in _members] hosts = {str(slot.get("host", "")) for slot in slots} if len(hosts) > 1: @@ -648,23 +939,46 @@ def activate_replica( replica_world_size=len(_members), ) with self._hybrid_resource_guard(): - proc = subprocess.Popen(command, shell=True, env=os.environ.copy()) - self._hybrid_worker_processes[_replica_id] = proc + task_dir = Path(self.config.global_resource_planner.hybrid_worker_task_dir) + for ready in task_dir.glob(f"{_replica_id}.rank_*.ready"): + ready.unlink(missing_ok=True) + (task_dir / f"{_replica_id}.ready").unlink(missing_ok=True) + host = next(iter(hosts), "") + if self._hybrid_worker_remote_control_enabled: + meta = self._request_remote_hybrid_worker( + worker_id=_replica_id, target_core_id=target_core_id, + snapshot_path=snapshot_path, command=command, + slot={"host": host, "gpus": gpus}, + ) + pid = meta.pid + proc = None + else: + proc = subprocess.Popen(command, shell=True, env=os.environ.copy(), start_new_session=True) + self._hybrid_worker_processes[_replica_id] = proc + pid = proc.pid for rank, worker_id in enumerate(_members): - self._hybrid_worker_processes[worker_id] = proc + if proc is not None: + self._hybrid_worker_processes[worker_id] = proc self._hybrid_worker_meta[worker_id] = ManagedHybridWorkerProcess( worker_id=worker_id, target_core_id=target_core_id, command=command, - pid=proc.pid, + pid=pid, snapshot_path=snapshot_path, - host=next(iter(hosts), ""), + host=host, gpus=gpus, replica_id=_replica_id, replica_rank=rank, replica_world_size=len(_members), ) - self._wait_hybrid_worker_ready(_replica_id) + receipt = self._wait_hybrid_worker_ready( + _replica_id, state_version=state_version, + replica_world_size=len(_members), snapshot_path=snapshot_path, + target_core_id=target_core_id, + timeout_s=float(getattr(self.config.global_resource_planner, + "elastic_hybrid_join_timeout_s", 180.0)), + ) + return self._queue_activation(_replica_id, receipt) with self._hybrid_resource_guard(): handle = self.elastic_pool.join_replica( @@ -708,6 +1022,7 @@ def activate_replica( actions.append(f"join_replica_capacity_shortfall:{delta - len(groups)}") elif delta < 0: for replica_id in active_replica_ids[: abs(delta)]: + self.wait_for_hybrid_training_tasks(replica_id) with self._pending_hybrid_join_lock: pending = self._pending_hybrid_joins.pop(replica_id, None) self._pending_hybrid_join_started_at.pop(replica_id, None) @@ -735,6 +1050,7 @@ def activate_replica( in {"hybrid_training", "hybrid_joining"} ] for worker_id in active_workers[: abs(delta)]: + self.wait_for_hybrid_training_tasks(worker_id) with self._pending_hybrid_join_lock: pending = self._pending_hybrid_joins.pop(worker_id, None) self._pending_hybrid_join_started_at.pop(worker_id, None) @@ -964,16 +1280,31 @@ def _release_all_hybrid_workers( reason: str, result: RuntimeReconfigurationResult | None = None, ) -> None: + with self._lifecycle_lock: + return self._release_all_hybrid_workers_locked(reason=reason, result=result) + + def _release_all_hybrid_workers_locked(self, *, reason, result=None): + self.wait_for_hybrid_training_tasks() + self._snapshot_cleanup_requested = True + with self._pending_activation_lock: + for request in self._pending_activations.values(): + request["error"] = f"EHP activation cancelled: {reason}" + request["event"].set() + self._pending_activations.clear() with self._pending_hybrid_join_lock: pending_items = list(self._pending_hybrid_joins.items()) self._pending_hybrid_joins.clear() self._pending_hybrid_join_started_at.clear() - for replica_id, handle in pending_items: + # Fence every removed generation before stopping any one worker: stop + # may fail, but all surviving loaders must remain tracked for cleanup. + for _replica_id, handle in pending_items: handle.cancel() + self._retired_hybrid_joins.append(handle) + for replica_id, handle in pending_items: self._publish_hybrid_membership() members = tuple(getattr(handle, "member_worker_ids", (replica_id,))) for worker_id in members: - self._stop_hybrid_worker_process(worker_id, result) + self._stop_snapshot_loader(worker_id, result) if result is not None: result.training_actions.append( f"join_replica_cancelled:{replica_id}:{reason}" @@ -999,7 +1330,7 @@ def _release_all_hybrid_workers( continue self._publish_hybrid_membership() for worker in members: - self._stop_hybrid_worker_process(worker.worker_id, result) + self._stop_snapshot_loader(worker.worker_id, result) if result is not None: result.training_actions.append( f"release_replica_to_rollout:{replica_id}:{reason}" @@ -1016,13 +1347,19 @@ def _release_all_hybrid_workers( except RuntimeError: continue self._publish_hybrid_membership() - self._stop_hybrid_worker_process(worker_id, result) + self._stop_snapshot_loader(worker_id, result) def _reconcile_pending_hybrid_joins( self, result: RuntimeReconfigurationResult | None = None, ) -> None: - """Collect completed background joins without blocking the train loop.""" + """Collect joins and request cleanup without issuing Core collectives.""" + with self._lifecycle_lock: + return self._reconcile_pending_hybrid_joins_locked(result) + + def _reconcile_pending_hybrid_joins_locked(self, result=None): + for worker_id in list(self._snapshot_stop_pending): + self._stop_snapshot_loader(worker_id, result) with self._pending_hybrid_join_lock: pending_items = list(self._pending_hybrid_joins.items()) for replica_id, handle in pending_items: @@ -1035,8 +1372,16 @@ def _reconcile_pending_hybrid_joins( 180.0, ) ) - if time.time() - started > timeout: - handle.cancel() + if timeout > 0 and time.time() - started > timeout: + # Serialize the cold-timeout decision with the boundary + # handoff; ready joins now wait for their Core boundary. + with self._lifecycle_lock: + with self._pending_activation_lock: + if replica_id in self._pending_activations: + continue + handle.cancel() + self._retired_hybrid_joins.append(handle) + self._snapshot_cleanup_requested = True self._publish_hybrid_membership() with self._pending_hybrid_join_lock: self._pending_hybrid_joins.pop(replica_id, None) @@ -1044,7 +1389,7 @@ def _reconcile_pending_hybrid_joins( for worker_id in tuple( getattr(handle, "member_worker_ids", (replica_id,)) ): - self._stop_hybrid_worker_process(worker_id, result) + self._stop_snapshot_loader(worker_id, result) if result is not None: result.training_actions.append( f"join_replica_timeout_cancelled:{replica_id}" @@ -1064,10 +1409,11 @@ def _reconcile_pending_hybrid_joins( for worker_id in tuple( getattr(handle, "member_worker_ids", (replica_id,)) ): - self._stop_hybrid_worker_process(worker_id, result) + self._stop_snapshot_loader(worker_id, result) if result is not None: result.errors.append(f"hybrid_replica_join_failed:{replica_id}:{exc}") finally: + self._snapshot_cleanup_requested = True with self._pending_hybrid_join_lock: self._pending_hybrid_joins.pop(replica_id, None) self._pending_hybrid_join_started_at.pop(replica_id, None) @@ -1081,13 +1427,16 @@ def _activate_hybrid_worker( replica_id: str = "", replica_rank: int = 0, replica_world_size: int = 1, - ) -> None: + ) -> dict[str, Any]: """Activation barrier executed by the pool's background join thread.""" + from RL_Framework.infra.elastic.readiness import clear_ready_receipts snapshot_path = self._snapshot_path_for_version(state_version) + task_dir = Path(self.config.global_resource_planner.hybrid_worker_task_dir) slot = self._cluster_swap_training_slot(worker_id) if not slot and self._hybrid_worker_remote_control_enabled: slot = self._remote_hybrid_worker_slot(worker_id) if not self._is_hybrid_worker_running(worker_id): + clear_ready_receipts(task_dir, worker_id) command = self._build_hybrid_worker_command( worker_id=worker_id, target_core_id=target_core_id, @@ -1113,7 +1462,7 @@ def _activate_hybrid_worker( meta.replica_world_size = replica_world_size self._hybrid_worker_meta[worker_id] = meta else: - proc = subprocess.Popen(command, shell=True, env=os.environ.copy()) + proc = subprocess.Popen(command, shell=True, env=os.environ.copy(), start_new_session=True) self._hybrid_worker_processes[worker_id] = proc self._hybrid_worker_meta[worker_id] = ManagedHybridWorkerProcess( worker_id=worker_id, @@ -1139,6 +1488,7 @@ def _activate_hybrid_worker( # it in the background before the domain marks it active. with self._hybrid_resource_guard(): self._stop_hybrid_worker_process(worker_id) + clear_ready_receipts(task_dir, worker_id) command = self._build_hybrid_worker_command( worker_id=worker_id, target_core_id=target_core_id, @@ -1164,7 +1514,7 @@ def _activate_hybrid_worker( meta.replica_world_size = replica_world_size self._hybrid_worker_meta[worker_id] = meta else: - proc = subprocess.Popen(command, shell=True, env=os.environ.copy()) + proc = subprocess.Popen(command, shell=True, env=os.environ.copy(), start_new_session=True) self._hybrid_worker_processes[worker_id] = proc self._hybrid_worker_meta[worker_id] = ManagedHybridWorkerProcess( worker_id=worker_id, @@ -1176,7 +1526,13 @@ def _activate_hybrid_worker( replica_rank=replica_rank, replica_world_size=replica_world_size, ) - self._wait_hybrid_worker_ready(worker_id) + receipt = self._wait_hybrid_worker_ready( + worker_id, state_version=state_version, replica_world_size=replica_world_size, + snapshot_path=snapshot_path, target_core_id=target_core_id, + timeout_s=float(getattr(self.config.global_resource_planner, + "elastic_hybrid_join_timeout_s", 180.0)), + ) + return self._queue_activation(replica_id or worker_id, receipt) @property def _hybrid_worker_remote_control_enabled(self) -> bool: @@ -1318,6 +1674,10 @@ def _runtime_plan_already_applied( if reconfigure_training and target_train_gpus != current_train_gpus: return False + current_tp = [int(instance.tp) for instance in self.config.heterogeneous_rollout.instances] + if [int(tp) for tp in plan.rollout_tp_list] != current_tp: + return False + desired = { meta.instance_id: meta for meta in self._desired_rollout_processes(plan) @@ -1556,6 +1916,7 @@ def _apply_plan_to_runtime_config(self, plan: GlobalResourcePlan) -> None: "train_gpus": self.config.train_gpus, "train_tp_size": self.config.train_tp_size, "train_ep_size": self.config.train_ep_size, + "train_cp_size": self.config.train_cp_size, "tp_size": self.config.tp_size, "train_pp_size": self.config.train_pp_size, "train_dp_size": self.config.train_dp_size, @@ -1646,12 +2007,14 @@ def _write_peer_rank_reconfiguration_state( ) -> Path: coord_dir = self._coordination_dir() coord_dir.mkdir(parents=True, exist_ok=True) + from RL_Framework.infra.elastic.rollout_lifecycle import unavailable_rollouts if coord_id is None: coord_id = getattr(self, "_active_runtime_coord_id", "") payload = { "coord_id": coord_id, "run_id": self.runtime_run_id, "phase": phase, + "unavailable_rollout_ids": unavailable_rollouts(self), "job_id": os.environ.get("SLURM_JOB_ID", os.environ.get("JOB_ID", "")), "updated_at": time.time(), "plan": plan.to_dict(), @@ -1850,7 +2213,7 @@ def _coordinate_peer_rank_drain( 3600.0, ) ) - deadline = time.time() + timeout + deadline = time.time() + timeout if timeout > 0 else float("inf") expected_ranks = self._expected_peer_ready_ranks() expected = [ready_dir / f"rank_{rank}.json" for rank in expected_ranks] job_id = os.environ.get("SLURM_JOB_ID", os.environ.get("JOB_ID", "")) @@ -1890,7 +2253,8 @@ def _expected_peer_ready_ranks(self) -> list[int]: if world_size <= 1: return [] planner_cfg = self.config.global_resource_planner - if not getattr(planner_cfg, "runtime_coordinate_batch_source_only", True): + if (getattr(planner_cfg, "runtime_reconfigure_training", False) + or not getattr(planner_cfg, "runtime_coordinate_batch_source_only", True)): return list(range(1, world_size)) model_parallel = ( @@ -1940,6 +2304,10 @@ def _initialize_elastic_pool(self): ] def fetch_snapshot(worker_id: str, target_core_id: str) -> int: + if self._hybrid_worker_launch_enabled: + if not self._snapshot_manifests: + raise RuntimeError("EHP join has no collective live-state snapshot at this boundary") + return max(self._snapshot_manifests) if hasattr(self.train_engine, "capture_elastic_state_snapshot"): return int( self.train_engine.capture_elastic_state_snapshot( @@ -2117,7 +2485,7 @@ def _launch_prewarmed_hybrid_worker( ) -> ManagedHybridWorkerProcess: with self._hybrid_resource_guard(): env = os.environ.copy() - proc = subprocess.Popen(command, shell=True, env=env) + proc = subprocess.Popen(command, shell=True, env=env, start_new_session=True) meta = ManagedHybridWorkerProcess( worker_id=worker_id, target_core_id=target_core_id, @@ -2161,6 +2529,36 @@ def _adopt_prewarmed_hybrid_worker( result.started_hybrid_workers.append(meta) return True + def wait_for_hybrid_training_tasks(self, worker_id: str | None = None) -> None: + """Finish assigned updates before a normal boundary returns their GPUs.""" + with self._lifecycle_lock: + if worker_id is None: + worker_ids = {meta.replica_id or meta.worker_id + for meta in self._hybrid_worker_meta.values()} + else: + meta = self._hybrid_worker_meta.get(worker_id) + worker_ids = {(meta.replica_id or meta.worker_id) if meta else worker_id} + cfg = self.config.global_resource_planner + directory = Path(getattr(cfg, "hybrid_worker_task_dir", "./logs/elastic_training_tasks")) + timeout = float(getattr(cfg, "hybrid_update_timeout_s", 300.0)) + deadline = time.monotonic() + timeout if timeout > 0 else float("inf") + for identity in worker_ids: + for task in sorted(directory.glob(f"{identity}.step_*.pt")): + while not task.with_suffix(task.suffix + ".done").exists(): + if task.with_suffix(task.suffix + ".failed").exists(): + raise RuntimeError(f"Hybrid training task failed before completion: {task.name}") + if time.monotonic() >= deadline: + raise TimeoutError(f"Hybrid training task completion timed out: {task.name}") + time.sleep(.05) + + def _retire_hybrid_worker_tasks(self, worker_id: str) -> None: + # Called only after the owned process group exited, while its launch + # gate remains locked. Keep cancellation distinct from successful work. + directory = Path(getattr(self.config.global_resource_planner, + "hybrid_worker_task_dir", "./logs/elastic_training_tasks")) + for task in directory.glob(f"{worker_id}.*.pt"): + task.rename(task.with_suffix(task.suffix + ".cancelled")) + def _stop_hybrid_worker_process( self, worker_id: str, @@ -2170,38 +2568,56 @@ def _stop_hybrid_worker_process( self._stop_hybrid_worker_process_locked(worker_id, result) def _stop_hybrid_worker_process_locked( - self, - worker_id: str, - result: RuntimeReconfigurationResult | None = None, + self, worker_id: str, result: RuntimeReconfigurationResult | None = None, ) -> None: proc = self._hybrid_worker_processes.get(worker_id) meta = self._hybrid_worker_meta.get(worker_id) if proc is None: if meta is not None and self._hybrid_worker_remote_control_enabled: + from RL_Framework.infra.elastic.rollout_lending import wait_hybrid_stopped control_dir = self._hybrid_worker_control_dir() control_dir.mkdir(parents=True, exist_ok=True) - # Withdraw an unclaimed start before queuing its stop. The - # controller may already own a .running request; no remote - # exit acknowledgement is implied by this local publish. - (control_dir / f".launch_{worker_id}.json").unlink(missing_ok=True) - self._write_json_atomic( - control_dir / f".stop_{worker_id}.json", - {"worker_id": worker_id, "host": meta.host}, - ) - self._hybrid_worker_meta.pop(worker_id, None) + launch_id = meta.worker_id + replica = meta.replica_id + if replica and replica in self._hybrid_worker_meta: + candidate = self._hybrid_worker_meta[replica] + if candidate.worker_id == replica: + launch_id, meta = replica, candidate + (control_dir / f".launch_{launch_id}.json").unlink(missing_ok=True) + request_id = getattr(meta, "stop_request_id", None) or uuid.uuid4().hex + meta.stop_request_id = request_id + outstanding = False + for suffix in ("json", "running", "done"): + try: + record = json.loads((control_dir / f".stop_{launch_id}.{suffix}").read_text(encoding="utf-8")) + outstanding |= record.get("request_id") == request_id + except FileNotFoundError: + pass + if not outstanding: + self._write_json_atomic(control_dir / f".stop_{launch_id}.json", { + "worker_id": launch_id, "host": meta.host, "pid": meta.pid, + "request_id": request_id, + }) + acknowledgement = wait_hybrid_stopped(self, launch_id, request_id) + if acknowledgement.get("request_id") != request_id or acknowledgement.get("exited") is not True: + raise RuntimeError("Hybrid stop did not acknowledge the owned request") + self._retire_hybrid_worker_tasks(launch_id) + # Retain every ownership record until the controller proves the + # entire recorded worker process group exited. + for alias, owned in list(self._hybrid_worker_meta.items()): + if owned is meta or (owned.worker_id == launch_id and owned.pid == meta.pid) or ( + owned.replica_id == launch_id and owned.host == meta.host + ): + self._hybrid_worker_meta.pop(alias, None) + elif meta is not None: + raise RuntimeError(f"Hybrid worker {worker_id} has metadata but no owned process or remote controller") if result is not None and meta is not None: result.training_actions.append(f"stop_hybrid_worker_meta:{worker_id}") return - if proc.poll() is None: - proc.terminate() - try: - proc.wait(timeout=5) - except subprocess.TimeoutExpired: - proc.kill() - proc.wait(timeout=5) - # A replica process is registered under both replica and member IDs. - # Remove all aliases only after stop succeeds, retaining retryable - # ownership if terminate/wait or remote publication fails. + from RL_Framework.infra.elastic.process_lifecycle import stop_local_worker + stop_local_worker(proc) + if meta is not None: + self._retire_hybrid_worker_tasks(meta.replica_id or meta.worker_id) for alias, owned in list(self._hybrid_worker_processes.items()): if owned is proc: self._hybrid_worker_processes.pop(alias, None) @@ -2281,7 +2697,7 @@ def _launch_hybrid_worker_after_join( ) with self._hybrid_resource_guard(handle): env = os.environ.copy() - proc = subprocess.Popen(command, shell=True, env=env) + proc = subprocess.Popen(command, shell=True, env=env, start_new_session=True) meta = ManagedHybridWorkerProcess( worker_id=worker.worker_id, target_core_id=handle.target_core_id, @@ -2301,55 +2717,60 @@ def _launch_hybrid_worker_after_join( result.errors.append(f"hybrid_worker_launch_failed:{exc}") return None - def _wait_hybrid_worker_ready(self, worker_id: str) -> None: - task_dir = Path( - getattr( - self.config.global_resource_planner, - "hybrid_worker_task_dir", - "./logs/elastic_training_tasks", - ) - ) - ready_path = task_dir / f"{worker_id}.ready" + def _wait_hybrid_worker_ready( + self, worker_id: str, *, state_version: int | None = None, + replica_world_size: int | None = None, snapshot_path: str | None = None, + target_core_id: str | None = None, timeout_s: float | None = None, + ) -> dict[str, Any]: + from RL_Framework.infra.elastic.readiness import read_ready_receipts, snapshot_contract + task_dir = Path(getattr(self.config.global_resource_planner, "hybrid_worker_task_dir", "./logs/elastic_training_tasks")) + meta = self._hybrid_worker_meta.get(worker_id) + if meta is None: + meta = next((value for value in self._hybrid_worker_meta.values() if value.replica_id == worker_id), None) + if snapshot_path is None and meta is not None: + snapshot_path = meta.snapshot_path + if target_core_id is None and meta is not None: + target_core_id = meta.target_core_id + width = int(replica_world_size if replica_world_size is not None else getattr(meta, "replica_world_size", 1)) + if width <= 0: + raise ValueError("Hybrid readiness requires positive model-lane coverage") + contract = snapshot_contract(snapshot_path, state_version=state_version, target_core_id=target_core_id) control_dir = self._hybrid_worker_control_dir() - launch_error_path = control_dir / f".launch_{worker_id}.error" - launch_started_path = control_dir / f".launch_{worker_id}.started" - timeout = float( - getattr( - self.config.global_resource_planner, - "hybrid_worker_ready_timeout_s", - 60.0, - ) - ) - deadline = time.time() + max(timeout, 0.0) - while time.time() < deadline: + error_path = control_dir / f".launch_{worker_id}.error" + started_path = control_dir / f".launch_{worker_id}.started" + # Cold startup and its first installation use the overall join budget, + # still bounded by the pool's original deadline through the guard below. + # A running worker's catch-up installation keeps the shorter ACK budget. + timeout = float(timeout_s if timeout_s is not None else getattr( + self.config.global_resource_planner, "hybrid_worker_ready_timeout_s", 60.)) + deadline = time.monotonic() + timeout if timeout > 0 else None + while True: + if self._closing.is_set(): + from RL_Framework.infra.elastic.hybrid_pool import JoinCancelledError + raise JoinCancelledError(f"Hybrid readiness cancelled during cleanup: {worker_id}") with self._hybrid_resource_guard(): - if ready_path.exists(): - meta = self._hybrid_worker_meta.get(worker_id) - if meta is not None and launch_started_path.exists(): - try: - payload = json.loads( - launch_started_path.read_text(encoding="utf-8") - ) - meta.pid = int(payload.get("pid", meta.pid)) - except Exception: - pass - return - if launch_error_path.exists(): - raise RuntimeError( - f"remote hybrid worker {worker_id} launch failed: " - f"{launch_error_path.read_text(encoding='utf-8')[-1000:]}" - ) + if error_path.exists(): + raise RuntimeError(f"remote hybrid worker {worker_id} launch failed: {error_path.read_text(encoding='utf-8')[-1000:]}") proc = self._hybrid_worker_processes.get(worker_id) if proc is not None and proc.poll() is not None: - raise RuntimeError( - f"hybrid worker {worker_id} exited before ready" - ) - self._closing.wait(0.2) - raise TimeoutError( - f"hybrid worker {worker_id} did not become ready after {timeout:.0f}s" - ) + raise RuntimeError(f"hybrid worker {worker_id} exited before ready") + receipt = read_ready_receipts(task_dir, worker_id, width, contract) + if receipt is not None: + if meta is not None and started_path.exists(): + started = json.loads(started_path.read_text(encoding="utf-8")) + if started.get("worker_id", worker_id) != worker_id: + raise RuntimeError("Hybrid launch controller returned a different worker identity") + meta.pid = int(started["pid"]) + if meta is not None: + meta.snapshot_path = str(snapshot_path) + return receipt + if deadline is not None and time.monotonic() >= deadline: + raise TimeoutError(f"hybrid worker {worker_id} did not acknowledge snapshot {contract['snapshot_id']} revision {contract['optimizer_revision']} on all {width} lanes after {timeout:.0f}s") + self._closing.wait(.05) def _snapshot_path_for_version(self, version: int) -> str: + if int(version) in self._snapshot_manifests: + return self._snapshot_manifests[int(version)] if self.train_engine is not None and hasattr( self.train_engine, "get_elastic_state_snapshot_path", @@ -2376,6 +2797,13 @@ def _build_hybrid_worker_command( replica_world_size: int = 1, ) -> str: cfg = self.config.global_resource_planner + if self._hybrid_worker_launch_enabled and hasattr(self.config, "to_yaml"): + config_path = self._coordination_dir() / f"hybrid-config-{self._active_runtime_coord_id}.yaml" + config_path.parent.mkdir(parents=True, exist_ok=True) + temporary = config_path.with_name(f".{config_path.name}.{uuid.uuid4().hex}.tmp") + self.config.to_yaml(str(temporary)) + temporary.replace(config_path) + cfg.hybrid_worker_config_path = str(config_path) endpoint = self.gradient_server.endpoint gradient_host = getattr(cfg, "gradient_server_public_host", "") or endpoint.host rl_root = Path(__file__).resolve().parents[2] @@ -2400,7 +2828,10 @@ def _build_hybrid_worker_command( getattr(cfg, "hybrid_worker_config_path", "") or os.environ.get("RUNTIME_CONFIG", os.environ.get("CONFIG_PATH", "")) ), - "gradient_endpoint_dir": shlex.quote(getattr(cfg, "hybrid_worker_endpoint_dir", getattr(cfg, "hybrid_worker_task_dir", "")) + "/core_endpoints"), + "gradient_endpoint_dir": shlex.quote( + (getattr(cfg, "hybrid_worker_endpoint_dir", "") + or getattr(cfg, "hybrid_worker_task_dir", "")) + "/core_endpoints" + ), "replica_gpus": len(assigned_gpus) or int(replica_world_size), "gradient_host": shlex.quote(gradient_host), "gradient_port": endpoint.port, @@ -2419,7 +2850,7 @@ def _build_hybrid_worker_command( env_prefix = f"CUDA_VISIBLE_DEVICES={values['cuda_visible_devices']} " worker_script = f"{values['rl_framework_path']}/scripts/elastic_hybrid_worker.py" launcher = "" - if values["worker_mode"] == "'megatron_core'" and values["replica_gpus"] > 1: + if getattr(cfg, "hybrid_worker_mode", "megatron_core") == "megatron_core": launcher = f"{values['python']} -m torch.distributed.run --standalone --nproc_per_node={values['replica_gpus']} " else: launcher = f"{values['python']} " @@ -2452,28 +2883,8 @@ def _role_name(role: Any) -> str: def _stop_managed_rollout_processes(self) -> list[ManagedRolloutProcess]: stopped: list[ManagedRolloutProcess] = [] - timeout = float(self.config.global_resource_planner.vllm_stop_timeout_s) - for instance_id, meta in list(self._process_meta.items()): - proc = self._processes.get(instance_id) - if meta.adopted: - if meta.pid > 0: - self._terminate_external_pid(meta) - else: - self._run_external_stop(meta) - self._processes.pop(instance_id, None) - self._process_meta.pop(instance_id, None) - stopped.append(meta) - continue - if proc is not None and proc.poll() is None: - proc.send_signal(signal.SIGTERM) - try: - proc.wait(timeout=timeout) - except subprocess.TimeoutExpired: - proc.kill() - proc.wait(timeout=5) - stopped.append(meta) - self._processes.pop(instance_id, None) - self._process_meta.pop(instance_id, None) + for instance_id in list(self._process_meta): + stopped.extend(self._stop_rollout_process(instance_id)) return stopped def _adopt_existing_rollout_processes(self) -> None: @@ -2737,7 +3148,7 @@ def _blue_green_reconfigure_rollout_processes( self, plan: GlobalResourcePlan, ) -> tuple[list[ManagedRolloutProcess], list[ManagedRolloutProcess]]: - """Hot-swap rollout processes by starting replacements before cutover.""" + """Release overlapping GPU owners before starting their replacements.""" desired = { meta.instance_id: meta @@ -2770,6 +3181,7 @@ def _blue_green_reconfigure_rollout_processes( wanted.command = self._command_for_instance_id(instance_id, wanted.port) to_stop.append(instance_id) + starting = [] for instance_id, wanted in desired.items(): current = self._process_meta.get(instance_id) if current is not None and self._same_rollout_process(current, wanted): @@ -2781,6 +3193,21 @@ def _blue_green_reconfigure_rollout_processes( self._set_config_instance_port(instance_id, wanted.port) wanted.command = self._command_for_instance_id(instance_id, wanted.port) used_ports.add(wanted.port) + starting.append(wanted) + + starting_slots = {(meta.host, gpu) for meta in starting for gpu in meta.gpus} + overlapping = [identity for identity in to_stop if any( + (self._process_meta[identity].host, gpu) in starting_slots + for gpu in self._process_meta[identity].gpus + )] + # A spare HTTP port does not provide spare GPU memory. The existing + # stop operation waits for process exit before releasing its ownership. + for identity in overlapping: + stopped.extend(self._stop_rollout_process(identity)) + to_stop = [identity for identity in to_stop if identity not in overlapping] + self._sleep_after_rollout_stop(stopped) + + for wanted in starting: started.append(self._start_one_rollout_process(wanted)) self._wait_rollout_process_start_ready(wanted) @@ -2867,12 +3294,11 @@ def _start_one_rollout_process( log_handle.flush() stdout = log_handle try: - proc = subprocess.Popen( - meta.command, - shell=True, - env=env, + from RL_Framework.infra.elastic.rollout_process import launch_rollout_process + proc = launch_rollout_process( + meta, env, stdout=stdout, - stderr=subprocess.STDOUT if stdout is not None else None, + control_dir=self.config.rollout_weight_sync_control_dir, ) finally: if log_handle is not None: @@ -2896,8 +3322,22 @@ def _rollout_launch_env(self, meta: ManagedRolloutProcess) -> dict[str, str]: env = os.environ.copy() for key in self._TRAIN_DISTRIBUTED_ENV_KEYS: env.pop(key, None) - env["PLANNED_CUDA_VISIBLE_DEVICES"] = ",".join(str(g) for g in meta.gpus) + for key in tuple(env): + if key.startswith('TORCHELASTIC_'): + env.pop(key) + devices = ",".join(str(gpu) for gpu in meta.gpus) + env["PLANNED_CUDA_VISIBLE_DEVICES"] = devices + env["CUDA_VISIBLE_DEVICES"] = devices env["RL_FRAMEWORK_ROLE"] = "runtime_rollout" + env["LIBRA_INSTANCE_ID"] = meta.instance_id + env["LIBRA_SERVICE_LABEL"] = meta.instance_id + env["LIBRA_EXPECTED_TP_SIZE"] = str(meta.tp) + env.setdefault("LIBRA_CHECK_PREFIX_CACHE", "1") + env.setdefault("VLLM_ENABLE_V1_MULTIPROCESSING", "1") + backend = str(getattr(self.config.heterogeneous_rollout.scheduling, "cmlfq_kv_backend", "recompute")) + if backend == "cpu_offload": + directory = str(Path(self.config.rollout_weight_sync_control_dir) / "kv-store") + env["LIBRA_KV_STORAGE_PATH"] = os.environ.get("LIBRA_KV_STORAGE_PATH", directory) return env def _stop_rollout_process(self, instance_id: str) -> list[ManagedRolloutProcess]: @@ -2905,21 +3345,9 @@ def _stop_rollout_process(self, instance_id: str) -> list[ManagedRolloutProcess] if meta is None: return [] proc = self._processes.get(instance_id) - if meta.adopted or proc is None: - if meta.pid > 0: - self._terminate_external_pid(meta) - self._run_external_stop(meta) - self._process_meta.pop(instance_id, None) - self._processes.pop(instance_id, None) - return [meta] - timeout = float(self.config.global_resource_planner.vllm_stop_timeout_s) - if proc.poll() is None: - proc.send_signal(signal.SIGTERM) - try: - proc.wait(timeout=timeout) - except subprocess.TimeoutExpired: - proc.kill() - proc.wait(timeout=5) + self._stop_process_meta(meta, proc) + # Retain ownership on a failed stop: a retry must still be able to find + # the live service, and its GPUs cannot yet be lent to training. self._processes.pop(instance_id, None) self._process_meta.pop(instance_id, None) return [meta] @@ -2930,18 +3358,18 @@ def _stop_process_meta( proc: subprocess.Popen | None, ) -> list[ManagedRolloutProcess]: if meta.adopted or proc is None: - if meta.pid > 0: + template = getattr(self.config.global_resource_planner, "vllm_stop_command_template", "") + if template: + self._run_external_stop(meta) + elif meta.pid > 0: self._terminate_external_pid(meta) - self._run_external_stop(meta) + else: + raise RuntimeError(f"No process ownership or stop command for rollout {meta.instance_id}") return [meta] + from RL_Framework.infra.elastic.rollout_process import stop_rollout_process timeout = float(self.config.global_resource_planner.vllm_stop_timeout_s) - if proc.poll() is None: - proc.send_signal(signal.SIGTERM) - try: - proc.wait(timeout=timeout) - except subprocess.TimeoutExpired: - proc.kill() - proc.wait(timeout=5) + stop_rollout_process(proc, meta, timeout=timeout, + control_dir=self.config.rollout_weight_sync_control_dir) return [meta] @staticmethod @@ -3033,23 +3461,21 @@ def _run_external_stop(self, meta: ManagedRolloutProcess) -> None: check=False, stdout=log_handle, stderr=subprocess.STDOUT, - timeout=float( - getattr( - self.config.global_resource_planner, - "vllm_stop_timeout_s", - 60.0, - ) - ), + timeout=(float(getattr(self.config.global_resource_planner, "vllm_stop_timeout_s", 60.0)) or None), ) log_handle.write( f"[{time.strftime('%Y-%m-%d %H:%M:%S')}] " f"stop exit_code={completed.returncode}\n" ) - except subprocess.TimeoutExpired: + if completed.returncode != 0: + raise RuntimeError(f"rollout stop failed for {meta.instance_id}: exit_code={completed.returncode}") + except subprocess.TimeoutExpired as exc: log_handle.write( f"[{time.strftime('%Y-%m-%d %H:%M:%S')}] stop timeout\n" ) + raise TimeoutError(f"rollout stop timed out for {meta.instance_id}") from exc + def _wait_rollout_process_start_ready(self, meta: ManagedRolloutProcess) -> None: if not getattr( self.config.global_resource_planner, @@ -3123,46 +3549,12 @@ def _sleep_before_rollout_stop(self) -> None: time.sleep(grace_s) def _terminate_external_pid(self, meta: ManagedRolloutProcess) -> None: - timeout = float( - getattr( - self.config.global_resource_planner, - "vllm_stop_timeout_s", - 30.0, - ) + from RL_Framework.infra.elastic.rollout_process import stop_adopted_rollout + stop_adopted_rollout( + meta, env=self._rollout_launch_env(meta), + timeout=float(self.config.global_resource_planner.vllm_stop_timeout_s), + control_dir=self.config.rollout_weight_sync_control_dir, ) - log_path = meta.log_path or self._rollout_log_path(meta.instance_id) - Path(log_path).parent.mkdir(parents=True, exist_ok=True) - with open(log_path, "a", encoding="utf-8") as log_handle: - log_handle.write( - f"\n[{time.strftime('%Y-%m-%d %H:%M:%S')}] " - f"terminating adopted launcher pid={meta.pid} " - f"instance={meta.instance_id} host={meta.host} port={meta.port}\n" - ) - log_handle.flush() - try: - os.kill(meta.pid, signal.SIGTERM) - except ProcessLookupError: - log_handle.write("launcher pid already exited\n") - return - except PermissionError as exc: - log_handle.write(f"launcher pid terminate permission error: {exc}\n") - self._run_external_stop(meta) - return - - deadline = time.time() + timeout - while time.time() < deadline: - try: - os.kill(meta.pid, 0) - except ProcessLookupError: - log_handle.write("launcher pid exited after SIGTERM\n") - return - time.sleep(0.5) - - log_handle.write("launcher pid did not exit after SIGTERM; sending SIGKILL\n") - try: - os.kill(meta.pid, signal.SIGKILL) - except ProcessLookupError: - log_handle.write("launcher pid exited before SIGKILL\n") def _write_rollout_manifest( self, @@ -3205,62 +3597,49 @@ def _build_vllm_command(self, inst: Any, port: int) -> str: planner_cfg = self.config.global_resource_planner template = getattr(planner_cfg, "vllm_launch_command_template", "") host = inst.host or self.config.heterogeneous_rollout.vllm_host - max_model_len = ( - self.config.heterogeneous_rollout.max_model_len - or self.config.max_seq_length - ) + max_model_len = self.config.heterogeneous_rollout.max_model_len or self.config.max_seq_length public_host = host if host != "0.0.0.0" else "127.0.0.1" - device_visible_name = "CUDA_VISIBLE_DEVICES" - vllm_module = ( - os.environ.get("VLLM_SERVER_MODULE") - or "vllm.entrypoints.openai.api_server" - ) - values = { - "python": shlex.quote(sys.executable), - "instance_id": shlex.quote(str(inst.instance_id or f"tp{inst.tp}_{port}")), - "model_path": shlex.quote(self.config.model_path), - "initial_model": shlex.quote(self.config.model_path), - "host": shlex.quote(host), - "public_host": shlex.quote(public_host), - "port": port, - "tp": inst.tp, - "gpus": ",".join(str(g) for g in inst.gpus), - "devices": ",".join(str(g) for g in inst.gpus), - "device_visible_env": device_visible_name, - "vllm_module": shlex.quote(vllm_module), - "max_model_len": max_model_len, - "gpu_memory_utilization": self.config.heterogeneous_rollout.gpu_memory_utilization, - "health_url": shlex.quote(f"http://{public_host}:{port}/health"), - "local_health_url": shlex.quote(f"http://127.0.0.1:{port}/health"), - "control_dir": shlex.quote( - getattr( - self.config, - "rollout_weight_sync_control_dir", - "", - ) - or os.path.join(self.config.log_dir, "rollout_weight_sync") - ), - "log_path": shlex.quote( - self._rollout_log_path(str(inst.instance_id or f"tp{inst.tp}_{port}")) - ), - } + module = os.environ.get("VLLM_SERVER_MODULE") or "RL_Framework.scripts.vllm_hot_reload_api_server" + control = str(getattr(self.config, "rollout_weight_sync_control_dir", "") or + Path(self.config.log_dir) / "rollout_weight_sync") + instance_id = str(inst.instance_id or f"tp{inst.tp}_{port}") + devices = ",".join(str(gpu) for gpu in inst.gpus) + values = dict(python=shlex.quote(sys.executable), instance_id=shlex.quote(instance_id), + model_path=shlex.quote(self.config.model_path), initial_model=shlex.quote(self.config.model_path), + host=shlex.quote(host), public_host=shlex.quote(public_host), port=port, tp=inst.tp, + gpus=devices, devices=devices, device_visible_env="CUDA_VISIBLE_DEVICES", + vllm_module=shlex.quote(module), max_model_len=max_model_len, + gpu_memory_utilization=self.config.heterogeneous_rollout.gpu_memory_utilization, + health_url=shlex.quote(f"http://{public_host}:{port}/health"), + local_health_url=shlex.quote(f"http://127.0.0.1:{port}/health"), + control_dir=shlex.quote(control), log_path=shlex.quote(self._rollout_log_path(instance_id))) if template: + # Keep explicit site templates, but placement remains the launch + # helper's responsibility for default local commands. return template.format(**values) - env_prefix = "" - if values["devices"]: - env_prefix = f"{device_visible_name}={shlex.quote(values['devices'])} " - return ( - f"{env_prefix}python3 -m {values['vllm_module']} " - f"--model {values['model_path']} " - f"--host {values['host']} " - f"--port {port} " - f"--tensor-parallel-size {inst.tp} " - f"--max-model-len {max_model_len} " - "--dtype bfloat16 " - f"--gpu-memory-utilization {values['gpu_memory_utilization']} " - "--enforce-eager " - "--trust-remote-code" - ) + if (len(inst.gpus) != int(inst.tp) or len(set(inst.gpus)) != len(inst.gpus) + or any(int(gpu) < 0 for gpu in inst.gpus)): + raise ValueError("Rollout TP workers must own distinct assigned node GPU IDs") + server = [sys.executable, "-m", module, "--model", "__MODEL_PATH__", + "--served-model-name", self.config.model_path, "--host", "0.0.0.0", "--port", str(port), + "--tensor-parallel-size", str(inst.tp), "--max-model-len", str(max_model_len), + "--dtype", "bfloat16", "--gpu-memory-utilization", str(values["gpu_memory_utilization"]), + "--max-num-seqs", str(self.config.max_concurrent_rollouts), "--seed", str(self.config.seed), + "--worker-extension-cls", "RL_Framework.vllm_hot_reload.InplaceReloadWorkerExtension", + "--enforce-eager", "--trust-remote-code", "--disable-log-requests"] + backend = str(getattr(self.config.heterogeneous_rollout.scheduling, "cmlfq_kv_backend", "recompute")) + if backend == "cpu_offload": + storage = os.environ.get("LIBRA_KV_STORAGE_PATH") or str(Path(control) / "kv-store") + server += ["--kv-transfer-config", json.dumps(dict( + kv_connector="LibraCPUOffloadConnector", kv_role="kv_both", + kv_connector_module_path="RL_Framework.infra.scheduling.vllm_cpu_offload", + kv_connector_extra_config=dict(storage_path=storage)), separators=(",", ":"))] + supervisor = [sys.executable, "-m", "RL_Framework.scripts.restartable_vllm_server", + "--instance-id", instance_id, "--control-dir", control, + "--health-url", f"http://127.0.0.1:{port}/health", "--initial-model", self.config.model_path, + "--ready-timeout", str(planner_cfg.vllm_ready_timeout_s), + "--reload-method", "inplace", "--reload-strategy", "parallel", "--", *server] + return shlex.join(["env", "CUDA_VISIBLE_DEVICES=" + devices, *supervisor]) def close(self): """Fence late callbacks, then stop every resource already committed. diff --git a/infra/elastic/state_transfer.py b/infra/elastic/state_transfer.py new file mode 100644 index 0000000..9d77d78 --- /dev/null +++ b/infra/elastic/state_transfer.py @@ -0,0 +1,148 @@ +"""Transient, bounded-frame CPU tensor transport for live EHP state. + +No model or optimizer checkpoint is written to disk. A capability URL names an +immutable in-memory snapshot; owners explicitly release it after installation. +""" +from __future__ import annotations + +import io +import secrets +import socket +import struct +import threading +from urllib.parse import urlparse + +import torch + +_FRAME = struct.Struct('!Q') +_CHUNK = 4 * 1024 * 1024 + + +def _read(sock, count): + result = bytearray(count) + view = memoryview(result) + offset = 0 + while offset < count: + received = sock.recv_into(view[offset:]) + if not received: + raise ConnectionError('live state stream ended early') + offset += received + return result + + +def _pack_tree(value, tensors): + if isinstance(value, torch.Tensor): + index = len(tensors) + tensors.append(value.detach().cpu().contiguous()) + return ('tensor', index) + if isinstance(value, dict): + return ('dict', [(_pack_tree(k, tensors), _pack_tree(v, tensors)) for k, v in value.items()]) + if isinstance(value, (list, tuple)): + return ('tuple' if isinstance(value, tuple) else 'list', [_pack_tree(v, tensors) for v in value]) + return ('value', value) + + +def _unpack_tree(node, tensors): + kind, value = node + if kind == 'tensor': + return tensors[value] + if kind == 'dict': + return {_unpack_tree(k, tensors): _unpack_tree(v, tensors) for k, v in value} + if kind in ('tuple', 'list'): + result = [_unpack_tree(v, tensors) for v in value] + return tuple(result) if kind == 'tuple' else result + return value + + +class LiveStateServer: + def __init__(self, host='0.0.0.0', advertised_host=None): + self._snapshots = {} + self._lock = threading.RLock() + self._stop = threading.Event() + self._sock = socket.socket() + self._sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1) + self._sock.bind((host, 0)) + self._sock.listen() + self._sock.settimeout(0.2) + self.host = advertised_host or socket.gethostbyname(socket.gethostname()) + self.port = self._sock.getsockname()[1] + self._thread = threading.Thread(target=self._serve, daemon=True) + self._thread.start() + + def publish(self, payload): + token = secrets.token_hex(24) + with self._lock: + self._snapshots[token] = payload + return f'ehp://{self.host}:{self.port}/{token}' + + def release(self, uri): + with self._lock: + self._snapshots.pop(urlparse(uri).path.lstrip('/'), None) + + def close(self): + self._stop.set() + self._sock.close() + self._thread.join(timeout=2) + with self._lock: + self._snapshots.clear() + + def _serve(self): + while not self._stop.is_set(): + try: + conn, _ = self._sock.accept() + except socket.timeout: + continue + except OSError: + return + threading.Thread(target=self._send, args=(conn,), daemon=True).start() + + def _send(self, conn): + with conn: + conn.settimeout(300) + try: + length, = _FRAME.unpack(_read(conn, _FRAME.size)) + if length > 256: + raise ValueError('invalid state capability') + token = _read(conn, length).decode() + with self._lock: + payload = self._snapshots.get(token) + if payload is None: + conn.sendall(_FRAME.pack(0)) + return + tensors = [] + tree = _pack_tree(payload, tensors) + metadata = {'tree': tree, 'tensors': [(str(t.dtype).removeprefix('torch.'), tuple(t.shape)) for t in tensors]} + buffer = io.BytesIO() + torch.save(metadata, buffer) + blob = buffer.getvalue() + conn.sendall(_FRAME.pack(len(blob))) + conn.sendall(blob) + for tensor in tensors: + flat = tensor.reshape(-1).view(torch.uint8) + for offset in range(0, flat.numel(), _CHUNK): + conn.sendall(flat[offset:offset + _CHUNK].numpy().tobytes()) + except (OSError, ValueError): + return + + +def receive_live_state(uri, timeout=300): + target = urlparse(uri) + if target.scheme != 'ehp': + raise ValueError('not a live EHP state endpoint') + with socket.create_connection((target.hostname, target.port), timeout=timeout or None) as conn: + token = target.path.lstrip('/').encode() + conn.sendall(_FRAME.pack(len(token))) + conn.sendall(token) + length, = _FRAME.unpack(_read(conn, _FRAME.size)) + if not length: + raise RuntimeError('live state snapshot expired or capability is invalid') + metadata = torch.load(io.BytesIO(_read(conn, length)), weights_only=False) + tensors = [] + for dtype_name, shape in metadata['tensors']: + tensor = torch.empty(shape, dtype=getattr(torch, dtype_name)) + flat = tensor.reshape(-1).view(torch.uint8) + for offset in range(0, flat.numel(), _CHUNK): + count = min(_CHUNK, flat.numel() - offset) + flat[offset:offset + count].copy_(torch.frombuffer(_read(conn, count), dtype=torch.uint8)) + tensors.append(tensor) + return _unpack_tree(metadata['tree'], tensors) diff --git a/infra/execution/batch_dispatcher.py b/infra/execution/batch_dispatcher.py index cbf99c3..a9bff5a 100644 --- a/infra/execution/batch_dispatcher.py +++ b/infra/execution/batch_dispatcher.py @@ -308,6 +308,12 @@ def wait_results( return selected + def peek_ready_data(self) -> list[Any]: + """Observe admitted results without consuming or resubmitting them.""" + with self._result_cv: + ready = sorted(self._pending_results.values(), key=lambda result: result.create_time) + return [result.data for result in ready] + def active_submit_and_wait( self, input_generator: Generator[TInput, None, None], diff --git a/scripts/elastic_hybrid_controller.py b/scripts/elastic_hybrid_controller.py index 7d42b19..d67a919 100644 --- a/scripts/elastic_hybrid_controller.py +++ b/scripts/elastic_hybrid_controller.py @@ -14,6 +14,7 @@ import os import re import signal +import shlex import subprocess import time from pathlib import Path @@ -37,6 +38,36 @@ def _write_json(path: Path, payload: dict) -> None: tmp.replace(path) +def _remote_stop_code(pid): + """Terminate the owned session and verify no live worker descendants remain.""" + return f"""import os, signal, time +pid = {int(pid)} +def live(): + result = [] + for entry in os.listdir('/proc'): + if not entry.isdigit(): + continue + try: + fields = open('/proc/' + entry + '/stat').read().rsplit(')', 1)[1].split() + if (int(fields[2]) == pid or int(entry) == pid) and fields[0] != 'Z': + result.append(int(entry)) + except (FileNotFoundError, ProcessLookupError, PermissionError): + pass + return result +for sig in (signal.SIGTERM, signal.SIGKILL): + for owned in live(): + try: os.kill(owned, sig) + except ProcessLookupError: pass + deadline = time.monotonic() + 10 + while live() and time.monotonic() < deadline: + time.sleep(.1) + if not live(): + break +if live(): + raise RuntimeError('owned worker session still alive') +""" + + class Controller: def __init__(self, args: argparse.Namespace): self.args = args @@ -76,7 +107,7 @@ def _handle_launch_requests(self) -> bool: command = str(request["command"]) log_path = self.task_dir / f"{worker_id}.controller.log" remote_command = ( - f"nohup {command} >{log_path} 2>&1 & echo $!" + f"nohup setsid sh -c {shlex.quote(command)} >{shlex.quote(str(log_path))} 2>&1 & echo $!" ) env = os.environ.copy() env.setdefault("INTERNAL_SSH_TIMEOUT", "30") @@ -108,21 +139,44 @@ def _handle_stop_requests(self) -> bool: handled = False for request_path in sorted(self.task_dir.glob(".stop_*.json")): worker_id = request_path.stem[len(".stop_") :] - request_path.unlink(missing_ok=True) - handled = True - record = self.remote_pids.pop(worker_id, None) - if record is None: + running = request_path.with_suffix(".running") + try: + request_path.replace(running) + except FileNotFoundError: continue - host, pid = record - helper = str(Path(self.args.project_dir) / "scripts" / "internal_ssh.sh") - subprocess.run( - [helper, host, "--", f"kill {int(pid)} 2>/dev/null || true"], - env=os.environ.copy(), - stdout=subprocess.DEVNULL, - stderr=subprocess.DEVNULL, - timeout=30, - check=False, - ) + handled = True + request = _read_json(running) + request_id = request.get("request_id") + try: + # Withdraw launches that the controller has not claimed. + (self.task_dir / f".launch_{worker_id}.json").unlink(missing_ok=True) + record = self.remote_pids.get(worker_id) + if record is None: + started = self.task_dir / f".launch_{worker_id}.started" + if started.exists(): + raise RuntimeError("controller lost process ownership; refusing a stale PID") + host, pid = record if record is not None else (str(request.get("host", "")), -1) + if pid > 0: + helper = str(Path(self.args.project_dir) / "scripts" / "internal_ssh.sh") + code = _remote_stop_code(pid) + proc = subprocess.run( + [helper, host, "--", "python3 -c " + shlex.quote(code)], + env=os.environ.copy(), capture_output=True, text=True, + timeout=45, check=False, + ) + if proc.returncode != 0: + raise RuntimeError(f"remote worker did not exit rc={proc.returncode}: {proc.stderr[-500:]}") + self.remote_pids.pop(worker_id, None) + _write_json(self.task_dir / f".stop_{worker_id}.done", { + "worker_id": worker_id, "host": host, "pid": pid, + "request_id": request_id, "exited": True, "finished_at": time.time(), + }) + except Exception as exc: + _write_json(self.task_dir / f".stop_{worker_id}.error", { + "worker_id": worker_id, "request_id": request_id, "error": str(exc), + }) + finally: + running.unlink(missing_ok=True) return handled def _stop_all_workers(self) -> None: diff --git a/scripts/elastic_hybrid_worker.py b/scripts/elastic_hybrid_worker.py index 992d237..b317cd1 100644 --- a/scripts/elastic_hybrid_worker.py +++ b/scripts/elastic_hybrid_worker.py @@ -64,13 +64,21 @@ def _load_payload_file(path: Path, *, worker_id: str, target_core_id: str, args) ) -def _touch_ready(task_dir: Path, worker_id: str, snapshot_path: str): +def _touch_ready(task_dir: Path, worker_id: str, snapshot_path: str, receipt=None): task_dir.mkdir(parents=True, exist_ok=True) rank = int(os.environ.get("RANK", "0")) ready = task_dir / f"{worker_id}.rank_{rank}.ready" - ready.write_text(f"snapshot={snapshot_path}\npid={os.getpid()}\nrank={rank}\n", encoding="utf-8") + contents = f"snapshot={snapshot_path}\npid={os.getpid()}\nrank={rank}\n" + if receipt: + contents += "".join(f"{key}={value}\n" for key, value in receipt.items()) + temporary = ready.with_suffix(f".tmp.{os.getpid()}") + temporary.write_text(contents, encoding="utf-8") + temporary.replace(ready) if rank == 0: - (task_dir / f"{worker_id}.ready").write_text(ready.read_text(encoding="utf-8"), encoding="utf-8") + summary = task_dir / f"{worker_id}.ready" + temporary = summary.with_suffix(f".tmp.{os.getpid()}") + temporary.write_text(contents, encoding="utf-8") + temporary.replace(summary) def _initialize_megatron_worker(args): @@ -87,10 +95,29 @@ def _initialize_megatron_worker(args): torch.cuda.set_device(int(os.environ.get("LOCAL_RANK", "0"))) if config.train_backend != "megatron_core": raise ValueError("elastic hybrid workers currently require Megatron-Core") + model_width = (int(getattr(config, "train_tp_size", 1)) + * int(getattr(config, "train_pp_size", 1)) + * int(getattr(config, "train_cp_size", 1))) + expert_width = int(getattr(config, "train_ep_size", 1)) + replica_world = model_width * expert_width + launched_world = int(os.environ.get("WORLD_SIZE", getattr(args, "replica_world_size", replica_world))) + if launched_world != replica_world: + raise ValueError(f"Hybrid launch must contain one complete model replica: {launched_world} != {replica_world}") + # In MCore, EP subdivides the dense-DP dimension. One dense replica uses + # DP1; a complete expert replica retains its EP workers and expert-DP1. + config.train_dp_size = expert_width + config.train_gpus = replica_world + config.tp_size = int(getattr(config, "train_tp_size", 1)) engine = create_train_engine(config) engine.initialize(max_seq_length=config.max_seq_length) + actual = engine.get_parallel_state() + if int(actual["train_dp"]) != expert_width or int(getattr(engine, "world_size", launched_world)) != launched_world: + raise RuntimeError("Hybrid factory initialized a different replica topology") + for key in ("tp", "pp", "cp", "ep"): + if int(actual[f"train_{key}"]) != int(getattr(config, f"train_{key}_size", 1)): + raise RuntimeError(f"Hybrid factory changed the requested {key} model parallelism") if not args.initialize_only: - engine.load_elastic_state_snapshot(args.snapshot_path) + engine.load_elastic_state_snapshot(args.snapshot_path, target_core_id=args.target_core_id, replica_rank=int(os.environ.get("RANK", "0"))) engine.hybrid_lockstep_gradient_sync = bool(getattr(config.global_resource_planner, "hybrid_lockstep_gradient_sync", True)) return engine @@ -121,15 +148,16 @@ def _gradient_client(args, endpoint_data): def main(): args = parse_args() - snapshot = Path(args.snapshot_path) - if not snapshot.exists() and not args.initialize_only: + snapshot = args.snapshot_path + if not snapshot.startswith("ehp://") and not Path(snapshot).exists() and not args.initialize_only: raise FileNotFoundError(f"snapshot does not exist: {snapshot}") engine = _initialize_megatron_worker(args) if args.worker_mode == "megatron_core" else None if args.initialize_only: _touch_ready(Path(args.task_dir), args.worker_id, str(snapshot)) return if engine is None: - torch.load(snapshot, map_location="cpu", weights_only=False) + if not snapshot.startswith("ehp://"): + torch.load(snapshot, map_location="cpu", weights_only=False) endpoint = GradientEndpoint(args.gradient_host, args.gradient_port, args.authkey) if engine is not None: client = _gradient_client(args, _endpoint_for_lane(args, engine)) @@ -163,7 +191,7 @@ def main(): if not args.task_dir: raise ValueError("--task-dir or --payload-file is required") task_dir = Path(args.task_dir) - _touch_ready(task_dir, args.worker_id, str(snapshot)) + _touch_ready(task_dir, args.worker_id, str(snapshot), getattr(engine, "_elastic_state_receipt", None)) while True: tasks = sorted(task_dir.glob(f"{args.worker_id}.*.pt")) if not tasks: @@ -175,11 +203,20 @@ def main(): if isinstance(data, dict) and data.get("shutdown"): task_path.rename(task_path.with_suffix(task_path.suffix + ".done")) return + if engine is not None and isinstance(data, dict) and data.get("install_state"): + snapshot = str(data["snapshot_path"]) + receipt = engine.load_elastic_state_snapshot(snapshot, target_core_id=args.target_core_id, replica_rank=int(os.environ.get("RANK", "0"))) + _touch_ready(task_dir, args.worker_id, snapshot, receipt) + if torch.distributed.is_initialized(): + torch.distributed.barrier() + if int(os.environ.get("RANK", "0")) == 0: + task_path.rename(task_path.with_suffix(task_path.suffix + ".done")) + continue if engine is not None and isinstance(data, dict) and "trajectories" in data: task_snapshot = str(data.get("snapshot_path", "")) if task_snapshot and task_snapshot != str(snapshot): - engine.load_elastic_state_snapshot(task_snapshot) - snapshot = Path(task_snapshot) + engine.load_elastic_state_snapshot(task_snapshot, target_core_id=args.target_core_id, replica_rank=int(os.environ.get("RANK", "0"))) + snapshot = task_snapshot engine.recompute_logprobs(data["trajectories"]) payload = engine.compute_elastic_gradient_payload(data["trajectories"], worker_id=args.replica_id or args.worker_id, target_core_id=args.target_core_id, step=int(data["step"]), state_version=int(data["state_version"]), membership_epoch=int(data["membership_epoch"]), global_sample_count=int(data.get("global_sample_count", 0))) else: @@ -188,7 +225,7 @@ def main(): expect_update = bool(engine is not None and getattr(engine, "hybrid_lockstep_gradient_sync", True)) update = client.send(payload, expect_update=expect_update) if update is not None: - if update.replica_id != payload.replica_id or update.step != payload.step or update.membership_epoch != payload.membership_epoch: + if update.replica_id != payload.replica_id or update.step != payload.step or update.membership_epoch != payload.membership_epoch or update.state_version != payload.state_version + 1: raise RuntimeError("received mismatched post-AllReduce update") engine.apply_elastic_gradient_update(update.tensors, state_version=update.state_version) if torch.distributed.is_initialized(): diff --git a/scripts/test_cpu_ci.sh b/scripts/test_cpu_ci.sh index 7df1cfb..f820ef6 100644 --- a/scripts/test_cpu_ci.sh +++ b/scripts/test_cpu_ci.sh @@ -37,4 +37,29 @@ python -m pytest -q --strict-markers --junitxml=reports/cpu-tests.xml \ tests/test_hybrid_optimizer_cuda_cache.py \ tests/test_torch_optimizer_memory.py \ tests/test_megatron_core_train_engine.py \ - tests/test_hybrid_worker_batch_protocol.py + tests/test_hybrid_worker_batch_protocol.py \ + tests/test_core_reconfiguration_protocol.py \ + tests/test_core_subset_publication.py \ + tests/test_ehp_membership_release.py \ + tests/test_ehp_readiness_stop.py \ + tests/test_hybrid_join_readiness_budget.py \ + tests/test_hybrid_boundary_activation.py \ + tests/test_hybrid_sample_statistics.py \ + tests/test_hybrid_task_lifecycle.py \ + tests/test_elastic_hybrid_pool.py \ + tests/test_grp_decision_gates.py \ + tests/test_rollout_lending.py \ + tests/test_rollout_lifecycle.py \ + tests/test_rollout_process.py \ + tests/test_rollout_supply.py \ + tests/test_runtime_boundary_integration.py \ + tests/test_runtime_core_capabilities.py \ + tests/test_runtime_elastic_executor.py \ + tests/test_runtime_blue_green_gpu_ownership.py \ + tests/test_runtime_repartition_eligibility.py \ + tests/test_runtime_executor_shutdown.py \ + tests/test_runtime_rollout_ownership.py \ + tests/test_runtime_failed_snapshot_cleanup.py \ + tests/test_scheduling_audit.py \ + tests/test_publication_preserves_hybrids.py \ + tests/test_grpo_grouping.py diff --git a/tests/test_core_reconfiguration_protocol.py b/tests/test_core_reconfiguration_protocol.py new file mode 100644 index 0000000..f36d234 --- /dev/null +++ b/tests/test_core_reconfiguration_protocol.py @@ -0,0 +1,188 @@ +"""Every parked training rank must execute fixed-topology live-state commands.""" + +import threading +import time +from concurrent.futures import ThreadPoolExecutor +from types import SimpleNamespace + +import pytest + +from RL_Framework.infra.elastic import core_reconfiguration as protocol +from RL_Framework.infra.elastic.runtime_boundary import read_record, write_record + + +class Engine: + def __init__(self, rank, calls, barrier=None, failing_rank=None): + self.rank, self.world_size = rank, 2 + self.calls, self.barrier, self.failing_rank = calls, barrier, failing_rank + + def capture_elastic_state_snapshot(self, worker, target): + self.calls.append((self.rank, "capture", worker, target)) + if self.barrier: + self.barrier.wait(timeout=3) + if self.rank == self.failing_rank: + raise RuntimeError("rank-local capture failed") + return 17 + + def get_elastic_state_snapshot_path(self, revision): + return f"ehp://rank-{self.rank}/revision-{revision}" + + def get_elastic_lane_state(self): + return {"global_rank": self.rank, "tensor_parallel_rank": self.rank, + "data_parallel_rank": 0} + + def release_elastic_state_snapshot(self, revision): + self.calls.append((self.rank, "release", revision)) + + + +class Trainer: + def __init__(self, engine, calls): + self.train_engine = engine + self.calls = calls + self.config = SimpleNamespace(train_tp_size=2, train_dp_size=1, tp_size=2) + self._elastic_gradient_server = SimpleNamespace( + close=lambda: calls.append((engine.rank, "close-gradient-server"))) + + def _runtime_state_is_current(self, command): + return command.get("run_id") == "run-current" + + def _initialize_elastic_communication_domains(self): + self.calls.append((self.train_engine.rank, "new-comm-domain")) + + def _init_nonblocking_elastic_core(self): + self.calls.append((self.train_engine.rank, "new-gradient-server")) + + def _on_train_engine_reconfigured(self): + self.calls.append((self.train_engine.rank, "new-topology-hook")) + + +def executor(directory, engine, trainer=None, timeout=1): + return SimpleNamespace( + train_engine=engine, trainer=trainer, + _coordination_dir=lambda: directory, _active_runtime_coord_id="coord-current", + runtime_run_id="run-current", _closing=threading.Event(), + config=SimpleNamespace(global_resource_planner=SimpleNamespace(runtime_drain_timeout_s=timeout)), + ) + + +def follow_one(trainer, directory): + deadline = time.monotonic() + 4 + while time.monotonic() < deadline: + if read_record(directory / "core_command.json").get("command_id"): + protocol.follow_core_command(trainer, directory, "coord-current") + return + time.sleep(0.001) + raise AssertionError("leader did not publish command") + + +@pytest.mark.parametrize("operation", ["capture", "release"]) +def test_every_training_rank_executes_command(tmp_path, operation): + calls = [] + barrier = threading.Barrier(2) if operation != "release" else None + engines = [Engine(rank, calls, barrier) for rank in range(2)] + trainers = [Trainer(engine, calls) for engine in engines] + with ThreadPoolExecutor(max_workers=1) as pool: + peer = pool.submit(follow_one, trainers[1], tmp_path) + receipts = protocol.dispatch_core_command( + executor(tmp_path, engines[0], trainers[0]), operation, + target={"train_tp_size": 1, "train_dp_size": 2}, revisions=(16, 17), + ) + peer.result(timeout=4) + assert {item["rank"] for item in receipts} == {0, 1} + assert all(item["status"] == "success" for item in receipts) + assert len({item["command_id"] for item in receipts}) == 1 + if operation == "capture": + assert {item["snapshot_uri"] for item in receipts} == { + "ehp://rank-0/revision-17", "ehp://rank-1/revision-17"} + assert all(item["optimizer_revision"] == 17 for item in receipts) + elif operation == "release": + assert {(item[0], item[2]) for item in calls if item[1] == "release"} == { + (0, 16), (0, 17), (1, 16), (1, 17)} + + +def test_one_rank_failure_reaches_dispatch_and_follower(tmp_path): + calls, barrier = [], threading.Barrier(2) + engines = [Engine(rank, calls, barrier, failing_rank=1) for rank in range(2)] + trainers = [Trainer(engine, calls) for engine in engines] + with ThreadPoolExecutor(max_workers=1) as pool: + peer = pool.submit(follow_one, trainers[1], tmp_path) + with pytest.raises(RuntimeError, match="collective core capture failed"): + protocol.dispatch_core_command(executor(tmp_path, engines[0], trainers[0]), "capture", + target={"train_tp_size": 1, "train_dp_size": 2}) + with pytest.raises(RuntimeError, match="rank-local capture failed"): + peer.result(timeout=4) + failed = read_record(tmp_path / "core_receipts" / "rank_1.json") + assert failed["status"] == "failed" + assert trainers[1].config.train_tp_size == 2 + + +def test_followers_ignore_stale_or_repeated_commands(tmp_path): + calls = [] + trainer = Trainer(Engine(1, calls), calls) + command = {"command_id": "release-once", "coord_id": "coord-current", + "run_id": "run-stale", "operation": "release", "revisions": [17]} + write_record(tmp_path / "core_command.json", command) + protocol.follow_core_command(trainer, tmp_path, "coord-current") + assert calls == [] + command["run_id"] = "run-current" + write_record(tmp_path / "core_command.json", command) + protocol.follow_core_command(trainer, tmp_path, "wrong-coord") + assert calls == [] + protocol.follow_core_command(trainer, tmp_path, "coord-current") + protocol.follow_core_command(trainer, tmp_path, "coord-current") + assert calls == [(1, "release", 17)] + + +@pytest.mark.parametrize("wrong", [{"coord_id": "coord-stale"}, {"rank": 0}, {"run_id": "run-stale"}]) +def test_rank_coverage_requires_current_coord_and_exact_rank(tmp_path, monkeypatch, wrong): + original = protocol.execute_core_command + + def local_and_invalid_peer(engine, command, directory, trainer): + receipt = original(engine, command, directory, trainer) + stale = dict(receipt, rank=1) + stale.update(wrong) + write_record(directory / "core_receipts" / "rank_1.json", stale) + return receipt + + monkeypatch.setattr(protocol, "execute_core_command", local_and_invalid_peer) + with pytest.raises(TimeoutError, match="missing rank receipts"): + protocol.dispatch_core_command(executor(tmp_path, Engine(0, []), timeout=0.02), "release") + + +def test_zero_timeout_waits_for_delayed_rank(tmp_path, monkeypatch): + original = protocol.execute_core_command + peer_thread = [] + + def delayed_peer(engine, command, directory, trainer): + receipt = original(engine, command, directory, trainer) + + def publish(): + time.sleep(0.05) + write_record(directory / "core_receipts" / "rank_1.json", dict(receipt, rank=1)) + + thread = threading.Thread(target=publish) + thread.start() + peer_thread.append(thread) + return receipt + + monkeypatch.setattr(protocol, "execute_core_command", delayed_peer) + receipts = protocol.dispatch_core_command(executor(tmp_path, Engine(0, []), timeout=0), "release") + assert len(receipts) == 2 + peer_thread[0].join(timeout=1) + + +def test_closing_executor_does_not_report_partial_success(tmp_path): + instance = executor(tmp_path, Engine(0, [])) + instance._closing.set() + with pytest.raises(RuntimeError, match="executor closed"): + protocol.dispatch_core_command(instance, "release") + + +def test_topology_rebuild_is_outside_fixed_core_protocol(tmp_path): + receipt = protocol.execute_core_command(Engine(0, []), { + "command_id": "unsupported", "coord_id": "coord-current", + "run_id": "run-current", "operation": "rebuild", + }, tmp_path) + assert receipt["status"] == "failed" + assert "unknown collective core operation: rebuild" in receipt["error"] diff --git a/tests/test_core_subset_publication.py b/tests/test_core_subset_publication.py new file mode 100644 index 0000000..9b9e33f --- /dev/null +++ b/tests/test_core_subset_publication.py @@ -0,0 +1,68 @@ +"""Rollback publication targets only newly restored services, preserving loans.""" +from types import SimpleNamespace as NS + +import pytest + +from RL_Framework.infra.elastic.core_reconfiguration import _run_operation +from RL_Framework.infra.elastic.runtime_executor import RuntimeElasticExecutor + + +@pytest.mark.parametrize('fail', [False, True]) +@pytest.mark.parametrize('preexisting', [False, True]) +def test_subset_publication_restores_temporary_scope_and_keeps_old_loan(fail, preexisting): + borrowed = {'old-hybrid-engine'} + unavailable = borrowed | {'returned-engine'} + topology = object() + engine = NS(optimizer_revision=7) + trainer = NS(config=topology, _rollout_unavailable_ids=unavailable, + rollout_engine=NS(scheduler=NS(_instances=[]))) + previous_targets, previous_pending = {'previous-target'}, {'previous-pending'} + if preexisting: + trainer._publication_target_ids = previous_targets + trainer._restored_rollouts_pending_publication = previous_pending + observed = [] + + def sync(*, target_version): + observed.append((target_version, set(trainer._publication_target_ids), + set(trainer._restored_rollouts_pending_publication))) + if fail: + raise RuntimeError('receiver failed') + # The real NCCL trainer invokes this helper after sender ACK. + from RL_Framework.infra.elastic.rollout_lifecycle import reopen_after_publication + reopen_after_publication(trainer) + + trainer._sync_weights = sync + command = dict(operation='publish', target=dict(scope='returned_subset', + instances=[dict(instance_id='returned-engine')])) + if fail: + with pytest.raises(RuntimeError, match='receiver failed'): + _run_operation(engine, command, trainer) + else: + assert _run_operation(engine, command, trainer)['acknowledged'] + assert observed == [(7, {'returned-engine'}, {'returned-engine'})] + assert trainer.config is topology + assert trainer._rollout_unavailable_ids is unavailable + if preexisting: + assert trainer._publication_target_ids is previous_targets + assert trainer._restored_rollouts_pending_publication is previous_pending + else: + assert not hasattr(trainer, '_publication_target_ids') + assert not hasattr(trainer, '_restored_rollouts_pending_publication') + + +def test_runtime_subset_targets_are_selected_from_current_configuration(monkeypatch): + from RL_Framework.infra.elastic import core_reconfiguration as commands + runtime = object.__new__(RuntimeElasticExecutor) + instances = [NS(instance_id=name, tp=1, gpus=[index], host='node', port=8000+index) + for index, name in enumerate(('old-hybrid-engine', 'returned-engine', 'unaffected-engine'))] + runtime.config = NS(heterogeneous_rollout=NS(instances=instances)) + calls = [] + monkeypatch.setattr(commands, 'dispatch_core_command', + lambda *args, **kwargs: calls.append((args, kwargs))) + runtime._publish_current_rollout_weights(instance_ids=['returned-engine']) + target = calls[0][1]['target'] + assert target['scope'] == 'returned_subset' + assert [entry['instance_id'] for entry in target['instances']] == ['returned-engine'] + with pytest.raises(ValueError, match='unknown instances'): + runtime._publish_current_rollout_weights(instance_ids=['no-longer-configured']) + assert len(calls) == 1 diff --git a/tests/test_ehp_batch_partition.py b/tests/test_ehp_batch_partition.py index 86b5978..9e41efc 100644 --- a/tests/test_ehp_batch_partition.py +++ b/tests/test_ehp_batch_partition.py @@ -122,7 +122,7 @@ def collect(data_gen, batch_size): assert sum(len(task["trajectories"]) for task in tasks) + expected == 16 assert all(task["global_sample_count"] == 16 for task in tasks) assert all(task["sample_count"] == len(task["trajectories"]) for task in tasks) - assert all(task["state_version"] == 0 for task in tasks) + assert all(task["state_version"] == step for task in tasks) all_groups = trainer.train_engine.updated[step] + [x["grpo_group_id"] for task in tasks for x in task["trajectories"]] assert len(all_groups) == 16 assert len(set(all_groups)) == 4 @@ -208,3 +208,17 @@ def test_rebind_preserves_admitted_results_and_incomplete_group_members(tmp_path trainer._reset_rollout_pipeline_after_reconfigure() assert trainer._pending_grpo_groups is pending assert len(pending["group"]) == 3 + + +def test_dispatch_rejects_membership_generation_change_before_update(tmp_path): + trainer = make_trainer(tmp_path, hybrids=["hybrid0"]) + trainer.train_engine.elastic_gradient_domain = SimpleNamespace( + membership_state=lambda: {"membership_epochs": {"hybrid0": 7}}) + (tmp_path / "membership").mkdir() + (tmp_path / "membership" / "hybrid0.json").write_text(json.dumps( + dict(run_id="fixed-b-test", role="hybrid_training", membership_epoch=8))) + trainer._collect_complete_grpo_batch = lambda generator, count: trajectories("batch", count) + with pytest.raises(RuntimeError, match="generation changed"): + trainer.train(None, dataset=[1]) + assert not trainer.train_engine.updated + assert not (tmp_path / "hybrid0.step_0.pt").exists() diff --git a/tests/test_ehp_membership_release.py b/tests/test_ehp_membership_release.py index 44cb4b1..c4b0d05 100644 --- a/tests/test_ehp_membership_release.py +++ b/tests/test_ehp_membership_release.py @@ -7,6 +7,7 @@ from types import SimpleNamespace import pytest +import torch from RL_Framework.config import AsyncRLConfig from RL_Framework.engine.megatron_core_train_engine import MegatronCoreTrainEngine @@ -51,8 +52,14 @@ def make(run_id="run-a"): engine.get_data_parallel_rank = lambda: 0 engine._pending_hybrid_gradients = [] engine._hybrid_gradient_condition = threading.Condition() - engine._elastic_active_gradient_timeout_s = 0.0 + engine._elastic_active_gradient_timeout_s = 0.05 engine._elastic_step_hybrid_workers = () + engine.optimizer_revision = 0 + engine.current_version = 0 + engine._elastic_seen_payloads = set() + engine._elastic_batch_partition = None + engine._elastic_gradient_update_callback = None + engine._model_main_gradients = lambda: [torch.zeros(1)] peer = object.__new__(AsyncRLTrainer) peer.config = config peer.train_engine = engine @@ -69,8 +76,14 @@ def make(run_id="run-a"): pool.close() +def _installed(*args): + # No real model/worker in these membership-ordering tests. Simulate the + # successful installation receipt; the readiness suite verifies real files. + return {"acknowledged": True, "state_version": 0, "optimizer_revision": 0} + + def activate(harness, replica="replica0", worker="rollout0"): - harness.pool.join_replica(replica, [worker], "dp0").result(timeout=3) + harness.pool.join_replica(replica, [worker], "dp0", replica_activation_barrier=_installed).result(timeout=3) harness.executor.hybrid_runtime_state() harness.peer._prepare_elastic_training_step(0) @@ -172,9 +185,9 @@ def activating(*args): exit_record = record(h, worker) if rejoin: if replica: - newer = h.pool.join_replica(worker, ["rollout0"], "dp0") + newer = h.pool.join_replica(worker, ["rollout0"], "dp0", replica_activation_barrier=_installed) else: - newer = h.pool.join_training(worker, "dp0") + newer = h.pool.join_training(worker, "dp0", activation_barrier=_installed) newer.result(timeout=3) h.executor.hybrid_runtime_state() expected_record = record(h, worker) @@ -210,9 +223,9 @@ def pause_after_domain_activation(worker): monkeypatch.setattr(h.pool.gradient_domain, "mark_active", pause_after_domain_activation) worker = "replica0" if replica else "rollout0" if replica: - handle = h.pool.join_replica(worker, ["rollout0"], "dp0") + handle = h.pool.join_replica(worker, ["rollout0"], "dp0", replica_activation_barrier=_installed) else: - handle = h.pool.join_training(worker, "dp0") + handle = h.pool.join_training(worker, "dp0", activation_barrier=_installed) h.executor._pending_hybrid_joins[worker] = handle h.executor._pending_hybrid_join_started_at[worker] = time.time() with ThreadPoolExecutor(max_workers=1) as workers: diff --git a/tests/test_ehp_readiness_stop.py b/tests/test_ehp_readiness_stop.py new file mode 100644 index 0000000..e8c886e --- /dev/null +++ b/tests/test_ehp_readiness_stop.py @@ -0,0 +1,231 @@ +from contextlib import nullcontext +import json +from pathlib import Path +from types import SimpleNamespace +import subprocess +import threading +import time + +import pytest + +from RL_Framework.infra.elastic.readiness import clear_ready_receipts, read_ready_receipts, snapshot_contract +from RL_Framework.infra.elastic.runtime_executor import ( + ManagedHybridWorkerProcess, ManagedRolloutProcess, RuntimeElasticExecutor, +) + + +def _executor(tmp_path, **options): + engine = RuntimeElasticExecutor.__new__(RuntimeElasticExecutor) + cfg = dict(hybrid_worker_task_dir=str(tmp_path), hybrid_worker_remote_control_dir=str(tmp_path), + hybrid_worker_ready_timeout_s=.15, hybrid_worker_stop_timeout_s=.15, + hybrid_worker_remote_control_enabled=False, vllm_stop_timeout_s=.1, + vllm_stop_command_template='stop {instance_id}') + cfg.update(options) + engine.config = SimpleNamespace(global_resource_planner=SimpleNamespace(**cfg)) + engine._hybrid_worker_meta = {} + engine._hybrid_worker_processes = {} + engine._closing = threading.Event() + engine._lifecycle_lock = threading.RLock() + engine._hybrid_resource_guard = lambda *args, **kwargs: nullcontext() + return engine + + +def _manifest(tmp_path): + path = tmp_path / 'snapshot.json' + path.write_text(json.dumps({'format': 'libra-ehp-lanes-v1', 'optimizer_revision': 9, + 'snapshot_id': 'actual-boundary', 'lanes': {'dp1': {'0': 'ehp://a', '1': 'ehp://b'}}})) + return str(path) + + +def _ready(tmp_path, snapshot, lane_rank, **overrides): + fields = dict(snapshot=snapshot, snapshot_id='actual-boundary', lane_snapshot_id=f'lane-{lane_rank}', + acknowledged='True', state_version=9, optimizer_revision=9, target_core_id='dp1', + rank=lane_rank, pid=123 + lane_rank, source_uri='ehp://a' if lane_rank == 0 else 'ehp://b') + fields.update(overrides) + (tmp_path / f'h.rank_{lane_rank}.ready').write_text(''.join(f'{key}={value}\n' for key, value in fields.items())) + + +def test_requires_every_model_lane_and_rejects_generic_virtual_ready(tmp_path): + snapshot = _manifest(tmp_path) + engine = _executor(tmp_path) + (tmp_path / 'h.ready').write_text('ready') + _ready(tmp_path, snapshot, 0) + with pytest.raises(TimeoutError, match='all 2 lanes'): + engine._wait_hybrid_worker_ready('h', state_version=9, replica_world_size=2, + snapshot_path=snapshot, target_core_id='dp1') + _ready(tmp_path, snapshot, 1) + receipt = engine._wait_hybrid_worker_ready('h', state_version=9, replica_world_size=2, + snapshot_path=snapshot, target_core_id='dp1') + assert receipt['acknowledged'] and len(receipt['rank_receipts']) == 2 + + +@pytest.mark.parametrize('change', [dict(snapshot_id='previous-boundary'), dict(optimizer_revision=8), + dict(state_version=8), dict(rank=5), dict(acknowledged='False'), + dict(target_core_id='dp0'), dict(lane_snapshot_id=''), + dict(source_uri='ehp://wrong-lane'), dict(pid=123)]) +def test_stale_or_wrong_lane_receipt_cannot_release_activation(tmp_path, change): + snapshot = _manifest(tmp_path) + contract = snapshot_contract(snapshot, state_version=9, target_core_id='dp1') + _ready(tmp_path, snapshot, 0) + _ready(tmp_path, snapshot, 1, **change) + assert read_ready_receipts(tmp_path, 'h', 2, contract) is None + + +def test_timeout_zero_waits_for_actual_replacement_ack(tmp_path): + snapshot = _manifest(tmp_path) + engine = _executor(tmp_path, hybrid_worker_ready_timeout_s=0) + _ready(tmp_path, snapshot, 0, snapshot_id='old') + result = [] + worker = threading.Thread(target=lambda: result.append(engine._wait_hybrid_worker_ready( + 'h', state_version=9, replica_world_size=1, snapshot_path=snapshot, target_core_id='dp1'))) + worker.start() + time.sleep(.06) + assert worker.is_alive() and not result + _ready(tmp_path, snapshot, 0) + worker.join(2) + assert not worker.is_alive() + assert result[0]['optimizer_revision'] == 9 + + +def test_clear_ready_before_restart_removes_same_snapshot_files(tmp_path): + snapshot = _manifest(tmp_path) + _ready(tmp_path, snapshot, 0) + (tmp_path / 'h.ready').write_text('old') + clear_ready_receipts(tmp_path, 'h') + assert not list(tmp_path.glob('*.ready')) + + +def test_singleton_activation_returns_boundary_approved_receipt(tmp_path): + engine = _executor(tmp_path) + snapshot = _manifest(tmp_path) + engine._snapshot_path_for_version = lambda version: snapshot + engine._cluster_swap_training_slot = lambda worker: {} + engine._is_hybrid_worker_running = lambda worker: True + engine._hybrid_worker_meta['h'] = ManagedHybridWorkerProcess('h', 'dp1', 'existing', 1, snapshot, replica_id='h') + initial = {'acknowledged': True, 'optimizer_revision': 9, 'state_version': 9} + approved = dict(initial, optimizer_revision=10, state_version=10) + engine._wait_hybrid_worker_ready = lambda *args, **kwargs: initial + queued = [] + engine._queue_activation = lambda replica, receipt: queued.append((replica, receipt)) or approved + result = engine._activate_hybrid_worker('h', 'dp1', 9) + assert result is approved and queued == [('h', initial)] + + +def test_remote_stop_retains_ownership_until_matching_exit_ack_and_ignores_stale_done(tmp_path): + engine = _executor(tmp_path, hybrid_worker_remote_control_enabled=True, hybrid_worker_stop_timeout_s=0) + meta = ManagedHybridWorkerProcess('h', 'dp1', 'remote', 1234, 'state', host='node') + engine._hybrid_worker_meta = {'h': meta, 'member': meta} + (tmp_path / '.stop_h.done').write_text(json.dumps({'request_id': 'stale', 'worker_id': 'h', 'exited': True})) + observed = [] + def controller(): + deadline = time.monotonic() + 2 + while not (tmp_path / '.stop_h.json').exists() and time.monotonic() < deadline: + time.sleep(.01) + request = json.loads((tmp_path / '.stop_h.json').read_text()) + observed.append(request) + assert engine._hybrid_worker_meta['h'] is meta + time.sleep(.05) + (tmp_path / '.stop_h.done').write_text(json.dumps(dict(request, exited=True))) + thread = threading.Thread(target=controller) + thread.start() + engine._stop_hybrid_worker_process('member') + thread.join(2) + assert observed[0]['request_id'] != 'stale' + assert not engine._hybrid_worker_meta + + +def test_stop_timeout_keeps_ownership_and_retry_uses_same_request(tmp_path): + engine = _executor(tmp_path, hybrid_worker_remote_control_enabled=True) + meta = ManagedHybridWorkerProcess('h', 'dp1', 'remote', 1234, 'state', host='node') + engine._hybrid_worker_meta['h'] = meta + with pytest.raises(TimeoutError): + engine._stop_hybrid_worker_process('h') + assert engine._hybrid_worker_meta['h'] is meta + request = json.loads((tmp_path / '.stop_h.json').read_text()) + (tmp_path / '.stop_h.done').write_text(json.dumps(dict(request, exited=True))) + engine._stop_hybrid_worker_process('h') + assert meta.stop_request_id == request['request_id'] + assert not engine._hybrid_worker_meta + + +@pytest.mark.parametrize('timeout', [False, True]) +def test_external_rollout_stop_failure_propagates(tmp_path, monkeypatch, timeout): + engine = _executor(tmp_path, vllm_stop_timeout_s=0) + meta = ManagedRolloutProcess('engine', 'command', 1, [0], 8000, log_path=str(tmp_path / 'stop.log')) + def run(*args, **kwargs): + assert kwargs['timeout'] is None + if timeout: + raise subprocess.TimeoutExpired('stop', 1) + return SimpleNamespace(returncode=17) + monkeypatch.setattr(subprocess, 'run', run) + with pytest.raises(TimeoutError if timeout else RuntimeError): + engine._run_external_stop(meta) + + +def test_local_stop_waits_for_launcher_and_child_processes(): + import sys + from RL_Framework.infra.elastic.process_lifecycle import _live_group_members, stop_local_worker + script = "import subprocess,sys,time; child=subprocess.Popen([sys.executable,'-c','import time;time.sleep(60)']); print(child.pid,flush=True); time.sleep(60)" + proc = subprocess.Popen([sys.executable, '-c', script], stdout=subprocess.PIPE, text=True, start_new_session=True) + try: + child_pid = int(proc.stdout.readline()) + assert child_pid in _live_group_members(proc.pid) + stop_local_worker(proc, timeout=1) + assert proc.poll() is not None + assert _live_group_members(proc.pid) == [] + finally: + if proc.poll() is None: + proc.kill() + proc.wait() + + +def test_hybrid_factory_partitions_core_dp2_into_one_complete_model_replica(monkeypatch): + import torch + from RL_Framework.config import AsyncRLConfig + from RL_Framework.engine import train_factory + from RL_Framework.scripts.elastic_hybrid_worker import _initialize_megatron_worker + config = SimpleNamespace(train_backend='megatron_core', max_seq_length=32768, + train_tp_size=2, train_pp_size=1, train_cp_size=1, train_ep_size=1, + train_dp_size=2, train_gpus=4, micro_batch_size=1, learning_rate=1e-6, + global_resource_planner=SimpleNamespace(hybrid_lockstep_gradient_sync=True)) + observed = [] + engine = SimpleNamespace(initialize=lambda **kwargs: None, world_size=2, + get_parallel_state=lambda: dict(train_tp=2, train_pp=1, train_cp=1, train_ep=1, train_dp=1)) + def factory(cfg): + observed.append((cfg.train_tp_size, cfg.train_dp_size, cfg.train_gpus, cfg.learning_rate)) + return engine + monkeypatch.setattr(AsyncRLConfig, 'from_yaml', lambda path: config) + monkeypatch.setattr(train_factory, 'create_train_engine', factory) + monkeypatch.setattr(torch.cuda, 'is_available', lambda: True) + monkeypatch.setattr(torch.cuda, 'set_device', lambda device: None) + monkeypatch.setenv('WORLD_SIZE', '2') + result = _initialize_megatron_worker(SimpleNamespace(config='current-core-config.yaml', initialize_only=True)) + assert result is engine + assert observed == [(2, 1, 2, 1e-6)] + monkeypatch.setenv('WORLD_SIZE', '4') + with pytest.raises(ValueError, match='one complete model replica'): + _initialize_megatron_worker(SimpleNamespace(config='current-core-config.yaml', initialize_only=True)) + + +def test_remote_member_alias_stops_actual_replica_controller_key(tmp_path, monkeypatch): + from RL_Framework.infra.elastic import rollout_lending + engine = _executor(tmp_path, hybrid_worker_remote_control_enabled=True) + canonical = ManagedHybridWorkerProcess('replica0', 'dp1', 'remote group', 1234, 'state', host='node') + engine._hybrid_worker_meta['replica0'] = canonical + for rank in range(2): + worker = f'rollout{rank}' + engine._hybrid_worker_meta[worker] = ManagedHybridWorkerProcess( + worker, 'dp1', 'remote group', -1, 'state', host='node', replica_id='replica0', replica_rank=rank, + replica_world_size=2, + ) + stopped = [] + def acknowledge(executor, worker, request_id=None): + stopped.append(worker) + assert len(executor._hybrid_worker_meta) == 3 + return {'worker_id': worker, 'request_id': request_id, 'exited': True} + monkeypatch.setattr(rollout_lending, 'wait_hybrid_stopped', acknowledge) + engine._stop_hybrid_worker_process('rollout1') + assert stopped == ['replica0'] + assert (tmp_path / '.stop_replica0.json').exists() + assert not (tmp_path / '.stop_rollout1.json').exists() + assert not engine._hybrid_worker_meta diff --git a/tests/test_elastic_gradient_idle_path.py b/tests/test_elastic_gradient_idle_path.py index b1f6e25..ca345c7 100644 --- a/tests/test_elastic_gradient_idle_path.py +++ b/tests/test_elastic_gradient_idle_path.py @@ -38,7 +38,7 @@ def test_mcore_dp1_reuses_main_gradient_for_weighted_hybrid_update(): engine._hybrid_gradient_condition = threading.Condition() engine._elastic_batch_partition = (2, 4, {'hybrid': 2}) engine._elastic_training_step = 3 - engine.current_version = 2 + engine._elastic_expected_state_version = 2 # This is the already weighted, finalized Core contribution. Passing it # through an out-of-place divide-by-one would waste a full model buffer. grad = torch.full((73,), 3.) @@ -59,7 +59,7 @@ def checked_reduce(**kwargs): assert not engine._pending_hybrid_gradients assert len(acknowledgements) == 1 assert acknowledgements[0].tensors[0].data_ptr() != grad.data_ptr() - assert acknowledgements[0].state_version == 2 + assert acknowledgements[0].state_version == 3 grad.zero_() assert torch.equal(acknowledgements[0].tensors[0], torch.full((73,), 6.5)) diff --git a/tests/test_elastic_hybrid_pool.py b/tests/test_elastic_hybrid_pool.py index a6f6099..164369d 100644 --- a/tests/test_elastic_hybrid_pool.py +++ b/tests/test_elastic_hybrid_pool.py @@ -109,6 +109,7 @@ def slow_fetch(worker_id, target_core_id): return 12 pool = ElasticHybridPool( + require_state_ack=False, # Standalone state-machine simulation, no engine. core_train_workers=["core0"], core_rollout_workers=["rollout0"], snapshot_fetcher=slow_fetch, @@ -137,6 +138,7 @@ def slow_fetch(worker_id, target_core_id): pool.close() + def test_transport_falls_back_to_local_clone_without_distributed_init(): transport = TorchDistributedRDMATransport(chunk_bytes=8) tensor = torch.arange(16, dtype=torch.float32) @@ -260,6 +262,7 @@ def activation_barrier(worker_id, target_core_id, version): assert release.wait(timeout=2.0) pool = ElasticHybridPool( + require_state_ack=False, # Standalone state-machine simulation, no engine. core_train_workers=["core0"], core_rollout_workers=["rollout0"], snapshot_fetcher=lambda *_args: 5, @@ -286,6 +289,7 @@ def activation_barrier(worker_id, target_core_id, version): pool.close() + def test_complete_dp_replica_rejoins_non_blocking_and_activates_atomically(): release = threading.Event() @@ -295,6 +299,7 @@ def fetch(worker_id, _target_core_id): return 23 pool = ElasticHybridPool( + require_state_ack=False, # Standalone state-machine simulation, no engine. core_train_workers=["core0", "core1"], core_rollout_workers=["rollout0", "rollout1"], snapshot_fetcher=fetch, @@ -324,6 +329,7 @@ def fetch(worker_id, _target_core_id): pool.close() + def test_dp_replica_join_failure_rolls_back_every_member(): def fetch(worker_id, _target_core_id): if worker_id == "rollout1": @@ -331,6 +337,7 @@ def fetch(worker_id, _target_core_id): return 7 pool = ElasticHybridPool( + require_state_ack=False, # Standalone state-machine simulation, no engine. core_train_workers=["core0"], core_rollout_workers=["rollout0", "rollout1"], snapshot_fetcher=fetch, @@ -353,8 +360,10 @@ def fetch(worker_id, _target_core_id): pool.close() + def test_released_dp_replica_can_rejoin_with_a_new_generation(): pool = ElasticHybridPool( + require_state_ack=False, # Standalone state-machine simulation, no engine. core_train_workers=["core0"], core_rollout_workers=["rollout0", "rollout1"], snapshot_fetcher=lambda *_args: 4, @@ -376,6 +385,7 @@ def test_released_dp_replica_can_rejoin_with_a_new_generation(): pool.close() + def test_cancelled_join_cannot_become_active_after_grp_reversal(): release = threading.Event() @@ -384,6 +394,7 @@ def fetch(*_args): return 9 pool = ElasticHybridPool( + require_state_ack=False, # Standalone state-machine simulation, no engine. core_train_workers=["core0"], core_rollout_workers=["rollout0"], snapshot_fetcher=fetch, @@ -401,11 +412,13 @@ def fetch(*_args): pool.close() + def test_failed_join_rolls_worker_back_to_rollout(): def fetch(*_args): raise RuntimeError("snapshot unavailable") pool = ElasticHybridPool( + require_state_ack=False, # Standalone state-machine simulation, no engine. core_train_workers=["core0"], core_rollout_workers=["rollout0"], snapshot_fetcher=fetch, @@ -421,6 +434,7 @@ def fetch(*_args): pool.close() + def test_required_isolated_domain_rejects_missing_elastic_group(monkeypatch): hybrid_pool = importlib.import_module("RL_Framework.infra.elastic.hybrid_pool") @@ -489,3 +503,33 @@ def test_decoupled_domain_membership_changes_do_not_mutate_core_domain(): ] == [1, 2, 3] assert domains.core_process_group == "fixed-megatron-dp" assert domains.hybrid_process_group == "elastic-side-channel" + + +def test_unacknowledged_snapshot_never_activates(): + pool = ElasticHybridPool(core_train_workers=['dp0'], hybrid_workers=['h'], snapshot_fetcher=lambda *_: 8) + try: + handle = pool.join_training('h', 'dp0', activation_barrier=lambda *_: None) + with pytest.raises(RuntimeError, match='ACK'): + handle.result(timeout=2) + assert pool.snapshot()['h'].role == ReplicaRole.HYBRID_ROLLOUT + finally: + pool.close() + + + +def test_actual_install_ack_can_catch_up_beyond_snapshot(): + installed = threading.Event() + def activate(*_): + installed.wait(2) + return {'acknowledged': True, 'state_version': 11, 'optimizer_revision': 11} + pool = ElasticHybridPool(core_train_workers=['dp0'], hybrid_workers=['h'], snapshot_fetcher=lambda *_: 8) + try: + handle = pool.join_training('h', 'dp0', activation_barrier=activate) + time.sleep(.03) + assert not handle.done() + assert pool.gradient_domain.is_joining('h') + installed.set() + assert handle.result(timeout=2).state_version == 11 + finally: + pool.close() + diff --git a/tests/test_fixed_batch_gradients.py b/tests/test_fixed_batch_gradients.py index 0eb3246..bfd2dfa 100644 --- a/tests/test_fixed_batch_gradients.py +++ b/tests/test_fixed_batch_gradients.py @@ -267,3 +267,46 @@ def test_full_engine_gloo_prescale_finalize_elastic_sum_matches_full_batch_with_ if process.is_alive(): process.terminate() process.join() + + +@pytest.mark.parametrize('change,message', [ + ({'step': -1}, 'step'), ({'state_version': -1}, 'revision'), + ({'membership_epoch': 6}, 'epoch'), ({'global_sample_count': 16}, 'batch size'), + ({'zero_placeholder': True}, 'placeholder'), +]) +def test_rejects_stale_or_incomplete_gradient_identity(change, message): + with pytest.raises(ValueError, match=message): + _domain().reduce_core_gradients( + core_gradients={'dp0': (torch.ones(1),), 'dp1': (torch.ones(1),)}, + core_sample_counts={'dp0': 4, 'dp1': 2}, global_sample_count=8, + hybrid_payloads=[_payload(**change)], step=3, state_version=2, + ) + + + +def test_rejects_duplicate_payload_before_double_counting(): + domain = _domain() + kwargs = dict(core_gradients={'dp0': (torch.ones(1),), 'dp1': (torch.ones(1),)}, + hybrid_payloads=[_payload(), _payload()], step=3, state_version=2) + with pytest.raises(ValueError, match='duplicate'): + domain.reduce_core_gradients(**kwargs) + + + +def test_engine_ingress_rejects_unknown_revision_and_duplicate_without_waking_waiters(): + engine = MegatronCoreTrainEngine.__new__(MegatronCoreTrainEngine) + engine.elastic_gradient_domain = _domain() + engine.optimizer_revision = 2 + engine._elastic_seen_payloads = set() + engine._pending_hybrid_gradients = [] + engine._hybrid_gradient_condition = threading.Condition() + engine.set_elastic_training_step(3, ['hybrid']) + engine.set_elastic_batch_partition(4, 8, {'hybrid': 2}) + for bad in (_payload(state_version=-1), _payload(step=-1), _payload(membership_epoch=6)): + with pytest.raises(ValueError): + engine.enqueue_hybrid_gradient_payload(bad) + assert not engine._pending_hybrid_gradients + engine.enqueue_hybrid_gradient_payload(_payload()) + with pytest.raises(ValueError, match='duplicate'): + engine.enqueue_hybrid_gradient_payload(_payload()) + assert len(engine._pending_hybrid_gradients) == 1 diff --git a/tests/test_grp_decision_gates.py b/tests/test_grp_decision_gates.py new file mode 100644 index 0000000..e03df56 --- /dev/null +++ b/tests/test_grp_decision_gates.py @@ -0,0 +1,103 @@ +"""EHP demand cannot authorize a rejected full resource repartition.""" + +from unittest.mock import Mock + +import pytest + +from RL_Framework.config import AsyncRLConfig +from RL_Framework.infra.cost_model.global_resource_planner import ( + GlobalResourcePlan, + GlobalResourcePlanner, + PlannerDecision, +) +from RL_Framework.infra.cost_model.model import RolloutClusterConfig, TrainParallelConfig +from RL_Framework.infra.cost_model.optimizer import OptimizationResult + + +def planner_fixture(monkeypatch, reconfiguration_cost): + for variable in ("GRP_FORCE_RUNTIME_RECONFIGURE", "GRP_FORCE_ROLLOUT_TP_LIST", "GRP_FORCE_TRAIN_GPUS"): + monkeypatch.delenv(variable, raising=False) + config = AsyncRLConfig( + model_path="not-loaded", train_gpus=2, rollout_gpus=2, + train_tp_size=1, train_dp_size=2, micro_batch_size=1, + batch_size=8, n_total_gpus=4, + ) + cfg = config.global_resource_planner + cfg.enabled = True + cfg.train_backend = cfg.rollout_backend = "analytic" + cfg.plan_interval = 1 + cfg.warmup_steps = 0 + cfg.min_history_size = 1 + cfg.reconfiguration_cost_s = reconfiguration_cost + cfg.elastic_hybrid_planning_enabled = True + planner = GlobalResourcePlanner.from_config(config) + train = TrainParallelConfig(tp=1, dp=2, b_micro=1) + current = GlobalResourcePlan( + train_config=train, rollout_config=RolloutClusterConfig([1, 1]), + t_train=100, t_rollout=90, t_global=100, + n_total_gpus=4, max_concurrent_rollouts=16, + ) + planner._build_current_plan = Mock(return_value=current) + # Freeze the optimizer result so these tests exercise the real decision + # gates and EHP policy independently of the search's selected candidate. + planner.optimizer.optimize = Mock(return_value=OptimizationResult( + train_config=train, rollout_config=RolloutClusterConfig([2]), + t_train=80, t_rollout=70, t_global=80, + )) + return config, planner, current + + +@pytest.mark.parametrize("resource_allowed,cost", [(False, 1e9), (True, 0)]) +@pytest.mark.parametrize("elastic_adjust,train_time", [(False, 1), (True, 10)]) +def test_full_repartition_and_ehp_have_independent_gates( + monkeypatch, resource_allowed, cost, elastic_adjust, train_time, +): + config, planner, current = planner_fixture(monkeypatch, cost) + metrics = planner.observe_runtime( + step=1, dispatcher_metrics={}, + step_stats={"train_time": train_time, "rollout_time": 1}, + ) + decision = planner.plan_if_needed( + 1, config, batch=[{"input_len": 128, "output_len": 4096}], + runtime_metrics=metrics, + ) + assert decision.resource_should_reconfigure is resource_allowed + assert decision.elastic_hybrid_signal.should_adjust is elastic_adjust + assert decision.should_reconfigure is (resource_allowed or elastic_adjust) + assert decision.candidate_plan.rollout_tp_list == [2] + assert decision.current_plan.rollout_tp_list == [1, 1] + assert decision.candidate_plan.expected_gain_s == 20 - cost + assert planner.active_plan is (decision.candidate_plan if resource_allowed else current) + wire = decision.to_dict() + assert wire["resource_should_reconfigure"] is resource_allowed + assert wire["elastic_should_adjust"] is elastic_adjust + assert wire["elastic_hybrid_signal"]["desired_workers"] == (2 if elastic_adjust else 0) + assert config.train_dp_size == 2 + assert config.rollout_gpus == 2 + + +def test_release_also_leaves_rejected_full_candidate_inactive(monkeypatch): + config, planner, current = planner_fixture(monkeypatch, 1e9) + metrics = planner.observe_runtime( + step=1, dispatcher_metrics={}, + step_stats={"train_time": 1, "rollout_time": 10, "active_hybrid_replicas": 1}, + ) + decision = planner.plan_if_needed( + 1, config, batch=[{"input_len": 128, "output_len": 4096}], runtime_metrics=metrics, + ) + assert decision.should_reconfigure + assert decision.resource_should_reconfigure is False + assert decision.elastic_hybrid_signal.action == "release" + assert planner.active_plan is current + + +def test_skipped_real_decision_and_legacy_manual_decision_are_distinguishable(monkeypatch): + config, planner, _ = planner_fixture(monkeypatch, 0) + planner.warmup_steps = 2 + decision = planner.plan_if_needed(1, config) + assert decision.resource_should_reconfigure is False + assert not decision.should_reconfigure + assert decision.to_dict()["elastic_should_adjust"] is False + manual = PlannerDecision(step=1, should_reconfigure=True, reason="legacy") + assert manual.resource_should_reconfigure is None + assert manual.to_dict()["resource_should_reconfigure"] is None diff --git a/tests/test_grpo_grouping.py b/tests/test_grpo_grouping.py index f1efe32..b7ec844 100644 --- a/tests/test_grpo_grouping.py +++ b/tests/test_grpo_grouping.py @@ -1,8 +1,13 @@ from itertools import islice from types import SimpleNamespace +import asyncio +import pytest import torch +from RL_Framework.infra.execution.async_runner import AsyncTaskRunner, TaskResult +from RL_Framework.infra.execution.batch_dispatcher import BatchTaskDispatcher, TaskInput +from RL_Framework.infra.sync.staleness import StalenessManager from RL_Framework.trainer.async_rl_trainer import AsyncRLTrainer @@ -92,3 +97,72 @@ def on_batch_consumed(self, count): assert [trajectory["grpo_group_id"] for trajectory in batch] == ["b"] * 4 assert trainer.staleness_manager.consumed == 8 + + +@pytest.mark.parametrize('ready_results', [False, True]) +def test_collector_consumes_only_accepted_packets_with_real_timeout_accounting(monkeypatch, ready_results): + trainer = _trainer(n_samples=4) + trainer.config.max_head_offpolicyness = 3 + trainer.config.enable_rollout_tracing = False + trainer.is_main_process = True + trainer._pending_grpo_groups = {} + trainer._consumed_grpo_group_ids = {'consumed'} + trainer.rollout_engine = object() + trainer._refresh_rollout_engine_from_manifest_for_task = lambda: None + manager = trainer.staleness_manager = StalenessManager( + trainer.train_engine, max_concurrent_rollouts=16, consumer_batch_size=4, max_staleness=3) + + async def episode(engine, data, **kwargs): + if data.get('timeout'): + await asyncio.Future() + return {} + + trainer.workflow = SimpleNamespace(run_episode=episode) + monkeypatch.setenv('ROLLOUT_TASK_TIMEOUT_S', '.001') + # One timeout never entered accepted. Successful duplicate members and a + # late result for a consumed group did, and must still release their credit. + tasks = [('incomplete', 0, False), ('incomplete', 1, False), + ('incomplete', 3, True), ('consumed', 0, False), + ('incomplete', 1, False), ('incomplete', 2, False)] + tasks += [('complete', index, False) for index in range(4)] + + async def produce(): + results = [] + for task_id, (group_id, index, timeout) in enumerate(tasks): + manager.on_rollout_enqueued() + manager.on_rollout_submitted() + task = TaskInput(task_id, {'timeout': timeout}, version=3, + group_id=group_id, rollout_index=index) + results.append(await trainer._create_rollout_task(task)()) + return results + + packets = asyncio.run(produce()) + assert packets[2]['status'] == 'failed' and packets[2]['failure_reason'] == 'task_timeout' + stats = manager.get_stats() + assert (stats.accepted, stats.running, stats.rejected, stats.enqueued) == (9, 0, 1, 0) + if ready_results: + runner = AsyncTaskRunner(max_queue_size=32) + dispatcher = BatchTaskDispatcher(runner, manager, trainer._create_rollout_task) + for task_id, packet in enumerate(packets): + dispatcher._pending_results[task_id] = TaskResult(packet, task_id, task_id, task_id) + dispatcher._active_task_ids.add(task_id) + dispatcher.active_submit_and_wait = lambda **kwargs: pytest.fail('Ready results must not submit work') + trainer.dispatcher = dispatcher + else: + remaining = list(packets) + + def collect(**kwargs): + count = kwargs['batch_size'] + chunk, remaining[:] = remaining[:count], remaining[count:] + return chunk + + trainer.dispatcher = SimpleNamespace(active_submit_and_wait=collect) + batch = trainer._collect_complete_grpo_batch(iter(()), batch_size=4) + assert [member['grpo_group_id'] for member in batch] == ['complete'] * 4 + assert [member['grpo_rollout_index'] for member in batch] == list(range(4)) + assert set(trainer._pending_grpo_groups) == {'incomplete'} + assert len(trainer._pending_grpo_groups['incomplete']) == 3 + assert packets[2]['status'] == 'failed' # Failure notification is untouched. + stats = manager.get_stats() + assert (stats.accepted, stats.running, stats.rejected, stats.enqueued) == (0, 0, 1, 0) + assert manager.get_capacity() == 16 diff --git a/tests/test_hybrid_boundary_activation.py b/tests/test_hybrid_boundary_activation.py new file mode 100644 index 0000000..6a995e6 --- /dev/null +++ b/tests/test_hybrid_boundary_activation.py @@ -0,0 +1,197 @@ +"""Ready joins wait for Core boundaries; installation keeps its ACK deadline.""" +from concurrent.futures import TimeoutError as FutureTimeout +import json +import queue +import threading +from types import SimpleNamespace as NS + +import pytest + +from RL_Framework.config import AsyncRLConfig +from RL_Framework.infra.elastic import hybrid_pool as pool_module +from RL_Framework.infra.elastic import runtime_executor as runtime_module +from RL_Framework.infra.elastic.hybrid_pool import ElasticHybridPool +from RL_Framework.infra.elastic.runtime_executor import RuntimeElasticExecutor + + +@pytest.fixture +def boundary_join(tmp_path, monkeypatch): + clock = NS(now=0.0) + timing = NS(monotonic=lambda: clock.now, time=lambda: clock.now) + monkeypatch.setattr(runtime_module, "time", timing) + monkeypatch.setattr(pool_module, "time", timing) + config = AsyncRLConfig(model_path="/not-loaded", train_gpus=1, rollout_gpus=1, + n_total_gpus=2, train_tp_size=1, train_dp_size=1) + cfg = config.global_resource_planner + cfg.elastic_hybrid_join_timeout_s = 180 + cfg.hybrid_worker_ready_timeout_s = 60 + cfg.hybrid_worker_task_dir = str(tmp_path) + pool = ElasticHybridPool(core_train_workers=["dp0"], + core_rollout_workers=["rollout0"], snapshot_fetcher=lambda *_: 0) + runtime = RuntimeElasticExecutor(config=config, planner=NS(), elastic_pool=pool, + train_engine=NS(optimizer_revision=14)) + checkpoints = queue.Queue() + requests = [] + + class ControlledEvent: + """Expose each real queue poll without relying on thread scheduling.""" + def __init__(self): + self.condition = threading.Condition() + self.ready = False + self.polls = 0 + requests.append(self) + + def set(self): + with self.condition: + self.ready = True + self.condition.notify_all() + + def wait(self, _timeout): + with self.condition: + if not self.ready: + checkpoints.put("waiting") + assert self.condition.wait_for( + lambda: self.ready or self.polls, timeout=3), "test did not release the queue poll" + if self.polls: + self.polls -= 1 + return self.ready + + def poll(self): + with self.condition: + self.polls += 1 + self.condition.notify_all() + + # Runtime locks and pool futures retain real threading primitives. Only + # the request Event created inside _queue_activation is controlled. + monkeypatch.setattr(runtime_module, "threading", NS(Event=ControlledEvent)) + closing = threading.Event() + state = NS(capture_seconds=0, install_seconds=0, checkpoint=None) + manifest = tmp_path / "state14.json" + manifest.write_text(json.dumps({"format": "libra-ehp-lanes-v1", + "snapshot_id": "capture14", "optimizer_revision": 14, + "lanes": {"dp0": {"0": "ehp://source/lane14"}}})) + + def advance(seconds): + clock.now += seconds + requests[0].poll() + state.checkpoint = checkpoints.get(timeout=3) + + def acknowledge(): + receipt = dict(snapshot=str(manifest), rank=0, pid=123, acknowledged=True, + snapshot_id="capture14", optimizer_revision=14, state_version=14, + lane_snapshot_id="lane14", target_core_id="dp0", source_uri="ehp://source/lane14") + (tmp_path / "replica.rank_0.ready").write_text( + "".join(f"{key}={value}\n" for key, value in receipt.items())) + + def capture(): + if state.capture_seconds: + advance(state.capture_seconds) + return str(manifest) + + def install(_seconds): + advance(state.install_seconds) + acknowledge() + return closing.is_set() + + runtime._closing = NS(is_set=closing.is_set, set=closing.set, wait=install) + monkeypatch.setattr(runtime, "_capture_core_snapshot", capture) + monkeypatch.setattr(runtime, "_publish_hybrid_membership", lambda: None) + monkeypatch.setattr(runtime, "_stop_managed_rollout_processes", lambda: None) + + def activate(replica_id, members, target_core_id, version): + return runtime._queue_activation(replica_id, + dict(acknowledged=True, state_version=version, optimizer_revision=version)) + + handle = pool.join_replica("replica", ["rollout0"], "dp0", + replica_activation_barrier=activate, timeout_s=180) + handle.future.add_done_callback(lambda future: checkpoints.put("finished")) + runtime._pending_hybrid_joins["replica"] = handle + runtime._pending_hybrid_join_started_at["replica"] = 0 + assert checkpoints.get(timeout=3) == "waiting" + assert runtime.has_pending_activations() and not handle.done() + + def finish(): + if not state.install_seconds: + acknowledge() + runtime._finish_pending_activations() + return handle.result(timeout=3) + + h = NS(runtime=runtime, pool=pool, clock=clock, state=state, + advance=advance, finish=finish, handle=handle) + try: + yield h + finally: + closing.set() + for event in requests: + event.set() + pool.close() + try: + handle.result(timeout=3) + except FutureTimeout: + if not handle.done(): + raise + except RuntimeError: + pass + + +def test_ready_join_waits_past_180_seconds_until_a_core_boundary(boundary_join): + h = boundary_join + h.advance(181) + workers = h.finish() + assert h.clock.now == 181 + assert all(worker.join_state.value == "active" and worker.state_version == 14 + for worker in workers) + + +def test_capture_and_ack_can_cross_the_old_queue_deadline(boundary_join): + h = boundary_join + h.advance(50) + h.state.capture_seconds = 96 + h.state.install_seconds = 50 + workers = h.finish() + assert h.clock.now == 196 + assert all(worker.join_state.value == "active" and worker.state_version == 14 + for worker in workers) + assert h.pool.snapshot()["rollout0"].state_version == 14 + + +@pytest.mark.parametrize("action", ["cancel", "close"]) +def test_cancel_and_cleanup_still_wake_ready_join(boundary_join, action): + h = boundary_join + if action == "cancel": + h.runtime._release_all_hybrid_workers(reason="test cancellation") + else: + h.runtime.close() + h.advance(1) # the production Event's next 50ms poll observes closing + with pytest.raises(RuntimeError, match="cancelled"): + h.handle.result(timeout=3) + worker = h.pool.snapshot()["rollout0"] + assert worker.join_state.value == "cancelled" + assert worker.role.value != "hybrid_training" + + +@pytest.mark.parametrize("cold,seconds", [(True, 180), (False, 60)]) +def test_worker_installation_ack_still_has_its_existing_deadline( + tmp_path, monkeypatch, cold, seconds): + clock = NS(now=0.0) + monkeypatch.setattr(runtime_module, "time", NS(monotonic=lambda: clock.now)) + config = AsyncRLConfig(model_path="/not-loaded", train_gpus=1, rollout_gpus=1, + n_total_gpus=2, train_tp_size=1, train_dp_size=1) + config.global_resource_planner.hybrid_worker_task_dir = str(tmp_path) + config.global_resource_planner.hybrid_worker_ready_timeout_s = 60 + runtime = RuntimeElasticExecutor(config=config, planner=NS()) + manifest = tmp_path / "manifest.json" + manifest.write_text(json.dumps({"format": "libra-ehp-lanes-v1", + "snapshot_id": "state", "optimizer_revision": 2, + "lanes": {"dp0": {"0": "ehp://source/lane0"}}})) + + def advance(_seconds): + clock.now += 1 + return False + + runtime._closing = NS(is_set=lambda: False, wait=advance) + with pytest.raises(TimeoutError, match=f"after {seconds}s"): + runtime._wait_hybrid_worker_ready("replica", state_version=2, + replica_world_size=1, snapshot_path=str(manifest), target_core_id="dp0", + **({"timeout_s": 180} if cold else {})) + assert clock.now == seconds diff --git a/tests/test_hybrid_join_readiness_budget.py b/tests/test_hybrid_join_readiness_budget.py new file mode 100644 index 0000000..7047cb0 --- /dev/null +++ b/tests/test_hybrid_join_readiness_budget.py @@ -0,0 +1,221 @@ +"""Cold startup shares the join budget; live state catch-up keeps its ACK budget.""" + +import json +import threading +from types import SimpleNamespace + +import pytest + +from RL_Framework.config import AsyncRLConfig +from RL_Framework.infra.elastic import hybrid_pool as pool_module +from RL_Framework.infra.elastic import runtime_executor as runtime_module +from RL_Framework.infra.elastic.hybrid_pool import ElasticHybridPool +from RL_Framework.infra.elastic.runtime_executor import ( + RuntimeElasticExecutor, + RuntimeReconfigurationResult, +) + + +@pytest.fixture +def bootstrap(tmp_path, monkeypatch): + clock = SimpleNamespace(now=0.0) + timing = SimpleNamespace(monotonic=lambda: clock.now, time=lambda: clock.now) + monkeypatch.setattr(runtime_module, "time", timing) + monkeypatch.setattr(pool_module, "time", timing) + config = AsyncRLConfig(model_path="/not-loaded", train_gpus=1, rollout_gpus=1, + n_total_gpus=2, train_tp_size=1, train_dp_size=1) + cfg = config.global_resource_planner + cfg.hybrid_worker_task_dir = str(tmp_path) + cfg.hybrid_worker_launch_enabled = True + cfg.hybrid_worker_remote_control_enabled = True + cfg.elastic_hybrid_join_timeout_s = 180 + cfg.hybrid_worker_ready_timeout_s = 60 + pool = ElasticHybridPool(core_train_workers=["dp0"], + core_rollout_workers=["rollout0"], snapshot_fetcher=lambda *args: 0) + executor = RuntimeElasticExecutor(config=config, planner=SimpleNamespace(), elastic_pool=pool) + manifest = tmp_path / "snapshot.json" + manifest.write_text(json.dumps({"format": "libra-ehp-lanes-v1", + "optimizer_revision": 0, "snapshot_id": "capture0", + "lanes": {"dp0": {"0": "ehp://source/lane0"}}})) + executor._snapshot_manifests[0] = str(manifest) + monkeypatch.setattr(executor, "_capture_core_snapshot", lambda: str(manifest)) + monkeypatch.setattr(executor, "_build_hybrid_worker_command", lambda **kw: "not executed") + monkeypatch.setattr(executor, "_cluster_swap_training_slot", + lambda *args: {"host": "allocated-node", "gpus": [1]}) + monkeypatch.setattr(executor, "_desired_hybrid_workers", lambda *args: 1) + monkeypatch.setattr(executor, "_training_reconfiguration_targets", + lambda *args: (["dp0"], "dp0", ["rollout0"])) + # The loader's initial installation is the subject here, before the real + # activation queue hands control back to a Core training boundary. + monkeypatch.setattr(executor, "_queue_activation", lambda identity, receipt: receipt) + ready_at = 70.0 + identity = SimpleNamespace(worker="ehp_replica_rollout0") + + def advance(_seconds): + clock.now += 1 + if clock.now >= ready_at: + receipt = dict(snapshot=str(manifest), rank=0, pid=123, acknowledged=True, + snapshot_id="capture0", optimizer_revision=0, state_version=0, + lane_snapshot_id="lane0", target_core_id="dp0", source_uri="ehp://source/lane0") + (tmp_path / f"{identity.worker}.rank_0.ready").write_text( + "".join(f"{key}={value}\n" for key, value in receipt.items())) + return False + + executor._closing = SimpleNamespace(is_set=lambda: False, wait=advance) + yield SimpleNamespace(executor=executor, pool=pool, clock=clock, + manifest=manifest, identity=identity) + pool.close() + + +@pytest.mark.parametrize("replica", [True, False], ids=["replica", "singleton"]) +def test_initial_installation_can_be_ready_after_seventy_seconds(bootstrap, replica): + h = bootstrap + if replica: + h.executor._reconfigure_training_pool(1, SimpleNamespace(), + RuntimeReconfigurationResult(False, "test")) + handle = h.executor._pending_hybrid_joins["ehp_replica_rollout0"] + else: + h.identity.worker = "rollout0" + handle = h.pool.join_training("rollout0", "dp0", timeout_s=180, + activation_barrier=h.executor._activate_hybrid_worker) + handle.result(timeout=3) + assert h.clock.now == 70 + assert h.pool.snapshot()["rollout0"].role.value == "hybrid_training" + + +def test_live_state_catchup_still_times_out_after_sixty_seconds(bootstrap): + h = bootstrap + with pytest.raises(TimeoutError, match="after 60s"): + h.executor._wait_hybrid_worker_ready(h.identity.worker, state_version=0, + replica_world_size=1, snapshot_path=str(h.manifest), target_core_id="dp0") + assert h.clock.now == 60 + + +@pytest.mark.parametrize("preparation_seconds", [0, 20]) +def test_initial_installation_does_not_extend_the_overall_join_deadline( + bootstrap, monkeypatch, preparation_seconds): + h = bootstrap + h.executor.config.global_resource_planner.elastic_hybrid_join_timeout_s = 65 + + def prepare(**kwargs): + h.clock.now += preparation_seconds + return "not executed" + + monkeypatch.setattr(h.executor, "_build_hybrid_worker_command", prepare) + h.executor._reconfigure_training_pool(1, SimpleNamespace(), + RuntimeReconfigurationResult(False, "test")) + handle = h.executor._pending_hybrid_joins["ehp_replica_rollout0"] + with pytest.raises(TimeoutError): + handle.result(timeout=3) + # The pool checks its original deadline on every poll; preparation cannot + # restart it. The clock advances in one-second increments in this fixture. + assert h.clock.now <= 66 + + +@pytest.fixture +def queued_activation(bootstrap, monkeypatch): + h = bootstrap + executor = h.executor + monkeypatch.setattr(executor, "_queue_activation", + RuntimeElasticExecutor._queue_activation.__get__(executor)) + monkeypatch.setattr(executor, "_publish_hybrid_membership", lambda: None) + monkeypatch.setattr(executor, "_stop_hybrid_worker_process", lambda *args: None) + monkeypatch.setattr(executor, "_stop_managed_rollout_processes", lambda: None) + from RL_Framework.infra.elastic import core_reconfiguration + monkeypatch.setattr(core_reconfiguration, "dispatch_core_command", lambda *args, **kwargs: []) + closing = threading.Event() + cold_wait = executor._closing.wait + state = SimpleNamespace(catching_up=False, ack_at=0, capture_seconds=0) + latest = h.manifest.with_name("snapshot14.json") + latest.write_text(json.dumps({"format": "libra-ehp-lanes-v1", + "optimizer_revision": 14, "snapshot_id": "capture14", + "lanes": {"dp0": {"0": "ehp://source/lane14"}}})) + + def acknowledge(): + receipt = dict(snapshot=str(latest), rank=0, pid=123, acknowledged=True, + snapshot_id="capture14", optimizer_revision=14, state_version=14, + lane_snapshot_id="lane14", target_core_id="dp0", source_uri="ehp://source/lane14") + (latest.parent / f"{h.identity.worker}.rank_0.ready").write_text( + "".join(f"{key}={value}\n" for key, value in receipt.items())) + + def advance(seconds): + if not state.catching_up: + return cold_wait(seconds) + h.clock.now += 1 + if h.clock.now >= state.ack_at: + acknowledge() + return closing.is_set() + + def capture(): + h.clock.now += state.capture_seconds + return str(latest) + + executor._closing = SimpleNamespace(is_set=closing.is_set, set=closing.set, wait=advance) + executor._reconfigure_training_pool(1, SimpleNamespace(), + RuntimeReconfigurationResult(False, "test")) + handle = executor._pending_hybrid_joins[h.identity.worker] + for _ in range(200): + if executor.has_pending_activations(): + break + threading.Event().wait(.01) + assert executor.has_pending_activations() + assert h.clock.now == 70 + assert not handle.done() + # Submission occurred at t=0, before the fake clock's cold-start polls. + executor._pending_hybrid_join_started_at[h.identity.worker] = 0 + executor.train_engine = SimpleNamespace(optimizer_revision=14) + monkeypatch.setattr(executor, "_capture_core_snapshot", capture) + state.catching_up = True + h.handle, h.catchup, h.acknowledge = handle, state, acknowledge + try: + yield h + finally: + closing.set() + with executor._pending_activation_lock: + for request in executor._pending_activations.values(): + request["error"] = "test cleanup" + request["event"].set() + + +@pytest.mark.parametrize("boundary_at,capture_seconds", [(140, 0), (70, 90)], + ids=["ack-at-190", "cold70-capture90-install50"]) +def test_activation_queue_can_finish_after_cold_deadline( + queued_activation, boundary_at, capture_seconds): + h = queued_activation + h.clock.now = boundary_at + h.catchup.capture_seconds = capture_seconds + h.catchup.ack_at = boundary_at + capture_seconds + 50 + h.executor._finish_pending_activations() + workers = h.handle.result(timeout=3) + assert h.clock.now == h.catchup.ack_at + assert all(worker.join_state.value == "active" and worker.state_version == 14 + for worker in workers) + assert h.pool.snapshot()["rollout0"].state_version == 14 + + +def test_reconcile_does_not_cancel_a_join_waiting_in_activation_queue(queued_activation): + h = queued_activation + h.clock.now = 190 + h.executor._reconcile_pending_hybrid_joins() + assert h.executor._pending_hybrid_joins[h.identity.worker] is h.handle + assert not h.handle.done() + assert h.pool.snapshot()["rollout0"].join_state.value == "activating" + + +@pytest.mark.parametrize("action", ["cancel", "close"]) +def test_queued_activation_late_ack_cannot_activate_cancelled_join(queued_activation, action): + h = queued_activation + if action == "cancel": + h.handle.cancel() + else: + h.executor.close() + h.clock.now = 90 + h.catchup.ack_at = 90 + h.acknowledge() + with pytest.raises(RuntimeError): + h.executor._finish_pending_activations() + with pytest.raises(RuntimeError): + h.handle.result(timeout=3) + worker = h.pool.snapshot()["rollout0"] + assert worker.join_state.value == "cancelled" + assert worker.role.value != "hybrid_training" diff --git a/tests/test_hybrid_optimizer_cuda_cache.py b/tests/test_hybrid_optimizer_cuda_cache.py index 34a3d7c..0ec61ad 100644 --- a/tests/test_hybrid_optimizer_cuda_cache.py +++ b/tests/test_hybrid_optimizer_cuda_cache.py @@ -34,7 +34,7 @@ def cuda_engine(monkeypatch): calls = [] gradients = [DeviceGradient(calls), DeviceGradient(calls)] engine = MegatronCoreTrainEngine.__new__(MegatronCoreTrainEngine) - engine.rank, engine.current_version = 0, 1 + engine.rank, engine.optimizer_revision = 0, 1 engine._model_main_gradients = lambda: gradients engine.optimizer = SimpleNamespace(step=lambda: (calls.append('step') or (True, 0., 0))) monkeypatch.setattr(torch.cuda, 'is_available', lambda: True) @@ -56,13 +56,13 @@ def test_first_cuda_hybrid_update_releases_after_copy_and_before_step(monkeypatc sources = (Source([1.25, 2.5]), Source([3., 4.])) engine.apply_elastic_gradient_update(sources, state_version=2) assert calls == ['copy', 'copy', 'enter', 'empty', 'exit', 'step'] - assert engine.current_version == 2 and engine._hybrid_optimizer_cuda_cache_released + assert engine.optimizer_revision == 2 and engine._hybrid_optimizer_cuda_cache_released for gradient, source in zip(gradients, sources): assert torch.equal(gradient.tensor, source.tensor.bfloat16()) assert not capsys.readouterr().out calls.clear() engine.apply_elastic_gradient_update(sources, state_version=3) - assert calls == ['copy', 'copy', 'step'] and engine.current_version == 3 + assert calls == ['copy', 'copy', 'step'] and engine.optimizer_revision == 3 assert not capsys.readouterr().out @@ -77,11 +77,11 @@ def failed_step(): engine.optimizer.step = failed_step with pytest.raises(RuntimeError) as error: engine.apply_elastic_gradient_update((Source([1., 2.]), Source([3., 4.])), state_version=2) - assert error.value is failure and engine.current_version == 1 + assert error.value is failure and engine.optimizer_revision == 1 assert not getattr(engine, '_hybrid_optimizer_cuda_cache_released', False) engine.optimizer.step = lambda: (True, 0., 0) engine.apply_elastic_gradient_update((Source([1., 2.]), Source([3., 4.])), state_version=2) - assert calls.count('empty') == 2 and engine.current_version == 2 + assert calls.count('empty') == 2 and engine.optimizer_revision == 2 def test_rejected_optimizer_update_does_not_mark_cache_or_revision_complete(monkeypatch): @@ -89,7 +89,7 @@ def test_rejected_optimizer_update_does_not_mark_cache_or_revision_complete(monk engine.optimizer.step = lambda: (False, 0., 0) with pytest.raises(FloatingPointError, match='rejected the update'): engine.apply_elastic_gradient_update((Source([1., 2.]), Source([3., 4.])), state_version=2) - assert calls.count('empty') == 1 and engine.current_version == 1 + assert calls.count('empty') == 1 and engine.optimizer_revision == 1 assert not getattr(engine, '_hybrid_optimizer_cuda_cache_released', False) @@ -100,7 +100,7 @@ def test_cpu_gradients_keep_original_copy_dtype_and_optimizer_update(monkeypatch source = torch.tensor([1.25, -2.5], dtype=torch.float32) source_before = source.clone() weight = torch.tensor([10., 20.]) - engine.current_version = 1 + engine.optimizer_revision = 1 engine._model_main_gradients = lambda: [gradient] monkeypatch.setattr(torch.cuda, 'is_available', lambda: cuda_available) monkeypatch.setattr(torch.cuda, 'empty_cache', lambda: pytest.fail('CPU path released CUDA cache')) @@ -114,7 +114,7 @@ def step(): assert torch.equal(source, source_before) assert torch.equal(gradient, source.bfloat16()) assert torch.equal(weight, torch.tensor([10., 20.]) - source.bfloat16().float() * .1) - assert engine.current_version == 2 + assert engine.optimizer_revision == 2 assert not getattr(engine, '_hybrid_optimizer_cuda_cache_released', False) @@ -122,4 +122,4 @@ def test_tensor_count_failure_happens_before_copy_or_cache_release(monkeypatch): engine, _gradients, calls = cuda_engine(monkeypatch) with pytest.raises(ValueError, match='tensor count mismatch'): engine.apply_elastic_gradient_update((Source([1., 2.]),), state_version=2) - assert calls == [] and engine.current_version == 1 + assert calls == [] and engine.optimizer_revision == 1 diff --git a/tests/test_hybrid_sample_statistics.py b/tests/test_hybrid_sample_statistics.py new file mode 100644 index 0000000..88c2ea4 --- /dev/null +++ b/tests/test_hybrid_sample_statistics.py @@ -0,0 +1,44 @@ +"""Core reports the full batch only after the activation protocol accepts gradients.""" +from dataclasses import replace + +import pytest + +from hybrid_statistics_helpers import compute, engine, rows + + +@pytest.mark.parametrize('core_count,zero_mask',[(8,False),(8,True),(0,False)]) +def test_real_payload_and_grpo_report_all_samples_without_counting_extra_core_updates( + monkeypatch,core_count,zero_mask): + core_rows=rows(0,core_count) + hybrid_count=4 if core_count else 12 + hybrid_rows=rows(8,hybrid_count,zero_mask=zero_mask) + _,payload=compute(monkeypatch,hybrid_rows) + core,steps=engine(monkeypatch) + core.set_elastic_batch_partition(core_count,12,{'hybrid':hybrid_count}) + core.enqueue_hybrid_gradient_payload(payload) + reference=core._merge_stats([core._loss(batch,core._forward(batch))[1] + for batch in core._iter_micro_batches(core_rows+hybrid_rows)]) + result=core.grpo_update(core_rows) + assert result['n_samples']==12 + assert result['reward_sum']==reference['reward_sum'] + for key in ('loss','pg_loss','kl','reward_mean'): + assert result[key]==pytest.approx(reference[key],abs=1e-6) + assert result['n_updates']==(core_count+2)//3 + if core_count: + assert result['grad_norm']==7. + assert len(steps)==1 and core.optimizer_revision==3 + + +def test_only_this_reduction_matching_lane_and_target_statistics_are_used(monkeypatch): + _,payload=compute(monkeypatch,rows(8,4),lane=1,width=2) + core,_=engine(monkeypatch,lane=1,width=2) + core.set_elastic_batch_partition(8,12,{'hybrid':4}) + core.enqueue_hybrid_gradient_payload(payload) + assert hasattr(payload,'stats'), 'Hybrid compute discarded the loss statistics' + ignored=replace(payload,replica_rank=0,stats=dict(payload.stats,n_samples=999.)) + stale=replace(payload,replica_id='old',step=2,stats=dict(payload.stats,n_samples=888.)) + prior=(stale.replica_id,stale.replica_rank,stale.step,stale.state_version,stale.membership_epoch) + core.elastic_gradient_domain._consumed_payloads.add(prior) + core._pending_hybrid_gradients.extend([ignored,stale]) + result=core.grpo_update(rows(0,8)) + assert result['n_samples']==12 and result['n_updates']==3 diff --git a/tests/test_hybrid_task_lifecycle.py b/tests/test_hybrid_task_lifecycle.py new file mode 100644 index 0000000..f73d533 --- /dev/null +++ b/tests/test_hybrid_task_lifecycle.py @@ -0,0 +1,124 @@ +"""Hybrid completion receipts fence normal return; exit fences cancellation.""" +import threading +from types import SimpleNamespace + +import pytest + +from RL_Framework.infra.elastic.runtime_executor import ManagedHybridWorkerProcess, RuntimeElasticExecutor + + +def _executor(tmp_path): + executor = RuntimeElasticExecutor.__new__(RuntimeElasticExecutor) + executor.config = SimpleNamespace(global_resource_planner=SimpleNamespace( + hybrid_worker_task_dir=str(tmp_path), hybrid_worker_remote_control_dir=str(tmp_path), + hybrid_worker_remote_control_enabled=True, hybrid_update_timeout_s=.3, + hybrid_worker_stop_timeout_s=.1)) + executor._lifecycle_lock = threading.RLock() + executor._hybrid_worker_processes = {} + canonical = ManagedHybridWorkerProcess('h', 'dp0', 'worker', 123, 'snapshot', host='node') + member = ManagedHybridWorkerProcess('member', 'dp0', 'worker', -1, 'snapshot', host='node', replica_id='h') + executor._hybrid_worker_meta = {'h': canonical, 'member': member} + return executor + + +def test_normal_return_waits_for_actual_task_completion_before_detach(tmp_path): + executor = _executor(tmp_path) + executor._snapshot_cleanup_requested = False + executor._pending_activation_lock = threading.RLock() + executor._pending_activations = {} + executor._pending_hybrid_join_lock = threading.RLock() + executor._pending_hybrid_joins = {} + executor._pending_hybrid_join_started_at = {} + executor._retired_hybrid_joins = [] + executor._publish_hybrid_membership = lambda: None + worker = SimpleNamespace(worker_id='member', role='hybrid_training') + detached = [] + executor.elastic_pool = SimpleNamespace(replica_snapshot=lambda: {'h': (worker,)}, + release_replica_to_rollout=lambda replica: detached.append(replica)) + executor._stop_snapshot_loader = lambda *_: None + task = tmp_path / 'h.step_14.pt' + task.write_text('already consumed by Core, still applying on Hybrid') + finished = threading.Event() + thread = threading.Thread(target=lambda: (executor._release_all_hybrid_workers(reason='planner'), finished.set())) + thread.start() + try: + assert not finished.wait(.08) + assert not detached + task.rename(str(task) + '.done') + assert finished.wait(2) + assert detached == ['h'] + finally: + task.rename(str(task) + '.done') if task.exists() else None + thread.join(2) + + +def test_failed_or_missing_completion_is_not_success(tmp_path): + executor = _executor(tmp_path) + executor.config.global_resource_planner.hybrid_update_timeout_s = .05 + task = tmp_path / 'h.step_14.pt'; task.touch() + with pytest.raises(TimeoutError, match='completion'): + executor.wait_for_hybrid_training_tasks('member') + assert task.exists() and not (tmp_path / 'h.step_14.pt.done').exists() + + +def test_worker_failure_while_waiting_propagates_without_completion(tmp_path): + executor = _executor(tmp_path) + task = tmp_path / 'h.step_14.pt'; task.touch() + worker = threading.Timer(.02, lambda: task.rename(str(task) + '.failed')) + worker.start() + try: + with pytest.raises(RuntimeError, match='failed before completion'): + executor.wait_for_hybrid_training_tasks('member') + finally: + worker.join(2) + assert not (tmp_path / 'h.step_14.pt.done').exists() + + +def test_confirmed_exit_retires_only_old_owner_tasks_without_faking_done(tmp_path, monkeypatch): + from RL_Framework.infra.elastic import rollout_lending + executor = _executor(tmp_path) + task = tmp_path / 'h.step_14.pt'; task.write_text('membership_epoch=7;state_version=14') + other = tmp_path / 'other.step_14.pt'; other.touch() + (tmp_path / 'h.step_13.pt.done').touch() + observed = [] + def acknowledge(executor, worker, request_id=None): + observed.append(worker) + assert task.exists() and 'h' in executor._hybrid_worker_meta + return dict(worker_id=worker, request_id=request_id, exited=True) + monkeypatch.setattr(rollout_lending, 'wait_hybrid_stopped', acknowledge) + executor._stop_hybrid_worker_process('member') + assert observed == ['h'] + assert not task.exists() and not (tmp_path / 'h.step_14.pt.done').exists() + assert (tmp_path / 'h.step_14.pt.cancelled').read_text() == 'membership_epoch=7;state_version=14' + assert other.exists() and (tmp_path / 'h.step_13.pt.done').exists() + assert not executor._hybrid_worker_meta + + +def test_failed_stop_retains_pending_task_for_owned_retry(tmp_path, monkeypatch): + from RL_Framework.infra.elastic import rollout_lending + executor = _executor(tmp_path) + task = tmp_path / 'h.step_14.pt'; task.touch() + def fail(*args, **kwargs): + raise TimeoutError('old process has not exited') + monkeypatch.setattr(rollout_lending, 'wait_hybrid_stopped', fail) + with pytest.raises(TimeoutError): + executor._stop_hybrid_worker_process('member') + assert task.exists() and executor._hybrid_worker_meta + assert not (tmp_path / 'h.step_14.pt.cancelled').exists() + + +def test_local_exit_also_retires_unfinished_task(tmp_path, monkeypatch): + from RL_Framework.infra.elastic import process_lifecycle + executor = _executor(tmp_path) + proc = object() + executor._hybrid_worker_processes = {'h': proc, 'member': proc} + task = tmp_path / 'h.step_14.pt'; task.touch() + exited = [] + def stop(owned): + assert owned is proc and task.exists() + exited.append(True) + monkeypatch.setattr(process_lifecycle, 'stop_local_worker', stop) + executor._stop_hybrid_worker_process('member') + assert exited and not task.exists() + assert (tmp_path / 'h.step_14.pt.cancelled').exists() + assert not executor._hybrid_worker_processes and not executor._hybrid_worker_meta diff --git a/tests/test_publication_preserves_hybrids.py b/tests/test_publication_preserves_hybrids.py new file mode 100644 index 0000000..9076226 --- /dev/null +++ b/tests/test_publication_preserves_hybrids.py @@ -0,0 +1,179 @@ +"""Routine rollout publication preserves the current Hybrid GPU loan.""" +import io +import json +from types import SimpleNamespace as NS + +import pytest +import torch + +from RL_Framework.infra.sync import nccl_weight_sync, publication_session +from RL_Framework.trainer.async_rl_trainer import AsyncRLTrainer + + +@pytest.fixture +def publication(tmp_path, monkeypatch): + instances = [NS(instance_id=name, host=name, port=8100 + index, tp=index + 1) + for index, name in enumerate(('a', 'b', 'c'))] + endpoints = {inst.instance_id: f'http://{inst.host}:{inst.port}' for inst in instances} + calls, requests, status_urls, exported = [], [], [], [] + members = ('hybrid-b',) + client = NS(scheduler=NS(_instances=[NS(instance_id=inst.instance_id, is_ready=True) + for inst in instances])) + config = NS(weight_sync_mode='nccl', heterogeneous_rollout=NS(instances=instances), + rollout_weight_sync_control_dir=str(tmp_path), rollout_nccl_host='core', + rollout_nccl_port=29620, rollout_nccl_chunk_mb=1, rollout_nccl_rate_limit_gbps=0, + rollout_weight_sync_timeout_s=2, rollout_weight_sync_poll_interval_s=.001, + enable_rollout_tracing=False) + dispatcher = NS(paused=False) + dispatcher.pause = lambda: setattr(dispatcher, 'paused', True) + dispatcher.resume = lambda: setattr(dispatcher, 'paused', False) + dispatcher.cancel_queued = lambda: dict(dispatcher_queue=0, runner_queue=0) + dispatcher.get_runtime_metrics = lambda: dict(staleness_running=0) + dispatcher.wait_until_idle = lambda **kwargs: None + + def weights(): + exported.append(True) + yield 'weight', torch.ones(2) + + trainer = NS(config=config, is_main_process=True, control=None, dispatcher=dispatcher, + rollout_engine=client, _use_heterogeneous=True, + _rollout_unavailable_ids={'b'}, + runtime_elastic_executor=NS(_lent_rollout={'engines': [instances[1]]}), + train_engine=NS(train_backend='megatron_core', optimizer_revision=7, + _elastic_step_hybrid_workers=members, stream_rollout_weights=weights, + set_version=lambda version: calls.append(('version', version))), + _restore_lent_rollouts_for_publication=lambda: calls.append(('restore',)), + _write_json_atomic=lambda path, value: path.write_text(json.dumps(value)), + _record_rollout_weight_version=lambda **kwargs: None, + _record_weight_sync_phase=lambda phase: None, + _trace_train_phase=lambda *args: None, + _wait_for_all_ranks_before_weight_sync=lambda step: None) + + def rebind(**kwargs): + calls.append(('rebind',)) + trainer.rollout_engine = NS(rebound=True) + + trainer._rebind_rollout_engine_from_config = rebind + for name in ('_nccl_rollout_topology', '_sync_weights_nccl', '_sync_weights', + '_reopen_restored_rollouts', '_resume_rollout_pipeline'): + setattr(trainer, name, getattr(AsyncRLTrainer, name).__get__(trainer)) + h = NS(trainer=trainer, instances=instances, endpoints=endpoints, client=client, + calls=calls, requests=requests, status_urls=status_urls, exported=exported, + expected=['a', 'c'], directory=tmp_path, members=members) + + def status(url, timeout): + status_urls.append(url) + identity = next(identity for identity, endpoint in endpoints.items() + if url == endpoint + '/publication/status') + return io.StringIO(json.dumps(dict(server_instance_id='service-' + identity))) + + def send(weights, spec, **kwargs): + tensors = list(weights) + request = json.loads((tmp_path / 'reload_request.json').read_text()) + requests.append(request) + expected_offsets, offset = {}, 0 + for inst in instances: + if inst.instance_id in h.expected: + expected_offsets[inst.instance_id] = offset + offset += inst.tp + assert request['nccl_rank_offsets'] == expected_offsets + assert spec.world_size == request['nccl_world_size'] == 1 + offset + assert spec.session_id == request['nccl_session_id'] + for name, tensor in tensors: + kwargs['on_tensor'](name, tensor.numel() * tensor.element_size()) + for inst in instances: + identity = inst.instance_id + if identity not in expected_offsets: + continue + workers = [dict(rank=rank, nccl_rank=expected_offsets[identity] + rank + 1, + received_tensors=len(tensors), receive_complete=True, loaded_parameters=1, + service_generation=1) for rank in range(inst.tp)] + record = dict(version=request['version'], snapshot_id=request['snapshot_id'], + optimizer_revision=request['optimizer_revision'], server_response=dict( + service_instance_id='service-' + identity, service_generation=1, + prefix_cache_cleared=True, workers=workers)) + (tmp_path / f'ack_{identity}_{request["version"]}.json').write_text(json.dumps(record)) + + monkeypatch.setattr(publication_session.urllib.request, 'urlopen', status) + monkeypatch.setattr(nccl_weight_sync, 'send_weights', send) + monkeypatch.setattr(torch.cuda, 'current_device', lambda: 'cpu') + return h + + +def test_periodic_publication_keeps_hybrid_members_and_existing_rollout_client(publication): + h = publication + # Exercise the real periodic + ordinary publication lifecycle separately + # from the sender's receiver selection, covered below. + h.trainer._sync_weights_nccl = lambda version: h.calls.append(('publish', version)) + AsyncRLTrainer._periodic_weight_sync(h.trainer, 7) + assert ('publish', 8) in h.calls + assert ('restore',) not in h.calls + assert ('rebind',) not in h.calls + assert h.trainer.rollout_engine is h.client + assert h.trainer.train_engine._elastic_step_hybrid_workers == h.members + assert h.trainer._rollout_unavailable_ids == {'b'} + assert not h.trainer.dispatcher.paused + + +@pytest.mark.parametrize('borrowed,restored', [('a', False), ('b', False), ('b', True)], + ids=['first-lent', 'middle-lent', 'middle-explicitly-returned']) +def test_native_publication_targets_only_live_or_explicitly_returned_receivers( + publication, borrowed, restored): + h = publication + h.trainer._rollout_unavailable_ids = {borrowed} + if restored: + h.trainer._restored_rollouts_pending_publication = {borrowed} + h.expected = [inst.instance_id for inst in h.instances + if inst.instance_id != borrowed or restored] + h.trainer._sync_weights(8) + assert len(h.requests) == 1 + assert h.status_urls == [h.endpoints[identity] + '/publication/status' for identity in h.expected] + assert [(row[0], row[1], row[2]) for row in h.trainer._nccl_session_members] == [ + (inst.instance_id, h.endpoints[inst.instance_id], inst.tp) + for inst in h.instances if inst.instance_id in h.expected] + assert ('restore',) not in h.calls + assert h.trainer.rollout_engine is h.client + assert h.trainer.train_engine._elastic_step_hybrid_workers == h.members + + +@pytest.mark.parametrize('borrowed', ['a', 'b'], ids=['first-lent', 'middle-lent']) +def test_publication_session_resolves_filtered_receivers_by_instance_identity(publication, borrowed): + h = publication + offsets, offset = {}, 0 + for inst in h.instances: + if inst.instance_id != borrowed: + offsets[inst.instance_id] = offset + offset += inst.tp + publication_session.prepare_nccl_session(h.trainer, offsets) + assert h.status_urls == [h.endpoints[identity] + '/publication/status' for identity in offsets] + assert [(row[0], row[1], row[2], row[3]) for row in h.trainer._nccl_session_members] == [ + (inst.instance_id, h.endpoints[inst.instance_id], inst.tp, offsets[inst.instance_id]) + for inst in h.instances if inst.instance_id in offsets] + + +def test_publication_with_every_receiver_lent_keeps_admission_paused_without_sender_only_group(publication): + h = publication + h.trainer._rollout_unavailable_ids = set(h.endpoints) + h.expected = [] + AsyncRLTrainer._periodic_weight_sync(h.trainer, 7) + assert not h.requests and not h.status_urls and not h.exported + assert not (h.directory / 'reload_request.json').exists() + assert h.trainer.dispatcher.paused + assert h.trainer.rollout_engine is h.client + assert h.trainer.train_engine._elastic_step_hybrid_workers == h.members + assert ('restore',) not in h.calls + + +def test_extension_rollback_publishes_only_explicit_receivers(): + trainer = AsyncRLTrainer.__new__(AsyncRLTrainer) + trainer.config = NS(heterogeneous_rollout=NS(instances=[ + NS(instance_id='old-loan', tp=1), NS(instance_id='returned', tp=2), + NS(instance_id='still-serving', tp=1)])) + trainer.rollout_engine = None + trainer._use_heterogeneous = True + trainer._rollout_unavailable_ids = {'old-loan', 'returned'} + trainer._restored_rollouts_pending_publication = {'returned'} + trainer._publication_target_ids = {'returned'} + assert trainer._nccl_rollout_topology() == (3, {'returned': 0}) + assert trainer._current_rollout_instance_ids() == ['returned'] + assert trainer._rollout_unavailable_ids == {'old-loan', 'returned'} diff --git a/tests/test_rollout_lending.py b/tests/test_rollout_lending.py new file mode 100644 index 0000000..1a66e37 --- /dev/null +++ b/tests/test_rollout_lending.py @@ -0,0 +1,221 @@ +"""Loans grow without disturbing existing Hybrid ownership or exposing unsafe GPUs.""" +from dataclasses import dataclass +from types import SimpleNamespace as NS + +import pytest + +from RL_Framework.infra.elastic.rollout_lending import _choose_engines, lend_rollout, return_rollout +from RL_Framework.infra.elastic.rollout_lifecycle import unavailable_rollouts + + +@dataclass +class Meta: + instance_id: str + command: str + pid: int + gpus: list + port: int + host: str = 'node' + tp: int = 1 + adopted: bool = False + + +def fixture(count=4, reserve=1): + original = [Meta(chr(97+i), 'serve', 10+i, [i], 8000+i) for i in range(count)] + handles = [NS(instance_id=m.instance_id, is_ready=True) for m in original] + events = [] + dispatcher = NS(paused=False) + dispatcher.is_paused = lambda: dispatcher.paused + dispatcher.pause = lambda: setattr(dispatcher, 'paused', True) + dispatcher.resume = lambda: setattr(dispatcher, 'paused', False) + dispatcher.wait_until_idle = lambda **kw: events.append('drained') + cfg = NS(train_tp_size=1, train_ep_size=1, train_pp_size=1, train_cp_size=1, + global_resource_planner=NS(elastic_hybrid_min_rollout_gpus=reserve, + runtime_drain_timeout_s=1)) + ex = NS(config=cfg, dispatcher=dispatcher, _process_meta={m.instance_id:m for m in original}, + _hybrid_worker_meta={}, rollout_engine=NS(scheduler=NS(_instances=handles), + wait_until_idle=lambda **kw: None)) + ex._adopt_existing_rollout_processes = lambda: None + ex._elastic_replica_size_gpus = lambda: cfg.train_tp_size*cfg.train_ep_size*cfg.train_pp_size*cfg.train_cp_size + def stop(identity): + assert not next(h for h in handles if h.instance_id == identity).is_ready + events.append(('stop',identity)) + return [ex._process_meta.pop(identity)] + ex._stop_rollout_process = stop + ex._command_for_instance_id = lambda identity,port: 'restart '+identity + def start(meta): + assert meta.instance_id not in ex._process_meta + events.append(('start',meta.instance_id)) + ex._process_meta[meta.instance_id] = meta + return meta + ex._start_one_rollout_process = start + ex._wait_rollout_process_start_ready = lambda meta: events.append(('ready',meta.instance_id)) + def publish(instance_ids=None): + assert instance_ids is not None + assert all(not h.is_ready for h in handles if h.instance_id in instance_ids) + events.append(('publish',tuple(instance_ids))) + ex._publish_current_rollout_weights = publish + return ex,events,handles + + +def test_grow_total_loan_preserves_existing_slots_generation_and_hybrid(): + ex,events,handles = fixture() + first = lend_rollout(ex,1) + saved = ex._lent_rollout + generation = saved['generation'] + hybrid = NS(host='node',gpus=[0]) + ex._hybrid_worker_meta['old-hybrid'] = hybrid + second = lend_rollout(ex,2) + assert ex._lent_rollout is saved and saved['generation'] == generation + assert first == {'rollout0':{'host':'node','gpus':[0]}} + assert second == {**first,'rollout1':{'host':'node','gpus':[1]}} + assert ex._hybrid_worker_meta == {'old-hybrid':hybrid} + assert [e for e in events if isinstance(e,tuple) and e[0]=='stop'] == [('stop','a'),('stop','b')] + assert not handles[0].is_ready and not handles[1].is_ready + assert lend_rollout(ex,2) == second + + +def test_reserve_counts_only_rollout_gpus_not_existing_loan(): + ex,events,_ = fixture(reserve=2) + first = lend_rollout(ex,1) + second = lend_rollout(ex,2) + assert len(second) == 2 + before = list(events) + assert lend_rollout(ex,3) == second + assert events == before and len(ex._process_meta) == 2 + assert first.items() <= second.items() + + +def test_selection_uses_full_runtime_replica_width(): + ex,_,_ = fixture(count=3,reserve=0) + ex.config.train_ep_size = 2 + for i,meta in enumerate(ex._process_meta.values()): + meta.host = 'node'+str(i) + assert _choose_engines(ex,2) == () + + +def test_failed_expansion_restores_only_new_engines_after_targeted_ack(): + ex,events,handles = fixture(reserve=1) + first = lend_rollout(ex,1) + old = ex._lent_rollout + snapshot = (dict(old['slots']),list(old['engines']),old['generation']) + stop = ex._stop_rollout_process + once = {'fail':True} + def fail_second(identity): + if identity=='c' and once['fail']: + once['fail']=False + raise RuntimeError('unconfirmed stop') + return stop(identity) + ex._stop_rollout_process = fail_second + with pytest.raises(RuntimeError,match='unconfirmed stop'): + lend_rollout(ex,3) + assert ex._lent_rollout is old + assert (old['slots'],old['engines'],old['generation']) == snapshot + assert lend_rollout(ex,1)==first + assert not handles[0].is_ready and all(h.is_ready for h in handles[1:]) + assert ('publish',('b','c')) in events + assert not any(e in [('start','a'),('publish',('a',))] for e in events) + assert getattr(ex,'_rollout_loan_rollback',None) is None + + +def test_failed_rollback_keeps_new_ownership_isolated_and_old_loan_usable(): + ex,events,handles = fixture() + first = lend_rollout(ex,1) + old = ex._lent_rollout + def fail_publish(instance_ids=None): + events.append(('publish-failed',tuple(instance_ids))) + raise RuntimeError('current weights not acknowledged') + ex._publish_current_rollout_weights=fail_publish + stop=ex._stop_rollout_process + once={'fail':True} + def fail_stop(identity): + if identity=='c' and once['fail']: + once['fail']=False + raise RuntimeError('stop lost acknowledgement') + return stop(identity) + ex._stop_rollout_process=fail_stop + with pytest.raises(RuntimeError): + lend_rollout(ex,3) + assert ex._lent_rollout is old and old['slots']==first + assert ex._rollout_loan_rollback + assert set(unavailable_rollouts(ex))=={'a','b','c'} + assert not any(h.is_ready for h in handles[:3]) + assert lend_rollout(ex,1)==first + before=list(events) + with pytest.raises(RuntimeError,match='incomplete'): + lend_rollout(ex,3) + assert events==before + + +def test_rollback_reopens_selected_services_that_were_never_stopped(): + ex,events,handles=fixture(count=6,reserve=1) + lend_rollout(ex,1) + stop=ex._stop_rollout_process + once={'fail':True} + def fail(identity): + if identity=='c' and once['fail']: + once['fail']=False + raise RuntimeError('middle stop failed') + return stop(identity) + ex._stop_rollout_process=fail + with pytest.raises(RuntimeError,match='middle stop failed'): + lend_rollout(ex,4) + assert not handles[0].is_ready + assert all(h.is_ready for h in handles[1:]) + assert ('stop','d') not in events and ('start','d') not in events + assert ('publish',('b','c')) in events + + +def test_return_readiness_retry_keeps_started_process_ownership(tmp_path): + ex,events,_=fixture(count=2,reserve=0) + lend_rollout(ex,1) + ex._hybrid_worker_control_dir=lambda:tmp_path + once={'fail':True} + def wait(meta): + if once['fail']: + once['fail']=False + raise RuntimeError('service still warming') + ex._wait_rollout_process_start_ready=wait + with pytest.raises(RuntimeError,match='still warming'): + return_rollout(ex) + owned=ex._lent_rollout['restored'][0] + assert ex._process_meta['a'] is owned + assert return_rollout(ex)==[owned] + assert events.count(('start','a'))==1 + + +def test_wide_replica_hands_off_only_after_every_engine_exit_ack(): + ex,events,_=fixture(count=4,reserve=0) + ex.config.train_tp_size=2 + ex.config.train_pp_size=2 + stop=ex._stop_rollout_process + def checked(identity): + assert not getattr(ex,'_lent_rollout',None) + return stop(identity) + ex._stop_rollout_process=checked + slots=lend_rollout(ex,4) + assert len(slots)==4 and not ex._process_meta + assert len([e for e in events if isinstance(e,tuple) and e[0]=='stop'])==4 + + +def test_one_wide_engine_cannot_borrow_through_rollout_reserve(): + ex,events,_=fixture(count=1,reserve=2) + ex.config.train_tp_size=2 + ex._process_meta['a'].gpus=[0,1,2,3] + ex._process_meta['a'].tp=4 + assert lend_rollout(ex,2)=={} + assert events==[] + + +def test_remote_exit_ack_is_bound_to_the_owned_request(tmp_path): + import json + from RL_Framework.infra.elastic.rollout_lending import wait_hybrid_stopped + ex,_,_=fixture() + ex._hybrid_worker_control_dir=lambda:tmp_path + ex.config.global_resource_planner.hybrid_worker_stop_timeout_s=.01 + path=tmp_path/'.stop_worker.done' + path.write_text(json.dumps(dict(worker_id='worker',request_id='old',exited=True))) + with pytest.raises(TimeoutError): + wait_hybrid_stopped(ex,'worker','current') + path.write_text(json.dumps(dict(worker_id='worker',request_id='current',exited=True))) + assert wait_hybrid_stopped(ex,'worker','current')['exited'] diff --git a/tests/test_rollout_lifecycle.py b/tests/test_rollout_lifecycle.py new file mode 100644 index 0000000..9af6d80 --- /dev/null +++ b/tests/test_rollout_lifecycle.py @@ -0,0 +1,333 @@ +"""Lent GPUs return before supply waits and stay closed until publication.""" + +import copy +import threading +from concurrent.futures import ThreadPoolExecutor +from types import SimpleNamespace + +import pytest + +from RL_Framework.infra.elastic import rollout_lifecycle as lifecycle +from RL_Framework.infra.elastic import rollout_lending + + +class Collectives: + def __init__(self): + self.barrier = threading.Barrier(2) + self.packet = None + self.values = [None, None] + self.calls = [] + + @staticmethod + def is_available(): + return True + + @staticmethod + def is_initialized(): + return True + + @staticmethod + def get_world_size(group=None): + return 2 + + @staticmethod + def get_rank(group=None): + return group.rank if group is not None else 0 + + def broadcast_object_list(self, packet, src, group): + assert src == 0 + self.calls.append((group.rank, "broadcast")) + if group.rank == 0: + self.packet = copy.deepcopy(packet[0]) + self.barrier.wait(timeout=3) + packet[0] = copy.deepcopy(self.packet) + self.barrier.wait(timeout=3) + + def all_gather_object(self, records, value, group): + self.calls.append((group.rank, "gather")) + self.values[group.rank] = value + self.barrier.wait(timeout=3) + records[:] = self.values + self.barrier.wait(timeout=3) + + +def trainer(rank, *, ready_groups=2, loan=True, publish_failure=False): + instances = [SimpleNamespace(instance_id=name) for name in ("engine-a", "engine-b")] + calls = [] + handles = [SimpleNamespace(instance_id=instance.instance_id, is_ready=not loan) + for instance in instances] + executor = SimpleNamespace( + _lent_rollout={"engines": instances} if loan else None, + _release_all_hybrid_workers=lambda **kwargs: calls.append("release-hybrids"), + _rollout_lending_resume_pending=True, + ) + state = SimpleNamespace( + is_main_process=rank == 0, control=SimpleNamespace(rank=rank), + runtime_elastic_executor=executor if rank == 0 else None, + config=SimpleNamespace( + n_samples=4, heterogeneous_rollout=SimpleNamespace(instances=instances), + global_resource_planner=SimpleNamespace( + enabled=True, runtime_reconfigure_training=True, + hybrid_worker_launch_enabled=True, runtime_training_pool_plan_only=False, + ), + ), + rollout_engine=SimpleNamespace(scheduler=SimpleNamespace(_instances=handles)), + dispatcher=SimpleNamespace(peek_ready_data=lambda: []), + _pending_grpo_groups={f"group-{i}": [ + {"grpo_group_id": f"group-{i}", "grpo_rollout_index": index} + for index in range(4)] for i in range(ready_groups)}, + _consumed_grpo_group_ids=set(), _source_batch_sample_count=lambda: 8, + _is_rollout_source=lambda: rank == 0, + _pause_rollout_pipeline=lambda: calls.append("pause"), + _resume_rollout_pipeline=lambda: calls.append("resume"), + _refresh_nonblocking_elastic_membership=lambda: calls.append("refresh-membership"), + _rebind_rollout_engine_from_config=lambda **kwargs: calls.append("rebind"), + train_engine=SimpleNamespace( + optimizer_revision=23, _elastic_step_hybrid_workers=("hybrid-old",), + set_version=lambda version: calls.append(("set-version", version)), + ), + calls=calls, handles=handles, + ) + + def publish(*, target_version): + assert target_version == 23 + assert all(not handle.is_ready for handle in handles) + assert state._restored_rollouts_pending_publication == {"engine-a", "engine-b"} + calls.append("publish") + if publish_failure: + raise RuntimeError("weight install failed") + + state._sync_weights = publish + return state + + +def run_ranks(function, trainers): + with ThreadPoolExecutor(max_workers=2) as pool: + futures = [pool.submit(function, state) for state in trainers] + results = [] + for future in futures: + try: + results.append(future.result(timeout=5)) + except Exception as exc: + results.append(exc) + return results + + +def install_transport(monkeypatch): + collectives = Collectives() + monkeypatch.setattr(lifecycle, "dist", collectives) + + def return_loan(executor): + executor._lent_rollout = None + return ["new-service-a", "new-service-b"] + + monkeypatch.setattr(rollout_lending, "return_rollout", return_loan) + return collectives + + +def test_keep_whole_pool_lent_while_complete_batch_ready(monkeypatch): + collectives = install_transport(monkeypatch) + ranks = [trainer(rank) for rank in range(2)] + assert run_ranks(lifecycle.ensure_batch_supply, ranks) == [None, None] + assert all(state.calls == [] for state in ranks) + assert ranks[0].runtime_elastic_executor._lent_rollout is not None + assert len([call for call in collectives.calls if call[1] == "gather"]) == 2 + + +def test_return_pool_before_supply_wait_then_publish_before_reopen(monkeypatch): + install_transport(monkeypatch) + ranks = [trainer(rank, ready_groups=1) for rank in range(2)] + assert run_ranks(lifecycle.ensure_batch_supply, ranks) == [None, None] + for state in ranks: + assert state.calls.index("pause") < state.calls.index("publish") + assert state.calls.index("refresh-membership") < state.calls.index("publish") + assert state.calls.index("publish") < state.calls.index("rebind") < state.calls.index("resume") + assert all(handle.is_ready for handle in state.handles) + assert state._restored_rollouts_pending_publication == set() + assert state.train_engine._elastic_step_hybrid_workers == () + assert ("set-version", 23) in state.calls + assert ranks[0].calls.index("release-hybrids") < ranks[0].calls.index("publish") + + +def test_supply_restore_uses_publication_hook_on_every_rank(monkeypatch): + install_transport(monkeypatch) + ranks = [trainer(rank, ready_groups=1) for rank in range(2)] + for state in ranks: + def restore(state=state): + state.calls.append("invalidate-old-publication") + return lifecycle.restore_for_publication(state) + state._restore_lent_rollouts_for_publication = restore + assert run_ranks(lifecycle.ensure_batch_supply, ranks) == [None, None] + for state in ranks: + assert state.calls.index("pause") < state.calls.index("invalidate-old-publication") + assert state.calls.index("invalidate-old-publication") < state.calls.index("publish") + + +def test_restore_failure_is_broadcast_to_every_rank_without_reopen(monkeypatch): + install_transport(monkeypatch) + + def failed_return(executor): + raise RuntimeError("service restart failed") + + monkeypatch.setattr(rollout_lending, "return_rollout", failed_return) + ranks = [trainer(rank, ready_groups=1) for rank in range(2)] + results = run_ranks(lifecycle.ensure_batch_supply, ranks) + assert all(isinstance(result, RuntimeError) and "service restart failed" in str(result) + for result in results) + assert all("publish" not in state.calls and "resume" not in state.calls for state in ranks) + assert all(not handle.is_ready for state in ranks for handle in state.handles) + + +def test_failed_publication_keeps_restored_engines_closed(monkeypatch): + install_transport(monkeypatch) + ranks = [trainer(rank, ready_groups=1, publish_failure=True) for rank in range(2)] + results = run_ranks(lifecycle.ensure_batch_supply, ranks) + assert all(isinstance(result, RuntimeError) and "weight install failed" in str(result) + for result in results) + assert all("resume" not in state.calls and "rebind" not in state.calls for state in ranks) + assert all(not handle.is_ready for state in ranks for handle in state.handles) + + +def test_one_rank_publication_failure_prevents_every_rank_reopening(monkeypatch): + install_transport(monkeypatch) + ranks = [trainer(0, ready_groups=1, publish_failure=True), + trainer(1, ready_groups=1, publish_failure=False)] + results = run_ranks(lifecycle.ensure_batch_supply, ranks) + assert all(isinstance(result, RuntimeError) and "weight install failed" in str(result) + for result in results) + assert all("resume" not in state.calls and "rebind" not in state.calls for state in ranks) + assert all(not handle.is_ready for state in ranks for handle in state.handles) + + +def test_any_dp_source_shortage_returns_pool(monkeypatch): + install_transport(monkeypatch) + ranks = [trainer(0), trainer(1, ready_groups=1)] + ranks[1]._is_rollout_source = lambda: True + assert run_ranks(lifecycle.ensure_batch_supply, ranks) == [None, None] + assert all("publish" in state.calls for state in ranks) + + +def test_partial_loan_keeps_available_rollout_supply_running(monkeypatch): + collectives = install_transport(monkeypatch) + ranks = [trainer(rank, ready_groups=0) for rank in range(2)] + ranks[0].runtime_elastic_executor._lent_rollout["engines"] = [SimpleNamespace(instance_id="engine-a")] + assert run_ranks(lifecycle.ensure_batch_supply, ranks) == [None, None] + assert all(state.calls == [] for state in ranks) + assert not any(call[1] == "gather" for call in collectives.calls) + + +def test_disabled_resource_management_avoids_collectives(monkeypatch): + state = trainer(0, loan=False) + state.config.global_resource_planner.enabled = False + + class ForbiddenCollectives: + def __getattr__(self, name): + raise AssertionError(f"disabled resource management entered dist.{name}") + + monkeypatch.setattr(lifecycle, "dist", ForbiddenCollectives()) + lifecycle.ensure_batch_supply(state) + assert state.calls == [] + + +def test_reopen_without_pending_install_is_noop(): + state = trainer(0) + lifecycle.reopen_after_publication(state) + assert all(not handle.is_ready for handle in state.handles) + + +def test_loan_publication_does_not_reenable_unrelated_unhealthy_engine(): + state = trainer(0) + unhealthy = SimpleNamespace(instance_id="unhealthy-other", is_ready=False) + state.handles.append(unhealthy) + lifecycle.apply_availability(state, ("engine-a", "engine-b")) + state._restored_rollouts_pending_publication = {"engine-a", "engine-b"} + lifecycle.reopen_after_publication(state) + assert all(handle.is_ready for handle in state.handles[:2]) + assert unhealthy.is_ready is False + + +def test_pending_hybrid_activation_gets_post_boundary_catchup_even_without_grp_gain(): + from RL_Framework.trainer.async_rl_trainer import AsyncRLTrainer + events = [] + current_plan = object() + state = SimpleNamespace( + config=object(), + runtime_elastic_executor=SimpleNamespace(has_pending_boundary_work=lambda: True), + global_resource_planner=SimpleNamespace( + _history=[], _build_current_plan=lambda config, history: current_plan), + _run_global_resource_planner=lambda *args: pytest.fail("pending join skipped catchup"), + _runtime_reconfiguration_boundary=lambda *args: events.append(("noop", args)), + ) + + def apply(step, decision, stats): + events.append((step, decision)) + state._runtime_boundary_published = True + + state._apply_global_resource_planner_decision = apply + AsyncRLTrainer._run_global_resource_planner_at_boundary(state, 9, [], {}) + assert len(events) == 1 + step, decision = events[0] + assert step == 9 + assert decision.reason == "hybrid_boundary_maintenance" + assert decision.should_reconfigure is True + assert decision.resource_should_reconfigure is False + assert decision.current_plan is decision.candidate_plan is current_plan + assert state._runtime_post_boundary is None + + +def test_no_pending_activation_keeps_normal_grp_post_boundary(): + from RL_Framework.trainer.async_rl_trainer import AsyncRLTrainer + events = [] + state = SimpleNamespace( + runtime_elastic_executor=SimpleNamespace(has_pending_boundary_work=lambda: False), + _run_global_resource_planner=lambda *args: events.append("planner"), + _runtime_reconfiguration_boundary=lambda step, point: events.append((step, point)), + ) + AsyncRLTrainer._run_global_resource_planner_at_boundary(state, 9, [], {}) + assert events == ["planner", (9, "post")] + assert state._runtime_post_boundary is None + + +def test_repartition_publication_uses_new_topology_before_client_reopen(): + from RL_Framework.trainer.async_rl_trainer import AsyncRLTrainer + state = SimpleNamespace( + _use_heterogeneous=True, + rollout_engine=SimpleNamespace(instance_configs=[ + {"instance_id": "old-engine-a", "tp_degree": 2}, + {"instance_id": "old-engine-b", "tp_degree": 2}, + ]), + config=SimpleNamespace(heterogeneous_rollout=SimpleNamespace(instances=[ + SimpleNamespace(instance_id="new-engine", tp=4), + ])), + ) + world_size, offsets = AsyncRLTrainer._nccl_rollout_topology(state) + assert world_size == 5 + assert offsets == {"new-engine": 0} + + +def test_resource_repartition_publishes_current_config_targets(monkeypatch): + from RL_Framework.infra.elastic.runtime_executor import RuntimeElasticExecutor + from RL_Framework.infra.elastic import core_reconfiguration + commands = [] + monkeypatch.setattr(core_reconfiguration, "dispatch_core_command", lambda executor, operation, **kwargs: + commands.append((operation, kwargs))) + state = SimpleNamespace(config=SimpleNamespace(heterogeneous_rollout=SimpleNamespace(instances=[ + SimpleNamespace(instance_id="new-engine", tp=4, gpus=[0, 1, 2, 3], host="rollout-node", port=8200), + ]))) + RuntimeElasticExecutor._publish_current_rollout_weights(state) + assert commands == [("publish", {"target": {"instances": [ + {"instance_id": "new-engine", "tp": 4, "gpus": [0, 1, 2, 3], + "host": "rollout-node", "port": 8200}, + ]}})] + + +def test_subset_publication_leaves_routing_to_rollback_owner(): + state = trainer(0) + lifecycle.apply_availability(state, ('engine-a','engine-b')) + state._publication_target_ids = {'engine-b'} + state._restored_rollouts_pending_publication = {'engine-b'} + lifecycle.reopen_after_publication(state) + assert state._rollout_unavailable_ids == {'engine-a','engine-b'} + assert state._restored_rollouts_pending_publication == {'engine-b'} + assert all(not handle.is_ready for handle in state.handles) + assert state.runtime_elastic_executor._rollout_lending_resume_pending diff --git a/tests/test_rollout_process.py b/tests/test_rollout_process.py new file mode 100644 index 0000000..30f3996 --- /dev/null +++ b/tests/test_rollout_process.py @@ -0,0 +1,200 @@ +"""Placement uses the existing Slurm allocation and precise owned process IDs.""" +import json +import os +from pathlib import Path +import shlex +import signal +import subprocess +import sys +import time +from types import SimpleNamespace as NS + +import pytest + +from RL_Framework.config import AsyncRLConfig +from RL_Framework.infra.elastic import rollout_process as placement + + +def test_remote_command_targets_current_allocation_and_preserves_gpu_ids(monkeypatch): + def query(argv, **kwargs): + assert argv == ['scontrol', 'show', 'hostnames', 'node[1-2]'] + return NS(stdout='node1\nnode2\n') + monkeypatch.setattr(placement.subprocess, 'run', query) + env = dict(SLURM_JOB_ID='123', SLURM_JOB_NODELIST='node[1-2]', SLURMD_NODENAME='node1') + command = placement.slurm_command('exec python serve.py', 'node2', [2, 3], env, step_name='owned') + args = shlex.split(command) + assert args[args.index('--jobid') + 1] == '123' + assert args[args.index('--nodelist') + 1] == 'node2' + assert '--overlap' in args and not any(arg.startswith('--gpus') for arg in args) + assert 'CUDA_VISIBLE_DEVICES=2,3' in args and 'PLANNED_CUDA_VISIBLE_DEVICES=2,3' in args + assert args[-3:] == ['bash', '-c', 'exec python serve.py'] + with pytest.raises(RuntimeError, match='not one unambiguous node'): + placement.slurm_command('serve', 'outside-job', [0], env) + + +def test_remote_host_without_allocation_is_rejected_and_local_stays_local(monkeypatch): + with pytest.raises(RuntimeError, match='current Slurm allocation'): + placement.placement_command('serve', 'remote-node', [0], {}) + assert placement.placement_command('serve', '127.0.0.1', [0], {}) == 'serve' + assert placement.is_local_host('node1.cluster', dict(SLURMD_NODENAME='node1')) + assert not placement.is_local_host('node2', dict(SLURMD_NODENAME='node1')) + + +def test_launch_remote_tracks_ownership_and_uses_a_dedicated_local_session(tmp_path, monkeypatch): + monkeypatch.setattr(placement, 'allocation_nodes', lambda env: ['node1', 'node2']) + seen = [] + def popen(command, **kwargs): + seen.append((command, kwargs)) + return NS(pid=123) + monkeypatch.setattr(placement.subprocess, 'Popen', popen) + meta = NS(instance_id='current', host='node2', gpus=[0, 1], command='serve') + proc = placement.launch_rollout_process(meta, dict(SLURM_JOB_ID='77', SLURMD_NODENAME='node1'), control_dir=tmp_path) + assert seen[0][1]['start_new_session'] is True + assert 'RL_Framework.infra.elastic.rollout_process' in seen[0][0] + assert proc._libra_rollout_owner['job_id'] == '77' + assert Path(proc._libra_rollout_owner['state_path']).parent == tmp_path / 'process-owners' + + +def test_remote_stop_cancels_only_receipted_step_and_zero_waits_for_verified_exit(tmp_path, monkeypatch): + path = tmp_path / 'owned.json' + path.write_text(json.dumps(dict(token='token', job_id='77', step_id='9', status='running'))) + owner = dict(remote=True, token='token', job_id='77', state_path=str(path), step_name='owned') + proc = NS(_libra_rollout_owner=owner) + proc.poll = lambda: 0 + stopped, commands = [], [] + monkeypatch.setenv('SLURM_JOB_ID', '77') + monkeypatch.setattr(placement, 'stop_owned_session', lambda proc, timeout: stopped.append(timeout)) + def run(argv, **kwargs): + commands.append(argv) + if argv[0] == 'scancel': + assert argv == ['scancel', '--signal=TERM', '77.9'] + path.write_text(json.dumps(dict(token='token', job_id='77', step_id='9', status='exited', owned_processes_exited=True))) + return NS(stdout='') + monkeypatch.setattr(placement.subprocess, 'run', run) + placement.stop_rollout_process(proc, timeout=0) + assert stopped == [0] and commands == [['scancel', '--signal=TERM', '77.9']] + path.write_text(json.dumps(dict(token='other-owner', job_id='77', step_id='9'))) + with pytest.raises(RuntimeError, match='does not match'): + placement.stop_rollout_process(proc, timeout=0) + assert commands == [['scancel', '--signal=TERM', '77.9']] + + +def test_owned_stop_covers_nested_descendant_sessions(tmp_path): + child_pid = tmp_path / 'child.pid' + script = ('import subprocess,sys,time,pathlib; ' + 'p=subprocess.Popen([sys.executable,"-c","import signal,time; signal.signal(signal.SIGTERM, signal.SIG_IGN); time.sleep(60)"],start_new_session=True); ' + f'pathlib.Path({str(child_pid)!r}).write_text(str(p.pid)); time.sleep(60)') + proc = subprocess.Popen([sys.executable, '-c', script], start_new_session=True) + try: + deadline = time.monotonic() + 3 + while not child_pid.exists() and time.monotonic() < deadline: + time.sleep(.01) + assert child_pid.exists() + pid = int(child_pid.read_text()) + placement.stop_owned_session(proc, timeout=.05) + assert proc.poll() is not None + record = placement._process_record(pid) + assert record is None or record['state'] == 'Z' + finally: + if proc.poll() is None: + os.killpg(proc.pid, signal.SIGKILL) + proc.wait() + if child_pid.exists(): + try: + os.kill(int(child_pid.read_text()), signal.SIGKILL) + except ProcessLookupError: + pass + + +@pytest.mark.parametrize('change', ['wrong-job', 'wrong-instance', 'wrong-port', 'wrong-role', 'wrong-control']) +def test_adopted_identity_mismatch_never_signals_any_pid(monkeypatch, change): + monkeypatch.setenv('SLURM_JOB_ID', '777') + contract = dict(pid=120, host='127.0.0.1', instance_id='wanted', port=9000, job_id='777', control_dir='/fyp/run/control') + argv = [sys.executable, '-m', 'RL_Framework.scripts.restartable_vllm_server', + '--instance-id', 'wanted', '--health-url', 'http://127.0.0.1:9000/health', + '--control-dir', '/fyp/run/control'] + environment = dict(SLURM_JOB_ID='777') + if change == 'wrong-job': + environment['SLURM_JOB_ID'] = 'other-job' + elif change == 'wrong-instance': + argv[argv.index('--instance-id')+1] = 'unrelated' + elif change == 'wrong-port': + argv[argv.index('--health-url')+1] = 'http://127.0.0.1:9100/health' + elif change == 'wrong-role': + argv[2] = 'unrelated.application' + elif change == 'wrong-control': + argv[-1] = '/fyp/another-run/control' + monkeypatch.setattr(placement, '_process_record', lambda pid: dict(pid=pid, state='S', start='1', parent=1, group=1)) + monkeypatch.setattr(placement, '_process_command', lambda pid: (argv, environment)) + monkeypatch.setattr(placement.os, 'kill', lambda *args: pytest.fail('unverified process must never be signalled')) + with pytest.raises(RuntimeError): + placement._stop_adopted_on_host(contract, timeout=0) + + +def test_remote_adopted_pid_is_verified_and_stopped_on_target_allocation_node(monkeypatch): + monkeypatch.setattr(placement, 'allocation_nodes', lambda env: ['train', 'rollout']) + monkeypatch.setattr(placement, '_verify_adopted_identity', lambda *args, **kwargs: (_ for _ in ()).throw(RuntimeError('not local launcher'))) + monkeypatch.setattr(placement.os, 'kill', lambda *args: pytest.fail('remote PID must never be signalled locally')) + calls = [] + def run(argv, **kwargs): + calls.append(argv) + assert argv[0] == 'srun' + assert argv[argv.index('--nodelist')+1] == 'rollout' + assert kwargs['timeout'] is None + command = shlex.split(argv[-1]) + contract = json.loads(command[command.index('--stop-adopted')+1]) + assert contract['pid'] == 999 and contract['job_id'] == '777' + assert not contract.get('discover_service') + return NS(stdout=json.dumps(dict(instance_id='wanted', port=9000, job_id='777', owned_processes_exited=True))) + monkeypatch.setattr(placement.subprocess, 'run', run) + meta = NS(pid=999, host='rollout', instance_id='wanted', port=9000, gpus=[0, 1]) + receipt = placement.stop_adopted_rollout(meta, dict(SLURM_JOB_ID='777', SLURMD_NODENAME='train'), timeout=0) + assert receipt['owned_processes_exited'] and len(calls) == 1 + + +def test_verified_local_srun_registry_pid_authorizes_unique_remote_service_lookup(monkeypatch): + monkeypatch.setattr(placement, 'allocation_nodes', lambda env: ['train', 'rollout']) + record = dict(pid=111, start='birth', parent=1, state='S', group=1) + monkeypatch.setattr(placement, '_verify_adopted_identity', lambda pid, contract, launcher=False: record) + monkeypatch.setattr(placement, '_still_owned', lambda record: False) + def run(argv, **kwargs): + command = shlex.split(argv[-1]) + contract = json.loads(command[command.index('--stop-adopted')+1]) + assert contract['discover_service'] is True + return NS(stdout=json.dumps(dict(instance_id='wanted', port=9000, job_id='777', owned_processes_exited=True))) + monkeypatch.setattr(placement.subprocess, 'run', run) + meta = NS(pid=111, host='rollout', instance_id='wanted', port=9000, gpus=[0]) + assert placement.stop_adopted_rollout(meta, dict(SLURM_JOB_ID='777', SLURMD_NODENAME='train'))['owned_processes_exited'] + + +def test_local_adopted_service_stops_only_verified_tree_and_preserves_other_instance(tmp_path, monkeypatch): + monkeypatch.setenv('SLURM_JOB_ID', '777') + script = tmp_path / 'restartable_vllm_server.py' + script.write_text('import time; time.sleep(60)') + control = str(tmp_path / 'control') + commands = [[sys.executable, str(script), '--instance-id', name, '--health-url', + f'http://127.0.0.1:{port}/health', '--control-dir', control] + for name, port in [('wanted', 9000), ('unrelated', 9100)]] + processes = [subprocess.Popen(command) for command in commands] + try: + time.sleep(.05) + meta = NS(pid=processes[0].pid, host='127.0.0.1', instance_id='wanted', port=9000) + receipt = placement.stop_adopted_rollout(meta, dict(os.environ), timeout=.05, control_dir=control) + assert receipt['owned_processes_exited'] and receipt['pid'] == processes[0].pid + processes[0].wait(timeout=1) + assert processes[1].poll() is None + finally: + for proc in processes: + if proc.poll() is None: + proc.terminate() + proc.wait(timeout=1) + + +def test_remote_adopted_exit_without_matching_owner_receipt_is_rejected(monkeypatch): + monkeypatch.setattr(placement, 'allocation_nodes', lambda env: ['train', 'rollout']) + monkeypatch.setattr(placement, '_verify_adopted_identity', lambda *args, **kwargs: (_ for _ in ()).throw(RuntimeError())) + monkeypatch.setattr(placement.subprocess, 'run', lambda *args, **kwargs: NS(stdout=json.dumps( + dict(instance_id='another', port=9000, job_id='777', owned_processes_exited=True)))) + meta = NS(pid=999, host='rollout', instance_id='wanted', port=9000, gpus=[0]) + with pytest.raises(RuntimeError, match='differs from the requested owner'): + placement.stop_adopted_rollout(meta, dict(SLURM_JOB_ID='777', SLURMD_NODENAME='train')) diff --git a/tests/test_rollout_supply.py b/tests/test_rollout_supply.py new file mode 100644 index 0000000..fbe6ae1 --- /dev/null +++ b/tests/test_rollout_supply.py @@ -0,0 +1,132 @@ +"""Already generated complete-group supply survives all-engine EHP loans.""" +from types import SimpleNamespace + +from RL_Framework.infra.elastic.rollout_supply import ready_complete_samples +from RL_Framework.infra.execution.async_runner import AsyncTaskRunner, TaskResult +from RL_Framework.infra.execution.batch_dispatcher import BatchTaskDispatcher +from RL_Framework.trainer.async_rl_trainer import AsyncRLTrainer + + +def member(group, index, **kwargs): + return dict(grpo_group_id=group, grpo_rollout_index=index, **kwargs) + + +def trainer_fixture(paused=False): + runner = AsyncTaskRunner(max_queue_size=32) + if paused: + runner.pause() + dispatcher = BatchTaskDispatcher(runner, SimpleNamespace(), lambda task: None) + trainer = AsyncRLTrainer.__new__(AsyncRLTrainer) + trainer.config = SimpleNamespace(n_samples=4, global_resource_planner=SimpleNamespace( + runtime_batch_collection_timeout_s=.05, runtime_batch_collection_max_retries=0)) + trainer.is_main_process = True + trainer._pending_grpo_groups = {} + trainer.dispatcher = dispatcher + trainer.staleness_manager = SimpleNamespace(on_batch_consumed=lambda count: None) + trainer._reset_rollout_pipeline_after_reconfigure = lambda: None + return trainer + + +def add_result(dispatcher, task_id, data, created=None): + dispatcher._pending_results[task_id] = TaskResult(data, task_id, created or float(task_id), float(task_id)) + dispatcher._active_task_ids.add(task_id) + + +def test_snapshot_is_non_consuming_and_matches_real_wait_order(): + trainer = trainer_fixture() + dispatcher = trainer.dispatcher + add_result(dispatcher, 10, member("b", 1), created=2) + add_result(dispatcher, 20, member("a", 0), created=1) + original = list(dispatcher._pending_results) + data = dispatcher.peek_ready_data() + assert [item["grpo_group_id"] for item in data] == ["a", "b"] + data.clear() + assert list(dispatcher._pending_results) == original + assert [result.data["grpo_group_id"] for result in dispatcher.wait_results(2)] == ["a", "b"] + + +def test_ready_supply_combines_partial_groups_and_ignores_failed_duplicate_members(): + trainer = trainer_fixture(paused=True) + trainer._pending_grpo_groups = {"a": [member("a", 0), member("a", 1)], + "incomplete": [member("incomplete", 0)]} + for task_id, value in enumerate([ + member("a", 1), member("a", 2), member("a", 3), + member("b", 0), member("b", 1), member("b", 2), member("b", 3), + member("incomplete", 1), member("incomplete", 2, status="failed"), + member("incomplete", 0), None, + ]): + add_result(trainer.dispatcher, task_id, value) + before = list(trainer.dispatcher._pending_results) + pending_before = {group: list(values) for group, values in trainer._pending_grpo_groups.items()} + assert ready_complete_samples(trainer) == 8 + assert list(trainer.dispatcher._pending_results) == before + assert trainer._pending_grpo_groups == pending_before + + +def test_collector_returns_existing_complete_group_without_dispatcher_call(): + trainer = trainer_fixture(paused=True) + group = [member("ready", index) for index in range(4)] + trainer._pending_grpo_groups = {"ready": group} + trainer.dispatcher = SimpleNamespace() + assert trainer._collect_complete_grpo_batch(iter(()), 4) == group + assert trainer._pending_grpo_groups == {} + + +def test_paused_real_dispatcher_consumes_enough_partial_and_ready_data_without_submitting(): + trainer = trainer_fixture(paused=True) + trainer._pending_grpo_groups = {"a": [member("a", 0), member("a", 1)]} + add_result(trainer.dispatcher, 1, member("a", 2)) + add_result(trainer.dispatcher, 2, member("a", 3)) + add_result(trainer.dispatcher, 3, member("future", 0)) + def never_submit(**kwargs): + raise AssertionError("paused ready data must not submit rollout work") + trainer.dispatcher.active_submit_and_wait = never_submit + batch = trainer._collect_complete_grpo_batch(iter(()), 4) + assert [value["grpo_rollout_index"] for value in batch] == [0, 1, 2, 3] + assert not trainer.dispatcher._pending_inputs + # Reading extra already-ready members preserves them for the next group. + assert trainer._pending_grpo_groups["future"] == [member("future", 0)] + assert ready_complete_samples(trainer) == 0 + + +def test_collector_deduplicates_ready_members_and_requests_only_deficit(): + trainer = trainer_fixture() + trainer._pending_grpo_groups = {"a": [member("a", 0), member("a", 1), member("a", 2)]} + requested = [] + def missing(**kwargs): + requested.append(kwargs["batch_size"]) + return [member("a", 3)] + trainer.dispatcher.active_submit_and_wait = missing + batch = trainer._collect_complete_grpo_batch(iter(()), 4) + assert requested == [1] + assert len(batch) == 4 + # Retransmissions of a previously consumed group cannot create another batch. + for index in range(4): + add_result(trainer.dispatcher, index, member("a", index)) + assert ready_complete_samples(trainer) == 0 + + +def test_duplicate_and_invalid_member_indices_do_not_complete_a_group(): + trainer = trainer_fixture() + trainer._pending_grpo_groups = {"a": [member("a", 0), member("a", 1), member("a", 2), + member("a", 2), member("a", 8)]} + assert ready_complete_samples(trainer) == 0 + + +def test_duplicate_task_identity_with_missing_member_index_is_not_counted_twice(): + trainer = trainer_fixture() + trainer._pending_grpo_groups = {"a": [member("a", 0, task_id=10), member("a", 1, task_id=11), + member("a", 2, task_id=12)]} + add_result(trainer.dispatcher, 10, dict(grpo_group_id="a", task_id=10)) + assert ready_complete_samples(trainer) == 0 + + +def test_paused_insufficient_supply_never_enqueues_new_rollouts(): + trainer = trainer_fixture(paused=True) + trainer._pending_grpo_groups = {"a": [member("a", 0), member("a", 1)]} + def cannot_submit(**kwargs): + raise AssertionError("paused dispatcher submitted new work") + trainer.dispatcher.active_submit_and_wait = cannot_submit + assert trainer._collect_complete_grpo_batch(iter(()), 4) == [] + assert not trainer.dispatcher._pending_inputs + assert len(trainer._pending_grpo_groups["a"]) == 2 diff --git a/tests/test_runtime_blue_green_gpu_ownership.py b/tests/test_runtime_blue_green_gpu_ownership.py new file mode 100644 index 0000000..7948e81 --- /dev/null +++ b/tests/test_runtime_blue_green_gpu_ownership.py @@ -0,0 +1,122 @@ +"""Release occupied rollout GPUs before replacement model initialization.""" +from dataclasses import replace +from types import SimpleNamespace as NS +import pytest +from RL_Framework.infra.elastic.runtime_executor import ManagedRolloutProcess, RuntimeElasticExecutor + +def process(name, gpus, port, *, host='node-a', pid=11): + return ManagedRolloutProcess(name, 'unused', pid, list(gpus), port, host, len(gpus)) + + + +class Executor(RuntimeElasticExecutor): + def __init__(self, current, desired, *, fail_stop=None): + config = NS(global_resource_planner=NS(runtime_drain_before_reconfigure=True), + heterogeneous_rollout=NS(vllm_base_port=8000, instances=[ + NS(instance_id=meta.instance_id, port=meta.port) for meta in desired])) + super().__init__(config=config, planner=None) + self.desired = desired + self.events = [] + self.fail_stop = fail_stop + self._process_meta = {meta.instance_id: meta for meta in current} + self._processes = {meta.instance_id: object() for meta in current} + + def _desired_rollout_processes(self, plan): + return self.desired + + def _command_for_instance_id(self, instance_id, port): + return f'start {instance_id} {port}' + + def _start_one_rollout_process(self, meta): + self.events.append(('start', meta.instance_id)) + old = self._process_meta.get(meta.instance_id) + if old is not None and not self._same_rollout_process(old, meta): + self._blue_green_old_processes = getattr(self, '_blue_green_old_processes', []) + [ + (old, self._processes.get(meta.instance_id))] + started = replace(meta, pid=100 + len(self.events)) + self._process_meta[meta.instance_id] = started + self._processes[meta.instance_id] = object() + return started + + def _wait_rollout_process_start_ready(self, meta): + self.events.append(('ready', meta.instance_id)) + + def _stop_process_meta(self, meta, proc): + self.events.append(('stop', meta.instance_id)) + if meta.instance_id == self.fail_stop: + raise RuntimeError('old owner did not exit') + return [meta] + + def _reconfigure_rollout_engine(self, plan, *, wait_ready): + self.events.append(('cutover', wait_ready)) + + def _sleep_before_rollout_stop(self): + self.events.append(('grace-before',)) + + def _sleep_after_rollout_stop(self, stopped): + if stopped: + self.events.append(('grace-after', tuple(meta.instance_id for meta in stopped))) + + + +def test_merge_releases_both_old_tp1_owners_before_new_tp2_and_preserves_survivor(): + keep = process('keep', [1], 8000) + executor = Executor([keep, process('old-b', [2], 8001), process('old-c', [3], 8002)], + [replace(keep), process('merged', [2, 3], 8001)]) + stopped, started = executor._blue_green_reconfigure_rollout_processes(None) + events = executor.events + assert events.index(('stop', 'old-b')) < events.index(('start', 'merged')) + assert events.index(('stop', 'old-c')) < events.index(('start', 'merged')) + assert events.index(('grace-after', ('old-b', 'old-c'))) < events.index(('start', 'merged')) + assert executor._process_meta['keep'] is keep and ('stop', 'keep') not in events + assert {meta.instance_id for meta in stopped} == {'old-b', 'old-c'} + assert [meta.instance_id for meta in started] == ['merged'] + assert started[0].port == 8003 + + + +def test_split_releases_tp2_once_before_either_tp1_starts(): + keep = process('keep', [1], 8000) + executor = Executor([keep, process('merged', [2, 3], 8003)], + [replace(keep), process('new-b', [2], 8001), process('new-c', [3], 8002)]) + stopped, started = executor._blue_green_reconfigure_rollout_processes(None) + assert executor.events.count(('stop', 'merged')) == 1 + for name in ('new-b', 'new-c'): + assert executor.events.index(('stop', 'merged')) < executor.events.index(('start', name)) + assert [meta.instance_id for meta in stopped] == ['merged'] + assert [meta.instance_id for meta in started] == ['new-b', 'new-c'] + assert executor._process_meta['keep'] is keep + + + +def test_nonoverlapping_gpu_keeps_start_ready_cutover_before_old_stop(): + executor = Executor([process('old', [2], 8000)], [process('new', [3], 8001)]) + executor._blue_green_reconfigure_rollout_processes(None) + events = executor.events + assert events.index(('start', 'new')) < events.index(('ready', 'new')) + assert events.index(('ready', 'new')) < events.index(('cutover', True)) + assert events.index(('cutover', True)) < events.index(('stop', 'old')) + assert all(event[0] != 'grace-after' for event in events) + + + +def test_overlapping_stop_failure_preserves_owner_and_prevents_any_launch(): + old = process('old', [2], 8000) + executor = Executor([old], [process('new', [2], 8001)], fail_stop='old') + owned_process = executor._processes['old'] + with pytest.raises(RuntimeError, match='old owner did not exit'): + executor._blue_green_reconfigure_rollout_processes(None) + assert executor._process_meta['old'] is old + assert executor._processes['old'] is owned_process + assert executor.events == [('stop', 'old')] + + + +def test_same_id_replacement_stops_old_once_and_retains_new_process(): + old = process('same', [2], 8000) + executor = Executor([old], [process('same', [2, 3], 8000)]) + stopped, started = executor._blue_green_reconfigure_rollout_processes(None) + assert executor.events.index(('stop', 'same')) < executor.events.index(('start', 'same')) + assert executor.events.count(('stop', 'same')) == 1 + assert stopped == [old] and executor._process_meta['same'] is started[0] + assert started[0].pid != old.pid and started[0].port == 8001 diff --git a/tests/test_runtime_boundary_integration.py b/tests/test_runtime_boundary_integration.py index dddd26f..2e3998a 100644 --- a/tests/test_runtime_boundary_integration.py +++ b/tests/test_runtime_boundary_integration.py @@ -6,6 +6,7 @@ from types import SimpleNamespace import pytest +import torch from RL_Framework.config import AsyncRLConfig from RL_Framework.engine.megatron_core_train_engine import MegatronCoreTrainEngine @@ -28,6 +29,7 @@ class NoReconfigurationDecision: should_reconfigure = False candidate_plan = None + current_plan = None reason = "same_plan" num_requests = 1 @@ -133,12 +135,22 @@ def test_release_only_post_is_visible_at_next_pre_before_gradient_freeze(tmp_pat engine._hybrid_gradient_condition = threading.Condition() engine._elastic_active_gradient_timeout_s = 0.0 engine._elastic_step_hybrid_workers = () + engine.optimizer_revision = 0 + engine.current_version = 0 + engine._elastic_seen_payloads = set() + engine._elastic_batch_partition = None + engine._elastic_gradient_update_callback = None + engine._model_main_gradients = lambda: [torch.zeros(1)] follower.train_engine = engine follower._elastic_membership_run_id = "membership-run" follower._elastic_membership_epochs = {} try: - pool.join_replica("replica0", ["rollout0"], "dp0").result(timeout=2) + pool.join_replica( + "replica0", ["rollout0"], "dp0", + replica_activation_barrier=lambda *args: {"acknowledged": True, "state_version": 0, + "optimizer_revision": 0}, + ).result(timeout=2) executor.hybrid_runtime_state() assert follower._prepare_elastic_training_step(0) == ["replica0"] membership_path = tmp_path / "tasks/membership/replica0.json" @@ -204,15 +216,18 @@ def test_final_boundary_waits_for_leader_close_before_follower_teardown(tmp_path assert follower._elastic_gradient_server is None -def test_cleanup_releases_cached_rollout_weight_communicators(tmp_path): +def test_cleanup_releases_cached_rollout_weight_communicators(tmp_path, monkeypatch): """Catch interpreter-shutdown aborts from lingering TCPStore owners.""" leader, _follower = make_cleanup_trainers(tmp_path) cache_key = ("127.0.0.1", 29620, 2, 0, "cuda:0") - nccl_weight_sync._COMMUNICATOR_CACHE[cache_key] = object() + communicator = SimpleNamespace(comm=None, disabled=False, available=True, group=object()) + monkeypatch.setitem(nccl_weight_sync._COMMUNICATOR_CACHE, cache_key, communicator) leader._cleanup() assert nccl_weight_sync._COMMUNICATOR_CACHE == {} + assert communicator.disabled and not communicator.available + assert communicator.group is None def test_final_draining_transaction_reaches_terminal_then_done(tmp_path): diff --git a/tests/test_runtime_core_capabilities.py b/tests/test_runtime_core_capabilities.py new file mode 100644 index 0000000..3480f86 --- /dev/null +++ b/tests/test_runtime_core_capabilities.py @@ -0,0 +1,218 @@ +"""Live Core ranks constrain executable GRP layouts while EHP stays elastic.""" +from types import SimpleNamespace as NS +from unittest.mock import Mock + +import pytest + +from RL_Framework.config import AsyncRLConfig, HeterogeneousInstanceConfig +from RL_Framework.infra.cost_model.global_resource_planner import ( + CoreRuntimeCapabilities, ElasticHybridSignal, GlobalResourcePlan, + GlobalResourcePlanner, PlannerDecision, RuntimeResourceMetrics, +) +from RL_Framework.infra.cost_model.model import RolloutClusterConfig, TrainParallelConfig +from RL_Framework.infra.cost_model.optimizer import OptimizationResult +from RL_Framework.infra.cost_model.preflight_planner import PreflightPlanner +from RL_Framework.infra.elastic.runtime_executor import RuntimeElasticExecutor + + +ROWS = [dict(input_len=1024, output_len=2048)] * 16 + + +def config(): + result = AsyncRLConfig(model_path='not-loaded') + result.n_total_gpus, result.train_gpus, result.rollout_gpus = 12, 4, 8 + result.train_tp_size, result.train_dp_size = 4, 1 + result.batch_size, result.micro_batch_size, result.n_samples = 16, 1, 4 + result.heterogeneous_rollout.enabled = True + result.heterogeneous_rollout.instances = [HeterogeneousInstanceConfig( + instance_id=f'rollout-{index}', tp=2, gpus=[index % 2 * 2, index % 2 * 2 + 1], + host=f'rollout-{index // 2}', port=8000 + index) for index in range(4)] + planner = result.global_resource_planner + planner.enabled = True + planner.initial_allocation_strategy = 'grp' + planner.allowed_train_tp, planner.allowed_train_pp = [1, 2, 4], [1] + planner.allowed_train_ep, planner.allowed_rollout_tp = [1], [1, 2, 4] + planner.micro_batch_sizes = [1] + planner.allocation_granularity_gpus = 4 + planner.min_train_gpus = planner.min_rollout_gpus = 4 + planner.elastic_hybrid_planning_enabled = True + planner.runtime_reconfigure_training = True + planner.runtime_training_pool_only = True + planner.runtime_training_pool_plan_only = False + planner.runtime_manage_rollout_processes = planner.hybrid_worker_launch_enabled = True + planner.elastic_hybrid_require_planner_signal = True + planner.elastic_hybrid_min_rollout_gpus = 4 + planner.reconfiguration_cost_s = planner.min_gain_ratio = 0 + planner.warmup_steps = 0 + planner.min_history_size = planner.plan_interval = 1 + return result + + +def engine(world=4, tp=4): + return NS(world_size=world, train_tp_size=tp, train_pp_size=1, + train_cp_size=1, train_ep_size=1, train_dp_size=world // tp, + get_parallel_state=lambda: dict(train_tp=tp, train_pp=1, train_cp=1, + train_ep=1, train_dp=world // tp), + reconfigure_parallelism=Mock()) + + +def result(train_gpus=8, tp=4): + return OptimizationResult(train_config=TrainParallelConfig(tp=tp, pp=1, + dp=train_gpus // tp, b_micro=1), + rollout_config=RolloutClusterConfig(tp_list=[4] if train_gpus == 8 else [2, 2, 2, 2]), + t_train=.1, t_rollout=.1, t_global=.1) + + +def as_plan(value): + return GlobalResourcePlan(train_config=value.train_config, rollout_config=value.rollout_config, + t_train=value.t_train, t_rollout=value.t_rollout, t_global=value.t_global, + n_total_gpus=12, max_concurrent_rollouts=16) + + +def test_runtime_replaces_unlaunchable_core_resize_but_preserves_hybrid_demand(): + cfg = config() + planner = GlobalResourcePlanner.from_config(cfg) + planner.runtime_capabilities = CoreRuntimeCapabilities.from_engine(engine()) + planner.optimizer.optimize = Mock(side_effect=[result(8), result(4)]) + metrics = RuntimeResourceMetrics(step=5, train_time_s=1, rollout_time_s=10) + decision = planner.plan_if_needed(step=5, config=cfg, batch=ROWS, runtime_metrics=metrics) + assert decision.candidate_plan.train_gpus == 4 + assert decision.candidate_plan.train_config.tp == 4 + assert decision.candidate_plan.train_config.dp == 1 + assert decision.candidate_plan.rollout_gpus == 8 + assert decision.elastic_hybrid_signal.desired_replicas == 1 + assert decision.elastic_hybrid_signal.reason == 'candidate_requests_more_training_capacity' + assert not decision.resource_should_reconfigure + assert decision.should_reconfigure # Real EHP membership still changes. + first, second = planner.optimizer.optimize.call_args_list + assert first.kwargs['fixed_train_gpus'] == 0 + assert second.kwargs['fixed_train_gpus'] == 4 + assert second.kwargs['allowed_train_tp'] == [4] + assert second.kwargs['allowed_train_ep'] == second.kwargs['allowed_train_pp'] == [1] + assert cfg.global_resource_planner.fixed_train_gpus == 0 # No configuration workaround. + proof = decision.candidate_plan.metadata['runtime_candidate_filter'] + assert proof['reason'] == 'physical_core_rank_set_changed' + assert proof['unconstrained_candidate']['train']['n_gpus'] == 8 + assert proof['elastic_demand_preserved'] + + +def test_even_a_forced_resize_uses_ehp_instead_of_inventing_core_ranks(monkeypatch): + cfg = config() + planner = GlobalResourcePlanner.from_config(cfg) + planner.runtime_capabilities = CoreRuntimeCapabilities.from_engine(engine()) + monkeypatch.setenv('GRP_FORCE_TRAIN_GPUS', '8') + monkeypatch.setenv('GRP_FORCE_ROLLOUT_TP_LIST', '4') + decision = planner.plan_if_needed(step=5, config=cfg, batch=ROWS, + runtime_metrics=RuntimeResourceMetrics(step=5, train_time_s=1, rollout_time_s=10)) + assert decision.candidate_plan.train_gpus == 4 + assert not decision.resource_should_reconfigure + assert decision.should_reconfigure and decision.elastic_hybrid_signal.desired_replicas == 1 + assert decision.reason == 'elastic_hybrid_adjustment' + + +@pytest.mark.parametrize('initial_gpus', [4, 8]) +def test_startup_grp_retains_freedom_to_choose_the_initial_physical_world(initial_gpus): + planner = PreflightPlanner(config()) + assert planner.planner.runtime_capabilities is None + planner.planner.optimizer.optimize = Mock(return_value=result(initial_gpus)) + selected = planner.run(ROWS) + assert selected.applied and selected.planned_config.train_gpus == initial_gpus + assert selected.planned_config.rollout_gpus == 12 - initial_gpus + assert planner.planner.optimizer.optimize.call_args.kwargs['fixed_train_gpus'] == 0 + + +@pytest.mark.parametrize('axis', ['pp', 'cp', 'ep']) +def test_live_core_rebuild_rejects_unsupported_model_axes(axis): + candidate = as_plan(result(4)) + setattr(candidate.train_config, axis, 2) + candidate.train_config.dp = 1 + candidate.train_config.tp = 2 + capability = CoreRuntimeCapabilities.from_engine(engine()) + assert capability.incompatibility(candidate) == f'core_{axis}_rebuild_unsupported' + + +def test_fixed_core_keeps_rollout_repartition_and_actual_ehp_enabled(): + cfg = config() + planner = GlobalResourcePlanner.from_config(cfg) + planner.runtime_capabilities = CoreRuntimeCapabilities.from_engine(engine()) + candidate = result(4) + candidate.rollout_config = RolloutClusterConfig(tp_list=[4, 4]) + planner.optimizer.optimize = Mock(return_value=candidate) + decision = planner.plan_if_needed(step=5, config=cfg, batch=ROWS) + assert decision.resource_should_reconfigure + assert decision.candidate_plan.train_config.tp == 4 + assert decision.candidate_plan.train_config.dp == 1 + assert decision.candidate_plan.rollout_tp_list == [4, 4] + executor = RuntimeElasticExecutor(config=cfg, planner=planner, train_engine=engine()) + prepared = executor.prepare_transaction(decision) + assert prepared.should_execute and prepared.resource_reconfiguration + assert cfg.global_resource_planner.runtime_reconfigure_training + assert cfg.global_resource_planner.runtime_training_pool_only + assert not cfg.global_resource_planner.runtime_training_pool_plan_only + assert cfg.global_resource_planner.hybrid_worker_launch_enabled + + +def test_fixed_core_rejects_tp_dp_change_even_if_adapter_exposes_old_rebuild_hook(): + capability = CoreRuntimeCapabilities.from_engine(engine()) + assert not capability.rebuild_tp_dp + assert capability.incompatibility(as_plan(result(4, tp=2))) == "core_tp_dp_rebuild_unsupported" + + +def test_executor_rejects_stale_resize_before_any_generation_or_core_side_effect(): + cfg, backend = config(), engine() + dispatcher = NS(pause=Mock(), wait_until_idle=Mock(), resume=Mock()) + planner = GlobalResourcePlanner.from_config(cfg) + executor = RuntimeElasticExecutor(config=cfg, planner=planner, train_engine=backend, dispatcher=dispatcher) + executor._adopt_existing_rollout_processes = Mock() + decision = PlannerDecision(step=5, should_reconfigure=True, reason='stale_async_plan', + candidate_plan=as_plan(result(8)), resource_should_reconfigure=True) + prepared = executor.prepare_transaction(decision) + assert not prepared.should_execute and not prepared.drain_required + assert prepared.reason == 'runtime_candidate_rejected:physical_core_rank_set_changed' + assert prepared.coord_id == '' + executor._adopt_existing_rollout_processes.assert_not_called() + dispatcher.pause.assert_not_called() + dispatcher.wait_until_idle.assert_not_called() + backend.reconfigure_parallelism.assert_not_called() + + +def test_executor_preserves_valid_ehp_signal_from_a_stale_unlaunchable_core_candidate(): + cfg = config() + current = as_plan(result(4)) + signal = ElasticHybridSignal(step=5, desired_workers=1, current_workers=0, + max_workers=1, action='join', reason='training_bottleneck', + train_rollout_ratio=2, rollout_pressure=0, ttl_steps=20, replica_size_gpus=4) + decision = PlannerDecision(step=5, should_reconfigure=True, reason='reconfigure', + current_plan=current, candidate_plan=as_plan(result(8)), + elastic_hybrid_signal=signal, resource_should_reconfigure=True) + executor = RuntimeElasticExecutor(config=cfg, planner=GlobalResourcePlanner.from_config(cfg), train_engine=engine()) + prepared = executor.prepare_transaction(decision) + assert prepared.should_execute and not prepared.resource_reconfiguration + assert prepared.candidate_plan is current + assert prepared.effective_planner_signal is signal + + +def test_training_backend_without_rebuild_keeps_its_actual_tp_dp(): + backend = engine() + backend.reconfigure_parallelism = None + capability = CoreRuntimeCapabilities.from_engine(backend) + assert capability.incompatibility(as_plan(result(4))) == '' + assert capability.incompatibility(as_plan(result(4, tp=2))) == 'core_tp_dp_rebuild_unsupported' + + +def test_native_trainer_installs_live_capabilities_without_a_config_freeze(): + from RL_Framework.trainer.async_rl_trainer import AsyncRLTrainer + trainer = object.__new__(AsyncRLTrainer) + trainer.config = config() + planner = trainer.config.global_resource_planner + planner.runtime_coordinate_reconfiguration_ranks = False + planner.runtime_async_planning = False + trainer.is_main_process = True + trainer.train_engine = engine() + trainer.rollout_engine = trainer.dispatcher = None + trainer._runtime_boundary_protocol = trainer._elastic_gradient_server = None + trainer._elastic_membership_run_id = '' + trainer._is_rollout_source = lambda: True + trainer._init_global_resource_planner() + assert trainer.global_resource_planner.runtime_capabilities.world_size == 4 + assert trainer.config.global_resource_planner.fixed_train_gpus == 0 diff --git a/tests/test_runtime_elastic_executor.py b/tests/test_runtime_elastic_executor.py index df144e2..207d7e4 100644 --- a/tests/test_runtime_elastic_executor.py +++ b/tests/test_runtime_elastic_executor.py @@ -1,5 +1,9 @@ import json +from pathlib import Path +import shlex +import sys import time +from types import SimpleNamespace import pytest @@ -94,6 +98,83 @@ def _decision(cfg: AsyncRLConfig): return planner, decision +@pytest.mark.parametrize('slots', [ + {}, + {'rollout0': {'host': 'a', 'gpus': [0]}}, + {'rollout0': {'host': 'a', 'gpus': [0]}, 'rollout1': {'host': 'b', 'gpus': [1]}}, + {'rollout0': {'host': 'a', 'gpus': [0]}, 'rollout1': {'host': 'a', 'gpus': [0]}}, +]) +@pytest.mark.parametrize('existing_replica', [False, True]) +def test_real_hybrid_join_requires_safely_lent_complete_single_host_slots( + monkeypatch, slots, existing_replica, +): + from types import SimpleNamespace + from unittest.mock import Mock + + cfg = _config() + cfg.train_gpus, cfg.rollout_gpus = 4, 6 + cfg.train_tp_size, cfg.train_dp_size = 2, 2 + cfg.global_resource_planner.hybrid_worker_launch_enabled = True + executor = RuntimeElasticExecutor(config=cfg, planner=GlobalResourcePlanner.from_config(cfg)) + core = {f'dp{i}': SimpleNamespace(role='core_training') for i in range(2)} + candidates = {f'rollout{i}': SimpleNamespace(role='hybrid_rollout') for i in range(4)} + active = ([SimpleNamespace(worker_id='old0', role='hybrid_training'), + SimpleNamespace(worker_id='old1', role='hybrid_training')] + if existing_replica else []) + active_before = [(worker.worker_id, worker.role) for worker in active] + replicas = {'existing': active} if active else {} + join = Mock(side_effect=AssertionError('unsafe physical allocation reached join')) + executor.elastic_pool = SimpleNamespace( + snapshot=lambda: {**core, **candidates}, replica_snapshot=lambda: replicas, + join_replica=join, + ) + executor._cluster_swap_training_slots = slots + monkeypatch.setattr(executor, '_desired_hybrid_workers', lambda *_: 4 if active else 2) + capture = Mock(side_effect=AssertionError('unsafe physical allocation reached capture')) + monkeypatch.setattr(executor, '_capture_core_snapshot', capture) + + actions = executor._reconfigure_training_pool(4, None, SimpleNamespace()) + + assert actions == ['join_replica_capacity_shortfall:1'] + capture.assert_not_called() + join.assert_not_called() + assert [(worker.worker_id, worker.role) for worker in replicas.get('existing', [])] == active_before + assert not executor._pending_hybrid_joins + + +def test_real_hybrid_join_groups_confirmed_slots_by_host_before_capture(monkeypatch): + from contextlib import nullcontext + from types import SimpleNamespace + from unittest.mock import Mock + + cfg = _config() + cfg.train_gpus, cfg.rollout_gpus = 4, 4 + cfg.train_tp_size, cfg.train_dp_size = 2, 2 + cfg.global_resource_planner.hybrid_worker_launch_enabled = True + executor = RuntimeElasticExecutor(config=cfg, planner=GlobalResourcePlanner.from_config(cfg)) + core = {f'dp{i}': SimpleNamespace(role='core_training') for i in range(2)} + candidates = {f'rollout{i}': SimpleNamespace(role='hybrid_rollout') for i in range(4)} + join = Mock(return_value=SimpleNamespace(generation=1)) + executor.elastic_pool = SimpleNamespace( + snapshot=lambda: {**core, **candidates}, replica_snapshot=lambda: {}, join_replica=join, + activation_guard=lambda *_: nullcontext(), + ) + executor._cluster_swap_training_slots = { + f'rollout{i}': {'host': 'a' if i % 2 == 0 else 'b', 'gpus': [i // 2]} + for i in range(4) + } + monkeypatch.setattr(executor, '_desired_hybrid_workers', lambda *_: 2) + capture = Mock() + monkeypatch.setattr(executor, '_capture_core_snapshot', capture) + + actions = executor._reconfigure_training_pool(4, None, SimpleNamespace()) + + capture.assert_called_once() + assert join.call_args.args[1] == ('rollout0', 'rollout2') + assert any(action.startswith('join_replica_submitted:') for action in actions) + assert not any('shortfall' in action for action in actions) + + def test_runtime_executor_applies_rollout_reconfiguration(): cfg = _config() planner, decision = _decision(cfg) @@ -377,6 +458,7 @@ def test_runtime_executor_does_not_rejoin_hybrids_after_target_is_reached(): cfg.global_resource_planner.runtime_training_pool_target_gpus = 3 cfg.global_resource_planner.runtime_training_pool_plan_only = False pool = ElasticHybridPool( + require_state_ack=False, # Pure scheduling-state simulation; no worker/backend. core_train_workers=["core0"], core_rollout_workers=["rollout0", "rollout1"], zero_sync_steps=0, @@ -545,6 +627,7 @@ def test_runtime_executor_adopts_prewarmed_training_worker_before_pause(tmp_path # Exercise ownership/ordering without leaving a real shell/sleep process. process = Mock(pid=12345) process.poll.return_value = None + process.terminate.side_effect = lambda: setattr(process.poll, "return_value", 0) monkeypatch.setattr( "RL_Framework.infra.elastic.runtime_executor.subprocess.Popen", lambda *args, **kwargs: process, @@ -554,6 +637,11 @@ def test_runtime_executor_adopts_prewarmed_training_worker_before_pause(tmp_path class FakeTrainEngine: current_version = 3 + optimizer_revision = 3 + + def get_elastic_lane_state(self): + return {"data_parallel_rank": 0, "tensor_parallel_rank": 0, + "pipeline_parallel_rank": 0, "context_parallel_rank": 0} def __init__(self): self.snapshots = [] @@ -567,9 +655,7 @@ def capture_elastic_state_snapshot(self, worker_id, target_core_id): return self.current_version def get_elastic_state_snapshot_path(self, version): - path = tmp_path / f"rank0_v{version}.pt" - path.write_bytes(b"snapshot") - return path + return f"ehp://fake-state:12345/revision-{version}" def set_elastic_gradient_domain(self, domain): self.domain = domain @@ -588,6 +674,18 @@ def __init__(self, events, **kwargs): super().__init__(**kwargs) self.events = events + def _capture_core_snapshot(self): + # This test isolates prewarm ordering; simulate the collective + # capture of every configured core DP lane without real workers. + import json + path = tmp_path / "capture.json" + path.write_text(json.dumps({"format": "libra-ehp-lanes-v1", "optimizer_revision": 3, + "snapshot_id": "simulated-collective", + "lanes": {f"dp{i}": {"0": f"ehp://fake/dp{i}"} + for i in range(self.config.train_dp_size)}})) + self._snapshot_manifests[3] = str(path) + return str(path) + def _ensure_gradient_server(self): class Endpoint: host = "127.0.0.1" @@ -623,6 +721,7 @@ def _is_hybrid_worker_running(self, worker_id): return worker_id in self._hybrid_worker_meta cfg = _config() + cfg.global_resource_planner.hybrid_worker_task_dir = str(tmp_path / "tasks") cfg.global_resource_planner.runtime_reconfigure_training = True cfg.global_resource_planner.runtime_training_pool_plan_only = False cfg.global_resource_planner.runtime_training_pool_target_gpus = 4 @@ -639,6 +738,12 @@ def _is_hybrid_worker_running(self, worker_id): train_engine=train_engine, dispatcher=OrderedDispatcher(events), ) + # The fake prewarm process does not execute a physical rollout stop. Model + # the confirmed loan explicitly instead of relying on an unowned fallback. + monkeypatch.setattr( + 'RL_Framework.infra.elastic.rollout_lending.lend_rollout', + lambda *_: {'rollout0': {'host': '', 'gpus': [0]}}, + ) result = executor.execute(decision) @@ -796,6 +901,7 @@ def _reconfigure_training_pool(self, current_train_gpus, plan, result): decision.candidate_plan.train_config.dp = 2 decision.candidate_plan.rollout_config.tp_list = [1, 1] pool = ElasticHybridPool( + require_state_ack=False, # Pure scheduling-state simulation; no worker/backend. core_train_workers=["core0", "core1"], core_rollout_workers=["rollout0", "rollout1"], zero_sync_steps=0, @@ -1157,6 +1263,7 @@ def _run_external_stop(self, meta): self.events.append(f"external_stop:{meta.instance_id}") cfg = _config() + cfg.global_resource_planner.vllm_stop_command_template = "test-owned-stop {instance_id}" cfg.global_resource_planner.runtime_manage_rollout_processes = True cfg.global_resource_planner.runtime_rollout_reconfigure_strategy = "prewarm" cfg.global_resource_planner.runtime_wait_rollout_process_start_ready = True @@ -1238,6 +1345,7 @@ def _run_external_stop(self, meta): self.events.append(f"external_stop:{meta.instance_id}") cfg = _config() + cfg.global_resource_planner.vllm_stop_command_template = "test-owned-stop {instance_id}" cfg.global_resource_planner.runtime_manage_rollout_processes = True cfg.global_resource_planner.runtime_rollout_reconfigure_strategy = "prewarm" cfg.global_resource_planner.runtime_prewarm_no_spare_fallback_strategy = "restart_all" @@ -1445,3 +1553,105 @@ def get_queue_sizes(self): assert stats.enqueued == 0 assert stats.running == 0 assert stats.accepted == 0 + + +@pytest.mark.parametrize('mode, replica_gpus', [ + ('megatron_core', 1), + ('megatron_core', 2), + ('transport', 2), +]) +def test_hybrid_worker_command_launches_mcore_in_own_process_group(mode, replica_gpus): + cfg = _config() + cfg.global_resource_planner.hybrid_worker_mode = mode + cfg.global_resource_planner.hybrid_worker_python = '/gpu env/bin/python' + executor = RuntimeElasticExecutor(config=cfg, planner=GlobalResourcePlanner.from_config(cfg)) + executor.gradient_server = SimpleNamespace( + endpoint=SimpleNamespace(host='127.0.0.1', port=12345, authkey='auth'), + close=lambda: None, + ) + try: + argv = shlex.split(executor._build_hybrid_worker_command( + worker_id='hybrid0', target_core_id='dp0', snapshot_path='/snapshot', + gpus=list(range(2, 2 + replica_gpus)), replica_id='hybrid-replica', + replica_world_size=replica_gpus, + )) + finally: + executor.close() + + assert argv[0] == 'CUDA_VISIBLE_DEVICES=' + ','.join(str(gpu) for gpu in range(2, 2 + replica_gpus)) + assert argv[1] == '/gpu env/bin/python' + if mode == 'megatron_core': + assert argv[2:6] == [ + '-m', 'torch.distributed.run', '--standalone', f'--nproc_per_node={replica_gpus}', + ] + script_index = 6 + else: + script_index = 2 + assert argv[script_index].endswith('/scripts/elastic_hybrid_worker.py') + assert argv[argv.index('--worker-mode') + 1] == mode + assert argv[argv.index('--replica-world-size') + 1] == str(replica_gpus) + assert not any(arg.startswith(('RANK=', 'WORLD_SIZE=', 'MASTER_ADDR=', 'MASTER_PORT=')) for arg in argv) + + + +def test_hybrid_worker_command_reads_published_lane_endpoint(tmp_path, monkeypatch): + from RL_Framework.infra.elastic.gradient_ipc import GradientEndpoint + from RL_Framework.scripts import elastic_hybrid_worker as worker + from RL_Framework.trainer import async_rl_trainer as trainer_module + + cfg = _config() + task_root = tmp_path / 'worker task space' + published_root = task_root + cfg.global_resource_planner.elastic_hybrid_planning_enabled = True + cfg.global_resource_planner.hybrid_worker_task_dir = str(task_root) + lane = dict(data_parallel_rank=1, tensor_parallel_rank=0, + pipeline_parallel_rank=0, context_parallel_rank=0) + engine = SimpleNamespace( + configure_elastic_training=lambda *args, **kwargs: None, + get_elastic_core_replica_ids=lambda: ['dp1'], + get_elastic_lane_state=lambda: lane, + enqueue_hybrid_gradient_payload=lambda payload: None, + ) + endpoint = GradientEndpoint('core-host', 12345, 'lane-auth') + monkeypatch.setattr(trainer_module, 'ElasticGradientServer', + lambda **kwargs: SimpleNamespace(endpoint=endpoint, start=lambda: endpoint, close=lambda: None)) + executor = RuntimeElasticExecutor(config=cfg, planner=GlobalResourcePlanner.from_config(cfg)) + producer = object.__new__(trainer_module.AsyncRLTrainer) + producer.config, producer.train_engine = cfg, engine + producer.rank, producer.is_main_process = 3, True + producer.runtime_elastic_executor = executor + producer._elastic_membership_epochs = {} + try: + producer._init_nonblocking_elastic_core() + receipt = published_root / 'core_endpoints' / 'rank_3.json' + assert receipt.is_file() + argv = shlex.split(executor._build_hybrid_worker_command( + worker_id='hybrid0', target_core_id='dp1', snapshot_path='/snapshot', gpus=[2], + )) + script_index = next(i for i, arg in enumerate(argv) if arg.endswith('/elastic_hybrid_worker.py')) + monkeypatch.setattr(sys, 'argv', argv[script_index:]) + args = worker.parse_args() + resolved = worker._endpoint_for_lane(args, engine) + assert resolved == json.loads(receipt.read_text()) + assert (resolved['port'], resolved['authkey']) == (12345, 'lane-auth') + assert Path(args.gradient_endpoint_dir) == receipt.parent + receipt.unlink() + with pytest.raises(RuntimeError, match='no core gradient endpoint matches'): + worker._endpoint_for_lane(args, engine) + finally: + executor.close() + + +@pytest.mark.parametrize("explicit_coord_id", [None, "explicit-transaction"]) +def test_peer_rank_state_keeps_the_selected_transaction_id(tmp_path, explicit_coord_id): + cfg = _config() + cfg.global_resource_planner.runtime_reconfiguration_coord_dir = str(tmp_path) + planner, decision = _decision(cfg) + executor = RuntimeElasticExecutor(config=cfg, planner=planner) + executor._active_runtime_coord_id = "active-transaction" + executor._lent_rollout = {"engines": [SimpleNamespace(instance_id="loan-a")]} + path = executor._write_peer_rank_reconfiguration_state( + decision.candidate_plan, phase="request", coord_id=explicit_coord_id) + record = json.loads(path.read_text()) + assert record["coord_id"] == (explicit_coord_id or "active-transaction") + assert record["unavailable_rollout_ids"] == ["loan-a"] diff --git a/tests/test_runtime_executor_shutdown.py b/tests/test_runtime_executor_shutdown.py index 132d6f1..0c9fc0d 100644 --- a/tests/test_runtime_executor_shutdown.py +++ b/tests/test_runtime_executor_shutdown.py @@ -66,7 +66,21 @@ def make(workers=3, background=3): monkeypatch.setattr(executor, "_build_hybrid_worker_command", lambda **kw: "fake worker") monkeypatch.setattr(executor, "_cluster_swap_training_slot", lambda wid: {"host": "fake-host", "gpus": [0]}) # Most tests isolate launch ordering. A separate test keeps the real poll. - monkeypatch.setattr(executor, "_wait_hybrid_worker_ready", lambda wid: None) + # Launch-order tests simulate a real installation and core-boundary ACK; + # dedicated readiness tests exercise actual files/identity/TP coverage. + receipt = {"acknowledged": True, "optimizer_revision": 0, "state_version": 0, + "snapshot_id": "simulated-state"} + manifest = tmp_path / "capture.json" + manifest.write_text(json.dumps({"format": "libra-ehp-lanes-v1", "optimizer_revision": 0, + "snapshot_id": "simulated-state", "lanes": {"dp0": {"0": "ehp://fake/lane0"}}})) + monkeypatch.setattr(executor, "_snapshot_path_for_version", lambda version: str(manifest)) + monkeypatch.setattr(executor, "_capture_core_snapshot", lambda: str(manifest)) + monkeypatch.setattr(executor, "_wait_hybrid_worker_ready", lambda wid, **kw: dict(receipt)) + monkeypatch.setattr(executor, "_queue_activation", lambda replica, ack: ack) + from RL_Framework.infra.elastic import rollout_lending + monkeypatch.setattr(rollout_lending, "wait_hybrid_stopped", lambda executor, wid, request_id=None: { + "worker_id": wid, "request_id": request_id, "exited": True, + }) h = SimpleNamespace(executor=executor, pool=pool, processes=processes, popen=popen, path=tmp_path) instances.append(h) return h @@ -252,7 +266,7 @@ def prepare(**kwargs): try: assert entered.wait(3) old.cancel() - newer = join(h, replica=replica, callback=lambda *args: None) + newer = join(h, replica=replica, callback=lambda *args: {"acknowledged": True, "state_version": 0, "optimizer_revision": 0}) newer.result(timeout=3) finally: resume.set() @@ -305,9 +319,9 @@ def test_readiness_poll_is_cancelled_without_waiting_for_timeout(harness, monkey entered = threading.Event() h.executor.config.global_resource_planner.hybrid_worker_ready_timeout_s = 3600 - def wait(worker): + def wait(worker, **kwargs): entered.set() - RuntimeElasticExecutor._wait_hybrid_worker_ready(h.executor, worker) + return RuntimeElasticExecutor._wait_hybrid_worker_ready(h.executor, worker, **kwargs) monkeypatch.setattr(h.executor, "_wait_hybrid_worker_ready", wait) handle = join(h) @@ -401,7 +415,7 @@ def test_close_failure_keeps_ownership_and_cleans_other_workers(harness): def test_prewarm_and_legacy_launch_are_fenced(harness): h = harness() - handle = join(h, callback=lambda *args: None) + handle = join(h, callback=lambda *args: {"acknowledged": True, "state_version": 0, "optimizer_revision": 0}) handle.result(timeout=3) h.executor.close() with pytest.raises(JoinCancelledError): @@ -417,9 +431,10 @@ def test_restart_committed_before_close_is_also_reaped(harness, monkeypatch): ) entered, resume = threading.Event(), threading.Event() - def wait(worker): + def wait(worker, **kwargs): entered.set() assert resume.wait(3) + return {"acknowledged": True, "state_version": 0, "optimizer_revision": 0} monkeypatch.setattr(h.executor, "_wait_hybrid_worker_ready", wait) handle = join(h) @@ -436,10 +451,10 @@ def wait(worker): def test_legacy_completed_join_generation_is_checked_at_launch(harness): h = harness() - old = join(h, callback=lambda *args: None) + old = join(h, callback=lambda *args: {"acknowledged": True, "state_version": 0, "optimizer_revision": 0}) old.result(timeout=3) old.cancel() - join(h, callback=lambda *args: None).result(timeout=3) + join(h, callback=lambda *args: {"acknowledged": True, "state_version": 0, "optimizer_revision": 0}).result(timeout=3) assert h.executor._launch_hybrid_worker_after_join(old) is None assert not h.processes diff --git a/tests/test_runtime_failed_snapshot_cleanup.py b/tests/test_runtime_failed_snapshot_cleanup.py new file mode 100644 index 0000000..9d20b32 --- /dev/null +++ b/tests/test_runtime_failed_snapshot_cleanup.py @@ -0,0 +1,248 @@ +"""Snapshot ownership survives join failure until a common Core boundary.""" +from concurrent.futures import Future +from types import SimpleNamespace as NS + +import pytest +import torch + +from RL_Framework.config import AsyncRLConfig +from RL_Framework.infra.elastic.core_reconfiguration import execute_core_command +from RL_Framework.infra.elastic.runtime_executor import RuntimeElasticExecutor +from RL_Framework.infra.elastic.state_transfer import LiveStateServer + + +class SnapshotEngine: + rank = 0 + world_size = 1 + optimizer_revision = 1 + + def __init__(self): + self.server = LiveStateServer(host='127.0.0.1', advertised_host='127.0.0.1') + self.urls, self.released = {}, [] + + def capture_elastic_state_snapshot(self, *args): + revision = self.optimizer_revision + if revision not in self.urls: + self.urls[revision] = self.server.publish({'model': torch.zeros(32)}) + return revision + + def get_elastic_state_snapshot_path(self, revision): + return self.urls[revision] + + def get_elastic_lane_state(self): + return dict(data_parallel_rank=0, tensor_parallel_rank=0, + context_parallel_rank=0, pipeline_parallel_rank=0) + + def release_elastic_state_snapshot(self, revision): + uri = self.urls.pop(revision, None) + if uri: + self.server.release(uri) + self.released.append(revision) + + +@pytest.fixture +def runtime(tmp_path): + config = AsyncRLConfig(model_path='/not-loaded', train_gpus=1, rollout_gpus=1, + n_total_gpus=2, train_tp_size=1, train_dp_size=1, log_dir=str(tmp_path)) + engine = SnapshotEngine() + executor = RuntimeElasticExecutor(config=config, planner=NS(), train_engine=engine) + executor._active_runtime_coord_id = 'snapshot-test' + yield executor, engine + engine.server.close() + + +def failed_join(): + future = Future() + future.set_exception(TimeoutError('initial installation failed')) + return NS(member_worker_ids=('rollout0',), done=future.done, result=future.result) + + +def boundary(executor, monkeypatch, *, execute=True): + prepared = NS(coord_id='common-boundary', should_execute=execute, candidate_plan=None, + decision=None, result=NS(applied=True), effective_planner_signal=None) + monkeypatch.setattr(executor, 'accept_planner_signal', lambda *args, **kwargs: None) + return executor.execute_prepared(prepared) + + +def fail_snapshot(executor): + path = executor._capture_core_snapshot() + executor._pending_hybrid_joins['replica'] = failed_join() + executor._reconcile_pending_hybrid_joins() + return path + + +def test_failed_join_releases_only_at_next_common_boundary(runtime, monkeypatch): + executor, engine = runtime + path = fail_snapshot(executor) + assert engine.released == [] # metrics/background reconciliation is not collective + boundary(executor, monkeypatch) + assert engine.released == [1] + assert not executor._snapshot_manifests + assert not engine.server._snapshots + from pathlib import Path + assert not Path(path).exists() + + +@pytest.mark.parametrize('busy', ['join', 'activation', 'cancelled_loader', 'stop']) +def test_release_waits_for_every_consumer(runtime, monkeypatch, busy): + executor, engine = runtime + fail_snapshot(executor) + if busy == 'join': + executor._pending_hybrid_joins['other'] = NS(done=lambda: False) + elif busy == 'activation': + executor._pending_activations['other'] = {} + elif busy == 'cancelled_loader': + executor._retired_hybrid_joins.append(NS(done=lambda: False)) + else: + executor._snapshot_stop_pending.add('rollout0') + monkeypatch.setattr(executor, '_stop_hybrid_worker_process', + lambda *args: (_ for _ in ()).throw(RuntimeError('exit not acknowledged'))) + if busy == 'stop': + with pytest.raises(RuntimeError, match='exit not acknowledged'): + boundary(executor, monkeypatch) + else: + boundary(executor, monkeypatch) + assert engine.released == [] + assert executor._snapshot_manifests + + +def test_no_transaction_cannot_issue_release(runtime, monkeypatch): + executor, engine = runtime + fail_snapshot(executor) + boundary(executor, monkeypatch, execute=False) + assert engine.released == [] + + +def test_collective_release_failure_retains_manifest_for_retry(runtime, monkeypatch): + executor, engine = runtime + path = fail_snapshot(executor) + import RL_Framework.infra.elastic.core_reconfiguration as commands + original = commands.dispatch_core_command + monkeypatch.setattr(commands, 'dispatch_core_command', + lambda *args, **kwargs: (_ for _ in ()).throw(RuntimeError('rank1 receipt missing'))) + with pytest.raises(RuntimeError, match='rank1 receipt missing'): + boundary(executor, monkeypatch) + assert executor._snapshot_manifests == {1: path} + monkeypatch.setattr(commands, 'dispatch_core_command', original) + boundary(executor, monkeypatch) + assert not executor._snapshot_manifests + + +def test_repeated_failed_revisions_do_not_accumulate(runtime, monkeypatch): + executor, engine = runtime + for revision in (1, 2, 3): + engine.optimizer_revision = revision + fail_snapshot(executor) + boundary(executor, monkeypatch) + assert not engine.server._snapshots + assert engine.released == [1, 2, 3] + + +def test_every_core_rank_releases_before_boundary_returns(runtime, monkeypatch): + import threading + from RL_Framework.infra.elastic.runtime_boundary import read_record + executor, engine = runtime + fail_snapshot(executor) + peer = SnapshotEngine() + peer.rank = 1 + peer.capture_elastic_state_snapshot() + engine.world_size = 2 + errors = [] + finished = threading.Event() + + def follow(): + try: + while not finished.wait(.005): + command = read_record(executor._coordination_dir() / 'core_command.json') + if command.get('operation') == 'release': + execute_core_command(peer, command, executor._coordination_dir()) + return + except Exception as exc: + errors.append(exc) + + thread = threading.Thread(target=follow) + thread.start() + try: + boundary(executor, monkeypatch) + assert not errors + assert engine.released == peer.released == [1] + assert not executor._snapshot_manifests + finally: + finished.set() + thread.join(timeout=2) + peer.server.close() + + +def test_pending_cleanup_requests_membership_boundary_without_planner_resize(runtime): + executor, _engine = runtime + fail_snapshot(executor) + assert executor.has_pending_boundary_work() + current = NS() + decision = NS(should_reconfigure=False, current_plan=current, candidate_plan=None, + resource_should_reconfigure=False, reason='hold', elastic_hybrid_signal=None) + prepared = executor.prepare_transaction(decision) + assert prepared.should_execute + assert not prepared.resource_reconfiguration + assert prepared.candidate_plan is current + assert prepared.coord_id + + +def test_cancelled_loader_remains_owned_until_future_finishes(runtime, monkeypatch): + import threading + executor, engine = runtime + executor._capture_core_snapshot() + future = Future() + cancelled = [] + handle = NS(member_worker_ids=('rollout0',), done=future.done, + cancel=lambda: cancelled.append(True)) + executor._pending_hybrid_joins['replica'] = handle + event = threading.Event() + request = {'event': event} + executor._pending_activations['replica'] = request + executor._release_all_hybrid_workers(reason='pressure') + assert cancelled and event.is_set() and request['error'] + assert not executor._pending_hybrid_joins + boundary(executor, monkeypatch) + assert not engine.released + future.set_exception(RuntimeError('cancelled generation')) + boundary(executor, monkeypatch) + assert engine.released == [1] + + +def test_partial_cancel_stop_failure_tracks_all_removed_loaders(runtime, monkeypatch): + executor, engine = runtime + executor._capture_core_snapshot() + futures = [Future(), Future()] + cancelled = [] + for index, future in enumerate(futures): + handle = NS(member_worker_ids=(f'rollout{index}',), done=future.done, + cancel=lambda index=index: cancelled.append(index)) + executor._pending_hybrid_joins[str(index)] = handle + monkeypatch.setattr(executor, '_stop_hybrid_worker_process', + lambda *args: (_ for _ in ()).throw(RuntimeError('stop failed'))) + with pytest.raises(RuntimeError, match='stop failed'): + executor._release_all_hybrid_workers(reason='pressure') + assert cancelled == [0, 1] + assert len(executor._retired_hybrid_joins) == 2 + assert not engine.released + assert executor._snapshot_stop_pending == {'rollout0'} + + +def test_cold_loader_does_not_request_repeated_cleanup_boundaries(runtime): + executor, _engine = runtime + executor._capture_core_snapshot() + future = Future() + executor._pending_hybrid_joins['replica'] = NS(done=future.done) + assert not executor.has_pending_boundary_work() + future.set_exception(TimeoutError('cold load failed')) + assert executor.has_pending_boundary_work() + + +def test_cancelled_fetch_requests_cleanup_only_after_it_stops(runtime): + executor, _engine = runtime + executor._capture_core_snapshot() + future = Future() + executor._retired_hybrid_joins.append(NS(done=future.done)) + assert not executor.has_pending_boundary_work() + future.set_result(None) + assert executor.has_pending_boundary_work() diff --git a/tests/test_runtime_repartition_eligibility.py b/tests/test_runtime_repartition_eligibility.py new file mode 100644 index 0000000..fce5c70 --- /dev/null +++ b/tests/test_runtime_repartition_eligibility.py @@ -0,0 +1,64 @@ +"""Candidate TP layouts are checked before mutating runtime ownership.""" + +from copy import deepcopy + +import pytest + +from RL_Framework.config import AsyncRLConfig, HeterogeneousInstanceConfig +from RL_Framework.infra.cost_model.global_resource_planner import ( + GlobalResourcePlan, + GlobalResourcePlanner, + PlannerDecision, +) +from RL_Framework.infra.cost_model.model import RolloutClusterConfig, TrainParallelConfig +from RL_Framework.infra.elastic.runtime_executor import RuntimeElasticExecutor + + +@pytest.mark.parametrize('initial,target,should_execute', [ + ([1, 1, 1], [1, 2], True), + ([1, 2], [1, 1, 1], True), + ([1, 1, 1], [1, 1, 1], False), + ([1, 2], [1, 2], False), +]) +def test_prepare_recognizes_candidate_layout_without_changing_live_resources( + monkeypatch, initial, target, should_execute): + config = AsyncRLConfig(model_path='/not-loaded', train_gpus=1, rollout_gpus=3, + train_tp_size=1, train_dp_size=1, n_total_gpus=4) + config.global_resource_planner.runtime_manage_rollout_processes = True + config.heterogeneous_rollout.instances = [ + HeterogeneousInstanceConfig(instance_id=f'seed{index}', host='shared', + tp=1, gpus=[index + 1], port=8000 + index) for index in range(3) + ] + planner = GlobalResourcePlanner.from_config(config) + + def plan(degrees): + return GlobalResourcePlan(train_config=TrainParallelConfig(tp=1, dp=1, b_micro=1), + rollout_config=RolloutClusterConfig(tp_list=degrees), + t_train=1., t_rollout=1., t_global=1., n_total_gpus=4, + max_concurrent_rollouts=config.max_concurrent_rollouts) + + planner.apply_plan_to_config(plan(initial), config) + decision = PlannerDecision(step=2, should_reconfigure=True, reason='test_repartition', + current_plan=plan(initial), candidate_plan=plan(target), resource_should_reconfigure=True) + executor = RuntimeElasticExecutor(config=config, planner=planner) + executor._process_meta = {meta.instance_id: meta + for meta in executor._desired_rollout_processes(decision.current_plan)} + before_config = deepcopy(config.to_dict()) + before_owners = deepcopy(executor._process_meta) + + def forbidden(*args, **kwargs): + pytest.fail('Eligibility must not launch, stop, wait for or reconfigure a process') + + for name in ('_start_one_rollout_process', '_stop_rollout_process', + '_reconfigure_rollout_engine', '_wait_rollout_process_start_ready'): + monkeypatch.setattr(executor, name, forbidden) + prepared = executor.prepare_transaction(decision) + + assert prepared.should_execute is should_execute + assert prepared.resource_reconfiguration is True + assert prepared.candidate_plan is decision.candidate_plan + assert prepared.result.applied is False and not prepared.result.actions + assert config.to_dict() == before_config + assert executor._process_meta == before_owners and not executor._processes + if not should_execute: + assert prepared.reason == 'runtime_plan_already_applied' diff --git a/tests/test_runtime_rollout_ownership.py b/tests/test_runtime_rollout_ownership.py new file mode 100644 index 0000000..cddcc23 --- /dev/null +++ b/tests/test_runtime_rollout_ownership.py @@ -0,0 +1,59 @@ +"""Executor never lends a rollout GPU before its owned process has exited.""" +from types import SimpleNamespace as NS + +import pytest + +from RL_Framework.infra.elastic import rollout_process +from RL_Framework.infra.elastic.runtime_executor import ManagedRolloutProcess, RuntimeElasticExecutor + + +def executor(tmp_path): + value = object.__new__(RuntimeElasticExecutor) + value.config = NS(rollout_weight_sync_control_dir=str(tmp_path), + global_resource_planner=NS(vllm_stop_timeout_s=2., vllm_stop_command_template='')) + value._processes, value._process_meta = {}, {} + value._rollout_launch_env = lambda meta: {'SLURM_JOB_ID': '123'} + return value + + +def test_start_uses_node_placement_and_stop_retains_ownership_on_failure(tmp_path, monkeypatch): + value = executor(tmp_path) + meta = ManagedRolloutProcess('rollout', 'service-command', 0, [2, 3], 8000, host='node2', tp=2) + proc = NS(pid=7) + launches = [] + monkeypatch.setattr(rollout_process, 'launch_rollout_process', + lambda meta, env, **kw: launches.append((meta.host, env, kw)) or proc) + started = value._start_one_rollout_process(meta) + assert launches[0][0] == 'node2' and launches[0][2]['control_dir'] == str(tmp_path) + + def fail(*args, **kwargs): + raise TimeoutError('service has not exited') + monkeypatch.setattr(rollout_process, 'stop_rollout_process', fail) + with pytest.raises(TimeoutError): + value._stop_rollout_process('rollout') + assert value._processes['rollout'] is proc and value._process_meta['rollout'] is started + monkeypatch.setattr(rollout_process, 'stop_rollout_process', lambda *args, **kwargs: None) + assert value._stop_rollout_process('rollout') == [started] + assert value._process_meta == value._processes == {} + assert value._stop_rollout_process('rollout') == [] + + +def test_adopted_stop_passes_remote_identity_without_local_signal(tmp_path, monkeypatch): + value = executor(tmp_path) + meta = ManagedRolloutProcess('initial', 'service', 321, [0, 1], 8001, + host='node2', tp=2, adopted=True) + value._process_meta['initial'] = meta + stopped = [] + monkeypatch.setattr(rollout_process, 'stop_adopted_rollout', + lambda meta, **kw: stopped.append((meta, kw))) + assert value._stop_rollout_process('initial') == [meta] + assert stopped[0][0].host == 'node2' and stopped[0][1]['env']['SLURM_JOB_ID'] == '123' + + +def test_unknown_process_cannot_be_reported_stopped(tmp_path): + value = executor(tmp_path) + meta = ManagedRolloutProcess('unknown', 'service', -1, [0, 1], 8000, adopted=True) + value._process_meta['unknown'] = meta + with pytest.raises(RuntimeError, match='No process ownership'): + value._stop_rollout_process('unknown') + assert value._process_meta['unknown'] is meta diff --git a/trainer/async_rl_trainer.py b/trainer/async_rl_trainer.py index 351bd11..9ae41f8 100644 --- a/trainer/async_rl_trainer.py +++ b/trainer/async_rl_trainer.py @@ -33,7 +33,7 @@ from RL_Framework.engine.heterogeneous_engine import HeterogeneousRolloutEngine from RL_Framework.engine.train_engine import TrainEngine from RL_Framework.engine.train_factory import create_train_engine -from RL_Framework.infra.cost_model.global_resource_planner import GlobalResourcePlanner +from RL_Framework.infra.cost_model.global_resource_planner import CoreRuntimeCapabilities, GlobalResourcePlanner from RL_Framework.infra.cost_model.startup_profile import load_length_profile_records from RL_Framework.infra.elastic.runtime_executor import RuntimeElasticExecutor from RL_Framework.infra.elastic.runtime_boundary import RuntimeBoundaryProtocol, read_record @@ -271,11 +271,11 @@ def setup(self, workflow): self._init_nonblocking_elastic_core() - if self.is_main_process and self.train_engine.is_batch_source(): + if self.is_main_process and self._is_rollout_source(): print("Connecting to vLLM inference services...") - if self.train_engine.is_batch_source(): + if self._is_rollout_source(): hetero_cfg = getattr(self.config, "heterogeneous_rollout", None) if hetero_cfg and hetero_cfg.enabled and hetero_cfg.instances: self._use_heterogeneous = True @@ -284,7 +284,7 @@ def setup(self, workflow): self._use_heterogeneous = False self._setup_homogeneous_engine() - if self.is_main_process and self.train_engine.is_batch_source(): + if self.is_main_process and self._is_rollout_source(): self.rollout_engine.wait_for_ready(timeout=300) n_inst = self.rollout_engine.num_instances print(f"vLLM inference services are ready ({n_inst} instances)") @@ -303,13 +303,11 @@ def setup(self, workflow): ) - if self.train_engine.is_batch_source(): + if self._is_rollout_source(): self.staleness_manager = StalenessManager( version_provider=self.train_engine, max_concurrent_rollouts=self.config.max_concurrent_rollouts, - consumer_batch_size=self.train_engine.get_local_batch_size( - self.config.batch_size - ), + consumer_batch_size=max(1, self._source_batch_sample_count()), max_staleness=self.config.max_head_offpolicyness, ) @@ -846,6 +844,9 @@ def train(self, workflow, dataset=None): self._trace_train_phase(step, "follow_reconfig_start") self._runtime_reconfiguration_boundary(step, "pre") self._trace_train_phase(step, "follow_reconfig_done") + if getattr(self.config.global_resource_planner, "runtime_reconfigure_training", False): + from RL_Framework.infra.elastic.rollout_lifecycle import ensure_batch_supply + ensure_batch_supply(self) active_hybrid_ids = self._prepare_elastic_training_step(step) # Reconfiguration can change the number of DP data sources. local_batch_size = self._source_batch_sample_count() @@ -928,75 +929,9 @@ def train(self, workflow, dataset=None): ) - weight_sync_time = 0.0 - rollout_sync_drain = None + weight_sync_time, rollout_sync_drain = 0.0, None if self.config.sync_interval > 0 and (step + 1) % self.config.sync_interval == 0: - self._trace_train_phase(step, "weight_sync_start") - if self.dispatcher is not None: - drain_started = time.time() - self.dispatcher.pause() - cancelled = self.dispatcher.cancel_queued() - metrics = self.dispatcher.get_runtime_metrics() - active_to_drain = int(metrics.get("staleness_running", 0)) - self.dispatcher.wait_until_idle( - timeout=float( - getattr( - self.config, - "rollout_weight_sync_timeout_s", - 1200.0, - ) - ) - ) - rollout_sync_drain = { - "cancelled_dispatcher_queue": int(cancelled["dispatcher_queue"]), - "cancelled_runner_queue": int(cancelled["runner_queue"]), - "active_to_drain": active_to_drain, - "drain_seconds": time.time() - drain_started, - } - if self.is_main_process: - print( - "[RolloutDrain] " - f"step={step} active={active_to_drain} " - f"cancelled_dispatcher={cancelled['dispatcher_queue']} " - f"cancelled_runner={cancelled['runner_queue']} " - f"seconds={rollout_sync_drain['drain_seconds']:.3f}", - flush=True, - ) - - self._wait_for_all_ranks_before_weight_sync(step) - sync_start = time.time() - skip_rollout_sync = ( - str(getattr(self.config, "rollout_weight_sync_mode", "")).lower() - == "none" - and not bool( - getattr(self.config, "require_rollout_weight_sync", True) - ) - ) - if skip_rollout_sync: - self._record_weight_sync_phase("sync_skipped") - else: - self._sync_weights(target_version=step + 1) - weight_sync_time = time.time() - sync_start - - new_version = step + 1 - self.train_engine.set_version(new_version) - self._rebind_rollout_engine_from_config(reason="after_weight_sync") - - if self.dispatcher is not None: - self.dispatcher.resume() - self._trace_train_phase(step, "weight_sync_done") - - if self.is_main_process and self.config.enable_rollout_tracing: - sm_stats = ( - self.staleness_manager.get_stats() - if self.staleness_manager is not None - else None - ) - print( - f"Weight synchronization complete: version={new_version}, " - f"running={getattr(sm_stats, 'running', 0)}, " - f"accepted={getattr(sm_stats, 'accepted', 0)}" - ) + weight_sync_time, rollout_sync_drain = self._periodic_weight_sync(step) step_total_time = time.time() - step_start @@ -1079,8 +1014,82 @@ def train(self, workflow, dataset=None): print(f"Training complete. Total time: {total_time:.1f}s") print("=" * 60) + if self.is_main_process and self.runtime_elastic_executor is not None: + self.runtime_elastic_executor.wait_for_hybrid_training_tasks() self._cleanup() + def _periodic_weight_sync(self, step): + """Global publication boundary; specialized publication may overlap it.""" + rollout_sync_drain = None + self._trace_train_phase(step, "weight_sync_start") + if self.dispatcher is not None: + drain_started = time.time() + self.dispatcher.pause() + cancelled = self.dispatcher.cancel_queued() + metrics = self.dispatcher.get_runtime_metrics() + active_to_drain = int(metrics.get("staleness_running", 0)) + self.dispatcher.wait_until_idle( + timeout=float( + getattr( + self.config, + "rollout_weight_sync_timeout_s", + 1200.0, + ) + ) + ) + rollout_sync_drain = { + "cancelled_dispatcher_queue": int(cancelled["dispatcher_queue"]), + "cancelled_runner_queue": int(cancelled["runner_queue"]), + "active_to_drain": active_to_drain, + "drain_seconds": time.time() - drain_started, + } + if self.is_main_process: + print( + "[RolloutDrain] " + f"step={step} active={active_to_drain} " + f"cancelled_dispatcher={cancelled['dispatcher_queue']} " + f"cancelled_runner={cancelled['runner_queue']} " + f"seconds={rollout_sync_drain['drain_seconds']:.3f}", + flush=True, + ) + + self._wait_for_all_ranks_before_weight_sync(step) + sync_start = time.time() + skip_rollout_sync = ( + str(getattr(self.config, "rollout_weight_sync_mode", "")).lower() + == "none" + and not bool( + getattr(self.config, "require_rollout_weight_sync", True) + ) + ) + if skip_rollout_sync: + self._record_weight_sync_phase("sync_skipped") + else: + self._sync_weights(target_version=step + 1) + weight_sync_time = time.time() - sync_start + + new_version = step + 1 + self.train_engine.set_version(new_version) + if not getattr(self, "_rollout_unavailable_ids", ()): + self._rebind_rollout_engine_from_config(reason="after_weight_sync") + + self._resume_rollout_pipeline() + self._trace_train_phase(step, "weight_sync_done") + + if self.is_main_process and self.config.enable_rollout_tracing: + sm_stats = ( + self.staleness_manager.get_stats() + if self.staleness_manager is not None + else None + ) + print( + f"Weight synchronization complete: version={new_version}, " + f"running={getattr(sm_stats, 'running', 0)}, " + f"accepted={getattr(sm_stats, 'accepted', 0)}" + ) + + return weight_sync_time, rollout_sync_drain + def _sync_aware_rollout_limit( self, *, @@ -1219,10 +1228,19 @@ def _run_periodic_evaluation(self, workflow, dataset, step: int): if self.dispatcher is not None: self.dispatcher.resume() + def _restore_lent_rollouts_for_publication(self): + from RL_Framework.infra.elastic.rollout_lifecycle import restore_for_publication + return restore_for_publication(self) + + def _reopen_restored_rollouts(self): + from RL_Framework.infra.elastic.rollout_lifecycle import reopen_after_publication + reopen_after_publication(self) + def _sync_weights(self, target_version: int): """Sync weights.""" if self.config.weight_sync_mode == "nccl": self._sync_weights_nccl(target_version) + self._reopen_restored_rollouts() return if self.config.weight_sync_mode == "disk": if self.is_main_process: @@ -1253,13 +1271,14 @@ def _nccl_rollout_topology(self) -> tuple[int, dict[str, int]]: # Batch-source ranks own a live rollout client, while other training # ranks intentionally keep ``self.rollout_engine`` unset. NCCL weight # synchronization still runs on every training rank, so derive the - # topology from the planner's config when no local client exists. + # topology from the committed planner config. A drained old client may + # still describe the previous engine layout during publication. rollout_engine = self.rollout_engine hetero_cfg = getattr(self.config, "heterogeneous_rollout", None) configured_instances = getattr(hetero_cfg, "instances", []) if hetero_cfg else [] if self._use_heterogeneous or configured_instances: - instance_configs = getattr(rollout_engine, "instance_configs", None) - if not instance_configs: + instance_configs = None + if configured_instances: instance_configs = [ { "instance_id": str(getattr(cfg, "instance_id", f"instance_{index}")), @@ -1267,6 +1286,8 @@ def _nccl_rollout_topology(self) -> tuple[int, dict[str, int]]: } for index, cfg in enumerate(configured_instances) ] + else: + instance_configs = getattr(rollout_engine, "instance_configs", []) entries = [ (str(cfg["instance_id"]), int(cfg.get("tp_degree", 1))) for cfg in instance_configs @@ -1284,7 +1305,15 @@ def _nccl_rollout_topology(self) -> tuple[int, dict[str, int]]: ] offsets: dict[str, int] = {} next_rank = 0 + borrowed = set(getattr(self, "_rollout_unavailable_ids", ())) - set( + getattr(self, "_restored_rollouts_pending_publication", ()) + ) + selected = getattr(self, "_publication_target_ids", None) for instance_id, tp_degree in entries: + if selected is not None and instance_id not in selected: + continue + if instance_id in borrowed: + continue offsets[instance_id] = next_rank next_rank += max(1, tp_degree) return 1 + next_rank, offsets @@ -1312,6 +1341,8 @@ def _sync_weights_nccl(self, target_version: int) -> None: # pay NCCL/TCPStore bootstrap on every training version. port = int(getattr(self.config, "rollout_nccl_port", 29620)) world_size, offsets = self._nccl_rollout_topology() + if not offsets: + return request = { "mode": "nccl", "version": int(target_version), @@ -1699,15 +1730,8 @@ def _reload_rollout_from_checkpoint( ) def _current_rollout_instance_ids(self) -> list[str]: - if self._use_heterogeneous: - return [ - str(cfg["instance_id"]) - for cfg in self.rollout_engine.instance_configs - ] - return [ - f"instance_{index}" - for index in range(self.rollout_engine.num_instances) - ] + _, offsets = self._nccl_rollout_topology() + return list(offsets) def _record_rollout_weight_version( self, @@ -2005,10 +2029,11 @@ def _init_global_resource_planner(self): planner_cfg = getattr(self.config, "global_resource_planner", None) if planner_cfg is None or not getattr(planner_cfg, "enabled", False): return - if not self.train_engine or not self.train_engine.is_batch_source(): + if not self.train_engine or not self._is_rollout_source(): return self.global_resource_planner = GlobalResourcePlanner.from_config(self.config) + self.global_resource_planner.runtime_capabilities = CoreRuntimeCapabilities.from_engine(self.train_engine) self.runtime_elastic_executor = RuntimeElasticExecutor( config=self.config, planner=self.global_resource_planner, @@ -2017,6 +2042,7 @@ def _init_global_resource_planner(self): dispatcher=self.dispatcher, membership_run_id=self._elastic_membership_run_id or None, ) + self.runtime_elastic_executor.trainer = self if self._runtime_boundary_protocol is not None: self.runtime_elastic_executor.runtime_run_id = self._runtime_boundary_protocol.run_id if self._elastic_gradient_server is not None: @@ -2141,7 +2167,7 @@ def _start_runtime_reconfiguration_follower(self) -> None: return if not getattr(planner_cfg, "runtime_coordinate_reconfiguration_ranks", True): return - if self.train_engine is not None and not self.train_engine.is_batch_source(): + if self.train_engine is not None and not self._is_rollout_source(): return if self._runtime_follow_thread is not None: return @@ -2255,6 +2281,18 @@ def _run_global_resource_planner_at_boundary(self, step: int, batch: list, stats self._runtime_post_boundary = step self._runtime_boundary_published = False try: + executor = self.runtime_elastic_executor + if executor is not None and executor.has_pending_boundary_work(): + from RL_Framework.infra.cost_model.global_resource_planner import PlannerDecision + planner = self.global_resource_planner + plan = planner._build_current_plan(self.config, planner._history) + decision = PlannerDecision( + step=step, should_reconfigure=True, reason="hybrid_boundary_maintenance", + current_plan=plan, candidate_plan=plan, + resource_should_reconfigure=False, + ) + self._apply_global_resource_planner_decision(step, decision, stats) + return # The only result-consumption location remains the existing POST path. self._run_global_resource_planner(step, batch, stats) finally: @@ -2314,6 +2352,7 @@ def _consume_runtime_reconfiguration_request(self, *, expected_coord_id: str = " return paused = False + state = None try: if self.dispatcher is not None and hasattr(self.dispatcher, "pause"): self.dispatcher.pause() @@ -2353,13 +2392,14 @@ def _consume_runtime_reconfiguration_request(self, *, expected_coord_id: str = " ) tmp_ready.replace(ready_path) - deadline = time.monotonic() + float( - getattr(planner_cfg, "runtime_drain_timeout_s", 3600.0) - ) + timeout = float(getattr(planner_cfg, "runtime_drain_timeout_s", 3600.0)) + deadline = time.monotonic() + timeout applied_path = coord_dir / "applied.json" aborted_path = coord_dir / "aborted.json" state: dict[str, Any] | None = None while time.monotonic() < deadline: + from RL_Framework.infra.elastic.core_reconfiguration import follow_core_command + follow_core_command(self, coord_dir, coord_id) if aborted_path.exists(): aborted = json.loads(aborted_path.read_text(encoding="utf-8")) aborted_job_id = str(aborted.get("job_id", "")) @@ -2390,8 +2430,8 @@ def _consume_runtime_reconfiguration_request(self, *, expected_coord_id: str = " self._runtime_reconfiguration_coord_id = coord_id self._runtime_handled_request_ids.add(coord_id) finally: - if paused and self.dispatcher is not None and hasattr(self.dispatcher, "resume"): - self.dispatcher.resume() + if paused and state is not None: + self._resume_rollout_pipeline() def _follow_applied_runtime_reconfiguration_if_available(self, coord_dir: Path) -> None: """Apply a completed non-draining runtime reconfiguration on peer ranks.""" @@ -2610,6 +2650,9 @@ def _rebind_rollout_engine_from_config(self, *, reason: str) -> None: ) ) ) + executor = getattr(self, "runtime_elastic_executor", None) + if executor is not None: + executor.rollout_engine = self.rollout_engine self._reset_rollout_pipeline_after_reconfigure() if self.is_main_process: instances = getattr(self.config.heterogeneous_rollout, "instances", []) @@ -2631,6 +2674,9 @@ def _rebind_rollout_engine_from_config(self, *, reason: str) -> None: def _apply_runtime_reconfiguration_state(self, state: dict[str, Any]) -> None: self._apply_training_reconfiguration_state(state.get("training") or {}) + from RL_Framework.infra.elastic.rollout_lifecycle import apply_availability + unavailable = state.get("unavailable_rollout_ids", []) + apply_availability(self, unavailable) instances = state.get("instances") or [] if not instances: return @@ -2646,7 +2692,7 @@ def _apply_runtime_reconfiguration_state(self, state: dict[str, Any]) -> None: ) for idx, inst in enumerate(instances) ] - if ( + if (not unavailable and self.rollout_engine is not None and hasattr(self.rollout_engine, "reconfigure_from_plan") ): @@ -2736,6 +2782,143 @@ def _prepare_elastic_training_step(self, step: int) -> list[str]: set_step(step, active_hybrid_ids) return active_hybrid_ids + def _initialize_rollout_source_layout(self) -> None: + """Register one stable data producer per initialized Core DP replica. + + Keep source identities and admitted groups across rollout client + rebindings and elastic Hybrid membership changes. + """ + source = None + if self.train_engine.is_batch_source(): + source = (int(self.train_engine.get_data_parallel_rank()), int(self.rank)) + if dist.is_available() and dist.is_initialized(): + records = [None] * dist.get_world_size() + dist.all_gather_object(records, source) + else: + records = [source] + sources = sorted(item for item in records if item is not None) + expected = int(self.train_engine.get_data_parallel_world_size()) + if [item[0] for item in sources] != list(range(expected)): + raise RuntimeError("initial topology requires one logical source per DP replica") + self._rollout_source_ranks = tuple(item[1] for item in sources) + self._rollout_source_world_size = len(sources) + self._rollout_source_rank = ( + self._rollout_source_ranks.index(self.rank) + if self.rank in self._rollout_source_ranks else None + ) + + def _is_rollout_source(self) -> bool: + if hasattr(self, "_rollout_source_rank"): + return self._rollout_source_rank is not None + return self.train_engine.is_batch_source() + + def _source_batch_sample_count(self) -> int: + source_size = int(getattr(self, "_rollout_source_world_size", self.train_engine.get_data_parallel_world_size())) + source_rank = getattr(self, "_rollout_source_rank", self.train_engine.get_data_parallel_rank()) + if source_rank is None: + return 0 + counts = complete_group_sample_counts( + int(self.config.batch_size), int(self.config.n_samples), source_size, + ) + return counts[int(source_rank)] + + def _collect_fixed_global_batch(self, data_gen, source_count, active_hybrid_ids): + """Gather distinct DP sources before assigning the fixed batch to compute. + + TP/PP followers participate in the same collective but cannot overwrite + another DP source's complete groups. + """ + record = None + if self._is_rollout_source(): + try: + trajectories = self._collect_complete_grpo_batch(data_gen, source_count) + error = "" if len(trajectories) == source_count else ( + f"incomplete fixed batch: expected {source_count}, got {len(trajectories)}" + ) + except Exception as exc: + trajectories, error = [], f"{type(exc).__name__}: {exc}" + record = { + "source_rank": int(getattr(self, "_rollout_source_rank", self.train_engine.get_data_parallel_rank())), + "trajectories": trajectories, "error": error, + } + # Rollout may run longer than the NCCL process-group timeout. Wait + # through the existing out-of-band readiness protocol before entering + # the data collective, including on a source that collected an error. + self._wait_for_all_ranks_after_rollout( + int(getattr(self, "global_step", 0)), + record["trajectories"] if record is not None else [], + ) + if dist.is_available() and dist.is_initialized(): + records = [None] * dist.get_world_size() + dist.all_gather_object(records, { + "source": record, + "core_id": int(self.train_engine.get_data_parallel_rank()), + "hybrids": sorted(active_hybrid_ids), + }) + else: + records = [{"source": record, "core_id": int(self.train_engine.get_data_parallel_rank()), + "hybrids": sorted(active_hybrid_ids)}] + sources, hybrid_targets = {}, {} + for item in records: + for worker_id in item["hybrids"]: + if worker_id in hybrid_targets and hybrid_targets[worker_id] != item["core_id"]: + raise RuntimeError(f"hybrid assigned to multiple core replicas: {worker_id}") + hybrid_targets[worker_id] = item["core_id"] + source = item["source"] + if source is None: + continue + if source["error"]: + raise RuntimeError(f"source {source['source_rank']} batch collection failed: {source['error']}") + if source["source_rank"] in sources: + raise RuntimeError(f"multiple logical batch sources for {source['source_rank']}") + sources[source["source_rank"]] = source + source_size = int(getattr(self, "_rollout_source_world_size", self.train_engine.get_data_parallel_world_size())) + if set(sources) != set(range(source_size)): + raise RuntimeError("missing logical batch source") + dp_size = int(self.train_engine.get_data_parallel_world_size()) + global_batch = [] + for source_rank in range(source_size): + global_batch.extend(sources[source_rank]["trajectories"]) + hybrids = sorted(hybrid_targets) + if len(global_batch) != int(self.config.batch_size): + raise RuntimeError("collected batch does not match fixed global batch size") + participants = [f"dp{i}" for i in range(dp_size)] + sorted(hybrids) + partition_complete_grpo_groups(global_batch, int(self.config.n_samples), participants) + return global_batch, participants + + def _assign_fixed_global_batch(self, step, global_batch, participants): + assignments = partition_complete_grpo_groups( + global_batch, int(self.config.n_samples), participants, + ) + core_id = f"dp{self.train_engine.get_data_parallel_rank()}" + core_batch = assignments[core_id] + core_ids = {f"dp{i}" for i in range(self.train_engine.get_data_parallel_world_size())} + hybrid_counts = { + worker_id: len(assignments[worker_id]) + for worker_id in participants if worker_id not in core_ids + } + setter = getattr(self.train_engine, "set_elastic_batch_partition", None) + if callable(setter): + setter(len(core_batch), len(global_batch), hybrid_counts) + elif hybrid_counts: + raise RuntimeError("train engine cannot normalize an elastic fixed batch") + dispatch_error = "" + if self.is_main_process and hybrid_counts: + self._pending_hybrid_training_batches = { + worker_id: assignments[worker_id] for worker_id in hybrid_counts + } + try: + self._dispatch_hybrid_training_batches(step) + except Exception as exc: + dispatch_error = f"{type(exc).__name__}: {exc}" + if hybrid_counts and dist.is_available() and dist.is_initialized(): + errors = [None] * dist.get_world_size() + dist.all_gather_object(errors, dispatch_error) + dispatch_error = next((error for error in errors if error), "") + if dispatch_error: + raise RuntimeError(f"fixed-batch hybrid dispatch failed: {dispatch_error}") + return core_batch + def _refresh_nonblocking_elastic_membership(self) -> None: planner_cfg = getattr(self.config, "global_resource_planner", None) domain = getattr(self.train_engine, "elastic_gradient_domain", None) @@ -2774,6 +2957,9 @@ def _dispatch_hybrid_training_batches(self, step: int) -> None: task_dir = Path(getattr(planner_cfg, "hybrid_worker_task_dir", "./logs/elastic_training_tasks")) task_dir.mkdir(parents=True, exist_ok=True) lockstep = bool(getattr(planner_cfg, "hybrid_lockstep_gradient_sync", True)) + domain = getattr(self.train_engine, "elastic_gradient_domain", None) + frozen_epochs = domain.membership_state().get("membership_epochs", {}) if domain is not None else {} + state_version = int(getattr(self.train_engine, "optimizer_revision", self.train_engine.get_version())) for worker_id, trajectories in self._pending_hybrid_training_batches.items(): membership_path = task_dir / "membership" / f"{worker_id}.json" if not membership_path.exists(): @@ -2782,17 +2968,19 @@ def _dispatch_hybrid_training_batches(self, step: int) -> None: if (state.get("run_id") != self._elastic_membership_run_id or str(state.get("role")) != "hybrid_training"): raise RuntimeError(f"active hybrid membership changed before dispatch: {worker_id}") + if worker_id in frozen_epochs and int(state.get("membership_epoch", 0)) != int(frozen_epochs[worker_id]): + raise RuntimeError(f"active hybrid generation changed before dispatch: {worker_id}") task_path = task_dir / f"{worker_id}.step_{step}.pt" tmp_path = task_path.with_suffix(task_path.suffix + ".tmp") torch.save( { "trajectories": trajectories, "step": int(step), - "state_version": int(self.train_engine.get_version()), + "state_version": state_version, "membership_epoch": int(state.get("membership_epoch", 0)), "sample_count": len(trajectories), "global_sample_count": int(self.config.batch_size), - "snapshot_path": "" if lockstep else str(self.train_engine.get_elastic_state_snapshot_path(int(self.train_engine.get_version()))), + "snapshot_path": "" if lockstep else str(self.train_engine.get_elastic_state_snapshot_path(state_version)), }, tmp_path, ) @@ -3103,7 +3291,7 @@ def _apply_global_resource_planner_decision(self, step: int, decision, stats: di False, ) ) - if rollout_runtime_enabled: + if rollout_runtime_enabled and not getattr(self, "_rollout_unavailable_ids", ()): self._rebind_rollout_engine_from_config( reason="after_runtime_reconfiguration" ) @@ -3148,6 +3336,9 @@ def _pause_rollout_pipeline(self) -> None: self.dispatcher.pause() def _resume_rollout_pipeline(self) -> None: + instances = getattr(getattr(self.config, "heterogeneous_rollout", None), "instances", ()) + if instances and len(getattr(self, "_rollout_unavailable_ids", ())) == len(instances): + return if self.dispatcher is not None and hasattr(self.dispatcher, "resume"): self.dispatcher.resume() @@ -3336,144 +3527,3 @@ def _record_history_step( resource_config=self._resource_config_snapshot, control_plane=control_plane, ) - - def _initialize_rollout_source_layout(self) -> None: - """Register one stable data producer per initialized Core DP replica. - - Keep source identities and admitted groups across rollout client - rebindings and elastic Hybrid membership changes. - """ - source = None - if self.train_engine.is_batch_source(): - source = (int(self.train_engine.get_data_parallel_rank()), int(self.rank)) - if dist.is_available() and dist.is_initialized(): - records = [None] * dist.get_world_size() - dist.all_gather_object(records, source) - else: - records = [source] - sources = sorted(item for item in records if item is not None) - expected = int(self.train_engine.get_data_parallel_world_size()) - if [item[0] for item in sources] != list(range(expected)): - raise RuntimeError("initial topology requires one logical source per DP replica") - self._rollout_source_ranks = tuple(item[1] for item in sources) - self._rollout_source_world_size = len(sources) - self._rollout_source_rank = ( - self._rollout_source_ranks.index(self.rank) - if self.rank in self._rollout_source_ranks else None - ) - - - def _is_rollout_source(self) -> bool: - if hasattr(self, "_rollout_source_rank"): - return self._rollout_source_rank is not None - return self.train_engine.is_batch_source() - - - def _source_batch_sample_count(self) -> int: - source_size = int(getattr(self, "_rollout_source_world_size", self.train_engine.get_data_parallel_world_size())) - source_rank = getattr(self, "_rollout_source_rank", self.train_engine.get_data_parallel_rank()) - if source_rank is None: - return 0 - counts = complete_group_sample_counts( - int(self.config.batch_size), int(self.config.n_samples), source_size, - ) - return counts[int(source_rank)] - - - def _collect_fixed_global_batch(self, data_gen, source_count, active_hybrid_ids): - """Gather distinct DP sources before assigning the fixed batch to compute. - - TP/PP followers participate in the same collective but cannot overwrite - another DP source's complete groups. - """ - record = None - if self._is_rollout_source(): - try: - trajectories = self._collect_complete_grpo_batch(data_gen, source_count) - error = "" if len(trajectories) == source_count else ( - f"incomplete fixed batch: expected {source_count}, got {len(trajectories)}" - ) - except Exception as exc: - trajectories, error = [], f"{type(exc).__name__}: {exc}" - record = { - "source_rank": int(getattr(self, "_rollout_source_rank", self.train_engine.get_data_parallel_rank())), - "trajectories": trajectories, "error": error, - } - # Rollout may run longer than the NCCL process-group timeout. Wait - # through the existing out-of-band readiness protocol before entering - # the data collective, including on a source that collected an error. - self._wait_for_all_ranks_after_rollout( - int(getattr(self, "global_step", 0)), - record["trajectories"] if record is not None else [], - ) - if dist.is_available() and dist.is_initialized(): - records = [None] * dist.get_world_size() - dist.all_gather_object(records, { - "source": record, - "core_id": int(self.train_engine.get_data_parallel_rank()), - "hybrids": sorted(active_hybrid_ids), - }) - else: - records = [{"source": record, "core_id": int(self.train_engine.get_data_parallel_rank()), - "hybrids": sorted(active_hybrid_ids)}] - sources, hybrid_targets = {}, {} - for item in records: - for worker_id in item["hybrids"]: - if worker_id in hybrid_targets and hybrid_targets[worker_id] != item["core_id"]: - raise RuntimeError(f"hybrid assigned to multiple core replicas: {worker_id}") - hybrid_targets[worker_id] = item["core_id"] - source = item["source"] - if source is None: - continue - if source["error"]: - raise RuntimeError(f"source {source['source_rank']} batch collection failed: {source['error']}") - if source["source_rank"] in sources: - raise RuntimeError(f"multiple logical batch sources for {source['source_rank']}") - sources[source["source_rank"]] = source - source_size = int(getattr(self, "_rollout_source_world_size", self.train_engine.get_data_parallel_world_size())) - if set(sources) != set(range(source_size)): - raise RuntimeError("missing logical batch source") - dp_size = int(self.train_engine.get_data_parallel_world_size()) - global_batch = [] - for source_rank in range(source_size): - global_batch.extend(sources[source_rank]["trajectories"]) - hybrids = sorted(hybrid_targets) - if len(global_batch) != int(self.config.batch_size): - raise RuntimeError("collected batch does not match fixed global batch size") - participants = [f"dp{i}" for i in range(dp_size)] + sorted(hybrids) - partition_complete_grpo_groups(global_batch, int(self.config.n_samples), participants) - return global_batch, participants - - - def _assign_fixed_global_batch(self, step, global_batch, participants): - assignments = partition_complete_grpo_groups( - global_batch, int(self.config.n_samples), participants, - ) - core_id = f"dp{self.train_engine.get_data_parallel_rank()}" - core_batch = assignments[core_id] - core_ids = {f"dp{i}" for i in range(self.train_engine.get_data_parallel_world_size())} - hybrid_counts = { - worker_id: len(assignments[worker_id]) - for worker_id in participants if worker_id not in core_ids - } - setter = getattr(self.train_engine, "set_elastic_batch_partition", None) - if callable(setter): - setter(len(core_batch), len(global_batch), hybrid_counts) - elif hybrid_counts: - raise RuntimeError("train engine cannot normalize an elastic fixed batch") - dispatch_error = "" - if self.is_main_process and hybrid_counts: - self._pending_hybrid_training_batches = { - worker_id: assignments[worker_id] for worker_id in hybrid_counts - } - try: - self._dispatch_hybrid_training_batches(step) - except Exception as exc: - dispatch_error = f"{type(exc).__name__}: {exc}" - if hybrid_counts and dist.is_available() and dist.is_initialized(): - errors = [None] * dist.get_world_size() - dist.all_gather_object(errors, dispatch_error) - dispatch_error = next((error for error in errors if error), "") - if dispatch_error: - raise RuntimeError(f"fixed-batch hybrid dispatch failed: {dispatch_error}") - return core_batch From 482a4aa7be6c2db6b4fd5d867bae8f45f25883be Mon Sep 17 00:00:00 2001 From: cytsoi <154662279+tcyjohn@users.noreply.github.com> Date: Thu, 22 Oct 2026 14:25:10 +0800 Subject: [PATCH 9/9] test: use library timeout exceptions across Python versions --- tests/test_ehp_membership_release.py | 3 ++- tests/test_runtime_boundary_integration.py | 5 +++-- tests/test_runtime_executor_shutdown.py | 3 ++- tests/test_serving_protocol.py | 2 +- 4 files changed, 8 insertions(+), 5 deletions(-) diff --git a/tests/test_ehp_membership_release.py b/tests/test_ehp_membership_release.py index c4b0d05..91c9a1e 100644 --- a/tests/test_ehp_membership_release.py +++ b/tests/test_ehp_membership_release.py @@ -4,6 +4,7 @@ import threading import time from concurrent.futures import ThreadPoolExecutor +from concurrent.futures import TimeoutError as FutureTimeout from types import SimpleNamespace import pytest @@ -236,7 +237,7 @@ def pause_after_domain_activation(worker): # worker-role update. With the old gap it completed in this window. try: releasing.result(timeout=0.1) - except TimeoutError: + except FutureTimeout: pass finally: resume.set() diff --git a/tests/test_runtime_boundary_integration.py b/tests/test_runtime_boundary_integration.py index 2e3998a..dc67c61 100644 --- a/tests/test_runtime_boundary_integration.py +++ b/tests/test_runtime_boundary_integration.py @@ -3,6 +3,7 @@ import json import threading from concurrent.futures import ThreadPoolExecutor +from concurrent.futures import TimeoutError as FutureTimeout from types import SimpleNamespace import pytest @@ -206,7 +207,7 @@ def test_final_boundary_waits_for_leader_close_before_follower_teardown(tmp_path with ThreadPoolExecutor(max_workers=1) as threads: waiting = threads.submit(follower._cleanup) - with pytest.raises(TimeoutError): + with pytest.raises(FutureTimeout): waiting.result(timeout=0.05) leader._cleanup() waiting.result(timeout=2) @@ -268,7 +269,7 @@ def test_final_draining_transaction_reaches_terminal_then_done(tmp_path): following.result(timeout=2) waiting = threads.submit(follower._cleanup) - with pytest.raises(TimeoutError): + with pytest.raises(FutureTimeout): waiting.result(timeout=0.05) leader._cleanup() waiting.result(timeout=2) diff --git a/tests/test_runtime_executor_shutdown.py b/tests/test_runtime_executor_shutdown.py index 0c9fc0d..a1013a5 100644 --- a/tests/test_runtime_executor_shutdown.py +++ b/tests/test_runtime_executor_shutdown.py @@ -3,6 +3,7 @@ import json import threading from concurrent.futures import CancelledError, ThreadPoolExecutor +from concurrent.futures import TimeoutError as FutureTimeout from types import SimpleNamespace import pytest @@ -166,7 +167,7 @@ def close(): assert created.wait(3) finished = runner.submit(close) assert closing.wait(3) - with pytest.raises(TimeoutError): + with pytest.raises(FutureTimeout): finished.result(timeout=0.05) finally: resume.set() diff --git a/tests/test_serving_protocol.py b/tests/test_serving_protocol.py index eaa805d..1363f91 100644 --- a/tests/test_serving_protocol.py +++ b/tests/test_serving_protocol.py @@ -85,7 +85,7 @@ def test_positive_drain_timeout_reopens_before_any_mutation(): async def run(): state = ServingState(1) await state.enter() - with pytest.raises(TimeoutError): + with pytest.raises(asyncio.TimeoutError): await state.begin_reload(.001) assert not state.closed and state.failed is None and state.server_generation == 0 await state.leave()