| 1 | #!/usr/bin/env bash |
| 2 | |
| 3 | set -ex |
| 4 | |
| 5 | # Generate and upload wheel indices for all wheels in the commit directory. |
| 6 | # This script should run once after all wheels have been built and uploaded. |
| 7 | |
| 8 | # ======== setup ======== |
| 9 | |
| 10 | BUCKET="vllm-wheels" |
| 11 | INDICES_OUTPUT_DIR="indices" |
| 12 | DEFAULT_VARIANT_ALIAS="cu130" # align with vLLM_MAIN_CUDA_VERSION in vllm/envs.py |
| 13 | SUBPATH=$BUILDKITE_COMMIT |
| 14 | S3_COMMIT_PREFIX="s3://$BUCKET/$SUBPATH/" |
| 15 | |
| 16 | # Select python3 (>= 3.12) -- local if available, else a docker fallback. |
| 17 | # shellcheck source=lib/select-python.sh |
| 18 | source .buildkite/scripts/lib/select-python.sh |
| 19 | select_python |
| 20 | |
| 21 | # ======== generate and upload indices ======== |
| 22 | |
| 23 | # list all wheels in the commit directory |
| 24 | echo "Existing wheels on S3:" |
| 25 | aws s3 ls "$S3_COMMIT_PREFIX" |
| 26 | obj_json="objects.json" |
| 27 | aws s3api list-objects-v2 --bucket "$BUCKET" --prefix "$SUBPATH/" --delimiter / --output json > "$obj_json" |
| 28 | mkdir -p "$INDICES_OUTPUT_DIR" |
| 29 | |
| 30 | # call script to generate indices for all existing wheels |
| 31 | # these indices have relative paths that work as long as they are next to the wheel directory in s3 |
| 32 | # i.e., the wheels are always in s3://vllm-wheels/<commit>/ |
| 33 | # and indices can be placed in /<commit>/, or /nightly/, or /<version>/ |
| 34 | alias_args=() |
| 35 | if [[ -n "$DEFAULT_VARIANT_ALIAS" ]]; then |
| 36 | alias_args=(--alias-to-default "$DEFAULT_VARIANT_ALIAS") |
| 37 | fi |
| 38 | |
| 39 | # HACK: we do not need regex module here, but it is required by pre-commit hook |
| 40 | # To avoid any external dependency, we simply replace it back to the stdlib re module |
| 41 | sed -i 's/import regex as re/import re/g' .buildkite/scripts/generate-nightly-index.py |
| 42 | $PYTHON .buildkite/scripts/generate-nightly-index.py --version "$SUBPATH" --current-objects "$obj_json" --output-dir "$INDICES_OUTPUT_DIR" --comment "commit $BUILDKITE_COMMIT" "${alias_args[@]}" |
| 43 | |
| 44 | # copy indices to /<commit>/ unconditionally |
| 45 | echo "Uploading indices to $S3_COMMIT_PREFIX" |
| 46 | aws s3 cp --recursive "$INDICES_OUTPUT_DIR/" "$S3_COMMIT_PREFIX" |
| 47 | |
| 48 | # copy to /nightly/ only when enabled for a main branch build that is not a PR |
| 49 | if [[ "${UPDATE_NIGHTLY_INDEX:-1}" == "1" && \ |
| 50 | "$BUILDKITE_BRANCH" == "main" && \ |
| 51 | "$BUILDKITE_PULL_REQUEST" == "false" ]]; then |
| 52 | echo "Uploading indices to overwrite /nightly/" |
| 53 | aws s3 cp --recursive "$INDICES_OUTPUT_DIR/" "s3://$BUCKET/nightly/" |
| 54 | fi |
| 55 | |
| 56 | # detect version from a vLLM wheel in the commit directory |
| 57 | first_wheel_key=$($PYTHON -c "import json; obj=json.load(open('$obj_json')); print(next((c['Key'] for c in obj.get('Contents', []) if c['Key'].rsplit('/', 1)[-1].startswith('vllm-') and c['Key'].endswith('.whl')), ''))") |
| 58 | if [[ -z "$first_wheel_key" ]]; then |
| 59 | echo "Error: No vLLM wheel found in $S3_COMMIT_PREFIX" |
| 60 | exit 1 |
| 61 | fi |
| 62 | first_wheel=$(basename "$first_wheel_key") |
| 63 | aws s3 cp "s3://$BUCKET/${first_wheel_key}" "/tmp/${first_wheel}" |
| 64 | version=$(unzip -p "/tmp/${first_wheel}" '**/METADATA' | grep '^Version: ' | cut -d' ' -f2) |
| 65 | rm -f "/tmp/${first_wheel}" |
| 66 | echo "Version in wheel: $version" |
| 67 | pure_version="${version%%+*}" |
| 68 | echo "Pure version (without variant): $pure_version" |
| 69 | |
| 70 | # re-generate and copy to /<pure_version>/ only if it does not have "dev" in the version |
| 71 | if [[ "${UPDATE_VERSION_INDEX:-1}" == "1" && "$version" != *"dev"* ]]; then |
| 72 | echo "Re-generating indices for /$pure_version/" |
| 73 | rm -rf "${INDICES_OUTPUT_DIR:?}" |
| 74 | mkdir -p "$INDICES_OUTPUT_DIR" |
| 75 | # wheel-dir is overridden to be the commit directory, so that the indices point to the correct wheel path |
| 76 | $PYTHON .buildkite/scripts/generate-nightly-index.py --version "$pure_version" --wheel-dir "$SUBPATH" --current-objects "$obj_json" --output-dir "$INDICES_OUTPUT_DIR" --comment "version $pure_version" "${alias_args[@]}" |
| 77 | aws s3 cp --recursive "$INDICES_OUTPUT_DIR/" "s3://$BUCKET/$pure_version/" |
| 78 | fi |