watcher.rs 31.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 dashmap::DashSet;
10
use futures::StreamExt;
Ryan Olson's avatar
Ryan Olson committed
11

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

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

48
use super::ModelManager;
49
use crate::namespace::is_global_namespace;
50

51
#[derive(Debug, Clone)]
52
pub enum ModelUpdate {
53
54
    Added(ModelDeploymentCard),
    Removed(ModelDeploymentCard),
55
56
}

57
pub struct ModelWatcher {
58
    manager: Arc<ModelManager>,
59
    drt: DistributedRuntime,
60
    router_config: RouterConfig,
61
    migration_limit: u32,
62
    notify_on_model: Notify,
63
    model_update_tx: Option<Sender<ModelUpdate>>,
64
    engine_factory: Option<EngineFactoryCallback>,
65
    metrics: Arc<Metrics>,
66
    registering_models: DashSet<String>,
Ryan Olson's avatar
Ryan Olson committed
67
68
}

69
70
71
72
const ALL_MODEL_TYPES: &[ModelType] = &[
    ModelType::Chat,
    ModelType::Completions,
    ModelType::Embedding,
73
    ModelType::TensorBased,
74
    ModelType::Images,
75
    ModelType::Prefill,
76
];
77

78
impl ModelWatcher {
79
    pub fn new(
80
        runtime: DistributedRuntime,
81
        model_manager: Arc<ModelManager>,
82
        router_config: RouterConfig,
83
        migration_limit: u32,
84
        engine_factory: Option<EngineFactoryCallback>,
85
        metrics: Arc<Metrics>,
86
87
    ) -> ModelWatcher {
        Self {
88
            manager: model_manager,
89
            drt: runtime,
90
            router_config,
91
            migration_limit,
92
            notify_on_model: Notify::new(),
93
            model_update_tx: None,
94
            engine_factory,
95
            metrics,
96
            registering_models: DashSet::new(),
97
        }
98
    }
Ryan Olson's avatar
Ryan Olson committed
99

100
101
102
103
    pub fn set_notify_on_model_update(&mut self, tx: Sender<ModelUpdate>) {
        self.model_update_tx = Some(tx);
    }

104
105
106
107
108
109
110
111
112
113
114
    /// 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
        }
    }

115
    /// Common watch logic with optional namespace filtering
116
117
118
119
120
    pub async fn watch(
        &self,
        mut discovery_stream: DiscoveryStream,
        target_namespace: Option<&str>,
    ) {
121
        let global_namespace = target_namespace.is_none_or(is_global_namespace);
122

123
124
125
126
127
128
129
130
131
        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;
                }
            };

132
            match event {
133
                DiscoveryEvent::Added(instance) => {
134
135
                    // Extract ModelCardInstanceId and card from the discovery instance
                    let (mcid, mut card) = match &instance {
136
137
138
139
140
                        DiscoveryInstance::Model {
                            namespace,
                            component,
                            endpoint,
                            instance_id,
141
                            model_suffix,
142
143
                            ..
                        } => {
144
                            let mcid = ModelCardInstanceId {
145
146
                                namespace: namespace.clone(),
                                component: component.clone(),
147
148
149
                                endpoint: endpoint.clone(),
                                instance_id: *instance_id,
                                model_suffix: model_suffix.clone(),
150
151
152
                            };

                            match instance.deserialize_model::<ModelDeploymentCard>() {
153
                                Ok(card) => (mcid, card),
154
155
156
157
158
159
160
161
162
163
                                Err(err) => {
                                    tracing::error!(%err, instance_id, "Failed to deserialize model card");
                                    continue;
                                }
                            }
                        }
                        _ => {
                            tracing::error!(
                                "Unexpected discovery instance type (expected ModelCard)"
                            );
164
165
166
                            continue;
                        }
                    };
167
168
169
170

                    // Filter by namespace if target_namespace is specified
                    if !global_namespace
                        && let Some(target_ns) = target_namespace
171
                        && mcid.namespace != target_ns
172
173
                    {
                        tracing::debug!(
174
                            model_namespace = mcid.namespace,
175
176
177
178
179
180
                            target_namespace = target_ns,
                            "Skipping model from different namespace"
                        );
                        continue;
                    }

181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
                    // 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;
                    }
202

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

234
                    match self
235
                        .handle_delete(model_card_instance_id, target_namespace, global_namespace)
236
237
238
239
240
241
242
243
244
245
246
                        .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");
                        }
