Repository: dask/dask
Pull request #7194 (2021-02-09 10:33:22)

Fix regression in Delayed._length
- Fixes regression introduced by #7177
- Closes dask/distributed#4493
- Supersedes #7192

FYI @jakirkham 

jakirkham: Would it make sense to have a test covering the Distributed use case or is that already covered here?
jrbourbeau: Since `persist` here and `client.persist` in `distributed` use the `__dask_postpersist__` and `__dask_keys__` interface, which is where the regression in https://github.com/dask/dask/pull/7177 was introduced, I think the tests added here which test those interfaces for `DelayedLeaf` and `DelayedAttr` should encompass the `distributed` case too

Patch:
diff --git a/dask/delayed.py b/dask/delayed.py
index 9ee580e16..342f027fe 100644
--- a/dask/delayed.py
+++ b/dask/delayed.py
@@ -473,13 +473,21 @@ class Delayed(DaskMethodsMixin, OperatorMethodMixin):
     Equivalent to the output from a single key in a dask graph.
     """
 
-    __slots__ = ("_key", "dask", "_length")
+    __slots__ = ("_key", "_dask", "_length")
 
     def __init__(self, key, dsk, length=None):
         self._key = key
-        self.dask = dsk
+        self._dask = dsk
         self._length = length
 
+    @property
+    def key(self):
+        return self._key
+
+    @property
+    def dask(self):
+        return self._dask
+
     def __dask_graph__(self):
         return self.dask
 
@@ -510,17 +518,6 @@ class Delayed(DaskMethodsMixin, OperatorMethodMixin):
         key = replace_name_in_key(self.key, name) if name else self.key
         return Delayed(key, dsk, self._length)
 
-    def __getstate__(self):
-        return tuple(getattr(self, i) for i in self.__slots__)
-
-    def __setstate__(self, state):
-        for k, v in zip(self.__slots__, state):
-            setattr(self, k, v)
-
-    @property
-    def key(self):
-        return self._key
-
     def __repr__(self):
         return "Delayed({0})".format(repr(self.key))
 
@@ -532,21 +529,27 @@ class Delayed(DaskMethodsMixin, OperatorMethodMixin):
 
     def __getattr__(self, attr):
         if attr.startswith("_"):
-            raise AttributeError("Attribute {0} not found".format(attr))
+            raise AttributeError(f"Attribute {attr} not found")
 
         if attr == "visualise":
             # added to warn users in case of spelling error
             # for more details: https://github.com/dask/dask/issues/5721
             warnings.warn(
-                "dask.delayed objects have no `visualise` method. Perhaps you meant `visualize`?"
+                "dask.delayed objects have no `visualise` method. "
+                "Perhaps you meant `visualize`?"
             )
 
         return DelayedAttr(self, attr)
 
     def __setattr__(self, attr, val):
-        if attr in self.__slots__:
+        try:
             object.__setattr__(self, attr, val)
-        else:
+        except AttributeError:
+            # attr is neither in type(self).__slots__ nor in the __slots__ of any of its
+            # parent classes, and all the parent classes define __slots__ too.
+            # This last bit needs to be unit tested: if any of the parent classes omit
+            # the __slots__ declaration, self will gain a __dict__ and this branch will
+            # become unreachable.
             raise TypeError("Delayed objects are immutable")
 
     def __setitem__(self, index, val):
@@ -619,11 +622,11 @@ def call_function(func, func_token, args, kwargs, pure=None, nout=None):
 
 
 class DelayedLeaf(Delayed):
-    __slots__ = ("_obj", "_key", "_pure", "_nout")
+    __slots__ = ("_obj", "_pure", "_nout")
 
     def __init__(self, obj, key, pure=None, nout=None):
+        super().__init__(key, None)
         self._obj = obj
-        self._key = key
         self._pure = pure
         self._nout = nout
 
@@ -640,12 +643,13 @@ class DelayedLeaf(Delayed):
 
 
 class DelayedAttr(Delayed):
-    __slots__ = ("_obj", "_attr", "_key")
+    __slots__ = ("_obj", "_attr")
 
     def __init__(self, obj, attr):
+        key = "getattr-%s" % tokenize(obj, attr, pure=True)
+        super().__init__(key, None)
         self._obj = obj
         self._attr = attr
-        self._key = "getattr-%s" % tokenize(obj, attr, pure=True)
 
     def __getattr__(self, attr):
         # Calling np.dtype(dask.delayed(...)) used to result in a segfault, as
@@ -654,7 +658,7 @@ class DelayedAttr(Delayed):
         # `x.dtype().dtype()` is called (which shouldn't ever show up in real
         # code). See https://github.com/dask/dask/pull/4374#issuecomment-454381465
         if attr == "dtype" and self._attr == "dtype":
-            raise AttributeError("Attribute %s not found" % attr)
+            raise AttributeError("Attribute dtype not found")
         return super().__getattr__(attr)
 
     @property
diff --git a/dask/utils.py b/dask/utils.py
index 56286575b..54dd7e6ef 100644
--- a/dask/utils.py
+++ b/dask/utils.py
@@ -1056,6 +1056,8 @@ def ensure_dict(d: Mapping[K, V], *, copy: bool = False) -> Dict[K, V]:
 class OperatorMethodMixin(object):
     """A mixin for dynamically implementing operators"""
 
+    __slots__ = ()
+
     @classmethod
     def _bind_operator(cls, op):
         """ bind operator to this class """


Test patch:
diff --git a/dask/tests/test_base.py b/dask/tests/test_base.py
index 1d349481a..dacedca80 100644
--- a/dask/tests/test_base.py
+++ b/dask/tests/test_base.py
@@ -1148,6 +1148,23 @@ def test_persist_delayed_change_name(key, new_name, new_key):
     assert dict(dp.dask) == {new_key: 2}
 
 
+def test_persist_delayedleaf():
+    x = delayed(1)
+    (xx,) = persist(x)
+    assert isinstance(xx, Delayed)
+    assert xx.compute() == 1
+
+
+def test_persist_delayedattr():
+    class C:
+        x = 1
+
+    x = delayed(C).x
+    (xx,) = persist(x)
+    assert isinstance(xx, Delayed)
+    assert xx.compute() == 1
+
+
 @pytest.mark.skipif("not da or not db")
 def test_persist_array_bag():
     x = da.arange(5, chunks=2) + 1