@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
sh 156 lines 5.47 KB
Raw
1 #!/usr/bin/env bash
2 # SPDX-License-Identifier: Apache-2.0
3 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
4 #
5 # Upload ROCm wheels to S3 with proper index generation
6 #
7 # Required environment variables:
8 # AWS_ACCESS_KEY_ID, AWS_SECRET_ACCESS_KEY (or IAM role)
9 # S3_BUCKET (default: vllm-wheels)
10 #
11 # S3 path structure:
12 # s3://vllm-wheels/rocm/{commit}/ - All wheels for this commit
13 # s3://vllm-wheels/rocm/nightly/ - Index pointing to latest nightly
14 # s3://vllm-wheels/rocm/{version}/ - Index for release versions
15
16 set -ex
17
18 # ======== Configuration ========
19 BUCKET="${S3_BUCKET:-vllm-wheels}"
20 ROCM_SUBPATH="rocm/${BUILDKITE_COMMIT}"
21 S3_COMMIT_PREFIX="s3://$BUCKET/$ROCM_SUBPATH/"
22 INDICES_OUTPUT_DIR="rocm-indices"
23
24 echo "========================================"
25 echo "ROCm Wheel Upload Configuration"
26 echo "========================================"
27 echo "S3 Bucket: $BUCKET"
28 echo "S3 Path: $ROCM_SUBPATH"
29 echo "Commit: $BUILDKITE_COMMIT"
30 echo "Branch: $BUILDKITE_BRANCH"
31 echo "========================================"
32
33 # ======== Part 0: Setup Python and helpers ========
34
35 # Pick a Python interpreter for index generation -- local if recent
36 # enough, else a one-shot docker fallback.
37 # shellcheck source=lib/select-python.sh
38 source .buildkite/scripts/lib/select-python.sh
39 select_python
40
41 # Set up auditwheel-in-a-container for the manylinux retagging step.
42 # Distinct from select_python: ``manylinux.sh`` deliberately pins both
43 # the Python and auditwheel versions (the script reads auditwheel
44 # internals) and so always runs in a known-good container regardless
45 # of what's on the agent.
46 # shellcheck source=lib/manylinux.sh
47 source .buildkite/scripts/lib/manylinux.sh
48
49 # ======== Part 1: Collect and prepare wheels ========
50
51 # Collect all wheels
52 mkdir -p all-rocm-wheels
53 cp artifacts/rocm-base-wheels/*.whl all-rocm-wheels/ 2>/dev/null || true
54 cp artifacts/rocm-vllm-wheel/*.whl all-rocm-wheels/ 2>/dev/null || true
55
56 WHEEL_COUNT=$(find all-rocm-wheels -maxdepth 1 -name '*.whl' 2>/dev/null | wc -l)
57 echo "Total wheels to upload: $WHEEL_COUNT"
58
59 if [ "$WHEEL_COUNT" -eq 0 ]; then
60 echo "ERROR: No wheels found to upload!"
61 exit 1
62 fi
63
64 # Detect the appropriate manylinux platform tag for any wheel that still
65 # carries the generic ``linux_<arch>`` tag, and rename it in place. We use
66 # auditwheel via ``apply_manylinux_tag`` (see lib/manylinux.sh) rather than
67 # a hard-coded ``manylinux_2_35`` string so that the label tracks the actual
68 # glibc symbol versions used by the binaries (and stays correct if the
69 # rocm_base image is rebased).
70 #
71 # The ``linux``/``manylinux`` filter below skips both pre-tagged wheels
72 # (e.g. upstream torch) and pure-Python ``-any.whl`` wheels.
73 for wheel in all-rocm-wheels/*.whl; do
74 if [[ "$wheel" == *"linux"* ]] && [[ "$wheel" != *"manylinux"* ]]; then
75 new_wheel="$(apply_manylinux_tag "$wheel")"
76 echo "Renamed: $(basename "$wheel") -> $(basename "$new_wheel")"
77 fi
78 done
79
80 echo ""
81 echo "Wheels to upload:"
82 ls -lh all-rocm-wheels/
83
84 # ======== Part 2: Upload wheels to S3 ========
85
86 echo ""
87 echo "Uploading wheels to $S3_COMMIT_PREFIX"
88 for wheel in all-rocm-wheels/*.whl; do
89 aws s3 cp "$wheel" "$S3_COMMIT_PREFIX"
90 done
91
92 # ======== Part 3: Generate and upload indices ========
93
94 # List existing wheels in commit directory
95 echo ""
96 echo "Generating indices..."
97 obj_json="rocm-objects.json"
98 aws s3api list-objects-v2 --bucket "$BUCKET" --prefix "$ROCM_SUBPATH/" --delimiter / --output json > "$obj_json"
99
100 mkdir -p "$INDICES_OUTPUT_DIR"
101
102 # Use the existing generate-nightly-index.py
103 # HACK: Replace regex module with stdlib re (same as CUDA script)
104 sed -i 's/import regex as re/import re/g' .buildkite/scripts/generate-nightly-index.py
105
106 $PYTHON .buildkite/scripts/generate-nightly-index.py \
107 --version "$ROCM_SUBPATH" \
108 --current-objects "$obj_json" \
109 --output-dir "$INDICES_OUTPUT_DIR" \
110 --comment "ROCm commit $BUILDKITE_COMMIT"
111
112 # Upload indices to commit directory
113 echo "Uploading indices to $S3_COMMIT_PREFIX"
114 aws s3 cp --recursive "$INDICES_OUTPUT_DIR/" "$S3_COMMIT_PREFIX"
115
116 # Only scheduled nightly builds should update the moving nightly index.
117 if [[ "${NIGHTLY:-0}" == "1" ]]; then
118 echo "Updating rocm/nightly/ index..."
119 aws s3 cp --recursive "$INDICES_OUTPUT_DIR/" "s3://$BUCKET/rocm/nightly/"
120 fi
121
122 # Extract version from vLLM wheel and update version-specific index
123 VLLM_WHEEL=$(find all-rocm-wheels -maxdepth 1 -name 'vllm*.whl' 2>/dev/null | head -1)
124 if [ -n "$VLLM_WHEEL" ]; then
125 VERSION=$(unzip -p "$VLLM_WHEEL" '**/METADATA' | grep '^Version: ' | cut -d' ' -f2)
126 echo "Version in wheel: $VERSION"
127 PURE_VERSION="${VERSION%%+*}"
128 PURE_VERSION="${PURE_VERSION%%.rocm}"
129 echo "Pure version: $PURE_VERSION"
130
131 if [[ "$VERSION" != *"dev"* ]]; then
132 echo "Updating rocm/$PURE_VERSION/ index..."
133 aws s3 cp --recursive "$INDICES_OUTPUT_DIR/" "s3://$BUCKET/rocm/$PURE_VERSION/"
134 fi
135 fi
136
137 # ======== Part 4: Summary ========
138
139 echo ""
140 echo "========================================"
141 echo "ROCm Wheel Upload Complete!"
142 echo "========================================"
143 echo ""
144 echo "Wheels available at:"
145 echo " s3://$BUCKET/$ROCM_SUBPATH/"
146 echo ""
147 echo "Install command (by commit):"
148 echo " pip install vllm --extra-index-url https://${BUCKET}.s3.amazonaws.com/$ROCM_SUBPATH/"
149 echo ""
150 if [[ "${NIGHTLY:-0}" == "1" ]]; then
151 echo "Install command (nightly):"
152 echo " pip install vllm --extra-index-url https://${BUCKET}.s3.amazonaws.com/rocm/nightly/"
153 fi
154 echo ""
155 echo "Wheel count: $WHEEL_COUNT"
156 echo "========================================"