[Fix] Preserve concurrent trace sessions in disaggregated RL

[matrix72c]

Summary

This PR fixes the ownership and cleanup of agentic rollout trace sessions in disaggregated RL training.

• Release only the trace sessions consumed by the current learner batch.
• Preserve sessions created concurrently by the background producer and sessions that remain retryable in the replay buffer.
• Release terminal failed/filtered rollouts and non-retryable expired rollouts that will never reach the learner.
• Add regression tests for trainer, producer, and replay-buffer cleanup paths.

Root Cause

RLDisaggregatedTrainer intentionally keeps its background producer running while the learner trains the current batch. During that overlap, the RolloutTraceStore contains sessions with different owners:

1. sessions belonging to the batch already consumed by the learner; and
2. sessions that the producer has concurrently created for future replay-buffer batches.

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-backed routed_experts while 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:

• the learner clears the store, the producer inserts a new session, and the learner's subsequent empty-store assertion observes the new session; or
• a future rollout/replay entry keeps references to data cleared by the learner, later causing trace prompt mismatches or stale ObjectRef failures 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:

• failed or filtered producer groups; and
• expired replay-buffer groups that are marked non-retryable.

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.
• Ray concurrency: no actor concurrency groups or scheduling policy are changed. Cleanup is one lightweight actor RPC over the session IDs represented by the batch/discarded group.
• Performance: no model forward, backward, or weight-sync hot path is changed. The added work is linear in the small number of session IDs being released, so a separate performance benchmark is not expected to be meaningful.

Reproduction

The issue was observed in a three-node disaggregated run with 8 learner GPUs and 16 rollout GPUs:

1. the learner starts training a completed rollout batch;
2. the background producer continues opening trace sessions for the next batch;
3. the learner completes backward and performs post-batch trace cleanup;
4. global cleanup removes sessions that were not part of the learner batch, leading to a non-empty-store assertion, prompt mismatch, or a stale Ray object later in replay/checkpoint handling depending on cross-caller timing.

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

• Targeted regression tests for:
  • consumed-only cleanup in RLDisaggregatedTrainer;
  • failed/filtered producer cleanup;
  • terminal expired replay cleanup; and
  • preservation of retryable expired sessions.
• Result: 5 passed.
• Ruff checks on all changed files: passed.
• Python byte-compilation of the changed production modules: passed.
• git diff --check: passed.
• Three-node, 24-GPU disaggregated smoke:
  • verified that the consumed batch was released while five concurrent sessions remained live;
  • completed learner forward/backward, producer pause, DCP/HF checkpoint, 16-rank weight synchronization, Ray shutdown, and all three replicas successfully.

The full repository matrix is left to upstream CI.

Out of Scope

This PR does not change replay-checkpoint serialization. The stale replay ObjectRef symptom can be caused by the same premature trace release, but serialization policy and timeout diagnostics are separate concerns.

[YanhuiDua]

@claude review

[github-actions]

Claude:

Summary

本 PR 将 disaggregated RL 的 trace session 清理从 release_all() 改为按已消费 batch 选择性释放，并为 FAILED/FILTERED 生产组和不可重试的 EXPIRED replay 组补上终态释放路径；方向正确，但 disagg 的 eval 清理路径仍保留全量释放，会重新引入本 PR 要修的问题。

ProduceBatchResult impact: 字段语义未改变；但 refresh_staleness 持锁远程调用可能拉长 produce_time_s。
RoutedExperts impact: 已消费与终态丢弃的 refs 仅释放一次；但 disagg eval 路径的 release_all() 仍可能提前释放 retryable/leftover replay 条目持有的 refs。
Ray concurrency impact: not affected（未改动 concurrency group 或调度策略）。

Main Flowchart after this PR

flowchart 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

Loading

核心原理实现与单测

核心实现是三处所有权划分：RolloutTraceStore.release_sessions() 在单次 actor RPC 内完成“查存在 + 释放 + 返回实际释放集合”，消除了客户端 list/intersect/release 的 TOCTOU 窗口；trainer 侧按 batch 内 session_id 集合选择性释放；producer 与 replay buffer 在丢弃终态样本前先释放 trace，并对已释放 session 清空 routed_experts 以避免二次 free。这几处对 routed_experts 单次释放的推理是正确的：session id 与 rollout 一一对应，同一 sandbox session 拆出的多个 segment 同属一个 group，选择性释放不会误伤兄弟 group。

单测方面，三个新增/修改测试全部替换了项目内部 seam，导致本 PR 的核心新代码没有被真实执行：RolloutTraceStore.release_sessions 的去重、跳过缺失 id、返回值与 _free_ray_refs 行为，以及 release_existing_sessions 的 ray.is_initialized() / store is None 兜底均无覆盖。

抽象与信息隐藏评估

• Warning xtuner/v1/train/rl_trainer.py:885 — 同一条“batch 后谁拥有 session”的策略被 release_only_consumed_trace_sessions 标志与 train_batch=None 隐式约定编码两次，disagg 专有策略泄漏进基类 Interface 并散落到三个调用点。

单测建议

• Warning xtuner/v1/rl/rollout/trace_store.py:338 — 新增的 release_sessions / release_existing_sessions 全程被 mock 掉，核心行为缺少真实代码路径覆盖（tests/rl/test_replay_buffer.py 已有本地 Ray，可做真实 actor 级测试）。

其他 Issues

• Warning xtuner/v1/train/rl_trainer.py:967-973 — disagg 的 _run_evaluation 仍在 finally 中执行全量 release_all()，会释放 pause_produce 刚放回 replay buffer 的 retryable ABORTED 与 leftover session，本 PR 的修复在该路径上不完整。
• Warning xtuner/v1/rl/replay_buffer.py:562 — refresh_staleness 在持有 self._lock 期间按组 await trace store RPC，把 producer/trainer 的全部 buffer 操作串行化在远程调用之后（put 在锁外调用，属不对称实现）。

Verdict

REQUEST_CHANGES