fix: patching to match mkl_fft

Author: jharlow-intel

mkl_random/src/_patch.pyx

@@ -35,7 +35,7 @@ replace NumPy's `Generator`/`default_rng()` unless mkl_random provides fully
 compatible replacements.
 """
 
-from threading import local as threading_local
+from threading import Lock, local
 from contextlib import ContextDecorator
 
 import numpy as _np
@@ -89,51 +89,30 @@ cdef tuple _DEFAULT_NAMES = (
 )
 
 
-cdef class patch:
-    cdef bint _is_patched
-    cdef object _numpy_module
-    cdef object _originals   # dict: name -> original object
-    cdef object _patched     # list of names actually patched
-
-    def __cinit__(self):
-        self._is_patched = False
+class _GlobalPatch:
+    def __init__(self):
+        self._lock = Lock()
+        self._patch_count = 0
         self._numpy_module = None
+        self._requested_names = None
         self._originals = {}
-        self._patched = []
-
-    def do_patch(self, numpy_module=None, names=None, bint strict=False):
-        """
-        Patch the given numpy module (default: imported numpy) in-place.
-
-        Parameters
-        ----------
-        numpy_module : module, optional
-            The numpy module to patch (e.g. `import numpy as np; use_in_numpy(np)`).
-        names : iterable[str], optional
-            Attributes under `numpy_module.random` to patch. Defaults to _DEFAULT_NAMES.
-        strict : bool
-            If True, raise if any requested symbol cannot be patched.
-        """
-        if numpy_module is None:
-            numpy_module = _np
+        self._patched = ()
+        self._tls = local()
+
+    def _normalize_names(self, names):
         if names is None:
             names = _DEFAULT_NAMES
+        return tuple(names)
 
+    def _validate_module(self, numpy_module):
         if not hasattr(numpy_module, "random"):
             raise TypeError("Expected a numpy-like module with a `.random` attribute.")
 
-        # If already patched, only allow idempotent re-entry for the same numpy module.
-        if self._is_patched:
-            if self._numpy_module is numpy_module:
-                return
-            raise RuntimeError("Already patched a different numpy module; call restore() first.")
-
+    def _apply_patch(self, numpy_module, names, strict):
         np_random = numpy_module.random
-
         originals = {}
         patched = []
         missing = []
-
         for name in names:
             if not hasattr(np_random, name) or not hasattr(_mr, name):
                 missing.append(name)
@@ -143,58 +122,75 @@ cdef class patch:
             patched.append(name)
 
         if strict and missing:
-            # revert partial patch before raising
-            for n, v in originals.items():
-                setattr(np_random, n, v)
+            for name, value in originals.items():
+                setattr(np_random, name, value)
             raise AttributeError(
                 "Could not patch these names (missing on numpy.random or mkl_random.mklrand): "
                 + ", ".join([str(x) for x in missing])
             )
 
         self._numpy_module = numpy_module
+        self._requested_names = names
         self._originals = originals
-        self._patched = patched
-        self._is_patched = True
-
-    def do_unpatch(self):
-        """
-        Restore the previously patched numpy module.
-        """
-        if not self._is_patched:
-            return
-        numpy_module = self._numpy_module
-        np_random = numpy_module.random
-        for n, v in self._originals.items():
-            setattr(np_random, n, v)
+        self._patched = tuple(patched)
 
-        self._numpy_module = None
-        self._originals = {}
-        self._patched = []
-        self._is_patched = False
+    def do_patch(self, numpy_module=None, names=None, strict=False, verbose=False):
+        if numpy_module is None:
+            numpy_module = _np
+        names = self._normalize_names(names)
+        self._validate_module(numpy_module)
+        strict = bool(strict)
+
+        with self._lock:
+            local_count = getattr(self._tls, "local_count", 0)
+            if self._patch_count == 0:
+                self._apply_patch(numpy_module, names, strict)
+            else:
+                if self._numpy_module is not numpy_module:
+                    raise RuntimeError(
+                        "Already patched a different numpy module; call restore() first."
+                    )
+                if names != self._requested_names:
+                    raise RuntimeError(
+                        "Already patched with a different names set; call restore() first."
+                    )
+            self._patch_count += 1
+            self._tls.local_count = local_count + 1
+
+    def do_restore(self, verbose=False):
+        with self._lock:
+            local_count = getattr(self._tls, "local_count", 0)
+            if local_count <= 0:
+                if verbose:
+                    print(
+                        "Warning: restore called more times than monkey_patch in this thread."
+                    )
+                return
+
+            self._tls.local_count = local_count - 1
+            self._patch_count -= 1
+            if self._patch_count == 0:
+                np_random = self._numpy_module.random
+                for name, value in self._originals.items():
+                    setattr(np_random, name, value)
+                self._numpy_module = None
+                self._requested_names = None
+                self._originals = {}
+                self._patched = ()
 
     def is_patched(self):
-        return self._is_patched
+        with self._lock:
+            return self._patch_count > 0
 
     def patched_names(self):
-        """
-        Returns list of names that were actually patched.
-        """
-        return list(self._patched)
-
-
-_tls = threading_local()
-
-
-def _is_tls_initialized():
-    return (getattr(_tls, "initialized", None) is not None) and (_tls.initialized is True)
+        with self._lock:
+            return list(self._patched)
 
 
-def _initialize_tls():
-    _tls.patch = patch()
-    _tls.initialized = True
+_patch = _GlobalPatch()
 
 
-def monkey_patch(numpy_module=None, names=None, strict=False):
+def monkey_patch(numpy_module=None, names=None, strict=False, verbose=False):
     """
     Enables using mkl_random in the given NumPy module by patching `numpy.random`.
 
