Enhancement Request: Environment Health Check

[MrEdwards007]

This is for those who need help knowing whether they have their environment configured to support an RTX 5090
With some minor modifications, this will support other cards and environments.

import os
import sys
import subprocess
import torch
import triton

def get_nvcc_info():
    try:
        nvcc_path = subprocess.check_output(["which", "nvcc"]).decode("utf-8").strip()
        nvcc_ver = subprocess.check_output(["nvcc", "--version"]).decode("utf-8")
        version_line = [line for line in nvcc_ver.split('\n') if "release" in line][0]
        return nvcc_path, version_line
    except Exception:
        return None, "❌ NVCC NOT FOUND (Required for local compilation)"

def check_env_vars():
    vars = ['CUDA_HOME', 'LD_LIBRARY_PATH', 'PATH']
    return {v: os.environ.get(v, "Not Set") for v in vars}

print(f"--- 🛡️ Blackwell (RTX 5090) Environment Health Check ---")
print(f"Python Version:    {sys.version.split()[0]}")
print(f"Triton Version:    {triton.__version__}")
print(f"PyTorch Version:   {torch.__version__}")
print(f"PyTorch CUDA:      {torch.version.cuda}")

if torch.cuda.is_available():
    prop = torch.cuda.get_device_properties(0)
    print(f"Hardware Detected: {prop.name} (Arch: {prop.major}.{prop.minor})")
else:
    print("Hardware Detected: ❌ NO GPU DETECTED")

# Compiler Check
nv_path, nv_ver = get_nvcc_info()
print(f"Compiler Path:     {nv_path}")
print(f"Compiler Version:  {nv_ver}")

# Path Check
env_vars = check_env_vars()
print(f"CUDA_HOME:         {env_vars['CUDA_HOME']}")

print("\n" + "="*50)
print("--- 📋 EVALUATION ---")

# 1. Check for Architecture Support (The Blackwell Gap)
is_blackwell_supported = "12.8" in nv_ver or "13.0" in nv_ver
if not is_blackwell_supported:
    print("❌ CRITICAL: Your system compiler (nvcc) is too old for the RTX 5090.")
    print("   Action: Install CUDA Toolkit 12.8+ globally or via Conda.")
else:
    print("✅ System compiler supports Blackwell architecture.")

# 2. Check for CUDA_HOME Mismatch
if env_vars['CUDA_HOME'] != "Not Set" and nv_path and env_vars['CUDA_HOME'] not in nv_path:
    print("⚠️  WARNING: CUDA_HOME points to a different folder than your active nvcc.")
    print(f"   CUDA_HOME: {env_vars['CUDA_HOME']} vs NVCC: {nv_path}")
else:
    print("✅ Pathing logic is consistent.")

# 3. Check for Triton version
if triton.__version__ < '3.5.1':
    print("⚠️  WARNING: Triton is < 3.5.1. You may face performance issues on 50-series.")
else:
    print("✅ Triton version is optimized for Blackwell.")

print("="*50)

Output for my machine
.

--- 🛡️ Blackwell (RTX 5090) Environment Health Check ---
Python Version:    3.12.12
Triton Version:    3.5.1
PyTorch Version:   2.9.1+cu128
PyTorch CUDA:      12.8
Hardware Detected: NVIDIA GeForce RTX 5090 (Arch: 12.0)
Compiler Path:     /home/your_path/anaconda3/envs/turbodiffusion/bin/nvcc
Compiler Version:  Cuda compilation tools, release 12.8, V12.8.61
CUDA_HOME:         /home/your_path/anaconda3/envs/turbodiffusion

==================================================
--- 📋 EVALUATION ---
✅ System compiler supports Blackwell architecture.
✅ Pathing logic is consistent.
✅ Triton version is optimized for Blackwell.
==================================================