fix(sgl-kernel): align wheel METADATA/WHEEL with +cu filename (#21437)

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."