diff --git a/.github/workflows/build.yaml b/.github/workflows/_build.yaml similarity index 67% rename from .github/workflows/build.yaml rename to .github/workflows/_build.yaml index 604ebd7a..ba9ef4d8 100644 --- a/.github/workflows/build.yaml +++ b/.github/workflows/_build.yaml @@ -1,102 +1,21 @@ -name: PJRT GPU library +name: Build PJRT GPU libraries + 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: required: false type: string - default: "main" rocm_xla_commit: required: false 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' }} - # rocm-jaxlib-v0.8.0 - ROCM_XLA_COMMIT: ${{ inputs.rocm_xla_commit || '06402b44669c52956732678772104dcb85c53806' }} + 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: @@ -143,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' @@ -154,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: @@ -175,20 +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: 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 - with: - name: xla-rocm-${{ needs.setup_openxla.outputs.rocm_xla_commit }} - path: xla - - uses: bazel-contrib/setup-bazel@0.15.0 with: bazelisk-version: 1.26.0 @@ -234,18 +158,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 7b89fcd2..b2f77bf3 100644 --- a/.github/workflows/nightly.yaml +++ b/.github/workflows/nightly.yaml @@ -1,21 +1,17 @@ -name: nightly - PJRT CUDA library +name: Nightly + 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: - contents: write - jobs: nightly-pjrt-artifacts: - if: github.event_name != 'pull_request' || contains(join(github.event.pull_request.labels.*.name, ','), 'nightly') - 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/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