Repository: dask/dask

Issue:
Performance regression in Parquet I/O
In #9175 we added some logic to bump up the default number of retries when writing parquet to remote filesystems. The reasoning was that remote I/O is usually somewhat flakier than local I/O, and we wanted to provide a better default user experience for that case without having to recommend everyone explicitly tune `retries` themselves.

Unfortunately, this attempt to increase reliability seems to have introduced a regression in performance (at least in some cases). The reason is that setting retries on specific operations goes through the dask annotation system, i.e.:
```python
with dask.annotate(retries=5):
   ... # do some stuff
```
Blockwise optimization fuses adjacent blockwise layers so that they can be reliably run as a single task, reducing intermediate results and task overhead. However, the optimization process bails if the annotations of adjacent layers don't match (the idea being: how do we choose to annotate the combined layer?): https://github.com/dask/dask/blob/097decd5f4914d7cd85b3babd6ecdd20a633bf8d/dask/blockwise.py#L1354-L1360

So by annotating-by-default parquet I/O layers, we accidentally prevented them from being fused with adjacent blockwise layers.

To reproduce, the following workflow does a very simple `to_parquet` call and measures the average memory usage of a local cluster:

```python
import time

import dask
import dask.datasets
import distributed
from distributed.diagnostics import MemorySampler

cluster = distributed.LocalCluster()
client = distributed.Client(cluster)

# Fake something that looks like a remote filesystem
# for the purpose of local demonstration. In reality, this would
# be something like s3://bucket/path/to/key
remote_fn = f"simplecache://timeser"
storage_options = {"target_protocol": "file"}

sampler = MemorySampler()
with sampler.sample("pq"):
    ddf = dask.datasets.timeseries(freq="10ms", partition_freq="12H")
    scalar = ddf.to_parquet(remote_fn, storage_options=storage_options, compute=False)
    scalar.compute()

df = sampler.to_pandas()
print(df.pq.mean() / 1024**3)
```

On `main`, this uses ~4.5 GiB average memory. But if I revert #9175 it uses an average of ~3.5 GiB. See also [here](https://github.com/coiled/coiled-runtime/pull/247#issuecomment-1218511200) for some evidence using larger, cloud-based cluster. We can inspect the high level graph post-optimization to confirm the failure to fuse blockwise layers. On `main`:

![image](https://user-images.githubusercontent.com/5728311/185263627-22d949f0-3900-4179-9b91-96dd67903e6b.png)

And with #9175 reverted:

![image](https://user-images.githubusercontent.com/5728311/185263663-210df8e8-069e-41cd-9a41-485f4ba0f79c.png)

This regression should only be an issue when there are adjacent blockwise layers to our parquet I/O. But I expect that to be quite common in practice. To me, this is a big enough deal that we should consider reverting #9175 and considering whether there is a better way to mark some set of tasks as possibly flakier than others.

Finally, I'll note that the interaction between annotations and graph optimization has been a persistent problem, and has led to a number of surprising behaviors (see at least #7036 and it's linked issues). That interaction may be due for some serious reconsideration.


Discussion:
An alternative to just reverting #9175 could be to actually do the work of merging annotations of adjacent layers based on some simple rules. They could be something like
* **priority**: max of layer priorities
* **workers**: intersection of layer workers
* **resources**: max of layer resources
* **retries**: max of layer retries
* Any custom annotations: bail on fusing layers

The above *could* work, but (a) I'm not sure if it's a great idea to put more resources towards the layer+annotation interactions which have proven so difficult, and (b) it may be worth doing more carefully after reverting #9175
Have we tried using retries on IO level? Naively, I would expect boto/fsspec/whatever to do this for us instead. 
>Have we tried using retries on IO level? Naively, I would expect boto/fsspec/whatever to do this for us instead.

Without understanding the problem fully, I am also hoping this is a possibility :)

>The above could work, but (a) I'm not sure if it's a great idea to put more resources towards the layer+annotation interactions which have proven so difficult, and (b) it may be worth doing more carefully after reverting https://github.com/dask/dask/pull/9175

I seems like it could be possible to implement the rules you described above in a relatively clear and simple way. If so, I'd personally be supportive of that solution (even if the initial PR only addresses the `retries` case for now).
> Have we tried using retries on IO level? Naively, I would expect boto/fsspec/whatever to do this for us instead.

Both boto and s3fs/gcsfs/fsspec have some functionality around retries (this is discussed in the original issue), but even then odd behavior can bubble up to the dask level The work here was basically a belt-and-suspenders approach.

I'm sure there is also work that could be done in s3fs et al (perhaps @martindurant has some insight). The retry logic there currently retries for some error codes, but not for the whole universe of things that might go wrong.
> I seems like it could be possible to implement the rules you described above in a relatively clear and simple way. If so, I'd personally be supportive of that solution (even if the initial PR only addresses the retries case for now).

I also think this could be fairly doable (regardless of whether this particular approach to retries is a good idea). I'm just not sure what our stance towards feature development on HLG/layers is right now...
> I'm sure there is also work that could be done in s3fs et al 

To be sure, there are some things that should *not* be retried, such as not-found and permission violations.
>I'm just not sure what our stance towards feature development on HLG/layers is right now...

Your hesitation is warranted.  We know we want to change many things in HLG, and we want to avoid further-complicating the current state. However, my understanding is that this only applies to `Blockwise` fusion (there is no other way to fuse HLG Layers), and so the "fix" could be somewhat distinct from the larger HLG design. I'm imagining a simple/modular `check_fusable_annotations` check in the blockwise-fusion logic that we may or may not throw away in the future.

Resolving patch:
diff --git a/dask/blockwise.py b/dask/blockwise.py
--- a/dask/blockwise.py
+++ b/dask/blockwise.py
@@ -9,6 +9,7 @@
 
 import tlz as toolz
 
+import dask
 from dask.base import clone_key, get_name_from_key, tokenize
 from dask.core import flatten, keys_in_tasks, reverse_dict
 from dask.delayed import unpack_collections
@@ -1351,10 +1352,9 @@ def _optimize_blockwise(full_graph, keys=()):
                 ):
                     stack.append(dep)
                     continue
-                if (
-                    blockwise_layers
-                    and layers[next(iter(blockwise_layers))].annotations
-                    != layers[dep].annotations
+                if blockwise_layers and not _can_fuse_annotations(
+                    layers[next(iter(blockwise_layers))].annotations,
+                    layers[dep].annotations,
                 ):
                     stack.append(dep)
                     continue
@@ -1412,6 +1412,60 @@ def _unique_dep(dep, ind):
     return dep + "_" + "_".join(str(i) for i in list(ind))
 
 
+def _can_fuse_annotations(a: dict | None, b: dict | None) -> bool:
+    """
+    Treat the special annotation keys, as fusable since we can apply simple
+    rules to capture their intent in a fused layer.
+    """
+    if a == b:
+        return True
+
+    if dask.config.get("optimization.annotations.fuse") is False:
+        return False
+
+    fusable = {"retries", "priority", "resources", "workers", "allow_other_workers"}
+    if (not a or all(k in fusable for k in a)) and (
+        not b or all(k in fusable for k in b)
+    ):
+        return True
+
+    return False
+
+
+def _fuse_annotations(*args: dict) -> dict:
+    """
+    Given an iterable of annotations dictionaries, fuse them according
+    to some simple rules.
+    """
+    # First, do a basic dict merge -- we are presuming that these have already
+    # been gated by `_can_fuse_annotations`.
+    annotations = toolz.merge(*args)
+    # Max of layer retries
+    retries = [a["retries"] for a in args if "retries" in a]
+    if retries:
+        annotations["retries"] = max(retries)
+    # Max of layer priorities
+    priorities = [a["priority"] for a in args if "priority" in a]
+    if priorities:
+        annotations["priority"] = max(priorities)
+    # Max of all the layer resources
+    resources = [a["resources"] for a in args if "resources" in a]
+    if resources:
+        annotations["resources"] = toolz.merge_with(max, *resources)
+    # Intersection of all the worker restrictions
+    workers = [a["workers"] for a in args if "workers" in a]
+    if workers:
+        annotations["workers"] = list(set.intersection(*[set(w) for w in workers]))
+    # More restrictive of allow_other_workers
+    allow_other_workers = [
+        a["allow_other_workers"] for a in args if "allow_other_workers" in a
+    ]
+    if allow_other_workers:
+        annotations["allow_other_workers"] = all(allow_other_workers)
+
+    return annotations
+
+
 def rewrite_blockwise(inputs):
     """Rewrite a stack of Blockwise expressions into a single blockwise expression
 
@@ -1435,6 +1489,9 @@ def rewrite_blockwise(inputs):
         # Fast path: if there's only one input we can just use it as-is.
         return inputs[0]
 
+    fused_annotations = _fuse_annotations(
+        *[i.annotations for i in inputs if i.annotations]
+    )
     inputs = {inp.output: inp for inp in inputs}
     dependencies = {
         inp.output: {d for d, v in inp.indices if v is not None and d in inputs}
@@ -1560,7 +1617,7 @@ def rewrite_blockwise(inputs):
         numblocks=numblocks,
         new_axes=new_axes,
         concatenate=concatenate,
-        annotations=inputs[root].annotations,
+        annotations=fused_annotations,
         io_deps=io_deps,
     )
 


Test patch:
diff --git a/dask/array/tests/test_atop.py b/dask/array/tests/test_atop.py
--- a/dask/array/tests/test_atop.py
+++ b/dask/array/tests/test_atop.py
@@ -16,7 +16,7 @@
     rewrite_blockwise,
 )
 from dask.highlevelgraph import HighLevelGraph
-from dask.utils_test import dec, inc
+from dask.utils_test import dec, hlg_layer_topological, inc
 
 a, b, c, d, e, f, g = "a", "b", "c", "d", "e", "f", "g"
 _0, _1, _2, _3, _4, _5, _6, _7, _8, _9 = (
@@ -324,7 +324,54 @@ def test_optimize_blockwise():
     )
 
 
-def test_optimize_blockwise_annotations():
+def test_optimize_blockwise_control_annotations():
+    """
+    Can we fuse blockwise layers with different, but compatible
+    annotations for retries, priority, etc.
+    """
+
+    a = da.ones(10, chunks=(5,))
+    b = a + 1
+
+    with dask.annotate(retries=5, workers=["a", "b", "c"], allow_other_workers=False):
+        c = b + 2
+
+    with dask.annotate(priority=2, workers=["b", "c", "d"], allow_other_workers=True):
+        d = c + 3
+
+    with dask.annotate(retries=3, resources={"GPU": 2, "Memory": 10}):
+        e = d + 4
+
+    with dask.annotate(priority=4, resources={"GPU": 5, "Memory": 4}):
+        f = e + 5
+
+    # This one will not be fused due to the custom annotation, nor will the one below
+    with dask.annotate(foo="bar"):
+        g = f + 6
+
+    h = g + 6
+
+    dsk = da.optimization.optimize_blockwise(h.dask)
+
+    # The layers and their annotations should be fusable until the custom one
+    assert len(dsk.layers) == 3
+    layer = hlg_layer_topological(dsk, 0)  # First layer is the fused one
+    annotations = layer.annotations
+
+    assert len(annotations) == 5
+    assert annotations["priority"] == 4  # max
+    assert annotations["retries"] == 5  # max
+    assert annotations["allow_other_workers"] is False  # More restrictive
+    assert set(annotations["workers"]) == {"b", "c"}  # intersection
+    assert annotations["resources"] == {"GPU": 5, "Memory": 10}  # Max of resources
+
+    # If we disable blockwise annotation fusion, we can only fuse the first two layers.
+    with dask.config.set({"optimization.annotations.fuse": False}):
+        dsk = da.optimization.optimize_blockwise(h.dask)
+        assert len(dsk.layers) == 7
+
+
+def test_optimize_blockwise_custom_annotations():
     a = da.ones(10, chunks=(5,))
     b = a + 1