rusty_stack/benchmarks/

mod.rs

//! Performance benchmark infrastructure for Rusty-Stack.
//! Heavy lifting is delegated to a small embedded Python helper that runs
//! ROCm-enabled PyTorch kernels to exercise the real hardware.

use serde::{Deserialize, Serialize};
use serde_json;
use std::env;
use std::fs;
use std::path::{Path, PathBuf};
use std::process::Command;
use std::time::Instant;

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct BenchmarkResult {
    pub name: String,
    pub success: bool,
    pub execution_time_ms: u128,
    pub metrics: serde_json::Value,
    pub errors: Vec<String>,
}

fn helper_dir_candidates() -> Vec<PathBuf> {
    let mut dirs = Vec::new();
    if let Some(home) = dirs::home_dir() {
        dirs.push(home.join(".rusty-stack").join("tmp"));
    }
    dirs.push(PathBuf::from("/tmp").join("rusty-stack-tmp"));
    dirs
}

fn ensure_helper_script() -> Result<PathBuf, String> {
    // Single source (Tenet 3): inject the iGPU token list generated from the
    // Rust `crate::gpu` consts so the benchmark Python filter cannot drift from
    // the install-path filter.
    let body = PY_HELPER
        .trim_start()
        .replace("__INTEGRATED_TOKENS__", &integrated_tokens_py());
    let mut last_error = String::new();

    for dir in helper_dir_candidates() {
        if let Err(err) = fs::create_dir_all(&dir) {
            last_error = format!("failed to create helper dir {}: {}", dir.display(), err);
            continue;
        }
        let path = dir.join("rusty_bench.py");
        match fs::write(&path, &body) {
            Ok(_) => return Ok(path),
            Err(err) => {
                last_error = format!("failed to write helper script {}: {}", path.display(), err);
                continue;
            }
        }
    }

    Err(last_error)
}

/// Build the Python iGPU-token list literal from the Rust `crate::gpu` consts
/// (Tenet 3 single source — the benchmark Python filter derives from the same
/// list as the install path, so it cannot drift).
fn integrated_tokens_py() -> String {
    let mut tokens: Vec<String> = crate::gpu::INTEGRATED_NAME_PATTERNS
        .iter()
        .map(|t| t.to_lowercase())
        .collect();
    tokens.extend(
        crate::gpu::INTEGRATED_GFX_ARCHS
            .iter()
            .map(|t| t.to_lowercase()),
    );
    let mut s = String::from("[");
    for (i, t) in tokens.iter().enumerate() {
        if i > 0 {
            s.push(',');
        }
        let escaped = t.replace('\\', "\\\\").replace('"', "\\\"");
        s.push_str(&format!("\"{escaped}\""));
    }
    s.push(']');
    s
}

fn resolve_benchmark_python() -> String {
    for key in [
        "MLSTACK_BENCHMARK_PYTHON",
        "MLSTACK_PYTHON_BIN",
        "UV_PYTHON",
    ] {
        if let Ok(value) = env::var(key) {
            let value = value.trim();
            if !value.is_empty() {
                return value.to_string();
            }
        }
    }

    {
        let candidate = "/usr/local/bin/python3";
        let path = Path::new(candidate);
        if path.exists() {
            return candidate.to_string();
        }
    }

    for candidate in ["python3", "python"] {
        if Command::new(candidate).arg("--version").output().is_ok() {
            return candidate.to_string();
        }
    }

    "python3".to_string()
}

fn extract_helper_payload(stdout: &str) -> Option<serde_json::Value> {
    let marker = "---BENCHMARK_RESULTS_START---";
    let search = if let Some(pos) = stdout.rfind(marker) {
        &stdout[pos + marker.len()..]
    } else {
        stdout
    };

    if let Some(start) = search.find('{') {
        let mut depth = 0usize;
        let mut in_string = false;
        let mut escaped = false;
        for (idx, ch) in search[start..].char_indices() {
            if in_string {
                if escaped {
                    escaped = false;
                    continue;
                }
                match ch {
                    '\\' => escaped = true,
                    '"' => in_string = false,
                    _ => {}
                }
                continue;
            }

            match ch {
                '"' => in_string = true,
                '{' => depth += 1,
                '}' => {
                    if depth == 0 {
                        return None;
                    }
                    depth -= 1;
                    if depth == 0 {
                        let end = start + idx + 1;
                        return serde_json::from_str::<serde_json::Value>(&search[start..end]).ok();
                    }
                }
                _ => {}
            }
        }
    }

    None
}

fn run_python_benchmark(name: &str) -> BenchmarkResult {
    let start = Instant::now();
    let python_bin = resolve_benchmark_python();
    let helper = match ensure_helper_script() {
        Ok(path) => path,
        Err(err) => {
            return BenchmarkResult {
                name: name.to_string(),
                success: false,
                execution_time_ms: start.elapsed().as_millis(),
                metrics: serde_json::json!({}),
                errors: vec![format!("Unable to create benchmark helper script: {}", err)],
            };
        }
    };

    let mut command = Command::new(&python_bin);
    command.arg(helper).arg(name).arg("--json");
    if env::var_os("PYTHONHASHSEED").is_none() {
        command.env("PYTHONHASHSEED", "0");
    }
    let output = command.output();

    let parse_payload_result = |parsed: serde_json::Value| BenchmarkResult {
        name: parsed
            .get("name")
            .and_then(|v| v.as_str())
            .unwrap_or(name)
            .to_string(),
        success: parsed
            .get("success")
            .and_then(|v| v.as_bool())
            .unwrap_or(false),
        execution_time_ms: parsed
            .get("execution_time_ms")
            .and_then(|v| v.as_u64())
            .unwrap_or(start.elapsed().as_millis() as u64) as u128,
        metrics: parsed.get("metrics").cloned().unwrap_or_default(),
        errors: parsed
            .get("errors")
            .and_then(|v| v.as_array())
            .map(|arr| {
                arr.iter()
                    .filter_map(|e| e.as_str().map(|s| s.to_string()))
                    .collect()
            })
            .unwrap_or_default(),
    };

    match output {
        Ok(out) => {
            let stdout_str = String::from_utf8_lossy(&out.stdout);
            if let Some(parsed) = extract_helper_payload(&stdout_str) {
                let mut result = parse_payload_result(parsed);
                if !result.success {
                    let stderr = String::from_utf8_lossy(&out.stderr);
                    let stderr_trim = stderr.trim();
                    if !stderr_trim.is_empty() {
                        let stderr_tail = stderr_trim
                            .lines()
                            .rev()
                            .take(24)
                            .collect::<Vec<_>>()
                            .into_iter()
                            .rev()
                            .collect::<Vec<_>>()
                            .join(" | ");
                        if !stderr_tail.is_empty() {
                            result
                                .errors
                                .push(format!("Helper stderr tail: {}", stderr_tail));
                        }
                    }
                }
                return result;
            }

            if out.status.success() {
                let stderr = String::from_utf8_lossy(&out.stderr).trim().to_string();
                let parse_error = if stderr.is_empty() {
                    format!("Failed to parse helper JSON output using {}", python_bin)
                } else {
                    format!(
                        "Failed to parse helper JSON output using {} (stderr: {})",
                        python_bin, stderr
                    )
                };
                return BenchmarkResult {
                    name: name.to_string(),
                    success: false,
                    execution_time_ms: start.elapsed().as_millis(),
                    metrics: serde_json::json!({}),
                    errors: vec![parse_error],
                };
            }

            BenchmarkResult {
                name: name.to_string(),
                success: false,
                execution_time_ms: start.elapsed().as_millis(),
                metrics: serde_json::json!({}),
                errors: {
                    let stderr = String::from_utf8_lossy(&out.stderr).trim().to_string();
                    let stdout = String::from_utf8_lossy(&out.stdout).trim().to_string();
                    let detail = if !stderr.is_empty() {
                        stderr
                    } else if !stdout.is_empty() {
                        stdout
                    } else {
                        "no output".to_string()
                    };
                    vec![format!(
                        "Helper failed with status {} using {}: {}",
                        out.status, python_bin, detail
                    )]
                },
            }
        }
        Err(err) => BenchmarkResult {
            name: name.to_string(),
            success: false,
            execution_time_ms: start.elapsed().as_millis(),
            metrics: serde_json::json!({}),
            errors: vec![format!("Failed to invoke {}: {}", python_bin, err)],
        },
    }
}

pub fn run_gpu_capability_benchmark() -> BenchmarkResult {
    run_python_benchmark("gpu-info")
}
pub fn run_memory_bandwidth_benchmark() -> BenchmarkResult {
    run_python_benchmark("memory-bandwidth")
}
pub fn run_tensor_core_benchmark() -> BenchmarkResult {
    run_python_benchmark("tensor-core")
}
pub fn run_gemm_benchmark() -> BenchmarkResult {
    run_python_benchmark("gemm")
}
pub fn run_pytorch_benchmark() -> BenchmarkResult {
    run_python_benchmark("pytorch")
}
pub fn run_llama_cpp_benchmark() -> BenchmarkResult {
    run_python_benchmark("llama-cpp")
}
pub fn run_flash_attention_benchmark() -> BenchmarkResult {
    run_python_benchmark("flash-attention")
}
pub fn run_vllm_benchmark() -> BenchmarkResult {
    run_python_benchmark("vllm")
}
pub fn run_deepspeed_benchmark() -> BenchmarkResult {
    run_python_benchmark("deepspeed")
}
pub fn run_megatron_benchmark() -> BenchmarkResult {
    run_python_benchmark("megatron")
}
pub fn run_onnx_benchmark() -> BenchmarkResult {
    run_python_benchmark("onnx")
}

