mistralrs_core/utils/
memory_usage.rs

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
50
51
52
53
54
55
56
57
58
use candle_core::{Device, Result};
use sysinfo::System;

const KB_TO_BYTES: usize = 1024;

pub struct MemoryUsage;

impl MemoryUsage {
    /// Amount of available memory in bytes.
    pub fn get_memory_available(&self, device: &Device) -> Result<usize> {
        match device {
            Device::Cpu => {
                let mut sys = System::new_all();
                sys.refresh_cpu();
                Ok(usize::try_from(sys.free_memory())? * KB_TO_BYTES)
            }
            #[cfg(feature = "cuda")]
            Device::Cuda(_) => {
                use candle_core::cuda_backend::WrapErr;
                Ok(candle_core::cuda::cudarc::driver::result::mem_get_info()
                    .w()?
                    .0)
            }
            #[cfg(not(feature = "cuda"))]
            Device::Cuda(_) => {
                candle_core::bail!("Cannot get memory available for CUDA device")
            }
            Device::Metal(_) => {
                candle_core::bail!("Cannot get memory available for Metal device")
            }
        }
    }

    /// Amount of total memory in bytes.
    pub fn get_total_memory(&self, device: &Device) -> Result<usize> {
        match device {
            Device::Cpu => {
                let mut sys = System::new_all();
                sys.refresh_cpu();
                Ok(usize::try_from(sys.total_memory())? * KB_TO_BYTES)
            }
            #[cfg(feature = "cuda")]
            Device::Cuda(_) => {
                use candle_core::cuda_backend::WrapErr;
                Ok(candle_core::cuda::cudarc::driver::result::mem_get_info()
                    .w()?
                    .1)
            }
            #[cfg(not(feature = "cuda"))]
            Device::Cuda(_) => {
                candle_core::bail!("Cannot get total memory for CUDA device")
            }
            Device::Metal(_) => {
                candle_core::bail!("Cannot get total memory for Metal device")
            }
        }
    }
}