Apply code suggestions from review

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!

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

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: