Propagate dataflow cleanup cancellation to rollout Ray tasks (#1699)

Author: YanhuiDua

xtuner/v1/ray/dataflow/flow.py

@@ -12,7 +12,11 @@
 from tqdm.auto import tqdm
 from typing_extensions import Annotated
 
-from xtuner.v1.data_proto.rl_data import MultimodalTrainInfo, RLDataFlowItem, RolloutState
+from xtuner.v1.data_proto.rl_data import (
+    MultimodalTrainInfo,
+    RLDataFlowItem,
+    RolloutState,
+)
 from xtuner.v1.ray.environment import SingleTurnEnvironment
 from xtuner.v1.ray.rollout.controller import SampleParams
 from xtuner.v1.ray.utils import create_task
@@ -180,6 +184,7 @@ def __init__(
             )
         self.logger.info(f"DataFlowConfig:\n{self.config.model_dump_json(indent=2)}")
         self.cleanup_task_time = 5 * 60  # 5 minutes
+        self.cancel_response_timeout = 5.0
 
     def _reset_internal_states(
         self,
@@ -253,11 +258,22 @@ async def worker_task(self, group_samples_for_retry: Optional[List[RLDataFlowIte
         assert len(group_data_items) > 0, "Sampled empty group data items from replay buffer."
         action_id = group_data_items[0].uid.action_id
         # step 2: env generate
-        group_data_items = await self.env_controller.run.remote(  # type: ignore[attr-defined]
+        env_run_ref = self.env_controller.run.remote(  # type: ignore[attr-defined]
             group_data_items,
             sample_params=self.sample_params,
             extra_params=self.extra_params,
         )
+        try:
+            group_data_items = await asyncio.shield(env_run_ref)
+        except asyncio.CancelledError as exc:
+            ray.cancel(env_run_ref, recursive=True)
+            try:
+                group_data_items = await asyncio.wait_for(
+                    asyncio.shield(env_run_ref),
+                    timeout=self.cancel_response_timeout,
+                )
+            except BaseException:
+                raise exc
 
         # Step 3: Determine the sample's state and act accordingly.
         group_state = determine_group_state(group_data_items)

xtuner/v1/ray/environment/single_turn_env.py

@@ -67,6 +67,7 @@ def __init__(
         # This should be longer than the controller's internal timeout (`rollout_timeout`)
         # to account for potential queuing delays and other overheads.
         self.timeout_multiplier = 2.0
+        self.cancel_response_timeout = 5.0
         self.rollout_cfg = rollout_cfg
 
     async def generate(  # type: ignore[override]
@@ -140,10 +141,25 @@ async def generate(  # type: ignore[override]
 
                 response_future.append(fut)
             try:
+                response_gather = asyncio.gather(*response_future)
                 rollout_responses = await asyncio.wait_for(
-                    asyncio.gather(*response_future), timeout=self.rollout_timeout * self.timeout_multiplier
+                    asyncio.shield(response_gather), timeout=self.rollout_timeout * self.timeout_multiplier
                 )
+            except asyncio.CancelledError as exc:
+                for fut in response_future:
+                    ray.cancel(fut, recursive=True)
+                try:
+                    rollout_responses = await asyncio.wait_for(
+                        asyncio.gather(*response_future, return_exceptions=True),
+                        timeout=self.cancel_response_timeout,
+                    )
+                except BaseException:
+                    raise exc
+                if not all(isinstance(response, RLRolloutResponseItem) for response in rollout_responses):
+                    raise exc
             except asyncio.TimeoutError:
+                for fut in response_future:
+                    ray.cancel(fut, recursive=True)
                 self.logger.error("Get rollout controller response timeout and return the failed response.")
                 rollout_responses = [RLRolloutResponseItem(state="skipped") for _ in group_data_items]
             group_data_items = update_rollout_item(group_data_items, rollout_responses)
@@ -171,12 +187,17 @@ async def run(  # type: ignore[override]
         group_data_items = await self.generate(group_data_items, sample_params, extra_params)  # type: ignore[assignment]
         continue_judger = is_valid_for_training(group_data_items)
         if self.judger_controller and continue_judger:
+            judger_response_ref = self.judger_controller.run.remote(group_data_items)
             try:
                 judger_responses: List[RLJudgerResponseItem] = await asyncio.wait_for(
-                    self.judger_controller.run.remote(group_data_items),
+                    judger_response_ref,
                     timeout=self.judger_timeout * self.timeout_multiplier,
                 )
+            except asyncio.CancelledError:
+                ray.cancel(judger_response_ref, recursive=True)
+                raise
             except asyncio.TimeoutError:
+                ray.cancel(judger_response_ref, recursive=True)
                 self.logger.error("Get judger controller response timeout and return the failed response.")
                 judger_responses = [
                     RLJudgerResponseItem(

xtuner/v1/ray/rollout/controller.py

@@ -163,6 +163,7 @@ def __init__(
         # This should be longer than the controller's internal timeout (`rollout_timeout`)
         # to account for potential queuing delays and other overheads.
         self.timeout_multiplier = 2.0
+        self.cancel_response_timeout = 5.0
 
     def _get_worker_status_for_router(self) -> Dict[RolloutWorker, bool]:
         """Helper to generate the status dict required by the SessionRouter."""
@@ -403,15 +404,22 @@ async def rollout(
         try:
             selected_worker_info = self.workers_info[server_url]
             response = await asyncio.wait_for(
-                response_ref, timeout=self.config.rollout_timeout * self.timeout_multiplier
+                asyncio.shield(response_ref), timeout=self.config.rollout_timeout * self.timeout_multiplier
             )
             selected_worker_info.success_count += 1
             if response.state == "failed" or response.state == "skipped":
                 selected_worker_info.failure_count += 1
                 self.logger.error(f"Get failed/skipped response from rollout worker {worker}, deactivate it.")
                 self.deactivate_worker_by_url(server_url)
             return response
+        except asyncio.CancelledError as exc:
+            ray.cancel(response_ref, recursive=True)
+            try:
+                return await asyncio.wait_for(asyncio.shield(response_ref), timeout=self.cancel_response_timeout)
+            except BaseException:
+                raise exc
         except asyncio.TimeoutError:
+            ray.cancel(response_ref, recursive=True)
             selected_worker_info.failure_count += 1
             self.logger.error(f"Get response from rollout worker {worker} timeout and return skip this sample.")
             self.deactivate_worker_by_url(server_url)

xtuner/v1/ray/rollout/worker.py

@@ -331,6 +331,7 @@ def _check_infer_engine_version(self, return_token_ids: bool):
             self.check_flag = False
 
     async def _safe_post_request(self, url, headers, payload) -> HttpRequestResult:
+        send_task = None
         try:
             if self.receive_abort_request.is_set():
                 self.logger.debug(f"Request to {url} was cancelled before sending due to an abort signal.")
@@ -341,10 +342,18 @@ async def _safe_post_request(self, url, headers, payload) -> HttpRequestResult:
                 headers=headers,
                 json=payload,
             )
-            r = await self.client.send(req)
+            send_task = asyncio.create_task(self.client.send(req))
+            r = await send_task
             r.raise_for_status()
             return HttpRequestResult(response=r)
 
+        except asyncio.CancelledError:
+            self.logger.debug(f"Request to {url} was cancelled while waiting for the response.")
+            if send_task is not None and not send_task.done():
+                send_task.cancel()
+                await asyncio.gather(send_task, return_exceptions=True)
+            self.receive_abort_request.set()
+            return HttpRequestResult(error_type=HttpRequestErrorType.REQUEST_ABORTED, url=url, payload=payload)
         except Exception as e:
             error_type = HttpRequestErrorType.from_exception(e)
             result = HttpRequestResult(error_type=error_type, exception=e, url=url, payload=payload)