From 73388c07a42233a1010595edaab73a9e7ab8d9a4 Mon Sep 17 00:00:00 2001 From: Woosuk Kwon Date: Thu, 8 Aug 2024 20:24:58 -0700 Subject: [PATCH] [TPU] Fix dockerfile.tpu (#7331) --- Dockerfile.tpu | 14 ++++---------- 1 file changed, 4 insertions(+), 10 deletions(-) diff --git a/Dockerfile.tpu b/Dockerfile.tpu index adebb8ab5a..ef0a53a7af 100644 --- a/Dockerfile.tpu +++ b/Dockerfile.tpu @@ -4,20 +4,14 @@ ARG BASE_IMAGE="us-central1-docker.pkg.dev/tpu-pytorch-releases/docker/xla:night FROM $BASE_IMAGE WORKDIR /workspace -# Install aiohttp separately to avoid build errors. -RUN pip install aiohttp -# Install NumPy 1 instead of NumPy 2. -RUN pip install "numpy<2" # Install the TPU and Pallas dependencies. -RUN pip install torch_xla[tpu] -f https://storage.googleapis.com/libtpu-releases/index.html -RUN pip install torch_xla[pallas] -f https://storage.googleapis.com/jax-releases/jax_nightly_releases.html -f https://storage.googleapis.com/jax-releases/jaxlib_nightly_releases.html - -# Fix FastAPI dependence -RUN pip install "starlette<0.38.0" +RUN python3 -m pip install torch_xla[tpu] -f https://storage.googleapis.com/libtpu-releases/index.html +RUN python3 -m pip install torch_xla[pallas] -f https://storage.googleapis.com/jax-releases/jax_nightly_releases.html -f https://storage.googleapis.com/jax-releases/jaxlib_nightly_releases.html # Build vLLM. COPY . /workspace/vllm ENV VLLM_TARGET_DEVICE="tpu" -RUN cd /workspace/vllm && python setup.py develop +RUN cd /workspace/vllm && python3 -m pip install -r requirements-tpu.txt +RUN cd /workspace/vllm && python3 setup.py develop CMD ["/bin/bash"]