introduce total_memory attr

Author: fjetter

distributed/scheduler.py

@@ -1778,6 +1778,7 @@ def __init__(
         self.task_groups = {}
         self.task_prefixes = {}
         self.task_metadata = {}
+        self.total_memory = 0
         self.total_nthreads = 0
         self.total_nthreads_history = [(time(), 0)]
         self.queued = queued
@@ -4087,7 +4088,7 @@ def identity(self, n_workers: int = -1) -> dict[str, Any]:
             "started": self.time_started,
             "n_workers": len(self.workers),
             "total_threads": self.total_nthreads,
-            "total_memory": sum(ws.memory_limit for ws in self.workers.values()),
+            "total_memory": self.total_memory,
             "workers": {
                 worker.address: worker.identity()
                 for worker in itertools.islice(self.workers.values(), n_workers)
@@ -4541,6 +4542,7 @@ async def add_worker(
         dh_addresses.add(address)
         dh["nthreads"] += nthreads
 
+        self.total_memory += ws.memory_limit
         self.total_nthreads += nthreads
         self.total_nthreads_history.append((time(), self.total_nthreads))
         self.aliases[name] = address
@@ -5452,6 +5454,7 @@ async def remove_worker(
         dh_addresses: set = dh["addresses"]
         dh_addresses.remove(address)
         dh["nthreads"] -= ws.nthreads
+        self.total_memory -= ws.memory_limit
         self.total_nthreads -= ws.nthreads
         self.total_nthreads_history.append((time(), self.total_nthreads))
         if not dh_addresses: