Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
34 changes: 33 additions & 1 deletion tests/rl/test_rollout_logic.py
Original file line number Diff line number Diff line change
Expand Up @@ -21,6 +21,7 @@

from xtuner.v1.data_proto.rl_data import RolloutState, SampleParams, Status
from xtuner.v1.rl.agent_loop import AgentLoopConfig
from xtuner.v1.rl.rollout.constants import ROLLOUT_RAY_GENERATE_MAX_CONCURRENCY
from xtuner.v1.rl.rollout.controller import RolloutController
from xtuner.v1.rl.rollout.health_manager import RolloutHealthManager
from xtuner.v1.rl.rollout.lmdeploy import LMDeployWorker
Expand All @@ -33,7 +34,13 @@
)
from xtuner.v1.rl.rollout.sglang import SGLangWorker
from xtuner.v1.rl.rollout.utils import PartialRolloutHandler, SessionRouter
from xtuner.v1.rl.rollout.worker import RolloutWorker, RolloutWorkerInitResult
from xtuner.v1.rl.rollout.worker import (
ROLLOUT_CONCURRENCY_GROUP_CONTROL,
ROLLOUT_CONCURRENCY_GROUP_GENERATE,
RolloutConfig,
RolloutWorker,
RolloutWorkerInitResult,
)
from xtuner.v1.rl.utils.misc import delete_from_routedapiproxy
from xtuner.v1.rl.weight_update.data import RolloutWeightUpdateInfo
from xtuner.v1.train.rl_trainer import BaseRLTrainer, _agent_loop_manager_requires_rollout_proxy
Expand Down Expand Up @@ -316,6 +323,31 @@ def test_sglang_tp16_cross_node_weight_update_targets_match_legacy_mesh_and_url_


class TestRolloutController(unittest.IsolatedAsyncioTestCase):
def test_pause_generation_uses_dedicated_control_concurrency_group(self):
self.assertEqual(
getattr(RolloutController.pause_generation, "__ray_concurrency_group__", None),
ROLLOUT_CONCURRENCY_GROUP_CONTROL,
)

def test_rollout_config_build_registers_control_concurrency_group(self):
remote_actor_cls = MagicMock()
remote_actor_cls.options.return_value.remote.return_value = "rollout-controller"

with (
patch("xtuner.v1.rl.rollout.worker.register_cpu_resources"),
patch("xtuner.v1.rl.trace.get_trace_env_vars", return_value={}),
patch("ray.remote", return_value=lambda _: remote_actor_cls) as ray_remote,
):
result = RolloutConfig.build(MagicMock(), placement_group=object())

self.assertEqual(result, "rollout-controller")
ray_remote.assert_called_once_with(
concurrency_groups={
ROLLOUT_CONCURRENCY_GROUP_GENERATE: ROLLOUT_RAY_GENERATE_MAX_CONCURRENCY,
ROLLOUT_CONCURRENCY_GROUP_CONTROL: 1,
},
)

def _state(self, uid: int, session_id: int) -> RolloutState:
return RolloutState(
rollout_id=uid,
Expand Down
2 changes: 2 additions & 0 deletions xtuner/v1/rl/rollout/controller.py
Original file line number Diff line number Diff line change
Expand Up @@ -18,6 +18,7 @@
from .rollout_topology import RolloutTopology
from .utils import SessionRouter
from .worker import (
ROLLOUT_CONCURRENCY_GROUP_CONTROL,
ROLLOUT_CONCURRENCY_GROUP_GENERATE,
RolloutConfig,
get_rollout_worker_base_cls,
Expand Down Expand Up @@ -131,6 +132,7 @@ def set_enable_partial_rollout(self, enable: bool) -> None:
]
)

@ray.method(concurrency_group=ROLLOUT_CONCURRENCY_GROUP_CONTROL)
def pause_generation(self):
self.health_manager.pause()
active_workers = self.registry.active_workers()
Expand Down
2 changes: 2 additions & 0 deletions xtuner/v1/rl/rollout/worker.py
Original file line number Diff line number Diff line change
Expand Up @@ -54,6 +54,7 @@

infer_group = Group("inference", help="Inference worker configuration.")
ROLLOUT_CONCURRENCY_GROUP_GENERATE = "generate"
ROLLOUT_CONCURRENCY_GROUP_CONTROL = "control"


@dataclass(frozen=True)
Expand Down Expand Up @@ -500,6 +501,7 @@ def build(self, placement_group: "PlacementGroup"):
ray.remote(
concurrency_groups={
ROLLOUT_CONCURRENCY_GROUP_GENERATE: ROLLOUT_RAY_GENERATE_MAX_CONCURRENCY,
ROLLOUT_CONCURRENCY_GROUP_CONTROL: 1,
},
)(RolloutController)
.options(**actor_options)
Expand Down