test_arg_utils.py 20.2 KB
Newer Older
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 AttentionConfig, CompilationConfig, config
13
14
from vllm.engine.arg_utils import (
    EngineArgs,
15
    _expand_json_human_readable_numbers,
16
17
18
19
20
21
22
23
24
25
    contains_type,
    get_kwargs,
    get_type,
    get_type_hints,
    is_not_builtin,
    is_type,
    literal_to_kwargs,
    optional_type,
    parse_type,
)
26
from vllm.utils.argparse_utils import FlexibleArgumentParser
27
28


29
30
31
32
33
34
35
36
37
@pytest.mark.parametrize(
    ("type", "value", "expected"),
    [
        (int, "42", 42),
        (float, "3.14", 3.14),
        (str, "Hello World!", "Hello World!"),
        (json.loads, '{"foo":1,"bar":2}', {"foo": 1, "bar": 2}),
    ],
)
38
39
def test_parse_type(type, value, expected):
    parse_type_func = parse_type(type)
40
    assert parse_type_func(value) == expected
41
42
43
44
45
46


def test_optional_type():
    optional_type_func = optional_type(int)
    assert optional_type_func("None") is None
    assert optional_type_func("42") == 42
47
48


49
50
51
52
53
54
55
56
57
58
@pytest.mark.parametrize(
    ("type_hint", "type", "expected"),
    [
        (int, int, True),
        (int, float, False),
        (list[int], list, True),
        (list[int], tuple, False),
        (Literal[0, 1], Literal, True),
    ],
)
59
60
61
62
def test_is_type(type_hint, type, expected):
    assert is_type(type_hint, type) == expected


63
64
65
66
67
68
69
70
71
72
73
74
@pytest.mark.parametrize(
    ("type_hints", "type", "expected"),
    [
        ({float, int}, int, True),
        ({int, tuple}, int, True),
        ({int, tuple[int]}, int, True),
        ({int, tuple[int, ...]}, int, True),
        ({int, tuple[int]}, float, False),
        ({int, tuple[int, ...]}, float, False),
        ({str, Literal["x", "y"]}, Literal, True),
    ],
)
75
76
77
78
def test_contains_type(type_hints, type, expected):
    assert contains_type(type_hints, type) == expected


79
80
81
82
83
84
85
86
@pytest.mark.parametrize(
    ("type_hints", "type", "expected"),
    [
        ({int, float}, int, int),
        ({int, float}, str, None),
        ({str, Literal["x", "y"]}, Literal, Literal["x", "y"]),
    ],
)
87
88
89
90
def test_get_type(type_hints, type, expected):
    assert get_type(type_hints, type) == expected


91
92
93
94
95
96
97
98
@pytest.mark.parametrize(
    ("type_hints", "expected"),
    [
        ({Literal[1, 2]}, {"type": int, "choices": [1, 2]}),
        ({str, Literal["x", "y"]}, {"type": str, "metavar": ["x", "y"]}),
        ({Literal[1, "a"]}, Exception),
    ],
)
99
def test_literal_to_kwargs(type_hints, expected):
100
    context: AbstractContextManager[object] = nullcontext()
101
102
103
104
105
106
    if expected is Exception:
        context = pytest.raises(expected)
    with context:
        assert literal_to_kwargs(type_hints) == expected


107
@config
108
109
110
111
112
113
114
class NestedConfig:
    field: int = 1
    """field"""


@config
class DummyConfig:
115
116
    regular_bool: bool = True
    """Regular bool with default True"""
117
    optional_bool: bool | None = None
118
    """Optional bool with default None"""
119
    optional_literal: Literal["x", "y"] | None = None
120
    """Optional literal with default None"""
121
    tuple_n: tuple[int, ...] = Field(default_factory=lambda: (1, 2, 3))
122
    """Tuple with variable length"""
123
    tuple_2: tuple[int, int] = Field(default_factory=lambda: (1, 2))
124
    """Tuple with fixed length"""
