executorlib 0.0.1__py3-none-any.whl

executorlib/shared/communication.py

@@ -0,0 +1,195 @@
+from socket import gethostname
+from typing import Optional
+
+import cloudpickle
+import zmq
+
+
+class SocketInterface(object):
+    """
+    The SocketInterface is an abstraction layer on top of the zero message queue.
+
+    Args:
+        interface (executorlib.shared.interface.BaseInterface): Interface for starting the parallel process
+    """
+
+    def __init__(self, interface=None):
+        self._context = zmq.Context()
+        self._socket = self._context.socket(zmq.PAIR)
+        self._process = None
+        self._interface = interface
+
+    def send_dict(self, input_dict: dict):
+        """
+        Send a dictionary with instructions to a connected client process.
+
+        Args:
+            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))
+
+    def receive_dict(self):
+        """
+        Receive a dictionary from a connected client process.
+
+        Returns:
+            dict: dictionary with response received from the connected client
+        """
+        output = cloudpickle.loads(self._socket.recv())
+        if "result" in output.keys():
+            return output["result"]
+        else:
+            error_type = output["error_type"].split("'")[1]
+            raise eval(error_type)(output["error"])
+
+    def send_and_receive_dict(self, input_dict: dict) -> dict:
+        """
+        Combine both the send_dict() and receive_dict() function in a single call.
+
+        Args:
+            input_dict (dict): dictionary of commands to be communicated. The key "shutdown" is reserved to stop the
+                connected client from listening.
+
+        Returns:
+            dict: dictionary with response received from the connected client
+        """
+        self.send_dict(input_dict=input_dict)
+        return self.receive_dict()
+
+    def bind_to_random_port(self):
+        """
+        Identify a random port typically in the range from 49152 to 65536 to bind the SocketInterface instance to. Other
+        processes can then connect to this port to receive instructions and send results.
+
+        Returns:
+            int: port the SocketInterface instance is bound to.
+        """
+        return self._socket.bind_to_random_port("tcp://*")
+
+    def bootup(
+        self,
+        command_lst: list[str],
+        prefix_name: Optional[str] = None,
+        prefix_path: Optional[str] = None,
+    ):
+        """
+        Boot up the client process to connect to the SocketInterface.
+
+        Args:
+            command_lst (list): list of strings to start the client process
+            prefix_name (str): name of the conda environment to initialize
+            prefix_path (str): path of the conda environment to initialize
+        """
+        self._interface.bootup(
+            command_lst=command_lst, prefix_name=prefix_name, prefix_path=prefix_path
+        )
+
+    def shutdown(self, wait: bool = True):
+        result = None
+        if self._interface.poll():
+            result = self.send_and_receive_dict(
+                input_dict={"shutdown": True, "wait": wait}
+            )
+            self._interface.shutdown(wait=wait)
+        if self._socket is not None:
+            self._socket.close()
+        if self._context is not None:
+            self._context.term()
+        self._process = None
+        self._socket = None
+        self._context = None
+        return result
+
+    def __del__(self):
+        self.shutdown(wait=True)
+
+
+def interface_bootup(
+    command_lst: list[str],
+    connections,
+    hostname_localhost: bool = False,
+    prefix_name: Optional[str] = None,
+    prefix_path: Optional[str] = None,
+):
+    """
+    Start interface for ZMQ communication
+
+    Args:
+        command_lst (list): List of commands as strings
+        connections (executorlib.shared.interface.BaseInterface): Interface to start parallel process, like MPI, SLURM
+                                                                  or Flux
+        hostname_localhost (boolean): use localhost instead of the hostname to establish the zmq connection. In the
+                                      context of an HPC cluster this essential to be able to communicate to an
+                                      Executor running on a different compute node within the same allocation. And
+                                      in principle any computer should be able to resolve that their own hostname
+                                      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
+        prefix_name (str): name of the conda environment to initialize
+        prefix_path (str): path of the conda environment to initialize
+
+    Returns:
+         executorlib.shared.communication.SocketInterface: socket interface for zmq communication
+    """
+    if not hostname_localhost:
+        command_lst += [
+            "--host",
+            gethostname(),
+        ]
+    interface = SocketInterface(interface=connections)
+    command_lst += [
+        "--zmqport",
+        str(interface.bind_to_random_port()),
+    ]
+    interface.bootup(
+        command_lst=command_lst, prefix_name=prefix_name, prefix_path=prefix_path
+    )
+    return interface
+
+
+def interface_connect(host: str, port: str):
+    """
+    Connect to an existing SocketInterface instance by providing the hostname and the port as strings.
+
+    Args:
+        host (str): hostname of the host running the SocketInterface instance to connect to.
+        port (str): port on the host the SocketInterface instance is running on.
+    """
+    context = zmq.Context()
+    socket = context.socket(zmq.PAIR)
+    socket.connect("tcp://" + host + ":" + port)
+    return context, socket
+
+
+def interface_send(socket: zmq.Socket, result_dict: dict):
+    """
+    Send results to a SocketInterface instance.
+
+    Args:
+        socket (zmq.Socket): socket for the connection
+        result_dict (dict): dictionary to be sent, supported keys are result, error and error_type.
+    """
+    socket.send(cloudpickle.dumps(result_dict))
+
+
+def interface_receive(socket: zmq.Socket):
+    """
+    Receive instructions from a SocketInterface instance.
+
+    Args:
+        socket (zmq.Socket): socket for the connection
+    """
+    return cloudpickle.loads(socket.recv())
+
+
+def interface_shutdown(socket: zmq.Socket, context: zmq.Context):
+    """
+    Close the connection to a SocketInterface instance.
+
+    Args:
+        socket (zmq.Socket): socket for the connection
+        context (zmq.sugar.context.Context): context for the connection
+    """
+    socket.close()
+    context.term()