redvault-ai/llama_proxy_man/src/state.rs
2025-02-20 02:13:21 +01:00

38 lines
1.1 KiB
Rust

use std::{collections::HashMap, sync::Arc};
use tokio::sync::Mutex;
use crate::{config::AppConfig, inference_process::InferenceProcess, util::parse_size};
#[derive(Clone, Debug)]
pub struct ResourceManager {
pub total_ram: u64,
pub total_vram: u64,
pub used_ram: u64,
pub used_vram: u64,
pub processes: HashMap<u16, InferenceProcess>,
}
pub type ResourceManagerHandle = Arc<Mutex<ResourceManager>>;
#[derive(Clone, Debug, derive_more::Deref)]
pub struct AppState(pub ResourceManagerHandle);
impl AppState {
pub fn from_config(config: AppConfig) -> Self {
let total_ram =
parse_size(&config.system_resources.ram).expect("Invalid RAM size in config");
let total_vram =
parse_size(&config.system_resources.vram).expect("Invalid VRAM size in config");
let resource_manager = ResourceManager {
total_ram,
total_vram,
used_ram: 0,
used_vram: 0,
processes: HashMap::new(),
};
Self(Arc::new(Mutex::new(resource_manager)))
}
}