_build.yml

name: xla-buld
on:
  workflow_call:
    inputs:
      gcr-docker-image:
        required: true
        type: string
        description: Base image for builds
      ecr-docker-image-base:
        required: true
        type: string
        description: Container registry to upload image to
      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 }}
    timeout-minutes: 240
    outputs:
      docker-image: ${{ steps.upload-docker-image.outputs.docker-image }}
    env:
      ECR_DOCKER_IMAGE_BASE: ${{ inputs.ecr-docker-image-base }}
      GCR_DOCKER_IMAGE: ${{ inputs.gcr-docker-image }}
      WORKDIR: /var/lib/jenkins/workspace
      SCCACHE_BUCKET: ossci-compiler-cache-circleci-v2
      GCLOUD_SERVICE_KEY: ${{ secrets.gcloud-service-key }}
      XLA_CUDA: ${{ inputs.cuda }}
      BAZEL_JOBS: 16
    steps:
      - name: Setup Linux
        uses: pytorch/test-infra/.github/actions/setup-linux@main
      - name: Setup SSH (Click me for login details)
        uses: pytorch/test-infra/.github/actions/setup-ssh@main
        with:
          github-secret: ${{ secrets.GITHUB_TOKEN }}
          instructions: |
            Build is  done inside the container, to start an interactive session run:
              docker exec -it $(docker container ps --format '{{.ID}}') bash
      - name: Checkout repo
        uses: actions/checkout@v3
      - name: Download docker image from GCR
        shell: bash
        run: docker pull "${GCR_DOCKER_IMAGE}"
      - name: Stage image to ECR
        shell: bash
        run: |
            # This is to stage PyTorch/XLA base image for use in the upstream.
            # To allow the upstream workflow to access PyTorch/XLA build images, we
            # need to have them in the ECR. This is not expensive, and only pushes it
            # if image layers are not present in the repo.
            # Note: disable the following 2 lines while testing a new image, so we do not
            # push to the upstream.
            docker tag "${GCR_DOCKER_IMAGE}" "${ECR_DOCKER_IMAGE_BASE}:v1.1-lite" >/dev/null
            docker push "${ECR_DOCKER_IMAGE_BASE}:v1.1-lite" >/dev/null
      - name: Start the container
        shell: bash
        run: |
          pid=$(docker run --privileged -t -d -w "$WORKDIR" "${GCR_DOCKER_IMAGE}")
          docker exec -u jenkins "${pid}" sudo chown -R jenkins "${WORKDIR}"
          docker cp "${GITHUB_WORKSPACE}/." "$pid:$WORKDIR"
          echo "pid=${pid}" >> "${GITHUB_ENV}"

      - name: Prepare build env
        shell: bash
        run: |
          echo "declare -x SCCACHE_BUCKET=${SCCACHE_BUCKET}" | docker exec -i "${pid}" sh -c "cat >> env"
          echo "declare -x XLA_CUDA=${XLA_CUDA}" | docker exec -i "${pid}" sh -c "cat >> xla_env"
          echo "declare -x BAZEL_JOBS=${BAZEL_JOBS}" | docker exec -i "${pid}" sh -c "cat >> xla_env"
          echo "declare -x BAZEL_REMOTE_CACHE=1" | docker exec -i "${pid}" sh -c "cat >> xla_env"
          echo "${GCLOUD_SERVICE_KEY}" | docker exec -i "${pid}" sh -c "cat >> default_credentials.json"

      - name: Build
        shell: bash
        run: |
          docker exec --privileged -u jenkins "${pid}" bash -c ".circleci/build.sh"
      - name: Cleanup build env
        shell: bash
        run: |
           docker exec "${pid}" rm default_credentials.json /tmp/pytorch/xla/default_credentials.json

      - name: Push built docker image to ECR
        id: upload-docker-image
        shell: bash
        run: |
          export COMMIT_DOCKER_IMAGE="${ECR_DOCKER_IMAGE_BASE}:latest-${GITHUB_SHA}"
          time docker commit "${pid}" "${COMMIT_DOCKER_IMAGE}"
          time docker push "${COMMIT_DOCKER_IMAGE}"
          echo "docker-image=${COMMIT_DOCKER_IMAGE}" >> "${GITHUB_OUTPUT}"

      - name: Teardown Linux
        uses: pytorch/test-infra/.github/actions/teardown-linux@main
        if: always()