model_manager.rs 37.1 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
use std::{collections::HashSet, sync::Arc};
5

6
use dashmap::{DashMap, mapref::entry::Entry};
7
use dynamo_kv_router::{PrefillLoadEstimator, config::KvRouterConfig, protocols::WorkerId};
8
use tokio::sync::oneshot;
9

10
use super::worker_monitor::LoadThresholdConfig;
11
use super::{KvWorkerMonitor, Model, RuntimeConfigWatch, WorkerSet, runtime_config_watch};
12

13
use dynamo_runtime::{
14
    component::{Endpoint, build_transport_type},
15
    discovery::DiscoverySpec,
16
17
18
    prelude::DistributedRuntimeProvider,
    protocols::EndpointId,
};
19
20

use crate::{
21
    kv_router::{KvRouter, router_endpoint_id, scheduler::DefaultWorkerSelector},
22
    local_model::runtime_config::DisaggregatedEndpoint,
23
24
25
26
    model_card::ModelDeploymentCard,
    types::{
        generic::tensor::TensorStreamingEngine,
        openai::{
27
            audios::OpenAIAudiosStreamingEngine,
28
29
            chat_completions::OpenAIChatCompletionsStreamingEngine,
            completions::OpenAICompletionsStreamingEngine,
30
            embeddings::OpenAIEmbeddingsStreamingEngine, images::OpenAIImagesStreamingEngine,
31
            videos::OpenAIVideosStreamingEngine,
32
33
        },
    },
34
};
35

36
37
38
39
40
41
42
43
/// 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>),
}

44
45
46
47
48
49
50
51
52
#[derive(Debug, thiserror::Error)]
pub enum ModelManagerError {
    #[error("Model not found: {0}")]
    ModelNotFound(String),

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

53
54
/// Central manager for model engines, routing, and configuration.
///
55
56
/// Models are stored hierarchically: ModelManager → Model → WorkerSet.
/// Each WorkerSet owns a complete pipeline built from its specific configuration.
57
58
///
/// Note: Don't implement Clone for this, put it in an Arc instead.
59
pub struct ModelManager {
60
61
    /// Model name → Model (which contains WorkerSets with engines)
    models: DashMap<String, Arc<Model>>,
62

63
    /// Per-instance model cards, keyed by instance path. Used for cleanup on worker removal.
64
    cards: DashMap<String, ModelDeploymentCard>,
65
66

    /// Prefill router activation rendezvous, keyed by "model_name:namespace".
67
    prefill_router_activators: DashMap<String, PrefillActivationState>,
68

69
    /// Per-endpoint runtime config watchers. Keyed by EndpointId (includes namespace).
70
    runtime_configs: DashMap<EndpointId, RuntimeConfigWatch>,
71
72
73
74
75
76
77
78
79
80
81
}

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

impl ModelManager {
    pub fn new() -> Self {
        Self {
82
            models: DashMap::new(),
83
84
            cards: DashMap::new(),
            prefill_router_activators: DashMap::new(),
85
            runtime_configs: DashMap::new(),
86
87
88
        }
    }

89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
    // -- Model access --

    /// Get or create a Model for the given name.
    pub fn get_or_create_model(&self, model_name: &str) -> Arc<Model> {
        self.models
            .entry(model_name.to_string())
            .or_insert_with(|| Arc::new(Model::new(model_name.to_string())))
            .clone()
    }

    /// Get an existing Model, if it exists.
    pub fn get_model(&self, model_name: &str) -> Option<Arc<Model>> {
        self.models
            .get(model_name)
            .map(|entry| entry.value().clone())
    }

    /// Remove a Model if it has no remaining WorkerSets.
    /// Uses atomic remove_if to avoid TOCTOU race between checking is_empty and removing.
    pub fn remove_model_if_empty(&self, model_name: &str) {
        if self
            .models
            .remove_if(model_name, |_, model| model.is_empty())
            .is_some()
        {
            tracing::info!(model_name, "Removed empty model from manager");
        }
    }

    /// Add a WorkerSet to a Model. Creates the Model if it doesn't exist.
119
    pub fn add_worker_set(&self, model_name: &str, namespace: &str, worker_set: WorkerSet) {
120
        let model = self.get_or_create_model(model_name);
121
        model.add_worker_set(namespace.to_string(), Arc::new(worker_set));
122
123
124
125
126
127
128
129
130
    }

    /// Remove a WorkerSet from a Model. Removes the Model if it becomes empty.
    pub fn remove_worker_set(&self, model_name: &str, namespace: &str) -> Option<Arc<WorkerSet>> {
        let model = self.models.get(model_name)?;
        let removed = model.remove_worker_set(namespace);
        drop(model);
        self.remove_model_if_empty(model_name);
        removed
131
132
    }

133
134
    // -- Model cards --

135
    pub fn get_model_cards(&self) -> Vec<ModelDeploymentCard> {
136
        self.cards.iter().map(|r| r.value().clone()).collect()
137
138
    }

139
140
141
142
143
144
145
146
147
148
149
150
151
152
    /// Save a ModelDeploymentCard from an instance's key so we can fetch it later when the key is
    /// deleted.
    pub fn save_model_card(&self, key: &str, card: ModelDeploymentCard) -> anyhow::Result<()> {
        self.cards.insert(key.to_string(), card);
        Ok(())
    }

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

