fix pause

Author: YanhuiDua

xtuner/v1/data_proto/rl_data.py

@@ -82,18 +82,29 @@ def update(self, other: "RLRolloutResponseItem") -> None:
         if not isinstance(other, RLRolloutResponseItem):
             raise TypeError("Can only update with another RLRolloutResponseItem instance.")
 
-        if self.response_ids is not None and other.response_ids:
-            self.response_ids.extend(other.response_ids)
+        if self.response_ids is not None:
+            init_response_len = len(self.response_ids)
+            if other.response_ids is not None:
+                self.response_ids.extend(other.response_ids)
+                logger.info(f"Updated response_ids from {init_response_len} to {len(self.response_ids)}")
+            else:
+                self.response_ids = self.response_ids
         else:
             self.response_ids = other.response_ids
 
-        if self.logprobs is not None and other.logprobs:
-            self.logprobs.extend(other.logprobs)
+        if self.logprobs is not None:
+            if other.logprobs is not None:
+                self.logprobs.extend(other.logprobs)
+            else:
+                self.logprobs = self.logprobs
         else:
             self.logprobs = other.logprobs
 
-        if self.response is not None and other.response:
-            self.response += other.response
+        if self.response is not None:
+            if other.response is not None:
+                self.response + other.response
+            else:
+                self.response = self.response
         else:
             self.response = other.response
         self.num_return_tokens += other.num_return_tokens
@@ -197,14 +208,14 @@ def check_valid_dataflow_item(group_data_items: List[RLDataFlowItem]) -> bool:
     """
     for item in group_data_items:
         rollout_info = item.env.rollout
-        response_valid = bool(rollout_info.response)
-        ids_valid = bool(rollout_info.response_ids)
-        logprobs_valid = bool(rollout_info.logprobs)
+        response_valid = True if rollout_info.response is not None and len(rollout_info.response) > 0 else False
+        ids_valid = True if rollout_info.response_ids is not None and len(rollout_info.response_ids) > 0 else False
+        logprobs_valid = True if rollout_info.logprobs is not None and len(rollout_info.logprobs) > 0 else False
         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 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}")
+            logger.info(f"Invalid dataflow item found: no response or response_ids. UID:{item.uid.action_id} 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}")

xtuner/v1/ray/dataflow/flow.py

@@ -221,7 +221,7 @@ async def worker_task(self, group_samples_for_retry: Optional[List[RLDataFlowIte
 
         # Step 3: Determine the sample's state and act accordingly.
         group_state = determine_group_state(group_data_items)
-        self.logger.debug(f"Determined replay state for {action_id}: {group_state}")
+        self.logger.info(f"Determined replay state for {action_id}: {group_state}")
         if group_state == "completed":
             group_data_items = await self.replay_buffer.post_processor.remote(group_data_items)  # type: ignore[attr-defined]
             if len(group_data_items) > 0:
@@ -260,6 +260,7 @@ async def concurrent_task_runner(self):
             before completing.
         """
         waiting_tasks = set()
+        start_time = time.monotonic()
         with tqdm(total=self.target_batch_size, desc="rollout_controller for training samples") as pbar:
             update_step = max(1, int(self.target_batch_size * 0.01))
             next_update_threshold = update_step
@@ -286,8 +287,12 @@ async def concurrent_task_runner(self):
             pbar.n = self.finished_samples_count
             pbar.refresh()
 
+        elapsed_time = time.monotonic() - start_time
+        self.logger.info(f"Sample collection finished. Time taken: {elapsed_time:.2f} seconds.")
+
         # NOTE: Directly send pause requests to rollout workers because calling `rollout_controller.pause()`
         # would be queued behind many worker tasks, causing a significant delay.
+        start_time = time.monotonic()
         if self.enable_partial_rollout:
             await self.pause()
             cleanup_start_time = time.monotonic()
@@ -311,21 +316,26 @@ async def concurrent_task_runner(self):
                     await self.pause()
                 waiting_tasks = pending_tasks
             self.logger.info("All worker tasks have completed after pausing env controller.")
-
+        elapsed_time = time.monotonic() - start_time
+        self.logger.info(f"Pause generation. Time taken: {elapsed_time:.2f} seconds.")
         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]
