Repository: dask/dask
Pull request #11906 (2025-04-24 12:02:12)

Ensure only HLGs are probited reuse
- [ ] Closes https://github.com/dask/dask/issues/11902


Patch:
diff --git a/dask/_expr.py b/dask/_expr.py
index bb12458de..b96467c40 100644
--- a/dask/_expr.py
+++ b/dask/_expr.py
@@ -1076,31 +1076,24 @@ class _HLGExprSequence(Expr):
         from dask.highlevelgraph import HighLevelGraph
 
         groups = toolz.groupby(
-            lambda x: (
-                x.low_level_optimizer if isinstance(x, HLGExpr) else None,
-                x.postcompute,
-            ),
+            lambda x: x.low_level_optimizer if isinstance(x, HLGExpr) else None,
             self.operands,
         )
         exprs = []
         changed = False
-        for (optimizer, postcompute), group in groups.items():
+        for optimizer, group in groups.items():
             if len(group) > 1:
-                changed = True
-                graphs = []
-                for expr in group:
-                    graphs.append(expr.hlg)
+                graphs = [expr.hlg for expr in group]
 
+                changed = True
                 dsk = HighLevelGraph.merge(*graphs)
-                keys = [v.__dask_keys__() for v in group]
-                exprs.append(
-                    _HLGExprGroup(
-                        dsk=dsk,
-                        low_level_optimizer=optimizer,
-                        output_keys=keys,
-                        postcompute=postcompute,
-                    )
+                hlg_group = _HLGExprGroup(
+                    dsk=dsk,
+                    low_level_optimizer=optimizer,
+                    output_keys=[v.__dask_keys__() for v in group],
+                    postcompute=[g.postcompute for g in group],
                 )
+                exprs.append(hlg_group)
             else:
                 exprs.append(group[0])
         if not changed:
@@ -1302,15 +1295,22 @@ class HLGFinalizeCompute(HLGExpr):
         layers = expr.dsk.layers.copy()
         deps = expr.dsk.dependencies.copy()
         keys = expr.__dask_keys__()
-        func, extra_args = expr.postcompute
-        t = Task(self._name, func, _convert_dask_keys(keys), *extra_args)
+        if isinstance(expr.postcompute, list):
+            postcomputes = expr.postcompute
+        else:
+            postcomputes = [expr.postcompute]
+        tasks = [
+            Task(self._name, func, _convert_dask_keys(keys), *extra_args)
+            for func, extra_args in postcomputes
+        ]
         from dask.highlevelgraph import HighLevelGraph, MaterializedLayer
 
-        layers[t.key] = MaterializedLayer({t.key: t})
         leafs = set(deps)
         for val in deps.values():
             leafs -= val
-        deps[t.key] = leafs
+        for t in tasks:
+            layers[t.key] = MaterializedLayer({t.key: t})
+            deps[t.key] = leafs
         return HighLevelGraph(layers, dependencies=deps)
 
     def __dask_keys__(self):
@@ -1324,6 +1324,7 @@ class ProhibitReuse(Expr):
     """
 
     _parameters = ["expr"]
+    _ALLOWED_TYPES = [HLGExpr, LLGExpr, HLGFinalizeCompute, _HLGExprSequence]
 
     def __dask_keys__(self):
         return self._modify_keys(self.expr.__dask_keys__())
@@ -1345,18 +1346,41 @@ class ProhibitReuse(Expr):
             k = str(k)
         return f"{k}-{self._suffix}"
 
+    def _simplify_down(self):
+        # FIXME: Shuffling cannot be rewritten since the barrier key is
+        # hardcoded. Skipping this here should do the trick most of the time
+        if not isinstance(
+            self.expr,
+            tuple(self._ALLOWED_TYPES),
+        ):
+            return self.expr
+
     def __dask_graph__(self):
+        try:
+            from distributed.shuffle._core import P2PBarrierTask
+        except ModuleNotFoundError:
+            P2PBarrierTask = type(None)
         dsk = convert_legacy_graph(self.expr.__dask_graph__())
 
         subs = {old_key: self._modify_keys(old_key) for old_key in dsk}
-        dsk2 = {
-            new_key: Task(
+        dsk2 = {}
+        for old_key, new_key in subs.items():
+            t = dsk[old_key]
+            if isinstance(t, P2PBarrierTask):
+                warnings.warn(
+                    "Cannot block reusing for graphs including a "
+                    "P2PBarrierTask. This may cause unexpected results. "
+                    "This typically happens when converting a dask "
+                    "DataFrame to delayed objects.",
+                    UserWarning,
+                )
+                return dsk
+            dsk2[new_key] = Task(
                 new_key,
                 ProhibitReuse._identity,
-                dsk.pop(old_key).substitute(subs),
+                t.substitute(subs),
             )
-            for old_key, new_key in subs.items()
-        }
+
         dsk2.update(dsk)
         return dsk2
 
diff --git a/dask/_task_spec.py b/dask/_task_spec.py
index f83011862..89b0a8b4b 100644
--- a/dask/_task_spec.py
+++ b/dask/_task_spec.py
@@ -54,6 +54,7 @@ should contain the dependencies of the task.
     t({"a": 1, "b": 2}) == 3
 
 """
+import functools
 import itertools
 import sys
 from collections import defaultdict
@@ -642,7 +643,6 @@ class Task(GraphNode):
         /,
         *args: Any,
         _data_producer: bool = False,
-        _dependencies: set | frozenset | None = None,
         **kwargs: Any,
     ):
         self.key = key
