Package: safe-gpu
safe-gpuA module for safe acquisition of GPUs in exclusive mode.
Relevant mainly in clusters with a purely declarative gpu resource, such as many versions of SGE.Features:toolkit independence (PyTorch/TensorFlow/pycuda/...), this just setsCUDA_VISIBLE_DEVICESproperlyincluded support for PyTorch and TensorFlow2 backends, open to othersmultiple GPUs acquisitionworkaround for machines with a single GPU used for display and computation alikeopen to implementation in different languagesDownsides:in order to really prevent the race condition, everyone on your cluster has to use thisInstalationIn addition to manual installation,safe-gpuis on PyPi, so you can simply:pip install safe-gpuNote thatsafe-gpudoes not formally depend on any backend, giving you, the user, the freedom to pick one of your liking.UsagePrior to initializing CUDA (typically happens in lazy fashion when you place something on GPU), callclaim_gpus.from safe_gpu import safe_gpu

safe_gpu.claim_gpus()If you want multiple GPUs, pass the desired number toclaim_gpus:safe_gpu.claim_gpus(nb_gpus)Internally,claim_gpus()constructs aGPUOwnerand stores it insafe_gpu.gpu_owner.
If preferred, user code can constructGPUOwneritself, but care should be taken to keep it alive until actual data is placed on said GPUs.Usage with HorovodTypical Horovod usage includes starting your script in several processes, one per GPU.
Therefore, only ask for one GPU in each process:safe_gpu.claim_gpus()  # 1 GPU is the default, can be ommited
hvd.init()Common errorsIn order to properly setup GPUs for your process,claim_gpusreally needs be called before CUDA is initialized.
When CUDA does get initialized, it fixes your logical devices (e.g. PyTorchcuda:1etc.) to actual GPUs in your system.
IfCUDA_VISIBLE_DEVICESare not set at that moment, CUDA will happily offer your process all of the visible GPUs, including those already occupied.Most commonly, this issue occurs for users who try to play it safe and check CUDA availability beforehand:if torch.cuda.is_available():  # This already initializes CUDA
  safe_gpu.claim_gpus(nb_gpus)  # So this can fail easilyIf your workflow mandates on-the-fly checking of GPU availability, instead use:try:
  safe_gpu.claim_gpus(nb_gpus)
except safe_gpu.NvidiasmiError:
  ...Also, horovod users can be at risk:hvd.init()
torch.cuda.set_device(hvd.local_rank())  # This initializes CUDA, too
safe_gpu.claim_gpus()  # Thus this is likely to failSee above for proper solution.Other backendsThe default implementation uses a PyTorch tensor to claim a GPU.
Additionally, a TensorFlow2 placeholder is provided assafe_gpu.tensorflow_placeholder.If you don't want to / can't use that, provide your own GPU memory allocating function asclaim_gpus's parameterplaceholder_fn.
It has to accept one parameterdevice_no, occupy a (preferably negligible) piece of memory on that device, and return a pointer to it.Pull requests for other backends are welcome.Checking that it worksTogether with this package, a small testing script is provided.
It exaggerates the time needed to acquire the GPU after polling nvidia-smi, making the race condition technically sure to happen.To run the following example, get to a machine with 3 free GPUs and run two instances of the script in parallel as shown.
You should see in the output that one of them really waited for the faster one to fully acquire the GPU.This script is not distributed along in the pip package, so please download it separately if needed.$ python3 gpu-acquisitor.py --backend pytorch --id 1 --nb-gpus 1 & python3 gpu-acquisitor.py --backend pytorch --id 2 --nb-gpus 2
GPUOwner1 2020-11-30 14:29:33,315 [INFO] acquiring lock
GPUOwner1 2020-11-30 14:29:33,315 [INFO] lock acquired
GPUOwner2 2020-11-30 14:29:33,361 [INFO] acquiring lock
GPUOwner1 2020-11-30 14:29:34,855 [INFO] Set CUDA_VISIBLE_DEVICES=2
GPUOwner2 2020-11-30 14:29:45,447 [INFO] lock acquired
GPUOwner1 2020-11-30 14:29:45,447 [INFO] lock released
GPUOwner2 2020-11-30 14:29:48,926 [INFO] Set CUDA_VISIBLE_DEVICES=4,5
GPUOwner1 2020-11-30 14:29:54,492 [INFO] Finished
GPUOwner2 2020-11-30 14:30:00,525 [INFO] lock released
GPUOwner2 2020-11-30 14:30:09,571 [INFO] Finished