From f109df6b81b41bcf7deabcfb1c900f713aea3f9c Mon Sep 17 00:00:00 2001 From: fmh Date: Thu, 21 May 2026 17:22:21 +0800 Subject: [PATCH] rollout: validate GPU placement before engine startup --- slime/ray/rollout.py | 11 +++++++ slime/ray/rollout_validation.py | 32 ++++++++++++++++++ tests/test_rollout_validation.py | 56 ++++++++++++++++++++++++++++++++ 3 files changed, 99 insertions(+) create mode 100644 slime/ray/rollout_validation.py create mode 100644 tests/test_rollout_validation.py diff --git a/slime/ray/rollout.py b/slime/ray/rollout.py index b71a633f9c..a2d7d9bccc 100644 --- a/slime/ray/rollout.py +++ b/slime/ray/rollout.py @@ -27,6 +27,7 @@ from slime.utils.types import Sample from ..utils.metric_utils import has_repetition +from .rollout_validation import validate_server_group_gpu_indices from .utils import NOSET_VISIBLE_DEVICES_ENV_VARS_LIST, Lock logging.getLogger("httpx").setLevel(logging.WARNING) @@ -87,6 +88,16 @@ def start_engines(self, port_cursors: dict[int, int] | None = None) -> tuple[lis num_gpu_per_engine = min(self.num_gpus_per_engine, self.args.num_gpus_per_node) pg, reordered_bundle_indices, reordered_gpu_ids = self.pg + validate_server_group_gpu_indices( + worker_type=self.worker_type, + gpu_offset=self.gpu_offset, + num_gpus_per_engine=self.num_gpus_per_engine, + num_gpu_per_engine=num_gpu_per_engine, + num_engines=len(self.all_engines), + num_available_gpus=len(reordered_gpu_ids), + rollout_num_gpus=self.args.rollout_num_gpus, + rollout_num_gpus_per_engine=self.args.rollout_num_gpus_per_engine, + ) RolloutRayActor = ray.remote(SGLangEngine) diff --git a/slime/ray/rollout_validation.py b/slime/ray/rollout_validation.py new file mode 100644 index 0000000000..46ec9f892c --- /dev/null +++ b/slime/ray/rollout_validation.py @@ -0,0 +1,32 @@ +def validate_server_group_gpu_indices( + *, + worker_type: str, + gpu_offset: int, + num_gpus_per_engine: int, + num_gpu_per_engine: int, + num_engines: int, + num_available_gpus: int, + rollout_num_gpus: int, + rollout_num_gpus_per_engine: int, +) -> None: + if num_engines == 0: + return + + required_gpu_slots = gpu_offset + num_engines * num_gpu_per_engine + if gpu_offset >= 0 and num_gpu_per_engine > 0 and required_gpu_slots <= num_available_gpus: + return + + raise ValueError( + "Invalid rollout server group GPU placement: " + f"worker_type={worker_type}, " + f"gpu_offset={gpu_offset}, " + f"num_gpus_per_engine={num_gpus_per_engine}, " + f"num_gpu_per_engine_on_node={num_gpu_per_engine}, " + f"num_engines={num_engines}, " + f"required_gpu_slots={required_gpu_slots}, " + f"len(reordered_gpu_ids)={num_available_gpus}, " + f"rollout_num_gpus={rollout_num_gpus}, " + f"rollout_num_gpus_per_engine={rollout_num_gpus_per_engine}. " + "Please align --rollout-num-gpus, --rollout-num-gpus-per-engine, " + "and --sglang-config server_groups." + ) diff --git a/tests/test_rollout_validation.py b/tests/test_rollout_validation.py new file mode 100644 index 0000000000..28c6954b1e --- /dev/null +++ b/tests/test_rollout_validation.py @@ -0,0 +1,56 @@ +import pytest + +from slime.ray.rollout_validation import validate_server_group_gpu_indices + + +@pytest.mark.unit +def test_validate_server_group_gpu_indices_accepts_valid_config(): + validate_server_group_gpu_indices( + worker_type="regular", + gpu_offset=2, + num_gpus_per_engine=1, + num_gpu_per_engine=1, + num_engines=2, + num_available_gpus=4, + rollout_num_gpus=4, + rollout_num_gpus_per_engine=1, + ) + + +@pytest.mark.unit +def test_validate_server_group_gpu_indices_allows_empty_group(): + validate_server_group_gpu_indices( + worker_type="placeholder", + gpu_offset=4, + num_gpus_per_engine=1, + num_gpu_per_engine=1, + num_engines=0, + num_available_gpus=4, + rollout_num_gpus=4, + rollout_num_gpus_per_engine=1, + ) + + +@pytest.mark.unit +def test_validate_server_group_gpu_indices_reports_config_context(): + with pytest.raises(ValueError) as exc_info: + validate_server_group_gpu_indices( + worker_type="regular", + gpu_offset=3, + num_gpus_per_engine=2, + num_gpu_per_engine=2, + num_engines=1, + num_available_gpus=4, + rollout_num_gpus=4, + rollout_num_gpus_per_engine=2, + ) + + message = str(exc_info.value) + assert "worker_type=regular" in message + assert "gpu_offset=3" in message + assert "num_gpus_per_engine=2" in message + assert "num_engines=1" in message + assert "required_gpu_slots=5" in message + assert "len(reordered_gpu_ids)=4" in message + assert "rollout_num_gpus=4" in message + assert "rollout_num_gpus_per_engine=2" in message