Repository: dask/dask
Pull request #11917 (2025-04-28 10:39:36)

Fix delayed parsing for futures
- Closes https://github.com/dask/dask/issues/11910
- Closes https://github.com/dask/dask/issues/11908
- Supersedes https://github.com/dask/dask/pull/11912
- Supersedes https://github.com/dask/dask/pull/11911

@TomAugspurger I hope you don't mind. This shoudl fix the two issues you have PRs open for. I figured patching `unpack_collections` properly is more durable

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

     5 files  ± 0       5 suites  ±0   1h 35m 17s ⏱️ +3s
18 015 tests + 2  16 798 ✅ ±0  1 215 💤 ±0  2 ❌ +2 
89 590 runs  +10  83 110 ✅ +4  6 474 💤 ±0  6 ❌ +6 

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

Results for commit 1e27aad7. ± Comparison against base commit d4994ade.

[test-results]:data:application/gzip;base64,H4sIAPhgD2gC/03MTQ7CIBCG4asY1i5gCgx4GYOFSYj9MRRWxrsLUdou32cm35tRnMLGbhd1vbCtxLyHL8nluC4tUWCVesvtKgwXqvd9K+PYUKM1Bz7jqyGcHsnFqRrsEFJaUxVeJZWlTRurbM++bAYhDvsNa4my039X9z7Pjus8x1yDiQDonMcBSHhNUjpFYEETWv4IdgDjyXIA9vkC4G7Z8RQBAAA=


Patch:
diff --git a/dask/_task_spec.py b/dask/_task_spec.py
index 89b0a8b4b..bf404f166 100644
--- a/dask/_task_spec.py
+++ b/dask/_task_spec.py
@@ -67,8 +67,6 @@ from dask.typing import Key as KeyType
 from dask.utils import funcname, is_namedtuple_instance
 
 _T = TypeVar("_T")
-_T_GraphNode = TypeVar("_T_GraphNode", bound="GraphNode")
-_T_Iterable = TypeVar("_T_Iterable", bound=Iterable)
 
 
 # Ported from more-itertools
@@ -817,7 +815,7 @@ class Task(GraphNode):
             )
 
 
-class NestedContainer(Task):
+class NestedContainer(Task, Iterable):
     constructor: Callable
     klass: type
     __slots__ = tuple(__annotations__)
@@ -888,6 +886,9 @@ class NestedContainer(Task):
     def to_container(*args, constructor):
         return constructor(args)
 
+    def __iter__(self):
+        yield from self.args
+
 
 class List(NestedContainer):
     constructor = klass = list
@@ -901,15 +902,33 @@ class Set(NestedContainer):
     constructor = klass = set
 
 
-class Dict(NestedContainer):
+class Dict(NestedContainer, Mapping):
     klass = dict
 
     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())))
+            assert not kwargs
+            if len(args) == 1:
+                args = args[0]
+                if isinstance(args, dict):  # type: ignore
+                    args = tuple(itertools.chain(*args.items()))  # type: ignore
+                elif isinstance(args, (list, tuple)):
+                    if all(
+                        len(el) == 2 if isinstance(el, (list, tuple)) else False
+                        for el in args
+                    ):
+                        args = tuple(itertools.chain(*args))
+                else:
+                    raise ValueError("Invalid argument provided")
+
+            if len(args) % 2 != 0:
+                raise ValueError("Invalid number of arguments provided")
+
+        elif kwargs:
+            assert not args
+            args = tuple(itertools.chain(*kwargs.items()))
+
+        super().__init__(*args)
 
     def __repr__(self):
         values = ", ".join(f"{k}: {v}" for k, v in batched(self.args, 2, strict=True))
@@ -924,11 +943,25 @@ class Dict(NestedContainer):
         if not subs_filtered:
             return self
 
-        new = {}
+        new_args = []
+        for arg in self.args:
+            new_arg = (
+                arg.substitute(subs_filtered) if isinstance(arg, GraphNode) else arg
+            )
+            new_args.append(new_arg)
+        return type(self)(new_args)
+
+    def __iter__(self):
+        yield from self.args[::2]
+
+    def __getitem__(self, key):
         for k, v in batched(self.args, 2, strict=True):
-            v = v.substitute(subs_filtered) if isinstance(v, GraphNode) else v
-            new[k] = v
-        return type(self)(new)
+            if k == key:
+                return v
+        raise KeyError(key)
+
+    def __len__(self):
+        return len(self.args) // 2
 
     @staticmethod
     def constructor(args):
diff --git a/dask/delayed.py b/dask/delayed.py
index b05d0d1ac..6a5482053 100644
--- a/dask/delayed.py
+++ b/dask/delayed.py
@@ -15,6 +15,7 @@ from dask import base, config, utils
 from dask._expr import FinalizeCompute, ProhibitReuse, _ExprSequence
 from dask._task_spec import (
     DataNode,
+    Dict,
     GraphNode,
     List,
     Task,
@@ -158,7 +159,7 @@ def unpack_collections(expr, _return_collections=True):
 
     >>> task, collections = unpack_collections({a: 1, b: 2})
     >>> task
-    <Task None dict(List((List((Alias('a'), 1)), List((Alias('b'), 2)))))>
+    Dict(Alias('a'): 1, Alias('b'): 2)
     >>> collections
     (Delayed('a'), Delayed('b'))
     """
@@ -208,9 +209,10 @@ def unpack_collections(expr, _return_collections=True):
         )
 
         collections = tuple(toolz.unique(toolz.concat(collections), key=id))
-        if not collections:
-            return expr, ()
+        # The List constructor also checks for futures
         args = List(*args)
+        if not collections and not args.dependencies:
+            return expr, ()
         if _return_collections:
             args, collections = _finalize_args_collections(args, collections)
         # Ensure output type matches input type
@@ -219,20 +221,25 @@ def unpack_collections(expr, _return_collections=True):
         return args, collections
 
     if typ is dict:
-        args, collections = unpack_collections(
-            [[k, v] for k, v in expr.items()], _return_collections=False
+        keyargs, kcollections = unpack_collections(
+            [k for k in expr.keys()], _return_collections=False
         )
-        if not collections:
+        valargs, valcollections = unpack_collections(
+            [v for v in expr.values()], _return_collections=False
+        )
+        collections = kcollections + valcollections
+        args = Dict([[k, v] for k, v in zip(keyargs, valargs)])
+        if not collections and not args.dependencies:
             return expr, ()
         if _return_collections:
             args, collections = _finalize_args_collections(args, collections)
-        return Task(None, dict, args), collections
+        return args, collections
 
     if typ is slice:
         args, collections = unpack_collections(
             [expr.start, expr.stop, expr.step], _return_collections=False
         )
-        if not collections:
+        if not collections and not isinstance(args, GraphNode):
             return expr, ()
 
         if _return_collections:
@@ -248,7 +255,7 @@ def unpack_collections(expr, _return_collections=True):
             ],
             _return_collections=False,
         )
-        if not collections:
+        if not collections and not isinstance(args, GraphNode):
             return expr, ()
 
         if _return_collections:
@@ -620,7 +627,9 @@ def delayed(obj, name=None, pure=None, nout=None, traverse=True):
     if not (nout is None or (type(nout) is int and nout >= 0)):
         raise ValueError("nout must be None or a non-negative integer, got %s" % nout)
     if task is obj:
-        if not name:
+        if isinstance(obj, TaskRef):
+            name = obj.key
+        elif not name:
             try:
                 prefix = obj.__name__
             except AttributeError:
@@ -647,10 +656,6 @@ def right(method):
     return partial(_swap, method)
 
 
-def _apply2(func, *args, dask_kwargs):
-    return func(*args, **dask_kwargs)
-
-
 def optimize(dsk, keys, **kwargs):
     if not isinstance(keys, (list, set)):
         keys = [keys]
@@ -815,10 +820,7 @@ def call_function(func, func_token, args, kwargs, pure=None, nout=None):
 
     dask_kwargs, collections2 = unpack_collections(kwargs)
     collections.extend(collections2)
-    if dask_kwargs:
-        task = Task(name, _apply2, func, *args2, dask_kwargs=dask_kwargs)
-    else:
-        task = Task(name, func, *args2, **kwargs)
+    task = Task(name, func, *args2, **dask_kwargs)
 
     graph = HighLevelGraph.from_collections(
         name, {name: task}, dependencies=collections
@@ -838,9 +840,11 @@ class DelayedLeaf(Delayed):
 
     @property
     def dask(self):
-        return HighLevelGraph.from_collections(
-            self._key, {self._key: DataNode(self._key, self._obj)}, dependencies=()
-        )
+        if isinstance(self._obj, (TaskRef, GraphNode)):
+            dsk = {self._key: self._obj}
+        else:
+            dsk = {self._key: DataNode(self._key, self._obj)}
+        return HighLevelGraph.from_collections(self._key, dsk, dependencies=())
 
     def __call__(self, *args, **kwargs):
         return call_function(


Test patch:
diff --git a/dask/tests/test_distributed.py b/dask/tests/test_distributed.py
index 63635577e..ea14cf4f4 100644
--- a/dask/tests/test_distributed.py
+++ b/dask/tests/test_distributed.py
@@ -1063,6 +1063,14 @@ async def test_release_persisted_futures_without_gc(c, s, a, b):
         gc.enable()
 
 
+@gen_cluster(client=True)
+async def test_delayed_future_with_kwargs(c, s, a, b):
+    fut = await c.scatter(1)
+
+    result = dask.delayed(inc)(x=fut)
+    assert await c.compute(result) == 2
+
+
 @gen_cluster(client=True)
 async def test_fusion_barrier_task(c, s, a, b):
     np = pytest.importorskip("numpy")
@@ -1074,3 +1082,12 @@ async def test_fusion_barrier_task(c, s, a, b):
     x2 = da.rechunk(x, chunks=new, method="p2p")
     result = await c.compute(x2)
     da.assert_eq(result, a)
+
+
+@gen_cluster(client=True)
+async def test_delayed_future(c, s, a, b):
+    fut = await c.scatter(1)
+    result = dask.delayed(fut)
+    res = await c.compute(result)
+    assert res == 1
+    assert result.key == fut.key
diff --git a/dask/tests/test_task_spec.py b/dask/tests/test_task_spec.py
index 4d349c149..becbefca8 100644
--- a/dask/tests/test_task_spec.py
+++ b/dask/tests/test_task_spec.py
@@ -965,6 +965,8 @@ def test_nested_containers():
     assert t.dependencies == {"b"}
     assert t({"b": "b"}) == ["a-b", "c-d"]
 
+
+def test_dict_class():
     t = Dict(k=Task("key-1", func, "a", "b"), v=Task("key-2", func, "c", "d"))
     assert not t.dependencies
     assert t() == {"k": "a-b", "v": "c-d"}
@@ -997,6 +999,21 @@ def test_nested_containers():
     assert t == t2
     assert tokenize(t) == tokenize(t2)
 
+    d = Dict(
+        [
+            ["k", Task("key-1", func, "a", TaskRef("b"))],
+            ["v", Task("key-2", func, TaskRef("c"), "d")],
+        ]
+    )
+    assert d.dependencies == {"c", "b"}
+    assert d({"b": "b", "c": "c"}) == {"k": "a-b", "v": "c-d"}
+
+    d = Dict([("columns", ["a", "b"])])
+    assert d() == {"columns": ["a", "b"]}
+
+    d = Dict([["columns", ["a", "b"]]])
+    assert d() == {"columns": ["a", "b"]}
+
 
 def test_block_io_fusion():