rusty_stack/

component_status.rs

use crate::benchmark_logs;
use crate::state::{Category, Component};
use std::collections::HashSet;
use std::env;
use std::fs;
use std::path::{Path, PathBuf};
use std::process::{Command, Stdio};

#[derive(Debug, Clone)]
pub struct VerificationCommand {
    pub label: String,
    pub target_id: String,
    pub program: String,
    pub args: Vec<String>,
    pub modules: Vec<String>,
}

pub fn python_interpreters() -> Vec<String> {
    let mut candidates = Vec::new();
    if let Ok(value) = env::var("MLSTACK_PYTHON_BIN") {
        push_python_candidate(&mut candidates, value);
    }
    if let Ok(value) = env::var("AITER_VENV_PYTHON") {
        push_python_candidate(&mut candidates, value);
    }
    if let Ok(value) = env::var("UV_PYTHON") {
        push_python_candidate(&mut candidates, value);
    }
    for key in [
        "WANDB_VENV_PYTHON",
        "MEGATRON_VENV_PYTHON",
        "MPI4PY_VENV_PYTHON",
        "FLASH_ATTENTION_VENV_PYTHON",
        "PYTORCH_VENV_PYTHON",
        "DEEPSPEED_VENV_PYTHON",
        "VLLM_VENV_PYTHON",
    ] {
        if let Ok(value) = env::var(key) {
            push_python_candidate(&mut candidates, value);
        }
    }
    let home = resolve_component_user_home();
    let venv = Path::new(&home)
        .join("rocm_venv")
        .join("bin")
        .join("python");
    if venv.exists() {
        candidates.push(venv.to_string_lossy().to_string());
    }

    // Anaconda/Miniconda candidates
    let conda_path = Path::new(&home)
        .join("anaconda3")
        .join("bin")
        .join("python");
    if conda_path.exists() {
        candidates.push(conda_path.to_string_lossy().to_string());
    }
    let miniconda_path = Path::new(&home)
        .join("miniconda3")
        .join("bin")
        .join("python");
    if miniconda_path.exists() {
        candidates.push(miniconda_path.to_string_lossy().to_string());
    }

    if Path::new("/usr/local/bin/python3").exists() {
        candidates.push("/usr/local/bin/python3".to_string());
    }
    if command_exists("python3.13") {
        candidates.push("python3.13".to_string());
    }
    if command_exists("python3.12") {
        candidates.push("python3.12".to_string());
    }
    if command_exists("python3") {
        candidates.push("python3".to_string());
    }
    if command_exists("python") {
        candidates.push("python".to_string());
    }

    let mut roots = vec![PathBuf::from(&home)];
    if let Ok(cwd) = env::current_dir() {
        roots.push(cwd.clone());
        roots.push(cwd.join("rusty-stack"));
    }
    for root in roots {
        push_component_venv_candidates(&mut candidates, &root);
    }

    dedupe_python_candidates_in_order(candidates)
}

fn dedupe_python_candidates_in_order(candidates: Vec<String>) -> Vec<String> {
    let mut seen = HashSet::new();
    let mut ordered = Vec::new();
    for candidate in candidates {
        if seen.insert(candidate.clone()) {
            ordered.push(candidate);
        }
    }
    ordered
}

fn push_component_venv_candidates(candidates: &mut Vec<String>, root: &Path) {
    for rel in [
        "wandb_venv/bin/python",
        "megatron_rocm_venv/bin/python",
        "Megatron-LM/megatron_rocm_venv/bin/python",
        "mpi4py_venv/bin/python",
        "flash_attention_venv/bin/python",
        "vllm_venv/bin/python",
        ".mlstack/venvs/aiter/bin/python",
        ".mlstack/venvs/vllm/bin/python",
        "pytorch_rocm_venv/bin/python",
    ] {
        let candidate = root.join(rel);
        if candidate.exists() {
            candidates.push(candidate.to_string_lossy().to_string());
        }
    }
}

fn push_python_candidate(candidates: &mut Vec<String>, value: String) {
    let trimmed = value.trim();
    if trimmed.is_empty() {
        return;
    }
    let candidate = trimmed.to_string();
    if Path::new(&candidate).exists() || command_exists(&candidate) {
        candidates.push(candidate);
    }
}

pub fn is_component_installed(component: &Component, python_candidates: &[String]) -> bool {
    if component.category == Category::Verification {
        return false;
    }
    is_component_installed_by_id(&component.id, python_candidates)
}

