fix: fix bug in diffing logic in list_and_watch (#5318)

Signed-off-by: mohammedabdulwahhab <[email protected]>

Author: mohammedabdulwahhab

lib/bindings/c/src/lib.rs

@@ -1306,12 +1306,12 @@ fn spawn_prefill_watcher(
                         }
                     }
                 }
-                DiscoveryEvent::Removed(instance_id) => {
+                DiscoveryEvent::Removed(id) => {
                     // Log removal for observability
                     // Note: The PrefillRouter remains active - worker availability
                     // is handled dynamically by the underlying Client's instance tracking
                     tracing::debug!(
-                        instance_id = instance_id,
+                        instance_id = id.instance_id(),
                         "Prefill worker instance removed from discovery"
                     );
                 }

lib/llm/src/discovery/model_manager.rs

@@ -320,14 +320,14 @@ impl ModelManager {
             .ok_or(ModelManagerError::ModelNotFound(model.to_string()))
     }
 
-    /// Save a ModelDeploymentCard from an instance's ModelDeploymentCard key so we can fetch it later when the key is
+    /// Save a ModelDeploymentCard from an instance's key so we can fetch it later when the key is
     /// deleted.
     pub fn save_model_card(&self, key: &str, card: ModelDeploymentCard) -> anyhow::Result<()> {
         self.cards.lock().insert(key.to_string(), card);
         Ok(())
     }
 
-    /// Remove and return model card for this instance's etcd key. We do this when the instance stops.
+    /// Remove and return model card for this instance's key. We do this when the instance stops.
     pub fn remove_model_card(&self, key: &str) -> Option<ModelDeploymentCard> {
         self.cards.lock().remove(key)
     }

lib/llm/src/discovery/watcher.rs

@@ -10,7 +10,10 @@ use futures::StreamExt;
 
 use dynamo_runtime::{
     DistributedRuntime,
-    discovery::{DiscoveryEvent, DiscoveryInstance, DiscoveryQuery, DiscoveryStream},
+    discovery::{
+        DiscoveryEvent, DiscoveryInstance, DiscoveryInstanceId, DiscoveryQuery, DiscoveryStream,
+        ModelCardInstanceId,
+    },
     pipeline::{
         ManyOut, Operator, RouterMode, SegmentSource, ServiceBackend, SingleIn, Source,
         network::egress::push_router::PushRouter,
@@ -119,23 +122,26 @@ impl ModelWatcher {
 
             match event {
                 DiscoveryEvent::Added(instance) => {
-                    // Extract EndpointId, instance_id, and card from the discovery instance
-                    let (endpoint_id, instance_id, mut card) = match &instance {
+                    // Extract ModelCardInstanceId and card from the discovery instance
+                    let (mcid, mut card) = match &instance {
                         DiscoveryInstance::Model {
                             namespace,
                             component,
                             endpoint,
                             instance_id,
+                            model_suffix,
                             ..
                         } => {
-                            let eid = EndpointId {
+                            let mcid = ModelCardInstanceId {
                                 namespace: namespace.clone(),
                                 component: component.clone(),
-                                name: endpoint.clone(),
+                                endpoint: endpoint.clone(),
+                                instance_id: *instance_id,
+                                model_suffix: model_suffix.clone(),
                             };
 
                             match instance.deserialize_model::<ModelDeploymentCard>() {
-                                Ok(card) => (eid, *instance_id, card),
+                                Ok(card) => (mcid, card),
                                 Err(err) => {
                                     tracing::error!(%err, instance_id, "Failed to deserialize model card");
                                     continue;
@@ -153,10 +159,10 @@ impl ModelWatcher {
                     // Filter by namespace if target_namespace is specified
                     if !global_namespace
                         && let Some(target_ns) = target_namespace
-                        && endpoint_id.namespace != target_ns
+                        && mcid.namespace != target_ns
                     {
                         tracing::debug!(
-                            model_namespace = endpoint_id.namespace,
+                            model_namespace = mcid.namespace,
                             target_namespace = target_ns,
                             "Skipping model from different namespace"
                         );
@@ -185,34 +191,39 @@ impl ModelWatcher {
                         continue;
                     }
 
-                    // Use instance_id as the HashMap key (simpler and sufficient since keys are opaque)
-                    let key = format!("{:x}", instance_id);
-
-                    match self.handle_put(&key, &endpoint_id, &mut card).await {
+                    match self.handle_put(&mcid, &mut card).await {
                         Ok(()) => {
                             tracing::info!(
                                 model_name = card.name(),
-                                namespace = endpoint_id.namespace,
+                                namespace = mcid.namespace,
                                 "added model"
                             );
                             self.notify_on_model.notify_waiters();
                         }
                         Err(err) => {
                             tracing::error!(
                                 model_name = card.name(),
-                                namespace = endpoint_id.namespace,
+                                namespace = mcid.namespace,
                                 error = format!("{err:#}"),
                                 "Error adding model from discovery",
                             );
                         }
                     }
                 }
-                DiscoveryEvent::Removed(instance_id) => {
-                    // Use instance_id hex as the HashMap key (matches what we saved with)
-                    let key = format!("{:x}", instance_id);
+                DiscoveryEvent::Removed(id) => {
+                    // Extract ModelCardInstanceId from the removal event
+                    let model_card_instance_id = match &id {
+                        DiscoveryInstanceId::Model(mcid) => mcid,
+                        DiscoveryInstanceId::Endpoint(_) => {
+                            tracing::error!(
+                                "Unexpected discovery instance type in removal (expected Model)"
+                            );
+                            continue;
+                        }
+                    };
 
                     match self
-                        .handle_delete(&key, target_namespace, global_namespace)
+                        .handle_delete(model_card_instance_id, target_namespace, global_namespace)
                         .await
                     {
                         Ok(Some(model_name)) => {
@@ -234,14 +245,15 @@ impl ModelWatcher {
     /// Returns the name of the model we just deleted, if any.
     async fn handle_delete(
         &self,
-        key: &str,
+        mcid: &ModelCardInstanceId,
         target_namespace: Option<&str>,
         is_global_namespace: bool,
     ) -> anyhow::Result<Option<String>> {
-        let card = match self.manager.remove_model_card(key) {
+        let key = mcid.to_path();
+        let card = match self.manager.remove_model_card(&key) {
             Some(card) => card,
             None => {
-                anyhow::bail!("Missing ModelDeploymentCard for {key}");
+                anyhow::bail!("Missing ModelDeploymentCard for {}", key);
             }
         };
         let model_name = card.name().to_string();
@@ -325,20 +337,20 @@ impl ModelWatcher {
     // models.
     async fn handle_put(
         &self,
-        key: &str,
-        endpoint_id: &EndpointId,
+        mcid: &ModelCardInstanceId,
         card: &mut ModelDeploymentCard,
     ) -> anyhow::Result<()> {
         card.download_config().await?;
 
         let component = self
             .drt
-            .namespace(&endpoint_id.namespace)?
-            .component(&endpoint_id.component)?;
-        let endpoint = component.endpoint(&endpoint_id.name);
+            .namespace(&mcid.namespace)?
+            .component(&mcid.component)?;
+        let endpoint = component.endpoint(&mcid.endpoint);
         let client = endpoint.client().await?;
         tracing::debug!(model_name = card.name(), "adding model");
-        self.manager.save_model_card(key, card.clone())?;
+        self.manager
+            .save_model_card(&mcid.to_path(), card.clone())?;
 
         // Skip duplicate registrations based on model type.
         // Prefill and decode models are tracked separately, so registering one
@@ -352,7 +364,7 @@ impl ModelWatcher {
         if already_registered {
             tracing::debug!(
                 model_name = card.name(),
-                namespace = endpoint_id.namespace,
+                namespace = mcid.namespace,
                 model_type = %card.model_type,
                 "Model already registered, skipping"
             );
@@ -372,7 +384,7 @@ impl ModelWatcher {
             // A model that expects pre-processed requests meaning it's up to us whether we
             // handle Chat or Completions requests, so handle whatever the model supports.
 
-            let endpoint = component.endpoint(&endpoint_id.name);
+            let endpoint = component.endpoint(&mcid.endpoint);
             let kv_chooser = if self.router_config.router_mode == RouterMode::KV {
                 Some(
                     self.manager

lib/llm/src/kv_router/subscriber.rs

@@ -564,10 +564,12 @@ pub async fn start_kv_router_background(
                         continue;
                     };
 
-                    let DiscoveryEvent::Removed(worker_id) = discovery_event else {
+                    let DiscoveryEvent::Removed(id) = discovery_event else {
                         continue;
                     };
 
+                    let worker_id = id.instance_id();
+
                     tracing::warn!(
                         "DISCOVERY: Generate endpoint instance removed, removing worker {worker_id}"
                     );
@@ -642,11 +644,13 @@ pub async fn start_kv_router_background(
                         continue;
                     };
 
-                    let DiscoveryEvent::Removed(router_instance_id) = router_event else {
+                    let DiscoveryEvent::Removed(id) = router_event else {
                         // We only care about removals for cleaning up consumers
                         continue;
                     };
 
+                    let router_instance_id = id.instance_id();
+
                     // The consumer UUID is the instance_id in hex format
                     let consumer_to_delete = router_instance_id.to_string();
 
@@ -708,7 +712,8 @@ async fn handle_worker_discovery(
                 }
             }
         }
-        DiscoveryEvent::Removed(worker_id) => {
+        DiscoveryEvent::Removed(id) => {
+            let worker_id = id.instance_id();
             tracing::warn!("DISCOVERY: Worker {worker_id} removed, removing from router indexer");
 
             if let Err(e) = remove_worker_tx.send(worker_id).await {

lib/runtime/src/component/client.rs

@@ -9,7 +9,7 @@ use arc_swap::ArcSwap;
 use futures::StreamExt;
 use tokio::net::unix::pipe::Receiver;
 
-use crate::discovery::{DiscoveryEvent, DiscoveryInstance};
+use crate::discovery::{DiscoveryEvent, DiscoveryInstance, DiscoveryInstanceId};
 use crate::{
     component::{Endpoint, Instance},
     pipeline::async_trait,
@@ -255,8 +255,8 @@ impl Client {
                                 map.insert(instance.instance_id, instance);
                         }
                     }
-                    DiscoveryEvent::Removed(instance_id) => {
-                        map.remove(&instance_id);
+                    DiscoveryEvent::Removed(id) => {
+                        map.remove(&id.instance_id());
                     }
                 }