watcher.rs 28.6 KB
Newer Older
1
// SPDX-FileCopyrightText: Copyright (c) 2024-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
Ryan Olson's avatar
Ryan Olson committed
2
3
// SPDX-License-Identifier: Apache-2.0

Ryan Olson's avatar
Ryan Olson committed
4
use std::sync::Arc;
5
use tokio::sync::Notify;
6
use tokio::sync::mpsc::Sender;
7

8
use anyhow::Context as _;
9
use futures::StreamExt;
Ryan Olson's avatar
Ryan Olson committed
10

Neelay Shah's avatar
Neelay Shah committed
11
use dynamo_runtime::{
12
    DistributedRuntime,
13
14
15
16
    discovery::{
        DiscoveryEvent, DiscoveryInstance, DiscoveryInstanceId, DiscoveryQuery, DiscoveryStream,
        ModelCardInstanceId,
    },
17
    pipeline::{
18
19
        ManyOut, Operator, RouterMode, SegmentSource, ServiceBackend, SingleIn, Source,
        network::egress::push_router::PushRouter,
20
    },
21
    protocols::{EndpointId, annotated::Annotated},
22
};
Ryan Olson's avatar
Ryan Olson committed
23

24
25
use crate::{
    backend::Backend,
26
    entrypoint::{self, EngineFactoryCallback, RouterConfig},
27
    http::service::metrics::Metrics,
28
    kv_router::PrefillRouter,
29
    model_card::ModelDeploymentCard,
30
31
    model_type::{ModelInput, ModelType},
    preprocessor::{OpenAIPreprocessor, PreprocessedEmbeddingRequest, prompt::PromptFormatter},
32
33
34
35
36
37
38
39
40
    protocols::{
        common::llm_backend::EmbeddingsEngineOutput,
        openai::{
            chat_completions::{
                NvCreateChatCompletionRequest, NvCreateChatCompletionStreamResponse,
            },
            completions::{NvCreateCompletionRequest, NvCreateCompletionResponse},
            embeddings::{NvCreateEmbeddingRequest, NvCreateEmbeddingResponse},
        },
41
        tensor::{NvCreateTensorRequest, NvCreateTensorResponse},
42
43
    },
};
44

45
use super::ModelManager;
46
use crate::namespace::is_global_namespace;
47

48
#[derive(Debug, Clone)]
49
pub enum ModelUpdate {
50
51
    Added(ModelDeploymentCard),
    Removed(ModelDeploymentCard),
52
53
}

54
pub struct ModelWatcher {
55
    manager: Arc<ModelManager>,
56
    drt: DistributedRuntime,
57
    router_config: RouterConfig,
58
    notify_on_model: Notify,
59
    model_update_tx: Option<Sender<ModelUpdate>>,
60
    engine_factory: Option<EngineFactoryCallback>,
61
    metrics: Arc<Metrics>,
Ryan Olson's avatar
Ryan Olson committed
62
63
}

64
65
66
67
const ALL_MODEL_TYPES: &[ModelType] = &[
    ModelType::Chat,
    ModelType::Completions,
    ModelType::Embedding,
68
    ModelType::TensorBased,
69
    ModelType::Prefill,
70
];
71