247
                    }
248
                }
249
            }
Ryan Olson's avatar
Ryan Olson committed
250
251
252
        }
    }

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

283
        // Ignore the errors because model could be either type
284
285
286
        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);
287
        let tensor_model_remove_err = self.manager.remove_tensor_model(&model_name);
288
        let images_model_remove_err = self.manager.remove_images_model(&model_name);
289
        let prefill_model_remove_err = self.manager.remove_prefill_model(&model_name);
290
291
292
293

        let mut chat_model_removed = false;
        let mut completions_model_removed = false;
        let mut embeddings_model_removed = false;
294
        let mut tensor_model_removed = false;
295
        let mut images_model_removed = false;
296
        let mut prefill_model_removed = false;
297
298
299
300
301
302
303
304
305
306
307

        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;
        }
308
309
310
        if tensor_model_remove_err.is_ok() && self.manager.list_tensor_models().is_empty() {
            tensor_model_removed = true;
        }
311
312
313
        if images_model_remove_err.is_ok() && self.manager.list_images_models().is_empty() {
            images_model_removed = true;
        }
314
315
316
        if prefill_model_remove_err.is_ok() && self.manager.list_prefill_models().is_empty() {
            prefill_model_removed = true;
        }
317

318
319
320
321
        if !chat_model_removed
            && !completions_model_removed
            && !embeddings_model_removed
            && !tensor_model_removed
322
            && !images_model_removed
323
            && !prefill_model_removed
324
        {
325
            tracing::debug!(
326
                "No updates to send for model {}: chat_model_removed: {}, completions_model_removed: {}, embeddings_model_removed: {}, tensor_model_removed: {}, images_model_removed: {}, prefill_model_removed: {}",
327
328
329
                model_name,
                chat_model_removed,
                completions_model_removed,
330
                embeddings_model_removed,
331
                tensor_model_removed,
332
                images_model_removed,
333
                prefill_model_removed
334
335
336
            );
        } else {
            for model_type in ALL_MODEL_TYPES {
337
                if ((chat_model_removed && *model_type == ModelType::Chat)
338
                    || (completions_model_removed && *model_type == ModelType::Completions)
339
                    || (embeddings_model_removed && *model_type == ModelType::Embedding)
340
                    || (tensor_model_removed && *model_type == ModelType::TensorBased)
341
                    || (images_model_removed && *model_type == ModelType::Images)
342
                    || (prefill_model_removed && *model_type == ModelType::Prefill))
343
                    && let Some(tx) = &self.model_update_tx
344
                {
345
                    tx.send(ModelUpdate::Removed(card.clone())).await.ok();
346
347
348
                }
            }
        }
349

350
        Ok(Some(model_name))
351
    }
Ryan Olson's avatar
Ryan Olson committed
352

353
    // Handles a PUT event from store, this usually means adding a new model to the list of served
354
    // models.
