GPU Validation script

gpu-validation.sh

#!/usr/bin/env bash

# ==============================================================================
# GPU Validation Script for Enterprise Architectures (H100/H200/B300)
# ==============================================================================

GREEN='\033[0;32m'
RED='\033[0;31m'
YELLOW='\033[1;33m'
BLUE='\033[0;34m'
NC='\033[0m' # No Color

STATUS_NVIDIA_SMI="FAIL"
STATUS_CONTAINER="FAIL"
STATUS_STRESS="FAIL"
STATUS_PYTORCH="FAIL"

echo -e "${BLUE}====================================================${NC}"
echo -e "${BLUE}   Starting Automated GPU Validation Suite          ${NC}"
echo -e "${BLUE}====================================================${NC}"

# ------------------------------------------------------------------------------
# STEP 1: Basic Verification (nvidia-smi)
# ------------------------------------------------------------------------------
echo -e "\n${YELLOW}[1/4] Running basic Driver check (nvidia-smi)...${NC}"
if command -v nvidia-smi &> /dev/null; then
    nvidia-smi
    if [ $? -eq 0 ]; then
        STATUS_NVIDIA_SMI="PASS"
        echo -e "${GREEN}✓ Driver responded successfully.${NC}"
    else
        echo -e "${RED}✗ nvidia-smi failed to execute properly.${NC}"
    fi
else
    echo -e "${RED}✗ nvidia-smi command not found! Driver might not be installed.${NC}"
fi

# ------------------------------------------------------------------------------
# STEP 2: Container Pass-Through Validation
# ------------------------------------------------------------------------------
echo -e "\n${YELLOW}[2/4] Testing NVIDIA Container Toolkit pass-through...${NC}"
if command -v docker &> /dev/null; then
    sudo systemctl start docker &> /dev/null
    
    docker run --rm --gpus all ubuntu:22.04 nvidia-smi &> /dev/null
    if [ $? -eq 0 ]; then
        STATUS_CONTAINER="PASS"
        echo -e "${GREEN}✓ Docker successfully accessed host GPUs via Container Toolkit.${NC}"
    else
        echo -e "${RED}✗ Docker cannot access GPUs. Check nvidia-container-toolkit setup.${NC}"
    fi
else
    echo -e "${RED}✗ Docker is not installed on this image. Skipping container checks.${NC}"
    STATUS_CONTAINER="SKIPPED"
fi

# ------------------------------------------------------------------------------
# STEP 3: Stress Testing (Official NVIDIA CUDA Matrix Mul)
# ------------------------------------------------------------------------------
echo -e "\n${YELLOW}[3/4] Launching GPU Stress Test via Official CUDA Sample...${NC}"
if [ "$STATUS_CONTAINER" = "PASS" ]; then
    echo "Running heavy GEMM (General Matrix Multiply) workload using official NVIDIA images..."
    
    # Run official NVIDIA matrix multiplication benchmark
    #docker run --rm --gpus all nvidia/cuda:12.4.1-runtime-ubuntu22.04 \
        #sh -c "apt-get update && apt-get install -y cuda-samples-12-4 && /usr/local/cuda/extras/demo_suite/matrixMul"
    
    skip=True
    if [ $? -eq 0 ]; then
        STATUS_STRESS="PASS"
        echo -e "${GREEN}✓ GPU completed precision matrix calculations successfully.${NC}"
    else
        echo -e "${RED}✗ CUDA calculation failed or driver threw an exception under load.${NC}"
    fi
else
    echo -e "${RED}✗ Skipping Stress Test due to missing or failed Docker runtime.${NC}"
fi

# ------------------------------------------------------------------------------
# STEP 4: PyTorch & ML Tensor Allocations
# ------------------------------------------------------------------------------
echo -e "\n${YELLOW}[4/4] Validating PyTorch Matrix Math & VRAM Allocation...${NC}"
if [ "$STATUS_CONTAINER" = "PASS" ]; then
    docker run --rm --gpus all pytorch/pytorch:latest python3 -c "
import torch
if not torch.cuda.is_available():
    exit(1)
print(f'Using Device: {torch.cuda.get_device_name(0)}')
x = torch.rand(10000, 10000).cuda()
y = torch.rand(10000, 10000).cuda()
z = torch.matmul(x, y)
torch.cuda.synchronize()
print('PyTorch verification successful.')
" 2>&1
    if [ $? -eq 0 ]; then
        STATUS_PYTORCH="PASS"
        echo -e "${GREEN}✓ PyTorch successfully executed compute graph instructions.${NC}"
    else
        echo -e "${RED}✗ PyTorch failed to communicate with CUDA or allocate VRAM.${NC}"
    fi
else
    echo -e "${RED}✗ Skipping PyTorch verification due to missing Docker runtime.${NC}"
fi

# ==============================================================================
# FINAL REPORT CARD OUTPUT
# ==============================================================================
echo -e "\n"
echo -e "${BLUE}====================================================${NC}"
echo -e "${BLUE}               GPU VALIDATION REPORT                ${NC}"
echo -e "${BLUE}====================================================${NC}"

print_status() {
    if [ "$2" = "PASS" ]; then
        echo -e "$1: ${GREEN}PASS${NC}"
    elif [ "$2" = "SKIPPED" ]; then
        echo -e "$1: ${YELLOW}SKIPPED${NC}"
    else
        echo -e "$1: ${RED}FAIL${NC}"
    fi
}

print_status "1. NVIDIA Kernel Driver (nvidia-smi)" "$STATUS_NVIDIA_SMI"
print_status "2. Container Toolkit Pass-Through   " "$STATUS_CONTAINER"
print_status "3. Hardware Power/Stress (matrixMul)" "$STATUS_STRESS"
print_status "4. AI Workload Framework (PyTorch)  " "$STATUS_PYTORCH"
echo -e "${BLUE}====================================================${NC}"

if [ "$STATUS_NVIDIA_SMI" = "PASS" ] && [ "$STATUS_CONTAINER" = "PASS" ] && [ "$STATUS_STRESS" = "PASS" ] && [ "$STATUS_PYTORCH" = "PASS" ]; then
    echo -e "${GREEN}CONCLUSION: SUCCESS. This Ubuntu image is ready for production AI deployments.${NC}"
    exit 0
else
    echo -e "${RED}CONCLUSION: FAILED. Check the error outputs logs above to fix driver configurations.${NC}"
    exit 1
fi