test_gpu.sh

#!/usr/bin/env bash
# GPU test suite for slurm-docker-cluster
# Tests GPU node availability and GRES functionality

set -euo pipefail

# Colors for output
RED='\033[0;31m'
GREEN='\033[0;32m'
YELLOW='\033[1;33m'
NC='\033[0m' # No Color

# Test counter
TESTS_RUN=0
TESTS_PASSED=0
TESTS_FAILED=0

# Helper functions
log_info() {
    echo -e "${YELLOW}[INFO]${NC} $1"
}

log_success() {
    echo -e "${GREEN}[PASS]${NC} $1"
    ((++TESTS_PASSED))
}

log_error() {
    echo -e "${RED}[FAIL]${NC} $1"
    ((++TESTS_FAILED))
}

run_test() {
    ((++TESTS_RUN))
    echo ""
    echo "=========================================="
    echo "Test $TESTS_RUN: $1"
    echo "=========================================="
}

# Check if GPU profile is enabled
check_gpu_enabled() {
    run_test "Check GPU profile enabled"

    if ! grep -q "^GPU_ENABLE=true" .env 2>/dev/null; then
        log_error "GPU_ENABLE not set to true in .env"
        echo "To enable GPU support:"
        echo "  1. Add 'GPU_ENABLE=true' to .env"
        echo "  2. Ensure nvidia-container-toolkit is installed on host"
        echo "  3. Run: make rebuild"
        return 1
    fi

    log_success "GPU_ENABLE=true found in .env"
    return 0
}

# Get the first gpu-worker container name
get_first_gpu_worker() {
    docker compose ps gpu-worker --format '{{.Names}}' 2>/dev/null | head -1
}

# Get the first gpu-worker node name from Slurm
get_first_gpu_node() {
    docker exec slurmctld scontrol show nodes 2>/dev/null | grep -o 'NodeName=g[0-9]*' | head -1 | cut -d= -f2
}

# Check gpu-worker container is running
check_gpu_worker_running() {
    run_test "Check gpu-worker container running"

    local gpu_worker
    gpu_worker=$(get_first_gpu_worker)

    if [[ -z "$gpu_worker" ]]; then
        log_error "No gpu-worker containers running"
        echo "Start GPU workers with: make up (with GPU_ENABLE=true in .env)"
        return 1
    fi

    local worker_count
    worker_count=$(docker compose ps gpu-worker --format '{{.Names}}' 2>/dev/null | wc -l)

    log_success "$worker_count gpu-worker container(s) running"
    return 0
}

# Check gpu-worker node registered in Slurm
check_gpu_worker_registered() {
    run_test "Check gpu-worker node registered in Slurm"

    local gpu_node
    gpu_node=$(get_first_gpu_node)

    if [[ -z "$gpu_node" ]]; then
        log_error "No gpu-worker nodes registered in Slurm"
        return 1
    fi

    log_success "gpu-worker node registered in Slurm: $gpu_node"
    return 0
}

# Check gpu-worker node has GPU GRES
check_gres_configured() {
    run_test "Check GPU GRES configured on gpu-worker"

    local gpu_node
    gpu_node=$(get_first_gpu_node)

    if [[ -z "$gpu_node" ]]; then
        log_error "No gpu-worker nodes found"
        return 1
    fi

    local gres_output
    gres_output=$(docker exec slurmctld scontrol show node "$gpu_node" | grep "Gres=" || true)

    if [[ ! "$gres_output" =~ gpu:nvidia ]]; then
        log_error "GPU GRES not configured correctly on $gpu_node"
        echo "Expected: Gres containing gpu:nvidia"
        echo "Got: $gres_output"
        return 1
    fi

    log_success "GPU GRES configured on $gpu_node: $gres_output"
    return 0
}

# Check GPU partition exists
check_gpu_partition() {
    run_test "Check 'gpu' partition exists"

    if ! docker exec slurmctld scontrol show partition gpu &>/dev/null; then
        log_error "'gpu' partition not found"
        return 1
    fi

    local partition_nodes
    partition_nodes=$(docker exec slurmctld scontrol show partition gpu | grep -E "^\s+Nodes=" | cut -d= -f2)

    if [[ -z "$partition_nodes" ]] || [[ "$partition_nodes" == "(null)" ]]; then
        log_error "GPU partition has no nodes"
        return 1
    fi

    log_success "'gpu' partition configured with nodes: $partition_nodes"
    return 0
}

# Check nvidia-smi works in gpu-worker container
check_nvidia_smi() {
    run_test "Check nvidia-smi works in gpu-worker container"

    local gpu_worker
    gpu_worker=$(get_first_gpu_worker)

    if [[ -z "$gpu_worker" ]]; then
        log_error "No gpu-worker container found"
        return 1
    fi

    if ! docker exec "$gpu_worker" nvidia-smi &>/dev/null; then
        log_error "nvidia-smi not available in $gpu_worker"
        echo "Ensure nvidia-container-toolkit is installed on host"
        return 1
    fi

    local gpu_count
    gpu_count=$(docker exec "$gpu_worker" nvidia-smi --query-gpu=count --format=csv,noheader | head -n1)

    log_success "nvidia-smi detected $gpu_count GPU(s)"
    docker exec "$gpu_worker" nvidia-smi
    return 0
}