355
356
    async fn handle_put(
        &self,
357
        mcid: &ModelCardInstanceId,
358
359
        card: &mut ModelDeploymentCard,
    ) -> anyhow::Result<()> {
360
361
362
        // Check if model is already registered before downloading config.
        // This prevents duplicate HuggingFace API calls when multiple workers register
        // the same model.
363
364
365
366
367
368
369
370
371
        // 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 {
372
373
            self.manager
                .save_model_card(&mcid.to_path(), card.clone())?;
374
375
            tracing::debug!(
                model_name = card.name(),
376
                namespace = mcid.namespace,
377
                model_type = %card.model_type,
378
379
380
381
382
383
384
385
386
387
388
389
390
391
                "Model already registered, skipping config download"
            );
            return Ok(());
        }

        // Use registering_models set to prevent concurrent registrations.
        let model_key = card.name().to_string();
        if !self.registering_models.insert(model_key.clone()) {
            self.manager
                .save_model_card(&mcid.to_path(), card.clone())?;
            tracing::debug!(
                model_name = card.name(),
                namespace = mcid.namespace,
                "Model registration in progress by another worker, skipping"
392
393
394
395
            );
            return Ok(());
        }

396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
        // We acquired the registration lock. Use a helper to ensure cleanup on all exit paths.
        let result = self.do_model_registration(mcid, card).await;

        // Always remove from registering set, whether success or failure
        self.registering_models.remove(&model_key);

        result
    }

    /// Inner function that performs the actual model registration.
    /// Called by handle_put after acquiring the registration lock.
    async fn do_model_registration(
        &self,
        mcid: &ModelCardInstanceId,
        card: &mut ModelDeploymentCard,
    ) -> anyhow::Result<()> {
        card.download_config().await?;

        let component = self
            .drt
            .namespace(&mcid.namespace)?
            .component(&mcid.component)?;
        let endpoint = component.endpoint(&mcid.endpoint);
        let client = endpoint.client().await?;
        tracing::debug!(model_name = card.name(), "adding model");
        self.manager
            .save_model_card(&mcid.to_path(), card.clone())?;

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

428
        let checksum = card.mdcsum();
429
430
431

        if card.model_input == ModelInput::Tokens
            && (card.model_type.supports_chat() || card.model_type.supports_completions())
432
433
434
435
436
        {
            // 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.

437
            let endpoint = component.endpoint(&mcid.endpoint);
438
            let kv_chooser = if self.router_config.router_mode == RouterMode::KV {
439
440
                Some(
                    self.manager
441
442
443
444
                        .kv_chooser_for(
                            &endpoint,
                            card.kv_cache_block_size,
                            Some(self.router_config.kv_router_config),
445
                            WORKER_TYPE_DECODE, // This is the decode router
446
                        )
447
448
449
450
451
                        .await?,
                )
            } else {
                None
            };
452

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

456
457
            // Create prefill chooser once if we're building pipelines
            // Both chat and completions will share the same prefill chooser instance
458
            let model_name = card.name().to_string();
459
460
            let prefill_chooser = self
                .manager
461
                .register_prefill_router(model_name.clone())
462
463
                .map(|rx| {
                    // Create prefill-specific config with track_active_blocks disabled
464
                    let mut prefill_config = self.router_config.kv_router_config;
465
466
467
468
469
                    prefill_config.router_track_active_blocks = false;

                    PrefillRouter::new(
                        rx,
                        self.manager.clone(),
470
                        self.router_config.router_mode,
471
472
                        card.kv_cache_block_size,
                        Some(prefill_config),
473
                        self.router_config.enforce_disagg,
474
                        model_name.clone(), // Pass model name for worker monitor lookup
475
476
477
                    )
                });

478
479
480
481
482
483
484
485
486
            // Get or create the worker monitor for this model.
            // Always create the monitor for Prometheus metrics (active_decode_blocks, active_prefill_tokens,
            // worker TTFT/ITL cleanup). The thresholds control busy detection behavior only.
            // LoadThresholdConfig allows dynamic threshold updates via the ModelManager.
            let worker_monitor = Some(self.manager.get_or_create_worker_monitor(
                card.name(),
                client.clone(),
                self.router_config.load_threshold_config.clone(),
            ));
487

488
            // Add chat engine only if the model supports chat
489
            if card.model_type.supports_chat() {
490
491
492
493
494
495
496
497
498
499
500
501
                // 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,
502
                        self.manager.clone(),
503
504
505
506
507
508
                        self.router_config.router_mode,
                        worker_monitor.clone(),
                        kv_chooser.clone(),
                        tokenizer_hf.clone(),
                        prefill_chooser.clone(),
                        self.router_config.enforce_disagg,
509
                        self.migration_limit,
510
                        self.metrics.clone(),
511
512
513
514
                    )
                    .await
                    .context("build_routed_pipeline")?
                };
515
                self.manager
516
                    .add_chat_completions_model(card.name(), checksum, chat_engine)
517
                    .context("add_chat_completions_model")?;
518
                tracing::info!("Chat completions is ready");
519
            }