72
impl ModelWatcher {
73
    pub fn new(
74
        runtime: DistributedRuntime,
75
        model_manager: Arc<ModelManager>,
76
        router_config: RouterConfig,
77
        engine_factory: Option<EngineFactoryCallback>,
78
        metrics: Arc<Metrics>,
79
80
    ) -> ModelWatcher {
        Self {
81
            manager: model_manager,
82
            drt: runtime,
83
            router_config,
84
            notify_on_model: Notify::new(),
85
            model_update_tx: None,
86
            engine_factory,
87
            metrics,
88
        }
89
    }
Ryan Olson's avatar
Ryan Olson committed
90

91
92
93
94
    pub fn set_notify_on_model_update(&mut self, tx: Sender<ModelUpdate>) {
        self.model_update_tx = Some(tx);
    }

95
96
97
98
99
100
101
102
103
104
105
    /// Wait until we have at least one chat completions model and return it's name.
    pub async fn wait_for_chat_model(&self) -> String {
        // Loop in case it gets added and immediately deleted
        loop {
            if let Some(model_name) = self.manager.list_chat_completions_models().first() {
                return model_name.to_owned();
            }
            self.notify_on_model.notified().await
        }
    }

106
    /// Common watch logic with optional namespace filtering
107
108
109
110
111
    pub async fn watch(
        &self,
        mut discovery_stream: DiscoveryStream,
        target_namespace: Option<&str>,
    ) {
112
        let global_namespace = target_namespace.is_none_or(is_global_namespace);
113

114
115
116
117
118
119
120
121
122
        while let Some(result) = discovery_stream.next().await {
            let event = match result {
                Ok(event) => event,
                Err(err) => {
                    tracing::error!(%err, "Error in discovery stream");
                    continue;
                }
            };

123
            match event {
124
                DiscoveryEvent::Added(instance) => {
125
126
                    // Extract ModelCardInstanceId and card from the discovery instance
                    let (mcid, mut card) = match &instance {
127
128
129
130
131
                        DiscoveryInstance::Model {
                            namespace,
                            component,
                            endpoint,
                            instance_id,
132
                            model_suffix,
133
134
                            ..
                        } => {
135
                            let mcid = ModelCardInstanceId {
136
137
                                namespace: namespace.clone(),
                                component: component.clone(),
138
139
140
                                endpoint: endpoint.clone(),
                                instance_id: *instance_id,
                                model_suffix: model_suffix.clone(),
141
142
143
                            };

                            match instance.deserialize_model::<ModelDeploymentCard>() {
144
                                Ok(card) => (mcid, card),
145
146
147
148
149
150
151
152
153
154
                                Err(err) => {
                                    tracing::error!(%err, instance_id, "Failed to deserialize model card");
                                    continue;
                                }
                            }
                        }
                        _ => {
                            tracing::error!(
                                "Unexpected discovery instance type (expected ModelCard)"
                            );
155
156
157
                            continue;
                        }
                    };
158
159
160
161

                    // Filter by namespace if target_namespace is specified
                    if !global_namespace
                        && let Some(target_ns) = target_namespace
162
                        && mcid.namespace != target_ns
163
164
                    {
                        tracing::debug!(
165
                            model_namespace = mcid.namespace,
166
167
168
169
170
171
                            target_namespace = target_ns,
                            "Skipping model from different namespace"
                        );
                        continue;
                    }

172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
                    // If we already have a worker for this model, and the ModelDeploymentCard
                    // cards don't match, alert, and don't add the new instance
                    let can_add =
                        self.manager
                            .is_valid_checksum(card.model_type, card.name(), card.mdcsum());
                    if can_add.is_some_and(|is_valid| !is_valid) {
                        tracing::error!(
                            model_name = card.name(),
                            "Checksum for new model does not match existing model."
                        );

                        // TODO: mark that instance down in clients
                        // Not obvious how to do that given the current design
                        // Instances come from an `InstanceSource` in a `Client` in a `PushRouter`.
                        // Calling `report_instance_down` on the Client should do it (although
                        // needs more testing).
                        // The `PushRouter` is in `ModelMananger` (`self.manager` here), but inside
                        // interface `AsyncEngine` which only has a `generate` method.

                        continue;
                    }
193

194
                    match self.handle_put(&mcid, &mut card).await {
195
                        Ok(()) => {
196
                            tracing::info!(
197
                                model_name = card.name(),
198
                                namespace = mcid.namespace,
199
200
                                "added model"
                            );
201
                            self.notify_on_model.notify_waiters();
202
                        }
203
204
                        Err(err) => {
                            tracing::error!(
205
                                model_name = card.name(),
206
                                namespace = mcid.namespace,
207
                                error = format!("{err:#}"),
208
                                "Error adding model from discovery",
209
                            );
210
211
212
                        }
                    }
                }
213
214
215
216
                DiscoveryEvent::Removed(id) => {
                    // Extract ModelCardInstanceId from the removal event
                    let model_card_instance_id = match &id {
                        DiscoveryInstanceId::Model(mcid) => mcid,
217
                        DiscoveryInstanceId::Endpoint(_) | DiscoveryInstanceId::EventChannel(_) => {
218
219
220
221
222
223
                            tracing::error!(
                                "Unexpected discovery instance type in removal (expected Model)"
                            );
                            continue;
                        }
                    };
224

225
                    match self
226
                        .handle_delete(model_card_instance_id, target_namespace, global_namespace)
227
228
229
230
231
232
233
234
235
236
237
                        .await
                    {
                        Ok(Some(model_name)) => {
                            tracing::info!(model_name, "removed model");
                        }
                        Ok(None) => {
                            // There are other instances running this model, nothing to do
                        }
                        Err(e) => {
                            tracing::error!(error = %e, "error removing model");
                        }
238
                    }
239
                }
240
            }
Ryan Olson's avatar
Ryan Olson committed
241
242
243
        }
    }

244
245
    /// If the last instance running this model has gone delete it.
    /// Returns the name of the model we just deleted, if any.
246
247
    async fn handle_delete(
        &self,
248
        mcid: &ModelCardInstanceId,
249
250
251
        target_namespace: Option<&str>,
        is_global_namespace: bool,
    ) -> anyhow::Result<Option<String>> {
252
253
        let key = mcid.to_path();
        let card = match self.manager.remove_model_card(&key) {
254
            Some(card) => card,
255
            None => {
256
                anyhow::bail!("Missing ModelDeploymentCard for {}", key);
257
258
            }
        };
259
        let model_name = card.name().to_string();
260
        let active_instances = self
261
            .cards_for_model(&model_name, target_namespace, is_global_namespace)
262
263
264
            .await
            .with_context(|| model_name.clone())?;
        if !active_instances.is_empty() {
265
266
267
268
269
270
            tracing::debug!(
                model_name,
                target_namespace = ?target_namespace,
                active_instance_count = active_instances.len(),
                "Model has other active instances, not removing"
            );
271
272
            return Ok(None);
        }
273

274
        // Ignore the errors because model could be either type
275
276
277
        let chat_model_remove_err = self.manager.remove_chat_completions_model(&model_name);
        let completions_model_remove_err = self.manager.remove_completions_model(&model_name);
        let embeddings_model_remove_err = self.manager.remove_embeddings_model(&model_name);
278
        let tensor_model_remove_err = self.manager.remove_tensor_model(&model_name);
279
        let prefill_model_remove_err = self.manager.remove_prefill_model(&model_name);
280
281
282
283

        let mut chat_model_removed = false;
        let mut completions_model_removed = false;
        let mut embeddings_model_removed = false;
284
        let mut tensor_model_removed = false;
285
        let mut prefill_model_removed = false;
286
287
288
289
290
291
292
293
294
295
296

        if chat_model_remove_err.is_ok() && self.manager.list_chat_completions_models().is_empty() {
            chat_model_removed = true;
        }
        if completions_model_remove_err.is_ok() && self.manager.list_completions_models().is_empty()
        {
            completions_model_removed = true;
        }
        if embeddings_model_remove_err.is_ok() && self.manager.list_embeddings_models().is_empty() {
            embeddings_model_removed = true;
        }
297
298
299
        if tensor_model_remove_err.is_ok() && self.manager.list_tensor_models().is_empty() {
            tensor_model_removed = true;
        }
300
301
302
        if prefill_model_remove_err.is_ok() && self.manager.list_prefill_models().is_empty() {
            prefill_model_removed = true;
        }
303

304
305
306
307
        if !chat_model_removed
            && !completions_model_removed
            && !embeddings_model_removed
            && !tensor_model_removed
308
            && !prefill_model_removed
309
        {
310
            tracing::debug!(
311
                "No updates to send for model {}: chat_model_removed: {}, completions_model_removed: {}, embeddings_model_removed: {}, tensor_model_removed: {}, prefill_model_removed: {}",
312
313
314
                model_name,
                chat_model_removed,
                completions_model_removed,
315
                embeddings_model_removed,
316
317
                tensor_model_removed,
                prefill_model_removed
318
319
320
            );
        } else {
            for model_type in ALL_MODEL_TYPES {
321
                if ((chat_model_removed && *model_type == ModelType::Chat)
322
                    || (completions_model_removed && *model_type == ModelType::Completions)
323
                    || (embeddings_model_removed && *model_type == ModelType::Embedding)
324
325
                    || (tensor_model_removed && *model_type == ModelType::TensorBased)
                    || (prefill_model_removed && *model_type == ModelType::Prefill))
326
                    && let Some(tx) = &self.model_update_tx
327
                {
328
                    tx.send(ModelUpdate::Removed(card.clone())).await.ok();
329
330
331
                }
            }
        }
332

333
        Ok(Some(model_name))
334
    }
Ryan Olson's avatar
Ryan Olson committed
335

336
    // Handles a PUT event from store, this usually means adding a new model to the list of served
337
    // models.
338
339
    async fn handle_put(
        &self,
340
        mcid: &ModelCardInstanceId,
341
342
        card: &mut ModelDeploymentCard,
    ) -> anyhow::Result<()> {
343
344
        card.download_config().await?;

345
        let component = self
346
            .drt
347
348
349
            .namespace(&mcid.namespace)?
            .component(&mcid.component)?;
        let endpoint = component.endpoint(&mcid.endpoint);
350
        let client = endpoint.client().await?;
351
        tracing::debug!(model_name = card.name(), "adding model");
352
353
        self.manager
            .save_model_card(&mcid.to_path(), card.clone())?;
354

355
356
357
358
359
360
361
362
363
364
        // Skip duplicate registrations based on model type.
        // Prefill and decode models are tracked separately, so registering one
        // doesn't block the other (they can arrive in any order).
        let already_registered = if card.model_type.supports_prefill() {
            self.manager.has_prefill_model(card.name())
        } else {
            self.manager.has_decode_model(card.name())
        };

        if already_registered {
365
366
            tracing::debug!(
                model_name = card.name(),
367
                namespace = mcid.namespace,
368
                model_type = %card.model_type,
369
                "Model already registered, skipping"
370
371
372
373
374
375
376
            );
            return Ok(());
        }

        if let Some(tx) = &self.model_update_tx {
            tx.send(ModelUpdate::Added(card.clone())).await.ok();
        }
377

378
        let checksum = card.mdcsum();
379
380
381

        if card.model_input == ModelInput::Tokens
            && (card.model_type.supports_chat() || card.model_type.supports_completions())
382
383
384
385
386
        {
            // Case 1: Tokens + (Chat OR Completions OR Both)
            // A model that expects pre-processed requests meaning it's up to us whether we
            // handle Chat or Completions requests, so handle whatever the model supports.

387
            let endpoint = component.endpoint(&mcid.endpoint);
388
            let kv_chooser = if self.router_config.router_mode == RouterMode::KV {
389
390
                Some(
                    self.manager
391
392
393
394
395
                        .kv_chooser_for(
                            &endpoint,
                            card.kv_cache_block_size,
                            Some(self.router_config.kv_router_config),
                        )
396
397
398
399
400
                        .await?,
                )
            } else {
                None
            };
401

402
            // This is expensive, we are loading ~10MiB JSON, so only do it once
403
            let tokenizer_hf = card.tokenizer_hf().context("tokenizer_hf")?;
404

405
406
407
408
409
410
411
            // Create prefill chooser once if we're building pipelines
            // Both chat and completions will share the same prefill chooser instance
            let prefill_chooser = self
                .manager
                .register_prefill_router(card.name().to_string())
                .map(|rx| {
                    // Create prefill-specific config with track_active_blocks disabled
412
                    let mut prefill_config = self.router_config.kv_router_config;
413
414
415
416
417
                    prefill_config.router_track_active_blocks = false;

                    PrefillRouter::new(
                        rx,
                        self.manager.clone(),
418
                        self.router_config.router_mode,
419
420
                        card.kv_cache_block_size,
                        Some(prefill_config),
421
                        self.router_config.enforce_disagg,
422
423
424
                    )
                });

425
426
            // Get or create the worker monitor for this model
            // This allows dynamic threshold updates via the ModelManager
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
            // Create monitor if either threshold is configured
            let worker_monitor = if self.router_config.active_decode_blocks_threshold.is_some()
                || self.router_config.active_prefill_tokens_threshold.is_some()
            {
                // Default thresholds: active_decode_blocks=1.0 (disabled), active_prefill_tokens=1000000 (effectively disabled)
                let active_decode_blocks = self
                    .router_config
                    .active_decode_blocks_threshold
                    .unwrap_or(1.0);
                let active_prefill_tokens = self
                    .router_config
                    .active_prefill_tokens_threshold
                    .unwrap_or(1000000);
                Some(self.manager.get_or_create_worker_monitor(
                    card.name(),
                    client.clone(),
                    active_decode_blocks,
                    active_prefill_tokens,
                ))
            } else {
                None
            };
449

450
            // Add chat engine only if the model supports chat
451
            if card.model_type.supports_chat() {
452
453
454
455
456
457
458
459
460
461
462
463
                // Work in progress. This will allow creating  a chat_engine from Python.
                let chat_engine = if let Some(ref factory) = self.engine_factory {
                    factory(card.clone())
                        .await
                        .context("python engine_factory")?
                } else {
                    entrypoint::build_routed_pipeline::<
                        NvCreateChatCompletionRequest,
                        NvCreateChatCompletionStreamResponse,
                    >(
                        card,
                        &client,
464
                        self.manager.clone(),
465
466
467
468
469
470
                        self.router_config.router_mode,
                        worker_monitor.clone(),
                        kv_chooser.clone(),
                        tokenizer_hf.clone(),
                        prefill_chooser.clone(),
                        self.router_config.enforce_disagg,
471
                        self.metrics.clone(),
472
473
474
475
                    )
                    .await
                    .context("build_routed_pipeline")?
                };
476
                self.manager
477
                    .add_chat_completions_model(card.name(), checksum, chat_engine)
478
                    .context("add_chat_completions_model")?;
479
                tracing::info!("Chat completions is ready");
480
            }
481

482
            // Add completions engine only if the model supports completions
483
            if card.model_type.supports_completions() {
484
485
                let formatter = PromptFormatter::no_op();
                let PromptFormatter::OAI(formatter) = formatter;
486
487
488
489
                let preprocessor = OpenAIPreprocessor::new_with_parts(
                    card.clone(),
                    formatter,
                    tokenizer_hf.clone(),
490
491
                )
                .context("OpenAIPreprocessor::new_with_parts")?;
492
                let completions_engine = entrypoint::build_routed_pipeline_with_preprocessor::<
493
494
495
                    NvCreateCompletionRequest,
                    NvCreateCompletionResponse,
                >(
496
                    card,
497
                    &client,
498
                    self.manager.clone(),
499
                    self.router_config.router_mode,
500
                    worker_monitor,
501
                    kv_chooser,
502
                    preprocessor,
503
                    tokenizer_hf,
504
                    prefill_chooser,
505
                    self.router_config.enforce_disagg,
506
                    self.metrics.clone(),
507
                )
508
509
                .await
                .context("build_routed_pipeline_with_preprocessor")?;
510
                self.manager
511
                    .add_completions_model(card.name(), checksum, completions_engine)
512
                    .context("add_completions_model")?;
513
                tracing::info!("Completions is ready");
514
            }
515
516
517
518
519
520
        } else if card.model_input == ModelInput::Text && card.model_type.supports_embedding() {
            // Case: Text + Embeddings
            let push_router = PushRouter::<
                NvCreateEmbeddingRequest,
                Annotated<NvCreateEmbeddingResponse>,
            >::from_client_with_threshold(
521
                client, self.router_config.router_mode, None, None
522
523
524
525
            )
            .await?;
            let engine = Arc::new(push_router);
            self.manager
526
                .add_embeddings_model(card.name(), checksum, engine)?;
527
        } else if card.model_input == ModelInput::Text && card.model_type.supports_chat() {
528
            // Case 3: Text + Chat
529
530
531
532
533
534
535
            let push_router = PushRouter::<
                NvCreateChatCompletionRequest,
                Annotated<NvCreateChatCompletionStreamResponse>,
            >::from_client_with_threshold(
                client, self.router_config.router_mode, None, None
            )
            .await?;
536
537
            let engine = Arc::new(push_router);
            self.manager
538
                .add_chat_completions_model(card.name(), checksum, engine)?;
539
        } else if card.model_input == ModelInput::Text && card.model_type.supports_completions() {
540
541
542
543
544
            // Case 2: Text + Completions
            let push_router = PushRouter::<
                NvCreateCompletionRequest,
                Annotated<NvCreateCompletionResponse>,
            >::from_client_with_threshold(
545
                client, self.router_config.router_mode, None, None
546
547
548
549
            )
            .await?;
            let engine = Arc::new(push_router);
            self.manager
550
                .add_completions_model(card.name(), checksum, engine)?;
551
        } else if card.model_input == ModelInput::Tokens && card.model_type.supports_embedding() {
552
553
554
555
556
557
558
559
            // Case 4: Tokens + Embeddings

            // Create preprocessing pipeline similar to Backend
            let frontend = SegmentSource::<
                SingleIn<NvCreateEmbeddingRequest>,
                ManyOut<Annotated<NvCreateEmbeddingResponse>>,
            >::new();

560
            let preprocessor = OpenAIPreprocessor::new(card.clone())?.into_operator();
561
            let backend = Backend::from_mdc(card).into_operator();
562
563
564
565
566

            let router = PushRouter::<
                PreprocessedEmbeddingRequest,
                Annotated<EmbeddingsEngineOutput>,
            >::from_client_with_threshold(
567
                client, self.router_config.router_mode, None, None
568
569
570
            )
            .await?;

571
            // Note: Embeddings don't need KV routing complexity or load monitoring
572
573
574
575
576
577
578
579
580
581
582
583
            let service_backend = ServiceBackend::from_engine(Arc::new(router));

            // Link the pipeline: frontend -> preprocessor -> backend -> service_backend -> backend -> preprocessor -> frontend
            let embedding_engine = frontend
                .link(preprocessor.forward_edge())?
                .link(backend.forward_edge())?
                .link(service_backend)?
                .link(backend.backward_edge())?
                .link(preprocessor.backward_edge())?
                .link(frontend)?;

            self.manager
584
                .add_embeddings_model(card.name(), checksum, embedding_engine)?;
585
        } else if card.model_input == ModelInput::Tensor && card.model_type.supports_tensor() {
586
            // Case 5: Tensor + Tensor (non-LLM)
587
            // No KV cache concepts - not an LLM model
588
589
590
591
            let push_router = PushRouter::<
                NvCreateTensorRequest,
                Annotated<NvCreateTensorResponse>,
            >::from_client_with_threshold(
592
                client, self.router_config.router_mode, None, None
593
594
595
            )
            .await?;
            let engine = Arc::new(push_router);
596
597
            self.manager
                .add_tensor_model(card.name(), checksum, engine)?;
598
599
600
601
602
603
604
605
606
607
608
609
        } else if card.model_type.supports_prefill() {
            // Case 6: Prefill
            // Guardrail: Verify model_input is Tokens
            if card.model_input != ModelInput::Tokens {
                anyhow::bail!(
                    "Prefill models must use ModelInput::Tokens, got {}",
                    card.model_input.as_str()
                );
            }

            tracing::info!(
                model_name = card.name(),
610
611
612
613
614
615
616
617
618
619
620
621
622
623
624
625
626
627
628
629
                "Prefill model detected, registering and activating prefill router"
            );

            // Register prefill model for tracking (no engine needed, just lifecycle)
            self.manager
                .add_prefill_model(card.name(), checksum)
                .context("add_prefill_model")?;

            // Activate the prefill router with the endpoint for this prefill model
            let Ok(()) = self.manager.activate_prefill_router(card.name(), endpoint) else {
                tracing::warn!(
                    model_name = card.name(),
                    "Failed to activate prefill router - prefill model may already be activated"
                );
                return Ok(());
            };

            tracing::info!(
                model_name = card.name(),
                "Prefill model registered and router activated successfully"
630
            );
631
632
633
634
        } else {
            // Reject unsupported combinations
            anyhow::bail!(
                "Unsupported model configuration: {} with {} input. Supported combinations: \
635
                Tokens+(Chat|Completions|Prefill), Text+Chat, Text+Completions, Tokens+Embeddings, Tensor+TensorBased",
636
637
                card.model_type,
                card.model_input.as_str()
638
            );
639
        }
Ryan Olson's avatar
Ryan Olson committed
640

641
642
        Ok(())
    }
643

644
    /// All the registered ModelDeploymentCard with the EndpointId they are attached to, one per instance
645
    async fn all_cards(&self) -> anyhow::Result<Vec<(EndpointId, ModelDeploymentCard)>> {
646
647
        let discovery = self.drt.discovery();
        let instances = discovery.list(DiscoveryQuery::AllModels).await?;
648

649
650
651
        let mut results = Vec::with_capacity(instances.len());
        for instance in instances {
            match instance.deserialize_model::<ModelDeploymentCard>() {
652
                Ok(card) => {
653
654
655
656
657
658
659
660
661
662
663
664
665
666
667
668
                    // Extract EndpointId from the instance
                    let endpoint_id = match &instance {
                        dynamo_runtime::discovery::DiscoveryInstance::Model {
                            namespace,
                            component,
                            endpoint,
                            ..
                        } => EndpointId {
                            namespace: namespace.clone(),
                            component: component.clone(),
                            name: endpoint.clone(),
                        },
                        _ => {
                            tracing::error!(
                                "Unexpected discovery instance type (expected ModelCard)"
                            );
669
670
671
                            continue;
                        }
                    };
672
                    results.push((endpoint_id, card));
673
                }
674
                Err(err) => {
675
                    tracing::error!(%err, "Failed to deserialize model card");
676
677
                    continue;
                }
678
            }
679
        }
680
        Ok(results)
681
682
    }

683
    pub async fn cards_for_model(
684
685
686
687
        &self,
        model_name: &str,
        target_namespace: Option<&str>,
        is_global_namespace: bool,
688
689
690
691
    ) -> anyhow::Result<Vec<ModelDeploymentCard>> {
        let mut all = self.all_cards().await?;
        all.retain(|(endpoint_id, card)| {
            let matches_name = card.name() == model_name;
692
693
694
            let matches_namespace = match (is_global_namespace, target_namespace) {
                (true, _) => true,
                (false, None) => true,
695
                (false, Some(target_ns)) => endpoint_id.namespace == target_ns,
696
697
698
            };
            matches_name && matches_namespace
        });
699
700
701
        Ok(all.into_iter().map(|(_eid, card)| card).collect())
    }
}