pub fn is_component_installed_by_id(component_id: &str, python_candidates: &[String]) -> bool {
    let home = resolve_component_user_home();
    let has_benchmark_logs = has_benchmark_log_dirs();
    match component_id {
        // ROCm requires BOTH version file AND functional rocminfo (not just binary existence)
        // This prevents false positives from partial downloads
        "rocm" => {
            // Check for version file in multiple locations
            let version_exists = path_exists("/opt/rocm/.info/version")
                || path_exists("/opt/rocm/version")
                || path_exists("/usr/lib/rocm/.info/version");

            // Check if rocminfo actually works (not just exists)
            let rocminfo_path = crate::installers::common::utils::resolve_rocminfo_path();
            let rocminfo_functional = if let Ok(output) = std::process::Command::new(&rocminfo_path)
                .arg("--version")
                .output()
            {
                output.status.success()
            } else {
                false
            };

            // Require both: version file AND functional rocminfo
            // This ensures we don't report partial installs as complete
            version_exists && rocminfo_functional
        }
        "pytorch" => python_any(python_candidates, &["torch"]),
        "triton" => python_any(python_candidates, &["triton"]),
        "mpi4py" => python_any(python_candidates, &["mpi4py"]),
        "deepspeed" => python_any(python_candidates, &["deepspeed"]),
        "ml-stack-core" => {
            python_any(python_candidates, &["stans_ml_stack"]) || repo_has_ml_stack_core()
        }
        "flash-attn" => {
            python_any(python_candidates, &["flash_attn", "flash_attn_2"])
                || path_exists(home_path(&home, &["ml_stack", "flash_attn_amd"]))
                || path_exists(home_path(&home, &["ml_stack", "flash_attn_amd_direct"]))
        }
        "megatron" => {
            path_exists(home_path(&home, &["megatron", "Megatron-LM"]))
                || python_any(python_candidates, &["megatron"])
        }
        "vllm" => {
            python_candidates
                .iter()
                .any(|python| python_exec(python, vllm_runtime_check_snippet()))
                || path_exists(home_path(&home, &["vllm_build"]))
                || path_exists(home_path(&home, &["vllm_py313"]))
        }
        "aiter" => python_any(python_candidates, &["aiter"]),
        "vllm-studio" => {
            command_exists("vllm-studio") || path_exists(home_path(&home, &["vllm-studio"]))
        }
        "onnx" => {
            python_any(python_candidates, &["onnxruntime"])
                || path_exists(home_path(&home, &["onnxruntime_build"]))
        }
        "bitsandbytes" => {
            python_any(python_candidates, &["bitsandbytes"])
                || path_exists(home_path(&home, &["ml_stack", "bitsandbytes"]))
        }
        // rocm-smi - check multiple locations and verify it works
        "rocm-smi" => {
            let rocm_smi_paths = [
                "/opt/rocm/bin/rocm-smi",
                "/usr/bin/rocm-smi",
                "/usr/local/bin/rocm-smi",
            ];
            for path in &rocm_smi_paths {
                if path_exists(path) {
                    // Verify it actually works
                    if let Ok(output) = std::process::Command::new(path).arg("--showuse").output() {
                        if output.status.success() {
                            return true;
                        }
                    }
                }
            }
            // Also check if rocm-smi is in PATH and works
            if command_exists("rocm-smi") {
                if let Ok(output) = std::process::Command::new("rocm-smi")
                    .arg("--showuse")
                    .output()
                {
                    return output.status.success();
                }
            }
            false
        }
        "migraphx" => {
            python_any(python_candidates, &["migraphx"])
                || path_exists(home_path(&home, &["migraphx_build"]))
        }
        "pytorch-profiler" => python_any(python_candidates, &["torch"]),
        "wandb" => python_any(python_candidates, &["wandb"]),
        "fastvideo" => python_any(python_candidates, &["fastvideo"]),
        "llama-cpp" => {
            // Use the detection contract from llama_cpp.rs: llama-cli --help must succeed
            crate::installers::components::llama_cpp::is_llama_cli_functional(&home)
        }
        "permanent-env" => env_file_has_permanent(&home_path(&home, &[".mlstack_env"])),
        "basic-env" => path_exists(home_path(&home, &[".mlstack_env"])),
        "enhanced-env" => env_file_has_enhanced(&home_path(&home, &[".mlstack_env"])),
        "vllm-performance" => has_benchmark_logs,
        "deepspeed-performance" => has_benchmark_logs,
        "megatron-performance" => has_benchmark_logs,
        "all-benchmarks" => has_benchmark_logs,
        "comfyui" => {
            // Check for ComfyUI installation
            // Default location: $HOME/ComfyUI
            // User location: /mnt/e0f7c1a8-b834-4827-b579-0251b006bc1f/ComfyUI/
            // Check for .git folder and main.py
            let default_path = home_path(&home, &["ComfyUI"]);
            path_exists(default_path.join(".git")) || path_exists(default_path.join("main.py"))
        }
        "textgen" => {
            // Check for text-generation-webui installation
            // Default location: $HOME/text-generation-webui
            let default_path = home_path(&home, &["text-generation-webui"]);
            path_exists(default_path.join(".git")) || path_exists(default_path.join("server.py"))
        }
        _ => false,
    }
}

pub fn verification_commands(
    component_id: &str,
    python_candidates: &[String],
) -> Vec<VerificationCommand> {
    match component_id {
        "verify-basic" => basic_verification_commands(python_candidates),
        "verify-enhanced" => enhanced_verification_commands(python_candidates),
        "verify-build" => build_verification_commands(python_candidates),
        _ => Vec::new(),
    }
}

