diff --git a/sgl-kernel/rename_wheels.sh b/sgl-kernel/rename_wheels.sh index 915f069e6..c14492bb3 100755 --- a/sgl-kernel/rename_wheels.sh +++ b/sgl-kernel/rename_wheels.sh @@ -1,34 +1,85 @@ #!/usr/bin/env bash +# Align CUDA wheel filenames (+cu124/+cu128/+cu130) with internal METADATA Version and +# WHEEL tags after build (fixes pip "inconsistent version" when only the .whl name changed). +# Unpack → patch WHEEL/METADATA → wheel pack (RECORD regenerated; no hand-editing). set -ex WHEEL_DIR="dist" -wheel_files=($WHEEL_DIR/*.whl) -for wheel in "${wheel_files[@]}"; do - intermediate_wheel="${wheel/linux/manylinux2014}" - - # Extract the current python version from the wheel name - if [[ $intermediate_wheel =~ -cp([0-9]+)- ]]; then - cp_version="${BASH_REMATCH[1]}" +detect_cuda_suffix() { + if ls /usr/local/ 2>/dev/null | grep -q "12.4"; then + echo "+cu124" + elif ls /usr/local/ 2>/dev/null | grep -q "12.8"; then + echo "+cu128" + elif ls /usr/local/ 2>/dev/null | grep -q "13.0"; then + echo "+cu130" else - echo "Could not extract Python version from wheel name: $intermediate_wheel" + echo "" + fi +} + +CUDA_SUFFIX=$(detect_cuda_suffix) + +patch_wheel_platform_tags() { + local wheel_file="$1" + # Line-end anchors: "linux_x86_64" is a substring of "manylinux2014_x86_64", so + # unanchored global replace corrupts tags on a second run. + sed -i \ + -e 's/-linux_x86_64$/-manylinux2014_x86_64/' \ + -e 's/-linux_aarch64$/-manylinux2014_aarch64/' \ + "$wheel_file" +} + +wheel_files=("$WHEEL_DIR"/*.whl) +for wheel in "${wheel_files[@]}"; do + [[ -f "$wheel" ]] || continue + + intermediate_wheel="$wheel" + case "$wheel" in + *-linux_x86_64.whl) + intermediate_wheel="${wheel%-linux_x86_64.whl}-manylinux2014_x86_64.whl" + ;; + *-linux_aarch64.whl) + intermediate_wheel="${wheel%-linux_aarch64.whl}-manylinux2014_aarch64.whl" + ;; + esac + if [[ "$wheel" != "$intermediate_wheel" ]]; then + mv -- "$wheel" "$intermediate_wheel" + wheel="$intermediate_wheel" + fi + + if [[ -z "$CUDA_SUFFIX" ]]; then continue fi - # Detect CUDA version and add appropriate suffix - if ls /usr/local/ | grep -q "12.4"; then - new_wheel="${intermediate_wheel/-cp${cp_version}/+cu124-cp${cp_version}}" - elif ls /usr/local/ | grep -q "12.8"; then - new_wheel="${intermediate_wheel/-cp${cp_version}/+cu128-cp${cp_version}}" - elif ls /usr/local/ | grep -q "13.0"; then - new_wheel="${intermediate_wheel/-cp${cp_version}/+cu130-cp${cp_version}}" - else - new_wheel="$intermediate_wheel" - fi + TMPDIR=$(mktemp -d) + trap 'rm -rf -- "$TMPDIR"' ERR - if [[ "$wheel" != "$new_wheel" ]]; then - echo "Renaming $wheel to $new_wheel" - mv -- "$wheel" "$new_wheel" + python3 -m wheel unpack "$wheel" --dest "$TMPDIR" + UNPACKED=$(find "$TMPDIR" -mindepth 1 -maxdepth 1 -type d | head -1) + DIST_INFO=$(find "$UNPACKED" -maxdepth 1 -type d -name "*.dist-info" | head -1) + WHEEL_META="${DIST_INFO}/WHEEL" + METADATA_FILE="${DIST_INFO}/METADATA" + + patch_wheel_platform_tags "$WHEEL_META" + + ORIG_VERSION=$(grep '^Version:' "$METADATA_FILE" | head -1 | sed 's/^Version:[[:space:]]*//') + if [[ "$ORIG_VERSION" == *"$CUDA_SUFFIX"* ]]; then + echo "Skipping $wheel: version in METADATA is already suffixed." + rm -rf "$TMPDIR" + trap - ERR + continue fi + NEW_VERSION="${ORIG_VERSION}${CUDA_SUFFIX}" + sed -i "s/^Version:.*/Version: ${NEW_VERSION}/" "$METADATA_FILE" + + OLD_BASE=$(basename "$DIST_INFO") + NEW_BASE="${OLD_BASE/${ORIG_VERSION}/${NEW_VERSION}}" + mv "$DIST_INFO" "${UNPACKED}/${NEW_BASE}" + + rm -f "$wheel" + python3 -m wheel pack "$UNPACKED" --dest-dir "$WHEEL_DIR" + rm -rf "$TMPDIR" + trap - ERR done echo "Wheel renaming completed."