singe_nvml/
lib.rs

1//! Safe NVIDIA Management Library wrappers for monitoring, device management,
2//! MIG, and vGPU APIs.
3//!
4//! This crate wraps NVML initialization and shutdown, system queries, device and
5//! unit handles, MIG GPU and compute instances, vGPU types and instances, event
6//! sets, clocks, power, thermals, memory, utilization, and other management
7//! interfaces exposed by NVIDIA drivers.
8//!
9//! # Setup
10//!
11//! NVML is supplied by the NVIDIA driver rather than the CUDA Toolkit. The build
12//! and runtime environment must be able to locate the driver NVML library, for
13//! example through the platform library path.
14//!
15//! # Example
16//!
17//! ```
18//! use singe_nvml::library::Library;
19//!
20//! let nvml = Library::create()?;
21//! let count = nvml.device_count()?;
22//!
23//! for index in 0..count {
24//!     let device = nvml.device(index)?;
25//!     let memory = device.memory_info()?;
26//!     println!("{}: {} MiB free", device.name()?, memory.free / 1_048_576);
27//! }
28//!
29//! println!("NVML version: {}", singe_nvml::version()?);
30//! # Ok::<(), singe_nvml::error::Error>(())
31//! ```
32
33pub mod compute_instance;
34pub mod device;
35pub mod error;
36pub mod gpu_instance;
37pub mod library;
38pub mod types;
39pub mod unit;
40pub mod vgpu_instance;
41pub mod vgpu_type;
42
43pub(crate) mod utility;
44
45use singe_core::string_from_c_chars;
46use singe_nvml_sys as sys;
47
48use crate::error::Result;
49
50/// Returns the version of the NVML library.
51///
52/// The version identifier is an alphanumeric string and does not require an
53/// initialized [`library::Library`] handle.
54///
55/// # Errors
56///
57/// Returns an error if the internal version buffer is too small or if NVML
58/// rejects the output argument.
59pub fn version() -> Result<String> {
60    let mut buffer = [0i8; sys::NVML_SYSTEM_NVML_VERSION_BUFFER_SIZE as usize];
61    unsafe {
62        try_ffi!(sys::nvmlSystemGetNVMLVersion(
63            buffer.as_mut_ptr(),
64            buffer.len() as u32,
65        ))?;
66    }
67    Ok(string_from_c_chars(&buffer))
68}
69
70#[cfg(test)]
71mod tests {
72    use crate::{
73        error::{Error, Status},
74        library::Library,
75        types::{InitFlags, SystemEventTypes},
76    };
77
78    #[test]
79    fn test_system_info_is_available_when_nvml_initializes() {
80        match Library::create_with_flags(InitFlags::NO_GPUS) {
81            Ok(nvml) => {
82                assert!(!crate::version().unwrap().is_empty());
83                assert!(!nvml.driver_version().unwrap().is_empty());
84                assert_ne!(nvml.cuda_driver_version().unwrap().raw, 0);
85                assert_ne!(nvml.cuda_driver_version_fallback().unwrap().raw, 0);
86                let _ = nvml.vgpu_version();
87                let _ = nvml.excluded_devices();
88                let _ = nvml.system_event_set().and_then(|set| {
89                    let _ = set.register_events(
90                        SystemEventTypes::GPU_DRIVER_BIND | SystemEventTypes::GPU_DRIVER_UNBIND,
91                    );
92                    match set.wait(0, 4) {
93                        Ok(_) => Ok(()),
94                        Err(Error::Nvml { code, .. }) if code == Status::Timeout => Ok(()),
95                        Err(error) => Err(error),
96                    }
97                });
98            }
99            Err(error) => eprintln!("error initializing nvml: {error:?}"),
100        }
101    }
102}
singe_nvml/lib.rs

singe_nvml/
lib.rs
