Reduce size of scheduler_info (#9045)

Author: fjetter

distributed/client.py

@@ -1298,7 +1298,7 @@ def dashboard_link(self):
         try:
             return self.cluster.dashboard_link
         except AttributeError:
-            scheduler, info = self._get_scheduler_info()
+            scheduler, info = self._get_scheduler_info(n_workers=0)
             if scheduler is None:
                 return None
             else:
@@ -1312,20 +1312,20 @@ def dashboard_link(self):
 
             return format_dashboard_link(host, port)
 
-    def _get_scheduler_info(self):
+    def _get_scheduler_info(self, n_workers):
         from distributed.scheduler import Scheduler
 
         if (
             self.cluster
             and hasattr(self.cluster, "scheduler")
             and isinstance(self.cluster.scheduler, Scheduler)
         ):
-            info = self.cluster.scheduler.identity()
+            info = self.cluster.scheduler.identity(n_workers=n_workers)
             scheduler = self.cluster.scheduler
         elif (
             self._loop_runner.is_started() and self.scheduler and not self.asynchronous
         ):
-            info = sync(self.loop, self.scheduler.identity)
+            info = sync(self.loop, self.scheduler.identity, n_workers=n_workers)
             scheduler = self.scheduler
         else:
             info = self._scheduler_identity
@@ -1368,7 +1368,7 @@ def _repr_html_(self):
         except PackageNotFoundError:
             JUPYTERLAB = False
 
-        scheduler, info = self._get_scheduler_info()
+        scheduler, info = self._get_scheduler_info(n_workers=5)
 
         return get_template("client.html.j2").render(
             id=self.id,
@@ -1585,18 +1585,20 @@ async def _ensure_connected(self, timeout=None):
 
         logger.debug("Started scheduling coroutines. Synchronized")
 
-    async def _update_scheduler_info(self):
+    async def _update_scheduler_info(self, n_workers=5):
         if self.status not in ("running", "connecting") or self.scheduler is None:
             return
         try:
-            self._scheduler_identity = SchedulerInfo(await self.scheduler.identity())
+            self._scheduler_identity = SchedulerInfo(
+                await self.scheduler.identity(n_workers=n_workers)
+            )
         except OSError:
             logger.debug("Not able to query scheduler for identity")
 
     async def _wait_for_workers(
         self, n_workers: int, timeout: float | None = None
     ) -> None:
-        info = await self.scheduler.identity()
+        info = await self.scheduler.identity(n_workers=-1)
         self._scheduler_identity = SchedulerInfo(info)
         if timeout:
             deadline = time() + parse_timedelta(timeout)
@@ -1619,7 +1621,7 @@ def running_workers(info):
                     % (running_workers(info), n_workers, timeout)
                 )
             await asyncio.sleep(0.1)
-            info = await self.scheduler.identity()
+            info = await self.scheduler.identity(n_workers=-1)
             self._scheduler_identity = SchedulerInfo(info)
 
     def wait_for_workers(self, n_workers: int, timeout: float | None = None) -> None:
@@ -4407,11 +4409,14 @@ async def _profile(
         else:
             return state
 
-    def scheduler_info(self, **kwargs):
+    def scheduler_info(self, n_workers: int = 5, **kwargs: Any) -> SchedulerInfo:
         """Basic information about the workers in the cluster
 
         Parameters
         ----------
+        n_workers: int
+            The number of workers for which to fetch information. To fetch all,
+            use -1.
         **kwargs : dict
             Optional keyword arguments for the remote function
 
@@ -4429,7 +4434,7 @@ def scheduler_info(self, **kwargs):
                                          'time-delay': 0.0061032772064208984}}}
         """
         if not self.asynchronous:
-            self.sync(self._update_scheduler_info)
+            self.sync(self._update_scheduler_info, n_workers=n_workers)
         return self._scheduler_identity
 
     def dump_cluster_state(

distributed/scheduler.py

@@ -1663,6 +1663,8 @@ class SchedulerState:
     idle_task_count: set[WorkerState]
     #: Workers that are fully utilized. May include non-running workers.
     saturated: set[WorkerState]
+    #: Current total memory across all workers (sum over memory_limit)
+    total_memory: int
     #: Current number of threads across all workers
     total_nthreads: int
     #: History of number of threads
@@ -1787,6 +1789,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
@@ -4084,16 +4087,22 @@ def _repr_html_(self) -> str:
             tasks=self.tasks,
         )
 
-    def identity(self) -> dict[str, Any]:
+    def identity(self, n_workers: int = -1) -> dict[str, Any]:
         """Basic information about ourselves and our cluster"""
+        if n_workers == -1:
+            n_workers = len(self.workers)
         d = {
             "type": type(self).__name__,
             "id": str(self.id),
             "address": self.address,
             "services": {key: v.port for (key, v) in self.services.items()},
             "started": self.time_started,
+            "n_workers": len(self.workers),
+            "total_threads": self.total_nthreads,
+            "total_memory": self.total_memory,
             "workers": {
-                worker.address: worker.identity() for worker in self.workers.values()
+                worker.address: worker.identity()
+                for worker in itertools.islice(self.workers.values(), n_workers)
             },
         }
         return d
@@ -4544,6 +4553,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
@@ -5455,6 +5465,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:

distributed/tests/test_client.py

@@ -3923,6 +3923,10 @@ def test_scheduler_info(c):
     assert isinstance(info, dict)
     assert len(info["workers"]) == 2
     assert isinstance(info["started"], float)
+    info = c.scheduler_info(n_workers=1)
+    assert len(info["workers"]) == 1
+    info = c.scheduler_info(n_workers=-1)
+    assert len(info["workers"]) == 2
 
 
 def test_write_scheduler_file(c, loop):

distributed/tests/test_scheduler.py

@@ -4301,6 +4301,9 @@ async def test_Scheduler__to_dict(c, s, a):
         "extensions",
         "services",
         "started",
+        "n_workers",
+        "total_threads",
+        "total_memory",
         "workers",
         "status",
         "thread_id",

distributed/widgets/templates/scheduler_info.html.j2

@@ -10,23 +10,23 @@
                         <strong>Comm:</strong> {{ address }}
                     </td>
                     <td style="text-align: left;">
-                        <strong>Workers:</strong> {{ workers | length }}
+                        <strong>Workers:</strong> {{ n_workers }} {% if n_workers > workers | length %} (shown below: {{ workers | length }}) {% endif %}
                     </td>
                 </tr>
                 <tr>
                     <td style="text-align: left;">
                         <strong>Dashboard:</strong> <a href="{{ scheduler | format_dashboard_address }}" target="_blank">{{ scheduler | format_dashboard_address }}</a>
                     </td>
                     <td style="text-align: left;">
-                        <strong>Total threads:</strong> {{ workers.values() | map(attribute='nthreads') | sum }}
+                        <strong>Total threads:</strong> {{ total_threads }}
                     </td>
                 </tr>
                 <tr>
                     <td style="text-align: left;">
                         <strong>Started:</strong> {{ started | datetime_from_timestamp | format_time_ago }}
                     </td>
                     <td style="text-align: left;">
-                        <strong>Total memory:</strong> {{ workers.values() | map(attribute='memory_limit') | sum | format_bytes }}
+                        <strong>Total memory:</strong> {{ total_memory | format_bytes }}
                     </td>
                 </tr>
             </table>