Skip to main content

rusty_stack/
component_status.rs

1use crate::benchmark_logs;
2use crate::state::{Category, Component};
3use std::collections::HashSet;
4use std::env;
5use std::fs;
6use std::path::{Path, PathBuf};
7use std::process::{Command, Stdio};
8
9#[derive(Debug, Clone)]
10pub struct VerificationCommand {
11    pub label: String,
12    pub target_id: String,
13    pub program: String,
14    pub args: Vec<String>,
15    pub modules: Vec<String>,
16}
17
18pub fn python_interpreters() -> Vec<String> {
19    let mut candidates = Vec::new();
20    if let Ok(value) = env::var("MLSTACK_PYTHON_BIN") {
21        push_python_candidate(&mut candidates, value);
22    }
23    if let Ok(value) = env::var("AITER_VENV_PYTHON") {
24        push_python_candidate(&mut candidates, value);
25    }
26    if let Ok(value) = env::var("UV_PYTHON") {
27        push_python_candidate(&mut candidates, value);
28    }
29    for key in [
30        "WANDB_VENV_PYTHON",
31        "MEGATRON_VENV_PYTHON",
32        "MPI4PY_VENV_PYTHON",
33        "FLASH_ATTENTION_VENV_PYTHON",
34        "PYTORCH_VENV_PYTHON",
35        "DEEPSPEED_VENV_PYTHON",
36        "VLLM_VENV_PYTHON",
37    ] {
38        if let Ok(value) = env::var(key) {
39            push_python_candidate(&mut candidates, value);
40        }
41    }
42    let home = resolve_component_user_home();
43    let venv = Path::new(&home)
44        .join("rocm_venv")
45        .join("bin")
46        .join("python");
47    if venv.exists() {
48        candidates.push(venv.to_string_lossy().to_string());
49    }
50
51    // Anaconda/Miniconda candidates
52    let conda_path = Path::new(&home)
53        .join("anaconda3")
54        .join("bin")
55        .join("python");
56    if conda_path.exists() {
57        candidates.push(conda_path.to_string_lossy().to_string());
58    }
59    let miniconda_path = Path::new(&home)
60        .join("miniconda3")
61        .join("bin")
62        .join("python");
63    if miniconda_path.exists() {
64        candidates.push(miniconda_path.to_string_lossy().to_string());
65    }
66
67    if Path::new("/usr/local/bin/python3").exists() {
68        candidates.push("/usr/local/bin/python3".to_string());
69    }
70    if command_exists("python3.13") {
71        candidates.push("python3.13".to_string());
72    }
73    if command_exists("python3.12") {
74        candidates.push("python3.12".to_string());
75    }
76    if command_exists("python3") {
77        candidates.push("python3".to_string());
78    }
79    if command_exists("python") {
80        candidates.push("python".to_string());
81    }
82
83    let mut roots = vec![PathBuf::from(&home)];
84    if let Ok(cwd) = env::current_dir() {
85        roots.push(cwd.clone());
86        roots.push(cwd.join("rusty-stack"));
87    }
88    for root in roots {
89        push_component_venv_candidates(&mut candidates, &root);
90    }
91
92    dedupe_python_candidates_in_order(candidates)
93}
94
95fn dedupe_python_candidates_in_order(candidates: Vec<String>) -> Vec<String> {
96    let mut seen = HashSet::new();
97    let mut ordered = Vec::new();
98    for candidate in candidates {
99        if seen.insert(candidate.clone()) {
100            ordered.push(candidate);
101        }
102    }
103    ordered
104}
105
106fn push_component_venv_candidates(candidates: &mut Vec<String>, root: &Path) {
107    for rel in [
108        "wandb_venv/bin/python",
109        "megatron_rocm_venv/bin/python",
110        "Megatron-LM/megatron_rocm_venv/bin/python",
111        "mpi4py_venv/bin/python",
112        "flash_attention_venv/bin/python",
113        "vllm_venv/bin/python",
114        ".mlstack/venvs/aiter/bin/python",
115        ".mlstack/venvs/vllm/bin/python",
116        "pytorch_rocm_venv/bin/python",
117    ] {
118        let candidate = root.join(rel);
119        if candidate.exists() {
120            candidates.push(candidate.to_string_lossy().to_string());
121        }
122    }
123}
124
125fn push_python_candidate(candidates: &mut Vec<String>, value: String) {
126    let trimmed = value.trim();
127    if trimmed.is_empty() {
128        return;
129    }
130    let candidate = trimmed.to_string();
131    if Path::new(&candidate).exists() || command_exists(&candidate) {
132        candidates.push(candidate);
133    }
134}
135
136pub fn is_component_installed(component: &Component, python_candidates: &[String]) -> bool {
137    if component.category == Category::Verification {
138        return false;
139    }
140    is_component_installed_by_id(&component.id, python_candidates)
141}
142
143pub fn is_component_installed_by_id(component_id: &str, python_candidates: &[String]) -> bool {
144    let home = resolve_component_user_home();
145    let has_benchmark_logs = has_benchmark_log_dirs();
146    match component_id {
147        // ROCm requires BOTH version file AND functional rocminfo (not just binary existence)
148        // This prevents false positives from partial downloads
149        "rocm" => {
150            // Check for version file in multiple locations
151            let version_exists = path_exists("/opt/rocm/.info/version")
152                || path_exists("/opt/rocm/version")
153                || path_exists("/usr/lib/rocm/.info/version");
154
155            // Check if rocminfo actually works (not just exists)
156            let rocminfo_path = crate::installers::common::utils::resolve_rocminfo_path();
157            let rocminfo_functional = if let Ok(output) = std::process::Command::new(&rocminfo_path)
158                .arg("--version")
159                .output()
160            {
161                output.status.success()
162            } else {
163                false
164            };
165
166            // Require both: version file AND functional rocminfo
167            // This ensures we don't report partial installs as complete
168            version_exists && rocminfo_functional
169        }
170        "pytorch" => python_any(python_candidates, &["torch"]),
171        "triton" => python_any(python_candidates, &["triton"]),
172        "mpi4py" => python_any(python_candidates, &["mpi4py"]),
173        "deepspeed" => python_any(python_candidates, &["deepspeed"]),
174        "ml-stack-core" => {
175            python_any(python_candidates, &["stans_ml_stack"]) || repo_has_ml_stack_core()
176        }
177        "flash-attn" => {
178            python_any(python_candidates, &["flash_attn", "flash_attn_2"])
179                || path_exists(home_path(&home, &["ml_stack", "flash_attn_amd"]))
180                || path_exists(home_path(&home, &["ml_stack", "flash_attn_amd_direct"]))
181        }
182        "megatron" => {
183            path_exists(home_path(&home, &["megatron", "Megatron-LM"]))
184                || python_any(python_candidates, &["megatron"])
185        }
186        "vllm" => {
187            python_candidates
188                .iter()
189                .any(|python| python_exec(python, vllm_runtime_check_snippet()))
190                || path_exists(home_path(&home, &["vllm_build"]))
191                || path_exists(home_path(&home, &["vllm_py313"]))
192        }
193        "aiter" => python_any(python_candidates, &["aiter"]),
194        "vllm-studio" => {
195            command_exists("vllm-studio") || path_exists(home_path(&home, &["vllm-studio"]))
196        }
197        "onnx" => {
198            python_any(python_candidates, &["onnxruntime"])
199                || path_exists(home_path(&home, &["onnxruntime_build"]))
200        }
201        "bitsandbytes" => {
202            python_any(python_candidates, &["bitsandbytes"])
203                || path_exists(home_path(&home, &["ml_stack", "bitsandbytes"]))
204        }
205        // rocm-smi - check multiple locations and verify it works
206        "rocm-smi" => {
207            let rocm_smi_paths = [
208                "/opt/rocm/bin/rocm-smi",
209                "/usr/bin/rocm-smi",
210                "/usr/local/bin/rocm-smi",
211            ];
212            for path in &rocm_smi_paths {
213                if path_exists(path) {
214                    // Verify it actually works
215                    if let Ok(output) = std::process::Command::new(path).arg("--showuse").output() {
216                        if output.status.success() {
217                            return true;
218                        }
219                    }
220                }
221            }
222            // Also check if rocm-smi is in PATH and works
223            if command_exists("rocm-smi") {
224                if let Ok(output) = std::process::Command::new("rocm-smi")
225                    .arg("--showuse")
226                    .output()
227                {
228                    return output.status.success();
229                }
230            }
231            false
232        }
233        "migraphx" => {
234            python_any(python_candidates, &["migraphx"])
235                || path_exists(home_path(&home, &["migraphx_build"]))
236        }
237        "pytorch-profiler" => python_any(python_candidates, &["torch"]),
238        "wandb" => python_any(python_candidates, &["wandb"]),
239        "fastvideo" => python_any(python_candidates, &["fastvideo"]),
240        "llama-cpp" => {
241            // Use the detection contract from llama_cpp.rs: llama-cli --help must succeed
242            crate::installers::components::llama_cpp::is_llama_cli_functional(&home)
243        }
244        "permanent-env" => env_file_has_permanent(&home_path(&home, &[".mlstack_env"])),
245        "basic-env" => path_exists(home_path(&home, &[".mlstack_env"])),
246        "enhanced-env" => env_file_has_enhanced(&home_path(&home, &[".mlstack_env"])),
247        "vllm-performance" => has_benchmark_logs,
248        "deepspeed-performance" => has_benchmark_logs,
249        "megatron-performance" => has_benchmark_logs,
250        "all-benchmarks" => has_benchmark_logs,
251        "comfyui" => {
252            // Check for ComfyUI installation
253            // Default location: $HOME/ComfyUI
254            // User location: /mnt/e0f7c1a8-b834-4827-b579-0251b006bc1f/ComfyUI/
255            // Check for .git folder and main.py
256            let default_path = home_path(&home, &["ComfyUI"]);
257            path_exists(default_path.join(".git")) || path_exists(default_path.join("main.py"))
258        }
259        "textgen" => {
260            // Check for text-generation-webui installation
261            // Default location: $HOME/text-generation-webui
262            let default_path = home_path(&home, &["text-generation-webui"]);
263            path_exists(default_path.join(".git")) || path_exists(default_path.join("server.py"))
264        }
265        _ => false,
266    }
267}
268
269pub fn verification_commands(
270    component_id: &str,
271    python_candidates: &[String],
272) -> Vec<VerificationCommand> {
273    match component_id {
274        "verify-basic" => basic_verification_commands(python_candidates),
275        "verify-enhanced" => enhanced_verification_commands(python_candidates),
276        "verify-build" => build_verification_commands(python_candidates),
277        _ => Vec::new(),
278    }
279}
280
281pub fn component_verification_commands(
282    component_id: &str,
283    python_candidates: &[String],
284) -> Vec<VerificationCommand> {
285    match component_id {
286        "rocm" => vec![shell_command("ROCm info", "rocm", "rocminfo", &[])],
287        "pytorch" => vec![python_command(
288            "PyTorch",
289            "pytorch",
290            &["torch"],
291            python_candidates,
292            pytorch_diagnostic_snippet(),
293        )],
294        "triton" => vec![python_command(
295            "Triton",
296            "triton",
297            &["triton"],
298            python_candidates,
299            "import triton,sys; import triton.language as tl; import importlib; ok=False\nfor _m in ('triton.runtime.driver','triton.runtime.backends.amd','triton.backends.amd'):\n    try:\n        importlib.import_module(_m); ok=True; break\n    except Exception:\n        pass\nprint(f'Triton {triton.__version__} backend_loaded={ok}'); sys.exit(0 if ok else 1)",
300        )],
301        "mpi4py" => vec![python_command(
302            "MPI4Py",
303            "mpi4py",
304            &["mpi4py"],
305            python_candidates,
306            "from mpi4py import MPI; import sys\ntry:\n    rank=MPI.COMM_WORLD.Get_rank(); ok=rank>=0\nexcept Exception:\n    rank=-1; ok=False\nprint(f'MPI4Py rank={rank} ok={ok}'); sys.exit(0 if ok else 1)",
307        )],
308        "deepspeed" => vec![python_command(
309            "DeepSpeed",
310            "deepspeed",
311            &["deepspeed"],
312            python_candidates,
313            "import deepspeed,sys; ok=False\ntry:\n    from deepspeed.ops.adam import DeepSpeedCPUAdam; ok=True\nexcept Exception:\n    pass\nprint(f'DeepSpeed {deepspeed.__version__} ops_loaded={ok}'); sys.exit(0 if ok else 1)",
314        )],
315        "ml-stack-core" => vec![python_command(
316            "ML Stack Core",
317            "ml-stack-core",
318            &["stans_ml_stack"],
319            python_candidates,
320            "import stans_ml_stack,sys; import importlib; ok=True\nfor _s in ('core','utils','cli','installers'):\n    try:\n        importlib.import_module(f'stans_ml_stack.{_s}')\n    except Exception:\n        ok=False; break\nver=getattr(stans_ml_stack,'__version__','ok')\nprint(f'ML Stack Core {ver} submodule_loaded={ok}'); sys.exit(0 if ok else 1)",
321        )],
322        "flash-attn" => vec![python_command(
323            "Flash Attention",
324            "flash-attn",
325            &["flash_attn", "flash_attention_amd", "flash_attn_2"],
326            python_candidates,
327            "import importlib, sys;\nmodules=[('flash_attn','flash_attn.flash_attn_func'),('flash_attention_amd',None),('flash_attn_2',None)];\nfor mod_name, func_path in modules:\n    try:\n        m = importlib.import_module(mod_name);\n        ver = getattr(m, '__version__', 'unknown');\n        if func_path:\n            parts = func_path.split('.');\n            obj = m;\n            [obj := getattr(obj, p) for p in parts[1:]];\n            print(f'{mod_name} {ver} — {func_path} OK');\n        else:\n            print(f'{mod_name} {ver} — imported OK');\n        sys.exit(0);\n    except SystemExit:\n        raise\n    except Exception:\n        continue;\nprint('No flash attention module found'); sys.exit(1)",
328        )],
329        "megatron" => vec![python_command(
330            "Megatron-LM",
331            "megatron",
332            &["megatron"],
333            python_candidates,
334            "import megatron,sys; ok=False\ntry:\n    from megatron.core import tensor_parallel; ok=True\nexcept Exception:\n    pass\nprint(f'Megatron-LM functional={ok}'); sys.exit(0 if ok else 1)",
335        )],
336        "vllm" => vec![python_command(
337            "vLLM",
338            "vllm",
339            &["vllm"],
340            python_candidates,
341            vllm_runtime_check_snippet(),
342        )],
343        "aiter" => vec![python_command(
344            "AITER",
345            "aiter",
346            &["aiter"],
347            python_candidates,
348            "import aiter,sys; ok=False\ntry:\n    from aiter import op as _op; ok=True\nexcept Exception:\n    pass\nmods=[x for x in dir(aiter) if not x.startswith('_')]\nver=getattr(aiter,'__version__','?')\nprint(f'AITER {ver} op_loaded={ok} exports={len(mods)}'); sys.exit(0 if ok else 1)",
349        )],
350        "vllm-studio" => vec![shell_command(
351            "vLLM Studio",
352            "vllm-studio",
353            "bun",
354            &["--version"],
355        )],
356        "onnx" => vec![python_command(
357            "ONNX Runtime",
358            "onnx",
359            &["onnxruntime"],
360            python_candidates,
361            "import onnxruntime as ort; import sys; print('Version:', ort.__version__); providers=ort.get_available_providers(); print('Providers:', providers); gpu=[p for p in providers if p in ('ROCMExecutionProvider','MIGraphXExecutionProvider','CUDAExecutionProvider')]; sys.exit(0 if gpu else 1)",
362        )],
363        "bitsandbytes" => vec![python_command(
364            "BITSANDBYTES",
365            "bitsandbytes",
366            &["bitsandbytes"],
367            python_candidates,
368            "import bitsandbytes as bnb; import pathlib; import sys; print('Version:', getattr(bnb, '__version__', 'unknown')); path=pathlib.Path(bnb.__file__).parent; libs=list(path.glob('libbitsandbytes_rocm*.so')); print('ROCm Libs:', libs); sys.exit(0 if libs else 1)",
369        )],
370        "rocm-smi" => vec![shell_command(
371            "ROCm SMI",
372            "rocm-smi",
373            "rocm-smi",
374            &["--showproductname"],
375        )],
376        "migraphx" => vec![python_command(
377            "MIGraphX",
378            "migraphx",
379            &["migraphx"],
380            python_candidates,
381            "import migraphx,sys; ok=hasattr(migraphx,'parse_onnx')\nver=getattr(migraphx,'__version__','ok')\nprint(f'MIGraphX {ver} parse_onnx={ok}'); sys.exit(0 if ok else 1)",
382        )],
383        "pytorch-profiler" => vec![python_command(
384            "PyTorch Profiler",
385            "pytorch-profiler",
386            &["torch"],
387            python_candidates,
388            "from torch.profiler import profile; print(profile)",
389        )],
390        "wandb" => vec![python_command(
391            "Weights & Biases",
392            "wandb",
393            &["wandb"],
394            python_candidates,
395            "import wandb,sys; ok=False\ntry:\n    _r=wandb.init(mode='disabled'); ok=_r is not None; _r.finish()\nexcept Exception:\n    pass\nprint(f'wandb {wandb.__version__} init_run_ok={ok}'); sys.exit(0 if ok else 1)",
396        )],
397        "fastvideo" => vec![python_command(
398            "FastVideo",
399            "fastvideo",
400            &["fastvideo"],
401            python_candidates,
402            "import fastvideo,sys; import importlib; ok=False\nfor _m in ('fastvideo.inference','fastvideo.engine','fastvideo'):\n    try:\n        importlib.import_module(_m); ok=True; break\n    except Exception:\n        pass\nver=getattr(fastvideo,'__version__','?')\nprint(f'FastVideo {ver} module_loaded={ok}'); sys.exit(0 if ok else 1)",
403        )],
404        "llama-cpp" => vec![shell_command(
405            "llama.cpp",
406            "llama-cpp",
407            "llama-cli",
408            &["--help"],
409        )],
410        "basic-env" => vec![shell_command(
411            "Environment file",
412            "basic-env",
413            "bash",
414            &[
415                "-c",
416                "test -f \"$HOME/.mlstack_env\" && echo OK || exit 1",
417            ],
418        )],
419        "enhanced-env" => vec![shell_command(
420            "Enhanced env",
421            "enhanced-env",
422            "bash",
423            &[
424                "-c",
425                "grep -q \"Enhanced ML Stack Environment Setup Script\" \"$HOME/.mlstack_env\" && echo OK || exit 1",
426            ],
427        )],
428        "permanent-env" => vec![shell_command(
429            "Permanent env",
430            "permanent-env",
431            "bash",
432            &[
433                "-c",
434                "grep -qiE \"Permanent ROCm Environment|Permanent ROCm Env|MLSTACK_PYTHON_BIN\" \"$HOME/.mlstack_env\" && echo OK || exit 1",
435            ],
436        )],
437        // Benchmark components - verify by checking log files exist
438        "mlperf-inference" => vec![benchmark_log_check_command(
439            "MLPerf benchmark logs",
440            "mlperf-inference",
441            "mlperf_inference",
442            "Benchmark logs found",
443            "No benchmark logs yet",
444        )],
445        "rocm-benchmarks" => vec![benchmark_log_check_command(
446            "ROCm benchmark logs",
447            "rocm-benchmarks",
448            "rocm_benchmarks",
449            "Benchmark logs found",
450            "No benchmark logs yet",
451        )],
452        "gpu-memory-bandwidth" => vec![benchmark_log_check_command(
453            "Memory bandwidth logs",
454            "gpu-memory-bandwidth",
455            "gpu_memory_bandwidth",
456            "Benchmark logs found",
457            "No benchmark logs yet",
458        )],
459        "rocm-smi-bench" => vec![benchmark_log_check_command(
460            "ROCm SMI logs",
461            "rocm-smi-bench",
462            "rocm_smi_benchmarks",
463            "Benchmark logs found",
464            "No benchmark logs yet",
465        )],
466        "pytorch-performance" => vec![benchmark_log_check_command(
467            "PyTorch benchmark logs",
468            "pytorch-performance",
469            "pytorch_performance",
470            "Benchmark logs found",
471            "No benchmark logs yet",
472        )],
473        "vllm-performance" => vec![benchmark_log_check_command(
474            "vLLM benchmark logs",
475            "vllm-performance",
476            "vllm_benchmarks",
477            "Benchmark logs found",
478            "No benchmark logs yet",
479        )],
480        "deepspeed-performance" => vec![benchmark_log_check_command(
481            "DeepSpeed benchmark logs",
482            "deepspeed-performance",
483            "deepspeed_benchmarks",
484            "Benchmark logs found",
485            "No benchmark logs yet",
486        )],
487        "megatron-performance" => vec![benchmark_log_check_command(
488            "Megatron benchmark logs",
489            "megatron-performance",
490            "megatron_benchmarks",
491            "Benchmark logs found",
492            "No benchmark logs yet",
493        )],
494        "all-benchmarks" => vec![benchmark_log_check_command(
495            "Full suite logs",
496            "all-benchmarks",
497            "full_benchmarks",
498            "Suite logs found",
499            "No suite logs yet",
500        )],
501        "comfyui" => vec![{
502            // F2: verify against the env's interpreter, not a hardcoded python3.
503            // Resolve the same python the installers pin (MLSTACK_PYTHON_BIN /
504            // UV_PYTHON / first candidate), then import from the install dir.
505            let py = python_candidates
506                .first()
507                .cloned()
508                .or_else(|| env::var("MLSTACK_PYTHON_BIN").ok().filter(|s| !s.is_empty()))
509                .or_else(|| env::var("UV_PYTHON").ok().filter(|s| !s.is_empty()))
510                .unwrap_or_else(|| "python3".to_string());
511            shell_command(
512                "ComfyUI",
513                "comfyui",
514                "bash",
515                &[
516                    "-c",
517                    &format!(
518                        "cd \"$HOME/ComfyUI\" 2>/dev/null && {py} -c 'import folder_paths' 2>/dev/null && echo 'comfyui functional' || exit 1"
519                    ),
520                ],
521            )
522        }],
523        "textgen" => vec![{
524            // F2: verify against the env's interpreter, not a hardcoded python3.
525            let py = python_candidates
526                .first()
527                .cloned()
528                .or_else(|| env::var("MLSTACK_PYTHON_BIN").ok().filter(|s| !s.is_empty()))
529                .or_else(|| env::var("UV_PYTHON").ok().filter(|s| !s.is_empty()))
530                .unwrap_or_else(|| "python3".to_string());
531            shell_command(
532                "text-generation-webui",
533                "textgen",
534                "bash",
535                &[
536                    "-c",
537                    &format!(
538                        "cd \"$HOME/text-generation-webui\" 2>/dev/null && {py} -c 'import server' 2>/dev/null && echo 'textgen functional' || exit 1"
539                    ),
540                ],
541            )
542        }],
543        _ => Vec::new(),
544    }
545}
546
547pub fn modules_available(modules: &[String], python_candidates: &[String]) -> bool {
548    if modules.is_empty() {
549        return true;
550    }
551    let module_refs: Vec<&str> = modules.iter().map(|s| s.as_str()).collect();
552    python_any(python_candidates, &module_refs)
553}
554
555fn basic_verification_commands(python_candidates: &[String]) -> Vec<VerificationCommand> {
556    vec![
557        shell_command("ROCm info", "rocm", "rocminfo", &[]),
558        python_command(
559            "PyTorch",
560            "pytorch",
561            &["torch"],
562            python_candidates,
563            pytorch_diagnostic_snippet(),
564        ),
565        python_command(
566            "Triton",
567            "triton",
568            &["triton"],
569            python_candidates,
570            "import triton; import triton._C; print(f'Triton {triton.__version__} ok')",
571        ),
572        python_command(
573            "MPI4Py",
574            "mpi4py",
575            &["mpi4py"],
576            python_candidates,
577            "from mpi4py import MPI; import sys\ntry:\n    rank=MPI.COMM_WORLD.Get_rank(); ok=rank>=0\nexcept Exception:\n    rank=-1; ok=False\nprint(f'MPI4Py rank={rank} ok={ok}'); sys.exit(0 if ok else 1)",
578        ),
579        python_command(
580            "DeepSpeed",
581            "deepspeed",
582            &["deepspeed", "einops"],
583            python_candidates,
584            "import deepspeed; import einops; print(deepspeed.__version__)",
585        ),
586    ]
587}
588
589fn enhanced_verification_commands(python_candidates: &[String]) -> Vec<VerificationCommand> {
590    let mut steps = basic_verification_commands(python_candidates);
591    steps.extend(vec![
592        python_command(
593            "Flash Attention",
594            "flash-attn",
595            &["flash_attn", "flash_attention_amd", "flash_attn_2"],
596            python_candidates,
597            "import importlib, sys;\nmodules=[('flash_attn','flash_attn.flash_attn_func'),('flash_attention_amd',None),('flash_attn_2',None)];\nfor mod_name, func_path in modules:\n    try:\n        m = importlib.import_module(mod_name);\n        ver = getattr(m, '__version__', 'unknown');\n        if func_path:\n            parts = func_path.split('.');\n            obj = m;\n            [obj := getattr(obj, p) for p in parts[1:]];\n            print(f'{mod_name} {ver} — {func_path} OK');\n        else:\n            print(f'{mod_name} {ver} — imported OK');\n        sys.exit(0);\n    except SystemExit:\n        raise\n    except Exception:\n        continue;\nprint('No flash attention module found'); sys.exit(1)",
598        ),
599        python_command(
600            "vLLM",
601            "vllm",
602            &["vllm"],
603            python_candidates,
604            vllm_runtime_check_snippet(),
605        ),
606        python_command(
607            "AITER",
608            "aiter",
609            &["aiter"],
610            python_candidates,
611            "import aiter; print(getattr(aiter, '__version__', 'ok'))",
612        ),
613        python_command(
614            "ONNX Runtime",
615            "onnx",
616            &["onnxruntime"],
617            python_candidates,
618            "import onnxruntime as ort; import pathlib; import os; import sys; print('Version:', ort.__version__); base=pathlib.Path(ort.__file__).parent; libs=list(base.rglob('libonnxruntime_providers_rocm.so')); [os.environ.update({'ORT_ROCM_EP_PROVIDER_PATH': str(l)}) for l in libs[:1]]; providers=ort.get_available_providers(); print('Providers:', providers); sys.exit(0 if 'ROCMExecutionProvider' in providers else 1)",
619        ),
620        python_command(
621            "BITSANDBYTES",
622            "bitsandbytes",
623            &["bitsandbytes"],
624            python_candidates,
625            "import bitsandbytes as bnb; import pathlib; import sys; print('Version:', getattr(bnb, '__version__', 'unknown')); path=pathlib.Path(bnb.__file__).parent; libs=list(path.glob('libbitsandbytes_rocm*.so')); print('ROCm Libs:', libs); sys.exit(0 if libs else 1)",
626        ),
627        shell_command("ROCm SMI", "rocm-smi", "rocm-smi", &["--showproductname"]),
628        python_command(
629            "MIGraphX",
630            "migraphx",
631            &["migraphx"],
632            python_candidates,
633            "import migraphx,sys; ok=hasattr(migraphx,'parse_onnx')\nver=getattr(migraphx,'__version__','ok')\nprint(f'MIGraphX {ver} parse_onnx={ok}'); sys.exit(0 if ok else 1)",
634        ),
635        python_command(
636            "PyTorch Profiler",
637            "pytorch-profiler",
638            &["torch"],
639            python_candidates,
640            "from torch.profiler import profile; print(profile)",
641        ),
642        python_command(
643            "Weights & Biases",
644            "wandb",
645            &["wandb"],
646            python_candidates,
647            "import wandb,sys; ok=False\ntry:\n    _r=wandb.init(mode='disabled'); ok=_r is not None; _r.finish()\nexcept Exception:\n    pass\nprint(f'wandb {wandb.__version__} init_run_ok={ok}'); sys.exit(0 if ok else 1)",
648        ),
649        python_command(
650            "FastVideo",
651            "fastvideo",
652            &["fastvideo"],
653            python_candidates,
654            "import fastvideo,sys; import importlib; ok=False\nfor _m in ('fastvideo.inference','fastvideo.engine','fastvideo'):\n    try:\n        importlib.import_module(_m); ok=True; break\n    except Exception:\n        pass\nver=getattr(fastvideo,'__version__','?')\nprint(f'FastVideo {ver} module_loaded={ok}'); sys.exit(0 if ok else 1)",
655        ),
656        shell_command("vLLM Studio", "vllm-studio", "bun", &["--version"]),
657    ]);
658    steps
659}
660
661fn build_verification_commands(python_candidates: &[String]) -> Vec<VerificationCommand> {
662    let mut steps = basic_verification_commands(python_candidates);
663    steps.extend(vec![
664        python_command(
665            "ONNX Runtime",
666            "onnx",
667            &["onnxruntime"],
668            python_candidates,
669            "import onnxruntime as ort; import pathlib; import os; import sys; print('Version:', ort.__version__); base=pathlib.Path(ort.__file__).parent; libs=list(base.rglob('libonnxruntime_providers_rocm.so')); [os.environ.update({'ORT_ROCM_EP_PROVIDER_PATH': str(l)}) for l in libs[:1]]; providers=ort.get_available_providers(); print('Providers:', providers); sys.exit(0 if 'ROCMExecutionProvider' in providers else 1)",
670        ),
671        python_command(
672            "Flash Attention (build)",
673            "flash-attn",
674            &["flash_attn", "flash_attention_amd", "flash_attn_2"],
675            python_candidates,
676            "import importlib, sys;\nmodules=[('flash_attn','flash_attn.flash_attn_func'),('flash_attention_amd',None),('flash_attn_2',None)];\nfor mod_name, func_path in modules:\n    try:\n        m = importlib.import_module(mod_name);\n        ver = getattr(m, '__version__', 'unknown');\n        if func_path:\n            parts = func_path.split('.');\n            obj = m;\n            [obj := getattr(obj, p) for p in parts[1:]];\n            print(f'{mod_name} {ver} — {func_path} OK');\n        else:\n            print(f'{mod_name} {ver} — imported OK');\n        sys.exit(0);\n    except SystemExit:\n        raise\n    except Exception:\n        continue;\nprint('No flash attention module found'); sys.exit(1)",
677        ),
678        python_command(
679            "MIGraphX",
680            "migraphx",
681            &["migraphx"],
682            python_candidates,
683            "import migraphx,sys; ok=hasattr(migraphx,'parse_onnx')\nver=getattr(migraphx,'__version__','ok')\nprint(f'MIGraphX {ver} parse_onnx={ok}'); sys.exit(0 if ok else 1)",
684        ),
685    ]);
686    steps
687}
688
689fn python_command(
690    label: &str,
691    target_id: &str,
692    modules: &[&str],
693    python_candidates: &[String],
694    code: &str,
695) -> VerificationCommand {
696    let program = select_python_for_modules(modules, python_candidates)
697        .or_else(|| python_candidates.first().cloned())
698        .unwrap_or_else(|| "python3".to_string());
699
700    VerificationCommand {
701        label: label.to_string(),
702        target_id: target_id.to_string(),
703        program,
704        args: vec!["-c".to_string(), code.to_string()],
705        modules: modules.iter().map(|m| m.to_string()).collect(),
706    }
707}
708
709/// PyTorch diagnostic verification snippet.
710///
711/// This performs a multi-tier check:
712/// 1. Can we import torch? (basic install check)
713/// 2. Is torch.version.hip set? (ROCm build check)
714/// 3. Is torch.cuda.is_available()? (runtime HIP check)
715/// 4. Diagnostics for why HIP might not be available
716///
717/// Verification reflects FUNCTIONALITY, not just import: it exits non-zero if
718/// torch is not a ROCm build (`torch.version.hip` is None) OR the HIP runtime
719/// is unavailable (`torch.cuda.is_available()` is False). The diagnostic tiers
720/// print actionable guidance. Whether a pytorch-verify failure should block
721/// sibling components is an orchestrator-layer concern — NOT a reason to lie
722/// about pytorch's status. (Stage 4 tenet: verify functionality, fail loud.)
723fn pytorch_diagnostic_snippet() -> &'static str {
724    r#"import torch, sys, os, ctypes.util
725
726# Tier 1: Basic torch import (already verified by module check)
727print(f'torch version: {torch.__version__}')
728
729# Tier 2: Check if this is a ROCm build
730hip_version = getattr(torch.version, 'hip', None)
731cuda_version = getattr(torch.version, 'cuda', None)
732print(f'hip: {hip_version}')
733print(f'cuda: {cuda_version}')
734
735# Tier 3: Runtime HIP/CUDA availability
736try:
737    hip_available = torch.cuda.is_available()
738    print(f'torch.cuda.is_available(): {hip_available}')
739    if hip_available:
740        device_count = torch.cuda.device_count()
741        print(f'device_count: {device_count}')
742        if device_count > 0:
743            print(f'device_name: {torch.cuda.get_device_name(0)}')
744except Exception as e:
745    hip_available = False
746    print(f'torch.cuda.is_available() error: {e}')
747
748# Tier 4: Diagnostics if HIP not available
749if not hip_available:
750    print('--- ROCm HIP Diagnostics ---')
751
752    # Check LD_LIBRARY_PATH
753    ld_path = os.environ.get('LD_LIBRARY_PATH', '')
754    print(f'LD_LIBRARY_PATH: {ld_path if ld_path else "(not set)"}')
755
756    # Check ROCm_PATH
757    rocm_path = os.environ.get('ROCM_PATH', '')
758    print(f'ROCM_PATH: {rocm_path if rocm_path else "(not set)"}')
759
760    # Check if ROCm libraries exist
761    rocm_lib_dirs = ['/opt/rocm/lib', '/opt/rocm/hip/lib']
762    for d in rocm_lib_dirs:
763        exists = os.path.isdir(d)
764        print(f'{d}: {"exists" if exists else "not found"}')
765
766    # Check if libamdhip64.so can be found
767    hip_lib = ctypes.util.find_library('amdhip64')
768    print(f'libamdhip64: {hip_lib if hip_lib else "not found"}')
769
770    # Check Python executable path
771    print(f'python: {sys.executable}')
772
773    # Provide actionable guidance
774    if hip_version is None and cuda_version is None:
775        print('HINT: torch.version.hip and torch.version.cuda are both None.')
776        print('  This torch build may not have GPU support.')
777        print('  Reinstall with ROCm index: pip install torch --index-url https://repo.radeon.com/rocm/manylinux/rocm-rel-7.2/')
778    elif hip_version is not None and not hip_available:
779        print('HINT: torch.version.hip is set but HIP runtime not available.')
780        print('  Possible causes:')
781        print('  1. LD_LIBRARY_PATH does not include /opt/rocm/lib')
782        print('  2. Running in a container/environment without GPU access')
783        print('  3. ROCm drivers not loaded (check with: rocminfo)')
784        print('  Fix: source ~/.mlstack_env or set LD_LIBRARY_PATH=/opt/rocm/lib')
785
786# Stage 4: exit code reflects FUNCTIONAL HIP availability on this ROCm stack.
787# A CPU-only/CUDA torch build, or a ROCm build whose HIP runtime is
788# unreachable, must FAIL verification — not silently pass. The diagnostics
789# above guide the fix; downstream gating belongs in the orchestrator layer,
790# not a lying exit code.
791ok = (hip_version is not None) and hip_available
792sys.exit(0 if ok else 1)"#
793}
794
795fn vllm_runtime_check_snippet() -> &'static str {
796    r#"import importlib
797import vllm
798import cachetools
799import cbor2
800import gguf
801import pybase64
802import ijson
803import mistral_common
804import openai_harmony
805
806loaded = []
807errs = []
808for name in ("vllm._C", "vllm._rocm_C"):
809    try:
810        importlib.import_module(name)
811        loaded.append(name)
812    except Exception as exc:
813        errs.append(f"{name}: {exc}")
814
815if not loaded:
816    raise SystemExit("vLLM native extension load failure: " + " | ".join(errs))
817
818print(vllm.__version__)"#
819}
820
821fn shell_command(
822    label: &str,
823    target_id: &str,
824    program: &str,
825    args: &[&str],
826) -> VerificationCommand {
827    VerificationCommand {
828        label: label.to_string(),
829        target_id: target_id.to_string(),
830        program: program.to_string(),
831        args: args.iter().map(|s| s.to_string()).collect(),
832        modules: Vec::new(),
833    }
834}
835
836fn has_benchmark_log_dirs() -> bool {
837    benchmark_logs::benchmark_log_directories()
838        .iter()
839        .any(|dir| {
840            fs::read_dir(dir)
841                .ok()
842                .map(|entries| {
843                    entries
844                        .filter_map(|entry| entry.ok())
845                        .any(|entry| entry.path().is_file())
846                })
847                .unwrap_or(false)
848        })
849}
850
851fn benchmark_log_check_command(
852    label: &str,
853    target_id: &str,
854    pattern: &str,
855    success_message: &str,
856    failure_message: &str,
857) -> VerificationCommand {
858    let script = benchmark_log_check_script(pattern, success_message, failure_message);
859    shell_command(label, target_id, "bash", &["-c", &script])
860}
861
862fn benchmark_log_check_script(
863    pattern: &str,
864    success_message: &str,
865    failure_message: &str,
866) -> String {
867    format!(
868        "dirs=(\"${{MLSTACK_LOG_DIR:-}}\" \"$HOME/.rusty-stack/logs\" \"${{TMPDIR:-/tmp}}/rusty-stack/logs\")\n\
869         found=0\n\
870         for dir in \"${{dirs[@]}}\"; do\n\
871             if [ -n \"$dir\" ] && [ -d \"$dir\" ]; then\n\
872                 if ls \"$dir\" | grep -q \"{}\"; then\n\
873                     echo \"{}\"\n\
874                     found=1\n\
875                     break\n\
876                 fi\n\
877             fi\n\
878         done\n\
879         if [ \"$found\" -ne 1 ]; then\n\
880             echo \"{}\"\n\
881         fi",
882        pattern, success_message, failure_message
883    )
884}
885
886fn python_any(python_candidates: &[String], modules: &[&str]) -> bool {
887    python_candidates.iter().any(|python| {
888        modules
889            .iter()
890            .any(|module| python_has_module(python, module))
891    })
892}
893
894fn python_has_module(python: &str, module: &str) -> bool {
895    python_exec(
896        python,
897        &format!(
898            "import importlib.util, sys; sys.exit(0 if importlib.util.find_spec('{module}') else 1)"
899        ),
900    )
901}
902
903fn python_exec(python: &str, code: &str) -> bool {
904    let mut cmd = Command::new(python);
905    cmd.arg("-c")
906        .arg(code)
907        .stdout(Stdio::null())
908        .stderr(Stdio::null());
909
910    if let Ok(existing) = env::var("PYTHONPATH") {
911        let mut path = existing;
912        let extra_paths = python_search_paths();
913        if !extra_paths.is_empty() {
914            if !path.is_empty() {
915                path.push(':');
916            }
917            path.push_str(&extra_paths.join(":"));
918        }
919        cmd.env("PYTHONPATH", path);
920    } else {
921        let extra_paths = python_search_paths();
922        if !extra_paths.is_empty() {
923            cmd.env("PYTHONPATH", extra_paths.join(":"));
924        }
925    }
926
927    let home = resolve_component_user_home();
928    let ld_path = component_status_ld_library_path(&home);
929    if !ld_path.is_empty() {
930        cmd.env("LD_LIBRARY_PATH", ld_path);
931    }
932
933    cmd.status().map(|status| status.success()).unwrap_or(false)
934}
935
936fn push_unique_path(paths: &mut Vec<String>, value: &str) {
937    let value = value.trim();
938    if value.is_empty() {
939        return;
940    }
941    if paths.iter().any(|existing| existing == value) {
942        return;
943    }
944    paths.push(value.to_string());
945}
946
947fn component_user_home_from_passwd(user_name: &str) -> Option<String> {
948    let target = user_name.trim();
949    if target.is_empty() {
950        return None;
951    }
952    let passwd = fs::read_to_string("/etc/passwd").ok()?;
953    passwd.lines().find_map(|line| {
954        if line.trim().is_empty() || line.starts_with('#') {
955            return None;
956        }
957        let fields: Vec<&str> = line.split(':').collect();
958        if fields.len() < 6 || fields[0] != target {
959            return None;
960        }
961        let home = fields[5].trim();
962        if home.is_empty() {
963            None
964        } else {
965            Some(home.to_string())
966        }
967    })
968}
969
970fn component_passwd_homes_with_mlstack() -> Vec<String> {
971    let mut homes = Vec::new();
972    let Ok(passwd) = fs::read_to_string("/etc/passwd") else {
973        return homes;
974    };
975
976    for line in passwd.lines() {
977        if line.trim().is_empty() || line.starts_with('#') {
978            continue;
979        }
980        let fields: Vec<&str> = line.split(':').collect();
981        if fields.len() < 6 {
982            continue;
983        }
984        let home = fields[5].trim();
985        if home.is_empty() {
986            continue;
987        }
988        if Path::new(home).join(".mlstack").is_dir()
989            || Path::new(home).join(".mlstack_env").is_file()
990        {
991            push_unique_path(&mut homes, home);
992        }
993    }
994
995    homes
996}
997
998fn component_candidate_homes(preferred_home: &str) -> Vec<String> {
999    let mut homes = Vec::new();
1000    push_unique_path(&mut homes, preferred_home);
1001
1002    if let Ok(value) = env::var("MLSTACK_USER_HOME") {
1003        push_unique_path(&mut homes, &value);
1004    }
1005    if let Ok(value) = env::var("HOME") {
1006        push_unique_path(&mut homes, &value);
1007    }
1008
1009    for key in ["SUDO_USER", "USER", "LOGNAME"] {
1010        if let Ok(user_name) = env::var(key) {
1011            if let Some(home) = component_user_home_from_passwd(&user_name) {
1012                push_unique_path(&mut homes, &home);
1013            }
1014        }
1015    }
1016
1017    for home in component_passwd_homes_with_mlstack() {
1018        push_unique_path(&mut homes, &home);
1019    }
1020
1021    homes
1022}
1023
1024fn resolve_component_user_home() -> String {
1025    let fallback = env::var("HOME").unwrap_or_else(|_| "/tmp".to_string());
1026    let candidates = component_candidate_homes(&fallback);
1027
1028    for home in &candidates {
1029        if Path::new(home).join(".mlstack_env").is_file() {
1030            return home.clone();
1031        }
1032    }
1033    for home in &candidates {
1034        if Path::new(home).join(".mlstack").is_dir() {
1035            return home.clone();
1036        }
1037    }
1038
1039    candidates.into_iter().next().unwrap_or(fallback)
1040}
1041
1042fn component_status_mpi_compat_dirs(home: &str) -> Vec<String> {
1043    let mut dirs = Vec::new();
1044    for home in component_candidate_homes(home) {
1045        let mlstack_dir = Path::new(&home).join(".mlstack");
1046        let primary = mlstack_dir.join("libmpi-compat");
1047        if primary.is_dir() {
1048            dirs.push(primary.to_string_lossy().to_string());
1049        }
1050
1051        if let Ok(entries) = fs::read_dir(&mlstack_dir) {
1052            for entry in entries.flatten() {
1053                let path = entry.path();
1054                if !path.is_dir() {
1055                    continue;
1056                }
1057                let Some(name) = path.file_name().and_then(|value| value.to_str()) else {
1058                    continue;
1059                };
1060                if name.starts_with("libmpi-compat-user-") {
1061                    dirs.push(path.to_string_lossy().to_string());
1062                }
1063            }
1064        }
1065    }
1066
1067    dirs.sort();
1068    dirs.dedup();
1069    dirs
1070}
1071
1072fn component_status_ld_library_path(home: &str) -> String {
1073    let mut paths = Vec::new();
1074
1075    for dir in component_status_mpi_compat_dirs(home) {
1076        push_unique_path(&mut paths, &dir);
1077    }
1078
1079    for rocm_path in ["/opt/rocm/lib", "/opt/rocm/hip/lib", "/opt/rocm/opencl/lib"] {
1080        if Path::new(rocm_path).exists() {
1081            push_unique_path(&mut paths, rocm_path);
1082        }
1083    }
1084
1085    if let Ok(existing) = env::var("LD_LIBRARY_PATH") {
1086        for part in existing.split(':') {
1087            push_unique_path(&mut paths, part);
1088        }
1089    }
1090
1091    paths.join(":")
1092}
1093
1094fn command_exists(command: &str) -> bool {
1095    Command::new("bash")
1096        .arg("-c")
1097        .arg(format!("command -v {}", command))
1098        .stdout(Stdio::null())
1099        .stderr(Stdio::null())
1100        .status()
1101        .map(|status| status.success())
1102        .unwrap_or(false)
1103}
1104
1105pub fn python_search_paths() -> Vec<String> {
1106    let home = resolve_component_user_home();
1107    let mut paths = vec![
1108        format!("{}/pytorch", home),
1109        format!("{}/ml_stack/flash_attn_amd_direct", home),
1110        format!("{}/ml_stack/flash_attn_amd", home),
1111        format!(
1112            "{}/ml_stack/flash_attn_amd/build/lib.linux-x86_64-cpython-313",
1113            home
1114        ),
1115        format!(
1116            "{}/ml_stack/flash_attn_amd/build/lib.linux-x86_64-cpython-312",
1117            home
1118        ),
1119        format!("{}/.local/lib/python3.13/site-packages", home),
1120        format!("{}/.local/lib/python3.12/site-packages", home),
1121        format!("{}/rocm_venv/lib/python3.13/site-packages", home),
1122        format!("{}/rocm_venv/lib/python3.12/site-packages", home),
1123        format!("{}/megatron/Megatron-LM", home),
1124        format!("{}/ml_stack/bitsandbytes/bitsandbytes", home),
1125    ];
1126
1127    // Dynamically add /opt/rocm/lib/python*/site-packages and /opt/rocm/lib
1128    let rocm_lib = Path::new("/opt/rocm/lib");
1129    if rocm_lib.exists() {
1130        paths.push(rocm_lib.to_string_lossy().to_string());
1131        if let Ok(entries) = fs::read_dir(rocm_lib) {
1132            for entry in entries.flatten() {
1133                let path = entry.path();
1134                if path.is_dir() {
1135                    if let Some(name) = path.file_name().and_then(|n| n.to_str()) {
1136                        if name.starts_with("python") {
1137                            let sp = path.join("site-packages");
1138                            if sp.exists() {
1139                                paths.push(sp.to_string_lossy().to_string());
1140                            }
1141                        }
1142                    }
1143                }
1144            }
1145        }
1146    }
1147
1148    paths
1149}
1150
1151fn select_python_for_modules(modules: &[&str], python_candidates: &[String]) -> Option<String> {
1152    for key in ["MLSTACK_PYTHON_BIN", "UV_PYTHON"] {
1153        if let Ok(value) = env::var(key) {
1154            let value = value.trim();
1155            if !value.is_empty() {
1156                if modules.is_empty() {
1157                    return Some(value.to_string());
1158                }
1159                for module in modules {
1160                    if python_has_module(value, module) {
1161                        return Some(value.to_string());
1162                    }
1163                }
1164            }
1165        }
1166    }
1167
1168    for python in python_candidates {
1169        for module in modules {
1170            if python_has_module(python, module) {
1171                return Some(python.clone());
1172            }
1173        }
1174    }
1175    None
1176}
1177
1178fn path_exists<P: AsRef<Path>>(path: P) -> bool {
1179    path.as_ref().exists()
1180}
1181
1182fn home_path(home: &str, parts: &[&str]) -> PathBuf {
1183    let mut path = PathBuf::from(home);
1184    for part in parts {
1185        path.push(part);
1186    }
1187    path
1188}
1189
1190fn env_file_has_enhanced(path: &Path) -> bool {
1191    if let Ok(contents) = fs::read_to_string(path) {
1192        return contents.contains("Enhanced ML Stack Environment Setup Script");
1193    }
1194    false
1195}
1196
1197fn env_file_has_permanent(path: &Path) -> bool {
1198    if let Ok(contents) = fs::read_to_string(path) {
1199        return contents.contains("Permanent ROCm Environment")
1200            || contents.contains("Permanent ROCm Env")
1201            || (contents.contains("ML Stack Environment File")
1202                && contents.contains("MLSTACK_PYTHON_BIN"));
1203    }
1204    false
1205}
1206
1207fn repo_has_ml_stack_core() -> bool {
1208    if let Ok(root) = env::var("MLSTACK_REPO_ROOT") {
1209        return path_exists(Path::new(&root).join("stans_ml_stack"));
1210    }
1211    false
1212}
1213
1214// ===========================================================================
1215// Tests
1216// ===========================================================================
1217
1218#[cfg(test)]
1219mod tests {
1220    use super::*;
1221
1222    #[test]
1223    fn test_pytorch_diagnostic_snippet_exits_strict() {
1224        // Stage 4: the snippet must FAIL (exit 1) when torch is not a ROCm build
1225        // or the HIP runtime is unavailable — it must NOT soft-pass with an
1226        // unconditional exit 0. Downstream gating is an orchestrator concern.
1227        let snippet = pytorch_diagnostic_snippet();
1228        assert!(
1229            snippet.contains("(hip_version is not None) and hip_available"),
1230            "Should gate the exit on functional HIP availability"
1231        );
1232        assert!(
1233            snippet.contains("sys.exit(0 if ok else 1)"),
1234            "Should exit non-zero when HIP is unavailable (fail loud)"
1235        );
1236        assert!(
1237            !snippet.contains("\nsys.exit(0)\n"),
1238            "Should NOT have an unconditional soft-pass sys.exit(0)"
1239        );
1240    }
1241
1242    #[test]
1243    fn test_pytorch_diagnostic_snippet_checks_hip() {
1244        // The snippet should check for HIP availability to provide diagnostics
1245        let snippet = pytorch_diagnostic_snippet();
1246        assert!(
1247            snippet.contains("torch.version"),
1248            "Should check torch.version for HIP info"
1249        );
1250        assert!(
1251            snippet.contains("torch.cuda.is_available()"),
1252            "Should check torch.cuda.is_available() for runtime HIP"
1253        );
1254    }
1255
1256    #[test]
1257    fn test_pytorch_diagnostic_snippet_provides_hints() {
1258        // The snippet should provide actionable hints when HIP is not available
1259        let snippet = pytorch_diagnostic_snippet();
1260        assert!(
1261            snippet.contains("HINT:"),
1262            "Should provide HINT messages for troubleshooting"
1263        );
1264        assert!(
1265            snippet.contains("LD_LIBRARY_PATH"),
1266            "Should mention LD_LIBRARY_PATH as a possible cause"
1267        );
1268        assert!(
1269            snippet.contains("/opt/rocm/lib"),
1270            "Should mention ROCm lib path"
1271        );
1272    }
1273
1274    #[test]
1275    fn test_pytorch_diagnostic_snippet_checks_rocm_libs() {
1276        // The snippet should check for ROCm library existence
1277        let snippet = pytorch_diagnostic_snippet();
1278        assert!(
1279            snippet.contains("libamdhip64"),
1280            "Should check for libamdhip64 (core HIP library)"
1281        );
1282        assert!(
1283            snippet.contains("/opt/rocm/lib"),
1284            "Should check /opt/rocm/lib directory"
1285        );
1286    }
1287
1288    #[test]
1289    fn test_pytorch_component_verification_uses_diagnostic() {
1290        // Verify that the pytorch component verification command uses
1291        // the diagnostic snippet (not the old strict check)
1292        let candidates = vec!["python3".to_string()];
1293        let cmds = component_verification_commands("pytorch", &candidates);
1294        assert_eq!(
1295            cmds.len(),
1296            1,
1297            "Should have exactly one verification command"
1298        );
1299        let cmd = &cmds[0];
1300        assert_eq!(cmd.label, "PyTorch");
1301        assert_eq!(cmd.target_id, "pytorch");
1302        // Stage 4: must use the strict FUNCTIONAL check — fail when torch has
1303        // no ROCm build or the HIP runtime is unavailable (no soft-pass).
1304        let empty = String::new();
1305        let code = cmd.args.get(1).unwrap_or(&empty);
1306        assert!(
1307            code.contains("(hip_version is not None) and hip_available"),
1308            "Should gate the exit on functional HIP availability (Stage 4)"
1309        );
1310        assert!(
1311            code.contains("sys.exit(0 if ok else 1)"),
1312            "Should exit non-zero when HIP is unavailable (Stage 4: fail loud)"
1313        );
1314    }
1315
1316    #[test]
1317    fn test_pytorch_basic_verification_uses_diagnostic() {
1318        // Verify that basic verification also uses the strict functional check
1319        let candidates = vec!["python3".to_string()];
1320        let cmds = basic_verification_commands(&candidates);
1321        let pytorch_cmd = cmds.iter().find(|c| c.target_id == "pytorch");
1322        assert!(
1323            pytorch_cmd.is_some(),
1324            "Should have a PyTorch verification command"
1325        );
1326        let cmd = pytorch_cmd.unwrap();
1327        let empty = String::new();
1328        let code = cmd.args.get(1).unwrap_or(&empty);
1329        assert!(
1330            code.contains("(hip_version is not None) and hip_available"),
1331            "Basic verification should gate exit on functional HIP (Stage 4)"
1332        );
1333        assert!(
1334            code.contains("sys.exit(0 if ok else 1)"),
1335            "Basic verification should fail loud when HIP unavailable (Stage 4)"
1336        );
1337    }
1338
1339    #[test]
1340    fn test_pytorch_installation_check_only_imports() {
1341        // The is_component_installed_by_id check for pytorch should only
1342        // check if the torch module is importable (not check HIP).
1343        // This is the correct basic detection — HIP availability is a
1344        // verification concern, not an installation concern.
1345        let candidates = vec!["python3".to_string()];
1346        // We can't assert the result since it depends on the system state,
1347        // but we can verify the function doesn't panic
1348        let _ = is_component_installed_by_id("pytorch", &candidates);
1349    }
1350
1351    #[test]
1352    fn test_verification_command_structure() {
1353        // Verify the VerificationCommand struct for pytorch is well-formed
1354        let candidates = vec!["python3".to_string()];
1355        let cmds = component_verification_commands("pytorch", &candidates);
1356        assert_eq!(cmds.len(), 1);
1357        let cmd = &cmds[0];
1358        assert!(!cmd.label.is_empty());
1359        assert!(!cmd.target_id.is_empty());
1360        assert!(!cmd.program.is_empty());
1361        assert_eq!(cmd.args.len(), 2, "Should have -c and the code");
1362        assert_eq!(cmd.args[0], "-c");
1363        assert!(!cmd.args[1].is_empty());
1364        assert!(cmd.modules.contains(&"torch".to_string()));
1365    }
1366}
1367
1368#[cfg(test)]
1369mod component_status_tests;