From 33adb01dc15c5170b0946fb317d9a9cbd74d8647 Mon Sep 17 00:00:00 2001 From: Vincent Date: Tue, 18 Aug 2026 15:47:58 +0800 Subject: [PATCH 1/2] feat(gpu): add Ascend NPU telemetry backend Add DCMI and npu-smi collection with bounded fallback behavior so Ascend devices integrate with existing GPU diagnostics without affecting CUDA-only workloads. Co-authored-by: Cursor --- Cargo.lock | 1 + docs/src/reference/env-vars.md | 4 +- probing/extensions/gpu/Cargo.toml | 5 +- .../gpu/src/extensions/backend/apple/mod.rs | 4 + .../gpu/src/extensions/backend/cuda/mod.rs | 11 +- .../gpu/src/extensions/backend/mod.rs | 6 + .../gpu/src/extensions/backend/npu/dcmi.rs | 239 +++++++ .../gpu/src/extensions/backend/npu/hccs.rs | 238 +++++++ .../gpu/src/extensions/backend/npu/mod.rs | 131 ++++ .../gpu/src/extensions/backend/npu/npu_smi.rs | 599 ++++++++++++++++++ .../gpu/src/extensions/backend/registry.rs | 54 +- .../gpu/src/extensions/backend/traits.rs | 12 + .../gpu/src/extensions/collector.rs | 74 ++- .../gpu/src/extensions/extension.rs | 2 +- python/probing_hook.py | 8 +- 15 files changed, 1355 insertions(+), 33 deletions(-) create mode 100644 probing/extensions/gpu/src/extensions/backend/npu/dcmi.rs create mode 100644 probing/extensions/gpu/src/extensions/backend/npu/hccs.rs create mode 100644 probing/extensions/gpu/src/extensions/backend/npu/mod.rs create mode 100644 probing/extensions/gpu/src/extensions/backend/npu/npu_smi.rs diff --git a/Cargo.lock b/Cargo.lock index 7487360e..1ba8aa58 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -3453,6 +3453,7 @@ dependencies = [ "async-trait", "cudarc", "libc", + "libloading 0.8.6", "log", "metal", "once_cell", diff --git a/docs/src/reference/env-vars.md b/docs/src/reference/env-vars.md index 8905a3ff..a5fc9131 100644 --- a/docs/src/reference/env-vars.md +++ b/docs/src/reference/env-vars.md @@ -105,7 +105,9 @@ Non-PROBING-prefixed aliases are also recognized for Megatron compatibility: |----------|---------|-------------| | `PROBING_GPU` | enabled | Set to `0`, `off`, `false`, or `no` to disable GPU sampling. | | `PROBING_GPU_SAMPLE_MS` | — | GPU sampling interval in milliseconds. | -| `PROBING_GPU_BACKEND` | `auto` | GPU backend filter: `auto`, `cuda`, `rocm`, `metal`. | +| `PROBING_GPU_BACKEND` | `auto` | GPU backend filter: `auto`, `cuda`, `npu`/`ascend`, `rocm`, `metal`. | +| `PROBING_NPU_SOURCE` | `auto` | Ascend metrics source: `auto` prefers DCMI and falls back to `npu-smi`; use `smi` to force the CLI path. | +| `PROBING_DCMI_LIB` | — | Explicit path to Ascend `libdcmi.so`. | ## NCCL & HCCL diff --git a/probing/extensions/gpu/Cargo.toml b/probing/extensions/gpu/Cargo.toml index 6a9a441a..8701f6c2 100644 --- a/probing/extensions/gpu/Cargo.toml +++ b/probing/extensions/gpu/Cargo.toml @@ -1,6 +1,6 @@ [package] name = "probing-gpu" -description = "GPU diagnostics for probing (CUDA + Apple Silicon, multi-backend)" +description = "GPU/NPU diagnostics for probing (CUDA, Ascend, and Apple Silicon)" version.workspace = true authors.workspace = true edition.workspace = true @@ -30,6 +30,9 @@ cudarc = { version = "0.19.7", default-features = false, features = [ "cuda-12080", ], optional = true } +[target.'cfg(all(not(target_os = "macos"), not(target_os = "windows")))'.dependencies] +libloading = "0.8" + # macOS-only: Metal + sysctl + ioreg. Not resolved when building for Linux. [target.'cfg(target_os = "macos")'.dependencies] metal = "0.31.0" diff --git a/probing/extensions/gpu/src/extensions/backend/apple/mod.rs b/probing/extensions/gpu/src/extensions/backend/apple/mod.rs index 8253ad0f..b19964fd 100644 --- a/probing/extensions/gpu/src/extensions/backend/apple/mod.rs +++ b/probing/extensions/gpu/src/extensions/backend/apple/mod.rs @@ -132,6 +132,10 @@ impl GpuBackend for AppleSiliconBackend { renderer_util_pct: perf.as_ref().and_then(|p| p.renderer_util_pct), tiler_util_pct: perf.as_ref().and_then(|p| p.tiler_util_pct), driver_mem_bytes, + vector_util_pct: None, + mem_bandwidth_util_pct: None, + temperature_c: None, + power_w: None, }) } } diff --git a/probing/extensions/gpu/src/extensions/backend/cuda/mod.rs b/probing/extensions/gpu/src/extensions/backend/cuda/mod.rs index 03a5808c..91791756 100644 --- a/probing/extensions/gpu/src/extensions/backend/cuda/mod.rs +++ b/probing/extensions/gpu/src/extensions/backend/cuda/mod.rs @@ -93,6 +93,10 @@ impl GpuBackend for CudaBackend { renderer_util_pct: None, tiler_util_pct: None, driver_mem_bytes: None, + vector_util_pct: None, + mem_bandwidth_util_pct: None, + temperature_c: None, + power_w: None, }) } } @@ -103,10 +107,9 @@ fn probe_cuda_backend() -> Option { return None; } - let count = - std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| CudaContext::device_count())) - .ok() - .and_then(|r| r.ok())?; + let count = std::panic::catch_unwind(std::panic::AssertUnwindSafe(CudaContext::device_count)) + .ok() + .and_then(|r| r.ok())?; if count <= 0 { log::debug!("CUDA backend unavailable (zero devices)"); diff --git a/probing/extensions/gpu/src/extensions/backend/mod.rs b/probing/extensions/gpu/src/extensions/backend/mod.rs index bcbb4783..8ad50b23 100644 --- a/probing/extensions/gpu/src/extensions/backend/mod.rs +++ b/probing/extensions/gpu/src/extensions/backend/mod.rs @@ -7,8 +7,14 @@ mod apple; #[cfg(feature = "cuda")] mod cuda; +#[cfg(all(not(target_os = "macos"), not(target_os = "windows")))] +mod npu; + pub use registry::{discover_backends, selected_backends}; pub use traits::{GpuBackend, GpuBackendKind, GpuDeviceInfo, GpuMemoryModel, GpuMemorySample}; #[cfg(feature = "cuda")] pub use cuda::read_utilization_by_index; + +#[cfg(all(not(target_os = "macos"), not(target_os = "windows")))] +pub use npu::read_utilization_by_index as read_npu_utilization_by_index; diff --git a/probing/extensions/gpu/src/extensions/backend/npu/dcmi.rs b/probing/extensions/gpu/src/extensions/backend/npu/dcmi.rs new file mode 100644 index 00000000..26cedc54 --- /dev/null +++ b/probing/extensions/gpu/src/extensions/backend/npu/dcmi.rs @@ -0,0 +1,239 @@ +//! Ascend DCMI (`libdcmi.so`) — in-process device management API (NVML analogue). +//! +//! Loaded dynamically so non-Ascend hosts still build/link. Prefer this over +//! `npu-smi` subprocesses for utilization / HBM / temperature / power. + +use std::collections::HashMap; +use std::os::raw::{c_int, c_uint}; +use std::path::Path; +use std::sync::Mutex; + +use libloading::Library; +use once_cell::sync::Lazy; + +use super::npu_smi::NpuDeviceStats; + +const DCMI_UTILIZATION_RATE_AICORE: c_int = 2; +const DCMI_UTILIZATION_RATE_HBM: c_int = 6; +const DCMI_UTILIZATION_RATE_NPU: c_int = 13; + +/// Empirical: `dcmi_get_device_power_info` returns tenths of a watt (1640 → 164.0 W). +const POWER_SCALE: f32 = 0.1; + +#[repr(C)] +#[derive(Clone, Copy, Default)] +struct DcmiHbmInfo { + memory_size: u64, + freq: u32, + memory_usage: u64, + temp: c_int, + bandwith_util_rate: u32, +} + +type FnInit = unsafe extern "C" fn() -> c_int; +type FnGetCardList = unsafe extern "C" fn(*mut c_int, *mut c_int, c_int) -> c_int; +type FnGetDeviceNumInCard = unsafe extern "C" fn(c_int, *mut c_int) -> c_int; +type FnGetDeviceUtilizationRate = unsafe extern "C" fn(c_int, c_int, c_int, *mut c_uint) -> c_int; +type FnGetDeviceTemperature = unsafe extern "C" fn(c_int, c_int, *mut c_int) -> c_int; +type FnGetDevicePowerInfo = unsafe extern "C" fn(c_int, c_int, *mut c_int) -> c_int; +type FnGetDeviceHbmInfo = unsafe extern "C" fn(c_int, c_int, *mut DcmiHbmInfo) -> c_int; + +struct DcmiApi { + _lib: Library, + init: FnInit, + get_card_list: FnGetCardList, + get_device_num_in_card: FnGetDeviceNumInCard, + get_device_utilization_rate: FnGetDeviceUtilizationRate, + get_device_temperature: FnGetDeviceTemperature, + get_device_power_info: FnGetDevicePowerInfo, + get_device_hbm_info: FnGetDeviceHbmInfo, +} + +static DCMI: Lazy>> = Lazy::new(|| Mutex::new(load_dcmi())); + +fn candidate_paths() -> Vec { + let mut paths = Vec::new(); + if let Ok(p) = std::env::var("PROBING_DCMI_LIB") { + if !p.is_empty() { + paths.push(p); + } + } + paths.extend([ + "/usr/local/Ascend/driver/lib64/driver/libdcmi.so".to_string(), + "/usr/local/Ascend/driver/lib64/libdcmi.so".to_string(), + "libdcmi.so".to_string(), + ]); + paths +} + +fn load_dcmi() -> Option { + if matches!( + std::env::var("PROBING_NPU_SOURCE").ok().as_deref(), + Some(v) if matches!(v.trim().to_ascii_lowercase().as_str(), "smi" | "npu-smi" | "cli") + ) { + log::info!("DCMI disabled via PROBING_NPU_SOURCE"); + return None; + } + + for path in candidate_paths() { + if path != "libdcmi.so" && !Path::new(&path).exists() { + continue; + } + // SAFETY: we only call Ascend DCMI C ABI symbols after successful load. + let result = unsafe { try_load_path(&path) }; + match result { + Ok(api) => { + let rc = unsafe { (api.init)() }; + if rc != 0 { + log::warn!("dcmi_init failed on {path}: rc={rc}"); + continue; + } + log::info!("DCMI loaded from {path}"); + return Some(api); + } + Err(e) => log::debug!("DCMI load skip {path}: {e}"), + } + } + None +} + +unsafe fn try_load_path(path: &str) -> Result { + let lib = Library::new(path).map_err(|e| e.to_string())?; + Ok(DcmiApi { + init: *lib + .get::(b"dcmi_init\0") + .map_err(|e| e.to_string())?, + get_card_list: *lib + .get::(b"dcmi_get_card_list\0") + .map_err(|e| e.to_string())?, + get_device_num_in_card: *lib + .get::(b"dcmi_get_device_num_in_card\0") + .map_err(|e| e.to_string())?, + get_device_utilization_rate: *lib + .get::(b"dcmi_get_device_utilization_rate\0") + .map_err(|e| e.to_string())?, + get_device_temperature: *lib + .get::(b"dcmi_get_device_temperature\0") + .map_err(|e| e.to_string())?, + get_device_power_info: *lib + .get::(b"dcmi_get_device_power_info\0") + .map_err(|e| e.to_string())?, + get_device_hbm_info: *lib + .get::(b"dcmi_get_device_hbm_info\0") + .map_err(|e| e.to_string())?, + _lib: lib, + }) +} + +fn util_or_neg1(api: &DcmiApi, card: c_int, device: c_int, kind: c_int) -> f32 { + let mut rate: c_uint = 0; + let rc = unsafe { (api.get_device_utilization_rate)(card, device, kind, &mut rate) }; + if rc == 0 { + (rate as f32).clamp(0.0, 100.0) + } else { + -1.0 + } +} + +/// Batch-read via DCMI. Returns `None` if library unavailable or no devices. +pub fn read_utilization_by_index() -> Option> { + let guard = DCMI.lock().ok()?; + let api = guard.as_ref()?; + + let mut card_num: c_int = 0; + let mut cards = [0_i32; 64]; + let rc = + unsafe { (api.get_card_list)(&mut card_num, cards.as_mut_ptr(), cards.len() as c_int) }; + if rc != 0 || card_num <= 0 { + return None; + } + + let mut map = HashMap::new(); + for &card_id in cards.iter().take(card_num as usize) { + let mut device_num: c_int = 0; + if unsafe { (api.get_device_num_in_card)(card_id, &mut device_num) } != 0 || device_num <= 0 + { + continue; + } + let chip_count = device_num.max(1); + for device_id in 0..device_num { + let mut aicore = util_or_neg1(api, card_id, device_id, DCMI_UTILIZATION_RATE_AICORE); + let npu_util = util_or_neg1(api, card_id, device_id, DCMI_UTILIZATION_RATE_NPU); + let hbm_util = util_or_neg1(api, card_id, device_id, DCMI_UTILIZATION_RATE_HBM); + if npu_util >= 0.0 { + aicore = npu_util; + } else if aicore < 0.0 { + aicore = 0.0; + } + + let mut temp: c_int = 0; + let temp_c = + if unsafe { (api.get_device_temperature)(card_id, device_id, &mut temp) } == 0 { + Some(temp as f32) + } else { + None + }; + + let mut power_raw: c_int = 0; + let power_w = if unsafe { + (api.get_device_power_info)(card_id, device_id, &mut power_raw) + } == 0 + { + Some((power_raw as f32) * POWER_SCALE) + } else { + None + }; + + let mut hbm = DcmiHbmInfo::default(); + let (hbm_total_bytes, hbm_used_bytes, hbm_util_final) = + if unsafe { (api.get_device_hbm_info)(card_id, device_id, &mut hbm) } == 0 { + // Empirically MB on A3 (65536 → 64 GiB), matching npu-smi. + let total = hbm.memory_size.saturating_mul(1024 * 1024); + let used = hbm.memory_usage.saturating_mul(1024 * 1024); + let util = if hbm_util >= 0.0 { + hbm_util + } else if total > 0 { + (used as f64 / total as f64 * 100.0) as f32 + } else { + 0.0 + }; + (total, used, util) + } else { + (0, 0, hbm_util.max(0.0)) + }; + + // DCMI field is historically misspelled `bandwith_util_rate` in the API. + let hbm_bw = if hbm.bandwith_util_rate > 0 || hbm.memory_size > 0 { + Some(hbm.bandwith_util_rate as f32) + } else { + None + }; + + let phy_id = card_id * chip_count + device_id; + map.insert( + phy_id, + NpuDeviceStats { + ai_core_util_pct: aicore, + hbm_util_pct: hbm_util_final, + hbm_total_bytes, + hbm_used_bytes, + aivector_util_pct: None, // DCMI has no Aivector util kind today + hbm_bw_util_pct: hbm_bw, + temp_c, + power_w, + source: "dcmi", + }, + ); + } + } + + if map.is_empty() { + None + } else { + Some(map) + } +} + +/// Keep the Library last-to-drop by declaring it first (Rust drops fields in reverse order). +#[allow(dead_code)] +fn _ensure_drop_order_docs() {} diff --git a/probing/extensions/gpu/src/extensions/backend/npu/hccs.rs b/probing/extensions/gpu/src/extensions/backend/npu/hccs.rs new file mode 100644 index 00000000..5cbe399e --- /dev/null +++ b/probing/extensions/gpu/src/extensions/backend/npu/hccs.rs @@ -0,0 +1,238 @@ +//! Parse Ascend HCCS (NVLink analogue) counters via `npu-smi -t hccs`. + +use super::npu_smi::{list_npu_card_ids, run_npu_smi_text}; + +/// One chip's HCCS snapshot (cumulative counters + health). +#[derive(Debug, Clone, Default)] +pub struct HccsChipSample { + pub npu_id: i32, + pub chip_id: i32, + /// Logical Phy-ID matching `gpu.utilization.device_id` (NPU*chip_count + chip). + pub device_id: i32, + pub tx_bytes: u64, + pub rx_bytes: u64, + pub tx_packets: u64, + pub rx_packets: u64, + pub error_count: u64, + pub retry_count: u64, + pub health_ok: bool, +} + +/// Optional instantaneous bandwidth from `npu-smi -t hccs-bw` (expensive). +#[derive(Debug, Clone, Default)] +pub struct HccsBwSample { + pub tx_gbs: f32, + pub rx_gbs: f32, +} + +/// Sum a bracket list like `[a b c]` or space-separated numbers on a value line. +fn sum_u64_list(value: &str) -> u64 { + let cleaned = value + .trim() + .trim_start_matches('[') + .trim_end_matches(']') + .replace(',', " "); + cleaned + .split_whitespace() + .filter_map(|t| t.parse::().ok()) + .sum() +} + +fn parse_kv(line: &str) -> Option<(&str, &str)> { + let (k, v) = line.split_once(':')?; + Some((k.trim(), v.trim())) +} + +/// Discover (npu_id, chip_id, device_id) triples once; reuse across ticks. +pub fn discover_hccs_targets() -> Vec<(i32, i32, i32)> { + let mut out = Vec::new(); + for npu_id in list_npu_card_ids() { + let chip_ids = discover_chip_ids(npu_id); + let chip_count = chip_ids.len().max(1) as i32; + for chip_id in chip_ids { + out.push((npu_id, chip_id, npu_id * chip_count + chip_id)); + } + } + out +} + +/// Sample HCCS counters for one chip. +pub fn sample_hccs_chip(npu_id: i32, chip_id: i32, device_id: i32) -> Option { + let text = run_npu_smi_text(&[ + "info", + "-t", + "hccs", + "-i", + &npu_id.to_string(), + "-c", + &chip_id.to_string(), + ])?; + let mut sample = HccsChipSample { + npu_id, + chip_id, + device_id, + health_ok: true, + ..Default::default() + }; + let mut saw = false; + for line in text.lines() { + let Some((k, v)) = parse_kv(line) else { + continue; + }; + let kl = k.to_ascii_lowercase(); + if kl.contains("hccs tx bytes") { + sample.tx_bytes = sum_u64_list(v); + saw = true; + } else if kl.contains("hccs rx bytes") { + sample.rx_bytes = sum_u64_list(v); + saw = true; + } else if kl.contains("hccs tx packets") { + sample.tx_packets = sum_u64_list(v); + saw = true; + } else if kl.contains("hccs rx packets") { + sample.rx_packets = sum_u64_list(v); + saw = true; + } else if kl.contains("hccs error count") { + sample.error_count = sum_u64_list(v); + saw = true; + } else if kl.contains("hccs retry count") { + sample.retry_count = sum_u64_list(v); + saw = true; + } else if kl.contains("hccs health") { + sample.health_ok = v.to_ascii_uppercase().contains("OK"); + saw = true; + } + } + saw.then_some(sample) +} + +/// Parse `hccs-bw` total row (GB/s). Slow (~1s/chip); call sparingly. +pub fn sample_hccs_bw(npu_id: i32, chip_id: i32) -> Option { + let text = run_npu_smi_text(&[ + "info", + "-t", + "hccs-bw", + "-i", + &npu_id.to_string(), + "-c", + &chip_id.to_string(), + ])?; + let mut tx = -1.0_f32; + let mut rx = -1.0_f32; + for line in text.lines() { + let parts: Vec<&str> = line.split_whitespace().collect(); + if parts.len() >= 3 && parts[0].eq_ignore_ascii_case("total") { + if let (Ok(t), Ok(r)) = (parts[1].parse::(), parts[2].parse::()) { + tx = t; + rx = r; + } + } + } + if tx < 0.0 && rx < 0.0 { + None + } else { + Some(HccsBwSample { + tx_gbs: tx.max(0.0), + rx_gbs: rx.max(0.0), + }) + } +} + +fn discover_chip_ids(npu_id: i32) -> Vec { + let Some(text) = run_npu_smi_text(&["info", "-t", "usages", "-i", &npu_id.to_string()]) else { + return vec![0]; + }; + let mut chips: Vec = text + .lines() + .filter_map(|line| { + let (k, v) = parse_kv(line)?; + if k.eq_ignore_ascii_case("Chip ID") { + v.split_whitespace().next()?.parse().ok() + } else { + None + } + }) + .collect(); + chips.sort_unstable(); + chips.dedup(); + if chips.is_empty() { + vec![0] + } else { + chips + } +} + +/// Parse a multi-line HCCS block (unit tests / offline). +pub fn parse_hccs_text( + text: &str, + npu_id: i32, + chip_id: i32, + device_id: i32, +) -> Option { + let mut sample = HccsChipSample { + npu_id, + chip_id, + device_id, + health_ok: true, + ..Default::default() + }; + let mut saw = false; + for line in text.lines() { + let Some((k, v)) = parse_kv(line) else { + continue; + }; + let kl = k.to_ascii_lowercase(); + if kl.contains("hccs tx bytes") { + sample.tx_bytes = sum_u64_list(v); + saw = true; + } else if kl.contains("hccs rx bytes") { + sample.rx_bytes = sum_u64_list(v); + saw = true; + } else if kl.contains("hccs tx packets") { + sample.tx_packets = sum_u64_list(v); + saw = true; + } else if kl.contains("hccs rx packets") { + sample.rx_packets = sum_u64_list(v); + saw = true; + } else if kl.contains("hccs error count") { + sample.error_count = sum_u64_list(v); + saw = true; + } else if kl.contains("hccs retry count") { + sample.retry_count = sum_u64_list(v); + saw = true; + } else if kl.contains("hccs health") { + sample.health_ok = v.to_ascii_uppercase().contains("OK"); + saw = true; + } + } + saw.then_some(sample) +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn sum_bracket_list() { + assert_eq!(sum_u64_list("[10 20 30]"), 60); + assert_eq!(sum_u64_list("5"), 5); + } + + #[test] + fn parse_real_hccs_block() { + let text = r#" + hccs health status : OK + hccs tx packets : [1993002168 1434439505 1588783520] + hccs tx bytes : [39860043360 28688790100 31775670400] + hccs rx packets : [2091328898 760745014 4126456722] + hccs rx bytes : [41826577960 15214900280 82529134440] + hccs retry count : [0 0 0] + hccs error count : [0 0 0] + "#; + let s = parse_hccs_text(text, 0, 0, 0).expect("parse"); + assert!(s.health_ok); + assert_eq!(s.tx_bytes, 39860043360 + 28688790100 + 31775670400); + assert_eq!(s.rx_packets, 2091328898 + 760745014 + 4126456722); + assert_eq!(s.error_count, 0); + } +} diff --git a/probing/extensions/gpu/src/extensions/backend/npu/mod.rs b/probing/extensions/gpu/src/extensions/backend/npu/mod.rs new file mode 100644 index 00000000..f12ad03b --- /dev/null +++ b/probing/extensions/gpu/src/extensions/backend/npu/mod.rs @@ -0,0 +1,131 @@ +//! Huawei Ascend NPU backend (DCMI preferred, `npu-smi` fallback). + +mod dcmi; +#[allow(dead_code)] +mod hccs; +mod npu_smi; + +use std::sync::Arc; + +use once_cell::sync::Lazy; + +use super::traits::{GpuBackend, GpuBackendKind, GpuDeviceInfo, GpuMemoryModel, GpuMemorySample}; + +pub use npu_smi::read_utilization_by_index; + +#[derive(Clone, Debug, PartialEq, Eq)] +pub struct NpuBackend { + device_ids: Arc>, +} + +static NPU_BACKEND: Lazy> = Lazy::new(probe_npu_backend); + +impl NpuBackend { + pub fn try_load() -> Option { + NPU_BACKEND.clone() + } + + pub fn device_count(&self) -> usize { + self.device_ids.len() + } +} + +impl GpuBackend for NpuBackend { + fn kind(&self) -> GpuBackendKind { + GpuBackendKind::Npu + } + + fn probe_devices(&self) -> Vec { + self.device_ids + .iter() + .map(|&ordinal| { + let detail = npu_smi::read_device_detail(ordinal); + let chip = detail.chip_name.clone(); + GpuDeviceInfo { + backend: GpuBackendKind::Npu, + ordinal, + name: chip + .clone() + .unwrap_or_else(|| format!("ascend-npu:{ordinal}")), + uuid: None, + compute_capability: None, + total_mem_bytes: detail.total_mem_bytes, + memory_model: GpuMemoryModel::Dedicated, + chip, + registry_id: None, + } + }) + .collect() + } + + fn sample_memory(&self, ordinal: i32) -> Option { + if !self.device_ids.contains(&ordinal) { + return None; + } + let detail = npu_smi::read_device_detail(ordinal); + let stats = npu_smi::read_utilization_by_index().get(&ordinal).copied(); + let chip = detail.chip_name.clone(); + let (free_bytes, total_bytes) = if detail.total_mem_bytes > 0 { + (detail.free_mem_bytes, detail.total_mem_bytes) + } else if let Some(s) = stats { + ( + s.hbm_total_bytes.saturating_sub(s.hbm_used_bytes), + s.hbm_total_bytes, + ) + } else { + (0, 0) + }; + Some(GpuMemorySample { + backend: self.kind(), + ordinal, + name: chip + .clone() + .unwrap_or_else(|| format!("ascend-npu:{ordinal}")), + free_bytes, + total_bytes, + memory_model: GpuMemoryModel::Dedicated, + chip, + // Also set here so callers that skip collector merge still see util. + gpu_util_pct: stats.map(|s| s.ai_core_util_pct), + mem_controller_util_pct: stats.map(|s| s.hbm_util_pct), + // Preserve compatibility with the existing web snapshot columns. + renderer_util_pct: stats.and_then(|s| s.aivector_util_pct), + tiler_util_pct: stats.and_then(|s| s.hbm_bw_util_pct), + driver_mem_bytes: None, + vector_util_pct: stats.and_then(|s| s.aivector_util_pct), + mem_bandwidth_util_pct: stats.and_then(|s| s.hbm_bw_util_pct), + temperature_c: stats.and_then(|s| s.temp_c), + power_w: stats.and_then(|s| s.power_w), + }) + } +} + +fn probe_npu_backend() -> Option { + let device_ids = npu_smi::list_device_ids(); + if device_ids.is_empty() { + log::debug!("NPU backend unavailable (DCMI/npu-smi missing or no devices)"); + return None; + } + let source = npu_smi::read_utilization_by_index() + .values() + .next() + .map(|s| s.source) + .unwrap_or("unknown"); + log::info!( + "NPU backend probe: {} device(s) via {source}", + device_ids.len() + ); + Some(NpuBackend { + device_ids: Arc::new(device_ids), + }) +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn npu_backend_probe_does_not_panic() { + let _ = NpuBackend::try_load(); + } +} diff --git a/probing/extensions/gpu/src/extensions/backend/npu/npu_smi.rs b/probing/extensions/gpu/src/extensions/backend/npu/npu_smi.rs new file mode 100644 index 00000000..837aca74 --- /dev/null +++ b/probing/extensions/gpu/src/extensions/backend/npu/npu_smi.rs @@ -0,0 +1,599 @@ +use std::collections::HashMap; +use std::io::Read; +use std::process::{Command, Stdio}; +use std::sync::Mutex; +use std::thread; +use std::time::{Duration, Instant}; + +/// Per-logical-device stats (Phy-ID / torch device ordinal on Ascend A3). +#[allow(dead_code)] +#[derive(Debug, Clone, Copy)] +pub struct NpuDeviceStats { + /// Prefer `NPU Utilization(%)`, else `Aicore Usage Rate(%)`. + pub ai_core_util_pct: f32, + pub hbm_util_pct: f32, + pub hbm_total_bytes: u64, + pub hbm_used_bytes: u64, + /// `Aivector Usage Rate(%)` — closest Ascend proxy for Minder "tensor activity". + pub aivector_util_pct: Option, + /// `HBM Bandwidth Usage Rate(%)`. + pub hbm_bw_util_pct: Option, + pub temp_c: Option, + pub power_w: Option, + /// `"dcmi"` or `"smi"`. + pub source: &'static str, +} + +impl Default for NpuDeviceStats { + fn default() -> Self { + Self { + ai_core_util_pct: 0.0, + hbm_util_pct: 0.0, + hbm_total_bytes: 0, + hbm_used_bytes: 0, + aivector_util_pct: None, + hbm_bw_util_pct: None, + temp_c: None, + power_w: None, + source: "smi", + } + } +} + +/// `npu-smi info -t usages` is ~0.5–1s per card. Without caching, one collector +/// tick calls it once in `sample_all` plus again per `sample_memory` → tens of +/// seconds before the first `gpu.utilization` row appears. +const UTIL_CACHE_TTL: Duration = Duration::from_millis(500); +/// Per-invocation timeout so a stuck `npu-smi` cannot block the collector forever. +const NPU_SMI_TIMEOUT: Duration = Duration::from_secs(5); + +struct UtilCache { + at: Instant, + map: HashMap, +} + +static UTIL_CACHE: Mutex> = Mutex::new(None); + +#[derive(Debug, Clone, Default)] +pub struct NpuDeviceDetail { + pub chip_name: Option, + pub total_mem_bytes: u64, + pub free_mem_bytes: u64, +} + +fn npu_smi_cmd() -> Command { + let mut cmd = Command::new("npu-smi"); + // Ascend images often need driver libs on LD_LIBRARY_PATH for npu-smi. + let extra = "/usr/local/Ascend/driver/lib64/common:/usr/local/Ascend/driver/lib64/driver:/usr/local/Ascend/driver/lib64"; + let ld = match std::env::var("LD_LIBRARY_PATH") { + Ok(v) if v.contains("Ascend/driver") => v, + Ok(v) if !v.is_empty() => format!("{extra}:{v}"), + _ => extra.to_string(), + }; + cmd.env("LD_LIBRARY_PATH", ld); + cmd +} + +/// Public wrapper for sibling modules (`hccs`). +pub(crate) fn run_npu_smi_text(args: &[&str]) -> Option { + let output = run_npu_smi(args)?; + if !output.status.success() { + return None; + } + Some(String::from_utf8_lossy(&output.stdout).into_owned()) +} + +fn run_npu_smi(args: &[&str]) -> Option { + let mut child = npu_smi_cmd() + .args(args) + .stdout(Stdio::piped()) + .stderr(Stdio::piped()) + .spawn() + .ok()?; + let stdout = child.stdout.take()?; + let stderr = child.stderr.take()?; + // Drain both pipes while the process runs. Waiting for exit before reading + // can deadlock when npu-smi emits more than the OS pipe buffer. + let stdout_reader = thread::spawn(move || { + let mut bytes = Vec::new(); + let _ = std::io::BufReader::new(stdout).read_to_end(&mut bytes); + bytes + }); + let stderr_reader = thread::spawn(move || { + let mut bytes = Vec::new(); + let _ = std::io::BufReader::new(stderr).read_to_end(&mut bytes); + bytes + }); + let deadline = Instant::now() + NPU_SMI_TIMEOUT; + let status = loop { + match child.try_wait() { + Ok(Some(status)) => break Some(status), + Ok(None) => { + if Instant::now() >= deadline { + let _ = child.kill(); + let _ = child.wait(); + break None; + } + std::thread::sleep(Duration::from_millis(20)); + } + Err(_) => { + let _ = child.kill(); + let _ = child.wait(); + break None; + } + } + }; + let stdout = stdout_reader.join().unwrap_or_default(); + let stderr = stderr_reader.join().unwrap_or_default(); + status.map(|status| std::process::Output { + status, + stdout, + stderr, + }) +} + +/// Return visible logical device ordinals (Phy-ID style: NPU×chip). +pub fn list_device_ids() -> Vec { + let util = read_utilization_by_index(); + if !util.is_empty() { + let mut ids: Vec = util.keys().copied().collect(); + ids.sort_unstable(); + return ids; + } + // Fallback: NPU ID list only (single-chip cards). + if !npu_smi_available() { + return Vec::new(); + } + let Some(output) = run_npu_smi(&["info", "-l"]) else { + return Vec::new(); + }; + if !output.status.success() { + return Vec::new(); + } + let mut ids = Vec::new(); + for line in String::from_utf8_lossy(&output.stdout).lines() { + let Some(id) = parse_kv_i32(line, &["NPU ID", "NpuID", "Device ID"]) else { + continue; + }; + if !ids.contains(&id) { + ids.push(id); + } + } + ids.sort_unstable(); + ids +} + +pub fn read_device_detail(ordinal: i32) -> NpuDeviceDetail { + let mut detail = NpuDeviceDetail { + chip_name: Some("Ascend910".to_string()), + ..Default::default() + }; + if let Some(stats) = read_utilization_by_index().get(&ordinal) { + detail.total_mem_bytes = stats.hbm_total_bytes; + detail.free_mem_bytes = stats.hbm_total_bytes.saturating_sub(stats.hbm_used_bytes); + return detail; + } + // Fallback: query by assuming ordinal == NPU ID (single-chip). + if !npu_smi_available() { + return detail; + } + if let Some(output) = run_npu_smi(&["info", "-t", "board", "-i", &ordinal.to_string()]) { + if output.status.success() { + for line in String::from_utf8_lossy(&output.stdout).lines() { + if let Some(name) = parse_kv_str( + line, + &["Chip Name", "Product Name", "Product Type", "Model"], + ) { + if name != "NA" { + detail.chip_name = Some(name); + break; + } + } + } + } + } + if let Some(output) = run_npu_smi(&["info", "-t", "memory", "-i", &ordinal.to_string()]) { + if output.status.success() { + let text = String::from_utf8_lossy(&output.stdout); + let total_mb = parse_memory_mb(&text, &["HBM Capacity", "Total Capacity", "HBM Total"]); + if total_mb > 0 { + detail.total_mem_bytes = total_mb.saturating_mul(1024 * 1024); + // memory -t often has no Used; leave free==total if unknown. + detail.free_mem_bytes = detail.total_mem_bytes; + } + } + } + detail +} + +/// Batch-read utilization keyed by logical Phy-ID (NPU_ID * chips + Chip_ID). +/// +/// Prefers in-process DCMI; falls back to `npu-smi`. Cached + single-flight so +/// collector and `sample_memory` share one sweep per TTL window. +pub fn read_utilization_by_index() -> HashMap { + // Hold the lock for the whole refresh (single-flight). Concurrent + // collector + devices/SQL probes used to launch overlapping `npu-smi` + // sweeps; on Ascend that can hang and leave `gpu.utilization` empty forever. + let Ok(mut guard) = UTIL_CACHE.lock() else { + return HashMap::new(); + }; + if let Some(cache) = guard.as_ref() { + if cache.at.elapsed() < UTIL_CACHE_TTL { + return cache.map.clone(); + } + } + + let map = super::dcmi::read_utilization_by_index() + .filter(|m| !m.is_empty()) + .unwrap_or_else(read_utilization_by_index_uncached); + *guard = Some(UtilCache { + at: Instant::now(), + map: map.clone(), + }); + map +} + +fn read_utilization_by_index_uncached() -> HashMap { + if !npu_smi_available() { + return HashMap::new(); + } + + let npu_ids = list_npu_card_ids(); + let mut map = HashMap::new(); + for npu_id in npu_ids { + let Some(output) = run_npu_smi(&["info", "-t", "usages", "-i", &npu_id.to_string()]) else { + continue; + }; + if !output.status.success() { + continue; + } + let text = String::from_utf8_lossy(&output.stdout); + let chip_count = parse_kv_i32_from_text(&text, &["Chip Count"]) + .unwrap_or(1) + .max(1); + for chip in parse_usage_chips(&text) { + let phy_id = npu_id * chip_count + chip.chip_id; + let hbm_total = chip.hbm_capacity_mb.saturating_mul(1024 * 1024); + let hbm_used = if chip.hbm_usage_rate >= 0.0 { + ((chip.hbm_capacity_mb as f32) * chip.hbm_usage_rate / 100.0).round() as u64 + * 1024 + * 1024 + } else { + 0 + }; + let util = if chip.npu_util_pct >= 0.0 { + chip.npu_util_pct + } else if chip.aicore_util_pct >= 0.0 { + chip.aicore_util_pct + } else { + 0.0 + }; + map.insert( + phy_id, + NpuDeviceStats { + ai_core_util_pct: util, + hbm_util_pct: chip.hbm_usage_rate.max(0.0), + hbm_total_bytes: hbm_total, + hbm_used_bytes: hbm_used, + aivector_util_pct: (chip.aivector_util_pct >= 0.0) + .then_some(chip.aivector_util_pct), + hbm_bw_util_pct: (chip.hbm_bw_util_pct >= 0.0).then_some(chip.hbm_bw_util_pct), + temp_c: None, + power_w: None, + source: "smi", + }, + ); + } + // Power is a separate `-t power` call; fill chips belonging to this card. + if let Some(power_text) = + run_npu_smi_text(&["info", "-t", "power", "-i", &npu_id.to_string()]) + { + let powers = parse_power_by_chip(&power_text); + let card_power = powers + .iter() + .find(|(chip, _)| *chip < 0) + .map(|(_, w)| *w) + .or_else(|| powers.first().map(|(_, w)| *w)); + for (phy_id, stats) in map.iter_mut() { + if *phy_id / chip_count != npu_id { + continue; + } + let chip_id = *phy_id % chip_count; + stats.power_w = powers + .iter() + .find(|(c, _)| *c == chip_id) + .map(|(_, w)| *w) + .or(card_power); + } + } + } + map +} + +/// Parse `npu-smi info -t power`. Chip ID `< 0` means card-level aggregate. +fn parse_power_by_chip(text: &str) -> Vec<(i32, f32)> { + let mut out = Vec::new(); + let mut cur_chip: Option = None; + let mut saw_chip = false; + for line in text.lines() { + if let Some(id) = parse_kv_i32(line, &["Chip ID"]) { + cur_chip = Some(id); + saw_chip = true; + continue; + } + if let Some(w) = parse_power_watts_line(line) { + let chip = if saw_chip { cur_chip.unwrap_or(0) } else { -1 }; + out.push((chip, w)); + } + } + out +} + +fn parse_power_watts_line(line: &str) -> Option { + let (key, value) = split_kv(line)?; + if !key_matches( + key, + &[ + "NPU Real-time Power", + "NPU Power", + "Power Dissipation", + "Power", + ], + ) { + return None; + } + // Avoid matching "Power Limit" style keys when possible. + let kl = key.to_ascii_lowercase(); + if kl.contains("limit") || kl.contains("cap") { + return None; + } + value + .split_whitespace() + .next()? + .trim_end_matches('W') + .parse::() + .ok() + .filter(|v| *v >= 0.0) +} + +#[derive(Debug, Default)] +struct ChipUsage { + chip_id: i32, + aicore_util_pct: f32, + npu_util_pct: f32, + aivector_util_pct: f32, + hbm_usage_rate: f32, + hbm_bw_util_pct: f32, + hbm_capacity_mb: u64, +} + +fn parse_usage_chips(text: &str) -> Vec { + let mut chips = Vec::new(); + let mut cur = ChipUsage { + aicore_util_pct: -1.0, + npu_util_pct: -1.0, + aivector_util_pct: -1.0, + hbm_usage_rate: -1.0, + hbm_bw_util_pct: -1.0, + ..Default::default() + }; + let mut saw = false; + + for line in text.lines() { + if let Some(v) = parse_percent_line( + line, + &[ + "Aicore Usage Rate", + "Ai Core Usage Rate", + "AICore Usage Rate", + ], + ) { + cur.aicore_util_pct = v; + saw = true; + } else if let Some(v) = parse_percent_line(line, &["NPU Utilization"]) { + cur.npu_util_pct = v; + saw = true; + } else if let Some(v) = + parse_percent_line(line, &["Aivector Usage Rate", "AI Vector Usage Rate"]) + { + cur.aivector_util_pct = v; + saw = true; + } else if let Some(v) = parse_percent_line(line, &["HBM Bandwidth Usage Rate"]) { + cur.hbm_bw_util_pct = v; + saw = true; + } else if let Some(v) = parse_percent_line(line, &["HBM Usage Rate", "Memory Usage Rate"]) { + // Prefer the capacity fill %; do not overwrite with bandwidth %. + if !key_matches( + split_kv(line).map(|(k, _)| k).unwrap_or(""), + &["HBM Bandwidth Usage Rate"], + ) { + cur.hbm_usage_rate = v; + saw = true; + } + } else if let Some(mb) = parse_memory_mb_line(line, &["HBM Capacity"]) { + cur.hbm_capacity_mb = mb; + saw = true; + } else if let Some(id) = parse_kv_i32(line, &["Chip ID"]) { + cur.chip_id = id; + if saw { + chips.push(std::mem::take(&mut cur)); + cur = ChipUsage { + aicore_util_pct: -1.0, + npu_util_pct: -1.0, + aivector_util_pct: -1.0, + hbm_usage_rate: -1.0, + hbm_bw_util_pct: -1.0, + ..Default::default() + }; + saw = false; + } + } + } + if saw { + chips.push(cur); + } + chips +} + +pub(super) fn list_npu_card_ids() -> Vec { + let Some(output) = run_npu_smi(&["info", "-l"]) else { + return Vec::new(); + }; + if !output.status.success() { + return Vec::new(); + } + let mut ids = Vec::new(); + for line in String::from_utf8_lossy(&output.stdout).lines() { + if let Some(id) = parse_kv_i32(line, &["NPU ID", "NpuID", "Device ID"]) { + if !ids.contains(&id) { + ids.push(id); + } + } + } + ids.sort_unstable(); + ids +} + +fn npu_smi_available() -> bool { + run_npu_smi(&["--help"]) + .map(|o| o.status.success() || !o.stdout.is_empty() || !o.stderr.is_empty()) + .unwrap_or(false) +} + +fn normalize_key(key: &str) -> String { + key.trim() + .trim_end_matches("(%)") + .trim_end_matches("(MB)") + .trim_end_matches("(page)") + .trim_end_matches("(C)") + .trim_end_matches("(MHz)") + .trim() + .to_string() +} + +fn key_matches(key: &str, patterns: &[&str]) -> bool { + let norm = normalize_key(key); + patterns.iter().any(|p| { + norm.eq_ignore_ascii_case(p) || norm.to_ascii_lowercase().contains(&p.to_ascii_lowercase()) + }) +} + +fn parse_kv_i32(line: &str, keys: &[&str]) -> Option { + let (key, value) = split_kv(line)?; + if !key_matches(key, keys) { + return None; + } + value.split_whitespace().next()?.parse().ok() +} + +fn parse_kv_i32_from_text(text: &str, keys: &[&str]) -> Option { + text.lines().find_map(|line| parse_kv_i32(line, keys)) +} + +fn parse_kv_str(line: &str, keys: &[&str]) -> Option { + let (key, value) = split_kv(line)?; + if !key_matches(key, keys) { + return None; + } + let trimmed = value.trim(); + if trimmed.is_empty() { + None + } else { + Some(trimmed.to_string()) + } +} + +fn split_kv(line: &str) -> Option<(&str, &str)> { + let (key, value) = line.split_once(':')?; + Some((key.trim(), value.trim())) +} + +fn parse_percent_line(line: &str, keys: &[&str]) -> Option { + let (key, value) = split_kv(line)?; + if !key_matches(key, keys) { + return None; + } + let num = value.split_whitespace().next()?; + num.trim_end_matches('%') + .parse::() + .ok() + .map(|v| v.clamp(0.0, 100.0)) +} + +fn parse_memory_mb_line(line: &str, keys: &[&str]) -> Option { + let (key, value) = split_kv(line)?; + if !key_matches(key, keys) { + return None; + } + let num = value.split_whitespace().next()?; + num.parse::().ok().map(|v| v.round() as u64) +} + +fn parse_memory_mb(text: &str, keys: &[&str]) -> u64 { + text.lines() + .find_map(|line| parse_memory_mb_line(line, keys)) + .unwrap_or(0) +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn parse_real_a3_usages_block() { + let text = r#" + NPU ID : 0 + Chip Count : 2 + + DDR Capacity(MB) : 0 + HBM Capacity(MB) : 65536 + HBM Usage Rate(%) : 16 + HBM Bandwidth Usage Rate(%) : 27 + Aicore Usage Rate(%) : 2 + Aivector Usage Rate(%) : 31 + NPU Utilization(%) : 65 + Chip ID : 0 + + HBM Capacity(MB) : 65536 + HBM Usage Rate(%) : 16 + Aicore Usage Rate(%) : 12 + NPU Utilization(%) : 58 + Chip ID : 1 + "#; + let chips = parse_usage_chips(text); + assert_eq!(chips.len(), 2); + assert_eq!(chips[0].chip_id, 0); + assert_eq!(chips[0].npu_util_pct, 65.0); + assert_eq!(chips[0].aicore_util_pct, 2.0); + assert_eq!(chips[0].hbm_usage_rate, 16.0); + assert_eq!(chips[0].hbm_bw_util_pct, 27.0); + assert_eq!(chips[0].aivector_util_pct, 31.0); + assert_eq!(chips[0].hbm_capacity_mb, 65536); + assert_eq!(chips[1].chip_id, 1); + assert_eq!(chips[1].npu_util_pct, 58.0); + } + + #[test] + fn key_matches_strips_units() { + assert!(key_matches("Aicore Usage Rate(%)", &["Aicore Usage Rate"])); + assert!(key_matches("HBM Capacity(MB)", &["HBM Capacity"])); + assert!(key_matches("NPU Utilization(%)", &["NPU Utilization"])); + } + + #[test] + fn parse_npu_id_line() { + assert_eq!( + parse_kv_i32(" NPU ID : 3", &["NPU ID"]), + Some(3) + ); + } + + #[test] + fn parse_power_ignores_limit() { + assert_eq!( + parse_power_watts_line("NPU Real-time Power : 164.5 W"), + Some(164.5) + ); + assert_eq!(parse_power_watts_line("Power Limit : 300 W"), None); + } +} diff --git a/probing/extensions/gpu/src/extensions/backend/registry.rs b/probing/extensions/gpu/src/extensions/backend/registry.rs index 957392ff..ca874ba9 100644 --- a/probing/extensions/gpu/src/extensions/backend/registry.rs +++ b/probing/extensions/gpu/src/extensions/backend/registry.rs @@ -11,6 +11,9 @@ use super::apple::AppleSiliconBackend; #[cfg(feature = "cuda")] use super::cuda::CudaBackend; +#[cfg(all(not(target_os = "macos"), not(target_os = "windows")))] +use super::npu::NpuBackend; + static BACKEND_FILTER: Lazy>>> = Lazy::new(|| Mutex::new(None)); /// Restrict which backends are active (`None` = auto-discover all available). @@ -18,7 +21,14 @@ pub fn set_backend_filter(kinds: Option>) { *lock_mutex(&BACKEND_FILTER, "GPU backend filter") = kinds; } -#[cfg_attr(not(any(feature = "cuda", target_os = "macos")), allow(dead_code))] +#[cfg_attr( + not(any( + feature = "cuda", + target_os = "macos", + all(not(target_os = "macos"), not(target_os = "windows")) + )), + allow(dead_code) +)] fn filter_allows(kind: GpuBackendKind) -> bool { match lock_mutex(&BACKEND_FILTER, "GPU backend filter").as_ref() { None => true, @@ -78,10 +88,40 @@ pub fn discover_backends() -> Vec> { } }; - #[cfg(not(any(feature = "cuda", target_os = "macos")))] + let npu: Option> = { + #[cfg(all(not(target_os = "macos"), not(target_os = "windows")))] + { + if filter_allows(GpuBackendKind::Npu) { + let npu = + std::panic::catch_unwind(std::panic::AssertUnwindSafe(NpuBackend::try_load)) + .ok() + .flatten(); + if let Some(npu) = npu { + let count = npu.device_count(); + log::info!("GPU backend loaded: ascend/npu ({count} device(s))"); + Some(Box::new(npu)) + } else { + log::debug!("Ascend NPU backend unavailable (DCMI/npu-smi missing)"); + None + } + } else { + None + } + } + #[cfg(any(target_os = "macos", target_os = "windows"))] + { + None + } + }; + + #[cfg(not(any( + feature = "cuda", + target_os = "macos", + all(not(target_os = "macos"), not(target_os = "windows")) + )))] log::debug!("probing-gpu: no GPU backends available for this target"); - [cuda, apple].into_iter().flatten().collect() + [cuda, apple, npu].into_iter().flatten().collect() } /// Backends selected by env `PROBING_GPU_BACKEND` (default: auto = all discovered). @@ -116,11 +156,13 @@ mod platform_tests { assert_eq!(first.len(), second.len()); } - /// Linux/Windows CI: crate must compile with zero GPU backends when cuda is off. + /// A no-CUDA build must not discover CUDA; platform-native backends may remain. #[test] #[cfg(all(not(target_os = "macos"), not(feature = "cuda")))] - fn non_mac_build_has_no_backends_without_cuda() { - assert!(discover_backends().is_empty()); + fn non_mac_build_has_no_cuda_backend_without_cuda() { + assert!(!discover_backends() + .iter() + .any(|backend| backend.kind() == GpuBackendKind::Cuda)); } /// macOS: Apple backend module is linked; discovery may return devices. diff --git a/probing/extensions/gpu/src/extensions/backend/traits.rs b/probing/extensions/gpu/src/extensions/backend/traits.rs index c3eaaaa4..05a3bc9e 100644 --- a/probing/extensions/gpu/src/extensions/backend/traits.rs +++ b/probing/extensions/gpu/src/extensions/backend/traits.rs @@ -5,6 +5,8 @@ use std::fmt; pub enum GpuBackendKind { /// NVIDIA CUDA (discrete or datacenter GPU). Cuda, + /// Huawei Ascend NPU (via DCMI / npu-smi). + Npu, /// AMD ROCm / HIP (reserved). Rocm, /// Apple Metal / AGX (M1–M4 integrated GPU). @@ -15,6 +17,7 @@ impl GpuBackendKind { pub fn as_str(self) -> &'static str { match self { Self::Cuda => "cuda", + Self::Npu => "npu", Self::Rocm => "rocm", Self::Metal => "metal", } @@ -23,6 +26,7 @@ impl GpuBackendKind { pub fn parse(s: &str) -> Option { match s.trim().to_ascii_lowercase().as_str() { "cuda" | "nvidia" => Some(Self::Cuda), + "npu" | "ascend" | "huawei" | "hccl" => Some(Self::Npu), "rocm" | "hip" | "amd" => Some(Self::Rocm), "metal" | "mps" | "apple" | "agx" | "apple-silicon" => Some(Self::Metal), _ => None, @@ -87,6 +91,14 @@ pub struct GpuMemorySample { pub tiler_util_pct: Option, /// Unified-memory only: bytes attributed to GPU driver (IORegistry). pub driver_mem_bytes: Option, + /// Backend-specific vector/tensor execution utilization. + pub vector_util_pct: Option, + /// Device-memory bandwidth utilization. + pub mem_bandwidth_util_pct: Option, + /// Device temperature in degrees Celsius. + pub temperature_c: Option, + /// Device power draw in watts. + pub power_w: Option, } impl GpuMemorySample { diff --git a/probing/extensions/gpu/src/extensions/collector.rs b/probing/extensions/gpu/src/extensions/collector.rs index 42125780..e5fb442b 100644 --- a/probing/extensions/gpu/src/extensions/collector.rs +++ b/probing/extensions/gpu/src/extensions/collector.rs @@ -17,6 +17,9 @@ use super::backend::{discover_backends, selected_backends, GpuBackend, GpuMemory #[cfg(feature = "cuda")] use super::backend::read_utilization_by_index; +#[cfg(all(not(target_os = "macos"), not(target_os = "windows")))] +use super::backend::read_npu_utilization_by_index; + const CHUNK_SIZE: u32 = 4096; const NUM_CHUNKS: u32 = 8; const DEFAULT_SAMPLE_INTERVAL_MS: u64 = 1000; @@ -55,7 +58,7 @@ pub fn autostart_interval_ms() -> Option { return Some(DEFAULT_SAMPLE_INTERVAL_MS); } - if discover_backends().is_empty() { + if selected_backends().is_empty() { return None; } @@ -114,6 +117,10 @@ fn utilization_schema() -> Schema { .col("tiler_util_pct", DType::F32) .col("driver_mem_bytes", DType::I64) .col("wall_ns", DType::I64) + .col("vector_util_pct", DType::F32) + .col("mem_bandwidth_util_pct", DType::F32) + .col("temperature_c", DType::F32) + .col("power_w", DType::F32) } #[derive(Debug, Clone)] @@ -160,6 +167,10 @@ fn push_utilization_row(table: &mut ExposedTable, ts: i64, wall_ns: u64, sample: Value::F32(opt_f32(sample.tiler_util_pct)), Value::I64(sample.driver_mem_bytes.unwrap_or(0) as i64), Value::I64(wall_ns as i64), + Value::F32(opt_f32(sample.vector_util_pct)), + Value::F32(opt_f32(sample.mem_bandwidth_util_pct)), + Value::F32(opt_f32(sample.temperature_c)), + Value::F32(opt_f32(sample.power_w)), ]) { log::warn!("gpu collector: push_row failed for gpu.utilization"); } @@ -173,27 +184,38 @@ fn sample_all(backends: &[Box]) -> Vec { #[cfg(feature = "cuda")] let nvidia_utils = read_utilization_by_index(); + #[cfg(all(not(target_os = "macos"), not(target_os = "windows")))] + let npu_utils = if backends + .iter() + .any(|backend| backend.kind() == super::backend::GpuBackendKind::Npu) + { + read_npu_utilization_by_index() + } else { + Default::default() + }; + let mut samples = Vec::new(); for backend in backends { for device in backend.probe_devices() { - if let Some(sample) = backend.sample_memory(device.ordinal) { - let sample = { - #[cfg(feature = "cuda")] - { - let mut s = sample; - if s.backend == super::backend::GpuBackendKind::Cuda { - if let Some(u) = nvidia_utils.get(&s.ordinal) { - s.gpu_util_pct = Some(u.gpu_util_pct); - s.mem_controller_util_pct = Some(u.mem_controller_util_pct); - } - } - s + if let Some(mut sample) = backend.sample_memory(device.ordinal) { + #[cfg(feature = "cuda")] + if sample.backend == super::backend::GpuBackendKind::Cuda { + if let Some(u) = nvidia_utils.get(&sample.ordinal) { + sample.gpu_util_pct = Some(u.gpu_util_pct); + sample.mem_controller_util_pct = Some(u.mem_controller_util_pct); } - #[cfg(not(feature = "cuda"))] - { - sample + } + + #[cfg(all(not(target_os = "macos"), not(target_os = "windows")))] + if sample.backend == super::backend::GpuBackendKind::Npu { + if let Some(u) = npu_utils.get(&sample.ordinal) { + sample.gpu_util_pct = Some(u.ai_core_util_pct); + sample.mem_controller_util_pct = Some(u.hbm_util_pct); + sample.renderer_util_pct = u.aivector_util_pct; + sample.tiler_util_pct = u.hbm_bw_util_pct; } - }; + } + samples.push(sample); } } @@ -325,6 +347,7 @@ mod tests { fn clear_gpu_env() { std::env::remove_var("PROBING_GPU"); std::env::remove_var("PROBING_GPU_SAMPLE_MS"); + std::env::remove_var("PROBING_GPU_BACKEND"); } #[test] @@ -354,6 +377,15 @@ mod tests { clear_gpu_env(); } + #[test] + fn autostart_respects_backend_filter() { + let _guard = ENV_TEST_LOCK.lock().unwrap(); + clear_gpu_env(); + std::env::set_var("PROBING_GPU_BACKEND", "unsupported"); + assert!(autostart_interval_ms().is_none()); + clear_gpu_env(); + } + #[test] #[cfg(target_os = "macos")] fn autostart_auto_when_backend_present() { @@ -368,10 +400,14 @@ mod tests { #[test] #[cfg(all(not(target_os = "macos"), not(feature = "cuda")))] - fn autostart_auto_off_without_backend() { + fn autostart_auto_matches_native_backend_discovery() { let _guard = ENV_TEST_LOCK.lock().unwrap(); clear_gpu_env(); - assert!(autostart_interval_ms().is_none()); + let has_backend = !discover_backends().is_empty(); + assert_eq!( + autostart_interval_ms(), + has_backend.then_some(DEFAULT_SAMPLE_INTERVAL_MS) + ); } #[test] diff --git a/probing/extensions/gpu/src/extensions/extension.rs b/probing/extensions/gpu/src/extensions/extension.rs index 4b3cb851..14928f6b 100644 --- a/probing/extensions/gpu/src/extensions/extension.rs +++ b/probing/extensions/gpu/src/extensions/extension.rs @@ -13,7 +13,7 @@ pub struct GpuProbeExtension { #[option(aliases = ["sample_interval", "interval", "gpu.interval"])] gpu_sample_interval_ms: Maybe, - /// Backend filter: `auto`, `cuda`, `rocm`, `metal`, or comma-separated list. + /// Backend filter: `auto`, `cuda`, `npu`, `rocm`, `metal`, or comma-separated list. #[option(aliases = ["backend", "gpu.backend"])] gpu_backend: Maybe, } diff --git a/python/probing_hook.py b/python/probing_hook.py index 7e3eda15..a9366bec 100644 --- a/python/probing_hook.py +++ b/python/probing_hook.py @@ -6,6 +6,12 @@ import sys +def _probing_requests_activation() -> bool: + """Avoid loading the native runtime in warmup and helper processes.""" + token = os.environ.get("PROBING", "0").strip().lower() + return token not in ("", "0", "off", "false", "no") + + def _is_elastic_supervisor() -> bool: """Match Rust ``is_elastic_supervisor`` without importing ``probing``.""" if os.environ.get("LOCAL_RANK") is not None or os.environ.get("RANK") is not None: @@ -34,7 +40,7 @@ def _is_elastic_supervisor() -> bool: _sup = _is_elastic_supervisor() -if not _sup: +if _probing_requests_activation() and not _sup: from probing.site_hook import run_site_hook run_site_hook() From 69baddc0ce8b4a45bff5f11581bbfd4acff65195 Mon Sep 17 00:00:00 2001 From: Vincent Date: Tue, 18 Aug 2026 16:39:49 +0800 Subject: [PATCH 2/2] fix(gpu): allow cfg-dependent sample mutability Suppress the unused mutability warning only on targets where neither CUDA nor NPU enrichment can modify samples. Co-authored-by: Cursor --- probing/extensions/gpu/src/extensions/collector.rs | 7 +++++++ 1 file changed, 7 insertions(+) diff --git a/probing/extensions/gpu/src/extensions/collector.rs b/probing/extensions/gpu/src/extensions/collector.rs index e5fb442b..6f6ac396 100644 --- a/probing/extensions/gpu/src/extensions/collector.rs +++ b/probing/extensions/gpu/src/extensions/collector.rs @@ -180,6 +180,13 @@ fn opt_f32(v: Option) -> f32 { v.unwrap_or(-1.0) } +#[cfg_attr( + not(any( + feature = "cuda", + all(not(target_os = "macos"), not(target_os = "windows")) + )), + allow(unused_mut) +)] fn sample_all(backends: &[Box]) -> Vec { #[cfg(feature = "cuda")] let nvidia_utils = read_utilization_by_index();