Autobump tokamax pin

.buildkite/integration_tokamax_promote.yml

@@ -0,0 +1,38 @@
+# Copyright 2026 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+steps:
+  - wait: ~
+  # -----------------------------------------------------------------
+  # Bump the tokamax pin on main. The wait above means this only runs once
+  # every uploaded step passed; the per-TPU "TPU Test Notification" steps
+  # turn the soft_fail test steps into a hard build failure, so a red test
+  # really does block this promote.
+  # -----------------------------------------------------------------
+  - label: "Promote validated tokamax nightly to main"
+    key: "promote_tokamax"
+    agents:
+      queue: "cpu"
+    commands:
+      - |
+        TOKAMAX_VERSION=$(buildkite-agent meta-data get "TOKAMAX_VERSION")
+        TOKAMAX_PREVIOUS_VERSION=$(buildkite-agent meta-data get "TOKAMAX_PREVIOUS_VERSION" --default "unknown")
+        bash .buildkite/scripts/update_tokamax_pin.sh "$$TOKAMAX_VERSION"
+        buildkite-agent annotate ":white_check_mark: tokamax bumped: $$TOKAMAX_PREVIOUS_VERSION -> $$TOKAMAX_VERSION" --style "success"
+
+notify:
+  - email: "ullm-test-notifications-external@google.com"
+    if: build.state == "failed" && build.source == "schedule"
+  - slack: "vllm#tpu-ci-notifications"
+    if: build.state == "failed" && build.source == "schedule"

.buildkite/scripts/bootstrap_tokamax.sh

