dask / dask/distributed

Support `task-group-key` and `task-prefix-key` annotations on the scheduler

Open
#5,742 2 comments 0 reactions 0 assignees View on GitHub
enhancement good second issue
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

Open the contributing guide

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.