build lightgbm gpu from source again

Author: djherbis

Dockerfile.tmpl

@@ -1,4 +1,11 @@
-FROM us-docker.pkg.dev/colab-images/public/runtime:latest
+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
 
@@ -50,11 +57,15 @@ RUN update-alternatives --set cuda /usr/local/cuda-$CUDA_MAJOR_VERSION.$CUDA_MIN
 RUN uv pip install --system "pycuda"
 
 # b/381256047 Remove once installed in Colabs base image.
-RUN pip install --force-reinstall --no-deps --no-binary lightgbm \
-  lightgbm~=4.5.0 \
-  --config-settings=cmake.define.USE_CUDA=ON \
-  --config-settings=cmake.define.OpenCL_INCLUDE_DIR="/usr/local/cuda/include/" \
-  --config-settings=cmake.define.OpenCL_LIBRARY="/usr/local/cuda/lib64/libOpenCL.so"
+# Install LightGBM
+COPY --from=lightgbm_whl /tmp/whl/*.whl /tmp/lightgbm/
+# Install OpenCL (required by LightGBM GPU version)
+RUN apt-get install -y ocl-icd-libopencl1 clinfo && \
+    mkdir -p /etc/OpenCL/vendors && \
+    echo "libnvidia-opencl.so.1" > /etc/OpenCL/vendors/nvidia.icd && \
+    uv pip install --system /tmp/lightgbm/*.whl && \
+    rm -rf /tmp/lightgbm && \
+    /tmp/clean-layer.sh
 
 # Remove CUDA_VERSION from non-GPU image.
 {{ else }}

Jenkinsfile

@@ -21,6 +21,27 @@ pipeline {
   }
 
   stages {
+    stage('Pre-build Packages from Source') {
+      parallel {
+        stage('lightgbm') {
+          options {
+            timeout(time: 10, unit: 'MINUTES')
+          }
+          steps {
+            sh '''#!/bin/bash
+              set -exo pipefail
+              source config.txt
+              cd packages/
+              ./build_package --base-image $BASE_IMAGE:$BASE_IMAGE_TAG \
+                --package lightgbm \
+                --version $LIGHTGBM_VERSION \
+                --build-arg CUDA_MAJOR_VERSION=$CUDA_MAJOR_VERSION \
+                --build-arg CUDA_MINOR_VERSION=$CUDA_MINOR_VERSION \
+                --push
+            '''
+          }
+        }
+    }
     stage('Build/Test/Diff') {
       parallel {
         stage('CPU') {

config.txt

@@ -1,2 +1,5 @@
+BASE_IMAGE=us-docker.pkg.dev/colab-images/public/runtime
+BASE_IMAGE_TAG=release-colab_20240920-060127_RC00
+LIGHTGBM_VERSION=4.5.0
 CUDA_MAJOR_VERSION=12
 CUDA_MINOR_VERSION=2

packages/build_package

@@ -115,12 +115,8 @@ if [[ -z "$DOCKERFILE" ]]; then
     exit 1
 fi
 
-# Keep only `tf2-gpu.2-6:m80` in `gcr.io/deeplearning-platform-release/tf2-gpu.2-6:m80` 
-TAG=${BASE_IMAGE/gcr.io\/deeplearning-platform-release\//}
-# Keep only `python:v108` in `gcr.io/kaggle-images/python:v108`
-TAG=${TAG/gcr.io\/kaggle-images\//}
-# Replace the `:` in `tf2-gpu.2-6:m80` by `-`
-TAG=${TAG/:/-}
+# Keep only `release-colab_20240920-060127_RC00` in `us-docker.pkg.dev/colab-images/public/runtime:release-colab_20240920-060127_RC00` 
+TAG=$(echo $BASE_IMAGE | cut -d ':' -f 2)
 # Append the package version
 TAG=$TAG-$PACKAGE_VERSION
 # Add the gcr repo.

tests/test_lightgbm.py

@@ -49,7 +49,7 @@ def test_gpu(self):
             'bagging_freq': 5,
             'force_row_wise': True,
             'verbose': 1,
-            'device': 'cuda'
+            'device': 'gpu'
         }
 
         # Run only one round for faster test