Debugging GKE GPU configuration

Running PyTorch workloads on GKE GPU nodepools is straightforward once you know what to look for. The diagram below shows the high-level dependencies that must align.

PyTorch installs most necessary dependencies (CUDA) ¹ during installation. A daemonset installs the Nvidia drivers when each GPU node starts up, before any workload is scheduled on it. You can configure the drivers as DEFAULT or LATEST when creating the nodepool, or manually install the daemonset. Finally, you request the GPU and its quantity when creating the nodepool.

Debug GPU configuration backwards, starting from the GPU itself. You will see an error like below if there is a compatibility issue between the nvidia drivers and the Python (PyTorch) code. The driver enables your code to run on the hardware. Forcing the nodepool to install LATEST drivers usually fixes this problem, but you might have to manually install the correct drivers on the daemonset.

main ERROR: Error processing dataset: The NVIDIA driver on your system is too old
(found version 11040). Please update your GPU driver by downloading and installing
 a new version from the URL: http://www.nvidia.com/Download/index.aspx
 Alternatively, go to: https://pytorch.org to install a PyTorch version that has
 been compiled with your version of the CUDA driver.

You do not need the nvidia container toolkit in your own image on GKE. The driver installer daemonset and the GPU device plugin mount the drivers into the container at /usr/local/nvidia, so make sure /usr/local/nvidia/lib64 is on your LD_LIBRARY_PATH (the official CUDA and PyTorch images already handle this).

Once the drivers are configured on the nodepool, verify the container can access the GPU. Run nvidia-smi on the command line and you should see the GPU(s) and associated information. If that command fails, you haven’t configured it correctly, so verify container compatibility and check whether any GPUs are assigned to that nodepool.

After the interface between the container and the GPU(s) is confirmed then check whether PyTorch is installed and can connect to the GPUs. Exec in and run the below commands on the Python terminal. If a number greater than zero is returned then you should be able to run your workload as expected - so long as you were hoping to use the number of GPUs that were found! Otherwise, you’ll need to check Python version and the software dependencies it was installed with.

>>> import torch
>>> torch.cuda.device_count()