Repository: joblib/loky

Issue:
Initializer example fails
The example from the [Advanced Executor Setup](https://loky.readthedocs.io/en/stable/auto_examples/reuseable_executor.html#sphx-glr-auto-examples-reuseable-executor-py) docs fails with loky 2.5.1.

Output:
```
$ python loky_init.py
[LokyProcess-1] init
Traceback (most recent call last):
  File "loky_init.py", line 27, in <module>
    assert executor.submit(return_initializer_status).result() == 'initialized'
AssertionError
[LokyProcess-2] init
```

Any help would be appreciated.
```
$ python --version
Python 3.7.3
>>> loky.__version__
'2.5.1'
$ uname -a
Darwin dstuebe 18.7.0 Darwin Kernel Version 18.7.0: Tue Aug 20 16:57:14 PDT 2019; root:xnu-4903.271.2~2/RELEASE_X86_64 x86_64
```

Discussion:

Patch:
diff --git a/CHANGES.md b/CHANGES.md
index 39ba8a9..3515946 100644
--- a/CHANGES.md
+++ b/CHANGES.md
@@ -1,5 +1,13 @@
-### 2.6.1 - XXXX-YY-ZZ - Bugfix release
-
+### 2.7.0 - XXXX-YY-ZZ
+
+- The ``resource_tracker`` now comes with built-in support for tracking files
+  in all OSes.  In addition, Python processes can now signal they do not need a
+  shared resource anymore by using the
+  ``resource_tracker.maybe_unlink(resource_name, resource_type)`` method.  After
+  all processes having access to the said resource have called this method, the
+  ``resource_tracker`` will proceed to unlink the resource. Previously, resource
+  unlinking by the ``resource_tracker`` was only done for leaked resources at
+  interpreter exit (#228).
 
 ### 2.6.0 - 2019-09-18
 
diff --git a/examples/reuse_executor.py b/examples/reuse_executor.py
index cecc8fd..c8e4c9f 100644
--- a/examples/reuse_executor.py
+++ b/examples/reuse_executor.py
@@ -8,14 +8,14 @@ The factory ``get_reusable_executor`` provides an executor. As long as this
 executor is not in a broken state, it is reused for all the computation.
 """
 import os
-import argparse
 from loky import get_reusable_executor
 
 
 def func_async(i):
     import os
     pid = os.getpid()
-    return (2*i, pid)
+    return (2 * i, pid)
+
 
 def test_1():
     executor = get_reusable_executor(max_workers=1)
@@ -31,6 +31,7 @@ def test_3():
     executor = get_reusable_executor(max_workers=1)
     return executor.submit(func_async, 3)
 
+
 f1 = test_1()
 f2 = test_2()
 f3 = test_3()
@@ -38,7 +39,7 @@ f3 = test_3()
 main_pid = os.getpid()
 results = [f1.result(), f2.result(), f3.result()]
 
-pids = [v[1] for v in results]
+pids = [pid for _, pid in results]
 
 for i, (val, pid) in enumerate(results):
     assert val == 2 * (i + 1)
diff --git a/examples/reuseable_executor.py b/examples/reuseable_executor.py
index da53239..76599dd 100644
--- a/examples/reuseable_executor.py
+++ b/examples/reuseable_executor.py
@@ -8,29 +8,29 @@ arguments stay the same.
 """
 from time import sleep
 import multiprocessing as mp
+import loky
 from loky import get_reusable_executor
 
-INITIALIZER_STATUS = "uninitialized"
+# Store the initialization status in a global variable of a module.
+loky._INITIALIZER_STATUS = "uninitialized"
 
 
 def initializer(x):
-    global INITIALIZER_STATUS
     print("[{}] init".format(mp.current_process().name))
-    INITIALIZER_STATUS = x
+    loky._INITIALIZER_STATUS = x
 
 
 def return_initializer_status(delay=0):
     sleep(delay)
 
-    global INITIALIZER_STATUS
-    return INITIALIZER_STATUS
+    return getattr(loky, '_INITIALIZER_STATUS', 'uninitialized')
 
 
 executor = get_reusable_executor(
     max_workers=2, initializer=initializer, initargs=('initialized',),
     context="loky", timeout=1000)
 
-assert INITIALIZER_STATUS == "uninitialized"
+assert loky._INITIALIZER_STATUS == "uninitialized"
 executor.submit(return_initializer_status).result()
 assert executor.submit(return_initializer_status).result() == 'initialized'
 
diff --git a/loky/backend/resource_tracker.py b/loky/backend/resource_tracker.py
index ce519fe..1b0b276 100644
--- a/loky/backend/resource_tracker.py
+++ b/loky/backend/resource_tracker.py
@@ -13,15 +13,32 @@
 #
 # On Unix we run a server process which keeps track of unlinked
 # resources. The server ignores SIGINT and SIGTERM and reads from a
-# pipe.  Every other process of the program has a copy of the writable
-# end of the pipe, so we get EOF when all other processes have exited.
-# Then the server process unlinks any remaining resources.
-#
+# pipe. The resource_tracker implements a reference counting scheme: each time
+# a Python process anticipates the shared usage of a resource by another
+# process, it signals the resource_tracker of this shared usage, and in return,
+# the resource_tracker increments the resource's reference count by 1.
+# Similarly, when access to a resource is closed by a Python process, the
+# process notifies the resource_tracker by asking it to decrement the
+# resource's reference count by 1.  When the reference count drops to 0, the
+# resource_tracker attempts to clean up the underlying resource.
+
+# Finally, every other process connected to the resource tracker has a copy of
+# the writable end of the pipe used to communicate with it, so the resource
+# tracker gets EOF when all other processes have exited. Then the
+# resource_tracker process unlinks any remaining leaked resources (with
+# reference count above 0)
+
 # For semaphores, this is important because the system only supports a limited
 # number of named semaphores, and they will not be automatically removed till
 # the next reboot.  Without this resource tracker process, "killall python"
 # would probably leave unlinked semaphores.
 
+# Note that this behavior differs from CPython's resource_tracker, which only
+# implements list of shared resources, and not a proper refcounting scheme.
+# Also, CPython's resource tracker will only attempt to cleanup those shared
+# resources once all procsses connected to the resouce tracker have exited.
+
+
 import os
 import shutil
 import sys
@@ -52,7 +69,8 @@ _HAVE_SIGMASK = hasattr(signal, 'pthread_sigmask')
 _IGNORED_SIGNALS = (signal.SIGINT, signal.SIGTERM)
 
 _CLEANUP_FUNCS = {
-    'folder': shutil.rmtree
+    'folder': shutil.rmtree,
+    'file': os.unlink
 }
 
 if os.name == "posix":
@@ -167,7 +185,7 @@ class ResourceTracker(object):
             return True
 
     def register(self, name, rtype):
-        '''Register a named resource with resource tracker.'''
+        '''Register a named resource, and increment its refcount.'''
         self.ensure_running()
         self._send('REGISTER', name, rtype)
 
@@ -176,6 +194,11 @@ class ResourceTracker(object):
         self.ensure_running()
         self._send('UNREGISTER', name, rtype)
 
+    def maybe_unlink(self, name, rtype):
+        '''Decrement the refcount of a resource, and delete it if it hits 0'''
+        self.ensure_running()
+        self._send("MAYBE_UNLINK", name, rtype)
+
     def _send(self, cmd, name, rtype):
         msg = '{0}:{1}:{2}\n'.format(cmd, name, rtype).encode('ascii')
         if len(name) > 512:
@@ -189,6 +212,7 @@ class ResourceTracker(object):
 _resource_tracker = ResourceTracker()
 ensure_running = _resource_tracker.ensure_running
 register = _resource_tracker.register
+maybe_unlink = _resource_tracker.maybe_unlink
 unregister = _resource_tracker.unregister
 getfd = _resource_tracker.getfd
 
@@ -196,6 +220,9 @@ getfd = _resource_tracker.getfd
 def main(fd, verbose=0):
     '''Run resource tracker.'''
     # protect the process from ^C and "killall python" etc
+    if verbose:
+        util.log_to_stderr(level=util.DEBUG)
+
     signal.signal(signal.SIGINT, signal.SIG_IGN)
     signal.signal(signal.SIGTERM, signal.SIG_IGN)
 
@@ -208,17 +235,19 @@ def main(fd, verbose=0):
         except Exception:
             pass
 
-    if verbose:  # pragma: no cover
-        sys.stderr.write("Main resource tracker is running\n")
-        sys.stderr.flush()
+    if verbose:
+        util.debug("Main resource tracker is running")
 
-    cache = {rtype: set() for rtype in _CLEANUP_FUNCS.keys()}
+    registry = {rtype: dict() for rtype in _CLEANUP_FUNCS.keys()}
     try:
         # keep track of registered/unregistered resources
         if sys.platform == "win32":
             fd = msvcrt.open_osfhandle(fd, os.O_RDONLY)
         with open(fd, 'rb') as f:
-            for line in f:
+            while True:
+                line = f.readline()
+                if line == b'':  # EOF
+                    break
                 try:
                     splitted = line.strip().decode('ascii').split(':')
                     # name can potentially contain separator symbols (for
@@ -237,18 +266,42 @@ def main(fd, verbose=0):
                                 name, rtype, list(_CLEANUP_FUNCS.keys())))
 
                     if cmd == 'REGISTER':
