Repository: dask/distributed

Issue:
Should worker_client execute rejoin as well
We have a task which calls a function executing worker_client. As this function gets called few times as part of the task execution, looks like we get an error:

File "e:\anaconda\python\lib\site-packages\distributed\worker_client.py", line 46, in worker_client
    secede()  # have this thread secede from the thread pool
  File "e:\anaconda\python\lib\site-packages\distributed\threadpoolexecutor.py", line 104, in secede
    thread_state.executor._threads.remove(threading.current_thread())
KeyError: <Thread(ThreadPool worker 0, started daemon 8556)>

Added rejoin after the "with" block and then it works fine. Should the work_client rejoin or have option to rejoin for such cases. Please let me know if this should be handled in any other way.

**Note** - Have left separate_thread option as True as not sure if setting it to False can result in the deadlock descibed in the documentation which is why we are utilizing worker_client. We are running single worker for each core with single thread.


Discussion:
Yes, I think that it probably makes sense for the worker_client context
manager to rejoin.  My apologies for my comment in the other issue where I
erroneously assumed that this was already the case.

On Thu, May 31, 2018 at 10:25 PM, ameetshah1983 <notifications@github.com>
wrote:

> We have a task which calls a function executing worker_client. As this
> function gets called few times as part of the task execution, looks like we
> get an error:
>
> File "e:\anaconda\python\lib\site-packages\distributed\worker_client.py",
> line 46, in worker_client
> secede() # have this thread secede from the thread pool
> File "e:\anaconda\python\lib\site-packages\distributed\threadpoolexecutor.py",
> line 104, in secede
> thread_state.executor._threads.remove(threading.current_thread())
> KeyError: <Thread(ThreadPool worker 0, started daemon 8556)>
>
> Added rejoin after the "with" block and then it works fine. Should the
> work_client rejoin or have option to rejoin for such cases. Please let me
> know if this should be handled in any other way.
>
> —
> You are receiving this because you are subscribed to this thread.
> Reply to this email directly, view it on GitHub
> <https://github.com/dask/distributed/issues/2021>, or mute the thread
> <https://github.com/notifications/unsubscribe-auth/AASszJSvnO0QHUTC-PH9CjpQmpyb1_Auks5t4KYGgaJpZM4UWACo>
> .
>


Resolving patch:
diff --git a/distributed/worker_client.py b/distributed/worker_client.py
--- a/distributed/worker_client.py
+++ b/distributed/worker_client.py
@@ -3,7 +3,7 @@
 from contextlib import contextmanager
 import warnings
 
-from .threadpoolexecutor import secede
+from .threadpoolexecutor import secede, rejoin
 from .worker import thread_state, get_client, get_worker
 
 
@@ -48,6 +48,9 @@ def worker_client(timeout=3, separate_thread=True):
 
     yield client
 
+    if separate_thread:
+        rejoin()
+
 
 def local_client(*args, **kwargs):
     warnings.warn("local_client has moved to worker_client")


Test patch:
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
@@ -107,3 +107,15 @@ def f():
 
         future = e.submit(f)
         result = future.result()
+
+
+def test_rejoin_idempotent():
+    with ThreadPoolExecutor(2) as e:
+        def f():
+            secede()
+            for i in range(5):
+                rejoin()
+            return 1
+
+        future = e.submit(f)
+        result = future.result()
diff --git a/distributed/tests/test_worker_client.py b/distributed/tests/test_worker_client.py
--- a/distributed/tests/test_worker_client.py
+++ b/distributed/tests/test_worker_client.py
@@ -1,6 +1,7 @@
 from __future__ import print_function, division, absolute_import
 
 import random
+import threading
 from time import sleep
 import warnings
 
@@ -284,3 +285,15 @@ def f():
 
     result = yield c.compute(f())
     assert result == 1
+
+
+@gen_cluster(client=True)
+def test_worker_client_rejoins(c, s, a, b):
+    def f():
+        with worker_client():
+            pass
+
+        return threading.current_thread() in get_worker().executor._threads
+
+    result = yield c.submit(f)
+    assert result