[Refactor] receive_dict() returns the status and the error as dict (#1029)

* [Refactor] receive_dict() returns the status and the error as dict

* fix test

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Fix ruff

* return error rather than crashing

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>

Author: jan-janssen

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

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

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)

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):