py-torchaudio, py-torchtext: rpath fixup (#47404)

* py-torchaudio, py-torchtext: rpath fixup

also add missing dependency on Spack ffmpeg to torchaudio.

* py-torchaudio: add torio rpath

2 files changed, 40 insertions, 0 deletions

diff --git a/var/spack/repos/builtin/packages/py-torchaudio/package.py b/var/spack/repos/builtin/packages/py-torchaudio/package.py
index 2ac903565f..59a7e2825b 100644
--- a/var/spack/repos/builtin/packages/py-torchaudio/package.py
+++ b/var/spack/repos/builtin/packages/py-torchaudio/package.py
@@ -104,6 +104,9 @@ class PyTorchaudio(PythonPackage):
     depends_on("cmake@3.5:", when="@0.8:", type="build")
     depends_on("ninja", when="@0.8:", type="build")
 
+    # prior to 2.1 ffmpeg was vendored
+    depends_on("ffmpeg@:6", when="@2.1:")
+
     # setup.py
     depends_on("py-setuptools", type="build")
     depends_on("py-pybind11", when="@0.12:", type=("build", "link"))
@@ -118,6 +121,22 @@ class PyTorchaudio(PythonPackage):
     )
     conflicts("^cuda@12.5:", when="@:2.1")
 
+    def patch(self):
+        # Add missing rpaths, which requires patching due to hardcoded cmake_args
+        if self.spec.satisfies("@0.8:"):
+            rpaths = [f"{python_platlib}/torchaudio/lib", f"{python_platlib}/torio/lib"]
+            cmake_args = [
+                f"-DCMAKE_INSTALL_RPATH={';'.join(rpaths)}",
+                "-DCMAKE_INSTALL_RPATH_USE_LINK_PATH=ON",
+            ]
+            cmake_str = ", ".join(f"'{arg}'" for arg in cmake_args)
+            filter_file(
+                "cmake_args = [",
+                f"cmake_args = [{cmake_str},",
+                "tools/setup_helpers/extension.py",
+                string=True,
+            )
+
     def flag_handler(self, name, flags):
         # https://github.com/pytorch/vision/issues/8653
         if name == "ldflags":
@@ -129,6 +148,12 @@ class PyTorchaudio(PythonPackage):
         # tools/setup_helpers/extension.py
         env.set("BUILD_SOX", 0)
 
+        if self.spec.satisfies("@2.1:"):
+            env.set("FFMPEG_ROOT", self.spec["ffmpeg"].prefix)
+        else:
+            # a specific ffmpeg is built but not installed, so just disable
+            env.set("USE_FFMPEG", "0")
+
         if "+cuda" in self.spec["py-torch"]:
             env.set("USE_CUDA", 1)
         else:
diff --git a/var/spack/repos/builtin/packages/py-torchtext/package.py b/var/spack/repos/builtin/packages/py-torchtext/package.py
index 95fde8068a..de5c66f20e 100644
--- a/var/spack/repos/builtin/packages/py-torchtext/package.py
+++ b/var/spack/repos/builtin/packages/py-torchtext/package.py
@@ -93,3 +93,18 @@ class PyTorchtext(PythonPackage):
     depends_on("py-pybind11", when="@0.8:", type=("build", "link"))
     depends_on("py-six", when="@:0.6", type=("build", "run"))
     depends_on("py-sentencepiece", when="@:0.7", type=("build", "run"))
+
+    def patch(self):
+        # Add missing rpaths, which requires patching due to hardcoded cmake_args
+        if self.spec.satisfies("@0.13:"):
+            cmake_args = [
+                f"-DCMAKE_INSTALL_RPATH={python_platlib}/torchtext/lib",
+                "-DCMAKE_INSTALL_RPATH_USE_LINK_PATH=ON",
+            ]
+            cmake_str = ", ".join(f"'{arg}'" for arg in cmake_args)
+            filter_file(
+                "cmake_args = [",
+                f"cmake_args = [{cmake_str},",
+                "tools/setup_helpers/extension.py",
+                string=True,
+            )