dynamo_runtime/discovery/
mod.rs

1// SPDX-FileCopyrightText: Copyright (c) 2024-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
2// SPDX-License-Identifier: Apache-2.0
3
4use anyhow::{Context, Result};
5use async_trait::async_trait;
6use futures::Stream;
7use serde::{Deserialize, Serialize};
8use std::pin::Pin;
9use tokio_util::sync::CancellationToken;
10
11mod metadata;
12pub use metadata::{DiscoveryMetadata, MetadataSnapshot};
13
14mod mock;
15pub use mock::{MockDiscovery, SharedMockRegistry};
16mod kv_store;
17pub use kv_store::KVStoreDiscovery;
18
19mod kube;
20pub use kube::{KubeDiscoveryClient, hash_pod_name};
21
22pub mod utils;
23use crate::component::{DeviceType, TransportType};
24pub use utils::watch_and_extract_field;
25
26/// Transport kind for event plane - used for configuration and env var selection.
27///
28/// This enum represents the *type* of transport without connection details.
29/// Use `EventTransport` when you need the full transport configuration.
30#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize, Default)]
31#[serde(rename_all = "snake_case")]
32pub enum EventTransportKind {
33    /// NATS Core pub/sub
34    Nats,
35    /// ZMQ pub/sub
36    #[default]
37    Zmq,
38}
39
40impl EventTransportKind {
41    /// Parse from environment variable `DYN_EVENT_PLANE`.
42    ///
43    /// Returns `Zmq` if the variable is not set or is empty: ZMQ is the default
44    /// event plane for all backends. NATS remains available as an explicit opt-in
45    /// (`DYN_EVENT_PLANE=nats`). When you have access to a runtime, prefer
46    /// [`DistributedRuntime::default_event_transport_kind`], which resolves the same
47    /// default through the configured discovery backend.
48    ///
49    /// Returns an error for unrecognised values.
50    pub fn from_env() -> Result<Self> {
51        match std::env::var(crate::config::environment_names::event_plane::DYN_EVENT_PLANE)
52            .as_deref()
53        {
54            Ok("nats") => Ok(Self::Nats),
55            Ok("zmq") | Ok("") | Err(_) => Ok(Self::Zmq),
56            Ok(other) => anyhow::bail!(
57                "Invalid DYN_EVENT_PLANE value '{}'. Valid values: 'nats', 'zmq'",
58                other
59            ),
60        }
61    }
62
63    /// Logs a warning if an invalid value is encountered.
64    pub fn from_env_or_default() -> Self {
65        Self::from_env().unwrap_or_else(|e| {
66            tracing::warn!("{e}, defaulting to ZMQ");
67            Self::Zmq
68        })
69    }
70
71    /// Get the default codec for this transport kind.
72    /// NATS defaults to JSON, ZMQ defaults to MsgPack.
73    pub fn default_codec(&self) -> EventCodecKind {
74        match self {
75            Self::Nats => EventCodecKind::Json,
76            Self::Zmq => EventCodecKind::Msgpack,
77        }
78    }
79}
80
81/// Codec kind for event plane serialization.
82///
83/// This enum represents the serialization format for event envelopes and payloads.
84#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)]
85#[serde(rename_all = "snake_case")]
86pub enum EventCodecKind {
87    /// JSON codec - human-readable, good for debugging
88    Json,
89    /// MessagePack codec - compact binary format
90    Msgpack,
91}
92
93impl EventCodecKind {
94    /// Parse from environment variable `DYN_EVENT_PLANE_CODEC`.
95    /// Returns None if not set, allowing transport to select default.
96    /// Returns error for invalid values.
97    pub fn from_env() -> Result<Option<Self>> {
98        match std::env::var(crate::config::environment_names::event_plane::DYN_EVENT_PLANE_CODEC)
99            .as_deref()
100        {
101            Err(_) => Ok(None), // Not set
102            Ok("") => Ok(None), // Empty
103            Ok("json") => Ok(Some(Self::Json)),
104            Ok("msgpack") => Ok(Some(Self::Msgpack)),
105            Ok(other) => anyhow::bail!(
106                "Invalid DYN_EVENT_PLANE_CODEC value '{}'. Valid values: 'json', 'msgpack'",
107                other
108            ),
109        }
110    }
111
112    /// Parse from environment variable with transport-specific default.
113    /// Logs a warning if an invalid value is encountered.
114    pub fn from_env_or_transport_default(transport: EventTransportKind) -> Self {
115        Self::from_env()
116            .unwrap_or_else(|e| {
117                tracing::warn!(
118                    "{}, defaulting to {:?} for {:?}",
119                    e,
120                    transport.default_codec(),
121                    transport
122                );
123                None
124            })
125            .unwrap_or_else(|| transport.default_codec())
126    }
127}
128
129/// Transport configuration for event plane channels.
130///
131/// This enum carries both the transport kind and its connection configuration.
132/// Kept separate from `TransportType` (request plane) to distinguish event semantics.
133#[derive(Debug, Clone, PartialEq, Eq, Hash, Serialize, Deserialize)]
134#[serde(tag = "kind", content = "config")]
135pub enum EventTransport {
136    /// NATS Core pub/sub - subject prefix for the channel
137    Nats {
138        /// Subject prefix (e.g., "namespace.dynamo.component.backend")
139        subject_prefix: String,
140    },
141    /// ZMQ pub/sub - endpoint address (direct mode)
142    Zmq {
143        /// ZMQ endpoint (e.g., "tcp://host:port")
144        endpoint: String,
145    },
146    /// ZMQ broker endpoints (broker mode) - for discovery of brokers
147    ZmqBroker {
148        /// XSUB endpoints (publishers connect here)
149        xsub_endpoints: Vec<String>,
150        /// XPUB endpoints (subscribers connect here)
151        xpub_endpoints: Vec<String>,
152    },
153}
154
155impl EventTransport {
156    /// Get the transport kind
157    pub fn kind(&self) -> EventTransportKind {
158        match self {
159            Self::Nats { .. } => EventTransportKind::Nats,
160            Self::Zmq { .. } | Self::ZmqBroker { .. } => EventTransportKind::Zmq,
161        }
162    }
163
164    /// Create a NATS transport with the given subject prefix
165    pub fn nats(subject_prefix: impl Into<String>) -> Self {
166        Self::Nats {
167            subject_prefix: subject_prefix.into(),
168        }
169    }
170
171    /// Create a ZMQ transport with the given endpoint
172    pub fn zmq(endpoint: impl Into<String>) -> Self {
173        Self::Zmq {
174            endpoint: endpoint.into(),
175        }
176    }
177
178    /// Get the subject prefix (NATS) or endpoint (ZMQ)
179    /// For ZmqBroker, returns the first XSUB endpoint
180    pub fn address(&self) -> &str {
181        match self {
182            Self::Nats { subject_prefix } => subject_prefix,
183            Self::Zmq { endpoint } => endpoint,
184            Self::ZmqBroker { xsub_endpoints, .. } => {
185                xsub_endpoints.first().map(|s| s.as_str()).unwrap_or("")
186            }
187        }
188    }
189}
190
191/// Query key for prefix-based discovery queries
192/// Supports hierarchical queries from all endpoints down to specific endpoints
193#[derive(Debug, Clone, PartialEq, Eq, Hash)]
194pub enum DiscoveryQuery {
195    /// Query all endpoints in the system
196    AllEndpoints,
197    /// Query all endpoints in a specific namespace
198    NamespacedEndpoints {
199        namespace: String,
200    },
201    /// Query all endpoints in a namespace/component
202    ComponentEndpoints {
203        namespace: String,
204        component: String,
205    },
206    /// Query a specific endpoint
207    Endpoint {
208        namespace: String,
209        component: String,
210        endpoint: String,
211    },
212    AllModels,
213    NamespacedModels {
214        namespace: String,
215    },
216    ComponentModels {
217        namespace: String,
218        component: String,
219    },
220    EndpointModels {
221        namespace: String,
222        component: String,
223        endpoint: String,
224    },
225    /// Unified event channel query with optional scope filters
226    EventChannels(EventChannelQuery),
227}
228
229/// Unified query for event channels with optional scope filters
230#[derive(Debug, Clone, PartialEq, Eq, Hash)]
231pub struct EventChannelQuery {
232    /// Optional namespace filter
233    pub namespace: Option<String>,
234    /// Optional component filter (requires namespace to be meaningful)
235    pub component: Option<String>,
236    /// Optional topic filter (requires namespace and component to be meaningful)
237    pub topic: Option<String>,
238}
239
240impl EventChannelQuery {
241    /// Query all event channels (no filters)
242    pub fn all() -> Self {
243        Self {
244            namespace: None,
245            component: None,
246            topic: None,
247        }
248    }
249
250    /// Query event channels in a specific namespace
251    pub fn namespace(namespace: impl Into<String>) -> Self {
252        Self {
253            namespace: Some(namespace.into()),
254            component: None,
255            topic: None,
256        }
257    }
258
259    /// Query event channels for a specific component
260    pub fn component(namespace: impl Into<String>, component: impl Into<String>) -> Self {
261        Self {
262            namespace: Some(namespace.into()),
263            component: Some(component.into()),
264            topic: None,
265        }
266    }
267
268    /// Query event channels for a specific topic
269    pub fn topic(
270        namespace: impl Into<String>,
271        component: impl Into<String>,
272        topic: impl Into<String>,
273    ) -> Self {
274        Self {
275            namespace: Some(namespace.into()),
276            component: Some(component.into()),
277            topic: Some(topic.into()),
278        }
279    }
280
281    /// Get the scope level (0=all, 1=namespace, 2=component, 3=topic)
282    pub fn scope_level(&self) -> u8 {
283        if self.topic.is_some() {
284            3
285        } else if self.component.is_some() {
286            2
287        } else if self.namespace.is_some() {
288            1
289        } else {
290            0
291        }
292    }
293}
294
295/// Specification for registering objects in the discovery plane
296/// Represents the input to the register() operation
297#[derive(Debug, Clone, PartialEq, Eq)]
298pub enum DiscoverySpec {
299    /// Endpoint specification for registration
300    Endpoint {
301        namespace: String,
302        component: String,
303        endpoint: String,
304        /// Transport type and routing information
305        transport: TransportType,
306        /// Optional execution device for this endpoint instance.
307        /// Used by hetero routing to distinguish CPU and CUDA workers.
308        device_type: Option<DeviceType>,
309    },
310    Model {
311        namespace: String,
312        component: String,
313        endpoint: String,
314        /// ModelDeploymentCard serialized as JSON
315        /// This allows lib/runtime to remain independent of lib/llm types
316        /// DiscoverySpec.from_model() and DiscoveryInstance.deserialize_model() are ergonomic helpers to create and deserialize the model card.
317        card_json: serde_json::Value,
318        /// Optional suffix appended after instance_id in the key path (e.g., for LoRA adapters)
319        /// Key format: {namespace}/{component}/{endpoint}/{instance_id}[/{model_suffix}]
320        model_suffix: Option<String>,
321    },
322    /// Event plane channel specification
323    /// Used for registering event publishers/subscribers for discovery
324    EventChannel {
325        namespace: String,
326        component: String,
327        /// Topic name for this channel (e.g., "kv-events", "kv-metrics")
328        topic: String,
329        /// Unique identity of this publisher incarnation.
330        ///
331        /// A process can host multiple publishers for the same topic, so event
332        /// channels cannot use the process-level discovery instance ID.
333        publisher_id: u64,
334        /// Event transport type (NATS subject prefix or ZMQ endpoint)
335        transport: EventTransport,
336    },
337}
338
339impl DiscoverySpec {
340    /// Creates a Model discovery spec from a serializable type
341    /// The card will be serialized to JSON to avoid cross-crate dependencies
342    pub fn from_model<T>(
343        namespace: String,
344        component: String,
345        endpoint: String,
346        card: &T,
347    ) -> Result<Self>
348    where
349        T: Serialize,
350    {
351        Self::from_model_with_suffix(namespace, component, endpoint, card, None)
352    }
353
354    /// Creates a Model discovery spec with an optional suffix (e.g., for LoRA adapters)
355    /// The suffix is appended after the instance_id in the key path
356    pub fn from_model_with_suffix<T>(
357        namespace: String,
358        component: String,
359        endpoint: String,
360        card: &T,
361        model_suffix: Option<String>,
362    ) -> Result<Self>
363    where
364        T: Serialize,
365    {
366        let card_json = serde_json::to_value(card)?;
367        Ok(Self::Model {
368            namespace,
369            component,
370            endpoint,
371            card_json,
372            model_suffix,
373        })
374    }
375
376    /// Converts this registration spec into a discovery instance.
377    ///
378    /// Endpoint and model specs use `default_instance_id`, normally the
379    /// discovery client's process-level ID. Event channel specs already carry
380    /// a publisher-level ID, so they use that instead.
381    pub fn into_instance(self, default_instance_id: u64) -> DiscoveryInstance {
382        match self {
383            Self::Endpoint {
384                namespace,
385                component,
386                endpoint,
387                transport,
388                device_type,
389            } => DiscoveryInstance::Endpoint(crate::component::Instance {
390                namespace,
391                component,
392                endpoint,
393                instance_id: default_instance_id,
394                transport,
395                device_type,
396            }),
397            Self::Model {
398                namespace,
399                component,
400                endpoint,
401                card_json,
402                model_suffix,
403            } => DiscoveryInstance::Model {
404                namespace,
405                component,
406                endpoint,
407                instance_id: default_instance_id,
408                card_json,
409                model_suffix,
410            },
411            Self::EventChannel {
412                namespace,
413                component,
414                topic,
415                publisher_id,
416                transport,
417            } => DiscoveryInstance::EventChannel {
418                namespace,
419                component,
420                topic,
421                instance_id: publisher_id,
422                transport,
423            },
424        }
425    }
426
427    /// Compatibility alias for [`DiscoverySpec::into_instance`].
428    pub fn with_instance_id(self, default_instance_id: u64) -> DiscoveryInstance {
429        self.into_instance(default_instance_id)
430    }
431}
432
433/// Registered instances in the discovery plane
434/// Represents objects that have been successfully registered with an instance ID
435#[derive(Debug, Clone, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
436#[serde(tag = "type")]
437pub enum DiscoveryInstance {
438    /// Registered endpoint instance - wraps the component::Instance directly
439    Endpoint(crate::component::Instance),
440    Model {
441        namespace: String,
442        component: String,
443        endpoint: String,
444        instance_id: u64,
445        /// ModelDeploymentCard serialized as JSON
446        /// This allows lib/runtime to remain independent of lib/llm types
447        card_json: serde_json::Value,
448        /// Optional suffix appended after instance_id in the key path (e.g., for LoRA adapters)
449        #[serde(default, skip_serializing_if = "Option::is_none")]
450        model_suffix: Option<String>,
451    },
452    /// Registered event channel instance for event plane pub/sub
453    EventChannel {
454        namespace: String,
455        component: String,
456        /// Topic name for this channel (e.g., "kv-events", "kv-metrics")
457        topic: String,
458        instance_id: u64,
459        /// Event transport type (NATS subject prefix or ZMQ endpoint)
460        transport: EventTransport,
461    },
462}
463
464impl DiscoveryInstance {
465    /// Returns the instance ID for this discovery instance
466    pub fn instance_id(&self) -> u64 {
467        match self {
468            Self::Endpoint(inst) => inst.instance_id,
469            Self::Model { instance_id, .. } => *instance_id,
470            Self::EventChannel { instance_id, .. } => *instance_id,
471        }
472    }
473
474    /// Deserializes the model JSON into the specified type T
475    /// Returns an error if this is not a Model instance or if deserialization fails
476    pub fn deserialize_model<T>(&self) -> Result<T>
477    where
478        T: for<'de> Deserialize<'de>,
479    {
480        match self {
481            Self::Model { card_json, .. } => Ok(serde_json::from_value(card_json.clone())?),
482            Self::Endpoint(_) => {
483                anyhow::bail!("Cannot deserialize model from Endpoint instance")
484            }
485            Self::EventChannel { .. } => {
486                anyhow::bail!("Cannot deserialize model from EventChannel instance")
487            }
488        }
489    }
490
491    /// Extracts the unique identifier for this discovery instance
492    /// Used for tracking, diffing, and removal events
493    pub fn id(&self) -> DiscoveryInstanceId {
494        match self {
495            Self::Endpoint(inst) => DiscoveryInstanceId::Endpoint(EndpointInstanceId {
496                namespace: inst.namespace.clone(),
497                component: inst.component.clone(),
498                endpoint: inst.endpoint.clone(),
499                instance_id: inst.instance_id,
500            }),
501            Self::Model {
502                namespace,
503                component,
504                endpoint,
505                instance_id,
506                model_suffix,
507                ..
508            } => DiscoveryInstanceId::Model(ModelCardInstanceId {
509                namespace: namespace.clone(),
510                component: component.clone(),
511                endpoint: endpoint.clone(),
512                instance_id: *instance_id,
513                model_suffix: model_suffix.clone(),
514            }),
515            Self::EventChannel {
516                namespace,
517                component,
518                topic,
519                instance_id,
520                ..
521            } => DiscoveryInstanceId::EventChannel(EventChannelInstanceId {
522                namespace: namespace.clone(),
523                component: component.clone(),
524                topic: topic.clone(),
525                instance_id: *instance_id,
526            }),
527        }
528    }
529}
530
531/// Unique identifier for an endpoint instance
532#[derive(Debug, Clone, PartialEq, Eq, Hash, Serialize, Deserialize)]
533pub struct EndpointInstanceId {
534    pub namespace: String,
535    pub component: String,
536    pub endpoint: String,
537    pub instance_id: u64,
538}
539
540impl EndpointInstanceId {
541    /// Converts to a path string: `{namespace}/{component}/{endpoint}/{instance_id:x}`
542    pub fn to_path(&self) -> String {
543        format!(
544            "{}/{}/{}/{:x}",
545            self.namespace, self.component, self.endpoint, self.instance_id
546        )
547    }
548
549    /// Parses from a path string: `{namespace}/{component}/{endpoint}/{instance_id:x}`
550    pub fn from_path(path: &str) -> Result<Self> {
551        let parts: Vec<&str> = path.split('/').collect();
552        if parts.len() != 4 {
553            anyhow::bail!(
554                "Invalid EndpointInstanceId path: expected 4 parts, got {}",
555                parts.len()
556            );
557        }
558        Ok(Self {
559            namespace: parts[0].to_string(),
560            component: parts[1].to_string(),
561            endpoint: parts[2].to_string(),
562            instance_id: u64::from_str_radix(parts[3], 16)
563                .map_err(|e| anyhow::anyhow!("Invalid instance_id hex: {}", e))?,
564        })
565    }
566}
567
568/// Unique identifier for a model card instance
569/// The combination of (namespace, component, endpoint, instance_id, model_suffix) uniquely identifies a model card
570#[derive(Debug, Clone, PartialEq, Eq, Hash, Serialize, Deserialize)]
571pub struct ModelCardInstanceId {
572    pub namespace: String,
573    pub component: String,
574    pub endpoint: String,
575    pub instance_id: u64,
576    /// None for base models, Some(slug) for LoRA adapters
577    pub model_suffix: Option<String>,
578}
579
580/// Unique identifier for an event channel instance
581#[derive(Debug, Clone, PartialEq, Eq, Hash, Serialize, Deserialize)]
582pub struct EventChannelInstanceId {
583    pub namespace: String,
584    pub component: String,
585    /// Topic name for this channel (e.g., "kv-events", "kv-metrics")
586    pub topic: String,
587    pub instance_id: u64,
588}
589
590impl EventChannelInstanceId {
591    /// Converts to a path string: `{namespace}/{component}/{topic}/{instance_id:x}`
592    pub fn to_path(&self) -> String {
593        format!(
594            "{}/{}/{}/{:x}",
595            self.namespace, self.component, self.topic, self.instance_id
596        )
597    }
598
599    /// Parses from a path string: `{namespace}/{component}/{topic}/{instance_id:x}`
600    pub fn from_path(path: &str) -> Result<Self> {
601        let parts: Vec<&str> = path.split('/').collect();
602        if parts.len() != 4 {
603            anyhow::bail!(
604                "Invalid EventChannelInstanceId path: expected 4 parts, got {}",
605                parts.len()
606            );
607        }
608        Ok(Self {
609            namespace: parts[0].to_string(),
610            component: parts[1].to_string(),
611            topic: parts[2].to_string(),
612            instance_id: u64::from_str_radix(parts[3], 16)
613                .map_err(|e| anyhow::anyhow!("Invalid instance_id hex: {}", e))?,
614        })
615    }
616}
617
618impl ModelCardInstanceId {
619    /// Converts to a path string: `{namespace}/{component}/{endpoint}/{instance_id:x}[/{model_suffix}]`
620    pub fn to_path(&self) -> String {
621        match &self.model_suffix {
622            Some(suffix) => format!(
623                "{}/{}/{}/{:x}/{}",
624                self.namespace, self.component, self.endpoint, self.instance_id, suffix
625            ),
626            None => format!(
627                "{}/{}/{}/{:x}",
628                self.namespace, self.component, self.endpoint, self.instance_id
629            ),
630        }
631    }
632
633    /// Parses from a path string: `{namespace}/{component}/{endpoint}/{instance_id:x}[/{model_suffix}]`
634    pub fn from_path(path: &str) -> Result<Self> {
635        let parts: Vec<&str> = path.split('/').collect();
636        if parts.len() < 4 || parts.len() > 5 {
637            anyhow::bail!(
638                "Invalid ModelCardInstanceId path: expected 4 or 5 parts, got {}",
639                parts.len()
640            );
641        }
642        Ok(Self {
643            namespace: parts[0].to_string(),
644            component: parts[1].to_string(),
645            endpoint: parts[2].to_string(),
646            instance_id: u64::from_str_radix(parts[3], 16)
647                .map_err(|e| anyhow::anyhow!("Invalid instance_id hex: {}", e))?,
648            model_suffix: parts.get(4).map(|s| s.to_string()),
649        })
650    }
651}
652
653/// Union of instance identifiers for different discovery object types
654#[derive(Debug, Clone, PartialEq, Eq, Hash, Serialize, Deserialize)]
655pub enum DiscoveryInstanceId {
656    Endpoint(EndpointInstanceId),
657    Model(ModelCardInstanceId),
658    EventChannel(EventChannelInstanceId),
659}
660
661impl DiscoveryInstanceId {
662    /// Returns the raw instance_id regardless of variant type
663    pub fn instance_id(&self) -> u64 {
664        match self {
665            Self::Endpoint(eid) => eid.instance_id,
666            Self::Model(mid) => mid.instance_id,
667            Self::EventChannel(ecid) => ecid.instance_id,
668        }
669    }
670
671    /// Extracts the EndpointInstanceId, returning an error if this is a Model or EventChannel variant
672    pub fn extract_endpoint_id(&self) -> Result<&EndpointInstanceId> {
673        match self {
674            Self::Endpoint(eid) => Ok(eid),
675            Self::Model(_) => anyhow::bail!("Expected Endpoint variant, got Model"),
676            Self::EventChannel(_) => anyhow::bail!("Expected Endpoint variant, got EventChannel"),
677        }
678    }
679
680    /// Extracts the ModelCardInstanceId, returning an error if this is an Endpoint or EventChannel variant
681    pub fn extract_model_id(&self) -> Result<&ModelCardInstanceId> {
682        match self {
683            Self::Model(mid) => Ok(mid),
684            Self::Endpoint(_) => anyhow::bail!("Expected Model variant, got Endpoint"),
685            Self::EventChannel(_) => anyhow::bail!("Expected Model variant, got EventChannel"),
686        }
687    }
688
689    /// Extracts the EventChannelInstanceId, returning an error if this is an Endpoint or Model variant
690    pub fn extract_event_channel_id(&self) -> Result<&EventChannelInstanceId> {
691        match self {
692            Self::EventChannel(ecid) => Ok(ecid),
693            Self::Endpoint(_) => anyhow::bail!("Expected EventChannel variant, got Endpoint"),
694            Self::Model(_) => anyhow::bail!("Expected EventChannel variant, got Model"),
695        }
696    }
697}
698
699/// Events emitted by the discovery watch stream
700#[derive(Debug, Clone, PartialEq, Eq)]
701pub enum DiscoveryEvent {
702    /// A new instance was added
703    Added(DiscoveryInstance),
704    /// An instance was removed (identified by its unique ID)
705    Removed(DiscoveryInstanceId),
706}
707
708/// Stream type for discovery events
709pub type DiscoveryStream = Pin<Box<dyn Stream<Item = Result<DiscoveryEvent>> + Send>>;
710
711#[derive(Clone, Debug, PartialEq, Eq)]
712struct ModelRegistrationIdentity {
713    display_name: String,
714    source_path: Option<String>,
715    is_lora: bool,
716}
717
718impl ModelRegistrationIdentity {
719    fn base_identity(&self) -> &str {
720        self.source_path.as_deref().unwrap_or(&self.display_name)
721    }
722
723    fn is_compatible_with(&self, other: &Self) -> bool {
724        if self.is_lora || other.is_lora {
725            self.base_identity() == other.base_identity()
726        } else {
727            self.display_name == other.display_name
728        }
729    }
730}
731
732fn extract_model_registration_identity(
733    card_json: &serde_json::Value,
734    model_suffix: Option<&str>,
735) -> Result<ModelRegistrationIdentity> {
736    let display_name = card_json
737        .get("display_name")
738        .and_then(serde_json::Value::as_str)
739        .map(str::to_owned)
740        .ok_or_else(|| {
741            anyhow::anyhow!("failed to deserialize model display_name from card_json")
742        })?;
743    let source_path = card_json
744        .get("source_path")
745        .and_then(serde_json::Value::as_str)
746        .map(str::to_owned);
747    let is_lora =
748        model_suffix.is_some() || card_json.get("lora").is_some_and(|value| !value.is_null());
749
750    Ok(ModelRegistrationIdentity {
751        display_name,
752        source_path,
753        is_lora,
754    })
755}
756
757fn find_conflicting_model_name(
758    instances: &[DiscoveryInstance],
759    requested_identity: &ModelRegistrationIdentity,
760) -> Result<Option<String>> {
761    for instance in instances {
762        if let DiscoveryInstance::Model {
763            card_json,
764            model_suffix,
765            ..
766        } = instance
767        {
768            let existing_identity =
769                extract_model_registration_identity(card_json, model_suffix.as_deref())?;
770            if !requested_identity.is_compatible_with(&existing_identity) {
771                return Ok(Some(existing_identity.display_name));
772            }
773        }
774    }
775
776    Ok(None)
777}
778
779/// Discovery trait for service discovery across different backends
780#[async_trait]
781pub trait Discovery: Send + Sync {
782    /// Returns a unique identifier for this worker (e.g lease id if using etcd or generated id for memory store)
783    /// Endpoint and model objects created by this worker use this ID. Event
784    /// channels use a publisher-level ID because a worker can own more than one
785    /// publisher for the same topic.
786    fn instance_id(&self) -> u64;
787
788    /// Registers an object in the discovery plane with the instance id
789    async fn register(&self, spec: DiscoverySpec) -> Result<DiscoveryInstance> {
790        let (namespace, component, endpoint, requested_identity) = match &spec {
791            DiscoverySpec::Model {
792                namespace,
793                component,
794                endpoint,
795                card_json,
796                model_suffix,
797                ..
798            } => (
799                namespace.clone(),
800                component.clone(),
801                endpoint.clone(),
802                extract_model_registration_identity(card_json, model_suffix.as_deref())?,
803            ),
804            _ => return self.register_internal(spec).await,
805        };
806
807        let query = DiscoveryQuery::EndpointModels {
808            namespace: namespace.clone(),
809            component: component.clone(),
810            endpoint: endpoint.clone(),
811        };
812
813        if let Some(conflicting_name) =
814            find_conflicting_model_name(&self.list(query.clone()).await?, &requested_identity)?
815        {
816            let requested_name = &requested_identity.display_name;
817            anyhow::bail!(
818                "Cannot register model '{requested_name}' on endpoint '{namespace}/{component}/{endpoint}': a different model '{conflicting_name}' is already registered there"
819            );
820        }
821
822        let instance = self.register_internal(spec).await?;
823
824        if let Some(conflicting_name) =
825            find_conflicting_model_name(&self.list(query).await?, &requested_identity)?
826        {
827            let requested_name = &requested_identity.display_name;
828            if let Err(unregister_err) = self.unregister(instance.clone()).await {
829                return Err(anyhow::anyhow!(
830                    "Cannot register model '{requested_name}' on endpoint '{namespace}/{component}/{endpoint}': a different model '{conflicting_name}' is already registered there"
831                ))
832                .context(format!(
833                    "failed to roll back conflicting model registration for instance {instance_id}: {unregister_err}",
834                    instance_id = instance.instance_id()
835                ));
836            }
837
838            anyhow::bail!(
839                "Cannot register model '{requested_name}' on endpoint '{namespace}/{component}/{endpoint}': a different model '{conflicting_name}' is already registered there"
840            );
841        }
842
843        Ok(instance)
844    }
845
846    /// Backend-specific raw registration implementation.
847    async fn register_internal(&self, spec: DiscoverySpec) -> Result<DiscoveryInstance>;
848
849    /// Unregisters an instance from the discovery plane
850    async fn unregister(&self, instance: DiscoveryInstance) -> Result<()>;
851
852    /// Returns a list of currently registered instances for the given discovery query
853    /// This is a one-time snapshot without watching for changes
854    async fn list(&self, query: DiscoveryQuery) -> Result<Vec<DiscoveryInstance>>;
855
856    /// Returns a stream of discovery events (Added/Removed) for the given discovery query
857    /// The optional cancellation token can be used to stop the watch stream
858    async fn list_and_watch(
859        &self,
860        query: DiscoveryQuery,
861        cancel_token: Option<CancellationToken>,
862    ) -> Result<DiscoveryStream>;
863
864    /// Clean up resources held by this discovery backend.
865    /// For KV store backends, this deletes owned registrations immediately rather than
866    /// waiting for TTL expiry. Default is a no-op for backends that don't need cleanup.
867    fn shutdown(&self) {}
868}
dynamo_runtime/discovery/mod.rs

dynamo_runtime/discovery/
mod.rs
