✅ enable most of Ruff rules

dagster_ray/__init__.py

@@ -9,11 +9,11 @@
 
 
 __all__ = [
+    "LocalRay",
+    "PipesRayJobClient",
+    "PipesRayJobMessageReader",
+    "RayIOManager",
     "RayResource",
     "RayRunLauncher",
-    "RayIOManager",
     "ray_executor",
-    "PipesRayJobMessageReader",
-    "PipesRayJobClient",
-    "LocalRay",
 ]

dagster_ray/_base/resources.py

@@ -1,12 +1,11 @@
+from __future__ import annotations
+
 import uuid
 from abc import ABC, abstractmethod
-from typing import TYPE_CHECKING, Optional, Union, cast
+from typing import TYPE_CHECKING, Union, cast
 
 from dagster import AssetExecutionContext, ConfigurableResource, InitResourceContext, OpExecutionContext
 from pydantic import Field, PrivateAttr
-
-# yes, `python-client` is actually the KubeRay package name
-# https://github.com/ray-project/kuberay/issues/2078
 from requests.exceptions import ConnectionError
 from tenacity import retry, retry_if_exception_type, stop_after_delay
 from typing_extensions import TypeAlias
@@ -15,42 +14,46 @@
 from dagster_ray.config import RayDataExecutionOptions
 
 if TYPE_CHECKING:
-    from ray._private.worker import BaseContext as RayBaseContext  # noqa
+    from ray._private.worker import BaseContext as RayBaseContext
 
 OpOrAssetExecutionContext: TypeAlias = Union[OpExecutionContext, AssetExecutionContext]
 
 
 class BaseRayResource(ConfigurableResource, ABC):
-    """
-    Base class for Ray Resources.
+    """Base class for Ray Resources.
     Defines the common interface and some utility methods.
     """
 
     data_execution_options: RayDataExecutionOptions = Field(default_factory=RayDataExecutionOptions)
     redis_port: int = Field(
-        default=10001, description="Redis port for connection. Make sure to match with the actual available port."
+        default=10001,
+        description="Redis port for connection. Make sure to match with the actual available port.",
     )
     dashboard_port: int = Field(
-        default=8265, description="Dashboard port for connection. Make sure to match with the actual available port."
+        default=8265,
+        description="Dashboard port for connection. Make sure to match with the actual available port.",
     )
 
-    _context: Optional["RayBaseContext"] = PrivateAttr()
+    _context: RayBaseContext | None = PrivateAttr()
 
     def setup_for_execution(self, context: InitResourceContext) -> None:
-        raise NotImplementedError(
+        msg = (
             "This is an abstract resource, it's not meant to be provided directly. "
             "Use a backend-specific resource instead."
         )
+        raise NotImplementedError(
+            msg,
+        )
 
     @property
-    def context(self) -> "RayBaseContext":
+    def context(self) -> RayBaseContext:
         assert self._context is not None, "RayClusterResource not initialized"
         return self._context
 
     @property
     @abstractmethod
     def host(self) -> str:
-        raise NotImplementedError()
+        raise NotImplementedError
 
     @property
     def ray_address(self) -> str:
@@ -62,16 +65,15 @@ def dashboard_url(self) -> str:
 
     @property
     def runtime_job_id(self) -> str:
-        """
-        Returns the Ray Job ID for the current job which was created with `ray.init()`.
+        """Returns the Ray Job ID for the current job which was created with `ray.init()`.
         :return:
         """
         import ray
 
         return ray.get_runtime_context().get_job_id()
 
     @retry(stop=stop_after_delay(120), retry=retry_if_exception_type(ConnectionError), reraise=True)
-    def init_ray(self, context: Union[OpOrAssetExecutionContext, InitResourceContext]) -> "RayBaseContext":
+    def init_ray(self, context: OpOrAssetExecutionContext | InitResourceContext) -> RayBaseContext:
         assert context.log is not None
 
         import ray
