"components/backends/trtllm/README.md" did not exist on "58df5aca4c3f58feaa903f76b06b09b87a6f6f28"
watcher.rs 39.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
215
216
217
                    // If a WorkerSet already exists for this (model, namespace, type),
                    // validate that the new worker's checksum matches. Different
                    // WorkerSets (different namespaces) are allowed to have different checksums to support rolling updates.
                    let ws_key = worker_set_key(&mcid.namespace, card.model_type);
                    if let Some(model) = self.manager.get_model(card.name())
                        && !model.is_checksum_compatible(&ws_key, card.mdcsum())
                    {
218
219
                        tracing::error!(
                            model_name = card.name(),
220
                            namespace = mcid.namespace,
221
222
223
                            new_checksum = card.mdcsum(),
                            "Checksum for new worker does not match existing WorkerSet's checksum. \
                             Drain all old workers in this namespace before deploying a new version."
224
225
226
227
228
229
230
231
232
233
                        );
                        // 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;
                    }
234

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

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

285
286
    /// 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.
287
288
    async fn handle_delete(
        &self,
289
        mcid: &ModelCardInstanceId,
290
        namespace_filter: &NamespaceFilter,
291
    ) -> anyhow::Result<Option<String>> {
292
293
        let key = mcid.to_path();
        let card = match self.manager.remove_model_card(&key) {
294
            Some(card) => card,
295
            None => {
296
                anyhow::bail!("Missing ModelDeploymentCard for {}", key);
297
298
            }
        };
299
        let model_name = card.name().to_string();
300
301
302
303
304
        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
305
        let active_instances = self
306
            .cards_for_model_with_endpoints(&model_name, namespace_filter)
307
308
            .await
            .with_context(|| model_name.clone())?;
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332

        // 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
333
        if !active_instances.is_empty() {
334
335
336
            tracing::debug!(
                model_name,
                active_instance_count = active_instances.len(),
337
                "Model has other active instances in other namespaces"
338
            );
339
340
            return Ok(None);
        }
341

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

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

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

358
    // Handles a PUT event from store, this usually means adding a new model to the list of served
359
    // models.
360
361
    async fn handle_put(
        &self,
362
        mcid: &ModelCardInstanceId,
363
364
        card: &mut ModelDeploymentCard,
    ) -> anyhow::Result<()> {
365
366
367
368
369
        // 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);
370

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

384
385
386
387
388
389
        // 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())
        {
390
391
392
393
            self.manager
                .save_model_card(&mcid.to_path(), card.clone())?;
            tracing::debug!(
                model_name = card.name(),
394
395
                namespace = namespace,
                "WorkerSet registration in progress, skipping"
396
397
398
399
            );
            return Ok(());
        }

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

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

        result
    }

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

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

436
        let checksum = card.mdcsum();
437
438
439
440
441
442
        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);
443
444
445

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

451
            let endpoint = component.endpoint(&mcid.endpoint);
452
453
454
455
456
457
458
459
460
            // 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)
            {
461
462
                Some(
                    self.manager
463
464
465
                        .kv_chooser_for(
                            &endpoint,
                            card.kv_cache_block_size,
466
                            Some(self.router_config.kv_router_config.clone()),
467
                            WORKER_TYPE_DECODE, // This is the decode router
468
                            Some(card.display_name.clone()),
469
                            card.runtime_config.enable_eagle,
470
                        )
471
472
473
474
475
                        .await?,
                )
            } else {
                None
            };
476

477
478
479
480
481
482
483
484
485
486
487
488
489
            // Loading the tokenizer is expensive (~10 MiB JSON), so only do it
            // once and only when a local pipeline actually needs it.  Models
            // without tokenizer.json (e.g. Qwen3-Omni) set tokenizer = None;
            // they rely on a Python chat_engine_factory for tokenization.
            // When a chat_engine_factory handles chat and no completions are
            // needed, skip tokenizer loading entirely — even if the file exists.
            let needs_rust_tokenizer =
                needs_local_chat_pipeline || needs_local_completions_pipeline;
            let tokenizer = if needs_rust_tokenizer && card.has_tokenizer() {
                Some(card.tokenizer().context("tokenizer")?)
            } else {
                None
            };
