docker-rocm.yaml

name: TritonBench Nightly ROCM Docker Build
on:
  pull_request:
    paths:
      - .github/workflows/docker-rocm.yaml
      - docker/tritonbench-rocm-nightly.dockerfile
  workflow_dispatch:
    inputs:
      nightly_date:
        description: "PyTorch nightly version"
        required: false
env:
  CONDA_ENV: "tritonbench"
  SETUP_SCRIPT: "/workspace/setup_instance.sh"

jobs:
  build-push-docker:
    if: ${{ github.repository_owner == 'pytorch-labs' }}
    runs-on: 32-core-ubuntu
    environment: docker-s3-upload
    steps:
      - name: Checkout
        uses: actions/checkout@v3
        with:
          path: tritonbench
      - name: Login to GitHub Container Registry
        if: github.event_name != 'pull_request'
        uses: docker/login-action@v2
        with:
          registry: ghcr.io
          username: pytorch-labs
          password: ${{ secrets.TRITONBENCH_ACCESS_TOKEN }}
      - name: Build TritonBench nightly docker
        run: |
          set -x
          export NIGHTLY_DATE="${{ github.event.inputs.nightly_date }}"
          cd tritonbench/docker
          # branch name is github.head_ref when triggered by pull_request
          # and it is github.ref_name when triggered by workflow_dispatch
          branch_name=${{ github.head_ref || github.ref_name }}
          docker build . --build-arg TRITONBENCH_BRANCH="${branch_name}" --build-arg FORCE_DATE="${NIGHTLY_DATE}" \
              -f tritonbench-rocm-nightly.dockerfile -t ghcr.io/pytorch-labs/tritonbench:rocm-latest
          # Extract pytorch version from the docker
          PYTORCH_VERSION=$(docker run -e SETUP_SCRIPT="${SETUP_SCRIPT}" ghcr.io/pytorch-labs/tritonbench:rocm-latest bash -c '. "${SETUP_SCRIPT}"; python -c "import torch; print(torch.__version__)"')
          export DOCKER_TAG=$(awk '{match($0, /dev[0-9]+/, arr); print arr[0]}' <<< "${PYTORCH_VERSION}")
          docker tag ghcr.io/pytorch-labs/tritonbench:rocm-latest ghcr.io/pytorch-labs/tritonbench:rocm-${DOCKER_TAG}
      - name: Push docker to remote
        if: github.event_name != 'pull_request'
        run: |
          # Extract pytorch version from the docker
          PYTORCH_VERSION=$(docker run -e SETUP_SCRIPT="${SETUP_SCRIPT}" ghcr.io/pytorch-labs/tritonbench:latest bash -c '. "${SETUP_SCRIPT}"; python -c "import torch; print(torch.__version__)"')
          export DOCKER_TAG=$(awk '{match($0, /dev[0-9]+/, arr); print arr[0]}' <<< "${PYTORCH_VERSION}")
          docker push ghcr.io/pytorch-labs/tritonbench:rocm-${DOCKER_TAG}
          docker push ghcr.io/pytorch-labs/tritonbench:rocm-latest
concurrency:
  group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.sha }}-${{ github.event_name == 'workflow_dispatch' }}
  cancel-in-progress: true