125
    list_n: list[int] = Field(default_factory=lambda: [1, 2, 3])
126
    """List with variable length"""
127
    list_literal: list[Literal[1, 2]] = Field(default_factory=list)
128
    """List with literal choices"""
129
    list_union: list[str | type[object]] = Field(default_factory=list)
130
    """List with union type"""
131
    set_n: set[int] = Field(default_factory=lambda: {1, 2, 3})
132
    """Set with variable length"""
133
134
    literal_literal: Literal[Literal[1], Literal[2]] = 1
    """Literal of literals with default 1"""
135
    json_tip: dict = Field(default_factory=dict)
136
    """Dict which will be JSON in CLI"""
137
    nested_config: NestedConfig = Field(default_factory=NestedConfig)
138
    """Nested config"""
139
140


141
142
143
144
145
146
147
@pytest.mark.parametrize(
    ("type_hint", "expected"),
    [
        (int, False),
        (DummyConfig, True),
    ],
)
148
149
150
151
def test_is_not_builtin(type_hint, expected):
    assert is_not_builtin(type_hint) == expected


152
@pytest.mark.parametrize(
153
154
    ("type_hint", "expected"),
    [
155
        (Annotated[int, "annotation"], {int}),
156
157
158
        (int | None, {int, type(None)}),
        (Annotated[int | None, "annotation"], {int, type(None)}),
        (Annotated[int, "annotation"] | None, {int, type(None)}),
159
    ],
160
    ids=["Annotated", "or_None", "Annotated_or_None", "or_None_Annotated"],
161
)
162
163
164
165
def test_get_type_hints(type_hint, expected):
    assert get_type_hints(type_hint) == expected


166
def test_get_kwargs():
167
    kwargs = get_kwargs(DummyConfig)
168
169
170
171
172
173
174
175
176
177
178
179
180
    print(kwargs)

    # bools should not have their type set
    assert kwargs["regular_bool"].get("type") is None
    assert kwargs["optional_bool"].get("type") is None
    # optional literals should have None as a choice
    assert kwargs["optional_literal"]["choices"] == ["x", "y", "None"]
    # tuples should have the correct nargs
    assert kwargs["tuple_n"]["nargs"] == "+"
    assert kwargs["tuple_2"]["nargs"] == 2
    # lists should work
    assert kwargs["list_n"]["type"] is int
    assert kwargs["list_n"]["nargs"] == "+"
181
182
183
184
    # lists with literals should have the correct choices
    assert kwargs["list_literal"]["type"] is int
    assert kwargs["list_literal"]["nargs"] == "+"
    assert kwargs["list_literal"]["choices"] == [1, 2]
185
186
187
    # lists with unions should become str type.
    # If not, we cannot know which type to use for parsing
    assert kwargs["list_union"]["type"] is str
188
189
190
    # sets should work like lists
    assert kwargs["set_n"]["type"] is int
    assert kwargs["set_n"]["nargs"] == "+"
191
192
    # literals of literals should have merged choices
    assert kwargs["literal_literal"]["choices"] == [1, 2]
193
    # dict should have json tip in help
194
195
    json_tip = "Should either be a valid JSON string or JSON keys"
    assert json_tip in kwargs["json_tip"]["help"]
196
    # nested config should construct the nested config
197
    assert kwargs["nested_config"]["type"]('{"field": 2}') == NestedConfig(2)  # type: ignore[call-arg]
198
199


200
201
202
203
@pytest.mark.parametrize(
    ("arg", "expected"),
    [
        (None, dict()),
204
        ('{"video": {"num_frames": 123} }', {"video": {"num_frames": 123}}),
205
206
207
        (
            '{"video": {"num_frames": 123, "fps": 1.0, "foo": "bar"}, "image": {"foo": "bar"} }',  # noqa
            {
208
209
210
211
212
213
                "video": {"num_frames": 123, "fps": 1.0, "foo": "bar"},
                "image": {"foo": "bar"},
            },
        ),
    ],
)
214
215
216
217
218
219
220
221
222
223
def test_media_io_kwargs_parser(arg, expected):
    parser = EngineArgs.add_cli_args(FlexibleArgumentParser())
    if arg is None:
        args = parser.parse_args([])
    else:
        args = parser.parse_args(["--media-io-kwargs", arg])

    assert args.media_io_kwargs == expected


