runtime_context: record the publishing process role (#31813)

python/sglang/srt/disaggregation/encode_server.py

@@ -66,7 +66,6 @@ from sglang.srt.observability.trace import (
 from sglang.srt.server_args import (
     PortArgs,
     ServerArgs,
-    set_global_server_args_for_scheduler,
 )
 from sglang.srt.utils import (
     add_prometheus_middleware,
@@ -262,7 +261,9 @@ class MMEncoder:
     ):
         logger.info(f"init MMEncoder {rank}/{server_args.tp_size}")
         self.server_args = server_args
-        set_global_server_args_for_scheduler(server_args)
+        from sglang.srt.runtime_context import publish
+
+        publish(server_args, role="encoder")
         self.rank = rank
         # DP rank for metric labels; overridden by run_dp_worker in DP mode.
         # 0 in the single-instance (non-DP) path.

python/sglang/srt/elastic_ep/expert_backup_manager.py

@@ -20,7 +20,6 @@ from sglang.srt.model_loader.utils import set_default_torch_dtype
 from sglang.srt.server_args import (
     PortArgs,
     ServerArgs,
-    set_global_server_args_for_scheduler,
 )
 from sglang.srt.utils.network import get_local_ip_auto
 
@@ -159,7 +158,9 @@ def run_expert_backup_manager_process(
     server_args: ServerArgs,
     port_args: PortArgs,
 ):
-    set_global_server_args_for_scheduler(server_args)
+    from sglang.srt.runtime_context import publish
+
+    publish(server_args, role="expert_backup")
     from sglang.srt.distributed.device_communicators.mooncake_transfer_engine import (
         init_mooncake_transfer_engine,
     )

python/sglang/srt/managers/data_parallel_controller.py

@@ -815,6 +815,12 @@ def run_data_parallel_controller_process(
     kill_itself_when_parent_died()
     parent_process = psutil.Process().parent()
 
+    # Publish the resolved config at DP-controller process entry: this process
+    # reads config namespaces (e.g. get_exec().moe.*) in its own address space
+    # before spawning schedulers.
+    from sglang.srt.runtime_context import publish
+
+    publish(server_args, role="scheduler")
     configure_logger(server_args)
     if server_args.enable_trace:
         process_tracing_init(

python/sglang/srt/managers/scheduler.py

@@ -4627,6 +4627,13 @@ def run_scheduler_process(
         display_dp_rank=display_dp_rank,
         display_moe_ep_rank=display_moe_ep_rank,
     )
+    # Publish the resolved config at scheduler process entry so the config
+    # namespaces (get_serving()/get_device()/get_exec()/...) are available to
+    # Scheduler.__init__ and its init_* helpers, which read them before the
+    # model worker's own publish. ModelRunner re-publishes idempotently.
+    from sglang.srt.runtime_context import publish
+
+    publish(server_args, role="scheduler")
     parent_process = psutil.Process().parent()
 
     # Set up tracing

python/sglang/srt/runtime_context.py

@@ -574,9 +574,17 @@ class _ConfigBag:
     writers are ``get_context().override(source, ...)`` (permanent) and
     the scoped ``.override(**kw)`` context manager (tests). Sub-namespaces
     (e.g. ``exec.moe``) are nested ``_ConfigBag`` instances reached by attribute.
-    """
 
-    __slots__ = ("_path", "_fields", "_subs")
+    Leaves and sub-bags are stored as **real instance attributes** (in
+    ``__dict__``), so ``bag.leaf`` / ``bag.sub`` is a plain attribute load that
+    ``torch.compile`` / dynamo can trace — config reads inside a compiled model
+    forward (e.g. ``get_exec().comm.enable_symm_mem`` in the embedding layer)
+    must not graph-break. ``_fields`` / ``_subs`` keep the authoritative
+    name→value maps used for override routing, membership, and scoped restore;
+    ``__getattr__`` is only a fallback for genuinely absent names. (Deliberately
+    no ``__slots__``: leaves are dynamic, and the ``__dict__`` is what makes the
+    reads traceable.)
+    """
 
     def __init__(self, path: str):
         object.__setattr__(self, "_path", path)
@@ -584,7 +592,9 @@ class _ConfigBag:
         object.__setattr__(self, "_subs", {})  # {subname: _ConfigBag}
 
     def __getattr__(self, name: str) -> Any:
-        # Reached only when ``name`` is not a real attribute (slot).
+        # Fallback only: real leaves/sub-bags resolve via __dict__ before this
+        # runs. Uses object.__getattribute__ (not self._fields) to stay safe if
+        # invoked before __init__ populates the bookkeeping dicts.
         fields = object.__getattribute__(self, "_fields")
         if name in fields:
             return fields[name]
@@ -601,8 +611,16 @@ class _ConfigBag:
         )
 
     def _set(self, name: str, value: Any) -> None:
-        """Internal write (publish + override) that bypasses the read-only guard."""
+        """Internal write (publish + override) that bypasses the read-only guard.
+        Updates both the bookkeeping map and the real attribute (traceable read)."""
         object.__getattribute__(self, "_fields")[name] = value
+        object.__setattr__(self, name, value)
+
+    def _set_sub(self, name: str, sub: _ConfigBag) -> None:
+        """Register a nested bag as both a bookkeeping entry and a real
+        attribute (so ``bag.sub`` is a plain, traceable attribute load)."""
+        object.__getattribute__(self, "_subs")[name] = sub
+        object.__setattr__(self, name, sub)
 
     def __contains__(self, name: str) -> bool:
         return name in object.__getattribute__(self, "_fields")
@@ -617,11 +635,13 @@ class _ConfigBag:
             path = object.__getattribute__(self, "_path")
             raise ValueError(f"unknown config leaf for {path!r}: {sorted(unknown)}")
         saved = {name: fields[name] for name in kwargs}
-        fields.update(kwargs)
+        for name, value in kwargs.items():
+            self._set(name, value)
         try:
             yield self
         finally:
-            fields.update(saved)
+            for name, value in saved.items():
+                self._set(name, value)
 
 
 def _build_config_bags(server_args: Any) -> dict:
@@ -660,7 +680,8 @@ def _build_config_bags(server_args: Any) -> dict:
             subs = object.__getattribute__(bag, "_subs")
             child = subs.get(name)
             if child is None:
-                child = subs[name] = _ConfigBag(".".join(parts[: depth + 1]))
+                child = _ConfigBag(".".join(parts[: depth + 1]))
+                bag._set_sub(name, child)
             bag = child
         if field in object.__getattribute__(bag, "_subs"):
             raise ValueError(
@@ -681,6 +702,7 @@ class RuntimeContext:
         "_server_args",
         "_config_bags",
         "_overrides_log",
+        "_publish_role",
         "flags",
         "resources",
         "forward",
@@ -691,6 +713,7 @@ class RuntimeContext:
         self._server_args: ServerArgs | None = None
         self._config_bags: dict | None = None
         self._overrides_log: list = []
+        self._publish_role: str | None = None
         self.flags = Flags()
         self.resources = Resources()
         self.forward = ForwardFlags()
@@ -815,8 +838,11 @@ class RuntimeContext:
         self._overrides_log.append((source, dict(fields)))
 
     def overrides_log(self) -> list:
-        """Provenance of post-publish ``override`` calls: ``[(source, {field: value})]``."""
-        return list(self._overrides_log)
+        """Provenance of post-publish ``override`` calls: ``[(source, {field: value})]``.
+
+        Returns deep-ish copies (source, dict(fields)) so callers inspecting the
+        log cannot mutate the recorded provenance in place."""
+        return [(source, dict(fields)) for source, fields in self._overrides_log]
 
     def resolved_server_args_dict(self, base: dict | None = None) -> dict:
         """Serialize the *resolved* config: the pristine ``server_args`` fields
@@ -858,6 +884,35 @@ class RuntimeContext:
         """
         return _ServerArgsOverride(self, fields)
 
+    @contextmanager
+    def preserve_config(self):
+        """Snapshot the full config lifecycle and reinstate it verbatim on exit.
+
+        Used when a nested construction step must leave the process-wide config
+        exactly as it found it — notably ``build_draft_tp_worker``, which builds
+        a draft worker off a private ``ServerArgs`` copy and must not disturb the
+        target's published config. Unlike ``set_server_args`` (which re-projects
+        the bags from a pristine record and so *discards* every post-publish
+        override made during target loading, e.g. ``kv_cache_dtype`` or
+        ``disable_shared_experts_fusion``), this restores the resolved bags
+        as-is, so namespace readers keep the target's resolved values afterward.
+        """
+        prev_server_args = self._server_args
+        prev_bags = self._config_bags
+        prev_overrides_log = self._overrides_log
+        prev_publish_role = self._publish_role
+        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
+            self._overrides_log = prev_overrides_log
+            self._publish_role = prev_publish_role
+            self.parallel._config = prev_parallel_config
+            self.flags.capture.enable_torch_compile = prev_capture
+
 
 class _ServerArgsOverride:
     """Scoped config override (see ``RuntimeContext.override_server_args``).
@@ -869,13 +924,21 @@ class _ServerArgsOverride:
     nondeterministic point.
     """
 
-    __slots__ = ("_context", "_fields", "_previous", "_previous_capture", "_installed")
+    __slots__ = (
+        "_context",
+        "_fields",
+        "_prev_server_args",
+        "_prev_bags",
+        "_prev_overrides_log",
+        "_prev_publish_role",
+        "_prev_parallel_config",
+        "_prev_capture",
+        "_installed",
+    )
 
     def __init__(self, context: RuntimeContext, fields: dict):
         self._context = context
         self._fields = fields
-        self._previous: ServerArgs | None = None
-        self._previous_capture = False
         self._installed = False
 
     def install(self) -> ServerArgs:
@@ -885,8 +948,18 @@ class _ServerArgsOverride:
         from sglang.srt.server_args import ServerArgs
 
         assert not self._installed, "override_server_args already installed"
-        self._previous = self._context._server_args
-        self._previous_capture = self._context.flags.capture.enable_torch_compile
+        # Snapshot the ENTIRE pre-install lifecycle state so restore() reinstates
+        # it verbatim: reseeding only ``_server_args`` would leave the projected
+        # bags / parallel leaves / provenance from this override live after the
+        # scope (violating fail-closed and leaking config into later tests), and
+        # would also drop any outer override that was active before this one.
+        ctx = self._context
+        self._prev_server_args = ctx._server_args
+        self._prev_bags = ctx._config_bags
+        self._prev_overrides_log = ctx._overrides_log
+        self._prev_publish_role = ctx._publish_role
+        self._prev_parallel_config = ctx.parallel._config
+        self._prev_capture = ctx.flags.capture.enable_torch_compile
         server_args = ServerArgs(model_path="dummy")
         if self._fields:
             server_args.override(source="test-override", **self._fields)
@@ -895,24 +968,26 @@ class _ServerArgsOverride:
         # materialized so bare post-publish writes raise like they do on a
         # fully resolved config.
         object.__setattr__(server_args, "_declarations_materialized", True)
-        self._context.set_server_args(server_args)
+        ctx.set_server_args(server_args)
         self._installed = True
         return server_args
 
     def restore(self) -> None:
-        """Reinstate the previously published config (or the empty slot)."""
+        """Reinstate the exact pre-install lifecycle state (or the empty slot)."""
         if not self._installed:
             return
         self._installed = False
-        previous, self._previous = self._previous, None
-        if previous is None:
-            self._context._server_args = None
-        else:
-            self._context.set_server_args(previous)
-        # set_server_args reseeds the capture tier from the published object
-        # (and the empty-slot path does not touch it at all); the snapshot
-        # puts back the exact pre-install runtime state either way.
-        self._context.flags.capture.enable_torch_compile = self._previous_capture
+        ctx = self._context
+        ctx._server_args = self._prev_server_args
+        ctx._config_bags = self._prev_bags
+        ctx._overrides_log = self._prev_overrides_log
+        ctx._publish_role = self._prev_publish_role
+        ctx.parallel._config = self._prev_parallel_config
+        ctx.flags.capture.enable_torch_compile = self._prev_capture
+        self._prev_server_args = None
+        self._prev_bags = None
+        self._prev_overrides_log = None
+        self._prev_parallel_config = None
 
     def __enter__(self) -> ServerArgs:
         return self.install()
@@ -998,6 +1073,26 @@ def get_observability() -> _ConfigBag:
     return _CONTEXT.config_bag("observability")
 
 
+def publish(server_args, *, role: str, hf_config: Any = None) -> RuntimeContext:
+    """Install process-wide config for this OS process.
+
+    Records the process ``role`` (``tokenizer`` / ``scheduler`` / ``encoder`` /
+    ``expert_backup`` / ``launcher`` / ``test``) and projects the config bags.
+    One call per process; draft workers skip publish (they must not clobber the
+    target). ``role`` is provenance today — per-role namespace projection and
+    fail-closed enforcement is a later unit. ``hf_config`` is accepted for
+    forward-compat and currently unused.
+    """
+    _CONTEXT._publish_role = role
+    _CONTEXT.set_server_args(server_args)
+    return _CONTEXT
+
+
+def publish_role() -> str | None:
+    """The role recorded by the last ``publish`` (None for a legacy set)."""
+    return _CONTEXT._publish_role
+
+
 def get_stream(name: str) -> Any:
     return _CONTEXT.get_stream(name)
 
@@ -1031,6 +1126,7 @@ def reset_context() -> None:
     _CONTEXT._server_args = None
     _CONTEXT._config_bags = None
     _CONTEXT._overrides_log = []
+    _CONTEXT._publish_role = None
     _CONTEXT.parallel._config = None
     _CONTEXT.flags = Flags()
     _CONTEXT.resources = Resources()

python/sglang/srt/server_args.py

@@ -8569,14 +8569,19 @@ class ServerArgs:
 # (decrease-only) by test/registered/unit/test_legacy_global_ratchet.py.
 # Imports are in-function so the two modules stay cycle-free at import time.
 def set_global_server_args_for_scheduler(server_args: ServerArgs):
-    """Legacy publish shim — prefer ``get_context().set_server_args()`` from
-    ``sglang.srt.runtime_context`` in new code."""
-    from sglang.srt.runtime_context import get_context
+    """Legacy publish shim (role=scheduler) — prefer
+    ``runtime_context.publish(server_args, role=...)`` in new code."""
+    from sglang.srt.runtime_context import publish
 
-    get_context().set_server_args(server_args)
+    publish(server_args, role="scheduler")
 
 
-set_global_server_args_for_tokenizer = set_global_server_args_for_scheduler
+def set_global_server_args_for_tokenizer(server_args: ServerArgs):
+    """Legacy publish shim (role=tokenizer). Not aliased to the scheduler shim:
+    the process role differs."""
+    from sglang.srt.runtime_context import publish
+
+    publish(server_args, role="tokenizer")
 
 
 def get_global_server_args() -> ServerArgs:

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
@@ -90,8 +90,17 @@ def build_draft_tp_worker(
         context_length=target_model_config.context_len,
     )
 
-    saved_server_args = get_server_args()
-    try:
+    # Publish the draft copy for the duration of the build so the draft's model
+    # layers resolve config (e.g. kv_cache_dtype) from the draft's own bags, not
+    # the target's -- an independently configured draft can resolve a different
+    # KV-cache dtype than the target, and reading the target-global bag would
+    # make draft attention record the wrong dtype. ``preserve_config`` snapshots
+    # the target's resolved config on entry and reinstates it verbatim on exit
+    # (post-publish overrides intact), so the target is undisturbed afterwards --
+    # unlike a plain ``set_server_args(saved)`` restore, which re-projects the
+    # bags from the pristine record and drops those overrides.
+    with get_context().preserve_config():
+        get_context().set_server_args(draft_server_args)
         draft_worker = TpModelWorker(
             server_args=draft_server_args,
             gpu_id=gpu_id,
@@ -99,8 +108,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_runtime_context_override.py

@@ -109,6 +109,27 @@ class TestContextOverride(CustomTestCase):
         with self.assertRaises(AttributeError):
             sa.page_size = 999
 
+    def test_publish_records_role(self):
+        rc.publish(ServerArgs(model_path="dummy"), role="scheduler")
+        self.assertEqual(rc.publish_role(), "scheduler")
+
+    def test_legacy_shims_record_roles(self):
+        # Unit 2a: the legacy setters publish with their process role.
+        from sglang.srt.server_args import (
+            set_global_server_args_for_scheduler,
+            set_global_server_args_for_tokenizer,
+        )
+
+        set_global_server_args_for_scheduler(ServerArgs(model_path="dummy"))
+        self.assertEqual(rc.publish_role(), "scheduler")
+        set_global_server_args_for_tokenizer(ServerArgs(model_path="dummy"))
+        self.assertEqual(rc.publish_role(), "tokenizer")
+
+    def test_reset_clears_role(self):
+        rc.publish(ServerArgs(model_path="dummy"), role="test")
+        rc.reset_context()
+        self.assertIsNone(rc.publish_role())
+
 
 if __name__ == "__main__":
     unittest.main()