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

Author: jan-janssen

src/executorlib/standalone/interactive/communication.py

@@ -86,11 +86,7 @@ def receive_dict(self) -> dict:
             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 +150,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,14 @@ 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):
+    output = interface.send_and_receive_dict(input_dict=task_dict)
+    if "result" in output:
+        future_obj.set_result(output["result"])
+    else:
+        if isinstance(output["error"], ExecutorlibSocketError):
             return False
         else:
-            future_obj.set_exception(exception=thread_exception)
+            future_obj.set_exception(exception=output["error"])
     return True
 
 
@@ -141,18 +142,18 @@ 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):
+        else:
+            if isinstance(output["error"], ExecutorlibSocketError):
                 return False
             else:
-                future_obj.set_exception(exception=thread_exception)
+                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)
 
@@ -195,7 +195,7 @@ def test_interface_serial_with_stopped_process(self):
         interface.send_dict(input_dict=task_dict)
         interface._spawner._process.terminate()
         with self.assertRaises(ExecutorlibSocketError):
-            interface.receive_dict()
+            raise interface.receive_dict()["error"]
 
 
 class TestZMQ(unittest.TestCase):