watcher.rs 41.4 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 dynamo_kv_router::PrefillLoadEstimator;
11
use futures::StreamExt;
Ryan Olson's avatar
Ryan Olson committed
12

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

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

51
use super::ModelManager;
52
53
54
55
56
57
58
59
60
61
62
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()
    }
}
63

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

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

85
86
87
88
const ALL_MODEL_TYPES: &[ModelType] = &[
    ModelType::Chat,
    ModelType::Completions,
    ModelType::Embedding,
89
    ModelType::Images,
90
    ModelType::Audios,
91
92
    ModelType::Videos,
    ModelType::TensorBased,
93
    ModelType::Prefill,
94
];
95

96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
/// 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
    }
}

117
impl ModelWatcher {
118
    #[allow(clippy::too_many_arguments)]
119
    pub fn new(
120
        runtime: DistributedRuntime,
121
        model_manager: Arc<ModelManager>,
122
        router_config: RouterConfig,
123
        migration_limit: u32,
124
        migration_max_seq_len: Option<u32>,
125
        chat_engine_factory: Option<ChatEngineFactoryCallback>,
126
        prefill_load_estimator: Option<Arc<dyn PrefillLoadEstimator>>,
127
        metrics: Arc<Metrics>,
128
129
    ) -> ModelWatcher {
        Self {
130
            manager: model_manager,
131
            drt: runtime,
132
            router_config,
133
            migration_limit,
134
            migration_max_seq_len,
135
            notify_on_model: Notify::new(),
136
            model_update_tx: None,
137
            chat_engine_factory,
138
            prefill_load_estimator,
139
            metrics,
140
            registering_worker_sets: DashSet::new(),
141
        }
142
    }
Ryan Olson's avatar
Ryan Olson committed
143

144
145
146
147
    pub fn set_notify_on_model_update(&mut self, tx: Sender<ModelUpdate>) {
        self.model_update_tx = Some(tx);
    }

148
149
150
151
152
153
154
155
156
157
158
    /// 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
        }
    }

159
    /// Common watch logic with optional namespace filtering
160
161
162
    pub async fn watch(
        &self,
        mut discovery_stream: DiscoveryStream,
163
        namespace_filter: NamespaceFilter,
164
165
166
167
168
169
170
171
172
173
    ) {
        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;
                }
            };

174
            match event {
175
                DiscoveryEvent::Added(instance) => {
176
177
                    // Extract ModelCardInstanceId and card from the discovery instance
                    let (mcid, mut card) = match &instance {
178
179
180
181
182
                        DiscoveryInstance::Model {
                            namespace,
                            component,
                            endpoint,
                            instance_id,
183
                            model_suffix,
184
185
                            ..
                        } => {
186
                            let mcid = ModelCardInstanceId {
187
188
                                namespace: namespace.clone(),
                                component: component.clone(),
189
190
191
                                endpoint: endpoint.clone(),
                                instance_id: *instance_id,
                                model_suffix: model_suffix.clone(),
192
193
194
                            };

                            match instance.deserialize_model::<ModelDeploymentCard>() {
195
                                Ok(card) => (mcid, card),
196
197
198
199
200
201
202
203
204
205
                                Err(err) => {
                                    tracing::error!(%err, instance_id, "Failed to deserialize model card");
                                    continue;
                                }
                            }
                        }
                        _ => {
                            tracing::error!(
                                "Unexpected discovery instance type (expected ModelCard)"
                            );
206
207
208
                            continue;
                        }
                    };
209

210
211
                    // Filter by namespace using the configured filter
                    if !namespace_filter.matches(&mcid.namespace) {
212
                        tracing::debug!(
213
                            model_namespace = mcid.namespace,
214
215
                            namespace_filter = ?namespace_filter,
                            "Skipping model due to namespace filter"
216
217
218
219
                        );
                        continue;
                    }

220
221
222
223
224
225
226
                    // 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())
                    {
227
228
                        tracing::error!(
                            model_name = card.name(),
229
                            namespace = mcid.namespace,
230
231
232
                            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."
233
234
235
236
237
238
239
240
241
242
                        );
                        // 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;
                    }
243

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

275
                    match self
276
                        .handle_delete(model_card_instance_id, &namespace_filter)
277
278
279
280
281
282
283
284
285
286
287
                        .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");
                        }
288
                    }
289
                }
290
            }
Ryan Olson's avatar
Ryan Olson committed
291
292
293
        }
    }

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

        // 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)"
                );
            }
339
340
341
342
343
344
345
346
347

            // If the removed component was a prefill worker, deactivate the decode-side
            // prefill router so requests fall back to aggregated mode (or fail cleanly
            // with enforce_disagg). The decode WorkerSet's namespace matches the
            // deployment namespace, not the ws_key.
            if card.model_type.supports_prefill() {
                self.manager
                    .deactivate_prefill_router_for_decode(&model_name, worker_namespace);
            }
