ARG BASE_IMAGE \

BASE_IMAGE_TAG \

LIGHTGBM_VERSION

{{ if eq .Accelerator "gpu" }}

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

{{ end }}

FROM ${BASE_IMAGE}:${BASE_IMAGE_TAG}

ADD kaggle_requirements.txt /kaggle_requirements.txt

# Freeze existing requirements from base image for critical packages:

RUN pip freeze | grep -E 'tensorflow|keras|torch|jax|lightgbm' > /colab_requirements.txt

# Merge requirements files:

RUN cat /colab_requirements.txt >> /requirements.txt

RUN cat /kaggle_requirements.txt >> /requirements.txt

# TODO: GPU requirements.txt

# TODO: merge them better (override matching ones).

# Install uv & Kaggle packages

RUN curl -LsSf https://astral.sh/uv/install.sh | sh

RUN export PATH="${HOME}/.local/bin:${PATH}" && uv pip install --system -r /requirements.txt

ENV PATH="~/.local/bin:${PATH}"

# Install manual packages:

# b/183041606#comment5: the Kaggle data proxy doesn't support these APIs. If the library is missing, it falls back to using a regular BigQuery query to fetch data.

RUN uv pip uninstall --system google-cloud-bigquery-storage

# b/394382016: sigstore (dependency of kagglehub) requires a prerelease packages, installing separate.

RUN uv pip install --system --force-reinstall --prerelease=allow kagglehub[pandas-datasets,hf-datasets,signing]>=0.3.9

# b/408284143: google-cloud-automl 2.0.0 introduced incompatible API changes, need to pin to 1.0.1

# b/408284435: Keras 3.6 broke test_keras.py > test_train > keras.datasets.mnist.load_data()

# See https://github.com/keras-team/keras/commit/dcefb139863505d166dd1325066f329b3033d45a

# Colab base is on Keras 3.8, we have to install the package separately

RUN uv pip install --system google-cloud-automl==1.0.1 google-cloud-aiplatform google-cloud-translate==3.12.1 \

google-cloud-videointelligence google-cloud-vision google-genai "keras<3.6"

# uv cannot install this in requirements.txt without --no-build-isolation

# to avoid affecting the larger build, we'll post-install it.

RUN uv pip install --no-build-isolation --system "git+https://github.com/Kaggle/learntools"

# b/408281617: Torch is adamant that it can not install cudnn 9.3.x, only 9.1.x, but Tensorflow can only support 9.3.x.

# This conflict causes a number of package downgrades, which are handled in this command

# b/302136621: Fix eli5 import for learntools

RUN uv pip install --system --force-reinstall --extra-index-url https://pypi.nvidia.com "cuml-cu12==25.2.1" \

"nvidia-cudnn-cu12==9.3.0.75" scipy tsfresh scikit-learn==1.2.2 category-encoders eli5

RUN uv pip install --system --force-reinstall "pynvjitlink-cu12==0.5.2"

# b/385145217 Latest Colab lacks mkl numpy, install it.

RUN uv pip install --system --force-reinstall -i https://pypi.anaconda.org/intel/simple numpy

# newer daal4py requires tbb>=2022, but libpysal is downgrading it for some reason

RUN uv pip install --system "tbb>=2022" "libpysal==4.9.2"

# b/404590350: Ray and torchtune have conflicting tune cli, we will prioritize torchtune.

RUN uv pip install --system --force-reinstall --no-deps torchtune

# Adding non-package dependencies:

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

# /opt/conda/lib/python3.11/site-packages

ARG PACKAGE_PATH=/usr/local/lib/python3.11/dist-packages

# Install GPU-specific non-pip packages.

{{ if eq .Accelerator "gpu" }}

ARG CUDA_MAJOR_VERSION \

CUDA_MINOR_VERSION

ENV CUDA_MAJOR_VERSION=${CUDA_MAJOR_VERSION} \

CUDA_MINOR_VERSION=${CUDA_MINOR_VERSION}

# Make sure we are on the right version of CUDA

RUN update-alternatives --set cuda /usr/local/cuda-$CUDA_MAJOR_VERSION.$CUDA_MINOR_VERSION

RUN uv pip install --system "pycuda"

# b/381256047 Remove once installed in Colabs base image.

# Install LightGBM

COPY --from=lightgbm_whl /tmp/whl/*.whl /tmp/lightgbm/