model_manager.rs 37.2 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
            device_type: None,
594
595
596
597
        };

        discovery.register(discovery_spec).await?;

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

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

618
619
620
621
622
623
624
625
    // -- 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)
626
627
    }

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

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

726
727
728
729
730
731
732
733
734
735
736
    /// 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"
            );
        }
737
738
    }

739
    // -- Worker monitoring --
740

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

752
753
    /// Gets an existing worker monitor for a specific namespace of a model.
    pub fn get_worker_monitor_for_namespace(
754
755
        &self,
        model: &str,
756
757
758
759
760
761
762
763
764
765
766
767
768
        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));
            }
769
        }
770
        result
771
772
    }

773
774
    // -- Runtime configs --

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

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

788
789
790
791
792
793
        // 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(),
794
            Entry::Vacant(e) => {
795
796
                e.insert(rx.clone());
                rx
797
798
799
            }
        };

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

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

815
816
817
818
819
820
821
822
823
824
825
#[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(),
        )
826
    }
827

828
    // -- CRUD delegation tests --
829

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

        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);
841
842
    }

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

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

852
853
854
    #[test]
    fn test_remove_worker_set_removes_empty_model() {
        let mm = ModelManager::new();
855
        mm.add_worker_set("llama", "ns1", make_worker_set("ns1", "abc"));
856
857
858
859
860
861
862
863
        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());
864
865
    }

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

        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);
879
880
    }

881
882
883
884
885
886
887
888
889
    #[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();
890
        mm.add_worker_set("llama", "ns1", make_worker_set("ns1", "abc"));
891
892
893
894
895
896
897
898
899
        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();
900
        mm.add_worker_set("llama", "ns1", make_worker_set("ns1", "abc"));
901
902
903

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

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

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

        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
942
        mm.add_worker_set("llama", "ns1", make_worker_set("ns1", "abc"));
943
944
945
946
947
948
949
950
        assert!(!mm.has_decode_model("llama"));
    }

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

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

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

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

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

        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();
982
983
        mm.add_worker_set("llama", "ns1", make_worker_set("ns1", "abc"));
        mm.add_worker_set("gpt", "ns1", make_worker_set("ns1", "def"));
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
1072

        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());
1073
1074
    }

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

1082
1083
1084
1085
1086
1087
1088
1089
1090
1091
    #[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"
        );
1092
    }
1093
}