fix comments

Author: YanhuiDua

xtuner/v1/ray/base/accelerator.py

@@ -404,7 +404,10 @@ def from_placement_group(cls, worker_cls, worker_config, pg: PlacementGroup):
         rank_bundle_idx_list = []
         for rank, bundle_idx in enumerate(sorted_bundle_idxs):
             worker = worker_cls.options(
-                placement_group=pg, placement_group_bundle_index=bundle_idx, **pg_options
+                max_concurrency=int(os.environ.get("RAY_MAX_CONCURRENCY", 1000)),
+                placement_group=pg,
+                placement_group_bundle_index=bundle_idx,
+                **pg_options,
             ).remote(worker_config, rank, master_addr, master_port, world_size, device_type)
             workers_list.append(worker)
             rank_bundle_idx_list.append((rank, bundle_idx))

xtuner/v1/ray/dataflow/flow.py

@@ -236,8 +236,8 @@ async def concurrent_task_runner(self):
             next_update_threshold = update_step
             while (
                 self.finished_samples_count < self.target_batch_size
-                and self.failed_samples_count < self.target_batch_size
-                and self.skipped_sample_count < self.target_batch_size
+                and self.failed_samples_count < self.target_batch_size * self.config.max_retry_times
+                and self.skipped_sample_count < self.target_batch_size * self.config.max_retry_times
             ):
                 if self.finished_samples_count >= next_update_threshold:
                     pbar.n = self.finished_samples_count
@@ -279,8 +279,10 @@ async def concurrent_task_runner(self):
 
         if self.finished_samples_count >= self.target_batch_size:
             self.logger.info("Target batch size reached. Pausing env controller.")
-        if self.failed_samples_count >= self.target_batch_size or self.skipped_sample_count >= self.target_batch_size:
+        if self.failed_samples_count >= self.target_batch_size * self.config.max_retry_times:
             self.logger.info("Max failed samples reached. Pausing env controller.")
+        if self.skipped_sample_count >= self.target_batch_size * self.config.max_retry_times:
+            self.logger.info("Max skipped samples reached. Pausing env controller.")
 
         # NOTE: Directly send pause requests to rollout workers because calling `rollout_controller.pause()`
         # would be queued behind many worker tasks, causing a significant delay.

xtuner/v1/ray/environment/base_env.py

@@ -1,3 +1,4 @@
+import os
 from abc import ABC, abstractmethod
 from typing import Any, List
 
@@ -70,7 +71,11 @@ def init_rollout_controller(self, rollout_cfg: Any, placement_group: Any):
 
         from xtuner.v1.ray.rollout.controller import RolloutController
 
-        rollout_controller = ray.remote(RolloutController).remote(rollout_cfg, placement_group)  # type: ignore[attr-defined]
+        rollout_controller = (
+            ray.remote(RolloutController)
+            .options(max_concurrency=int(os.environ.get("RAY_MAX_CONCURRENCY", 1000)))
+            .remote(rollout_cfg, placement_group)
+        )  # type: ignore[attr-defined]
         return rollout_controller
 
     def init_judger_controller(self, judger_cfg: Any, placement_group: Any):

xtuner/v1/ray/rollout/controller.py

@@ -331,6 +331,7 @@ async def rollout(
                 url = response.extra_info["url"]
                 if response.finish_reason == "failed":
                     self.deactivate_worker_by_url(url)
+                response.extra_info.pop("url", None)
             return response
         except asyncio.TimeoutError:
             self.logger.error("Get response from rollout worker timeout and return the failed response.")

xtuner/v1/ray/rollout/sglang.py

@@ -1,7 +1,6 @@
 import os
 from typing import Any, Dict, List, Union
 
-import ray
 import requests
 from urllib3.exceptions import NewConnectionError
 
@@ -11,7 +10,6 @@
 from .worker import RolloutWorker
 
 
-@ray.remote(max_concurrency=int(os.environ.get("RAY_MAX_CONCURRENCY", 1000)))
 class SGLangWorker(RolloutWorker):
     def __init__(
         self,

xtuner/v1/ray/rollout/vllm.py

@@ -1,7 +1,6 @@
 from argparse import Namespace
 from typing import Any, Dict, List, Union
 
-import ray
 import uvloop
 from vllm.entrypoints.openai.api_server import run_server
 from vllm.entrypoints.openai.cli_args import make_arg_parser
@@ -16,7 +15,6 @@ def run_vllm_server_wrapper(server_args):
     uvloop.run(run_server(server_args))
 
 
-@ray.remote
 class vLLMWorker(RolloutWorker):
     def __init__(
         self,