// ---------------------------------------------------------------------------
// Embedded Python helper
// ---------------------------------------------------------------------------
const PY_HELPER: &str = r#"
import argparse
import json
import os
import shutil
import subprocess
import sys
import time
from glob import glob

_GPU_RUNTIME_CACHE = None
_DEFAULT_TINY_SAFETENSORS_MODEL = "HuggingFaceTB/SmolLM2-135M-Instruct"


def _set_amd_gpu_id_table_env():
    if os.environ.get("AMDGPU_ASIC_ID_TABLE_PATH"):
        return
    candidate = "/usr/share/libdrm/amdgpu.ids"
    if os.path.isfile(candidate):
        os.environ["AMDGPU_ASIC_ID_TABLE_PATH"] = candidate
        os.environ.setdefault("AMDGPU_ASIC_ID_TABLE_PATHS", "/usr/share/libdrm")


def _env_or_default(name, default):
    value = os.environ.get(name)
    if value is None:
        return default
    value = value.strip()
    return value if value else default


def _resolve_gguf_model_path(raw):
    raw = (raw or "").strip()
    if not raw:
        return ""
    expanded = os.path.expanduser(raw)
    if os.path.isfile(expanded):
        return expanded
    if os.path.isdir(expanded):
        for entry in sorted(os.listdir(expanded)):
            if entry.lower().endswith(".gguf"):
                return os.path.join(expanded, entry)
    return ""


def _find_llama_cpp_binary():
    candidates = [
        os.path.expanduser("~/.mlstack/components/llama-cpp/bin/llama-bench"),
        shutil.which("llama-bench"),
    ]
    for candidate in candidates:
        if candidate and os.path.isfile(candidate) and os.access(candidate, os.X_OK):
            return candidate
    return ""


def _find_gguf_model():
    candidates = []
    candidates.extend(sorted(glob(os.path.expanduser("~/.mlstack/models/*.gguf"))))
    candidates.extend(sorted(glob(os.path.expanduser("~/.cache/huggingface/hub/*/snapshots/*/*.gguf"))))
    candidates.extend(sorted(glob(os.path.expanduser("~/.cache/**/*.gguf"), recursive=True)))
    for candidate in candidates:
        if os.path.isfile(candidate):
            return candidate
    return ""


def _detect_llama_cpp_gpus():
    gpu_names = []
    try:
        proc = subprocess.run(
            ["rocm-smi", "--showproductname"],
            check=False,
            capture_output=True,
            text=True,
        )
        output = (proc.stdout or "") + "\n" + (proc.stderr or "")
        for line in output.splitlines():
            line = line.strip()
            if not line:
                continue
            lowered = line.lower()
            if "gfx" in lowered or "amd" in lowered or "radeon" in lowered or "instinct" in lowered:
                gpu_names.append(line)
    except Exception:
        pass
    if gpu_names:
        return gpu_names

    try:
        proc = subprocess.run(
            "rocminfo | grep gfx",
            shell=True,
            check=False,
            capture_output=True,
            text=True,
        )
        output = (proc.stdout or "") + "\n" + (proc.stderr or "")
        for line in output.splitlines():
            line = line.strip()
            if "gfx" in line.lower():
                gpu_names.append(line)
    except Exception:
        pass
    return gpu_names


def _llama_cpp():
    start = time.perf_counter()
    bench = _find_llama_cpp_binary()
    if not bench:
        return False, {
            "name": "llama-cpp",
            "success": False,
            "execution_time_ms": 0,
            "metrics": {},
            "errors": ["llama-cpp not installed"],
        }, ["llama-cpp not installed"]

    model = _find_gguf_model()
    if not model:
        elapsed = int((time.perf_counter() - start) * 1000)
        return False, {
            "name": "llama-cpp",
            "success": False,
            "execution_time_ms": elapsed,
            "metrics": {},
            "errors": ["no GGUF model found"],
        }, ["no GGUF model found"]

    gpus = _detect_llama_cpp_gpus()
    if not gpus:
        elapsed = int((time.perf_counter() - start) * 1000)
        return False, {
            "name": "llama-cpp",
            "success": False,
            "execution_time_ms": elapsed,
            "metrics": {},
            "errors": ["no ROCm GPUs detected"],
        }, ["no ROCm GPUs detected"]

    metrics = {}
    errors = []
    for gpu_idx, _ in enumerate(gpus):
        env = os.environ.copy()
        env["ROCM_VISIBLE_DEVICES"] = str(gpu_idx)
        try:
            proc = subprocess.run(
                [bench, "-m", model, "-p", "512", "-p", "2048", "-p", "8192", "-n", "128", "-o", "json", "-r", "3"],
                check=False,
                capture_output=True,
                text=True,
                env=env,
            )
            raw = (proc.stdout or "").strip()
            parsed = json.loads(raw) if raw else []
            if isinstance(parsed, dict):
                parsed = [parsed]
            for entry in parsed:
                if not isinstance(entry, dict):
                    continue
                n_prompt = int(entry.get("n_prompt", 0) or 0)
                n_gen = int(entry.get("n_gen", 0) or 0)
                avg_ts = entry.get("avg_ts")
                if avg_ts is None:
                    avg_ts = entry.get("avg_tps")
                stddev_ts = entry.get("stddev_ts", 0.0)
                prefix = "prefill" if n_prompt > 0 and n_gen == 0 else "decode" if n_prompt == 0 and n_gen > 0 else "other"
                if prefix == "other" or avg_ts is None:
                    continue
                context = n_prompt if prefix == "prefill" else n_gen
                metrics[f"{prefix}_{context}_tps_gpu{gpu_idx}"] = float(avg_ts)
                metrics[f"{prefix}_{context}_stddev_tps_gpu{gpu_idx}"] = float(stddev_ts or 0.0)
        except Exception as exc:
            errors.append(str(exc))

    elapsed = int((time.perf_counter() - start) * 1000)
    success = bool(metrics) and not errors
    result = {
        "name": "llama-cpp",
        "success": success,
        "execution_time_ms": elapsed,
        "metrics": metrics,
        "errors": errors,
    }
    print("---BENCHMARK_RESULTS_START---")
    print(json.dumps(result, indent=2, sort_keys=True))
    return success, result, errors


def _resolve_vllm_model_candidates():
    safetensors_model = _env_or_default(
        "MLSTACK_BENCH_VLLM_SAFETENSORS_MODEL",
        _DEFAULT_TINY_SAFETENSORS_MODEL,
    )
    gguf_hint = os.environ.get("MLSTACK_BENCH_VLLM_GGUF_MODEL_PATH", "")
    gguf_tokenizer = _env_or_default("MLSTACK_BENCH_VLLM_GGUF_TOKENIZER", safetensors_model)
    gguf_path = _resolve_gguf_model_path(gguf_hint)

    candidates = []
    if gguf_path:
        candidates.append({
            "format": "gguf",
            "model": gguf_path,
            "tokenizer": gguf_tokenizer,
        })
    candidates.append({
        "format": "safetensors",
        "model": safetensors_model,
    })
    return candidates


def _probe_gpu_runtime():
    global _GPU_RUNTIME_CACHE
    if _GPU_RUNTIME_CACHE is not None:
        return _GPU_RUNTIME_CACHE

    _set_amd_gpu_id_table_env()
    torch = _load_torch()
    if isinstance(torch, tuple):
        _, errors = torch
        reason = "; ".join(errors) if errors else "unable to import torch"
        _GPU_RUNTIME_CACHE = (False, reason, {"available": False})
        return _GPU_RUNTIME_CACHE

    if not torch.cuda.is_available():
        _GPU_RUNTIME_CACHE = (False, "torch.cuda is not available", {"available": False})
        return _GPU_RUNTIME_CACHE

    arch = "unknown"
    try:
        props = torch.cuda.get_device_properties(torch.device("cuda:0"))
        arch = getattr(props, "gcnArchName", "unknown")
    except Exception:
        pass

    probe_code = (
        "import torch;"
        "x=torch.ones(1, device='cuda', dtype=torch.float32);"
        "y=x+1;"
        "torch.cuda.synchronize();"
        "print(float(y.item()))"
    )
    try:
        probe = subprocess.run(
            [sys.executable, "-c", probe_code],
            check=False,
            capture_output=True,
            text=True,
            timeout=20,
        )
        if probe.returncode == 0:
            _GPU_RUNTIME_CACHE = (True, "", {"available": True, "gcn_arch": arch})
            return _GPU_RUNTIME_CACHE

        detail = (probe.stderr or probe.stdout or "").strip()
        if not detail:
            detail = f"exit code {probe.returncode}"
        if probe.returncode < 0:
            detail = f"signal {-probe.returncode}: {detail}"
        reason = f"GPU runtime probe failed on {arch}: {detail}"
        _GPU_RUNTIME_CACHE = (False, reason, {"available": True, "gcn_arch": arch})
        return _GPU_RUNTIME_CACHE
    except subprocess.TimeoutExpired:
        reason = f"GPU runtime probe timed out on {arch}"
        _GPU_RUNTIME_CACHE = (False, reason, {"available": True, "gcn_arch": arch})
        return _GPU_RUNTIME_CACHE
    except Exception as exc:
        reason = f"GPU runtime probe error on {arch}: {exc}"
        _GPU_RUNTIME_CACHE = (False, reason, {"available": True, "gcn_arch": arch})
        return _GPU_RUNTIME_CACHE


def _degraded_metrics(name, reason, extra=None):
    payload = {
        "mode": "degraded",
        "component": name,
        "gpu_runtime_stable": False,
        "reason": reason,
    }
    if isinstance(extra, dict):
        payload.update(extra)
    return payload


