model_manager.rs 45.4 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
#[derive(Debug, thiserror::Error)]
pub enum ModelManagerError {
    #[error("Model not found: {0}")]
    ModelNotFound(String),

49
50
51
    #[error("Model unavailable: {0}")]
    ModelUnavailable(String),

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

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

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

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

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

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

impl ModelManager {
    pub fn new() -> Self {
        Self {
85
            models: DashMap::new(),
86
87
            cards: DashMap::new(),
            prefill_router_activators: DashMap::new(),
88
            runtime_configs: DashMap::new(),
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
119
120
121
    // -- 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.
122
    pub fn add_worker_set(&self, model_name: &str, namespace: &str, worker_set: WorkerSet) {
123
        let model = self.get_or_create_model(model_name);
124
        model.add_worker_set(namespace.to_string(), Arc::new(worker_set));
125
126
127
128
129
130
131
132
133
    }

    /// 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
134
135
    }

136
137
    // -- Model cards --

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

142
143
144
145
146
147
148
149
150
151
152
153
154
155
    /// 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) --

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

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

    /// 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)
171
172
    }

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

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

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

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

205
    pub fn list_tensor_models(&self) -> Vec<String> {
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
        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()
227
228
    }

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

237
238
239
240
241
242
243
244
    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()
245
246
    }

247
248
249
250
251
252
253
254
    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()
255
256
    }

257
    pub fn get_chat_completions_engine(
258
259
        &self,
        model: &str,
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
294
295
296
    ) -> 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()
    }

297
298
299
300
301
302
303
304
305
306
    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()
    }

307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
    // -- 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()
323
324
    }

325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
    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.

347
348
349
    pub fn add_chat_completions_model(
        &self,
        model: &str,
350
        card_checksum: &str,
351
352
        engine: OpenAIChatCompletionsStreamingEngine,
    ) -> Result<(), ModelManagerError> {
353
354
355
356
357
358
359
360
361
362
363
        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);
364
        model_entry.add_worker_set(namespace, Arc::new(ws));
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
        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);
385
        model_entry.add_worker_set(namespace, Arc::new(ws));
386
        Ok(())
387
388
389
390
391
    }

    pub fn add_embeddings_model(
        &self,
        model: &str,
392
        card_checksum: &str,
393
394
        engine: OpenAIEmbeddingsStreamingEngine,
    ) -> Result<(), ModelManagerError> {
395
396
397
398
399
400
401
402
403
404
405
        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);
406
        model_entry.add_worker_set(namespace, Arc::new(ws));
407
        Ok(())
408
409
    }

410
411
412
    pub fn add_tensor_model(
        &self,
        model: &str,
413
        card_checksum: &str,
414
415
        engine: TensorStreamingEngine,
    ) -> Result<(), ModelManagerError> {
416
417
418
419
420
421
422
423
424
425
426
        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);
427
        model_entry.add_worker_set(namespace, Arc::new(ws));
428
        Ok(())
429
430
    }

431
432
433
434
435
436
    pub fn add_images_model(
        &self,
        model: &str,
        card_checksum: &str,
        engine: OpenAIImagesStreamingEngine,
    ) -> Result<(), ModelManagerError> {
437
438
439
440
441
442
443
444
445
446
447
        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);
448
        model_entry.add_worker_set(namespace, Arc::new(ws));
449
        Ok(())
450
451
    }

452
453
454
455
456
457
    pub fn add_videos_model(
        &self,
        model: &str,
        card_checksum: &str,
        engine: OpenAIVideosStreamingEngine,
    ) -> Result<(), ModelManagerError> {
458
459
460
461
462
463
464
465
466
467
468
        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);
469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
        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);
490
        model_entry.add_worker_set(namespace, Arc::new(ws));
491
        Ok(())
492
493
    }

494
495
496
497
498
    pub fn add_prefill_model(
        &self,
        model: &str,
        card_checksum: &str,
    ) -> Result<(), ModelManagerError> {
499
500
501
502
503
504
505
506
507
508
        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(),
        );
509
        model_entry.add_worker_set(namespace, Arc::new(ws));
510
        Ok(())
511
512
    }

513
    // -- Model removal --
514

515
516
517
518
    /// 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)
519
520
    }

521
522
    // 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.
523

524
525
526
527
528
    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()))
529
530
    }

531
532
533
534
535
    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()))
536
537
    }

538
539
540
541
542
    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()))
543
544
    }

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

552
553
554
555
556
    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()))
557
558
    }

559
560
561
562
563
    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()))
564
565
    }

566
    // -- KV Router creation --
