[refactor] Add the resolved-flags tier + resolvable-field metadata (stack 4/15) (#30066)

Co-authored-by: Claude Fable 5 <noreply@anthropic.com>

python/sglang/srt/arg_groups/arg_utils.py

@@ -40,6 +40,7 @@ annotation is equivalent to ``Arg(help=that_string)``.
 from __future__ import annotations
 
 import dataclasses
+import functools
 import types
 from typing import (
     Annotated,
@@ -74,6 +75,24 @@ class Arg:
     # 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 resolved by model overrides: it is part of
+    # the whitelist accepted by the apply_model_overrides gate, and its
+    # resolved value lives on the flags tier (the server_args field itself
+    # stays the pristine user input).
+    model_overridable: bool = False
+
+
+@functools.lru_cache(maxsize=None)
+def model_overridable_fields(cls) -> frozenset:
+    """Names of ``cls`` dataclass fields whose ``Arg`` metadata declares
+    ``model_overridable=True`` — the whitelist for model-override resolution."""
+    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.model_overridable:
+            names.add(field.name)
+    return frozenset(names)
 
 
 # ---------------------------------------------------------------------------

python/sglang/srt/runtime_context.py

@@ -27,10 +27,19 @@ tier). The context owns the storage: publishing goes through
 ``set_global_server_args_for_scheduler`` / ``get_global_server_args`` in
 ``server_args.py`` are thin shims over this slot), and the object is returned
 by reference — the same live instance everywhere, never a copy.
+
+``get_flags()`` returns the resolved-flags tier: what the system *resolved*
+the configuration to (``server_args`` stays the pristine user input). Flags
+live in typed dataclass groups (``flags.attn`` / ``flags.moe`` / flat generic
+leaves on ``flags`` itself); reads and writes are plain attribute access.
+Static groups are writable during resolution and locked by ``freeze()``;
+``flags.capture`` stays writable (capture-time state). Each group offers a
+transactional, test-only ``override(**kw)`` that also works on frozen groups.
 """
 
 from __future__ import annotations
 
+import dataclasses
 from contextlib import contextmanager
 from typing import TYPE_CHECKING, Any
 
@@ -215,15 +224,125 @@ class ParallelContext:
         return self._v("attn_cp_group", _ps().get_attn_cp_group)
 
 
-class RuntimeContext:
-    """Container for the structured runtime accessors; exposes ``parallel`` and
-    ``server_args``."""
+class _FlagGroupBase:
+    """Shared flag-group behavior: typo-safe writes + transactional ``override()``.
 
