Better support ProcessPoolExecutors (#5063)

Previously we sent non-serializable state, like locks and the worker
itself into the function that was sent to the executor

Now we only do this if we think tha the executor is a
ThreadPoolExecutor

* Remove Worker.executor_submit

Fixes #3938

This is old and was a workaround for a Tornado issue

* remove gen.coroutine from worker.py

We're generally trying to clean this up from the codebase

* add Worker.active_keys

Author: mrocklin

distributed/tests/test_worker.py

@@ -5,7 +5,7 @@
 import sys
 import threading
 import traceback
-from concurrent.futures import ThreadPoolExecutor
+from concurrent.futures import ProcessPoolExecutor, ThreadPoolExecutor
 from numbers import Number
 from operator import add
 from time import sleep
@@ -508,17 +508,14 @@ async def f(dask_worker=None):
 @gen_cluster(client=True, nthreads=[])
 async def test_Executor(c, s):
     with ThreadPoolExecutor(2) as e:
-        w = Worker(s.address, executor=e)
-        assert w.executor is e
-        w = await w
+        async with Worker(s.address, executor=e) as w:
+            assert w.executor is e
 
-        future = c.submit(inc, 1)
-        result = await future
-        assert result == 2
-
-        assert e._threads  # had to do some work
+            future = c.submit(inc, 1)
+            result = await future
+            assert result == 2
 
-        await w.close()
+            assert e._threads  # had to do some work
 
 
 @pytest.mark.skip(
@@ -2028,6 +2025,34 @@ def get_thread_name():
                 assert "Dask-GPU-Threads" in gpu_result
 
 
+@gen_cluster(client=True)
+async def test_process_executor(c, s, a, b):
+    with ProcessPoolExecutor() as e:
+        a.executors["processes"] = e
+        b.executors["processes"] = e
+
+        future = c.submit(os.getpid, pure=False)
+        assert (await future) == os.getpid()
+
+        with dask.annotate(executor="processes"):
+            future = c.submit(os.getpid, pure=False)
+
+        assert (await future) != os.getpid()
+
+
+@gen_cluster(client=True)
+async def test_process_executor_kills_process(c, s, a, b):
+    with ProcessPoolExecutor() as e:
+        a.executors["processes"] = e
+        b.executors["processes"] = e
+
+        with dask.annotate(executor="processes", retries=1):
+            future = c.submit(sys.exit, 1)
+
+        exc = await future.exception()
+        assert "SystemExit(1)" in repr(exc)
+
+
 def assert_task_states_on_worker(expected, worker):
     for dep_key, expected_state in expected.items():
         assert dep_key in worker.tasks, (worker.name, dep_key, worker.tasks)

distributed/worker.py

@@ -19,7 +19,6 @@
 from typing import Dict, Iterable, Optional
 
 from tlz import first, keymap, merge, pluck  # noqa: F401
-from tornado import gen
 from tornado.ioloop import IOLoop, PeriodicCallback
 
 import dask
@@ -438,6 +437,7 @@ def __init__(
 
         self.active_threads_lock = threading.Lock()
         self.active_threads = dict()
+        self.active_keys = set()
         self.profile_keys = defaultdict(profile.create)
         self.profile_keys_history = deque(maxlen=3600)
         self.profile_recent = profile.create()
@@ -968,16 +968,14 @@ async def heartbeat(self):
         logger.debug("Heartbeat: %s", self.address)
         try:
             start = time()
-            with self.active_threads_lock:
-                active_keys = list(self.active_threads.values())
             response = await retry_operation(
                 self.scheduler.heartbeat_worker,
                 address=self.contact_address,
                 now=start,
                 metrics=await self.get_metrics(),
                 executing={
                     key: start - self.tasks[key].start_time
-                    for key in active_keys
+                    for key in self.active_keys
                     if key in self.tasks
                 },
             )
@@ -2686,41 +2684,6 @@ def release_key(
     # Execute Task #
     ################
 
-    # FIXME: this breaks if changed to async def...
-    # xref: https://github.com/dask/distributed/issues/3938
-    @gen.coroutine
-    def executor_submit(self, key, function, args=(), kwargs=None, executor=None):
-        """Safely run function in thread pool executor
-
-        We've run into issues running concurrent.future futures within
-        tornado.  Apparently it's advantageous to use timeouts and periodic
-        callbacks to ensure things run smoothly.  This can get tricky, so we
-        pull it off into an separate method.
-        """
-        executor = executor or self.executors["default"]
-        job_counter[0] += 1
-        # logger.info("%s:%d Starts job %d, %s", self.ip, self.port, i, key)
-        kwargs = kwargs or {}
-        future = executor.submit(function, *args, **kwargs)
-        pc = PeriodicCallback(
-            lambda: logger.debug("future state: %s - %s", key, future._state), 1000
-        )
-        ts = self.tasks.get(key)
-        if ts is not None:
-            ts.start_time = time()
-        pc.start()
-        try:
-            yield future
-        finally:
-            pc.stop()
-            if ts is not None:
-                ts.stop_time = time()
-
-        result = future.result()
-
-        # logger.info("Finish job %d, %s", i, key)
-        raise gen.Return(result)
-
     def run(self, comm, function, args=(), wait=True, kwargs=None):
         return run(self, comm, function=function, args=args, kwargs=kwargs, wait=wait)
 
@@ -2782,19 +2745,16 @@ async def actor_execute(
             if iscoroutinefunction(func):
                 result = await func(*args, **kwargs)
             elif separate_thread:
-                result = await self.executor_submit(
-                    name,
+                result = await self.loop.run_in_executor(
+                    self.executors["actor"],
                     apply_function_actor,
-                    args=(
-                        func,
-                        args,
-                        kwargs,
-                        self.execution_state,
-                        name,
-                        self.active_threads,
-                        self.active_threads_lock,
-                    ),
-                    executor=self.executors["actor"],
+                    func,
+                    args,
+                    kwargs,
+                    self.execution_state,
+                    name,
+                    self.active_threads,
+                    self.active_threads_lock,
                 )
             else:
                 result = func(*args, **kwargs)
@@ -2946,11 +2906,14 @@ async def execute(self, key, report=False):
                 executor,
             )  # TODO: comment out?
             assert key == ts.key
+            self.active_keys.add(ts.key)
             try:
-                result = await self.executor_submit(
-                    ts.key,
-                    apply_function,
-                    args=(
+                e = self.executors[executor]
+                ts.start_time = time()
+                if "ThreadPoolExecutor" in str(type(e)):
+                    result = await self.loop.run_in_executor(
+                        e,
+                        apply_function,
                         function,
                         args2,
                         kwargs2,
@@ -2959,12 +2922,32 @@ async def execute(self, key, report=False):
                         self.active_threads,
                         self.active_threads_lock,
                         self.scheduler_delay,
-                    ),
-                    executor=self.executors[executor],
-                )
+                    )
+                else:
+                    try:
+                        start = time() + self.scheduler_delay
+                        result = await self.loop.run_in_executor(
+                            e,
+                            apply_function_simple,
+                            function,
+                            args2,
+                            kwargs2,
+                            self.scheduler_delay,
+                        )
+                    except BaseException as e:
+                        msg = error_message(e)
+                        msg["op"] = "task-erred"
+                        msg["actual-exception"] = e
+                        msg["start"] = start
+                        msg["stop"] = time() + self.scheduler_delay
+                        msg["thread"] = None
+                        result = msg
+
             except RuntimeError as e:
                 executor_error = e
                 raise
+            finally:
+                self.active_keys.discard(ts.key)
 
             # We'll need to check again for the task state since it may have
             # changed since the execution was kicked off. In particular, it may
@@ -3854,6 +3837,27 @@ def apply_function(
     thread_state.start_time = time()
     thread_state.execution_state = execution_state
     thread_state.key = key
+
+    msg = apply_function_simple(function, args, kwargs, time_delay)
+
+    with active_threads_lock:
+        del active_threads[ident]
+    return msg
+
+
+def apply_function_simple(
+    function,
+    args,
+    kwargs,
+    time_delay,
+):
+    """Run a function, collect information
+
+    Returns
+    -------
+    msg: dictionary with status, result/error, timings, etc..
+    """
+    ident = threading.get_ident()
     start = time()
     try:
         result = function(*args, **kwargs)
@@ -3874,8 +3878,6 @@ def apply_function(
     msg["start"] = start + time_delay
     msg["stop"] = end + time_delay
     msg["thread"] = ident
-    with active_threads_lock:
-        del active_threads[ident]
     return msg
 
 
@@ -4020,9 +4022,8 @@ async def run(server, comm, function, args=(), kwargs=None, is_coro=None, wait=T
     pass
 else:
 
-    @gen.coroutine
-    def gpu_metric(worker):
-        result = yield offload(nvml.real_time)
+    async def gpu_metric(worker):
+        result = await offload(nvml.real_time)
         return result
 
     DEFAULT_METRICS["gpu"] = gpu_metric