Expose details about worker start timeout in the exception message (#9092)

so that calling code can have more precise logic about how to handle the error

Author: nocnokneo

distributed/client.py

@@ -86,6 +86,7 @@
     WorkerPlugin,
     _get_plugin_name,
 )
+from distributed.exceptions import WorkerStartTimeoutError
 from distributed.metrics import time
 from distributed.objects import HasWhat, SchedulerInfo, WhoHas
 from distributed.protocol import serialize, to_serialize
@@ -1631,10 +1632,8 @@ def running_workers(info):
 
         while running_workers(info) < n_workers:
             if deadline and time() > deadline:
-                raise TimeoutError(
-                    "Only %d/%d workers arrived after %s"
-                    % (running_workers(info), n_workers, timeout)
-                )
+                assert timeout is not None
+                raise WorkerStartTimeoutError(running_workers(info), n_workers, timeout)
             await asyncio.sleep(0.1)
             info = await self.scheduler.identity(n_workers=-1)
             self._scheduler_identity = SchedulerInfo(info)

distributed/deploy/cluster.py

@@ -18,6 +18,7 @@
 from distributed.compatibility import PeriodicCallback
 from distributed.core import Status
 from distributed.deploy.adaptive import Adaptive
+from distributed.exceptions import WorkerStartTimeoutError
 from distributed.metrics import time
 from distributed.objects import SchedulerInfo
 from distributed.utils import (
@@ -610,9 +611,8 @@ def running_workers(info):
 
         while n_workers and running_workers(self.scheduler_info) < n_workers:
             if deadline and time() > deadline:
-                raise TimeoutError(
-                    "Only %d/%d workers arrived after %s"
-                    % (running_workers(self.scheduler_info), n_workers, timeout)
+                raise WorkerStartTimeoutError(
+                    running_workers(self.scheduler_info), n_workers, timeout
                 )
             await asyncio.sleep(0.1)
 

distributed/exceptions.py

@@ -1,5 +1,7 @@
 from __future__ import annotations
 
+from asyncio import TimeoutError
+
 
 class Reschedule(Exception):
     """Reschedule this task
@@ -13,3 +15,31 @@ class Reschedule(Exception):
     load across the cluster has significantly changed since first scheduling
     the task.
     """
+
+
+class WorkerStartTimeoutError(TimeoutError):
+    """Raised when the expected number of workers to not start within the timeout period."""
+
+    #: Number of workers that are available.
+    available_workers: int
+
+    #: Number of workers that were expected to be available.
+    expected_workers: int
+
+    #: Timeout period in seconds.
+    timeout: float
+
+    def __init__(
+        self, available_workers: int, expected_workers: int, timeout: float
+    ) -> None:
+        self.available_workers = available_workers
+        self.expected_workers = expected_workers
+        self.timeout = timeout
+        super().__init__(available_workers, expected_workers, timeout)
+
+    def __str__(self) -> str:
+        return "Only %d/%d workers arrived after %s" % (
+            self.available_workers,
+            self.expected_workers,
+            self.timeout,
+        )