-    __slots__ = ("parallel", "_server_args")
+    Groups are plain dataclasses; ``__dataclass_fields__`` is the single source
+    of truth for which leaves exist, so a mistyped name fails loudly instead of
+    creating a stray attribute.
+    """
+
+    def __setattr__(self, name: str, value: Any) -> None:
+        if name not in type(self).__dataclass_fields__:
+            raise AttributeError(
+                f"{type(self).__name__} has no flag '{name}' (leaves are "
+                "declared as dataclass fields; check for typos)"
+            )
+        if getattr(self, "_frozen", False):
+            raise RuntimeError(
+                f"{type(self).__name__} is frozen; cannot write '{name}'. "
+                "Test-scoped changes go through override()."
+            )
+        object.__setattr__(self, name, value)
+
+    @contextmanager
+    def override(self, **kwargs):
+        """Temporarily force flag values, restoring on exit. Transactional
+        (keys validated before any write) and usable on frozen groups — this
+        is the test-only injection primitive."""
+        fields = type(self).__dataclass_fields__
+        unknown = set(kwargs) - set(fields)
+        if unknown:
+            raise ValueError(
+                f"unknown flag(s) for {type(self).__name__}: {sorted(unknown)}"
+            )
+        saved = {name: getattr(self, name) for name in kwargs}
+        for name, value in kwargs.items():
+            object.__setattr__(self, name, value)
+        try:
+            yield self
+        finally:
+            for name, value in saved.items():
+                object.__setattr__(self, name, value)
+
+
+class _StaticFlags(_FlagGroupBase):
+    """Static flag-group: writable during resolution, locked by ``freeze()``."""
+
+    def freeze(self) -> None:
+        object.__setattr__(self, "_frozen", True)
+
+    @property
+    def frozen(self) -> bool:
+        return getattr(self, "_frozen", False)
+
+
+@dataclasses.dataclass
+class AttnFlags(_StaticFlags):
+    """Attention-family resolved flags (leaves arrive with the V3 sweeps)."""
+
+
+@dataclasses.dataclass
+class MoeFlags(_StaticFlags):
+    """MoE-family resolved flags (leaves arrive with the V3 sweeps)."""
+
+
+@dataclasses.dataclass
+class CaptureFlags(_FlagGroupBase):
+    """Capture-time flags; never frozen (written during cuda-graph capture)."""
+
+
+@dataclasses.dataclass
+class Flags(_StaticFlags):
+    """Root of the resolved-flags tier.
+
+    Family groups hang off it (``flags.attn`` / ``flags.moe`` / ``flags.capture``);
+    single generic flags live flat on this container, declared as fields here.
+    ``freeze()`` locks the container and every static sub-group; ``capture``
+    stays writable.
+    """
+
+    attn: AttnFlags = dataclasses.field(default_factory=AttnFlags)
+    moe: MoeFlags = dataclasses.field(default_factory=MoeFlags)
+    capture: CaptureFlags = dataclasses.field(default_factory=CaptureFlags)
+
+    def freeze(self) -> None:
+        for field in dataclasses.fields(self):
+            value = getattr(self, field.name)
+            if isinstance(value, _StaticFlags):
+                value.freeze()
+        super().freeze()
+
+
+# Resolved-config field name → dotted flag-leaf path (e.g. a V3 sweep adds
+# "use_mla_backend": "attn.use_mla_backend"). Fields not listed default to a
+# flat leaf of the same name on the Flags container. Populated per field
+# family as readers migrate; empty in the skeleton.
+FLAG_LEAF_MAP: dict[str, str] = {}
+
+
+def resolve_flag_leaf(
+    flags: Flags, field: str, *, leaf_map: dict[str, str] | None = None
+) -> tuple[Any, str]:
+    """Return ``(owning group, leaf attribute name)`` for a resolved-config field."""
+    path = (FLAG_LEAF_MAP if leaf_map is None else leaf_map).get(field, field)
+    owner: Any = flags
+    *groups, leaf = path.split(".")
+    for part in groups:
+        owner = getattr(owner, part)
+    return owner, leaf
+
+
+class RuntimeContext:
+    """Container for the structured runtime accessors; exposes ``parallel``,
+    ``server_args``, and ``flags``."""
+
+    __slots__ = ("parallel", "_server_args", "flags")
 
     def __init__(self, parallel: ParallelContext):
         self.parallel = parallel
         self._server_args: ServerArgs | None = None
+        self.flags = Flags()
 
     @property
     def server_args(self) -> ServerArgs:
@@ -260,9 +379,15 @@ def get_server_args() -> ServerArgs:
     return _CONTEXT.server_args
 
 
+def get_flags() -> Flags:
+    return _CONTEXT.flags
+
+
 def reset_context() -> None:
-    """Clear the context-owned store (unit-test teardown).
+    """Clear the context-owned store (unit-test teardown): drop the published
+    ``server_args`` and install a fresh, unfrozen ``Flags``.
 
     Wrapper subsystems (``parallel``) hold no state and are unaffected.
     """
     _CONTEXT._server_args = None
+    _CONTEXT.flags = Flags()

test/registered/unit/test_model_overrides.py

@@ -0,0 +1,42 @@
+"""Unit tests for the model-override machinery: whitelist metadata (V3a)."""
+
+from sglang.test.ci.ci_register import register_cpu_ci
+
+register_cpu_ci(est_time=5, suite="base-a-test-cpu")
+
+import dataclasses
+import unittest
+from typing import Optional
+
+from sglang.srt.arg_groups.arg_utils import A, Arg, model_overridable_fields
+from sglang.test.test_utils import CustomTestCase
+
+
+@dataclasses.dataclass
+class _FakeArgs:
+    plain: A[int, "help text only"] = 0
+    resolved_by_model: A[str, Arg(help="x", model_overridable=True)] = "auto"
+    also_resolved: A[Optional[int], Arg(help="y", model_overridable=True)] = None
+    metadata_but_not_overridable: A[bool, Arg(help="z")] = False
+
+
+class TestModelOverridableWhitelist(CustomTestCase):
+    def test_arg_defaults_to_not_overridable(self):
+        self.assertFalse(Arg().model_overridable)
+
+    def test_whitelist_derivation_from_annotated_metadata(self):
+        self.assertEqual(
+            model_overridable_fields(_FakeArgs),
+            frozenset({"resolved_by_model", "also_resolved"}),
+        )
+
+    def test_server_args_whitelist_empty_at_skeleton(self):
+        # No ServerArgs field is tagged yet: the V3 sweeps whitelist fields
+        # one family at a time. This pin makes accidental tagging visible.
+        from sglang.srt.server_args import ServerArgs
+
+        self.assertEqual(model_overridable_fields(ServerArgs), frozenset())
+
+
+if __name__ == "__main__":
+    unittest.main()

