Measure GPU memory usage

measure-gpu-memory-usage.py

#!/usr/bin/env python3
# Usage: measure-gpu-memory-usage.py command args...

import cupy
import subprocess
import sys

initial_memory = None
for device_id in range(cupy.cuda.runtime.getDeviceCount()):
    cupy.cuda.runtime.setDevice(device_id)
    free, _ = cupy.cuda.runtime.memGetInfo()
    if initial_memory is None or free < initial_memory:
        initial_memory = free

sys.stderr.write("Initial memory size: %d bytes\n" % initial_memory)

def allocate(target):
    pool = []
    for device_id in range(cupy.cuda.runtime.getDeviceCount()):
        cupy.cuda.runtime.setDevice(device_id)
        while True:
            free, _ = cupy.cuda.runtime.memGetInfo()
            if free <= target:
                break
            pool.append(cupy.zeros(((free - target) // 32 + 1, ), dtype="f"))
    return pool

def release():
    for device_id in range(cupy.cuda.runtime.getDeviceCount()):
        cupy.cuda.runtime.setDevice(device_id)
        cupy.get_default_memory_pool().free_all_blocks()

def run():
    result = subprocess.run(
        sys.argv[1:], stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)
    return result.returncode == 0

def main():
    if not run():
        print("Failed to execute the command")
        sys.exit(1)
    max_memory = initial_memory
    min_memory = 0
    for _ in range(20):
        if max_memory < 16 * 1024 * 1024:
            print("0 MB")
            sys.exit(0)
        pool = None
        release()
        target = (max_memory + min_memory) // 2
        sys.stderr.write("Testing with %d bytes...\n" % target)
        pool = allocate(target)
        if run():
            max_memory = target
        else:
            min_memory = target
    print("%d MB" % (max_memory // 1024 // 1024 + 1))

main()