Update GitHub workflows

Author: raphaelsty

.github/workflows/pypi.yaml

@@ -36,7 +36,7 @@ jobs:
       - name: Install build dependencies
         run: |
           apt-get update
-          apt-get install -y --no-install-recommends build-essential curl libssl-dev pkg-config unzip
+          apt-get install -y --no-install-recommends build-essential curl libssl-dev pkg-config unzip zip
         shell: bash
 
       - name: Install Rust
@@ -61,26 +61,48 @@ jobs:
           echo "CUDA_LIB_DIR=$LIB_DIR" >> $GITHUB_ENV
         shell: bash
 
-      - name: Build wheel with auditwheel exclusions
+      - name: Build wheel
         env:
           LD_LIBRARY_PATH: ${{ env.CUDA_LIB_DIR }}:${LD_LIBRARY_PATH}
           CUDA_COMPUTE_CAP: ${{ matrix.variant.compute-caps }}
-        run: |
-          # Explicitly export AUDITWHEEL_EXCLUDE to ensure it's picked up by child processes.
-          # This prevents bundling large CUDA libraries into the wheel.
-          export AUDITWHEEL_EXCLUDE="libcuda.so,libnvJitLink.so,libcudart.so,libcublas.so,libcufft.so,libcurand.so,libcusolver.so,libcusparse.so,libcudnn.so"
-          echo "AUDITWHEEL_EXCLUDE is set to: $AUDITWHEEL_EXCLUDE"
-          maturin build --release --out dist
+        run: maturin build --release --out dist
         shell: bash
 
-      - name: Verify wheel contents and size
+      - name: Strip CUDA libraries and Verify
         run: |
-          echo "--- Verifying wheel contents ---"
+          set -e
+          echo "--- Stripping CUDA libraries from the wheel ---"
           WHEEL_FILE=$(ls dist/*.whl)
-          # The grep command should now only show your project's .so file and system libs like libcrypto
-          unzip -l $WHEEL_FILE | grep '\.so'
-          echo "--- Verifying wheel size ---"
-          ls -lh dist
+          echo "Original wheel size:"
+          ls -lh "$WHEEL_FILE"
+
+          UNPACK_DIR=$(mktemp -d)
+          unzip -q "$WHEEL_FILE" -d "$UNPACK_DIR"
+
+          echo "Deleting CUDA libraries..."
+          find "$UNPACK_DIR" -type f \( \
+            -name 'libcuda.so*' -o \
+            -name 'libcublas.so*' -o \
+            -name 'libcublasLt.so*' -o \
+            -name 'libnvJitLink.so*' -o \
+            -name 'libcudart.so*' -o \
+            -name 'libcufft.so*' -o \
+            -name 'libcurand.so*' -o \
+            -name 'libcusolver.so*' -o \
+            -name 'libcusparse.so*' -o \
+            -name 'libcudnn.so*' \
+          \) -print -delete
+
+          echo "Re-packing the wheel..."
+          (cd "$UNPACK_DIR" && zip -r -q "$WHEEL_FILE" .)
+          rm -rf "$UNPACK_DIR"
+
+          echo "--- Verification ---"
+          echo "Final wheel size:"
+          ls -lh "$WHEEL_FILE"
+          echo "Final wheel contents:"
+          unzip -l "$WHEEL_FILE" | grep '\.so'
+        shell: bash
 
       - name: Upload wheel artifact
         uses: actions/upload-artifact@v4