test/registered/unit/test_runtime_context.py

@@ -4,17 +4,23 @@ from sglang.test.ci.ci_register import register_cpu_ci
 
 register_cpu_ci(est_time=5, suite="base-a-test-cpu")
 
+import dataclasses
 import unittest
 from unittest.mock import patch
 
 import sglang.srt.server_args as server_args_module
 from sglang.srt.runtime_context import (
+    Flags,
     ParallelContext,
     RuntimeContext,
+    _FlagGroupBase,
+    _StaticFlags,
     get_context,
+    get_flags,
     get_parallel,
     get_server_args,
     reset_context,
+    resolve_flag_leaf,
 )
 from sglang.test.test_utils import CustomTestCase
 
@@ -208,5 +214,92 @@ class TestServerArgsOwnership(_IsolatedServerArgs):
         self.assertFalse(hasattr(server_args_module, "_global_server_args"))
 
 
+@dataclasses.dataclass
+class _FakeStaticGroup(_StaticFlags):
+    alpha: int = 1
+    beta: str = "b"
+
+
+@dataclasses.dataclass
+class _FakeCaptureGroup(_FlagGroupBase):
+    gamma: int = 0
+
+
+class TestFlagsTier(_IsolatedServerArgs):
+    """V3a skeleton: typed dataclass groups, freeze guard, override primitive."""
+
+    def test_wiring_and_groups(self):
+        flags = get_flags()
+        self.assertIs(flags, get_context().flags)
+        self.assertIsInstance(flags, Flags)
+        for group in ("attn", "moe", "capture"):
+            self.assertTrue(hasattr(flags, group))
+        self.assertFalse(flags.frozen)
+
+    def test_typo_safety(self):
+        group = _FakeStaticGroup()
+        with self.assertRaises(AttributeError):
+            group.alpha_misspelled = 2  # undeclared leaf
+        with self.assertRaises(AttributeError):
+            get_flags().not_a_flag = 1
+
+    def test_static_group_writable_until_freeze(self):
+        group = _FakeStaticGroup()
+        group.alpha = 5
+        self.assertEqual(group.alpha, 5)
+        group.freeze()
+        with self.assertRaises(RuntimeError):
+            group.alpha = 6
+        self.assertEqual(group.alpha, 5)
+
+    def test_override_is_transactional_and_works_on_frozen(self):
+        group = _FakeStaticGroup()
+        group.freeze()
+        with group.override(alpha=99, beta="x"):
+            self.assertEqual(group.alpha, 99)
+            self.assertEqual(group.beta, "x")
+        self.assertEqual(group.alpha, 1)
+        self.assertEqual(group.beta, "b")
+        with self.assertRaises(ValueError):
+            with group.override(alpha=2, gamma=3):  # gamma undeclared
+                pass
+        self.assertEqual(group.alpha, 1)  # validated before any write
+
+    def test_non_static_group_has_no_freeze(self):
+        group = _FakeCaptureGroup()
+        group.gamma = 42
+        self.assertEqual(group.gamma, 42)
+        self.assertFalse(hasattr(group, "freeze"))
+
+    def test_container_freeze_cascades_except_capture(self):
+        flags = Flags()  # fresh container, not the process singleton
+        flags.freeze()
+        self.assertTrue(flags.frozen)
+        self.assertTrue(flags.attn.frozen)
+        self.assertTrue(flags.moe.frozen)
+        with self.assertRaises(RuntimeError):
+            flags.attn = flags.attn  # container leaves lock too
+        self.assertFalse(getattr(flags.capture, "_frozen", False))
+
+    def test_resolve_flag_leaf_flat_default_and_mapped(self):
+        flags = Flags()
+        owner, leaf = resolve_flag_leaf(flags, "some_field")
+        self.assertIs(owner, flags)
+        self.assertEqual(leaf, "some_field")
+        owner, leaf = resolve_flag_leaf(flags, "x", leaf_map={"x": "attn.x"})
+        self.assertIs(owner, flags.attn)
+        self.assertEqual(leaf, "x")
+
+    def test_reset_context_installs_fresh_unfrozen_flags(self):
+        try:
+            old = get_flags()
+            old.freeze()
+            reset_context()
+            self.assertIsNot(get_flags(), old)
+            self.assertFalse(get_flags().frozen)
+        finally:
+            reset_context()  # never leave the singleton frozen for other tests
+
+
 if __name__ == "__main__":
     unittest.main()