feat(kv-events): expose structured KV-event publisher block on /server_info (#25844)
Co-authored-by: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
This commit is contained in:
co-authored by
Claude Opus 4.7
parent
06c23d55b5
commit
085777210c
@@ -256,10 +256,15 @@ class ZmqEventPublisher(EventPublisher):
|
||||
self._pub = self._ctx.socket(zmq.PUB)
|
||||
self._pub.set_hwm(self._hwm)
|
||||
# Heuristic: bind if wildcard / * present, else connect.
|
||||
# bind stable, connect volatile convention
|
||||
# bind stable, connect volatile convention.
|
||||
# ``0.0.0.0`` is the IPv4 bind-all wildcard alongside ``*``
|
||||
# and ``::``; ``/server_info`` advertises it as a wildcard,
|
||||
# so the publisher must bind it for the advertised endpoint
|
||||
# to actually be listening.
|
||||
if (
|
||||
"*" in self._endpoint
|
||||
or "::" in self._endpoint
|
||||
or "0.0.0.0" in self._endpoint
|
||||
or self._endpoint.startswith("ipc://")
|
||||
or self._endpoint.startswith("inproc://")
|
||||
):
|
||||
|
||||
@@ -636,12 +636,18 @@ async def server_info():
|
||||
await _global_state.tokenizer_manager.get_internal_state()
|
||||
)
|
||||
|
||||
server_args = _global_state.tokenizer_manager.server_args
|
||||
|
||||
# server_args.model_config is not serializable but should be excluded by asdict.
|
||||
return {
|
||||
**dataclasses.asdict(_global_state.tokenizer_manager.server_args),
|
||||
**dataclasses.asdict(server_args),
|
||||
**_global_state.scheduler_info,
|
||||
"internal_states": internal_states,
|
||||
"version": __version__,
|
||||
# Structured KV-event publisher descriptor for KV-aware routers.
|
||||
# `None` when publishing is disabled or misconfigured; see
|
||||
# `ServerArgs.describe_kv_events_publisher` for the precise contract.
|
||||
"kv_events": server_args.describe_kv_events_publisher(),
|
||||
}
|
||||
|
||||
|
||||
|
||||
@@ -7430,6 +7430,92 @@ class ServerArgs:
|
||||
else:
|
||||
return False
|
||||
|
||||
def describe_kv_events_publisher(self) -> 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": <page_size>, # subscribers MUST hash
|
||||
# prompts at this size
|
||||
"dp_size": <dp_size>, # number of SUB sockets
|
||||
# to open
|
||||
}
|
||||
|
||||
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, or port outside
|
||||
``1..65535``).
|
||||
|
||||
Reuses ``KVEventsConfig.from_cli`` for JSON parsing; the inline
|
||||
``rfind(":")`` endpoint split mirrors
|
||||
``ZmqEventPublisher.offset_endpoint_port`` rather than adding a
|
||||
new module-level helper.
|
||||
"""
|
||||
# Lazy import so loading ``server_args`` doesn't pull in
|
||||
# disaggregation / msgspec / zmq at module top level.
|
||||
from sglang.srt.disaggregation.kv_events import KVEventsConfig
|
||||
|
||||
raw = self.kv_events_config
|
||||
page_size = self.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
|
||||
if not cfg.endpoint.startswith("tcp://"):
|
||||
return None
|
||||
body = cfg.endpoint[len("tcp://") :]
|
||||
last_colon = body.rfind(":")
|
||||
if last_colon < 0:
|
||||
return None
|
||||
host = body[:last_colon]
|
||||
try:
|
||||
port = int(body[last_colon + 1 :])
|
||||
except ValueError:
|
||||
return None
|
||||
if not host or not (0 < port < 65536):
|
||||
return None
|
||||
return {
|
||||
"publisher": cfg.publisher,
|
||||
"endpoint_host": host,
|
||||
"endpoint_port_base": port,
|
||||
"topic": cfg.topic,
|
||||
"block_size": page_size,
|
||||
"dp_size": self.dp_size,
|
||||
}
|
||||
|
||||
|
||||
# NOTE: This is a global variable to hold the server args for scheduler.
|
||||
_global_server_args: Optional[ServerArgs] = None
|
||||
|
||||
@@ -0,0 +1,296 @@
|
||||
"""Endpoint-level tests for `/server_info`.
|
||||
|
||||
`/server_info` is the introspection surface that external consumers
|
||||
(SGLang's own deprecated `/get_server_info` alias, monitoring tools,
|
||||
KV-aware routers) scrape to learn about the running server's
|
||||
configuration. New `/server_info` behaviours should add their test
|
||||
classes to this file as the surface grows.
|
||||
|
||||
Current coverage:
|
||||
|
||||
* `TestServerInfoKvEventsField` — the `kv_events` publisher descriptor
|
||||
surfaced by `_build_kv_events_block`. Covers the full input matrix
|
||||
end-to-end (happy path / disabled / malformed JSON / inproc endpoint /
|
||||
port edge cases / missing-or-non-positive page_size) because the
|
||||
helper has no separate test target; the handler is its only caller.
|
||||
|
||||
* `TestServerInfoExistingFieldsPreserved` — regression guard that no
|
||||
field existing consumers depend on is silently dropped: every
|
||||
`ServerArgs` dataclass field, `internal_states`, `version`, and the
|
||||
pre-existing flat `kv_events_config` string all remain visible.
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
import dataclasses
|
||||
import unittest
|
||||
from types import SimpleNamespace
|
||||
|
||||
from sglang.srt.entrypoints import http_server
|
||||
from sglang.srt.server_args import ServerArgs
|
||||
from sglang.test.ci.ci_register import register_cpu_ci
|
||||
from sglang.test.test_utils import CustomTestCase
|
||||
|
||||
register_cpu_ci(est_time=5, suite="base-a-test-cpu")
|
||||
|
||||
|
||||
def _call_server_info_with(server_args: ServerArgs) -> dict:
|
||||
"""Invoke `http_server.server_info()` against a stub global state.
|
||||
|
||||
Bypasses the FastAPI HTTP layer (no TestClient): the handler is an
|
||||
`async def` that reads module-level `_global_state`, so wiring a
|
||||
`SimpleNamespace` stub via `set_global_state` and awaiting the
|
||||
coroutine directly is enough to exercise the handler logic without
|
||||
booting a model server.
|
||||
"""
|
||||
|
||||
async def _fake_internal_state():
|
||||
return [{"max_req_input_len": 1024}]
|
||||
|
||||
stub_state = SimpleNamespace(
|
||||
tokenizer_manager=SimpleNamespace(
|
||||
server_args=server_args,
|
||||
get_internal_state=_fake_internal_state,
|
||||
),
|
||||
scheduler_info={"max_req_input_len": 1024},
|
||||
)
|
||||
prior_state = http_server.get_global_state()
|
||||
http_server.set_global_state(stub_state)
|
||||
try:
|
||||
return asyncio.run(http_server.server_info())
|
||||
finally:
|
||||
# Restore so a later test in the same process isn't surprised.
|
||||
http_server._global_state = prior_state
|
||||
|
||||
|
||||
class TestServerInfoKvEventsField(CustomTestCase):
|
||||
"""The new `kv_events` field is wired correctly across the full
|
||||
`_build_kv_events_block` input matrix.
|
||||
"""
|
||||
|
||||
# ----- happy path --------------------------------------------------
|
||||
|
||||
def test_kv_events_key_present_when_publishing_enabled(self):
|
||||
args = ServerArgs(
|
||||
model_path="dummy",
|
||||
kv_events_config=(
|
||||
'{"publisher": "zmq", "endpoint": "tcp://*:5557", "topic": "kv"}'
|
||||
),
|
||||
page_size=64,
|
||||
dp_size=2,
|
||||
)
|
||||
|
||||
info = _call_server_info_with(args)
|
||||
|
||||
self.assertIn("kv_events", info)
|
||||
self.assertEqual(
|
||||
info["kv_events"],
|
||||
{
|
||||
"publisher": "zmq",
|
||||
"endpoint_host": "*",
|
||||
"endpoint_port_base": 5557,
|
||||
"topic": "kv",
|
||||
"block_size": 64,
|
||||
"dp_size": 2,
|
||||
},
|
||||
)
|
||||
|
||||
def test_kv_events_descriptor_carries_specific_host_and_topic(self):
|
||||
args = ServerArgs(
|
||||
model_path="dummy",
|
||||
kv_events_config=(
|
||||
'{"publisher": "zmq", "endpoint": "tcp://0.0.0.0:7777", "topic": "kv"}'
|
||||
),
|
||||
page_size=128,
|
||||
dp_size=1,
|
||||
)
|
||||
|
||||
info = _call_server_info_with(args)
|
||||
|
||||
self.assertIsNotNone(info["kv_events"])
|
||||
self.assertEqual(info["kv_events"]["endpoint_host"], "0.0.0.0")
|
||||
self.assertEqual(info["kv_events"]["endpoint_port_base"], 7777)
|
||||
self.assertEqual(info["kv_events"]["topic"], "kv")
|
||||
self.assertEqual(info["kv_events"]["block_size"], 128)
|
||||
self.assertEqual(info["kv_events"]["dp_size"], 1)
|
||||
|
||||
# ----- disabled / unconfigured -------------------------------------
|
||||
|
||||
def test_kv_events_is_null_when_no_publisher_configured(self):
|
||||
args = ServerArgs(model_path="dummy") # no --kv-events-config
|
||||
|
||||
info = _call_server_info_with(args)
|
||||
|
||||
# The key must still be present so consumers can detect
|
||||
# "publishing disabled" via a single shape check.
|
||||
self.assertIn("kv_events", info)
|
||||
self.assertIsNone(info["kv_events"])
|
||||
|
||||
def test_kv_events_is_null_when_publisher_explicitly_null(self):
|
||||
args = ServerArgs(
|
||||
model_path="dummy",
|
||||
kv_events_config='{"publisher": "null"}',
|
||||
page_size=64,
|
||||
)
|
||||
|
||||
info = _call_server_info_with(args)
|
||||
|
||||
self.assertIsNone(info["kv_events"])
|
||||
|
||||
# ----- malformed config --------------------------------------------
|
||||
|
||||
def test_kv_events_is_null_for_malformed_json(self):
|
||||
# Not JSON — the publisher would have failed at server startup,
|
||||
# but /server_info must keep working.
|
||||
args = ServerArgs(
|
||||
model_path="dummy",
|
||||
kv_events_config="not-json",
|
||||
page_size=64,
|
||||
)
|
||||
|
||||
info = _call_server_info_with(args)
|
||||
|
||||
self.assertIsNone(info["kv_events"])
|
||||
|
||||
# ----- unreachable endpoints ---------------------------------------
|
||||
|
||||
def test_kv_events_is_null_for_inproc_endpoint(self):
|
||||
# `inproc://` is not reachable across process boundaries, so the
|
||||
# descriptor must hide it from external routers.
|
||||
args = ServerArgs(
|
||||
model_path="dummy",
|
||||
kv_events_config=(
|
||||
'{"publisher": "zmq", "endpoint": "inproc://cache", "topic": ""}'
|
||||
),
|
||||
page_size=64,
|
||||
)
|
||||
|
||||
info = _call_server_info_with(args)
|
||||
|
||||
self.assertIsNone(info["kv_events"])
|
||||
|
||||
def test_kv_events_is_null_when_endpoint_missing_port(self):
|
||||
args = ServerArgs(
|
||||
model_path="dummy",
|
||||
kv_events_config=(
|
||||
'{"publisher": "zmq", "endpoint": "tcp://0.0.0.0", "topic": ""}'
|
||||
),
|
||||
page_size=64,
|
||||
)
|
||||
|
||||
info = _call_server_info_with(args)
|
||||
|
||||
self.assertIsNone(info["kv_events"])
|
||||
|
||||
def test_kv_events_is_null_when_port_not_integer(self):
|
||||
args = ServerArgs(
|
||||
model_path="dummy",
|
||||
kv_events_config=(
|
||||
'{"publisher": "zmq", "endpoint": "tcp://0.0.0.0:abc", "topic": ""}'
|
||||
),
|
||||
page_size=64,
|
||||
)
|
||||
|
||||
info = _call_server_info_with(args)
|
||||
|
||||
self.assertIsNone(info["kv_events"])
|
||||
|
||||
def test_kv_events_is_null_for_port_out_of_range(self):
|
||||
# TCP ports are 1..65535; values outside the range can't bind, so
|
||||
# the descriptor refuses to advertise them rather than handing
|
||||
# subscribers a non-dialable address.
|
||||
for bad_port in (0, -1, 65536, 1_000_000):
|
||||
with self.subTest(port=bad_port):
|
||||
args = ServerArgs(
|
||||
model_path="dummy",
|
||||
kv_events_config=(
|
||||
f'{{"publisher": "zmq", "endpoint": "tcp://0.0.0.0:{bad_port}", "topic": ""}}'
|
||||
),
|
||||
page_size=64,
|
||||
)
|
||||
info = _call_server_info_with(args)
|
||||
self.assertIsNone(info["kv_events"])
|
||||
|
||||
# ----- bad scheduler context ---------------------------------------
|
||||
|
||||
def test_kv_events_is_null_when_page_size_missing_or_non_positive(self):
|
||||
# Without a real positive `page_size` the descriptor's
|
||||
# `block_size` would be a misleading placeholder; subscribers
|
||||
# would hash prompts at the wrong granularity and miss every
|
||||
# cache entry. Refuse to advertise instead.
|
||||
good_cfg = '{"publisher": "zmq", "endpoint": "tcp://*:5557", "topic": ""}'
|
||||
for bad_page_size in (None, 0, -1):
|
||||
with self.subTest(page_size=bad_page_size):
|
||||
args = ServerArgs(
|
||||
model_path="dummy",
|
||||
kv_events_config=good_cfg,
|
||||
page_size=bad_page_size,
|
||||
)
|
||||
info = _call_server_info_with(args)
|
||||
self.assertIsNone(info["kv_events"])
|
||||
|
||||
|
||||
class TestServerInfoExistingFieldsPreserved(CustomTestCase):
|
||||
"""Regression guard: the new `kv_events` field is additive — none of
|
||||
the fields existing consumers depend on may be silently dropped.
|
||||
|
||||
Existing `/server_info` consumers in the wild include:
|
||||
* SGLang's own deprecated `/get_server_info` (forwards to the
|
||||
same handler).
|
||||
* External monitoring tools that scrape the full ServerArgs.
|
||||
* KV-aware routers reading `kv_events_config`, `page_size`,
|
||||
`dp_size` directly to derive subscription info (this is the
|
||||
path the new `kv_events` block enriches but does not replace).
|
||||
"""
|
||||
|
||||
def test_every_server_args_field_appears_in_response(self):
|
||||
# `dataclasses.asdict(server_args)` is spread into the response;
|
||||
# asserting every dataclass field surfaces is the strongest
|
||||
# backward-compat guarantee that's still implementation-agnostic.
|
||||
args = ServerArgs(model_path="dummy")
|
||||
|
||||
info = _call_server_info_with(args)
|
||||
|
||||
for field in dataclasses.fields(ServerArgs):
|
||||
self.assertIn(
|
||||
field.name,
|
||||
info,
|
||||
f"existing ServerArgs field '{field.name}' missing from "
|
||||
f"/server_info response — kv_events patch must not "
|
||||
f"shadow or drop ServerArgs fields",
|
||||
)
|
||||
|
||||
def test_internal_states_and_version_keys_preserved(self):
|
||||
# These two top-level keys predate the kv_events patch and are
|
||||
# named individually (not spread from a dataclass), so a stray
|
||||
# edit could remove them without breaking syntax. Lock them down.
|
||||
args = ServerArgs(model_path="dummy")
|
||||
|
||||
info = _call_server_info_with(args)
|
||||
|
||||
self.assertIn("internal_states", info)
|
||||
self.assertIn("version", info)
|
||||
|
||||
def test_kv_events_config_raw_field_still_surfaced(self):
|
||||
# The new structured `kv_events` block sits alongside the
|
||||
# pre-existing flat `kv_events_config` field (the raw CLI string
|
||||
# already on ServerArgs). Both must remain visible so
|
||||
# consumers that hand-parse the raw config keep working.
|
||||
raw_cfg = '{"publisher": "zmq", "endpoint": "tcp://*:5557", "topic": ""}'
|
||||
args = ServerArgs(
|
||||
model_path="dummy",
|
||||
kv_events_config=raw_cfg,
|
||||
page_size=64,
|
||||
dp_size=1,
|
||||
)
|
||||
|
||||
info = _call_server_info_with(args)
|
||||
|
||||
self.assertIn("kv_events_config", info)
|
||||
self.assertEqual(info["kv_events_config"], raw_cfg)
|
||||
# And the new structured block is separately present:
|
||||
self.assertIn("kv_events", info)
|
||||
self.assertIsNotNone(info["kv_events"])
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
Reference in New Issue
Block a user