name: docker-builds

on:
  workflow_dispatch:
  push:
    branches:
      - main
      - release/*
    # A pull request opts in by applying the ciflow/docker label, which pushes
    # this tag. Building on every PR that touches .ci/docker would hold an
    # image builder for each one.
    tags:
      - ciflow/docker/*
    paths:
      - .ci/docker/**
      - .github/workflows/docker-builds.yml
  schedule:
    - cron: 1 3 * * 3

concurrency:
  group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.sha }}-${{ github.event_name == 'workflow_dispatch' }}
  cancel-in-progress: true

# The build assumes an AWS role by OIDC to push to ECR, which needs a token the
# workflow has to ask for. contents is named explicitly because naming any
# permission zeroes the ones left out, and checkout needs it.
permissions:
  id-token: write
  contents: read

jobs:
  docker-build:
    # Skip scheduled runs on forks, where they would only fail and email the fork owner
    if: github.repository_owner == 'pytorch' || github.event_name != 'schedule'
    strategy:
      fail-fast: false
      matrix:
        include:
          - docker-image-name: torchtitan-ubuntu-22.04-clang12
          - docker-image-name: torchtitan-ubuntu-22.04-clang12
            image-tag-prefix: rl
          - docker-image-name: torchtitan-rocm-ubuntu-22.04-clang12
    # BuildKit runs out of cluster on the OSDC BuildKit pool, so this runner only
    # streams the (small) .ci/docker context and waits.
    runs-on: mt-l-x86iavx512-8-64
    container:
      image: ghcr.io/actions/actions-runner:latest
    timeout-minutes: 240
    env:
      DOCKER_IMAGE: 308535385114.dkr.ecr.us-east-1.amazonaws.com/torchtitan/${{ matrix.docker-image-name }}
    steps:
      - name: Checkout the repo
        uses: actions/checkout@v7

      - name: Configure AWS credentials
        uses: aws-actions/configure-aws-credentials@v6
        with:
          role-to-assume: arn:aws:iam::308535385114:role/arc
          aws-region: us-east-1
          role-duration-seconds: 18000

      # buildx forwards the client's registry auth to the remote builder, so the
      # push at the end of build.sh authenticates with what this step writes.
      - name: Login to ECR
        uses: aws-actions/amazon-ecr-login@v2
        with:
          registries: "308535385114"

      - name: Compute the image tag
        id: tag
        shell: bash
        run: |
          set -eux
          git config --global --add safe.directory "${GITHUB_WORKSPACE}"
          # The tag set-matrix derives for consumers, so the test jobs resolve
          # the image this build pushes.
          DOCKER_TAG=$(git rev-parse HEAD:.ci/docker)
          DOCKER_TAG="${{ matrix.image-tag-prefix && 'rl-' || '' }}${DOCKER_TAG}"
          echo "docker-image=${DOCKER_IMAGE}:${DOCKER_TAG}" >> "${GITHUB_OUTPUT}"

      - name: Build and push to ECR
        uses: pytorch/test-infra/.github/actions/docker-build-remote-buildkit@main
        with:
          command: |
            cd .ci/docker
            REMOTE_BUILDKIT=1 ./build.sh ${{ matrix.docker-image-name }}${{ matrix.image-tag-prefix && ':rl' || '' }} \
              -t ${{ steps.tag.outputs.docker-image }}