490

491
492
            // Create prefill chooser once if we're building pipelines
            // Both chat and completions will share the same prefill chooser instance
493
            let model_name = card.name().to_string();
494
495
            let prefill_chooser = self
                .manager
496
                .register_prefill_router(&model_name, &namespace)
497
498
                .map(|rx| {
                    // Create prefill-specific config with track_active_blocks disabled
499
                    let mut prefill_config = self.router_config.kv_router_config.clone();
500
501
502
503
504
                    prefill_config.router_track_active_blocks = false;

                    PrefillRouter::new(
                        rx,
                        self.manager.clone(),
505
                        self.router_config.router_mode,
506
507
                        card.kv_cache_block_size,
                        Some(prefill_config),
508
                        self.router_config.enforce_disagg,
509
510
                        model_name.clone(),
                        namespace.clone(),
511
                        card.runtime_config.enable_eagle,
512
513
514
                    )
                });

515
516
517
            // 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,
518
            // worker TTFT/ITL cleanup). The thresholds control busy detection behavior only.
519
520
521
522
523
524
525
526
527
528
            //
            // IMPORTANT: When KV routing is active, the monitor must use the KvRouter's Client
            // so that busy-state updates (via update_free_instances) are visible to the
            // PushRouter, which also uses the KvRouter's Client (see common.rs:258-263).
            // Using a different Client instance would cause the PushRouter to never see
            // busy workers, since each Client::new() creates independent ArcSwap state.
            let monitor_client = kv_chooser
                .as_ref()
                .map(|chooser| chooser.client().clone())
                .unwrap_or_else(|| client.clone());
529
            let worker_monitor = Some(KvWorkerMonitor::new(
530
                monitor_client,
531
532
                self.router_config.load_threshold_config.clone(),
            ));
533

534
535
536
537
            // Store KV router and worker monitor on the WorkerSet
            worker_set.kv_router = kv_chooser.clone();
            worker_set.worker_monitor = worker_monitor.clone();

538
            // Add chat engine only if the model supports chat
539
            if card.model_type.supports_chat() {
540
541
542
543
544
545
546
547
548
549
550
                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
551
                } else {
552
553
554
555
556
557
558
                    let tk = tokenizer.clone().ok_or_else(|| {
                        anyhow::anyhow!(
                            "Model has no supported Rust tokenizer and no chat_engine_factory. \
                             Use --dyn-chat-processor vllm/sglang or provide a supported \
                             tokenizer file (tokenizer.json, tiktoken.model, or *.tiktoken)."
                        )
                    })?;
559
560
561
562
563
564
                    entrypoint::build_routed_pipeline::<
                        NvCreateChatCompletionRequest,
                        NvCreateChatCompletionStreamResponse,
                    >(
                        card,
                        &client,
565
                        self.manager.clone(),
566
567
568
                        self.router_config.router_mode,
                        worker_monitor.clone(),
                        kv_chooser.clone(),
569
                        tk,
570
                        prefill_chooser.clone(),
571
                        self.router_config.enforce_disagg,
572
                        self.migration_limit,
573
                        self.metrics.clone(),
574
575
576
577
                    )
                    .await
                    .context("build_routed_pipeline")?
                };
578
                worker_set.chat_engine = Some(chat_engine);
579
                tracing::info!("Chat completions is ready");
580
            }
581

582
583
            // Add completions engine only if the model supports completions
            // and we have a tokenizer (completions always uses the Rust preprocessor).
