diff --git a/.github/workflows/pr-test.yml b/.github/workflows/pr-test.yml index 5d9db6a23..d88c34097 100644 --- a/.github/workflows/pr-test.yml +++ b/.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 }} diff --git a/.github/workflows/pr-test.yml.j2 b/.github/workflows/pr-test.yml.j2 index def8dc6bf..38a48bf2f 100644 --- a/.github/workflows/pr-test.yml.j2 +++ b/.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}, diff --git a/tests/test_rollout_validation.py b/tests/test_rollout_validation.py new file mode 100644 index 000000000..ef710c72a --- /dev/null +++ b/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__])) diff --git a/vime/ray/rollout.py b/vime/ray/rollout.py index 60b21a746..398a6791a 100644 --- a/vime/ray/rollout.py +++ b/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). diff --git a/vime/ray/rollout_validation.py b/vime/ray/rollout_validation.py new file mode 100644 index 000000000..d6618e293 --- /dev/null +++ b/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." + )