    // -- Engine accessors (delegate through Model → WorkerSet) --

153
154
    /// Check if a decode model (chat or completions) is registered
    pub fn has_decode_model(&self, model: &str) -> bool {
155
156
157
        self.models
            .get(model)
            .is_some_and(|m| m.has_decode_engine())
158
159
160
161
    }

    /// Check if a prefill model is registered
    pub fn has_prefill_model(&self, model: &str) -> bool {
162
        self.models.get(model).is_some_and(|m| m.has_prefill())
163
164
165
166
167
    }

    /// Check if any model (decode or prefill) is registered.
    pub fn has_model_any(&self, model: &str) -> bool {
        self.has_decode_model(model) || self.has_prefill_model(model)
168
169
    }

170
    pub fn model_display_names(&self) -> HashSet<String> {
171
172
173
174
175
        self.models
            .iter()
            .filter(|entry| entry.value().is_displayable())
            .map(|entry| entry.key().clone())
            .collect()
176
177
    }

178
    pub fn list_chat_completions_models(&self) -> Vec<String> {
179
180
181
182
183
        self.models
            .iter()
            .filter(|entry| entry.value().has_chat_engine())
            .map(|entry| entry.key().clone())
            .collect()
184
185
186
    }

    pub fn list_completions_models(&self) -> Vec<String> {
187
188
189
190
191
        self.models
            .iter()
            .filter(|entry| entry.value().has_completions_engine())
            .map(|entry| entry.key().clone())
            .collect()
192
193
194
    }

    pub fn list_embeddings_models(&self) -> Vec<String> {
195
196
197
198
199
        self.models
            .iter()
            .filter(|entry| entry.value().has_embeddings_engine())
            .map(|entry| entry.key().clone())
            .collect()
200
201
    }

202
    pub fn list_tensor_models(&self) -> Vec<String> {
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
        self.models
            .iter()
            .filter(|entry| entry.value().has_tensor_engine())
            .map(|entry| entry.key().clone())
            .collect()
    }

    pub fn list_images_models(&self) -> Vec<String> {
        self.models
            .iter()
            .filter(|entry| entry.value().has_images_engine())
            .map(|entry| entry.key().clone())
            .collect()
    }

    pub fn list_videos_models(&self) -> Vec<String> {
        self.models
            .iter()
            .filter(|entry| entry.value().has_videos_engine())
            .map(|entry| entry.key().clone())
            .collect()
224
225
    }

226
    pub fn list_prefill_models(&self) -> Vec<String> {
227
228
229
230
231
        self.models
            .iter()
            .filter(|entry| entry.value().has_prefill())
            .map(|entry| entry.key().clone())
            .collect()
232
233
    }

234
235
236
237
238
239
240
241
    pub fn get_embeddings_engine(
        &self,
        model: &str,
    ) -> Result<OpenAIEmbeddingsStreamingEngine, ModelManagerError> {
        self.models
            .get(model)
            .ok_or_else(|| ModelManagerError::ModelNotFound(model.to_string()))?
            .get_embeddings_engine()
242
243
    }

244
245
246
247
248
249
250
251
    pub fn get_completions_engine(
        &self,
        model: &str,
    ) -> Result<OpenAICompletionsStreamingEngine, ModelManagerError> {
        self.models
            .get(model)
            .ok_or_else(|| ModelManagerError::ModelNotFound(model.to_string()))?
            .get_completions_engine()
252
253
    }

254
    pub fn get_chat_completions_engine(
255
256
        &self,
        model: &str,
257
258
259
260
261
262
263
264
265
266
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
    ) -> Result<OpenAIChatCompletionsStreamingEngine, ModelManagerError> {
        self.models
            .get(model)
            .ok_or_else(|| ModelManagerError::ModelNotFound(model.to_string()))?
            .get_chat_engine()
    }

    pub fn get_tensor_engine(
        &self,
        model: &str,
    ) -> Result<TensorStreamingEngine, ModelManagerError> {
        self.models
            .get(model)
            .ok_or_else(|| ModelManagerError::ModelNotFound(model.to_string()))?
            .get_tensor_engine()
    }

    pub fn get_images_engine(
        &self,
        model: &str,
    ) -> Result<OpenAIImagesStreamingEngine, ModelManagerError> {
        self.models
            .get(model)
            .ok_or_else(|| ModelManagerError::ModelNotFound(model.to_string()))?
            .get_images_engine()
    }

    pub fn get_videos_engine(
        &self,
        model: &str,
    ) -> Result<OpenAIVideosStreamingEngine, ModelManagerError> {
        self.models
            .get(model)
            .ok_or_else(|| ModelManagerError::ModelNotFound(model.to_string()))?
            .get_videos_engine()
    }

294
295
296
297
298
299
300
301
302
303
    pub fn get_audios_engine(
        &self,
        model: &str,
    ) -> Result<OpenAIAudiosStreamingEngine, ModelManagerError> {
        self.models
            .get(model)
            .ok_or_else(|| ModelManagerError::ModelNotFound(model.to_string()))?
            .get_audios_engine()
    }

304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
    // -- Combined engine + parsing options (atomically from one WorkerSet) --

