neuronbox_runtime/
gpu_manager.rs

1//! In-memory session registry (soft VRAM accounting per registered PID).
2
3use std::collections::HashMap;
4use std::sync::Arc;
5
6use tokio::sync::RwLock;
7
8use crate::protocol::SessionInfo;
9
10#[derive(Default)]
11pub struct GpuManagerState {
12    sessions: HashMap<u32, SessionInfo>,
13}
14
15#[derive(Clone)]
16pub struct GpuManager {
17    inner: Arc<RwLock<GpuManagerState>>,
18}
19
20impl Default for GpuManager {
21    fn default() -> Self {
22        Self::new()
23    }
24}
25
26impl GpuManager {
27    pub fn new() -> Self {
28        Self {
29            inner: Arc::new(RwLock::new(GpuManagerState::default())),
30        }
31    }
32
33    pub async fn register(&self, info: SessionInfo) {
34        let mut g = self.inner.write().await;
35        g.sessions.insert(info.pid, info);
36    }
37
38    pub async fn unregister(&self, pid: u32) -> bool {
39        let mut g = self.inner.write().await;
40        g.sessions.remove(&pid).is_some()
41    }
42
43    pub async fn list(&self) -> Vec<SessionInfo> {
44        let g = self.inner.read().await;
45        g.sessions.values().cloned().collect()
46    }
47
48    pub async fn total_estimated_vram_mb(&self) -> u64 {
49        let g = self.inner.read().await;
50        g.sessions.values().map(|s| s.estimated_vram_mb).sum()
51    }
52}
neuronbox_runtime/gpu_manager.rs

neuronbox_runtime/
gpu_manager.rs
