fix annotations

Author: fjetter

distributed/client.py

@@ -39,11 +39,15 @@
 
 from packaging.version import parse as parse_version
 from tlz import first, groupby, merge, partition_all, valmap
+from tornado import gen
+from tornado.ioloop import IOLoop
 
 import dask
+from dask._expr import Expr, HLGExpr, LLGExpr
+from dask._task_spec import DataNode, GraphNode, List, Task, TaskRef, parse_input
 from dask.base import collections_to_dsk
 from dask.core import flatten, validate_key
-from dask.layers import Layer
+from dask.highlevelgraph import HighLevelGraph
 from dask.tokenize import tokenize
 from dask.typing import Key, NestedKeys, NoDefault, no_default
 from dask.utils import (
@@ -57,19 +61,6 @@
 )
 from dask.widgets import get_template
 
-from distributed.core import OKMessage
-from distributed.protocol.serialize import _is_dumpable
-from distributed.utils import Deadline, wait_for
-
-try:
-    from dask.delayed import single_key
-except ImportError:
-    single_key = first
-from tornado import gen
-from tornado.ioloop import IOLoop
-
-from dask._task_spec import DataNode, GraphNode, List, Task, TaskRef, parse_input
-
 import distributed.utils
 from distributed import cluster_dump, preloading
 from distributed import versions as version_module
@@ -79,6 +70,7 @@
 from distributed.core import (
     CommClosedError,
     ConnectionPool,
+    OKMessage,
     PooledRPCCall,
     Status,
     clean_exception,
@@ -98,6 +90,7 @@
 from distributed.objects import HasWhat, SchedulerInfo, WhoHas
 from distributed.protocol import to_serialize
 from distributed.protocol.pickle import dumps, loads
+from distributed.protocol.serialize import _is_dumpable
 from distributed.publish import Datasets
 from distributed.pubsub import PubSubClientExtension
 from distributed.security import Security
@@ -106,6 +99,7 @@
 from distributed.threadpoolexecutor import rejoin
 from distributed.utils import (
     CancelledError,
+    Deadline,
     LoopRunner,
     NoOpAwaitable,
     SyncMethodMixin,
@@ -117,6 +111,7 @@
     nbytes,
     sync,
     thread_state,
+    wait_for,
 )
 from distributed.utils_comm import (
     gather_from_workers,
@@ -834,51 +829,32 @@ def _is_nested(iterable):
     return False
 
 
-class _MapLayer(Layer):
+class _MapExpr(Expr):
     func: Callable
-    iterables: Iterable[Any]
-    key: str | Iterable[str] | None
+    iterables: Iterable
+    key: Key
     pure: bool
-    annotations: dict[str, Any] | None
-
-    def __init__(
-        self,
-        func: Callable,
-        iterables: Iterable[Any],
-        key: str | Iterable[str] | None = None,
-        pure: bool = True,
-        annotations: dict[str, Any] | None = None,
-        **kwargs,
-    ):
-        self.func: Callable = func
-        self.iterables = [tuple(map(parse_input, iterable)) for iterable in iterables]
-        self.key: str | Iterable[str] | None = key
-        self.pure: bool = pure
-        self.kwargs = {k: parse_input(v) for k, v in kwargs.items()}
-        super().__init__(annotations=annotations)
-
-    def __repr__(self) -> str:
-        return f"{type(self).__name__} <func='{funcname(self.func)}'>"
+    annotations: dict
+    kwargs: dict
+    _cached_keys: Iterable[Key] | None
+    _parameters = [
+        "func",
+        "iterables",
+        "key",
+        "pure",
+        "annotations",
+        "kwargs",
+        "_cached_keys",
+    ]
+    _defaults = {"_cached_keys": None}
 
     @property
-    def _dict(self) -> _T_LowLevelGraph:
-        self._cached_dict: _T_LowLevelGraph
-        dsk: _T_LowLevelGraph
-
-        if hasattr(self, "_cached_dict"):
-            return self._cached_dict
-        else:
-            dsk = self._construct_graph()
-            self._cached_dict = dsk
-        return self._cached_dict
-
-    @property
-    def _keys(self) -> Iterable[Key]:
-        if hasattr(self, "_cached_keys"):
+    def keys(self) -> Iterable[Key]:
+        if self._cached_keys is not None:
             return self._cached_keys
         else:
             if isinstance(self.key, Iterable) and not isinstance(self.key, str):
-                self._cached_keys: Iterable[Key] = self.key
+                self.operands[-1] = self.key
                 return self.key
 
             else:
@@ -898,34 +874,19 @@ def _keys(self) -> Iterable[Key]:
                         if self.iterables
                         else []
                     )
-                self._cached_keys = keys
+                self.operands[-1] = keys
                 return keys
 
-    def get_output_keys(self) -> set[Key]:
-        return set(self._keys)
-
-    def get_ordered_keys(self):
-        return list(self._keys)
-
-    def is_materialized(self) -> bool:
-        return hasattr(self, "_cached_dict")
-
-    def __getitem__(self, key: Key) -> GraphNode:
-        return self._dict[key]
+    def _meta(self):
+        return []
 
-    def __iter__(self) -> Iterator[Key]:
-        return iter(self._dict)
-
-    def __len__(self) -> int:
-        return len(self._dict)
-
-    def _construct_graph(self) -> _T_LowLevelGraph:
+    def _layer(self):
         dsk: _T_LowLevelGraph = {}
 
         if not self.kwargs:
             dsk = {
                 key: Task(key, self.func, *args)
-                for key, args in zip(self._keys, zip(*self.iterables))
+                for key, args in zip(self.keys, zip(*self.iterables))
             }
 
         else:
@@ -937,12 +898,12 @@ def _construct_graph(self) -> _T_LowLevelGraph:
                     kwargs2[k] = vv.ref()
                     dsk[vv.key] = vv
                 else:
-                    kwargs2[k] = v
+                    kwargs2[k] = parse_input(v)
 
                 dsk.update(
                     {
                         key: Task(key, self.func, *args, **kwargs2)
-                        for key, args in zip(self._keys, zip(*self.iterables))
+                        for key, args in zip(self.keys, zip(*self.iterables))
                     }
                 )
         return dsk
@@ -2162,16 +2123,19 @@ def submit(
 
         if isinstance(workers, (str, Number)):
             workers = [workers]
-        dsk = {
-            key: Task(
-                key,
-                func,
-                *(parse_input(a) for a in args),
-                **{k: parse_input(v) for k, v in kwargs.items()},
-            )
-        }
+
+        expr = LLGExpr(
+            {
+                key: Task(
+                    key,
+                    func,
+                    *(parse_input(a) for a in args),
+                    **{k: parse_input(v) for k, v in kwargs.items()},
+                )
+            }
+        )
         futures = self._graph_to_futures(
-            dsk,
+            expr,
             [key],
             workers=workers,
             allow_other_workers=allow_other_workers,
@@ -2331,14 +2295,16 @@ def map(
         if allow_other_workers and workers is None:
             raise ValueError("Only use allow_other_workers= if using workers=")
 
-        dsk = _MapLayer(
+        expr = _MapExpr(
             func,
             iterables,
             key=key,
             pure=pure,
-            **kwargs,
+            # FIXME: this doesn't look right
+            annotations={},
+            kwargs=kwargs,
         )
-        keys = dsk.get_ordered_keys()
+        keys = list(expr.keys)
         if isinstance(workers, (str, Number)):
             workers = [workers]
         if workers is not None and not isinstance(workers, (list, set)):
@@ -2347,7 +2313,7 @@ def map(
         internal_priority = dict(zip(keys, range(len(keys))))
 
         futures = self._graph_to_futures(
-            dsk,
+            expr,
             keys,
             workers=workers,
             allow_other_workers=allow_other_workers,
@@ -2361,7 +2327,6 @@ def map(
         )
 
         # make sure the graph is not materialized
-        assert not dsk.is_materialized(), "Graph must be non-materialized"
         logger.debug("map(%s, ...)", funcname(func))
         return [futures[k] for k in keys]
 
@@ -3464,8 +3429,12 @@ def get(
         --------
         Client.compute : Compute asynchronous collections
         """
+        if isinstance(dsk, dict):
+            dsk = LLGExpr(dsk)
+        elif isinstance(dsk, HighLevelGraph):
+            dsk = HLGExpr(dsk)
         futures = self._graph_to_futures(
-            dsk,
+            expr=dsk,
             keys=set(flatten([keys])),
             workers=workers,
             allow_other_workers=allow_other_workers,
@@ -3667,7 +3636,6 @@ def compute(
         expr = FinalizeCompute(expr)
 
         expr = expr.optimize()
-        # FIXME: Is this actually required?
         names = list(flatten(expr.__dask_keys__()))
 
         futures_dict = self._graph_to_futures(

distributed/scheduler.py

@@ -9389,7 +9389,8 @@ def _materialize_graph(
         annotations_by_type[annotations_type].update(
             {k: (value(k) if callable(value) else value) for k in dsk}
         )
-    annotations_by_type.update(expr.__dask_annotations__())
+    for annotations_type, value in expr.__dask_annotations__().items():
+        annotations_by_type[annotations_type].update(value)
 
     dsk2 = convert_legacy_graph(dsk)
     # FIXME: There should be no need to fully materialize and copy this but some

distributed/tests/test_client.py

@@ -42,7 +42,7 @@
 import dask
 import dask.bag as db
 from dask import delayed
-from dask.tokenize import tokenize
+from dask.tokenize import TokenizationError, tokenize
 from dask.utils import get_default_shuffle_method, parse_timedelta, tmpfile
 
 from distributed import (
@@ -4639,9 +4639,9 @@ def test_recreate_error_sync(c):
     y0 = c.submit(dec, 1)
     x = c.submit(div, 1, x0)
     y = c.submit(div, 1, y0)
-    tot = c.submit(sum, x, y)
-    f = c.compute(tot)
-
+    f = c.submit(sum, x, y)
+    wait(f)
+    assert f.status == "error"
     with pytest.raises(ZeroDivisionError):
         c.recreate_error_locally(f)
     assert f.status == "error"
@@ -4658,8 +4658,8 @@ def test_recreate_task_sync(c):
     y0 = c.submit(dec, 2)
     x = c.submit(div, 1, x0)
     y = c.submit(div, 1, y0)
-    tot = c.submit(sum, [x, y])
-    f = c.compute(tot)
+    f = c.submit(sum, [x, y])
+    wait(f)
 
     assert c.recreate_task_locally(f) == 2
 
@@ -4814,7 +4814,17 @@ class Foo:
         def __getstate__(self):
             raise MyException()
 
-    with pytest.raises(TypeError, match="Could not serialize"):
+    with pytest.raises((TypeError, TokenizationError), match="serialize"):
+        future = c.submit(identity, Foo())
+
+    class Foo:
+        def __dask_tokenize__(self):
+            return 1
+
+        def __getstate__(self):
+            raise MyException()
+
+    with pytest.raises((TypeError, TokenizationError), match="serialize"):
         future = c.submit(identity, Foo())
 
     futures = c.map(inc, range(10))
@@ -4830,6 +4840,9 @@ class Foo:
         def __getstate__(self):
             return 1
 
+        def __dask_tokenize__(self):
+            return 1
+
         def __setstate__(self, state):
             raise MyException("hello")
 
@@ -4855,6 +4868,9 @@ def __getstate__(self):
         def __setstate__(self, state):
             raise MyException("hello")
 
+        def __dask_tokenize__(self):
+            return 1
+
         def __call__(self, *args):
             return 1