    pub fn get_chat_completions_engine_with_parsing(
        &self,
        model: &str,
    ) -> Result<
        (
            OpenAIChatCompletionsStreamingEngine,
            crate::protocols::openai::ParsingOptions,
        ),
        ModelManagerError,
    > {
        self.models
            .get(model)
            .ok_or_else(|| ModelManagerError::ModelNotFound(model.to_string()))?
            .get_chat_engine_with_parsing()
320
321
    }

322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
    pub fn get_completions_engine_with_parsing(
        &self,
        model: &str,
    ) -> Result<
        (
            OpenAICompletionsStreamingEngine,
            crate::protocols::openai::ParsingOptions,
        ),
        ModelManagerError,
    > {
        self.models
            .get(model)
            .ok_or_else(|| ModelManagerError::ModelNotFound(model.to_string()))?
            .get_completions_engine_with_parsing()
    }

    // -- Convenience methods for in-process models (http.rs, grpc.rs) --
    // These create a WorkerSet with a default namespace for local models.
    // TODO: These methods use ModelDeploymentCard::default() for the WorkerSet, which means
    // parsing_options() returns defaults (no tool_call_parser/reasoning_parser). Pass the real
    // MDC from callers so ParsingOptions reflect the model's actual configuration.

344
345
346
    pub fn add_chat_completions_model(
        &self,
        model: &str,
347
        card_checksum: &str,
348
349
        engine: OpenAIChatCompletionsStreamingEngine,
    ) -> Result<(), ModelManagerError> {
350
351
352
353
354
355
356
357
358
359
360
        let model_entry = self.get_or_create_model(model);
        if model_entry.has_chat_engine() {
            return Err(ModelManagerError::ModelAlreadyExists(model.to_string()));
        }
        let namespace = format!("__local_chat_{}", model);
        let mut ws = WorkerSet::new(
            namespace.clone(),
            card_checksum.to_string(),
            ModelDeploymentCard::default(),
        );
        ws.chat_engine = Some(engine);
361
        model_entry.add_worker_set(namespace, Arc::new(ws));
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
        Ok(())
    }

    pub fn add_completions_model(
        &self,
        model: &str,
        card_checksum: &str,
        engine: OpenAICompletionsStreamingEngine,
    ) -> Result<(), ModelManagerError> {
        let model_entry = self.get_or_create_model(model);
        if model_entry.has_completions_engine() {
            return Err(ModelManagerError::ModelAlreadyExists(model.to_string()));
        }
        let namespace = format!("__local_completions_{}", model);
        let mut ws = WorkerSet::new(
            namespace.clone(),
            card_checksum.to_string(),
            ModelDeploymentCard::default(),
        );
        ws.completions_engine = Some(engine);
382
        model_entry.add_worker_set(namespace, Arc::new(ws));
383
        Ok(())
384
385
386
387
388
    }

    pub fn add_embeddings_model(
        &self,
        model: &str,
389
        card_checksum: &str,
390
391
        engine: OpenAIEmbeddingsStreamingEngine,
    ) -> Result<(), ModelManagerError> {
392
393
394
395
396
397
398
399
400
401
402
        let model_entry = self.get_or_create_model(model);
        if model_entry.has_embeddings_engine() {
            return Err(ModelManagerError::ModelAlreadyExists(model.to_string()));
        }
        let namespace = format!("__local_embeddings_{}", model);
        let mut ws = WorkerSet::new(
            namespace.clone(),
            card_checksum.to_string(),
            ModelDeploymentCard::default(),
        );
        ws.embeddings_engine = Some(engine);
403
        model_entry.add_worker_set(namespace, Arc::new(ws));
404
        Ok(())
405
406
    }

407
408
409
    pub fn add_tensor_model(
        &self,
        model: &str,
410
        card_checksum: &str,
411
412
        engine: TensorStreamingEngine,
    ) -> Result<(), ModelManagerError> {
413
414
415
416
417
418
419
420
421
422
423
        let model_entry = self.get_or_create_model(model);
        if model_entry.has_tensor_engine() {
            return Err(ModelManagerError::ModelAlreadyExists(model.to_string()));
        }
        let namespace = format!("__local_tensor_{}", model);
        let mut ws = WorkerSet::new(
            namespace.clone(),
            card_checksum.to_string(),
            ModelDeploymentCard::default(),
        );
        ws.tensor_engine = Some(engine);
424
        model_entry.add_worker_set(namespace, Arc::new(ws));
425
        Ok(())
426
427
    }

428
429
430
431
432
433
    pub fn add_images_model(
        &self,
        model: &str,
        card_checksum: &str,
        engine: OpenAIImagesStreamingEngine,
    ) -> Result<(), ModelManagerError> {
434
435
436
437
438
439
440
441
442
443
444
        let model_entry = self.get_or_create_model(model);
        if model_entry.has_images_engine() {
            return Err(ModelManagerError::ModelAlreadyExists(model.to_string()));
        }
        let namespace = format!("__local_images_{}", model);
        let mut ws = WorkerSet::new(
            namespace.clone(),
            card_checksum.to_string(),
            ModelDeploymentCard::default(),
        );
        ws.images_engine = Some(engine);
445
        model_entry.add_worker_set(namespace, Arc::new(ws));
446
        Ok(())
447
448
    }

449
450
451
452
453
454
    pub fn add_videos_model(
        &self,
        model: &str,
        card_checksum: &str,
        engine: OpenAIVideosStreamingEngine,
    ) -> Result<(), ModelManagerError> {
455
456
457
458
459
460
461
462
463
464
465
        let model_entry = self.get_or_create_model(model);
        if model_entry.has_videos_engine() {
            return Err(ModelManagerError::ModelAlreadyExists(model.to_string()));
        }
        let namespace = format!("__local_videos_{}", model);
        let mut ws = WorkerSet::new(
            namespace.clone(),
            card_checksum.to_string(),
            ModelDeploymentCard::default(),
        );
        ws.videos_engine = Some(engine);
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
486
        model_entry.add_worker_set(namespace, Arc::new(ws));
        Ok(())
    }

