model_manager.rs 22.8 KB
Newer Older
1
// SPDX-FileCopyrightText: Copyright (c) 2024-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
2
3
// SPDX-License-Identifier: Apache-2.0

4
5
use std::{
    collections::{HashMap, HashSet},
6
    sync::Arc,
7
8
};

9
use dashmap::{DashMap, mapref::entry::Entry};
10
use parking_lot::RwLock;
11
use tokio::sync::oneshot;
12

13
14
use super::worker_monitor::LoadThresholdConfig;
use super::{KvWorkerMonitor, RuntimeConfigs};
15

16
use dynamo_runtime::{
17
    component::{Client, Endpoint, build_transport_type},
18
    discovery::DiscoverySpec,
19
20
21
    prelude::DistributedRuntimeProvider,
    protocols::EndpointId,
};
22
23

use crate::{
24
25
26
27
    kv_router::{
        KvRouter, KvRouterConfig, protocols::WorkerId, router_endpoint_id,
        scheduler::DefaultWorkerSelector,
    },
28
    local_model::runtime_config::DisaggregatedEndpoint,
29
    model_card::ModelDeploymentCard,
30
    model_type::ModelType,
31
32
33
34
35
36
37
38
    types::{
        generic::tensor::TensorStreamingEngine,
        openai::{
            chat_completions::OpenAIChatCompletionsStreamingEngine,
            completions::OpenAICompletionsStreamingEngine,
            embeddings::OpenAIEmbeddingsStreamingEngine,
        },
    },
39
};
40

41
42
43
44
45
46
47
48
/// State for prefill router activation rendezvous
enum PrefillActivationState {
    /// Decode model registered, waiting for prefill endpoint
    DecodeWaiting(oneshot::Sender<Endpoint>),
    /// Prefill endpoint arrived, waiting for decode model to register
    PrefillReady(oneshot::Receiver<Endpoint>),
}

49
50
51
52
53
54
55
56
57
#[derive(Debug, thiserror::Error)]
pub enum ModelManagerError {
    #[error("Model not found: {0}")]
    ModelNotFound(String),

    #[error("Model already exists: {0}")]
    ModelAlreadyExists(String),
}

58
59
60
61
62
63
/// Central manager for model engines, routing, and configuration.
///
/// Manages model lifecycle including engines, KV routers, prefill coordination,
/// and per-model busy thresholds for load-based request rejection.
///
/// Note: Don't implement Clone for this, put it in an Arc instead.
64
65
66
67
68
pub struct ModelManager {
    // We read a lot and write rarely, so these three are RwLock
    completion_engines: RwLock<ModelEngines<OpenAICompletionsStreamingEngine>>,
    chat_completion_engines: RwLock<ModelEngines<OpenAIChatCompletionsStreamingEngine>>,
    embeddings_engines: RwLock<ModelEngines<OpenAIEmbeddingsStreamingEngine>>,
69
    tensor_engines: RwLock<ModelEngines<TensorStreamingEngine>>,
70
71
    // Prefill models don't have engines - they're only tracked for discovery/lifecycle
    prefill_engines: RwLock<ModelEngines<()>>,
72

73
74
75
    cards: DashMap<String, ModelDeploymentCard>,
    kv_choosers: DashMap<EndpointId, Arc<KvRouter>>,
    prefill_router_activators: DashMap<String, PrefillActivationState>,
76

77
78
    // Per-model monitoring: worker_monitors for load-based rejection, runtime_configs for KvScheduler
    worker_monitors: DashMap<String, KvWorkerMonitor>,
79
    runtime_configs: DashMap<EndpointId, Arc<RuntimeConfigs>>,
80
81
82
83
84
85
86
87
88
89
90
91
92
93
}

impl Default for ModelManager {
    fn default() -> Self {
        Self::new()
    }
}

