Merge pull request #2690 from AI-Hypercomputer:post_training_docker_fix

PiperOrigin-RevId: 832376901

Author: Google-ML-Automation

dependencies/dockerfiles/maxtext_post_training_dependencies.Dockerfile

@@ -20,19 +20,18 @@ ENV MODE=$MODE
 
 RUN echo "Installing Post-Training dependencies (vLLM, tpu-common, tunix) with MODE=${MODE}"
 
-
 # Uninstall existing jax to avoid conflicts
-RUN uv pip uninstall -y jax jaxlib libtpu
+RUN pip uninstall -y jax jaxlib libtpu
 
-RUN uv pip install aiohttp==3.12.15
+RUN pip install aiohttp==3.12.15
 
-RUN uv pip install numba==0.61.2
+RUN pip install numba==0.61.2
 
 # Install vLLM for Jax and TPUs
-RUN uv pip install vllm-tpu
+RUN pip install vllm-tpu
 
 RUN if [ "$MODE" = "post-training-experimental" ]; then \
-    uv pip uninstall -y jax jaxlib libtpu && \
-    uv pip install --pre -U jax jaxlib -i https://us-python.pkg.dev/ml-oss-artifacts-published/jax/simple/ && \
-    uv pip install -U --pre libtpu -f https://storage.googleapis.com/jax-releases/libtpu_releases.html; \
+    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; \
     fi

dependencies/scripts/docker_build_dependency_image.sh

@@ -150,7 +150,7 @@ if [[ -z ${LIBTPU_GCS_PATH+x} ]] ; then
                    -f "$MAXTEXT_REPO_ROOT"'/dependencies/dockerfiles/maxtext_db_dependencies.Dockerfile' \
                    -t ${LOCAL_IMAGE_NAME} .
     elif [[ ${INSTALL_POST_TRAINING} -eq 1 && ${DEVICE} == "tpu" ]]; then
-      echo "Installing MaxText stable mode dependencies for GRPO"
+      echo "Installing MaxText stable mode dependencies for post-training"
       docker build --network host --build-arg MODE=stable --build-arg JAX_VERSION=$JAX_VERSION \
                    --build-arg LIBTPU_GCS_PATH=$LIBTPU_GCS_PATH --build-arg DEVICE=$DEVICE \
                    -f "$MAXTEXT_REPO_ROOT"'/dependencies/dockerfiles/maxtext_dependencies.Dockerfile' \