Repository: dask/distributed

Issue:
Co-assign neighboring tasks to neighboring workers
**Edit:** This proposal is now defunct(ish).  Read below for updated proposal.


Sometimes we assign tasks that are closely related to different machines, resulting in unnecessary communication.  The most simple case is when we have a couple of siblings

```
   c
  / \
 /   \
a     b
```

Because neither `a` nor `b` have dependencies, they are assigned somewhat randomly onto workers.  However because `c` requires them both one or the other will have to move across the network.  This can cause excess communication which can bog down the system.  

This has come up several times, notably in https://github.com/dask/distributed/issues/2602 and https://github.com/ocean-transport/coiled_collaboration/issues/3 and https://github.com/dask/distributed/pull/4864 .  There are a few different things that we can do to solve these issues.  This issue contains one part of a solution, assigning sibling tasks to the same worker.  

### Straightforward (but slow) solution

@fjetter provided a solution here https://github.com/dask/distributed/pull/4864#issuecomment-855228497

which involves iterating up through all dependents of a task, and then back down through all of those tasks dependencies, in order to find similar machines.  This doesn't work in the general case (we have to navigate through k-squared links, where k is the number of dependencies/dependents) but it might be ok in a restricted case. 

### Restricted case

In particular we seem to be interested in the following conditions:

1.  The task has no dependencies (otherwise we would schedule wherever those are)
2.  The task has very few dependents, possibly only one
3.  That dependent has very few dependencies, maybe less than five? (ten?)
4.  There are many more tasks than workers, so that we're comfortable losing some concurrency here
5.  The task we're thinking about is sizeable, and not something that we would want to move, relative to it's computation time

I suspect that there is an order by which we could check these conditions that is fast in all of the cases where it matters.  Additionally, if we have several dependencies to check then we might consider just accepting the location of the very first dependency that is already assigned, and skip checking the rest.  



Discussion:
I hit send on this too early.  I've updated the original comment.
A (failed) attempt at something like this in https://github.com/dask/distributed/pull/2940. I didn't post it, but I vaguely recall that approach really slowing down the scheduling.
Yeah, the implementation you had hit the full `N^2` case.  See https://github.com/dask/distributed/pull/2940/files#r311737887

I'm hopeful that if we are able to restrict things suitably that we can avoid significant slowdown 
Conversely though, the question of the approach in this issue is whether or not this set of constraints will let us actually solve important problems.  It could be that although these constraints make things fast, they exclude important cases.

For example, this wouldn't work with a rechunk situation (although, to be fair, that graph is fully connected in just two layers, and so not a case where we would want to schedule on the same node as our sibling). 
OK, @gjoseph92 and I were talking about this for a while.  I'll try to summarize some of the thoughts that went around there. 

First, Gabe brought up the same concern I brought up before with Tom around checking for siblings, namely uncles/aunts, cousins, second cousins twice removed, and so on.  Graphs get messy and so restrictive rules are less fun.

We then went on a journey thinking about switching decide_worker behavior in the case where we are assigning many tasks (where many is defined relative to the number of workers/cores that we have).  If we're in this situation, and if we assume that tasks are being cycled through in the order determined by `dask.order` (which I *think* should be true?) then we might do better by adding a bit of inertia to our choice of which worker to decide.  This would be a change in our policy:

-  **Today**: we choose the worker with least occupancy.  This results in tasks that are neighbors in terms of dask.order being very likely to be on *different* machines.  This is good if the graph is generic/asymmetric, because the task ordered 0 and the task ordered 1 are the two most important tasks to start, and so they should both run as soon as possible
- **Proposed**: we try to allocate batches of neighboring tasks (neighboring in dask.order space) to similar workers.  For example if we have 1000 tasks and 10 workers then we might allocate tasks ordered 0..99 to worker 1 and tasks ordered 100..199 to worker 2.  This is good if the graph is highly symmetric, such as is the case with I/O tasks that show up in update_graph sometimes.  This is good because task 0 and task 1 are very likely to have shared dependents (directly or transitively (grandparents))

To be clear, the proposed behavior, of bulk/batch assignment is *bad* in many cases, but *good* in some others.   How do we identify the situation?  

First we thought about identifying this in update_graph, by looking for a case where we had many tasks with no dependencies.   Unfortunately, this breaks in situations like `from_zarr` where there is a single dependency claimed by thousands of dependents.  It's also somewhat odd to change a policy in update_graph. 

### Where we ended up

So we still want to support a batch/bulk task assignment policy (where neighboring tasks are assigned to the same worker, rather than spread around), and we think that we can identify this situation deep within `decide_worker` by comparing the size of the task state's group with the number of threads we have going and checking that the number of dependencies of that group is quite small.  

If those hold then we want to glom together neighboring tasks rather than spread them apart.

One way to achieve that glomming together would be for each TaskGroup to keep track of the most recent worker to which state of that group was assigned.  Then, in `decide_worker` we would figure out whatever worker we would normally choose, and then if it's not much better than the previously assigned worker we would use the previously assigned worker.  

```python
ws = decide_worker(ts, ...)
if not_much_better(ws, ts.group._last_assigned_worker):
    ws = ts.group._last_assigned_worker
else:
    ts.group._last_assigned_worker = ws
```

The definition of "much better" here would have to depend on the ratio of tasks in that group to cores in the cluster.  This ensures that we still respect general load allocation (if a worker is busy for other reasons then it won't be used here) while still providing some inertia that we only shift workers somewhat slowly.

cc @fjetter 
I don't mind the batched assignment and think this is in general a good idea as long as it requires minimal state and clean APIs. I wouldn't want to have something like `Worker.gather_dep / Worker.select_keys_for_gather` again. imho, that feels overly complicated. If we can pull this off with minimal state and clean APIs, that's great. However, our entire state machine is currently working using individual task transitions which is why batched assignments do not fit nicely to the concept. The `_last_assigned_worker` _might_ be just enough to pull this off, idk. 

To make this somewhat stateless, one could structure the batch assignment more explicitly, similar to the worker for execution/fetch but a bit more explicit. We'd introduce a new state `Assigning` for tasks and the transition path is then `waiting->Assigning->Processing` where the two transition would not do much. However, this intermediare state would give us a transition barrier to accumulate tasks and would mark a natural transition between single-task state machine and bulk processing.

```python
def transition_waiting_assigned(self, ts: TaskState):
    self.to_be_assigned.add(ts)
    ts.state = "assigning"

def transition_assigned_processing(self, ts: TaskState, ws: WorkerState):
    ts.processing_on = ws
    submit_task_to_worker(ts, ws)
    ...

def periodic_callback_assign_worker(self):
    assignments = decide_worker(self.to_be_assigned, self.workers.values())
    for ts, ws in assignments:
        transition_assigned_processing(ts, ws)

def decide_worker(tasks: Set[TaskState], workers: Set[WorkerState]) -> List[Tuple[TaskState, WorkerState]]:
    # This function performs the actual optimization problem given the above input and this input alone (please do not mutate anything here ;)
    # We do not keep any additional state but only optimize on the provided input information
    # Input, e.g. allowed workers could be further reduced by applying filters up front (based on occupancy / take n smallest, allowed resources, etc.)
    # Optimization could be task group based, account for occupancy, apply some generic weighted 
    # bin packing or just magic but would not require additional state, e.g. `_last_assigned_worker` would be local variable for this decision
```

I'm sure we can find something where we can leverage synergies if we're optimizing a bunch of tasks at the same time. I would argue this can be done in `O(#Tasks * #Dependents)` which would be sufficient, wouldn't it?

we might also want to peak towards our most recent rebalancing algorithm which, I believe, is a related problem, algorithmically speaking

Why I believe we can pull this off without holding on to state is because I think most of the assignments where this matters, or at least the assignments where I suspect the biggest impact, happen during the same stimulus.

----

Taking a step back, this is somehow what I would envision to be solved by the speculative task assignment. if we had already assigned the dependencies, the decision should become simpler. I'm wondering if we should try to investigate this approach since it is also not straight forward how this speculative assignment would work. fwiw, we could implement the speculative assignment on scheduler side without actually implementing it on the worker. for this use case it would still give us benefits and would allow us to implement this incrementally, letting us deal with stealing and worker states at a later point in time.

I don't think that code complexity will be a big issue here.  I think that we can accomplish this by storing the most recently assigned worker, and making a ten-line change to decide worker


```python
def decide_worker(...):
    ws = ...  # current best guess for worker

    # Group is large relative to cluster size and there are few dependencies
    if len(ts._group) > self._total_nthreads and sum(map(len, ts._group.dependencies)) < 5:
        old = ts._group._last_worker  # what we chose last time
        ratio = math.ceil(len(ts._group) / self._total_nthreads)  # how many tasks we expect per worker

        # Can we still fit a few more of these tasks on the old worker?
        if ts._prefix._duration_average + old.occupancy < ws.occupancy + ratio * ts._prefix._duration_average:
            ws = old
        else: 
            ts._group._last_worker = ws
    ...
```

That's probably optimistic, but I think that something similar might work
@gjoseph92 my understanding is that you're trying this path out.  Is that correct?
It's also worth noting that all of this is built on the assumption that we're iterating over tasks in a way that respects dask.order.  It might be worth verifying that assumption at some point.  (or not, if this is very easy to test)
@mrocklin yes, I'm working on this.

I think that a clever metric for when to switch to a new worker (like `ts._prefix._duration_average + old.occupancy < ws.occupancy + ratio * ts._prefix._duration_average` in your example) might not be better than just counting how many tasks have been assigned to `last_worker`, and when that's exceeded our `ratio`, finding a new worker. If we trust that our initial checks prove we're in a situation with lots of root-like tasks which saturate the cluster (and that we're indeed being called in priority order), then occupancy doesn't really matter—we're going to have to schedule all the tasks anyway, and the workers will all be over-saturated with tasks. Especially we're still learning `_duration_average` (say, on the second task submission), trying to estimate occupancy seems likely to do something we don't expect.

> all of this is built on the assumption that we're iterating over tasks in a way that respects dask.order

The main way I think this will fail is if the tasks aren't strictly root tasks, but have a few small dependencies (like zarr arrays). If not all tasks require all those dependencies, then `decide_workers` will be called in whatever order those dependencies finish, which may not follow priority.
> I think that a clever metric for when to switch to a new worker (like ts._prefix._duration_average + old.occupancy < ws.occupancy + ratio * ts._prefix._duration_average in your example) might not be better than just counting how many tasks have been assigned to last_worker, and when that's exceeded our ratio, finding a new worker

I don't disagree.  However then we would need to keep track of how many tasks we've assigned to each worker.  We're *already* doing this implicitly with occupancy, so I'm reusing that counter by multiplying by the duration_average (which is what gets added to occupancy each time we add a task).

I think that what I wrote down achieves the result that you want, but tracks less new things.
Regardless, I think that the thing to do now is to try this out and see how it performs.  I'm hopeful that this is a small enough change that running this experiment takes a small number of hours.
> I think that what I wrote down achieves the result that you want, but tracks less new things.

Unfortunately `ts._prefix._duration_average` is -1 for all the root `sum` tasks in the example from #4864. I'd expect this to be the case in general, since we should be scheduling root tasks much faster than they complete—when we're scheduling the second root task moments after scheduling the first, we wouldn't have any timing information about that task prefix yet.
Here are initial results from #4899. I'm using a `_last_worker_tasks_left` field on the TaskGroup to keep track of when to switch to a new worker for the reason above.

Performance report: https://cloud.coiled.io/gjoseph92/reports/16. The task stream looks great; no transfers just like in #4864.

I'm confused why worker memory usage increases like this, as though tasks aren't being released from memory:

![Screen Shot 2021-06-09 at 2 41 53 PM](https://user-images.githubusercontent.com/3309802/121429029-b27d4000-c933-11eb-9586-a913e9a1d93a.png)

The dashboard indicates tasks are being released as usual though:

![Screen Shot 2021-06-09 at 3 12 38 PM](https://user-images.githubusercontent.com/3309802/121430692-a4c8ba00-c935-11eb-8a73-c6a4211401e0.png)

Those look like excellent initial results?  :) 
I'm curious if any of the Pangeo workloads benefit from this as well .  Is it too early to ping Julius?
I think it is. I'll try running some Pangeo cases once this is further along, but I think this monotonically-increasing-memory issue is important to track down first.
> I'm confused why worker memory usage increases like this, as though tasks aren't being released from memory:

This feels like an orthogonal problem to me.  We can see that tasks are being assigned to workers well.  Extra memory sticking around seems to be a lower level issue to me.
Using jemalloc on macOS (with `DYLD_INSERT_LIBRARIES=$(brew --prefix jemalloc)/lib/libjemalloc.dylib mprof run --multiprocess test.py`) gives this—so yup, it was the same malloc issue as usual:

![jemalloc](https://user-images.githubusercontent.com/3309802/121445141-17de2a80-c94e-11eb-9440-9ff1180d5746.png)


Resolving patch:
diff --git a/distributed/scheduler.py b/distributed/scheduler.py
--- a/distributed/scheduler.py
+++ b/distributed/scheduler.py
@@ -921,6 +921,16 @@ class TaskGroup:
 
        The result types of this TaskGroup
 
+    .. attribute:: last_worker: WorkerState
+
+       The worker most recently assigned a task from this group, or None when the group
+       is not identified to be root-like by `SchedulerState.decide_worker`.
+
+    .. attribute:: last_worker_tasks_left: int
+
+       If `last_worker` is not None, the number of times that worker should be assigned
+       subsequent tasks until a new worker is chosen.
+
     See also
     --------
     TaskPrefix
@@ -936,6 +946,8 @@ class TaskGroup:
     _start: double
     _stop: double
     _all_durations: object
+    _last_worker: WorkerState
+    _last_worker_tasks_left: Py_ssize_t
 
     def __init__(self, name: str):
         self._name = name
@@ -949,6 +961,8 @@ def __init__(self, name: str):
         self._start = 0.0
         self._stop = 0.0
         self._all_durations = defaultdict(float)
+        self._last_worker = None
+        self._last_worker_tasks_left = 0
 
     @property
     def name(self):
@@ -990,6 +1004,14 @@ def start(self):
     def stop(self):
         return self._stop
 
+    @property
+    def last_worker(self):
+        return self._last_worker
+
+    @property
+    def last_worker_tasks_left(self):
+        return self._last_worker_tasks_left
+
     @ccall
     def add(self, o):
         ts: TaskState = o
@@ -2309,21 +2331,60 @@ def transition_no_worker_waiting(self, key):
     @exceptval(check=False)
     def decide_worker(self, ts: TaskState) -> WorkerState:
         """
-        Decide on a worker for task *ts*.  Return a WorkerState.
+        Decide on a worker for task *ts*. Return a WorkerState.
+
+        If it's a root or root-like task, we place it with its relatives to
+        reduce future data tansfer.
+
+        If it has dependencies or restrictions, we use
+        `decide_worker_from_deps_and_restrictions`.
+
+        Otherwise, we pick the least occupied worker, or pick from all workers
+        in a round-robin fashion.
         """
+        if not self._workers_dv:
+            return None
+
         ws: WorkerState = None
+        group: TaskGroup = ts._group
         valid_workers: set = self.valid_workers(ts)
 
         if (
             valid_workers is not None
             and not valid_workers
             and not ts._loose_restrictions
-            and self._workers_dv
         ):
             self._unrunnable.add(ts)
             ts.state = "no-worker"
             return ws
 
+        # Group is larger than cluster with few dependencies? Minimize future data transfers.
+        if (
+            valid_workers is None
+            and len(group) > self._total_nthreads * 2
+            and sum(map(len, group._dependencies)) < 5
+        ):
+            ws: WorkerState = group._last_worker
+
+            if not (
+                ws and group._last_worker_tasks_left and ws._address in self._workers_dv
+            ):
+                # Last-used worker is full or unknown; pick a new worker for the next few tasks
+                ws = min(
+                    (self._idle_dv or self._workers_dv).values(),
+                    key=partial(self.worker_objective, ts),
+                )
+                group._last_worker_tasks_left = math.floor(
+                    (len(group) / self._total_nthreads) * ws._nthreads
+                )
+
+            # Record `last_worker`, or clear it on the final task
+            group._last_worker = (
+                ws if group.states["released"] + group.states["waiting"] > 1 else None
+            )
+            group._last_worker_tasks_left -= 1
+            return ws
+
         if ts._dependencies or valid_workers is not None:
             ws = decide_worker(
                 ts,
@@ -2332,6 +2393,7 @@ def decide_worker(self, ts: TaskState) -> WorkerState:
                 partial(self.worker_objective, ts),
             )
         else:
+            # Fastpath when there are no related tasks or restrictions
             worker_pool = self._idle or self._workers
             worker_pool_dv = cast(dict, worker_pool)
             wp_vals = worker_pool.values()


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
@@ -17,7 +17,7 @@
 
 import dask
 from dask import delayed
-from dask.utils import apply, parse_timedelta
+from dask.utils import apply, parse_timedelta, stringify
 
 from distributed import Client, Nanny, Worker, fire_and_forget, wait
 from distributed.comm import Comm
@@ -126,6 +126,114 @@ async def test_decide_worker_with_restrictions(client, s, a, b, c):
     assert x.key in a.data or x.key in b.data
 
 
+@pytest.mark.parametrize("ndeps", [0, 1, 4])
+@pytest.mark.parametrize(
+    "nthreads",
+    [
+        [("127.0.0.1", 1)] * 5,
+        [("127.0.0.1", 3), ("127.0.0.1", 2), ("127.0.0.1", 1)],
+    ],
+)
+def test_decide_worker_coschedule_order_neighbors(ndeps, nthreads):
+    @gen_cluster(
+        client=True,
+        nthreads=nthreads,
+        config={"distributed.scheduler.work-stealing": False},
+    )
+    async def test(c, s, *workers):
+        r"""
+        Ensure that sibling root tasks are scheduled to the same node, reducing future data transfer.
+
+        We generate a wide layer of "root" tasks (random NumPy arrays). All of those tasks share 0-5
+        trivial dependencies. The ``ndeps=0`` and ``ndeps=1`` cases are most common in real-world use
+        (``ndeps=1`` is basically ``da.from_array(..., inline_array=False)`` or ``da.from_zarr``).
+        The graph is structured like this (though the number of tasks and workers is different):
+
+            |-W1-|  |-W2-| |-W3-|  |-W4-|   < ---- ideal task scheduling
+
+              q       r       s       t      < --- `sum-aggregate-`
+             / \     / \     / \     / \
+            i   j   k   l   m   n   o   p    < --- `sum-`
+            |   |   |   |   |   |   |   |
+            a   b   c   d   e   f   g   h    < --- `random-`
+            \   \   \   |   |   /   /   /
+                   TRIVIAL * 0..5
+
+        Neighboring `random-` tasks should be scheduled on the same worker. We test that generally,
+        only one worker holds each row of the array, that the `random-` tasks are never transferred,
+        and that there are few transfers overall.
+        """
+        da = pytest.importorskip("dask.array")
+        np = pytest.importorskip("numpy")
+
+        if ndeps == 0:
+            x = da.random.random((100, 100), chunks=(10, 10))
+        else:
+
+            def random(**kwargs):
+                assert len(kwargs) == ndeps
+                return np.random.random((10, 10))
+
+            trivial_deps = {f"k{i}": delayed(object()) for i in range(ndeps)}
+
+            # TODO is there a simpler (non-blockwise) way to make this sort of graph?
+            x = da.blockwise(
+                random,
+                "yx",
+                new_axes={"y": (10,) * 10, "x": (10,) * 10},
+                dtype=float,
+                **trivial_deps,
+            )
+
+        xx, xsum = dask.persist(x, x.sum(axis=1, split_every=20))
+        await xsum
+
+        # Check that each chunk-row of the array is (mostly) stored on the same worker
+        primary_worker_key_fractions = []
+        secondary_worker_key_fractions = []
+        for i, keys in enumerate(x.__dask_keys__()):
+            # Iterate along rows of the array.
+            keys = set(stringify(k) for k in keys)
+
+            # No more than 2 workers should have any keys
+            assert sum(any(k in w.data for k in keys) for w in workers) <= 2
+
+            # What fraction of the keys for this row does each worker hold?
+            key_fractions = [
+                len(set(w.data).intersection(keys)) / len(keys) for w in workers
+            ]
+            key_fractions.sort()
+            # Primary worker: holds the highest percentage of keys
+            # Secondary worker: holds the second highest percentage of keys
+            primary_worker_key_fractions.append(key_fractions[-1])
+            secondary_worker_key_fractions.append(key_fractions[-2])
+
+        # There may be one or two rows that were poorly split across workers,
+        # but the vast majority of rows should only be on one worker.
+        assert np.mean(primary_worker_key_fractions) >= 0.9
+        assert np.median(primary_worker_key_fractions) == 1.0
+        assert np.mean(secondary_worker_key_fractions) <= 0.1
+        assert np.median(secondary_worker_key_fractions) == 0.0
+
+        # Check that there were few transfers
+        unexpected_transfers = []
+        for worker in workers:
+            for log in worker.incoming_transfer_log:
+                keys = log["keys"]
+                # The root-ish tasks should never be transferred
+                assert not any(k.startswith("random") for k in keys), keys
+                # `object-` keys (the trivial deps of the root random tasks) should be transferred
+                if any(not k.startswith("object") for k in keys):
+                    # But not many other things should be
+                    unexpected_transfers.append(list(keys))
+
+        # A transfer at the very end to move aggregated results is fine (necessary with unbalanced workers in fact),
+        # but generally there should be very very few transfers.
+        assert len(unexpected_transfers) <= 3, unexpected_transfers
+
+    test()
+
+
 @gen_cluster(client=True, nthreads=[("127.0.0.1", 1)] * 3)
 async def test_move_data_over_break_restrictions(client, s, a, b, c):
     [x] = await client.scatter([1], workers=b.address)
diff --git a/distributed/tests/test_steal.py b/distributed/tests/test_steal.py
--- a/distributed/tests/test_steal.py
+++ b/distributed/tests/test_steal.py
@@ -146,21 +146,18 @@ async def test_steal_related_tasks(e, s, a, b, c):
 
 @gen_cluster(client=True, nthreads=[("127.0.0.1", 1)] * 10, timeout=1000)
 async def test_dont_steal_fast_tasks_compute_time(c, s, *workers):
-    np = pytest.importorskip("numpy")
-    x = c.submit(np.random.random, 10000000, workers=workers[0].address)
-
     def do_nothing(x, y=None):
         pass
 
-    # execute and measure runtime once
-    await wait(c.submit(do_nothing, 1))
+    xs = c.map(do_nothing, range(10), workers=workers[0].address)
+    await wait(xs)
 
-    futures = c.map(do_nothing, range(1000), y=x)
+    futures = c.map(do_nothing, range(1000), y=xs)
 
     await wait(futures)
 
-    assert len(s.who_has[x.key]) == 1
-    assert len(s.has_what[workers[0].address]) == 1001
+    assert len(set.union(*(s.who_has[x.key] for x in xs))) == 1
+    assert len(s.has_what[workers[0].address]) == len(xs) + len(futures)
 
 
 @gen_cluster(client=True)
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
@@ -1590,12 +1590,12 @@ async def test_lifetime(cleanup):
     async with Scheduler() as s:
         async with Worker(s.address) as a, Worker(s.address, lifetime="1 seconds") as b:
             async with Client(s.address, asynchronous=True) as c:
-                futures = c.map(slowinc, range(200), delay=0.1)
+                futures = c.map(slowinc, range(200), delay=0.1, worker=[b.address])
                 await asyncio.sleep(1.5)
                 assert b.status != Status.running
                 await b.finished()
 
-                assert set(b.data).issubset(a.data)  # successfully moved data over
+                assert set(b.data) == set(a.data)  # successfully moved data over
 
 
 @gen_cluster(client=True, worker_kwargs={"lifetime": "10s", "lifetime_stagger": "2s"})