-        self.worker_url_list = list(rollout_info["server_url_dict"].values())
-
+        self.logger.info("Sending abort requests to all rollout workers.")
+        # rollout_info = ray.get(self.env_controller.get_rollout_info.remote())  # type: ignore[attr-defined]
+        # self.worker_url_list = list(rollout_info["server_url_dict"].values())
+        self.logger.info("get self.worker_url_list from env_controller: ", self.worker_url_list)
         if not self.worker_url_list:
             self.logger.info("No active rollout workers to pause.")
             return
 
         async with httpx.AsyncClient() as client:
-            tasks = [self._send_abort_request(client, url, timeout=timeout) for url in self.worker_url_list]
+            tasks = []
+            for url in self.worker_url_list:
+                self.logger.info(f"Sending abort request to worker at {url}")
+                tasks.append(self._send_abort_request(client, url, timeout=timeout))
             results = await asyncio.gather(*tasks)
 
         failed_workers = [url for url, success in results if not success]
@@ -337,7 +347,7 @@ async def pause(self, timeout: float = 60.0):
                 f"Failed: {len(failed_workers)}. Failed workers: {failed_workers}"
             )
         else:
-            self.logger.debug(f"All {succeeded_count} abort requests sent successfully.")
+            self.logger.info(f"All {succeeded_count} abort requests sent successfully.")
 
     async def run(
         self,
@@ -397,7 +407,7 @@ async def _send_abort_request(self, client, url, timeout):
         try:
             response = await client.post(worker_url, json={"abort_all": True}, timeout=timeout)
             response.raise_for_status()
-            self.logger.debug(f"Successfully sent abort request to {url}")
+            self.logger.info(f"Successfully sent abort request to {url}")
             return url, True
         except Exception as e:
             self.logger.error(f"Failed to send abort request to {url}: {e}")

xtuner/v1/ray/dataflow/replay_buffer.py

@@ -117,7 +117,7 @@ def mapping_dataitem_to_replaymeta(grouped_dataitem: List[RLDataFlowItem]) -> Re
     group_state = determine_group_state(grouped_dataitem)
 
     replay_state = ReplayState.from_str(group_state)
-    logger.debug(f"determined group_state: {group_state}, replay_state: {replay_state}")
+    logger.info(f"determined group_state: {group_state}, replay_state: {replay_state}, version: {version}")
     replay_meta = ReplayMeta(
         env=env_str,
         root_id=root_id,
@@ -337,18 +337,19 @@ 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.info(f"Existing root_id: {root_id} found. Incrementing version to {replay_meta.version}.")
+            replay_meta.version += 1 if partial_rollout_step > 0 else 0
+            self.logger.info(f"Existing root_id: {root_id} with action_id {action_id} found. Incrementing version to {replay_meta.version}.")
             self._root2actions[root_id].append(action_id)
         else:
             self._root2actions[root_id] = [action_id]
+
         self._actions[action_id] = replay_meta
 
         # 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:
+        if replay_meta.state == ReplayState.INTERRUPTED and (replay_meta.version < partial_rollout_step or partial_rollout_step == 0):
             self._interrupted_actions[replay_meta.version].append(action_id)
             self.logger.info(
-                f"Add aborted sample with root_id: {root_id}, action_id: {action_id} to _interrupted_actions."
+                f"Add aborted sample with action_id: {action_id} version: {replay_meta.version} to _interrupted_actions."
             )
         elif replay_meta.state == ReplayState.INTERRUPTED and replay_meta.version >= partial_rollout_step:
             self._expired_actions.append(action_id)
@@ -359,7 +360,7 @@ def add(self, grouped_dataitem: List[RLDataFlowItem], partial_rollout_step: int
             )
         elif replay_meta.state == ReplayState.COMPLETED:
             self._completed_actions[replay_meta.version].append(action_id)
-            self.logger.debug(f"Add sample with root_id: {root_id}, action_id: {action_id} to finished_actions.")
+            self.logger.info(f"Add sample with root_id: {root_id}, action_id: {action_id} to finished_actions.")
         elif replay_meta.state == ReplayState.FAILED:
             assert False, "Currently, failed samples are not supported in the replay buffer."
 
@@ -541,7 +542,7 @@ def sample_from_expired_storage(self) -> List[RLDataFlowItem]:
 
         # update env for expired samples
         for sample in group_samples:
-            sample.data.input_ids = sample.data.input_ids[: sample.data.num_tokens]
+            # sample.data.input_ids = sample.data.input_ids[: sample.data.num_tokens]
             sample.env = RLEnvDataItem()
             sample.uid.version = 0
             sample.extra_info.state = str(ReplayState.INIT)
@@ -560,20 +561,20 @@ def sample_from_interrupted_storage(self, tokenizer) -> List[RLDataFlowItem]:
         # update env for interrupted samples
         for sample in group_samples:
             assert sample.data.input_ids and sample.data.num_tokens, "input_ids or num_tokens is empty!"
-            sample.data.input_ids = sample.data.input_ids[: sample.data.num_tokens]
+            # sample.data.input_ids = sample.data.input_ids[: sample.data.num_tokens]
             sample.uid.action_id = int(uuid4().int)
             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))
