watcher.rs 36.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::{KvWorkerMonitor, WORKER_TYPE_DECODE, WorkerSet},
28
    entrypoint::{self, ChatEngineFactoryCallback, 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
            videos::{NvCreateVideoRequest, NvVideosResponse},
44
        },
45
        tensor::{NvCreateTensorRequest, NvCreateTensorResponse},
46
47
    },
};
48

49
use super::ModelManager;
50
51
52
53
54
55
56
57
58
59
60
use crate::namespace::NamespaceFilter;

/// Constructs the WorkerSet storage key. Prefill and decode workers in the same
/// namespace get different keys so they don't block each other's registration.
fn worker_set_key(namespace: &str, model_type: ModelType) -> String {
    if model_type.supports_prefill() {
        format!("{}:prefill", namespace)
    } else {
        namespace.to_string()
    }
}
61

62
#[derive(Debug, Clone)]
63
pub enum ModelUpdate {
64
65
    Added(ModelDeploymentCard),
    Removed(ModelDeploymentCard),
66
67
}

68
pub struct ModelWatcher {
69
    manager: Arc<ModelManager>,
70
    drt: DistributedRuntime,
71
    router_config: RouterConfig,
72
    migration_limit: u32,
73
    notify_on_model: Notify,
74
    model_update_tx: Option<Sender<ModelUpdate>>,
75
    chat_engine_factory: Option<ChatEngineFactoryCallback>,
76
    metrics: Arc<Metrics>,
77
78
    /// Guards against concurrent pipeline construction for the same (model, namespace).
    registering_worker_sets: DashSet<String>,
Ryan Olson's avatar
Ryan Olson committed
79
80
}

81
82
83
84
const ALL_MODEL_TYPES: &[ModelType] = &[
    ModelType::Chat,
    ModelType::Completions,
    ModelType::Embedding,
85
    ModelType::Images,
86
    ModelType::Audios,
87
88
    ModelType::Videos,
    ModelType::TensorBased,
89
    ModelType::Prefill,
90
];
91

92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
/// Returns true if no models in the manager support the given model type.
fn is_model_type_list_empty(manager: &ModelManager, model_type: ModelType) -> bool {
    if model_type == ModelType::Chat {
        manager.list_chat_completions_models().is_empty()
    } else if model_type == ModelType::Completions {
        manager.list_completions_models().is_empty()
    } else if model_type == ModelType::Embedding {
        manager.list_embeddings_models().is_empty()
    } else if model_type == ModelType::Images {
        manager.list_images_models().is_empty()
    } else if model_type == ModelType::Videos {
        manager.list_videos_models().is_empty()
    } else if model_type == ModelType::TensorBased {
        manager.list_tensor_models().is_empty()
    } else if model_type == ModelType::Prefill {
        manager.list_prefill_models().is_empty()
    } else {
        true
    }
}

113
impl ModelWatcher {
114
    pub fn new(
115
        runtime: DistributedRuntime,
116
        model_manager: Arc<ModelManager>,
117
        router_config: RouterConfig,
118
        migration_limit: u32,
119
        chat_engine_factory: Option<ChatEngineFactoryCallback>,
120
        metrics: Arc<Metrics>,
121
122
    ) -> ModelWatcher {
        Self {
123
            manager: model_manager,
124
            drt: runtime,
125
            router_config,
126
            migration_limit,
127
            notify_on_model: Notify::new(),
128
            model_update_tx: None,
129
            chat_engine_factory,
130
            metrics,
131
            registering_worker_sets: DashSet::new(),
132
        }
133
    }
Ryan Olson's avatar
Ryan Olson committed
134

135
136
137
138
    pub fn set_notify_on_model_update(&mut self, tx: Sender<ModelUpdate>) {
        self.model_update_tx = Some(tx);
    }

139
140
141
142
143
144
145
146
147
148
149
    /// 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
        }
    }

150
    /// Common watch logic with optional namespace filtering
151
152
153
    pub async fn watch(
        &self,
        mut discovery_stream: DiscoveryStream,
154
        namespace_filter: NamespaceFilter,
155
156
157
158
159
160
161
162
163
164
    ) {
        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;
                }
            };