224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
@pytest.mark.parametrize(
    ("args", "expected"),
    [
        (["-O", "1"], "1"),
        (["-O", "2"], "2"),
        (["-O", "3"], "3"),
        (["-O0"], "0"),
        (["-O1"], "1"),
        (["-O2"], "2"),
        (["-O3"], "3"),
    ],
)
def test_optimization_level(args, expected):
    """
    Test space-separated optimization levels (-O 1, -O 2, -O 3) map to
    optimization_level.
    """
241
    parser = EngineArgs.add_cli_args(FlexibleArgumentParser())
242
243
244
    parsed_args = parser.parse_args(args)
    assert parsed_args.optimization_level == expected
    assert parsed_args.compilation_config.mode is None
245
246


247
248
249
@pytest.mark.parametrize(
    ("args", "expected"),
    [
250
251
252
253
        (["-cc.mode=0"], 0),
        (["-cc.mode=1"], 1),
        (["-cc.mode=2"], 2),
        (["-cc.mode=3"], 3),
254
255
256
257
    ],
)
def test_mode_parser(args, expected):
    """
258
    Test compilation config modes (-cc.mode=int) map to compilation_config.
259
260
261
262
    """
    parser = EngineArgs.add_cli_args(FlexibleArgumentParser())
    parsed_args = parser.parse_args(args)
    assert parsed_args.compilation_config.mode == expected
263
264


265
266
def test_compilation_config():
    parser = EngineArgs.add_cli_args(FlexibleArgumentParser())
267

268
269
270
    # default value
    args = parser.parse_args([])
    assert args.compilation_config == CompilationConfig()
271

272
    # set to string form of a dict
273
274
    args = parser.parse_args(
        [
275
            "-cc",
276
            '{"mode": 3, "cudagraph_capture_sizes": [1, 2, 4, 8], "backend": "eager"}',
277
278
279
        ]
    )
    assert (
280
        args.compilation_config.mode == 3
281
        and args.compilation_config.cudagraph_capture_sizes == [1, 2, 4, 8]
282
        and args.compilation_config.backend == "eager"
283
    )
284

285
    # set to string form of a dict
286
287
288
    args = parser.parse_args(
        [
            "--compilation-config="
289
            '{"mode": 3, "cudagraph_capture_sizes": [1, 2, 4, 8], '
290
            '"backend": "inductor"}',
291
292
293
        ]
    )
    assert (
294
        args.compilation_config.mode == 3
295
        and args.compilation_config.cudagraph_capture_sizes == [1, 2, 4, 8]
296
        and args.compilation_config.backend == "inductor"
297
    )
298
299


