Dynamically scale heartbeat and scheduler_info intervals (#9046)

Author: fjetter

distributed/client.py

@@ -1164,6 +1164,7 @@ def __init__(
             "restart": self._handle_restart,
             "error": self._handle_error,
             "event": self._handle_event,
+            "adjust-heartbeat-interval": self._adjust_heartbeat_intervals,
         }
 
         self._state_handlers = {
@@ -1645,6 +1646,11 @@ def wait_for_workers(self, n_workers: int, timeout: float | None = None) -> None
 
         return self.sync(self._wait_for_workers, n_workers, timeout=timeout)
 
+    def _adjust_heartbeat_intervals(self, interval):
+        """Adjust the heartbeat intervals for the client and scheduler"""
+        self._periodic_callbacks["heartbeat"].callback_time = interval * 1000
+        self._periodic_callbacks["scheduler-info"].callback_time = interval * 1000
+
     def _heartbeat(self):
         # Don't send heartbeat if scheduler comm or cluster are already closed
         if self.scheduler_comm is not None and not (

distributed/scheduler.py

@@ -5701,6 +5701,14 @@ def client_heartbeat(self, client: str) -> None:
         """Handle heartbeats from Client"""
         cs = self.clients[client]
         cs.last_seen = time()
+        self.client_comms[client].send(
+            {
+                "op": "adjust-heartbeat-interval",
+                # heartbeat_interval is used for workers
+                # We don't require the clients to heartbeat this often
+                "interval": heartbeat_interval(len(self.clients)) * 10,
+            }
+        )
 
     ###################
     # Task Validation #

distributed/tests/test_client.py

@@ -8257,3 +8257,27 @@ def f(x, ev):
     ev.set()
     with pytest.raises(FutureCancelledError):
         future.result()
+
+
+@gen_cluster(client=True, nthreads=[])
+async def test_adjust_heartbeat(c, s):
+    assert "heartbeat" in c._periodic_callbacks
+    heartbeat_pc = c._periodic_callbacks["heartbeat"]
+    scheduler_info_pc = c._periodic_callbacks["scheduler-info"]
+    assert 0 < (initial_value := heartbeat_pc.callback_time) < 10_000
+
+    clients = []
+    for _ in range(20):
+        clients.append(Client(s.address, asynchronous=True))
+    await asyncio.gather(*clients)
+
+    while initial_value == heartbeat_pc.callback_time:
+        await asyncio.sleep(0.1)
+    assert (newvalue := heartbeat_pc.callback_time) > initial_value
+    assert heartbeat_pc.callback_time == scheduler_info_pc.callback_time
+
+    await asyncio.gather(*[cc.close() for cc in clients])
+    while newvalue == heartbeat_pc.callback_time:
+        await asyncio.sleep(0.1)
+    assert heartbeat_pc.callback_time == initial_value
+    assert heartbeat_pc.callback_time == scheduler_info_pc.callback_time