chore: import upstream snapshot with attribution

python/sglang/srt/arg_groups/arg_utils.py

@@ -0,0 +1,300 @@
+# Copyright 2023-2024 SGLang Team
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Utilities for auto-deriving argparse CLI arguments from dataclass fields.
+
+Usage::
+
+    from sglang.srt.arg_groups.arg_utils import A, Arg, add_cli_args_from_dataclass
+
+    @dataclasses.dataclass
+    class ServerArgs:
+        # Simple fields — bare string is the help text:
+        host: A[str, "The host of the HTTP server."] = "127.0.0.1"
+        port: A[int, "The port of the HTTP server."] = 30000
+        trust_remote_code: A[bool, "Whether to allow custom models."] = False
+        tokenizer_path: A[Optional[str], "The path of the tokenizer."] = None
+
+        # Fields with extra metadata — use Arg(...):
+        model_path: A[str, Arg(help="Path to model weights.", aliases=["--model"])]
+        load_format: A[str, Arg(help="Format.", choices=CHOICES)] = "auto"
+
+        @staticmethod
+        def add_cli_args(parser):
+            add_cli_args_from_dataclass(parser, ServerArgs)
+
+``A`` is a short alias for ``typing.Annotated``. A bare ``str`` inside the
+annotation is equivalent to ``Arg(help=that_string)``.
+"""
+
+from __future__ import annotations
+
+import dataclasses
+import functools
+import types
+from typing import (
+    Annotated,
+    Any,
+    Callable,
+    List,
+    Literal,
+    Optional,
+    Union,
+    get_args,
+    get_origin,
+    get_type_hints,
+)
+
+A = Annotated
+
+
+@dataclasses.dataclass(frozen=True)
+class Arg:
+    """CLI argument metadata attached to a dataclass field via ``Annotated``."""
+
+    help: str = ""
+    choices: Optional[list] = None
+    aliases: Optional[List[str]] = None
+    cli_name: Optional[str] = None
+    type_parser: Optional[Callable] = None
+    nargs: Optional[str] = None
+    required: Optional[bool] = None
+    action: Optional[Any] = None
+    action_kwargs: Optional[dict] = None
+    const: Optional[Any] = None
+    # When True, this field is skipped by add_cli_args_from_dataclass.
+    # Use for fields that have no CLI surface (e.g. injected via Python only).
+    no_cli: bool = False
+    # When True, this field may be written by config resolution (model
+    # overrides and post-process passes): it is part of the whitelist accepted
+    # by the declaration stash, and its resolved value materializes onto the
+    # field at the end of __post_init__.
+    resolvable: bool = False
+
+
+@functools.lru_cache(maxsize=None)
+def resolvable_fields(cls) -> frozenset:
+    """Names of ``cls`` dataclass fields whose ``Arg`` metadata declares
+    ``resolvable=True`` — the whitelist for config resolution.
+
+    Non-dataclass types (e.g. mock config objects in tests) have no Arg
+    metadata and yield an empty whitelist."""
+    if not dataclasses.is_dataclass(cls):
+        return frozenset()
+    hints = get_type_hints(cls, include_extras=True)
+    names = set()
+    for field in dataclasses.fields(cls):
+        _, arg = _unwrap_annotated(hints.get(field.name, field.type))
+        if arg is not None and arg.resolvable:
+            names.add(field.name)
+    return frozenset(names)
+
+
+# ---------------------------------------------------------------------------
+# Internal helpers
+# ---------------------------------------------------------------------------
+
+_MISSING = dataclasses.MISSING
+
+
+def _unwrap_annotated(tp):
+    """Return (inner_type, Arg | None) from ``Annotated[T, Arg(...)]``.
+
+    Also accepts a bare string as shorthand: ``Annotated[T, "help text"]``
+    is equivalent to ``Annotated[T, Arg(help="help text")]``.
+    """
+    origin = get_origin(tp)
+    if origin is Annotated:
+        args = get_args(tp)
+        inner = args[0]
+        for a in args[1:]:
+            if isinstance(a, Arg):
+                return inner, a
+            if isinstance(a, str):
+                return inner, Arg(help=a)
+        return inner, None
+    return tp, None
+
+
+def _unwrap_optional(tp):
+    """If tp is Optional[X] (i.e. Union[X, None]), return (X, True). Else (tp, False)."""
+    origin = get_origin(tp)
+    is_union = origin is Union or (
+        hasattr(types, "UnionType") and origin is types.UnionType
+    )
+    if is_union:
+        args = get_args(tp)
+        non_none = [a for a in args if a is not type(None)]
+        if len(non_none) == 1:
+            return non_none[0], True
+    return tp, False
+
+
+def _unwrap_literal(tp):
+    """If tp is Literal[...], return list of values. Else None."""
+    origin = get_origin(tp)
+    if origin is Literal:
+        return list(get_args(tp))
+    return None
+
+
+def _infer_type_func(tp):
+    """Map a Python type annotation to an argparse ``type=`` callable."""
+    if tp is str:
+        return str
+    if tp is int:
+        return int
+    if tp is float:
+        return float
+    return str
+
+
+def _field_default(field):
+    """Return the default value for a dataclass field, or _MISSING."""
+    if field.default is not _MISSING:
+        return field.default
+    if field.default_factory is not _MISSING:
+        return field.default_factory()
+    return _MISSING
+
+
+def _field_to_cli_name(name: str) -> str:
+    """Convert a field name like ``model_path`` to ``--model-path``."""
+    return "--" + name.replace("_", "-")
+
+
+# ---------------------------------------------------------------------------
+# Public API
+# ---------------------------------------------------------------------------
+
+
+def add_cli_args_from_dataclass(parser, cls, *, fields: Optional[List[str]] = None):
+    """Add argparse arguments for every ``A[T, "help"]`` or ``A[T, Arg(...)]`` field.
+
+    Fields without an ``Arg`` or bare-string annotation are silently skipped —
+    they must still be registered manually (this allows incremental migration).
+
+    Parameters
+    ----------
+    parser : argparse.ArgumentParser
+    cls : dataclass type
+    fields : optional list of field names to include. If None, all fields with
+        ``Arg`` annotations are included.
+    """
+    hints = get_type_hints(cls, include_extras=True)
+
+    for field in dataclasses.fields(cls):
+        if fields is not None and field.name not in fields:
+            continue
+
+        hint = hints.get(field.name)
+        if hint is None:
+            continue
+
+        raw_type, arg_meta = _unwrap_annotated(hint)
+        if arg_meta is None:
+            continue
+        if arg_meta.no_cli:
+            continue
+
+        cli_name = arg_meta.cli_name or _field_to_cli_name(field.name)
+        names = [cli_name] + (arg_meta.aliases or [])
+        default = _field_default(field)
+        # Anchor dest to the field name so argparse stores the value
+        # under the dataclass attribute directly, even when cli_name
+        # differs (e.g. --tensor-parallel-size → tp_size).
+        auto_dest = cli_name.lstrip("-").replace("-", "_")
+        dest_kwarg = {"dest": field.name} if field.name != auto_dest else {}
+
+        # Handle custom action
+        if arg_meta.action is not None:
+            kwargs = {
+                "action": arg_meta.action,
+                "help": arg_meta.help,
+                **dest_kwarg,
+            }
+            if default is not _MISSING:
+                kwargs["default"] = default
+            if arg_meta.action_kwargs:
+                kwargs.update(arg_meta.action_kwargs)
+            parser.add_argument(*names, **kwargs)
+            continue
+
+        # Unwrap Optional
+        inner_type, is_optional = _unwrap_optional(raw_type)
+
+        # Check for Literal — auto-derive choices
+        literal_vals = _unwrap_literal(inner_type)
+        if literal_vals is not None:
+            choices = arg_meta.choices or literal_vals
+            # Infer type from first literal value
+            val_type = type(literal_vals[0]) if literal_vals else str
+            type_func = arg_meta.type_parser or _infer_type_func(val_type)
+            kwargs = dict(
+                type=type_func, choices=choices, help=arg_meta.help, **dest_kwarg
+            )
+            if default is not _MISSING:
+                kwargs["default"] = default
+            if arg_meta.const is not None:
+                kwargs["const"] = arg_meta.const
+            parser.add_argument(*names, **kwargs)
+            continue
+
+        # Check for List[X] — but skip if type_parser is set (the parser
+        # handles the whole value as a single string, e.g. json_list_type).
+        origin = get_origin(inner_type)
+        if (origin is list or origin is List) and arg_meta.type_parser is None:
+            elem_args = get_args(inner_type)
+            elem_type = elem_args[0] if elem_args else str
+            type_func = _infer_type_func(elem_type)
+            nargs = arg_meta.nargs or "+"
+            kwargs = dict(
+                type=type_func,
+                nargs=nargs,
+                help=arg_meta.help,
+                **dest_kwarg,
+            )
+            if arg_meta.choices:
+                kwargs["choices"] = arg_meta.choices
+            if default is not _MISSING:
+                kwargs["default"] = default
+            if arg_meta.const is not None:
+                kwargs["const"] = arg_meta.const
+            parser.add_argument(*names, **kwargs)
+            continue
+
+        # Bool → store_true
+        if inner_type is bool:
+            kwargs = dict(action="store_true", help=arg_meta.help, **dest_kwarg)
+            if default is not _MISSING:
+                kwargs["default"] = default
+            parser.add_argument(*names, **kwargs)
+            continue
+
+        # Scalar types (str, int, float, etc.)
+        type_func = arg_meta.type_parser or _infer_type_func(inner_type)
+        kwargs = dict(type=type_func, help=arg_meta.help, **dest_kwarg)
+        if arg_meta.choices:
+            kwargs["choices"] = arg_meta.choices
+        if arg_meta.nargs:
+            kwargs["nargs"] = arg_meta.nargs
+        if default is not _MISSING:
+            kwargs["default"] = default
+        if arg_meta.const is not None:
+            kwargs["const"] = arg_meta.const
+        if (
+            arg_meta.required is True
+            or (arg_meta.required is None and default is _MISSING)
+        ) and any(name.startswith("-") for name in names):
+            kwargs["required"] = True
+        parser.add_argument(*names, **kwargs)