    pub fn add_audios_model(
        &self,
        model: &str,
        card_checksum: &str,
        engine: OpenAIAudiosStreamingEngine,
    ) -> Result<(), ModelManagerError> {
        let model_entry = self.get_or_create_model(model);
        if model_entry.has_audios_engine() {
            return Err(ModelManagerError::ModelAlreadyExists(model.to_string()));
        }
        let namespace = format!("__local_audios_{}", model);
        let mut ws = WorkerSet::new(
            namespace.clone(),
            card_checksum.to_string(),
            ModelDeploymentCard::default(),
        );
        ws.audios_engine = Some(engine);
487
        model_entry.add_worker_set(namespace, Arc::new(ws));
488
        Ok(())
489
490
    }

491
492
493
494
495
    pub fn add_prefill_model(
        &self,
        model: &str,
        card_checksum: &str,
    ) -> Result<(), ModelManagerError> {
496
497
498
499
500
501
502
503
504
505
        let model_entry = self.get_or_create_model(model);
        if model_entry.has_prefill() {
            return Err(ModelManagerError::ModelAlreadyExists(model.to_string()));
        }
        let namespace = format!("__local_prefill_{}", model);
        let ws = WorkerSet::new(
            namespace.clone(),
            card_checksum.to_string(),
            ModelDeploymentCard::default(),
        );
506
        model_entry.add_worker_set(namespace, Arc::new(ws));
507
        Ok(())
508
509
    }

510
    // -- Model removal --
511

512
513
514
515
    /// Remove a model entirely (all its WorkerSets).
    /// Returns the removed Model, or None if not found.
    pub fn remove_model(&self, model: &str) -> Option<Arc<Model>> {
        self.models.remove(model).map(|(_, m)| m)
516
517
    }

518
519
    // Per-type remove methods for in-process models (used by Python bindings).
    // These remove the specific synthetic WorkerSet created by the corresponding add_*_model method.
520

521
522
523
524
525
    pub fn remove_chat_completions_model(&self, model: &str) -> Result<(), ModelManagerError> {
        let namespace = format!("__local_chat_{}", model);
        self.remove_worker_set(model, &namespace)
            .map(|_| ())
            .ok_or_else(|| ModelManagerError::ModelNotFound(model.to_string()))
526
527
    }

528
529
530
531
532
    pub fn remove_completions_model(&self, model: &str) -> Result<(), ModelManagerError> {
        let namespace = format!("__local_completions_{}", model);
        self.remove_worker_set(model, &namespace)
            .map(|_| ())
            .ok_or_else(|| ModelManagerError::ModelNotFound(model.to_string()))
533
534
    }

535
536
537
538
539
    pub fn remove_tensor_model(&self, model: &str) -> Result<(), ModelManagerError> {
        let namespace = format!("__local_tensor_{}", model);
        self.remove_worker_set(model, &namespace)
            .map(|_| ())
            .ok_or_else(|| ModelManagerError::ModelNotFound(model.to_string()))
540
541
    }

542
543
544
545
546
    pub fn remove_embeddings_model(&self, model: &str) -> Result<(), ModelManagerError> {
        let namespace = format!("__local_embeddings_{}", model);
        self.remove_worker_set(model, &namespace)
            .map(|_| ())
            .ok_or_else(|| ModelManagerError::ModelNotFound(model.to_string()))
547
548
    }

549
550
551
552
553
    pub fn remove_images_model(&self, model: &str) -> Result<(), ModelManagerError> {
        let namespace = format!("__local_images_{}", model);
        self.remove_worker_set(model, &namespace)
            .map(|_| ())
            .ok_or_else(|| ModelManagerError::ModelNotFound(model.to_string()))
554
555
    }

556
557
558
559
560
    pub fn remove_videos_model(&self, model: &str) -> Result<(), ModelManagerError> {
        let namespace = format!("__local_videos_{}", model);
        self.remove_worker_set(model, &namespace)
            .map(|_| ())
            .ok_or_else(|| ModelManagerError::ModelNotFound(model.to_string()))
561
562
    }

563
    // -- KV Router creation --
564

565
    #[allow(clippy::too_many_arguments)]
566
567
    pub async fn kv_chooser_for(
        &self,
568
        endpoint: &Endpoint,
569
        kv_cache_block_size: u32,
570
        kv_router_config: Option<KvRouterConfig>,
571
        prefill_load_estimator: Option<Arc<dyn PrefillLoadEstimator>>,
572
        worker_type: &'static str,
573
        model_name: Option<String>,
574
        is_eagle: bool,
575
    ) -> anyhow::Result<Arc<KvRouter>> {
576
        let client = endpoint.client().await?;
577

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

582
        // Build transport for router endpoint based on request plane mode
583
584
585
        // Use the worker's component name so each target pool gets its own router discovery group
        let router_endpoint_id =
            router_endpoint_id(endpoint.id().namespace, endpoint.id().component);
586
        let transport = build_transport_type(endpoint, &router_endpoint_id, instance_id).await?;
587
588
589
590
591

        let discovery_spec = DiscoverySpec::Endpoint {
            namespace: router_endpoint_id.namespace.clone(),
            component: router_endpoint_id.component.clone(),
            endpoint: router_endpoint_id.name.clone(),
592
            transport,
593
594
595
596
        };

        discovery.register(discovery_spec).await?;

597
        // Get of create runtime config watcher for this endpoint
598
599
        let workers_with_configs = self.get_or_create_runtime_config_watcher(endpoint).await?;

600
        let selector = DefaultWorkerSelector::new(kv_router_config.clone(), worker_type);
601
        let chooser = KvRouter::new(
602
603
            endpoint.clone(),
            client,
604
            workers_with_configs,
605
            kv_cache_block_size,
606
            selector,
607
            kv_router_config,
608
            prefill_load_estimator,
609
            worker_type,
610
            model_name,
611
            is_eagle,
612
613
        )
        .await?;
614
        Ok(Arc::new(chooser))
615
    }
616

617
618
619
620
621
622
623
624
    // -- Prefill router coordination --
    // Keyed by "model_name:namespace" so each namespace's decode WorkerSet gets its own
    // prefill router activated by same-namespace prefill workers.

