Merge pull request #3071 from devitocodes/fix-complex-header-race

compiler: Write the complex arithmetic header atomically

Author: ggorman

devito/arch/compiler.py

@@ -22,7 +22,8 @@
 from devito.logger import debug, warning
 from devito.parameters import configuration
 from devito.tools import (
-    Signer, as_list, change_directory, filter_ordered, make_tempdir, memoized_func
+    Signer, as_list, atomic_write, change_directory, filter_ordered, make_tempdir,
+    memoized_func
 )
 
 __all__ = ['compiler_registry', 'sniff_mpi_distro']
@@ -316,10 +317,17 @@ def save_header(self, filename, code):
             The name of the header file (w/o the suffix).
         code : str
             The source code to be stored.
+
+        Returns
+        -------
+        Path
+            The path of the header file.
         """
         hfile = self.get_jit_dir().joinpath(filename).with_suffix('.h')
-        with open(str(hfile), 'w') as f:
-            f.write(code)
+        # The JIT dir is shared with other processes, e.g. the workers of a
+        # parallel test run, so the header is written atomically: a concurrent
+        # compiler never reads a truncated or half-written header
+        return atomic_write(hfile, code)
 
     def save(self, soname, binary):
         """

devito/passes/iet/dtypes.py

@@ -28,10 +28,11 @@ def _complex_includes(iet: Callable, langbb: type[LangBB], compiler: Compiler,
 
     # Some languages such as c++11 need some extra arithmetic definitions
     if langbb.get('def-complex'):
-        dest = compiler.get_jit_dir()
-        hfile = dest.joinpath('complex_arith.h')
-        with open(str(hfile), 'w') as ff:
-            ff.write(str(langbb['def-complex']))
+        # The definitions differ across languages (e.g., CUDA vs C++), hence
+        # the language-specific header name
+        langname = langbb['name'].lower()
+        hfile = compiler.save_header(f'{langname}_complex_arith',
+                                     str(langbb['def-complex']))
         lib += (str(hfile),)
 
     metadata['includes'] = lib

devito/tools/os_helper.py

@@ -1,8 +1,9 @@
 import os
+from contextlib import suppress
 from pathlib import Path
-from tempfile import gettempdir
+from tempfile import gettempdir, mkstemp
 
-__all__ = ['change_directory', 'make_tempdir']
+__all__ = ['atomic_write', 'change_directory', 'make_tempdir']
 
 
 class change_directory:
@@ -34,3 +35,32 @@ def make_tempdir(prefix=None):
     tmpdir = Path(gettempdir()).joinpath(name)
     tmpdir.mkdir(parents=True, exist_ok=True)
     return tmpdir
+
+
+def atomic_write(path, data, mode='w'):
+    """
+    Write `data` to `path` through a private temporary file in the same
+    directory, then rename it into place. The rename is atomic, so a concurrent
+    reader never sees a partial file: it reads either the previous content or
+    the new one, and a reader that already opened the previous file keeps
+    reading that one in full. The temporary file is removed if the write fails.
+
+    Note: the resulting file has the restrictive mode of the temporary file
+    (0600, further restricted by the process umask).
+    """
+    path = Path(path)
+    fd, tmpfile = mkstemp(dir=path.parent, prefix=f'{path.name}.', suffix='.tmp')
+    try:
+        try:
+            f = os.fdopen(fd, mode)
+        except BaseException:
+            os.close(fd)
+            raise
+        with f:
+            f.write(data)
+        os.replace(tmpfile, path)
+    except BaseException:
+        with suppress(FileNotFoundError):
+            os.unlink(tmpfile)
+        raise
+    return path

tests/test_arch.py

@@ -1,3 +1,6 @@
+from concurrent.futures import ThreadPoolExecutor
+from contextlib import suppress
+
 import pytest
 
 from devito import configuration, switchconfig
@@ -29,3 +32,37 @@ def test_switcharch():
     assert isinstance(tmp_comp, old_compiler.__class__)
     assert old_compiler.suffix == tmp_comp.suffix
     assert old_compiler.name == tmp_comp.name
+
+
+def test_save_header_concurrent(tmp_path, monkeypatch):
+    """
+    The JIT dir is shared across processes, so a header being rewritten must
+    never be seen truncated or half-written by a concurrent reader.
+    """
+    compiler = configuration['compiler']
+    monkeypatch.setattr(compiler, 'get_jit_dir', lambda: tmp_path)
+
+    hfile = tmp_path.joinpath('header.h')
+    codes = ['a'*2**20, 'b'*2**20]
+
+    def write(code):
+        for _ in range(20):
+            compiler.save_header('header', code)
+
+    def read(writers):
+        seen = set()
+        while not all(w.done() for w in writers):
+            with suppress(FileNotFoundError):
+                seen.add(hfile.read_text())
+        return seen
+
+    with ThreadPoolExecutor(max_workers=8) as tpe:
+        writers = [tpe.submit(write, codes[i % 2]) for i in range(4)]
+        readers = [tpe.submit(read, writers) for _ in range(4)]
+        for w in writers:
+            w.result()
+        seen = set.union(*[r.result() for r in readers])
+
+    # Only complete headers are ever read, and no temporary file is left behind
+    assert seen and seen <= set(codes)
+    assert list(tmp_path.iterdir()) == [hfile]

tests/test_dtypes.py

@@ -1,3 +1,5 @@
+from pathlib import Path
+
 import numpy as np
 import pytest
 import sympy
@@ -13,8 +15,9 @@
 from devito.ir.cgen.printer import BasePrinter, get_printer
 from devito.passes.iet.langbase import LangBB
 from devito.passes.iet.languages.C import CBB, CPrinter
+from devito.passes.iet.languages.CXX import CXXPrinter
 from devito.passes.iet.languages.openacc import AccBB, AccPrinter
-from devito.passes.iet.languages.openmp import OmpBB
+from devito.passes.iet.languages.openmp import CXXOmpBB, OmpBB
 from devito.symbolics.extended_dtypes import ctypes_vector_mapper
 from devito.tools import dtype_to_cstr
 from devito.types.basic import Basic, Scalar, Symbol
@@ -25,12 +28,14 @@
 _languages: dict[str, type[LangBB]] = {
     'C': CBB,
     'openmp': OmpBB,
+    'CXXopenmp': CXXOmpBB,
     'openacc': AccBB
 }
 
 _printers: dict[str, type[BasePrinter]] = {
     'C': CPrinter,
     'openmp': CPrinter,
+    'CXXopenmp': CXXPrinter,
     'openacc': AccPrinter
 }
 
@@ -147,6 +152,30 @@ def test_complex_headers(dtype: np.dtype[np.inexact], kwargs: dict[str, str]) ->
         assert header not in op._includes
 
 
+@pytest.mark.parametrize('kwargs', _configs + [_config_kwargs('cpu64', 'CXXopenmp')],
+                         ids=kw_id)
+def test_complex_arith_header(kwargs: dict[str, str]) -> None:
+    """
+    Tests that the extra complex arithmetic definitions some languages need are
+    saved in full to a language-specific header the operator includes.
+    """
+    langbb = _get_language(**kwargs)
+    code = langbb.get('def-complex')
+    if code is None:
+        pytest.skip(f"`{kwargs['language']}` needs no complex arithmetic header")
+
+    # Set up an operator
+    grid = Grid(shape=(3, 3))
+    x, y = grid.dimensions
+
+    u = Function(name='u', grid=grid, dtype=np.complex64)
+    op = Operator(Eq(u, x*y + 1j), **kwargs)
+
+    hfile, = [Path(i) for i in op._includes if 'complex_arith' in i]
+    assert hfile.name == f"{langbb['name'].lower()}_complex_arith.h"
+    assert hfile.read_text() == str(code)
+
+
 @pytest.mark.parametrize('dtype', [np.complex64, np.complex128])
 @pytest.mark.parametrize('kwargs', _configs, ids=kw_id)
 def test_imag_unit(dtype: np.complexfloating, kwargs: dict[str, str]) -> None:

tests/test_tools.py

@@ -7,8 +7,8 @@
 
 from devito import Eq, Operator, switchenv
 from devito.tools import (
-    CacheInstances, DefaultFrozenDict, UnboundedMultiTuple, UnboundTuple, ctypes_to_cstr,
-    filter_ordered, memoized_meth, toposort, transitive_closure
+    CacheInstances, DefaultFrozenDict, UnboundedMultiTuple, UnboundTuple, atomic_write,
+    ctypes_to_cstr, filter_ordered, memoized_meth, toposort, transitive_closure
 )
 from devito.types.basic import Symbol
 
@@ -338,3 +338,16 @@ def foo():
     finally:
         os.environ.clear()
         os.environ.update(previous_environ)
+
+
+def test_atomic_write(tmp_path):
+    path = tmp_path.joinpath('data.txt')
+
+    assert atomic_write(path, 'foo') == path
+    assert path.read_text() == 'foo'
+
+    # A failed write leaves the previous file untouched and no temporary file behind
+    with pytest.raises(TypeError):
+        atomic_write(path, 1)
+    assert path.read_text() == 'foo'
+    assert list(tmp_path.iterdir()) == [path]