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]>=0.3.7

# 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/385161357 Latest Colab uses tf 2.17.1, but tf decision forests only has a version for 2.17.0.

# Instead, we'll install tfdf with its deps and hope that 2.17.0 compat tfdf works with tf 2.17.1.

RUN uv pip install --system --no-deps tensorflow-decision-forests==1.10.0 wurlitzer==3.1.1 ydf==0.9.0

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

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

# b/328788268 We install an incompatible pair of libs (shapely<2, libpysal==4.9.2) so we can't put this one in the requirements.txt

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

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

# 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.10/site-packages

ARG PACKAGE_PATH=/usr/local/lib/python3.10/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/