From ab2553284aec27a09a1d8a878869e5573859bc22 Mon Sep 17 00:00:00 2001 From: Cheng Wan <54331508+ch-wan@users.noreply.github.com> Date: Fri, 31 Jul 2026 15:05:41 -0700 Subject: [PATCH] config: preserve resolved config across nested publishes + mutation ratchets (#33011) --- python/sglang/srt/runtime_context.py | 53 ++++++++++++ .../srt/speculative/draft_worker_common.py | 7 +- .../unit/test_migration_deferral_ratchet.py | 53 ++++++++++++ .../unit/test_runtime_context_override.py | 55 +++++++++++++ .../unit/test_server_args_writer_ratchet.py | 81 +++++++++++++++++++ 5 files changed, 244 insertions(+), 5 deletions(-) create mode 100644 test/registered/unit/test_migration_deferral_ratchet.py create mode 100644 test/registered/unit/test_server_args_writer_ratchet.py diff --git a/python/sglang/srt/runtime_context.py b/python/sglang/srt/runtime_context.py index 5281ef432..d1d29615f 100644 --- a/python/sglang/srt/runtime_context.py +++ b/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``). diff --git a/python/sglang/srt/speculative/draft_worker_common.py b/python/sglang/srt/speculative/draft_worker_common.py index a42d74bf8..4c65d300f 100644 --- a/python/sglang/srt/speculative/draft_worker_common.py +++ b/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 diff --git a/test/registered/unit/test_migration_deferral_ratchet.py b/test/registered/unit/test_migration_deferral_ratchet.py new file mode 100644 index 000000000..eefce0e8c --- /dev/null +++ b/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/ -> 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() diff --git a/test/registered/unit/test_runtime_context_override.py b/test/registered/unit/test_runtime_context_override.py index 2b9b5ba7f..d946ffbc7 100644 --- a/test/registered/unit/test_runtime_context_override.py +++ b/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() diff --git a/test/registered/unit/test_server_args_writer_ratchet.py b/test/registered/unit/test_server_args_writer_ratchet.py new file mode 100644 index 000000000..3082b5654 --- /dev/null +++ b/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(``, +# ``.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()