Repository: pytorch/pytorch

Issue:
[jit cpp extensions] Possible file corruption in distributed setting
I believe there might be some sort of race condition that might happen when compiling jit cpp extensions in a distributed setting.

Here a summary of what happens:
- I launch my code from multiple different processes (to use with DistributedDataParallel, via `torch.distributed.launch`)
- each process starts to compile the extension
- they all try to save the extension to the same file
- sometimes, the file gets corrupted because of that

Here is part of the stack trace of the error I obtain (the rest is just the same error repeated):
```
  1 0: Traceback (most recent call last):
  2 0:   File "/private/home/fmassa/github/detectron.pytorch/torch_detectron/training_faster_rcnn.py", line 190, in <module>
  3 0:     from lib.faster_rcnn import fasterrcnn_resnet18, C2ResNetFasterRCNN
  4 0:   File "/private/home/fmassa/github/detectron.pytorch/torch_detectron/lib/faster_rcnn.py", line 13, in <module>
  5 0:     from lib.layers import ROIAlign, ROIPool, FixedBatchNorm2d
  6 0:   File "/private/home/fmassa/github/detectron.pytorch/torch_detectron/lib/layers/__init__.py", line 28, in <module>
  7 0:     _C = _load_C_extensions()
  8 0:   File "/private/home/fmassa/github/detectron.pytorch/torch_detectron/lib/layers/__init__.py", line 26, in _load_C_extensions
  9 0:     '-gencode arch=compute_70,code=sm_70'])
 10 0:   File "/private/home/fmassa/.conda/envs/detectron_v2/lib/python3.6/site-packages/torch/utils/cpp_extension.py", line 330, in load
 11 0:     build_directory = _get_build_directory(name, verbose)
 12 0:   File "/private/home/fmassa/.conda/envs/detectron_v2/lib/python3.6/site-packages/torch/utils/cpp_extension.py", line 392, in _get_build_directory
 13 0:     os.makedirs(build_directory)
 14 0:   File "/private/home/fmassa/.conda/envs/detectron_v2/lib/python3.6/os.py", line 220, in makedirs
 15 0:     mkdir(name, mode)
 16 0: FileExistsError: [Errno 17] File exists: '/tmp/torch_extensions/detectron_modules'
 17 0: Traceback (most recent call last):
 18 0:   File "/private/home/fmassa/github/detectron.pytorch/torch_detectron/training_faster_rcnn.py", line 190, in <module>
 19 0:     from lib.faster_rcnn import fasterrcnn_resnet18, C2ResNetFasterRCNN
 20 0:   File "/private/home/fmassa/github/detectron.pytorch/torch_detectron/lib/faster_rcnn.py", line 13, in <module>
 21 0:     from lib.layers import ROIAlign, ROIPool, FixedBatchNorm2d
 22 0:   File "/private/home/fmassa/github/detectron.pytorch/torch_detectron/lib/layers/__init__.py", line 28, in <module>
 23 0:     _C = _load_C_extensions()
 24 0:   File "/private/home/fmassa/github/detectron.pytorch/torch_detectron/lib/layers/__init__.py", line 26, in _load_C_extensions
 25 0:     '-gencode arch=compute_70,code=sm_70'])
 26 0:   File "/private/home/fmassa/.conda/envs/detectron_v2/lib/python3.6/site-packages/torch/utils/cpp_extension.py", line 361, in load
 27 0:     return _import_module_from_library(name, build_directory)
 28 0:   File "/private/home/fmassa/.conda/envs/detectron_v2/lib/python3.6/site-packages/torch/utils/cpp_extension.py", line 414, in _import_module_from_library
 29 0:     return imp.load_module(module_name, file, path, description)
 30 0:   File "/private/home/fmassa/.conda/envs/detectron_v2/lib/python3.6/imp.py", line 243, in load_module
 31 0:     return load_dynamic(name, filename, file)
 32 0:   File "/private/home/fmassa/.conda/envs/detectron_v2/lib/python3.6/imp.py", line 343, in load_dynamic
 33 0:     return _load(spec)
 34 0: ImportError: /tmp/torch_extensions/detectron_modules/detectron_modules.so: file too short
...
```
We can see that two types of errors pop out: `FileExistsError` and `ImportError: modules.so: file too short`.

Let me know if you need a minimal reproducible example.

pytorch version `'0.4.0a0+b21e135'`


Discussion:
Here is a `mwe.py`  (note: you need to modify the path to `pytorch_dir`)

```python
import argparse

def compile_ext():
    import os.path
    from torch.utils.cpp_extension import load as load_ext
    pytorch_dir = '/private/home/fmassa/github/pytorch/test/cpp_extensions'
    source = [
        'extension.cpp',
    ]
    source = [os.path.join(pytorch_dir, s) for s in source]
    return load_ext('multiprocess_bug', source)


_C = compile_ext()

if __name__ == "__main__":
    parser = argparse.ArgumentParser(description='Bug!')
    parser.add_argument('--local_rank', default=0)

    args = parser.parse_args()
    print(args)
```

Run it with
```
python -m torch.distributed.launch --nproc_per_node=8 mwe.py
```

On the first run I tried, I got:
```
Traceback (most recent call last):
  File "mwe.py", line 16, in <module>
    _C = compile_ext()
  File "mwe.py", line 13, in compile_ext
    return load_ext('multiprocess_bug', source)
  File "/private/home/fmassa/.conda/envs/detectron_v2/lib/python3.6/site-packages/torch/utils/cpp_extension.py", line 361, in load
    return _import_module_from_library(name, build_directory)
  File "/private/home/fmassa/.conda/envs/detectron_v2/lib/python3.6/site-packages/torch/utils/cpp_extension.py", line 411, in _import_module_from_library
    file, path, description = imp.find_module(module_name, [path])
  File "/private/home/fmassa/.conda/envs/detectron_v2/lib/python3.6/imp.py", line 297, in find_module
    raise ImportError(_ERR_MSG.format(name), name=name)
ImportError: No module named 'multiprocess_bug'
Traceback (most recent call last):
  File "/private/home/fmassa/.conda/envs/detectron_v2/lib/python3.6/site-packages/torch/utils/cpp_extension.py", line 400, in _build_extension_module
    ['ninja', '-v'], stderr=subprocess.STDOUT, cwd=build_directory)
  File "/private/home/fmassa/.conda/envs/detectron_v2/lib/python3.6/subprocess.py", line 336, in check_output
    **kwargs).stdout
  File "/private/home/fmassa/.conda/envs/detectron_v2/lib/python3.6/subprocess.py", line 418, in run
    output=stdout, stderr=stderr)
subprocess.CalledProcessError: Command '['ninja', '-v']' returned non-zero exit status 1.

During handling of the above exception, another exception occurred:

Traceback (most recent call last):
  File "mwe.py", line 16, in <module>
    _C = compile_ext()
  File "mwe.py", line 13, in compile_ext
    return load_ext('multiprocess_bug', source)
  File "/private/home/fmassa/.conda/envs/detectron_v2/lib/python3.6/site-packages/torch/utils/cpp_extension.py", line 357, in load
    _build_extension_module(name, build_directory)
  File "/private/home/fmassa/.conda/envs/detectron_v2/lib/python3.6/site-packages/torch/utils/cpp_extension.py", line 406, in _build_extension_module
    name, error.output.decode()))
RuntimeError: Error building extension 'multiprocess_bug': [1/2] c++ -MMD -MF extension.o.d -DTORCH_EXTENSION_NAME=multiprocess_bug -I/private/home/fmassa/.conda/envs/detectron_v2/lib/python3.6/site-packages/torch/lib/include -I/private/home/fmassa/.conda/envs/detectron_v2/lib/python3.6/site-packages/torch/lib/include/TH -I/private/home/fmassa/.conda/envs/detectron_v2/lib/python3.6/site-packages/torch/lib/include/THC -I/private/home/fmassa/.conda/envs/detectron_v2/include/python3.6m -fPIC -std=c++11 -c /private/home/fmassa/github/pytorch/test/cpp_extensions/extension.cpp -o extension.o
[2/2] c++ -shared extension.o -o multiprocess_bug.so
FAILED: multiprocess_bug.so
c++ -shared extension.o -o multiprocess_bug.so
extension.o: file not recognized: File truncated
collect2: error: ld returned 1 exit status
ninja: build stopped: subcommand failed.

Traceback (most recent call last):
  File "mwe.py", line 16, in <module>
    _C = compile_ext()
  File "mwe.py", line 13, in compile_ext
    return load_ext('multiprocess_bug', source)
  File "/private/home/fmassa/.conda/envs/detectron_v2/lib/python3.6/site-packages/torch/utils/cpp_extension.py", line 361, in load
    return _import_module_from_library(name, build_directory)
  File "/private/home/fmassa/.conda/envs/detectron_v2/lib/python3.6/site-packages/torch/utils/cpp_extension.py", line 411, in _import_module_from_library
    file, path, description = imp.find_module(module_name, [path])
  File "/private/home/fmassa/.conda/envs/detectron_v2/lib/python3.6/imp.py", line 297, in find_module
    raise ImportError(_ERR_MSG.format(name), name=name)
ImportError: No module named 'multiprocess_bug'
Traceback (most recent call last):
  File "mwe.py", line 16, in <module>
    _C = compile_ext()
  File "mwe.py", line 13, in compile_ext
    return load_ext('multiprocess_bug', source)
  File "/private/home/fmassa/.conda/envs/detectron_v2/lib/python3.6/site-packages/torch/utils/cpp_extension.py", line 361, in load
    return _import_module_from_library(name, build_directory)
  File "/private/home/fmassa/.conda/envs/detectron_v2/lib/python3.6/site-packages/torch/utils/cpp_extension.py", line 411, in _import_module_from_library
    file, path, description = imp.find_module(module_name, [path])
  File "/private/home/fmassa/.conda/envs/detectron_v2/lib/python3.6/imp.py", line 297, in find_module
    raise ImportError(_ERR_MSG.format(name), name=name)
ImportError: No module named 'multiprocess_bug'
```