300
def test_attention_config():
301
    from vllm.v1.attention.backends.registry import AttentionBackendEnum
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432

    parser = EngineArgs.add_cli_args(FlexibleArgumentParser())

    # default value
    args = parser.parse_args([])
    assert args is not None
    engine_args = EngineArgs.from_cli_args(args)
    assert engine_args.attention_config == AttentionConfig()

    # set backend via dot notation
    args = parser.parse_args(["--attention-config.backend", "FLASH_ATTN"])
    assert args is not None
    engine_args = EngineArgs.from_cli_args(args)
    assert engine_args.attention_config.backend is not None
    assert engine_args.attention_config.backend.name == "FLASH_ATTN"

    # set backend via --attention-backend shorthand
    args = parser.parse_args(["--attention-backend", "FLASHINFER"])
    assert args is not None
    engine_args = EngineArgs.from_cli_args(args)
    assert engine_args.attention_backend is not None
    assert engine_args.attention_backend == "FLASHINFER"

    # set all fields via dot notation
    args = parser.parse_args(
        [
            "--attention-config.backend",
            "FLASH_ATTN",
            "--attention-config.flash_attn_version",
            "3",
            "--attention-config.use_prefill_decode_attention",
            "true",
            "--attention-config.flash_attn_max_num_splits_for_cuda_graph",
            "16",
            "--attention-config.use_cudnn_prefill",
            "true",
            "--attention-config.use_trtllm_ragged_deepseek_prefill",
            "true",
            "--attention-config.use_trtllm_attention",
            "true",
            "--attention-config.disable_flashinfer_prefill",
            "true",
            "--attention-config.disable_flashinfer_q_quantization",
            "true",
        ]
    )
    assert args is not None
    engine_args = EngineArgs.from_cli_args(args)
    assert engine_args.attention_config.backend is not None
    assert engine_args.attention_config.backend.name == "FLASH_ATTN"
    assert engine_args.attention_config.flash_attn_version == 3
    assert engine_args.attention_config.use_prefill_decode_attention is True
    assert engine_args.attention_config.flash_attn_max_num_splits_for_cuda_graph == 16
    assert engine_args.attention_config.use_cudnn_prefill is True
    assert engine_args.attention_config.use_trtllm_ragged_deepseek_prefill is True
    assert engine_args.attention_config.use_trtllm_attention is True
    assert engine_args.attention_config.disable_flashinfer_prefill is True
    assert engine_args.attention_config.disable_flashinfer_q_quantization is True

    # set to string form of a dict with all fields
    args = parser.parse_args(
        [
            "--attention-config="
            '{"backend": "FLASHINFER", "flash_attn_version": 2, '
            '"use_prefill_decode_attention": false, '
            '"flash_attn_max_num_splits_for_cuda_graph": 8, '
            '"use_cudnn_prefill": false, '
            '"use_trtllm_ragged_deepseek_prefill": false, '
            '"use_trtllm_attention": false, '
            '"disable_flashinfer_prefill": false, '
            '"disable_flashinfer_q_quantization": false}',
        ]
    )
    assert args is not None
    engine_args = EngineArgs.from_cli_args(args)
    assert engine_args.attention_config.backend is not None
    assert engine_args.attention_config.backend.name == "FLASHINFER"
    assert engine_args.attention_config.flash_attn_version == 2
    assert engine_args.attention_config.use_prefill_decode_attention is False
    assert engine_args.attention_config.flash_attn_max_num_splits_for_cuda_graph == 8
    assert engine_args.attention_config.use_cudnn_prefill is False
    assert engine_args.attention_config.use_trtllm_ragged_deepseek_prefill is False
    assert engine_args.attention_config.use_trtllm_attention is False
    assert engine_args.attention_config.disable_flashinfer_prefill is False
    assert engine_args.attention_config.disable_flashinfer_q_quantization is False

    # test --attention-backend flows into VllmConfig.attention_config
    args = parser.parse_args(
        [
            "--model",
            "facebook/opt-125m",
            "--attention-backend",
            "FLASH_ATTN",
        ]
    )
    assert args is not None
    engine_args = EngineArgs.from_cli_args(args)
    vllm_config = engine_args.create_engine_config()
    assert vllm_config.attention_config.backend == AttentionBackendEnum.FLASH_ATTN

    # test --attention-config.backend flows into VllmConfig.attention_config
    args = parser.parse_args(
        [
            "--model",
            "facebook/opt-125m",
            "--attention-config.backend",
            "FLASHINFER",
        ]
    )
    assert args is not None
    engine_args = EngineArgs.from_cli_args(args)
    vllm_config = engine_args.create_engine_config()
    assert vllm_config.attention_config.backend == AttentionBackendEnum.FLASHINFER

    # test --attention-backend and --attention-config.backend are mutually exclusive
    args = parser.parse_args(
        [
            "--model",
            "facebook/opt-125m",
            "--attention-backend",
            "FLASH_ATTN",
            "--attention-config.backend",
            "FLASHINFER",
        ]
    )
    assert args is not None
    engine_args = EngineArgs.from_cli_args(args)
    with pytest.raises(ValueError, match="mutually exclusive"):
        engine_args.create_engine_config()


433
434
435
436
def test_prefix_cache_default():
    parser = EngineArgs.add_cli_args(FlexibleArgumentParser())
    args = parser.parse_args([])

437
    # should be None by default (depends on model).
438
    engine_args = EngineArgs.from_cli_args(args=args)
439
    assert engine_args.enable_prefix_caching is None
440
441
442
443
444
445
446
447
448
449
450
451

    # with flag to turn it on.
    args = parser.parse_args(["--enable-prefix-caching"])
    engine_args = EngineArgs.from_cli_args(args=args)
    assert engine_args.enable_prefix_caching

    # with disable flag to turn it off.
    args = parser.parse_args(["--no-enable-prefix-caching"])
    engine_args = EngineArgs.from_cli_args(args=args)
    assert not engine_args.enable_prefix_caching


452
453
454
455
456
457
458
459
460
@pytest.mark.parametrize(
    ("arg", "expected", "option"),
    [
        (None, None, "mm-processor-kwargs"),
        ("{}", {}, "mm-processor-kwargs"),
        ('{"num_crops": 4}', {"num_crops": 4}, "mm-processor-kwargs"),
        ('{"foo": {"bar": "baz"}}', {"foo": {"bar": "baz"}}, "mm-processor-kwargs"),
    ],
)
461
def test_composite_arg_parser(arg, expected, option):
462
463
464
465
    parser = EngineArgs.add_cli_args(FlexibleArgumentParser())
    if arg is None:
        args = parser.parse_args([])
    else:
466
467
        args = parser.parse_args([f"--{option}", arg])
    assert getattr(args, option.replace("-", "_")) == expected
468
469
470
471


def test_human_readable_model_len():
    # `exit_on_error` disabled to test invalid values below
472
    parser = EngineArgs.add_cli_args(FlexibleArgumentParser(exit_on_error=False))
473
474
475
476
477
478
479
480
481
482
483
484

    args = parser.parse_args([])
    assert args.max_model_len is None

    args = parser.parse_args(["--max-model-len", "1024"])
    assert args.max_model_len == 1024

    # Lower
    args = parser.parse_args(["--max-model-len", "1m"])
    assert args.max_model_len == 1_000_000
    args = parser.parse_args(["--max-model-len", "10k"])
    assert args.max_model_len == 10_000
485
486
487
488
    args = parser.parse_args(["--max-model-len", "2g"])
    assert args.max_model_len == 2_000_000_000
    args = parser.parse_args(["--max-model-len", "2t"])
    assert args.max_model_len == 2_000_000_000_000
489
490
491

    # Capital
    args = parser.parse_args(["--max-model-len", "3K"])
492
    assert args.max_model_len == 2**10 * 3
493
494
    args = parser.parse_args(["--max-model-len", "10M"])
    assert args.max_model_len == 2**20 * 10
495
496
497
498
    args = parser.parse_args(["--max-model-len", "4G"])
    assert args.max_model_len == 2**30 * 4
    args = parser.parse_args(["--max-model-len", "4T"])
    assert args.max_model_len == 2**40 * 4
499
500
501
502
503

    # Decimal values
    args = parser.parse_args(["--max-model-len", "10.2k"])
    assert args.max_model_len == 10200
    # ..truncated to the nearest int
504
    args = parser.parse_args(["--max-model-len", "10.2123451234567k"])
505
    assert args.max_model_len == 10212
506
507
508
509
510
511
    args = parser.parse_args(["--max-model-len", "10.2123451234567m"])
    assert args.max_model_len == 10212345
    args = parser.parse_args(["--max-model-len", "10.2123451234567g"])
    assert args.max_model_len == 10212345123
    args = parser.parse_args(["--max-model-len", "10.2123451234567t"])
    assert args.max_model_len == 10212345123456
