diff --git a/.dev_scripts/test_qwen35_vl_moe_async_hf_recovery_e2e.py b/.dev_scripts/test_qwen35_vl_moe_async_hf_recovery_e2e.py new file mode 100644 index 000000000..cf016b0f1 --- /dev/null +++ b/.dev_scripts/test_qwen35_vl_moe_async_hf_recovery_e2e.py @@ -0,0 +1,504 @@ +"""Real Qwen3.5 VLM MoE async-HF recovery E2E test. + +This test focuses only on the recovery protocol: + +1. train step 1 starts an asynchronous recovery-HF export; +2. the test waits until that export has been published as ready; +3. while train step 2 rollout is running, rank 0's backend is crashed; +4. RolloutHealthManager detects the failure and reloads the ready HF; +5. train step 2 and the post-recovery train step 3 both complete. + +Run in the same 8-GPU environment used by the Qwen3.5 VLM MoE +async-training E2E test. +""" + +from __future__ import annotations + +import asyncio +import os +import threading +import time +import unittest +from concurrent.futures import Future +from pathlib import Path +from typing import Any, Callable + + +# These values are consumed while XTuner modules are imported. +os.environ["XTUNER_DETERMINISTIC"] = "false" +os.environ["XTUNER_USE_LMDEPLOY"] = "1" +os.environ["XTUNER_USE_SGLANG"] = "0" +os.environ["XTUNER_USE_VLLM"] = "0" +os.environ["XTUNER_TEST_IMMEDIATE_RECOVERY"] = "1" + +import ray + +from xtuner.v1.config import AdamWConfig, FSDPConfig, LRConfig +from xtuner.v1.data_proto.rl_data import SampleParams +from xtuner.v1.datasets.config import DataloaderConfig, DatasetConfig +from xtuner.v1.datasets.rl_tokenize_fn import RLQwen3VLTokenizeFnConfig +from xtuner.v1.model import Qwen3_5_VLMoE35BA3Config +from xtuner.v1.rl.advantage import GRPOAdvantageConfig +from xtuner.v1.rl.agent_loop import SingleTurnAgentLoopConfig +from xtuner.v1.rl.agent_loop_manager import ( + AgentLoopManagerConfig, + AsyncProduceStrategyConfig, + SamplerConfig, + TaskSpecConfig, +) +from xtuner.v1.rl.judger import GEO3KJudgerConfig +from xtuner.v1.rl.loss import GRPOLossConfig +from xtuner.v1.rl.replay_buffer import AsyncReplayBufferConfig +from xtuner.v1.rl.rollout.worker import RolloutConfig +from xtuner.v1.rl.rollout.worker_registry import WorkerLifecycleState +from xtuner.v1.rl.trainer import RolloutImportanceSampling, WorkerConfig +from xtuner.v1.rl.utils import AcceleratorResourcesConfig, CPUResourcesConfig +from xtuner.v1.train.rl_trainer import RLColocateTrainerConfig + + +EXPERIMENT_NAME = "qwen35_vl_moe_async_hf_recovery_e2e" +TOTAL_TRAIN_STEPS = 3 +TRAIN_BATCH_SIZE_BY_STEP = {1: 8, 2: 128, 3: 8} +PROMPT_REPEAT_K = 2 +MAX_PROMPT_LENGTH = 4096 +MAX_RESPONSE_LENGTH = 2048 +PACK_MAX_LENGTH = 8192 +RECOVERY_TIMEOUT_S = 600.0 +RAY_GET_TIMEOUT_S = 600.0 +POLL_INTERVAL_S = 0.5 + + +class TestQwen35VLMoEAsyncHFRecoveryE2E(unittest.TestCase): + def setUp(self) -> None: + self.model_path = self._required_path("QWEN3_5_MOE_PATH") + self.media_root = self._required_path("GEO3K_MEDIA_ROOT") + self.data_path = self._required_path("GEO3K_LONGTAIL_DATA_PATH") + + default_work_dir = ( + Path.cwd() / "work_dirs" / f"{EXPERIMENT_NAME}_{time.strftime('%Y%m%d%H%M%S')}_{os.getpid()}" + ) + self.work_dir = Path(os.environ.get("WORK_DIR", str(default_work_dir))) + self.work_dir.mkdir(parents=True, exist_ok=True) + + self._events: list[str] = [] + self._events_lock = threading.Lock() + self._pending_export_captured = threading.Event() + self._rollout_step_2_started = threading.Event() + self._rollout_step_2_finished = threading.Event() + self._recovery_finished = threading.Event() + self._first_recovery_hf_export: Future[Path | None] | None = None + self._fault_injection_error: Exception | None = None + self._recovery_hf_observation: dict[str, Any] | None = None + self._rank_0_lifecycle_states: list[str] = [] + self._produce_calls: list[dict[str, int]] = [] + + self._patch_env( + { + "XTUNER_USE_LMDEPLOY": "1", + "XTUNER_USE_SGLANG": "0", + "XTUNER_USE_VLLM": "0", + "XTUNER_USE_FA3": "1", + "XTUNER_DETERMINISTIC": "false", + "XTUNER_TEST_IMMEDIATE_RECOVERY": "1", + "PYTORCH_CUDA_ALLOC_CONF": "expandable_segments:True", + }, + unset=("RAY_ADDRESS",), + ) + ray.init(address="local", num_cpus=128, num_gpus=8, ignore_reinit_error=True) + + def tearDown(self) -> None: + if ray.is_initialized(): + ray.shutdown() + if hasattr(self, "_old_env"): + self._restore_env() + + def test_async_hf_save_and_backend_failure_recovery(self) -> None: + trainer = self._build_config().build() + self._install_rollout_probe(trainer) + self._install_async_hf_probe(trainer) + + fault_injection_thread = threading.Thread( + target=self._inject_failure_after_recovery_hf_ready, + args=(trainer,), + name="async-hf-recovery-fault-injector", + daemon=True, + ) + fault_injection_thread.start() + + try: + trainer.fit() + finally: + fault_injection_thread.join(timeout=10) + + self.assertFalse(fault_injection_thread.is_alive(), "Fault-injection coordinator did not exit.") + if self._fault_injection_error is not None: + raise AssertionError("Fault-injection coordinator failed.") from self._fault_injection_error + + observation = self._recovery_hf_observation + self.assertIsNotNone(observation) + assert observation is not None + self.assertTrue(observation["path_existed"]) + self.assertGreater(observation["file_count"], 0) + unavailable_states = { + WorkerLifecycleState.INACTIVE.value, + WorkerLifecycleState.RECOVERING.value, + } + self.assertTrue(unavailable_states.intersection(self._rank_0_lifecycle_states)) + self.assertEqual(self._rank_0_lifecycle_states[-1], WorkerLifecycleState.ACTIVE.value) + self.assertEqual( + [call["train_step"] for call in self._produce_calls], + [1, 2, 3], + ) + self.assertEqual( + [call["batch_size"] for call in self._produce_calls], + [TRAIN_BATCH_SIZE_BY_STEP[step] for step in range(1, TOTAL_TRAIN_STEPS + 1)], + ) + self._assert_recovery_event_order() + + def _install_rollout_probe(self, trainer: Any) -> None: + original_produce_batch = trainer.agent_loop_manager.produce_batch + + async def produce_batch_wrapper(batch_size: int, train_step: int, *, model_step: int) -> Any: + batch_size = TRAIN_BATCH_SIZE_BY_STEP.get(train_step, batch_size) + self._record_event(f"rollout_{train_step}_started") + if train_step == 2: + self._rollout_step_2_started.set() + + try: + result = await original_produce_batch(batch_size, train_step, model_step=model_step) + self._produce_calls.append( + { + "batch_size": batch_size, + "train_step": train_step, + "model_step": model_step, + } + ) + if train_step == 2: + recovered = await asyncio.to_thread( + self._recovery_finished.wait, + RECOVERY_TIMEOUT_S, + ) + if not recovered: + raise TimeoutError("Timed out waiting for rank 0 recovery during train step 2 rollout.") + return result + finally: + if train_step == 2: + self._rollout_step_2_finished.set() + self._record_event(f"rollout_{train_step}_finished") + + trainer.agent_loop_manager.produce_batch = produce_batch_wrapper + + def _install_async_hf_probe(self, trainer: Any) -> None: + original_maybe_save_recovery_hf = trainer._maybe_save_recovery_hf + + def maybe_save_recovery_hf_wrapper(cur_step: int) -> None: + original_maybe_save_recovery_hf(cur_step) + if cur_step != 1: + return + + pending_export = trainer._pending_hf_export + if pending_export is None: + raise AssertionError("Train step 1 did not schedule an asynchronous recovery-HF export.") + self._first_recovery_hf_export = pending_export + self._record_event("async_hf_1_scheduled") + self._pending_export_captured.set() + + trainer._maybe_save_recovery_hf = maybe_save_recovery_hf_wrapper + + def _inject_failure_after_recovery_hf_ready(self, trainer: Any) -> None: + try: + if not self._pending_export_captured.wait(timeout=RECOVERY_TIMEOUT_S): + raise TimeoutError("Timed out waiting for the train step 1 recovery-HF export to be scheduled.") + + pending_export = self._first_recovery_hf_export + if pending_export is None: + raise AssertionError("Recovery-HF export event was set without a Future.") + recovery_hf_path = pending_export.result(timeout=RECOVERY_TIMEOUT_S) + if recovery_hf_path is None: + raise RuntimeError("Train step 1 recovery-HF export failed.") + + self._recovery_hf_observation = { + "path": str(recovery_hf_path), + "path_existed": recovery_hf_path.is_dir(), + "file_count": sum(1 for path in recovery_hf_path.rglob("*") if path.is_file()), + } + self._record_event("async_hf_1_ready") + + if not self._rollout_step_2_started.wait(timeout=RECOVERY_TIMEOUT_S): + raise TimeoutError("Timed out waiting for train step 2 rollout to start.") + if self._rollout_step_2_finished.is_set(): + raise RuntimeError("Train step 2 rollout finished before backend failure injection.") + + initial_state = self._get_rank_0_lifecycle_state(trainer) + if initial_state != WorkerLifecycleState.ACTIVE.value: + raise RuntimeError(f"Rank 0 was not active before fault injection: state={initial_state}.") + self._record_rank_0_state(initial_state) + + ray.get( + trainer.rollout_controller.inject_backend_crash_for_test.remote(rank=0), + timeout=RAY_GET_TIMEOUT_S, + ) + self._record_event("backend_crash_injected") + + self._wait_for_rank_0_state( + trainer, + expected=lambda state: state != WorkerLifecycleState.ACTIVE.value, + description="become inactive", + ) + self._record_event("rank_0_unavailable") + self._wait_for_rank_0_state( + trainer, + expected=lambda state: state == WorkerLifecycleState.ACTIVE.value, + description="recover to active", + ) + self._record_event("rank_0_recovered") + except Exception as error: + self._fault_injection_error = error + finally: + self._recovery_finished.set() + + def _wait_for_rank_0_state( + self, + trainer: Any, + *, + expected: Callable[[str], bool], + description: str, + ) -> str: + deadline = time.monotonic() + RECOVERY_TIMEOUT_S + while time.monotonic() < deadline: + state = self._get_rank_0_lifecycle_state(trainer) + self._record_rank_0_state(state) + if expected(state): + return state + time.sleep(POLL_INTERVAL_S) + raise TimeoutError( + f"Timed out waiting for rank 0 to {description}; observed states={self._rank_0_lifecycle_states}." + ) + + @staticmethod + def _get_rank_0_lifecycle_state(trainer: Any) -> str: + targets = ray.get( + trainer.rollout_controller.get_weight_update_targets.remote(), + timeout=RAY_GET_TIMEOUT_S, + ) + for target in targets: + if target.endpoint_rank == 0: + return target.lifecycle_state + raise RuntimeError(f"Rank 0 weight-update target was not found: targets={targets}.") + + def _record_rank_0_state(self, state: str) -> None: + if not self._rank_0_lifecycle_states or self._rank_0_lifecycle_states[-1] != state: + self._rank_0_lifecycle_states.append(state) + + def _assert_recovery_event_order(self) -> None: + required_events = ( + "async_hf_1_scheduled", + "async_hf_1_ready", + "rollout_2_started", + "backend_crash_injected", + "rank_0_unavailable", + "rank_0_recovered", + "rollout_2_finished", + "rollout_3_started", + "rollout_3_finished", + ) + for event in required_events: + self.assertEqual(self._events.count(event), 1, f"Unexpected event count for {event}: {self._events}") + + positions = {event: self._events.index(event) for event in required_events} + ordered_pairs = ( + ("async_hf_1_scheduled", "async_hf_1_ready"), + ("async_hf_1_ready", "backend_crash_injected"), + ("rollout_2_started", "backend_crash_injected"), + ("backend_crash_injected", "rank_0_unavailable"), + ("rank_0_unavailable", "rank_0_recovered"), + ("rank_0_recovered", "rollout_2_finished"), + ("rollout_2_finished", "rollout_3_started"), + ("rollout_3_started", "rollout_3_finished"), + ) + for first, second in ordered_pairs: + self.assertLess(positions[first], positions[second], f"Expected {first} before {second}: {self._events}") + + def _record_event(self, event: str) -> None: + with self._events_lock: + self._events.append(event) + + def _build_config(self) -> RLColocateTrainerConfig: + resources = AcceleratorResourcesConfig( + accelerator="GPU", + num_workers=8, + num_cpus_per_worker=12, + cpu_memory_per_worker=16 * 1024**3, + ) + rollout_config = RolloutConfig( + env=EXPERIMENT_NAME, + device=resources.accelerator, + model_path=str(self.model_path), + tokenizer_path=str(self.model_path), + dtype="bfloat16", + tensor_parallel_size=1, + expert_parallel_size=2, + gpu_memory_utilization=0.8, + context_length=MAX_PROMPT_LENGTH + MAX_RESPONSE_LENGTH, + rollout_max_batch_size_per_instance=128, + allow_over_concurrency_ratio=1.0, + enable_return_routed_experts=True, + health_check_interval_seconds=5.0, + health_check_failure_threshold=1, + extra_rollout_config={ + "lmdeploy_backend": "pytorch", + "lmdeploy_log_level": "ERROR", + "lmdeploy_uvicorn_log_level": "ERROR", + }, + ) + + train_worker_cfg = WorkerConfig( + model_cfg=Qwen3_5_VLMoE35BA3Config(freeze_vision=True, freeze_projector=True), + load_from=str(self.model_path), + optim_cfg=AdamWConfig( + lr=1e-6, + betas=(0.9, 0.999), + max_grad_norm=1.0, + weight_decay=0.1, + foreach=False, + ), + loss_cfg=GRPOLossConfig( + policy_loss_cfg={ + "cliprange_high": 0.28, + "cliprange_low": 0.2, + "loss_type": "vanilla", + "clip_ratio_c": 10.0, + "log_prob_diff_min": -20, + "log_prob_diff_max": 20, + }, + ignore_idx=-100, + use_kl_loss=False, + kl_loss_coef=0.0, + kl_loss_type="low_var_kl", + mode="chunk", + chunk_size=512, + rollout_is=RolloutImportanceSampling( + rollout_is_level="token", + rollout_is_mode="both", + rollout_is_threshold=(5, 0.5), + rollout_is_mask_threshold=(5, 0.5), + rollout_is_veto_threshold=(20, 0), + ), + ), + lr_cfg=LRConfig(lr_type="constant", warmup_ratio=0, lr_min=1e-6), + fsdp_cfg=FSDPConfig(torch_compile=False, cpu_offload=False, ep_size=1, fp32_lm_head=True), + sp_size=1, + optimizer_steps=8, + pack_max_length=PACK_MAX_LENGTH, + ) + + dataloader_cfg = DataloaderConfig( + dataset_config_list=[ + { + "dataset": DatasetConfig( + name=EXPERIMENT_NAME, + anno_path=self.data_path, + class_name="VLMJsonlDataset", + media_root=str(self.media_root), + ), + "tokenize_fn": RLQwen3VLTokenizeFnConfig( + processor_path=str(self.model_path), + max_length=MAX_PROMPT_LENGTH, + chat_template="qwen3.5-vl", + add_generation_prompt=True, + enable_thinking=True, + ), + } + ], + pack_max_length=PACK_MAX_LENGTH, + collator="fake_collator", + pack_level="none", + ) + agent_loop_manager_cfg = AgentLoopManagerConfig( + tasks=[ + TaskSpecConfig( + task_name="geo3k_longtail", + agent_loop_config=SingleTurnAgentLoopConfig( + hf_checkpoint=str(self.model_path), + sample_params=SampleParams( + max_tokens=MAX_RESPONSE_LENGTH, + top_k=0, + top_p=1.0, + temperature=0.0, + min_tokens=0, + return_logprob=True, + return_token_ids=True, + return_routed_experts=True, + ), + ), + judger_config=GEO3KJudgerConfig( + judger_name="hiyouga/geometry3k", + cpu_resources=CPUResourcesConfig(num_workers=1, num_cpus_per_worker=1), + ), + produce_strategy_config=AsyncProduceStrategyConfig( + over_sample_threshold=1.0, + enable_partial_rollout=True, + max_staleness=1, + ), + sampler_config=SamplerConfig( + dataloader_cfg=dataloader_cfg, + prompt_repeat_k=PROMPT_REPEAT_K, + ), + ) + ], + ) + + return RLColocateTrainerConfig( + resources=resources, + train_worker_cfg=train_worker_cfg, + rollout_config=rollout_config, + tokenizer_path=str(self.model_path), + replay_buffer_config=AsyncReplayBufferConfig(), + agent_loop_manager_cfg=agent_loop_manager_cfg, + load_from=str(self.model_path), + total_train_steps=TOTAL_TRAIN_STEPS, + train_batch_size=TRAIN_BATCH_SIZE_BY_STEP[1], + advantage_estimator_config=GRPOAdvantageConfig(eps=1e-8), + sync_weights_interval=1, + enable_evaluate=False, + enable_initial_evaluate=False, + evaluate_step=1, + work_dir=str(self.work_dir), + checkpoint_interval=-1, + checkpoint_maxkeep=-1, + hf_interval=-1, + hf_max_keep=-1, + enable_immediate_recovery=True, + seed=123, + debug_rollout=False, + exp_tracker="jsonl", + ) + + @staticmethod + def _required_path(env_name: str) -> Path: + value = os.environ.get(env_name) + if not value: + raise RuntimeError(f"{env_name} must be set for the async-HF recovery E2E test.") + path = Path(value) + if not path.exists(): + raise FileNotFoundError(f"{env_name} does not exist: {path}") + return path + + def _patch_env(self, updates: dict[str, str], *, unset: tuple[str, ...] = ()) -> None: + keys = set(updates) | set(unset) + self._old_env = {key: os.environ.get(key) for key in keys} + for key, value in updates.items(): + os.environ[key] = value + for key in unset: + os.environ.pop(key, None) + + def _restore_env(self) -> None: + for key, value in self._old_env.items(): + if value is None: + os.environ.pop(key, None) + else: + os.environ[key] = value + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/rl/test_rl_colocate_trainer.py b/tests/rl/test_rl_colocate_trainer.py index 69f902d6d..6e3c5acb6 100644 --- a/tests/rl/test_rl_colocate_trainer.py +++ b/tests/rl/test_rl_colocate_trainer.py @@ -123,6 +123,7 @@ def _make_trainer(self, agent_loop_manager, *, total_train_steps: int = 1, sync_ trainer._enable_evaluate = False trainer._enable_initial_evaluate = False trainer._evaluate_step = 1 + trainer._hf_export_executor = None trainer._cpu_resource_manager = None trainer._train_worker_cfg = SimpleNamespace(pack_max_length=16) trainer._meta = SimpleNamespace( diff --git a/tests/rl/test_rl_disaggregated_trainer.py b/tests/rl/test_rl_disaggregated_trainer.py index 80b96c6da..2f9e483a9 100644 --- a/tests/rl/test_rl_disaggregated_trainer.py +++ b/tests/rl/test_rl_disaggregated_trainer.py @@ -113,6 +113,9 @@ def _make_trainer(self, agent_loop_manager): trainer._enable_initial_evaluate = False trainer._evaluate_step = 1 trainer._debug_rollout = False + trainer._enable_immediate_recovery = False + trainer._hf_export_executor = None + trainer._pending_hf_export = None trainer._display_all_workers_log = False trainer._num_workers = 1.0 trainer._rollout_num_workers = 1.0 @@ -150,7 +153,7 @@ def _make_trainer(self, agent_loop_manager): check_and_shutdown_inactive_workers=SimpleNamespace( remote=MagicMock(return_value="rollout_inactive_workers_shutdown") ), - restart_inactive_workers=SimpleNamespace(remote=MagicMock(return_value="rollout_restarted")), + restart_inactive_workers=SimpleNamespace(remote=AsyncMock(return_value=None)), 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")), @@ -160,7 +163,7 @@ def _make_trainer(self, agent_loop_manager): return trainer def _run_fit(self, trainer): - with patch("xtuner.v1.train.rl_trainer.asyncio_run", side_effect=asyncio.run): + with patch("xtuner.v1.train.rl_trainer.ray.get", side_effect=lambda obj, timeout=None: obj): trainer.fit() def _minimal_train_info(self, *, training_samples: int, training_tokens: int, benchmark_end_time_s: float = 108.0): @@ -202,7 +205,6 @@ def test_fit_persists_checkpoint_for_completed_model_step(self): ) 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), patch("xtuner.v1.train.rl_trainer.bind_train_rollout"), ): @@ -263,10 +265,7 @@ def test_fit_rebinds_weight_update_with_rollout_update_address(self): trainer = self._make_trainer(manager) trainer._rollout_config = SimpleNamespace(weight_update_host="10.0.0.1", weight_update_port=23456) - with ( - patch("xtuner.v1.train.rl_trainer.asyncio_run", side_effect=asyncio.run), - patch("xtuner.v1.train.rl_trainer.bind_train_rollout") as bind_train_rollout_mock, - ): + with patch("xtuner.v1.train.rl_trainer.bind_train_rollout") as bind_train_rollout_mock: trainer.fit() bind_train_rollout_mock.assert_called_once_with( diff --git a/tests/rl/test_rl_recovery_hf.py b/tests/rl/test_rl_recovery_hf.py new file mode 100644 index 000000000..41a329533 --- /dev/null +++ b/tests/rl/test_rl_recovery_hf.py @@ -0,0 +1,534 @@ +"""RL immediate recovery 的故障恢复与慢导出回退测试。""" + +import asyncio +import inspect +import os +import tempfile +import unittest +from concurrent.futures import Future +from pathlib import Path +from types import SimpleNamespace +from unittest.mock import AsyncMock, MagicMock, patch + +from xtuner.v1.rl.rollout.controller import RolloutController +from xtuner.v1.rl.rollout.health_manager import RolloutHealthManager +from xtuner.v1.rl.rollout.rollout_topology import RolloutEngine, RolloutServerProcess, RolloutTopology +from xtuner.v1.rl.rollout.worker import RolloutWorker, RolloutWorkerInitResult +from xtuner.v1.rl.rollout.worker_registry import RolloutWorkerRegistry +from xtuner.v1.train.rl_trainer import RLColocateTrainer, RLDisaggregatedTrainer + + +def _identity_ray_get(value, *, timeout=None): + del timeout + return value + + +class _LocalRemoteMethod: + """Expose a local callable through the small Ray method surface used here.""" + + def __init__(self, func): + self._func = func + self.calls = [] + + def remote(self, *args, **kwargs): + self.calls.append((args, kwargs)) + return self._func(*args, **kwargs) + + +class _LocalAsyncRemoteMethod: + """Expose a local callable through the async ``.remote()`` surface used by Ray actors.""" + + def __init__(self, func): + self._func = func + self.calls = [] + + def remote(self, *args, **kwargs): + self.calls.append((args, kwargs)) + + async def invoke(): + return self._func(*args, **kwargs) + + return invoke() + + +class _RecoverableRolloutActor: + """A fake actor whose backend health changes across crash and reinit.""" + + def __init__(self, rank: int, *, healthy: bool = True): + self.rank = rank + self.server_url = f"http://worker-{rank}" + self.session_url = f"http://session-{rank}" + self.healthy = healthy + self.loaded_model_path = None + self.loaded_tokenizer_path = None + self.loaded_skip_load_weights = None + + worker = RolloutWorker.__new__(RolloutWorker) + worker.rank = rank + worker.server_url = self.server_url + worker.server_task = object() + worker.server_process = None + worker.logger = MagicMock() + self.worker = worker + + self.inject_backend_crash_for_test = _LocalAsyncRemoteMethod(worker.inject_backend_crash_for_test) + self.check_health = _LocalAsyncRemoteMethod(lambda: self.healthy) + self.shutdown = _LocalAsyncRemoteMethod(self._shutdown) + self.reinit = _LocalAsyncRemoteMethod(self._reinit) + self.offload = _LocalAsyncRemoteMethod(lambda: None) + + def _shutdown(self) -> None: + self.healthy = False + + def _reinit( + self, + *, + model_path: str | None = None, + tokenizer_path: str | None = None, + skip_load_weights: bool | None = None, + ) -> RolloutWorkerInitResult: + self.loaded_model_path = model_path + self.loaded_tokenizer_path = tokenizer_path + self.loaded_skip_load_weights = skip_load_weights + self.healthy = True + return RolloutWorkerInitResult( + rank=self.rank, + server_url=self.server_url, + session_url=self.session_url, + ) + + +def _local_ray_get(value, *, timeout=None): + del timeout + + def resolve(item): + if inspect.isawaitable(item): + return asyncio.run(item) + return item + + if isinstance(value, list): + return [resolve(item) for item in value] + return resolve(value) + + +class TestRLImmediateRecovery(unittest.TestCase): + _MODEL_PATH = "/tmp/ready-recovery-hf" + _TOKENIZER_PATH = "/tmp/ready-recovery-tokenizer" + + def _build_registry(self, actors: list[_RecoverableRolloutActor]) -> RolloutWorkerRegistry: + topology = RolloutTopology( + engines=tuple( + RolloutEngine( + engine_ranks=(actor.rank,), + dist_init_addr=f"addr-{actor.rank}", + server_processes=( + RolloutServerProcess( + worker_rank=actor.rank, + placement_group_bundle_idxs=(actor.rank,), + accepts_rollout_requests=True, + weight_update_ranks=(actor.rank,), + ), + ), + ) + for actor in actors + ) + ) + registry = RolloutWorkerRegistry(rollout_topology=topology) + registry.register_started_servers( + init_results=tuple( + RolloutWorkerInitResult( + rank=actor.rank, + server_url=actor.server_url, + session_url=actor.session_url, + ) + for actor in actors + ), + workers_by_rank=tuple(actors), + ) + return registry + + def _build_health_manager( + self, + registry: RolloutWorkerRegistry, + *, + listeners=(), + ) -> RolloutHealthManager: + config = SimpleNamespace( + health_check_interval_seconds=10, + health_check_timeout_seconds=1.0, + health_check_failure_threshold=1, + ) + return RolloutHealthManager( + config=config, + registry=registry, + worker_lifecycle_listeners=listeners, + ) + + def _build_controller( + self, + registry: RolloutWorkerRegistry, + health_manager: RolloutHealthManager, + ) -> RolloutController: + controller = RolloutController.__new__(RolloutController) + controller.logger = MagicMock() + controller.registry = registry + controller.health_manager = health_manager + return controller + + def test_injected_single_server_failure_recovers_immediately_from_ready_hf(self): + """测试内容：注入单个 rollout server 故障后，使用 ready HF 立即恢复 active 状态。 + + 测试流程： + 1. 注册一个健康 worker，并向 HealthManager 发布已经完成的 ready recovery HF。 + 2. 通过 Controller 调用 Worker 的测试注入接口强制停止 backend server。 + 3. 执行一次健康检查，由 HealthManager 检测并立即恢复该 worker。 + 4. 验证 reinit 使用 ready HF，且 worker 重新变为 active。 + """ + actor = _RecoverableRolloutActor(rank=0) + registry = self._build_registry([actor]) + health_manager = self._build_health_manager(registry) + health_manager.set_ready_recovery_hf( + model_path=self._MODEL_PATH, + tokenizer_path=self._TOKENIZER_PATH, + ) + controller = self._build_controller(registry, health_manager) + server_task = actor.worker.server_task + + def cancel_backend(*args, **kwargs): + del args, kwargs + actor.healthy = False + + with ( + patch.dict(os.environ, {"XTUNER_TEST_IMMEDIATE_RECOVERY": "1"}), + patch("xtuner.v1.rl.rollout.controller.ray.get", side_effect=_local_ray_get), + patch("xtuner.v1.rl.rollout.health_manager.ray.get", side_effect=_local_ray_get), + patch("xtuner.v1.rl.rollout.worker.ray.cancel", side_effect=cancel_backend) as ray_cancel, + ): + controller.inject_backend_crash_for_test(rank=0) + health_manager.run_once() + + ray_cancel.assert_called_once_with(server_task, force=True, recursive=True) + self.assertIsNone(actor.worker.server_task) + self.assertEqual( + actor.reinit.calls, + [ + ( + (), + { + "model_path": self._MODEL_PATH, + "tokenizer_path": self._TOKENIZER_PATH, + "skip_load_weights": False, + }, + ) + ], + ) + self.assertTrue(registry.active_entrypoint_by_rank(0).is_active()) + self.assertEqual(actor.offload.calls, []) + + def test_all_server_groups_failure_recovers_every_group_from_same_ready_hf(self): + """测试内容：所有 rollout server group 同时失效后，全部使用同一 ready HF 立即恢复。 + + 测试流程： + 1. 注册三个独立 lifecycle group，并预先发布同一个 ready recovery HF。 + 2. 将三个 fake backend 同时置为不健康，执行一次健康检查模拟全部 worker 挂掉。 + 3. HealthManager 检测故障后自动恢复全部 inactive group。 + 4. 验证每组都用同一路径 reinit、恢复通知完整，最终所有 worker 重新 active。 + """ + actors = [_RecoverableRolloutActor(rank=rank) for rank in range(3)] + registry = self._build_registry(actors) + inactive_groups = [] + recovered_groups = [] + lifecycle_events = [] + listener = SimpleNamespace( + on_worker_group_inactive=lambda group: ( + inactive_groups.append(group), + lifecycle_events.append(("inactive", group.ranks)), + ), + on_worker_group_recovered=lambda group: ( + recovered_groups.append(group), + lifecycle_events.append(("recovered", group.ranks)), + ), + ) + health_manager = self._build_health_manager(registry, listeners=[listener]) + health_manager.set_ready_recovery_hf( + model_path=self._MODEL_PATH, + tokenizer_path=self._TOKENIZER_PATH, + ) + for actor in actors: + actor.healthy = False + + with patch("xtuner.v1.rl.rollout.health_manager.ray.get", side_effect=_local_ray_get): + health_manager.run_once() + + expected_reinit_call = [ + ( + (), + { + "model_path": self._MODEL_PATH, + "tokenizer_path": self._TOKENIZER_PATH, + "skip_load_weights": False, + }, + ) + ] + self.assertEqual([group.ranks for group in inactive_groups], [(0,), (1,), (2,)]) + self.assertEqual([group.ranks for group in recovered_groups], [(0,), (1,), (2,)]) + self.assertEqual( + lifecycle_events, + [ + ("inactive", (0,)), + ("inactive", (1,)), + ("inactive", (2,)), + ("recovered", (0,)), + ("recovered", (1,)), + ("recovered", (2,)), + ], + ) + self.assertTrue(all(worker.is_active() for worker in registry.all_workers())) + for actor in actors: + self.assertEqual(actor.reinit.calls, expected_reinit_call) + self.assertEqual(actor.loaded_model_path, self._MODEL_PATH) + self.assertEqual(actor.loaded_tokenizer_path, self._TOKENIZER_PATH) + self.assertFalse(actor.loaded_skip_load_weights) + self.assertEqual(actor.offload.calls, []) + + def test_inactive_group_recovers_on_health_run_after_ready_hf_is_published(self): + """测试内容：ready HF 只唤醒 HealthManager，由下一次 health run 恢复 inactive group。""" + actor = _RecoverableRolloutActor(rank=0, healthy=False) + registry = self._build_registry([actor]) + health_manager = self._build_health_manager(registry) + + with patch("xtuner.v1.rl.rollout.health_manager.ray.get", side_effect=_local_ray_get): + health_manager.run_once() + self.assertIsNone(registry.active_entrypoint_by_rank(0)) + self.assertEqual(actor.reinit.calls, []) + + health_manager.set_ready_recovery_hf( + model_path=self._MODEL_PATH, + tokenizer_path=self._TOKENIZER_PATH, + ) + self.assertIsNone(registry.active_entrypoint_by_rank(0)) + self.assertEqual(actor.reinit.calls, []) + self.assertTrue(health_manager._health_loop_wakeup_event.is_set()) + + # Simulate the background loop consuming the wakeup and owning the + # next health/recovery workflow. + health_manager._health_loop_wakeup_event.clear() + health_manager.run_once() + + self.assertTrue(registry.active_entrypoint_by_rank(0).is_active()) + self.assertEqual( + actor.reinit.calls, + [ + ( + (), + { + "model_path": self._MODEL_PATH, + "tokenizer_path": self._TOKENIZER_PATH, + "skip_load_weights": False, + }, + ) + ], + ) + +class TestImmediateRecoveryHFInterval(unittest.TestCase): + def test_hf_interval_reuses_regular_hf_without_async_export(self): + """测试内容：命中 hf_interval 时复用常规 HF，并跳过 recovery HF 的异步导出。 + + 测试流程： + 1. 构造 step 2 命中 hf_interval=2 的 trainer，并准备常规保存产生的 hf-step-2。 + 2. 调用 recovery HF 调度入口，模拟常规 HF 保存完成后的 immediate-recovery 发布阶段。 + 3. 验证该路径直接发布 hf-step-2，并将其记录为当前 ready recovery HF。 + 4. 验证 start_hf_export 从未调用，且没有创建 pending 异步导出 Future。 + """ + with tempfile.TemporaryDirectory() as temp_dir: + trainer = RLColocateTrainer.__new__(RLColocateTrainer) + trainer._enable_immediate_recovery = True + trainer._hf_interval = 2 + trainer._total_train_steps = 4 + trainer._rollout_config = SimpleNamespace( + tokenizer_path="/tmp/tokenizer", + model_path="/tmp/model", + ) + trainer._ready_recovery_hf_path = None + trainer._pending_hf_export = None + trainer._meta = SimpleNamespace( + latest_exp=SimpleNamespace( + exp_dir=temp_dir, + hf_checkpoint_list=[], + ), + ) + regular_hf_path = trainer.exp_dir / trainer._HF_DIR / "hf-step-2" + regular_hf_path.mkdir(parents=True) + trainer._meta.latest_exp.hf_checkpoint_list.append(str(regular_hf_path)) + + clear_ready_hf = _LocalRemoteMethod(lambda: None) + set_ready_hf = _LocalRemoteMethod(lambda **_kwargs: None) + trainer.rollout_controller = SimpleNamespace( + clear_ready_recovery_hf=clear_ready_hf, + set_ready_recovery_hf=set_ready_hf, + ) + start_hf_export = MagicMock() + trainer.train_controller = SimpleNamespace(start_hf_export=start_hf_export) + + with patch("xtuner.v1.train.rl_trainer.ray.get", side_effect=_identity_ray_get): + trainer._maybe_save_recovery_hf(cur_step=2) + + self.assertEqual(clear_ready_hf.calls, [((), {})]) + self.assertEqual( + set_ready_hf.calls, + [ + ( + (), + { + "model_path": str(regular_hf_path), + "tokenizer_path": "/tmp/tokenizer", + }, + ) + ], + ) + start_hf_export.assert_not_called() + self.assertIsNone(trainer._pending_hf_export) + self.assertEqual(trainer._ready_recovery_hf_path, regular_hf_path) + + +class TestImmediateRecoverySlowExportFallback(unittest.TestCase): + def setUp(self): + self.temp_dir = tempfile.TemporaryDirectory() + + def tearDown(self): + self.temp_dir.cleanup() + + def _build_pending_export(self, events: list[str], name: str) -> tuple[MagicMock, Path]: + recovery_hf_path = Path(self.temp_dir.name) / name + recovery_hf_path.mkdir(parents=True) + pending = MagicMock(spec=Future) + pending.done.return_value = False + + def finish_export(): + events.append("wait_for_pending_export") + return recovery_hf_path + + pending.result.side_effect = finish_export + return pending, recovery_hf_path + + def test_colocate_slow_export_falls_back_before_weight_update(self): + """测试内容：colocate 下慢于下一次同步点的异步 HF 必须回退到权重更新恢复。 + + 测试流程： + 1. 构造尚未完成的 recovery HF Future，模拟异步保存跨过下一次同步点。 + 2. 执行完整的 colocate save/sync 编排，等待旧 Future 收尾并清空 ready HF。 + 3. 确认 immediate recovery 被关闭且没有启动下一次异步保存。 + 4. 确认 inactive group 的恢复发生在 bind/update_weights 之前，走权重更新基线。 + """ + events: list[str] = [] + pending, _ = self._build_pending_export(events, "colocate-hf-step-1") + trainer = RLColocateTrainer.__new__(RLColocateTrainer) + trainer.logger = MagicMock() + trainer._enable_immediate_recovery = True + trainer._pending_hf_export = pending + trainer._ready_recovery_hf_path = None + trainer._sync_weights_interval = 1 + trainer._enable_evaluate = False + trainer._evaluate_step = 1 + trainer._total_train_steps = 2 + trainer._rollout_config = SimpleNamespace() + trainer._maybe_save_checkpoint = AsyncMock(side_effect=lambda _step: events.append("save_checkpoint")) + trainer._maybe_save_hf = MagicMock(side_effect=lambda _step: events.append("save_regular_hf")) + trainer.train_controller = SimpleNamespace( + offload=MagicMock(side_effect=lambda *, target: events.append(f"offload_{target}")), + update_weights=MagicMock(side_effect=lambda: events.append("update_weights")), + ) + trainer.rollout_controller = SimpleNamespace( + clear_ready_recovery_hf=_LocalRemoteMethod(lambda: events.append("clear_ready_hf")), + restart_inactive_workers=_LocalRemoteMethod(lambda: events.append("restart_inactive")), + onload_weights=_LocalRemoteMethod(lambda: events.append("onload_weights")), + onload_kvcache=_LocalRemoteMethod(lambda: events.append("onload_kvcache")), + ) + + with ( + patch("xtuner.v1.train.rl_trainer.asyncio_run", side_effect=asyncio.run), + patch("xtuner.v1.train.rl_trainer.ray.get", side_effect=_identity_ray_get), + patch( + "xtuner.v1.train.rl_trainer.bind_train_rollout", + side_effect=lambda **_kwargs: events.append("bind"), + ), + ): + synced = trainer._sync_weights_and_save(train_step=1, step_timer_dict={}) + + self.assertTrue(synced) + self.assertFalse(trainer._enable_immediate_recovery) + self.assertIsNone(trainer._pending_hf_export) + self.assertIsNone(trainer._ready_recovery_hf_path) + self.assertEqual( + events, + [ + "wait_for_pending_export", + "clear_ready_hf", + "offload_optimizer", + "save_checkpoint", + "save_regular_hf", + "restart_inactive", + "bind", + "onload_weights", + "update_weights", + "offload_model", + "onload_kvcache", + ], + ) + + def test_disaggregated_slow_export_falls_back_before_weight_update(self): + """测试内容：disaggregated 下慢异步 HF 也必须在权重同步前切回恢复基线。 + + 测试流程： + 1. 构造跨过下一次同步点的 pending recovery HF Future。 + 2. 执行 disaggregated 的 save -> restart -> bind -> update_weights 编排。 + 3. 确认同步点先等待并撤销 ready HF，然后关闭 immediate recovery。 + 4. 确认只恢复 inactive group 并直接进入本轮权重更新，不再启动 recovery HF。 + """ + events: list[str] = [] + pending, _ = self._build_pending_export(events, "disaggregated-hf-step-1") + trainer = RLDisaggregatedTrainer.__new__(RLDisaggregatedTrainer) + trainer.logger = MagicMock() + trainer._enable_immediate_recovery = True + trainer._pending_hf_export = pending + trainer._ready_recovery_hf_path = None + trainer._rollout_config = SimpleNamespace(weight_update_host="host", weight_update_port=1234) + trainer._maybe_save_checkpoint = AsyncMock(side_effect=lambda _step: events.append("save_checkpoint")) + trainer._maybe_save_hf = MagicMock(side_effect=lambda _step: events.append("save_regular_hf")) + trainer.update_weights = MagicMock(side_effect=lambda: events.append("update_weights")) + trainer.train_controller = MagicMock() + trainer.rollout_controller = SimpleNamespace( + clear_ready_recovery_hf=_LocalRemoteMethod(lambda: events.append("clear_ready_hf")), + restart_inactive_workers=SimpleNamespace( + remote=AsyncMock(side_effect=lambda: events.append("restart_inactive")) + ), + ) + + with ( + patch("xtuner.v1.train.rl_trainer.ray.get", side_effect=_identity_ray_get), + patch( + "xtuner.v1.train.rl_trainer.bind_train_rollout", + side_effect=lambda **_kwargs: events.append("bind"), + ), + ): + asyncio.run(trainer._sync_weights_and_save(model_step=1, step_timer_dict={})) + + self.assertFalse(trainer._enable_immediate_recovery) + self.assertIsNone(trainer._pending_hf_export) + self.assertIsNone(trainer._ready_recovery_hf_path) + self.assertEqual( + events, + [ + "wait_for_pending_export", + "clear_ready_hf", + "save_checkpoint", + "save_regular_hf", + "restart_inactive", + "bind", + "update_weights", + ], + ) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/rl/test_rl_trainer_checkpoint.py b/tests/rl/test_rl_trainer_checkpoint.py index cb2977b6c..afccb8b52 100644 --- a/tests/rl/test_rl_trainer_checkpoint.py +++ b/tests/rl/test_rl_trainer_checkpoint.py @@ -101,7 +101,7 @@ def __init__(self): self.continue_generation = _RemoteMethod(async_result=True) self.offload = _RemoteMethod(return_value="rollout_offloaded") self.check_and_shutdown_inactive_workers = _RemoteMethod(return_value="rollout_inactive_workers_shutdown") - self.restart_inactive_workers = _RemoteMethod(return_value="rollout_restarted") + self.restart_inactive_workers = _RemoteMethod(return_value=_AwaitableValue(None)) self.onload_weights = _RemoteMethod(return_value="weights_loaded") self.onload_kvcache = _RemoteMethod(return_value="kvcache_loaded") self.get_weight_update_targets = _RemoteMethod(return_value=()) diff --git a/tests/rl/test_rollout_logic.py b/tests/rl/test_rollout_logic.py index 5b0aa5ad8..df3cd5372 100644 --- a/tests/rl/test_rollout_logic.py +++ b/tests/rl/test_rollout_logic.py @@ -1253,14 +1253,11 @@ def test_wait_until_next_check_waits_for_resume_when_paused_during_interval(self manager._check_interval = 0.01 manager._pause_event = threading.Event() - class _FakeStopEvent: + class _FakeWakeupEvent: def __init__(self): self.wait_calls = [] self._paused_once = False - def is_set(self): - return False - def wait(self, timeout=None): self.wait_calls.append(timeout) if timeout == manager._check_interval and not self._paused_once: @@ -1270,11 +1267,14 @@ def wait(self, timeout=None): manager._pause_event.clear() return False - stop_event = _FakeStopEvent() - manager._stop_event = stop_event + def clear(self): + return None + + wakeup_event = _FakeWakeupEvent() + manager._health_loop_wakeup_event = wakeup_event self.assertTrue(manager._wait_until_next_check()) - self.assertEqual(stop_event.wait_calls, [manager._check_interval, 0.5, manager._check_interval]) + self.assertEqual(wakeup_event.wait_calls, [manager._check_interval, 0.5]) def test_shutdown_barrier_keeps_failed_shutdown_group_inactive(self): actor = SimpleNamespace(check_health=_FakeAsyncRemoteMethod(True)) @@ -1349,7 +1349,8 @@ def test_restart_barrier_cleans_claimed_groups_when_stopping(self): ) manager, registry = self._build_manager({0: worker_info}) - def stop_after_restart(groups): + def stop_after_restart(groups, *, ready_recovery_hf): + self.assertIsNone(ready_recovery_hf) manager._stop_event.set() return {group.ranks: False for group in groups} @@ -1395,12 +1396,10 @@ def test_restart_worker_group_uses_reinit(self): session_url="http://session-0", ) actor = SimpleNamespace( - set_skip_load_weights=_FakeAsyncRemoteMethod(None), init=_FakeAsyncRemoteMethod(init_result), reinit=_FakeAsyncRemoteMethod(init_result), check_health=_FakeAsyncRemoteMethod(True), offload=_FakeAsyncRemoteMethod(None), - restore_skip_load_weights=_FakeAsyncRemoteMethod(None), ) worker_info = WorkerSnapshot( rank=0, @@ -1419,16 +1418,22 @@ def fake_ray_get(refs, timeout=None): with ( patch.object(manager, "_shutdown_worker_group", return_value=True), patch("xtuner.v1.rl.rollout.health_manager.ray.get", side_effect=fake_ray_get), + patch("xtuner.v1.rl.rollout.health_manager.logger.info") as log_info, ): - result = manager._restart_worker_group(group) + result = manager._restart_worker_group(group, ready_recovery_hf=None) self.assertTrue(result) - self.assertEqual(actor.set_skip_load_weights.calls, [(True,)]) - self.assertEqual(actor.reinit.calls, [()]) + self.assertEqual(actor.reinit.calls, [((), {"skip_load_weights": True})]) self.assertEqual(actor.init.calls, []) self.assertEqual(actor.check_health.calls, [()]) self.assertEqual(actor.offload.calls, [()]) - self.assertEqual(actor.restore_skip_load_weights.calls, [()]) + self.assertTrue( + any( + "without ready recovery HF" in call.args[0] + and "recovery will complete through the next weight update" in call.args[0] + for call in log_info.call_args_list + ) + ) def test_recovered_listener_runs_outside_lifecycle_operation_lock(self): actor = SimpleNamespace(check_health=_FakeAsyncRemoteMethod(True)) diff --git a/xtuner/v1/rl/rollout/controller.py b/xtuner/v1/rl/rollout/controller.py index cee7d7c6e..949b2edd7 100644 --- a/xtuner/v1/rl/rollout/controller.py +++ b/xtuner/v1/rl/rollout/controller.py @@ -161,6 +161,34 @@ 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) + def set_ready_recovery_hf( + self, + *, + model_path: str, + tokenizer_path: str | None = None, + ) -> None: + self.health_manager.set_ready_recovery_hf( + model_path=model_path, + tokenizer_path=tokenizer_path, + ) + + def clear_ready_recovery_hf(self) -> None: + self.health_manager.clear_ready_recovery_hf() + + 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 continue_generation(self): self._broadcast_to_active_workers("continue_generation") self.health_manager.resume() diff --git a/xtuner/v1/rl/rollout/health_manager.py b/xtuner/v1/rl/rollout/health_manager.py index 2026cd655..aaede020e 100644 --- a/xtuner/v1/rl/rollout/health_manager.py +++ b/xtuner/v1/rl/rollout/health_manager.py @@ -101,6 +101,12 @@ def mark_failed_ranks(self, worker_health_results: dict[int, bool]) -> set[int]: return failed_ranks +@dataclass(frozen=True) +class _ReadyRecoveryHF: + model_path: str + tokenizer_path: str | None = None + + class RolloutHealthManager: """Own worker health state and recovery after controller startup. @@ -124,13 +130,33 @@ def __init__( self._stop_event = threading.Event() self._pause_event = threading.Event() self._pause_event.set() + self._health_loop_wakeup_event = threading.Event() self._thread: threading.Thread | None = None self._lifecycle_operation_lock = threading.Lock() self._worker_health_failure_tracker = _WorkerHealthFailureTracker(threshold=self._check_failure_threshold) + self._ready_recovery_hf: _ReadyRecoveryHF | None = None # ------------------------------------------------------------------ # Public lifecycle # ------------------------------------------------------------------ + def set_ready_recovery_hf( + self, + *, + model_path: str, + tokenizer_path: str | None = None, + ) -> None: + self._ready_recovery_hf = _ReadyRecoveryHF( + model_path=model_path, + tokenizer_path=tokenizer_path, + ) + logger.info(f"Ready rollout recovery HF updated: model_path={model_path}, tokenizer_path={tokenizer_path}.") + # The background health loop owns automatic recovery. Wake it so an + # already-inactive group does not wait for the next periodic check. + self._health_loop_wakeup_event.set() + + def clear_ready_recovery_hf(self) -> None: + self._ready_recovery_hf = None + logger.info("Ready rollout recovery HF cleared.") def start(self) -> None: health_thread_alive = self._thread is not None and self._thread.is_alive() @@ -138,6 +164,7 @@ def start(self) -> None: return self._stop_event.clear() + self._health_loop_wakeup_event.clear() self._pause_event.set() if not self._periodic_health_checks_enabled: logger.info("Rollout worker periodic health check is disabled.") @@ -149,6 +176,7 @@ def start(self) -> None: def stop(self) -> None: self._stop_event.set() + self._health_loop_wakeup_event.set() self._pause_event.clear() thread = self._thread if not thread: @@ -169,10 +197,12 @@ def stop(self) -> None: def pause(self) -> None: self._pause_event.set() + self._health_loop_wakeup_event.set() logger.info("RolloutHealthManager paused.") def resume(self) -> None: self._pause_event.clear() + self._health_loop_wakeup_event.set() logger.info("RolloutHealthManager resumed.") # ------------------------------------------------------------------ @@ -193,13 +223,12 @@ def run_once(self) -> None: try: worker_health_results = self._check_active_workers_health() failed_ranks = self._worker_health_failure_tracker.update_failed_ranks(worker_health_results) - if not failed_ranks: - return - try: - self._checkpoint_not_stopping() - except _HealthManagerStopping: - return - failed_groups = self._registry.mark_unhealthy_ranks(failed_ranks) + if failed_ranks: + try: + self._checkpoint_not_stopping() + except _HealthManagerStopping: + return + failed_groups = self._registry.mark_unhealthy_ranks(failed_ranks) finally: self._lifecycle_operation_lock.release() @@ -210,23 +239,42 @@ def run_once(self) -> None: event_name="inactive", notify_listener=lambda listener, group: listener.on_worker_group_inactive(group), ) + # TODO：Recovery runs synchronously on the health-check thread, so the next + # periodic health check waits until this restart finishes. Move restart to another thread. + self._restart_inactive_workers(require_ready_recovery_hf=True) def restart_inactive_workers(self) -> None: """Synchronously restart inactive groups before the next sync-step weight update.""" + self._restart_inactive_workers(require_ready_recovery_hf=False) + + def _restart_inactive_workers( + self, + *, + require_ready_recovery_hf: bool, + ) -> None: recovered_groups: list[WorkerGroup] = [] groups_to_recover: tuple[WorkerGroup, ...] = () try: with self._paused_lifecycle_operation(): + ready_recovery_hf = self._ready_recovery_hf + if require_ready_recovery_hf and ready_recovery_hf is None: + logger.info( + "Ready recovery HF is unavailable; deferring rollout worker restart until " + "a recovery HF is published or the next weight update begins." + ) + return groups_to_recover = self._registry.claim_inactive_groups_for_recovery() if groups_to_recover: - recovered_groups = self._restart_claimed_recovery_groups(groups_to_recover) + recovered_groups = self._restart_claimed_recovery_groups( + groups_to_recover, + ready_recovery_hf=ready_recovery_hf, + ) except _HealthManagerStopping: return if not groups_to_recover: - logger.info("No failed rollout workers detected during recovery.") return self._notify_worker_lifecycle_listeners( @@ -285,15 +333,20 @@ def _run_loop(self) -> None: def _wait_until_next_check(self) -> bool: while True: - while self._pause_event.is_set() and not self._stop_event.is_set(): - self._stop_event.wait(timeout=0.5) - if self._stop_event.is_set(): return False - if self._stop_event.wait(self._check_interval): - return False + if self._pause_event.is_set(): + self._health_loop_wakeup_event.wait(timeout=0.5) + self._health_loop_wakeup_event.clear() + if self._stop_event.is_set(): + return False + if not self._pause_event.is_set(): + return True + continue + self._health_loop_wakeup_event.wait(timeout=self._check_interval) + self._health_loop_wakeup_event.clear() if not self._pause_event.is_set() and not self._stop_event.is_set(): return True @@ -309,12 +362,12 @@ def _checkpoint_not_stopping(self) -> None: def _background_health_checks_paused(self): was_paused = self._pause_event.is_set() if not was_paused: - self.pause() + self._pause_event.set() try: yield finally: if not was_paused: - self.resume() + self._pause_event.clear() @contextmanager def _paused_lifecycle_operation(self): @@ -387,11 +440,19 @@ def _mark_unhealthy_worker_groups_inactive(self, worker_health_results: dict[int # Worker group recovery state # ------------------------------------------------------------------ - def _restart_claimed_recovery_groups(self, groups: tuple[WorkerGroup, ...]) -> list[WorkerGroup]: + def _restart_claimed_recovery_groups( + self, + groups: tuple[WorkerGroup, ...], + *, + ready_recovery_hf: _ReadyRecoveryHF | None, + ) -> list[WorkerGroup]: groups_needing_cleanup = {group.ranks: group for group in groups} try: - group_recovery_results = self._restart_worker_groups(groups) + group_recovery_results = self._restart_worker_groups( + groups, + ready_recovery_hf=ready_recovery_hf, + ) self._checkpoint_not_stopping() recovered_groups: list[WorkerGroup] = [] @@ -434,6 +495,8 @@ def _cleanup_unfinalized_recovery_groups(self, groups: tuple[WorkerGroup, ...]) def _restart_worker_groups( self, groups_to_recover: tuple[WorkerGroup, ...], + *, + ready_recovery_hf: _ReadyRecoveryHF | None, ) -> dict[tuple[int, ...], bool]: logger.info( f"Restarting rollout worker groups in parallel: " @@ -450,6 +513,7 @@ def _restart_worker_groups( pool.submit( self._restart_worker_group, group, + ready_recovery_hf=ready_recovery_hf, ): group for group in groups_to_recover } @@ -465,56 +529,75 @@ def _restart_worker_groups( def _restart_worker_group( self, group: WorkerGroup, + *, + ready_recovery_hf: _ReadyRecoveryHF | None, ) -> bool: - """Shutdown, restart with empty-init, and health-check one complete - worker group.""" + """Shutdown, restart, and health-check one complete worker group.""" if not group.workers or len(group.workers) != len(group.ranks): logger.error(f"Cannot restart incomplete rollout worker group: ranks={group.ranks}.") return False restart_cleanup_needed = False + recovery_started_at = time.perf_counter() try: self._checkpoint_not_stopping() + if not self._shutdown_worker_group(group): return False restart_cleanup_needed = True self._checkpoint_not_stopping() - with self._skip_load_weights_during_restart(group): - self._checkpoint_not_stopping() - ray.get( - [ - # reinit() reuses the server launch spec bound during - # controller startup. - worker.actor.reinit.remote() # type: ignore[attr-defined] - for worker in group.workers - ], - timeout=ROLLOUT_RAY_GET_TIMEOUT, + if ready_recovery_hf is None: + reinit_kwargs: dict[str, object] = {"skip_load_weights": True} + logger.info( + "Restarting rollout worker group without ready recovery HF: " + f"ranks={group.ranks}, skip_load_weights=True; " + "recovery will complete through the next weight update." + ) + else: + reinit_kwargs = { + "model_path": ready_recovery_hf.model_path, + "tokenizer_path": ready_recovery_hf.tokenizer_path, + "skip_load_weights": False, + } + logger.info( + "Restarting rollout worker group by loading ready recovery HF: " + f"ranks={group.ranks}, model_path={ready_recovery_hf.model_path}, " + f"tokenizer_path={ready_recovery_hf.tokenizer_path}." ) - self._checkpoint_not_stopping() - health_results = self._check_workers_health(group.workers) - unhealthy_ranks = [ - worker.rank for worker in group.workers if not health_results.get(worker.rank, False) - ] - if unhealthy_ranks: - logger.error( - f"Restarted rollout worker group ranks={group.ranks} has unhealthy ranks={unhealthy_ranks}." - ) - self._shutdown_worker_group(group, wait_server_down=False) - return False + ray.get( + [ + worker.actor.reinit.remote(**reinit_kwargs) # type: ignore[attr-defined] + for worker in group.workers + ], + timeout=ROLLOUT_RAY_GET_TIMEOUT, + ) + + self._checkpoint_not_stopping() + health_results = self._check_workers_health(group.workers) + unhealthy_ranks = [worker.rank for worker in group.workers if not health_results.get(worker.rank, False)] + if unhealthy_ranks: + logger.error( + f"Restarted rollout worker group ranks={group.ranks} has unhealthy ranks={unhealthy_ranks}." + ) + self._shutdown_worker_group(group, wait_server_down=False) + return False + if ready_recovery_hf is None: self._checkpoint_not_stopping() - # Newly restarted workers should return to the same offloaded/sleep - # baseline as the other colocated rollout workers before the sync - # path wakes weights/KV back up. + # Weight-update recovery returns to the offloaded baseline + # before the sync path wakes weights and KV cache back up. ray.get( [worker.actor.offload.remote() for worker in group.workers], # type: ignore[attr-defined] timeout=ROLLOUT_RAY_GET_TIMEOUT, ) - logger.info(f"Successfully restarted rollout worker group ranks={group.ranks}.") + logger.info( + f"Successfully restarted rollout worker group ranks={group.ranks} " + f"in {time.perf_counter() - recovery_started_at:.3f}s." + ) return True except _HealthManagerStopping: if restart_cleanup_needed: @@ -526,31 +609,6 @@ def _restart_worker_group( self._shutdown_worker_group(group, wait_server_down=False) return False - @contextmanager - def _skip_load_weights_during_restart(self, group: WorkerGroup): - try: - ray.get( - [ - worker.actor.set_skip_load_weights.remote(True) # type: ignore[attr-defined] - for worker in group.workers - ], - timeout=ROLLOUT_RAY_GET_TIMEOUT, - ) - yield - finally: - try: - ray.get( - [ - worker.actor.restore_skip_load_weights.remote() # type: ignore[attr-defined] - for worker in group.workers - ], - timeout=ROLLOUT_RAY_GET_TIMEOUT, - ) - except Exception: - logger.exception( - f"Failed to restore rollout worker skip_load_weights after restart: group_ranks={group.ranks}." - ) - def _shutdown_worker_group( self, group: WorkerGroup, diff --git a/xtuner/v1/rl/rollout/worker.py b/xtuner/v1/rl/rollout/worker.py index f6118fbdf..14dfd30ac 100644 --- a/xtuner/v1/rl/rollout/worker.py +++ b/xtuner/v1/rl/rollout/worker.py @@ -536,7 +536,6 @@ def __init__( Defaults to "GPU". """ self.config = config - self._default_skip_load_weights = config.skip_load_weights self.rank = rank self.master_addr = master_addr # ray master self.master_port = master_port @@ -594,9 +593,25 @@ def init(self, server_launch_spec: ServerLaunchSpec) -> RolloutWorkerInitResult: self._bind_server_launch_spec(server_launch_spec) return self._init_server() - def reinit(self) -> RolloutWorkerInitResult: + def reinit( + self, + *, + model_path: str | Path | None = None, + tokenizer_path: str | Path | None = None, + skip_load_weights: bool | None = None, + ) -> RolloutWorkerInitResult: """Reinitialize the rollout server using the previously bound launch spec.""" + config_updates: dict[str, object] = {} + if model_path is not None: + config_updates["model_path"] = str(model_path) + if tokenizer_path is not None: + config_updates["tokenizer_path"] = str(tokenizer_path) + if skip_load_weights is not None: + config_updates["skip_load_weights"] = skip_load_weights + + if config_updates: + self.config = self.config.model_copy(update=config_updates) return self._init_server() def _init_server(self) -> RolloutWorkerInitResult: @@ -616,12 +631,6 @@ def _init_server(self) -> RolloutWorkerInitResult: session_url=self.session_server_url, ) - def set_skip_load_weights(self, skip_load_weights: bool) -> None: - self.config = self.config.model_copy(update={"skip_load_weights": skip_load_weights}) - - def restore_skip_load_weights(self) -> None: - self.config = self.config.model_copy(update={"skip_load_weights": self._default_skip_load_weights}) - def init_dist_port(self) -> tuple[int, str]: """Initialize distributed communication ports. @@ -682,6 +691,46 @@ def shutdown(self, *, stop_session_server: bool = False): self.logger.debug(f"Worker {self.rank} server process and its children terminated.") return + def inject_backend_crash_for_test(self) -> bool: + """Force-stop the backend server for the immediate-recovery test.""" + if os.environ.get("XTUNER_TEST_IMMEDIATE_RECOVERY", "0") != "1": + raise RuntimeError("Rollout test fault injection requires XTUNER_TEST_IMMEDIATE_RECOVERY=1.") + self.logger.warning( + f"[ImmediateRecoveryExperiment] crashing_backend_server rank={self.rank} url={self.server_url}" + ) + + if self.server_task is not None: + server_task = self.server_task + ray.cancel(server_task, force=True, recursive=True) + try: + ray.get(server_task, timeout=60) + except ray.exceptions.GetTimeoutError: + self.logger.warning(f"Worker {self.rank} server task did not stop within crash timeout.") + raise + except Exception as e: + self.logger.debug(f"Worker {self.rank} server task stopped after injected crash: {e}") + self.server_task = None + return True + + if self.server_process is not None: + import psutil + + try: + parent = psutil.Process(self.server_process.pid) + except psutil.NoSuchProcess: + self.server_process = None + return True + children = parent.children(recursive=True) + for child in children: + child.kill() + parent.kill() + parent.wait(timeout=5) + self.server_process = None + self.logger.debug(f"Worker {self.rank} server process and its children killed.") + return True + + return False + def _start_session_server(self) -> None: """Start the per-worker SessionServer proxy.""" assert self.server_launch_spec is not None diff --git a/xtuner/v1/rl/trainer/controller.py b/xtuner/v1/rl/trainer/controller.py index b87da88ce..9fe4a89d8 100644 --- a/xtuner/v1/rl/trainer/controller.py +++ b/xtuner/v1/rl/trainer/controller.py @@ -328,6 +328,22 @@ def save_hf(self, hf_dir: str, save_dtype: torch.dtype = torch.bfloat16): ray.get(handles, timeout=TRAIN_RAY_GET_TIMEOUT) return + def start_hf_export( + self, + hf_dir: str, + save_dtype: torch.dtype = torch.bfloat16, + ) -> None: + handles = [worker.start_hf_export.remote(hf_dir, save_dtype) for worker in self.workers] + ray.get(handles, timeout=TRAIN_RAY_GET_TIMEOUT) + + def is_hf_export_done(self) -> bool: + handles = [worker.is_hf_export_done.remote() for worker in self.workers] + return all(ray.get(handles, timeout=TRAIN_RAY_GET_TIMEOUT)) + + def wait_hf_export(self) -> str: + handles = [worker.wait_hf_export.remote() for worker in self.workers] + return ray.get(handles, timeout=TRAIN_RAY_GET_TIMEOUT)[0] + def resume(self, load_checkpoint_cfg: LoadCheckpointConfig): """Resume the training workers from the checkpoint.""" handles = [worker.resume.remote(load_checkpoint_cfg) for worker in self.workers] # type: ignore diff --git a/xtuner/v1/rl/trainer/worker.py b/xtuner/v1/rl/trainer/worker.py index f7e0bc487..c72a723fa 100644 --- a/xtuner/v1/rl/trainer/worker.py +++ b/xtuner/v1/rl/trainer/worker.py @@ -3,6 +3,7 @@ import math import os import time +from concurrent.futures import Future from contextlib import contextmanager from pathlib import Path from typing import ( @@ -242,6 +243,7 @@ def __init__( if not worker_cfg.fsdp_cfg.torch_compile: worker_cfg.model_cfg.compile_cfg = False self._engine = self._build_engine(worker_cfg) + self._pending_hf_export: Future[Path] | None = None self._has_ref = False if worker_cfg.loss_cfg.use_kl_loss: @@ -942,6 +944,28 @@ def _reduce_number_across_rank(self, rank_number: int) -> int: def save_hf(self, hf_dir: str, save_dtype: torch.dtype = torch.bfloat16): self._engine.save_hf(hf_dir, save_dtype) + @ray_method + def start_hf_export( + self, + hf_dir: str, + save_dtype: torch.dtype = torch.bfloat16, + ) -> None: + self._pending_hf_export = self._engine.async_save_hf(hf_dir, save_dtype) + + @ray_method + def is_hf_export_done(self) -> bool: + pending = cast(Future[Path], self._pending_hf_export) + return pending.done() + + @ray_method + def wait_hf_export(self) -> str: + pending = cast(Future[Path], self._pending_hf_export) + try: + finalized_path = pending.result() + finally: + self._pending_hf_export = None + return str(finalized_path) + @ray_method def get_data_replicate_size(self) -> int: """Get the data replicate size for the training worker.""" diff --git a/xtuner/v1/train/rl_trainer.py b/xtuner/v1/train/rl_trainer.py index 9e810ad16..b4dfcdfcf 100644 --- a/xtuner/v1/train/rl_trainer.py +++ b/xtuner/v1/train/rl_trainer.py @@ -4,6 +4,7 @@ import random import re import time +from concurrent.futures import Future, ThreadPoolExecutor from dataclasses import asdict, dataclass from pathlib import Path from shutil import rmtree @@ -63,6 +64,7 @@ # TODO: Move DEVICE to `xtuner.utils.device` PG_READY_TIMEOUT = 30 RL_TRAINER_RAY_GET_TIMEOUT = 3600 +HF_EXPORT_POLL_INTERVAL_S = 1.0 DEVICE = get_device() DEVICE_MODULE = get_torch_device_module() @@ -354,6 +356,7 @@ class BaseRLTrainerConfig(BaseModel): checkpoint_maxkeep: int | None = -1 hf_interval: int | None = -1 hf_max_keep: int | None = -1 + enable_immediate_recovery: bool = False checkpoint_no_save_optimizer: bool = False checkpoint_no_save_replay_buffer: bool = False log_dir: Path | str | None = None @@ -439,6 +442,9 @@ class RLColocateTrainerConfig(BaseRLTrainerConfig): Defaults to -1. hf_max_keep (int | None): Maximum number of Hugging Face checkpoints to keep. Defaults to -1. + enable_immediate_recovery (bool): Whether rollout worker recovery may + use fresh ready HF exports before the next weight update. Defaults + to False. checkpoint_no_save_optimizer (bool): Whether to skip optimizer states when saving checkpoints. Defaults to False. checkpoint_no_save_replay_buffer (bool): Whether to skip replay buffer @@ -527,6 +533,9 @@ class RLDisaggregatedTrainerConfig(BaseRLTrainerConfig): Defaults to -1. hf_max_keep (int | None): Maximum number of Hugging Face checkpoints to keep. Defaults to -1. + enable_immediate_recovery (bool): Whether rollout worker recovery may + use fresh ready HF exports before the next weight update. Defaults + to False. checkpoint_no_save_optimizer (bool): Whether to skip optimizer states when saving checkpoints. Defaults to False. checkpoint_no_save_replay_buffer (bool): Whether to skip replay buffer @@ -622,6 +631,14 @@ def _init_load_source(self, cfg: BaseRLTrainerConfig) -> None: def _init_save_config(self, cfg: BaseRLTrainerConfig) -> None: self._hf_max_keep = cfg.hf_max_keep self._hf_interval = cfg.hf_interval + self._enable_immediate_recovery = cfg.enable_immediate_recovery + self._hf_export_executor = ( + ThreadPoolExecutor(max_workers=1, thread_name_prefix="rl-hf-export") + if self._enable_immediate_recovery + else None + ) + self._pending_hf_export: Future[Path | None] | None = None + self._ready_recovery_hf_path: Path | None = None self._checkpoint_interval = cfg.checkpoint_interval self._checkpoint_maxkeep = cfg.checkpoint_maxkeep @@ -858,6 +875,96 @@ def _maybe_save_hf(self, cur_step: int): if isinstance(self.tokenizer, (PreTrainedTokenizer, PreTrainedTokenizerFast)): self.tokenizer.save_pretrained(str(save_hf_path)) + def _maybe_save_recovery_hf(self, cur_step: int) -> None: + if not self._enable_immediate_recovery: + return + + reuse_regular_hf = ( + self._hf_interval is not None + and self._hf_interval != -1 + and (cur_step % self._hf_interval == 0 or cur_step == self._total_train_steps) + ) + tokenizer_path = self._rollout_config.tokenizer_path or self._rollout_config.model_path + previous_ready_hf_path = self._ready_recovery_hf_path + ray.get( + self.rollout_controller.clear_ready_recovery_hf.remote(), + timeout=RL_TRAINER_RAY_GET_TIMEOUT, + ) + self._ready_recovery_hf_path = None + if ( + previous_ready_hf_path is not None + and str(previous_ready_hf_path) not in self._meta.latest_exp.hf_checkpoint_list + ): + rmtree(previous_ready_hf_path, ignore_errors=True) + + save_hf_path = self.exp_dir / self._HF_DIR / f"hf-step-{cur_step}" + if reuse_regular_hf: + try: + ray.get( + self.rollout_controller.set_ready_recovery_hf.remote( + model_path=str(save_hf_path), + tokenizer_path=str(tokenizer_path), + ), + timeout=RL_TRAINER_RAY_GET_TIMEOUT, + ) + except Exception: + self.logger.exception(f"Failed to publish recovery HF: path={save_hf_path}.") + return + + self._ready_recovery_hf_path = save_hf_path + return + + save_hf_path.mkdir(parents=True, exist_ok=True) + self.logger.info(f"Starting async recovery HF export: path={save_hf_path}.") + self.train_controller.start_hf_export(str(save_hf_path)) + executor = cast(ThreadPoolExecutor, self._hf_export_executor) + + def wait_and_publish_recovery_hf() -> Path | None: + try: + while not self.train_controller.is_hf_export_done(): + time.sleep(HF_EXPORT_POLL_INTERVAL_S) + finalized_hf_path = Path(self.train_controller.wait_hf_export()) + ray.get( + self.rollout_controller.set_ready_recovery_hf.remote( + model_path=str(finalized_hf_path), + tokenizer_path=str(tokenizer_path), + ), + timeout=RL_TRAINER_RAY_GET_TIMEOUT, + ) + self.logger.info(f"Async recovery HF export completed and published: path={finalized_hf_path}.") + except Exception: + self.logger.exception(f"Async recovery HF export failed: path={save_hf_path}.") + return None + + return finalized_hf_path + + self._pending_hf_export = executor.submit(wait_and_publish_recovery_hf) + + def _wait_or_disable_immediate_recovery(self) -> None: + pending = self._pending_hf_export + if pending is None: + return + + disable_immediate_recovery = not pending.done() + if disable_immediate_recovery: + self._enable_immediate_recovery = False + self.logger.warning( + "Disable immediate recovery because the previous recovery HF " + "export did not finish before the next weight sync." + ) + + finalized_hf_path = pending.result() + self._pending_hf_export = None + if not disable_immediate_recovery: + self._ready_recovery_hf_path = finalized_hf_path + return + + ray.get( + self.rollout_controller.clear_ready_recovery_hf.remote(), + timeout=RL_TRAINER_RAY_GET_TIMEOUT, + ) + self._ready_recovery_hf_path = None + async def _run_initial_evaluate(self) -> None: try: eval_produce_result = await self.eval_agent_loop_manager.produce_batch( @@ -1639,6 +1746,8 @@ def fit(self): self._fit() finally: self._exp_tracker.close() + if self._hf_export_executor is not None: + self._hf_export_executor.shutdown(wait=True) close_trace() def _fit(self): @@ -1748,6 +1857,8 @@ def _get_colocate_rollout_model_step(self, train_step: int) -> int: def _sync_weights_and_save(self, train_step: int, step_timer_dict: dict) -> bool: """保存后切回共卡 rollout 资源。""" + self._wait_or_disable_immediate_recovery() + should_sync_weights = train_step % self._sync_weights_interval == 0 will_evaluate = self._enable_evaluate and train_step % self._evaluate_step == 0 needs_rollout_ready = train_step < self._total_train_steps or will_evaluate @@ -1781,6 +1892,7 @@ def _sync_weights_and_save(self, train_step: int, step_timer_dict: dict) -> bool ) self.train_controller.update_weights() self.logger.info("Rollout workers update weights successfully in colocate mode") + self._maybe_save_recovery_hf(train_step) self.train_controller.offload(target="model") else: self.train_controller.offload(target="model") @@ -1870,6 +1982,8 @@ def fit(self): return asyncio_run(self._fit()) finally: self._exp_tracker.close() + if self._hf_export_executor is not None: + self._hf_export_executor.shutdown(wait=True) close_trace() async def _get_batch_or_raise_producer_failure( @@ -2007,12 +2121,16 @@ async def _fit(self): async def _sync_weights_and_save(self, model_step: int, step_timer_dict: dict): # producer 已暂停；保持 save -> bind -> update 顺序。 + self._wait_or_disable_immediate_recovery() + with timer("save_ckpt", step_timer_dict): await self._maybe_save_checkpoint(model_step) self._maybe_save_hf(model_step) - # TODO: 非共卡需要额外加健康检查恢复worker的逻辑，共卡是在训练之前恢复，但是非共卡不需要在训练之前恢复,挂掉就恢复或者更新权重前恢复，需要评估一下哪种方式更合理。 with timer("sync_weight", step_timer_dict): + # 非共卡在权重更新前恢复 inactive workers；如果没有 ready + # recovery HF，HealthManager 会用空权重启动并等待本次 update。 + await self.rollout_controller.restart_inactive_workers.remote() # type: ignore[attr-defined] bind_train_rollout( train_controller=self.train_controller, rollout_controller=self.rollout_controller, @@ -2022,6 +2140,7 @@ async def _sync_weights_and_save(self, model_step: int, step_timer_dict: dict): weight_update_port=self._rollout_config.weight_update_port, ) self.update_weights() + self._maybe_save_recovery_hf(model_step) def update_weights(self): # rollout 恢复由 AgentLoopManager 控制。