mistralrs_core/utils/
memory_usage.rs

1use candle_core::{Device, Result};
2use sysinfo::System;
3
4pub struct MemoryUsage;
5
6/// Returns the memory fraction to use for integrated CUDA GPUs.
7/// Defaults to 0.75, configurable via MISTRALRS_IGPU_MEMORY_FRACTION.
8#[cfg(feature = "cuda")]
9fn igpu_memory_fraction() -> f64 {
10    std::env::var("MISTRALRS_IGPU_MEMORY_FRACTION")
11        .ok()
12        .and_then(|s| s.parse::<f64>().ok())
13        .and_then(|f| {
14            if (0.0..=1.0).contains(&f) {
15                Some(f)
16            } else {
17                None
18            }
19        })
20        .unwrap_or(0.75)
21}
22
23impl MemoryUsage {
24    /// Amount of available memory in bytes.
25    #[allow(clippy::cast_possible_truncation, clippy::cast_precision_loss)]
26    pub fn get_memory_available(&self, device: &Device) -> Result<usize> {
27        match device {
28            Device::Cpu => {
29                let mut sys = System::new_all();
30                sys.refresh_cpu_all();
31                Ok(usize::try_from(sys.available_memory())?)
32            }
33            #[cfg(feature = "cuda")]
34            Device::Cuda(dev) => {
35                if super::normal::is_integrated_gpu(device) {
36                    // For integrated GPUs with unified memory, use system memory
37                    // scaled by a configurable fraction (default 75%)
38                    let mut sys = System::new_all();
39                    sys.refresh_cpu_all();
40                    let avail = usize::try_from(sys.available_memory())?;
41                    let fraction = igpu_memory_fraction();
42                    Ok((avail as f64 * fraction) as usize)
43                } else {
44                    use candle_core::cuda::cudarc::driver::result;
45                    use candle_core::cuda_backend::WrapErr;
46
47                    dev.cuda_stream().context().bind_to_thread().w()?;
48                    let (free, _total) = result::mem_get_info().w()?;
49                    Ok(free)
50                }
51            }
52            #[cfg(not(feature = "cuda"))]
53            Device::Cuda(_) => {
54                candle_core::bail!("Cannot get memory available for CUDA device")
55            }
56            #[cfg(feature = "metal")]
57            Device::Metal(dev) => {
58                let max = dev.device().recommended_max_working_set_size();
59                let alloc = dev.current_allocated_size();
60                let avail = max.saturating_sub(alloc);
61
62                #[allow(clippy::cast_possible_truncation)]
63                Ok(avail)
64            }
65            #[cfg(not(feature = "metal"))]
66            Device::Metal(_) => {
67                candle_core::bail!("Cannot get memory available for Metal device")
68            }
69        }
70    }
71
72    /// Amount of total memory in bytes.
73    #[allow(clippy::cast_possible_truncation, clippy::cast_precision_loss)]
74    pub fn get_total_memory(&self, device: &Device) -> Result<usize> {
75        match device {
76            Device::Cpu => {
77                let mut sys = System::new_all();
78                sys.refresh_cpu_all();
79                Ok(usize::try_from(sys.total_memory())?)
80            }
81            #[cfg(feature = "cuda")]
82            Device::Cuda(dev) => {
83                if super::normal::is_integrated_gpu(device) {
84                    // For integrated GPUs with unified memory, use system total memory
85                    // scaled by a configurable fraction (default 75%)
86                    let mut sys = System::new_all();
87                    sys.refresh_cpu_all();
88                    let total = usize::try_from(sys.total_memory())?;
89                    let fraction = igpu_memory_fraction();
90                    Ok((total as f64 * fraction) as usize)
91                } else {
92                    use candle_core::cuda::cudarc::driver::result;
93                    use candle_core::cuda_backend::WrapErr;
94
95                    dev.cuda_stream().context().bind_to_thread().w()?;
96                    let (_free, total) = result::mem_get_info().w()?;
97                    Ok(total)
98                }
99            }
100            #[cfg(not(feature = "cuda"))]
101            Device::Cuda(_) => {
102                candle_core::bail!("Cannot get total memory for CUDA device")
103            }
104            #[cfg(feature = "metal")]
105            #[allow(clippy::cast_possible_truncation)]
106            Device::Metal(dev) => {
107                const SIZE_IN_MB: usize = 1024 * 1024;
108
109                // Get system RAM in MB
110                let system_ram_mb = {
111                    let mut sys = System::new_all();
112                    sys.refresh_cpu_all();
113                    usize::try_from(sys.total_memory())? / SIZE_IN_MB
114                };
115
116                // Check for Metal GPU wired limit
117                let metal_cap_mb = std::process::Command::new("sysctl")
118                    .arg("-n")
119                    .arg("iogpu.wired_limit_mb")
120                    .output()
121                    .ok()
122                    .and_then(|o| String::from_utf8(o.stdout).ok())
123                    .and_then(|s| s.trim().parse::<usize>().ok());
124
125                // Apply default cap based on system RAM if not set or 0
126                let default_cap = match system_ram_mb {
127                    x if x <= 36 * 1024 => (system_ram_mb * 2) / 3,
128                    x if x > 36 * 1024 => (system_ram_mb * 3) / 4,
129                    x => {
130                        return Err(candle_core::Error::Msg(format!(
131                            "Invalid system ram mb value {x}."
132                        )))
133                    }
134                };
135
136                let metal_cap_mb = match metal_cap_mb {
137                    Some(0) => default_cap,
138                    Some(x) => x,
139                    None => default_cap,
140                };
141
142                let device_max = dev.recommended_max_working_set_size();
143                let metal_cap_bytes = metal_cap_mb * SIZE_IN_MB;
144
145                Ok(device_max.min(metal_cap_bytes))
146            }
147            #[cfg(not(feature = "metal"))]
148            Device::Metal(_) => {
149                candle_core::bail!("Cannot get memory available for Metal device")
150            }
151        }
152    }
153}
mistralrs_core/utils/memory_usage.rs

mistralrs_core/utils/
memory_usage.rs