@@ -655,11 +655,12 @@ class Task(GraphNode):
         self.kwargs = {
             k: Alias(v.key) if isinstance(v, TaskRef) else v for k, v in kwargs.items()
         }
-        if _dependencies is None:
-            _dependencies = set()
-            for a in itertools.chain(self.args, self.kwargs.values()):
-                if isinstance(a, GraphNode):
-                    _dependencies.update(a.dependencies)
+        _dependencies: set[KeyType] | None = None
+        for a in itertools.chain(self.args, self.kwargs.values()):
+            if isinstance(a, GraphNode):
+                if _dependencies is None:
+                    _dependencies = set()
+                _dependencies.update(a.dependencies)
         if _dependencies:
             self._dependencies = frozenset(_dependencies)
         else:
@@ -681,7 +682,6 @@ class Task(GraphNode):
             self.key,
             self.func,
             *self.args,
-            _dependencies=self._dependencies,
             **self.kwargs,
         )
 
@@ -784,6 +784,10 @@ class Task(GraphNode):
         subs_filtered = {
             k: v for k, v in subs.items() if k in self.dependencies and k != v
         }
+        extras = _extra_args(type(self))  # type: ignore
+        extra_kwargs = {
+            name: getattr(self, name) for name in extras if name not in {"key", "func"}
+        }
         if subs_filtered:
             new_args = tuple(
                 a.substitute(subs_filtered) if isinstance(a, GraphNode) else a
@@ -797,8 +801,8 @@ class Task(GraphNode):
                 key or self.key,
                 self.func,
                 *new_args,
-                _data_producer=self.data_producer,
                 **new_kwargs,
+                **extra_kwargs,
             )
         elif key is None or key == self.key:
             return self
@@ -808,8 +812,8 @@ class Task(GraphNode):
                 key,
                 self.func,
                 *self.args,
-                _data_producer=self.data_producer,
                 **self.kwargs,
+                **extra_kwargs,
             )
 
 
@@ -822,7 +826,6 @@ class NestedContainer(Task):
         self,
         /,
         *args: Any,
-        _dependencies: set | frozenset | None = None,
         **kwargs: Any,
     ):
         if len(args) == 1 and isinstance(args[0], self.klass):
@@ -832,7 +835,6 @@ class NestedContainer(Task):
             self.to_container,
             *args,
             constructor=self.constructor,
-            _dependencies=_dependencies,
             **kwargs,
         )
 
@@ -902,16 +904,12 @@ class Set(NestedContainer):
 class Dict(NestedContainer):
     klass = dict
 
-    def __init__(
-        self, /, *args: Any, _dependencies: set | frozenset | None = None, **kwargs: Any
-    ):
+    def __init__(self, /, *args: Any, **kwargs: Any):
         if args:
             if len(args) > 1:
                 raise ValueError("Dict can only take one positional argument")
             kwargs = args[0]
-        super().__init__(
-            *(itertools.chain(*kwargs.items())), _dependencies=_dependencies
-        )
+        super().__init__(*(itertools.chain(*kwargs.items())))
 
     def __repr__(self):
         values = ", ".join(f"{k}: {v}" for k, v in batched(self.args, 2, strict=True))
