focus on task resources

rjzamora · Aug 2, 2024 · e494e88 · e494e88
1 parent 308111d
commit e494e88
Show file tree

Hide file tree

Showing 2 changed files with 9 additions and 7 deletions.
diff --git a/distributed/client.py b/distributed/client.py
@@ -3313,7 +3313,7 @@ def _graph_to_futures(
         retries=None,
         fifo_timeout=0,
         actors=None,
-        annotations_by_type=None,
+        task_resources=None,
     ):
         with self._refcount_lock:
             if actors is not None and actors is not True and actors is not False:
@@ -3381,7 +3381,7 @@ def _graph_to_futures(
                     "actors": actors,
                     "code": ToPickle(computations),
                     "annotations": ToPickle(annotations),
-                    "annotations_by_type": ToPickle(annotations_by_type or {}),
+                    "task_resources": ToPickle(task_resources),
                     "span_metadata": ToPickle(span_metadata),
                 }
             )
@@ -3471,7 +3471,7 @@ def get(
             user_priority=priority,
             actors=actors,
             span_metadata=SpanMetadata(collections=[{"type": "low-level-graph"}]),
-            annotations_by_type=kwargs.get("annotations_by_type", None),
+            task_resources=kwargs.get("task_resources"),
         )
         packed = pack_data(keys, futures)
         if sync:

diff --git a/distributed/scheduler.py b/distributed/scheduler.py
@@ -4753,7 +4753,7 @@ async def update_graph(
         fifo_timeout: float = 0.0,
         code: tuple[SourceCode, ...] = (),
         annotations: dict | None = None,
-        annotations_by_type: dict | None = None,
+        task_resources: dict | None = None,
         stimulus_id: str | None = None,
     ) -> None:
         start = time()
@@ -4772,15 +4772,17 @@ async def update_graph(
             (
                 dsk,
                 dependencies,
-                _annotations_by_type,
+                annotations_by_type,
             ) = await offload(
                 _materialize_graph,
                 graph=graph,
                 global_annotations=annotations or {},
             )
             del graph
-            annotations_by_type = _annotations_by_type.update(annotations_by_type or {})
-            del _annotations_by_type
+
+            if task_resources:
+                annotations_by_type["resources"].update(task_resources)
+
             if not internal_priority:
                 # Removing all non-local keys before calling order()
                 dsk_keys = set(