Measure maximum VRAM consumption of a program on linux for Nvidia devices

gpumeter.sh

#!/bin/bash
# Usage ./gpumeter.sh my-gpu-hungry-script.sh


# Check if the command argument is provided
if [ "$#" -ne 1 ]; then
    echo "Usage: $0 <script_to_run>"
    exit 1
fi

# Start monitoring GPU memory usage in the background
nvidia-smi --query-gpu=memory.used --format=csv,noheader,nounits -l 1 > gpu_usage.log &
PID=$!

# Execute the provided script
./$1

# Kill the nvidia-smi monitoring process
kill $PID

# Calculate and display the max GPU memory usage
MAX_USAGE=$(awk '{ if(max < $1) { max=$1 } } END { print max }' gpu_usage.log)
echo "Maximum GPU VRAM usage: ${MAX_USAGE}MiB"