update nvcr dockefile to multistage to save space

Signed-off-by: Dushyant Behl <[email protected]>

Author: dushyantbehl

.github/workflows/image.yaml

@@ -2,8 +2,6 @@ name: Image
 on:
   push:
     branches: [ "main" ]
-  pull_request:
-    branches: [ "main" ]
 
 jobs:
   build:
@@ -16,13 +14,11 @@ jobs:
         run: |
           docker build -t fms-hf-tuning:main-nvcr-latest . -f build/nvcr.Dockerfile
       - name: Login to Quay.io
-        if: github.event_name == 'push'
         uses: docker/login-action@v3
         with:
           registry: quay.io
           username: ${{ secrets.QUAY_USERNAME }}
           password: ${{ secrets.QUAY_ROBOT_TOKEN }}
       - name: Push docker image for every commit to Quay.io as dev images
-        if: github.event_name == 'push'
         run: |
             docker push ${{ env.QUAY_REPOSITORY }}/fms-hf-tuning:main-nvcr-latest

build/nvcr.Dockerfile

@@ -20,37 +20,36 @@ ARG NVCR_IMAGE_VERSION=25.02-py3
 # This is based on what is inside the NVCR image already
 ARG PYTHON_VERSION=3.12
 
-## Base Layer ##################################################################
-FROM nvcr.io/nvidia/pytorch:${NVCR_IMAGE_VERSION} AS dev
+######################## BUILDER ########################
+FROM nvcr.io/nvidia/pytorch:${NVCR_IMAGE_VERSION} AS builder
 
 ARG USER=root
 ARG USER_UID=0
 ARG WORKDIR=/app
 ARG SOURCE_DIR=${WORKDIR}/fms-hf-tuning
 
 ARG ENABLE_FMS_ACCELERATION=true
-ARG ENABLE_AIM=true
-ARG ENABLE_MLFLOW=true
-ARG ENABLE_SCANNER=true
+ARG ENABLE_AIM=false
+ARG ENABLE_MLFLOW=false
+ARG ENABLE_SCANNER=false
 ARG ENABLE_CLEARML=true
 ARG ENABLE_TRITON_KERNELS=true
-ARG ENABLE_MAMBA_SUPPORT=true
 
 # Ensures to always build mamba_ssm from source
 ENV PIP_NO_BINARY=mamba-ssm,mamba_ssm
 
-RUN python -m pip install --upgrade pip
-
 # upgrade torch as the base layer contains only torch 2.7
-RUN pip install --upgrade --force-reinstall torch torchaudio torchvision --index-url https://download.pytorch.org/whl/cu128
+RUN python -m pip install --upgrade pip && \
+    pip install --upgrade setuptools && \
+    pip install --upgrade --force-reinstall torch torchaudio torchvision --index-url https://download.pytorch.org/whl/cu128
 
 # Install main package + flash attention
 COPY . ${SOURCE_DIR}
 RUN cd ${SOURCE_DIR}
 
-RUN pip install --upgrade pip setuptools wheel
-RUN pip install --no-cache-dir ${SOURCE_DIR}
-RUN pip install --user --no-build-isolation ${SOURCE_DIR}[flash-attn]
+RUN pip install --no-cache-dir ${SOURCE_DIR} && \
+    pip install --user --no-build-isolation ${SOURCE_DIR}[flash-attn] && \
+    pip install --no-cache-dir --no-build-isolation ${SOURCE_DIR}[mamba]
 
 # Optional extras
 RUN if [[ "${ENABLE_FMS_ACCELERATION}" == "true" ]]; then \
@@ -62,6 +61,12 @@ RUN if [[ "${ENABLE_FMS_ACCELERATION}" == "true" ]]; then \
         python -m fms_acceleration.cli install fms_acceleration_odm; \
     fi
 
+RUN if [[ "${ENABLE_TRITON_KERNELS}" == "true" ]]; then \
+        pip install --no-cache-dir "git+https://github.com/triton-lang/triton.git@main#subdirectory=python/triton_kernels"; \
+    fi
+RUN if [[ "${ENABLE_CLEARML}" == "true" ]]; then \
+        pip install --no-cache-dir ${SOURCE_DIR}[clearml]; \
+    fi
 RUN if [[ "${ENABLE_AIM}" == "true" ]]; then \
         pip install --no-cache-dir ${SOURCE_DIR}[aim]; \
     fi
@@ -71,15 +76,22 @@ RUN if [[ "${ENABLE_MLFLOW}" == "true" ]]; then \
 RUN if [[ "${ENABLE_SCANNER}" == "true" ]]; then \
         pip install --no-cache-dir ${SOURCE_DIR}[scanner-dev]; \
     fi
-RUN if [[ "${ENABLE_CLEARML}" == "true" ]]; then \
-        pip install --no-cache-dir ${SOURCE_DIR}[clearml]; \
-    fi
-RUN if [[ "${ENABLE_MAMBA_SUPPORT}" == "true" ]]; then \
-        pip install --no-cache-dir ${SOURCE_DIR}[mamba]; \
-    fi
-RUN if [[ "${ENABLE_TRITON_KERNELS}" == "true" ]]; then \
-        pip install --no-cache-dir "git+https://github.com/triton-lang/triton.git@main#subdirectory=python/triton_kernels"; \
-    fi
+
+# cleanup
+RUN rm -rf /root/.cache /tmp/* /opt/pytorch
+
+######################## RUNTIME ########################
+FROM nvcr.io/nvidia/pytorch:${NVCR_IMAGE_VERSION}
+
+WORKDIR ${WORKDIR}
+
+# Copy only Python site-packages + app
+COPY --from=builder /usr/local/lib/python3.12/dist-packages \
+                    /usr/local/lib/python3.12/dist-packages
+COPY --from=builder ${SOURCE_DIR} ${SOURCE_DIR}
+
+# Runtime cleanup
+RUN rm -rf /opt/pytorch /root/.cache /tmp/*
 
 RUN chmod -R g+rwX $WORKDIR /tmp
 RUN mkdir -p /.cache && chmod -R 777 /.cache