567

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

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

585
        // Build transport for router endpoint based on request plane mode
586
587
588
        // 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);
589
        let transport = build_transport_type(endpoint, &router_endpoint_id, instance_id).await?;
590
591
592
593
594

        let discovery_spec = DiscoverySpec::Endpoint {
            namespace: router_endpoint_id.namespace.clone(),
            component: router_endpoint_id.component.clone(),
            endpoint: router_endpoint_id.name.clone(),
595
            transport,
596
            device_type: None,
597
598
599
600
        };

        discovery.register(discovery_spec).await?;

601
        // Get of create runtime config watcher for this endpoint
602
603
        let workers_with_configs = self.get_or_create_runtime_config_watcher(endpoint).await?;

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

621
622
623
624
625
626
627
628
    // -- 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)
629
630
    }

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

    /// Activate a prefill router by sending the endpoint through the oneshot channel.
677
    /// The namespace must match the decode WorkerSet's namespace.
678
679
680
    pub fn activate_prefill_router(
        &self,
        model_name: &str,
681
        namespace: &str,
682
683
        endpoint: Endpoint,
    ) -> anyhow::Result<()> {
684
685
        let key = Self::model_namespace_key(model_name, namespace);
        match self.prefill_router_activators.remove(&key) {
686
            Some((_, PrefillActivationState::DecodeWaiting(sender))) => {
687
688
                sender.send(endpoint).map_err(|_| {
                    anyhow::anyhow!(
689
690
691
                        "Failed to send endpoint to prefill router activator for {}:{}",
                        model_name,
                        namespace
692
693
694
695
                    )
                })?;
                tracing::info!(
                    model_name = %model_name,
696
697
                    namespace = %namespace,
                    "Activated prefill router for decode WorkerSet"
698
699
700
                );
                Ok(())
            }
701
            Some((_, PrefillActivationState::PrefillReady(_))) => {
702
703
704
705
706
                anyhow::bail!(
                    "Prefill router for {}:{} already activated",
                    model_name,
                    namespace
                );
707
708
            }
            None => {
709
710
711
712
713
714
715
716
717
718
719
720
721
722
723
724
725
726
727
728
729
730
731
732
733
734
735
736
737
738
739
740
741
                // Try to reactivate an existing deactivated router first.
                // This handles prefill rejoin after a transient failure: the decode
                // WorkerSet's PrefillRouter already exists but is deactivated.
                if let Some(model) = self.get_model(model_name)
                    && let Some(ws) = model.get_worker_set(namespace)
                    && let Some(ref pr) = ws.prefill_router
                    && pr.is_deactivated()
                {
                    pr.reactivate();
                    // Store the endpoint so that if the decode WorkerSet is rebuilt
                    // (removed and re-added), a subsequent register_prefill_router call
                    // finds PrefillReady instead of falling back to DecodeWaiting and
                    // stalling.
                    let (tx, rx) = oneshot::channel();
                    tx.send(endpoint).map_err(|_| {
                        anyhow::anyhow!(
                            "Failed to send endpoint for prefill model {}:{}",
                            model_name,
                            namespace
                        )
                    })?;
                    self.prefill_router_activators
                        .insert(key, PrefillActivationState::PrefillReady(rx));
                    tracing::info!(
                        model_name = %model_name,
                        namespace = %namespace,
                        "Reactivated existing prefill router for decode WorkerSet (prefill rejoin)"
                    );
                    return Ok(());
                }

                // No existing deactivated router -- store endpoint for a future decode
                // registration.
742
743
                let (tx, rx) = oneshot::channel();
                tx.send(endpoint).map_err(|_| {
744
745
746
747
748
                    anyhow::anyhow!(
                        "Failed to send endpoint for prefill model {}:{}",
                        model_name,
                        namespace
                    )
749
                })?;
750
751
                self.prefill_router_activators
                    .insert(key, PrefillActivationState::PrefillReady(rx));
752
753
                tracing::info!(
                    model_name = %model_name,
754
755
                    namespace = %namespace,
                    "Stored prefill endpoint for future decode WorkerSet registration"
756
757
758
759
                );
                Ok(())
            }
        }
760
761
    }

762
763
764
765
766
767
768
769
770
771
772
773
    /// Deactivate the prefill router on the decode WorkerSet for the given model/namespace.
    /// Called by the watcher when all prefill workers in a namespace are removed.
    /// After deactivation, requests fall back to aggregated mode (or fail if enforce_disagg).
    pub fn deactivate_prefill_router_for_decode(&self, model_name: &str, namespace: &str) {
        if let Some(model) = self.get_model(model_name)
            && let Some(ws) = model.get_worker_set(namespace)
            && let Some(ref pr) = ws.prefill_router
        {
            pr.deactivate();
        }
    }

