revamp the workflow in preparation for nightly runs

Author: jbdalido

.github/workflows/build.yaml

@@ -4,6 +4,12 @@ on:
     tags:
       - "*"
   pull_request:
+  workflow_dispatch:
+    inputs:
+      xla_commit:
+        required: true
+        type: string
+        default: "main"
 
 concurrency:
   group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.ref }}
@@ -12,98 +18,49 @@ concurrency:
 permissions:
   contents: write
 
-jobs:
-  pjrt-osx-artifacts:
-    strategy:
-      matrix:
-        pjrt:
-          - target: cpu
-            config: "--config=release_macos_base"
-            artifact: bazel-bin/xla/pjrt/c/libpjrt_c_api_cpu_plugin.dylib
-            runs_on: ["macos-14"]
-            platform: darwin-amd64
-            bazel_target: //xla/pjrt/c:pjrt_c_api_cpu_plugin
-          - target: cpu
-            config: "--config=release_macos_arm64"
-            artifact: bazel-bin/xla/pjrt/c/libpjrt_c_api_cpu_plugin.dylib
-            runs_on: ["self-hosted", "macOS"]
-            platform: darwin-arm64
-            bazel_target: //xla/pjrt/c:pjrt_c_api_cpu_plugin
+env:
+  XLA_COMMIT: ${{ inputs.xla_commit || 'cc075beb6148c2777da2b6749c63830856ee6c2a' }}
 
-    runs-on: ${{ matrix.pjrt.runs_on }}
+jobs:
+  setup_openxla:
+    runs-on: ubuntu-latest
     steps:
       - name: "Checking out repository"
         uses: actions/checkout@v4
         with:
           path: "pjrt-artifacts"
-
-      - name: Extract OpenXLA commit id
-        id: extract
-        run: |
-          commitid=`grep "OPENXLA_COMMIT = " ./pjrt-artifacts/third_party/openxla/xla.bzl | awk -F ' = ' '{print $2}' | sed 's/"//g'`
-          echo "commitid=$commitid" >> $GITHUB_OUTPUT
-
       - name: "Checking out openxla repository"
         uses: actions/checkout@v4
         with:
+          ref: ${{ env.XLA_COMMIT }}
           repository: openxla/xla
           path: "xla"
-          ref: ${{ steps.extract.outputs.commitid }}
-
-      - name: Apply patches
+      - name: Apply patches to openxla
         working-directory: ./xla
         run: |
