diff --git a/components/src/dynamo/sglang/args.py b/components/src/dynamo/sglang/args.py index 4d0970255c5d..7e94c9f0fe09 100644 --- a/components/src/dynamo/sglang/args.py +++ b/components/src/dynamo/sglang/args.py @@ -44,6 +44,7 @@ from dynamo.sglang.capacity import sglang_dp_layout from dynamo.sglang.elastic_ep_preflight import check_elastic_ep_backend from dynamo.sglang.gateway import effective_gateway_workers +from dynamo.sglang.state_agent import validate_state_agent_worker configure_dynamo_logging() PREFILL_DECODE_DISAGGREGATION_MODE = "pd" @@ -105,6 +106,7 @@ def __init__( # SGLang chose the value or the user did. self.attention_backend_from_cli = attention_backend_from_cli self.serving_mode = self._set_serving_strategy() + validate_state_agent_worker(self) if getattr(dynamo_args, "freeze_gc_after_init", False) and ( self.serving_mode == DisaggregationMode.PREFILL or any( diff --git a/components/src/dynamo/sglang/backend_args.py b/components/src/dynamo/sglang/backend_args.py index 6a51d992f737..9cbb2fe35c12 100644 --- a/components/src/dynamo/sglang/backend_args.py +++ b/components/src/dynamo/sglang/backend_args.py @@ -56,6 +56,15 @@ def add_arguments(self, parser: argparse.ArgumentParser) -> None: g = parser.add_argument_group("Dynamo SGLang Options") + add_argument( + g, + flag_name="--kvcr-state-agent-owner", + env_var="DYN_SGL_KVCR_STATE_AGENT_OWNER", + default=None, + arg_type=str, + help="Attach an aggregated direct KVCR linker to the external KV state agent using this stable CacheOwnerId.", + ) + add_negatable_bool_argument( g, flag_name="--use-sglang-tokenizer", @@ -216,6 +225,7 @@ class DynamoSGLangConfig(ConfigBase): use_sglang_tokenizer: bool gateway_workers: Optional[int] = None + kvcr_state_agent_owner: Optional[str] = None # Internal roles derived from the canonical multimodal arguments in args.py. multimodal_encode_worker: bool = False multimodal_worker: bool = False diff --git a/components/src/dynamo/sglang/init_llm.py b/components/src/dynamo/sglang/init_llm.py index 098cf64be0da..882a4d533c2f 100644 --- a/components/src/dynamo/sglang/init_llm.py +++ b/components/src/dynamo/sglang/init_llm.py @@ -34,6 +34,7 @@ ) from dynamo.sglang.register import register_model_with_readiness_gate from dynamo.sglang.request_handlers import DecodeWorkerHandler, PrefillWorkerHandler +from dynamo.sglang.state_agent import StateAgentLifecycle async def _warmup_prefill_engine( @@ -201,7 +202,9 @@ async def init_decode( ) body_failed = True + state_agent = StateAgentLifecycle() try: + await state_agent.start(config, generate_endpoint) gather_tasks = [ generate_endpoint.serve_endpoint( handler.generate, @@ -246,12 +249,15 @@ async def init_decode( else: body_failed = False finally: - await finish_worker_teardown( - metrics_task, - handler.cleanup, - run_deferred_handlers, - body_failed=body_failed, - ) + try: + await state_agent.close() + finally: + await finish_worker_teardown( + metrics_task, + handler.cleanup, + run_deferred_handlers, + body_failed=body_failed, + ) async def init_prefill( diff --git a/components/src/dynamo/sglang/kv_hints.py b/components/src/dynamo/sglang/kv_hints.py index 85c5540f1ed3..160b8481c471 100644 --- a/components/src/dynamo/sglang/kv_hints.py +++ b/components/src/dynamo/sglang/kv_hints.py @@ -104,20 +104,22 @@ def kvcr_hint_endpoints( } -def publish_kvcr_hint_capabilities( - runtime_config, engine, server_args, worker_type -) -> None: - """Opt in only for the supported single-host KVCR linker or HiCache path.""" +def kvcr_hint_sources( + server_args: Any, worker_type: Any +) -> tuple[str, dict[str, str]] | None: + """The role and per-DP-rank source-control endpoints this worker advertises. + + None when the worker runs no supported single-host KVCR linker or HiCache + path with remote hints enabled. + """ server_args = resolved_server_args(server_args) mode = kvcr_mode(server_args) if ( - engine is None - or mode is None + mode is None or getattr(server_args, "nnodes", 1) != 1 or getattr(server_args, "pp_size", 1) != 1 - or not filter_supported_async_generate_kwargs(engine, {"kv_hints": None}) ): - return + return None extra = server_args.hicache_storage_backend_extra_config if isinstance(extra, str): extra = extra.strip() @@ -133,7 +135,7 @@ def publish_kvcr_hint_capabilities( raise ValueError("KVCR HiCache extra configuration must be an object") # Legacy linker revisions have no switch; newer ones can disable hints. if extra.get("enable_remote_hint", True) is not True: - return + return None if ( mode == "linker" and extra.get("advertise_host") is None @@ -141,7 +143,7 @@ def publish_kvcr_hint_capabilities( ): # The newer control_advertise_host dialect defaults remote hints off. # The legacy advertise_host runtime has no switch and enables peers. - return + return None if mode == "linker" and ( "control_port" not in extra or ( @@ -151,14 +153,14 @@ def publish_kvcr_hint_capabilities( ): # Local-only/default linkers can still serve, but cannot advertise a # guessed source-control endpoint to another worker. - return + return None # WorkerType is a PyO3 enum, not enum.Enum: its wire role is __str__, # and accessing .value aborts registration after the model has loaded. role = str(worker_type) if role == "agg": role = "aggregated" if role not in ("aggregated", "prefill", "decode"): - return + return None # DP attention partitions TP into independent attention/cache groups. dp_size = server_args.attn_dp_size attn_cp_size = server_args.attn_cp_size @@ -171,6 +173,21 @@ def publish_kvcr_hint_capabilities( pp_size=server_args.pp_size, mode=mode, ) + return role, endpoints + + +def publish_kvcr_hint_capabilities( + runtime_config, engine, server_args, worker_type +) -> None: + """Opt in only for the supported single-host KVCR linker or HiCache path.""" + if engine is None or not filter_supported_async_generate_kwargs( + engine, {"kv_hints": None} + ): + return + sources = kvcr_hint_sources(server_args, worker_type) + if sources is None: + return + role, endpoints = sources for key, value in ( (KV_HINT_TRANSFER_WORKER_TYPE_RUNTIME_KEY, role), (KV_HINT_TRANSFER_SOURCE_CONTROL_ENDPOINTS_RUNTIME_KEY, endpoints), diff --git a/components/src/dynamo/sglang/publisher.py b/components/src/dynamo/sglang/publisher.py index 64375ba4127d..9593049eac2b 100644 --- a/components/src/dynamo/sglang/publisher.py +++ b/components/src/dynamo/sglang/publisher.py @@ -396,7 +396,13 @@ def init_kv_event_publish(self) -> List[KvEventPublisher]: List of KvEventPublisher instances if KV event publishing is enabled, empty list otherwise. """ - if self.dynamo_args.use_kv_events and not publishes_kv_events(self.server_args): + if getattr(self.dynamo_args, "kvcr_state_agent_owner", None) is not None: + # The state agent consumes the raw stream; legacy normalization + # rejects KVCR-owned residency events. Keep use_kv_events advertised. + pass + elif self.dynamo_args.use_kv_events and not publishes_kv_events( + self.server_args + ): logging.info( "Non-leader node (node_rank=%s) shares the leader's single KV " "rank slice; skipping KV event publishing so the router sees " diff --git a/components/src/dynamo/sglang/register.py b/components/src/dynamo/sglang/register.py index 7a0a7efd9a39..ee3de4056ca6 100644 --- a/components/src/dynamo/sglang/register.py +++ b/components/src/dynamo/sglang/register.py @@ -562,6 +562,8 @@ async def get_runtime_config( dynamo_args.enable_local_indexer and not is_decode_worker ) runtime_config.kv_event_publishing_enabled = dynamo_args.use_kv_events + if getattr(dynamo_args, "kvcr_state_agent_owner", None) is not None: + runtime_config.kv_event_source_mode = "state_agent_v2" start_dp_rank, end_dp_rank = model_card_dp_rank_bounds(server_args) registered_dp_size = end_dp_rank - start_dp_rank diff --git a/components/src/dynamo/sglang/state_agent.py b/components/src/dynamo/sglang/state_agent.py new file mode 100644 index 000000000000..21ae102711f7 --- /dev/null +++ b/components/src/dynamo/sglang/state_agent.py @@ -0,0 +1,109 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Opt-in state-agent attachment for an aggregated KVCR linker worker.""" + +import asyncio +import json +import logging +import re +from typing import Any +from urllib.parse import urlparse + +from dynamo.common.constants import DisaggregationMode +from dynamo.llm import KvStateAttachmentOwner, WorkerType +from dynamo.sglang._compat import resolved_server_args +from dynamo.sglang.capacity import kv_event_block_size +from dynamo.sglang.kv_hints import kvcr_hint_sources, kvcr_mode + +_CACHE_OWNER_PATTERN = re.compile( + r"^[0-9a-f]{32}:[0-9a-f]{32}/[0-9a-f]{16}/[0-9a-f]{32}$" +) +_CLEANUP_TIMEOUT_SECS = 30.0 + + +class StateAgentLifecycle: + """Start one attachment after engine init and close it once on shutdown.""" + + def __init__(self) -> None: + self._owner: KvStateAttachmentOwner | None = None + + async def start(self, config: Any, endpoint: Any) -> None: + descriptor = state_agent_descriptor(config) + if descriptor is None: + return + owner = KvStateAttachmentOwner(endpoint, endpoint.connection_id(), [descriptor]) + # Cancelling this await drops the native waiter, which withdraws its + # own registrations; nothing is left for close() to undo. + await owner.start() + self._owner = owner + + async def close(self) -> None: + owner, self._owner = self._owner, None + if owner is None: + return + try: + await asyncio.wait_for(owner.close(), _CLEANUP_TIMEOUT_SECS) + except asyncio.TimeoutError: + logging.warning( + "KV state attachment cleanup timed out; native cleanup continues" + ) + + +def validate_state_agent_worker(config: Any) -> str | None: + """Validate launch intent without forcing SGLang's engine resolution.""" + dynamo_args = config.dynamo_args + owner = getattr(dynamo_args, "kvcr_state_agent_owner", None) + if owner is None: + return None + if not isinstance(owner, str) or not _CACHE_OWNER_PATTERN.fullmatch(owner): + raise ValueError("--kvcr-state-agent-owner requires a canonical CacheOwnerId") + server_args = resolved_server_args(config.server_args) + if config.serving_mode != DisaggregationMode.AGGREGATED: + raise ValueError("KVCR state agent supports only aggregated LLM workers") + if kvcr_mode(server_args) != "linker" or not dynamo_args.use_kv_events: + raise ValueError( + "KVCR state agent requires the direct KVCR linker and KV events" + ) + return owner + + +def state_agent_descriptor(config: Any) -> dict[str, Any] | None: + """Read final event and linker coordinates after engine initialization.""" + owner = validate_state_agent_worker(config) + if owner is None: + return None + server_args = resolved_server_args(config.server_args) + dynamo_args = config.dynamo_args + events = server_args.kv_events_config + events = json.loads(events) if isinstance(events, str) else events + if not isinstance(events, dict) or events.get("publisher") != "zmq": + raise ValueError("KVCR state agent requires a ZMQ event publisher") + if events.get("topic", "") != "": + raise ValueError("KVCR state agent requires an empty event topic") + endpoint = urlparse(events.get("endpoint", "")) + if endpoint.scheme != "tcp" or not endpoint.hostname or not endpoint.port: + raise ValueError("KVCR state agent requires an explicit TCP event port") + # Registration advertises the same source; its resolver already encodes + # the linker dialects, and the raw event stream shares the advertise host. + sources = kvcr_hint_sources(server_args, WorkerType.Aggregated) + if sources is None: + raise ValueError("KVCR state agent requires remote linker hints") + hint_endpoint = sources[1]["0"] + host = urlparse(hint_endpoint).netloc.rsplit(":", 1)[0] + return { + "cache_owner_id": owner, + "global_dp_rank": 0, + "kv_state_endpoint": dynamo_args.kv_state_endpoint + or f"{dynamo_args.namespace}/{dynamo_args.component}/{dynamo_args.endpoint}", + "indexer_domain_id": owner.split("/", 1)[0], + "kv_block_size": kv_event_block_size(server_args), + "raw_zmq_endpoint": f"tcp://{host}:{endpoint.port}", + "raw_topic": "", + "image_token_id": None, + "video_token_id": None, + "router_hint_source": { + "source_control_endpoint": hint_endpoint, + "worker_type": "aggregated", + }, + } diff --git a/components/src/dynamo/sglang/tests/test_sglang_state_agent.py b/components/src/dynamo/sglang/tests/test_sglang_state_agent.py new file mode 100644 index 000000000000..3925b54e5056 --- /dev/null +++ b/components/src/dynamo/sglang/tests/test_sglang_state_agent.py @@ -0,0 +1,211 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +import json +from types import SimpleNamespace +from unittest.mock import AsyncMock, Mock + +import pytest + +from dynamo.common.constants import DisaggregationMode +from dynamo.sglang.state_agent import StateAgentLifecycle, state_agent_descriptor + +OWNER = f"{'01' * 16}:{'02' * 16}/{'03' * 8}/{'04' * 16}" + +pytestmark = [ + pytest.mark.unit, + pytest.mark.sglang, + pytest.mark.router, + pytest.mark.gpu_0, + pytest.mark.profiled_vram_gib(0), + pytest.mark.pre_merge, +] + + +def _config(owner=None): + return SimpleNamespace( + serving_mode=DisaggregationMode.AGGREGATED, + dynamo_args=SimpleNamespace( + kvcr_state_agent_owner=owner, + namespace="ns", + component="backend", + endpoint="generate", + kv_state_endpoint=None, + use_kv_events=True, + ), + server_args=SimpleNamespace( + enable_unified_cache_external_linker=True, + unified_cache_external_linker_backend="kvcr", + hicache_storage_backend_extra_config=json.dumps( + {"advertise_host": "10.0.0.8", "control_port": 19500} + ), + kv_events_config=json.dumps( + {"publisher": "zmq", "endpoint": "tcp://*:5557"} + ), + tp_size=1, + attn_dp_size=1, + attn_cp_size=1, + pp_size=1, + nnodes=1, + page_size=64, + ), + ) + + +@pytest.mark.asyncio +async def test_opt_out_touches_neither_kvcr_nor_the_native_attachment(monkeypatch): + factory = Mock() + monkeypatch.setattr("dynamo.sglang.state_agent.KvStateAttachmentOwner", factory) + config = _config() + config.server_args = None + assert state_agent_descriptor(config) is None + lifecycle = StateAgentLifecycle() + await lifecycle.start(config, Mock()) + await lifecycle.close() + factory.assert_not_called() + + +def test_descriptor_preserves_owner_and_exact_source_coordinates(): + # These ports are protocol fixtures; no sockets are opened. + assert state_agent_descriptor(_config(OWNER)) == { + "cache_owner_id": OWNER, + "global_dp_rank": 0, + "kv_state_endpoint": "ns/backend/generate", + "indexer_domain_id": OWNER.split("/", 1)[0], + "kv_block_size": 64, + "raw_zmq_endpoint": "tcp://10.0.0.8:5557", + "raw_topic": "", + "image_token_id": None, + "video_token_id": None, + "router_hint_source": { + "source_control_endpoint": "tcp://10.0.0.8:19500", + "worker_type": "aggregated", + }, + } + + +_HINTS_OFF = ( + '{"advertise_host":"10.0.0.8","control_port":19500,"enable_remote_hint":false}' +) +_TOPIC = '{"publisher":"zmq","endpoint":"tcp://*:5557","topic":"other"}' + + +@pytest.mark.parametrize( + "target,field,value,match", + [ + ("config", "serving_mode", DisaggregationMode.DECODE, "aggregated"), + ("server", "pp_size", 2, "linker hints"), + ("server", "nnodes", 2, "linker hints"), + ("server", "enable_unified_cache_external_linker", False, "linker"), + ("dynamo", "use_kv_events", False, "KV events"), + ("server", "hicache_storage_backend_extra_config", _HINTS_OFF, "linker hints"), + *[ + ( + "server", + "kv_events_config", + json.dumps({"publisher": "zmq", "endpoint": e}), + "event", + ) + for e in ( + "ipc:///tmp/events", + "tcp://*:0", + ) + ], + ("server", "kv_events_config", _TOPIC, "topic"), + *[ + ("dynamo", "kvcr_state_agent_owner", o, "CacheOwnerId") + for o in ("", "not-an-owner", OWNER.replace("01", "AA")) + ], + ], +) +def test_rejects_invalid_configuration(target, field, value, match): + config = _config(OWNER) + holder = { + "config": config, + "server": config.server_args, + "dynamo": config.dynamo_args, + } + setattr(holder[target], field, value) + with pytest.raises(ValueError, match=match): + state_agent_descriptor(config) + + +@pytest.mark.asyncio +async def test_lifecycle_starts_one_owner_and_closes_once(monkeypatch): + owner = SimpleNamespace(start=AsyncMock(), close=AsyncMock()) + factory = Mock(return_value=owner) + monkeypatch.setattr("dynamo.sglang.state_agent.KvStateAttachmentOwner", factory) + endpoint = Mock() + endpoint.connection_id.return_value = 123 + lifecycle = StateAgentLifecycle() + await lifecycle.start(_config(OWNER), endpoint) + factory.assert_called_once_with( + endpoint, 123, [state_agent_descriptor(_config(OWNER))] + ) + owner.start.assert_awaited_once() + await lifecycle.close() + await lifecycle.close() + owner.close.assert_awaited_once() + + +def test_config_rejects_bad_owner_before_engine_startup(): + from dynamo.sglang.args import Config + + config = _config("invalid") + config.server_args.disaggregation_mode = "null" + with pytest.raises(ValueError, match="CacheOwnerId"): + Config(config.server_args, config.dynamo_args) + + +def test_early_launch_validation_does_not_require_final_page_size(): + from sglang.srt.server_args import ServerArgs + + from dynamo.sglang.args import Config + + config = _config(OWNER) + raw = ServerArgs( + model_path="dummy", + page_size=None, + attn_cp_size=None, + enable_unified_cache_external_linker=True, + unified_cache_external_linker_backend="kvcr", + kv_events_config=config.server_args.kv_events_config, + ) + assert Config(raw, config.dynamo_args).server_args is raw + + +@pytest.mark.asyncio +@pytest.mark.parametrize("owner", [None, OWNER]) +async def test_registration_advertises_state_agent_only_when_opted_in(owner): + from sglang.srt.server_args import ServerArgs + + from dynamo.sglang.register import DynamoConfig, get_runtime_config + + args = DynamoConfig() + args.kvcr_state_agent_owner = owner + args.use_kv_events = True + args.enable_local_indexer = False + runtime_config = await get_runtime_config( + None, ServerArgs(model_path="dummy", page_size=64), args + ) + assert runtime_config.kv_event_publishing_enabled is True + assert runtime_config.kv_event_source_mode == ("state_agent_v2" if owner else None) + + +@pytest.mark.parametrize("opt_in", [False, True]) +def test_only_opt_in_suppresses_legacy_publisher(opt_in, monkeypatch): + import dynamo.sglang.publisher as module + + config = _config(OWNER if opt_in else None) + config.dynamo_args.enable_local_indexer = False + publisher = module.DynamoSglangPublisher.__new__(module.DynamoSglangPublisher) + publisher.dynamo_args = config.dynamo_args + publisher.server_args = config.server_args + publisher.generate_endpoint = Mock() + publisher.kv_worker_id = 123 + publisher.kv_publishers = [] + factory = Mock() + monkeypatch.setattr(module, "KvEventPublisher", factory) + assert len(publisher.init_kv_event_publish()) == (0 if opt_in else 1) + assert factory.call_count == (0 if opt_in else 1) + assert config.dynamo_args.use_kv_events is True