165
            match event {
166
                DiscoveryEvent::Added(instance) => {
167
168
                    // Extract ModelCardInstanceId and card from the discovery instance
                    let (mcid, mut card) = match &instance {
169
170
171
172
173
                        DiscoveryInstance::Model {
                            namespace,
                            component,
                            endpoint,
                            instance_id,
174
                            model_suffix,
175
176
                            ..
                        } => {
177
                            let mcid = ModelCardInstanceId {
178
179
                                namespace: namespace.clone(),
                                component: component.clone(),
180
181
182
                                endpoint: endpoint.clone(),
                                instance_id: *instance_id,
                                model_suffix: model_suffix.clone(),
183
184
185
                            };

                            match instance.deserialize_model::<ModelDeploymentCard>() {
186
                                Ok(card) => (mcid, card),
187
188
189
190
191
192
193
194
195
196
                                Err(err) => {
                                    tracing::error!(%err, instance_id, "Failed to deserialize model card");
                                    continue;
                                }
                            }
                        }
                        _ => {
                            tracing::error!(
                                "Unexpected discovery instance type (expected ModelCard)"
                            );
197
198
199
                            continue;
                        }
                    };
200

201
202
                    // Filter by namespace using the configured filter
                    if !namespace_filter.matches(&mcid.namespace) {
203
                        tracing::debug!(
204
                            model_namespace = mcid.namespace,
205
206
                            namespace_filter = ?namespace_filter,
                            "Skipping model due to namespace filter"
207
208
209
210
                        );
                        continue;
                    }

211
212
213
214
                    // If we already have a WorkerSet for this model and the checksums
                    // don't match, reject the new worker. All WorkerSets of a model
                    // must share the same checksum.
                    let can_add = self.manager.is_valid_checksum(card.name(), card.mdcsum());
215
216
217
                    if can_add.is_some_and(|is_valid| !is_valid) {
                        tracing::error!(
                            model_name = card.name(),
218
219
220
221
                            namespace = mcid.namespace,
                            "Checksum for new worker does not match model's canonical checksum. \
                             All WorkerSets must share the same checksum. \
                             Drain all old workers before deploying a new version."
222
223
224
225
226
227
228
229
230
231
232
                        );

                        // 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;
                    }
233

234
                    match self.handle_put(&mcid, &mut card).await {
235
                        Ok(()) => {
236
                            tracing::info!(
237
                                model_name = card.name(),
238
                                namespace = mcid.namespace,
239
240
                                "added model"
                            );
241
                            self.notify_on_model.notify_waiters();
242
                        }
243
244
                        Err(err) => {
                            tracing::error!(
245
                                model_name = card.name(),
246
                                namespace = mcid.namespace,
247
                                error = format!("{err:#}"),
248
                                "Error adding model from discovery",
249
                            );
250
251
252
                        }
                    }
                }
253
254
255
256
                DiscoveryEvent::Removed(id) => {
                    // Extract ModelCardInstanceId from the removal event
                    let model_card_instance_id = match &id {
                        DiscoveryInstanceId::Model(mcid) => mcid,
257
                        DiscoveryInstanceId::Endpoint(_) | DiscoveryInstanceId::EventChannel(_) => {
258
259
260
261
262
263
                            tracing::error!(
                                "Unexpected discovery instance type in removal (expected Model)"
                            );
                            continue;
                        }
                    };
264

265
                    match self
266
                        .handle_delete(model_card_instance_id, &namespace_filter)
267
268
269
270
271
272
273
274
275
276
277
                        .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");
                        }
278
                    }
279
                }
280
            }
Ryan Olson's avatar
Ryan Olson committed
281
282
283
        }
    }

284
285
    /// Handle a worker removal. Cleans up per-namespace WorkerSets and the Model itself
    /// when no instances remain. Returns the model name if the entire Model was removed.