-                        cache[rtype].add(name)
-                        if verbose:  # pragma: no cover
-                            sys.stderr.write("[ResourceTracker] register {}"
-                                             " {}\n" .format(rtype, name))
-                            sys.stderr.flush()
+                        if name not in registry[rtype]:
+                            registry[rtype][name] = 1
+                        else:
+                            registry[rtype][name] += 1
+
+                        if verbose:
+                            util.debug(
+                                "[ResourceTracker] incremented refcount of {} "
+                                "{} (current {})".format(
+                                    rtype, name, registry[rtype][name]))
                     elif cmd == 'UNREGISTER':
-                        cache[rtype].remove(name)
-                        if verbose:  # pragma: no cover
-                            sys.stderr.write("[ResourceTracker] unregister {}"
-                                             " {}: cache({})\n"
-                                             .format(name, rtype, len(cache)))
-                            sys.stderr.flush()
+                        del registry[rtype][name]
+                        if verbose:
+                            util.debug(
+                                "[ResourceTracker] unregister {} {}: "
+                                "registry({})".format(name, rtype, len(registry)))
+                    elif cmd == 'MAYBE_UNLINK':
+                        registry[rtype][name] -= 1
+                        if verbose:
+                            util.debug(
+                                "[ResourceTracker] decremented refcount of {} "
+                                "{} (current {})".format(
+                                    rtype, name, registry[rtype][name]))
+
+                        if registry[rtype][name] == 0:
+                            del registry[rtype][name]
+                            try:
+                                if verbose:
+                                    util.debug(
+                                            "[ResourceTracker] unlink {}"
+                                            .format(name))
+                                _CLEANUP_FUNCS[rtype](name)
+                            except Exception as e:
+                                warnings.warn(
+                                    'resource_tracker: %s: %r' % (name, e))
+
                     else:
                         raise RuntimeError('unrecognized command %r' % cmd)
                 except BaseException:
@@ -258,30 +311,43 @@ def main(fd, verbose=0):
                         pass
     finally:
         # all processes have terminated; cleanup any remaining resources
-        for rtype, rtype_cache in cache.items():
-            if rtype_cache:
+        def _unlink_resources(rtype_registry, rtype):
+            if rtype_registry:
                 try:
                     warnings.warn('resource_tracker: There appear to be %d '
                                   'leaked %s objects to clean up at shutdown' %
-                                  (len(rtype_cache), rtype))
+                                  (len(rtype_registry), rtype))
                 except Exception:
                     pass
-            for name in rtype_cache:
+            for name in rtype_registry:
                 # For some reason the process which created and registered this
                 # resource has failed to unregister it. Presumably it has
                 # died.  We therefore clean it up.
                 try:
                     _CLEANUP_FUNCS[rtype](name)
-                    if verbose:  # pragma: no cover
-                        sys.stderr.write("[ResourceTracker] unlink {}\n"
+                    if verbose:
+                        util.debug("[ResourceTracker] unlink {}"
                                          .format(name))
-                        sys.stderr.flush()
                 except Exception as e:
                     warnings.warn('resource_tracker: %s: %r' % (name, e))
 
-    if verbose:  # pragma: no cover
-        sys.stderr.write("resource tracker shut down\n")
-        sys.stderr.flush()
+        for rtype, rtype_registry in registry.items():
+            if rtype == "folder":
+                continue
+            else:
+                _unlink_resources(rtype_registry, rtype)
+
+        # The default cleanup routine for folders deletes everything inside
+        # those folders recursively, which can include other resources tracked
+        # by the resource tracker). To limit the risk of the resource tracker
+        # attempting to delete twice a resource (once as part of a tracked
+        # folder, and once as a resource), we delete the folders after all
+        # other resource types.
+        if "folder" in registry:
+            _unlink_resources(registry["folder"], "folder")
+
+    if verbose:
+        util.debug("resource tracker shut down")
 
 
 #


Test patch:
diff --git a/tests/test_resource_tracker.py b/tests/test_resource_tracker.py
index f04401a..e3569dc 100644
--- a/tests/test_resource_tracker.py
+++ b/tests/test_resource_tracker.py
@@ -6,7 +6,9 @@ import os
 import pytest
 import re
 import signal
+import subprocess
 import sys
+import tempfile
 import time
 import warnings
 import weakref
@@ -38,51 +40,58 @@ class TestResourceTracker:
         # child process. If the two processes do not share the same
         # resource_tracker, a cache KeyError should be printed in stderr.
         import subprocess
-        folder_name = 'loky_tempfolder'
         cmd = '''if 1:
         import os, sys
 
         from loky import ProcessPoolExecutor
         from loky.backend import resource_tracker
         from loky.backend.semlock import SemLock
+        from tempfile import NamedTemporaryFile
 
-        resource_tracker.VERBOSE=True
-        folder_name = "{}"
 
-        # We don't need to create the semaphore as registering / unregistering
-        # operations simply add / remove entries from a cache, but do not
-        # manipulate the actual semaphores.
-        resource_tracker.register(folder_name, "folder")
+        tmpfile = NamedTemporaryFile(delete=False)
+        tmpfile.close()
+        filename = tmpfile.name
+        resource_tracker.VERBOSE = True
 
-        def unregister(name, rtype):
-            # resource_tracker.unregister is actually a bound method of the
+        resource_tracker.register(filename, "file")
+
+        def maybe_unlink(name, rtype):
+            # resource_tracker.maybe_unlink is actually a bound method of the
             # ResourceTracker. We need a custom wrapper to avoid object
             # serialization.
             from loky.backend import resource_tracker
-            resource_tracker.unregister(folder_name, rtype)
+            resource_tracker.maybe_unlink(name, rtype)
 
+        print(filename)
         e = ProcessPoolExecutor(1)
-        e.submit(unregister, folder_name, "folder").result()
+        e.submit(maybe_unlink, filename, "file").result()
         e.shutdown()
         '''
         try:
             p = subprocess.Popen(
-                [sys.executable, '-E', '-c', cmd.format(folder_name)],
-                stderr=subprocess.PIPE)
+                [sys.executable, '-E', '-c', cmd],
+                stderr=subprocess.PIPE,
+                stdout=subprocess.PIPE)
             p.wait()
 