520

521
            // Add completions engine only if the model supports completions
522
            if card.model_type.supports_completions() {
523
524
                let formatter = PromptFormatter::no_op();
                let PromptFormatter::OAI(formatter) = formatter;
525
526
527
528
                let preprocessor = OpenAIPreprocessor::new_with_parts(
                    card.clone(),
                    formatter,
                    tokenizer_hf.clone(),
529
530
                )
                .context("OpenAIPreprocessor::new_with_parts")?;
531
                let completions_engine = entrypoint::build_routed_pipeline_with_preprocessor::<
532
533
534
                    NvCreateCompletionRequest,
                    NvCreateCompletionResponse,
                >(
535
                    card,
536
                    &client,
537
                    self.manager.clone(),
538
                    self.router_config.router_mode,
539
                    worker_monitor,
540
                    kv_chooser,
541
                    preprocessor,
542
                    tokenizer_hf,
543
                    prefill_chooser,
544
                    self.router_config.enforce_disagg,
545
                    self.migration_limit,
546
                    self.metrics.clone(),
547
                )
548
549
                .await
                .context("build_routed_pipeline_with_preprocessor")?;
550
                self.manager
551
                    .add_completions_model(card.name(), checksum, completions_engine)
552
                    .context("add_completions_model")?;
553
                tracing::info!("Completions is ready");
554
            }
555
556
557
558
559
560
        } 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(
561
                client, self.router_config.router_mode, None, None
562
563
564
565
            )
            .await?;
            let engine = Arc::new(push_router);
            self.manager
566
                .add_embeddings_model(card.name(), checksum, engine)?;
