From 302ebbcaf936e54cdfbd3555ab77ce18fc738f02 Mon Sep 17 00:00:00 2001 From: Ziqi Fan Date: Wed, 28 Oct 2026 12:09:16 -0700 Subject: [PATCH 1/9] feat(kv-router): support process-local HiCache inventory owners Signed-off-by: Ziqi Fan --- .../python/rust/llm/kv_state_agent.rs | 5 +- lib/llm/src/discovery/kv_state_agent.rs | 9 ++ .../kv_router/indexer/recovery/state_agent.rs | 49 ++++++-- .../kv_router/publisher/attachment_owner.rs | 2 +- .../src/kv_router/publisher/state_agent.rs | 107 ++++++++++++++---- .../kv_router/publisher/state_agent_host.rs | 8 +- 6 files changed, 147 insertions(+), 33 deletions(-) diff --git a/lib/bindings/python/rust/llm/kv_state_agent.rs b/lib/bindings/python/rust/llm/kv_state_agent.rs index a1e645a6881f..c88f70c7b81a 100644 --- a/lib/bindings/python/rust/llm/kv_state_agent.rs +++ b/lib/bindings/python/rust/llm/kv_state_agent.rs @@ -96,6 +96,8 @@ struct AttachmentDescriptorInput { kv_state_endpoint: String, indexer_domain_id: String, kv_block_size: u32, + #[serde(default)] + ingress_protocol: Option, raw_zmq_endpoint: String, #[serde(default)] raw_topic: String, @@ -133,8 +135,9 @@ impl KvStateAttachmentOwner { kv_state_endpoint: input.kv_state_endpoint.parse()?, indexer_domain_id: input.indexer_domain_id.parse()?, kv_block_size: input.kv_block_size, - ingress_protocol: + ingress_protocol: input.ingress_protocol.unwrap_or( llm_rs::discovery::kv_state_agent::KvStateIngressProtocol::VllmResidencyV1, + ), raw_zmq_endpoint: input.raw_zmq_endpoint, raw_topic: input.raw_topic, image_token_id: input.image_token_id, diff --git a/lib/llm/src/discovery/kv_state_agent.rs b/lib/llm/src/discovery/kv_state_agent.rs index 8e10c9906092..4834d2571f83 100644 --- a/lib/llm/src/discovery/kv_state_agent.rs +++ b/lib/llm/src/discovery/kv_state_agent.rs @@ -59,6 +59,15 @@ pub fn attachment_record_id(publisher_id: PublisherId, generation: u64) -> Publi pub enum KvStateIngressProtocol { FrameworkV1, VllmResidencyV1, + /// HiCache backend memory is owned by the attached engine, not a Guard. + #[serde(rename = "sglang_hicache_v1")] + SglangHiCacheV1, +} + +impl KvStateIngressProtocol { + pub fn supports_owned_residency(self) -> bool { + matches!(self, Self::VllmResidencyV1 | Self::SglangHiCacheV1) + } } #[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] diff --git a/lib/llm/src/kv_router/indexer/recovery/state_agent.rs b/lib/llm/src/kv_router/indexer/recovery/state_agent.rs index b90063bd1a64..c411f5dbefcf 100644 --- a/lib/llm/src/kv_router/indexer/recovery/state_agent.rs +++ b/lib/llm/src/kv_router/indexer/recovery/state_agent.rs @@ -1424,6 +1424,23 @@ fn update_advertisements( } } +fn compatible_state_source( + discovery_id: u64, + source: &KvStateSourceAdvertisement, + endpoint: &dynamo_runtime::protocols::EndpointId, + block_size: u32, +) -> bool { + source.publisher_id == discovery_id + && source.protocol_version == KvStateProtocolVersion::V2 + && source.event_topic == KV_STATE_EVENT_TOPIC_V2 + && source.kv_state_endpoint == *endpoint + && source.kv_block_size == block_size + // The router consumes canonical V2 residency, not either engine's raw + // ingress. Source/attachment protocol equality is checked separately. + && source.ingress_protocol.supports_owned_residency() + && source.cache_owner_id.pool().indexer_domain() == source.indexer_domain_id +} + #[allow(clippy::too_many_arguments)] async fn schedule_state_sources( indexer: &Indexer, @@ -1454,17 +1471,8 @@ async fn schedule_state_sources( let mut source_by_owner: HashMap> = HashMap::new(); for (discovery_id, source) in sources { - let expected_ingress_protocol = - crate::discovery::kv_state_agent::KvStateIngressProtocol::VllmResidencyV1; let expected_domain = source.cache_owner_id.pool().indexer_domain(); - if source.publisher_id != *discovery_id - || source.protocol_version != KvStateProtocolVersion::V2 - || source.event_topic != KV_STATE_EVENT_TOPIC_V2 - || source.kv_state_endpoint != *endpoint - || source.kv_block_size != expected_block_size - || source.ingress_protocol != expected_ingress_protocol - || expected_domain != source.indexer_domain_id - { + if !compatible_state_source(*discovery_id, source, endpoint, expected_block_size) { tracing::warn!( discovery_id = *discovery_id, advertised_publisher_id = source.publisher_id, @@ -1477,7 +1485,7 @@ async fn schedule_state_sources( advertised_block_size = source.kv_block_size, expected_block_size, advertised_ingress_protocol = ?source.ingress_protocol, - expected_ingress_protocol = ?expected_ingress_protocol, + expected_ingress_protocol = "owned-residency", advertised_domain = %source.indexer_domain_id, expected_domain = %expected_domain, owner = %source.cache_owner_id, @@ -2483,6 +2491,25 @@ mod tests { } } + #[test] + fn source_compatibility_accepts_owned_ingress_contracts() { + let mut source = source_advertisement(owner(4)); + let endpoint = source.kv_state_endpoint.clone(); + for protocol in [ + KvStateIngressProtocol::VllmResidencyV1, + KvStateIngressProtocol::SglangHiCacheV1, + KvStateIngressProtocol::FrameworkV1, + ] { + source.ingress_protocol = protocol; + assert_eq!( + compatible_state_source(41, &source, &endpoint, 4), + protocol.supports_owned_residency() + ); + assert!(!compatible_state_source(42, &source, &endpoint, 4)); + assert!(!compatible_state_source(41, &source, &endpoint, 8)); + } + } + fn attachment_advertisement( owner: CacheOwnerId, generation: u64, diff --git a/lib/llm/src/kv_router/publisher/attachment_owner.rs b/lib/llm/src/kv_router/publisher/attachment_owner.rs index e6608a85f4bd..9be267718010 100644 --- a/lib/llm/src/kv_router/publisher/attachment_owner.rs +++ b/lib/llm/src/kv_router/publisher/attachment_owner.rs @@ -176,7 +176,7 @@ fn validate_descriptors(descriptors: &[KvStateAttachmentDescriptor]) -> Result<( if descriptor.cache_owner_id.pool().indexer_domain() != descriptor.indexer_domain_id { anyhow::bail!("CacheOwnerId and indexer domain disagree"); } - if descriptor.ingress_protocol != KvStateIngressProtocol::VllmResidencyV1 { + if !descriptor.ingress_protocol.supports_owned_residency() { anyhow::bail!("unsupported state-agent raw ingress protocol"); } if descriptor.raw_zmq_endpoint.is_empty() { diff --git a/lib/llm/src/kv_router/publisher/state_agent.rs b/lib/llm/src/kv_router/publisher/state_agent.rs index 2abc9cdb882a..ca540716cd6a 100644 --- a/lib/llm/src/kv_router/publisher/state_agent.rs +++ b/lib/llm/src/kv_router/publisher/state_agent.rs @@ -367,7 +367,7 @@ impl Coordinator

{ self.attachment = Some(attachment.clone()); self.publish_status(); - let barrier_cursor = match self + let mut barrier_cursor = match self .apply_and_publish_reset(attachment.worker, ResidencyDomain::Worker) .await { @@ -379,6 +379,19 @@ impl Coordinator

{ .context("failed to establish Worker reset barrier before attach"); } }; + if attachment.ingress_protocol == KvStateIngressProtocol::SglangHiCacheV1 { + barrier_cursor = match self + .apply_and_publish_reset(attachment.worker, ResidencyDomain::CacheOwner) + .await + { + Ok(cursor) => cursor, + Err(error) => { + self.fail_cache_owner(attachment.generation, "HiCache attach reset failed"); + return Err(error).context("failed to reset process-local KVCR inventory"); + } + }; + self.cache_owner_domain_failed = false; + } attachment.ready_at_outbound_cursor = barrier_cursor; self.ingress_generation = Some(attachment.generation); attachment.ready = true; @@ -400,6 +413,9 @@ impl Coordinator

{ anyhow::bail!("stale attachment generation {generation}"); } attachment.ready = false; + if attachment.ingress_protocol == KvStateIngressProtocol::SglangHiCacheV1 { + self.cache_owner_domain_failed = true; + } self.publish_status(); Ok(()) } @@ -420,7 +436,7 @@ impl Coordinator

{ // residency truth and must be drained or fail CacheOwner closed. self.drain_detaching_cache_owner(generation).await; - let cursor = match self + let mut cursor = match self .apply_and_publish_reset(attachment.worker, ResidencyDomain::Worker) .await { @@ -431,6 +447,12 @@ impl Coordinator

{ return Err(error).context("failed to publish final Worker detach barrier"); } }; + if attachment.ingress_protocol == KvStateIngressProtocol::SglangHiCacheV1 { + cursor = self + .apply_and_publish_reset(attachment.worker, ResidencyDomain::CacheOwner) + .await + .context("failed to clear process-local KVCR inventory on detach")?; + } self.attachment = None; self.source_cursor = SourceCursorState::default(); self.publish_status(); @@ -553,6 +575,11 @@ impl Coordinator

{ if let Some(metrics) = kv_publisher_metrics() { metrics.increment_engines_dropped_events(missing); } + if self.attachment.as_ref().is_some_and(|attachment| { + attachment.ingress_protocol == KvStateIngressProtocol::SglangHiCacheV1 + }) { + self.fail_source(batch.generation, "process-local inventory cursor gap"); + } self.source_cursor.accept_after_loss(batch); } @@ -796,7 +823,11 @@ impl Coordinator

{ self.status.store(KvStateAgentStatus { identity: self.identity.clone(), attachment: self.attachment.as_ref().map(Attachment::status), - cache_owner_ready: !self.cache_owner_domain_failed, + cache_owner_ready: !self.cache_owner_domain_failed + && self.attachment.as_ref().is_none_or(|attachment| { + attachment.ingress_protocol != KvStateIngressProtocol::SglangHiCacheV1 + || attachment.ready + }), outbound_cursor: self.outbound_cursor(), }); } @@ -846,8 +877,8 @@ impl KvStateAgent { WORKER_KV_INDEXER_BUFFER_SIZE, )); - if config.ingress_protocol != KvStateIngressProtocol::VllmResidencyV1 { - anyhow::bail!("external state-agent slots require vllm_residency_v1 ingress"); + if !config.ingress_protocol.supports_owned_residency() { + anyhow::bail!("external state-agent slots require an owned-residency ingress"); } let event_publisher = match &config.sockets.event_socket { Some(socket) => { @@ -877,7 +908,7 @@ impl KvStateAgent { let status = Arc::new(KvStateAgentStatusCell::new(KvStateAgentStatus { identity: identity.clone(), attachment: None, - cache_owner_ready: true, + cache_owner_ready: config.ingress_protocol != KvStateIngressProtocol::SglangHiCacheV1, outbound_cursor: 0, })); @@ -1761,9 +1792,7 @@ fn normalize_raw_batch( continue; } }; - if ownership == KvEventOwnership::Kvcr - && ingress_protocol != KvStateIngressProtocol::VllmResidencyV1 - { + if ownership == KvEventOwnership::Kvcr && !ingress_protocol.supports_owned_residency() { source_fault.get_or_insert("KVCR event on framework-only raw protocol"); continue; } @@ -2724,6 +2753,16 @@ mod tests { #[tokio::test] async fn detach_accepts_cache_owner_ingress_after_readiness_withdrawal() { + check_detach_lifetime(KvStateIngressProtocol::VllmResidencyV1).await; + } + + #[tokio::test] + async fn sglang_detach_clears_process_local_owner() { + check_detach_lifetime(KvStateIngressProtocol::SglangHiCacheV1).await; + } + + async fn check_detach_lifetime(protocol: KvStateIngressProtocol) { + let process_local = protocol == KvStateIngressProtocol::SglangHiCacheV1; let owner = cache_owner_id(); let worker = WorkerWithDpRank::new(17, 3); let identity = KvStateAgentIdentity { @@ -2768,9 +2807,11 @@ mod tests { source_failed: false, }; let mut attachment = test_attachment(worker, "tcp://framework"); + attachment.ingress_protocol = protocol; coordinator.attach(&mut attachment).await.unwrap(); coordinator.begin_detach(attachment.generation).unwrap(); + assert_eq!(coordinator.status.load().cache_owner_ready, !process_local); coordinator .handle_ingress(IngressBatch { generation: attachment.generation, @@ -2791,7 +2832,7 @@ mod tests { .complete_detach(attachment.generation) .await .unwrap(), - 4 + if process_local { 5 } else { 4 } ); let events = output.lock().unwrap().clone(); @@ -2804,12 +2845,22 @@ mod tests { matches!(event.event.data, KvCacheEventData::Cleared), )) .collect::>(), - vec![ - (1, ResidencyDomain::Worker, true), - (2, ResidencyDomain::Worker, false), - (3, ResidencyDomain::CacheOwner, false), - (4, ResidencyDomain::Worker, true), - ] + if process_local { + vec![ + (1, ResidencyDomain::Worker, true), + (2, ResidencyDomain::CacheOwner, true), + (3, ResidencyDomain::Worker, false), + (4, ResidencyDomain::Worker, true), + (5, ResidencyDomain::CacheOwner, true), + ] + } else { + vec![ + (1, ResidencyDomain::Worker, true), + (2, ResidencyDomain::Worker, false), + (3, ResidencyDomain::CacheOwner, false), + (4, ResidencyDomain::Worker, true), + ] + } ); let WorkerKvQueryResponse::TreeDump { events, .. } = local_indexer.get_events_in_id_range(None, None).await @@ -2821,10 +2872,11 @@ mod tests { .iter() .filter(|event| event.state_source == Some(owner)) .count(), - 1 + usize::from(!process_local) ); let mut replacement = test_attachment(worker, "tcp://framework-2"); + replacement.ingress_protocol = protocol; coordinator.attach(&mut replacement).await.unwrap(); coordinator .handle_ingress(IngressBatch { @@ -2838,14 +2890,31 @@ mod tests { }) .await .unwrap(); + if process_local { + coordinator + .handle_ingress(IngressBatch { + generation: replacement.generation, + source_cursor: 3, + chunk_index: 0, + final_chunk: true, + events: vec![], + source_fault: None, + cache_owner_fault: None, + }) + .await + .unwrap(); + assert!(!coordinator.status.load().cache_owner_ready); + assert!(!coordinator.status.load().attachment.as_ref().unwrap().ready); + } coordinator.begin_detach(replacement.generation).unwrap(); coordinator .complete_detach(replacement.generation) .await .unwrap(); - assert!( + assert_eq!( coordinator.status.load().cache_owner_ready, - "live-stream continuity remains advisory until journal replay is implemented" + !process_local, + "only Guard-backed inventory may outlive a detached engine" ); } diff --git a/lib/llm/src/kv_router/publisher/state_agent_host.rs b/lib/llm/src/kv_router/publisher/state_agent_host.rs index d53a00441114..41f8bfc105b6 100644 --- a/lib/llm/src/kv_router/publisher/state_agent_host.rs +++ b/lib/llm/src/kv_router/publisher/state_agent_host.rs @@ -582,7 +582,7 @@ fn validate_intent( if intent.cache_owner_id.pool().indexer_domain() != intent.indexer_domain_id { anyhow::bail!("attachment intent indexer domain disagrees with CacheOwnerId"); } - if intent.ingress_protocol != KvStateIngressProtocol::VllmResidencyV1 { + if !intent.ingress_protocol.supports_owned_residency() { anyhow::bail!("unsupported state-agent raw ingress protocol"); } if intent.raw_zmq_endpoint.is_empty() { @@ -622,6 +622,12 @@ async fn reconcile_owner( 0 => { if let Some(slot) = slots.get_mut(&owner) { detach_slot(slot).await?; + if slot.ingress_protocol == KvStateIngressProtocol::SglangHiCacheV1 { + if let Some(agent) = slot.agent.as_ref() { + agent.shutdown().await?; + } + slot.agent = None; + } if slot.agent.is_none() { slots.remove(&owner); } From 5b6315538bff50992e25e48b03811576b21fb108 Mon Sep 17 00:00:00 2001 From: Ziqi Fan Date: Wed, 28 Oct 2026 12:09:16 -0700 Subject: [PATCH 2/9] feat(sglang): attach KVCR HiCache inventory to state agents Signed-off-by: Ziqi Fan --- components/src/dynamo/sglang/README.md | 33 ++++ components/src/dynamo/sglang/init_llm.py | 35 ++-- components/src/dynamo/sglang/publisher.py | 5 + components/src/dynamo/sglang/register.py | 3 + components/src/dynamo/sglang/state_agent.py | 174 ++++++++++++++++++ .../dynamo/sglang/tests/test_state_agent.py | 97 ++++++++++ 6 files changed, 335 insertions(+), 12 deletions(-) create mode 100644 components/src/dynamo/sglang/state_agent.py create mode 100644 components/src/dynamo/sglang/tests/test_state_agent.py diff --git a/components/src/dynamo/sglang/README.md b/components/src/dynamo/sglang/README.md index f68cbfdc654e..25f7430ad2ed 100644 --- a/components/src/dynamo/sglang/README.md +++ b/components/src/dynamo/sglang/README.md @@ -4,3 +4,36 @@ # SGLang See [docs/backends/sglang/](../../../../docs/fern/pages/developer-guide/knowledge-base/modular-components/backends/sglang/overview.md) for documentation. + +## KVCR-owned HiCache inventory + +KVCR inventory publication is opt-in for the in-process `python -m dynamo.sglang` +HiCacheStorage path. It requires a matching SGLang backend that reports the +`kvcr_inventory` capability and engine-lifetime cache-owner metadata. + +Keep KV events enabled and configure the KVCR backend with +`"publish_inventory": true`, remote hints enabled, and an explicit reachable +`control_advertise_host`. Use a ZMQ KV-event publisher with `"format": "dynamo"`. +Start a state-agent host in the workers' namespace before starting the workers: + +```bash +DYN_NAMESPACE= python -m dynamo.kv_state_agent +``` + +The worker reads the raw-event endpoint, incarnation-specific topic, cache owner, +and peer-control endpoint from SGLang startup metadata. It attaches before model +registration and uses `state_agent_v2` instead of the legacy event publisher. +The router can then distinguish framework GPU/HiCache-host residency from actual +KVCR-owned DRAM residency, even though both host caches use CPU memory. + +Upgrade the Dynamo frontend, worker, and state-agent host together when enabling +this mode. Existing deployments without the capability keep the legacy path; +vLLM attachment descriptors without an ingress protocol keep their existing +default. The new attachment requires a single-host LLM engine, PP=1, and one +tokenizer worker; incomplete source metadata is rejected at startup. + +This is a live process-local inventory, not durable recovery: detaching the +engine clears its KVCR owner, and a detected raw-event sequence gap withdraws +readiness until the worker is restarted. Router hints remain advisory; storage +lookup and completed delivery still determine usable pages. This integration +does not enable direct-device linker mode. diff --git a/components/src/dynamo/sglang/init_llm.py b/components/src/dynamo/sglang/init_llm.py index 098cf64be0da..d416f0b9874f 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(engine, 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( @@ -354,7 +360,9 @@ async def init_prefill( ready_event = asyncio.Event() body_failed = True + state_agent = StateAgentLifecycle() try: + await state_agent.start(engine, config, generate_endpoint) await asyncio.gather( generate_endpoint.serve_endpoint( handler.generate, @@ -406,9 +414,12 @@ async def init_prefill( 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, + ) diff --git a/components/src/dynamo/sglang/publisher.py b/components/src/dynamo/sglang/publisher.py index 64375ba4127d..11600b11b596 100644 --- a/components/src/dynamo/sglang/publisher.py +++ b/components/src/dynamo/sglang/publisher.py @@ -49,6 +49,7 @@ owns_engine_metrics, ) from dynamo.sglang.nixl_telemetry import scheduler_metrics_ipc_name +from dynamo.sglang.state_agent import inventory_sources def get_local_dp_rank_range(server_args) -> range: @@ -396,6 +397,10 @@ def init_kv_event_publish(self) -> List[KvEventPublisher]: List of KvEventPublisher instances if KV event publishing is enabled, empty list otherwise. """ + if inventory_sources(self.engine): + # The versioned state agent consumes the mixed framework/KVCR + # stream. The legacy normalizer rejects independently owned events. + return [] if 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 " diff --git a/components/src/dynamo/sglang/register.py b/components/src/dynamo/sglang/register.py index dacf6087b4b3..ffb3f15783c5 100644 --- a/components/src/dynamo/sglang/register.py +++ b/components/src/dynamo/sglang/register.py @@ -64,6 +64,7 @@ gateway_engine_id, ) from dynamo.sglang.kv_hints import publish_kvcr_hint_capabilities +from dynamo.sglang.state_agent import inventory_sources from dynamo.sglang.video_routing import publish_sglang_qwen_video_processor_contract SGLANG_HICACHE_MOONCAKE_RUNTIME_KEY = "sglang_hicache_mooncake" @@ -212,6 +213,8 @@ async def _register_model_with_runtime_config( runtime_config = await get_runtime_config(engine, server_args, dynamo_args) if runtime_config is not None and forwards_kv_hints: publish_kvcr_hint_capabilities(runtime_config, engine, server_args, worker_type) + if inventory_sources(engine): + runtime_config.kv_event_source_mode = "state_agent_v2" if dynamo_args.use_sglang_tokenizer and not ( output_type.supports_embedding() or output_type.supports_rerank() diff --git a/components/src/dynamo/sglang/state_agent.py b/components/src/dynamo/sglang/state_agent.py new file mode 100644 index 000000000000..f2e151e49a90 --- /dev/null +++ b/components/src/dynamo/sglang/state_agent.py @@ -0,0 +1,174 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Attach actual, process-local HiCache backend inventory to the KV state agent.""" + +from __future__ import annotations + +import asyncio +import ipaddress +import logging +import re +from collections.abc import Mapping +from urllib.parse import urlsplit + +from dynamo.common.constants import DisaggregationMode +from dynamo.llm import KvStateAttachmentOwner +from dynamo.sglang._compat import resolved_server_args +from dynamo.sglang.capacity import model_card_dp_rank_bounds +from dynamo.sglang.kv_hints import kvcr_hint_endpoints + +logger = logging.getLogger(__name__) +_OWNER_PATTERN = re.compile(r"^[0-9a-f]{32}:[0-9a-f]{32}/[0-9a-f]{16}/[0-9a-f]{32}$") + + +def inventory_sources(engine) -> dict[int, dict]: + """Read capabilities from scheduler metadata, never backend configuration.""" + initialization = getattr(engine, "_scheduler_init_result", None) + if initialization is None: + return {} + sources = {} + for info in initialization.scheduler_infos: + hints = info.get("kv_hint_sources", []) + if not isinstance(hints, list): + raise TypeError("SGLang KV hint sources must be a list") + for source in hints: + if not isinstance(source, Mapping): + raise TypeError("SGLang KV hint source must be an object") + capabilities = source.get("capabilities", []) + if not isinstance(capabilities, list): + raise TypeError("SGLang KV hint capabilities must be a list") + if "kvcr_inventory" not in capabilities: + continue + rank = source.get("dp_rank") + if type(rank) is not int or rank < 0: + raise ValueError("KVCR inventory source requires a nonnegative DP rank") + if rank in sources and sources[rank] != source: + raise ValueError("Conflicting KVCR inventory sources for one DP rank") + sources[rank] = source + return sources + + +def state_agent_descriptors(engine, config) -> list[dict]: + sources = inventory_sources(engine) + if not sources: + return [] + server_args = resolved_server_args(config.server_args) + if ( + not server_args.enable_hierarchical_cache + or getattr(server_args, "nnodes", 1) != 1 + or getattr(server_args, "pp_size", 1) != 1 + or getattr(server_args, "tokenizer_worker_num", 1) != 1 + or not config.dynamo_args.use_kv_events + ): + raise ValueError( + "KVCR inventory requires a single-host HiCache worker, one tokenizer " + "worker, PP=1 and KV event emission" + ) + roles = { + DisaggregationMode.AGGREGATED: "aggregated", + DisaggregationMode.PREFILL: "prefill", + DisaggregationMode.DECODE: "decode", + } + role = roles.get(config.serving_mode) + if role is None: + raise ValueError("KVCR inventory requires a hint-forwarding LLM worker") + bounds = model_card_dp_rank_bounds(server_args) + if set(sources) != set(range(*bounds)): + raise ValueError("KVCR inventory must cover every managed DP rank") + infos = engine._scheduler_init_result.scheduler_infos + endpoints = kvcr_hint_endpoints(infos, bounds) + raw_sources = {} + for info in infos: + for source in info.get("kv_event_sources", []): + rank = source.get("dp_rank") + if type(rank) is not int or rank not in sources: + raise ValueError("KVCR raw event source has an invalid DP rank") + if rank in raw_sources and raw_sources[rank] != source: + raise ValueError("Conflicting raw KV event sources for one DP rank") + raw_sources[rank] = source + if set(raw_sources) != set(sources): + raise ValueError("KVCR inventory requires a reported raw event source per rank") + descriptors = [] + owners = set() + for rank, source in sorted(sources.items()): + owner = source.get("cache_owner_id") + if not isinstance(owner, str) or _OWNER_PATTERN.fullmatch(owner) is None: + raise ValueError("KVCR inventory source requires a canonical CacheOwnerId") + if owner in owners: + raise ValueError("KVCR inventory owners must be unique across DP ranks") + owners.add(owner) + raw = raw_sources[rank] + if raw.get("format") != "dynamo": + raise ValueError("KVCR inventory requires Dynamo-format KV events") + endpoint = raw.get("endpoint") + if not isinstance(endpoint, str): + raise TypeError("KVCR inventory requires a reported TCP event endpoint") + parsed = urlsplit(endpoint) + if ( + parsed.scheme != "tcp" + or not parsed.hostname + or parsed.port is None + or not 0 < parsed.port <= 65535 + or parsed.username is not None + or parsed.password is not None + or parsed.path + or parsed.query + or parsed.fragment + or any(char.isspace() for char in parsed.hostname) + or parsed.hostname == "*" + ): + raise ValueError("Invalid reported KVCR raw event endpoint") + try: + address = ipaddress.ip_address(parsed.hostname) + except ValueError: + address = None + if address is not None and address.is_unspecified: + raise ValueError("KVCR raw event source cannot advertise a wildcard") + block_size = raw.get("block_size") + if type(block_size) is not int or block_size <= 0: + raise ValueError("KVCR raw event source requires a positive block size") + topic = raw.get("topic", "") + if not isinstance(topic, str) or not topic.strip(): + raise ValueError( + "KVCR raw event topic must identify this engine incarnation" + ) + descriptors.append( + { + "cache_owner_id": owner, + "global_dp_rank": rank, + "kv_state_endpoint": config.dynamo_args.kv_state_endpoint + or f"{config.dynamo_args.namespace}/{config.dynamo_args.component}/" + f"{config.dynamo_args.endpoint}", + "indexer_domain_id": owner.split("/", 1)[0], + "kv_block_size": block_size, + "ingress_protocol": "sglang_hicache_v1", + "raw_zmq_endpoint": endpoint, + "raw_topic": topic, + "image_token_id": None, + "video_token_id": None, + "router_hint_source": { + "source_control_endpoint": endpoints[str(rank)], + "worker_type": role, + }, + } + ) + return descriptors + + +class StateAgentLifecycle: + def __init__(self): + self._owner = None + + async def start(self, engine, config, endpoint): + descriptors = state_agent_descriptors(engine, config) + if not descriptors: + return + owner = KvStateAttachmentOwner(endpoint, endpoint.connection_id(), descriptors) + await owner.start() + self._owner = owner + logger.info("Attached %d KVCR HiCache inventory owners", len(descriptors)) + + async def close(self): + owner, self._owner = self._owner, None + if owner is not None: + await asyncio.wait_for(owner.close(), timeout=30.0) diff --git a/components/src/dynamo/sglang/tests/test_state_agent.py b/components/src/dynamo/sglang/tests/test_state_agent.py new file mode 100644 index 000000000000..e380bfc98377 --- /dev/null +++ b/components/src/dynamo/sglang/tests/test_state_agent.py @@ -0,0 +1,97 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from types import SimpleNamespace + +import pytest + +from dynamo.common.constants import DisaggregationMode +from dynamo.sglang.state_agent import state_agent_descriptors + +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 fixture(): + source = { + "dp_rank": 0, + "capabilities": ["router_hint", "kvcr_inventory"], + "source_control_endpoint": "tcp://worker:24100", + "cache_owner_id": f"{'01' * 16}:{'02' * 16}/{'00' * 8}/{'03' * 16}", + } + info = { + "kv_hint_sources": [source], + "kv_event_sources": [ + { + "dp_rank": 0, + "endpoint": "tcp://worker:24020", + "topic": "raw", + "block_size": 64, + "format": "dynamo", + } + ], + } + engine = SimpleNamespace( + _scheduler_init_result=SimpleNamespace(scheduler_infos=[info]) + ) + config = SimpleNamespace( + server_args=SimpleNamespace(enable_hierarchical_cache=True, dp_size=1), + dynamo_args=SimpleNamespace( + use_kv_events=True, + kv_state_endpoint=None, + namespace="ns", + component="backend", + endpoint="generate", + ), + serving_mode=DisaggregationMode.AGGREGATED, + ) + return engine, config, source, info + + +def test_descriptors_use_only_reported_sources(): + engine, config, source, _ = fixture() + (result,) = state_agent_descriptors(engine, config) + assert result["cache_owner_id"] == source["cache_owner_id"] + assert ( + result["router_hint_source"]["source_control_endpoint"] == "tcp://worker:24100" + ) + assert result["raw_zmq_endpoint"] == "tcp://worker:24020" + assert result["ingress_protocol"] == "sglang_hicache_v1" + assert result["kv_state_endpoint"] == "ns/backend/generate" + source["capabilities"] = ["router_hint"] + assert state_agent_descriptors(engine, config) == [] + + +@pytest.mark.parametrize( + "change,match", + [ + ("missing_events", "raw event source per rank"), + ("wildcard", "wildcard"), + ("missing_rank", "every managed DP rank"), + ("events_disabled", "KV event emission"), + ("wrong_format", "Dynamo-format"), + ("empty_topic", "engine incarnation"), + ], +) +def test_inventory_fails_closed_for_incomplete_startup_metadata(change, match): + engine, config, _, info = fixture() + if change == "missing_events": + info["kv_event_sources"] = [] + elif change == "wildcard": + info["kv_event_sources"][0]["endpoint"] = "tcp://0.0.0.0:24020" + elif change == "missing_rank": + config.server_args.dp_size = 2 + elif change == "wrong_format": + info["kv_event_sources"][0]["format"] = "default" + elif change == "empty_topic": + info["kv_event_sources"][0]["topic"] = "" + else: + config.dynamo_args.use_kv_events = False + with pytest.raises(ValueError, match=match): + state_agent_descriptors(engine, config) From ddc38b2bb8c6cc1a4bc887cd7dfbe58a811bbb09 Mon Sep 17 00:00:00 2001 From: Ziqi Fan Date: Wed, 28 Oct 2026 13:17:19 -0700 Subject: [PATCH 3/9] fix(sglang): close inventory startup and socket lifecycle gaps Validate router-hint endpoint coverage before building inventory descriptors. Keep native registration alive across Python cancellation and withdraw it after the binding has initialized its owner. Return shared publisher allocations after successful HiCache slot shutdown so attachment churn reuses sockets. Signed-off-by: Ziqi Fan --- components/src/dynamo/sglang/state_agent.py | 28 +++++- .../dynamo/sglang/tests/test_state_agent.py | 54 ++++++++++- .../kv_router/publisher/state_agent_host.rs | 96 ++++++++++++++----- 3 files changed, 146 insertions(+), 32 deletions(-) diff --git a/components/src/dynamo/sglang/state_agent.py b/components/src/dynamo/sglang/state_agent.py index f2e151e49a90..bd1b9d073c21 100644 --- a/components/src/dynamo/sglang/state_agent.py +++ b/components/src/dynamo/sglang/state_agent.py @@ -55,9 +55,9 @@ def state_agent_descriptors(engine, config) -> list[dict]: server_args = resolved_server_args(config.server_args) if ( not server_args.enable_hierarchical_cache - or getattr(server_args, "nnodes", 1) != 1 - or getattr(server_args, "pp_size", 1) != 1 - or getattr(server_args, "tokenizer_worker_num", 1) != 1 + or server_args.nnodes != 1 + or server_args.pp_size != 1 + or server_args.tokenizer_worker_num != 1 or not config.dynamo_args.use_kv_events ): raise ValueError( @@ -77,6 +77,10 @@ def state_agent_descriptors(engine, config) -> list[dict]: raise ValueError("KVCR inventory must cover every managed DP rank") infos = engine._scheduler_init_result.scheduler_infos endpoints = kvcr_hint_endpoints(infos, bounds) + if set(endpoints) != {str(rank) for rank in sources}: + raise ValueError( + "KVCR inventory requires a router_hint source for every managed DP rank" + ) raw_sources = {} for info in infos: for source in info.get("kv_event_sources", []): @@ -158,17 +162,33 @@ def state_agent_descriptors(engine, config) -> list[dict]: class StateAgentLifecycle: def __init__(self): self._owner = None + self._cleanup_task = None async def start(self, engine, config, endpoint): descriptors = state_agent_descriptors(engine, config) if not descriptors: return owner = KvStateAttachmentOwner(endpoint, endpoint.connection_id(), descriptors) - await owner.start() + startup = asyncio.ensure_future(owner.start()) + try: + await asyncio.shield(startup) + except asyncio.CancelledError: + self._cleanup_task = asyncio.create_task( + self._close_after_start(owner, startup) + ) + raise self._owner = owner logger.info("Attached %d KVCR HiCache inventory owners", len(descriptors)) + async def _close_after_start(self, owner, startup): + # Let the native binding initialize its OnceCell before withdrawal. + await startup + await owner.close() + async def close(self): + if self._cleanup_task is not None: + await asyncio.wait_for(asyncio.shield(self._cleanup_task), timeout=30.0) + return owner, self._owner = self._owner, None if owner is not None: await asyncio.wait_for(owner.close(), timeout=30.0) diff --git a/components/src/dynamo/sglang/tests/test_state_agent.py b/components/src/dynamo/sglang/tests/test_state_agent.py index e380bfc98377..178aba6a7dc0 100644 --- a/components/src/dynamo/sglang/tests/test_state_agent.py +++ b/components/src/dynamo/sglang/tests/test_state_agent.py @@ -1,12 +1,15 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 +import asyncio from types import SimpleNamespace +from unittest.mock import AsyncMock import pytest from dynamo.common.constants import DisaggregationMode -from dynamo.sglang.state_agent import state_agent_descriptors +from dynamo.sglang import state_agent +from dynamo.sglang.state_agent import StateAgentLifecycle, state_agent_descriptors pytestmark = [ pytest.mark.unit, @@ -41,7 +44,13 @@ def fixture(): _scheduler_init_result=SimpleNamespace(scheduler_infos=[info]) ) config = SimpleNamespace( - server_args=SimpleNamespace(enable_hierarchical_cache=True, dp_size=1), + server_args=SimpleNamespace( + enable_hierarchical_cache=True, + dp_size=1, + nnodes=1, + pp_size=1, + tokenizer_worker_num=1, + ), dynamo_args=SimpleNamespace( use_kv_events=True, kv_state_endpoint=None, @@ -74,19 +83,22 @@ def test_descriptors_use_only_reported_sources(): ("missing_events", "raw event source per rank"), ("wildcard", "wildcard"), ("missing_rank", "every managed DP rank"), + ("missing_router_hint", "router_hint source for every managed DP rank"), ("events_disabled", "KV event emission"), ("wrong_format", "Dynamo-format"), ("empty_topic", "engine incarnation"), ], ) def test_inventory_fails_closed_for_incomplete_startup_metadata(change, match): - engine, config, _, info = fixture() + engine, config, source, info = fixture() if change == "missing_events": info["kv_event_sources"] = [] elif change == "wildcard": info["kv_event_sources"][0]["endpoint"] = "tcp://0.0.0.0:24020" elif change == "missing_rank": config.server_args.dp_size = 2 + elif change == "missing_router_hint": + source["capabilities"] = ["kvcr_inventory"] elif change == "wrong_format": info["kv_event_sources"][0]["format"] = "default" elif change == "empty_topic": @@ -95,3 +107,39 @@ def test_inventory_fails_closed_for_incomplete_startup_metadata(change, match): config.dynamo_args.use_kv_events = False with pytest.raises(ValueError, match=match): state_agent_descriptors(engine, config) + + +@pytest.mark.parametrize("native_finished_first", [False, True]) +def test_cancelled_start_withdraws_registered_owner(monkeypatch, native_finished_first): + async def exercise(): + engine, config, _, _ = fixture() + started = asyncio.Event() + native_result = asyncio.get_running_loop().create_future() + + def start(): + started.set() + return native_result + + async def close(): + assert native_result.done() and not native_result.cancelled() + + owner = SimpleNamespace(start=start, close=AsyncMock(side_effect=close)) + monkeypatch.setattr(state_agent, "KvStateAttachmentOwner", lambda *args: owner) + lifecycle = StateAgentLifecycle() + endpoint = SimpleNamespace(connection_id=lambda: 17) + task = asyncio.create_task(lifecycle.start(engine, config, endpoint)) + await asyncio.wait_for(started.wait(), timeout=1) + if native_finished_first: + native_result.set_result(None) + task.cancel() + with pytest.raises(asyncio.CancelledError): + await task + assert not native_result.cancelled() + if not native_finished_first: + owner.close.assert_not_awaited() + native_result.set_result(None) + await asyncio.wait_for(lifecycle.close(), timeout=1) + owner.close.assert_awaited_once() + assert lifecycle._owner is None + + asyncio.run(exercise()) diff --git a/lib/llm/src/kv_router/publisher/state_agent_host.rs b/lib/llm/src/kv_router/publisher/state_agent_host.rs index 41f8bfc105b6..8167e42487d3 100644 --- a/lib/llm/src/kv_router/publisher/state_agent_host.rs +++ b/lib/llm/src/kv_router/publisher/state_agent_host.rs @@ -60,6 +60,7 @@ struct HostSockets { slots_per_socket: usize, /// Brokered and NATS transports have no per-publisher socket to share. share_publishers: bool, + // Entries stay in place so hosted slots can retain their publisher index. publishers: Vec<(SharedZmqPublisherSocket, usize)>, raw_pool: Option, } @@ -82,8 +83,8 @@ impl HostSockets { }) } - fn for_new_slot(&mut self) -> Result { - let event_socket = if self.share_publishers { + fn for_new_slot(&mut self) -> Result<(KvStateAgentSockets, Option)> { + let publisher_index = if self.share_publishers { let index = match self .publishers .iter() @@ -97,27 +98,28 @@ impl HostSockets { self.publishers.len() - 1 } }; - let (socket, slots) = &mut self.publishers[index]; - *slots += 1; - Some(socket.clone()) + self.publishers[index].1 += 1; + Some(index) } else { None }; - Ok(KvStateAgentSockets { - event_socket, - raw_pool: self.raw_pool.clone(), - }) + Ok(( + KvStateAgentSockets { + event_socket: publisher_index.map(|index| self.publishers[index].0.clone()), + raw_pool: self.raw_pool.clone(), + }, + publisher_index, + )) + } + fn release(&mut self, publisher_index: Option) { + if let Some(index) = publisher_index { + self.publishers[index].1 -= 1; + } } - /// Returns the share of a slot that failed to start. - fn release(&mut self, sockets: &KvStateAgentSockets) { - if let Some(socket) = &sockets.event_socket - && let Some((_, slots)) = self - .publishers - .iter_mut() - .find(|(shared, _)| shared.endpoint() == socket.endpoint()) - { - *slots -= 1; + fn remove_slot(&mut self, owner: CacheOwnerId, slots: &mut HashMap) { + if let Some(slot) = slots.remove(&owner) { + self.release(slot.publisher_index); } } } @@ -134,6 +136,7 @@ struct HostedSlot { producer_instance: Option, lifecycle: SlotLifecycle, agent: Option>, + publisher_index: Option, global_dp_rank: u32, kv_state_endpoint: dynamo_runtime::protocols::EndpointId, indexer_domain_id: dynamo_kv_router::identity::IndexerDomainId, @@ -629,7 +632,7 @@ async fn reconcile_owner( slot.agent = None; } if slot.agent.is_none() { - slots.remove(&owner); + sockets.remove_slot(owner, slots); } } return Ok(()); @@ -664,6 +667,7 @@ async fn reconcile_owner( producer_instance: Some(intent.producer_instance.clone()), lifecycle: SlotLifecycle::Failed, agent: None, + publisher_index: None, global_dp_rank: intent.worker.dp_rank, kv_state_endpoint: intent.kv_state_endpoint.clone(), indexer_domain_id: intent.indexer_domain_id, @@ -671,7 +675,7 @@ async fn reconcile_owner( ingress_protocol: intent.ingress_protocol, }; slots.insert(owner, slot_shape); - let slot_sockets = sockets.for_new_slot()?; + let (slot_sockets, publisher_index) = sockets.for_new_slot()?; let agent = match KvStateAgent::start(KvStateAgentConfig { endpoint: component.endpoint(slot_endpoint_name(owner)), kv_state_endpoint: intent.kv_state_endpoint.clone(), @@ -692,11 +696,13 @@ async fn reconcile_owner( // It remains capacity-accounted and avoids implicit construction // retries under an unchanged lease-owned request. Its socket share // is returned now; the slot never published on it. - sockets.release(&slot_sockets); + sockets.release(publisher_index); return Err(error).context("failed to construct stable state-agent slot"); } }; - slots.get_mut(&owner).expect("inserted above").agent = Some(agent); + let slot = slots.get_mut(&owner).expect("inserted above"); + slot.agent = Some(agent); + slot.publisher_index = publisher_index; } let slot = slots.get_mut(&owner).expect("slot exists"); @@ -820,11 +826,50 @@ mod tests { raw_pool: None, }; sockets.for_new_slot().unwrap(); - let failed = sockets.for_new_slot().unwrap(); - sockets.release(&failed); - sockets.for_new_slot().unwrap(); + let (_, failed) = sockets.for_new_slot().unwrap(); + sockets.release(failed); + assert_eq!(sockets.for_new_slot().unwrap().1, failed); assert_eq!(sockets.publishers.len(), 1); + assert_eq!(sockets.publishers[0].1, 2); + } + + #[tokio::test] + async fn removed_slots_reuse_their_shared_publisher() { + let mut sockets = HostSockets { + slots_per_socket: 2, + share_publishers: true, + publishers: Vec::new(), + raw_pool: None, + }; + let mut slots = HashMap::new(); + for _ in 0..4 { + for id in [1, 2] { + let owner = owner(id); + slots.insert( + owner, + HostedSlot { + intent_incarnation: None, + producer_instance: None, + lifecycle: SlotLifecycle::Detached, + agent: None, + publisher_index: sockets.for_new_slot().unwrap().1, + global_dp_rank: 0, + kv_state_endpoint: EndpointId::from("ns.backend.generate"), + indexer_domain_id: owner.pool().indexer_domain(), + kv_block_size: 64, + ingress_protocol: KvStateIngressProtocol::SglangHiCacheV1, + }, + ); + } + assert_eq!(sockets.publishers.len(), 1); + assert_eq!(sockets.publishers[0].1, 2); + sockets.remove_slot(owner(1), &mut slots); + assert_eq!(sockets.publishers[0].1, 1); + sockets.remove_slot(owner(2), &mut slots); + assert!(slots.is_empty()); + assert_eq!(sockets.publishers[0].1, 0); + } } fn instance(component: &str, instance_id: u64) -> Instance { @@ -893,6 +938,7 @@ mod tests { producer_instance: None, lifecycle: SlotLifecycle::Detached, agent: None, + publisher_index: None, global_dp_rank: original.worker.dp_rank, kv_state_endpoint: original.kv_state_endpoint.clone(), indexer_domain_id: original.indexer_domain_id, From 09fb8df03960b6444f53cb2e9a7720b1e92f5e11 Mon Sep 17 00:00:00 2001 From: Ziqi Fan Date: Wed, 28 Oct 2026 13:25:51 -0700 Subject: [PATCH 4/9] fix(sglang): annotate HiCache inventory metadata mappings Signed-off-by: Ziqi Fan --- components/src/dynamo/sglang/state_agent.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/components/src/dynamo/sglang/state_agent.py b/components/src/dynamo/sglang/state_agent.py index bd1b9d073c21..43a3a31107b1 100644 --- a/components/src/dynamo/sglang/state_agent.py +++ b/components/src/dynamo/sglang/state_agent.py @@ -21,12 +21,12 @@ _OWNER_PATTERN = re.compile(r"^[0-9a-f]{32}:[0-9a-f]{32}/[0-9a-f]{16}/[0-9a-f]{32}$") -def inventory_sources(engine) -> dict[int, dict]: +def inventory_sources(engine) -> dict[int, Mapping[str, object]]: """Read capabilities from scheduler metadata, never backend configuration.""" initialization = getattr(engine, "_scheduler_init_result", None) if initialization is None: return {} - sources = {} + sources: dict[int, Mapping[str, object]] = {} for info in initialization.scheduler_infos: hints = info.get("kv_hint_sources", []) if not isinstance(hints, list): @@ -81,7 +81,7 @@ def state_agent_descriptors(engine, config) -> list[dict]: raise ValueError( "KVCR inventory requires a router_hint source for every managed DP rank" ) - raw_sources = {} + raw_sources: dict[int, Mapping[str, object]] = {} for info in infos: for source in info.get("kv_event_sources", []): rank = source.get("dp_rank") From efb1bc335ad3d9874e750572354484779e3e9cf3 Mon Sep 17 00:00:00 2001 From: Ziqi Fan Date: Wed, 28 Oct 2026 15:47:53 -0700 Subject: [PATCH 5/9] fix(helpers-prep,non_mechanical_provable): normalize KV attachment helpers Prepare shared cleanup and endpoint validation interfaces in place. Preserve cancellation and bound shutdown waits; reject loopback aliases without DNS lookup. Seed common logging before relocation. Signed-off-by: Ziqi Fan --- .../src/dynamo/common/kv_state_agent.py | 5 +-- components/src/dynamo/vllm/state_agent.py | 33 ++++++++++++------- .../src/dynamo/vllm/tests/test_state_agent.py | 9 +++-- 3 files changed, 31 insertions(+), 16 deletions(-) diff --git a/components/src/dynamo/common/kv_state_agent.py b/components/src/dynamo/common/kv_state_agent.py index 37e5c873bf99..4d2f09b2915b 100644 --- a/components/src/dynamo/common/kv_state_agent.py +++ b/components/src/dynamo/common/kv_state_agent.py @@ -4,10 +4,11 @@ """Attachment helpers shared by KV state-agent workers.""" import hashlib +import logging from collections.abc import Iterable -# Fixed width, so no worker's topic is a prefix of another's: SUB subscriptions -# are prefix filters, and the state agent routes pooled raw streams by topic. +logger = logging.getLogger(__name__) + _RAW_TOPIC_HEX_DIGITS = 16 diff --git a/components/src/dynamo/vllm/state_agent.py b/components/src/dynamo/vllm/state_agent.py index 570d2a11aaff..c749f7e07bad 100644 --- a/components/src/dynamo/vllm/state_agent.py +++ b/components/src/dynamo/vllm/state_agent.py @@ -7,6 +7,8 @@ import ipaddress import logging import re +import socket +from collections.abc import Awaitable from dataclasses import dataclass from typing import TYPE_CHECKING, Any, Mapping @@ -24,7 +26,7 @@ logger = logging.getLogger(__name__) -_CACHE_OWNER_PATTERN = re.compile( +CACHE_OWNER_PATTERN = re.compile( r"^(?P[0-9a-f]{32}:[0-9a-f]{32})/[0-9a-f]{16}/[0-9a-f]{32}$" ) _STATE_AGENT_CLOSE_TIMEOUT_SECS = 30.0 @@ -55,7 +57,7 @@ async def install(self, owner: KvStateAttachmentOwner) -> None: except BaseException: # A started owner already holds discovery registrations. Its Rust # cleanup task survives cancellation of this Python await. - await asyncio.shield(_close_attachment_owner(owner)) + await asyncio.shield(bounded_close(owner.close())) raise async def close(self) -> None: @@ -63,12 +65,13 @@ async def close(self) -> None: self._closed = True owner, self._owner = self._owner, None if owner is not None: - await _close_attachment_owner(owner) + await bounded_close(owner.close()) -async def _close_attachment_owner(owner: KvStateAttachmentOwner) -> None: +async def bounded_close(cleanup: Awaitable[None]) -> None: + """Bound cleanup without replacing a worker's error or cancellation on timeout.""" try: - await asyncio.wait_for(owner.close(), timeout=_STATE_AGENT_CLOSE_TIMEOUT_SECS) + await asyncio.wait_for(cleanup, timeout=_STATE_AGENT_CLOSE_TIMEOUT_SECS) except asyncio.TimeoutError: logger.warning( "KV state attachment cleanup timed out after %.0fs; continuing worker shutdown", @@ -96,7 +99,7 @@ def state_agent_settings_from_engine_args( host = section.get("raw_advertise_host") if not isinstance(host, str) or not host: raise ValueError("dynamo_state_agent.raw_advertise_host is required") - _validate_advertise_host(host) + validate_advertise_host(host) raw_owners = section.get("cache_owner_ids") if not isinstance(raw_owners, dict) or not raw_owners: @@ -111,7 +114,7 @@ def state_agent_settings_from_engine_args( ) from error if str(rank) != str(raw_rank): raise ValueError(f"non-canonical global DP rank key: {raw_rank!r}") - if not isinstance(owner, str) or _CACHE_OWNER_PATTERN.fullmatch(owner) is None: + if not isinstance(owner, str) or CACHE_OWNER_PATTERN.fullmatch(owner) is None: raise ValueError(f"invalid canonical CacheOwnerId for global rank {rank}") owners[rank] = owner if len(set(owners.values())) != len(owners): @@ -252,8 +255,9 @@ def _advertised_raw_endpoint(base: str, rank: int, host: str) -> str: return f"tcp://{host}:{port}" -def _validate_advertise_host(host: str) -> None: - if host in {"*", "localhost", "0.0.0.0", "::", "[::]"}: +def validate_advertise_host(host: str) -> None: + """Require a connectable hostname, IPv4 address, or bracketed IPv6 literal.""" + if host.rstrip(".").lower() in {"*", "localhost", "0.0.0.0", "::", "[::]"}: raise ValueError("raw_advertise_host must be a non-loopback connectable host") if "://" in host or "/" in host: raise ValueError("raw_advertise_host must not include a scheme, port, or path") @@ -278,7 +282,14 @@ def _validate_advertise_host(host: str) -> None: try: address = ipaddress.ip_address(literal) except ValueError: - return - + # POSIX accepts abbreviated/integer IPv4 literals (e.g.127.1). + # Check them without resolving arbitrary DNS names at startup. + try: + address = ipaddress.IPv4Address(socket.inet_aton(literal)) + except OSError: + return + + if address.version == 6 and address.ipv4_mapped is not None: + address = address.ipv4_mapped if address.is_loopback or address.is_unspecified: raise ValueError("raw_advertise_host must be a non-loopback connectable host") diff --git a/components/src/dynamo/vllm/tests/test_state_agent.py b/components/src/dynamo/vllm/tests/test_state_agent.py index 924eaa334a9e..0d2d083d86c9 100644 --- a/components/src/dynamo/vllm/tests/test_state_agent.py +++ b/components/src/dynamo/vllm/tests/test_state_agent.py @@ -10,6 +10,7 @@ pytest.importorskip("vllm.config") +from dynamo.vllm import state_agent as kv_state_agent # noqa: E402 # isort: skip from dynamo.common.kv_state_agent import worker_raw_topic # noqa: E402 from dynamo.vllm import state_agent # noqa: E402 from dynamo.vllm.cache_info import DYNAMO_KV_EVENT_BLOCK_SIZE_KEY # noqa: E402 @@ -17,6 +18,7 @@ pytestmark = [ pytest.mark.unit, pytest.mark.vllm, + pytest.mark.router, pytest.mark.gpu_0, pytest.mark.pre_merge, pytest.mark.usefixtures("vllm_cpu_platform_when_no_accelerator"), @@ -102,7 +104,7 @@ async def test_closed_lifecycle_bounds_late_owner_cleanup(monkeypatch): never_finishes = asyncio.Event() lifecycle = state_agent.StateAgentLifecycle() owner = SimpleNamespace(close=AsyncMock(side_effect=never_finishes.wait)) - monkeypatch.setattr(state_agent, "_STATE_AGENT_CLOSE_TIMEOUT_SECS", 0.01) + monkeypatch.setattr(kv_state_agent, "_STATE_AGENT_CLOSE_TIMEOUT_SECS", 0.01) await lifecycle.close() with pytest.raises(RuntimeError, match="closed"): @@ -112,16 +114,17 @@ async def test_closed_lifecycle_bounds_late_owner_cleanup(monkeypatch): @pytest.mark.asyncio -async def test_lifecycle_close_has_a_bounded_wait(monkeypatch): +async def test_lifecycle_close_has_a_bounded_wait(monkeypatch, caplog): never_finishes = asyncio.Event() lifecycle = state_agent.StateAgentLifecycle() owner = SimpleNamespace(close=AsyncMock(side_effect=never_finishes.wait)) - monkeypatch.setattr(state_agent, "_STATE_AGENT_CLOSE_TIMEOUT_SECS", 0.01) + monkeypatch.setattr(kv_state_agent, "_STATE_AGENT_CLOSE_TIMEOUT_SECS", 0.01) await lifecycle.install(owner) await lifecycle.close() owner.close.assert_awaited_once() + assert "cleanup timed out" in caplog.text def test_cli_validates_state_agent_after_vllm_engine_config(): From b22d1a7ac980c0f34d0548303bfacac35820a03e Mon Sep 17 00:00:00 2001 From: Ziqi Fan Date: Wed, 28 Oct 2026 15:51:15 -0700 Subject: [PATCH 6/9] fix(helpers-move,mechanical_provable): share KV attachment helpers Relocate the prepared validator, owner pattern and bounded cleanup verbatim into common code. Repoint vLLM imports and test timeout configuration without changing their behavior. Signed-off-by: Ziqi Fan --- .../src/dynamo/common/kv_state_agent.py | 63 ++++++++++++++++- components/src/dynamo/vllm/state_agent.py | 67 ++----------------- .../src/dynamo/vllm/tests/test_state_agent.py | 2 +- 3 files changed, 69 insertions(+), 63 deletions(-) diff --git a/components/src/dynamo/common/kv_state_agent.py b/components/src/dynamo/common/kv_state_agent.py index 4d2f09b2915b..f42a422e844e 100644 --- a/components/src/dynamo/common/kv_state_agent.py +++ b/components/src/dynamo/common/kv_state_agent.py @@ -3,15 +3,76 @@ """Attachment helpers shared by KV state-agent workers.""" +import asyncio import hashlib +import ipaddress import logging -from collections.abc import Iterable +import re +import socket +from collections.abc import Awaitable, Iterable logger = logging.getLogger(__name__) +CACHE_OWNER_PATTERN = re.compile( + r"^(?P[0-9a-f]{32}:[0-9a-f]{32})/[0-9a-f]{16}/[0-9a-f]{32}$" +) + +_STATE_AGENT_CLOSE_TIMEOUT_SECS = 30.0 + _RAW_TOPIC_HEX_DIGITS = 16 +async def bounded_close(cleanup: Awaitable[None]) -> None: + """Bound cleanup without replacing a worker's error or cancellation on timeout.""" + try: + await asyncio.wait_for(cleanup, timeout=_STATE_AGENT_CLOSE_TIMEOUT_SECS) + except asyncio.TimeoutError: + logger.warning( + "KV state attachment cleanup timed out after %.0fs; continuing worker shutdown", + _STATE_AGENT_CLOSE_TIMEOUT_SECS, + ) + + +def validate_advertise_host(host: str) -> None: + """Require a connectable hostname, IPv4 address, or bracketed IPv6 literal.""" + if host.rstrip(".").lower() in {"*", "localhost", "0.0.0.0", "::", "[::]"}: + raise ValueError("raw_advertise_host must be a non-loopback connectable host") + if "://" in host or "/" in host: + raise ValueError("raw_advertise_host must not include a scheme, port, or path") + if host.startswith("["): + if not host.endswith("]"): + raise ValueError("raw_advertise_host has mismatched IPv6 brackets") + literal = host[1:-1] + try: + address = ipaddress.ip_address(literal) + except ValueError as error: + raise ValueError( + "bracketed raw_advertise_host must be an IPv6 literal" + ) from error + if address.version != 6: + raise ValueError("only IPv6 literals may use brackets") + else: + if "[" in host or "]" in host or ":" in host or any(c.isspace() for c in host): + raise ValueError( + "raw_advertise_host must be a hostname or bracketed IPv6 literal without a port" + ) + literal = host + try: + address = ipaddress.ip_address(literal) + except ValueError: + # POSIX accepts abbreviated/integer IPv4 literals (e.g.127.1). + # Check them without resolving arbitrary DNS names at startup. + try: + address = ipaddress.IPv4Address(socket.inet_aton(literal)) + except OSError: + return + + if address.version == 6 and address.ipv4_mapped is not None: + address = address.ipv4_mapped + if address.is_loopback or address.is_unspecified: + raise ValueError("raw_advertise_host must be a non-loopback connectable host") + + def worker_raw_topic(cache_owner_ids: Iterable[str]) -> str: """Raw KV event topic of one worker, published by every DP rank it owns. diff --git a/components/src/dynamo/vllm/state_agent.py b/components/src/dynamo/vllm/state_agent.py index c749f7e07bad..30bdc0437505 100644 --- a/components/src/dynamo/vllm/state_agent.py +++ b/components/src/dynamo/vllm/state_agent.py @@ -4,18 +4,19 @@ """Opt-in vLLM attachment configuration for the external KV state agent.""" import asyncio -import ipaddress import logging -import re -import socket -from collections.abc import Awaitable from dataclasses import dataclass from typing import TYPE_CHECKING, Any, Mapping if TYPE_CHECKING: from vllm.config import VllmConfig -from dynamo.common.kv_state_agent import worker_raw_topic +from dynamo.common.kv_state_agent import ( + CACHE_OWNER_PATTERN, + bounded_close, + validate_advertise_host, + worker_raw_topic, +) from dynamo.llm import KvStateAttachmentOwner, WorkerType from dynamo.runtime import Endpoint @@ -26,11 +27,6 @@ logger = logging.getLogger(__name__) -CACHE_OWNER_PATTERN = re.compile( - r"^(?P[0-9a-f]{32}:[0-9a-f]{32})/[0-9a-f]{16}/[0-9a-f]{32}$" -) -_STATE_AGENT_CLOSE_TIMEOUT_SECS = 30.0 - @dataclass(frozen=True) class StateAgentSettings: @@ -68,17 +64,6 @@ async def close(self) -> None: await bounded_close(owner.close()) -async def bounded_close(cleanup: Awaitable[None]) -> None: - """Bound cleanup without replacing a worker's error or cancellation on timeout.""" - try: - await asyncio.wait_for(cleanup, timeout=_STATE_AGENT_CLOSE_TIMEOUT_SECS) - except asyncio.TimeoutError: - logger.warning( - "KV state attachment cleanup timed out after %.0fs; continuing worker shutdown", - _STATE_AGENT_CLOSE_TIMEOUT_SECS, - ) - - def state_agent_settings(config: Any) -> StateAgentSettings | None: return state_agent_settings_from_engine_args(config.engine_args) @@ -253,43 +238,3 @@ def _advertised_raw_endpoint(base: str, rank: int, host: str) -> str: if not port.isdigit(): raise ValueError(f"vLLM KV event endpoint has no numeric port: {endpoint!r}") return f"tcp://{host}:{port}" - - -def validate_advertise_host(host: str) -> None: - """Require a connectable hostname, IPv4 address, or bracketed IPv6 literal.""" - if host.rstrip(".").lower() in {"*", "localhost", "0.0.0.0", "::", "[::]"}: - raise ValueError("raw_advertise_host must be a non-loopback connectable host") - if "://" in host or "/" in host: - raise ValueError("raw_advertise_host must not include a scheme, port, or path") - if host.startswith("["): - if not host.endswith("]"): - raise ValueError("raw_advertise_host has mismatched IPv6 brackets") - literal = host[1:-1] - try: - address = ipaddress.ip_address(literal) - except ValueError as error: - raise ValueError( - "bracketed raw_advertise_host must be an IPv6 literal" - ) from error - if address.version != 6: - raise ValueError("only IPv6 literals may use brackets") - else: - if "[" in host or "]" in host or ":" in host or any(c.isspace() for c in host): - raise ValueError( - "raw_advertise_host must be a hostname or bracketed IPv6 literal without a port" - ) - literal = host - try: - address = ipaddress.ip_address(literal) - except ValueError: - # POSIX accepts abbreviated/integer IPv4 literals (e.g.127.1). - # Check them without resolving arbitrary DNS names at startup. - try: - address = ipaddress.IPv4Address(socket.inet_aton(literal)) - except OSError: - return - - if address.version == 6 and address.ipv4_mapped is not None: - address = address.ipv4_mapped - if address.is_loopback or address.is_unspecified: - raise ValueError("raw_advertise_host must be a non-loopback connectable host") diff --git a/components/src/dynamo/vllm/tests/test_state_agent.py b/components/src/dynamo/vllm/tests/test_state_agent.py index 0d2d083d86c9..6780d2ce11b3 100644 --- a/components/src/dynamo/vllm/tests/test_state_agent.py +++ b/components/src/dynamo/vllm/tests/test_state_agent.py @@ -10,7 +10,7 @@ pytest.importorskip("vllm.config") -from dynamo.vllm import state_agent as kv_state_agent # noqa: E402 # isort: skip +from dynamo.common import kv_state_agent # noqa: E402 from dynamo.common.kv_state_agent import worker_raw_topic # noqa: E402 from dynamo.vllm import state_agent # noqa: E402 from dynamo.vllm.cache_info import DYNAMO_KV_EVENT_BLOCK_SIZE_KEY # noqa: E402 From 5ba144a9a7c05f20ee70ae476941619de130bc90 Mon Sep 17 00:00:00 2001 From: Ziqi Fan Date: Wed, 28 Oct 2026 15:51:15 -0700 Subject: [PATCH 7/9] fix(hicache,non_mechanical_provable): address inventory lifecycle review Release failed engine-local slots without reclaiming durable Guard slots. Fence both event domains after a terminal source gap, share SGLang cleanup/endpoint checks, centralize lifetime policy, and retain all seven validator invariants in table-driven coverage. Validated old-code failures and fixed Python/native regressions, followed by all eight TP1/TP2 Python/Rust Qwen3-0.6B E2E cases on the retained two-B200 nScale allocation, with non-loopback endpoints and independently audited byte-level transfers. No lossless replay or performance claim. Signed-off-by: Ziqi Fan --- .../src/dynamo/common/kv_state_agent.py | 3 +- components/src/dynamo/sglang/state_agent.py | 22 +-- .../dynamo/sglang/tests/test_state_agent.py | 138 ++++++++++--- lib/llm/src/discovery/kv_state_agent.rs | 4 + .../src/kv_router/publisher/state_agent.rs | 94 +++++---- .../kv_router/publisher/state_agent_host.rs | 187 ++++++++++++++++-- 6 files changed, 363 insertions(+), 85 deletions(-) diff --git a/components/src/dynamo/common/kv_state_agent.py b/components/src/dynamo/common/kv_state_agent.py index f42a422e844e..a6e5d12fa02e 100644 --- a/components/src/dynamo/common/kv_state_agent.py +++ b/components/src/dynamo/common/kv_state_agent.py @@ -16,9 +16,10 @@ CACHE_OWNER_PATTERN = re.compile( r"^(?P[0-9a-f]{32}:[0-9a-f]{32})/[0-9a-f]{16}/[0-9a-f]{32}$" ) - _STATE_AGENT_CLOSE_TIMEOUT_SECS = 30.0 +# Fixed width, so no worker's topic is a prefix of another's: SUB subscriptions +# are prefix filters, and the state agent routes pooled raw streams by topic. _RAW_TOPIC_HEX_DIGITS = 16 diff --git a/components/src/dynamo/sglang/state_agent.py b/components/src/dynamo/sglang/state_agent.py index 43a3a31107b1..bc212d9fd51c 100644 --- a/components/src/dynamo/sglang/state_agent.py +++ b/components/src/dynamo/sglang/state_agent.py @@ -5,20 +5,22 @@ from __future__ import annotations import asyncio -import ipaddress import logging -import re from collections.abc import Mapping from urllib.parse import urlsplit from dynamo.common.constants import DisaggregationMode +from dynamo.common.kv_state_agent import ( + CACHE_OWNER_PATTERN, + bounded_close, + validate_advertise_host, +) from dynamo.llm import KvStateAttachmentOwner from dynamo.sglang._compat import resolved_server_args from dynamo.sglang.capacity import model_card_dp_rank_bounds from dynamo.sglang.kv_hints import kvcr_hint_endpoints logger = logging.getLogger(__name__) -_OWNER_PATTERN = re.compile(r"^[0-9a-f]{32}:[0-9a-f]{32}/[0-9a-f]{16}/[0-9a-f]{32}$") def inventory_sources(engine) -> dict[int, Mapping[str, object]]: @@ -96,7 +98,7 @@ def state_agent_descriptors(engine, config) -> list[dict]: owners = set() for rank, source in sorted(sources.items()): owner = source.get("cache_owner_id") - if not isinstance(owner, str) or _OWNER_PATTERN.fullmatch(owner) is None: + if not isinstance(owner, str) or CACHE_OWNER_PATTERN.fullmatch(owner) is None: raise ValueError("KVCR inventory source requires a canonical CacheOwnerId") if owner in owners: raise ValueError("KVCR inventory owners must be unique across DP ranks") @@ -122,12 +124,8 @@ def state_agent_descriptors(engine, config) -> list[dict]: or parsed.hostname == "*" ): raise ValueError("Invalid reported KVCR raw event endpoint") - try: - address = ipaddress.ip_address(parsed.hostname) - except ValueError: - address = None - if address is not None and address.is_unspecified: - raise ValueError("KVCR raw event source cannot advertise a wildcard") + host = parsed.hostname + validate_advertise_host(f"[{host}]" if ":" in host else host) block_size = raw.get("block_size") if type(block_size) is not int or block_size <= 0: raise ValueError("KVCR raw event source requires a positive block size") @@ -187,8 +185,8 @@ async def _close_after_start(self, owner, startup): async def close(self): if self._cleanup_task is not None: - await asyncio.wait_for(asyncio.shield(self._cleanup_task), timeout=30.0) + await bounded_close(asyncio.shield(self._cleanup_task)) return owner, self._owner = self._owner, None if owner is not None: - await asyncio.wait_for(owner.close(), timeout=30.0) + await bounded_close(owner.close()) diff --git a/components/src/dynamo/sglang/tests/test_state_agent.py b/components/src/dynamo/sglang/tests/test_state_agent.py index 178aba6a7dc0..77765cd1da40 100644 --- a/components/src/dynamo/sglang/tests/test_state_agent.py +++ b/components/src/dynamo/sglang/tests/test_state_agent.py @@ -7,6 +7,7 @@ import pytest +from dynamo.common import kv_state_agent from dynamo.common.constants import DisaggregationMode from dynamo.sglang import state_agent from dynamo.sglang.state_agent import StateAgentLifecycle, state_agent_descriptors @@ -78,37 +79,128 @@ def test_descriptors_use_only_reported_sources(): @pytest.mark.parametrize( - "change,match", + "target,updates,match", [ - ("missing_events", "raw event source per rank"), - ("wildcard", "wildcard"), - ("missing_rank", "every managed DP rank"), - ("missing_router_hint", "router_hint source for every managed DP rank"), - ("events_disabled", "KV event emission"), - ("wrong_format", "Dynamo-format"), - ("empty_topic", "engine incarnation"), + ("info", {"kv_event_sources": []}, "raw event source per rank"), + ("raw", {"endpoint": "tcp://0.0.0.0:24020"}, "non-loopback"), + ("server_args", {"dp_size": 2}, "every managed DP rank"), + ( + "source", + {"capabilities": ["kvcr_inventory"]}, + "router_hint source for every managed DP rank", + ), + ("dynamo_args", {"use_kv_events": False}, "KV event emission"), + ("raw", {"format": "default"}, "Dynamo-format"), + ("raw", {"topic": ""}, "engine incarnation"), ], ) -def test_inventory_fails_closed_for_incomplete_startup_metadata(change, match): +def test_inventory_fails_closed_for_incomplete_startup_metadata(target, updates, match): engine, config, source, info = fixture() - if change == "missing_events": - info["kv_event_sources"] = [] - elif change == "wildcard": - info["kv_event_sources"][0]["endpoint"] = "tcp://0.0.0.0:24020" - elif change == "missing_rank": - config.server_args.dp_size = 2 - elif change == "missing_router_hint": - source["capabilities"] = ["kvcr_inventory"] - elif change == "wrong_format": - info["kv_event_sources"][0]["format"] = "default" - elif change == "empty_topic": - info["kv_event_sources"][0]["topic"] = "" - else: - config.dynamo_args.use_kv_events = False + targets = { + "info": info, + "source": source, + "raw": info["kv_event_sources"][0], + "server_args": vars(config.server_args), + "dynamo_args": vars(config.dynamo_args), + } + targets[target].update(updates) with pytest.raises(ValueError, match=match): state_agent_descriptors(engine, config) +@pytest.mark.parametrize( + "host", + [ + "localhost", + "127.0.0.1", + "127.1", + "2130706433", + "0", + "[::1]", + "[::ffff:127.0.0.1]", + ], +) +def test_raw_event_endpoint_rejects_loopback_and_unspecified_hosts(host): + engine, config, _, info = fixture() + info["kv_event_sources"][0]["endpoint"] = f"tcp://{host}:24020" + with pytest.raises(ValueError, match="non-loopback"): + state_agent_descriptors(engine, config) + + +def test_raw_event_endpoint_accepts_connectable_ipv6(): + engine, config, _, info = fixture() + endpoint = "tcp://[2001:db8::1]:24020" + info["kv_event_sources"][0]["endpoint"] = endpoint + (descriptor,) = state_agent_descriptors(engine, config) + assert descriptor["raw_zmq_endpoint"] == endpoint + + +@pytest.mark.asyncio +@pytest.mark.parametrize("error_type", [RuntimeError, asyncio.CancelledError]) +async def test_cleanup_timeout_preserves_body_exception( + monkeypatch, caplog, error_type +): + engine, config, _, _ = fixture() + owner = SimpleNamespace( + start=AsyncMock(), close=AsyncMock(side_effect=asyncio.Event().wait) + ) + monkeypatch.setattr(state_agent, "KvStateAttachmentOwner", lambda *args: owner) + monkeypatch.setattr(kv_state_agent, "_STATE_AGENT_CLOSE_TIMEOUT_SECS", 0.01) + lifecycle = StateAgentLifecycle() + await lifecycle.start(engine, config, SimpleNamespace(connection_id=lambda: 17)) + original = error_type("worker body failed") + + with pytest.raises(error_type) as caught: + try: + raise original + finally: + await lifecycle.close() + + assert caught.value is original + assert "cleanup timed out" in caplog.text + await lifecycle.close() + owner.close.assert_awaited_once() + + +@pytest.mark.asyncio +async def test_cancelled_start_cleanup_timeout_preserves_withdrawal( + monkeypatch, caplog +): + engine, config, _, _ = fixture() + started = asyncio.Event() + native_result = asyncio.get_running_loop().create_future() + + def start(): + started.set() + return native_result + + async def close(): + assert native_result.done() and not native_result.cancelled() + + owner = SimpleNamespace(start=start, close=AsyncMock(side_effect=close)) + monkeypatch.setattr(state_agent, "KvStateAttachmentOwner", lambda *args: owner) + monkeypatch.setattr(kv_state_agent, "_STATE_AGENT_CLOSE_TIMEOUT_SECS", 0.01) + lifecycle = StateAgentLifecycle() + task = asyncio.create_task( + lifecycle.start(engine, config, SimpleNamespace(connection_id=lambda: 17)) + ) + await asyncio.wait_for(started.wait(), timeout=1) + task.cancel() + with pytest.raises(asyncio.CancelledError): + try: + await task + finally: + await lifecycle.close() + + assert "cleanup timed out" in caplog.text + assert not native_result.done() + owner.close.assert_not_awaited() + native_result.set_result(None) + await asyncio.wait_for(lifecycle.close(), timeout=1) + owner.close.assert_awaited_once() + assert lifecycle._owner is None + + @pytest.mark.parametrize("native_finished_first", [False, True]) def test_cancelled_start_withdraws_registered_owner(monkeypatch, native_finished_first): async def exercise(): diff --git a/lib/llm/src/discovery/kv_state_agent.rs b/lib/llm/src/discovery/kv_state_agent.rs index 4834d2571f83..912d38de66cc 100644 --- a/lib/llm/src/discovery/kv_state_agent.rs +++ b/lib/llm/src/discovery/kv_state_agent.rs @@ -68,6 +68,10 @@ impl KvStateIngressProtocol { pub fn supports_owned_residency(self) -> bool { matches!(self, Self::VllmResidencyV1 | Self::SglangHiCacheV1) } + + pub fn has_process_local_inventory(self) -> bool { + matches!(self, Self::SglangHiCacheV1) + } } #[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] diff --git a/lib/llm/src/kv_router/publisher/state_agent.rs b/lib/llm/src/kv_router/publisher/state_agent.rs index ca540716cd6a..30e669444adb 100644 --- a/lib/llm/src/kv_router/publisher/state_agent.rs +++ b/lib/llm/src/kv_router/publisher/state_agent.rs @@ -379,7 +379,7 @@ impl Coordinator

{ .context("failed to establish Worker reset barrier before attach"); } }; - if attachment.ingress_protocol == KvStateIngressProtocol::SglangHiCacheV1 { + if attachment.ingress_protocol.has_process_local_inventory() { barrier_cursor = match self .apply_and_publish_reset(attachment.worker, ResidencyDomain::CacheOwner) .await @@ -413,7 +413,7 @@ impl Coordinator

{ anyhow::bail!("stale attachment generation {generation}"); } attachment.ready = false; - if attachment.ingress_protocol == KvStateIngressProtocol::SglangHiCacheV1 { + if attachment.ingress_protocol.has_process_local_inventory() { self.cache_owner_domain_failed = true; } self.publish_status(); @@ -447,7 +447,7 @@ impl Coordinator

{ return Err(error).context("failed to publish final Worker detach barrier"); } }; - if attachment.ingress_protocol == KvStateIngressProtocol::SglangHiCacheV1 { + if attachment.ingress_protocol.has_process_local_inventory() { cursor = self .apply_and_publish_reset(attachment.worker, ResidencyDomain::CacheOwner) .await @@ -468,7 +468,7 @@ impl Coordinator

{ } attachment.ready = false; self.ingress_generation = None; - tracing::warn!(generation, reason, "Quarantining vLLM KV ingress"); + tracing::warn!(generation, reason, "Quarantining KV ingress"); self.publish_status(); } @@ -490,7 +490,9 @@ impl Coordinator

{ anyhow::bail!("oversized ingress envelope"); } - self.observe_source_cursor_lossy(&batch, "live ingress"); + if !self.observe_source_cursor(&batch, "live ingress") { + anyhow::bail!("raw KV source retired after a sequence gap"); + } if let Some(reason) = batch.source_fault { self.fail_source(batch.generation, reason); @@ -522,7 +524,9 @@ impl Coordinator

{ continue; } - self.observe_source_cursor_lossy(&batch, "detach drain"); + if !self.observe_source_cursor(&batch, "detach drain") { + continue; + } if let Some(reason) = batch.source_fault { self.fail_source(generation, reason); continue; @@ -544,10 +548,7 @@ impl Coordinator

{ } } if self.source_cursor.active.is_some() { - tracing::warn!( - generation, - "Detaching after an incomplete raw source batch; CacheOwner remains advisory" - ); + tracing::warn!(generation, "Detaching after an incomplete raw source batch"); if let Some(metrics) = kv_publisher_metrics() { metrics.increment_engines_dropped_events(1); } @@ -555,15 +556,34 @@ impl Coordinator

{ } } - fn observe_source_cursor_lossy(&mut self, batch: &IngressBatch, phase: &'static str) { + fn observe_source_cursor(&mut self, batch: &IngressBatch, phase: &'static str) -> bool { if self.source_cursor.observe(batch) { - return; + return true; } let previous_cursor = self.source_cursor.completed; let missing = previous_cursor .and_then(|previous| batch.source_cursor.checked_sub(previous.checked_add(1)?)) .filter(|missing| *missing > 0) .unwrap_or(1); + if let Some(metrics) = kv_publisher_metrics() { + metrics.increment_engines_dropped_events(missing); + } + if self + .attachment + .as_ref() + .is_some_and(|attachment| attachment.ingress_protocol.has_process_local_inventory()) + { + tracing::error!( + generation = batch.generation, + source_cursor = batch.source_cursor, + ?previous_cursor, + chunk_index = batch.chunk_index, + phase, + "Raw KV source cursor gap; process-local inventory retired until engine restart" + ); + self.fail_source(batch.generation, "process-local inventory cursor gap"); + return false; + } tracing::warn!( generation = batch.generation, source_cursor = batch.source_cursor, @@ -572,15 +592,8 @@ impl Coordinator

{ phase, "Raw KV source cursor gap or regression; continuing with advisory residency" ); - if let Some(metrics) = kv_publisher_metrics() { - metrics.increment_engines_dropped_events(missing); - } - if self.attachment.as_ref().is_some_and(|attachment| { - attachment.ingress_protocol == KvStateIngressProtocol::SglangHiCacheV1 - }) { - self.fail_source(batch.generation, "process-local inventory cursor gap"); - } self.source_cursor.accept_after_loss(batch); + true } fn fail_source(&mut self, generation: u64, reason: &'static str) { @@ -825,8 +838,7 @@ impl Coordinator

{ attachment: self.attachment.as_ref().map(Attachment::status), cache_owner_ready: !self.cache_owner_domain_failed && self.attachment.as_ref().is_none_or(|attachment| { - attachment.ingress_protocol != KvStateIngressProtocol::SglangHiCacheV1 - || attachment.ready + !attachment.ingress_protocol.has_process_local_inventory() || attachment.ready }), outbound_cursor: self.outbound_cursor(), }); @@ -908,7 +920,7 @@ impl KvStateAgent { let status = Arc::new(KvStateAgentStatusCell::new(KvStateAgentStatus { identity: identity.clone(), attachment: None, - cache_owner_ready: config.ingress_protocol != KvStateIngressProtocol::SglangHiCacheV1, + cache_owner_ready: !config.ingress_protocol.has_process_local_inventory(), outbound_cursor: 0, })); @@ -2757,7 +2769,7 @@ mod tests { } #[tokio::test] - async fn sglang_detach_clears_process_local_owner() { + async fn sglang_inventory_is_cleared_on_detach_and_retires_on_gaps() { check_detach_lifetime(KvStateIngressProtocol::SglangHiCacheV1).await; } @@ -2890,21 +2902,31 @@ mod tests { }) .await .unwrap(); + let cursor_before_gap = coordinator.outbound_cursor(); + let result = coordinator + .handle_ingress(IngressBatch { + generation: replacement.generation, + source_cursor: 3, + chunk_index: 0, + final_chunk: true, + events: vec![ + store(worker, ResidencyDomain::Worker, None, 401, 41), + store(worker, ResidencyDomain::CacheOwner, None, 402, 42), + ], + source_fault: None, + cache_owner_fault: None, + }) + .await; if process_local { - coordinator - .handle_ingress(IngressBatch { - generation: replacement.generation, - source_cursor: 3, - chunk_index: 0, - final_chunk: true, - events: vec![], - source_fault: None, - cache_owner_fault: None, - }) - .await - .unwrap(); + assert!(result.is_err()); + assert_eq!(coordinator.outbound_cursor(), cursor_before_gap); assert!(!coordinator.status.load().cache_owner_ready); assert!(!coordinator.status.load().attachment.as_ref().unwrap().ready); + } else { + result.unwrap(); + assert_eq!(coordinator.outbound_cursor(), cursor_before_gap + 2); + assert!(coordinator.status.load().cache_owner_ready); + assert!(coordinator.status.load().attachment.as_ref().unwrap().ready); } coordinator.begin_detach(replacement.generation).unwrap(); coordinator diff --git a/lib/llm/src/kv_router/publisher/state_agent_host.rs b/lib/llm/src/kv_router/publisher/state_agent_host.rs index 8167e42487d3..6c7ec3da458d 100644 --- a/lib/llm/src/kv_router/publisher/state_agent_host.rs +++ b/lib/llm/src/kv_router/publisher/state_agent_host.rs @@ -624,12 +624,17 @@ async fn reconcile_owner( let intent = match matching.map(HashMap::len).unwrap_or(0) { 0 => { if let Some(slot) = slots.get_mut(&owner) { - detach_slot(slot).await?; - if slot.ingress_protocol == KvStateIngressProtocol::SglangHiCacheV1 { - if let Some(agent) = slot.agent.as_ref() { - agent.shutdown().await?; + if slot.ingress_protocol.has_process_local_inventory() { + if let Err(error) = detach_slot(slot).await { + tracing::warn!(%owner, %error, "Failed to detach engine-local KV state-agent slot"); + } + if let Some(agent) = slot.agent.take() + && let Err(error) = agent.shutdown().await + { + tracing::warn!(%owner, %error, "Failed to stop engine-local KV state agent"); } - slot.agent = None; + } else { + detach_slot(slot).await?; } if slot.agent.is_none() { sockets.remove_slot(owner, slots); @@ -650,10 +655,10 @@ async fn reconcile_owner( let intent = intent.expect("one intent was counted"); if !slots.contains_key(&owner) { - // NOTE: Detached agents deliberately retain CacheOwner/index/recovery - // state and count against max_slots. This is a durable-slot resource - // budget, not an active-attachment count. Failed construction remains - // reserved only until its matching intent disappears. + // NOTE: Detached persistent agents retain CacheOwner/index/recovery state + // and count against max_slots. Engine-local slots release their capacity + // when their last intent disappears. Failed construction remains reserved + // only until its matching intent disappears. // TODO(#13044): require explicit decommission/migration authority // before reclaiming a durable owner from a long-lived host. if slots.len() >= max_slots { @@ -814,11 +819,12 @@ mod tests { use dynamo_runtime::component::{Instance, TransportType}; use dynamo_runtime::discovery::DiscoveryInstanceId; use dynamo_runtime::protocols::EndpointId; + use dynamo_runtime::{DistributedRuntime, Runtime, distributed::DistributedConfig}; use super::*; - #[tokio::test] - async fn failed_slot_returns_its_shared_socket_share() { + #[test] + fn failed_slot_returns_its_shared_socket_share() { let mut sockets = HostSockets { slots_per_socket: 2, share_publishers: true, @@ -834,8 +840,8 @@ mod tests { assert_eq!(sockets.publishers[0].1, 2); } - #[tokio::test] - async fn removed_slots_reuse_their_shared_publisher() { + #[test] + fn removed_slots_reuse_their_shared_publisher() { let mut sockets = HostSockets { slots_per_socket: 2, share_publishers: true, @@ -872,6 +878,161 @@ mod tests { } } + async fn test_component() -> Component { + DistributedRuntime::new( + Runtime::from_current().unwrap(), + DistributedConfig::process_local(), + ) + .await + .unwrap() + .namespace("state-agent-host-test") + .unwrap() + .component("host") + .unwrap() + } + + #[test] + fn missing_intent_releases_failed_engine_local_slot_but_retains_durable_slot() { + for (protocol, has_process_local_inventory) in [ + (KvStateIngressProtocol::SglangHiCacheV1, true), + (KvStateIngressProtocol::VllmResidencyV1, false), + ] { + let agent_runtime = tokio::runtime::Builder::new_current_thread() + .enable_all() + .build() + .unwrap(); + let (mut slots, mut sockets, mut request) = agent_runtime.block_on(async { + let component = test_component().await; + let raw = SharedZmqPublisherSocket::bind("raw-kv").unwrap(); + let mut request = intent( + &instance("kv_state_agent", 1), + instance("producer", 11), + owner(4), + WorkerWithDpRank::new(17, 4), + 101, + ); + request.ingress_protocol = protocol; + request.raw_zmq_endpoint = raw.endpoint().to_string(); + let mut slots = HashMap::new(); + let mut sockets = HostSockets { + slots_per_socket: 2, + share_publishers: true, + publishers: Vec::new(), + raw_pool: None, + }; + reconcile_owner( + &component, + request.cache_owner_id, + &HashMap::from([( + request.cache_owner_id, + HashMap::from([(request.intent_incarnation, request.clone())]), + )]), + &mut slots, + 1, + &mut sockets, + &mut 0, + &HostMetrics::new(&component).unwrap(), + ) + .await + .unwrap(); + assert_eq!( + slots[&request.cache_owner_id].lifecycle, + SlotLifecycle::Active + ); + (slots, sockets, request) + }); + let agent = slots[&request.cache_owner_id] + .agent + .as_ref() + .unwrap() + .clone(); + let generation = agent.status().attachment.as_ref().unwrap().generation; + let publisher_index = slots[&request.cache_owner_id].publisher_index; + // Abort the real coordinator and recovery endpoint: detach encounters + // a closed channel, and shutdown encounters the endpoint's JoinError. + drop(agent_runtime); + + let host_runtime = tokio::runtime::Builder::new_current_thread() + .enable_all() + .build() + .unwrap(); + host_runtime.block_on(async { + let component = test_component().await; + let metrics = HostMetrics::new(&component).unwrap(); + let mut rejected = 0; + assert!(agent.detach(generation).await.is_err()); + let result = reconcile_owner( + &component, + request.cache_owner_id, + &HashMap::new(), + &mut slots, + 1, + &mut sockets, + &mut rejected, + &metrics, + ) + .await; + if has_process_local_inventory { + result.unwrap(); + assert!(slots.is_empty()); + assert_eq!(Arc::strong_count(&agent), 1); + assert_eq!(sockets.publishers[0].1, 0); + // Shutdown must already have consumed the aborted endpoint. + agent.shutdown().await.unwrap(); + } else { + assert!(result.is_err()); + assert_eq!( + slots[&request.cache_owner_id].lifecycle, + SlotLifecycle::Failed + ); + assert!(Arc::ptr_eq( + slots[&request.cache_owner_id].agent.as_ref().unwrap(), + &agent + )); + assert_eq!(sockets.publishers[0].1, 1); + } + + let raw = SharedZmqPublisherSocket::bind("replacement-kv").unwrap(); + request.cache_owner_id = owner(5); + request.raw_zmq_endpoint = raw.endpoint().to_string(); + reconcile_owner( + &component, + request.cache_owner_id, + &HashMap::from([( + request.cache_owner_id, + HashMap::from([(request.intent_incarnation, request.clone())]), + )]), + &mut slots, + 1, + &mut sockets, + &mut rejected, + &metrics, + ) + .await + .unwrap(); + assert_eq!(slots.len(), 1); + assert_eq!(sockets.publishers.len(), 1); + assert_eq!(sockets.publishers[0].1, 1); + if has_process_local_inventory { + let replacement = &slots[&request.cache_owner_id]; + assert_eq!(replacement.lifecycle, SlotLifecycle::Active); + assert_eq!(replacement.publisher_index, publisher_index); + assert_eq!(rejected, 0); + replacement + .agent + .as_ref() + .unwrap() + .shutdown() + .await + .unwrap(); + } else { + assert!(!slots.contains_key(&request.cache_owner_id)); + assert_eq!(rejected, 1); + } + }); + } + } + fn instance(component: &str, instance_id: u64) -> Instance { Instance { namespace: "ns".to_string(), From 681d43db8ef9c9311371bf879f727a0b8ecab9c5 Mon Sep 17 00:00:00 2001 From: Ziqi Fan Date: Wed, 28 Oct 2026 16:58:25 -0700 Subject: [PATCH 8/9] test(sglang): skip state agent tests without SGLang Signed-off-by: Ziqi Fan --- components/src/dynamo/sglang/tests/test_state_agent.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/components/src/dynamo/sglang/tests/test_state_agent.py b/components/src/dynamo/sglang/tests/test_state_agent.py index 77765cd1da40..2a0bf71343ac 100644 --- a/components/src/dynamo/sglang/tests/test_state_agent.py +++ b/components/src/dynamo/sglang/tests/test_state_agent.py @@ -7,6 +7,8 @@ import pytest +pytest.importorskip("sglang") + from dynamo.common import kv_state_agent from dynamo.common.constants import DisaggregationMode from dynamo.sglang import state_agent From 25dd503929a4bcfe05f28ce2d272bdd86c072538 Mon Sep 17 00:00:00 2001 From: Ziqi Fan Date: Wed, 28 Oct 2026 17:00:20 -0700 Subject: [PATCH 9/9] test(sglang): align optional import guard with lint Signed-off-by: Ziqi Fan --- .../src/dynamo/sglang/tests/test_state_agent.py | 12 ++++++++---- 1 file changed, 8 insertions(+), 4 deletions(-) diff --git a/components/src/dynamo/sglang/tests/test_state_agent.py b/components/src/dynamo/sglang/tests/test_state_agent.py index 2a0bf71343ac..4f2e43377587 100644 --- a/components/src/dynamo/sglang/tests/test_state_agent.py +++ b/components/src/dynamo/sglang/tests/test_state_agent.py @@ -7,12 +7,16 @@ import pytest -pytest.importorskip("sglang") - from dynamo.common import kv_state_agent from dynamo.common.constants import DisaggregationMode -from dynamo.sglang import state_agent -from dynamo.sglang.state_agent import StateAgentLifecycle, state_agent_descriptors + +pytest.importorskip("sglang") + +from dynamo.sglang import state_agent # noqa: E402 +from dynamo.sglang.state_agent import ( # noqa: E402 + StateAgentLifecycle, + state_agent_descriptors, +) pytestmark = [ pytest.mark.unit,