## Coverage Report

#### Created: 2026-01-25 15:05

next uncovered line (L), next uncovered region (R), next uncovered branch (B)

    /home/noah/src/realizar/src/registry.rs

    Line
    Count
    Source
    1
    //! Model registry for multi-model serving
    2
    //!
    3
    //! Provides a central registry to manage multiple models in production environments.
    4
    //! Supports dynamic model loading/unloading, caching, and thread-safe concurrent access.
    5
    //!
    6
    //! ## Features
    7
    //!
    8
    //! - Thread-safe model registry with concurrent access
    9
    //! - Integration with `ModelCache` for efficient memory management
    10
    //! - Support for multiple model formats (GGUF, Safetensors)
    11
    //! - Model metadata and configuration
    12
    //! - Graceful loading/unloading with error handling
    13
    //!
    14
    //! ## Example
    15
    //!
    16
    //! ```rust,ignore
    17
    //! use realizar::registry::{ModelRegistry, ModelConfig};
    18
    //!
    19
    //! let registry = ModelRegistry::new(5); // Max 5 models cached
    20
    //! registry.register("llama-7b", model, tokenizer)?;
    21
    //! let model = registry.get("llama-7b")?;
    22
    //! ```
    23
    24
    use std::{
    25
        collections::HashMap,
    26
        sync::{Arc, Mutex},
    27
    };
    28
    29
    use arc_swap::ArcSwap;
    30
    31
    use crate::{
    32
        cache::ModelCache,
    33
        error::{RealizarError, Result},
    34
        layers::Model,
    35
        tokenizer::BPETokenizer,
    36
    };
    37
    38
    /// Information about a registered model
    39
    #[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
    40
    pub struct ModelInfo {
    41
        /// Unique model identifier
    42
        pub id: String,
    43
        /// Human-readable model name
    44
        pub name: String,
    45
        /// Model description
    46
        pub description: String,
    47
        /// Model format (GGUF, Safetensors, etc.)
    48
        pub format: String,
    49
        /// Whether the model is currently loaded
    50
        pub loaded: bool,
    51
    }
    52
    53
    /// Entry in the model registry
    54
    #[derive(Clone)]
    55
    struct ModelEntry {
    56
        /// Model instance
    57
        model: Arc<Model>,
    58
        /// Tokenizer for this model
    59
        tokenizer: Arc<BPETokenizer>,
    60
        /// Model metadata
    61
        info: ModelInfo,
    62
    }
    63
    64
    /// Type alias for the models map (immutable snapshot)
    65
    type ModelsMap = HashMap<String, ModelEntry>;
    66
    67
    /// Type alias for model and tokenizer tuple
    68
    type ModelTuple = (Arc<Model>, Arc<BPETokenizer>);
    69
    70
    /// Central registry for managing multiple models
    71
    ///
    72
    /// The `ModelRegistry` provides thread-safe access to multiple models,
    73
    /// with automatic caching and lifecycle management.
    74
    ///
    75
    /// Uses `ArcSwap` for lock-free reads (per `McKenney` 2011).
    76
    pub struct ModelRegistry {
    77
        /// Registry of loaded models (lock-free reads via `ArcSwap`)
    78
        models: ArcSwap<ModelsMap>,
    79
        /// Write lock to serialize modifications
    80
        write_lock: Mutex<()>,
    81
        /// Model cache for memory management (reserved for future use)
    82
        #[allow(dead_code)]
    83
        cache: Arc<ModelCache>,
    84
    }
    85
    86
    impl ModelRegistry {
    87
        /// Create a new model registry
    88
        ///
    89
        /// # Arguments
    90
        ///
    91
        /// * `cache_capacity` - Maximum number of models to keep in cache
    92
        #[must_use]
    93
    18
        pub fn new(cache_capacity: usize) -> Self {
    94
    18
            Self {
    95
    18
                models: ArcSwap::from_pointee(HashMap::new()),
    96
    18
                write_lock: Mutex::new(()),
    97
    18
                cache: Arc::new(ModelCache::new(cache_capacity)),
    98
    18
            }
    99
    18
        }
    100
    101
        /// Register a new model
    102
        ///
    103
        /// # Arguments
    104
        ///
    105
        /// * `id` - Unique identifier for the model
    106
        /// * `model` - Model instance
    107
        /// * `tokenizer` - Tokenizer for the model
    108
        ///
    109
        /// # Errors
    110
        ///
    111
        /// Returns error if model ID already exists
    112
    18
        pub fn register(&self, id: &str, model: Model, tokenizer: BPETokenizer) -> Result<()> {
    113
    18
            let _guard = self.write_lock.lock().map_err(|_| {0
    114
    0
                RealizarError::RegistryError("Failed to acquire write lock".to_string())
    115
    0
            })?;
    116
    117
    18
            let current = self.models.load();
    118
    18
            if current.contains_key(id) {
    119
    1
                return Err(RealizarError::ModelAlreadyExists(id.to_string()));
    120
    17
            }
    121
    122
    17
            let entry = ModelEntry {
    123
    17
                model: Arc::new(model),
    124
    17
                tokenizer: Arc::new(tokenizer),
    125
    17
                info: ModelInfo {
    126
    17
                    id: id.to_string(),
    127
    17
                    name: id.to_string(),
    128
    17
                    description: String::new(),
    129
    17
                    format: "unknown".to_string(),
    130
    17
                    loaded: true,
    131
    17
                },
    132
    17
            };
    133
    134
    17
            let mut new_map: ModelsMap = (**current).clone();
    135
    17
            new_map.insert(id.to_string(), entry);
    136
    17
            self.models.store(Arc::new(new_map));
    137
    17
            Ok(())
    138
    18
        }
    139
    140
        /// Register a model with full metadata
    141
        ///
    142
        /// # Arguments
    143
        ///
    144
        /// * `info` - Model metadata
    145
        /// * `model` - Model instance
    146
        /// * `tokenizer` - Tokenizer for the model
    147
        ///
    148
        /// # Errors
    149
        ///
    150
        /// Returns error if model ID already exists
    151
    1
        pub fn register_with_info(
    152
    1
            &self,
    153
    1
            mut info: ModelInfo,
    154
    1
            model: Model,
    155
    1
            tokenizer: BPETokenizer,
    156
    1
        ) -> Result<()> {
    157
    1
            let _guard = self.write_lock.lock().map_err(|_| {0
    158
    0
                RealizarError::RegistryError("Failed to acquire write lock".to_string())
    159
    0
            })?;
    160
    161
    1
            let current = self.models.load();
    162
    1
            if current.contains_key(&info.id) {
    163
    0
                return Err(RealizarError::ModelAlreadyExists(info.id));
    164
    1
            }
    165
    166
    1
            info.loaded = true;
    167
    1
            let entry = ModelEntry {
    168
    1
                model: Arc::new(model),
    169
    1
                tokenizer: Arc::new(tokenizer),
    170
    1
                info,
    171
    1
            };
    172
    173
    1
            let id = entry.info.id.clone();
    174
    1
            let mut new_map: ModelsMap = (**current).clone();
    175
    1
            new_map.insert(id, entry);
    176
    1
            self.models.store(Arc::new(new_map));
    177
    1
            Ok(())
    178
    1
        }
    179
    180
        /// Get a model by ID
    181
        ///
    182
        /// # Arguments
    183
        ///
    184
        /// * `id` - Model identifier
    185
        ///
    186
        /// # Errors
    187
        ///
    188
        /// Returns error if model not found
    189
    5
        pub fn get(&self, id: &str) -> Result<ModelTuple> {
    190
            // Lock-free read via ArcSwap::load()
    191
    5
            let models = self.models.load();
    192
    193
    5
            let entry4 = models
    194
    5
                .get(id)
    195
    5
                .ok_or_else(|| RealizarError::ModelNotFound(id1.to_string1()))?1;
    196
    197
    4
            Ok((Arc::clone(&entry.model), Arc::clone(&entry.tokenizer)))
    198
    5
        }
    199
    200
        /// Get model info by ID
    201
        ///
    202
        /// # Arguments
    203
        ///
    204
        /// * `id` - Model identifier
    205
        ///
    206
        /// # Errors
    207
        ///
    208
        /// Returns error if model not found
    209
    2
        pub fn get_info(&self, id: &str) -> Result<ModelInfo> {
    210
    2
            let models = self.models.load();
    211
    212
    2
            let entry1 = models
    213
    2
                .get(id)
    214
    2
                .ok_or_else(|| RealizarError::ModelNotFound(id1.to_string1()))?1;
    215
    216
    1
            Ok(entry.info.clone())
    217
    2
        }
    218
    219
        /// List all registered models (lock-free)
    220
        #[must_use]
    221
    1
        pub fn list(&self) -> Vec<ModelInfo> {
    222
    1
            let models = self.models.load();
    223
    2
            models.values()1.map1(|entry| entry.info.clone()).collect1()
    224
    1
        }
    225
    226
        /// Unregister a model
    227
        ///
    228
        /// # Arguments
    229
        ///
    230
        /// * `id` - Model identifier
    231
        ///
    232
        /// # Errors
    233
        ///
    234
        /// Returns error if model not found
    235
    2
        pub fn unregister(&self, id: &str) -> Result<()> {
    236
    2
            let _guard = self.write_lock.lock().map_err(|_| {0
    237
    0
                RealizarError::RegistryError("Failed to acquire write lock".to_string())
    238
    0
            })?;
    239
    240
    2
            let current = self.models.load();
    241
    2
            if !current.contains_key(id) {
    242
    1
                return Err(RealizarError::ModelNotFound(id.to_string()));
    243
    1
            }
    244
    245
    1
            let mut new_map: ModelsMap = (**current).clone();
    246
    1
            new_map.remove(id);
    247
    1
            self.models.store(Arc::new(new_map));
    248
    1
            Ok(())
    249
    2
        }
    250
    251
        /// Atomically replace a model (for hot-reload)
    252
        ///
    253
        /// # Arguments
    254
        ///
    255
        /// * `id` - Model identifier to replace
    256
        /// * `model` - New model instance
    257
        /// * `tokenizer` - New tokenizer instance
    258
        ///
    259
        /// # Errors
    260
        ///
    261
        /// Returns error if model not found or lock acquisition fails
    262
    2
        pub fn replace(&self, id: &str, model: Model, tokenizer: BPETokenizer) -> Result<()> {
    263
    2
            let _guard = self.write_lock.lock().map_err(|_| {0
    264
    0
                RealizarError::RegistryError("Failed to acquire write lock".to_string())
    265
    0
            })?;
    266
    267
    2
            let current = self.models.load();
    268
    2
            if !current.contains_key(id) {
    269
    1
                return Err(RealizarError::ModelNotFound(id.to_string()));
    270
    1
            }
    271
    272
            // Get existing info to preserve metadata
    273
    1
            let existing_info = current.get(id).map_or_else(
    274
                || ModelInfo {
    275
    0
                    id: id.to_string(),
    276
    0
                    name: id.to_string(),
    277
    0
                    description: String::new(),
    278
    0
                    format: "unknown".to_string(),
    279
                    loaded: true,
    280
    0
                },
    281
    1
                |e| e.info.clone(),
    282
            );
    283
    284
    1
            let entry = ModelEntry {
    285
    1
                model: Arc::new(model),
    286
    1
                tokenizer: Arc::new(tokenizer),
    287
    1
                info: existing_info,
    288
    1
            };
    289
    290
    1
            let mut new_map: ModelsMap = (**current).clone();
    291
    1
            new_map.insert(id.to_string(), entry);
    292
    1
            self.models.store(Arc::new(new_map));
    293
    1
            Ok(())
    294
    2
        }
    295
    296
        /// Check if a model is registered (lock-free)
    297
        #[must_use]
    298
    5
        pub fn contains(&self, id: &str) -> bool {
    299
    5
            let models = self.models.load();
    300
    5
            models.contains_key(id)
    301
    5
        }
    302
    303
        /// Get the number of registered models (lock-free)
    304
        #[must_use]
    305
    15
        pub fn len(&self) -> usize {
    306
    15
            let models = self.models.load();
    307
    15
            models.len()
    308
    15
        }
    309
    310
        /// Check if the registry is empty (lock-free)
    311
        #[must_use]
    312
    4
        pub fn is_empty(&self) -> bool {
    313
    4
            self.len() == 0
    314
    4
        }
    315
    }
    316
    317
    #[cfg(test)]
    318
    mod tests {
    319
        use super::*;
    320
        use crate::layers::ModelConfig;
    321
    322
    21
        fn create_test_model() -> (Model, BPETokenizer) {
    323
    21
            let config = ModelConfig {
    324
    21
                vocab_size: 100,
    325
    21
                hidden_dim: 32,
    326
    21
                num_heads: 1,
    327
    21
                num_layers: 1,
    328
    21
                intermediate_dim: 64,
    329
    21
                eps: 1e-5,
    330
    21
            };
    331
    21
            let model = Model::new(config).expect("test");
    332
    333
    21
            let vocab: Vec<String> = (0..100)
    334
    2.10k
                .map21(|i| {
    335
    2.10k
                    if i == 0 {
    336
    21
                        "<unk>".to_string()
    337
                    } else {
    338
    2.07k
                        format!("token{i}")
    339
                    }
    340
    2.10k
                })
    341
    21
                .collect();
    342
    21
            let tokenizer = BPETokenizer::new(vocab, vec![], "<unk>").expect("test");
    343
    344
    21
            (model, tokenizer)
    345
    21
        }
    346
    347
        #[test]
    348
    1
        fn test_registry_creation() {
    349
    1
            let registry = ModelRegistry::new(5);
    350
    1
            assert_eq!(registry.len(), 0);
    351
    1
            assert!(registry.is_empty());
    352
    1
        }
    353
    354
        #[test]
    355
    1
        fn test_register_model() {
    356
    1
            let registry = ModelRegistry::new(5);
    357
    1
            let (model, tokenizer) = create_test_model();
    358
    359
    1
            registry
    360
    1
                .register("test-model", model, tokenizer)
    361
    1
                .expect("test");
    362
    363
    1
            assert_eq!(registry.len(), 1);
    364
    1
            assert!(!registry.is_empty());
    365
    1
            assert!(registry.contains("test-model"));
    366
    1
        }
    367
    368
        #[test]
    369
    1
        fn test_register_duplicate_error() {
    370
    1
            let registry = ModelRegistry::new(5);
    371
    1
            let (model1, tokenizer1) = create_test_model();
    372
    1
            let (model2, tokenizer2) = create_test_model();
    373
    374
    1
            registry
    375
    1
                .register("test-model", model1, tokenizer1)
    376
    1
                .expect("test");
    377
    1
            let result = registry.register("test-model", model2, tokenizer2);
    378
    379
    1
            assert!(result.is_err());
    380
    1
            assert_eq!(registry.len(), 1);
    381
    1
        }
    382
    383
        #[test]
    384
    1
        fn test_get_model() {
    385
    1
            let registry = ModelRegistry::new(5);
    386
    1
            let (model, tokenizer) = create_test_model();
    387
    388
    1
            registry
    389
    1
                .register("test-model", model, tokenizer)
    390
    1
                .expect("test");
    391
    392
    1
            let (retrieved_model, retrieved_tokenizer) = registry.get("test-model").expect("test");
    393
    1
            assert!(Arc::strong_count(&retrieved_model) >= 2); // Registry + local
    394
    1
            assert!(Arc::strong_count(&retrieved_tokenizer) >= 2);
    395
    1
        }
    396
    397
        #[test]
    398
    1
        fn test_get_nonexistent_model() {
    399
    1
            let registry = ModelRegistry::new(5);
    400
    1
            let result = registry.get("nonexistent");
    401
    402
    1
            assert!(result.is_err());
    403
    1
        }
    404
    405
        #[test]
    406
    1
        fn test_register_with_info() {
    407
    1
            let registry = ModelRegistry::new(5);
    408
    1
            let (model, tokenizer) = create_test_model();
    409
    410
    1
            let info = ModelInfo {
    411
    1
                id: "llama-7b".to_string(),
    412
    1
                name: "Llama 7B".to_string(),
    413
    1
                description: "7B parameter Llama model".to_string(),
    414
    1
                format: "GGUF".to_string(),
    415
    1
                loaded: false,
    416
    1
            };
    417
    418
    1
            registry
    419
    1
                .register_with_info(info.clone(), model, tokenizer)
    420
    1
                .expect("test");
    421
    422
    1
            let retrieved_info = registry.get_info("llama-7b").expect("test");
    423
    1
            assert_eq!(retrieved_info.id, "llama-7b");
    424
    1
            assert_eq!(retrieved_info.name, "Llama 7B");
    425
    1
            assert_eq!(retrieved_info.description, "7B parameter Llama model");
    426
    1
            assert_eq!(retrieved_info.format, "GGUF");
    427
    1
            assert!(retrieved_info.loaded); // Should be set to true
    428
    1
        }
    429
    430
        #[test]
    431
    1
        fn test_list_models() {
    432
    1
            let registry = ModelRegistry::new(5);
    433
    1
            let (model1, tokenizer1) = create_test_model();
    434
    1
            let (model2, tokenizer2) = create_test_model();
    435
    436
    1
            registry
    437
    1
                .register("model-1", model1, tokenizer1)
    438
    1
                .expect("test");
    439
    1
            registry
    440
    1
                .register("model-2", model2, tokenizer2)
    441
    1
                .expect("test");
    442
    443
    1
            let model_list = registry.list();
    444
    1
            assert_eq!(model_list.len(), 2);
    445
    446
    2
            let ids1: Vec<String>1 = model_list.iter()1.map1(|m| m.id.clone()).collect1();
    447
    1
            assert!(ids.contains(&"model-1".to_string()));
    448
    1
            assert!(ids.contains(&"model-2".to_string()));
    449
    1
        }
    450
    451
        #[test]
    452
    1
        fn test_unregister_model() {
    453
    1
            let registry = ModelRegistry::new(5);
    454
    1
            let (model, tokenizer) = create_test_model();
    455
    456
    1
            registry
    457
    1
                .register("test-model", model, tokenizer)
    458
    1
                .expect("test");
    459
    1
            assert_eq!(registry.len(), 1);
    460
    461
    1
            registry.unregister("test-model").expect("test");
    462
    1
            assert_eq!(registry.len(), 0);
    463
    1
            assert!(!registry.contains("test-model"));
    464
    1
        }
    465
    466
        #[test]
    467
    1
        fn test_unregister_nonexistent() {
    468
    1
            let registry = ModelRegistry::new(5);
    469
    1
            let result = registry.unregister("nonexistent");
    470
    471
    1
            assert!(result.is_err());
    472
    1
        }
    473
    474
        #[test]
    475
    1
        fn test_concurrent_access() {
    476
            use std::thread;
    477
    478
    1
            let registry = Arc::new(ModelRegistry::new(10));
    479
    1
            let mut handles = vec![];
    480
    481
            // Register models from multiple threads
    482
    6
            for i5 in 0..5 {
    483
    5
                let registry_clone = Arc::clone(&registry);
    484
    5
                let handle = thread::spawn(move || {
    485
    5
                    let (model, tokenizer) = create_test_model();
    486
    5
                    registry_clone
    487
    5
                        .register(&format!("model-{i}"), model, tokenizer)
    488
    5
                        .expect("test");
    489
    5
                });
    490
    5
                handles.push(handle);
    491
            }
    492
    493
    6
            for handle5 in handles {
    494
    5
                handle.join().expect("test");
    495
    5
            }
    496
    497
    1
            assert_eq!(registry.len(), 5);
    498
    1
        }
    499
    500
        #[test]
    501
    1
        fn test_multiple_get_same_model() {
    502
    1
            let registry = ModelRegistry::new(5);
    503
    1
            let (model, tokenizer) = create_test_model();
    504
    505
    1
            registry
    506
    1
                .register("test-model", model, tokenizer)
    507
    1
                .expect("test");
    508
    509
            // Get the same model multiple times
    510
    1
            let (model1, _) = registry.get("test-model").expect("test");
    511
    1
            let (model2, _) = registry.get("test-model").expect("test");
    512
    513
            // Both should point to the same underlying model
    514
    1
            assert!(Arc::ptr_eq(&model1, &model2));
    515
    1
        }
    516
    517
        #[test]
    518
    1
        fn test_replace_model() {
    519
    1
            let registry = ModelRegistry::new(5);
    520
    1
            let (model1, tokenizer1) = create_test_model();
    521
    1
            let (model2, tokenizer2) = create_test_model();
    522
    523
            // Register initial model
    524
    1
            registry
    525
    1
                .register("test-model", model1, tokenizer1)
    526
    1
                .expect("test");
    527
    1
            assert_eq!(registry.len(), 1);
    528
    529
            // Replace with new model
    530
    1
            registry
    531
    1
                .replace("test-model", model2, tokenizer2)
    532
    1
                .expect("test");
    533
    1
            assert_eq!(registry.len(), 1);
    534
    535
            // Verify replacement worked
    536
    1
            let (retrieved, _) = registry.get("test-model").expect("test");
    537
    1
            assert!(Arc::strong_count(&retrieved) >= 2);
    538
    1
        }
    539
    540
        #[test]
    541
    1
        fn test_replace_nonexistent_model() {
    542
    1
            let registry = ModelRegistry::new(5);
    543
    1
            let (model, tokenizer) = create_test_model();
    544
    545
            // Try to replace a model that doesn't exist
    546
    1
            let result = registry.replace("nonexistent", model, tokenizer);
    547
    1
            assert!(result.is_err());
    548
    1
        }
    549
    550
        #[test]
    551
    1
        fn test_contains_method() {
    552
    1
            let registry = ModelRegistry::new(5);
    553
    1
            let (model, tokenizer) = create_test_model();
    554
    555
    1
            assert!(!registry.contains("test-model"));
    556
    1
            registry
    557
    1
                .register("test-model", model, tokenizer)
    558
    1
                .expect("test");
    559
    1
            assert!(registry.contains("test-model"));
    560
    1
        }
    561
    562
        #[test]
    563
    1
        fn test_is_empty_method() {
    564
    1
            let registry = ModelRegistry::new(5);
    565
    1
            assert!(registry.is_empty());
    566
    567
    1
            let (model, tokenizer) = create_test_model();
    568
    1
            registry
    569
    1
                .register("test-model", model, tokenizer)
    570
    1
                .expect("test");
    571
    1
            assert!(!registry.is_empty());
    572
    1
        }
    573
    574
        #[test]
    575
    1
        fn test_get_info_nonexistent() {
    576
    1
            let registry = ModelRegistry::new(5);
    577
    1
            let result = registry.get_info("nonexistent");
    578
    1
            assert!(result.is_err());
    579
    1
        }
    580
    581
        #[test]
    582
    1
        fn test_len_method() {
    583
    1
            let registry = ModelRegistry::new(5);
    584
    1
            assert_eq!(registry.len(), 0);
    585
    586
    1
            let (model1, tokenizer1) = create_test_model();
    587
    1
            registry
    588
    1
                .register("model-1", model1, tokenizer1)
    589
    1
                .expect("test");
    590
    1
            assert_eq!(registry.len(), 1);
    591
    592
    1
            let (model2, tokenizer2) = create_test_model();
    593
    1
            registry
    594
    1
                .register("model-2", model2, tokenizer2)
    595
    1
                .expect("test");
    596
    1
            assert_eq!(registry.len(), 2);
    597
    1
        }
    598
    }