348
349
350
        }

        // Check if the Model still has instances in any namespace
351
        if !active_instances.is_empty() {
352
353
354
            tracing::debug!(
                model_name,
                active_instance_count = active_instances.len(),
355
                "Model has other active instances in other namespaces"
356
            );
357
358
            return Ok(None);
        }
359

360
361
        // No instances remain anywhere — remove the entire Model
        let _ = self.manager.remove_model(&model_name);
362

363
        if let Some(tx) = &self.model_update_tx {
364
            for model_type in ALL_MODEL_TYPES {
365
366
                if card.model_type.intersects(*model_type)
                    && is_model_type_list_empty(&self.manager, *model_type)
367
                {
368
                    tx.send(ModelUpdate::Removed(card.clone())).await.ok();
369
370
371
                }
            }
        }
372

373
        Ok(Some(model_name))
374
    }
Ryan Olson's avatar
Ryan Olson committed
375

376
    // Handles a PUT event from store, this usually means adding a new model to the list of served
377
    // models.
378
379
    async fn handle_put(
        &self,
380
        mcid: &ModelCardInstanceId,
381
382
        card: &mut ModelDeploymentCard,
    ) -> anyhow::Result<()> {
383
384
385
386
387
        // 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);
388

389
390
391
        if let Some(model) = self.manager.get_model(&model_name)
            && model.has_worker_set(&ws_key)
        {
392
393
            self.manager
                .save_model_card(&mcid.to_path(), card.clone())?;
394
395
            tracing::debug!(
                model_name = card.name(),
396
397
                namespace = namespace,
                "Worker joined existing WorkerSet, skipping pipeline build"
398
399
400
401
            );
            return Ok(());
        }

402
403
404
405
406
407
        // 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())
        {
408
409
410
411
            self.manager
                .save_model_card(&mcid.to_path(), card.clone())?;
            tracing::debug!(
                model_name = card.name(),
412
413
                namespace = namespace,
                "WorkerSet registration in progress, skipping"
414
415
416
417
            );
            return Ok(());
        }

418
        let result = self.do_worker_set_registration(mcid, card).await;
419

420
421
        // Always remove from registering set
        self.registering_worker_sets.remove(&registration_key);
422
423
424
425

        result
    }

426
427
428
    /// Build a complete WorkerSet with all engines for this (model, namespace)
    /// and add it to the Model.
    async fn do_worker_set_registration(
429
430
431
432
433
434
435
436
437
438
439
440
        &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?;
441
442
443
444
445
446
        let instance_watcher = client.instance_avail_watcher();
        tracing::debug!(
            model_name = card.name(),
            namespace = mcid.namespace,
            "building worker set pipeline"
        );
447
448
449
        self.manager
            .save_model_card(&mcid.to_path(), card.clone())?;

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

454
        let checksum = card.mdcsum();
455
456
457
458
459
460
        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);
461
462
463

        if card.model_input == ModelInput::Tokens
            && (card.model_type.supports_chat() || card.model_type.supports_completions())
464
465
466
467
468
        {
            // 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.

469
            let endpoint = component.endpoint(&mcid.endpoint);
470
471
472
473
474
475
476
477
478
            // 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)
            {
479
480
                Some(
                    self.manager
481
482
483
                        .kv_chooser_for(
                            &endpoint,
                            card.kv_cache_block_size,
484
                            Some(self.router_config.kv_router_config.clone()),
485
                            self.prefill_load_estimator.clone(),
486
                            WORKER_TYPE_DECODE, // This is the decode router
487
                            Some(card.display_name.clone()),
488
                            card.runtime_config.enable_eagle,
489
                        )
490
491
492
493
494
                        .await?,
                )
            } else {
                None
            };
495

496
497
498
499
500
501
502
503
504
505
506
507
508
            // 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
            };
509

510
511
            // Create prefill chooser once if we're building pipelines
            // Both chat and completions will share the same prefill chooser instance
512
            let model_name = card.name().to_string();
513
514
            let prefill_chooser = self
                .manager
515
                .register_prefill_router(&model_name, &namespace)
516
517
                .map(|rx| {
                    // Create prefill-specific config with track_active_blocks disabled
518
                    let mut prefill_config = self.router_config.kv_router_config.clone();
519
520
521
522
523
                    prefill_config.router_track_active_blocks = false;

                    PrefillRouter::new(
                        rx,
                        self.manager.clone(),
524
                        self.router_config.router_mode,
525
526
                        card.kv_cache_block_size,
                        Some(prefill_config),
527
                        self.prefill_load_estimator.clone(),
528
                        self.router_config.enforce_disagg,
529
530
                        model_name.clone(),
                        namespace.clone(),
531
                        card.runtime_config.enable_eagle,
532
533
534
                    )
                });