@@ -1140,3 +1138,20 @@ def cull(
         dsk2[k] = v = dsk[k]
         wupdate(v.dependencies)
     return dsk2
+
+
+@functools.cache
+def _extra_args(typ: type) -> set[str]:
+    import inspect
+
+    sig = inspect.signature(typ)
+    extras = set()
+    for name, param in sig.parameters.items():
+        if param.kind in (
+            inspect.Parameter.VAR_POSITIONAL,
+            inspect.Parameter.VAR_KEYWORD,
+        ):
+            continue
+        if name in typ.get_all_slots():  # type: ignore
+            extras.add(name)
+    return extras


Test patch:
diff --git a/dask/array/tests/test_xarray.py b/dask/array/tests/test_xarray.py
index 007b4d8ad..f35e1bd42 100644
--- a/dask/array/tests/test_xarray.py
+++ b/dask/array/tests/test_xarray.py
@@ -125,3 +125,32 @@ def test_xarray_blockwise_fusion_store(compute):
             scheduler = partial(custom_scheduler_get, expected=expected)
             stored = y.to_zarr(dirname, compute=False)
             dask.compute(stored, scheduler=scheduler)
+
+
+@pytest.mark.parametrize("wrap_xarray", [False, True])
+def test_shared_tasks(wrap_xarray):
+    total_calls = 0
+
+    def my_func(a: np.ndarray) -> np.ndarray:
+        nonlocal total_calls
+        total_calls += 1
+        return a
+
+    in1 = da.zeros((5, 5), chunks=2)
+    res = da.map_blocks(
+        my_func, in1, meta=np.array((), dtype=in1.dtype), dtype=in1.dtype
+    )
+    out1 = res + 1
+    out2 = res + 2
+    if wrap_xarray:
+        out1 = xr.DataArray(out1)
+        out2 = xr.DataArray(out2)
+
+    comp_res = dask.compute(out1, out2)
+
+    if wrap_xarray:
+        assert isinstance(comp_res[0], xr.DataArray)
+        assert isinstance(comp_res[0].data, np.ndarray)
+    else:
+        assert isinstance(comp_res[0], np.ndarray)
+    assert total_calls == in1.blocks.size
diff --git a/dask/tests/test_delayed.py b/dask/tests/test_delayed.py
index 0972b7eaa..bc0b28a0a 100644
--- a/dask/tests/test_delayed.py
+++ b/dask/tests/test_delayed.py
@@ -929,3 +929,35 @@ def test_delayed_fusion():
         result = dask.compute(obj2)
     assert len(dsk2) == 2
     assert dask.compute(obj) == result
+
+
+def test_p2p_as_delayed():
+    pytest.importorskip("distributed")
+    pytest.importorskip("numpy")
+    pytest.importorskip("dask.dataframe")
+    from distributed.utils_test import gen_cluster
+
+    import dask.dataframe as dd
+
+    @gen_cluster(client=True)
+    async def _test(c, s, *workers):
+        test_df = dd.from_dict(
+            {
+                "partition": [0, 1, 2, 3, 0, 1, 2, 3],
+                "value": [1, 2, 3, 4, 5, 6, 7, 8],
+            },
+            npartitions=2,
+        )
+        part_df = test_df.shuffle("partition", force=True, shuffle_method="p2p")
+
+        @delayed
+        def delayed_func(x):
+            return x
+
+        delay_df = delayed_func(part_df)
+        dd.utils.assert_eq(
+            await c.gather(c.compute(delay_df)),
+            await c.gather(c.compute(part_df)),
+        )
+
+    _test()
diff --git a/dask/tests/test_expr.py b/dask/tests/test_expr.py
index af90c712e..0d0aa30a8 100644
--- a/dask/tests/test_expr.py
+++ b/dask/tests/test_expr.py
@@ -151,13 +151,17 @@ class FooExpr(Expr):
 
 def test_prohibit_reuse():
     once = FooExpr()
-    dsk = _ExprSequence(once, ProhibitReuse(once)).optimize().__dask_graph__()
-
-    assert len(dsk) == 2
-    first = dsk.pop("foo")()
-    key, val = dsk.popitem()
-    assert key.startswith("foo") and key != "foo"
-    # We don't want to chain anything but actually _hide_ the task
-    assert not val.dependencies
-    # Task is wrapped
-    assert val() is first
+    ProhibitReuse._ALLOWED_TYPES.append(FooExpr)
+    try:
+        dsk = _ExprSequence(once, ProhibitReuse(once)).optimize().__dask_graph__()
+
+        assert len(dsk) == 2
+        first = dsk.pop("foo")()
+        key, val = dsk.popitem()
+        assert key.startswith("foo") and key != "foo"
+        # We don't want to chain anything but actually _hide_ the task
+        assert not val.dependencies
+        # Task is wrapped
+        assert val() is first
+    finally:
+        ProhibitReuse._ALLOWED_TYPES.remove(FooExpr)
diff --git a/dask/tests/test_task_spec.py b/dask/tests/test_task_spec.py
index 1b7667638..4d349c149 100644
--- a/dask/tests/test_task_spec.py
+++ b/dask/tests/test_task_spec.py
@@ -1095,3 +1095,18 @@ def test_substitute_nested():
 @pytest.mark.parametrize("Container", [Dict, List, Set, Tuple])
 def test_nested_containers_empty(Container):
     assert Container(Container.klass())() == Container.klass()
+
+
+class MySubclass(Task):
+    __slots__ = ("custom_kwarg_only",)
+
+    def __init__(self, key, func, /, *args, custom_kwarg_only, **kwargs):
+        self.custom_kwarg_only = custom_kwarg_only
+        super().__init__(key, func, *args, **kwargs)
+
+
+def test_substitute_subclasses():
+    t = MySubclass("key", func, "a", TaskRef("b"), custom_kwarg_only="foo")
+    t2 = t.substitute({"b": "c"})
+    assert t2.custom_kwarg_only == "foo"
+    assert t2({"a": "a", "c": "b"}) == t({"a": "a", "b": "b"})