From 0537accea501919ec09f70aacb466e19d97451b2 Mon Sep 17 00:00:00 2001 From: Nitin Gangahar Date: Tue, 21 Jul 2026 16:55:34 -0700 Subject: [PATCH] Update tpu-inference and vllm pinned commits. PiperOrigin-RevId: 951775202 --- Dockerfile | 34 +++++++++++++++++++++++---- examples/deepswe/train_deepswe_nb.py | 10 ++++---- requirements/requirements.txt | 2 +- requirements/special_requirements.txt | 2 +- 4 files changed, 38 insertions(+), 10 deletions(-) diff --git a/Dockerfile b/Dockerfile index ed2b09519..94954d3cc 100644 --- a/Dockerfile +++ b/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 build-essential git python3 python3-pip && \ + apt-get install -y build-essential git curl python3 python3-pip && \ rm -rf /var/lib/apt/lists/* # Upgrade pip @@ -15,7 +15,7 @@ RUN python3 -m pip install --upgrade pip # Create a virtual environment RUN python3.12 -m venv /opt/venv -ENV PATH="/opt/venv/bin:$PATH" +ENV PATH="/opt/venv/bin:/root/.local/bin:/root/.cargo/bin:$PATH" # Upgrade pip RUN pip install --upgrade pip @@ -25,17 +25,43 @@ RUN pip install git+https://github.com/ayaka14732/jax-smi.git # RUN pip install git+https://github.com/AI-Hypercomputer/pathways-utils.git@b72729bb152b7b3426299405950b3af300d765a9#egg=pathwaysutils RUN pip install gcsfs RUN pip install wandb +RUN pip install kubernetes + +# Copy the project files to a temp directory +COPY . /app/g3_sources + +# Restructure to match OSS layout +RUN mkdir -p /app/tunix && \ + cp -a /app/g3_sources/* /app/tunix/ && \ + mv /app/tunix/oss/pyproject.toml /app/pyproject.toml && \ + mv /app/tunix/oss/requirements /app/requirements && \ + mv /app/tunix/oss/scripts /app/scripts && \ + mv /app/tunix/oss/examples /app/examples && \ + mv /app/tunix/tests /app/tests && \ + rm -rf /app/tunix/pyproject.toml /app/tunix/requirements /app/tunix/scripts /app/tunix/examples && \ + rm -rf /app/g3_sources # Set the working directory WORKDIR /app -# Copy the project files to the image -COPY . . +# Apply OSS replacements to avoid internal imports on GCP +RUN python3 /app/scripts/apply_oss_replacements.py /app/tunix # Install the project in editable mode RUN pip install -e . RUN bash /app/scripts/install_tunix_vllm_requirement.sh +RUN pip install gym swebench==3.0.2 +RUN pip install --no-deps git+https://github.com/r2e-gym/r2e-gym.git@0d94c4eb9431cd195c55a7ea3abd54006c9a1735 +RUN pip install --upgrade flax +RUN pip install aqtp tokamax +RUN pip install git+https://github.com/google/maxtext.git + +# Patch r2egym bugs +RUN sed -i 's/create_repo, upload_folder, HfFolder/create_repo, upload_folder/' /opt/venv/lib/python3.12/site-packages/r2egym/agenthub/utils/utils.py +RUN sed -i 's/self.commit = ParsedCommit(\*\*json.loads(self.commit_json))/self.commit = ParsedCommit(\*\*(json.loads(self.commit_json) if isinstance(self.commit_json, str) else self.commit_json))/' /opt/venv/lib/python3.12/site-packages/r2egym/agenthub/runtime/docker.py +RUN sed -i 's/"karpenter.sh\/nodepool": "bigcpu-standby"/"cloud.google.com\/gke-nodepool": "cpu-np"/' /opt/venv/lib/python3.12/site-packages/r2egym/agenthub/runtime/docker.py + # Set the default command to bash CMD ["bash"] diff --git a/examples/deepswe/train_deepswe_nb.py b/examples/deepswe/train_deepswe_nb.py index 32ff61bc7..8f5a12710 100644 --- a/examples/deepswe/train_deepswe_nb.py +++ b/examples/deepswe/train_deepswe_nb.py @@ -535,6 +535,8 @@ # Each explicitly-provided dim becomes an axis in the mesh; unspecified dims are # dropped (not defaulted to 1), so passing only --rollout_mesh_fsdp yields a 1D mesh. # If nothing is provided, fall back to the split-fraction heuristic (2D: fsdp+tp). +tp_axis_name = "tensor" if MODEL_SOURCE == "maxtext" else "tp" + rollout_fsdp = args.rollout_mesh_fsdp rollout_tp = args.rollout_mesh_tp if rollout_fsdp is not None or rollout_tp is not None: @@ -542,12 +544,12 @@ if rollout_fsdp is not None: rollout_dims.append(("fsdp", rollout_fsdp)) if rollout_tp is not None: - rollout_dims.append(("tp", rollout_tp)) + rollout_dims.append((tp_axis_name, rollout_tp)) else: num_rollout_devices = int(total_devices * args.rollout_split_fraction) rollout_tp = int(np.gcd(num_rollout_devices, config.num_kv_heads)) rollout_fsdp = num_rollout_devices // rollout_tp - rollout_dims = [("fsdp", rollout_fsdp), ("tp", rollout_tp)] + rollout_dims = [("fsdp", rollout_fsdp), (tp_axis_name, rollout_tp)] num_rollout_devices = int(np.prod([d for _, d in rollout_dims])) # 2. Resolve Train Mesh Dimensions @@ -562,14 +564,14 @@ train_dims.append(("fsdp", train_fsdp if train_fsdp is not None else 1)) if train_sp is not None: train_dims.append(("sp", train_sp)) - train_dims.append(("tp", train_tp if train_tp is not None else 1)) + train_dims.append((tp_axis_name, train_tp if train_tp is not None else 1)) else: num_train_devices = total_devices - num_rollout_devices train_fsdp = int( np.gcd(num_train_devices, TRAIN_MICRO_BATCH_SIZE * NUM_GENERATIONS) ) train_tp = num_train_devices // train_fsdp - train_dims = [("fsdp", train_fsdp), ("tp", train_tp)] + train_dims = [("fsdp", train_fsdp), (tp_axis_name, train_tp)] num_train_devices = int(np.prod([d for _, d in train_dims])) # 3. Sanity Check diff --git a/requirements/requirements.txt b/requirements/requirements.txt index dfa7b948a..9adff6081 100644 --- a/requirements/requirements.txt +++ b/requirements/requirements.txt @@ -1 +1 @@ -vllm @ git+https://github.com/vllm-project/vllm.git@5e584ce9ecb3cce63f1caab86177aef5c831690f +vllm @ git+https://github.com/vllm-project/vllm.git@85f638a2b8f9838275b911c44fb0c2abdf604476 diff --git a/requirements/special_requirements.txt b/requirements/special_requirements.txt index c5f144e91..0884ebebd 100644 --- a/requirements/special_requirements.txt +++ b/requirements/special_requirements.txt @@ -2,4 +2,4 @@ # --find-links https://storage.googleapis.com/jax-releases/libtpu_releases.html # --pre -tpu-inference @ git+https://github.com/vllm-project/tpu-inference.git@8dd3107ac1d50b823c2a4edd036dd181775b1c94 +tpu-inference @ git+https://github.com/vllm-project/tpu-inference.git@5caa2dcdab134771758beae1a6f28de9099efc9f