From 8a09f34b0f859da9feb33f20eee24de5d7b7f1c7 Mon Sep 17 00:00:00 2001
From: Nicolas Hug <contact@nicolas-hug.com>
Date: Thu, 5 Dec 2024 15:32:21 +0000
Subject: [PATCH] Repair wheel with auditwheel to ship relevant so files (#12)

---
 .github/workflows/linux_wheel.yaml |  9 +++++----
 packaging/fake_smoke_test.py       |  4 ++++
 packaging/post_build_script.sh     | 25 ++++++++++++++++++++++++-
 packaging/pre_build_script.sh      |  1 +
 4 files changed, 34 insertions(+), 5 deletions(-)
 create mode 100644 packaging/fake_smoke_test.py

diff --git a/.github/workflows/linux_wheel.yaml b/.github/workflows/linux_wheel.yaml
index 059d072..57ee717 100644
--- a/.github/workflows/linux_wheel.yaml
+++ b/.github/workflows/linux_wheel.yaml
@@ -26,12 +26,12 @@ defaults:
 jobs:
 
   generate-matrix:
-    uses: pytorch/test-infra/.github/workflows/generate_binary_build_matrix.yml@main
+    uses: pytorch/test-infra/.github/workflows/generate_binary_build_matrix.yml@release/2.5
     with:
       package-type: wheel
       os: linux
       test-infra-repository: pytorch/test-infra
-      test-infra-ref: main
+      test-infra-ref: release/2.5
       with-xpu: disable
       with-rocm: disable
       with-cuda: disable
@@ -42,14 +42,15 @@ jobs:
     strategy:
       fail-fast: false
     name: Build and Upload Linux wheel
-    uses: pytorch/test-infra/.github/workflows/build_wheels_linux.yml@main
+    uses: pytorch/test-infra/.github/workflows/build_wheels_linux.yml@release/2.5
     with:
       repository: pytorch-labs/torchvision-extra-decoders
       ref: ""
       test-infra-repository: pytorch/test-infra
-      test-infra-ref: main
+      test-infra-ref: release/2.5
       build-matrix: ${{ needs.generate-matrix.outputs.matrix }}
       pre-script: packaging/pre_build_script.sh
       post-script: packaging/post_build_script.sh
+      smoke-test-script: packaging/fake_smoke_test.py
       package-name: torchvision_extra_decoders
       trigger-event: ${{ github.event_name }}
diff --git a/packaging/fake_smoke_test.py b/packaging/fake_smoke_test.py
new file mode 100644
index 0000000..ec6a5b9
--- /dev/null
+++ b/packaging/fake_smoke_test.py
@@ -0,0 +1,4 @@
+# This is a fake smoke test that runs on the test-infra instances after we build
+# a wheel.
+
+print("Success")
diff --git a/packaging/post_build_script.sh b/packaging/post_build_script.sh
index 2386988..2b84e85 100755
--- a/packaging/post_build_script.sh
+++ b/packaging/post_build_script.sh
@@ -2,7 +2,30 @@
 
 set -ex
 
-old="linux_x86_64"
+ls dist
+wheel_path=`find dist -type f -name "*.whl"`
+
+echo Found $wheel_path
+
+
+# Put all the .so files we want to ship with the wheels within the lib_to_bundle
+# dir. Maybe it's enough to LD_LIBRARY_PATH $CONDA_PREFIX/lib but I fear this
+# might include some un-desired libc or libcxx so files?
+mkdir libs_to_bundle
+for f in `find $CONDA_PREFIX/lib | grep -e libavif.so -e libdav1d.so -e librav1e.so -e libSvtAv1Enc.so -e libaom.so -e libheif.so -e libx265.so -e libde265.so`; do cp $f libs_to_bundle; done
+echo `ls libs_to_bundle`
+export LD_LIBRARY_PATH="$PWD/libs_to_bundle:$LD_LIBRARY_PATH"
+
+${CONDA_RUN} auditwheel -v repair --plat manylinux_2_34_x86_64 $wheel_path  --exclude libtorch_python.so --exclude libc10.so --exclude libtorch.so --exclude libtorch_cpu.so --wheel-dir dist
+
+# mv original wheel to a different dir, the 'repaired' wheel outputed by
+# auditwheel is still in dist/
+mkdir original_wheel
+mv $wheel_path original_wheel/
+
+
+# This is absolutely disgusting.
+old="manylinux_2_34_x86_64"
 new="manylinux_2_17_x86_64.manylinux2014_x86_64"
 echo "Replacing ${old} with ${new} in wheel name"
 mv dist/*${old}*.whl $(echo dist/*${old}*.whl | sed "s/${old}/${new}/")
diff --git a/packaging/pre_build_script.sh b/packaging/pre_build_script.sh
index 52df13b..3e5ea3e 100644
--- a/packaging/pre_build_script.sh
+++ b/packaging/pre_build_script.sh
@@ -1,3 +1,4 @@
 #!/bin/bash
 
+pip install auditwheel
 conda install libavif libheif -c conda-forge -yq