@@ -211,43 +207,45 @@ def monkey_patch(numpy_module=None, names=None, strict=False):
     >>> mkl_random.is_patched()
     False
     """
-    if not _is_tls_initialized():
-        _initialize_tls()
-    _tls.patch.do_patch(numpy_module=numpy_module, names=names, strict=bool(strict))
+    _patch.do_patch(
+        numpy_module=numpy_module,
+        names=names,
+        strict=bool(strict),
+        verbose=bool(verbose),
+    )
 
 
-def use_in_numpy(numpy_module=None, names=None, strict=False):
+def use_in_numpy(numpy_module=None, names=None, strict=False, verbose=False):
     """
     Backward-compatible alias for monkey_patch().
     """
-    monkey_patch(numpy_module=numpy_module, names=names, strict=strict)
+    monkey_patch(
+        numpy_module=numpy_module,
+        names=names,
+        strict=strict,
+        verbose=verbose,
+    )
 
 
-def restore():
+def restore(verbose=False):
     """
     Disables using mkl_random in NumPy by restoring the original `numpy.random` symbols.
     """
-    if not _is_tls_initialized():
-        _initialize_tls()
-    _tls.patch.do_unpatch()
+    _patch.do_restore(verbose=bool(verbose))
 
 
 def is_patched():
     """
     Returns whether NumPy has been patched with mkl_random.
     """
-    if not _is_tls_initialized():
-        _initialize_tls()
-    return bool(_tls.patch.is_patched())
+    return _patch.is_patched()
 
 
 def patched_names():
     """
     Returns the names actually patched in `numpy.random`.
     """
-    if not _is_tls_initialized():
-        _initialize_tls()
-    return _tls.patch.patched_names()
+    return _patch.patched_names()
 
 
 class mkl_random(ContextDecorator):

mkl_random/tests/test_patch.py

@@ -93,3 +93,41 @@ def test_patched_names():
         assert "RandomState" in names
     finally:
         mkl_random.restore()
+
+
+def test_patch_redundant_patching():
+    orig_normal = np.random.normal
+    assert not mkl_random.is_patched()
+
+    mkl_random.monkey_patch(np)
+    mkl_random.monkey_patch(np)
+
+    assert mkl_random.is_patched()
+    assert np.random.normal is mkl_random.mklrand.normal
+
+    mkl_random.restore()
+    assert mkl_random.is_patched()
+    assert np.random.normal is mkl_random.mklrand.normal
+
+    mkl_random.restore()
+    assert not mkl_random.is_patched()
+    assert np.random.normal is orig_normal
+
+
+def test_patch_reentrant():
+    orig_uniform = np.random.uniform
+    assert not mkl_random.is_patched()
+
+    with mkl_random.mkl_random(np):
+        assert mkl_random.is_patched()
+        assert np.random.uniform is not orig_uniform
+
+        with mkl_random.mkl_random(np):
+            assert mkl_random.is_patched()
+            assert np.random.uniform is not orig_uniform
+
+        assert mkl_random.is_patched()
+        assert np.random.uniform is not orig_uniform
+
+    assert not mkl_random.is_patched()
+    assert np.random.uniform is orig_uniform