Dynamic annotations

[madsbk]

As part of the discussion of Heterogeneous Computing Design, this PR implements dynamic annotations inspired by @mrocklin's option 4 in #4656.
The idea is that a user can specify a function that updates the annotations and restrictions of a task after its execution:

1. The worker runs the user specified annotations functions on the task output.
2. The worker sends the updated attributes back to the scheduler using a new "annotate-task" message.
3. The scheduler updates the attributes of the task and all tasks depending on it.

The following in an example of an dynamic annotation function that sets the worker executor to "gpu" when the task output is a CUDA device object. Together with a default GPU executor (#5084), this will make sure that all tasks preceding the first CUDA task will use the GPU executor.

def set_gpu_executor(ts: TaskState, value: object) -> bool:
    if dask_cuda.is_device_object(value):
        ts.annotations["executor"] = "gpu"
        return True
    return False

Notice, this will not annotate the first use of a GPU task, only its dependent tasks are annotated automatically. The user can of cause still annotate tasks manually.

• Tests added / passed
• Passes black distributed / flake8 distributed / isort distributed

[mrocklin]

I recommend making restrictions a separate route. I don't think that there is any reason for annotations and restrictions to be intertwined.

However, I do think that it would be useful to support a variety of different kinds of restrictions. You should probably roll this into the Scheduler.set_restrictions method. I think that it has a worker= keyword today as a hint that host= and resource= should come later.

[madsbk]

Good point.
A related question, what is the different between regular handlers and stream handlers?
Can I use self.batched_stream.send for a "op": "set_restrictions" message when set_restrictions is in Scheduler.handlers?

[mrocklin]

Stream handlers are used with BatchedSend objects. They are strictly fire and forget. They're also very cheap to send. We accumulate several small messages in a list and then send off that list periodically (with a very short period)

The handlers are what receive the await rpc.foo(...) commands. They wait for and collect a response.

If you need a response then use await rpc.foo(...). If you want to send lots of little small messages without really affecting overhead then try to use a stream handler. You can certainly put a method in both stream handlers and handlers, although if the receiving end needs a response then this might require some cleverness.

Does that provide enough information?

[madsbk]

Yes, thanks!

[mrocklin]

I'm curious about the motivation behind this keyword. Can you expand here?

[madsbk]

I am considering if an annotator function should be able to enable/disable annotations of task dependents. Alternativevly, we can always apply annotations to task dependents.

[mrocklin]

I agree that it's possibly in scope. I'm curious if there are active use cases for it today that you have in mind. If so, I would be curious to hear them (for example, maybe you need the results of gpu tasks are likely to also be gpu tasks). If not, then I would recommend waiting on this until we know more.

My gut reaction is not to do this until we have a motivating use case.

[madsbk]

for example, maybe you need the results of gpu tasks are likely to also be gpu tasks

Yes, to implement automatic use of a GPU ThreadPoolExecutor (#5084). If the initial array or dataframe creation is annotated with the GPU executor, all the following operations are as well.

This will also be very helpful when mixing data types. By tracking the GPU use of each individual partition, we can utilize GPU and CPU workers simultaneously.

[mrocklin]

Some small comments. In general this looks fine to me though.

[mrocklin]

The fact that you had to use get/compute here is intersting. There are interesting challenges with doing this with futures / dependencies that might change. I still think that this is a useful feature to have, but it seems like this approach might not be comprehensive if we want to solve things across update_graph calls. Agree or disagree?

[madsbk]

Agree! I hadn't thought about this issue before today. As you say, it is still useful but I think we should put this PR on hold for a bit and see if we can come up with a better approach.

[mrocklin]

Do you have thoughts on better names here?

[madsbk]

Heh, yeah I can do that :)