[Feature] Retain Agent tasks across complete-request weight updates

xtuner/v1/data_proto/rl_data.py

@@ -250,7 +250,7 @@ def reset_response(self) -> None:
             None
         """
         routed_experts = self.routed_experts
-        if routed_experts is not None:
+        if routed_experts is not None and not self.extra_fields.get("request_ids"):
             from ray import ObjectRef
 
             from xtuner.v1.rl.utils.ray_utils import free_object_refs
@@ -342,6 +342,10 @@ def free_rollout_state_refs(rollout_state: RolloutState) -> None:
 
     from xtuner.v1.rl.utils.ray_utils import free_object_refs
 
+    # Request-aware trace routes are shared with sibling segments and workers.
+    # Keep the ordinary rollout's explicit free contract when provenance is off.
+    if rollout_state.extra_fields.get("request_ids"):
+        rollout_state.routed_experts = None
     refs: list[ObjectRef] = []
 
     def clear_object_refs(value: Any) -> Any:

xtuner/v1/rl/agent_loop/sandbox_agent_loop/agent_in_sandbox_loop.py

@@ -195,6 +195,8 @@ class AgentInSandboxLoopConfig(AgentLoopConfig):
     max_concurrent_samples: int | None = None
     mode: Literal["train", "eval"] = "train"
     requires_rollout_proxy: bool = True
+    require_turn_provenance: bool = False
+    retire_route_on_completion: bool = False
 
     def build_local(
         self, rollout_controller: RolloutController | None = None, judger: Judger | None = None, logger=None
@@ -207,6 +209,8 @@ def build_local(
             logger=logger,
             max_concurrent_samples=self.max_concurrent_samples,
             mode=self.mode,
+            require_turn_provenance=self.require_turn_provenance,
+            retire_route_on_completion=self.retire_route_on_completion,
         )
 
 
@@ -220,13 +224,17 @@ def __init__(
         logger=None,
         max_concurrent_samples: int | None = None,
         mode: Literal["train", "eval"] = "train",
+        require_turn_provenance: bool = False,
+        retire_route_on_completion: bool = False,
     ):
         if hf_checkpoint is None:
             raise ValueError("hf_checkpoint must be provided for AgentInSandboxLoop.")
         super().__init__(rollout_ctl, sample_params, hf_checkpoint, judger, logger)
         self.max_concurrent_samples = max_concurrent_samples
         self._sample_semaphore = asyncio.Semaphore(max_concurrent_samples) if max_concurrent_samples else None
         self.mode = mode
+        self.require_turn_provenance = require_turn_provenance
+        self.retire_route_on_completion = retire_route_on_completion
 
     async def generate_group(self, rollout_state: list[RolloutState], **kwargs) -> list[RolloutState]:
         async def generate_one(state: RolloutState) -> list[RolloutState]:
@@ -283,6 +291,12 @@ async def generate_sample(  # type: ignore[override]
             rollout_state.error_msg = f"{type(exc).__name__}: {exc}"
             self.logger.error(f"[AgentInSandboxLoop] failed: {exc}\n{traceback.format_exc()}")
             return [rollout_state]
+        finally:
+            # This loop returns terminal Agent executions, never partial continuations.
+            # Retire HTTP ownership after runner completion/cancellation, independently
+            # of later filtering or replay consumption; keep trace ownership unchanged.
+            if self.retire_route_on_completion and self.rollout_ctl is not None and rollout_state.session_id is not None:
+                await self.rollout_ctl.retire_route_sessions.remote([str(rollout_state.session_id)])
 
     async def _run_item(self, item: AgentRolloutItem) -> AgentRolloutItem:
         runner = _resolve_runner(item.pipeline, str(item.uid))
@@ -365,6 +379,31 @@ async def _build_rollout_states(self, rollout_state: RolloutState, item: AgentRo
             # indistinguishable from a dense trace; catching that would need the MoE-vs-dense signal the loop lacks
             # (session_server gates on ``enable_return_routed_experts``).
             segment_state.routed_experts = data["routed_experts"]
+            if self.require_turn_provenance:
+                # main's response_ids is the full suffix, including tool/history
+                # positions. Versions must follow the same positions, not only
+                # the supervised labels. Preserve each original request's lease.
+                versions = data.get("token_model_steps", [])[prompt_len:]
+                if len(versions) != len(segment_state.response_ids) or any(v is None for v in versions):
+                    raise ValueError("Missing request-time response model versions")
+                if not data["routed_experts"]:
+                    raise ValueError("turn-resume MoE training requires routed experts on every node")
+                records = _load_messages_artifact(item.artifacts)
+                matching = [
+                    record for record in records if record.get("messages") == messages and record.get("tools") == tools
+                ]
+                recorded_ids = matching[0].get("request_ids", []) if len(matching) == 1 else []
+                path_ids = data.get("request_ids", [])
+                if (
+                    not recorded_ids
+                    or not path_ids
+                    or recorded_ids[-1] != path_ids[-1]
+                    or not set(path_ids).issubset(recorded_ids)
+                    or not await trace_store.contains_request_ids.remote(str(rollout_state.session_id), recorded_ids)
+                ):
+                    raise ValueError("Agent transcript request IDs do not match committed trace")
+                segment_state.response_model_steps = versions
+                segment_state.extra_fields["request_ids"] = data["request_ids"]
             if segment_state.response_ids:
                 segment_state.response = self.tokenizer.decode(segment_state.response_ids)
             else:

xtuner/v1/rl/agent_loop_manager/agent_loop_manager.py

@@ -127,6 +127,10 @@ class AgentLoopManagerConfig(BaseModel):
     model_config = ConfigDict(extra="forbid", arbitrary_types_allowed=True)
 
     tasks: list[TaskSpecConfig] | TaskSpecConfig
+    turn_keepalive: bool = False
+    turn_keepalive_freeze_timeout: float = Field(default=30, gt=0)
+    turn_keepalive_drain_timeout: float = Field(default=450, gt=0)
+    turn_keepalive_quiesce_timeout: float = Field(default=3600, gt=0)
 
     def build(
         self,
@@ -137,6 +141,15 @@ def build(
         sync_weights_interval: int = 1,
         rollout_teacher_scorer_config: RolloutTeacherScorerConfig | None = None,
     ) -> AgentLoopManager:
+        if self.turn_keepalive:
+            from .colocate_keepalive import build_keepalive_manager
+
+            return build_keepalive_manager(
+                self, rollout_controller=rollout_controller, tokenizer=tokenizer,
+                replay_buffer=replay_buffer, logger=logger,
+                sync_weights_interval=sync_weights_interval,
+                rollout_teacher_scorer_config=rollout_teacher_scorer_config,
+            )
         tasks = self.tasks if isinstance(self.tasks, list) else [self.tasks]
         if not tasks:
             raise ValueError("AgentLoopManagerConfig requires at least one task config.")
@@ -221,6 +234,7 @@ async def produce_batch(
         train_step: int,
         *,
         model_step: int,
+        quiesce: bool = False,
     ) -> ProduceBatchResult:
         # 共卡同步入口：生产入 buffer -> pause/drain 本轮 pending -> 取非空训练 batch。
         if batch_size <= 0:
@@ -233,7 +247,12 @@ async def produce_batch(
         active_tasks = [task for task in self.task_runners if current_sizes[task.task_name] > 0]
         assert active_tasks, "No active tasks found"
 
-        await self._rollout_controller.continue_generation.remote()  # type: ignore[attr-defined]
+        keepalive = getattr(self, "_keepalive", None)
+        if keepalive is None:
+            await self._rollout_controller.continue_generation.remote()
+        else:
+            await keepalive.resume(self._rollout_controller, model_step)
+            await keepalive.check_requests(self._rollout_controller)
         local_progress = ProduceProgress.build(
             task_names=self.task_names,
             target_samples=current_sizes,
@@ -271,12 +290,37 @@ async def produce_batch(
         await asyncio.gather(*produce_futures)
         local_progress.add_produce_time(time.perf_counter() - produce_start)
 
-        # pause 只收尾本轮本地 pending。
-        await self._rollout_controller.pause_generation.remote()  # type: ignore[attr-defined]
+        # Checkpoint/final quiesce happens while gates are open. Never wait for
+        # an entire episode behind a sealed request gate.
+        if keepalive is not None:
+            if quiesce:
+                async def finish_live_tasks():
+                    for task in self.task_runners:
+                        ctx = ProduceContext(
+                            agent_loop=task.agent_loop, sampler=task.sampler,
+                            replay_buffer=self.replay_buffer, task_batch_size=0,
+                            task_name=task.task_name, train_step=train_step,
+                            model_step=model_step, progress=local_progress,
+                            stale_threshold=task.stale_threshold,
+                            token_stale_threshold=task.token_stale_threshold,
+                            expired_groups_retryable=task.expired_groups_retryable,
+                        )
+                        owner = task.produce_strategy._pending_tasks
+                        while owner.count():
+                            for ready in await owner.wait_and_claim(timeout_s=1):
+                                await ctx.put_generated_group(ready.result())
+                await asyncio.wait_for(finish_live_tasks(), keepalive.quiesce_timeout)
+            await keepalive.pause(self._rollout_controller)
+        else:
+            await self._rollout_controller.pause_generation.remote()
 
         pause_time_s = 0.0
         for task in active_tasks:
             produce_strategy = cast(ProduceStrategy, task.produce_strategy)
+            if keepalive is not None:
+                # Live futures stay in the strategy's existing _PendingTasks.
+                # Done futures are harvested on the next window (or quiesce).
+                continue
             pause_time_s += await produce_strategy.pause_produce(
                 ProduceContext(
                     agent_loop=task.agent_loop,
@@ -357,6 +401,8 @@ async def resume(self, checkpoint_path: Path | str) -> int:
         state_path = manager_state_path(checkpoint_path)
         with state_path.open("r") as f:
             manager_state = json.load(f)
+        if getattr(self, "_keepalive", None) is not None and not manager_state.get("replay_buffer_saved", True):
+            raise RuntimeError("Full colocate resume requires a persisted replay buffer")
         if manager_state.get("replay_buffer_saved", True):
             # replay buffer 恢复是 async I/O，不能在已有 event loop 中再次嵌套 asyncio_run。
             await self.replay_buffer.resume(checkpoint_path)

xtuner/v1/rl/agent_loop_manager/agent_turn_pause_policy.py

@@ -0,0 +1,50 @@
+"""Disaggregated pause policy; PendingTasks remains the sole live-task owner."""
+
+from .produce_utils import pause_pending_tasks, retain_pending_tasks
+
+
+class DrainPendingPausePolicy:
+    pending_cleanup = staticmethod(pause_pending_tasks)
+
+    async def check_requests(self, controller):
+        return
+
+    async def before_sync(self, manager, checkpoint_due):
+        return
+
+    async def pause(self, controller):
+        await controller.pause_generation.remote()
+
+    async def resume(self, controller, model_step):
+        await controller.continue_generation.remote()
+
+
+class AgentTurnPausePolicy:
+    pending_cleanup = staticmethod(retain_pending_tasks)
+
+    def __init__(self, *, freeze_timeout: float, drain_timeout: float, quiesce_timeout: float = 1800):
+        self.quiesce_timeout = quiesce_timeout
+        self.freeze_timeout = freeze_timeout
+        self.drain_timeout = drain_timeout
+        self.barrier = None
+
+    async def check_requests(self, controller):
+        await controller.assert_request_gates_healthy.remote()
+
+    async def before_sync(self, manager, checkpoint_due):
+        if checkpoint_due:
+            await manager.quiesce(self.quiesce_timeout)
+
+    async def pause(self, controller):
+        self.barrier = await controller.prepare_pause.remote(
+            freeze_timeout=self.freeze_timeout, drain_timeout=self.drain_timeout
+        )
+        await controller.pause_generation.remote()
+
+    async def resume(self, controller, model_step):
+        # Initial construction has no admitted requests and no barrier yet.
+        if self.barrier is None:
+            await controller.continue_generation.remote()
+        else:
+            await controller.resume_requests.remote(self.barrier, model_step, timeout=self.freeze_timeout)
+            self.barrier = None

xtuner/v1/rl/agent_loop_manager/colocate_keepalive.py

@@ -0,0 +1,127 @@
+"""Opt-in persistent asyncio owner for alternating colocated RL.
+
+The native asyncio_run loop is stable but stops during synchronous training.
+This adapter keeps tool/client tasks progressing without enabling GPU inference.
+Only the strategy's existing _PendingTasks owns live generation tasks.
+"""
+
+import asyncio
+import threading
+
+
+class LoopThread:
+    def __init__(self):
+        self.loop = asyncio.new_event_loop()
+        self._ready = threading.Event()
+        self.thread = threading.Thread(target=self._run, name="xt-colocate-agent-loop", daemon=True)
+        self.thread.start()
+        self._ready.wait()
+
+    def _run(self):
+        asyncio.set_event_loop(self.loop)
+        self._ready.set()
+        self.loop.run_forever()
+        self.loop.close()
+
+    def submit(self, coro):
+        if not self.thread.is_alive():
+            coro.close()
+            raise RuntimeError("Colocate Agent loop is closed")
+        return asyncio.run_coroutine_threadsafe(coro, self.loop)
+
+    async def call(self, coro):
+        return await asyncio.wrap_future(self.submit(coro))
+
+    def stop(self):
+        self.loop.call_soon_threadsafe(self.loop.stop)
+        self.thread.join(timeout=30)
+        if self.thread.is_alive():
+            raise TimeoutError("Colocate Agent loop thread did not stop")
+
+
+class ColocateKeepaliveManager:
+    def __init__(self, inner, owner):
+        self._inner, self._owner = inner, owner
+        self._closed = False
+
+    def __getattr__(self, name):
+        return getattr(self._inner, name)
+
+    async def produce_batch(self, *args, **kwargs):
+        return await self._owner.call(self._inner.produce_batch(*args, **kwargs))
+
+    async def prepare_initial_sync(self):
+        # Servers were born closed, before their HTTP endpoint was published.
+        # Build the ordinary fixed-topology barrier; no group is scheduled.
+        return await self._owner.call(self._inner._keepalive.pause(self._inner._rollout_controller))
+
+    async def publish_initial_version(self, model_step):
+        return await self._owner.call(self._inner._keepalive.resume(self._inner._rollout_controller, model_step))
+
+    async def save(self, *args, **kwargs):
+        return await self._owner.call(self._inner.save(*args, **kwargs))
+
+    async def resume(self, *args, **kwargs):
+        return await self._owner.call(self._inner.resume(*args, **kwargs))
+
+    async def close(self):
+        if self._closed:
+            return
+        self._closed = True
+
+        async def cleanup():
+            cancellations = [
+                asyncio.create_task(task.produce_strategy._pending_tasks.cancel_all())
+                for task in self._inner.task_runners
+            ]
+            if cancellations:
+                _, pending = await asyncio.wait(cancellations, timeout=120)
+                if pending:
+                    raise TimeoutError("Pending Agent cleanup exceeded 120 seconds")
+            # Local AgentInSandboxLoop cancellation reaches original Runner
+            # finally/pool cleanup. No duplicate sandbox owner is introduced.
+            remaining = [t for t in asyncio.all_tasks() if t is not asyncio.current_task()]
+            for task in remaining:
+                task.cancel()
+            if remaining:
+                _, pending = await asyncio.wait(remaining, timeout=30)
+                if pending:
+                    raise TimeoutError("Agent client teardown exceeded 30 seconds")
+
+        try:
+            await self._owner.call(cleanup())
+        finally:
+            self._owner.stop()
+
+
+def build_keepalive_manager(config, **kwargs):
+    from .agent_turn_pause_policy import AgentTurnPausePolicy
+    from .producer import AsyncProduceStrategyConfig
+
+    tasks = config.tasks if isinstance(config.tasks, list) else [config.tasks]
+    for task in tasks:
+        strategy = task.produce_strategy_config
+        if not isinstance(strategy, AsyncProduceStrategyConfig) or strategy.enable_partial_rollout:
+            raise ValueError("Colocate turn keepalive requires AsyncProduceStrategy and token partial disabled")
+        if not getattr(task.agent_loop_config, "require_turn_provenance", False):
+            raise ValueError("Colocate turn keepalive requires real request provenance")
+    owner = LoopThread()
+
+    async def build():
+        # Build clients, locks, sampler and tasks on the one persistent loop.
+        inner = config.model_copy(update={"turn_keepalive": False}).build(**kwargs)
+        inner._keepalive = AgentTurnPausePolicy(
+            freeze_timeout=config.turn_keepalive_freeze_timeout,
+            drain_timeout=config.turn_keepalive_drain_timeout,
+            quiesce_timeout=config.turn_keepalive_quiesce_timeout,
+        )
+        for task in inner.task_runners:
+            task.produce_strategy.retain_live_tasks = True
+        return inner
+
+    try:
+        inner = owner.submit(build()).result()
+    except BaseException:
+        owner.stop()
+        raise
+    return ColocateKeepaliveManager(inner, owner)