Repository: dask/distributed

Issue:
`BaseException` in task leads to task never completing
**What happened**:

If a task raises a `BaseException` (`KeyboardInterrupt`, `SystemExit`, etc.), the task will appear to be `processing` forever.

**What you expected to happen**:

The task should definitely not deadlock. But what should actually happen, I'm not sure. Could go two ways:
1. The entire worker should shut down gracefully.
2. We should catch them just like any other exceptions and error the task.

More discussion in comments.

**Minimal Complete Verifiable Example**:

```python
In [1]: import distributed

In [2]: client = distributed.Client(n_workers=1)

In [3]: def raiser():
   ...:     raise BaseException("this could be a KeyboardInterrupt!")
   ...: 

In [4]: f = client.submit(raiser)

In [5]: Exception in callback IOLoop.add_future.<locals>.<lambda>(<Task finishe...dInterrupt!')>) at /Users/gabe/miniconda3/envs/dask-distributed/lib/python3.9/site-packages/tornado/ioloop.py:688
handle: <Handle IOLoop.add_future.<locals>.<lambda>(<Task finishe...dInterrupt!')>) at /Users/gabe/miniconda3/envs/dask-distributed/lib/python3.9/site-packages/tornado/ioloop.py:688>
Traceback (most recent call last):
  File "/Users/gabe/miniconda3/envs/dask-distributed/lib/python3.9/asyncio/events.py", line 80, in _run
    self._context.run(self._callback, *self._args)
  File "/Users/gabe/miniconda3/envs/dask-distributed/lib/python3.9/site-packages/tornado/ioloop.py", line 688, in <lambda>
    lambda f: self._run_callback(functools.partial(callback, future))
  File "/Users/gabe/miniconda3/envs/dask-distributed/lib/python3.9/site-packages/tornado/ioloop.py", line 741, in _run_callback
    ret = callback()
  File "/Users/gabe/miniconda3/envs/dask-distributed/lib/python3.9/site-packages/tornado/ioloop.py", line 765, in _discard_future_result
    future.result()
  File "/Users/gabe/dev/distributed/distributed/worker.py", line 3504, in execute
    result = await self.loop.run_in_executor(
  File "/Users/gabe/dev/distributed/distributed/_concurrent_futures_thread.py", line 65, in run
    result = self.fn(*self.args, **self.kwargs)
  File "/Users/gabe/dev/distributed/distributed/worker.py", line 4503, in apply_function
    msg = apply_function_simple(function, args, kwargs, time_delay)
  File "/Users/gabe/dev/distributed/distributed/worker.py", line 4525, in apply_function_simple
    result = function(*args, **kwargs)
  File "<ipython-input-3-7f701e80695b>", line 2, in raiser
BaseException: this could be a KeyboardInterrupt!
In [5]: 

In [5]: client.processing()
Out[5]: {'tcp://127.0.0.1:58316': ('raiser-e3b7ab59305f9e4ddb4ecddd75c55f85',)}

In [6]: client.call_stack()
Out[6]: 
{'tcp://127.0.0.1:58316': {'raiser-e3b7ab59305f9e4ddb4ecddd75c55f85': ('  File "/Users/gabe/miniconda3/envs/dask-distributed/lib/python3.9/threading.py", line 912, in _bootstrap\n\tself._bootstrap_inner()\n',
   '  File "/Users/gabe/miniconda3/envs/dask-distributed/lib/python3.9/threading.py", line 954, in _bootstrap_inner\n\tself.run()\n',
   '  File "/Users/gabe/miniconda3/envs/dask-distributed/lib/python3.9/threading.py", line 892, in run\n\tself._target(*self._args, **self._kwargs)\n',
   '  File "/Users/gabe/dev/distributed/distributed/threadpoolexecutor.py", line 51, in _worker\n\ttask = work_queue.get(timeout=1)\n',
   '  File "/Users/gabe/miniconda3/envs/dask-distributed/lib/python3.9/queue.py", line 180, in get\n\tself.not_empty.wait(remaining)\n',
   '  File "/Users/gabe/miniconda3/envs/dask-distributed/lib/python3.9/threading.py", line 316, in wait\n\tgotit = waiter.acquire(True, timeout)\n')}}

In [7]: client.submit(lambda: 1).result(timeout=5)  # the worker still works fine; just that task is stuck now
Out[7]: 1
```

**Anything else we need to know?**:

**Environment**:

- Dask version: 2022.2.1
- Python version: 3.9.5
- Operating System: macOS
- Install method (conda, pip, source): source

