Support set in the CLI generation (#27031)

Signed-off-by: Harry Mellor <19981378+hmellor@users.noreply.github.com>

tests/engine/test_arg_utils.py

@@ -129,6 +129,8 @@ class DummyConfig:
     """List with literal choices"""
     list_union: list[str | type[object]] = field(default_factory=list)
     """List with union type"""
+    set_n: set[int] = field(default_factory=lambda: {1, 2, 3})
+    """Set with variable length"""
     literal_literal: Literal[Literal[1], Literal[2]] = 1
     """Literal of literals with default 1"""
     json_tip: dict = field(default_factory=dict)
@@ -184,6 +186,9 @@ def test_get_kwargs():
     # 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
+    # sets should work like lists
+    assert kwargs["set_n"]["type"] is int
+    assert kwargs["set_n"]["nargs"] == "+"
     # literals of literals should have merged choices
     assert kwargs["literal_literal"]["choices"] == [1, 2]
     # dict should have json tip in help

vllm/engine/arg_utils.py

@@ -162,6 +162,31 @@ def literal_to_kwargs(type_hints: set[TypeHint]) -> dict[str, Any]:
     return {"type": option_type, kwarg: sorted(options)}
 
 
+def collection_to_kwargs(type_hints: set[TypeHint], type: TypeHint) -> dict[str, Any]:
+    type_hint = get_type(type_hints, type)
+    types = get_args(type_hint)
+    elem_type = types[0]
+
+    # Handle Ellipsis
+    assert all(t is elem_type for t in types if t is not Ellipsis), (
+        f"All non-Ellipsis elements must be of the same type. Got {types}."
+    )
+
+    # Handle Union types
+    if get_origin(elem_type) in {Union, UnionType}:
+        # Union for Union[X, Y] and UnionType for X | Y
+        assert str in get_args(elem_type), (
+            "If element can have multiple types, one must be 'str' "
+            f"(i.e. 'list[int | str]'). Got {elem_type}."
+        )
+        elem_type = str
+
+    return {
+        "type": elem_type,
+        "nargs": "+" if type is not tuple or Ellipsis in types else len(types),
+    }
+
+
 def is_not_builtin(type_hint: TypeHint) -> bool:
     """Check if the class is not a built-in type."""
     return type_hint.__module__ != "builtins"
@@ -251,26 +276,11 @@ def _compute_kwargs(cls: ConfigType) -> dict[str, dict[str, Any]]:
         elif contains_type(type_hints, Literal):
             kwargs[name].update(literal_to_kwargs(type_hints))
         elif contains_type(type_hints, tuple):
-            type_hint = get_type(type_hints, tuple)
+            kwargs[name].update(collection_to_kwargs(type_hints, tuple))
-            types = get_args(type_hint)
-            tuple_type = types[0]
-            assert all(t is tuple_type for t in types if t is not Ellipsis), (
-                "All non-Ellipsis tuple elements must be of the same "
-                f"type. Got {types}."
-            )
-            kwargs[name]["type"] = tuple_type
-            kwargs[name]["nargs"] = "+" if Ellipsis in types else len(types)
         elif contains_type(type_hints, list):
-            type_hint = get_type(type_hints, list)
+            kwargs[name].update(collection_to_kwargs(type_hints, list))
-            types = get_args(type_hint)
+        elif contains_type(type_hints, set):
-            list_type = types[0]
+            kwargs[name].update(collection_to_kwargs(type_hints, set))
-            if get_origin(list_type) in {Union, UnionType}:
-                # Union for Union[X, Y] and UnionType for X | Y
-                msg = "List type must contain str if it is a Union."
-                assert str in get_args(list_type), msg
-                list_type = str
-            kwargs[name]["type"] = list_type
-            kwargs[name]["nargs"] = "+"
         elif contains_type(type_hints, int):
             kwargs[name]["type"] = int
             # Special case for large integers