configs.py

from __future__ import annotations

from collections.abc import Mapping
from typing import Any, Literal

import dagster as dg
from packaging.version import Version
from pydantic import Field, model_validator

USER_DEFINED_RAY_KEY = "dagster-ray/config"
DAGSTER_RAY_NAMESPACES_ENV_VAR = "DAGSTER_RAY_NAMESPACES"
DAGSTER_RAY_NAMESPACES_DEFAULT_VALUE = "ray"
DAGSTER_RAY_CLUSTER_EXPIRATION_SECONDS_ENV_VAR = "DAGSTER_RAY_CLUSTER_EXPIRATION_SECONDS"
DAGSTER_RAY_CLUSTER_EXPIRATION_SECONDS_DEFAULT_VALUE = str(4 * 60 * 60)


class Lifecycle(dg.Config):
    create: bool = Field(
        default=True,
        description="Whether to create the resource. If set to `False`, the user can manually call `.create` instead.",
    )
    wait: bool = Field(
        default=True,
        description="Whether to wait for the remote Ray cluster to become ready to accept connections. If set to `False`, the user can manually call `.wait` instead.",
    )
    connect: bool = Field(
        default=True,
        description="Whether to run `ray.init` against the remote Ray cluster. If set to `False`, the user can manually call `.connect` instead.",
    )
    cleanup: Literal["never", "always", "on_exception"] = Field(
        default="always",
        description="Resource cleanup policy. Determines when the resource should be deleted after Dagster step execution or during interruption.",
    )


class RayExecutionConfig(dg.Config):
    runtime_env: dict[str, Any] | None = Field(default=None, description="The runtime environment to use.")
    num_cpus: float | None = Field(default=None, description="The number of CPUs to allocate.")
    num_gpus: float | None = Field(default=None, description="The number of GPUs to allocate.")
    memory: int | None = Field(default=None, description="The amount of memory in bytes to allocate.")
    resources: dict[str, float] | None = Field(default=None, description="Custom resources to allocate.")

    @classmethod
    def from_tags(cls, tags: Mapping[str, str]) -> RayExecutionConfig:
        if USER_DEFINED_RAY_KEY in tags:
            return cls.model_validate_json(tags[USER_DEFINED_RAY_KEY])
        else:
            return cls()


class RayJobSubmissionClientConfig(dg.Config):
    address: str | None = Field(default=None, description="The address of the Ray cluster to connect to.")
    metadata: dict[str, str] | None = Field(
        default=None,
        description="""Arbitrary metadata to store along with all jobs. New metadata
            specified per job will be merged with the global metadata provided here
            via a simple dict update.""",
    )
    headers: dict[str, Any] | None = Field(
        default=None,
        description="""Headers to use when sending requests to the HTTP job server, used
            for cases like authentication to a remote cluster.""",
    )
    cookies: dict[str, Any] | None = Field(
        default=None, description="Cookies to use when sending requests to the HTTP job server."
    )


class ExecutionOptionsConfig(dg.Config):
    cpu: int | None = None
    gpu: int | None = None
    object_store_memory: int | None = None


class RayDataExecutionOptions(dg.Config):
    execution_options: ExecutionOptionsConfig = Field(default_factory=ExecutionOptionsConfig)

    cpu_limit: int = 5000
    gpu_limit: int = 0
    verbose_progress: bool = True
    use_polars_sort: bool = Field(
        default=True,
        description="Whether to use Polars for sort operations in Ray Data. Forwarded to `ray.data.DatasetContext.use_polars_sort`.",
    )
    use_polars: bool | None = Field(
        default=None,
        description="Deprecated. Use `use_polars_sort` instead. Ray 2.55 renamed the underlying `DatasetContext` attribute to `use_polars_sort`.",
        deprecated="`use_polars` is deprecated; use `use_polars_sort` instead.",
    )

    @model_validator(mode="before")
    @classmethod
    def _forward_deprecated_use_polars(cls, data: Any) -> Any:
        if not isinstance(data, dict):
            return data
        legacy = data.get("use_polars")
        new = data.get("use_polars_sort")
        if legacy is not None and new is not None and legacy != new:
            raise ValueError(
                "`use_polars` and `use_polars_sort` were set to contradicting values "
                f"({legacy!r} vs {new!r}). `use_polars` is deprecated; set only "
                "`use_polars_sort`."
            )
        if legacy is not None and new is None:
            data = {**data, "use_polars_sort": legacy}
        return data

    def apply(self):
        import ray
        from ray.data import ExecutionResources

        ctx = ray.data.DatasetContext.get_current()

        ctx.execution_options.resource_limits = ExecutionResources.for_limits(
            cpu=self.execution_options.cpu,
            gpu=self.execution_options.gpu,
            object_store_memory=self.execution_options.object_store_memory,
        )

        ctx.verbose_progress = self.verbose_progress
        # Ray >=2.50 renamed the sort-time Polars toggle from `use_polars` to
        # `use_polars_sort`; read the effective value (the legacy field wins when
        # both are set only because the validator has already unified them).
        effective = self.use_polars if self.use_polars is not None else self.use_polars_sort
        if Version(ray.__version__) >= Version("2.50"):
            ctx.use_polars_sort = effective
        else:
            ctx.use_polars = effective

    def apply_remote(self):
        import ray

        @ray.remote
        def apply():
            self.apply()

        ray.get(apply.remote())