def _load_torch():
    try:
        import torch  # noqa: F401
        return torch
    except Exception as exc:  # pragma: no cover
        return None, [f"Unable to import torch: {exc}"]


def _parse_visible_gpu_indices():
    for key in ("HIP_VISIBLE_DEVICES", "ROCR_VISIBLE_DEVICES", "CUDA_VISIBLE_DEVICES"):
        raw = os.environ.get(key, "").strip()
        if not raw:
            continue

        values = []
        valid = True
        for token in raw.split(","):
            token = token.strip()
            if not token:
                continue
            if not token.isdigit():
                valid = False
                break
            values.append(int(token))

        if valid and values:
            return values
    return []


def _coerce_smi_value(entry, key):
    raw = entry.get(key)
    if isinstance(raw, dict):
        value = raw.get("value")
        if value is None:
            return None
        return value
    return raw


def _discrete_hint_from_name(name):
    if not name:
        return None

    lowered = str(name).lower()
    if any(token in lowered for token in ("radeon rx", " rx ", "radeon pro", "instinct", "firepro")):
        return True
    if "ryzen" in lowered:
        return False
    if "radeon graphics" in lowered and "rx" not in lowered and "pro" not in lowered:
        return False
    if "integrated" in lowered or "igpu" in lowered or "apu" in lowered:
        return False
    return None


def _is_integrated_name(name):
    hint = _discrete_hint_from_name(name)
    if hint is not None:
        return not hint
    n = str(name or "")
    lowered = n.lower()
    up = n.upper()
    # Token set GENERATED from crate::gpu::INTEGRATED_NAME_PATTERNS at write
    # time (single source — no hand-maintained copy that can drift from Rust).
    for token in __INTEGRATED_TOKENS__:
        if token in lowered:
            return True
    # "Ryzen" without "RX" => APU/iGPU (discrete cards never contain "Ryzen").
    if "RYZEN" in up and "RX" not in up:
        return True
    return False


def _gpu_info():
    torch = _load_torch()
    if isinstance(torch, tuple):
        _, errors = torch
        return False, {}, errors
    if not torch.cuda.is_available():
        return False, {}, ["torch.cuda is not available"]

    visible_indices = _parse_visible_gpu_indices()
    smi_available = shutil.which("rocm-smi") is not None
    devices = []
    for idx in range(torch.cuda.device_count()):
        global_idx = visible_indices[idx] if idx < len(visible_indices) else idx
        device = torch.device(f"cuda:{idx}")
        props = torch.cuda.get_device_properties(device)
        info = {
            "index": global_idx,
            "local_index": idx,
            "gpu_model": getattr(props, "name", "Unknown"),
            "gcn_arch": getattr(props, "gcnArchName", "unknown"),
            "vram_gb": round(getattr(props, "total_memory", 0) / 1e9, 2),
            "compute_units": getattr(props, "multi_processor_count", 0),
            "max_clock_mhz": getattr(props, "max_frequency", 0) / 1e6,
            "tensor_cores": bool(getattr(props, "gcnArchName", "").startswith("gfx11")),
        }

        smi_discrete_hint = None
        if smi_available:
            try:
                smi = subprocess.run(
                    [
                        "rocm-smi",
                        "--showproductname",
                        "--showtemp",
                        "--showpower",
                        "--showuse",
                        "--showclocks",
                        "--json",
                        "-d",
                        str(global_idx),
                    ],
                    check=False,
                    capture_output=True,
                    text=True,
                )
                if smi.returncode == 0 and smi.stdout.strip().startswith("{"):
                    data = json.loads(smi.stdout)
                    first_key = next(iter(data.keys()))
                    entry = data.get(first_key, {})

                    smi_name = (
                        _coerce_smi_value(entry, "Card Series")
                        or _coerce_smi_value(entry, "Card series")
                        or _coerce_smi_value(entry, "Card Model")
                        or _coerce_smi_value(entry, "Card model")
                    )
                    if smi_name:
                        info["gpu_model"] = str(smi_name).strip()
                        smi_discrete_hint = _discrete_hint_from_name(smi_name)

                    temp = entry.get("Temperature (Sensor die)", {}).get("value")
                    power = entry.get("Average Graphics Package Power", {}).get("value")
                    if temp is not None:
                        info["temperature_c"] = float(temp)
                    if power is not None:
                        info["power_watts"] = float(power)
                    
                    # capture additional metrics requested
                    util = entry.get("GPU use (%)", {}).get("value")
                    if util is not None:
                        info["utilization_percent"] = float(util)
                    
                    mem_util = entry.get("GPU memory use (%)", {}).get("value")
                    if mem_util is not None:
                        info["memory_percent"] = float(mem_util)
                    
                    sclk = entry.get("GFX Clock (MHz)", {}).get("value")
                    if sclk is not None:
                        info["sclk_mhz"] = int(float(sclk))
                    
                    mclk = entry.get("Memory Clock (MHz)", {}).get("value")
                    if mclk is not None:
                        info["mclk_mhz"] = int(float(mclk))
            except Exception:
                pass

        if smi_discrete_hint is False:
            continue
        if smi_discrete_hint is None and _is_integrated_name(info.get("gpu_model")):
            continue
        devices.append(info)

    probe_ok, probe_reason, probe_meta = _probe_gpu_runtime()
    metrics = {"gpus": devices}
    metrics.update({"gpu_runtime_stable": probe_ok})
    if isinstance(probe_meta, dict):
        metrics.update({f"probe_{k}": v for k, v in probe_meta.items()})
    if probe_ok:
        return True, metrics, []
    return True, metrics, [probe_reason]


def _time_fn(fn, warmup=1, repeat=5):
    import torch

    torch.cuda.synchronize()
    for _ in range(warmup):
        fn()
    times = []
    for _ in range(repeat):
        torch.cuda.synchronize()
        start = time.perf_counter()
        fn()
        torch.cuda.synchronize()
        times.append(time.perf_counter() - start)
    return sum(times) / len(times)


def _memory_bandwidth():
    torch = _load_torch()
    if isinstance(torch, tuple):
        _, errors = torch
        return False, {}, errors
    if not torch.cuda.is_available():
        return False, {}, ["torch.cuda not available"]
    probe_ok, probe_reason, probe_meta = _probe_gpu_runtime()
    if not probe_ok:
        extra = {
            "hbm_peak_gb_s": 0.0,
            "system_peak_gb_s": 0.0,
            "hbm_ratio": 0.0,
            "hbm_samples_gbps": [],
            "system_samples_gbps": [],
        }
        if isinstance(probe_meta, dict):
            extra.update(probe_meta)
        return True, _degraded_metrics("memory-bandwidth", probe_reason, extra), [probe_reason]

    device = torch.device("cuda:0")
    sizes_mb = [64, 128, 256, 512]
    hbm_samples = []
    system_samples = []

    for size_mb in sizes_mb:
        numel = size_mb * 1024 * 1024 // 4
        a = torch.empty(numel, device=device, dtype=torch.float32)
        def op():
            a.add_(1.0)
        t = _time_fn(op, warmup=1, repeat=4)
        bytes_moved = a.numel() * a.element_size() * 2  # read + write
        hbm_samples.append(bytes_moved / t / 1e9)

        host = torch.empty_like(a.cpu())
        def h2d():
            _ = host.to(device, non_blocking=True)
        t_h2d = _time_fn(h2d, warmup=1, repeat=3)
        system_samples.append(bytes_moved / t_h2d / 1e9 if t_h2d > 0 else 0)

    hbm_peak = max(hbm_samples) if hbm_samples else 0
    sys_peak = max(system_samples) if system_samples else 0
    hbm_ratio = hbm_peak / sys_peak if sys_peak > 0 else 0

    return True, {
        "hbm_peak_gb_s": round(hbm_peak, 2),
        "system_peak_gb_s": round(sys_peak, 2),
        "hbm_ratio": round(hbm_ratio, 2),
        "hbm_samples_gbps": [round(x, 2) for x in hbm_samples],
        "system_samples_gbps": [round(x, 2) for x in system_samples],
    }, []


def _tensor_core():
    torch = _load_torch()
    if isinstance(torch, tuple):
        _, errors = torch
        return False, {}, errors
    if not torch.cuda.is_available():
        return False, {}, ["torch.cuda not available"]
    probe_ok, probe_reason, probe_meta = _probe_gpu_runtime()
    if not probe_ok:
        extra = {
            "fp16_tflops": 0.0,
            "bf16_tflops": 0.0,
            "tf32_tflops": 0.0,
            "fp32_tflops": 0.0,
            "fp16_samples_tflops": [],
        }
        if isinstance(probe_meta, dict):
            extra.update(probe_meta)
        return True, _degraded_metrics("tensor-core", probe_reason, extra), [probe_reason]

    device = torch.device("cuda:0")
    sizes = [512, 1024, 2048]
    
    results = {}
    
    # FP16
    fp16_samples = []
    for n in sizes:
        a = torch.randn((n, n), device=device, dtype=torch.float16)
        b = torch.randn((n, n), device=device, dtype=torch.float16)
        t = _time_fn(lambda: torch.matmul(a, b), warmup=2, repeat=3)
        flops = 2 * n * n * n
        fp16_samples.append(flops / t / 1e12)
    results["fp16_tflops"] = round(max(fp16_samples), 2)
    results["fp16_samples"] = [round(x, 2) for x in fp16_samples]

    # BF16
    bf16_samples = []
    try:
        for n in sizes:
            a = torch.randn((n, n), device=device, dtype=torch.bfloat16)
            b = torch.randn((n, n), device=device, dtype=torch.bfloat16)
            t = _time_fn(lambda: torch.matmul(a, b), warmup=2, repeat=3)
            flops = 2 * n * n * n
            bf16_samples.append(flops / t / 1e12)
        results["bf16_tflops"] = round(max(bf16_samples), 2)
    except Exception:
        results["bf16_tflops"] = 0.0

    # FP32
    fp32_samples = []
    for n in sizes:
        a = torch.randn((n, n), device=device, dtype=torch.float32)
        b = torch.randn((n, n), device=device, dtype=torch.float32)
        t = _time_fn(lambda: torch.matmul(a, b), warmup=2, repeat=3)
        flops = 2 * n * n * n
        fp32_samples.append(flops / t / 1e12)
    results["fp32_tflops"] = round(max(fp32_samples), 2)
    
    # TF32 is NVIDIA specific, on AMD we'll report 0 or skip
    results["tf32_tflops"] = 0.0

    return True, {
        "fp16_tflops": results["fp16_tflops"],
        "bf16_tflops": results["bf16_tflops"],
        "tf32_tflops": results["tf32_tflops"],
        "fp32_tflops": results["fp32_tflops"],
        "fp16_samples_tflops": results["fp16_samples"],
    }, []


def _gemm():
    torch = _load_torch()
    if isinstance(torch, tuple):
        _, errors = torch
        return False, {}, errors
    if not torch.cuda.is_available():
        return False, {}, ["torch.cuda not available"]
    probe_ok, probe_reason, probe_meta = _probe_gpu_runtime()
    if not probe_ok:
        extra = {
            "fp16_peak_gflops": 0.0,
            "fp16_samples_gflops": [],
        }
        if isinstance(probe_meta, dict):
            extra.update(probe_meta)
        return True, _degraded_metrics("gemm", probe_reason, extra), [probe_reason]

    device = torch.device("cuda:0")
    shapes = [(1024, 1024, 1024), (1536, 1536, 1536), (2048, 2048, 2048)]
    samples = []
    for m, k, n in shapes:
        a = torch.randn((m, k), device=device, dtype=torch.float16)
        b = torch.randn((k, n), device=device, dtype=torch.float16)
        def op():
            return torch.matmul(a, b)
        t = _time_fn(op, warmup=1, repeat=3)
        flops = 2 * m * n * k
        samples.append(flops / t / 1e9)
    peak = max(samples) if samples else 0
    return True, {
        "fp16_peak_gflops": round(peak, 2),
        "fp16_samples_gflops": [round(x, 2) for x in samples],
    }, []


def _pytorch():
    torch = _load_torch()
    if isinstance(torch, tuple):
        _, errors = torch
        return False, {}, errors
    if not torch.cuda.is_available():
        return False, {}, ["torch.cuda not available"]
    probe_ok, probe_reason, probe_meta = _probe_gpu_runtime()
    if not probe_ok:
        extra = {
            "gemm_gflops": 0.0,
            "convolution_gflops": 0.0,
            "gemm_samples_gflops": [],
            "conv_samples_gflops": [],
            "autograd_overhead_percent": 0.0,
        }
        if isinstance(probe_meta, dict):
            extra.update(probe_meta)
        return True, _degraded_metrics("pytorch", probe_reason, extra), [probe_reason]

    device = torch.device("cuda:0")
    m = k = n = 1024
    a = torch.randn((m, k), device=device, dtype=torch.float32, requires_grad=True)
    b = torch.randn((k, n), device=device, dtype=torch.float32, requires_grad=True)
    gemm_samples = []
    for _ in range(3):
        gemm_t = _time_fn(lambda: torch.matmul(a, b), warmup=1, repeat=1)
        gemm_samples.append(2 * m * n * k / gemm_t / 1e9)
    gemm_peak = max(gemm_samples) if gemm_samples else 0

    conv = torch.nn.Conv2d(64, 64, kernel_size=3, padding=1).to(device)
    x = torch.randn((32, 64, 64, 64), device=device)
    conv_samples = []
    for _ in range(3):
        conv_t = _time_fn(lambda: conv(x), warmup=1, repeat=1)
        conv_ops = 2 * 32 * 64 * 64 * 64 * 64 * 3 * 3
        conv_samples.append(conv_ops / conv_t / 1e9)
    conv_peak = max(conv_samples) if conv_samples else 0

    y = a @ b
    y.sum().backward()
    torch.cuda.synchronize()
    start = time.perf_counter()
    y = a @ b
    loss = y.sum()
    loss.backward()
    torch.cuda.synchronize()
    full = time.perf_counter() - start

    start = time.perf_counter()
    _ = a @ b
    torch.cuda.synchronize()
    fwd = time.perf_counter() - start
    overhead = (full - fwd) / full * 100 if full > 0 else 0

    return True, {
        "gemm_gflops": round(gemm_peak, 2),
        "convolution_gflops": round(conv_peak, 2),
        "gemm_samples_gflops": [round(x, 2) for x in gemm_samples],
        "conv_samples_gflops": [round(x, 2) for x in conv_samples],
        "autograd_overhead_percent": round(overhead, 2),
    }, []


def _flash_attention():
    torch = _load_torch()
    if isinstance(torch, tuple):
        _, errors = torch
        return False, {}, errors
    if not torch.cuda.is_available():
        return False, {}, ["torch.cuda not available"]
    probe_ok, probe_reason, probe_meta = _probe_gpu_runtime()
    if not probe_ok:
        extra = {
            "standard_attention_speed": 0.0,
            "flash_attention_speed": 0.0,
            "standard_samples_tok_s": [],
            "flash_samples_tok_s": [],
            "speedup": 0.0,
            "memory_savings_gb": 0.0,
        }
        if isinstance(probe_meta, dict):
            extra.update(probe_meta)
        return True, _degraded_metrics("flash-attention", probe_reason, extra), [probe_reason]

    try:
        from torch.backends.cuda import sdp_kernel
    except Exception as exc:  # pragma: no cover
        return False, {}, [f"flash attention control unavailable: {exc}"]

    device = torch.device("cuda:0")
    seqs = [128, 256, 512, 1024]
    bsz, heads, dim = 8, 8, 64
    flash_samples = []
    std_samples = []
    mem_std = 0
    mem_flash = 0

    for seqlen in seqs:
        q = torch.randn((bsz, heads, seqlen, dim), device=device, dtype=torch.float16)
        k = torch.randn_like(q)
        v = torch.randn_like(q)

        def run_flash():
            with sdp_kernel(enable_flash=True, enable_math=True, enable_mem_efficient=True):
                return torch.nn.functional.scaled_dot_product_attention(q, k, v, is_causal=False)

        def run_standard():
            with sdp_kernel(enable_flash=False, enable_math=True, enable_mem_efficient=False):
                return torch.nn.functional.scaled_dot_product_attention(q, k, v, is_causal=False)

        torch.cuda.reset_peak_memory_stats()
        _ = run_standard()
        mem_std = max(mem_std, torch.cuda.max_memory_allocated() / 1e9)

        torch.cuda.reset_peak_memory_stats()
        _ = run_flash()
        mem_flash = max(mem_flash, torch.cuda.max_memory_allocated() / 1e9)

        flash_t = _time_fn(run_flash, warmup=1, repeat=2)
        std_t = _time_fn(run_standard, warmup=1, repeat=2)

        tokens = bsz * seqlen
        flash_samples.append(tokens / flash_t)
        std_samples.append(tokens / std_t)

    flash_peak = max(flash_samples) if flash_samples else 0
    std_peak = max(std_samples) if std_samples else 0

    return True, {
        "standard_attention_speed": round(std_peak, 2),
        "flash_attention_speed": round(flash_peak, 2),
        "standard_samples_tok_s": [round(x, 2) for x in std_samples],
        "flash_samples_tok_s": [round(x, 2) for x in flash_samples],
        "speedup": round(flash_peak / std_peak, 3) if std_peak > 0 else 0,
        "memory_savings_gb": round(max(0, mem_std - mem_flash), 4),
    }, []


