[Config] Round 5.1: the published-side readers ask the bags, and a platform fact gets one address (#37086)

This commit is contained in:
Cheng Wan
2026-08-30 02:18:33 -07:00
committed by GitHub
parent a6e4021368
commit 7e751153eb
148 changed files with 1585 additions and 1114 deletions
+293 -15
View File
@@ -54,7 +54,7 @@ import math
import os
import sys
from contextlib import contextmanager
from typing import TYPE_CHECKING, Any
from typing import TYPE_CHECKING, Any, Dict, Optional
if TYPE_CHECKING:
from sglang.srt.server_args import ServerArgs
@@ -833,6 +833,25 @@ def _build_config_bags(server_args: Any) -> dict:
return tops
def _resolved_or_field(server_args: Any, name: str, default: Any) -> Any:
"""What resolution decided for `name`, falling back to the field.
Publishes that carry no config at all (sentinels, mocks) have neither, and
answer with `default`.
"""
if server_args is None:
return default
from sglang.srt.arg_groups.overrides import resolution_result
decided = resolution_result(server_args, name)
if decided is not None:
return decided
# The default is for the callers that hand over something record-shaped but
# not a record -- the fake configs the context tests publish, and `object()`
# for the sentinel publish. A real ServerArgs always has the field.
return getattr(server_args, name, default)
class RuntimeContext:
"""Container for the structured runtime accessors; exposes ``parallel``,
``server_args``, the resolved config namespace bags, ``flags``,
@@ -914,9 +933,10 @@ class RuntimeContext:
stash, which is what the bags are projected from.
"""
# Seed the capture tier for the new lifecycle (defaults for sentinel
# and mock publishes, which carry no config).
self.flags.capture.enable_torch_compile = getattr(
server_args, "enable_torch_compile", False
# and mock publishes, which carry no config). Through the resolution,
# not the field: the field is the operator's input.
self.flags.capture.enable_torch_compile = bool(
_resolved_or_field(server_args, "enable_torch_compile", False)
)
self._server_args = server_args
# The adaptive draft-token bound memoizes on the config *path*, so a new
@@ -1145,7 +1165,6 @@ class _ServerArgsOverride:
self._prev_parallel_config = ctx.parallel._config
self._prev_capture = ctx.flags.capture.enable_torch_compile
from sglang.srt.arg_groups.overrides import (
_apply_fields,
declare_late_resolution,
)
@@ -1162,18 +1181,20 @@ class _ServerArgsOverride:
)
# Declared so the projection sees it; late, because the record is
# resolved already and not yet published.
# Underscore names are not fields at all (they seed private property
# caches), so they stay a direct write.
declared = {
name: value for name, value in self._fields.items() if name[0] != "_"
}
# Split on whether the name is a field, not on whether it starts with
# an underscore: `_speculative_draft_quantization_explicitly_set` is a
# real field, and seeding it as a raw attribute would leave the earlier
# declaration authoritative, so `resolution_result` and the bag would
# both keep answering the pre-override value.
fields = set(type(server_args).__dataclass_fields__)
declared = {n: v for n, v in self._fields.items() if n in fields}
if declared:
declare_late_resolution(server_args, "override_server_args", **declared)
# This hook stands in for a launch: the caller's values are both what
# the operator passed and what resolution decided, so they go on the
# record as well as into the stash. Production late resolution declares
# only -- there the record stays the operator's input.
_apply_fields(server_args, self._fields)
# What is left seeds the record's own private caches (`_model_config`
# and friends), which are not configuration and never were.
seeds = {n: v for n, v in self._fields.items() if n not in fields}
for name, value in seeds.items():
object.__setattr__(server_args, name, value)
ctx.set_server_args(server_args)
self._installed = True
return server_args
@@ -1723,6 +1744,142 @@ def pre_capture_activation_reserve_mb(gpu_mem: float | None) -> float:
return reserved_mem
# --- Platform facts -----------------------------------------------------------
#
# One address for what kind of machine this is, so a reader asks
# `get_platform().is_sm100` and an override is stated once instead of patched
# into every module that imported a probe. True before publish, so the context
# probes when no override is installed; `utils.common` holds the implementation.
_PLATFORM_PROBES: Dict[str, str] = {
"is_cuda": "is_cuda",
"is_hip": "is_hip",
"is_npu": "is_npu",
"is_xpu": "is_xpu",
"is_musa": "is_musa",
"is_sm90": "is_sm90_supported",
"is_sm100": "is_sm100_supported",
"is_sm100_or_sm110": "is_sm100_or_sm110_supported",
"is_sm120": "is_sm120_supported",
"is_blackwell": "is_blackwell_supported",
"is_hopper_with_cuda_12_3": "is_hopper_with_cuda_12_3",
"has_amx": "cpu_has_amx_support",
"has_flashinfer": "is_flashinfer_available",
}
# Not yes/no facts, same address.
_PLATFORM_VALUES: Dict[str, str] = {
"device_sm": "get_device_sm",
"device_capability": "get_device_capability",
}
class PlatformContext:
"""The machine's own facts, with one place to override them.
Every name maps to a probe in `utils.common`; the probes are
`lru_cache`-d, so reading through here costs a call and a dict lookup
(~26 ns) rather than a device query.
"""
__slots__ = ("_overrides",)
def __init__(self) -> None:
object.__setattr__(self, "_overrides", {})
def __getattr__(self, name: str) -> Any:
probe = _PLATFORM_PROBES.get(name) or _PLATFORM_VALUES.get(name)
if probe is None:
known = sorted(set(_PLATFORM_PROBES) | set(_PLATFORM_VALUES))
raise AttributeError(
f"unknown platform fact {name!r}; known: {', '.join(known)}"
)
overrides = object.__getattribute__(self, "_overrides")
if name in overrides:
return overrides[name]
from sglang.srt.utils import common as _common
return getattr(_common, probe)()
def __setattr__(self, name: str, value: Any) -> None:
raise AttributeError(
"platform facts are not assigned; use "
"`sglang.srt.runtime_context.override_platform(...)` so every "
"reader agrees"
)
def _install(self, **facts: Any) -> Dict[str, Any]:
unknown = set(facts) - set(_PLATFORM_PROBES) - set(_PLATFORM_VALUES)
if unknown:
raise ValueError(f"unknown platform fact(s): {sorted(unknown)}")
overrides = object.__getattribute__(self, "_overrides")
previous = {k: overrides[k] for k in facts if k in overrides}
missing = [k for k in facts if k not in overrides]
overrides.update(facts)
return {"previous": previous, "missing": missing}
def _restore(self, saved: Dict[str, Any]) -> None:
overrides = object.__getattribute__(self, "_overrides")
overrides.update(saved["previous"])
for k in saved["missing"]:
overrides.pop(k, None)
_PLATFORM = PlatformContext()
def get_platform() -> PlatformContext:
"""The machine's facts. Answers before publish, unlike a config bag."""
return _PLATFORM
class _PlatformOverride:
"""Scoped platform override: `with override_platform(is_sm100=True): ...`"""
__slots__ = ("_facts", "_saved")
def __init__(self, **facts: Any) -> None:
self._facts = facts
self._saved = None
def install(self) -> PlatformContext:
self._saved = _PLATFORM._install(**self._facts)
return _PLATFORM
def restore(self) -> None:
if self._saved is not None:
_PLATFORM._restore(self._saved)
self._saved = None
def __enter__(self) -> PlatformContext:
return self.install()
def __exit__(self, *exc: Any) -> None:
self.restore()
def __call__(self, fn: Any) -> Any:
"""Also usable as a decorator, like the `patch` it replaces.
A fresh scope per call: the same object decorating two tests must not
share one saved state.
"""
import functools
facts = dict(self._facts)
@functools.wraps(fn)
def wrapper(*args: Any, **kwargs: Any) -> Any:
with _PlatformOverride(**facts):
return fn(*args, **kwargs)
return wrapper
def override_platform(**facts: Any) -> _PlatformOverride:
"""Say what kind of machine this is, once, for every reader."""
return _PlatformOverride(**facts)
# --- Derived config accessors ------------------------------------------------
#
# A few values are computed from several config fields plus the HF config, so
@@ -1873,3 +2030,124 @@ def is_ep_joiner() -> bool:
def is_ep_scale_joiner() -> bool:
"""True in a process launched as an elastic-EP scale-up joiner."""
return get_exec().moe.ep_join_mode == "scale"
def describe_kv_events_publisher(server_args: Any) -> Optional[dict]:
"""Return a structured description of this server's KV-event
publisher, or `None` if publishing is disabled / misconfigured.
This is the wire contract surfaced under the `kv_events` key on
`/server_info` so KV-aware routers (e.g. the SGLang model
gateway) can subscribe per-worker without operator-supplied port
coordination. The router constructs the per-DP-rank SUB endpoint
as tcp://<worker_host>:<endpoint_port_base + dp_rank> for
every rank reported in dp_size.
Returned descriptor shape:
{
"publisher": "zmq",
"endpoint_host": "*", # may be a ZMQ wildcard
# ("*", "0.0.0.0", "::");
# subscribers MUST substitute
# the worker URL's host when
# dialing
"endpoint_port_base": 5557, # base TCP port; per-rank
# port = base + dp_rank
"topic": "", # ZMQ topic prefix on the
# SUB filter (empty =
# subscribe-all)
"block_size": <kv_event_block_size>, # subscribers MUST
# hash prompts at this size
"dp_size": <dp_size>, # number of SUB sockets to
# open; not DCP-scaled, as
# DCP shards within a rank
# rather than adding
# publishers
"load_endpoint_port_base": <resolved>,
# base TCP port of the load
# range (load rank r = base
# + r). Consumers MUST read
# this key, not re-derive
# it; present only when
# --load-publish-endpoint
# opted in and a range
# resolved
"load_topic": "load", # SUB filter for the load
# socket; present iff
# load_endpoint_port_base
# is present
}
Returns None (i.e. "no publisher to describe") when any of:
* --kv-events-config is unset / empty / malformed JSON,
* the configured publisher is "null",
* page_size is missing or non-positive (a placeholder
block_size would cause silent KV-cache misses by hashing
prompts at the wrong granularity on the router side),
* the endpoint is not a routable TCP address (inproc:// /
ipc://, missing port, non-integer port, port outside
1..65535, or a bare unbracketed IPv6 host, which is
ambiguous).
NOTE for load-socket consumers: pair the load port with the worker's
own URL host, as with the KV SUB endpoints — endpoint_host is a
wildcard ("*", "0.0.0.0", "::") whenever the default packing applies,
so splicing it yields tcp://*:PORT and connects to nothing.
Reuses parse_advertisable_tcp and resolve_load_pub_range — the same
helpers the scheduler binds through — so the advertisement cannot
drift from the sockets.
"""
from sglang.srt.arg_groups.overrides import resolving_view
# Lazy import so loading server_args doesn't pull in
# disaggregation / msgspec / zmq at module top level.
from sglang.srt.disaggregation.kv_events import (
LOAD_TOPIC,
KVEventsConfig,
parse_advertisable_tcp,
resolve_load_pub_range,
)
resolved = resolving_view(server_args)
raw = resolved.kv_events_config
page_size = resolved.page_size
if not raw or page_size is None or page_size <= 0:
return None
try:
cfg = KVEventsConfig.from_cli(raw)
except Exception:
# Malformed JSON / schema mismatch. The publisher would
# have failed at server startup; /server_info must
# keep working, so just report "no publisher" to consumers.
return None
if cfg.publisher == "null" or not cfg.endpoint:
return None
resolved_kv = parse_advertisable_tcp(cfg.endpoint)
if resolved_kv is None:
return None
host, port = resolved_kv
descriptor = {
"publisher": cfg.publisher,
"endpoint_host": host,
"endpoint_port_base": port,
"topic": cfg.topic,
"block_size": resolved.kv_event_block_size,
"dp_size": resolved.dp_size,
}
# Load range, from the same resolver SchedulerLoadPublisher binds
# with (so the two can't drift). The decline reason is logged once at
# startup, not here — this runs per /server_info request.
resolved_range, _reason = resolve_load_pub_range(
kv_endpoint=cfg.endpoint,
replay_endpoint=cfg.replay_endpoint,
dp_size=resolved.dp_size,
load_publish_endpoint=resolved.load_publish_endpoint,
)
if resolved_range is not None:
descriptor["load_endpoint_port_base"] = resolved_range[1]
descriptor["load_topic"] = LOAD_TOPIC
return descriptor