From 8c329824d4313ee721569e813412739c21c88662 Mon Sep 17 00:00:00 2001
From: Chuck Tang <chuck.tang@databricks.com>
Date: Wed, 25 Sep 2024 16:34:52 -0700
Subject: [PATCH] commit change

---
 .github/workflows/docker.yaml | 10 ++++++++--
 Dockerfile                    | 21 +++++++++++++++++++++
 2 files changed, 29 insertions(+), 2 deletions(-)
 create mode 100644 Dockerfile

diff --git a/.github/workflows/docker.yaml b/.github/workflows/docker.yaml
index ecdbc036..164b5ccb 100644
--- a/.github/workflows/docker.yaml
+++ b/.github/workflows/docker.yaml
@@ -17,14 +17,20 @@ jobs:
     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
+          base_image: mosaicml/pytorch:2.4.0_cu124-python3.11-ubuntu20.04-aws
           dep_groups: "[all]"
     steps:
-    
+
     - name: Checkout
       uses: actions/checkout@v3
 
diff --git a/Dockerfile b/Dockerfile
new file mode 100644
index 00000000..2496ddb8
--- /dev/null
+++ b/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