Repository: dask/distributed

Issue:
Scheduler gather should warn or abort requests if data is too large
By default, `Client.compute` or `Client.gather` will proxy the data fetching over the scheduler via [`Scheduler.gather`](https://github.com/dask/distributed/blob/fca4b35e03ec25e81878c37d8f63837080c0d94f/distributed/scheduler.py#L5894C52-L5941).

Particularly for novice users (but also for the careless or under-caffeinated veterans) an accidental `compute` call can kill the scheduler since the scheduler will collect data without any further checks.

The scheduler knows at this point already how large the requested keys are and could abort the request by raising an appropriate exception instead of fetching the data until it ultimately runs out of memory.

While the end result in both situations is that the computation is lost, the UX is substantially better since the user receives immediate and hopefully actionable feedback to resolve the situation without killing the cluster.

Even for large requests that will make it through the scheduler, I would like to see an informative warning message on the client side. For instance, whenever data exceeds a certain threshold, a warning is issued on the client side informing the user that a fetch of `XGB` is currently in progress.



Discussion:
I thought about this a little more. I think a rather stereotypical issue is something like this

```python
ddf = read_parquet("a_couple_of_TB.parquet")
ddf.compute()
```
which would not be caught because this would generate a graph with a couple read-parquet tasks and a subsequent finalize-* task which concats all dataframes before fetching the data. Only after this finalize task is computed, the Scheduler.gather would kick in. Chances are however high that the worker trying to execute the finalize task dies before this happens.

Implementing the check in gather would still be nice for situations when fetching multiple keys but this is likely much rarer.


However, this also introduces the idea of another interesting case: Why would we even assign `finalize-*` to a worker if we know that this worker cannot hold it's data? We have the memory limit and the nbytes and could abort during transitions with an OOM error. Implementing this will likely be a little harder
Important interactions:
- #7997 
- #7993
So this ticket discusses two different features

1. Issue a warning / exception during Scheduler.gather

```python
class Scheduler:
    def gather(self, ...):
        if sum(ts.nbytes for ts in to_gather) > threshold1(client_memory, worker_memory, scheduler_memory):
             warn.warn("Are you sure?"). # use dask client to log the warning
        if sum(ts.nbytes for ts in to_gather) > threshold2:
             raise Exception

````

2. Detect that we are about to assign a task that is basically guaranteed to fail due to OOM. In this case we should transition the task already to error without even trying (trying would mean trying the task N times (default 3) and eventually it raises a KilledWorkerException)
```python
def decide_worker(...): # ??
     if sum(dts.nbytes for dts in ts.dependencies) > ws.memory_limit * some_factor:
        recommendation[ts] = "error"

```


2) can be further refined later on but for a first iteration I'd keep this simple
3. same as 1, but in Client.gather?
  a. to cover direct=True
  b. in case e.g. your scheduler mounts 32 GiB RAM but your laptop only 8
> 2. can be further refined later on but for a first iteration I'd keep this simple

I think I'd like to create a setting `distributed.worker.memory.max_dependencies_size: 0.7`. I see a lot of value in power users lowering this a lot to debug issues in their reductions. e.g. when the user is convinced that they're working with 128 MiB chunks and they have reduce_every=4.
I never realized, 
- `client.gather(client.compute(collection))` has a finalizer, but
- `collection.compute()` doesn't.

Resolving patch:
diff --git a/distributed/scheduler.py b/distributed/scheduler.py
--- a/distributed/scheduler.py
+++ b/distributed/scheduler.py
@@ -2078,7 +2078,6 @@ def transition_released_waiting(self, key: str, stimulus_id: str) -> RecsMsgs:
 
     def transition_no_worker_processing(self, key: str, stimulus_id: str) -> RecsMsgs:
         ts = self.tasks[key]
-        worker_msgs: Msgs = {}
 
         if self.validate:
             assert not ts.actor, f"Actors can't be in `no-worker`: {ts}"
@@ -2086,10 +2085,10 @@ def transition_no_worker_processing(self, key: str, stimulus_id: str) -> RecsMsg
 
         if ws := self.decide_worker_non_rootish(ts):
             self.unrunnable.discard(ts)
-            worker_msgs = self._add_to_processing(ts, ws)
+            return self._add_to_processing(ts, ws, stimulus_id=stimulus_id)
         # If no worker, task just stays in `no-worker`
 
-        return {}, {}, worker_msgs
+        return {}, {}, {}
 
     def decide_worker_rootish_queuing_disabled(
         self, ts: TaskState
@@ -2295,8 +2294,7 @@ def transition_waiting_processing(self, key: str, stimulus_id: str) -> RecsMsgs:
             if not (ws := self.decide_worker_non_rootish(ts)):
                 return {ts.key: "no-worker"}, {}, {}
 
-        worker_msgs = self._add_to_processing(ts, ws)
-        return {}, {}, worker_msgs
+        return self._add_to_processing(ts, ws, stimulus_id=stimulus_id)
 
     def transition_waiting_memory(
         self,
@@ -2751,8 +2749,6 @@ def transition_queued_released(self, key: str, stimulus_id: str) -> RecsMsgs:
 
     def transition_queued_processing(self, key: str, stimulus_id: str) -> RecsMsgs:
         ts = self.tasks[key]
-        recommendations: Recs = {}
-        worker_msgs: Msgs = {}
 
         if self.validate:
             assert not ts.actor, f"Actors can't be queued: {ts}"
@@ -2760,10 +2756,9 @@ def transition_queued_processing(self, key: str, stimulus_id: str) -> RecsMsgs:
 
         if ws := self.decide_worker_rootish_queuing_enabled():
             self.queued.discard(ts)
-            worker_msgs = self._add_to_processing(ts, ws)
+            return self._add_to_processing(ts, ws, stimulus_id=stimulus_id)
         # If no worker, task just stays `queued`
-
-        return recommendations, {}, worker_msgs
+        return {}, {}, {}
 
     def _remove_key(self, key: str) -> None:
         ts = self.tasks.pop(key)
@@ -3144,7 +3139,9 @@ def _validate_ready(self, ts: TaskState) -> None:
         assert ts not in self.queued
         assert all(dts.who_has for dts in ts.dependencies)
 
-    def _add_to_processing(self, ts: TaskState, ws: WorkerState) -> Msgs:
+    def _add_to_processing(
+        self, ts: TaskState, ws: WorkerState, stimulus_id: str
+    ) -> RecsMsgs:
         """Set a task as processing on a worker and return the worker messages to send"""
         if self.validate:
             self._validate_ready(ts)
@@ -3161,7 +3158,45 @@ def _add_to_processing(self, ts: TaskState, ws: WorkerState) -> Msgs:
         if ts.actor:
             ws.actors.add(ts)
 
-        return {ws.address: [self._task_to_msg(ts)]}
+        ndep_bytes = sum(dts.nbytes for dts in ts.dependencies)
+        if (
+            ws.memory_limit
+            and ndep_bytes > ws.memory_limit
+            and dask.config.get("distributed.worker.memory.terminate")
+        ):
+            # Note
+            # ----
+            # This is a crude safety system, only meant to prevent order-of-magnitude
+            # fat-finger errors.
+            #
+            # For collection finalizers and in general most concat operations, it takes
+            # a lot less to kill off the worker; you'll just need
+            # ndep_bytes * 2 > ws.memory_limit * terminate threshold.
+            #
+            # In heterogeneous clusters with workers mounting different amounts of
+            # memory, the user is expected to manually set host/worker/resource
+            # restrictions.
+            msg = (
+                f"Task {ts.key} has {format_bytes(ndep_bytes)} worth of input "
+                f"dependencies, but worker {ws.address} has memory_limit set to "
+                f"{format_bytes(ws.memory_limit)}."
+            )
+            if ts.prefix.name == "finalize":
+                msg += (
+                    " It seems like you called client.compute() on a huge collection. "
+                    "Consider writing to distributed storage or slicing/reducing first."
+                )
+            logger.error(msg)
+            return self._transition(
+                ts.key,
+                "erred",
+                exception=pickle.dumps(MemoryError(msg)),
+                cause=ts.key,
+                stimulus_id=stimulus_id,
+                worker=ws.address,
+            )
+
+        return {}, {}, {ws.address: [self._task_to_msg(ts)]}
 
     def _exit_processing_common(self, ts: TaskState) -> WorkerState | None:
         """Remove *ts* from the set of processing tasks.


Test patch:
diff --git a/distributed/tests/test_scheduler.py b/distributed/tests/test_scheduler.py
--- a/distributed/tests/test_scheduler.py
+++ b/distributed/tests/test_scheduler.py
@@ -7,6 +7,7 @@
 import math
 import operator
 import pickle
+import random
 import re
 import sys
 from collections.abc import Collection
@@ -22,7 +23,7 @@
 from tornado.ioloop import IOLoop
 
 import dask
-from dask import delayed
+from dask import bag, delayed
 from dask.core import flatten
 from dask.highlevelgraph import HighLevelGraph, MaterializedLayer
 from dask.utils import parse_timedelta, tmpfile, typename
@@ -4472,3 +4473,39 @@ async def test_scatter_creates_ts(c, s, a, b):
         await a.close()
         assert await x2 == 2
     assert s.tasks["x"].run_spec is not None
+
+
+@pytest.mark.parametrize("finalize", [False, True])
+@gen_cluster(
+    client=True,
+    nthreads=[("", 1)] * 4,
+    worker_kwargs={"memory_limit": "100 kB"},
+    config={
+        "distributed.worker.memory.target": False,
+        "distributed.worker.memory.spill": False,
+        "distributed.worker.memory.pause": False,
+    },
+)
+async def test_refuse_to_schedule_huge_task(c, s, *workers, finalize):
+    """If the total size of a task's input grossly exceed the memory available on the
+    worker, the scheduler must refuse to compute it
+    """
+    bg = bag.from_sequence(
+        [random.randbytes(30_000) for _ in range(4)],
+        npartitions=4,
+    )
+    match = r"worth of input dependencies, but worker .* has memory_limit set to"
+    if finalize:
+        fut = c.compute(bg)
+        match += r".* you called client.compute()"
+    else:
+        bg = bg.repartition(npartitions=1).persist()
+        fut = list(c.futures_of(bg))[0]
+
+    with pytest.raises(MemoryError, match=match):
+        await fut
+
+    # The task never reached the workers
+    for w in workers:
+        for ev in w.state.log:
+            assert fut.key not in ev
diff --git a/distributed/tests/test_worker_memory.py b/distributed/tests/test_worker_memory.py
--- a/distributed/tests/test_worker_memory.py
+++ b/distributed/tests/test_worker_memory.py
@@ -637,6 +637,7 @@ def f(ev):
             "distributed.worker.memory.target": False,
             "distributed.worker.memory.spill": False,
             "distributed.worker.memory.pause": False,
+            "distributed.worker.memory.terminate": False,
         },
     ),
 )