feat: Support caching nixl build stage (#147)

container/Dockerfile.vllm

@@ -4,6 +4,16 @@
 ARG BASE_IMAGE="nvcr.io/nvidia/cuda-dl-base"
 ARG BASE_IMAGE_TAG="25.01-cuda12.8-devel-ubuntu24.04"
 
+FROM ${BASE_IMAGE}:${BASE_IMAGE_TAG} AS nixl_base
+WORKDIR /opt/nixl
+# Add a cache hint that only changes when the nixl commit changes
+ARG NIXL_COMMIT
+# This line acts as a cache key - it only changes when NIXL_COMMIT changes
+RUN echo "NIXL commit: ${NIXL_COMMIT}" > /opt/nixl/commit.txt
+# Copy the nixl source
+COPY --from=nixl . .
+
+# Main build stage
 FROM ${BASE_IMAGE}:${BASE_IMAGE_TAG} AS dev
 
 USER root
@@ -117,8 +127,9 @@ ENV CPATH=/usr/local/ompi/include:$CPATH
 ENV PATH=/usr/local/ompi/bin:$PATH
 ENV PKG_CONFIG_PATH=/usr/local/ompi/lib/pkgconfig:$PKG_CONFIG_PATH
 
-COPY --from=nixl . /opt/nixl
-
+# Copy nixl source, and use commit hash as cache hint
+COPY --from=nixl_base /opt/nixl /opt/nixl
+COPY --from=nixl_base /opt/nixl/commit.txt /opt/nixl/commit.txt
 RUN cd /opt/nixl && \
     mkdir build && \
     meson setup build/ --prefix=/usr/local/nixl && \

container/build.sh

@@ -218,7 +218,7 @@ get_options() {
 
     if [ -z "$TAG" ]; then
         TAG="--tag dynamo:${VERSION}-${FRAMEWORK,,}"
-        if [ ! -z ${TARGET} ]; then
+        if [ ! -z "${TARGET}" ]; then
             TAG="${TAG}-${TARGET}"
         fi
     fi
@@ -284,28 +284,34 @@ elif [[ $FRAMEWORK == "TENSORRTLLM" ]]; then
 fi
 
 if [[ $FRAMEWORK == "VLLM" ]]; then
-    TEMP_DIR=$(mktemp -d)
-
-    # Clean up temp directory on script exit
-    trap 'rm -rf "$TEMP_DIR"' EXIT
+    NIXL_DIR="/tmp/nixl/nixl_src"
 
     # Clone original NIXL to temp directory
-
+    if [ -d "$NIXL_DIR" ]; then
+        echo "Warning: $NIXL_DIR already exists, skipping clone"
+    else
         if [ ! -z ${GITHUB_TOKEN} ]; then
-        git clone https://oauth2:${GITHUB_TOKEN}@github.com/${NIXL_REPO} "$TEMP_DIR/nixl_src"
+            git clone https://oauth2:${GITHUB_TOKEN}@github.com/${NIXL_REPO} "$NIXL_DIR"
         else
             # Try HTTPS first with credential prompting disabled, fall back to SSH if it fails
-        if ! GIT_TERMINAL_PROMPT=0 git clone https://github.com/${NIXL_REPO} "$TEMP_DIR/nixl_src"; then
+            if ! GIT_TERMINAL_PROMPT=0 git clone https://github.com/${NIXL_REPO} "$NIXL_DIR"; then
                 echo "HTTPS clone failed, falling back to SSH..."
-            git clone git@github.com:${NIXL_REPO} "$TEMP_DIR/nixl_src"
+                git clone git@github.com:${NIXL_REPO} "$NIXL_DIR"
+            fi
         fi
     fi
 
-    cd "$TEMP_DIR/nixl_src"
+    cd "$NIXL_DIR"
+    if ! git checkout ${NIXL_COMMIT}; then
+        echo "ERROR: Failed to checkout NIXL commit ${NIXL_COMMIT}. The cached directory may be out of date."
+        echo "Please delete $NIXL_DIR and re-run the build script."
+        exit 1
+    fi
 
-    git checkout ${NIXL_COMMIT}
+    BUILD_CONTEXT_ARG+=" --build-context nixl=$NIXL_DIR"
 
-    BUILD_CONTEXT_ARG+=" --build-context nixl=$TEMP_DIR/nixl_src"
+    # Add NIXL_COMMIT as a build argument to enable caching
+    BUILD_ARGS+=" --build-arg NIXL_COMMIT=${NIXL_COMMIT} "
 fi
 
 # BUILD DEV IMAGE