diff --git a/src/executorlib/standalone/interactive/communication.py b/src/executorlib/standalone/interactive/communication.py index 665d71f9..a5623dd2 100644 --- a/src/executorlib/standalone/interactive/communication.py +++ b/src/executorlib/standalone/interactive/communication.py @@ -78,19 +78,17 @@ def receive_dict(self) -> dict: while len(response_lst) == 0: response_lst = self._poller.poll(self._time_out_ms) if not self._spawner.poll(): - raise ExecutorlibSocketError( - "SocketInterface crashed during execution." - ) + return { + "error": ExecutorlibSocketError( + "SocketInterface crashed during execution." + ), + } data = self._socket.recv(zmq.NOBLOCK) if self._logger is not None: self._logger.warning( "Received dictionary of size: " + str(sys.getsizeof(data)) ) - output = cloudpickle.loads(data) - if "result" in output: - return output["result"] - else: - raise output["error"] + return cloudpickle.loads(data) def send_and_receive_dict(self, input_dict: dict) -> dict: """ @@ -154,7 +152,7 @@ def shutdown(self, wait: bool = True): if self._spawner.poll(): result = self.send_and_receive_dict( input_dict={"shutdown": True, "wait": wait} - ) + )["result"] self._spawner.shutdown(wait=wait) self._reset_socket() return result diff --git a/src/executorlib/task_scheduler/interactive/blockallocation.py b/src/executorlib/task_scheduler/interactive/blockallocation.py index 1f7474df..02501e78 100644 --- a/src/executorlib/task_scheduler/interactive/blockallocation.py +++ b/src/executorlib/task_scheduler/interactive/blockallocation.py @@ -379,10 +379,9 @@ def _set_init_function( ) -> Optional[Exception]: interface_initialization_exception = None if init_function is not None and interface.status: - try: - _ = interface.send_and_receive_dict( - input_dict={"init": True, "fn": init_function, "args": (), "kwargs": {}} - ) - except Exception as init_exception: - interface_initialization_exception = init_exception + output = interface.send_and_receive_dict( + input_dict={"init": True, "fn": init_function, "args": (), "kwargs": {}} + ) + if "error" in output: + interface_initialization_exception = output["error"] return interface_initialization_exception diff --git a/src/executorlib/task_scheduler/interactive/shared.py b/src/executorlib/task_scheduler/interactive/shared.py index e4084222..badca5ee 100644 --- a/src/executorlib/task_scheduler/interactive/shared.py +++ b/src/executorlib/task_scheduler/interactive/shared.py @@ -101,13 +101,13 @@ def _execute_task_without_cache( Returns: bool: True if the task was submitted successfully, False otherwise. """ - try: - future_obj.set_result(interface.send_and_receive_dict(input_dict=task_dict)) - except Exception as thread_exception: - if isinstance(thread_exception, ExecutorlibSocketError): - return False - else: - future_obj.set_exception(exception=thread_exception) + 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 + else: + future_obj.set_exception(exception=output["error"]) return True @@ -141,18 +141,17 @@ def _execute_task_with_cache( ) 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): - try: - time_start = time.time() - result = interface.send_and_receive_dict(input_dict=task_dict) - data_dict["output"] = result + 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(result) - except Exception as thread_exception: - if isinstance(thread_exception, ExecutorlibSocketError): - return False - else: - future_obj.set_exception(exception=thread_exception) + 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) diff --git a/tests/unit/standalone/interactive/test_communication.py b/tests/unit/standalone/interactive/test_communication.py index 4ed5f35a..6c0832d4 100644 --- a/tests/unit/standalone/interactive/test_communication.py +++ b/tests/unit/standalone/interactive/test_communication.py @@ -63,7 +63,7 @@ def test_interface_mpi(self): ) self.assertTrue(interface.status) self.assertEqual( - interface.send_and_receive_dict(input_dict=task_dict), np.array(4) + interface.send_and_receive_dict(input_dict=task_dict)["result"], np.array(4) ) interface.shutdown(wait=True) @@ -96,7 +96,7 @@ def test_interface_serial_without_debug(self): ) self.assertTrue(interface.status) self.assertEqual( - interface.send_and_receive_dict(input_dict=task_dict), np.array(4) + interface.send_and_receive_dict(input_dict=task_dict)["result"], np.array(4) ) interface.shutdown(wait=True) @@ -129,7 +129,7 @@ def test_interface_serial_with_debug(self): ) self.assertTrue(interface.status) self.assertEqual( - interface.send_and_receive_dict(input_dict=task_dict), np.array(4) + interface.send_and_receive_dict(input_dict=task_dict)["result"], np.array(4) ) interface.shutdown(wait=True) @@ -194,8 +194,8 @@ def test_interface_serial_with_stopped_process(self): self.assertTrue(interface.status) interface.send_dict(input_dict=task_dict) interface._spawner._process.terminate() - with self.assertRaises(ExecutorlibSocketError): - interface.receive_dict() + output = interface.receive_dict() + self.assertIsInstance(output["error"], ExecutorlibSocketError) class TestZMQ(unittest.TestCase):