Repository: dask/distributed
Pull request #9045: Reduce size of scheduler_info
Author: login: fjetter
Created: 2025-04-16T11:46:32Z | Merged: 2025-04-16T15:53:58Z

See https://github.com/dask/distributed/issues/9043

The scheduler_info is used for this view

![image](https://github.com/user-attachments/assets/0d23a854-8a28-49fb-b49d-a5c8d938122d)

I am introducing a limit to the number of workers being returned.

I cannot imagine that somebody is using this for actual diagnostics (but it's cute and we should leave it) and I just randomly hard coded this to `5`. With this change the total numbers in the upper right corner is still the max and indicates that the below view is only a subset. I'm open to suggestions for making this prettier but that has only low priority.

![image](https://github.com/user-attachments/assets/9ad534ef-dca6-46a0-99fb-477b92cd4b0b)

cc @jacobtomlinson I think you introduced this HTML view and might have an opinion about details or the hard coded number, etc.


Diff:
diff --git a/distributed/client.py b/distributed/client.py
index 1c1c0b46706..1632a52b548 100644
--- a/distributed/client.py
+++ b/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,7 +1312,7 @@ 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 (
@@ -1320,12 +1320,12 @@ def _get_scheduler_info(self):
             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(
diff --git a/distributed/scheduler.py b/distributed/scheduler.py
index a026e897ed6..27d23e3c69d 100644
--- a/distributed/scheduler.py
+++ b/distributed/scheduler.py
@@ -1654,6 +1654,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
@@ -1778,6 +1780,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
@@ -4075,16 +4078,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
@@ -4535,6 +4544,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
@@ -5446,6 +5456,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:
diff --git a/distributed/tests/test_client.py b/distributed/tests/test_client.py
index ac781625917..6ae578db79d 100644
--- a/distributed/tests/test_client.py
+++ b/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):
diff --git a/distributed/tests/test_scheduler.py b/distributed/tests/test_scheduler.py
index 956c886f1f5..5cafb8c7dd8 100644
--- a/distributed/tests/test_scheduler.py
+++ b/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",
diff --git a/distributed/widgets/templates/scheduler_info.html.j2 b/distributed/widgets/templates/scheduler_info.html.j2
index 47501d1a6fe..9f1c697825e 100644
--- a/distributed/widgets/templates/scheduler_info.html.j2
+++ b/distributed/widgets/templates/scheduler_info.html.j2
@@ -10,7 +10,7 @@
                         <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>
@@ -18,7 +18,7 @@
                         <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>
@@ -26,7 +26,7 @@
                         <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>


Reviews:
user: login: fjetter
state: COMMENTED
author_association: MEMBER
submitted_at: 2025-04-16T11:49:43Z
user: login: fjetter
state: COMMENTED
author_association: MEMBER
submitted_at: 2025-04-16T11:50:33Z
user: login: fjetter
state: COMMENTED
author_association: MEMBER
submitted_at: 2025-04-16T11:51:04Z
user: login: copilot-pull-request-reviewer[bot]
body: 

Copilot reviewed 3 out of 4 changed files in this pull request and generated 1 comment.

<details>
<summary>Files not reviewed (1)</summary>

* **distributed/widgets/templates/scheduler_info.html.j2**: Language not supported
</details>

<details>
<summary>Comments suppressed due to low confidence (1)</summary>

**distributed/client.py:4412**
* [nitpick] There are varying conventions for n_workers across client calls (e.g. 0 in dashboard_link, 5 in scheduler_info, -1 for full info). Consider clarifying/documenting the intent behind each value to avoid confusion.
```
def scheduler_info(self, n_workers: int = 5, **kwargs: Any) -> SchedulerInfo:
```
</details>


state: COMMENTED
author_association: NONE
submitted_at: 2025-04-16T11:55:21Z
user: login: fjetter
state: COMMENTED
author_association: MEMBER
submitted_at: 2025-04-16T11:55:54Z

Review comments:
diff_hunk: @@ -1778,6 +1778,7 @@ def __init__(
         self.task_groups = {}
         self.task_prefixes = {}
         self.task_metadata = {}
+        self.total_memory = 0
path: distributed/scheduler.py
user: login: fjetter
body: It's almost overkill to maintain additional state for this but if we have to call `identity` often I want it to be constant time
created_at: 2025-04-16T11:49:43Z
author_association: MEMBER
diff_hunk: @@ -4075,16 +4076,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]:
path: distributed/scheduler.py
user: login: fjetter
body: the default of -1 is to maintain backwards compat (no idea who is calling this in the wild). All callers in this code base are setting this explicitly
created_at: 2025-04-16T11:50:33Z
author_association: MEMBER
diff_hunk: @@ -1778,6 +1778,7 @@ def __init__(
         self.task_groups = {}
         self.task_prefixes = {}
         self.task_metadata = {}
+        self.total_memory = 0
path: distributed/scheduler.py
user: login: fjetter
body: Also, I'm surprised we are not tracking this already. Seems like a useful metric for adaptive scaling or smth like that 🤷 
created_at: 2025-04-16T11:51:04Z
author_association: MEMBER
diff_hunk: @@ -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)
path: distributed/client.py
user: login: Copilot
body: Using n_workers=0 in this context results in an empty workers list. If the intention is to provide a default view for diagnostics, consider using a value such as -1 (to fetch all) or the same default limit (e.g. 5) used elsewhere.
```suggestion
            scheduler, info = self._get_scheduler_info(n_workers=-1)
```
created_at: 2025-04-16T11:55:21Z
author_association: NONE
diff_hunk: @@ -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)
path: distributed/client.py
user: login: fjetter
body: that's intentional since the workers info is not even accessed here
created_at: 2025-04-16T11:55:54Z
author_association: MEMBER

PR comments:
user: login: github-actions[bot]
created_at: 2025-04-16T12:01:48Z
body: ## Unit Test Results
_See [test report](https://dask.github.io/distributed/test_report.html) for an extended history of previous test failures. This is useful for diagnosing flaky tests._

    27 files  ±0      27 suites  ±0   11h 17m 19s ⏱️ + 1m 9s
 4 100 tests ±0   3 980 ✅  - 2    112 💤 ±0  7 ❌ +1  1 🔥 +1 
51 401 runs  +1  49 085 ✅ ±0  2 308 💤  - 1  7 ❌ +1  1 🔥 +1 

For more details on these failures and errors, see [this check](https://github.com/dask/distributed/runs/40656379546).

Results for commit 41ce938f. ± Comparison against base commit 16aa1890.

[test-results]:data:application/gzip;base64,H4sIAEqv/2cC/03Pyw7CIBCF4VdpWLuAgeHiy5iCkBB7MRRWxnd3qqF2+X9DTsKLpTzFjV0HMJeBbS3Xf91bGWteF2rFtXREdK37XQnOe962FgKZdPZkj/wkEwIOSmOeiMwBsZS17I9ISlv2XRSK9+y7ynGLh/12QXLb6bT77fNsWOc5VwqmRIhO2iQQtNHOq6jRAShrPGAyKOlLPmBk7w9os9+8FAEAAA==

:recycle: This comment has been updated with latest results.
author_association: CONTRIBUTOR
user: login: fjetter
created_at: 2025-04-16T15:53:54Z
body: I'll move forward with merging. Test failures appear to be unrelated and I think the benefits of this outweight possible minor UX things. I'm happy to follow up on this if there is feedback but I'd like to get the functional fix into the release today.
author_association: MEMBER
user: login: rcheatham-q
created_at: 2025-04-28T21:20:39Z
body: For the record, this broke my automated alert that would notify me when I didn't have all of my hundreds of expected workers connected to the scheduler. If there is a better method of doing this, I'm happy to hear it, but this method seemed to be exactly what I needed.
author_association: NONE
user: login: fjetter
created_at: 2025-04-29T08:01:36Z
body: > For the record, this broke my automated alert that would notify me when I didn't have all of my hundreds of expected workers connected to the scheduler

There are now new fields

* "n_workers"
* "total_threads"
* "total_memory"

Alternatively, you can also query this yourself with `client.scheduler_info(n_workers=-1)`
author_association: MEMBER
user: login: rcheatham-q
created_at: 2025-05-14T20:37:32Z
body: > There are now new fields
> 
> * "n_workers"
> * "total_threads"
> * "total_memory"
> 
> Alternatively, you can also query this yourself with `client.scheduler_info(n_workers=-1)`

The fact that a workaround exists doesn't change the fact that this was a backwards incompatible change that broke a bunch of things. We can already see three of them listed in this PR, and I have encountered another issue caused by this change that means I have to pin distributed to the latest version without this change.

The proper implementation for this change would have been to set the default value of `n_workers` to `-1` and changed the view to specify the number of workers to pull.

At the very least this should have been mentioned as a breaking change in the [changelog](https://docs.dask.org/en/stable/changelog.html#v2025-4-0).
author_association: NONE