Repository: dask/dask
Pull request #11569 (2024-11-28 17:13:22)

Convert `Bag` graphs to TaskSpec graphs during optimization
This optimization seems to be important. There are some interesting shananigans going on inside of the bag implementation with iterators so not lazifying  this can break fusion it seems.

I stumbled over this in https://github.com/dask/dask/pull/11568 because it looks like I'll have to touch bag optimizers to get this over the finishing line, similar to what we did with arrays.

~This is a minimal step in that direction~

Edit: I had to follow through and add the conversion layer to the optimization step as well. I killed off one of the optimizations (inlining lists). Many of these types of tasks will vanish when moving to https://github.com/dask/dask/pull/11568 so I didn't bother with the migration. If this becomes an issue, we'll have to look at how those tasks look like afterwards since it's not entirely clear to me at the moment.

From what I can tell, the most important test is `test_map_releases_element_references_as_soon_as_possible` which tests that tasks are indeed properly pipelined and intermediate memory is released as required. This is also what tripped me in https://github.com/dask/dask/pull/11568

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._

     15 files  ± 0       15 suites  ±0   4h 11m 11s ⏱️ + 1m 51s
 13 258 tests + 2   12 196 ✅ + 2   1 062 💤 ±0  0 ❌ ±0 
164 476 runs  +30  141 558 ✅ +34  22 918 💤  - 4  0 ❌ ±0 

Results for commit 882f39e6. ± Comparison against base commit 0f3e5ff6.

[test-results]:data:application/gzip;base64,H4sIAGauSGcC/03NTQ7CIBCG4asY1i7KFAbwMqaBISH2x1BYGe8upZV2+T6TfPNhPoy0sseNy/uNrTmks1yOQwrLXLtTvFC5pnrvQep/P9ds7YbADZ74Cu8NO4Rmfghjsa4BxbjEQ2Ke6zQKofDoNi24rA933KcBDG92ma59XbbLNIVUgmkNvjeESgN6I7UjhVZYo3RvtTKePDlnOLHvD8pcxKsaAQAA


Patch:
diff --git a/dask/_task_spec.py b/dask/_task_spec.py
index 82b896b61..8a1712137 100644
--- a/dask/_task_spec.py
+++ b/dask/_task_spec.py
@@ -915,3 +915,22 @@ def fuse_linear_task_spec(dsk, keys):
                 # 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
+
+
+def cull(
+    dsk: dict[KeyType, GraphNode], keys: Iterable[KeyType]
+) -> dict[KeyType, GraphNode]:
+    work = set(keys)
+    seen: set[KeyType] = set()
+    dsk2 = {}
+    wpop = work.pop
+    wupdate = work.update
+    sadd = seen.add
+    while work:
+        k = wpop()
+        if k in seen or k not in dsk:
+            continue
+        sadd(k)
+        dsk2[k] = v = dsk[k]
+        wupdate(v.dependencies)
+    return dsk2
diff --git a/dask/bag/core.py b/dask/bag/core.py
index 8ee3efafe..4bd31a258 100644
--- a/dask/bag/core.py
+++ b/dask/bag/core.py
@@ -38,7 +38,15 @@ from tlz import (
 )
 
 from dask import config
-from dask._task_spec import GraphNode
+from dask._task_spec import (
+    GraphNode,
+    List,
+    Task,
+    _execute_subgraph,
+    convert_legacy_graph,
+    cull,
+    fuse_linear_task_spec,
+)
 from dask.bag import chunk
 from dask.bag.avro import to_avro
 from dask.base import (
@@ -50,17 +58,15 @@ from dask.base import (
 )
 from dask.blockwise import blockwise
 from dask.context import globalmethod
-from dask.core import flatten, get_dependencies, istask, quote, reverse_dict
+from dask.core import flatten, istask, quote
 from dask.delayed import Delayed, unpack_collections
 from dask.highlevelgraph import HighLevelGraph
-from dask.optimization import cull, fuse, inline
 from dask.sizeof import sizeof
 from dask.typing import Graph, NestedKeys, no_default
 from dask.utils import (
     apply,
     digit,
     ensure_bytes,
-    ensure_dict,
     ensure_unicode,
     funcname,
     get_default_shuffle_method,
@@ -96,17 +102,47 @@ def lazify_task(task, start=True):
     (<built-in function sum>, (<class 'map'>, <function inc at ...>, [1, 2, 3]))
     """
     if isinstance(task, GraphNode):
-        return task
-    if type(task) is list and len(task) < 50:
-        return [lazify_task(arg, False) for arg in task]
-    if not istask(task):
-        return task
-    head, tail = task[0], task[1:]
-    if not start and head in (list, reify):
-        task = task[1]
-        return lazify_task(*tail, start=False)
+        if isinstance(task, List) and len(task.args) < 50:
+            return List(*[lazify_task(arg, False) for arg in task.args])
+        if not isinstance(task, Task):
+            return task
+        if not start and task.func in (list, reify) and isinstance(task.args[0], Task):
+            assert len(task.args) == 1
+            task = task.args[0]
+        if task.func is _execute_subgraph:
+            subgraph = task.args[0]
+            outkey = task.args[1]
+            # If there is a reify at the output of the subgraph we don't want to act
+            final_task = lazify_task(subgraph[outkey], True)
+            subgraph = {
+                k: lazify_task(v, False) for k, v in subgraph.items() if k != outkey
+            }
+            subgraph[outkey] = final_task
+            return Task(
+                task.key,
+                _execute_subgraph,
+                subgraph,
+                outkey,
+                *task.args[2:],
+                **task.kwargs,
+            )
+        return Task(
+            task.key,
+            task.func,
+            *[lazify_task(arg, False) for arg in task.args],
+            **task.kwargs,
+        )
     else:
-        return (head,) + tuple(lazify_task(arg, False) for arg in tail)
+        if type(task) is list and len(task) < 50:
+            return [lazify_task(arg, False) for arg in task]
+        if not istask(task):
+            return task
+        head, tail = task[0], task[1:]
+        if not start and head in (list, reify):
+            task = task[1]
+            return lazify_task(*tail, start=False)
+        else:
+            return (head,) + tuple(lazify_task(arg, False) for arg in tail)
 
 
 def lazify(dsk):
@@ -120,47 +156,14 @@ def lazify(dsk):
     return valmap(lazify_task, dsk)
 
 
-def inline_singleton_lists(dsk, keys, dependencies=None):
-    """Inline lists that are only used once.
-
-    >>> d = {'b': (list, 'a'),
-    ...      'c': (sum, 'b', 1)}
-    >>> inline_singleton_lists(d, 'c')
-    {'c': (<built-in function sum>, (<class 'list'>, 'a'), 1)}
-
-    Pairs nicely with lazify afterwards.
-    """
-    if dependencies is None:
-        dependencies = {k: get_dependencies(dsk, task=v) for k, v in dsk.items()}
-    dependents = reverse_dict(dependencies)
-
-    inline_keys = {
-        k
-        for k, v in dsk.items()
-        if istask(v)
-        and not isinstance(v, GraphNode)
-        and v
-        and v[0] is list
-        and len(dependents[k]) == 1
-    }
-    inline_keys.difference_update(flatten(keys))
-    dsk = inline(dsk, inline_keys, inline_constants=False)
-    for k in inline_keys:
-        del dsk[k]
-    return dsk
-
-
-def optimize(dsk, keys, fuse_keys=None, rename_fused_keys=None, **kwargs):
+def optimize(dsk, keys, fuse_keys=None, **kwargs):
     """Optimize a dask from a dask Bag."""
-    dsk = ensure_dict(dsk)
-    dsk2, dependencies = cull(dsk, keys)
-    kwargs = {}
-    if rename_fused_keys is not None:
-        kwargs["rename_keys"] = rename_fused_keys
-    dsk3, dependencies = fuse(dsk2, keys + (fuse_keys or []), dependencies, **kwargs)
-    dsk4 = inline_singleton_lists(dsk3, keys, dependencies)
-    dsk5 = lazify(dsk4)
-    return dsk5
+    dsk = convert_legacy_graph(dsk)
+    keys = list(flatten(keys))
+    dsk2 = cull(dsk, keys)
+    dsk3 = fuse_linear_task_spec(dsk2, keys + (fuse_keys or []))
+    dsk4 = lazify(dsk3)
+    return dsk4
 
 
 def _to_textfiles_chunk(data, lazy_file, last_endline):


Test patch:
diff --git a/dask/bag/tests/test_bag.py b/dask/bag/tests/test_bag.py
index fe1b5ffc9..20f43b89f 100644
--- a/dask/bag/tests/test_bag.py
+++ b/dask/bag/tests/test_bag.py
@@ -19,14 +19,13 @@ from tlz import groupby, identity, join, merge, pluck, unique, valmap
 
 import dask
 import dask.bag as db
+from dask._task_spec import List, Task
 from dask.bag.core import (
     Bag,
     collect,
     from_delayed,
-    inline_singleton_lists,
     lazify,
     lazify_task,
-    optimize,
     partition,
     reduceby,
     reify,
@@ -574,6 +573,47 @@ def test_random_sample_random_state():
 
 
 def test_lazify_task():
+    task = Task(None, sum, Task(None, reify, Task(None, map, inc, [1, 2, 3])))
+    assert lazify_task(task) == Task(None, sum, Task(None, map, inc, [1, 2, 3]))
+
+    task = Task(None, reify, Task(None, map, inc, [1, 2, 3]))
+    assert lazify_task(task) == task
+
+    task = Task(
+        None,
+        reify,
+        Task(None, map, inc, Task(None, reify, Task(None, filter, iseven, "y"))),
+    )
+    expected = Task(None, reify, Task(None, map, inc, Task(None, filter, iseven, "y")))
+    assert lazify_task(task) == expected
+
+    task = Task(
+        None,
+        sum,
+        List(
+            Task(
+                None,
+                reify,
+                Task(None, map, inc, Task(None, reify, Task(None, filter, iseven, 1))),
+            ),
+            Task(
+                None,
+                reify,
+                Task(None, map, inc, Task(None, reify, Task(None, filter, iseven, 2))),
+            ),
+        ),
+    )
+    assert lazify_task(task) == Task(
+        None,
+        sum,
+        List(
+            Task(None, map, inc, Task(None, filter, iseven, 1)),
+            Task(None, map, inc, Task(None, filter, iseven, 2)),
+        ),
+    )
+
+
+def test_lazify_task_legacy():
     task = (sum, (reify, (map, inc, [1, 2, 3])))
     assert lazify_task(task) == (sum, (map, inc, [1, 2, 3]))
 
@@ -598,40 +638,6 @@ def test_lazify():
     assert lazify(a) == b
 
 
-def test_inline_singleton_lists():
-    inp = {"b": (list, "a"), "c": (f, "b", 1)}
-    out = {"c": (f, (list, "a"), 1)}
-    assert inline_singleton_lists(inp, ["c"]) == out
-
-    out = {"c": (f, "a", 1)}
-    assert optimize(inp, ["c"], rename_fused_keys=False) == out
-
-    # If list is an output key, don't fuse it
-    assert inline_singleton_lists(inp, ["b", "c"]) == inp
-    assert optimize(inp, ["b", "c"], rename_fused_keys=False) == inp
-
-    inp = {"b": (list, "a"), "c": (f, "b", 1), "d": (f, "b", 2)}
-    assert inline_singleton_lists(inp, ["c", "d"]) == inp
-
-    # Doesn't inline constants
-    inp = {"b": (4, 5), "c": (f, "b")}
-    assert inline_singleton_lists(inp, ["c"]) == inp
-
-
-def test_rename_fused_keys_bag():
-    inp = {"b": (list, "a"), "c": (f, "b", 1)}
-
-    outp = optimize(inp, ["c"], rename_fused_keys=False)
-    assert outp.keys() == {"c"}
-    assert outp["c"][1:] == ("a", 1)
-
-    with dask.config.set({"optimization.fuse.rename-keys": False}):
-        assert optimize(inp, ["c"]) == outp
-
-    # By default, fused keys are renamed
-    assert optimize(inp, ["c"]) != outp
-
-
 def test_take():
     assert list(b.take(2)) == [0, 1]
     assert b.take(2) == (0, 1)
@@ -1156,32 +1162,55 @@ def test_to_delayed():
     assert t.compute() == 21
 
 
-def test_to_delayed_optimize_graph(tmpdir):
+def test_to_delayed_optimize_graph(tmpdir, monkeypatch):
+    calls = 0
+    from dask.bag.core import reify
+
+    def count_reify(*args, **kwargs):
+        nonlocal calls
+        calls += 1
+        return reify(*args, **kwargs)
+
+    monkeypatch.setattr(dask.bag.core, "reify", count_reify)
     b = db.from_sequence([1, 2, 3, 4, 5, 6], npartitions=1)
     b2 = b.map(inc).map(inc).map(inc)
 
     [d] = b2.to_delayed()
-    text = str(dict(d.dask))
-    assert text.count("reify") == 1
+
+    calls = 0
+    d.compute()
+    assert calls == 1
+
+    calls = 0
+    b = db.from_sequence([1, 2, 3, 4, 5, 6], npartitions=1)
+    b2 = b.map(inc).map(inc).map(inc)
+    [d] = b2.to_delayed()
     assert d.__dask_layers__() != b2.__dask_layers__()
     [d2] = b2.to_delayed(optimize_graph=False)
     assert dict(d2.dask) == dict(b2.dask)
     assert d2.__dask_layers__() == b2.__dask_layers__()
-    assert d.compute() == d2.compute()
+    calls = 0
+    result = d.compute()
+    assert calls == 1
+    calls = 0
+    result2 = d2.compute()
+    assert calls == 3
+    assert result == result2
 
     x = b2.sum()
     d = x.to_delayed()
-    text = str(dict(d.dask))
     assert d.__dask_layers__() == x.__dask_layers__()
-    assert text.count("reify") == 0
     d2 = x.to_delayed(optimize_graph=False)
     assert dict(d2.dask) == dict(x.dask)
     assert d2.__dask_layers__() == x.__dask_layers__()
+
+    calls = 0
     assert d.compute() == d2.compute()
+    assert calls == 3
 
+    calls = 0
     [d] = b2.to_textfiles(str(tmpdir), compute=False)
-    text = str(dict(d.dask))
-    assert text.count("reify") <= 0
+    assert calls == 0
 
 
 def test_from_delayed():
diff --git a/dask/tests/test_base.py b/dask/tests/test_base.py
index f03a30042..752288805 100644
--- a/dask/tests/test_base.py
+++ b/dask/tests/test_base.py
@@ -371,7 +371,7 @@ def test_compute_no_opt():
     # Otherwise, the lengths below would be 4 and 0.
     assert len([k for k in keys if "mul" in k[0]]) == 8
     assert len([k for k in keys if "add" in k[0]]) == 4
-    assert len([k for k in keys if "add-mul" in k[0]]) == 4  # See? Renamed
+    assert len([k for k in keys if "add-from_sequence-mul" in k[0]]) == 4
 
 
 @pytest.mark.skipif("not da")