Repository: dask/dask
Pull request #11525 (2024-11-15 11:06:24)

Implement linear_fusion for task spec class


- optimize_slices is now unused, but not sure how public this is
- same for inline_functions

phofl: I checked optimize_slices and there is no usage on Github

inline_functions is not as trivial since that's a more common name :(
github-actions[bot]: ## Unit Test Results
_See [test report](https://dask.github.io/dask/test_report.html) for an extended history of previous test failures. This is useful for diagnosing flaky tests._

     13 files  ± 0       13 suites  ±0   3h 20m 46s ⏱️ -8s
 13 292 tests + 3   12 228 ✅ + 1   1 062 💤 ±0   2 ❌ + 2 
138 713 runs  +37  119 706 ✅ +23  18 995 💤 +2  12 ❌ +12 

For more details on these failures, see [this check](https://github.com/dask/dask/runs/33043800442).

Results for commit a699a31b. ± Comparison against base commit 98331028.

[test-results]:data:application/gzip;base64,H4sIACBHN2cC/02NSw7DIAwFrxKx7gJMAriXqYCAhJpPRWAV9e4lX7qc8dN4JT4MbiHPhvFHQ5YcUqU+R53CPG0MtBVFlWs67oBw8WvJ1u4jAFXlO3w2SUUdeh2G4qpwMc6xGFpMzNORVnJ/v/GdZiipuOWZVojd5c40g0v8p+08jiEVIFogas5MD8pyNMZzCs51hnoJ3vSsdUJa5g35/gD47gI5GwEAAA==


Patch:
diff --git a/dask/_task_spec.py b/dask/_task_spec.py
index 64b700b3c..2d4683f62 100644
--- a/dask/_task_spec.py
+++ b/dask/_task_spec.py
@@ -656,7 +656,6 @@ class Task(GraphNode):
 
     @staticmethod
     def fuse(*tasks: Task, key: KeyType | None = None) -> Task:
-        leafs = set()
         all_keys = set()
         all_deps: set[KeyType] = set()
         for t in tasks:
@@ -829,3 +828,92 @@ def execute_graph(
                 del cache[dep]
 
     return cache
+
+
+def fuse_linear_task_spec(dsk, keys):
+    """
+    keys are the keys from the graph that are requested by a computation. We
+    can't fuse those together.
+    """
+    from dask.core import reverse_dict
+    from dask.optimization import default_fused_keys_renamer
+
+    keys = set(keys)
+    dependencies = DependenciesMapping(dsk)
+    dependents = reverse_dict(dependencies)
+
+    seen = set()
+    result = {}
+
+    def _check_data_node(tsk):
+        return isinstance(tsk, DataNode) and not isinstance(tsk.value, str)
+
+    for key in dsk:
+        if key in seen:
+            continue
+
+        seen.add(key)
+
+        deps = dependencies[key]
+        dependents_key = dependents[key]
+
+        if len(deps) != 1 and len(dependents_key) != 1 or _check_data_node(dsk[key]):
+            result[key] = dsk[key]
+            continue
+
+        linear_chain = [dsk[key]]
+        top_key = key
+
+        # TODO: The old implementation doesn't fuse data nodes
+        # It claims that a data node is easier to serialize as a raw value
+
+        # Walk towards the leafs as long as the nodes have a single dependency
+        # and a single dependent, we can't fuse two nodes of an intermediate node
+        # is the source for 2 dependents
+        while len(deps) == 1:
+            (new_key,) = deps
+            if new_key in seen:
+                break
+            seen.add(new_key)
+            if new_key not in dsk:
+                # This can happen if a future is in the graph, the dependency mapping
+                # adds the key that is referenced by the future as a dependency
+                # see test_futures_to_delayed_array
+                break
+            if (
+                len(dependents[new_key]) != 1
+                or _check_data_node(dsk[new_key])
+                or new_key in keys
+            ):
+                result[new_key] = dsk[new_key]
+                break
+            # backwards comp for new names, temporary until is_rootish is removed
+            linear_chain.insert(0, dsk[new_key])
+            deps = dependencies[new_key]
+
+        # Walk the tree towards the root as long as the nodes have a single dependent
+        # and a single dependency, we can't fuse two nodes if node has multiple
+        # dependencies
+        while len(dependents_key) == 1 and top_key not in keys:
+            new_key = dependents_key.pop()
+            if new_key in seen:
+                break
+            seen.add(new_key)
+            if len(dependencies[new_key]) != 1:
+                # Exit if the dependent has multiple dependencies, triangle
+                result[new_key] = dsk[new_key]
+                break
+            linear_chain.append(dsk[new_key])
+            top_key = new_key
+            dependents_key = dependents[new_key]
+
+        if len(linear_chain) == 1:
+            result[top_key] = linear_chain[0]
+        else:
+            # Renaming the keys is necessary to preserve the rootish detection for now
+            renamed_key = default_fused_keys_renamer([tsk.key for tsk in linear_chain])
+            result[renamed_key] = Task.fuse(*linear_chain, key=renamed_key)
+            if renamed_key != top_key:
+                # Having the same prefixes can result in the same key, i.e. getitem-hash -> getitem-hash
+                result[top_key] = Alias(top_key, target=renamed_key)
+    return result
diff --git a/dask/array/optimization.py b/dask/array/optimization.py
index 101cd80e0..3ec56353a 100644
--- a/dask/array/optimization.py
+++ b/dask/array/optimization.py
@@ -8,12 +8,13 @@ from typing import Any
 import numpy as np
 
 from dask import config
+from dask._task_spec import convert_legacy_graph, fuse_linear_task_spec
 from dask.array.chunk import getitem
 from dask.array.core import getter, getter_inline, getter_nofancy
 from dask.blockwise import fuse_roots, optimize_blockwise
-from dask.core import flatten, reverse_dict
+from dask.core import flatten
 from dask.highlevelgraph import HighLevelGraph
-from dask.optimization import SubgraphCallable, fuse, inline_functions
+from dask.optimization import SubgraphCallable
 from dask.utils import ensure_dict
 
 # All get* functions the optimizations know about
@@ -24,15 +25,7 @@ GETTERS = (getter, getter_nofancy, getter_inline, getitem)
 GETNOREMOVE = (getter, getter_nofancy)
 
 
-def optimize(
-    dsk,
-    keys,
-    fuse_keys=None,
-    fast_functions=None,
-    inline_functions_fast_functions=(getter_inline,),
-    rename_fused_keys=True,
-    **kwargs,
-):
+def optimize(dsk, keys, **kwargs):
     """Optimize dask for array computation
 
     1.  Cull tasks not necessary to evaluate keys
@@ -55,68 +48,12 @@ def optimize(
     if config.get("optimization.fuse.active") is False:
         return dsk
 
-    dependencies = dsk.get_all_dependencies()
     dsk = ensure_dict(dsk)
 
-    # Low level task optimizations
-    if fast_functions is not None:
-        inline_functions_fast_functions = fast_functions
+    dsk = convert_legacy_graph(dsk)
+    dsk = fuse_linear_task_spec(dsk, keys=keys)
 
-    hold = hold_keys(dsk, dependencies)
-
-    dsk, dependencies = fuse(
-        dsk,
-        hold + keys + (fuse_keys or []),
-        dependencies,
-        rename_keys=rename_fused_keys,
-    )
-    if inline_functions_fast_functions:
-        dsk = inline_functions(
-            dsk,
-            keys,
-            dependencies=dependencies,
-            fast_functions=inline_functions_fast_functions,
-        )
-
-    return optimize_slices(dsk)
-
-
-def hold_keys(dsk, dependencies):
-    """Find keys to avoid fusion
-
-    We don't want to fuse data present in the graph because it is easier to
-    serialize as a raw value.
-
-    We don't want to fuse chains after getitem/GETTERS because we want to
-    move around only small pieces of data, rather than the underlying arrays.
-    """
-    dependents = reverse_dict(dependencies)
-    data = {k for k, v in dsk.items() if type(v) not in (tuple, str)}
-
-    hold_keys = list(data)
-    for dat in data:
-        deps = dependents[dat]
-        for dep in deps:
-            task = dsk[dep]
-            # If the task is a get* function, we walk up the chain, and stop
-            # when there's either more than one dependent, or the dependent is
-            # no longer a get* function or an alias. We then add the final
-            # key to the list of keys not to fuse.
-            if _is_getter_task(task):
-                try:
-                    while len(dependents[dep]) == 1:
-                        new_dep = next(iter(dependents[dep]))
-                        new_task = dsk[new_dep]
-                        # If the task is a get* or an alias, continue up the
-                        # linear chain
-                        if _is_getter_task(new_task) or new_task in dsk:
-                            dep = new_dep
-                        else:
-                            break
-                except (IndexError, TypeError):
-                    pass
-                hold_keys.append(dep)
-    return hold_keys
+    return dsk
 
 
 def _is_getter_task(
diff --git a/dask/dataframe/core.py b/dask/dataframe/core.py
index 8d09d520a..e50035131 100644
--- a/dask/dataframe/core.py
+++ b/dask/dataframe/core.py
@@ -8283,9 +8283,15 @@ def new_dd_object(dsk, name, meta, divisions, parent_meta=None):
                 layer.new_axes["j"] = chunks[1][0]
                 layer.output_indices = layer.output_indices + ("j",)
             else:
+                from dask._task_spec import Alias, Task
+
                 suffix = (0,) * (len(chunks) - 1)
                 for i in range(len(chunks[0])):
-                    layer[(name, i) + suffix] = layer.pop((name, i))
+                    task = layer.get((name, i))
+                    new_key = (name, i) + suffix
+                    if isinstance(task, Task):
+                        task = Alias(new_key, task.key)
+                    layer[new_key] = task
         return da.Array(dsk, name=name, chunks=chunks, dtype=meta.dtype)
     else:
         return get_parallel_type(meta)(dsk, name, meta, divisions)


Test patch:
diff --git a/dask/array/tests/test_array_core.py b/dask/array/tests/test_array_core.py
index 5a49c19be..e3d8b8a2d 100644
--- a/dask/array/tests/test_array_core.py
+++ b/dask/array/tests/test_array_core.py
@@ -8,6 +8,8 @@ import xml.etree.ElementTree
 
 import pytest
 
+from dask._task_spec import Task
+
 np = pytest.importorskip("numpy")
 
 import itertools
@@ -2040,7 +2042,12 @@ def test_store_locks():
     v = store([a, b], [at, bt], compute=False, lock=lock)
     assert isinstance(v, Delayed)
     dsk = v.dask
-    locks = {vv for v in dsk.values() for vv in v if isinstance(vv, _Lock)}
+    locks = {
+        vv
+        for v in dsk.values()
+        for vv in (v.args if isinstance(v, Task) else v)
+        if isinstance(vv, _Lock)
+    }
     assert locks == {lock}
 
     # Ensure same lock applies over multiple stores
diff --git a/dask/array/tests/test_optimization.py b/dask/array/tests/test_optimization.py
index fc9bd5b79..cf8583bb5 100644
--- a/dask/array/tests/test_optimization.py
+++ b/dask/array/tests/test_optimization.py
@@ -2,6 +2,8 @@ from __future__ import annotations
 
 import pytest
 
+from dask._task_spec import Alias, DataNode
+
 pytest.importorskip("numpy")
 
 import numpy as np
@@ -12,7 +14,6 @@ from dask.array.chunk import getitem as da_getitem
 from dask.array.core import getter as da_getter
 from dask.array.core import getter_nofancy as da_getter_nofancy
 from dask.array.optimization import (
-    _is_getter_task,
     fuse_slice,
     optimize,
     optimize_blockwise,
@@ -257,8 +258,7 @@ def test_optimize_with_getitem_fusion(getter):
     }
 
     result = optimize(dsk, ["c"])
-    expected_task = (getter, "some-array", (15, slice(150, 160)))
-    assert any(_check_get_task_eq(v, expected_task) for v in result.values())
+    assert isinstance(result["c"], Alias)
     assert len(result) < len(dsk)
 
 
@@ -355,7 +355,13 @@ def test_dont_fuse_numpy_arrays():
         y = da.from_array(x, chunks=(10,))
 
         dsk = y.__dask_optimize__(y.dask, y.__dask_keys__())
-        assert sum(isinstance(v, np.ndarray) for v in dsk.values()) == 1
+        assert (
+            sum(
+                isinstance(v, DataNode) and isinstance(v.value, np.ndarray)
+                for v in dsk.values()
+            )
+            == 1
+        )
 
 
 def test_fuse_slices_with_alias(getter, getitem):
@@ -368,8 +374,6 @@ def test_fuse_slices_with_alias(getter, getitem):
     keys = [("dx2", 0)]
     dsk2 = optimize(dsk, keys)
     assert len(dsk2) == 3
-    fused_key = (dsk2.keys() - {"x", ("dx2", 0)}).pop()
-    assert _check_get_task_eq(dsk2[fused_key], (getter, "x", (slice(0, 4), 0)))
 
 
 def test_dont_fuse_fancy_indexing_in_getter_nofancy(getitem, getter_nofancy):
@@ -392,17 +396,10 @@ def test_fuse_getter_with_asarray(chunks):
     y = da.ones(10, chunks=chunks)
     z = x + y
     dsk = z.__dask_optimize__(z.dask, z.__dask_keys__())
-    for v in dsk.values():
-        s = str(v)
-        assert s.count("getitem") + s.count("getter") <= 1
-        if v is not x:
-            assert "1234567890" not in s
-    n_getters = len([v for v in dsk.values() if _is_getter_task(v)])
-    if y.npartitions > 1:
-        assert n_getters == y.npartitions
-    else:
-        assert n_getters == 0
-
+    assert any(
+        isinstance(v, DataNode) and isinstance(v.value, np.ndarray)
+        for v in dsk.values()
+    )
     assert_eq(z, x + 1)
 
 
diff --git a/dask/tests/test_base.py b/dask/tests/test_base.py
index 919d9e987..f03a30042 100644
--- a/dask/tests/test_base.py
+++ b/dask/tests/test_base.py
@@ -1007,7 +1007,6 @@ def test_optimizations_ctd():
     dsk1 = collections_to_dsk([x])
     with dask.config.set({"optimizations": [lambda dsk, keys: dsk]}):
         dsk2 = collections_to_dsk([x])
-
     assert dsk1 == dsk2
 
 
diff --git a/dask/tests/test_task_spec.py b/dask/tests/test_task_spec.py
index 0fae036fd..9ee6051e3 100644
--- a/dask/tests/test_task_spec.py
+++ b/dask/tests/test_task_spec.py
@@ -20,6 +20,7 @@ from dask._task_spec import (
     _get_dependencies,
     convert_legacy_graph,
     execute_graph,
+    fuse_linear_task_spec,
     parse_input,
     resolve_aliases,
 )
@@ -27,6 +28,7 @@ from dask.base import tokenize
 from dask.core import keys_in_tasks, reverse_dict
 from dask.optimization import SubgraphCallable
 from dask.sizeof import sizeof
+from dask.utils import funcname
 
 
 def convert_and_verify_keys(dsk):
@@ -871,9 +873,108 @@ def test_subgraph_dont_hold_in_memory_too_long_legacy():
     assert OnlyTwice.total == 10
 
 
+def test_linear_fusion():
+    tasks = [
+        io := Task("foo", func, 1),
+        second := Task("second", func, io.ref()),
+        Task("third", func, second.ref()),
+    ]
+    dsk = {t.key: t for t in tasks}
+    result = fuse_linear_task_spec(dsk, {"third"})
+    assert len(result) == 2
+    assert isinstance(result["third"], Alias)
+    assert (
+        isinstance(result["foo-second-third"], Task)
+        and funcname(result["foo-second-third"].func) == "_execute_subgraph"
+    )
+
+    # Data Nodes don't get fused
+
+    tasks = [
+        io := DataNode("foo", 1),
+        second := Task("second", func, io.ref()),
+        third := Task("third", func, second.ref()),
+        Task("fourth", func, third.ref()),
+    ]
+    dsk = {t.key: t for t in tasks}
+    result = fuse_linear_task_spec(dsk, {"fourth"})
+    assert len(result) == 3
+    assert isinstance(result["fourth"], Alias)
+    assert (
+        isinstance(result["second-third-fourth"], Task)
+        and funcname(result["second-third-fourth"].func) == "_execute_subgraph"
+    )
+    assert isinstance(result["foo"], DataNode)
+
+    # Branch, so no fusion
+    tasks.append(Task("branch", func, third.ref()))
+    dsk = {t.key: t for t in tasks}
+    result = fuse_linear_task_spec(dsk, {"fourth"})
+    assert len(result) == 5
+    assert "second-third" in result
+    assert isinstance(result["third"], Alias)
+
+    # Branch, so no fusion at all
+    tasks.append(Task("branch2", func, second.ref()))
+    dsk = {t.key: t for t in tasks}
+    result = fuse_linear_task_spec(dsk, {"fourth"})
+    assert len(result) == 6
+    assert not any("-" in k for k in dsk)
+
+
+def test_linear_fusion_intermediate_branch():
+    tasks = [
+        io := DataNode("foo", 1),
+        second := Task("second", func, io.ref()),
+        third := Task("third", func, second.ref()),
+        other := Task("other", func, 1),
+        Task("fourth", func, third.ref(), other.ref()),
+    ]
+    dsk = {t.key: t for t in tasks}
+    result = fuse_linear_task_spec(dsk, {"fourth"})
+    assert "second-third" in result
+    assert isinstance(result["fourth"], Task)
+
+
+def test_linear_fusion_two_branches():
+    tasks = [
+        left_one := DataNode("left_one", "a"),
+        left_two := Task("left_two", func, left_one.ref()),
+        right_one := DataNode("right_one", "a"),
+        right_two := Task("right_two", func, right_one.ref()),
+        middle := Task("middle", func, right_two.ref(), left_two.ref()),
+        third := Task("third", func, middle.ref()),
+        Task("fourth", func, third.ref()),
+    ]
+    dsk = {t.key: t for t in tasks}
+    result = fuse_linear_task_spec(dsk, {"fourth"})
+    assert "left_one-left_two" in result
+    assert "right_one-right_two" in result
+    assert isinstance(result["right_two"], Alias)
+    assert isinstance(result["left_two"], Alias)
+    assert isinstance(result["fourth"], Alias)
+    assert isinstance(result["third-fourth"], Task)
+
+
+def test_linear_fusion_multiple_outputs():
+    tasks = [
+        first := DataNode("first", "a"),
+        second := Task("second", func, first.ref()),
+        third := Task("third", func, second.ref()),
+        Task("fourth", func, third.ref()),
+    ]
+    dsk = {t.key: t for t in tasks}
+    result = fuse_linear_task_spec(dsk, {"fourth", "second"})
+    assert "first-second" in result
+    assert "second" in result
+    assert isinstance(result["second"], Alias)
+    assert isinstance(result["fourth"], Alias)
+    assert isinstance(result["first-second"], Task)
+    assert isinstance(result["third-fourth"], Task)
+
+
 def test_nested_containers():
     t = List(Task("key-1", func, "a", "b"), Task("key-2", func, "c", "d"))
-    assert not t.dependencies
     assert t() == ["a-b", "c-d"]
 
     t = List(Task("key-1", func, "a", TaskRef("b")), Task("key-2", func, "c", "d"))