From 29cbb0b5ca17cf7dc63a95511903729fd68488c0 Mon Sep 17 00:00:00 2001 From: YanhuiDua Date: Fri, 22 May 2026 03:29:48 +0000 Subject: [PATCH 1/3] Add trace store lifecycle management --- xtuner/v1/rl/rollout/session_server.py | 14 +- xtuner/v1/rl/rollout/trace_store.py | 264 ++++++++++++++++++++++--- 2 files changed, 250 insertions(+), 28 deletions(-) diff --git a/xtuner/v1/rl/rollout/session_server.py b/xtuner/v1/rl/rollout/session_server.py index 365c67e8df..736c850385 100644 --- a/xtuner/v1/rl/rollout/session_server.py +++ b/xtuner/v1/rl/rollout/session_server.py @@ -99,6 +99,7 @@ async def on_response(self, worker_resp: dict) -> dict: else: output_logprobs = [0.0] * len(output_token_ids) raw_routed_expert = choice.get("routed_experts") # 本次 call 的 raw routed_expert，可为 None + response_expert_ref = None # 2. Store 把 input_delta / assistant_output 两个节点补齐字段。 old_prompt = self.tokenizer.apply_chat_template( @@ -133,13 +134,14 @@ async def on_response(self, worker_resp: dict) -> dict: response_expert = raw_routed_expert[prefix_len + delta_len :] if delta_len > 0: - delta_node_val.expert_key = ray.put(delta_expert) # update delta node in store - await self.store.insert.remote(session_id, old_prompt, delta_node_val) + await self.store.insert.remote( + session_id, old_prompt, delta_node_val, routed_experts=ray.put(delta_expert) + ) - raw_routed_expert = ray.put(response_expert) - else: - raw_routed_expert = ray.put(raw_routed_expert) + raw_routed_expert = response_expert + + response_expert_ref = ray.put(raw_routed_expert) await self.store.insert.remote( session_id, @@ -149,9 +151,9 @@ async def on_response(self, worker_resp: dict) -> dict: token_ids=output_token_ids, logprobs=output_logprobs, labels=output_token_ids, - expert_key=raw_routed_expert, length=len(output_token_ids), ), + routed_experts=response_expert_ref, ) # 3. 返回标准 OpenAI response，session_id 由 SessionClient 层再剥 diff --git a/xtuner/v1/rl/rollout/trace_store.py b/xtuner/v1/rl/rollout/trace_store.py index 4f30453151..d5c039c4ee 100644 --- a/xtuner/v1/rl/rollout/trace_store.py +++ b/xtuner/v1/rl/rollout/trace_store.py @@ -1,10 +1,12 @@ import time from dataclasses import dataclass, field +from enum import Enum from typing import Any, Dict, List, Optional, Tuple import ray -from pydantic import BaseModel, Field +from pydantic import BaseModel, Field, StrictStr +from xtuner.v1.data_proto.rl_data import Status from xtuner.v1.utils import get_logger @@ -12,6 +14,22 @@ _handle_cache: Any = None +class TraceState(str, Enum): + ROLLOUT_RUNNING = "RolloutRunning" + ROLLOUT_FINISHED = "RolloutFinished" + TRAIN_RUNNING = "TrainRunning" + TRAIN_FINISHED = "TrainFinished" + TO_BE_RELEASED = "ToBeReleased" + RELEASED = "Released" + + +_ROLLOUT_RELEASE_STATUSES = {Status.FAILED, Status.FILTERED, Status.EXPIRED} +_ROLLOUT_KEEP_RUNNING_STATUSES = { + Status.ABORTED, + Status.INIT, +} + + def _free_ray_refs(obj: Any): """Recursively free ray.ObjectRef instances trapped inside an object. @@ -40,12 +58,17 @@ def _free_ray_refs(obj: Any): _free_ray_refs(v) +def make_expert_key(session_id: str) -> str: + """Build a stable key for a routed experts object.""" + return f"{session_id}:routed_experts" + + class TokenizedSegment(BaseModel): text: str token_ids: List[int] labels: List[int] | None = Field(default=None, repr=False) logprobs: List[float] | None = Field(default=None, repr=False) - expert_key: Any = Field(default=None, repr=False) + expert_key: StrictStr | None = Field(default=None, repr=False) length: int | None = None def model_post_init(self, _): @@ -79,6 +102,13 @@ class Trie: def __init__(self): """Initialize the prefix tree (Trie).""" self.root = TreeNode(value=None, parent=None) + self.state = TraceState.ROLLOUT_RUNNING + self.expert_key: str | None = None + self.updated_at = time.time() + + def touch(self) -> None: + """Record that this session was updated.""" + self.updated_at = time.time() def keys(self) -> List[str]: """Get all keys (i.e., strings) stored in the Trie.""" @@ -117,6 +147,7 @@ def insert(self, key: str, value: Any) -> None: break node.value = value + self.touch() def search(self, text: str, filter_none: bool = False) -> Tuple[str, List["TreeNode"]]: """Search for the longest prefix matching the input text. @@ -173,6 +204,7 @@ def _free_subtree(node: TreeNode): if key is None: _free_subtree(self.root) + self.touch() return node = self.root @@ -193,6 +225,7 @@ def _free_subtree(node: TreeNode): if node.value is not None: _free_ray_refs(node.value) node.value = None + self.touch() for parent, token in reversed(path): child_node = parent.children[token] @@ -211,7 +244,6 @@ def __init__(self): """Initialize the rollout trace store actor.""" self.sessions: Dict[str, Trie] = {} self.objects: Dict[str, ray.ObjectRef] = {} - self.updated_at: Dict[str, float] = {} def get_or_create(self, session_id: str) -> Trie: """Get the Trie for a session, or create one if it doesn't exist. @@ -226,6 +258,127 @@ def get_or_create(self, session_id: str) -> Trie: self.sessions[session_id] = Trie() return self.sessions[session_id] + def get_state(self, session_id: str) -> dict | None: + """Get lifecycle metadata for a session. + + Args: + session_id (str): The session identifier. + + Returns: + dict | None: A snapshot of session metadata, or None when the + session does not exist. + """ + trie = self.sessions.get(session_id) + if trie is None: + return None + return { + "session_id": session_id, + "state": trie.state.value, + "updated_at": trie.updated_at, + "has_object_ref": trie.expert_key in self.objects if trie.expert_key is not None else False, + } + + def list_sessions(self, state: str | None = None) -> list[dict]: + """List current session metadata snapshots, optionally filtered by state.""" + snapshots = [] + for session_id in sorted(self.sessions): + snapshot = self.get_state(session_id) + if snapshot is None: + continue + if state is not None and snapshot["state"] != state: + continue + snapshots.append(snapshot) + return snapshots + + def _set_state( + self, + session_id: str, + next_state: TraceState, + ) -> TraceState: + """Set a session state and trigger release when needed.""" + trie = self.sessions.get(session_id) + if trie is None: + raise KeyError(f"Trace session {session_id!r} does not exist.") + + trie.state = next_state + trie.touch() + self._maybe_release(session_id) + + released = session_id not in self.sessions + return TraceState.RELEASED if released else next_state + + def _maybe_release(self, session_id: str) -> None: + """Physically release a session once it reaches ToBeReleased.""" + trie = self.sessions.get(session_id) + if trie is None or trie.state != TraceState.TO_BE_RELEASED: + return + self._release_session(session_id, trie) + + def _release_session(self, session_id: str, trie: Trie) -> None: + """Release trie data and routed expert refs for one session.""" + if trie.expert_key is not None: + obj_ref = self.objects.pop(trie.expert_key, None) + if obj_ref is not None: + _free_ray_refs(obj_ref) + trie.release() + self.sessions.pop(session_id, None) + + def mark_rollout_status( + self, + session_id: str, + status: Status, + *, + enable_partial_rollout: bool = False, + ) -> str: + """Apply a rollout-side status event to one trace session.""" + release_like = status in _ROLLOUT_RELEASE_STATUSES or ( + status == Status.ABORTED and not enable_partial_rollout + ) + + trie = self.sessions.get(session_id) + if trie is None: + if release_like: + return TraceState.RELEASED.value + raise KeyError(f"Trace session {session_id!r} does not exist.") + if trie.state != TraceState.ROLLOUT_RUNNING: + raise RuntimeError( + f"Cannot handle mark_rollout_status for trace session {session_id!r} " + f"in state {trie.state.value}." + ) + + if release_like: + return self._set_state(session_id, TraceState.TO_BE_RELEASED).value + if status == Status.COMPLETED: + return self._set_state(session_id, TraceState.ROLLOUT_FINISHED).value + if status in _ROLLOUT_KEEP_RUNNING_STATUSES: + trie.touch() + return TraceState.ROLLOUT_RUNNING.value + raise AssertionError(f"Unhandled rollout status: {status!r}") + + def mark_commit_failed(self, session_id: str) -> str: + """Release a rollout session whose response commit failed.""" + trie = self.sessions.get(session_id) + if trie is None: + return TraceState.RELEASED.value + if trie.state != TraceState.ROLLOUT_RUNNING: + raise RuntimeError( + f"Cannot handle mark_commit_failed for trace session {session_id!r} " + f"in state {trie.state.value}." + ) + return self._set_state(session_id, TraceState.TO_BE_RELEASED).value + + def mark_rollout_discarded(self, session_id: str) -> str: + """Release a rollout session that external scheduling has discarded.""" + trie = self.sessions.get(session_id) + if trie is None: + return TraceState.RELEASED.value + if trie.state not in (TraceState.ROLLOUT_RUNNING, TraceState.ROLLOUT_FINISHED): + raise RuntimeError( + f"Cannot handle mark_rollout_discarded for trace session {session_id!r} " + f"in state {trie.state.value}." + ) + return self._set_state(session_id, TraceState.TO_BE_RELEASED).value + def keys(self, session_id: str) -> List[str]: """Get all keys (i.e., strings) stored in a session's Trie. @@ -236,18 +389,39 @@ def keys(self, session_id: str) -> List[str]: List[str]: A list of all keys in the session's Trie. """ trie = self.get_or_create(session_id) + if trie.state == TraceState.TO_BE_RELEASED: + get_logger().error(f"Trace session {session_id!r} is pending release; skip keys.") + return [] return trie.keys() - def insert(self, session_id: str, key: str, value: Any): + def insert( + self, + session_id: str, + key: str, + value: TokenizedSegment, + routed_experts: ray.ObjectRef | None = None, + ): """Insert a (key, value) pair into a session's Trie. Args: session_id (str): The session identifier. key (str): The key string. - value (Any): The trace segment/value to store. + value (TokenizedSegment): The trace segment to store. + routed_experts (ray.ObjectRef | None): Optional routed experts + object for this session. """ trie = self.get_or_create(session_id) - return trie.insert(key, value) + if trie.state != TraceState.ROLLOUT_RUNNING: + get_logger().error( + f"Cannot insert into trace session {session_id!r} in state {trie.state.value}; skip insert." + ) + return + if routed_experts is not None: + expert_key = make_expert_key(session_id) + self.objects[expert_key] = routed_experts + value.expert_key = expert_key + trie.expert_key = expert_key + trie.insert(key, value) def search(self, session_id: str, text: str, filter_none: bool = False): """Search the longest prefix in a session's Trie. @@ -261,19 +435,11 @@ def search(self, session_id: str, text: str, filter_none: bool = False): Tuple[str, List["TreeNode"]]: The matched prefix and matched nodes. """ trie = self.get_or_create(session_id) + if trie.state == TraceState.TO_BE_RELEASED: + get_logger().error(f"Trace session {session_id!r} is pending release; skip search.") + return "", [] return trie.search(text, filter_none) - def release(self, session_id: str): - """Release the Trie and free associated resources for a specific - session. - - Args: - session_id (str): The session identifier. - """ - assert session_id in self.sessions, f"Session ID '{session_id}' not found for release." - trie = self.sessions.pop(session_id) - trie.release() - def export_training_trace(self, session_id: str, prompt_text: str) -> dict: """Export the stored training trace given a complete prompt text. @@ -283,26 +449,73 @@ def export_training_trace(self, session_id: str, prompt_text: str) -> dict: Returns: dict: The trace dictionary containing `input_ids`, `labels`, `logprobs`, - and `routed_experts`. + and the session-level `routed_experts` object key. Raises: + KeyError: If the session does not exist. + RuntimeError: If the session is not ready for training export. ValueError: If the prompt_text does not completely match the trace keys in the session. """ - trie = self.get_or_create(session_id) + trie = self.sessions.get(session_id) + if trie is None: + raise KeyError(f"Trace session {session_id!r} does not exist.") + if trie.state != TraceState.ROLLOUT_FINISHED: + raise RuntimeError( + f"Cannot export training trace for session {session_id!r} in state {trie.state.value}." + ) + key, nodes = trie.search(prompt_text, filter_none=True) if prompt_text != key: + self._set_state(session_id, TraceState.TO_BE_RELEASED) raise ValueError( f"Prompt text '{prompt_text}' does not match any trace key '{key}' in session '{session_id}'." ) - trace = {"input_ids": [], "labels": [], "logprobs": [], "routed_experts": []} + trace = {"input_ids": [], "labels": [], "logprobs": [], "routed_experts": trie.expert_key} for node in nodes: node_val: TokenizedSegment = node.value trace["input_ids"].extend(node_val.token_ids) trace["labels"].extend(node_val.labels) trace["logprobs"].extend(node_val.logprobs) - trace["routed_experts"].append(node_val.expert_key) + self._set_state(session_id, TraceState.TRAIN_RUNNING) return trace + def mark_train_finished(self, session_id: str) -> str: + """Release a session after trainer consumers have finished using it.""" + trie = self.sessions.get(session_id) + if trie is None: + return TraceState.RELEASED.value + if trie.state != TraceState.TRAIN_RUNNING: + raise RuntimeError( + f"Cannot handle mark_train_finished for trace session {session_id!r} " + f"in state {trie.state.value}." + ) + self._set_state(session_id, TraceState.TRAIN_FINISHED) + return self._set_state(session_id, TraceState.TO_BE_RELEASED).value + + def mark_train_abandoned(self, session_id: str) -> str: + """Release a training session that trainer will no longer consume.""" + trie = self.sessions.get(session_id) + if trie is None: + return TraceState.RELEASED.value + if trie.state != TraceState.TRAIN_RUNNING: + raise RuntimeError( + f"Cannot handle mark_train_abandoned for trace session {session_id!r} " + f"in state {trie.state.value}." + ) + return self._set_state(session_id, TraceState.TO_BE_RELEASED).value + + def gc_stale_sessions(self, ttl_seconds: float) -> list[str]: + """Release stale RolloutRunning sessions older than the given TTL.""" + now = time.time() + stale_session_ids = [ + session_id + for session_id, trie in self.sessions.items() + if trie.state == TraceState.ROLLOUT_RUNNING and (now - trie.updated_at) > ttl_seconds + ] + for session_id in stale_session_ids: + self._set_state(session_id, TraceState.TO_BE_RELEASED) + return stale_session_ids + def get_objects(self, keys: list[str]) -> list[ray.ObjectRef]: """Fetch ray.ObjectRef elements by their keys. @@ -312,7 +525,14 @@ def get_objects(self, keys: list[str]) -> list[ray.ObjectRef]: Returns: list[ray.ObjectRef]: The mapped ray.ObjectRefs. """ - return [self.objects[key] for key in keys if key in self.objects] + object_refs: list[ray.ObjectRef] = [] + for key in keys: + if not isinstance(key, str) or not key: + raise KeyError(f"Invalid trace object key: {key!r}") + if key not in self.objects: + raise KeyError(f"Trace object key {key!r} does not exist.") + object_refs.append(self.objects[key]) + return object_refs def get_store(): From a2f66912d3e7532d7cdef84c030fc0f701a3012d Mon Sep 17 00:00:00 2001 From: YanhuiDua Date: Fri, 22 May 2026 09:07:31 +0000 Subject: [PATCH 2/3] Report rollout statuses to trace store --- .../rl/test_multi_task_agent_loop_manager.py | 16 +- tests/rl/test_producer.py | 16 +- tests/rl/test_replay_buffer.py | 21 ++- tests/rl/test_trace_store.py | 159 ++++++++++++++++++ .../agent_loop_manager/agent_loop_manager.py | 21 ++- xtuner/v1/rl/agent_loop_manager/producer.py | 12 ++ xtuner/v1/rl/replay_buffer.py | 21 ++- xtuner/v1/rl/rollout/trace_store.py | 28 +++ 8 files changed, 272 insertions(+), 22 deletions(-) create mode 100644 tests/rl/test_trace_store.py diff --git a/tests/rl/test_multi_task_agent_loop_manager.py b/tests/rl/test_multi_task_agent_loop_manager.py index bbff8aeef9..68caecbb01 100644 --- a/tests/rl/test_multi_task_agent_loop_manager.py +++ b/tests/rl/test_multi_task_agent_loop_manager.py @@ -14,6 +14,7 @@ _TaskRunner, ) from xtuner.v1.rl.agent_loop_manager.producer import GROUP_GENERATE_TIME_KEY, ProduceBatchStatus +from xtuner.v1.rl.replay_buffer import RefreshStalenessResult from xtuner.v1.rl.utils import calculate_seq_staleness @@ -144,7 +145,11 @@ async def produce_batch(self, ctx) -> ProduceBatchStatus: class _FakeReplayBuffer: - def __init__(self, rollout_states_by_task: dict[str, list[list[str]]], leftover_counts: dict[tuple[str, Status], int]): + def __init__( + self, + rollout_states_by_task: dict[str, list[list[str]]], + leftover_counts: dict[tuple[str, Status], int], + ): self._rollout_states_by_task = rollout_states_by_task self._leftover_counts = leftover_counts self.saved_paths: list[Path] = [] @@ -168,7 +173,7 @@ async def refresh_staleness( current_train_step: int, statuses: list[Status] | None = None, ): - expired_counts = {} + refresh_results = {} for task_name, stale_threshold in task_stale_thresholds.items(): self.refresh_staleness_calls.append( (task_name, current_train_step, stale_threshold, tuple(statuses or ())) @@ -180,8 +185,11 @@ async def refresh_staleness( state.seq_staleness = calculate_seq_staleness( min(response_model_steps), current_train_step ) - expired_counts[task_name] = 0 - return expired_counts + refresh_results[task_name] = RefreshStalenessResult( + expired_count=0, + expired_session_ids=[], + ) + return refresh_results async def is_ready(self, task_batch_sizes: dict[str, int], *, group_status: Status = Status.COMPLETED): for task_name, batch_size in task_batch_sizes.items(): diff --git a/tests/rl/test_producer.py b/tests/rl/test_producer.py index 5e6adf0d4e..cffe9b5ff5 100644 --- a/tests/rl/test_producer.py +++ b/tests/rl/test_producer.py @@ -16,9 +16,17 @@ class MockRolloutState: - def __init__(self, id, seq_staleness=1, status=Status.COMPLETED, reward_score=None): + def __init__( + self, + id, + seq_staleness=1, + status=Status.COMPLETED, + reward_score=None, + session_uid=None, + ): self.id = id self.uid = id + self.session_uid = session_uid self.status = status self.seq_staleness = seq_staleness self.response_ids = [] @@ -114,6 +122,7 @@ def _build_context( progress=progress, is_valid_sample_fn=strategy.is_valid_sample_fn, stale_threshold=getattr(strategy, "stale_threshold", None), + enable_partial_rollout=getattr(strategy, "enable_partial_rollout", False), ) def test_produce_progress_methods_keep_absolute_window(self): @@ -827,12 +836,13 @@ async def test_refresh_staleness_refreshes_before_expire_check(self): stale_item.response_model_steps = [3] await self.replay_buffer.put([stale_item], task_name) - expired_counts = await self.replay_buffer.refresh_staleness( + refresh_results = await self.replay_buffer.refresh_staleness( task_stale_thresholds={task_name: 2}, current_train_step=6, ) expired_groups = await self.replay_buffer.get(10, task_name, Status.EXPIRED) - self.assertEqual(expired_counts, {task_name: 1}) + self.assertEqual(refresh_results[task_name].expired_count, 1) + self.assertEqual(refresh_results[task_name].expired_session_ids, []) self.assertEqual(len(expired_groups), 1) self.assertEqual(expired_groups[0][0].seq_staleness, 2) diff --git a/tests/rl/test_replay_buffer.py b/tests/rl/test_replay_buffer.py index d90dc11f84..97308d3550 100644 --- a/tests/rl/test_replay_buffer.py +++ b/tests/rl/test_replay_buffer.py @@ -15,6 +15,7 @@ def __init__( status=Status.COMPLETED, response_model_steps=None, response_ids=None, + session_uid=None, ): self.id = state_id self.seq_staleness = staleness @@ -22,6 +23,7 @@ def __init__( self.input_ids = input_ids if input_ids is not None else [state_id] self.response_model_steps = response_model_steps self.response_ids = response_ids if response_ids is not None else [] + self.session_uid = session_uid class TestReplayBuffer(unittest.IsolatedAsyncioTestCase): @@ -195,12 +197,13 @@ async def test_refresh_staleness_expires_completed_in_place(self): "task", ) - expired_counts = await replay_buffer.refresh_staleness( + refresh_results = await replay_buffer.refresh_staleness( task_stale_thresholds={"task": 2}, current_train_step=6, ) - self.assertEqual(expired_counts, {"task": 1}) + self.assertEqual(refresh_results["task"].expired_count, 1) + self.assertEqual(refresh_results["task"].expired_session_ids, []) self.assertEqual(await replay_buffer.count("task", Status.COMPLETED), 1) self.assertEqual(await replay_buffer.count("task", Status.EXPIRED), 1) expired = await replay_buffer.get(1, "task", Status.EXPIRED) @@ -212,7 +215,7 @@ async def test_refresh_staleness_expires_completed_in_place(self): async def test_refresh_staleness_expires_aborted_in_place(self): replay_buffer = AsyncReplayBufferConfig().build() await replay_buffer.put( - [MockState("stale-aborted", response_model_steps=[3], status=Status.ABORTED)], + [MockState("stale-aborted", response_model_steps=[3], status=Status.ABORTED, session_uid="sid-aborted")], "task", ) await replay_buffer.put( @@ -220,12 +223,13 @@ async def test_refresh_staleness_expires_aborted_in_place(self): "task", ) - expired_counts = await replay_buffer.refresh_staleness( + refresh_results = await replay_buffer.refresh_staleness( task_stale_thresholds={"task": 2}, current_train_step=6, ) - self.assertEqual(expired_counts, {"task": 1}) + self.assertEqual(refresh_results["task"].expired_count, 1) + self.assertEqual(refresh_results["task"].expired_session_ids, ["sid-aborted"]) self.assertEqual(await replay_buffer.count("task", Status.ABORTED), 1) self.assertEqual(await replay_buffer.count("task", Status.EXPIRED), 1) expired = await replay_buffer.get(1, "task", Status.EXPIRED) @@ -244,17 +248,18 @@ async def test_refresh_staleness_respects_status_filter(self): "task", ) await replay_buffer.put( - [MockState("stale-aborted", response_model_steps=[3], status=Status.ABORTED)], + [MockState("stale-aborted", response_model_steps=[3], status=Status.ABORTED, session_uid="sid-filter")], "task", ) - expired_counts = await replay_buffer.refresh_staleness( + refresh_results = await replay_buffer.refresh_staleness( task_stale_thresholds={"task": 2}, current_train_step=6, statuses=[Status.ABORTED], ) - self.assertEqual(expired_counts, {"task": 1}) + self.assertEqual(refresh_results["task"].expired_count, 1) + self.assertEqual(refresh_results["task"].expired_session_ids, ["sid-filter"]) self.assertEqual(await replay_buffer.count("task", Status.COMPLETED), 1) self.assertEqual(await replay_buffer.count("task", Status.ABORTED), 0) self.assertEqual(await replay_buffer.count("task", Status.EXPIRED), 1) diff --git a/tests/rl/test_trace_store.py b/tests/rl/test_trace_store.py new file mode 100644 index 0000000000..eb9ee752e9 --- /dev/null +++ b/tests/rl/test_trace_store.py @@ -0,0 +1,159 @@ +import asyncio +import unittest +from unittest.mock import AsyncMock, MagicMock, patch + +from xtuner.v1.data_proto.rl_data import Status +from xtuner.v1.rl.agent_loop_manager import ProduceContext, ProduceProgress +from xtuner.v1.rl.agent_loop_manager.agent_loop_manager import AgentLoopManager, _TaskRunner +from xtuner.v1.rl.replay_buffer import AsyncReplayBufferConfig, RefreshStalenessResult +from xtuner.v1.rl.rollout.trace_store import RolloutTraceStore, TokenizedSegment, TraceState + + +class _TraceRolloutState: + def __init__( + self, + uid: int | str, + *, + status: Status = Status.COMPLETED, + reward_score: float | None = None, + session_uid: int | str | None = None, + ): + self.uid = uid + self.id = uid + self.session_uid = session_uid + self.status = status + self.seq_staleness = 0 + self.response_ids = [] + self.extra_fields = {} + self.reward = {"score": reward_score} if reward_score is not None else None + + +class _TraceRefreshReplayBuffer: + def __init__(self, expired_session_ids_by_task: dict[str, list[int | str]]): + self._expired_session_ids_by_task = expired_session_ids_by_task + + async def refresh_staleness( + self, + *, + task_stale_thresholds: dict[str, int], + current_train_step: int, + statuses: list[Status] | None = None, + ): + return { + task_name: RefreshStalenessResult( + expired_count=len(self._expired_session_ids_by_task.get(task_name, [])), + expired_session_ids=self._expired_session_ids_by_task.get(task_name, []), + ) + for task_name in task_stale_thresholds + } + + +class _TraceProduceStrategy: + stale_threshold = 5 + enable_partial_rollout = False + + +class TestRolloutTraceStoreRolloutStatus(unittest.TestCase): + def setUp(self): + store_cls = RolloutTraceStore.__ray_metadata__.modified_class + self.store = store_cls() + + def _insert_segment(self, session_id: str): + self.store.insert(session_id, "prompt", TokenizedSegment(text="prompt", token_ids=[1])) + + def test_mark_rollout_statuses_marks_completed_and_releases_filtered(self): + self._insert_segment("completed") + self._insert_segment("filtered") + + results = self.store.mark_rollout_statuses( + [ + ("completed", Status.COMPLETED), + ("filtered", Status.FILTERED), + ] + ) + + self.assertEqual(results["completed"], TraceState.ROLLOUT_FINISHED.value) + self.assertEqual(results["filtered"], TraceState.RELEASED.value) + self.assertEqual( + self.store.get_state("completed")["state"], + TraceState.ROLLOUT_FINISHED.value, + ) + self.assertIsNone(self.store.get_state("filtered")) + + def test_mark_rollout_statuses_discards_expired_finished_session(self): + self._insert_segment("expired") + self.store.mark_rollout_status("expired", Status.COMPLETED) + + results = self.store.mark_rollout_statuses([("expired", Status.EXPIRED)]) + + self.assertEqual(results["expired"], TraceState.RELEASED.value) + self.assertIsNone(self.store.get_state("expired")) + + +class TestTraceStoreProducerReporting(unittest.IsolatedAsyncioTestCase): + async def test_put_generated_group_reports_final_status_to_trace_store(self): + task_name = "test_trace_status" + progress = ProduceProgress.build([task_name]) + replay_buffer = AsyncReplayBufferConfig().build() + ctx = ProduceContext( + agent_loop=MagicMock(), + sampler=MagicMock(), + replay_buffer=replay_buffer, + task_batch_size=1, + task_name=task_name, + train_step=0, + update_event=asyncio.Event(), + model_step=0, + progress=progress, + is_valid_sample_fn=lambda samples: False, + ) + store = MagicMock() + store.mark_rollout_statuses.remote = AsyncMock(return_value={}) + + completed_group = [ + _TraceRolloutState( + 1, + status=Status.COMPLETED, + reward_score=1.0, + session_uid="trace-session", + ) + ] + with patch("xtuner.v1.rl.agent_loop_manager.producer.get_store", return_value=store): + self.assertFalse(await ctx.put_generated_group(completed_group)) + + store.mark_rollout_statuses.remote.assert_awaited_once_with( + [("trace-session", Status.FILTERED)], + enable_partial_rollout=False, + ) + + +class TestTraceStoreManagerReporting(unittest.IsolatedAsyncioTestCase): + async def test_refresh_for_all_tasks_reports_expired_sessions_to_trace_store(self): + replay_buffer = _TraceRefreshReplayBuffer({"task_a": ["sid-a", "sid-b"]}) + manager = AgentLoopManager( + task_runners=[ + _TaskRunner( + task_name="task_a", + agent_loop=MagicMock(), + produce_strategy=_TraceProduceStrategy(), + sampler=MagicMock(), + weight=1.0, + order=0, + ), + ], + replay_buffer=replay_buffer, + ) + store = MagicMock() + store.mark_rollout_statuses.remote = AsyncMock(return_value={}) + + with patch("xtuner.v1.rl.agent_loop_manager.agent_loop_manager.get_store", return_value=store): + await manager._refresh_for_all_tasks(9, [Status.COMPLETED, Status.ABORTED]) + + store.mark_rollout_statuses.remote.assert_awaited_once_with( + [("sid-a", Status.EXPIRED), ("sid-b", Status.EXPIRED)], + enable_partial_rollout=False, + ) + + +if __name__ == "__main__": + unittest.main() 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 7663ecc85e..466b5f09ec 100644 --- a/xtuner/v1/rl/agent_loop_manager/agent_loop_manager.py +++ b/xtuner/v1/rl/agent_loop_manager/agent_loop_manager.py @@ -14,6 +14,7 @@ from xtuner.v1.rl.judger import ComposedJudgerConfig, JudgerConfig, build_judger from xtuner.v1.rl.replay_buffer import ReplayBuffer from xtuner.v1.rl.rollout import RolloutController +from xtuner.v1.rl.rollout.trace_store import get_store from xtuner.v1.rl.utils import asyncio_run from xtuner.v1.utils import get_logger @@ -209,6 +210,7 @@ def _build_produce_context( progress=progress, is_valid_sample_fn=getattr(task_runner.produce_strategy, "is_valid_sample_fn", default_is_valid_sample_fn), stale_threshold=getattr(task_runner.produce_strategy, "stale_threshold", None), + enable_partial_rollout=getattr(task_runner.produce_strategy, "enable_partial_rollout", False), ) @@ -461,15 +463,28 @@ async def _refresh_for_all_tasks(self, train_step: int, statuses: list[Status]) stale_threshold = getattr(task.produce_strategy, "stale_threshold", 1) task_stale_thresholds[task.task_name] = stale_threshold - expired_counts = await self.replay_buffer.refresh_staleness( + refresh_results = await self.replay_buffer.refresh_staleness( task_stale_thresholds=task_stale_thresholds, current_train_step=train_step, statuses=statuses, ) - for task_name, expired_count in expired_counts.items(): + for task in self.task_runners: + task_name = task.task_name + refresh_result = refresh_results[task_name] self.logger.info( - f"[AgentLoopManager][{self.name}] Refresh staleness for task {task_name}: expired_count={expired_count}" + f"[AgentLoopManager][{self.name}] Refresh staleness for task {task_name}: " + f"expired_count={refresh_result.expired_count}" ) + if refresh_result.expired_session_ids: + trace_events = [(session_id, Status.EXPIRED) for session_id in refresh_result.expired_session_ids] + try: + store = get_store() + await store.mark_rollout_statuses.remote( + trace_events, + enable_partial_rollout=getattr(task.produce_strategy, "enable_partial_rollout", False), + ) + except Exception as exc: + self.logger.error(f"Failed to report trace store expired rollout status events: {exc}") def _get_task_batch_sizes_for_step(self, batch_size: int, train_step: int) -> dict[str, int]: if len(self.task_runners) == 1: diff --git a/xtuner/v1/rl/agent_loop_manager/producer.py b/xtuner/v1/rl/agent_loop_manager/producer.py index 35db08f27d..a7892ddaf9 100644 --- a/xtuner/v1/rl/agent_loop_manager/producer.py +++ b/xtuner/v1/rl/agent_loop_manager/producer.py @@ -22,6 +22,7 @@ ) from xtuner.v1.rl.agent_loop import AgentLoopSpec, get_agent_loop_rollout_ctl from xtuner.v1.rl.replay_buffer import ReplayBuffer +from xtuner.v1.rl.rollout.trace_store import get_store from xtuner.v1.rl.utils import calculate_seq_staleness, create_task from xtuner.v1.utils import get_logger @@ -307,6 +308,7 @@ class ProduceContext: progress: ProduceProgress is_valid_sample_fn: IsValidSampleFn = default_is_valid_sample_fn stale_threshold: int | None = None + enable_partial_rollout: bool = False @property def consumer_step(self) -> int: @@ -383,6 +385,16 @@ async def put_generated_group(self, group: list[RolloutState]) -> bool: current_train_step=self.consumer_step, stale_threshold=self.stale_threshold, ) + trace_events = [(item.session_uid, item.status) for item in group if item.session_uid is not None] + if trace_events: + try: + store = get_store() + await store.mark_rollout_statuses.remote( + trace_events, + enable_partial_rollout=self.enable_partial_rollout, + ) + except Exception as exc: + logger.error(f"Failed to report trace store rollout status events: {exc}") produced_tokens = 0 for item in group: response_ids = getattr(item, "response_ids", None) diff --git a/xtuner/v1/rl/replay_buffer.py b/xtuner/v1/rl/replay_buffer.py index d5f9f1336a..a29a2e2c64 100644 --- a/xtuner/v1/rl/replay_buffer.py +++ b/xtuner/v1/rl/replay_buffer.py @@ -59,6 +59,12 @@ class StorageItem: staleness: int +@dataclass +class RefreshStalenessResult: + expired_count: int + expired_session_ids: list[int | str] + + @dataclass class SerializedRayObjectRef: value: Any @@ -502,14 +508,14 @@ async def refresh_staleness( task_stale_thresholds: dict[str, int], current_train_step: int, statuses: list[Status] | None = None, - ) -> dict[str, int]: + ) -> dict[str, RefreshStalenessResult]: # 刷新可复用样本的 staleness；completed / aborted 都可能来自旧权重，需要按 train_step 淘汰。 for task_name, stale_threshold in task_stale_thresholds.items(): if stale_threshold <= 0: raise ValueError(f"stale_threshold must be positive, got {stale_threshold}.") if statuses is None: statuses = [Status.COMPLETED, Status.ABORTED] - expired_counts: dict[str, int] = {} + results: dict[str, RefreshStalenessResult] = {} async with self._lock: updated_records: list[StorageItem] = [] for task_name, stale_threshold in task_stale_thresholds.items(): @@ -521,6 +527,7 @@ async def refresh_staleness( } records = await self._storage.get(query_dsl) expired_count = 0 + expired_session_ids: list[int | str] = [] for record in records: refresh_seq_staleness(record.item, current_train_step) staleness = max((getattr(item, "seq_staleness", 0) for item in record.item), default=0) @@ -532,12 +539,18 @@ async def refresh_staleness( item.status = Status.EXPIRED status = Status.EXPIRED expired_count += 1 + for item in record.item: + if item.session_uid is not None: + expired_session_ids.append(item.session_uid) else: status = get_group_status(record.item) updated_records.append(replace(record, status=status, staleness=staleness)) - expired_counts[task_name] = expired_count + results[task_name] = RefreshStalenessResult( + expired_count=expired_count, + expired_session_ids=expired_session_ids, + ) await self._storage.update(updated_records) - return expired_counts + return results async def is_ready( self, diff --git a/xtuner/v1/rl/rollout/trace_store.py b/xtuner/v1/rl/rollout/trace_store.py index d5c039c4ee..43560a1801 100644 --- a/xtuner/v1/rl/rollout/trace_store.py +++ b/xtuner/v1/rl/rollout/trace_store.py @@ -355,6 +355,34 @@ def mark_rollout_status( return TraceState.ROLLOUT_RUNNING.value raise AssertionError(f"Unhandled rollout status: {status!r}") + def mark_rollout_statuses( + self, + events: list[tuple[Any, Status]], + *, + enable_partial_rollout: bool = False, + ) -> dict[Any, str]: + """Apply final rollout status events reported by the producer.""" + results: dict[Any, str] = {} + for session_id, status in events: + try: + trie = self.sessions.get(session_id) + if trie is None: + continue + if status == Status.EXPIRED and trie.state == TraceState.ROLLOUT_FINISHED: + results[session_id] = self.mark_rollout_discarded(session_id) + else: + results[session_id] = self.mark_rollout_status( + session_id, + status, + enable_partial_rollout=enable_partial_rollout, + ) + except Exception as exc: + get_logger().error( + f"Failed to mark trace store rollout status for session {session_id!r} " + f"with status {status}: {exc}" + ) + return results + def mark_commit_failed(self, session_id: str) -> str: """Release a rollout session whose response commit failed.""" trie = self.sessions.get(session_id) From fb60fb4ce7321727d3431dd854ffadd977c6e13e Mon Sep 17 00:00:00 2001 From: YanhuiDua Date: Tue, 26 May 2026 03:23:51 +0000 Subject: [PATCH 3/3] Integrate TraceStore training trace materialization --- docs/design/trace_store_lifecycle_impl.md | 589 ++++++++++++++++++++ docs/design/trace_store_lifecycle_simple.md | 431 ++++++++++++++ xtuner/v1/data_proto/rl_data.py | 3 + xtuner/v1/data_proto/sequence_context.py | 4 +- xtuner/v1/rl/agent_loop_manager/sampler.py | 3 +- xtuner/v1/rl/trainer/worker.py | 12 +- xtuner/v1/rl/utils/misc.py | 56 +- xtuner/v1/train/rl_trainer.py | 217 ++++++-- 8 files changed, 1235 insertions(+), 80 deletions(-) create mode 100644 docs/design/trace_store_lifecycle_impl.md create mode 100644 docs/design/trace_store_lifecycle_simple.md diff --git a/docs/design/trace_store_lifecycle_impl.md b/docs/design/trace_store_lifecycle_impl.md new file mode 100644 index 0000000000..09a2c30994 --- /dev/null +++ b/docs/design/trace_store_lifecycle_impl.md @@ -0,0 +1,589 @@ +# Trace Store Lifecycle Implementation Design + +## 1. 目标和边界 + +本文是 `trace_store_lifecycle_simple.md` 的实现讨论稿，用于把 Rollout Trace Store 的生命周期设计落到 +`xtuner/v1/rl/rollout/trace_store.py`。 + +本文件记录当前实现决策。后续每次讨论后，直接更新本文，而不是立即修改代码。 + +第一版实现边界： + +1. 只设计 Trace Store actor 内部的 session 级生命周期管理。 +2. 核心只围绕 `state`。 +3. 不引入 `created_at`；`updated_at` 保留，但从 `RolloutTraceStore.updated_at` 下沉为每个 session / `Trie` 自己的属性。 +4. 不引入 trainer rank / materialize 计数。trainer 自己汇总消费状态，并在确认不再依赖 Store 后调用 Trace Store 状态转换 API。 +5. 不引入新的 staged / committed object registry。 +6. routed experts Ray `ObjectRef` 由 `RolloutTraceStore.objects: dict[str, ray.ObjectRef]` 统一持有；`TokenizedSegment.expert_key` 中只保留对应 object key。 +7. 一个 session 只记录一个 routed experts object key；释放 session 时删除这个 key 对应的 object ref。 + +非目标： + +1. 不修改 ReplayBuffer / producer 的采样和重试策略。 +2. 不实现 failed prefix resume。 +3. 不支持失败重试复用旧 `session_id`。 +4. 不做节点级 partial release。 +5. 不设计 trainer 内部多 rank ack 协议。 + +## 2. 当前实现理解 + +旧实现里 `RolloutTraceStore` 主要有三组 actor 级状态： + +```python +self.sessions: Dict[str, Trie] +self.objects: Dict[str, ray.ObjectRef] +self.updated_at: Dict[str, float] +``` + +其中 `self.objects` 已有查询接口，但没有完整写入和释放链路；`self.updated_at` 与 `sessions` 分离。新实现保留 `self.objects`，但把它收敛为正式的 routed experts object registry；`updated_at` 下沉到 `Trie`。 + +当前实际使用路径里，`sessions` 是核心： + +1. `SessionServer.on_request` 根据 `session_id` 在 store 中做 prefix `search`。 +2. prompt delta 会写入对应 session 的 `Trie`。 +3. `SessionServer.on_response` 会把 assistant output 写入同一棵 `Trie`。 +4. routed experts 通过 `ray.put(...)` 后得到 `ObjectRef`。 +5. 旧实现存在把 `ObjectRef` 直接放进 `TokenizedSegment.expert_key` 的路径；新实现禁止这种写法。 + +新的实现方向： + +1. `RolloutTraceStore.sessions` 继续保持 `dict[str, Trie]`。 +2. `Trie` 不改名，作为单条 session 的数据和生命周期承载对象。 +3. actor 级 `self.objects` 保留，作为唯一持有 routed experts `ObjectRef` 的 registry。 +4. actor 级 `self.updated_at` 下沉为 `Trie.updated_at`。 +5. routed experts 的实际 `ray.ObjectRef` 存到 `RolloutTraceStore.objects[expert_key]`。 +6. `TokenizedSegment.expert_key` 保存 object key，而不是直接保存 routed experts 的 `ObjectRef`。 +7. session 删除时，必须按该 session 记录的 object key 从 `self.objects` 中删除并释放 ref。 + +这样保持原有 actor 级 object lookup 形态，同时要求所有 object 写入、替换和释放都经过 `RolloutTraceStore` 的 helper，避免 500M 级别 routed experts object 因漏删 key 而常驻。 + +## 3. 核心状态模型 + +新增 `TraceState`，与上游设计保持一致： + +```python +class TraceState(str, Enum): + ROLLOUT_RUNNING = "RolloutRunning" + ROLLOUT_FINISHED = "RolloutFinished" + TRAIN_RUNNING = "TrainRunning" + TRAIN_FINISHED = "TrainFinished" + TO_BE_RELEASED = "ToBeReleased" + RELEASED = "Released" +``` + +状态含义： + +| 状态 | 含义 | 能否物理释放 | +| --- | --- | --- | +| `RolloutRunning` | rollout 还在运行、继续写入，或 `ABORTED + enable_partial_rollout=True` 等待续跑 | 不能 | +| `RolloutFinished` | rollout 已完成且未被过滤，等待训练侧 export | 不能 | +| `TrainRunning` | 训练侧已 export，仍可能依赖 Store 中的 refs | 不能 | +| `TrainFinished` | trainer 已确认不再依赖 Store | 不能；下一步进入 `ToBeReleased` | +| `ToBeReleased` | 已不再服务 rollout 或 training，等待统一释放 | 可以 | +| `Released` | 概念终态；实现上表现为 session metadata 已删除 | 终态 | + +## 4. Trie 作为 session 数据结构 + +不新增 `TraceSession` 包装类。`Trie` 名字保持不变，在 `Trie` 内新增 session 级属性： + +```python +class Trie: + root: TreeNode + state: TraceState = TraceState.ROLLOUT_RUNNING + expert_key: str | None = None + updated_at: float +``` + +对应地，`RolloutTraceStore.sessions` 保持： + +```python +self.sessions: Dict[str, Trie] +``` + +不调整为 `Dict[str, TraceSession]`。 + +字段含义： + +1. `root`：原有 prefix tree 根节点。 +2. `state`：该 session 当前生命周期状态。 +3. `expert_key`：该 session 唯一 routed experts object key。 +4. `updated_at`：该 session 最近一次写入、状态转换或 routed experts object key 变更时间。 + +实现约束： + +1. `get_or_create(session_id)` 仍返回 `Trie`。 +2. 现有 `keys` / `insert` / `search` 的外部接口尽量保持兼容。 +3. `Trie` 可以新增内部 helper，例如 `touch()`。 + +### 4.1 routed experts object 存储 + +`TokenizedSegment.expert_key` 继续保留这个字段名，但语义调整为 object key。 + +推荐实现方式： + +```python +def make_expert_key(session_id: str) -> str: + return f"{session_id}:routed_experts" + +def insert( + session_id: str, + key: str, + value: TokenizedSegment, + routed_experts: ray.ObjectRef | None = None, +) -> None: + if routed_experts is not None: + expert_key = make_expert_key(session_id) + self.objects[expert_key] = routed_experts + value.expert_key = expert_key + trie.expert_key = expert_key + trie.insert(key, value) +``` + +写入 `TokenizedSegment` 前后需要满足： + +1. routed experts 的实际 `ObjectRef` 在 `RolloutTraceStore.objects` 中。 +2. `TokenizedSegment.expert_key` 是能找回该 `ObjectRef` 的 key。 +3. `TokenizedSegment.expert_key` 只允许 `str | None`，禁止写入 `ray.ObjectRef`。 +4. object key 是 `TokenizedSegment` 里保存的轻量索引，用来从 session 找回实际 `ray.ObjectRef`。它不是新的业务协议，只是避免把 `ObjectRef` 直接塞进 trie node value。 +5. object key 从 `session_id` 派生，一个 session 只有一个 routed experts object key。 +6. `TokenizedSegment.expert_key` 类型语义为 `str | None`。`None` 表示该 segment 没有 routed experts object。 +7. 一个 session 只有一个 object key。 +8. 调用方通过 `insert(..., routed_experts=obj_ref)` 写入真实 ref；`insert` 负责生成 `expert_key` 并写回 `TokenizedSegment`。 +9. 第一版假设同一个 expert key 不会重复写入；`insert` 覆盖 dict 时不负责释放旧 ref。 + +## 5. 状态转换 API 草案 + +### 5.1 查询状态 + +```python +def get_state(session_id: str) -> dict | None: + ... +``` + +返回示例： + +```python +{ + "session_id": session_id, + "state": "RolloutRunning", + "updated_at": 1234567890.0, + "has_object_ref": True, +} +``` + +session 不存在时返回 `None`。 + +### 5.2 rollout 状态上报 + +```python +def mark_rollout_status( + session_id: str, + status: Status, + *, + enable_partial_rollout: bool = False, +) -> str: + ... +``` + +状态映射： + +| rollout status | 条件 | Trace Store 状态 | +| --- | --- | --- | +| `COMPLETED` | - | `RolloutFinished` | +| `ABORTED` | `enable_partial_rollout=True` | 保持 `RolloutRunning` | +| `ABORTED` | `enable_partial_rollout=False` | `ToBeReleased` | +| `FAILED` | - | `ToBeReleased` | +| `FILTERED` | - | `ToBeReleased` | +| `EXPIRED` | - | `ToBeReleased` | +| `INIT` | - | 保持 `RolloutRunning` | +| `ARCHIVED` | - | 第一版暂不处理，传入 `mark_rollout_status` 时抛错 | + +进入 `ToBeReleased` 后立即调用 `_maybe_release(session_id)`。 + +状态约束： + +1. `mark_rollout_status` 不创建 session。 +2. `COMPLETED` 和 `ABORTED + enable_partial_rollout=True` 必须要求 session 已存在且当前为 `RolloutRunning`；否则抛 `KeyError` 或 `RuntimeError`。 +3. `FAILED` / `FILTERED` / `EXPIRED` / `ABORTED + enable_partial_rollout=False` 是 release-like 事件；session 不存在时返回 `Released`，不创建空 session。 +4. 如果 session 已经不是 `RolloutRunning`，新的 rollout status 不能覆盖后续状态，必须抛 `RuntimeError`。 + +### 5.3 rollout 放弃和 commit 失败 + +```python +def mark_commit_failed(session_id: str) -> str: + ... + +def mark_rollout_discarded(session_id: str) -> str: + ... +``` + +两者都进入 `ToBeReleased`，然后调用 `_maybe_release(session_id)`。 + +`mark_rollout_discarded` 用于表达 skipped、timeout、final cancelled、旧 session 被新 session 替换等语义事件。 + +状态约束： + +1. `mark_commit_failed` 只允许从 `RolloutRunning` 进入 `ToBeReleased`。 +2. `mark_rollout_discarded` 允许从 `RolloutRunning` 或 `RolloutFinished` 进入 `ToBeReleased`。 +3. 两者都是 release-like 事件；session 不存在时返回 `Released`，不创建空 session。 + +### 5.4 training export + +```python +def export_training_trace(session_id: str, prompt_text: str) -> dict: + ... +``` + +成功条件： + +1. session 存在。 +2. session 处于 `RolloutFinished`。 +3. `prompt_text` 能完整命中 session trie。 +4. token-level 字段可以组成训练 trace。 + +返回的 `routed_experts` 是 session 级 `expert_key` 或 `None`。第一版不再返回每个 segment 的 routed experts key +列表，因为 routed experts 在 Store 中按 session 级别维护。 + +成功后： + +```text +RolloutFinished -> TrainRunning +``` + +失败后： + +```text +RolloutFinished -> ToBeReleased -> Released +``` + +失败时进入 `ToBeReleased` 并继续抛出 `ValueError`。 + +约束：`export_training_trace` 必须要求 session 已经是 `RolloutFinished`。不能为了兼容当前代码从 `RolloutRunning` +直接导出；调用侧必须先通过 rollout 完成事件把 session 推进到 `RolloutFinished`。 + +失败行为： + +1. session 不存在时抛 `KeyError`，不能自动创建 session。 +2. session 不是 `RolloutFinished` 时抛 `RuntimeError`，不能导出训练 trace。 +3. session 是 `RolloutFinished` 但 trace 不完整时，先进入 `ToBeReleased`，再抛 `ValueError`。 + +### 5.5 trainer 完成或放弃 + +```python +def mark_train_finished(session_id: str) -> str: + ... + +def mark_train_abandoned(session_id: str) -> str: + ... +``` + +`mark_train_finished` 由 trainer 在确认所有训练消费者都不再依赖 Store 后调用： + +```text +TrainRunning -> TrainFinished -> ToBeReleased -> Released +``` + +`mark_train_abandoned` 由 trainer 在取消、不可恢复 materialize 失败、batch 被替换等情况下调用。调用前提仍然是 +trainer 已经确认不会再访问 Store： + +```text +TrainRunning -> ToBeReleased -> Released +``` + +Trace Store 不维护 rank 级 materialize 状态。 + +状态约束： + +1. `mark_train_finished` 只允许从 `TrainRunning` 调用。 +2. `mark_train_abandoned` 只允许从 `TrainRunning` 调用。 +3. 两者都是训练侧消费结束事件；session 不存在时返回 `Released`，不创建空 session，用于兼容重复上报或释放后的迟到事件。 + +### 5.6 不提供外部 release API + +外部模块不能直接调用 `release(session_id)`。物理释放只能由 Trace Store 内部在状态进入 `ToBeReleased` 后触发。 + +当前已有的 actor 方法： + +```python +def release(session_id: str): + ... +``` + +实现时删除 public actor method `release(session_id)`，改成内部私有方法 `_release_session(session_id)`。对外暴露的 API +只能是语义事件，例如 `mark_rollout_status`、`mark_commit_failed`、`mark_rollout_discarded`、`mark_train_finished`、 +`mark_train_abandoned`。 + +## 6. 状态转换和释放触发 + +所有状态写入必须通过统一 helper 完成，避免某条路径只改状态但漏掉 release。第一版暂不在 `_set_state` +里实现模块级目标状态转换表；具体语义 API 在调用 `_set_state` 前用自己的入口状态校验表达约束。后续如果状态转换路径继续变多， +可以再引入 `next_state -> allowed_previous_states` 的全局表。 + +```python +def _set_state( + self, + session_id: str, + next_state: TraceState, +) -> TraceState: + trie = self.sessions.get(session_id) + if trie is None: + raise KeyError(f"Trace session {session_id!r} does not exist.") + trie.state = next_state + trie.touch() + self._maybe_release(session_id) + return next_state +``` + +约束： + +1. 任何进入 `ToBeReleased` 的 API 都必须走 `_set_state(...)` 或等价 helper。 +2. `_set_state` 每次状态更新后都调用 `_maybe_release(session_id)`。 +3. `_maybe_release` 内部只在状态为 `ToBeReleased` 时释放，所以可以在每次状态变更后安全调用。 +4. `_set_state` 不能调用 `get_or_create`，状态事件不能创建空 session。 +5. 第一版不在 `_set_state` 内实现模块级目标状态转换表；每个语义 API 自己校验允许的入口状态。 +6. release-like 事件如果允许 missing session no-op，必须在外层语义 API 中处理，不进入 `_set_state`。 +7. 正常路径 `mark_train_finished` 需要先记录 `TrainFinished`，再进入 `ToBeReleased`： + +```python +def mark_train_finished(self, session_id: str) -> str: + if session_id not in self.sessions: + return TraceState.RELEASED.value + self._set_state( + session_id, + TraceState.TRAIN_FINISHED, + ) + return self._set_state( + session_id, + TraceState.TO_BE_RELEASED, + ).value +``` + +8. 异常放弃路径直接进入 `ToBeReleased`： + +```python +def mark_train_abandoned(self, session_id: str) -> str: + if session_id not in self.sessions: + return TraceState.RELEASED.value + return self._set_state( + session_id, + TraceState.TO_BE_RELEASED, + ).value +``` + +### 6.1 合法状态转换表 + +下表是设计约束，不表示当前代码中有模块级全局转换表。第一版由各语义 API 在入口处校验来源状态。 + +| 事件/API | 允许来源状态 | 目标状态 | +| --- | --- | --- | +| 首次 `insert` / `search` / `keys` | session 不存在 | `RolloutRunning` | +| `insert` | `RolloutRunning` | `RolloutRunning` | +| `mark_rollout_status(COMPLETED)` | `RolloutRunning` | `RolloutFinished` | +| `mark_rollout_status(ABORTED, enable_partial_rollout=True)` | `RolloutRunning` | `RolloutRunning` | +| `mark_rollout_status(ABORTED, enable_partial_rollout=False)` | `RolloutRunning` | `ToBeReleased` | +| `mark_rollout_status(FAILED/FILTERED/EXPIRED)` | `RolloutRunning` | `ToBeReleased` | +| `mark_commit_failed` | `RolloutRunning` | `ToBeReleased` | +| `mark_rollout_discarded` | `RolloutRunning` / `RolloutFinished` | `ToBeReleased` | +| `export_training_trace` 成功 | `RolloutFinished` | `TrainRunning` | +| `export_training_trace` trace 不完整 | `RolloutFinished` | `ToBeReleased` | +| `mark_train_finished` | `TrainRunning` | `TrainFinished` -> `ToBeReleased` | +| `mark_train_abandoned` | `TrainRunning` | `ToBeReleased` | +| `_maybe_release` | `ToBeReleased` | `Released`，实现上删除 session | + +## 7. 释放语义 + +唯一物理释放入口： + +```python +def _maybe_release(session_id: str) -> None: + trie = self.sessions.get(session_id) + if trie is None: + return + if trie.state != TraceState.TO_BE_RELEASED: + return + self._release_session(session_id, trie) +``` + +`_release_session(session_id, trie)` 是内部物理释放方法，负责： + +1. 从 `trie.expert_key` 获取该 session 唯一 routed experts object key。 +2. 删除该 key 在 `self.objects` 中对应的 ref，并调用 `_free_ray_refs`。 +3. 调用 `trie.release()` 释放 session tree 中其他可能残留的 `ObjectRef`。 +4. 从 `self.sessions` 删除 `session_id`。 + +约束： + +1. `_maybe_release` 必须幂等。 +2. session 不存在时直接返回。 +3. 非 `ToBeReleased` 不释放。 +4. 释放粒度是整个 session tree。 +5. routed experts object 的释放由 `RolloutTraceStore` 负责，`Trie.release()` 不直接访问 `self.objects`。 +6. `Trie.release()` 仍保留对 tree node value 的递归 `_free_ray_refs`，作为防御性清理，避免非 routed experts 字段里仍残留 `ObjectRef`。 +7. 删除 session metadata 前必须保证该 session 记录的 object key 已经从 `self.objects` 中删除。 +8. `Released` 不需要持久保存；删除 session metadata 就代表已释放。 +9. lifecycle 释放只允许 full-session release。现有 `Trie.release(key=None)` 的 `key` 参数需要删除，或保留为非 lifecycle 私有能力；生命周期路径不能做 subtree release。 + +## 8. 现有方法兼容策略 + +### 8.1 `insert` + +`insert(session_id, key, value, routed_experts=None)` 当前会自动创建 session。 + +第一版保留该行为： + +1. session 不存在时创建 `RolloutRunning` session。 +2. session 为 `RolloutRunning` 时允许写入。 +3. 如果 `value` 是 `TokenizedSegment`，`value.expert_key` 必须已经是 `str | None`，不能是 `ray.ObjectRef`。 +4. routed experts `ObjectRef` 的写入只能走 `insert(..., routed_experts=obj_ref)`。 +5. `insert` 不处理旧 `TokenizedSegment` 覆盖后的 object ref 释放；第一版先假设写入路径不会产生需要清理的旧 ref。 +6. session 已进入 `RolloutFinished` / `TrainRunning` / `ToBeReleased` 后，`insert` 必须记录 error 日志并跳过写入，避免已完成或待释放 session 被继续污染，同时不打断上层训练/rollout 主流程。 + +### 8.2 `search` / `keys` + +当前 `search` / `keys` 也会通过 `get_or_create` 自动创建 session。第一版继续保留这个兼容行为。 + +实现约束： + +1. session 不存在时创建空 `RolloutRunning` session。 +2. session 处于 `RolloutRunning` / `RolloutFinished` / `TrainRunning` 时允许读取。 +3. session 处于 `ToBeReleased` 时不应继续读取；实现上 `_maybe_release` 在进入该状态后立即触发释放，故此情形窗口极窄。防御路径记录 error 日志，`keys` 返回 `[]`，`search` 返回 `("", [])`。 + +### 8.3 `get_objects` + +`get_objects(keys)` 的实现从 actor 级 `self.objects` 取对象。 + +object key 的含义： + +1. `self.objects[object_key]` 保存实际 `ray.ObjectRef`。 +2. `TokenizedSegment.expert_key` 保存 object key 字符串。 +3. `export_training_trace` 返回 session 级 routed experts object key 或 `None`。 +4. 训练侧如果需要实际 ref，再调用 `get_objects([key])`。 +5. `None` 是合法的 routed experts 占位，表示该 session 没有 routed experts object；训练侧不能把 `None` 传给 `get_objects`。 + +`get_objects` 行为： + +1. 只接受非空字符串 object key。 +2. object key 必须存在于 `self.objects`。 +3. 缺失 key 必须抛 `KeyError`，不能静默跳过。 + +## 9. 孤儿 Session TTL 清理 + +### 9.1 问题背景 + +rollout worker crash 后，若未调用任何生命周期 API，session 会永远卡在 `RolloutRunning`，`self.objects` 中该 session 引用的 `ray.ObjectRef` 持续占用内存。这类 session 称为**孤儿 session（orphan session）**。 + +### 9.2 TTL 机制设计 + +在 `Trie` 中使用 `updated_at` 字段记录最近一次写入或状态变更时间。`RolloutTraceStore` 暴露一个 `gc_stale_sessions(ttl_seconds: float)` 方法，供外部按需调用（例如由 train controller 或 rollout coordinator 周期性触发）： + +```python +def gc_stale_sessions(self, ttl_seconds: float) -> list[str]: + """释放超过 ttl_seconds 未更新且仍处于 RolloutRunning 的孤儿 session。 + + Returns: + list[str]: 被释放的 session_id 列表，用于日志和告警。 + """ + now = time.time() + stale = [ + sid + for sid, trie in self.sessions.items() + if trie.state == TraceState.ROLLOUT_RUNNING + and (now - trie.updated_at) > ttl_seconds + ] + for sid in stale: + self._set_state( + sid, + TraceState.TO_BE_RELEASED, + ) + return stale +``` + +### 9.3 约束 + +1. `gc_stale_sessions` 只清理 `RolloutRunning` 状态的 session，不触碰 `RolloutFinished` / `TrainRunning`（这些状态可能正在被训练侧使用）。 +2. TTL 值建议由调用方配置（例如 `ttl_seconds = 600`），Trace Store 不内置 TTL 常量。 +3. `gc_stale_sessions` 本身不启动后台线程，不依赖 Ray actor 内部 timer；调用方负责周期性调用。 +4. 返回被释放的 session_id 列表，调用方应记录 WARNING 日志，以便排查 worker crash。 +5. `gc_stale_sessions` 另提供诊断能力：被观测到频繁有孤儿 session 说明 rollout worker crash 率异常，需告警。 + +### 9.4 可观测性补充 + +新增 `list_sessions(state: str | None = None) -> list[dict]` 诊断 API，返回当前所有 session 的状态快照（`session_id`、`state`、`updated_at`、`has_object_ref`）。 +`state` 参数可过滤特定状态，例如 `list_sessions(state="RolloutRunning")` 快速定位孤儿 session。用于外部监控和调试，不触发任何状态变更。 + +## 10. Trainer 接入点 + +trainer 在 train controller 的训练结束位置调用 Trace Store 状态转换 API： + +1. 正常训练消费结束，确认后续不会再访问 Store 后，调用 `mark_train_finished(session_id)`。 +2. 训练取消、batch 放弃、不可恢复 materialize 失败，并确认后续不会再访问 Store 后，调用 `mark_train_abandoned(session_id)`。 +3. Trace Store 不在 trainer 内部分 rank 维度做判断，只接收 train controller 汇总后的事件。 + +### 10.1 主流程接入顺序 + +Trace Store actor 内部 API 完成后，主流程接入按以下顺序推进。 + +1. Rollout 状态接入：在 `ProduceContext.put_generated_group` 中等待业务过滤、ReplayBuffer 过期处理完成后，按每个 `RolloutState` 的最终 `status` 调用 `mark_rollout_status`。这里是 rollout 结果进入系统状态机的统一入口。 +2. Training trace materialize 接入：在 `BaseRLTrainer._prepare_train_data` 中默认对每个 sample 调用 `export_training_trace`。如果 trace 返回 routed experts object key，trainer 只保留 `get_objects([key])` 这个 Ray ref，不在 trainer 里 `ray.get` 实体；真正的 routed experts tensor 只在 `TrainingWorker` 内部消费时 `ray.get`。ReplayBuffer 中的轻量 `RolloutState` 可以不携带 token-heavy 字段，进入 `_prepare_train_data` 后由 Trace Store 回填 `prompt_ids` / `response_ids` / `response_mask` / `logprobs` / `routed_experts`，再复用原 batch 构造逻辑。 +3. 训练完成释放接入：在 `BaseRLTrainer._train_one_batch` 调用 `self.train_controller.fit(...)` 成功返回后，对本 batch 中已 export 的 session 调用 `mark_train_finished`。 +4. 训练放弃释放接入：如果 export 成功后 `_prepare_train_data`、packing 或 `train_controller.fit` 出现不可恢复异常，并且确认后续不会再访问 Store，则对已 export 的 session 调用 `mark_train_abandoned`。 +5. 孤儿 session GC 接入：trainer 主循环按配置频率调用 `gc_stale_sessions(ttl_seconds)`，并对返回的 stale session_id 记录 warning 日志。 + +### 10.2 接入约束 + +1. Trace Store 状态必须以最终进入 ReplayBuffer 或最终被 trainer 消费的状态为准，不能在 rollout worker 单次 generate 内提前标最终状态。 +2. `export_training_trace` 只能由 trainer/data prepare 层调用一次，成功后 session 进入 `TrainRunning`。 +3. TrainingWorker rank 内部不直接调用 Trace Store lifecycle API，也不分别 release session。 +4. `mark_train_finished` 必须在所有 TrainingWorker 完成消费后调用。 +5. `mark_train_abandoned` 只用于确认后续不会再访问 Store 的失败路径。 +6. `gc_stale_sessions` 只清理 `RolloutRunning`，不能作为训练完成释放的替代机制。 + +### 10.3 Trace Store 训练样本的轻量字段契约 + +第一版先不新增独立训练样本结构，仍复用 `RolloutState` 进入 ReplayBuffer 和 trainer。ReplayBuffer 中的 `RolloutState` 保持轻量，token 级训练数据留在 `RolloutTraceStore`： + +1. `RolloutState.session_uid` 是 Trace Store 的 `session_id`。 +2. trainer 默认从 Trace Store materialize 训练数据，不再用 `extra_fields["use_trace_store"]` 在 trainer 内部分支。 +3. `RolloutState.extra_fields["trace_store_prompt_text"]` 是可选字段；如果调用方能拿到与 Trace Store key 完全一致的最终 prompt text，应写入这里，trainer 优先使用它。 +4. trainer 中的 group 是 advantage group，不是 Trace Store session group；group 内每个 rollout sample 都是一次独立 request，必须有独立 `session_uid`。 +5. `RolloutState.response_ids` / `logprobs` / `routed_experts` 在 ReplayBuffer 中可以保持 `None`，不能写成 `session_id` 或 object key。 +6. `_prepare_train_data` 从 Trace Store 导出后，把 token 字段回填到当前 `RolloutState`，让后续 batch 构造继续按 `RolloutState` 字段读取。 +7. trainer 不 `ray.get` routed experts 实体，只把 routed experts 的 Ray ref 放进 `SequenceContext.rollout_routed_experts`。 +8. `RolloutState.response`、`message`、`tools`、`tool_calls` 仍可保留，供 reward、日志、prompt_text 重建或调试使用。 +9. 如果 `trace_store_prompt_text` 不存在，trainer 可以用 `message + response/tool_calls + tools` 重新渲染完整 prompt text；这要求调用方保证这些轻量字段和 SessionServer 写入 Trace Store 时使用的 chat template 一致。 + +这只是过渡方案。后续如果 Trace Store 成为唯一训练数据来源，应新增专门的轻量训练 envelope，替代在 `RolloutState.extra_fields` 中携带协议字段。 + +## 11. 第一版测试计划 + +设计稳定并进入代码实现后，新增 `tests/rl/test_trace_store_lifecycle.py`。 + +测试覆盖： + +1. first insert 创建 `RolloutRunning` session。 +2. `COMPLETED` 进入 `RolloutFinished`。 +3. `FAILED` / `FILTERED` / `EXPIRED` 进入 `ToBeReleased` 并释放。 +4. `ABORTED + enable_partial_rollout=True` 保持 `RolloutRunning`，已有 trie 内容仍可 search。 +5. `ABORTED + enable_partial_rollout=False` 释放 session。 +6. `export_training_trace` 成功后进入 `TrainRunning`。 +7. `export_training_trace` prefix 不完整时释放。 +9. `mark_train_finished` 触发 `TrainFinished -> ToBeReleased -> Released`。 +10. `mark_train_abandoned` 触发释放。 +11. routed experts `ObjectRef` 写入后进入 `RolloutTraceStore.objects`，`TokenizedSegment.expert_key` 保存 object key。 +12. `get_objects` 能从 object key 定位真实 `ObjectRef`。 +13. `routed_experts` 允许 `None`，但 `get_objects([None])` 或缺失 key 必须报错。 +14. session lifecycle release 会删除该 session 记录的 object key，并释放对应 ref。 +15. 非法状态转换报错，例如 `RolloutRunning -> TrainFinished`。 +16. `_maybe_release` 对不存在 session 幂等 no-op；release-like 语义事件对已释放 session 返回 `Released`。 +17. `search` / `keys` 对不存在 session 自动创建空 `RolloutRunning` session。 +18. session 处于 `ToBeReleased` 时调用 `search` / `keys` 记录 error 日志并返回空结果。 +19. `gc_stale_sessions(ttl)` 释放超过 TTL 的 `RolloutRunning` session，不影响 `RolloutFinished` / `TrainRunning` session。 +20. `gc_stale_sessions` 返回被释放的 session_id 列表。 +21. `list_sessions()` 返回所有 session 快照；`list_sessions(state="RolloutRunning")` 过滤指定状态。 + +建议验证命令： + +```bash +python -m unittest tests.rl.test_trace_store_lifecycle +python -m compileall -q xtuner/v1/rl/rollout/trace_store.py tests/rl/test_trace_store_lifecycle.py +``` diff --git a/docs/design/trace_store_lifecycle_simple.md b/docs/design/trace_store_lifecycle_simple.md new file mode 100644 index 0000000000..8399ccce6e --- /dev/null +++ b/docs/design/trace_store_lifecycle_simple.md @@ -0,0 +1,431 @@ +# Rollout Trace Store 生命周期管理简化版 + +## 1. 文档目标 + +本文是 `trace_store_lifecycle.md` 的简化版设计。旧文档保留，用于记录更细的事件和异常分支；本文只解释 Trace Store session 在一次 RL 训练链路中的主生命周期。 + +这一版进一步简化：生命周期核心只保留 `state`。之前讨论过的 `write_txns`、`outcome`、`lease`、`retention` 不作为核心字段。 + +一句话： + +```text +state 决定这条 session 当前能做什么、不能做什么。 +``` + +本版按当前实现边界收敛三个规则： + +1. `RolloutState.status == ABORTED` 且 `enable_partial_rollout=True` 是唯一允许复用老 `session_id` 的续跑路径。它不是失败重试，也不是新 attempt，而是同一条 session 在已有历史轨迹上继续 rollout。 +2. 除上述 partial resume 外，失败重试、超时后重启、worker 失败后重新调度，都必须使用新的 `session_id`。因此当前版本不要求在所有事件里额外携带 `attempt_id` / generation；如果未来允许失败重试复用 `session_id`，必须重新引入 generation fencing。 +3. 多轮 LLM call 在 Trace Store 里按 session 级别处理。一次 agent loop 内任意一轮 LLM call 出现不可恢复失败、commit 失败或协议校验失败，都认为整条 session 不再可用，进入 `ToBeReleased`，由 Store 释放整棵 session tree。当前版本不做节点级 partial release，也不把 failed trace 的 prefix 用于失败重试。 + +## 2. 主状态 + +Trace Store session 保留 6 个互斥主状态。 + +| 状态 | 含义 | 能否物理释放 | +| --- | --- | --- | +| `RolloutRunning` | rollout 还没有形成最终可训练结果；包括 agent loop 正在运行、多轮 LLM call 继续写入，或 `ABORTED + enable_partial_rollout=True` 后等待 resume | 不能 | +| `RolloutFinished` | rollout 已经完成且未被过滤，等待训练侧 export | 不能 | +| `TrainRunning` | 训练侧已经 export trace，TrainingWorker 仍可能依赖 Store object refs | 不能 | +| `TrainFinished` | 训练侧已经完成 materialize / 消费，不再依赖 Store object refs | 不能；下一步进入 `ToBeReleased` | +| `ToBeReleased` | session 已经不再服务 rollout 或 training，等待 Trace Store 统一释放 | 可以由 Store 尝试释放 | +| `Released` | Store 已释放 object refs，并删除 session metadata；实现上可以表现为 session 不再存在 | 终态 | + +几个关键约定： + +- `ABORTED + enable_partial_rollout=True` 保持 `RolloutRunning`，并复用老 `session_id` 做 partial resume。 +- `ABORTED + enable_partial_rollout=False`、任一不可恢复 LLM call 失败、`commit_response` 失败都会使整个 session 进入 `ToBeReleased`。 +- `RolloutFinished` 在本文里不是泛指“agent loop 停了”，而是“这条 rollout 已经完成且未被过滤”。 +- `TrainRunning` 本身就表达了训练侧还在使用 Store 对象，因此不再单独维护 `lease = active`。 +- `ToBeReleased` 是 release pending。它不是资源已经释放，只表示可以进入 Store 统一清理流程。 +- Trace Store 的状态只决定是否保留 / 释放历史轨迹，不决定 ReplayBuffer / producer 是否继续采样或重试这个样本。 +- 如果 ReplayBuffer / producer 因失败、超时、worker 崩溃等原因决定重跑同一个样本，必须创建新的 `session_id`；旧 session 只负责按自身状态完成释放。 + +## 3. 状态机 + +```mermaid +stateDiagram-v2 + [*] --> RolloutRunning: 首次请求进入 Store + + RolloutRunning --> RolloutRunning: 多轮 LLM call 继续写入 + RolloutRunning --> RolloutRunning: ABORTED 且 enable_partial_rollout=True + RolloutRunning --> RolloutFinished: COMPLETED 且未被过滤 + RolloutRunning --> ToBeReleased: FAILED / FILTERED / EXPIRED + RolloutRunning --> ToBeReleased: ABORTED 且 enable_partial_rollout=False + RolloutRunning --> ToBeReleased: 任一不可恢复 LLM call 或 commit 失败 + + RolloutFinished --> TrainRunning: export_training_trace 成功 + RolloutFinished --> ToBeReleased: 训练前被取消或旧 session 被新 session 替换 + + TrainRunning --> TrainFinished: 所有 TrainingWorker 已脱离 Store 依赖 + TrainRunning --> ToBeReleased: 训练侧放弃消费且不会再访问 Store + + TrainFinished --> ToBeReleased: 训练消费已结束 + ToBeReleased --> Released: Store 释放 session tree + Released --> [*] +``` + +状态机里的线表示状态转移原因。进入 `ToBeReleased` 不代表已经释放，只代表这条 session 已经不应该继续参与 rollout 或 training。 + +## 4. Rollout 阶段 + +Rollout 阶段拆成两段： + +```text +RolloutRunning -> RolloutFinished +RolloutFinished -> TrainRunning +``` + +第一段由 rollout 侧返回结果和 ReplayBuffer / producer 的判定共同决定；第二段由训练侧 export 是否成功决定。 + +### 4.1 RolloutRunning -> RolloutFinished / ToBeReleased + +```mermaid +flowchart TD + A["RolloutRunning"] --> B{"RolloutState.status"} + B -- "COMPLETED 且未被过滤" --> C["RolloutFinished"] + B -- "COMPLETED 后被过滤" --> D["ToBeReleased"] + B -- "ABORTED 且 enable_partial_rollout=True" --> A + B -- "ABORTED 且 enable_partial_rollout=False" --> D + B -- "EXPIRED" --> D + B -- "FAILED / FILTERED" --> D + A -- "任一不可恢复 LLM call / commit 失败" --> D +``` + +`RolloutRunning` 表示 session 仍属于 rollout 生产阶段。多轮 agent loop 中，多个 LLM call 会持续写入同一棵 session tree。当前简化版采用 session 级 all-or-nothing 语义：只要其中任意一轮 LLM call 出现不可恢复失败，或者某次 `commit_response` 失败，整条 session 都进入 `ToBeReleased`。 + +注意：下面的状态映射只影响 Rollout Trace Store 是否保留历史轨迹，不影响 ReplayBuffer / producer 是否继续采样。`ABORTED + enable_partial_rollout=True` 是唯一会复用老 `session_id` 的路径；其他失败、超时、worker 崩溃后的重跑，都必须从 ReplayBuffer 重新采样原始样本并使用新的 `session_id`。 + +`RolloutState.status` 和 Trace Store 状态的关系如下。 + +| `RolloutState.status` | Trace Store 状态转移 | 说明 | +| --- | --- | --- | +| `COMPLETED` 且未被过滤 | `RolloutRunning -> RolloutFinished` | completed rollout 可以进入训练读取阶段。 | +| `FILTERED` | `RolloutRunning -> ToBeReleased` | `producer.put_generated_group` 会对 completed group 做 `is_valid_sample_fn`，不通过时会置为 `FILTERED`；Trace Store 不需要继续保留历史轨迹。 | +| `ABORTED` 且 `enable_partial_rollout=True` | 保持 `RolloutRunning` | 复用老 `session_id`，基于已有历史轨迹继续 partial rollout。 | +| `ABORTED` 且 `enable_partial_rollout=False` | `RolloutRunning -> ToBeReleased` | 不做 partial resume；Trace Store 历史轨迹不再复用，可以释放。 | +| `EXPIRED` | `RolloutRunning -> ToBeReleased` | 过期样本不应继续复用旧 Trace Store 历史轨迹；采样侧是否重跑由 ReplayBuffer / producer 决定。 | +| `FAILED` | `RolloutRunning -> ToBeReleased` | failed rollout 不形成可训练 trace，全部进入释放等待。 | +| `FILTERED` | `RolloutRunning -> ToBeReleased` | 样本业务上不可训练。 | +| `INIT` | 保持 `RolloutRunning` | 尚未形成有效 rollout 返回，不触发释放。 | +| `ARCHIVED` | 第一版暂不处理 | 如果传入 Trace Store rollout 状态 API，直接报错。 | + +因此，`COMPLETED` 表示未被过滤的可训练结果，进入 `RolloutFinished`；过滤后的样本应以 `FILTERED` 上报，进入 `ToBeReleased`。`FAILED` 和 `EXPIRED` 对 Trace Store 来说也直接进入 `ToBeReleased`。 + +多轮 LLM call 的失败规则如下。 + +| 事件 | Trace Store 状态转移 | 说明 | +| --- | --- | --- | +| 第 1 到第 N-1 轮 LLM call 成功，第 N 轮 LLM call 出现不可恢复失败 | `RolloutRunning -> ToBeReleased` | 已写入的历史轨迹不再作为失败重试的 prefix 复用，释放整棵 session tree。 | +| 第 N 轮 rollout 返回 `ABORTED`，且 `enable_partial_rollout=True` | 保持 `RolloutRunning` | 这是 partial resume 路径，后续继续使用老 `session_id`。 | +| 某轮 LLM call 返回字段不满足 completed contract | `RolloutRunning -> ToBeReleased` | rollout worker 应将状态置为 `FAILED`；Trace Store 按 failed session 处理。 | +| 某轮 `commit_response` 失败 | `RolloutRunning -> ToBeReleased` | 这条 session 不能再进入 `RolloutFinished`；需要清理已 staged 或已部分挂载的对象。 | +| producer 因失败、超时、worker 崩溃决定重跑同一个样本 | 新建 `session_id`，旧 session 进入 `ToBeReleased` | 新旧 session 通过不同 `session_id` 隔离，旧事件不能影响新 session。 | + +#### 4.1.1 Worker completed contract + +rollout worker 返回 `Status.COMPLETED` 前必须满足 completed contract： + +| 字段 / 条件 | 要求 | +| --- | --- | +| `response_ids` | 当 `return_token_ids=True` 时必须存在且非空 | +| `logprobs` | 当 `return_logprob=True` 时必须存在且和 `response_ids` 对齐 | +| `routed_experts` | 当启用 routed experts 采集时必须存在 | +| `response` | 正常 completed 输出应有可用文本 | +| 长度一致性 | token、logprob、routed experts 等 token-level 字段必须能和本次生成对齐 | + +如果这些条件不满足，应该由 rollout worker 把 `RolloutState.status` 标记为 `FAILED`，并设置 `error_msg`。Trace Store 可以保留防御性校验，但不应该把字段缺失当成正常 commit 分支处理。 + +注意：这个 contract 约束的是 rollout worker / SessionServer 写入 Trace Store 前的生成结果。进入 ReplayBuffer 的轻量 `RolloutState` 可以不再携带 `response_ids` / `logprobs` / `routed_experts`；trainer 默认通过 `session_uid` 从 Store 导出训练 trace，并在 `_prepare_train_data` 内回填这些训练字段。 + +#### 4.1.2 ABORTED + +`Status.ABORTED` 是否保留历史轨迹，取决于 `enable_partial_rollout`。 + +```mermaid +flowchart LR + A["RolloutRunning"] --> B["rollout worker 返回 ABORTED"] + B --> C{"enable_partial_rollout"} + C -- "True" --> D["RolloutRunning"] + C -- "False" --> E["ToBeReleased"] +``` + +`enable_partial_rollout=True` 时，`ABORTED` 是同一条 session 的 partial resume 路径，Trace Store 保留历史轨迹，后续继续使用老 `session_id`。 + +`enable_partial_rollout=False` 时，`ABORTED` 不再复用历史轨迹，进入 `ToBeReleased`。ReplayBuffer / producer 如果要重跑，必须基于原始样本创建新的 `session_id`。 + +注意：只有 `RolloutState.status == ABORTED` 且 `enable_partial_rollout=True` 才能复用老 `session_id`。`FAILED`、`EXPIRED`、`commit_response` 失败、协议校验失败、worker 崩溃后的失败重试，都不能复用旧 `session_id`。 + +### 4.2 RolloutFinished -> TrainRunning + +`RolloutFinished` 表示 rollout 侧已经提供了一条可训练样本，但训练侧还没有开始消费 Trace Store。 + +```mermaid +flowchart TD + A["RolloutFinished"] --> B{"训练侧是否 export 成功"} + B -- "export_training_trace 成功" --> C["TrainRunning"] + B -- "等待训练侧拉取" --> A + B -- "trace 不完整 / LLMTrace 缺失 / prompt 不匹配" --> D["ToBeReleased"] + B -- "训练前取消 / 旧 session 被替换" --> D +``` + +`RolloutFinished -> TrainRunning` 的条件： + +1. 训练侧可以从 LLMTrace 渲染出完整 `prompt_text`。 +2. `export_training_trace(session_id, prompt_text)` 能完整命中 Trace Store。 +3. token-level 字段可以组成训练样本。 +4. TrainingWorker 即将消费 Store 返回的 object keys / object refs。 + +如果 export 前发现 LLMTrace 缺失、prompt render 失败、Trace Store prefix 不完整、object key 缺失，说明这条样本已经无法训练，应进入 `ToBeReleased`。 + +## 5. 训练阶段 + +训练阶段用 `TrainRunning` 和 `TrainFinished` 表达原来 lease 想表达的保护语义。 + +`TrainRunning` 的含义是：训练侧已经通过 `export_training_trace` 接管了这条 session，并且 TrainingWorker 可能还会通过 object keys / object refs 从 Trace Store 读取 token-level 数据。只要还处在这个状态，Trace Store 就不能释放 session tree。 + +`TrainFinished` 是 Trace Store 视角的完成态：训练侧已经把需要的数据 materialize 到本地 batch / tensor，后续不会再访问 Trace Store 里的 object refs。它不要求 optimizer step、backward 或参数更新已经完成，只要求 Store 对这条 session 已经没有被训练侧继续读取的可能。 + +```mermaid +sequenceDiagram + participant DB as Data Builder + participant TS as Trace Store + participant TW0 as TrainingWorker rank0 + participant TW1 as TrainingWorker rank1 + + DB->>TS: export_training_trace(session_id, prompt_text) + TS->>TS: state = TrainRunning + TS-->>DB: TrainingTrace + object keys + + DB->>TW0: batch + object keys + DB->>TW1: batch + object keys + + TW0->>TS: get_objects(keys) + TW1->>TS: get_objects(keys) + TW0->>TS: ack_materialized(rank0) + TW1->>TS: ack_materialized(rank1) + + TS->>TS: state = TrainFinished + TS->>TS: state = ToBeReleased + TS->>TS: _maybe_release(session_id) +``` + +### 5.1 TrainRunning -> TrainFinished / ToBeReleased + +```mermaid +flowchart TD + A["TrainRunning"] --> B{"训练侧是否仍可能访问 Store"} + B -- "仍可能 get_objects / retry materialize" --> A + B -- "所有 rank 已 materialize 并 ack" --> C["TrainFinished"] + B -- "训练任务取消且消费者已撤销" --> D["ToBeReleased"] + B -- "materialize 不可恢复失败且消费者已停止" --> D + B -- "batch / session 被替换且旧消费者已停止" --> D + C --> E["ToBeReleased"] +``` + +正常路径是： + +```text +TrainRunning -> TrainFinished -> ToBeReleased +``` + +异常放弃路径是： + +```text +TrainRunning -> ToBeReleased +``` + +两条路径的区别是：`TrainFinished` 表示训练侧成功拿走了这条 session 需要的数据；`TrainRunning -> ToBeReleased` 表示训练侧已经放弃这条 session，且系统确认不会再有 TrainingWorker 访问 Store。 + +| 事件 | 状态转移 | 说明 | +| --- | --- | --- | +| 所有 TrainingWorker rank 都完成 `get_objects`，并确认本地 batch / tensor 已 materialize | `TrainRunning -> TrainFinished` | 这是训练阶段的正常完成路径。 | +| 部分 rank 还没 ack materialized | 保持 `TrainRunning` | 任意 rank 仍可能读取 Store，不能释放。 | +| `get_objects` 临时失败，但训练侧会重试 | 保持 `TrainRunning` | 可恢复失败不应该触发释放，否则重试会读不到 object refs。 | +| 训练任务取消，并且 Data Builder / Trainer 已撤销这条 batch，确认没有 worker 会继续读取 Store | `TrainRunning -> ToBeReleased` | 这条 trace 不再训练，也不会被继续读取。 | +| TrainingWorker materialize 发现不可恢复错误，并且训练侧决定放弃该 session | `TrainRunning -> ToBeReleased` | 例如 object key 永久缺失、object ref 已损坏、反序列化失败、字段 shape 无法组成 batch。 | +| 当前 batch / session 被新 batch / session 替换，旧消费者已全部停止 | `TrainRunning -> ToBeReleased` | 旧 session 不应继续占用 Store 资源。 | +| 训练进程崩溃但调度器会用同一条 trace 重试 | 保持 `TrainRunning` | 因为后续重试仍依赖 Store。 | +| 训练进程崩溃且调度器明确放弃该 session | `TrainRunning -> ToBeReleased` | 前提仍然是确认不会再有消费者访问 Store。 | + +进入 `ToBeReleased` 的判断条件不是“训练出现错误”，而是“训练侧已经不再依赖 Store”。因此，训练阶段的释放判断必须先回答两个问题： + +1. 是否还有 TrainingWorker 可能持有 object keys / object refs？ +2. 是否还有 Data Builder / Trainer 可能基于这条 session 重试 materialize？ + +只要任一问题答案是“是”，状态都必须保持 `TrainRunning`。只有两个答案都是否，才允许进入 `TrainFinished` 或 `ToBeReleased`。 + +### 5.2 训练阶段进入 ToBeReleased 的情况 + +`TrainRunning` 下需要进入 `ToBeReleased` 的情况可以归纳为三类。 + +| 类别 | 典型情况 | +| --- | --- | +| 训练侧主动放弃 | step 取消、trainer 停止、batch 被撤销、旧 session 被替换 | +| materialize 不可恢复失败 | object key 永久缺失、object ref 损坏、反序列化失败、shape / length 校验无法组成 batch | +| 重试窗口结束 | worker 崩溃后不再重试、materialize 超时后被调度器放弃 | + +这些情况有一个共同前提：**所有训练侧消费者都已经停止，或者被调度器明确撤销**。如果还有任何消费者可能继续读取 Store，即使训练看起来已经失败，也必须保持 `TrainRunning`。 + +### 5.3 训练阶段不应该释放的情况 + +- `TrainRunning` 表示训练侧可能还没有 materialize 完，Trace Store 不能释放。 +- 单个 worker 还在 `get_objects`、等待 `get_objects` 或准备重试时，不能释放。 +- 部分 rank 已经 materialize，但其他 rank 还没有确认时，不能释放。 +- 训练侧只是 backpressure、队列堆积或临时不可用时，不能释放。 +- materialize 失败但调度器还会用同一条 session 重试时，不能释放。 +- 只有确认 TrainingWorker 已经 materialize 完，或者训练消费失败但所有消费者都不会再依赖 Store，才能离开 `TrainRunning`。 + +## 6. 何时进入 ToBeReleased + +`ToBeReleased` 是统一释放等待状态。正常完成和异常退出都会进入这里；具体原因先由调用方日志或上游状态表达。 + +```mermaid +flowchart TD + RR["RolloutRunning"] --> TR["ToBeReleased"] + RF["RolloutFinished"] --> TR + TRun["TrainRunning"] --> TR + TF["TrainFinished"] --> TR + TR --> R["Released"] +``` + +从不同状态进入 `ToBeReleased` 的原因如下。 + +| 来源状态 | 进入 `ToBeReleased` 的情况 | +| --- | --- | +| `RolloutRunning` | rollout failed / skipped / final cancelled / timeout | +| `RolloutRunning` | rollout worker 返回 `FAILED` | +| `RolloutRunning` | rollout worker 返回 `ABORTED`，且 `enable_partial_rollout=False` | +| `RolloutRunning` | rollout_state 变为 `EXPIRED` | +| `RolloutRunning` | `commit_response` 失败，staged objects 无法形成完整 committed trace | +| `RolloutRunning` | 失败、超时或 worker 崩溃后重跑，并创建了新的 `session_id`，旧 session 不再继续写入 | +| `RolloutRunning` | completed rollout 未通过 reward / filter / rule 判定，不能进入训练 | +| `RolloutFinished` | 训练前被取消、过期且被明确放弃、旧 session 被替换 | +| `RolloutFinished` | LLMTrace 缺失、prompt render 失败、Trace Store 字段不完整，导致无法构造训练样本 | +| `TrainRunning` | TrainingWorker materialize 不可恢复失败，且训练侧确认放弃这条 session | +| `TrainRunning` | object key 永久缺失、object ref 损坏、反序列化失败，且不会再重试 | +| `TrainRunning` | 训练任务取消、batch 被撤销、旧 session 被替换，且所有消费者都不会再依赖 Store | +| `TrainRunning` | materialize 超时或 worker 崩溃后，调度器明确放弃该 session | +| `TrainFinished` | 训练侧已经完成 materialize，不再依赖 Store object refs | + +以下情况不能直接进入 `ToBeReleased`： + +| 情况 | 处理方式 | +| --- | --- | +| `Status.ABORTED` 且 `enable_partial_rollout=True` | 保持 `RolloutRunning`，继续使用老 `session_id` resume | +| TrainingWorker 仍可能依赖 Store object refs | 保持 `TrainRunning` | +| materialize / get_objects 失败但训练侧还会重试 | 保持 `TrainRunning` | +| 只有部分 TrainingWorker rank 完成 materialize | 保持 `TrainRunning` | +| 失败、超时、worker 崩溃后的重跑复用旧 `session_id` | 不允许；必须新建 `session_id` | + +## 7. 释放判断 + +`Released` 不是外部模块直接设置的状态，而是 `_maybe_release(session_id)` 的结果。 + +```mermaid +flowchart TD + A["_maybe_release(session_id)"] --> B{"session 是否存在"} + B -- "否" --> N["直接返回"] + B -- "是" --> C{"state 是否为 ToBeReleased"} + C -- "否" --> K["继续保留"] + C -- "是" --> R1["收集 session tree 和 staged refs"] + R1 --> R2["释放 staged / orphan object refs"] + R2 --> R3["释放 committed object refs"] + R3 --> R4["删除 session tree 和 metadata"] + R4 --> R5["概念上进入 Released"] +``` + +对应伪代码： + +```python +def _maybe_release(session_id: str) -> None: + session = sessions.get(session_id) + if session is None: + return + if session.state != "ToBeReleased": + return + _release_staged_objects(session_id) + _release_committed_objects(session) + _delete_session_tree(session) + _delete_session_metadata(session_id) +``` + +这里不再检查 `lease` 或 `retention`，因为它们已经被状态表达： + +- 训练侧仍可能依赖 Store object refs 时，状态必须停留在 `TrainRunning`。 +- 多轮 LLM call 仍在继续写入，或 `ABORTED + enable_partial_rollout=True` 等待 resume 时，状态必须停留在 `RolloutRunning`。 +- 只有进入 `ToBeReleased`，才表示这些保留理由都已经消失。 + +释放必须覆盖两类对象： + +1. 已经挂到 session tree 上的 committed objects。 +2. `commit_response` 过程中已经 staged / pinned、但还没有完整挂到 session tree 上的 objects。 + +因此实现上需要有 staged object registry，或者保证 `commit_response` 具备原子 rollback 能力。否则一次半提交失败可能留下 session tree 看不到的孤儿 object refs。 + +## 8. 常见路径 + +### 8.1 正常训练路径 + +```mermaid +flowchart LR + A["RolloutRunning"] --> B["确认可训练"] + B --> C["RolloutFinished"] + C --> D["export_training_trace"] + D --> E["TrainRunning"] + E --> F["worker 全部 materialize"] + F --> G["TrainFinished"] + G --> H["ToBeReleased"] + H --> I["_maybe_release"] + I --> J["Released"] +``` + +### 8.2 rollout 阶段不可恢复失败 + +```mermaid +flowchart LR + A["RolloutRunning"] --> B["failed / skipped / final cancelled"] + B --> C["ToBeReleased"] + C --> D["_maybe_release"] + D --> E["Released"] +``` + +### 8.3 rollout 完成后不可训练 + +```mermaid +flowchart LR + A["RolloutRunning"] --> B["确认可训练失败"] + B --> C["ToBeReleased"] + C --> D["_maybe_release"] + D --> E["Released"] +``` + +### 8.4 aborted 处理 + +```mermaid +flowchart LR + A["RolloutRunning"] --> B["ABORTED"] + B --> C{"enable_partial_rollout"} + C -- "True" --> D["RolloutRunning"] + C -- "False" --> E["ToBeReleased"] + E --> F["_maybe_release"] + F --> G["Released"] +``` + +`enable_partial_rollout=True` 时，`ABORTED` 是 partial resume 路径，继续使用老 `session_id`。`enable_partial_rollout=False` 时，`ABORTED` 对 Trace Store 是释放路径；采样侧如果要重跑，必须重新采样原始样本并创建新的 `session_id`。 + +## 9. 设计约束 + +1. `state` 是互斥枚举，同一 session 同一时刻只能处于一个主状态。 +2. Trace Store 的释放粒度是整个 session tree，不做单个 LLM call 粒度释放。 +3. 外部模块不能直接调用物理 `release(session_id)`。 +4. 外部模块只能上报语义事件，例如 rollout 可训练、rollout 丢弃、训练消费完成、训练消费失败。 +5. Trace Store actor 是唯一能执行物理释放的地方。 +6. `_maybe_release` 必须幂等，可以在每个事件后调用。 +7. `ABORTED + enable_partial_rollout=True` 是唯一允许复用老 `session_id` 的 partial resume 路径。 +8. 失败重试、超时重跑、worker 崩溃后重跑都必须使用新的 `session_id`。 +9. 任一不可恢复 LLM call 失败、`FAILED`、`EXPIRED`、`ABORTED + enable_partial_rollout=False` 或 `commit_response` 失败，都会使整个 session 进入 `ToBeReleased`。 +10. `EXPIRED` 对 Trace Store 是释放路径，进入 `ToBeReleased`；是否重新采样由 ReplayBuffer / producer 决定。 +11. `TrainRunning` 不能释放；离开 `TrainRunning` 的事件必须保证 TrainingWorker 不再依赖 Store object refs。 +12. `ToBeReleased` 是唯一允许进入物理释放流程的状态。 +13. 当前版本只支持 `ABORTED + enable_partial_rollout=True` 这一种 partial resume；如果后续要支持 failed prefix resume，需要重新设计可恢复 prefix、resume TTL 和 generation fencing。 diff --git a/xtuner/v1/data_proto/rl_data.py b/xtuner/v1/data_proto/rl_data.py index 1e0091749a..8ccd5caeb6 100644 --- a/xtuner/v1/data_proto/rl_data.py +++ b/xtuner/v1/data_proto/rl_data.py @@ -22,6 +22,9 @@ logger = get_logger() +USE_TRACE_STORE_KEY = "use_trace_store" +TRACE_STORE_PROMPT_TEXT_KEY = "trace_store_prompt_text" + class SampleParams(BaseModel): model_config = ConfigDict(extra="forbid") diff --git a/xtuner/v1/data_proto/sequence_context.py b/xtuner/v1/data_proto/sequence_context.py index 51914f9d56..3f88b96a2c 100644 --- a/xtuner/v1/data_proto/sequence_context.py +++ b/xtuner/v1/data_proto/sequence_context.py @@ -1,5 +1,5 @@ # Copyright (c) OpenMMLab. All rights reserved. -from typing import cast +from typing import Any, cast import torch from torch.distributed.device_mesh import DeviceMesh @@ -48,7 +48,7 @@ class SequenceContext: num_img_tokens: list[list[int]] | None # moe routed_experts - rollout_routed_experts: torch.Tensor | None + rollout_routed_experts: Any # Private backing attributes for SP shard reconstruction _raw_input_ids: torch.LongTensor | None diff --git a/xtuner/v1/rl/agent_loop_manager/sampler.py b/xtuner/v1/rl/agent_loop_manager/sampler.py index e5b0170bbe..188737a345 100644 --- a/xtuner/v1/rl/agent_loop_manager/sampler.py +++ b/xtuner/v1/rl/agent_loop_manager/sampler.py @@ -107,9 +107,10 @@ def sample_from_dataloader(self) -> list[RolloutState]: if XTUNER_DETERMINISTIC: new_data.message_uid = message_uid new_data.uid = uid_base + item_idx - new_data.session_uid = new_data.uid else: new_data.uid = uuid4().int + if new_data.session_uid is None: + new_data.session_uid = new_data.uid group_data.append(new_data) self._consumed_samples += 1 return cast(list[RolloutState], group_data) diff --git a/xtuner/v1/rl/trainer/worker.py b/xtuner/v1/rl/trainer/worker.py index 4558cd6dc8..abecf8036f 100644 --- a/xtuner/v1/rl/trainer/worker.py +++ b/xtuner/v1/rl/trainer/worker.py @@ -406,7 +406,9 @@ def compute_ref_logprobs( return ref_logprobs_list def _add_rollout_routed_experts( - self, seq_ctx: SequenceContext, rollout_routed_experts: torch.Tensor | list[torch.Tensor | ray.ObjectRef] + self, + seq_ctx: SequenceContext, + rollout_routed_experts: torch.Tensor | list[torch.Tensor | ray.ObjectRef] | ray.ObjectRef, ): language_cfg = ( self.config.model_cfg.text_config @@ -415,6 +417,8 @@ def _add_rollout_routed_experts( ) to_free_routed_expert_refs: list[ray.ObjectRef] = [] + if isinstance(rollout_routed_experts, ray.ObjectRef): + rollout_routed_experts = ray.get(rollout_routed_experts) if isinstance(rollout_routed_experts, list): # list[n,l,e] out_rollout_routed_expert = [] @@ -433,6 +437,12 @@ def _add_rollout_routed_experts( else: rollout_routed_expert_refs = rollout_routed_expert rollout_routed_expert = ray.get(rollout_routed_expert_refs) + if isinstance(rollout_routed_expert, list): + assert len(rollout_routed_expert) == 1, ( + f"Expected one routed experts ref, got {len(rollout_routed_expert)}" + ) + rollout_routed_expert_refs = rollout_routed_expert[0] + rollout_routed_expert = ray.get(rollout_routed_expert_refs) # free obj store explicitly if self.sp_mesh is None or self.sp_mesh.size() == 1: ray.internal.free(rollout_routed_expert_refs, local_only=False) diff --git a/xtuner/v1/rl/utils/misc.py b/xtuner/v1/rl/utils/misc.py index 79b23290fa..ecd9d42119 100644 --- a/xtuner/v1/rl/utils/misc.py +++ b/xtuner/v1/rl/utils/misc.py @@ -15,7 +15,7 @@ import requests import torch.nn.functional as F -from xtuner.v1.data_proto.rl_data import RolloutState, Status +from xtuner.v1.data_proto.rl_data import TRACE_STORE_PROMPT_TEXT_KEY, USE_TRACE_STORE_KEY, RolloutState, Status from xtuner.v1.data_proto.utils import calculate_seq_staleness as calculate_seq_staleness from xtuner.v1.utils.logger import get_logger @@ -283,13 +283,26 @@ def chat_trace_records_to_rollout_states( states: list[RolloutState] = [] for index, record in enumerate(normalized_records): + combined_extra_fields = { + **deepcopy(rollout_state.extra_fields), + "gateway_trace_index": index, + "gateway_trace_count": trace_count, + "gateway_trace_records": deepcopy(trace_summary), + "gateway_request_id": record.get("request_id"), + "gateway_request_snapshot": record.get("request_snapshot"), + "gateway_response_snapshot": record.get("response_snapshot"), + **deepcopy(extra_fields or {}), + } + use_trace_store = bool(combined_extra_fields.get(USE_TRACE_STORE_KEY, True)) + prompt_ids = record.get("prompt_ids") response_ids = record.get("response_ids") - if not prompt_ids or not response_ids: + if not use_trace_store and (not prompt_ids or not response_ids): raise RuntimeError(f"Gateway trace record {index} is missing prompt_ids or response_ids.") + response_ids_list = response_ids if isinstance(response_ids, list) else [] logprobs = record.get("logprobs") - if not isinstance(logprobs, list) or len(logprobs) != len(response_ids): + if not isinstance(logprobs, list) or len(logprobs) != len(response_ids_list): logprobs = None status_value = record.get("status") @@ -310,34 +323,37 @@ def chat_trace_records_to_rollout_states( uid = None response = record.get("output_text") - if response is None and tokenizer is not None: + if response is None and tokenizer is not None and response_ids_list: try: - response = tokenizer.decode(response_ids) + response = tokenizer.decode(response_ids_list) except Exception: response = None normalized = rollout_state.model_copy(deep=True) normalized.uid = uid - normalized.prompt_ids = list(prompt_ids) - normalized.tokens = list(prompt_ids) - normalized.response_ids = list(response_ids) - normalized.response_mask = [1] * len(response_ids) - normalized.logprobs = logprobs + if record.get("session_uid") is not None: + normalized.session_uid = record.get("session_uid") + normalized.prompt_ids = list(prompt_ids or []) + normalized.tokens = list(prompt_ids or []) + if use_trace_store: + if isinstance(record.get("internal_messages"), list): + normalized.message = deepcopy(record["internal_messages"]) + normalized.response_ids = None + normalized.response_mask = None + normalized.logprobs = None + normalized.routed_experts = None + if record.get("trace_store_prompt_text") is not None: + combined_extra_fields[TRACE_STORE_PROMPT_TEXT_KEY] = record.get("trace_store_prompt_text") + else: + normalized.response_ids = list(response_ids_list) + normalized.response_mask = [1] * len(response_ids_list) + normalized.logprobs = logprobs normalized.response = response normalized.finish_reason = record.get("finish_reason") normalized.status = status normalized.error_msg = None if status == Status.COMPLETED else f"Gateway trace status={status.value}" normalized.reward = None - normalized.extra_fields = { - **deepcopy(rollout_state.extra_fields), - "gateway_trace_index": index, - "gateway_trace_count": trace_count, - "gateway_trace_records": deepcopy(trace_summary), - "gateway_request_id": record.get("request_id"), - "gateway_request_snapshot": record.get("request_snapshot"), - "gateway_response_snapshot": record.get("response_snapshot"), - **deepcopy(extra_fields or {}), - } + normalized.extra_fields = combined_extra_fields states.append(normalized) return states diff --git a/xtuner/v1/train/rl_trainer.py b/xtuner/v1/train/rl_trainer.py index 1fde1eb3aa..5bbba1172c 100644 --- a/xtuner/v1/train/rl_trainer.py +++ b/xtuner/v1/train/rl_trainer.py @@ -17,7 +17,11 @@ from transformers import AutoTokenizer, PreTrainedTokenizer, PreTrainedTokenizerFast from xtuner.v1._writer import get_writer -from xtuner.v1.data_proto.rl_data import RolloutState, Status +from xtuner.v1.data_proto.rl_data import ( + TRACE_STORE_PROMPT_TEXT_KEY, + RolloutState, + Status, +) from xtuner.v1.data_proto.sequence_context import SequenceContext from xtuner.v1.patch import patch_default_save_plan from xtuner.v1.rl.advantage import BaseAdvantageConfig, GRPOAdvantageConfig @@ -36,6 +40,7 @@ _snapshot_nested_objectrefs, ) from xtuner.v1.rl.rollout.controller import RolloutControllerProxy +from xtuner.v1.rl.rollout.trace_store import get_store from xtuner.v1.rl.rollout.worker import RolloutConfig from xtuner.v1.rl.trainer.controller import TrainingController from xtuner.v1.rl.trainer.worker import WorkerConfig, WorkerLogItem @@ -187,9 +192,8 @@ def is_valid_for_training(group_data_items: list[RolloutState], logger) -> bool: - 'aborted': These items represent rollouts that were stopped prematurely. Using such partial data could lead the model to learn undesirable behaviors (e.g., stopping generation too early). - - Empty response/response_ids: The model's generated response is the core - of the training data for RL algorithms like PPO. If the response is - missing, there is nothing to compute rewards on or to train the model with. + - Token-level training data lives in Trace Store. The trainer validates the + exported trace before building the training batch. """ is_abort = any(item.status == Status.ABORTED for item in group_data_items) is_filtered = any(item.status == Status.FILTERED for item in group_data_items) @@ -199,19 +203,6 @@ def is_valid_for_training(group_data_items: list[RolloutState], logger) -> bool: f"Invalid dataflow group found during training, rollout state filtered: {is_filtered}, failed: {is_failed}, aborted: {is_abort}." ) return False - for item in group_data_items: - response_valid = item.response is not None and len(item.response) > 0 - ids_valid = item.response_ids is not None and len(item.response_ids) > 0 - if not ids_valid: - # NOTE: `response_ids` is the critical field for token-in-token-out mode, so we ensure it's not empty. - logger.warning( - "Invalid dataflow item found during training: no response or response_ids and skip this item." - ) - return False - if not response_valid: - # NOTE: check valid response string for judger inputs - logger.warning("Invalid dataflow item found during training: empty response string and skip this item.") - return False return True @@ -872,6 +863,118 @@ def _load_debug_rollout_batch(self, train_step: int) -> list[list[RolloutState]] self.logger.info(f"Loaded debug rollout batch for step {train_step} from {debug_file}") return cast(list[list[RolloutState]], train_batch) + def _render_trace_store_prompt_text(self, data: RolloutState) -> str | None: + """Build the exact Trace Store lookup key for a rollout item. + + The preferred path is an explicit `trace_store_prompt_text` carried by + the producer. If it is absent, this reconstructs the final chat-template + text from the lightweight `RolloutState` fields. + """ + extra_fields = data.extra_fields or {} + prompt_text = extra_fields.get(TRACE_STORE_PROMPT_TEXT_KEY) + if prompt_text is not None: + if isinstance(prompt_text, str) and prompt_text: + return prompt_text + self.logger.error(f"Invalid trace store prompt text for session {data.session_uid}: {prompt_text!r}") + return None + + if data.response is None: + self.logger.error(f"Cannot render trace store prompt text for session {data.session_uid}: response is None.") + return None + + messages = list(data.message) + assistant_message: dict[str, Any] = {"role": "assistant", "content": data.response} + if data.tool_calls: + assistant_message["tool_calls"] = [ + tool_call.model_dump() if hasattr(tool_call, "model_dump") else tool_call + for tool_call in data.tool_calls + ] + messages.append(assistant_message) + try: + rendered = self.tokenizer.apply_chat_template( + messages, + tools=data.tools, + add_generation_prompt=False, + tokenize=False, + ) + except Exception as exc: + self.logger.error(f"Failed to render trace store prompt text for session {data.session_uid}: {exc}") + return None + if not isinstance(rendered, str) or not rendered: + self.logger.error(f"Rendered empty trace store prompt text for session {data.session_uid}.") + return None + return rendered.rstrip() + + def _mark_trace_store_train_abandoned(self, session_ids: list[Any]) -> None: + """Best-effort release for sessions exported from Trace Store but not consumed by training.""" + if not session_ids: + return + store = get_store() + for session_id in session_ids: + try: + ray.get(store.mark_train_abandoned.remote(session_id)) + except Exception as exc: + self.logger.error(f"Failed to mark trace store train abandoned for session {session_id}: {exc}") + + def _materialize_trace_store_training_data(self, data: RolloutState) -> bool: + """Export token-level Trace Store data back into one RolloutState. + + The trainer keeps using the existing `RolloutState` fields in + `_prepare_train_data`: `prompt_ids`, `response_ids`, `response_mask`, + `logprobs`, and `routed_experts`. Routed experts stay as Ray refs and + are dereferenced later inside `TrainingWorker`. + """ + prompt_text = self._render_trace_store_prompt_text(data) + if prompt_text is None: + return False + + store = get_store() + session_id = data.session_uid + exported = False + try: + trace = ray.get(store.export_training_trace.remote(session_id, prompt_text)) + exported = True + routed_experts_key = trace.get("routed_experts") + rollout_routed_experts = None + if routed_experts_key is not None: + rollout_routed_experts = store.get_objects.remote([routed_experts_key]) + + input_ids = list(trace["input_ids"]) + labels = list(trace["labels"]) + logprobs = list(trace["logprobs"]) if trace.get("logprobs") is not None else None + if len(input_ids) < 2 or len(input_ids) != len(labels) or ( + logprobs is not None and len(logprobs) != len(input_ids) + ): + raise ValueError( + f"Invalid trace store training trace lengths: " + f"input_ids={len(input_ids)}, labels={len(labels)}, " + f"logprobs={None if logprobs is None else len(logprobs)}" + ) + + response_start = next((idx for idx, label in enumerate(labels) if label != -100), None) + if response_start is None or response_start == 0: + raise ValueError("Trace store training trace has no trainable response tokens.") + + prompt_ids = input_ids[:response_start] + response_ids = input_ids[response_start:] + response_labels = labels[response_start:] + response_mask = [0 if label == -100 else 1 for label in response_labels] + if not response_ids or not any(response_mask): + raise ValueError("Trace store training trace has no valid response ids.") + + data.prompt_ids = prompt_ids + data.tokens = list(prompt_ids) + data.response_ids = response_ids + data.response_mask = response_mask + data.logprobs = logprobs[response_start:] if logprobs is not None else None + data.routed_experts = rollout_routed_experts + return True + except Exception as exc: + self.logger.error(f"Failed to materialize trace store training data for session {session_id}: {exc}") + if exported: + self._mark_trace_store_train_abandoned([session_id]) + return False + # TODO: simplify with Packer.pack_pad_dispatch() def _prepare_train_data( self, @@ -890,18 +993,29 @@ def _prepare_train_data( for j, group in enumerate(data_groups): if not is_valid_for_training(group, self.logger): - self.logger.error(f"Skip one data group {group} due to rollout failed or empty response.") + self.logger.error(f"Skip one data group {group} due to invalid rollout status.") + continue + + session_ids = [data.session_uid for data in group] + if any(session_id is None for session_id in session_ids): + self.logger.error(f"Skip one data group {group} due to missing trace store session id.") + continue + if len(set(session_ids)) != len(session_ids): + self.logger.error(f"Skip one data group {group} due to duplicated trace store session ids: {session_ids}.") + continue + + exported_trace_store_session_ids: list[Any] = [] + trace_store_failed = False + for data in group: + if not self._materialize_trace_store_training_data(data): + trace_store_failed = True + break + exported_trace_store_session_ids.append(data.session_uid) + if trace_store_failed: + self._mark_trace_store_train_abandoned(exported_trace_store_session_ids) + self.logger.error(f"Skip one data group {group} due to trace store materialization failure.") continue - is_vlm_model = "train_prompt_ids" in group[0].extra_fields - if is_vlm_model: - # TODO(hha): VLM, 不好的设计，后续要去掉 - prompt_ids = group[0].extra_fields["train_prompt_ids"] - else: - prompt_ids = group[0].prompt_ids - assert prompt_ids is not None and len(prompt_ids) > 0, ( - f"Prompt ids cannot be None or empty in data: {group[0]}" - ) rewards = [] for data in group: assert data.reward is not None and "score" in data.reward, ( @@ -915,36 +1029,25 @@ def _prepare_train_data( prompt_repeat_k = len(group) for i in range(prompt_repeat_k): - item = group[i].response - logprobs: list[float] | None = None - - response_ids: List[int] = [] - if group[i].response_ids is not None: - resp_ids_raw = group[i].response_ids - if isinstance(resp_ids_raw, torch.Tensor): - response_ids = resp_ids_raw.flatten().tolist() - else: - response_ids = cast(List[int], resp_ids_raw) + prompt_ids = group[i].prompt_ids + response_ids = group[i].response_ids + assert prompt_ids is not None and len(prompt_ids) > 0, ( + f"Prompt ids cannot be None or empty in data: {group[i]}" + ) + assert response_ids is not None and len(response_ids) > 0, ( + f"Response ids cannot be None or empty in data: {group[i]}" + ) - logprobs = group[i].logprobs - if logprobs is not None: - assert len(logprobs) == len(response_ids), ( - f"{len(logprobs)} vs {len(response_ids)}, data: {group[i]}" - ) - # 只有 response 部分有 logprobs, 需要前面追加 - logprobs = [0.0] * (len(prompt_ids) - 1) + logprobs # type: ignore[arg-type] - else: - assert item is not None, "response item cannot be None" - response_ids = self.tokenizer(item, return_tensors="pt")["input_ids"].flatten().tolist() + logprobs = group[i].logprobs + if logprobs is not None: + assert len(logprobs) == len(response_ids), ( + f"{len(logprobs)} vs {len(response_ids)}, data: {group[i]}" + ) + logprobs = [0.0] * (len(prompt_ids) - 1) + logprobs - # 返回的 routed_experts 不包括 eos 的值，实际上也不需要，需要减一 - # TODO: verl tool agent loop 是否需要？ input_ids = prompt_ids + response_ids[:-1] - - prompt_len_list.append(len(prompt_ids)) - response_len_list.append(len(response_ids)) - - # 根据 response_mask 计算 response_ids 对应的shifted_labels + prompt_len = len(prompt_ids) + response_len = len(response_ids) if not group[i].response_mask: response_mask = [1] * len(response_ids) response_labels = response_ids @@ -958,9 +1061,11 @@ def _prepare_train_data( for response_id, mask_id in zip(response_ids, response_mask) ] shifted_labels = [-100] * (len(prompt_ids) - 1) + response_labels + + prompt_len_list.append(prompt_len) + response_len_list.append(response_len) shifted_labels_t = torch.tensor(shifted_labels, dtype=torch.int64).unsqueeze(0) - # 根据 response_mask 计算新的 advantages advatnages_val = advantages[i].item() actual_advantages = [advatnages_val] * len(prompt_ids) + [ 0.0 if mask == 0 else advatnages_val for mask in response_mask @@ -982,7 +1087,7 @@ def _prepare_train_data( position_ids = group[i].position_ids multimodal_train_info = group[i].mm_info multi_info_cast = cast(dict | None, multimodal_train_info) - seq_ctx = get_train_seq_ctx(input_ids_t, position_ids, multi_info_cast, len(response_ids) - 1) # type: ignore[arg-type] + seq_ctx = get_train_seq_ctx(input_ids_t, position_ids, multi_info_cast, response_len - 1) # type: ignore[arg-type] data_dict = { "seq_ctx": seq_ctx,