P2P shuffling and queuing combined may cause high memory usage with `dask.dataframe.merge`

[hendrikmakait]

When P2P shuffling and queueing are used together in a workload where two dataframes are merged and both dataframes are generated by rootish tasks, they can cause unnecessarily high memory usage and spilling:

Queuing enabled:

Queuing disabled:

What I expect to happen
We should first finish all transfers of one shuffle as well as its barrier, followed by a single unpack and then the transfers of the second shuffle and the second barrier as prescribed by the topological sort. Once both shuffle barriers are finished, the individual output pairs of the shuffle get read from disk, merged and handed to downstream tasks.

This is what we see without queueing.

What happens instead
All upstream tasks of both shuffles are queued. This causes the scheduler to finish one shuffle transfer followed by its barrier and immediately materialize all of its outputs. The scheduler is forced to ignore the topological sort since the inputs of the second transfer are still queued and therefore ignored while other tasks - the output tasks of the shuffle - are available for execution.

With larger-than-cluster-memory data, this can cause significant spilling, excessive disk I/O, and consequently longer runtime.

Reproducer
The plots above were generated by running test_join_big from coiled/benchmarks#645 with a size factor of 1x the cluster memory using a cluster of 10 workers. Any workload generating two dataframes using rootish tasks (e.g. by generating random partitions or reading from Parquet) and then merging those two should work.

from dask.datasets import timeseries
with dask.config.set(shuffle="p2p"):
    df1_big = timeseries(
        "2000-01-01", 
        "2001-01-01", 
        dtypes={str(i): float for i in range(100)}
    )
    df1_big["predicate"] = df1_big["0"] * 1e9
    df1_big = df1_big.astype({"predicate": "int"})

    df2_big = timeseries(
        "2000-01-01", 
        "2001-01-01", 
        dtypes={str(i): float for i in range(100)}
    )

    df2_big["predicate"] = df2_big["0"] * 1e9
    df2_big = df2_big.astype({"predicate": "int"})

    join = dd.merge(df1_big, df2_big, on="predicate", how="inner")
    await client.compute(join.size)

[fjetter]

Edit: Sorry, I believe I misread your comment. I thought you were suggesting that the topological sort is wrong

one after another as prescribed by the topological sort.

This is not what the topological sort suggests. Given two simple dataframes

df = pd.DataFrame({
    "A":[2001,2002,2003],
    "B":['Apple','Banana','Coconut']
})
ddf1 = dd.from_pandas(df, 2)
ddf2 = dd.from_pandas(df, 3)
# This join is nonsense, of course but dask doesn't know that ;)
join = dd.merge(ddf1, ddf2, how="inner", shuffle='p2p')

we can inspect the join (I used different partition sizes to allow us to distinguish the DFs in the graph
This graph is merely for reference. It shows how we generate two barriers, etc.

If we look at the ordering, we can clearly see that the first DF is transferred entirely, followed by the barrier and a single unpack. only then it starts to generate and transfer the second dataframe (right hand side). We can further see that once the second dataframe is transferred, the first unpack (18) will be merged by the previously unpacked partition of the first dataframe. from here on, the tasks will be unpacked in corresponding pairs.
I would argue this is perfect sorting for this case.

If the unpack task pairs do not have the same worker assingment, this will require another data transfer, of course so I would be more interested in knowing how the worker assignment looks like for this case.
My naive expectation given our current implementation is that the output should already be assigned properly but we should definitely check this.

FWIW I believe tasks with restrictions, i.e. the unpack tasks are not using queuing, yet so this may be a weird interaction between the queued root tasks and the not-queued unpack tasks

[fjetter]

This can be easily reproduced

from dask.datasets import timeseries
ddf1 = timeseries(
    "2020-01-01",
    "2030-01-02",
)
ddf2 = timeseries(
    "2020-01-01",
    "2030-01-02",
)
ddf1.merge(ddf2, shuffle='p2p').compute()

screenshot of the tipping point.

[hendrikmakait]

one after another as prescribed by the topological sort.

This is not what the topological sort suggests. [...]

If we look at the ordering, we can clearly see that the first DF is transferred entirely, followed by the barrier and a single unpack. only then it starts to generate and transfer the second dataframe (right hand side). We can further see that once the second dataframe is transferred, the first unpack (18) will be merged by the previously unpacked partition of the first dataframe. from here on, the tasks will be unpacked in corresponding pairs.

Sorry, I adjusted the description to be more precise. I skipped over the single unpack to focus on the important part: We should transfer all shuffle data before starting to read (a lot of) outputs.

I would argue this is perfect sorting for this case.

Agreed, as I said, what happens without queueing is what the topological sort prescribes and what I would expect to happen.

[gjoseph92]

Interesting. I think this is happening because the output tasks use worker restrictions, so they can't be queued.

Queuing would handle this graph structure fine. Both the transfer tasks and the unpack tasks should be root-ish, so everything should be queued and processed in priority order.

But, because the unpack tasks use worker restrictions, they don't meet the is_rootish heuristic, and aren't queued. So while the input tasks wait in an orderly queue, the output tasks all get submitted at once after the barrier (aka root task overproduction). When this happens, the workers become fully saturated, preventing the remaining input tasks in the queue from being scheduled.

(as @fjetter showed) here's the graph in priority order:

• Queuing generally doesn't play well with task restrictions. Normal tasks are always going to defer to restricted tasks, even if numerically, their priorities indicate they should be scheduled first.

  That is, queueing only works assuming all tasks could be queued. If some tasks get to bypass the queue, we'll always execute them first, and traverse the graph way out of priority order.

  More broadly, I'd say scheduling with restrictions is still second-class in dask IMO. I don't know what we want to do about this.

• Fixing this well would be hard. We'd need to work out better scheduling mechanisms for restricted tasks, including a way of queuing them on the scheduler.

• The easiest, hacky fix is probably a way to prevent the scheduler from queuing the transfer tasks. Maybe like:

diff --git a/distributed/scheduler.py b/distributed/scheduler.py
index fa167efc..ec04b7f2 100644
--- a/distributed/scheduler.py
+++ b/distributed/scheduler.py
@@ -2211,7 +2211,10 @@ class SchedulerState:
             # removed, there should only be one, which combines co-assignment and
             # queuing. Eventually, special-casing root tasks might be removed entirely,
             # with better heuristics.
-            if math.isinf(self.WORKER_SATURATION):
+            if (
+                math.isinf(self.WORKER_SATURATION)
+                or ts.annotations.get("queue") is False
+            ):
                 if not (ws := self.decide_worker_rootish_queuing_disabled(ts)):
                     return {ts.key: "no-worker"}, {}, {}
             else:
diff --git a/distributed/shuffle/_shuffle.py b/distributed/shuffle/_shuffle.py
index 2f07b56e..3fbf5336 100644
--- a/distributed/shuffle/_shuffle.py
+++ b/distributed/shuffle/_shuffle.py
@@ -135,7 +135,7 @@ class P2PShuffleLayer(SimpleShuffleLayer):
         annotations: dict | None = None,
     ):
         annotations = annotations or {}
-        annotations.update({"shuffle": lambda key: key[1]})
+        annotations.update({"shuffle": lambda key: key[1], "queue": False})
         super().__init__(
             name,
             column,

[hendrikmakait]

• The easiest, hacky fix is probably a way to prevent the scheduler from queuing the transfer tasks.

We'd have to propagate the queue annotation upstream, which might very well have unintended side effects. Otherwise, the issue would be that the assign tasks upstream of the transfer still get queued.

[gjoseph92]

Yup, setting a queue=False annotation on the P2PShuffleLayer doesn't work. It does prevent the transfer tasks from getting queued—but the upstream tasks from that are still queued. You'd have to set queue=False on your entire pipeline up until the shuffle.

I suppose we could write a HLG optimization to do this for you. This might even be the best option. Implementing the scheduling to make queueing work alongside task restrictions seems like a significant amount of work.

[hendrikmakait]

While queue=False won't work, maybe we could do the opposite, add an annotation that lets the scheduler queue tasks that have worker-restrictions? I'm not sure what the consequences of that would be performance-wise, but that might be a quick hack.

[gjoseph92]

Unfortunately we can't do that. Because when a slot opens up on a worker, we just pop the first task off the queue. If that task in fact needs to run on a different worker, we can't use it. We'd end up having to search the queue for a task that can run on the worker. This is at least O(n) per task (sort plus linear search)!

The whole design of queuing right now is around tasks being agnostic to which worker they run on (not just regarding restrictions; they don't even have dependencies we need to schedule near). Making queuing performant with tasks that are heterogeneous in any way will require some thought.

[hendrikmakait]

@gjoseph92: Yeah, I feared that was the answer. Thanks for reminding me of the implementation!

[gjoseph92]

I'll note that worker restrictions are probably the easiest case of restrictions to add to queueing. Compared to arbitrary (combinations of) resource annotations, for example, worker-restricted tasks are pretty straightforward. You just need a queue per worker, and you don't need to worry about rebalancing those queues (except when restrictions on tasks are changed...).

The trickier part is hooking it into the scheduling process. There are now multiple queues from which to source tasks. So you need to take the top K tasks from both queues. That's not hard with just 2 queues, but making it performant when expanding into other types of task restrictions (and potentially many queues) could be more complex.

(The disappointing thing is that worker restrictions are probably the least-useful type of restriction for other purposes IMO. Resource restrictions are pretty useful and people actually use them sometimes, worker restrictions I imagine less so in both regards.)