| 1 | #!/usr/bin/env bash |
| 2 | # SPDX-License-Identifier: Apache-2.0 |
| 3 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 4 | |
| 5 | set -euo pipefail |
| 6 | |
| 7 | readonly BUCKET="vllm-wheels" |
| 8 | readonly PREFIX="xpu" |
| 9 | readonly WHEEL_URL="https://github.com/intel/intel-xpu-backend-for-triton/releases/download/v3.8.0/triton-3.8.0+xpu-py3-none-any.whl" |
| 10 | readonly WHEEL_SHA256="f16b24efcc789e3987ee602b7442e11cac0b9d3bc78f96e3a2ec3cdd5a60b7bb" |
| 11 | readonly WHEEL_FILENAME="${WHEEL_URL##*/}" |
| 12 | readonly ENCODED_WHEEL_FILENAME="${WHEEL_FILENAME/+/%2B}" |
| 13 | readonly S3_PREFIX="s3://${BUCKET}/${PREFIX}/" |
| 14 | readonly COMMIT="${BUILDKITE_COMMIT:-}" |
| 15 | readonly DRY_RUN="${DRY_RUN:-0}" |
| 16 | |
| 17 | if [[ "$DRY_RUN" != "0" && "$DRY_RUN" != "1" ]]; then |
| 18 | echo "DRY_RUN must be 0 or 1" >&2 |
| 19 | exit 2 |
| 20 | fi |
| 21 | if [[ "$DRY_RUN" == "0" && ! "$COMMIT" =~ ^[0-9a-f]{40}$ ]]; then |
| 22 | echo "BUILDKITE_COMMIT must be a full lowercase commit hash" >&2 |
| 23 | exit 2 |
| 24 | fi |
| 25 | |
| 26 | cd "$(dirname "${BASH_SOURCE[0]}")/../../.." |
| 27 | |
| 28 | # Keep temporary files under the repository so the Docker Python fallback, |
| 29 | # which mounts the repository at /app, can access them. |
| 30 | work_dir=$(mktemp -d "$PWD/.xpu-triton-index.XXXXXX") |
| 31 | work_dir=${work_dir#"$PWD"/} |
| 32 | trap 'rm -rf "$work_dir"' EXIT |
| 33 | |
| 34 | wheel_path="$work_dir/$WHEEL_FILENAME" |
| 35 | objects_path="$work_dir/objects.json" |
| 36 | index_output_dir="$work_dir/$PREFIX" |
| 37 | index_generator="$work_dir/generate-nightly-index.py" |
| 38 | |
| 39 | curl --fail --location --retry 3 --output "$wheel_path" "$WHEEL_URL" |
| 40 | |
| 41 | if command -v sha256sum >/dev/null 2>&1; then |
| 42 | actual_sha=$(sha256sum "$wheel_path") |
| 43 | else |
| 44 | actual_sha=$(shasum -a 256 "$wheel_path") |
| 45 | fi |
| 46 | actual_sha=${actual_sha%% *} |
| 47 | if [[ "$actual_sha" != "$WHEEL_SHA256" ]]; then |
| 48 | echo "SHA256 mismatch for $WHEEL_FILENAME" >&2 |
| 49 | echo "Expected: $WHEEL_SHA256" >&2 |
| 50 | echo "Actual: $actual_sha" >&2 |
| 51 | exit 1 |
| 52 | fi |
| 53 | |
| 54 | if [[ "$DRY_RUN" == "1" ]]; then |
| 55 | printf '{"Contents":[{"Key":"%s/%s"}]}' \ |
| 56 | "$PREFIX" "$WHEEL_FILENAME" > "$objects_path" |
| 57 | else |
| 58 | aws s3 cp "$wheel_path" "$S3_PREFIX" |
| 59 | aws s3api list-objects-v2 \ |
| 60 | --bucket "$BUCKET" \ |
| 61 | --prefix "$PREFIX/" \ |
| 62 | --delimiter / \ |
| 63 | --output json > "$objects_path" |
| 64 | fi |
| 65 | |
| 66 | # Pick Python >= 3.12 locally or use the container fallback. |
| 67 | # shellcheck source=../lib/select-python.sh |
| 68 | source .buildkite/scripts/lib/select-python.sh |
| 69 | select_python |
| 70 | |
| 71 | # The index generator only needs stdlib re here. Keep the tracked source intact. |
| 72 | sed 's/import regex as re/import re/' \ |
| 73 | .buildkite/scripts/generate-nightly-index.py > "$index_generator" |
| 74 | |
| 75 | # shellcheck disable=SC2086 |
| 76 | $PYTHON "$index_generator" \ |
| 77 | --version "$PREFIX" \ |
| 78 | --wheel-dir "$work_dir/$PREFIX" \ |
| 79 | --current-objects "$objects_path" \ |
| 80 | --output-dir "$work_dir" \ |
| 81 | --comment "XPU Triton shim" |
| 82 | |
| 83 | grep -Fq 'href="triton/"' "$index_output_dir/index.html" |
| 84 | grep -Fq "href=\"../$ENCODED_WHEEL_FILENAME\"" \ |
| 85 | "$index_output_dir/triton/index.html" |
| 86 | grep -Fq "\"filename\": \"$WHEEL_FILENAME\"" \ |
| 87 | "$index_output_dir/triton/metadata.json" |
| 88 | |
| 89 | if [[ "$DRY_RUN" == "1" ]]; then |
| 90 | echo "Dry run succeeded; generated files:" |
| 91 | find "$index_output_dir" -type f -print | sort |
| 92 | else |
| 93 | aws s3 cp --recursive "$index_output_dir/" "$S3_PREFIX" |
| 94 | echo "Published XPU Triton shim index to https://wheels.vllm.ai/$PREFIX/" |
| 95 | aws s3 cp "$S3_PREFIX$WHEEL_FILENAME" \ |
| 96 | "s3://$BUCKET/$COMMIT/$WHEEL_FILENAME" |
| 97 | echo "Staged XPU Triton shim for https://wheels.vllm.ai/$COMMIT/xpu/" |
| 98 | fi |