Repository: dask/dask
Pull request: Implement fuse option for delayed objects
- [ ] Closes #11436
- [ ] Tests added / passed
- [ ] Passes `pre-commit run --all-files`


Problem statement:
Make it easier to enable linear fusion for delayed objects
Delayed can be used to express embarrassingly parallel chains of operations. Certain configurations benefit from processing one chain completely before starting a new one (high memory usage for example). This is currently not possible by default because other tasks will almost certainly run before follow up tasks are scheduled. Linear fusion is the obvious solution for this which ensures that there is only a single task in the dask sense to express the whole chain. Enabling this is currently a real PITA:

```

import dask
from dask.optimization import cull, fuse
from dask.utils import ensure_dict
from dask import config, core

def fuse_delayed(dsk, keys, *args, **kwargs):
    dependencies = dsk.get_all_dependencies()
    dsk = ensure_dict(dsk)

    fuse_subgraphs = config.get("optimization.fuse.subgraphs")
    if fuse_subgraphs is None:
        fuse_subgraphs = True
    dsk, _ = fuse(
        dsk,
        keys,
        dependencies=dependencies,
        fuse_subgraphs=fuse_subgraphs,
    )
    dsk, _ = cull(dsk, keys)
    return dsk
    
dask.config.set({"delayed_optimize": fuse_delayed})
```

This is pretty ugly. Even if we ignore this, you still have to put a function into the dask config which isn't great either.

Ideally, we add a configuration like

```
dask.config.set({"enable-delayed-fusion": True})
```

Even better, we turn this on by default, but one step at a time

Hints from issue discussion:

Commit message: Fixup

Patch to review:
diff --git a/dask/dask-schema.yaml b/dask/dask-schema.yaml
index 4be821187d6..267bbeccf07 100644
--- a/dask/dask-schema.yaml
+++ b/dask/dask-schema.yaml
@@ -244,6 +244,13 @@ properties:
               False, then the top-most key will be used. For advanced usage, a function
               to create the new name is also accepted.
 
+          delayed:
+            type: boolean
+            description: |
+              Whether to fuse adjacent delayed calls together. This can be useful to ensure 
+              that an independent chain of delayed calls is executed in a single task to
+              reduce the memory footprint on a single worker.
+
   admin:
     type: object
     properties:
diff --git a/dask/dask.yaml b/dask/dask.yaml
index 5ba27ce89d7..ce6b93cfeb6 100644
--- a/dask/dask.yaml
+++ b/dask/dask.yaml
@@ -42,6 +42,7 @@ optimization:
     max-depth-new-edges: null  # ave_width * 1.5
     subgraphs: null  # true for dask.dataframe, false for everything else
     rename-keys: true
+    delayed: false
 
 admin:
   traceback:
diff --git a/dask/delayed.py b/dask/delayed.py
index 1ff6a2dba56..e43fab1da88 100644
--- a/dask/delayed.py
+++ b/dask/delayed.py
@@ -22,10 +22,12 @@
 from dask.context import globalmethod
 from dask.core import flatten, quote
 from dask.highlevelgraph import HighLevelGraph
+from dask.optimization import fuse
 from dask.typing import Graph, NestedKeys
 from dask.utils import (
     OperatorMethodMixin,
     apply,
+    ensure_dict,
     funcname,
     is_namedtuple_instance,
     methodcaller,
@@ -522,9 +524,29 @@ def right(method):
     return partial(_swap, method)
 
 
+def _fuse_delayed(dsk, keys, **kwargs):
+    dependencies = dsk.get_all_dependencies()
+    dsk = ensure_dict(dsk)
+
+    fuse_subgraphs = config.get("optimization.fuse.subgraphs")
+    if fuse_subgraphs is None:
+        fuse_subgraphs = True
+    dsk, _ = fuse(
+        dsk,
+        keys,
+        dependencies=dependencies,
+        fuse_subgraphs=fuse_subgraphs,
+    )
+    return dsk
+
+
 def optimize(dsk, keys, **kwargs):
     if not isinstance(keys, (list, set)):
         keys = [keys]
+
+    if config.get("optimization.fuse.delayed"):
+        dsk = _fuse_delayed(dsk, keys, **kwargs)
+
     if not isinstance(dsk, HighLevelGraph):
         dsk = HighLevelGraph.from_collections(id(dsk), dsk, dependencies=())
     dsk = dsk.cull(set(flatten(keys)))
diff --git a/dask/tests/test_delayed.py b/dask/tests/test_delayed.py
index b9ae588cf67..c318488f925 100644
--- a/dask/tests/test_delayed.py
+++ b/dask/tests/test_delayed.py
@@ -16,6 +16,7 @@
 import dask
 import dask.bag as db
 from dask import compute
+from dask.base import collections_to_dsk
 from dask.delayed import Delayed, delayed, to_task_dask
 from dask.highlevelgraph import HighLevelGraph
 from dask.utils_test import inc
@@ -855,3 +856,27 @@ def add(x, y):
     assert add.__name__ == "add"
     assert add.__doc__ == "This is a docstring"
     assert add.__wrapped__(1, 2) == 3
+
+
+def test_delayed_fusion():
+    @delayed
+    def test(i):
+        return i + 1
+
+    @delayed
+    def test2(i):
+        return i + 2
+
+    @delayed
+    def test3(i):
+        return i + 3
+
+    obj = test3(test2(test(10)))
+    dsk = dict(collections_to_dsk([obj]))
+    assert len(dsk) == 3
+
+    obj2 = test3(test2(test(10)))
+    with dask.config.set({"optimization.fuse.delayed": True}):
+        dsk2 = dict(collections_to_dsk([obj]))
+    assert len(dsk2) == 2
+    assert dask.compute(obj) == dask.compute(obj2)


Code review comments:
diff_hunk: @@ -855,3 +856,27 @@ def add(x, y):
     assert add.__name__ == "add"
     assert add.__doc__ == "This is a docstring"
     assert add.__wrapped__(1, 2) == 3
+
+
+def test_delayed_fusion():
+    @delayed
+    def test(i):
+        return i + 1
+
+    @delayed
+    def test2(i):
+        return i + 2
+
+    @delayed
+    def test3(i):
+        return i + 3
+
+    obj = test3(test2(test(10)))
+    dsk = dict(collections_to_dsk([obj]))
+    assert len(dsk) == 3
+
+    obj2 = test3(test2(test(10)))
+    with dask.config.set({"optimization.fuse.delayed": True}):
+        dsk2 = dict(collections_to_dsk([obj]))
+    assert len(dsk2) == 2
+    assert dask.compute(obj) == dask.compute(obj2)
path: dask/tests/test_delayed.py
text: @user1:
Curious what you are testing here? (both `obj` and `obj2`) are defined as `test3(test2(test(10)))`.

@author:
good point, I intended to compute the first one in the with statement, not sure why I didn't do that

@author:
done