Add Ray executor (#687)

* Add Ray executor

* Mypy fixes

Author: tomwhite

.github/workflows/ray-tests.yml

@@ -0,0 +1,45 @@
+name: Ray tests
+
+on:
+  push:
+    branches:
+      - "main"
+  pull_request:
+  workflow_dispatch:
+
+concurrency:
+  group: ${{ github.workflow }}-${{ github.ref }}
+  cancel-in-progress: true
+
+jobs:
+  test:
+    runs-on: ${{ matrix.os }}
+    strategy:
+      fail-fast: false
+      matrix:
+        os: ["ubuntu-latest"]
+        python-version: ["3.12"]
+
+    steps:
+      - name: Checkout source
+        uses: actions/checkout@v3
+        with:
+          fetch-depth: 0
+
+      - name: Set up Python
+        uses: actions/setup-python@v3
+        with:
+          python-version: ${{ matrix.python-version }}
+          architecture: x64
+
+      - name: Setup Graphviz
+        uses: ts-graphviz/setup-graphviz@v2
+
+      - name: Install
+        run: |
+          python -m pip install --upgrade pip
+          python -m pip install -e '.[test]' 'ray[default]'
+
+      - name: Run tests
+        run: |
+          pytest -vs

cubed/runtime/create.py

@@ -30,6 +30,10 @@ def create_executor(name: str, executor_options: Optional[dict] = None) -> Execu
         from cubed.runtime.executors.local import ProcessesExecutor
 
         return ProcessesExecutor(**executor_options)
+    elif name == "ray":
+        from cubed.runtime.executors.ray import RayExecutor
+
+        return RayExecutor(**executor_options)
     elif name == "single-threaded":
         from cubed.runtime.executors.local import SingleThreadedExecutor
 

cubed/runtime/executors/ray.py

@@ -0,0 +1,86 @@
+import asyncio
+from typing import Optional, Sequence
+
+import ray
+from networkx import MultiDiGraph
+
+from cubed.runtime.asyncio import async_map_dag
+from cubed.runtime.backup import use_backups_default
+from cubed.runtime.types import Callback, DagExecutor
+from cubed.runtime.utils import asyncio_run, execute_with_stats
+from cubed.spec import Spec
+
+
+class RayExecutor(DagExecutor):
+    """An execution engine that uses Ray."""
+
+    def __init__(self, **kwargs):
+        self.kwargs = kwargs
+
+    @property
+    def name(self) -> str:
+        return "ray"
+
+    def execute_dag(
+        self,
+        dag: MultiDiGraph,
+        callbacks: Optional[Sequence[Callback]] = None,
+        resume: Optional[bool] = None,
+        spec: Optional[Spec] = None,
+        compute_id: Optional[str] = None,
+        **kwargs,
+    ) -> None:
+        merged_kwargs = {**self.kwargs, **kwargs}
+
+        ray_init = merged_kwargs.pop("ray_init", None)
+        if ray_init is not None:
+            ray.init(**ray_init)
+
+        asyncio_run(
+            self._async_execute_dag(
+                dag,
+                callbacks=callbacks,
+                resume=resume,
+                spec=spec,
+                compute_id=compute_id,
+                **merged_kwargs,
+            )
+        )
+
+    async def _async_execute_dag(
+        self,
+        dag: MultiDiGraph,
+        callbacks: Optional[Sequence[Callback]] = None,
+        resume: Optional[bool] = None,
+        spec: Optional[Spec] = None,
+        compute_arrays_in_parallel: Optional[bool] = None,
+        **kwargs,
+    ) -> None:
+        if spec is not None:
+            if "use_backups" not in kwargs and use_backups_default(spec):
+                kwargs["use_backups"] = True
+
+        allowed_mem = spec.allowed_mem if spec is not None else 2_000_000_000
+        retries = kwargs.pop("retries", 2)
+
+        # note we can define the remote function here (doesn't need to be at top-level), and pass in memory and retries
+        @ray.remote(memory=allowed_mem, max_retries=retries, retry_exceptions=True)
+        def run_remotely(input, func=None, config=None, name=None, compute_id=None):
+            # note we can't use the execution_stat decorator since it doesn't work with ray decorators
+            result, stats = execute_with_stats(func, input, config=config)
+            return result, stats
+
+        def create_futures_func(input, **kwargs):
+            return [
+                (i, asyncio.wrap_future(run_remotely.remote(i, **kwargs).future()))
+                for i in input
+            ]
+
+        await async_map_dag(
+            create_futures_func,
+            dag=dag,
+            callbacks=callbacks,
+            resume=resume,
+            compute_arrays_in_parallel=compute_arrays_in_parallel,
+            **kwargs,
+        )

cubed/tests/utils.py

@@ -51,6 +51,12 @@
 except ImportError:
     pass
 
+try:
+    ALL_EXECUTORS.append(create_executor("ray"))
+    MAIN_EXECUTORS.append(create_executor("ray"))
+except ImportError:
+    pass
+
 MODAL_EXECUTORS = []
 
 try:

setup.cfg

@@ -64,6 +64,8 @@ ignore_missing_imports = True
 ignore_missing_imports = True
 [mypy-pytest.*]
 ignore_missing_imports = True
+[mypy-ray.*]
+ignore_missing_imports = True
 [mypy-rechunker.*]
 ignore_missing_imports = True
 [mypy-seaborn.*]