spk921 / waitGPU

Small Python library to automatically set CUDA_VISIBLE_DEVICES to the least loaded device on multi-GPU systems.

waitGPU

A small Python library that waits for GPU conditions to be satisfied, and then setting CUDA_VISIBLE_DEVICES to the qualifying GPU on multi-GPU systems.

Installation

pip install waitGPU

Usage

Use this library before any import that will use a GPU like torch or tensorflow.

import waitGPU
waitGPU.wait(utilization=50, memory_ratio=0.5, available_memory=300,
             gpu_ids=[1,2], interval=10)

Specifying keyword arguments to wait will determine the criteria to wait for:

• utilization will wait until GPU utilization is at most the given value
• memory_ratio will wait until the GPU memory utilization is at most the given value
• available_memory will wait until the available memory is at least the given value
• gpu_ids will only consider GPUs with the given IDs
• interval is the number of seconds to wait before checking conditions

Dependencies

• Jongwook Choi's gpustat library.

Licensing

This code is in the public domain.