chore: Use KeyValueStoreManager instead of etcd::Client (#3822)

Signed-off-by: Graham King <[email protected]>

Author: grahamking

lib/bindings/c/src/lib.rs

@@ -154,7 +154,7 @@ fn dynamo_create_kv_publisher(
     {
         Ok(drt) => {
             let backend = drt.namespace(namespace)?.component(component)?;
-            KvEventPublisher::new(backend, worker_id, kv_block_size, None)
+            KvEventPublisher::new(backend, worker_id as u64, kv_block_size, None)
         }
         Err(e) => Err(e),
     }

lib/bindings/python/rust/lib.rs

@@ -771,7 +771,7 @@ impl Endpoint {
         })
     }
 
-    fn lease_id(&self) -> i64 {
+    fn lease_id(&self) -> u64 {
         self.inner
             .drt()
             .primary_lease()
@@ -807,7 +807,7 @@ impl Namespace {
 impl Client {
     /// Get list of current instances.
     /// Replaces endpoint_ids.
-    fn instance_ids(&self) -> Vec<i64> {
+    fn instance_ids(&self) -> Vec<u64> {
         self.router.client.instance_ids()
     }
 
@@ -819,7 +819,7 @@ impl Client {
             inner
                 .wait_for_instances()
                 .await
-                .map(|v| v.into_iter().map(|cei| cei.id()).collect::<Vec<i64>>())
+                .map(|v| v.into_iter().map(|cei| cei.id()).collect::<Vec<u64>>())
                 .map_err(to_pyerr)
         })
     }
@@ -920,7 +920,7 @@ impl Client {
         &self,
         py: Python<'p>,
         request: PyObject,
-        instance_id: i64,
+        instance_id: u64,
         annotated: Option<bool>,
         context: Option<context::Context>,
     ) -> PyResult<Bound<'p, PyAny>> {

lib/bindings/python/rust/llm/kv.rs

@@ -322,7 +322,7 @@ pub(crate) struct OverlapScores {
 #[pymethods]
 impl OverlapScores {
     #[getter]
-    fn scores(&self) -> HashMap<(i64, u32), u32> {
+    fn scores(&self) -> HashMap<(u64, u32), u32> {
         // Return scores with full WorkerWithDpRank granularity as (worker_id, dp_rank) tuples
         self.inner
             .scores

lib/llm/src/block_manager/controller/client.rs

@@ -59,7 +59,10 @@ impl ControlClient {
     }
 
     async fn execute<T: DeserializeOwned>(&self, message: ControlMessage) -> Result<T> {
-        let mut stream = self.client.direct(message.into(), self.instance_id).await?;
+        let mut stream = self
+            .client
+            .direct(message.into(), self.instance_id as u64)
+            .await?;
         let resp = stream
             .next()
             .await

lib/llm/src/discovery/watcher.rs

@@ -567,19 +567,22 @@ impl ModelWatcher {
     }
 
     /// All the registered ModelDeploymentCard with the EndpointId they are attached to, one per instance
-    pub async fn all_cards(&self) -> anyhow::Result<Vec<(EndpointId, ModelDeploymentCard)>> {
-        let Some(etcd_client) = self.drt.etcd_client() else {
-            anyhow::bail!("all_cards: Missing etcd client");
+    async fn all_cards(&self) -> anyhow::Result<Vec<(EndpointId, ModelDeploymentCard)>> {
+        let store = self.drt.store();
+
+        //let kvs = etcd_client.kv_get_prefix(model_card::ROOT_PATH).await?;
+        let Some(card_bucket) = store.get_bucket(model_card::ROOT_PATH).await? else {
+            // no cards
+            return Ok(vec![]);
         };
-        let kvs = etcd_client.kv_get_prefix(model_card::ROOT_PATH).await?;
-        let mut results = Vec::with_capacity(kvs.len());
-        for kv in kvs {
-            let maybe_convert = serde_json::from_slice::<ModelDeploymentCard>(kv.value());
-            let r = match maybe_convert {
+        let entries = card_bucket.entries().await?;
+
+        let mut results = Vec::with_capacity(entries.len());
+        for (key, card_bytes) in entries {
+            let r = match serde_json::from_slice::<ModelDeploymentCard>(&card_bytes) {
                 Ok(card) => {
-                    let maybe_endpoint_id = kv.key_str().map_err(|err| err.into()).and_then(|k| {
-                        etcd_key_extract(k).map(|(endpoint_id, _instance_id)| endpoint_id)
-                    });
+                    let maybe_endpoint_id =
+                        etcd_key_extract(&key).map(|(endpoint_id, _instance_id)| endpoint_id);
                     let endpoint_id = match maybe_endpoint_id {
                         Ok(eid) => eid,
                         Err(err) => {
@@ -590,14 +593,8 @@ impl ModelWatcher {
                     (endpoint_id, card)
                 }
                 Err(err) => {
-                    match kv.value_str() {
-                        Ok(value) => {
-                            tracing::error!(%err, value, "Invalid JSON in model card");
-                        }
-                        Err(value_str_err) => {
-                            tracing::error!(original_error=%err, %value_str_err, "Invalid UTF-8 string in model card, expected JSON");
-                        }
-                    }
+                    let value = String::from_utf8_lossy(&card_bytes);
+                    tracing::error!(%err, %value, "Invalid JSON in model card");
                     continue;
                 }
             };

lib/llm/src/discovery/worker_monitor.rs

@@ -52,7 +52,7 @@ impl WorkerLoadState {
 /// Worker monitor for tracking KV cache usage and busy states
 pub struct KvWorkerMonitor {
     client: Arc<Client>,
-    worker_load_states: Arc<RwLock<HashMap<i64, WorkerLoadState>>>,
+    worker_load_states: Arc<RwLock<HashMap<u64, WorkerLoadState>>>,
     busy_threshold: f64,
 }
 
@@ -67,7 +67,7 @@ impl KvWorkerMonitor {
     }
 
     /// Get the worker load states for external access
-    pub fn load_states(&self) -> Arc<RwLock<HashMap<i64, WorkerLoadState>>> {
+    pub fn load_states(&self) -> Arc<RwLock<HashMap<u64, WorkerLoadState>>> {
         self.worker_load_states.clone()
     }
 }
@@ -154,7 +154,7 @@ impl WorkerLoadMonitor for KvWorkerMonitor {
 
                             // Recalculate all busy instances and update
                             let states = worker_load_states.read().unwrap();
-                            let busy_instances: Vec<i64> = states
+                            let busy_instances: Vec<u64> = states
                                 .iter()
                                 .filter_map(|(&id, state)| {
                                     state.is_busy(busy_threshold).then_some(id)

lib/llm/src/entrypoint/input/http.rs

@@ -64,10 +64,10 @@ pub async fn run(runtime: Runtime, engine_config: EngineConfig) -> anyhow::Resul
     let http_service = match engine_config {
         EngineConfig::Dynamic(_) => {
             let distributed_runtime = DistributedRuntime::from_settings(runtime.clone()).await?;
-            let etcd_client = distributed_runtime.etcd_client();
             // This allows the /health endpoint to query etcd for active instances
-            http_service_builder = http_service_builder.with_etcd_client(etcd_client.clone());
+            http_service_builder = http_service_builder.store(distributed_runtime.store().clone());
             let http_service = http_service_builder.build()?;
+            let etcd_client = distributed_runtime.etcd_client();
             match etcd_client {
                 Some(ref etcd_client) => {
                     let router_config = engine_config.local_model().router_config();
@@ -241,17 +241,7 @@ pub async fn run(runtime: Runtime, engine_config: EngineConfig) -> anyhow::Resul
             http_service.custom_backend_registry.as_ref(),
         ) {
             // Create DistributedRuntime for polling, matching the engine's mode
-            // Check if we have etcd_client to determine if we're in dynamic or static mode
-            let drt = if http_service.state().etcd_client().is_some() {
-                // Dynamic mode: use from_settings() which respects environment (includes etcd)
-                DistributedRuntime::from_settings(runtime.clone()).await?
-            } else {
-                // Static mode: no etcd
-                let dst_config =
-                    dynamo_runtime::distributed::DistributedConfig::from_settings(true);
-                DistributedRuntime::new(runtime.clone(), dst_config).await?
-            };
-
+            let drt = DistributedRuntime::from_settings(runtime.clone()).await?;
             tracing::info!(
                 namespace_component_endpoint=%namespace_component_endpoint,
                 polling_interval_secs=polling_interval,

lib/llm/src/http/service/service_v2.rs

@@ -20,10 +20,7 @@ use axum_server::tls_rustls::RustlsConfig;
 use derive_builder::Builder;
 use dynamo_runtime::logging::make_request_span;
 use dynamo_runtime::metrics::prometheus_names::name_prefix;
-use dynamo_runtime::storage::key_value_store::EtcdStore;
-use dynamo_runtime::storage::key_value_store::KeyValueStore;
-use dynamo_runtime::storage::key_value_store::MemoryStore;
-use dynamo_runtime::transports::etcd;
+use dynamo_runtime::storage::key_value_store::KeyValueStoreManager;
 use std::net::SocketAddr;
 use tokio::task::JoinHandle;
 use tokio_util::sync::CancellationToken;
@@ -33,8 +30,7 @@ use tower_http::trace::TraceLayer;
 pub struct State {
     metrics: Arc<Metrics>,
     manager: Arc<ModelManager>,
-    etcd_client: Option<etcd::Client>,
-    store: Arc<dyn KeyValueStore>,
+    store: KeyValueStoreManager,
     flags: StateFlags,
 }
 
@@ -75,12 +71,11 @@ impl StateFlags {
 }
 
 impl State {
-    pub fn new(manager: Arc<ModelManager>) -> Self {
+    pub fn new(manager: Arc<ModelManager>, store: KeyValueStoreManager) -> Self {
         Self {
             manager,
             metrics: Arc::new(Metrics::default()),
-            etcd_client: None,
-            store: Arc::new(MemoryStore::new()),
+            store,
             flags: StateFlags {
                 chat_endpoints_enabled: AtomicBool::new(false),
                 cmpl_endpoints_enabled: AtomicBool::new(false),
@@ -90,20 +85,6 @@ impl State {
         }
     }
 
-    pub fn new_with_etcd(manager: Arc<ModelManager>, etcd_client: etcd::Client) -> Self {
-        Self {
-            manager,
-            metrics: Arc::new(Metrics::default()),
-            store: Arc::new(EtcdStore::new(etcd_client.clone())),
-            etcd_client: Some(etcd_client),
-            flags: StateFlags {
-                chat_endpoints_enabled: AtomicBool::new(false),
-                cmpl_endpoints_enabled: AtomicBool::new(false),
-                embeddings_endpoints_enabled: AtomicBool::new(false),
-                responses_endpoints_enabled: AtomicBool::new(false),
-            },
-        }
-    }
     /// Get the Prometheus [`Metrics`] object which tracks request counts and inflight requests
     pub fn metrics_clone(&self) -> Arc<Metrics> {
         self.metrics.clone()
@@ -117,12 +98,8 @@ impl State {
         self.manager.clone()
     }
 
-    pub fn etcd_client(&self) -> Option<&etcd::Client> {
-        self.etcd_client.as_ref()
-    }
-
-    pub fn store(&self) -> Arc<dyn KeyValueStore> {
-        self.store.clone()
+    pub fn store(&self) -> &KeyValueStoreManager {
+        &self.store
     }
 
     // TODO
@@ -186,8 +163,8 @@ pub struct HttpServiceConfig {
     #[builder(default = "None")]
     request_template: Option<RequestTemplate>,
 
-    #[builder(default = "None")]
-    etcd_client: Option<etcd::Client>,
+    #[builder(default)]
+    store: KeyValueStoreManager,
 
     // DEPRECATED: To be removed after custom backends migrate to Dynamo backend.
     #[builder(default = "None")]
@@ -335,10 +312,7 @@ impl HttpServiceConfigBuilder {
         let config: HttpServiceConfig = self.build_internal()?;
 
         let model_manager = Arc::new(ModelManager::new());
-        let state = match config.etcd_client {
-            Some(etcd_client) => Arc::new(State::new_with_etcd(model_manager, etcd_client)),
-            None => Arc::new(State::new(model_manager)),
-        };
+        let state = Arc::new(State::new(model_manager, config.store));
         state
             .flags
             .set(&EndpointType::Chat, config.enable_chat_endpoints);
@@ -422,11 +396,6 @@ impl HttpServiceConfigBuilder {
         self
     }
 
-    pub fn with_etcd_client(mut self, etcd_client: Option<etcd::Client>) -> Self {
-        self.etcd_client = Some(etcd_client);
-        self
-    }
-
     // DEPRECATED: To be removed after custom backends migrate to Dynamo backend.
     pub fn with_custom_backend_config(
         mut self,

lib/llm/src/kv_router/protocols.rs

@@ -6,7 +6,7 @@ use serde::{Deserialize, Serialize};
 use uuid::Uuid;
 
 /// A worker identifier.
-pub type WorkerId = i64;
+pub type WorkerId = u64;
 
 /// A data parallel rank identifier.
 pub type DpRank = u32;

lib/llm/src/kv_router/publisher.rs

@@ -97,7 +97,7 @@ pub struct KvEventPublisher {
 impl KvEventPublisher {
     pub fn new(
         component: Component,
-        worker_id: i64,
+        worker_id: u64,
         kv_block_size: u32,
         source_config: Option<KvEventSourceConfig>,
     ) -> Result<Self> {
@@ -174,7 +174,7 @@ impl Drop for KvEventPublisher {
 
 async fn start_event_processor<P: EventPublisher + Send + Sync + 'static>(
     publisher: P,
-    worker_id: i64,
+    worker_id: u64,
     cancellation_token: CancellationToken,
     mut rx: mpsc::UnboundedReceiver<KvCacheEvent>,
 ) {
@@ -801,7 +801,7 @@ impl WorkerMetricsPublisher {
     ///
     /// This task monitors metric changes (specifically kv_active_blocks and num_requests_waiting)
     /// and publishes stable metrics to NATS after they've been unchanged for 1ms.
-    fn start_nats_metrics_publishing(&self, namespace: Namespace, worker_id: i64) {
+    fn start_nats_metrics_publishing(&self, namespace: Namespace, worker_id: u64) {
         let nats_rx = self.rx.clone();
 
         tokio::spawn(async move {