Integrate Upstream Engine Changes. Closes #157 (#160)

* WIP- Cluster state changes

* Integrate Upstream Changes. Closes #157

* Wait for 10 seconds

* OBJC_DISABLE_INITIALIZE_FORK_SAFETY=YES

* Fix init_services

.github/workflows/CI.yaml

@@ -48,6 +48,7 @@ jobs:
         run: |
           brew install portaudio
           echo "MANUAL_OS_SET=Darwin" >> $GITHUB_ENV
+          echo "OBJC_DISABLE_INITIALIZE_FORK_SAFETY=YES" >> $GITHUB_ENV
 
       - uses: actions/checkout@v3
         name: Checkout ChimeraPy-Engine (Bleeding Edge)

chimerapy/orchestrator/init_services.py

@@ -42,10 +42,9 @@ def get(name):
 
 def teardown():
     """Teardown the services."""
-    manager = available_services.get("cluster_manager")
+    manager: ClusterManager = available_services.get("cluster_manager")
     if manager and not manager.has_shutdown():
         manager.shutdown()
 
     for worker in available_services.get("workers", []):
-        if not worker.has_shutdown:
-            worker.shutdown()
+        worker.shutdown()

chimerapy/orchestrator/models/cluster_models.py

@@ -14,17 +14,50 @@
 )
 
 
+class RegisteredMethod(BaseModel):
+    name: str
+    style: Literal["concurrent", "blocking", "reset"] = "concurrent"
+    params: Dict[str, str] = Field(default_factory=dict)
+
+
+class NodeDiagnostics(BaseModel):
+    timestamp: str
+    latency: float
+    payload_size: float
+    memory_usage: float
+    cpu_usage: float
+    num_of_steps: int
+
+
 class NodeState(BaseModel):
     id: str
     name: str = ""
-    init: bool = False
-    connected: bool = False
-    ready: bool = False
-    finished: bool = False
     port: int = 0
+    fsm: Literal[
+        "NULL",
+        "INITIALIZED",
+        "CONNECTED",
+        "READY",
+        "PREVIEWING",
+        "RECORDING",
+        "STOPPED",
+        "SAVED",
+        "SHUTDOWN",
+    ]
+    registered_methods: Dict[str, RegisteredMethod] = Field(
+        default_factory=dict
+    )
+    logdir: Optional[str] = None
+    diagnostics: NodeDiagnostics
 
     @classmethod
     def from_cp_node_state(cls, node_state: _NodeState):
+        node_state_dict = node_state.to_dict()
+        node_state_dict["logdir"] = (
+            str(node_state_dict["logdir"])
+            if node_state_dict["logdir"] is not None
+            else None
+        )
         return cls(**node_state.to_dict())
 
     model_config: ClassVar[ConfigDict] = ConfigDict(frozen=True, extra="forbid")
@@ -36,10 +69,12 @@ class WorkerState(BaseModel):
     port: int = 0
     ip: str = ""
     nodes: Dict[str, NodeState] = Field(default_factory=dict)
+    tempfolder: str = ""
 
     @classmethod
     def from_cp_worker_state(cls, worker_state: _WorkerState):
         state_dict = worker_state.to_dict()
+        state_dict["tempfolder"] = str(state_dict["tempfolder"])
         return cls(
             **state_dict,
         )
@@ -54,16 +89,16 @@ class ClusterState(BaseModel):
     workers: Dict[str, WorkerState] = Field(default_factory=dict)
 
     logs_subscription_port: Optional[int] = None
-    running: bool = False
-    collecting: bool = False
-    collection_status: Optional[Literal["PASS", "FAIL"]] = None
+    log_sink_enabled: bool = False
     zeroconf_discovery: bool = False
+    logdir: str = None
 
     @classmethod
     def from_cp_manager_state(
         cls, state: ManagerState, zeroconf_discovery: bool
     ):
         state_dict = state.to_dict()
+        state_dict["logdir"] = str(state_dict["logdir"])
         return cls(**state_dict, zeroconf_discovery=zeroconf_discovery)
 
     model_config: ClassVar[ConfigDict] = ConfigDict(frozen=True, extra="forbid")

chimerapy/orchestrator/services/cluster_service/cluster_manager.py

@@ -137,7 +137,7 @@ async def update_network_status(self) -> None:
 
     def is_zeroconf_discovery_enabled(self) -> bool:
         """Check if zeroconf discovery is enabled."""
