dask / dask/distributed

Task co-assignment logic is worst-case for binary operations like `a + b`

Open
#6,597 1 comment 0 reactions 0 assignees View on GitHub
bug memory performance scheduling
Dominant language
Python
Stars
1.7k
Forks
778
Avg merge
2h 50m
Merged PRs (30d)
3

Description

The root task co-assignment logic does the exact opposite of what it should for operations combining two different datasets, like `a + b`.

```
x x x x
/\ /\ /\ /\
a b a b a b a b
1 2 3 4 5 6 7 8 <-- priority
```

It assigns all the `a`s to one worker, and all the `b`s to another. Each `x` then requires transferring an `a` or a `b`. So 50% of the data gets transferred. This could have been 0% if we had co-assigned properly.

The reason for this is that co-assignment selects a worker to re-use _per task group_. So it goes something like (recall that we're iterating through root tasks in priority order):
1. Assign `a1`. It has no `last_worker` set, so pick the least busy worker: `w1`.
2. Assign `b1`. It has no `last_worker` set, so pick the least busy worker. `w1` already has a task assigned to it (`a1`), so we pick `w2`.
3. Assign the next 3 `a`s to `w1`, and the next 3 `b`s to `w2` (they come through interleaved, since they're interleaved in priority order)
4. Time to pick a new worker for `a5`. They're both equally busy; say we pick `w2`.
5. Time to pick a new worker for `b5`. We just made `w2` slightly busier than `w1`, so pick `w1`.
6. Pattern continues. Each time we flip-flop, sending the tasks to opposite workers

The last-used worker should be global state (well, global to a particular sequence of `transitions` caused by `update_graph`). Each subsequent task in priority should re-use this worker until it's filled up, regardless of what task group the task belongs to.

The tricky part is calculating what "filled up" means. We currently use the size of the task group to decide how many root tasks in total there are, which we then divide by nthreads to decide how many to assign per worker. But of course, that's not actually the total number of root tasks. I'm not sure yet how to figure out the total number of root tasks in constant time within `decide_worker`.

Broadly speaking, this stateful and kinda hacky co-assignment logic is a bit of a pain to integrate into https://github.com/dask/distributed/issues/6560. I've been able to do it, but maintaining good assignment while rebalancing tasks when adding and removing workers is difficult. Our co-assignment logic is too reliant on statefulness and getting to iterate through all the tasks at once in priority order, we can't actually re-co-assign things when workers change. If we had a data structure/mechanism to efficiently identify "which tasks are siblings of this one", or maybe even "which worker holds the task nearest in priority to this one", it might make solving both problems easier.

----

As a simple test that fails on main (each worker has transferred 4 keys):
```python
@gen_cluster(
client=True,
nthreads=[("", 1), ("", 1)],
)
async def test_decide_worker_coschedule_order_binary_op(c, s, a, b):
xs = [delayed(i, name=f"x-{i}") for i in range(8)]
ys = [delayed(i, name=f"y-{i}") for i in range(8)]
zs = [x + y for x, y in zip(xs, ys)]

await c.gather(c.compute(zs))

assert not a.transfer_incoming_log, [l["keys"] for l in a.transfer_incoming_log]
assert not b.transfer_incoming_log, [l["keys"] for l in b.transfer_incoming_log]
```

Note that this case occurs in @TomNicholas's example workload: https://github.com/dask/distributed/issues/6571

cc @fjetter @mrocklin

Contributor guide

Open the contributing guide

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.