CUDA / ML check environment script using rich

check_env.py

#!/usr/bin/env python3

import subprocess
import importlib
import platform
import sys
import locale
from rich.console import Console
from rich.table import Table
from rich.panel import Panel
from rich import box
from datetime import datetime
import warnings
warnings.filterwarnings('ignore')

console = Console()

def run_command_with_encoding(command):
    try:
        result = subprocess.run(
            command,
            stdout=subprocess.PIPE,
            stderr=subprocess.PIPE,
            text=True,
            encoding='utf-8',
            errors='replace',
            check=True
        )
        return result.stdout
    except subprocess.CalledProcessError as e:
        return "not_installed"
    except Exception as e:
        return f"Error: {str(e)}"

def check_import(module_name, display_name=None):
    try:
        module = importlib.import_module(module_name)
        version = getattr(module, '__version__', 'Version unknown')
        return "[green]✓[/green]", version
    except ImportError:
        return "[red]✗[/red]", "Not installed"

def check_command(command):
    try:
        result = run_command_with_encoding(command)
        version = result.split('\n')[0]
        return "[green]✓[/green]", version
    except:
        return "[red]✗[/red]", "Not installed"

def run_gpu_tests():
    results = []

    # Test PyTorch
    try:
        import torch
        results.append(("PyTorch GPU", "[green]✓[/green]", f"CUDA {torch.version.cuda} ({torch.cuda.get_device_name(0) if torch.cuda.is_available() else 'Not available'})"))
    except Exception as e:
        results.append(("PyTorch GPU", "[red]✗[/red]", str(e)))

    # Test TensorFlow
    try:
        import tensorflow as tf
        gpus = tf.config.list_physical_devices('GPU')
        tf_status = "[green]✓[/green]" if gpus else "[red]✗[/red]"
        tf_info = f"{len(gpus)} GPU(s) detected" if gpus else "No GPUs found"
        results.append(("TensorFlow GPU", tf_status, tf_info))
    except Exception as e:
        results.append(("TensorFlow GPU", "[red]✗[/red]", str(e)))

    # Test RAPIDS (cuDF)
    try:
        import cudf
        results.append(("RAPIDS cuDF", "[green]✓[/green]", cudf.__version__))
    except Exception as e:
        results.append(("RAPIDS cuDF", "[red]✗[/red]", "Not available"))

    # Test CuPy with memory info
    try:
        import cupy as cp
        x_gpu = cp.array([1, 2, 3, 4, 5])
        mean = float(x_gpu.mean())
        mem_info = cp.cuda.runtime.memGetInfo()
        free_mb = mem_info[0] / (1024**2)
        total_mb = mem_info[1] / (1024**2)
        results.append(("CuPy", "[green]✓[/green]", f"v{cp.__version__} (Test: {mean}, GPU Memory: {free_mb:.0f}MB free / {total_mb:.0f}MB total)"))
    except Exception as e:
        results.append(("CuPy", "[red]✗[/red]", str(e)))

    # Test cuML
    try:
        import cuml
        results.append(("RAPIDS cuML", "[green]✓[/green]", cuml.__version__))
    except Exception as e:
        results.append(("RAPIDS cuML", "[red]✗[/red]", "Not available"))

    return results

def check_cuda_toolkit():
    try:
        result = run_command_with_encoding(['nvcc', '--version'])
        version = result.split('release')[1].split(',')[0].strip()
        return "[green]✓[/green]", f"Version {version}"
    except:
        return "[red]✗[/red]", "Not installed"

def check_environment():
    # Python packages table
    python_table = Table(title="[bold green]Python Packages", box=box.ROUNDED)
    python_table.add_column("Package", style="cyan", justify="right")
    python_table.add_column("Status", justify="center")
    python_table.add_column("Version", style="magenta")

    python_packages = [
        ("numpy", "NumPy"),
        ("pandas", "Pandas"),
        ("tensorflow", "TensorFlow"),
        ("torch", "PyTorch"),
        ("transformers", "Transformers"),
        ("cv2", "OpenCV"),
        ("cupy", "CuPy"),
        ("cudf", "cuDF"),
        ("cuml", "cuML"),
        ("xgboost", "XGBoost"),
        ("lightgbm", "LightGBM"),
        ("rich", "Rich"),
    ]

    for pkg, display_name in python_packages:
        status, version = check_import(pkg)
        python_table.add_row(display_name, status, version)

    # R packages table
    r_table = Table(title="[bold blue]R Packages", box=box.ROUNDED)
    r_table.add_column("Package", style="cyan", justify="right")
    r_table.add_column("Status", justify="center")

    r_packages = [
        "ggplot2", "dplyr", "tidyr", "shiny", "caret",
        "randomForest", "rmarkdown", "tensorflow",
        "xgboost", "lightgbm", "RMariaDB"
    ]

    for pkg in r_packages:
        output = run_command_with_encoding(["Rscript", "-e", f"if(require('{pkg}')) cat('installed') else cat('not_installed')"])
        status = "[green]✓[/green]" if "installed" in output.lower() else "[red]✗[/red]"
        r_table.add_row(pkg, status)

    # System commands table
    command_table = Table(title="[bold yellow]System Commands", box=box.ROUNDED)
    command_table.add_column("Command", style="cyan", justify="right")
    command_table.add_column("Status", justify="center")
    command_table.add_column("Version", style="magenta")

    commands = [
        (["nvidia-smi"], "NVIDIA SMI"),
        (["python3", "--version"], "Python3"),
        (["R", "--version"], "R"),
    ]

    for cmd, name in commands:
        status, version = check_command(cmd)
        command_table.add_row(name, status, version)

    return python_table, r_table, command_table

def create_system_info_panel():
    info_table = Table(show_header=False, box=box.ROUNDED)
    info_table.add_column("Property", style="cyan")
    info_table.add_column("Value", style="yellow")

    info_table.add_row("System", platform.system())
    info_table.add_row("Release", platform.release())
    info_table.add_row("Python Version", sys.version.split('\n')[0])
    info_table.add_row("Machine", platform.machine())

    try:
        result = run_command_with_encoding(['nvidia-smi', '--query-gpu=driver_version', '--format=csv,noheader'])
        driver_version = result.strip()
        info_table.add_row("NVIDIA Driver", f"[green]{driver_version}[/green]")
    except:
        info_table.add_row("NVIDIA Driver", "[red]Not found[/red]")

    cuda_status, cuda_version = check_cuda_toolkit()
    info_table.add_row("CUDA Toolkit", f"{cuda_status} {cuda_version}")

    return Panel(info_table, title="[bold blue]System Information", border_style="blue")

def create_gpu_test_panel():
    table = Table(title="[bold green]GPU Acceleration Tests", box=box.ROUNDED)
    table.add_column("Component", style="cyan", justify="right")
    table.add_column("Status", justify="center")
    table.add_column("Details", style="magenta")

    for name, status, details in run_gpu_tests():
        table.add_row(name, status, str(details))

    return table

def main():
    console.clear()
    console.print("[bold cyan]Environment Check Tool[/bold cyan]", justify="center")
    console.print(f"[grey]Run at: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}[/grey]\n", justify="center")

    console.print(create_system_info_panel())
    console.print("")
    console.print(create_gpu_test_panel())
    console.print("")

    python_table, r_table, command_table = check_environment()
    console.print(python_table)
    console.print("")
    console.print(r_table)
    console.print("")
    console.print(command_table)

if __name__ == "__main__":
    main()