build extensions in parallel

yuantailing · yuantailing · commit f0df092f4f04 · 2025-02-20T04:33:16.000Z
diff --git a/README.md b/README.md
@@ -130,7 +130,7 @@ CUDA and C++ extensions via
 git clone https://github.com/NVIDIA/apex
 cd apex
 # if pip >= 23.1 (ref: https://pip.pypa.io/en/stable/news/#v23-1) which supports multiple `--config-settings` with the same key... 
-pip install -v --disable-pip-version-check --no-cache-dir --no-build-isolation --config-settings "--build-option=--cpp_ext" --config-settings "--build-option=--cuda_ext" ./
+pip install -v --disable-pip-version-check --no-cache-dir --no-build-isolation --config-settings "--build-option=--cpp_ext --cuda_ext --parallel 4" ./
 # otherwise
 pip install -v --disable-pip-version-check --no-cache-dir --no-build-isolation --global-option="--cpp_ext" --global-option="--cuda_ext" ./
 ```
diff --git a/setup.py b/setup.py
@@ -1,6 +1,7 @@
 import sys
 import warnings
 import os
+import threading
 import glob
 from packaging.version import parse, Version
 
@@ -859,6 +860,39 @@ def check_cudnn_version_and_warn(global_option: str, required_cudnn_version: int
     )
 
 
+# Patch because `setup.py bdist_wheel` does not accept the `parallel` option
+parallel = None
+if "--parallel" in sys.argv:
+    idx = sys.argv.index("--parallel")
+    parallel = int(sys.argv[idx + 1])
+    sys.argv.pop(idx + 1)
+    sys.argv.pop(idx)
+
+
+# Prevent file conflicts when multiple extensions are compiled simultaneously
+class BuildExtensionSeparateDir(BuildExtension):
+    build_extension_patch_lock = threading.Lock()
+    thread_ext_name_map = {}
+
+    def build_extension(self, ext):
+        with self.build_extension_patch_lock:
+            if not getattr(self.compiler, "_compile_separate_output_dir", False):
+                compile_orig = self.compiler.compile
+
+                def compile_new(*args, **kwargs):
+                    return compile_orig(*args, **{
+                        **kwargs,
+                        "output_dir": os.path.join(
+                            kwargs["output_dir"],
+                            self.thread_ext_name_map[threading.current_thread().ident]),
+                    })
+                self.compiler.compile = compile_new
+                self.compiler._compile_separate_output_dir = True
+        self.thread_ext_name_map[threading.current_thread().ident] = ext.name
+        objects = super().build_extension(ext)
+        return objects
+
+
 setup(
     name="apex",
     version="0.1",
@@ -868,6 +902,6 @@ def check_cudnn_version_and_warn(global_option: str, required_cudnn_version: int
     install_requires=["packaging>20.6"],
     description="PyTorch Extensions written by NVIDIA",
     ext_modules=ext_modules,
-    cmdclass={"build_ext": BuildExtension} if ext_modules else {},
+    cmdclass={"build_ext": BuildExtensionSeparateDir.with_options(parallel=parallel)} if ext_modules else {},
     extras_require=extras,
 )
