diff --git a/python/sglang/srt/arg_groups/arg_utils.py b/python/sglang/srt/arg_groups/arg_utils.py index d4af742a6..c7e20720f 100644 --- a/python/sglang/srt/arg_groups/arg_utils.py +++ b/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) # --------------------------------------------------------------------------- diff --git a/python/sglang/srt/runtime_context.py b/python/sglang/srt/runtime_context.py index ca62cb3c7..4612635dd 100644 --- a/python/sglang/srt/runtime_context.py +++ b/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() diff --git a/test/registered/unit/test_model_overrides.py b/test/registered/unit/test_model_overrides.py new file mode 100644 index 000000000..b87923476 --- /dev/null +++ b/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() diff --git a/test/registered/unit/test_runtime_context.py b/test/registered/unit/test_runtime_context.py index 6d4c7f024..47b556bd5 100644 --- a/test/registered/unit/test_runtime_context.py +++ b/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()