Repository: cubed-dev/cubed (python)

Issue:
Add a progress bar context manager
It would be nice to be able to do the following (like Dask):

```python
with ProgressBar():
    res = a.compute()
```

(Related to #538)

Pull request:
Make callbacks context managers
Fixes #541 

Patch:
diff --git a/cubed/array/nan_functions.py b/cubed/array/nan_functions.py
index 1801fa5b..fe0d7b66 100644
--- a/cubed/array/nan_functions.py
+++ b/cubed/array/nan_functions.py
@@ -8,9 +8,9 @@ from cubed.core import reduction
 # https://github.com/data-apis/array-api/issues/621
 
 
-def nanmean(x, /, *, axis=None, keepdims=False, split_every=None):
+def nanmean(x, /, *, axis=None, dtype=None, keepdims=False, split_every=None):
     """Compute the arithmetic mean along the specified axis, ignoring NaNs."""
-    dtype = x.dtype
+    dtype = dtype or x.dtype
     intermediate_dtype = [("n", nxp.int64), ("total", nxp.float64)]
     return reduction(
         x,
diff --git a/cubed/array_api/array_object.py b/cubed/array_api/array_object.py
index 94121d7e..83b1ee94 100644
--- a/cubed/array_api/array_object.py
+++ b/cubed/array_api/array_object.py
@@ -367,9 +367,9 @@ class Array(CoreArray):
             "2023.12",
         ):
             raise ValueError(f"Unrecognized array API version: {api_version!r}")
-        import cubed.array_api as array_api
+        import cubed
 
-        return array_api
+        return cubed
 
     def __bool__(self, /):
         if self.ndim != 0:
diff --git a/cubed/core/array.py b/cubed/core/array.py
index 7730f47f..a161927b 100644
--- a/cubed/core/array.py
+++ b/cubed/core/array.py
@@ -278,10 +278,18 @@ def compute(
 
             executor = ThreadsExecutor()
 
+    # combine any callbacks specified as args with any active callbacks from the context manager
+    if callbacks is None and len(Callback.active) == 0:
+        all_callbacks = None
+    else:
+        all_callbacks = list(Callback.active)
+        if callbacks is not None:
+            all_callbacks.extend(callbacks)
+
     _return_in_memory_array = kwargs.pop("_return_in_memory_array", True)
     plan.execute(
         executor=executor,
-        callbacks=callbacks,
+        callbacks=all_callbacks,
         optimize_graph=optimize_graph,
         optimize_function=optimize_function,
         resume=resume,
diff --git a/cubed/diagnostics/__init__.py b/cubed/diagnostics/__init__.py
index e69de29b..d1aadeba 100644
--- a/cubed/diagnostics/__init__.py
+++ b/cubed/diagnostics/__init__.py
@@ -0,0 +1,3 @@
+from .rich import RichProgressBar as ProgressBar
+
+__all__ = ["ProgressBar"]
diff --git a/cubed/runtime/create.py b/cubed/runtime/create.py
index d3c9637b..4fa802a2 100644
--- a/cubed/runtime/create.py
+++ b/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
 
diff --git a/cubed/runtime/executors/ray.py b/cubed/runtime/executors/ray.py
new file mode 100644
index 00000000..00e50788
--- /dev/null
+++ b/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,
+        )
diff --git a/cubed/runtime/types.py b/cubed/runtime/types.py
index 8aa06fca..8bf9fe9c 100644
--- a/cubed/runtime/types.py
+++ b/cubed/runtime/types.py
@@ -1,5 +1,6 @@
 from dataclasses import dataclass
-from typing import Any, Callable, Iterable, Optional
+from types import TracebackType
+from typing import Any, Callable, ClassVar, Iterable, Optional
 
 from networkx import MultiDiGraph
 
@@ -96,6 +97,26 @@ class TaskEndEvent:
 class Callback:
     """Object to receive callback events during array computation."""
 
+    active: ClassVar[set["Callback"]] = set()
+
+    def __enter__(self):
+        self.register()
+        return self
+
+    def __exit__(
+        self,
+        exc_type: type[BaseException] | None,
+        exc_value: BaseException | None,
+        traceback: TracebackType | None,
+    ) -> None:
+        self.unregister()
+
+    def register(self) -> None:
+        Callback.active.add(self)
+
+    def unregister(self) -> None:
+        Callback.active.remove(self)
+
     def on_compute_start(self, event):
         """Called when the computation is about to start.
 
diff --git a/examples/add-random.py b/examples/add-random.py
index ba0d6bbb..eff172b2 100644
--- a/examples/add-random.py
+++ b/examples/add-random.py
@@ -3,8 +3,8 @@ import logging
 import cubed
 import cubed.array_api as xp
 import cubed.random
+from cubed.diagnostics import ProgressBar
 from cubed.diagnostics.history import HistoryCallback
-from cubed.diagnostics.rich import RichProgressBar
 from cubed.diagnostics.timeline import TimelineVisualizationCallback
 
 # suppress harmless connection pool warnings
@@ -16,12 +16,6 @@ if __name__ == "__main__":
     b = cubed.random.random((50000, 50000), chunks=(5000, 5000))
     c = xp.add(a, b)
 
-    progress = RichProgressBar()
-    hist = HistoryCallback()
-    timeline_viz = TimelineVisualizationCallback()
     # use store=None to write to temporary zarr
-    cubed.to_zarr(
-        c,
-        store=None,
-        callbacks=[progress, hist, timeline_viz],
-    )
+    with ProgressBar(), HistoryCallback(), TimelineVisualizationCallback():
+        cubed.to_zarr(c, store=None)
diff --git a/setup.cfg b/setup.cfg
index 0577370d..aee60eb7 100644
--- a/setup.cfg
+++ b/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.*]


Test patch:
diff --git a/.github/workflows/ray-tests.yml b/.github/workflows/ray-tests.yml
new file mode 100644
index 00000000..e8967567
--- /dev/null
+++ b/.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
diff --git a/cubed/tests/test_executor_features.py b/cubed/tests/test_executor_features.py
index dfc4e1ac..e29d486d 100644
--- a/cubed/tests/test_executor_features.py
+++ b/cubed/tests/test_executor_features.py
@@ -12,6 +12,7 @@ from numpy.testing import assert_array_equal
 import cubed
 import cubed.array_api as xp
 import cubed.random
+from cubed.diagnostics import ProgressBar
 from cubed.diagnostics.history import HistoryCallback
 from cubed.diagnostics.mem_warn import MemoryWarningCallback
 from cubed.diagnostics.rich import RichProgressBar
@@ -115,6 +116,21 @@ def test_callbacks(spec, executor):
     assert task_counter.value == num_created_arrays + 4
 
 
+def test_callbacks_as_context_managers(spec, executor):
+    with TaskCounter() as task_counter, ProgressBar():
+        assert task_counter is not None
+        a = xp.asarray([[1, 2, 3], [4, 5, 6], [7, 8, 9]], chunks=(2, 2), spec=spec)
+        b = xp.asarray([[1, 1, 1], [1, 1, 1], [1, 1, 1]], chunks=(2, 2), spec=spec)
+        c = xp.add(a, b)
+        assert_array_equal(
+            c.compute(executor=executor),
+            np.array([[2, 3, 4], [5, 6, 7], [8, 9, 10]]),
+        )
+
+        num_created_arrays = 1
+        assert task_counter.value == num_created_arrays + 4
+
+
 def test_rich_progress_bar(spec, executor):
     # test indirectly by checking it doesn't cause a failure
     progress = RichProgressBar()
diff --git a/cubed/tests/utils.py b/cubed/tests/utils.py
index e3d266bb..a5a98f74 100644
--- a/cubed/tests/utils.py
+++ b/cubed/tests/utils.py
@@ -51,6 +51,12 @@ try:
 except ImportError:
     pass
 
+try:
+    ALL_EXECUTORS.append(create_executor("ray"))
+    MAIN_EXECUTORS.append(create_executor("ray"))
+except ImportError:
+    pass
+
 MODAL_EXECUTORS = []
 
 try: