Skip to content

Commit

Permalink
[Bugfix][TPU] Fix outlines installation in TPU Dockerfile (#6256)
Browse files Browse the repository at this point in the history
  • Loading branch information
WoosukKwon authored Jul 9, 2024
1 parent 5d5b4c5 commit 08c5bde
Showing 1 changed file with 7 additions and 3 deletions.
10 changes: 7 additions & 3 deletions Dockerfile.tpu
Original file line number Diff line number Diff line change
Expand Up @@ -2,18 +2,22 @@ ARG NIGHTLY_DATE="20240601"
ARG BASE_IMAGE="us-central1-docker.pkg.dev/tpu-pytorch-releases/docker/xla:nightly_3.10_tpuvm_$NIGHTLY_DATE"

FROM $BASE_IMAGE

WORKDIR /workspace
COPY . /workspace/vllm

ENV VLLM_TARGET_DEVICE="tpu"
# Install aiohttp separately to avoid build errors.
RUN pip install aiohttp
# 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

# Build vLLM.
COPY . /workspace/vllm
ENV VLLM_TARGET_DEVICE="tpu"
RUN cd /workspace/vllm && python setup.py develop

# Re-install outlines to avoid dependency errors.
# The outlines version must follow requirements-common.txt.
RUN pip uninstall outlines -y
RUN pip install "outlines>=0.0.43"

CMD ["/bin/bash"]

0 comments on commit 08c5bde

Please sign in to comment.