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)