Merge pull request #2748 from AI-Hypercomputer:anisha-update-local-docker-image

PiperOrigin-RevId: 836347842

Author: Google-ML-Automation

dependencies/dockerfiles/maxtext_post_training_local_dependencies.Dockerfile

@@ -33,22 +33,11 @@ RUN pip install -e /tunix --no-cache-dir
 
 
 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 VLLM_TARGET_DEVICE="tpu" pip install -e /vllm --no-cache-dir
 
 
 COPY tpu-inference /tpu-inference
-RUN pip install -e /tpu-inference --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
+RUN pip install -e /tpu-inference --no-cache-dir
 
 RUN pip install --no-deps qwix==0.1.4