@@ -84,8 +86,7 @@ def init_ray(self, context: Union[OpOrAssetExecutionContext, InitResourceContext
         return cast("RayBaseContext", self._context)
 
     def get_dagster_tags(self, context: InitResourceContext) -> dict[str, str]:
-        tags = get_dagster_tags(context)
-        return tags
+        return get_dagster_tags(context)
 
     def _get_step_key(self, context: InitResourceContext) -> str:
         # just return a random string

dagster_ray/_base/utils.py

@@ -1,15 +1,16 @@
-import os
-from typing import Union, cast
+from __future__ import annotations
 
-from dagster import AssetExecutionContext, InitResourceContext, OpExecutionContext
+import os
+from typing import TYPE_CHECKING, cast
 
 from dagster_ray._base.constants import DEFAULT_DEPLOYMENT_NAME
 
+if TYPE_CHECKING:
+    from dagster import AssetExecutionContext, InitResourceContext, OpExecutionContext
+
 
-def get_dagster_tags(context: Union[InitResourceContext, OpExecutionContext, AssetExecutionContext]) -> dict[str, str]:
-    """
-    Returns a dictionary with common Dagster tags.
-    """
+def get_dagster_tags(context: InitResourceContext | OpExecutionContext | AssetExecutionContext) -> dict[str, str]:
+    """Returns a dictionary with common Dagster tags."""
     assert context.dagster_run is not None
 
     labels = {

dagster_ray/config.py

@@ -1,11 +1,13 @@
 from __future__ import annotations
 
-from collections.abc import Mapping
-from typing import Any
+from typing import TYPE_CHECKING, Any
 
 from dagster import Config
 from pydantic import Field
 
+if TYPE_CHECKING:
+    from collections.abc import Mapping
+
 USER_DEFINED_RAY_KEY = "dagster-ray/config"
 
 
@@ -20,8 +22,7 @@ class RayExecutionConfig(Config):
     def from_tags(cls, tags: Mapping[str, str]) -> RayExecutionConfig:
         if USER_DEFINED_RAY_KEY in tags:
             return cls.parse_raw(tags[USER_DEFINED_RAY_KEY])
-        else:
-            return cls()
+        return cls()
 
 
 class RayJobSubmissionClientConfig(Config):
@@ -38,7 +39,8 @@ class RayJobSubmissionClientConfig(Config):
             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."
+        default=None,
+        description="Cookies to use when sending requests to the HTTP job server.",
     )
 
 
@@ -56,7 +58,7 @@ class RayDataExecutionOptions(Config):
     verbose_progress: bool = True
     use_polars: bool = True
 
-    def apply(self):
+    def apply(self) -> None:
         import ray
         from ray.data import ExecutionResources
 
@@ -71,11 +73,11 @@ def apply(self):
         ctx.verbose_progress = self.verbose_progress
         ctx.use_polars = self.use_polars
 
-    def apply_remote(self):
+    def apply_remote(self) -> None:
         import ray
 
         @ray.remote
-        def apply():
+        def apply() -> None:
             self.apply()
 
         ray.get(apply.remote())

dagster_ray/executor.py

@@ -1,4 +1,5 @@
-from collections.abc import Iterator
+from __future__ import annotations
+
 from typing import TYPE_CHECKING, Any, Optional, cast
 
 import dagster
@@ -13,8 +14,6 @@
 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,
@@ -32,15 +31,20 @@
 from dagster_ray.utils import resolve_env_vars_list
 
 if TYPE_CHECKING:
+    from collections.abc import Iterator
+
+    from dagster._core.executor.base import Executor
+    from dagster._core.executor.init import InitExecutorContext
     from ray.job_submission import JobSubmissionClient
 
 
 class RayExecutorConfig(RayExecutionConfig, RayJobSubmissionClientConfig):
-    env_vars: Optional[list[str]] = Field(
+    env_vars: list[str] | None = Field(
         default=None,
-        description="A list of environment variables to inject into the Job. Each can be of the form KEY=VALUE or just KEY (in which case the value will be pulled from the current process).",
+        description="""A list of environment variables to inject into the Job. Each can be of the form KEY=VALUE
+            or just KEY (in which case the value will be pulled from the current process).""",
     )
-    address: Optional[str] = Field(default=None, description="The address of the Ray cluster to connect to.")  # type: ignore
+    address: str | None = Field(default=None, description="The address of the Ray cluster to connect to.")  # type: ignore
     # sorry for the long name, but it has to be very clear what this is doing
     inherit_job_submission_client_from_ray_run_launcher: bool = True
 
@@ -71,13 +75,17 @@ def ray_executor(init_context: InitExecutorContext) -> Executor:
     ray_cfg = RayExecutorConfig(**exc_cfg["ray"])  # type: ignore
 
     if ray_cfg.inherit_job_submission_client_from_ray_run_launcher and isinstance(
-        init_context.instance.run_launcher, RayRunLauncher
+        init_context.instance.run_launcher,
+        RayRunLauncher,
     ):
         # TODO: some RunLauncher config values can be automatically passed to the executor
         client = init_context.instance.run_launcher.client
     else:
         client = JobSubmissionClient(
-            ray_cfg.address, metadata=ray_cfg.metadata, headers=ray_cfg.headers, cookies=ray_cfg.cookies
+            ray_cfg.address,
+            metadata=ray_cfg.metadata,
+            headers=ray_cfg.headers,
+            cookies=ray_cfg.cookies,
         )
 
     return StepDelegatingExecutor(
@@ -99,19 +107,19 @@ def ray_executor(init_context: InitExecutorContext) -> Executor:
 
 class RayStepHandler(StepHandler):
     @property
-    def name(self):
+    def name(self) -> str:
         return "RayStepHandler"
 
     def __init__(
         self,
-        client: "JobSubmissionClient",
-        env_vars: Optional[list[str]],
-        runtime_env: Optional[dict[str, Any]],
-        num_cpus: Optional[float],
-        num_gpus: Optional[float],
-        memory: Optional[int],
-        resources: Optional[dict[str, float]],
-    ):
+        client: JobSubmissionClient,
+        env_vars: list[str] | None,
+        runtime_env: dict[str, Any] | None,
+        num_cpus: float | None,
+        num_gpus: float | None,
+        memory: int | None,
+        resources: dict[str, float] | None,
+    ) -> None:
         super().__init__()
 
         self.client = client
@@ -127,7 +135,7 @@ def _get_step_key(self, step_handler_context: StepHandlerContext) -> str:
         assert len(step_keys_to_execute) == 1, "Launching multiple steps is not currently supported"
         return step_keys_to_execute[0]
 
-    def _get_ray_job_submission_id(self, step_handler_context: StepHandlerContext):
+    def _get_ray_job_submission_id(self, step_handler_context: StepHandlerContext) -> str:
         step_key = self._get_step_key(step_handler_context)
 
         name_key = get_k8s_object_name(
@@ -195,7 +203,7 @@ def launch_step(self, step_handler_context: StepHandlerContext) -> Iterator[Dags
 
         self.client.submit_job(
             entrypoint=" ".join(
-                step_handler_context.execute_step_args.get_command_args(skip_serialized_namedtuple=True)
+                step_handler_context.execute_step_args.get_command_args(skip_serialized_namedtuple=True),
             ),
             submission_id=submission_id,
             metadata=labels,
@@ -217,7 +225,7 @@ def check_step_health(self, step_handler_context: StepHandlerContext) -> CheckSt
             status = self.client.get_job_status(submission_id)
         except RuntimeError:
             return CheckStepHealthResult.unhealthy(
-                reason=f"Ray job {submission_id} for step {step_key} could not be found."
+                reason=f"Ray job {submission_id} for step {step_key} could not be found.",
             )
 
         if status == JobStatus.FAILED: