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/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/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/test_cpu_ci.sh b/scripts/test_cpu_ci.sh index ee6430f..d73bb53 100644 --- a/scripts/test_cpu_ci.sh +++ b/scripts/test_cpu_ci.sh @@ -6,8 +6,17 @@ 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_rollout_credit.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/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_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_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_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/tests/test_serving_protocol.py b/tests/test_serving_protocol.py new file mode 100644 index 0000000..1363f91 --- /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(asyncio.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/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 48d89a8..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 @@ -951,7 +955,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 +1119,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) @@ -1319,8 +1323,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 +1353,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 +1406,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 +1414,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)