Fuse compatible annotations (#9402)

Ian Rose · web-flow · commit 44cd9bc7dccd · 2022-08-19T15:24:08.000-05:00
diff --git 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
 
diff --git 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,
     )
 
diff --git a/dask/dask-schema.yaml b/dask/dask-schema.yaml
@@ -92,6 +92,21 @@ properties:
     type: object
     properties:
 
+      annotations:
+        type: object
+        properties:
+          fuse:
+            type: boolean
+            description: |
+              If adjacent blockwise layers have different annotations (e.g., one has
+              retries=3 and another has retries=4), Dask can make an attempt to merge
+              those annotations according to some simple rules. ``retries`` is set to
+              the max of the layers, ``priority`` is set to the max of the layers,
+              ``resources`` are set to the max of all the resources, ``workers`` is
+              set to the intersection of the requested workers. If this setting is
+              disabled, then adjacent blockwise layers with different annotations
+              will *not* be fused.
+
       fuse:
         type: object
         description: Options for Dask's task fusion optimizations
diff --git a/dask/dask.yaml b/dask/dask.yaml
@@ -19,6 +19,8 @@ array:
     split-large-chunks: null  # How to handle large output chunks in slicing. Warns by default.
 
 optimization:
+  annotations:
+    fuse: true  # Automatically fuse compatible annotations on layers
   fuse:
     active: null  # Treat as false for dask.dataframe, true for everything else
     ave-width: 1
