Skip to main content

neuronbox_runtime/
gpu.rs

1//! GPU detection and soft VRAM control — delegates probing to [`crate::host`].
2
3use thiserror::Error;
4
5pub use crate::host::GpuRecord as GpuDevice;
6use crate::host::HostProbe;
7
8#[derive(Debug, Error)]
9pub enum GpuError {
10    #[error("nvidia-smi failed: {0}")]
11    NvidiaSmi(String),
12    #[error("could not parse GPU listing")]
13    Parse,
14}
15
16/// List GPUs. Order: NVIDIA, then AMD ROCm, then Apple Silicon.
17pub fn detect_gpus() -> Result<Vec<GpuDevice>, GpuError> {
18    Ok(HostProbe::snapshot().gpus)
19}
20
21/// Best-effort total VRAM of the primary GPU (NVIDIA index 0, or ROCm).
22pub fn primary_gpu_vram_mb() -> Option<u64> {
23    HostProbe::snapshot().primary_vram_mb()
24}
25
26/// Soft check: returns `Ok(())` if `available_mb >= required_mb`, else an error message for the user.
27pub fn soft_vram_check(
28    available_mb: u64,
29    required_mb: u64,
30    project_hint: &str,
31) -> Result<(), String> {
32    if available_mb >= required_mb {
33        return Ok(());
34    }
35    Err(format!(
36        "[neuronbox] Error: {project_hint} needs {required_mb} MB VRAM, only {available_mb} MB available (estimate).\n\
37         Try: use a smaller model, enable quantization (e.g. q4_k_m), or free other GPU processes."
38    ))
39}