Avoid constructing NestedContainers in case of trivial inputs (#11600)

dask/_task_spec.py

@@ -84,7 +84,7 @@ def _identity_cast(*args, typ):
 def parse_input(obj: Any) -> object:
     """Tokenize user input into GraphNode objects
 
-    Note: This is similar to `convert_old_style_task` but does not
+    Note: This is similar to `convert_legacy_task` but does not
     - compare any values to a global set of known keys to infer references/futures
     - parse tuples and interprets them as runnable tasks
     - Deal with SubgraphCallables
@@ -101,20 +101,27 @@ def parse_input(obj: Any) -> object:
     """
     if isinstance(obj, GraphNode):
         return obj
-    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 Tuple(*(parse_input(o) for o in obj))
-    if isinstance(obj, dict):
-        return Dict({k: parse_input(v) for k, v in obj.items()})
 
     if isinstance(obj, TaskRef):
         return Alias(obj.key)
 
+    if isinstance(obj, dict):
+        parsed_dict = {k: parse_input(v) for k, v in obj.items()}
+        if any(isinstance(v, GraphNode) for v in parsed_dict.values()):
+            return Dict(parsed_dict)
+
+    if isinstance(obj, (list, set, tuple)):
+        parsed_collection = tuple(parse_input(o) for o in obj)
+        if any(isinstance(o, GraphNode) for o in parsed_collection):
+            if isinstance(obj, list):
+                return List(*parsed_collection)
+            if isinstance(obj, set):
+                return Set(*parsed_collection)
+            if isinstance(obj, tuple):
+                if is_namedtuple_instance(obj):
+                    return _wrap_namedtuple_task(None, obj, parse_input)
+                return Tuple(*parsed_collection)
+
     return obj
 
 

dask/tests/test_task_spec.py

@@ -538,6 +538,11 @@ def test_pickle_literals():
     assert roundtripped == obj
 
 
+@pytest.mark.parametrize("obj", [set, {0}, [], [1], {}, {2: 3}, (), (4,)])
+def test_parse_non_task_inputs(obj):
+    assert parse_input(obj) == obj
+
+
 def test_resolve_aliases():
     tasks = [
         Alias("bar", "foo"),