[unsupervised AI] fix: propagate task annotations during worker execution

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():

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"]