286
287
    async fn handle_delete(
        &self,
288
        mcid: &ModelCardInstanceId,
289
        namespace_filter: &NamespaceFilter,
290
    ) -> anyhow::Result<Option<String>> {
291
292
        let key = mcid.to_path();
        let card = match self.manager.remove_model_card(&key) {
293
            Some(card) => card,
294
            None => {
295
                anyhow::bail!("Missing ModelDeploymentCard for {}", key);
296
297
            }
        };
298
        let model_name = card.name().to_string();
299
300
301
302
303
        let worker_namespace = &mcid.namespace;
        let worker_component = &mcid.component;
        let ws_key = worker_set_key(&mcid.namespace, card.model_type);

        // Query discovery for all remaining instances of this model
304
        let active_instances = self
305
            .cards_for_model_with_endpoints(&model_name, namespace_filter)
306
307
            .await
            .with_context(|| model_name.clone())?;
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331

        // Check if instances of the SAME component remain in this namespace.
        // In disaggregated deployments, prefill and decode are different components
        // in the same namespace, so we must check at the component level to avoid
        // removing one type's WorkerSet while the other still has workers.
        let component_has_instances = active_instances.iter().any(|(eid, _)| {
            eid.namespace == *worker_namespace && eid.component == *worker_component
        });

        if !component_has_instances {
            // No more workers of this component in this namespace — remove its WorkerSet
            if let Some(_removed_ws) = self.manager.remove_worker_set(&model_name, &ws_key) {
                // remove_prefill_activator uses deployment namespace (not ws_key)
                self.manager
                    .remove_prefill_activator(&model_name, worker_namespace);
                tracing::info!(
                    model_name,
                    namespace = %worker_namespace,
                    "Removed WorkerSet (no remaining instances in namespace)"
                );
            }
        }

        // Check if the Model still has instances in any namespace
332
        if !active_instances.is_empty() {
333
334
335
            tracing::debug!(
                model_name,
                active_instance_count = active_instances.len(),
336
                "Model has other active instances in other namespaces"
337
            );
338
339
            return Ok(None);
        }
340

341
342
        // No instances remain anywhere — remove the entire Model
        let _ = self.manager.remove_model(&model_name);
343

344
        if let Some(tx) = &self.model_update_tx {
345
            for model_type in ALL_MODEL_TYPES {
346
347
                if card.model_type.intersects(*model_type)
                    && is_model_type_list_empty(&self.manager, *model_type)
348
                {
349
                    tx.send(ModelUpdate::Removed(card.clone())).await.ok();
350
351
352
                }
            }
        }
353

354
        Ok(Some(model_name))
355
    }
Ryan Olson's avatar
Ryan Olson committed
356

357
    // Handles a PUT event from store, this usually means adding a new model to the list of served
358
    // models.
359
360
    async fn handle_put(
        &self,
361
        mcid: &ModelCardInstanceId,
362
363
        card: &mut ModelDeploymentCard,
    ) -> anyhow::Result<()> {
364
365
366
367
368
        // Check if this specific (model, namespace, type) WorkerSet already exists.
        // If so, this is just another worker joining an existing set — no pipeline build needed.
        let model_name = card.name().to_string();
        let namespace = mcid.namespace.clone();
        let ws_key = worker_set_key(&namespace, card.model_type);
369

370
371
372
        if let Some(model) = self.manager.get_model(&model_name)
            && model.has_worker_set(&ws_key)
        {
373
374
            self.manager
                .save_model_card(&mcid.to_path(), card.clone())?;
375
376
            tracing::debug!(
                model_name = card.name(),
377
378
                namespace = namespace,
                "Worker joined existing WorkerSet, skipping pipeline build"
379
380
381
382
            );
            return Ok(());
        }

383
384
385
386
387
388
        // Guard against concurrent pipeline construction for the same (model, namespace, type)
        let registration_key = ModelManager::model_namespace_key(&model_name, &ws_key);
        if !self
            .registering_worker_sets
            .insert(registration_key.clone())
        {
389
390
391
392
            self.manager
                .save_model_card(&mcid.to_path(), card.clone())?;
            tracing::debug!(
                model_name = card.name(),
393
394
                namespace = namespace,
                "WorkerSet registration in progress, skipping"
395
396
397
398
            );
            return Ok(());
        }

399
        let result = self.do_worker_set_registration(mcid, card).await;
400

401
402
        // Always remove from registering set
        self.registering_worker_sets.remove(&registration_key);
403
404
405
406

        result
    }

407
408
409
    /// Build a complete WorkerSet with all engines for this (model, namespace)
    /// and add it to the Model.
    async fn do_worker_set_registration(
410
411
412
413
414
415
416
417
418
419
420
421
        &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?;
422
423
424
425
426
427
        let instance_watcher = client.instance_avail_watcher();
        tracing::debug!(
            model_name = card.name(),
            namespace = mcid.namespace,
            "building worker set pipeline"
        );
428
429
430
        self.manager
            .save_model_card(&mcid.to_path(), card.clone())?;

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

435
        let checksum = card.mdcsum();
436
437
438
439
440
441
        let namespace = mcid.namespace.clone();
        let ws_key = worker_set_key(&namespace, card.model_type);

        // Build the WorkerSet with all applicable engines
        let mut worker_set = WorkerSet::new(namespace.clone(), checksum.to_string(), card.clone());
        worker_set.set_instance_watcher(instance_watcher);
442
443
444

        if card.model_input == ModelInput::Tokens
            && (card.model_type.supports_chat() || card.model_type.supports_completions())
445
446
447
448
449
        {
            // 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.

450
            let endpoint = component.endpoint(&mcid.endpoint);
451
452
453
454
455
456
457
458
459
            // Create the KV router whenever any local routed pipeline will be built.
            // The chat factory builds its own router, but completions currently always
            // uses the local routed pipeline and therefore still needs a chooser.
            let needs_local_chat_pipeline =
                card.model_type.supports_chat() && self.chat_engine_factory.is_none();
            let needs_local_completions_pipeline = card.model_type.supports_completions();
            let kv_chooser = if self.router_config.router_mode == RouterMode::KV
                && (needs_local_chat_pipeline || needs_local_completions_pipeline)
            {
460
461
                Some(
                    self.manager
462
463
464
465
                        .kv_chooser_for(
                            &endpoint,
                            card.kv_cache_block_size,
                            Some(self.router_config.kv_router_config),
466
                            WORKER_TYPE_DECODE, // This is the decode router
467
                        )
468
469
470
471
472
                        .await?,
                )
            } else {
                None
            };
473

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

477
478
            // Create prefill chooser once if we're building pipelines
            // Both chat and completions will share the same prefill chooser instance
479
            let model_name = card.name().to_string();
480
481
            let prefill_chooser = self
                .manager
482
                .register_prefill_router(&model_name, &namespace)
483
484
                .map(|rx| {
                    // Create prefill-specific config with track_active_blocks disabled
485
                    let mut prefill_config = self.router_config.kv_router_config;
486
487
488
489
490
                    prefill_config.router_track_active_blocks = false;

                    PrefillRouter::new(
                        rx,
                        self.manager.clone(),
491
                        self.router_config.router_mode,
492
493
                        card.kv_cache_block_size,
                        Some(prefill_config),
494
                        self.router_config.enforce_disagg,
495
496
                        model_name.clone(),
                        namespace.clone(),
497
498
499
                    )
                });

500
501
502
            // Create a new worker monitor for this WorkerSet. Each WorkerSet gets its own
            // monitor (1-to-1) since each monitor is scoped to this WorkerSet's Client/namespace.
            // The monitor tracks Prometheus metrics (active_decode_blocks, active_prefill_tokens,
503
            // worker TTFT/ITL cleanup). The thresholds control busy detection behavior only.
504
            let worker_monitor = Some(KvWorkerMonitor::new(
505
506
507
                client.clone(),
                self.router_config.load_threshold_config.clone(),
            ));
508

509
510
511
512
            // Store KV router and worker monitor on the WorkerSet
            worker_set.kv_router = kv_chooser.clone();
            worker_set.worker_monitor = worker_monitor.clone();

513
            // Add chat engine only if the model supports chat
514
            if card.model_type.supports_chat() {
515
516
517
518
519
520
521
522
523
524
525
                let factory_engine = if let Some(ref factory) = self.chat_engine_factory {
                    match factory(mcid.clone(), card.clone()).await {
                        Ok(engine) => Some(engine),
                        Err(err) => return Err(err).context("python chat_engine_factory"),
                    }
                } else {
                    None
                };

                let chat_engine = if let Some(engine) = factory_engine {
                    engine
526
527
528
529
530
531
532
                } else {
                    entrypoint::build_routed_pipeline::<
                        NvCreateChatCompletionRequest,
                        NvCreateChatCompletionStreamResponse,
                    >(
                        card,
                        &client,
533
                        self.manager.clone(),
534
535
536
537
538
539
                        self.router_config.router_mode,
                        worker_monitor.clone(),
                        kv_chooser.clone(),
                        tokenizer_hf.clone(),
                        prefill_chooser.clone(),
                        self.router_config.enforce_disagg,
540
                        self.migration_limit,
541
                        self.metrics.clone(),
542
543
544
545
                    )
                    .await
                    .context("build_routed_pipeline")?
                };
546
                worker_set.chat_engine = Some(chat_engine);
547
                tracing::info!("Chat completions is ready");
548
            }
549

550
            // Add completions engine only if the model supports completions.
551
            if card.model_type.supports_completions() {
552
553
                let formatter = PromptFormatter::no_op();
                let PromptFormatter::OAI(formatter) = formatter;
554
555
556
557
                let preprocessor = OpenAIPreprocessor::new_with_parts(
                    card.clone(),
                    formatter,
                    tokenizer_hf.clone(),
558
559
                )
                .context("OpenAIPreprocessor::new_with_parts")?;
560
                let completions_engine = entrypoint::build_routed_pipeline_with_preprocessor::<
561
562
563
                    NvCreateCompletionRequest,
                    NvCreateCompletionResponse,
                >(
564
                    card,
565
                    &client,
566
                    self.manager.clone(),
567
                    self.router_config.router_mode,
568
                    worker_monitor,
569
                    kv_chooser,
570
                    preprocessor,
571
                    tokenizer_hf,
572
                    prefill_chooser,
573
                    self.router_config.enforce_disagg,
574
                    self.migration_limit,
575
                    self.metrics.clone(),
576
                )
577
578
                .await
                .context("build_routed_pipeline_with_preprocessor")?;
579
                worker_set.completions_engine = Some(completions_engine);
580
                tracing::info!("Completions is ready");
581
            }
582
583
584
585
586
587
        } 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(
588
                client, self.router_config.router_mode, None, None
589
590
            )
            .await?;
591
            worker_set.embeddings_engine = Some(Arc::new(push_router));
592
593
594
595
596
597
598
599
600
601
        }
        // Case: Text + (Images, Audio, Videos)
        // Must come before the plain Text+Chat / Text+Completions branches because
        // diffusion models often set both Images and Chat flags. The branch below
        // handles the chat registration internally when supports_chat() is true.
        else if card.model_input == ModelInput::Text
            && (card.model_type.supports_images()
                || card.model_type.supports_audios()
                || card.model_type.supports_videos())
        {
602
            // Image/Audio/Video models can also support chat completions (vLLM omni way)
603
604
605
606
607
608
609
610
611
612
613
            if card.model_type.supports_chat() {
                let chat_router = PushRouter::<
                    NvCreateChatCompletionRequest,
                    Annotated<NvCreateChatCompletionStreamResponse>,
                >::from_client_with_threshold(
                    client.clone(),
                    self.router_config.router_mode,
                    None,
                    None,
                )
                .await?;
614
                worker_set.chat_engine = Some(Arc::new(chat_router));
615
616
617
618
619
620
621
622
623
624
            }

            if card.model_type.supports_images() {
                let images_router = PushRouter::<
                    NvCreateImageRequest,
                    Annotated<NvImagesResponse>,
                >::from_client_with_threshold(
                    client.clone(), self.router_config.router_mode, None, None
                )
                .await?;
625
                worker_set.images_engine = Some(Arc::new(images_router));
626
627
628
629
630
631
632
633
634
635
            }

            if card.model_type.supports_videos() {
                let videos_router = PushRouter::<
                    NvCreateVideoRequest,
                    Annotated<NvVideosResponse>,
                >::from_client_with_threshold(
                    client.clone(), self.router_config.router_mode, None, None
                )
                .await?;
636
                worker_set.videos_engine = Some(Arc::new(videos_router));
637
638
639
            }

            // TODO: add audio models support
640
        } else if card.model_input == ModelInput::Text && card.model_type.supports_chat() {
641
            // Case: Text + Chat (pure text-to-text, no diffusion)
642
643
644
645
646
647
648
            let push_router = PushRouter::<
                NvCreateChatCompletionRequest,
                Annotated<NvCreateChatCompletionStreamResponse>,
            >::from_client_with_threshold(
                client, self.router_config.router_mode, None, None
            )
            .await?;
649
            worker_set.chat_engine = Some(Arc::new(push_router));
650
        } else if card.model_input == ModelInput::Text && card.model_type.supports_completions() {
651
            // Case: Text + Completions
652
653
654
655
            let push_router = PushRouter::<
                NvCreateCompletionRequest,
                Annotated<NvCreateCompletionResponse>,
            >::from_client_with_threshold(
656
                client, self.router_config.router_mode, None, None
657
658
            )
            .await?;
659
            worker_set.completions_engine = Some(Arc::new(push_router));
660
        } else if card.model_input == ModelInput::Tokens && card.model_type.supports_embedding() {
661
662
663
664
665
666
667
            // Case 4: Tokens + Embeddings
            // Create preprocessing pipeline similar to Backend
            let frontend = SegmentSource::<
                SingleIn<NvCreateEmbeddingRequest>,
                ManyOut<Annotated<NvCreateEmbeddingResponse>>,
            >::new();

668
            let preprocessor = OpenAIPreprocessor::new(card.clone())?.into_operator();
669
            let backend = Backend::from_mdc(card).into_operator();
670
671
672
673
674

            let router = PushRouter::<
                PreprocessedEmbeddingRequest,
                Annotated<EmbeddingsEngineOutput>,
            >::from_client_with_threshold(
675
                client, self.router_config.router_mode, None, None
676
677
678
            )
            .await?;

679
            // Note: Embeddings don't need KV routing complexity or load monitoring
680
681
682
683
684
685
686
687
688
689
690
            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)?;

