A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 15, 12:15 UTC
Improve conversion from dataclass configs to argparse arguments (#17303)
Signed-off-by: Harry Mellor <19981378+hmellor@users.noreply.github.com>
Harry Mellor committed
Apr 28, 2025 at 17:22 UTC
f94886946e4c610158c40d3b5909c9ba9da279fe
4 files changed
+245
-154
tests/engine/test_arg_utils.py
+106
-2
@@ -1,14 +1,118 @@
1
# SPDX-License-Identifier: Apache-2.0
2
3
+import json
4
from argparse import ArgumentError, ArgumentTypeError
5
+from contextlib import nullcontext
6
+from dataclasses import dataclass, field
7
+from typing import Literal, Optional
8
9
import pytest
10
7
-from vllm.config import PoolerConfig
8
-from vllm.engine.arg_utils import EngineArgs, nullable_kvs
11
+from vllm.config import PoolerConfig, config
12
+from vllm.engine.arg_utils import (EngineArgs, contains_type, get_kwargs,
13
+ get_type, is_not_builtin, is_type,
14
+ nullable_kvs, optional_type)
15
from vllm.utils import FlexibleArgumentParser
16
17
18
+@pytest.mark.parametrize(("type", "value", "expected"), [
19
+ (int, "42", 42),
20
+ (int, "None", None),
21
+ (float, "3.14", 3.14),
22
+ (float, "None", None),
23
+ (str, "Hello World!", "Hello World!"),
24
+ (str, "None", None),
25
+ (json.loads, '{"foo":1,"bar":2}', {
26
+ "foo": 1,
27
+ "bar": 2
28
+ }),
29
+ (json.loads, "foo=1,bar=2", {
30
+ "foo": 1,
31
+ "bar": 2
32
+ }),
33
+ (json.loads, "None", None),
34
+])
35
+def test_optional_type(type, value, expected):
36
+ optional_type_func = optional_type(type)
37
+ context = nullcontext()
38
+ if value == "foo=1,bar=2":
39
+ context = pytest.warns(DeprecationWarning)
40
+ with context:
41
+ assert optional_type_func(value) == expected
42
+
43
+
44
+@pytest.mark.parametrize(("type_hint", "type", "expected"), [
45
+ (int, int, True),
46
+ (int, float, False),
47
+ (list[int], list, True),
48
+ (list[int], tuple, False),
49
+ (Literal[0, 1], Literal, True),
50
+])
51
+def test_is_type(type_hint, type, expected):
52
+ assert is_type(type_hint, type) == expected
53
+
54
+
55
+@pytest.mark.parametrize(("type_hints", "type", "expected"), [
56
+ ({float, int}, int, True),
57
+ ({int, tuple[int]}, int, True),
58
+ ({int, tuple[int]}, float, False),
59
+ ({str, Literal["x", "y"]}, Literal, True),
60
+])
61
+def test_contains_type(type_hints, type, expected):
62
+ assert contains_type(type_hints, type) == expected
63
+
64
+
65
+@pytest.mark.parametrize(("type_hints", "type", "expected"), [
66
+ ({int, float}, int, int),
67
+ ({int, float}, str, None),
68
+ ({str, Literal["x", "y"]}, Literal, Literal["x", "y"]),
69
+])
70
+def test_get_type(type_hints, type, expected):
71
+ assert get_type(type_hints, type) == expected
72
+
73
+
74
+@config
75
+@dataclass
76
+class DummyConfigClass:
77
+ regular_bool: bool = True
78
+ """Regular bool with default True"""
79
+ optional_bool: Optional[bool] = None
80
+ """Optional bool with default None"""
81
+ optional_literal: Optional[Literal["x", "y"]] = None
82
+ """Optional literal with default None"""
83
+ tuple_n: tuple[int, ...] = field(default_factory=lambda: (1, 2, 3))
84
+ """Tuple with default (1, 2, 3)"""
85
+ tuple_2: tuple[int, int] = field(default_factory=lambda: (1, 2))
86
+ """Tuple with default (1, 2)"""
87
+ list_n: list[int] = field(default_factory=lambda: [1, 2, 3])
88
+ """List with default [1, 2, 3]"""
89
+
90
+
91
+@pytest.mark.parametrize(("type_hint", "expected"), [
92
+ (int, False),
93
+ (DummyConfigClass, True),
94
+])
95
+def test_is_not_builtin(type_hint, expected):
96
+ assert is_not_builtin(type_hint) == expected
97
+
98
+
99
+def test_get_kwargs():
100
+ kwargs = get_kwargs(DummyConfigClass)
101
+ print(kwargs)
102
+
103
+ # bools should not have their type set
104
+ assert kwargs["regular_bool"].get("type") is None
105
+ assert kwargs["optional_bool"].get("type") is None
106
+ # optional literals should have None as a choice
107
+ assert kwargs["optional_literal"]["choices"] == ["x", "y", "None"]
108
+ # tuples should have the correct nargs
109
+ assert kwargs["tuple_n"]["nargs"] == "+"
110
+ assert kwargs["tuple_2"]["nargs"] == 2
111
+ # lists should work
112
+ assert kwargs["list_n"]["type"] is int
113
+ assert kwargs["list_n"]["nargs"] == "+"
114
+
115
+
116
@pytest.mark.parametrize(("arg", "expected"), [
117
(None, dict()),
118
("image=16", {
vllm/engine/arg_utils.py
+125
-138
@@ -11,7 +11,7 @@ from typing import (Any, Callable, Dict, List, Literal, Optional, Type,
11
TypeVar, Union, cast, get_args, get_origin)
12
13
import torch
14
-from typing_extensions import TypeIs
14
+from typing_extensions import TypeIs, deprecated
15
16
import vllm.envs as envs
17
from vllm import version
@@ -48,33 +48,29 @@ TypeHint = Union[type[Any], object]
48
TypeHintT = Union[type[T], object]
49
50
51
-def optional_arg(val: str, return_type: Callable[[str], T]) -> Optional[T]:
52
- if val == "" or val == "None":
53
- return None
54
- try:
55
- return return_type(val)
56
- except ValueError as e:
57
- raise argparse.ArgumentTypeError(
58
- f"Value {val} cannot be converted to {return_type}.") from e
59
-
60
-
61
-def optional_str(val: str) -> Optional[str]:
62
- return optional_arg(val, str)
63
-
64
-
65
-def optional_int(val: str) -> Optional[int]:
66
- return optional_arg(val, int)
51
+def optional_type(
52
+ return_type: Callable[[str], T]) -> Callable[[str], Optional[T]]:
53
54
+ def _optional_type(val: str) -> Optional[T]:
55
+ if val == "" or val == "None":
56
+ return None
57
+ try:
58
+ if return_type is json.loads and not re.match("^{.*}$", val):
59
+ return cast(T, nullable_kvs(val))
60
+ return return_type(val)
61
+ except ValueError as e:
62
+ raise argparse.ArgumentTypeError(
63
+ f"Value {val} cannot be converted to {return_type}.") from e
64
69
-def optional_float(val: str) -> Optional[float]:
70
- return optional_arg(val, float)
65
+ return _optional_type
66
67
73
-def nullable_kvs(val: str) -> Optional[dict[str, int]]:
74
- """NOTE: This function is deprecated, args should be passed as JSON
75
- strings instead.
76
-
77
- Parses a string containing comma separate key [str] to value [int]
68
+@deprecated(
69
+ "Passing a JSON argument as a string containing comma separated key=value "
70
+ "pairs is deprecated. This will be removed in v0.10.0. Please use a JSON "
71
+ "string instead.")
72
+def nullable_kvs(val: str) -> dict[str, int]:
73
+ """Parses a string containing comma separate key [str] to value [int]
74
pairs into a dictionary.
75
76
Args:
@@ -83,10 +79,7 @@ def nullable_kvs(val: str) -> Optional[dict[str, int]]:
79
Returns:
80
Dictionary with parsed values.
81
"""
86
- if len(val) == 0:
87
- return None
88
-
89
- out_dict: Dict[str, int] = {}
82
+ out_dict: dict[str, int] = {}
83
for item in val.split(","):
84
kv_parts = [part.lower().strip() for part in item.split("=")]
85
if len(kv_parts) != 2:
@@ -108,15 +101,103 @@ def nullable_kvs(val: str) -> Optional[dict[str, int]]:
101
return out_dict
102
103
111
-def optional_dict(val: str) -> Optional[dict[str, int]]:
112
- if re.match("^{.*}$", val):
113
- return optional_arg(val, json.loads)
104
+def is_type(type_hint: TypeHint, type: TypeHintT) -> TypeIs[TypeHintT]:
105
+ """Check if the type hint is a specific type."""
106
+ return type_hint is type or get_origin(type_hint) is type
107
+
108
+
109
+def contains_type(type_hints: set[TypeHint], type: TypeHintT) -> bool:
110
+ """Check if the type hints contain a specific type."""
111
+ return any(is_type(type_hint, type) for type_hint in type_hints)
112
+
113
115
- logger.warning(
116
- "Failed to parse JSON string. Attempting to parse as "
117
- "comma-separated key=value pairs. This will be deprecated in a "
118
- "future release.")
119
- return nullable_kvs(val)
114
+def get_type(type_hints: set[TypeHint], type: TypeHintT) -> TypeHintT:
115
+ """Get the specific type from the type hints."""
116
+ return next((th for th in type_hints if is_type(th, type)), None)
117
+
118
+
119
+def is_not_builtin(type_hint: TypeHint) -> bool:
120
+ """Check if the class is not a built-in type."""
121
+ return type_hint.__module__ != "builtins"
122
+
123
+
124
+def get_kwargs(cls: ConfigType) -> dict[str, Any]:
125
+ cls_docs = get_attr_docs(cls)
126
+ kwargs = {}
127
+ for field in fields(cls):
128
+ # Get the default value of the field
129
+ default = field.default
130
+ if field.default_factory is not MISSING:
131
+ default = field.default_factory()
132
+
133
+ # Get the help text for the field
134
+ name = field.name
135
+ help = cls_docs[name]
136
+ # Escape % for argparse
137
+ help = help.replace("%", "%%")
138
+
139
+ # Initialise the kwargs dictionary for the field
140
+ kwargs[name] = {"default": default, "help": help}
141
+
142
+ # Get the set of possible types for the field
143
+ type_hints: set[TypeHint] = set()
144
+ if get_origin(field.type) is Union:
145
+ type_hints.update(get_args(field.type))
146
+ else:
147
+ type_hints.add(field.type)
148
+
149
+ # Set other kwargs based on the type hints
150
+ if contains_type(type_hints, bool):
151
+ # Creates --no-<name> and --<name> flags
152
+ kwargs[name]["action"] = argparse.BooleanOptionalAction
153
+ elif contains_type(type_hints, Literal):
154
+ # Creates choices from Literal arguments
155
+ type_hint = get_type(type_hints, Literal)
156
+ choices = sorted(get_args(type_hint))
157
+ kwargs[name]["choices"] = choices
158
+ choice_type = type(choices[0])
159
+ assert all(type(c) is choice_type for c in choices), (
160
+ "All choices must be of the same type. "
161
+ f"Got {choices} with types {[type(c) for c in choices]}")
162
+ kwargs[name]["type"] = choice_type
163
+ elif contains_type(type_hints, tuple):
164
+ type_hint = get_type(type_hints, tuple)
165
+ types = get_args(type_hint)
166
+ tuple_type = types[0]
167
+ assert all(t is tuple_type for t in types if t is not Ellipsis), (
168
+ "All non-Ellipsis tuple elements must be of the same "
169
+ f"type. Got {types}.")
170
+ kwargs[name]["type"] = tuple_type
171
+ kwargs[name]["nargs"] = "+" if Ellipsis in types else len(types)
172
+ elif contains_type(type_hints, list):
173
+ type_hint = get_type(type_hints, list)
174
+ types = get_args(type_hint)
175
+ assert len(types) == 1, (
176
+ "List type must have exactly one type. Got "
177
+ f"{type_hint} with types {types}")
178
+ kwargs[name]["type"] = types[0]
179
+ kwargs[name]["nargs"] = "+"
180
+ elif contains_type(type_hints, int):
181
+ kwargs[name]["type"] = int
182
+ elif contains_type(type_hints, float):
183
+ kwargs[name]["type"] = float
184
+ elif contains_type(type_hints, dict):
185
+ # Dict arguments will always be optional
186
+ kwargs[name]["type"] = optional_type(json.loads)
187
+ elif (contains_type(type_hints, str)
188
+ or any(is_not_builtin(th) for th in type_hints)):
189
+ kwargs[name]["type"] = str
190
+ else:
191
+ raise ValueError(
192
+ f"Unsupported type {type_hints} for argument {name}.")
193
+
194
+ # If None is in type_hints, make the argument optional.
195
+ # But not if it's a bool, argparse will handle this better.
196
+ if type(None) in type_hints and not contains_type(type_hints, bool):
197
+ kwargs[name]["type"] = optional_type(kwargs[name]["type"])
198
+ if kwargs[name].get("choices"):
199
+ kwargs[name]["choices"].append("None")
200
+ return kwargs
201
202
203
@dataclass
@@ -279,100 +360,6 @@ class EngineArgs:
360
def add_cli_args(parser: FlexibleArgumentParser) -> FlexibleArgumentParser:
361
"""Shared CLI arguments for vLLM engine."""
362
282
- def is_type_in_union(cls: TypeHint, type: TypeHint) -> bool:
283
- """Check if the class is a type in a union type."""
284
- is_union = get_origin(cls) is Union
285
- type_in_union = type in [get_origin(a) or a for a in get_args(cls)]
286
- return is_union and type_in_union
287
-
288
- def get_type_from_union(cls: TypeHint, type: TypeHintT) -> TypeHintT:
289
- """Get the type in a union type."""
290
- for arg in get_args(cls):
291
- if (get_origin(arg) or arg) is type:
292
- return arg
293
- raise ValueError(f"Type {type} not found in union type {cls}.")
294
-
295
- def is_optional(cls: TypeHint) -> TypeIs[Union[Any, None]]:
296
- """Check if the class is an optional type."""
297
- return is_type_in_union(cls, type(None))
298
-
299
- def can_be_type(cls: TypeHint, type: TypeHintT) -> TypeIs[TypeHintT]:
300
- """Check if the class can be of type."""
301
- return cls is type or get_origin(cls) is type or is_type_in_union(
302
- cls, type)
303
-
304
- def is_custom_type(cls: TypeHint) -> bool:
305
- """Check if the class is a custom type."""
306
- return cls.__module__ != "builtins"
307
-
308
- def get_kwargs(cls: ConfigType) -> dict[str, Any]:
309
- cls_docs = get_attr_docs(cls)
310
- kwargs = {}
311
- for field in fields(cls):
312
- # Get the default value of the field
313
- default = field.default
314
- if field.default_factory is not MISSING:
315
- default = field.default_factory()
316
-
317
- # Get the help text for the field
318
- name = field.name
319
- help = cls_docs[name]
320
- # Escape % for argparse
321
- help = help.replace("%", "%%")
322
-
323
- # Initialise the kwargs dictionary for the field
324
- kwargs[name] = {"default": default, "help": help}
325
-
326
- # Make note of if the field is optional and get the actual
327
- # type of the field if it is
328
- optional = is_optional(field.type)
329
- field_type = get_args(
330
- field.type)[0] if optional else field.type
331
-
332
- # Set type, action and choices for the field depending on the
333
- # type of the field
334
- if can_be_type(field_type, bool):
335
- # Creates --no-<name> and --<name> flags
336
- kwargs[name]["action"] = argparse.BooleanOptionalAction
337
- kwargs[name]["type"] = bool
338
- elif can_be_type(field_type, Literal):
339
- # Creates choices from Literal arguments
340
- if is_type_in_union(field_type, Literal):
341
- field_type = get_type_from_union(field_type, Literal)
342
- choices = get_args(field_type)
343
- kwargs[name]["choices"] = choices
344
- choice_type = type(choices[0])
345
- assert all(type(c) is choice_type for c in choices), (
346
- "All choices must be of the same type. "
347
- f"Got {choices} with types {[type(c) for c in choices]}"
348
- )
349
- kwargs[name]["type"] = choice_type
350
- elif can_be_type(field_type, tuple):
351
- if is_type_in_union(field_type, tuple):
352
- field_type = get_type_from_union(field_type, tuple)
353
- dtypes = get_args(field_type)
354
- dtype = dtypes[0]
355
- assert all(
356
- d is dtype for d in dtypes if d is not Ellipsis
357
- ), ("All non-Ellipsis tuple elements must be of the same "
358
- f"type. Got {dtypes}.")
359
- kwargs[name]["type"] = dtype
360
- kwargs[name]["nargs"] = "+"
361
- elif can_be_type(field_type, int):
362
- kwargs[name]["type"] = optional_int if optional else int
363
- elif can_be_type(field_type, float):
364
- kwargs[name][
365
- "type"] = optional_float if optional else float
366
- elif can_be_type(field_type, dict):
367
- kwargs[name]["type"] = optional_dict
368
- elif (can_be_type(field_type, str)
369
- or is_custom_type(field_type)):
370
- kwargs[name]["type"] = optional_str if optional else str
371
- else:
372
- raise ValueError(
373
- f"Unsupported type {field.type} for argument {name}. ")
374
- return kwargs
375
-
363
# Model arguments
364
parser.add_argument(
365
'--model',
@@ -390,13 +377,13 @@ class EngineArgs:
377
'which task to use.')
378
parser.add_argument(
379
'--tokenizer',
393
- type=optional_str,
380
+ type=optional_type(str),
381
default=EngineArgs.tokenizer,
382
help='Name or path of the huggingface tokenizer to use. '
383
'If unspecified, model name or path will be used.')
384
parser.add_argument(
385
"--hf-config-path",
399
- type=optional_str,
386
+ type=optional_type(str),
387
default=EngineArgs.hf_config_path,
388
help='Name or path of the huggingface config to use. '
389
'If unspecified, model name or path will be used.')
@@ -408,21 +395,21 @@ class EngineArgs:
395
'the input. The generated output will contain token ids.')
396
parser.add_argument(
397
'--revision',
411
- type=optional_str,
398
+ type=optional_type(str),
399
default=None,
400
help='The specific model version to use. It can be a branch '
401
'name, a tag name, or a commit id. If unspecified, will use '
402
'the default version.')
403
parser.add_argument(
404
'--code-revision',
418
- type=optional_str,
405
+ type=optional_type(str),
406
default=None,
407
help='The specific revision to use for the model code on '
408
'Hugging Face Hub. It can be a branch name, a tag name, or a '
409
'commit id. If unspecified, will use the default version.')
410
parser.add_argument(
411
'--tokenizer-revision',
425
- type=optional_str,
412
+ type=optional_type(str),
413
default=None,
414
help='Revision of the huggingface tokenizer to use. '
415
'It can be a branch name, a tag name, or a commit id. '
@@ -513,7 +500,7 @@ class EngineArgs:
500
501
parser.add_argument(
502
'--logits-processor-pattern',
516
- type=optional_str,
503
+ type=optional_type(str),
504
default=None,
505
help='Optional regex pattern specifying valid logits processor '
506
'qualified names that can be passed with the `logits_processors` '
@@ -612,7 +599,7 @@ class EngineArgs:
599
# Quantization settings.
600
parser.add_argument('--quantization',
601
'-q',
615
- type=optional_str,
602
+ type=optional_type(str),
603
choices=[*QUANTIZATION_METHODS, None],
604
default=EngineArgs.quantization,
605
help='Method used to quantize the weights. If '
@@ -921,7 +908,7 @@ class EngineArgs:
908
'class without changing the existing functions.')
909
parser.add_argument(
910
"--generation-config",
924
- type=optional_str,
911
+ type=optional_type(str),
912
default="auto",
913
help="The folder path to the generation config. "
914
"Defaults to 'auto', the generation config will be loaded from "
vllm/entrypoints/openai/cli_args.py
+12
-12
@@ -11,7 +11,7 @@ import ssl
11
from collections.abc import Sequence
12
from typing import Optional, Union, get_args
13
14
-from vllm.engine.arg_utils import AsyncEngineArgs, optional_str
14
+from vllm.engine.arg_utils import AsyncEngineArgs, optional_type
15
from vllm.entrypoints.chat_utils import (ChatTemplateContentFormatOption,
16
validate_chat_template)
17
from vllm.entrypoints.openai.serving_models import (LoRAModulePath,
@@ -79,7 +79,7 @@ class PromptAdapterParserAction(argparse.Action):
79
80
def make_arg_parser(parser: FlexibleArgumentParser) -> FlexibleArgumentParser:
81
parser.add_argument("--host",
82
- type=optional_str,
82
+ type=optional_type(str),
83
default=None,
84
help="Host name.")
85
parser.add_argument("--port", type=int, default=8000, help="Port number.")
@@ -108,13 +108,13 @@ def make_arg_parser(parser: FlexibleArgumentParser) -> FlexibleArgumentParser:
108
default=["*"],
109
help="Allowed headers.")
110
parser.add_argument("--api-key",
111
- type=optional_str,
111
+ type=optional_type(str),
112
default=None,
113
help="If provided, the server will require this key "
114
"to be presented in the header.")
115
parser.add_argument(
116
"--lora-modules",
117
- type=optional_str,
117
+ type=optional_type(str),
118
default=None,
119
nargs='+',
120
action=LoRAParserAction,
@@ -126,14 +126,14 @@ def make_arg_parser(parser: FlexibleArgumentParser) -> FlexibleArgumentParser:
126
"\"base_model_name\": \"id\"}``")
127
parser.add_argument(
128
"--prompt-adapters",
129
- type=optional_str,
129
+ type=optional_type(str),
130
default=None,
131
nargs='+',
132
action=PromptAdapterParserAction,
133
help="Prompt adapter configurations in the format name=path. "
134
"Multiple adapters can be specified.")
135
parser.add_argument("--chat-template",
136
- type=optional_str,
136
+ type=optional_type(str),
137
default=None,
138
help="The file path to the chat template, "
139
"or the template in single-line form "
@@ -151,20 +151,20 @@ def make_arg_parser(parser: FlexibleArgumentParser) -> FlexibleArgumentParser:
151
'similar to OpenAI schema. '
152
'Example: ``[{"type": "text", "text": "Hello world!"}]``')
153
parser.add_argument("--response-role",
154
- type=optional_str,
154
+ type=optional_type(str),
155
default="assistant",
156
help="The role name to return if "
157
"``request.add_generation_prompt=true``.")
158
parser.add_argument("--ssl-keyfile",
159
- type=optional_str,
159
+ type=optional_type(str),
160
default=None,
161
help="The file path to the SSL key file.")
162
parser.add_argument("--ssl-certfile",
163
- type=optional_str,
163
+ type=optional_type(str),
164
default=None,
165
help="The file path to the SSL cert file.")
166
parser.add_argument("--ssl-ca-certs",
167
- type=optional_str,
167
+ type=optional_type(str),
168
default=None,
169
help="The CA certificates file.")
170
parser.add_argument(
@@ -180,13 +180,13 @@ def make_arg_parser(parser: FlexibleArgumentParser) -> FlexibleArgumentParser:
180
)
181
parser.add_argument(
182
"--root-path",
183
- type=optional_str,
183
+ type=optional_type(str),
184
default=None,
185
help="FastAPI root_path when app is behind a path based routing proxy."
186
)
187
parser.add_argument(
188
"--middleware",
189
- type=optional_str,
189
+ type=optional_type(str),
190
action="append",
191
default=[],
192
help="Additional ASGI middleware to apply to the app. "
vllm/entrypoints/openai/run_batch.py
+2
-2
@@ -12,7 +12,7 @@ import torch
12
from prometheus_client import start_http_server
13
from tqdm import tqdm
14
15
-from vllm.engine.arg_utils import AsyncEngineArgs, optional_str
15
+from vllm.engine.arg_utils import AsyncEngineArgs, optional_type
16
from vllm.engine.async_llm_engine import AsyncLLMEngine
17
from vllm.entrypoints.logger import RequestLogger, logger
18
# yapf: disable
@@ -61,7 +61,7 @@ def parse_args():
61
"to the output URL.",
62
)
63
parser.add_argument("--response-role",
64
- type=optional_str,
64
+ type=optional_type(str),
65
default="assistant",
66
help="The role name to return if "
67
"`request.add_generation_prompt=True`.")