567
        } else if card.model_input == ModelInput::Text && card.model_type.supports_chat() {
568
            // Case 3: Text + Chat
569
570
571
572
573
574
575
            let push_router = PushRouter::<
                NvCreateChatCompletionRequest,
                Annotated<NvCreateChatCompletionStreamResponse>,
            >::from_client_with_threshold(
                client, self.router_config.router_mode, None, None
            )
            .await?;
576
577
            let engine = Arc::new(push_router);
            self.manager
578
                .add_chat_completions_model(card.name(), checksum, engine)?;
579
        } else if card.model_input == ModelInput::Text && card.model_type.supports_completions() {
580
581
582
583
584
            // Case 2: Text + Completions
            let push_router = PushRouter::<
                NvCreateCompletionRequest,
                Annotated<NvCreateCompletionResponse>,
            >::from_client_with_threshold(
585
                client, self.router_config.router_mode, None, None
586
587
588
589
            )
            .await?;
            let engine = Arc::new(push_router);
            self.manager
590
                .add_completions_model(card.name(), checksum, engine)?;
591
        } else if card.model_input == ModelInput::Tokens && card.model_type.supports_embedding() {
592
593
594
595
596
597
598
599
            // Case 4: Tokens + Embeddings

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

600
            let preprocessor = OpenAIPreprocessor::new(card.clone())?.into_operator();
601
            let backend = Backend::from_mdc(card).into_operator();
602
603
604
605
606

            let router = PushRouter::<
                PreprocessedEmbeddingRequest,
                Annotated<EmbeddingsEngineOutput>,
            >::from_client_with_threshold(
607
                client, self.router_config.router_mode, None, None
608
609
610
            )
            .await?;

611
            // Note: Embeddings don't need KV routing complexity or load monitoring
612
613
614
615
616
617
618
619
620
621
622
623
            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
624
                .add_embeddings_model(card.name(), checksum, embedding_engine)?;
625
        } else if card.model_input == ModelInput::Tensor && card.model_type.supports_tensor() {
626
            // Case 6: Tensor + TensorBased (non-LLM)
627
            // No KV cache concepts - not an LLM model
628
629
630
631
            let push_router = PushRouter::<
                NvCreateTensorRequest,
                Annotated<NvCreateTensorResponse>,
            >::from_client_with_threshold(
632
                client, self.router_config.router_mode, None, None
633
634
635
            )
            .await?;
            let engine = Arc::new(push_router);
636
637
            self.manager
                .add_tensor_model(card.name(), checksum, engine)?;
638
        } else if card.model_input == ModelInput::Text && card.model_type.supports_images() {
639
640
641
642
            // Case: Text + Images (e.g. vLLM-Omni, diffusion models)
            // Takes text prompts as input, generates images. Images models also support
            // chat completions (see model_type.rs as_endpoint_types).
            let images_router = PushRouter::<
643
644
645
                NvCreateImageRequest,
                Annotated<NvImagesResponse>,
            >::from_client_with_threshold(
646
                client.clone(), self.router_config.router_mode, None, None
647
648
649
            )
            .await?;
            self.manager
650
651
652
653
654
655
656
657
658
659
660
661
662
663
                .add_images_model(card.name(), checksum, Arc::new(images_router))?;

            let chat_router = PushRouter::<
                NvCreateChatCompletionRequest,
                Annotated<NvCreateChatCompletionStreamResponse>,
            >::from_client_with_threshold(
                client, self.router_config.router_mode, None, None
            )
            .await?;
            self.manager.add_chat_completions_model(
                card.name(),
                checksum,
                Arc::new(chat_router),
            )?;
664
665
666
667
668
669
670
671
672
673
674
675
        } 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(),
676
677
678
679
680
681
682
683
684
685
686
687
688
689
690
691
692
693
694
695
                "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"
696
            );
697
698
699
700
        } else {
            // Reject unsupported combinations
            anyhow::bail!(
                "Unsupported model configuration: {} with {} input. Supported combinations: \
701
                Tokens+(Chat|Completions|Prefill), Text+(Chat|Completions|Images), Tokens+Embeddings, Tensor+TensorBased",
702
703
                card.model_type,
                card.model_input.as_str()
704
            );
705
        }
Ryan Olson's avatar
Ryan Olson committed
706

707
708
        Ok(())
    }
709

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

715
716
717
        let mut results = Vec::with_capacity(instances.len());
        for instance in instances {
            match instance.deserialize_model::<ModelDeploymentCard>() {
718
                Ok(card) => {
719
720
721
722
723
724
725
726
727
728
729
730
731
732
733
734
                    // 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)"
                            );
735
736
737
                            continue;
                        }
                    };
738
                    results.push((endpoint_id, card));
739
                }
740
                Err(err) => {
741
                    tracing::error!(%err, "Failed to deserialize model card");
742
743
                    continue;
                }
744
            }
745
        }
746
        Ok(results)
747
748
    }

749
    pub async fn cards_for_model(
750
751
752
753
        &self,
        model_name: &str,
        target_namespace: Option<&str>,
        is_global_namespace: bool,
754
755
756
757
    ) -> anyhow::Result<Vec<ModelDeploymentCard>> {
        let mut all = self.all_cards().await?;
        all.retain(|(endpoint_id, card)| {
            let matches_name = card.name() == model_name;
758
759
760
            let matches_namespace = match (is_global_namespace, target_namespace) {
                (true, _) => true,
                (false, None) => true,
761
                (false, Some(target_ns)) => endpoint_id.namespace == target_ns,
762
763
764
            };
            matches_name && matches_namespace
        });
765
766
767
        Ok(all.into_iter().map(|(_eid, card)| card).collect())
    }
}