    /// Build a key for a (model, namespace) pair. Used for prefill router activators
    /// and registration guards.
    pub(crate) fn model_namespace_key(model_name: &str, namespace: &str) -> String {
        format!("{}:{}", model_name, namespace)
625
626
    }

627
628
629
    /// Register a prefill router for a decode WorkerSet. Returns a receiver that will be
    /// activated when the corresponding prefill model in the same namespace is discovered.
    /// Returns None if a decode WorkerSet in this namespace was already registered.
630
631
    pub fn register_prefill_router(
        &self,
632
633
        model_name: &str,
        namespace: &str,
634
    ) -> Option<oneshot::Receiver<Endpoint>> {
635
636
        let key = Self::model_namespace_key(model_name, namespace);
        match self.prefill_router_activators.remove(&key) {
637
            Some((_, PrefillActivationState::PrefillReady(rx))) => {
638
639
640
                // Prefill endpoint already arrived - rx will immediately resolve
                tracing::debug!(
                    model_name = %model_name,
641
642
                    namespace = %namespace,
                    "Prefill endpoint already available for namespace, returning receiver"
643
644
645
                );
                Some(rx)
            }
646
            Some((key, PrefillActivationState::DecodeWaiting(tx))) => {
647
648
649
                // Decode already registered - this shouldn't happen, restore state and return None
                tracing::error!(
                    model_name = %model_name,
650
651
                    namespace = %namespace,
                    "Decode WorkerSet already registered for this prefill router"
652
                );
653
654
                self.prefill_router_activators
                    .insert(key, PrefillActivationState::DecodeWaiting(tx));
655
656
657
658
659
                None
            }
            None => {
                // New registration: create tx/rx pair, store sender and return receiver
                let (tx, rx) = oneshot::channel();
660
661
                self.prefill_router_activators
                    .insert(key, PrefillActivationState::DecodeWaiting(tx));
662
663
                tracing::debug!(
                    model_name = %model_name,
664
665
                    namespace = %namespace,
                    "No prefill endpoint for namespace yet, storing sender for future activation"
666
667
668
669
670
671
672
                );
                Some(rx)
            }
        }
    }

