add check pypi version

Signed-off-by: Ylang Tsou <[email protected]>

Author: Ylang Tsou

.buildkite/pipeline_pypi.yml

@@ -2,21 +2,10 @@ steps:
   # -----------------------------------------------------------------
   # TEST STEPS - Calling wrapper
   # -----------------------------------------------------------------
-   - label: "Wait for 20 mins"
-     if: build.env("NIGHTLY") == "1"
-     key: "wait_20_minutes"
-     depends_on: "record_verified_commit_hashes"
-     agents:
-      queue: cpu
-     commands:
-      - "echo 'Starting 20 minute delay...'"
-      - "sleep 1200"
-      - "echo 'Delay finished, starting benchmarks.'"
-
    - label: "Performance benchmarks for meta-llama/Llama-3.1-8B-Instruct"
      key: "meta-llama_Llama-3_1-8B-Instruct_Benchmark"
-     if: build.env("NIGHTLY") == "1"
-     depends_on: "wait_20_minutes"
+     #if: build.env("NIGHTLY") == "1"
+     depends_on: "record_verified_commit_hashes"
      agents:
       queue: tpu_v6e_queue
      env:
@@ -34,9 +23,9 @@ steps:
         .buildkite/scripts/run_with_pypi.sh bash /workspace/tpu_inference/tests/e2e/benchmarking/benchmark.sh
 
    - label: "Performance benchmarks for Qwen/Qwen3-4B"
-     if: build.env("NIGHTLY") == "1"
+     #if: build.env("NIGHTLY") == "1"
      key: "Qwen_Qwen3-4B_Benchmark"
-     depends_on: "wait_20_minutes"
+     depends_on: "record_verified_commit_hashes"
      agents:
       queue: tpu_v6e_queue
      env:

.buildkite/scripts/run_with_pypi.sh

@@ -2,6 +2,43 @@
 # Exit on error, exit on unset variable, fail on pipe errors.
 set -euo pipefail
 
+if ! git rev-parse --is-inside-work-tree >/dev/null 2>&1; then
+    echo "Error: Not a git repository. Cannot calculate version from tags."
+    exit 1
+fi
+# Get the TPU_INFERENCE_VERSION based on the latest stable tag and current date.
+LATEST_STABLE_TAG=$(git tag --sort=-v:refname | grep -E '^v[0-9]+\.[0-9]+\.[0-9]+$' | head -n 1)
+BASE_VERSION=${LATEST_STABLE_TAG#v}
+DATETIME_STR=$(date -d 'yesterday' +%Y%m%d)
+TPU_INFERENCE_VERSION="${BASE_VERSION}.dev${DATETIME_STR}"
+
+echo "Target Nightly Version: ${TPU_INFERENCE_VERSION}"
+
+PACKAGE_NAME="tpu-inference"
+MAX_RETRIES=20
+SLEEP_SEC=60
+FOUND_VERSION=false
+
+echo "Checking PyPI for ${PACKAGE_NAME} == ${TPU_INFERENCE_VERSION}..."
+
+for ((i=1; i<=MAX_RETRIES; i++)); do
+    if pip index versions "${PACKAGE_NAME}" --pre 2>/dev/null | grep -q "${TPU_INFERENCE_VERSION}"; then
+        echo "✅ Success! Found version ${TPU_INFERENCE_VERSION} on PyPI."
+        FOUND_VERSION=true
+        break
+    fi
+
+    echo "[Attempt $i/$MAX_RETRIES] Version not found yet. Waiting ${SLEEP_SEC} seconds..."
+    if [ "$i" -lt "$MAX_RETRIES" ]; then
+        sleep "$SLEEP_SEC"
+    fi
+done
+
+if [ "$FOUND_VERSION" = "false" ]; then
+    echo "The version ${TPU_INFERENCE_VERSION} was not found on PyPI."
+    exit 1
+fi
+
 # Build vllm-tpu with nightly tpu-inference from PyPI (using docker/Dockerfile.pypi instead of docker/Dockerfile).
 export RUN_WITH_PYPI="true"
 SCRIPT_DIR=$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" &>/dev/null && pwd)

docker/Dockerfile.pypi

@@ -5,6 +5,8 @@ ARG VLLM_COMMIT_HASH=""
 
 FROM $BASE_IMAGE
 
+ARG IS_FOR_V7X="false"
+
 # Remove existing versions of dependencies
 RUN pip uninstall -y torch torch_xla torchvision
 
@@ -17,7 +19,11 @@ RUN apt-get update && apt-get install -y \
 WORKDIR /workspace/tpu_inference
 COPY requirements_benchmarking.txt .
 # These are needed for the E2E benchmarking tests (i.e. tests/e2e/benchmarking/mlperf.sh)
-RUN pip install -r requirements_benchmarking.txt --retries 3
+RUN --mount=type=cache,target=/root/.cache/pip pip install -r requirements_benchmarking.txt --retries 3
+COPY requirements_v7x.txt .
+RUN --mount=type=cache,target=/root/.cache/pip if [ "$IS_FOR_V7X" = "true" ]; then \
+        pip install -r requirements_v7x.txt; \
+    fi
 COPY . .
 
 # Build vllm-tpu wheel