I'm curious why these were removed.

I was generally trying to avoid any expansion of keys/collections on the client, instead preferring to only allow a worker-or-iterable-of-workers. With that API restriction, the unpacking/expanding logic on the scheduler is much simpler. So I view this as philosophically aligned with removing the ability to expand { collection: worker } on the client.

I was generally trying to avoid any expansion of keys/collections on the client, instead preferring to only allow a worker-or-iterable-of-workers. With that API restriction, the unpacking/expanding logic on the scheduler is much simpler. So I view this as philosophically aligned with removing the ability to expand { collection: worker } on the client.

I understand this change provides backwards compatibility for the existing per-key Client.{submit, persist} kwarg taxonomy (retries, priority, workers)?

The functionality in _materialized_layer_unpack mirrors Layer.__dask_distributed_unpack__ and it all it's subclasses. They also need to expand and unpack annotations.

I wonder if it would be better to place this functionality in Layer.expand_annotations so that it automatically is applied through the Layer class hierarchy? It would require a sister dask PR though.

Thanks for pointing this out. This is actually not for backwards compatibility of per-key/collection annotations, but instead is needed to handle the case where different HLG Layers have different retries/priority/workers.

For example, if we are unpacking a two-layer HLG, and Layer A has retries:2 while Layer B has retries:4, we need to make sure that the unpacking keeps both of those annotations. The current implementation has a shallow merge, which makes later-unpacked layers overwrite the annotations of previously-unpacked layers. To make the above case work, we need a bit deeper of a merge.

For the same reason, I also think this logic needs to be one step above Layer.expand_annotations.

Yes, you're right. Thanks for correcting my understanding.

new_annotations[k] = {**annotations[k], **v}

may be faster than tlz.merge now that dask supports >= Python 3.6

cc'ing @madsbk

The functionality in _materialized_layer_unpack mirrors Layer.__dask_distributed_unpack__ and it all it's subclasses. They also need to expand and unpack annotations.

I think this is still an issue, we need to apply the same changes to Blockwise.__dask_distributed_unpack__ and SimpleShuffleLayer.__dask_distributed_unpack__.

Maybe it makes more sense to redesign the annotations by removing the annotations argument from Layer.__dask_distributed_unpack__ and then have highlevelgraph_unpack() expand and unpack annotations just after unpack_func(layer["state"], dsk, deps).
This will simplify all the __dask_distributed_unpack__ implementations and remove duplicate code. The only downside is that it limits Layers to do a specialized annotation expansion, which no one is doing currently. Unless I am missing something, I think that is a great trade-off.

Thanks for the detailed discussion @sjperkins and @madsbk. I've tried to do something like a compromise solution in dask/dask#7102, and I'd love to hear your thoughts on it. It takes a somewhat minimal approach to API changes (so __dask_distributed_unpack__() still takes the annotations), and avoids further imports of distributed there. Overall the strategy I'm trying out is

1. Move the default implementation of layer unpacking into Layer.
2. Move the logic of doing a deeper merge on annotations into a new Layer.merge_annotations, similar to Layer.expand_annotations. I am not overly attached to this division of responsibility, so if there are suggestions for how to make it cleaner, I'd be happy to hear them.
3. Ensure that Layer, Blockwise, and ShuffleLayer all unpack annotations properly (previously the latter two failed).

LGTM @ian-r-rose but I think @sjperkins has a point, we should unify everything in Dask and change the API such that layers can delegate the annotation unpacking to Layer.__dask_distributed_unpack__().

I suggest that we split this work into multiple PRs. Let's rollback this PR to the state that didn't require any significant change to the API and get it merged. It greatly simplifies the code and is a win in all cases.
A followup PR can then modify the API based on the design thoughts we have been discussing here. I will be happy to initiate the followup PR, I think I understand @sjperkins's design ideas and concerns now :)

Thanks for the discussion @madsbk and @ian-r-rose . I agree with moving Layer API changes forward to a new PR, so don't let my concerns hold this one up. LGTM

@madsbk I'd be happy to have this be across multiple PRs. Can you specify what API changes you would want to roll back? This now depends on dask/dask#7102, but that is somewhat more of a bugfix than a major API change.

I could bring the default implementation of Layer.__dask_distributed_unpack__() back here and then xfail the test that catches dask/dask#7102 , if that's what you mean.

Sorry my bad, I thought you changed the API in the last commit. I am totally fine with this PR+dask/dask#7102

To me this implies that Layer annotations always take precedence over the existing Client.{submit, persist} kwarg taxonomy. I think this makes sense in terms of backwards compatibility (specific annotations override global annotations)

The functionality in _materialized_layer_unpack mirrors Layer.__dask_distributed_unpack__ and it all it's subclasses. They also need to expand and unpack annotations.

I think this is still an issue, we need to apply the same changes to Blockwise.__dask_distributed_unpack__ and SimpleShuffleLayer.__dask_distributed_unpack__.

Maybe it makes more sense to redesign the annotations by removing the annotations argument from Layer.__dask_distributed_unpack__ and then have highlevelgraph_unpack() expand and unpack annotations just after unpack_func(layer["state"], dsk, deps).
This will simplify all the __dask_distributed_unpack__ implementations and remove duplicate code. The only downside is that it limits Layers to do a specialized annotation expansion, which no one is doing currently. Unless I am missing something, I think that is a great trade-off.

Since the workers= annotation also supports worker addresses in addition to worker hostnames, I think we can update this to remove the pytest.mark.skipif while testing the same behavior:

diff --git a/distributed/tests/test_client.py b/distributed/tests/test_client.py
index eb1c362d..0a271596 100644
--- a/distributed/tests/test_client.py
+++ b/distributed/tests/test_client.py
@@ -957,15 +957,12 @@ async def test_restrictions_get(c, s, a, b):
     assert len(b.data) == 0


-@pytest.mark.skipif(
-    not sys.platform.startswith("linux"), reason="Need 127.0.0.2 to mean localhost"
-)
-@gen_cluster([("127.0.0.1", 1), ("127.0.0.2", 2)], client=True)
+@gen_cluster(client=True)
 async def test_restrictions_get_annotate(c, s, a, b):
     x = 1
-    with dask.annotate(workers=a.ip):
+    with dask.annotate(workers=a.address):
         y = delayed(inc)(x)
-    with dask.annotate(workers=b.ip):
+    with dask.annotate(workers=b.address):
         z = delayed(inc)(y)

     futures = c.get(z.__dask_graph__(), [y.key, z.key], sync=False)

This is particularly nice given our current linux testing situation on Travis

Just for my own understanding, it looks like test_persist_workers and test_compute_workers have been updated to specifically not run tasks on a worker, which is earlier to specify with the workers= keyword, instead of running certain collections on certain workers. We're now covering the certain collections on certain workers functionality with new test_persist_workers_annotate and test_compute_workers_annotate tests. Is that correct?

I've updated them to persist/compute on a given set of workers, but it is no longer possible to delegate certain tasks/keys to a specific worker in a persist/compute call. In order to do that, you now need to use the annotation framework.

We're now covering the certain collections on certain workers functionality with new test_persist_workers_annotate and test_compute_workers_annotate tests. Is that correct?

Correct