    /// Activate a prefill router by sending the endpoint through the oneshot channel.
673
    /// The namespace must match the decode WorkerSet's namespace.
674
675
676
    pub fn activate_prefill_router(
        &self,
        model_name: &str,
677
        namespace: &str,
678
679
        endpoint: Endpoint,
    ) -> anyhow::Result<()> {
680
681
        let key = Self::model_namespace_key(model_name, namespace);
        match self.prefill_router_activators.remove(&key) {
682
            Some((_, PrefillActivationState::DecodeWaiting(sender))) => {
683
684
                sender.send(endpoint).map_err(|_| {
                    anyhow::anyhow!(
685
686
687
                        "Failed to send endpoint to prefill router activator for {}:{}",
                        model_name,
                        namespace
688
689
690
691
                    )
                })?;
                tracing::info!(
                    model_name = %model_name,
692
693
                    namespace = %namespace,
                    "Activated prefill router for decode WorkerSet"
694
695
696
                );
                Ok(())
            }
697
            Some((_, PrefillActivationState::PrefillReady(_))) => {
698
699
700
701
702
                anyhow::bail!(
                    "Prefill router for {}:{} already activated",
                    model_name,
                    namespace
                );
703
704
705
706
            }
            None => {
                let (tx, rx) = oneshot::channel();
                tx.send(endpoint).map_err(|_| {
707
708
709
710
711
                    anyhow::anyhow!(
                        "Failed to send endpoint for prefill model {}:{}",
                        model_name,
                        namespace
                    )
712
                })?;
713
714
                self.prefill_router_activators
                    .insert(key, PrefillActivationState::PrefillReady(rx));
715
716
                tracing::info!(
                    model_name = %model_name,
717
718
                    namespace = %namespace,
                    "Stored prefill endpoint for future decode WorkerSet registration"
719
720
721
722
                );
                Ok(())
            }
        }
723
724
    }

725
726
727
728
729
730
731
732
733
734
735
    /// Remove the prefill router activator for a (model, namespace) pair.
    /// Called when a WorkerSet is removed to prevent stale activators.
    pub fn remove_prefill_activator(&self, model_name: &str, namespace: &str) {
        let key = Self::model_namespace_key(model_name, namespace);
        if self.prefill_router_activators.remove(&key).is_some() {
            tracing::debug!(
                model_name = %model_name,
                namespace = %namespace,
                "Cleaned up prefill router activator for removed WorkerSet"
            );
        }
736
737
    }

738
    // -- Worker monitoring --
739

740
    /// Gets or sets the load threshold config for a model's worker monitor.
741
    /// Checks across all WorkerSets for the model.
742
    pub fn load_threshold_config(
743
744
        &self,
        model: &str,
745
746
        config: Option<&LoadThresholdConfig>,
    ) -> Option<LoadThresholdConfig> {
747
748
        let model_entry = self.models.get(model)?;
        model_entry.load_threshold_config(config)
749
750
    }

751
752
    /// Gets an existing worker monitor for a specific namespace of a model.
    pub fn get_worker_monitor_for_namespace(
753
754
        &self,
        model: &str,
755
756
757
758
759
760
761
762
763
764
765
766
767
        namespace: &str,
    ) -> Option<KvWorkerMonitor> {
        let model_entry = self.models.get(model)?;
        model_entry.get_worker_monitor_for_namespace(namespace)
    }

    /// Lists all models with worker monitors configured.
    pub fn list_busy_thresholds(&self) -> Vec<(String, LoadThresholdConfig)> {
        let mut result = Vec::new();
        for entry in self.models.iter() {
            if let Some(config) = entry.value().load_threshold_config(None) {
                result.push((entry.key().clone(), config));
            }
768
        }
769
        result
770
771
    }

772
773
    // -- Runtime configs --

774
    /// Get or create a runtime config watcher for an endpoint.
775
776
    /// Spawns a background task that joins instance availability and config discovery.
    /// Returns a `watch::Receiver` with the latest `HashMap<WorkerId, ModelRuntimeConfig>`.
777
778
779
    pub async fn get_or_create_runtime_config_watcher(
        &self,
        endpoint: &Endpoint,
780
    ) -> anyhow::Result<RuntimeConfigWatch> {
781
782
783
784
785
786
        let endpoint_id = endpoint.id();

        if let Some(existing) = self.runtime_configs.get(&endpoint_id) {
            return Ok(existing.clone());
        }

787
788
789
790
791
792
        // Slow path: create the watch (spawns a background task).
        // If another caller raced us, the entry() below picks up the winner;
        // the loser's background task stops once its receivers are dropped.
        let rx = runtime_config_watch(endpoint).await?;
        let result = match self.runtime_configs.entry(endpoint_id) {
            Entry::Occupied(e) => e.get().clone(),
793
            Entry::Vacant(e) => {
794
795
                e.insert(rx.clone());
                rx
796
797
798
            }
        };

799
        Ok(result)
800
801
802
803
804
805
806
807
    }

    /// Get disaggregated endpoint for a specific worker.
    pub fn get_disaggregated_endpoint(
        &self,
        endpoint_id: &EndpointId,
        worker_id: WorkerId,
    ) -> Option<DisaggregatedEndpoint> {
808
809
810
        let rx = self.runtime_configs.get(endpoint_id)?;
        let configs = rx.borrow();
        configs.get(&worker_id)?.disaggregated_endpoint.clone()
811
    }
812
}
813

814
815
816
817
818
819
820
821
822
823
824
#[cfg(test)]
mod tests {
    use super::*;
    use crate::model_card::ModelDeploymentCard;