+            # 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))
         self.logger.info(
             f"Sampling interrupted action_id: {action_id} from replay buffer, remain interrupted samples: {self.get_interrupted_samples()}"
         )

xtuner/v1/ray/environment/single_turn_env.py

@@ -87,8 +87,16 @@ 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.info(
-                        f"Set num_return_tokens: {sample.env.rollout.num_return_tokens} for sample {sample.uid}."
+                    sample.data.extra_info["response_ids"] = sample.env.rollout.response_ids
+                    sample.data.extra_info["response"] = sample.env.rollout.response
+                    sample.data.extra_info["logprobs"] = sample.env.rollout.logprobs
+                    assert len(sample.env.rollout.response_ids) == len(sample.env.rollout.logprobs), (
+                        f"num_return_tokens {sample.env.rollout.num_return_tokens} mismatch "
+                        f"len of response_ids {len(sample.env.rollout.response_ids)} and "
+                        f"len of logprobs {len(sample.env.rollout.logprobs)} for sample {sample.uid}."
+                    )
+                    self.logger.debug(
+                        f"Set num_return_tokens: {sample.env.rollout.num_return_tokens} and len of response_ids {len(sample.env.rollout.response_ids)} for sample {sample.uid}."
                     )
                 fut = self.rollout_controller.rollout.remote(
                     prompt=sample.data.messages,

xtuner/v1/ray/rollout/controller.py

@@ -390,11 +390,6 @@ 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)
         try:
             response = await asyncio.wait_for(response_ref, timeout=self.config.rollout_timeout * 2)
             self.workers_info[server_url].success_count += 1

xtuner/v1/ray/rollout/lmdeploy.py

@@ -127,11 +127,12 @@ 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
+            init_input_len = len(input_ids) if input_ids else 0
+            payload["input_ids"] += extra_info["response_ids"]
             self.logger.info(
-                f"Set max_tokens to {max_return_tokens} based on num_return_tokens {extra_info['num_return_tokens']}"
+                f"Set max_tokens to {max_return_tokens} based on num_return_tokens {extra_info['num_return_tokens']}, init input_len: {init_input_len} and payload input len {len(payload['input_ids'])}."
             )
-
-        if self.enable_return_routed_experts:
+        if self.enable_return_routed_experts:   
             extra_params["return_routed_experts"] = True
 
         lmdeploy_sample_params = self._transform_sample_params(sample_params, extra_params)

xtuner/v1/ray/rollout/worker.py

@@ -1,4 +1,3 @@
-import asyncio
 import copy
 import json
 import multiprocessing
@@ -333,7 +332,13 @@ async def rollout_task(
             endpoint_url = f"{self.server_url}/{self.endpoints['v1/chat/completions']}"
 
         while True:
-            if extra_info.get("num_return_tokens", None) is not None and (sample_params["max_tokens"] - extra_info["num_return_tokens"]) == 0:
+            if (
+                extra_info.get("num_return_tokens", None) is not None
+                and (sample_params["max_tokens"] - extra_info["num_return_tokens"]) == 0
+            ):  
+                self.logger.info(
+                    f"rollout request {uid} reached max tokens {sample_params['max_tokens']}, returning length finish_reason"
+                )
                 return RLRolloutResponseItem(
                     response="",
                     response_ids=[],
@@ -484,6 +489,9 @@ async def _handle_non_stream_response(self, uid, sample_params, extra_params, re
                         routed_experts = ray.put(routed_experts)
                     extra_info = {"routed_experts": routed_experts}
 
+                if finish_reason != "abort" and len(last_token_ids) == 0:
+                    self.logger.error(f"rollout request {uid} returned zero tokens with finish_reason {finish_reason}")
+                    
                 rollout_response = RLRolloutResponseItem(
                     response=response["text"],
                     response_ids=last_token_ids if len(last_token_ids) > 0 else None,