feat(kuberay): support resolve_hostname hook in PipesKubeRayJobClient (#340)

Author: danielgafni

.zed/settings.json

@@ -0,0 +1,7 @@
+{
+  "languages": {
+    "Python": {
+      "language_servers": ["basedpyright", "ruff", "!pyrefly", "!ty", "!pyright", "!pylsp"]
+    }
+  }
+}

src/dagster_ray/core/pipes.py

@@ -227,11 +227,10 @@ class PipesRayJobClient(dg.PipesClient, TreatAsResourceParam):
     Args:
         address: Ray dashboard HTTP address.
             If unspecified, connects to a local Ray cluster or uses the ``RAY_ADDRESS`` environment variable.
-        headers: HTTP headers for dashboard requests, e.g. ``{"Authorization": "Bearer token"}``.
-        verify: TLS certificate verification. ``True`` uses system certs, ``False`` disables
-            verification, or a path to a CA bundle.
-        cookies: Cookies to send with dashboard requests.
-        metadata: Arbitrary metadata stored alongside all submitted jobs.
+        headers: HTTP headers for Ray Dashboard requests. Passed to [`JobSubmissionClient`][ray.job_submission.JobSubmissionClient].
+        verify: Whether to verify TLS certificate. Passed to [`JobSubmissionClient`][ray.job_submission.JobSubmissionClient].
+        cookies: HTTP cookies for Ray Dashboard requests. Passed to [`JobSubmissionClient`][ray.job_submission.JobSubmissionClient].
+        metadata: Ray Job metadata. Passed to [`JobSubmissionClient`][ray.job_submission.JobSubmissionClient].
         context_injector: A context injector to use to inject
             context into the Ray job. Defaults to [`PipesEnvContextInjector`][dagster.PipesEnvContextInjector].
         message_reader: A message reader to use when reading Pipes messages.
@@ -289,9 +288,9 @@ def run(  # type: ignore
         Execute a RayJob, enriched with the Pipes protocol.
 
         Args:
-            context (OpExecutionContext): Current Dagster op or asset context.
-            submit_job_params (Dict[str, Any]): Parameters for [`JobSubmissionClient.submit_job`][ray.job_submission.JobSubmissionClient.submit_job].
-            extras (Optional[Dict[str, Any]]): Additional information to pass to the Pipes session.
+            context: Current Dagster op or asset context.
+            submit_job_params: Parameters for [`JobSubmissionClient.submit_job`][ray.job_submission.JobSubmissionClient.submit_job].
+            extras: Additional information to pass to the Pipes session, retrievable via [`PipesContext.get_extras`](https://docs.dagster.io/integrations/libraries/pipes/dagster-pipes#dagster_pipes.PipesContext.get_extra).
         """
 
         with open_pipes_session(

src/dagster_ray/kuberay/pipes.py

@@ -26,7 +26,7 @@
 )
 from dagster_ray.kuberay.client import RayJobClient
 from dagster_ray.kuberay.client.rayjob.client import RayJobStatus
-from dagster_ray.kuberay.utils import normalize_k8s_label_values
+from dagster_ray.kuberay.utils import k8s_service_fqdn, normalize_k8s_label_values
 from dagster_ray.types import OpOrAssetExecutionContext
 
 if TYPE_CHECKING:
@@ -36,6 +36,8 @@
 class PipesKubeRayJobClient(dg.PipesClient, TreatAsResourceParam):
     """A pipes client for running ``RayJob`` on Kubernetes.
 
+    Most constructor arguments can be overridden in the [`run`][dagster_ray.kuberay.pipes.PipesKubeRayJobClient.run] method.
+
     Args:
         context_injector: A context injector to use to inject
             context into the ``RayJob``. Defaults to [`PipesEnvContextInjector`][dagster.PipesEnvContextInjector].
@@ -48,18 +50,12 @@ class PipesKubeRayJobClient(dg.PipesClient, TreatAsResourceParam):
         poll_interval: Interval at which to poll Kubernetes for status updates.
         port_forward: Whether to use Kubernetes port-forwarding to connect to the KubeRay cluster.
             Is useful when running in a local environment.
-        address: Ray dashboard address (e.g., "https://ray-cluster.example.com").
-            When provided, connects directly to this address instead of using port-forwarding or in-cluster service IPs.
-            Can be overridden per-job in the `run()` method.
-        headers: HTTP headers for dashboard requests, e.g. ``{"Authorization": "Bearer token"}``.
-            Can be overridden per-job in the `run()` method.
-        verify: TLS certificate verification. ``True`` uses system certs, ``False`` disables
-            verification, or a path to a CA bundle.
-            Can be overridden per-job in the `run()` method.
-        cookies: Arbitrary metadata stored alongside all submitted jobs.
-            Can be overridden per-job in the `run()` method.
-        metadata: Arbitrary metadata to store along with all jobs.
-            Will be merged with per-job metadata. Can be overridden per-job in the `run()` method.
+        address: Ray dashboard address. Passed to [`JobSubmissionClient`][ray.job_submission.JobSubmissionClient].
+            Generated by [`resolve_hostname`][dagster_ray.kuberay.pipes.PipesKubeRayJobClient.resolve_hostname] if not provided.
+        headers: HTTP headers for Ray Dashboard requests. Passed to [`JobSubmissionClient`][ray.job_submission.JobSubmissionClient].
+        verify: Whether to verify TLS certificate. Passed to [`JobSubmissionClient`][ray.job_submission.JobSubmissionClient].
+        cookies: HTTP cookies for Ray Dashboard requests. Passed to [`JobSubmissionClient`][ray.job_submission.JobSubmissionClient].
+        metadata: Ray Job metadata. Passed to [`JobSubmissionClient`][ray.job_submission.JobSubmissionClient].
 
     Info:
         Image defaults to `dagster/image` run tag.
@@ -130,15 +126,18 @@ def run(  # type: ignore
         """
         Execute a RayJob, enriched with the Pipes protocol.
 
+        Parameters passed to [`JobSubmissionClient`][ray.job_submission.JobSubmissionClient] take precedence over the constructor arguments.
+
         Args:
             context: Current Dagster op or asset context.
-            ray_job: RayJob specification. `API reference <https://ray-project.github.io/kuberay/reference/api/#rayjob>`_.
-            extras: Additional information to pass to the Pipes session.
-            address: Ray dashboard address override for this specific job.
-            headers: HTTP headers override for this specific job.
-            verify: TLS verification override for this specific job.
-            cookies: Cookies override for this specific job.
-            metadata: Metadata override for this specific job.
+            ray_job: `RayJob` specification. See [API reference](https://ray-project.github.io/kuberay/reference/api/#rayjob).
+            extras: Additional information to pass to the Pipes session, retrievable via [`PipesContext.get_extras`](https://docs.dagster.io/integrations/libraries/pipes/dagster-pipes#dagster_pipes.PipesContext.get_extra).
+            address: Ray dashboard address. Passed to [`JobSubmissionClient`][ray.job_submission.JobSubmissionClient].
+                Generated by [`resolve_hostname`][dagster_ray.kuberay.pipes.PipesKubeRayJobClient.resolve_hostname] if not provided.
+            headers: HTTP headers for Ray Dashboard requests. Passed to [`JobSubmissionClient`][ray.job_submission.JobSubmissionClient].
+            verify: Whether to verify TLS certificate. Passed to [`JobSubmissionClient`][ray.job_submission.JobSubmissionClient].
+            cookies: HTTP cookies for Ray Dashboard requests. Passed to [`JobSubmissionClient`][ray.job_submission.JobSubmissionClient].
+            metadata: Ray Job metadata. Passed to [`JobSubmissionClient`][ray.job_submission.JobSubmissionClient].
         """
         with open_pipes_session(
             context=context,
@@ -154,13 +153,24 @@ def run(  # type: ignore
             name = ray_job["metadata"]["name"]
             namespace = ray_job["metadata"]["namespace"]
 
+            cluster_name = self.client.get_ray_cluster_name(
+                name=name, namespace=namespace, timeout=self.timeout, poll_interval=self.poll_interval
+            )
+
+            resolved_address = address or self.address
+            if resolved_address is None and not self.port_forward:
+                cluster_status = self.client.ray_cluster_client.get_status(
+                    name=cluster_name, namespace=namespace, timeout=self.timeout, poll_interval=self.poll_interval
+                )
+                service_name = cluster_status["head"]["serviceName"]  # type: ignore
+                dashboard_port = int(cluster_status["endpoints"]["dashboard"])  # type: ignore
+                resolved_address = self.resolve_hostname(service_name, namespace, dashboard_port)
+
             with self.client.ray_cluster_client.job_submission_client(
-                name=self.client.get_ray_cluster_name(
-                    name=name, namespace=namespace, timeout=self.timeout, poll_interval=self.poll_interval
-                ),
+                name=cluster_name,
                 namespace=namespace,
                 port_forward=self.port_forward,
-                address=address or self.address,
+                address=resolved_address,
                 headers=headers or self.headers,
                 verify=verify if verify is not None else self.verify,
                 cookies=cookies or self.cookies,
@@ -203,14 +213,41 @@ def run(  # type: ignore
                         self._terminate(
                             context,
                             start_response,
-                            address=address,
+                            address=resolved_address,
                             headers=headers,
                             verify=verify,
                             cookies=cookies,
                             metadata=metadata,
                         )
                     raise
 
+    def resolve_hostname(self, service_name: str, namespace: str, dashboard_port: int) -> str:
+        """Resolve the Ray dashboard address from a Kubernetes service name, namespace, and port.
+
+        By default, returns an HTTP URL using the in-cluster FQDN
+        (e.g. `http://my-svc.namespace.svc.cluster.local:8265`).
+
+        Override this method to customize the address, for example to use a proxy
+        or a different protocol. The `dashboard_port` can be ignored if not needed.
+
+        Args:
+            service_name: The Kubernetes service name for the Ray head node.
+            namespace: The Kubernetes namespace.
+            dashboard_port: The Ray dashboard port reported by the KubeRay operator.
+
+        Returns:
+            The full Ray dashboard URL (including protocol and optional port).
+
+        Example:
+            ```python
+            class MyPipesKubeRayJobClient(PipesKubeRayJobClient):
+                def resolve_hostname(self, service_name: str, namespace: str, dashboard_port: int) -> str:
+                    return f"https://{service_name}.{namespace}.company.com"
+            ```
+        """
+        hostname = k8s_service_fqdn(service_name, namespace)
+        return f"http://{hostname}:{dashboard_port}"
+
     def get_dagster_tags(self, context: OpOrAssetExecutionContext) -> dict[str, str]:
         tags = get_dagster_tags(context)
         return tags
@@ -320,7 +357,7 @@ def _terminate(
             name=name,
             namespace=namespace,
             port_forward=self.port_forward,
-            address=address or self.address,
+            address=address,
             headers=headers or self.headers,
             verify=verify if verify is not None else self.verify,
             cookies=cookies or self.cookies,

src/dagster_ray/kuberay/resources/base.py

@@ -35,7 +35,17 @@ 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:
+        Override this method to customize the host address, for example to use a custom domain.
+
+        The returned hostname is used to construct both the Ray Client address (`ray://<host>:<redis_port>`)
+        and the dashboard URL (`http://<host>:<dashboard_port>`).
+
+        Args:
+            service_name: The Kubernetes service name for the Ray head node.
+            namespace: The Kubernetes namespace.
+
+        Returns:
+            The hostname to use for connecting to the Ray cluster.
 
         Example:
             ```python