watcher.rs 47.5 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 std::time::Duration;
6
use tokio::sync::Notify;
7
use tokio::sync::mpsc::Sender;
8
use tokio::task::JoinHandle;
9

10
use anyhow::Context as _;
11
use dashmap::{DashMap, DashSet};
12
use dynamo_kv_router::PrefillLoadEstimator;
13
use futures::StreamExt;
Ryan Olson's avatar
Ryan Olson committed
14

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

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

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

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

72
pub struct ModelWatcher {
73
    manager: Arc<ModelManager>,
74
    drt: DistributedRuntime,
75
    router_config: RouterConfig,
76
    migration_limit: u32,
77
    migration_max_seq_len: Option<u32>,
78
    notify_on_model: Notify,
79
    model_update_tx: Option<Sender<ModelUpdate>>,
80
    chat_engine_factory: Option<ChatEngineFactoryCallback>,
81
    prefill_load_estimator: Option<Arc<dyn PrefillLoadEstimator>>,
82
    metrics: Arc<Metrics>,
83
84
    /// Guards against concurrent pipeline construction for the same (model, namespace).
    registering_worker_sets: DashSet<String>,
85
86
87
    /// Tracks in-flight `handle_put` tasks by instance path so that `handle_delete`
    /// can await a racing put before proceeding with cleanup.
    pending_puts: DashMap<String, JoinHandle<()>>,
Ryan Olson's avatar
Ryan Olson committed
88
89
}

90
91
92
93
const ALL_MODEL_TYPES: &[ModelType] = &[
    ModelType::Chat,
    ModelType::Completions,
    ModelType::Embedding,
94
    ModelType::Images,
95
    ModelType::Audios,
96
97
    ModelType::Videos,
    ModelType::TensorBased,
98
    ModelType::Prefill,
99
];
100

101
102
103
104
105
106
107
108
109
110
/// 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()
111
112
    } else if model_type == ModelType::Audios {
        manager.list_audios_models().is_empty()
113
114
115
116
117
118
119
120
121
122
123
    } 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
    }
}

124
125
126
127
128
129
130
131
132
133
134
135
136
137
/// RAII guard that removes a key from a `DashSet` on drop.
/// Ensures `registering_worker_sets` is cleaned up even if the registration
/// task panics, preventing permanent poisoning of the registration key.
struct RegistrationGuard<'a> {
    set: &'a DashSet<String>,
    key: String,
}

impl Drop for RegistrationGuard<'_> {
    fn drop(&mut self) {
        self.set.remove(&self.key);
    }
}

138
impl ModelWatcher {
139
    #[allow(clippy::too_many_arguments)]
140
    pub fn new(
141
        runtime: DistributedRuntime,
142
        model_manager: Arc<ModelManager>,
143
        router_config: RouterConfig,
144
        migration_limit: u32,
145
        migration_max_seq_len: Option<u32>,
146
        chat_engine_factory: Option<ChatEngineFactoryCallback>,
147
        prefill_load_estimator: Option<Arc<dyn PrefillLoadEstimator>>,
148
        metrics: Arc<Metrics>,
149
150
    ) -> ModelWatcher {
        Self {
151
            manager: model_manager,
152
            drt: runtime,
153
            router_config,
154
            migration_limit,
155
            migration_max_seq_len,
156
            notify_on_model: Notify::new(),
157
            model_update_tx: None,
158
            chat_engine_factory,
159
            prefill_load_estimator,
160
            metrics,
161
            registering_worker_sets: DashSet::new(),
162
            pending_puts: DashMap::new(),
163
        }
164
    }
Ryan Olson's avatar
Ryan Olson committed
165

166
167
168
169
    pub fn set_notify_on_model_update(&mut self, tx: Sender<ModelUpdate>) {
        self.model_update_tx = Some(tx);
    }

170
171
172
173
174
175
176
177
178
179
180
    /// 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
        }
    }

181
182
183
184
185
    /// Common watch logic with optional namespace filtering.
    ///
    /// Takes `Arc<Self>` so that each `handle_put` call can be spawned into its own
    /// tokio task, preventing a slow HuggingFace config download for one model from
    /// blocking discovery events for all subsequent models.
186
    pub async fn watch(
187
        self: Arc<Self>,
188
        mut discovery_stream: DiscoveryStream,
189
        namespace_filter: NamespaceFilter,
190
191
192
193
194
195
196
197
198
199
    ) {
        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;
                }
            };

200
            match event {
201
                DiscoveryEvent::Added(instance) => {
202
203
                    // Extract ModelCardInstanceId and card from the discovery instance
                    let (mcid, mut card) = match &instance {
204
205
206
207
208
                        DiscoveryInstance::Model {
                            namespace,
                            component,
                            endpoint,
                            instance_id,
209
                            model_suffix,
210
211
                            ..
                        } => {
212
                            let mcid = ModelCardInstanceId {
213
214
                                namespace: namespace.clone(),
                                component: component.clone(),
215
216
217
                                endpoint: endpoint.clone(),
                                instance_id: *instance_id,
                                model_suffix: model_suffix.clone(),
218
219
220
                            };

                            match instance.deserialize_model::<ModelDeploymentCard>() {
221
                                Ok(card) => (mcid, card),
222
223
224
225
226
227
228
229
230
231
                                Err(err) => {
                                    tracing::error!(%err, instance_id, "Failed to deserialize model card");
                                    continue;
                                }
                            }
                        }
                        _ => {
                            tracing::error!(
                                "Unexpected discovery instance type (expected ModelCard)"
                            );
232
233
234
                            continue;
                        }
                    };
235

236
237
                    // Filter by namespace using the configured filter
                    if !namespace_filter.matches(&mcid.namespace) {
238
                        tracing::debug!(
239
                            model_namespace = mcid.namespace,
240
241
                            namespace_filter = ?namespace_filter,
                            "Skipping model due to namespace filter"
242
243
244
245
                        );
                        continue;
                    }

246
247
248
249
250
251
252
                    // 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())
                    {
253
254
                        tracing::error!(
                            model_name = card.name(),
255
                            namespace = mcid.namespace,
256
257
258
                            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."
259
260
261
262
263
264
265
266
267
268
                        );
                        // 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;
                    }
269

270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
                    // Spawn each handle_put into its own task so that a slow
                    // HuggingFace config download for one model cannot block
                    // discovery events for all subsequent models.
                    //
                    // The JoinHandle is stored in `pending_puts` so that a
                    // subsequent `handle_delete` for the same instance can
                    // await the in-flight put before attempting cleanup.
                    let instance_key = mcid.to_path();
                    let watcher = Arc::clone(&self);
                    let key_for_cleanup = instance_key.clone();
                    let handle = tokio::spawn(async move {
                        match watcher.handle_put(&mcid, &mut card).await {
                            Ok(()) => {
                                tracing::info!(
                                    model_name = card.name(),
                                    namespace = mcid.namespace,
                                    "added model"
                                );
                                watcher.notify_on_model.notify_waiters();
                            }
                            Err(err) => {
                                tracing::error!(
                                    model_name = card.name(),
                                    namespace = mcid.namespace,
                                    error = format!("{err:#}"),
                                    "Error adding model from discovery",
                                );
                            }
298
                        }
299
300
301
302
303
304
305
306
                        // Remove ourselves from pending_puts once complete
                        watcher.pending_puts.remove(&key_for_cleanup);
                    });
                    // If a duplicate Added event arrives while the first task is still
                    // in-flight, abort the old task to prevent its cleanup from
                    // removing the new task's handle from pending_puts.
                    if let Some((_, old_handle)) = self.pending_puts.remove(&instance_key) {
                        old_handle.abort();
307
                    }
308
                    self.pending_puts.insert(instance_key, handle);
309
                }
310
311
312
313
                DiscoveryEvent::Removed(id) => {
                    // Extract ModelCardInstanceId from the removal event
                    let model_card_instance_id = match &id {
                        DiscoveryInstanceId::Model(mcid) => mcid,
314
                        DiscoveryInstanceId::Endpoint(_) | DiscoveryInstanceId::EventChannel(_) => {
315
316
317
318
319
320
                            tracing::error!(
                                "Unexpected discovery instance type in removal (expected Model)"
                            );
                            continue;
                        }
                    };
321

322
                    match self
323
                        .handle_delete(model_card_instance_id, &namespace_filter)
324
325
326
327
328
329
330
331
332
333
334
                        .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");
                        }
335
                    }
336
                }
337
            }
Ryan Olson's avatar
Ryan Olson committed
338
339
340
        }
    }

341
342
    /// 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.
