fix: expand discovery interface to support model types (#4090)

mohammedabdulwahhab · web-flow · commit e1547e24f2e6 · 2025-11-06T22:25:12.000Z
Signed-off-by: mohammedabdulwahhab &lt;furkhan324@berkeley.edu&gt;
diff --git a/lib/runtime/src/discovery/mock.rs b/lib/runtime/src/discovery/mock.rs
@@ -46,37 +46,79 @@ impl MockDiscoveryClient {
 /// Helper function to check if an instance matches a discovery key query
 fn matches_key(instance: &DiscoveryInstance, key: &DiscoveryKey) -> bool {
     match (instance, key) {
-        (DiscoveryInstance::Endpoint { .. }, DiscoveryKey::AllEndpoints) => true,
+        // Endpoint matching
+        (DiscoveryInstance::Endpoint(_), DiscoveryKey::AllEndpoints) => true,
+        (DiscoveryInstance::Endpoint(inst), DiscoveryKey::NamespacedEndpoints { namespace }) => {
+            &inst.namespace == namespace
+        }
+        (
+            DiscoveryInstance::Endpoint(inst),
+            DiscoveryKey::ComponentEndpoints {
+                namespace,
+                component,
+            },
+        ) => &inst.namespace == namespace && &inst.component == component,
+        (
+            DiscoveryInstance::Endpoint(inst),
+            DiscoveryKey::Endpoint {
+                namespace,
+                component,
+                endpoint,
+            },
+        ) => {
+            &inst.namespace == namespace
+                && &inst.component == component
+                && &inst.endpoint == endpoint
+        }
+
+        // ModelCard matching
+        (DiscoveryInstance::ModelCard { .. }, DiscoveryKey::AllModelCards) => true,
         (
-            DiscoveryInstance::Endpoint {
-                namespace: ins_ns, ..
+            DiscoveryInstance::ModelCard {
+                namespace: inst_ns, ..
             },
-            DiscoveryKey::NamespacedEndpoints { namespace },
-        ) => ins_ns == namespace,
+            DiscoveryKey::NamespacedModelCards { namespace },
+        ) => inst_ns == namespace,
         (
-            DiscoveryInstance::Endpoint {
-                namespace: ins_ns,
-                component: ins_comp,
+            DiscoveryInstance::ModelCard {
+                namespace: inst_ns,
+                component: inst_comp,
                 ..
             },
-            DiscoveryKey::ComponentEndpoints {
+            DiscoveryKey::ComponentModelCards {
                 namespace,
                 component,
             },
-        ) => ins_ns == namespace && ins_comp == component,
+        ) => inst_ns == namespace && inst_comp == component,
         (
-            DiscoveryInstance::Endpoint {
-                namespace: ins_ns,
-                component: ins_comp,
-                endpoint: ins_ep,
+            DiscoveryInstance::ModelCard {
+                namespace: inst_ns,
+                component: inst_comp,
+                endpoint: inst_ep,
                 ..
             },
-            DiscoveryKey::Endpoint {
+            DiscoveryKey::EndpointModelCards {
                 namespace,
                 component,
                 endpoint,
             },
-        ) => ins_ns == namespace && ins_comp == component && ins_ep == endpoint,
+        ) => inst_ns == namespace && inst_comp == component && inst_ep == endpoint,
+
+        // Cross-type matches return false
+        (
+            DiscoveryInstance::Endpoint(_),
+            DiscoveryKey::AllModelCards
+            | DiscoveryKey::NamespacedModelCards { .. }
+            | DiscoveryKey::ComponentModelCards { .. }
+            | DiscoveryKey::EndpointModelCards { .. },
+        ) => false,
+        (
+            DiscoveryInstance::ModelCard { .. },
+            DiscoveryKey::AllEndpoints
+            | DiscoveryKey::NamespacedEndpoints { .. }
+            | DiscoveryKey::ComponentEndpoints { .. }
+            | DiscoveryKey::Endpoint { .. },
+        ) => false,
     }
 }
 
@@ -98,6 +140,15 @@ impl DiscoveryClient for MockDiscoveryClient {
         Ok(instance)
     }
 
+    async fn list(&self, key: DiscoveryKey) -> Result<Vec<DiscoveryInstance>> {
+        let instances = self.registry.instances.lock().unwrap();
+        Ok(instances
+            .iter()
+            .filter(|instance| matches_key(instance, &key))
+            .cloned()
+            .collect())
+    }
+
     async fn list_and_watch(&self, key: DiscoveryKey) -> Result<DiscoveryStream> {
         use std::collections::HashSet;
 
@@ -118,14 +169,16 @@ impl DiscoveryClient for MockDiscoveryClient {
 
                 let current_ids: HashSet<_> = current.iter().map(|i| {
                     match i {
-                        DiscoveryInstance::Endpoint { instance_id, .. } => *instance_id,
+                        DiscoveryInstance::Endpoint(inst) => inst.instance_id,
+                        DiscoveryInstance::ModelCard { instance_id, .. } => *instance_id,
                     }
                 }).collect();
 
                 // Emit Added events for new instances
                 for instance in current {
                     let id = match &instance {
-                        DiscoveryInstance::Endpoint { instance_id, .. } => *instance_id,
+                        DiscoveryInstance::Endpoint(inst) => inst.instance_id,
+                        DiscoveryInstance::ModelCard { instance_id, .. } => *instance_id,
                     };
                     if known_instances.insert(id) {
                         yield Ok(DiscoveryEvent::Added(instance));
@@ -161,6 +214,7 @@ mod tests {
             namespace: "test-ns".to_string(),
             component: "test-comp".to_string(),
             endpoint: "test-ep".to_string(),
+            transport: crate::component::TransportType::NatsTcp("test-subject".to_string()),
         };
 
         let key = DiscoveryKey::Endpoint {
@@ -177,8 +231,8 @@ mod tests {
 
         let event = stream.next().await.unwrap().unwrap();
         match event {
-            DiscoveryEvent::Added(DiscoveryInstance::Endpoint { instance_id, .. }) => {
-                assert_eq!(instance_id, 1);
+            DiscoveryEvent::Added(DiscoveryInstance::Endpoint(inst)) => {
+                assert_eq!(inst.instance_id, 1);
             }
             _ => panic!("Expected Added event for instance-1"),
         }
@@ -188,15 +242,16 @@ mod tests {
 
         let event = stream.next().await.unwrap().unwrap();
         match event {
-            DiscoveryEvent::Added(DiscoveryInstance::Endpoint { instance_id, .. }) => {
-                assert_eq!(instance_id, 2);
+            DiscoveryEvent::Added(DiscoveryInstance::Endpoint(inst)) => {
+                assert_eq!(inst.instance_id, 2);
             }
             _ => panic!("Expected Added event for instance-2"),
         }
 
         // Remove first instance
         registry.instances.lock().unwrap().retain(|i| match i {
-            DiscoveryInstance::Endpoint { instance_id, .. } => *instance_id != 1,
+            DiscoveryInstance::Endpoint(inst) => inst.instance_id != 1,
+            DiscoveryInstance::ModelCard { instance_id, .. } => *instance_id != 1,
         });
 
         let event = stream.next().await.unwrap().unwrap();
diff --git a/lib/runtime/src/discovery/mod.rs b/lib/runtime/src/discovery/mod.rs
@@ -2,6 +2,7 @@
 // SPDX-License-Identifier: Apache-2.0
 
 use crate::Result;
+use crate::component::TransportType;
 use async_trait::async_trait;
 use futures::Stream;
 use serde::{Deserialize, Serialize};
@@ -10,14 +11,19 @@ use std::pin::Pin;
 mod mock;
 pub use mock::{MockDiscoveryClient, SharedMockRegistry};
 
+pub mod utils;
+pub use utils::watch_and_extract_field;
+
 /// Query key for prefix-based discovery queries
 /// Supports hierarchical queries from all endpoints down to specific endpoints
 #[derive(Debug, Clone, PartialEq, Eq, Hash)]
 pub enum DiscoveryKey {
     /// Query all endpoints in the system
     AllEndpoints,
     /// Query all endpoints in a specific namespace
-    NamespacedEndpoints { namespace: String },
+    NamespacedEndpoints {
+        namespace: String,
+    },
     /// Query all endpoints in a namespace/component
     ComponentEndpoints {
         namespace: String,
@@ -29,59 +35,136 @@ pub enum DiscoveryKey {
         component: String,
         endpoint: String,
     },
-    // TODO: Extend to support ModelCard queries:
-    // - AllModels
-    // - NamespacedModels { namespace }
-    // - ComponentModels { namespace, component }
-    // - Model { namespace, component, model_name }
+    AllModelCards,
+    NamespacedModelCards {
+        namespace: String,
+    },
+    ComponentModelCards {
+        namespace: String,
+        component: String,
+    },
+    EndpointModelCards {
+        namespace: String,
+        component: String,
+        endpoint: String,
+    },
 }
 
 /// Specification for registering objects in the discovery plane
 /// Represents the input to the register() operation
-#[derive(Debug, Clone, PartialEq, Eq, Hash)]
+#[derive(Debug, Clone, PartialEq, Eq)]
 pub enum DiscoverySpec {
     /// Endpoint specification for registration
     Endpoint {
         namespace: String,
         component: String,
         endpoint: String,
+        /// Transport type and routing information
+        transport: TransportType,
+    },
+    ModelCard {
+        namespace: String,
+        component: String,
+        endpoint: String,
+        /// ModelDeploymentCard serialized as JSON
+        /// This allows lib/runtime to remain independent of lib/llm types
+        /// DiscoverySpec.from_model_card() and DiscoveryInstance.deserialize_model_card() are ergonomic helpers to create and deserialize the model card.
+        card_json: serde_json::Value,
     },
-    // TODO: Add ModelCard variant:
-    // - ModelCard { namespace, component, model_name, card: ModelDeploymentCard }
 }
 
 impl DiscoverySpec {
+    /// Creates a ModelCard discovery spec from a serializable type
+    /// The card will be serialized to JSON to avoid cross-crate dependencies
+    pub fn from_model_card<T>(
+        namespace: String,
+        component: String,
+        endpoint: String,
+        card: &T,
+    ) -> crate::Result<Self>
+    where
+        T: Serialize,
+    {
+        let card_json = serde_json::to_value(card)?;
+        Ok(Self::ModelCard {
+            namespace,
+            component,
+            endpoint,
+            card_json,
+        })
+    }
+
     /// Attaches an instance ID to create a DiscoveryInstance
     pub fn with_instance_id(self, instance_id: u64) -> DiscoveryInstance {
         match self {
             Self::Endpoint {
                 namespace,
                 component,
                 endpoint,
-            } => DiscoveryInstance::Endpoint {
+                transport,
+            } => DiscoveryInstance::Endpoint(crate::component::Instance {
+                namespace,
+                component,
+                endpoint,
+                instance_id,
+                transport,
+            }),
+            Self::ModelCard {
+                namespace,
+                component,
+                endpoint,
+                card_json,
+            } => DiscoveryInstance::ModelCard {
                 namespace,
                 component,
                 endpoint,
                 instance_id,
+                card_json,
             },
         }
     }
 }
 
 /// Registered instances in the discovery plane
 /// Represents objects that have been successfully registered with an instance ID
-#[derive(Debug, Clone, PartialEq, Eq, Hash, serde::Serialize, serde::Deserialize)]
+#[derive(Debug, Clone, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
 #[serde(tag = "type")]
 pub enum DiscoveryInstance {
-    /// Registered endpoint instance
-    Endpoint {
+    /// Registered endpoint instance - wraps the component::Instance directly
+    Endpoint(crate::component::Instance),
+    ModelCard {
         namespace: String,
         component: String,
         endpoint: String,
         instance_id: u64,
+        /// ModelDeploymentCard serialized as JSON
+        /// This allows lib/runtime to remain independent of lib/llm types
+        card_json: serde_json::Value,
     },
-    // TODO: Add ModelCard variant:
-    // - ModelCard { namespace, component, model_name, instance_id, card: ModelDeploymentCard }
+}
+
+impl DiscoveryInstance {
+    /// Returns the instance ID for this discovery instance
+    pub fn instance_id(&self) -> u64 {
+        match self {
+            Self::Endpoint(inst) => inst.instance_id,
+            Self::ModelCard { instance_id, .. } => *instance_id,
+        }
+    }
+
+    /// Deserializes the model card JSON into the specified type T
+    /// Returns an error if this is not a ModelCard instance or if deserialization fails
+    pub fn deserialize_model_card<T>(&self) -> crate::Result<T>
+    where
+        T: for<'de> Deserialize<'de>,
+    {
+        match self {
+            Self::ModelCard { card_json, .. } => Ok(serde_json::from_value(card_json.clone())?),
+            Self::Endpoint(_) => {
+                crate::raise!("Cannot deserialize model card from Endpoint instance")
+            }
+        }
+    }
 }
 
 /// Events emitted by the discovery client watch stream
@@ -106,6 +189,10 @@ pub trait DiscoveryClient: Send + Sync {
     /// Registers an object in the discovery plane with the instance id
     async fn register(&self, spec: DiscoverySpec) -> Result<DiscoveryInstance>;
 
+    /// Returns a list of currently registered instances for the given discovery key
+    /// This is a one-time snapshot without watching for changes
+    async fn list(&self, key: DiscoveryKey) -> Result<Vec<DiscoveryInstance>>;
+
     /// Returns a stream of discovery events (Added/Removed) for the given discovery key
     async fn list_and_watch(&self, key: DiscoveryKey) -> Result<DiscoveryStream>;
 }
diff --git a/lib/runtime/src/discovery/utils.rs b/lib/runtime/src/discovery/utils.rs
