diff --git a/src/executorlib/standalone/hdf.py b/src/executorlib/standalone/hdf.py index ff7be59c..8c61b659 100644 --- a/src/executorlib/standalone/hdf.py +++ b/src/executorlib/standalone/hdf.py @@ -31,18 +31,22 @@ def dump(file_name: Optional[str], data_dict: dict) -> None: """ if file_name is not None: file_name_abs = os.path.abspath(file_name) + # Serialize all values before opening the file, so a serialization error does + # not leave an incomplete HDF5 file behind. + binary_dict = { + group_dict[data_key]: cloudpickle.dumps(data_value) + for data_key, data_value in data_dict.items() + if data_key in group_dict + } os.makedirs(os.path.dirname(file_name_abs), exist_ok=True) with h5py.File(file_name_abs, "a") as fname: - for data_key, data_value in data_dict.items(): - if data_key in group_dict: - with contextlib.suppress(ValueError): - fname.create_dataset( - name="/" + group_dict[data_key], - data=np.frombuffer( - cloudpickle.dumps(data_value), dtype=np.uint8 - ), - compression="gzip", - ) + for group_name, binary in binary_dict.items(): + with contextlib.suppress(ValueError): + fname.create_dataset( + name="/" + group_name, + data=np.frombuffer(binary, dtype=np.uint8), + compression="gzip", + ) def load(file_name: str) -> dict: diff --git a/src/executorlib/standalone/select.py b/src/executorlib/standalone/select.py index e3e3f575..2df2f19c 100644 --- a/src/executorlib/standalone/select.py +++ b/src/executorlib/standalone/select.py @@ -21,9 +21,10 @@ def __init__(self, future: Future, selector: int | str): future (Future): The underlying future whose result is a collection. selector (int | str): Index or key used to extract an element from the result. """ + # Future.__init__() is not called, as its attributes would be forwarded to the + # wrapped future by __setattr__() and reset the state of an already completed future. self._future = future self._selector = selector - super().__init__() def __getattr__(self, attr: str) -> Any: """Delegate attribute access to the wrapped future.""" diff --git a/src/executorlib/task_scheduler/file/shared.py b/src/executorlib/task_scheduler/file/shared.py index c65865fa..75749dff 100644 --- a/src/executorlib/task_scheduler/file/shared.py +++ b/src/executorlib/task_scheduler/file/shared.py @@ -1,7 +1,7 @@ import contextlib import os import queue -from concurrent.futures import Future +from concurrent.futures import CancelledError, Future from time import sleep from typing import Any, Callable, Optional @@ -118,83 +118,93 @@ def execute_tasks_h5( future_queue.join() break elif task_dict is not None: - task_args, task_kwargs, future_wait_key_lst = _convert_args_and_kwargs( - task_dict=task_dict, - memory_dict=memory_dict, - file_name_dict=file_name_dict, - ) - task_resource_dict, cache_key, cache_directory, error_log_file = ( - _get_task_input( - task_resource_dict=task_dict["resource_dict"].copy(), - executor_kwargs=executor_kwargs, + try: + task_args, task_kwargs, future_wait_key_lst = _convert_args_and_kwargs( + task_dict=task_dict, + memory_dict=memory_dict, + file_name_dict=file_name_dict, + duplicate_dict=duplicate_dict, ) - ) - task_key, data_dict, serialize_exception = serialize_funct( - fn=task_dict["fn"], - fn_args=task_args, - fn_kwargs=task_kwargs, - resource_dict=task_resource_dict, - cache_key=cache_key, - ) - data_dict["error_log_file"] = error_log_file - if serialize_exception is not None: - task_dict["future"].set_exception(serialize_exception) - elif task_key not in memory_dict: - if os.path.join( - cache_directory, task_key + "_o.h5" - ) not in get_cache_files(cache_directory=cache_directory): - file_name = os.path.join(cache_directory, task_key + "_i.h5") - if not disable_dependencies: - task_dependent_lst = [ - process_dict[k] - for k in future_wait_key_lst - if k in process_dict - ] - else: - if len(future_wait_key_lst) > 0: - task_dict["future"].set_exception( - ValueError( - "Future objects are not supported as input if disable_dependencies=True." + task_resource_dict, cache_key, cache_directory, error_log_file = ( + _get_task_input( + task_resource_dict=task_dict["resource_dict"].copy(), + executor_kwargs=executor_kwargs, + ) + ) + task_key, data_dict, serialize_exception = serialize_funct( + fn=task_dict["fn"], + fn_args=task_args, + fn_kwargs=task_kwargs, + resource_dict=task_resource_dict, + cache_key=cache_key, + ) + data_dict["error_log_file"] = error_log_file + if serialize_exception is not None: + task_dict["future"].set_exception(serialize_exception) + elif task_key not in memory_dict: + if os.path.join( + cache_directory, task_key + "_o.h5" + ) not in get_cache_files(cache_directory=cache_directory): + file_name = os.path.join(cache_directory, task_key + "_i.h5") + if not disable_dependencies: + task_dependent_lst = [ + process_dict[k] + for k in future_wait_key_lst + if k in process_dict + ] + else: + if len(future_wait_key_lst) > 0: + task_dict["future"].set_exception( + ValueError( + "Future objects are not supported as input if disable_dependencies=True." + ) ) - ) - task_dependent_lst = [] - try: - process_dict[task_key] = execute_function( - command=get_cache_execute_command( + continue + task_dependent_lst = [] + try: + process_dict[task_key] = execute_function( + command=get_cache_execute_command( + file_name=file_name, + cores=task_resource_dict["cores"], + backend=backend, + exclusive=task_resource_dict.get( + "exclusive", False + ), + openmpi_oversubscribe=task_resource_dict.get( + "openmpi_oversubscribe", False + ), + pmi_mode=pmi_mode, + ), file_name=file_name, - cores=task_resource_dict["cores"], + data_dict=data_dict, + task_dependent_lst=task_dependent_lst, + resource_dict=task_resource_dict, + config_directory=pysqa_config_directory, backend=backend, - exclusive=task_resource_dict.get("exclusive", False), - openmpi_oversubscribe=task_resource_dict.get( - "openmpi_oversubscribe", False - ), - pmi_mode=pmi_mode, - ), - file_name=file_name, - data_dict=data_dict, - task_dependent_lst=task_dependent_lst, - resource_dict=task_resource_dict, - config_directory=pysqa_config_directory, - backend=backend, - cache_directory=cache_directory, - ) - except Exception as e: - task_dict["future"].set_exception(e) - memory_dict[task_key] = task_dict["future"] - future_queue.task_done() - continue - file_name = os.path.join(cache_directory, task_key + "_o.h5") - file_name_dict[task_key] = file_name - queue_id = get_queue_id(file_name=file_name) - if queue_id is not None: - process_dict[task_key] = queue_id - memory_dict[task_key] = task_dict["future"] - cache_dir_dict[task_key] = cache_directory - elif memory_dict[task_key] != task_dict["future"]: - if task_key not in duplicate_dict: - duplicate_dict[task_key] = [] - duplicate_dict[task_key].append(task_dict["future"]) - future_queue.task_done() + cache_directory=cache_directory, + ) + except Exception as e: + task_dict["future"].set_exception(e) + memory_dict[task_key] = task_dict["future"] + continue + file_name = os.path.join(cache_directory, task_key + "_o.h5") + file_name_dict[task_key] = file_name + queue_id = get_queue_id(file_name=file_name) + if queue_id is not None: + process_dict[task_key] = queue_id + memory_dict[task_key] = task_dict["future"] + cache_dir_dict[task_key] = cache_directory + elif memory_dict[task_key] != task_dict["future"]: + if task_key not in duplicate_dict: + duplicate_dict[task_key] = [] + duplicate_dict[task_key].append(task_dict["future"]) + except Exception as e: + # Fail only the current task, never the scheduler thread, e.g. when + # one of the input futures of this task failed before submission. + if not task_dict["future"].done(): + task_dict["future"].set_exception(e) + finally: + future_queue.task_done() else: memory_dict = _refresh_memory_dict( memory_dict=memory_dict, @@ -308,7 +318,10 @@ def _update_future( def _convert_args_and_kwargs( - task_dict: dict, memory_dict: dict, file_name_dict: dict + task_dict: dict, + memory_dict: dict, + file_name_dict: dict, + duplicate_dict: Optional[dict] = None, ) -> tuple[list, dict, list]: """ Convert the arguments and keyword arguments in a task dictionary to the appropriate types. @@ -317,59 +330,84 @@ def _convert_args_and_kwargs( task_dict (dict): The task dictionary containing the arguments and keyword arguments. memory_dict (dict): The dictionary mapping future objects to their associated task keys. file_name_dict (dict): The dictionary mapping task keys to their corresponding file names. + duplicate_dict (dict, optional): The dictionary mapping task keys to their associated duplicate future objects. Returns: Tuple[list, dict, list]: A tuple containing the converted arguments, converted keyword arguments, and a list of future wait keys. + Raises: + Exception: The exception of an input future which already failed, so the dependent task fails as well. + """ - task_args = [] - task_kwargs = {} - future_wait_key_lst = [] - for arg in task_dict["args"]: - selector = None - if isinstance(arg, Future): - if hasattr(arg, "_future") and hasattr(arg, "_selector"): - selector = arg._selector - future = arg._future - else: - future = arg - match_found = False - for k, v in memory_dict.items(): - if future == v: - task_args.append( - FutureItem(file_name=file_name_dict[k], selector=selector) - ) - future_wait_key_lst.append(k) - match_found = True - break - if not match_found: - task_args.append(future.result()) - else: - task_args.append(arg) - for key, arg in task_dict["kwargs"].items(): - selector = None - if isinstance(arg, Future): - if hasattr(arg, "_future") and hasattr(arg, "_selector"): - selector = arg._selector - future = arg._future - else: - future = arg - match_found = False - for k, v in memory_dict.items(): - if future == v: - task_kwargs[key] = FutureItem( - file_name=file_name_dict[k], selector=selector - ) - future_wait_key_lst.append(k) - match_found = True - break - if not match_found: - task_kwargs[key] = future.result() - else: - task_kwargs[key] = arg + if duplicate_dict is None: + duplicate_dict = {} + future_wait_key_lst: list = [] + task_args = [ + _convert_future( + arg=arg, + memory_dict=memory_dict, + file_name_dict=file_name_dict, + duplicate_dict=duplicate_dict, + future_wait_key_lst=future_wait_key_lst, + ) + for arg in task_dict["args"] + ] + task_kwargs = { + key: _convert_future( + arg=arg, + memory_dict=memory_dict, + file_name_dict=file_name_dict, + duplicate_dict=duplicate_dict, + future_wait_key_lst=future_wait_key_lst, + ) + for key, arg in task_dict["kwargs"].items() + } return task_args, task_kwargs, future_wait_key_lst +def _convert_future( + arg: Any, + memory_dict: dict, + file_name_dict: dict, + duplicate_dict: dict, + future_wait_key_lst: list, +) -> Any: + """ + Convert a single argument, replacing future objects by FutureItem objects or their results. + + Args: + arg (Any): The argument to convert. + memory_dict (dict): The dictionary mapping task keys to their associated future objects. + file_name_dict (dict): The dictionary mapping task keys to their corresponding file names. + duplicate_dict (dict): The dictionary mapping task keys to their associated duplicate future objects. + future_wait_key_lst (list): List of task keys the current task depends on, extended in place. + + Returns: + Any: The converted argument. + + Raises: + Exception: The exception of the input future, if it already failed. + + """ + if not isinstance(arg, Future): + return arg + if hasattr(arg, "_future") and hasattr(arg, "_selector"): + selector = arg._selector + future = arg._future + else: + selector = None + future = arg + if future.done(): + exception = CancelledError() if future.cancelled() else future.exception() + if exception is not None: + raise exception + for k, v in memory_dict.items(): + if k in file_name_dict and (future == v or future in duplicate_dict.get(k, [])): + future_wait_key_lst.append(k) + return FutureItem(file_name=file_name_dict[k], selector=selector) + return arg.result() + + def _refresh_memory_dict( memory_dict: dict, cache_dir_dict: dict, diff --git a/tests/unit/standalone/test_hdf.py b/tests/unit/standalone/test_hdf.py index ea9aa0df..ef74c463 100644 --- a/tests/unit/standalone/test_hdf.py +++ b/tests/unit/standalone/test_hdf.py @@ -184,5 +184,20 @@ def test_hdf_error(self): self.assertFalse(no_error) self.assertTrue(isinstance(output, error.__class__)) + def test_hdf_unpicklable(self): + class Unpicklable: + def __reduce__(self): + raise TypeError("cannot pickle Unpicklable") + + cache_directory = os.path.abspath("executorlib_cache") + os.makedirs(cache_directory, exist_ok=True) + file_name = os.path.join(cache_directory, "test_unpicklable.h5") + with self.assertRaises(TypeError): + dump( + file_name=file_name, + data_dict={"fn": my_funct, "args": [Unpicklable()], "kwargs": {}}, + ) + self.assertFalse(os.path.exists(file_name)) + def tearDown(self): shutil.rmtree("executorlib_cache", ignore_errors=True) diff --git a/tests/unit/standalone/test_select.py b/tests/unit/standalone/test_select.py index f6cc365b..9e20e69d 100644 --- a/tests/unit/standalone/test_select.py +++ b/tests/unit/standalone/test_select.py @@ -56,6 +56,20 @@ def test_integration_exception(self): with self.assertRaises(RuntimeError): f3.result() + def test_split_completed_future(self): + f1 = Future() + f1.set_result([1, 2]) + fs1, fs2 = split_future(future=f1, n=2) + self.assertTrue(f1.done()) + self.assertTrue(fs1.done()) + self.assertEqual(fs2.result(timeout=1), 2) + f2 = Future() + f2.set_exception(TypeError()) + fs3 = get_item_from_future(future=f2, key="a") + self.assertTrue(f2.done()) + self.assertTrue(fs3.done()) + self.assertEqual(type(fs3.exception(timeout=1)), TypeError) + def test_split_future_object(self): f1 = Future() fs1 = FutureSelector(future=f1, selector=1) diff --git a/tests/unit/task_scheduler/file/test_serial.py b/tests/unit/task_scheduler/file/test_serial.py index 76f7470e..86d6a044 100644 --- a/tests/unit/task_scheduler/file/test_serial.py +++ b/tests/unit/task_scheduler/file/test_serial.py @@ -1,4 +1,4 @@ -from concurrent.futures import Future +from concurrent.futures import CancelledError, Future import os from queue import Queue import shutil @@ -13,6 +13,7 @@ ) from executorlib.task_scheduler.file.task_scheduler import FileTaskScheduler, create_file_executor from executorlib.task_scheduler.file.shared import execute_tasks_h5, _convert_args_and_kwargs + from executorlib.standalone.select import FutureSelector skip_h5py_test = False except ImportError: @@ -36,10 +37,50 @@ def __reduce__(self): raise TypeError("cannot pickle Unpicklable") +UNPICKLABLE = Unpicklable() + + +def use_unpicklable_global(a): + return UNPICKLABLE is not None and [a, a] + + @unittest.skipIf( skip_h5py_test, "h5py is not installed, so the h5py tests are skipped." ) class TestCacheExecutorSerial(unittest.TestCase): + def test_submit_dependency_on_failed_task(self): + for resource_dict in [{"cache_key": "bad_dependency"}, {}]: + with self.subTest(resource_dict=resource_dict): + with FileTaskScheduler(execute_function=subprocess_execute) as exe: + fs_bad = exe.submit( + use_unpicklable_global, 1, resource_dict=resource_dict + ) + fs_args = exe.submit(my_funct, fs_bad, b=1) + fs_kwargs = exe.submit(my_funct, 1, b=fs_bad) + fs_selector = exe.submit( + my_funct, FutureSelector(future=fs_bad, selector=0), b=1 + ) + fs_downstream = exe.submit(my_funct, fs_args, b=1) + fs_good = exe.submit(my_funct, 1, b=2) + for fs in [fs_bad, fs_args, fs_kwargs, fs_selector, fs_downstream]: + with self.assertRaises(TypeError): + fs.result(timeout=15) + self.assertEqual(fs_good.result(timeout=15), 3) + self.assertFalse( + os.path.exists( + os.path.join("executorlib_cache", "bad_dependency_i.h5") + ) + ) + + def test_submit_dependency_on_duplicate_task(self): + with FileTaskScheduler(execute_function=subprocess_execute) as exe: + fs1 = exe.submit(my_funct, 1, b=2) + fs2 = exe.submit(my_funct, 1, b=2) + fs3 = exe.submit(my_funct, fs2, b=1) + self.assertEqual(fs3.result(timeout=15), 4) + self.assertEqual(fs1.result(timeout=15), 3) + self.assertEqual(fs2.result(timeout=15), 3) + def test_submit_with_positional_and_keyword_args(self): with FileTaskScheduler(execute_function=subprocess_execute) as exe: fs1 = exe.submit(my_funct, 1, b=2) @@ -356,5 +397,43 @@ def test_convert_args_and_kwargs(self): self.assertEqual(task_kwargs, {"a": 2}) self.assertTrue(len(future_wait_key_lst) == 0) + def test_convert_args_and_kwargs_selector(self): + f1 = Future() + f1.set_result([1, 2]) + task_args, task_kwargs, future_wait_key_lst = _convert_args_and_kwargs( + task_dict={ + "fn": 1, + "args": (FutureSelector(future=f1, selector=1),), + "kwargs": {"a": FutureSelector(future=f1, selector=0)}, + }, + memory_dict={}, + file_name_dict={}, + ) + self.assertEqual(task_args, [2]) + self.assertEqual(task_kwargs, {"a": 1}) + self.assertTrue(len(future_wait_key_lst) == 0) + + def test_convert_args_and_kwargs_failed_future(self): + f1 = Future() + f1.set_exception(TypeError("failed before submission")) + f2 = Future() + f2.cancel() + for task_dict, error in [ + ({"fn": 1, "args": (f1,), "kwargs": {}}, TypeError), + ({"fn": 1, "args": (), "kwargs": {"a": f1}}, TypeError), + ( + {"fn": 1, "args": (FutureSelector(future=f1, selector=0),), "kwargs": {}}, + TypeError, + ), + ({"fn": 1, "args": (f2,), "kwargs": {}}, CancelledError), + ]: + with self.subTest(task_dict=task_dict): + with self.assertRaises(error): + _convert_args_and_kwargs( + task_dict=task_dict, + memory_dict={"failed": f1}, + file_name_dict={}, + ) + def tearDown(self): shutil.rmtree("executorlib_cache", ignore_errors=True)