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
41
42
43
44
45
46
47
48
49

//! Soft monitoring: if VRAM usage reported by `nvidia-smi` for a PID exceeds
//! ~115% of the registered estimate, send SIGKILL (Linux only).

use std::process::Command;
use std::time::Duration;

use tokio::time::sleep;

use crate::gpu_manager::GpuManager;
use crate::host::compute_apps_pid_memory_mb;

pub async fn run_soft_vram_enforcement(gm: GpuManager) {
    loop {
        sleep(Duration::from_secs(5)).await;
        let usage =
            match tokio::task::spawn_blocking(|| compute_apps_pid_memory_mb().ok_or(())).await {
                Ok(Ok(m)) => m,
                _ => continue,
            };
        let sessions = gm.list().await;
        for s in sessions {
            let Some(used_mb) = usage.get(&s.pid).copied() else {
                continue;
            };
            let limit = s.estimated_vram_mb.saturating_mul(115) / 100;
            if used_mb > limit && limit > 0 {
                tracing::warn!(
                    "VRAM soft limit: pid {} ({}) used {} MiB > {} MiB — SIGKILL",
                    s.pid,
                    s.name,
                    used_mb,
                    limit
                );
                kill_pid_hard(s.pid);
                gm.unregister(s.pid).await;
            }
        }
    }
}

#[cfg(unix)]
fn kill_pid_hard(pid: u32) {
    let _ = Command::new("kill")
        .args(["-KILL", &pid.to_string()])
        .status();
}

#[cfg(not(unix))]
fn kill_pid_hard(_pid: u32) {}