impl ModelManager {
    pub fn new() -> Self {
        Self {
            completion_engines: RwLock::new(ModelEngines::default()),
            chat_completion_engines: RwLock::new(ModelEngines::default()),
            embeddings_engines: RwLock::new(ModelEngines::default()),
94
            tensor_engines: RwLock::new(ModelEngines::default()),
95
            prefill_engines: RwLock::new(ModelEngines::default()),
96
97
98
            cards: DashMap::new(),
            kv_choosers: DashMap::new(),
            prefill_router_activators: DashMap::new(),
99
            worker_monitors: DashMap::new(),
100
            runtime_configs: DashMap::new(),
101
102
103
        }
    }

104
105
106
107
108
109
110
111
112
113
114
115
116
    pub fn is_valid_checksum(
        &self,
        model_type: ModelType,
        model_name: &str,
        candidate_checksum: &str,
    ) -> Option<bool> {
        let mut results = vec![];
        for unit in model_type.units() {
            let maybe_valid_checksum = match unit {
                ModelType::Chat => self.chat_completion_engines.read().checksum(model_name),
                ModelType::Completions => self.completion_engines.read().checksum(model_name),
                ModelType::Embedding => self.embeddings_engines.read().checksum(model_name),
                ModelType::TensorBased => self.tensor_engines.read().checksum(model_name),
117
                ModelType::Prefill => self.prefill_engines.read().checksum(model_name),
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
                _ => {
                    continue;
                }
            };
            if let Some(is_valid) = maybe_valid_checksum.map(|valid_checksum| {
                tracing::debug!(
                    model_name,
                    valid_checksum,
                    candidate_checksum,
                    "is_valid_checksum: check case"
                );
                valid_checksum == candidate_checksum
            }) {
                results.push(is_valid)
            }
        }
        if results.is_empty() {
            None
        } else {
            // The checksum is valid if it is correct for all the ModelType in the bitflag.
            Some(results.into_iter().all(|x| x))
        }
    }

142
    pub fn get_model_cards(&self) -> Vec<ModelDeploymentCard> {
143
        self.cards.iter().map(|r| r.value().clone()).collect()
144
145
    }

146
147
    /// Check if a decode model (chat or completions) is registered
    pub fn has_decode_model(&self, model: &str) -> bool {
148
149
        self.chat_completion_engines.read().contains(model)
            || self.completion_engines.read().contains(model)
150
151
152
153
154
155
156
157
158
159
160
    }

    /// Check if a prefill model is registered
    pub fn has_prefill_model(&self, model: &str) -> bool {
        self.prefill_engines.read().contains(model)
    }

    /// Check if any model (decode or prefill) is registered.
    /// Note: For registration skip-checks, use has_decode_model() or has_prefill_model() instead.
    pub fn has_model_any(&self, model: &str) -> bool {
        self.has_decode_model(model) || self.has_prefill_model(model)
161
162
    }

163
164
165
166
167
    pub fn model_display_names(&self) -> HashSet<String> {
        self.list_chat_completions_models()
            .into_iter()
            .chain(self.list_completions_models())
            .chain(self.list_embeddings_models())
168
            .chain(self.list_tensor_models())
169
            .chain(self.list_prefill_models())
170
171
172
            .collect()
    }

173
    pub fn list_chat_completions_models(&self) -> Vec<String> {
174
        self.chat_completion_engines.read().list()
175
176
177
    }

    pub fn list_completions_models(&self) -> Vec<String> {
178
        self.completion_engines.read().list()
179
180
181
    }

    pub fn list_embeddings_models(&self) -> Vec<String> {
182
        self.embeddings_engines.read().list()
183
184
    }

185
186
187
188
    pub fn list_tensor_models(&self) -> Vec<String> {
        self.tensor_engines.read().list()
    }

189
190
191
192
    pub fn list_prefill_models(&self) -> Vec<String> {
        self.prefill_engines.read().list()
    }

193
194
195
    pub fn add_completions_model(
        &self,
        model: &str,
196
        card_checksum: &str,
197
198
        engine: OpenAICompletionsStreamingEngine,
    ) -> Result<(), ModelManagerError> {
199
        let mut clients = self.completion_engines.write();
200
        clients.add(model, card_checksum, engine)
201
202
203
204
205
    }

    pub fn add_chat_completions_model(
        &self,
        model: &str,
206
        card_checksum: &str,
207
208
        engine: OpenAIChatCompletionsStreamingEngine,
    ) -> Result<(), ModelManagerError> {
209
        let mut clients = self.chat_completion_engines.write();
210
        clients.add(model, card_checksum, engine)
211
212
213
214
215
    }

