metrale_storage/
probe.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-09-26: Storage capability probe (`run_probe`): loads cuFile and opens its driver
4//! when it can, checks for the nvidia-fs module, runs the `bench` read benchmarks
5//! (cuFile's only when its driver opened), and recommends a `Backend` from the results
6//! (`decide`). A benchmark that fails is logged and left out of the result.
7//!
8//! Owner: storage.
9//! Invariants: none beyond the types.
10
11use anyhow::{Context, Result};
12use serde::Serialize;
13use std::fs::File;
14use std::io::Write;
15use std::os::fd::AsRawFd;
16use std::path::{Path, PathBuf};
17
18use crate::bench::{
19    BenchResult, RAND_IO_BYTES, RAND_ITERS, SEQ_IO_BYTES, SEQ_ITERS, bench_cufile, bench_io_uring,
20    bench_posix, open_test_file,
21};
22use crate::cuda_min::{CudaCtx, DeviceBuffer, PinnedBuffer};
23
24#[derive(Debug, Clone, Copy, Serialize, PartialEq, Eq)]
25pub enum Backend {
26    CuFileDirect,
27    CuFileCompat,
28    IoUring,
29    PosixOnly,
30    None,
31}
32
33#[derive(Debug, Clone)]
34pub struct ProbeConfig {
35    pub dir: PathBuf,
36    pub test_file_bytes: u64,
37}
38
39#[derive(Debug, Clone, Serialize)]
40pub struct ProbeResult {
41    pub libcufile_loaded: bool,
42    pub libcufile_load_error: Option<String>,
43    pub cufile_version: Option<i32>,
44    pub nvidia_fs_kmod_loaded: bool,
45    pub cufile_driver_open_ok: bool,
46    pub cufile_driver_open_error: Option<String>,
47    pub cufile_seq_4mib: Option<f64>,
48    pub cufile_rand_64kib: Option<f64>,
49    pub io_uring_seq_4mib: Option<f64>,
50    pub io_uring_rand_64kib: Option<f64>,
51    pub posix_seq_4mib: Option<f64>,
52    pub posix_rand_64kib: Option<f64>,
53    pub recommended: Backend,
54    pub recommendation_reason: String,
55}
56
57fn fill_test_file(path: &Path, bytes: u64) -> Result<()> {
58    if path.exists() && std::fs::metadata(path)?.len() == bytes {
59        return Ok(());
60    }
61    let mut f = File::create(path).with_context(|| format!("create {}", path.display()))?;
62    let chunk = vec![0xA5u8; 1 << 20];
63    let mut written = 0u64;
64    while written < bytes {
65        let n = ((bytes - written) as usize).min(chunk.len());
66        f.write_all(&chunk[..n])?;
67        written += n as u64;
68    }
69    f.sync_all()?;
70    Ok(())
71}
72
73fn run_one_bench<F>(label: &str, f: F) -> Option<f64>
74where
75    F: FnOnce() -> Result<BenchResult>,
76{
77    match f() {
78        Ok(r) => {
79            tracing::info!(
80                "{label}: {:.1} MiB/s ({} iters @ {} bytes)",
81                r.mib_per_sec,
82                r.iters,
83                r.bytes_per_io
84            );
85            Some(r.mib_per_sec)
86        }
87        Err(e) => {
88            tracing::warn!("{label}: SKIPPED ({e:#})");
89            None
90        }
91    }
92}
93
94pub fn run_probe(cfg: &ProbeConfig) -> Result<ProbeResult> {
95    std::fs::create_dir_all(&cfg.dir).with_context(|| format!("mkdir {}", cfg.dir.display()))?;
96    let test_path = cfg.dir.join("probe-test.bin");
97    fill_test_file(&test_path, cfg.test_file_bytes)?;
98
99    let nvidia_fs = metrale_gpu_sys::cufile::nvidia_fs_loaded();
100    let cufile_load = metrale_gpu_sys::cufile::CuFile::load();
101    let (libcufile_loaded, libcufile_err, cufile) = match cufile_load {
102        Ok(c) => (true, None, Some(c)),
103        Err(e) => (false, Some(e), None),
104    };
105    let mut version: Option<i32> = None;
106    let mut driver_open_ok = false;
107    let mut driver_open_err: Option<String> = None;
108    if let Some(cufile) = cufile.as_ref() {
109        let mut v = 0i32;
110        let r = unsafe { (cufile.get_version)(&mut v) };
111        if r.err == metrale_gpu_sys::cufile::CU_FILE_SUCCESS {
112            version = Some(v);
113        }
114        let r = unsafe { (cufile.driver_open)() };
115        if r.err == metrale_gpu_sys::cufile::CU_FILE_SUCCESS {
116            driver_open_ok = true;
117        } else {
118            driver_open_err = Some(format!(
119                "{} ({})",
120                r.err,
121                metrale_gpu_sys::cufile::err_to_str(r.err)
122            ));
123        }
124    }
125
126    let cuda = CudaCtx::new(0).context("init CUDA context for probe")?;
127    let dev = DeviceBuffer::new(SEQ_IO_BYTES)?;
128    let pinned = PinnedBuffer::new(SEQ_IO_BYTES)?;
129    let (file, file_bytes) = open_test_file(&test_path)?;
130    let fd = file.as_raw_fd();
131
132    let cufile_seq = cufile.as_ref().filter(|_| driver_open_ok).and_then(|c| {
133        run_one_bench("cuFile seq 4MiB", || {
134            bench_cufile(c, fd, file_bytes, &dev, SEQ_IO_BYTES, SEQ_ITERS, true)
135        })
136    });
137    let cufile_rand = cufile.as_ref().filter(|_| driver_open_ok).and_then(|c| {
138        run_one_bench("cuFile rand 64KiB", || {
139            bench_cufile(c, fd, file_bytes, &dev, RAND_IO_BYTES, RAND_ITERS, false)
140        })
141    });
142    let iou_seq = run_one_bench("io_uring seq 4MiB", || {
143        bench_io_uring(
144            &cuda,
145            fd,
146            file_bytes,
147            &pinned,
148            &dev,
149            SEQ_IO_BYTES,
150            SEQ_ITERS,
151            true,
152        )
153    });
154    let iou_rand = run_one_bench("io_uring rand 64KiB", || {
155        bench_io_uring(
156            &cuda,
157            fd,
158            file_bytes,
159            &pinned,
160            &dev,
161            RAND_IO_BYTES,
162            RAND_ITERS,
163            false,
164        )
165    });
166    let posix_seq = run_one_bench("posix seq 4MiB", || {
167        bench_posix(
168            &cuda,
169            fd,
170            file_bytes,
171            &pinned,
172            &dev,
173            SEQ_IO_BYTES,
174            SEQ_ITERS,
175            true,
176        )
177    });
178    let posix_rand = run_one_bench("posix rand 64KiB", || {
179        bench_posix(
180            &cuda,
181            fd,
182            file_bytes,
183            &pinned,
184            &dev,
185            RAND_IO_BYTES,
186            RAND_ITERS,
187            false,
188        )
189    });
190    let _ = (file, dev, pinned, cuda);
191
192    let (recommended, reason) = decide(nvidia_fs, driver_open_ok, cufile_seq, iou_seq, posix_seq);
193
194    Ok(ProbeResult {
195        libcufile_loaded,
196        libcufile_load_error: libcufile_err,
197        cufile_version: version,
198        nvidia_fs_kmod_loaded: nvidia_fs,
199        cufile_driver_open_ok: driver_open_ok,
200        cufile_driver_open_error: driver_open_err,
201        cufile_seq_4mib: cufile_seq,
202        cufile_rand_64kib: cufile_rand,
203        io_uring_seq_4mib: iou_seq,
204        io_uring_rand_64kib: iou_rand,
205        posix_seq_4mib: posix_seq,
206        posix_rand_64kib: posix_rand,
207        recommended,
208        recommendation_reason: reason,
209    })
210}
211
212fn decide(
213    nvfs: bool,
214    cufile_open: bool,
215    cufile: Option<f64>,
216    iou: Option<f64>,
217    posix: Option<f64>,
218) -> (Backend, String) {
219    let cf = cufile.unwrap_or(0.0);
220    let iu = iou.unwrap_or(0.0);
221    let px = posix.unwrap_or(0.0);
222    if nvfs && cufile_open && cf >= iu * 1.05 {
223        return (
224            Backend::CuFileDirect,
225            format!("nvidia-fs loaded; cuFile {cf:.0} MiB/s ≥ io_uring {iu:.0} MiB/s"),
226        );
227    }
228    if cufile_open && cf >= iu.max(px) * 1.05 {
229        return (
230            Backend::CuFileCompat,
231            format!("cuFile compat-mode {cf:.0} MiB/s beats io_uring {iu:.0} / posix {px:.0}"),
232        );
233    }
234    if iu >= px && iu > 0.0 {
235        return (
236            Backend::IoUring,
237            format!("io_uring {iu:.0} MiB/s ≥ posix {px:.0} MiB/s"),
238        );
239    }
240    if px > 0.0 {
241        return (
242            Backend::PosixOnly,
243            format!("posix-only fallback {px:.0} MiB/s"),
244        );
245    }
246    (Backend::None, "no backend produced bandwidth".into())
247}
metrale_storage/probe.rs

metrale_storage/
probe.rs