def _vllm():
    # Some shell setups can propagate an empty/whitespace-only target device.
    # vLLM rejects empty device strings, so normalize aggressively before import.
    target_device = _env_or_default("VLLM_TARGET_DEVICE", "rocm")
    if target_device not in {"rocm", "cuda", "cpu"}:
        target_device = "rocm"
    os.environ["VLLM_TARGET_DEVICE"] = target_device

    def _ensure_writable_triton_cache_env():
        import tempfile

        def _writable_dir(path):
            if not path:
                return False
            try:
                os.makedirs(path, exist_ok=True)
                probe = os.path.join(path, f".mlstack_probe_{os.getpid()}")
                with open(probe, "w", encoding="utf-8") as fp:
                    fp.write("ok")
                os.remove(probe)
                return True
            except Exception:
                return False

        triton_cache = (os.environ.get("TRITON_CACHE_DIR") or "").strip()
        if triton_cache and _writable_dir(triton_cache):
            triton_home = (os.environ.get("TRITON_HOME") or os.path.dirname(triton_cache)).strip()
            os.environ.setdefault("TRITON_HOME", triton_home or os.path.expanduser("~/.cache/mlstack/triton"))
            os.environ.setdefault("TRITON_DUMP_DIR", os.path.join(os.environ["TRITON_HOME"], "dump"))
            os.environ.setdefault("TRITON_OVERRIDE_DIR", os.path.join(os.environ["TRITON_HOME"], "override"))
            _writable_dir(os.environ.get("TRITON_DUMP_DIR", ""))
            _writable_dir(os.environ.get("TRITON_OVERRIDE_DIR", ""))
            return

        home = os.path.expanduser("~")
        candidates = [
            os.environ.get("MLSTACK_TRITON_HOME", "").strip(),
            os.environ.get("TRITON_HOME", "").strip(),
            os.path.join(home, ".cache", "mlstack", "triton"),
            os.path.join(tempfile.gettempdir(), f"mlstack-triton-{os.getuid()}"),
        ]
        for root in candidates:
            if not root:
                continue
            cache_dir = os.path.join(root, "cache")
            dump_dir = os.path.join(root, "dump")
            override_dir = os.path.join(root, "override")
            if _writable_dir(cache_dir) and _writable_dir(dump_dir) and _writable_dir(override_dir):
                os.environ["MLSTACK_TRITON_HOME"] = root
                os.environ["TRITON_HOME"] = root
                os.environ["TRITON_CACHE_DIR"] = cache_dir
                os.environ["TRITON_DUMP_DIR"] = dump_dir
                os.environ["TRITON_OVERRIDE_DIR"] = override_dir
                return

    _ensure_writable_triton_cache_env()

    def _normalize_visible_devices():
        visible = os.environ.get("HIP_VISIBLE_DEVICES") or os.environ.get("CUDA_VISIBLE_DEVICES") or ""
        candidates = []
        seen = set()
        for token in visible.split(","):
            item = token.strip()
            if not item.isdigit() or item in seen:
                continue
            seen.add(item)
            candidates.append(item)
        if not candidates:
            candidates = ["0"]

        normalized = ",".join(candidates)
        primary = candidates[0]
        os.environ["HIP_VISIBLE_DEVICES"] = normalized
        os.environ["CUDA_VISIBLE_DEVICES"] = normalized
        if not str(os.environ.get("PYTORCH_ROCM_DEVICE", "")).strip():
            os.environ["PYTORCH_ROCM_DEVICE"] = primary
        return normalized, primary

    # Preserve the full discrete-GPU set. Do not collapse to a single GPU.
    visible_devices, primary_visible = _normalize_visible_devices()
    if target_device == "rocm":
        os.environ.setdefault("VLLM_WORKER_MULTIPROC_METHOD", "spawn")

    candidates = _resolve_vllm_model_candidates()
    candidate_names = [f"{c.get('format')}:{c.get('model')}" for c in candidates]
    torch = _load_torch()
    if isinstance(torch, tuple):
        _, errors = torch
        reason = "; ".join(errors) if errors else "unable to import torch"
        return True, _degraded_metrics("vllm", reason, {
            "model": "vLLM (Unavailable)",
            "model_format": "unavailable",
            "throughput_tokens_per_sec": 0.0,
            "latency_ms": 0.0,
            "throughput_samples": [],
            "target_device": target_device,
            "candidate_models": candidate_names,
        }), [reason]
    runtime_probe_hint = ""
    if target_device in {"rocm", "cuda"}:
        try:
            if not torch.cuda.is_available():
                runtime_probe_hint = "ROCm torch reports no available HIP GPUs in benchmark environment"
            else:
                torch.cuda.get_device_properties(0)
        except Exception as exc:
            runtime_probe_hint = f"ROCm torch GPU initialization failed before vLLM load: {exc}"
    def _patch_amdsmi_for_vllm_arch_probe():
        try:
            import sys
            import types

            if os.environ.get("MLSTACK_VLLM_DISABLE_AMDSMI_SHIM", "").strip() in {"1", "true", "TRUE"}:
                return

            arch_hint = (
                os.environ.get("GPU_ARCH")
                or os.environ.get("PYTORCH_ROCM_ARCH")
                or os.environ.get("HSA_OVERRIDE_GFX_VERSION")
                or ""
            ).strip()
            if not arch_hint:
                return
            if not arch_hint.startswith("gfx") and "." in arch_hint:
                parts = [p for p in arch_hint.split(".") if p]
                if len(parts) >= 2:
                    arch_hint = f"gfx{parts[0]}{parts[1]}"

            shim = types.ModuleType("amdsmi")

            class AmdSmiException(Exception):
                pass

            def amdsmi_init():
                return None

            def amdsmi_shut_down():
                return None

            def amdsmi_get_processor_handles():
                return [0]

            def amdsmi_get_gpu_asic_info(_handle):
                return {"target_graphics_version": arch_hint}

            def amdsmi_topo_get_link_type(*_args, **_kwargs):
                return (0, 0)

            shim.AmdSmiException = AmdSmiException
            shim.amdsmi_init = amdsmi_init
            shim.amdsmi_shut_down = amdsmi_shut_down
            shim.amdsmi_get_processor_handles = amdsmi_get_processor_handles
            shim.amdsmi_get_gpu_asic_info = amdsmi_get_gpu_asic_info
            shim.amdsmi_topo_get_link_type = amdsmi_topo_get_link_type

            sys.modules["amdsmi"] = shim
        except Exception:
            # Best-effort compatibility shim.
            pass

    _patch_amdsmi_for_vllm_arch_probe()
    try:
        import vllm
        from vllm import LLM, SamplingParams
        import logging
        import importlib
        # Suppress vLLM and related logging to avoid breaking JSON parsing
        logging.getLogger("vllm").setLevel(logging.ERROR)
    except Exception as exc:
        err = f"vLLM not available: {exc}"
        return True, _degraded_metrics("vllm", err, {
            "model": "vLLM (Unavailable)",
            "model_format": "unavailable",
            "throughput_tokens_per_sec": 0.0,
            "latency_ms": 0.0,
            "throughput_samples": [],
            "target_device": target_device,
            "candidate_models": candidate_names,
        }), [err]
    abi_errors = []
    for ext_name in ("vllm._C", "vllm._rocm_C"):
        try:
            importlib.import_module(ext_name)
        except Exception as exc:
            abi_errors.append(f"{ext_name} import failed: {exc}")
    if len(abi_errors) == 2:
        reason = "vLLM ROCm native extensions failed to load; " + "; ".join(abi_errors)
        return True, _degraded_metrics("vllm", reason, {
            "model": "vLLM (Unavailable)",
            "model_format": "unavailable",
            "throughput_tokens_per_sec": 0.0,
            "latency_ms": 0.0,
            "throughput_samples": [],
            "target_device": target_device,
            "candidate_models": candidate_names,
        }), [reason]
    platform_device_type = ""
    platform_class = "unknown"
    try:
        from vllm.platforms import current_platform
        platform_device_type = str(getattr(current_platform, "device_type", "") or "").strip()
        platform_class = current_platform.__class__.__name__
    except Exception:
        pass
    if not platform_device_type:
        platform_reason = (
            "vLLM platform detection returned an empty device_type (UnspecifiedPlatform). "
            "This commonly means amdsmi is missing, so ROCm platform detection failed. "
            "Install amdsmi in the benchmark/runtime Python environment and rerun."
        )
        return True, _degraded_metrics("vllm", platform_reason, {
            "model": "vLLM (Detected)",
            "model_format": "unavailable",
            "throughput_tokens_per_sec": 0.0,
            "latency_ms": 0.0,
            "throughput_samples": [],
            "target_device": target_device,
            "candidate_models": candidate_names,
            "vllm_platform_class": platform_class,
            "vllm_platform_device_type": platform_device_type,
        }), [platform_reason]
    probe_ok, probe_reason, probe_meta = _probe_gpu_runtime()
    if not probe_ok:
        runtime_probe_hint = probe_reason if not runtime_probe_hint else f"{runtime_probe_hint}; {probe_reason}"

    prompts = [
        "The future of AI is",
        "ROCm performance on AMD is",
        "vLLM is a fast",
    ]
    max_tokens = 24
    attempt_errors = []

    def _run_vllm_attempt_subprocess(llm_kwargs, prompts, max_tokens, env_overrides=None):
        payload = {
            "llm_kwargs": llm_kwargs,
            "prompts": prompts,
            "max_tokens": int(max_tokens),
        }
        runner = r"""
import json
import logging
import os
import sys
import time
import traceback

def _emit(obj):
    print(json.dumps(obj))

try:
    payload = json.loads(sys.stdin.read())
    llm_kwargs = payload.get("llm_kwargs", {})
    prompts = payload.get("prompts", [])
    max_tokens = int(payload.get("max_tokens", 24))
    try:
        import types
        if os.environ.get("MLSTACK_VLLM_DISABLE_AMDSMI_SHIM", "").strip() not in {"1", "true", "TRUE"}:
            arch_hint = (
                os.environ.get("GPU_ARCH")
                or os.environ.get("PYTORCH_ROCM_ARCH")
                or os.environ.get("HSA_OVERRIDE_GFX_VERSION")
                or ""
            ).strip()
            if arch_hint and not arch_hint.startswith("gfx") and "." in arch_hint:
                parts = [p for p in arch_hint.split(".") if p]
                if len(parts) >= 2:
                    arch_hint = f"gfx{parts[0]}{parts[1]}"
            if arch_hint:
                shim = types.ModuleType("amdsmi")
                class AmdSmiException(Exception):
                    pass
                def amdsmi_init():
                    return None
                def amdsmi_shut_down():
                    return None
                def amdsmi_get_processor_handles():
                    return [0]
                def amdsmi_get_gpu_asic_info(_handle):
                    return {"target_graphics_version": arch_hint}
                def amdsmi_topo_get_link_type(*_args, **_kwargs):
                    return (0, 0)
                shim.AmdSmiException = AmdSmiException
                shim.amdsmi_init = amdsmi_init
                shim.amdsmi_shut_down = amdsmi_shut_down
                shim.amdsmi_get_processor_handles = amdsmi_get_processor_handles
                shim.amdsmi_get_gpu_asic_info = amdsmi_get_gpu_asic_info
                shim.amdsmi_topo_get_link_type = amdsmi_topo_get_link_type
                sys.modules["amdsmi"] = shim
    except Exception:
        pass
    import vllm
    from vllm import LLM, SamplingParams
    logging.getLogger("vllm").setLevel(logging.ERROR)
    sampling_params = SamplingParams(temperature=0.0, max_tokens=max_tokens)
    llm = LLM(**llm_kwargs)
    start = time.perf_counter()
    outputs = llm.generate(prompts, sampling_params)
    elapsed = time.perf_counter() - start

    throughput_samples = []
    for prompt in prompts:
        sample_start = time.perf_counter()
        sample_out = llm.generate([prompt], sampling_params)
        sample_elapsed = time.perf_counter() - sample_start
        tokens = len(sample_out[0].outputs[0].token_ids)
        throughput_samples.append(tokens / sample_elapsed if sample_elapsed > 0 else 0.0)

    total_tokens = sum(len(output.outputs[0].token_ids) for output in outputs)
    throughput = total_tokens / elapsed if elapsed > 0 else 0.0
    _emit({
        "ok": True,
        "throughput_tokens_per_sec": round(throughput, 2),
        "latency_ms": round((elapsed / max(len(prompts), 1)) * 1000, 2),
        "throughput_samples": [round(x, 2) for x in throughput_samples],
    })
except Exception as exc:
    tb = traceback.format_exc().strip()
    tail_lines = [line.strip() for line in tb.splitlines()[-12:] if line.strip()]
    msg = str(exc).strip() or repr(exc)
    if tail_lines:
        msg = f"{msg} :: traceback_tail: {' | '.join(tail_lines)}"
    _emit({"ok": False, "error": msg})
"""
        env = os.environ.copy()
        if isinstance(env_overrides, dict):
            for key, value in env_overrides.items():
                if value is None:
                    env.pop(key, None)
                else:
                    env[key] = str(value)
        try:
            proc = subprocess.run(
                [sys.executable, "-c", runner],
                input=json.dumps(payload),
                capture_output=True,
                text=True,
                timeout=240,
                env=env,
            )
        except Exception as exc:
            return {"ok": False, "error": f"subprocess launch failed: {exc}"}

        merged = (proc.stdout or "").splitlines()
        if proc.stderr:
            merged.extend(proc.stderr.splitlines())
        for line in reversed(merged):
            text = line.strip()
            if not text.startswith("{") or not text.endswith("}"):
                continue
            try:
                parsed = json.loads(text)
            except Exception:
                continue
            if isinstance(parsed, dict) and "ok" in parsed:
                if not parsed.get("ok", False):
                    log_tail = []
                    for raw in merged:
                        msg = raw.strip()
                        if not msg:
                            continue
                        if msg.startswith("{") and msg.endswith("}"):
                            continue
                        lower = msg.lower()
                        if any(
                            marker in lower
                            for marker in (
                                "error",
                                "exception",
                                "traceback",
                                "runtimeerror",
                                "valueerror",
                                "fatal",
                                "enginecore",
                                "failed",
                            )
                        ):
                            log_tail.append(msg)
                    if log_tail:
                        tail = " | ".join(log_tail[-12:])[:3000]
                        base = str(parsed.get("error") or "unknown vLLM subprocess failure")
                        if tail and tail not in base:
                            parsed["error"] = f"{base}; engine_log_tail: {tail}"
                return parsed

        detail = " ".join([line.strip() for line in merged if line.strip()])[:4000]
        if not detail:
            detail = f"exit code {proc.returncode}"
        return {"ok": False, "error": detail}

    for candidate in candidates:
        model_name = candidate.get("model")
        model_format = candidate.get("format", "unknown")
        if not model_name:
            continue

        llm_kwargs = {
            "model": model_name,
            "trust_remote_code": True,
            "enforce_eager": True,
            "gpu_memory_utilization": 0.65,
            "disable_log_stats": True,
            "tensor_parallel_size": 1,
            "dtype": "float16",
        }
        if model_format == "gguf":
            tokenizer_name = candidate.get("tokenizer")
            if tokenizer_name:
                llm_kwargs["tokenizer"] = tokenizer_name

        attempt = _run_vllm_attempt_subprocess(llm_kwargs, prompts, max_tokens)
        if attempt.get("ok"):
            return True, {
                "model": model_name,
                "model_format": model_format,
                "target_device": target_device,
                "visible_devices": os.environ.get("HIP_VISIBLE_DEVICES", primary_visible),
                "throughput_tokens_per_sec": float(attempt.get("throughput_tokens_per_sec", 0.0)),
                "latency_ms": float(attempt.get("latency_ms", 0.0)),
                "throughput_samples": [
                    float(x) for x in (attempt.get("throughput_samples") or [])
                ],
                "candidate_models": candidate_names,
            }, []
        err_msg = str(attempt.get("error") or "unknown vLLM execution error")
        engine_core_failure = (
            "Engine core initialization failed" in err_msg
            or "Failed core proc" in err_msg
            or "No HIP GPUs are available" in err_msg
            or "torch.cuda is not available" in err_msg
        )
        # Defensive fallback for intermittent vLLM device/runtime failures.
        # Preserve full visible GPU list for multi-GPU hosts.
        if "Device string must not be empty" in err_msg or engine_core_failure:
            retry_profiles = [
                (
                    "spawn-v1-default",
                    {
                        "VLLM_WORKER_MULTIPROC_METHOD": "spawn",
                        "VLLM_ENABLE_V1_MULTIPROCESSING": "1",
                        "VLLM_USE_V1": "1",
                    },
                ),
                (
                    "fork-v1-default",
                    {
                        "VLLM_WORKER_MULTIPROC_METHOD": "fork",
                        "VLLM_ENABLE_V1_MULTIPROCESSING": "1",
                        "VLLM_USE_V1": "1",
                    },
                ),
                (
                    "spawn-v0-singleproc",
                    {
                        "VLLM_WORKER_MULTIPROC_METHOD": "spawn",
                        "VLLM_ENABLE_V1_MULTIPROCESSING": "0",
                        "VLLM_USE_V1": "0",
                    },
                ),
                (
                    "spawn-v0-no-aiter",
                    {
                        "VLLM_WORKER_MULTIPROC_METHOD": "spawn",
                        "VLLM_ENABLE_V1_MULTIPROCESSING": "0",
                        "VLLM_USE_V1": "0",
                        "VLLM_ROCM_USE_AITER": "0",
                    },
                ),
            ]
            retry_failures = []
            for profile_name, profile_overrides in retry_profiles:
                retry_overrides = {
                    "HIP_VISIBLE_DEVICES": visible_devices,
                    "CUDA_VISIBLE_DEVICES": visible_devices,
                    "PYTORCH_ROCM_DEVICE": primary_visible,
                    "VLLM_TARGET_DEVICE": target_device,
                }
                retry_overrides.update(profile_overrides)
                retry = _run_vllm_attempt_subprocess(
                    llm_kwargs,
                    prompts,
                    max_tokens,
                    env_overrides=retry_overrides,
                )
                if retry.get("ok"):
                    return True, {
                        "model": model_name,
                        "model_format": model_format,
                        "target_device": target_device,
                        "throughput_tokens_per_sec": float(retry.get("throughput_tokens_per_sec", 0.0)),
                        "latency_ms": float(retry.get("latency_ms", 0.0)),
                        "throughput_samples": [float(x) for x in (retry.get("throughput_samples") or [])],
                        "candidate_models": candidate_names,
                        "visible_devices_after_retry": visible_devices,
                        "fallback_mode": profile_name,
                    }, []
                retry_failures.append(f"{profile_name}: {retry.get('error', 'unknown retry failure')}")
            if retry_failures:
                err_msg = f"{err_msg}; retry_profiles_failed: {' || '.join(retry_failures)}"
        if "Entry Not Found" in err_msg or "not found on the Hugging Face Hub" in err_msg:
            err_msg = (
                f"Model {model_name} not found. Download/copy it first or override "
                "MLSTACK_BENCH_VLLM_SAFETENSORS_MODEL / MLSTACK_BENCH_VLLM_GGUF_MODEL_PATH"
            )
        elif "authentication" in err_msg.lower():
            err_msg = "Hugging Face authentication required. Run 'huggingface-cli login' first."
        attempt_errors.append(f"{model_format} model {model_name} failed: {err_msg}")

    failure_reason = "No tiny vLLM benchmark model could be loaded"
    if runtime_probe_hint:
        failure_reason = f"{failure_reason}; runtime_probe={runtime_probe_hint}"
    if attempt_errors:
        failure_reason = f"{failure_reason}; {attempt_errors[0]}"

    return True, _degraded_metrics("vllm", failure_reason, {
        "model": "vLLM (Detected)",
        "model_format": "unavailable",
        "target_device": target_device,
        "visible_devices": os.environ.get("HIP_VISIBLE_DEVICES", primary_visible),
        "throughput_tokens_per_sec": 0.0,
        "latency_ms": 0.0,
        "throughput_samples": [],
        "candidate_models": candidate_names,
    }), attempt_errors or [failure_reason]


