Repository: dask/dask
Pull request #11438 (2024-10-17 13:40:42)

TaskSpec: Ensure dependencies are not mutated
Fun fact: `deps -= self._removed` is the only operation on sets that is mutating the set but is still working on a frozenset. The frozenset is falling back to an implementation that is creating a new set. Note: This wasn't intended to be inplace

phofl: this is interesting...

Patch:
diff --git a/dask/_task_spec.py b/dask/_task_spec.py
index 77c388735..3b563aced 100644
--- a/dask/_task_spec.py
+++ b/dask/_task_spec.py
@@ -407,7 +407,7 @@ _func_cache_reverse: MutableMapping = LRU(maxsize=1000)
 
 class GraphNode:
     key: KeyType
-    dependencies: set | frozenset
+    _dependencies: frozenset
 
     __slots__ = tuple(__annotations__)
 
@@ -417,6 +417,10 @@ class GraphNode:
     def copy(self):
         raise NotImplementedError
 
+    @property
+    def dependencies(self) -> frozenset:
+        return self._dependencies
+
     def _verify_values(self, values: tuple | dict) -> None:
         if not self.dependencies:
             return
@@ -442,7 +446,6 @@ _no_deps: frozenset = frozenset()
 
 class Alias(GraphNode):
     __weakref__: Any = None
-    _dependencies: set | None
     target: TaskRef
     __slots__ = tuple(__annotations__)
 
@@ -455,13 +458,7 @@ class Alias(GraphNode):
         if not isinstance(target, TaskRef):
             target = TaskRef(target)
         self.target = target
-        self._dependencies = None
-
-    @property
-    def dependencies(self):
-        if self._dependencies is None:
-            self._dependencies = {self.target.key}
-        return self._dependencies
+        self._dependencies = frozenset([target.key])
 
     def copy(self):
         return Alias(self.key, self.target)
@@ -509,7 +506,7 @@ class DataNode(GraphNode):
         self.key = key
         self.value = value
         self.typ = type(value)
-        self.dependencies = _no_deps
+        self._dependencies = _no_deps
 
     def inline(self, dsk) -> DataNode:
         return self
@@ -610,9 +607,9 @@ class Task(GraphNode):
         dependencies.update(_get_dependencies(self.args))
         dependencies.update(_get_dependencies(tuple(self.kwargs.values())))
         if dependencies:
-            self.dependencies = dependencies
+            self._dependencies = frozenset(dependencies)
         else:
-            self.dependencies = _no_deps
+            self._dependencies = _no_deps
         self._is_coro = None
         self._token = None
 
@@ -771,7 +768,7 @@ class Task(GraphNode):
     def __setstate__(self, state):
         self.key = state["key"]
         self.packed_func = state["packed_func"]
-        self.dependencies = state["dependencies"]
+        self._dependencies = state["dependencies"]
         self.kwargs = state["kwargs"]
         self.args = state["args"]
         self._is_coro = state["_is_coro"]
@@ -823,7 +820,9 @@ class DependenciesMapping(MutableMapping):
         else:
             deps = self.dsk[key].dependencies
         if self._removed:
-            deps -= self._removed
+            # deps is a frozenset but for good measure, let's not use -= since
+            # that _may_ perform an inplace mutation
+            deps = deps - self._removed
         return deps
 
     def __iter__(self):


Test patch:
diff --git a/dask/tests/test_task_spec.py b/dask/tests/test_task_spec.py
index da361701e..f9c993dec 100644
--- a/dask/tests/test_task_spec.py
+++ b/dask/tests/test_task_spec.py
@@ -762,3 +762,18 @@ def test_execute_tasks_in_graph():
 def test_deterministic_tokenization_respected():
     with pytest.raises(RuntimeError, match="deterministic"):
         tokenize(Task("key", func, object()), ensure_deterministic=True)
+
+
+def test_dependencies_mapping_doesnt_mutate_task():
+
+    t = Task("key", func, "a", "b")
+    t2 = Task("key2", func, "a", t.ref())
+
+    assert t2.dependencies == {"key"}
+
+    dsk = {t.key: t, t2.key: t2}
+    deps = DependenciesMapping(dsk)
+    del deps[t.key]
+    # The getitem is doing weird stuff and could mutate the task state
+    deps[t2.key]
+    assert t2.dependencies == {"key"}