test_arg_utils.py 11.8 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
7
from contextlib import nullcontext
from dataclasses import dataclass, field
8
from typing import Annotated, Literal, Optional, Union
9

10
11
import pytest

12
from vllm.config import CompilationConfig, config
13
from vllm.engine.arg_utils import (EngineArgs, contains_type, get_kwargs,
14
                                   get_type, get_type_hints, is_not_builtin,
15
16
                                   is_type, literal_to_kwargs, optional_type,
                                   parse_type)
17
18
19
from vllm.utils import FlexibleArgumentParser


20
21
22
23
24
25
26
27
28
@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
    }),
])
29
30
def test_parse_type(type, value, expected):
    parse_type_func = parse_type(type)
31
    assert parse_type_func(value) == expected
32
33
34
35
36
37


def test_optional_type():
    optional_type_func = optional_type(int)
    assert optional_type_func("None") is None
    assert optional_type_func("42") == 42
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69


@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),
])
def test_is_type(type_hint, type, expected):
    assert is_type(type_hint, type) == expected


@pytest.mark.parametrize(("type_hints", "type", "expected"), [
    ({float, int}, int, True),
    ({int, tuple[int]}, int, True),
    ({int, tuple[int]}, float, False),
    ({str, Literal["x", "y"]}, Literal, True),
])
def test_contains_type(type_hints, type, expected):
    assert contains_type(type_hints, type) == expected


@pytest.mark.parametrize(("type_hints", "type", "expected"), [
    ({int, float}, int, int),
    ({int, float}, str, None),
    ({str, Literal["x", "y"]}, Literal, Literal["x", "y"]),
])
def test_get_type(type_hints, type, expected):
    assert get_type(type_hints, type) == expected


70
71
72
73
74
@pytest.mark.parametrize(("type_hints", "expected"), [
    ({Literal[1, 2]}, {
        "type": int,
        "choices": [1, 2]
    }),
75
76
77
78
    ({str, Literal["x", "y"]}, {
        "type": str,
        "metavar": ["x", "y"]
    }),
79
80
81
82
83
84
85
86
87
88
    ({Literal[1, "a"]}, Exception),
])
def test_literal_to_kwargs(type_hints, expected):
    context = nullcontext()
    if expected is Exception:
        context = pytest.raises(expected)
    with context:
        assert literal_to_kwargs(type_hints) == expected


89
90
@config
@dataclass
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
class NestedConfig:
    field: int = 1
    """field"""


@config
@dataclass
class FromCliConfig1:
    field: int = 1
    """field"""

    @classmethod
    def from_cli(cls, cli_value: str):
        inst = cls(**json.loads(cli_value))
        inst.field += 1
        return inst


@config
@dataclass
class FromCliConfig2:
    field: int = 1
    """field"""

    @classmethod
    def from_cli(cls, cli_value: str):
        inst = cls(**json.loads(cli_value))
        inst.field += 2
        return inst


@config
@dataclass
class DummyConfig:
125
126
127
128
129
130
131
    regular_bool: bool = True
    """Regular bool with default True"""
    optional_bool: Optional[bool] = None
    """Optional bool with default None"""
    optional_literal: Optional[Literal["x", "y"]] = None
    """Optional literal with default None"""
    tuple_n: tuple[int, ...] = field(default_factory=lambda: (1, 2, 3))
132
    """Tuple with variable length"""
133
    tuple_2: tuple[int, int] = field(default_factory=lambda: (1, 2))
134
    """Tuple with fixed length"""
135
    list_n: list[int] = field(default_factory=lambda: [1, 2, 3])
136
137
138
    """List with variable length"""
    list_literal: list[Literal[1, 2]] = field(default_factory=list)
    """List with literal choices"""
139
140
    list_union: list[Union[str, type[object]]] = field(default_factory=list)
    """List with union type"""
141
142
    literal_literal: Literal[Literal[1], Literal[2]] = 1
    """Literal of literals with default 1"""
143
144
    json_tip: dict = field(default_factory=dict)
    """Dict which will be JSON in CLI"""
145
146
147
148
149
150
    nested_config: NestedConfig = field(default_factory=NestedConfig)
    """Nested config"""
    from_cli_config1: FromCliConfig1 = field(default_factory=FromCliConfig1)
    """Config with from_cli method"""
    from_cli_config2: FromCliConfig2 = field(default_factory=FromCliConfig2)
    """Different config with from_cli method"""
151
152
153
154


@pytest.mark.parametrize(("type_hint", "expected"), [
    (int, False),
155
    (DummyConfig, True),
156
157
158
159
160
])
def test_is_not_builtin(type_hint, expected):
    assert is_not_builtin(type_hint) == expected


161
162
163
164
165
166
167
168
169
170
171
172
@pytest.mark.parametrize(
    ("type_hint", "expected"), [
        (Annotated[int, "annotation"], {int}),
        (Optional[int], {int, type(None)}),
        (Annotated[Optional[int], "annotation"], {int, type(None)}),
        (Optional[Annotated[int, "annotation"]], {int, type(None)}),
    ],
    ids=["Annotated", "Optional", "Annotated_Optional", "Optional_Annotated"])
