Mark fast tasks of tasks-based shuffle when dask-expr is enabled

.github/workflows/tests.yaml

@@ -154,7 +154,7 @@ jobs:
           # Increase this value to reset cache if
           # continuous_integration/environment-${{ matrix.environment }}.yaml has not
           # changed. See also same variable in .pre-commit-config.yaml
-          CACHE_NUMBER: 1
+          CACHE_NUMBER: 2
         id: cache
 
       - name: Update environment

distributed/distributed.yaml

@@ -29,6 +29,8 @@ distributed:
     default-task-durations: # How long we expect function names to run ("1h", "1s") (helps for long tasks)
       rechunk-split: 1us
       split-shuffle: 1us
+      split-taskshuffle: 1us
+      split-stage: 1us
     validate: False         # Check scheduler state at every step for debugging
     dashboard:
       status:

distributed/stealing.py

@@ -542,4 +542,8 @@ def _get_thief(
     return min(potential_thieves, key=partial(scheduler.worker_objective, ts))
 
 
-fast_tasks = {"split-shuffle"}
+fast_tasks = {
+    k
+    for k, v in dask.config.get("distributed.scheduler.default-task-durations").items()
+    if parse_timedelta(v) <= 0.001
+}

distributed/tests/test_scheduler.py

@@ -2783,12 +2783,13 @@ async def test_default_task_duration_splits(c, s, a, b):
     await wait(fut)
 
     split_prefix = [pre for pre in s.task_prefixes.keys() if "split" in pre]
-    assert len(split_prefix) == 1
-    split_prefix = split_prefix[0]
-    default_time = parse_timedelta(
-        dask.config.get("distributed.scheduler.default-task-durations")[split_prefix]
-    )
-    assert default_time <= 1e-6
+    # dask-expr enabled: ['split-taskshuffle', 'split-stage']
+    # dask-expr disabled: ['split-shuffle']
+    assert split_prefix
+    default_times = dask.config.get("distributed.scheduler.default-task-durations")
+    for p in split_prefix:
+        default_time = parse_timedelta(default_times[p])
+        assert default_time <= 1e-6
 
 
 @gen_test()

distributed/tests/test_steal.py

@@ -1027,10 +1027,10 @@ async def test_blocklist_shuffle_split(c, s, a, b):
 
     while not s.tasks:
         await asyncio.sleep(0.005)
-    prefixes = set(s.task_prefixes.keys())
+
     from distributed.stealing import fast_tasks
 
-    blocked = fast_tasks & prefixes
+    blocked = fast_tasks & s.task_prefixes.keys()
     assert blocked
     assert any(["split" in prefix for prefix in blocked])