def _deepspeed():
    _set_amd_gpu_id_table_env()
    try:
        import deepspeed
        import torch
        import os
    except Exception as exc:
        err = f"DeepSpeed or Torch not available: {exc}"
        return True, _degraded_metrics("deepspeed", err, {
            "throughput_samples_per_sec": 0.0,
            "avg_latency_ms": 0.0,
            "stage": 1,
            "accelerator": "rocm",
            "samples": [],
        }), [err]
    probe_ok, probe_reason, probe_meta = _probe_gpu_runtime()
    if not probe_ok:
        extra = {
            "throughput_samples_per_sec": 0.0,
            "avg_latency_ms": 0.0,
            "stage": 1,
            "accelerator": "rocm",
            "samples": [],
        }
        if isinstance(probe_meta, dict):
            extra.update(probe_meta)
        return True, _degraded_metrics("deepspeed", probe_reason, extra), [probe_reason]

    # Ensure ROCm environment
    os.environ["DS_ACCELERATOR"] = "rocm"
    
    try:
        from deepspeed.accelerator import get_accelerator

        accelerator = get_accelerator()
        device_name = accelerator.device_name()
        device = torch.device(device_name if ":" in device_name else f"{device_name}:0")

        class SimpleModel(torch.nn.Module):
            def __init__(self):
                super().__init__()
                self.net = torch.nn.Sequential(
                    torch.nn.Linear(1024, 1024),
                    torch.nn.ReLU(),
                    torch.nn.Linear(1024, 1024),
                )

            def forward(self, x):
                return self.net(x)

        model = SimpleModel().to(device).half()
        batch_size = 4
        input_data = torch.randn(batch_size, 1024, device=device, dtype=torch.float16)

        ds_config = {
            "train_batch_size": batch_size,
            "train_micro_batch_size_per_gpu": batch_size,
            "gradient_accumulation_steps": 1,
            "zero_optimization": {"stage": 0},
            "fp16": {"enabled": True},
            # DeepSpeed internally computes `step % steps_per_print`; keep this non-zero.
            "steps_per_print": 1,
            "wall_clock_breakdown": False,
        }
        optimizer = torch.optim.AdamW(model.parameters(), lr=0.001)

        try:
            model_engine, _, _, _ = deepspeed.initialize(
                config=ds_config,
                model=model,
                model_parameters=model.parameters(),
                optimizer=optimizer,
            )
            mode = "train"
        except Exception as init_exc:
            init_err = str(init_exc)
            if "fused_adam" in init_err.lower() or "cuda_runtime_api.h" in init_err:
                model_engine = deepspeed.init_inference(
                    model,
                    mp_size=1,
                    dtype=torch.float16,
                    replace_with_kernel_inject=False,
                )
                mode = "inference"
            else:
                raise

        times = []
        if mode == "train":
            for _ in range(2):
                out = model_engine(input_data)
                model_engine.backward(out.sum())
                model_engine.step()

            for _ in range(5):
                accelerator.synchronize()
                start = time.perf_counter()
                out = model_engine(input_data)
                model_engine.backward(out.sum())
                model_engine.step()
                accelerator.synchronize()
                times.append(time.perf_counter() - start)
        else:
            model_engine.eval()
            for _ in range(2):
                with torch.no_grad():
                    _ = model_engine(input_data)
                accelerator.synchronize()

            for _ in range(5):
                accelerator.synchronize()
                start = time.perf_counter()
                with torch.no_grad():
                    _ = model_engine(input_data)
                accelerator.synchronize()
                times.append(time.perf_counter() - start)

        avg_time = sum(times) / len(times)
        throughput = batch_size / avg_time

        return True, {
            "throughput_samples_per_sec": round(throughput, 2),
            "avg_latency_ms": round(avg_time * 1000, 2),
            "stage": 0 if mode == "train" else 0,
            "accelerator": "rocm",
            "mode": mode,
            "samples": [round(batch_size / t, 2) for t in times],
        }, []
    except Exception as exc:
        err = f"DeepSpeed benchmark failed: {exc}"
        return True, _degraded_metrics("deepspeed", err, {
            "throughput_samples_per_sec": 0.0,
            "avg_latency_ms": 0.0,
            "stage": 1,
            "accelerator": "rocm",
            "samples": [],
        }), [err]


