Repository: dask/dask

Issue:
[REVIEW] Prioritize getitem tasks in shuffle
This PR implements custom task prioritization of `getitem()` tasks in shuffle operations. See https://github.com/dask/dask/pull/6051 for a detailed discussion of the performance benefits.  

```python

    @staticmethod
    def __dask_distributed_annotations_unpack__(
        annotations: MutableMapping[str, Any],
        new_annotations: Optional[Mapping[str, Any]],
        keys: Iterable[Hashable],
    ) -> None:
        """Implement custom task prioritization of `getitem()` tasks.

        The scheduling policies of Dask is depth-first generally speaking, which works
        great in most cases. However, in case of shuffle, it increases the memory usage
        significantly. This is because depth-first delays the freeing of the result of
        `shuffle_group()` until the end of the shuffling.

        We address this by forcing a breadth-first scheduling of the `getitem()` tasks
        that follows the `shuffle_group()` task. This way all the `getitem()` tasks
        finish immediately and the output of shuffle_group()` can be freed before
        continuing to the next phase of the operation.

        See https://github.com/dask/dask/pull/6051 for a detailed discussion
        """
```

- [ ] Closes #6051
- [ ] Tests added / passed
- [x] Passes `black dask` / `flake8 dask` / `isort dask`



Discussion:
Is there something we can do here that would be a good test?  The implementation here with `"split-"` seems somewhat brittle to me (future devs could change that name and this functionality would stop without us knowing).  

Maybe we look at the Scheduler.transition_log and see that some set of tasks occurs before some other set of tasks?
Also, I think that doing this in unpack is fine, but I'm also curious if it would make sense to put these annotations on the Layer itself and then trust the general machinery to move those annotations up to the scheduler.  This feels somewhat less efficient, but also somewhat more general. 
> Is there something we can do here that would be a good test? The implementation here with `"split-"` seems somewhat brittle to me (future devs could change that name and this functionality would stop without us knowing).
> 
> Maybe we look at the Scheduler.transition_log and see that some set of tasks occurs before some other set of tasks?

I have been thinking of a more general solution but that will require some API redesign, I think. 
One idea is to make [`order()`](https://github.com/dask/dask/blob/main/dask/order.py#L85) understand annotations and prioritizes  fast-tasks such as `getitem()`. But I think it makes sense to start with this specialized approach and get some experience before redesigning the task ordering API.


> Also, I think that doing this in unpack is fine, but I'm also curious if it would make sense to put these annotations on the Layer itself and then trust the general machinery to move those annotations up to the scheduler. This feels somewhat less efficient, but also somewhat more general.

Notice, before layer unpack we don't know the task keys. We could do this in `__dask_distributed_unpack__` instead of `__dask_distributed_annotations_unpack__` but I like to keep annotation specializations to the annotation method.
There is a middle ground.  We can put annotations on the layer instance.  This would be equivalent to what happens when we call 

```python
with dask.annotate(priority=...):
    ... create some layers ...
```

We can just populate the layer instance with an `annotations = {"priority": ...}` dict.  Again, this isn't as efficient (we have to pack and unpack annotations) but it does reduce the number of code paths that we use for this sort of thing.  While layers and packing are in flux it might be wise to stick to client-side / blessed API for a while.
> We can just populate the layer instance with an `annotations = {"priority": ...}` dict. Again, this isn't as efficient (we have to pack and unpack annotations) but it does reduce the number of code paths that we use for this sort of thing. While layers and packing are in flux it might be wise to stick to client-side / blessed API for a while.

We only want to prioritize `getitem()` tasks, the rest of the shuffle tasks should keep there original priority. Basically, it is the Shuffle.Layer that tells the scheduler that some of it tasks needs priority.
> We only want to prioritize getitem() tasks, the rest of the shuffle tasks should keep there original priority. Basically, it is the Shuffle.Layer that tells the scheduler that some of it tasks needs priority.

Agreed.  I believe that you can put the information you placed in the unpack method directly on the Layer and that the Client will send that information up.  
> Agreed. I believe that you can put the information you placed in the unpack method directly on the Layer and that the Client will send that information up.

I am not sure I follow, do you suggest that we set annotations when creating `SimpleShuffleLayer` and `ShuffleLayer` like here: https://github.com/dask/dask/blob/main/dask/dataframe/shuffle.py#L602 ?

In that case we need to use a `callable` annotation in order to avoid graph materialization, which could also work.
Yes there, or even within the layer's constructor 

> In that case we need to use a callable annotation in order to avoid graph materialization, which could also work.

My hope is that we could figure out the relevant keys without generating the entire graph.  

What is in this PR is fine too.  I just think that if we can keep this logic on the client side then it's less likely to need maintenance going forward.  
> My hope is that we could figure out the relevant keys without generating the entire graph.

 I will have to think about this :)
@rjzamora I'm not sure if @madsbk mentioned this to you or not, but if it was easy to generate a set of keys for shuffle layers without generating the full graph then that would make it easier for us to specify nicer priorities at the layer level to shuffles.  I would love to get this PR in, one way or another, by next release.

My hope is that you understand the shuffle system well enough to be able to quickly generate a `def keys(self)` method.  I think that if we add that then we can move the callable here directly onto the Layer instance, rather than push it into the unpack method.  

If we can't do that this week then let's punt and find a way to merge this in.
>My hope is that you understand the shuffle system well enough to be able to quickly generate a def keys(self) method.

Makes sense - I’ll look through this PR and give you my thoughts tomorrow.  There is no doubt that we can generate a set of keys without “materializing” the graph, but the algorithm will have the same complexity as _construct_graph.
Oh interesting.  I would have expected it to be fairly cheap to generate just the keys. 
>Oh interesting. I would have expected it to be fairly cheap to generate just the keys.

I should say *may* have - I'll need to think about it :). 
Also, we only need to generate the keys that need to have a set priority here.

Resolving patch:
diff --git a/dask/layers.py b/dask/layers.py
--- a/dask/layers.py
+++ b/dask/layers.py
@@ -411,6 +411,34 @@ def __init__(
         self.name_input = name_input
         self.meta_input = meta_input
         self.parts_out = parts_out or range(npartitions)
+        self.split_name = "split-" + self.name
+
+        # The scheduling policy of Dask is generally depth-first,
+        # which works great in most cases. However, in case of shuffle,
+        # it increases the memory usage significantly. This is because
+        # depth-first delays the freeing of the result of `shuffle_group()`
+        # until the end of the shuffling.
+        #
+        # We address this by manually setting a high "prioroty" to the
+        # `getitem()` ("split") tasks, using annotations. This forces a
+        # breadth-first scheduling of the tasks tath directly depend on
+        # the `shuffle_group()` output, allowing that data to be freed
+        # much earlier.
+        #
+        # See https://github.com/dask/dask/pull/6051 for a detailed discussion.
+        self.annotations = self.annotations or {}
+        if "priority" not in self.annotations:
+            self.annotations["priority"] = {}
+        self.annotations["priority"]["__expanded_annotations__"] = None
+        self.annotations["priority"].update({_key: 1 for _key in self.get_split_keys()})
+
+    def get_split_keys(self):
+        # Return SimpleShuffleLayer "split" keys
+        return [
+            stringify((self.split_name, part_out, part_in))
+            for part_in in range(self.npartitions_input)
+            for part_out in self.parts_out
+        ]
 
     def get_output_keys(self):
         return {(self.name, part) for part in self.parts_out}
@@ -557,7 +585,6 @@ def _construct_graph(self, deserializing=False):
         """Construct graph for a simple shuffle operation."""
 
         shuffle_group_name = "group-" + self.name
-        shuffle_split_name = "split-" + self.name
 
         if deserializing:
             # Use CallableLazyImport objects to avoid importing dataframe
@@ -574,7 +601,7 @@ def _construct_graph(self, deserializing=False):
         dsk = {}
         for part_out in self.parts_out:
             _concat_list = [
-                (shuffle_split_name, part_out, part_in)
+                (self.split_name, part_out, part_in)
                 for part_in in range(self.npartitions_input)
             ]
             dsk[(self.name, part_out)] = (
@@ -583,7 +610,7 @@ def _construct_graph(self, deserializing=False):
                 self.ignore_index,
             )
             for _, _part_out, _part_in in _concat_list:
-                dsk[(shuffle_split_name, _part_out, _part_in)] = (
+                dsk[(self.split_name, _part_out, _part_in)] = (
                     operator.getitem,
                     (shuffle_group_name, _part_in),
                     _part_out,
@@ -655,6 +682,9 @@ def __init__(
         parts_out=None,
         annotations=None,
     ):
+        self.inputs = inputs
+        self.stage = stage
+        self.nsplits = nsplits
         super().__init__(
             name,
             column,
@@ -666,9 +696,23 @@ def __init__(
             parts_out=parts_out or range(len(inputs)),
             annotations=annotations,
         )
-        self.inputs = inputs
-        self.stage = stage
-        self.nsplits = nsplits
+
+    def get_split_keys(self):
+        # Return ShuffleLayer "split" keys
+        keys = []
+        for part in self.parts_out:
+            out = self.inputs[part]
+            for i in range(self.nsplits):
+                keys.append(
+                    stringify(
+                        (
+                            self.split_name,
+                            out[self.stage],
+                            insert(out, self.stage, i),
+                        )
+                    )
+                )
+        return keys
 
     def __repr__(self):
         return "ShuffleLayer<name='{}', stage={}, nsplits={}, npartitions={}>".format(
@@ -738,7 +782,6 @@ def _construct_graph(self, deserializing=False):
         """Construct graph for a "rearrange-by-column" stage."""
 
         shuffle_group_name = "group-" + self.name
-        shuffle_split_name = "split-" + self.name
 
         if deserializing:
             # Use CallableLazyImport objects to avoid importing dataframe
@@ -763,7 +806,7 @@ def _construct_graph(self, deserializing=False):
                 # Get out each individual dataframe piece from the dicts
                 _inp = insert(out, self.stage, i)
                 _idx = out[self.stage]
-                _concat_list.append((shuffle_split_name, _idx, _inp))
+                _concat_list.append((self.split_name, _idx, _inp))
 
             # concatenate those pieces together, with their friends
             dsk[(self.name, part)] = (
@@ -773,7 +816,7 @@ def _construct_graph(self, deserializing=False):
             )
 
             for _, _idx, _inp in _concat_list:
-                dsk[(shuffle_split_name, _idx, _inp)] = (
+                dsk[(self.split_name, _idx, _inp)] = (
                     operator.getitem,
                     (shuffle_group_name, _inp),
                     _idx,


Test patch:
diff --git a/dask/tests/test_distributed.py b/dask/tests/test_distributed.py
--- a/dask/tests/test_distributed.py
+++ b/dask/tests/test_distributed.py
@@ -528,6 +528,35 @@ async def test_futures_in_subgraphs(c, s, a, b):
     ddf = await c.submit(dd.categorical.categorize, ddf, columns=["day"], index=False)
 
 
+@gen_cluster(client=True)
+async def test_shuffle_priority(c, s, a, b):
+    pd = pytest.importorskip("pandas")
+    np = pytest.importorskip("numpy")
+    dd = pytest.importorskip("dask.dataframe")
+
+    df = pd.DataFrame({"a": range(1000)})
+    ddf = dd.from_pandas(df, npartitions=10)
+    ddf2 = ddf.shuffle("a", shuffle="tasks", max_branch=32)
+    await c.compute(ddf2)
+
+    # Parse transition log for processing tasks
+    log = [
+        eval(l[0])[0]
+        for l in s.transition_log
+        if l[1] == "processing" and "simple-shuffle-" in l[0]
+    ]
+
+    # Make sure most "split" tasks are processing before
+    # any "combine" tasks begin
+    late_split = np.quantile(
+        [i for i, st in enumerate(log) if st.startswith("split")], 0.75
+    )
+    early_combine = np.quantile(
+        [i for i, st in enumerate(log) if st.startswith("simple")], 0.25
+    )
+    assert late_split < early_combine
+
+
 @gen_cluster(client=True)
 async def test_map_partitions_da_input(c, s, a, b):
     """Check that map_partitions can handle a dask array input"""