[TPU] Add docker support

[Michaelvll]

The following dockerfile (from vLLM) can be used on TPU

ARG NIGHTLY_DATE="20240726"
ARG BASE_IMAGE="us-central1-docker.pkg.dev/tpu-pytorch-releases/docker/xla:nightly_3.10_tpuvm_$NIGHTLY_DATE"

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"

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

CMD ["/bin/bash"]

Tested (run the relevant ones):

• Code formatting: bash format.sh
• Any manual or new tests for this PR (please specify below)
  • sky launch -c test-tpu-docker --gpus tpu-v3-8 --image-id docker:michaelvll/tpu-test
• All smoke tests: pytest tests/test_smoke.py
• Relevant individual smoke tests: pytest tests/test_smoke.py::test_fill_in_the_name
• Backward compatibility tests: conda deactivate; bash -i tests/backward_compatibility_tests.sh

[Michaelvll]

Moved to #4115. Closing now.