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::collections::{HashMap, HashSet};
11use std::pin::Pin;
12use tokio_util::sync::CancellationToken;
13
14mod metadata;
15pub use metadata::{DiscoveryMetadata, MetadataSnapshot};
16
17mod registration;
18pub use registration::EndpointRegistrationLease;
19pub(crate) use registration::EndpointRegistrationManager;
20
21mod mock;
22pub use mock::{MockDiscovery, SharedMockRegistry};
23mod kv_store;
24pub use kv_store::KVStoreDiscovery;
25
26mod kube;
27pub use kube::{KubeDiscoveryClient, hash_container_name, hash_pod_name};
28
29pub mod utils;
30use crate::{
31    component::{DeviceType, Instance, TransportType},
32    pipeline::network::RequestPlanePayloadCodec,
33};
34pub use utils::watch_and_extract_field;
35
36/// Largest publisher ID exactly representable by float64-backed JSON metadata.
37pub(crate) const MAX_JSON_SAFE_PUBLISHER_ID: u64 = (1 << 53) - 1;
38
39/// Transport kind for event plane - used for configuration and env var selection.
40///
41/// This enum represents the *type* of transport without connection details.
42/// Use `EventTransport` when you need the full transport configuration.
43#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize, Default)]
44#[serde(rename_all = "snake_case")]
45pub enum EventTransportKind {
46    /// NATS Core pub/sub
47    Nats,
48    /// ZMQ pub/sub
49    #[default]
50    Zmq,
51}
52
53impl EventTransportKind {
54    /// Parse from environment variable `DYN_EVENT_PLANE`.
55    ///
56    /// Returns `Zmq` if the variable is not set or is empty: ZMQ is the default
57    /// event plane for all backends. NATS remains available as an explicit opt-in
58    /// (`DYN_EVENT_PLANE=nats`). When you have access to a runtime, prefer
59    /// `DistributedRuntime::default_event_transport_kind`, which resolves the same
60    /// default through the configured discovery backend.
61    ///
62    /// Returns an error for unrecognised values.
63    pub fn from_env() -> Result<Self> {
64        match std::env::var(crate::config::environment_names::event_plane::DYN_EVENT_PLANE)
65            .as_deref()
66        {
67            Ok("nats") => Ok(Self::Nats),
68            Ok("zmq") | Ok("") | Err(_) => Ok(Self::Zmq),
69            Ok(other) => anyhow::bail!(
70                "Invalid DYN_EVENT_PLANE value '{}'. Valid values: 'nats', 'zmq'",
71                other
72            ),
73        }
74    }
75
76    /// Logs a warning if an invalid value is encountered.
77    pub fn from_env_or_default() -> Self {
78        Self::from_env().unwrap_or_else(|e| {
79            tracing::warn!("{e}, defaulting to ZMQ");
80            Self::Zmq
81        })
82    }
83
84    /// Get the default codec for this transport kind.
85    /// NATS defaults to JSON, ZMQ defaults to MsgPack.
86    pub fn default_codec(&self) -> EventCodecKind {
87        match self {
88            Self::Nats => EventCodecKind::Json,
89            Self::Zmq => EventCodecKind::Msgpack,
90        }
91    }
92}
93
94/// Codec kind for event plane serialization.
95///
96/// This enum represents the serialization format for event envelopes and payloads.
97#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)]
98#[serde(rename_all = "snake_case")]
99pub enum EventCodecKind {
100    /// JSON codec - human-readable, good for debugging
101    Json,
102    /// MessagePack codec - compact binary format
103    Msgpack,
104}
105
106impl EventCodecKind {
107    /// Parse from environment variable `DYN_EVENT_PLANE_CODEC`.
108    /// Returns None if not set, allowing transport to select default.
109    /// Returns error for invalid values.
110    pub fn from_env() -> Result<Option<Self>> {
111        match std::env::var(crate::config::environment_names::event_plane::DYN_EVENT_PLANE_CODEC)
112            .as_deref()
113        {
114            Err(_) => Ok(None), // Not set
115            Ok("") => Ok(None), // Empty
116            Ok("json") => Ok(Some(Self::Json)),
117            Ok("msgpack") => Ok(Some(Self::Msgpack)),
118            Ok(other) => anyhow::bail!(
119                "Invalid DYN_EVENT_PLANE_CODEC value '{}'. Valid values: 'json', 'msgpack'",
120                other
121            ),
122        }
123    }
124
125    /// Parse from environment variable with transport-specific default.
126    /// Logs a warning if an invalid value is encountered.
127    pub fn from_env_or_transport_default(transport: EventTransportKind) -> Self {
128        Self::from_env()
129            .unwrap_or_else(|e| {
130                tracing::warn!(
131                    "{}, defaulting to {:?} for {:?}",
132                    e,
133                    transport.default_codec(),
134                    transport
135                );
136                None
137            })
138            .unwrap_or_else(|| transport.default_codec())
139    }
140}
141
142/// Transport configuration for event plane channels.
143///
144/// This enum carries both the transport kind and its connection configuration.
145/// Kept separate from `TransportType` (request plane) to distinguish event semantics.
146#[derive(Debug, Clone, PartialEq, Eq, Hash, Serialize, Deserialize)]
147#[serde(tag = "kind", content = "config")]
148pub enum EventTransport {
149    /// NATS Core pub/sub - subject prefix for the channel
150    Nats {
151        /// Subject prefix (e.g., "namespace.dynamo.component.backend")
152        subject_prefix: String,
153    },
154    /// ZMQ pub/sub - endpoint address (direct mode)
155    Zmq {
156        /// ZMQ endpoint (e.g., "tcp://host:port")
157        endpoint: String,
158    },
159    /// ZMQ broker endpoints (broker mode) - for discovery of brokers
160    ZmqBroker {
161        /// XSUB endpoints (publishers connect here)
162        xsub_endpoints: Vec<String>,
163        /// XPUB endpoints (subscribers connect here)
164        xpub_endpoints: Vec<String>,
165    },
166}
167
168impl EventTransport {
169    /// Get the transport kind
170    pub fn kind(&self) -> EventTransportKind {
171        match self {
172            Self::Nats { .. } => EventTransportKind::Nats,
173            Self::Zmq { .. } | Self::ZmqBroker { .. } => EventTransportKind::Zmq,
174        }
175    }
176
177    /// Create a NATS transport with the given subject prefix
178    pub fn nats(subject_prefix: impl Into<String>) -> Self {
179        Self::Nats {
180            subject_prefix: subject_prefix.into(),
181        }
182    }
183
184    /// Create a ZMQ transport with the given endpoint
185    pub fn zmq(endpoint: impl Into<String>) -> Self {
186        Self::Zmq {
187            endpoint: endpoint.into(),
188        }
189    }
190
191    /// Get the subject prefix (NATS) or endpoint (ZMQ)
192    /// For ZmqBroker, returns the first XSUB endpoint
193    pub fn address(&self) -> &str {
194        match self {
195            Self::Nats { subject_prefix } => subject_prefix,
196            Self::Zmq { endpoint } => endpoint,
197            Self::ZmqBroker { xsub_endpoints, .. } => {
198                xsub_endpoints.first().map(|s| s.as_str()).unwrap_or("")
199            }
200        }
201    }
202}
203
204/// Query key for prefix-based discovery queries
205/// Supports hierarchical queries from all endpoints down to specific endpoints
206#[derive(Debug, Clone, PartialEq, Eq, Hash)]
207pub enum DiscoveryQuery {
208    /// Query all endpoints in the system
209    AllEndpoints,
210    /// Query all endpoints in a specific namespace
211    NamespacedEndpoints {
212        namespace: String,
213    },
214    /// Query all endpoints in a namespace/component
215    ComponentEndpoints {
216        namespace: String,
217        component: String,
218    },
219    /// Query a specific endpoint
220    Endpoint {
221        namespace: String,
222        component: String,
223        endpoint: String,
224    },
225    AllModels,
226    NamespacedModels {
227        namespace: String,
228    },
229    ComponentModels {
230        namespace: String,
231        component: String,
232    },
233    EndpointModels {
234        namespace: String,
235        component: String,
236        endpoint: String,
237    },
238    /// Unified event channel query with optional scope filters
239    EventChannels(EventChannelQuery),
240    /// Semantic event source query with optional scope filters.
241    EventSources(EventSourceQuery),
242}
243
244/// Scope of an event channel.
245///
246/// Event scopes are exact ownership domains. A namespace-scoped query does not
247/// match component- or endpoint-scoped publishers in that namespace.
248#[derive(Debug, Clone, PartialEq, Eq, Hash, Serialize, Deserialize)]
249#[serde(tag = "kind", rename_all = "snake_case")]
250pub enum EventScope {
251    Namespace {
252        name: String,
253    },
254    Component {
255        namespace: String,
256        component: String,
257    },
258    Endpoint {
259        endpoint: EndpointId,
260    },
261}
262
263impl EventScope {
264    pub fn namespace(&self) -> &str {
265        match self {
266            Self::Namespace { name } => name,
267            Self::Component { namespace, .. } => namespace,
268            Self::Endpoint { endpoint } => &endpoint.namespace,
269        }
270    }
271
272    pub fn component(&self) -> Option<&str> {
273        match self {
274            Self::Namespace { .. } => None,
275            Self::Component { component, .. } => Some(component),
276            Self::Endpoint { endpoint } => Some(&endpoint.component),
277        }
278    }
279
280    pub fn endpoint(&self) -> Option<&EndpointId> {
281        match self {
282            Self::Endpoint { endpoint } => Some(endpoint),
283            Self::Namespace { .. } | Self::Component { .. } => None,
284        }
285    }
286
287    /// Canonical NATS/ZMQ-broker subject prefix for this scope.
288    pub fn subject_prefix(&self) -> String {
289        match self {
290            Self::Namespace { name } => {
291                format!("namespace.{}", encode_event_segment(name))
292            }
293            Self::Component {
294                namespace,
295                component,
296            } => format!(
297                "namespace.{}.component.{}",
298                encode_event_segment(namespace),
299                encode_event_segment(component)
300            ),
301            Self::Endpoint { endpoint } => format!(
302                "namespace.{}.component.{}.endpoint.{}",
303                encode_event_segment(&endpoint.namespace),
304                encode_event_segment(&endpoint.component),
305                encode_event_segment(&endpoint.name)
306            ),
307        }
308    }
309
310    /// Canonical subject/routing key for a topic in this exact scope.
311    pub fn subject(&self, topic: &str) -> String {
312        format!("{}.{}", self.subject_prefix(), encode_event_segment(topic))
313    }
314
315    pub(crate) fn path_prefix(&self) -> String {
316        match self {
317            Self::Namespace { name } => {
318                format!("namespace/{}", encode_event_segment(name))
319            }
320            Self::Component {
321                namespace,
322                component,
323            } => format!(
324                "namespace/{}/component/{}",
325                encode_event_segment(namespace),
326                encode_event_segment(component)
327            ),
328            Self::Endpoint { endpoint } => format!(
329                "namespace/{}/component/{}/endpoint/{}",
330                encode_event_segment(&endpoint.namespace),
331                encode_event_segment(&endpoint.component),
332                encode_event_segment(&endpoint.name)
333            ),
334        }
335    }
336}
337
338/// Percent-encode a subject/path segment while keeping common identifier
339/// characters readable. The encoding is reversible and prevents NATS wildcard
340/// or discovery path delimiters from changing the channel identity.
341pub(crate) fn encode_event_segment(value: &str) -> String {
342    let mut encoded = String::with_capacity(value.len());
343    for byte in value.bytes() {
344        if byte.is_ascii_alphanumeric() || matches!(byte, b'-' | b'_') {
345            encoded.push(char::from(byte));
346        } else {
347            use std::fmt::Write as _;
348            write!(encoded, "%{byte:02X}").expect("writing to String cannot fail");
349        }
350    }
351    encoded
352}
353
354fn decode_event_segment(value: &str) -> Result<String> {
355    let bytes = value.as_bytes();
356    let mut decoded = Vec::with_capacity(bytes.len());
357    let mut index = 0;
358    while index < bytes.len() {
359        if bytes[index] != b'%' {
360            decoded.push(bytes[index]);
361            index += 1;
362            continue;
363        }
364        if index + 2 >= bytes.len() {
365            anyhow::bail!("invalid percent-encoded event segment: {value}");
366        }
367        let hex = std::str::from_utf8(&bytes[index + 1..index + 3])?;
368        decoded
369            .push(u8::from_str_radix(hex, 16).map_err(|error| {
370                anyhow::anyhow!("invalid event segment escape %{hex}: {error}")
371            })?);
372        index += 3;
373    }
374    String::from_utf8(decoded)
375        .map_err(|error| anyhow::anyhow!("event segment is not valid UTF-8: {error}"))
376}
377
378/// Unified query for event channels with an optional exact scope and topic.
379#[derive(Debug, Clone, PartialEq, Eq, Hash)]
380pub struct EventChannelQuery {
381    /// Exact scope. `None` is reserved for administrative all-channel queries.
382    scope: Option<EventScope>,
383    topic: Option<String>,
384}
385
386impl EventChannelQuery {
387    /// Query all event channels (no filters)
388    pub fn all() -> Self {
389        Self {
390            scope: None,
391            topic: None,
392        }
393    }
394
395    /// Query event channels in a specific namespace
396    pub fn namespace(namespace: impl Into<String>) -> Self {
397        Self {
398            scope: Some(EventScope::Namespace {
399                name: namespace.into(),
400            }),
401            topic: None,
402        }
403    }
404
405    pub fn namespace_topic(namespace: impl Into<String>, topic: impl Into<String>) -> Self {
406        Self {
407            scope: Some(EventScope::Namespace {
408                name: namespace.into(),
409            }),
410            topic: Some(topic.into()),
411        }
412    }
413
414    /// Query event channels for a specific component
415    pub fn component(namespace: impl Into<String>, component: impl Into<String>) -> Self {
416        Self {
417            scope: Some(EventScope::Component {
418                namespace: namespace.into(),
419                component: component.into(),
420            }),
421            topic: None,
422        }
423    }
424
425    /// Query event channels for a specific topic
426    pub fn topic(
427        namespace: impl Into<String>,
428        component: impl Into<String>,
429        topic: impl Into<String>,
430    ) -> Self {
431        Self {
432            scope: Some(EventScope::Component {
433                namespace: namespace.into(),
434                component: component.into(),
435            }),
436            topic: Some(topic.into()),
437        }
438    }
439
440    pub fn endpoint(endpoint: EndpointId) -> Self {
441        Self {
442            scope: Some(EventScope::Endpoint { endpoint }),
443            topic: None,
444        }
445    }
446
447    pub fn endpoint_topic(endpoint: EndpointId, topic: impl Into<String>) -> Self {
448        Self {
449            scope: Some(EventScope::Endpoint { endpoint }),
450            topic: Some(topic.into()),
451        }
452    }
453
454    /// Get the query specificity (0=all, 1=scope, 2=scope+topic).
455    pub fn scope_level(&self) -> u8 {
456        if self.topic.is_some() {
457            2
458        } else if self.scope.is_some() {
459            1
460        } else {
461            0
462        }
463    }
464}
465
466/// Unified query for semantic event sources with an optional exact scope and topic.
467#[derive(Debug, Clone, PartialEq, Eq, Hash)]
468pub struct EventSourceQuery {
469    /// Exact scope. `None` is reserved for administrative all-source queries.
470    scope: Option<EventScope>,
471    topic: Option<String>,
472}
473
474impl EventSourceQuery {
475    pub fn all() -> Self {
476        Self {
477            scope: None,
478            topic: None,
479        }
480    }
481
482    pub fn namespace(namespace: impl Into<String>) -> Self {
483        Self {
484            scope: Some(EventScope::Namespace {
485                name: namespace.into(),
486            }),
487            topic: None,
488        }
489    }
490
491    pub fn namespace_topic(namespace: impl Into<String>, topic: impl Into<String>) -> Self {
492        Self {
493            scope: Some(EventScope::Namespace {
494                name: namespace.into(),
495            }),
496            topic: Some(topic.into()),
497        }
498    }
499
500    pub fn component(namespace: impl Into<String>, component: impl Into<String>) -> Self {
501        Self {
502            scope: Some(EventScope::Component {
503                namespace: namespace.into(),
504                component: component.into(),
505            }),
506            topic: None,
507        }
508    }
509
510    pub fn topic(
511        namespace: impl Into<String>,
512        component: impl Into<String>,
513        topic: impl Into<String>,
514    ) -> Self {
515        Self {
516            scope: Some(EventScope::Component {
517                namespace: namespace.into(),
518                component: component.into(),
519            }),
520            topic: Some(topic.into()),
521        }
522    }
523
524    pub fn endpoint(endpoint: EndpointId) -> Self {
525        Self {
526            scope: Some(EventScope::Endpoint { endpoint }),
527            topic: None,
528        }
529    }
530
531    pub fn endpoint_topic(endpoint: EndpointId, topic: impl Into<String>) -> Self {
532        Self {
533            scope: Some(EventScope::Endpoint { endpoint }),
534            topic: Some(topic.into()),
535        }
536    }
537
538    /// Get the query specificity (0=all, 1=scope, 2=scope+topic).
539    pub fn scope_level(&self) -> u8 {
540        if self.topic.is_some() {
541            2
542        } else if self.scope.is_some() {
543            1
544        } else {
545            0
546        }
547    }
548}
549
550/// Specification for registering objects in the discovery plane
551/// Represents the input to the register() operation
552#[derive(Debug, Clone, PartialEq, Eq)]
553pub enum DiscoverySpec {
554    /// Endpoint specification for registration
555    Endpoint {
556        namespace: String,
557        component: String,
558        endpoint: String,
559        /// Transport type and routing information
560        transport: TransportType,
561        /// Optional execution device for this endpoint instance.
562        /// Used by hetero routing to distinguish CPU and CUDA workers.
563        device_type: Option<DeviceType>,
564        /// Payload codec accepted by this endpoint's request-plane worker.
565        /// `None` represents a legacy JSON-only worker.
566        request_plane_codec: Option<RequestPlanePayloadCodec>,
567    },
568    Model {
569        namespace: String,
570        component: String,
571        endpoint: String,
572        /// ModelDeploymentCard serialized as JSON
573        /// This allows lib/runtime to remain independent of lib/llm types
574        /// DiscoverySpec.from_model() and DiscoveryInstance.deserialize_model() are ergonomic helpers to create and deserialize the model card.
575        card_json: serde_json::Value,
576        /// Optional suffix appended after instance_id in the key path (e.g., for LoRA adapters)
577        /// Key format: {namespace}/{component}/{endpoint}/{instance_id}[/{model_suffix}]
578        model_suffix: Option<String>,
579    },
580    /// Event plane channel specification
581    /// Used for registering event publishers/subscribers for discovery
582    EventChannel {
583        scope: EventScope,
584        /// Topic name for this channel (e.g., "kv-events", "kv-metrics")
585        topic: String,
586        /// Unique identity of this publisher incarnation.
587        ///
588        /// A process can host multiple publishers for the same topic, so event
589        /// channels cannot use the process-level discovery instance ID.
590        publisher_id: u64,
591        /// Event transport type (NATS subject prefix or ZMQ endpoint)
592        transport: EventTransport,
593    },
594    /// Semantic source of events, independent of event transport discovery.
595    EventSource {
596        scope: EventScope,
597        topic: String,
598        /// Unique identity of this source incarnation.
599        publisher_id: u64,
600        /// Domain-specific source descriptor owned by the consuming crate.
601        metadata: serde_json::Value,
602    },
603}
604
605impl DiscoverySpec {
606    /// Creates a Model discovery spec from a serializable type
607    /// The card will be serialized to JSON to avoid cross-crate dependencies
608    pub fn from_model<T>(
609        namespace: String,
610        component: String,
611        endpoint: String,
612        card: &T,
613    ) -> Result<Self>
614    where
615        T: Serialize,
616    {
617        Self::from_model_with_suffix(namespace, component, endpoint, card, None)
618    }
619
620    /// Creates a Model discovery spec with an optional suffix (e.g., for LoRA adapters)
621    /// The suffix is appended after the instance_id in the key path
622    pub fn from_model_with_suffix<T>(
623        namespace: String,
624        component: String,
625        endpoint: String,
626        card: &T,
627        model_suffix: Option<String>,
628    ) -> Result<Self>
629    where
630        T: Serialize,
631    {
632        let card_json = serde_json::to_value(card)?;
633        Ok(Self::Model {
634            namespace,
635            component,
636            endpoint,
637            card_json,
638            model_suffix,
639        })
640    }
641
642    /// Converts this registration spec into a discovery instance.
643    ///
644    /// Endpoint and model specs use `default_instance_id`, normally the
645    /// discovery client's process-level ID. Event channel and source specs
646    /// already carry a publisher-level ID, so they use that instead.
647    pub fn into_instance(self, default_instance_id: u64) -> DiscoveryInstance {
648        match self {
649            Self::Endpoint {
650                namespace,
651                component,
652                endpoint,
653                transport,
654                device_type,
655                request_plane_codec,
656            } => DiscoveryInstance::Endpoint(crate::component::Instance {
657                namespace,
658                component,
659                endpoint,
660                instance_id: default_instance_id,
661                transport,
662                device_type,
663                request_plane_codec,
664            }),
665            Self::Model {
666                namespace,
667                component,
668                endpoint,
669                card_json,
670                model_suffix,
671            } => DiscoveryInstance::Model {
672                namespace,
673                component,
674                endpoint,
675                instance_id: default_instance_id,
676                card_json,
677                model_suffix,
678            },
679            Self::EventChannel {
680                scope,
681                topic,
682                publisher_id,
683                transport,
684            } => DiscoveryInstance::EventChannel {
685                scope,
686                topic,
687                instance_id: publisher_id,
688                transport,
689            },
690            Self::EventSource {
691                scope,
692                topic,
693                publisher_id,
694                metadata,
695            } => DiscoveryInstance::EventSource {
696                scope,
697                topic,
698                publisher_id,
699                metadata,
700            },
701        }
702    }
703
704    /// Compatibility alias for [`DiscoverySpec::into_instance`].
705    pub fn with_instance_id(self, default_instance_id: u64) -> DiscoveryInstance {
706        self.into_instance(default_instance_id)
707    }
708}
709
710/// Registered instances in the discovery plane
711/// Represents objects that have been successfully registered with an instance ID
712#[derive(Debug, Clone, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
713#[serde(tag = "type")]
714pub enum DiscoveryInstance {
715    /// Registered endpoint instance - wraps the component::Instance directly
716    Endpoint(crate::component::Instance),
717    Model {
718        namespace: String,
719        component: String,
720        endpoint: String,
721        instance_id: u64,
722        /// ModelDeploymentCard serialized as JSON
723        /// This allows lib/runtime to remain independent of lib/llm types
724        card_json: serde_json::Value,
725        /// Optional suffix appended after instance_id in the key path (e.g., for LoRA adapters)
726        #[serde(default, skip_serializing_if = "Option::is_none")]
727        model_suffix: Option<String>,
728    },
729    /// Registered event channel instance for event plane pub/sub
730    EventChannel {
731        scope: EventScope,
732        /// Topic name for this channel (e.g., "kv-events", "kv-metrics")
733        topic: String,
734        instance_id: u64,
735        /// Event transport type (NATS subject prefix or ZMQ endpoint)
736        transport: EventTransport,
737    },
738    /// Registered semantic event source.
739    EventSource {
740        scope: EventScope,
741        topic: String,
742        publisher_id: u64,
743        metadata: serde_json::Value,
744    },
745}
746
747/// Validate an idempotent registration for one semantic event-source incarnation.
748///
749/// NOTE: Descriptor immutability belongs to the generic discovery contract. Backends still
750/// perform their own atomic lookup/insert, but all of them use this comparison so an identical
751/// registration succeeds and a changed descriptor preserves the original record.
752pub(crate) fn validate_event_source_reregistration(
753    existing: &DiscoveryInstance,
754    candidate: &DiscoveryInstance,
755) -> Result<()> {
756    let DiscoveryInstanceId::EventSource(existing_id) = existing.id() else {
757        anyhow::bail!("existing discovery record is not an event source")
758    };
759    if candidate.id() != DiscoveryInstanceId::EventSource(existing_id.clone()) {
760        anyhow::bail!("event source re-registration changed its identity")
761    }
762    if existing != candidate {
763        anyhow::bail!(
764            "Event source incarnation '{}' cannot change its descriptor",
765            existing_id.to_path()
766        )
767    }
768    Ok(())
769}
770
771impl DiscoveryInstance {
772    /// Returns the instance ID for this discovery instance
773    pub fn instance_id(&self) -> u64 {
774        match self {
775            Self::Endpoint(inst) => inst.instance_id,
776            Self::Model { instance_id, .. } => *instance_id,
777            Self::EventChannel { instance_id, .. } => *instance_id,
778            Self::EventSource { publisher_id, .. } => *publisher_id,
779        }
780    }
781
782    /// Deserializes the model JSON into the specified type T
783    /// Returns an error if this is not a Model instance or if deserialization fails
784    pub fn deserialize_model<T>(&self) -> Result<T>
785    where
786        T: for<'de> Deserialize<'de>,
787    {
788        match self {
789            Self::Model { card_json, .. } => Ok(serde_json::from_value(card_json.clone())?),
790            Self::Endpoint(_) => {
791                anyhow::bail!("Cannot deserialize model from Endpoint instance")
792            }
793            Self::EventChannel { .. } => {
794                anyhow::bail!("Cannot deserialize model from EventChannel instance")
795            }
796            Self::EventSource { .. } => {
797                anyhow::bail!("Cannot deserialize model from EventSource instance")
798            }
799        }
800    }
801
802    /// Extracts the unique identifier for this discovery instance
803    /// Used for tracking, diffing, and removal events
804    pub fn id(&self) -> DiscoveryInstanceId {
805        match self {
806            Self::Endpoint(inst) => DiscoveryInstanceId::Endpoint(EndpointInstanceId {
807                namespace: inst.namespace.clone(),
808                component: inst.component.clone(),
809                endpoint: inst.endpoint.clone(),
810                instance_id: inst.instance_id,
811            }),
812            Self::Model {
813                namespace,
814                component,
815                endpoint,
816                instance_id,
817                model_suffix,
818                ..
819            } => DiscoveryInstanceId::Model(ModelCardInstanceId {
820                namespace: namespace.clone(),
821                component: component.clone(),
822                endpoint: endpoint.clone(),
823                instance_id: *instance_id,
824                model_suffix: model_suffix.clone(),
825            }),
826            Self::EventChannel {
827                scope,
828                topic,
829                instance_id,
830                ..
831            } => DiscoveryInstanceId::EventChannel(EventChannelInstanceId {
832                scope: scope.clone(),
833                topic: topic.clone(),
834                instance_id: *instance_id,
835            }),
836            Self::EventSource {
837                scope,
838                topic,
839                publisher_id,
840                ..
841            } => DiscoveryInstanceId::EventSource(EventSourceInstanceId {
842                scope: scope.clone(),
843                topic: topic.clone(),
844                publisher_id: *publisher_id,
845            }),
846        }
847    }
848}
849
850/// Unique identifier for an endpoint instance
851#[derive(Debug, Clone, PartialEq, Eq, Hash, Serialize, Deserialize)]
852pub struct EndpointInstanceId {
853    pub namespace: String,
854    pub component: String,
855    pub endpoint: String,
856    pub instance_id: u64,
857}
858
859impl EndpointInstanceId {
860    /// Converts to a path string.
861    pub fn to_path(&self) -> String {
862        format!(
863            "{}/{}/{}/{:x}",
864            self.namespace, self.component, self.endpoint, self.instance_id
865        )
866    }
867
868    /// Parses an endpoint instance path.
869    pub fn from_path(path: &str) -> Result<Self> {
870        let parts: Vec<&str> = path.split('/').collect();
871        if parts.len() != 4 {
872            anyhow::bail!(
873                "Invalid EndpointInstanceId path: expected 4 parts, got {}",
874                parts.len()
875            );
876        }
877        Ok(Self {
878            namespace: parts[0].to_string(),
879            component: parts[1].to_string(),
880            endpoint: parts[2].to_string(),
881            instance_id: u64::from_str_radix(parts[3], 16)
882                .map_err(|e| anyhow::anyhow!("Invalid instance_id hex: {}", e))?,
883        })
884    }
885}
886
887/// Unique identifier for a model card instance
888/// The combination of (namespace, component, endpoint, instance_id, model_suffix) uniquely identifies a model card
889#[derive(Debug, Clone, PartialEq, Eq, Hash, Serialize, Deserialize)]
890pub struct ModelCardInstanceId {
891    pub namespace: String,
892    pub component: String,
893    pub endpoint: String,
894    pub instance_id: u64,
895    /// None for base models, Some(slug) for LoRA adapters
896    pub model_suffix: Option<String>,
897}
898
899/// Unique identifier for an event channel instance
900#[derive(Debug, Clone, PartialEq, Eq, Hash, Serialize, Deserialize)]
901pub struct EventChannelInstanceId {
902    pub scope: EventScope,
903    /// Topic name for this channel (e.g., "kv-events", "kv-metrics")
904    pub topic: String,
905    pub instance_id: u64,
906}
907
908impl EventChannelInstanceId {
909    /// Converts to a delimiter-safe path containing the exact channel scope.
910    pub fn to_path(&self) -> String {
911        format!(
912            "{}/topic/{}/{:x}",
913            self.scope.path_prefix(),
914            encode_event_segment(&self.topic),
915            self.instance_id
916        )
917    }
918
919    /// Parses a path produced by [`Self::to_path`].
920    pub fn from_path(path: &str) -> Result<Self> {
921        let parts: Vec<&str> = path.split('/').collect();
922        let (scope, topic_index, instance_index) = match parts.as_slice() {
923            ["namespace", namespace, "topic", _, _] => (
924                EventScope::Namespace {
925                    name: decode_event_segment(namespace)?,
926                },
927                3,
928                4,
929            ),
930            [
931                "namespace",
932                namespace,
933                "component",
934                component,
935                "topic",
936                _,
937                _,
938            ] => (
939                EventScope::Component {
940                    namespace: decode_event_segment(namespace)?,
941                    component: decode_event_segment(component)?,
942                },
943                5,
944                6,
945            ),
946            [
947                "namespace",
948                namespace,
949                "component",
950                component,
951                "endpoint",
952                endpoint,
953                "topic",
954                _,
955                _,
956            ] => (
957                EventScope::Endpoint {
958                    endpoint: EndpointId {
959                        namespace: decode_event_segment(namespace)?,
960                        component: decode_event_segment(component)?,
961                        name: decode_event_segment(endpoint)?,
962                    },
963                },
964                7,
965                8,
966            ),
967            _ => anyhow::bail!("invalid EventChannelInstanceId path: {path}"),
968        };
969        Ok(Self {
970            scope,
971            topic: decode_event_segment(parts[topic_index])?,
972            instance_id: u64::from_str_radix(parts[instance_index], 16)
973                .map_err(|e| anyhow::anyhow!("Invalid instance_id hex: {}", e))?,
974        })
975    }
976}
977
978/// Unique identifier for a semantic event source incarnation.
979#[derive(Debug, Clone, PartialEq, Eq, Hash, Serialize, Deserialize)]
980pub struct EventSourceInstanceId {
981    pub scope: EventScope,
982    pub topic: String,
983    pub publisher_id: u64,
984}
985
986impl EventSourceInstanceId {
987    /// Converts to a delimiter-safe path containing the exact source scope.
988    pub fn to_path(&self) -> String {
989        format!(
990            "{}/topic/{}/{:x}",
991            self.scope.path_prefix(),
992            encode_event_segment(&self.topic),
993            self.publisher_id
994        )
995    }
996
997    /// Parses a path produced by [`Self::to_path`].
998    pub fn from_path(path: &str) -> Result<Self> {
999        let channel_id = EventChannelInstanceId::from_path(path)
1000            .with_context(|| format!("invalid EventSourceInstanceId path: {path}"))?;
1001        Ok(Self {
1002            scope: channel_id.scope,
1003            topic: channel_id.topic,
1004            publisher_id: channel_id.instance_id,
1005        })
1006    }
1007}
1008
1009impl ModelCardInstanceId {
1010    /// Converts to a path string: `{namespace}/{component}/{endpoint}/{instance_id:x}[/{model_suffix}]`
1011    pub fn to_path(&self) -> String {
1012        match &self.model_suffix {
1013            Some(suffix) => format!(
1014                "{}/{}/{}/{:x}/{}",
1015                self.namespace, self.component, self.endpoint, self.instance_id, suffix
1016            ),
1017            None => format!(
1018                "{}/{}/{}/{:x}",
1019                self.namespace, self.component, self.endpoint, self.instance_id
1020            ),
1021        }
1022    }
1023
1024    /// Parses from a path string: `{namespace}/{component}/{endpoint}/{instance_id:x}[/{model_suffix}]`
1025    pub fn from_path(path: &str) -> Result<Self> {
1026        let parts: Vec<&str> = path.split('/').collect();
1027        if parts.len() < 4 || parts.len() > 5 {
1028            anyhow::bail!(
1029                "Invalid ModelCardInstanceId path: expected 4 or 5 parts, got {}",
1030                parts.len()
1031            );
1032        }
1033        Ok(Self {
1034            namespace: parts[0].to_string(),
1035            component: parts[1].to_string(),
1036            endpoint: parts[2].to_string(),
1037            instance_id: u64::from_str_radix(parts[3], 16)
1038                .map_err(|e| anyhow::anyhow!("Invalid instance_id hex: {}", e))?,
1039            model_suffix: parts.get(4).map(|s| s.to_string()),
1040        })
1041    }
1042}
1043
1044/// Union of instance identifiers for different discovery object types
1045#[derive(Debug, Clone, PartialEq, Eq, Hash, Serialize, Deserialize)]
1046pub enum DiscoveryInstanceId {
1047    Endpoint(EndpointInstanceId),
1048    Model(ModelCardInstanceId),
1049    EventChannel(EventChannelInstanceId),
1050    EventSource(EventSourceInstanceId),
1051}
1052
1053impl DiscoveryInstanceId {
1054    /// Returns the raw instance_id regardless of variant type
1055    pub fn instance_id(&self) -> u64 {
1056        match self {
1057            Self::Endpoint(eid) => eid.instance_id,
1058            Self::Model(mid) => mid.instance_id,
1059            Self::EventChannel(ecid) => ecid.instance_id,
1060            Self::EventSource(esid) => esid.publisher_id,
1061        }
1062    }
1063
1064    /// Extracts the EndpointInstanceId, returning an error if this is a Model or EventChannel variant
1065    pub fn extract_endpoint_id(&self) -> Result<&EndpointInstanceId> {
1066        match self {
1067            Self::Endpoint(eid) => Ok(eid),
1068            Self::Model(_) => anyhow::bail!("Expected Endpoint variant, got Model"),
1069            Self::EventChannel(_) => anyhow::bail!("Expected Endpoint variant, got EventChannel"),
1070            Self::EventSource(_) => anyhow::bail!("Expected Endpoint variant, got EventSource"),
1071        }
1072    }
1073
1074    /// Extracts the ModelCardInstanceId, returning an error if this is an Endpoint or EventChannel variant
1075    pub fn extract_model_id(&self) -> Result<&ModelCardInstanceId> {
1076        match self {
1077            Self::Model(mid) => Ok(mid),
1078            Self::Endpoint(_) => anyhow::bail!("Expected Model variant, got Endpoint"),
1079            Self::EventChannel(_) => anyhow::bail!("Expected Model variant, got EventChannel"),
1080            Self::EventSource(_) => anyhow::bail!("Expected Model variant, got EventSource"),
1081        }
1082    }
1083
1084    /// Extracts the EventChannelInstanceId, returning an error if this is an Endpoint or Model variant
1085    pub fn extract_event_channel_id(&self) -> Result<&EventChannelInstanceId> {
1086        match self {
1087            Self::EventChannel(ecid) => Ok(ecid),
1088            Self::Endpoint(_) => anyhow::bail!("Expected EventChannel variant, got Endpoint"),
1089            Self::Model(_) => anyhow::bail!("Expected EventChannel variant, got Model"),
1090            Self::EventSource(_) => {
1091                anyhow::bail!("Expected EventChannel variant, got EventSource")
1092            }
1093        }
1094    }
1095
1096    /// Extracts the EventSourceInstanceId, returning an error for other variants.
1097    pub fn extract_event_source_id(&self) -> Result<&EventSourceInstanceId> {
1098        match self {
1099            Self::EventSource(esid) => Ok(esid),
1100            Self::Endpoint(_) => anyhow::bail!("Expected EventSource variant, got Endpoint"),
1101            Self::Model(_) => anyhow::bail!("Expected EventSource variant, got Model"),
1102            Self::EventChannel(_) => {
1103                anyhow::bail!("Expected EventSource variant, got EventChannel")
1104            }
1105        }
1106    }
1107}
1108
1109/// Events emitted by the discovery watch stream
1110#[derive(Debug, Clone, PartialEq, Eq)]
1111pub enum DiscoveryEvent {
1112    /// A new instance was added.
1113    ///
1114    /// Endpoint watches also emit this event when the endpoint data changes without changing its
1115    /// [`DiscoveryInstanceId`]. Consumers of endpoint watches must replace the previous value.
1116    Added(DiscoveryInstance),
1117    /// The complete normalized taint set for an existing model card changed.
1118    ModelTaintsUpdated(ModelTaintsUpdate),
1119    /// An instance was removed (identified by its unique ID)
1120    Removed(DiscoveryInstanceId),
1121}
1122
1123/// A scoped, idempotent update to an existing model card's routing taints.
1124#[derive(Debug, Clone, PartialEq, Eq)]
1125pub struct ModelTaintsUpdate {
1126    pub id: ModelCardInstanceId,
1127    pub taints: Vec<String>,
1128}
1129
1130/// Stream type for discovery events
1131pub type DiscoveryStream = Pin<Box<dyn Stream<Item = Result<DiscoveryEvent>> + Send>>;
1132
1133#[derive(Clone, Debug, PartialEq, Eq)]
1134struct ModelRegistrationIdentity {
1135    display_name: String,
1136    aliases: Vec<String>,
1137    source_path: Option<String>,
1138    is_lora: bool,
1139}
1140
1141impl ModelRegistrationIdentity {
1142    fn base_identity(&self) -> &str {
1143        self.source_path.as_deref().unwrap_or(&self.display_name)
1144    }
1145
1146    fn is_compatible_with(&self, other: &Self) -> bool {
1147        if self.is_lora != other.is_lora {
1148            let (adapter, base) = if self.is_lora {
1149                (self, other)
1150            } else {
1151                (other, self)
1152            };
1153            adapter.base_identity() == base.base_identity()
1154                && adapter.display_name != base.display_name
1155                && !base.aliases.contains(&adapter.display_name)
1156        } else if self.is_lora {
1157            self.base_identity() == other.base_identity()
1158        } else {
1159            // Preserve existing same-name registration compatibility across local model paths.
1160            self.display_name == other.display_name
1161                || self.source_path.as_deref().is_some_and(|source| {
1162                    !source.is_empty()
1163                        && other.source_path.as_deref() == Some(source)
1164                        && !self.aliases.contains(&other.display_name)
1165                        && !other.aliases.contains(&self.display_name)
1166                        && !self
1167                            .aliases
1168                            .iter()
1169                            .any(|alias| other.aliases.contains(alias))
1170                })
1171        }
1172    }
1173}
1174
1175fn extract_model_registration_identity(
1176    card_json: &serde_json::Value,
1177    model_suffix: Option<&str>,
1178) -> Result<ModelRegistrationIdentity> {
1179    let display_name = card_json
1180        .get("display_name")
1181        .and_then(serde_json::Value::as_str)
1182        .map(str::to_owned)
1183        .ok_or_else(|| {
1184            anyhow::anyhow!("failed to deserialize model display_name from card_json")
1185        })?;
1186    let source_path = card_json
1187        .get("source_path")
1188        .and_then(serde_json::Value::as_str)
1189        .map(str::to_owned);
1190    let aliases = card_json
1191        .get("aliases")
1192        .and_then(serde_json::Value::as_array)
1193        .into_iter()
1194        .flatten()
1195        .filter_map(serde_json::Value::as_str)
1196        .map(str::to_owned)
1197        .collect();
1198    let is_lora =
1199        model_suffix.is_some() || card_json.get("lora").is_some_and(|value| !value.is_null());
1200
1201    Ok(ModelRegistrationIdentity {
1202        display_name,
1203        aliases,
1204        source_path,
1205        is_lora,
1206    })
1207}
1208
1209fn find_conflicting_model_name(
1210    instances: &[DiscoveryInstance],
1211    requested_identity: &ModelRegistrationIdentity,
1212) -> Result<Option<String>> {
1213    for instance in instances {
1214        if let DiscoveryInstance::Model {
1215            card_json,
1216            model_suffix,
1217            ..
1218        } = instance
1219        {
1220            let existing_identity =
1221                extract_model_registration_identity(card_json, model_suffix.as_deref())?;
1222            if !requested_identity.is_compatible_with(&existing_identity) {
1223                return Ok(Some(existing_identity.display_name));
1224            }
1225        }
1226    }
1227
1228    Ok(None)
1229}
1230
1231const TOPOLOGY_TAINT_PREFIX: &str = "dynamo.topology/";
1232
1233fn model_card_without_taints(
1234    instance: &DiscoveryInstance,
1235) -> Result<(serde_json::Value, HashSet<String>)> {
1236    let DiscoveryInstance::Model { card_json, .. } = instance else {
1237        anyhow::bail!("model update requires a model discovery instance")
1238    };
1239
1240    let mut card = card_json.clone();
1241    let runtime_config = card
1242        .get_mut("runtime_config")
1243        .and_then(serde_json::Value::as_object_mut)
1244        .context("model card is missing runtime_config")?;
1245    let taints = runtime_config
1246        .remove("taints")
1247        .unwrap_or_else(|| serde_json::Value::Array(Vec::new()));
1248    let taints = taints
1249        .as_array()
1250        .context("model card runtime_config.taints must be an array")?
1251        .iter()
1252        .map(|value| {
1253            value
1254                .as_str()
1255                .map(str::to_owned)
1256                .context("model card runtime_config.taints entries must be strings")
1257        })
1258        .collect::<Result<HashSet<_>>>()?;
1259
1260    Ok((card, taints))
1261}
1262
1263fn expected_topology_taints(card: &serde_json::Value) -> Result<HashSet<String>> {
1264    let Some(domains) = card
1265        .pointer("/runtime_config/topology_domains")
1266        .and_then(serde_json::Value::as_object)
1267    else {
1268        return Ok(HashSet::new());
1269    };
1270
1271    domains
1272        .iter()
1273        .map(|(domain, value)| {
1274            let value = value
1275                .as_str()
1276                .context("model card runtime_config.topology_domains values must be strings")?;
1277            Ok(format!(
1278                "{TOPOLOGY_TAINT_PREFIX}{}={}",
1279                domain.trim(),
1280                value.trim()
1281            ))
1282        })
1283        .collect()
1284}
1285
1286/// Validate the discovery-layer mutable boundary for model-card updates.
1287///
1288/// Model cards remain immutable after registration except for worker-managed
1289/// `runtime_config.taints`. Reserved topology taints are derived from the
1290/// immutable `topology_domains` map and must stay canonical.
1291#[derive(Debug)]
1292struct ValidatedModelTaintUpdate {
1293    existing_taints: HashSet<String>,
1294    candidate_taints: HashSet<String>,
1295}
1296
1297fn validate_model_taint_update(
1298    existing: &DiscoveryInstance,
1299    candidate: &DiscoveryInstance,
1300) -> Result<ValidatedModelTaintUpdate> {
1301    if existing.id() != candidate.id() {
1302        anyhow::bail!("model update cannot change discovery identity")
1303    }
1304
1305    let (existing_card, existing_taints) = model_card_without_taints(existing)?;
1306    let (candidate_card, candidate_taints) = model_card_without_taints(candidate)?;
1307    if existing_card != candidate_card {
1308        anyhow::bail!("model update can only change runtime_config.taints")
1309    }
1310
1311    let expected_topology = expected_topology_taints(&candidate_card)?;
1312    let actual_topology = candidate_taints
1313        .iter()
1314        .filter(|taint| taint.starts_with(TOPOLOGY_TAINT_PREFIX))
1315        .cloned()
1316        .collect::<HashSet<_>>();
1317    if actual_topology != expected_topology {
1318        anyhow::bail!(
1319            "reserved {TOPOLOGY_TAINT_PREFIX} taints must match runtime_config.topology_domains"
1320        )
1321    }
1322
1323    Ok(ValidatedModelTaintUpdate {
1324        existing_taints,
1325        candidate_taints,
1326    })
1327}
1328
1329/// Validate a same-ID model registration replay without replacing authoritative taints.
1330pub(crate) fn validate_model_reregistration(
1331    existing: &DiscoveryInstance,
1332    candidate: &DiscoveryInstance,
1333) -> Result<()> {
1334    validate_model_taint_update(existing, candidate).map(|_| ())
1335}
1336
1337fn sorted_taints(taints: HashSet<String>) -> Vec<String> {
1338    let mut taints = taints.into_iter().collect::<Vec<_>>();
1339    taints.sort_unstable();
1340    taints
1341}
1342
1343/// Classify a discovery value transition without widening `Added` into an upsert.
1344///
1345/// Model cards are immutable after registration except for `runtime_config.taints`.
1346/// Endpoints retain their existing replacement-as-Added behavior. Same-ID changes
1347/// to other discovery object types remain ignored.
1348pub(crate) fn classify_discovery_change(
1349    existing: Option<&DiscoveryInstance>,
1350    candidate: &DiscoveryInstance,
1351) -> Result<Option<DiscoveryEvent>> {
1352    let Some(existing) = existing else {
1353        return Ok(Some(DiscoveryEvent::Added(candidate.clone())));
1354    };
1355
1356    if existing == candidate {
1357        return Ok(None);
1358    }
1359
1360    if matches!(existing, DiscoveryInstance::Model { .. })
1361        && matches!(candidate, DiscoveryInstance::Model { .. })
1362    {
1363        let ValidatedModelTaintUpdate {
1364            existing_taints,
1365            candidate_taints,
1366        } = validate_model_taint_update(existing, candidate)?;
1367        if existing_taints == candidate_taints {
1368            return Ok(None);
1369        }
1370
1371        let DiscoveryInstanceId::Model(id) = candidate.id() else {
1372            unreachable!("model discovery instance must have a model id")
1373        };
1374        return Ok(Some(DiscoveryEvent::ModelTaintsUpdated(
1375            ModelTaintsUpdate {
1376                id,
1377                taints: sorted_taints(candidate_taints),
1378            },
1379        )));
1380    }
1381
1382    if matches!(candidate, DiscoveryInstance::Endpoint(_)) {
1383        Ok(Some(DiscoveryEvent::Added(candidate.clone())))
1384    } else {
1385        Ok(None)
1386    }
1387}
1388
1389/// Reconcile an authoritative snapshot while retaining the last valid value for
1390/// any model card that attempts an immutable mutation.
1391pub(crate) fn reconcile_discovery_snapshot(
1392    known: &HashMap<DiscoveryInstanceId, DiscoveryInstance>,
1393    current: HashMap<DiscoveryInstanceId, DiscoveryInstance>,
1394) -> (
1395    Vec<DiscoveryEvent>,
1396    HashMap<DiscoveryInstanceId, DiscoveryInstance>,
1397) {
1398    let mut events = known
1399        .keys()
1400        .filter(|id| !current.contains_key(*id))
1401        .cloned()
1402        .map(DiscoveryEvent::Removed)
1403        .collect::<Vec<_>>();
1404    let mut next = HashMap::with_capacity(current.len());
1405
1406    for (id, candidate) in current {
1407        match classify_discovery_change(known.get(&id), &candidate) {
1408            Ok(Some(event)) => {
1409                events.push(event);
1410                next.insert(id, candidate);
1411            }
1412            Ok(None) => {
1413                let retained = known.get(&id).cloned().unwrap_or(candidate);
1414                next.insert(id, retained);
1415            }
1416            Err(error) => {
1417                tracing::error!(
1418                    ?id,
1419                    %error,
1420                    "Rejecting immutable discovery model-card mutation"
1421                );
1422                if let Some(existing) = known.get(&id) {
1423                    next.insert(id, existing.clone());
1424                }
1425            }
1426        }
1427    }
1428
1429    (events, next)
1430}
1431
1432fn model_with_updated_taints(
1433    existing: &DiscoveryInstance,
1434    mut taints: HashSet<String>,
1435) -> Result<DiscoveryInstance> {
1436    if let Some(taint) = taints
1437        .iter()
1438        .find(|taint| taint.starts_with(TOPOLOGY_TAINT_PREFIX))
1439    {
1440        anyhow::bail!("taint '{taint}' uses reserved prefix '{TOPOLOGY_TAINT_PREFIX}'")
1441    }
1442
1443    let (card_without_taints, existing_taints) = model_card_without_taints(existing)?;
1444    taints.extend(expected_topology_taints(&card_without_taints)?);
1445    if taints == existing_taints {
1446        return Ok(existing.clone());
1447    }
1448
1449    let mut candidate = existing.clone();
1450    let DiscoveryInstance::Model { card_json, .. } = &mut candidate else {
1451        anyhow::bail!("model taint update requires a model discovery instance")
1452    };
1453    let runtime_config = card_json
1454        .get_mut("runtime_config")
1455        .and_then(serde_json::Value::as_object_mut)
1456        .context("model card is missing runtime_config")?;
1457    runtime_config.insert(
1458        "taints".to_string(),
1459        serde_json::Value::Array(
1460            sorted_taints(taints)
1461                .into_iter()
1462                .map(serde_json::Value::String)
1463                .collect(),
1464        ),
1465    );
1466    Ok(candidate)
1467}
1468
1469/// Discovery trait for service discovery across different backends
1470#[async_trait]
1471pub trait Discovery: Send + Sync {
1472    /// Returns a unique identifier for this worker (e.g lease id if using etcd or generated id for memory store)
1473    /// Endpoint and model objects created by this worker use this ID. Event
1474    /// channels and sources use a publisher-level ID because a worker can own
1475    /// more than one publisher for the same topic.
1476    fn instance_id(&self) -> u64;
1477
1478    /// Registers an object in the discovery plane with the instance id
1479    async fn register(&self, spec: DiscoverySpec) -> Result<DiscoveryInstance> {
1480        let (namespace, component, endpoint, requested_identity) = match &spec {
1481            DiscoverySpec::Model {
1482                namespace,
1483                component,
1484                endpoint,
1485                card_json,
1486                model_suffix,
1487                ..
1488            } => (
1489                namespace.clone(),
1490                component.clone(),
1491                endpoint.clone(),
1492                extract_model_registration_identity(card_json, model_suffix.as_deref())?,
1493            ),
1494            _ => return self.register_internal(spec).await,
1495        };
1496
1497        let query = DiscoveryQuery::EndpointModels {
1498            namespace: namespace.clone(),
1499            component: component.clone(),
1500            endpoint: endpoint.clone(),
1501        };
1502
1503        if let Some(conflicting_name) =
1504            find_conflicting_model_name(&self.list(query.clone()).await?, &requested_identity)?
1505        {
1506            let requested_name = &requested_identity.display_name;
1507            anyhow::bail!(
1508                "Cannot register model '{requested_name}' on endpoint '{namespace}/{component}/{endpoint}': a different model '{conflicting_name}' is already registered there"
1509            );
1510        }
1511
1512        let instance = self.register_internal(spec).await?;
1513
1514        if let Some(conflicting_name) =
1515            find_conflicting_model_name(&self.list(query).await?, &requested_identity)?
1516        {
1517            let requested_name = &requested_identity.display_name;
1518            if let Err(unregister_err) = self.unregister(instance.clone()).await {
1519                return Err(anyhow::anyhow!(
1520                    "Cannot register model '{requested_name}' on endpoint '{namespace}/{component}/{endpoint}': a different model '{conflicting_name}' is already registered there"
1521                ))
1522                .context(format!(
1523                    "failed to roll back conflicting model registration for instance {instance_id}: {unregister_err}",
1524                    instance_id = instance.instance_id()
1525                ));
1526            }
1527
1528            anyhow::bail!(
1529                "Cannot register model '{requested_name}' on endpoint '{namespace}/{component}/{endpoint}': a different model '{conflicting_name}' is already registered there"
1530            );
1531        }
1532
1533        Ok(instance)
1534    }
1535
1536    /// Backend-specific raw registration implementation.
1537    async fn register_internal(&self, spec: DiscoverySpec) -> Result<DiscoveryInstance>;
1538
1539    /// Replace the caller-managed taints of this worker's existing base model card.
1540    async fn update_model_taints(
1541        &self,
1542        id: ModelCardInstanceId,
1543        taints: HashSet<String>,
1544    ) -> Result<()> {
1545        if id.instance_id != self.instance_id() {
1546            anyhow::bail!(
1547                "cannot update model taints for worker {}; this discovery client owns worker {}",
1548                id.instance_id,
1549                self.instance_id()
1550            )
1551        }
1552        if id.model_suffix.is_some() {
1553            anyhow::bail!("model taint updates are supported only for base model cards")
1554        }
1555        self.update_model_taints_internal(id, taints).await
1556    }
1557
1558    /// Backend-specific authoritative read, taint-only mutation, and persistence.
1559    async fn update_model_taints_internal(
1560        &self,
1561        _id: ModelCardInstanceId,
1562        _taints: HashSet<String>,
1563    ) -> Result<()> {
1564        anyhow::bail!("model taint updates are not supported by this discovery backend")
1565    }
1566
1567    /// Unregisters an instance from the discovery plane
1568    async fn unregister(&self, instance: DiscoveryInstance) -> Result<()>;
1569
1570    /// Returns a list of currently registered instances for the given discovery query
1571    /// This is a one-time snapshot without watching for changes
1572    async fn list(&self, query: DiscoveryQuery) -> Result<Vec<DiscoveryInstance>>;
1573
1574    /// Returns a stream of discovery events (Added/Removed) for the given discovery query
1575    /// The optional cancellation token can be used to stop the watch stream
1576    async fn list_and_watch(
1577        &self,
1578        query: DiscoveryQuery,
1579        cancel_token: Option<CancellationToken>,
1580    ) -> Result<DiscoveryStream>;
1581
1582    /// Clean up resources held by this discovery backend.
1583    /// For KV store backends, this deletes owned registrations immediately rather than
1584    /// waiting for TTL expiry. Default is a no-op for backends that don't need cleanup.
1585    fn shutdown(&self) {}
1586}
1587
1588#[cfg(test)]
1589mod tests {
1590    use super::*;
1591
1592    #[test]
1593    fn endpoint_channel_id_path_round_trips_reserved_segments() {
1594        let id = EventChannelInstanceId {
1595            scope: EventScope::Endpoint {
1596                endpoint: EndpointId {
1597                    namespace: "ns.with/slash".to_string(),
1598                    component: "component.*".to_string(),
1599                    name: "endpoint.>/%".to_string(),
1600                },
1601            },
1602            topic: "kv.events/>".to_string(),
1603            instance_id: 0xfeed,
1604        };
1605
1606        let path = id.to_path();
1607        assert!(!path.contains("ns.with/slash"));
1608        assert_eq!(EventChannelInstanceId::from_path(&path).unwrap(), id);
1609    }
1610
1611    #[test]
1612    fn endpoint_source_id_path_round_trips_reserved_segments() {
1613        let id = EventSourceInstanceId {
1614            scope: EventScope::Endpoint {
1615                endpoint: EndpointId {
1616                    namespace: "ns.with/slash".to_string(),
1617                    component: "component.*".to_string(),
1618                    name: "endpoint.>/%".to_string(),
1619                },
1620            },
1621            topic: "kv.events/>".to_string(),
1622            publisher_id: 0xfeed,
1623        };
1624
1625        let path = id.to_path();
1626        assert!(!path.contains("ns.with/slash"));
1627        assert_eq!(EventSourceInstanceId::from_path(&path).unwrap(), id);
1628    }
1629
1630    #[test]
1631    fn endpoint_codec_metadata_round_trips_and_defaults_when_omitted() {
1632        let instance = DiscoverySpec::Endpoint {
1633            namespace: "default".to_string(),
1634            component: "worker".to_string(),
1635            endpoint: "generate".to_string(),
1636            transport: TransportType::Nats("worker.generate".to_string()),
1637            device_type: None,
1638            request_plane_codec: Some(RequestPlanePayloadCodec::Msgpack),
1639        }
1640        .into_instance(42);
1641
1642        let mut metadata = serde_json::to_value(&instance).unwrap();
1643        assert_eq!(metadata["request_plane_codec"], "msgpack");
1644        let round_trip: DiscoveryInstance = serde_json::from_value(metadata.clone()).unwrap();
1645        match round_trip {
1646            DiscoveryInstance::Endpoint(instance) => assert_eq!(
1647                instance.request_plane_codec,
1648                Some(RequestPlanePayloadCodec::Msgpack)
1649            ),
1650            _ => panic!("expected endpoint discovery metadata"),
1651        }
1652
1653        metadata
1654            .as_object_mut()
1655            .unwrap()
1656            .remove("request_plane_codec");
1657        let legacy: DiscoveryInstance = serde_json::from_value(metadata).unwrap();
1658        match legacy {
1659            DiscoveryInstance::Endpoint(instance) => {
1660                assert_eq!(instance.request_plane_codec, None)
1661            }
1662            _ => panic!("expected endpoint discovery metadata"),
1663        }
1664    }
1665}
1666
1667#[cfg(test)]
1668mod model_taint_update_tests {
1669    use super::*;
1670
1671    fn model_instance(taints: &[&str]) -> DiscoveryInstance {
1672        DiscoveryInstance::Model {
1673            namespace: "ns".to_string(),
1674            component: "worker".to_string(),
1675            endpoint: "generate".to_string(),
1676            instance_id: 7,
1677            card_json: serde_json::json!({
1678                "display_name": "model",
1679                "runtime_config": {
1680                    "taints": taints,
1681                    "topology_domains": {"zone": "west"}
1682                }
1683            }),
1684            model_suffix: None,
1685        }
1686    }
1687
1688    #[test]
1689    fn model_update_accepts_only_caller_managed_taint_changes() {
1690        let existing = model_instance(&["old", "dynamo.topology/zone=west"]);
1691        let candidate = model_instance(&["new", "dynamo.topology/zone=west"]);
1692
1693        validate_model_taint_update(&existing, &candidate).unwrap();
1694    }
1695
1696    #[test]
1697    fn model_update_rejects_immutable_card_changes() {
1698        let existing = model_instance(&["dynamo.topology/zone=west"]);
1699        let mut candidate = model_instance(&["dynamo.topology/zone=west"]);
1700        let DiscoveryInstance::Model { card_json, .. } = &mut candidate else {
1701            unreachable!()
1702        };
1703        card_json["display_name"] = serde_json::json!("other-model");
1704
1705        let error = validate_model_taint_update(&existing, &candidate).unwrap_err();
1706        assert!(
1707            error
1708                .to_string()
1709                .contains("can only change runtime_config.taints")
1710        );
1711    }
1712
1713    #[test]
1714    fn model_update_rejects_reserved_topology_taint_changes() {
1715        let existing = model_instance(&["dynamo.topology/zone=west"]);
1716        let candidate = model_instance(&["dynamo.topology/zone=east"]);
1717
1718        let error = validate_model_taint_update(&existing, &candidate).unwrap_err();
1719        assert!(
1720            error
1721                .to_string()
1722                .contains("must match runtime_config.topology_domains")
1723        );
1724    }
1725
1726    #[test]
1727    fn changed_taints_are_classified_as_a_scoped_normalized_event() {
1728        let existing = model_instance(&["old", "dynamo.topology/zone=west"]);
1729        let candidate = model_instance(&["gpu", "blue", "dynamo.topology/zone=west"]);
1730        let DiscoveryInstanceId::Model(id) = candidate.id() else {
1731            unreachable!()
1732        };
1733
1734        assert_eq!(
1735            classify_discovery_change(Some(&existing), &candidate).unwrap(),
1736            Some(DiscoveryEvent::ModelTaintsUpdated(ModelTaintsUpdate {
1737                id,
1738                taints: vec![
1739                    "blue".to_string(),
1740                    "dynamo.topology/zone=west".to_string(),
1741                    "gpu".to_string(),
1742                ],
1743            }))
1744        );
1745    }
1746
1747    #[test]
1748    fn taint_order_only_changes_are_no_ops() {
1749        let existing = model_instance(&["gpu", "dynamo.topology/zone=west"]);
1750        let candidate = model_instance(&["dynamo.topology/zone=west", "gpu"]);
1751
1752        assert_eq!(
1753            classify_discovery_change(Some(&existing), &candidate).unwrap(),
1754            None
1755        );
1756    }
1757
1758    #[test]
1759    fn update_api_derives_topology_taints_and_rejects_reserved_input() {
1760        let existing = model_instance(&["old", "dynamo.topology/zone=west"]);
1761        let updated =
1762            model_with_updated_taints(&existing, HashSet::from(["new".to_string()])).unwrap();
1763        let (_, taints) = model_card_without_taints(&updated).unwrap();
1764        assert_eq!(
1765            taints,
1766            HashSet::from(["new".to_string(), "dynamo.topology/zone=west".to_string()])
1767        );
1768        assert!(
1769            model_with_updated_taints(
1770                &existing,
1771                HashSet::from(["dynamo.topology/zone=east".to_string()])
1772            )
1773            .is_err()
1774        );
1775    }
1776}
dynamo_runtime/discovery/mod.rs

dynamo_runtime/discovery/
mod.rs