774
775
776
777
778
779
780
781
782
783
784
    /// 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"
            );
        }
785
786
    }

787
    // -- Worker monitoring --
788

789
    /// Gets or sets the load threshold config for a model's worker monitor.
790
    /// Checks across all WorkerSets for the model.
791
    pub fn load_threshold_config(
792
793
        &self,
        model: &str,
794
795
        config: Option<&LoadThresholdConfig>,
    ) -> Option<LoadThresholdConfig> {
796
797
        let model_entry = self.models.get(model)?;
        model_entry.load_threshold_config(config)
798
799
    }

800
801
    /// Gets an existing worker monitor for a specific namespace of a model.
    pub fn get_worker_monitor_for_namespace(
802
803
        &self,
        model: &str,
804
805
806
807
808
809
810
811
812
813
814
815
816
        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));
            }
817
        }
818
        result
819
820
    }

821
822
    // -- Runtime configs --

823
    /// Get or create a runtime config watcher for an endpoint.
824
825
    /// Spawns a background task that joins instance availability and config discovery.
    /// Returns a `watch::Receiver` with the latest `HashMap<WorkerId, ModelRuntimeConfig>`.
826
827
828
    pub async fn get_or_create_runtime_config_watcher(
        &self,
        endpoint: &Endpoint,
829
    ) -> anyhow::Result<RuntimeConfigWatch> {
830
831
832
833
834
835
        let endpoint_id = endpoint.id();

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

836
837
838
839
840
841
        // 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(),
842
            Entry::Vacant(e) => {
843
844
                e.insert(rx.clone());
                rx
845
846
847
            }
        };

848
        Ok(result)
849
850
851
852
853
854
855
856
    }

    /// Get disaggregated endpoint for a specific worker.
    pub fn get_disaggregated_endpoint(
        &self,
        endpoint_id: &EndpointId,
        worker_id: WorkerId,
    ) -> Option<DisaggregatedEndpoint> {
857
858
859
        let rx = self.runtime_configs.get(endpoint_id)?;
        let configs = rx.borrow();
        configs.get(&worker_id)?.disaggregated_endpoint.clone()
860
    }
861
}
862

863
864
865
866
867
868
869
870
871
872
873
#[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(),
        )
874
    }
875

876
    // -- CRUD delegation tests --
877

878
879
880
881
    #[test]
    fn test_add_and_get_worker_set() {
        let mm = ModelManager::new();
        let ws = make_worker_set("ns1", "abc");
882
        mm.add_worker_set("llama", "ns1", ws);
883
884
885
886
887
888

        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);
889
890
    }

891
892
893
894
895
    #[test]
    fn test_add_worker_set_creates_model() {
        let mm = ModelManager::new();
        assert!(mm.get_model("llama").is_none());

896
        mm.add_worker_set("llama", "ns1", make_worker_set("ns1", "abc"));
897
        assert!(mm.get_model("llama").is_some());
898
899
    }

900
901
902
    #[test]
    fn test_remove_worker_set_removes_empty_model() {
        let mm = ModelManager::new();
903
        mm.add_worker_set("llama", "ns1", make_worker_set("ns1", "abc"));
904
905
906
907
908
909
910
911
        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());
912
913
    }

914
915
916
    #[test]
    fn test_remove_worker_set_keeps_model_with_remaining() {
        let mm = ModelManager::new();
917
918
        mm.add_worker_set("llama", "ns1", make_worker_set("ns1", "abc"));
        mm.add_worker_set("llama", "ns2", make_worker_set("ns2", "abc"));
919
920
921
922
923
924
925
926

        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);
927
928
    }

929
930
931
932
933
934
935
936
937
    #[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();
938
        mm.add_worker_set("llama", "ns1", make_worker_set("ns1", "abc"));
939
940
941
942
943
944
945
946
947
        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();
948
        mm.add_worker_set("llama", "ns1", make_worker_set("ns1", "abc"));
949
950
951

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

954
955
956
957
    #[test]
    fn test_remove_model_if_empty_noop_when_missing() {
        let mm = ModelManager::new();
        mm.remove_model_if_empty("nonexistent"); // Should not panic
958
959
    }

