GPU not detected in Docker

[rdwight]

I'm closely following the easy version of installing with CUDA, which amounts to just a pip install -U jax[cuda12]. I'm running this in a Docker container utilizing the NVIDIA Container Toolkit, and starting the container with the nvidia runtime, so that the GPUs are available inside the container. This is with the latest versions of everything, from drivers up. Observations:

1. The pip call is successful but jax.devices() reports CPU only.
2. If I install pytorch in the same container, the GPUs are detected and usable.
3. If I install JAX in the base OS the GPUs are detected and usable.

Running in the base OS is fine for me for now, but not ideal - and I thought you guys might be interested in this issue, and maybe have some tips? Some general tips for debugging non-detected GPUs would be helpful too.

A side comment, on the installation page it states: "The JAX team strongly recommends installing CUDA and cuDNN using the pip wheels, since it is much easier!". It would be helpful if it was made clear that this is done automatically as dependencies of JAX, and no prior CUDA installation is required. This was not obvious to me.

There follows my system information, and the Dockerfile. Please let me know if you need anything else.
Richard

# uname -a
Linux Bruteforce 5.14.0-1056-oem #63-Ubuntu SMP Fri Dec 16 14:32:59 UTC 2022 x86_64 x86_64 x86_64 GNU/Linux

# nvidia-smi  #  (this works from inside and outside the container)
+-----------------------------------------------------------------------------------------+
| NVIDIA-SMI 560.35.03              Driver Version: 560.35.03      CUDA Version: 12.6     |
|-----------------------------------------+------------------------+----------------------+
| GPU  Name                 Persistence-M | Bus-Id          Disp.A | Volatile Uncorr. ECC |
| Fan  Temp   Perf          Pwr:Usage/Cap |           Memory-Usage | GPU-Util  Compute M. |
|                                         |                        |               MIG M. |
|=========================================+========================+======================|
|   0  Quadro P620                    Off |   00000000:65:00.0  On |                  N/A |
| 38%   51C    P0             N/A /  N/A  |     525MiB /   2048MiB |      1%      Default |
|                                         |                        |                  N/A |
+-----------------------------------------+------------------------+----------------------+
|   1  NVIDIA TITAN V                 Off |   00000000:B3:00.0 Off |                  N/A |
| 32%   46C    P8             29W /  250W |       6MiB /  12288MiB |      0%      Default |
|                                         |                        |                  N/A |
+-----------------------------------------+------------------------+----------------------+

# cat /proc/driver/nvidia/version 
NVRM version: NVIDIA UNIX x86_64 Kernel Module  560.35.03  Fri Aug 16 21:39:15 UTC 2024
GCC version:  gcc version 11.4.0 (Ubuntu 11.4.0-1ubuntu1~22.04) 

# docker -v
Docker version 27.0.3, build 7d4bcd8

My Dockerfile:

# syntax=docker/dockerfile:1

FROM ubuntu:latest
RUN useradd -ms /bin/bash jax
WORKDIR /home/jax

USER root
RUN apt-get update && apt-get install -y wget git
USER jax

# Miniconda
ENV PATH="~/miniconda3/bin:$PATH"
ARG PATH="~/miniconda3/bin:$PATH"
RUN mkdir -p miniconda3 \
    && wget https://repo.anaconda.com/miniconda/Miniconda3-latest-Linux-x86_64.sh \
            -O ~/miniconda3/miniconda.sh \
    && bash ~/miniconda3/miniconda.sh -b -u -p ~/miniconda3 \
    && rm -f ~/miniconda3/miniconda.sh \
    && ~/miniconda3/bin/conda init bash \
    && ~/miniconda3/bin/pip install --upgrade pip
RUN ~/miniconda3/bin/conda install -y ipython

# JAX
RUN ~/miniconda3/bin/pip install --upgrade "jax[cuda12]"
RUN ~/miniconda3/bin/python -c "import jax; print(jax.devices())"

Which I executed with:

docker build -t jax-v007 .
docker run -it --runtime=nvidia --gpus '"device=1"' jax-v007    # Just using the one GPU

[bmaxdk]

can you try:

docker run --gpus all --runtime=nvidia jax-v007

check this configuring docker

[rdwight]

Thanks, that worked!:

docker run -it --gpus all --runtime=nvidia jax-v007
[CudaDevice(id=0), CudaDevice(id=1)]

So it appears that just the --gpus '"device=1"' functionality of docker doesn't work with JAX. So problem solved, still any insight into why would be appreciated!

[yashk2810]

Cam you try with the latest jax and jaxlib release?

[rdwight]

This was with jax-0.4.31 jaxlib-0.4.31, which are the latest afaik.