Merge pull request #2438 from Xarthisius/parallel_build

Use all available cores for both cythonize and build_ext

setup.py

@@ -1,4 +1,7 @@
 import os
+import multiprocessing
+import platform
+from concurrent.futures import ThreadPoolExecutor as Pool
 import glob
 import sys
 from sys import platform as _platform
@@ -12,9 +15,36 @@
     read_embree_location, \
     in_conda_env
 from distutils.version import LooseVersion
+from distutils.ccompiler import CCompiler
 import pkg_resources
 
 
+def _get_cpu_count():
+    if platform.system() != "Windows":
+        return os.cpu_count()
+
+def _compile(
+    self, sources, output_dir=None, macros=None, include_dirs=None,
+    debug=0, extra_preargs=None, extra_postargs=None, depends=None,
+):
+    """Function to monkey-patch distutils.ccompiler.CCompiler"""
+    macros, objects, extra_postargs, pp_opts, build = self._setup_compile(
+        output_dir, macros, include_dirs, sources, depends, extra_postargs
+    )
+    cc_args = self._get_cc_args(pp_opts, debug, extra_preargs)
+
+    for obj in objects:
+        try:
+            src, ext = build[obj]
+        except KeyError:
+            continue
+        self._compile(obj, src, ext, cc_args, extra_postargs, pp_opts)
+
+    # Return *all* object filenames, not just the ones we just built.
+    return objects
+
+CCompiler.compile = _compile
+
 if sys.version_info < (3, 5):
     print("yt currently supports versions newer than Python 3.5")
     print("certain features may fail unexpectedly and silently with older "
@@ -328,7 +358,9 @@ def finalize_options(self):
         from Cython.Build import cythonize
         self.distribution.ext_modules[:] = cythonize(
             self.distribution.ext_modules,
-            compiler_directives={'language_level': 2})
+            compiler_directives={'language_level': 2},
+            nthreads=_get_cpu_count(),
+        )
         _build_ext.finalize_options(self)
         # Prevent numpy from thinking it is still in its setup process
         # see http://stackoverflow.com/a/21621493/1382869
@@ -341,6 +373,16 @@ def finalize_options(self):
         import numpy
         self.include_dirs.append(numpy.get_include())
 
+    def build_extensions(self):
+        self.check_extensions_list(self.extensions)
+
+        if _get_cpu_count():
+            with Pool(_get_cpu_count()) as pool:
+                pool.map(self.build_extension, self.extensions)
+        else:
+            super().build_extensions()
+
+
 class sdist(_sdist):
     # subclass setuptools source distribution builder to ensure cython
     # generated C files are included in source distribution.
@@ -351,6 +393,7 @@ def run(self):
         cythonize(
             cython_extensions,
             compiler_directives={'language_level': 2},
+            nthreads=_get_cpu_count(),
         )
         _sdist.run(self)