Move connection backend config into Manager __init__

Author: dav0dea

src/goofi/connection.py

@@ -48,15 +48,16 @@ def set_backend(backend: str, mp_manager: BaseManager) -> None:
         `mp_manager` : multiprocessing.Manager
             The multiprocessing manager to use for creating shared objects.
         """
+        # initialize a shared set of connection ids
+        assert Connection._CONNECTION_IDS is None, "Connection._CONNECTION_IDS is not None."
+        Connection._CONNECTION_IDS = mp_manager.list()
+
+        # set the backend if discovering connection IDs was successful (prevents overwriting the backend)
         assert (
             backend in Connection.get_backends().keys()
         ), f"Invalid backend: {backend}. Choose from {list(Connection.get_backends().keys())}"
         Connection._BACKEND = backend
 
-        # initialize a shared set of connection ids
-        assert Connection._CONNECTION_IDS is None, "Connection._CONNECTION_IDS is not None."
-        Connection._CONNECTION_IDS = mp_manager.list()
-
     @staticmethod
     def create() -> Tuple["Connection", "Connection"]:
         """

src/goofi/manager.py

@@ -112,11 +112,29 @@ class Manager:
         same process as the manager.
     `duration` : float
         The duration to run the manager for. If `0`, runs indefinitely.
+    `communication_backend` : str
+        The communication backend to use for node communication. Default is "mp" (multiprocessing).
     """
 
     def __init__(
-        self, filepath: Optional[str] = None, headless: bool = True, use_multiprocessing: bool = True, duration: float = 0
+        self,
+        filepath: Optional[str] = None,
+        headless: bool = True,
+        use_multiprocessing: bool = True,
+        duration: float = 0,
+        communication_backend: str = "mp",
     ) -> None:
+        # create a multiprocessing manager
+        self._mp_manager = MPManager()
+
+        try:
+            # set up communication backend
+            Connection.set_backend(communication_backend, self._mp_manager)
+        except AssertionError:
+            print("Connection backend already set. Skipping.")
+        # make sure the connection backend is set correctly
+        assert Connection._BACKEND == communication_backend
+
         # TODO: add proper logging
         print("Starting goofi-pipe...")
         # preload all nodes to avoid delays
@@ -319,7 +337,9 @@ def remove_node(self, name: str, notify_gui: bool = True, **gui_kwargs) -> None:
             Window().remove_node(name, **gui_kwargs)
 
     @mark_unsaved_changes
-    def add_link(self, node_out: str, node_in: str, slot_out: str, slot_in: str, notify_gui: bool = True, **gui_kwargs) -> None:
+    def add_link(
+        self, node_out: str, node_in: str, slot_out: str, slot_in: str, notify_gui: bool = True, **gui_kwargs
+    ) -> None:
         """
         Adds a link between two nodes.
 
@@ -403,6 +423,9 @@ def terminate(self, notify_gui: bool = True) -> None:
         for node in self.nodes:
             self.nodes[node].terminate()
 
+        # close the communication backend
+        self._mp_manager.shutdown()
+
     def save(self, filepath: Optional[str] = None, overwrite: bool = False, timeout: float = 3.0) -> None:
         """
         Saves the state of the manager to a file.
@@ -585,21 +608,14 @@ def main(duration: float = 0, args=None):
         docs()
         return
 
-    with MPManager() as manager:
-        # set the communication backend
-        try:
-            Connection.set_backend(args.comm, manager)
-        except AssertionError:
-            # connection backend is already set (occurrs when running tests)
-            pass
-
-        # create and run the manager (this blocks until the manager is terminated)
-        Manager(
-            filepath=args.filepath,
-            headless=args.headless,
-            use_multiprocessing=not args.no_multiprocessing,
-            duration=duration,
-        )
+    # create and run the manager (this blocks until the manager is terminated)
+    Manager(
+        filepath=args.filepath,
+        headless=args.headless,
+        use_multiprocessing=not args.no_multiprocessing,
+        duration=duration,
+        communication_backend=args.comm,
+    )
 
 
 def docs():

tests/test_manager.py

@@ -1,6 +1,5 @@
 import platform
 import time
-from multiprocessing import Manager as MPManager
 from os import path
 
 import pytest
@@ -10,16 +9,22 @@
 from goofi.manager import Manager
 from goofi.message import Message, MessageType
 
+MANAGER_TEST_DURATION = 0.1
 
-def create_simple_manager() -> Manager:
+
+def create_simple_manager(comm_backend: str = "mp") -> Manager:
     """
     Creates a simple manager with a constant node, a sine node, and an add node.
 
+    ### Parameters
+    `comm_backend` : str
+        The communication backend to use. Choose from "mp", "zmq-tcp" or "zmq-ipc".
+
     ### Returns
     manager : Manager
         The manager object.
     """
-    manager = Manager()
+    manager = Manager(duration=MANAGER_TEST_DURATION, communication_backend=comm_backend)
     manager.add_node("ConstantArray", "inputs")
     manager.add_node("Sine", "inputs")
     manager.add_node("Operation", "array")
@@ -30,11 +35,11 @@ def create_simple_manager() -> Manager:
 
 
 def test_creation():
-    Manager()
+    Manager(duration=MANAGER_TEST_DURATION)
 
 
 def test_main():
-    goofi.manager.main(1, ["--headless"])
+    goofi.manager.main(MANAGER_TEST_DURATION, ["--headless"])
 
 
 @pytest.mark.skipif(platform.system() == "Windows", reason="Multiprocessing is very slow on Windows.")
@@ -44,14 +49,7 @@ def test_simple(comm_backend):
         # TODO: make sure zmq backend works
         pytest.skip("ZeroMQ backend still has some issues.")
 
-    try:
-        mp_manager = MPManager()
-        Connection.set_backend("mp", mp_manager)
-    except AssertionError:
-        # connection backend is already set
-        pass
-
-    manager = create_simple_manager()
+    manager = create_simple_manager(comm_backend=comm_backend)
     my_conn, node_conn = Connection.create()
     manager.nodes["operation0"].connection.send(
         Message(MessageType.ADD_OUTPUT_PIPE, {"slot_name_out": "out", "slot_name_in": "in", "node_connection": my_conn})
@@ -80,7 +78,7 @@ def test_simple(comm_backend):
 
 
 def test_save_empty(tmpdir):
-    manager = Manager()
+    manager = Manager(duration=MANAGER_TEST_DURATION)
 
     # if path is a file, save to that file
     manager.save(path.join(tmpdir, "test.gfi"))