Internal review: Feature/concurrent build

cppimport/__init__.py

@@ -8,9 +8,11 @@
 from cppimport.find import _check_first_line_contains_cppimport
 
 settings = dict(
-    force_rebuild=False,
+    force_rebuild=False,  # `force_rebuild` with multiple processes is not supported
     file_exts=[".cpp", ".c"],
     rtld_flags=ctypes.RTLD_LOCAL,
+    lock_suffix='.lock',
+    lock_timeout=10*60,
     remove_strict_prototypes=True,
     release_mode=os.getenv("CPPIMPORT_RELEASE_MODE", "0").lower()
     in ("true", "yes", "1"),
@@ -60,16 +62,15 @@ def imp_from_filepath(filepath, fullname=None):
         is_build_needed,
         load_module,
         setup_module_data,
-        template_and_build,
-        try_load,
+        build_safely,
     )
-
+    filepath = os.path.abspath(filepath)
     if fullname is None:
         fullname = os.path.splitext(os.path.basename(filepath))[0]
     module_data = setup_module_data(fullname, filepath)
-    if is_build_needed(module_data) or not try_load(module_data):
-        template_and_build(filepath, module_data)
-        load_module(module_data)
+    if is_build_needed(module_data):
+        build_safely(filepath, module_data)
+    load_module(module_data)
     return module_data["module"]
 
 
@@ -110,15 +111,16 @@ def build_filepath(filepath, fullname=None):
     from cppimport.importer import (
         is_build_needed,
         setup_module_data,
-        template_and_build,
+        build_safely,
+        load_module,
     )
-
+    filepath = os.path.abspath(filepath)
     if fullname is None:
         fullname = os.path.splitext(os.path.basename(filepath))[0]
     module_data = setup_module_data(fullname, filepath)
     if is_build_needed(module_data):
-        template_and_build(filepath, module_data)
-
+        build_safely(filepath, module_data)
+    load_module(module_data)
     # Return the path to the built module
     return module_data["ext_path"]
 

cppimport/importer.py

@@ -3,6 +3,9 @@
 import os
 import sys
 import sysconfig
+from contextlib import suppress
+from time import time, sleep
+import filelock
 
 import cppimport
 from cppimport.build_module import build_module
@@ -12,6 +15,36 @@
 logger = logging.getLogger(__name__)
 
 
+def build_safely(filepath, module_data):
+    """Protect against race conditions when multiple processes executing `template_and_build`"""
+    binary_path = module_data['ext_path']
+    lock_path = binary_path + cppimport.settings['lock_suffix']
+
+    build_completed = lambda: os.path.exists(binary_path) and is_checksum_valid(module_data)
+
+    t = time()
+
+    # Race to obtain the lock and build. Other processes can wait
+    while not build_completed() and time() - t < cppimport.settings['lock_timeout']:
+        try:
+            with filelock.FileLock(lock_path, timeout=1):
+                if build_completed():
+                    break
+                template_and_build(filepath, module_data)
+        except filelock.Timeout:
+            logging.debug(f'Could not obtain lock (pid {os.getpid()})')
+            sleep(1)
+
+    if os.path.exists(lock_path):
+        with suppress(OSError):
+            os.remove(lock_path)
+
+    if not build_completed():
+        raise Exception(
+            f'Could not compile binary as lock already taken and timed out. Try increasing the timeout setting if '
+            f'the build time is longer (pid {os.getpid()}).')
+
+
 def template_and_build(filepath, module_data):
     logger.debug(f"Compiling {filepath}.")
     run_templating(module_data)
@@ -77,13 +110,3 @@ def is_build_needed(module_data):
     logger.debug(f"Matching checksum for {module_data['filepath']} --> not compiling")
     return False
 
-
-def try_load(module_data):
-    try:
-        load_module(module_data)
-        return True
-    except ImportError as e:
-        logger.info(
-            f"ImportError during import with matching checksum: {e}. Trying to rebuild."
-        )
-        return False

environment.yml

@@ -10,3 +10,4 @@ dependencies:
   - pytest
   - pytest-cov
   - pre-commit
+  - filelock

setup.py

@@ -6,7 +6,7 @@
     use_scm_version={"version_scheme": "post-release"},
     setup_requires=["setuptools_scm"],
     packages=["cppimport"],
-    install_requires=["mako", "pybind11"],
+    install_requires=["mako", "pybind11", "filelock"],
     zip_safe=False,
     name="cppimport",
     description="Import C++ files directly from Python!",

tests/test_cppimport.py

@@ -2,8 +2,11 @@
 import copy
 import logging
 import os
+import shutil
 import subprocess
 import sys
+from multiprocessing import Process
+from tempfile import TemporaryDirectory
 
 import cppimport
 import cppimport.build_module
@@ -45,6 +48,20 @@ def subprocess_check(test_code, returncode=0):
     assert p.returncode == returncode
 
 
+@contextlib.contextmanager
+def tmp_dir(files=None):
+    """Create a temporary directory and copy `files` into it. `files` can also include directories."""
+    files = files if files else []
+
+    with TemporaryDirectory() as tmp_path:
+        for f in files:
+            if os.path.isdir(f):
+                shutil.copytree(f, os.path.join(tmp_path, os.path.basename(f)))
+            else:
+                shutil.copyfile(f, os.path.join(tmp_path, os.path.basename(f)))
+        yield tmp_path
+
+
 def test_find_module_cpppath():
     mymodule_loc = find_module_cpppath("mymodule")
     mymodule_dir = os.path.dirname(mymodule_loc)
@@ -170,3 +187,22 @@ def test_import_hook():
 
     cppimport.force_rebuild(False)
     hook_test
+
+
+def test_multiple_processes():
+    with tmp_dir(['hook_test.cpp']) as tmp_path:
+        test_code = f"""
+import os;
+os.chdir('{tmp_path}');
+import cppimport.import_hook;
+import hook_test;
+        """
+        processes = [Process(target=subprocess_check, args=(test_code, )) for i in range(100)]
+
+        for p in processes:
+            p.start()
+
+        for p in processes:
+            p.join()
+
+        assert all(p.exitcode == 0 for p in processes)