Do not fully materialize global annotations by type (#9035)

Author: fjetter

distributed/scheduler.py

@@ -4687,7 +4687,6 @@ def _create_taskstate_from_graph(
             # FIXME: This is kind of inconsistent since it only includes global
             # annotations.
             computation.annotations.update(global_annotations)
-        del global_annotations
         (
             runnable,
             touched_tasks,
@@ -4709,6 +4708,7 @@ def _create_taskstate_from_graph(
         keys_with_annotations = self._apply_annotations(
             tasks=new_tasks,
             annotations_by_type=annotations_by_type,
+            global_annotations=global_annotations,
         )
 
         self._set_priorities(
@@ -4872,7 +4872,6 @@ async def update_graph(
             ) = await offload(
                 _materialize_graph,
                 expr=expr,
-                global_annotations=annotations or {},
                 validate=self.validate,
             )
 
@@ -4927,8 +4926,6 @@ async def update_graph(
                 code=code,
                 span_metadata=span_metadata,
                 annotations_by_type=annotations_by_type,
-                # FIXME: This is just used to attach to Computation
-                # objects. This should be removed
                 global_annotations=annotations,
                 start=start,
                 stimulus_id=stimulus_id,
@@ -5085,6 +5082,7 @@ def _apply_annotations(
         self,
         tasks: Iterable[TaskState],
         annotations_by_type: dict[str, dict[Key, Any]],
+        global_annotations: dict[str, Any] | None = None,
     ) -> set[Key]:
         """Apply the provided annotations to the provided `TaskState` objects.
 
@@ -5104,13 +5102,18 @@ def _apply_annotations(
         keys_with_annotations
         """
         keys_with_annotations: set[Key] = set()
-        if not annotations_by_type:
+        if not annotations_by_type and not global_annotations:
             return keys_with_annotations
 
         for ts in tasks:
             key = ts.key
 
             ts_annotations = {}
+            if global_annotations:
+                for annot, value in global_annotations.items():
+                    if callable(value):
+                        value = value(ts.key)
+                    ts_annotations[annot] = value
             for annot, key_value in annotations_by_type.items():
                 if (value := key_value.get(key)) is not None:
                     ts_annotations[annot] = value
@@ -9429,18 +9432,14 @@ def transition(
 
 def _materialize_graph(
     expr: Expr,
-    global_annotations: dict[str, Any],
     validate: bool,
 ) -> tuple[dict[Key, T_runspec], dict[Key, set[Key]], dict[str, dict[Key, Any]]]:
     dsk: dict = expr.__dask_graph__()
     if validate:
         for k in dsk:
             validate_key(k)
     annotations_by_type: defaultdict[str, dict[Key, Any]] = defaultdict(dict)
-    for annotations_type, value in global_annotations.items():
-        annotations_by_type[annotations_type].update(
-            {k: (value(k) if callable(value) else value) for k in dsk}
-        )
+
     for annotations_type, value in expr.__dask_annotations__().items():
         annotations_by_type[annotations_type].update(value)