Add docker images for diffusion repo

.github/workflows/docker.yaml

@@ -0,0 +1,86 @@
+name: Docker
+on:
+  push:
+    branches:
+    - main
+  pull_request:
+    branches:
+    - main
+    paths:
+    - Dockerfile
+    - .github/workflows/docker.yaml
+  workflow_dispatch: {}
+jobs:
+  docker-build:
+    runs-on: mosaic-8wide
+    if: github.repository_owner == 'mosaicml'
+    strategy:
+      matrix:
+        include:
+        - name: "2.1.2_cu121"
+          base_image: mosaicml/pytorch:2.1.2_cu121-python3.10-ubuntu20.04
+          dep_groups: "[all]"
+        - name: "2.1.2_cu121_aws"
+          base_image: mosaicml/pytorch:2.1.2_cu121-python3.10-ubuntu20.04-aws
+          dep_groups: "[all]"
+        - name: "2.4.0_cu124"
+          base_image: mosaicml/pytorch:2.4.0_cu124-python3.11-ubuntu20.04
+          dep_groups: "[all]"
+        - name: "2.4.0_cu124_aws"
+          base_image: mosaicml/pytorch:2.4.0_cu124-python3.11-ubuntu20.04-aws
+          dep_groups: "[all]"
+    steps:
+
+    - 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/diffusion"
+          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 }}

.github/workflows/pytest-cpu.yaml

@@ -62,3 +62,4 @@ jobs:
         with:
           name: coverage-${{ github.sha }}-${{ inputs.name }}
           path: .coverage
+          include-hidden-files: true

Dockerfile

@@ -0,0 +1,21 @@
+# Copyright 2022 MosaicML Diffusion authors
+# SPDX-License-Identifier: Apache-2.0
+
+ARG BASE_IMAGE
+FROM $BASE_IMAGE
+
+ARG BRANCH_NAME
+ARG DEP_GROUPS
+
+ENV TORCH_CUDA_ARCH_LIST="8.0 8.6 8.7 8.9 9.0"
+
+# Check for changes in setup.py.
+# If there are changes, the docker cache is invalidated and a fresh pip installation is triggered.
+ADD https://raw.githubusercontent.com/mosaicml/diffusion/$BRANCH_NAME/setup.py setup.py
+RUN rm setup.py
+
+# Install and uninstall diffusion to cache diffusion requirements
+RUN git clone -b $BRANCH_NAME https://github.com/mosaicml/diffusion.git
+RUN pip install --no-cache-dir "./diffusion${DEP_GROUPS}"
+RUN pip uninstall -y diffusion
+RUN rm -rf diffusion