watcher.rs 47.3 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
111
112
113
114
115
116
117
118
119
120
121
/// 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
    }
}

122
123
124
125
126
127
128
129
130
131
132
133
134
135
/// 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);
    }
}

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

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

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

179
180
181
182
183
    /// 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.
184
    pub async fn watch(
185
        self: Arc<Self>,
186
        mut discovery_stream: DiscoveryStream,
187
        namespace_filter: NamespaceFilter,
188
189
190
191
192
193
194
195
196
197
    ) {
        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;
                }
            };

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

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

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

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

268
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
                    // 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",
                                );
                            }
296
                        }
297
298
299
300
301
302
303
304
                        // 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();
305
                    }
306
                    self.pending_puts.insert(instance_key, handle);
307
                }
308
309
310
311
                DiscoveryEvent::Removed(id) => {
                    // Extract ModelCardInstanceId from the removal event
                    let model_card_instance_id = match &id {
                        DiscoveryInstanceId::Model(mcid) => mcid,
312
                        DiscoveryInstanceId::Endpoint(_) | DiscoveryInstanceId::EventChannel(_) => {
313
314
315
316
317
318
                            tracing::error!(
                                "Unexpected discovery instance type in removal (expected Model)"
                            );
                            continue;
                        }
                    };
319

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

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

        // 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;
        }
358
        let card = match self.manager.remove_model_card(&key) {
359
            Some(card) => card,
360
            None => {
361
                anyhow::bail!("Missing ModelDeploymentCard for {}", key);
362
363
            }
        };
364
        let model_name = card.name().to_string();
365
366
367
368
369
        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
370
        let active_instances = self
371
            .cards_for_model_with_endpoints(&model_name, namespace_filter)
372
373
            .await
            .with_context(|| model_name.clone())?;
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394

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

            // 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);
            }
404
405
406
        }

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

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

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

429
        Ok(Some(model_name))
430
    }
Ryan Olson's avatar
Ryan Olson committed
431

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

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

458
459
460
461
462
        // 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())
463
464
465
466
467
468
469
470
471
472
            && !self
                .recover_concurrent_registration(
                    mcid,
                    card,
                    &model_name,
                    &namespace,
                    &ws_key,
                    &registration_key,
                )
                .await?
473
        {
474
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
            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!(
516
                model_name = card.name(),
517
                namespace = namespace,
518
519
520
                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."
521
            );
522
            return Ok(false);
523
524
        }

525
526
527
528
529
530
531
532
533
534
535
536
537
538
539
540
541
        // 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);
        }
542

543
544
545
546
547
548
549
550
        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)
551
552
    }

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

970
971
        Ok(())
    }
972

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

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

1011
    pub async fn cards_for_model(
1012
1013
        &self,
        model_name: &str,
1014
        namespace_filter: &NamespaceFilter,
1015
    ) -> anyhow::Result<Vec<ModelDeploymentCard>> {
1016
1017
1018
1019
1020
1021
1022
1023
1024
1025
1026
1027
1028
1029
1030
        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)>> {
1031
1032
1033
        let mut all = self.all_cards().await?;
        all.retain(|(endpoint_id, card)| {
            let matches_name = card.name() == model_name;
1034
            let matches_namespace = namespace_filter.matches(&endpoint_id.namespace);
1035
1036
            matches_name && matches_namespace
        });
1037
1038
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
1064
1065
1066
1067
1068
1069
1070
        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
1071
        mm.add_worker_set("model-a", "ns1", make_worker_set("ns1"));
1072
1073
1074
1075
1076
1077
1078
1079
1080
1081
1082
1083
1084
1085

        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();
1086
        mm.add_worker_set("model-a", "ns1", make_worker_set("ns1"));
1087
1088
1089
1090
1091
1092
1093
1094
1095
        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();
1096
1097
        mm.add_worker_set("model-a", "ns1", make_worker_set("ns1"));
        mm.add_worker_set("model-b", "ns1", make_worker_set("ns1"));
1098
1099
1100
1101
1102
1103
1104
1105

        // 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));
1106
1107
    }
}