neuronbox-runtime 0.1.0

Local ML runner: declarative neuron.yaml, model store, daemon, and terminal dashboard
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

//! Tracks logical active model for swap / display (weights are loaded by Python / user code, not in Rust).

use std::sync::Arc;

use tokio::sync::RwLock;

#[derive(Clone, Default)]
pub struct ActiveModel {
    pub model_ref: String,
    pub quantization: Option<String>,
}

#[derive(Clone)]
pub struct ModelLoader {
    active: Arc<RwLock<ActiveModel>>,
}

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

impl ModelLoader {
    pub fn new() -> Self {
        Self {
            active: Arc::new(RwLock::new(ActiveModel::default())),
        }
    }

    pub async fn swap(&self, model_ref: String, quantization: Option<String>) {
        let mut a = self.active.write().await;
        a.model_ref = model_ref;
        a.quantization = quantization;
    }

    pub async fn get(&self) -> ActiveModel {
        self.active.read().await.clone()
    }
}