960
961
962
    #[test]
    fn test_remove_model() {
        let mm = ModelManager::new();
963
964
        mm.add_worker_set("llama", "ns1", make_worker_set("ns1", "abc"));
        mm.add_worker_set("llama", "ns2", make_worker_set("ns2", "abc"));
965
966
967
968
969
970
971
972
973
974
975
976
977
978
979
980
981
982
983
984
985
986
987
988
989

        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
990
        mm.add_worker_set("llama", "ns1", make_worker_set("ns1", "abc"));
991
992
993
994
995
996
997
998
        assert!(!mm.has_decode_model("llama"));
    }

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

        // Prefill set = no engines
999
        mm.add_worker_set("llama", "ns1", make_worker_set("ns1", "abc"));
1000
1001
1002
1003
1004
1005
1006
1007
        assert!(mm.has_prefill_model("llama"));
    }

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

1008
        mm.add_worker_set("llama", "ns1", make_worker_set("ns1", "abc"));
1009
1010
1011
1012
1013
1014
        assert!(mm.has_model_any("llama")); // has prefill
    }

    #[test]
    fn test_model_display_names_includes_prefill() {
        let mm = ModelManager::new();
1015
        mm.add_worker_set("llama", "ns1", make_worker_set("ns1", "abc"));
1016
1017
1018
1019
1020
1021
1022
1023
1024
1025
1026
1027
1028
1029

        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();
1030
1031
        mm.add_worker_set("llama", "ns1", make_worker_set("ns1", "abc"));
        mm.add_worker_set("gpt", "ns1", make_worker_set("ns1", "def"));
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
1073
1074
1075
1076
1077
1078
1079
1080
1081
1082
1083
1084
1085
1086
1087
1088
1089
1090
1091
1092
1093
1094
1095
1096
1097
1098
1099
1100
1101
1102
1103
1104
1105
1106
1107
1108
1109
1110
1111
1112
1113
1114
1115
1116
1117
1118
1119
1120

        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());
1121
1122
    }

1123
1124
1125
1126
1127
    #[test]
    fn test_prefill_router_remove_nonexistent_noop() {
        let mm = ModelManager::new();
        // Should not panic
        mm.remove_prefill_activator("llama", "ns1");
1128
    }
1129

1130
1131
1132
1133
1134
1135
1136
1137
1138
1139
    #[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"
        );
1140
    }