    fn make_worker_set(namespace: &str, mdcsum: &str) -> WorkerSet {
        WorkerSet::new(
            namespace.to_string(),
            mdcsum.to_string(),
            ModelDeploymentCard::default(),
        )
825
    }
826

827
    // -- CRUD delegation tests --
828

829
830
831
832
    #[test]
    fn test_add_and_get_worker_set() {
        let mm = ModelManager::new();
        let ws = make_worker_set("ns1", "abc");
833
        mm.add_worker_set("llama", "ns1", ws);
834
835
836
837
838
839

        let model = mm.get_model("llama");
        assert!(model.is_some());
        let model = model.unwrap();
        assert!(model.has_worker_set("ns1"));
        assert_eq!(model.worker_set_count(), 1);
840
841
    }

842
843
844
845
846
    #[test]
    fn test_add_worker_set_creates_model() {
        let mm = ModelManager::new();
        assert!(mm.get_model("llama").is_none());

847
        mm.add_worker_set("llama", "ns1", make_worker_set("ns1", "abc"));
848
        assert!(mm.get_model("llama").is_some());
849
850
    }

851
852
853
    #[test]
    fn test_remove_worker_set_removes_empty_model() {
        let mm = ModelManager::new();
854
        mm.add_worker_set("llama", "ns1", make_worker_set("ns1", "abc"));
855
856
857
858
859
860
861
862
        assert!(mm.get_model("llama").is_some());

        let removed = mm.remove_worker_set("llama", "ns1");
        assert!(removed.is_some());
        assert_eq!(removed.unwrap().namespace(), "ns1");

        // Model should be auto-removed since it's now empty
        assert!(mm.get_model("llama").is_none());
863
864
    }

865
866
867
    #[test]
    fn test_remove_worker_set_keeps_model_with_remaining() {
        let mm = ModelManager::new();
868
869
        mm.add_worker_set("llama", "ns1", make_worker_set("ns1", "abc"));
        mm.add_worker_set("llama", "ns2", make_worker_set("ns2", "abc"));
870
871
872
873
874
875
876
877

        mm.remove_worker_set("llama", "ns1");

        // Model should still exist with ns2
        let model = mm.get_model("llama").unwrap();
        assert!(!model.has_worker_set("ns1"));
        assert!(model.has_worker_set("ns2"));
        assert_eq!(model.worker_set_count(), 1);
878
879
    }

880
881
882
883
884
885
886
887
888
    #[test]
    fn test_remove_worker_set_nonexistent_model() {
        let mm = ModelManager::new();
        assert!(mm.remove_worker_set("llama", "ns1").is_none());
    }

    #[test]
    fn test_remove_worker_set_nonexistent_namespace() {
        let mm = ModelManager::new();
889
        mm.add_worker_set("llama", "ns1", make_worker_set("ns1", "abc"));
890
891
892
893
894
895
896
897
898
        assert!(mm.remove_worker_set("llama", "ns2").is_none());

        // Model should still exist (ns1 still there)
        assert!(mm.get_model("llama").is_some());
    }

    #[test]
    fn test_remove_model_if_empty_noop_when_not_empty() {
        let mm = ModelManager::new();
899
        mm.add_worker_set("llama", "ns1", make_worker_set("ns1", "abc"));
900
901
902

        mm.remove_model_if_empty("llama");
        assert!(mm.get_model("llama").is_some()); // Still has ns1
903
904
    }

905
906
907
908
    #[test]
    fn test_remove_model_if_empty_noop_when_missing() {
        let mm = ModelManager::new();
        mm.remove_model_if_empty("nonexistent"); // Should not panic
909
910
    }

911
912
913
    #[test]
    fn test_remove_model() {
        let mm = ModelManager::new();
914
915
        mm.add_worker_set("llama", "ns1", make_worker_set("ns1", "abc"));
        mm.add_worker_set("llama", "ns2", make_worker_set("ns2", "abc"));
916
917
918
919
920
921
922
923
924
925
926
927
928
929
930
931
932
933
934
935
936
937
938
939
940

        let removed = mm.remove_model("llama");
        assert!(removed.is_some());
        assert!(mm.get_model("llama").is_none());
    }

    #[test]
    fn test_get_or_create_model_idempotent() {
        let mm = ModelManager::new();
        let m1 = mm.get_or_create_model("llama");
        let m2 = mm.get_or_create_model("llama");
        // Both should point to the same Model (same Arc)
        assert!(Arc::ptr_eq(&m1, &m2));
    }

    // -- Model listing and filtering tests --

    #[test]
    fn test_has_decode_model() {
        let mm = ModelManager::new();

        // No model → false
        assert!(!mm.has_decode_model("llama"));

        // Prefill-only set (no engines) → false
941
        mm.add_worker_set("llama", "ns1", make_worker_set("ns1", "abc"));
942
943
944
945
946
947
948
949
        assert!(!mm.has_decode_model("llama"));
    }

    #[test]
    fn test_has_prefill_model() {
        let mm = ModelManager::new();

        // Prefill set = no engines
950
        mm.add_worker_set("llama", "ns1", make_worker_set("ns1", "abc"));
951
952
953
954
955
956
957
958
        assert!(mm.has_prefill_model("llama"));
    }

