scikit-learn · thomasjpfan · Jun 4, 2024 · May 29, 2024 · May 29, 2024 · May 29, 2024
diff --git a/.github/workflows/cuda-gpu-ci.yml b/.github/workflows/cuda-gpu-ci.yml
@@ -0,0 +1,47 @@
+name: CUDA GPU
+on:
+  workflow_dispatch:
+    inputs:
+      pr_id:
+        description: Test the contents of this Pull Request
+        required: true
+
+permissions: read-all
+
+jobs:
+  tests:
+    runs-on:
+      group: cuda-gpu-runner-group
+    name: Run Array API unit tests
+    steps:
+      - uses: actions/setup-python@v4
+        with:
+          python-version: '3.12'
+      - name: Checkout main repository
+        uses: actions/checkout@v2
+      - name: Checkout a particular Pull Request
+        if: inputs.pr_id
+        env:
+          PR_ID: ${{ inputs.pr_id }}
+          GH_TOKEN: ${{ github.token }}
+        run: |
+          gh pr checkout ${{ env.PR_ID }}
+      - name: Cache conda environment
+        id: cache-conda
+        uses: actions/cache@v3
+        with:
+          path: ~/conda
+          key: ${{ runner.os }}-build-${{ hashFiles('build_tools/github/create_gpu_environment.sh') }}-${{ hashFiles('build_tools/github/pylatest_conda_forge_cuda_array-api_linux-64_conda.lock') }}
+      - name: Install miniforge
+        if: ${{ steps.cache-conda.outputs.cache-hit != 'true' }}
+        run: bash build_tools/github/create_gpu_environment.sh
+      - name: Install scikit-learn
+        run: |
+          source "${HOME}/conda/etc/profile.d/conda.sh"
+          conda activate sklearn
+          pip install --verbose --no-build-isolation --config-settings editable-verbose=true --editable .
+      - name: Run array API tests
+        run: |
+          source "${HOME}/conda/etc/profile.d/conda.sh"
+          conda activate sklearn
+          pytest -k 'array_api'
diff --git a/.github/workflows/update-lock-files.yml b/.github/workflows/update-lock-files.yml
@@ -6,6 +6,10 @@ on:
   schedule:
     - cron: '0 5 * * 1'
 
+# XXX Set the right permissions, per step??
+# Can we set read only at the global level here and then elevate to write for some steps?
+#permissions: read-all
+
 jobs:
   update_lock_files:
     if: github.repository == 'scikit-learn/scikit-learn'
@@ -25,6 +29,8 @@ jobs:
           - name: cirrus-arm
             update_script_args: "--select-tag arm"
             additional_commit_message: "[cirrus arm]"
+          - name: array API
+            update_script_args: "--select-tag cuda"
 
     steps:
       - uses: actions/checkout@v4
@@ -56,6 +62,14 @@ jobs:
             ### Note
             If the CI tasks fail, create a new branch based on this PR and add the required fixes to that branch.
 
+      # The CUDA workflow needs to be triggered explicitly as it uses an expensive runner
+      - name: Trigger additional tests
+        if: steps.cpr.outputs.pull-request-number != '' && matrix.name == 'array API'
+        env:
+          GH_TOKEN: ${{ github.token }}
+        run: |
+          gh workflow run .github/workflows/cuda-gpu-ci.yml -f pr_id=${{steps.cpr.outputs.pull-request-number}}
+
       - name: Check Pull Request
         if: steps.cpr.outputs.pull-request-number != ''
         run: |

diff --git a/build_tools/github/create_gpu_environment.sh b/build_tools/github/create_gpu_environment.sh
@@ -0,0 +1,17 @@
+#!/bin/bash
+
+set -e
+set -x
+
+curl -L -O "https://github.com/conda-forge/miniforge/releases/latest/download/Miniforge3-$(uname)-$(uname -m).sh"
+bash Miniforge3-$(uname)-$(uname -m).sh -b -p "${HOME}/conda"
+source "${HOME}/conda/etc/profile.d/conda.sh"
+
+
+# defines the get_dep and show_installed_libraries functions
+source build_tools/shared.sh
+conda activate base
+
+# XXX switch once https://github.com/scikit-learn/scikit-learn/pull/29176 is merged
+conda install -c conda-forge "$(get_dep conda-lock min)" -y
+conda-lock install --name sklearn build_tools/github/pylatest_conda_forge_cuda_array-api_linux-64_conda.lock