584
            if card.model_type.supports_completions() {
585
586
587
588
589
590
591
592
593
594
595
596
597
598
599
600
601
602
603
604
605
606
607
608
609
610
611
612
613
614
615
616
617
618
619
620
621
622
623
624
625
626
627
628
629
                if let Some(tk) = tokenizer {
                    let formatter = PromptFormatter::no_op();
                    let PromptFormatter::OAI(formatter) = formatter;
                    let preprocessor =
                        OpenAIPreprocessor::new_with_parts(card.clone(), formatter, tk.clone())
                            .context("OpenAIPreprocessor::new_with_parts")?;
                    let completions_engine = entrypoint::build_routed_pipeline_with_preprocessor::<
                        NvCreateCompletionRequest,
                        NvCreateCompletionResponse,
                    >(
                        card,
                        &client,
                        self.manager.clone(),
                        self.router_config.router_mode,
                        worker_monitor,
                        kv_chooser,
                        preprocessor,
                        tk,
                        prefill_chooser,
                        self.router_config.enforce_disagg,
                        self.migration_limit,
                        self.metrics.clone(),
                    )
                    .await
                    .context("build_routed_pipeline_with_preprocessor")?;
                    worker_set.completions_engine = Some(completions_engine);
                    tracing::info!("Completions is ready");
                } else {
                    tracing::warn!(
                        "Skipping completions engine: no Rust tokenizer available for this model"
                    );
                }
            }

            // Verify we built at least one serving engine. A Tokens model that
            // ends up with no chat AND no completions engine (e.g. completions-only
            // model with no tokenizer) should fail fast rather than register an
            // empty WorkerSet that can't serve any requests.
            if !worker_set.has_decode_engine() {
                anyhow::bail!(
                    "Model '{}' requires frontend tokenization/preprocessing (ModelInput::Tokens) \
                     but no serving engine could be built. Provide a working tokenizer config or \
                     perform tokenization in the backend (ModelInput::Text).",
                    card.name()
                );
630
            }
631
632
633
634
635
636
        } 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(
637
                client, self.router_config.router_mode, None, None
638
639
            )
            .await?;
640
            worker_set.embeddings_engine = Some(Arc::new(push_router));
641
642
643
644
645
646
647
648
649
650
        }
        // 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())
        {
651
            // Image/Audio/Video models can also support chat completions (vLLM omni way)
652
653
654
655
656
657
658
659
660
661
662
            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?;
663
                worker_set.chat_engine = Some(Arc::new(chat_router));
664
665
666
667
668
669
670
671
672
673
            }

            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?;
674
                worker_set.images_engine = Some(Arc::new(images_router));
675
676
677
678
679
680
681
682
683
684
            }

            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?;
685
                worker_set.videos_engine = Some(Arc::new(videos_router));
686
687
688
            }

            // TODO: add audio models support
689
        } else if card.model_input == ModelInput::Text && card.model_type.supports_chat() {
690
            // Case: Text + Chat (pure text-to-text, no diffusion)
691
692
693
694
695
696
697
            let push_router = PushRouter::<
                NvCreateChatCompletionRequest,
                Annotated<NvCreateChatCompletionStreamResponse>,
            >::from_client_with_threshold(
                client, self.router_config.router_mode, None, None
            )
            .await?;
698
            worker_set.chat_engine = Some(Arc::new(push_router));
699
        } else if card.model_input == ModelInput::Text && card.model_type.supports_completions() {
700
            // Case: Text + Completions
701
702
703
704
            let push_router = PushRouter::<
                NvCreateCompletionRequest,
                Annotated<NvCreateCompletionResponse>,
            >::from_client_with_threshold(
705
                client, self.router_config.router_mode, None, None
706
707
            )
            .await?;
708
            worker_set.completions_engine = Some(Arc::new(push_router));
709
        } else if card.model_input == ModelInput::Tokens && card.model_type.supports_embedding() {
710
711
712
713
714
715
716
            // Case 4: Tokens + Embeddings
            // Create preprocessing pipeline similar to Backend
            let frontend = SegmentSource::<
                SingleIn<NvCreateEmbeddingRequest>,
                ManyOut<Annotated<NvCreateEmbeddingResponse>>,
            >::new();

717
            let preprocessor = OpenAIPreprocessor::new(card.clone())?.into_operator();
718
            let backend = Backend::from_mdc(card).into_operator();
719
720
721
722
723

            let router = PushRouter::<
                PreprocessedEmbeddingRequest,
                Annotated<EmbeddingsEngineOutput>,
            >::from_client_with_threshold(
724
                client, self.router_config.router_mode, None, None
725
726
727
            )
            .await?;

728
            // Note: Embeddings don't need KV routing complexity or load monitoring
729
730
731
732
733
734
735
736
737
738
739
            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)?;

