Skip to main content

mistralrs_core/
diagnostics.rs

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    /// CUDA compute capability (major, minor) - None for non-CUDA devices
38    #[serde(skip_serializing_if = "Option::is_none")]
39    pub compute_capability: Option<(u32, u32)>,
40    /// Whether this GPU supports Flash Attention v2 (compute capability >= 8.0)
41    #[serde(skip_serializing_if = "Option::is_none")]
42    pub flash_attn_compatible: Option<bool>,
43    /// Whether this GPU supports Flash Attention v3 (compute capability == 9.0, Hopper only)
44    #[serde(skip_serializing_if = "Option::is_none")]
45    pub flash_attn_v3_compatible: Option<bool>,
46    /// Whether this device uses unified memory (GPU and CPU share the same physical RAM)
47    #[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    /// Whether HuggingFace is reachable
66    pub reachable: bool,
67    /// Latency in milliseconds (if reachable)
68    #[serde(skip_serializing_if = "Option::is_none")]
69    pub latency_ms: Option<u64>,
70    /// Whether the token is valid for gated models
71    #[serde(skip_serializing_if = "Option::is_none")]
72    pub token_valid_for_gated: Option<bool>,
73    /// Error message if not reachable
74    #[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    // CPU device
129    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                    // Get compute capability
155                    let compute_cap = get_cuda_compute_capability(ord);
156                    let flash_attn_v2_ok = compute_cap.map(|(major, _minor)| {
157                        // Flash Attention v2 requires compute capability >= 8.0 (Ampere+)
158                        major >= 8
159                    });
160                    let flash_attn_v3_ok = compute_cap.map(|(major, minor)| {
161                        // Flash Attention v3 requires compute capability == 9.0 (Hopper only)
162                        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), // Metal always supports flash attention
201                    flash_attn_v3_compatible: None,    // Flash Attn v3 is CUDA Hopper only
202                    unified_memory: Some(true),        // Apple Silicon always uses unified memory
203                });
204            }
205        }
206    }
207
208    devices
209}
210
211/// Get CUDA compute capability for a device ordinal
212#[cfg(feature = "cuda")]
213fn get_cuda_compute_capability(ordinal: usize) -> Option<(u32, u32)> {
214    // Use nvidia-smi to query compute capability
215    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    // Parse "8.9" format
232    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
248/// Detect CPU extensions (AVX, AVX2, AVX-512, FMA)
249fn 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/// Check HuggingFace connectivity and token validity by accessing a gated model
300#[allow(clippy::cast_possible_truncation)]
301pub fn check_hf_gated_access() -> HfConnectivityInfo {
302    let start = Instant::now();
303
304    // Try to access a gated model (google/gemma-3-4b-it)
305    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            // Check if it's an auth error vs network error
322            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                // Network works, but token is invalid/missing
331                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                // Network/other error
339                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    // CPU extensions check (ARM-aware)
372    {
373        let is_arm = cfg!(any(target_arch = "aarch64", target_arch = "arm"));
374
375        if is_arm {
376            // ARM CPUs use NEON, not AVX - no warning needed
377            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            // x86/x86_64 - check for AVX extensions
385            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    // Binary vs hardware mismatch check
424    {
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    // Unified memory detection
457    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    // CUDA compute capability + Flash Attention v2/v3 check
477    #[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                // Build status strings with emojis
485                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                // Warn if hardware supports flash attn v2 but binary doesn't have it
503                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                // Warn if hardware supports flash attn v3 but binary doesn't have it
518                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    // HuggingFace connectivity + gated model access check
564    {
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}