343
344
    async fn handle_delete(
        &self,
345
        mcid: &ModelCardInstanceId,
346
        namespace_filter: &NamespaceFilter,
347
    ) -> anyhow::Result<Option<String>> {
348
        let key = mcid.to_path();
349
350
351
352
353
354
355
356
357
358
359

        // If there is an in-flight handle_put for this instance, wait for it
        // to complete before we attempt cleanup. Without this, a Removed event
        // arriving while handle_put is still downloading HF config would fail
        // to find the model card, leaving a stale registration.
        if let Some((_, handle)) = self.pending_puts.remove(&key) {
            tracing::debug!(key = %key, "awaiting in-flight handle_put before delete");
            // Ignore join errors (panic in the spawned task) — we still proceed
            // with cleanup since the put may have partially registered the model.
            let _ = handle.await;
        }
360
        let card = match self.manager.remove_model_card(&key) {
361
            Some(card) => card,
362
            None => {
363
                anyhow::bail!("Missing ModelDeploymentCard for {}", key);
364
365
            }
        };
366
        let model_name = card.name().to_string();
367
368
369
370
371
        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
372
        let active_instances = self
373
            .cards_for_model_with_endpoints(&model_name, namespace_filter)
374
375
            .await
            .with_context(|| model_name.clone())?;
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396

        // 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)"
                );
            }
397
398
399
400
401
402
403
404
405

            // 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);
            }
406
407
408
        }

        // Check if the Model still has instances in any namespace
409
        if !active_instances.is_empty() {
410
411
412
            tracing::debug!(
                model_name,
                active_instance_count = active_instances.len(),
413
                "Model has other active instances in other namespaces"
414
            );
415
416
            return Ok(None);
        }
417

418
419
        // No instances remain anywhere — remove the entire Model
        let _ = self.manager.remove_model(&model_name);
420

421
        if let Some(tx) = &self.model_update_tx {
422
            for model_type in ALL_MODEL_TYPES {
423
424
                if card.model_type.intersects(*model_type)
                    && is_model_type_list_empty(&self.manager, *model_type)
425
                {
426
                    tx.send(ModelUpdate::Removed(card.clone())).await.ok();
427
428
429
                }
            }
        }
430

431
        Ok(Some(model_name))
432
    }
Ryan Olson's avatar
Ryan Olson committed
433

434
    // Handles a PUT event from store, this usually means adding a new model to the list of served
435
    // models.
436
437
    async fn handle_put(
        &self,
438
        mcid: &ModelCardInstanceId,
439
440
        card: &mut ModelDeploymentCard,
    ) -> anyhow::Result<()> {
441
442
443
444
445
        // 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);
446

447
448
449
        if let Some(model) = self.manager.get_model(&model_name)
            && model.has_worker_set(&ws_key)
        {
450
451
            self.manager
                .save_model_card(&mcid.to_path(), card.clone())?;
452
453
            tracing::debug!(
                model_name = card.name(),
454
455
                namespace = namespace,
                "Worker joined existing WorkerSet, skipping pipeline build"
456
457
458
459
            );
            return Ok(());
        }

460
461
462
463
464
        // 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())
465
466
467
468
469
470
471
472
473
474
            && !self
                .recover_concurrent_registration(
                    mcid,
                    card,
                    &model_name,
                    &namespace,
                    &ws_key,
                    &registration_key,
                )
                .await?
475
        {
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
491
492
493
494
495
496
497
498
499
500
501
502
503
504
505
506
507
508
509
510
511
512
513
514
515
516
517
            return Ok(());
        }

        // RAII guard ensures the registration key is removed even if
        // do_worker_set_registration panics, preventing permanent poisoning.
        let _guard = RegistrationGuard {
            set: &self.registering_worker_sets,
            key: registration_key,
        };

        self.do_worker_set_registration(mcid, card).await
    }

    /// Handle the case where another task is already building the pipeline for this
    /// (model, namespace, type). This is a recovery path — it waits for the in-flight
    /// registration to finish, then either joins the resulting WorkerSet or retries.
    ///
    /// Returns `true` if the caller should proceed with its own registration
    /// (i.e. the other task failed), `false` if the worker was handled (joined or rejected).
    async fn recover_concurrent_registration(
        &self,
        mcid: &ModelCardInstanceId,
        card: &mut ModelDeploymentCard,
        model_name: &str,
        namespace: &str,
        ws_key: &str,
        registration_key: &str,
    ) -> anyhow::Result<bool> {
        // Wait for the in-flight registration to complete so we can validate
        // the new worker's checksum. Without this, a concurrent worker with a
        // mismatched checksum could sneak past the early check in `watch`.
        let mut attempts = 0;
        while self.registering_worker_sets.contains(registration_key) && attempts < 300 {
            tokio::time::sleep(Duration::from_millis(100)).await;
            attempts += 1;
        }

        // Validate checksum against the registered model
        if let Some(model) = self.manager.get_model(model_name)
            && !model.is_checksum_compatible(ws_key, card.mdcsum())
        {
            tracing::error!(
518
                model_name = card.name(),
519
                namespace = namespace,
520
521
522
                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."
523
            );
524
            return Ok(false);
525
526
        }

527
528
529
530
531
532
533
534
535
536
537
538
539
540
541
542
543
        // If the first registration failed or timed out, no WorkerSet exists.
        // Fall through to do_worker_set_registration instead of becoming a ghost
        // worker (registered in cards but with no serving pipeline).
        if self
            .manager
            .get_model(model_name)
            .is_none_or(|m| !m.has_worker_set(ws_key))
        {
            tracing::warn!(
                model_name = card.name(),
                namespace = namespace,
                "Concurrent registration produced no WorkerSet, retrying"
            );
            self.registering_worker_sets
                .insert(registration_key.to_string());
            return Ok(true);
        }
544

545
546
547
548
549
550
551
552
        self.manager
            .save_model_card(&mcid.to_path(), card.clone())?;
        tracing::debug!(
            model_name = card.name(),
            namespace = namespace,
            "Worker joined existing WorkerSet, skipping pipeline build"
        );
        Ok(false)
553
554
    }

