Repository: dask/dask
Pull request #11867 (2025-04-03 15:05:33)

Ensure dictionaries tokenize deterministically
In https://github.com/dask/dask/pull/11373 I changed this sorting key to a hash to pick up a little speed. However, this broke determinism and at least for trivial data types like these, determinism would be nice. I don't think the additional speed from saved str comparison is truly worth is

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

      9 files  ±0        9 suites  ±0   3h 29m 5s ⏱️ - 4m 14s
 18 002 tests ±0   16 788 ✅ ±0   1 214 💤 ±0  0 ❌ ±0 
161 092 runs  ±0  148 988 ✅ ±0  12 104 💤 ±0  0 ❌ ±0 

Results for commit 31b8dc87. ± Comparison against base commit 21b648b8.

[test-results]:data:application/gzip;base64,H4sIAGGt7mcC/02Myw6DIBBFf8Ww7mIGlEd/ptFBElIfDcKq6b8XqVKX59zc82bOT+PG7o25NWxLPlawKfTRr0tG5F3bZZXHuM+oAfjJjy0R7VIqrf/y6V/liW11rvdTdlDFGMIaDhPSUtISwfCDa7rVprR/sqahPd0lXfhapnWefczABA7aklaDkdAbGC2SkxyFEo6EQSsIQCpJ7PMFS4PlghgBAAA=


Patch:
diff --git a/dask/local.py b/dask/local.py
index d701615da..66b72f5ea 100644
--- a/dask/local.py
+++ b/dask/local.py
@@ -110,15 +110,16 @@ See the function ``inline_functions`` for more information.
 from __future__ import annotations
 
 import os
+from collections import defaultdict
 from collections.abc import Mapping, Sequence
 from concurrent.futures import Executor, Future
 from functools import partial
 from queue import Empty, Queue
 
 from dask import config
-from dask._task_spec import DataNode, DependenciesMapping, convert_legacy_graph
+from dask._task_spec import DataNode, convert_legacy_graph
 from dask.callbacks import local_callbacks, unpack_callbacks
-from dask.core import flatten, get_dependencies, reverse_dict
+from dask.core import flatten, get_dependencies
 from dask.order import order
 from dask.typing import Key
 
@@ -140,7 +141,7 @@ else:
         return q.get()
 
 
