Repository: dask/dask
Pull request #11477 (2024-10-31 15:02:25)

Remove recursion from TaskSpec
Recursing into all data structures in the graph is the devil I wanted to get rid of. However, especially with merging https://github.com/dask/dask/pull/11377 I did in fact introduce a lot of recursion and this is very heavily impacting runtime performance. I decided to remove recursion now completely.

This means that Tasks will only be recognized top-level. I think this is a fine compromise. There are changes in distributed and dask-expr necessary, unfortunately.

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 36m 25s ⏱️ - 8m 8s
 13 134 tests  -   121   12 070 ✅  - 123   1 044 💤  -  18   18 ❌ + 18   2 🔥 + 2 
137 162 runs   - 1 136  118 608 ✅  - 698  18 328 💤  - 664  146 ❌ +146  80 🔥 +80 

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

Results for commit 925b40ef. ± Comparison against base commit 0cb1cb1f.

<details>
  <summary>This pull request <b>removes</b> 121 tests.</summary>

```
dask.tests.test_distributed ‑ test_annotations_blockwise_unpack
dask.tests.test_distributed ‑ test_await
dask.tests.test_distributed ‑ test_bag_groupby_default
dask.tests.test_distributed ‑ test_blockwise_array_creation[False-full]
dask.tests.test_distributed ‑ test_blockwise_array_creation[False-ones]
dask.tests.test_distributed ‑ test_blockwise_array_creation[False-zeros]
dask.tests.test_distributed ‑ test_blockwise_array_creation[None-full]
dask.tests.test_distributed ‑ test_blockwise_array_creation[None-ones]
dask.tests.test_distributed ‑ test_blockwise_array_creation[None-zeros]
dask.tests.test_distributed ‑ test_blockwise_array_creation[True-full]
…
```
</details>

[test-results]:data:application/gzip;base64,H4sIAPSkI2cC/0WOSw4CIQxAr2JYuyjl72UMMpAQ52MYWBnvLjjCLN9r89o3CXH2O7ldKLteyF5iPmkqyea4rY3RaFFVneZjThnvfN+Lc78lUHDKZ3w1CfxcDDbOzelhfEpbqgqrSWU92orKzqNNtQQ95L+tGQ7X21x209u6PeW2ZYm5EjEoHhx8MExYIad6zCBgQGWk8c4GUMqBQkU+X4+zbFoeAQAA

phofl: I like the wrapper containers, they should make it easier to migrate things and are generally easier to work with. 

The parse_input function seems fine as well, this should be good from my side

Patch:
diff --git a/dask/_task_spec.py b/dask/_task_spec.py
index d53298d7b..64b700b3c 100644
--- a/dask/_task_spec.py
+++ b/dask/_task_spec.py
@@ -40,26 +40,6 @@ Referencing other tasks is possible by using either one of `Alias` or a
     # If a task is still in scope, the method `ref` can be used for convenience
     t2 = Task("key2", func2, t.ref())
 
-Nested tasks and inlining
--------------------------
-
-Tasks can be nested in a dict or list. This is useful for expressing more complex tasks.
-
-.. code-block:: python
-
-    func(func2("a", "b"), "c") ~ Task("key", func, [Task("key-1", func2, "a", "b"), "c"])
-
-    {"a": func("b")} ~ {"a": Task("a", func, "b")}
-
-This kind of nesting is often a byproduct of an optimizer step that inlines the
-tasks explicitly. This is done by calling the `inline` method on a task.
-
-.. code-block:: python
-
-    t1 = Task("key-1", func, "a")
-    t2 = Task("key-2", func, TaskRef("key-1"))
-
-    t2.inline({"key-1": t1}) ~ Task("key-2", func, Task("key-1", func, "a"))
 
 Executing a task
 ----------------
@@ -79,7 +59,7 @@ import sys
 from collections import defaultdict
 from collections.abc import Callable, Container, Iterable, Mapping, MutableMapping
 from functools import partial
-from typing import Any, TypeVar, cast, overload
+from typing import Any, TypeVar, cast
 
 from dask.sizeof import sizeof
 from dask.typing import Key as KeyType