-        return self._manager.services.zeroconf.enabled
+        return self._manager.zeroconf_service.enabled
 
     async def instantiate_pipeline(
         self, pipeline_id

chimerapy/orchestrator/services/cluster_service/updates_broadcaster.py

@@ -110,7 +110,7 @@ async def add_client(
         await self.updater.add_client(q)
 
         if message is not None:
-            await q.put(message.dict())
+            await q.put(message.model_dump(mode="json"))
 
     async def remove_client(self, q: asyncio.Queue) -> None:
         """Remove a client queue from the broadcaster."""
@@ -139,7 +139,7 @@ async def broadcast_updates(self) -> None:
                 else:
                     msg = None
                 if msg is not None:
-                    msg_dict = msg.dict()
+                    msg_dict = msg.model_dump(mode="json")
                     await self.updater.put_update(msg_dict)
                 if msg and msg.signal is UpdateMessageType.SHUTDOWN:
                     break
@@ -165,7 +165,7 @@ async def put_update(self, msg: Dict[str, Any]) -> None:
             UpdateMessageType.NETWORK_UPDATE,
             self.zeroconf_enabled,
         )
-        await self.updater.put_update(update_msg.dict())
+        await self.updater.put_update(update_msg.model_dump(mode="json"))
 
     @staticmethod
     def is_cluster_update_message(msg: Dict[str, Any]) -> bool:

chimerapy/orchestrator/tests/models/test_cluster_models.py

@@ -0,0 +1,223 @@
+from pathlib import Path
+
+import pytest
+
+from chimerapy.engine.states import (
+    ManagerState as _ManagerState,
+)
+from chimerapy.engine.states import (
+    NodeState as _NodeState,
+)
+from chimerapy.engine.states import (
+    RegisteredMethod as _RegisteredMethod,
+)
+from chimerapy.engine.states import (
+    WorkerState as _WorkerState,
+)
+from chimerapy.orchestrator.models.cluster_models import (
+    ClusterState,
+    NodeState,
+    RegisteredMethod,
+    WorkerState,
+)
+from chimerapy.orchestrator.tests.base_test import BaseTest
+
+
+class TestClusterModels(BaseTest):
+    @pytest.fixture(scope="class")
+    def m_empty(self):
+        return _ManagerState(
+            id="manager1",
+        )
+
+    @pytest.fixture(scope="class")
+    def m_populated(self):
+        return _ManagerState(
+            id="manager1",
+            logdir=Path("/tmp24"),
+            ip="192.168.2.0",
+            port=55000,
+            workers={},
+            log_sink_enabled=True,
+            logs_subscription_port=55001,
+        )
+
+    @pytest.fixture(scope="class")
+    def m_w_empty(self):
+        return _ManagerState(
+            id="manager1",
+            logdir=Path("/tmp24"),
+            ip="192.168.2.0",
+            port=55000,
+            log_sink_enabled=True,
+            logs_subscription_port=55001,
+            workers={
+                "w1": _WorkerState(
+                    id="w1",
+                    name="worker1",
+                    port=55002,
+                    ip="192.168.2.1",
+                    nodes={},
+                    tempfolder=Path("/tmp24/w1"),
+                ),
+                "w2": _WorkerState(
+                    id="w2",
+                    name="worker2",
+                    port=55003,
+                    ip="192.168.2.2",
+                    nodes={},
+                    tempfolder=Path("/tmp24/w2"),
+                ),
+            },
+        )
+
+    @pytest.fixture(scope="class")
+    def m_w_populated(self):
+        return _ManagerState(
+            id="manager1",
+            logdir=Path("/tmp24"),
+            ip="192.168.2.0",
+            port=55000,
+            log_sink_enabled=True,
+            logs_subscription_port=55001,
+            workers={
+                "w1": _WorkerState(
+                    id="w1",
+                    name="worker1",
+                    port=55002,
+                    ip="192.168.2.1",
+                    nodes={
+                        "n1": _NodeState(
+                            id="n1",
+                            name="node1",
+                            port=55004,
+                            fsm="NULL",
+                            registered_methods={
+                                "m1": _RegisteredMethod(
+                                    name="func1",
+                                    style="concurrent",
+                                    params={
+                                        "p1": "int",
+                                    },
+                                )
+                            },
+                        )
+                    },
+                    tempfolder=Path("/tmp24/w1"),
+                ),
+                "w2": _WorkerState(
+                    id="w2",
+                    name="worker2",
+                    port=55003,
+                    ip="192.168.2.2",
+                    nodes={
+                        "n2": _NodeState(
+                            id="n2",
+                            name="node2",
+                            port=55005,
+                            fsm="INITIALIZED",
+                            registered_methods={
+                                "m2": _RegisteredMethod(
+                                    name="func2",
+                                    style="concurrent",
+                                    params={
+                                        "p2": "float",
+                                    },
+                                )
+                            },
+                        )
+                    },
+                    tempfolder=Path("/tmp24/w2"),
+                ),
+            },
+        )
+
+    def test_m_empty(self, m_empty):
+        manager_state = ClusterState.from_cp_manager_state(
+            m_empty, zeroconf_discovery=False
+        )
+        assert manager_state.id == "manager1"
+        assert manager_state.logdir == str(Path.cwd())
+        assert manager_state.log_sink_enabled is False
+        assert manager_state.logs_subscription_port is None
+        assert manager_state.ip == "0.0.0.0"
+        assert manager_state.port == 0
+        assert manager_state.workers == {}
+
+    def test_m_populated(self, m_populated):
+        manager_state = ClusterState.from_cp_manager_state(
+            m_populated, zeroconf_discovery=True
+        )
+        assert manager_state.id == "manager1"
+        assert manager_state.logdir == "/tmp24"
+        assert manager_state.log_sink_enabled is True
+        assert manager_state.logs_subscription_port == 55001
+        assert manager_state.ip == "192.168.2.0"
+        assert manager_state.port == 55000
+        assert manager_state.workers == {}
+        assert manager_state.zeroconf_discovery is True
+
+    def test_m_w_empty(self, m_w_empty):
+        manager_state = ClusterState.from_cp_manager_state(
+            m_w_empty, zeroconf_discovery=True
+        )
+        assert len(manager_state.workers) == 2
+        w1 = manager_state.workers["w1"]
+        assert w1.id == "w1"
+        assert w1.name == "worker1"
+        assert w1.port == 55002
+        assert w1.ip == "192.168.2.1"
+        assert w1.tempfolder == "/tmp24/w1"
+        assert w1.nodes == {}
+        assert isinstance(w1, WorkerState)
+
+        w2 = manager_state.workers["w2"]
+        assert w2.id == "w2"
+        assert w2.name == "worker2"
+        assert w2.port == 55003
+        assert w2.ip == "192.168.2.2"
+        assert w2.tempfolder == "/tmp24/w2"
+        assert w2.nodes == {}
+        assert isinstance(w2, WorkerState)
+
+    def test_m_w_populated(self, m_w_populated):
+        manager_state = ClusterState.from_cp_manager_state(
+            m_w_populated, zeroconf_discovery=True
+        )
+        assert len(manager_state.workers) == 2
+        w1 = manager_state.workers["w1"]
+
+        n1 = w1.nodes["n1"]
+        assert isinstance(n1, NodeState)
+        assert n1.id == "n1"
+        assert n1.name == "node1"
+        assert n1.port == 55004
+        assert n1.fsm == "NULL"
+
+        assert n1.registered_methods == {
+            "m1": RegisteredMethod(
+                name="func1",
+                style="concurrent",
+                params={
+                    "p1": "int",
+                },
+            )
+        }
+
+        w2 = manager_state.workers["w2"]
+        n2 = w2.nodes["n2"]
+        assert isinstance(n2, NodeState)
+        assert n2.id == "n2"
+        assert n2.name == "node2"
+        assert n2.port == 55005
+        assert n2.fsm == "INITIALIZED"
+
+        assert n2.registered_methods == {
+            "m2": RegisteredMethod(
+                name="func2",
+                style="concurrent",
+                params={
+                    "p2": "float",
+                },
+            )
+        }

chimerapy/orchestrator/tests/models/test_pipeline_config.py

@@ -17,7 +17,7 @@ def dummy_pipeline_config(self) -> ChimeraPyPipelineConfig:
 
     def test_pipeline_config(self, dummy_pipeline_config):
         assert dummy_pipeline_config.name == "Pipeline"
-        assert dummy_pipeline_config.description == "A pipeline"
+        assert dummy_pipeline_config.description == ""
         assert dummy_pipeline_config.runtime == 2000
 
     def test_worker_config(self, dummy_pipeline_config):