691
            worker_set.embeddings_engine = Some(embedding_engine);
692
        } else if card.model_input == ModelInput::Tensor && card.model_type.supports_tensor() {
693
            // Case 6: Tensor + TensorBased (non-LLM)
694
            // No KV cache concepts - not an LLM model
695
696
697
698
            let push_router = PushRouter::<
                NvCreateTensorRequest,
                Annotated<NvCreateTensorResponse>,
            >::from_client_with_threshold(
699
                client, self.router_config.router_mode, None, None
700
701
            )
            .await?;
702
            worker_set.tensor_engine = Some(Arc::new(push_router));
703
704
705
706
707
708
709
710
711
712
713
714
        } 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(),
715
716
717
                "Prefill model detected, registering and activating prefill router"
            );

718
719
            // Prefill sets have no engines — we add the WorkerSet first for tracking,
            // then activate the prefill router.
720
            self.manager
721
                .add_worker_set(card.name(), &ws_key, worker_set)?;
722

723
724
725
726
727
728
729
            // Note: activate_prefill_router is keyed by deployment namespace (not ws_key)
            // because it coordinates between decode and prefill WorkerSets that share
            // the same deployment namespace but have different ws_keys ("ns" vs "ns:prefill").
            let Ok(()) = self
                .manager
                .activate_prefill_router(card.name(), &namespace, endpoint)
            else {
730
731
732
733
734
735
736
737
738
739
                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"
740
            );