@@ -101,18 +81,6 @@ def _identity_cast(*args, typ):
 _anom_count = itertools.count()
 
 
-@overload
-def parse_input(obj: _T_Iterable) -> _T_Iterable: ...
-
-
-@overload
-def parse_input(obj: TaskRef) -> Alias: ...
-
-
-@overload
-def parse_input(obj: _T_GraphNode) -> _T_GraphNode: ...
-
-
 def parse_input(obj: Any) -> object:
     """Tokenize user input into GraphNode objects
 
@@ -133,12 +101,16 @@ def parse_input(obj: Any) -> object:
     """
     if isinstance(obj, GraphNode):
         return obj
-    elif isinstance(obj, dict):
-        return {k: parse_input(v) for k, v in obj.items()}
-    elif isinstance(obj, (list, tuple, set, frozenset)):
+    if isinstance(obj, list):
+        return List(*(parse_input(o) for o in obj))
+    if isinstance(obj, set):
+        return Set(*(parse_input(o) for o in obj))
+    if isinstance(obj, tuple):
         if is_namedtuple_instance(obj):
             return _wrap_namedtuple_task(None, obj, parse_input)
-        return type(obj)(map(parse_input, obj))
+        return Tuple(*(parse_input(o) for o in obj))
+    if isinstance(obj, dict):
+        return Dict(**{k: parse_input(v) for k, v in obj.items()})  # type: ignore
 
     if isinstance(obj, TaskRef):
         return Alias(obj.key)
@@ -154,14 +126,14 @@ def _wrap_namedtuple_task(k, obj, parser):
         new_args = obj.__getnewargs__()
         kwargs = {}
 
-    args_converted = type(new_args)(map(parser, new_args))
+    args_converted = parse_input(type(new_args)(map(parser, new_args)))
 
     return Task(
-        k, partial(_instantiate_named_tuple, type(obj)), *args_converted, **kwargs
+        k, partial(_instantiate_named_tuple, type(obj)), args_converted, Dict(kwargs)
     )
 
 
-def _instantiate_named_tuple(typ, *args, **kwargs):
+def _instantiate_named_tuple(typ, args, kwargs):
     return typ(*args, **kwargs)
 
 
@@ -186,6 +158,8 @@ def _to_dict(*args: Any) -> dict:
 
 
 def _wrap_dict_in_task(d: dict) -> Task:
+    """Wrap a dictionary in a task such that all values will be properly
+    recognized as runnable tasks"""
     return Task(None, _to_dict, *d.values(), d.keys())
 
 
@@ -204,11 +178,12 @@ def _execute_subgraph(inner_dsk, outkey, inkeys, external_deps):
 
 
 def convert_legacy_task(
-    k: KeyType | None,
+    key: KeyType | None,
     task: _T,
     all_keys: Container,
 ) -> GraphNode | _T:
     global SubgraphType
+
     if SubgraphType is None:
         from dask.optimization import SubgraphCallable
 
