feat: validate rollout server-group GPU placement (slime #1934, #1944)

.github/workflows/pr-test.yml

@@ -476,7 +476,7 @@ jobs:
     strategy:
       fail-fast: false
       matrix:
-        info: [{"num_gpus": 0, "test_file": "test_megatron_argument_validation.py"}, {"num_gpus": 0, "test_file": "plugin_contracts/test_plugin_rollout_contracts.py"}, {"num_gpus": 0, "test_file": "plugin_contracts/test_plugin_runtime_hook_contracts.py"}, {"num_gpus": 0, "test_file": "plugin_contracts/test_plugin_path_loading_contracts.py"}, {"num_gpus": 0, "test_file": "plugin_contracts/test_plugin_generate_contracts.py"}]
+        info: [{"num_gpus": 0, "test_file": "test_megatron_argument_validation.py"}, {"num_gpus": 0, "test_file": "test_rollout_validation.py"}, {"num_gpus": 0, "test_file": "plugin_contracts/test_plugin_rollout_contracts.py"}, {"num_gpus": 0, "test_file": "plugin_contracts/test_plugin_runtime_hook_contracts.py"}, {"num_gpus": 0, "test_file": "plugin_contracts/test_plugin_path_loading_contracts.py"}, {"num_gpus": 0, "test_file": "plugin_contracts/test_plugin_generate_contracts.py"}]
     defaults:
       run:
         working-directory: ${{ github.workspace }}

.github/workflows/pr-test.yml.j2

@@ -55,6 +55,7 @@
       'cpu': True,
       'tests': [
         {'test_file': 'test_megatron_argument_validation.py', 'num_gpus': 0},
+        {'test_file': 'test_rollout_validation.py', 'num_gpus': 0},
         {'test_file': 'plugin_contracts/test_plugin_rollout_contracts.py', 'num_gpus': 0},
         {'test_file': 'plugin_contracts/test_plugin_runtime_hook_contracts.py', 'num_gpus': 0},
         {'test_file': 'plugin_contracts/test_plugin_path_loading_contracts.py', 'num_gpus': 0},

tests/test_rollout_validation.py

@@ -0,0 +1,60 @@
+import pytest
+
+from vime.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
+
+
+if __name__ == "__main__":
+    raise SystemExit(pytest.main([__file__]))

vime/ray/rollout.py

@@ -30,6 +30,7 @@
 from vime.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)
@@ -115,6 +116,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,
+        )
 
         from vime.backends.vllm_utils.vllm_engine import VLLMEngine
 
@@ -1053,7 +1064,9 @@ def start_rollout_servers(args, pg) -> dict[str, RolloutServer]:
         model_cfg.resolve(args)
 
         has_pd = model_cfg.has_pd_disaggregation
-        router_ip, router_port, prom_port = _start_router(args, has_pd_disaggregation=has_pd, force_new=(model_idx > 0))
+        router_ip, router_port, prom_port = _start_router(
+            args, has_pd_disaggregation=has_pd, force_new=(model_idx > 0)
+        )
 
         # Write back so downstream readers (vllm_rollout, vllm_engine) see the
         # router we just started (only relevant for first model in multi-model setups).

vime/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 and --rollout-num-gpus-per-engine "
+        "with the rollout engine GPU placement."
+    )