pub fn component_verification_commands(
    component_id: &str,
    python_candidates: &[String],
) -> Vec<VerificationCommand> {
    match component_id {
        "rocm" => vec![shell_command("ROCm info", "rocm", "rocminfo", &[])],
        "pytorch" => vec![python_command(
            "PyTorch",
            "pytorch",
            &["torch"],
            python_candidates,
            pytorch_diagnostic_snippet(),
        )],
        "triton" => vec![python_command(
            "Triton",
            "triton",
            &["triton"],
            python_candidates,
            "import triton; import triton.language as tl; print(f'Triton {triton.__version__} — language OK')",
        )],
        "mpi4py" => vec![python_command(
            "MPI4Py",
            "mpi4py",
            &["mpi4py"],
            python_candidates,
            "import importlib.metadata as m; import mpi4py; print(getattr(mpi4py, '__version__', m.version('mpi4py')))",
        )],
        "deepspeed" => vec![python_command(
            "DeepSpeed",
            "deepspeed",
            &["deepspeed"],
            python_candidates,
            "import deepspeed; print(deepspeed.__version__)",
        )],
        "ml-stack-core" => vec![python_command(
            "ML Stack Core",
            "ml-stack-core",
            &["stans_ml_stack"],
            python_candidates,
            "import stans_ml_stack; print(getattr(stans_ml_stack, '__version__', 'ok'))",
        )],
        "flash-attn" => vec![python_command(
            "Flash Attention",
            "flash-attn",
            &["flash_attn", "flash_attention_amd", "flash_attn_2"],
            python_candidates,
            "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)",
        )],
        "megatron" => vec![python_command(
            "Megatron-LM",
            "megatron",
            &["megatron"],
            python_candidates,
            "import megatron; print('megatron ok')",
        )],
        "vllm" => vec![python_command(
            "vLLM",
            "vllm",
            &["vllm"],
            python_candidates,
            vllm_runtime_check_snippet(),
        )],
        "aiter" => vec![python_command(
            "AITER",
            "aiter",
            &["aiter"],
            python_candidates,
            "import aiter; mods=[x for x in dir(aiter) if not x.startswith('_')]; print(f'AITER imported — {len(mods)} exports available')",
        )],
        "vllm-studio" => vec![shell_command(
            "vLLM Studio",
            "vllm-studio",
            "bun",
            &["--version"],
        )],
        "onnx" => vec![python_command(
            "ONNX Runtime",
            "onnx",
            &["onnxruntime"],
            python_candidates,
            "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)",
        )],
        "bitsandbytes" => vec![python_command(
            "BITSANDBYTES",
            "bitsandbytes",
            &["bitsandbytes"],
            python_candidates,
            "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)",
        )],
        "rocm-smi" => vec![shell_command(
            "ROCm SMI",
            "rocm-smi",
            "rocm-smi",
            &["--showproductname"],
        )],
        "migraphx" => vec![python_command(
            "MIGraphX",
            "migraphx",
            &["migraphx"],
            python_candidates,
            "import migraphx; print(getattr(migraphx, '__version__', 'ok'))",
        )],
        "pytorch-profiler" => vec![python_command(
            "PyTorch Profiler",
            "pytorch-profiler",
            &["torch"],
            python_candidates,
            "from torch.profiler import profile; print(profile)",
        )],
        "wandb" => vec![python_command(
            "Weights & Biases",
            "wandb",
            &["wandb"],
            python_candidates,
            "import wandb; print(wandb.__version__)",
        )],
        "fastvideo" => vec![python_command(
            "FastVideo",
            "fastvideo",
            &["fastvideo"],
            python_candidates,
            "import fastvideo; print(fastvideo.__version__)",
        )],
        "llama-cpp" => vec![shell_command(
            "llama.cpp",
            "llama-cpp",
            "llama-cli",
            &["--help"],
        )],
        "basic-env" => vec![shell_command(
            "Environment file",
            "basic-env",
            "bash",
            &[
                "-c",
                "test -f \"$HOME/.mlstack_env\" && echo OK || exit 1",
            ],
        )],
        "enhanced-env" => vec![shell_command(
            "Enhanced env",
            "enhanced-env",
            "bash",
            &[
                "-c",
                "grep -q \"Enhanced ML Stack Environment Setup Script\" \"$HOME/.mlstack_env\" && echo OK || exit 1",
            ],
        )],
        "permanent-env" => vec![shell_command(
            "Permanent env",
            "permanent-env",
            "bash",
            &[
                "-c",
                "grep -qiE \"Permanent ROCm Environment|Permanent ROCm Env|MLSTACK_PYTHON_BIN\" \"$HOME/.mlstack_env\" && echo OK || exit 1",
            ],
        )],
        // Benchmark components - verify by checking log files exist
        "mlperf-inference" => vec![benchmark_log_check_command(
            "MLPerf benchmark logs",
            "mlperf-inference",
            "mlperf_inference",
            "Benchmark logs found",
            "No benchmark logs yet",
        )],
        "rocm-benchmarks" => vec![benchmark_log_check_command(
            "ROCm benchmark logs",
            "rocm-benchmarks",
            "rocm_benchmarks",
            "Benchmark logs found",
            "No benchmark logs yet",
        )],
        "gpu-memory-bandwidth" => vec![benchmark_log_check_command(
            "Memory bandwidth logs",
            "gpu-memory-bandwidth",
            "gpu_memory_bandwidth",
            "Benchmark logs found",
            "No benchmark logs yet",
        )],
        "rocm-smi-bench" => vec![benchmark_log_check_command(
            "ROCm SMI logs",
            "rocm-smi-bench",
            "rocm_smi_benchmarks",
            "Benchmark logs found",
            "No benchmark logs yet",
        )],
        "pytorch-performance" => vec![benchmark_log_check_command(
            "PyTorch benchmark logs",
            "pytorch-performance",
            "pytorch_performance",
            "Benchmark logs found",
            "No benchmark logs yet",
        )],
        "vllm-performance" => vec![benchmark_log_check_command(
            "vLLM benchmark logs",
            "vllm-performance",
            "vllm_benchmarks",
            "Benchmark logs found",
            "No benchmark logs yet",
        )],
        "deepspeed-performance" => vec![benchmark_log_check_command(
            "DeepSpeed benchmark logs",
            "deepspeed-performance",
            "deepspeed_benchmarks",
            "Benchmark logs found",
            "No benchmark logs yet",
        )],
        "megatron-performance" => vec![benchmark_log_check_command(
            "Megatron benchmark logs",
            "megatron-performance",
            "megatron_benchmarks",
            "Benchmark logs found",
            "No benchmark logs yet",
        )],
        "all-benchmarks" => vec![benchmark_log_check_command(
            "Full suite logs",
            "all-benchmarks",
            "full_benchmarks",
            "Suite logs found",
            "No suite logs yet",
        )],
        "comfyui" => vec![shell_command(
            "ComfyUI",
            "comfyui",
            "bash",
            &["-c", "test -f \"$HOME/ComfyUI/main.py\" && echo 'ComfyUI installed' || exit 1"],
        )],
        "textgen" => vec![shell_command(
            "text-generation-webui",
            "textgen",
            "bash",
            &["-c", "test -f \"$HOME/text-generation-webui/server.py\" && echo 'text-generation-webui installed' || exit 1"],
        )],
        _ => Vec::new(),
    }
}

pub fn modules_available(modules: &[String], python_candidates: &[String]) -> bool {
    if modules.is_empty() {
        return true;
    }
    let module_refs: Vec<&str> = modules.iter().map(|s| s.as_str()).collect();
    python_any(python_candidates, &module_refs)
}

fn basic_verification_commands(python_candidates: &[String]) -> Vec<VerificationCommand> {
    vec![
        shell_command("ROCm info", "rocm", "rocminfo", &[]),
        python_command(
            "PyTorch",
            "pytorch",
            &["torch"],
            python_candidates,
            pytorch_diagnostic_snippet(),
        ),
        python_command(
            "Triton",
            "triton",
            &["triton"],
            python_candidates,
            "import triton; import triton._C; print(f'Triton {triton.__version__} ok')",
        ),
        python_command(
            "MPI4Py",
            "mpi4py",
            &["mpi4py"],
            python_candidates,
            "import importlib.metadata as m; import mpi4py; print(getattr(mpi4py, '__version__', m.version('mpi4py')))",
        ),
        python_command(
            "DeepSpeed",
            "deepspeed",
            &["deepspeed", "einops"],
            python_candidates,
            "import deepspeed; import einops; print(deepspeed.__version__)",
        ),
    ]
}

fn enhanced_verification_commands(python_candidates: &[String]) -> Vec<VerificationCommand> {
    let mut steps = basic_verification_commands(python_candidates);
    steps.extend(vec![
        python_command(
            "Flash Attention",
            "flash-attn",
            &["flash_attn", "flash_attention_amd", "flash_attn_2"],
            python_candidates,
            "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)",
        ),
        python_command(
            "vLLM",
            "vllm",
            &["vllm"],
            python_candidates,
            vllm_runtime_check_snippet(),
        ),
        python_command(
            "AITER",
            "aiter",
            &["aiter"],
            python_candidates,
            "import aiter; print(getattr(aiter, '__version__', 'ok'))",
        ),
        python_command(
            "ONNX Runtime",
            "onnx",
            &["onnxruntime"],
            python_candidates,
            "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)",
        ),
        python_command(
            "BITSANDBYTES",
            "bitsandbytes",
            &["bitsandbytes"],
            python_candidates,
            "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)",
        ),
        shell_command("ROCm SMI", "rocm-smi", "rocm-smi", &["--showproductname"]),
        python_command(
            "MIGraphX",
            "migraphx",
            &["migraphx"],
            python_candidates,
            "import migraphx; print(getattr(migraphx, '__version__', 'ok'))",
        ),
        python_command(
            "PyTorch Profiler",
            "pytorch-profiler",
            &["torch"],
            python_candidates,
            "from torch.profiler import profile; print(profile)",
        ),
        python_command(
            "Weights & Biases",
            "wandb",
            &["wandb"],
            python_candidates,
            "import wandb; print(wandb.__version__)",
        ),
        python_command(
            "FastVideo",
            "fastvideo",
            &["fastvideo"],
            python_candidates,
            "import fastvideo; print(fastvideo.__version__)",
        ),
        shell_command("vLLM Studio", "vllm-studio", "bun", &["--version"]),
    ]);
    steps
}

fn build_verification_commands(python_candidates: &[String]) -> Vec<VerificationCommand> {
    let mut steps = basic_verification_commands(python_candidates);
    steps.extend(vec![
        python_command(
            "ONNX Runtime",
            "onnx",
            &["onnxruntime"],
            python_candidates,
            "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)",
        ),
        python_command(
            "Flash Attention (build)",
            "flash-attn",
            &["flash_attn", "flash_attention_amd", "flash_attn_2"],
            python_candidates,
            "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)",
        ),
        python_command(
            "MIGraphX",
            "migraphx",
            &["migraphx"],
            python_candidates,
            "import migraphx; print(getattr(migraphx, '__version__', 'ok'))",
        ),
    ]);
    steps
}

