diff --git a/pyop2/caching.py b/pyop2/caching.py index 96c64de75..107f1e445 100644 --- a/pyop2/caching.py +++ b/pyop2/caching.py @@ -43,7 +43,8 @@ from warnings import warn # noqa F401 from collections import defaultdict from itertools import count -from functools import partial, wraps +from functools import wraps +from tempfile import mkstemp from pyop2.configuration import configuration from pyop2.exceptions import CachingError, HashError # noqa: F401 @@ -319,10 +320,21 @@ def __setitem__(self, key, value): basedir = Path(self.cachedir, k1) basedir.mkdir(parents=True, exist_ok=True) - tempfile = basedir.joinpath(f"{k2}_p{os.getpid()}.tmp") - filepath = basedir.joinpath(k2) + # Care must be taken here to ensure that the file is created safely as + # the filesystem may be network based. `mkstemp` does so securely without + # race conditions: + # https://docs.python.org/3/library/tempfile.html#tempfile.mkstemp + _, tempfile = mkstemp(suffix=".tmp", prefix=k2, dir=basedir, text=False) + tempfile = Path(tempfile) + # Open using `tempfile` (the filename) rather than the file descriptor + # to allow redefining `self.open` with self.open(tempfile, mode="wb") as fh: self.write(fh, value) + + # Renaming (moving) the file is guaranteed by any POSIX compliant + # filesystem to be atomic. This may fail if somehow the destination is + # on another filesystem, but that shouldn't happen here. + filepath = basedir.joinpath(k2) tempfile.rename(filepath.with_suffix(self.extension)) def __delitem__(self, key): @@ -417,7 +429,9 @@ class DEFAULT_CACHE(dict): # Example of how to instrument and use different default caches: -EXOTIC_CACHE = partial(instrument(cachetools.LRUCache), maxsize=100) +# from functools import partial +# EXOTIC_CACHE = partial(instrument(cachetools.LRUCache), maxsize=100) + # Turn on cache measurements if printing cache info is enabled if configuration["print_cache_info"] or _running_on_ci: DEFAULT_CACHE = instrument(DEFAULT_CACHE) @@ -562,14 +576,3 @@ def memory_and_disk_cache(*args, cachedir=configuration["cache_dir"], **kwargs): def decorator(func): return memory_cache(*args, **kwargs)(disk_only_cache(*args, cachedir=cachedir, **kwargs)(func)) return decorator - -# JBTODO: (Wishlist) -# * Try more exotic caches ie: memory_cache = partial(parallel_cache, cache_factory=lambda: cachetools.LRUCache(maxsize=1000)) ✓ -# * Add some sort of cache reporting ✓ -# * Add some sort of cache statistics ✓ -# * Refactor compilation.py to use @mem_and_disk_cached, where get_so is just uses DictLikeDiskAccess with an overloaded self.write() method ✓ -# * Systematic investigation into cache sizes/types for Firedrake -# - Is a mem cache needed for DLLs? ~~No~~ Yes!! -# - Is LRUCache better than a simple dict? (memory profile test suite) No -# - What is the optimal maxsize? ∞ -# * Add some docstrings and maybe some exposition! diff --git a/pyop2/compilation.py b/pyop2/compilation.py index 9fb654d9c..08e35c14a 100644 --- a/pyop2/compilation.py +++ b/pyop2/compilation.py @@ -46,8 +46,8 @@ from functools import partial from pathlib import Path from contextlib import contextmanager -from tempfile import gettempdir -from uuid import uuid4 +from tempfile import gettempdir, mkstemp +from random import randint from pyop2 import mpi @@ -575,17 +575,6 @@ def make_so(compiler, jitmodule, extension, comm, filename=None): :arg filename: Optional Returns a :class:`ctypes.CDLL` object of the resulting shared library.""" - if filename is None: - # A UUID should ensure we have a unique path - uuid = uuid4().hex - # Taking the first two characters avoids using excessive filesystem inodes - tempdir = MEM_TMP_DIR.joinpath(f"{uuid[:2]}") - # This path + filename should be unique - filename = tempdir.joinpath(f"{uuid[2:]}.{extension}") - else: - tempdir = None - filename = Path(filename).absolute() - # Compilation communicators are reference counted on the PyOP2 comm icomm = mpi.internal_comm(comm, compiler) ccomm = mpi.compilation_comm(icomm, compiler) @@ -598,39 +587,40 @@ def make_so(compiler, jitmodule, extension, comm, filename=None): exe = compiler.cc compiler_flags = compiler.cflags - # TODO: Do we still need to worry about atomic file renaming in this function? - base = filename.stem - path = filename.parent - pid = os.getpid() - cname = filename.with_name(f"{base}_p{pid}.{extension}") - oname = filename.with_name(f"{base}_p{pid}.o") - # Link into temporary file, then rename to shared library atomically (avoiding races). - tempname = filename.with_stem(f"{base}_p{pid}.so") - soname = filename.with_suffix(".so") - # Compile on compilation communicator (ccomm) rank 0 + soname = None if ccomm.rank == 0: - if tempdir is None: - filename.parent.mkdir(exist_ok=True) - else: + if filename is None: + # Adding random 2-digit hexnum avoids using excessive filesystem inodes + tempdir = MEM_TMP_DIR.joinpath(f"{randint(0, 255):02x}") tempdir.mkdir(parents=True, exist_ok=True) - logfile = path.joinpath(f"{base}_p{pid}.log") - errfile = path.joinpath(f"{base}_p{pid}.err") + # This path + filename should be unique + _, filename = mkstemp(suffix=f".{extension}", dir=tempdir, text=True) + filename = Path(filename) + else: + filename.parent.mkdir(exist_ok=True) + + cname = filename + oname = filename.with_suffix(".o") + soname = filename.with_suffix(".so") + logfile = filename.with_suffix(".log") + errfile = filename.with_suffix(".err") with progress(INFO, 'Compiling wrapper'): + # Write source code to disk with open(cname, "w") as fh: fh.write(jitmodule.code_to_compile) - # Compiler also links + if not compiler.ld: - cc = (exe,) + compiler_flags + ('-o', str(tempname), str(cname)) + compiler.ldflags + # Compile and link + cc = (exe,) + compiler_flags + ('-o', str(soname), str(cname)) + compiler.ldflags _run(cc, logfile, errfile) else: + # Compile cc = (exe,) + compiler_flags + ('-c', '-o', oname, cname) _run(cc, logfile, errfile) - # Extract linker specific "cflags" from ldflags - ld = tuple(shlex.split(compiler.ld)) + ('-o', str(tempname), str(oname)) + tuple(expandWl(compiler.ldflags)) + # Extract linker specific "cflags" from ldflags and link + ld = tuple(shlex.split(compiler.ld)) + ('-o', str(soname), str(oname)) + tuple(expandWl(compiler.ldflags)) _run(ld, logfile, errfile, step="Linker", filemode="a") - # Atomically ensure soname exists - tempname.rename(soname) return ccomm.bcast(soname, root=0) @@ -664,10 +654,10 @@ def _run(cc, logfile, errfile, step="Compilation", filemode="w"): def _add_profiling_events(dll, events): """ - If PyOP2 is in profiling mode, events are attached to dll to profile the local linear algebra calls. - The event is generated here in python and then set in the shared library, - so that memory is not allocated over and over again in the C kernel. The naming - convention is that the event ids are named by the event name prefixed by "ID_". + If PyOP2 is in profiling mode, events are attached to dll to profile the local linear algebra calls. + The event is generated here in python and then set in the shared library, + so that memory is not allocated over and over again in the C kernel. The naming + convention is that the event ids are named by the event name prefixed by "ID_". """ if PETSc.Log.isActive(): # also link the events from the linear algebra callables diff --git a/pyop2/global_kernel.py b/pyop2/global_kernel.py index 7e313a5e8..ae13dc1c5 100644 --- a/pyop2/global_kernel.py +++ b/pyop2/global_kernel.py @@ -367,7 +367,6 @@ def code_to_compile(self): from pyop2.codegen.rep2loopy import generate wrapper = generate(self.builder) - # JBTODO: Expensive? Can this be wrapped with a cache? code = lp.generate_code_v2(wrapper) if self.local_kernel.cpp: