split vllm dependencies and code

Author: A9isha

docker_build_dependency_image.sh

@@ -142,18 +142,14 @@ if [[ ${INSTALL_POST_TRAINING} -eq 1 ]] ; then
     exit 1
   fi
 
-  # To install tpu_commons from a local path, we copy it into the build context, excluding __pycache__.
-  # This assumes vllm, tunix, tpu_commons is a sibling directory to the current one (maxtext).
+  # To install from local paths, we copy vllm and tpu_commons into the build context.
+  # This assumes vllm and tpu_commons are sibling directories to the current one (maxtext).
+  echo "Copying local vllm and tpu_commons directories into the build context..."
   rsync -a --exclude='__pycache__' ../tpu_commons .
-  # To install vllm from a local path, we copy it into the build context, excluding __pycache__.
-  # This assumes vllm is a sibling directory to the current one (maxtext).
   rsync -a --exclude='__pycache__' ../vllm .
 
-  # rsync -a --exclude='__pycache__' ../tunix .
-
-  # The cleanup is set to run even if the build fails to remove the copied directory.
-  # trap "rm -rf ./tpu_commons ./vllm ./tunix" EXIT INT TERM
-  trap "rm -rf ./tpu_commons ./vllm " EXIT INT TERM
+  # The cleanup is set to run even if the build fails to remove the copied directories.
+  trap "echo 'Cleaning up copied directories...' && rm -rf ./tpu_commons ./vllm" EXIT INT TERM
 
   docker build \
     --network host \

maxtext_post_training_dependencies.Dockerfile

@@ -15,7 +15,6 @@
 ARG BASEIMAGE
 FROM ${BASEIMAGE}
 ARG MODE
-
 ENV MODE=$MODE
 
 RUN echo "Installing Post-Training dependencies (vLLM, tpu-common, tunix) with MODE=${MODE}"
@@ -24,53 +23,52 @@ RUN echo "Installing Post-Training dependencies (vLLM, tpu-common, tunix) with M
 # Uninstall existing jax to avoid conflicts
 # RUN pip uninstall -y jax jaxlib libtpu
 
-RUN pip install aiohttp==3.12.15
+# --- STAGE 1: Install Static Dependencies ---
+# Install any packages *not* defined in your project dependency files
+RUN --mount=type=cache,target=/root/.cache/pip pip install \
+    aiohttp==3.12.15\
+    keyring \
+    keyrings.google-artifactregistry-auth \
+    numba==0.61.2
 
-# Install Python packages that enable pip to authenticate with Google Artifact Registry automatically.
-RUN pip install keyring keyrings.google-artifactregistry-auth
+# --- STAGE 2: Install Project Dependencies (The Main Cached Layer) ---
 
-RUN pip install numba==0.61.2
+# Copy *only* the dependency definition files.
+# This assumes vllm and tpu_commons are in the build context, copied from the parent directory.
+COPY vllm/requirements/tpu.txt /tmp/
+COPY vllm/requirements/build.txt /tmp/
+COPY vllm/requirements/common.txt /tmp/
+COPY tpu_commons/requirements.txt /tmp/
 
-COPY vllm /vllm
-RUN VLLM_TARGET_DEVICE="tpu" pip install -e /vllm --no-cache-dir --pre \
-    --extra-index-url https://pypi.org/simple/ \
-    --extra-index-url https://us-python.pkg.dev/ml-oss-artifacts-published/jax/simple/ \
-    --extra-index-url https://download.pytorch.org/whl/nightly/cpu \
-    --find-links https://storage.googleapis.com/jax-releases/libtpu_releases.html \
-    --find-links https://storage.googleapis.com/libtpu-wheels/index.html \
-    --find-links https://storage.googleapis.com/libtpu-releases/index.html \
-    --find-links https://storage.googleapis.com/jax-releases/jax_nightly_releases.html \
-    --find-links https://storage.googleapis.com/jax-releases/jaxlib_nightly_releases.html 
+# Run the full dependency installation.
+# This entire layer is cached and will *only* be rebuilt if
+# these .txt files change.
+RUN --mount=type=cache,target=/root/.cache/pip bash -c ' \
+    # Set the target device so pip installs the right JAX/libtpu
+    export VLLM_TARGET_DEVICE="tpu" && \
+    pip install -r /tmp/tpu.txt -r /tmp/build.txt -r /tmp/common.txt -r /tmp/requirements.txt --no-cache-dir --pre \
+        --extra-index-url https://pypi.org/simple/ \
+        --extra-index-url https://us-python.pkg.dev/ml-oss-artifacts-published/jax/simple/ \
+        --extra-index-url https://download.pytorch.org/whl/nightly/cpu \
+        --find-links https://storage.googleapis.com/jax-releases/libtpu_releases.html \
+        --find-links https://storage.googleapis.com/libtpu-wheels/index.html \
+        --find-links https://storage.googleapis.com/libtpu-releases/index.html \
+        --find-links https://storage.googleapis.com/jax-releases/jax_nightly_releases.html \
+        --find-links https://storage.googleapis.com/jax-releases/jaxlib_nightly_releases.html'
 
-# Install tpu-commons from local source
-COPY tpu_commons /tpu_commons
-RUN pip install -e /tpu_commons --no-cache-dir --pre \
-    --extra-index-url https://pypi.org/simple/ \
-    --extra-index-url https://us-python.pkg.dev/ml-oss-artifacts-published/jax/simple/ \
-    --find-links https://storage.googleapis.com/jax-releases/libtpu_releases.html
+# --- STAGE 3: Install Project Source Code ---
 
-# # Install vLLM for Jax and TPUs from the artifact registry
-# RUN VLLM_TARGET_DEVICE="tpu" pip install --no-cache-dir --pre \
-#     --index-url https://us-python.pkg.dev/cloud-tpu-images/maxtext-rl/simple/ \
-#     --extra-index-url https://pypi.org/simple/ \
-#     --extra-index-url https://us-python.pkg.dev/ml-oss-artifacts-published/jax/simple/ \
-#     --extra-index-url https://download.pytorch.org/whl/nightly/cpu \
-#     --find-links https://storage.googleapis.com/jax-releases/libtpu_releases.html \
-#     --find-links https://storage.googleapis.com/libtpu-wheels/index.html \
-#     --find-links https://storage.googleapis.com/libtpu-releases/index.html \
-#     --find-links https://storage.googleapis.com/jax-releases/jax_nightly_releases.html \
-#     --find-links https://storage.googleapis.com/jax-releases/jaxlib_nightly_releases.html \
-#     vllm==0.11.1rc1.dev292+g1b86bd8e1.tpu
+# Now, copy the full source code. This invalidates cache frequently,
+# but the next step is fast.
+COPY vllm /vllm/
+COPY tpu_commons /tpu_commons/
 
-# # Install tpu-commons from the artifact registry
-# RUN pip install --no-cache-dir --pre \
-#     --index-url https://us-python.pkg.dev/cloud-tpu-images/maxtext-rl/simple/ \
-#     --extra-index-url https://pypi.org/simple/ \
-#     --extra-index-url https://us-python.pkg.dev/ml-oss-artifacts-published/jax/simple/ \
-#     --find-links https://storage.googleapis.com/jax-releases/libtpu_releases.html \
-#     tpu-commons==0.1.2
+# Install in editable mode. This is lightning-fast because all
+# dependencies were installed and cached in STAGE 2.
+RUN --mount=type=cache,target=/root/.cache/pip VLLM_TARGET_DEVICE="tpu" pip install -e /vllm/ -e /tpu_commons/
 
 RUN if [ "$MODE" = "post-training-experimental" ]; then \
+    echo "MODE=grpo-experimental: Re-installing JAX/libtpu"; \
     pip uninstall -y jax jaxlib libtpu && \
     pip install --pre -U jax jaxlib -i https://us-python.pkg.dev/ml-oss-artifacts-published/jax/simple/ && \
     pip install -U --pre libtpu -f https://storage.googleapis.com/jax-releases/libtpu_releases.html; \