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 14, 05:30 UTC
Support `set` in the CLI generation (#27031)
Signed-off-by: Harry Mellor <19981378+hmellor@users.noreply.github.com>
Harry Mellor committed
Oct 16, 2025 at 19:07 UTC
aa255ff55aecfb1a90f1359a2ead315f4d29cff7
2 files changed
+34
-19
tests/engine/test_arg_utils.py
+5
@@ -129,6 +129,8 @@ class DummyConfig:
129
"""List with literal choices"""
130
list_union: list[str | type[object]] = field(default_factory=list)
131
"""List with union type"""
132
+ set_n: set[int] = field(default_factory=lambda: {1, 2, 3})
133
+ """Set with variable length"""
134
literal_literal: Literal[Literal[1], Literal[2]] = 1
135
"""Literal of literals with default 1"""
136
json_tip: dict = field(default_factory=dict)
@@ -184,6 +186,9 @@ def test_get_kwargs():
186
# lists with unions should become str type.
187
# If not, we cannot know which type to use for parsing
188
assert kwargs["list_union"]["type"] is str
189
+ # sets should work like lists
190
+ assert kwargs["set_n"]["type"] is int
191
+ assert kwargs["set_n"]["nargs"] == "+"
192
# literals of literals should have merged choices
193
assert kwargs["literal_literal"]["choices"] == [1, 2]
194
# dict should have json tip in help
vllm/engine/arg_utils.py
+29
-19
@@ -162,6 +162,31 @@ def literal_to_kwargs(type_hints: set[TypeHint]) -> dict[str, Any]:
162
return {"type": option_type, kwarg: sorted(options)}
163
164
165
+def collection_to_kwargs(type_hints: set[TypeHint], type: TypeHint) -> dict[str, Any]:
166
+ type_hint = get_type(type_hints, type)
167
+ types = get_args(type_hint)
168
+ elem_type = types[0]
169
+
170
+ # Handle Ellipsis
171
+ assert all(t is elem_type for t in types if t is not Ellipsis), (
172
+ f"All non-Ellipsis elements must be of the same type. Got {types}."
173
+ )
174
+
175
+ # Handle Union types
176
+ if get_origin(elem_type) in {Union, UnionType}:
177
+ # Union for Union[X, Y] and UnionType for X | Y
178
+ assert str in get_args(elem_type), (
179
+ "If element can have multiple types, one must be 'str' "
180
+ f"(i.e. 'list[int | str]'). Got {elem_type}."
181
+ )
182
+ elem_type = str
183
+
184
+ return {
185
+ "type": elem_type,
186
+ "nargs": "+" if type is not tuple or Ellipsis in types else len(types),
187
+ }
188
+
189
+
190
def is_not_builtin(type_hint: TypeHint) -> bool:
191
"""Check if the class is not a built-in type."""
192
return type_hint.__module__ != "builtins"
@@ -251,26 +276,11 @@ def _compute_kwargs(cls: ConfigType) -> dict[str, dict[str, Any]]:
276
elif contains_type(type_hints, Literal):
277
kwargs[name].update(literal_to_kwargs(type_hints))
278
elif contains_type(type_hints, tuple):
254
- type_hint = get_type(type_hints, tuple)
255
- types = get_args(type_hint)
256
- tuple_type = types[0]
257
- assert all(t is tuple_type for t in types if t is not Ellipsis), (
258
- "All non-Ellipsis tuple elements must be of the same "
259
- f"type. Got {types}."
260
- )
261
- kwargs[name]["type"] = tuple_type
262
- kwargs[name]["nargs"] = "+" if Ellipsis in types else len(types)
279
+ kwargs[name].update(collection_to_kwargs(type_hints, tuple))
280
elif contains_type(type_hints, list):
264
- type_hint = get_type(type_hints, list)
265
- types = get_args(type_hint)
266
- list_type = types[0]
267
- if get_origin(list_type) in {Union, UnionType}:
268
- # Union for Union[X, Y] and UnionType for X | Y
269
- msg = "List type must contain str if it is a Union."
270
- assert str in get_args(list_type), msg
271
- list_type = str
272
- kwargs[name]["type"] = list_type
273
- kwargs[name]["nargs"] = "+"
281
+ kwargs[name].update(collection_to_kwargs(type_hints, list))
282
+ elif contains_type(type_hints, set):
283
+ kwargs[name].update(collection_to_kwargs(type_hints, set))
284
elif contains_type(type_hints, int):
285
kwargs[name]["type"] = int
286
# Special case for large integers