✨ add run launcher and executor for KubeRay

Author: danielgafni

src/dagster_ray/_base/utils.py

@@ -1,12 +1,17 @@
 from __future__ import annotations
 
 import dagster as dg
+from dagster._core.executor.step_delegating import (
+    StepHandlerContext,
+)
+from dagster._core.launcher.base import LaunchRunContext
 
 from dagster_ray._base.constants import DEFAULT_DEPLOYMENT_NAME
+from dagster_ray.types import AnyDagsterContext
 
 
 def get_dagster_tags(
-    context: dg.InitResourceContext | dg.OpExecutionContext | dg.AssetExecutionContext,
+    context: dg.InitResourceContext | AnyDagsterContext | StepHandlerContext | LaunchRunContext,
     extra_tags: dict[str, str] | None = None,
 ) -> dict[str, str]:
     """
@@ -29,9 +34,24 @@ def get_dagster_tags(
                 # inject the resource key used for this KubeRay resource
                 # this enables e.g reusing the same `RayCluster` across Dagster steps that require it without recreating the cluster
                 labels["dagster/resource-key"] = resource_key
-    else:
+
+    elif isinstance(context, StepHandlerContext):
+        labels.update(
+            **context.dagster_run.dagster_execution_info,
+        )
+    elif isinstance(context, LaunchRunContext):
+        labels.update(
+            **context.dagster_run.dagster_execution_info,
+        )
+    elif isinstance(context, dg.OpExecutionContext):
+        labels.update(
+            **context.run.dagster_execution_info,
+        )
+    elif isinstance(context, dg.AssetExecutionContext):
         labels.update(
             **context.run.dagster_execution_info,
         )
+    else:
+        raise ValueError(f"Unexpected context type: {type(context)}")
 
     return labels

src/dagster_ray/core/executor.py

@@ -1,5 +1,3 @@
-from __future__ import annotations
-
 from collections.abc import Iterator
 from typing import TYPE_CHECKING, Any, cast
 
@@ -23,14 +21,14 @@
 except ImportError:
     # for new versions of dagster > 1.11.6
     from dagster._core.remote_origin import RemoteJobOrigin  # pyright: ignore[reportMissingImports]
+
 from dagster._utils.merger import merge_dicts
 from packaging.version import Version
 from pydantic import Field
 
 from dagster_ray.configs import RayExecutionConfig, RayJobSubmissionClientConfig
 from dagster_ray.core.run_launcher import RayRunLauncher
-from dagster_ray.kuberay.utils import get_k8s_object_name
-from dagster_ray.utils import resolve_env_vars_list
+from dagster_ray.utils import get_k8s_object_name, resolve_env_vars_list
 
 if TYPE_CHECKING:
     from ray.job_submission import JobSubmissionClient
@@ -129,7 +127,7 @@ def name(self):
 
     def __init__(
         self,
-        client: JobSubmissionClient,
+        client: "JobSubmissionClient",
         env_vars: list[str] | None,
         runtime_env: dict[str, Any] | None,
         num_cpus: float | None,

src/dagster_ray/kuberay/__init__.py

@@ -1,4 +1,5 @@
 from dagster_ray.kuberay.configs import RayClusterConfig, RayClusterSpec, RayJobConfig, RayJobSpec
+from dagster_ray.kuberay.executor import kuberay_executor
 from dagster_ray.kuberay.jobs import cleanup_kuberay_clusters, delete_kuberay_clusters
 from dagster_ray.kuberay.ops import cleanup_kuberay_clusters_op, delete_kuberay_clusters_op
 from dagster_ray.kuberay.pipes import PipesKubeRayJobClient
@@ -8,6 +9,7 @@
     KubeRayInteractiveJob,
     KubeRayJobClientResource,
 )
+from dagster_ray.kuberay.run_launcher import KubeRayRunLauncher
 from dagster_ray.kuberay.schedules import cleanup_kuberay_clusters_daily
 
 __all__ = [
@@ -25,4 +27,6 @@
     "RayJobSpec",
     "KubeRayInteractiveJob",
     "KubeRayJobClientResource",
+    "kuberay_executor",
+    "KubeRayRunLauncher",
 ]

src/dagster_ray/kuberay/configs.py

@@ -164,13 +164,11 @@ def namespace(self) -> str:
 
     def to_k8s(
         self,
-        context: AnyDagsterContext,
         image: str | None = None,  # is injected into headgroup and workergroups, unless already specified there
         labels: Mapping[str, str] | None = None,
         annotations: Mapping[str, str] | None = None,
         env_vars: Mapping[str, str] | None = None,
     ) -> dict[str, Any]:
-        assert context.log is not None
         """Convert into Kubernetes manifests in camelCase format and inject additional information"""
 
         labels = labels or {}
@@ -186,7 +184,7 @@ def to_k8s(
                     "annotations": {**self.metadata.get("annotations", {}), **annotations},
                 }
             ),
-            "spec": self.spec.to_k8s(context=context, image=image, env_vars=env_vars),
+            "spec": self.spec.to_k8s(image=image, env_vars=env_vars),
         }
 
 
@@ -216,7 +214,6 @@ class RayJobSpec(dg.PermissiveConfig):
 
     def to_k8s(
         self,
-        context: AnyDagsterContext,
         image: str | None = None,  # is injected into headgroup and workergroups, unless already specified there
         env_vars: Mapping[str, str] | None = None,
     ) -> dict[str, Any]:
@@ -240,7 +237,7 @@ def to_k8s(
                 "ttlSecondsAfterFinished": self.ttl_seconds_after_finished,
                 "shutdownAfterJobFinishes": self.shutdown_after_job_finishes,
                 "suspend": self.suspend,
-                "rayClusterSpec": self.ray_cluster_spec.to_k8s(context=context, image=image, env_vars=env_vars)
+                "rayClusterSpec": self.ray_cluster_spec.to_k8s(image=image, env_vars=env_vars)
                 if self.ray_cluster_spec is not None
                 else None,
             }
@@ -262,7 +259,6 @@ def namespace(self) -> str:
 
     def to_k8s(
         self,
-        context: AnyDagsterContext,
         image: str | None = None,  # is injected into headgroup and workergroups, unless already specified there
         labels: Mapping[str, str] | None = None,
         annotations: Mapping[str, str] | None = None,
@@ -284,7 +280,6 @@ def to_k8s(
                 }
             ),
             "spec": self.spec.to_k8s(
-                context=context,
                 image=image,
                 env_vars=env_vars,
             ),