def _megatron():
    try:
        import torch
        import megatron
    except Exception as exc:
        err = f"Megatron not available: {exc}"
        return True, _degraded_metrics("megatron", err, {
            "megatron_backend": "unavailable",
            "megatron_throughput_samples_per_sec": 0.0,
            "megatron_avg_latency_ms": 0.0,
            "megatron_samples": [],
        }), [err]

    probe_ok, probe_reason, probe_meta = _probe_gpu_runtime()
    if not probe_ok:
        extra = {
            "megatron_backend": "megatron-core",
            "megatron_throughput_samples_per_sec": 0.0,
            "megatron_avg_latency_ms": 0.0,
            "megatron_samples": [],
        }
        if isinstance(probe_meta, dict):
            extra.update(probe_meta)
        return True, _degraded_metrics("megatron", probe_reason, extra), [probe_reason]

    try:
        device = torch.device("cuda:0")
        model = torch.nn.Sequential(
            torch.nn.Linear(4096, 4096),
            torch.nn.GELU(),
            torch.nn.Linear(4096, 4096),
        ).to(device).half()

        input_data = torch.randn(8, 4096, device=device, dtype=torch.float16)
        samples = []

        for _ in range(2):
            out = model(input_data)
            _ = out.sum()
            torch.cuda.synchronize()

        for _ in range(5):
            torch.cuda.synchronize()
            start = time.perf_counter()
            out = model(input_data)
            _ = out.sum()
            torch.cuda.synchronize()
            elapsed = time.perf_counter() - start
            samples.append(8.0 / elapsed if elapsed > 0 else 0.0)

        avg_tput = sum(samples) / len(samples) if samples else 0.0
        avg_latency_ms = (1000.0 / avg_tput * 8.0) if avg_tput > 0 else 0.0
        return True, {
            "megatron_backend": getattr(megatron, "__name__", "megatron"),
            "megatron_throughput_samples_per_sec": round(avg_tput, 2),
            "megatron_avg_latency_ms": round(avg_latency_ms, 2),
            "megatron_samples": [round(x, 2) for x in samples],
        }, []
    except Exception as exc:
        err = f"Megatron benchmark failed: {exc}"
        return True, _degraded_metrics("megatron", err, {
            "megatron_backend": "megatron-core",
            "megatron_throughput_samples_per_sec": 0.0,
            "megatron_avg_latency_ms": 0.0,
            "megatron_samples": [],
        }), [err]


