dynamo_runtime/discovery/
metadata.rs

1// SPDX-FileCopyrightText: Copyright (c) 2024-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
2// SPDX-License-Identifier: Apache-2.0
3
4use anyhow::Result;
5use serde::Deserialize as _;
6use std::collections::{HashMap, HashSet};
7use std::sync::Arc;
8
9use super::{
10    DiscoveryInstance, DiscoveryInstanceId, DiscoveryQuery, ModelCardInstanceId,
11    model_with_updated_taints, validate_event_source_reregistration, validate_model_reregistration,
12};
13
14/// Deserializes a JSON `null` or missing field as `T::default()`.
15///
16/// Kubernetes Server-Side Apply with `schema = "disabled"` can write an empty
17/// object `{}` as `null` for nested free-form fields. Without this helper, the
18/// daemon fails to deserialize the `DynamoWorkerMetadata` CR, and the worker is
19/// excluded from the `MetadataSnapshot` (i.e. invisible to service discovery),
20/// causing `KubeDiscoveryClient::list` to return 0 instances and all inference
21/// requests to 404. One concrete example is vLLM elastic EP scaling:
22/// `scale_elastic_ep` reinitializes event plane sockets, which triggers
23/// `unregister_event_channel()`, leaving `event_channels` as an empty map `{}`.
24/// SSA then writes it back as `null`, breaking deserialization until this helper
25/// treats `null` as an empty map. The issue applies to any event plane
26/// implementation, not only a specific transport.
27fn deserialize_null_default<'de, D, T>(deserializer: D) -> Result<T, D::Error>
28where
29    D: serde::Deserializer<'de>,
30    T: Default + serde::Deserialize<'de>,
31{
32    Ok(Option::<T>::deserialize(deserializer)?.unwrap_or_default())
33}
34
35/// Metadata stored on each pod and exposed via HTTP endpoint
36/// This struct holds all discovery registrations for this pod instance
37#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
38pub struct DiscoveryMetadata {
39    /// Registered endpoint instances (key: path string from EndpointInstanceId::to_path())
40    #[serde(default, deserialize_with = "deserialize_null_default")]
41    endpoints: HashMap<String, DiscoveryInstance>,
42    /// Registered model card instances (key: path string from ModelCardInstanceId::to_path())
43    #[serde(default, deserialize_with = "deserialize_null_default")]
44    model_cards: HashMap<String, DiscoveryInstance>,
45    /// Registered event channel instances (key: path string from EventChannelInstanceId::to_path())
46    #[serde(default, deserialize_with = "deserialize_null_default")]
47    event_channels: HashMap<String, DiscoveryInstance>,
48    /// Registered event source instances (key: path string from EventSourceInstanceId::to_path())
49    #[serde(default, deserialize_with = "deserialize_null_default")]
50    event_sources: HashMap<String, DiscoveryInstance>,
51}
52
53impl DiscoveryMetadata {
54    /// Create a new empty metadata store
55    pub fn new() -> Self {
56        Self {
57            endpoints: HashMap::new(),
58            model_cards: HashMap::new(),
59            event_channels: HashMap::new(),
60            event_sources: HashMap::new(),
61        }
62    }
63
64    /// Register an endpoint instance
65    pub fn register_endpoint(&mut self, instance: DiscoveryInstance) -> Result<()> {
66        match instance.id() {
67            DiscoveryInstanceId::Endpoint(key) => {
68                self.endpoints.insert(key.to_path(), instance);
69                Ok(())
70            }
71            DiscoveryInstanceId::Model(_) => {
72                anyhow::bail!("Cannot register non-endpoint instance as endpoint")
73            }
74            DiscoveryInstanceId::EventChannel(_) => {
75                anyhow::bail!("Cannot register EventChannel instance as endpoint")
76            }
77            DiscoveryInstanceId::EventSource(_) => {
78                anyhow::bail!("Cannot register EventSource instance as endpoint")
79            }
80        }
81    }
82
83    /// Register a model card instance
84    pub fn register_model_card(
85        &mut self,
86        instance: DiscoveryInstance,
87    ) -> Result<DiscoveryInstance> {
88        match instance.id() {
89            DiscoveryInstanceId::Model(key) => match self.model_cards.entry(key.to_path()) {
90                std::collections::hash_map::Entry::Vacant(entry) => {
91                    entry.insert(instance.clone());
92                    Ok(instance)
93                }
94                std::collections::hash_map::Entry::Occupied(entry) => {
95                    validate_model_reregistration(entry.get(), &instance)?;
96                    Ok(entry.get().clone())
97                }
98            },
99            DiscoveryInstanceId::Endpoint(_) => {
100                anyhow::bail!("Cannot register non-model-card instance as model card")
101            }
102            DiscoveryInstanceId::EventChannel(_) => {
103                anyhow::bail!("Cannot register EventChannel instance as model card")
104            }
105            DiscoveryInstanceId::EventSource(_) => {
106                anyhow::bail!("Cannot register EventSource instance as model card")
107            }
108        }
109    }
110
111    /// Update one authoritative model card while the caller holds the metadata lock.
112    pub fn update_model_taints(
113        &mut self,
114        id: &ModelCardInstanceId,
115        taints: HashSet<String>,
116    ) -> Result<bool> {
117        let path = id.to_path();
118        let existing = self
119            .model_cards
120            .get_mut(&path)
121            .ok_or_else(|| anyhow::anyhow!("model discovery record {path} is not registered"))?;
122        let candidate = model_with_updated_taints(existing, taints)?;
123        if candidate == *existing {
124            return Ok(false);
125        }
126        *existing = candidate;
127        Ok(true)
128    }
129
130    /// Unregister an endpoint instance
131    pub fn unregister_endpoint(&mut self, instance: &DiscoveryInstance) -> Result<()> {
132        match instance.id() {
133            DiscoveryInstanceId::Endpoint(key) => {
134                self.endpoints.remove(&key.to_path());
135                Ok(())
136            }
137            DiscoveryInstanceId::Model(_) => {
138                anyhow::bail!("Cannot unregister non-endpoint instance as endpoint")
139            }
140            DiscoveryInstanceId::EventChannel(_) => {
141                anyhow::bail!("Cannot unregister EventChannel instance as endpoint")
142            }
143            DiscoveryInstanceId::EventSource(_) => {
144                anyhow::bail!("Cannot unregister EventSource instance as endpoint")
145            }
146        }
147    }
148
149    /// Unregister a model card instance
150    pub fn unregister_model_card(&mut self, instance: &DiscoveryInstance) -> Result<()> {
151        match instance.id() {
152            DiscoveryInstanceId::Model(key) => {
153                self.model_cards.remove(&key.to_path());
154                Ok(())
155            }
156            DiscoveryInstanceId::Endpoint(_) => {
157                anyhow::bail!("Cannot unregister non-model-card instance as model card")
158            }
159            DiscoveryInstanceId::EventChannel(_) => {
160                anyhow::bail!("Cannot unregister EventChannel instance as model card")
161            }
162            DiscoveryInstanceId::EventSource(_) => {
163                anyhow::bail!("Cannot unregister EventSource instance as model card")
164            }
165        }
166    }
167
168    /// Register an event channel instance
169    pub fn register_event_channel(&mut self, instance: DiscoveryInstance) -> Result<()> {
170        match instance.id() {
171            DiscoveryInstanceId::EventChannel(key) => {
172                self.event_channels.insert(key.to_path(), instance);
173                Ok(())
174            }
175            DiscoveryInstanceId::Endpoint(_) => {
176                anyhow::bail!("Cannot register Endpoint instance as event channel")
177            }
178            DiscoveryInstanceId::Model(_) => {
179                anyhow::bail!("Cannot register Model instance as event channel")
180            }
181            DiscoveryInstanceId::EventSource(_) => {
182                anyhow::bail!("Cannot register EventSource instance as event channel")
183            }
184        }
185    }
186
187    /// Unregister an event channel instance
188    pub fn unregister_event_channel(&mut self, instance: &DiscoveryInstance) -> Result<()> {
189        match instance.id() {
190            DiscoveryInstanceId::EventChannel(key) => {
191                self.event_channels.remove(&key.to_path());
192                Ok(())
193            }
194            DiscoveryInstanceId::Endpoint(_) => {
195                anyhow::bail!("Cannot unregister Endpoint instance as event channel")
196            }
197            DiscoveryInstanceId::Model(_) => {
198                anyhow::bail!("Cannot unregister Model instance as event channel")
199            }
200            DiscoveryInstanceId::EventSource(_) => {
201                anyhow::bail!("Cannot unregister EventSource instance as event channel")
202            }
203        }
204    }
205
206    /// Register a semantic event source instance.
207    pub fn register_event_source(&mut self, instance: DiscoveryInstance) -> Result<()> {
208        match instance.id() {
209            DiscoveryInstanceId::EventSource(key) => {
210                let path = key.to_path();
211                match self.event_sources.entry(path) {
212                    std::collections::hash_map::Entry::Vacant(entry) => {
213                        entry.insert(instance);
214                        Ok(())
215                    }
216                    std::collections::hash_map::Entry::Occupied(entry) => {
217                        validate_event_source_reregistration(entry.get(), &instance)
218                    }
219                }
220            }
221            DiscoveryInstanceId::Endpoint(_) => {
222                anyhow::bail!("Cannot register Endpoint instance as event source")
223            }
224            DiscoveryInstanceId::Model(_) => {
225                anyhow::bail!("Cannot register Model instance as event source")
226            }
227            DiscoveryInstanceId::EventChannel(_) => {
228                anyhow::bail!("Cannot register EventChannel instance as event source")
229            }
230        }
231    }
232
233    /// Unregister one exact semantic event source incarnation.
234    pub fn unregister_event_source(&mut self, instance: &DiscoveryInstance) -> Result<()> {
235        match instance.id() {
236            DiscoveryInstanceId::EventSource(key) => {
237                self.event_sources.remove(&key.to_path());
238                Ok(())
239            }
240            DiscoveryInstanceId::Endpoint(_) => {
241                anyhow::bail!("Cannot unregister Endpoint instance as event source")
242            }
243            DiscoveryInstanceId::Model(_) => {
244                anyhow::bail!("Cannot unregister Model instance as event source")
245            }
246            DiscoveryInstanceId::EventChannel(_) => {
247                anyhow::bail!("Cannot unregister EventChannel instance as event source")
248            }
249        }
250    }
251
252    /// Get all registered endpoints
253    pub fn get_all_endpoints(&self) -> Vec<DiscoveryInstance> {
254        self.endpoints.values().cloned().collect()
255    }
256
257    /// Get all registered model cards
258    pub fn get_all_model_cards(&self) -> Vec<DiscoveryInstance> {
259        self.model_cards.values().cloned().collect()
260    }
261
262    /// Get all registered event channels
263    pub fn get_all_event_channels(&self) -> Vec<DiscoveryInstance> {
264        self.event_channels.values().cloned().collect()
265    }
266
267    /// Get all registered semantic event sources.
268    pub fn get_all_event_sources(&self) -> Vec<DiscoveryInstance> {
269        self.event_sources.values().cloned().collect()
270    }
271
272    /// Get all registered instances (endpoints, model cards, and event channels)
273    pub fn get_all(&self) -> Vec<DiscoveryInstance> {
274        self.endpoints
275            .values()
276            .chain(self.model_cards.values())
277            .chain(self.event_channels.values())
278            .chain(self.event_sources.values())
279            .cloned()
280            .collect()
281    }
282
283    /// Filter this metadata by query
284    pub fn filter(&self, query: &DiscoveryQuery) -> Vec<DiscoveryInstance> {
285        let all_instances = match query {
286            DiscoveryQuery::AllEndpoints
287            | DiscoveryQuery::NamespacedEndpoints { .. }
288            | DiscoveryQuery::ComponentEndpoints { .. }
289            | DiscoveryQuery::Endpoint { .. } => self.get_all_endpoints(),
290
291            DiscoveryQuery::AllModels
292            | DiscoveryQuery::NamespacedModels { .. }
293            | DiscoveryQuery::ComponentModels { .. }
294            | DiscoveryQuery::EndpointModels { .. } => self.get_all_model_cards(),
295
296            // EventChannel queries now return actual event channels
297            DiscoveryQuery::EventChannels(_) => self.get_all_event_channels(),
298            DiscoveryQuery::EventSources(_) => self.get_all_event_sources(),
299        };
300
301        filter_instances(all_instances, query)
302    }
303}
304
305impl Default for DiscoveryMetadata {
306    fn default() -> Self {
307        Self::new()
308    }
309}
310
311/// Filter instances by query predicate
312fn filter_instances(
313    instances: Vec<DiscoveryInstance>,
314    query: &DiscoveryQuery,
315) -> Vec<DiscoveryInstance> {
316    instances.into_iter().filter(|i| i.matches(query)).collect()
317}
318
319/// Snapshot of all discovered instances and their metadata
320#[derive(Clone, Debug)]
321pub struct MetadataSnapshot {
322    /// Map of instance_id -> metadata
323    pub instances: HashMap<u64, Arc<DiscoveryMetadata>>,
324    /// Map of instance_id -> CR generation for change detection
325    /// Keys match `instances` keys exactly - only ready pods with CRs are included
326    pub generations: HashMap<u64, i64>,
327    /// Sequence number for debugging
328    pub sequence: u64,
329    /// Timestamp for observability
330    pub timestamp: std::time::Instant,
331}
332
333impl MetadataSnapshot {
334    pub fn empty() -> Self {
335        Self {
336            instances: HashMap::new(),
337            generations: HashMap::new(),
338            sequence: 0,
339            timestamp: std::time::Instant::now(),
340        }
341    }
342
343    /// Compare with previous snapshot and return true if changed.
344    /// Logs diagnostic info about what changed.
345    /// This is done on the basis of the generation of the DynamoWorkerMetadata CRs that are owned by ready workers
346    pub fn has_changes_from(&self, prev: &MetadataSnapshot) -> bool {
347        if self.generations == prev.generations {
348            tracing::trace!(
349                "Snapshot (seq={}): no changes, {} instances",
350                self.sequence,
351                self.instances.len()
352            );
353            return false;
354        }
355
356        // Compute diff for logging
357        let curr_ids: HashSet<u64> = self.generations.keys().copied().collect();
358        let prev_ids: HashSet<u64> = prev.generations.keys().copied().collect();
359
360        let added: Vec<_> = curr_ids
361            .difference(&prev_ids)
362            .map(|id| format!("{:x}", id))
363            .collect();
364        let removed: Vec<_> = prev_ids
365            .difference(&curr_ids)
366            .map(|id| format!("{:x}", id))
367            .collect();
368        let updated: Vec<_> = self
369            .generations
370            .iter()
371            .filter(|(k, v)| prev.generations.get(*k).is_some_and(|pv| pv != *v))
372            .map(|(k, _)| format!("{:x}", k))
373            .collect();
374
375        tracing::info!(
376            "Snapshot (seq={}): {} instances, added={:?}, removed={:?}, updated={:?}",
377            self.sequence,
378            self.instances.len(),
379            added,
380            removed,
381            updated
382        );
383
384        true
385    }
386
387    /// Filter all instances in the snapshot by query
388    pub fn filter(&self, query: &DiscoveryQuery) -> Vec<DiscoveryInstance> {
389        self.instances
390            .values()
391            .flat_map(|metadata| metadata.filter(query))
392            .collect()
393    }
394}
395
396#[cfg(test)]
397mod tests {
398    use super::*;
399    use crate::component::{Instance, TransportType};
400    use crate::discovery::{EventChannelQuery, EventSourceQuery};
401
402    #[test]
403    fn authoritative_model_taint_updates_do_not_read_stale_snapshots() {
404        let mut metadata = DiscoveryMetadata::new();
405        let model = DiscoveryInstance::Model {
406            namespace: "ns".to_string(),
407            component: "worker".to_string(),
408            endpoint: "generate".to_string(),
409            instance_id: 7,
410            card_json: serde_json::json!({
411                "runtime_config": {"taints": ["a"]}
412            }),
413            model_suffix: None,
414        };
415        let DiscoveryInstanceId::Model(id) = model.id() else {
416            unreachable!()
417        };
418        metadata.register_model_card(model.clone()).unwrap();
419
420        assert!(
421            metadata
422                .update_model_taints(&id, HashSet::from(["b".to_string()]))
423                .unwrap()
424        );
425        assert!(
426            metadata
427                .update_model_taints(&id, HashSet::from(["a".to_string()]))
428                .unwrap()
429        );
430
431        let stored = metadata.get_all_model_cards().pop().unwrap();
432        let DiscoveryInstance::Model { card_json, .. } = stored else {
433            unreachable!()
434        };
435        assert_eq!(
436            card_json["runtime_config"]["taints"],
437            serde_json::json!(["a"])
438        );
439
440        metadata.unregister_model_card(&model).unwrap();
441        let error = metadata
442            .update_model_taints(&id, HashSet::from(["b".to_string()]))
443            .unwrap_err();
444        assert!(error.to_string().contains("is not registered"));
445    }
446
447    #[test]
448    fn same_id_model_registration_preserves_authoritative_taints() {
449        let mut metadata = DiscoveryMetadata::new();
450        let model = DiscoveryInstance::Model {
451            namespace: "ns".to_string(),
452            component: "worker".to_string(),
453            endpoint: "generate".to_string(),
454            instance_id: 7,
455            card_json: serde_json::json!({
456                "display_name": "model",
457                "runtime_config": {"taints": ["initial"]}
458            }),
459            model_suffix: None,
460        };
461        let DiscoveryInstanceId::Model(id) = model.id() else {
462            unreachable!()
463        };
464        metadata.register_model_card(model.clone()).unwrap();
465        metadata
466            .update_model_taints(&id, HashSet::from(["updated".to_string()]))
467            .unwrap();
468
469        let replayed = metadata.register_model_card(model).unwrap();
470
471        let DiscoveryInstance::Model { card_json, .. } = replayed else {
472            unreachable!()
473        };
474        assert_eq!(
475            card_json["runtime_config"]["taints"],
476            serde_json::json!(["updated"])
477        );
478        assert_eq!(metadata.get_all_model_cards().len(), 1);
479    }
480
481    #[test]
482    fn test_metadata_serde() {
483        let mut metadata = DiscoveryMetadata::new();
484
485        // Add an endpoint
486        let instance = DiscoveryInstance::Endpoint(Instance {
487            namespace: "test".to_string(),
488            component: "comp1".to_string(),
489            endpoint: "ep1".to_string(),
490            instance_id: 123,
491            transport: TransportType::Nats("nats://localhost:4222".to_string()),
492            device_type: None,
493            request_plane_codec: None,
494        });
495
496        metadata.register_endpoint(instance).unwrap();
497
498        // Serialize
499        let json = serde_json::to_string(&metadata).unwrap();
500
501        // Deserialize
502        let deserialized: DiscoveryMetadata = serde_json::from_str(&json).unwrap();
503
504        assert_eq!(deserialized.endpoints.len(), 1);
505        assert_eq!(deserialized.model_cards.len(), 0);
506    }
507
508    #[tokio::test]
509    async fn test_metadata_accessors() {
510        let mut metadata = DiscoveryMetadata::new();
511
512        // Register endpoints
513        for i in 0..3 {
514            let instance = DiscoveryInstance::Endpoint(Instance {
515                namespace: "test".to_string(),
516                component: "comp1".to_string(),
517                endpoint: format!("ep{}", i),
518                instance_id: i,
519                transport: TransportType::Nats("nats://localhost:4222".to_string()),
520                device_type: None,
521                request_plane_codec: None,
522            });
523            metadata.register_endpoint(instance).unwrap();
524        }
525
526        // Register model cards
527        for i in 0..2 {
528            let instance = DiscoveryInstance::Model {
529                namespace: "test".to_string(),
530                component: "comp1".to_string(),
531                endpoint: format!("ep{}", i),
532                instance_id: i,
533                card_json: serde_json::json!({"model": "test"}),
534                model_suffix: None,
535            };
536            metadata.register_model_card(instance).unwrap();
537        }
538
539        assert_eq!(metadata.get_all_endpoints().len(), 3);
540        assert_eq!(metadata.get_all_model_cards().len(), 2);
541        assert_eq!(metadata.get_all().len(), 5);
542    }
543
544    #[test]
545    fn event_source_registration_filters_and_removes_exact_incarnation() {
546        use crate::discovery::EventScope;
547        use crate::protocols::EndpointId;
548
549        let mut metadata = DiscoveryMetadata::new();
550        let endpoint = EndpointId {
551            namespace: "test".to_string(),
552            component: "worker".to_string(),
553            name: "decode".to_string(),
554        };
555        let source = |publisher_id| DiscoveryInstance::EventSource {
556            scope: EventScope::Endpoint {
557                endpoint: endpoint.clone(),
558            },
559            topic: "kv-events".to_string(),
560            publisher_id,
561            metadata: serde_json::json!({"worker_id": 7, "dp_rank": 0}),
562        };
563
564        let old = source(100);
565        let current = source(205);
566        metadata.register_event_source(old.clone()).unwrap();
567        metadata.register_event_source(old.clone()).unwrap();
568        let mut conflicting = old.clone();
569        let DiscoveryInstance::EventSource {
570            metadata: descriptor,
571            ..
572        } = &mut conflicting
573        else {
574            unreachable!()
575        };
576        *descriptor = serde_json::json!({"worker_id": 7, "dp_rank": 0, "changed": true});
577        assert!(
578            metadata
579                .register_event_source(conflicting)
580                .unwrap_err()
581                .to_string()
582                .contains("cannot change its descriptor")
583        );
584        metadata.register_event_source(current.clone()).unwrap();
585
586        let query =
587            DiscoveryQuery::EventSources(EventSourceQuery::endpoint_topic(endpoint, "kv-events"));
588        assert_eq!(metadata.filter(&query).len(), 2);
589
590        metadata.unregister_event_source(&old).unwrap();
591        assert_eq!(metadata.filter(&query), vec![current]);
592    }
593
594    #[test]
595    fn event_channel_queries_match_exact_scope() {
596        use crate::discovery::{EventScope, EventTransport};
597        use crate::protocols::EndpointId;
598
599        let mut metadata = DiscoveryMetadata::new();
600        let component_scope = EventScope::Component {
601            namespace: "test".to_string(),
602            component: "worker".to_string(),
603        };
604        let endpoint_a = EndpointId {
605            namespace: "test".to_string(),
606            component: "worker".to_string(),
607            name: "a".to_string(),
608        };
609        let endpoint_b = EndpointId {
610            name: "b".to_string(),
611            ..endpoint_a.clone()
612        };
613
614        for (instance_id, scope) in [
615            (1, component_scope),
616            (
617                2,
618                EventScope::Endpoint {
619                    endpoint: endpoint_a.clone(),
620                },
621            ),
622            (
623                3,
624                EventScope::Endpoint {
625                    endpoint: endpoint_b.clone(),
626                },
627            ),
628        ] {
629            metadata
630                .register_event_channel(DiscoveryInstance::EventChannel {
631                    scope,
632                    topic: "kv-events".to_string(),
633                    instance_id,
634                    transport: EventTransport::zmq(format!("tcp://localhost:{instance_id}")),
635                })
636                .unwrap();
637        }
638
639        assert_eq!(metadata.get_all_event_channels().len(), 3);
640        assert_eq!(metadata.get_all().len(), 3);
641        assert_eq!(
642            metadata
643                .filter(&DiscoveryQuery::EventChannels(EventChannelQuery::all()))
644                .len(),
645            3
646        );
647
648        let component = metadata.filter(&DiscoveryQuery::EventChannels(EventChannelQuery::topic(
649            "test",
650            "worker",
651            "kv-events",
652        )));
653        assert_eq!(component.len(), 1);
654
655        let endpoint_a_instances = metadata.filter(&DiscoveryQuery::EventChannels(
656            EventChannelQuery::endpoint_topic(endpoint_a, "kv-events"),
657        ));
658        assert_eq!(endpoint_a_instances.len(), 1);
659        assert_eq!(endpoint_a_instances[0].instance_id(), 2);
660
661        let endpoint_b_instances = metadata.filter(&DiscoveryQuery::EventChannels(
662            EventChannelQuery::endpoint_topic(endpoint_b, "kv-events"),
663        ));
664        assert_eq!(endpoint_b_instances.len(), 1);
665        assert_eq!(endpoint_b_instances[0].instance_id(), 3);
666
667        metadata
668            .unregister_event_channel(&endpoint_a_instances[0])
669            .unwrap();
670        assert_eq!(metadata.get_all_event_channels().len(), 2);
671        assert!(
672            metadata
673                .filter(&DiscoveryQuery::EventChannels(
674                    EventChannelQuery::component("other", "worker")
675                ))
676                .is_empty()
677        );
678    }
679
680    #[tokio::test]
681    async fn test_mixed_instances() {
682        use crate::discovery::{EventScope, EventTransport};
683
684        let mut metadata = DiscoveryMetadata::new();
685
686        // Register one of each type
687        let endpoint = DiscoveryInstance::Endpoint(Instance {
688            namespace: "test".to_string(),
689            component: "comp1".to_string(),
690            endpoint: "ep1".to_string(),
691            instance_id: 1,
692            transport: TransportType::Nats("nats://localhost:4222".to_string()),
693            device_type: None,
694            request_plane_codec: None,
695        });
696        metadata.register_endpoint(endpoint).unwrap();
697
698        let model = DiscoveryInstance::Model {
699            namespace: "test".to_string(),
700            component: "comp1".to_string(),
701            endpoint: "ep1".to_string(),
702            instance_id: 2,
703            card_json: serde_json::json!({"model": "test"}),
704            model_suffix: None,
705        };
706        metadata.register_model_card(model).unwrap();
707
708        let event_channel = DiscoveryInstance::EventChannel {
709            scope: EventScope::Component {
710                namespace: "test".to_string(),
711                component: "comp1".to_string(),
712            },
713            topic: "test-topic".to_string(),
714            instance_id: 3,
715            transport: EventTransport::zmq("tcp://localhost:5000"),
716        };
717        metadata.register_event_channel(event_channel).unwrap();
718
719        // Verify get_all returns all three
720        assert_eq!(metadata.get_all().len(), 3);
721        assert_eq!(metadata.get_all_endpoints().len(), 1);
722        assert_eq!(metadata.get_all_model_cards().len(), 1);
723        assert_eq!(metadata.get_all_event_channels().len(), 1);
724    }
725}
dynamo_runtime/discovery/metadata.rs

dynamo_runtime/discovery/
metadata.rs
