diff --git a/components/src/dynamo/sglang/register.py b/components/src/dynamo/sglang/register.py index 794604f3135f..275854c097e1 100644 --- a/components/src/dynamo/sglang/register.py +++ b/components/src/dynamo/sglang/register.py @@ -8,7 +8,6 @@ from typing import Any, List, Optional import sglang as sgl -from sglang.srt.environ import envs from sglang.srt.server_args import ServerArgs from sglang.srt.speculative.spec_info import SpeculativeAlgorithm @@ -232,33 +231,6 @@ def _get_mooncake_runtime_data(server_args: ServerArgs) -> Optional[dict[str, An getattr(server_args, "hicache_storage_backend_extra_config", None) ) - try: - from sglang.srt.mem_cache.storage.mooncake_store.mooncake_store import ( - MooncakeStoreConfig, - ) - except ImportError as e: - logging.warning(f"MooncakeStoreConfig import unavailable: {e}") - return None - - # Graceful degradation: Mooncake runtime metadata is optional. If config - # resolution fails for any reason (file not found, malformed env vars, - # upstream API change), skip publishing the metadata rather than crashing - # the worker -- the worker still serves requests, just without HiCache - # router hints. Broad catch is intentional per python-guidelines.md. - try: - if extra_config and ( - extra_config.get("master_server_address") is not None - or extra_config.get("client_server_address") is not None - ): - mooncake_config = MooncakeStoreConfig.load_from_extra_config(extra_config) - elif envs.SGLANG_HICACHE_MOONCAKE_CONFIG_PATH.is_set(): - mooncake_config = MooncakeStoreConfig.from_file() - else: - mooncake_config = MooncakeStoreConfig.load_from_env() - except Exception as e: - logging.warning(f"Failed to resolve Mooncake config for runtime metadata: {e}") - return None - tp_size = int(getattr(server_args, "tp_size", 1) or 1) pp_size = int(getattr(server_args, "pp_size", 1) or 1) @@ -296,10 +268,6 @@ def _get_mooncake_runtime_data(server_args: ServerArgs) -> Optional[dict[str, An if not isinstance(extra_backend_tag, str) or not extra_backend_tag: extra_backend_tag = None - master_server_address = getattr(mooncake_config, "master_server_address", None) - if not isinstance(master_server_address, str) or not master_server_address: - master_server_address = None - return { "backend": "mooncake", "page_size": int(getattr(server_args, "page_size", 1) or 1), @@ -310,10 +278,7 @@ def _get_mooncake_runtime_data(server_args: ServerArgs) -> Optional[dict[str, An "tp_lcm_size": tp_lcm_size, "should_split_heads": should_split_heads, "extra_backend_tag": extra_backend_tag, - "master_server_address": master_server_address, - "master_metrics_port": int( - getattr(mooncake_config, "master_metrics_port", 9003) - ), + "kv_events_endpoint": os.getenv("DYN_MOONCAKE_KV_EVENTS_ENDPOINT") or None, } diff --git a/docs/fern/pages/developer-guide/knowledge-base/modular-components/backends/sglang/hicache.md b/docs/fern/pages/developer-guide/knowledge-base/modular-components/backends/sglang/hicache.md index 35da2be0a3fb..9956e8c398a1 100644 --- a/docs/fern/pages/developer-guide/knowledge-base/modular-components/backends/sglang/hicache.md +++ b/docs/fern/pages/developer-guide/knowledge-base/modular-components/backends/sglang/hicache.md @@ -17,7 +17,7 @@ SGLang HiCache extends RadixAttention with a multi-tier KV cache that transparen What Dynamo adds on top of HiCache: - **Tier-aware routing.** The KV router tracks which cache tier each block lives on (GPU / Host / External) and uses that when scoring candidate workers — not just device overlap. -- **Shared-pool awareness.** When an external backend such as Mooncake is configured, the router queries the shared pool in parallel with its own indexer so it can discount prefill cost for blocks any worker can fetch, not just blocks the candidate holds locally. +- **Shared-pool awareness.** When an external backend such as Mooncake is configured, the router tracks Mooncake object events so it can discount prefill cost for blocks any worker can fetch, not just blocks the candidate holds locally. If you are running a single worker with HiCache and no shared pool, no Dynamo-side configuration is required — the worker reports KV events to the router as usual. @@ -84,15 +84,17 @@ flowchart LR Worker -- "KV events (store/remove + medium)" --> Router Worker -- "writes pages" --> Mooncake - Router -- "batch_query on each request" --> Mooncake + Mooncake -- "object events (stored/removed)" --> Router ``` -On every request the router runs two lookups in parallel: +The router maintains two indexes: - Its own radix tree, built from worker KV events (per-tier). -- A batch query to the Mooncake master for blocks reachable from the shared pool. +- A set of Mooncake object keys and verified logical groups, built from the Mooncake master's event stream. -If the shared-pool query fails, the router falls back to indexer-only scoring and logs a warning. The request still succeeds. +Request routing checks both indexes locally. When an event includes SGLang's `group_id`, the first lookup verifies every physical object in the group and caches the result. Later requests use one group lookup per page; any event for that group invalidates the cached result. Events without `group_id` use exact object-key checks. + +If the event stream starts after Mooncake already contains objects, the shared-pool index starts empty and learns about subsequent events. A sequence gap clears the shared-pool index to avoid stale hits. ### Scoring @@ -143,7 +145,8 @@ Earlier SGLang versions do not emit `medium=CPU_PINNED` for Host-tier residency, You also need: - Dynamo router started with `--shared-cache-type hicache` (see [Configuration](#configuration)). -- A Mooncake master reachable from the Dynamo frontend host. Worker-side Mooncake config (master address, page size, TP/PP layout, split-head layout) is published automatically via each worker's registration metadata when the worker is started with `--hicache-storage-backend mooncake`. +- A Mooncake master with KV events enabled. This requires the event publisher introduced by [kvcache-ai/Mooncake#2214](https://github.com/kvcache-ai/Mooncake/pull/2214) until that change is available in a Mooncake release. +- A Mooncake KV event endpoint reachable from the Dynamo frontend host. Worker-side Mooncake config (event endpoint, page size, TP/PP layout, and split-head layout) is published automatically through each worker's registration metadata. ## Setup @@ -160,15 +163,16 @@ python -m dynamo.sglang \ --hicache-ratio 2 \ --hicache-write-policy write_through \ --hicache-storage-backend mooncake \ - --hicache-storage-backend-extra-config '{"master_server_address": "mooncake-master.internal:50051"}' \ + --hicache-storage-backend-extra-config '{"master_server_address": "mooncake-master.internal:50051", "enable_group_semantics": true}' \ --skip-tokenizer-init ``` Launch additional workers on other GPUs / hosts with the same Mooncake config so they back to the same cluster. -**Dynamo frontend** — enable tier-aware routing: +**Dynamo frontend** — configure the Mooncake event endpoint and enable tier-aware routing: ```bash +DYN_MOONCAKE_KV_EVENTS_ENDPOINT=tcp://mooncake-master.internal:5557 \ python -m dynamo.frontend \ --http-port 8000 \ --router-mode kv \ @@ -180,12 +184,15 @@ python -m dynamo.frontend \ | Flag | Env var | Default | Description | | --------------------------- | ----------------------------- | ------- | ------------------------------------------------------------------------------------------------------------------------------------------------------------------ | -| `--shared-cache-type` | `DYN_SHARED_CACHE_TYPE` | `none` | `none` disables shared-pool lookups; `hicache` enables Mooncake queries. | -| `--shared-cache-multiplier` | `DYN_SHARED_CACHE_MULTIPLIER` | `0.5` | Discount factor for shared-pool hits. `0.0` queries but ignores them; `0.5` treats a shared hit as half a device hit; `1.0` treats shared and device hits equally. | +| `--shared-cache-type` | `DYN_SHARED_CACHE_TYPE` | `none` | `none` disables shared-pool tracking; `hicache` enables the Mooncake event-backed index. | +| `--shared-cache-multiplier` | `DYN_SHARED_CACHE_MULTIPLIER` | `0.5` | Discount factor for shared-pool hits. `0.0` ignores them; `0.5` treats a shared hit as half a device hit; `1.0` treats shared and device hits equally. | +| — | `DYN_MOONCAKE_KV_EVENTS_ENDPOINT` | unset | Mooncake PUB endpoint consumed by the frontend. When unset, Dynamo uses one consistently advertised worker endpoint. | Per-request overrides are available via `RouterConfigOverride.shared_cache_multiplier` for A/B experimentation without restarting the router. -No extra flags are required on the worker. When `--hicache-storage-backend mooncake` is set, Dynamo publishes the required metadata (page size, TP/PP layout, master address) via the worker's `ModelRuntimeConfig.engine_specific` blob under the key `sglang_hicache_mooncake`. +Set `DYN_MOONCAKE_KV_EVENTS_ENDPOINT` on the frontend to the Mooncake PUB endpoint, such as `tcp://mooncake-master.internal:5557`. The endpoint must be reachable from the frontend. Workers can advertise the same variable as a fallback, but a missing worker value does not disable shared-cache routing. + +Set `enable_group_semantics` to `true` in `--hicache-storage-backend-extra-config` to include SGLang logical group IDs in Mooncake metadata. Dynamo falls back to exact physical-key checks when group metadata is unavailable. ## Verification @@ -199,12 +206,13 @@ python -m dynamo.sglang ... --log-level debug 2>&1 | grep -E 'BlockStored|BlockR If `medium` is missing or Host-tier transitions never report `CPU_PINNED`, confirm that the worker runs SGLang 0.5.11 or later (or a custom build that includes PR #22894). -**Router sees the shared pool.** Two new histograms are exposed on the frontend's Prometheus endpoint: +**Router sees the shared pool.** Shared-cache metrics are exposed on the frontend's Prometheus endpoint: -| Metric | Meaning | -| ----------------------------------- | ------------------------------------------------------------------------ | -| `router_shared_cache_hit_rate` | Fraction of request blocks found in the shared pool (0.0–1.0). | -| `router_shared_cache_beyond_blocks` | Blocks in the shared pool _beyond_ the selected worker's device overlap. | +| Metric | Meaning | +| ----------------------------------------- | ------------------------------------------------------------------------ | +| `router_shared_cache_hit_rate` | Fraction of request blocks found in the shared pool (0.0–1.0). | +| `router_shared_cache_beyond_blocks` | Blocks in the shared pool _beyond_ the selected worker's device overlap. | +| `dynamo_router_shared_cache_errors_total` | Shared-cache query and Mooncake subscriber failures. | ```bash curl -s localhost:8000/metrics | grep shared_cache @@ -214,10 +222,10 @@ curl -s localhost:8000/metrics | grep shared_cache | Symptom | Likely cause | Fix | | -------------------------------------------------------- | ---------------------------------------------------------------------------- | --------------------------------------------------------------------------------------------- | -| `shared_cache_hit_rate` is always 0 | Mooncake master unreachable from the router host | Check network path; the router logs `Shared cache query failed` when it can't reach Mooncake. | -| Events only ever carry `medium=GPU` | SGLang older than 0.5.11 or a custom build missing [PR #22894](https://github.com/sgl-project/sglang/pull/22894) | Upgrade to SGLang 0.5.11 or later. | -| Workers registered but router never queries shared cache | `--shared-cache-type` left at default `none` | Set `--shared-cache-type hicache` on the frontend. | -| Queries issued but winning worker rarely changes | `--shared-cache-multiplier 0.0` | Raise the multiplier — typical starting range is `0.3`–`0.7`. | +| `shared_cache_hit_rate` is always 0 | Event endpoint missing, unreachable, or connected after objects were stored | Set `DYN_MOONCAKE_KV_EVENTS_ENDPOINT` on the frontend and inspect `dynamo_router_shared_cache_errors_total`. | +| Events only ever carry `medium=GPU` | SGLang older than 0.5.11 or a custom build missing [PR #22894](https://github.com/sgl-project/sglang/pull/22894) | Upgrade to SGLang 0.5.11 or later. | +| Workers registered but router never tracks shared cache | `--shared-cache-type` left at default `none` | Set `--shared-cache-type hicache` on the frontend. | +| Shared hits do not affect worker selection | `--shared-cache-multiplier 0.0` | Raise the multiplier — typical starting range is `0.3`–`0.7`. | | Page-size mismatch warnings | Router `--page-size` doesn't match worker `--page-size` | They must agree; the router hashes pages using the worker's page size. | | Router logs "no workers have HiCache enabled" | No worker published `sglang_hicache_mooncake` metadata | Confirm workers started with `--hicache-storage-backend mooncake`. | diff --git a/lib/llm/src/discovery/model_manager.rs b/lib/llm/src/discovery/model_manager.rs index 5ebb92d8e14b..ee97ecb0e2fc 100644 --- a/lib/llm/src/discovery/model_manager.rs +++ b/lib/llm/src/discovery/model_manager.rs @@ -148,6 +148,9 @@ pub struct ModelManager { /// Per-endpoint runtime config watchers. Keyed by EndpointId (includes namespace). runtime_configs: DashMap, + /// Per-endpoint HiCache state and its one Mooncake event subscriber. + hicache_caches: DashMap, + /// Shared KV-source membership coordinators, scoped by exact serving endpoint. /// Weak ownership lets the discovery loop stop when its last consumer goes away. kv_source_memberships: DashMap>, @@ -188,6 +191,7 @@ impl ModelManager { prefill_router_activators: DashMap::new(), encoder_router_activators: DashMap::new(), runtime_configs: DashMap::new(), + hicache_caches: DashMap::new(), kv_source_memberships: DashMap::new(), lora_domains: DashMap::new(), lora_enabled: crate::lora::lora_serving_enabled(), @@ -1050,6 +1054,45 @@ impl ModelManager { lora_enabled && worker_type == crate::protocols::common::timing::WORKER_TYPE_DECODE } + fn hicache_cache_for( + &self, + endpoint: &Endpoint, + runtime_configs: RuntimeConfigWatch, + ) -> HicacheSharedKvCache { + self.hicache_caches + .entry(endpoint.id()) + .or_insert_with(|| { + let frontend_kv_events_endpoint = std::env::var("DYN_MOONCAKE_KV_EVENTS_ENDPOINT") + .ok() + .filter(|endpoint| !endpoint.is_empty()); + let cache = HicacheSharedKvCache::new_with_cancellation_and_endpoint( + runtime_configs, + endpoint.component().drt().child_token(), + frontend_kv_events_endpoint, + ); + cache.start_subscriber(); + cache + }) + .clone() + } + + pub fn remove_hicache_caches(&self, namespace: &str, component: &str) { + let endpoint_ids = self + .hicache_caches + .iter() + .filter(|entry| { + entry.key().namespace == namespace && entry.key().component == component + }) + .map(|entry| entry.key().clone()) + .collect::>(); + + for endpoint_id in endpoint_ids { + if let Some((_, cache)) = self.hicache_caches.remove(&endpoint_id) { + cache.shutdown(); + } + } + } + #[allow(clippy::too_many_arguments)] pub async fn kv_chooser_for( &self, @@ -1127,9 +1170,9 @@ impl ModelManager { worker_component = worker_component_name, "Using HiCache shared KV cache" ); - Some(Box::new(HicacheSharedKvCache::new( - workers_with_configs.clone(), - ))) + Some(Box::new( + self.hicache_cache_for(endpoint, workers_with_configs.clone()), + )) } }; @@ -2378,6 +2421,38 @@ mod tests { assert!(mm.get_model("llama").is_some()); } + #[test] + fn remove_hicache_caches_cancels_only_the_removed_component() { + let manager = ModelManager::new(); + let (_tx, runtime_configs) = + tokio::sync::watch::channel(HashMap::::new()); + let cancelled = tokio_util::sync::CancellationToken::new(); + let retained = tokio_util::sync::CancellationToken::new(); + manager.hicache_caches.insert( + EndpointId::from("ns.worker.generate"), + HicacheSharedKvCache::new_with_cancellation(runtime_configs.clone(), cancelled.clone()), + ); + manager.hicache_caches.insert( + EndpointId::from("ns.other.generate"), + HicacheSharedKvCache::new_with_cancellation(runtime_configs, retained.clone()), + ); + + manager.remove_hicache_caches("ns", "worker"); + + assert!(cancelled.is_cancelled()); + assert!(!retained.is_cancelled()); + assert!( + !manager + .hicache_caches + .contains_key(&EndpointId::from("ns.worker.generate")) + ); + assert!( + manager + .hicache_caches + .contains_key(&EndpointId::from("ns.other.generate")) + ); + } + #[test] fn test_alias_resolution_maps_to_primary() { let mm = ModelManager::new(); diff --git a/lib/llm/src/discovery/watcher.rs b/lib/llm/src/discovery/watcher.rs index 1bd5eb378840..dcfd364bab93 100644 --- a/lib/llm/src/discovery/watcher.rs +++ b/lib/llm/src/discovery/watcher.rs @@ -105,6 +105,16 @@ fn model_card_endpoint_id(mcid: &ModelCardInstanceId) -> EndpointId { } } +fn has_live_endpoint_card( + cards: &[(EndpointId, ModelDeploymentCard)], + namespace: &str, + component: &str, +) -> bool { + cards + .iter() + .any(|(endpoint, _)| endpoint.namespace == namespace && endpoint.component == component) +} + fn model_card_instance_id(instance: &DiscoveryInstance) -> anyhow::Result { match instance { DiscoveryInstance::Model { @@ -1025,10 +1035,13 @@ impl ModelWatcher { // state. If discovery is temporarily unavailable, retaining the card // lets the next reconciliation pass retry this stale entry instead of // losing the key while leaving its WorkerSet behind. - let active_instances = self - .cards_for_model_with_endpoints(&model_name, namespace_filter) - .await - .with_context(|| model_name.clone())?; + let all_cards = self.all_cards().await.with_context(|| model_name.clone())?; + let active_instances = all_cards + .iter() + .filter(|(endpoint_id, card)| { + card.name() == model_name && namespace_filter.matches(&endpoint_id.namespace) + }) + .collect::>(); let card = match self.manager.remove_model_card(&key) { Some(card) => card, @@ -1090,6 +1103,8 @@ impl ModelWatcher { && eid.component == *worker_component && worker_set_key(eid, other_card.model_type, other_card.worker_type) == ws_key }); + let endpoint_has_instances = + has_live_endpoint_card(&all_cards, worker_namespace, worker_component); if !component_has_instances { // No more workers of this component in this namespace — remove its WorkerSet @@ -1111,6 +1126,11 @@ impl ModelWatcher { } } + if !endpoint_has_instances { + self.manager + .remove_hicache_caches(worker_namespace, worker_component); + } + // Activator-state cleanup depends on which component just went away. // // PREFILL teardown (cached endpoint is stale): drop everything for @@ -2185,6 +2205,18 @@ mod tests { } } + #[test] + fn endpoint_liveness_considers_all_models() { + // This is the discovery snapshot after the last adapter card was removed: its base + // model remains on the same worker endpoint, so the endpoint is still live. + let cards = vec![( + test_endpoint_id("generate"), + ModelDeploymentCard::with_name_only("base-model"), + )]; + + assert!(has_live_endpoint_card(&cards, "ns1", "workers")); + } + #[test] fn vllm_generate_requires_explicit_worker_capability() { let mut card = ModelDeploymentCard::with_name_only("model"); diff --git a/lib/llm/src/kv_router/metrics.rs b/lib/llm/src/kv_router/metrics.rs index fdcaf168b114..a7c225636cd1 100644 --- a/lib/llm/src/kv_router/metrics.rs +++ b/lib/llm/src/kv_router/metrics.rs @@ -716,7 +716,7 @@ impl RoutingOverheadMetrics { routing_overhead::SHARED_CACHE_ERRORS_TOTAL ); prometheus::IntCounter::with_opts( - Opts::new(name, "Total shared cache query errors") + Opts::new(name, "Total shared cache failures") .const_label(labels::ROUTER_ID, &router_id), ) .expect("shared_cache_errors_total") diff --git a/lib/llm/src/kv_router/shared_cache.rs b/lib/llm/src/kv_router/shared_cache.rs index 482605c5cfa8..b9c4ceb00e2d 100644 --- a/lib/llm/src/kv_router/shared_cache.rs +++ b/lib/llm/src/kv_router/shared_cache.rs @@ -3,36 +3,45 @@ //! HiCache shared KV cache client for SGLang + Mooncake. //! -//! Instead of querying a worker endpoint over the request plane, this client: +//! This client: //! 1. Reads Mooncake HiCache metadata published by SGLang workers in runtime config. //! 2. Recomputes the logical HiCache page hashes from request tokens using the //! same token -> page-hash logic as SGLang. //! 3. Expands those logical page hashes into the concrete Mooncake object keys //! SGLang uses for the configured TP/PP/MLA layout. -//! 4. Queries the Mooncake master HTTP service directly via `/batch_query_keys`. +//! 4. Tracks those object keys from the Mooncake master's KV event stream. -use std::collections::HashMap; +use std::sync::{ + Arc, + atomic::{AtomicBool, AtomicU64, Ordering}, +}; use std::time::Duration; +use arc_swap::ArcSwapOption; use async_trait::async_trait; -use reqwest::Url; +use dashmap::{DashMap, DashSet}; +use futures::StreamExt; use serde::{Deserialize, Serialize}; use sha2::{Digest, Sha256}; +use tokio_util::sync::CancellationToken; -const MOONCAKE_HTTP_TIMEOUT: Duration = Duration::from_secs(2); - +use crate::{ + discovery::RuntimeConfigWatch, + kv_router::metrics::RoutingOverheadMetrics, + local_model::runtime_config::ModelRuntimeConfig, + utils::zmq::{connect_sub_socket, multipart_message}, +}; use dynamo_kv_router::{ SharedKvCache, indexer::KvRouterError, protocols::{SharedCacheHits, WorkerId}, }; -use crate::{discovery::RuntimeConfigWatch, local_model::runtime_config::ModelRuntimeConfig}; - const SGLANG_HICACHE_MOONCAKE_RUNTIME_KEY: &str = "sglang_hicache_mooncake"; -const MOONCAKE_BATCH_QUERY_KEYS_CHUNK_SIZE: usize = 128; +const MOONCAKE_EVENT_RECONNECT_DELAY: Duration = Duration::from_secs(1); +const MAX_MOONCAKE_INDEX_ENTRIES: usize = 1_000_000; -#[derive(Debug, Clone, Deserialize, Serialize, PartialEq, Eq)] +#[derive(Debug, Clone, Deserialize, Serialize)] struct SglangHicacheMooncakeConfig { backend: String, page_size: u32, @@ -40,51 +49,121 @@ struct SglangHicacheMooncakeConfig { pp_size: u32, is_mla_model: bool, is_eagle: bool, + #[serde(default)] tp_lcm_size: Option, should_split_heads: bool, + #[serde(default)] extra_backend_tag: Option, - master_server_address: Option, - master_metrics_port: u16, + #[serde(default)] + kv_events_endpoint: Option, } -#[derive(Debug, Deserialize)] -struct MooncakeBatchQueryKeysResponse { - success: bool, - #[serde(default)] - data: HashMap, +impl SglangHicacheMooncakeConfig { + fn has_same_layout(&self, other: &Self) -> bool { + self.backend == other.backend + && self.page_size == other.page_size + && self.tp_size == other.tp_size + && self.pp_size == other.pp_size + && self.is_mla_model == other.is_mla_model + && self.is_eagle == other.is_eagle + && self.tp_lcm_size == other.tp_lcm_size + && self.should_split_heads == other.should_split_heads + && self.extra_backend_tag == other.extra_backend_tag + } } -#[derive(Debug, Deserialize, Default)] -struct MooncakeBatchQueryKeyResult { +#[derive(Debug, Deserialize, Serialize)] +struct MooncakeObjectEvent { + event_type: String, + #[serde(default)] + object_key: Option, #[serde(default)] - ok: bool, + tenant_id: String, + #[serde(default)] + group_id: Option, } +/// Mooncake KV publisher payload: `(timestamp_ms, events, dp_rank)` after the empty-topic and +/// big-endian sequence ZMQ frames. +type MooncakeEventBatch = (i64, Vec, u32); + #[derive(Debug, Clone, Copy)] enum QueryToken { Single(u32), Bigram(u32, u32), } -/// Shared KV cache client that queries the Mooncake master HTTP service for -/// SGLang HiCache (L3) state. +/// Event-driven shared KV cache index for SGLang HiCache (L3) state. +#[derive(Clone)] pub struct HicacheSharedKvCache { runtime_configs: RuntimeConfigWatch, - http_client: reqwest::Client, + present_keys: Arc>, + group_states: Arc>, + last_sequence: Arc, + has_sequence: Arc, + last_layout: Arc>, + cancellation_token: CancellationToken, + frontend_kv_events_endpoint: Option, } impl HicacheSharedKvCache { pub fn new(runtime_configs: RuntimeConfigWatch) -> Self { + Self::new_with_cancellation_and_endpoint(runtime_configs, CancellationToken::new(), None) + } + + pub fn new_with_cancellation( + runtime_configs: RuntimeConfigWatch, + cancellation_token: CancellationToken, + ) -> Self { + Self::new_with_cancellation_and_endpoint(runtime_configs, cancellation_token, None) + } + + pub fn new_with_cancellation_and_endpoint( + runtime_configs: RuntimeConfigWatch, + cancellation_token: CancellationToken, + frontend_kv_events_endpoint: Option, + ) -> Self { Self { runtime_configs, - http_client: reqwest::Client::builder() - .timeout(MOONCAKE_HTTP_TIMEOUT) - .build() - .expect("failed to build reqwest client"), + present_keys: Arc::new(DashSet::new()), + group_states: Arc::new(DashMap::new()), + last_sequence: Arc::new(AtomicU64::new(0)), + has_sequence: Arc::new(AtomicBool::new(false)), + last_layout: Arc::new(ArcSwapOption::empty()), + cancellation_token, + frontend_kv_events_endpoint, } } - fn resolve_mooncake_config(&self) -> Option { + pub fn start_subscriber(&self) { + let cache = self.clone(); + let cancellation_token = self.cancellation_token.clone(); + tokio::spawn(async move { cache.run_subscriber(cancellation_token).await }); + } + + pub fn shutdown(&self) { + self.cancellation_token.cancel(); + self.clear(); + } + + fn clear_on_layout_change(&self, layout: &SglangHicacheMooncakeConfig) { + let last_layout = self.last_layout.load(); + if last_layout + .as_ref() + .is_some_and(|previous| previous.has_same_layout(layout)) + { + return; + } + if last_layout.is_some() { + self.clear(); + tracing::warn!("SGLang Mooncake HiCache layout changed; cleared shared-cache state"); + } + self.last_layout.store(Some(Arc::new(layout.clone()))); + } + + fn resolve_mooncake_config_and_endpoint( + &self, + ) -> Option<(SglangHicacheMooncakeConfig, String)> { let workers = self.runtime_configs.borrow(); let mut configs = Vec::new(); @@ -96,7 +175,10 @@ impl HicacheSharedKvCache { let (_, first) = configs.first()?; - if configs.iter().any(|(_, config)| config != first) { + if configs + .iter() + .any(|(_, config)| !config.has_same_layout(first)) + { tracing::warn!( workers = ?configs.iter().map(|(worker_id, _)| *worker_id).collect::>(), "SGLang Mooncake HiCache runtime configs differ across workers; skipping shared-cache lookup" @@ -104,63 +186,205 @@ impl HicacheSharedKvCache { return None; } - Some(first.clone()) + self.clear_on_layout_change(first); + + if let Some(endpoint) = &self.frontend_kv_events_endpoint { + return Some((first.clone(), endpoint.clone())); + } + + let mut endpoints = configs + .iter() + .filter_map(|(_, config)| config.kv_events_endpoint.as_deref()) + .filter(|endpoint| !endpoint.is_empty()); + let endpoint = endpoints.next()?; + if endpoints.any(|candidate| candidate != endpoint) { + tracing::warn!( + "SGLang Mooncake KV event endpoints differ across workers; skipping shared-cache lookup" + ); + return None; + } + Some((first.clone(), endpoint.to_string())) } - async fn fetch_key_presence( - &self, - endpoint: &Url, - actual_keys: &[String], - ) -> Result, KvRouterError> { - let mut key_presence = HashMap::with_capacity(actual_keys.len()); - - for chunk in actual_keys.chunks(MOONCAKE_BATCH_QUERY_KEYS_CHUNK_SIZE) { - let joined_keys = chunk.join(","); - - let mut url = endpoint.clone(); - // Mooncake expects a raw comma-separated `keys=` list. If commas are - // percent-encoded (`%2C`), Mooncake treats the entire value as one key. - url.set_query(Some(&format!("keys={joined_keys}"))); - - let response = self.http_client.get(url.clone()).send().await.map_err(|e| { - tracing::warn!(error = %e, url = %url, "Mooncake batch_query_keys request failed"); - KvRouterError::IndexerOffline - })?; - - let status = response.status(); - if !status.is_success() { - tracing::warn!( - status = %status, - url = %url, - "Mooncake batch_query_keys returned non-success status" - ); - return Err(KvRouterError::IndexerOffline); + fn kv_events_endpoint(&self) -> Option { + self.resolve_mooncake_config_and_endpoint() + .map(|(_, endpoint)| endpoint) + } + + fn apply_batch(&self, sequence: u64, events: Vec) { + // SGLang's ZmqEventPublisher increments this sequence once per published batch, so a + // non-consecutive value means one or more whole batches were missed. + let has_previous = self.has_sequence.swap(true, Ordering::AcqRel); + let previous = self.last_sequence.swap(sequence, Ordering::AcqRel); + if has_previous && sequence == previous { + return; + } + if has_previous && sequence != previous.wrapping_add(1) { + self.present_keys.clear(); + self.group_states.clear(); + tracing::warn!( + previous, + sequence, + "Mooncake KV event sequence gap; cleared shared-cache state" + ); + } + + for event in events { + // The shared-cache query contract currently has no tenant input and + // historically queried Mooncake's default tenant only. + if !event.tenant_id.is_empty() && event.tenant_id != "default" { + continue; + } + let Some(object_key) = event.object_key else { + continue; + }; + let group_id = event.group_id.filter(|id| !id.is_empty()); + match event.event_type.as_str() { + "stored" => { + self.present_keys.insert(object_key); + if let Some(group_id) = group_id { + self.group_states.insert(group_id, (sequence, false)); + } + } + "removed" => { + self.present_keys.remove(&object_key); + if let Some(group_id) = group_id { + self.group_states.remove(&group_id); + } else { + // An older Mooncake publisher may omit `group_id` on removal. Clearing all + // verified groups is conservative and prevents a stale group fast-path hit. + self.group_states.clear(); + } + } + _ => {} } + } + + self.clear_if_index_too_large(MAX_MOONCAKE_INDEX_ENTRIES); + } - let body: MooncakeBatchQueryKeysResponse = response.json().await.map_err(|e| { - tracing::warn!( - error = %e, - url = %url, - "Failed to decode Mooncake batch_query_keys response" - ); - KvRouterError::IndexerOffline - })?; - - if !body.success { - tracing::warn!(url = %url, "Mooncake batch_query_keys reported failure"); - return Err(KvRouterError::IndexerOffline); + fn clear(&self) { + self.present_keys.clear(); + self.group_states.clear(); + self.last_sequence.store(0, Ordering::Release); + self.has_sequence.store(false, Ordering::Release); + } + + fn clear_if_index_too_large(&self, max_entries: usize) { + let present_keys = self.present_keys.len(); + let group_states = self.group_states.len(); + if present_keys.saturating_add(group_states) > max_entries { + self.clear(); + tracing::warn!( + present_keys, + group_states, + max_entries, + "Mooncake KV event index exceeded its size limit; cleared shared-cache state" + ); + } + } + + fn record_subscriber_error(&self) { + if let Some(metrics) = RoutingOverheadMetrics::get() { + metrics.inc_shared_cache_errors(); + } + } + + async fn run_subscriber(mut self, cancellation_token: CancellationToken) { + loop { + let endpoint = loop { + if let Some(endpoint) = self.kv_events_endpoint() { + break endpoint; + } + + tokio::select! { + _ = cancellation_token.cancelled() => return, + result = self.runtime_configs.changed() => { + if result.is_err() { + self.clear(); + return; + } + } + } + }; + + self.clear(); + let mut socket = match connect_sub_socket(&endpoint, None).await { + Ok(socket) => socket, + Err(error) => { + self.record_subscriber_error(); + tracing::warn!(%endpoint, %error, "Failed to connect to Mooncake KV events; retrying"); + tokio::select! { + _ = cancellation_token.cancelled() => return, + _ = tokio::time::sleep(MOONCAKE_EVENT_RECONNECT_DELAY) => continue, + } + } + }; + tracing::info!(%endpoint, "Connected to Mooncake KV events"); + + loop { + tokio::select! { + _ = cancellation_token.cancelled() => return, + result = self.runtime_configs.changed() => { + if result.is_err() { + self.clear(); + return; + } + let next_endpoint = self.kv_events_endpoint(); + if next_endpoint.as_deref() != Some(endpoint.as_str()) { + tracing::info!(%endpoint, next_endpoint = ?next_endpoint, "Mooncake KV event endpoint changed; reconnecting"); + break; + } + } + message = socket.next() => { + let frames = match message { + Some(Ok(frames)) => multipart_message(frames), + Some(Err(error)) => { + self.record_subscriber_error(); + tracing::warn!(%endpoint, %error, "Mooncake KV event stream failed; reconnecting"); + break; + } + None => { + self.record_subscriber_error(); + tracing::warn!(%endpoint, "Mooncake KV event stream ended; reconnecting"); + break; + } + }; + match parse_mooncake_event_frames(&frames) { + Ok((sequence, events)) => self.apply_batch(sequence, events), + Err(error) => { + self.record_subscriber_error(); + tracing::warn!(%error, "Dropping invalid Mooncake KV event frame"); + } + } + } + } } - for key in chunk { - let exists = body.data.get(key).map(|entry| entry.ok).unwrap_or(false); - key_presence.insert(key.clone(), exists); + tokio::select! { + _ = cancellation_token.cancelled() => return, + _ = tokio::time::sleep(MOONCAKE_EVENT_RECONNECT_DELAY) => {} } } - - Ok(key_presence) } } +fn parse_mooncake_event_frames( + frames: &[Vec], +) -> anyhow::Result<(u64, Vec)> { + let [_, sequence, payload] = frames else { + anyhow::bail!("expected three frames, got {}", frames.len()); + }; + let sequence = u64::from_be_bytes( + sequence + .as_slice() + .try_into() + .map_err(|_| anyhow::anyhow!("expected an 8-byte sequence frame"))?, + ); + let (_, events, _) = rmp_serde::from_slice::(payload)?; + Ok((sequence, events)) +} + #[async_trait] impl SharedKvCache for HicacheSharedKvCache { async fn check_blocks( @@ -177,7 +401,7 @@ impl SharedKvCache for HicacheSharedKvCache { return Ok(SharedCacheHits::default()); } - let Some(config) = self.resolve_mooncake_config() else { + let Some((config, _endpoint)) = self.resolve_mooncake_config_and_endpoint() else { tracing::debug!("No SGLang Mooncake HiCache runtime config available"); return Ok(SharedCacheHits::default()); }; @@ -205,28 +429,34 @@ impl SharedKvCache for HicacheSharedKvCache { return Ok(SharedCacheHits::default()); } - let Some(endpoint) = mooncake_batch_query_endpoint(&config) else { - tracing::debug!("Mooncake master HTTP endpoint is unavailable"); - return Ok(SharedCacheHits::default()); - }; - let page_hashes = logical_page_hashes(tokens, config.page_size, config.is_eagle); if page_hashes.is_empty() { return Ok(SharedCacheHits::default()); } - let page_query_keys = build_page_query_keys(&page_hashes, &config); - let all_actual_keys = page_query_keys - .iter() - .flat_map(|keys| keys.iter().cloned()) - .collect::>(); - - let key_presence = self.fetch_key_presence(&endpoint, &all_actual_keys).await?; - let page_hits = page_query_keys + let page_hits = page_hashes .iter() - .map(|keys| { - keys.iter() - .all(|key| key_presence.get(key).copied().unwrap_or(false)) + .map(|page_hash| { + let group_id = sglang_group_id(page_hash, &config); + let generation = self.group_states.get(&group_id).map(|state| *state); + if generation.is_some_and(|(_, verified)| verified) { + return true; + } + + let hit = expand_actual_query_keys(page_hash, &config) + .iter() + .all(|key| self.present_keys.contains(key)); + if hit + && let Some((generation, _)) = generation + && let Some(mut state) = self.group_states.get_mut(&group_id) + { + // A concurrent stored event invalidates verification, not the physical key + // check that already proved this request is a hit. + if state.0 == generation { + state.1 = true; + } + } + hit }) .collect::>(); @@ -255,33 +485,6 @@ fn mooncake_config_from_runtime( } } -fn mooncake_batch_query_endpoint(config: &SglangHicacheMooncakeConfig) -> Option { - let master_server_address = config.master_server_address.as_deref()?; - - let mut url = Url::parse(&format!("http://{master_server_address}")) - .inspect_err(|error| { - tracing::warn!( - master_server_address, - %error, - "Failed to parse Mooncake master address" - ); - }) - .ok()?; - - if url.set_port(Some(config.master_metrics_port)).is_err() { - tracing::warn!( - master_server_address, - master_metrics_port = config.master_metrics_port, - "Failed to set Mooncake master HTTP port" - ); - return None; - } - - url.set_path("/batch_query_keys"); - url.set_query(None); - Some(url) -} - fn logical_page_hashes(tokens: &[u32], page_size: u32, is_eagle: bool) -> Vec { let page_size = page_size as usize; if page_size == 0 { @@ -347,14 +550,15 @@ fn hex_encode(bytes: &[u8]) -> String { output } -fn build_page_query_keys( - page_hashes: &[String], - config: &SglangHicacheMooncakeConfig, -) -> Vec> { - page_hashes - .iter() - .map(|page_hash| expand_actual_query_keys(page_hash, config)) - .collect() +fn sglang_group_id(logical_page_hash: &str, config: &SglangHicacheMooncakeConfig) -> String { + match config + .extra_backend_tag + .as_deref() + .filter(|tag| !tag.is_empty()) + { + Some(tag) => format!("sglang-hicache:{tag}_{logical_page_hash}"), + None => format!("sglang-hicache:{logical_page_hash}"), + } } fn expand_actual_query_keys( @@ -410,11 +614,9 @@ fn maybe_prefix_key(logical_key: &str, extra_backend_tag: Option<&str>) -> Strin #[cfg(test)] mod tests { - use std::ops::Range; + use std::{collections::HashMap, ops::Range}; use super::*; - use mockito::{Matcher, Server}; - use serde_json::json; use tokio::sync::watch; fn mooncake_config() -> SglangHicacheMooncakeConfig { @@ -428,12 +630,16 @@ mod tests { tp_lcm_size: None, should_split_heads: false, extra_backend_tag: None, - master_server_address: Some("127.0.0.1:50051".to_string()), - master_metrics_port: 9003, + kv_events_endpoint: Some("tcp://127.0.0.1:5557".to_string()), } } - fn runtime_watch_with_config(config: SglangHicacheMooncakeConfig) -> RuntimeConfigWatch { + fn runtime_watch_with_config_and_sender( + config: SglangHicacheMooncakeConfig, + ) -> ( + RuntimeConfigWatch, + watch::Sender>, + ) { let mut runtime_config = ModelRuntimeConfig::new(); runtime_config .set_engine_specific(SGLANG_HICACHE_MOONCAKE_RUNTIME_KEY, config) @@ -442,8 +648,12 @@ mod tests { let mut workers = HashMap::new(); workers.insert(1, runtime_config); - let (_tx, rx) = watch::channel(workers); - rx + let (tx, rx) = watch::channel(workers); + (rx, tx) + } + + fn runtime_watch_with_config(config: SglangHicacheMooncakeConfig) -> RuntimeConfigWatch { + runtime_watch_with_config_and_sender(config).0 } #[test] @@ -531,65 +741,408 @@ mod tests { ); } + #[test] + fn test_sglang_group_id_uses_extra_backend_tag() { + let config = SglangHicacheMooncakeConfig { + extra_backend_tag: Some("tag".to_string()), + ..mooncake_config() + }; + + assert_eq!(sglang_group_id("hash", &config), "sglang-hicache:tag_hash"); + } + + #[test] + fn test_parse_mooncake_event_frames() { + let payload = rmp_serde::to_vec(&( + 0_i64, + vec![MooncakeObjectEvent { + event_type: "stored".to_string(), + object_key: Some("key-0".to_string()), + tenant_id: "default".to_string(), + group_id: Some("group-0".to_string()), + }], + 0_u32, + )) + .unwrap(); + + let (sequence, events) = + parse_mooncake_event_frames(&[Vec::new(), 7_u64.to_be_bytes().to_vec(), payload]) + .unwrap(); + + assert_eq!(sequence, 7); + assert_eq!(events[0].object_key.as_deref(), Some("key-0")); + assert_eq!(events[0].group_id.as_deref(), Some("group-0")); + } + + #[test] + fn test_kv_events_endpoint_tracks_runtime_config_updates() { + let (runtime_configs, tx) = runtime_watch_with_config_and_sender(mooncake_config()); + let cache = HicacheSharedKvCache::new(runtime_configs); + assert_eq!( + cache.kv_events_endpoint().as_deref(), + Some("tcp://127.0.0.1:5557") + ); + + let mut updated = mooncake_config(); + updated.kv_events_endpoint = Some("tcp://127.0.0.1:5558".to_string()); + let mut runtime_config = ModelRuntimeConfig::new(); + runtime_config + .set_engine_specific(SGLANG_HICACHE_MOONCAKE_RUNTIME_KEY, updated) + .unwrap(); + tx.send(HashMap::from([(1, runtime_config)])).unwrap(); + + assert_eq!( + cache.kv_events_endpoint().as_deref(), + Some("tcp://127.0.0.1:5558") + ); + } + + #[test] + fn test_kv_events_endpoint_tolerates_worker_metadata_omission() { + let mut advertised = mooncake_config(); + advertised.kv_events_endpoint = Some("tcp://127.0.0.1:5557".to_string()); + let mut missing = advertised.clone(); + missing.kv_events_endpoint = None; + let mut advertised_runtime = ModelRuntimeConfig::new(); + advertised_runtime + .set_engine_specific(SGLANG_HICACHE_MOONCAKE_RUNTIME_KEY, advertised) + .unwrap(); + let mut missing_runtime = ModelRuntimeConfig::new(); + missing_runtime + .set_engine_specific(SGLANG_HICACHE_MOONCAKE_RUNTIME_KEY, missing) + .unwrap(); + let (_tx, runtime_configs) = watch::channel(HashMap::from([ + (1, advertised_runtime), + (2, missing_runtime), + ])); + let cache = HicacheSharedKvCache::new(runtime_configs); + + assert_eq!( + cache.kv_events_endpoint().as_deref(), + Some("tcp://127.0.0.1:5557") + ); + } + + #[test] + fn test_frontend_kv_events_endpoint_overrides_worker_metadata() { + let mut worker_config = mooncake_config(); + worker_config.kv_events_endpoint = None; + let cache = HicacheSharedKvCache::new_with_cancellation_and_endpoint( + runtime_watch_with_config(worker_config), + CancellationToken::new(), + Some("tcp://frontend-config:5557".to_string()), + ); + + assert_eq!( + cache.kv_events_endpoint().as_deref(), + Some("tcp://frontend-config:5557") + ); + } + #[tokio::test] - async fn test_check_blocks_queries_mooncake_master() { - let mut server = Server::new_async().await; - let server_url = Url::parse(&server.url()).unwrap(); + async fn test_layout_change_clears_cached_hits() { + let config = mooncake_config(); + let (runtime_configs, tx) = runtime_watch_with_config_and_sender(config.clone()); + let cache = HicacheSharedKvCache::new(runtime_configs); + let hash = logical_page_hashes(&[1, 2, 3, 4], config.page_size, config.is_eagle) + .pop() + .unwrap(); + let group_id = sglang_group_id(&hash, &config); + cache.apply_batch( + 1, + expand_actual_query_keys(&hash, &config) + .into_iter() + .map(|object_key| MooncakeObjectEvent { + event_type: "stored".to_string(), + object_key: Some(object_key), + tenant_id: "default".to_string(), + group_id: Some(group_id.clone()), + }) + .collect(), + ); + assert_eq!( + cache + .check_blocks(&[1, 2, 3, 4], config.page_size, None) + .await + .unwrap() + .total_hits, + 1 + ); + let mut new_layout = config; + new_layout.tp_size = 2; + let mut runtime_config = ModelRuntimeConfig::new(); + runtime_config + .set_engine_specific(SGLANG_HICACHE_MOONCAKE_RUNTIME_KEY, new_layout) + .unwrap(); + tx.send(HashMap::from([(1, runtime_config)])).unwrap(); + + assert_eq!( + cache + .check_blocks(&[1, 2, 3, 4], 4, None) + .await + .unwrap() + .total_hits, + 0 + ); + assert!(cache.present_keys.is_empty()); + assert!(cache.group_states.is_empty()); + } + + #[tokio::test] + async fn test_check_blocks_uses_mooncake_events() { let hash0 = "cf97adeedb59e05bfd73a2b4c2a8885708c4f4f70c84c64b27120e72ab733b72".to_string(); let hash1 = "4ebfa8a1f3c341517621838c6e1b9aa350307e3f00b3cbd1a07ef740f54396d6".to_string(); + let cache = HicacheSharedKvCache::new(runtime_watch_with_config(mooncake_config())); + cache.apply_batch( + 1, + vec![ + MooncakeObjectEvent { + event_type: "stored".to_string(), + object_key: Some(format!("{hash0}_0_k")), + tenant_id: "default".to_string(), + group_id: None, + }, + MooncakeObjectEvent { + event_type: "stored".to_string(), + object_key: Some(format!("{hash0}_0_v")), + tenant_id: "default".to_string(), + group_id: None, + }, + MooncakeObjectEvent { + event_type: "stored".to_string(), + object_key: Some(format!("{hash1}_0_k")), + tenant_id: "default".to_string(), + group_id: None, + }, + ], + ); + let hits = cache + .check_blocks(&[1, 2, 3, 4, 5, 6, 7, 8], 4, None) + .await + .unwrap(); - let response = json!({ - "success": true, - "data": { - format!("{hash0}_0_k"): {"ok": true, "values": []}, - format!("{hash0}_0_v"): {"ok": true, "values": []}, - format!("{hash1}_0_k"): {"ok": true, "values": []}, - format!("{hash1}_0_v"): {"ok": false, "error": "not found"}, - } - }); - - let mock = server - .mock("GET", "/batch_query_keys") - .match_query(Matcher::Exact(format!( - "keys={hash0}_0_k,{hash0}_0_v,{hash1}_0_k,{hash1}_0_v" - ))) - .with_status(200) - .with_header("content-type", "application/json") - .with_body(response.to_string()) - .create_async() - .await; - - let config = SglangHicacheMooncakeConfig { - master_server_address: Some(format!("{}:50051", server_url.host_str().unwrap())), - master_metrics_port: server_url.port().unwrap(), - ..mooncake_config() - }; + assert_eq!(hits.ranges, vec![Range { start: 0, end: 1 }]); + assert_eq!(hits.total_hits, 1); - let cache = HicacheSharedKvCache::new(runtime_watch_with_config(config)); + cache.apply_batch( + 2, + vec![MooncakeObjectEvent { + event_type: "removed".to_string(), + object_key: Some(format!("{hash0}_0_v")), + tenant_id: "default".to_string(), + group_id: None, + }], + ); let hits = cache .check_blocks(&[1, 2, 3, 4, 5, 6, 7, 8], 4, None) .await .unwrap(); + assert_eq!(hits.total_hits, 0); + } - assert_eq!(hits.ranges, vec![Range { start: 0, end: 1 }]); + #[tokio::test] + async fn test_check_blocks_invalidates_group_on_unlabeled_removal() { + let hash = "cf97adeedb59e05bfd73a2b4c2a8885708c4f4f70c84c64b27120e72ab733b72"; + let group_id = format!("sglang-hicache:{hash}"); + let cache = HicacheSharedKvCache::new(runtime_watch_with_config(mooncake_config())); + cache.apply_batch( + 1, + vec![ + MooncakeObjectEvent { + event_type: "stored".to_string(), + object_key: Some(format!("{hash}_0_k")), + tenant_id: "default".to_string(), + group_id: Some(group_id.clone()), + }, + MooncakeObjectEvent { + event_type: "stored".to_string(), + object_key: Some(format!("{hash}_0_v")), + tenant_id: "default".to_string(), + group_id: Some(group_id.clone()), + }, + ], + ); + + let hits = cache.check_blocks(&[1, 2, 3, 4], 4, None).await.unwrap(); assert_eq!(hits.total_hits, 1); + assert!(cache.group_states.get(&group_id).is_some_and(|v| v.1)); + + cache.apply_batch( + 2, + vec![MooncakeObjectEvent { + event_type: "removed".to_string(), + object_key: Some(format!("{hash}_0_v")), + tenant_id: "default".to_string(), + group_id: None, + }], + ); + assert!(cache.group_states.is_empty()); + let hits = cache.check_blocks(&[1, 2, 3, 4], 4, None).await.unwrap(); + assert_eq!(hits.total_hits, 0); + } + + #[tokio::test] + async fn test_labeled_removal_preserves_other_verified_groups() { + let hash0 = "cf97adeedb59e05bfd73a2b4c2a8885708c4f4f70c84c64b27120e72ab733b72"; + let hash1 = "4ebfa8a1f3c341517621838c6e1b9aa350307e3f00b3cbd1a07ef740f54396d6"; + let group0 = format!("sglang-hicache:{hash0}"); + let group1 = format!("sglang-hicache:{hash1}"); + let cache = HicacheSharedKvCache::new(runtime_watch_with_config(mooncake_config())); + cache.apply_batch( + 1, + [(&hash0, &group0), (&hash1, &group1)] + .into_iter() + .flat_map(|(hash, group_id)| { + ["k", "v"].into_iter().map(move |kind| MooncakeObjectEvent { + event_type: "stored".to_string(), + object_key: Some(format!("{hash}_0_{kind}")), + tenant_id: "default".to_string(), + group_id: Some(group_id.clone()), + }) + }) + .collect(), + ); + assert_eq!( + cache + .check_blocks(&[1, 2, 3, 4, 5, 6, 7, 8], 4, None) + .await + .unwrap() + .total_hits, + 2 + ); + + cache.apply_batch( + 2, + vec![MooncakeObjectEvent { + event_type: "removed".to_string(), + object_key: Some(format!("{hash0}_0_k")), + tenant_id: "default".to_string(), + group_id: Some(group0.clone()), + }], + ); - mock.assert_async().await; + assert!(!cache.group_states.contains_key(&group0)); + assert!(cache.group_states.get(&group1).is_some_and(|state| state.1)); + } + + #[test] + fn test_duplicate_sequence_preserves_shared_cache_state() { + let cache = HicacheSharedKvCache::new(runtime_watch_with_config(mooncake_config())); + cache.apply_batch( + 0, + vec![MooncakeObjectEvent { + event_type: "stored".to_string(), + object_key: Some("key-0".to_string()), + tenant_id: "default".to_string(), + group_id: None, + }], + ); + cache.apply_batch( + 0, + vec![MooncakeObjectEvent { + event_type: "removed".to_string(), + object_key: Some("key-0".to_string()), + tenant_id: "default".to_string(), + group_id: None, + }], + ); + + assert!(cache.present_keys.contains("key-0")); + } + + #[test] + fn test_index_size_limit_clears_shared_cache_state() { + let cache = HicacheSharedKvCache::new(runtime_watch_with_config(mooncake_config())); + cache.apply_batch( + 1, + vec![MooncakeObjectEvent { + event_type: "stored".to_string(), + object_key: Some("key-0".to_string()), + tenant_id: "default".to_string(), + group_id: Some("group-0".to_string()), + }], + ); + + cache.clear_if_index_too_large(1); + + assert!(cache.present_keys.is_empty()); + assert!(cache.group_states.is_empty()); } #[tokio::test] - async fn test_check_blocks_skips_mooncake_for_cache_namespace() { - let server = Server::new_async().await; - let server_url = Url::parse(&server.url()).unwrap(); + async fn test_subscriber_retries_failed_connection_until_cancelled() { + let mut config = mooncake_config(); + config.kv_events_endpoint = Some("invalid://mooncake-events".to_string()); + let (runtime_configs, _tx) = runtime_watch_with_config_and_sender(config); + let cache = HicacheSharedKvCache::new(runtime_configs); + let cancellation_token = CancellationToken::new(); + let task = tokio::spawn(cache.run_subscriber(cancellation_token.clone())); + + tokio::time::sleep(Duration::from_millis(20)).await; + assert!(!task.is_finished()); + cancellation_token.cancel(); + tokio::time::timeout(Duration::from_secs(1), task) + .await + .unwrap() + .unwrap(); + } - let config = SglangHicacheMooncakeConfig { - master_server_address: Some(format!("{}:50051", server_url.host_str().unwrap())), - master_metrics_port: server_url.port().unwrap(), - ..mooncake_config() - }; + #[tokio::test] + async fn test_subscriber_clears_state_when_runtime_config_watch_closes() { + let (tx, runtime_configs) = watch::channel(HashMap::::new()); + let cache = HicacheSharedKvCache::new(runtime_configs); + cache.apply_batch( + 1, + vec![MooncakeObjectEvent { + event_type: "stored".to_string(), + object_key: Some("key-0".to_string()), + tenant_id: "default".to_string(), + group_id: None, + }], + ); + let task = tokio::spawn(cache.clone().run_subscriber(CancellationToken::new())); - let cache = HicacheSharedKvCache::new(runtime_watch_with_config(config)); + drop(tx); + tokio::time::timeout(Duration::from_secs(1), task) + .await + .unwrap() + .unwrap(); + + assert!(cache.present_keys.is_empty()); + } + + #[test] + fn test_sequence_gap_clears_stale_keys() { + let cache = HicacheSharedKvCache::new(runtime_watch_with_config(mooncake_config())); + cache.apply_batch( + 1, + vec![MooncakeObjectEvent { + event_type: "stored".to_string(), + object_key: Some("old_0_k".to_string()), + tenant_id: "default".to_string(), + group_id: Some("old-group".to_string()), + }], + ); + assert!(cache.group_states.contains_key("old-group")); + cache.apply_batch( + 3, + vec![MooncakeObjectEvent { + event_type: "stored".to_string(), + object_key: Some("new_0_k".to_string()), + tenant_id: "default".to_string(), + group_id: None, + }], + ); + + assert!(!cache.present_keys.contains("old_0_k")); + assert!(cache.group_states.is_empty()); + assert!(cache.present_keys.contains("new_0_k")); + } + + #[tokio::test] + async fn test_check_blocks_skips_mooncake_for_cache_namespace() { + let cache = HicacheSharedKvCache::new(runtime_watch_with_config(mooncake_config())); let hits = cache .check_blocks(&[1, 2, 3, 4, 5, 6, 7, 8], 4, Some("tenant-a")) .await diff --git a/lib/runtime/src/metrics/prometheus_names.rs b/lib/runtime/src/metrics/prometheus_names.rs index 42148314deb0..000cda5c8400 100644 --- a/lib/runtime/src/metrics/prometheus_names.rs +++ b/lib/runtime/src/metrics/prometheus_names.rs @@ -592,7 +592,7 @@ pub mod routing_overhead { /// Time spent querying the shared KV cache (Mooncake) pub const SHARED_CACHE_QUERY_MS: &str = "overhead_shared_cache_query_ms"; - /// Total shared cache query errors (timeouts, HTTP failures) + /// Total shared cache failures (query and subscriber failures) pub const SHARED_CACHE_ERRORS_TOTAL: &str = "shared_cache_errors_total"; }