fix deactivate rollout worker

Author: YanhuiDua

xtuner/v1/ray/rollout/controller.py

@@ -4,6 +4,7 @@
 import threading
 import time
 from collections import OrderedDict
+from dataclasses import dataclass
 from itertools import cycle
 from typing import Any, Dict, List, Optional, Union
 from uuid import uuid4
@@ -22,6 +23,15 @@
 from .worker import RolloutWorker
 
 
+@dataclass
+class WorkerInfo:
+    """A data class to hold all state information for a single worker."""
+
+    actor: RolloutWorker
+    is_active: bool = True
+    failure_count: int = 0
+
+
 class SessionRouter:
     def __init__(
         self,
@@ -38,6 +48,7 @@ def __init__(
         self._map: OrderedDict[int, tuple[Any, float]] = OrderedDict()
         self._worker_cycler = cycle(self._workers)
         self._lock = asyncio.Lock()
+        self.logger = get_logger()
 
     def _now(self) -> float:
         return time.time()
@@ -60,6 +71,11 @@ def _evict_lru_to_capacity(self):
         while len(self._map) > self._max_sessions:
             self._map.popitem(last=False)
 
+    def update_active_workers(self, worker_status: Dict[Any, bool]):
+        self._workers = list(worker_status.items())
+        self.logger.debug(f"SessionRouter update active workers: {self._workers}")
+        self._worker_cycler = cycle(self._workers)
+
     async def get_worker(self, session_id: int) -> Any:
         async with self._lock:
             self._evict_expired()
@@ -103,21 +119,18 @@ def __init__(
         )
         self.logger = get_logger(log_dir=infer_config.worker_log_dir, tag="RolloutController")
         self.num_workers = 0
+        self.workers_info: Dict[str, WorkerInfo] = {}  # url -> WorkerInfo
         self.worker_server_urls: List[str] = []
         self.active_rollout_workers: List[RolloutWorker] = []
-        self.active_workers_to_status: Dict[RolloutWorker, bool] = {}
-        self.active_url_to_workers: Dict[str, RolloutWorker] = {}
-        self.url_failed_counts: Dict[str, int] = {}
         self.tokenizer = AutoTokenizer.from_pretrained(infer_config.tokenizer_path, trust_remote_code=True)
         self.workers, self.rank_bundle_idx_list = AutoAcceleratorWorkers.from_placement_group(
             self._get_worker_cls(), infer_config, placement_group
         )
-        self.engine_mesh_list, self.server_url_dict = self.init_workers()
+        self.engine_mesh_list, self.worker_server_urls_map = self.init_workers()
         self.start_api_server()
         # todo(@duanyanhui): add router to replace native round robin
-        self.router = SessionRouter(self.active_workers_to_status)
+        self.router = SessionRouter(self._get_worker_status_for_router())
         self.sample_params = SampleParams().dict()
-        # note: 目前默认使用return_token_ids和return_logprob，并且不使用流式
         self.extra_params = dict(
             RolloutExtraParams(
                 stream=False,
@@ -132,6 +145,10 @@ def __init__(
         )
         self.print_params_flag = True
 
+    def _get_worker_status_for_router(self) -> Dict[RolloutWorker, bool]:
+        """Helper to generate the status dict required by the SessionRouter."""
+        return {info.actor: info.is_active for info in self.workers_info.values()}
+
     def _get_worker_cls(self):
         if os.environ.get("XTUNER_USE_LMDEPLOY") == "1":
             from .lmdeploy import LMDeployWorker
@@ -161,7 +178,7 @@ def _is_port_in_use(self, host: str, port: int) -> bool:
             except OSError:
                 return True
 
-    def _update_active_workers_and_urls(self):
+    def _update_active_workers_and_urls_map(self, active_rollout_workers, worker_server_urls_map):
         """Update the list of active rollout workers and their server URLs.
 
         When the inference engine is launched across nodes (rollout_cross_node_comm=True), only the worker with
@@ -170,13 +187,12 @@ def _update_active_workers_and_urls(self):
         workers and their corresponding URLs.
         """
         if self.config.rollout_cross_node_comm or self.num_gpus_per_engine < self.config.gpus_per_node:
-            return
+            return active_rollout_workers, worker_server_urls_map
         else:
             active_worker_interval = self.num_gpus_per_engine // self.config.gpus_per_node
-            self.active_rollout_workers = self.active_rollout_workers[::active_worker_interval]
-            active_rank = list(self.worker_server_urls_map.keys())[::active_worker_interval]
-            active_worker_server_urls = list(self.worker_server_urls_map.values())[::active_worker_interval]
-            self.worker_server_urls_map = dict(zip(active_rank, active_worker_server_urls))
+            active_rank = list(worker_server_urls_map.keys())[::active_worker_interval]
+            active_worker_server_urls = list(worker_server_urls_map.values())[::active_worker_interval]
+            return active_rollout_workers[::active_worker_interval], dict(zip(active_rank, active_worker_server_urls))
 
     def get_rollout_info(self):
         """Get information about the current rollout setup.
@@ -185,10 +201,12 @@ def get_rollout_info(self):
             dict: A dictionary containing the engine mesh list, server URL
                 dictionary, and the rollout configuration.
         """
+        worker_server_urls_status = {url: info.is_active for url, info in self.workers_info.items()}
         return dict(
             engine_mesh_list=self.engine_mesh_list,
-            server_url_dict=self.server_url_dict,
+            server_url_dict=self.worker_server_urls_map,
             rollout_config=self.config,
+            worker_server_urls_status=worker_server_urls_status,
         )
 
     def init_workers(self):
@@ -210,13 +228,13 @@ def init_workers(self):
         """
         active_servers_count, nodes_per_engine = self._get_active_servers_count(self.config, len(self.workers))
         interval = len(self.workers) // active_servers_count
-        self.active_rollout_workers = self.workers[::interval]
-        self.num_workers = len(self.active_rollout_workers)
+        active_rollout_workers = self.workers[::interval]
+        self.num_workers = len(active_rollout_workers)
 
         set_bundle_idxs_objectref = []
         engine_mesh_list = []
         activate_worker_idx = 0
-        for active_worker in self.active_rollout_workers:
+        for active_worker in active_rollout_workers:
             head_rank, _ = self.rank_bundle_idx_list[activate_worker_idx]
             engine_workers_meta = self.rank_bundle_idx_list[head_rank : head_rank + interval]
             engine_bundle_idxs = [meta[1] for meta in engine_workers_meta]  # meta: (rank, bundle_idx)
@@ -225,24 +243,35 @@ def init_workers(self):
             activate_worker_idx += interval
         ray.get(set_bundle_idxs_objectref)
         # init dist_init_addr for each worker according to parallel settings
-        init_dist_init_addrs = ray.get([worker.init_dist_port.remote() for worker in self.active_rollout_workers])  # type: ignore[attr-defined]
+        init_dist_init_addrs = ray.get([worker.init_dist_port.remote() for worker in active_rollout_workers])  # type: ignore[attr-defined]
         dist_init_addrs = self._update_dist_init_addr(nodes_per_engine, init_dist_init_addrs, self.num_gpus_per_engine)
         # launch rollout servers
-        self.worker_server_urls_map = dict(
-            ray.get(
-                [
-                    worker.init.remote(dist_init_addrs[i])  # type: ignore[attr-defined]
-                    for i, worker in enumerate(self.active_rollout_workers)
-                ]
-            )
+        worker_server_urls_map = dict(  # rank -> url
+            ray.get([worker.init.remote(dist_init_addrs[i]) for i, worker in enumerate(active_rollout_workers)])
         )
-        self._update_active_workers_and_urls()
-        self.worker_server_urls = list(self.worker_server_urls_map.values())
-        self.logger.info(f"Rollout worker server URLs: {self.worker_server_urls}")
-        self.active_workers_to_status = {worker: True for worker in self.active_rollout_workers}
-        self.active_url_to_workers = dict(zip(self.worker_server_urls, self.active_rollout_workers))
-        self.url_failed_counts = {url: 0 for url in self.worker_server_urls}
-        return engine_mesh_list, self.worker_server_urls_map
+        active_rollout_workers, worker_server_urls_map = self._update_active_workers_and_urls_map(
+            active_rollout_workers, worker_server_urls_map
+        )
+        self.workers_info = {
+            url: WorkerInfo(actor=worker)
+            for url, worker in zip(worker_server_urls_map.values(), active_rollout_workers)
+        }
+        self.logger.info(f"Rollout worker server URLs: {list(self.workers_info.keys())}")
+        return engine_mesh_list, worker_server_urls_map
+
+    def _deactivate_worker(self, url: str):
+        """A helper function to deactivate a worker, update all related states,
+        and shut it down."""
+        worker_info = self.workers_info.get(url)
+        if not worker_info or not worker_info.is_active:
+            return
+
+        self.logger.warning(f"Deactivating rollout worker {worker_info.actor} with URL {url} due to failures.")
+        worker_info.is_active = False
+        self.router.update_active_workers(self._get_worker_status_for_router())
+
+        ray.get(worker_info.actor.offload.remote())  # type: ignore[attr-defined]
+        ray.get(worker_info.actor.shutdown.remote())  # type: ignore[attr-defined]
 
     def check_active_workers(self):
         """Check the health of all active rollout workers.
@@ -251,28 +280,41 @@ def check_active_workers(self):
             List[bool]: A list of booleans indicating the health status of
                 each active rollout worker.
         """
+        active_workers = [(url, info) for url, info in self.workers_info.items() if info.is_active]
+        if not active_workers:
+            return
 
-        active_worker_response = ray.get(
-            [worker.check_health.remote() for worker in self.active_rollout_workers]  # type: ignore[attr-defined]
-        )
-        for idx, status in enumerate(active_worker_response):
-            if not status:
-                self.logger.info(
-                    f"Rollout worker {self.worker_server_urls[idx]} is unhealthy. Removing it from active workers."
-                )
-                self.active_workers_to_status[self.active_rollout_workers[idx]] = False
-
-    def deactivate_worker_by_url(self, url):
-        self.url_failed_counts[url] += 1
-        if self.url_failed_counts[url] < self.config.max_retry_per_worker:
+        urls, infos = zip(*active_workers)
+        actors = [info.actor for info in infos]
+
+        health_statuses = ray.get([actor.check_health.remote() for actor in actors])
+
+        count = 0
+        for url, is_healthy in zip(urls, health_statuses):
+            if count == 3:
+                is_healthy = False
+            count += 1
+            if not is_healthy:
+                self._deactivate_worker(url)
+
+    def deactivate_worker_by_url(self, url: str):
+        """Deactivates a worker identified by its URL after it exceeds the
+        maximum retry count."""
+        worker_info = self.workers_info.get(url)
+        if not worker_info or not worker_info.is_active:
+            return
+
+        worker_info.failure_count += 1
+        if (
+            self.config.max_retry_per_worker is not None
+            and worker_info.failure_count < self.config.max_retry_per_worker
+        ):
             self.logger.warning(
-                f"Rollout worker {url} failed {self.url_failed_counts[url]} times, but not deactivated yet."
+                f"Rollout worker {url} failed {worker_info.failure_count} times, but not deactivated yet."
             )
             return
-        inactive_workers = self.active_url_to_workers.get(url)
-        if inactive_workers:
-            self.logger.warning(f"Deactivating rollout worker {url} due to repeated failures.")
-            self.active_workers_to_status[inactive_workers] = False
+
+        self._deactivate_worker(url)
 
     async def rollout(
         self,
@@ -429,9 +471,11 @@ def _broadcast_to_active_workers(self, method_name: str, block: bool):
             A list of futures if `block` is False, otherwise a list of results.
         """
         futures = []
-        for worker, status in self.active_workers_to_status.items():
-            if status:
-                futures.append(getattr(worker, method_name).remote())
+        for info in self.workers_info.values():
+            if info.is_active:
+                futures.append(getattr(info.actor, method_name).remote())
+            else:
+                self.logger.warning(f"Skipping {method_name} for inactive worker {info.actor}.")
 
         if not block:
             return futures

xtuner/v1/rl/base/worker.py

@@ -527,7 +527,11 @@ def onload_optimizer(self):
         self._engine.put_optimizer_to_device(DEVICE)
 
     def update_rollout_info(
-        self, engine_mesh_list: DeviceMeshRaw, server_url_dict: ServiceUrlMap, rollout_config: RolloutConfig
+        self,
+        engine_mesh_list: DeviceMeshRaw,
+        server_url_dict: ServiceUrlMap,
+        rollout_config: RolloutConfig,
+        worker_server_urls_status: Dict[str, bool],
     ):
         """Update the rollout information for the training worker."""
         tp = rollout_config.tensor_parallel_size
@@ -537,6 +541,9 @@ def update_rollout_info(
             "cpu", mesh=engine_mesh_list, mesh_dim_names=("engine_instance", "engine_parallel")
         )
         self.rollout_url = server_url_dict.get(self.rank, "")
+        if worker_server_urls_status.get(self.rollout_url, False) is False:
+            self.logger.error(f"Rollout server url {self.rollout_url} is not available.")
+            self.rollout_url = None
         self.rollout_cfg_info["tp"] = tp
         self.rollout_cfg_info["ep"] = ep
         self.rollout_cfg_info["api_key"] = rollout_config.api_key
@@ -858,7 +865,9 @@ def request_update_params(self, state_dict, finished=False):
         cpu_mesh = self.rollout_device_mesh["engine_parallel"]
         cpu_group = cpu_mesh.get_group()
         head_rank = cpu_mesh.mesh[0].item()
-
+        if self.rollout_url is None:
+            self.logger.error(f"rank {self.rank} url in None, cannot update weights and skip")
+            return
         if self.rollout_cfg_info["backend"] == "pytorch":
             # TODO(chenchiyu): remove lmdeploy related code
             from lmdeploy.utils import serialize_state_dict
@@ -989,7 +998,6 @@ def request_update_params(self, state_dict, finished=False):
 
                 if use_flattened_tensor_bucket:
                     data["load_format"] = "flattened_bucket"
-
                 response = requests.post(
                     f"{self.rollout_url}/{self.endpoints['update_weights']}", headers=headers, json=data
                 )

xtuner/v1/train/rl_trainer.py

@@ -426,6 +426,9 @@ def fit(self):
             with timer("saving and sync_weight", step_timer_dict):
                 ray.get(self._train_controller.offload.remote(target="optimizer"))
                 self._maybe_save_hf()
+                bind_train_rollout(
+                    train_controller=self._train_controller, env_controller=self._rollout_env_controller
+                )
                 ray.get(self._rollout_env_controller.onload_weights.remote())
                 ray.get(self._train_controller.update_weights.remote())
                 self.logger.info("Model weights synchronized successfully.")