ghostlink-core 1.17.0

Core runtime primitives for Ghostlink, including planning, routing, health monitoring, and transport logic for distributed inference workloads.
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113

//! Model Registry and Asset Management
//!
//! This module tracks locally available model assets, their metadata,
//! and provides hooks for loading/unloading from accelerators.

use serde::{Deserialize, Serialize};
use std::collections::HashMap;
use std::path::{Path, PathBuf};
use std::sync::{Arc, Mutex};

/// Metadata for a locally stored model.
#[derive(Clone, Debug, Serialize, Deserialize, PartialEq)]
pub struct ModelMetadata {
    pub id: String,
    pub name: String,
    pub architecture: String,
    pub total_layers: usize,
    pub parameter_count: u64,
    pub precision: String,
    pub vram_required_gb: f32,
    pub local_path: PathBuf,
}

/// Status of a model in the registry.
#[derive(Clone, Copy, Debug, Serialize, Deserialize, PartialEq, Eq)]
pub enum ModelStatus {
    Downloaded,
    Loading,
    Loaded,
    Unloading,
}

/// A registry of model assets available on the local node.
pub struct ModelRegistry {
    models: Arc<Mutex<HashMap<String, ModelMetadata>>>,
    status: Arc<Mutex<HashMap<String, ModelStatus>>>,
    base_dir: PathBuf,
}

impl ModelRegistry {
    /// Create a new model registry using the specified base directory.
    pub fn new<P: AsRef<Path>>(base_dir: P) -> Self {
        Self {
            models: Arc::new(Mutex::new(HashMap::new())),
            status: Arc::new(Mutex::new(HashMap::new())),
            base_dir: base_dir.as_ref().to_path_buf(),
        }
    }

    /// Register a model in the registry.
    pub fn register(&self, metadata: ModelMetadata) {
        let mut models = self.models.lock().unwrap();
        models.insert(metadata.id.clone(), metadata);
    }

    /// Get metadata for a specific model.
    pub fn get_metadata(&self, id: &str) -> Option<ModelMetadata> {
        let models = self.models.lock().unwrap();
        models.get(id).cloned()
    }

    /// List all models in the registry.
    pub fn list_models(&self) -> Vec<ModelMetadata> {
        let models = self.models.lock().unwrap();
        models.values().cloned().collect()
    }

    /// Set the status of a model.
    pub fn set_status(&self, id: &str, status: ModelStatus) {
        let mut statuses = self.status.lock().unwrap();
        statuses.insert(id.to_string(), status);
    }

    /// Get the status of a model.
    pub fn get_status(&self, id: &str) -> ModelStatus {
        let statuses = self.status.lock().unwrap();
        statuses.get(id).copied().unwrap_or(ModelStatus::Downloaded)
    }

    /// Return the local storage base directory.
    pub fn base_dir(&self) -> &Path {
        &self.base_dir
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use tempfile::tempdir;

    #[test]
    fn test_model_registration() {
        let dir = tempdir().unwrap();
        let registry = ModelRegistry::new(dir.path());
        let meta = ModelMetadata {
            id: "llama3-8b".to_string(),
            name: "Llama 3 8B".to_string(),
            architecture: "llama".to_string(),
            total_layers: 32,
            parameter_count: 8_000_000_000,
            precision: "f16".to_string(),
            vram_required_gb: 16.0,
            local_path: dir.path().join("llama3-8b"),
        };

        registry.register(meta.clone());
        assert_eq!(registry.get_metadata("llama3-8b"), Some(meta));
        assert_eq!(registry.get_status("llama3-8b"), ModelStatus::Downloaded);

        registry.set_status("llama3-8b", ModelStatus::Loaded);
        assert_eq!(registry.get_status("llama3-8b"), ModelStatus::Loaded);
    }
}