ARG BASE_IMAGE_REPO

ARG BASE_IMAGE_TAG

ARG CPU_BASE_IMAGE_NAME

ARG GPU_BASE_IMAGE_NAME

ARG LIGHTGBM_VERSION

ARG TORCH_VERSION

ARG TORCHAUDIO_VERSION

ARG TORCHTEXT_VERSION

ARG TORCHVISION_VERSION

{{ if eq .Accelerator "gpu" }}

FROM gcr.io/kaggle-images/python-lightgbm-whl:${GPU_BASE_IMAGE_NAME}-${BASE_IMAGE_TAG}-${LIGHTGBM_VERSION} AS lightgbm_whl

FROM gcr.io/kaggle-images/python-torch-whl:${GPU_BASE_IMAGE_NAME}-${BASE_IMAGE_TAG}-${TORCH_VERSION} AS torch_whl

FROM ${BASE_IMAGE_REPO}/${GPU_BASE_IMAGE_NAME}:${BASE_IMAGE_TAG}

ARG CUDA_MAJOR_VERSION

ARG CUDA_MINOR_VERSION

ENV CUDA_MAJOR_VERSION=${CUDA_MAJOR_VERSION}

ENV CUDA_MINOR_VERSION=${CUDA_MINOR_VERSION}

# NVIDIA binaries from the host are mounted to /opt/bin.

ENV PATH=/opt/bin:${PATH}

# Add CUDA stubs to LD_LIBRARY_PATH to support building the GPU image on a CPU machine.

ENV LD_LIBRARY_PATH_NO_STUBS="$LD_LIBRARY_PATH"

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

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

{{ else }}

FROM ${BASE_IMAGE_REPO}/${CPU_BASE_IMAGE_NAME}:${BASE_IMAGE_TAG}

{{ end }}

# Keep these variables in sync if base image is updated.

ENV TENSORFLOW_VERSION=2.6.4

# We need to redefine the ARG here to get the ARG value defined above the FROM instruction.

# See: https://docs.docker.com/engine/reference/builder/#understand-how-arg-and-from-interact

ARG LIGHTGBM_VERSION

ARG TORCH_VERSION

ARG TORCHAUDIO_VERSION

ARG TORCHTEXT_VERSION

ARG TORCHVISION_VERSION

# Disable pesky logs like: KMP_AFFINITY: pid 6121 tid 6121 thread 0 bound to OS proc set 0

# See: https://stackoverflow.com/questions/57385766/disable-tensorflow-log-information

ENV KMP_WARNINGS=0

# Also make the KMP logs noverbose.

# https://stackoverflow.com/questions/70250304/stop-tensorflow-from-printing-warning-message

ENV KMP_SETTINGS=false

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

ADD patches/nbconvert-extensions.tpl /opt/kaggle/nbconvert-extensions.tpl

ADD patches/template_conf.json /opt/kaggle/conf.json

{{ if eq .Accelerator "gpu" }}

# b/200968891 Keeps horovod once torch is upgraded.

RUN pip uninstall -y horovod && \

/tmp/clean-layer.sh

{{ end }}

# Update GPG key per documentation at https://cloud.google.com/compute/docs/troubleshooting/known-issues

RUN curl https://packages.cloud.google.com/apt/doc/apt-key.gpg | sudo apt-key add -

RUN curl https://packages.cloud.google.com/apt/doc/apt-key.gpg | sudo apt-key --keyring /usr/share/keyrings/cloud.google.gpg add -

# Use a fixed apt-get repo to stop intermittent failures due to flaky httpredir connections,

# as described by Lionel Chan at http://stackoverflow.com/a/37426929/5881346

RUN sed -i "s/httpredir.debian.org/debian.uchicago.edu/" /etc/apt/sources.list && \

apt-get update && \

# Needed by lightGBM (GPU build)

# https://lightgbm.readthedocs.io/en/latest/GPU-Tutorial.html#build-lightgbm

apt-get install -y build-essential unzip cmake libboost-dev libboost-system-dev libboost-filesystem-dev p7zip-full && \

# b/182601974: ssh client was removed from the base image but is required for packages such as stable-baselines.

apt-get install -y openssh-client && \

/tmp/clean-layer.sh

# b/128333086: Set PROJ_LIB to points to the proj4 cartographic library.

ENV PROJ_LIB=/opt/conda/share/proj

# Install conda packages not available on pip.

# 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/

RUN conda config --add channels nvidia && \

conda config --add channels rapidsai && \

# Base image channel order: conda-forge (highest priority), defaults.

# End state: rapidsai (highest priority), nvidia, conda-forge, defaults.

conda install mkl cartopy=0.19 imagemagick=7.1 pyproj==3.1.0 && \

/tmp/clean-layer.sh

{{ if eq .Accelerator "gpu" }}

# b/232247930: uninstall pyarrow to avoid double installation with the GPU specific version.

RUN pip uninstall -y pyarrow && \

conda install cudf=21.10 cuml=21.10 cudatoolkit=$CUDA_MAJOR_VERSION.$CUDA_MINOR_VERSION && \

/tmp/clean-layer.sh

{{ end }}

# Install implicit

{{ if eq .Accelerator "gpu" }}

RUN conda install implicit implicit-proc=*=gpu && \

/tmp/clean-layer.sh

{{ else }}

RUN conda install implicit && \

/tmp/clean-layer.sh

{{ end}}

# Install PyTorch

{{ if eq .Accelerator "gpu" }}

COPY --from=torch_whl /tmp/whl/*.whl /tmp/torch/