Skip to main content

xlog_cuda_tests/harness/
provider.rs

1//! Test context and provider setup for CUDA certification tests.
2
3use cudarc::driver::sys;
4use cudarc::driver::DeviceRepr;
5use fs2::FileExt;
6use std::sync::atomic::{AtomicU64, Ordering};
7use std::sync::{Arc, Mutex, OnceLock};
8use std::{fs::OpenOptions, path::Path};
9use xlog_core::{resolve_bool, MemoryBudget, Result, XlogError};
10use xlog_cuda::device_runtime::{LogRecord, LoggingSink, SinkError, XlogDeviceRuntime};
11use xlog_cuda::memory::{DeviceRead, DeviceWrite};
12use xlog_cuda::{
13    CudaBuffer, CudaDevice, CudaKernelProvider, CudaProviderBuilder, GpuMemoryManager,
14};
15
16#[derive(Default)]
17struct TransferCounters {
18    htod_bytes: AtomicU64,
19    dtoh_bytes: AtomicU64,
20}
21
22impl TransferCounters {
23    fn reset(&self) {
24        self.htod_bytes.store(0, Ordering::SeqCst);
25        self.dtoh_bytes.store(0, Ordering::SeqCst);
26    }
27
28    fn add_htod(&self, bytes: u64) {
29        self.htod_bytes.fetch_add(bytes, Ordering::SeqCst);
30    }
31
32    fn add_dtoh(&self, bytes: u64) {
33        self.dtoh_bytes.fetch_add(bytes, Ordering::SeqCst);
34    }
35
36    fn snapshot(&self) -> (u64, u64) {
37        (
38            self.htod_bytes.load(Ordering::SeqCst),
39            self.dtoh_bytes.load(Ordering::SeqCst),
40        )
41    }
42}
43
44const GPU_TEST_LOCK_PATH: &str = "/tmp/xlog-gpu-tests.lock";
45
46fn gpu_test_lock() -> std::sync::MutexGuard<'static, ()> {
47    static LOCK: OnceLock<Mutex<()>> = OnceLock::new();
48    LOCK.get_or_init(|| Mutex::new(()))
49        .lock()
50        .unwrap_or_else(|e| e.into_inner())
51}
52
53fn gpu_test_lock_file() -> Result<std::fs::File> {
54    let path = Path::new(GPU_TEST_LOCK_PATH);
55    let file = OpenOptions::new()
56        .create(true)
57        .truncate(false)
58        .read(true)
59        .write(true)
60        .open(path)
61        .map_err(|e| XlogError::Kernel(format!("Failed to open GPU test lock file: {}", e)))?;
62    file.lock_exclusive()
63        .map_err(|e| XlogError::Kernel(format!("Failed to acquire GPU test lock file: {}", e)))?;
64    Ok(file)
65}
66
67/// Sink that drops every log record. The certification harness has
68/// no need for a live log stream during kernel testing, but `LoggingResource`
69/// requires a sink. Errors from the sink would propagate as
70/// allocation failures, which would mask kernel issues — drop
71/// silently here.
72struct DiscardSink;
73
74impl LoggingSink for DiscardSink {
75    fn emit(&self, _record: LogRecord) -> std::result::Result<(), SinkError> {
76        Ok(())
77    }
78}
79
80/// Test context providing CUDA resources for certification tests.
81pub struct TestContext {
82    // Hold a process-wide mutex for the lifetime of the context so GPU tests run serially.
83    // This prevents timing-based certification tests from producing false positives under
84    // parallel `cargo test` execution.
85    _lock: std::sync::MutexGuard<'static, ()>,
86    _file_lock: std::fs::File,
87    pub provider: Arc<CudaKernelProvider>,
88    pub device: Arc<CudaDevice>,
89    pub memory: Arc<GpuMemoryManager>,
90    /// Provider-owned runtime retained for explicit lifecycle checks.
91    _runtime: Arc<XlogDeviceRuntime>,
92    transfer: Arc<TransferCounters>,
93}
94
95/// Release-gate hardening: when `XLOG_REQUIRE_CUDA=1`, CUDA/test-context
96/// initialization failures must fail loudly instead of returning `Err`.
97/// Tests use `let Ok(ctx) = ... else { return; }` skip patterns that pass
98/// vacuously without a GPU; `scripts/validate_release_gpu.sh` exports this
99/// variable so a CPU-only machine can never satisfy the certification gate.
100pub fn enforce_cuda_required(context: &str, err: &XlogError) {
101    if resolve_bool(None, "XLOG_REQUIRE_CUDA", false)
102        .unwrap_or_else(|error| panic!("invalid CUDA test configuration: {error}"))
103    {
104        panic!("XLOG_REQUIRE_CUDA=1 but CUDA is unavailable ({context}): {err}");
105    }
106}
107
108impl TestContext {
109    /// Create test context with specific memory budget in bytes.
110    ///
111    /// When `XLOG_REQUIRE_CUDA=1`, any initialization failure panics via
112    /// [`enforce_cuda_required`] so callers' skip-on-error paths cannot turn
113    /// a missing GPU into a vacuous pass.
114    pub fn with_budget(budget_bytes: usize) -> Result<Self> {
115        let result = Self::with_budget_inner(budget_bytes);
116        if let Err(err) = &result {
117            enforce_cuda_required("TestContext::with_budget", err);
118        }
119        result
120    }
121
122    fn with_budget_inner(budget_bytes: usize) -> Result<Self> {
123        let lock = gpu_test_lock();
124        let file_lock = gpu_test_lock_file()?;
125
126        // cudarc may panic on driver init failures in restricted containers; treat as a normal error.
127        let device_count = std::panic::catch_unwind(CudaDevice::count)
128            .map_err(|_| {
129                XlogError::Kernel(
130                    "Failed to get device count: cudarc panicked during driver initialization"
131                        .to_string(),
132                )
133            })?
134            .map_err(|e| XlogError::Kernel(format!("Failed to get device count: {}", e)))?;
135
136        if device_count == 0 {
137            return Err(XlogError::Kernel("No CUDA devices available".to_string()));
138        }
139
140        let budget = MemoryBudget::with_limit(budget_bytes as u64);
141        let transfer = Arc::new(TransferCounters::default());
142        let provider = Arc::new(
143            CudaProviderBuilder::new(0, budget)
144                .with_logging_sink(Arc::new(DiscardSink) as Arc<dyn LoggingSink>)
145                .build()?,
146        );
147        let device = Arc::clone(provider.device());
148        let memory = Arc::clone(provider.memory());
149        let runtime = Arc::clone(memory.runtime().ok_or_else(|| {
150            XlogError::Kernel("provider builder did not attach a runtime".to_string())
151        })?);
152
153        Ok(Self {
154            _lock: lock,
155            _file_lock: file_lock,
156            provider,
157            device,
158            memory,
159            _runtime: runtime,
160            transfer,
161        })
162    }
163
164    /// Drain pending async frees on the provider-owned runtime allocator.
165    /// Cert harnesses call
166    /// this between categories so the `GlobalDeviceBudget`
167    /// reservation bookkeeping releases bytes that have been
168    /// freed via `cuMemFreeAsync` but whose owning stream has
169    /// not yet been synchronized — without this, a long
170    /// sequence of small allocate-then-drop tests piles up
171    /// pending frees and exhausts the reservation pool even
172    /// though the GPU has plenty of free memory.
173    pub fn reap_pending(&self) {
174        // Best-effort: a transient driver error during reap is recoverable on
175        // the next iteration, while the certification suite already runs
176        // categories sequentially under the harness lock.
177        let _ = self._runtime.reap_pending();
178    }
179
180    /// Create test context with default 1GB memory budget.
181    pub fn new() -> Result<Self> {
182        Self::with_budget(1024 * 1024 * 1024)
183    }
184
185    /// Create test context with large 4GB memory budget for stress tests.
186    pub fn large() -> Result<Self> {
187        Self::with_budget(4 * 1024 * 1024 * 1024)
188    }
189
190    /// Force device synchronization and check for async errors.
191    pub fn sync_and_check(&self) -> Result<()> {
192        self.device
193            .inner()
194            .synchronize()
195            .map_err(|e| XlogError::Kernel(format!("Sync failed: {}", e)))?;
196        Ok(())
197    }
198
199    /// Get current memory usage in bytes.
200    pub fn memory_used(&self) -> usize {
201        self.memory.allocated_bytes() as usize
202    }
203
204    /// Get memory budget in bytes.
205    pub fn memory_budget(&self) -> usize {
206        self.memory.budget().device_bytes as usize
207    }
208
209    /// Reset host/device transfer counters (in bytes).
210    pub fn reset_transfer_counters(&self) {
211        self.transfer.reset();
212    }
213
214    /// Snapshot transfer counters (htod_bytes, dtoh_bytes).
215    pub fn transfer_counters(&self) -> (u64, u64) {
216        self.transfer.snapshot()
217    }
218
219    /// Track a synchronous host-to-device copy.
220    pub fn htod_sync_copy_into<T: DeviceRepr, Dst: DeviceWrite<T>>(
221        &self,
222        src: &[T],
223        dst: &mut Dst,
224    ) -> Result<()> {
225        let bytes = std::mem::size_of::<T>()
226            .checked_mul(src.len())
227            .ok_or_else(|| XlogError::Kernel("htod byte count overflow".to_string()))?;
228        self.transfer.add_htod(bytes as u64);
229        self.device
230            .inner()
231            .htod_sync_copy_into(src, dst)
232            .map_err(|e| XlogError::Kernel(format!("Failed htod copy: {}", e)))?;
233        Ok(())
234    }
235
236    /// Track a synchronous device-to-host copy returning a Vec.
237    pub fn dtoh_sync_copy<T: DeviceRepr, Src: DeviceRead<T>>(&self, src: &Src) -> Result<Vec<T>> {
238        let bytes = std::mem::size_of::<T>()
239            .checked_mul(src.device_view().len())
240            .ok_or_else(|| XlogError::Kernel("dtoh byte count overflow".to_string()))?;
241        self.transfer.add_dtoh(bytes as u64);
242        self.device
243            .inner()
244            .dtoh_sync_copy(src)
245            .map_err(|e| XlogError::Kernel(format!("Failed dtoh copy: {}", e)))
246    }
247
248    /// Track a synchronous device-to-host copy into a slice.
249    pub fn dtoh_sync_copy_into<T: DeviceRepr, Src: DeviceRead<T>>(
250        &self,
251        src: &Src,
252        dst: &mut [T],
253    ) -> Result<()> {
254        let bytes = std::mem::size_of::<T>()
255            .checked_mul(dst.len())
256            .ok_or_else(|| XlogError::Kernel("dtoh byte count overflow".to_string()))?;
257        self.transfer.add_dtoh(bytes as u64);
258        self.device
259            .inner()
260            .dtoh_sync_copy_into(src, dst)
261            .map_err(|e| XlogError::Kernel(format!("Failed dtoh copy: {}", e)))?;
262        Ok(())
263    }
264
265    /// Measure GPU time (in milliseconds) for work enqueued on the device stream.
266    pub fn measure_gpu_ms<F>(&self, f: F) -> Result<f32>
267    where
268        F: FnOnce() -> Result<()>,
269    {
270        let stream = self.device.inner().stream();
271        let start = stream
272            .context()
273            .new_event(Some(sys::CUevent_flags::CU_EVENT_DEFAULT))
274            .map_err(|e| XlogError::Kernel(format!("Failed to create start event: {}", e)))?;
275        let end = stream
276            .context()
277            .new_event(Some(sys::CUevent_flags::CU_EVENT_DEFAULT))
278            .map_err(|e| XlogError::Kernel(format!("Failed to create end event: {}", e)))?;
279
280        start
281            .record(stream)
282            .map_err(|e| XlogError::Kernel(format!("Failed to record start event: {}", e)))?;
283        f()?;
284        end.record(stream)
285            .map_err(|e| XlogError::Kernel(format!("Failed to record end event: {}", e)))?;
286        self.sync_and_check()?;
287        let elapsed = start.elapsed_ms(&end).map_err(|e| {
288            XlogError::Kernel(format!("Failed to measure event elapsed time: {}", e))
289        })?;
290        Ok(elapsed)
291    }
292
293    /// Check if multi-GPU is available.
294    pub fn multi_gpu_available(&self) -> bool {
295        match std::panic::catch_unwind(CudaDevice::count) {
296            Ok(Ok(n)) => n > 1,
297            _ => false,
298        }
299    }
300
301    /// Get compute capability of current device.
302    pub fn compute_capability(&self) -> Result<(u32, u32)> {
303        use cudarc::driver::sys::CUdevice_attribute;
304
305        let major = self
306            .device
307            .inner()
308            .attribute(CUdevice_attribute::CU_DEVICE_ATTRIBUTE_COMPUTE_CAPABILITY_MAJOR)
309            .map_err(|e| {
310                XlogError::Kernel(format!("Failed to query compute capability major: {}", e))
311            })?;
312        let minor = self
313            .device
314            .inner()
315            .attribute(CUdevice_attribute::CU_DEVICE_ATTRIBUTE_COMPUTE_CAPABILITY_MINOR)
316            .map_err(|e| {
317                XlogError::Kernel(format!("Failed to query compute capability minor: {}", e))
318            })?;
319
320        let major_u32: u32 = major.try_into().map_err(|_| {
321            XlogError::Kernel(format!(
322                "Failed to convert compute capability major {} to u32",
323                major
324            ))
325        })?;
326        let minor_u32: u32 = minor.try_into().map_err(|_| {
327            XlogError::Kernel(format!(
328                "Failed to convert compute capability minor {} to u32",
329                minor
330            ))
331        })?;
332
333        Ok((major_u32, minor_u32))
334    }
335
336    /// Read device-resident row count for a buffer.
337    ///
338    /// Panics if the count cannot be read; certification tests treat this as fatal.
339    pub fn device_row_count(&self, buffer: &CudaBuffer) -> u64 {
340        let mut host_rows = [0u32];
341        self.dtoh_sync_copy_into(buffer.num_rows_device(), &mut host_rows)
342            .unwrap_or_else(|e| panic!("Failed to read device row count: {}", e));
343        host_rows[0] as u64
344    }
345}
346
347/// Macro for tests requiring CUDA device - panics if not available.
348#[macro_export]
349macro_rules! gpu_test {
350    ($name:ident, $body:expr) => {
351        #[test]
352        fn $name() {
353            let ctx =
354                $crate::harness::TestContext::new().expect("CUDA device required for this test");
355            $body(&ctx);
356        }
357    };
358}
359
360/// Macro for tests requiring CUDA device - skips if not available.
361#[macro_export]
362macro_rules! gpu_test_skip {
363    ($name:ident, $body:expr) => {
364        #[test]
365        fn $name() {
366            let ctx = match $crate::harness::TestContext::new() {
367                Ok(ctx) => ctx,
368                Err(_) => {
369                    eprintln!("Skipping {}: no CUDA device", stringify!($name));
370                    return;
371                }
372            };
373            $body(&ctx);
374        }
375    };
376}
377
378#[cfg(test)]
379mod tests {
380    use super::*;
381    use std::process::Command;
382
383    #[test]
384    fn test_gpu_test_lock_file_is_exclusive() {
385        let ctx = match TestContext::new() {
386            Ok(ctx) => ctx,
387            Err(_) => {
388                eprintln!("Skipping test: no CUDA device available");
389                return;
390            }
391        };
392
393        let status = Command::new("flock")
394            .arg("-n")
395            .arg(GPU_TEST_LOCK_PATH)
396            .arg("-c")
397            .arg("true")
398            .status();
399
400        match status {
401            Ok(status) => assert!(
402                !status.success(),
403                "GPU test lock file should be held exclusively while TestContext is alive"
404            ),
405            Err(_) => {
406                eprintln!("Skipping test: flock command not available");
407            }
408        }
409
410        drop(ctx);
411    }
412}