    #[test]
    fn test_has_model_any() {
        let mm = ModelManager::new();
        assert!(!mm.has_model_any("llama"));

959
        mm.add_worker_set("llama", "ns1", make_worker_set("ns1", "abc"));
960
961
962
963
964
965
        assert!(mm.has_model_any("llama")); // has prefill
    }

    #[test]
    fn test_model_display_names_includes_prefill() {
        let mm = ModelManager::new();
966
        mm.add_worker_set("llama", "ns1", make_worker_set("ns1", "abc"));
967
968
969
970
971
972
973
974
975
976
977
978
979
980

        let names = mm.model_display_names();
        assert!(names.contains("llama"));
    }

    #[test]
    fn test_model_display_names_empty() {
        let mm = ModelManager::new();
        assert!(mm.model_display_names().is_empty());
    }

    #[test]
    fn test_list_prefill_models() {
        let mm = ModelManager::new();
981
982
        mm.add_worker_set("llama", "ns1", make_worker_set("ns1", "abc"));
        mm.add_worker_set("gpt", "ns1", make_worker_set("ns1", "def"));
983
984
985
986
987
988
989
990
991
992
993
994
995
996
997
998
999
1000
1001
1002
1003
1004
1005
1006
1007
1008
1009
1010
1011
1012
1013
1014
1015
1016
1017
1018
1019
1020
1021
1022
1023
1024
1025
1026
1027
1028
1029
1030
1031
1032
1033
1034
1035
1036
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
1071

        let prefill = mm.list_prefill_models();
        assert_eq!(prefill.len(), 2);
        assert!(prefill.contains(&"llama".to_string()));
        assert!(prefill.contains(&"gpt".to_string()));
    }

    // -- Model card tests --

    #[test]
    fn test_save_and_remove_model_card() {
        let mm = ModelManager::new();
        let card = ModelDeploymentCard::default();
        mm.save_model_card("instance/key/1", card.clone()).unwrap();

        let cards = mm.get_model_cards();
        assert_eq!(cards.len(), 1);

        let removed = mm.remove_model_card("instance/key/1");
        assert!(removed.is_some());
        assert!(mm.get_model_cards().is_empty());
    }

    #[test]
    fn test_remove_model_card_nonexistent() {
        let mm = ModelManager::new();
        assert!(mm.remove_model_card("nonexistent").is_none());
    }

    // -- Prefill router rendezvous tests --
    // Note: activate_prefill_router requires an Endpoint (needs DistributedRuntime),
    // so we test the registration state machine and cleanup only.

    #[test]
    fn test_prefill_router_register_new() {
        let mm = ModelManager::new();

        // First registration for a (model, namespace) returns Some(rx)
        let rx = mm.register_prefill_router("llama", "ns1");
        assert!(rx.is_some());
    }

    #[test]
    fn test_prefill_router_double_register_returns_none() {
        let mm = ModelManager::new();

        let rx1 = mm.register_prefill_router("llama", "ns1");
        assert!(rx1.is_some());

        // Second registration for the same (model, namespace) returns None
        let rx2 = mm.register_prefill_router("llama", "ns1");
        assert!(rx2.is_none());
    }

    #[test]
    fn test_prefill_router_different_namespaces_independent() {
        let mm = ModelManager::new();

        // Different namespaces should be independent
        let rx1 = mm.register_prefill_router("llama", "ns1");
        let rx2 = mm.register_prefill_router("llama", "ns2");
        assert!(rx1.is_some());
        assert!(rx2.is_some());
    }

    #[test]
    fn test_prefill_router_different_models_independent() {
        let mm = ModelManager::new();

        // Different models should be independent
        let rx1 = mm.register_prefill_router("llama", "ns1");
        let rx2 = mm.register_prefill_router("gpt", "ns1");
        assert!(rx1.is_some());
        assert!(rx2.is_some());
    }

    #[test]
    fn test_prefill_router_remove_allows_reregister() {
        let mm = ModelManager::new();

        let rx = mm.register_prefill_router("llama", "ns1");
        assert!(rx.is_some());

        // Remove the activator
        mm.remove_prefill_activator("llama", "ns1");

        // Should be able to register again
        let rx2 = mm.register_prefill_router("llama", "ns1");
        assert!(rx2.is_some());
1072
1073
    }

1074
1075
1076
1077
1078
    #[test]
    fn test_prefill_router_remove_nonexistent_noop() {
        let mm = ModelManager::new();
        // Should not panic
        mm.remove_prefill_activator("llama", "ns1");
1079
    }
1080

1081
1082
1083
1084
1085
1086
1087
1088
1089
1090
    #[test]
    fn test_model_namespace_key_format() {
        assert_eq!(
            ModelManager::model_namespace_key("llama", "ns1"),
            "llama:ns1"
        );
        assert_eq!(
            ModelManager::model_namespace_key("gpt-4", "default-abc"),
            "gpt-4:default-abc"
        );
1091
    }
1092
}