fn python_command(
    label: &str,
    target_id: &str,
    modules: &[&str],
    python_candidates: &[String],
    code: &str,
) -> VerificationCommand {
    let program = select_python_for_modules(modules, python_candidates)
        .or_else(|| python_candidates.first().cloned())
        .unwrap_or_else(|| "python3".to_string());

    VerificationCommand {
        label: label.to_string(),
        target_id: target_id.to_string(),
        program,
        args: vec!["-c".to_string(), code.to_string()],
        modules: modules.iter().map(|m| m.to_string()).collect(),
    }
}

/// PyTorch diagnostic verification snippet.
///
/// This performs a multi-tier check:
/// 1. Can we import torch? (basic install check)
/// 2. Is torch.version.hip set? (ROCm build check)
/// 3. Is torch.cuda.is_available()? (runtime HIP check)
/// 4. Diagnostics for why HIP might not be available
///
/// Unlike the old strict check that failed when `torch.version.hip` was None,
/// this approach succeeds if torch is importable and provides diagnostic
/// information about HIP availability. This prevents cascading failures
/// to downstream components (DeepSpeed, Flash Attention, etc.) when PyTorch
/// is correctly installed but HIP libraries aren't accessible from the
/// current Python environment (e.g., system Python vs venv, LD_LIBRARY_PATH
/// issues, or running without GPU access).
fn pytorch_diagnostic_snippet() -> &'static str {
    r#"import torch, sys, os, ctypes.util

# Tier 1: Basic torch import (already verified by module check)
print(f'torch version: {torch.__version__}')

# Tier 2: Check if this is a ROCm build
hip_version = getattr(torch.version, 'hip', None)
cuda_version = getattr(torch.version, 'cuda', None)
print(f'hip: {hip_version}')
print(f'cuda: {cuda_version}')

# Tier 3: Runtime HIP/CUDA availability
try:
    hip_available = torch.cuda.is_available()
    print(f'torch.cuda.is_available(): {hip_available}')
    if hip_available:
        device_count = torch.cuda.device_count()
        print(f'device_count: {device_count}')
        if device_count > 0:
            print(f'device_name: {torch.cuda.get_device_name(0)}')
except Exception as e:
    hip_available = False
    print(f'torch.cuda.is_available() error: {e}')

# Tier 4: Diagnostics if HIP not available
if not hip_available:
    print('--- ROCm HIP Diagnostics ---')

    # Check LD_LIBRARY_PATH
    ld_path = os.environ.get('LD_LIBRARY_PATH', '')
    print(f'LD_LIBRARY_PATH: {ld_path if ld_path else "(not set)"}')

    # Check ROCm_PATH
    rocm_path = os.environ.get('ROCM_PATH', '')
    print(f'ROCM_PATH: {rocm_path if rocm_path else "(not set)"}')

    # Check if ROCm libraries exist
    rocm_lib_dirs = ['/opt/rocm/lib', '/opt/rocm/hip/lib']
    for d in rocm_lib_dirs:
        exists = os.path.isdir(d)
        print(f'{d}: {"exists" if exists else "not found"}')

    # Check if libamdhip64.so can be found
    hip_lib = ctypes.util.find_library('amdhip64')
    print(f'libamdhip64: {hip_lib if hip_lib else "not found"}')

    # Check Python executable path
    print(f'python: {sys.executable}')

    # Provide actionable guidance
    if hip_version is None and cuda_version is None:
        print('HINT: torch.version.hip and torch.version.cuda are both None.')
        print('  This torch build may not have GPU support.')
        print('  Reinstall with ROCm index: pip install torch --index-url https://repo.radeon.com/rocm/manylinux/rocm-rel-7.2/')
    elif hip_version is not None and not hip_available:
        print('HINT: torch.version.hip is set but HIP runtime not available.')
        print('  Possible causes:')
        print('  1. LD_LIBRARY_PATH does not include /opt/rocm/lib')
        print('  2. Running in a container/environment without GPU access')
        print('  3. ROCm drivers not loaded (check with: rocminfo)')
        print('  Fix: source ~/.mlstack_env or set LD_LIBRARY_PATH=/opt/rocm/lib')

# Exit successfully if torch is importable (even without HIP)
# This prevents cascading failures to downstream components.
# The diagnostic output above tells the user if HIP is available.
sys.exit(0)"#
}

fn vllm_runtime_check_snippet() -> &'static str {
    r#"import importlib
import vllm
import cachetools
import cbor2
import gguf
import pybase64
import ijson
import mistral_common
import openai_harmony

loaded = []
errs = []
for name in ("vllm._C", "vllm._rocm_C"):
    try:
        importlib.import_module(name)
        loaded.append(name)
    except Exception as exc:
        errs.append(f"{name}: {exc}")

if not loaded:
    raise SystemExit("vLLM native extension load failure: " + " | ".join(errs))

print(vllm.__version__)"#
}

fn shell_command(
    label: &str,
    target_id: &str,
    program: &str,
    args: &[&str],
) -> VerificationCommand {
    VerificationCommand {
        label: label.to_string(),
        target_id: target_id.to_string(),
        program: program.to_string(),
        args: args.iter().map(|s| s.to_string()).collect(),
        modules: Vec::new(),
    }
}

fn has_benchmark_log_dirs() -> bool {
    benchmark_logs::benchmark_log_directories()
        .iter()
        .any(|dir| {
            fs::read_dir(dir)
                .ok()
                .map(|entries| {
                    entries
                        .filter_map(|entry| entry.ok())
                        .any(|entry| entry.path().is_file())
                })
                .unwrap_or(false)
        })
}

