feat(kuberay): add resolve_hostname hook to KubeRayResource (#335)

Author: danielgafni

docs/api/kuberay.md

@@ -114,6 +114,7 @@ These resources initialize Ray client connection with a remote cluster.
         - "image"
         - "deployment_name"
         - "poll_interval"
+        - "resolve_hostname"
 
 ---
 

src/dagster_ray/kuberay/resources/base.py

@@ -7,7 +7,7 @@
 
 from dagster_ray._base.constants import DEFAULT_DEPLOYMENT_NAME
 from dagster_ray._base.resources import RayResource
-from dagster_ray.kuberay.utils import get_k8s_object_name
+from dagster_ray.kuberay.utils import get_k8s_object_name, k8s_service_fqdn
 from dagster_ray.types import AnyDagsterContext
 
 
@@ -31,6 +31,21 @@ class BaseKubeRayResource(RayResource, ABC):
     def namespace(self) -> str:
         raise NotImplementedError
 
+    def resolve_hostname(self, service_name: str, namespace: str) -> str:
+        """Resolve the host address from a Kubernetes service name and namespace.
+
+        By default, returns the in-cluster FQDN (e.g. `my-svc.namespace.svc.cluster.local`).
+        Override this method to customize the host address, for example to use a custom domain:
+
+        Example:
+            ```python
+            class MyKubeRayCluster(KubeRayCluster):
+                def resolve_hostname(self, service_name: str, namespace: str) -> str:
+                    return f"{service_name}.{namespace}.company.com"
+            ```
+        """
+        return k8s_service_fqdn(service_name, namespace)
+
     def _get_step_name(self, context: AnyDagsterContext) -> str:
         assert isinstance(context.run_id, str)
         assert context.dagster_run is not None

src/dagster_ray/kuberay/resources/raycluster.py

@@ -12,7 +12,7 @@
 from dagster_ray.kuberay.configs import ClusterSharing, RayClusterConfig
 from dagster_ray.kuberay.leader_election import LeaderElection
 from dagster_ray.kuberay.resources.base import BaseKubeRayResource
-from dagster_ray.kuberay.utils import k8s_service_fqdn, normalize_k8s_label_values
+from dagster_ray.kuberay.utils import normalize_k8s_label_values
 from dagster_ray.types import AnyDagsterContext
 from dagster_ray.utils import get_dagster_run
 
@@ -272,7 +272,7 @@ def wait(self, context: AnyDagsterContext):
         )[  # pyright: ignore
             "head"
         ]["serviceName"]
-        self._host = k8s_service_fqdn(service_name, self.namespace)
+        self._host = self.resolve_hostname(service_name, self.namespace)
 
     @override
     def on_ready(self, context: AnyDagsterContext):

src/dagster_ray/kuberay/resources/rayjob.py

@@ -8,7 +8,7 @@
 from dagster_ray.kuberay.client import RayJobClient
 from dagster_ray.kuberay.configs import RayJobConfig, RayJobSpec
 from dagster_ray.kuberay.resources.base import BaseKubeRayResource
-from dagster_ray.kuberay.utils import k8s_service_fqdn, normalize_k8s_label_values
+from dagster_ray.kuberay.utils import normalize_k8s_label_values
 from dagster_ray.types import AnyDagsterContext
 
 if TYPE_CHECKING:
@@ -163,7 +163,7 @@ def wait(self, context: AnyDagsterContext):
         )[  # pyright: ignore
             "head"
         ]["serviceName"]
-        self._host = k8s_service_fqdn(service_name, self.namespace)
+        self._host = self.resolve_hostname(service_name, self.namespace)
 
     @override
     def on_ready(self, context: AnyDagsterContext):