submit multiple tasks to a broker but execute only one

pympipool/interfaces/fluxbroker.py

@@ -5,7 +5,6 @@
 from time import sleep
 
 from pympipool.shared.broker import (
-    MetaExecutorFuture,
     _get_future_done,
     _execute_task_dict,
 )
@@ -219,16 +218,13 @@ def _get_executor_list(
     cwd=None,
     executor=None,
 ):
-    return [
-        MetaExecutorFuture(
-            future=_get_future_done(),
-            executor=SingleTaskExecutor(
-                cores=cores_per_worker,
-                gpus_per_task=int(gpus_per_worker / cores_per_worker),
-                init_function=init_function,
-                cwd=cwd,
-                executor=executor,
-            ),
+    return {
+        _get_future_done(): SingleTaskExecutor(
+            cores=cores_per_worker,
+            gpus_per_task=int(gpus_per_worker / cores_per_worker),
+            init_function=init_function,
+            cwd=cwd,
+            executor=executor,
         )
         for _ in range(max_workers)
-    ]
+    }

pympipool/shared/broker.py

@@ -5,31 +5,6 @@
 from pympipool.interfaces.taskexecutor import Executor
 
 
-class MetaExecutorFuture(object):
-    def __init__(self, future, executor):
-        self.future = future
-        self.executor = executor
-
-    @property
-    def _condition(self):
-        return self.future._condition
-
-    @property
-    def _state(self):
-        return self.future._state
-
-    @property
-    def _waiters(self):
-        return self.future._waiters
-
-    def done(self):
-        return self.future.done()
-
-    def submit(self, task_dict):
-        self.future = task_dict["future"]
-        self.executor._future_queue.put(task_dict)
-
-
 def executor_broker(
     future_queue,
     max_workers,
@@ -71,12 +46,14 @@ def executor_broker(
 
 def _execute_task_dict(task_dict, meta_future_lst):
     if "fn" in task_dict.keys():
-        meta_future = next(as_completed(meta_future_lst))
-        meta_future.submit(task_dict=task_dict)
+        meta_future = next(as_completed(meta_future_lst.keys()))
+        executor = meta_future_lst.pop(meta_future)
+        executor._future_queue.put(task_dict)
+        meta_future_lst[task_dict["future"]] = executor
         return True
     elif "shutdown" in task_dict.keys() and task_dict["shutdown"]:
-        for meta in meta_future_lst:
-            meta.executor.shutdown(wait=task_dict["wait"])
+        for executor in meta_future_lst.values():
+            executor.shutdown(wait=task_dict["wait"])
         return False
     else:
         raise ValueError("Unrecognized Task in task_dict: ", task_dict)
@@ -94,23 +71,20 @@ def _get_executor_list(
     queue_adapter=None,
     queue_adapter_kwargs=None,
 ):
-    return [
-        MetaExecutorFuture(
-            future=_get_future_done(),
-            executor=Executor(
-                cores=cores_per_worker,
-                gpus_per_task=int(gpus_per_worker / cores_per_worker),
-                oversubscribe=oversubscribe,
-                enable_flux_backend=enable_flux_backend,
-                enable_slurm_backend=enable_slurm_backend,
-                init_function=init_function,
-                cwd=cwd,
-                queue_adapter=queue_adapter,
-                queue_adapter_kwargs=queue_adapter_kwargs,
-            ),
+    return {
+        _get_future_done(): Executor(
+            cores=cores_per_worker,
+            gpus_per_task=int(gpus_per_worker / cores_per_worker),
+            oversubscribe=oversubscribe,
+            enable_flux_backend=enable_flux_backend,
+            enable_slurm_backend=enable_slurm_backend,
+            init_function=init_function,
+            cwd=cwd,
+            queue_adapter=queue_adapter,
+            queue_adapter_kwargs=queue_adapter_kwargs,
         )
         for _ in range(max_workers)
-    ]
+    }
 
 
 def _get_future_done():

tests/test_meta.py

@@ -31,12 +31,14 @@ def test_get_future_done(self):
 
 class TestMetaExecutorFuture(unittest.TestCase):
     def test_meta_executor_future(self):
-        meta_future = _get_executor_list(max_workers=1)[0]
-        self.assertTrue(isinstance(meta_future.future, Future))
-        self.assertTrue(isinstance(meta_future.executor, Executor))
-        self.assertTrue(meta_future.done())
-        self.assertEqual(meta_future, next(as_completed([meta_future])))
-        meta_future.submit(task_dict={"shutdown": True, "wait": True, "future": _get_future_done()})
+        meta_future = _get_executor_list(max_workers=1)
+        future_obj = list(meta_future.keys())[0]
+        executor_obj = list(meta_future.values())[0]
+        self.assertTrue(isinstance(future_obj, Future))
+        self.assertTrue(isinstance(executor_obj, Executor))
+        self.assertTrue(future_obj.done())
+        self.assertEqual(future_obj, next(as_completed(meta_future.keys())))
+        executor_obj.shutdown(wait=True)
 
     def test_execute_task_dict(self):
         meta_future_lst = _get_executor_list(max_workers=1)
@@ -52,6 +54,12 @@ def test_execute_task_dict(self):
             meta_future_lst=meta_future_lst
         ))
 
+    def test_execute_task_dict_error(self):
+        meta_future_lst = _get_executor_list(max_workers=1)
+        with self.assertRaises(ValueError):
+            _execute_task_dict(task_dict={}, meta_future_lst=meta_future_lst)
+        list(meta_future_lst.values())[0].shutdown(wait=True)
+
     def test_executor_broker(self):
         q = Queue()
         f = Future()
@@ -73,6 +81,15 @@ def test_meta_executor_serial(self):
             self.assertTrue(fs_1.done())
             self.assertTrue(fs_2.done())
 
+    def test_meta_executor_single(self):
+        with HPCExecutor(max_workers=1) as exe:
+            fs_1 = exe.submit(calc, 1)
+            fs_2 = exe.submit(calc, 2)
+            self.assertEqual(fs_1.result(), 1)
+            self.assertEqual(fs_2.result(), 2)
+            self.assertTrue(fs_1.done())
+            self.assertTrue(fs_2.done())
+
     def test_meta_executor_parallel(self):
         with HPCExecutor(max_workers=1, cores_per_worker=2) as exe:
             fs_1 = exe.submit(mpi_funct, 1)

tests/test_worker.py

@@ -40,6 +40,15 @@ def test_pool_serial(self):
             self.assertEqual(len(p), 0)
         self.assertEqual(output.result(), np.array(4))
 
+    def test_executor_multi_submission(self):
+        with Executor(cores=1) as p:
+            fs_1 = p.submit(calc, i=2)
+            fs_2 = p.submit(calc, i=2)
+            self.assertEqual(fs_1.result(), np.array(4))
+            self.assertEqual(fs_2.result(), np.array(4))
+            self.assertTrue(fs_1.done())
+            self.assertTrue(fs_2.done())
+
     def test_shutdown(self):
         p = Executor(cores=1)
         fs1 = p.submit(sleep_one, i=2)