fix abort state

Author: YanhuiDua

xtuner/v1/data_proto/rl_data.py

@@ -3,7 +3,9 @@
 from cyclopts import Parameter
 from pydantic import BaseModel, ConfigDict, Field
 from typing_extensions import Annotated
+from xtuner.v1.utils.logger import get_logger   
 
+logger = get_logger()
 
 # ====================================
 # ====== DataFlow 数据流 ==============
@@ -84,10 +86,12 @@ def update(self, other: "RLRolloutResponseItem") -> None:
             self.response_ids.extend(other.response_ids)
         else:
             self.response_ids = other.response_ids
+
         if self.logprobs is not None and other.logprobs:
             self.logprobs.extend(other.logprobs)
         else:
             self.logprobs = other.logprobs
+            
         if self.response is not None and other.response:
             self.response += other.response
         else:
@@ -197,10 +201,13 @@ def check_valid_dataflow_item(group_data_items: List[RLDataFlowItem]) -> bool:
         ids_valid = bool(rollout_info.response_ids)
         logprobs_valid = bool(rollout_info.logprobs)
         if item.env.rollout.state in ["skipped", "failed"]:
+            logger.info(f"Invalid dataflow item found: rollout state is {item.env.rollout.state}. UID: {item.uid}")
             return False
-        if not response_valid and not ids_valid:
+        if not response_valid and not ids_valid and item.env.rollout.state != "interrupted":
+            logger.info(f"Invalid dataflow item found: no response or response_ids. UID:{item.data.uid} with rollout response {item.env.rollout}")
             return False
         if ids_valid and logprobs_valid and len(rollout_info.logprobs) != len(rollout_info.response_ids):  # type: ignore[arg-type]
+            logger.info(f"Invalid dataflow item found: logprobs and response_ids length mismatch. UID: {item.uid}")
             return False
     return True
 

xtuner/v1/ray/dataflow/flow.py

