bge_m3_embedding_server/gpu_stats.rs
1// Copyright (c) 2026 J. Patrick Fulton
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License at
6//
7// http://www.apache.org/licenses/LICENSE-2.0
8//
9// Unless required by applicable law or agreed to in writing, software
10// distributed under the License is distributed on an "AS IS" BASIS,
11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12// See the License for the specific language governing permissions and
13// limitations under the License.
14
15//! Per-device GPU VRAM, utilization, and temperature heartbeat logging.
16//!
17//! On GPU builds (`cuda` or `tensorrt` feature), this module initialises an
18//! NVML handle at startup and emits one structured `INFO` log event per CUDA
19//! device on each heartbeat tick. On CPU builds the module compiles to a
20//! zero-cost stub so the rest of the codebase can call it unconditionally
21//! without any `#[cfg]` noise at the call site.
22
23// ---------------------------------------------------------------------------
24// GPU build — real NVML implementation
25// ---------------------------------------------------------------------------
26
27#[cfg(feature = "nvml")]
28mod inner {
29 use nvml_wrapper::Nvml;
30 use nvml_wrapper::enum_wrappers::device::TemperatureSensor;
31 use tracing::{debug, info, warn};
32
33 enum NvmlState {
34 Ready { nvml: Box<Nvml>, gpu_count: usize },
35 Unavailable,
36 }
37
38 /// Collects per-device VRAM and GPU utilization stats via NVML and emits
39 /// them as structured log events.
40 ///
41 /// Instantiate once with [`GpuStatsCollector::init`], then call
42 /// [`GpuStatsCollector::emit_heartbeat`] on each heartbeat tick.
43 pub struct GpuStatsCollector {
44 state: NvmlState,
45 }
46
47 impl GpuStatsCollector {
48 /// Attempts to initialise NVML.
49 ///
50 /// If NVML is unavailable (driver not present, permission denied, etc.)
51 /// a single `WARN` is logged and the collector enters a no-op state for
52 /// the remainder of the process lifetime.
53 pub fn init(gpu_count: usize) -> Self {
54 match Nvml::init() {
55 Ok(nvml) => {
56 info!(gpu_count, "NVML initialised; GPU heartbeat stats enabled");
57 Self {
58 state: NvmlState::Ready {
59 nvml: Box::new(nvml),
60 gpu_count,
61 },
62 }
63 }
64 Err(e) => {
65 warn!(
66 error = %e,
67 "NVML unavailable — GPU heartbeat stats disabled for this process"
68 );
69 Self {
70 state: NvmlState::Unavailable,
71 }
72 }
73 }
74 }
75
76 /// Emits one `INFO` log event per CUDA device with VRAM, utilization,
77 /// and temperature statistics.
78 ///
79 /// Per-device errors are logged at `DEBUG` and skipped; the loop
80 /// continues for remaining devices. This method never panics.
81 pub fn emit_heartbeat(&self) {
82 let (nvml, gpu_count) = match &self.state {
83 NvmlState::Ready { nvml, gpu_count } => (nvml, *gpu_count),
84 NvmlState::Unavailable => return,
85 };
86
87 for device_idx in 0..gpu_count {
88 #[allow(clippy::cast_possible_truncation)]
89 let device_idx_u32 = device_idx as u32;
90
91 let device = match nvml.device_by_index(device_idx_u32) {
92 Ok(d) => d,
93 Err(e) => {
94 debug!(
95 gpu_device = device_idx_u32,
96 error = %e,
97 "NVML: could not open device"
98 );
99 continue;
100 }
101 };
102
103 let mem = match device.memory_info() {
104 Ok(m) => m,
105 Err(e) => {
106 debug!(
107 gpu_device = device_idx_u32,
108 error = %e,
109 "NVML: could not read memory info"
110 );
111 continue;
112 }
113 };
114
115 let utilization = match device.utilization_rates() {
116 Ok(u) => u,
117 Err(e) => {
118 debug!(
119 gpu_device = device_idx_u32,
120 error = %e,
121 "NVML: could not read utilization rates"
122 );
123 continue;
124 }
125 };
126
127 let gpu_temp_c = match device.temperature(TemperatureSensor::Gpu) {
128 Ok(t) => t,
129 Err(e) => {
130 debug!(
131 gpu_device = device_idx_u32,
132 error = %e,
133 "NVML: could not read GPU temperature"
134 );
135 continue;
136 }
137 };
138
139 let vram_used_mb = mem.used / (1024 * 1024);
140 let vram_total_mb = mem.total / (1024 * 1024);
141 #[allow(clippy::cast_precision_loss)]
142 let vram_utilization_pct = if mem.total > 0 {
143 mem.used as f32 / mem.total as f32 * 100.0
144 } else {
145 0.0
146 };
147 let gpu_utilization_pct = utilization.gpu;
148 let gpu_temp_f = gpu_temp_c * 9 / 5 + 32;
149
150 info!(
151 gpu_device = device_idx_u32,
152 vram_used_mb,
153 vram_total_mb,
154 vram_utilization_pct,
155 gpu_utilization_pct,
156 gpu_temp_c,
157 gpu_temp_f,
158 "gpu heartbeat"
159 );
160 }
161 }
162 }
163}
164
165// ---------------------------------------------------------------------------
166// CPU build — zero-cost stub
167// ---------------------------------------------------------------------------
168
169#[cfg(not(feature = "nvml"))]
170mod inner {
171 /// No-op GPU stats collector for CPU builds.
172 ///
173 /// All methods compile away completely; no NVML dependency is pulled in.
174 pub struct GpuStatsCollector;
175
176 impl GpuStatsCollector {
177 /// Returns a no-op collector. The `gpu_count` argument is accepted for
178 /// API compatibility with GPU builds but is otherwise ignored.
179 pub fn init(_gpu_count: usize) -> Self {
180 Self
181 }
182
183 /// No-op on CPU builds.
184 #[allow(clippy::unused_self)]
185 pub fn emit_heartbeat(&self) {}
186 }
187}
188
189pub(crate) use inner::GpuStatsCollector;