ghostlink_core/
models.rs

1//! Model Registry and Asset Management
2//!
3//! This module tracks locally available model assets, their metadata,
4//! and provides hooks for loading/unloading from accelerators.
5
6use serde::{Deserialize, Serialize};
7use std::collections::HashMap;
8use std::path::{Path, PathBuf};
9use std::sync::{Arc, Mutex};
10
11/// Metadata for a locally stored model.
12#[derive(Clone, Debug, Serialize, Deserialize, PartialEq)]
13pub struct ModelMetadata {
14    pub id: String,
15    pub name: String,
16    pub architecture: String,
17    pub total_layers: usize,
18    pub parameter_count: u64,
19    pub precision: String,
20    pub vram_required_gb: f32,
21    pub local_path: PathBuf,
22}
23
24/// Status of a model in the registry.
25#[derive(Clone, Copy, Debug, Serialize, Deserialize, PartialEq, Eq)]
26pub enum ModelStatus {
27    Downloaded,
28    Loading,
29    Loaded,
30    Unloading,
31}
32
33/// A registry of model assets available on the local node.
34pub struct ModelRegistry {
35    models: Arc<Mutex<HashMap<String, ModelMetadata>>>,
36    status: Arc<Mutex<HashMap<String, ModelStatus>>>,
37    base_dir: PathBuf,
38}
39
40impl ModelRegistry {
41    /// Create a new model registry using the specified base directory.
42    pub fn new<P: AsRef<Path>>(base_dir: P) -> Self {
43        Self {
44            models: Arc::new(Mutex::new(HashMap::new())),
45            status: Arc::new(Mutex::new(HashMap::new())),
46            base_dir: base_dir.as_ref().to_path_buf(),
47        }
48    }
49
50    /// Register a model in the registry.
51    pub fn register(&self, metadata: ModelMetadata) {
52        let mut models = self.models.lock().unwrap();
53        models.insert(metadata.id.clone(), metadata);
54    }
55
56    /// Get metadata for a specific model.
57    pub fn get_metadata(&self, id: &str) -> Option<ModelMetadata> {
58        let models = self.models.lock().unwrap();
59        models.get(id).cloned()
60    }
61
62    /// List all models in the registry.
63    pub fn list_models(&self) -> Vec<ModelMetadata> {
64        let models = self.models.lock().unwrap();
65        models.values().cloned().collect()
66    }
67
68    /// Set the status of a model.
69    pub fn set_status(&self, id: &str, status: ModelStatus) {
70        let mut statuses = self.status.lock().unwrap();
71        statuses.insert(id.to_string(), status);
72    }
73
74    /// Get the status of a model.
75    pub fn get_status(&self, id: &str) -> ModelStatus {
76        let statuses = self.status.lock().unwrap();
77        statuses.get(id).copied().unwrap_or(ModelStatus::Downloaded)
78    }
79
80    /// Return the local storage base directory.
81    pub fn base_dir(&self) -> &Path {
82        &self.base_dir
83    }
84}
85
86#[cfg(test)]
87mod tests {
88    use super::*;
89    use tempfile::tempdir;
90
91    #[test]
92    fn test_model_registration() {
93        let dir = tempdir().unwrap();
94        let registry = ModelRegistry::new(dir.path());
95        let meta = ModelMetadata {
96            id: "llama3-8b".to_string(),
97            name: "Llama 3 8B".to_string(),
98            architecture: "llama".to_string(),
99            total_layers: 32,
100            parameter_count: 8_000_000_000,
101            precision: "f16".to_string(),
102            vram_required_gb: 16.0,
103            local_path: dir.path().join("llama3-8b"),
104        };
105
106        registry.register(meta.clone());
107        assert_eq!(registry.get_metadata("llama3-8b"), Some(meta));
108        assert_eq!(registry.get_status("llama3-8b"), ModelStatus::Downloaded);
109
110        registry.set_status("llama3-8b", ModelStatus::Loaded);
111        assert_eq!(registry.get_status("llama3-8b"), ModelStatus::Loaded);
112    }
113}
ghostlink_core/models.rs

ghostlink_core/
models.rs