-def start_state_from_dask(dsk, cache=None, sortkey=None):
+def start_state_from_dask(dsk, cache=None, sortkey=None, keys=None):
     """Start state from a dask
 
     Examples
@@ -166,35 +167,57 @@ def start_state_from_dask(dsk, cache=None, sortkey=None):
         cache = config.get("cache", None)
     if cache is None:
         cache = dict()
-
+    if keys is None:
+        keys = list(set(dsk) - set(cache))
     dsk = convert_legacy_graph(dsk, all_keys=set(dsk) | set(cache))
-    data_keys = set()
-    for k, v in dsk.items():
-        if isinstance(v, DataNode):
-            cache[k] = v()
-            data_keys.add(k)
-
-    dsk2 = dsk.copy()
-    dsk2.update(cache)
-
-    dependencies = DependenciesMapping(dsk)
-    waiting = {k: set(v) for k, v in dependencies.items() if k not in data_keys}
-
-    dependents = reverse_dict(dependencies)
-    for a in cache:
-        for b in dependents.get(a, ()):
-            waiting[b].remove(a)
-    waiting_data = {k: v.copy() for k, v in dependents.items() if v}
+    stack = list(keys)
+    dependencies = defaultdict(set)
+    dependents = defaultdict(set)
+    waiting = defaultdict(set)
+    waiting_data = defaultdict(set)
+    ready_set = set()
+    seen = set()
+    while stack:
+        key = stack.pop()
+        if key in seen:
+            continue
+        seen.add(key)
+        dependents[key]
+        waiting_data[key]
+        dependencies[key]
+        task = dsk.get(key, None)
+        if task is None:
+            continue
+        elif isinstance(task, DataNode):
+            cache[key] = task()
+            dependencies[key]
+            for d in dependents[key]:
+                if d in waiting:
+                    waiting[d].remove(key)
+                    if not waiting[d]:
+                        del waiting[d]
+                        ready_set.add(d)
+                else:
+                    ready_set.add(d)
+        else:
+            _wait = task.dependencies - set(cache)
+            if not _wait:
+                ready_set.add(key)
+            else:
+                waiting[key] = set(_wait)
+            for dep in task.dependencies:
+                dependencies[key].add(dep)
+                dependents[dep].add(key)
+                waiting_data[dep].add(key)
+                stack.append(dep)
 
-    ready_set = {k for k, v in waiting.items() if not v}
     ready = sorted(ready_set, key=sortkey, reverse=True)
-    waiting = {k: v for k, v in waiting.items() if v}
 
     state = {
-        "dependencies": dependencies,
-        "dependents": dependents,
-        "waiting": waiting,
-        "waiting_data": waiting_data,
+        "dependencies": dict(dependencies),
+        "dependents": dict(dependents),
+        "waiting": dict(waiting),
+        "waiting_data": dict(waiting_data),
         "cache": cache,
         "ready": ready,
         "running": set(),
@@ -444,7 +467,9 @@ def get_async(
 
             keyorder = order(dsk)
 
-            state = start_state_from_dask(dsk, cache=cache, sortkey=keyorder.get)
+            state = start_state_from_dask(
+                dsk, keys=results, cache=cache, sortkey=keyorder.get
+            )
 
             for _, start_state, _, _, _ in callbacks:
                 if start_state:
@@ -480,7 +505,7 @@ def get_async(
 
                     # Prep args to send
                     data = {
-                        dep: state["cache"][dep] for dep in get_dependencies(dsk, key)
+                        dep: state["cache"][dep] for dep in state["dependencies"][key]
                     }
                     args.append(
                         (
diff --git a/dask/tokenize.py b/dask/tokenize.py
index fb3b81aa1..ab07f363a 100644
--- a/dask/tokenize.py
+++ b/dask/tokenize.py
@@ -120,7 +120,7 @@ def normalize_dict(d):
         _SEEN[id(d)] = len(_SEEN), d
         try:
             return "dict", _normalize_seq_func(
-                sorted(d.items(), key=lambda kv: hash(kv[0]))
+                sorted(d.items(), key=lambda kv: str(kv[0]))
             )
         finally:
             _SEEN.pop(id(d), None)


Test patch:
diff --git a/dask/tests/test_local.py b/dask/tests/test_local.py
index a1bf96d4f..684e43949 100644
--- a/dask/tests/test_local.py
+++ b/dask/tests/test_local.py
@@ -3,6 +3,7 @@ from __future__ import annotations
 import pytest
 
 import dask
+from dask._task_spec import Task
 from dask.local import finish_task, get_sync, sortkey, start_state_from_dask
 from dask.order import order
 from dask.utils_test import GetFunctionTestMixin, add, inc
@@ -28,7 +29,7 @@ def test_start_state():
         "running": set(),
         "ready": ["z"],
         "waiting": {"w": {"z"}},
-        "waiting_data": {"x": {"z"}, "y": {"w"}, "z": {"w"}},
+        "waiting_data": {"w": set(), "x": {"z"}, "y": {"w"}, "z": {"w"}},
     }
     assert result == expected
 
@@ -87,7 +88,7 @@ def test_finish_task():
         "dependents": {"w": set(), "x": {"z"}, "y": {"w"}, "z": {"w"}},
         "ready": ["w"],
         "waiting": {},
-        "waiting_data": {"y": {"w"}, "z": {"w"}},
+        "waiting_data": {"w": set(), "y": {"w"}, "z": {"w"}},
     }
 
 
@@ -190,3 +191,25 @@ def test_complex_ordering():
     with Callback(pretask=track_order):
         get_sync(dsk, exp_order[-1])
     assert actual_order == exp_order
+
+
+def test_ensure_calculate_only_whats_needed():
+    counter = 0
+
+    def only_once():
+        nonlocal counter
+        if counter > 0:
+            raise RuntimeError("Should only be called once")
+        counter += 1
+        return 1
+
+    def add(x, y):
+        return x + y
+
+    tasks = [
+        Task("x1", only_once),
+        (t1 := Task("x2", only_once)),
+        Task("y", add, 1, t1.ref()),
+    ]
+    dsk = {t.key: t for t in tasks}
+    assert get_sync(dsk, "y") == 2
diff --git a/dask/tests/test_tokenize.py b/dask/tests/test_tokenize.py
index 51c6750fa..08da401c5 100644
--- a/dask/tests/test_tokenize.py
+++ b/dask/tests/test_tokenize.py
@@ -848,6 +848,7 @@ def test_nested_tokenize_seen():
 
 def test_tokenize_dict():
     # Insertion order is ignored. Keys can be an unsortable mix of types.
+    assert check_tokenize({"x": 1, 1: "x"}) == "ba2498a4c5583cc3ec540865185cdafd"
     assert check_tokenize({"x": 1, 1: "x"}) == check_tokenize({1: "x", "x": 1})
     assert check_tokenize({"x": 1, 1: "x"}) != check_tokenize({"x": 1, 2: "x"})
     assert check_tokenize({"x": 1, 1: "x"}) != check_tokenize({"x": 2, 1: "x"})