oxirs_embed/
model_registry.rs

1//! Model Registry and Versioning System
2//!
3//! This module provides a comprehensive model lifecycle management system including
4//! versioning, deployment, performance tracking, and A/B testing capabilities.
5
6use crate::{EmbeddingModel, ModelConfig};
7use anyhow::{anyhow, Result};
8use chrono::{DateTime, Utc};
9// Removed unused import
10use serde::{Deserialize, Serialize};
11use std::collections::HashMap;
12use std::path::PathBuf;
13use std::sync::Arc;
14use tokio::sync::RwLock;
15// Removed unused import
16// Removed unused import
17use uuid::Uuid;
18
19/// Model version information
20#[derive(Debug, Clone, Serialize, Deserialize)]
21pub struct ModelVersion {
22    pub version_id: Uuid,
23    pub model_id: Uuid,
24    pub version_number: String,
25    pub created_at: DateTime<Utc>,
26    pub created_by: String,
27    pub description: String,
28    pub tags: Vec<String>,
29    pub metrics: HashMap<String, f64>,
30    pub config: ModelConfig,
31    pub is_production: bool,
32    pub is_deprecated: bool,
33}
34
35/// Model deployment status
36#[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq)]
37pub enum DeploymentStatus {
38    NotDeployed,
39    Deploying,
40    Deployed,
41    Failed,
42    Retiring,
43    Retired,
44}
45
46/// Model deployment information
47#[derive(Debug, Clone, Serialize, Deserialize)]
48pub struct ModelDeployment {
49    pub deployment_id: Uuid,
50    pub version_id: Uuid,
51    pub status: DeploymentStatus,
52    pub deployed_at: Option<DateTime<Utc>>,
53    pub endpoint: Option<String>,
54    pub resource_allocation: ResourceAllocation,
55    pub health_check_url: Option<String>,
56    pub rollback_version: Option<Uuid>,
57}
58
59/// Resource allocation for model deployment
60#[derive(Debug, Clone, Serialize, Deserialize)]
61pub struct ResourceAllocation {
62    pub cpu_cores: f32,
63    pub memory_gb: f32,
64    pub gpu_count: u32,
65    pub gpu_memory_gb: f32,
66    pub max_concurrent_requests: usize,
67}
68
69impl Default for ResourceAllocation {
70    fn default() -> Self {
71        Self {
72            cpu_cores: 2.0,
73            memory_gb: 4.0,
74            gpu_count: 0,
75            gpu_memory_gb: 0.0,
76            max_concurrent_requests: 100,
77        }
78    }
79}
80
81/// A/B test configuration
82#[derive(Debug, Clone, Serialize, Deserialize)]
83pub struct ABTestConfig {
84    pub test_id: Uuid,
85    pub name: String,
86    pub description: String,
87    pub version_a: Uuid,
88    pub version_b: Uuid,
89    pub traffic_split: f32, // Percentage going to version B (0.0 - 1.0)
90    pub started_at: DateTime<Utc>,
91    pub ends_at: Option<DateTime<Utc>>,
92    pub metrics_to_track: Vec<String>,
93    pub is_active: bool,
94}
95
96/// Model performance metrics
97#[derive(Debug, Clone, Serialize, Deserialize)]
98pub struct PerformanceMetrics {
99    pub timestamp: DateTime<Utc>,
100    pub latency_p50_ms: f64,
101    pub latency_p95_ms: f64,
102    pub latency_p99_ms: f64,
103    pub throughput_qps: f64,
104    pub error_rate: f64,
105    pub cpu_utilization: f64,
106    pub memory_utilization: f64,
107    pub gpu_utilization: Option<f64>,
108    pub cache_hit_rate: f64,
109}
110
111/// Model registry for managing model lifecycle
112pub struct ModelRegistry {
113    models: Arc<RwLock<HashMap<Uuid, ModelMetadata>>>,
114    versions: Arc<RwLock<HashMap<Uuid, ModelVersion>>>,
115    deployments: Arc<RwLock<HashMap<Uuid, ModelDeployment>>>,
116    ab_tests: Arc<RwLock<HashMap<Uuid, ABTestConfig>>>,
117    performance_history: Arc<RwLock<HashMap<Uuid, Vec<PerformanceMetrics>>>>,
118    storage_path: PathBuf,
119}
120
121/// Model metadata
122#[derive(Debug, Clone, Serialize, Deserialize)]
123pub struct ModelMetadata {
124    pub model_id: Uuid,
125    pub name: String,
126    pub model_type: String,
127    pub created_at: DateTime<Utc>,
128    pub updated_at: DateTime<Utc>,
129    pub owner: String,
130    pub description: String,
131    pub versions: Vec<Uuid>,
132    pub production_version: Option<Uuid>,
133    pub staging_version: Option<Uuid>,
134}
135
136impl ModelRegistry {
137    /// Create a new model registry
138    pub fn new(storage_path: PathBuf) -> Self {
139        Self {
140            models: Arc::new(RwLock::new(HashMap::new())),
141            versions: Arc::new(RwLock::new(HashMap::new())),
142            deployments: Arc::new(RwLock::new(HashMap::new())),
143            ab_tests: Arc::new(RwLock::new(HashMap::new())),
144            performance_history: Arc::new(RwLock::new(HashMap::new())),
145            storage_path,
146        }
147    }
148
149    /// Base directory this registry's model artifacts are stored under.
150    pub fn storage_path(&self) -> &std::path::Path {
151        &self.storage_path
152    }
153
154    /// Register a new model
155    pub async fn register_model(
156        &self,
157        name: String,
158        model_type: String,
159        owner: String,
160        description: String,
161    ) -> Result<Uuid> {
162        let model_id = Uuid::new_v4();
163        let metadata = ModelMetadata {
164            model_id,
165            name,
166            model_type,
167            created_at: Utc::now(),
168            updated_at: Utc::now(),
169            owner,
170            description,
171            versions: Vec::new(),
172            production_version: None,
173            staging_version: None,
174        };
175
176        self.models.write().await.insert(model_id, metadata);
177        Ok(model_id)
178    }
179
180    /// Register a new model version
181    pub async fn register_version(
182        &self,
183        model_id: Uuid,
184        version_number: String,
185        created_by: String,
186        description: String,
187        config: ModelConfig,
188        metrics: HashMap<String, f64>,
189    ) -> Result<Uuid> {
190        let version_id = Uuid::new_v4();
191
192        // Verify model exists
193        let mut models = self.models.write().await;
194        let model = models
195            .get_mut(&model_id)
196            .ok_or_else(|| anyhow!("Model not found: {}", model_id))?;
197
198        let version = ModelVersion {
199            version_id,
200            model_id,
201            version_number,
202            created_at: Utc::now(),
203            created_by,
204            description,
205            tags: Vec::new(),
206            metrics,
207            config,
208            is_production: false,
209            is_deprecated: false,
210        };
211
212        model.versions.push(version_id);
213        model.updated_at = Utc::now();
214
215        self.versions.write().await.insert(version_id, version);
216        Ok(version_id)
217    }
218
219    /// Deploy a model version
220    pub async fn deploy_version(
221        &self,
222        version_id: Uuid,
223        resource_allocation: ResourceAllocation,
224    ) -> Result<Uuid> {
225        // Verify version exists
226        if !self.versions.read().await.contains_key(&version_id) {
227            return Err(anyhow!("Version not found: {}", version_id));
228        }
229
230        let deployment_id = Uuid::new_v4();
231        let deployment = ModelDeployment {
232            deployment_id,
233            version_id,
234            status: DeploymentStatus::Deploying,
235            deployed_at: None,
236            endpoint: None,
237            resource_allocation,
238            health_check_url: None,
239            rollback_version: None,
240        };
241
242        self.deployments
243            .write()
244            .await
245            .insert(deployment_id, deployment);
246
247        // Start deployment process (in real implementation)
248        self.start_deployment(deployment_id).await?;
249
250        Ok(deployment_id)
251    }
252
253    /// Start deployment process
254    async fn start_deployment(&self, deployment_id: Uuid) -> Result<()> {
255        // In a real implementation, this would:
256        // 1. Allocate resources
257        // 2. Load model weights
258        // 3. Start serving infrastructure
259        // 4. Configure load balancer
260        // 5. Run health checks
261
262        // For now, simulate deployment
263        tokio::time::sleep(tokio::time::Duration::from_secs(2)).await;
264
265        let mut deployments = self.deployments.write().await;
266        if let Some(deployment) = deployments.get_mut(&deployment_id) {
267            deployment.status = DeploymentStatus::Deployed;
268            deployment.deployed_at = Some(Utc::now());
269            deployment.endpoint = Some(format!("https://api.oxirs.ai/v1/embed/{deployment_id}"));
270            deployment.health_check_url = Some(format!(
271                "https://api.oxirs.ai/v1/embed/{deployment_id}/health"
272            ));
273        }
274
275        Ok(())
276    }
277
278    /// Promote version to production
279    pub async fn promote_to_production(&self, version_id: Uuid) -> Result<()> {
280        let versions = self.versions.read().await;
281        let version = versions
282            .get(&version_id)
283            .ok_or_else(|| anyhow!("Version not found: {}", version_id))?;
284
285        let model_id = version.model_id;
286        drop(versions);
287
288        let mut models = self.models.write().await;
289        let model = models
290            .get_mut(&model_id)
291            .ok_or_else(|| anyhow!("Model not found: {}", model_id))?;
292
293        // Mark previous production version as non-production
294        if let Some(prev_prod) = model.production_version {
295            let mut versions = self.versions.write().await;
296            if let Some(prev_version) = versions.get_mut(&prev_prod) {
297                prev_version.is_production = false;
298            }
299        }
300
301        model.production_version = Some(version_id);
302        model.updated_at = Utc::now();
303
304        let mut versions = self.versions.write().await;
305        if let Some(version) = versions.get_mut(&version_id) {
306            version.is_production = true;
307        }
308
309        Ok(())
310    }
311
312    /// Create A/B test
313    pub async fn create_ab_test(
314        &self,
315        name: String,
316        description: String,
317        version_a: Uuid,
318        version_b: Uuid,
319        traffic_split: f32,
320        duration_hours: Option<u32>,
321    ) -> Result<Uuid> {
322        // Verify both versions exist
323        let versions = self.versions.read().await;
324        if !versions.contains_key(&version_a) {
325            return Err(anyhow!("Version A not found: {}", version_a));
326        }
327        if !versions.contains_key(&version_b) {
328            return Err(anyhow!("Version B not found: {}", version_b));
329        }
330        drop(versions);
331
332        if !(0.0..=1.0).contains(&traffic_split) {
333            return Err(anyhow!("Traffic split must be between 0.0 and 1.0"));
334        }
335
336        let test_id = Uuid::new_v4();
337        let ab_test = ABTestConfig {
338            test_id,
339            name,
340            description,
341            version_a,
342            version_b,
343            traffic_split,
344            started_at: Utc::now(),
345            ends_at: duration_hours.map(|h| Utc::now() + chrono::Duration::hours(h as i64)),
346            metrics_to_track: vec![
347                "latency_p95".to_string(),
348                "accuracy".to_string(),
349                "error_rate".to_string(),
350            ],
351            is_active: true,
352        };
353
354        self.ab_tests.write().await.insert(test_id, ab_test);
355        Ok(test_id)
356    }
357
358    /// Record performance metrics
359    pub async fn record_performance(
360        &self,
361        version_id: Uuid,
362        metrics: PerformanceMetrics,
363    ) -> Result<()> {
364        let mut history = self.performance_history.write().await;
365        history
366            .entry(version_id)
367            .or_insert_with(Vec::new)
368            .push(metrics);
369
370        // Keep only last 1000 metrics per version
371        if let Some(vec) = history.get_mut(&version_id) {
372            if vec.len() > 1000 {
373                vec.drain(0..vec.len() - 1000);
374            }
375        }
376
377        Ok(())
378    }
379
380    /// Get model metadata
381    pub async fn get_model(&self, model_id: Uuid) -> Result<ModelMetadata> {
382        self.models
383            .read()
384            .await
385            .get(&model_id)
386            .cloned()
387            .ok_or_else(|| anyhow!("Model not found: {}", model_id))
388    }
389
390    /// Get version info
391    pub async fn get_version(&self, version_id: Uuid) -> Result<ModelVersion> {
392        self.versions
393            .read()
394            .await
395            .get(&version_id)
396            .cloned()
397            .ok_or_else(|| anyhow!("Version not found: {}", version_id))
398    }
399
400    /// Get deployment info
401    pub async fn get_deployment(&self, deployment_id: Uuid) -> Result<ModelDeployment> {
402        self.deployments
403            .read()
404            .await
405            .get(&deployment_id)
406            .cloned()
407            .ok_or_else(|| anyhow!("Deployment not found: {}", deployment_id))
408    }
409
410    /// Get performance history
411    pub async fn get_performance_history(
412        &self,
413        version_id: Uuid,
414        limit: Option<usize>,
415    ) -> Result<Vec<PerformanceMetrics>> {
416        let history = self.performance_history.read().await;
417        let metrics = history
418            .get(&version_id)
419            .ok_or_else(|| anyhow!("No performance history for version: {}", version_id))?;
420
421        let limit = limit.unwrap_or(100);
422        let start = metrics.len().saturating_sub(limit);
423
424        Ok(metrics[start..].to_vec())
425    }
426
427    /// Rollback deployment
428    pub async fn rollback_deployment(&self, deployment_id: Uuid) -> Result<()> {
429        let (rollback_version, resource_allocation) = {
430            let deployments = self.deployments.read().await;
431            let deployment = deployments
432                .get(&deployment_id)
433                .ok_or_else(|| anyhow!("Deployment not found: {}", deployment_id))?;
434
435            if let Some(rollback_version) = deployment.rollback_version {
436                (rollback_version, deployment.resource_allocation.clone())
437            } else {
438                return Err(anyhow!("No rollback version configured"));
439            }
440        };
441
442        // Deploy the rollback version
443        self.deploy_version(rollback_version, resource_allocation)
444            .await?;
445
446        // Mark current deployment as retired
447        let mut deployments = self.deployments.write().await;
448        if let Some(deployment) = deployments.get_mut(&deployment_id) {
449            deployment.status = DeploymentStatus::Retired;
450        }
451
452        Ok(())
453    }
454
455    /// List all models
456    pub async fn list_models(&self) -> Vec<ModelMetadata> {
457        self.models.read().await.values().cloned().collect()
458    }
459
460    /// List versions for a model
461    pub async fn list_versions(&self, model_id: Uuid) -> Result<Vec<ModelVersion>> {
462        let models = self.models.read().await;
463        let model = models
464            .get(&model_id)
465            .ok_or_else(|| anyhow!("Model not found: {}", model_id))?;
466
467        let version_ids = model.versions.clone();
468        drop(models);
469
470        let versions = self.versions.read().await;
471        let mut result = Vec::new();
472
473        for version_id in version_ids {
474            if let Some(version) = versions.get(&version_id) {
475                result.push(version.clone());
476            }
477        }
478
479        Ok(result)
480    }
481
482    /// Get active A/B tests
483    pub async fn get_active_ab_tests(&self) -> Vec<ABTestConfig> {
484        self.ab_tests
485            .read()
486            .await
487            .values()
488            .filter(|test| test.is_active)
489            .cloned()
490            .collect()
491    }
492
493    /// End A/B test
494    pub async fn end_ab_test(&self, test_id: Uuid) -> Result<()> {
495        let mut ab_tests = self.ab_tests.write().await;
496        let test = ab_tests
497            .get_mut(&test_id)
498            .ok_or_else(|| anyhow!("A/B test not found: {}", test_id))?;
499
500        test.is_active = false;
501        test.ends_at = Some(Utc::now());
502
503        Ok(())
504    }
505}
506
507/// A model, loaded into memory and shared across concurrent callers of
508/// [`ModelServer`].
509type SharedEmbeddingModel = Arc<Box<dyn EmbeddingModel>>;
510
511/// Registry of in-memory models, keyed by version ID, backing
512/// [`ModelServer::loaded_models`].
513type LoadedModelMap = Arc<RwLock<HashMap<Uuid, SharedEmbeddingModel>>>;
514
515/// Model serving infrastructure
516pub struct ModelServer {
517    registry: Arc<ModelRegistry>,
518    /// On-disk artifact repository (see [`crate::persistence::ModelRepository`]),
519    /// rooted under the owning [`ModelRegistry`]'s storage path. Artifacts are
520    /// looked up by `version_id.to_string()`; callers must persist a version's
521    /// artifact there (via [`crate::persistence::ModelRepository::save_model`])
522    /// before [`load_model`](Self::load_model) can find it.
523    repository: tokio::sync::Mutex<crate::persistence::ModelRepository>,
524    loaded_models: LoadedModelMap,
525    warm_up_cache: Arc<RwLock<HashMap<Uuid, Vec<String>>>>,
526}
527
528impl ModelServer {
529    /// Create a new model server backed by an on-disk artifact repository
530    /// rooted at `registry`'s storage path.
531    pub fn new(registry: Arc<ModelRegistry>) -> Result<Self> {
532        let repo_path = registry.storage_path().join("artifacts");
533        let repository = crate::persistence::ModelRepository::new(&repo_path)?;
534        Ok(Self {
535            registry,
536            repository: tokio::sync::Mutex::new(repository),
537            loaded_models: Arc::new(RwLock::new(HashMap::new())),
538            warm_up_cache: Arc::new(RwLock::new(HashMap::new())),
539        })
540    }
541
542    /// Load a model version into memory from the on-disk artifact repository.
543    ///
544    /// Looks up the artifact previously saved under the key
545    /// `version_id.to_string()` and caches the reconstructed model for
546    /// subsequent [`get_model`](Self::get_model) calls. Returns an error if no
547    /// artifact has been saved for this version yet.
548    pub async fn load_model(&self, version_id: Uuid) -> Result<()> {
549        let model = {
550            let repo = self.repository.lock().await;
551            repo.load_model(&version_id.to_string())?
552        };
553        self.loaded_models
554            .write()
555            .await
556            .insert(version_id, Arc::new(model));
557        Ok(())
558    }
559
560    /// Warm up model with sample inputs.
561    ///
562    /// If the model is already loaded, actually runs `encode` over the
563    /// samples once so its internal caches are genuinely primed; otherwise
564    /// the samples are recorded for a future load.
565    pub async fn warm_up_model(&self, version_id: Uuid, samples: Vec<String>) -> Result<()> {
566        self.warm_up_cache
567            .write()
568            .await
569            .insert(version_id, samples.clone());
570
571        if let Some(model) = self.loaded_models.read().await.get(&version_id).cloned() {
572            model.encode(&samples).await?;
573        }
574
575        Ok(())
576    }
577
578    /// Get a previously loaded model for inference. Returns an error if
579    /// [`load_model`](Self::load_model) has not been called successfully for
580    /// this version.
581    pub async fn get_model(&self, version_id: Uuid) -> Result<SharedEmbeddingModel> {
582        self.loaded_models
583            .read()
584            .await
585            .get(&version_id)
586            .cloned()
587            .ok_or_else(|| {
588                anyhow!("Model version {version_id} is not loaded; call load_model first")
589            })
590    }
591
592    /// Route request based on A/B test
593    pub async fn route_request(&self, test_id: Uuid) -> Result<Uuid> {
594        let ab_tests = self.registry.ab_tests.read().await;
595        let test = ab_tests
596            .get(&test_id)
597            .ok_or_else(|| anyhow!("A/B test not found: {}", test_id))?;
598
599        // Simple random routing based on traffic split
600        let random = {
601            use scirs2_core::random::{Random, RngExt};
602            let mut random = Random::default();
603            random.random::<f32>()
604        };
605        Ok(if random < test.traffic_split {
606            test.version_b
607        } else {
608            test.version_a
609        })
610    }
611}
612
613#[cfg(test)]
614mod tests {
615    use super::*;
616    use tempfile::tempdir;
617
618    #[tokio::test]
619    async fn test_model_registry_lifecycle() {
620        let temp_dir = tempdir().expect("should succeed");
621        let registry = ModelRegistry::new(temp_dir.path().to_path_buf());
622
623        // Register model
624        let model_id = registry
625            .register_model(
626                "test-model".to_string(),
627                "TransformerEmbedding".to_string(),
628                "test-user".to_string(),
629                "Test model".to_string(),
630            )
631            .await
632            .expect("should succeed");
633
634        // Register version
635        let config = ModelConfig::default();
636        let mut metrics = HashMap::new();
637        metrics.insert("accuracy".to_string(), 0.95);
638
639        let version_id = registry
640            .register_version(
641                model_id,
642                "1.0.0".to_string(),
643                "test-user".to_string(),
644                "Initial version".to_string(),
645                config,
646                metrics,
647            )
648            .await
649            .expect("should succeed");
650
651        // Deploy version
652        let deployment_id = registry
653            .deploy_version(version_id, ResourceAllocation::default())
654            .await
655            .expect("should succeed");
656
657        // Wait for deployment
658        tokio::time::sleep(tokio::time::Duration::from_secs(3)).await;
659
660        // Check deployment status
661        let deployment = registry
662            .get_deployment(deployment_id)
663            .await
664            .expect("should succeed");
665        assert_eq!(deployment.status, DeploymentStatus::Deployed);
666        assert!(deployment.endpoint.is_some());
667
668        // Promote to production
669        registry
670            .promote_to_production(version_id)
671            .await
672            .expect("should succeed");
673
674        let model = registry.get_model(model_id).await.expect("should succeed");
675        assert_eq!(model.production_version, Some(version_id));
676    }
677
678    #[tokio::test]
679    async fn test_ab_testing() {
680        let temp_dir = tempdir().expect("should succeed");
681        let registry = ModelRegistry::new(temp_dir.path().to_path_buf());
682
683        // Register model and two versions
684        let model_id = registry
685            .register_model(
686                "ab-test-model".to_string(),
687                "GNNEmbedding".to_string(),
688                "test-user".to_string(),
689                "AB test model".to_string(),
690            )
691            .await
692            .expect("should succeed");
693
694        let version_a = registry
695            .register_version(
696                model_id,
697                "1.0.0".to_string(),
698                "test-user".to_string(),
699                "Version A".to_string(),
700                ModelConfig::default(),
701                HashMap::new(),
702            )
703            .await
704            .expect("should succeed");
705
706        let version_b = registry
707            .register_version(
708                model_id,
709                "1.1.0".to_string(),
710                "test-user".to_string(),
711                "Version B".to_string(),
712                ModelConfig::default(),
713                HashMap::new(),
714            )
715            .await
716            .expect("should succeed");
717
718        // Create A/B test
719        let test_id = registry
720            .create_ab_test(
721                "Performance test".to_string(),
722                "Testing new model version".to_string(),
723                version_a,
724                version_b,
725                0.3,      // 30% traffic to version B
726                Some(24), // 24 hour test
727            )
728            .await
729            .expect("should succeed");
730
731        // Check active tests
732        let active_tests = registry.get_active_ab_tests().await;
733        assert_eq!(active_tests.len(), 1);
734        assert_eq!(active_tests[0].test_id, test_id);
735
736        // End test
737        registry.end_ab_test(test_id).await.expect("should succeed");
738
739        let active_tests = registry.get_active_ab_tests().await;
740        assert_eq!(active_tests.len(), 0);
741    }
742
743    #[tokio::test]
744    async fn test_model_server_get_model_before_load_errors() {
745        let temp_dir = tempdir().expect("should succeed");
746        let registry = Arc::new(ModelRegistry::new(temp_dir.path().to_path_buf()));
747        let server = ModelServer::new(registry).expect("server should construct");
748
749        let result = server.get_model(Uuid::new_v4()).await;
750        assert!(result.is_err(), "expected error before load_model");
751    }
752
753    #[tokio::test]
754    async fn test_model_server_load_model_missing_artifact_errors() {
755        let temp_dir = tempdir().expect("should succeed");
756        let registry = Arc::new(ModelRegistry::new(temp_dir.path().to_path_buf()));
757        let server = ModelServer::new(registry).expect("server should construct");
758
759        let result = server.load_model(Uuid::new_v4()).await;
760        assert!(
761            result.is_err(),
762            "expected an error when no artifact has been saved for this version"
763        );
764    }
765
766    #[tokio::test]
767    async fn test_model_server_load_and_get_model_round_trip() {
768        let temp_dir = tempdir().expect("should succeed");
769        let registry_path = temp_dir.path().to_path_buf();
770        let version_id = Uuid::new_v4();
771
772        // Seed the artifact repository directly (as a real persistence layer
773        // would after training) before the server ever starts.
774        let artifacts_path = registry_path.join("artifacts");
775        let mut seed_repo =
776            crate::persistence::ModelRepository::new(&artifacts_path).expect("repo should open");
777        let model = crate::models::TransE::new(ModelConfig::default());
778        seed_repo
779            .save_model(&model, &version_id.to_string(), None)
780            .expect("save should succeed");
781
782        let registry = Arc::new(ModelRegistry::new(registry_path));
783        let server = ModelServer::new(registry).expect("server should construct");
784
785        server
786            .load_model(version_id)
787            .await
788            .expect("load should succeed");
789        let loaded = server
790            .get_model(version_id)
791            .await
792            .expect("get should succeed");
793        assert_eq!(loaded.model_type(), "TransE");
794    }
795}
oxirs_embed/model_registry.rs

oxirs_embed/
model_registry.rs
