Skip to content

Commit

Permalink
Adding docker images for diffusion repo (#174)
Browse files Browse the repository at this point in the history
  • Loading branch information
j316chuck authored Sep 26, 2024
1 parent c1f953f commit 6acffcd
Show file tree
Hide file tree
Showing 3 changed files with 108 additions and 0 deletions.
86 changes: 86 additions & 0 deletions .github/workflows/docker.yaml
Original file line number Diff line number Diff line change
@@ -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 }}
1 change: 1 addition & 0 deletions .github/workflows/pytest-cpu.yaml
Original file line number Diff line number Diff line change
Expand Up @@ -62,3 +62,4 @@ jobs:
with:
name: coverage-${{ github.sha }}-${{ inputs.name }}
path: .coverage
include-hidden-files: true
21 changes: 21 additions & 0 deletions Dockerfile
Original file line number Diff line number Diff line change
@@ -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

0 comments on commit 6acffcd

Please sign in to comment.