1141
1142
1143
1144
1145
1146
1147
1148
1149
1150
1151
1152
1153
1154
1155
1156
1157
1158
1159
1160
1161
1162
1163
1164
1165
1166
1167
1168
1169
1170
1171
1172
1173
1174
1175
1176
1177
1178
1179
1180
1181
1182
1183
1184
1185
1186
1187
1188
1189
1190
1191
1192
1193
1194
1195
1196
1197
1198
1199
1200
1201
1202
1203
1204
1205
1206
1207
1208
1209
1210
1211
1212
1213
1214
1215
1216
1217
1218
1219
1220
1221
1222
1223
1224
1225
1226
1227
1228
1229
1230
1231
1232
1233
1234
1235
1236
1237
1238
1239
1240
1241
1242
1243
1244
1245
1246
1247
1248
1249
1250
1251
1252
1253
1254
1255
1256
1257
1258
1259
1260
1261
1262
1263
1264
1265
1266
1267
1268
1269
1270
1271
1272
1273
1274
1275
1276
1277
1278
1279
1280
1281
1282
1283
1284
1285
1286
1287
1288
1289
1290
1291
1292
1293
1294
1295
1296
1297
1298

    // -- deactivate_prefill_router_for_decode tests --

    use crate::kv_router::PrefillRouter;

    /// Helper: make a WorkerSet with an activated PrefillRouter attached.
    /// The router is marked as activated to simulate a real deployment where
    /// the prefill endpoint has already rendezvoused with the decode side.
    fn make_worker_set_with_prefill_router(
        namespace: &str,
        mdcsum: &str,
        enforce_disagg: bool,
    ) -> WorkerSet {
        let mut ws = make_worker_set(namespace, mdcsum);
        let pr = PrefillRouter::disabled(
            std::sync::Arc::new(ModelManager::new()),
            dynamo_runtime::pipeline::RouterMode::RoundRobin,
            enforce_disagg,
        );
        pr.mark_activated_for_test();
        ws.prefill_router = Some(pr);
        ws
    }

    /// Calling deactivate on a non-existent model must not panic.
    #[test]
    fn test_deactivate_prefill_router_for_decode_noop_missing_model() {
        let mm = ModelManager::new();
        mm.deactivate_prefill_router_for_decode("nonexistent", "ns1");
    }

    /// Calling deactivate on a WorkerSet without a prefill_router must not panic.
    #[test]
    fn test_deactivate_prefill_router_for_decode_noop_no_router() {
        let mm = ModelManager::new();
        mm.add_worker_set("llama", "ns1", make_worker_set("ns1", "abc"));
        mm.deactivate_prefill_router_for_decode("llama", "ns1");
    }

    /// Full pipeline test: deactivate finds the WorkerSet, calls deactivate() on its
    /// PrefillRouter, and the model is hidden from model_display_names() when
    /// enforce_disagg=true.
    #[test]
    fn test_deactivate_prefill_router_for_decode_hides_model() {
        let mm = ModelManager::new();
        mm.add_worker_set(
            "llama",
            "ns1",
            make_worker_set_with_prefill_router("ns1", "abc", true),
        );

        // Model is visible before deactivation.
        assert!(mm.model_display_names().contains("llama"));

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

        // Model must be hidden after deactivation with enforce_disagg=true.
        assert!(
            !mm.model_display_names().contains("llama"),
            "model must be hidden after prefill deactivation with enforce_disagg=true"
        );

        // Idempotent: calling again must not panic.
        mm.deactivate_prefill_router_for_decode("llama", "ns1");
        assert!(!mm.model_display_names().contains("llama"));
    }

    /// Full disagg lifecycle with enforce_disagg=true:
    /// decode registers -> prefill registers -> prefill dies -> model hidden.
    #[test]
    fn test_disagg_lifecycle_prefill_death_hides_model() {
        let mm = ModelManager::new();

        // Step 1: Decode WorkerSet with a PrefillRouter (not yet deactivated).
        mm.add_worker_set(
            "llama",
            "decode-ns",
            make_worker_set_with_prefill_router("decode-ns", "abc", true),
        );
        assert!(
            mm.model_display_names().contains("llama"),
            "step 1: model must be visible with active prefill router"
        );

        // Step 2: Prefill WorkerSet registers (same model, different namespace key).
        mm.add_worker_set("llama", "prefill-ns", make_worker_set("prefill-ns", "abc"));
        assert!(
            mm.model_display_names().contains("llama"),
            "step 2: model must be visible with both decode and prefill"
        );

        // Step 3: Prefill WorkerSet removed (engine dies).
        mm.remove_worker_set("llama", "prefill-ns");

        // Step 4: Deactivate the prefill router on the decode side.
        mm.deactivate_prefill_router_for_decode("llama", "decode-ns");
        assert!(
            !mm.model_display_names().contains("llama"),
            "step 4: model must be hidden after prefill death with enforce_disagg=true"
        );
    }

    /// Full disagg lifecycle with enforce_disagg=false (fallback allowed).
    #[test]
    fn test_disagg_lifecycle_prefill_death_keeps_model_no_enforce() {
        let mm = ModelManager::new();

        mm.add_worker_set(
            "llama",
            "decode-ns",
            make_worker_set_with_prefill_router("decode-ns", "abc", false),
        );
        assert!(mm.model_display_names().contains("llama"));

        // Deactivate -- model stays visible (enforce_disagg=false, fallback allowed).
        mm.deactivate_prefill_router_for_decode("llama", "decode-ns");
        assert!(
            mm.model_display_names().contains("llama"),
            "model must remain visible (enforce_disagg=false, fallback allowed)"
        );
    }

    /// Full disagg lifecycle including prefill rejoin after transient failure.
    /// decode registers -> prefill dies -> model hidden -> prefill rejoins -> model visible.
    #[test]
    fn test_disagg_lifecycle_prefill_rejoin_restores_model() {
        let mm = ModelManager::new();

        // Decode WorkerSet with enforce_disagg=true.
        mm.add_worker_set(
            "llama",
            "decode-ns",
            make_worker_set_with_prefill_router("decode-ns", "abc", true),
        );
        assert!(mm.model_display_names().contains("llama"));

        // Prefill dies -> deactivate.
        mm.deactivate_prefill_router_for_decode("llama", "decode-ns");
        assert!(
            !mm.model_display_names().contains("llama"),
            "model must be hidden after prefill death"
        );

        // Prefill rejoins -> reactivate via the WorkerSet's PrefillRouter.
        if let Some(model) = mm.get_model("llama")
            && let Some(ws) = model.get_worker_set("decode-ns")
            && let Some(ref pr) = ws.prefill_router
        {
            pr.reactivate();
        } else {
            panic!("decode WorkerSet or prefill_router not found");
        }

        assert!(
            mm.model_display_names().contains("llama"),
            "model must be visible again after prefill rejoin"
        );
    }
1299
}