From c5f1dc4f44d2c1633886cac9974cd377cb4e9d13 Mon Sep 17 00:00:00 2001
From: khluu <khluu000@gmail.com>
Date: Sun, 17 May 2026 03:07:03 -0700
Subject: [PATCH] Add PyPI publish step to release pipeline

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
---
 .buildkite/release-pipeline.yaml              | 14 +++++
 .../scripts/publish-release-wheels-pypi.sh    | 52 +++++++++++++++++++
 2 files changed, 66 insertions(+)
 create mode 100755 .buildkite/scripts/publish-release-wheels-pypi.sh

diff --git a/.buildkite/release-pipeline.yaml b/.buildkite/release-pipeline.yaml
index 5f2dbdf0d..c560fefe1 100644
--- a/.buildkite/release-pipeline.yaml
+++ b/.buildkite/release-pipeline.yaml
@@ -26,6 +26,20 @@ steps:
     commands:
       - "bash .buildkite/scripts/generate-and-upload-nightly-index.sh"
 
+  - block: "Upload release wheels to PyPI"
+    depends_on: "build-wheel"
+    key: block-upload-release-wheels
+
+  - label: "Upload release wheels to PyPI"
+    depends_on:
+      - block-upload-release-wheels
+      - input-release-version
+    key: upload-release-wheels
+    agents:
+      queue: small_cpu_queue_release
+    commands:
+      - "bash .buildkite/scripts/publish-release-wheels-pypi.sh"
+
   - block: "Unblock to build release Docker images"
     depends_on: ~
     key: block-build-release-images
diff --git a/.buildkite/scripts/publish-release-wheels-pypi.sh b/.buildkite/scripts/publish-release-wheels-pypi.sh
new file mode 100755
index 000000000..523d73069
--- /dev/null
+++ b/.buildkite/scripts/publish-release-wheels-pypi.sh
@@ -0,0 +1,52 @@
+#!/usr/bin/env bash
+# SPDX-License-Identifier: Apache-2.0
+
+set -e
+
+BUCKET="vllm-wheels"
+SUBPATH="omni/$BUILDKITE_COMMIT"
+S3_COMMIT_PREFIX="s3://$BUCKET/$SUBPATH/"
+
+RELEASE_VERSION=$(buildkite-agent meta-data get release-version | sed 's/^v//')
+if [[ -z "$RELEASE_VERSION" ]]; then
+  echo "[FATAL] release-version metadata not set."
+  exit 1
+fi
+echo "Release version: $RELEASE_VERSION"
+
+if [[ -z "$PYPI_TOKEN" ]]; then
+  echo "[FATAL] PYPI_TOKEN is not set."
+  exit 1
+else
+  export TWINE_USERNAME="__token__"
+  export TWINE_PASSWORD="$PYPI_TOKEN"
+fi
+
+set -x
+
+if ! command -v uv &> /dev/null; then
+  curl -LsSf https://astral.sh/uv/install.sh | sh
+  export PATH="$HOME/.local/bin:$PATH"
+fi
+
+uv venv --python 3.12 /tmp/vllm-omni-release-env
+source /tmp/vllm-omni-release-env/bin/activate
+uv pip install twine
+
+DIST_DIR=/tmp/vllm-omni-release-dist
+mkdir -p "$DIST_DIR"
+
+echo "Downloading wheels from S3:"
+aws s3 ls "$S3_COMMIT_PREFIX"
+aws s3 cp --recursive --exclude "*" --include "vllm_omni-${RELEASE_VERSION}*.whl" --exclude "*dev*" --exclude "*rc[0-9]*" "$S3_COMMIT_PREFIX" "$DIST_DIR"
+ls -la "$DIST_DIR"
+
+PYPI_WHEEL_FILES=$(find "$DIST_DIR" -name "vllm_omni-${RELEASE_VERSION}*.whl" -not -name "*+*")
+if [[ -z "$PYPI_WHEEL_FILES" ]]; then
+  echo "[FATAL] No wheels found for version ${RELEASE_VERSION}"
+  exit 1
+fi
+
+python3 -m twine check $PYPI_WHEEL_FILES
+python3 -m twine upload --non-interactive --verbose $PYPI_WHEEL_FILES
+echo "Wheels uploaded to PyPI"