docker.yaml

name: Docker
on:
  push:
    branches:
    - main
  pull_request:
    branches:
    - main
    paths:
    - Dockerfile
    - .github/workflows/docker.yaml
  workflow_dispatch: {}
jobs:
  docker-build:
    runs-on: ubuntu-latest
    if: github.repository_owner == 'mosaicml'
    strategy:
      matrix:
        include:
        - name: "2.1.0_cu121_flash2"
          base_image: mosaicml/pytorch:2.1.2_cu121-python3.10-ubuntu20.04
          dep_groups: "[gpu-flash2]"
        - name: "2.1.0_cu121_flash2_aws"
          base_image: mosaicml/pytorch:2.1.2_cu121-python3.10-ubuntu20.04-aws
          dep_groups: "[gpu-flash2]"
    steps:
    - name: Maximize Build Space on Worker
      uses: easimon/maximize-build-space@v4
      with:
        overprovision-lvm: true
        remove-dotnet: true
        remove-android: true
        remove-haskell: true

    - name: Checkout
      uses: actions/checkout@v3

    - name: Setup QEMU
      uses: docker/setup-qemu-action@v2

    - name: Setup Docker Buildx
      uses: docker/setup-buildx-action@v2

    - name: Login to DockerHub
      uses: docker/login-action@v2
      with:
        username: ${{ secrets.DOCKER_HUB_USERNAME }}
        password: ${{ secrets.DOCKER_HUB_PASSWORD }}

    - name: Calculate Docker Image Variables
      run: |
        set -euxo pipefail

        ###################
        # Calculate the tag
        ###################
        GIT_SHA=$(echo ${{ github.sha }} | cut -c1-7)
        echo "IMAGE_TAG=${GIT_SHA}" >> ${GITHUB_ENV}

        if [ "${{ github.event_name }}" == "pull_request" ]; then
          echo "Triggered by pull_request event."
          STAGING_REPO="mosaicml/ci-staging"
          IMAGE_TAG="${STAGING_REPO}:${{matrix.name}}-${GIT_SHA}"
          IMAGE_CACHE="${STAGING_REPO}:${{matrix.name}}-buildcache"
        else
          # Triggered by push or workflow_dispatch event
          echo "Triggered by ${{ github.event_name }} event, releasing to prod"
          PROD_REPO="mosaicml/llm-foundry"
          IMAGE_TAG="${PROD_REPO}:${{matrix.name}}-${GIT_SHA},${PROD_REPO}:${{matrix.name}}-latest"
          IMAGE_CACHE="${PROD_REPO}:${{matrix.name}}-buildcache"
        fi

        echo "IMAGE_TAG=${IMAGE_TAG}" >> ${GITHUB_ENV}
        echo "IMAGE_CACHE=${IMAGE_CACHE}" >> ${GITHUB_ENV}

    - name: Build and Push the Docker Image
      uses: docker/build-push-action@v3
      with:
        context: .
        tags: ${{ env.IMAGE_TAG }}
        push: true
        cache-from: type=registry,ref=${{ env.IMAGE_CACHE }}
        cache-to: type=registry,ref=${{ env.IMAGE_CACHE }},mode=max
        build-args: |
          BRANCH_NAME=${{ github.head_ref || github.ref_name }}
          BASE_IMAGE=${{ matrix.base_image }}
          DEP_GROUPS=${{ matrix.dep_groups }}