From 5f2978a55343e6653db946ca7d4f9f04a2c46068 Mon Sep 17 00:00:00 2001 From: Corentin Kerisit Date: Mon, 22 Dec 2025 17:04:02 +0100 Subject: [PATCH 1/6] Proper split between nightly and pull request Fixes deadlock because of concurrency key duplicated. --- .github/workflows/build.yaml | 22 ---------------------- .github/workflows/nightly.yaml | 6 ++---- .github/workflows/pull_request.yaml | 20 ++++++++++++++++++++ 3 files changed, 22 insertions(+), 26 deletions(-) create mode 100644 .github/workflows/pull_request.yaml diff --git a/.github/workflows/build.yaml b/.github/workflows/build.yaml index 604ebd7a..eb8682a6 100644 --- a/.github/workflows/build.yaml +++ b/.github/workflows/build.yaml @@ -1,19 +1,4 @@ -name: PJRT GPU library on: - push: - tags: - - "*" - pull_request: - workflow_dispatch: - inputs: - xla_commit: - required: true - type: string - default: "main" - rocm_xla_commit: - required: true - type: string - default: "main" workflow_call: inputs: xla_commit: @@ -25,13 +10,6 @@ on: type: string default: "main" -concurrency: - group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.ref }} - cancel-in-progress: true - -permissions: - contents: write - env: # From master 2025-12-12 XLA_COMMIT: ${{ inputs.xla_commit || 'f238b48769d2ab8d62eeb09b5d31a972dfa4841a' }} diff --git a/.github/workflows/nightly.yaml b/.github/workflows/nightly.yaml index 7b89fcd2..68ad4a09 100644 --- a/.github/workflows/nightly.yaml +++ b/.github/workflows/nightly.yaml @@ -1,12 +1,11 @@ name: nightly - PJRT CUDA library on: workflow_dispatch: - pull_request: schedule: - - cron: '0 */4 * * *' + - cron: '0 0 * * *' concurrency: - group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.ref }} + group: ${{ github.workflow_sha }}-${{ github.ref }} cancel-in-progress: true permissions: @@ -14,7 +13,6 @@ permissions: jobs: nightly-pjrt-artifacts: - if: github.event_name != 'pull_request' || contains(join(github.event.pull_request.labels.*.name, ','), 'nightly') uses: ./.github/workflows/build.yaml with: xla_commit: main diff --git a/.github/workflows/pull_request.yaml b/.github/workflows/pull_request.yaml new file mode 100644 index 00000000..50271298 --- /dev/null +++ b/.github/workflows/pull_request.yaml @@ -0,0 +1,20 @@ +name: PJRT GPU library +on: + push: + tags: + - "*" + pull_request: + +concurrency: + group: ${{ github.workflow_sha }}-${{ github.event.pull_request.number || github.ref }} + cancel-in-progress: true + +permissions: + contents: write + +jobs: + pjrt-artifacts: + uses: ./.github/workflows/build.yaml + with: + xla_commit: main + rocm_xla_commit: rocm-jaxlib-v0.8.0 From 7f24f0cc2b5445bf160887a023641152ed809a1c Mon Sep 17 00:00:00 2001 From: Corentin Kerisit Date: Mon, 22 Dec 2025 17:07:04 +0100 Subject: [PATCH 2/6] fix --- .github/workflows/build.yaml | 6 ++---- .github/workflows/pull_request.yaml | 4 ++-- 2 files changed, 4 insertions(+), 6 deletions(-) diff --git a/.github/workflows/build.yaml b/.github/workflows/build.yaml index eb8682a6..df3e971c 100644 --- a/.github/workflows/build.yaml +++ b/.github/workflows/build.yaml @@ -11,10 +11,8 @@ on: default: "main" env: - # From master 2025-12-12 - XLA_COMMIT: ${{ inputs.xla_commit || 'f238b48769d2ab8d62eeb09b5d31a972dfa4841a' }} - # rocm-jaxlib-v0.8.0 - ROCM_XLA_COMMIT: ${{ inputs.rocm_xla_commit || '06402b44669c52956732678772104dcb85c53806' }} + XLA_COMMIT: ${{ inputs.xla_commit }} + ROCM_XLA_COMMIT: ${{ inputs.rocm_xla_commit }} TF_ROCM_AMDGPU_TARGETS: "gfx900,gfx906,gfx908,gfx90a,gfx942,gfx1030,gfx1100" jobs: diff --git a/.github/workflows/pull_request.yaml b/.github/workflows/pull_request.yaml index 50271298..3a222e32 100644 --- a/.github/workflows/pull_request.yaml +++ b/.github/workflows/pull_request.yaml @@ -16,5 +16,5 @@ jobs: pjrt-artifacts: uses: ./.github/workflows/build.yaml with: - xla_commit: main - rocm_xla_commit: rocm-jaxlib-v0.8.0 + xla_commit: f238b48769d2ab8d62eeb09b5d31a972dfa4841a # main + rocm_xla_commit: 06402b44669c52956732678772104dcb85c53806 # rocm-jaxlib-v0.8.0 From efe1dbbcaefde8aa73e3d5fb5ec87603f2f475bc Mon Sep 17 00:00:00 2001 From: Corentin Kerisit Date: Tue, 23 Dec 2025 14:05:46 +0100 Subject: [PATCH 3/6] revamp --- .github/workflows/{build.yaml => _build.yaml} | 117 ++++-------------- .github/workflows/ci.yaml | 11 ++ .github/workflows/nightly.yaml | 8 +- .github/workflows/pull_request.yaml | 20 --- .github/workflows/release.yaml | 31 +++++ 5 files changed, 68 insertions(+), 119 deletions(-) rename .github/workflows/{build.yaml => _build.yaml} (69%) create mode 100644 .github/workflows/ci.yaml delete mode 100644 .github/workflows/pull_request.yaml create mode 100644 .github/workflows/release.yaml diff --git a/.github/workflows/build.yaml b/.github/workflows/_build.yaml similarity index 69% rename from .github/workflows/build.yaml rename to .github/workflows/_build.yaml index df3e971c..06dfedf1 100644 --- a/.github/workflows/build.yaml +++ b/.github/workflows/_build.yaml @@ -1,78 +1,21 @@ +name: Build PJRT GPU libraries + on: workflow_call: inputs: xla_commit: required: false type: string - default: "main" rocm_xla_commit: required: false type: string - default: "main" env: - XLA_COMMIT: ${{ inputs.xla_commit }} - ROCM_XLA_COMMIT: ${{ inputs.rocm_xla_commit }} + XLA_COMMIT: ${{ inputs.xla_commit || 'f238b48769d2ab8d62eeb09b5d31a972dfa4841a' }} # main from 2025-12-12 + ROCM_XLA_COMMIT: ${{ inputs.rocm_xla_commit || '06402b44669c52956732678772104dcb85c53806' }} # rocm-jaxlib-v0.8.0 TF_ROCM_AMDGPU_TARGETS: "gfx900,gfx906,gfx908,gfx90a,gfx942,gfx1030,gfx1100" jobs: - setup_openxla: - runs-on: ubuntu-latest - outputs: - xla_commit: ${{ steps.patches.outputs.XLA_COMMIT_ID }} - rocm_xla_commit: ${{ steps.rocm_patches.outputs.XLA_COMMIT_ID }} - steps: - - name: "Checking out repository" - uses: actions/checkout@v4 - with: - path: "pjrt-artifacts" - - name: "Checking out openxla repository" - uses: actions/checkout@v4 - with: - ref: ${{ env.XLA_COMMIT }} - repository: openxla/xla - path: "xla" - - name: Apply patches to openxla - id: patches - working-directory: ./xla - run: | - xla_commit=$(git rev-parse HEAD) - echo "XLA_COMMIT_ID=$xla_commit" >> $GITHUB_OUTPUT - echo ::notice::Applying patches to openxla $xla_commit - for patch in $(ls ../pjrt-artifacts/openxla/patches/upstream/*.patch | sort); do - echo "Applying patch $patch" - git apply "$patch" - done - - name: Upload openxla repository artifact - uses: actions/upload-artifact@v4 - with: - include-hidden-files: true - name: xla-${{ steps.patches.outputs.XLA_COMMIT_ID }} - path: ./xla - - name: "Checking out ROCm xla repository" - uses: actions/checkout@v4 - with: - ref: ${{ env.ROCM_XLA_COMMIT }} - repository: ROCm/xla - path: "xla-rocm" - - name: Apply patches to ROCm openxla - id: rocm_patches - working-directory: ./xla-rocm - run: | - xla_commit=$(git rev-parse HEAD) - echo "XLA_COMMIT_ID=$xla_commit" >> $GITHUB_OUTPUT - echo ::notice::Applying patches to ROCm openxla $xla_commit - for patch in $(ls ../pjrt-artifacts/openxla/patches/rocm/*.patch | sort); do - echo "Applying patch $patch" - git apply "$patch" - done - - name: Upload ROCm openxla repository artifact - uses: actions/upload-artifact@v4 - with: - include-hidden-files: true - name: xla-rocm-${{ steps.rocm_patches.outputs.XLA_COMMIT_ID }} - path: ./xla-rocm - pjrt-artifacts: runs-on: ${{ matrix.pjrt.runs_on }} strategy: @@ -119,8 +62,6 @@ jobs: runs_on: ["runs-on", "runner=32cpu-linux-x64", "image=ubuntu22-amd64"] platform: linux-amd64 bazel_target: //xla/pjrt/c:pjrt_c_api_cpu_plugin - - needs: ["setup_openxla"] steps: - uses: runs-on/action@v1 if: matrix.pjrt.platform != 'darwin-arm64' && matrix.pjrt.platform != 'darwin-amd64' @@ -151,19 +92,25 @@ jobs: sudo apt install ./amdgpu-install_7.1.1.70101-1_all.deb -y sudo amdgpu-install --usecase=rocm,rocmdev,hiplibsdk -y --no-dkms - - name: Download xla artifact - if: matrix.pjrt.target != 'rocm' - uses: actions/download-artifact@v4 - with: - name: xla-${{ needs.setup_openxla.outputs.xla_commit }} - path: xla - - - name: Download ROCm xla artifact - if: matrix.pjrt.target == 'rocm' - uses: actions/download-artifact@v4 + - name: "Checking out openxla repository" + uses: actions/checkout@v4 with: - name: xla-rocm-${{ needs.setup_openxla.outputs.rocm_xla_commit }} - path: xla + ref: ${{ matrix.pjrt.target == 'rocm' && env.ROCM_XLA_COMMIT || env.XLA_COMMIT }} + repository: ${{ matrix.pjrt.repository }} + path: "xla" + - name: Apply patches to openxla + id: patches + working-directory: ./xla + env: + PATCH_DIRECTORY: ${{ matrix.pjrt.target == 'rocm' && 'rocm' || 'upstream' }} + run: | + xla_commit=$(git rev-parse HEAD) + echo "XLA_COMMIT_ID=$xla_commit" >> $GITHUB_OUTPUT + echo ::notice::Applying patches to openxla $xla_commit + for patch in $(ls ../pjrt-artifacts/openxla/patches/$PATCH_DIRECTORY/*.patch | sort); do + echo "Applying patch $patch" + git apply "$patch" + done - uses: bazel-contrib/setup-bazel@0.15.0 with: @@ -174,13 +121,10 @@ jobs: common --color=yes common --show_timestamps - - name: "Setup bazelrc for openxla" - run: | - cp pjrt-artifacts/openxla/bazelrc/${{ matrix.pjrt.target }}.bazelrc xla/xla_configure.bazelrc - - name: "Build ${{ matrix.pjrt.target }} ${{ matrix.pjrt.platform }}" working-directory: ./xla run: | + cp pjrt-artifacts/openxla/bazelrc/${{ matrix.pjrt.target }}.bazelrc xla/xla_configure.bazelrc bazel ${{ matrix.pjrt.bazel_opts }} build ${{ matrix.pjrt.config }} ${{ matrix.pjrt.bazel_target }} - name: Set file at the right path @@ -210,18 +154,3 @@ jobs: 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"] - 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: | - *.tar.gz/*.tar.gz diff --git a/.github/workflows/ci.yaml b/.github/workflows/ci.yaml new file mode 100644 index 00000000..ae0689b6 --- /dev/null +++ b/.github/workflows/ci.yaml @@ -0,0 +1,11 @@ +name: CI +on: + pull_request: + +concurrency: + group: ${{ github.workflow_sha }}-${{ github.event.pull_request.number }} + cancel-in-progress: true + +jobs: + pjrt-artifacts: + uses: ./.github/workflows/_build.yaml diff --git a/.github/workflows/nightly.yaml b/.github/workflows/nightly.yaml index 68ad4a09..b2f77bf3 100644 --- a/.github/workflows/nightly.yaml +++ b/.github/workflows/nightly.yaml @@ -1,4 +1,5 @@ -name: nightly - PJRT CUDA library +name: Nightly + on: workflow_dispatch: schedule: @@ -8,12 +9,9 @@ concurrency: group: ${{ github.workflow_sha }}-${{ github.ref }} cancel-in-progress: true -permissions: - contents: write - jobs: nightly-pjrt-artifacts: - uses: ./.github/workflows/build.yaml + uses: ./.github/workflows/_build.yaml with: xla_commit: main rocm_xla_commit: rocm-jaxlib-v0.8.0 diff --git a/.github/workflows/pull_request.yaml b/.github/workflows/pull_request.yaml deleted file mode 100644 index 3a222e32..00000000 --- a/.github/workflows/pull_request.yaml +++ /dev/null @@ -1,20 +0,0 @@ -name: PJRT GPU library -on: - push: - tags: - - "*" - pull_request: - -concurrency: - group: ${{ github.workflow_sha }}-${{ github.event.pull_request.number || github.ref }} - cancel-in-progress: true - -permissions: - contents: write - -jobs: - pjrt-artifacts: - uses: ./.github/workflows/build.yaml - with: - xla_commit: f238b48769d2ab8d62eeb09b5d31a972dfa4841a # main - rocm_xla_commit: 06402b44669c52956732678772104dcb85c53806 # rocm-jaxlib-v0.8.0 diff --git a/.github/workflows/release.yaml b/.github/workflows/release.yaml new file mode 100644 index 00000000..760bf0d6 --- /dev/null +++ b/.github/workflows/release.yaml @@ -0,0 +1,31 @@ +name: Release + +on: + push: + tags: + - "*" + +concurrency: + group: ${{ github.workflow_sha }}-${{ github.ref_name }} + cancel-in-progress: true + +permissions: + contents: write + +jobs: + pjrt-artifacts: + uses: ./.github/workflows/_build.yaml + + release: + name: "Release PJRT Artifacts" + needs: ["pjrt-artifacts"] + runs-on: ubuntu-latest + steps: + - run: rm -rf pjrt*.tar.gz + - name: Download all artifacts + uses: actions/download-artifact@v4 + - name: Release + uses: softprops/action-gh-release@v2 + with: + files: | + *.tar.gz/*.tar.gz From 024e3d7fecb3dba64ec8321ac0476584c6c3dd38 Mon Sep 17 00:00:00 2001 From: Corentin Kerisit Date: Tue, 23 Dec 2025 14:08:08 +0100 Subject: [PATCH 4/6] fix --- .github/workflows/_build.yaml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/.github/workflows/_build.yaml b/.github/workflows/_build.yaml index 06dfedf1..ab13dc87 100644 --- a/.github/workflows/_build.yaml +++ b/.github/workflows/_build.yaml @@ -96,7 +96,7 @@ jobs: uses: actions/checkout@v4 with: ref: ${{ matrix.pjrt.target == 'rocm' && env.ROCM_XLA_COMMIT || env.XLA_COMMIT }} - repository: ${{ matrix.pjrt.repository }} + repository: ${{ matrix.pjrt.target == 'rocm' && 'ROCm/xla' || 'openxla/xla' }} path: "xla" - name: Apply patches to openxla id: patches From b11d45c951680296bdbdbcae92d96cab1a757991 Mon Sep 17 00:00:00 2001 From: Corentin Kerisit Date: Tue, 23 Dec 2025 14:11:36 +0100 Subject: [PATCH 5/6] fix --- .github/workflows/_build.yaml | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/.github/workflows/_build.yaml b/.github/workflows/_build.yaml index ab13dc87..a6b43ef0 100644 --- a/.github/workflows/_build.yaml +++ b/.github/workflows/_build.yaml @@ -121,10 +121,13 @@ jobs: common --color=yes common --show_timestamps + - name: "Setup bazelrc for openxla" + run: | + cp pjrt-artifacts/openxla/bazelrc/${{ matrix.pjrt.target }}.bazelrc xla/xla_configure.bazelrc + - name: "Build ${{ matrix.pjrt.target }} ${{ matrix.pjrt.platform }}" working-directory: ./xla run: | - cp pjrt-artifacts/openxla/bazelrc/${{ matrix.pjrt.target }}.bazelrc xla/xla_configure.bazelrc bazel ${{ matrix.pjrt.bazel_opts }} build ${{ matrix.pjrt.config }} ${{ matrix.pjrt.bazel_target }} - name: Set file at the right path From a3a5201d9b70958b510c22680e19a86362a2bfd9 Mon Sep 17 00:00:00 2001 From: Corentin Kerisit Date: Tue, 23 Dec 2025 14:35:45 +0100 Subject: [PATCH 6/6] First checkout and patch --- .github/workflows/_build.yaml | 41 ++++++++++++++++++----------------- 1 file changed, 21 insertions(+), 20 deletions(-) diff --git a/.github/workflows/_build.yaml b/.github/workflows/_build.yaml index a6b43ef0..ba9ef4d8 100644 --- a/.github/workflows/_build.yaml +++ b/.github/workflows/_build.yaml @@ -71,6 +71,27 @@ jobs: with: path: "pjrt-artifacts" + - name: "Checking out openxla repository" + uses: actions/checkout@v4 + with: + ref: ${{ matrix.pjrt.target == 'rocm' && env.ROCM_XLA_COMMIT || env.XLA_COMMIT }} + repository: ${{ matrix.pjrt.target == 'rocm' && 'ROCm/xla' || 'openxla/xla' }} + path: "xla" + + - name: Apply patches to openxla + id: patches + working-directory: ./xla + env: + PATCH_DIRECTORY: ${{ matrix.pjrt.target == 'rocm' && 'rocm' || 'upstream' }} + run: | + xla_commit=$(git rev-parse HEAD) + echo "XLA_COMMIT_ID=$xla_commit" >> $GITHUB_OUTPUT + echo ::notice::Applying patches to openxla $xla_commit + for patch in $(ls ../pjrt-artifacts/openxla/patches/$PATCH_DIRECTORY/*.patch | sort); do + echo "Applying patch $patch" + git apply "$patch" + done + - uses: runs-on/snapshot@v1 if: matrix.pjrt.platform != 'darwin-arm64' && matrix.pjrt.platform != 'darwin-amd64' with: @@ -92,26 +113,6 @@ jobs: sudo apt install ./amdgpu-install_7.1.1.70101-1_all.deb -y sudo amdgpu-install --usecase=rocm,rocmdev,hiplibsdk -y --no-dkms - - name: "Checking out openxla repository" - uses: actions/checkout@v4 - with: - ref: ${{ matrix.pjrt.target == 'rocm' && env.ROCM_XLA_COMMIT || env.XLA_COMMIT }} - repository: ${{ matrix.pjrt.target == 'rocm' && 'ROCm/xla' || 'openxla/xla' }} - path: "xla" - - name: Apply patches to openxla - id: patches - working-directory: ./xla - env: - PATCH_DIRECTORY: ${{ matrix.pjrt.target == 'rocm' && 'rocm' || 'upstream' }} - run: | - xla_commit=$(git rev-parse HEAD) - echo "XLA_COMMIT_ID=$xla_commit" >> $GITHUB_OUTPUT - echo ::notice::Applying patches to openxla $xla_commit - for patch in $(ls ../pjrt-artifacts/openxla/patches/$PATCH_DIRECTORY/*.patch | sort); do - echo "Applying patch $patch" - git apply "$patch" - done - - uses: bazel-contrib/setup-bazel@0.15.0 with: bazelisk-version: 1.26.0