sync with xla 2024_09_19 #133
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: | |
| branches: | |
| - master | |
| tags: | |
| - '*' | |
| pull_request: | |
| concurrency: | |
| group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.ref }} | |
| cancel-in-progress: true | |
| permissions: | |
| contents: write | |
| jobs: | |
| pjrt-osx-artifacts: | |
| strategy: | |
| matrix: | |
| pjrt: | |
| - 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 | |
| runs-on: ${{ matrix.pjrt.runs_on }} | |
| steps: | |
| - name: "Checking out openxla repository" | |
| uses: actions/checkout@v4 | |
| with: | |
| repository: zml/xla | |
| token: ${{ secrets.OPENXLA_READ_ONLY }} | |
| ref: 336122e2fb0e3d7dd93b5a4a30f7551d8e1a21b6 | |
| - name: Setup Bazel | |
| run: | | |
| wget https://github.com/bazelbuild/bazelisk/releases/download/v1.20.0/bazelisk-linux-amd64 -O /usr/local/bin/bazel | |
| chmod +x /usr/local/bin/bazel | |
| - name: "Build ${{ matrix.pjrt.target }} ${{ matrix.pjrt.platform }}" | |
| run: | | |
| bazelisk \ | |
| build \ | |
| --disk_cache=/tmp/cache \ | |
| ${{ matrix.pjrt.config }} --repo_env=HERMETIC_PYTHON_VERSION=3.11 \ | |
| ${{ matrix.pjrt.bazel_target }} | |
| - name: Strip binary and rpath | |
| 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/[email protected] | |
| with: | |
| command: c | |
| cwd: . | |
| files: | | |
| ./libpjrt_cpu.dylib | |
| 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 | |
| pjrt-artifacts: | |
| 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: "" | |
| artifact: 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 | |
| runs-on: ${{ matrix.pjrt.runs_on }} | |
| steps: | |
| - name: "Checking out repository" | |
| uses: actions/checkout@v4 | |
| - name: Setup Bazel | |
| run: | | |
| wget https://github.com/bazelbuild/bazelisk/releases/download/v1.20.0/bazelisk-linux-amd64 -O ~/bin/bazel | |
| chmod +x ~/bin/bazel | |
| - name: Force remove container | |
| run: | | |
| docker rm -f pjrt-${{ matrix.pjrt.target }} | |
| - name: "Build ${{ matrix.pjrt.target }} ${{ matrix.pjrt.platform }} image" | |
| run: | | |
| ~/bin/bazel run //${{ matrix.pjrt.target }}:${{ matrix.pjrt.target }}_stripped_tarball | |
| - name: "Run ${{ matrix.pjrt.target }} ${{ matrix.pjrt.platform }} image" | |
| run: | | |
| docker run \ | |
| --name=pjrt-${{ matrix.pjrt.target }} \ | |
| -w=/xla \ | |
| --net=host \ | |
| -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 }} --repo_env=HERMETIC_PYTHON_VERSION=3.11 \ | |
| ${{ matrix.pjrt.bazel_target }} | |
| - name: "Retrieve ${{ matrix.pjrt.target }} ${{ matrix.pjrt.platform }} artifacts" | |
| 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: 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","pjrt-osx-artifacts"] | |
| runs-on: ['runs-on','runner=2cpu-linux-x64'] | |
| 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-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 |