    pub fn add_embeddings_model(
        &self,
        model: &str,
216
        card_checksum: &str,
217
218
        engine: OpenAIEmbeddingsStreamingEngine,
    ) -> Result<(), ModelManagerError> {
219
        let mut clients = self.embeddings_engines.write();
220
        clients.add(model, card_checksum, engine)
221
222
    }

223
224
225
    pub fn add_tensor_model(
        &self,
        model: &str,
226
        card_checksum: &str,
227
228
229
        engine: TensorStreamingEngine,
    ) -> Result<(), ModelManagerError> {
        let mut clients = self.tensor_engines.write();
230
        clients.add(model, card_checksum, engine)
231
232
    }

233
234
235
236
237
238
    pub fn add_prefill_model(
        &self,
        model: &str,
        card_checksum: &str,
    ) -> Result<(), ModelManagerError> {
        let mut clients = self.prefill_engines.write();
239
        clients.add(model, card_checksum, ())
240
241
    }

242
    pub fn remove_completions_model(&self, model: &str) -> Result<(), ModelManagerError> {
243
        let mut clients = self.completion_engines.write();
244
245
246
247
        clients.remove(model)
    }

    pub fn remove_chat_completions_model(&self, model: &str) -> Result<(), ModelManagerError> {
248
        let mut clients = self.chat_completion_engines.write();
249
250
251
252
        clients.remove(model)
    }

    pub fn remove_embeddings_model(&self, model: &str) -> Result<(), ModelManagerError> {
253
        let mut clients = self.embeddings_engines.write();
254
255
256
        clients.remove(model)
    }

257
258
259
260
261
    pub fn remove_tensor_model(&self, model: &str) -> Result<(), ModelManagerError> {
        let mut clients = self.tensor_engines.write();
        clients.remove(model)
    }

262
263
264
265
266
    pub fn remove_prefill_model(&self, model: &str) -> Result<(), ModelManagerError> {
        let mut clients = self.prefill_engines.write();
        clients.remove(model)
    }

267
    pub fn get_embeddings_engine(
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
296
297
298
299
        &self,
        model: &str,
    ) -> Result<OpenAIEmbeddingsStreamingEngine, ModelManagerError> {
        self.embeddings_engines
            .read()
            .get(model)
            .cloned()
            .ok_or(ModelManagerError::ModelNotFound(model.to_string()))
    }

    pub fn get_completions_engine(
        &self,
        model: &str,
    ) -> Result<OpenAICompletionsStreamingEngine, ModelManagerError> {
        self.completion_engines
            .read()
            .get(model)
            .cloned()
            .ok_or(ModelManagerError::ModelNotFound(model.to_string()))
    }

    pub fn get_chat_completions_engine(
        &self,
        model: &str,
    ) -> Result<OpenAIChatCompletionsStreamingEngine, ModelManagerError> {
        self.chat_completion_engines
            .read()
            .get(model)
            .cloned()
            .ok_or(ModelManagerError::ModelNotFound(model.to_string()))
    }

300
301
302
303
304
305
306
307
308
309
310
    pub fn get_tensor_engine(
        &self,
        model: &str,
    ) -> Result<TensorStreamingEngine, ModelManagerError> {
        self.tensor_engines
            .read()
            .get(model)
            .cloned()
            .ok_or(ModelManagerError::ModelNotFound(model.to_string()))
    }

311
    /// Save a ModelDeploymentCard from an instance's key so we can fetch it later when the key is
312
313
    /// deleted.
    pub fn save_model_card(&self, key: &str, card: ModelDeploymentCard) -> anyhow::Result<()> {
314
        self.cards.insert(key.to_string(), card);
315
        Ok(())
316
317
    }

318
    /// Remove and return model card for this instance's key. We do this when the instance stops.
319
    pub fn remove_model_card(&self, key: &str) -> Option<ModelDeploymentCard> {
320
        self.cards.remove(key).map(|(_, v)| v)
321
322
323
324
    }

