Rollback changes to set_resources and fix fixture

distributed/utils_test.py

@@ -2449,7 +2449,8 @@ def ws_with_running_task(ws, request):
 
     The task may or may not raise secede(); the tests using this fixture runs twice.
     """
-    ws.set_resources(R=1)
+    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

@@ -1752,7 +1752,14 @@ def update_data(
         return {"nbytes": {k: sizeof(v) for k, v in data.items()}, "status": "OK"}
 
     async def set_resources(self, **resources: float) -> None:
-        self.state.set_resources(**resources)
+        for r, quantity in resources.items():
+            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.state.total_resources[r] = quantity
 
         await retry_operation(
             self.scheduler.set_resources,

distributed/worker_state_machine.py

@@ -3116,14 +3116,6 @@ def validate_state(self) -> None:
         for tss in self.data_needed.values():
             assert len({ts.key for ts in tss}) == len(tss)
 
-    def set_resources(self, **resources: float) -> None:
-        for r, quantity in resources.items():
-            if r in self.total_resources:
-                self.available_resources[r] += quantity - self.total_resources[r]
-            else:
-                self.available_resources[r] = quantity
-            self.total_resources[r] = quantity
-
 
 class BaseWorker(abc.ABC):
     """Wrapper around the :class:`WorkerState` that implements instructions handling.