-          for patch in $(ls ../pjrt-artifacts/third_party/openxla/patches/*.patch | sort); do
+          echo "::notice Applying patches to openxla ${{ github.sha }}"
+
+          for patch in $(ls ../pjrt-artifacts/openxla/patches/*.patch | sort); do
               echo "Applying patch $patch"
               git apply "$patch"
           done
-
-      - name: Setup Bazel
-        if: ${{ contains(matrix.pjrt.runs_on, 'self-hosted') }} # Run only for self-hosted macOS
-        run: |
-          wget https://github.com/bazelbuild/bazelisk/releases/download/v1.25.0/bazelisk-linux-amd64 -O /usr/local/bin/bazel
-          chmod +x /usr/local/bin/bazel
-
-      - name: "Build ${{ matrix.pjrt.target }} ${{ matrix.pjrt.platform }}"
-        working-directory: ./xla
-        run: |
-          bazelisk \
-              build \
-              --disk_cache=/tmp/cache \
-              ${{ matrix.pjrt.config }} \
-              ${{ matrix.pjrt.bazel_target }}
-
-      - name: Strip binary and rpath
-        working-directory: ./xla
-        run: |
-          strip -x -o libpjrt_cpu.dylib ${{ matrix.pjrt.artifact }}
-          install_name_tool -id '@rpath/libpjrt_cpu.dylib' libpjrt_cpu.dylib
-
-      - name: Create compressed ${{ matrix.pjrt.target }} ${{ matrix.pjrt.platform }} release file
-        uses: a7ul/tar-action@v1.2.0
-        with:
-          command: c
-          cwd: ./xla
-          files: |
-            ./libpjrt_cpu.dylib
-          outPath: pjrt-${{ matrix.pjrt.target }}_${{ matrix.pjrt.platform }}.tar.gz
-
-      - name: Upload ${{ matrix.pjrt.target }} ${{ matrix.pjrt.platform }} artifacts
+      - name: Upload openxla repository artifact
         uses: actions/upload-artifact@v4
         with:
-          name: pjrt-${{ matrix.pjrt.target }}_${{ matrix.pjrt.platform }}.tar.gz
-          path: pjrt-${{ matrix.pjrt.target }}_${{ matrix.pjrt.platform }}.tar.gz
+          include-hidden-files: true
+          name: xla-${{ env.XLA_COMMIT }}
+          path: ./xla
 
   pjrt-artifacts:
+    runs-on: ${{ matrix.pjrt.runs_on }}
     strategy:
       matrix:
         pjrt:
           - target: cuda
             config: "--config=cuda"
             artifact: libpjrt_c_api_gpu_plugin.so
             renamed_artifact: libpjrt_cuda.so
-            runs_on:
-              ["runs-on", "runner=32cpu-linux-x64", "image=ubuntu24-amd64"]
+            runs_on: ["self-hosted", "Linux", "X64"]
             platform: linux-amd64
             bazel_target: //xla/pjrt/c:pjrt_c_api_gpu_plugin
           - target: rocm
@@ -114,6 +71,13 @@ jobs:
               ["runs-on", "runner=32cpu-linux-x64", "image=ubuntu24-amd64"]
             platform: linux-amd64
             bazel_target: //xla/pjrt/c:pjrt_c_api_gpu_plugin
+          - target: cpu
+            config: "--config=release_macos_arm64"
+            artifact: bazel-bin/xla/pjrt/c/libpjrt_c_api_cpu_plugin.dylib
+            runs_on: ["self-hosted", "macOS"]
+            renamed_artifact: libpjrt_cpu.dylib
+            platform: darwin-arm64
+            bazel_target: //xla/pjrt/c:pjrt_c_api_cpu_plugin
           - target: cpu
             config: ""
             artifact: libpjrt_c_api_cpu_plugin.so
@@ -122,66 +86,105 @@ jobs:
               ["runs-on", "runner=32cpu-linux-x64", "image=ubuntu24-amd64"]
             platform: linux-amd64
             bazel_target: //xla/pjrt/c:pjrt_c_api_cpu_plugin
-    runs-on: ${{ matrix.pjrt.runs_on }}
+    needs: ["setup_openxla"]
     steps:
       - name: "Checking out repository"
         uses: actions/checkout@v4
+        with:
+          path: "pjrt-artifacts"
+
+      - name: Download xla artifact
+        uses: actions/download-artifact@v4
+        with:
+          name: xla-${{ env.XLA_COMMIT }}
+          path: xla
+
       - name: Setup Bazel
         run: |
-          wget https://github.com/bazelbuild/bazelisk/releases/download/v1.25.0/bazelisk-linux-amd64 -O ~/bin/bazel
+          mkdir -p ~/bin/
+          wget https://github.com/bazelbuild/bazelisk/releases/download/v1.25.0/bazelisk-${{ matrix.pjrt.platform }} -O ~/bin/bazel
           chmod +x ~/bin/bazel
-      - name: Force remove container
+
+      - name: "Setup bazelrc for openxla"
         run: |
-          docker rm -f pjrt-${{ matrix.pjrt.target }}
+          cp pjrt-artifacts/openxla/bazelrc/common.bazelrc xla/
+          cp pjrt-artifacts/openxla/bazelrc/${{ matrix.pjrt.target }}.bazelrc xla/xla_configure.bazelrc
+
       - name: "Build ${{ matrix.pjrt.target }} ${{ matrix.pjrt.platform }} image"
+        if: matrix.pjrt.target != 'cpu'
+        working-directory: ./pjrt-artifacts
         run: |
+          docker rm -f pjrt-${{ matrix.pjrt.target }}
           ~/bin/bazel run //${{ matrix.pjrt.target }}:${{ matrix.pjrt.target }}_stripped_tarball
+
       - name: "Run ${{ matrix.pjrt.target }} ${{ matrix.pjrt.platform }} image"
+        if: matrix.pjrt.target != 'cpu'
         run: |
           docker run \
             --name=pjrt-${{ matrix.pjrt.target }} \
             -w=/xla \
             --net=host \
+            -v=$(pwd)/xla:/xla \
             -v=/tmp/cache:/tmp/cache \
             distroless/${{ matrix.pjrt.target }}_builder:latest \
               bazelisk \
               build \
-              --remote_cache=grpc://127.0.0.1:15501 \
               --disk_cache=/tmp/cache \
               ${{ matrix.pjrt.config }} \
               ${{ matrix.pjrt.bazel_target }}
+
       - name: "Retrieve ${{ matrix.pjrt.target }} ${{ matrix.pjrt.platform }} artifacts"
+        if: matrix.pjrt.target != 'cpu'
         run: |
           docker cp pjrt-${{ matrix.pjrt.target }}:/xla/bazel-bin/xla/pjrt/c/${{ matrix.pjrt.artifact }} ${{ matrix.pjrt.renamed_artifact }}
           docker rm -f pjrt-${{ matrix.pjrt.target }}
+
+      - name: "Build ${{ matrix.pjrt.target }} ${{ matrix.pjrt.platform }}"
+        if: matrix.pjrt.target == 'cpu'
+        working-directory: ./xla
+        run: |
+          ~/bin/bazel \
+              build \
+              --disk_cache=/tmp/cache \
+              ${{ matrix.pjrt.config }} \
+              ${{ matrix.pjrt.bazel_target }}
+
+      - name: Strip binary and rpath
+        if: matrix.pjrt.target == 'cpu'
+        working-directory: ./xla
+        run: |
+          strip -x -o libpjrt_cpu.dylib ${{ matrix.pjrt.artifact }}
+          install_name_tool -id '@rpath/libpjrt_cpu.dylib' libpjrt_cpu.dylib
+
       - name: Create compressed ${{ matrix.pjrt.target }} ${{ matrix.pjrt.platform }} release file
         uses: a7ul/tar-action@v1.2.0
         with:
           command: c
-          cwd: .
+          cwd: ./
           files: |
             ./${{ matrix.pjrt.renamed_artifact }}
           outPath: pjrt-${{ matrix.pjrt.target }}_${{ matrix.pjrt.platform }}.tar.gz
+
       - name: Upload ${{ matrix.pjrt.target }} ${{ matrix.pjrt.platform }} artifacts
         uses: actions/upload-artifact@v4
         with:
           name: pjrt-${{ matrix.pjrt.target }}_${{ matrix.pjrt.platform }}.tar.gz
           path: pjrt-${{ matrix.pjrt.target }}_${{ matrix.pjrt.platform }}.tar.gz
 
   release:
-    needs: ["pjrt-artifacts", "pjrt-osx-artifacts"]
-    runs-on: ["runs-on", "runner=2cpu-linux-x64"]
-    if: startsWith(github.ref, 'refs/tags/')
+    needs: ["pjrt-artifacts"]
+    runs-on: ubuntu-latest
     steps:
       - run: rm -rf pjrt*.tar.gz
       - name: Download all artifacts
         uses: actions/download-artifact@v4
+        if: startsWith(github.ref, 'refs/tags/')
       - name: Release
         uses: softprops/action-gh-release@v2
+        if: startsWith(github.ref, 'refs/tags/')
         with:
           files: |
-            pjrt-cpu_darwin-arm64.tar.gz/pjrt-cpu_darwin-arm64.tar.gz
-            pjrt-cpu_darwin-amd64.tar.gz/pjrt-cpu_darwin-amd64.tar.gz
             pjrt-rocm_linux-amd64.tar.gz/pjrt-rocm_linux-amd64.tar.gz
             pjrt-cuda_linux-amd64.tar.gz/pjrt-cuda_linux-amd64.tar.gz
             pjrt-cpu_linux-amd64.tar.gz/pjrt-cpu_linux-amd64.tar.gz
+            pjrt-cpu_darwin-arm64.tar.gz/pjrt-cpu_darwin-arm64.tar.gz

cpu/BUILD.bazel

@@ -65,8 +65,6 @@ oci_image(
         "//base/jammy:packages",
         ":strip_packages",
         ":symlinks",
-        "//third_party/openxla:openxla_tar",
-        "//third_party/openxla:cpu_configuration_tar",
     ],
     env = {
         "MAX_NUM_WORKERS": -1,

cuda/BUILD.bazel

@@ -67,8 +67,6 @@ oci_image(
         "//base/noble:bazelisk",
         ":strip_packages",
         ":symlinks",
-        "//third_party/openxla:openxla_tar",
-        "//third_party/openxla:cuda_configuration_tar",
     ],
     env = {
         "MAX_NUM_WORKERS": -1,