py-torchvision: fix linking to -lavcodec (#18093)

3 files changed, 48 insertions, 0 deletions

diff --git a/var/spack/repos/builtin/packages/ffmpeg/package.py b/var/spack/repos/builtin/packages/ffmpeg/package.py
index cc3a0216e8..f268040ac2 100644
--- a/var/spack/repos/builtin/packages/ffmpeg/package.py
+++ b/var/spack/repos/builtin/packages/ffmpeg/package.py
@@ -104,6 +104,16 @@ class Ffmpeg(AutotoolsPackage):
     conflicts('+libssh', when='@2.0.999:')
     conflicts('+libzmq', when='@:1.999.999')
 
+    @property
+    def libs(self):
+        return find_libraries('*', self.prefix, recursive=True)
+
+    @property
+    def headers(self):
+        headers = find_all_headers(self.prefix.include)
+        headers.directories = self.prefix.include
+        return headers
+
     def enable_or_disable_meta(self, variant, options):
         switch = 'enable' if '+{0}'.format(variant) in self.spec else 'disable'
         return ['--{0}-{1}'.format(switch, option) for option in options]
diff --git a/var/spack/repos/builtin/packages/py-torch/package.py b/var/spack/repos/builtin/packages/py-torch/package.py
index 1e2978dc54..71b71fe210 100644
--- a/var/spack/repos/builtin/packages/py-torch/package.py
+++ b/var/spack/repos/builtin/packages/py-torch/package.py
@@ -195,6 +195,22 @@ class PyTorch(PythonPackage, CudaPackage):
     # Only run once to speed up build times
     phases = ['install']
 
+    @property
+    def libs(self):
+        root = join_path(
+            self.prefix, self.spec['python'].package.site_packages_dir,
+            'torch', 'lib')
+        return find_libraries('libtorch', root)
+
+    @property
+    def headers(self):
+        root = join_path(
+            self.prefix, self.spec['python'].package.site_packages_dir,
+            'torch', 'include')
+        headers = find_all_headers(root)
+        headers.directories = root
+        return headers
+
     def setup_build_environment(self, env):
         def enable_or_disable(variant, keyword='USE', var=None, newer=False):
             """Set environment variable to enable or disable support for a
diff --git a/var/spack/repos/builtin/packages/py-torchvision/package.py b/var/spack/repos/builtin/packages/py-torchvision/package.py
index c94f06ef6f..ae6402d1c1 100644
--- a/var/spack/repos/builtin/packages/py-torchvision/package.py
+++ b/var/spack/repos/builtin/packages/py-torchvision/package.py
@@ -37,11 +37,13 @@ class PyTorchvision(PythonPackage):
     depends_on('python@2.7:2.8,3.5:3.7', when='@:0.4', type=('build', 'run'))
 
     depends_on('py-setuptools', type='build')
+    depends_on('ninja', type='build')
     depends_on('py-numpy', type=('build', 'run'))
     depends_on('py-six', when='@:0.5', type=('build', 'run'))
 
     # See README.rst
     depends_on('py-torch@1.6.0', when='@0.7.0', type=('build', 'link', 'run'))
+    depends_on('py-torch@1.5.1', when='@0.6.1', type=('build', 'link', 'run'))
     depends_on('py-torch@1.5.0', when='@0.6.0', type=('build', 'link', 'run'))
     depends_on('py-torch@1.4.0', when='@0.5.0', type=('build', 'link', 'run'))
     depends_on('py-torch@1.3.1', when='@0.4.2', type=('build', 'link', 'run'))
@@ -66,3 +68,23 @@ class PyTorchvision(PythonPackage):
     depends_on('py-scipy', type='test')
 
     depends_on('ffmpeg@3.1:', when='@0.4.2:')
+
+    def setup_build_environment(self, env):
+        include = []
+        library = []
+        for dep in self.spec.dependencies(deptype='link'):
+            query = self.spec[dep.name]
+            include.extend(query.headers.directories)
+            library.extend(query.libs.directories)
+
+        # README says to use TORCHVISION_INCLUDE and TORCHVISION_LIBRARY,
+        # but these do not work. Build uses a mix of Spack's compiler wrapper
+        # and the actual compiler, so this is needed to get parts of the build
+        # working. See https://github.com/pytorch/vision/issues/2591
+        env.set('CPATH', ':'.join(include))
+        env.set('LIBRARY_PATH', ':'.join(library))
+
+        if '+cuda' in self.spec['py-torch']:
+            env.set('FORCE_CUDA', 1)
+        else:
+            env.set('FORCE_CUDA', 0)