🐛 fir working_dir and 📖 add docker example (#42)

Author: danielgafni

README.md

@@ -45,6 +45,7 @@ There are different options available for running Dagster code on Ray. The follo
 | Enabled per-asset  | ❌ | ❌ | ✅ | ✅ |
 | Configurable per-asset | ❌ | ✅ | ✅ | ✅ |
 | Runs asset/op body on Ray | ✅ | ✅ | ❌ | ❌ |
+| Requires configuring Dagster in the Ray cluster | ❌ | ❌ | ✅ | ✅ |
 
 # Examples
 
@@ -78,8 +79,9 @@ run_launcher:
   module: dagster_ray
   class: RayRunLauncher
   config:
-    num_cpus: 1
-    num_gpus: 0
+    ray:
+      num_cpus: 1
+      num_gpus: 0
 ```
 
 Individual Runs can **override** Ray configuration:

dagster_ray/executor.py

@@ -25,12 +25,17 @@
 from dagster_ray.config import RayExecutionConfig, RayJobSubmissionClientConfig
 from dagster_ray.kuberay.resources import get_k8s_object_name
 from dagster_ray.run_launcher import RayRunLauncher
+from dagster_ray.utils import resolve_env_vars_list
 
 if TYPE_CHECKING:
     from ray.job_submission import JobSubmissionClient
 
 
 class RayExecutorConfig(RayExecutionConfig, RayJobSubmissionClientConfig):
+    env_vars: Optional[List[str]] = 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).",
+    )
     address: Optional[str] = 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
@@ -74,6 +79,7 @@ def ray_executor(init_context: InitExecutorContext) -> Executor:
     return StepDelegatingExecutor(
         RayStepHandler(
             client=client,
+            env_vars=ray_cfg.env_vars,
             runtime_env=ray_cfg.runtime_env,
             num_cpus=ray_cfg.num_cpus,
             num_gpus=ray_cfg.num_gpus,
@@ -95,6 +101,7 @@ def name(self):
     def __init__(
         self,
         client: "JobSubmissionClient",
+        env_vars: Optional[List[str]],
         runtime_env: Optional[Dict[str, Any]],
         num_cpus: Optional[float],
         num_gpus: Optional[float],
@@ -104,6 +111,7 @@ def __init__(
         super().__init__()
 
         self.client = client
+        self.env_vars = env_vars or []
         self.runtime_env = runtime_env or {}
         self.num_cpus = num_cpus
         self.num_gpus = num_gpus
@@ -148,6 +156,7 @@ def launch_step(self, step_handler_context: StepHandlerContext) -> Iterator[Dags
 
         user_provided_config = RayExecutionConfig.from_tags({**step_handler_context.step_tags[step_key]})
 
+        # note! ray modifies the user-provided runtime_env, so we copy it
         runtime_env = (user_provided_config.runtime_env or self.runtime_env).copy()
 
         dagster_env_vars = {
@@ -158,6 +167,8 @@ def launch_step(self, step_handler_context: StepHandlerContext) -> Iterator[Dags
 
         runtime_env["env_vars"] = {**dagster_env_vars, **runtime_env.get("env_vars", {})}  # type: ignore
 
+        runtime_env["env_vars"].update(resolve_env_vars_list(self.env_vars))
+
         num_cpus = self.num_cpus or user_provided_config.num_cpus
         num_gpus = self.num_gpus or user_provided_config.num_gpus
         memory = self.memory or user_provided_config.memory

dagster_ray/io_manager.py

@@ -78,7 +78,7 @@ def handle_output(self, context: OutputContext, obj):
 
         object_map.set.remote(storage_key, ref)
 
-        context.log.debug(f"Stored object with key {storage_key} as {ref}")
+        context.log.debug(f"[RayIOManager] Stored object with key {storage_key} as {ref}")
 
     def load_input(self, context: InputContext):
         import ray
@@ -100,11 +100,11 @@ def load_input(self, context: InputContext):
         else:
             storage_key = self._get_single_key(context)
 
-        context.log.debug(f"Loading object with key {storage_key}")
+        context.log.debug(f"[RayIOManager] Loading object with key {storage_key}")
 
         ref = object_map.get.remote(storage_key)
 
-        assert ref is not None, f"Object with key {storage_key} not found in RayObjectMap"
+        assert ref is not None, f"[RayIOManager] Object with key {storage_key} not found in RayObjectMap"
 
         return ray.get(ref)
 
@@ -120,4 +120,6 @@ def _get_multiple_keys(self, context: InputContext) -> Dict[str, str]:
                 partition_key: "/".join(asset_path + [partition_key]) for partition_key in context.asset_partition_keys
             }
         else:
-            raise RuntimeError("This method can only be called with an InputContext that has multiple partitions")
+            raise RuntimeError(
+                "[RayIOManager] This method can only be called with an InputContext that has multiple partitions"
+            )

dagster_ray/kuberay/client/base.py

@@ -1,10 +1,8 @@
 import time
 from typing import TYPE_CHECKING, Any, Dict, Optional
 
-from kubernetes import client
-from kubernetes.client import ApiException
-
 if TYPE_CHECKING:
+    from kubernetes import client
     from kubernetes.client.models.v1_endpoints import V1Endpoints
 
 
@@ -27,8 +25,10 @@ def __init__(
         version: str,
         kind: str,
         plural: str,
-        api_client: Optional[client.ApiClient] = None,
+        api_client: Optional["client.ApiClient"] = None,
     ):
+        from kubernetes import client
+
         self.group = group
         self.version = version
         self.kind = kind
@@ -38,6 +38,8 @@ def __init__(
         self._core_v1_api = client.CoreV1Api(api_client=api_client)
 
     def wait_for_service_endpoints(self, service_name: str, namespace: str, poll_interval: int = 5, timeout: int = 60):
+        from kubernetes.client import ApiException
+
         start_time = time.time()
 
         while True:
@@ -62,6 +64,8 @@ def wait_for_service_endpoints(self, service_name: str, namespace: str, poll_int
             time.sleep(poll_interval)
 
     def get_status(self, name: str, namespace: str, timeout: int = 60, poll_interval: int = 5) -> Dict[str, Any]:
+        from kubernetes.client import ApiException
+
         while timeout > 0:
             try:
                 resource: Any = self._api.get_namespaced_custom_object_status(
@@ -83,6 +87,8 @@ def get_status(self, name: str, namespace: str, timeout: int = 60, poll_interval
         raise TimeoutError(f"Timed out waiting for status of {self.kind} {name} in namespace {namespace}")
 
     def list(self, namespace: str, label_selector: str = "", async_req: bool = False) -> Dict[str, Any]:
+        from kubernetes.client import ApiException
+
         try:
             resource: Any = self._api.list_namespaced_custom_object(
                 group=self.group,
@@ -103,6 +109,8 @@ def list(self, namespace: str, label_selector: str = "", async_req: bool = False
             raise
 
     def get(self, name: str, namespace: str) -> Dict[str, Any]:
+        from kubernetes.client import ApiException
+
         try:
             resource: Any = self._api.get_namespaced_custom_object(
                 group=self.group,

dagster_ray/kuberay/client/raycluster/client.py

@@ -17,7 +17,6 @@
     cast,
 )
 
-from kubernetes import watch
 from typing_extensions import NotRequired
 
 from dagster_ray.kuberay.client.base import BaseKubeRayClient
@@ -104,6 +103,8 @@ def wait_until_ready(
         timeout: int,
         image: Optional[str] = None,
     ) -> Tuple[str, Dict[str, str]]:
+        from kubernetes import watch
+
         """
         If ready, returns service ip address and a dictionary of ports.
         Dictionary keys: ["client", "dashboard", "metrics", "redis", "serve"]

dagster_ray/kuberay/resources.py

@@ -4,10 +4,9 @@
 import re
 import string
 import sys
-from typing import Any, Dict, Generator, Optional, cast
+from typing import TYPE_CHECKING, Any, Dict, Generator, Optional, cast
 
 import dagster._check as check
-import kubernetes
 from dagster import ConfigurableResource, InitResourceContext
 from dagster._annotations import experimental
 from pydantic import Field, PrivateAttr
@@ -30,15 +29,18 @@
 from dagster_ray._base.resources import BaseRayResource
 from dagster_ray.kuberay.client.base import load_kubeconfig
 
+if TYPE_CHECKING:
+    import kubernetes
+
 
 @experimental
 class RayClusterClientResource(ConfigurableResource):
     kube_context: Optional[str] = None
     kubeconfig_file: Optional[str] = None
 
     _raycluster_client: RayClusterClient = PrivateAttr()
-    _k8s_api: kubernetes.client.CustomObjectsApi = PrivateAttr()
-    _k8s_core_api: kubernetes.client.CoreV1Api = PrivateAttr()
+    _k8s_api: "kubernetes.client.CustomObjectsApi" = PrivateAttr()
+    _k8s_core_api: "kubernetes.client.CoreV1Api" = PrivateAttr()
 
     @property
     def client(self) -> RayClusterClient:
@@ -47,18 +49,20 @@ def client(self) -> RayClusterClient:
         return self._raycluster_client
 
     @property
-    def k8s(self) -> kubernetes.client.CustomObjectsApi:
+    def k8s(self) -> "kubernetes.client.CustomObjectsApi":
         if self._k8s_api is None:
             raise ValueError(f"{self.__class__.__name__} not initialized")
         return self._k8s_api
 
     @property
-    def k8s_core(self) -> kubernetes.client.CoreV1Api:
+    def k8s_core(self) -> "kubernetes.client.CoreV1Api":
         if self._k8s_core_api is None:
             raise ValueError(f"{self.__class__.__name__} not initialized")
         return self._k8s_core_api
 
     def setup_for_execution(self, context: InitResourceContext) -> None:
+        import kubernetes
+
         load_kubeconfig(context=self.kube_context, config_file=self.kubeconfig_file)
 
         self._raycluster_client = RayClusterClient(context=self.kube_context, config_file=self.kubeconfig_file)
@@ -218,6 +222,8 @@ def update_group_spec(group_spec: Dict[str, Any]):
         }
 
     def _wait_raycluster_ready(self):
+        import kubernetes
+
         self.client.client.wait_until_ready(self.cluster_name, namespace=self.namespace, timeout=self.timeout)
 
         # the above code only checks for RayCluster creation

dagster_ray/run_launcher.py

@@ -1,6 +1,6 @@
 import logging
 import sys
-from typing import Any, Dict, Optional
+from typing import TYPE_CHECKING, Any, Dict, List, Optional
 
 from dagster import _check as check
 from dagster._cli.api import ExecuteRunArgs  # type: ignore
@@ -12,15 +12,24 @@
 from dagster._grpc.types import ResumeRunArgs
 from dagster._serdes import ConfigurableClass, ConfigurableClassData
 from dagster._utils.error import serializable_error_info_from_exc_info
+from pydantic import Field
 
 from dagster_ray.config import RayExecutionConfig, RayJobSubmissionClientConfig
+from dagster_ray.utils import resolve_env_vars_list
+
+if TYPE_CHECKING:
+    from ray.job_submission import JobSubmissionClient
 
 
 def get_job_submission_id_from_run_id(run_id: str, resume_attempt_number=None):
     return f"dagster-run-{run_id}" + ("" if not resume_attempt_number else f"-{resume_attempt_number}")
 
 
-class RayLauncherConfig(RayExecutionConfig, RayJobSubmissionClientConfig): ...
+class RayLauncherConfig(RayExecutionConfig, RayJobSubmissionClientConfig):
+    env_vars: Optional[List[str]] = 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).",
+    )
 
 
 class RayRunLauncher(RunLauncher, ConfigurableClass):
@@ -30,6 +39,7 @@ def __init__(
         metadata: Optional[Dict[str, Any]] = None,
         headers: Optional[Dict[str, Any]] = None,
         cookies: Optional[Dict[str, Any]] = None,
+        env_vars: Optional[List[str]] = None,
         runtime_env: Optional[Dict[str, Any]] = None,
         num_cpus: Optional[int] = None,
         num_gpus: Optional[int] = None,
@@ -55,17 +65,13 @@ def __init__(
         Fields such as `num_cpus` set via `dagster-ray/config` Run tag will override the yaml configuration.
 
         """
-
-        from ray.job_submission import JobSubmissionClient
-
         self._inst_data = check.opt_inst_param(inst_data, "inst_data", ConfigurableClassData)
 
-        self.client = JobSubmissionClient(address, metadata=metadata, headers=headers, cookies=cookies)
-
         self.address = address
         self.metadata = metadata
         self.headers = headers
         self.cookies = cookies
+        self.env_vars = env_vars
         self.runtime_env = runtime_env
         self.num_cpus = num_cpus
         self.num_gpus = num_gpus
@@ -74,17 +80,23 @@ def __init__(
 
         super().__init__()
 
+    @property
+    def client(self) -> "JobSubmissionClient":  # note: this must be a property
+        from ray.job_submission import JobSubmissionClient
+
+        return JobSubmissionClient(self.address, metadata=self.metadata, headers=self.headers, cookies=self.cookies)
+
     @property
     def inst_data(self) -> Optional[ConfigurableClassData]:
         return self._inst_data
 
     @classmethod
     def config_type(cls) -> UserConfigSchema:
-        return RayLauncherConfig.to_fields_dict()
+        return {"ray": RayLauncherConfig.to_config_schema().as_field()}
 
     @classmethod
     def from_config_value(cls, inst_data, config_value):
-        return cls(inst_data=inst_data, **config_value)
+        return cls(inst_data=inst_data, **config_value["ray"])
 
     @property
     def supports_resume_run(self):
@@ -132,13 +144,16 @@ def _launch_ray_job(self, submission_id: str, entrypoint: str, run: DagsterRun):
 
         cfg_from_tags = RayLauncherConfig.from_tags(run.tags)
 
-        runtime_env = cfg_from_tags.runtime_env or self.runtime_env or {}
+        env_vars = cfg_from_tags.env_vars or self.env_vars or []
+        # note! ray modifies the user-provided runtime_env, so we copy it
+        runtime_env = (cfg_from_tags.runtime_env or self.runtime_env or {}).copy()
         num_cpus = cfg_from_tags.num_cpus or self.num_cpus
         num_gpus = cfg_from_tags.num_gpus or self.num_gpus
         memory = cfg_from_tags.memory or self.memory
         resources = cfg_from_tags.resources or self.resources
 
         runtime_env["env_vars"] = runtime_env.get("env_vars", {})
+        runtime_env["env_vars"].update(resolve_env_vars_list(env_vars))
         runtime_env["env_vars"].update(
             {
                 "DAGSTER_RUN_JOB_NAME": job_origin.job_name,

dagster_ray/utils.py

@@ -0,0 +1,17 @@
+import os
+from typing import Dict, List, Optional
+
+
+def resolve_env_vars_list(env_vars: Optional[List[str]]) -> Dict[str, str]:
+    res = {}
+
+    if env_vars is not None:
+        for env_var in env_vars:
+            if "=" in env_var:
+                var, value = env_var.split("=", 1)
+                res[var] = value
+            else:
+                if value := os.getenv(env_var):
+                    res[env_var] = value
+
+    return res