[Refactor] refactor http request error

Author: YanhuiDua

xtuner/v1/data_proto/rl_data.py

@@ -69,7 +69,7 @@ class RLRolloutResponseItem(BaseModel):
     response: Optional[str] = None
     response_ids: Optional[List[int]] = None
     num_return_tokens: Optional[int] = None
-    finish_reason: Optional[str] = None
+    finish_reason: Optional[str] = None  # "stop", "length", "abort", "failed", "skipped"
     logprobs: Optional[List[float]] = None
     extra_info: Dict[str, Any] = dict()
 

xtuner/v1/ray/config/worker.py

@@ -198,6 +198,13 @@ class RolloutConfig(BaseModel):
             help='Extra configuration for different rollout worker. vllm parameters will start with prefix "vllm", etc.',
         ),
     ] = {"lmdeploy_log_level": "CRITICAL", "lmdeploy_uvicorn_log_level": "CRITICAL"}
+    max_retry_per_worker: Annotated[
+        Optional[int],
+        Parameter(
+            group=infer_group,
+            help="Maximum number of retries per rollout worker before deactivation.",
+        ),
+    ] = None
     worker_log_dir: Annotated[Path, Parameter(help="Directory to save worker logs.")] = Path.cwd() / "work_dir"
 
     def __init__(self, **kwargs):
@@ -268,6 +275,9 @@ def __init__(self, **kwargs):
             else:
                 kwargs["rollout_max_batch_size_per_instance"] = 128
 
+        if "max_retry_per_worker" not in kwargs:
+            kwargs["max_retry_per_worker"] = int(kwargs["rollout_max_batch_size_per_instance"] * 0.1)
+
         super().__init__(**kwargs)
         self.worker_log_dir.mkdir(parents=True, exist_ok=True)
 

xtuner/v1/ray/dataflow/flow.py

@@ -193,6 +193,11 @@ async def worker_task(self, group_samples_for_retry: Optional[List[RLDataFlowIte
                 f"Dataflow item check failed for {group_data_items[0].uid.action_id} response. Returning meta for retry."
             )
             return group_data_items
+        if any(item.env.rollout.finish_reason == "skipped" for item in group_data_items):
+            self.logger.warning(
+                f"Bad request for {group_data_items[0].uid.action_id} response. Skipping this request."
+            )
+            return
 
         # step 3: filter
         filtered_group_data_items = await self.replay_buffer.post_processor.remote(group_data_items)  # type: ignore[attr-defined]

xtuner/v1/ray/environment/single_turn_env.py

@@ -62,18 +62,18 @@ async def generate(
                 and state from the rollout controller.
         """
         if self.rollout_controller:
-            # 在env中对输入的数据进行转换，是为了支持rollout_controller单独作为rollout engine使用，使各个模块进行解耦
-            # 每个模块返回独立的data item, 在env中进行更新
-            response_future = [
-                self.rollout_controller.rollout.remote(
+            response_future = []
+            for sample in group_data_items:
+                sample.data.extra_info["root_id"] = sample.uid.root_id
+                sample.data.extra_info["action_id"] = sample.uid.action_id
+                fut = self.rollout_controller.rollout.remote(
                     prompt=sample.data.messages,
                     input_ids=sample.data.input_ids,
                     sample_params=sample_params,
                     extra_params=extra_params,
                     extra_info=sample.data.extra_info,
                 )
-                for sample in group_data_items
-            ]
+                response_future.append(fut)
             try:
                 rollout_responses = await asyncio.wait_for(
                     asyncio.gather(*response_future), timeout=self.rollout_timeout
@@ -109,8 +109,7 @@ async def run(
         """
         group_data_items = await self.generate(group_data_items, sample_params, extra_params)  # type: ignore[assignment]
         skip_judger = any(
-            item.env.rollout.finish_reason == "abort" or item.env.rollout.finish_reason == "failed"
-            for item in group_data_items
+            item.env.rollout.finish_reason in ["failed", "skipped", "abort"] for item in group_data_items
         )
         if self.judger_controller and not skip_judger:
             try:

xtuner/v1/ray/rollout/controller.py

@@ -106,15 +106,17 @@ def __init__(
         self.num_workers = 0
         self.worker_server_urls: List[str] = []
         self.active_rollout_workers: List[RolloutWorker] = []
-        self.active_rollout_workers_status: Dict = {}
+        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.start_api_server()
         # todo(@duanyanhui): add router to replace native round robin
-        self.router = SessionRouter(self.active_rollout_workers_status)
+        self.router = SessionRouter(self.active_workers_to_status)
         self.sample_params = SampleParams().dict()
         # note: 目前默认使用return_token_ids和return_logprob，并且不使用流式
         self.extra_params = dict(
@@ -237,7 +239,10 @@ def init_workers(self):
         )
         self._update_active_workers_and_urls()
         self.worker_server_urls = list(self.worker_server_urls_map.values())
-        self.active_rollout_workers_status = {worker: True for worker in self.active_rollout_workers}
+        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
 
     def check_active_workers(self):
@@ -254,9 +259,19 @@ def check_active_workers(self):
         for idx, status in enumerate(active_worker_response):
             if not status:
                 self.logger.info(
-                    f"Rollout worker {self.active_rollout_workers[idx]} is unhealthy. Removing it from active workers."
+                    f"Rollout worker {self.worker_server_urls[idx]} is unhealthy. Removing it from active workers."
                 )
-                self.active_rollout_workers_status[self.active_rollout_workers[idx]] = False
+                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:
+            self.logger.warning(
+                f"Rollout worker {url} failed {self.url_failed_counts[url]} times, but not deactivated yet."
+            )
+            return
+        inactive_workers = self.active_url_to_workers.get(url)
+        self.active_workers_to_status[inactive_workers] = False
 
     async def rollout(
         self,
@@ -296,7 +311,6 @@ async def rollout(
         self.sample_params.update(sample_params.dict() if sample_params else {})
         self.extra_params.update(extra_params if extra_params else {})
         if self.print_params_flag:
-            # 通过print_params_flag控制只打印一次参数
             self.logger.info(f"Rollout with sample params: {self.sample_params}, extra params: {self.extra_params}")
             self.print_params_flag = False
         assert prompt is not None or input_ids is not None, "Either prompt or input_ids must be provided."
@@ -311,8 +325,18 @@ async def rollout(
             extra_info=extra_info,
         )
         try:
-            response = await asyncio.wait_for(response_ref, timeout=self.config.rollout_timeout)
-            return response
+            response, http_result = await asyncio.wait_for(response_ref, timeout=self.config.rollout_timeout)
+            if http_result.is_success:
+                return response
+            elif http_result.is_retryable or http_result.is_server_error:
+                response.finish_reason = "failed"
+                return response
+            elif http_result.is_client_error:
+                response.finish_reason = "skipped"
+                return response
+            else:  # unknown error
+                raise RuntimeError("Unknown error occurred during rollout. Error message: ", http_result.error_message)
+
         except asyncio.TimeoutError:
             self.logger.error("Get response from rollout worker timeout and return the failed response.")
             failed_response = RLRolloutResponseItem(
@@ -409,7 +433,7 @@ 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_rollout_workers_status.items():
+        for worker, status in self.active_workers_to_status.items():
             if status:
                 futures.append(getattr(worker, method_name).remote())
 

xtuner/v1/ray/rollout/worker.py

@@ -21,6 +21,7 @@
 from xtuner.v1.ray.base import AutoAcceleratorWorkers, SingleAcceleratorWorker
 from xtuner.v1.ray.config import RolloutConfig
 from xtuner.v1.utils import get_logger
+from xtuner.v1.utils.httpx_utils import HttpRequestErrorType, HttpRequestResult
 
 
 class RolloutWorker(SingleAcceleratorWorker):
@@ -285,9 +286,8 @@ def _check_infer_engine_version(self, return_token_ids: bool):
                     )
             self.check_flag = False
 
-    async def _safe_post_request(self, url, headers, payload) -> Tuple[Optional[httpx.Response], bool, Optional[str]]:
+    async def _safe_post_request(self, url, headers, payload) -> HttpRequestResult:
         try:
-            # new_url = self.server_url[-2] + str(int(self.server_url[-1]) + 1) + "'"
             req = self.client.build_request(
                 "POST",
                 url,
@@ -296,34 +296,11 @@ async def _safe_post_request(self, url, headers, payload) -> Tuple[Optional[http
             )
             r = await self.client.send(req)
             r.raise_for_status()
-            return r, True, None
-        # NOTE(@duanyanhui): 目前只有TimeoutException时，第二个返回值为True ，即continue_rollout=True，不影响主程序正常运行
-        # 其他错误都认为是请求失败，会通过assert进行报错，并且根据错误类型返回不同的error msg.
-        except httpx.TimeoutException as e:
-            error_msg = f"create_request error: Request to {url} timed out: {e}"
-            self.logger.warning(error_msg)
-            return None, True, None
-        except httpx.HTTPStatusError as e:
-            if e.response.status_code == 400:
-                log_payload = copy.deepcopy(payload)
-                if "input_ids" in log_payload and log_payload["input_ids"] is not None:
-                    log_payload["input_ids"] = str(log_payload["input_ids"])
-                error_msg = (
-                    f"Bad Request (400) Error for {url} with payload {log_payload}. Server response: {e.response.text}"
-                )
-                return None, False, error_msg
-            else:
-                error_msg = f"HTTP error occurred for {url}: {e.response.status_code} - {e.response.text}"
-                return None, False, error_msg
-        except httpx.RequestError as e:
-            log_payload = copy.deepcopy(payload)
-            if "input_ids" in log_payload and log_payload["input_ids"] is not None:
-                log_payload["input_ids"] = str(log_payload["input_ids"])
-            error_msg = f"Request Error occurred while requesting {payload} to {url}: {e}"
-            return None, False, error_msg
+            return HttpRequestResult(response=r)
         except Exception as e:
-            error_msg = f"Unexpected Error occurred: {e} with traceback: \n {traceback.format_exc()}"
-            return None, False, error_msg
+            error_type = HttpRequestErrorType.from_exception(e)
+            result = HttpRequestResult(error_type=error_type, exception=e, url=url, payload=payload)
+            return result
 
     async def rollout_task(
         self,
@@ -335,58 +312,59 @@ async def rollout_task(
         extra_params: dict,
         format: str,
         extra_info: dict,
-    ) -> RLRolloutResponseItem:
-        uid = str(uuid.uuid4())
+    ) -> Tuple[RLRolloutResponseItem, HttpRequestResult]:
+        uid = extra_info.get("action_id", str(uuid.uuid4()))
         response = None
-        failed_rollout_response = RLRolloutResponseItem(
-            finish_reason="failed",
-        )
+        failed_rollout_response = RLRolloutResponseItem(finish_reason="failed")
         self._check_infer_engine_version("return_token_ids" in extra_params and extra_params["return_token_ids"])
 
         if format == "openai":
             openai_prompts, openai_tools = prompts, tools
         else:
             openai_prompts, openai_tools = self._adapt_input_to_openai_spec(prompts, tools, tool_choice)
+
         if "return_token_ids" in extra_params and extra_params["return_token_ids"]:
-            response, continue_rollout, error_msg = await self._create_request(
-                f"{self.server_url}/{self.endpoints['generate']}",
-                openai_prompts,
-                input_ids,
-                openai_tools,
-                tool_choice,
-                sample_params=sample_params,
-                extra_params=extra_params,
-                extra_info=extra_info,
-            )
+            endpoint_url = f"{self.server_url}/{self.endpoints['generate']}"
         else:
-            assert prompts is not None, "prompts should not be None when you call v1/chat/completions API"
-            response, continue_rollout, error_msg = await self._create_request(
-                f"{self.server_url}/{self.endpoints['v1/chat/completions']}",
-                openai_prompts,
-                None,
-                openai_tools,
-                tool_choice,
-                sample_params=sample_params,
-                extra_params=extra_params,
-                extra_info=extra_info,
-            )
-        assert continue_rollout, (
-            f"Unhandled error occurred during rollout request creation, You should check infer engine or input params. \n Error message: {error_msg}"
+            endpoint_url = f"{self.server_url}/{self.endpoints['v1/chat/completions']}"
+
+        http_result = await self._create_request(
+            endpoint_url,
+            openai_prompts,
+            None,
+            openai_tools,
+            tool_choice,
+            sample_params=sample_params,
+            extra_params=extra_params,
+            extra_info=extra_info,
         )
-        if response:
+
+        if http_result.response is not None:
             try:
                 rollout_response = (
                     await self._handle_stream_response(uid, sample_params, extra_params, response)
                     if extra_params["stream"]
                     else await self._handle_non_stream_response(uid, sample_params, extra_params, response)
                 )
             finally:
-                if hasattr(response, "aclose"):
-                    await response.aclose()
-            return rollout_response
+                if hasattr(http_result.response, "aclose"):
+                    await http_result.response.aclose()
+            return rollout_response, http_result
         else:
-            self.logger.warning(f"Retrying rollout for {uid} due to httpx timeout")
-            return failed_rollout_response
+            if http_result.is_retryable:
+                self.logger.warning(f"Retryable error occurred during rollout request {uid} to {http_result.url}")
+                return failed_rollout_response, http_result
+            elif http_result.is_server_error:
+                self.logger.error(
+                    f"Server error during rollout request {uid} to {http_result.url}, please check the server logs."
+                )
+                http_result.url = self.server_url
+                return failed_rollout_response, http_result
+            else:  # http_result.is_client_error:
+                self.logger.error(
+                    f"Client error during rollout request {uid} to {http_result.url} and skip this request."
+                )
+                return failed_rollout_response, http_result
 
     async def _handle_stream_response(self, uid, sample_params, extra_params, response) -> RLRolloutResponseItem:
         last_trajectory = ""
@@ -556,7 +534,7 @@ async def rollout(
         extra_params: dict = dict(),
         format: str = "openai",
         extra_info: dict = dict(),
-    ) -> RLRolloutResponseItem:
+    ) -> Tuple[RLRolloutResponseItem, HttpRequestResult]:
         """Public method to initiate a rollout.
 
         Args: