Repository: tomwhite/cubed
Pull request #263 (2023-07-14 14:26:26)

Add `compute_arrays_in_parallel` option to Modal and Python async executors
This is #80 for Modal.

(See also #259.)


Patch:
diff --git a/cubed/extensions/timeline.py b/cubed/extensions/timeline.py
index 1b8a269f..310379d9 100644
--- a/cubed/extensions/timeline.py
+++ b/cubed/extensions/timeline.py
@@ -31,6 +31,10 @@ class TimelineVisualizationCallback(Callback):
 def create_timeline(stats, start_tstamp, end_tstamp, dst=None):
     stats_df = pd.DataFrame(stats)
 
+    stats_df = stats_df.sort_values(
+        by=["task_create_tstamp", "array_name"], ascending=True
+    )
+
     total_calls = len(stats_df)
 
     palette = sns.color_palette("deep", 6)
diff --git a/cubed/runtime/executors/modal_async.py b/cubed/runtime/executors/modal_async.py
index dd1704e1..c327b47f 100644
--- a/cubed/runtime/executors/modal_async.py
+++ b/cubed/runtime/executors/modal_async.py
@@ -2,15 +2,17 @@ import asyncio
 import copy
 import time
 from asyncio.exceptions import TimeoutError
+from functools import partial
 from typing import Any, AsyncIterator, Dict, Iterable, Optional, Sequence
 
+from aiostream import stream
 from modal.exception import ConnectionError
 from modal.functions import Function
 from networkx import MultiDiGraph
 from tenacity import retry, retry_if_exception_type, stop_after_attempt
 
 from cubed.core.array import Callback
-from cubed.core.plan import visit_nodes
+from cubed.core.plan import visit_node_generations, visit_nodes
 from cubed.runtime.backup import should_launch_backup
 from cubed.runtime.executors.modal import Container, run_remotely, stub
 from cubed.runtime.types import DagExecutor
@@ -115,6 +117,29 @@ async def map_unordered(
                     backups[new_task] = task
 
 
+def pipeline_to_stream(app_function, name, pipeline, **kwargs):
+    if any([stage for stage in pipeline.stages if stage.mappable is None]):
+        raise NotImplementedError("All stages must be mappable in pipelines")
+    it = stream.iterate(
+        [
+            partial(
+                map_unordered,
+                app_function,
+                stage.mappable,
+                return_stats=True,
+                name=name,
+                func=stage.function,
+                config=pipeline.config,
+                **kwargs,
+            )
+            for stage in pipeline.stages
+            if stage.mappable is not None
+        ]
+    )
+    # concat stages, running only one stage at a time
+    return stream.concatmap(it, lambda f: f(), task_limit=1)
+
+
 # This just retries the initial connection attempt, not the function calls
 @retry(
     reraise=True,
@@ -127,6 +152,7 @@ async def async_execute_dag(
     array_names: Optional[Sequence[str]] = None,
     resume: Optional[bool] = None,
     cloud: Optional[str] = None,
+    compute_arrays_in_parallel: Optional[bool] = None,
     **kwargs,
 ) -> None:
     async with stub.run():
@@ -137,23 +163,24 @@ async def async_execute_dag(
             app_function = Container().run_remotely
         else:
             raise ValueError(f"Unrecognized cloud: {cloud}")
-        for name, node in visit_nodes(dag, resume=resume):
-            pipeline = node["pipeline"]
-
-            for stage in pipeline.stages:
-                if stage.mappable is not None:
-                    async for _, stats in map_unordered(
-                        app_function,
-                        stage.mappable,
-                        return_stats=True,
-                        name=name,
-                        func=stage.function,
-                        config=pipeline.config,
-                        **kwargs,
-                    ):
+        if not compute_arrays_in_parallel:
+            # run one pipeline at a time
+            for name, node in visit_nodes(dag, resume=resume):
+                st = pipeline_to_stream(app_function, name, node["pipeline"], **kwargs)
+                async with st.stream() as streamer:
+                    async for _, stats in streamer:
+                        handle_callbacks(callbacks, stats)
+        else:
+            for gen in visit_node_generations(dag, resume=resume):
+                # run pipelines in the same topological generation in parallel by merging their streams
+                streams = [
+                    pipeline_to_stream(app_function, name, node["pipeline"], **kwargs)
+                    for name, node in gen
+                ]
+                merged_stream = stream.merge(*streams)
+                async with merged_stream.stream() as streamer:
+                    async for _, stats in streamer:
                         handle_callbacks(callbacks, stats)
-                else:
-                    raise NotImplementedError()
 
 
 class AsyncModalDagExecutor(DagExecutor):
diff --git a/cubed/runtime/executors/python_async.py b/cubed/runtime/executors/python_async.py
index 37faaff6..3f5351a5 100644
--- a/cubed/runtime/executors/python_async.py
+++ b/cubed/runtime/executors/python_async.py
@@ -10,7 +10,7 @@ from networkx import MultiDiGraph
 from tenacity import Retrying, stop_after_attempt
 
 from cubed.core.array import Callback
-from cubed.core.plan import visit_node_generations
+from cubed.core.plan import visit_node_generations, visit_nodes
 from cubed.primitive.types import CubedPipeline
 from cubed.runtime.types import DagExecutor
 from cubed.runtime.utils import execution_stats, handle_callbacks
@@ -99,21 +99,32 @@ async def async_execute_dag(
     callbacks: Optional[Sequence[Callback]] = None,
     array_names: Optional[Sequence[str]] = None,
     resume: Optional[bool] = None,
+    compute_arrays_in_parallel: Optional[bool] = None,
     **kwargs,
 ) -> None:
     with ThreadPoolExecutor() as concurrent_executor:
-        for gen in visit_node_generations(dag, resume=resume):
-            # run pipelines in the same topological generation in parallel by merging their streams
-            streams = [
-                pipeline_to_stream(
+        if not compute_arrays_in_parallel:
+            # run one pipeline at a time
+            for name, node in visit_nodes(dag, resume=resume):
+                st = pipeline_to_stream(
                     concurrent_executor, name, node["pipeline"], **kwargs
                 )
-                for name, node in gen
-            ]
-            merged_stream = stream.merge(*streams)
-            async with merged_stream.stream() as streamer:
-                async for _, stats in streamer:
-                    handle_callbacks(callbacks, stats)
+                async with st.stream() as streamer:
+                    async for _, stats in streamer:
+                        handle_callbacks(callbacks, stats)
+        else:
+            for gen in visit_node_generations(dag, resume=resume):
+                # run pipelines in the same topological generation in parallel by merging their streams
+                streams = [
+                    pipeline_to_stream(
+                        concurrent_executor, name, node["pipeline"], **kwargs
+                    )
+                    for name, node in gen
+                ]
+                merged_stream = stream.merge(*streams)
+                async with merged_stream.stream() as streamer:
+                    async for _, stats in streamer:
+                        handle_callbacks(callbacks, stats)
 
 
 class AsyncPythonDagExecutor(DagExecutor):


Test patch:
diff --git a/cubed/tests/test_core.py b/cubed/tests/test_core.py
index 6a9cc333..f4519815 100644
--- a/cubed/tests/test_core.py
+++ b/cubed/tests/test_core.py
@@ -15,7 +15,13 @@ from cubed.extensions.timeline import TimelineVisualizationCallback
 from cubed.extensions.tqdm import TqdmProgressBar
 from cubed.primitive.blockwise import apply_blockwise
 from cubed.runtime.types import DagExecutor
-from cubed.tests.utils import MAIN_EXECUTORS, MODAL_EXECUTORS, TaskCounter, create_zarr
+from cubed.tests.utils import (
+    ALL_EXECUTORS,
+    MAIN_EXECUTORS,
+    MODAL_EXECUTORS,
+    TaskCounter,
+    create_zarr,
+)
 
 
 @pytest.fixture()
@@ -28,6 +34,11 @@ def executor(request):
     return request.param
 
 
+@pytest.fixture(scope="module", params=ALL_EXECUTORS)
+def any_executor(request):
+    return request.param
+
+
 @pytest.fixture(scope="module", params=MODAL_EXECUTORS)
 def modal_executor(request):
     return request.param
@@ -540,3 +551,45 @@ def test_plan_scaling(tmp_path, factor):
 
     assert c.plan.num_tasks() > 0
     c.visualize(filename=tmp_path / "c")
+
+
+@pytest.mark.parametrize("compute_arrays_in_parallel", [True, False])
+def test_compute_arrays_in_parallel(spec, any_executor, compute_arrays_in_parallel):
+    from cubed.runtime.executors.python_async import AsyncPythonDagExecutor
+
+    if not isinstance(any_executor, AsyncPythonDagExecutor):
+        pytest.skip(f"{type(any_executor)} does not support compute_arrays_in_parallel")
+
+    a = cubed.random.random((10, 10), chunks=(5, 5), spec=spec)
+    b = cubed.random.random((10, 10), chunks=(5, 5), spec=spec)
+    c = xp.add(a, b)
+
+    c.compute(
+        executor=any_executor, compute_arrays_in_parallel=compute_arrays_in_parallel
+    )
+
+
+@pytest.mark.cloud
+@pytest.mark.parametrize("compute_arrays_in_parallel", [True, False])
+def test_compute_arrays_in_parallel_modal(modal_executor, compute_arrays_in_parallel):
+    from cubed.runtime.executors.modal_async import AsyncModalDagExecutor
+
+    if not isinstance(modal_executor, AsyncModalDagExecutor):
+        pytest.skip(
+            f"{type(modal_executor)} does not support compute_arrays_in_parallel"
+        )
+
+    tmp_path = "s3://cubed-unittest/parallel_pipelines"
+    spec = cubed.Spec(tmp_path, allowed_mem=100000)
+    try:
+        a = cubed.random.random((10, 10), chunks=(5, 5), spec=spec)
+        b = cubed.random.random((10, 10), chunks=(5, 5), spec=spec)
+        c = xp.add(a, b)
+
+        c.compute(
+            executor=modal_executor,
+            compute_arrays_in_parallel=compute_arrays_in_parallel,
+        )
+    finally:
+        fs = fsspec.open(tmp_path).fs
+        fs.rm(tmp_path, recursive=True)