ARG BASE_TAG=staging

FROM nvidia/cuda:10.2-cudnn7-devel-ubuntu18.04 AS nvidia

FROM gcr.io/kaggle-images/python-tensorflow-whl:2.4.1-py37 as tensorflow_whl

FROM gcr.io/kaggle-images/python:${BASE_TAG}

ADD clean-layer.sh /tmp/clean-layer.sh

COPY --from=nvidia /etc/apt/sources.list.d/cuda.list /etc/apt/sources.list.d/

COPY --from=nvidia /etc/apt/sources.list.d/nvidia-ml.list /etc/apt/sources.list.d/

COPY --from=nvidia /etc/apt/trusted.gpg /etc/apt/trusted.gpg.d/cuda.gpg

RUN sed -i 's/deb https:\/\/developer.download.nvidia.com/deb http:\/\/developer.download.nvidia.com/' /etc/apt/sources.list.d/*.list

ENV CUDA_MAJOR_VERSION=10

ENV CUDA_MINOR_VERSION=2

ENV CUDA_PATCH_VERSION=89

ENV CUDA_VERSION=$CUDA_MAJOR_VERSION.$CUDA_MINOR_VERSION.$CUDA_PATCH_VERSION

ENV CUDA_PKG_VERSION=$CUDA_MAJOR_VERSION-$CUDA_MINOR_VERSION=$CUDA_VERSION-1

LABEL com.nvidia.volumes.needed="nvidia_driver"

LABEL com.nvidia.cuda.version="${CUDA_VERSION}"

ENV PATH=/usr/local/nvidia/bin:/usr/local/cuda/bin:/opt/bin:${PATH}

ENV LD_LIBRARY_PATH_NO_STUBS="/usr/local/nvidia/lib64:/usr/local/cuda/lib64:$LD_LIBRARY_PATH"

ENV LD_LIBRARY_PATH="/usr/local/nvidia/lib64:/usr/local/cuda/lib64:/usr/local/cuda/lib64/stubs:$LD_LIBRARY_PATH"

ENV NVIDIA_VISIBLE_DEVICES=all

ENV NVIDIA_DRIVER_CAPABILITIES=compute,utility

ENV NVIDIA_REQUIRE_CUDA="cuda>=$CUDA_MAJOR_VERSION.$CUDA_MINOR_VERSION"

RUN apt-get update && apt-get install -y --no-install-recommends \

cuda-cupti-$CUDA_PKG_VERSION \

cuda-cudart-$CUDA_PKG_VERSION \

cuda-cudart-dev-$CUDA_PKG_VERSION \

cuda-libraries-$CUDA_PKG_VERSION \

cuda-libraries-dev-$CUDA_PKG_VERSION \

cuda-nvml-dev-$CUDA_PKG_VERSION \

cuda-minimal-build-$CUDA_PKG_VERSION \

cuda-command-line-tools-$CUDA_PKG_VERSION \

libcudnn7=7.6.5.32-1+cuda$CUDA_MAJOR_VERSION.$CUDA_MINOR_VERSION \

libcudnn7-dev=7.6.5.32-1+cuda$CUDA_MAJOR_VERSION.$CUDA_MINOR_VERSION \

libnccl2=2.5.6-1+cuda$CUDA_MAJOR_VERSION.$CUDA_MINOR_VERSION \

libnccl-dev=2.5.6-1+cuda$CUDA_MAJOR_VERSION.$CUDA_MINOR_VERSION && \

ln -s /usr/local/cuda-$CUDA_MAJOR_VERSION.$CUDA_MINOR_VERSION /usr/local/cuda && \

ln -s /usr/local/cuda/lib64/stubs/libcuda.so /usr/local/cuda/lib64/stubs/libcuda.so.1 && \

/tmp/clean-layer.sh

RUN apt-get install -y ocl-icd-libopencl1 clinfo libboost-all-dev && \

mkdir -p /etc/OpenCL/vendors && \

echo "libnvidia-opencl.so.1" > /etc/OpenCL/vendors/nvidia.icd && \

/tmp/clean-layer.sh

RUN conda remove --force -y pytorch torchvision torchaudio cpuonly && \

conda install "pytorch=1.7" "torchvision=0.8" "torchaudio=0.7" "torchtext=0.8" "cudf=0.16" "cuml=0.16" cudatoolkit=$CUDA_VERSION && \

/tmp/clean-layer.sh

RUN pip uninstall -y lightgbm && \

cd /usr/local/src && \

git clone --recursive https://github.com/microsoft/LightGBM && \

cd LightGBM && \

git checkout tags/v3.1.1 && \

mkdir build && cd build && \

cmake -DUSE_GPU=1 -DOpenCL_LIBRARY=/usr/local/cuda/lib64/libOpenCL.so -DOpenCL_INCLUDE_DIR=/usr/local/cuda/include/ .. && \

make -j$(nproc) && \

cd /usr/local/src/LightGBM/python-package && \

python setup.py install --precompile && \

mkdir -p /etc/OpenCL/vendors && \

echo "libnvidia-opencl.so.1" > /etc/OpenCL/vendors/nvidia.icd && \

/tmp/clean-layer.sh

RUN pip install jax==0.2.6 jaxlib==0.1.57+cuda$CUDA_MAJOR_VERSION$CUDA_MINOR_VERSION -f https://storage.googleapis.com/jax-releases/jax_releases.html && \

/tmp/clean-layer.sh

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 && \

pip uninstall -y mxnet && \

pip install mxnet-cu$CUDA_MAJOR_VERSION$CUDA_MINOR_VERSION && \

/tmp/clean-layer.sh

# Reinstall TensorFlow addons (TFA) with GPU support.

COPY --from=tensorflow_whl /tmp/tfa_gpu/*.whl /tmp/tfa_gpu/

RUN pip install /tmp/tfa_gpu/tensorflow*.whl && \

rm -rf /tmp/tfa_gpu/ && \

/tmp/clean-layer.sh

RUN pip install pycuda && \

pip install cupy-cuda$CUDA_MAJOR_VERSION$CUDA_MINOR_VERSION && \

pip install pynvrtc && \

# b/175638062 remove pin once we update to cuDNN 8.x

pip install nnabla-ext-cuda$CUDA_MAJOR_VERSION$CUDA_MINOR_VERSION==1.13.0 && \

/tmp/clean-layer.sh

ENV LD_LIBRARY_PATH="$LD_LIBRARY_PATH_NO_STUBS"