@@ -218,9 +193,6 @@ def convert_legacy_task(
         return task
 
     if type(task) is tuple and task and callable(task[0]):
-        if k is None:
-            k = ("anom", next(_anom_count))
-
         func, args = task[0], task[1:]
         if isinstance(func, SubgraphType):
             subgraph = func
@@ -228,36 +200,64 @@ def convert_legacy_task(
                 subgraph.dsk, set(subgraph.inkeys), all_keys
             )
             sub_dsk = subgraph.dsk
-            converted_graph = convert_legacy_graph(sub_dsk, all_keys_inner)
-            external_deps = (
-                _get_dependencies(converted_graph)
-                - set(converted_graph)
-                - set(subgraph.inkeys)
-            )
-
-            converted_sub_dsk = DataNode(None, converted_graph)
-
+            deps: set[KeyType] = set()
+            converted_subgraph = convert_legacy_graph(sub_dsk, all_keys_inner)
+            for v in converted_subgraph.values():
+                if isinstance(v, GraphNode):
+                    deps.update(v.dependencies)
+
+            # There is an explicit and implicit way to provide dependencies /
+            # data to a SubgraphCallable. Since we're not recursing into any
+            # containers any more we'll have to provide those arguments in a
+            # more explicit way as a separate argument.
+
+            # The explicit way are arguments to the subgraph callable. Those can
+            # again be tasks.
+
+            explicit_inkeys = dict()
+            for k, target in zip(subgraph.inkeys, args):
+                explicit_inkeys[k] = t = convert_legacy_task(None, target, all_keys)
+                if isinstance(t, GraphNode):
+                    deps.update(t.dependencies)
+            explicit_inkeys_wrapped = Dict(explicit_inkeys)
+            deps.update(explicit_inkeys_wrapped.dependencies)
+
+            # The implicit way is when the tasks inside of the subgraph are
+            # referencing keys that are not part of the subgraph but are part of
+            # the outer graph.
+
+            deps -= set(subgraph.dsk)
+            deps -= set(subgraph.inkeys)
+
+            implicit_inkeys = dict()
+            for k in deps - explicit_inkeys_wrapped.dependencies:
+                assert k is not None
+                implicit_inkeys[k] = Alias(k)
             return Task(
-                k,
+                key,
                 _execute_subgraph,
-                converted_sub_dsk,
+                converted_subgraph,
                 func.outkey,
-                {
-                    k: convert_legacy_task(None, target, all_keys)
-                    for k, target in zip(subgraph.inkeys, args)
-                },
-                {ext_dep: Alias(ext_dep) for ext_dep in external_deps},
+                explicit_inkeys_wrapped,
+                Dict(implicit_inkeys),
             )
         else:
-            new_args = tuple(convert_legacy_task(None, a, all_keys) for a in args)
-            return Task(k, func, *new_args)
+            new_args = []
+            new: object
+            for a in args:
+                if isinstance(a, dict):
+                    new = Dict(a)
+                else:
+                    new = convert_legacy_task(None, a, all_keys)
+                new_args.append(new)
+            return Task(key, func, *new_args)
     try:
         if isinstance(task, (bytes, int, float, str, tuple)):
             if task in all_keys:
-                if k is None:
+                if key is None:
                     return Alias(task)
                 else:
-                    return Alias(k, target=task)
+                    return Alias(key, target=task)
     except TypeError:
         # Unhashable
         pass
@@ -265,40 +265,41 @@ def convert_legacy_task(
     if isinstance(task, (list, tuple, set, frozenset)):
         if is_namedtuple_instance(task):
             return _wrap_namedtuple_task(
-                k,
+                key,
                 task,
-                partial(convert_legacy_task, None, all_keys=all_keys),
+                partial(
+                    convert_legacy_task,
+                    None,
+                    all_keys=all_keys,
+                ),
             )
         else:
             parsed_args = tuple(convert_legacy_task(None, t, all_keys) for t in task)
             if any(isinstance(a, GraphNode) for a in parsed_args):
                 return Task(
-                    k, _identity_cast, *parsed_args, typ=DataNode(None, type(task))
+                    key, _identity_cast, *parsed_args, typ=DataNode(None, type(task))
                 )
             else:
                 return cast(_T, type(task)(parsed_args))
-    elif isinstance(task, dict):
-        return _wrap_dict_in_task(
-            {k: convert_legacy_task(k, v, all_keys) for k, v in task.items()}
-        )
-
     elif isinstance(task, TaskRef):
-        if k is None:
+        if key is None:
             return Alias(task.key)
         else:
-            return Alias(k, target=task.key)
+            return Alias(key, target=task.key)
     else:
         return task
 
 
-def convert_legacy_graph(dsk: Mapping, all_keys: Container | None = None):
-    if not all_keys:
+def convert_legacy_graph(
+    dsk: Mapping,
+    all_keys: Container | None = None,
+):
+    if all_keys is None:
         all_keys = set(dsk)
     new_dsk = {}
     for k, arg in dsk.items():
         t = convert_legacy_task(k, arg, all_keys)
-        if isinstance(t, Alias) and isinstance(arg, TaskRef) and t.key == arg.key:
-            # This detects cycles?
+        if isinstance(t, Alias) and t.target.key == k:
             continue
         new_dsk[k] = t
     new_dsk2 = {
@@ -380,6 +381,11 @@ class TaskRef:
     def __hash__(self) -> int:
         return hash(self.key)
 
+    def __eq__(self, value: object) -> bool:
+        if not isinstance(value, TaskRef):
+            return False
+        return self.key == value.key
+
 
 class GraphNode:
     key: KeyType
@@ -403,9 +409,6 @@ class GraphNode:
         if missing := set(self.dependencies) - set(values):
             raise RuntimeError(f"Not enough arguments provided: missing keys {missing}")
 
-    def inline(self, dsk) -> GraphNode:
-        raise NotImplementedError("Not implemented")
-
     def __call__(self, values) -> Any:
         raise NotImplementedError("Not implemented")
 
@@ -443,7 +446,7 @@ class Alias(GraphNode):
         if not isinstance(target, TaskRef):
             target = TaskRef(target)
         self.target = target
-        self._dependencies = frozenset([target.key])
+        self._dependencies = frozenset((target.key,))
 
     def copy(self):
         return Alias(self.key, self.target)
@@ -452,18 +455,6 @@ class Alias(GraphNode):
         self._verify_values(values)
         return values[self.target.key]
 
-    def inline(self, dsk) -> GraphNode:
-        if self.key in dsk:
-            # This can otherwise cause recursion errors
-            new_dsk = dsk.copy()
-            new_dsk.pop(self.key)
-            if isinstance(dsk[self.key], GraphNode):
-                return dsk[self.key].inline(new_dsk)
-            else:
-                return dsk[self.key]
-        else:
-            return self
-
     def __repr__(self):
         return f"Alias({self.key}->{self.target})"
 
@@ -490,9 +481,6 @@ class DataNode(GraphNode):
         self.typ = type(value)
         self._dependencies = _no_deps
 
-    def inline(self, dsk) -> DataNode:
-        return self
-
     def copy(self):
         return DataNode(self.key, self.value)
 
@@ -514,20 +502,6 @@ class DataNode(GraphNode):
         return iter(self.value)
 
 
-def _inline_recursively(obj, values, seen=None):
-    seen = seen or set()
-    if id(obj) in seen:
-        return obj
-    seen.add(id(obj))
-    if isinstance(obj, GraphNode):
-        return obj.inline(values)
-    elif isinstance(obj, dict):
-        return {k: _inline_recursively(v, values, seen=seen) for k, v in obj.items()}
-    elif isinstance(obj, (list, tuple, frozenset, set)):
-        return type(obj)(_inline_recursively(v, values, seen=seen) for v in obj)
-    return obj
-
-
 def _call_recursively(obj, values):
     if isinstance(obj, GraphNode):
         return obj(values)
@@ -568,17 +542,24 @@ class Task(GraphNode):
         func: Callable,
         /,
         *args: Any,
+        _dependencies: set | frozenset | None = None,
         **kwargs: Any,
     ):
         self.key = key
         self.func = func
-        self.args = parse_input(args)
-        self.kwargs = parse_input(kwargs)
-        dependencies: set = set()
-        dependencies.update(_get_dependencies(self.args))
-        dependencies.update(_get_dependencies(tuple(self.kwargs.values())))
-        if dependencies:
-            self._dependencies = frozenset(dependencies)
+        self.args = tuple(
+            Alias(obj.key) if isinstance(obj, TaskRef) else obj for obj in args
+        )
+        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)
+        if _dependencies:
+            self._dependencies = frozenset(_dependencies)
         else:
             self._dependencies = _no_deps
         self._is_coro = None
@@ -586,7 +567,13 @@ class Task(GraphNode):
         self._repr = None
 
     def copy(self):
-        return Task(self.key, self.func, *self.args, **self.kwargs)
+        return type(self)(
+            self.key,
+            self.func,
+            *self.args,
+            _dependencies=self._dependencies,
+            **self.kwargs,
+        )
 
     def __hash__(self):
         return hash(self._get_token())
@@ -644,18 +631,17 @@ class Task(GraphNode):
 
     def __call__(self, values=()):
         self._verify_values(values)
-        new_argspec = _call_recursively(self.args, values)
+        new_argspec = tuple(
+            a(values) if isinstance(a, GraphNode) else a for a in self.args
+        )
         if self.kwargs:
-            kwargs = _call_recursively(self.kwargs, values)
+            kwargs = {
+                k: kw(values) if isinstance(kw, GraphNode) else kw
+                for k, kw in self.kwargs.items()
+            }
             return self.func(*new_argspec, **kwargs)
         return self.func(*new_argspec)
 
-    def inline(self, dsk) -> Task:
-        new_args = _inline_recursively(self.args, dsk)
-        new_kwargs = _inline_recursively(self.kwargs, dsk)
-        assert self.func is not None
-        return Task(self.key, self.func, *new_args, **new_kwargs)
-
     @property
     def is_coro(self):
         if self._is_coro is None:
@@ -686,13 +672,66 @@ class Task(GraphNode):
         return Task(
             key or outkey,
             _execute_subgraph,
-            DataNode(None, {t.key: t for t in tasks}),
+            {t.key: t for t in tasks},
             outkey,
-            {k: Alias(k) for k in external_deps},
+            Dict({k: Alias(k) for k in external_deps}),
             {},
         )
 
 
+class NestedContainer(Task):
+    klass: type
+    __slots__ = tuple(__annotations__)
+
+    def __init__(
+        self,
+        /,
+        *args: Any,
+        _dependencies: set | frozenset | None = None,
+        **kwargs: Any,
+    ):
+        super().__init__(
+            None,
+            self.to_container,
+            *args,
+            _dependencies=_dependencies,
+            **kwargs,
+        )
+
+    @classmethod
+    def to_container(cls, *args, **kwargs):
+        return cls.klass(args)
+
+
+class List(NestedContainer):
+    klass = list
+
+
+class Tuple(NestedContainer):
+    klass = tuple
+
+
+class Set(NestedContainer):
+    klass = set
+
+
+class Dict(NestedContainer):
+    klass = dict
+
+    def __init__(
+        self, /, *args: Any, _dependencies: set | frozenset | None = None, **kwargs: Any
+    ):
+        if args:
+            if len(args) > 1:
+                raise ValueError("Dict can only take one positional argument")
+            kwargs = args[0]
+        elif not kwargs:
+            raise ValueError("Dict needs at least one argument")
+        super().__init__(
+            *(Tuple(*it) for it in kwargs.items()), _dependencies=_dependencies
+        )
+
+
 class DependenciesMapping(MutableMapping):
     def __init__(self, dsk):
         self.dsk = dsk


Test patch:
diff --git a/dask/tests/test_task_spec.py b/dask/tests/test_task_spec.py
index e9b99cd6b..0fae036fd 100644
--- a/dask/tests/test_task_spec.py
+++ b/dask/tests/test_task_spec.py
@@ -11,11 +11,16 @@ from dask._task_spec import (
     Alias,
     DataNode,
     DependenciesMapping,
+    Dict,
+    List,
+    Set,
     Task,
     TaskRef,
+    Tuple,
     _get_dependencies,
     convert_legacy_graph,
     execute_graph,
+    parse_input,
     resolve_aliases,
 )
 from dask.base import tokenize
@@ -24,6 +29,13 @@ from dask.optimization import SubgraphCallable
 from dask.sizeof import sizeof
 
 
+def convert_and_verify_keys(dsk):
+    new_dsk = convert_legacy_graph(dsk)
+    vals = [(k, v.key) for k, v in new_dsk.items()]
+    assert all(v[0] == v[1] for v in vals), vals
+    return new_dsk
+
+
 def identity(x):
     return x
 
@@ -47,7 +59,7 @@ def test_convert_legacy_dsk_skip_new():
     dsk = {
         "key-1": Task("key-1", func, "a", "b"),
     }
-    converted = convert_legacy_graph(dsk)
+    converted = convert_and_verify_keys(dsk)
     assert converted["key-1"] is dsk["key-1"]
     assert converted == dsk
 
@@ -102,11 +114,6 @@ def test_unpickled_repr():
     assert repr(t) == "<Task 'kwarg' use_kwargs('foo', kwarg='kwarg_value')>"
 
 
-def _assert_dsk_conversion(new_dsk):
-    vals = [(k, v.key) for k, v in new_dsk.items()]
-    assert all(v[0] == v[1] for v in vals), vals
-
-
 def test_convert_legacy_dsk():
     def func(*args):
         return "-".join(args)
@@ -126,8 +133,7 @@ def test_convert_legacy_dsk():
             (func, "const", "bar"),
         ],
     }
