query_runtime/runtime/gpu/
inventory.rs

1//! Per-device GPU facts and live diagnostics telemetry.
2
3#[cfg(feature = "nvml")]
4use std::sync::OnceLock;
5use std::time::{SystemTime, UNIX_EPOCH};
6
7use query_planning::error::Result;
8use serde::Serialize;
9
10use crate::native_error::cudf_query_error;
11
12/// Availability of one live device-telemetry substrate.
13#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize)]
14#[non_exhaustive]
15pub enum DeviceTelemetryStatus {
16    /// The substrate returned a live sample.
17    Available,
18    /// The optional telemetry feature was not compiled.
19    FeatureDisabled,
20    /// The substrate or device sample was unavailable.
21    Unavailable,
22}
23
24impl DeviceTelemetryStatus {
25    /// Stable diagnostics label.
26    pub const fn as_str(self) -> &'static str {
27        match self {
28            Self::Available => "available",
29            Self::FeatureDisabled => "feature_disabled",
30            Self::Unavailable => "unavailable",
31        }
32    }
33}
34
35/// Live CUDA and NVML telemetry for one device.
36///
37/// These values are diagnostics. Admission and memory accounting read the
38/// device ledger, query-domain counters and pool statistics instead. Each
39/// `Option` is `None` when its substrate produced no sample for this snapshot.
40#[derive(Debug, Clone, PartialEq, Eq, Serialize)]
41pub struct DeviceTelemetrySnapshot {
42    /// Wall-clock sample time for correlating diagnostics with external telemetry.
43    pub(crate) sampled_at_unix_ms: u64,
44    /// CUDA ordinal sampled by this diagnostic snapshot.
45    pub device_ordinal: usize,
46    /// Availability of the CUDA memory sample.
47    pub cuda_status: DeviceTelemetryStatus,
48    pub(crate) cuda_total_bytes: Option<u64>,
49    /// CUDA-visible free bytes, or `None` when the CUDA sample is unavailable.
50    pub cuda_free_bytes: Option<u64>,
51    /// Availability of optional NVML telemetry, independent of CUDA availability.
52    pub nvml_status: DeviceTelemetryStatus,
53    pub(crate) nvml_total_bytes: Option<u64>,
54    pub(crate) nvml_reserved_bytes: Option<u64>,
55    /// NVML-used framebuffer bytes, or `None` when memory telemetry is unavailable.
56    pub nvml_used_bytes: Option<u64>,
57    pub(crate) nvml_free_bytes: Option<u64>,
58    /// NVML GPU utilization percentage, or `None` when unavailable.
59    pub gpu_utilization_pct: Option<u32>,
60    /// NVML memory-controller utilization percentage, or `None` when unavailable.
61    pub memory_utilization_pct: Option<u32>,
62}
63
64/// Fakeable provider for live device telemetry.
65///
66/// Telemetry is diagnostics only; it is not a substitute data source for the
67/// admission ledger.
68pub(crate) trait GpuTelemetry: Send + Sync {
69    /// Samples live telemetry for one device.
70    ///
71    /// Sampling does not fail. An unavailable substrate, a disabled telemetry
72    /// feature, or an unresolvable ordinal is reported through the snapshot's
73    /// status fields with the affected values left `None`, and the CUDA and
74    /// NVML substrates can be unavailable independently.
75    fn snapshot(&self, device_ordinal: usize) -> DeviceTelemetrySnapshot;
76}
77
78pub(crate) fn device_telemetry_snapshots_for_ordinals(
79    telemetry: &dyn GpuTelemetry,
80    device_ordinals: &[usize],
81) -> Vec<DeviceTelemetrySnapshot> {
82    device_ordinals.iter().map(|ordinal| telemetry.snapshot(*ordinal)).collect()
83}
84
85/// Production live telemetry backed by CUDA plus optional NVML.
86#[derive(Debug, Default, Clone, Copy)]
87pub(crate) struct AlgeonNvmlTelemetry;
88
89impl GpuTelemetry for AlgeonNvmlTelemetry {
90    fn snapshot(&self, device_ordinal: usize) -> DeviceTelemetrySnapshot {
91        let mut snapshot = DeviceTelemetrySnapshot {
92            sampled_at_unix_ms: sampled_at_unix_ms(),
93            device_ordinal,
94            cuda_status: DeviceTelemetryStatus::Unavailable,
95            cuda_total_bytes: None,
96            cuda_free_bytes: None,
97            nvml_status: nvml_status_without_sample(),
98            nvml_total_bytes: None,
99            nvml_reserved_bytes: None,
100            nvml_used_bytes: None,
101            nvml_free_bytes: None,
102            gpu_utilization_pct: None,
103            memory_utilization_pct: None,
104        };
105        let Ok(device) = cudf::DeviceId::new(device_ordinal) else {
106            return snapshot;
107        };
108        if let Ok(memory) = device.available_memory() {
109            snapshot.cuda_status = DeviceTelemetryStatus::Available;
110            snapshot.cuda_total_bytes = Some(memory.total_bytes);
111            snapshot.cuda_free_bytes = Some(memory.free_bytes);
112        }
113        sample_nvml(device, &mut snapshot);
114        snapshot
115    }
116}
117
118fn sampled_at_unix_ms() -> u64 {
119    SystemTime::now()
120        .duration_since(UNIX_EPOCH)
121        .ok()
122        .and_then(|duration| u64::try_from(duration.as_millis()).ok())
123        .unwrap_or(0)
124}
125
126/// Snapshot of the stable facts algeon needs for one CUDA device.
127#[derive(Debug, Clone, PartialEq, Eq)]
128pub struct DeviceFacts {
129    /// CUDA device ordinal sampled by backend construction.
130    pub device: usize,
131    /// NVML product name, when the optional NVML diagnostics substrate is available.
132    pub model_name: Option<String>,
133    /// NVML UUID, when the optional NVML diagnostics substrate is available.
134    pub uuid: Option<String>,
135    /// NVML-reported physical framebuffer capacity, when the optional NVML diagnostics substrate is available.
136    pub physical_total_bytes: Option<u64>,
137    /// CUDA-visible allocation capacity. This can exclude NVML-reserved framebuffer capacity.
138    pub total_bytes: u64,
139    /// Free bytes at sampling time. The registry's post-install capacity
140    /// freeze reads this once per device; it is never consulted during
141    /// admission.
142    pub(crate) free_bytes: u64,
143    pub(crate) sm_count: Option<u32>,
144}
145
146/// Query GPU facts behind a fakeable interface.
147pub(crate) trait GpuInventory: Send + Sync {
148    /// Facts for one device ordinal.
149    ///
150    /// `Err` means the required CUDA facts for the ordinal are unavailable.
151    /// Optional fields are `None` when their diagnostics substrate is
152    /// unavailable, and the call still returns `Ok`.
153    fn facts(&self, device_ordinal: usize) -> Result<DeviceFacts>;
154}
155
156/// Production GPU inventory backed by cudf.
157#[derive(Debug, Default, Clone, Copy)]
158pub(crate) struct AlgeonGpuInventory;
159
160impl GpuInventory for AlgeonGpuInventory {
161    fn facts(&self, device_ordinal: usize) -> Result<DeviceFacts> {
162        let device = cudf::DeviceId::new(device_ordinal).map_err(|error| {
163            cudf_query_error(
164                "query_engine.gpu_inventory.device_id",
165                "failed to resolve CUDA device id while collecting inventory facts",
166                error,
167            )
168            .with_fact("device_ordinal", device_ordinal)
169        })?;
170        let memory = device.available_memory().map_err(|error| {
171            cudf_query_error("query_engine.gpu_inventory.memory", "failed to query CUDA device memory inventory", error)
172                .with_fact("device_ordinal", device.ordinal())
173        })?;
174        let sm_count = device.multiprocessor_count().ok().filter(|count| *count > 0);
175        let (model_name, uuid, physical_total_bytes) = nvml_device_identity(device);
176        Ok(DeviceFacts {
177            device: device_ordinal,
178            model_name,
179            uuid,
180            physical_total_bytes,
181            total_bytes: memory.total_bytes,
182            free_bytes: memory.free_bytes,
183            sm_count,
184        })
185    }
186}
187
188#[cfg(feature = "nvml")]
189fn nvml_device_identity(device: cudf::DeviceId) -> (Option<String>, Option<String>, Option<u64>) {
190    let Some(nvml) = shared_nvml() else {
191        return (None, None, None);
192    };
193    let Ok(pci_bus_id) = device.pci_bus_id() else {
194        return (None, None, None);
195    };
196    let Ok(device) = nvml.device_by_pci_bus_id(pci_bus_id) else {
197        return (None, None, None);
198    };
199    (device.name().ok(), device.uuid().ok(), device.memory_info().ok().map(|memory| memory.total))
200}
201
202#[cfg(not(feature = "nvml"))]
203fn nvml_device_identity(_device: cudf::DeviceId) -> (Option<String>, Option<String>, Option<u64>) {
204    (None, None, None)
205}
206
207/// Convert stable device capacity facts into a bounded query budget.
208#[cfg(test)]
209pub(crate) fn device_budget_bytes_from_facts(facts: &DeviceFacts, headroom_bytes: u64, fallback_bytes: u64) -> u64 {
210    facts.total_bytes.checked_sub(headroom_bytes).filter(|value| *value > 0).unwrap_or(fallback_bytes)
211}
212
213#[cfg(feature = "nvml")]
214fn shared_nvml() -> Option<&'static nvml_wrapper::Nvml> {
215    static NVML: OnceLock<Option<nvml_wrapper::Nvml>> = OnceLock::new();
216    NVML.get_or_init(|| nvml_wrapper::Nvml::init().ok()).as_ref()
217}
218
219#[cfg(feature = "nvml")]
220fn nvml_status_without_sample() -> DeviceTelemetryStatus {
221    DeviceTelemetryStatus::Unavailable
222}
223
224#[cfg(not(feature = "nvml"))]
225fn nvml_status_without_sample() -> DeviceTelemetryStatus {
226    DeviceTelemetryStatus::FeatureDisabled
227}
228
229#[cfg(feature = "nvml")]
230fn sample_nvml(device: cudf::DeviceId, snapshot: &mut DeviceTelemetrySnapshot) {
231    let Some(nvml) = shared_nvml() else {
232        return;
233    };
234    let Ok(pci_bus_id) = device.pci_bus_id() else {
235        return;
236    };
237    let Ok(device) = nvml.device_by_pci_bus_id(pci_bus_id) else {
238        return;
239    };
240    let mut sampled = false;
241    if let Ok(memory) = device.memory_info() {
242        sampled = true;
243        snapshot.nvml_total_bytes = Some(memory.total);
244        snapshot.nvml_reserved_bytes = Some(memory.reserved);
245        snapshot.nvml_used_bytes = Some(memory.used);
246        snapshot.nvml_free_bytes = Some(memory.free);
247    }
248    if let Ok(utilization) = device.utilization_rates() {
249        sampled = true;
250        snapshot.gpu_utilization_pct = Some(utilization.gpu);
251        snapshot.memory_utilization_pct = Some(utilization.memory);
252    }
253    if sampled {
254        snapshot.nvml_status = DeviceTelemetryStatus::Available;
255    }
256}
257
258#[cfg(not(feature = "nvml"))]
259fn sample_nvml(_device: cudf::DeviceId, _snapshot: &mut DeviceTelemetrySnapshot) {}
260
261#[cfg(test)]
262mod tests {
263    use super::{
264        DeviceTelemetrySnapshot, DeviceTelemetryStatus, GpuTelemetry, device_telemetry_snapshots_for_ordinals,
265    };
266
267    struct FakeTelemetry;
268
269    impl GpuTelemetry for FakeTelemetry {
270        fn snapshot(&self, device_ordinal: usize) -> DeviceTelemetrySnapshot {
271            let nvml_available = device_ordinal == 3;
272            DeviceTelemetrySnapshot {
273                sampled_at_unix_ms: 42,
274                device_ordinal,
275                cuda_status: DeviceTelemetryStatus::Available,
276                cuda_total_bytes: Some(1_000 + device_ordinal as u64),
277                cuda_free_bytes: Some(500),
278                nvml_status: if nvml_available {
279                    DeviceTelemetryStatus::Available
280                } else {
281                    DeviceTelemetryStatus::Unavailable
282                },
283                nvml_total_bytes: nvml_available.then_some(1_000),
284                nvml_reserved_bytes: None,
285                nvml_used_bytes: nvml_available.then_some(250),
286                nvml_free_bytes: nvml_available.then_some(750),
287                gpu_utilization_pct: nvml_available.then_some(12),
288                memory_utilization_pct: nvml_available.then_some(34),
289            }
290        }
291    }
292
293    #[test]
294    fn telemetry_provider_samples_every_configured_ordinal_in_order() {
295        let snapshots = device_telemetry_snapshots_for_ordinals(&FakeTelemetry, &[3, 1]);
296
297        assert_eq!(snapshots.len(), 2);
298        assert_eq!(snapshots[0].device_ordinal, 3);
299        assert_eq!(snapshots[0].cuda_total_bytes, Some(1_003));
300        assert_eq!(snapshots[0].nvml_status, DeviceTelemetryStatus::Available);
301        assert_eq!(snapshots[0].gpu_utilization_pct, Some(12));
302        assert_eq!(snapshots[1].device_ordinal, 1);
303        assert_eq!(snapshots[1].nvml_status, DeviceTelemetryStatus::Unavailable);
304    }
305
306    #[cfg(feature = "nvml")]
307    #[test]
308    fn live_telemetry_covers_every_cuda_device_with_valid_ranges() {
309        use super::AlgeonNvmlTelemetry;
310
311        let device_count = cudf::DeviceId::count().expect("CUDA device count");
312        assert!(device_count > 0, "GPU test lane requires at least one CUDA device");
313        let ordinals = (0..device_count).collect::<Vec<_>>();
314
315        let snapshots = device_telemetry_snapshots_for_ordinals(&AlgeonNvmlTelemetry, &ordinals);
316
317        assert_eq!(snapshots.len(), device_count);
318        for (ordinal, snapshot) in snapshots.iter().enumerate() {
319            assert_eq!(snapshot.device_ordinal, ordinal);
320            assert_eq!(snapshot.cuda_status, DeviceTelemetryStatus::Available);
321            assert!(snapshot.cuda_free_bytes.unwrap() <= snapshot.cuda_total_bytes.unwrap());
322            assert_eq!(snapshot.nvml_status, DeviceTelemetryStatus::Available);
323            assert!(snapshot.nvml_free_bytes.unwrap() <= snapshot.nvml_total_bytes.unwrap());
324            assert!(snapshot.gpu_utilization_pct.unwrap() <= 100);
325            assert!(snapshot.memory_utilization_pct.unwrap() <= 100);
326        }
327    }
328}
query_runtime/runtime/gpu/inventory.rs

query_runtime/runtime/gpu/
inventory.rs