Resolving patch:
diff --git a/torch/utils/cpp_extension.py b/torch/utils/cpp_extension.py
--- a/torch/utils/cpp_extension.py
+++ b/torch/utils/cpp_extension.py
@@ -11,6 +11,7 @@
 import warnings
 
 import torch
+from .file_baton import FileBaton
 
 from setuptools.command.build_ext import build_ext
 
@@ -30,7 +31,8 @@ def _find_cuda_home():
             # Guess #3
             try:
                 which = 'where' if sys.platform == 'win32' else 'which'
-                nvcc = subprocess.check_output([which, 'nvcc']).decode().rstrip('\r\n')
+                nvcc = subprocess.check_output(
+                    [which, 'nvcc']).decode().rstrip('\r\n')
                 cuda_home = os.path.dirname(os.path.dirname(nvcc))
             except Exception:
                 cuda_home = None
@@ -71,7 +73,8 @@ def check_compiler_abi_compatibility(compiler):
     '''
     try:
         check_cmd = '{}' if sys.platform == 'win32' else '{} --version'
-        info = subprocess.check_output(check_cmd.format(compiler).split(), stderr=subprocess.STDOUT)
+        info = subprocess.check_output(
+            check_cmd.format(compiler).split(), stderr=subprocess.STDOUT)
     except Exception:
         _, error, _ = sys.exc_info()
         warnings.warn('Error checking compiler version: {}'.format(error))
@@ -93,7 +96,8 @@ def check_compiler_abi_compatibility(compiler):
             version = re.search(r'(\d+)\.(\d+)\.(\d+)', info)
             if version is not None:
                 major, minor, revision = version.groups()
-                if (int(major), int(minor), int(revision)) >= MINIMUM_MSVC_VERSION:
+                if (int(major), int(minor),
+                        int(revision)) >= MINIMUM_MSVC_VERSION:
                     return True
                 else:
                     # Append the detected version for the warning.
@@ -156,9 +160,14 @@ def unix_wrap_compile(obj, src, ext, cc_args, extra_postargs, pp_opts):
                 # Put the original compiler back in place.
                 self.compiler.set_executable('compiler_so', original_compiler)
 
-        def win_wrap_compile(sources, output_dir=None, macros=None,
-                             include_dirs=None, debug=0, extra_preargs=None,
-                             extra_postargs=None, depends=None):
+        def win_wrap_compile(sources,
+                             output_dir=None,
+                             macros=None,
+                             include_dirs=None,
+                             debug=0,
+                             extra_preargs=None,
+                             extra_postargs=None,
+                             depends=None):
 
             self.cflags = copy.deepcopy(extra_postargs)
             extra_postargs = None
@@ -168,16 +177,22 @@ def spawn(cmd):
                 # Using regex to match src, obj and include files
 
                 src_regex = re.compile('/T(p|c)(.*)')
-                src_list = [m.group(2) for m in (
-                    src_regex.match(elem) for elem in cmd) if m]
+                src_list = [
+                    m.group(2) for m in (src_regex.match(elem) for elem in cmd)
+                    if m
+                ]
 
                 obj_regex = re.compile('/Fo(.*)')
-                obj_list = [m.group(1) for m in (
-                    obj_regex.match(elem) for elem in cmd) if m]
+                obj_list = [
+                    m.group(1) for m in (obj_regex.match(elem) for elem in cmd)
+                    if m
+                ]
 
                 include_regex = re.compile(r'((\-|\/)I.*)')
-                include_list = [m.group(1) for m in (
-                    include_regex.match(elem) for elem in cmd) if m]
+                include_list = [
+                    m.group(1)
+                    for m in (include_regex.match(elem) for elem in cmd) if m
+                ]
 
                 if len(src_list) >= 1 and len(obj_list) >= 1:
                     src = src_list[0]
@@ -190,8 +205,10 @@ def spawn(cmd):
                             cflags = self.cflags
                         else:
                             cflags = []
-                        cmd = [nvcc, '-c', src, '-o', obj, '-Xcompiler',
-                               '/wd4819', '-Xcompiler', '/MD'] + include_list + cflags
+                        cmd = [
+                            nvcc, '-c', src, '-o', obj, '-Xcompiler',
+                            '/wd4819', '-Xcompiler', '/MD'
+                        ] + include_list + cflags
                     elif isinstance(self.cflags, dict):
                         cflags = self.cflags['cxx']
                         cmd += cflags
@@ -203,9 +220,9 @@ def spawn(cmd):
 
             try:
                 self.compiler.spawn = spawn
-                return original_compile(sources,
-                                        output_dir, macros, include_dirs, debug,
-                                        extra_preargs, extra_postargs, depends)
+                return original_compile(sources, output_dir, macros,
+                                        include_dirs, debug, extra_preargs,
+                                        extra_postargs, depends)
             finally:
                 self.compiler.spawn = original_spawn
 
@@ -454,7 +471,57 @@ def load(name,
     if build_directory is None:
         build_directory = _get_build_directory(name, verbose)
 
-    extra_ldflags = extra_ldflags or []
+    baton = FileBaton(os.path.join(build_directory, 'lock'))
+
+    if baton.try_acquire():
+        try:
+            with_cuda = any(map(_is_cuda_file, sources))
+            extra_ldflags = _prepare_ldflags(
+                extra_ldflags or [],
+                with_cuda,
+                verbose)
+            build_file_path = os.path.join(build_directory, 'build.ninja')
+            if verbose:
+                print(
+                    'Emitting ninja build file {}...'.format(build_file_path))
+            # NOTE: Emitting a new ninja build file does not cause re-compilation if
+            # the sources did not change, so it's ok to re-emit (and it's fast).
+            _write_ninja_file(
+                path=build_file_path,
+                name=name,
+                sources=sources,
+                extra_cflags=extra_cflags or [],
+                extra_cuda_cflags=extra_cuda_cflags or [],
+                extra_ldflags=extra_ldflags or [],
+                extra_include_paths=extra_include_paths or [],
+                with_cuda=with_cuda)
+
+            if verbose:
+                print('Building extension module {}...'.format(name))
+            _build_extension_module(name, build_directory)
+        finally:
+            baton.release()
+    else:
+        baton.wait()
+
+    if verbose:
+        print('Loading extension module {}...'.format(name))
+    return _import_module_from_library(name, build_directory)
+
+
+def verify_ninja_availability():
+    '''
+    Returns ``True`` if the `ninja <https://ninja-build.org/>`_ build system is
+    available on the system.
+    '''
+    with open(os.devnull, 'wb') as devnull:
+        try:
+            subprocess.check_call('ninja --version'.split(), stdout=devnull)
+        except OSError:
+            raise RuntimeError("Ninja is required to load C++ extensions")
+
+
+def _prepare_ldflags(extra_ldflags, with_cuda, verbose):
     if sys.platform == 'win32':
         python_path = os.path.dirname(sys.executable)
         python_lib_path = os.path.join(python_path, 'libs')
@@ -468,51 +535,18 @@ def load(name,
         extra_ldflags.append('/LIBPATH:{}'.format(python_lib_path))
         extra_ldflags.append('/LIBPATH:{}'.format(lib_path))
 
-    with_cuda = any(map(_is_cuda_file, sources))
     if with_cuda:
         if verbose:
             print('Detected CUDA files, patching ldflags')
         if sys.platform == 'win32':
-            extra_ldflags.append('/LIBPATH:{}'.format(_join_cuda_home('lib/x64')))
+            extra_ldflags.append('/LIBPATH:{}'.format(
+                _join_cuda_home('lib/x64')))
             extra_ldflags.append('cudart.lib')
         else:
             extra_ldflags.append('-L{}'.format(_join_cuda_home('lib64')))
             extra_ldflags.append('-lcudart')
 
-    build_file_path = os.path.join(build_directory, 'build.ninja')
-    if verbose:
-        print('Emitting ninja build file {}...'.format(build_file_path))
-    # NOTE: Emitting a new ninja build file does not cause re-compilation if
-    # the sources did not change, so it's ok to re-emit (and it's fast).
-    _write_ninja_file(
-        path=build_file_path,
-        name=name,
-        sources=sources,
-        extra_cflags=extra_cflags or [],
-        extra_cuda_cflags=extra_cuda_cflags or [],
-        extra_ldflags=extra_ldflags or [],
-        extra_include_paths=extra_include_paths or [],
-        with_cuda=with_cuda)
-
-    if verbose:
-        print('Building extension module {}...'.format(name))
-    _build_extension_module(name, build_directory)
-
-    if verbose:
-        print('Loading extension module {}...'.format(name))
-    return _import_module_from_library(name, build_directory)
-
-
-def verify_ninja_availability():
-    '''
-    Returns ``True`` if the `ninja <https://ninja-build.org/>`_ build system is
-    available on the system.
-    '''
-    with open(os.devnull, 'wb') as devnull:
-        try:
-            subprocess.check_call('ninja --version'.split(), stdout=devnull)
-        except OSError:
-            raise RuntimeError("Ninja is required to load C++ extensions")
+    return extra_ldflags
 
 
 def _get_build_directory(name, verbose):
@@ -631,12 +665,15 @@ def _write_ninja_file(path,
 
     link_rule = ['rule link']
     if sys.platform == 'win32':
-        cl_paths = subprocess.check_output(['where', 'cl']).decode().split('\r\n')
+        cl_paths = subprocess.check_output(['where',
+                                            'cl']).decode().split('\r\n')
         if len(cl_paths) >= 1:
             cl_path = os.path.dirname(cl_paths[0]).replace(':', '$:')
         else:
             raise RuntimeError("MSVC is required to load C++ extensions")
-        link_rule.append('  command = "{}/link.exe" $in /nologo $ldflags /out:$out'.format(cl_path))
+        link_rule.append(
+            '  command = "{}/link.exe" $in /nologo $ldflags /out:$out'.format(
+                cl_path))
     else:
         link_rule.append('  command = $cxx $ldflags $in -o $out')
 
diff --git a/torch/utils/file_baton.py b/torch/utils/file_baton.py
new file mode 100644
--- /dev/null
+++ b/torch/utils/file_baton.py
@@ -0,0 +1,47 @@
+import os
+import time
+
+
+class FileBaton:
+    '''A primitive, file-based synchronization utility.'''
+
+    def __init__(self, lock_file_path, wait_seconds=0.1):
+        '''
+        Creates a new :class:`FileBaton`.
+
+        Args:
+            lock_file_path: The path to the file used for locking.
+            wait_seconds: The seconds to periorically sleep (spin) when
+                calling ``wait()``.
+        '''
+        self.lock_file_path = lock_file_path
+        self.wait_seconds = wait_seconds
+        self.fd = None
+
+    def try_acquire(self):
+        '''
+        Tries to atomically create a file under exclusive access.
+
+        Returns:
+            True if the file could be created, else False.
+        '''
+        try:
+            self.fd = os.open(self.lock_file_path, os.O_CREAT | os.O_EXCL)
+            return True
+        except FileExistsError:
+            return False
+
+    def wait(self):
+        '''
+        Periodically sleeps for a certain amount until the baton is released.
+
+        The amount of time slept depends on the ``wait_seconds`` parameter
+        passed to the constructor.
+        '''
+        while os.path.exists(self.lock_file_path):
+            time.sleep(self.wait_seconds)
+
+    def release(self):
+        '''Releaes the baton and removes its file.'''
+        os.close(self.fd)
+        os.remove(self.lock_file_path)


Test patch: