Jaxlib error due to convolutions

[rdilip]

I'm having an issue where jaxlib exits with the following errors: jaxlib.xla_extension.XlaRuntimeError: UNKNOWN: Failed to determine best cudnn convolution algorithm for... and Original error: UNIMPLEMENTED: DNN library is not found. I've seen this error in previous issues with a variety of causes, such as the cudnn version and the GPU memory. Is there a current recommendation on how to resolve this? I'm running everything inside a docker container (Dockerfile below) on Google cloud, so I'm hoping it's possible to resolve this, as I haven't had any luck yet.

FROM tensorflow/tensorflow:2.9.0-gpu

# add the user's .local/bin to path so that our --user pip installed packages will appear
ENV PATH=/home/jax/.local/bin:$PATH

WORKDIR /notebooks
RUN pip3 install --upgrade "jax[cuda]" -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html
RUN pip3 install "jax[cuda11_cudnn82]" -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html

RUN apt-get update && apt-get install -y \
    git && \
    rm -rf /var/lib/apt/lists/*

RUN apt install --allow-change-held-packages libcudnn8=8.1.0.77-1+cuda11.2

COPY requirements.txt ./

RUN pip install --no-cache-dir -r requirements.txt

COPY . .

CMD ["jupyter", "lab", "--ip=0.0.0.0", "--allow-root"]

[cgarciae]

This is tricky to debug and should probably go to the JAX repo, but if it helps this is what I used recently:

FROM nvidia/cuda:11.7.1-cudnn8-devel-ubuntu22.04

RUN rm -f /etc/apt/sources.list.d/cuda.list && rm -f /etc/apt/sources.list.d/nvidia-ml.list
RUN apt-get update
RUN apt-get install -y python3 python3-pip python-is-python3 python3-tk git
RUN pip install jax[cuda] -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html

# test runtime
RUN python -c "import jax; print(jax.devices())"

[rdilip]

This works! Many thanks, much appreciated. I know NVIDIA is releasing official images, so hopefully this all becomes moot soon.

In case others have this issue, I had to make one small adjustment to the script above, since configuring tzdata on a docker container will cause the image building to hang. I added two lines to the Dockerfile suggested by cgarciae to fix this.

FROM nvidia/cuda:11.7.1-cudnn8-devel-ubuntu22.04
ENV TZ=America/Los_Angeles
RUN ln -snf /usr/share/zoneinfo/$TZ /etc/localtime && echo $TZ > /etc/timezone

RUN rm -f /etc/apt/sources.list.d/cuda.list && rm -f /etc/apt/sources.list.d/nvidia-ml.list
RUN apt-get update
RUN apt-get install -y python3 python3-pip python-is-python3 python3-tk git
RUN pip install jax[cuda] -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html

# test runtime
RUN python -c "import jax; print(jax.devices())"

[cgarciae]

Ah right, sorry about that, I shared is a stripped down version of my actual docker file, I removed the TZ stuff because I didn't remember why I added it 😅