name: Docker
on:
  push:
    branches:
    - main
  workflow_dispatch: {}
jobs:
  docker-build:
    runs-on: ubuntu-latest
    if: github.repository_owner == 'mosaicml'
    strategy:
      matrix:
        include:
        - name: '1.13.1_cu117'
          base_image: mosaicml/pytorch:1.13.1_cu117-python3.10-ubuntu20.04
        - name: '2.0.1_cu118'
          base_image: mosaicml/pytorch:2.0.1_cu118-python3.10-ubuntu20.04

    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}

    - name: Build and Push the Docker Image
      uses: docker/build-push-action@v3
      with:
        context: .
        tags: mosaicml/llm-foundry:${{ matrix.name }}-latest,
          mosaicml/llm-foundry:${{ matrix.name }}-${{ env.IMAGE_TAG }}
        push: true
        cache-from: type=registry,ref=mosaicml/llm-foundry:${{ matrix.name }}-buildcache
        cache-to: type=registry,ref=mosaicml/llm-foundry:${{ matrix.name }}-buildcache,mode=max
        build-args: BASE_IMAGE=${{ matrix.base_image }}