cc @fjetter @graingert @crusaderky 


Discussion:
Arguments for the two possible ways we could handle `BaseException`s from tasks:

### The entire worker should shut down gracefully.

The `BaseException` should be propagated back to the event loop, and handled as usual there (which _hopefully_ would mean shutting down gracefully; we're known to not handle KeyboardInterrupts that well though).

Rationale: Python explicitly states that `BaseExceptions` should not be considered user errors:
> **`BaseException`**: The base class for all built-in exceptions. It is not meant to be directly inherited by user-defined classes (for that, use [Exception](https://docs.python.org/3/library/exceptions.html#Exception)).
https://docs.python.org/3/library/exceptions.html#BaseException

Furthermore, Python says this about normal `Exception`:
> **`Exception`**: All built-in, non-system-exiting exceptions are derived from this class. All user-defined exceptions should also be derived from this class.
https://docs.python.org/3/library/exceptions.html#Exception

That seems to imply that `BaseException` is considered a "system-exiting exception".

So a `BaseException` indicates a system/process-level problem, therefore the process overall should handle the exception. We should propagate this from the thread and re-raise it in the event loop.

### We should catch them just like any other exceptions

Just change this to `BaseException`: https://github.com/dask/distributed/blob/1b4d13fe4f6898bb825bd9e58dc3e795a7230be5/distributed/worker.py#L4526-L4529 (It used to be this; @fjetter changed it in https://github.com/dask/distributed/pull/5128.)

Rationale: if a `BaseException` is coming from a OS signal, the event loop would be the one to receive it. So if it's happening in a task, that means user code must be causing it. Which is silly, so we should treat it like any other user error.

These are the possible `BaseException`s that could be raised (based on https://docs.python.org/3/library/exceptions.html#exception-hierarchy):
* `KeyboardInterrupt`: main one I'd worry about. Raised when the process receives SIGINT, by a [`signal`](https://docs.python.org/3/library/signal.html) hander that's registered by default.

    However, if the process actually gets SIGINT'd by the OS, `KeyboardInterrupt` will raised in the main thread (unlike [in C](https://stackoverflow.com/a/22005820/17100540)):
    > Python signal handlers are always executed in the main Python thread of the main interpreter, even if the signal was received in another thread
    https://docs.python.org/3/library/signal.html#signals-and-threadsd

    So if a `KeyboardInterrupt` is coming from a task, I _think_ that means user code is actually calling `raise KeyboardInterrupt`, so in principle it's a user exception.
* `SystemExit`: raised by `sys.exit()`. [`sys.exit()`](https://docs.python.org/3/library/sys.html#sys.exit) is basically just a wrapper around `raise SystemExit`. You could argue this is equivalent to user code intentionally saying `raise SystemExit`, so it's a user exception. (If code within the worker event loop called `sys.exit()`, I don't think the exception would show up within a thread.) Whether or not user code should be able to shut down the whole worker via `sys.exit()` is a reasonable debate.
* `GeneratorExit`: no reason why this should even come out of user code, but it's definitely not a reason to shut down the whole worker.
* Some user-defined subclass of `BaseException` (or plain `BaseException` itself): clearly, this is a user raising this, so we can consider it a user error.
asyncio.CancelledError is also a BaseException and should be considered carefully.
A synchronous task running in a tread raising `CancelledError` would be... odd. But likely not a reason to re-raise `CancelledError` in the main worker event loop.
> A synchronous task running in a tread raising CancelledError would be... odd.

It's not necessarily about the task itself but about the entire code surrounding it. Right now, we have a big block around the entire [`Worker.execute`](https://github.com/dask/distributed/blob/2e98ae97c9f1b7487abb309207f6a455e3c7ed37/distributed/worker.py#L3402-L3519) and handle any `Exception` in this block the same.

`CancelledError` must be dealt with separately. I also think `SystemExit` needs special treatment. I'm not entirely sure about `KeyboardInterrupt`
> It's not necessarily about the task itself but about the entire code surrounding it

I think it is just about the task. We're talking about an error raised by task code within the threadpool here, within `apply_function_simple`: https://github.com/dask/distributed/blob/1b4d13fe4f6898bb825bd9e58dc3e795a7230be5/distributed/worker.py#L4526-L4529

I'd be against special-casing certain `BaseException`s coming from tasks (`KeyboardInterrupt` vs `CancelledError`, etc.). I think we should either propagate _all_ `BaseException`s from the threadpool back to `Worker.execute`, or catch all of them and treat them just like any other exception raised in a task (error the task).

I'm currently leaning towards treating `BaseException`s just like any other user error and catching them. I feel confident that they all indicate users are just doing something weird, not that there's any external signal that the worker should stop. `SystemExit` could be the one exception, but task code _really_ should not be able to expect `sys.exit()` will shut down the whole worker (what if you're using a `ProcessPoolExecutor` instead of threads?).
> I think it is just about the task. We're talking about an error raised by task code within the threadpool here, within apply_function_simple:

yes, my point is that we need the other exception handling in there as well. we previously missed a lot of the code and that also caused us to deadlock

>  (what if you're using a ProcessPoolExecutor instead of threads?).

ProcessPoolExecutor are weird and should not be our primary concern, see

https://github.com/dask/distributed/pull/5128#issue-954126420

> In real world examples if a process gets killed this is typically done by the OS using a signal like SIGINT or SIGTERM and for these situations the behaviour of such a processpool is actually quite different. the processpool ends up completely broken but I don't see anything we can do about this right now. See also https://github.com/dask/distributed/issues/5075

> my point is that we need the other exception handling in there as well

To me this try/except in [`Worker.execute`](https://github.com/dask/distributed/blob/2e98ae97c9f1b7487abb309207f6a455e3c7ed37/distributed/worker.py#L3402-L3519) is a separate topic. I'm surprised it's so broad actually. I get that there are a couple places that could raise errors due to user input (`self.executors[executor]`, `_maybe_deserialize_task`), but I'd prefer to see tighter-scoped try/excepts around those specific places. Having that one big try/except could also incorrectly catch state-machine errors and treat them like user errors.

I wonder if part of the reason we have this giant try/except is because `Worker.executing` only ever gets called via `loop.add_callback` [here](https://github.com/dask/distributed/blob/5c7d5550aa535ad525a842ce7796badec7ff9453/distributed/worker.py#L2349) and [here](https://github.com/dask/distributed/blob/5c7d5550aa535ad525a842ce7796badec7ff9453/distributed/worker.py#L2331). So if something goes wrong, the callback will just silently fail and the exception will be lost to the void of Tornado. These are the [consequences of unstructured concurrency](https://vorpus.org/blog/notes-on-structured-concurrency-or-go-statement-considered-harmful/) 😁 .

Reading the traceback I posted again, I realize this is the actual problem. `await self.loop.run_in_executor` is properly re-raising the `BaseException` back into `Worker.execute`. And `Worker.execute` is properly letting that `BaseException` propagate (just as we'd want to let a `SystemExit` or `KeyboardInterrupt` propagate). The problem is that because we're using Tornado callbacks, there's nothing above `Worker.execute` to handle the error. Tornado just prints that `Exception in callback` happened and moves on as though nothing went wrong. cc @graingert—another reason why Trio is great for not letting you do stuff like this!

-----

For this issue, I still think we should handle `BaseException` within `apply_function_simple` though and treat it as user error. That's kind of what [`concurrent.futures.ThreadPoolExecutor` does](https://github.com/python/cpython/blob/1e5d33e9b9b8631b36f061103a30208b206fd03a/Lib/concurrent/futures/thread.py#L51-L56).

Resolving patch:
diff --git a/distributed/worker.py b/distributed/worker.py
--- a/distributed/worker.py
+++ b/distributed/worker.py
@@ -2293,8 +2293,24 @@ async def execute(self, key: str, *, stimulus_id: str) -> StateMachineEvent:
                     stimulus_id=f"task-finished-{time()}",
                 )
 
-            if isinstance(result["actual-exception"], Reschedule):
+            task_exc = result["actual-exception"]
+            if isinstance(task_exc, Reschedule):
                 return RescheduleEvent(key=ts.key, stimulus_id=f"reschedule-{time()}")
+            if (
+                self.status == Status.closing
+                and isinstance(task_exc, asyncio.CancelledError)
+                and iscoroutinefunction(function)
+            ):
+                # `Worker.cancel` will cause async user tasks to raise `CancelledError`.
+                # Since we cancelled those tasks, we shouldn't treat them as failures.
+                # This is just a heuristic; it's _possible_ the task happened to
+                # fail independently with `CancelledError`.
+                logger.info(
+                    f"Async task {key!r} cancelled during worker close; rescheduling."
+                )
+                return RescheduleEvent(
+                    key=ts.key, stimulus_id=f"cancelled-by-worker-close-{time()}"
+                )
 
             logger.warning(
                 "Compute Failed\n"
@@ -3014,7 +3030,18 @@ def apply_function_simple(
     start = time()
     try:
         result = function(*args, **kwargs)
-    except Exception as e:
+    except (SystemExit, KeyboardInterrupt):
+        # Special-case these, just like asyncio does all over the place. They will pass
+        # through `fail_hard` and `_handle_stimulus_from_task`, and eventually be caught
+        # by special-case logic in asyncio:
+        # https://github.com/python/cpython/blob/v3.9.4/Lib/asyncio/events.py#L81-L82
+        # Any other `BaseException` types would ultimately be ignored by asyncio if
+        # raised here, after messing up the worker state machine along their way.
+        raise
+    except BaseException as e:
+        # Users _shouldn't_ use `BaseException`s, but if they do, we can assume they
+        # aren't a reason to shut down the whole system (since we allow the
+        # system-shutting-down `SystemExit` and `KeyboardInterrupt` to pass through)
         msg = error_message(e)
         msg["op"] = "task-erred"
         msg["actual-exception"] = e
@@ -3050,7 +3077,20 @@ async def apply_function_async(
     start = time()
     try:
         result = await function(*args, **kwargs)
-    except Exception as e:
+    except (SystemExit, KeyboardInterrupt):
+        # Special-case these, just like asyncio does all over the place. They will pass
+        # through `fail_hard` and `_handle_stimulus_from_task`, and eventually be caught
+        # by special-case logic in asyncio:
+        # https://github.com/python/cpython/blob/v3.9.4/Lib/asyncio/events.py#L81-L82
+        # Any other `BaseException` types would ultimately be ignored by asyncio if
+        # raised here, after messing up the worker state machine along their way.
+        raise
+    except BaseException as e:
+        # NOTE: this includes `CancelledError`! Since it's a user task, that's _not_ a
+        # reason to shut down the worker.
+        # Users _shouldn't_ use `BaseException`s, but if they do, we can assume they
+        # aren't a reason to shut down the whole system (since we allow the
+        # system-shutting-down `SystemExit` and `KeyboardInterrupt` to pass through)
         msg = error_message(e)
         msg["op"] = "task-erred"
         msg["actual-exception"] = e


Test patch:
diff --git a/distributed/tests/test_worker.py b/distributed/tests/test_worker.py
--- a/distributed/tests/test_worker.py
+++ b/distributed/tests/test_worker.py
@@ -53,7 +53,7 @@
 )
 from distributed.metrics import time
 from distributed.protocol import pickle
-from distributed.scheduler import Scheduler
+from distributed.scheduler import KilledWorker, Scheduler
 from distributed.utils_test import (
     NO_AMM,
     BlockedExecute,
@@ -421,6 +421,44 @@ def __str__(self):
         assert "Bar" in str(e.__cause__)
 
 
+@pytest.mark.slow
+@pytest.mark.parametrize("sync", [True, False])
+@pytest.mark.parametrize(
+    "exc_type", [BaseException, SystemExit, KeyboardInterrupt, asyncio.CancelledError]
+)
+@gen_cluster(
+    nthreads=[("", 1)],
+    client=True,
+    Worker=Nanny,
+    config={"distributed.scheduler.allowed-failures": 0},
+)
+async def test_base_exception_in_task(c, s, a, sync, exc_type):
+    if sync:
+
+        def raiser():
+            raise exc_type(f"this is a {exc_type}")
+
+    else:
+
+        async def raiser():
+            raise exc_type(f"this is a {exc_type}")
+
+    f = c.submit(raiser)
+
+    try:
+        with pytest.raises(
+            KilledWorker if exc_type in (SystemExit, KeyboardInterrupt) else exc_type
+        ):
+            await f
+    except BaseException as e:
+        # Prevent test failure from killing the whole pytest process
+        traceback.print_exc()
+        pytest.fail(f"BaseException propagated back to test: {e!r}. See stdout.")
+
+    # Nanny restarts it
+    await c.wait_for_workers(1)
+
+
 @gen_test()
 async def test_plugin_exception():
     class MyPlugin:
@@ -1547,7 +1585,7 @@ async def f(ev):
         task for task in asyncio.all_tasks() if "execute(f1)" in task.get_name()
     )
     await a.close()
-    assert task.cancelled()
+    assert task.done()
     assert s.tasks["f1"].state in ("queued", "no-worker")