555
556
557
    /// Build a complete WorkerSet with all engines for this (model, namespace)
    /// and add it to the Model.
    async fn do_worker_set_registration(
558
559
560
561
562
563
564
565
566
567
568
569
        &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?;
570
571
572
573
574
575
        let instance_watcher = client.instance_avail_watcher();
        tracing::debug!(
            model_name = card.name(),
            namespace = mcid.namespace,
            "building worker set pipeline"
        );
576
577
578
        self.manager
            .save_model_card(&mcid.to_path(), card.clone())?;

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

583
        let checksum = card.mdcsum();
584
585
586
587
588
589
        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);
590
591
592

        if card.model_input == ModelInput::Tokens
            && (card.model_type.supports_chat() || card.model_type.supports_completions())
593
594
595
596
597
        {
            // 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.

598
            let endpoint = component.endpoint(&mcid.endpoint);
599
600
601
602
603
604
605
606
607
            // 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)
            {
608
609
                Some(
                    self.manager
610
611
612
                        .kv_chooser_for(
                            &endpoint,
                            card.kv_cache_block_size,
613
                            Some(self.router_config.kv_router_config.clone()),
614
                            self.prefill_load_estimator.clone(),
615
                            WORKER_TYPE_DECODE, // This is the decode router
616
                            Some(card.display_name.clone()),
617
                            card.runtime_config.enable_eagle,
618
                        )
619
620
621
622
623
                        .await?,
                )
            } else {
                None
            };
624

625
626
627
628
629
630
631
632
633
634
635
636
637
            // 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
            };
638

639
640
            // Create prefill chooser once if we're building pipelines
            // Both chat and completions will share the same prefill chooser instance
641
            let model_name = card.name().to_string();
642
643
            let prefill_chooser = self
                .manager
644
                .register_prefill_router(&model_name, &namespace)
645
646
                .map(|rx| {
                    // Create prefill-specific config with track_active_blocks disabled
647
                    let mut prefill_config = self.router_config.kv_router_config.clone();
648
649
650
651
652
                    prefill_config.router_track_active_blocks = false;

                    PrefillRouter::new(
                        rx,
                        self.manager.clone(),
653
                        self.router_config.router_mode,
654
655
                        card.kv_cache_block_size,
                        Some(prefill_config),
656
                        self.prefill_load_estimator.clone(),
657
                        self.router_config.enforce_disagg,
658
659
                        model_name.clone(),
                        namespace.clone(),
660
                        card.runtime_config.enable_eagle,
661
662
663
                    )
                });

664
665
666
            // 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,
667
            // worker TTFT/ITL cleanup). The thresholds control busy detection behavior only.
668
669
670
671
672
673
674
675
676
677
            //
            // 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());
678
            let worker_monitor = Some(KvWorkerMonitor::new(
679
                monitor_client,
680
681
                self.router_config.load_threshold_config.clone(),
            ));
