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());
}
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::Maintenance {
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" => {
let version_exists = path_exists("/opt/rocm/.info/version")
|| path_exists("/opt/rocm/version")
|| path_exists("/usr/lib/rocm/.info/version");
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
};
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"]),
"flash-attn-triton" => {
python_any(python_candidates, &["flash_attn", "flash_attn_2"])
&& backend_marker_is("triton")
}
"flash-attn-ck" => {
python_any(python_candidates, &["flash_attn", "flash_attn_2"])
&& backend_marker_is("ck")
}
"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_candidates
.iter()
.any(|python| python_exec(python, onnxruntime_present_snippet()))
|| path_exists(home_path(&home, &["onnxruntime_build"]))
}
"bitsandbytes" => {
python_any(python_candidates, &["bitsandbytes"])
|| path_exists(home_path(&home, &["ml_stack", "bitsandbytes"]))
}
"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) {
if let Ok(output) = std::process::Command::new(path).arg("--showuse").output() {
if output.status.success() {
return true;
}
}
}
}
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"])
|| crate::installers::components::migraphx_multi::migraphx_driver_functional()
}
"pytorch-profiler" => python_any(python_candidates, &["torch"]),
"wandb" => python_any(python_candidates, &["wandb"]),
"fastvideo" => python_candidates.iter().any(|python| {
python_exec(
python,
crate::installers::components::fastvideo::fastvideo_verification_snippet(),
)
}),
"llama-cpp" => {
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" => {
let default_path = home_path(&home, &["ComfyUI"]);
path_exists(default_path.join(".git")) || path_exists(default_path.join("main.py"))
}
"textgen" => {
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,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)",
)],
"mpi4py" => vec![python_command(
"MPI4Py",
"mpi4py",
&["mpi4py"],
python_candidates,
"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)",
)],
"deepspeed" => vec![python_command(
"DeepSpeed",
"deepspeed",
&["deepspeed"],
python_candidates,
"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)",
)],
"ml-stack-core" => vec![python_command(
"ML Stack Core",
"ml-stack-core",
&["stans_ml_stack"],
python_candidates,
"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)",
)],
"flash-attn" => vec![flash_attention_verification_command(
"Flash Attention",
"flash-attn",
None,
python_candidates,
)],
"flash-attn-triton" => vec![flash_attention_verification_command(
"Flash Attention (Triton)",
"flash-attn-triton",
Some("triton"),
python_candidates,
)],
"flash-attn-ck" => vec![flash_attention_verification_command(
"Flash Attention (CK)",
"flash-attn-ck",
Some("ck"),
python_candidates,
)],
"megatron" => vec![python_command(
"Megatron-LM",
"megatron",
&["megatron"],
python_candidates,
"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)",
)],
"vllm" => vec![python_command(
"vLLM",
"vllm",
&["vllm"],
python_candidates,
vllm_runtime_check_snippet(),
)],
"aiter" => vec![python_command(
"AITER",
"aiter",
&["aiter"],
python_candidates,
"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)",
)],
"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,
migraphx_verify_snippet(),
)],
"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,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)",
)],
"fastvideo" => vec![python_command(
"FastVideo",
"fastvideo",
&["fastvideo"],
python_candidates,
crate::installers::components::fastvideo::fastvideo_verification_snippet(),
)],
"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",
],
)],
"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![{
let py = python_candidates
.first()
.cloned()
.or_else(|| env::var("MLSTACK_PYTHON_BIN").ok().filter(|s| !s.is_empty()))
.or_else(|| env::var("UV_PYTHON").ok().filter(|s| !s.is_empty()))
.unwrap_or_else(|| "python3".to_string());
shell_command(
"ComfyUI",
"comfyui",
"bash",
&[
"-c",
&format!(
"cd \"$HOME/ComfyUI\" 2>/dev/null && {py} -c 'import folder_paths' 2>/dev/null && echo 'comfyui functional' || exit 1"
),
],
)
}],
"textgen" => vec![{
let py = python_candidates
.first()
.cloned()
.or_else(|| env::var("MLSTACK_PYTHON_BIN").ok().filter(|s| !s.is_empty()))
.or_else(|| env::var("UV_PYTHON").ok().filter(|s| !s.is_empty()))
.unwrap_or_else(|| "python3".to_string());
shell_command(
"text-generation-webui",
"textgen",
"bash",
&[
"-c",
&format!(
"cd \"$HOME/text-generation-webui\" 2>/dev/null && {py} -c 'import server' 2>/dev/null && echo 'textgen functional' || 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,
"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)",
),
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![
flash_attention_verification_command(
"Flash Attention",
"flash-attn",
None,
python_candidates,
),
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,
migraphx_verify_snippet(),
),
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,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)",
),
python_command(
"FastVideo",
"fastvideo",
&["fastvideo"],
python_candidates,
crate::installers::components::fastvideo::fastvideo_verification_snippet(),
),
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)",
),
flash_attention_verification_command(
"Flash Attention (build)",
"flash-attn",
None,
python_candidates,
),
python_command(
"MIGraphX",
"migraphx",
&["migraphx"],
python_candidates,
migraphx_verify_snippet(),
),
]);
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(),
}
}
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_build: {cuda_version}')
# Tier 3: Runtime HIP/CUDA availability
try:
hip_available = torch.cuda.is_available()
# 'cuda' here is AVAILABILITY (CUDA-via-HIP on ROCm), NOT torch.version.cuda
# (which is None on a ROCm build — that's the CUDA *toolkit* version torch was
# built against, not whether CUDA-style code can run). is_available()==True
# means CUDA works through HIP, which is the user-facing answer.
print(f'cuda: {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://download.pytorch.org/whl/rocm7.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')
# Stage 4: exit code reflects FUNCTIONAL HIP availability on this ROCm stack.
# A CPU-only/CUDA torch build, or a ROCm build whose HIP runtime is
# unreachable, must FAIL verification — not silently pass. The diagnostics
# above guide the fix; downstream gating belongs in the orchestrator layer,
# not a lying exit code.
ok = (hip_version is not None) and hip_available
sys.exit(0 if ok else 1)"#
}
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 onnxruntime_present_snippet() -> &'static str {
"import sys, os, importlib.util; \
sys.path = [p for p in sys.path if p and os.path.abspath(p.rstrip('/')) != os.path.abspath(os.getcwd())]; \
s = importlib.util.find_spec('onnxruntime'); \
sys.exit(0 if (s is not None and s.origin is not None) 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 {
if let Ok(home) = env::var("MLSTACK_USER_HOME") {
return home;
}
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!("{}/.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),
];
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 backend_marker_is(name: &str) -> bool {
let home = resolve_component_user_home();
let marker = home_path(&home, &[".mlstack", "flash-attention", ".backend"]);
std::fs::read_to_string(&marker)
.ok()
.map(|s| s.trim() == name)
.unwrap_or(false)
}
fn migraphx_verify_snippet() -> &'static str {
r#"
import os, sys, subprocess
# (1) Drop CWD / '' entries so a repo-local migraphx stub can't shadow a real install.
sys.path = [p for p in sys.path if p and os.path.abspath(p.rstrip('/')) != os.path.abspath(os.getcwd())]
ok = False
ver = 'ok'
via = 'none'
try:
import migraphx
if hasattr(migraphx, 'parse_onnx'):
ok = True
via = 'python'
ver = getattr(migraphx, '__version__', 'ok')
else:
ver = getattr(migraphx, '__version__', 'stub')
except ImportError:
ver = 'absent'
# (2) Fall through to the C++ driver whenever python bindings are missing OR
# incomplete (no parse_onnx). On Arch/CachyOS the driver is the install target.
if not ok:
r = subprocess.run(['/opt/rocm/bin/migraphx-driver', '--version'], capture_output=True)
if r.returncode == 0:
ok = True
via = 'driver'
ver = 'driver'
print(f'MIGraphX {ver} parse_onnx={ok} via={via}')
sys.exit(0 if ok else 1)
"#
}
fn flash_attention_verification_command(
label: &str,
target_id: &str,
expected_backend: Option<&str>,
python_candidates: &[String],
) -> VerificationCommand {
let marker_block = match expected_backend {
Some(name) => format!(
"home=pathlib.Path(os.environ.get('HOME',''))\nmarker=home/'.mlstack'/'flash-attention'/'.backend'\nif not marker.exists() or marker.read_text().strip()!='{name}':\n print(f'backend marker mismatch (expected {name})');sys.exit(1)\n"
),
None => String::new(),
};
let code = format!(
"import importlib,sys,os,pathlib\nm=importlib.import_module('flash_attn')\nver=getattr(m,'__version__','unknown')\nhas_func=hasattr(m,'flash_attn_func') or hasattr(m,'flash_attn_qkvpacked_func')\n{marker_block}print(f'flash_attn {{ver}} func={{has_func}}')"
);
python_command(label, target_id, &["flash_attn"], python_candidates, &code)
}
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) {
let has_marker = contents.contains("Permanent ROCm Environment")
|| contents.contains("Permanent ROCm Env")
|| (contents.contains("ML Stack Environment File")
&& contents.contains("MLSTACK_PYTHON_BIN"));
if !has_marker {
return false;
}
let valid_arch_prefixes = ["gfx9", "gfx10", "gfx11", "gfx12"];
let mut has_gpu_arch_line = false;
for line in contents.lines() {
let trimmed = line.trim();
let assign = trimmed.strip_prefix("export ").unwrap_or(trimmed);
if assign.starts_with("GPU_ARCH=") {
has_gpu_arch_line = true;
if let Some(rest) = assign.split("GPU_ARCH=").nth(1) {
let arch = rest.split_whitespace().next().unwrap_or(rest);
if arch == "gfx000" {
return false;
}
if !valid_arch_prefixes
.iter()
.any(|prefix| arch.starts_with(prefix))
{
return false;
}
}
}
}
has_gpu_arch_line
} else {
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
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_fastvideo_verification_uses_shared_compiled_kernel_smoke_test() {
let candidates = vec!["python3".to_string()];
let commands = component_verification_commands("fastvideo", &candidates);
assert_eq!(commands.len(), 1);
let command = &commands[0];
assert_eq!(command.label, "FastVideo");
assert_eq!(command.target_id, "fastvideo");
assert_eq!(command.modules, vec!["fastvideo".to_string()]);
assert_eq!(
command.args,
vec![
"-c".to_string(),
crate::installers::components::fastvideo::fastvideo_verification_snippet()
.to_string(),
]
);
let snippet = &command.args[1];
for anchor in [
"from fastvideo_kernel import int8_quant",
"range(int(discovered.stdout.strip()))",
"int8_quant(source)",
"torch.cuda.synchronize()",
"raise SystemExit",
] {
assert!(
snippet.contains(anchor),
"missing verification anchor: {anchor}"
);
}
let enhanced = enhanced_verification_commands(&candidates);
let enhanced_fastvideo: Vec<_> = enhanced
.iter()
.filter(|candidate| candidate.target_id == "fastvideo")
.collect();
assert_eq!(enhanced_fastvideo.len(), 1);
assert_eq!(
enhanced_fastvideo[0].args[1],
crate::installers::components::fastvideo::fastvideo_verification_snippet()
);
assert_eq!(enhanced_fastvideo[0].modules, vec!["fastvideo".to_string()]);
}
#[test]
fn test_fastvideo_installation_status_requires_compiled_kernel_smoke_test() {
use std::os::unix::fs::PermissionsExt;
use std::time::{SystemTime, UNIX_EPOCH};
let unique = SystemTime::now()
.duration_since(UNIX_EPOCH)
.expect("system clock must be after Unix epoch")
.as_nanos();
let temp_dir = env::temp_dir().join(format!(
"rusty-stack-fastvideo-status-{}-{unique}",
std::process::id()
));
fs::create_dir_all(&temp_dir).expect("create temporary FastVideo status test directory");
let import_only = temp_dir.join("import-only-python");
fs::write(
&import_only,
"#!/bin/sh\ncase \"$2\" in\n *\"find_spec\"*\"fastvideo\"*) exit 0 ;;\n *) exit 1 ;;\nesac\n",
)
.expect("write import-only fake Python");
let mut permissions = fs::metadata(&import_only)
.expect("stat import-only fake Python")
.permissions();
permissions.set_mode(0o755);
fs::set_permissions(&import_only, permissions)
.expect("make import-only fake Python executable");
assert!(
!is_component_installed_by_id(
"fastvideo",
&[import_only.to_string_lossy().into_owned()]
),
"an import spec alone must not report FastVideo installed"
);
let functional = temp_dir.join("functional-python");
fs::write(
&functional,
"#!/bin/sh\ncase \"$2\" in\n *\"import fastvideo\"*\"from fastvideo_kernel import int8_quant\"*\"int8_quant(source)\"*) exit 0 ;;\n *) exit 1 ;;\nesac\n",
)
.expect("write functional fake Python");
let mut permissions = fs::metadata(&functional)
.expect("stat functional fake Python")
.permissions();
permissions.set_mode(0o755);
fs::set_permissions(&functional, permissions)
.expect("make functional fake Python executable");
assert!(
is_component_installed_by_id("fastvideo", &[functional.to_string_lossy().into_owned()]),
"the shared compiled-kernel smoke must report FastVideo installed"
);
fs::remove_dir_all(&temp_dir).expect("remove temporary FastVideo status test directory");
}
#[test]
fn test_pytorch_diagnostic_snippet_exits_strict() {
let snippet = pytorch_diagnostic_snippet();
assert!(
snippet.contains("(hip_version is not None) and hip_available"),
"Should gate the exit on functional HIP availability"
);
assert!(
snippet.contains("sys.exit(0 if ok else 1)"),
"Should exit non-zero when HIP is unavailable (fail loud)"
);
assert!(
!snippet.contains("\nsys.exit(0)\n"),
"Should NOT have an unconditional soft-pass sys.exit(0)"
);
}
#[test]
fn test_pytorch_diagnostic_snippet_checks_hip() {
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() {
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() {
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() {
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");
let empty = String::new();
let code = cmd.args.get(1).unwrap_or(&empty);
assert!(
code.contains("(hip_version is not None) and hip_available"),
"Should gate the exit on functional HIP availability (Stage 4)"
);
assert!(
code.contains("sys.exit(0 if ok else 1)"),
"Should exit non-zero when HIP is unavailable (Stage 4: fail loud)"
);
}
#[test]
fn test_pytorch_basic_verification_uses_diagnostic() {
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("(hip_version is not None) and hip_available"),
"Basic verification should gate exit on functional HIP (Stage 4)"
);
assert!(
code.contains("sys.exit(0 if ok else 1)"),
"Basic verification should fail loud when HIP unavailable (Stage 4)"
);
}
#[test]
fn test_pytorch_installation_check_only_imports() {
let candidates = vec!["python3".to_string()];
let _ = is_component_installed_by_id("pytorch", &candidates);
}
#[test]
fn test_verification_command_structure() {
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()));
}
#[test]
fn test_migraphx_snippet_excludes_cwd_and_falls_back_to_driver() {
let s = migraphx_verify_snippet();
assert!(
s.contains("getcwd()"),
"snippet must drop CWD from sys.path; got: {s}"
);
assert!(
s.contains("sys.path = [p for p in sys.path"),
"snippet must filter sys.path; got: {s}"
);
assert!(
s.contains("if not ok:"),
"snippet must fall back to the driver whenever python bindings are incomplete"
);
assert!(s.contains("migraphx-driver"));
assert!(s.contains("--version"));
assert!(s.contains("hasattr(migraphx, 'parse_onnx')"));
assert!(s.contains("sys.exit(0 if ok else 1)"));
}
#[test]
fn test_onnx_present_snippet_strips_cwd_and_requires_real_package() {
let s = onnxruntime_present_snippet();
assert!(
s.contains("getcwd()"),
"snippet must drop CWD from sys.path; got: {s}"
);
assert!(
s.contains("sys.path = [p for p in sys.path"),
"snippet must filter sys.path; got: {s}"
);
assert!(
s.contains("origin is not None"),
"snippet must require a real package (non-None origin), not a namespace shadow; got: {s}"
);
}
}
#[cfg(test)]
mod component_status_tests;