-    new_dsk = convert_legacy_graph(dsk)
-    _assert_dsk_conversion(new_dsk)
+    new_dsk = convert_and_verify_keys(dsk)
     t1 = new_dsk["key-1"]
     assert isinstance(t1, Task)
     assert t1.func == func
@@ -168,17 +174,19 @@ def test_task_executable():
 
 
 def test_task_nested_sequence():
+    # We do **not** recurse into structures since this kills performance
+    # we're only considering top-level task objects as arguments or kwargs
     t1 = Task("key-1", func, "a", "b")
     t2 = Task("key-2", func2, "c", "d")
 
     tseq = Task("seq", identity, [t1, t2])
-    assert tseq() == [func("a", "b"), func2("c", "d")]
+    assert tseq() == [t1, t2]
 
     tseq = Task("set", identity, {t1, t2})
-    assert tseq() == {func("a", "b"), func2("c", "d")}
+    assert tseq() == {t1, t2}
 
     tseq = Task("dict", identity, {"a": t1, "b": t2})
-    assert tseq() == {"a": func("a", "b"), "b": func2("c", "d")}
+    assert tseq() == {"a": t1, "b": t2}
 
 
 def test_reference_remote():
@@ -288,8 +296,7 @@ def test_parse_curry():
     dsk = {
         "key-1": (curry, func, "a", "b"),
     }
-    converted = convert_legacy_graph(dsk)
-    _assert_dsk_conversion(converted)
+    converted = convert_and_verify_keys(dsk)
     t = Task("key-1", curry, func, "a", "b")
     assert converted["key-1"]() == t()
     assert t() == "a-bc"
@@ -310,7 +317,7 @@ def test_avoid_cycles():
     dsk = {
         "key": TaskRef("key"),  # e.g. a persisted key
     }
-    new_dsk = convert_legacy_graph(dsk)
+    new_dsk = convert_and_verify_keys(dsk)
     assert not new_dsk
 
 
@@ -378,8 +385,7 @@ def test_subgraph_callable():
         "b": 2,
         "c": (subgraph, "a", "b"),
     }
-    new_dsk = convert_legacy_graph(dsk)
-    _assert_dsk_conversion(new_dsk)
+    new_dsk = convert_and_verify_keys(dsk)
     assert new_dsk["c"].dependencies == {"a", "b", "add-123"}
     assert (
         new_dsk["c"](
@@ -409,51 +415,14 @@ def test_redirect_to_subgraph_callable():
         "alias": "foo",
         "foo": (func, (subgraph, (func2, "bar", "baz")), "second_arg"),
     }
-    new_dsk = convert_legacy_graph(dsk)
-    from dask import get
+    new_dsk = convert_and_verify_keys(dsk)
 
-    expected = get(dsk, ["foo"])[0]
-    assert expected == "bar=baz//some=dict-second_arg"
     t = new_dsk["alias"]
     assert t.dependencies == {"foo"}
     t = new_dsk["foo"]
     assert not t.dependencies
 
-    assert t() == expected
-
-
-def test_inline():
-    t = Task("key-1", func, "a", TaskRef("b"))
-
-    assert t.inline({"b": DataNode(None, "b")})() == "a-b"
-
-    assert t.inline({"b": Task(None, func, "b", "c")})() == "a-b-c"
-
-    t2 = Task("key-1", func, TaskRef("a"), TaskRef("b"))
-    inline2 = t2.inline({"a": DataNode(None, "a")})
-    assert inline2.dependencies == {"b"}
-
-    assert inline2({"b": "foo"}) == "a-foo"
-
-    t = Task("foo", identity, [t, t2])
-    assert t.dependencies == {"b", "a"}
-    inlined = t.inline({"a": DataNode(None, "foo")})
-    assert inlined.dependencies == {"b"}
-    assert inlined({"b": "bar"}) == ["a-bar", "foo-bar"]
-
-    t = Task("key-2", identity, [t, t2])
-    t = Task("key-2", identity, {"foo": t, "bar": t2})
-    inlined = t.inline(
-        {
-            "a": DataNode(None, "foo"),
-            "bar": DataNode(None, "bar"),
-        }
-    )
-    assert inlined({"b": "bar"})
-
-    t = Task("key-1", identity, [TaskRef("a"), TaskRef("b")])
-    inlined = t.inline({"a": DataNode(None, "a"), "b": DataNode(None, "b")})
-    assert inlined() == ["a", "b"]
+    assert t() == "bar=baz//some=dict-second_arg"
 
 
 @pytest.mark.parametrize(
@@ -505,8 +474,7 @@ def test_parse_graph_namedtuple_legacy(typ, args, kwargs):
         return x
 
     dsk = {"foo": (func, typ(*args, **kwargs))}
-    new_dsk = convert_legacy_graph(dsk)
-    _assert_dsk_conversion(new_dsk)
+    new_dsk = convert_and_verify_keys(dsk)
 
     assert new_dsk["foo"]() == typ(*args, **kwargs)
 
@@ -524,7 +492,7 @@ def test_parse_namedtuple(typ, args, kwargs):
         return x
 
     obj = typ(*args, **kwargs)
-    t = Task("foo", func, obj)
+    t = Task("foo", func, parse_input(obj))
 
     assert t() == obj
 
@@ -532,7 +500,7 @@ def test_parse_namedtuple(typ, args, kwargs):
     if typ is PlainNamedTuple:
         args = tuple([TaskRef("b")] + list(args)[1:])
         obj = typ(*args, **kwargs)
-        t = Task("foo", func, obj)
+        t = Task("foo", func, parse_input(obj))
         assert t.dependencies == {"b"}
         assert t({"b": "foo"}) == typ(*tuple(["foo"] + list(args)[1:]), **kwargs)
 
@@ -630,7 +598,7 @@ def test_convert_resolve():
         "fourth": "third",
         "fifth": (func, "fourth"),
     }
-    dsk = convert_legacy_graph(dsk)
+    dsk = convert_and_verify_keys(dsk)
     assert len(dsk) == 5
 
     optimized = resolve_aliases(dsk, {"fifth"}, reverse_dict(DependenciesMapping(dsk)))
@@ -901,3 +869,45 @@ def test_subgraph_dont_hold_in_memory_too_long_legacy():
     converted = convert_legacy_graph(dsk)
     assert converted["bar"]()
     assert OnlyTwice.total == 10
+
+
+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"))
+    assert t.dependencies == {"b"}
+    assert t({"b": "b"}) == ["a-b", "c-d"]
+
+    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"}
+    t = Dict(k=Task("key-1", func, "a", TaskRef("b")), v=Task("key-2", func, "c", "d"))
+    assert t.dependencies == {"b"}
+    assert t({"b": "b"}) == {"k": "a-b", "v": "c-d"}
+
+    t = Dict(
+        {
+            "k": Task("key-1", func, "a", TaskRef("b")),
+            # Use a hashable key that isn't a string. This cannot be used as
+            # kwargs, for instance
+            ("v", 1): Task("key-2", func, "c", "d"),
+        }
+    )
+    assert t.dependencies == {"b"}
+    assert t({"b": "b"}) == {"k": "a-b", ("v", 1): "c-d"}
+
+
+@pytest.mark.parametrize(
+    "task_type, python_type",
+    [
+        (Tuple, tuple),
+        (List, list),
+        (Set, set),
+    ],
+)
+def test_nested_containers_different_types(task_type, python_type):
+    t = task_type(Task("key-1", func, "a", TaskRef("b")), Task("key-2", func, "c", "d"))
+    assert t.dependencies == {"b"}
+    assert t({"b": "b"}) == python_type(("a-b", "c-d"))