Small script to check the availability of GPU

check_gpu_availability.py

import subprocess

def check_gpu_availability(threshold=50):
    # Run the nvidia-smi command to get GPU information
    cmd = "nvidia-smi --query-gpu=index,memory.used,memory.total --format=csv,noheader"
    gpu_info = subprocess.check_output(cmd, shell=True).decode().strip().split("\n")

    # Parse the GPU information into a list of dictionaries
    gpu_list = []
    for gpu in gpu_info:
        index, mem_used, mem_total = gpu.split(",")
        gpu_list.append({
            "index": int(index.strip()),
            "memory_used": int(mem_used.strip().split()[0]),
            "memory_total": int(mem_total.strip().split()[0])
        })

    # Filter the GPUs by memory usage and availability
    available_gpus = []
    for gpu in gpu_list:
        if gpu["memory_used"] / gpu["memory_total"] * 100 < threshold:
            available_gpus.append(gpu)

    # Display the available GPUs in a table
    if len(available_gpus) > 0:
        print("Available GPUs:")
        print("--------------")
        print("| Index | Memory Used | Memory Total |")
        print("|-------|-------------|--------------|")
        for gpu in available_gpus:
            print(f"| {gpu['index']:5d} | {gpu['memory_used']:11d} | {gpu['memory_total']:12d} |")
    else:
        print("No available GPUs found.")

# Example usage
check_gpu_availability(threshold=50)