Repository: dask/distributed

Issue:
Use Layer annotations
In https://github.com/dask/dask/pull/6767 task annotations were added to `HighLevelGraph` `Layers`. We can now start supporting `priority`, `resources`, `retries`, etc. annotations here in the distributed scheduler.

@sjperkins you've been leading the task annotation effort and expressed an interested in incorporating them into `distributed`. Some colleagues and I have time to dedicate to this if now is a busy time for you, but I first wanted to check in with you to see if you're still interested in working on this


Discussion:
Thanks for the reminder @jrbourbeau. I've been on leave this past week, but this is first on my to-do list this week.
Awesome, thank you for continuing to lead this effort (also hope your leave was relaxing : ) )

Resolving patch:
diff --git a/distributed/client.py b/distributed/client.py
--- a/distributed/client.py
+++ b/distributed/client.py
@@ -11,10 +11,9 @@
 from functools import partial
 import html
 import inspect
-import itertools
 import json
 import logging
-from numbers import Number, Integral
+from numbers import Number
 import os
 import sys
 import uuid
@@ -1507,13 +1506,13 @@ def submit(
         pure: bool (defaults to True)
             Whether or not the function is pure.  Set ``pure=False`` for
             impure functions like ``np.random.random``.
-        workers: set, iterable of sets
-            A set of worker hostnames on which computations may be performed.
-            Leave empty to default to all workers (common case)
+        workers: string or iterable of strings
+            A set of worker addresses or hostnames on which computations may be
+            performed. Leave empty to default to all workers (common case)
         key: str
             Unique identifier for the task.  Defaults to function-name and hash
         allow_other_workers: bool (defaults to False)
-            Used with `workers`. Indicates whether or not the computations
+            Used with ``workers``. Indicates whether or not the computations
             may be performed on workers that are not in the `workers` set(s).
         retries: int (default to 0)
             Number of allowed automatic retries if the task fails
@@ -1523,9 +1522,10 @@ def submit(
         fifo_timeout: str timedelta (default '100ms')
             Allowed amount of time between calls to consider the same priority
         resources: dict (defaults to {})
-            Defines the `resources` this job requires on the worker; e.g.
-            ``{'GPU': 2}``. See :doc:`worker resources <resources>` for details
-            on defining resources.
+            Defines the ``resources`` each instance of this mapped task requires
+            on the worker; e.g. ``{'GPU': 2}``.
+            See :doc:`worker resources <resources>` for details on defining
+            resources.
         actor: bool (default False)
             Whether this task should exist on the worker as a stateful actor.
             See :doc:`actors` for additional details.
@@ -1571,12 +1571,6 @@ def submit(
 
         if isinstance(workers, (str, Number)):
             workers = [workers]
-        if workers is not None:
-            restrictions = {skey: workers}
-            loose_restrictions = [skey] if allow_other_workers else []
-        else:
-            restrictions = {}
-            loose_restrictions = []
 
         if kwargs:
             dsk = {skey: (apply, func, list(args), kwargs)}
@@ -1586,11 +1580,11 @@ def submit(
         futures = self._graph_to_futures(
             dsk,
             [skey],
-            restrictions,
-            loose_restrictions,
+            workers=workers,
+            allow_other_workers=allow_other_workers,
             priority={skey: 0},
             user_priority=priority,
-            resources={skey: resources} if resources else None,
+            resources=resources,
             retries=retries,
             fifo_timeout=fifo_timeout,
             actors=actor,
@@ -1631,7 +1625,7 @@ def map(
         pure: bool (defaults to True)
             Whether or not the function is pure.  Set ``pure=False`` for
             impure functions like ``np.random.random``.
-        workers: set, iterable of sets
+        workers: string or iterable of strings
             A set of worker hostnames on which computations may be performed.
             Leave empty to default to all workers (common case)
         allow_other_workers: bool (defaults to False)
@@ -1646,8 +1640,8 @@ def map(
             Allowed amount of time between calls to consider the same priority
         resources: dict (defaults to {})
             Defines the `resources` each instance of this mapped task requires
-            on the worker; e.g. ``{'GPU': 2}``. See
-            :doc:`worker resources <resources>` for details on defining
+            on the worker; e.g. ``{'GPU': 2}``.
+            See :doc:`worker resources <resources>` for details on defining
             resources.
         actor: bool (default False)
             Whether these tasks should exist on the worker as stateful actors.
@@ -1763,39 +1757,16 @@ def map(
 
         if isinstance(workers, (str, Number)):
             workers = [workers]
-        if isinstance(workers, (list, set)):
-            if workers and isinstance(first(workers), (list, set)):
-                if len(workers) != len(keys):
-                    raise ValueError(
-                        "You only provided %d worker restrictions"
-                        " for a sequence of length %d" % (len(workers), len(keys))
-                    )
-                restrictions = dict(zip(keys, workers))
-            else:
-                restrictions = {k: workers for k in keys}
-        elif workers is None:
-            restrictions = {}
-        else:
+        if workers is not None and not isinstance(workers, (list, set)):
             raise TypeError("Workers must be a list or set of workers or None")
-        if allow_other_workers not in (True, False, None):
-            raise TypeError("allow_other_workers= must be True or False")
-        if allow_other_workers is True:
-            loose_restrictions = set(keys)
-        else:
-            loose_restrictions = set()
 
         internal_priority = dict(zip(keys, range(len(keys))))
 
-        if resources:
-            resources = {k: resources for k in keys}
-        else:
-            resources = None
-
         futures = self._graph_to_futures(
             dsk,
             keys,
-            restrictions,
-            loose_restrictions,
+            workers=workers,
+            allow_other_workers=allow_other_workers,
             priority=internal_priority,
             resources=resources,
             retries=retries,
@@ -2551,8 +2522,8 @@ def _graph_to_futures(
         self,
         dsk,
         keys,
-        restrictions=None,
-        loose_restrictions=None,
+        workers=None,
+        allow_other_workers=None,
         priority=None,
         user_priority=0,
         resources=None,
@@ -2561,38 +2532,35 @@ def _graph_to_futures(
         actors=None,
     ):
         with self._refcount_lock:
-            if resources:
-                resources = self._expand_resources(
-                    resources, all_keys=itertools.chain(dsk, keys)
-                )
-                resources = {stringify(k): v for k, v in resources.items()}
-
-            if retries:
-                retries = self._expand_retries(
-                    retries, all_keys=itertools.chain(dsk, keys)
-                )
-
             if actors is not None and actors is not True and actors is not False:
                 actors = list(self._expand_key(actors))
 
-            if restrictions:
-                restrictions = keymap(stringify, restrictions)
-                restrictions = valmap(list, restrictions)
-
-            if loose_restrictions is not None:
-                loose_restrictions = list(map(stringify, loose_restrictions))
-
             keyset = set(keys)
 
             # Make sure `dsk` is a high level graph
             if not isinstance(dsk, HighLevelGraph):
                 dsk = HighLevelGraph.from_collections(id(dsk), dsk, dependencies=())
 
-            if isinstance(retries, Number) and retries > 0:
-                retries = {k: retries for k in dsk}
-
             dsk = highlevelgraph_pack(dsk, self, keyset)
 
+            annotations = {}
+            if user_priority:
+                annotations["priority"] = user_priority
+            if workers:
+                if not isinstance(workers, (list, tuple, set)):
+                    workers = [workers]
+                annotations["workers"] = workers
+            if retries:
+                annotations["retries"] = retries
+            if allow_other_workers not in (True, False, None):
+                raise TypeError("allow_other_workers= must be True, False, or None")
+            if allow_other_workers:
+                annotations["allow_other_workers"] = allow_other_workers
+            if resources:
+                annotations["resources"] = resources
+
+            annotations = merge(dask.config.get("annotations", {}), annotations)
+
             # Create futures before sending graph (helps avoid contention)
             futures = {key: Future(key, self, inform=False) for key in keyset}
             self._send_to_scheduler(
@@ -2600,14 +2568,10 @@ def _graph_to_futures(
                     "op": "update-graph-hlg",
                     "hlg": dsk,
                     "keys": list(map(stringify, keys)),
-                    "restrictions": restrictions or {},
-                    "loose_restrictions": loose_restrictions,
                     "priority": priority,
-                    "user_priority": user_priority,
-                    "resources": resources,
                     "submitting_task": getattr(thread_state, "key", None),
-                    "retries": retries,
                     "fifo_timeout": fifo_timeout,
+                    "annotations": annotations,
                     "actors": actors,
                 }
             )
@@ -2617,8 +2581,8 @@ def get(
         self,
         dsk,
         keys,
-        restrictions=None,
-        loose_restrictions=None,
+        workers=None,
+        allow_other_workers=None,
         resources=None,
         sync=True,
         asynchronous=None,
@@ -2635,14 +2599,22 @@ def get(
         ----------
         dsk: dict
         keys: object, or nested lists of objects
-        restrictions: dict (optional)
-            A mapping of {key: {set of worker hostnames}} that restricts where
-            jobs can take place
+        workers: string or iterable of strings
+            A set of worker addresses or hostnames on which computations may be
+            performed. Leave empty to default to all workers (common case)
+        allow_other_workers: bool (defaults to False)
+            Used with ``workers``. Indicates whether or not the computations
+            may be performed on workers that are not in the `workers` set(s).
         retries: int (default to 0)
             Number of allowed automatic retries if computing a result fails
         priority: Number
             Optional prioritization of task.  Zero is default.
             Higher priorities take precedence
+        resources: dict (defaults to {})
+            Defines the ``resources`` each instance of this mapped task requires
+            on the worker; e.g. ``{'GPU': 2}``.
+            See :doc:`worker resources <resources>` for details on defining
+            resources.
         sync: bool (optional)
             Returns Futures if False or concrete values if True (default).
         direct: bool
@@ -2664,8 +2636,8 @@ def get(
         futures = self._graph_to_futures(
             dsk,
             keys=set(flatten([keys])),
-            restrictions=restrictions,
-            loose_restrictions=loose_restrictions,
+            workers=workers,
+            allow_other_workers=allow_other_workers,
             resources=resources,
             fifo_timeout=fifo_timeout,
             retries=retries,
@@ -2770,15 +2742,12 @@ def compute(
             Returns Futures if False (default) or concrete values if True
         optimize_graph: bool
             Whether or not to optimize the underlying graphs
-        workers: str, list, dict
-            Which workers can run which parts of the computation
-            If a string or list then the output collections will run on the listed
-            workers, but other sub-computations can run anywhere
-            If a dict then keys should be (tuples of) collections or
-            task keys and values should be addresses or lists.
-        allow_other_workers: bool, list
-            If True then all restrictions in workers= are considered loose
-            If a list then only the keys for the listed collections are loose
+        workers: string or iterable of strings
+            A set of worker hostnames on which computations may be performed.
+            Leave empty to default to all workers (common case)
+        allow_other_workers: bool (defaults to False)
+            Used with `workers`. Indicates whether or not the computations
+            may be performed on workers that are not in the `workers` set(s).
         retries: int (default to 0)
             Number of allowed automatic retries if computing a result fails
         priority: Number
@@ -2792,9 +2761,8 @@ def compute(
             be expensive. If none of the arguments contain any dask objects,
             set ``traverse=False`` to avoid doing this traversal.
         resources: dict (defaults to {})
-            Defines the `resources` these tasks require on the worker. Can
-            specify global resources (``{'GPU': 2}``), or per-task resources
-            (``{'x': {'GPU': 1}, 'y': {'SSD': 4}}``), but not both.
+            Defines the `resources` each instance of this mapped task requires
+            on the worker; e.g. ``{'GPU': 2}``.
             See :doc:`worker resources <resources>` for details on defining
             resources.
         actors: bool or dict (default None)
@@ -2857,13 +2825,6 @@ def compute(
             else:
                 dsk2[name] = (func, keys) + extra_args
 
-        restrictions, loose_restrictions = self.get_restrictions(
-            collections, workers, allow_other_workers
-        )
-
-        if not isinstance(priority, Number):
-            priority = {k: p for c, p in priority.items() for k in self._expand_key(c)}
-
         if not isinstance(dsk, HighLevelGraph):
             dsk = HighLevelGraph.from_collections(id(dsk), dsk, dependencies=())
 
@@ -2878,8 +2839,8 @@ def compute(
         futures_dict = self._graph_to_futures(
             dsk,
             names,
-            restrictions,
-            loose_restrictions,
+            workers=workers,
+            allow_other_workers=allow_other_workers,
             resources=resources,
             retries=retries,
             user_priority=priority,
@@ -2931,15 +2892,12 @@ def persist(
             Collections like dask.array or dataframe or dask.value objects
         optimize_graph: bool
             Whether or not to optimize the underlying graphs
-        workers: str, list, dict
-            Which workers can run which parts of the computation
-            If a string or list then the output collections will run on the listed
-            workers, but other sub-computations can run anywhere
-            If a dict then keys should be (tuples of) collections or
-            task keys and values should be addresses or lists.
-        allow_other_workers: bool, list
-            If True then all restrictions in workers= are considered loose
-            If a list then only the keys for the listed collections are loose
+        workers: string or iterable of strings
+            A set of worker hostnames on which computations may be performed.
+            Leave empty to default to all workers (common case)
+        allow_other_workers: bool (defaults to False)
+            Used with `workers`. Indicates whether or not the computations
+            may be performed on workers that are not in the `workers` set(s).
         retries: int (default to 0)
             Number of allowed automatic retries if computing a result fails
         priority: Number
@@ -2948,9 +2906,8 @@ def persist(
         fifo_timeout: timedelta str (defaults to '60s')
             Allowed amount of time between calls to consider the same priority
         resources: dict (defaults to {})
-            Defines the `resources` these tasks require on the worker. Can
-            specify global resources (``{'GPU': 2}``), or per-task resources
-            (``{'x': {'GPU': 1}, 'y': {'SSD': 4}}``), but not both.
+            Defines the `resources` each instance of this mapped task requires
+            on the worker; e.g. ``{'GPU': 2}``.
             See :doc:`worker resources <resources>` for details on defining
             resources.
         actors: bool or dict (default None)
@@ -2985,18 +2942,11 @@ def persist(
 
         names = {k for c in collections for k in flatten(c.__dask_keys__())}
 
-        restrictions, loose_restrictions = self.get_restrictions(
-            collections, workers, allow_other_workers
-        )
-
-        if not isinstance(priority, Number):
-            priority = {k: p for c, p in priority.items() for k in self._expand_key(c)}
-
         futures = self._graph_to_futures(
             dsk,
             names,
-            restrictions,
-            loose_restrictions,
+            workers=workers,
+            allow_other_workers=allow_other_workers,
             resources=resources,
             retries=retries,
             user_priority=priority,
@@ -3881,90 +3831,6 @@ def _expand_key(cls, k):
             else:
                 yield stringify(kk)
 
-    @classmethod
-    def _expand_retries(cls, retries, all_keys):
-        """
-        Expand the user-provided "retries" specification
-        to a {task key: Integral} dictionary.
-        """
-        if retries and isinstance(retries, dict):
-            result = {
-                name: value
-                for key, value in retries.items()
-                for name in cls._expand_key(key)
-            }
-        elif isinstance(retries, Integral):
-            # Each task unit may potentially fail, allow retrying all of them
-            result = {name: retries for name in all_keys}
-        else:
-            raise TypeError(
-                "`retries` should be an integer or dict, got %r" % (type(retries))
-            )
-        return keymap(stringify, result)
-
-    def _expand_resources(cls, resources, all_keys):
-        """
-        Expand the user-provided "resources" specification
-        to a {task key: {resource name: Number}} dictionary.
-        """
-        # Resources can either be a single dict such as {'GPU': 2},
-        # indicating a requirement for all keys, or a nested dict
-        # such as {'x': {'GPU': 1}, 'y': {'SSD': 4}} indicating
-        # per-key requirements
-        if not isinstance(resources, dict):
-            raise TypeError("`resources` should be a dict, got %r" % (type(resources)))
-
-        per_key_reqs = {}
-        global_reqs = {}
-        all_keys = list(all_keys)
-        for k, v in resources.items():
-            if isinstance(v, dict):
-                # It's a per-key requirement
-                per_key_reqs.update((kk, v) for kk in cls._expand_key(k))
-            else:
-                # It's a global requirement
-                global_reqs.update((kk, {k: v}) for kk in all_keys)
-
-        if global_reqs and per_key_reqs:
-            raise ValueError(
-                "cannot have both per-key and all-key requirements "
-                "in resources dict %r" % (resources,)
-            )
-        return global_reqs or per_key_reqs
-
-    @classmethod
-    def get_restrictions(cls, collections, workers, allow_other_workers):
-        """ Get restrictions from inputs to compute/persist """
-        if isinstance(workers, (str, tuple, list)):
-            workers = {tuple(collections): workers}
-        if isinstance(workers, dict):
-            restrictions = {}
-            for colls, ws in workers.items():
-                if isinstance(ws, str):
-                    ws = [ws]
-                if dask.is_dask_collection(colls):
-                    keys = flatten(colls.__dask_keys__())
-                elif isinstance(colls, str):
-                    keys = [colls]
-                else:
-                    keys = list(
-                        {k for c in flatten(colls) for k in flatten(c.__dask_keys__())}
-                    )
-                restrictions.update({k: ws for k in keys})
-        else:
-            restrictions = {}
-
-        if allow_other_workers is True:
-            loose_restrictions = list(restrictions)
-        elif allow_other_workers:
-            loose_restrictions = list(
-                {k for c in flatten(allow_other_workers) for k in c.__dask_keys__()}
-            )
-        else:
-            loose_restrictions = []
-
-        return restrictions, loose_restrictions
-
     @staticmethod
     def collections_to_dsk(collections, *args, **kwargs):
         return collections_to_dsk(collections, *args, **kwargs)
diff --git a/distributed/protocol/highlevelgraph.py b/distributed/protocol/highlevelgraph.py
--- a/distributed/protocol/highlevelgraph.py
+++ b/distributed/protocol/highlevelgraph.py
@@ -124,18 +124,7 @@ def highlevelgraph_pack(hlg: HighLevelGraph, client, client_keys):
     return dumps_msgpack({"layers": layers})
 
 
-def _materialized_layer_unpack(state, dsk, dependencies, annotations):
-    dsk.update(state["dsk"])
-    for k, v in state["dependencies"].items():
-        dependencies[k] = list(set(dependencies.get(k, ())) | set(v))
-
-    if state["annotations"]:
-        annotations.update(
-            Layer.expand_annotations(state["annotations"], state["dsk"].keys())
-        )
-
-
-def highlevelgraph_unpack(dumped_hlg):
+def highlevelgraph_unpack(dumped_hlg, annotations: dict):
     """Unpack the high level graph for Scheduler -> Worker communication
 
     The approach is to delegate the packaging to each layer in the high
@@ -148,6 +137,11 @@ def highlevelgraph_unpack(dumped_hlg):
     dumped_hlg: list of header and payload
         Packed high level graph serialized by dumps_msgpack
 
+    annotations: dict
+        A top-level annotations object which may be partially populated,
+        and which may be further filled by annotations from the layers
+        of the dumped_hlg.
+
     Returns
     -------
     dsk: dict
@@ -157,18 +151,21 @@ def highlevelgraph_unpack(dumped_hlg):
     annotations: dict
         Annotations for `dsk`
     """
-
     hlg = loads_msgpack(*dumped_hlg)
 
     dsk = {}
     deps = {}
-    annotations = {}
+    out_annotations = {}
     for layer in hlg["layers"]:
+        if annotations:
+            if layer["state"]["annotations"] is None:
+                layer["state"]["annotations"] = {}
+            layer["state"]["annotations"].update(annotations)
         if layer["__module__"] is None:  # Default implementation
-            unpack_func = _materialized_layer_unpack
+            unpack_func = Layer.__dask_distributed_unpack__
         else:
             mod = import_allowed_module(layer["__module__"])
             unpack_func = getattr(mod, layer["__name__"]).__dask_distributed_unpack__
-        unpack_func(layer["state"], dsk, deps, annotations)
+        unpack_func(layer["state"], dsk, deps, out_annotations)
 
-    return dsk, deps, annotations
+    return dsk, deps, out_annotations
diff --git a/distributed/scheduler.py b/distributed/scheduler.py
--- a/distributed/scheduler.py
+++ b/distributed/scheduler.py
@@ -3612,9 +3612,10 @@ def update_graph_hlg(
         user_priority=0,
         actors=None,
         fifo_timeout=0,
+        annotations=None,
     ):
 
-        dsk, dependencies, annotations = highlevelgraph_unpack(hlg)
+        dsk, dependencies, annotations = highlevelgraph_unpack(hlg, annotations)
 
         # Remove any self-dependencies (happens on test_publish_bag() and others)
         for k, v in dependencies.items():


Test patch:
diff --git a/distributed/tests/test_client.py b/distributed/tests/test_client.py
--- a/distributed/tests/test_client.py
+++ b/distributed/tests/test_client.py
@@ -250,21 +250,31 @@ async def test_compute_retries(c, s, a, b):
     x = c.compute(delayed(varying(args))(), retries=2)
     assert await x == 3
 
+
+@gen_cluster(client=True)
+async def test_compute_retries_annotations(c, s, a, b):
     # Per-future retries
     xargs = [ZeroDivisionError("one"), ZeroDivisionError("two"), 30, 40]
     yargs = [ZeroDivisionError("five"), ZeroDivisionError("six"), 70]
     zargs = [80, 90, 100]
 
-    x, y = [delayed(varying(args))() for args in (xargs, yargs)]
-    x, y = c.compute([x, y], retries={x: 2})
+    with dask.annotate(retries=2):
+        x = delayed(varying(xargs))()
+    y = delayed(varying(yargs))()
+
+    x, y = c.compute([x, y], optimize_graph=False)
     gc.collect()
 
     assert await x == 30
     with pytest.raises(ZeroDivisionError, match="five"):
         await y
 
-    x, y, z = [delayed(varying(args))() for args in (xargs, yargs, zargs)]
-    x, y, z = c.compute([x, y, z], retries={(y, z): 2})
+    x = delayed(varying(xargs))()
+    with dask.annotate(retries=2):
+        y = delayed(varying(yargs))()
+        z = delayed(varying(zargs))()
+
+    x, y, z = c.compute([x, y, z], optimize_graph=False)
 
     with pytest.raises(ZeroDivisionError, match="one"):
         await x
@@ -322,13 +332,20 @@ async def test_persist_retries(c, s, a, b):
     x = c.compute(x)
     assert await x == 3
 
+
+@gen_cluster(client=True)
+async def test_persist_retries_annotations(c, s, a, b):
     # Per-key retries
     xargs = [ZeroDivisionError("one"), ZeroDivisionError("two"), 30, 40]
     yargs = [ZeroDivisionError("five"), ZeroDivisionError("six"), 70]
     zargs = [80, 90, 100]
 
-    x, y, z = [delayed(varying(args))() for args in (xargs, yargs, zargs)]
-    x, y, z = c.persist([x, y, z], retries={(y, z): 2})
+    x = delayed(varying(xargs))()
+    with dask.annotate(retries=2):
+        y = delayed(varying(yargs))()
+        z = delayed(varying(zargs))()
+
+    x, y, z = c.persist([x, y, z], optimize_graph=False)
     x, y, z = c.compute([x, y, z])
 
     with pytest.raises(ZeroDivisionError, match="one"):
@@ -910,31 +927,6 @@ async def test_restrictions_ip_port(c, s, a, b):
     assert y.key in b.data
 
 
-@gen_cluster(client=True)
-async def test_restrictions_ip_port_task_key(c, s, a, b):
-    # Create a long dependency list
-    tasks = [delayed(inc)(1)]
-    for _ in range(100):
-        tasks.append(delayed(add)(tasks[-1], random.choice(tasks)))
-
-    last_task = tasks[-1]
-
-    # calculate all dependency keys
-    all_tasks = list(last_task.__dask_graph__())
-    # only restrict to a single worker
-    workers = {d: a.address for d in all_tasks}
-    result = c.compute(last_task, workers=workers)
-    await result
-
-    # all tasks should have been calculated by the first worker
-    for task in tasks:
-        assert s.worker_restrictions[task.key] == {a.address}
-
-    # and the data should also be there
-    assert last_task.key in a.data
-    assert last_task.key not in b.data
-
-
 @pytest.mark.skipif(
     not sys.platform.startswith("linux"), reason="Need 127.0.0.2 to mean localhost"
 )
@@ -948,16 +940,6 @@ async def test_restrictions_map(c, s, a, b):
     for x in L:
         assert s.host_restrictions[x.key] == {a.ip}
 
-    L = c.map(inc, [10, 11, 12], workers=[{a.ip}, {a.ip, b.ip}, {b.ip}])
-    await wait(L)
-
-    assert s.host_restrictions[L[0].key] == {a.ip}
-    assert s.host_restrictions[L[1].key] == {a.ip, b.ip}
-    assert s.host_restrictions[L[2].key] == {b.ip}
-
-    with pytest.raises(ValueError):
-        c.map(inc, [10, 11, 12], workers=[{a.ip}])
-
 
 @pytest.mark.skipif(
     not sys.platform.startswith("linux"), reason="Need 127.0.0.2 to mean localhost"
@@ -965,13 +947,28 @@ async def test_restrictions_map(c, s, a, b):
 @gen_cluster([("127.0.0.1", 1), ("127.0.0.2", 2)], client=True)
 async def test_restrictions_get(c, s, a, b):
     dsk = {"x": 1, "y": (inc, "x"), "z": (inc, "y")}
-    restrictions = {"y": {a.ip}, "z": {b.ip}}
 
-    futures = c.get(dsk, ["y", "z"], restrictions, sync=False)
+    futures = c.get(dsk, ["y", "z"], workers=a.ip, sync=False)
     result = await c.gather(futures)
     assert result == [2, 3]
     assert "y" in a.data
-    assert "z" in b.data
+    assert "z" in a.data
+    assert len(b.data) == 0
+
+
+@gen_cluster(client=True)
+async def test_restrictions_get_annotate(c, s, a, b):
+    x = 1
+    with dask.annotate(workers=a.address):
+        y = delayed(inc)(x)
+    with dask.annotate(workers=b.address):
+        z = delayed(inc)(y)
+
+    futures = c.get(z.__dask_graph__(), [y.key, z.key], sync=False)
+    result = await c.gather(futures)
+    assert result == [2, 3]
+    assert y.key in a.data
+    assert z.key in b.data
 
 
 @gen_cluster(client=True)
@@ -4100,6 +4097,29 @@ async def run2():
     await asyncio.gather(run1(), run2())
 
 
+@nodebug  # test timing is fragile
+@gen_cluster(nthreads=[("127.0.0.1", 1)] * 3, client=True)
+async def test_persist_workers_annotate(e, s, a, b, c):
+    with dask.annotate(workers=a.address, allow_other_workers=False):
+        L1 = [delayed(inc)(i) for i in range(4)]
+    with dask.annotate(workers=b.address, allow_other_workers=False):
+        total = delayed(sum)(L1)
+    with dask.annotate(workers=c.address, allow_other_workers=True):
+        L2 = [delayed(add)(i, total) for i in L1]
+    with dask.annotate(workers=b.address, allow_other_workers=True):
+        total2 = delayed(sum)(L2)
+
+    # TODO: once annotations are faithfully forwarded upon graph optimization,
+    # we shouldn't need to disable that here.
+    out = e.persist(L1 + L2 + [total, total2], optimize_graph=False)
+
+    await wait(out)
+    assert all(v.key in a.data for v in L1)
+    assert total.key in b.data
+
+    assert s.loose_restrictions == {total2.key} | {v.key for v in L2}
+
+
 @nodebug  # test timing is fragile
 @gen_cluster(nthreads=[("127.0.0.1", 1)] * 3, client=True)
 async def test_persist_workers(e, s, a, b, c):
@@ -4110,20 +4130,40 @@ async def test_persist_workers(e, s, a, b, c):
 
     out = e.persist(
         L1 + L2 + [total, total2],
-        workers={
-            tuple(L1): a.address,
-            total: b.address,
-            tuple(L2): [c.address],
-            total2: b.address,
-        },
-        allow_other_workers=L2 + [total2],
+        workers=[a.address, b.address],
+        allow_other_workers=True,
     )
 
     await wait(out)
-    assert all(v.key in a.data for v in L1)
-    assert total.key in b.data
 
-    assert s.loose_restrictions == {total2.key} | {v.key for v in L2}
+    for v in L1 + L2 + [total, total2]:
+        assert s.worker_restrictions[v.key] == {a.address, b.address}
+    assert not any(c.address in r for r in s.worker_restrictions)
+
+    assert s.loose_restrictions == {total.key, total2.key} | {v.key for v in L1 + L2}
+
+
+@gen_cluster(nthreads=[("127.0.0.1", 1)] * 3, client=True)
+async def test_compute_workers_annotate(e, s, a, b, c):
+    with dask.annotate(workers=a.address, allow_other_workers=True):
+        L1 = [delayed(inc)(i) for i in range(4)]
+    with dask.annotate(workers=b.address, allow_other_workers=True):
+        total = delayed(sum)(L1)
+    with dask.annotate(workers=[c.address]):
+        L2 = [delayed(add)(i, total) for i in L1]
+
+    # TODO: once annotations are faithfully forwarded upon graph optimization,
+    # we shouldn't need to disable that here.
+    out = e.compute(L1 + L2 + [total], optimize_graph=False)
+
+    await wait(out)
+    for v in L1:
+        assert s.worker_restrictions[v.key] == {a.address}
+    for v in L2:
+        assert s.worker_restrictions[v.key] == {c.address}
+    assert s.worker_restrictions[total.key] == {b.address}
+
+    assert s.loose_restrictions == {total.key} | {v.key for v in L1}
 
 
 @gen_cluster(nthreads=[("127.0.0.1", 1)] * 3, client=True)
@@ -4134,18 +4174,16 @@ async def test_compute_workers(e, s, a, b, c):
 
     out = e.compute(
         L1 + L2 + [total],
-        workers={tuple(L1): a.address, total: b.address, tuple(L2): [c.address]},
-        allow_other_workers=L1 + [total],
+        workers=[a.address, b.address],
+        allow_other_workers=True,
     )
 
     await wait(out)
-    for v in L1:
-        assert s.worker_restrictions[v.key] == {a.address}
-    for v in L2:
-        assert s.worker_restrictions[v.key] == {c.address}
-    assert s.worker_restrictions[total.key] == {b.address}
+    for v in L1 + L2 + [total]:
+        assert s.worker_restrictions[v.key] == {a.address, b.address}
+    assert not any(c.address in r for r in s.worker_restrictions)
 
-    assert s.loose_restrictions == {total.key} | {v.key for v in L1}
+    assert s.loose_restrictions == {total.key} | {v.key for v in L1 + L2}
 
 
 @gen_cluster(client=True)
@@ -4162,28 +4200,6 @@ async def test_compute_nested_containers(c, s, a, b):
     assert result["y"] == 123
 
 
-def test_get_restrictions():
-    L1 = [delayed(inc)(i) for i in range(4)]
-    total = delayed(sum)(L1)
-    L2 = [delayed(add)(i, total) for i in L1]
-
-    r1, loose = Client.get_restrictions(L2, "127.0.0.1", False)
-    assert r1 == {d.key: ["127.0.0.1"] for d in L2}
-    assert not loose
-
-    r1, loose = Client.get_restrictions(L2, ["127.0.0.1"], True)
-    assert r1 == {d.key: ["127.0.0.1"] for d in L2}
-    assert set(loose) == {d.key for d in L2}
-
-    r1, loose = Client.get_restrictions(L2, {total: "127.0.0.1"}, True)
-    assert r1 == {total.key: ["127.0.0.1"]}
-    assert loose == [total.key]
-
-    r1, loose = Client.get_restrictions(L2, {(total,): "127.0.0.1"}, True)
-    assert r1 == {total.key: ["127.0.0.1"]}
-    assert loose == [total.key]
-
-
 @gen_cluster(client=True)
 async def test_scatter_type(c, s, a, b):
     [future] = await c.scatter([1])
@@ -6318,6 +6334,43 @@ async def test_annotations_task_state(c, s, a, b):
     )
 
 
+@pytest.mark.parametrize("fn", ["compute", "persist"])
+def test_annotations_compute_time(fn):
+    da = pytest.importorskip("dask.array")
+
+    @gen_cluster(client=True)
+    async def test(c, s, a, b):
+        x = da.ones(10, chunks=(5,))
+
+        with dask.annotate(foo="bar"):
+            # Turn off optimization to avoid rewriting layers and picking up annotations
+            # that way. Instead, we want `compute`/`persist` to be able to pick them up.
+            x = await getattr(c, fn)(x, optimize_graph=False)
+
+        assert all({"foo": "bar"} == ts.annotations for ts in s.tasks.values())
+
+    test()
+
+
+@pytest.mark.xfail(reason="https://github.com/dask/dask/issues/7036")
+@gen_cluster(client=True)
+async def test_annotations_survive_optimization(c, s, a, b):
+    da = pytest.importorskip("dask.array")
+
+    with dask.annotate(foo="bar"):
+        x = da.ones(10, chunks=(5,))
+
+    ann = x.__dask_graph__().layers[x.name].annotations
+    assert ann is not None
+    assert ann.get("foo", None) == "bar"
+
+    (xx,) = dask.optimize(x)
+
+    ann = xx.__dask_graph__().layers[x.name].annotations
+    assert ann is not None
+    assert ann.get("foo", None) == "bar"
+
+
 @gen_cluster(client=True)
 async def test_annotations_priorities(c, s, a, b):
     da = pytest.importorskip("dask.array")
@@ -6363,6 +6416,37 @@ async def test_annotations_retries(c, s, a, b):
     assert all(ts.annotations == {"retries": 2} for ts in s.tasks.values())
 
 
+@gen_cluster(client=True)
+async def test_annotations_blockwise_unpack(c, s, a, b):
+    da = pytest.importorskip("dask.array")
+    from dask.array.utils import assert_eq
+    import numpy as np
+
+    # A flaky doubling function -- need extra args because it is called before
+    # application to establish dtype/meta.
+    scale = varying([ZeroDivisionError("one"), ZeroDivisionError("two"), 2, 2])
+
+    def flaky_double(x):
+        return scale() * x
+
+    # A reliable double function.
+    def reliable_double(x):
+        return 2 * x
+
+    x = da.ones(10, chunks=(5,))
+
+    # The later annotations should not override the earlier annotations
+    with dask.annotate(retries=2):
+        y = x.map_blocks(flaky_double, meta=np.array((), dtype=np.float))
+    with dask.annotate(retries=0):
+        z = y.map_blocks(reliable_double, meta=np.array((), dtype=np.float))
+
+    with dask.config.set(optimization__fuse__active=False):
+        z = await c.compute(z)
+
+    assert_eq(z, np.ones(10) * 4.0)
+
+
 @gen_cluster(
     client=True,
     nthreads=[
@@ -6402,3 +6486,12 @@ async def test_annotations_loose_restrictions(c, s, a, b):
             for ts in s.tasks.values()
         ]
     )
+
+
+@gen_cluster(client=True)
+async def test_workers_collection_restriction(c, s, a, b):
+    da = pytest.importorskip("dask.array")
+
+    future = c.compute(da.arange(10), workers=a.address)
+    await future
+    assert a.data and not b.data
diff --git a/distributed/tests/test_priorities.py b/distributed/tests/test_priorities.py
--- a/distributed/tests/test_priorities.py
+++ b/distributed/tests/test_priorities.py
@@ -67,23 +67,27 @@ async def test_persist(c, s):
 
 
 @gen_cluster(client=True)
-async def test_expand_compute(c, s, a, b):
-    low = delayed(inc)(1)
+async def test_annotate_compute(c, s, a, b):
+    with dask.annotate(priority=-1):
+        low = delayed(inc)(1)
+    with dask.annotate(priority=1):
+        high = delayed(inc)(2)
     many = [delayed(slowinc)(i, delay=0.1) for i in range(10)]
-    high = delayed(inc)(2)
 
-    low, many, high = c.compute([low, many, high], priority={low: -1, high: 1})
+    low, many, high = c.compute([low, many, high], optimize_graph=False)
     await wait(high)
     assert s.tasks[low.key].state == "processing"
 
 
 @gen_cluster(client=True)
-async def test_expand_persist(c, s, a, b):
-    low = delayed(inc)(1, dask_key_name="low")
+async def test_annotate_persist(c, s, a, b):
+    with dask.annotate(priority=-1):
+        low = delayed(inc)(1, dask_key_name="low")
+    with dask.annotate(priority=1):
+        high = delayed(inc)(2, dask_key_name="high")
     many = [delayed(slowinc)(i, delay=0.1) for i in range(4)]
-    high = delayed(inc)(2, dask_key_name="high")
 
-    low, high, x, y, z, w = persist(low, high, *many, priority={low: -1, high: 1})
+    low, high, x, y, z, w = persist(low, high, *many, optimize_graph=False)
     await wait(high)
     assert s.tasks[low.key].state == "processing"
 
diff --git a/distributed/tests/test_resources.py b/distributed/tests/test_resources.py
--- a/distributed/tests/test_resources.py
+++ b/distributed/tests/test_resources.py
@@ -1,6 +1,7 @@
 import asyncio
 from time import time
 
+import dask
 from dask import delayed
 from dask.utils import stringify
 import pytest
@@ -131,10 +132,12 @@ async def test_map(c, s, a, b):
     ],
 )
 async def test_persist(c, s, a, b):
-    x = delayed(inc)(1)
-    y = delayed(inc)(x)
+    with dask.annotate(resources={"A": 1}):
+        x = delayed(inc)(1)
+    with dask.annotate(resources={"B": 1}):
+        y = delayed(inc)(x)
 
-    xx, yy = c.persist([x, y], resources={x: {"A": 1}, y: {"B": 1}})
+    xx, yy = c.persist([x, y], optimize_graph=False)
 
     await wait([xx, yy])
 
@@ -150,10 +153,12 @@ async def test_persist(c, s, a, b):
     ],
 )
 async def test_compute(c, s, a, b):
-    x = delayed(inc)(1)
-    y = delayed(inc)(x)
+    with dask.annotate(resources={"A": 1}):
+        x = delayed(inc)(1)
+    with dask.annotate(resources={"B": 1}):
+        y = delayed(inc)(x)
 
-    yy = c.compute(y, resources={x: {"A": 1}, y: {"B": 1}})
+    yy = c.compute(y, optimize_graph=False)
     await wait(yy)
 
     assert b.data
@@ -175,8 +180,10 @@ async def test_compute(c, s, a, b):
 async def test_get(c, s, a, b):
     dsk = {"x": (inc, 1), "y": (inc, "x")}
 
-    result = await c.get(dsk, "y", resources={"y": {"A": 1}}, sync=False)
+    result = await c.get(dsk, "y", resources={"A": 1}, sync=False)
     assert result == 3
+    assert "y" in a.data
+    assert not b.data
 
 
 @gen_cluster(
@@ -186,11 +193,12 @@ async def test_get(c, s, a, b):
         ("127.0.0.1", 1, {"resources": {"B": 1}}),
     ],
 )
-async def test_persist_tuple(c, s, a, b):
-    x = delayed(inc)(1)
-    y = delayed(inc)(x)
+async def test_persist_multiple_collections(c, s, a, b):
+    with dask.annotate(resources={"A": 1}):
+        x = delayed(inc)(1)
+        y = delayed(inc)(x)
 
-    xx, yy = c.persist([x, y], resources={(x, y): {"A": 1}})
+    xx, yy = c.persist([x, y], optimize_graph=False)
 
     await wait([xx, yy])
 
@@ -308,11 +316,12 @@ async def test_set_resources(c, s, a):
 async def test_persist_collections(c, s, a, b):
     da = pytest.importorskip("dask.array")
     x = da.arange(10, chunks=(5,))
-    y = x.map_blocks(lambda x: x + 1)
+    with dask.annotate(resources={"A": 1}):
+        y = x.map_blocks(lambda x: x + 1)
     z = y.map_blocks(lambda x: 2 * x)
     w = z.sum()
 
-    ww, yy = c.persist([w, y], resources={tuple(y.__dask_keys__()): {"A": 1}})
+    ww, yy = c.persist([w, y], optimize_graph=False)
 
     await wait([ww, yy])
 
@@ -360,6 +369,7 @@ async def test_full_collections(c, s, a, b):
     assert not b.log
 
 
+@pytest.mark.xfail(reason="atop fusion seemed to break this")
 @pytest.mark.parametrize(
     "optimize_graph",
     [
@@ -382,9 +392,10 @@ async def f(dask_worker):
 
     client.run(f, workers=[a["address"]])
 
-    x = da.random.random(100, chunks=(10,)) + 1
+    with dask.annotate(resources={"A": 1}):
+        x = da.random.random(100, chunks=(10,)) + 1
 
-    x.compute(resources={tuple(x.dask): {"A": 1}}, optimize_graph=optimize_graph)
+    x.compute(optimize_graph=optimize_graph)
 
     def g(dask_worker):
         return len(dask_worker.log)
diff --git a/distributed/tests/test_scheduler.py b/distributed/tests/test_scheduler.py
--- a/distributed/tests/test_scheduler.py
+++ b/distributed/tests/test_scheduler.py
@@ -1039,11 +1039,12 @@ async def test_no_workers_to_memory(c, s):
 
 @gen_cluster(client=True)
 async def test_no_worker_to_memory_restrictions(c, s, a, b):
-    x = delayed(slowinc)(1, delay=0.4)
-    y = delayed(slowinc)(x, delay=0.4)
-    z = delayed(slowinc)(y, delay=0.4)
+    with dask.annotate(workers="alice"):
+        x = delayed(slowinc)(1, delay=0.4)
+        y = delayed(slowinc)(x, delay=0.4)
+        z = delayed(slowinc)(y, delay=0.4)
 
-    yy, zz = c.persist([y, z], workers={(x, y, z): "alice"})
+    yy, zz = c.persist([y, z], optimize_graph=False)
 
     while not s.tasks:
         await asyncio.sleep(0.01)