From 9ce93727b79e096e3547fd5bed7ea68b2f061f2d Mon Sep 17 00:00:00 2001 From: rohitsalla <87137863+rohitsalla@users.noreply.github.com> Date: Tue, 4 Aug 2026 11:39:38 -0400 Subject: [PATCH 1/2] Propagate task annotations during worker execution --- distributed/tests/test_worker_client.py | 24 ++++++++++++++++++++++++ distributed/worker.py | 13 ++++++------- 2 files changed, 30 insertions(+), 7 deletions(-) diff --git a/distributed/tests/test_worker_client.py b/distributed/tests/test_worker_client.py index 0e854e2e098..bcd6c2ec596 100644 --- a/distributed/tests/test_worker_client.py +++ b/distributed/tests/test_worker_client.py @@ -16,6 +16,7 @@ as_completed, get_client, get_worker, + span, wait, worker_client, ) @@ -329,6 +330,29 @@ def f(): assert result == 1 +@gen_cluster(client=True) +async def test_nested_compute_inherits_task_annotations(c, s, a, b): + @dask.delayed + def read_annotations(): + return dask.get_annotations() + + @dask.delayed + async def read_annotations_async(): + return dask.get_annotations() + + def outer(value): + with span(str(value)), dask.annotate(custom=value): + expected = dask.get_annotations() + actual_sync = read_annotations().compute(optimize_graph=False) + actual_async = read_annotations_async().compute(optimize_graph=False) + return expected, actual_sync, actual_async + + for value in (1, 2): + expected, actual_sync, actual_async = await c.submit(outer, value, pure=False) + assert actual_sync == expected + assert actual_async == expected + + @gen_cluster(client=True) async def test_worker_client_rejoins(c, s, a, b): def f(): diff --git a/distributed/worker.py b/distributed/worker.py index fd0060bb2ea..c99a2113521 100644 --- a/distributed/worker.py +++ b/distributed/worker.py @@ -2164,14 +2164,13 @@ async def execute(self, key: Key, *, stimulus_id: str) -> StateMachineEvent: ) self.active_keys.add(key) - # Propagate span (see distributed.spans). This is useful when spawning - # more tasks using worker_client() and for logging. - span_ctx = ( - dask.annotate(span=ts.annotations["span"]) - if "span" in ts.annotations + # Make resolved task annotations visible to user code and nested tasks. + annotations_ctx = ( + dask.annotate(**ts.annotations) + if ts.annotations else contextlib.nullcontext() ) - span_ctx.__enter__() + annotations_ctx.__enter__() run_spec = ts.run_spec try: ts.start_time = time() @@ -2219,7 +2218,7 @@ async def execute(self, key: Key, *, stimulus_id: str) -> StateMachineEvent: ) finally: self.active_keys.discard(key) - span_ctx.__exit__(None, None, None) + annotations_ctx.__exit__(None, None, None) self.threads[key] = result["thread"] From 696800cc9151a13764673fe71345f9e37e74bf7b Mon Sep 17 00:00:00 2001 From: rohitsalla <87137863+rohitsalla@users.noreply.github.com> Date: Tue, 4 Aug 2026 22:54:55 -0400 Subject: [PATCH 2/2] Avoid propagating task control annotations --- distributed/tests/test_worker_client.py | 14 +++++++++++++ distributed/worker.py | 27 ++++++++++++++++++++++--- 2 files changed, 38 insertions(+), 3 deletions(-) diff --git a/distributed/tests/test_worker_client.py b/distributed/tests/test_worker_client.py index bcd6c2ec596..cded2901471 100644 --- a/distributed/tests/test_worker_client.py +++ b/distributed/tests/test_worker_client.py @@ -332,6 +332,9 @@ def f(): @gen_cluster(client=True) async def test_nested_compute_inherits_task_annotations(c, s, a, b): + def read_task_annotations(): + return dask.get_annotations() + @dask.delayed def read_annotations(): return dask.get_annotations() @@ -348,6 +351,17 @@ def outer(value): return expected, actual_sync, actual_async for value in (1, 2): + with dask.annotate(custom=value): + task_annotations = await c.submit( + read_task_annotations, + pure=False, + workers=[a.address], + allow_other_workers=True, + priority=10, + retries=1, + ) + assert task_annotations == {"custom": value} + expected, actual_sync, actual_async = await c.submit(outer, value, pure=False) assert actual_sync == expected assert actual_async == expected diff --git a/distributed/worker.py b/distributed/worker.py index c99a2113521..9364e050e13 100644 --- a/distributed/worker.py +++ b/distributed/worker.py @@ -154,6 +154,21 @@ logger = logging.getLogger(__name__) + +_TASK_CONTROL_ANNOTATIONS = frozenset( + { + "allow_other_workers", + "executor", + "loose_restrictions", + "priority", + "resources", + "restrictions", + "retries", + "shuffle_original_restrictions", + "workers", + } +) + LOG_PDB = dask.config.get("distributed.admin.pdb-on-err") DEFAULT_EXTENSIONS: dict[str, type] = { @@ -2164,10 +2179,16 @@ async def execute(self, key: Key, *, stimulus_id: str) -> StateMachineEvent: ) self.active_keys.add(key) - # Make resolved task annotations visible to user code and nested tasks. + # Propagate user annotations without applying the parent task's + # scheduling and execution controls to nested tasks. + annotations = { + key: value + for key, value in ts.annotations.items() + if key not in _TASK_CONTROL_ANNOTATIONS + } annotations_ctx = ( - dask.annotate(**ts.annotations) - if ts.annotations + dask.annotate(**annotations) + if annotations else contextlib.nullcontext() ) annotations_ctx.__enter__()