fn benchmark_log_check_command(
    label: &str,
    target_id: &str,
    pattern: &str,
    success_message: &str,
    failure_message: &str,
) -> VerificationCommand {
    let script = benchmark_log_check_script(pattern, success_message, failure_message);
    shell_command(label, target_id, "bash", &["-c", &script])
}

fn benchmark_log_check_script(
    pattern: &str,
    success_message: &str,
    failure_message: &str,
) -> String {
    format!(
        "dirs=(\"${{MLSTACK_LOG_DIR:-}}\" \"$HOME/.rusty-stack/logs\" \"${{TMPDIR:-/tmp}}/rusty-stack/logs\")\n\
         found=0\n\
         for dir in \"${{dirs[@]}}\"; do\n\
             if [ -n \"$dir\" ] && [ -d \"$dir\" ]; then\n\
                 if ls \"$dir\" | grep -q \"{}\"; then\n\
                     echo \"{}\"\n\
                     found=1\n\
                     break\n\
                 fi\n\
             fi\n\
         done\n\
         if [ \"$found\" -ne 1 ]; then\n\
             echo \"{}\"\n\
         fi",
        pattern, success_message, failure_message
    )
}

fn python_any(python_candidates: &[String], modules: &[&str]) -> bool {
    python_candidates.iter().any(|python| {
        modules
            .iter()
            .any(|module| python_has_module(python, module))
    })
}

fn python_has_module(python: &str, module: &str) -> bool {
    python_exec(
        python,
        &format!(
            "import importlib.util, sys; sys.exit(0 if importlib.util.find_spec('{module}') else 1)"
        ),
    )
}

fn python_exec(python: &str, code: &str) -> bool {
    let mut cmd = Command::new(python);
    cmd.arg("-c")
        .arg(code)
        .stdout(Stdio::null())
        .stderr(Stdio::null());

    if let Ok(existing) = env::var("PYTHONPATH") {
        let mut path = existing;
        let extra_paths = python_search_paths();
        if !extra_paths.is_empty() {
            if !path.is_empty() {
                path.push(':');
            }
            path.push_str(&extra_paths.join(":"));
        }
        cmd.env("PYTHONPATH", path);
    } else {
        let extra_paths = python_search_paths();
        if !extra_paths.is_empty() {
            cmd.env("PYTHONPATH", extra_paths.join(":"));
        }
    }

    let home = resolve_component_user_home();
    let ld_path = component_status_ld_library_path(&home);
    if !ld_path.is_empty() {
        cmd.env("LD_LIBRARY_PATH", ld_path);
    }

    cmd.status().map(|status| status.success()).unwrap_or(false)
}

fn push_unique_path(paths: &mut Vec<String>, value: &str) {
    let value = value.trim();
    if value.is_empty() {
        return;
    }
    if paths.iter().any(|existing| existing == value) {
        return;
    }
    paths.push(value.to_string());
}

fn component_user_home_from_passwd(user_name: &str) -> Option<String> {
    let target = user_name.trim();
    if target.is_empty() {
        return None;
    }
    let passwd = fs::read_to_string("/etc/passwd").ok()?;
    passwd.lines().find_map(|line| {
        if line.trim().is_empty() || line.starts_with('#') {
            return None;
        }
        let fields: Vec<&str> = line.split(':').collect();
        if fields.len() < 6 || fields[0] != target {
            return None;
        }
        let home = fields[5].trim();
        if home.is_empty() {
            None
        } else {
            Some(home.to_string())
        }
    })
}

fn component_passwd_homes_with_mlstack() -> Vec<String> {
    let mut homes = Vec::new();
    let Ok(passwd) = fs::read_to_string("/etc/passwd") else {
        return homes;
    };

    for line in passwd.lines() {
        if line.trim().is_empty() || line.starts_with('#') {
            continue;
        }
        let fields: Vec<&str> = line.split(':').collect();
        if fields.len() < 6 {
            continue;
        }
        let home = fields[5].trim();
        if home.is_empty() {
            continue;
        }
        if Path::new(home).join(".mlstack").is_dir()
            || Path::new(home).join(".mlstack_env").is_file()
        {
            push_unique_path(&mut homes, home);
        }
    }

    homes
}

fn component_candidate_homes(preferred_home: &str) -> Vec<String> {
    let mut homes = Vec::new();
    push_unique_path(&mut homes, preferred_home);

    if let Ok(value) = env::var("MLSTACK_USER_HOME") {
        push_unique_path(&mut homes, &value);
    }
    if let Ok(value) = env::var("HOME") {
        push_unique_path(&mut homes, &value);
    }

    for key in ["SUDO_USER", "USER", "LOGNAME"] {
        if let Ok(user_name) = env::var(key) {
            if let Some(home) = component_user_home_from_passwd(&user_name) {
                push_unique_path(&mut homes, &home);
            }
        }
    }

    for home in component_passwd_homes_with_mlstack() {
        push_unique_path(&mut homes, &home);
    }

    homes
}

