| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | import json |
| 5 | from argparse import ArgumentError |
| 6 | from contextlib import AbstractContextManager, nullcontext |
| 7 | from typing import Annotated, Literal |
| 8 | |
| 9 | import pytest |
| 10 | from pydantic import Field |
| 11 | |
| 12 | from vllm.config import ( |
| 13 | AttentionConfig, |
| 14 | CacheConfig, |
| 15 | CompilationConfig, |
| 16 | ModelConfig, |
| 17 | config, |
| 18 | ) |
| 19 | from vllm.engine.arg_utils import ( |
| 20 | EngineArgs, |
| 21 | _expand_json_human_readable_numbers, |
| 22 | contains_type, |
| 23 | get_kwargs, |
| 24 | get_type, |
| 25 | get_type_hints, |
| 26 | is_not_builtin, |
| 27 | is_type, |
| 28 | literal_to_kwargs, |
| 29 | optional_type, |
| 30 | parse_type, |
| 31 | ) |
| 32 | from vllm.utils.argparse_utils import FlexibleArgumentParser |
| 33 | |
| 34 | |
| 35 | @pytest.mark.parametrize( |
| 36 | ("type", "value", "expected"), |
| 37 | [ |
| 38 | (int, "42", 42), |
| 39 | (float, "3.14", 3.14), |
| 40 | (str, "Hello World!", "Hello World!"), |
| 41 | (json.loads, '{"foo":1,"bar":2}', {"foo": 1, "bar": 2}), |
| 42 | ], |
| 43 | ) |
| 44 | def test_parse_type(type, value, expected): |
| 45 | parse_type_func = parse_type(type) |
| 46 | assert parse_type_func(value) == expected |
| 47 | |
| 48 | |
| 49 | def test_optional_type(): |
| 50 | optional_type_func = optional_type(int) |
| 51 | assert optional_type_func("None") is None |
| 52 | assert optional_type_func("42") == 42 |
| 53 | |
| 54 | |
| 55 | def test_watermark_config_cli(): |
| 56 | parser = EngineArgs.add_cli_args(FlexibleArgumentParser()) |
| 57 | args = parser.parse_args( |
| 58 | [ |
| 59 | "--model", |
| 60 | "dummy", |
| 61 | "--watermark-config", |
| 62 | '{"algorithm":"dual_key_gumbel","key":42,"prf":"philox","alpha":0.25,' |
| 63 | '"allow_target_only_watermarking":true}', |
| 64 | ] |
| 65 | ) |
| 66 | |
| 67 | config = EngineArgs.from_cli_args(args).create_watermark_config() |
| 68 | |
| 69 | assert config is not None |
| 70 | assert config.algorithm == "dual_key_gumbel" |
| 71 | assert config.key == 42 |
| 72 | assert config.alpha == 0.25 |
| 73 | assert config.context_width == 4 |
| 74 | assert config.deduplicate_contexts == "single_turn" |
| 75 | assert config.deduplicate_contexts_max_history == 8192 |
| 76 | assert config.prf == "philox" |
| 77 | assert config.allow_target_only_watermarking |
| 78 | |
| 79 | args = parser.parse_args( |
| 80 | [ |
| 81 | "--model", |
| 82 | "dummy", |
| 83 | "--watermark-config", |
| 84 | '{"key":42,"deduplicate_contexts":"none",' |
| 85 | '"deduplicate_contexts_max_history":32}', |
| 86 | ] |
| 87 | ) |
| 88 | config = EngineArgs.from_cli_args(args).create_watermark_config() |
| 89 | |
| 90 | assert config is not None |
| 91 | assert config.deduplicate_contexts == "none" |
| 92 | assert config.deduplicate_contexts_max_history == 32 |
| 93 | |
| 94 | |
| 95 | @pytest.mark.parametrize( |
| 96 | "option", |
| 97 | ["--gpu-memory-utilization", "--device-memory-utilization"], |
| 98 | ) |
| 99 | def test_memory_utilization_cli_aliases(option): |
| 100 | parser = EngineArgs.add_cli_args(FlexibleArgumentParser()) |
| 101 | args = EngineArgs.from_cli_args(parser.parse_args([option, "0.8"])) |
| 102 | |
| 103 | assert args.gpu_memory_utilization == 0.8 |
| 104 | |
| 105 | |
| 106 | def test_device_memory_utilization_property(): |
| 107 | config = CacheConfig(gpu_memory_utilization=0.8) |
| 108 | |
| 109 | assert config.device_memory_utilization == 0.8 |
| 110 | |
| 111 | config.device_memory_utilization = 0.7 |
| 112 | assert config.gpu_memory_utilization == 0.7 |
| 113 | |
| 114 | |
| 115 | @pytest.mark.parametrize( |
| 116 | "options", |
| 117 | [ |
| 118 | [ |
| 119 | "--engram-config", |
| 120 | '{"cpu_offload": false, "embedding_across_dp": true}', |
| 121 | ], |
| 122 | [ |
| 123 | "--engram-config.cpu_offload", |
| 124 | "false", |
| 125 | "--engram-config.embedding_across_dp", |
| 126 | "true", |
| 127 | ], |
| 128 | ], |
| 129 | ) |
| 130 | def test_engram_config_cli(options): |
| 131 | """JSON and dotted CLI options independently control Engram settings.""" |
| 132 | parser = EngineArgs.add_cli_args(FlexibleArgumentParser()) |
| 133 | args = EngineArgs.from_cli_args(parser.parse_args(options)) |
| 134 | assert args.engram_config is not None |
| 135 | assert args.engram_config.cpu_offload is False |
| 136 | assert args.engram_config.embedding_across_dp is True |
| 137 | |
| 138 | |
| 139 | @pytest.mark.parametrize( |
| 140 | "options,provided,dp_shared_memory", |
| 141 | [ |
| 142 | ([], False, False), |
| 143 | (["--engram-config", "{}"], True, None), |
| 144 | ( |
| 145 | ["--engram-config", '{"dp_shared_memory": true}'], |
| 146 | True, |
| 147 | True, |
| 148 | ), |
| 149 | (["--engram-config.dp_shared_memory", "true"], True, True), |
| 150 | ], |
| 151 | ) |
| 152 | def test_engram_config_cli_optional(options, provided, dp_shared_memory): |
| 153 | """Explicit configs honor defaults and the JSON/dotted DP shared-memory flag.""" |
| 154 | parser = EngineArgs.add_cli_args(FlexibleArgumentParser()) |
| 155 | args = EngineArgs.from_cli_args(parser.parse_args(options)) |
| 156 | assert (args.engram_config is not None) == provided |
| 157 | if provided: |
| 158 | assert args.engram_config.cpu_offload is True |
| 159 | assert args.engram_config.dp_shared_memory is dp_shared_memory |
| 160 | |
| 161 | |
| 162 | @pytest.mark.parametrize( |
| 163 | ("type_hint", "type", "expected"), |
| 164 | [ |
| 165 | (int, int, True), |
| 166 | (int, float, False), |
| 167 | (list[int], list, True), |
| 168 | (list[int], tuple, False), |
| 169 | (Literal[0, 1], Literal, True), |
| 170 | ], |
| 171 | ) |
| 172 | def test_is_type(type_hint, type, expected): |
| 173 | assert is_type(type_hint, type) == expected |
| 174 | |
| 175 | |
| 176 | @pytest.mark.parametrize( |
| 177 | ("type_hints", "type", "expected"), |
| 178 | [ |
| 179 | ({float, int}, int, True), |
| 180 | ({int, tuple}, int, True), |
| 181 | ({int, tuple[int]}, int, True), |
| 182 | ({int, tuple[int, ...]}, int, True), |
| 183 | ({int, tuple[int]}, float, False), |
| 184 | ({int, tuple[int, ...]}, float, False), |
| 185 | ({str, Literal["x", "y"]}, Literal, True), |
| 186 | ], |
| 187 | ) |
| 188 | def test_contains_type(type_hints, type, expected): |
| 189 | assert contains_type(type_hints, type) == expected |
| 190 | |
| 191 | |
| 192 | @pytest.mark.parametrize( |
| 193 | ("type_hints", "type", "expected"), |
| 194 | [ |
| 195 | ({int, float}, int, int), |
| 196 | ({int, float}, str, None), |
| 197 | ({str, Literal["x", "y"]}, Literal, Literal["x", "y"]), |
| 198 | ], |
| 199 | ) |
| 200 | def test_get_type(type_hints, type, expected): |
| 201 | assert get_type(type_hints, type) == expected |
| 202 | |
| 203 | |
| 204 | @pytest.mark.parametrize( |
| 205 | ("type_hints", "expected"), |
| 206 | [ |
| 207 | ({Literal[1, 2]}, {"type": int, "choices": [1, 2]}), |
| 208 | ({str, Literal["x", "y"]}, {"type": str, "metavar": ["x", "y"]}), |
| 209 | ({Literal[1, "a"]}, Exception), |
| 210 | ], |
| 211 | ) |
| 212 | def test_literal_to_kwargs(type_hints, expected): |
| 213 | context: AbstractContextManager[object] = nullcontext() |
| 214 | if expected is Exception: |
| 215 | context = pytest.raises(expected) |
| 216 | with context: |
| 217 | assert literal_to_kwargs(type_hints) == expected |
| 218 | |
| 219 | |
| 220 | @config |
| 221 | class NestedConfig: |
| 222 | field: int = 1 |
| 223 | """field""" |
| 224 | |
| 225 | |
| 226 | @config |
| 227 | class DummyConfig: |
| 228 | regular_bool: bool = True |
| 229 | """Regular bool with default True""" |
| 230 | optional_bool: bool | None = None |
| 231 | """Optional bool with default None""" |
| 232 | |
| 233 | optional_bool_or_str: bool | str | None = None |
| 234 | """Optional bool-or-str with default None""" |
| 235 | |
| 236 | optional_literal: Literal["x", "y"] | None = None |
| 237 | """Optional literal with default None""" |
| 238 | tuple_n: tuple[int, ...] = Field(default_factory=lambda: (1, 2, 3)) |
| 239 | """Tuple with variable length""" |
| 240 | tuple_2: tuple[int, int] = Field(default_factory=lambda: (1, 2)) |
| 241 | """Tuple with fixed length""" |
| 242 | list_n: list[int] = Field(default_factory=lambda: [1, 2, 3]) |
| 243 | """List with variable length""" |
| 244 | list_literal: list[Literal[1, 2]] = Field(default_factory=list) |
| 245 | """List with literal choices""" |
| 246 | list_union: list[str | type[object]] = Field(default_factory=list) |
| 247 | """List with union type""" |
| 248 | set_n: set[int] = Field(default_factory=lambda: {1, 2, 3}) |
| 249 | """Set with variable length""" |
| 250 | literal_literal: Literal[Literal[1], Literal[2]] = 1 |
| 251 | """Literal of literals with default 1""" |
| 252 | json_tip: dict = Field(default_factory=dict) |
| 253 | """Dict which will be JSON in CLI""" |
| 254 | nested_config: NestedConfig = Field(default_factory=NestedConfig) |
| 255 | """Nested config""" |
| 256 | |
| 257 | |
| 258 | @pytest.mark.parametrize( |
| 259 | ("type_hint", "expected"), |
| 260 | [ |
| 261 | (int, False), |
| 262 | (DummyConfig, True), |
| 263 | ], |
| 264 | ) |
| 265 | def test_is_not_builtin(type_hint, expected): |
| 266 | assert is_not_builtin(type_hint) == expected |
| 267 | |
| 268 | |
| 269 | @pytest.mark.parametrize( |
| 270 | ("type_hint", "expected"), |
| 271 | [ |
| 272 | (Annotated[int, "annotation"], {int}), |
| 273 | (int | None, {int, type(None)}), |
| 274 | (Annotated[int | None, "annotation"], {int, type(None)}), |
| 275 | (Annotated[int, "annotation"] | None, {int, type(None)}), |
| 276 | ], |
| 277 | ids=["Annotated", "or_None", "Annotated_or_None", "or_None_Annotated"], |
| 278 | ) |
| 279 | def test_get_type_hints(type_hint, expected): |
| 280 | assert get_type_hints(type_hint) == expected |
| 281 | |
| 282 | |
| 283 | @pytest.fixture |
| 284 | def dummy_config_kwargs(): |
| 285 | return get_kwargs(DummyConfig) |
| 286 | |
| 287 | |
| 288 | def test_get_kwargs(dummy_config_kwargs): |
| 289 | kwargs = dummy_config_kwargs |
| 290 | # bools should not have their type set |
| 291 | assert kwargs["regular_bool"].get("type") is None |
| 292 | assert kwargs["optional_bool"].get("type") is None |
| 293 | # optional bool-or-str should accept an optional string value |
| 294 | assert kwargs["optional_bool_or_str"]["type"] is str |
| 295 | assert kwargs["optional_bool_or_str"]["nargs"] == "?" |
| 296 | assert kwargs["optional_bool_or_str"]["const"] is True |
| 297 | assert "action" not in kwargs["optional_bool_or_str"] |
| 298 | # optional literals should have None as a choice |
| 299 | assert kwargs["optional_literal"]["choices"] == ["x", "y", None] |
| 300 | # tuples should have the correct nargs |
| 301 | assert kwargs["tuple_n"]["nargs"] == "+" |
| 302 | assert kwargs["tuple_2"]["nargs"] == 2 |
| 303 | # lists should work |
| 304 | assert kwargs["list_n"]["type"] is int |
| 305 | assert kwargs["list_n"]["nargs"] == "+" |
| 306 | # lists with literals should have the correct choices |
| 307 | assert kwargs["list_literal"]["type"] is int |
| 308 | assert kwargs["list_literal"]["nargs"] == "+" |
| 309 | assert kwargs["list_literal"]["choices"] == [1, 2] |
| 310 | # lists with unions should become str type. |
| 311 | # If not, we cannot know which type to use for parsing |
| 312 | assert kwargs["list_union"]["type"] is str |
| 313 | # sets should work like lists |
| 314 | assert kwargs["set_n"]["type"] is int |
| 315 | assert kwargs["set_n"]["nargs"] == "+" |
| 316 | # literals of literals should have merged choices |
| 317 | assert kwargs["literal_literal"]["choices"] == [1, 2] |
| 318 | # dict should have json tip in help |
| 319 | json_tip = "Should either be a valid JSON string or JSON keys" |
| 320 | assert json_tip in kwargs["json_tip"]["help"] |
| 321 | # nested config should construct the nested config |
| 322 | assert kwargs["nested_config"]["type"]('{"field": 2}') == NestedConfig(2) # type: ignore[call-arg] |
| 323 | |
| 324 | |
| 325 | @pytest.mark.parametrize( |
| 326 | ("args", "expected"), |
| 327 | [ |
| 328 | (["--optional-literal", "None"], None), |
| 329 | (["--optional-literal", ""], None), |
| 330 | (["--optional-literal", "x"], "x"), |
| 331 | ], |
| 332 | ) |
| 333 | def test_optional_handling(args, expected, dummy_config_kwargs): |
| 334 | parser = FlexibleArgumentParser() |
| 335 | parser.add_argument("--optional-literal", **dummy_config_kwargs["optional_literal"]) |
| 336 | |
| 337 | assert parser.parse_args(args).optional_literal is expected |
| 338 | assert "None" in parser.format_help() |
| 339 | |
| 340 | |
| 341 | def test_jit_monitor_verbose_arg(): |
| 342 | parser = EngineArgs.add_cli_args(FlexibleArgumentParser()) |
| 343 | args = parser.parse_args(["--jit-monitor-verbose"]) |
| 344 | |
| 345 | assert args.jit_monitor_verbose |
| 346 | assert EngineArgs(model="test", jit_monitor_verbose=True).jit_monitor_verbose |
| 347 | |
| 348 | |
| 349 | @pytest.mark.parametrize("mode", ["warn", "error"]) |
| 350 | def test_jit_monitor_mode_arg(mode): |
| 351 | parser = EngineArgs.add_cli_args(FlexibleArgumentParser()) |
| 352 | args = parser.parse_args(["--jit-monitor-mode", mode]) |
| 353 | |
| 354 | assert args.jit_monitor_mode == mode |
| 355 | engine_args = EngineArgs(model="test", jit_monitor_mode=mode) |
| 356 | assert engine_args.jit_monitor_mode == mode |
| 357 | assert engine_args.create_observability_config().jit_monitor_mode == mode |
| 358 | |
| 359 | |
| 360 | def test_hf_token_get_kwargs(): |
| 361 | kwargs = get_kwargs(ModelConfig)["hf_token"] |
| 362 | |
| 363 | assert kwargs["type"] is str |
| 364 | assert kwargs["nargs"] == "?" |
| 365 | assert kwargs["const"] is True |
| 366 | assert "action" not in kwargs |
| 367 | |
| 368 | |
| 369 | @pytest.mark.parametrize( |
| 370 | ("cli_args", "expected"), |
| 371 | [ |
| 372 | ([], None), |
| 373 | (["--hf-token"], True), |
| 374 | (["--hf-token", "hf_secret"], "hf_secret"), |
| 375 | (["--hf-token", "None"], "None"), |
| 376 | ], |
| 377 | ) |
| 378 | def test_hf_token_cli_arg(cli_args, expected): |
| 379 | parser = EngineArgs.add_cli_args(FlexibleArgumentParser()) |
| 380 | |
| 381 | args = parser.parse_args(cli_args) |
| 382 | |
| 383 | assert args.hf_token == expected |
| 384 | |
| 385 | |
| 386 | @pytest.mark.parametrize( |
| 387 | ("arg", "expected"), |
| 388 | [ |
| 389 | (None, dict()), |
| 390 | ('{"video": {"num_frames": 123} }', {"video": {"num_frames": 123}}), |
| 391 | ( |
| 392 | '{"video": {"num_frames": 123, "fps": 1.0, "foo": "bar"}, "image": {"foo": "bar"} }', # noqa |
| 393 | { |
| 394 | "video": {"num_frames": 123, "fps": 1.0, "foo": "bar"}, |
| 395 | "image": {"foo": "bar"}, |
| 396 | }, |
| 397 | ), |
| 398 | ], |
| 399 | ) |
| 400 | def test_media_io_kwargs_parser(arg, expected): |
| 401 | parser = EngineArgs.add_cli_args(FlexibleArgumentParser()) |
| 402 | if arg is None: |
| 403 | args = parser.parse_args([]) |
| 404 | else: |
| 405 | args = parser.parse_args(["--media-io-kwargs", arg]) |
| 406 | |
| 407 | assert args.media_io_kwargs == expected |
| 408 | |
| 409 | |
| 410 | @pytest.mark.parametrize( |
| 411 | ("args", "expected"), |
| 412 | [ |
| 413 | (["-O", "1"], "1"), |
| 414 | (["-O", "2"], "2"), |
| 415 | (["-O", "3"], "3"), |
| 416 | (["-O0"], "0"), |
| 417 | (["-O1"], "1"), |
| 418 | (["-O2"], "2"), |
| 419 | (["-O3"], "3"), |
| 420 | ], |
| 421 | ) |
| 422 | def test_optimization_level(args, expected): |
| 423 | """Test space-separated optimization levels (-O 1, -O 2, -O 3) map to |
| 424 | optimization_level. |
| 425 | """ |
| 426 | parser = EngineArgs.add_cli_args(FlexibleArgumentParser()) |
| 427 | parsed_args = parser.parse_args(args) |
| 428 | assert parsed_args.optimization_level == expected |
| 429 | assert parsed_args.compilation_config.mode is None |
| 430 | |
| 431 | |
| 432 | @pytest.mark.parametrize( |
| 433 | ("args", "expected"), |
| 434 | [ |
| 435 | (["-cc.mode=0"], 0), |
| 436 | (["-cc.mode=1"], 1), |
| 437 | (["-cc.mode=2"], 2), |
| 438 | (["-cc.mode=3"], 3), |
| 439 | ], |
| 440 | ) |
| 441 | def test_mode_parser(args, expected): |
| 442 | """Test compilation config modes (-cc.mode=int) map to compilation_config.""" |
| 443 | parser = EngineArgs.add_cli_args(FlexibleArgumentParser()) |
| 444 | parsed_args = parser.parse_args(args) |
| 445 | assert parsed_args.compilation_config.mode == expected |
| 446 | |
| 447 | |
| 448 | def test_compilation_config(): |
| 449 | parser = EngineArgs.add_cli_args(FlexibleArgumentParser()) |
| 450 | |
| 451 | # default value |
| 452 | args = parser.parse_args([]) |
| 453 | assert args.compilation_config == CompilationConfig() |
| 454 | |
| 455 | # set to string form of a dict |
| 456 | args = parser.parse_args( |
| 457 | [ |
| 458 | "-cc", |
| 459 | '{"mode": 3, "cudagraph_capture_sizes": [1, 2, 4, 8], "backend": "eager"}', |
| 460 | ] |
| 461 | ) |
| 462 | assert ( |
| 463 | args.compilation_config.mode == 3 |
| 464 | and args.compilation_config.cudagraph_capture_sizes == [1, 2, 4, 8] |
| 465 | and args.compilation_config.backend == "eager" |
| 466 | ) |
| 467 | |
| 468 | # set to string form of a dict |
| 469 | args = parser.parse_args( |
| 470 | [ |
| 471 | "--compilation-config=" |
| 472 | '{"mode": 3, "cudagraph_capture_sizes": [1, 2, 4, 8], ' |
| 473 | '"backend": "inductor"}', |
| 474 | ] |
| 475 | ) |
| 476 | assert ( |
| 477 | args.compilation_config.mode == 3 |
| 478 | and args.compilation_config.cudagraph_capture_sizes == [1, 2, 4, 8] |
| 479 | and args.compilation_config.backend == "inductor" |
| 480 | ) |
| 481 | |
| 482 | |
| 483 | def test_trust_request_mm_kwargs_cli(): |
| 484 | from vllm.entrypoints.launchers.cli_args import FrontendArgs |
| 485 | |
| 486 | parser = FrontendArgs.add_cli_args(FlexibleArgumentParser()) |
| 487 | |
| 488 | args = parser.parse_args([]) |
| 489 | assert not args.trust_request_mm_kwargs |
| 490 | |
| 491 | args = parser.parse_args(["--trust-request-mm-kwargs"]) |
| 492 | assert args.trust_request_mm_kwargs |
| 493 | |
| 494 | |
| 495 | def test_attention_config(): |
| 496 | from vllm.v1.attention.backends.registry import AttentionBackendEnum |
| 497 | |
| 498 | parser = EngineArgs.add_cli_args(FlexibleArgumentParser()) |
| 499 | |
| 500 | # default value |
| 501 | args = parser.parse_args([]) |
| 502 | assert args is not None |
| 503 | engine_args = EngineArgs.from_cli_args(args) |
| 504 | assert engine_args.attention_config == AttentionConfig() |
| 505 | |
| 506 | # set backend via dot notation |
| 507 | args = parser.parse_args(["--attention-config.backend", "FLASH_ATTN"]) |
| 508 | assert args is not None |
| 509 | engine_args = EngineArgs.from_cli_args(args) |
| 510 | assert engine_args.attention_config.backend is not None |
| 511 | assert engine_args.attention_config.backend.name == "FLASH_ATTN" |
| 512 | |
| 513 | # set backend via --attention-backend shorthand |
| 514 | args = parser.parse_args(["--attention-backend", "FLASHINFER"]) |
| 515 | assert args is not None |
| 516 | engine_args = EngineArgs.from_cli_args(args) |
| 517 | assert engine_args.attention_backend is not None |
| 518 | assert engine_args.attention_backend == "FLASHINFER" |
| 519 | |
| 520 | # set all fields via dot notation |
| 521 | args = parser.parse_args( |
| 522 | [ |
| 523 | "--attention-config.backend", |
| 524 | "FLASH_ATTN", |
| 525 | "--attention-config.flash_attn_version", |
| 526 | "3", |
| 527 | "--attention-config.flash_attn_max_num_splits_for_cuda_graph", |
| 528 | "16", |
| 529 | "--attention-config.use_trtllm_attention", |
| 530 | "true", |
| 531 | "--attention-config.disable_flashinfer_q_quantization", |
| 532 | "true", |
| 533 | ] |
| 534 | ) |
| 535 | assert args is not None |
| 536 | engine_args = EngineArgs.from_cli_args(args) |
| 537 | assert engine_args.attention_config.backend is not None |
| 538 | assert engine_args.attention_config.backend.name == "FLASH_ATTN" |
| 539 | assert engine_args.attention_config.flash_attn_version == 3 |
| 540 | assert engine_args.attention_config.flash_attn_max_num_splits_for_cuda_graph == 16 |
| 541 | assert engine_args.attention_config.use_trtllm_attention is True |
| 542 | assert engine_args.attention_config.disable_flashinfer_q_quantization is True |
| 543 | |
| 544 | # set to string form of a dict with all fields |
| 545 | args = parser.parse_args( |
| 546 | [ |
| 547 | "--attention-config=" |
| 548 | '{"backend": "FLASHINFER", "flash_attn_version": 2, ' |
| 549 | '"flash_attn_max_num_splits_for_cuda_graph": 8, ' |
| 550 | '"use_trtllm_attention": false, ' |
| 551 | '"disable_flashinfer_q_quantization": false}', |
| 552 | ] |
| 553 | ) |
| 554 | assert args is not None |
| 555 | engine_args = EngineArgs.from_cli_args(args) |
| 556 | assert engine_args.attention_config.backend is not None |
| 557 | assert engine_args.attention_config.backend.name == "FLASHINFER" |
| 558 | assert engine_args.attention_config.flash_attn_version == 2 |
| 559 | assert engine_args.attention_config.flash_attn_max_num_splits_for_cuda_graph == 8 |
| 560 | assert engine_args.attention_config.use_trtllm_attention is False |
| 561 | assert engine_args.attention_config.disable_flashinfer_q_quantization is False |
| 562 | |
| 563 | # test --attention-backend flows into VllmConfig.attention_config |
| 564 | args = parser.parse_args( |
| 565 | [ |
| 566 | "--model", |
| 567 | "facebook/opt-125m", |
| 568 | "--attention-backend", |
| 569 | "FLASH_ATTN", |
| 570 | ] |
| 571 | ) |
| 572 | assert args is not None |
| 573 | engine_args = EngineArgs.from_cli_args(args) |
| 574 | vllm_config = engine_args.create_engine_config() |
| 575 | assert vllm_config.attention_config.backend == AttentionBackendEnum.FLASH_ATTN |
| 576 | |
| 577 | # test --attention-config.backend flows into VllmConfig.attention_config |
| 578 | args = parser.parse_args( |
| 579 | [ |
| 580 | "--model", |
| 581 | "facebook/opt-125m", |
| 582 | "--attention-config.backend", |
| 583 | "FLASHINFER", |
| 584 | ] |
| 585 | ) |
| 586 | assert args is not None |
| 587 | engine_args = EngineArgs.from_cli_args(args) |
| 588 | vllm_config = engine_args.create_engine_config() |
| 589 | assert vllm_config.attention_config.backend == AttentionBackendEnum.FLASHINFER |
| 590 | |
| 591 | # test --attention-backend and --attention-config.backend are mutually exclusive |
| 592 | args = parser.parse_args( |
| 593 | [ |
| 594 | "--model", |
| 595 | "facebook/opt-125m", |
| 596 | "--attention-backend", |
| 597 | "FLASH_ATTN", |
| 598 | "--attention-config.backend", |
| 599 | "FLASHINFER", |
| 600 | ] |
| 601 | ) |
| 602 | assert args is not None |
| 603 | engine_args = EngineArgs.from_cli_args(args) |
| 604 | with pytest.raises(ValueError, match="mutually exclusive"): |
| 605 | engine_args.create_engine_config() |
| 606 | |
| 607 | |
| 608 | def test_multi_node_world_size_includes_pcp(monkeypatch): |
| 609 | """PCP expands the process world size, so the --nnodes divisibility check |
| 610 | must include it. Without this, TP=1/PCP=2 over 2 nodes computes a world |
| 611 | size of 1 and the launch is rejected before the engine starts.""" |
| 612 | import vllm.config.vllm |
| 613 | |
| 614 | # PCP requires the V2 model runner, which is gated on Triton. |
| 615 | monkeypatch.setattr(vllm.config.vllm, "HAS_TRITON", True) |
| 616 | |
| 617 | engine_args = EngineArgs( |
| 618 | model="facebook/opt-125m", |
| 619 | tensor_parallel_size=1, |
| 620 | prefill_context_parallel_size=2, |
| 621 | nnodes=2, |
| 622 | ) |
| 623 | vllm_config = engine_args.create_engine_config() |
| 624 | assert vllm_config.parallel_config.world_size == 2 |
| 625 | |
| 626 | |
| 627 | def test_prefix_cache_default(): |
| 628 | parser = EngineArgs.add_cli_args(FlexibleArgumentParser()) |
| 629 | args = parser.parse_args([]) |
| 630 | |
| 631 | # should be None by default (depends on model). |
| 632 | engine_args = EngineArgs.from_cli_args(args=args) |
| 633 | assert engine_args.enable_prefix_caching is None |
| 634 | assert engine_args.prefix_cache_retention_interval == 0 |
| 635 | |
| 636 | # with flag to turn it on. |
| 637 | args = parser.parse_args(["--enable-prefix-caching"]) |
| 638 | engine_args = EngineArgs.from_cli_args(args=args) |
| 639 | assert engine_args.enable_prefix_caching |
| 640 | |
| 641 | # with disable flag to turn it off. |
| 642 | args = parser.parse_args(["--no-enable-prefix-caching"]) |
| 643 | engine_args = EngineArgs.from_cli_args(args=args) |
| 644 | assert not engine_args.enable_prefix_caching |
| 645 | |
| 646 | args = parser.parse_args(["--prefix-cache-retention-interval", "64"]) |
| 647 | engine_args = EngineArgs.from_cli_args(args=args) |
| 648 | assert engine_args.prefix_cache_retention_interval == 64 |
| 649 | |
| 650 | |
| 651 | @pytest.mark.parametrize( |
| 652 | ("arg", "expected", "option"), |
| 653 | [ |
| 654 | (None, None, "mm-processor-kwargs"), |
| 655 | ("{}", {}, "mm-processor-kwargs"), |
| 656 | ('{"num_crops": 4}', {"num_crops": 4}, "mm-processor-kwargs"), |
| 657 | ('{"foo": {"bar": "baz"}}', {"foo": {"bar": "baz"}}, "mm-processor-kwargs"), |
| 658 | ], |
| 659 | ) |
| 660 | def test_composite_arg_parser(arg, expected, option): |
| 661 | parser = EngineArgs.add_cli_args(FlexibleArgumentParser()) |
| 662 | if arg is None: |
| 663 | args = parser.parse_args([]) |
| 664 | else: |
| 665 | args = parser.parse_args([f"--{option}", arg]) |
| 666 | assert getattr(args, option.replace("-", "_")) == expected |
| 667 | |
| 668 | |
| 669 | def test_human_readable_model_len(): |
| 670 | # `exit_on_error` disabled to test invalid values below |
| 671 | parser = EngineArgs.add_cli_args(FlexibleArgumentParser(exit_on_error=False)) |
| 672 | |
| 673 | args = parser.parse_args([]) |
| 674 | assert args.max_model_len is None |
| 675 | |
| 676 | args = parser.parse_args(["--max-model-len", "1024"]) |
| 677 | assert args.max_model_len == 1024 |
| 678 | |
| 679 | # Lower |
| 680 | args = parser.parse_args(["--max-model-len", "1m"]) |
| 681 | assert args.max_model_len == 1_000_000 |
| 682 | args = parser.parse_args(["--max-model-len", "10k"]) |
| 683 | assert args.max_model_len == 10_000 |
| 684 | args = parser.parse_args(["--max-model-len", "2g"]) |
| 685 | assert args.max_model_len == 2_000_000_000 |
| 686 | args = parser.parse_args(["--max-model-len", "2t"]) |
| 687 | assert args.max_model_len == 2_000_000_000_000 |
| 688 | |
| 689 | # Capital |
| 690 | args = parser.parse_args(["--max-model-len", "3K"]) |
| 691 | assert args.max_model_len == 2**10 * 3 |
| 692 | args = parser.parse_args(["--max-model-len", "10M"]) |
| 693 | assert args.max_model_len == 2**20 * 10 |
| 694 | args = parser.parse_args(["--max-model-len", "4G"]) |
| 695 | assert args.max_model_len == 2**30 * 4 |
| 696 | args = parser.parse_args(["--max-model-len", "4T"]) |
| 697 | assert args.max_model_len == 2**40 * 4 |
| 698 | |
| 699 | # Decimal values |
| 700 | args = parser.parse_args(["--max-model-len", "10.2k"]) |
| 701 | assert args.max_model_len == 10200 |
| 702 | # ..truncated to the nearest int |
| 703 | args = parser.parse_args(["--max-model-len", "10.2123451234567k"]) |
| 704 | assert args.max_model_len == 10212 |
| 705 | args = parser.parse_args(["--max-model-len", "10.2123451234567m"]) |
| 706 | assert args.max_model_len == 10212345 |
| 707 | args = parser.parse_args(["--max-model-len", "10.2123451234567g"]) |
| 708 | assert args.max_model_len == 10212345123 |
| 709 | args = parser.parse_args(["--max-model-len", "10.2123451234567t"]) |
| 710 | assert args.max_model_len == 10212345123456 |
| 711 | |
| 712 | # Special value -1 for auto-fit to GPU memory |
| 713 | args = parser.parse_args(["--max-model-len", "-1"]) |
| 714 | assert args.max_model_len == -1 |
| 715 | |
| 716 | # 'auto' is an alias for -1 |
| 717 | args = parser.parse_args(["--max-model-len", "auto"]) |
| 718 | assert args.max_model_len == -1 |
| 719 | args = parser.parse_args(["--max-model-len", "AUTO"]) |
| 720 | assert args.max_model_len == -1 |
| 721 | |
| 722 | # Invalid (do not allow decimals with binary multipliers) |
| 723 | for invalid in ["1a", "pwd", "10.24", "1.23M", "1.22T"]: |
| 724 | with pytest.raises(ArgumentError): |
| 725 | parser.parse_args(["--max-model-len", invalid]) |
| 726 | |
| 727 | |
| 728 | def test_human_readable_other_args(): |
| 729 | # Test human-readable parsing for other integer args |
| 730 | # that were added to use human_readable_int parser |
| 731 | parser = EngineArgs.add_cli_args(FlexibleArgumentParser(exit_on_error=False)) |
| 732 | |
| 733 | # Test max_num_scheduled_tokens |
| 734 | args = parser.parse_args(["--max-num-scheduled-tokens", "1024"]) |
| 735 | assert args.max_num_scheduled_tokens == 1024 |
| 736 | args = parser.parse_args(["--max-num-scheduled-tokens", "2k"]) |
| 737 | assert args.max_num_scheduled_tokens == 2_000 |
| 738 | args = parser.parse_args(["--max-num-scheduled-tokens", "4K"]) |
| 739 | assert args.max_num_scheduled_tokens == 2**10 * 4 |
| 740 | args = parser.parse_args(["--max-num-scheduled-tokens", "10.5k"]) |
| 741 | assert args.max_num_scheduled_tokens == 10500 |
| 742 | |
| 743 | # Test kv_cache_memory_bytes (existing human-readable arg) |
| 744 | args = parser.parse_args(["--kv-cache-memory-bytes", "100000"]) |
| 745 | assert args.kv_cache_memory_bytes == 100000 |
| 746 | args = parser.parse_args(["--kv-cache-memory-bytes", "100k"]) |
| 747 | assert args.kv_cache_memory_bytes == 100_000 |
| 748 | args = parser.parse_args(["--kv-cache-memory-bytes", "1M"]) |
| 749 | assert args.kv_cache_memory_bytes == 2**20 |
| 750 | args = parser.parse_args(["--kv-cache-memory-bytes", "1m"]) |
| 751 | assert args.kv_cache_memory_bytes == 1_000_000 |
| 752 | |
| 753 | # Test max_num_batched_tokens (existing human-readable arg) |
| 754 | args = parser.parse_args(["--max-num-batched-tokens", "1024"]) |
| 755 | assert args.max_num_batched_tokens == 1024 |
| 756 | args = parser.parse_args(["--max-num-batched-tokens", "2k"]) |
| 757 | assert args.max_num_batched_tokens == 2_000 |
| 758 | args = parser.parse_args(["--max-num-batched-tokens", "4K"]) |
| 759 | assert args.max_num_batched_tokens == 2**10 * 4 |
| 760 | |
| 761 | |
| 762 | def test_numa_bind_args(): |
| 763 | parser = EngineArgs.add_cli_args(FlexibleArgumentParser()) |
| 764 | args = parser.parse_args( |
| 765 | [ |
| 766 | "--numa-bind", |
| 767 | "--numa-bind-nodes", |
| 768 | "0", |
| 769 | "0", |
| 770 | "1", |
| 771 | "1", |
| 772 | "--numa-bind-cpus", |
| 773 | "0-3", |
| 774 | "4-7", |
| 775 | "8-11", |
| 776 | "12-15", |
| 777 | ] |
| 778 | ) |
| 779 | engine_args = EngineArgs.from_cli_args(args=args) |
| 780 | assert engine_args.numa_bind is True |
| 781 | assert engine_args.numa_bind_nodes == [0, 0, 1, 1] |
| 782 | assert engine_args.numa_bind_cpus == ["0-3", "4-7", "8-11", "12-15"] |
| 783 | |
| 784 | |
| 785 | def test_ir_op_priority(): |
| 786 | from vllm.config.kernel import IrOpPriorityConfig, KernelConfig |
| 787 | |
| 788 | ir_op_priority = IrOpPriorityConfig(rms_norm=["vllm_c"]) |
| 789 | cfg1 = EngineArgs(ir_op_priority=ir_op_priority).create_engine_config() |
| 790 | cfg2 = EngineArgs( |
| 791 | kernel_config=KernelConfig(ir_op_priority=ir_op_priority) |
| 792 | ).create_engine_config() |
| 793 | assert cfg1.kernel_config.ir_op_priority == cfg2.kernel_config.ir_op_priority |
| 794 | |
| 795 | with pytest.raises(ValueError, match="rms_norm"): |
| 796 | _ = EngineArgs( |
| 797 | ir_op_priority=ir_op_priority, |
| 798 | kernel_config=KernelConfig(ir_op_priority=ir_op_priority), |
| 799 | ).create_engine_config() |
| 800 | |
| 801 | |
| 802 | @pytest.mark.parametrize( |
| 803 | ("input_json", "expected_json"), |
| 804 | [ |
| 805 | # Decimal suffixes (lowercase) |
| 806 | ('{"x": 80g}', '{"x": 80000000000}'), |
| 807 | ('{"x": 1k}', '{"x": 1000}'), |
| 808 | ('{"x": 5m}', '{"x": 5000000}'), |
| 809 | ('{"x": 2t}', '{"x": 2000000000000}'), |
| 810 | # Binary suffixes (uppercase) |
| 811 | ('{"x": 1K}', f'{{"x": {2**10}}}'), |
| 812 | ('{"x": 1G}', f'{{"x": {2**30}}}'), |
| 813 | # Decimal values |
| 814 | ('{"x": 1.5g}', '{"x": 1500000000}'), |
| 815 | # Quoted strings must NOT be modified |
| 816 | ('{"my_key": 80g}', '{"my_key": 80000000000}'), |
| 817 | ('{"name": "80g"}', '{"name": "80g"}'), |
| 818 | ('{"model_name": "foo_bar"}', '{"model_name": "foo_bar"}'), |
| 819 | # Multiple values |
| 820 | ('{"a": 1k, "b": 2m}', '{"a": 1000, "b": 2000000}'), |
| 821 | # Plain numbers are untouched |
| 822 | ('{"x": 42}', '{"x": 42}'), |
| 823 | # Nested JSON |
| 824 | ('{"outer": {"inner": 10g}}', '{"outer": {"inner": 10000000000}}'), |
| 825 | ], |
| 826 | ) |
| 827 | def test_expand_json_human_readable_numbers(input_json, expected_json): |
| 828 | assert _expand_json_human_readable_numbers(input_json) == expected_json |
| 829 | |
| 830 | |
| 831 | @pytest.mark.parametrize( |
| 832 | "uri", |
| 833 | ["s3://bucket/model", "gs://bucket/model", "az://container/model"], |
| 834 | ) |
| 835 | def test_cloud_storage_uri_skips_get_model_path(uri, monkeypatch): |
| 836 | """Cloud storage URIs should not be passed to get_model_path() |
| 837 | when HF_HUB_OFFLINE=1, as they are not valid HF repo IDs.""" |
| 838 | import huggingface_hub |
| 839 | |
| 840 | monkeypatch.setattr(huggingface_hub.constants, "HF_HUB_OFFLINE", True) |
| 841 | |
| 842 | args = EngineArgs(model=uri) |
| 843 | # model should remain the original cloud URI, not raise |
| 844 | assert args.model == uri |
| 845 | |
| 846 | |
| 847 | def test_cloud_storage_tokenizer_skips_get_model_path(monkeypatch): |
| 848 | """Cloud storage tokenizer URI should not be passed to |
| 849 | get_model_path() when HF_HUB_OFFLINE=1.""" |
| 850 | import huggingface_hub |
| 851 | |
| 852 | monkeypatch.setattr(huggingface_hub.constants, "HF_HUB_OFFLINE", True) |
| 853 | |
| 854 | args = EngineArgs(model="s3://bucket/model", tokenizer="s3://bucket/tokenizer") |
| 855 | assert args.model == "s3://bucket/model" |
| 856 | assert args.tokenizer == "s3://bucket/tokenizer" |
| 857 | |
| 858 | |
| 859 | class TestDeviceIds: |
| 860 | def test_device_ids_with_cvd_out_of_range(self, monkeypatch): |
| 861 | """--device-ids index beyond the CVD set raises ValueError.""" |
| 862 | from vllm.platforms import current_platform |
| 863 | |
| 864 | key = current_platform.device_control_env_var |
| 865 | monkeypatch.setenv(key, "4,5") |
| 866 | args = EngineArgs(model="m", device_ids=[0, 2]) |
| 867 | with pytest.raises(ValueError, match="out of range"): |
| 868 | args._resolve_device_ids() |
| 869 | |
| 870 | def test_device_ids_with_cvd_resolve_to_physical_ids(self, monkeypatch): |
| 871 | """--device-ids are CVD-local indices resolved to physical ids.""" |
| 872 | from vllm.platforms import current_platform |
| 873 | |
| 874 | key = current_platform.device_control_env_var |
| 875 | monkeypatch.setenv(key, "4,5") |
| 876 | args = EngineArgs(model="m", device_ids=[0, 1]) |
| 877 | assert args._resolve_device_ids() == [4, 5] |
| 878 | |
| 879 | def test_device_ids_with_uuid_cvd_resolve_to_physical_ids(self, monkeypatch): |
| 880 | """--device-ids support UUID CVD values resolved by the platform.""" |
| 881 | from vllm.platforms import current_platform |
| 882 | |
| 883 | key = current_platform.device_control_env_var |
| 884 | monkeypatch.setenv(key, "GPU-abcd1234,GPU-ef567890") |
| 885 | monkeypatch.setattr( |
| 886 | type(current_platform), |
| 887 | "device_control_id_to_physical_device_id", |
| 888 | classmethod( |
| 889 | lambda cls, device_id: {"GPU-abcd1234": 4, "GPU-ef567890": 5}[device_id] |
| 890 | ), |
| 891 | ) |
| 892 | |
| 893 | args = EngineArgs(model="m", device_ids=[0, 1]) |
| 894 | assert args._resolve_device_ids() == [4, 5] |
| 895 | |
| 896 | def test_device_ids_with_uuid_args_resolve_to_physical_ids(self, monkeypatch): |
| 897 | """UUID --device-ids are resolved to physical IDs immediately.""" |
| 898 | from vllm.platforms import current_platform |
| 899 | |
| 900 | monkeypatch.setattr( |
| 901 | type(current_platform), |
| 902 | "device_control_id_to_physical_device_id", |
| 903 | classmethod(lambda cls, device_id: {"GPU-abcd1234": 4}[device_id]), |
| 904 | ) |
| 905 | |
| 906 | args = EngineArgs(model="m", device_ids=["GPU-abcd1234"]) |
| 907 | assert args._resolve_device_ids() == [4] |
| 908 | |
| 909 | def test_device_ids_reject_mixed_integer_and_uuid_args(self): |
| 910 | """--device-ids must not mix CVD indices and UUIDs.""" |
| 911 | args = EngineArgs(model="m", device_ids=[0, "GPU-abcd1234"]) |
| 912 | with pytest.raises(ValueError, match="must not mix"): |
| 913 | args._resolve_device_ids() |
| 914 | |
| 915 | def test_no_device_ids(self): |
| 916 | """No --device-ids returns None.""" |
| 917 | args = EngineArgs(model="m") |
| 918 | assert args._resolve_device_ids() is None |
| 919 | |
| 920 | def test_cli_parsing(self): |
| 921 | """--device-ids parses comma-separated string from CLI.""" |
| 922 | parser = FlexibleArgumentParser() |
| 923 | EngineArgs.add_cli_args(parser) |
| 924 | parsed = parser.parse_args(["--model", "m", "--device-ids", "0,2,4"]) |
| 925 | assert parsed.device_ids == [0, 2, 4] |
| 926 | |
| 927 | def test_cli_parsing_uuid(self): |
| 928 | """--device-ids parses comma-separated UUID strings from CLI.""" |
| 929 | parser = FlexibleArgumentParser() |
| 930 | EngineArgs.add_cli_args(parser) |
| 931 | parsed = parser.parse_args( |
| 932 | ["--model", "m", "--device-ids", "GPU-abcd1234,GPU-ef567890"] |
| 933 | ) |
| 934 | assert parsed.device_ids == ["GPU-abcd1234", "GPU-ef567890"] |
| 935 | |
| 936 | def test_assigned_physical_gpu_ids_are_physical_with_cvd(self, monkeypatch): |
| 937 | """assigned_physical_gpu_ids are already physical and not composed with CVD.""" |
| 938 | import vllm.platforms.interface as platform_interface |
| 939 | from vllm.platforms import current_platform |
| 940 | |
| 941 | monkeypatch.setattr(platform_interface, "_assigned_physical_gpu_ids", [4, 5]) |
| 942 | monkeypatch.setenv(current_platform.device_control_env_var, "4,5") |
| 943 | |
| 944 | assert current_platform.device_id_to_physical_device_id(0) == 4 |
| 945 | assert current_platform.device_id_to_physical_device_id(1) == 5 |
| 946 | assert current_platform.logical_device_id_to_visible_device_id(0) == 0 |
| 947 | assert current_platform.logical_device_id_to_visible_device_id(1) == 1 |
| 948 | |
| 949 | def test_assigned_physical_gpu_ids_map_to_visible_uuid_cvd(self, monkeypatch): |
| 950 | """Physical IDs map back to visible ordinals when CVD uses UUIDs.""" |
| 951 | import vllm.platforms.interface as platform_interface |
| 952 | from vllm.platforms import current_platform |
| 953 | |
| 954 | monkeypatch.setattr(platform_interface, "_assigned_physical_gpu_ids", [5]) |
| 955 | monkeypatch.setenv( |
| 956 | current_platform.device_control_env_var, |
| 957 | "GPU-abcd1234,GPU-ef567890", |
| 958 | ) |
| 959 | monkeypatch.setattr( |
| 960 | type(current_platform), |
| 961 | "device_control_id_to_physical_device_id", |
| 962 | classmethod( |
| 963 | lambda cls, device_id: {"GPU-abcd1234": 4, "GPU-ef567890": 5}[device_id] |
| 964 | ), |
| 965 | ) |
| 966 | |
| 967 | assert current_platform.logical_device_id_to_visible_device_id(0) == 1 |
| 968 | |
| 969 | def test_device_ids_reject_duplicates(self): |
| 970 | """--device-ids must not contain duplicate entries.""" |
| 971 | args = EngineArgs(model="m", device_ids=[2, 2]) |
| 972 | with pytest.raises(ValueError, match="duplicates"): |
| 973 | args._resolve_device_ids() |
| 974 | |
| 975 | def test_cli_parsing_strips_whitespace(self): |
| 976 | """--device-ids tolerates whitespace around commas.""" |
| 977 | parser = FlexibleArgumentParser() |
| 978 | EngineArgs.add_cli_args(parser) |
| 979 | parsed = parser.parse_args(["--model", "m", "--device-ids", "0, 2, 4"]) |
| 980 | assert parsed.device_ids == [0, 2, 4] |
| 981 | |
| 982 | def test_visible_ordinal_to_physical_ignores_assigned_ids(self, monkeypatch): |
| 983 | """visible_device_id_to_physical_device_id maps torch device ordinals, |
| 984 | independent of the logical-to-physical mapping. |
| 985 | |
| 986 | Regression test: CustomAllreduce passes device.index (a visible |
| 987 | ordinal) and must not index into assigned_physical_gpu_ids, which |
| 988 | raised IndexError for non-identity --device-ids like [2, 3]. |
| 989 | """ |
| 990 | import vllm.platforms.interface as platform_interface |
| 991 | from vllm.platforms import current_platform |
| 992 | |
| 993 | monkeypatch.setattr(platform_interface, "_assigned_physical_gpu_ids", [2, 3]) |
| 994 | monkeypatch.delenv(current_platform.device_control_env_var, raising=False) |
| 995 | |
| 996 | # CVD unset: visible ordinal == physical ID, even beyond the |
| 997 | # assigned list's length. |
| 998 | assert current_platform.visible_device_id_to_physical_device_id(2) == 2 |
| 999 | assert current_platform.visible_device_id_to_physical_device_id(3) == 3 |
| 1000 | |
| 1001 | monkeypatch.setenv(current_platform.device_control_env_var, "4,5") |
| 1002 | assert current_platform.visible_device_id_to_physical_device_id(1) == 5 |
| 1003 | with pytest.raises(IndexError, match="out of range"): |
| 1004 | current_platform.visible_device_id_to_physical_device_id(2) |
| 1005 | |
| 1006 | |
| 1007 | class TestDpDeviceIdSharding: |
| 1008 | def test_dp_supervisor_device_ids_stay_env_relative(self): |
| 1009 | """Regression test: the DP supervisor must pass env-relative indices, |
| 1010 | not physical IDs, because each child re-resolves --device-ids |
| 1011 | against its inherited device-control env var.""" |
| 1012 | import argparse |
| 1013 | |
| 1014 | from vllm.entrypoints.launchers.dp_supervisor import _build_device_ids |
| 1015 | |
| 1016 | args = argparse.Namespace( |
| 1017 | tensor_parallel_size=2, pipeline_parallel_size=1, device_ids=None |
| 1018 | ) |
| 1019 | assert _build_device_ids(args, local_rank=0) == [0, 1] |
| 1020 | assert _build_device_ids(args, local_rank=1) == [2, 3] |
| 1021 | |
| 1022 | def test_dp_supervisor_shards_user_device_ids(self): |
| 1023 | """User-provided --device-ids are sharded across DP children.""" |
| 1024 | import argparse |
| 1025 | |
| 1026 | from vllm.entrypoints.launchers.dp_supervisor import _build_device_ids |
| 1027 | |
| 1028 | args = argparse.Namespace( |
| 1029 | tensor_parallel_size=2, pipeline_parallel_size=1, device_ids=[4, 5, 6, 7] |
| 1030 | ) |
| 1031 | assert _build_device_ids(args, local_rank=0) == [4, 5] |
| 1032 | assert _build_device_ids(args, local_rank=1) == [6, 7] |
| 1033 | with pytest.raises(ValueError, match="needs devices"): |
| 1034 | _build_device_ids(args, local_rank=2) |
| 1035 | |
| 1036 | def test_dp_rank_shards_user_assigned_gpu_ids(self): |
| 1037 | """get_physical_gpu_ids_for_local_dp_rank slices the user-provided |
| 1038 | --device-ids list instead of recomputing from the env var.""" |
| 1039 | from vllm.platforms import current_platform |
| 1040 | from vllm.v1.engine.utils import get_physical_gpu_ids_for_local_dp_rank |
| 1041 | |
| 1042 | evar = current_platform.device_control_env_var |
| 1043 | assert get_physical_gpu_ids_for_local_dp_rank( |
| 1044 | evar, local_dp_rank=1, world_size=2, user_assigned_gpu_ids=[4, 5, 6, 7] |
| 1045 | ) == [6, 7] |
| 1046 | with pytest.raises(ValueError, match="needs devices"): |
| 1047 | get_physical_gpu_ids_for_local_dp_rank( |
| 1048 | evar, local_dp_rank=2, world_size=2, user_assigned_gpu_ids=[4, 5, 6, 7] |
| 1049 | ) |