535
536
537
            // 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,
538
            // worker TTFT/ITL cleanup). The thresholds control busy detection behavior only.
539
540
541
542
543
544
545
546
547
548
            //
            // 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());
549
            let worker_monitor = Some(KvWorkerMonitor::new(
550
                monitor_client,
551
552
                self.router_config.load_threshold_config.clone(),
            ));
553

554
555
556
            // Store KV router, worker monitor, and prefill router on the WorkerSet.
            // The prefill router is stored so the watcher can deactivate/reactivate it
            // when prefill workers die or rejoin.
557
558
            worker_set.kv_router = kv_chooser.clone();
            worker_set.worker_monitor = worker_monitor.clone();
559
            worker_set.prefill_router = prefill_chooser.clone();
560

561
            // Add chat engine only if the model supports chat
562
            if card.model_type.supports_chat() {
563
564
565
566
567
568
569
570
571
572
573
                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
574
                } else {
575
576
577
578
579
580
581
                    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)."
                        )
                    })?;
582
583
584
585
586
587
                    entrypoint::build_routed_pipeline::<
                        NvCreateChatCompletionRequest,
                        NvCreateChatCompletionStreamResponse,
                    >(
                        card,
                        &client,
588
                        self.manager.clone(),
589
590
591
                        self.router_config.router_mode,
                        worker_monitor.clone(),
                        kv_chooser.clone(),
592
                        tk,
593
                        prefill_chooser.clone(),
594
                        self.router_config.enforce_disagg,
595
                        self.migration_limit,
596
                        self.migration_max_seq_len,
597
                        self.metrics.clone(),
598
599
600
601
                    )
                    .await
                    .context("build_routed_pipeline")?
                };
602
                worker_set.chat_engine = Some(chat_engine);
603
                tracing::info!("Chat completions is ready");
604
            }
605

606
607
            // Add completions engine only if the model supports completions
            // and we have a tokenizer (completions always uses the Rust preprocessor).
608
            if card.model_type.supports_completions() {
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,
630
                        self.migration_max_seq_len,
631
632
633
634
635
636
637
638
639
640
641
642
643
644
645
646
647
648
649
650
651
652
653
654
                        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()
                );
655
            }
656
657
658
659
660
661
        } 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(
662
                client, self.router_config.router_mode, None, None
663
664
            )
            .await?;
665
            worker_set.embeddings_engine = Some(Arc::new(push_router));
666
667
668
669
670
671
672
673
674
675
        }
        // 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())
        {
676
            // Image/Audio/Video models can also support chat completions (vLLM omni way)
677
678
679
680
681
682
683
684
685
686
687
            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?;
688
                worker_set.chat_engine = Some(Arc::new(chat_router));
689
690
691
692
693
694
695
696
697
698
            }

            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?;
699
                worker_set.images_engine = Some(Arc::new(images_router));
700
701
702
703
704
705
706
707
708
709
            }

            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?;
710
                worker_set.videos_engine = Some(Arc::new(videos_router));
711
712
            }

713
714
715
716
717
718
719
720
721
722
723
724
725
            if card.model_type.supports_audios() {
                let audios_router = PushRouter::<
                    NvCreateAudioSpeechRequest,
                    Annotated<NvAudioSpeechResponse>,
                >::from_client_with_threshold(
                    client.clone(),
                    self.router_config.router_mode,
                    None,
                    None,
                )
                .await?;
                worker_set.audios_engine = Some(Arc::new(audios_router));
            }
726
        } else if card.model_input == ModelInput::Text && card.model_type.supports_chat() {
727
            // Case: Text + Chat (pure text-to-text, no diffusion)
728
729
730
731
732
733
734
            let push_router = PushRouter::<
                NvCreateChatCompletionRequest,
                Annotated<NvCreateChatCompletionStreamResponse>,
            >::from_client_with_threshold(
                client, self.router_config.router_mode, None, None
            )
            .await?;
735
            worker_set.chat_engine = Some(Arc::new(push_router));
736
        } else if card.model_input == ModelInput::Text && card.model_type.supports_completions() {
737
            // Case: Text + Completions
738
739
740
741
            let push_router = PushRouter::<
                NvCreateCompletionRequest,
                Annotated<NvCreateCompletionResponse>,
            >::from_client_with_threshold(
742
                client, self.router_config.router_mode, None, None
743
744
            )
            .await?;
745
            worker_set.completions_engine = Some(Arc::new(push_router));
746
        } else if card.model_input == ModelInput::Tokens && card.model_type.supports_embedding() {
747
748
749
750
751
752
753
            // Case 4: Tokens + Embeddings
            // Create preprocessing pipeline similar to Backend
            let frontend = SegmentSource::<
                SingleIn<NvCreateEmbeddingRequest>,
                ManyOut<Annotated<NvCreateEmbeddingResponse>>,
            >::new();

754
            let preprocessor = OpenAIPreprocessor::new(card.clone())?.into_operator();
755
            let backend = Backend::from_mdc(card).into_operator();
756
757
758
759
760

            let router = PushRouter::<
                PreprocessedEmbeddingRequest,
                Annotated<EmbeddingsEngineOutput>,
            >::from_client_with_threshold(
761
                client, self.router_config.router_mode, None, None
762
763
764
            )
            .await?;

765
            // Note: Embeddings don't need KV routing complexity or load monitoring
766
767
768
769
770
771
772
773
774
775
776
            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)?;

