diff --git a/cubed/runtime/executors/local.py b/cubed/runtime/executors/local.py index 588cdd71..5a13ef42 100644 --- a/cubed/runtime/executors/local.py +++ b/cubed/runtime/executors/local.py @@ -180,7 +180,10 @@ async def async_execute_dag( check_runtime_memory(spec, max_workers) if use_processes: max_tasks_per_child = kwargs.pop("max_tasks_per_child", None) - context = multiprocessing.get_context("spawn") + if isinstance(use_processes, str): + context = multiprocessing.get_context(use_processes) + else: + context = multiprocessing.get_context("spawn") # max_tasks_per_child is only supported from Python 3.11 if max_tasks_per_child is None: concurrent_executor = ProcessPoolExecutor( diff --git a/examples/matmul-random.py b/examples/matmul-random.py index 106464f4..07d5c9ac 100644 --- a/examples/matmul-random.py +++ b/examples/matmul-random.py @@ -1,3 +1,4 @@ +import dragon import logging import cubed @@ -11,9 +12,9 @@ logging.getLogger("urllib3.connectionpool").setLevel(logging.ERROR) if __name__ == "__main__": - # 200MB chunks - a = cubed.random.random((50000, 50000), chunks=(5000, 5000)) - b = cubed.random.random((50000, 50000), chunks=(5000, 5000)) + # 250KB chunks tested on a smaller machine + a = cubed.random.random((5000, 5000), chunks=(500, 500)) + b = cubed.random.random((5000, 5000), chunks=(500, 500)) c = xp.astype(a, xp.float32) d = xp.astype(b, xp.float32) e = xp.matmul(c, d) @@ -26,4 +27,15 @@ e, store=None, callbacks=[progress, hist, timeline_viz], + use_processes="dragon", ) + +# Example output: +# (dragon311v09) tantalum:~/cray/cubed/examples$ dragon -s matmul-random.py +# create-arrays 5/5 ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 100.0% 0:00:00 +# op-007 astype 100/100 ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 100.0% 0:00:00 +# op-008 astype 100/100 ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 100.0% 0:00:00 +# op-009 matmul 1000/1000 ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 100.0% 0:00:05 +# op-010 matmul 300/300 ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 100.0% 0:00:02 +# op-013 to_zarr 100/100 ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 100.0% 0:00:00 +# +++ head proc exited, code 0