fn resolve_component_user_home() -> String {
    let fallback = env::var("HOME").unwrap_or_else(|_| "/tmp".to_string());
    let candidates = component_candidate_homes(&fallback);

    for home in &candidates {
        if Path::new(home).join(".mlstack_env").is_file() {
            return home.clone();
        }
    }
    for home in &candidates {
        if Path::new(home).join(".mlstack").is_dir() {
            return home.clone();
        }
    }

    candidates.into_iter().next().unwrap_or(fallback)
}

fn component_status_mpi_compat_dirs(home: &str) -> Vec<String> {
    let mut dirs = Vec::new();
    for home in component_candidate_homes(home) {
        let mlstack_dir = Path::new(&home).join(".mlstack");
        let primary = mlstack_dir.join("libmpi-compat");
        if primary.is_dir() {
            dirs.push(primary.to_string_lossy().to_string());
        }

        if let Ok(entries) = fs::read_dir(&mlstack_dir) {
            for entry in entries.flatten() {
                let path = entry.path();
                if !path.is_dir() {
                    continue;
                }
                let Some(name) = path.file_name().and_then(|value| value.to_str()) else {
                    continue;
                };
                if name.starts_with("libmpi-compat-user-") {
                    dirs.push(path.to_string_lossy().to_string());
                }
            }
        }
    }

    dirs.sort();
    dirs.dedup();
    dirs
}

fn component_status_ld_library_path(home: &str) -> String {
    let mut paths = Vec::new();

    for dir in component_status_mpi_compat_dirs(home) {
        push_unique_path(&mut paths, &dir);
    }

    for rocm_path in ["/opt/rocm/lib", "/opt/rocm/hip/lib", "/opt/rocm/opencl/lib"] {
        if Path::new(rocm_path).exists() {
            push_unique_path(&mut paths, rocm_path);
        }
    }

    if let Ok(existing) = env::var("LD_LIBRARY_PATH") {
        for part in existing.split(':') {
            push_unique_path(&mut paths, part);
        }
    }

    paths.join(":")
}

fn command_exists(command: &str) -> bool {
    Command::new("bash")
        .arg("-c")
        .arg(format!("command -v {}", command))
        .stdout(Stdio::null())
        .stderr(Stdio::null())
        .status()
        .map(|status| status.success())
        .unwrap_or(false)
}

pub fn python_search_paths() -> Vec<String> {
    let home = resolve_component_user_home();
    let mut paths = vec![
        format!("{}/pytorch", home),
        format!("{}/ml_stack/flash_attn_amd_direct", home),
        format!("{}/ml_stack/flash_attn_amd", home),
        format!(
            "{}/ml_stack/flash_attn_amd/build/lib.linux-x86_64-cpython-313",
            home
        ),
        format!(
            "{}/ml_stack/flash_attn_amd/build/lib.linux-x86_64-cpython-312",
            home
        ),
        format!("{}/.local/lib/python3.13/site-packages", home),
        format!("{}/.local/lib/python3.12/site-packages", home),
        format!("{}/rocm_venv/lib/python3.13/site-packages", home),
        format!("{}/rocm_venv/lib/python3.12/site-packages", home),
        format!("{}/megatron/Megatron-LM", home),
        format!("{}/ml_stack/bitsandbytes/bitsandbytes", home),
    ];

    // Dynamically add /opt/rocm/lib/python*/site-packages and /opt/rocm/lib
    let rocm_lib = Path::new("/opt/rocm/lib");
    if rocm_lib.exists() {
        paths.push(rocm_lib.to_string_lossy().to_string());
        if let Ok(entries) = fs::read_dir(rocm_lib) {
            for entry in entries.flatten() {
                let path = entry.path();
                if path.is_dir() {
                    if let Some(name) = path.file_name().and_then(|n| n.to_str()) {
                        if name.starts_with("python") {
                            let sp = path.join("site-packages");
                            if sp.exists() {
                                paths.push(sp.to_string_lossy().to_string());
                            }
                        }
                    }
                }
            }
        }
    }

    paths
}

fn select_python_for_modules(modules: &[&str], python_candidates: &[String]) -> Option<String> {
    for key in ["MLSTACK_PYTHON_BIN", "UV_PYTHON"] {
        if let Ok(value) = env::var(key) {
            let value = value.trim();
            if !value.is_empty() {
                if modules.is_empty() {
                    return Some(value.to_string());
                }
                for module in modules {
                    if python_has_module(value, module) {
                        return Some(value.to_string());
                    }
                }
            }
        }
    }

    for python in python_candidates {
        for module in modules {
            if python_has_module(python, module) {
                return Some(python.clone());
            }
        }
    }
    None
}

fn path_exists<P: AsRef<Path>>(path: P) -> bool {
    path.as_ref().exists()
}

fn home_path(home: &str, parts: &[&str]) -> PathBuf {
    let mut path = PathBuf::from(home);
    for part in parts {
        path.push(part);
    }
    path
}

fn env_file_has_enhanced(path: &Path) -> bool {
    if let Ok(contents) = fs::read_to_string(path) {
        return contents.contains("Enhanced ML Stack Environment Setup Script");
    }
    false
}

fn env_file_has_permanent(path: &Path) -> bool {
    if let Ok(contents) = fs::read_to_string(path) {
        return contents.contains("Permanent ROCm Environment")
            || contents.contains("Permanent ROCm Env")
            || (contents.contains("ML Stack Environment File")
                && contents.contains("MLSTACK_PYTHON_BIN"));
    }
    false
}

fn repo_has_ml_stack_core() -> bool {
    if let Ok(root) = env::var("MLSTACK_REPO_ROOT") {
        return path_exists(Path::new(&root).join("stans_ml_stack"));
    }
    false
}

// ===========================================================================
// Tests
// ===========================================================================

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_pytorch_diagnostic_snippet_exits_zero() {
        // The diagnostic snippet should always exit 0 (not fail) even when
        // HIP is not available. This prevents cascading failures to downstream
        // components like DeepSpeed, Flash Attention, etc.
        let snippet = pytorch_diagnostic_snippet();
        assert!(
            snippet.contains("sys.exit(0)"),
            "PyTorch diagnostic snippet should exit 0 (not block downstream components)"
        );
        assert!(
            !snippet.contains("sys.exit(1)"),
            "PyTorch diagnostic snippet should NOT exit 1 (was the old strict check)"
        );
    }

    #[test]
    fn test_pytorch_diagnostic_snippet_checks_hip() {
        // The snippet should check for HIP availability to provide diagnostics
        let snippet = pytorch_diagnostic_snippet();
        assert!(
            snippet.contains("torch.version"),
            "Should check torch.version for HIP info"
        );
        assert!(
            snippet.contains("torch.cuda.is_available()"),
            "Should check torch.cuda.is_available() for runtime HIP"
        );
    }

    #[test]
    fn test_pytorch_diagnostic_snippet_provides_hints() {
        // The snippet should provide actionable hints when HIP is not available
        let snippet = pytorch_diagnostic_snippet();
        assert!(
            snippet.contains("HINT:"),
            "Should provide HINT messages for troubleshooting"
        );
        assert!(
            snippet.contains("LD_LIBRARY_PATH"),
            "Should mention LD_LIBRARY_PATH as a possible cause"
        );
        assert!(
            snippet.contains("/opt/rocm/lib"),
            "Should mention ROCm lib path"
        );
    }

    #[test]
    fn test_pytorch_diagnostic_snippet_checks_rocm_libs() {
        // The snippet should check for ROCm library existence
        let snippet = pytorch_diagnostic_snippet();
        assert!(
            snippet.contains("libamdhip64"),
            "Should check for libamdhip64 (core HIP library)"
        );
        assert!(
            snippet.contains("/opt/rocm/lib"),
            "Should check /opt/rocm/lib directory"
        );
    }

    #[test]
    fn test_pytorch_component_verification_uses_diagnostic() {
        // Verify that the pytorch component verification command uses
        // the diagnostic snippet (not the old strict check)
        let candidates = vec!["python3".to_string()];
        let cmds = component_verification_commands("pytorch", &candidates);
        assert_eq!(
            cmds.len(),
            1,
            "Should have exactly one verification command"
        );
        let cmd = &cmds[0];
        assert_eq!(cmd.label, "PyTorch");
        assert_eq!(cmd.target_id, "pytorch");
        // The args should contain the diagnostic snippet, not the old strict check
        let empty = String::new();
        let code = cmd.args.get(1).unwrap_or(&empty);
        assert!(
            !code.contains("sys.exit(0 if getattr(torch.version, 'hip', None) else 1)"),
            "Should NOT use old strict HIP check that exits 1"
        );
        assert!(
            code.contains("sys.exit(0)"),
            "Should use new diagnostic snippet that always exits 0"
        );
    }

    #[test]
    fn test_pytorch_basic_verification_uses_diagnostic() {
        // Verify that basic verification also uses the diagnostic snippet
        let candidates = vec!["python3".to_string()];
        let cmds = basic_verification_commands(&candidates);
        let pytorch_cmd = cmds.iter().find(|c| c.target_id == "pytorch");
        assert!(
            pytorch_cmd.is_some(),
            "Should have a PyTorch verification command"
        );
        let cmd = pytorch_cmd.unwrap();
        let empty = String::new();
        let code = cmd.args.get(1).unwrap_or(&empty);
        assert!(
            !code.contains("sys.exit(0 if getattr(torch.version, 'hip', None) else 1)"),
            "Basic verification should NOT use old strict HIP check"
        );
        assert!(
            code.contains("sys.exit(0)"),
            "Basic verification should use new diagnostic snippet"
        );
    }

    #[test]
    fn test_pytorch_installation_check_only_imports() {
        // The is_component_installed_by_id check for pytorch should only
        // check if the torch module is importable (not check HIP).
        // This is the correct basic detection — HIP availability is a
        // verification concern, not an installation concern.
        let candidates = vec!["python3".to_string()];
        // We can't assert the result since it depends on the system state,
        // but we can verify the function doesn't panic
        let _ = is_component_installed_by_id("pytorch", &candidates);
    }

    #[test]
    fn test_verification_command_structure() {
        // Verify the VerificationCommand struct for pytorch is well-formed
        let candidates = vec!["python3".to_string()];
        let cmds = component_verification_commands("pytorch", &candidates);
        assert_eq!(cmds.len(), 1);
        let cmd = &cmds[0];
        assert!(!cmd.label.is_empty());
        assert!(!cmd.target_id.is_empty());
        assert!(!cmd.program.is_empty());
        assert_eq!(cmd.args.len(), 2, "Should have -c and the code");
        assert_eq!(cmd.args[0], "-c");
        assert!(!cmd.args[1].is_empty());
        assert!(cmd.modules.contains(&"torch".to_string()));
    }
}

#[cfg(test)]
mod component_status_tests;