diff --git a/.github/workflows/release.yaml b/.github/workflows/release.yaml index c09f9bb7a5..3617732c8f 100644 --- a/.github/workflows/release.yaml +++ b/.github/workflows/release.yaml @@ -50,3 +50,67 @@ jobs: 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: mosaic-8wide + if: github.repository_owner == 'mosaicml' + steps: + - name: Checkout source + uses: actions/checkout@v3 + + - 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: Define Docker tags + id: define-tags + run: | + BRANCH_NAME="${{ github.ref_name }}" + TAG_NAME=$(echo "${BRANCH_NAME}" | sed 's/\//_/g') + echo "BRANCH_NAME=${BRANCH_NAME}" >> $GITHUB_ENV + + 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 + + + - name: Build and push AWS Docker image + uses: docker/build-push-action@v3 + with: + context: . + file: 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 + DEP_GROUPS=[all] + KEEP_FOUNDRY=true + + - name: Build and push Docker image + uses: docker/build-push-action@v3 + with: + context: . + file: 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 + DEP_GROUPS=[all] + KEEP_FOUNDRY=true diff --git a/Dockerfile b/Dockerfile index ca52532395..a9d44bfa27 100644 --- a/Dockerfile +++ b/Dockerfile @@ -7,6 +7,7 @@ FROM $BASE_IMAGE ARG BRANCH_NAME ARG DEP_GROUPS ARG TE_COMMIT +ARG KEEP_FOUNDRY=false ENV TORCH_CUDA_ARCH_LIST="8.0 8.6 8.7 8.9 9.0" @@ -21,5 +22,9 @@ RUN NVTE_FRAMEWORK=pytorch CMAKE_BUILD_PARALLEL_LEVEL=4 MAX_JOBS=4 pip install g # 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${DEP_GROUPS}" -RUN pip uninstall -y llm-foundry -RUN rm -rf llm-foundry + +# Conditionally uninstall llm-foundry and remove its directory +RUN if [ "$KEEP_FOUNDRY" != "true" ]; then \ + pip uninstall -y llm-foundry && \ + rm -rf /llm-foundry; \ + fi