Merge pull request #782 from Kaggle/pin-jaxlib

Pin JAX and jaxlib

Author: rosbo

Dockerfile

@@ -39,9 +39,6 @@ ENV PROJ_LIB=/opt/conda/share/proj
 RUN conda install -c conda-forge matplotlib basemap cartopy python-igraph imagemagick pysal && \
     # b/142337634#comment22 pin required to avoid torchaudio downgrade.
     conda install -c pytorch pytorch torchvision "torchaudio>=0.4.0" cpuonly && \
-    # pixman 0.38.0 has a known issue causing issues with packages such as pycairo & openslide.
-    # See: https://gitlab.freedesktop.org/pixman/pixman/commit/8256c235d9b3854d039242356905eca854a890ba
-    conda install -c conda-forge --no-deps pixman==0.34 && \
     /tmp/clean-layer.sh
 
 # The anaconda base image includes outdated versions of these packages. Update them to include the latest version.

gpu.Dockerfile

@@ -56,6 +56,14 @@ RUN apt-get install -y ocl-icd-libopencl1 clinfo libboost-all-dev && \
     echo "libnvidia-opencl.so.1" > /etc/OpenCL/vendors/nvidia.icd && \
     /tmp/clean-layer.sh
 
+# When using pip in a conda environment, conda commands should be ran first and then
+# the remaining pip commands: https://www.anaconda.com/using-pip-in-a-conda-environment/
+# However, because this image is based on the CPU image, this isn't possible but better
+# to put them at the top of this file to minize conflicts.
+RUN conda remove --force -y pytorch torchvision torchaudio cpuonly && \
+    conda install -y pytorch torchvision torchaudio cudatoolkit=$CUDA_MAJOR_VERSION.$CUDA_MINOR_VERSION -c pytorch && \
+    /tmp/clean-layer.sh
+
 # Install LightGBM with GPU
 RUN pip uninstall -y lightgbm && \
     cd /usr/local/src && \
@@ -72,21 +80,22 @@ RUN pip uninstall -y lightgbm && \
     /tmp/clean-layer.sh
 
 # Install JAX
+# b/154150582#comment9: JAX 0.1.63 with jaxlib 0.1.43 is causing the GPU tests to hang.
+ENV JAX_VERSION=0.1.62
+ENV JAXLIB_VERSION=0.1.41
 ENV JAX_PYTHON_VERSION=cp37
 ENV JAX_CUDA_VERSION=cuda$CUDA_MAJOR_VERSION$CUDA_MINOR_VERSION
 ENV JAX_PLATFORM=linux_x86_64
 ENV JAX_BASE_URL="https://storage.googleapis.com/jax-releases"
 
-RUN  pip install --upgrade $JAX_BASE_URL/$JAX_CUDA_VERSION/jaxlib-0.1.43-$JAX_PYTHON_VERSION-none-$JAX_PLATFORM.whl && \
-     pip install --upgrade jax
+RUN  pip install $JAX_BASE_URL/$JAX_CUDA_VERSION/jaxlib-$JAXLIB_VERSION-$JAX_PYTHON_VERSION-none-$JAX_PLATFORM.whl && \
+     pip install jax==$JAX_VERSION
 
 # Reinstall packages with a separate version for GPU support.
 COPY --from=tensorflow_whl /tmp/tensorflow_gpu/*.whl /tmp/tensorflow_gpu/
 RUN pip uninstall -y tensorflow && \
     pip install /tmp/tensorflow_gpu/tensorflow*.whl && \
     rm -rf /tmp/tensorflow_gpu && \
-    conda remove --force -y pytorch torchvision torchaudio cpuonly && \
-    conda install -y pytorch torchvision torchaudio cudatoolkit=$CUDA_MAJOR_VERSION.$CUDA_MINOR_VERSION -c pytorch && \
     pip uninstall -y mxnet && \
     # b/126259508 --no-deps prevents numpy from being downgraded.
     pip install --no-deps mxnet-cu$CUDA_MAJOR_VERSION$CUDA_MINOR_VERSION && \