_build_plugin.yml

name: build-cuda-plugin
on:
  workflow_call:
    inputs:
      dev-image:
        required: true
        type: string
        description: Base image for builds
      runner:
        required: false
        type: string
        description: Runner type for the test
        default: linux.12xlarge
      cuda:
        required: false
        type: string
        description: Whether to build XLA with CUDA
        default: 1

    secrets:
      gcloud-service-key:
        required: true
        description: Secret to access Bazel build cache

    outputs:
      docker-image:
        value: ${{ jobs.build.outputs.docker-image }}
        description: The docker image containing the built PyTorch.
jobs:
  build:
    runs-on: ${{ inputs.runner }}
    container:
      image: ${{ inputs.dev-image }}
      volumes:
      - ${GITHUB_WORKSPACE}/dist/:/dist/
    env:
      GCLOUD_SERVICE_KEY: ${{ secrets.gcloud-service-key }}
      BAZEL_JOBS: 16
      BAZEL_REMOTE_CACHE: 1
    steps:
      - name: Checkout repo
        uses: actions/checkout@v3
      - name: Setup gcloud
        shell: bash
        run: |
          echo "${GCLOUD_SERVICE_KEY}" >> default_credentials.json
      - name: Build
        shell: bash
        run: |
          cd infra/ansible
          # TODO: don't clone everything again
          ansible-playbook playbook.yaml -e "stage=build_cuda arch=amd64 accelerator=cuda" --skip-tags=install_deps
      - name: actions/upload-artifact@v4
        with:
          name: cuda-plugin
          path: ${GITHUB_WORKSPACE}/dist/*.whl