support checkpoint engine[P2P] recovery in colocate mode

tests/rl/test_rl_colocate_trainer.py

@@ -19,6 +19,7 @@
 """
 
 import asyncio
+import threading
 import tempfile
 import unittest
 from pathlib import Path
@@ -112,6 +113,11 @@ def tearDown(self):
 
     def _make_trainer(self, agent_loop_manager, *, total_train_steps: int = 1, sync_weights_interval: int = 1):
         trainer = RLColocateTrainer.__new__(RLColocateTrainer)
+        trainer._rollout_resources_available = threading.Event()
+        trainer._rollout_weight_update_lock = threading.Lock()
+        trainer._pending_rollout_weight_update_stop_event = threading.Event()
+        trainer._pending_rollout_weight_update_thread: threading.Thread | None = None
+        trainer._rollout_config = SimpleNamespace(weight_transport_type='ipc')
         trainer.logger = MagicMock()
         trainer._total_train_steps = total_train_steps
         trainer._cur_step = 0

tests/rl/test_rollout_logic.py

@@ -651,6 +651,29 @@ def test_registry_filters_entrypoints_and_tracks_lifecycle(self):
         registry.set_group_recovery_result(claimed_groups[0], recovered=False)
         self.assertEqual(self._worker_by_rank(registry, 0).lifecycle_state, WorkerLifecycleState.INACTIVE)
 
+    def test_registry_sets_groups_state_with_source_filter(self):
+        runtime_layout = self._runtime_layout(engine_ranks=(0,))
+        registry = RolloutWorkerRegistry(rollout_topology=runtime_layout)
+        _register_started_servers(
+            registry,
+            ((0, object(), "http://worker-0", "http://session-0"),),
+            lifecycle_state=WorkerLifecycleState.PENDING_WEIGHTS,
+        )
+
+        pending_group = registry.get_target_state_worker_groups(WorkerLifecycleState.PENDING_WEIGHTS)[0]
+        updated_groups = registry.set_groups_state(
+            groups=[pending_group],
+            target_state=WorkerLifecycleState.ACTIVE,
+            source_state=WorkerLifecycleState.PENDING_WEIGHTS,
+        )
+
+        self.assertEqual(updated_groups[0].ranks, (0,))
+        self.assertEqual(registry.get_target_state_worker_groups(WorkerLifecycleState.PENDING_WEIGHTS), ())
+        self.assertEqual(
+            tuple(worker.rank for worker in registry.get_target_state_workers(WorkerLifecycleState.ACTIVE)),
+            (0,),
+        )
+
     def test_registry_projects_weight_update_targets_from_topology_and_runtime_state(self):
         runtime_layout = self._runtime_layout(engine_ranks=(0, 1))
         registry = RolloutWorkerRegistry(rollout_topology=runtime_layout)
@@ -668,7 +691,6 @@ def test_registry_projects_weight_update_targets_from_topology_and_runtime_state
         self.assertEqual(target.engine_size, 2)
         self.assertEqual(target.server_url, "http://worker-0")
         self.assertEqual(target.lifecycle_state, WorkerLifecycleState.ACTIVE.value)
-        self.assertTrue(target.is_active)
 
 
 class TestSessionRouter(unittest.IsolatedAsyncioTestCase):
@@ -1228,7 +1250,10 @@ def test_run_once_does_not_log_error_when_last_active_worker_becomes_inactive(se
         with patch("xtuner.v1.rl.rollout.health_manager.logger.error") as log_error:
             manager.run_once()
 
-        log_error.assert_not_called()
+        self.assertFalse(
+            any("No active rollout worker" in call.args[0] for call in log_error.call_args_list),
+            f"Expected no stale no-active-worker log, got: {log_error.call_args_list}",
+        )
         self.assertFalse(self._worker_by_rank(registry, 0).is_active())
         self.assertEqual(actor.check_health.calls, [()])
 
@@ -1325,7 +1350,7 @@ def test_restart_barrier_keeps_failed_recovery_group_inactive(self):
             f"Expected restart failure log to explain why it is non-fatal, got: {log_error.call_args_list}",
         )
 
-    def test_restart_barrier_notifies_recovered_group_after_success(self):
+    def test_restart_barrier_marks_recovered_group_pending_weights_after_success(self):
         actor = SimpleNamespace(check_health=_FakeAsyncRemoteMethod(True))
         worker_info = WorkerSnapshot(
             rank=0,
@@ -1334,10 +1359,10 @@ def test_restart_barrier_notifies_recovered_group_after_success(self):
             session_url="http://session-0",
             lifecycle_state=WorkerLifecycleState.INACTIVE,
         )
-        recovered_groups = []
+        pending_weights_groups = []
         listener = SimpleNamespace(
             on_worker_group_inactive=MagicMock(),
-            on_worker_group_recovered=recovered_groups.append,
+            on_worker_group_pending_weights=pending_weights_groups.append,
         )
         manager, registry = self._build_manager(
             {0: worker_info},
@@ -1347,9 +1372,14 @@ def test_restart_barrier_notifies_recovered_group_after_success(self):
         with patch.object(manager, "_restart_worker_group", return_value=True):
             manager.restart_inactive_workers()
 
-        self.assertTrue(self._worker_by_rank(registry, 0).is_active())
-        self.assertEqual([group.ranks for group in recovered_groups], [(0,)])
-        self.assertTrue(all(worker.is_active() for worker in recovered_groups[0].workers))
+        self.assertEqual(self._worker_by_rank(registry, 0).lifecycle_state, WorkerLifecycleState.PENDING_WEIGHTS)
+        self.assertEqual([group.ranks for group in pending_weights_groups], [(0,)])
+        self.assertTrue(
+            all(
+                worker.lifecycle_state is WorkerLifecycleState.PENDING_WEIGHTS
+                for worker in pending_weights_groups[0].workers
+            )
+        )
 
     def test_restart_barrier_cleans_claimed_groups_when_stopping(self):
         actor = SimpleNamespace(check_health=_FakeAsyncRemoteMethod(True))
@@ -1442,7 +1472,7 @@ def fake_ray_get(refs, timeout=None):
         self.assertEqual(actor.offload.calls, [()])
         self.assertEqual(actor.restore_skip_load_weights.calls, [()])
 
-    def test_recovered_listener_runs_outside_lifecycle_operation_lock(self):
+    def test_pending_weights_listener_runs_outside_lifecycle_operation_lock(self):
         actor = SimpleNamespace(check_health=_FakeAsyncRemoteMethod(True))
         worker_info = WorkerSnapshot(
             rank=0,
@@ -1453,7 +1483,7 @@ def test_recovered_listener_runs_outside_lifecycle_operation_lock(self):
         lock_acquired_by_listener = []
         manager, _ = self._build_manager({0: worker_info})
 
-        def on_worker_group_recovered(group):
+        def on_worker_group_pending_weights(group):
             acquired = manager._lifecycle_operation_lock.acquire(blocking=False)
             lock_acquired_by_listener.append(acquired)
             if acquired:
@@ -1462,7 +1492,7 @@ def on_worker_group_recovered(group):
         manager._worker_lifecycle_listeners = (
             SimpleNamespace(
                 on_worker_group_inactive=MagicMock(),
-                on_worker_group_recovered=on_worker_group_recovered,
+                on_worker_group_pending_weights=on_worker_group_pending_weights,
             ),
         )
 

tests/rl/test_update_weight_colocate.py

@@ -29,6 +29,7 @@
     clear_cpu_resource_manager,
     set_cpu_resource_manager,
 )
+from xtuner.v1.rl.rollout.worker_registry import WorkerLifecycleState
 
 MODEL_PATH = os.environ["QWEN3_5_MOE_PATH"]
 
@@ -164,7 +165,7 @@ def _setup_engines(self, *, weight_transport_type: str):
 
     def _check_sglang_weights(self, rollout_controller, action):
         targets = ray.get(rollout_controller.get_weight_update_targets.remote())
-        active_urls = [target.server_url for target in targets if target.is_active]
+        active_urls = [target.server_url for target in targets if target.lifecycle_state == WorkerLifecycleState.ACTIVE.value]
         self.assertGreater(len(active_urls), 0)
         results = []
         for url in active_urls:

tests/rl/test_update_weight_disaggregated.py

@@ -22,6 +22,7 @@
     clear_cpu_resource_manager,
     set_cpu_resource_manager,
 )
+from xtuner.v1.rl.rollout.worker_registry import WorkerLifecycleState
 
 TEST_TEXT_MESSAGES = [{"role": "user", "content": "Hello!"}]
 MODEL_PATH = os.environ["QWEN3_VL_DENSE_PATH"]
@@ -120,7 +121,7 @@ def init_config(self):
 
     def _check_sglang_weights(self, rollout_controller, action):
         targets = ray.get(rollout_controller.get_weight_update_targets.remote())
-        active_urls = [target.server_url for target in targets if target.is_active]
+        active_urls = [target.server_url for target in targets if target.lifecycle_state == WorkerLifecycleState.ACTIVE.value]
         self.assertGreater(len(active_urls), 0)
         results = []
         for url in active_urls:

xtuner/v1/rl/agent_loop_manager/producer.py

@@ -216,6 +216,9 @@ class AsyncProduceStrategyConfig(ProduceStrategyConfig):
             rerolls out immediately without entering tail-batch mode, and
             ``N > 0`` waits until the expired pool contains at least ``N``
             groups before entering tail-batch mode.
+        max_pending_tasks (int | None): Maximum number of concurrently pending
+            rollout groups in one produce_batch call. Defaults to None, which
+            keeps the existing unbounded scheduling behavior.
 
     **Examples:**
 
@@ -233,6 +236,7 @@ class AsyncProduceStrategyConfig(ProduceStrategyConfig):
     max_staleness: int = Field(default=0, ge=0)
     max_token_staleness: int | None = Field(default=None, ge=0)
     tail_batch_trigger_size: int = Field(default=-1, ge=-1)
+    max_pending_tasks: int | None = Field(default=None, gt=0)
 
     def build(
         self,
@@ -261,6 +265,7 @@ def build(
             max_token_staleness=self.max_token_staleness,
             sync_weights_interval=sync_weights_interval,
             tail_batch_trigger_size=self.tail_batch_trigger_size,
+            max_pending_tasks=self.max_pending_tasks,
             is_valid_sample_fn=self.is_valid_sample_fn,
             should_continue_fn=self.should_continue_fn,
         )
@@ -338,6 +343,7 @@ def __init__(
         over_sample_threshold: float,
         enable_partial_rollout: bool,
         tail_batch_trigger_size: int,
+        max_pending_tasks: int | None,
         max_staleness: int,
         max_token_staleness: int | None,
         sync_weights_interval: int,
@@ -368,6 +374,7 @@ def __init__(
             else calculate_stale_threshold(max_token_staleness, sync_weights_interval)
         )
         self.tail_batch_trigger_size = tail_batch_trigger_size
+        self.max_pending_tasks = max_pending_tasks
         self._local_pending_tasks: set[asyncio.Task] = set()
 
     def pending_task_count(self) -> int:
@@ -434,7 +441,9 @@ async def spawn_one() -> asyncio.Task:
 
             pending_count = len(self._local_pending_tasks)
             desired_pending = max(0, scheduled_target - available)
-            if available + pending_count < scheduled_target:
+            if self.max_pending_tasks is not None:
+                desired_pending = min(desired_pending, self.max_pending_tasks)
+            if pending_count < desired_pending:
                 while len(self._local_pending_tasks) < desired_pending:
                     self._local_pending_tasks.add(await spawn_one())
 

xtuner/v1/rl/rollout/controller.py

@@ -22,7 +22,7 @@
     RolloutConfig,
     get_rollout_worker_base_cls,
 )
-from .worker_registry import RolloutWorkerRegistry
+from .worker_registry import RolloutWorkerRegistry, WorkerLifecycleState
 
 
 # Keep this as a Ray actor because Ray AgentLoop actors need a shared, cross-process handle to the same controller
@@ -70,6 +70,28 @@ def get_weight_update_targets(self) -> tuple[RolloutWeightUpdateTarget, ...]:
         """Return rollout endpoints that can receive weight update requests."""
         return self.registry.weight_update_targets()
 
+    def get_pending_weight_update_targets(self) -> tuple[RolloutWeightUpdateTarget, ...]:
+        """Return recovered rollout endpoints waiting for weights."""
+        return tuple(
+            target
+            for target in self.registry.weight_update_targets()
+            if target.lifecycle_state == WorkerLifecycleState.PENDING_WEIGHTS
+        )
+
+    def inject_backend_crash_for_test(self, *, rank: int = 0) -> None:
+        """Crash one active rollout backend for the immediate-recovery test."""
+        worker = self.registry.active_entrypoint_by_rank(rank)
+        if worker is None:
+            raise RuntimeError(f"No active rollout request entrypoint found for test fault injection: rank={rank}.")
+
+        accepted = ray.get(
+            worker.actor.inject_backend_crash_for_test.remote(),  # type: ignore[attr-defined]
+            timeout=ROLLOUT_RAY_GET_TIMEOUT,
+        )
+        if not accepted:
+            raise RuntimeError(f"Rollout worker rejected test fault injection: rank={rank}, url={worker.url}.")
+        self.logger.warning(f"[ImmediateRecoveryExperiment] backend_crash_injected rank={rank} url={worker.url}")
+
     def register_active_workers_to_proxy(self) -> None:
         if self.proxy_manager is None:
             return
@@ -120,6 +142,11 @@ async def generate(self, rollout_state: RolloutState) -> RolloutState:
                 f"Rollout request timed out after {self.config.rollout_timeout * self.timeout_multiplier} seconds."
             )
             return rollout_state
+        except Exception as e:
+            self.logger.exception(f"RolloutController.generate failed: session_id={session_id}")
+            rollout_state.status = Status.FAILED
+            rollout_state.error_msg = f"Rollout request failed: {type(e).__name__}: {str(e)[:1024]}"
+            return rollout_state
 
     def set_enable_partial_rollout(self, enable: bool) -> None:
         """Propagate enable_partial_rollout flag to all active workers."""
@@ -159,24 +186,48 @@ async def check_and_shutdown_inactive_workers(self):
 
     async def restart_inactive_workers(self):
         """Restart inactive groups before a sync-step weight update."""
-        await asyncio.to_thread(self.health_manager.restart_inactive_workers)
+        groups = await asyncio.to_thread(self.health_manager.restart_inactive_workers)
+        return tuple(group.ranks for group in groups)
+
+    def mark_worker_groups_lifecycle_state(
+        self,
+        group_ranks: list[tuple[int, ...]],
+        source_state: WorkerLifecycleState,
+        target_state: WorkerLifecycleState,
+    ) -> None:
+        """Move selected worker groups from source_state to target_state.
+
+        Only groups whose current lifecycle state matches source_state are considered. When groups are moved to ACTIVE
+        or INACTIVE, the health manager is notified so routing and lifecycle listeners stay in sync.
+        """
+        groups_by_ranks = {group.ranks: group for group in self.registry.get_target_state_worker_groups(source_state)}
+        groups = tuple(groups_by_ranks[ranks] for ranks in group_ranks if ranks in groups_by_ranks)
+        updated_groups = self.registry.set_groups_state(
+            groups,
+            target_state,
+            source_state=source_state,
+        )
+        if target_state is WorkerLifecycleState.ACTIVE:
+            self.health_manager.notify_worker_group_recovered(updated_groups)
+        elif target_state is WorkerLifecycleState.INACTIVE:
+            self.health_manager.notify_worker_group_inactive(updated_groups)
 
     def continue_generation(self):
-        self._broadcast_to_active_workers("continue_generation")
+        self._broadcast_to_workers("continue_generation", WorkerLifecycleState.ACTIVE)
         self.health_manager.resume()
 
     def offload(self):
-        self._broadcast_to_active_workers("offload")
+        self._broadcast_to_workers("offload", WorkerLifecycleState.ACTIVE)
 
     def onload(self):
-        self._broadcast_to_active_workers("onload_weights")
-        self._broadcast_to_active_workers("onload_kvcache")
+        self._broadcast_to_workers("onload_weights", WorkerLifecycleState.ACTIVE)
+        self._broadcast_to_workers("onload_kvcache", WorkerLifecycleState.ACTIVE)
 
-    def onload_weights(self):
-        self._broadcast_to_active_workers("onload_weights")
+    def onload_weights(self, target_state: WorkerLifecycleState = WorkerLifecycleState.ACTIVE):
+        self._broadcast_to_workers("onload_weights", target_state)
 
-    def onload_kvcache(self):
-        self._broadcast_to_active_workers("onload_kvcache")
+    def onload_kvcache(self, target_state: WorkerLifecycleState = WorkerLifecycleState.ACTIVE):
+        self._broadcast_to_workers("onload_kvcache", target_state)
 
     def shutdown(self):
         """Shut down all rollout workers tracked by the controller."""
@@ -187,8 +238,8 @@ def shutdown(self):
             timeout=ROLLOUT_RAY_GET_TIMEOUT,
         )
 
-    def _broadcast_to_active_workers(self, method_name: str, **kwargs):
-        workers = self.registry.active_workers()
+    def _broadcast_to_workers(self, method_name: str, target_state: WorkerLifecycleState, **kwargs):
+        workers = self.registry.get_target_state_workers(target_state)
         futures = [getattr(worker.actor, method_name).remote(**kwargs) for worker in workers]
         return ray.get(futures, timeout=ROLLOUT_RAY_GET_TIMEOUT)