+            filename = p.stdout.readline().decode('utf-8').strip()
             err = p.stderr.read().decode('utf-8')
             p.stderr.close()
+            p.stdout.close()
+
+            pattern = "decremented refcount of file %s" % filename
+            assert pattern in err
+            assert "leaked" not in err
 
-            assert re.search("unregister %s" % folder_name, err) is not None
-            assert re.search("leaked", err) is None
-            assert re.search("KeyError: '%s'" % folder_name, err) is None
+            pattern = "KeyError: '%s'" % filename
+            assert pattern not in err
 
         finally:
             executor.shutdown()
 
     # The following four tests are inspired from cpython _test_multiprocessing
-    @pytest.mark.parametrize("rtype", ["folder", "semlock"])
+    @pytest.mark.parametrize("rtype", ["file", "folder", "semlock"])
     def test_resource_tracker(self, rtype):
         #
         # Check that killing process does not leak named resources
@@ -96,18 +105,7 @@ class TestResourceTracker:
 
             from loky.backend.semlock import SemLock
             from loky.backend import resource_tracker, reduction
-
-            def create_resource(rtype):
-                if rtype == "folder":
-                    return tempfile.mkdtemp(dir=os.getcwd())
-
-                elif rtype == "semlock":
-                    name = "/loky-%i-%s" % (os.getpid(), next(SemLock._rand))
-                    lock = SemLock(1, 1, 1, name)
-                    return name
-                else:
-                    raise ValueError(
-                        "Resource type %s not understood" % rtype)
+            from utils import create_resource
 
             for _ in range(2):
                 rname = create_resource("{rtype}")