    pub async fn kv_chooser_for(
        &self,
325
        endpoint: &Endpoint,
326
        kv_cache_block_size: u32,
327
        kv_router_config: Option<KvRouterConfig>,
328
    ) -> anyhow::Result<Arc<KvRouter>> {
329
        let endpoint_id = endpoint.id();
330

331
        if let Some(kv_chooser) = self.get_kv_chooser(&endpoint_id) {
332
333
334
            // Check if the existing router has a different block size
            if kv_chooser.block_size() != kv_cache_block_size {
                tracing::warn!(
335
                    endpoint = %endpoint_id,
336
337
                    existing_block_size = %kv_chooser.block_size(),
                    requested_block_size = %kv_cache_block_size,
338
                    "KV Router block size mismatch! Endpoint is requesting a different kv_cache_block_size than the existing router. \
339
340
341
                     This will cause routing to fail silently. Consider using the same block size or restarting the router."
                );
            }
342
343
344
            return Ok(kv_chooser);
        }

345
        let client = endpoint.client().await?;
346
347
348
349
350

        // Register router via discovery mechanism
        let discovery = endpoint.component().drt().discovery();
        let instance_id = discovery.instance_id();

351
        // Build transport for router endpoint based on request plane mode
352
353
        // Use KV_ROUTER_COMPONENT as the component name to distinguish from the generate endpoint's component
        let router_endpoint_id = router_endpoint_id(endpoint.id().namespace);
354
        let transport = build_transport_type(endpoint, &router_endpoint_id, instance_id).await?;
355
356
357
358
359

        let discovery_spec = DiscoverySpec::Endpoint {
            namespace: router_endpoint_id.namespace.clone(),
            component: router_endpoint_id.component.clone(),
            endpoint: router_endpoint_id.name.clone(),
360
            transport,
361
362
363
364
        };

        discovery.register(discovery_spec).await?;

365
366
367
        // Get or create runtime config watcher for this endpoint
        let workers_with_configs = self.get_or_create_runtime_config_watcher(endpoint).await?;

368
        let selector = Box::new(DefaultWorkerSelector::new(kv_router_config));
369
        let chooser = KvRouter::new(
370
371
            endpoint.clone(),
            client,
372
            workers_with_configs,
373
374
            kv_cache_block_size,
            Some(selector),
375
            kv_router_config,
376
            instance_id,
377
378
        )
        .await?;
379
        let new_kv_chooser = Arc::new(chooser);
380
        self.kv_choosers.insert(endpoint_id, new_kv_chooser.clone());
381
382
        Ok(new_kv_chooser)
    }
383

384
    fn get_kv_chooser(&self, id: &EndpointId) -> Option<Arc<KvRouter>> {
385
        self.kv_choosers.get(id).map(|r| r.value().clone())
386
387
388
389
390
391
392
393
394
    }

    /// Register a prefill router for a decode model. Returns a receiver that will be
    /// activated when the corresponding prefill model is discovered.
    /// Returns None if the decode model was already registered.
    pub fn register_prefill_router(
        &self,
        model_name: String,
    ) -> Option<oneshot::Receiver<Endpoint>> {
395
396
        match self.prefill_router_activators.remove(&model_name) {
            Some((_, PrefillActivationState::PrefillReady(rx))) => {
397
398
399
400
401
402
403
                // Prefill endpoint already arrived - rx will immediately resolve
                tracing::debug!(
                    model_name = %model_name,
                    "Prefill endpoint already available, returning receiver with endpoint"
                );
                Some(rx)
            }
404
            Some((key, PrefillActivationState::DecodeWaiting(tx))) => {
405
406
407
408
409
                // Decode already registered - this shouldn't happen, restore state and return None
                tracing::error!(
                    model_name = %model_name,
                    "Decode model already registered for this prefill router"
                );
410
411
                self.prefill_router_activators
                    .insert(key, PrefillActivationState::DecodeWaiting(tx));
412
413
414
415
416
                None
            }
            None => {
                // New registration: create tx/rx pair, store sender and return receiver
                let (tx, rx) = oneshot::channel();
417
                self.prefill_router_activators.insert(
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
                    model_name.clone(),
                    PrefillActivationState::DecodeWaiting(tx),
                );
                tracing::debug!(
                    model_name = %model_name,
                    "No prefill endpoint available yet, storing sender for future activation"
                );
                Some(rx)
            }
        }
    }

