From 1d0a6ee1781351c3819f64101dbe6ce2a62ed5ac Mon Sep 17 00:00:00 2001 From: Cheng Wan <54331508+ch-wan@users.noreply.github.com> Date: Wed, 22 Jul 2026 01:17:42 -0700 Subject: [PATCH] runtime_context: record the publishing process role (#31813) --- .../srt/disaggregation/encode_server.py | 5 +- .../srt/elastic_ep/expert_backup_manager.py | 5 +- .../srt/managers/data_parallel_controller.py | 6 + python/sglang/srt/managers/scheduler.py | 7 + python/sglang/srt/runtime_context.py | 146 +++++++++++++++--- python/sglang/srt/server_args.py | 15 +- .../srt/speculative/draft_worker_common.py | 17 +- .../unit/test_runtime_context_override.py | 21 +++ 8 files changed, 183 insertions(+), 39 deletions(-) diff --git a/python/sglang/srt/disaggregation/encode_server.py b/python/sglang/srt/disaggregation/encode_server.py index 626b75a1b..875a80f2d 100644 --- a/python/sglang/srt/disaggregation/encode_server.py +++ b/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. diff --git a/python/sglang/srt/elastic_ep/expert_backup_manager.py b/python/sglang/srt/elastic_ep/expert_backup_manager.py index 3ca4a7249..15a6536e4 100644 --- a/python/sglang/srt/elastic_ep/expert_backup_manager.py +++ b/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, ) diff --git a/python/sglang/srt/managers/data_parallel_controller.py b/python/sglang/srt/managers/data_parallel_controller.py index c43a4fc1b..35efe50d8 100644 --- a/python/sglang/srt/managers/data_parallel_controller.py +++ b/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( diff --git a/python/sglang/srt/managers/scheduler.py b/python/sglang/srt/managers/scheduler.py index 94340105e..0b96a4de5 100644 --- a/python/sglang/srt/managers/scheduler.py +++ b/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 diff --git a/python/sglang/srt/runtime_context.py b/python/sglang/srt/runtime_context.py index 5281ef432..e0daefeea 100644 --- a/python/sglang/srt/runtime_context.py +++ b/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() diff --git a/python/sglang/srt/server_args.py b/python/sglang/srt/server_args.py index b5ec57619..ebb675433 100644 --- a/python/sglang/srt/server_args.py +++ b/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: diff --git a/python/sglang/srt/speculative/draft_worker_common.py b/python/sglang/srt/speculative/draft_worker_common.py index 901aa88cf..5b4aaa0f8 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 @@ -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 diff --git a/test/registered/unit/test_runtime_context_override.py b/test/registered/unit/test_runtime_context_override.py index 2b9b5ba7f..12db59aa5 100644 --- a/test/registered/unit/test_runtime_context_override.py +++ b/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()