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}

#b/415358342: UV reports missing requirements files https://github.com/googlecolab/colabtools/issues/5237

ENV UV_CONSTRAINT= \

UV_BUILD_CONSTRAINT=

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 Kaggle packages

RUN uv pip install --system -r /requirements.txt

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

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

# installed outside of kaggle_requirements.txt due to requiring an incompatibile version of protobuf.

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

google-cloud-automl==1.0.1

# 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

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"

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.

# b/415358158: Gensim removed from Colab image to upgrade scipy

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

# 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" }}

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/