pyhf Docker images built on the NVIDIA CUDA enabled images for runtime use with the the NVIDIA Container Toolkit.
These images are meant to be used as base images for bespoke user application images to build on.
- Make sure that you have the
nvidia-container-toolkitinstalled on the host machine - Check the list of available tags on Docker Hub to find the tag you want
- Use
docker pullto pull down the image corresponding to the tag
Example:
docker pull pyhf/cuda:0.6.3-jax-cuda-11.1
To check that NVIDIA GPUS are being properly detected run
docker run --rm --gpus all pyhf/cuda:0.6.3-jax-cuda-11.1 'nvidia-smi'
and check if the nvidia-smi output appears correctly.
To run (interactively) using GPUs on the host machine:
docker run --rm -ti --gpus all pyhf/cuda:0.6.3-jax-cuda-11.1
To verify things are working on your host machine you can run
docker run --rm --gpus all -v $PWD:$PWD -w $PWD pyhf/cuda:0.6.3-jax-cuda-11.1 'bash tests/test_backend.sh jax'