    /// Activate a prefill router by sending the endpoint through the oneshot channel.
    /// If no decode model has registered yet, stores the endpoint for future retrieval.
    pub fn activate_prefill_router(
        &self,
        model_name: &str,
        endpoint: Endpoint,
    ) -> anyhow::Result<()> {
437
438
        match self.prefill_router_activators.remove(model_name) {
            Some((_, PrefillActivationState::DecodeWaiting(sender))) => {
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
                // Decode model already registered
                sender.send(endpoint).map_err(|_| {
                    anyhow::anyhow!(
                        "Failed to send endpoint to prefill router activator for model: {}",
                        model_name
                    )
                })?;

                tracing::info!(
                    model_name = %model_name,
                    "Activated prefill router for already-registered decode model"
                );

                Ok(())
            }
454
            Some((_, PrefillActivationState::PrefillReady(_))) => {
455
456
457
458
459
460
461
462
463
464
465
466
                // Prefill already activated - this shouldn't happen
                anyhow::bail!("Prefill router for model {} already activated", model_name);
            }
            None => {
                // Decode model not registered yet - create pair and immediately send endpoint
                let (tx, rx) = oneshot::channel();

                tx.send(endpoint).map_err(|_| {
                    anyhow::anyhow!("Failed to send endpoint for prefill model: {}", model_name)
                })?;

                // Store the receiver for when decode model registers
467
                self.prefill_router_activators.insert(
468
469
470
471
472
473
474
475
476
477
478
479
                    model_name.to_string(),
                    PrefillActivationState::PrefillReady(rx),
                );

                tracing::info!(
                    model_name = %model_name,
                    "Stored prefill endpoint for future decode model registration"
                );

                Ok(())
            }
        }
480
481
    }

482
    pub fn get_model_tool_call_parser(&self, model: &str) -> Option<String> {
483
        self.cards
484
485
486
            .iter()
            .find(|r| r.value().display_name == model)
            .and_then(|r| r.value().runtime_config.tool_call_parser.clone())
487
    }
488
489
490
491
492
493
494
495
496

    /// Creates parsing options with tool call parser and reasoning parser for the specified model.
    /// Currently reasoning parser is not implemented (returns None).
    pub fn get_parsing_options(&self, model: &str) -> crate::protocols::openai::ParsingOptions {
        let tool_call_parser = self.get_model_tool_call_parser(model);
        let reasoning_parser = None; // TODO: Implement reasoning parser

        crate::protocols::openai::ParsingOptions::new(tool_call_parser, reasoning_parser)
    }
497

498
499
500
    /// Gets or sets the load threshold config for a model's worker monitor.
    /// Pass `Some(config)` to update, `None` to get. Returns `None` if no monitor exists.
    pub fn load_threshold_config(
501
502
        &self,
        model: &str,
503
504
505
506
507
        config: Option<&LoadThresholdConfig>,
    ) -> Option<LoadThresholdConfig> {
        let monitor = self.worker_monitors.get(model)?;
        if let Some(cfg) = config {
            monitor.set_load_threshold_config(cfg);
508
        }
509
        Some(monitor.load_threshold_config())
510
511
    }

512
    /// Gets or creates a worker monitor for a model. Updates thresholds if monitor exists.
513
514
515
    pub fn get_or_create_worker_monitor(
        &self,
        model: &str,
516
        client: Client,
517
        config: LoadThresholdConfig,
518
    ) -> KvWorkerMonitor {
519
520
521
        if let Some(existing) = self.worker_monitors.get(model) {
            existing.set_load_threshold_config(&config);
            return existing.clone();
522
        }
523
524
525
526
        let monitor = KvWorkerMonitor::new(client, config);
        self.worker_monitors
            .insert(model.to_string(), monitor.clone());
        monitor
527
528
    }

529
    /// Get or create a runtime config watcher for an endpoint.
530
531
    /// Spawns a background task to watch for worker config changes.
    /// Returns a shared RuntimeConfigs that KvScheduler can use directly.
532
533
534
    pub async fn get_or_create_runtime_config_watcher(
        &self,
        endpoint: &Endpoint,
535
    ) -> anyhow::Result<Arc<RuntimeConfigs>> {
536
537
538
539
540
541
542
543
        let endpoint_id = endpoint.id();

        // Fast path: return existing if present
        if let Some(existing) = self.runtime_configs.get(&endpoint_id) {
            return Ok(existing.clone());
        }

        // Atomic get-or-insert to avoid TOCTOU race
544
        let inner = Arc::new(RuntimeConfigs::new());
545
        let (result, is_new) = match self.runtime_configs.entry(endpoint_id) {
546
547
            Entry::Occupied(e) => (e.get().clone(), false),
            Entry::Vacant(e) => {
548
549
                e.insert(inner.clone());
                (inner, true)
550
551
552
553
554
            }
        };

        // Only spawn watcher if we were the one who inserted
        if is_new {
555
            result.start_watcher(endpoint).await?;
556
557
        }

558
        Ok(result)
559
560
561
562
563
564
565
566
567
    }

    /// Get disaggregated endpoint for a specific worker.
    /// Used by PrefillRouter for bootstrap info - works for ANY routing mode.
    pub fn get_disaggregated_endpoint(
        &self,
        endpoint_id: &EndpointId,
        worker_id: WorkerId,
    ) -> Option<DisaggregatedEndpoint> {
568
569
        let inner = self.runtime_configs.get(endpoint_id)?;
        let config_ref = inner.configs.get(&worker_id)?;
570
571
572
        config_ref.as_ref()?.disaggregated_endpoint.clone()
    }

573
574
    /// Lists all models with worker monitors configured.
    pub fn list_busy_thresholds(&self) -> Vec<(String, LoadThresholdConfig)> {
575
576
        self.worker_monitors
            .iter()
577
            .map(|entry| (entry.key().clone(), entry.value().load_threshold_config()))
578
579
            .collect()
    }
580
581
582
583
584
585
}

pub struct ModelEngines<E> {
    /// Optional default model name
    default: Option<String>,
    engines: HashMap<String, E>,
586
587
588
    /// Key: Model name, value: Checksum of the ModelDeploymentCard. New instances must have the
    /// same card.
    checksums: HashMap<String, String>,
589
590
591
592
593
594
595
}

impl<E> Default for ModelEngines<E> {
    fn default() -> Self {
        Self {
            default: None,
            engines: HashMap::new(),
596
            checksums: HashMap::new(),
597
598
599
600
601
602
603
604
605
606
607
608
609
610
611
        }
    }
}

impl<E> ModelEngines<E> {
    #[allow(dead_code)]
    fn set_default(&mut self, model: &str) {
        self.default = Some(model.to_string());
    }

    #[allow(dead_code)]
    fn clear_default(&mut self) {
        self.default = None;
    }

612
    fn add(&mut self, model: &str, checksum: &str, engine: E) -> Result<(), ModelManagerError> {
613
614
615
616
        if self.engines.contains_key(model) {
            return Err(ModelManagerError::ModelAlreadyExists(model.to_string()));
        }
        self.engines.insert(model.to_string(), engine);
617
618
        self.checksums
            .insert(model.to_string(), checksum.to_string());
619
620
621
622
623
624
625
        Ok(())
    }

    fn remove(&mut self, model: &str) -> Result<(), ModelManagerError> {
        if self.engines.remove(model).is_none() {
            return Err(ModelManagerError::ModelNotFound(model.to_string()));
        }
626
        let _ = self.checksums.remove(model);
627
628
629
630
631
632
633
634
635
636
637
638
639
640
        Ok(())
    }

    fn get(&self, model: &str) -> Option<&E> {
        self.engines.get(model)
    }

    fn contains(&self, model: &str) -> bool {
        self.engines.contains_key(model)
    }

    pub fn list(&self) -> Vec<String> {
        self.engines.keys().map(|k| k.to_owned()).collect()
    }
641
642
643
644
645
646

    /// Returns a newly allocated String for called convenience. All the places I use
    /// this I need a String.
    pub fn checksum(&self, model: &str) -> Option<String> {
        self.checksums.get(model).map(|s| s.to_string())
    }
647
}