@@ -119,9 +117,12 @@ class TestResourceTracker:
             time.sleep(10)
         '''
         cmd = cmd.format(rtype=rtype, parent_pid=os.getpid())
-        p = subprocess.Popen([sys.executable, '-E', '-c', cmd],
+        env = os.environ.copy()
+        env['PYTHONPATH'] = os.path.dirname(__file__)
+        p = subprocess.Popen([sys.executable, '-c', cmd],
                              stderr=subprocess.PIPE,
-                             stdout=subprocess.PIPE)
+                             stdout=subprocess.PIPE,
+                             env=env)
         name1 = p.stdout.readline().rstrip().decode('ascii')
         name2 = p.stdout.readline().rstrip().decode('ascii')
 
@@ -149,14 +150,76 @@ class TestResourceTracker:
         # resource 1 is still registered, but was destroyed externally: the
         # tracker is expected to complain.
         if sys.platform == "win32":
-            expected = ("resource_tracker: %s: (WindowsError\\((%d)|"
-                        "FileNotFoundError)" % (re.escape(name1), errno.ESRCH))
+            errno_map = {'file': 2, 'folder': 3}
+            expected = (
+                "resource_tracker: %s: (WindowsError\\((%d)|"
+                "FileNotFoundError)" % (re.escape(name1), errno_map[rtype])
+            )
         else:
             expected = ("resource_tracker: %s: (OSError\\(%d|"
                         "FileNotFoundError)" % (re.escape(name1),
                                                 errno.ENOENT))
         assert re.search(expected, err) is not None
 
+    @pytest.mark.parametrize("rtype", ["file", "folder", "semlock"])
+    def test_resource_tracker_refcounting(self, rtype):
+        if sys.platform == "win32" and rtype == "semlock":
+            pytest.skip("no semlock on windows")
+
+        cmd = '''if 1:
+        import os
+        import tempfile
+        import time
+        from loky.backend.semlock import SemLock, _sem_open
+        from loky.backend import resource_tracker
+        from utils import resource_unlink, create_resource, resource_exists
+
+        resource_tracker.VERBOSE = True
+
+        try:
+            name = create_resource("{rtype}")
+            assert resource_exists(name, "{rtype}")
+
+            from loky.backend.resource_tracker import _resource_tracker
+            _resource_tracker.register(name, "{rtype}")
+            _resource_tracker.register(name, "{rtype}")
+
+            # Forget all information about the resource, but do not try to
+            # remove it
+            _resource_tracker.unregister(name, "{rtype}")
+            time.sleep(1)
+            assert resource_exists(name, "{rtype}")
+
+            _resource_tracker.register(name, "{rtype}")
+            _resource_tracker.register(name, "{rtype}")
+            _resource_tracker.maybe_unlink(name, "{rtype}")
+            time.sleep(1)
+            assert resource_exists(name, "{rtype}")
+
+            _resource_tracker.maybe_unlink(name, "{rtype}")
+            for _ in range(100):
+                if not resource_exists(name, "{rtype}"):
+                    break
+                time.sleep(.1)
+            else:
+                raise AssertionError("%s was not unlinked in time"  % name)
+        finally:
+            if resource_exists(name, "{rtype}"):
+                resource_unlink(name, "{rtype}")
+        '''
+
+        env = os.environ.copy()
+        env['PYTHONPATH'] = os.path.dirname(__file__)
+        p = subprocess.Popen(
+            [sys.executable, '-c', cmd.format(rtype=rtype)],
+            stderr=subprocess.PIPE,
+            stdout=subprocess.PIPE,
+            env=env
+        )
+        p.wait()
+        out, err = p.communicate()
+        assert p.returncode == 0, err
+
     def check_resource_tracker_death(self, signum, should_die):
         # bpo-31310: if the semaphore tracker process has died, it should
         # be restarted implicitly.
diff --git a/tests/utils.py b/tests/utils.py
index 2d5597d..34f4ce1 100644
--- a/tests/utils.py
+++ b/tests/utils.py
@@ -7,13 +7,58 @@ import warnings
 import threading
 import subprocess
 import contextlib
-from tempfile import mkstemp
+from tempfile import mkstemp, mkdtemp, NamedTemporaryFile
+from loky.backend import resource_tracker
+from loky.backend.semlock import SemLock, _sem_open
+
+try:
+    FileNotFoundError = FileNotFoundError
+except NameError:  # FileNotFoundError is Python 3-only
+    from loky.backend.semlock import FileNotFoundError
+
 
 if sys.version_info[0] == 2:
     class TimeoutError(OSError):
         pass
 
 
+def resource_unlink(name, rtype):
+    resource_tracker._CLEANUP_FUNCS[rtype](name)
+
+
+def create_resource(rtype):
+    if rtype == "folder":
+        return mkdtemp(dir=os.getcwd())
+
+    elif rtype == "semlock":
+        name = "loky-%i-%s" % (os.getpid(), next(SemLock._rand))
+        _ = SemLock(1, 1, 1, name)
+        return name
+    elif rtype == "file":
+        tmpfile = NamedTemporaryFile(delete=False)
+        tmpfile.close()
+        return tmpfile.name
+    else:
+        raise ValueError("Resource type %s not understood" % rtype)
+
+
+def resource_exists(name, rtype):
+    if rtype in ["folder", "file"]:
+        return os.path.exists(name)
+    elif rtype == "semlock":
+        # On OSX, semaphore are not visible in the file system, we must
+        # try to open the semaphore to check if it exists.
+        from loky.backend.semlock import pthread
+        try:
+            h = _sem_open(name.encode('ascii'))
+            pthread.sem_close(h)
+            return True
+        except FileNotFoundError:
+            return False
+    else:
+        raise ValueError("Resource type %s not understood" % rtype)
+
+
 @contextlib.contextmanager
 def captured_output(stream_name):
     """Return a context manager used by captured_stdout/stdin/stderr
@@ -156,7 +201,7 @@ except ImportError:
 
 # A decorator to run tests only when numpy is available
 try:
-    from ._openmp_test_helper.parallel_sum import parallel_sum  # noqa F401
+    from _openmp_test_helper.parallel_sum import parallel_sum  # noqa F401
 
     def with_parallel_sum(func):
         """A decorator to skip tests if parallel_sum is not compiled."""