1use std::sync::{Arc, OnceLock};
9use std::time::{Duration, Instant};
10
11use arc_swap::ArcSwapOption;
12
13use crate::protocol::NodeResources;
14
15const FAST_PROFILE_CACHE_TTL: Duration = Duration::from_secs(5);
16
17static FAST_PROFILE_CACHE: OnceLock<ArcSwapOption<CachedRuntimeProfileEntry>> = OnceLock::new();
20
21#[derive(Clone, Copy, Debug, Default, PartialEq, Eq, Hash)]
23pub enum AccelerationMode {
24 #[default]
26 Gpu,
27 Avx512,
29 Avx2,
31 Neon,
33 Generic,
35}
36
37#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)]
39pub enum GpuBackend {
40 Cuda,
41 Rocm,
42 Vulkan,
43 Directml,
44 Metal,
45 Npu,
46 Cpu,
47}
48
49impl GpuBackend {
50 pub const fn as_str(self) -> &'static str {
52 match self {
53 Self::Cuda => "cuda",
54 Self::Rocm => "rocm",
55 Self::Vulkan => "vulkan",
56 Self::Directml => "directml",
57 Self::Metal => "metal",
58 Self::Npu => "npu",
59 Self::Cpu => "cpu",
60 }
61 }
62}
63
64#[derive(Clone, Copy, Debug, PartialEq, Eq)]
66pub enum ProbeMode {
67 Fast,
69 Full,
71}
72
73impl ProbeMode {
74 pub const fn as_str(self) -> &'static str {
76 match self {
77 Self::Fast => "fast",
78 Self::Full => "full",
79 }
80 }
81}
82
83impl AccelerationMode {
84 pub const fn as_str(self) -> &'static str {
86 match self {
87 Self::Gpu => "GPU",
88 Self::Avx512 => "AVX-512",
89 Self::Avx2 => "AVX2",
90 Self::Neon => "NEON",
91 Self::Generic => "Generic CPU",
92 }
93 }
94}
95
96#[derive(Clone, Debug)]
98pub struct RuntimeProfile {
99 pub node_resources: NodeResources,
101 pub logical_cores: usize,
103 pub recommended_workers: usize,
105 pub acceleration_mode: AccelerationMode,
107 pub gpu_backend: GpuBackend,
109 pub xdp_supported: bool,
111 pub detection_source: String,
113 pub probe_mode: ProbeMode,
115}
116
117impl RuntimeProfile {
118 pub fn summary(&self) -> String {
120 format!(
121 "Host Runtime Profile\n\
122 ====================\n\
123 Node ID: {}\n\
124 Logical cores: {}\n\
125 Recommended workers: {}\n\
126 System memory: {:.1} GB\n\
127 GPU VRAM: {:.1} GB\n\
128 Compute capability: {}\n\
129 GPU: {}\n\
130 GPU Backend: {}\n\
131 Acceleration: {}\n\
132 XDP support: {}\n\
133 Detection source: {}\n\
134 Probe mode: {}\n",
135 self.node_resources.id,
136 self.logical_cores,
137 self.recommended_workers,
138 self.node_resources.system_memory_gb,
139 self.node_resources.vram_gb,
140 if self.node_resources.compute_capability.is_empty() {
141 "cpu"
142 } else {
143 &self.node_resources.compute_capability
144 },
145 self.node_resources
146 .gpu_name
147 .as_deref()
148 .unwrap_or("Not detected"),
149 self.gpu_backend.as_str(),
150 self.acceleration_mode.as_str(),
151 if self.xdp_supported {
152 "available"
153 } else {
154 "unavailable"
155 },
156 self.detection_source,
157 self.probe_mode.as_str()
158 )
159 }
160}
161
162#[derive(Clone, Debug)]
163struct CachedRuntimeProfileEntry {
164 captured_at: Instant,
165 profile: RuntimeProfile,
166}
167
168pub fn detect_runtime_profile(node_id: impl Into<String>) -> RuntimeProfile {
170 detect_runtime_profile_with_mode(node_id, ProbeMode::Fast)
171}
172
173pub fn detect_runtime_profile_with_mode(
177 node_id: impl Into<String>,
178 probe_mode: ProbeMode,
179) -> RuntimeProfile {
180 let node_id = node_id.into();
181
182 if let Some(mut profile) = load_cached_runtime_profile(probe_mode) {
183 profile.node_resources.id = node_id;
184 return profile;
185 }
186
187 let sp = match probe_mode {
188 ProbeMode::Fast => crate::system_profile::SystemProfile::detect_fast(),
189 ProbeMode::Full => crate::system_profile::SystemProfile::detect(),
190 };
191
192 let mut profile: RuntimeProfile = sp.into();
193 profile.node_resources.id = node_id;
194 profile.probe_mode = probe_mode;
195
196 store_cached_runtime_profile(&profile);
197 profile
198}
199
200pub fn detect_local_node_resources(node_id: impl Into<String>) -> NodeResources {
202 detect_runtime_profile(node_id).node_resources
203}
204
205pub fn detect_runtime_profile_full(node_id: impl Into<String>) -> RuntimeProfile {
207 detect_runtime_profile_with_mode(node_id, ProbeMode::Full)
208}
209
210fn load_cached_runtime_profile(probe_mode: ProbeMode) -> Option<RuntimeProfile> {
211 if probe_mode != ProbeMode::Fast {
212 return None;
213 }
214
215 let cache = FAST_PROFILE_CACHE.get_or_init(|| ArcSwapOption::from(None));
216 let guard = cache.load();
217 let entry = guard.as_ref()?;
218 if entry.captured_at.elapsed() > FAST_PROFILE_CACHE_TTL {
219 return None;
220 }
221
222 Some(entry.profile.clone())
223}
224
225fn store_cached_runtime_profile(profile: &RuntimeProfile) {
226 if profile.probe_mode != ProbeMode::Fast {
227 return;
228 }
229
230 let cache = FAST_PROFILE_CACHE.get_or_init(|| ArcSwapOption::from(None));
231 cache.store(Some(Arc::new(CachedRuntimeProfileEntry {
232 captured_at: Instant::now(),
233 profile: profile.clone(),
234 })));
235}
236
237pub fn infer_compute_capability_from_name(name: &str) -> String {
239 let lowered = name.to_ascii_lowercase();
240 if lowered.contains("rtx 50")
241 || lowered.contains("rtx50")
242 || lowered.contains("b100")
243 || lowered.contains("b200")
244 {
245 String::from("9.0")
246 } else if lowered.contains("rtx 40") || lowered.contains("rtx40") {
247 String::from("8.9")
248 } else if lowered.contains("rtx 30") || lowered.contains("rtx30") {
249 String::from("8.6")
250 } else if lowered.contains("arc") {
251 String::from("xe")
252 } else if lowered.contains("radeon")
253 || lowered.contains("amd")
254 || lowered.contains("advanced micro")
255 || lowered.contains("rx ")
256 || lowered.contains("w7900")
257 || lowered.contains("w6800")
258 || lowered.contains("mi250")
259 || lowered.contains("mi300")
260 || lowered.contains("mi350")
261 || lowered.contains("instinct")
262 {
263 if cfg!(feature = "rocm") {
264 String::from("rocm")
265 } else {
266 String::from("gpu")
267 }
268 } else if lowered.contains("intel") || lowered.contains("iris") || lowered.contains("uhd") {
269 String::from("xe")
270 } else if lowered.contains("apple")
271 || lowered.contains("m1")
272 || lowered.contains("m2")
273 || lowered.contains("m3")
274 || lowered.contains("m4")
275 {
276 String::from("apple_metal")
277 } else {
278 String::from("gpu")
279 }
280}
281
282#[cfg(target_os = "macos")]
284pub fn is_apple_silicon() -> bool {
285 std::process::Command::new("sysctl")
286 .arg("hw.optional.arm64")
287 .output()
288 .map(|output| String::from_utf8_lossy(&output.stdout).contains("1"))
289 .unwrap_or(false)
290}
291
292#[cfg(not(target_os = "macos"))]
294pub fn is_apple_silicon() -> bool {
295 false
296}
297
298#[cfg(test)]
299mod tests {
300 use super::*;
301
302 #[test]
303 fn summary_includes_tuning_fields() {
304 let profile = RuntimeProfile {
305 node_resources: NodeResources::new(
306 "node-a",
307 24.0,
308 64.0,
309 "8.9",
310 Some(String::from("RTX4090")),
311 ),
312 logical_cores: 16,
313 recommended_workers: 10,
314 acceleration_mode: AccelerationMode::Gpu,
315 gpu_backend: GpuBackend::Cuda,
316 xdp_supported: true,
317 detection_source: String::from("test"),
318 probe_mode: ProbeMode::Fast,
319 };
320
321 let summary = profile.summary();
322 assert!(summary.contains("Recommended workers: 10"));
323 assert!(summary.contains("Acceleration: GPU"));
324 assert!(summary.contains("Probe mode: fast"));
325 }
326
327 #[test]
328 fn infers_compute_capability_from_gpu_name() {
329 assert_eq!(infer_compute_capability_from_name("RTX 4090"), "8.9");
330 assert_eq!(infer_compute_capability_from_name("Arc A770"), "xe");
331 #[cfg(feature = "rocm")]
332 {
333 assert_eq!(
334 infer_compute_capability_from_name("AMD Radeon RX 7900 XTX"),
335 "rocm"
336 );
337 assert_eq!(
338 infer_compute_capability_from_name("AMD Instinct MI250"),
339 "rocm"
340 );
341 }
342 #[cfg(not(feature = "rocm"))]
343 {
344 assert_eq!(
345 infer_compute_capability_from_name("AMD Radeon RX 7900 XTX"),
346 "gpu"
347 );
348 }
349 }
350
351 #[test]
352 fn fast_profile_cache_reuses_detected_shape() {
353 let first = detect_runtime_profile_with_mode("node-a", ProbeMode::Fast);
354 let second = detect_runtime_profile_with_mode("node-b", ProbeMode::Fast);
355
356 assert_eq!(first.logical_cores, second.logical_cores);
357 assert_eq!(first.acceleration_mode, second.acceleration_mode);
358 assert_eq!(second.node_resources.id, "node-b");
359 }
360}