✨ introduce failure_tolerance_timeout parameter for KubeRay resources (#247)

Author: danielgafni

docs/api/kuberay.md

@@ -15,6 +15,7 @@ These resources initialize Ray client connection with a remote cluster.
         - "lifecycle"
         - "ray_job"
         - "client"
+        - "failure_tolerance_timeout"
         - "log_cluster_conditions"
 
 
@@ -25,6 +26,7 @@ These resources initialize Ray client connection with a remote cluster.
         - "lifecycle"
         - "ray_cluster"
         - "client"
+        - "failure_tolerance_timeout"
         - "log_cluster_conditions"
 
 ---

src/dagster_ray/kuberay/client/raycluster/client.py

@@ -73,8 +73,10 @@ class RayClusterEndpoints(TypedDict):  # these are ports
 
 
 class RayClusterHead(TypedDict):
-    podIP: str
-    serviceIP: str
+    podIP: NotRequired[str]
+    podName: NotRequired[str]
+    serviceIP: NotRequired[str]
+    serviceName: NotRequired[str]
 
 
 class RayClusterStatus(TypedDict):
@@ -112,12 +114,23 @@ def wait_until_ready(
         name: str,
         namespace: str,
         timeout: float,
-        image: str | None = None,
+        failure_tolerance_timeout: float = 0.0,
         poll_interval: float = 5.0,
         log_cluster_conditions: bool = False,
     ) -> tuple[str, RayClusterEndpoints]:
         """
         If ready, returns service ip address and a dictionary of ports.
+
+        Parameters:
+            name (str): The name of the `RayCluster` resource
+            namespace (str): The namespace of the `RayCluster` resource
+            timeout (float): The timeout in seconds to wait for the cluster to become ready.
+            failure_tolerance_timeout (float): The period in seconds to wait for the cluster to transition out of `failed` state if it reaches it. This state can be transient under certain conditions. With the default value of 0, the first `failed` state appearance will raise an exception immediately.
+            poll_interval (float): The interval in seconds to poll the cluster status.
+            log_cluster_conditions (bool): Whether to log cluster conditions. See [KubeRay docs](https://docs.ray.io/en/latest/cluster/kubernetes/user-guides/observability.html#raycluster-status-conditions)
+
+        Returns:
+            tuple[str, RayClusterEndpoints]: The service ip address and a dictionary of ports.
         """
         start_time = time.time()
 
@@ -150,14 +163,22 @@ def get_status_with_retry() -> RayClusterStatus:
                     logger.info(f"RayCluster {namespace}/{name} condition: {conditions[condition_index_to_log]}")
                     condition_index_to_log += 1
 
-            if state == "failed":
-                raise Exception(
-                    f"RayCluster {namespace}/{name} failed to start. Status:\n\n{status}\n\nMore details: `kubectl -n {namespace} describe RayCluster {name}`"
+            if state == "failed" and (time.time() - start_time > failure_tolerance_timeout):
+                raise RuntimeError(
+                    f"RayCluster {namespace}/{name} status transitioned into `failed`. Consider increasing `failure_tolerance_timeout` ({failure_tolerance_timeout:.1f}s). Status:\n\n{status}\n\nMore details: `kubectl -n {namespace} describe RayCluster {name}`"
                 )
 
-            if state == "ready" and status.get("head") and status.get("endpoints", {}).get("dashboard"):
+            if (
+                state == "ready"
+                and status.get("endpoints", {}).get("dashboard")
+                and (head := status.get("head"))
+                and head.get("serviceIP")
+                and head.get("serviceName")
+            ):
+                # TODO: this should return serviceName instead
+                # to support multi-cluster networking
                 logger.debug(f"RayCluster {namespace}/{name} is ready!")
-                return status["head"]["serviceIP"], status["endpoints"]  # pyright: ignore[reportTypedDictNotRequiredAccess]
+                return head["serviceIP"], status["endpoints"]  # pyright: ignore[reportTypedDictNotRequiredAccess]
 
             time.sleep(poll_interval)
         else:

src/dagster_ray/kuberay/client/rayjob/client.py

@@ -71,12 +71,24 @@ def wait_until_ready(
         name: str,
         namespace: str,
         timeout: float = 600,
+        failure_tolerance_timeout: float = 0.0,
         poll_interval: float = 1.0,
         log_cluster_conditions: bool = False,
     ) -> tuple[str, RayClusterEndpoints]:
         """Wait until the RayCluster attached to the RayJob is ready.
 
         This doesn't necessarily mean that the cluster has already taken a job, just that it is ready to accept connections.
+
+        Parameters:
+            name (str): The name of the `RayJob` resource
+            namespace (str): The namespace of the `RayJob` resource
+            timeout (float): The timeout in seconds to wait for the cluster to become ready.
+            failure_tolerance_timeout (float): The period in seconds to wait for the cluster to transition out of `failed` state if it reaches it. This state can be transient under certain conditions. With the default value of 0, the first `failed` state appearance will raise an exception immediately.
+            poll_interval (float): The interval in seconds to poll the cluster status.
+            log_cluster_conditions (bool): Whether to log cluster conditions. See [KubeRay docs](https://docs.ray.io/en/latest/cluster/kubernetes/user-guides/observability.html#raycluster-status-conditions)
+
+        Returns:
+            tuple[str, RayClusterEndpoints]: The service ip address and a dictionary of ports.
         """
         ray_cluster_name = self.get_ray_cluster_name(name, namespace, timeout=timeout, poll_interval=poll_interval)
         ray_cluster_client = self.ray_cluster_client
@@ -87,6 +99,7 @@ def wait_until_ready(
             ray_cluster_name,
             namespace=namespace,
             timeout=timeout,
+            failure_tolerance_timeout=failure_tolerance_timeout,
             poll_interval=poll_interval,
             log_cluster_conditions=log_cluster_conditions,
         )

src/dagster_ray/kuberay/resources/base.py

@@ -21,6 +21,10 @@ class BaseKubeRayResourceConfig(dg.Config):
         default=DEFAULT_DEPLOYMENT_NAME,
         description="Dagster deployment name. Is used as a prefix for the Kubernetes resource name. Dagster Cloud variables are used to determine the default value.",
     )
+    failure_tolerance_timeout: float = Field(
+        default=0.0,
+        description="The period in seconds to wait for the cluster to transition out of `failed` state if it reaches it. This state can be transient under certain conditions. With the default value of 0, the first `failed` state appearance will raise an exception immediately.",
+    )
     poll_interval: float = Field(default=1.0, description="Poll interval for various API requests")
 
     _host: str = PrivateAttr()

src/dagster_ray/kuberay/resources/raycluster.py

@@ -120,6 +120,7 @@ def wait(self, context: AnyDagsterContext):
             self.name,
             namespace=self.namespace,
             timeout=self.timeout,
+            failure_tolerance_timeout=self.failure_tolerance_timeout,
             poll_interval=self.poll_interval,
             log_cluster_conditions=self.log_cluster_conditions,
         )

src/dagster_ray/kuberay/resources/rayjob.py

@@ -150,6 +150,7 @@ def wait(self, context: AnyDagsterContext):
         self.client.wait_until_ready(
             name=self.name,
             namespace=self.namespace,
+            failure_tolerance_timeout=self.failure_tolerance_timeout,
             log_cluster_conditions=self.log_cluster_conditions,
             timeout=self.timeout,
             poll_interval=self.poll_interval,