def _onnx():
    import os
    try:
        import onnxruntime as ort
    except Exception as exc:
        err = f"onnxruntime not available: {exc}"
        return True, _degraded_metrics("onnx", err, {
            "ort_version": "unavailable",
            "provider": "none",
            "providers_available": [],
            "model_load_ms": 0.0,
            "session_create_ms": 0.0,
            "inference_latency_p50_ms": 0.0,
            "inference_latency_p95_ms": 0.0,
            "inference_latency_p99_ms": 0.0,
            "throughput_inf_per_sec": 0.0,
            "input_shape": [],
            "output_shape": [],
            "graph_opt_level": "unavailable",
            "inference_samples": [],
        }), [err]

    ort_version = ort.__version__
    all_providers = ort.get_available_providers()
    providers_available = list(all_providers)

    if "MIGraphXExecutionProvider" in all_providers:
        provider = "MIGraphXExecutionProvider"
    elif "DmlExecutionProvider" in all_providers:
        provider = "DmlExecutionProvider"
    elif "ROCMExecutionProvider" in all_providers:
        provider = "ROCMExecutionProvider"
    elif "CUDAExecutionProvider" in all_providers:
        provider = "CUDAExecutionProvider"
    else:
        provider = "CPUExecutionProvider"

    try:
        import numpy as np
    except Exception:
        err = "numpy not available for ONNX benchmark"
        return True, _degraded_metrics("onnx", err, {
            "ort_version": ort_version,
            "provider": provider,
            "providers_available": providers_available,
            "model_load_ms": 0.0,
            "session_create_ms": 0.0,
            "inference_latency_p50_ms": 0.0,
            "inference_latency_p95_ms": 0.0,
            "inference_latency_p99_ms": 0.0,
            "throughput_inf_per_sec": 0.0,
            "input_shape": [],
            "output_shape": [],
            "graph_opt_level": "ORT_ENABLE_ALL",
            "inference_samples": [],
        }), [err]

    try:
        import onnx
        from onnx import helper, TensorProto, numpy_helper
    except Exception:
        onnx = None

    batch_size = 8
    hidden = 512
    input_shape = [batch_size, hidden]

    def _build_test_model():
        X = helper.make_tensor_value_info("input", TensorProto.FLOAT, input_shape)
        Y = helper.make_tensor_value_info("output", TensorProto.FLOAT, input_shape)
        W_init = numpy_helper.from_array(
            np.random.randn(hidden, hidden).astype(np.float32) * 0.01, name="W"
        )
        B_init = numpy_helper.from_array(
            np.zeros(hidden, dtype=np.float32), name="B"
        )
        matmul = helper.make_node("MatMul", ["input", "W"], ["matmul_out"])
        add = helper.make_node("Add", ["matmul_out", "B"], ["add_out"])
        relu = helper.make_node("Relu", ["add_out"], ["output"])
        graph = helper.make_graph(
            [matmul, add, relu], "bench_graph", [X], [Y], initializer=[W_init, B_init]
        )
        model = helper.make_model(graph, opset_imports=[helper.make_opsetid("", 17)])
        model.ir_version = 8
        onnx.checker.validate = False
        return onnx._serialize(model)

    def _build_quantized_model():
        X = helper.make_tensor_value_info("input", TensorProto.FLOAT, input_shape)
        Y = helper.make_tensor_value_info("output", TensorProto.FLOAT, input_shape)
        W_init = numpy_helper.from_array(
            np.random.randn(hidden, hidden).astype(np.float32) * 0.01, name="W"
        )
        B_init = numpy_helper.from_array(
            np.zeros(hidden, dtype=np.float32), name="B"
        )
        dq = helper.make_node(
            "DynamicQuantizeLinear", ["input"], ["input_quant", "input_scale", "input_zp"]
        )
        matmul_int = helper.make_node(
            "MatMulInteger", ["input_quant", "W"], ["matmul_int_out"]
        )
        cast_out = helper.make_node(
            "Cast", ["matmul_int_out"], ["matmul_fp_out"], to=TensorProto.FLOAT
        )
        add = helper.make_node("Add", ["matmul_fp_out", "B"], ["output"])
        graph = helper.make_graph(
            [dq, matmul_int, cast_out, add], "quant_bench_graph",
            [X], [Y], initializer=[W_init, B_init]
        )
        model = helper.make_model(graph, opset_imports=[helper.make_opsetid("", 17)])
        model.ir_version = 8
        onnx.checker.validate = False
        return onnx._serialize(model)

    import tempfile
    sess_opts = ort.SessionOptions()
    sess_opts.graph_optimization_level = ort.GraphOptimizationLevel.ORT_ENABLE_ALL
    sess_opts.log_severity_level = 3

    results = {}
    errors = []

    # --- FP32 model benchmark ---
    try:
        if onnx is not None:
            model_bytes = _build_test_model()
        else:
            model_bytes = None

        if model_bytes is None:
            raise RuntimeError("onnx package required to build test model")

        with tempfile.NamedTemporaryFile(suffix=".onnx", delete=False) as f:
            f.write(model_bytes)
            fp32_path = f.name

        t0 = time.perf_counter()
        session = ort.InferenceSession(fp32_path, sess_opts, providers=[provider])
        session_create_ms = (time.perf_counter() - t0) * 1000.0

        input_meta = session.get_inputs()[0]
        output_meta = session.get_outputs()[0]
        input_name = input_meta.name
        actual_input_shape = input_meta.shape
        output_shape = output_meta.shape

        if isinstance(actual_input_shape[0], int):
            bench_batch = actual_input_shape[0]
        else:
            bench_batch = batch_size
        actual_input_shape_resolved = [bench_batch] + [
            d if isinstance(d, int) else hidden for d in actual_input_shape[1:]
        ]

        feed = {input_name: np.random.randn(*actual_input_shape_resolved).astype(np.float32)}

        for _ in range(3):
            session.run(None, feed)

        num_runs = 20
        latencies = []
        for _ in range(num_runs):
            t_start = time.perf_counter()
            session.run(None, feed)
            latencies.append((time.perf_counter() - t_start) * 1000.0)

        latencies.sort()
        p50 = latencies[len(latencies) // 2]
        p95 = latencies[int(len(latencies) * 0.95)]
        p99 = latencies[int(len(latencies) * 0.99)]
        avg_latency = sum(latencies) / len(latencies)
        throughput = bench_batch / (avg_latency / 1000.0) if avg_latency > 0 else 0.0

        mem_rss_mb = 0.0
        try:
            import resource
            mem_rss_mb = resource.getrusage(resource.RUSAGE_SELF).ru_maxrss / 1024.0
        except Exception:
            pass

        results.update({
            "fp32_session_create_ms": round(session_create_ms, 2),
            "fp32_inference_latency_p50_ms": round(p50, 3),
            "fp32_inference_latency_p95_ms": round(p95, 3),
            "fp32_inference_latency_p99_ms": round(p99, 3),
            "fp32_avg_latency_ms": round(avg_latency, 3),
            "fp32_throughput_inf_per_sec": round(throughput, 2),
            "fp32_inference_samples": [round(l, 3) for l in latencies],
            "input_shape": actual_input_shape_resolved,
            "output_shape": list(output_shape) if output_shape else [],
            "input_dtype": str(input_meta.type),
            "output_dtype": str(output_meta.type),
            "peak_rss_mb": round(mem_rss_mb, 1),
            "num_warmup": 3,
            "num_timed_runs": num_runs,
        })

        os.unlink(fp32_path)
        del session
    except Exception as exc:
        results["fp32_error"] = str(exc)
        errors.append(f"FP32 benchmark failed: {exc}")

    # --- Quantized model benchmark ---
    quantized_ok = False
    try:
        if onnx is not None:
            q_model_bytes = _build_quantized_model()
            if q_model_bytes is not None:
                with tempfile.NamedTemporaryFile(suffix=".onnx", delete=False) as f:
                    f.write(q_model_bytes)
                    quant_path = f.name

                t0 = time.perf_counter()
                q_session = ort.InferenceSession(
                    quant_path, sess_opts, providers=[provider]
                )
                q_create_ms = (time.perf_counter() - t0) * 1000.0

                q_input = q_session.get_inputs()[0]
                q_feed = {
                    q_input.name: np.random.randn(
                        *([batch_size] + [
                            d if isinstance(d, int) else hidden
                            for d in q_input.shape[1:]
                        ])
                    ).astype(np.float32)
                }

                for _ in range(3):
                    q_session.run(None, q_feed)

                q_latencies = []
                for _ in range(20):
                    t_start = time.perf_counter()
                    q_session.run(None, q_feed)
                    q_latencies.append((time.perf_counter() - t_start) * 1000.0)

                q_latencies.sort()
                q_p50 = q_latencies[len(q_latencies) // 2]
                q_avg = sum(q_latencies) / len(q_latencies)
                q_throughput = batch_size / (q_avg / 1000.0) if q_avg > 0 else 0.0

                results.update({
                    "quantized_session_create_ms": round(q_create_ms, 2),
                    "quantized_inference_latency_p50_ms": round(q_p50, 3),
                    "quantized_avg_latency_ms": round(q_avg, 3),
                    "quantized_throughput_inf_per_sec": round(q_throughput, 2),
                    "quantized_inference_samples": [round(l, 3) for l in q_latencies],
                    "quantized_ops": ["DynamicQuantizeLinear", "MatMulInteger"],
                })
                quantized_ok = True
                os.unlink(quant_path)
                del q_session
    except Exception as exc:
        results["quantized_error"] = str(exc)
        errors.append(f"Quantized benchmark failed: {exc}")

    p50_val = results.get("fp32_inference_latency_p50_ms", 0.0)
    p95_val = results.get("fp32_inference_latency_p95_ms", 0.0)
    p99_val = results.get("fp32_inference_latency_p99_ms", 0.0)
    tput_val = results.get("fp32_throughput_inf_per_sec", 0.0)

    metrics = {
        "ort_version": ort_version,
        "provider": provider,
        "providers_available": providers_available,
        "graph_opt_level": "ORT_ENABLE_ALL",
        "model_load_ms": results.get("fp32_session_create_ms", 0.0),
        "session_create_ms": results.get("fp32_session_create_ms", 0.0),
        "inference_latency_p50_ms": p50_val,
        "inference_latency_p95_ms": p95_val,
        "inference_latency_p99_ms": p99_val,
        "throughput_inf_per_sec": tput_val,
        "input_shape": results.get("input_shape", input_shape),
        "output_shape": results.get("output_shape", []),
        "input_dtype": results.get("input_dtype", "tensor(float)"),
        "output_dtype": results.get("output_dtype", "tensor(float)"),
        "peak_rss_mb": results.get("peak_rss_mb", 0.0),
        "num_warmup": results.get("num_warmup", 0),
        "num_timed_runs": results.get("num_timed_runs", 0),
        "quantized_supported": quantized_ok,
        "inference_samples": results.get("fp32_inference_samples", []),
    }

    for k in ("fp32_inference_samples", "fp32_session_create_ms", "fp32_inference_latency_p50_ms",
              "fp32_inference_latency_p95_ms", "fp32_inference_latency_p99_ms",
              "fp32_avg_latency_ms", "fp32_throughput_inf_per_sec",
              "quantized_session_create_ms", "quantized_inference_latency_p50_ms",
              "quantized_avg_latency_ms", "quantized_throughput_inf_per_sec",
              "quantized_inference_samples", "quantized_ops",
              "quantized_error", "fp32_error"):
        if k in results:
            metrics[k] = results[k]

    success = True
    return success, metrics, errors


BENCHES = {
    "gpu-info": _gpu_info,
    "memory-bandwidth": _memory_bandwidth,
    "tensor-core": _tensor_core,
    "gemm": _gemm,
    "pytorch": _pytorch,
    "flash-attention": _flash_attention,
    "vllm": _vllm,
    "deepspeed": _deepspeed,
    "megatron": _megatron,
    "onnx": _onnx,
    "llama-cpp": _llama_cpp,
}


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("bench", help=f"One of: {', '.join(BENCHES.keys())}")
    parser.add_argument("--json", action="store_true", help="Emit JSON")
    args = parser.parse_args()

    fn = BENCHES.get(args.bench)
    if fn is None:
        payload = {
            "name": args.bench,
            "success": False,
            "execution_time_ms": 0,
            "metrics": {},
            "errors": [f"Unknown benchmark {args.bench}"],
        }
        print(json.dumps(payload, indent=2))
        sys.exit(1)

    start = time.perf_counter()
    success, metrics, errors = fn()
    elapsed_ms = int((time.perf_counter() - start) * 1000)
    payload = {
        "name": args.bench,
        "success": success,
        "execution_time_ms": elapsed_ms,
        "metrics": metrics,
        "errors": errors,
    }
    
    # Print a clear marker to help the Rust wrapper find the JSON block
    print("\n---BENCHMARK_RESULTS_START---")
    print(json.dumps(payload, indent=2))
    sys.exit(0 if success else 1)


if __name__ == "__main__":
    main()
"#;