Support `task-group-key` and `task-prefix-key` annotations on the scheduler
- Dominant language
- Python
- Stars
- 1.7k
- Forks
- 778
- Avg merge
- 2h 50m
- Merged PRs (30d)
- 3
Description
Allowing tasks to specify a custom TaskGroup and TaskPrefix key—via annotations—would give us an easy path to work around the fact that Blockwise fusion produces task names that aren't meaningful to users, and can be confusing: https://github.com/dask/dask/issues/8635.
In [`update_graph`](https://github.com/dask/distributed/blob/9a4e0e2f097aa97a373fdd9b7c76f8d79ec746a9/distributed/scheduler.py#L4582):
1. If a `task-group-key` annotation is set for a task, support passing it in as the group key [instead of calling `key_split_group` on the task name](https://github.com/dask/distributed/blob/9a4e0e2f097aa97a373fdd9b7c76f8d79ec746a9/distributed/scheduler.py#L1517).
1. If a `task-prefix-key` annotation is set for a task, use its value as the prefix key [instead of calling `key_split` on the task name](https://github.com/dask/distributed/blob/9a4e0e2f097aa97a373fdd9b7c76f8d79ec746a9/distributed/scheduler.py#L2206).
1. (Maybe?) if only `task-group-key` is given, but not `task-prefix-key`, calculate the prefix key from `task-group-key` (this would require a new function like `key_split_from_group`). Not really necessary, just nice to not have to pass the semi-redundant annotation.
In `update_graph`, if the TaskState is already in memory, the `task-group-key` and `task-prefix-key` annotations can be ignored, whether or not they match the group/prefix already associated with that task.
Once this is implemented, then we can have `optimize_blockwise` set these annotations accordingly.
Contributor guide
Assessment
This issue has not been assessed yet.