track_gpu_usage.py

"""
Module for continuously printing the GPU usage after a set interval, until
it is terminated.
"""

import time
import GPUtil

# Set the interval for printing GPU memory (in seconds)
INTERVAL = 2

def print_gpu_memory():
    try:
        # Get the list of available GPUs
        gpus = GPUtil.getGPUs()

        for i, gpu in enumerate(gpus):
            print(f"GPU {i + 1}: {gpu.name}")
            print(f"    Memory Free: {gpu.memoryFree:.2f} MB | Memory Used: {gpu.memoryUsed:.2f} MB | Memory Total: {gpu.memoryTotal:.2f} MB")
        
        print("\n-----------------------------------------\n")

    except Exception as e:
        print(f"Error while getting GPU information: {e}")


def main():
    try:
        while True:
            print_gpu_memory()
            time.sleep(INTERVAL)

    except KeyboardInterrupt:
        print("Monitoring stopped by the user.")

if __name__ == "__main__":
    main()