Changes to setup.py invalidate docker cache. Use branch name in docke…

…rfile (#930)

Co-authored-by: Daniel King <43149077+dakinggg@users.noreply.github.com>

.github/workflows/docker.yaml

@@ -88,5 +88,6 @@ jobs:
         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 }}

Dockerfile

@@ -4,10 +4,16 @@
 ARG BASE_IMAGE
 FROM $BASE_IMAGE
 
+ARG BRANCH_NAME
 ARG DEP_GROUPS
 
+# 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/llm-foundry/$BRANCH_NAME/setup.py setup.py
+RUN rm setup.py
+
 # Install and uninstall foundry to cache foundry requirements
-RUN git clone -b main https://github.com/mosaicml/llm-foundry.git
+RUN git clone -b $BRANCH_NAME https://github.com/mosaicml/llm-foundry.git
 RUN pip install --no-cache-dir "./llm-foundry${DEP_GROUPS}"
 RUN pip uninstall -y llm-foundry
 RUN rm -rf llm-foundry