777
            worker_set.embeddings_engine = Some(embedding_engine);
778
        } else if card.model_input == ModelInput::Tensor && card.model_type.supports_tensor() {
779
            // Case 6: Tensor + TensorBased (non-LLM)
780
            // No KV cache concepts - not an LLM model
781
782
783
784
            let push_router = PushRouter::<
                NvCreateTensorRequest,
                Annotated<NvCreateTensorResponse>,
            >::from_client_with_threshold(
785
                client, self.router_config.router_mode, None, None
786
787
            )
            .await?;
788
            worker_set.tensor_engine = Some(Arc::new(push_router));
789
790
791
792
793
794
795
796
797
798
799
800
        } 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(),
801
802
803
                "Prefill model detected, registering and activating prefill router"
            );

804
805
            // Prefill sets have no engines — we add the WorkerSet first for tracking,
            // then activate the prefill router.
806
            self.manager
807
                .add_worker_set(card.name(), &ws_key, worker_set);
808

809
810
811
812
813
814
815
            // 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 {
816
817
818
819
820
821
822
823
824
825
                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"
826
            );
827
828

            return Ok(());
829
830
831
832
        } else {
            // Reject unsupported combinations
            anyhow::bail!(
                "Unsupported model configuration: {} with {} input. Supported combinations: \
833
                Tokens+(Chat|Completions|Prefill), Text+(Chat|Completions|Images), Tokens+Embeddings, Tensor+TensorBased",
834
835
                card.model_type,
                card.model_input.as_str()
836
            );
837
        }
Ryan Olson's avatar
Ryan Olson committed
838

839
840
        // Add the completed WorkerSet to the Model
        self.manager
841
            .add_worker_set(card.name(), &ws_key, worker_set);
842

843
844
        Ok(())
    }
845

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

851
852
853
        let mut results = Vec::with_capacity(instances.len());
        for instance in instances {
            match instance.deserialize_model::<ModelDeploymentCard>() {
854
                Ok(card) => {
855
856
857
858
859
860
861
862
863
864
865
866
867
868
869
                    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)"
                            );
870
871
872
                            continue;
                        }
                    };
873
                    results.push((endpoint_id, card));
874
                }
875
                Err(err) => {
876
                    tracing::error!(%err, "Failed to deserialize model card");
877
878
                    continue;
                }
879
            }
880
        }
881
        Ok(results)
882
883
    }

884
    pub async fn cards_for_model(
885
886
        &self,
        model_name: &str,
887
        namespace_filter: &NamespaceFilter,
888
    ) -> anyhow::Result<Vec<ModelDeploymentCard>> {
889
890
891
892
893
894
895
896
897
898
899
900
901
902
903
        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)>> {
904
905
906
        let mut all = self.all_cards().await?;
        all.retain(|(endpoint_id, card)| {
            let matches_name = card.name() == model_name;
907
            let matches_namespace = namespace_filter.matches(&endpoint_id.namespace);
908
909
            matches_name && matches_namespace
        });
910
911
912
913
914
915
916
917
918
919
920
921
922
923
924
925
926
927
928
929
930
931
932
933
934
935
936
937
938
939
940
941
942
943
        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
944
        mm.add_worker_set("model-a", "ns1", make_worker_set("ns1"));
945
946
947
948
949
950
951
952
953
954
955
956
957
958

        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();
959
        mm.add_worker_set("model-a", "ns1", make_worker_set("ns1"));
960
961
962
963
964
965
966
967
968
        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();
969
970
        mm.add_worker_set("model-a", "ns1", make_worker_set("ns1"));
        mm.add_worker_set("model-b", "ns1", make_worker_set("ns1"));
971
972
973
974
975
976
977
978

        // 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));
979
980
    }
}