512

513
514
515
516
517
518
519
520
521
522
    # Special value -1 for auto-fit to GPU memory
    args = parser.parse_args(["--max-model-len", "-1"])
    assert args.max_model_len == -1

    # 'auto' is an alias for -1
    args = parser.parse_args(["--max-model-len", "auto"])
    assert args.max_model_len == -1
    args = parser.parse_args(["--max-model-len", "AUTO"])
    assert args.max_model_len == -1

523
    # Invalid (do not allow decimals with binary multipliers)
524
    for invalid in ["1a", "pwd", "10.24", "1.23M", "1.22T"]:
525
        with pytest.raises(ArgumentError):
526
            parser.parse_args(["--max-model-len", invalid])
527
528


529
530
531
532
533
534
535
536
537
538
539
540
541
542
543
544
545
546
547
548
549
550
551
def test_numa_bind_args():
    parser = EngineArgs.add_cli_args(FlexibleArgumentParser())
    args = parser.parse_args(
        [
            "--numa-bind",
            "--numa-bind-nodes",
            "0",
            "0",
            "1",
            "1",
            "--numa-bind-cpus",
            "0-3",
            "4-7",
            "8-11",
            "12-15",
        ]
    )
    engine_args = EngineArgs.from_cli_args(args=args)
    assert engine_args.numa_bind is True
    assert engine_args.numa_bind_nodes == [0, 0, 1, 1]
    assert engine_args.numa_bind_cpus == ["0-3", "4-7", "8-11", "12-15"]


552
553
554
555
556
557
558
559
560
561
562
563
564
565
566
def test_ir_op_priority():
    from vllm.config.kernel import IrOpPriorityConfig, KernelConfig

    ir_op_priority = IrOpPriorityConfig(rms_norm=["vllm_c"])
    cfg1 = EngineArgs(ir_op_priority=ir_op_priority).create_engine_config()
    cfg2 = EngineArgs(
        kernel_config=KernelConfig(ir_op_priority=ir_op_priority)
    ).create_engine_config()
    assert cfg1.kernel_config.ir_op_priority == cfg2.kernel_config.ir_op_priority

    with pytest.raises(ValueError, match="rms_norm"):
        _ = EngineArgs(
            ir_op_priority=ir_op_priority,
            kernel_config=KernelConfig(ir_op_priority=ir_op_priority),
        ).create_engine_config()
567
568
569
570
571
572
573
574
575
576
577
578
579
580
581
582
583
584
585
586
587
588
589
590
591
592
593
594
595


@pytest.mark.parametrize(
    ("input_json", "expected_json"),
    [
        # Decimal suffixes (lowercase)
        ('{"x": 80g}', '{"x": 80000000000}'),
        ('{"x": 1k}', '{"x": 1000}'),
        ('{"x": 5m}', '{"x": 5000000}'),
        ('{"x": 2t}', '{"x": 2000000000000}'),
        # Binary suffixes (uppercase)
        ('{"x": 1K}', f'{{"x": {2**10}}}'),
        ('{"x": 1G}', f'{{"x": {2**30}}}'),
        # Decimal values
        ('{"x": 1.5g}', '{"x": 1500000000}'),
        # Quoted strings must NOT be modified
        ('{"my_key": 80g}', '{"my_key": 80000000000}'),
        ('{"name": "80g"}', '{"name": "80g"}'),
        ('{"model_name": "foo_bar"}', '{"model_name": "foo_bar"}'),
        # Multiple values
        ('{"a": 1k, "b": 2m}', '{"a": 1000, "b": 2000000}'),
        # Plain numbers are untouched
        ('{"x": 42}', '{"x": 42}'),
        # Nested JSON
        ('{"outer": {"inner": 10g}}', '{"outer": {"inner": 10000000000}}'),
    ],
)
def test_expand_json_human_readable_numbers(input_json, expected_json):
    assert _expand_json_human_readable_numbers(input_json) == expected_json