use std::sync::{Mutex, OnceLock};
use std::time::{Duration, Instant};
use crate::protocol::NodeResources;
const FAST_PROFILE_CACHE_TTL: Duration = Duration::from_secs(5);
static FAST_PROFILE_CACHE: OnceLock<Mutex<Option<CachedRuntimeProfileEntry>>> = OnceLock::new();
#[derive(Clone, Copy, Debug, Default, PartialEq, Eq, Hash)]
pub enum AccelerationMode {
#[default]
Gpu,
Avx512,
Avx2,
Neon,
Generic,
}
#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)]
pub enum GpuBackend {
Cuda,
Rocm,
Vulkan,
Directml,
Metal,
Npu,
Cpu,
}
impl GpuBackend {
pub const fn as_str(self) -> &'static str {
match self {
Self::Cuda => "cuda",
Self::Rocm => "rocm",
Self::Vulkan => "vulkan",
Self::Directml => "directml",
Self::Metal => "metal",
Self::Npu => "npu",
Self::Cpu => "cpu",
}
}
}
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum ProbeMode {
Fast,
Full,
}
impl ProbeMode {
pub const fn as_str(self) -> &'static str {
match self {
Self::Fast => "fast",
Self::Full => "full",
}
}
}
impl AccelerationMode {
pub const fn as_str(self) -> &'static str {
match self {
Self::Gpu => "GPU",
Self::Avx512 => "AVX-512",
Self::Avx2 => "AVX2",
Self::Neon => "NEON",
Self::Generic => "Generic CPU",
}
}
}
#[derive(Clone, Debug)]
pub struct RuntimeProfile {
pub node_resources: NodeResources,
pub logical_cores: usize,
pub recommended_workers: usize,
pub acceleration_mode: AccelerationMode,
pub gpu_backend: GpuBackend,
pub xdp_supported: bool,
pub detection_source: String,
pub probe_mode: ProbeMode,
}
impl RuntimeProfile {
pub fn summary(&self) -> String {
format!(
"Host Runtime Profile\n\
====================\n\
Node ID: {}\n\
Logical cores: {}\n\
Recommended workers: {}\n\
System memory: {:.1} GB\n\
GPU VRAM: {:.1} GB\n\
Compute capability: {}\n\
GPU: {}\n\
GPU Backend: {}\n\
Acceleration: {}\n\
XDP support: {}\n\
Detection source: {}\n\
Probe mode: {}\n",
self.node_resources.id,
self.logical_cores,
self.recommended_workers,
self.node_resources.system_memory_gb,
self.node_resources.vram_gb,
if self.node_resources.compute_capability.is_empty() {
"cpu"
} else {
&self.node_resources.compute_capability
},
self.node_resources
.gpu_name
.as_deref()
.unwrap_or("Not detected"),
self.gpu_backend.as_str(),
self.acceleration_mode.as_str(),
if self.xdp_supported {
"available"
} else {
"unavailable"
},
self.detection_source,
self.probe_mode.as_str()
)
}
}
#[derive(Clone, Debug)]
struct CachedRuntimeProfileEntry {
captured_at: Instant,
profile: RuntimeProfile,
}
pub fn detect_runtime_profile(node_id: impl Into<String>) -> RuntimeProfile {
detect_runtime_profile_with_mode(node_id, ProbeMode::Fast)
}
pub fn detect_runtime_profile_with_mode(
node_id: impl Into<String>,
probe_mode: ProbeMode,
) -> RuntimeProfile {
let node_id = node_id.into();
if let Some(mut profile) = load_cached_runtime_profile(probe_mode) {
profile.node_resources.id = node_id;
return profile;
}
let sp = match probe_mode {
ProbeMode::Fast => crate::system_profile::SystemProfile::detect_fast(),
ProbeMode::Full => crate::system_profile::SystemProfile::detect(),
};
let mut profile: RuntimeProfile = sp.into();
profile.node_resources.id = node_id;
profile.probe_mode = probe_mode;
store_cached_runtime_profile(&profile);
profile
}
pub fn detect_local_node_resources(node_id: impl Into<String>) -> NodeResources {
detect_runtime_profile(node_id).node_resources
}
pub fn detect_runtime_profile_full(node_id: impl Into<String>) -> RuntimeProfile {
detect_runtime_profile_with_mode(node_id, ProbeMode::Full)
}
fn load_cached_runtime_profile(probe_mode: ProbeMode) -> Option<RuntimeProfile> {
if probe_mode != ProbeMode::Fast {
return None;
}
let cache = FAST_PROFILE_CACHE.get_or_init(|| Mutex::new(None));
let guard = cache.lock().ok()?;
let entry = guard.as_ref()?;
if entry.captured_at.elapsed() > FAST_PROFILE_CACHE_TTL {
return None;
}
Some(entry.profile.clone())
}
fn store_cached_runtime_profile(profile: &RuntimeProfile) {
if profile.probe_mode != ProbeMode::Fast {
return;
}
let cache = FAST_PROFILE_CACHE.get_or_init(|| Mutex::new(None));
if let Ok(mut guard) = cache.lock() {
*guard = Some(CachedRuntimeProfileEntry {
captured_at: Instant::now(),
profile: profile.clone(),
});
}
}
pub fn infer_compute_capability_from_name(name: &str) -> String {
let lowered = name.to_ascii_lowercase();
if lowered.contains("rtx 50")
|| lowered.contains("rtx50")
|| lowered.contains("b100")
|| lowered.contains("b200")
{
String::from("9.0")
} else if lowered.contains("rtx 40") || lowered.contains("rtx40") {
String::from("8.9")
} else if lowered.contains("rtx 30") || lowered.contains("rtx30") {
String::from("8.6")
} else if lowered.contains("arc") {
String::from("xe")
} else if lowered.contains("radeon")
|| lowered.contains("amd")
|| lowered.contains("advanced micro")
|| lowered.contains("rx ")
|| lowered.contains("w7900")
|| lowered.contains("w6800")
|| lowered.contains("mi250")
|| lowered.contains("mi300")
|| lowered.contains("mi350")
|| lowered.contains("instinct")
{
if cfg!(feature = "rocm") {
String::from("rocm")
} else {
String::from("gpu")
}
} else if lowered.contains("intel") || lowered.contains("iris") || lowered.contains("uhd") {
String::from("xe")
} else if lowered.contains("apple")
|| lowered.contains("m1")
|| lowered.contains("m2")
|| lowered.contains("m3")
|| lowered.contains("m4")
{
String::from("apple_metal")
} else {
String::from("gpu")
}
}
#[cfg(target_os = "macos")]
pub fn is_apple_silicon() -> bool {
std::process::Command::new("sysctl")
.arg("hw.optional.arm64")
.output()
.map(|output| String::from_utf8_lossy(&output.stdout).contains("1"))
.unwrap_or(false)
}
#[cfg(not(target_os = "macos"))]
pub fn is_apple_silicon() -> bool {
false
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn summary_includes_tuning_fields() {
let profile = RuntimeProfile {
node_resources: NodeResources::new(
"node-a",
24.0,
64.0,
"8.9",
Some(String::from("RTX4090")),
),
logical_cores: 16,
recommended_workers: 10,
acceleration_mode: AccelerationMode::Gpu,
gpu_backend: GpuBackend::Cuda,
xdp_supported: true,
detection_source: String::from("test"),
probe_mode: ProbeMode::Fast,
};
let summary = profile.summary();
assert!(summary.contains("Recommended workers: 10"));
assert!(summary.contains("Acceleration: GPU"));
assert!(summary.contains("Probe mode: fast"));
}
#[test]
fn infers_compute_capability_from_gpu_name() {
assert_eq!(infer_compute_capability_from_name("RTX 4090"), "8.9");
assert_eq!(infer_compute_capability_from_name("Arc A770"), "xe");
#[cfg(feature = "rocm")]
{
assert_eq!(
infer_compute_capability_from_name("AMD Radeon RX 7900 XTX"),
"rocm"
);
assert_eq!(
infer_compute_capability_from_name("AMD Instinct MI250"),
"rocm"
);
}
#[cfg(not(feature = "rocm"))]
{
assert_eq!(
infer_compute_capability_from_name("AMD Radeon RX 7900 XTX"),
"gpu"
);
}
}
#[test]
fn fast_profile_cache_reuses_detected_shape() {
let first = detect_runtime_profile_with_mode("node-a", ProbeMode::Fast);
let second = detect_runtime_profile_with_mode("node-b", ProbeMode::Fast);
assert_eq!(first.logical_cores, second.logical_cores);
assert_eq!(first.acceleration_mode, second.acceleration_mode);
assert_eq!(second.node_resources.id, "node-b");
}
}