741
742

            return Ok(());
743
744
745
746
        } else {
            // Reject unsupported combinations
            anyhow::bail!(
                "Unsupported model configuration: {} with {} input. Supported combinations: \
747
                Tokens+(Chat|Completions|Prefill), Text+(Chat|Completions|Images), Tokens+Embeddings, Tensor+TensorBased",
748
749
                card.model_type,
                card.model_input.as_str()
750
            );
751
        }
Ryan Olson's avatar
Ryan Olson committed
752

753
754
755
756
        // Add the completed WorkerSet to the Model
        self.manager
            .add_worker_set(card.name(), &ws_key, worker_set)?;

757
758
        Ok(())
    }
759

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

765
766
767
        let mut results = Vec::with_capacity(instances.len());
        for instance in instances {
            match instance.deserialize_model::<ModelDeploymentCard>() {
768
                Ok(card) => {
769
770
771
772
773
774
775
776
777
778
779
780
781
782
783
                    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)"
                            );
784
785
786
                            continue;
                        }
                    };
787
                    results.push((endpoint_id, card));
788
                }
789
                Err(err) => {
790
                    tracing::error!(%err, "Failed to deserialize model card");
791
792
                    continue;
                }
793
            }
794
        }
795
        Ok(results)
796
797
    }

798
    pub async fn cards_for_model(
799
800
        &self,
        model_name: &str,
801
        namespace_filter: &NamespaceFilter,
802
    ) -> anyhow::Result<Vec<ModelDeploymentCard>> {
803
804
805
806
807
808
809
810
811
812
813
814
815
816
817
        Ok(self
            .cards_for_model_with_endpoints(model_name, namespace_filter)
            .await?
            .into_iter()
            .map(|(_, card)| card)
            .collect())
    }

    /// Like `cards_for_model` but also returns the EndpointId for each card,
    /// allowing callers to filter by namespace.
    async fn cards_for_model_with_endpoints(
        &self,
        model_name: &str,
        namespace_filter: &NamespaceFilter,
    ) -> anyhow::Result<Vec<(EndpointId, ModelDeploymentCard)>> {
818
819
820
        let mut all = self.all_cards().await?;
        all.retain(|(endpoint_id, card)| {
            let matches_name = card.name() == model_name;
821
            let matches_namespace = namespace_filter.matches(&endpoint_id.namespace);
822
823
            matches_name && matches_namespace
        });
824
825
826
827
828
829
830
831
832
833
834
835
836
837
838
839
840
841
842
843
844
845
846
847
848
849
850
851
852
853
854
855
856
857
858
859
860
861
862
863
864
865
866
867
868
869
870
871
872
873
874
875
876
877
878
879
880
881
882
883
884
885
886
887
888
889
890
891
892
893
894
895
896
        Ok(all)
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::discovery::WorkerSet;
    use crate::model_card::ModelDeploymentCard;

    fn make_worker_set(namespace: &str) -> WorkerSet {
        WorkerSet::new(
            namespace.to_string(),
            "test-checksum".to_string(),
            ModelDeploymentCard::default(),
        )
    }

    #[test]
    fn test_is_model_type_list_empty_on_empty_manager() {
        let mm = ModelManager::new();
        assert!(is_model_type_list_empty(&mm, ModelType::Chat));
        assert!(is_model_type_list_empty(&mm, ModelType::Completions));
        assert!(is_model_type_list_empty(&mm, ModelType::Embedding));
        assert!(is_model_type_list_empty(&mm, ModelType::Images));
        assert!(is_model_type_list_empty(&mm, ModelType::Videos));
        assert!(is_model_type_list_empty(&mm, ModelType::TensorBased));
        assert!(is_model_type_list_empty(&mm, ModelType::Prefill));
    }

    #[test]
    fn test_is_model_type_list_empty_prefill_present() {
        let mm = ModelManager::new();
        // A WorkerSet with no engines is treated as a prefill set
        mm.add_worker_set("model-a", "ns1", make_worker_set("ns1"))
            .unwrap();

        assert!(!is_model_type_list_empty(&mm, ModelType::Prefill));
        // Other types should still be empty since the WorkerSet has no engines
        assert!(is_model_type_list_empty(&mm, ModelType::Chat));
        assert!(is_model_type_list_empty(&mm, ModelType::Completions));
        assert!(is_model_type_list_empty(&mm, ModelType::Embedding));
        assert!(is_model_type_list_empty(&mm, ModelType::Images));
        assert!(is_model_type_list_empty(&mm, ModelType::Videos));
        assert!(is_model_type_list_empty(&mm, ModelType::TensorBased));
    }

    #[test]
    fn test_is_model_type_list_empty_after_removal() {
        let mm = ModelManager::new();
        mm.add_worker_set("model-a", "ns1", make_worker_set("ns1"))
            .unwrap();
        assert!(!is_model_type_list_empty(&mm, ModelType::Prefill));

        mm.remove_model("model-a");
        assert!(is_model_type_list_empty(&mm, ModelType::Prefill));
    }

    #[test]
    fn test_is_model_type_list_not_empty_when_other_model_remains() {
        let mm = ModelManager::new();
        mm.add_worker_set("model-a", "ns1", make_worker_set("ns1"))
            .unwrap();
        mm.add_worker_set("model-b", "ns1", make_worker_set("ns1"))
            .unwrap();

        // Remove one model — other still provides prefill
        mm.remove_model("model-a");
        assert!(!is_model_type_list_empty(&mm, ModelType::Prefill));

        // Remove the last model — now empty
        mm.remove_model("model-b");
        assert!(is_model_type_list_empty(&mm, ModelType::Prefill));
897
898
    }
}