Ensure scheduler updates task and worker states after successful worker data deletion
@@ -2952,6 +2952,30 @@ class Scheduler(ServerNode):
)
return d[worker]
+ async def _delete_worker_data(self, worker_address, keys):
+ """ Delete data from a worker and update the corresponding worker/task states
+
+ Parameters
+ ----------
+ worker_address: str
+ Worker address to delete keys from
+ keys: List[str]
+ List of keys to delete on the specified worker
+ """
+ await retry_operation(
+ self.rpc(addr=worker_address).delete_data, keys=list(keys), report=False,
+ )
+
+ ws = self.workers[worker_address]
+ tasks = {self.tasks[key] for key in keys}
+ ws.has_what -= tasks
+ for ts in tasks:
+ ts.who_has.remove(ws)
+ ws.nbytes -= ts.get_nbytes()
+ self.log_event(
+ ws.address, {"action": "remove-worker-data", "keys": keys},
+ )
+
async def rebalance(self, comm=None, keys=None, workers=None):
""" Rebalance keys so that each worker stores roughly equal bytes
@@ -3068,18 +3092,8 @@ class Scheduler(ServerNode):
)
await asyncio.gather(
- *(
- retry_operation(
- self.rpc(addr=r).delete_data, keys=v, report=False
+ *(self._delete_worker_data(r, v) for r, v in to_senders.items())
)
- for r, v in to_senders.items()
- )
- )
-
- for sender, recipient, ts in msgs:
- ts.who_has.remove(sender)
- sender.has_what.remove(ts)
- sender.nbytes -= ts.get_nbytes()
return {"status": "OK"}
@@ -3142,28 +3156,11 @@ class Scheduler(ServerNode):
await asyncio.gather(
*(
- retry_operation(
- self.rpc(addr=ws.address).delete_data,
- keys=[ts.key for ts in tasks],
- report=False,
- )
+ self._delete_worker_data(ws.address, [t.key for t in tasks])
for ws, tasks in del_worker_tasks.items()
)
)
- for ws, tasks in del_worker_tasks.items():
- ws.has_what -= tasks
- for ts in tasks:
- ts.who_has.remove(ws)
- ws.nbytes -= ts.get_nbytes()
- self.log_event(
- ws.address,
- {
- "action": "replicate-remove",
- "keys": [ts.key for ts in tasks],
- },
- )
-
# Copy not-yet-filled data
while tasks:
gathers = defaultdict(dict)