682

683
684
685
            // 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.
686
687
            worker_set.kv_router = kv_chooser.clone();
            worker_set.worker_monitor = worker_monitor.clone();
688
            worker_set.prefill_router = prefill_chooser.clone();
689

690
            // Add chat engine only if the model supports chat
691
            if card.model_type.supports_chat() {
692
693
694
695
696
697
698
699
700
701
702
                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
703
                } else {
704
705
706
707
708
709
710
                    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)."
                        )
                    })?;
711
712
713
714
715
716
                    entrypoint::build_routed_pipeline::<
                        NvCreateChatCompletionRequest,
                        NvCreateChatCompletionStreamResponse,
                    >(
                        card,
                        &client,
717
                        self.manager.clone(),
718
719
720
                        self.router_config.router_mode,
                        worker_monitor.clone(),
                        kv_chooser.clone(),
721
                        tk,
722
                        prefill_chooser.clone(),
723
                        self.router_config.enforce_disagg,
724
                        self.migration_limit,
725
                        self.migration_max_seq_len,
726
                        self.metrics.clone(),
727
728
729
730
                    )
                    .await
                    .context("build_routed_pipeline")?
                };
731
                worker_set.chat_engine = Some(chat_engine);
732
                tracing::info!("Chat completions is ready");
733
            }
734

735
736
            // Add completions engine only if the model supports completions
            // and we have a tokenizer (completions always uses the Rust preprocessor).
737
            if card.model_type.supports_completions() {
738
739
740
741
742
743
744
745
746
747
748
749
750
751
752
753
754
755
756
757
758
                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,
759
                        self.migration_max_seq_len,
760
761
762
763
764
765
766
767
768
769
770
771
772
773
774
775
776
777
778
779
780
781
782
783
                        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()
                );
784
            }
785
786
787
788
789
790
        } 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(
791
                client, self.router_config.router_mode, None, None
792
793
            )
            .await?;
794
            worker_set.embeddings_engine = Some(Arc::new(push_router));
795
796
797
798
799
800
801
802
803
804
        }
        // 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())
        {
805
            // Image/Audio/Video models can also support chat completions (vLLM omni way)
806
807
808
809
810
811
812
813
814
815
816
            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?;
817
                worker_set.chat_engine = Some(Arc::new(chat_router));
818
819
820
821
822
823
824
825
826
827
            }

            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?;
828
                worker_set.images_engine = Some(Arc::new(images_router));
829
830
831
832
833
834
835
836
837
838
            }

            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?;
839
                worker_set.videos_engine = Some(Arc::new(videos_router));
840
841
            }

842
843
844
845
846
847
848
849
850
851
852
853
854
            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));
            }
855
        } else if card.model_input == ModelInput::Text && card.model_type.supports_chat() {
856
            // Case: Text + Chat (pure text-to-text, no diffusion)
857
858
859
860
861
862
863
            let push_router = PushRouter::<
                NvCreateChatCompletionRequest,
                Annotated<NvCreateChatCompletionStreamResponse>,
            >::from_client_with_threshold(
                client, self.router_config.router_mode, None, None
            )
            .await?;
864
            worker_set.chat_engine = Some(Arc::new(push_router));
865
        } else if card.model_input == ModelInput::Text && card.model_type.supports_completions() {
866
            // Case: Text + Completions
867
868
869
870
            let push_router = PushRouter::<
                NvCreateCompletionRequest,
                Annotated<NvCreateCompletionResponse>,
            >::from_client_with_threshold(
871
                client, self.router_config.router_mode, None, None
872
873
            )
            .await?;
874
            worker_set.completions_engine = Some(Arc::new(push_router));
875
        } else if card.model_input == ModelInput::Tokens && card.model_type.supports_embedding() {
876
877
878
879
880
881
882
            // Case 4: Tokens + Embeddings
            // Create preprocessing pipeline similar to Backend
            let frontend = SegmentSource::<
                SingleIn<NvCreateEmbeddingRequest>,
                ManyOut<Annotated<NvCreateEmbeddingResponse>>,
            >::new();

883
            let preprocessor = OpenAIPreprocessor::new(card.clone())?.into_operator();
884
            let backend = Backend::from_mdc(card).into_operator();
885
886
887
888
889

            let router = PushRouter::<
                PreprocessedEmbeddingRequest,
                Annotated<EmbeddingsEngineOutput>,
            >::from_client_with_threshold(
890
                client, self.router_config.router_mode, None, None
891
892
893
            )
            .await?;

894
            // Note: Embeddings don't need KV routing complexity or load monitoring
895
896
897
898
899
900
901
902
903
904
905
            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)?;

