feat: ETCD high availability client failover - lease watch resilience (#3950)

Signed-off-by: Jacky <[email protected]>
Co-authored-by: Ryan McCormick <[email protected]>

Author: kthui

lib/bindings/python/rust/planner.rs

@@ -479,7 +479,7 @@ impl InnerClient {
     /// Wait for a new scaling decision. Use `get` when this returns to fetch the values.
     async fn wait(&self) -> anyhow::Result<()> {
         let watcher = self.etcd_client.kv_watch_prefix(&self.key).await?;
-        let (_prefix, _watcher, mut receiver) = watcher.dissolve();
+        let (_prefix, mut receiver) = watcher.dissolve();
         tokio::select! {
             _ = receiver.recv() => {
                 Ok(())

lib/llm/src/disagg_router.rs

@@ -0,0 +1,250 @@
+// SPDX-FileCopyrightText: Copyright (c) 2024-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+// SPDX-License-Identifier: Apache-2.0
+
+use serde::{Deserialize, Serialize};
+use std::sync::{Arc, Mutex};
+use tokio::sync::watch;
+use tracing;
+
+use dynamo_runtime::DistributedRuntime;
+use dynamo_runtime::transports::etcd::WatchEvent;
+
+#[derive(Clone, Debug, Serialize, Deserialize)]
+pub struct DisaggRouterConf {
+    pub max_local_prefill_length: i32,
+}
+
+impl Default for DisaggRouterConf {
+    fn default() -> Self {
+        Self {
+            max_local_prefill_length: 1000,
+        }
+    }
+}
+
+impl DisaggRouterConf {
+    pub async fn from_etcd_with_watcher(
+        drt: Arc<DistributedRuntime>,
+        model_name: &str,
+    ) -> anyhow::Result<(Self, watch::Receiver<Self>)> {
+        let etcd_key = format!("public/components/disagg_router/models/chat/{}", model_name);
+
+        // Get the initial value if it exists
+        let Some(etcd_client) = drt.etcd_client() else {
+            anyhow::bail!("Static components don't have an etcd client");
+        };
+        let initial_config = match etcd_client.kv_get_prefix(&etcd_key).await {
+            Ok(kvs) => {
+                if let Some(kv) = kvs.first() {
+                    match serde_json::from_slice::<DisaggRouterConf>(kv.value()) {
+                        Ok(config) => {
+                            tracing::debug!(
+                                "Found initial config for key {}: {:?}",
+                                etcd_key,
+                                config
+                            );
+                            config
+                        }
+                        Err(e) => {
+                            tracing::warn!(
+                                "Failed to parse initial config for key {}: {}",
+                                etcd_key,
+                                e
+                            );
+                            DisaggRouterConf::default()
+                        }
+                    }
+                } else {
+                    tracing::debug!(
+                        "No initial config found for key {}, using default",
+                        etcd_key
+                    );
+                    DisaggRouterConf::default()
+                }
+            }
+            Err(e) => {
+                tracing::warn!("Error fetching initial config for key {}: {}", etcd_key, e);
+                DisaggRouterConf::default()
+            }
+        };
+
+        // Create watch channel for config updates
+        let (watch_tx, watch_rx) = watch::channel(initial_config.clone());
+
+        // Set up the watcher after getting the initial value
+        let prefix_watcher = etcd_client.kv_get_and_watch_prefix(&etcd_key).await?;
+        let (key, mut kv_event_rx) = prefix_watcher.dissolve();
+
+        // Spawn background task to watch for config changes
+        drt.runtime().secondary().spawn(async move {
+            tracing::info!("Starting config watcher for disagg router key: {}", key);
+
+            loop {
+                let kv_event = tokio::select! {
+                    _ = watch_tx.closed() => {
+                        tracing::debug!("All watchers have closed; shutting down config watcher for key: {}", key);
+                        break;
+                    }
+                    kv_event = kv_event_rx.recv() => {
+                        match kv_event {
+                            Some(kv_event) => kv_event,
+                            None => {
+                                tracing::debug!("Watch stream has closed; shutting down config watcher for key: {}", key);
+                                break;
+                            }
+                        }
+                    }
+                };
+
+                tracing::debug!("Received watch event for key {}", key);
+
+                match kv_event {
+                    WatchEvent::Put(kv) => {
+                        let val = serde_json::from_slice::<DisaggRouterConf>(kv.value());
+                        if let Ok(config) = val {
+                            tracing::info!("Config updated for key {}: {:?}", key, config);
+                            // Broadcast the update
+                            if watch_tx.send(config).is_err() {
+                                tracing::debug!("Unable to send watch updates; shutting down config watcher for key: {}", key);
+                                break;
+                            }
+                        } else {
+                            tracing::error!("Unable to parse router config for key {}", key);
+                            break;
+                        }
+                    }
+                    WatchEvent::Delete(_) => {
+                        tracing::warn!("Config key was deleted: {}", key);
+                        // Reset to default values
+                        if watch_tx.send(DisaggRouterConf::default()).is_err() {
+                            tracing::debug!("Unable to send watch updates; shutting down config watcher for key: {}", key);
+                            break;
+                        }
+                    }
+                }
+            }
+
+            tracing::debug!("Completed config watcher for key: {}", key);
+        });
+
+        Ok((initial_config, watch_rx))
+    }
+}
+
+#[derive(Clone)]
+pub struct DisaggregatedRouter {
+    max_local_prefill_length: Arc<Mutex<i32>>,
+    model_name: String,
+    config_watcher: Option<watch::Receiver<DisaggRouterConf>>,
+}
+
+impl DisaggregatedRouter {
+    pub fn new(max_local_prefill_length: i32, model_name: String) -> Self {
+        DisaggregatedRouter {
+            max_local_prefill_length: Arc::new(Mutex::new(max_local_prefill_length)),
+            model_name,
+            config_watcher: None,
+        }
+    }
+
+    pub async fn new_with_etcd_and_default(
+        drt: Arc<DistributedRuntime>,
+        model_name: String,
+        default_max_local_prefill_length: i32,
+    ) -> anyhow::Result<Self> {
+        let (mut config, watcher) =
+            DisaggRouterConf::from_etcd_with_watcher(drt, &model_name).await?;
+
+        // Use the provided default if no etcd value was found (when config is the default value)
+        if config.max_local_prefill_length == DisaggRouterConf::default().max_local_prefill_length {
+            config.max_local_prefill_length = default_max_local_prefill_length;
+        }
+
+        let router = Self {
+            max_local_prefill_length: Arc::new(Mutex::new(config.max_local_prefill_length)),
+            model_name: model_name.clone(),
+            config_watcher: Some(watcher),
+        };
+
+        // Start background task to watch for config updates
+        router.start_config_watcher();
+
+        Ok(router)
+    }
+
+    fn start_config_watcher(&self) {
+        if let Some(watcher) = self.config_watcher.clone() {
+            let mut watcher = watcher;
+            // Create a clone for the task
+            let model_name = self.model_name.clone();
+            let max_local_prefill_length = self.max_local_prefill_length.clone();
+
+            tokio::spawn(async move {
+                tracing::info!("Starting config update watcher for model: {}", model_name);
+
+                while watcher.changed().await.is_ok() {
+                    let config = watcher.borrow().clone();
+                    let new_value = config.max_local_prefill_length;
+
+                    // Update the value using the mutex
+                    let mut current_value = max_local_prefill_length.lock().unwrap();
+                    let old_value = *current_value;
+                    if old_value != new_value {
+                        *current_value = new_value;
+                        tracing::info!(
+                            "Applied config update for model {}: max_local_prefill_length changed from {} to {}",
+                            model_name,
+                            old_value,
+                            new_value
+                        );
+                    }
+                }
+
+                tracing::debug!("Config watcher closed for model: {}", model_name);
+            });
+        }
+    }
+
+    pub fn check_for_updates(&self) {
+        if let Some(watcher) = &self.config_watcher
+            && watcher.has_changed().unwrap_or(false)
+        {
+            let config = watcher.borrow().clone();
+            let new_value = config.max_local_prefill_length;
+
+            // Update the value using the mutex
+            let mut current_value = self.max_local_prefill_length.lock().unwrap();
+            let old_value = *current_value;
+            if old_value != new_value {
+                *current_value = new_value;
+                tracing::info!(
+                    "Applied config update for model {}: max_local_prefill_length changed from {} to {}",
+                    self.model_name,
+                    old_value,
+                    new_value
+                );
+            }
+        }
+    }
+
+    pub fn prefill_remote(&self, prefill_length: i32, prefix_hit_length: i32) -> bool {
+        // Check for updates before making the decision
+        self.check_for_updates();
+
+        // Get the current value from the mutex
+        let max_local_prefill_length = *self.max_local_prefill_length.lock().unwrap();
+
+        // schedule the request purely based on the prefill length
+        // TODO: apply math models and compare local vs remote prefill TTFT
+        prefill_length - prefix_hit_length > max_local_prefill_length
+    }
+
+    pub fn update_value(&self, max_local_prefill_length: i32) {
+        let mut current = self.max_local_prefill_length.lock().unwrap();
+        *current = max_local_prefill_length;
+    }
+
+    pub fn get_model_name(&self) -> &str {
+        &self.model_name
+    }
+}

lib/llm/src/kv_router/subscriber.rs

@@ -274,14 +274,14 @@ pub async fn start_kv_router_background(
     cleanup_orphaned_consumers(&mut nats_queue, &etcd_client, &component, &consumer_uuid).await;
 
     // Watch for router deletions to clean up orphaned consumers
-    let (_prefix_str, _watcher, mut router_replicas_rx) = etcd_client
+    let (_prefix_str, mut router_replicas_rx) = etcd_client
         .kv_get_and_watch_prefix(&format!("{}/", KV_ROUTERS_ROOT_PATH))
         .await?
         .dissolve();
 
     // Get the generate endpoint and watch for instance deletions
     let generate_endpoint = component.endpoint("generate");
-    let (_instance_prefix, _instance_watcher, mut instance_event_rx) = etcd_client
+    let (_instance_prefix, mut instance_event_rx) = etcd_client
         .kv_get_and_watch_prefix(generate_endpoint.etcd_root())
         .await?
         .dissolve();

lib/runtime/src/component/client.rs

@@ -213,7 +213,7 @@ impl Client {
             .kv_get_and_watch_prefix(endpoint.etcd_root())
             .await?;
 
-        let (prefix, _watcher, mut kv_event_rx) = prefix_watcher.dissolve();
+        let (prefix, mut kv_event_rx) = prefix_watcher.dissolve();
 
         let (watch_tx, watch_rx) = tokio::sync::watch::channel(vec![]);