From 51a3672ba3d89861ff323308add27e5f78dcbaf2 Mon Sep 17 00:00:00 2001 From: danielgafni Date: Sat, 20 Sep 2025 00:49:21 +0300 Subject: [PATCH] :sparkles: add run launcher and executor for KubeRay --- docs/api/core.md | 4 +- docs/api/kuberay.md | 12 + justfile | 2 +- src/dagster_ray/_base/utils.py | 24 +- src/dagster_ray/core/executor.py | 8 +- src/dagster_ray/griffe_extensions.py | 428 +++++++++++++----- src/dagster_ray/kuberay/__init__.py | 4 + src/dagster_ray/kuberay/configs.py | 9 +- src/dagster_ray/kuberay/executor.py | 314 +++++++++++++ .../kuberay/resources/raycluster.py | 1 - src/dagster_ray/kuberay/resources/rayjob.py | 1 - src/dagster_ray/kuberay/run_launcher.py | 308 +++++++++++++ src/dagster_ray/utils.py | 23 + tests/kuberay/test_raycluster.py | 7 +- 14 files changed, 999 insertions(+), 146 deletions(-) create mode 100644 src/dagster_ray/kuberay/executor.py create mode 100644 src/dagster_ray/kuberay/run_launcher.py diff --git a/docs/api/core.md b/docs/api/core.md index 85b3fcb4..44f6a41f 100644 --- a/docs/api/core.md +++ b/docs/api/core.md @@ -26,7 +26,7 @@ The `LocalRay` can be used to connect to a local Ray cluster. ::: dagster_ray.core.run_launcher.RayRunLauncher options: - members: true + members: false --- @@ -34,7 +34,7 @@ The `LocalRay` can be used to connect to a local Ray cluster. ::: dagster_ray.core.executor.ray_executor options: - members: true + members: false --- diff --git a/docs/api/kuberay.md b/docs/api/kuberay.md index 23047f05..0c592f2e 100644 --- a/docs/api/kuberay.md +++ b/docs/api/kuberay.md @@ -4,6 +4,18 @@ KubeRay integration components for running Ray on Kubernetes. Learn how to use --- +## Run Launcher + +::: dagster_ray.kuberay.run_launcher.KubeRayRunLauncher + +--- + +# Executor + +::: dagster_ray.kuberay.executor.kuberay_executor + +--- + ## Client Mode Resources These resources initialize Ray client connection with a remote cluster. diff --git a/justfile b/justfile index 14ee13c6..5fd5469e 100644 --- a/justfile +++ b/justfile @@ -14,7 +14,7 @@ docs-build: uv run --group docs mkdocs build --clean --strict docs-serve: - uv run --group docs mkdocs serve + uv run --group docs mkdocs serve --clean docs-publish: uv run --group docs --all-extras mike deploy --push --update-aliases $(uv run dunamai from any --style pep440) diff --git a/src/dagster_ray/_base/utils.py b/src/dagster_ray/_base/utils.py index 173de651..67fce3f5 100644 --- a/src/dagster_ray/_base/utils.py +++ b/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 diff --git a/src/dagster_ray/core/executor.py b/src/dagster_ray/core/executor.py index bf976442..b9b2bf99 100644 --- a/src/dagster_ray/core/executor.py +++ b/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, diff --git a/src/dagster_ray/griffe_extensions.py b/src/dagster_ray/griffe_extensions.py index e8b2c2a2..e7a33bbb 100644 --- a/src/dagster_ray/griffe_extensions.py +++ b/src/dagster_ray/griffe_extensions.py @@ -2,7 +2,6 @@ from __future__ import annotations -import json from typing import Any import griffe @@ -35,8 +34,11 @@ def __init__(self, display_format: str = "json") -> None: def on_object(self, obj: Object, **kwargs: Any) -> None: """Process objects to extract configuration schemas.""" # Check if this is a dagster configurable object - if self._is_dagster_configurable(obj): + is_configurable = self._is_dagster_configurable(obj) + + if is_configurable: schema_info = self._extract_config_schema(obj) + if schema_info and schema_info.get("parameters"): # Add schema parameters as members self._add_schema_members(obj, schema_info) @@ -153,68 +155,174 @@ def _fallback_detection(self, obj: Object) -> bool: return False def _extract_config_schema(self, obj: Object) -> dict[str, Any] | None: - """Extract configuration schema from the object.""" + """Extract configuration schema from any configurable object.""" try: - # Use dynamic import to get live objects during doc build - if obj.name == "ray_executor": - return self._extract_ray_executor_schema() - elif obj.name == "RayRunLauncher": - return self._extract_ray_run_launcher_schema() + # Get the actual runtime object + runtime_obj = self._get_runtime_object(obj) + if runtime_obj is None: + return None + + # Try to extract schema using Dagster's config system + schema_dict = self._extract_schema_from_runtime_object(runtime_obj, obj) + if schema_dict: + return { + "title": f"Configuration Schema for {obj.name}", + "description": f"Configuration schema for {obj.name}", + "parameters": schema_dict, + } + except Exception as e: return {"title": f"Configuration Schema for {obj.name}", "error": f"Could not extract schema: {e}"} return None - def _extract_ray_executor_schema(self) -> dict[str, Any]: - """Extract schema for ray_executor.""" + def _extract_schema_from_runtime_object(self, runtime_obj, griffe_obj: Object) -> list[dict[str, Any]] | None: + """Extract schema from any Dagster configurable object using Dagster's official detection.""" + import inspect + try: - from dagster_ray.core.executor import _RAY_EXECUTOR_CONFIG_SCHEMA, RayExecutorConfig - - # Get the full Dagster config schema (includes ray + retries + tag_concurrency_limits) - dagster_schema = self._dagster_schema_to_dict(_RAY_EXECUTOR_CONFIG_SCHEMA) - - # Also get the JSON schema for the ray section - json_schema = RayExecutorConfig.model_json_schema() - - return { - "title": "Ray Executor Configuration", - "description": "Configuration schema for the ray_executor. This includes the main 'ray' section plus Dagster's built-in retry and concurrency configurations.", - "config_class": "RayExecutorConfig (ray section) + Dagster retries/concurrency", - "schema": json_schema, - "schema_json": json.dumps(json_schema, indent=2), - "schema_html": _highlight_json(json.dumps(json_schema, indent=2)), - "parameters": self._format_parameters(json_schema), - } - except ImportError as e: - return { - "title": "Ray Executor Configuration", - "description": "Configuration schema for the ray_executor", - "error": f"Could not import schema: {e}", - } + # Import Dagster's config detection utilities + from dagster import Field + from dagster._config.pythonic_config import ( + ConfigurableResource, + ConfigurableResourceFactory, + infer_schema_from_config_class, + ) + from dagster._core.definitions.configurable import ConfigurableDefinition + from dagster._serdes import ConfigurableClass - def _extract_ray_run_launcher_schema(self) -> dict[str, Any]: - """Extract schema for RayRunLauncher.""" + obj = runtime_obj + config_field = None + + # Follow Dagster's official detection pattern + if isinstance(obj, ConfigurableDefinition): + if obj.config_schema: + config_field = obj.config_schema.as_field() + elif inspect.isclass(obj) and ( + issubclass(obj, ConfigurableResource) or issubclass(obj, ConfigurableResourceFactory) + ): + config_field = infer_schema_from_config_class(obj) + elif isinstance(obj, type) and issubclass(obj, ConfigurableClass): + config_field = Field(obj.config_type()) + elif inspect.isfunction(obj): + # Handle function-based configurables (like @executor decorated functions) + try: + # Try to call with empty config to get the configurable object + instance = obj([]) + if hasattr(instance, "config_schema") and instance.config_schema: + config_field = instance.config_schema.as_field() + except: + pass + + if config_field: + # Convert the Dagster config field to our parameter format + return self._dagster_field_to_parameters(config_field) + + except Exception: + pass + + return None + + def _dagster_field_to_parameters(self, field) -> list[dict[str, Any]]: + """Convert a Dagster config field to parameter list using Dagster's approach.""" + parameters = [] + + def extract_field_info(field_obj, field_name=None, prefix=""): + """Recursively extract field information.""" + if field_name: + param_name = f"{prefix}.{field_name}" if prefix else field_name + + # Get type representation using Dagster's type_repr logic + type_str = self._dagster_type_repr(field_obj.config_type) + + param_data = { + "name": param_name, + "type": type_str, + "required": field_obj.is_required, + "description": field_obj.description or "", + } + + # Add default value if provided + if field_obj.default_provided: + param_data["default"] = field_obj.default_value + + parameters.append(param_data) + + # Recurse into subfields if they exist + if hasattr(field_obj.config_type, "fields") and field_obj.config_type.fields: + for name, subfield in field_obj.config_type.fields.items(): + subfield_prefix = f"{prefix}.{field_name}" if prefix and field_name else (field_name or "") + extract_field_info(subfield, name, subfield_prefix) + + # Start extraction + extract_field_info(field) + return parameters + + def _dagster_type_repr(self, config_type) -> str: + """Generate human-readable type representation following Dagster's approach.""" try: - from dagster_ray.core.run_launcher import RayLauncherConfig - - # Generate JSON schema - schema = RayLauncherConfig.model_json_schema() - - return { - "title": "Ray Run Launcher Configuration", - "description": "Configuration schema for the RayRunLauncher", - "config_class": "RayLauncherConfig", - "schema": schema, - "schema_json": json.dumps(schema, indent=2), - "schema_html": _highlight_json(json.dumps(schema, indent=2)), - "parameters": self._format_parameters(schema), - } - except ImportError: - return { - "title": "Ray Run Launcher Configuration", - "description": "Configuration schema for the RayRunLauncher", - "error": "Could not import RayLauncherConfig", - } + from dagster import BoolSource, IntSource, StringSource + from dagster._config.config_type import ( + ConfigTypeKind, + ) + + # Use given name if possible + if hasattr(config_type, "given_name") and config_type.given_name: + return config_type.given_name + + # Handle special source types + if config_type == StringSource: + return "str" + elif config_type == BoolSource: + return "bool" + elif config_type == IntSource: + return "int" + elif config_type.kind == ConfigTypeKind.ANY: + return "Any" + elif config_type.kind == ConfigTypeKind.SCALAR: + scalar_name = config_type.scalar_kind.name.lower() + # Fix common scalar type names + if scalar_name == "float": + return "float" + elif scalar_name == "int": + return "int" + elif scalar_name == "string": + return "str" + elif scalar_name == "bool": + return "bool" + return scalar_name + elif config_type.kind == ConfigTypeKind.ENUM: + values = ", ".join(str(val) for val in config_type.config_values) + return f"Enum[{values}]" + elif config_type.kind == ConfigTypeKind.ARRAY: + inner_type = self._dagster_type_repr(config_type.inner_type) + return f"list[{inner_type}]" + elif config_type.kind == ConfigTypeKind.SELECTOR: + # For selectors, show the available options if possible + if hasattr(config_type, "fields") and config_type.fields: + options = list(config_type.fields.keys()) + if len(options) <= 3: + return f"Literal[{', '.join(repr(opt) for opt in options)}]" + else: + return f"Literal[{', '.join(repr(opt) for opt in options[:3])}, ...]" + return "dict" # Fallback to dict since selectors are dict-like + elif config_type.kind == ConfigTypeKind.STRICT_SHAPE: + return "dict" + elif config_type.kind == ConfigTypeKind.PERMISSIVE_SHAPE: + return "dict" + elif config_type.kind == ConfigTypeKind.MAP: + return "dict" + elif config_type.kind == ConfigTypeKind.SCALAR_UNION: + scalar_type = self._dagster_type_repr(config_type.scalar_type) + non_scalar_type = self._dagster_type_repr(config_type.non_scalar_type) + return f"{scalar_type} | {non_scalar_type}" + elif config_type.kind == ConfigTypeKind.NONEABLE: + inner_type = self._dagster_type_repr(config_type.inner_type) + return f"{inner_type} | None" + else: + return "unknown" + except Exception: + return "unknown" def _format_parameters(self, schema: dict[str, Any]) -> list[dict[str, Any]]: """Format schema properties into a readable parameter list.""" @@ -474,71 +582,145 @@ def _extract_section_fields(self, fields: dict[str, Any], section_name: str, par self._extract_section_fields(field_info["nested_fields"], section_name, parameters, param_name) def _add_schema_members(self, obj: Object, schema_info: dict[str, Any]) -> None: - """Add configuration schema parameters as individual members of the object.""" - from griffe import Attribute, Kind - - # Add "Fields:" section to the main object's docstring (only once) - if schema_info["parameters"] and obj.docstring: - # Get existing docstring content - existing_content = obj.docstring.value if obj.docstring.value else "" - - # Only add Fields section if it doesn't already exist - if "Fields:" not in existing_content: - # Create Fields section as a vertical list with cross-references and types - fields_lines = ["", "Fields:", ""] - for param in schema_info["parameters"]: - field_name = param["name"] - field_type = self._format_type_with_crossrefs(param["type"]) - # Create cross-reference to the attribute and show type in parentheses - # Use bullet points to avoid code block formatting - fields_lines.append(f"- [{field_name}][{obj.path}.{field_name}] ({field_type})") - - fields_section = "\n".join(fields_lines) - - # Update the docstring - obj.docstring = griffe.Docstring(existing_content + fields_section) - - # Add each parameter as a direct member of the object - for param in schema_info["parameters"]: - param_attr = Attribute(name=param["name"], lineno=1, parent=obj) - param_attr.kind = Kind.ATTRIBUTE - param_attr.annotation = self._make_type_annotation_with_crossrefs(param["type"]) - - # Create clean docstring with just the description and default value - doc_parts = [] - - # Add description first - if param.get("description"): - doc_parts.append(param["description"]) - - # Only add default value if it exists (don't mention required/optional status) - if not param.get("required"): - default_val = param.get("default") - if default_val is not None: - doc_parts.append(f"Default: `{default_val}`") - - param_attr.docstring = griffe.Docstring("\n\n".join(doc_parts)) if doc_parts else None - - # Add as direct member of the main object - obj.members[param["name"]] = param_attr - - # Also add a JSON schema member if we have schema data for advanced users - if "schema" in schema_info: - schema_member = Attribute(name="__config_schema__", lineno=1, parent=obj) - schema_member.kind = Kind.ATTRIBUTE - schema_member.annotation = "dict" - - # Create docstring with foldable JSON schema - doc_parts = [] - doc_parts.append("") - doc_parts.append("
") - doc_parts.append("JSON Schema") - doc_parts.append("") - doc_parts.append("```json") - doc_parts.append(schema_info["schema_json"]) - doc_parts.append("```") - doc_parts.append("") - doc_parts.append("
") - - schema_member.docstring = griffe.Docstring("\n".join(doc_parts)) - obj.members["__config_schema__"] = schema_member + """Add configuration schema documentation to the object's docstring.""" + + if not schema_info["parameters"] or not obj.docstring: + return + + # Get existing docstring content + existing_content = obj.docstring.value if obj.docstring.value else "" + + # Only add configuration section if it doesn't already exist + # Look for the exact pattern we add: "Configuration:" as a standalone section + if "\nConfiguration:\n" not in existing_content: + # Create YAML-style nested configuration documentation + config_lines = ["", "Configuration:", "", "```yaml"] + + # Group parameters by their structure + structure = self._build_config_structure(schema_info["parameters"]) + + # Format as nested YAML-style structure + yaml_lines = self._format_yaml_structure(structure, 0) + config_lines.extend(yaml_lines) + + config_lines.append("```") + config_section = "\n".join(config_lines) + + # Update the docstring + obj.docstring = griffe.Docstring(existing_content + config_section) + + def _build_config_structure(self, parameters: list[dict[str, Any]]) -> dict: + """Build nested structure from flat parameter list.""" + structure = {} + + # First, identify ALL parameter paths that have nested children + has_children = set() + all_paths = [param["name"] for param in parameters] + + for path in all_paths: + # Check if any other path starts with this path + "." + for other_path in all_paths: + if other_path.startswith(path + "."): + has_children.add(path) + break + + for param in parameters: + param_name = param["name"] + parts = param_name.split(".") + + if param_name in has_children: + # This parameter has nested children, skip it - we'll show the nested structure instead + continue + + # Build the nested structure + current = structure + for i, part in enumerate(parts[:-1]): + if part not in current: + current[part] = {} + current = current[part] + + # Add the final parameter + final_key = parts[-1] + current[final_key] = param + + return structure + + def _format_yaml_structure(self, structure: dict, indent_level: int) -> list[str]: + """Format the structure as proper YAML with type and description comments.""" + lines = [] + indent = " " * indent_level + + for key, value in structure.items(): + if isinstance(value, dict) and "type" in value: + # This is a leaf parameter - show it with its value + param_type = value["type"] + required = value.get("required", True) + default = value.get("default") + description = value.get("description", "") + + # Format as YAML with type and description as comment (keep description intact on one line) + comment_parts = [] + comment_parts.append(f"({param_type})") + if description: + # Keep description intact but on one line + clean_desc = description.replace("\n", " ").strip() + comment_parts.append(clean_desc) + + comment = " # " + " - ".join(comment_parts) if comment_parts else "" + + # Show proper YAML value + if default is not None: + if isinstance(default, str): + yaml_value = f'"{default}"' + elif isinstance(default, bool): + yaml_value = str(default).lower() + elif default == {}: + yaml_value = "{}" + elif default == []: + yaml_value = "[]" + elif isinstance(default, dict): + yaml_value = "# has default configuration" + elif isinstance(default, list): + yaml_value = "# has default list" + else: + yaml_value = str(default) + elif not required: + yaml_value = "null # optional" + else: + yaml_value = "# required" + + lines.append(f"{indent}{key}: {yaml_value}{comment}") + + elif isinstance(value, dict): + # This is a nested section - always expand it + lines.append(f"{indent}{key}:") + lines.extend(self._format_yaml_structure(value, indent_level + 1)) + + return lines + + def _format_config_parameter_simple(self, param: dict[str, Any]) -> list[str]: + """Format a single configuration parameter for documentation.""" + lines = [] + + # Parameter name and type + param_name = param["name"] + param_type = param["type"] + optional_indicator = " (optional)" if not param.get("required") else "" + + lines.append(f"#### {param_name}{optional_indicator}") + lines.append("") + lines.append(f"Type: {param_type}") + lines.append("") + + # Description + if param.get("description"): + lines.append(param["description"]) + lines.append("") + + # Default value + if not param.get("required") and param.get("default") is not None: + default_val = param["default"] + lines.append(f"Default: `{default_val}`") + lines.append("") + + return lines diff --git a/src/dagster_ray/kuberay/__init__.py b/src/dagster_ray/kuberay/__init__.py index edd547f4..86c70810 100644 --- a/src/dagster_ray/kuberay/__init__.py +++ b/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", ] diff --git a/src/dagster_ray/kuberay/configs.py b/src/dagster_ray/kuberay/configs.py index 5375c1bd..0167ff1c 100644 --- a/src/dagster_ray/kuberay/configs.py +++ b/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, ), diff --git a/src/dagster_ray/kuberay/executor.py b/src/dagster_ray/kuberay/executor.py new file mode 100644 index 00000000..733c3448 --- /dev/null +++ b/src/dagster_ray/kuberay/executor.py @@ -0,0 +1,314 @@ +from collections.abc import Iterator, Mapping +from typing import Any, cast + +import dagster as dg +from dagster._core.definitions.executor_definition import multiple_process_executor_requirements +from dagster._core.definitions.metadata import MetadataValue +from dagster._core.events import DagsterEvent, EngineEventData +from dagster._core.execution.retries import RetryMode, get_retries_config +from dagster._core.execution.tags import get_tag_concurrency_limits_config +from dagster._core.executor.base import Executor +from dagster._core.executor.init import InitExecutorContext +from dagster._core.executor.step_delegating import ( + CheckStepHealthResult, + StepDelegatingExecutor, + StepHandler, + StepHandlerContext, +) + +from dagster_ray.kuberay.configs import RayJobConfig + +try: + from dagster._core.remote_representation.origin import RemoteJobOrigin # noqa: F401 +except ImportError: + # for new versions of dagster > 1.11.6 + from dagster._core.remote_origin import RemoteJobOrigin # noqa: F401 # pyright: ignore[reportMissingImports] + +from dagster._utils.merger import merge_dicts +from pydantic import Field + +from dagster_ray.configs import RayExecutionConfig +from dagster_ray.kuberay.client import RayJobClient +from dagster_ray.kuberay.client.base import load_kubeconfig +from dagster_ray.kuberay.resources.base import BaseKubeRayResourceConfig +from dagster_ray.kuberay.utils import normalize_k8s_label_values +from dagster_ray.utils import get_k8s_object_name + + +class KubeRayExecutorConfig(BaseKubeRayResourceConfig, RayExecutionConfig): + """Configuration for the KubeRay executor.""" + + ray_job: RayJobConfig = Field( + default_factory=RayJobConfig, + description="Configuration for the Kubernetes `RayJob` CR", + ) + + kube_context: str | None = Field( + default=None, + description="Kubernetes context to use. If not specified, uses the current context.", + ) + + kube_config: str | None = Field( + default=None, + description="Path to the Kubernetes config file. If not specified, uses the default config.", + ) + + log_cluster_conditions: bool = Field( + default=True, + description="Whether to log `RayCluster` conditions while waiting for the RayCluster to become ready. Learn more: [KubeRay docs](https://docs.ray.io/en/latest/cluster/kubernetes/user-guides/observability.html#raycluster-status-conditions).", + ) + + timeout: float = Field( + default=600.0, + description="Timeout for various Kubernetes operations in seconds.", + ) + + +_KUBERAY_CONFIG_SCHEMA = KubeRayExecutorConfig.to_config_schema().as_field() + +_KUBERAY_EXECUTOR_CONFIG_SCHEMA = merge_dicts( + {"kuberay": _KUBERAY_CONFIG_SCHEMA}, # type: ignore + {"retries": get_retries_config(), "tag_concurrency_limits": get_tag_concurrency_limits_config()}, +) + + +@dg.executor( + name="kuberay", + config_schema=_KUBERAY_EXECUTOR_CONFIG_SCHEMA, + requirements=multiple_process_executor_requirements(), +) +def kuberay_executor(init_context: InitExecutorContext) -> Executor: + """Executes steps by submitting them as KubeRay jobs. + + Each step is executed as a separate Ray Job on a Kubernetes cluster using the KubeRay operator. + This executor provides automatic cluster management and cleanup through Kubernetes native resources. + + Example: + Use `kuberay_executor` for the entire code location + ```python + import dagster as dg + from dagster_ray import kuberay_executor + + kuberay_executor = kuberay_executor.configured( + { + "kuberay": { + "ray_job": { + "spec": { + "rayClusterSpec": { + "headGroupSpec": { + "template": { + "spec": { + "containers": [ + { + "name": "ray-head", + "image": "rayproject/ray:2.0.0", + } + ] + } + } + } + } + } + } + } + } + ) + + defs = dg.Definitions(..., executor=kuberay_executor) + ``` + + Example: + Override configuration for a specific asset + ```python + import dagster as dg + + @dg.asset( + op_tags={"dagster-ray/config": {"num_cpus": 2}} + ) + def my_asset(): ... + ``` + """ + exc_cfg = init_context.executor_config + kuberay_cfg = KubeRayExecutorConfig(**exc_cfg["kuberay"]) # type: ignore + + load_kubeconfig(context=kuberay_cfg.kube_context, config_file=kuberay_cfg.kube_config) + client = RayJobClient(kube_context=kuberay_cfg.kube_context, kube_config=kuberay_cfg.kube_config) + + return StepDelegatingExecutor( + KubeRayStepHandler( + client=client, + config=kuberay_cfg, + ), + retries=RetryMode.from_config(exc_cfg["retries"]), # type: ignore + max_concurrent=dg._check.opt_int_elem(exc_cfg, "max_concurrent"), + tag_concurrency_limits=dg._check.opt_list_elem(exc_cfg, "tag_concurrency_limits"), + should_verify_step=True, + ) + + +class KubeRayStepHandler(StepHandler): + @property + def name(self): + return "KubeRayStepHandler" + + def __init__( + self, + client: RayJobClient, + config: KubeRayExecutorConfig, + ): + super().__init__() + self.client = client + self.config = config + + def _get_step_key(self, step_handler_context: StepHandlerContext) -> str: + step_keys_to_execute = cast(list[str], step_handler_context.execute_step_args.step_keys_to_execute) + assert len(step_keys_to_execute) == 1, "Launching multiple steps is not currently supported" + return step_keys_to_execute[0] + + def _get_ray_job_name(self, step_handler_context: StepHandlerContext) -> str: + step_key = self._get_step_key(step_handler_context) + + name_key = get_k8s_object_name( + step_handler_context.execute_step_args.run_id, + step_key, + ) + + if step_handler_context.execute_step_args.known_state: + retry_state = step_handler_context.execute_step_args.known_state.get_retry_state() + if retry_state.get_attempt_count(step_key): + return f"dagster-step-{name_key}-{retry_state.get_attempt_count(step_key)}" + + return f"dagster-step-{name_key}" + + def _get_dagster_tags(self, step_handler_context: StepHandlerContext) -> Mapping[str, str]: + """Get standardized Dagster tags for labeling Kubernetes resources.""" + return step_handler_context.dagster_run.dagster_execution_info + + def _create_ray_job_spec(self, step_handler_context: StepHandlerContext) -> dict[str, Any]: + """Create the RayJob specification for the step.""" + run = step_handler_context.dagster_run + step_key = self._get_step_key(step_handler_context) + + # Get user-provided step configuration - allow full RayJobConfig override + user_provided_config = {} + if "dagster-ray/config" in step_handler_context.step_tags[step_key]: + import json + + user_provided_config = json.loads(step_handler_context.step_tags[step_key]["dagster-ray/config"]) + + # Create a merged config by overlaying user config on base config + merged_ray_job_config = self._merge_ray_job_configs(self.config.ray_job, user_provided_config) + + # Create the RayJob spec + ray_job_name = self._get_ray_job_name(step_handler_context) + + ray_job_spec = merged_ray_job_config.to_k8s( + image=(self.config.image or run.tags.get("dagster/image")), + labels=normalize_k8s_label_values({**self._get_dagster_tags(step_handler_context)}), + env_vars=self._get_env_vars_to_inject(step_handler_context), + ) + + ray_job_spec["metadata"]["name"] = ray_job_name + + # Add the Dagster command to the entrypoint + command_args = step_handler_context.execute_step_args.get_command_args(skip_serialized_namedtuple=True) + ray_job_spec["spec"]["entrypoint"] = " ".join(command_args) + + return ray_job_spec + + def _merge_ray_job_configs(self, base_config: RayJobConfig, user_config: dict[str, Any]) -> RayJobConfig: + """Merge user-provided configuration with base configuration. + + This follows the same pattern as dagster-k8s where users can override + any field in the configuration via the dagster-ray/config tag. + """ + from dagster._utils.merger import deep_merge_dicts + + # Convert base config to dict + base_dict = base_config.model_dump() + + # Deep merge user config over base config + merged_dict = deep_merge_dicts(base_dict, user_config) + + # Create new config instance with merged values + return RayJobConfig(**merged_dict) # pyright: ignore[reportArgumentType] + + def _get_env_vars_to_inject(self, step_handler_context: StepHandlerContext) -> dict[str, str]: + """Get environment variables to inject into the Ray job.""" + run = step_handler_context.dagster_run + step_key = self._get_step_key(step_handler_context) + + env_vars = { + "DAGSTER_RUN_JOB_NAME": run.job_name, + "DAGSTER_RUN_STEP_KEY": step_key, + **{env["name"]: env["value"] for env in step_handler_context.execute_step_args.get_command_env()}, + } + + return env_vars + + def launch_step(self, step_handler_context: StepHandlerContext) -> Iterator[DagsterEvent]: + step_key = self._get_step_key(step_handler_context) + ray_job_name = self._get_ray_job_name(step_handler_context) + namespace = self.config.ray_job.namespace + + yield DagsterEvent.step_worker_starting( + step_handler_context.get_step_context(step_key), + message=f'Executing step "{step_key}" in KubeRay job {namespace}/{ray_job_name}.', + metadata={ + "RayJob Name": MetadataValue.text(ray_job_name), + "Namespace": MetadataValue.text(namespace), + }, + ) + + ray_job_spec = self._create_ray_job_spec(step_handler_context) + + self.client.create( + body=ray_job_spec, + namespace=namespace, + ) + + # Wait for the job to start running + self.client.wait_until_running( + name=ray_job_name, + namespace=namespace, + timeout=self.config.timeout, + poll_interval=self.config.poll_interval, + terminate_on_timeout=True, + port_forward=False, + log_cluster_conditions=self.config.log_cluster_conditions, + ) + + def check_step_health(self, step_handler_context: StepHandlerContext) -> CheckStepHealthResult: + step_key = self._get_step_key(step_handler_context) + ray_job_name = self._get_ray_job_name(step_handler_context) + namespace = self.config.ray_job.namespace + + try: + status = self.client.get_status(name=ray_job_name, namespace=namespace) + except Exception as e: + return CheckStepHealthResult.unhealthy( + reason=f"KubeRay job {namespace}/{ray_job_name} for step {step_key} could not be found: {e}" + ) + + job_status = status.get("jobStatus") + + if job_status in ["FAILED", "STOPPED"]: + message = status.get("message", "No message provided") + return CheckStepHealthResult.unhealthy( + reason=f"Discovered failed KubeRay job {namespace}/{ray_job_name} for step {step_key}. Status: {job_status}. Message: {message}" + ) + + return CheckStepHealthResult.healthy() + + def terminate_step(self, step_handler_context: StepHandlerContext) -> Iterator[DagsterEvent]: + step_key = self._get_step_key(step_handler_context) + ray_job_name = self._get_ray_job_name(step_handler_context) + namespace = self.config.ray_job.namespace + + yield DagsterEvent.engine_event( + step_handler_context.get_step_context(step_key), + message=f"Stopping KubeRay job {namespace}/{ray_job_name} for step", + event_specific_data=EngineEventData(), + ) + + self.client.terminate(name=ray_job_name, namespace=namespace, port_forward=False) diff --git a/src/dagster_ray/kuberay/resources/raycluster.py b/src/dagster_ray/kuberay/resources/raycluster.py index b3616a27..53a11aec 100644 --- a/src/dagster_ray/kuberay/resources/raycluster.py +++ b/src/dagster_ray/kuberay/resources/raycluster.py @@ -99,7 +99,6 @@ def create(self, context: AnyDagsterContext): namespace=self.namespace, ): k8s_manifest = self.ray_cluster.to_k8s( - context, image=(self.image or context.dagster_run.tags.get("dagster/image")), labels=normalize_k8s_label_values(self.get_dagster_tags(context)), env_vars=self.get_env_vars_to_inject(), diff --git a/src/dagster_ray/kuberay/resources/rayjob.py b/src/dagster_ray/kuberay/resources/rayjob.py index 896935fa..96a7256c 100644 --- a/src/dagster_ray/kuberay/resources/rayjob.py +++ b/src/dagster_ray/kuberay/resources/rayjob.py @@ -126,7 +126,6 @@ def create(self, context: AnyDagsterContext): self._name = self.ray_job.metadata.get("name") or self._get_step_name(context) k8s_manifest = self.ray_job.to_k8s( - context, image=(self.image or context.dagster_run.tags.get("dagster/image")), labels=normalize_k8s_label_values(self.get_dagster_tags(context)), ) diff --git a/src/dagster_ray/kuberay/run_launcher.py b/src/dagster_ray/kuberay/run_launcher.py new file mode 100644 index 00000000..50b1b058 --- /dev/null +++ b/src/dagster_ray/kuberay/run_launcher.py @@ -0,0 +1,308 @@ +from collections.abc import Mapping +from typing import Any + +import dagster as dg +from dagster._config import UserConfigSchema +from dagster._core.launcher.base import LaunchRunContext, RunLauncher, WorkerStatus +from dagster._core.storage.dagster_run import DagsterRun +from dagster._core.storage.tags import DOCKER_IMAGE_TAG +from dagster._grpc.types import ExecuteRunArgs +from dagster._serdes import ConfigurableClass +from pydantic import Field + +from dagster_ray._base.utils import get_dagster_tags +from dagster_ray.configs import RayExecutionConfig +from dagster_ray.kuberay.client import RayJobClient +from dagster_ray.kuberay.client.base import load_kubeconfig +from dagster_ray.kuberay.configs import RayJobConfig +from dagster_ray.kuberay.resources.base import BaseKubeRayResourceConfig +from dagster_ray.kuberay.utils import normalize_k8s_label_values + +try: + from dagster._core.remote_representation.origin import RemoteJobOrigin +except ImportError: + # for new versions of dagster > 1.11.6 + from dagster._core.remote_origin import RemoteJobOrigin # pyright: ignore[reportMissingImports] # noqa: F401 + + +def get_ray_job_id_from_run_id(run_id: str, resume_attempt_number=None): + """Generate a KubeRay job name from Dagster run ID.""" + return f"dagster-run-{run_id}" + ("" if not resume_attempt_number else f"-{resume_attempt_number}") + + +class KubeRayRunLauncherConfig(BaseKubeRayResourceConfig, RayExecutionConfig): + """Configuration for the KubeRay run launcher.""" + + ray_job: RayJobConfig = Field( + default_factory=RayJobConfig, + description="Configuration for the Kubernetes `RayJob` CR", + ) + + kube_context: str | None = Field( + default=None, + description="Kubernetes context to use. If not specified, uses the current context.", + ) + + kube_config: str | None = Field( + default=None, + description="Path to the Kubernetes config file. If not specified, uses the default config.", + ) + + log_cluster_conditions: bool = Field( + default=True, + description="Whether to log `RayCluster` conditions while waiting for the RayCluster to become ready. Learn more: [KubeRay docs](https://docs.ray.io/en/latest/cluster/kubernetes/user-guides/observability.html#raycluster-status-conditions).", + ) + + timeout: float = Field( + default=600.0, + description="Timeout for various Kubernetes operations in seconds.", + ) + + @property + def namespace(self) -> str: + return self.ray_job.namespace + + +class KubeRayRunLauncher(RunLauncher, ConfigurableClass): + """RunLauncher that submits Dagster runs as KubeRay jobs to a Kubernetes cluster. + + Each Dagster run is executed as a separate Ray Job on Kubernetes using the KubeRay operator. + This provides automatic cluster management and cleanup through Kubernetes native resources. + + Configuration can be provided via `dagster.yaml` and individual runs can override + settings using the `dagster-ray/config` tag. + + Example: + Configure via `dagster.yaml` + ```yaml + run_launcher: + module: dagster_ray.kuberay + class: KubeRayRunLauncher + config: + ray_job: + spec: + rayClusterSpec: + headGroupSpec: + template: + spec: + containers: + - name: ray-head + image: "rayproject/ray:2.0.0" + ``` + + Example: + Override settings per job + ```python + import dagster as dg + + @dg.job( + tags={ + "dagster-ray/config": { + "num_cpus": 16, + "num_gpus": 1, + "runtime_env": {"pip": {"packages": ["torch"]}}, + } + } + ) + def my_job(): + return my_op() + ``` + """ + + def __init__(self, **config): + self._config = KubeRayRunLauncherConfig(**config) + + # Initialize Kubernetes configuration + load_kubeconfig(context=self._config.kube_context, config_file=self._config.kube_config) + self._client = RayJobClient(kube_context=self._config.kube_context, kube_config=self._config.kube_config) + + super().__init__() + + @classmethod + def config_type(cls) -> UserConfigSchema: + return {"kuberay": KubeRayRunLauncherConfig.to_config_schema().as_field()} + + @classmethod + def from_config_value(cls, inst_data, config_value): + return cls(**config_value) + + def _get_dagster_tags(self, context: LaunchRunContext) -> Mapping[str, str]: + """Get standardized Dagster tags for labeling Kubernetes resources.""" + return get_dagster_tags(context) + + def _get_env_vars_to_inject(self, dagster_run: DagsterRun) -> dict[str, str]: + """Get environment variables to inject into the Ray job.""" + env_vars = { + "DAGSTER_RUN_JOB_NAME": dagster_run.job_name or "unknown", + } + + return env_vars + + def _create_ray_job_spec(self, context: LaunchRunContext, command_args: list[str]) -> dict[str, Any]: + """Create the RayJob specification for the run.""" + dagster_run = context.dagster_run + + # Get user-provided run configuration - allow full RayJobConfig override + user_provided_config = {} + if "dagster-ray/config" in dagster_run.tags: + import json + + user_provided_config = json.loads(dagster_run.tags["dagster-ray/config"]) + + # Create a merged config by overlaying user config on base config + merged_ray_job_config = self._merge_ray_job_configs(self._config.ray_job, user_provided_config) + + # Generate job name + ray_job_name = get_ray_job_id_from_run_id(dagster_run.run_id) + + # Create the RayJob spec with merged configuration + ray_job_spec = merged_ray_job_config.to_k8s( + image=(self._config.image or dagster_run.tags.get(DOCKER_IMAGE_TAG)), + labels=normalize_k8s_label_values({**self._get_dagster_tags(context)}), + env_vars=self._get_env_vars_to_inject(dagster_run), + ) + + ray_job_spec["metadata"]["name"] = ray_job_name + + # Add the Dagster command to the entrypoint + ray_job_spec["spec"]["entrypoint"] = " ".join(command_args) + + return ray_job_spec + + def _merge_ray_job_configs(self, base_config: RayJobConfig, user_config: dict[str, Any]) -> RayJobConfig: + """Merge user-provided configuration with base configuration. + + This follows the same pattern as dagster-k8s where users can override + any field in the configuration via the dagster-ray/config tag. + """ + from dagster._utils.merger import deep_merge_dicts + + # Convert base config to dict + base_dict = base_config.model_dump() + + # Deep merge user config over base config + merged_dict = deep_merge_dicts(base_dict, user_config) + + # Create new config instance with merged values + return RayJobConfig(**merged_dict) # pyright: ignore[reportArgumentType] + + def launch_run(self, context: LaunchRunContext) -> None: + """Launch a Dagster run as a KubeRay job.""" + dagster_run = context.dagster_run + ray_job_name = get_ray_job_id_from_run_id(dagster_run.run_id) + namespace = self._config.namespace + job_origin = dg._check.not_none(dagster_run.job_code_origin) + + args = list( + ExecuteRunArgs( + job_origin=job_origin, + run_id=dagster_run.run_id, + instance_ref=self._instance.get_ref(), + set_exit_code_on_failure=True, + ).get_command_args() + ) + + # wrap the json in quotes to prevent errors with shell commands + args[-1] = "'" + args[-1] + "'" + + ray_job_spec = self._create_ray_job_spec(context, args) + + self._instance.report_engine_event( + f"Launching run {dagster_run.run_id} as KubeRay job {namespace}/{ray_job_name}", + dagster_run, + cls=self.__class__, + ) + + try: + # Create the RayJob + self._client.create( + body=ray_job_spec, + namespace=namespace, + ) + + # Wait for the job to start running + self._client.wait_until_running( + name=ray_job_name, + namespace=namespace, + timeout=self._config.timeout, + poll_interval=self._config.poll_interval, + terminate_on_timeout=False, # Don't terminate on timeout for runs + port_forward=False, + log_cluster_conditions=self._config.log_cluster_conditions, + ) + + self._instance.report_engine_event( + f"KubeRay job {namespace}/{ray_job_name} started successfully", + dagster_run, + cls=self.__class__, + ) + + except Exception as e: + self._instance.report_engine_event( + f"Failed to launch KubeRay job {namespace}/{ray_job_name}: {str(e)}", + dagster_run, + cls=self.__class__, + ) + raise + + def can_terminate(self, run_id: str) -> bool: + """Check if a run can be terminated.""" + return True + + def terminate(self, run_id: str) -> bool: + """Terminate a running KubeRay job.""" + ray_job_name = get_ray_job_id_from_run_id(run_id) + namespace = self._config.ray_job.namespace + + try: + self._client.terminate( + name=ray_job_name, + namespace=namespace, + port_forward=False, + ) + return True + except Exception: + # If termination fails, the job might already be completed or not exist + return False + + def get_run_worker_debug_info(self, run: DagsterRun, include_container_logs: bool | None = True) -> str | None: + """Get debug information for a run.""" + ray_job_name = get_ray_job_id_from_run_id(run.run_id) + namespace = self._config.ray_job.namespace + + try: + # Try to get the RayJob status + ray_job = self._client.get(name=ray_job_name, namespace=namespace) + if ray_job: + status_info = f"RayJob {namespace}/{ray_job_name}:\n" + status_info += f"Status: {ray_job.get('status', {})}\n" + if include_container_logs: + # Could potentially include container logs here if needed + pass + return status_info + except Exception: + pass + + return None + + def get_run_worker_status(self, run_id: str) -> WorkerStatus: + """Get the status of a run worker.""" + ray_job_name = get_ray_job_id_from_run_id(run_id) + namespace = self._config.ray_job.namespace + + try: + status = self._client.get_status(name=ray_job_name, namespace=namespace, timeout=5.0) + job_status = status.get("jobStatus") + + if job_status == "RUNNING": + return WorkerStatus.RUNNING + elif job_status == "SUCCEEDED": + return WorkerStatus.SUCCESS + elif job_status in ["FAILED", "STOPPED"]: + return WorkerStatus.FAILED + else: + return WorkerStatus.UNKNOWN + except TimeoutError: + return WorkerStatus.NOT_FOUND + except Exception: + return WorkerStatus.UNKNOWN diff --git a/src/dagster_ray/utils.py b/src/dagster_ray/utils.py index 17dfed1f..a597de99 100644 --- a/src/dagster_ray/utils.py +++ b/src/dagster_ray/utils.py @@ -1,10 +1,14 @@ from __future__ import annotations +import hashlib import os +import random +import string from collections.abc import Mapping, Sequence from typing import Any import dagster as dg +import dagster._check as check from dagster._config.field_utils import IntEnvVar # yes, `python-client` is actually the KubeRay package name @@ -47,6 +51,25 @@ def _process_dagster_env_vars(config: Any) -> Any: return config +def get_k8s_object_name(run_id: str, step_key: str | None = None) -> str: + """Creates a unique (short!) identifier to name k8s objects based on run ID and step key(s). + + K8s Job names are limited to 63 characters, because they are used as labels. For more info, see: + + https://kubernetes.io/docs/concepts/overview/working-with-objects/names/ + """ + check.str_param(run_id, "run_id") + check.opt_str_param(step_key, "step_key") + if not step_key: + letters = string.ascii_lowercase + step_key = "".join(random.choice(letters) for i in range(20)) + + # Creates 32-bit signed int, so could be negative + name_hash = hashlib.md5((run_id + step_key).encode("utf-8")) + + return name_hash.hexdigest()[:8] + + def get_current_job_id() -> str: import ray diff --git a/tests/kuberay/test_raycluster.py b/tests/kuberay/test_raycluster.py index 97909cd1..e59e72e0 100644 --- a/tests/kuberay/test_raycluster.py +++ b/tests/kuberay/test_raycluster.py @@ -279,10 +279,9 @@ def my_asset(ray_cluster: RayResource) -> None: def test_ray_cluster_builder_debug(): kuberay_cluster = KubeRayCluster(enable_debug_post_mortem=True, image="test") kuberay_cluster._cluster_name = "test-cluster" - context = dg.build_init_resource_context() ray_cluster_config = kuberay_cluster.ray_cluster.to_k8s( - context, env_vars=kuberay_cluster.get_env_vars_to_inject(), labels={"foo": "bar"}, image="test" + env_vars=kuberay_cluster.get_env_vars_to_inject(), labels={"foo": "bar"}, image="test" ) assert ray_cluster_config["metadata"]["labels"]["foo"] == "bar" for group_spec in [ray_cluster_config["spec"]["headGroupSpec"], *ray_cluster_config["spec"]["workerGroupSpecs"]]: @@ -292,7 +291,7 @@ def test_ray_cluster_builder_debug(): kuberay_cluster = KubeRayCluster(enable_tracing=True) kuberay_cluster._cluster_name = "test-cluster" ray_cluster_config = kuberay_cluster.ray_cluster.to_k8s( - context, env_vars=kuberay_cluster.get_env_vars_to_inject(), image="test" + env_vars=kuberay_cluster.get_env_vars_to_inject(), image="test" ) for group_spec in [ray_cluster_config["spec"]["headGroupSpec"], *ray_cluster_config["spec"]["workerGroupSpecs"]]: for container in group_spec["template"]["spec"]["containers"]: @@ -301,7 +300,7 @@ def test_ray_cluster_builder_debug(): kuberay_cluster = KubeRayCluster(enable_actor_task_logging=True) kuberay_cluster._cluster_name = "test-cluster" ray_cluster_config = kuberay_cluster.ray_cluster.to_k8s( - context, env_vars=kuberay_cluster.get_env_vars_to_inject(), image="test" + env_vars=kuberay_cluster.get_env_vars_to_inject(), image="test" ) kuberay_cluster._cluster_name = "test-cluster" for group_spec in [ray_cluster_config["spec"]["headGroupSpec"], *ray_cluster_config["spec"]["workerGroupSpecs"]]: