[Fix] Preserve concurrent trace sessions in disaggregated RL - #2021
[Fix] Preserve concurrent trace sessions in disaggregated RL#2021matrix72c wants to merge 5 commits into
Conversation
|
@claude review |
|
Claude: Summary本 PR 将 disaggregated RL 的 trace session 清理从 ProduceBatchResult impact: 字段语义未改变；但 Main Flowchart after this PRflowchart TD
A[disagg _fit: get_batch] --> B[_train_one_batch]
B --> C{release_only_consumed_trace_sessions}
C -->|True| D[_release_trace_store train_batch<br/>release_sessions 已消费 session]
D --> E{need_sync}
E -->|Yes| F[pause_produce<br/>in-flight rollout 转为 ABORTED 入 replay buffer]
F --> G[_sync_weights_and_save]
G --> H{enable_evaluate}
H -->|Yes| I[_run_evaluation<br/>finally: _release_trace_store 无参数 → release_all]
I --> J[replay buffer 中 ABORTED/leftover session 被误释放<br/>routed_experts refs 失效]
H -->|No| K[continue_produce]
J --> K
style D fill:#cce5ff,stroke:#0366d6
style I fill:#f9c0c0,stroke:#d73a49
style J fill:#f9c0c0,stroke:#d73a49
核心原理实现与单测核心实现是三处所有权划分： 单测方面，三个新增/修改测试全部替换了项目内部 seam，导致本 PR 的核心新代码没有被真实执行： 抽象与信息隐藏评估
单测建议
其他 Issues
VerdictREQUEST_CHANGES |
![claude[bot]](https://avatars.githubusercontent.com/in/1236702?s=60&v=4){kind=small}
058c1e4 to
099b55a
Compare
| if keys: | ||
| self.logger.warning(f"Trace store keys not released after release_all: {keys}") | ||
|
|
||
| def _release_trace_sessions_after_train_batch(self, train_batch: list[list[RolloutState]]) -> None: |
There was a problem hiding this comment.
我理解，原来是只能清理所有，现在增加了train_batch粒度的清理。这里release_train_batch函数又调用了releae_all，这个是为什么呢
There was a problem hiding this comment.
基类/colocated trainer 没有后台 producer 并发持有 session，因此一个 train batch 结束后仍沿用原来的 release_all() 行为；RLDisaggregatedTrainer 会 override 这个方法，只调用 _release_trace_sessions(_trace_session_ids(train_batch))，保留 producer/replay buffer 持有的 session。
There was a problem hiding this comment.
看上去，训练batch后都会调用该函数，通过override走不同的路径的是吧
| ) | ||
| for item in group: | ||
| if item.session_id is not None and str(item.session_id) in released_session_ids: | ||
| # TraceStore.release_sessions() already freed these routed-expert refs. |
There was a problem hiding this comment.
这里不需要再额外将routed_experts置为None，discard_rollout_state函数中会把item的response都丢掉
There was a problem hiding this comment.
这里置 None 主要是为了避免重复显式 free。release_existing_sessions() 成功后，TraceStore.release_sessions() -> Trie.release() -> _free_ray_refs() 已经释放了该 session 的 routed_experts ObjectRef。随后 discard_rollout_state() 会先执行 free_rollout_state_refs()，递归扫描 item.routed_experts 并再次调用 free_object_refs()，之后才重置字段。
因此这里只对实际已由 TraceStore 释放的 session 将 routed_experts 置为 None，避免第二次 free；如果 session 在 TraceStore 中不存在，则保留该字段，让 discard_rollout_state() 负责释放。最终清空字段的职责仍然在 discard_rollout_state()。
| trie = self.sessions.pop(session_id) if key is None else self.sessions[session_id] | ||
| trie.release(key) | ||
|
|
||
| def release_sessions(self, session_ids: list[str]) -> list[str]: |
There was a problem hiding this comment.
release_sessions和release_all都是遍历list挨个清理，这里有可能合并成一个吗，根据参数决定哪些需要清理
There was a problem hiding this comment.
可以复用底层遍历实现。我倾向于保留 release_sessions(ids) 和 release_all() 两个显式接口，避免重新引入 None 表示全量清理的隐式语义，就是改成 release_all() 内部调用 release_sessions(list(self.sessions))，随后清理全局的 objects/updated_at。
| retried.""" | ||
|
|
||
| @staticmethod | ||
| def _reset_retryable_expired_group(group: list[RolloutState]) -> None: |
There was a problem hiding this comment.
这个函数不能完全代替 _cleanup_expired_group 的功能吧，为啥要把之前的函数删掉呢？用之前的_cleanup_expired_group 是不是就够了
There was a problem hiding this comment.
不是单独替代旧函数。retryable group 只在锁内同步 reset；terminal groups 在锁内完成 storage 删除和收集，出锁后合并成一次 trace-store RPC 再 discard。原 _cleanup_expired_group 如果原样保留并在循环中调用，会重新变成持锁逐 group await。put() 的单 group terminal 场景也复用批量 discard helper。
| if keys: | ||
| self.logger.warning(f"Trace store keys not released after release_all: {keys}") | ||
|
|
||
| def _release_trace_sessions_after_train_batch(self, train_batch: list[list[RolloutState]]) -> None: |
There was a problem hiding this comment.
这里封装成 hook 是为了让 _train_one_batch 保持一份实现，同时由 trainer 类型决定 session ownership。基类/colocated 沿用 release_all()；disaggregated override 后只释放 consumed batch。如果去掉这个 hook，就需要在基类判断子类类型或重新引入布尔参数。
|
总体看下来，trace store的释放散落在 producer, replaybuffer和trainer，每个调用者都需要提取、去重 这样的调用路径改动点会比较少，并且可以把每次调用的去重等操作集中在flush release中 @jayhenry @Harold-lkk 麻烦也review下这个实现 |
|
LGTM |
099b55a to
94e4d45
Compare
|
我重新梳理后先做了一版更小范围的收敛：新增统一的 release_and_discard_rollout_groups()，producer 和 replay buffer 不再分别提取、去重 session id、调用 release 以及处理 routed_experts 重复 free，这些细节现在都集中在同一个 helper 中；TraceStore actor 继续通过幂等的 release_sessions() 统一去重和释放，release_all() 也复用了同一底层实现。这里暂时没有引入 pending_release_session_ids + flush_release，主要是 deferred cleanup 会额外引入 eval、异常退出、最终 flush 以及跨 caller 顺序等生命周期边界，而当前释放 RPC 还不是已知性能瓶颈，所以这版保留即时释放，避免新增中间状态。Trainer/Eval 仍按各自实际消费的 batch 精确释放，retryable rollout 不受影响。分支已 rebase 到最新 main，并补跑了 terminal cleanup、锁外释放以及上游新增的 token-level/partial staleness 回归测试，不知道这样收敛设计行不行 @YanhuiDua |
| return Status.EXPIRED | ||
|
|
||
| @staticmethod | ||
| async def _discard_terminal_expired_groups(groups: list[list[RolloutState]]) -> None: |
There was a problem hiding this comment.
结构我觉得没问题了，能否将 terminal 改为 non-retryable ，这样比较符合 xtuner 中原本的定义
4 participants
Summary
This PR fixes the ownership and cleanup of agentic rollout trace sessions in disaggregated RL training.
Root Cause
RLDisaggregatedTrainerintentionally keeps its background producer running while the learner trains the current batch. During that overlap, theRolloutTraceStorecontains sessions with different owners:The inherited post-batch cleanup called
RolloutTraceStore.release_all(). That operation assumed a batch-synchronous lifecycle in which every live session belonged to the completed learner batch. The assumption is not valid for disaggregated training, so cleanup could delete producer-owned sessions and force-free their Ray-backedrouted_expertswhile future rollout states still referenced them.Ray preserves ordering for calls from one caller, but it does not provide a global ordering across the learner and producer callers. Consequently, the same race could surface in more than one form:
ObjectReffailures when the data is consumed or checkpointed.A short one-step asynchronous smoke test can miss the problem when cleanup happens before the next producer insertion, or when the prefetched batch is immediately aborted at shutdown and never consumed. The failure becomes reproducible when producer progress overlaps the learner's post-batch cleanup.
This is therefore an XTuner lifecycle bug in the combination of disaggregated training, agentic trace storage, and concurrent prefetch; it is not caused by a sandbox configuration.
Fix
Selective trainer cleanup
Add an atomic, idempotent
RolloutTraceStore.release_sessions(session_ids)actor method and use it after a disaggregated learner batch. The trainer extracts the sessions represented by the consumed batch and releases only those sessions. Colocated and evaluation cleanup retain their existing full-store behavior.The actor method ignores already-absent IDs and returns the IDs it actually released. Keeping the lookup and release in one actor RPC avoids a client-side list/intersection/release time-of-check/time-of-use window.
Terminal discard cleanup
Selective post-batch cleanup means a rollout that will never reach
train_batch()needs an explicit terminal cleanup path. This PR therefore releases trace sessions before discarding:Retryable expired groups retain their sessions because they may re-enter the training lifecycle. When trace cleanup has already freed the routed-expert objects, the corresponding local field is cleared before generic rollout-state disposal to avoid a second explicit free.
Impact
ProduceBatchResult: no status, reward, timing, rollout-state, or accounting semantics are changed.RoutedExperts: references for consumed or terminally discarded sessions are freed exactly once; references owned by concurrent or retryable rollouts remain valid until their owning lifecycle ends.Reproduction
The issue was observed in a three-node disaggregated run with 8 learner GPUs and 16 rollout GPUs:
The trainer regression test models this deterministically with one consumed session and one concurrently produced session, then verifies that only the consumed session is released.
Test Plan
RLDisaggregatedTrainer;5 passed.git diff --check: passed.The full repository matrix is left to upstream CI.
Out of Scope
This PR does not change replay-checkpoint serialization. The stale replay
ObjectRefsymptom can be caused by the same premature trace release, but serialization policy and timeout diagnostics are separate concerns.