diff --git a/src/executorlib/backend/executor_nested.py b/src/executorlib/backend/executor_nested.py new file mode 100644 index 000000000..2683aae1a --- /dev/null +++ b/src/executorlib/backend/executor_nested.py @@ -0,0 +1,113 @@ +from concurrent.futures import Future +from typing import Callable, Optional + +from executorlib.standalone.interactive.arguments import ( + get_future_objects_from_input, +) + + +class BackendExecutor: + def __init__(self): + self._tasks_dict = {} + self._future_dict = {} + + @property + def tasks(self): + return self._tasks_dict + + def batched( + self, + iterable: list[Future], + n: int, + ) -> list[Future]: + """ + Batch futures from the iterable into tuples of length n. The last batch may be shorter than n. + + Args: + iterable (list): list of future objects to batch based on which future objects finish first + n (int): badge size + + Returns: + list[Future]: list of future objects one for each batch + """ + raise NotImplementedError("The batched method is not implemented.") + + def map( + self, + fn: Callable, + *iterables, + timeout: Optional[float] = None, + chunksize: int = 1, + ): + """Returns an iterator equivalent to map(fn, iter). + + Args: + fn: A callable that will take as many arguments as there are + passed iterables. + timeout: The maximum number of seconds to wait. If None, then there + is no limit on the wait time. + chunksize: The size of the chunks the iterable will be broken into + before being passed to a child process. This argument is only + used by ProcessPoolExecutor; it is ignored by + ThreadPoolExecutor. + + Returns: + An iterator equivalent to: map(func, *iterables) but the calls may + be evaluated out-of-order. + + Raises: + TimeoutError: If the entire result iterator could not be generated + before the given timeout. + Exception: If fn(*args) raises for any values. + """ + raise NotImplementedError("The map method is not implemented.") + + def submit( # type: ignore + self, + fn: Callable, + /, + *args, + resource_dict: Optional[dict] = None, + **kwargs, + ) -> Future: + """ + Submits a callable to be executed with the given arguments. + + Schedules the callable to be executed as fn(*args, **kwargs) and returns + a Future instance representing the execution of the callable. + + Args: + fn (callable): function to submit for execution + args: arguments for the submitted function + kwargs: keyword arguments for the submitted function + resource_dict (dict): A dictionary of resources required by the task. With the following keys: + - cores (int): number of MPI cores to be used for each function call + - threads_per_core (int): number of OpenMP threads to be used for each function call + - gpus_per_core (int): number of GPUs per worker - defaults to 0 + - cwd (str/None): current working directory where the parallel python task is executed + - openmpi_oversubscribe (bool): adds the `--oversubscribe` command line flag (OpenMPI and + SLURM only) - default False + - slurm_cmd_args (list): Additional command line arguments for the srun call (SLURM only) + - error_log_file (str): Name of the error log file to use for storing exceptions raised + by the Python functions submitted to the Executor. + + Returns: + Future: A Future representing the given call. + """ + if resource_dict is None: + resource_dict = {} + f: Future = Future() + self._future_dict[f] = id(f) + self._tasks_dict[id(f)] = { + "fn": fn, + "args": args, # at the momemt the future objects are not removed from the args + "kwargs": kwargs, # at the momemt the future objects are not removed from the kwargs + "resource_dict": resource_dict, + "dependencies": [ + self._future_dict[future_dependency] + for future_dependency in get_future_objects_from_input( + args=args, kwargs=kwargs + ) + ], + } + return f diff --git a/src/executorlib/backend/interactive_parallel.py b/src/executorlib/backend/interactive_parallel.py index 7f968391d..ffd2e457a 100644 --- a/src/executorlib/backend/interactive_parallel.py +++ b/src/executorlib/backend/interactive_parallel.py @@ -6,6 +6,7 @@ import cloudpickle import zmq +from executorlib.backend.executor_nested import BackendExecutor from executorlib.standalone.error import backend_write_error_file from executorlib.standalone.interactive.backend import call_funct, parse_arguments from executorlib.standalone.interactive.communication import ( @@ -43,7 +44,11 @@ def main() -> None: host=argument_dict["host"], port=argument_dict["zmqport"] ) - memory = {"executorlib_worker_id": int(argument_dict["worker_id"])} + nested_executor = BackendExecutor() + memory = { + "executorlib_worker_id": int(argument_dict["worker_id"]), + "executorlib_executor": nested_executor, + } # required for flux interface - otherwise the current path is not included in the python path cwd = abspath(".") @@ -90,7 +95,13 @@ def main() -> None: else: # Send output if mpi_rank_zero: - interface_send(socket=socket, result_dict={"result": output_reply}) + interface_send( + socket=socket, + result_dict={ + "result": output_reply, + "tasks_nested": nested_executor.tasks, + }, + ) elif ( "init" in input_dict and input_dict["init"] diff --git a/src/executorlib/backend/interactive_serial.py b/src/executorlib/backend/interactive_serial.py index dc3971b0d..c694351bb 100644 --- a/src/executorlib/backend/interactive_serial.py +++ b/src/executorlib/backend/interactive_serial.py @@ -2,6 +2,7 @@ from os.path import abspath from typing import Optional +from executorlib.backend.executor_nested import BackendExecutor from executorlib.standalone.error import backend_write_error_file from executorlib.standalone.interactive.backend import call_funct, parse_arguments from executorlib.standalone.interactive.communication import ( @@ -29,7 +30,11 @@ def evaluate_cmd(argument_lst: Optional[list[str]] = None): host=argument_dict["host"], port=argument_dict["zmqport"] ) - memory = {"executorlib_worker_id": int(argument_dict["worker_id"])} + nested_executor = BackendExecutor() + memory = { + "executorlib_worker_id": int(argument_dict["worker_id"]), + "executorlib_executor": nested_executor, + } # required for flux interface - otherwise the current path is not included in the python path cwd = abspath(".") @@ -65,7 +70,13 @@ def evaluate_cmd(argument_lst: Optional[list[str]] = None): ) else: # Send output - interface_send(socket=socket, result_dict={"result": output}) + interface_send( + socket=socket, + result_dict={ + "result": output, + "tasks_nested": nested_executor.tasks, + }, + ) elif ( "init" in input_dict and input_dict["init"] diff --git a/src/executorlib/task_scheduler/base.py b/src/executorlib/task_scheduler/base.py index fd58e9586..124391fb9 100644 --- a/src/executorlib/task_scheduler/base.py +++ b/src/executorlib/task_scheduler/base.py @@ -50,6 +50,7 @@ def __init__( self._process_kwargs: dict = {} self._max_cores = max_cores self._future_queue: Optional[queue.Queue] = queue.Queue() + self._return_queue: Optional[queue.Queue] = queue.Queue() self._process: Optional[Union[Thread, list[Thread]]] = None self._validator = validator @@ -84,6 +85,8 @@ def info(self) -> Optional[dict]: meta_data_dict = self._process_kwargs.copy() if "future_queue" in meta_data_dict: del meta_data_dict["future_queue"] + if "return_queue" in meta_data_dict: + del meta_data_dict["return_queue"] if self._process is not None and isinstance(self._process, list): meta_data_dict["max_workers"] = len(self._process) return meta_data_dict @@ -102,6 +105,16 @@ def future_queue(self) -> Optional[queue.Queue]: """ return self._future_queue + @property + def return_queue(self) -> Optional[queue.Queue]: + """ + Get the return queue. + + Returns: + queue.Queue: The return queue. + """ + return self._return_queue + def batched( self, iterable: list[Future], @@ -238,8 +251,10 @@ def shutdown(self, wait: bool = True, *, cancel_futures: bool = False): if isinstance(self._process, Thread): self._process.join() self._future_queue.join() + self._return_queue.join() self._process = None self._future_queue = None + self._return_queue = None def _set_process(self, process: Thread): """ diff --git a/src/executorlib/task_scheduler/interactive/blockallocation.py b/src/executorlib/task_scheduler/interactive/blockallocation.py index 02501e786..b7bd5e7ef 100644 --- a/src/executorlib/task_scheduler/interactive/blockallocation.py +++ b/src/executorlib/task_scheduler/interactive/blockallocation.py @@ -75,6 +75,7 @@ def __init__( max_cores=executor_kwargs.get("max_cores"), validator=validator ) executor_kwargs["future_queue"] = self._future_queue + executor_kwargs["return_queue"] = self._return_queue executor_kwargs["spawner"] = spawner executor_kwargs["queue_join_on_shutdown"] = False executor_kwargs["restart_limit"] = restart_limit @@ -217,6 +218,7 @@ def _set_process(self, process: list[Thread]): # type: ignore def _execute_multiple_tasks( future_queue: queue.Queue, + return_queue: queue.Queue, cores: int = 1, spawner: type[BaseSpawner] = MpiExecSpawner, hostname_localhost: Optional[bool] = None, @@ -240,6 +242,7 @@ def _execute_multiple_tasks( Args: future_queue (queue.Queue): task queue of dictionary objects which are submitted to the parallel process + return_queue (queue.Queue): task queue of dictionary objects which are submitted to the parallel process cores (int): defines the total number of MPI ranks to use spawner (BaseSpawner): Spawner to start process on selected compute resources hostname_localhost (boolean): use localhost instead of the hostname to establish the zmq connection. In the @@ -317,7 +320,7 @@ def _execute_multiple_tasks( f.set_exception(exception=interface_initialization_exception) else: # The interface failed during the execution - interface.status = execute_task_dict( + interface.status, nested_task_dict = execute_task_dict( task_dict=task_dict, future_obj=f, interface=interface, @@ -329,6 +332,8 @@ def _execute_multiple_tasks( reset_task_dict( future_obj=f, future_queue=future_queue, task_dict=task_dict ) + elif len(nested_task_dict) > 0: + return_queue.put(nested_task_dict) task_done(future_queue=future_queue) diff --git a/src/executorlib/task_scheduler/interactive/dependency.py b/src/executorlib/task_scheduler/interactive/dependency.py index e781195d5..20ecdd8f0 100644 --- a/src/executorlib/task_scheduler/interactive/dependency.py +++ b/src/executorlib/task_scheduler/interactive/dependency.py @@ -257,6 +257,10 @@ def _execute_tasks_with_dependencies( task_dict = future_queue.get_nowait() except queue.Empty: task_dict = None + try: + task_return_dict = executor.return_queue.get_nowait() + except (queue.Empty, AttributeError): + task_return_dict = None if ( # shutdown the executor task_dict is not None and "shutdown" in task_dict and task_dict["shutdown"] ): @@ -324,6 +328,16 @@ def _execute_tasks_with_dependencies( executor_queue=executor_queue, refresh_rate=refresh_rate, ) + elif task_return_dict is not None: + # we need this future dict later on to resolve the dependencies + future_lookup_dict = {} + for k, v in task_return_dict.items(): + future_lookup_dict[k] = executor.submit( + fn=v["fn"], + *v["args"], + **v["kwargs"], + resource_dict=v["resource_dict"], + ) else: # If there is nothing else to do, sleep for a moment sleep(refresh_rate) diff --git a/src/executorlib/task_scheduler/interactive/onetoone.py b/src/executorlib/task_scheduler/interactive/onetoone.py index 52c06976e..3d2bb56a8 100644 --- a/src/executorlib/task_scheduler/interactive/onetoone.py +++ b/src/executorlib/task_scheduler/interactive/onetoone.py @@ -61,6 +61,7 @@ def __init__( executor_kwargs.update( { "future_queue": self._future_queue, + "return_queue": self._return_queue, "spawner": spawner, "max_cores": max_cores, "max_workers": max_workers, @@ -77,6 +78,7 @@ def __init__( def _execute_single_task( future_queue: queue.Queue, + return_queue: queue.Queue, spawner: type[BaseSpawner] = MpiExecSpawner, max_cores: Optional[int] = None, max_workers: Optional[int] = None, @@ -88,6 +90,7 @@ def _execute_single_task( Args: future_queue (queue.Queue): task queue of dictionary objects which are submitted to the parallel process + return_queue (queue.Queue): task queue of dictionary objects which are submitted to the parallel process spawner (BaseSpawner): Interface to start process on selected compute resources max_cores (int): defines the number cores which can be used in parallel max_workers (int): for backwards compatibility with the standard library, max_workers also defines the number of @@ -116,6 +119,7 @@ def _execute_single_task( elif "fn" in task_dict and "future" in task_dict: process, active_task_dict = _wrap_execute_task_in_separate_process( task_dict=task_dict, + return_queue=return_queue, active_task_dict=active_task_dict, spawner=spawner, executor_kwargs=kwargs, @@ -162,6 +166,7 @@ def _wait_for_free_slots( def _wrap_execute_task_in_separate_process( task_dict: dict, + return_queue: queue.Queue, active_task_dict: dict, spawner: type[BaseSpawner], executor_kwargs: dict, @@ -211,6 +216,7 @@ def _wrap_execute_task_in_separate_process( task_kwargs.update( { "task_dict": task_dict, + "return_queue": return_queue, "spawner": spawner, "hostname_localhost": hostname_localhost, "future_obj": f, @@ -226,6 +232,7 @@ def _wrap_execute_task_in_separate_process( def _execute_task_in_thread( task_dict: dict, + return_queue: queue.Queue, future_obj: Future, cores: int = 1, spawner: type[BaseSpawner] = MpiExecSpawner, @@ -262,7 +269,7 @@ def _execute_task_in_thread( worker_id (int): Communicate the worker which ID was assigned to it for future reference and resource distribution. """ - if not execute_task_dict( + status, nested_task_dict = execute_task_dict( task_dict=task_dict, future_obj=future_obj, interface=interface_bootup( @@ -277,7 +284,10 @@ def _execute_task_in_thread( cache_directory=cache_directory, cache_key=cache_key, error_log_file=error_log_file, - ): + ) + if status is False: future_obj.set_exception( ExecutorlibSocketError("SocketInterface crashed during execution.") ) + elif len(nested_task_dict) > 0: + return_queue.put(nested_task_dict) diff --git a/src/executorlib/task_scheduler/interactive/shared.py b/src/executorlib/task_scheduler/interactive/shared.py index badca5eed..a4abf9f65 100644 --- a/src/executorlib/task_scheduler/interactive/shared.py +++ b/src/executorlib/task_scheduler/interactive/shared.py @@ -20,7 +20,7 @@ def execute_task_dict( cache_directory: Optional[str] = None, cache_key: Optional[str] = None, error_log_file: Optional[str] = None, -) -> bool: +) -> tuple[bool, dict]: """ Execute the task in the task_dict by communicating it via the interface. @@ -37,6 +37,7 @@ def execute_task_dict( Returns: bool: True if the task was submitted successfully, False otherwise. + dict: Dictionary of nested tasks if any were returned by the executed function, otherwise an empty dictionary. """ if not future_obj.done() and future_obj.set_running_or_notify_cancel(): if error_log_file is not None: @@ -54,7 +55,7 @@ def execute_task_dict( future_obj=future_obj, ) else: - return True + return True, {} def task_done(future_queue: queue.Queue): @@ -88,7 +89,7 @@ def reset_task_dict(future_obj: Future, future_queue: queue.Queue, task_dict: di def _execute_task_without_cache( interface: SocketInterface, task_dict: dict, future_obj: Future -) -> bool: +) -> tuple[bool, dict]: """ Execute the task in the task_dict by communicating it via the interface. @@ -101,14 +102,17 @@ def _execute_task_without_cache( Returns: bool: True if the task was submitted successfully, False otherwise. """ + new_task_dict = {} output = interface.send_and_receive_dict(input_dict=task_dict) if "result" in output: future_obj.set_result(output["result"]) + if "tasks_nested" in output: + new_task_dict = output["tasks_nested"] elif isinstance(output["error"], ExecutorlibSocketError): - return False + return False, new_task_dict else: future_obj.set_exception(exception=output["error"]) - return True + return True, new_task_dict def _execute_task_with_cache( @@ -117,7 +121,7 @@ def _execute_task_with_cache( future_obj: Future, cache_directory: str, cache_key: Optional[str] = None, -) -> bool: +) -> tuple[bool, dict]: """ Execute the task in the task_dict by communicating it via the interface using the cache in the cache directory. @@ -139,6 +143,7 @@ def _execute_task_with_cache( resource_dict=task_dict.get("resource_dict", {}), cache_key=cache_key, ) + new_task_dict = {} file_name = os.path.abspath(os.path.join(cache_directory, task_key + "_o.h5")) if file_name not in get_cache_files(cache_directory=cache_directory): time_start = time.time() @@ -148,11 +153,13 @@ def _execute_task_with_cache( data_dict["runtime"] = time.time() - time_start dump(file_name=file_name, data_dict=data_dict) future_obj.set_result(output["result"]) + if "tasks_nested" in output: + new_task_dict = output["tasks_nested"] elif isinstance(output["error"], ExecutorlibSocketError): - return False + return False, new_task_dict else: future_obj.set_exception(exception=output["error"]) else: _, _, result = get_output(file_name=file_name) future_obj.set_result(result) - return True + return True, new_task_dict diff --git a/tests/unit/backend/test_interactive_serial.py b/tests/unit/backend/test_interactive_serial.py index 544bfc067..281da8786 100644 --- a/tests/unit/backend/test_interactive_serial.py +++ b/tests/unit/backend/test_interactive_serial.py @@ -52,7 +52,7 @@ def test_main_as_thread(self): t.start() submit(socket=socket) self.assertEqual(cloudpickle.loads(socket.recv()), {"result": True}) - self.assertEqual(cloudpickle.loads(socket.recv()), {"result": 7}) + self.assertEqual(cloudpickle.loads(socket.recv()), {'result': 7, 'tasks_nested': {}}) self.assertEqual(cloudpickle.loads(socket.recv()), {"result": True}) socket.close() context.term() @@ -96,7 +96,7 @@ def test_submit_as_thread(self): t.start() evaluate_cmd(argument_lst=["--zmqport", str(port)]) self.assertEqual(cloudpickle.loads(socket.recv()), {"result": True}) - self.assertEqual(cloudpickle.loads(socket.recv()), {"result": 7}) + self.assertEqual(cloudpickle.loads(socket.recv()), {'result': 7, 'tasks_nested': {}}) self.assertEqual(cloudpickle.loads(socket.recv()), {"result": True}) socket.close() context.term() diff --git a/tests/unit/executor/test_single_shell_executor.py b/tests/unit/executor/test_single_shell_executor.py index 1e60f3382..dc8963ec4 100644 --- a/tests/unit/executor/test_single_shell_executor.py +++ b/tests/unit/executor/test_single_shell_executor.py @@ -34,6 +34,7 @@ def test_execute_single_task(self): self.assertFalse(f.done()) _execute_multiple_tasks( future_queue=test_queue, + return_queue=queue.Queue(), cores=1, openmpi_oversubscribe=False, spawner=MpiExecSpawner, @@ -60,6 +61,7 @@ def test_wrong_error(self): with self.assertRaises(TypeError): _execute_multiple_tasks( future_queue=test_queue, + return_queue=queue.Queue(), cores=1, openmpi_oversubscribe=False, spawner=MpiExecSpawner, @@ -87,6 +89,7 @@ def test_broken_executable(self): with self.assertRaises(FileNotFoundError): _execute_multiple_tasks( future_queue=test_queue, + return_queue=queue.Queue(), cores=1, openmpi_oversubscribe=False, spawner=MpiExecSpawner, diff --git a/tests/unit/executor/test_single_shell_interactive.py b/tests/unit/executor/test_single_shell_interactive.py index 831cb06b1..991fca205 100644 --- a/tests/unit/executor/test_single_shell_interactive.py +++ b/tests/unit/executor/test_single_shell_interactive.py @@ -90,6 +90,7 @@ def test_execute_single_task(self): self.assertFalse(future_pattern.done()) _execute_multiple_tasks( future_queue=test_queue, + return_queue=queue.Queue(), cores=1, openmpi_oversubscribe=False, spawner=MpiExecSpawner, diff --git a/tests/unit/standalone/interactive/test_spawner.py b/tests/unit/standalone/interactive/test_spawner.py index 1af872cdc..d2365cead 100644 --- a/tests/unit/standalone/interactive/test_spawner.py +++ b/tests/unit/standalone/interactive/test_spawner.py @@ -263,6 +263,7 @@ def test_execute_task(self): cloudpickle_register(ind=1) _execute_multiple_tasks( future_queue=q, + return_queue=Queue(), cores=1, openmpi_oversubscribe=False, spawner=MpiExecSpawner, @@ -445,6 +446,7 @@ def test_execute_task_failed_no_argument(self): cloudpickle_register(ind=1) _execute_multiple_tasks( future_queue=q, + return_queue=Queue(), cores=1, openmpi_oversubscribe=False, spawner=MpiExecSpawner, @@ -461,6 +463,7 @@ def test_execute_task_failed_wrong_argument(self): cloudpickle_register(ind=1) _execute_multiple_tasks( future_queue=q, + return_queue=Queue(), cores=1, openmpi_oversubscribe=False, spawner=MpiExecSpawner, @@ -477,6 +480,7 @@ def test_execute_task(self): cloudpickle_register(ind=1) _execute_multiple_tasks( future_queue=q, + return_queue=Queue(), cores=1, openmpi_oversubscribe=False, spawner=MpiExecSpawner, @@ -495,6 +499,7 @@ def test_execute_task_parallel(self): cloudpickle_register(ind=1) _execute_multiple_tasks( future_queue=q, + return_queue=Queue(), cores=2, openmpi_oversubscribe=False, spawner=MpiExecSpawner, @@ -518,6 +523,7 @@ def test_execute_task_cache(self): cloudpickle_register(ind=1) _execute_multiple_tasks( future_queue=q, + return_queue=Queue(), cores=1, openmpi_oversubscribe=False, spawner=MpiExecSpawner, @@ -537,6 +543,7 @@ def test_execute_task_cache_failed_no_argument(self): cloudpickle_register(ind=1) _execute_multiple_tasks( future_queue=q, + return_queue=Queue(), cores=1, openmpi_oversubscribe=False, spawner=MpiExecSpawner, diff --git a/tests/unit/task_scheduler/interactive/test_shared.py b/tests/unit/task_scheduler/interactive/test_shared.py index bd691d3bf..cf8343a8e 100644 --- a/tests/unit/task_scheduler/interactive/test_shared.py +++ b/tests/unit/task_scheduler/interactive/test_shared.py @@ -43,7 +43,7 @@ def test_execute_task_sum(self): cache_directory=None, cache_key=None, error_log_file=None, - ) + )[0] self.assertTrue(result) self.assertTrue(f.done()) self.assertEqual(f.result(), 3) @@ -71,7 +71,7 @@ def test_execute_task_done(self): cache_directory=None, cache_key=None, error_log_file=None, - ) + )[0] self.assertTrue(result) self.assertTrue(f.done()) self.assertEqual(f.result(), 5) @@ -98,7 +98,7 @@ def test_execute_task_error(self): cache_directory=None, cache_key=None, error_log_file=None, - ) + )[0] self.assertFalse(result) self.assertFalse(f.done()) @@ -132,7 +132,7 @@ def test_execute_task_sum(self): cache_directory="cache_execute_task", cache_key=None, error_log_file=None, - ) + )[0] self.assertTrue(result) self.assertTrue(f.done()) self.assertEqual(f.result(), 3) @@ -160,7 +160,7 @@ def test_execute_task_done(self): cache_directory="cache_execute_task", cache_key=None, error_log_file=None, - ) + )[0] self.assertTrue(result) self.assertTrue(f.done()) self.assertEqual(f.result(), 5) @@ -187,6 +187,6 @@ def test_execute_task_error(self): cache_directory="cache_execute_task", cache_key=None, error_log_file=None, - ) + )[0] self.assertFalse(result) self.assertFalse(f.done()) \ No newline at end of file diff --git a/tests/unit/task_scheduler/interactive/test_spawner_flux.py b/tests/unit/task_scheduler/interactive/test_spawner_flux.py index 5825f0820..d31e7a89a 100644 --- a/tests/unit/task_scheduler/interactive/test_spawner_flux.py +++ b/tests/unit/task_scheduler/interactive/test_spawner_flux.py @@ -113,6 +113,7 @@ def test_execute_task(self): cloudpickle_register(ind=1) _execute_multiple_tasks( future_queue=q, + return_queue=Queue(), cores=1, flux_executor=self.flux_executor, spawner=FluxPythonSpawner, @@ -128,6 +129,7 @@ def test_execute_task_threads(self): cloudpickle_register(ind=1) _execute_multiple_tasks( future_queue=q, + return_queue=Queue(), cores=1, threads_per_core=1, flux_executor=self.flux_executor,