runtime_context: record the publishing process role (#33012)

This commit is contained in:
Cheng Wan
2026-07-31 15:06:22 -07:00
committed by GitHub
parent ab2553284a
commit 4862edc85f
12 changed files with 166 additions and 45 deletions
@@ -63,10 +63,10 @@ from sglang.srt.observability.trace import (
process_tracing_init,
trace_set_thread_info,
)
from sglang.srt.runtime_context import publish
from sglang.srt.server_args import (
PortArgs,
ServerArgs,
set_global_server_args_for_scheduler,
)
from sglang.srt.utils import (
CLIENT_MEDIA_EXCEPTIONS,
@@ -265,7 +265,7 @@ 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)
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.
@@ -17,10 +17,10 @@ from sglang.srt.managers.io_struct import (
)
from sglang.srt.model_loader.loader import DefaultModelLoader, get_model_loader
from sglang.srt.model_loader.utils import set_default_torch_dtype
from sglang.srt.runtime_context import publish
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 +159,7 @@ def run_expert_backup_manager_process(
server_args: ServerArgs,
port_args: PortArgs,
):
set_global_server_args_for_scheduler(server_args)
publish(server_args, role="expert_backup")
from sglang.srt.distributed.device_communicators.mooncake_transfer_engine import (
init_mooncake_transfer_engine,
)
@@ -48,6 +48,7 @@ from sglang.srt.managers.scheduler import run_scheduler_process
from sglang.srt.observability.cpu_monitor import start_cpu_monitor_thread
from sglang.srt.observability.req_time_stats import DPControllerReqTimeStats
from sglang.srt.observability.trace import process_tracing_init, trace_set_thread_info
from sglang.srt.runtime_context import publish
from sglang.srt.server_args import (
DP_ATTENTION_HANDSHAKE_PORT_DELTA,
PortArgs,
@@ -815,6 +816,8 @@ def run_data_parallel_controller_process(
kill_itself_when_parent_died()
parent_process = psutil.Process().parent()
# This process reads the config namespaces before spawning schedulers.
publish(server_args, role="dp_controller")
configure_logger(server_args)
if server_args.enable_trace:
process_tracing_init(
+4 -1
View File
@@ -250,7 +250,7 @@ from sglang.srt.observability.trace import process_tracing_init, trace_set_threa
from sglang.srt.parser.reasoning_parser import ReasoningParser
from sglang.srt.platforms import current_platform
from sglang.srt.plugins import load_plugins
from sglang.srt.runtime_context import get_context, get_parallel
from sglang.srt.runtime_context import get_context, get_parallel, publish
from sglang.srt.sampling.sampling_batch_info import SamplingBatchInfo
from sglang.srt.sampling.sampling_params import TOP_K_ALL
from sglang.srt.server_args import PortArgs, ServerArgs
@@ -4765,6 +4765,9 @@ def run_scheduler_process(
display_dp_rank=display_dp_rank,
display_moe_ep_rank=display_moe_ep_rank,
)
# Scheduler.__init__ reads the config namespaces before the model
# worker's own publish.
publish(server_args, role="scheduler")
parent_process = psutil.Process().parent()
# Set up tracing
+7 -4
View File
@@ -16,13 +16,12 @@
from __future__ import annotations
import logging
from typing import TYPE_CHECKING, Any, Dict, Optional
from typing import Any, Dict, Optional
import ray
if TYPE_CHECKING:
from sglang.srt.server_args import PortArgs, ServerArgs
from sglang.srt.runtime_context import publish
from sglang.srt.server_args import PortArgs, ServerArgs
logger = logging.getLogger(__name__)
@@ -100,6 +99,10 @@ class SchedulerActor:
f"[TP{tp_rank}] Bound to NUMA node {numa_node} for GPU {actual_gpu_id}"
)
# This actor constructs Scheduler directly (no run_scheduler_process),
# which reads the config namespaces before the model worker's publish.
publish(server_args, role="scheduler")
# Create scheduler (loads model into GPU, initializes NCCL)
self.scheduler = Scheduler(
server_args=server_args,
+101 -26
View File
@@ -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(
@@ -709,6 +730,7 @@ class RuntimeContext:
"_server_args",
"_config_bags",
"_overrides_log",
"_publish_role",
"flags",
"resources",
"forward",
@@ -719,6 +741,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()
@@ -784,8 +807,9 @@ class RuntimeContext:
# leaves like pp_max_micro_batch_size are served via ParallelContext
# __getattr__; live topology properties still win by name).
self.parallel._config = self._config_bags.get("parallel")
# Fresh config lifecycle: prior override provenance no longer applies.
# A direct install is roleless; ``publish`` assigns the role afterwards.
self._overrides_log = []
self._publish_role = None
def config_bag(self, name: str) -> _ConfigBag:
"""Return the top-level config namespace bag (``device`` / ``model`` /
@@ -843,8 +867,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
@@ -898,6 +925,7 @@ class RuntimeContext:
prev_bags = self._config_bags
prev_bag_values = _snapshot_bag_values(prev_bags)
prev_overrides_log = list(self._overrides_log)
prev_publish_role = self._publish_role
prev_parallel_config = self.parallel._config
prev_capture = self.flags.capture.enable_torch_compile
try:
@@ -908,6 +936,7 @@ class RuntimeContext:
if prev_bags is not None:
_restore_bag_values(prev_bags, prev_bag_values)
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
@@ -922,13 +951,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:
@@ -938,8 +975,13 @@ 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
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)
@@ -948,24 +990,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()
@@ -1051,6 +1095,36 @@ 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`` /
``dp_controller`` / ``encoder`` / ``expert_backup`` /
``weight_cache_daemon`` / ``launcher`` / ``test``) and
projects the config bags. 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.
Normally one call per process, but re-publish is allowed and is
**last-publish-wins** (bags re-projected, provenance reset, role
overwritten). Two sanctioned multi-publish shapes exist: the in-process
Engine builds its ``TokenizerManager`` inside the launcher process (the
process ends up with the tokenizer publish), and multiple Engines in one
process publish in sequence — which is exactly why per-instance managers
must read ``self.server_args`` for anything engine-specific rather than
the process-global bags.
"""
_CONTEXT.set_server_args(server_args)
_CONTEXT._publish_role = role
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)
@@ -1084,6 +1158,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()
+10 -5
View File
@@ -9007,14 +9007,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:
@@ -97,7 +97,9 @@ def build_draft_tp_worker(
context_length=target_model_config.context_len,
)
# The draft's layers must resolve config from the draft's own bags.
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,
+2 -2
View File
@@ -46,6 +46,7 @@ import torch.distributed as dist
from sglang.srt.configs.load_config import LoadConfig
from sglang.srt.platforms import current_platform
from sglang.srt.runtime_context import publish
from sglang.srt.utils import MultiprocessingSerializer
from .protocol import (
@@ -208,7 +209,6 @@ class WeightCacheDaemon:
from sglang.srt.configs.device_config import DeviceConfig
from sglang.srt.configs.model_config import ModelConfig
from sglang.srt.model_loader.loader import get_model_loader
from sglang.srt.runtime_context import get_context
from sglang.srt.server_args import ServerArgs
server_args = ServerArgs(
@@ -223,7 +223,7 @@ class WeightCacheDaemon:
load_format=self.load_format,
model_loader_extra_config=self.model_loader_extra_config,
)
get_context().set_server_args(server_args)
publish(server_args, role="weight_cache_daemon")
# Initialize distributed backend for model loading
# (must be done after server_args and model_config are available)