config: preserve resolved config across nested publishes + mutation ratchets (#33011)

python/sglang/srt/runtime_context.py

@@ -671,6 +671,34 @@ def _build_config_bags(server_args: Any) -> dict:
     return tops
 
 
+def _snapshot_bag_values(bags: dict | None) -> dict | None:
+    """Per-leaf value snapshot of a config-bag tree (bags are mutated in
+    place by ``override``, so reference snapshots alias live state)."""
+    if bags is None:
+        return None
+    snap: dict = {}
+
+    def walk(prefix: str, bag) -> None:
+        snap[prefix] = dict(object.__getattribute__(bag, "_fields"))
+        for name, sub in object.__getattribute__(bag, "_subs").items():
+            walk(f"{prefix}.{name}", sub)
+
+    for name, bag in bags.items():
+        walk(name, bag)
+    return snap
+
+
+def _restore_bag_values(bags: dict, snap: dict) -> None:
+    def walk(prefix: str, bag) -> None:
+        for key, value in snap[prefix].items():
+            bag._set(key, value)
+        for name, sub in object.__getattribute__(bag, "_subs").items():
+            walk(f"{prefix}.{name}", sub)
+
+    for name, bag in bags.items():
+        walk(name, bag)
+
+
 class RuntimeContext:
     """Container for the structured runtime accessors; exposes ``parallel``,
     ``server_args``, the resolved config namespace bags, ``flags``,
@@ -858,6 +886,31 @@ class RuntimeContext:
         """
         return _ServerArgsOverride(self, fields)
 
+    @contextmanager
+    def preserve_config(self):
+        """Snapshot the full config lifecycle and reinstate it verbatim on exit.
+
+        For nested construction steps that publish a private ``ServerArgs``
+        copy (e.g. a draft-worker build) and must leave the enclosing
+        lifecycle — including its post-publish overrides — untouched.
+        """
+        prev_server_args = self._server_args
+        prev_bags = self._config_bags
+        prev_bag_values = _snapshot_bag_values(prev_bags)
+        prev_overrides_log = list(self._overrides_log)
+        prev_parallel_config = self.parallel._config
+        prev_capture = self.flags.capture.enable_torch_compile
+        try:
+            yield
+        finally:
+            self._server_args = prev_server_args
+            self._config_bags = prev_bags
+            if prev_bags is not None:
+                _restore_bag_values(prev_bags, prev_bag_values)
+            self._overrides_log = prev_overrides_log
+            self.parallel._config = prev_parallel_config
+            self.flags.capture.enable_torch_compile = prev_capture
+
 
 class _ServerArgsOverride:
     """Scoped config override (see ``RuntimeContext.override_server_args``).

python/sglang/srt/speculative/draft_worker_common.py

@@ -10,7 +10,7 @@ import torch
 from sglang.srt.layers.logits_processor import LogitsProcessorOutput
 from sglang.srt.managers.tp_worker import TpModelWorker
 from sglang.srt.model_executor.forward_batch_info import CaptureHiddenMode
-from sglang.srt.runtime_context import get_context, get_server_args
+from sglang.srt.runtime_context import get_context
 from sglang.srt.server_args import ServerArgs
 from sglang.srt.speculative.dflash_info import DFlashVerifyInput
 from sglang.srt.speculative.dflash_info_v2 import DFlashDraftInputV2
@@ -97,8 +97,7 @@ def build_draft_tp_worker(
         context_length=target_model_config.context_len,
     )
 
-    saved_server_args = get_server_args()
-    try:
+    with get_context().preserve_config():
         draft_worker = TpModelWorker(
             server_args=draft_server_args,
             gpu_id=gpu_id,
@@ -106,8 +105,6 @@ def build_draft_tp_worker(
             nccl_port=nccl_port,
             is_draft_worker=True,
         )
-    finally:
-        get_context().set_server_args(saved_server_args)
 
     draft_model_runner = draft_worker.model_runner
     draft_worker.draft_runner = draft_model_runner

test/registered/unit/test_migration_deferral_ratchet.py

@@ -0,0 +1,53 @@
+"""Ratchet guard: config-namespace-migration test deferrals may only decrease.
+
+A set of unit tests was module-skipped ("Temporarily skipped during the
+ServerArgs config-namespace migration") because their fixtures inject config
+in ways the namespace accessors cannot see; they are recovered together with
+the reader migration. Nothing else enforces that list: without this pin a new
+skip can ride in unnoticed, and the already-skipped files keep growing test
+code nobody has ever run. The count is exact: un-skipping a file must lower
+the baseline to lock in the recovery, and no new deferral may appear.
+"""
+
+from sglang.test.ci.ci_register import register_cpu_ci
+
+register_cpu_ci(est_time=5, suite="base-a-test-cpu")
+
+import unittest
+from pathlib import Path
+
+from sglang.test.test_utils import CustomTestCase
+
+# test/registered/unit/<this file> -> test/
+_TEST_ROOT = Path(__file__).resolve().parents[2]
+
+_MARKER = "config-namespace migration"
+
+_BASELINE = 15
+
+
+class TestMigrationDeferralRatchet(CustomTestCase):
+    def test_deferred_test_files_match_the_baseline(self):
+        deferred = sorted(
+            p.relative_to(_TEST_ROOT).as_posix()
+            for p in _TEST_ROOT.rglob("*.py")
+            if p.name != Path(__file__).name and _MARKER in p.read_text(errors="ignore")
+        )
+        count = len(deferred)
+        if count > _BASELINE:
+            self.fail(
+                f"deferred (module-skipped) migration tests grew: {count} > "
+                f"baseline {_BASELINE}: {deferred}. Do not add new deferrals — "
+                "seed the context (override_server_args / publish a real "
+                "ServerArgs) instead of skipping the module."
+            )
+        if count < _BASELINE:
+            self.fail(
+                f"deferred migration tests shrank: {count} < baseline "
+                f"{_BASELINE}. Lower the baseline in this file to lock in the "
+                "recovery."
+            )
+
+
+if __name__ == "__main__":
+    unittest.main()

test/registered/unit/test_runtime_context_override.py

@@ -109,6 +109,61 @@ class TestContextOverride(CustomTestCase):
         with self.assertRaises(AttributeError):
             sa.page_size = 999
 
+    def test_preserve_config_keeps_post_publish_overrides(self):
+        # A nested build (e.g. a draft worker) publishes its own private copy;
+        # on exit the target's resolved bags — including post-publish
+        # overrides — must be reinstated verbatim, not re-projected from the
+        # pristine record (which would silently drop the overrides).
+        target = self._publish()
+        rc.get_context().override(
+            "ModelRunner.configure_kv_cache_dtype", kv_cache_dtype="fp8_e4m3"
+        )
+        draft = ServerArgs(model_path="dummy")
+        draft.override(source="draft-build", kv_cache_dtype="bf16")
+        with rc.get_context().preserve_config():
+            rc.get_context().set_server_args(draft)
+            # Inside the scope the draft's bags are live...
+            self.assertEqual(rc.get_model().kv_cache_dtype, "bf16")
+            # ...and its own post-publish overrides work as usual.
+            rc.get_context().override("draft-load", kv_cache_dtype="fp8_e5m2")
+            self.assertEqual(rc.get_model().kv_cache_dtype, "fp8_e5m2")
+        # Target slot, bags, and provenance restored verbatim.
+        self.assertIs(rc.get_context().server_args, target)
+        self.assertEqual(rc.get_model().kv_cache_dtype, "fp8_e4m3")
+        self.assertEqual(
+            rc.get_context().overrides_log(),
+            [
+                (
+                    "ModelRunner.configure_kv_cache_dtype",
+                    {"kv_cache_dtype": "fp8_e4m3"},
+                )
+            ],
+        )
+
+    def test_preserve_config_restores_in_scope_override_without_republish(self):
+        # An override inside the scope (no republish) mutates the live bags
+        # and provenance log in place; the scope must restore entry VALUES,
+        # not just reassign the aliased objects.
+        self._publish()
+        rc.get_context().override("srcA", page_size=16)
+        with rc.get_context().preserve_config():
+            rc.get_context().override("in-scope", page_size=64)
+            self.assertEqual(rc.get_schedule().page_size, 64)
+        self.assertEqual(rc.get_schedule().page_size, 16)
+        self.assertEqual(
+            rc.get_context().overrides_log(), [("srcA", {"page_size": 16})]
+        )
+
+    def test_preserve_config_restores_on_exception(self):
+        target = self._publish()
+        rc.get_context().override("srcA", page_size=16)
+        with self.assertRaises(RuntimeError):
+            with rc.get_context().preserve_config():
+                rc.get_context().set_server_args(ServerArgs(model_path="dummy"))
+                raise RuntimeError("nested build failed")
+        self.assertIs(rc.get_context().server_args, target)
+        self.assertEqual(rc.get_schedule().page_size, 16)
+
 
 if __name__ == "__main__":
     unittest.main()

test/registered/unit/test_server_args_writer_ratchet.py

@@ -0,0 +1,81 @@
+"""Ratchet guard: ``ServerArgs.override`` call-sites may only decrease.
+
+``ServerArgs.override(source, **fields)`` mutates a ``ServerArgs`` *instance*
+only — the resolved-config bags on the runtime context never see the write, so
+any consumer reading the namespace accessors (``get_exec()`` / ``get_memory()``
+/ …) desyncs from the writer. The migration end-state removes this primitive
+entirely: post-publish, process-global config changes go through
+``get_context().override(source, **fields)`` (which writes the bags), and
+per-runner resolved values live on the runner object rather than on a
+``ServerArgs`` copy.
+
+Until every call-site is rerouted together with its readers, this exact pin
+keeps the writer surface from growing unwatched: new writers must use
+``get_context().override``, and each rerouted batch lowers the baseline to
+lock in the progress. (The count is textual and includes docstring mentions
+and the test-kit's private-config use — the pin tracks growth, not the exact
+production-writer census.)
+"""
+
+from sglang.test.ci.ci_register import register_cpu_ci
+
+register_cpu_ci(est_time=5, suite="base-a-test-cpu")
+
+import re
+import unittest
+from pathlib import Path
+
+import sglang
+from sglang.test.test_utils import CustomTestCase
+
+_SGLANG_ROOT = Path(next(iter(sglang.__path__)))
+
+# ``server_args.override(`` also matches ``self.server_args.override(``,
+# ``<obj>.server_args.override(``, and the ``draft_server_args`` /
+# ``dp_server_args`` copies; ``args`` / ``sa`` are the aliases a few call-sites
+# bind first.
+_WRITER_PATTERNS = [
+    re.compile(r"server_args\.override\("),
+    re.compile(r"\bargs\.override\("),
+    re.compile(r"\bsa\.override\("),
+]
+
+# The resolution pipeline itself (its declare face forwards through
+# ``override`` by design) and multimodal_gen, whose ServerArgs is a different
+# class outside this contract.
+_EXCLUDED = (
+    "srt/server_args.py",
+    "srt/arg_groups",
+    "multimodal_gen",
+)
+
+_BASELINE = 49
+
+
+class TestServerArgsWriterRatchet(CustomTestCase):
+    def test_server_args_override_call_sites_match_the_baseline(self):
+        count = 0
+        for path in sorted(_SGLANG_ROOT.rglob("*.py")):
+            rel = path.relative_to(_SGLANG_ROOT).as_posix()
+            if rel.startswith(_EXCLUDED):
+                continue
+            source = path.read_text()
+            count += sum(len(p.findall(source)) for p in _WRITER_PATTERNS)
+        if count > _BASELINE:
+            self.fail(
+                f"ServerArgs.override call-sites grew: {count} > baseline "
+                f"{_BASELINE}. Instance writes never reach the resolved-config "
+                "bags, so namespace readers desync from the writer. Post-publish "
+                "process-global changes go through get_context().override(...); "
+                "per-runner resolved values belong on the runner object."
+            )
+        if count < _BASELINE:
+            self.fail(
+                f"ServerArgs.override call-sites shrank: {count} < baseline "
+                f"{_BASELINE}. Lower the baseline in this file to lock in the "
+                "progress."
+            )
+
+
+if __name__ == "__main__":
+    unittest.main()