def test_get_type_hints(type_hint, expected):
    assert get_type_hints(type_hint) == expected


173
def test_get_kwargs():
174
    kwargs = get_kwargs(DummyConfig)
175
176
177
178
179
180
181
182
183
184
185
186
187
    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"] == "+"
188
189
190
191
    # 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]
192
193
194
    # 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
195
196
    # literals of literals should have merged choices
    assert kwargs["literal_literal"]["choices"] == [1, 2]
197
    # dict should have json tip in help
198
199
    json_tip = "Should either be a valid JSON string or JSON keys"
    assert json_tip in kwargs["json_tip"]["help"]
200
201
202
203
204
    # nested config should should construct the nested config
    assert kwargs["nested_config"]["type"]('{"field": 2}') == NestedConfig(2)
    # from_cli configs should be constructed with the correct method
    assert kwargs["from_cli_config1"]["type"]('{"field": 2}').field == 3
    assert kwargs["from_cli_config2"]["type"]('{"field": 2}').field == 4
205
206


207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
@pytest.mark.parametrize(
    ("arg", "expected"),
    [
        (None, dict()),
        ('{"video": {"num_frames": 123} }', {
            "video": {
                "num_frames": 123
            }
        }),
        (
            '{"video": {"num_frames": 123, "fps": 1.0, "foo": "bar"}, "image": {"foo": "bar"} }',  # noqa
            {
                "video": {
                    "num_frames": 123,
                    "fps": 1.0,
                    "foo": "bar"
                },
                "image": {
                    "foo": "bar"
                }
            }),
    ])
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


239
240
241
242
243
def test_compilation_config():
    parser = EngineArgs.add_cli_args(FlexibleArgumentParser())

    # default value
    args = parser.parse_args([])
244
    assert args.compilation_config == CompilationConfig()
245
246

    # set to O3
247
248
    args = parser.parse_args(["-O0"])
    assert args.compilation_config.level == 0
249
250

    # set to O 3 (space)
251
252
    args = parser.parse_args(["-O", "1"])
    assert args.compilation_config.level == 1
253
254

    # set to O 3 (equals)
255
256
257
258
259
    args = parser.parse_args(["-O=2"])
    assert args.compilation_config.level == 2

    # set to O.level 3
    args = parser.parse_args(["-O.level", "3"])
260
261
    assert args.compilation_config.level == 3

262
    # set to string form of a dict
263
    args = parser.parse_args([
264
265
266
        "-O",
        '{"level": 3, "cudagraph_capture_sizes": [1, 2, 4, 8], '
        '"use_inductor": false}',
267
268
    ])
    assert (args.compilation_config.level == 3 and
269
270
            args.compilation_config.cudagraph_capture_sizes == [1, 2, 4, 8]
            and not args.compilation_config.use_inductor)
271

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


283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
def test_prefix_cache_default():
    parser = EngineArgs.add_cli_args(FlexibleArgumentParser())
    args = parser.parse_args([])

    engine_args = EngineArgs.from_cli_args(args=args)
    assert (not engine_args.enable_prefix_caching
            ), "prefix caching defaults to off."

    # 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


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
# yapf: disable
@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"
    ),
    (
        '{"cast_logits_dtype":"bfloat16","sequence_parallel_norm":true,"sequence_parallel_norm_threshold":2048}',
        {
            "cast_logits_dtype": "bfloat16",
            "sequence_parallel_norm": True,
            "sequence_parallel_norm_threshold": 2048,
        },
        "override-neuron-config"
    ),
332
])
333
334
# yapf: enable
def test_composite_arg_parser(arg, expected, option):
335
336
337
338
    parser = EngineArgs.add_cli_args(FlexibleArgumentParser())
    if arg is None:
        args = parser.parse_args([])
    else:
339
340
        args = parser.parse_args([f"--{option}", arg])
    assert getattr(args, option.replace("-", "_")) == expected
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


def test_human_readable_model_len():
    # `exit_on_error` disabled to test invalid values below
    parser = EngineArgs.add_cli_args(
        FlexibleArgumentParser(exit_on_error=False))

    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

    # Capital
    args = parser.parse_args(["--max-model-len", "3K"])
    assert args.max_model_len == 1024 * 3
    args = parser.parse_args(["--max-model-len", "10M"])
    assert args.max_model_len == 2**20 * 10

    # Decimal values
    args = parser.parse_args(["--max-model-len", "10.2k"])
    assert args.max_model_len == 10200
    # ..truncated to the nearest int
    args = parser.parse_args(["--max-model-len", "10.212345k"])
    assert args.max_model_len == 10212

    # Invalid (do not allow decimals with binary multipliers)
    for invalid in ["1a", "pwd", "10.24", "1.23M"]:
        with pytest.raises(ArgumentError):
            args = parser.parse_args(["--max-model-len", invalid])