Skip to main content

proof_stream/viz/
gpu.rs

1//! Helpers for building `ProofEvent::GpuStatus` from device info.
2
3use crate::events::{GpuSnapshot, ProofEvent};
4
5/// Build a `GpuStatus` event from a slice of device snapshots and progress
6/// counters.
7pub fn gpu_status_event(
8    devices: Vec<GpuSnapshot>,
9    matmul_done: usize,
10    matmul_total: usize,
11    layers_done: usize,
12    layers_total: usize,
13) -> ProofEvent {
14    ProofEvent::GpuStatus {
15        devices,
16        matmul_done,
17        matmul_total,
18        layers_done,
19        layers_total,
20    }
21}
22
23/// Create a `GpuSnapshot` with unknown utilization (useful when NVML is not
24/// available).
25pub fn unknown_gpu_snapshot(device_id: usize, name: impl Into<String>) -> GpuSnapshot {
26    GpuSnapshot {
27        device_id,
28        device_name: name.into(),
29        utilization: 0.0,
30        free_memory_bytes: None,
31    }
32}