diff --git a/pyproject.toml b/pyproject.toml index f130c9b4e..0c3e44c58 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -102,6 +102,8 @@ ignore = [ "PLR0912", # Too many arguments in function definition "PLR0913", + # Too many statements + "PLR0915", # Too many not-keyword-only arguments, "PLR0917", # Magic value used in comparison diff --git a/src/executorlib/standalone/serialize.py b/src/executorlib/standalone/serialize.py index 46c7b85a4..5d5c86dff 100644 --- a/src/executorlib/standalone/serialize.py +++ b/src/executorlib/standalone/serialize.py @@ -34,7 +34,7 @@ def serialize_funct( fn_kwargs: Optional[dict] = None, resource_dict: Optional[dict] = None, cache_key: Optional[str] = None, -) -> tuple[str, dict]: +) -> tuple[str, dict, Optional[Exception]]: """ Serialize a function and its arguments and keyword arguments into an HDF5 file. @@ -56,7 +56,8 @@ def serialize_funct( overwritten by setting the cache_key. Returns: - Tuple[str, dict]: A tuple containing the task key and the serialized data. + Tuple[str, dict, Optional[Exception]]: A tuple containing the task key, the serialized data and an + exception if the serialization failed. """ if fn_args is None: @@ -68,21 +69,25 @@ def serialize_funct( if cache_key is not None: task_key = cache_key else: - binary_all = cloudpickle.dumps( - { - "fn": fn, - "args": fn_args, - "kwargs": fn_kwargs, - } - ) - task_key = _get_function_name(fn=fn) + _get_hash(binary=binary_all) + try: + binary_all = cloudpickle.dumps( + { + "fn": fn, + "args": fn_args, + "kwargs": fn_kwargs, + } + ) + except Exception as e: + return "", {}, e + else: + task_key = _get_function_name(fn=fn) + _get_hash(binary=binary_all) data = { "fn": fn, "args": fn_args, "kwargs": fn_kwargs, "resource_dict": resource_dict, } - return task_key, data + return task_key, data, None def _get_hash(binary: bytes) -> str: diff --git a/src/executorlib/task_scheduler/file/shared.py b/src/executorlib/task_scheduler/file/shared.py index d6deb74b3..c65865faf 100644 --- a/src/executorlib/task_scheduler/file/shared.py +++ b/src/executorlib/task_scheduler/file/shared.py @@ -129,7 +129,7 @@ def execute_tasks_h5( executor_kwargs=executor_kwargs, ) ) - task_key, data_dict = serialize_funct( + task_key, data_dict, serialize_exception = serialize_funct( fn=task_dict["fn"], fn_args=task_args, fn_kwargs=task_kwargs, @@ -137,7 +137,9 @@ def execute_tasks_h5( cache_key=cache_key, ) data_dict["error_log_file"] = error_log_file - if task_key not in memory_dict: + 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): @@ -156,25 +158,31 @@ def execute_tasks_h5( ) ) task_dependent_lst = [] - process_dict[task_key] = execute_function( - command=get_cache_execute_command( + 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, - ) + 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) diff --git a/src/executorlib/task_scheduler/interactive/shared.py b/src/executorlib/task_scheduler/interactive/shared.py index badca5eed..abb633ae3 100644 --- a/src/executorlib/task_scheduler/interactive/shared.py +++ b/src/executorlib/task_scheduler/interactive/shared.py @@ -101,13 +101,17 @@ def _execute_task_without_cache( Returns: bool: True if the task was submitted successfully, False otherwise. """ - output = interface.send_and_receive_dict(input_dict=task_dict) - if "result" in output: - future_obj.set_result(output["result"]) - elif isinstance(output["error"], ExecutorlibSocketError): - return False + try: + output = interface.send_and_receive_dict(input_dict=task_dict) + except Exception as e: + future_obj.set_exception(exception=e) else: - future_obj.set_exception(exception=output["error"]) + if "result" in output: + future_obj.set_result(output["result"]) + elif isinstance(output["error"], ExecutorlibSocketError): + return False + else: + future_obj.set_exception(exception=output["error"]) return True @@ -132,27 +136,34 @@ def _execute_task_with_cache( """ from executorlib.standalone.hdf import dump, get_cache_files, get_output - task_key, data_dict = serialize_funct( + task_key, data_dict, serialize_exception = serialize_funct( fn=task_dict["fn"], fn_args=task_dict["args"], fn_kwargs=task_dict["kwargs"], resource_dict=task_dict.get("resource_dict", {}), cache_key=cache_key, ) - 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() - output = interface.send_and_receive_dict(input_dict=task_dict) - if "result" in output: - data_dict["output"] = output["result"] - data_dict["runtime"] = time.time() - time_start - dump(file_name=file_name, data_dict=data_dict) - future_obj.set_result(output["result"]) - elif isinstance(output["error"], ExecutorlibSocketError): - return False - else: - future_obj.set_exception(exception=output["error"]) + if serialize_exception is not None: + future_obj.set_exception(exception=serialize_exception) else: - _, _, result = get_output(file_name=file_name) - future_obj.set_result(result) + 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() + try: + output = interface.send_and_receive_dict(input_dict=task_dict) + except Exception as e: + future_obj.set_exception(exception=e) + return True + if "result" in output: + data_dict["output"] = output["result"] + data_dict["runtime"] = time.time() - time_start + dump(file_name=file_name, data_dict=data_dict) + future_obj.set_result(output["result"]) + elif isinstance(output["error"], ExecutorlibSocketError): + return False + else: + future_obj.set_exception(exception=output["error"]) + else: + _, _, result = get_output(file_name=file_name) + future_obj.set_result(result) return True diff --git a/tests/unit/executor/test_single_dependencies.py b/tests/unit/executor/test_single_dependencies.py index 83de011ae..2370eed80 100644 --- a/tests/unit/executor/test_single_dependencies.py +++ b/tests/unit/executor/test_single_dependencies.py @@ -47,6 +47,11 @@ def raise_error(parameter): raise RuntimeError +class Unpicklable: + def __reduce__(self): + raise TypeError("cannot pickle Unpicklable") + + class TestExecutorWithDependencies(unittest.TestCase): def test_future_chaining_resolves_dependency(self): with SingleNodeExecutor(max_cores=1) as exe: @@ -337,6 +342,29 @@ def test_block_allocation_true_two_workers(self): fs = exe.submit(raise_error, parameter=0) fs.result() + def test_unpicklable_argument_block_allocation_false(self): + with self.assertRaises(TypeError): + with SingleNodeExecutor(max_cores=1, block_allocation=False) as exe: + cloudpickle_register(ind=1) + fs = exe.submit(len, [Unpicklable()]) + fs.result(timeout=15) + + def test_unpicklable_argument_block_allocation_true(self): + with self.assertRaises(TypeError): + with SingleNodeExecutor(max_cores=1, block_allocation=True) as exe: + cloudpickle_register(ind=1) + fs = exe.submit(len, [Unpicklable()]) + fs.result(timeout=15) + + def test_unpicklable_argument_recovers_for_next_task(self): + with SingleNodeExecutor(max_cores=1, block_allocation=False) as exe: + cloudpickle_register(ind=1) + bad_fs = exe.submit(len, [Unpicklable()]) + good_fs = exe.submit(sum, [1, 2]) + with self.assertRaises(TypeError): + bad_fs.result(timeout=15) + self.assertEqual(good_fs.result(timeout=15), 3) + def test_block_allocation_false_one_worker_loop(self): with self.assertRaises(RuntimeError): with SingleNodeExecutor(max_cores=1, block_allocation=False) as exe: diff --git a/tests/unit/task_scheduler/file/test_backend.py b/tests/unit/task_scheduler/file/test_backend.py index 0ae08f189..35a62f7b4 100644 --- a/tests/unit/task_scheduler/file/test_backend.py +++ b/tests/unit/task_scheduler/file/test_backend.py @@ -55,7 +55,7 @@ class TestSharedFunctions(unittest.TestCase): def test_execute_function_mixed(self): cache_directory = os.path.abspath("executorlib_cache") os.makedirs(cache_directory, exist_ok=True) - task_key, data_dict = serialize_funct( + task_key, data_dict, _ = serialize_funct( fn=my_funct, fn_args=[1], fn_kwargs={"b": 2}, @@ -106,7 +106,7 @@ def test_backend_write_file_serialization_error(self): def test_execute_function_mixed_selector_convert(self): cache_directory = os.path.abspath("executorlib_cache") os.makedirs(cache_directory, exist_ok=True) - task_key_1, data_dict = serialize_funct( + task_key_1, data_dict, _ = serialize_funct( fn=return_dict, fn_args=[1], fn_kwargs={"b": 2}, @@ -118,7 +118,7 @@ def test_execute_function_mixed_selector_convert(self): _check_task_output( task_key=task_key_1, future_obj=f1, cache_directory=cache_directory, ) - task_key_2, data_dict = serialize_funct( + task_key_2, data_dict, _ = serialize_funct( fn=return_list, fn_args=[1], fn_kwargs={"b": 2}, @@ -147,7 +147,7 @@ def test_execute_function_mixed_selector_convert(self): def test_execute_function_args(self): cache_directory = os.path.abspath("executorlib_cache") os.makedirs(cache_directory, exist_ok=True) - task_key, data_dict = serialize_funct( + task_key, data_dict, _ = serialize_funct( fn=my_funct, fn_args=[1, 2], fn_kwargs=None, @@ -175,7 +175,7 @@ def test_execute_function_args(self): def test_execute_function_kwargs(self): cache_directory = os.path.abspath("executorlib_cache") os.makedirs(cache_directory, exist_ok=True) - task_key, data_dict = serialize_funct( + task_key, data_dict, _ = serialize_funct( fn=my_funct, fn_args=None, fn_kwargs={"a": 1, "b": 2}, @@ -202,7 +202,7 @@ def test_execute_function_kwargs(self): def test_execute_function_error(self): cache_directory = os.path.abspath("executorlib_cache") os.makedirs(cache_directory, exist_ok=True) - task_key, data_dict = serialize_funct( + task_key, data_dict, _ = serialize_funct( fn=get_error, fn_args=[], fn_kwargs={"a": 1}, @@ -241,7 +241,7 @@ def test_check_task_output_dead_job_without_output(self): # or an external scancel) must fail the future instead of leaving it pending forever. cache_directory = os.path.abspath("executorlib_cache") os.makedirs(cache_directory, exist_ok=True) - task_key, data_dict = serialize_funct(fn=my_funct, fn_args=[1], fn_kwargs={"b": 2}) + task_key, data_dict, _ = serialize_funct(fn=my_funct, fn_args=[1], fn_kwargs={"b": 2}) future_obj = Future() with patch( "executorlib.standalone.command_pysqa.pysqa_get_status_of_job", @@ -270,7 +270,7 @@ def test_check_task_output_dead_job_reported_as_error_status(self): # https://github.com/pyiron/executorlib/issues/1037. cache_directory = os.path.abspath("executorlib_cache") os.makedirs(cache_directory, exist_ok=True) - task_key, data_dict = serialize_funct(fn=my_funct, fn_args=[1], fn_kwargs={"b": 2}) + task_key, data_dict, _ = serialize_funct(fn=my_funct, fn_args=[1], fn_kwargs={"b": 2}) future_obj = Future() with patch( "executorlib.standalone.command_pysqa.pysqa_get_status_of_job", @@ -293,7 +293,7 @@ def test_check_task_output_dead_job_reported_as_error_status(self): def test_check_task_output_job_still_running(self): cache_directory = os.path.abspath("executorlib_cache") os.makedirs(cache_directory, exist_ok=True) - task_key, data_dict = serialize_funct(fn=my_funct, fn_args=[1], fn_kwargs={"b": 2}) + task_key, data_dict, _ = serialize_funct(fn=my_funct, fn_args=[1], fn_kwargs={"b": 2}) future_obj = Future() with patch( "executorlib.standalone.command_pysqa.pysqa_get_status_of_job", @@ -314,7 +314,7 @@ def test_check_task_output_job_still_running(self): def test_check_task_output_status_check_is_throttled(self): cache_directory = os.path.abspath("executorlib_cache") os.makedirs(cache_directory, exist_ok=True) - task_key, data_dict = serialize_funct(fn=my_funct, fn_args=[1], fn_kwargs={"b": 2}) + task_key, data_dict, _ = serialize_funct(fn=my_funct, fn_args=[1], fn_kwargs={"b": 2}) status_check_dict = {} with patch( "executorlib.standalone.command_pysqa.pysqa_get_status_of_job", @@ -337,7 +337,7 @@ def test_check_task_output_no_backend_never_queries_status(self): # subprocess-backed tasks (backend=None) must never trigger a queuing system status check. cache_directory = os.path.abspath("executorlib_cache") os.makedirs(cache_directory, exist_ok=True) - task_key, data_dict = serialize_funct(fn=my_funct, fn_args=[1], fn_kwargs={"b": 2}) + task_key, data_dict, _ = serialize_funct(fn=my_funct, fn_args=[1], fn_kwargs={"b": 2}) future_obj = Future() with patch( "executorlib.standalone.command_pysqa.pysqa_get_status_of_job", diff --git a/tests/unit/task_scheduler/file/test_serial.py b/tests/unit/task_scheduler/file/test_serial.py index bd89b6b66..76f7470e8 100644 --- a/tests/unit/task_scheduler/file/test_serial.py +++ b/tests/unit/task_scheduler/file/test_serial.py @@ -31,6 +31,11 @@ def get_error(a): raise ValueError(a) +class Unpicklable: + def __reduce__(self): + raise TypeError("cannot pickle Unpicklable") + + @unittest.skipIf( skip_h5py_test, "h5py is not installed, so the h5py tests are skipped." ) @@ -49,6 +54,18 @@ def test_submit_with_custom_cache_key(self): self.assertEqual(fs1.result(), 3) self.assertTrue(fs1.done()) + def test_submit_unpicklable_argument_with_custom_cache_key(self): + with FileTaskScheduler(execute_function=subprocess_execute) as exe: + fs_bad = exe.submit( + len, + [Unpicklable()], + resource_dict={"cache_key": "bad/key"}, + ) + fs_good = exe.submit(my_funct, 1, b=2) + with self.assertRaises(TypeError): + fs_bad.result(timeout=15) + self.assertEqual(fs_good.result(timeout=15), 3) + def test_submit_dependency_with_keyword_arg_future(self): with FileTaskScheduler(execute_function=subprocess_execute) as exe: fs1 = exe.submit(my_funct, 1, b=2) @@ -245,6 +262,86 @@ def test_execute_in_subprocess_errors(self): backend="flux", ) + def test_executor_function_unpicklable_argument(self): + fs_bad = Future() + fs_good = Future() + q = Queue() + q.put( + { + "fn": len, + "args": ([Unpicklable()],), + "kwargs": {}, + "future": fs_bad, + "resource_dict": {}, + } + ) + q.put( + { + "fn": my_funct, + "args": (), + "kwargs": {"a": 1, "b": 2}, + "future": fs_good, + "resource_dict": {}, + } + ) + cache_dir = os.path.abspath("executorlib_cache") + os.makedirs(cache_dir, exist_ok=True) + process = Thread( + target=execute_tasks_h5, + kwargs={ + "future_queue": q, + "execute_function": subprocess_execute, + "executor_kwargs": {"cores": 1, "cwd": None, "cache_directory": cache_dir}, + "terminate_function": subprocess_terminate, + }, + ) + process.start() + with self.assertRaises(TypeError): + fs_bad.result(timeout=15) + self.assertEqual(fs_good.result(timeout=15), 3) + q.put({"shutdown": True, "wait": True}) + process.join() + + def test_executor_function_unpicklable_argument_with_cache_key(self): + fs_bad = Future() + fs_good = Future() + q = Queue() + q.put( + { + "fn": len, + "args": ([Unpicklable()],), + "kwargs": {}, + "future": fs_bad, + "resource_dict": {"cache_key": "bad/key"}, + } + ) + q.put( + { + "fn": my_funct, + "args": (), + "kwargs": {"a": 1, "b": 2}, + "future": fs_good, + "resource_dict": {}, + } + ) + cache_dir = os.path.abspath("executorlib_cache") + os.makedirs(cache_dir, exist_ok=True) + process = Thread( + target=execute_tasks_h5, + kwargs={ + "future_queue": q, + "execute_function": subprocess_execute, + "executor_kwargs": {"cores": 1, "cwd": None, "cache_directory": cache_dir}, + "terminate_function": subprocess_terminate, + }, + ) + process.start() + with self.assertRaises(TypeError): + fs_bad.result(timeout=15) + self.assertEqual(fs_good.result(timeout=15), 3) + q.put({"shutdown": True, "wait": True}) + process.join() + def test_convert_args_and_kwargs(self): f1 = Future() f1.set_result(1) diff --git a/tests/unit/task_scheduler/interactive/test_shared.py b/tests/unit/task_scheduler/interactive/test_shared.py index bd691d3bf..51c2d1471 100644 --- a/tests/unit/task_scheduler/interactive/test_shared.py +++ b/tests/unit/task_scheduler/interactive/test_shared.py @@ -20,6 +20,11 @@ def get_error(): raise ExecutorlibSocketError() +class Unpicklable: + def __reduce__(self): + raise TypeError("cannot pickle Unpicklable") + + class TestExecuteTaskDictWithoutCache(unittest.TestCase): def test_execute_task_sum(self): cloudpickle_register(ind=1) @@ -102,6 +107,34 @@ def test_execute_task_error(self): self.assertFalse(result) self.assertFalse(f.done()) + def test_execute_task_unpicklable_argument(self): + cloudpickle_register(ind=1) + f = Future() + interface = interface_bootup( + command_lst=get_interactive_execute_command( + cores=1, + ), + connections=SubprocessSpawner(), + hostname_localhost=True, + log_obj_size=False, + worker_id=1, + stop_function=None, + ) + self.assertTrue(interface.status) + self.assertFalse(f.done()) + result = execute_task_dict( + task_dict={"fn": len, "args": ([Unpicklable()],), "kwargs": {}}, + future_obj=f, + interface=interface, + cache_directory=None, + cache_key=None, + error_log_file=None, + ) + self.assertTrue(result) + self.assertTrue(f.done()) + with self.assertRaises(TypeError): + f.result() + @unittest.skipIf( skip_h5py_test, "h5py is not installed, so the h5io tests are skipped." @@ -189,4 +222,45 @@ def test_execute_task_error(self): error_log_file=None, ) self.assertFalse(result) - self.assertFalse(f.done()) \ No newline at end of file + self.assertFalse(f.done()) + + def test_execute_task_unpicklable_argument_with_cache_key(self): + cloudpickle_register(ind=1) + bad_future = Future() + good_future = Future() + interface = interface_bootup( + command_lst=get_interactive_execute_command( + cores=1, + ), + connections=SubprocessSpawner(), + hostname_localhost=True, + log_obj_size=False, + worker_id=1, + stop_function=None, + ) + self.assertTrue(interface.status) + + bad_result = execute_task_dict( + task_dict={"fn": len, "args": ([Unpicklable()],), "kwargs": {}}, + future_obj=bad_future, + interface=interface, + cache_directory="cache_execute_task", + cache_key="cache/key", + error_log_file=None, + ) + self.assertTrue(bad_result) + self.assertTrue(bad_future.done()) + with self.assertRaises(TypeError): + bad_future.result() + + good_result = execute_task_dict( + task_dict={"fn": sum, "args": ([1, 2],), "kwargs": {}}, + future_obj=good_future, + interface=interface, + cache_directory="cache_execute_task", + cache_key="cache/key-next", + error_log_file=None, + ) + self.assertTrue(good_result) + self.assertTrue(good_future.done()) + self.assertEqual(good_future.result(), 3)