# Test GPU job submission
test_gpu_job_submission() {
    run_test "Submit test job requesting GPU GRES"

    # Clean up old job outputs
    docker exec slurmctld bash -c "cd /data && rm -f slurm-*.out gpu-test-*.out" || true

    # Submit job requesting gpu:1
    local job_script='#!/bin/bash
#SBATCH --job-name=gpu-test
#SBATCH --partition=gpu
#SBATCH --gres=gpu:1
#SBATCH --output=gpu-test-%j.out

echo "=== GPU Test Job ==="
echo "Hostname: $(hostname)"
echo "CUDA_VISIBLE_DEVICES: ${CUDA_VISIBLE_DEVICES:-not set}"
echo ""
echo "=== nvidia-smi output ==="
nvidia-smi
echo ""
echo "=== GPU Detection ==="
if command -v nvidia-smi &> /dev/null; then
    nvidia-smi --query-gpu=index,name,driver_version,memory.total --format=csv
else
    echo "nvidia-smi not found"
fi
'

    local jobid
    jobid=$(docker exec slurmctld bash -c "cd /data && cat > gpu-job.sh <<'EOF'
$job_script
EOF
chmod +x gpu-job.sh && sbatch gpu-job.sh | grep -oP '\\d+'")

    if [[ -z "$jobid" ]]; then
        log_error "Failed to submit GPU job"
        return 1
    fi

    log_info "Job $jobid submitted, waiting for completion..."

    # Wait for job to complete (max 60 seconds)
    local elapsed=0
    while [[ $elapsed -lt 60 ]]; do
        local state
        state=$(docker exec slurmctld scontrol show job "$jobid" | grep "JobState=" | sed -n 's/.*JobState=\([A-Z]*\).*/\1/p' || echo "UNKNOWN")

        if [[ "$state" == "COMPLETED" ]]; then
            log_success "GPU job $jobid completed successfully"

            # Show job output
            echo ""
            echo "========== Job Output =========="
            docker exec slurmctld bash -c "cd /data && cat gpu-test-${jobid}.out"
            echo "================================"

            return 0
        elif [[ "$state" == "FAILED" ]] || [[ "$state" == "CANCELLED" ]] || [[ "$state" == "TIMEOUT" ]]; then
            log_error "GPU job $jobid failed with state: $state"
            docker exec slurmctld bash -c "cd /data && cat gpu-test-${jobid}.out" || echo "No output file"
            return 1
        fi

        sleep 2
        ((elapsed+=2))
    done

    log_error "GPU job $jobid timed out after 60 seconds"
    docker exec slurmctld scontrol show job "$jobid"
    return 1
}

# Test GPU allocation in squeue
test_gpu_allocation() {
    run_test "Verify GPU allocation shows in job"

    # Submit a held job to check allocation
    local jobid
    jobid=$(docker exec slurmctld bash -c "cd /data && sbatch --hold --partition=gpu --gres=gpu:1 --wrap='sleep 10' | grep -oP '\\d+'")

    if [[ -z "$jobid" ]]; then
        log_error "Failed to submit held GPU job"
        return 1
    fi

    local gres_alloc
    gres_alloc=$(docker exec slurmctld scontrol show job "$jobid" | grep "TresPerNode=" || echo "")

    # Cancel the held job
    docker exec slurmctld scancel "$jobid" || true

    if [[ "$gres_alloc" =~ gres/gpu ]]; then
        log_success "GPU allocation visible in job: $gres_alloc"
        return 0
    else
        log_error "GPU allocation not visible in job"
        echo "Expected TresPerNode to contain gres/gpu"
        echo "Got: $gres_alloc"
        return 1
    fi
}

# Main test execution
main() {
    echo ""
    echo "=========================================="
    echo "GPU Test Suite for slurm-docker-cluster"
    echo "=========================================="
    echo ""

    # Check prerequisites
    if ! check_gpu_enabled; then
        echo ""
        echo "GPU profile not enabled. Exiting."
        exit 1
    fi

    if ! check_gpu_worker_running; then
        echo ""
        echo "No gpu-worker containers running. Exiting."
        exit 1
    fi

    # Run tests (|| true so set -e doesn't abort on first failure; failures are tracked via TESTS_FAILED)
    check_gpu_worker_registered || true
    check_gres_configured || true
    check_gpu_partition || true
    check_nvidia_smi || true
    test_gpu_job_submission || true
    test_gpu_allocation || true

    # Summary
    echo ""
    echo "=========================================="
    echo "Test Summary"
    echo "=========================================="
    echo "Total tests: $TESTS_RUN"
    echo -e "${GREEN}Passed: $TESTS_PASSED${NC}"
    if [[ $TESTS_FAILED -gt 0 ]]; then
        echo -e "${RED}Failed: $TESTS_FAILED${NC}"
        exit 1
    else
        echo -e "${GREEN}All tests passed!${NC}"
        exit 0
    fi
}

main "$@"