build.yaml

name: PJRT GPU library
on:
  push:
    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: 7018a84922e775445bcae422363703d2abe8d571

      - 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/tar-action@v1.2.0
        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/tar-action@v1.2.0
        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