From 865aa6c51202e90d53bcf0374979e6ccd8b11b6a Mon Sep 17 00:00:00 2001 From: Davin Potts Date: Sun, 18 Aug 2024 18:07:43 -0500 Subject: [PATCH 1/2] Permit specification of multiprocessing start methods other than "spawn" (e.g. "dragon" or "fork"). --- cubed/runtime/executors/local.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/cubed/runtime/executors/local.py b/cubed/runtime/executors/local.py index 588cdd71f..5a13ef42c 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( From cc9bbd4d129024c427e2cd0c657fc0825e0ee301 Mon Sep 17 00:00:00 2001 From: Davin Potts Date: Sun, 18 Aug 2024 18:47:10 -0500 Subject: [PATCH 2/2] Modified version of existing example to use "dragon" start method. --- examples/matmul-random.py | 18 +++++++++++++++--- 1 file changed, 15 insertions(+), 3 deletions(-) diff --git a/examples/matmul-random.py b/examples/matmul-random.py index 106464f43..07d5c9acf 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