Merge pull request #5161 from juliantaylor/parallel-distutils

ENH: support parallel compilation of extensions

5 files changed, 105 insertions, 16 deletions

diff --git a/numpy/distutils/ccompiler.py b/numpy/distutils/ccompiler.py
index 8484685c0..d7fe702a6 100644
--- a/numpy/distutils/ccompiler.py
+++ b/numpy/distutils/ccompiler.py
@@ -16,7 +16,7 @@ from distutils.version import LooseVersion
 from numpy.distutils import log
 from numpy.distutils.exec_command import exec_command
 from numpy.distutils.misc_util import cyg2win32, is_sequence, mingw32, \
-                                      quote_args
+                                      quote_args, get_num_build_jobs
 from numpy.distutils.compat import get_exception
 
 
@@ -165,9 +165,10 @@ def CCompiler_compile(self, sources, output_dir=None, macros=None,
         return []
     # FIXME:RELATIVE_IMPORT
     if sys.version_info[0] < 3:
-        from .fcompiler import FCompiler
+        from .fcompiler import FCompiler, is_f_file, has_f90_header
     else:
-        from numpy.distutils.fcompiler import FCompiler
+        from numpy.distutils.fcompiler import (FCompiler, is_f_file,
+                                               has_f90_header)
     if isinstance(self, FCompiler):
         display = []
         for fc in ['f77', 'f90', 'fix']:
@@ -189,20 +190,45 @@ def CCompiler_compile(self, sources, output_dir=None, macros=None,
         display += "\nextra options: '%s'" % (' '.join(extra_postargs))
     log.info(display)
 
-    # build any sources in same order as they were originally specified
-    #   especially important for fortran .f90 files using modules
+    def single_compile(args):
+        obj, (src, ext) = args
+        self._compile(obj, src, ext, cc_args, extra_postargs, pp_opts)
+
     if isinstance(self, FCompiler):
         objects_to_build = list(build.keys())
+        f77_objects, other_objects = [], []
         for obj in objects:
             if obj in objects_to_build:
                 src, ext = build[obj]
                 if self.compiler_type=='absoft':
                     obj = cyg2win32(obj)
                     src = cyg2win32(src)
-                self._compile(obj, src, ext, cc_args, extra_postargs, pp_opts)
+                if is_f_file(src) and not has_f90_header(src):
+                    f77_objects.append((obj, (src, ext)))
+                else:
+                    other_objects.append((obj, (src, ext)))
+
+        # f77 objects can be built in parallel
+        build_items = f77_objects
+        # build f90 modules serial, module files are generated during
+        # compilation and may be used by files later in the list so the
+        # ordering is important
+        for o in other_objects:
+            single_compile(o)
+    else:
+        build_items = build.items()
+
+    jobs = get_num_build_jobs()
+    if len(build) > 1 and jobs > 1:
+        # build parallel
+        import multiprocessing.pool
+        pool = multiprocessing.pool.ThreadPool(jobs)
+        pool.map(single_compile, build_items)
+        pool.close()
     else:
-        for obj, (src, ext) in build.items():
-            self._compile(obj, src, ext, cc_args, extra_postargs, pp_opts)
+        # build serial
+        for o in build_items:
+            single_compile(o)
 
     # Return *all* object filenames, not just the ones we just built.
     return objects
diff --git a/numpy/distutils/command/build.py b/numpy/distutils/command/build.py
index b6912be15..f7249ae81 100644
--- a/numpy/distutils/command/build.py
+++ b/numpy/distutils/command/build.py
@@ -16,6 +16,8 @@ class build(old_build):
     user_options = old_build.user_options + [
         ('fcompiler=', None,
          "specify the Fortran compiler type"),
+        ('jobs=', 'j',
+         "number of parallel jobs"),
         ]
 
     help_options = old_build.help_options + [
@@ -26,8 +28,14 @@ class build(old_build):
     def initialize_options(self):
         old_build.initialize_options(self)
         self.fcompiler = None
+        self.jobs = None
 
     def finalize_options(self):
+        if self.jobs:
+            try:
+                self.jobs = int(self.jobs)
+            except ValueError:
+                raise ValueError("--jobs/-j argument must be an integer")
         build_scripts = self.build_scripts
         old_build.finalize_options(self)
         plat_specifier = ".%s-%s" % (get_platform(), sys.version[0:3])
diff --git a/numpy/distutils/command/build_clib.py b/numpy/distutils/command/build_clib.py
index 84ca87250..6e65a3bfb 100644
--- a/numpy/distutils/command/build_clib.py
+++ b/numpy/distutils/command/build_clib.py
@@ -30,6 +30,8 @@ class build_clib(old_build_clib):
         ('fcompiler=', None,
          "specify the Fortran compiler type"),
         ('inplace', 'i', 'Build in-place'),
+        ('jobs=', 'j',
+         "number of parallel jobs"),
         ]
 
     boolean_options = old_build_clib.boolean_options + ['inplace']
@@ -38,7 +40,16 @@ class build_clib(old_build_clib):
         old_build_clib.initialize_options(self)
         self.fcompiler = None
         self.inplace = 0
-        return
+        self.jobs = None
+
+    def finalize_options(self):
+        if self.jobs:
+            try:
+                self.jobs = int(self.jobs)
+            except ValueError:
+                raise ValueError("--jobs/-j argument must be an integer")
+        old_build_clib.finalize_options(self)
+        self.set_undefined_options('build', ('jobs', 'jobs'))
 
     def have_f_sources(self):
         for (lib_name, build_info) in self.libraries:
diff --git a/numpy/distutils/command/build_ext.py b/numpy/distutils/command/build_ext.py
index 68aab21eb..59c453607 100644
--- a/numpy/distutils/command/build_ext.py
+++ b/numpy/distutils/command/build_ext.py
@@ -34,6 +34,8 @@ class build_ext (old_build_ext):
     user_options = old_build_ext.user_options + [
         ('fcompiler=', None,
          "specify the Fortran compiler type"),
+        ('jobs=', 'j',
+         "number of parallel jobs"),
         ]
 
     help_options = old_build_ext.help_options + [
@@ -44,12 +46,19 @@ class build_ext (old_build_ext):
     def initialize_options(self):
         old_build_ext.initialize_options(self)
         self.fcompiler = None
+        self.jobs = None
 
     def finalize_options(self):
+        if self.jobs:
+            try:
+                self.jobs = int(self.jobs)
+            except ValueError:
+                raise ValueError("--jobs/-j argument must be an integer")
         incl_dirs = self.include_dirs
         old_build_ext.finalize_options(self)
         if incl_dirs is not None:
             self.include_dirs.extend(self.distribution.include_dirs or [])
+        self.set_undefined_options('build', ('jobs', 'jobs'))
 
     def run(self):
         if not self.extensions:
diff --git a/numpy/distutils/misc_util.py b/numpy/distutils/misc_util.py
index c146178f0..cab0a4cba 100644
--- a/numpy/distutils/misc_util.py
+++ b/numpy/distutils/misc_util.py
@@ -13,6 +13,10 @@ import shutil
 
 import distutils
 from distutils.errors import DistutilsError
+try:
+    from threading import local as tlocal
+except ImportError:
+    from dummy_threading import local as tlocal
 
 try:
     set
@@ -31,7 +35,8 @@ __all__ = ['Configuration', 'get_numpy_include_dirs', 'default_config_dict',
            'get_script_files', 'get_lib_source_files', 'get_data_files',
            'dot_join', 'get_frame', 'minrelpath', 'njoin',
            'is_sequence', 'is_string', 'as_list', 'gpaths', 'get_language',
-           'quote_args', 'get_build_architecture', 'get_info', 'get_pkg_info']
+           'quote_args', 'get_build_architecture', 'get_info', 'get_pkg_info',
+           'get_num_build_jobs']
 
 class InstallableLib(object):
     """
@@ -60,6 +65,36 @@ class InstallableLib(object):
         self.build_info = build_info
         self.target_dir = target_dir
 
+
+def get_num_build_jobs():
+    """
+    Get number of parallel build jobs set by the --jobs command line argument
+    of setup.py
+    If the command did not receive a setting the environment variable
+    NPY_NUM_BUILD_JOBS checked and if that is unset it returns 1.
+
+    Returns
+    -------
+    out : int
+        number of parallel jobs that can be run
+
+    """
+    from numpy.distutils.core import get_distribution
+    envjobs = int(os.environ.get("NPY_NUM_BUILD_JOBS", 1))
+    dist = get_distribution()
+    # may be None during configuration
+    if dist is None:
+        return envjobs
+
+    # any of these three may have the job set, take the largest
+    cmdattr = (getattr(dist.get_command_obj('build'), 'jobs'),
+               getattr(dist.get_command_obj('build_ext'), 'jobs'),
+               getattr(dist.get_command_obj('build_clib'), 'jobs'))
+    if all(x is None for x in cmdattr):
+        return envjobs
+    else:
+        return max(x for x in cmdattr if x is not None)
+
 def quote_args(args):
     # don't used _nt_quote_args as it does not check if
     # args items already have quotes or not.
@@ -249,9 +284,9 @@ def gpaths(paths, local_path='', include_non_existing=True):
     return _fix_paths(paths, local_path, include_non_existing)
 
 
-_temporary_directory = None
 def clean_up_temporary_directory():
-    global _temporary_directory
+    tdata = tlocal()
+    _temporary_directory = getattr(tdata, 'tempdir', None)
     if not _temporary_directory:
         return
     try:
@@ -261,13 +296,13 @@ def clean_up_temporary_directory():
     _temporary_directory = None
 
 def make_temp_file(suffix='', prefix='', text=True):
-    global _temporary_directory
-    if not _temporary_directory:
-        _temporary_directory = tempfile.mkdtemp()
+    tdata = tlocal()
+    if not hasattr(tdata, 'tempdir'):
+        tdata.tempdir = tempfile.mkdtemp()
         atexit.register(clean_up_temporary_directory)
     fid, name = tempfile.mkstemp(suffix=suffix,
                                  prefix=prefix,
-                                 dir=_temporary_directory,
+                                 dir=tdata.tempdir,
                                  text=text)
     fo = os.fdopen(fid, 'w')
     return fo, name