Installation issue (JAX-GPU)

[EngineerKhan]

I have installed JAX (GPU version) in Docker (running on GPU machine with CUDA installed) and Docker was built successfully.

Here's the relevant Docker part:

#Install JAX GPU
RUN pip install --upgrade pip &&\
	pip install --upgrade "jax[cuda]" -f https://storage.googleapis.com/jax-releases/jax_releases.html 

# Install JAX Ecosystem	

RUN pip install git+https://github.com/deepmind/dm-haiku && pip install flax && pip install optax

Now I am running the JAX code and its giving me error

2021-12-28 15:00:57.693321: W external/org_tensorflow/tensorflow/stream_executor/platform/default/dso_loader.cc:64] Could not load dynamic library 'libcudart.so.11.0'; dlerror: libcudart.so.11.0: cannot open shared object file: No such file or directory

Can anyone please help there? CUDA version installed is 11.5

[hawkinsp]

Where is CUDA installed inside your Docker container? (You need CUDA inside the docker container.)

Since this is in docker, if that doesn't help, can you share a complete Dockerfile that we can use to reproduce?

[EngineerKhan]

It was my bad to assume CUDA doesn't need to be installed inside Docker.

[EngineerKhan]

I am getting this message (after CUDA installation, followed by JAX GPU installation - both in Docker) now:

2021-12-30 03:59:10.465166: E external/org_tensorflow/tensorflow/stream_executor/cuda/cuda_driver.cc:271] failed call to cuInit: CUDA_ERROR_NO_DEVICE: no CUDA-capable device is detected

Here's Docker file:

# Latest version of ubuntu as base image
FROM    ubuntu:latest

# Install the following packages
RUN     apt-get update && apt-get install software-properties-common -y &&\
        add-apt-repository ppa:deadsnakes/ppa && apt-get update &&\
        apt-get install python3.9 -y && apt install python3-pip -y &&\
        apt-get update && apt-get install -y git &&\
        pip3 install --upgrade pip && pip3 install jupyter &&\
        mkdir /usr/local/notebooks &&\
        rm ~/.cache/pip -rf

# Install any Python modules required in the notebook
RUN     pip3 install pandas && pip3 install numpy && pip3 install seaborn &&\
        pip3 install matplotlib && pip3 install --upgrade sklearn && pip3 install --upgrade scikit-learn &&\
        pip3 install jupyterlab &&\ 
	pip3 install config &&\
        pip3 install librosa &&\
        pip3 install tensorboard &&\
        pip3 install jupyter-server-proxy && jupyter serverextension enable --sys-prefix jupyter_server_proxy 

RUN apt-get install -y wget

#Install CUDA (JAX guy asked me to install CUDA in Docker as well)
RUN 	wget https://developer.download.nvidia.com/compute/cuda/repos/ubuntu2004/x86_64/cuda-ubuntu2004.pin 
RUN	mv cuda-ubuntu2004.pin /etc/apt/preferences.d/cuda-repository-pin-600 
RUN	apt-key adv --fetch-keys https://developer.download.nvidia.com/compute/cuda/repos/ubuntu2004/x86_64/7fa2af80.pub 
RUN	add-apt-repository "deb https://developer.download.nvidia.com/compute/cuda/repos/ubuntu2004/x86_64/ /" 
RUN	apt-get update
RUN 	apt-get -y install cuda

#Install JAX GPU
RUN pip install --upgrade pip &&\
	pip install --upgrade "jax[cuda]" -f https://storage.googleapis.com/jax-releases/jax_releases.html 

# Install JAX Ecosystem	

RUN pip install git+https://github.com/deepmind/dm-haiku && pip install flax && pip install optax

[hawkinsp]

My guess is this isn't related to JAX at this point, instead it's related to the Docker container having access to the GPU at all.

Does nvidia-smi, run within the Docker container, report the GPU being present? I think you may need to use nvidia-docker (https://github.com/NVIDIA/nvidia-docker), although I'm not a docker expert.

[EngineerKhan]

Yeah it can be as nvidia-smi doesn't report the GPU being present. I am also switching to trying out the nvidia's own docker image as base.

There's one thing related to JAX which I would like to share though is that when I open the terminal (after running the aforementioned docker), run python in the shell and run

import jax
jax.devices()

It does show me the GPU device (even without nvidia-smi reporting the GPU due to lack of .bashrc profiling) and this is something weird as when I run the code through JuPyter, etc. it gives this error, but within the terminal this way its showing GPU. This is something which maybe of interest for you for further analysis.

[EngineerKhan]