expose details about worker start timeout in the exception message

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

Author: nocnokneo

distributed/client.py

@@ -127,6 +127,7 @@
     unpack_remotedata,
 )
 from distributed.worker import get_client, get_worker, secede
+from distributed.exceptions import WorkerStartTimeoutError
 
 if TYPE_CHECKING:
     from typing_extensions import TypeAlias
@@ -1651,10 +1652,7 @@ 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)
-                )
+                raise WorkerStartTimeoutError(running_workers(info), n_workers, timeout)
             await asyncio.sleep(0.1)
             info = await self.scheduler.identity()
             self._scheduler_identity = SchedulerInfo(info)

distributed/deploy/cluster.py

@@ -29,6 +29,7 @@
     format_dashboard_link,
     log_errors,
 )
+from distributed.exceptions import WorkerStartTimeoutError
 
 logger = logging.getLogger(__name__)
 
@@ -610,10 +611,7 @@ 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)
 
             self.scheduler_info = SchedulerInfo(await self.scheduler_comm.identity())

distributed/exceptions.py

@@ -1,5 +1,7 @@
 from __future__ import annotations
 
+from asyncio import TimeoutError
+
 
 class Reschedule(Exception):
     """Reschedule this task
@@ -13,3 +15,32 @@ 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."""
+
+    def __init__(self, available_workers: int, expected_workers: int, timeout) -> None:
+        super().__init__(available_workers, expected_workers, timeout)
+
+    @property
+    def available_workers(self) -> int:
+        """Number of workers that are available."""
+        return self.args[0]
+
+    @property
+    def expected_workers(self) -> int:
+        """Number of workers that were expected to be available."""
+        return self.args[1]
+
+    @property
+    def timeout(self) -> float:
+        """Timeout period in seconds."""
+        return self.args[2]
+
+    def __str__(self) -> str:
+        return "Only %d/%d workers arrived after %s" % (
+            self.available_workers,
+            self.expected_workers,
+            self.timeout,
+        )