b/381256047 lightgbm gpu install

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
 
@@ -49,6 +56,17 @@ 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.
+# 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 }}
 ENV CUDA_VERSION=""

Jenkinsfile

@@ -21,6 +21,29 @@ pipeline {
   }
 
   stages {
+    stage('Pre-build Packages from Source') {
+      stages {
+        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

@@ -34,9 +34,7 @@ def test_cpu(self):
 
         self.assertEqual(1, gbm.best_iteration)
 
-    # TODO(b/381256047): Colab needs to install GPU-enabled lightgbm.
     @gpu_test
-    @unittest.skip("Skipping this test until b/381256047 is resolved.")
     def test_gpu(self):
         lgb_train, lgb_eval = self.load_datasets()