build: Optimize caching (#316)

Signed-off-by: Oliver Koenig <[email protected]>

Dockerfile

@@ -4,7 +4,7 @@
 #
 # To update NeMo-Aligner from a pre-built NeMo-Framework container:
 #
-#   docker buildx build --target=aligner-bump --build-arg=BASE_IMAGE=nvcr.io/nvidia/nemo:24.07 -t aligner:latest .
+#   docker buildx build --target=aligner-bump -t aligner:latest .
 #
 
 # Number of parallel threads for compute heavy build jobs
@@ -18,32 +18,26 @@ ARG MLM_TAG=a3fe0c75df82218901fa2c3a7c9e389aa5f53182  # On: core_r0.8.0
 ARG ALIGNER_COMMIT=main
 ARG TRTLLM_VERSION=v0.10.0
 ARG PROTOBUF_VERSION=4.24.4
-
 ARG BASE_IMAGE=nvcr.io/nvidia/pytorch:24.03-py3
 
 FROM ${BASE_IMAGE} AS aligner-bump
-
 ARG ALIGNER_COMMIT
-
 WORKDIR /opt
-
 # NeMo Aligner
 RUN <<"EOF" bash -exu
 if [[ ! -d NeMo-Aligner ]]; then
     git clone https://github.com/NVIDIA/NeMo-Aligner.git
     cd NeMo-Aligner
     git checkout $ALIGNER_COMMIT
-    pip install --no-deps -e .
-    cd -
+else
+    cd NeMo-Aligner
 fi
-cd NeMo-Aligner
-git fetch -a
-git checkout -f ${ALIGNER_COMMIT}
-git pull
-EOF
 
-FROM aligner-bump as final
+pip install --no-deps -e .
+EOF
 
+FROM ${BASE_IMAGE} as final
+WORKDIR /opt
 # needed in case git complains that it can't detect a valid email, this email is fake but works
 RUN git config --global user.email "[email protected]"
 # install TransformerEngine
@@ -107,6 +101,10 @@ RUN curl -s https://packagecloud.io/install/repositories/github/git-lfs/script.d
     apt-get install git-lfs && \
     git lfs install
 
+COPY --from=aligner-bump /opt/NeMo-Aligner /opt/NeMo-Aligner
+RUN cd /opt/NeMo-Aligner && \
+    pip install --no-deps -e .
+
 # TRTLLM
 ARG TRTLLM_VERSION
 RUN git clone https://github.com/NVIDIA/TensorRT-LLM.git && \