Refine rollout worker health check and recovery lifecycle (#1877)

* Refactor rollout server launch specs

* Refactor rollout health manager

* Restore rollout worker session URLs

* fix

* fix

* fix

* fix comments

* Delay rollout health checks after resume and log LMDeploy health errors

* fix ci

* delete rank2info and worker_server_urls_map

* Require fail-fast health checks before rollout offload

* Align rollout generate concurrency with request entrypoint topology

* Refactor colocate rollout recovery with fail-fast shutdown and pre-sync restart

* simplify code

* fix ci

* fix ci and comments

* Refactor rollout worker state into registry

Author: YanhuiDua

tests/rl/test_rl_colocate_trainer.py

@@ -52,7 +52,6 @@ def __init__(self, uid: int):
         self.extra_fields = {}
         self.response_model_steps = []
 
-
 class _FakeSampler:
     def __init__(self):
         self._next_id = 0
@@ -148,13 +147,18 @@ def _make_trainer(self, agent_loop_manager, *, total_train_steps: int = 1, sync_
         )
 
         trainer.rollout_controller = SimpleNamespace(
-            ensure_workers_healthy_before_training=SimpleNamespace(
-                remote=MagicMock(return_value="rollout_ready_for_training")
+            check_and_shutdown_inactive_workers=SimpleNamespace(
+                remote=MagicMock(return_value="rollout_inactive_workers_shutdown")
             ),
             offload=SimpleNamespace(remote=MagicMock(return_value="rollout_offloaded")),
+            restart_inactive_workers=SimpleNamespace(remote=MagicMock(return_value="rollout_restarted")),
+            onload_weights=SimpleNamespace(remote=MagicMock(return_value="weights_loaded")),
+            onload_kvcache=SimpleNamespace(remote=MagicMock(return_value="kvcache_loaded")),
         )
         trainer.train_controller = SimpleNamespace(
             onload=MagicMock(return_value="train_onloaded"),
+            offload=MagicMock(return_value="train_offloaded"),
+            update_weights=MagicMock(return_value="weights_updated"),
             fit=MagicMock(
                 return_value=[
                     {
@@ -220,15 +224,37 @@ async def _produce_empty(batch_size, train_step, **kwargs):
         trainer.train_controller.fit.assert_not_called()
         self.assertEqual(trainer._cur_step, 0)
 
+    def test_fit_does_not_onload_train_when_rollout_training_barrier_fails(self):
+        # 验证共卡训练进入训练前必须先通过 rollout phase-switch barrier；
+        # 失败时不能 onload 训练。
+        async def _produce_batch(batch_size, train_step, *, model_step):
+            return ProduceBatchResult(rollout_states=[[_FakeRolloutState(train_step)]])
+
+        trainer = self._make_trainer(SimpleNamespace(produce_batch=_produce_batch))
+        trainer.rollout_controller.check_and_shutdown_inactive_workers.remote.side_effect = RuntimeError(
+            "inactive rollout workers after recovery"
+        )
+
+        with (
+            patch("xtuner.v1.train.rl_trainer.asyncio_run", side_effect=asyncio.run),
+            patch("xtuner.v1.train.rl_trainer.ray.get", side_effect=lambda obj, timeout=None: obj),
+        ):
+            with self.assertRaisesRegex(RuntimeError, "inactive rollout workers"):
+                trainer.fit()
+
+        trainer.rollout_controller.check_and_shutdown_inactive_workers.remote.assert_called_once_with()
+        trainer.rollout_controller.offload.remote.assert_not_called()
+        trainer.train_controller.onload.assert_not_called()
+        trainer.train_controller.fit.assert_not_called()
+        self.assertEqual(trainer._cur_step, 0)
+
     def test_fit_uses_sync_interval_and_passes_rollout_model_step(self):
         # 验证 rollout 看到的是按 sync interval 推进后的 model_step。
         produce_calls = []
 
         async def _produce_batch(batch_size, train_step, *, model_step):
             produce_calls.append((batch_size, train_step, model_step))
-            return ProduceBatchResult(
-                rollout_states=[[SimpleNamespace(group_id=train_step, rollout_id=train_step)]]
-            )
+            return ProduceBatchResult(rollout_states=[[_FakeRolloutState(train_step)]])
 
         trainer = self._make_trainer(
             SimpleNamespace(produce_batch=_produce_batch),

tests/rl/test_rl_disaggregated_trainer.py

@@ -146,7 +146,10 @@ def _make_trainer(self, agent_loop_manager):
             update_weights=MagicMock(return_value="update"),
         )
         trainer.rollout_controller = SimpleNamespace(
-            recover_failed_workers=SimpleNamespace(remote=MagicMock(return_value="recover")),
+            check_and_shutdown_inactive_workers=SimpleNamespace(
+                remote=MagicMock(return_value="rollout_inactive_workers_shutdown")
+            ),
+            restart_inactive_workers=SimpleNamespace(remote=MagicMock(return_value="rollout_restarted")),
             pause_generation=SimpleNamespace(remote=MagicMock(return_value="pause")),
             continue_generation=SimpleNamespace(remote=MagicMock(return_value="continue")),
             onload_weights=SimpleNamespace(remote=MagicMock(return_value="onload_weights")),

tests/rl/test_rl_trainer_checkpoint.py

@@ -90,8 +90,8 @@ def __init__(self):
         self.pause_generation = _RemoteMethod(async_result=True)
         self.continue_generation = _RemoteMethod(async_result=True)
         self.offload = _RemoteMethod(return_value="rollout_offloaded")
-        self.ensure_workers_healthy_before_training = _RemoteMethod(return_value="rollout_ready_for_training")
-        self.recover_failed_workers = _RemoteMethod(return_value="rollout_recovered")
+        self.check_and_shutdown_inactive_workers = _RemoteMethod(return_value="rollout_inactive_workers_shutdown")
+        self.restart_inactive_workers = _RemoteMethod(return_value="rollout_restarted")
         self.onload_weights = _RemoteMethod(return_value="weights_loaded")
         self.onload_kvcache = _RemoteMethod(return_value="kvcache_loaded")
         self.get_rollout_metadata = _RemoteMethod(return_value={"server_url_dict": {}})
@@ -204,6 +204,7 @@ def build_rollout_controller(rollout_cfg, placement_group):
             return controller
 
         with (
+            patch("ray.get", side_effect=lambda obj, timeout=None: obj),
             patch("xtuner.v1.rl.utils.ray_accelerator_worker.ray.is_initialized", return_value=True),
             patch(
                 "xtuner.v1.rl.utils.ray_accelerator_worker.ray.available_resources",
@@ -217,6 +218,12 @@ def build_rollout_controller(rollout_cfg, placement_group):
             patch("xtuner.v1.train.rl_trainer.BaseRLTrainer._release_trace_store", return_value=None),
             patch.object(WorkerConfig, "build", autospec=True, side_effect=build_train_controller),
             patch.object(RolloutConfig, "build", autospec=True, side_effect=build_rollout_controller),
+            patch.object(
+                RolloutConfig,
+                "get_controller_generate_concurrency",
+                autospec=True,
+                side_effect=lambda rollout_cfg, placement_group: rollout_cfg.generate_concurrency_per_instance,
+            ),
         ):
             yield runtime
 
@@ -321,6 +328,7 @@ def _build_colocate_config(
             auto_resume=auto_resume,
             checkpoint_interval=1,
             checkpoint_maxkeep=None,
+            checkpoint_no_save_replay_buffer=True,
             hf_interval=-1,
             seed=42,
             exp_tracker="jsonl",
@@ -361,6 +369,7 @@ def _build_disaggregated_config(
             auto_resume=auto_resume,
             checkpoint_interval=1,
             checkpoint_maxkeep=None,
+            checkpoint_no_save_replay_buffer=True,
             hf_interval=-1,
             seed=42,
             exp_tracker="jsonl",