Add EasyBlock for PyTorch

easybuild/easyblocks/p/pytorch.py

@@ -0,0 +1,174 @@
+##
+# Copyright 2020-2020 Ghent University
+#
+# This file is part of EasyBuild,
+# originally created by the HPC team of Ghent University (http://ugent.be/hpc/en),
+# with support of Ghent University (http://ugent.be/hpc),
+# the Flemish Supercomputer Centre (VSC) (https://www.vscentrum.be),
+# Flemish Research Foundation (FWO) (http://www.fwo.be/en)
+# and the Department of Economy, Science and Innovation (EWI) (http://www.ewi-vlaanderen.be/en).
+#
+# https://github.com/easybuilders/easybuild
+#
+# EasyBuild is free software: you can redistribute it and/or modify
+# it under the terms of the GNU General Public License as published by
+# the Free Software Foundation v2.
+#
+# EasyBuild is distributed in the hope that it will be useful,
+# but WITHOUT ANY WARRANTY; without even the implied warranty of
+# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+# GNU General Public License for more details.
+#
+# You should have received a copy of the GNU General Public License
+# along with EasyBuild.  If not, see <http://www.gnu.org/licenses/>.
+##
+"""
+EasyBuild support for building and installing PyTorch, implemented as an easyblock
+
+@author: Alexander Grund (TU Dresden)
+"""
+
+import os
+import re
+import tempfile
+from easybuild.easyblocks.generic.pythonpackage import PythonPackage
+from easybuild.framework.easyconfig import CUSTOM
+from easybuild.tools.build_log import EasyBuildError
+from easybuild.tools.config import build_option
+import easybuild.tools.environment as env
+from easybuild.tools.modules import get_software_root
+from easybuild.tools.systemtools import POWER, get_cpu_architecture
+
+
+class EB_PyTorch(PythonPackage):
+    """Support for building/installing TensorFlow."""
+
+    @staticmethod
+    def extra_options():
+        extra_vars = PythonPackage.extra_options()
+        extra_vars.update({
+            # see https://developer.nvidia.com/cuda-gpus
+            'cuda_compute_capabilities': [[], 'List of CUDA compute capabilities to build with', CUSTOM],
+            'excluded_tests': [{}, 'Mapping of architecture strings to list of tests to be excluded', CUSTOM],
+            'extra_opts': [[], 'List of extra options, overwrites defaults from EasyBlock, e.g. USE_MKLDNN=0', CUSTOM]
+        })
+        extra_vars['download_dep_fail'][0] = True
+        extra_vars['sanity_pip_check'][0] = True
+
+        return extra_vars
+
+    def __init__(self, *args, **kwargs):
+        super(EB_PyTorch, self).__init__(*args, **kwargs)
+        self.options['modulename'] = 'torch'
+        # Test as-if pytorch was installed
+        self.testinstall = True
+        self.tmpdir = tempfile.mkdtemp(suffix='-pytorch-build')
+
+    def fetch_step(self, skip_checksums=False):
+        super(EB_PyTorch, self).fetch_step(skip_checksums)
+        # Resolve tests early to avoid failures later. Use obtain_file if path is not absolute
+        tests = [test if os.path.isabs(test) else self.obtain_file(test) for test in self.cfg['tests']]
+        self.cfg['tests'] = tests
+
+    def configure_step(self):
+        """Custom configure procedure for PyTorch."""
+        super(EB_PyTorch, self).configure_step()
+
+        options = ['PYTORCH_BUILD_VERSION=' + self.version, 'PYTORCH_BUILD_NUMBER=1']
+
+        # BLAS Interface
+        if get_software_root('imkl'):
+            options.append('BLAS=MKL')
+        else:
+            # This is what PyTorch defaults to if no MKL is found. Make this explicit here
+            options.append('BLAS=Eigen')
+
+        known_dependencies = ('FFmpeg', 'gflags', 'glog')
+        dependency_names = {dep['name'] for dep in self.cfg.dependencies()}
+        for dep in known_dependencies:
+            if dep in dependency_names:
+                options.append('USE_%s=1' % dep.upper())
+
+        # Some dependencies that are bundled with PyTorch can be used from the system
+        # No known ECs: 'cpuinfo', 'sleef', 'gloo', 'FP16', 'pthreadpool', 'psimd', 'FXDiv', 'ONNX', 'XNNPack'
+        #               'Eigen', 'GoogleBenchmark'
+        if 'protobuf' in dependency_names:
+            options.append('BUILD_CUSTOM_PROTOBUF=0')
+        # Always us Infiniband
+        options.append('USE_IBVERBS=1')
+
+        cudnn_root = get_software_root('cuDNN')
+        if cudnn_root:
+            options.append('CUDNN_LIB_DIR=%s/lib64' % cudnn_root)
+            options.append('CUDNN_INCLUDE_DIR=%s/include' % cudnn_root)
+
+        nccl_root = get_software_root('NCCL')
+        if nccl_root:
+            options.append('USE_SYSTEM_NCCL=1')
+            options.append('NCCL_INCLUDE_DIR=%s/include' % nccl_root)
+
+        # list of CUDA compute capabilities to use can be specifed in two ways (where (2) overrules (1)):
+        # (1) in the easyconfig file, via the custom cuda_compute_capabilities;
+        # (2) in the EasyBuild configuration, via --cuda-compute-capabilities configuration option;
+        cuda_cc = build_option('cuda_compute_capabilities') or self.cfg['cuda_compute_capabilities']
+        if not cuda_cc:
+            raise EasyBuildError('cuda_compute_capabilities must be specified')
+        self.log.info("Compiling with specified list of CUDA compute capabilities: %s", ', '.join(cuda_cc))
+        options.append('TORCH_CUDA_ARCH_LIST="%s"' % ';'.join(cuda_cc))
+
+        if get_cpu_architecture() == POWER:
+            # *NNPACK is not supported on Power, disable to avoid warnings
+            nnpacks = ('NNPACK', 'QNNPACK', 'USE_PYTORCH_QNNPACK', 'USE_XNNPACK')
+            options.extend('USE_%s=0' % nnpack for nnpack in nnpacks)
+
+        # Metal only supported on IOS which likely doesn't work with EB, so disabled
+        options.append('USE_METAL=0')
+
+        options.extend(self.cfg['extra_opts'])
+
+        self.cfg.update('prebuildopts', ' '.join(options) + ' ')
+        self.cfg.update('preinstallopts', ' '.join(options) + ' ')
+
+    def test_step(self):
+        """Run unit tests"""
+        # Make PyTorch tests not use the user home
+        env.setvar('XDG_CACHE_HOME', os.path.join(self.tmpdir, '.cache'))
+        # Parse excluded_tests and flatten into space separated string
+        excluded_tests = []
+        for arch, tests in self.cfg['excluded_tests'].items():
+            if not arch or arch == get_cpu_architecture():
+                excluded_tests.extend(tests)
+        # -x should not be used if there are no excluded tests
+        if excluded_tests:
+            excluded_tests = ['-x'] + excluded_tests
+        self.cfg.template_values.update({
+            'python': self.python_cmd,
+            'excluded_tests': ' '.join(excluded_tests)
+        })
+        super(EB_PyTorch, self).test_step()
+
+    def test_cases_step(self):
+        # Make PyTorch tests not use the user home
+        env.setvar('XDG_CACHE_HOME', os.path.join(self.tmpdir, '.cache'))
+        super(EB_PyTorch, self).test_cases_step()
+
+    def sanity_check_step(self, *args, **kwargs):
+        """Custom sanity check for PyTorch"""
+
+        super(EB_PyTorch, self).sanity_check_step(*args, **kwargs)
+
+        if self.cfg.get('download_dep_fail', True):
+            # CMake might mistakenly download dependencies during configure
+            pattern = r'^-- Downloading (\w+) to /'
+            downloaded_deps = re.findall(pattern, self.install_cmd_output, re.M)
+
+            if downloaded_deps:
+                fail_msg = "found one or more downloaded dependencies: %s" % ', '.join(downloaded_deps)
+                self.sanity_check_fail_msgs.append(fail_msg)
+
+    def make_module_req_guess(self):
+        """Set extra environment variables for PyTorch."""
+
+        guesses = super(EB_PyTorch, self).make_module_req_guess()
+        guesses['CMAKE_PREFIX_PATH'] = [os.path.join(self.pylibdir, 'torch')]
+        return guesses