From fb5c03f2065124d31a93092849d7039cc0e81ef3 Mon Sep 17 00:00:00 2001 From: Codex Date: Wed, 11 Nov 2026 12:18:02 +0000 Subject: [PATCH 1/5] feat(rl): preserve request provenance and shared routed-expert references Mechanical extraction from XTuner #2257 fixed head 43d62b6. No new implementation. Review commit only: protocol, affinity, lifecycle, Trainer guards and cold-start wiring form one atomic core merge group. Do not enable or merge this intermediate commit alone. Trie traversal overlaps existing #2252; preserves #2257 request-aware reference rules rather than replacing the file. --- xtuner/v1/data_proto/rl_data.py | 6 +- .../agent_in_sandbox_loop.py | 29 ++++++++ .../v1/rl/agent_loop_manager/produce_utils.py | 20 ++++-- xtuner/v1/rl/rollout/trace_store.py | 69 ++++++++++++++----- 4 files changed, 100 insertions(+), 24 deletions(-) diff --git a/xtuner/v1/data_proto/rl_data.py b/xtuner/v1/data_proto/rl_data.py index eb9dad2cfe..2ba5aaa0b8 100644 --- a/xtuner/v1/data_proto/rl_data.py +++ b/xtuner/v1/data_proto/rl_data.py @@ -250,7 +250,7 @@ def reset_response(self) -> None: None """ routed_experts = self.routed_experts - if routed_experts is not None: + if routed_experts is not None and not self.extra_fields.get("request_ids"): from ray import ObjectRef from xtuner.v1.rl.utils.ray_utils import free_object_refs @@ -342,6 +342,10 @@ def free_rollout_state_refs(rollout_state: RolloutState) -> None: from xtuner.v1.rl.utils.ray_utils import free_object_refs + # Request-aware trace routes are shared with sibling segments and workers. + # Keep the ordinary rollout's explicit free contract when provenance is off. + if rollout_state.extra_fields.get("request_ids"): + rollout_state.routed_experts = None refs: list[ObjectRef] = [] def clear_object_refs(value: Any) -> Any: diff --git a/xtuner/v1/rl/agent_loop/sandbox_agent_loop/agent_in_sandbox_loop.py b/xtuner/v1/rl/agent_loop/sandbox_agent_loop/agent_in_sandbox_loop.py index a1be70b247..74a20c0297 100644 --- a/xtuner/v1/rl/agent_loop/sandbox_agent_loop/agent_in_sandbox_loop.py +++ b/xtuner/v1/rl/agent_loop/sandbox_agent_loop/agent_in_sandbox_loop.py @@ -195,6 +195,7 @@ class AgentInSandboxLoopConfig(AgentLoopConfig): max_concurrent_samples: int | None = None mode: Literal["train", "eval"] = "train" requires_rollout_proxy: bool = True + require_turn_provenance: bool = False def build_local( self, rollout_controller: RolloutController | None = None, judger: Judger | None = None, logger=None @@ -207,6 +208,7 @@ def build_local( logger=logger, max_concurrent_samples=self.max_concurrent_samples, mode=self.mode, + require_turn_provenance=self.require_turn_provenance, ) @@ -220,6 +222,7 @@ def __init__( logger=None, max_concurrent_samples: int | None = None, mode: Literal["train", "eval"] = "train", + require_turn_provenance: bool = False, ): if hf_checkpoint is None: raise ValueError("hf_checkpoint must be provided for AgentInSandboxLoop.") @@ -227,6 +230,7 @@ def __init__( self.max_concurrent_samples = max_concurrent_samples self._sample_semaphore = asyncio.Semaphore(max_concurrent_samples) if max_concurrent_samples else None self.mode = mode + self.require_turn_provenance = require_turn_provenance async def generate_group(self, rollout_state: list[RolloutState], **kwargs) -> list[RolloutState]: async def generate_one(state: RolloutState) -> list[RolloutState]: @@ -365,6 +369,31 @@ async def _build_rollout_states(self, rollout_state: RolloutState, item: AgentRo # indistinguishable from a dense trace; catching that would need the MoE-vs-dense signal the loop lacks # (session_server gates on ``enable_return_routed_experts``). segment_state.routed_experts = data["routed_experts"] + if self.require_turn_provenance: + # main's response_ids is the full suffix, including tool/history + # positions. Versions must follow the same positions, not only + # the supervised labels. Preserve each original request's lease. + versions = data.get("token_model_steps", [])[prompt_len:] + if len(versions) != len(segment_state.response_ids) or any(v is None for v in versions): + raise ValueError("Missing request-time response model versions") + if not data["routed_experts"]: + raise ValueError("turn-resume MoE training requires routed experts on every node") + records = _load_messages_artifact(item.artifacts) + matching = [ + record for record in records if record.get("messages") == messages and record.get("tools") == tools + ] + recorded_ids = matching[0].get("request_ids", []) if len(matching) == 1 else [] + path_ids = data.get("request_ids", []) + if ( + not recorded_ids + or not path_ids + or recorded_ids[-1] != path_ids[-1] + or not set(path_ids).issubset(recorded_ids) + or not await trace_store.contains_request_ids.remote(str(rollout_state.session_id), recorded_ids) + ): + raise ValueError("Agent transcript request IDs do not match committed trace") + segment_state.response_model_steps = versions + segment_state.extra_fields["request_ids"] = data["request_ids"] if segment_state.response_ids: segment_state.response = self.tokenizer.decode(segment_state.response_ids) else: diff --git a/xtuner/v1/rl/agent_loop_manager/produce_utils.py b/xtuner/v1/rl/agent_loop_manager/produce_utils.py index 435bca7e6c..053cca332c 100644 --- a/xtuner/v1/rl/agent_loop_manager/produce_utils.py +++ b/xtuner/v1/rl/agent_loop_manager/produce_utils.py @@ -140,6 +140,7 @@ async def generate_group( for item in rollout_state: item.extra_fields["producer_future_step"] = self.train_step + inherited_tokens = sum(len(item.response_ids or []) for item in rollout_state) start = time.perf_counter() if isinstance(self.agent_loop, ray.actor.ActorHandle): result = await self.agent_loop.generate_group.remote( @@ -152,15 +153,24 @@ async def generate_group( enable_partial_rollout=enable_partial_rollout, ) elapsed = time.perf_counter() - start - if isinstance(result, list): - for item in result: - if isinstance(item, RolloutState): - item.extra_fields[GROUP_GENERATE_TIME_KEY] = elapsed + new_tokens = max(0, sum(len(item.response_ids or []) for item in result) - inherited_tokens) + for index, item in enumerate(result): + extra_fields = getattr(item, "extra_fields", None) + if extra_fields is None: + extra_fields = {} + setattr(item, "extra_fields", extra_fields) + extra_fields[GROUP_GENERATE_TIME_KEY] = elapsed + # Provenance may already cover every new token. Count generation + # separately, once per group, without mutating token versions. + extra_fields["new_response_tokens"] = new_tokens if index == 0 else 0 return result async def put_generated_group(self, group: list[RolloutState]) -> bool: produced_tokens = sum( - max(0, len(item.response_ids or []) - len(item.response_model_steps or [])) for item in group + item.extra_fields.get( + "new_response_tokens", max(0, len(item.response_ids or []) - len(item.response_model_steps or [])) + ) + for item in group ) initial_status = get_group_status(group) diff --git a/xtuner/v1/rl/rollout/trace_store.py b/xtuner/v1/rl/rollout/trace_store.py index 36ba406e0a..987a79ca26 100644 --- a/xtuner/v1/rl/rollout/trace_store.py +++ b/xtuner/v1/rl/rollout/trace_store.py @@ -98,6 +98,8 @@ def _resolve_routed_experts(expert_keys: List[Any], session_id: str) -> List[Any class TokenizedSegment(BaseModel): model_config = ConfigDict(extra="forbid") + request_id: str | None = None + model_step: int | None = None text: str token_ids: List[int] labels: List[int] | None = Field(default=None, repr=False) @@ -141,13 +143,12 @@ def keys(self) -> List[str]: """Get all keys (i.e., strings) stored in the Trie.""" result = [] - def _collect(node: TreeNode, path: List[str]): + stack = [(self.root, "")] + while stack: + node, key = stack.pop() if node.value is not None: - result.append("".join(path)) - for token, child in node.children.items(): - _collect(child, path + [token]) - - _collect(self.root, []) + result.append(key) + stack.extend((child, key + token) for token, child in reversed(list(node.children.items()))) return result def insert(self, key: str, value: Any) -> None: @@ -173,6 +174,13 @@ def insert(self, key: str, value: Any) -> None: node = node.children[key] break + if node.value is not None and ( + getattr(node.value, "request_id", None) is not None or getattr(value, "request_id", None) is not None + ): + raise ValueError( + f"Conflicting trace write: old_request_id={getattr(node.value, 'request_id', None)} " + f"new_request_id={getattr(value, 'request_id', None)}" + ) node.value = value def search(self, text: str, filter_none: bool = False) -> Tuple[str, List["TreeNode"]]: @@ -210,8 +218,7 @@ def search(self, text: str, filter_none: bool = False) -> Tuple[str, List["TreeN return "".join(key_path), matched_nodes def release(self, key: str | None = None): - """Release nodes and free associated resources (e.g. Ray objects) for a - given key. + """Drop trace ownership; exported and sibling route references remain valid. If the key is None, the entire Trie is released. Prunes empty branches bottom-up. @@ -220,16 +227,17 @@ def release(self, key: str | None = None): If None, releases the entire tree. """ - def _free_subtree(node: TreeNode): - for child in node.children.values(): - _free_subtree(child) - if node.value is not None: - _free_ray_refs(node.value) - node.value = None - node.children.clear() - if key is None: - _free_subtree(self.root) + # Drop only this store's ownership. Borrowers retain live R3 refs. + stack = [self.root] + while stack: + node = stack.pop() + stack.extend(node.children.values()) + if node.value is not None and getattr(node.value, "request_id", None) is None: + _free_ray_refs(node.value) + node.value = None + node.children.clear() + node.parent = None return node = self.root @@ -248,7 +256,8 @@ def _free_subtree(node: TreeNode): return if node.value is not None: - _free_ray_refs(node.value) + if getattr(node.value, "request_id", None) is None: + _free_ray_refs(node.value) node.value = None for parent, token in reversed(path): @@ -402,6 +411,9 @@ def export_training_trace(self, session_id: str, prompt_text: str) -> dict: ) trace: dict[str, Any] = {"input_ids": [], "labels": [], "logprobs": []} expert_keys: list[Any] = [] + response_model_steps = [] + token_model_steps = [] + request_ids = [] for node in nodes: node_val = node.value if not isinstance(node_val, TokenizedSegment): @@ -412,9 +424,30 @@ def export_training_trace(self, session_id: str, prompt_text: str) -> dict: trace["labels"].extend(node_val.labels) trace["logprobs"].extend(node_val.logprobs) expert_keys.append(node_val.expert_key) + response_model_steps.extend(node_val.model_step for label in node_val.labels if label != -100) + token_model_steps.extend([node_val.model_step] * len(node_val.token_ids)) + if node_val.request_id is not None and node_val.request_id not in request_ids: + request_ids.append(node_val.request_id) trace["routed_experts"] = _resolve_routed_experts(expert_keys, session_id) + trace["token_model_steps"] = token_model_steps + if request_ids: + trace["response_model_steps"] = response_model_steps + trace["request_ids"] = request_ids return trace + def contains_request_ids(self, session_id: str, request_ids: list[str]) -> bool: + """Validate transcript provenance across committed branches of one session.""" + trie = self.sessions.get(session_id) + if trie is None: + return False + missing = set(request_ids) + stack = [trie.root] + while stack and missing: + node = stack.pop() + missing.discard(getattr(node.value, "request_id", None)) + stack.extend(node.children.values()) + return not missing + def get_objects(self, keys: list[str]) -> list[ray.ObjectRef]: """Fetch ray.ObjectRef elements by their keys. From ff2ad68bde2506a66e30fdf566f3af68a04773b8 Mon Sep 17 00:00:00 2001 From: Codex Date: Wed, 11 Nov 2026 12:18:03 +0000 Subject: [PATCH 2/5] feat(rollout): settle complete requests behind version publication barriers Mechanical extraction from XTuner #2257 fixed head 43d62b6. No new implementation. Review commit only: protocol, affinity, lifecycle, Trainer guards and cold-start wiring form one atomic core merge group. Do not enable or merge this intermediate commit alone. --- xtuner/v1/rl/rollout/controller.py | 73 ++++++++++- xtuner/v1/rl/rollout/request_gate.py | 120 +++++++++++++++++ xtuner/v1/rl/rollout/session_server.py | 175 ++++++++++++++++++++++++- xtuner/v1/rl/rollout/worker.py | 10 +- 4 files changed, 372 insertions(+), 6 deletions(-) create mode 100644 xtuner/v1/rl/rollout/request_gate.py diff --git a/xtuner/v1/rl/rollout/controller.py b/xtuner/v1/rl/rollout/controller.py index f6fd1824c9..7e52dd5db4 100644 --- a/xtuner/v1/rl/rollout/controller.py +++ b/xtuner/v1/rl/rollout/controller.py @@ -1,4 +1,6 @@ import asyncio +import math +from time import monotonic from typing import Any, TypeAlias from uuid import uuid4 @@ -46,6 +48,8 @@ def __init__( placement_group (PlacementGroup): The placement group for the RolloutWorker actors. """ + self._request_barrier = None + self._request_epoch = 0 self.config = infer_config self.num_gpus_per_engine = self.config.num_gpus_per_engine self.logger = get_logger(log_dir=infer_config.worker_log_dir, tag="RolloutController") @@ -174,6 +178,70 @@ def set_enable_partial_rollout(self, enable: bool) -> None: ] ) + def assert_request_gates_healthy(self): + workers = self.registry.active_entrypoints() + servers = ray.get([w.actor.get_session_server.remote() for w in workers], timeout=ROLLOUT_RAY_GET_TIMEOUT) + if not servers or any(server is None for server in servers): + raise RuntimeError("A fixed Agent request entrypoint is missing") + states = ray.get([s.request_status.remote() for s in servers], timeout=ROLLOUT_RAY_GET_TIMEOUT) + failures = [state for state in states if state["failures"]] + if failures: + raise RuntimeError(f"Unsettled Agent requests stop this run: {failures}") + + def prepare_pause(self, *, freeze_timeout: float, drain_timeout: float): + """Freeze topology, seal every entrypoint, then drain with one deadline.""" + if self._request_barrier is not None: + raise RuntimeError("A request barrier already exists") + self.health_manager.pause() + if not self.health_manager.wait_recovery_done(timeout=freeze_timeout): + raise TimeoutError("Could not freeze rollout lifecycle") + workers = self.registry.active_entrypoints() + servers = ray.get([w.actor.get_session_server.remote() for w in workers], timeout=freeze_timeout) + if not servers or any(server is None for server in servers): + raise RuntimeError("Every fixed rollout entrypoint must own a SessionServer") + self._request_epoch += 1 + barrier = dict(epoch=self._request_epoch, endpoints=tuple(w.session_url for w in workers)) + self._request_barrier = (barrier, servers) + # Send all seals before waiting for any drain. + ray.get([s.seal_requests.remote(barrier["epoch"]) for s in servers], timeout=freeze_timeout) + ray.get([s.drain_requests.remote(drain_timeout) for s in servers], timeout=drain_timeout + freeze_timeout) + return barrier + + def resume_requests(self, barrier: dict, model_step: int, *, timeout: float = 60): + """Publish an updated model under the existing admission barrier. + + The topology control budget covers version installation, clearing backend + abort flags, and reopening. It does not bound training or weight transfer. + A failed install or backend resume cannot open a gate. Reopening multiple + entrypoints is not atomic; an RPC failure may leave some valid gates open. + """ + if not math.isfinite(timeout) or timeout <= 0: + raise ValueError("Request publication timeout must be finite and positive") + if self._request_barrier is None or barrier != self._request_barrier[0]: + raise RuntimeError("Unknown request barrier") + _, servers = self._request_barrier + endpoints = tuple(w.session_url for w in self.registry.active_entrypoints()) + if endpoints != barrier["endpoints"]: + raise RuntimeError("Fixed rollout topology changed under barrier") + deadline = monotonic() + timeout + + def wait_all(actors, method, *args): + remaining = deadline - monotonic() + if remaining <= 0: + raise TimeoutError(f"Request publication timed out before {method}") + refs = [getattr(actor, method).remote(*args) for actor in actors] + return ray.get(refs, timeout=remaining) + + # Install while sealed, then clear every backend abort flag before opening. + wait_all(servers, "install_request_version", barrier["epoch"], model_step) + workers = self.registry.get_target_state_workers(WorkerLifecycleState.ACTIVE) + results = wait_all([w.actor for w in workers], "continue_generation") + if any(result is False for result in results): + raise RuntimeError("Backend resume failed under request barrier") + wait_all(servers, "reopen_requests", barrier["epoch"]) + self._request_barrier = None + self.health_manager.resume() + def pause_generation(self): self.health_manager.pause() # Wait for the health manager to finish recovery before pausing generation. @@ -195,6 +263,8 @@ def pause_generation(self): failed_worker_urls = [worker.url for worker, result in zip(active_workers, results) if result is False] if succeeded_worker_urls: self.logger.info(f"Abort request sent successfully: count={len(succeeded_worker_urls)}") + if failed_worker_urls and self._request_barrier is not None: + raise RuntimeError(f"Backend pause failed: {failed_worker_urls}") if failed_worker_urls: self.logger.warning(f"Abort request failed: worker_urls={failed_worker_urls}") @@ -240,7 +310,8 @@ def mark_worker_groups_lifecycle_state( def continue_generation(self): self._broadcast_to_workers("continue_generation", WorkerLifecycleState.ACTIVE) - self.health_manager.resume() + if self._request_barrier is None: + self.health_manager.resume() def offload(self): self._broadcast_to_workers("offload", WorkerLifecycleState.ACTIVE) diff --git a/xtuner/v1/rl/rollout/request_gate.py b/xtuner/v1/rl/rollout/request_gate.py new file mode 100644 index 0000000000..a782b362cb --- /dev/null +++ b/xtuner/v1/rl/rollout/request_gate.py @@ -0,0 +1,120 @@ +"""Complete-turn admission barrier; forwarding alone certifies settlement. + +Admission-accounting reference: matrix72c/sandbox_rl before fd76fba, +src/sandbox_rl/agents/request_gate.py. This native condition/lease implementation +adds explicit backend/trace settlement; an unknown outcome poisons the barrier. +""" + +from __future__ import annotations + +import asyncio +from contextlib import asynccontextmanager +from dataclasses import dataclass +from uuid import uuid4 + + +@dataclass(frozen=True) +class SettlementReceipt: + backend_ended: bool + terminal_recorded: bool + outcome: str + + @property + def safe_for_weight_update(self) -> bool: + return self.backend_ended and self.terminal_recorded + + +@dataclass +class RequestLease: + request_id: str + session_id: str + epoch: int + model_step: int + settlement: SettlementReceipt | None = None + + +class RequestGate: + def __init__(self, *, wait_timeout: float, initially_open: bool = True): + if wait_timeout <= 0: + raise ValueError("wait_timeout must be positive") + self.wait_timeout = wait_timeout + self._condition = asyncio.Condition() + self._open = initially_open + self._epoch = 0 + self._installed_epoch = 0 + self._model_step = 0 + self._active: dict[str, RequestLease] = {} + self._failures: dict[str, str] = {} + + def snapshot(self) -> dict: + return dict( + open=self._open, + epoch=self._epoch, + model_step=self._model_step, + active=list(self._active), + failures=dict(self._failures), + ) + + @asynccontextmanager + async def request(self, session_id: str): + async with self._condition: + async with asyncio.timeout(self.wait_timeout): + await self._condition.wait_for(lambda: self._open or bool(self._failures)) + if self._failures: + raise RuntimeError(f"Request barrier failed: {self._failures}") + if any(lease.session_id == session_id for lease in self._active.values()): + raise RuntimeError("Concurrent generation in the same session is unsupported") + lease = RequestLease("chatcmpl-" + uuid4().hex, session_id, self._epoch, self._model_step) + self._active[lease.request_id] = lease + try: + yield lease + finally: + async with self._condition: + if lease.settlement is None or not lease.settlement.safe_for_weight_update: + self._open = False + self._failures[lease.request_id] = "backend_or_trace_unsettled" + else: + self._active.pop(lease.request_id) + self._condition.notify_all() + + async def seal(self, epoch: int): + async with self._condition: + if epoch < self._epoch: + raise ValueError("Stale barrier epoch") + if epoch == self._epoch and self._open: + raise ValueError("A reopened epoch cannot be sealed again") + self._open = False + self._epoch = epoch + + async def drain(self, timeout: float): + async with self._condition: + if self._open: + raise RuntimeError("Seal admission before draining") + try: + async with asyncio.timeout(timeout): + await self._condition.wait_for(lambda: not self._active or bool(self._failures)) + except TimeoutError: + self._failures["barrier"] = "drain_timeout" + self._condition.notify_all() + raise + if self._failures: + raise RuntimeError(f"Unsettled requests prevent weight update: {self._failures}") + return self.snapshot() + + async def install_version(self, epoch: int, model_step: int): + async with self._condition: + if self._open or epoch != self._epoch or self._active or self._failures: + raise RuntimeError("Version publication requires the drained current barrier") + if model_step < self._model_step: + raise ValueError("Model versions cannot move backwards") + if self._installed_epoch == epoch and model_step != self._model_step: + raise ValueError("Conflicting version for installed epoch") + self._model_step = model_step + self._installed_epoch = epoch + + async def reopen(self, epoch: int): + async with self._condition: + if epoch != self._epoch or self._installed_epoch != epoch or self._active or self._failures: + raise RuntimeError("Cannot reopen an uninstalled or unsafe barrier") + self._open = True + self._condition.notify_all() diff --git a/xtuner/v1/rl/rollout/session_server.py b/xtuner/v1/rl/rollout/session_server.py index 5e53ad15a6..ab85e40ff9 100644 --- a/xtuner/v1/rl/rollout/session_server.py +++ b/xtuner/v1/rl/rollout/session_server.py @@ -1,5 +1,7 @@ import copy import json +import re +import time from functools import reduce from http import HTTPStatus from operator import add @@ -13,6 +15,7 @@ from xtuner.v1.utils import get_logger from .chat_template import canonicalize_messages_for_chat_template +from .request_gate import RequestGate, RequestLease, SettlementReceipt from .trace_store import TokenizedSegment, get_store @@ -20,7 +23,7 @@ FMT_ANTHROPIC = "anthropic" # Fields the SessionServer consumes locally and never forwards upstream. -_SESSION_SERVER_ONLY_KEYS = {"session_id"} +_SESSION_SERVER_ONLY_KEYS = {"session_id", "_xt_request_id", "_xt_model_step"} def _normalize_path(req_path: str) -> str: @@ -45,6 +48,35 @@ def _is_error_payload(payload: dict) -> bool: return payload.get("error") is not None or payload.get("type") == "error" or payload.get("object") == "error" +def _is_lmdeploy_context_rejection(status: int, body: bytes) -> bool: + """Recognize the pinned LMDeploy's input check before engine generation. + + A finished HTTP error alone does not prove that backend execution ended. + LMDeploy bc2c8da's AsyncEngine rejects input_len >= session_len before + submitting generation. Its OpenAI response has no stable machine error code, + so match that specific payload conservatively; unknown errors stay unsafe. + """ + if status != HTTPStatus.BAD_REQUEST: + return False + try: + payload = json.loads(body) + except (ValueError, UnicodeDecodeError): + return False + if not isinstance(payload, dict) or ( + payload.get("object") != "error" + or payload.get("type") != "invalid_request_error" + or payload.get("code") != 400 + or "choices" in payload + or not isinstance(payload.get("message"), str) + ): + return False + match = re.fullmatch( + r"Input length \(([0-9]{1,20})\) must be smaller than the model context length \(([0-9]{1,20})\)\.", + payload["message"], + ) + return match is not None and int(match[1]) >= int(match[2]) > 0 + + class _UpstreamResponseError(RuntimeError): """Internal marker for an error explicitly returned by the upstream.""" @@ -267,6 +299,8 @@ class SessionServer: read_bufsize (int): Buffer limit for line reader in ClientSession. Default is 64MB (2**26). """ + request_gate: RequestGate | None = None + def __init__( self, worker_base_url: str, @@ -276,6 +310,8 @@ def __init__( request_timeout: float = 1200.0, read_bufsize: int = 2**26, enable_return_routed_experts: bool = False, + turn_resume_wait_timeout: float | None = None, + turn_resume_initially_closed: bool = False, ): self.worker_base_url = worker_base_url.rstrip("/") self.tokenizer = AutoTokenizer.from_pretrained(tokenizer_path, trust_remote_code=True) @@ -284,6 +320,12 @@ def __init__( self.request_timeout = request_timeout self.read_bufsize = read_bufsize self.enable_return_routed_experts = enable_return_routed_experts + self.turn_resume_wait_timeout = turn_resume_wait_timeout + self.turn_resume_initially_closed = turn_resume_initially_closed + self.request_gate = ( + RequestGate(wait_timeout=turn_resume_wait_timeout, initially_open=not turn_resume_initially_closed) + if turn_resume_wait_timeout is not None else None + ) self.store = get_store() self.stop_word = self.tokenizer.eos_token or "" @@ -414,6 +456,14 @@ async def on_response(self, worker_resp: dict, fmt: str, orig_req_body: dict) -> dict: The (possibly trimmed) response — currently unused by the caller but kept for symmetry. """ + if orig_req_body.get("_xt_request_id") is not None and ( + not worker_resp.get("choices") + or any(c.get("finish_reason") not in {"stop", "length", "tool_calls"} for c in worker_resp["choices"]) + ): + raise RuntimeError( + "Backend returned an unsuccessful generation terminal: " + + str([c.get("finish_reason") for c in worker_resp.get("choices", [])]) + ) session_id = orig_req_body["session_id"] if fmt == FMT_ANTHROPIC: @@ -487,8 +537,9 @@ async def on_response(self, worker_resp: dict, fmt: str, orig_req_body: dict) -> if delta_len > 0: delta_expert = ray.put(raw_routed_expert[prefix_len : prefix_len + delta_len]) response_expert = ray.put(raw_routed_expert[prefix_len + delta_len :]) - elif self.enable_return_routed_experts and _bool_request_value( - orig_req_body.get("return_routed_experts"), True + elif self.enable_return_routed_experts and ( + orig_req_body.get("_xt_request_id") is not None + or _bool_request_value(orig_req_body.get("return_routed_experts"), True) ): raise RuntimeError( "SessionServer response is missing routed_experts, but enable_return_routed_experts is True." @@ -503,7 +554,13 @@ async def on_response(self, worker_resp: dict, fmt: str, orig_req_body: dict) -> await self.store.insert.remote( session_id, key=old_prompt, - value=TokenizedSegment(text=delta_text, token_ids=delta_ids, expert_key=delta_expert), + value=TokenizedSegment( + text=delta_text, + token_ids=delta_ids, + expert_key=delta_expert, + request_id=orig_req_body.get("_xt_request_id"), + model_step=orig_req_body.get("_xt_model_step"), + ), ) await self.store.insert.remote( session_id, @@ -514,6 +571,8 @@ async def on_response(self, worker_resp: dict, fmt: str, orig_req_body: dict) -> logprobs=output_logprobs, labels=output_token_ids, expert_key=response_expert, + request_id=orig_req_body.get("_xt_request_id"), + model_step=orig_req_body.get("_xt_model_step"), length=len(output_token_ids), ), ) @@ -549,8 +608,40 @@ async def stop(self): get_logger().info("SessionServer stopped.") async def _handle_request(self, request: web.Request) -> web.Response: + path = request.match_info["path"] + if self.request_gate is None or not _is_generation_endpoint(request.method, path): + return await self._forward_request(request) + fmt = _detect_format(path) + if fmt != FMT_OPENAI: + return web.json_response( + _error_payload(fmt, "turn-resume supports OpenAI chat completions only", 400), status=400 + ) + try: + body = await request.json() + except (ValueError, UnicodeDecodeError): + return web.json_response(_error_payload(fmt, "Invalid generation JSON", 400), status=400) + if not isinstance(body, dict) or body.get("stream") or not _request_uses_trace_store(body): + return web.json_response( + _error_payload(fmt, "turn-resume requires nonstream traced generation", 400), status=400 + ) + async with self.request_gate.request(str(body["session_id"])) as lease: + started = time.perf_counter() + try: + return await self._forward_request(request, lease) + except BaseException as exc: + get_logger().error( + f"request_forward_failed id={lease.request_id} exception={type(exc).__name__} " + f"elapsed_s={time.perf_counter() - started:.6f} " + f"settlement_recorded={lease.settlement is not None}" + ) + raise + + async def _forward_request(self, request: web.Request, lease: RequestLease | None = None) -> web.Response: """Proxy handler: detect format, run hooks, forward, stream back.""" + request_start = time.perf_counter() + prepare_s = 0.0 + trace_s = 0.0 req_path = request.match_info["path"] is_generation_endpoint = _is_generation_endpoint(request.method, req_path) fmt = _detect_format(req_path) @@ -607,6 +698,9 @@ async def _handle_request(self, request: web.Request) -> web.Response: request_data["tools"] = filtered else: request_data.pop("tools", None) + if lease is not None: + request_data["_xt_request_id"] = lease.request_id + request_data["_xt_model_step"] = lease.model_step orig_req_body = copy.deepcopy(request_data) trace_enabled = _request_uses_trace_store(request_data) @@ -619,13 +713,19 @@ async def _handle_request(self, request: web.Request) -> web.Response: orig_return_token_ids = _bool_request_value(request_data.get("return_token_ids"), False) orig_return_routed_experts = _bool_request_value(request_data.get("return_routed_experts"), True) + prepare_start = time.perf_counter() request_data = await self.on_request(request_data, fmt, trace_enabled=trace_enabled) + prepare_s = time.perf_counter() - prepare_start request_body = json.dumps(request_data).encode("utf-8") except (json.JSONDecodeError, UnicodeDecodeError): pass except Exception as exc: message = f"SessionServer request hook failed: {type(exc).__name__}: {exc}" get_logger().error(message) + if lease is not None: + # Request preparation failed before any backend send. + get_logger().error(f"request_terminal id={lease.request_id} outcome=not_started") + lease.settlement = SettlementReceipt(True, True, "not_started") return web.json_response(_error_payload(fmt, message), status=500) trace_active = is_generation_endpoint and trace_enabled @@ -660,6 +760,7 @@ async def _handle_request(self, request: web.Request) -> web.Response: else None ) + forward_start = time.perf_counter() timeout = ClientTimeout(total=self.request_timeout, sock_connect=30) async with ClientSession(read_bufsize=self.read_bufsize, timeout=timeout) as client: async with client.request( @@ -750,6 +851,7 @@ async def _handle_request(self, request: web.Request) -> web.Response: body=final_raw_response, ) + forward_s = time.perf_counter() - forward_start # Apply abstract on_response processing response_data: Optional[dict] = None skip_done = bool(is_stream and not trace_response) @@ -782,10 +884,49 @@ async def _handle_request(self, request: web.Request) -> web.Response: if response_data is not None: try: + trace_start = time.perf_counter() await self.on_response(response_data, fmt, orig_req_body) + trace_s = time.perf_counter() - trace_start + if lease is not None and any( + c.get("finish_reason") not in {"stop", "length", "tool_calls"} + for c in response_data.get("choices", []) + ): + raise RuntimeError("Unsuccessful backend generation terminal") except Exception as exc: session_error_msg = f"SessionServer response hook failed: {type(exc).__name__}: {exc}" + if lease is not None: + complete = ( + not is_stream + and is_success_response + and response_data is not None + and not _is_error_payload(response_data) + and bool(response_data.get("choices")) + and all(c.get("finish_reason") is not None for c in response_data["choices"]) + ) + terminal = "trace_committed" if complete and not session_error_msg else "failed" + # A known pre-generation rejection is a safely ended failed attempt, + # not a successful trace and not an ambiguous in-flight request. + rejected_before_generation = not is_stream and _is_lmdeploy_context_rejection( + upstream_status, raw_response + ) + backend_ended = complete or rejected_before_generation + get_logger().info(f"request_terminal id={lease.request_id} outcome={terminal} backend_ended={backend_ended}") + if rejected_before_generation: + get_logger().info( + f"request_rejected_before_generation id={lease.request_id} reason=context_length_exceeded" + ) + get_logger().info( + f"request_timing id={lease.request_id} prepare_s={prepare_s:.6f} " + f"forward_s={forward_s:.6f} trace_s={trace_s:.6f} " + f"total_s={time.perf_counter() - request_start:.6f}" + ) + lease.settlement = SettlementReceipt(backend_ended, True, terminal) + if terminal == "trace_committed": + payload = json.loads(response.body) + payload["id"] = lease.request_id + response.body = json.dumps(payload).encode("utf-8") + if session_error_msg: get_logger().error(session_error_msg) @@ -1121,6 +1262,8 @@ def __init__( port: int, request_timeout: float, enable_return_routed_experts: bool = False, + turn_resume_wait_timeout: float | None = None, + turn_resume_initially_closed: bool = False, ): self.worker_base_url = worker_base_url self.tokenizer_path = tokenizer_path @@ -1128,6 +1271,8 @@ def __init__( self.port = port self.request_timeout = request_timeout self.enable_return_routed_experts = enable_return_routed_experts + self.turn_resume_wait_timeout = turn_resume_wait_timeout + self.turn_resume_initially_closed = turn_resume_initially_closed self.server: SessionServer | None = None @property @@ -1145,10 +1290,32 @@ async def start(self) -> str: port=self.port, request_timeout=self.request_timeout, enable_return_routed_experts=self.enable_return_routed_experts, + turn_resume_wait_timeout=self.turn_resume_wait_timeout, + turn_resume_initially_closed=self.turn_resume_initially_closed, ) await self.server.start() return self.server.url + def _gate(self) -> RequestGate: + if self.server is None or self.server.request_gate is None: + raise RuntimeError("turn-resume gate is not enabled") + return self.server.request_gate + + async def request_status(self): + return self._gate().snapshot() + + async def seal_requests(self, epoch: int): + await self._gate().seal(epoch) + + async def drain_requests(self, timeout: float): + return await self._gate().drain(timeout) + + async def install_request_version(self, epoch: int, model_step: int): + await self._gate().install_version(epoch, model_step) + + async def reopen_requests(self, epoch: int): + await self._gate().reopen(epoch) + async def stop(self) -> None: if self.server is not None: await self.server.stop() diff --git a/xtuner/v1/rl/rollout/worker.py b/xtuner/v1/rl/rollout/worker.py index f76b309139..3fa74345ac 100644 --- a/xtuner/v1/rl/rollout/worker.py +++ b/xtuner/v1/rl/rollout/worker.py @@ -16,7 +16,7 @@ import requests # type: ignore[import-untyped] from cyclopts import Group, Parameter from packaging.version import Version -from pydantic import BaseModel, ConfigDict +from pydantic import BaseModel, ConfigDict, Field from ray.util.scheduling_strategies import PlacementGroupSchedulingStrategy from typing_extensions import Annotated @@ -331,6 +331,9 @@ class RolloutConfig(BaseModel): help="Timeout duration (in seconds) for rollout requests.", ), ] = 1200.0 + turn_resume_wait_timeout: float | None = Field(default=None, gt=0) + turn_resume_initially_closed: bool = False + session_server_timeout: Annotated[ float, Parameter( @@ -796,6 +799,8 @@ def _start_session_server(self) -> None: host=self.host, port=self.session_server_port, request_timeout=self.config.session_server_timeout, + turn_resume_wait_timeout=self.config.turn_resume_wait_timeout, + turn_resume_initially_closed=self.config.turn_resume_initially_closed, enable_return_routed_experts=self.enable_return_routed_experts, ) ) @@ -808,6 +813,9 @@ def _start_session_server(self) -> None: f"Request-entrypoint rollout worker rank={self.rank} did not start session server during init." ) + def get_session_server(self): + return self.session_server_actor + def _stop_session_server(self) -> None: if self.session_server_actor is not None: try: From edab2ea5e559e3ba831761a240a59e0eee08b3ef Mon Sep 17 00:00:00 2001 From: Codex Date: Wed, 11 Nov 2026 12:18:03 +0000 Subject: [PATCH 3/5] feat(rollout): pin and retire external HTTP sessions with owned probes Mechanical extraction from XTuner #2257 fixed head 43d62b6. No new implementation. Review commit only: protocol, affinity, lifecycle, Trainer guards and cold-start wiring form one atomic core merge group. Do not enable or merge this intermediate commit alone. --- .../agent_in_sandbox_loop.py | 10 + xtuner/v1/rl/rollout/affinity_router.py | 200 ++++++++++++++++++ xtuner/v1/rl/rollout/controller.py | 18 +- xtuner/v1/rl/rollout/proxy_manager.py | 59 +++++- xtuner/v1/rl/rollout/worker.py | 4 + xtuner/v1/rl/utils/misc.py | 6 +- 6 files changed, 284 insertions(+), 13 deletions(-) create mode 100644 xtuner/v1/rl/rollout/affinity_router.py diff --git a/xtuner/v1/rl/agent_loop/sandbox_agent_loop/agent_in_sandbox_loop.py b/xtuner/v1/rl/agent_loop/sandbox_agent_loop/agent_in_sandbox_loop.py index 74a20c0297..67c1027d56 100644 --- a/xtuner/v1/rl/agent_loop/sandbox_agent_loop/agent_in_sandbox_loop.py +++ b/xtuner/v1/rl/agent_loop/sandbox_agent_loop/agent_in_sandbox_loop.py @@ -196,6 +196,7 @@ class AgentInSandboxLoopConfig(AgentLoopConfig): mode: Literal["train", "eval"] = "train" requires_rollout_proxy: bool = True require_turn_provenance: bool = False + retire_route_on_completion: bool = False def build_local( self, rollout_controller: RolloutController | None = None, judger: Judger | None = None, logger=None @@ -209,6 +210,7 @@ def build_local( max_concurrent_samples=self.max_concurrent_samples, mode=self.mode, require_turn_provenance=self.require_turn_provenance, + retire_route_on_completion=self.retire_route_on_completion, ) @@ -223,6 +225,7 @@ def __init__( max_concurrent_samples: int | None = None, mode: Literal["train", "eval"] = "train", require_turn_provenance: bool = False, + retire_route_on_completion: bool = False, ): if hf_checkpoint is None: raise ValueError("hf_checkpoint must be provided for AgentInSandboxLoop.") @@ -231,6 +234,7 @@ def __init__( self._sample_semaphore = asyncio.Semaphore(max_concurrent_samples) if max_concurrent_samples else None self.mode = mode self.require_turn_provenance = require_turn_provenance + self.retire_route_on_completion = retire_route_on_completion async def generate_group(self, rollout_state: list[RolloutState], **kwargs) -> list[RolloutState]: async def generate_one(state: RolloutState) -> list[RolloutState]: @@ -287,6 +291,12 @@ async def generate_sample( # type: ignore[override] rollout_state.error_msg = f"{type(exc).__name__}: {exc}" self.logger.error(f"[AgentInSandboxLoop] failed: {exc}\n{traceback.format_exc()}") return [rollout_state] + finally: + # This loop returns terminal Agent executions, never partial continuations. + # Retire HTTP ownership after runner completion/cancellation, independently + # of later filtering or replay consumption; keep trace ownership unchanged. + if self.retire_route_on_completion and self.rollout_ctl is not None and rollout_state.session_id is not None: + await self.rollout_ctl.retire_route_sessions.remote([str(rollout_state.session_id)]) async def _run_item(self, item: AgentRolloutItem) -> AgentRolloutItem: runner = _resolve_runner(item.pipeline, str(item.uid)) diff --git a/xtuner/v1/rl/rollout/affinity_router.py b/xtuner/v1/rl/rollout/affinity_router.py new file mode 100644 index 0000000000..43ddfa9d19 --- /dev/null +++ b/xtuner/v1/rl/rollout/affinity_router.py @@ -0,0 +1,200 @@ +"""Job-owned HTTP session affinity in front of the original SessionServers. + +The mapping names a concrete DP entrypoint, not an EP group. This module never +generates tokens, retries a POST, or changes cache/trace/version state. +""" + +from __future__ import annotations + +import asyncio +import logging +import threading +from collections import Counter + +import aiohttp +from aiohttp import web + + +_HOP_HEADERS = { + "connection", + "keep-alive", + "proxy-authenticate", + "proxy-authorization", + "te", + "trailer", + "transfer-encoding", + "upgrade", + "host", + "content-length", +} + + +def _headers(headers): + blocked = _HOP_HEADERS | {x.strip().lower() for x in headers.get("Connection", "").split(",")} + return [(k, v) for k, v in headers.items() if k.lower() not in blocked] + + +class SessionAffinityRouter: + def __init__(self, *, max_sessions=1000000): + self.backends: set[str] = set() + self.bindings: dict[str, str] = {} + # Tombstones prevent a late retry from starting a second Trial after cleanup. + self.retired: set[str] = set() + self.inflight = Counter() + self.assigned = Counter() + self.max_sessions = max_sessions + self.client = None + self.runner = None + self.url = None + self.logger = logging.getLogger(__name__) + + def replace_backends(self, urls): + self.backends = {url.rstrip("/") for url in urls} + # Do not discard bindings on pause or failure. A dead binding fails closed. + + def set_backend(self, url, active): + if url: + if active: + self.backends.add(url.rstrip("/")) + else: + self.backends.discard(url.rstrip("/")) + + def retire_sessions(self, session_ids): + for sid in map(str, session_ids): + url = self.bindings.pop(sid, None) + if url is not None: + self.assigned[url] -= 1 + self.retired.add(sid) + + def acquire(self, sid): + if sid in self.retired: + raise web.HTTPGone(text='{"error":{"code":"session_retired"}}', content_type="application/json") + url = self.bindings.get(sid) if sid is not None else None + if url is not None and url not in self.backends: + raise web.HTTPServiceUnavailable( + text='{"error":{"code":"session_backend_unavailable"}}', content_type="application/json" + ) + if url is None: + if not self.backends: + raise web.HTTPServiceUnavailable( + text='{"error":{"code":"no_active_backend"}}', content_type="application/json" + ) + if sid is not None and len(self.bindings) + len(self.retired) >= self.max_sessions: + raise web.HTTPServiceUnavailable( + text='{"error":{"code":"session_capacity_exhausted"}}', content_type="application/json" + ) + url = min(self.backends, key=lambda u: (self.inflight[u], self.assigned[u], u)) + if sid is not None: + self.bindings[sid] = url + self.assigned[url] += 1 + self.logger.info("Session affinity assigned session=%s backend=%s", sid, url) + self.inflight[url] += 1 + return url + + async def handle(self, request): + if request.method == "GET" and request.path == "/health": + return web.json_response({"active_backends": len(self.backends)}) + # Only forward inference APIs. Control-plane gate/version operations remain Ray calls. + if request.method == "GET" and request.path in {"/v1/models", "/models"}: + body, sid = b"", None + elif request.method == "POST" and request.path in { + "/v1/chat/completions", + "/v1/messages", + "/v1/messages/count_tokens", + }: + body = await request.read() + try: + data = await request.json() + sid = data.get("session_id") + if isinstance(sid, bool) or not isinstance(sid, (str, int)) or str(sid) == "": + raise ValueError("missing session_id") + sid = str(sid) + except (ValueError, TypeError, AttributeError): + raise web.HTTPBadRequest( + text='{"error":{"code":"valid_session_id_required"}}', content_type="application/json" + ) + else: + raise web.HTTPNotFound() + url = self.acquire(sid) # No await between selection, binding and load reservation. + downstream = None + try: + async with self.client.request( + request.method, + url + str(request.rel_url), + data=body, + headers=_headers(request.headers), + allow_redirects=False, + ) as upstream: + downstream = web.StreamResponse(status=upstream.status, headers=_headers(upstream.headers)) + await downstream.prepare(request) + async for chunk in upstream.content.iter_any(): + await downstream.write(chunk) + await downstream.write_eof() + return downstream + except (aiohttp.ClientError, OSError, asyncio.TimeoutError): + # Never infer worker death from one timeout and never replay an ambiguous POST. + if downstream is not None and downstream.prepared: + if request.transport is not None: + request.transport.close() + return downstream + raise web.HTTPBadGateway( + text='{"error":{"code":"session_backend_transport_error"}}', content_type="application/json" + ) + finally: + self.inflight[url] -= 1 + + async def start(self, host, port=0): + self.client = aiohttp.ClientSession( + timeout=aiohttp.ClientTimeout(total=None, sock_connect=30), + connector=aiohttp.TCPConnector(limit=0), + auto_decompress=False, + trust_env=False, + ) + app = web.Application(client_max_size=16 * 1024**2) + app.router.add_route("*", "/{path:.*}", self.handle) + self.runner = web.AppRunner(app, handler_cancellation=True, shutdown_timeout=10, access_log=None) + await self.runner.setup() + site = web.TCPSite(self.runner, host, port) + await site.start() + self.url = f"http://{host}:{self.runner.addresses[0][1]}" + return self.url + + async def stop(self): + if self.runner is not None: + await self.runner.cleanup() + if self.client is not None: + await self.client.close() + + +class SessionAffinityService: + """Dedicated event loop: controller's blocking Ray calls cannot stall HTTP.""" + + def __init__(self, host): + self.loop = asyncio.new_event_loop() + self.router = SessionAffinityRouter() + self.thread = threading.Thread(target=self.loop.run_forever, name="session-affinity", daemon=True) + self.thread.start() + try: + self.url = self.call(self.router.start(host)) + except BaseException: + self.close() + raise + + def call(self, coroutine): + return asyncio.run_coroutine_threadsafe(coroutine, self.loop).result(timeout=60) + + def invoke(self, method, *args): + async def run(): + return getattr(self.router, method)(*args) + + return self.call(run()) + + def close(self): + try: + self.call(self.router.stop()) + finally: + self.loop.call_soon_threadsafe(self.loop.stop) + self.thread.join(timeout=60) + if self.thread.is_alive(): + raise RuntimeError("session affinity thread did not stop") + self.loop.close() diff --git a/xtuner/v1/rl/rollout/controller.py b/xtuner/v1/rl/rollout/controller.py index 7e52dd5db4..a40148e016 100644 --- a/xtuner/v1/rl/rollout/controller.py +++ b/xtuner/v1/rl/rollout/controller.py @@ -326,14 +326,22 @@ def onload_weights(self, target_state: WorkerLifecycleState = WorkerLifecycleSta def onload_kvcache(self, target_state: WorkerLifecycleState = WorkerLifecycleState.ACTIVE): self._broadcast_to_workers("onload_kvcache", target_state) + def retire_route_sessions(self, session_ids: list[str]) -> None: + if self.proxy_manager is not None: + self.proxy_manager.retire_sessions(session_ids) + def shutdown(self): """Shut down all rollout workers tracked by the controller.""" self.health_manager.stop() - actors = self.registry.all_actors() - ray.get( - [actor.shutdown.remote(stop_session_server=True) for actor in actors], # type: ignore[attr-defined] - timeout=ROLLOUT_RAY_GET_TIMEOUT, - ) + try: + actors = self.registry.all_actors() + ray.get( + [actor.shutdown.remote(stop_session_server=True) for actor in actors], # type: ignore[attr-defined] + timeout=ROLLOUT_RAY_GET_TIMEOUT, + ) + finally: + if self.proxy_manager is not None: + self.proxy_manager.close() def _broadcast_to_workers(self, method_name: str, target_state: WorkerLifecycleState, **kwargs): workers = self.registry.get_target_state_workers(target_state) diff --git a/xtuner/v1/rl/rollout/proxy_manager.py b/xtuner/v1/rl/rollout/proxy_manager.py index 768f4acb6e..75697f74cc 100644 --- a/xtuner/v1/rl/rollout/proxy_manager.py +++ b/xtuner/v1/rl/rollout/proxy_manager.py @@ -1,6 +1,9 @@ import time from collections.abc import Iterable from typing import TYPE_CHECKING +from uuid import uuid4 + +import ray from xtuner.v1.rl.utils.misc import ( check_chat_completions, @@ -25,13 +28,22 @@ class RolloutProxyManager(RolloutWorkerLifecycleListener): _CHECK_INTERVAL_SECONDS = 3.0 def __init__(self, config: "RolloutConfig"): + self._turn_resume = getattr(config, "turn_resume_wait_timeout", None) is not None self._model_name = config.model_name self._routed_proxy_url = config.routed_proxy_url self._routed_proxy_admin_url = config.routed_proxy_admin_url self._registered_session_urls: set[str] = set() self._logger = get_logger(log_dir=config.worker_log_dir, tag="RolloutProxyManager") + self._affinity = None + if config.session_affinity: + from .affinity_router import SessionAffinityService + + self._affinity = SessionAffinityService(ray.util.get_node_ip_address()) def replace_registered_session_urls(self, session_urls: Iterable[str]) -> None: + if self._affinity is not None: + self._affinity.invoke("replace_backends", session_urls) + session_urls = [self._affinity.url] model_name = self._get_model_name() delete_from_routedapiproxy(self._routed_proxy_admin_url, model_name) self._registered_session_urls.clear() @@ -55,13 +67,32 @@ def on_worker_group_inactive(self, group: "WorkerGroup") -> None: """Remove inactive request entrypoints from routed API proxy.""" for worker in group.workers: if worker.is_request_entrypoint: - self._delete_session_url(worker.session_url) + if self._affinity is not None: + self._affinity.invoke("set_backend", worker.session_url, False) + else: + self._delete_session_url(worker.session_url) def on_worker_group_active(self, group: "WorkerGroup") -> None: """Register active request entrypoints to routed API proxy.""" for worker in group.workers: if worker.is_request_entrypoint: - self._register_session_url(worker.session_url) + if self._affinity is not None: + self._affinity.invoke("set_backend", worker.session_url, True) + else: + self._register_session_url(worker.session_url) + + def retire_sessions(self, session_ids) -> None: + if self._affinity is not None: + self._affinity.invoke("retire_sessions", session_ids) + + def close(self) -> None: + try: + for url in list(self._registered_session_urls): + self._delete_session_url(url, raise_on_error=True) + finally: + if self._affinity is not None: + self._affinity.close() + self._affinity = None def _register_session_url( self, @@ -107,10 +138,28 @@ def _validate_routed_proxy(self) -> None: self._logger.info( f"Validating routed API proxy chat completions: url={self._routed_proxy_url}, model={model_name}" ) - for attempt in range(1, self._CHECK_ATTEMPTS + 1): - if check_chat_completions(self._routed_proxy_url, model_name): + attempts = 1 if self._turn_resume else self._CHECK_ATTEMPTS + for attempt in range(1, attempts + 1): + if self._turn_resume: + from .trace_store import get_existing_store + + session_id = str(uuid4()) + try: + succeeded = check_chat_completions( + self._routed_proxy_url, model_name, max_tokens=1024, session_id=session_id + ) + finally: + store = get_existing_store() + try: + if store is not None: + ray.get(store.release_sessions.remote([session_id])) + finally: + self.retire_sessions([session_id]) + else: + succeeded = check_chat_completions(self._routed_proxy_url, model_name) + if succeeded: return - if attempt < self._CHECK_ATTEMPTS: + if attempt < attempts: self._logger.warning( f"check chat completions failed for routed API proxy {self._routed_proxy_url}, " f"retrying {attempt}/{self._CHECK_ATTEMPTS - 1} after {self._CHECK_INTERVAL_SECONDS}s" diff --git a/xtuner/v1/rl/rollout/worker.py b/xtuner/v1/rl/rollout/worker.py index 3fa74345ac..4a8954de06 100644 --- a/xtuner/v1/rl/rollout/worker.py +++ b/xtuner/v1/rl/rollout/worker.py @@ -434,6 +434,10 @@ class RolloutConfig(BaseModel): help="Number of consecutive health check failures required before marking a worker inactive.", ), ] = 3 + session_affinity: bool = Field( + default=False, + description="Pin external HTTP sessions to one concrete DP SessionServer.", + ) enable_proxy: Annotated[ bool, Parameter( diff --git a/xtuner/v1/rl/utils/misc.py b/xtuner/v1/rl/utils/misc.py index 3f8903c78e..d6f37e9924 100644 --- a/xtuner/v1/rl/utils/misc.py +++ b/xtuner/v1/rl/utils/misc.py @@ -305,7 +305,7 @@ def _post(url: str, payload: dict) -> dict: return json.loads(resp.read()) -def check_chat_completions(base_url: str, model: str) -> bool: +def check_chat_completions(base_url: str, model: str, *, max_tokens: int = 16, session_id: str | None = None) -> bool: normalized_base_url = base_url.rstrip("/") if normalized_base_url.endswith("/v1"): url = f"{normalized_base_url}/chat/completions" @@ -313,9 +313,9 @@ def check_chat_completions(base_url: str, model: str) -> bool: url = f"{normalized_base_url}/v1/chat/completions" payload = { "model": model, - "session_id": str(uuid.uuid4()), + "session_id": session_id or str(uuid.uuid4()), "messages": [{"role": "user", "content": "Reply with exactly: pong"}], - "max_tokens": 16, + "max_tokens": max_tokens, "temperature": 0.0, "extra_body": {"spaces_between_special_tokens": False}, } From e226e976cae49c6afe8b9070963d421804128c88 Mon Sep 17 00:00:00 2001 From: Codex Date: Wed, 11 Nov 2026 12:18:03 +0000 Subject: [PATCH 4/5] feat(rl): retain Agent tasks across colocated and disaggregated windows Mechanical extraction from XTuner #2257 fixed head 43d62b6. No new implementation. Review commit only: protocol, affinity, lifecycle, Trainer guards and cold-start wiring form one atomic core merge group. Do not enable or merge this intermediate commit alone. --- .../agent_loop_manager/agent_loop_manager.py | 50 ++++++- .../agent_turn_pause_policy.py | 50 +++++++ .../agent_loop_manager/colocate_keepalive.py | 127 ++++++++++++++++++ .../disagg_agent_loop_manager.py | 96 ++++++++++--- .../rl/agent_loop_manager/disagg_producer.py | 29 +++- .../v1/rl/agent_loop_manager/produce_utils.py | 8 ++ xtuner/v1/rl/agent_loop_manager/producer.py | 30 +++-- 7 files changed, 354 insertions(+), 36 deletions(-) create mode 100644 xtuner/v1/rl/agent_loop_manager/agent_turn_pause_policy.py create mode 100644 xtuner/v1/rl/agent_loop_manager/colocate_keepalive.py diff --git a/xtuner/v1/rl/agent_loop_manager/agent_loop_manager.py b/xtuner/v1/rl/agent_loop_manager/agent_loop_manager.py index 71b1f3caa8..eeb1f27752 100644 --- a/xtuner/v1/rl/agent_loop_manager/agent_loop_manager.py +++ b/xtuner/v1/rl/agent_loop_manager/agent_loop_manager.py @@ -127,6 +127,10 @@ class AgentLoopManagerConfig(BaseModel): model_config = ConfigDict(extra="forbid", arbitrary_types_allowed=True) tasks: list[TaskSpecConfig] | TaskSpecConfig + turn_keepalive: bool = False + turn_keepalive_freeze_timeout: float = Field(default=30, gt=0) + turn_keepalive_drain_timeout: float = Field(default=450, gt=0) + turn_keepalive_quiesce_timeout: float = Field(default=3600, gt=0) def build( self, @@ -137,6 +141,15 @@ def build( sync_weights_interval: int = 1, rollout_teacher_scorer_config: RolloutTeacherScorerConfig | None = None, ) -> AgentLoopManager: + if self.turn_keepalive: + from .colocate_keepalive import build_keepalive_manager + + return build_keepalive_manager( + self, rollout_controller=rollout_controller, tokenizer=tokenizer, + replay_buffer=replay_buffer, logger=logger, + sync_weights_interval=sync_weights_interval, + rollout_teacher_scorer_config=rollout_teacher_scorer_config, + ) tasks = self.tasks if isinstance(self.tasks, list) else [self.tasks] if not tasks: raise ValueError("AgentLoopManagerConfig requires at least one task config.") @@ -221,6 +234,7 @@ async def produce_batch( train_step: int, *, model_step: int, + quiesce: bool = False, ) -> ProduceBatchResult: # 共卡同步入口：生产入 buffer -> pause/drain 本轮 pending -> 取非空训练 batch。 if batch_size <= 0: @@ -233,7 +247,12 @@ async def produce_batch( active_tasks = [task for task in self.task_runners if current_sizes[task.task_name] > 0] assert active_tasks, "No active tasks found" - await self._rollout_controller.continue_generation.remote() # type: ignore[attr-defined] + keepalive = getattr(self, "_keepalive", None) + if keepalive is None: + await self._rollout_controller.continue_generation.remote() + else: + await keepalive.resume(self._rollout_controller, model_step) + await keepalive.check_requests(self._rollout_controller) local_progress = ProduceProgress.build( task_names=self.task_names, target_samples=current_sizes, @@ -271,12 +290,37 @@ async def produce_batch( await asyncio.gather(*produce_futures) local_progress.add_produce_time(time.perf_counter() - produce_start) - # pause 只收尾本轮本地 pending。 - await self._rollout_controller.pause_generation.remote() # type: ignore[attr-defined] + # Checkpoint/final quiesce happens while gates are open. Never wait for + # an entire episode behind a sealed request gate. + if keepalive is not None: + if quiesce: + async def finish_live_tasks(): + for task in self.task_runners: + ctx = ProduceContext( + agent_loop=task.agent_loop, sampler=task.sampler, + replay_buffer=self.replay_buffer, task_batch_size=0, + task_name=task.task_name, train_step=train_step, + model_step=model_step, progress=local_progress, + stale_threshold=task.stale_threshold, + token_stale_threshold=task.token_stale_threshold, + expired_groups_retryable=task.expired_groups_retryable, + ) + owner = task.produce_strategy._pending_tasks + while owner.count(): + for ready in await owner.wait_and_claim(timeout_s=1): + await ctx.put_generated_group(ready.result()) + await asyncio.wait_for(finish_live_tasks(), keepalive.quiesce_timeout) + await keepalive.pause(self._rollout_controller) + else: + await self._rollout_controller.pause_generation.remote() pause_time_s = 0.0 for task in active_tasks: produce_strategy = cast(ProduceStrategy, task.produce_strategy) + if keepalive is not None: + # Live futures stay in the strategy's existing _PendingTasks. + # Done futures are harvested on the next window (or quiesce). + continue pause_time_s += await produce_strategy.pause_produce( ProduceContext( agent_loop=task.agent_loop, diff --git a/xtuner/v1/rl/agent_loop_manager/agent_turn_pause_policy.py b/xtuner/v1/rl/agent_loop_manager/agent_turn_pause_policy.py new file mode 100644 index 0000000000..6b6ce4c7c5 --- /dev/null +++ b/xtuner/v1/rl/agent_loop_manager/agent_turn_pause_policy.py @@ -0,0 +1,50 @@ +"""Disaggregated pause policy; PendingTasks remains the sole live-task owner.""" + +from .produce_utils import pause_pending_tasks, retain_pending_tasks + + +class DrainPendingPausePolicy: + pending_cleanup = staticmethod(pause_pending_tasks) + + async def check_requests(self, controller): + return + + async def before_sync(self, manager, checkpoint_due): + return + + async def pause(self, controller): + await controller.pause_generation.remote() + + async def resume(self, controller, model_step): + await controller.continue_generation.remote() + + +class AgentTurnPausePolicy: + pending_cleanup = staticmethod(retain_pending_tasks) + + def __init__(self, *, freeze_timeout: float, drain_timeout: float, quiesce_timeout: float = 1800): + self.quiesce_timeout = quiesce_timeout + self.freeze_timeout = freeze_timeout + self.drain_timeout = drain_timeout + self.barrier = None + + async def check_requests(self, controller): + await controller.assert_request_gates_healthy.remote() + + async def before_sync(self, manager, checkpoint_due): + if checkpoint_due: + await manager.quiesce(self.quiesce_timeout) + + async def pause(self, controller): + self.barrier = await controller.prepare_pause.remote( + freeze_timeout=self.freeze_timeout, drain_timeout=self.drain_timeout + ) + await controller.pause_generation.remote() + + async def resume(self, controller, model_step): + # Initial construction has no admitted requests and no barrier yet. + if self.barrier is None: + await controller.continue_generation.remote() + else: + await controller.resume_requests.remote(self.barrier, model_step, timeout=self.freeze_timeout) + self.barrier = None diff --git a/xtuner/v1/rl/agent_loop_manager/colocate_keepalive.py b/xtuner/v1/rl/agent_loop_manager/colocate_keepalive.py new file mode 100644 index 0000000000..f147b9edaa --- /dev/null +++ b/xtuner/v1/rl/agent_loop_manager/colocate_keepalive.py @@ -0,0 +1,127 @@ +"""Opt-in persistent asyncio owner for alternating colocated RL. + +The native asyncio_run loop is stable but stops during synchronous training. +This adapter keeps tool/client tasks progressing without enabling GPU inference. +Only the strategy's existing _PendingTasks owns live generation tasks. +""" + +import asyncio +import threading + + +class LoopThread: + def __init__(self): + self.loop = asyncio.new_event_loop() + self._ready = threading.Event() + self.thread = threading.Thread(target=self._run, name="xt-colocate-agent-loop", daemon=True) + self.thread.start() + self._ready.wait() + + def _run(self): + asyncio.set_event_loop(self.loop) + self._ready.set() + self.loop.run_forever() + self.loop.close() + + def submit(self, coro): + if not self.thread.is_alive(): + coro.close() + raise RuntimeError("Colocate Agent loop is closed") + return asyncio.run_coroutine_threadsafe(coro, self.loop) + + async def call(self, coro): + return await asyncio.wrap_future(self.submit(coro)) + + def stop(self): + self.loop.call_soon_threadsafe(self.loop.stop) + self.thread.join(timeout=30) + if self.thread.is_alive(): + raise TimeoutError("Colocate Agent loop thread did not stop") + + +class ColocateKeepaliveManager: + def __init__(self, inner, owner): + self._inner, self._owner = inner, owner + self._closed = False + + def __getattr__(self, name): + return getattr(self._inner, name) + + async def produce_batch(self, *args, **kwargs): + return await self._owner.call(self._inner.produce_batch(*args, **kwargs)) + + async def prepare_initial_sync(self): + # Servers were born closed, before their HTTP endpoint was published. + # Build the ordinary fixed-topology barrier; no group is scheduled. + return await self._owner.call(self._inner._keepalive.pause(self._inner._rollout_controller)) + + async def publish_initial_version(self, model_step): + return await self._owner.call(self._inner._keepalive.resume(self._inner._rollout_controller, model_step)) + + async def save(self, *args, **kwargs): + return await self._owner.call(self._inner.save(*args, **kwargs)) + + async def resume(self, *args, **kwargs): + return await self._owner.call(self._inner.resume(*args, **kwargs)) + + async def close(self): + if self._closed: + return + self._closed = True + + async def cleanup(): + cancellations = [ + asyncio.create_task(task.produce_strategy._pending_tasks.cancel_all()) + for task in self._inner.task_runners + ] + if cancellations: + _, pending = await asyncio.wait(cancellations, timeout=120) + if pending: + raise TimeoutError("Pending Agent cleanup exceeded 120 seconds") + # Local AgentInSandboxLoop cancellation reaches original Runner + # finally/pool cleanup. No duplicate sandbox owner is introduced. + remaining = [t for t in asyncio.all_tasks() if t is not asyncio.current_task()] + for task in remaining: + task.cancel() + if remaining: + _, pending = await asyncio.wait(remaining, timeout=30) + if pending: + raise TimeoutError("Agent client teardown exceeded 30 seconds") + + try: + await self._owner.call(cleanup()) + finally: + self._owner.stop() + + +def build_keepalive_manager(config, **kwargs): + from .agent_turn_pause_policy import AgentTurnPausePolicy + from .producer import AsyncProduceStrategyConfig + + tasks = config.tasks if isinstance(config.tasks, list) else [config.tasks] + for task in tasks: + strategy = task.produce_strategy_config + if not isinstance(strategy, AsyncProduceStrategyConfig) or strategy.enable_partial_rollout: + raise ValueError("Colocate turn keepalive requires AsyncProduceStrategy and token partial disabled") + if not getattr(task.agent_loop_config, "require_turn_provenance", False): + raise ValueError("Colocate turn keepalive requires real request provenance") + owner = LoopThread() + + async def build(): + # Build clients, locks, sampler and tasks on the one persistent loop. + inner = config.model_copy(update={"turn_keepalive": False}).build(**kwargs) + inner._keepalive = AgentTurnPausePolicy( + freeze_timeout=config.turn_keepalive_freeze_timeout, + drain_timeout=config.turn_keepalive_drain_timeout, + quiesce_timeout=config.turn_keepalive_quiesce_timeout, + ) + for task in inner.task_runners: + task.produce_strategy.retain_live_tasks = True + return inner + + try: + inner = owner.submit(build()).result() + except BaseException: + owner.stop() + raise + return ColocateKeepaliveManager(inner, owner) diff --git a/xtuner/v1/rl/agent_loop_manager/disagg_agent_loop_manager.py b/xtuner/v1/rl/agent_loop_manager/disagg_agent_loop_manager.py index 70ae8731b3..bd4cc9eb23 100644 --- a/xtuner/v1/rl/agent_loop_manager/disagg_agent_loop_manager.py +++ b/xtuner/v1/rl/agent_loop_manager/disagg_agent_loop_manager.py @@ -17,6 +17,7 @@ from xtuner.v1.rl.rollout import RolloutController from xtuner.v1.utils import get_logger +from .agent_turn_pause_policy import AgentTurnPausePolicy, DrainPendingPausePolicy from .disagg_producer import ( DisaggAsyncProduceStrategyConfig, DisaggProduceContext, @@ -67,6 +68,9 @@ class DisaggAgentLoopManagerConfig(BaseModel): model_config = ConfigDict(extra="forbid", arbitrary_types_allowed=True) tasks: list[DisaggTaskSpecConfig] | DisaggTaskSpecConfig + turn_resume_drain_timeout: float | None = Field(default=None, gt=0) + turn_resume_freeze_timeout: float = Field(default=60, gt=0) + turn_resume_quiesce_timeout: float = Field(default=1800, gt=0) def build( self, @@ -112,6 +116,15 @@ def build( ) return DisaggAgentLoopManager( + pause_policy=( + AgentTurnPausePolicy( + quiesce_timeout=self.turn_resume_quiesce_timeout, + freeze_timeout=self.turn_resume_freeze_timeout, + drain_timeout=self.turn_resume_drain_timeout, + ) + if self.turn_resume_drain_timeout is not None + else DrainPendingPausePolicy() + ), task_runners=task_runners, replay_buffer=replay_buffer, rollout_controller=rollout_controller, @@ -174,12 +187,17 @@ def __init__( replay_buffer: ReplayBuffer, rollout_controller: RolloutController, logger=None, + pause_policy=None, ): if not task_runners: raise ValueError("DisaggAgentLoopManager requires at least one task runner.") if sum(task.weight for task in task_runners) <= 0: raise ValueError("At least one task weight must be positive for DisaggAgentLoopManager.") + self._pause_policy = pause_policy or DrainPendingPausePolicy() + if isinstance(self._pause_policy, AgentTurnPausePolicy): + for task in task_runners: + task.produce_strategy.configure_turn_resume(self._pause_policy.pending_cleanup) self.task_runners = task_runners self.replay_buffer = replay_buffer self._rollout_controller = rollout_controller @@ -195,6 +213,10 @@ def __init__( # consumer 同步权重前置位；producer / strategy 直接观察 event。 self._update_event = asyncio.Event() self._finish_event = asyncio.Event() + self._group_admission = asyncio.Event() + self._group_admission.set() + self._producer_idle = asyncio.Event() + self._producer_idle.set() # rollout 侧当前模型版本；pause 清空 pending 后才能更新。 self._model_step = 0 @@ -260,29 +282,57 @@ async def _produce_batch_to_buffer( progress.add_produce_time(time.perf_counter() - produce_start) return produce_status + def _pending_context(self, task): + return DisaggProduceContext( + agent_loop=task.agent_loop, + sampler=task.sampler, + replay_buffer=self.replay_buffer, + task_batch_size=0, + task_name=task.task_name, + train_step=self._produce_progress.producer_future_step, + model_step=self._model_step, + progress=self._produce_progress, + update_event=self._update_event, + stale_threshold=task.stale_threshold, + token_stale_threshold=task.token_stale_threshold, + expired_groups_retryable=task.expired_groups_retryable, + ) + + async def prepare_sync(self, checkpoint_due: bool): + await self._pause_policy.before_sync(self, checkpoint_due) + + async def quiesce(self, timeout: float): + """Stop new groups, join producer claims, finish episodes with gates open.""" + self._group_admission.clear() + for task in self.task_runners: + task.produce_strategy.scheduling_paused = True + async with asyncio.timeout(timeout): + await self._producer_idle.wait() + await asyncio.gather( + *[task.produce_strategy.quiesce(self._pending_context(task)) for task in self.task_runners] + ) + if get_pending_task_counts(self.task_runners): + raise RuntimeError("Quiesce left pending tasks") + + async def close(self): + self.shutdown() + results = await asyncio.gather( + *[task.produce_strategy.close() for task in self.task_runners], return_exceptions=True + ) + for result in results: + if isinstance(result, BaseException): + raise result + async def pause_produce(self) -> float: # 非共卡显式刹车；共卡没有 public pause。 self._status = AgentLoopManagerStatus.UPDATE_WEIGHT_AND_ABORT self._update_event.set() - await self._rollout_controller.pause_generation.remote() # type: ignore[attr-defined] + await self._pause_policy.pause(self._rollout_controller) pause_time_s = 0.0 for task in self.task_runners: produce_strategy = cast(DisaggProduceStrategy, task.produce_strategy) - ctx = DisaggProduceContext( - agent_loop=task.agent_loop, - sampler=task.sampler, - replay_buffer=self.replay_buffer, - task_batch_size=0, - task_name=task.task_name, - train_step=self._produce_progress.producer_future_step, - model_step=self._model_step, - progress=self._produce_progress, - update_event=self._update_event, - stale_threshold=task.stale_threshold, - token_stale_threshold=task.token_stale_threshold, - expired_groups_retryable=task.expired_groups_retryable, - ) + ctx = self._pending_context(task) pause_time_s += await produce_strategy.pause_produce(ctx) self._pause_time_s = pause_time_s return pause_time_s @@ -290,14 +340,18 @@ async def pause_produce(self) -> float: async def continue_produce(self, model_step: int) -> None: # 与 pause_produce 成对：同步/评测完成后，用新 model_step 恢复后台 producer。 self._model_step = model_step - await self._rollout_controller.continue_generation.remote() # type: ignore[attr-defined] + await self._pause_policy.resume(self._rollout_controller, model_step) self._status = AgentLoopManagerStatus.NORMAL self._update_event.clear() + for task in self.task_runners: + task.produce_strategy.scheduling_paused = False + self._group_admission.set() def shutdown(self) -> None: self._status = AgentLoopManagerStatus.FINISH self._update_event.set() self._finish_event.set() + self._group_admission.set() async def _wait_for_status_exit(self, blocked_status: AgentLoopManagerStatus) -> None: while not self._finish_event.is_set() and self._status == blocked_status: @@ -306,6 +360,7 @@ async def _wait_for_status_exit(self, blocked_status: AgentLoopManagerStatus) -> async def produce_loop(self, batch_size: int) -> None: # 后台持续生产；前台通过 get_batch 消费。 while not self._finish_event.is_set(): + await self._group_admission.wait() if self._status == AgentLoopManagerStatus.FINISH: break if self._status in (AgentLoopManagerStatus.UPDATE_WEIGHT_AND_ABORT, AgentLoopManagerStatus.EXPIRED_BATCH): @@ -322,12 +377,16 @@ async def produce_loop(self, batch_size: int) -> None: future_step, ), ) - produce_status = await self._produce_batch_to_buffer(task_batch_sizes, self._produce_progress) + self._producer_idle.clear() + try: + produce_status = await self._produce_batch_to_buffer(task_batch_sizes, self._produce_progress) + finally: + self._producer_idle.set() if produce_status == ProduceBatchStatus.EXPIRED_BATCH: # EXPIRED_BATCH 是 producer 自己检测出来的“立即停下”信号。 self._status = AgentLoopManagerStatus.EXPIRED_BATCH - elif produce_status == ProduceBatchStatus.NORMAL: + elif produce_status == ProduceBatchStatus.NORMAL and self._group_admission.is_set(): # 只有正常完成一轮生产时，producer 自己维护的 train_step 才前进一步。 self._produce_progress.advance_future_step() @@ -350,6 +409,7 @@ async def get_batch(self, batch_size: int, train_step: int) -> ProduceBatchResul current_model_step = train_step - 1 while not self._finish_event.is_set(): + await self._pause_policy.check_requests(self._rollout_controller) if self._status == AgentLoopManagerStatus.EXPIRED_BATCH: if current_model_step > self._model_step: pause_time_s = self._consume_pause_time() diff --git a/xtuner/v1/rl/agent_loop_manager/disagg_producer.py b/xtuner/v1/rl/agent_loop_manager/disagg_producer.py index 59a6762612..ca080206c3 100644 --- a/xtuner/v1/rl/agent_loop_manager/disagg_producer.py +++ b/xtuner/v1/rl/agent_loop_manager/disagg_producer.py @@ -321,12 +321,22 @@ async def produce_batch(self, ctx: DisaggProduceContext) -> ProduceBatchStatus: async def pause_produce(self, ctx: DisaggProduceContext) -> float: return 0.0 + async def close(self): + return + def is_model_expired(self, train_step: int, model_step: int) -> bool: return False def pending_task_count(self) -> int: return 0 + def configure_turn_resume(self, pending_cleanup) -> None: + """Opt in only when the strategy owns and can quiesce retained tasks.""" + raise ValueError(f"{type(self).__name__} does not support complete-turn resume") + + async def quiesce(self, ctx: DisaggProduceContext) -> None: + raise NotImplementedError(f"{type(self).__name__} does not support quiesce") + class DisaggAsyncProduceStrategy(DisaggProduceStrategy): """非共卡 async strategy；pending 跨后台生产轮次存在。""" @@ -361,21 +371,34 @@ def __init__( ) self.tail_batch_trigger_size = tail_batch_trigger_size self._pending_tasks = _PendingTasks() + self.pending_cleanup = pause_pending_tasks + self.scheduling_paused = False def is_model_expired(self, train_step: int, model_step: int) -> bool: staleness = calculate_seq_staleness(model_step, train_step) return staleness >= self.stale_threshold + def configure_turn_resume(self, pending_cleanup) -> None: + self.pending_cleanup = pending_cleanup + def pending_task_count(self) -> int: return self._pending_tasks.count() async def pause_produce(self, ctx: DisaggProduceContext) -> float: - return await pause_pending_tasks( + return await self.pending_cleanup( pending_tasks=self._pending_tasks, ctx=ctx, put_claimed_task=lambda task: ctx.put_generated_group(task.result()), ) + async def quiesce(self, ctx: DisaggProduceContext): + while self._pending_tasks.count(): + claimed = await self._pending_tasks.wait_and_claim(timeout_s=1) + await _put_claimed_tasks(claimed, ctx) + + async def close(self): + await self._pending_tasks.cancel_all() + async def produce_batch(self, ctx: DisaggProduceContext) -> ProduceBatchStatus: if ctx.task_name not in ctx.progress.consumed_samples: raise KeyError(f"DisaggProduceProgress.consumed_samples missing task_name={ctx.task_name!r}") @@ -439,6 +462,8 @@ async def spawn_one() -> asyncio.Task: ) produce_status = ProduceBatchStatus.NORMAL while True: + if self.scheduling_paused: + break if ctx.should_abort(): produce_status = ProduceBatchStatus.UPDATE_WEIGHT_AND_ABORT break @@ -456,7 +481,7 @@ async def spawn_one() -> asyncio.Task: if available + pending_count < scheduled_target: while await self._pending_tasks.schedule_one( max_pending=desired_pending, - should_abort=ctx.should_abort, + should_abort=lambda: self.scheduling_paused or ctx.should_abort(), spawn_one=spawn_one, ): pass diff --git a/xtuner/v1/rl/agent_loop_manager/produce_utils.py b/xtuner/v1/rl/agent_loop_manager/produce_utils.py index 053cca332c..00ff5db452 100644 --- a/xtuner/v1/rl/agent_loop_manager/produce_utils.py +++ b/xtuner/v1/rl/agent_loop_manager/produce_utils.py @@ -785,6 +785,14 @@ async def request_agent_loop_pause(ctx: BaseProduceContext, *, pending_count: in ) +async def retain_pending_tasks(*, pending_tasks, ctx, put_claimed_task) -> float: + """Claim completed groups once; live Agent tasks keep their original owner.""" + start = time.perf_counter() + for task in await pending_tasks.claim_ready(): + await put_claimed_task(task) + return time.perf_counter() - start + + async def pause_pending_tasks( *, pending_tasks: set[asyncio.Task] | _PendingTasks, diff --git a/xtuner/v1/rl/agent_loop_manager/producer.py b/xtuner/v1/rl/agent_loop_manager/producer.py index 7f6186e116..4ec6bdeff1 100644 --- a/xtuner/v1/rl/agent_loop_manager/producer.py +++ b/xtuner/v1/rl/agent_loop_manager/producer.py @@ -15,6 +15,7 @@ PERIODIC_ABORT_INTERVAL_S, BaseProduceContext, ShouldContinueFn, + _PendingTasks, _ProgressDisplayer, _put_claimed_tasks, calculate_stale_threshold, @@ -360,14 +361,15 @@ def __init__( ) self.tail_batch_trigger_size = tail_batch_trigger_size self.max_pending_tasks = max_pending_tasks - self._local_pending_tasks: set[asyncio.Task] = set() + self._pending_tasks = _PendingTasks() + self.retain_live_tasks = False def pending_task_count(self) -> int: - return len(self._local_pending_tasks) + return self._pending_tasks.count() async def pause_produce(self, ctx: ProduceContext) -> float: return await pause_pending_tasks( - pending_tasks=self._local_pending_tasks, + pending_tasks=self._pending_tasks, ctx=ctx, put_claimed_task=lambda task: ctx.put_generated_group(task.result()), ) @@ -376,8 +378,11 @@ async def produce_batch(self, ctx: ProduceContext) -> None: if ctx.task_name not in ctx.progress.target_samples: raise KeyError(f"ProduceProgress.target_samples missing task_name={ctx.task_name!r}") - # 共卡 async 的 pending 只属于本次 produce_batch。 - self._local_pending_tasks = set() + # Opt-in keepalive retains the same owner across production windows. + if not self.retain_live_tasks: + self._pending_tasks = _PendingTasks() + else: + await _put_claimed_tasks(await self._pending_tasks.claim_ready(), ctx) if ctx.batch_target <= 0: return @@ -424,22 +429,21 @@ async def spawn_one() -> asyncio.Task: if not self.should_continue_fn(available, batch_target): break - pending_count = len(self._local_pending_tasks) + pending_count = self._pending_tasks.count() desired_pending = max(0, scheduled_target - available) if self.max_pending_tasks is not None: desired_pending = min(desired_pending, self.max_pending_tasks) if pending_count < desired_pending: - while len(self._local_pending_tasks) < desired_pending: - self._local_pending_tasks.add(await spawn_one()) + while await self._pending_tasks.schedule_one( + max_pending=desired_pending, should_abort=lambda: False, spawn_one=spawn_one + ): + pass - if not self._local_pending_tasks: + if not self._pending_tasks.count(): logger.warning("All tasks are done but not enough samples collected.") break - done_tasks, _ = await asyncio.wait( - set(self._local_pending_tasks), timeout=1, return_when=asyncio.FIRST_COMPLETED - ) - self._local_pending_tasks.difference_update(done_tasks) + done_tasks = await self._pending_tasks.wait_and_claim(timeout_s=1) await _put_claimed_tasks( done_tasks, ctx, From f5ee2bebdfd466467c032f62f246cfc72afb2cb8 Mon Sep 17 00:00:00 2001 From: Codex Date: Wed, 11 Nov 2026 12:18:03 +0000 Subject: [PATCH 5/5] feat(rl): wire quiescent checkpoints and sealed cold publication Mechanical extraction from XTuner #2257 fixed head 43d62b6. No new implementation. Review commit only: protocol, affinity, lifecycle, Trainer guards and cold-start wiring form one atomic core merge group. Do not enable or merge this intermediate commit alone. --- .../agent_loop_manager/agent_loop_manager.py | 2 + xtuner/v1/train/rl_trainer.py | 137 +++++++++++++++++- 2 files changed, 131 insertions(+), 8 deletions(-) diff --git a/xtuner/v1/rl/agent_loop_manager/agent_loop_manager.py b/xtuner/v1/rl/agent_loop_manager/agent_loop_manager.py index eeb1f27752..03679d5b6a 100644 --- a/xtuner/v1/rl/agent_loop_manager/agent_loop_manager.py +++ b/xtuner/v1/rl/agent_loop_manager/agent_loop_manager.py @@ -401,6 +401,8 @@ async def resume(self, checkpoint_path: Path | str) -> int: state_path = manager_state_path(checkpoint_path) with state_path.open("r") as f: manager_state = json.load(f) + if getattr(self, "_keepalive", None) is not None and not manager_state.get("replay_buffer_saved", True): + raise RuntimeError("Full colocate resume requires a persisted replay buffer") if manager_state.get("replay_buffer_saved", True): # replay buffer 恢复是 async I/O，不能在已有 event loop 中再次嵌套 asyncio_run。 await self.replay_buffer.resume(checkpoint_path) diff --git a/xtuner/v1/train/rl_trainer.py b/xtuner/v1/train/rl_trainer.py index 00efb8c37e..32c8dede03 100644 --- a/xtuner/v1/train/rl_trainer.py +++ b/xtuner/v1/train/rl_trainer.py @@ -2,6 +2,7 @@ import json import os import re +import sys import time from dataclasses import asdict, dataclass from pathlib import Path @@ -446,6 +447,29 @@ class RLColocateTrainerConfig(BaseRLTrainerConfig): agent_loop_manager_cfg: AgentLoopManagerConfig resources: AcceleratorResourcesConfig + @model_validator(mode="after") + def _validate_turn_keepalive(self): + if self.agent_loop_manager_cfg.turn_keepalive != (self.rollout_config.turn_resume_wait_timeout is not None): + raise ValueError("Colocate turn keepalive requires both request gates and the keepalive manager") + if self.agent_loop_manager_cfg.turn_keepalive: + if self.rollout_config.health_check_failure_threshold != 0: + raise ValueError("Colocate turn keepalive requires fixed topology: disable automatic worker recovery") + if self.enable_evaluate or self.enable_initial_evaluate or self.eval_agent_loop_manager_cfg is not None: + raise ValueError("Colocate turn keepalive requires all online evaluation disabled") + if self.rollout_config.turn_resume_wait_timeout is None: + raise ValueError("Colocate turn keepalive requires request gates") + self.rollout_config.turn_resume_initially_closed = True + if (self.debug_rollout or self.debug_train) and ( + self.load_checkpoint_cfg.checkpoint_path is not None or self.auto_resume + ): + raise ValueError("Use full online run for colocate cold resume; debug modes start fresh") + if self.auto_resume and self.load_checkpoint_cfg.checkpoint_path is not None: + raise ValueError("Choose explicit checkpoint_path or auto_resume, not both") + if self.load_checkpoint_cfg.checkpoint_path is not None or self.auto_resume: + if not (self.load_checkpoint_cfg.load_optimizer_states and self.load_checkpoint_cfg.load_optimizer_args): + raise ValueError("Full colocate resume requires optimizer states") + return self + def build(self) -> "RLColocateTrainer": return RLColocateTrainer(self) @@ -538,6 +562,29 @@ class RLDisaggregatedTrainerConfig(BaseRLTrainerConfig): train_resources: AcceleratorResourcesConfig rollout_resources: AcceleratorResourcesConfig + @model_validator(mode="after") + def _validate_turn_resume(self): + enabled = self.rollout_config.turn_resume_wait_timeout is not None + manager_enabled = self.agent_loop_manager_cfg.turn_resume_drain_timeout is not None + if enabled != manager_enabled: + raise ValueError("turn-resume requires both request gates and Agent pause policy") + if enabled: + # A restored engine must never expose version-zero admission while + # actor/optimizer and manager state are still loading. + if self.load_checkpoint_cfg.checkpoint_path is not None or self.auto_resume: + self.rollout_config.turn_resume_initially_closed = True + if self.enable_evaluate or self.enable_initial_evaluate or self.eval_agent_loop_manager_cfg is not None: + raise ValueError("turn-resume first version rejects every online evaluation entrypoint") + if self.rollout_config.health_check_failure_threshold != 0: + raise ValueError("turn-resume requires fixed topology: disable automatic worker recovery") + tasks = self.agent_loop_manager_cfg.tasks + for task in tasks if isinstance(tasks, list) else [tasks]: + if not getattr(task.agent_loop_config, "require_turn_provenance", False): + raise ValueError("Agent turn-resume requires complete request provenance and MoE routes") + if getattr(task.produce_strategy_config, "enable_partial_rollout", False): + raise ValueError("Complete-turn pause cannot use token partial rollout") + return self + def build(self) -> "RLDisaggregatedTrainer": return RLDisaggregatedTrainer(self) @@ -771,12 +818,13 @@ async def _resume_agent_loop_manager(self, checkpoint_path: Path | str) -> int: saved_model_step = await self.agent_loop_manager.resume(checkpoint_path) return saved_model_step + def _checkpoint_due(self, cur_step: int) -> bool: + interval = self._checkpoint_interval + return interval is not None and interval > 0 and cur_step % interval == 0 + async def _maybe_save_checkpoint(self, cur_step: int) -> None: """Save checkpoint if interval condition is met.""" - ckp_interval = self._checkpoint_interval - if ckp_interval is None or ckp_interval == -1: - return - if cur_step % ckp_interval != 0: + if not self._checkpoint_due(cur_step): return checkpoint_path = self.exp_dir / self._CHECKPOINT_DIR / f"ckpt-step-{cur_step}" @@ -1430,12 +1478,14 @@ def _log_mini_batch_metrics(self, workers_log_item: List[WorkerLogItem]): class RLColocateTrainer(BaseRLTrainer): + _turn_keepalive = False _META_PATH = ".xtuner_rl_colocate_trainer" agent_loop_manager: AgentLoopManager rl_health_manager: RLHealthManager | _NoOpRLHealthManager # 共卡保留资源切换和权重同步流程；通用保存、日志在 BaseRLTrainer。 def __init__(self, cfg: RLColocateTrainerConfig): + self._turn_keepalive = cfg.agent_loop_manager_cfg.turn_keepalive self._init_common(cfg, meta_path=self._META_PATH, logger_tag="RLTrainer") self._num_workers = float(cfg.resources.num_workers) self._rollout_num_workers = float(cfg.resources.num_workers) @@ -1458,6 +1508,8 @@ def __init__(self, cfg: RLColocateTrainerConfig): self._build_agent_loop_components(cfg, replay_buffer) self._cpu_resource_manager.log_registered_summary() self.logger.warning("Debug rollout mode is enabled. Only rollout workers will be started.") + if self._turn_keepalive: + self._initialize_keepalive_rollout(None) return @@ -1504,14 +1556,59 @@ def __init__(self, cfg: RLColocateTrainerConfig): replay_buffer = cfg.replay_buffer_config.build() self._build_agent_loop_components(cfg, replay_buffer) - if checkpoint_path is not None: + if self._turn_keepalive: + self._initialize_keepalive_rollout(checkpoint_path) + elif checkpoint_path is not None: asyncio_run(self._resume_agent_loop_manager(checkpoint_path)) self._cpu_resource_manager.log_registered_summary() - if self._rollout_config.skip_load_weights: + if not self._turn_keepalive and self._rollout_config.skip_load_weights: self._sync_weights_from_train_workers() + def _init_work_dir_and_meta(self, cfg, meta_path): + if self._turn_keepalive and cfg.auto_resume: + # Reuse native metadata and native selection, but write the resumed + # run into a new native exp_dir. Old HF/log/trajectory evidence must + # not be overwritten when replaying from an earlier saved step. + work_dir = Path(cfg.work_dir) if cfg.work_dir else Path.cwd() / "work_dirs" + previous_meta_path = work_dir / meta_path + if previous_meta_path.exists(): + previous = XTunerMeta.model_validate_json(previous_meta_path.read_text()) + # A prior resumed exp may have emitted HF but no new DCP. + # Consult only recorded metadata, newest experiment first; + # retain each ExpInfo.latest_checkpoint's native selection. + checkpoint_exp = next( + (exp for exp in reversed(previous.exps) if exp.latest_checkpoint is not None), None + ) + if checkpoint_exp is not None: + selected_meta = previous.model_copy(update={"exps": [checkpoint_exp]}) + self._meta = selected_meta + self._resolve_load_checkpoint_cfg(True, cfg.load_checkpoint_cfg) + super()._init_work_dir_and_meta(cfg.model_copy(update={"auto_resume": False}), meta_path) + else: + super()._init_work_dir_and_meta(cfg, meta_path) + + def _initialize_keepalive_rollout(self, checkpoint_path): + # No new groups before restored sampler/replay and rollout publication. + # Gate construction already closed admission before external visibility. + try: + asyncio_run(self.agent_loop_manager.prepare_initial_sync()) + restored_step = 0 + if checkpoint_path is not None: + restored_step = asyncio_run(self._resume_agent_loop_manager(checkpoint_path)) + if restored_step != self._cur_step: + raise RuntimeError("Restored manager model_step disagrees with train state") + if checkpoint_path is not None or self._rollout_config.skip_load_weights: + self._sync_weights_from_train_workers() + asyncio_run(self.agent_loop_manager.publish_initial_version(restored_step)) + except BaseException: + try: + asyncio_run(self.agent_loop_manager.close()) + except Exception: + self.logger.exception("Keepalive startup cleanup failed") + raise + def _sync_weights_from_train_workers(self) -> None: if self._rollout_config.weight_transport_type == "checkpoint_engine": ray.get(self.rollout_controller.offload.remote(), timeout=RL_TRAINER_RAY_GET_TIMEOUT) @@ -1538,9 +1635,20 @@ def fit(self): try: self._fit() finally: + original_error = sys.exc_info()[0] is not None + cleanup_error = None + if self._turn_keepalive and not self._debug_train: + try: + asyncio_run(self.agent_loop_manager.close()) + except Exception as exc: + cleanup_error = exc + self._turn_keepalive_cleanup_error = repr(exc) + self.logger.exception("Colocate keepalive cleanup failed") self.rl_health_manager.stop() self._exp_tracker.close() close_trace() + if cleanup_error is not None and not original_error: + raise cleanup_error def _fit(self): self.logger.info("Start RL training") @@ -1582,6 +1690,8 @@ def _fit(self): self.train_batch_size, train_step=train_step, model_step=model_step, + **({"quiesce": self._checkpoint_due(train_step) or train_step == self._total_train_steps} + if self._turn_keepalive else {}), ) ) finally: @@ -1793,6 +1903,12 @@ def _resume_from_checkpoint(self, checkpoint_path: Path | str) -> None: saved_model_step = asyncio_run(self._resume_agent_loop_manager(checkpoint_path)) assert self._cur_step == saved_model_step + # A fresh SessionServer starts at version zero. Seal requests before + # replacing weights so continue_produce publishes the restored version + # through the same barrier used by ordinary updates. + asyncio_run(self.agent_loop_manager.prepare_sync(False)) + asyncio_run(self.agent_loop_manager.pause_produce()) + # Empty-init engines can retain GPU allocations while marked asleep. # Use the same offload/allocate/receive/warmup lifecycle as online sync; # NCCL transfer alone does not clear the engine's weights sleep tag. @@ -1917,7 +2033,8 @@ async def _fit(self): or sync_model_step == self._total_train_steps ) - if need_sync: + if need_sync and not self._debug_rollout: + await self.agent_loop_manager.prepare_sync(self._checkpoint_due(sync_model_step)) # 同步前暂停 producer，避免 pending rollout 继续写 buffer。 with timer("pause_produce", step_timer_dict): await self.agent_loop_manager.pause_produce() @@ -1943,7 +2060,11 @@ async def _fit(self): train_step = self._cur_step + 1 finally: self.agent_loop_manager.shutdown() - await producer_task + await asyncio.sleep(0) + if not producer_task.done(): + producer_task.cancel() + await asyncio.gather(producer_task, return_exceptions=True) + await self.agent_loop_manager.close() async def _sync_weights_and_save(self, model_step: int, step_timer_dict: dict): # producer 已暂停；保持 save -> bind -> update 顺序。