Refactor resource restriction handling in WorkerState (#6672)

distributed/scheduler.py

@@ -1860,7 +1860,7 @@ def transition_waiting_processing(self, key, stimulus_id):
             self._set_duration_estimate(ts, ws)
             ts.processing_on = ws
             ts.state = "processing"
-            self.consume_resources(ts, ws)
+            self.acquire_resources(ts, ws)
             self.check_idle_saturated(ws)
             self.n_tasks += 1
 
@@ -2675,7 +2675,7 @@ def valid_workers(self, ts: TaskState) -> set[WorkerState] | None:
 
         return s
 
-    def consume_resources(self, ts: TaskState, ws: WorkerState):
+    def acquire_resources(self, ts: TaskState, ws: WorkerState):
         for r, required in ts.resource_restrictions.items():
             ws.used_resources[r] += required
 

distributed/tests/test_resources.py

@@ -70,8 +70,8 @@ async def test_submit_many_non_overlapping_2(c, s, a, b):
         assert b.state.executing_count <= 1
 
     await wait(futures)
-    assert a.total_resources == a.state.available_resources
-    assert b.total_resources == b.state.available_resources
+    assert a.state.total_resources == a.state.available_resources
+    assert b.state.total_resources == b.state.available_resources
 
 
 @gen_cluster(
@@ -232,7 +232,7 @@ async def test_minimum_resource(c, s, a):
         assert a.state.executing_count <= 1
 
     await wait(futures)
-    assert a.total_resources == a.state.available_resources
+    assert a.state.total_resources == a.state.available_resources
 
 
 @gen_cluster(client=True, nthreads=[("127.0.0.1", 2, {"resources": {"A": 1}})])
@@ -271,7 +271,7 @@ async def test_balance_resources(c, s, a, b):
 @gen_cluster(client=True, nthreads=[("127.0.0.1", 2)])
 async def test_set_resources(c, s, a):
     await a.set_resources(A=2)
-    assert a.total_resources["A"] == 2
+    assert a.state.total_resources["A"] == 2
     assert a.state.available_resources["A"] == 2
     assert s.workers[a.address].resources == {"A": 2}
     lock = Lock()
@@ -281,7 +281,7 @@ async def test_set_resources(c, s, a):
             await asyncio.sleep(0.01)
 
     await a.set_resources(A=3)
-    assert a.total_resources["A"] == 3
+    assert a.state.total_resources["A"] == 3
     assert a.state.available_resources["A"] == 2
     assert s.workers[a.address].resources == {"A": 3}
 

distributed/utils_test.py

@@ -2450,6 +2450,7 @@ def ws_with_running_task(ws, request):
     The task may or may not raise secede(); the tests using this fixture runs twice.
     """
     ws.available_resources = {"R": 1}
+    ws.total_resources = {"R": 1}
     instructions = ws.handle_stimulus(
         ComputeTaskEvent.dummy(
             key="x", resource_restrictions={"R": 1}, stimulus_id="compute"

distributed/worker.py

@@ -406,7 +406,6 @@ class Worker(BaseWorker, ServerNode):
     _dashboard_address: str | None
     _dashboard: bool
     _http_prefix: str
-    total_resources: dict[str, float]
     death_timeout: float | None
     lifetime: float | None
     lifetime_stagger: float | None
@@ -628,7 +627,6 @@ def __init__(
         if resources is None:
             resources = dask.config.get("distributed.worker.resources")
             assert isinstance(resources, dict)
-        self.total_resources = resources.copy()
 
         self.death_timeout = parse_timedelta(death_timeout)
 
@@ -754,7 +752,7 @@ def __init__(
             data=self.memory_manager.data,
             threads=self.threads,
             plugins=self.plugins,
-            resources=self.total_resources,
+            resources=resources,
             total_out_connections=total_out_connections,
             validate=validate,
             transition_counter_max=transition_counter_max,
@@ -877,6 +875,7 @@ def data(self) -> MutableMapping[str, Any]:
     tasks = DeprecatedWorkerStateAttribute()
     target_message_size = DeprecatedWorkerStateAttribute()
     total_out_connections = DeprecatedWorkerStateAttribute()
+    total_resources = DeprecatedWorkerStateAttribute()
     transition_counter = DeprecatedWorkerStateAttribute()
     transition_counter_max = DeprecatedWorkerStateAttribute()
     validate = DeprecatedWorkerStateAttribute()
@@ -1100,7 +1099,7 @@ async def _register_with_scheduler(self) -> None:
                         },
                         types={k: typename(v) for k, v in self.data.items()},
                         now=time(),
-                        resources=self.total_resources,
+                        resources=self.state.total_resources,
                         memory_limit=self.memory_manager.memory_limit,
                         local_directory=self.local_directory,
                         services=self.service_ports,
@@ -1752,17 +1751,19 @@ def update_data(
         )
         return {"nbytes": {k: sizeof(v) for k, v in data.items()}, "status": "OK"}
 
-    async def set_resources(self, **resources) -> None:
+    async def set_resources(self, **resources: float) -> None:
         for r, quantity in resources.items():
-            if r in self.total_resources:
-                self.state.available_resources[r] += quantity - self.total_resources[r]
+            if r in self.state.total_resources:
+                self.state.available_resources[r] += (
+                    quantity - self.state.total_resources[r]
+                )
             else:
                 self.state.available_resources[r] = quantity
-            self.total_resources[r] = quantity
+            self.state.total_resources[r] = quantity
 
         await retry_operation(
             self.scheduler.set_resources,
-            resources=self.total_resources,
+            resources=self.state.total_resources,
             worker=self.contact_address,
         )
 

distributed/worker_state_machine.py

@@ -82,7 +82,6 @@
 }
 READY: set[TaskStateState] = {"ready", "constrained"}
 
-
 NO_VALUE = "--no-value-sentinel--"
 
 
@@ -1027,8 +1026,12 @@ class WorkerState:
     #: determining a last-in-first-out order between them.
     generation: int
 
+    #: ``{resource name: amount}``. Total resources available for task execution.
+    #: See :doc: `resources`.
+    total_resources: dict[str, float]
+
     #: ``{resource name: amount}``. Current resources that aren't being currently
-    #: consumed by task execution. Always less or equal to ``Worker.total_resources``.
+    #: consumed by task execution. Always less or equal to :attr:`total_resources`.
     #: See :doc:`resources`.
     available_resources: dict[str, float]
 
@@ -1102,7 +1105,8 @@ def __init__(
         self.data = data if data is not None else {}
         self.threads = threads if threads is not None else {}
         self.plugins = plugins if plugins is not None else {}
-        self.available_resources = dict(resources) if resources is not None else {}
+        self.total_resources = dict(resources) if resources is not None else {}
+        self.available_resources = self.total_resources.copy()
 
         self.validate = validate
         self.tasks = {}
@@ -1445,15 +1449,11 @@ def _ensure_computing(self) -> RecsInstrs:
             if ts in recs:
                 continue
 
-            if any(
-                self.available_resources[resource] < needed
-                for resource, needed in ts.resource_restrictions.items()
-            ):
+            if not self._resource_restrictions_satisfied(ts):
                 break
 
             self.constrained.popleft()
-            for resource, needed in ts.resource_restrictions.items():
-                self.available_resources[resource] -= needed
+            self._acquire_resources(ts)
 
             recs[ts] = "executing"
             self.executing.add(ts)
@@ -1734,8 +1734,7 @@ def _transition_executing_rescheduled(
             # Reschedule(), which is "cancelled"
             assert ts.state in ("executing", "long-running"), ts
 
-        for resource, quantity in ts.resource_restrictions.items():
-            self.available_resources[resource] += quantity
+        self._release_resources(ts)
         self.executing.discard(ts)
 
         return merge_recs_instructions(
@@ -1831,8 +1830,7 @@ def _transition_executing_error(
         *,
         stimulus_id: str,
     ) -> RecsInstrs:
-        for resource, quantity in ts.resource_restrictions.items():
-            self.available_resources[resource] += quantity
+        self._release_resources(ts)
         self.executing.discard(ts)
 
         return merge_recs_instructions(
@@ -1977,9 +1975,7 @@ def _transition_cancelled_released(
         self.executing.discard(ts)
         self.in_flight_tasks.discard(ts)
 
-        for resource, quantity in ts.resource_restrictions.items():
-            self.available_resources[resource] += quantity
-
+        self._release_resources(ts)
         return self._transition_generic_released(ts, stimulus_id=stimulus_id)
 
     def _transition_executing_released(
@@ -2006,10 +2002,7 @@ def _transition_generic_memory(
                 f"Tried to transition task {ts} to `memory` without data available"
             )
 
-        if ts.resource_restrictions is not None:
-            for resource, quantity in ts.resource_restrictions.items():
-                self.available_resources[resource] += quantity
-
+        self._release_resources(ts)
         self.executing.discard(ts)
         self.in_flight_tasks.discard(ts)
         ts.coming_from = None
@@ -2351,6 +2344,20 @@ def _transition(
         )
         return recs, instructions
 
+    def _resource_restrictions_satisfied(self, ts: TaskState) -> bool:
+        return all(
+            self.available_resources[resource] >= needed
+            for resource, needed in ts.resource_restrictions.items()
+        )
+
+    def _acquire_resources(self, ts: TaskState) -> None:
+        for resource, needed in ts.resource_restrictions.items():
+            self.available_resources[resource] -= needed
+
+    def _release_resources(self, ts: TaskState) -> None:
+        for resource, needed in ts.resource_restrictions.items():
+            self.available_resources[resource] += needed
+
     def _transitions(self, recommendations: Recs, *, stimulus_id: str) -> Instructions:
         """Process transitions until none are left