1use std::path::{Path, PathBuf};
2use std::time::Instant;
3
4use hf_hub::{api::sync::ApiBuilder, Cache};
5use serde::{Deserialize, Serialize};
6use sysinfo::{Disks, System};
7
8#[cfg(any(feature = "cuda", feature = "metal"))]
9use crate::MemoryUsage;
10#[cfg(any(feature = "cuda", feature = "metal"))]
11use candle_core::Device;
12
13#[derive(Debug, Clone, Serialize, Deserialize)]
14pub struct CpuInfo {
15 pub brand: Option<String>,
16 pub logical_cores: usize,
17 pub physical_cores: Option<usize>,
18 pub avx: bool,
19 pub avx2: bool,
20 pub avx512: bool,
21 pub fma: bool,
22}
23
24#[derive(Debug, Clone, Serialize, Deserialize)]
25pub struct MemoryInfo {
26 pub total_bytes: u64,
27 pub available_bytes: u64,
28}
29
30#[derive(Debug, Clone, Serialize, Deserialize)]
31pub struct DeviceInfo {
32 pub kind: String,
33 pub ordinal: Option<usize>,
34 pub name: Option<String>,
35 pub total_memory_bytes: Option<u64>,
36 pub available_memory_bytes: Option<u64>,
37 #[serde(skip_serializing_if = "Option::is_none")]
39 pub compute_capability: Option<(u32, u32)>,
40 #[serde(skip_serializing_if = "Option::is_none")]
42 pub flash_attn_compatible: Option<bool>,
43 #[serde(skip_serializing_if = "Option::is_none")]
45 pub flash_attn_v3_compatible: Option<bool>,
46 #[serde(skip_serializing_if = "Option::is_none")]
48 pub unified_memory: Option<bool>,
49}
50
51#[derive(Debug, Clone, Serialize, Deserialize)]
52pub struct BuildInfo {
53 pub cuda: bool,
54 pub metal: bool,
55 pub cudnn: bool,
56 pub flash_attn: bool,
57 pub flash_attn_v3: bool,
58 pub accelerate: bool,
59 pub mkl: bool,
60 pub git_revision: String,
61}
62
63#[derive(Debug, Clone, Serialize, Deserialize)]
64pub struct HfConnectivityInfo {
65 pub reachable: bool,
67 #[serde(skip_serializing_if = "Option::is_none")]
69 pub latency_ms: Option<u64>,
70 #[serde(skip_serializing_if = "Option::is_none")]
72 pub token_valid_for_gated: Option<bool>,
73 #[serde(skip_serializing_if = "Option::is_none")]
75 pub error: Option<String>,
76}
77
78#[derive(Debug, Clone, Serialize, Deserialize)]
79pub struct SystemInfo {
80 pub os: Option<String>,
81 pub kernel: Option<String>,
82 pub cpu: CpuInfo,
83 pub memory: MemoryInfo,
84 pub devices: Vec<DeviceInfo>,
85 pub build: BuildInfo,
86 pub hf_cache_path: Option<String>,
87}
88
89#[derive(Debug, Clone, Serialize, Deserialize)]
90#[serde(rename_all = "lowercase")]
91pub enum DoctorStatus {
92 Ok,
93 Warn,
94 Error,
95}
96
97#[derive(Debug, Clone, Serialize, Deserialize)]
98pub struct DoctorCheck {
99 pub name: String,
100 pub status: DoctorStatus,
101 pub message: String,
102 #[serde(skip_serializing_if = "Option::is_none")]
103 pub suggestion: Option<String>,
104}
105
106#[derive(Debug, Clone, Serialize, Deserialize)]
107pub struct DoctorReport {
108 pub system: SystemInfo,
109 pub checks: Vec<DoctorCheck>,
110}
111
112fn build_info() -> BuildInfo {
113 BuildInfo {
114 cuda: cfg!(feature = "cuda"),
115 metal: cfg!(feature = "metal"),
116 cudnn: cfg!(feature = "cudnn"),
117 flash_attn: cfg!(feature = "flash-attn"),
118 flash_attn_v3: cfg!(feature = "flash-attn-v3"),
119 accelerate: cfg!(feature = "accelerate"),
120 mkl: cfg!(feature = "mkl"),
121 git_revision: crate::MISTRALRS_GIT_REVISION.to_string(),
122 }
123}
124
125fn collect_devices(sys: &System) -> Vec<DeviceInfo> {
126 let mut devices = Vec::new();
127
128 let cpu_brand = sys.cpus().first().map(|c| c.brand().to_string());
130 devices.push(DeviceInfo {
131 kind: "cpu".to_string(),
132 ordinal: None,
133 name: cpu_brand,
134 total_memory_bytes: Some(sys.total_memory()),
135 available_memory_bytes: Some(sys.available_memory()),
136 compute_capability: None,
137 flash_attn_compatible: None,
138 flash_attn_v3_compatible: None,
139 unified_memory: None,
140 });
141
142 #[cfg(feature = "cuda")]
143 {
144 let mut ord = 0;
145 loop {
146 match Device::new_cuda(ord) {
147 Ok(dev) => {
148 let total = MemoryUsage.get_total_memory(&dev).ok().map(|v| v as u64);
149 let avail = MemoryUsage
150 .get_memory_available(&dev)
151 .ok()
152 .map(|v| v as u64);
153
154 let compute_cap = get_cuda_compute_capability(ord);
156 let flash_attn_v2_ok = compute_cap.map(|(major, _minor)| {
157 major >= 8
159 });
160 let flash_attn_v3_ok = compute_cap.map(|(major, minor)| {
161 major == 9 && minor == 0
163 });
164
165 devices.push(DeviceInfo {
166 kind: "cuda".to_string(),
167 ordinal: Some(ord),
168 name: None,
169 total_memory_bytes: total,
170 available_memory_bytes: avail,
171 compute_capability: compute_cap,
172 flash_attn_compatible: flash_attn_v2_ok,
173 flash_attn_v3_compatible: flash_attn_v3_ok,
174 unified_memory: Some(crate::utils::normal::is_integrated_gpu(&dev)),
175 });
176 ord += 1;
177 }
178 Err(_) => break,
179 }
180 }
181 }
182
183 #[cfg(feature = "metal")]
184 {
185 let total = candle_metal_kernels::metal::Device::all().len();
186 for ord in 0..total {
187 if let Ok(dev) = Device::new_metal(ord) {
188 let total = MemoryUsage.get_total_memory(&dev).ok().map(|v| v as u64);
189 let avail = MemoryUsage
190 .get_memory_available(&dev)
191 .ok()
192 .map(|v| v as u64);
193 devices.push(DeviceInfo {
194 kind: "metal".to_string(),
195 ordinal: Some(ord),
196 name: None,
197 total_memory_bytes: total,
198 available_memory_bytes: avail,
199 compute_capability: None,
200 flash_attn_compatible: Some(true), flash_attn_v3_compatible: None, unified_memory: Some(true), });
204 }
205 }
206 }
207
208 devices
209}
210
211#[cfg(feature = "cuda")]
213fn get_cuda_compute_capability(ordinal: usize) -> Option<(u32, u32)> {
214 let output = std::process::Command::new("nvidia-smi")
216 .args([
217 "--query-gpu=compute_cap",
218 "--format=csv,noheader",
219 &format!("-i={ordinal}"),
220 ])
221 .output()
222 .ok()?;
223
224 if !output.status.success() {
225 return None;
226 }
227
228 let stdout = String::from_utf8(output.stdout).ok()?;
229 let cap = stdout.trim();
230
231 let parts: Vec<&str> = cap.split('.').collect();
233 if parts.len() == 2 {
234 let major = parts[0].parse().ok()?;
235 let minor = parts[1].parse().ok()?;
236 Some((major, minor))
237 } else {
238 None
239 }
240}
241
242#[cfg(not(feature = "cuda"))]
243#[allow(dead_code)]
244fn get_cuda_compute_capability(_ordinal: usize) -> Option<(u32, u32)> {
245 None
246}
247
248fn detect_cpu_extensions() -> (bool, bool, bool, bool) {
250 #[cfg(any(target_arch = "x86", target_arch = "x86_64"))]
251 {
252 let avx = std::arch::is_x86_feature_detected!("avx");
253 let avx2 = std::arch::is_x86_feature_detected!("avx2");
254 let avx512 = std::arch::is_x86_feature_detected!("avx512f");
255 let fma = std::arch::is_x86_feature_detected!("fma");
256 (avx, avx2, avx512, fma)
257 }
258 #[cfg(not(any(target_arch = "x86", target_arch = "x86_64")))]
259 {
260 (false, false, false, false)
261 }
262}
263
264pub fn collect_system_info() -> SystemInfo {
265 let mut sys = System::new_all();
266 sys.refresh_all();
267
268 let (avx, avx2, avx512, fma) = detect_cpu_extensions();
269
270 let cpu = CpuInfo {
271 brand: sys.cpus().first().map(|c| c.brand().to_string()),
272 logical_cores: sys.cpus().len(),
273 physical_cores: System::physical_core_count(),
274 avx,
275 avx2,
276 avx512,
277 fma,
278 };
279
280 let memory = MemoryInfo {
281 total_bytes: sys.total_memory(),
282 available_bytes: sys.available_memory(),
283 };
284
285 let hf_cache = Cache::from_env();
286 let hf_cache_path = hf_cache.path().to_string_lossy().to_string();
287
288 SystemInfo {
289 os: System::name(),
290 kernel: System::kernel_version(),
291 cpu,
292 memory,
293 devices: collect_devices(&sys),
294 build: build_info(),
295 hf_cache_path: Some(hf_cache_path),
296 }
297}
298
299#[allow(clippy::cast_possible_truncation)]
301pub fn check_hf_gated_access() -> HfConnectivityInfo {
302 let start = Instant::now();
303
304 let api_result = ApiBuilder::from_env()
306 .with_progress(false)
307 .build()
308 .and_then(|api| api.model("google/gemma-3-4b-it".to_string()).info());
309
310 let latency_ms = start.elapsed().as_millis() as u64;
311
312 match api_result {
313 Ok(_) => HfConnectivityInfo {
314 reachable: true,
315 latency_ms: Some(latency_ms),
316 token_valid_for_gated: Some(true),
317 error: None,
318 },
319 Err(e) => {
320 let error_str = e.to_string();
321 let is_auth_error = error_str.contains("401")
323 || error_str.contains("403")
324 || error_str.contains("unauthorized")
325 || error_str.contains("Unauthorized")
326 || error_str.contains("Access denied")
327 || error_str.contains("gated");
328
329 if is_auth_error {
330 HfConnectivityInfo {
332 reachable: true,
333 latency_ms: Some(latency_ms),
334 token_valid_for_gated: Some(false),
335 error: Some("Token invalid or missing for gated models".to_string()),
336 }
337 } else {
338 HfConnectivityInfo {
340 reachable: false,
341 latency_ms: None,
342 token_valid_for_gated: None,
343 error: Some(error_str),
344 }
345 }
346 }
347 }
348}
349
350fn disk_usage_for(path: &Path) -> Option<(u64, u64)> {
351 let disks = Disks::new_with_refreshed_list();
352 let mut best: Option<(usize, u64, u64)> = None;
353 for disk in disks.list() {
354 let mount = disk.mount_point();
355 if path.starts_with(mount) {
356 let len = mount.as_os_str().len();
357 let avail = disk.available_space();
358 let total = disk.total_space();
359 if best.map(|b| len > b.0).unwrap_or(true) {
360 best = Some((len, avail, total));
361 }
362 }
363 }
364 best.map(|(_, avail, total)| (avail, total))
365}
366
367pub fn run_doctor() -> DoctorReport {
368 let system = collect_system_info();
369 let mut checks = Vec::new();
370
371 {
373 let is_arm = cfg!(any(target_arch = "aarch64", target_arch = "arm"));
374
375 if is_arm {
376 checks.push(DoctorCheck {
378 name: "cpu_extensions".to_string(),
379 status: DoctorStatus::Ok,
380 message: "CPU: ARM architecture (uses NEON)".to_string(),
381 suggestion: None,
382 });
383 } else {
384 let mut extensions = Vec::new();
386 if system.cpu.avx {
387 extensions.push("AVX");
388 }
389 if system.cpu.avx2 {
390 extensions.push("AVX2");
391 }
392 if system.cpu.fma {
393 extensions.push("FMA");
394 }
395 if system.cpu.avx512 {
396 extensions.push("AVX-512");
397 }
398
399 let has_avx2 = system.cpu.avx2;
400 let ext_str = if extensions.is_empty() {
401 "none detected".to_string()
402 } else {
403 extensions.join(", ")
404 };
405
406 checks.push(DoctorCheck {
407 name: "cpu_extensions".to_string(),
408 status: if has_avx2 {
409 DoctorStatus::Ok
410 } else {
411 DoctorStatus::Warn
412 },
413 message: format!("CPU extensions: {ext_str}"),
414 suggestion: if !has_avx2 {
415 Some("AVX2 is recommended for optimal GGML performance on x86.".to_string())
416 } else {
417 None
418 },
419 });
420 }
421 }
422
423 {
425 let has_cuda_device = system.devices.iter().any(|d| d.kind == "cuda");
426 let has_metal_device = system.devices.iter().any(|d| d.kind == "metal");
427
428 if has_cuda_device && !system.build.cuda {
429 checks.push(DoctorCheck {
430 name: "binary_hardware_match".to_string(),
431 status: DoctorStatus::Error,
432 message: "NVIDIA GPU detected but binary compiled without CUDA support."
433 .to_string(),
434 suggestion: Some("Reinstall with CUDA: cargo install --features cuda".to_string()),
435 });
436 } else if has_metal_device && !system.build.metal {
437 checks.push(DoctorCheck {
438 name: "binary_hardware_match".to_string(),
439 status: DoctorStatus::Error,
440 message: "Apple GPU detected but binary compiled without Metal support."
441 .to_string(),
442 suggestion: Some(
443 "Reinstall with Metal: cargo install --features metal".to_string(),
444 ),
445 });
446 } else {
447 checks.push(DoctorCheck {
448 name: "binary_hardware_match".to_string(),
449 status: DoctorStatus::Ok,
450 message: "Binary features match detected hardware.".to_string(),
451 suggestion: None,
452 });
453 }
454 }
455
456 for dev in system
458 .devices
459 .iter()
460 .filter(|d| d.unified_memory == Some(true))
461 {
462 let kind = &dev.kind;
463 let ord = dev.ordinal.map(|o| format!(" {o}")).unwrap_or_default();
464 checks.push(DoctorCheck {
465 name: format!("{}_{}_unified_memory", kind, dev.ordinal.unwrap_or(0)),
466 status: DoctorStatus::Ok,
467 message: format!(
468 "{}{}: unified memory detected. GPU and CPU share the same physical RAM.",
469 kind.to_uppercase(),
470 ord,
471 ),
472 suggestion: None,
473 });
474 }
475
476 #[cfg(feature = "cuda")]
478 {
479 for dev in system.devices.iter().filter(|d| d.kind == "cuda") {
480 if let (Some(ord), Some((major, minor))) = (dev.ordinal, dev.compute_capability) {
481 let fa_v2_ok = dev.flash_attn_compatible.unwrap_or(false);
482 let fa_v3_ok = dev.flash_attn_v3_compatible.unwrap_or(false);
483
484 let fa_v2_str = if fa_v2_ok { "✅" } else { "❌" };
486 let fa_v3_str = if fa_v3_ok {
487 "✅"
488 } else {
489 "❌ (requires Hopper/Compute 9.0)"
490 };
491
492 checks.push(DoctorCheck {
493 name: format!("cuda_{}_compute", ord),
494 status: DoctorStatus::Ok,
495 message: format!(
496 "GPU {}: compute {}.{} - Flash Attn v2 {}, v3 {}",
497 ord, major, minor, fa_v2_str, fa_v3_str
498 ),
499 suggestion: None,
500 });
501
502 if fa_v2_ok && !system.build.flash_attn {
504 checks.push(DoctorCheck {
505 name: format!("cuda_{}_flash_attn_v2_missing", ord),
506 status: DoctorStatus::Warn,
507 message: format!(
508 "GPU {} supports Flash Attention v2 but binary compiled without it.",
509 ord
510 ),
511 suggestion: Some(
512 "Reinstall with: cargo install --features flash-attn".to_string(),
513 ),
514 });
515 }
516
517 if fa_v3_ok && !system.build.flash_attn_v3 {
519 checks.push(DoctorCheck {
520 name: format!("cuda_{}_flash_attn_v3_missing", ord),
521 status: DoctorStatus::Warn,
522 message: format!(
523 "GPU {} supports Flash Attention v3 but binary compiled without it.",
524 ord
525 ),
526 suggestion: Some(
527 "Reinstall with: cargo install --features flash-attn-v3".to_string(),
528 ),
529 });
530 }
531 }
532 }
533 }
534
535 let hf_cache_path = system
536 .hf_cache_path
537 .as_ref()
538 .map(PathBuf::from)
539 .unwrap_or_else(|| Cache::from_env().path().clone());
540
541 if std::fs::create_dir_all(&hf_cache_path).is_err() {
542 checks.push(DoctorCheck {
543 name: "hf_cache_writable".to_string(),
544 status: DoctorStatus::Error,
545 message: format!(
546 "Cannot create or access Hugging Face cache dir at {}",
547 hf_cache_path.display()
548 ),
549 suggestion: Some("Set HF_HOME or fix permissions.".to_string()),
550 });
551 } else {
552 checks.push(DoctorCheck {
553 name: "hf_cache_writable".to_string(),
554 status: DoctorStatus::Ok,
555 message: format!(
556 "Hugging Face cache dir is writable: {}",
557 hf_cache_path.display()
558 ),
559 suggestion: None,
560 });
561 }
562
563 {
565 let hf_info = check_hf_gated_access();
566 if hf_info.reachable {
567 if hf_info.token_valid_for_gated == Some(true) {
568 checks.push(DoctorCheck {
569 name: "hf_connectivity".to_string(),
570 status: DoctorStatus::Ok,
571 message: format!(
572 "Hugging Face: connected ({}ms), token valid for allowed gated models.",
573 hf_info.latency_ms.unwrap_or(0)
574 ),
575 suggestion: None,
576 });
577 } else {
578 checks.push(DoctorCheck {
579 name: "hf_connectivity".to_string(),
580 status: DoctorStatus::Warn,
581 message: format!(
582 "Hugging Face: connected ({}ms), but token invalid/missing.",
583 hf_info.latency_ms.unwrap_or(0)
584 ),
585 suggestion: Some(
586 "Run `huggingface-cli login` or set HF_TOKEN to access gated models."
587 .to_string(),
588 ),
589 });
590 }
591 } else {
592 checks.push(DoctorCheck {
593 name: "hf_connectivity".to_string(),
594 status: DoctorStatus::Error,
595 message: format!(
596 "Hugging Face: unreachable - {}",
597 hf_info.error.unwrap_or_else(|| "unknown error".to_string())
598 ),
599 suggestion: Some(
600 "Check your internet connection and firewall settings.".to_string(),
601 ),
602 });
603 }
604 }
605
606 if let Some((avail, total)) = disk_usage_for(&hf_cache_path) {
607 let min_free = 10_u64 * 1024 * 1024 * 1024;
608 let status = if avail < min_free {
609 DoctorStatus::Warn
610 } else {
611 DoctorStatus::Ok
612 };
613 checks.push(DoctorCheck {
614 name: "disk_space".to_string(),
615 status,
616 #[allow(clippy::cast_precision_loss)]
617 message: format!(
618 "Disk free: {:.1} GB / {:.1} GB on the volume containing the HF cache at {}.",
619 avail as f64 / 1e9,
620 total as f64 / 1e9,
621 hf_cache_path.display()
622 ),
623 suggestion: if avail < min_free {
624 Some("Free up disk space or move HF cache.".to_string())
625 } else {
626 None
627 },
628 });
629 }
630
631 let has_cuda = system.devices.iter().any(|d| d.kind == "cuda");
632
633 if system.build.cuda && !has_cuda {
634 checks.push(DoctorCheck {
635 name: "cuda_devices".to_string(),
636 status: DoctorStatus::Warn,
637 message: "CUDA support is enabled but no CUDA devices were found.".to_string(),
638 suggestion: Some("Check NVIDIA driver installation.".to_string()),
639 });
640 }
641
642 DoctorReport { system, checks }
643}