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};
8
9use crate::protocols::EndpointId;
10use std::pin::Pin;
11use tokio_util::sync::CancellationToken;
12
13mod metadata;
14pub use metadata::{DiscoveryMetadata, MetadataSnapshot};
15
16mod registration;
17pub use registration::EndpointRegistrationLease;
18pub(crate) use registration::EndpointRegistrationManager;
19
20mod mock;
21pub use mock::{MockDiscovery, SharedMockRegistry};
22mod kv_store;
23pub use kv_store::KVStoreDiscovery;
24
25mod kube;
26pub use kube::{KubeDiscoveryClient, hash_pod_name};
27
28pub mod utils;
29use crate::{
30    component::{DeviceType, TransportType},
31    pipeline::network::RequestPlanePayloadCodec,
32};
33pub use utils::watch_and_extract_field;
34
35/// Largest publisher ID exactly representable by float64-backed JSON metadata.
36pub(crate) const MAX_JSON_SAFE_PUBLISHER_ID: u64 = (1 << 53) - 1;
37
38/// Transport kind for event plane - used for configuration and env var selection.
39///
40/// This enum represents the *type* of transport without connection details.
41/// Use `EventTransport` when you need the full transport configuration.
42#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize, Default)]
43#[serde(rename_all = "snake_case")]
44pub enum EventTransportKind {
45    /// NATS Core pub/sub
46    Nats,
47    /// ZMQ pub/sub
48    #[default]
49    Zmq,
50}
51
52impl EventTransportKind {
53    /// Parse from environment variable `DYN_EVENT_PLANE`.
54    ///
55    /// Returns `Zmq` if the variable is not set or is empty: ZMQ is the default
56    /// event plane for all backends. NATS remains available as an explicit opt-in
57    /// (`DYN_EVENT_PLANE=nats`). When you have access to a runtime, prefer
58    /// [`DistributedRuntime::default_event_transport_kind`], which resolves the same
59    /// default through the configured discovery backend.
60    ///
61    /// Returns an error for unrecognised values.
62    pub fn from_env() -> Result<Self> {
63        match std::env::var(crate::config::environment_names::event_plane::DYN_EVENT_PLANE)
64            .as_deref()
65        {
66            Ok("nats") => Ok(Self::Nats),
67            Ok("zmq") | Ok("") | Err(_) => Ok(Self::Zmq),
68            Ok(other) => anyhow::bail!(
69                "Invalid DYN_EVENT_PLANE value '{}'. Valid values: 'nats', 'zmq'",
70                other
71            ),
72        }
73    }
74
75    /// Logs a warning if an invalid value is encountered.
76    pub fn from_env_or_default() -> Self {
77        Self::from_env().unwrap_or_else(|e| {
78            tracing::warn!("{e}, defaulting to ZMQ");
79            Self::Zmq
80        })
81    }
82
83    /// Get the default codec for this transport kind.
84    /// NATS defaults to JSON, ZMQ defaults to MsgPack.
85    pub fn default_codec(&self) -> EventCodecKind {
86        match self {
87            Self::Nats => EventCodecKind::Json,
88            Self::Zmq => EventCodecKind::Msgpack,
89        }
90    }
91}
92
93/// Codec kind for event plane serialization.
94///
95/// This enum represents the serialization format for event envelopes and payloads.
96#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)]
97#[serde(rename_all = "snake_case")]
98pub enum EventCodecKind {
99    /// JSON codec - human-readable, good for debugging
100    Json,
101    /// MessagePack codec - compact binary format
102    Msgpack,
103}
104
105impl EventCodecKind {
106    /// Parse from environment variable `DYN_EVENT_PLANE_CODEC`.
107    /// Returns None if not set, allowing transport to select default.
108    /// Returns error for invalid values.
109    pub fn from_env() -> Result<Option<Self>> {
110        match std::env::var(crate::config::environment_names::event_plane::DYN_EVENT_PLANE_CODEC)
111            .as_deref()
112        {
113            Err(_) => Ok(None), // Not set
114            Ok("") => Ok(None), // Empty
115            Ok("json") => Ok(Some(Self::Json)),
116            Ok("msgpack") => Ok(Some(Self::Msgpack)),
117            Ok(other) => anyhow::bail!(
118                "Invalid DYN_EVENT_PLANE_CODEC value '{}'. Valid values: 'json', 'msgpack'",
119                other
120            ),
121        }
122    }
123
124    /// Parse from environment variable with transport-specific default.
125    /// Logs a warning if an invalid value is encountered.
126    pub fn from_env_or_transport_default(transport: EventTransportKind) -> Self {
127        Self::from_env()
128            .unwrap_or_else(|e| {
129                tracing::warn!(
130                    "{}, defaulting to {:?} for {:?}",
131                    e,
132                    transport.default_codec(),
133                    transport
134                );
135                None
136            })
137            .unwrap_or_else(|| transport.default_codec())
138    }
139}
140
141/// Transport configuration for event plane channels.
142///
143/// This enum carries both the transport kind and its connection configuration.
144/// Kept separate from `TransportType` (request plane) to distinguish event semantics.
145#[derive(Debug, Clone, PartialEq, Eq, Hash, Serialize, Deserialize)]
146#[serde(tag = "kind", content = "config")]
147pub enum EventTransport {
148    /// NATS Core pub/sub - subject prefix for the channel
149    Nats {
150        /// Subject prefix (e.g., "namespace.dynamo.component.backend")
151        subject_prefix: String,
152    },
153    /// ZMQ pub/sub - endpoint address (direct mode)
154    Zmq {
155        /// ZMQ endpoint (e.g., "tcp://host:port")
156        endpoint: String,
157    },
158    /// ZMQ broker endpoints (broker mode) - for discovery of brokers
159    ZmqBroker {
160        /// XSUB endpoints (publishers connect here)
161        xsub_endpoints: Vec<String>,
162        /// XPUB endpoints (subscribers connect here)
163        xpub_endpoints: Vec<String>,
164    },
165}
166
167impl EventTransport {
168    /// Get the transport kind
169    pub fn kind(&self) -> EventTransportKind {
170        match self {
171            Self::Nats { .. } => EventTransportKind::Nats,
172            Self::Zmq { .. } | Self::ZmqBroker { .. } => EventTransportKind::Zmq,
173        }
174    }
175
176    /// Create a NATS transport with the given subject prefix
177    pub fn nats(subject_prefix: impl Into<String>) -> Self {
178        Self::Nats {
179            subject_prefix: subject_prefix.into(),
180        }
181    }
182
183    /// Create a ZMQ transport with the given endpoint
184    pub fn zmq(endpoint: impl Into<String>) -> Self {
185        Self::Zmq {
186            endpoint: endpoint.into(),
187        }
188    }
189
190    /// Get the subject prefix (NATS) or endpoint (ZMQ)
191    /// For ZmqBroker, returns the first XSUB endpoint
192    pub fn address(&self) -> &str {
193        match self {
194            Self::Nats { subject_prefix } => subject_prefix,
195            Self::Zmq { endpoint } => endpoint,
196            Self::ZmqBroker { xsub_endpoints, .. } => {
197                xsub_endpoints.first().map(|s| s.as_str()).unwrap_or("")
198            }
199        }
200    }
201}
202
203/// Query key for prefix-based discovery queries
204/// Supports hierarchical queries from all endpoints down to specific endpoints
205#[derive(Debug, Clone, PartialEq, Eq, Hash)]
206pub enum DiscoveryQuery {
207    /// Query all endpoints in the system
208    AllEndpoints,
209    /// Query all endpoints in a specific namespace
210    NamespacedEndpoints {
211        namespace: String,
212    },
213    /// Query all endpoints in a namespace/component
214    ComponentEndpoints {
215        namespace: String,
216        component: String,
217    },
218    /// Query a specific endpoint
219    Endpoint {
220        namespace: String,
221        component: String,
222        endpoint: String,
223    },
224    AllModels,
225    NamespacedModels {
226        namespace: String,
227    },
228    ComponentModels {
229        namespace: String,
230        component: String,
231    },
232    EndpointModels {
233        namespace: String,
234        component: String,
235        endpoint: String,
236    },
237    /// Unified event channel query with optional scope filters
238    EventChannels(EventChannelQuery),
239    /// Semantic event source query with optional scope filters.
240    EventSources(EventSourceQuery),
241}
242
243/// Scope of an event channel.
244///
245/// Event scopes are exact ownership domains. A namespace-scoped query does not
246/// match component- or endpoint-scoped publishers in that namespace.
247#[derive(Debug, Clone, PartialEq, Eq, Hash, Serialize, Deserialize)]
248#[serde(tag = "kind", rename_all = "snake_case")]
249pub enum EventScope {
250    Namespace {
251        name: String,
252    },
253    Component {
254        namespace: String,
255        component: String,
256    },
257    Endpoint {
258        endpoint: EndpointId,
259    },
260}
261
262impl EventScope {
263    pub fn namespace(&self) -> &str {
264        match self {
265            Self::Namespace { name } => name,
266            Self::Component { namespace, .. } => namespace,
267            Self::Endpoint { endpoint } => &endpoint.namespace,
268        }
269    }
270
271    pub fn component(&self) -> Option<&str> {
272        match self {
273            Self::Namespace { .. } => None,
274            Self::Component { component, .. } => Some(component),
275            Self::Endpoint { endpoint } => Some(&endpoint.component),
276        }
277    }
278
279    pub fn endpoint(&self) -> Option<&EndpointId> {
280        match self {
281            Self::Endpoint { endpoint } => Some(endpoint),
282            Self::Namespace { .. } | Self::Component { .. } => None,
283        }
284    }
285
286    /// Canonical NATS/ZMQ-broker subject prefix for this scope.
287    pub fn subject_prefix(&self) -> String {
288        match self {
289            Self::Namespace { name } => {
290                format!("namespace.{}", encode_event_segment(name))
291            }
292            Self::Component {
293                namespace,
294                component,
295            } => format!(
296                "namespace.{}.component.{}",
297                encode_event_segment(namespace),
298                encode_event_segment(component)
299            ),
300            Self::Endpoint { endpoint } => format!(
301                "namespace.{}.component.{}.endpoint.{}",
302                encode_event_segment(&endpoint.namespace),
303                encode_event_segment(&endpoint.component),
304                encode_event_segment(&endpoint.name)
305            ),
306        }
307    }
308
309    /// Canonical subject/routing key for a topic in this exact scope.
310    pub fn subject(&self, topic: &str) -> String {
311        format!("{}.{}", self.subject_prefix(), encode_event_segment(topic))
312    }
313
314    pub(crate) fn path_prefix(&self) -> String {
315        match self {
316            Self::Namespace { name } => {
317                format!("namespace/{}", encode_event_segment(name))
318            }
319            Self::Component {
320                namespace,
321                component,
322            } => format!(
323                "namespace/{}/component/{}",
324                encode_event_segment(namespace),
325                encode_event_segment(component)
326            ),
327            Self::Endpoint { endpoint } => format!(
328                "namespace/{}/component/{}/endpoint/{}",
329                encode_event_segment(&endpoint.namespace),
330                encode_event_segment(&endpoint.component),
331                encode_event_segment(&endpoint.name)
332            ),
333        }
334    }
335}
336
337/// Percent-encode a subject/path segment while keeping common identifier
338/// characters readable. The encoding is reversible and prevents NATS wildcard
339/// or discovery path delimiters from changing the channel identity.
340pub(crate) fn encode_event_segment(value: &str) -> String {
341    let mut encoded = String::with_capacity(value.len());
342    for byte in value.bytes() {
343        if byte.is_ascii_alphanumeric() || matches!(byte, b'-' | b'_') {
344            encoded.push(char::from(byte));
345        } else {
346            use std::fmt::Write as _;
347            write!(encoded, "%{byte:02X}").expect("writing to String cannot fail");
348        }
349    }
350    encoded
351}
352
353fn decode_event_segment(value: &str) -> Result<String> {
354    let bytes = value.as_bytes();
355    let mut decoded = Vec::with_capacity(bytes.len());
356    let mut index = 0;
357    while index < bytes.len() {
358        if bytes[index] != b'%' {
359            decoded.push(bytes[index]);
360            index += 1;
361            continue;
362        }
363        if index + 2 >= bytes.len() {
364            anyhow::bail!("invalid percent-encoded event segment: {value}");
365        }
366        let hex = std::str::from_utf8(&bytes[index + 1..index + 3])?;
367        decoded
368            .push(u8::from_str_radix(hex, 16).map_err(|error| {
369                anyhow::anyhow!("invalid event segment escape %{hex}: {error}")
370            })?);
371        index += 3;
372    }
373    String::from_utf8(decoded)
374        .map_err(|error| anyhow::anyhow!("event segment is not valid UTF-8: {error}"))
375}
376
377/// Unified query for event channels with an optional exact scope and topic.
378#[derive(Debug, Clone, PartialEq, Eq, Hash)]
379pub struct EventChannelQuery {
380    /// Exact scope. `None` is reserved for administrative all-channel queries.
381    scope: Option<EventScope>,
382    topic: Option<String>,
383}
384
385impl EventChannelQuery {
386    /// Query all event channels (no filters)
387    pub fn all() -> Self {
388        Self {
389            scope: None,
390            topic: None,
391        }
392    }
393
394    /// Query event channels in a specific namespace
395    pub fn namespace(namespace: impl Into<String>) -> Self {
396        Self {
397            scope: Some(EventScope::Namespace {
398                name: namespace.into(),
399            }),
400            topic: None,
401        }
402    }
403
404    pub fn namespace_topic(namespace: impl Into<String>, topic: impl Into<String>) -> Self {
405        Self {
406            scope: Some(EventScope::Namespace {
407                name: namespace.into(),
408            }),
409            topic: Some(topic.into()),
410        }
411    }
412
413    /// Query event channels for a specific component
414    pub fn component(namespace: impl Into<String>, component: impl Into<String>) -> Self {
415        Self {
416            scope: Some(EventScope::Component {
417                namespace: namespace.into(),
418                component: component.into(),
419            }),
420            topic: None,
421        }
422    }
423
424    /// Query event channels for a specific topic
425    pub fn topic(
426        namespace: impl Into<String>,
427        component: impl Into<String>,
428        topic: impl Into<String>,
429    ) -> Self {
430        Self {
431            scope: Some(EventScope::Component {
432                namespace: namespace.into(),
433                component: component.into(),
434            }),
435            topic: Some(topic.into()),
436        }
437    }
438
439    pub fn endpoint(endpoint: EndpointId) -> Self {
440        Self {
441            scope: Some(EventScope::Endpoint { endpoint }),
442            topic: None,
443        }
444    }
445
446    pub fn endpoint_topic(endpoint: EndpointId, topic: impl Into<String>) -> Self {
447        Self {
448            scope: Some(EventScope::Endpoint { endpoint }),
449            topic: Some(topic.into()),
450        }
451    }
452
453    /// Get the query specificity (0=all, 1=scope, 2=scope+topic).
454    pub fn scope_level(&self) -> u8 {
455        if self.topic.is_some() {
456            2
457        } else if self.scope.is_some() {
458            1
459        } else {
460            0
461        }
462    }
463}
464
465/// Unified query for semantic event sources with an optional exact scope and topic.
466#[derive(Debug, Clone, PartialEq, Eq, Hash)]
467pub struct EventSourceQuery {
468    /// Exact scope. `None` is reserved for administrative all-source queries.
469    scope: Option<EventScope>,
470    topic: Option<String>,
471}
472
473impl EventSourceQuery {
474    pub fn all() -> Self {
475        Self {
476            scope: None,
477            topic: None,
478        }
479    }
480
481    pub fn namespace(namespace: impl Into<String>) -> Self {
482        Self {
483            scope: Some(EventScope::Namespace {
484                name: namespace.into(),
485            }),
486            topic: None,
487        }
488    }
489
490    pub fn namespace_topic(namespace: impl Into<String>, topic: impl Into<String>) -> Self {
491        Self {
492            scope: Some(EventScope::Namespace {
493                name: namespace.into(),
494            }),
495            topic: Some(topic.into()),
496        }
497    }
498
499    pub fn component(namespace: impl Into<String>, component: impl Into<String>) -> Self {
500        Self {
501            scope: Some(EventScope::Component {
502                namespace: namespace.into(),
503                component: component.into(),
504            }),
505            topic: None,
506        }
507    }
508
509    pub fn topic(
510        namespace: impl Into<String>,
511        component: impl Into<String>,
512        topic: impl Into<String>,
513    ) -> Self {
514        Self {
515            scope: Some(EventScope::Component {
516                namespace: namespace.into(),
517                component: component.into(),
518            }),
519            topic: Some(topic.into()),
520        }
521    }
522
523    pub fn endpoint(endpoint: EndpointId) -> Self {
524        Self {
525            scope: Some(EventScope::Endpoint { endpoint }),
526            topic: None,
527        }
528    }
529
530    pub fn endpoint_topic(endpoint: EndpointId, topic: impl Into<String>) -> Self {
531        Self {
532            scope: Some(EventScope::Endpoint { endpoint }),
533            topic: Some(topic.into()),
534        }
535    }
536
537    /// Get the query specificity (0=all, 1=scope, 2=scope+topic).
538    pub fn scope_level(&self) -> u8 {
539        if self.topic.is_some() {
540            2
541        } else if self.scope.is_some() {
542            1
543        } else {
544            0
545        }
546    }
547}
548
549/// Specification for registering objects in the discovery plane
550/// Represents the input to the register() operation
551#[derive(Debug, Clone, PartialEq, Eq)]
552pub enum DiscoverySpec {
553    /// Endpoint specification for registration
554    Endpoint {
555        namespace: String,
556        component: String,
557        endpoint: String,
558        /// Transport type and routing information
559        transport: TransportType,
560        /// Optional execution device for this endpoint instance.
561        /// Used by hetero routing to distinguish CPU and CUDA workers.
562        device_type: Option<DeviceType>,
563        /// Payload codec accepted by this endpoint's request-plane worker.
564        /// `None` represents a legacy JSON-only worker.
565        request_plane_codec: Option<RequestPlanePayloadCodec>,
566    },
567    Model {
568        namespace: String,
569        component: String,
570        endpoint: String,
571        /// ModelDeploymentCard serialized as JSON
572        /// This allows lib/runtime to remain independent of lib/llm types
573        /// DiscoverySpec.from_model() and DiscoveryInstance.deserialize_model() are ergonomic helpers to create and deserialize the model card.
574        card_json: serde_json::Value,
575        /// Optional suffix appended after instance_id in the key path (e.g., for LoRA adapters)
576        /// Key format: {namespace}/{component}/{endpoint}/{instance_id}[/{model_suffix}]
577        model_suffix: Option<String>,
578    },
579    /// Event plane channel specification
580    /// Used for registering event publishers/subscribers for discovery
581    EventChannel {
582        scope: EventScope,
583        /// Topic name for this channel (e.g., "kv-events", "kv-metrics")
584        topic: String,
585        /// Unique identity of this publisher incarnation.
586        ///
587        /// A process can host multiple publishers for the same topic, so event
588        /// channels cannot use the process-level discovery instance ID.
589        publisher_id: u64,
590        /// Event transport type (NATS subject prefix or ZMQ endpoint)
591        transport: EventTransport,
592    },
593    /// Semantic source of events, independent of event transport discovery.
594    EventSource {
595        scope: EventScope,
596        topic: String,
597        /// Unique identity of this source incarnation.
598        publisher_id: u64,
599        /// Domain-specific source descriptor owned by the consuming crate.
600        metadata: serde_json::Value,
601    },
602}
603
604impl DiscoverySpec {
605    /// Creates a Model discovery spec from a serializable type
606    /// The card will be serialized to JSON to avoid cross-crate dependencies
607    pub fn from_model<T>(
608        namespace: String,
609        component: String,
610        endpoint: String,
611        card: &T,
612    ) -> Result<Self>
613    where
614        T: Serialize,
615    {
616        Self::from_model_with_suffix(namespace, component, endpoint, card, None)
617    }
618
619    /// Creates a Model discovery spec with an optional suffix (e.g., for LoRA adapters)
620    /// The suffix is appended after the instance_id in the key path
621    pub fn from_model_with_suffix<T>(
622        namespace: String,
623        component: String,
624        endpoint: String,
625        card: &T,
626        model_suffix: Option<String>,
627    ) -> Result<Self>
628    where
629        T: Serialize,
630    {
631        let card_json = serde_json::to_value(card)?;
632        Ok(Self::Model {
633            namespace,
634            component,
635            endpoint,
636            card_json,
637            model_suffix,
638        })
639    }
640
641    /// Converts this registration spec into a discovery instance.
642    ///
643    /// Endpoint and model specs use `default_instance_id`, normally the
644    /// discovery client's process-level ID. Event channel and source specs
645    /// already carry a publisher-level ID, so they use that instead.
646    pub fn into_instance(self, default_instance_id: u64) -> DiscoveryInstance {
647        match self {
648            Self::Endpoint {
649                namespace,
650                component,
651                endpoint,
652                transport,
653                device_type,
654                request_plane_codec,
655            } => DiscoveryInstance::Endpoint(crate::component::Instance {
656                namespace,
657                component,
658                endpoint,
659                instance_id: default_instance_id,
660                transport,
661                device_type,
662                request_plane_codec,
663            }),
664            Self::Model {
665                namespace,
666                component,
667                endpoint,
668                card_json,
669                model_suffix,
670            } => DiscoveryInstance::Model {
671                namespace,
672                component,
673                endpoint,
674                instance_id: default_instance_id,
675                card_json,
676                model_suffix,
677            },
678            Self::EventChannel {
679                scope,
680                topic,
681                publisher_id,
682                transport,
683            } => DiscoveryInstance::EventChannel {
684                scope,
685                topic,
686                instance_id: publisher_id,
687                transport,
688            },
689            Self::EventSource {
690                scope,
691                topic,
692                publisher_id,
693                metadata,
694            } => DiscoveryInstance::EventSource {
695                scope,
696                topic,
697                publisher_id,
698                metadata,
699            },
700        }
701    }
702
703    /// Compatibility alias for [`DiscoverySpec::into_instance`].
704    pub fn with_instance_id(self, default_instance_id: u64) -> DiscoveryInstance {
705        self.into_instance(default_instance_id)
706    }
707}
708
709/// Registered instances in the discovery plane
710/// Represents objects that have been successfully registered with an instance ID
711#[derive(Debug, Clone, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
712#[serde(tag = "type")]
713pub enum DiscoveryInstance {
714    /// Registered endpoint instance - wraps the component::Instance directly
715    Endpoint(crate::component::Instance),
716    Model {
717        namespace: String,
718        component: String,
719        endpoint: String,
720        instance_id: u64,
721        /// ModelDeploymentCard serialized as JSON
722        /// This allows lib/runtime to remain independent of lib/llm types
723        card_json: serde_json::Value,
724        /// Optional suffix appended after instance_id in the key path (e.g., for LoRA adapters)
725        #[serde(default, skip_serializing_if = "Option::is_none")]
726        model_suffix: Option<String>,
727    },
728    /// Registered event channel instance for event plane pub/sub
729    EventChannel {
730        scope: EventScope,
731        /// Topic name for this channel (e.g., "kv-events", "kv-metrics")
732        topic: String,
733        instance_id: u64,
734        /// Event transport type (NATS subject prefix or ZMQ endpoint)
735        transport: EventTransport,
736    },
737    /// Registered semantic event source.
738    EventSource {
739        scope: EventScope,
740        topic: String,
741        publisher_id: u64,
742        metadata: serde_json::Value,
743    },
744}
745
746/// Validate an idempotent registration for one semantic event-source incarnation.
747///
748/// NOTE: Descriptor immutability belongs to the generic discovery contract. Backends still
749/// perform their own atomic lookup/insert, but all of them use this comparison so an identical
750/// registration succeeds and a changed descriptor preserves the original record.
751pub(crate) fn validate_event_source_reregistration(
752    existing: &DiscoveryInstance,
753    candidate: &DiscoveryInstance,
754) -> Result<()> {
755    let DiscoveryInstanceId::EventSource(existing_id) = existing.id() else {
756        anyhow::bail!("existing discovery record is not an event source")
757    };
758    if candidate.id() != DiscoveryInstanceId::EventSource(existing_id.clone()) {
759        anyhow::bail!("event source re-registration changed its identity")
760    }
761    if existing != candidate {
762        anyhow::bail!(
763            "Event source incarnation '{}' cannot change its descriptor",
764            existing_id.to_path()
765        )
766    }
767    Ok(())
768}
769
770impl DiscoveryInstance {
771    /// Returns the instance ID for this discovery instance
772    pub fn instance_id(&self) -> u64 {
773        match self {
774            Self::Endpoint(inst) => inst.instance_id,
775            Self::Model { instance_id, .. } => *instance_id,
776            Self::EventChannel { instance_id, .. } => *instance_id,
777            Self::EventSource { publisher_id, .. } => *publisher_id,
778        }
779    }
780
781    /// Deserializes the model JSON into the specified type T
782    /// Returns an error if this is not a Model instance or if deserialization fails
783    pub fn deserialize_model<T>(&self) -> Result<T>
784    where
785        T: for<'de> Deserialize<'de>,
786    {
787        match self {
788            Self::Model { card_json, .. } => Ok(serde_json::from_value(card_json.clone())?),
789            Self::Endpoint(_) => {
790                anyhow::bail!("Cannot deserialize model from Endpoint instance")
791            }
792            Self::EventChannel { .. } => {
793                anyhow::bail!("Cannot deserialize model from EventChannel instance")
794            }
795            Self::EventSource { .. } => {
796                anyhow::bail!("Cannot deserialize model from EventSource instance")
797            }
798        }
799    }
800
801    /// Extracts the unique identifier for this discovery instance
802    /// Used for tracking, diffing, and removal events
803    pub fn id(&self) -> DiscoveryInstanceId {
804        match self {
805            Self::Endpoint(inst) => DiscoveryInstanceId::Endpoint(EndpointInstanceId {
806                namespace: inst.namespace.clone(),
807                component: inst.component.clone(),
808                endpoint: inst.endpoint.clone(),
809                instance_id: inst.instance_id,
810            }),
811            Self::Model {
812                namespace,
813                component,
814                endpoint,
815                instance_id,
816                model_suffix,
817                ..
818            } => DiscoveryInstanceId::Model(ModelCardInstanceId {
819                namespace: namespace.clone(),
820                component: component.clone(),
821                endpoint: endpoint.clone(),
822                instance_id: *instance_id,
823                model_suffix: model_suffix.clone(),
824            }),
825            Self::EventChannel {
826                scope,
827                topic,
828                instance_id,
829                ..
830            } => DiscoveryInstanceId::EventChannel(EventChannelInstanceId {
831                scope: scope.clone(),
832                topic: topic.clone(),
833                instance_id: *instance_id,
834            }),
835            Self::EventSource {
836                scope,
837                topic,
838                publisher_id,
839                ..
840            } => DiscoveryInstanceId::EventSource(EventSourceInstanceId {
841                scope: scope.clone(),
842                topic: topic.clone(),
843                publisher_id: *publisher_id,
844            }),
845        }
846    }
847}
848
849/// Unique identifier for an endpoint instance
850#[derive(Debug, Clone, PartialEq, Eq, Hash, Serialize, Deserialize)]
851pub struct EndpointInstanceId {
852    pub namespace: String,
853    pub component: String,
854    pub endpoint: String,
855    pub instance_id: u64,
856}
857
858impl EndpointInstanceId {
859    /// Converts to a path string.
860    pub fn to_path(&self) -> String {
861        format!(
862            "{}/{}/{}/{:x}",
863            self.namespace, self.component, self.endpoint, self.instance_id
864        )
865    }
866
867    /// Parses an endpoint instance path.
868    pub fn from_path(path: &str) -> Result<Self> {
869        let parts: Vec<&str> = path.split('/').collect();
870        if parts.len() != 4 {
871            anyhow::bail!(
872                "Invalid EndpointInstanceId path: expected 4 parts, got {}",
873                parts.len()
874            );
875        }
876        Ok(Self {
877            namespace: parts[0].to_string(),
878            component: parts[1].to_string(),
879            endpoint: parts[2].to_string(),
880            instance_id: u64::from_str_radix(parts[3], 16)
881                .map_err(|e| anyhow::anyhow!("Invalid instance_id hex: {}", e))?,
882        })
883    }
884}
885
886/// Unique identifier for a model card instance
887/// The combination of (namespace, component, endpoint, instance_id, model_suffix) uniquely identifies a model card
888#[derive(Debug, Clone, PartialEq, Eq, Hash, Serialize, Deserialize)]
889pub struct ModelCardInstanceId {
890    pub namespace: String,
891    pub component: String,
892    pub endpoint: String,
893    pub instance_id: u64,
894    /// None for base models, Some(slug) for LoRA adapters
895    pub model_suffix: Option<String>,
896}
897
898/// Unique identifier for an event channel instance
899#[derive(Debug, Clone, PartialEq, Eq, Hash, Serialize, Deserialize)]
900pub struct EventChannelInstanceId {
901    pub scope: EventScope,
902    /// Topic name for this channel (e.g., "kv-events", "kv-metrics")
903    pub topic: String,
904    pub instance_id: u64,
905}
906
907impl EventChannelInstanceId {
908    /// Converts to a delimiter-safe path containing the exact channel scope.
909    pub fn to_path(&self) -> String {
910        format!(
911            "{}/topic/{}/{:x}",
912            self.scope.path_prefix(),
913            encode_event_segment(&self.topic),
914            self.instance_id
915        )
916    }
917
918    /// Parses a path produced by [`Self::to_path`].
919    pub fn from_path(path: &str) -> Result<Self> {
920        let parts: Vec<&str> = path.split('/').collect();
921        let (scope, topic_index, instance_index) = match parts.as_slice() {
922            ["namespace", namespace, "topic", _, _] => (
923                EventScope::Namespace {
924                    name: decode_event_segment(namespace)?,
925                },
926                3,
927                4,
928            ),
929            [
930                "namespace",
931                namespace,
932                "component",
933                component,
934                "topic",
935                _,
936                _,
937            ] => (
938                EventScope::Component {
939                    namespace: decode_event_segment(namespace)?,
940                    component: decode_event_segment(component)?,
941                },
942                5,
943                6,
944            ),
945            [
946                "namespace",
947                namespace,
948                "component",
949                component,
950                "endpoint",
951                endpoint,
952                "topic",
953                _,
954                _,
955            ] => (
956                EventScope::Endpoint {
957                    endpoint: EndpointId {
958                        namespace: decode_event_segment(namespace)?,
959                        component: decode_event_segment(component)?,
960                        name: decode_event_segment(endpoint)?,
961                    },
962                },
963                7,
964                8,
965            ),
966            _ => anyhow::bail!("invalid EventChannelInstanceId path: {path}"),
967        };
968        Ok(Self {
969            scope,
970            topic: decode_event_segment(parts[topic_index])?,
971            instance_id: u64::from_str_radix(parts[instance_index], 16)
972                .map_err(|e| anyhow::anyhow!("Invalid instance_id hex: {}", e))?,
973        })
974    }
975}
976
977/// Unique identifier for a semantic event source incarnation.
978#[derive(Debug, Clone, PartialEq, Eq, Hash, Serialize, Deserialize)]
979pub struct EventSourceInstanceId {
980    pub scope: EventScope,
981    pub topic: String,
982    pub publisher_id: u64,
983}
984
985impl EventSourceInstanceId {
986    /// Converts to a delimiter-safe path containing the exact source scope.
987    pub fn to_path(&self) -> String {
988        format!(
989            "{}/topic/{}/{:x}",
990            self.scope.path_prefix(),
991            encode_event_segment(&self.topic),
992            self.publisher_id
993        )
994    }
995
996    /// Parses a path produced by [`Self::to_path`].
997    pub fn from_path(path: &str) -> Result<Self> {
998        let channel_id = EventChannelInstanceId::from_path(path)
999            .with_context(|| format!("invalid EventSourceInstanceId path: {path}"))?;
1000        Ok(Self {
1001            scope: channel_id.scope,
1002            topic: channel_id.topic,
1003            publisher_id: channel_id.instance_id,
1004        })
1005    }
1006}
1007
1008impl ModelCardInstanceId {
1009    /// Converts to a path string: `{namespace}/{component}/{endpoint}/{instance_id:x}[/{model_suffix}]`
1010    pub fn to_path(&self) -> String {
1011        match &self.model_suffix {
1012            Some(suffix) => format!(
1013                "{}/{}/{}/{:x}/{}",
1014                self.namespace, self.component, self.endpoint, self.instance_id, suffix
1015            ),
1016            None => format!(
1017                "{}/{}/{}/{:x}",
1018                self.namespace, self.component, self.endpoint, self.instance_id
1019            ),
1020        }
1021    }
1022
1023    /// Parses from a path string: `{namespace}/{component}/{endpoint}/{instance_id:x}[/{model_suffix}]`
1024    pub fn from_path(path: &str) -> Result<Self> {
1025        let parts: Vec<&str> = path.split('/').collect();
1026        if parts.len() < 4 || parts.len() > 5 {
1027            anyhow::bail!(
1028                "Invalid ModelCardInstanceId path: expected 4 or 5 parts, got {}",
1029                parts.len()
1030            );
1031        }
1032        Ok(Self {
1033            namespace: parts[0].to_string(),
1034            component: parts[1].to_string(),
1035            endpoint: parts[2].to_string(),
1036            instance_id: u64::from_str_radix(parts[3], 16)
1037                .map_err(|e| anyhow::anyhow!("Invalid instance_id hex: {}", e))?,
1038            model_suffix: parts.get(4).map(|s| s.to_string()),
1039        })
1040    }
1041}
1042
1043/// Union of instance identifiers for different discovery object types
1044#[derive(Debug, Clone, PartialEq, Eq, Hash, Serialize, Deserialize)]
1045pub enum DiscoveryInstanceId {
1046    Endpoint(EndpointInstanceId),
1047    Model(ModelCardInstanceId),
1048    EventChannel(EventChannelInstanceId),
1049    EventSource(EventSourceInstanceId),
1050}
1051
1052impl DiscoveryInstanceId {
1053    /// Returns the raw instance_id regardless of variant type
1054    pub fn instance_id(&self) -> u64 {
1055        match self {
1056            Self::Endpoint(eid) => eid.instance_id,
1057            Self::Model(mid) => mid.instance_id,
1058            Self::EventChannel(ecid) => ecid.instance_id,
1059            Self::EventSource(esid) => esid.publisher_id,
1060        }
1061    }
1062
1063    /// Extracts the EndpointInstanceId, returning an error if this is a Model or EventChannel variant
1064    pub fn extract_endpoint_id(&self) -> Result<&EndpointInstanceId> {
1065        match self {
1066            Self::Endpoint(eid) => Ok(eid),
1067            Self::Model(_) => anyhow::bail!("Expected Endpoint variant, got Model"),
1068            Self::EventChannel(_) => anyhow::bail!("Expected Endpoint variant, got EventChannel"),
1069            Self::EventSource(_) => anyhow::bail!("Expected Endpoint variant, got EventSource"),
1070        }
1071    }
1072
1073    /// Extracts the ModelCardInstanceId, returning an error if this is an Endpoint or EventChannel variant
1074    pub fn extract_model_id(&self) -> Result<&ModelCardInstanceId> {
1075        match self {
1076            Self::Model(mid) => Ok(mid),
1077            Self::Endpoint(_) => anyhow::bail!("Expected Model variant, got Endpoint"),
1078            Self::EventChannel(_) => anyhow::bail!("Expected Model variant, got EventChannel"),
1079            Self::EventSource(_) => anyhow::bail!("Expected Model variant, got EventSource"),
1080        }
1081    }
1082
1083    /// Extracts the EventChannelInstanceId, returning an error if this is an Endpoint or Model variant
1084    pub fn extract_event_channel_id(&self) -> Result<&EventChannelInstanceId> {
1085        match self {
1086            Self::EventChannel(ecid) => Ok(ecid),
1087            Self::Endpoint(_) => anyhow::bail!("Expected EventChannel variant, got Endpoint"),
1088            Self::Model(_) => anyhow::bail!("Expected EventChannel variant, got Model"),
1089            Self::EventSource(_) => {
1090                anyhow::bail!("Expected EventChannel variant, got EventSource")
1091            }
1092        }
1093    }
1094
1095    /// Extracts the EventSourceInstanceId, returning an error for other variants.
1096    pub fn extract_event_source_id(&self) -> Result<&EventSourceInstanceId> {
1097        match self {
1098            Self::EventSource(esid) => Ok(esid),
1099            Self::Endpoint(_) => anyhow::bail!("Expected EventSource variant, got Endpoint"),
1100            Self::Model(_) => anyhow::bail!("Expected EventSource variant, got Model"),
1101            Self::EventChannel(_) => {
1102                anyhow::bail!("Expected EventSource variant, got EventChannel")
1103            }
1104        }
1105    }
1106}
1107
1108/// Events emitted by the discovery watch stream
1109#[derive(Debug, Clone, PartialEq, Eq)]
1110pub enum DiscoveryEvent {
1111    /// A new instance was added
1112    Added(DiscoveryInstance),
1113    /// An instance was removed (identified by its unique ID)
1114    Removed(DiscoveryInstanceId),
1115}
1116
1117/// Stream type for discovery events
1118pub type DiscoveryStream = Pin<Box<dyn Stream<Item = Result<DiscoveryEvent>> + Send>>;
1119
1120#[derive(Clone, Debug, PartialEq, Eq)]
1121struct ModelRegistrationIdentity {
1122    display_name: String,
1123    source_path: Option<String>,
1124    is_lora: bool,
1125}
1126
1127impl ModelRegistrationIdentity {
1128    fn base_identity(&self) -> &str {
1129        self.source_path.as_deref().unwrap_or(&self.display_name)
1130    }
1131
1132    fn is_compatible_with(&self, other: &Self) -> bool {
1133        if self.is_lora || other.is_lora {
1134            self.base_identity() == other.base_identity()
1135        } else {
1136            self.display_name == other.display_name
1137        }
1138    }
1139}
1140
1141fn extract_model_registration_identity(
1142    card_json: &serde_json::Value,
1143    model_suffix: Option<&str>,
1144) -> Result<ModelRegistrationIdentity> {
1145    let display_name = card_json
1146        .get("display_name")
1147        .and_then(serde_json::Value::as_str)
1148        .map(str::to_owned)
1149        .ok_or_else(|| {
1150            anyhow::anyhow!("failed to deserialize model display_name from card_json")
1151        })?;
1152    let source_path = card_json
1153        .get("source_path")
1154        .and_then(serde_json::Value::as_str)
1155        .map(str::to_owned);
1156    let is_lora =
1157        model_suffix.is_some() || card_json.get("lora").is_some_and(|value| !value.is_null());
1158
1159    Ok(ModelRegistrationIdentity {
1160        display_name,
1161        source_path,
1162        is_lora,
1163    })
1164}
1165
1166fn find_conflicting_model_name(
1167    instances: &[DiscoveryInstance],
1168    requested_identity: &ModelRegistrationIdentity,
1169) -> Result<Option<String>> {
1170    for instance in instances {
1171        if let DiscoveryInstance::Model {
1172            card_json,
1173            model_suffix,
1174            ..
1175        } = instance
1176        {
1177            let existing_identity =
1178                extract_model_registration_identity(card_json, model_suffix.as_deref())?;
1179            if !requested_identity.is_compatible_with(&existing_identity) {
1180                return Ok(Some(existing_identity.display_name));
1181            }
1182        }
1183    }
1184
1185    Ok(None)
1186}
1187
1188/// Discovery trait for service discovery across different backends
1189#[async_trait]
1190pub trait Discovery: Send + Sync {
1191    /// Returns a unique identifier for this worker (e.g lease id if using etcd or generated id for memory store)
1192    /// Endpoint and model objects created by this worker use this ID. Event
1193    /// channels and sources use a publisher-level ID because a worker can own
1194    /// more than one publisher for the same topic.
1195    fn instance_id(&self) -> u64;
1196
1197    /// Registers an object in the discovery plane with the instance id
1198    async fn register(&self, spec: DiscoverySpec) -> Result<DiscoveryInstance> {
1199        let (namespace, component, endpoint, requested_identity) = match &spec {
1200            DiscoverySpec::Model {
1201                namespace,
1202                component,
1203                endpoint,
1204                card_json,
1205                model_suffix,
1206                ..
1207            } => (
1208                namespace.clone(),
1209                component.clone(),
1210                endpoint.clone(),
1211                extract_model_registration_identity(card_json, model_suffix.as_deref())?,
1212            ),
1213            _ => return self.register_internal(spec).await,
1214        };
1215
1216        let query = DiscoveryQuery::EndpointModels {
1217            namespace: namespace.clone(),
1218            component: component.clone(),
1219            endpoint: endpoint.clone(),
1220        };
1221
1222        if let Some(conflicting_name) =
1223            find_conflicting_model_name(&self.list(query.clone()).await?, &requested_identity)?
1224        {
1225            let requested_name = &requested_identity.display_name;
1226            anyhow::bail!(
1227                "Cannot register model '{requested_name}' on endpoint '{namespace}/{component}/{endpoint}': a different model '{conflicting_name}' is already registered there"
1228            );
1229        }
1230
1231        let instance = self.register_internal(spec).await?;
1232
1233        if let Some(conflicting_name) =
1234            find_conflicting_model_name(&self.list(query).await?, &requested_identity)?
1235        {
1236            let requested_name = &requested_identity.display_name;
1237            if let Err(unregister_err) = self.unregister(instance.clone()).await {
1238                return Err(anyhow::anyhow!(
1239                    "Cannot register model '{requested_name}' on endpoint '{namespace}/{component}/{endpoint}': a different model '{conflicting_name}' is already registered there"
1240                ))
1241                .context(format!(
1242                    "failed to roll back conflicting model registration for instance {instance_id}: {unregister_err}",
1243                    instance_id = instance.instance_id()
1244                ));
1245            }
1246
1247            anyhow::bail!(
1248                "Cannot register model '{requested_name}' on endpoint '{namespace}/{component}/{endpoint}': a different model '{conflicting_name}' is already registered there"
1249            );
1250        }
1251
1252        Ok(instance)
1253    }
1254
1255    /// Backend-specific raw registration implementation.
1256    async fn register_internal(&self, spec: DiscoverySpec) -> Result<DiscoveryInstance>;
1257
1258    /// Unregisters an instance from the discovery plane
1259    async fn unregister(&self, instance: DiscoveryInstance) -> Result<()>;
1260
1261    /// Returns a list of currently registered instances for the given discovery query
1262    /// This is a one-time snapshot without watching for changes
1263    async fn list(&self, query: DiscoveryQuery) -> Result<Vec<DiscoveryInstance>>;
1264
1265    /// Returns a stream of discovery events (Added/Removed) for the given discovery query
1266    /// The optional cancellation token can be used to stop the watch stream
1267    async fn list_and_watch(
1268        &self,
1269        query: DiscoveryQuery,
1270        cancel_token: Option<CancellationToken>,
1271    ) -> Result<DiscoveryStream>;
1272
1273    /// Clean up resources held by this discovery backend.
1274    /// For KV store backends, this deletes owned registrations immediately rather than
1275    /// waiting for TTL expiry. Default is a no-op for backends that don't need cleanup.
1276    fn shutdown(&self) {}
1277}
1278
1279#[cfg(test)]
1280mod tests {
1281    use super::*;
1282
1283    #[test]
1284    fn endpoint_channel_id_path_round_trips_reserved_segments() {
1285        let id = EventChannelInstanceId {
1286            scope: EventScope::Endpoint {
1287                endpoint: EndpointId {
1288                    namespace: "ns.with/slash".to_string(),
1289                    component: "component.*".to_string(),
1290                    name: "endpoint.>/%".to_string(),
1291                },
1292            },
1293            topic: "kv.events/>".to_string(),
1294            instance_id: 0xfeed,
1295        };
1296
1297        let path = id.to_path();
1298        assert!(!path.contains("ns.with/slash"));
1299        assert_eq!(EventChannelInstanceId::from_path(&path).unwrap(), id);
1300    }
1301
1302    #[test]
1303    fn endpoint_source_id_path_round_trips_reserved_segments() {
1304        let id = EventSourceInstanceId {
1305            scope: EventScope::Endpoint {
1306                endpoint: EndpointId {
1307                    namespace: "ns.with/slash".to_string(),
1308                    component: "component.*".to_string(),
1309                    name: "endpoint.>/%".to_string(),
1310                },
1311            },
1312            topic: "kv.events/>".to_string(),
1313            publisher_id: 0xfeed,
1314        };
1315
1316        let path = id.to_path();
1317        assert!(!path.contains("ns.with/slash"));
1318        assert_eq!(EventSourceInstanceId::from_path(&path).unwrap(), id);
1319    }
1320
1321    #[test]
1322    fn endpoint_codec_metadata_round_trips_and_defaults_when_omitted() {
1323        let instance = DiscoverySpec::Endpoint {
1324            namespace: "default".to_string(),
1325            component: "worker".to_string(),
1326            endpoint: "generate".to_string(),
1327            transport: TransportType::Nats("worker.generate".to_string()),
1328            device_type: None,
1329            request_plane_codec: Some(RequestPlanePayloadCodec::Msgpack),
1330        }
1331        .into_instance(42);
1332
1333        let mut metadata = serde_json::to_value(&instance).unwrap();
1334        assert_eq!(metadata["request_plane_codec"], "msgpack");
1335        let round_trip: DiscoveryInstance = serde_json::from_value(metadata.clone()).unwrap();
1336        match round_trip {
1337            DiscoveryInstance::Endpoint(instance) => assert_eq!(
1338                instance.request_plane_codec,
1339                Some(RequestPlanePayloadCodec::Msgpack)
1340            ),
1341            _ => panic!("expected endpoint discovery metadata"),
1342        }
1343
1344        metadata
1345            .as_object_mut()
1346            .unwrap()
1347            .remove("request_plane_codec");
1348        let legacy: DiscoveryInstance = serde_json::from_value(metadata).unwrap();
1349        match legacy {
1350            DiscoveryInstance::Endpoint(instance) => {
1351                assert_eq!(instance.request_plane_codec, None)
1352            }
1353            _ => panic!("expected endpoint discovery metadata"),
1354        }
1355    }
1356}
dynamo_runtime/discovery/mod.rs

dynamo_runtime/discovery/
mod.rs
