diff --git a/components/src/dynamo/common/kv_state_agent.py b/components/src/dynamo/common/kv_state_agent.py index 37e5c873bf99..a6e5d12fa02e 100644 --- a/components/src/dynamo/common/kv_state_agent.py +++ b/components/src/dynamo/common/kv_state_agent.py @@ -3,14 +3,77 @@ """Attachment helpers shared by KV state-agent workers.""" +import asyncio import hashlib -from collections.abc import Iterable +import ipaddress +import logging +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 # 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 +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/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..bc212d9fd51c --- /dev/null +++ b/components/src/dynamo/sglang/state_agent.py @@ -0,0 +1,192 @@ +# 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 logging +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__) + + +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: dict[int, Mapping[str, object]] = {} + 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 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( + "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) + 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: dict[int, Mapping[str, object]] = {} + 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 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") + 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") + 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") + 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 + 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) + 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 bounded_close(asyncio.shield(self._cleanup_task)) + return + owner, self._owner = self._owner, None + if owner is not None: + 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 new file mode 100644 index 000000000000..4f2e43377587 --- /dev/null +++ b/components/src/dynamo/sglang/tests/test_state_agent.py @@ -0,0 +1,243 @@ +# 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 import kv_state_agent +from dynamo.common.constants import DisaggregationMode + +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, + 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, + nnodes=1, + pp_size=1, + tokenizer_worker_num=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( + "target,updates,match", + [ + ("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(target, updates, match): + engine, config, source, info = fixture() + 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(): + 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/components/src/dynamo/vllm/state_agent.py b/components/src/dynamo/vllm/state_agent.py index 570d2a11aaff..30bdc0437505 100644 --- a/components/src/dynamo/vllm/state_agent.py +++ b/components/src/dynamo/vllm/state_agent.py @@ -4,16 +4,19 @@ """Opt-in vLLM attachment configuration for the external KV state agent.""" import asyncio -import ipaddress import logging -import re 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 @@ -24,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: @@ -55,7 +53,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,17 +61,7 @@ async def close(self) -> None: self._closed = True owner, self._owner = self._owner, None if owner is not None: - await _close_attachment_owner(owner) - - -async def _close_attachment_owner(owner: KvStateAttachmentOwner) -> None: - try: - await asyncio.wait_for(owner.close(), 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, - ) + await bounded_close(owner.close()) def state_agent_settings(config: Any) -> StateAgentSettings | None: @@ -96,7 +84,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 +99,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): @@ -250,35 +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: - if host 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: - return - - 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..6780d2ce11b3 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.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 @@ -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(): 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 53b136d07d52..f508c974969a 100644 --- a/lib/llm/src/discovery/kv_state_agent.rs +++ b/lib/llm/src/discovery/kv_state_agent.rs @@ -59,6 +59,19 @@ 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) + } + + 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/indexer/recovery/state_agent.rs b/lib/llm/src/kv_router/indexer/recovery/state_agent.rs index 3d19422680ef..8a58d0b61b24 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 d4cef1a32b72..d4428c84767c 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 5a010221ffa6..6268b3008a00 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.has_process_local_inventory() { + 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.has_process_local_inventory() { + 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.has_process_local_inventory() { + 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(); @@ -446,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(); } @@ -468,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); @@ -500,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; @@ -522,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); } @@ -533,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, @@ -550,10 +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); - } self.source_cursor.accept_after_loss(batch); + true } fn fail_source(&mut self, generation: u64, reason: &'static str) { @@ -796,7 +836,10 @@ 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.has_process_local_inventory() || attachment.ready + }), outbound_cursor: self.outbound_cursor(), }); } @@ -846,8 +889,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 +920,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.has_process_local_inventory(), outbound_cursor: 0, })); @@ -1761,9 +1804,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 +2765,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_inventory_is_cleared_on_detach_and_retires_on_gaps() { + 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 +2819,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 +2844,7 @@ mod tests { .complete_detach(attachment.generation) .await .unwrap(), - 4 + if process_local { 5 } else { 4 } ); let events = output.lock().unwrap().clone(); @@ -2804,12 +2857,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 +2884,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 +2902,41 @@ 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 { + 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 .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 1dbe53b340a4..d4b08cb61992 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, @@ -582,7 +585,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() { @@ -621,9 +624,20 @@ 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.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"); + } + } else { + detach_slot(slot).await?; + } if slot.agent.is_none() { - slots.remove(&owner); + sockets.remove_slot(owner, slots); } } return Ok(()); @@ -641,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 { @@ -658,6 +672,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, @@ -665,7 +680,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(), @@ -686,11 +701,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"); @@ -802,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, @@ -814,11 +832,205 @@ 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); + } + + #[test] + 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); + } + } + + 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 { @@ -887,6 +1099,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,