Repository: dask/distributed

Issue:
Add rejoin operation to rejoin thread pool
Currently the `secede` function allows a task to remove itself from the worker's current thread pool, opening up a space for more tasks.

We might consider an inverse operation, `rejoin` that blocks until a new spot in the thread pool has opened up.  This would enable long-running task computations to avoid contention of many threads computing at once.

First suggested by @adamklein in https://github.com/dask/distributed/issues/1342

also cc @ogrisel


Discussion:

Resolving patch:
diff --git a/distributed/__init__.py b/distributed/__init__.py
--- a/distributed/__init__.py
+++ b/distributed/__init__.py
@@ -10,6 +10,7 @@
 from .nanny import Nanny
 from .queues import Queue
 from .scheduler import Scheduler
+from .threadpoolexecutor import rejoin
 from .utils import sync
 from .variable import Variable
 from .worker import Worker, get_worker, get_client, secede
diff --git a/distributed/threadpoolexecutor.py b/distributed/threadpoolexecutor.py
--- a/distributed/threadpoolexecutor.py
+++ b/distributed/threadpoolexecutor.py
@@ -26,7 +26,6 @@
 import logging
 import threading
 
-from .compatibility import get_thread_identity
 from .metrics import time
 
 logger = logging.getLogger(__name__)
@@ -40,6 +39,13 @@ def _worker(executor, work_queue):
 
     try:
         while thread_state.proceed:
+            with executor._rejoin_lock:
+                if executor._rejoin_list:
+                    rejoin_thread, rejoin_event = executor._rejoin_list.pop()
+                    executor._threads.add(rejoin_thread)
+                    executor._threads.remove(threading.current_thread())
+                    rejoin_event.set()
+                    break
             task = work_queue.get()
             if task is not None:  # sentinel
                 task.run()
@@ -56,6 +62,11 @@ def _worker(executor, work_queue):
 
 
 class ThreadPoolExecutor(thread.ThreadPoolExecutor):
+    def __init__(self, *args, **kwargs):
+        super(ThreadPoolExecutor, self).__init__(*args, **kwargs)
+        self._rejoin_list = []
+        self._rejoin_lock = threading.Lock()
+
     def _adjust_thread_count(self):
         if len(self._threads) < self._max_workers:
             t = threading.Thread(target=_worker,
@@ -80,16 +91,38 @@ def shutdown(self, wait=True, timeout=None):
                 t.join(timeout=timeout2)
 
 
-def secede():
-    """ Have this thread secede from the ThreadPoolExecutor """
+def secede(adjust=True):
+    """ Have this thread secede from the ThreadPoolExecutor
+
+    See Also
+    --------
+    rejoin: rejoin the thread pool
+    """
     thread_state.proceed = False
-    ident = get_thread_identity()
     with threads_lock:
-        for t in list(thread_state.executor._threads):
-            if t.ident == ident:
-                thread_state.executor._threads.remove(t)
-                break
-        thread_state.executor._adjust_thread_count()
+        thread_state.executor._threads.remove(threading.current_thread())
+        if adjust:
+            thread_state.executor._adjust_thread_count()
+
+
+def rejoin():
+    """ Have this thread rejoin the ThreadPoolExecutor
+
+    This will block until a new slot opens up in the executor.  The next thread
+    to finish a task will leave the pool to allow this one to join.
+
+    See Also
+    --------
+    secede: leave the thread pool
+    """
+    thread = threading.current_thread()
+    event = threading.Event()
+    e = thread_state.executor
+    with e._rejoin_lock:
+        e._rejoin_list.append((thread, event))
+    e.submit(lambda: None)
+    event.wait()
+    thread_state.proceed = True
 
 
 threads_lock = threading.Lock()


Test patch:
diff --git a/distributed/tests/test_client.py b/distributed/tests/test_client.py
--- a/distributed/tests/test_client.py
+++ b/distributed/tests/test_client.py
@@ -4469,6 +4469,7 @@ def f():
 @slow
 @gen_cluster(client=True, ncores=[('127.0.0.1', 1)] * 2, timeout=60)
 def test_secede_balances(c, s, a, b):
+    count = threading.active_count()
     def f(x):
         client = get_client()
         sleep(0.01)  # do some work
@@ -4481,7 +4482,7 @@ def f(x):
     start = time()
     while not all(f.status == 'finished' for f in futures):
         yield gen.sleep(0.01)
-        assert threading.active_count() < 50
+        assert threading.active_count() < count + 20
 
     # assert 0.005 < s.task_duration['f'] < 0.1
     assert len(a.log) < 2 * len(b.log)
diff --git a/distributed/tests/test_threadpoolexecutor.py b/distributed/tests/test_threadpoolexecutor.py
--- a/distributed/tests/test_threadpoolexecutor.py
+++ b/distributed/tests/test_threadpoolexecutor.py
@@ -1,29 +1,30 @@
 from time import sleep
+import threading
 
 from distributed.metrics import time
-from distributed.threadpoolexecutor import ThreadPoolExecutor, secede
+from distributed.threadpoolexecutor import ThreadPoolExecutor, secede, rejoin
 
 
 def test_tpe():
-    e = ThreadPoolExecutor(2)
-    list(e.map(sleep, [0.01] * 4))
+    with ThreadPoolExecutor(2) as e:
+        list(e.map(sleep, [0.01] * 4))
 
-    threads = e._threads.copy()
-    assert len(threads) == 2
+        threads = e._threads.copy()
+        assert len(threads) == 2
 
-    def f():
-        secede()
-        return 1
+        def f():
+            secede()
+            return 1
 
-    assert e.submit(f).result() == 1
+        assert e.submit(f).result() == 1
 
-    list(e.map(sleep, [0.01] * 4))
-    assert len(threads | e._threads) == 3
+        list(e.map(sleep, [0.01] * 4))
+        assert len(threads | e._threads) == 3
 
-    start = time()
-    while all(t.is_alive() for t in threads):
-        sleep(0.01)
-        assert time() < start + 1
+        start = time()
+        while all(t.is_alive() for t in threads):
+            sleep(0.01)
+            assert time() < start + 1
 
 
 def test_shutdown_timeout():
@@ -46,3 +47,52 @@ def test_shutdown_timeout_raises():
     e.shutdown(timeout=0.1)
     end = time()
     assert end - start > 0.05
+
+
+def test_secede_rejoin_busy():
+    with ThreadPoolExecutor(2) as e:
+
+        def f():
+            assert threading.current_thread() in e._threads
+            secede()
+            sleep(0.1)
+            assert threading.current_thread() not in e._threads
+            rejoin()
+            assert len(e._threads) == 2
+            assert threading.current_thread() in e._threads
+            return threading.current_thread()
+
+        future = e.submit(f)
+        L = [e.submit(sleep, 0.2) for i in range(10)]
+        start = time()
+        special_thread = future.result()
+        stop = time()
+
+        assert 0.1 < stop - start < 0.3
+
+        assert len(e._threads) == 2
+        assert special_thread in e._threads
+
+        def f():
+            sleep(0.01)
+            return threading.current_thread()
+
+        futures = [e.submit(f) for _ in range(10)]
+        assert special_thread in {future.result() for future in futures}
+
+
+def test_secede_rejoin_quiet():
+    with ThreadPoolExecutor(2) as e:
+
+        def f():
+            assert threading.current_thread() in e._threads
+            secede()
+            sleep(0.1)
+            assert threading.current_thread() not in e._threads
+            rejoin()
+            assert len(e._threads) == 2
+            assert threading.current_thread() in e._threads
+            return threading.current_thread()
+
+        future = e.submit(f)
+        result = future.result()