Repository: dask/distributed

Issue:
Unable to provide workers as dict for client.persist()
client.persist() API provided an option to provide workers as 'dict' in order to specify which keys/computation needs to be run on which worker. However in PR #4406 this was removed and I now get TypeError: unhashable type "dict'.

To provide more information on what I'm trying to do:
I have each worker mapped to a GPU in my system. When I run dask shuffle, I need partition 0 to got to GPU 0, partition 1 to GPU 1 and so on. Prior to the change I accomplished this by passing a dict mapping keys to intended workers/GPUs in client.persist. 
1) Is this feature permanently removed? 
2) Is there some other way in dask to accomplish what I'm trying to do?





Discussion:
Thanks for raising an issue @Iroy30. We recently moved specifying keywords like `workers=`, `priority=`, etc. with Dask collections to use [Dask's annotation machinery](https://docs.dask.org/en/latest/api.html#dask.annotate). If you need to specify a specific worker for each partition you can use something along the lines of:

```python
    with dask.annotate(workers=<callable>)):
        client.persist(objects)
```

where `<callable>` should be replaced with a callable that takes a given partition key as input and returns the corresponding worker you want the partition to be `persist`ed on 

cc @ian-r-rose for visibility
cc @pentschev @quasiben (for vis as well)
@jrbourbeau Thanks. I will try this.
Hi, I tried the annotation machinery. This is a small reproducer of the intended workflow:
```
    ddf['src'] = ddf['src']*2
    worker_addresses = Comms.get_workers()
    def key_to_worker(key):
        return tuple([worker_addresses[key[1]]])
    with dask.annotate(workers = key_to_worker):
        pe = client.persist(ddf)
    pa = futures_of(pe)
    wait(pa)
```
where ddf is a dask dataframe and worker_addresses is simply a list of worker addresses when client is initialized. 
The key_to_worker function takes in a key and returns the corresponding worker indexed by the partition number in the key.

This gives me the following error:
```
distributed.protocol.core - CRITICAL - Failed to Serialize
Traceback (most recent call last):
  File "/home/nfs/iroy/anaconda3/envs/env_dev/lib/python3.8/site-packages/distributed/protocol/core.py", line 39, in dumps
    small_header, small_payload = dumps_msgpack(msg, **compress_opts)
  File "/home/nfs/iroy/anaconda3/envs/env_dev/lib/python3.8/site-packages/distributed/protocol/core.py", line 184, in dumps_msgpack
    payload = msgpack.dumps(msg, default=msgpack_encode_default, use_bin_type=True)
  File "/home/nfs/iroy/anaconda3/envs/env_dev/lib/python3.8/site-packages/msgpack/__init__.py", line 35, in packb
    return Packer(**kwargs).pack(o)
  File "msgpack/_packer.pyx", line 292, in msgpack._cmsgpack.Packer.pack
  File "msgpack/_packer.pyx", line 298, in msgpack._cmsgpack.Packer.pack
  File "msgpack/_packer.pyx", line 295, in msgpack._cmsgpack.Packer.pack
  File "msgpack/_packer.pyx", line 264, in msgpack._cmsgpack.Packer._pack
  File "msgpack/_packer.pyx", line 231, in msgpack._cmsgpack.Packer._pack
  File "msgpack/_packer.pyx", line 231, in msgpack._cmsgpack.Packer._pack
  File "msgpack/_packer.pyx", line 289, in msgpack._cmsgpack.Packer._pack
TypeError: can not serialize 'function' object
```
Is this the correct way of using the dask annotation ?
What version of Dask + Distributed are you using?
It looks like this is an issue on the current `dask` and `distributed` `main` branches. Here's a minimal reproducer

```python

import dask
import dask.array as da
from dask.distributed import Client


if __name__ == "__main__":

    with Client() as client:
        worker = next(iter(client.scheduler_info()["workers"].keys()))
        
        def key_to_worker(key):
            return worker
        
        x = da.ones(100, chunks=100)

        with dask.annotate(workers=key_to_worker):
            xp = client.persist(x)
        
        print(xp.compute())
```

The `key_to_worker` function makes its way into [the `update-graph-hlg` message](https://github.com/dask/distributed/blob/42894418a7c63499eeb240a8cbd3c7cc548bea11/distributed/client.py#L2576-L2587) because we pull in annotations set by `dask.annotate`, which in this case is a non-msgpack serializable function, [here](https://github.com/dask/distributed/blob/42894418a7c63499eeb240a8cbd3c7cc548bea11/distributed/client.py#L2572).

@Iroy30 as a temporary workaround you can move the `dask.annotate` context manager around the creation of your Dask collection instead of around the `persist` call and things should work. I.e. instead of 

```python
        x = da.ones(100, chunks=100)

        with dask.annotate(workers=key_to_worker):
            xp = client.persist(x)
```

do

```python
        with dask.annotate(workers=key_to_worker):
            x = da.ones(100, chunks=100)

        xp = client.persist(x)
```
cc @madsbk

Resolving patch:
diff --git a/distributed/client.py b/distributed/client.py
--- a/distributed/client.py
+++ b/distributed/client.py
@@ -2545,14 +2545,10 @@ def _graph_to_futures(
             if actors is not None and actors is not True and actors is not False:
                 actors = list(self._expand_key(actors))
 
-            keyset = set(keys)
-
             # Make sure `dsk` is a high level graph
             if not isinstance(dsk, HighLevelGraph):
                 dsk = HighLevelGraph.from_collections(id(dsk), dsk, dependencies=())
 
-            dsk = dsk.__dask_distributed_pack__(self, keyset)
-
             annotations = {}
             if user_priority:
                 annotations["priority"] = user_priority
@@ -2569,8 +2565,13 @@ def _graph_to_futures(
             if resources:
                 annotations["resources"] = resources
 
+            # Merge global and local annotations
             annotations = merge(dask.config.get("annotations", {}), annotations)
 
+            # Pack the high level graph before sending it to the scheduler
+            keyset = set(keys)
+            dsk = dsk.__dask_distributed_pack__(self, keyset, annotations)
+
             # Create futures before sending graph (helps avoid contention)
             futures = {key: Future(key, self, inform=False) for key in keyset}
             self._send_to_scheduler(
@@ -2581,7 +2582,6 @@ def _graph_to_futures(
                     "priority": priority,
                     "submitting_task": getattr(thread_state, "key", None),
                     "fifo_timeout": fifo_timeout,
-                    "annotations": annotations,
                     "actors": actors,
                 }
             )
diff --git a/distributed/scheduler.py b/distributed/scheduler.py
--- a/distributed/scheduler.py
+++ b/distributed/scheduler.py
@@ -4054,9 +4054,8 @@ def update_graph_hlg(
         user_priority=0,
         actors=None,
         fifo_timeout=0,
-        annotations=None,
     ):
-        unpacked_graph = HighLevelGraph.__dask_distributed_unpack__(hlg, annotations)
+        unpacked_graph = HighLevelGraph.__dask_distributed_unpack__(hlg)
         dsk = unpacked_graph["dsk"]
         dependencies = unpacked_graph["deps"]
         annotations = unpacked_graph["annotations"]
@@ -4296,6 +4295,9 @@ def update_graph(
                     continue
                 ts._host_restrictions = set()
                 ts._worker_restrictions = set()
+                # Make sure `v` is a collection and not a single worker name / address
+                if not isinstance(v, (list, tuple, set)):
+                    v = [v]
                 for w in v:
                     try:
                         w = self.coerce_address(w)


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
@@ -4143,6 +4143,26 @@ async def test_persist_workers_annotate(e, s, a, b, c):
     assert s.loose_restrictions == {total2.key} | {v.key for v in L2}
 
 
+@gen_cluster(nthreads=[("127.0.0.1", 1)] * 3, client=True)
+async def test_persist_workers_annotate2(e, s, a, b, c):
+    def key_to_worker(key):
+        return a.address
+
+    L1 = [delayed(inc)(i) for i in range(4)]
+    for x in L1:
+        assert all(layer.annotations is None for layer in x.dask.layers.values())
+
+    with dask.annotate(workers=key_to_worker):
+        out = e.persist(L1, optimize_graph=False)
+        await wait(out)
+
+    for x in L1:
+        assert all(layer.annotations is None for layer in x.dask.layers.values())
+
+    for v in L1:
+        assert s.worker_restrictions[v.key] == {a.address}
+
+
 @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):