Release docker img

.github/workflows/release.yaml

@@ -10,43 +10,111 @@ jobs:
   code-quality:
     uses: ./.github/workflows/code-quality.yaml
 
-  pypi-packaging:
-    name: Build and Publish llm-foundry PyPI Package
+  # rm for testing
+  # pypi-packaging:
+  #   name: Build and Publish llm-foundry PyPI Package
+  #   needs:
+  #   - code-quality
+  #   runs-on: linux-ubuntu-latest
+  #   steps:
+  #   - name: Checkout source
+  #     uses: actions/checkout@v3
+
+  #   - name: Set up Python
+  #     uses: actions/setup-python@v3
+  #     with:
+  #       python-version: "3.9"
+
+  #   - name: Build source and wheel distributions
+  #     run: |
+  #       if [[ "${{ github.ref }}" =~ refs\/tags\/v ]]; then
+  #         PYPI_PACKAGE_NAME="llm-foundry"
+  #       else
+  #         PYPI_PACKAGE_NAME="llm-foundry-test-$(date +%Y%m%d%H%M%S)"
+  #       fi
+
+  #       python -m pip install --upgrade build twine
+  #       python -m build
+  #       twine check --strict dist/*
+
+  #   - name: Publish 📦 to PyPI
+  #     uses: pypa/gh-action-pypi-publish@release/v1
+  #     if: contains(github.ref, 'refs/tags/v')
+  #     with:
+  #       user: __token__
+  #       password: ${{ secrets.PROD_PYPI_API_TOKEN }}
+
+  #   - name: Publish distribution 📦 to Test PyPI
+  #     uses: pypa/gh-action-pypi-publish@release/v1
+  #     if: contains(github.ref, 'refs/heads/') || contains(github.ref, 'refs/pull/')
+  #     with:
+  #       user: __token__
+  #       password: ${{ secrets.TEST_PYPI_API_TOKEN }}
+  #       repository_url: https://test.pypi.org/legacy/
+
+  build-docker:
+    name: Build llm-foundry Release Docker Image
     needs:
     - code-quality
-    runs-on: linux-ubuntu-latest
+    runs-on: mosaic-8wide
+    if: github.repository_owner == 'mosaicml'
     steps:
     - name: Checkout source
       uses: actions/checkout@v3
 
-    - name: Set up Python
-      uses: actions/setup-python@v3
+    - name: Setup Docker Buildx
+      uses: docker/setup-buildx-action@v2
+
+    - name: Login to DockerHub
+      uses: docker/login-action@v2
       with:
-        python-version: "3.9"
+        username: ${{ secrets.DOCKER_HUB_USERNAME }}
+        password: ${{ secrets.DOCKER_HUB_PASSWORD }}
 
-    - name: Build source and wheel distributions
+    - name: Define Docker tags
+      id: define-tags
       run: |
-        if [[ "${{ github.ref }}" =~ refs\/tags\/v ]]; then
-          PYPI_PACKAGE_NAME="llm-foundry"
+        BRANCH_NAME="${{ github.ref_name }}"
+        TAG_NAME=$(echo "${BRANCH_NAME}" | sed 's/\//_/g')
+        echo "BRANCH_NAME=${BRANCH_NAME}" >> $GITHUB_ENV
+
+        if [[ "${{ github.event_name }}" == "push" && "${{ github.ref }}" == refs/tags/v* ]]; then
+          # Official release (tag)
+          echo "DOCKER_TAG=mosaicml/llm-foundry:release_${TAG_NAME}" >> $GITHUB_ENV
+          echo "AWS_DOCKER_TAG=mosaicml/llm-foundry:release_${TAG_NAME}_aws" >> $GITHUB_ENV
+          echo "LATEST_TAG=mosaicml/llm-foundry:release-latest" >> $GITHUB_ENV
+          echo "AWS_LATEST_TAG=mosaicml/llm-foundry:release_aws-latest" >> $GITHUB_ENV
         else
-          PYPI_PACKAGE_NAME="llm-foundry-test-$(date +%Y%m%d%H%M%S)"
+          # Testing release (manual)
+          echo "DOCKER_TAG=mosaicml/llm-foundry:test_release_${TAG_NAME}" >> $GITHUB_ENV
+          echo "AWS_DOCKER_TAG=mosaicml/llm-foundry:test_release_${TAG_NAME}_aws" >> $GITHUB_ENV
         fi
 
-        python -m pip install --upgrade build twine
-        python -m build
-        twine check --strict dist/*
 
-    - name: Publish 📦 to PyPI
-      uses: pypa/gh-action-pypi-publish@release/v1
-      if: contains(github.ref, 'refs/tags/v')
+    - name: Build and push AWS Docker image
+      uses: docker/build-push-action@v3
       with:
-        user: __token__
-        password: ${{ secrets.PROD_PYPI_API_TOKEN }}
+        context: .
+        file: release.Dockerfile
+        push: true
+        tags: |
+          ${{ env.AWS_DOCKER_TAG }}
+          ${{ env.AWS_LATEST_TAG }}
+        build-args: |
+          BASE_IMAGE=mosaicml/pytorch:2.4.0_cu124-python3.11-ubuntu20.04-aws
+          BRANCH_NAME=${{ env.BRANCH_NAME }}
+          TE_COMMIT=901e5d2
 
-    - name: Publish distribution 📦 to Test PyPI
-      uses: pypa/gh-action-pypi-publish@release/v1
-      if: contains(github.ref, 'refs/heads/') || contains(github.ref, 'refs/pull/')
+    - name: Build and push Docker image
+      uses: docker/build-push-action@v3
       with:
-        user: __token__
-        password: ${{ secrets.TEST_PYPI_API_TOKEN }}
-        repository_url: https://test.pypi.org/legacy/
+        context: .
+        file: release.Dockerfile
+        push: true
+        tags: |
+          ${{ env.DOCKER_TAG }}
+          ${{ env.LATEST_TAG }}
+        build-args: |
+          BASE_IMAGE=mosaicml/pytorch:2.4.0_cu124-python3.11-ubuntu20.04
+          BRANCH_NAME=${{ env.BRANCH_NAME }}
+          TE_COMMIT=901e5d2

release.Dockerfile

@@ -0,0 +1,22 @@
+# Copyright 2022 MosaicML LLM Foundry authors
+# SPDX-License-Identifier: Apache-2.0
+
+ARG BASE_IMAGE
+FROM $BASE_IMAGE
+
+ARG BRANCH_NAME
+ARG TE_COMMIT
+
+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/llm-foundry/$BRANCH_NAME/setup.py setup.py
+RUN rm setup.py
+
+# Install TransformerEngine
+RUN NVTE_FRAMEWORK=pytorch CMAKE_BUILD_PARALLEL_LEVEL=4 MAX_JOBS=4 pip install git+https://github.com/NVIDIA/TransformerEngine.git@$TE_COMMIT
+
+# Install and uninstall foundry to cache foundry requirements
+RUN git clone -b $BRANCH_NAME https://github.com/mosaicml/llm-foundry.git
+RUN pip install --no-cache-dir "./llm-foundry[all]"