Update tpu-inference and vllm pinned commits.

Dockerfile

@@ -7,15 +7,15 @@ 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
 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"]

examples/deepswe/train_deepswe_nb.py

@@ -535,19 +535,21 @@
 # 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:
   rollout_dims = []
   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

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

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