Allow nested worker_client calls

distributed/tests/test_worker_client.py

@@ -180,6 +180,27 @@ def mysum():
         assert time() < start + 3
 
 
+@gen_cluster(client=True)
+async def test_nested_worker_clients(c, s, a, b):
+    def mysum():
+        result = 0
+        with worker_client() as lc:
+            futures = lc.map(double, range(100))
+            for f in as_completed(futures):
+                result += f.result()
+        return result
+
+    def inc_mysum():
+        with worker_client() as lc:
+            future = lc.submit(inc, mysum())
+            result = future.result()
+        return result
+
+    future = c.submit(inc_mysum)
+    result = await future
+    assert result == 9901
+
+
 @gen_cluster(client=True, nthreads=[("127.0.0.1", 3)])
 async def test_separate_thread_false(c, s, a):
     a.count = 0

distributed/worker_client.py

@@ -68,16 +68,20 @@ def worker_client(timeout=None, separate_thread=True):
                 pass
             else:
                 duration = time() - thread_state.start_time
-                secede()  # have this thread secede from the thread pool
-                stack.callback(rejoin)
-                worker.loop.add_callback(
-                    worker.handle_stimulus,
-                    SecedeEvent(
-                        key=thread_state.key,
-                        compute_duration=duration,
-                        stimulus_id=f"worker-client-secede-{time()}",
-                    ),
-                )
+                try:
+                    secede()  # have this thread secede from the thread pool
+                except KeyError:  # already seceded
+                    pass
+                else:
+                    stack.callback(rejoin)
+                    worker.loop.add_callback(
+                        worker.handle_stimulus,
+                        SecedeEvent(
+                            key=thread_state.key,
+                            compute_duration=duration,
+                            stimulus_id=f"worker-client-secede-{time()}",
+                        ),
+                    )
 
         yield client