github-actions: revamp of the workflow #236
Workflow file for this run
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| name: PJRT GPU library | |
| on: | |
| push: | |
| 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 }} | |
| cancel-in-progress: true | |
| permissions: | |
| contents: write | |
| env: | |
| XLA_COMMIT: ${{ inputs.xla_commit || 'cc075beb6148c2777da2b6749c63830856ee6c2a' }} | |
| jobs: | |
| setup_openxla: | |
| runs-on: ubuntu-latest | |
| 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/*.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 | |
| 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"] | |
| platform: linux-amd64 | |
| bazel_target: //xla/pjrt/c:pjrt_c_api_gpu_plugin | |
| - target: rocm | |
| config: "--config=rocm" | |
| artifact: libpjrt_c_api_gpu_plugin.so | |
| renamed_artifact: libpjrt_rocm.so | |
| runs_on: ["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: "--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: "" | |
| artifact: bazel-bin/xla/pjrt/c/libpjrt_c_api_cpu_plugin.so | |
| renamed_artifact: libpjrt_cpu.so | |
| runs_on: ["runs-on", "runner=32cpu-linux-x64", "image=ubuntu24-amd64"] | |
| platform: linux-amd64 | |
| bazel_target: //xla/pjrt/c:pjrt_c_api_cpu_plugin | |
| 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: | | |
| 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: "Setup bazelrc for openxla" | |
| run: | | |
| 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 \ | |
| distroless/${{ matrix.pjrt.target }}_builder:latest \ | |
| bazelisk \ | |
| build \ | |
| --bes_results_url=https://app.buildbuddy.io/invocation/ \ | |
| --bes_backend=grpcs://remote.buildbuddy.io \ | |
| --remote_cache=grpcs://remote.buildbuddy.io \ | |
| --experimental_remote_build_event_upload=minimal \ | |
| --nolegacy_important_outputs \ | |
| --remote_download_toplevel \ | |
| --remote_timeout=3600 \ | |
| --remote_header=x-buildbuddy-api-key=${{ secrets.BUILDBUDDY_ORG_API_KEY }} \ | |
| --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 \ | |
| --bes_results_url=https://app.buildbuddy.io/invocation/ \ | |
| --bes_backend=grpcs://remote.buildbuddy.io \ | |
| --remote_cache=grpcs://remote.buildbuddy.io \ | |
| --experimental_remote_build_event_upload=minimal \ | |
| --nolegacy_important_outputs \ | |
| --remote_download_toplevel \ | |
| --remote_timeout=3600 \ | |
| --remote_header=x-buildbuddy-api-key=${{ secrets.BUILDBUDDY_ORG_API_KEY }} \ | |
| ${{ matrix.pjrt.config }} \ | |
| ${{ matrix.pjrt.bazel_target }} | |
| - name: Set file at the right path | |
| if: matrix.pjrt.target == 'cpu' && matrix.pjrt.platform == 'linux-amd64' | |
| working-directory: ./xla | |
| run: | | |
| cp ${{ matrix.pjrt.artifact }} ../${{ matrix.pjrt.renamed_artifact }} | |
| - name: Strip binary and rpath | |
| if: matrix.pjrt.platform == 'darwin-arm64' || matrix.pjrt.platform == 'darwin-amd64' | |
| working-directory: ./xla | |
| run: | | |
| strip -x -o ../${{ matrix.pjrt.renamed_artifact }} ${{ matrix.pjrt.artifact }} | |
| install_name_tool -id '@rpath/${{ matrix.pjrt.renamed_artifact }}' ../${{ matrix.pjrt.renamed_artifact }} | |
| - name: Create compressed ${{ matrix.pjrt.target }} ${{ matrix.pjrt.platform }} release file | |
| uses: a7ul/[email protected] | |
| with: | |
| command: c | |
| 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"] | |
| 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 |