Include vLLM and tpu inference to Tunix build.

Dockerfile

@@ -7,7 +7,7 @@ ENV TZ=Etc/UTC
 
 # Install system dependencies, including Python 3 and pip
 RUN apt-get update && \
-    apt-get install -y git python3 python3-pip && \
+    apt-get install -y build-essential git python3 python3-pip && \
     rm -rf /var/lib/apt/lists/*
 
 # Upgrade pip
@@ -35,8 +35,13 @@ WORKDIR /app
 # Copy the project files to the image
 COPY . .
 
+ENV VLLM_TARGET_DEVICE=tpu
+
+RUN pip install -r /app/requirements/requirements.txt
+RUN pip install --force-reinstall -r /app/requirements/special_requirements.txt
+
 # Install the project in editable mode
-RUN pip install  --force-reinstall .
+RUN pip install --force-reinstall .
 
 # Set the default command to bash
-CMD ["bash"]
+CMD ["bash"]

README.md

@@ -87,6 +87,19 @@ tuning workflows on top of the XLA and JAX infrastructure. See [Design Overview]
 **Installation:** Jump to [Installation](https://tunix.readthedocs.io/en/latest/quickstart.html#installation) to install Tunix and run your first training
 job.
 
+For TPU users integrating `vllm` and `tpu-inference`, there are two supported
+setup paths:
+
+- Docker image builds use [Dockerfile](/usr/local/google/home/lancewang/github/tunix/Dockerfile) and install
+    the pinned dependencies directly from `requirements/requirements.txt` and
+    `requirements/special_requirements.txt`.
+- Local TPU VM or developer-machine installs can use
+    [scripts/install_tunix_vllm_requirement.sh](/usr/local/google/home/lancewang/github/tunix/scripts/install_tunix_vllm_requirement.sh),
+    which installs the same requirement files outside Docker.
+
+These are separate entry points. If you are building the Docker image, you do
+not need to run the install script inside the container build.
+
 **Examples:** To get started, we have a number of detailed examples and tutorials. You can see [Quick Start](https://tunix.readthedocs.io/en/latest/quickstart.html) for a great set of starting examples and [Examples and Guides](https://tunix.readthedocs.io/en/latest/examples.html) for a comprehensive list of all the notebooks and examples we have.
 
 

build_docker.sh

@@ -31,8 +31,24 @@ build_ai_image() {
     echo "Building Tunix Image at commit hash ${COMMIT_HASH}..."
 
     DOCKER_COMMAND="docker"
-    if ! docker info >/dev/null 2>&1; then
-        DOCKER_COMMAND="sudo docker"
+    if docker info >/dev/null 2>&1; then
+        DOCKER_COMMAND="docker"
+    else
+        # Avoid invoking sudo interactively which can prompt for a password.
+        # Check whether non-interactive sudo would work (no password).
+        if sudo -n docker info >/dev/null 2>&1; then
+            DOCKER_COMMAND="sudo docker"
+        else
+            cat <<'MSG'
+Docker does not appear usable from this account and the build would prompt for a password.
+
+Run the build with sufficient privileges (will prompt): sudo bash build_docker.sh
+On Linux, add your user to the docker group so sudo isn't required (you must re-login):
+  sudo usermod -aG docker "$USER" && newgrp docker
+
+MSG
+            exit 1
+        fi
     fi
 
     $DOCKER_COMMAND build \
@@ -48,4 +64,4 @@ echo "*************************
 "
 
 echo "Built your docker image and named it ${LOCAL_IMAGE_NAME}.
-It only has the dependencies installed. "
+It now installs Tunix and the pinned vLLM and tpu-inference dependencies from requirements/requirements.txt. "

docs/quickstart.md

@@ -83,6 +83,7 @@ pip install "git+https://github.com/google/tunix#egg=google-tunix[prod]"
 # For GPU/CPU
 pip install git+https://github.com/google/tunix
 # Then install JAX for GPU or CPU as described above.
+
 ```
 
 ### Option 3: From Source
@@ -110,13 +111,39 @@ These need to be installed manually.
 
 **vLLM on TPU**
 
-The TPU-inference supported version of vLLM is not on PyPI yet. Please follow
-the instructions to install manually from
-[vLLM TPU docs](https://docs.vllm.ai/projects/tpu/en/latest/getting_started/installation/)
-or use the Docker image `vllm/vllm-tpu:v0.11.1` and then run:
+The TPU-inference supported version of `vllm` is not always available as a
+single PyPI release, and installing the TPU build sometimes requires extra pip flags
+so that `libtpu` wheels (hosted by the JAX project) can be resolved. You can
+install the pinned vLLM + TPU requirements from this repository using one of
+the raw requirement-file URLs below.
+
+Install from remote:
+
+```sh
+pip install -r https://github.com/google/tunix/raw/main/requirements/requirements.txt
+pip install -r https://github.com/google/tunix/raw/main/requirements/special_requirements.txt
+```
+
+Or (direct raw.githubusercontent URL):
+
+```sh
+pip install -r https://raw.githubusercontent.com/google/tunix/main/requirements/requirements.txt
+pip install -r https://raw.githubusercontent.com/google/tunix/main/requirements/special_requirements.txt
+```
+
+If you prefer a single-line install that directly overrides `tpu-inference`, you can also run:
+
+```sh
+pip install vllm @git+https://github.com/vllm-project/vllm.git@<commit>
+pip install --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 \\
+            --pre \\
+            tpu-inference@git+https://github.com/vllm-project/tpu-inference.git@<commit>
+```
 
+Or install from source:
 ```sh
-pip install tpu-inference
+bash scripts/install_tunix_vllm_requirements.sh
 ```
 
 **SGLang-Jax**

requirements/requirements.txt

@@ -0,0 +1 @@
+vllm @ git+https://github.com/vllm-project/vllm.git@f6983f01de2bf2e92ab468fa735ebac39cddd670

requirements/special_requirements.txt

@@ -0,0 +1,5 @@
+--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
+--pre
+
+tpu-inference @ git+https://github.com/vllm-project/tpu-inference.git@581c4d49959c30638a3bdcecb981fe83ab85830e

scripts/install_tunix_vllm_requirement.sh

@@ -1,4 +1,4 @@
-# !/bin/bash
+#!/bin/bash
 
 # Copyright 2023–2025 Google LLC
 #
@@ -16,29 +16,20 @@
 
 # This script installs the dependencies for running GRPO with MaxText+Tunix+vLLM on TPUs
 
-set -e
+set -euo pipefail
 set -x
 
-python -m ensurepip --default-pip
+ROOT_DIR=$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)
+REQ_FILE=${REQ_FILE:-"${ROOT_DIR}/requirements/requirements.txt"}
+SPECIAL_REQ_FILE=${SPECIAL_REQ_FILE:-"${ROOT_DIR}/requirements/special_requirements.txt"}
 
-pip uninstall -y jax jaxlib libtpu
+python3 -m ensurepip --default-pip
+python3 -m pip install --upgrade pip setuptools wheel
 
 pip install aiohttp==3.12.15
 
 # Install Python packages that enable pip to authenticate with Google Artifact Registry automatically.
 pip install keyring keyrings.google-artifactregistry-auth
 
-# Install vLLM for Jax and TPUs from GitHub repository
-git clone https://github.com/vllm-project/vllm.git /tmp/vllm
-cd /tmp/vllm
-VLLM_TARGET_DEVICE="tpu" pip install -e .
-
-# Install tpu-commons from the artifact registry
-pip install --no-cache-dir --pre \
-    --index-url https://us-python.pkg.dev/cloud-tpu-images/maxtext-rl/simple/ \
-    --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 \
-    tpu-commons==0.1.1
-
-pip install numba==0.61.2
+VLLM_TARGET_DEVICE="tpu" python3 -m pip install -r "${REQ_FILE}"
+python3 -m pip install -r "${SPECIAL_REQ_FILE}" --force-reinstall