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, nproc=1, ngpu=1)
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 -
nproc
will wait until the number of processes running on each GPU is at most the given value -
ngpu
will wait until the number of GPUs that satisfy all other criteria is at least the given value (default: 1)
Dependencies
- Jongwook Choi's gpustat library.
Licensing
This code is in the public domain.