906
            worker_set.embeddings_engine = Some(embedding_engine);
907
        } else if card.model_input == ModelInput::Tensor && card.model_type.supports_tensor() {
908
            // Case 6: Tensor + TensorBased (non-LLM)
909
            // No KV cache concepts - not an LLM model
910
911
912
913
            let push_router = PushRouter::<
                NvCreateTensorRequest,
                Annotated<NvCreateTensorResponse>,
            >::from_client_with_threshold(
914
                client, self.router_config.router_mode, None, None
915
916
            )
            .await?;
917
            worker_set.tensor_engine = Some(Arc::new(push_router));
918
919
920
921
922
923
924
925
926
927
928
929
        } 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(),
930
931
932
                "Prefill model detected, registering and activating prefill router"
            );

933
934
            // Prefill sets have no engines — we add the WorkerSet first for tracking,
            // then activate the prefill router.
935
            self.manager
936
                .add_worker_set(card.name(), &ws_key, worker_set);
937

938
939
940
941
942
943
944
            // 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 {
945
946
947
948
949
950
951
952
953
954
                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"
955
            );
956
957

            return Ok(());
958
959
960
961
        } else {
            // Reject unsupported combinations
            anyhow::bail!(
                "Unsupported model configuration: {} with {} input. Supported combinations: \
962
                Tokens+(Chat|Completions|Prefill), Text+(Chat|Completions|Images), Tokens+Embeddings, Tensor+TensorBased",
963
964
                card.model_type,
                card.model_input.as_str()
965
            );
966
        }
Ryan Olson's avatar
Ryan Olson committed
967

968
969
        // Add the completed WorkerSet to the Model
        self.manager
970
            .add_worker_set(card.name(), &ws_key, worker_set);
971

972
973
        Ok(())
    }
974

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

980
981
982
        let mut results = Vec::with_capacity(instances.len());
        for instance in instances {
            match instance.deserialize_model::<ModelDeploymentCard>() {
983
                Ok(card) => {
984
985
986
987
988
989
990
991
992
993
994
995
996
997
998
                    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)"
                            );
999
1000
1001
                            continue;
                        }
                    };
1002
                    results.push((endpoint_id, card));
1003
                }
1004
                Err(err) => {
1005
                    tracing::error!(%err, "Failed to deserialize model card");
1006
1007
                    continue;
                }
1008
            }
1009
        }
1010
        Ok(results)
1011
1012
    }

1013
    pub async fn cards_for_model(
1014
1015
        &self,
        model_name: &str,
1016
        namespace_filter: &NamespaceFilter,
1017
    ) -> anyhow::Result<Vec<ModelDeploymentCard>> {
1018
1019
1020
1021
1022
1023
1024
1025
1026
1027
1028
1029
1030
1031
1032
        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)>> {
1033
1034
1035
        let mut all = self.all_cards().await?;
        all.retain(|(endpoint_id, card)| {
            let matches_name = card.name() == model_name;
1036
            let matches_namespace = namespace_filter.matches(&endpoint_id.namespace);
1037
1038
            matches_name && matches_namespace
        });
1039
1040
1041
1042
1043
1044
1045
1046
1047
1048
1049
1050
1051
1052
1053
1054
1055
1056
1057
1058
1059
1060
1061
1062
1063
        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));
1064
        assert!(is_model_type_list_empty(&mm, ModelType::Audios));
1065
1066
1067
1068
1069
1070
1071
1072
1073
        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
1074
        mm.add_worker_set("model-a", "ns1", make_worker_set("ns1"));
1075
1076
1077
1078
1079
1080
1081

        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));
1082
        assert!(is_model_type_list_empty(&mm, ModelType::Audios));
1083
1084
1085
1086
1087
1088
1089
        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();
1090
        mm.add_worker_set("model-a", "ns1", make_worker_set("ns1"));
1091
1092
1093
1094
1095
1096
1097
1098
1099
        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();
1100
1101
        mm.add_worker_set("model-a", "ns1", make_worker_set("ns1"));
        mm.add_worker_set("model-b", "ns1", make_worker_set("ns1"));
1102
1103
1104
1105
1106
1107
1108
1109

        // 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));
1110
1111
    }
}