diff --git a/distributed/tests/test_worker_client.py b/distributed/tests/test_worker_client.py index 0e854e2e09..cded290147 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,43 @@ def f(): assert result == 1 +@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() + + @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): + 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 + + @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 fd0060bb2e..9364e050e1 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,14 +2179,19 @@ 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 + # 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(**annotations) + if annotations else contextlib.nullcontext() ) - span_ctx.__enter__() + annotations_ctx.__enter__() run_spec = ts.run_spec try: ts.start_time = time() @@ -2219,7 +2239,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"]