@@ -122,6 +122,7 @@ def __init__(
         self.finished_samples_count = 0
         self.failed_samples_count = 0
         self.skipped_sample_count = 0
+        self.filtered_sample_count = 0
         self.sample_from_expired_storage = False
         self.logger = get_logger(log_dir=self.config.worker_log_dir, tag="DataFlow")
         self.target_batch_size = self.config.global_batch_size
@@ -161,6 +162,7 @@ def _reset_internal_states_on_step(
         self.finished_samples_count = 0
         self.failed_samples_count = 0
         self.skipped_sample_count = 0
+        self.filtered_sample_count = 0
         self.logger.info(
             f"global_batch_size: {global_batch_size}, sample_params: {sample_params}, extra_params: {extra_params}"
         )
@@ -224,6 +226,8 @@ async def worker_task(self, group_samples_for_retry: Optional[List[RLDataFlowIte
             group_data_items = await self.replay_buffer.post_processor.remote(group_data_items)  # type: ignore[attr-defined]
             if len(group_data_items) > 0:
                 await self.replay_buffer.add.remote(group_data_items)  # type: ignore[attr-defined]
+            else:
+                self.filtered_sample_count += 1
             self.logger.debug(f"Worker task completed successfully for {action_id}.")
         elif group_state == "interrupted":
             await self.replay_buffer.add.remote(group_data_items)  # type: ignore[attr-defined]
@@ -289,17 +293,17 @@ async def concurrent_task_runner(self):
             cleanup_start_time = time.monotonic()
             cleanup_timeout = 10 * 60  # 10 minutes in seconds
             while len(waiting_tasks) > 0:
-                elapsed_time = time.monotonic() - cleanup_start_time
-                if elapsed_time > cleanup_timeout:
-                    self.logger.warning(
-                        f"Cleanup timeout of {cleanup_timeout}s reached. "
-                        f"Forcefully cancelling {len(waiting_tasks)} remaining tasks."
-                    )
-                    for task in waiting_tasks:
-                        task.cancel()
-                    # Wait for cancellations to complete
-                    await asyncio.gather(*waiting_tasks, return_exceptions=True)
-                    break  # Exit the cleanup loop
+                # elapsed_time = time.monotonic() - cleanup_start_time
+                # if elapsed_time > cleanup_timeout:
+                #     self.logger.warning(
+                #         f"Cleanup timeout of {cleanup_timeout}s reached. "
+                #         f"Forcefully cancelling {len(waiting_tasks)} remaining tasks."
+                #     )
+                #     for task in waiting_tasks:
+                #         task.cancel()
+                #     # Wait for cancellations to complete
+                #     await asyncio.gather(*waiting_tasks, return_exceptions=True)
+                #     break  # Exit the cleanup loop
                 done_tasks, pending_tasks = await asyncio.wait(
                     waiting_tasks, timeout=0.1, return_when=asyncio.FIRST_COMPLETED
                 )
@@ -309,7 +313,8 @@ async def concurrent_task_runner(self):
             self.logger.info("All worker tasks have completed after pausing env controller.")
 
         self.logging_replaybuffer_state()
-
+        self.logger.info(ray.get(self.env_controller.get_rollout_stats.remote()))
+        
     async def pause(self, timeout: float = 60.0):
         """Asynchronously sends abort requests to all rollout workers."""
         rollout_info = ray.get(self.env_controller.get_rollout_info.remote())  # type: ignore[attr-defined]
@@ -378,6 +383,7 @@ def logging_replaybuffer_state(self, logging_msg: Optional[str] = None):
         status = self.get_replaybuffer_status()
         logging_msg = logging_msg if logging_msg else ""
         logging_msg += f"ReplayBuffer Status: {status}"
+        logging_msg += f", Filtered samples count: {self.filtered_sample_count}"
         self.logger.info(logging_msg)
 
     def get_replaybuffer_status(self):

xtuner/v1/ray/dataflow/replay_buffer.py

@@ -336,8 +336,9 @@ def add(self, grouped_dataitem: List[RLDataFlowItem], partial_rollout_step: int
 
         # 1. 跟prompt相关的action_id记录
         if root_id in self._root2actions:
+            # TODO: version 更新需要 根据是否update_weights来判断，需要考虑到非共卡的情况
             replay_meta.version += 1
-            self.logger.debug(f"Existing root_id: {root_id} found. Incrementing version to {replay_meta.version}.")
+            self.logger.info(f"Existing root_id: {root_id} found. Incrementing version to {replay_meta.version}.")
             self._root2actions[root_id].append(action_id)
         else:
             self._root2actions[root_id] = [action_id]
@@ -346,14 +347,14 @@ def add(self, grouped_dataitem: List[RLDataFlowItem], partial_rollout_step: int
         # 2. 根据rollout状态加到finished, abort, abort_over_version队列中；Partial rollout is handled based on whether finish_reason is "abort".
         if replay_meta.state == ReplayState.INTERRUPTED and replay_meta.version < partial_rollout_step:
             self._interrupted_actions[replay_meta.version].append(action_id)
-            self.logger.debug(
+            self.logger.info(
                 f"Add aborted sample with root_id: {root_id}, action_id: {action_id} to _interrupted_actions."
             )
         elif replay_meta.state == ReplayState.INTERRUPTED and replay_meta.version >= partial_rollout_step:
             self._expired_actions.append(action_id)
             replay_meta.version = 0
             replay_meta.state = ReplayState.EXPIRED
-            self.logger.debug(
+            self.logger.info(
                 f"Action_id: {action_id} has exceeded partial_rollout_step {partial_rollout_step}. Add this sample with root_id: {root_id} to _expired_actions list."
             )
         elif replay_meta.state == ReplayState.COMPLETED:
@@ -386,8 +387,9 @@ def get(self, global_batch_size: int) -> Tuple[List[List[RLDataFlowItem]], List[
         """
         samples = []
         multimodal_train_infos = []
-        target_batch_size = min(global_batch_size, len(self._completed_actions))
-        for _ in range(global_batch_size):
+        target_batch_size = min(global_batch_size, self.get_completed_samples())
+        self.logger.info(f"Retrieving {target_batch_size} completed samples from the replay buffer.")
+        for _ in range(target_batch_size):
             action_id = self._pop_highest_version_action(self._completed_actions)
             replay_meta = self._actions[action_id]  # type: ignore[index]
             group_samples = mapping_replaymeta_to_dataitem(replay_meta)
@@ -563,16 +565,17 @@ def sample_from_interrupted_storage(self, tokenizer) -> List[RLDataFlowItem]:
             sample.uid.version = replay_meta.version
             sample.extra_info.state = str(ReplayState.INIT)
             if sample.env.rollout.response_ids and sample.data.input_ids:
+                # TODO： response_ids 累加 
                 if "train_prompt_ids" in sample.data.extra_info:
                     sample.data.input_ids = (
                         sample.data.extra_info["train_prompt_ids"] + sample.env.rollout.response_ids
                     )
                 else:
                     sample.data.input_ids.extend(sample.env.rollout.response_ids)
-            # elif sample.env.rollout.response:
-            #     sample.data.input_ids.extend(tokenizer.encode(sample.env.rollout.response, add_special_tokens=False))
+            elif sample.env.rollout.response:
+                sample.data.input_ids.extend(tokenizer.encode(sample.env.rollout.response, add_special_tokens=False))
         self.logger.info(
-            f"Sampling interrupted action_id: {action_id} from replay buffer, remain interrupted samples: {len(self._interrupted_actions)}"
+            f"Sampling interrupted action_id: {action_id} from replay buffer, remain interrupted samples: {self.get_interrupted_samples()}"
         )
         return group_samples
 

xtuner/v1/ray/environment/base_env.py

@@ -212,3 +212,11 @@ def check_active_workers(self, block=True):
             block (bool): Whether to block until the operation completes.
         """
         return self._call_rollout_func("check_active_workers", block)
+
+    def get_rollout_stats(self, block=True):
+        """Gets statistics from the rollout workers.
+
+        Args:
+            block (bool): Whether to block until the operation completes.
+        """
+        return self._call_rollout_func("get_rollout_stats", block)

xtuner/v1/ray/environment/single_turn_env.py

@@ -87,7 +87,7 @@ async def generate(
                 sample.data.extra_info["action_id"] = sample.uid.action_id
                 if sample.env.rollout.num_return_tokens > 0:
                     sample.data.extra_info["num_return_tokens"] = sample.env.rollout.num_return_tokens
-                    self.logger.debug(
+                    self.logger.info(
                         f"Set num_return_tokens: {sample.env.rollout.num_return_tokens} for sample {sample.uid}."
                     )
                 fut = self.rollout_controller.rollout.remote(
@@ -136,7 +136,7 @@ async def run(
         if self.judger_controller and continue_judger:
             try:
                 judger_responses: List[RLJudgerResponseItem] = await asyncio.wait_for(
-                    self.judger_controller.run.remote(group_data_items), timeout=self.judger_timeout
+                    self.judger_controller.run.remote(group_data_items), timeout=self.judger_timeout * 2 
                 )
             except asyncio.TimeoutError:
                 self.logger.error("Get judger controller response timeout and return the failed response.")

xtuner/v1/ray/rollout/controller.py

@@ -390,26 +390,41 @@ async def rollout(
             format=format,
             extra_info=extra_info,
         )
-        if self.workers_info[server_url].running_count % 100 == 0:
-            log_msg = ""
-            for _, info in self.workers_info.items():
-                log_msg += f"rank {info.rank} worker info: {info}"
-            self.logger.info(log_msg)
+        # if self.workers_info[server_url].running_count % 100 == 0:
+        #     log_msg = ""
+        #     for _, info in self.workers_info.items():
+        #         log_msg += f"rank {info.rank} worker info: {info}"
+        #     self.logger.info(log_msg)
         try:
             response = await asyncio.wait_for(response_ref, timeout=self.config.rollout_timeout * 2)
-            self.workers_info[server_url].running_count -= 1
             self.workers_info[server_url].success_count += 1
             if response.state == "failed" or response.state == "skipped":
                 self.logger.info(f"Rollout worker {worker} returned state {response.state}. Deactivating worker.")
+                self.workers_info[server_url].skipped_count += 1
                 self.deactivate_worker_by_url(server_url)
             return response
         except asyncio.TimeoutError:
-            self.workers_info[server_url].running_count -= 1
             self.workers_info[server_url].failure_count += 1
             # self.deactivate_worker_by_url(server_url) # do not deactivate on timeout, only on skipped state
             self.logger.error(f"Get response from rollout worker {worker} timeout and return skip this sample.")
+            self.deactivate_worker_by_url(server_url) 
             return RLRolloutResponseItem(state="skipped")
 
+    def get_rollout_stats(self) -> str:
+        """Get statistics about the rollout workers.
+        Returns:
+            str: A formatted string containing statistics about each rollout
+        """
+        log_parts = ["Rollout Worker Stats:"]
+        for url, info in self.workers_info.items():
+            log_parts.append(
+                f"  - URL: {url} | Rank: {info.rank} | Active: {info.is_active} | "
+                f"Running: {info.running_count} | Success: {info.success_count} | "
+                f"Failures: {info.failure_count}"
+            )
+        log_msg = "\n".join(log_parts)
+        return log_msg
+    
     def start_api_server(self, host: str = "0.0.0.0", port: int = 8000):
         """Starts the API server to expose the rollout functionality."""
         app = FastAPI()

xtuner/v1/ray/rollout/lmdeploy.py

@@ -127,7 +127,7 @@ async def _create_request(
         if "num_return_tokens" in extra_info:
             max_return_tokens = sample_params["max_tokens"] - extra_info["num_return_tokens"]
             sample_params["max_tokens"] = max_return_tokens
-            self.logger.debug(
+            self.logger.info(
                 f"Set max_tokens to {max_return_tokens} based on num_return_tokens {extra_info['num_return_tokens']}"
             )
 

xtuner/v1/ray/rollout/sglang.py

@@ -78,11 +78,11 @@ async def _create_request(
             payload["messages"] = prompt
             payload.update(sglang_sample_params)
             # note: chat completions 接口需要传入 max_tokens 和 min_tokens 参数
-            if "num_return_tokens" in extra_params:
-                max_return_tokens = sglang_sample_params["max_new_tokens"] - extra_params["num_return_tokens"]
+            if "num_return_tokens" in extra_info:
+                max_return_tokens = sglang_sample_params["max_new_tokens"] - extra_info["num_return_tokens"]
                 payload["max_tokens"] = max_return_tokens
                 self.logger.info(
-                    f"Set max_tokens to {max_return_tokens} based on num_return_tokens {extra_params['num_return_tokens']}"
+                    f"Set max_tokens to {max_return_tokens} based on num_return_tokens {extra_info['num_return_tokens']}"
                 )
             else:
                 payload["max_tokens"] = sglang_sample_params["max_new_tokens"]

xtuner/v1/ray/rollout/worker.py

@@ -333,16 +333,24 @@ async def rollout_task(
             endpoint_url = f"{self.server_url}/{self.endpoints['v1/chat/completions']}"
 
         while True:
-            http_result = await self._create_request(
-                endpoint_url,
-                openai_prompts,
-                input_ids,
-                openai_tools,
-                tool_choice,
-                sample_params=sample_params,
-                extra_params=extra_params,
-                extra_info=extra_info,
-            )
+            if extra_info.get("num_return_tokens", None) is not None and (sample_params["max_tokens"] - extra_info["num_return_tokens"]) == 0:
+                return RLRolloutResponseItem(
+                    response="",
+                    response_ids=[],
+                    num_return_tokens=0,
+                    finish_reason="length",
+                )
+            else:
+                http_result = await self._create_request(
+                    endpoint_url,
+                    openai_prompts,
+                    input_ids,
+                    openai_tools,
+                    tool_choice,
+                    sample_params=sample_params,
+                    extra_params=extra_params,
+                    extra_info=extra_info,
+                )
             # Case 1: Request was successful
             if http_result.response is not None:  # 推理完成：completed状态：finish_reason为abort/stop/length, 退出
                 response = await self._handle_non_stream_response(