Debugging GKE GPU configuration

Troubleshooting GPU configuration on GKE can be painful. I explain how to troubleshoot GPU configuration issues in Google Kubernetes Engine (GKE) when running PyTorch workloads, emphasizing the importance of ensuring compatibility between Nvidia drivers, CUDA, and PyTorch. I outline a systematic approach to debug and verify that the GPU is recognized within the container and that PyTorch can access the GPU resources correctly.

gpumlopsgkenvidiapytorchkubernetes

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.

Loading diagram…

PyTorch installs most necessary dependencies (CUDA) [^1] during installation. A daemonset injects the Nvidia drivers when a workload is scheduled. 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.

If you’ve used Docker to build your own container, you must install the nvidia container toolkit. This toolkit ensures your container is compatible with Nvidia GPUs.

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()