_lib/
lib.rs

1// lib.rs : Final script where to place wrapped functionality
2
3use pyo3::prelude::*;
4use std::env;
5
6use pyo3::exceptions::PyRuntimeError;
7use nvml_wrapper::error::NvmlError;
8use nvml_wrapper::{cuda_driver_version_major, cuda_driver_version_minor, Nvml};
9use pretty_bytes::converter::convert;
10
11#[pyfunction]
12fn device_info() -> PyResult<String>{
13// Initialize NVML
14    let nvml = Nvml::init().map_err(|e| PyRuntimeError::new_err(format!("NVML Init Error: {}", e)))?;
15
16    // Get the version of the cuda driver
17    let cuda_version = nvml.sys_cuda_driver_version().map_err(|e| PyRuntimeError::new_err(format!("CUDA Version Error: {}", e)))?;
18
19    // Index the first detected device
20    let device = nvml.device_by_index(0).map_err(|e| PyRuntimeError::new_err(format!("Device Index Error: {}", e)))?;
21
22    // Get general device information
23    let name = device.name().map_err(|e| PyRuntimeError::new_err(format!("Device Name Error: {}", e)))?;
24
25    Ok(name)
26}
27
28/// Formats the sum of two numbers as string.
29#[pyfunction]
30fn sum_as_string(a: usize, b: usize) -> PyResult<String> {
31    Ok((a + b).to_string())
32}
33
34// Substract two numbers and returns the result as a string
35#[pyfunction]
36fn sub_as_string(a: usize, b: usize) -> PyResult<String> {
37    Ok((a - b).to_string())
38}
39
40
41/// A Python module implemented in Rust. 
42/// The name of this function must match
43/// the `lib.name` setting in the `Cargo.toml`, 
44/// else Python will not be able to import the module.
45
46#[pymodule]
47fn _lib(_py: Python<'_>, m: &PyModule) -> PyResult<()> {
48    m.add_function(wrap_pyfunction!(sum_as_string, m)?)?;
49    m.add_function(wrap_pyfunction!(sub_as_string, m)?)?;
50    m.add_function(wrap_pyfunction!(device_info, m)?)?;
51    Ok(())
52}
_lib/lib.rs

_lib/
lib.rs