740
            worker_set.embeddings_engine = Some(embedding_engine);
741
        } else if card.model_input == ModelInput::Tensor && card.model_type.supports_tensor() {
742
            // Case 6: Tensor + TensorBased (non-LLM)
743
            // No KV cache concepts - not an LLM model
744
745
746
747
            let push_router = PushRouter::<
                NvCreateTensorRequest,
                Annotated<NvCreateTensorResponse>,
            >::from_client_with_threshold(
748
                client, self.router_config.router_mode, None, None
749
750
            )
            .await?;
751
            worker_set.tensor_engine = Some(Arc::new(push_router));
752
753
754
755
756
757
758
759
760
761
762
763
        } 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(),
764
765
766
                "Prefill model detected, registering and activating prefill router"
            );

767
768
            // Prefill sets have no engines — we add the WorkerSet first for tracking,
            // then activate the prefill router.
769
            self.manager
770
                .add_worker_set(card.name(), &ws_key, worker_set);
771

772
773
774
775
776
777
778
            // 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 {
779
780
781
782
783
784
785
786
787
788
                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"
789
            );
790
791

            return Ok(());
792
793
794
795
        } else {
            // Reject unsupported combinations
            anyhow::bail!(
                "Unsupported model configuration: {} with {} input. Supported combinations: \
796
                Tokens+(Chat|Completions|Prefill), Text+(Chat|Completions|Images), Tokens+Embeddings, Tensor+TensorBased",
797
798
                card.model_type,
                card.model_input.as_str()
799
            );
800
        }
Ryan Olson's avatar
Ryan Olson committed
801

802
803
        // Add the completed WorkerSet to the Model
        self.manager
804
            .add_worker_set(card.name(), &ws_key, worker_set);
805

806
807
        Ok(())
    }
808

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

814
815
816
        let mut results = Vec::with_capacity(instances.len());
        for instance in instances {
            match instance.deserialize_model::<ModelDeploymentCard>() {
817
                Ok(card) => {
818
819
820
821
822
823
824
825
826
827
828
829
830
831
832
                    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)"
                            );
833
834
835
                            continue;
                        }
                    };
836
                    results.push((endpoint_id, card));
837
                }
838
                Err(err) => {
839
                    tracing::error!(%err, "Failed to deserialize model card");
840
841
                    continue;
                }
842
            }
843
        }
844
        Ok(results)
845
846
    }

847
    pub async fn cards_for_model(
848
849
        &self,
        model_name: &str,
850
        namespace_filter: &NamespaceFilter,
851
    ) -> anyhow::Result<Vec<ModelDeploymentCard>> {
852
853
854
855
856
857
858
859
860
861
862
863
864
865
866
        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)>> {
867
868
869
        let mut all = self.all_cards().await?;
        all.retain(|(endpoint_id, card)| {
            let matches_name = card.name() == model_name;
870
            let matches_namespace = namespace_filter.matches(&endpoint_id.namespace);
871
872
            matches_name && matches_namespace
        });
873
874
875
876
877
878
879
880
881
882
883
884
885
886
887
888
889
890
891
892
893
894
895
896
897
898
899
900
901
902
903
904
905
906
        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
907
        mm.add_worker_set("model-a", "ns1", make_worker_set("ns1"));
908
909
910
911
912
913
914
915
916
917
918
919
920
921

        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();
922
        mm.add_worker_set("model-a", "ns1", make_worker_set("ns1"));
923
924
925
926
927
928
929
930
931
        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();
932
933
        mm.add_worker_set("model-a", "ns1", make_worker_set("ns1"));
        mm.add_worker_set("model-b", "ns1", make_worker_set("ns1"));
934
935
936
937
938
939
940
941

        // 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));
942
943
    }
}