@@ -0,0 +1,139 @@
+#!/bin/bash
+# Copyright 2026 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+# Bootstrap for the scheduled "tpu-tokamax-integration" pipeline.
+#
+# 1. Resolves the newest tokamax nightly on PyPI
+# 2. Runs the JAX test suites against it.
+# 3. Only if every test passes, it pushes the requirements.txt bump to main.
+
+set -euo pipefail
+
+SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
+# shellcheck source=/dev/null
+source "${SCRIPT_DIR}/configs/pipeline_config.sh"
+
+REQUIREMENTS_FILE="requirements.txt"
+JOB_PRIORITY="${PRIORITY_INTEGRATION}"
+export JOB_PRIORITY
+buildkite-agent meta-data set "JOB_PRIORITY" "${JOB_PRIORITY}"
+
+# Handles the environment state for different TPU generations.
+set_jax_envs() {
+    case $1 in
+        v6)
+            export TESTS_GROUP_LABEL="[jax] TPU6e Tests Group"
+            export TPU_VERSION="tpu6e"
+            export TPU_QUEUE_SINGLE="tpu_v6e_queue"
+            export TPU_QUEUE_MULTI="tpu_v6e_8_queue"
+            export TENSOR_PARALLEL_SIZE_SINGLE=1
+            export TENSOR_PARALLEL_SIZE_MULTI=8
+            ;;
+        v7)
+            export TESTS_GROUP_LABEL="[jax] TPU7x Tests Group"
+            export TPU_VERSION="tpu7x"
+            export TPU_QUEUE_SINGLE="tpu_v7x_2_queue"
+            export TPU_QUEUE_MULTI="tpu_v7x_8_queue"
+            export TENSOR_PARALLEL_SIZE_SINGLE=2
+            export TENSOR_PARALLEL_SIZE_MULTI=8
+            export COV_FAIL_UNDER="67"
+            ;;
+        unset)
+            unset TESTS_GROUP_LABEL TPU_VERSION TPU_QUEUE_SINGLE TPU_QUEUE_MULTI \
+                  TENSOR_PARALLEL_SIZE_SINGLE TENSOR_PARALLEL_SIZE_MULTI COV_FAIL_UNDER
+            ;;
+    esac
+}
+
+CURRENT_VERSION="$(sed -n 's/^tokamax==\(.*\)$/\1/p' "${REQUIREMENTS_FILE}")"
+if [[ -z "${CURRENT_VERSION}" ]]; then
+    echo "ERROR: no 'tokamax==' pin found in ${REQUIREMENTS_FILE}." >&2
+    exit 1
+fi
+
+# TOKAMAX_VERSION can be set in the build env from the Buildkite UI.
+NEW_VERSION="${TOKAMAX_VERSION:-}"
+if [[ -z "${NEW_VERSION}" ]]; then
+    echo "--- :package: Resolving the newest tokamax nightly from PyPI"
+    NEW_VERSION="$(python3 - <<'PYEOF'
+import json
+import re
+import sys
+import urllib.request
+
+with urllib.request.urlopen("https://pypi.org/pypi/tokamax/json", timeout=60) as resp:
+    data = json.load(resp)
+
+# Nightlies are published as <base>.devYYYYMMDD, e.g. 0.0.14.dev20260903.
+# info.version only tracks the newest *stable* release, so the nightly has to be
+# picked out of the full release list. Releases whose files were all removed or
+# yanked are skipped - pip cannot install those.
+pattern = re.compile(r"^(\d+(?:\.\d+)*)\.dev(\d{8})$")
+candidates = []
+for version, files in data["releases"].items():
+    match = pattern.match(version)
+    if not match:
+        continue
+    if not any(not f.get("yanked", False) for f in files):
+        continue
+    base = tuple(int(part) for part in match.group(1).split("."))
+    candidates.append(((base, int(match.group(2))), version))
+
+if not candidates:
+    print("No installable tokamax nightly (*.devYYYYMMDD) found on PyPI.", file=sys.stderr)
+    sys.exit(1)
+
+print(max(candidates)[1])
+PYEOF
+)"
+fi
+
+echo "Pinned tokamax version   : ${CURRENT_VERSION}"
+echo "Candidate tokamax version: ${NEW_VERSION}"
+
+if [[ "${CURRENT_VERSION}" == "${NEW_VERSION}" ]]; then
+    echo "Already on ${NEW_VERSION}. Nothing to bump; skipping the test run."
+    buildkite-agent annotate \
+        ":white_check_mark: tokamax already pinned to \`${NEW_VERSION}\` - no bump needed." \
+        --style "success"
+    exit 0
+fi
+
+buildkite-agent meta-data set "TOKAMAX_VERSION" "${NEW_VERSION}"
+buildkite-agent meta-data set "TOKAMAX_PREVIOUS_VERSION" "${CURRENT_VERSION}"
+# buildkite-agent annotate posts a markdown banner at the top of the build page
+buildkite-agent annotate \
+    ":arrow_up: Validating tokamax bump \`${CURRENT_VERSION}\` :arrow_right: \`${NEW_VERSION}\`. main is bumped only if every step below passes." \
+    --style "info"
+
+VLLM_COMMIT_HASH="$(get_vllm_commit_hash)"
+buildkite-agent meta-data set "VLLM_COMMIT_HASH" "${VLLM_COMMIT_HASH}"
+echo "Using vllm LKG commit hash: ${VLLM_COMMIT_HASH}"
+
+# Buildkite inserts uploaded steps in reverse order, so the promote has to be
+# uploaded first for it to end up last.
+upload_with_priority .buildkite/integration_tokamax_promote.yml "${JOB_PRIORITY}"
+
+set_jax_envs v7
+upload_with_priority .buildkite/pipeline_jax.yml "${JOB_PRIORITY}"
+set_jax_envs unset
+
+set_jax_envs v6
+upload_with_priority .buildkite/pipeline_jax.yml "${JOB_PRIORITY}"
+set_jax_envs unset
+
+upload_with_priority .buildkite/pipeline_build.yml "${JOB_PRIORITY}"
+
+echo "--- Tokamax Integration Bootstrap Finished"

.buildkite/scripts/setup_docker_env.sh

