[Feature] Log Communication overhead (#671)

* [Debug] Log Communication overhead

* Implement log_size_of_communicated_objects parameter

* Rename parameter to log_obj_size

* Format black

* fix test

---------

Co-authored-by: pyiron-runner <pyiron@mpie.de>

Author: jan-janssen

executorlib/executor/flux.py

@@ -61,6 +61,7 @@ class FluxJobExecutor(BaseExecutor):
         plot_dependency_graph (bool): Plot the dependencies of multiple future objects without executing them. For
                                       debugging purposes and to get an overview of the specified dependencies.
         plot_dependency_graph_filename (str): Name of the file to store the plotted graph in.
+        log_obj_size (bool): Enable debug mode which reports the size of the communicated objects.
 
     Examples:
         ```
@@ -100,6 +101,7 @@ def __init__(
         refresh_rate: float = 0.01,
         plot_dependency_graph: bool = False,
         plot_dependency_graph_filename: Optional[str] = None,
+        log_obj_size: bool = False,
     ):
         """
         The executorlib.FluxJobExecutor leverages either the message passing interface (MPI), the SLURM workload manager
@@ -144,6 +146,7 @@ def __init__(
             plot_dependency_graph (bool): Plot the dependencies of multiple future objects without executing them. For
                                           debugging purposes and to get an overview of the specified dependencies.
             plot_dependency_graph_filename (str): Name of the file to store the plotted graph in.
+            log_obj_size (bool): Enable debug mode which reports the size of the communicated objects.
 
         """
         default_resource_dict: dict = {
@@ -174,6 +177,7 @@ def __init__(
                         hostname_localhost=hostname_localhost,
                         block_allocation=block_allocation,
                         init_function=init_function,
+                        log_obj_size=log_obj_size,
                     ),
                     max_cores=max_cores,
                     refresh_rate=refresh_rate,
@@ -197,6 +201,7 @@ def __init__(
                     hostname_localhost=hostname_localhost,
                     block_allocation=block_allocation,
                     init_function=init_function,
+                    log_obj_size=log_obj_size,
                 )
             )
 
@@ -392,6 +397,7 @@ def create_flux_executor(
     hostname_localhost: Optional[bool] = None,
     block_allocation: bool = False,
     init_function: Optional[Callable] = None,
+    log_obj_size: bool = False,
 ) -> Union[OneProcessTaskScheduler, BlockAllocationTaskScheduler]:
     """
     Create a flux executor
@@ -427,6 +433,7 @@ def create_flux_executor(
                                     resources have to be defined on the executor, rather than during the submission
                                     of the individual function.
         init_function (None): optional function to preset arguments for functions which are submitted later
+        log_obj_size (bool): Enable debug mode which reports the size of the communicated objects.
 
     Returns:
         InteractiveStepExecutor/ InteractiveExecutor
@@ -441,6 +448,7 @@ def create_flux_executor(
     cores_per_worker = resource_dict.get("cores", 1)
     resource_dict["cache_directory"] = cache_directory
     resource_dict["hostname_localhost"] = hostname_localhost
+    resource_dict["log_obj_size"] = log_obj_size
     check_init_function(block_allocation=block_allocation, init_function=init_function)
     check_pmi(backend="flux_allocation", pmi=flux_executor_pmi_mode)
     check_oversubscribe(oversubscribe=resource_dict.get("openmpi_oversubscribe", False))

executorlib/executor/single.py

@@ -91,6 +91,7 @@ def __init__(
         refresh_rate: float = 0.01,
         plot_dependency_graph: bool = False,
         plot_dependency_graph_filename: Optional[str] = None,
+        log_obj_size: bool = False,
     ):
         """
         The executorlib.SingleNodeExecutor leverages either the message passing interface (MPI), the SLURM workload
@@ -131,6 +132,7 @@ def __init__(
             plot_dependency_graph (bool): Plot the dependencies of multiple future objects without executing them. For
                                           debugging purposes and to get an overview of the specified dependencies.
             plot_dependency_graph_filename (str): Name of the file to store the plotted graph in.
+            log_obj_size (bool): Enable debug mode which reports the size of the communicated objects.
 
         """
         default_resource_dict: dict = {
@@ -157,6 +159,7 @@ def __init__(
                         hostname_localhost=hostname_localhost,
                         block_allocation=block_allocation,
                         init_function=init_function,
+                        log_obj_size=log_obj_size,
                     ),
                     max_cores=max_cores,
                     refresh_rate=refresh_rate,
@@ -176,6 +179,7 @@ def __init__(
                     hostname_localhost=hostname_localhost,
                     block_allocation=block_allocation,
                     init_function=init_function,
+                    log_obj_size=log_obj_size,
                 )
             )
 
@@ -188,6 +192,7 @@ def create_single_node_executor(
     hostname_localhost: Optional[bool] = None,
     block_allocation: bool = False,
     init_function: Optional[Callable] = None,
+    log_obj_size: bool = False,
 ) -> Union[OneProcessTaskScheduler, BlockAllocationTaskScheduler]:
     """
     Create a single node executor
@@ -219,6 +224,7 @@ def create_single_node_executor(
                                     resources have to be defined on the executor, rather than during the submission
                                     of the individual function.
         init_function (None): optional function to preset arguments for functions which are submitted later
+        log_obj_size (bool): Enable debug mode which reports the size of the communicated objects.
 
     Returns:
         InteractiveStepExecutor/ InteractiveExecutor
@@ -228,6 +234,7 @@ def create_single_node_executor(
     cores_per_worker = resource_dict.get("cores", 1)
     resource_dict["cache_directory"] = cache_directory
     resource_dict["hostname_localhost"] = hostname_localhost
+    resource_dict["log_obj_size"] = log_obj_size
 
     check_init_function(block_allocation=block_allocation, init_function=init_function)
     check_gpus_per_worker(gpus_per_worker=resource_dict.get("gpus_per_core", 0))

executorlib/executor/slurm.py

@@ -236,6 +236,7 @@ class SlurmJobExecutor(BaseExecutor):
         plot_dependency_graph (bool): Plot the dependencies of multiple future objects without executing them. For
                                       debugging purposes and to get an overview of the specified dependencies.
         plot_dependency_graph_filename (str): Name of the file to store the plotted graph in.
+        log_obj_size (bool): Enable debug mode which reports the size of the communicated objects.
 
     Examples:
         ```
@@ -271,6 +272,7 @@ def __init__(
         refresh_rate: float = 0.01,
         plot_dependency_graph: bool = False,
         plot_dependency_graph_filename: Optional[str] = None,
+        log_obj_size: bool = False,
     ):
         """
         The executorlib.SlurmJobExecutor leverages either the message passing interface (MPI), the SLURM workload
@@ -315,6 +317,7 @@ def __init__(
             plot_dependency_graph (bool): Plot the dependencies of multiple future objects without executing them. For
                                           debugging purposes and to get an overview of the specified dependencies.
             plot_dependency_graph_filename (str): Name of the file to store the plotted graph in.
+            log_obj_size (bool): Enable debug mode which reports the size of the communicated objects.
 
         """
         default_resource_dict: dict = {
@@ -341,6 +344,7 @@ def __init__(
                         hostname_localhost=hostname_localhost,
                         block_allocation=block_allocation,
                         init_function=init_function,
+                        log_obj_size=log_obj_size,
                     ),
                     max_cores=max_cores,
                     refresh_rate=refresh_rate,
@@ -360,6 +364,7 @@ def __init__(
                     hostname_localhost=hostname_localhost,
                     block_allocation=block_allocation,
                     init_function=init_function,
+                    log_obj_size=log_obj_size,
                 )
             )
 
@@ -372,6 +377,7 @@ def create_slurm_executor(
     hostname_localhost: Optional[bool] = None,
     block_allocation: bool = False,
     init_function: Optional[Callable] = None,
+    log_obj_size: bool = False,
 ) -> Union[OneProcessTaskScheduler, BlockAllocationTaskScheduler]:
     """
     Create a SLURM executor
@@ -407,6 +413,7 @@ def create_slurm_executor(
                                     resources have to be defined on the executor, rather than during the submission
                                     of the individual function.
         init_function (None): optional function to preset arguments for functions which are submitted later
+        log_obj_size (bool): Enable debug mode which reports the size of the communicated objects.
 
     Returns:
         InteractiveStepExecutor/ InteractiveExecutor
@@ -416,6 +423,7 @@ def create_slurm_executor(
     cores_per_worker = resource_dict.get("cores", 1)
     resource_dict["cache_directory"] = cache_directory
     resource_dict["hostname_localhost"] = hostname_localhost
+    resource_dict["log_obj_size"] = log_obj_size
     check_init_function(block_allocation=block_allocation, init_function=init_function)
     if block_allocation:
         resource_dict["init_function"] = init_function

executorlib/standalone/interactive/communication.py

@@ -1,3 +1,4 @@
+import logging
 import sys
 from socket import gethostname
 from typing import Optional
@@ -12,9 +13,10 @@ class SocketInterface:
 
     Args:
         spawner (executorlib.shared.spawner.BaseSpawner): Interface for starting the parallel process
+        log_obj_size (boolean): Enable debug mode which reports the size of the communicated objects.
     """
 
-    def __init__(self, spawner=None):
+    def __init__(self, spawner=None, log_obj_size=False):
         """
         Initialize the SocketInterface.
 
@@ -24,6 +26,10 @@ def __init__(self, spawner=None):
         self._context = zmq.Context()
         self._socket = self._context.socket(zmq.PAIR)
         self._process = None
+        if log_obj_size:
+            self._logger = logging.getLogger("executorlib")
+        else:
+            self._logger = None
         self._spawner = spawner
 
     def send_dict(self, input_dict: dict):
@@ -34,7 +40,10 @@ def send_dict(self, input_dict: dict):
             input_dict (dict): dictionary of commands to be communicated. The key "shutdown" is reserved to stop the
                 connected client from listening.
         """
-        self._socket.send(cloudpickle.dumps(input_dict))
+        data = cloudpickle.dumps(input_dict)
+        if self._logger is not None:
+            self._logger.warning("Send dictionary of size: " + str(sys.getsizeof(data)))
+        self._socket.send(data)
 
     def receive_dict(self) -> dict:
         """
@@ -43,7 +52,12 @@ def receive_dict(self) -> dict:
         Returns:
             dict: dictionary with response received from the connected client
         """
-        output = cloudpickle.loads(self._socket.recv())
+        data = self._socket.recv()
+        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:
@@ -121,6 +135,7 @@ def interface_bootup(
     command_lst: list[str],
     connections,
     hostname_localhost: Optional[bool] = None,
+    log_obj_size: bool = False,
 ) -> SocketInterface:
     """
     Start interface for ZMQ communication
@@ -136,6 +151,7 @@ def interface_bootup(
                                       points to the same address as localhost. Still MacOS >= 12 seems to disable
                                       this look up for security reasons. So on MacOS it is required to set this
                                       option to true
+        log_obj_size (boolean): Enable debug mode which reports the size of the communicated objects.
 
     Returns:
          executorlib.shared.communication.SocketInterface: socket interface for zmq communication
@@ -149,7 +165,10 @@ def interface_bootup(
             "--host",
             gethostname(),
         ]
-    interface = SocketInterface(spawner=connections)
+    interface = SocketInterface(
+        spawner=connections,
+        log_obj_size=log_obj_size,
+    )
     command_lst += [
         "--zmqport",
         str(interface.bind_to_random_port()),

executorlib/task_scheduler/interactive/shared.py

@@ -23,6 +23,7 @@ def execute_tasks(
     init_function: Optional[Callable] = None,
     cache_directory: Optional[str] = None,
     queue_join_on_shutdown: bool = True,
+    log_obj_size: bool = False,
     **kwargs,
 ) -> None:
     """
@@ -42,13 +43,15 @@ def execute_tasks(
        init_function (Callable): optional function to preset arguments for functions which are submitted later
        cache_directory (str, optional): The directory to store cache files. Defaults to "cache".
        queue_join_on_shutdown (bool): Join communication queue when thread is closed. Defaults to True.
+       log_obj_size (bool): Enable debug mode which reports the size of the communicated objects.
     """
     interface = interface_bootup(
         command_lst=_get_backend_path(
             cores=cores,
         ),
         connections=spawner(cores=cores, **kwargs),
         hostname_localhost=hostname_localhost,
+        log_obj_size=log_obj_size,
     )
     if init_function is not None:
         interface.send_dict(

tests/test_singlenodeexecutor_dependencies.py

@@ -341,6 +341,7 @@ def setUp(self):
             'openmpi_oversubscribe': False,
             'cache_directory': None,
             'hostname_localhost': None,
+            'log_obj_size': False,
             'spawner': MpiExecSpawner,
             'max_cores': None,
             'max_workers': None,

tests/test_standalone_interactive_communication.py

@@ -56,11 +56,41 @@ def test_interface_mpi(self):
         )
         interface.shutdown(wait=True)
 
-    def test_interface_serial(self):
+    def test_interface_serial_without_debug(self):
         cloudpickle_register(ind=1)
         task_dict = {"fn": calc, "args": (), "kwargs": {"i": 2}}
         interface = SocketInterface(
-            spawner=MpiExecSpawner(cwd=None, cores=1, openmpi_oversubscribe=False)
+            spawner=MpiExecSpawner(cwd=None, cores=1, openmpi_oversubscribe=False),
+            log_obj_size=False,
+        )
+        interface.bootup(
+            command_lst=[
+                sys.executable,
+                os.path.abspath(
+                    os.path.join(
+                        __file__,
+                        "..",
+                        "..",
+                        "executorlib",
+                        "backend",
+                        "interactive_serial.py",
+                    )
+                ),
+                "--zmqport",
+                str(interface.bind_to_random_port()),
+            ]
+        )
+        self.assertEqual(
+            interface.send_and_receive_dict(input_dict=task_dict), np.array(4)
+        )
+        interface.shutdown(wait=True)
+
+    def test_interface_serial_with_debug(self):
+        cloudpickle_register(ind=1)
+        task_dict = {"fn": calc, "args": (), "kwargs": {"i": 2}}
+        interface = SocketInterface(
+            spawner=MpiExecSpawner(cwd=None, cores=1, openmpi_oversubscribe=False),
+            log_obj_size=True,
         )
         interface.bootup(
             command_lst=[