spark_storage/
probe.rs

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

spark_storage/
probe.rs