@@ -123,6 +123,31 @@ verify_image_vllm() {
   echo "[verify-vllm] OK: ${image_ref} contains expected vLLM ${expected_vllm}."
 }
 
+verify_image_tokamax() {
+  local image_ref="$1"
+  local expected_tokamax="${2:-}"
+  if [[ -z "${expected_tokamax}" ]]; then
+    return 0
+  fi
+  local actual_tokamax rc=0 err
+  err=$(mktemp)
+  actual_tokamax=$(docker run --rm --entrypoint python3 "${image_ref}" \
+    -c 'import importlib.metadata as m; print(m.version("tokamax"))' 2>"${err}") || rc=$?
+  if [[ ${rc} -ne 0 ]]; then
+    echo "[FATAL][verify-tokamax] Could not read the tokamax version from ${image_ref} (exit ${rc}):" >&2
+    cat "${err}" >&2
+    rm -f "${err}"
+    exit 1
+  fi
+  rm -f "${err}"
+  if [[ "${actual_tokamax}" != "${expected_tokamax}" ]]; then
+    echo "[FATAL][verify-tokamax] ${image_ref} contains tokamax ${actual_tokamax}," >&2
+    echo "[FATAL][verify-tokamax] but this build is validating ${expected_tokamax}." >&2
+    exit 1
+  fi
+  echo "[verify-tokamax] ${image_ref} contains tokamax ${actual_tokamax}."
+}
+
 setup_environment() {
   local image_name_param=${1:-"vllm-tpu"}
   local should_push=${2:-"false"}
@@ -179,12 +204,22 @@ setup_environment() {
       VLLM_COMMIT_HASH=$(buildkite-agent meta-data get "VLLM_COMMIT_HASH" --default "")
       TPU_INFERENCE_HASH="$BUILDKITE_COMMIT"
   fi
+
+  # Non-empty only in the tokamax integration pipeline, which validates a
+  # candidate nightly that is not yet pinned in requirements.txt. Every other
+  # pipeline leaves this empty and is unaffected by the two blocks below.
+  local TOKAMAX_VERSION=""
+  if [ -n "${BUILDKITE:-}" ]; then
+    TOKAMAX_VERSION=$(buildkite-agent meta-data get "TOKAMAX_VERSION" --default "")
+  fi
 
   # Include the vLLM commit in the cache tag so an image is uniquely identified
   # by BOTH its tpu-inference commit and its vLLM commit. Without this, the CI
   # pipeline (LKG vLLM) and the integration pipeline (HEAD vLLM) produce the same
   # tag for the same tpu-inference commit and can overwrite each other's image in
   # the registry -- a test could then silently run a different vLLM than intended.
+  # Similarly, we bake the tokamax commit in the cache tag below.
+  # The tag has 128 char limit.
   local CACHE_TAG="${TPU_INFERENCE_HASH}-${LOCAL_TPU_VERSION}"
   if [[ -n "${VLLM_COMMIT_HASH}" ]]; then
     CACHE_TAG="${TPU_INFERENCE_HASH}-${VLLM_COMMIT_HASH}-${LOCAL_TPU_VERSION}"
@@ -200,6 +235,10 @@ setup_environment() {
     exit 1
   fi
 
+  if [[ -n "${TOKAMAX_VERSION}" ]]; then
+    CACHE_TAG="${CACHE_TAG}-tkmx${TOKAMAX_VERSION//[^a-zA-Z0-9]/}"
+  fi
+
   # ==========================================
   # Pull-Only Mode for TPU execution nodes
   # ==========================================
@@ -208,13 +247,20 @@ setup_environment() {
     # -q: the layer-by-layer pull progress is several hundred lines per job.
     docker pull -q "${CI_IMAGE_REPO}:${CACHE_TAG}"
     verify_image_vllm "${CI_IMAGE_REPO}:${CACHE_TAG}" "${VLLM_COMMIT_HASH}"
+    verify_image_tokamax "${CI_IMAGE_REPO}:${CACHE_TAG}" "${TOKAMAX_VERSION}"
     docker tag "${CI_IMAGE_REPO}:${CACHE_TAG}" "${IMAGE_NAME}:${TPU_INFERENCE_HASH}"
     docker tag "${CI_IMAGE_REPO}:${CACHE_TAG}" "${IMAGE_NAME}:latest"
     # Export the computed CI cache image name so calling scripts can use it.
     export EXPORTED_CI_CACHE_IMAGE="${CI_IMAGE_REPO}:${CACHE_TAG}"
     return 0
   fi
 
+  if [[ -n "${TOKAMAX_VERSION}" ]]; then
+    echo "[tokamax] Pinning the build context to tokamax==${TOKAMAX_VERSION}"
+    sed -i "s/^tokamax==.*$/tokamax==${TOKAMAX_VERSION}/" requirements.txt
+    grep -n '^tokamax==' requirements.txt
+  fi
+
   # Build with specific hash and 'latest' tag for convenience
   docker build \
       --build-arg VLLM_COMMIT_HASH="${VLLM_COMMIT_HASH}" \
@@ -228,6 +274,7 @@ setup_environment() {
   # Fail fast if the freshly built image does not contain the expected vLLM
   # commit (guards against a mis-set VLLM_COMMIT_HASH build-arg).
   verify_image_vllm "${IMAGE_NAME}:${CACHE_TAG}" "${VLLM_COMMIT_HASH}"
+  verify_image_tokamax "${IMAGE_NAME}:${CACHE_TAG}" "${TOKAMAX_VERSION}"
 
   # ==========================================
   # Push to CI Image Registry (Executed by dedicate CPU builder)

.buildkite/scripts/update_tokamax_pin.sh

@@ -0,0 +1,61 @@
+#!/bin/bash
+# Copyright 2026 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+# Pushes the validated tokamax pin to main. Called from the promote step of the
+# tokamax integration pipeline, which only runs after every test step passed.
+# Mirrors update_lkg_version.sh.
+
+set -e
+
+NEW_VERSION=$1
+if [[ -z "$NEW_VERSION" ]]; then
+    echo "Error: No tokamax version provided."
+    exit 1
+fi
+
+# Configuration. TARGET_BRANCH is overridable so the flow can be rehearsed
+# against a scratch branch without writing to main.
+TARGET_BRANCH="${TARGET_BRANCH:-main}"
+REQUIREMENTS_FILE="requirements.txt"
+
+# Configure credentials
+git config user.name "vllm-ci-bot[bot]"
+git config user.email "vllm-ci-bot[bot]@users.noreply.github.com"
+
+# Fetch and checkout
+git fetch origin "${TARGET_BRANCH}"
+git checkout -f "${TARGET_BRANCH}"
+git reset --hard origin/"${TARGET_BRANCH}"
+
+# Update the tokamax pin
+echo "Updating $REQUIREMENTS_FILE to: tokamax==$NEW_VERSION"
+sed -i "s/^tokamax==.*$/tokamax==${NEW_VERSION}/" "$REQUIREMENTS_FILE"
+
+# Check in file
+git add "$REQUIREMENTS_FILE"
+
+# Check if we have changed anything
+if git diff --cached --quiet; then
+    echo "No changes in the tokamax pin. Skipping push."
+else
+    # "[skip ci]" matches update_lkg_version.sh: the v6e+v7x suite that just
+    # passed is the validation for this bump, so a post-merge re-run would
+    # largely duplicate it. Note the reset above lands the bump on whatever
+    # main is now, which may have moved past the commit this build tested.
+    git commit -s -m "[skip ci] Update tokamax pin to $NEW_VERSION"
+    echo "Pushing the tokamax bump to $TARGET_BRANCH..."
+    git push origin "$TARGET_BRANCH"
+    echo "Successfully bumped tokamax to $NEW_VERSION"
+fi