fix: manage config changes even when actor disabled

Author: d-perl

bec_server/bec_server/actors/builtin_actor_manager.py

@@ -1,4 +1,5 @@
 from threading import Event, Thread
+from typing import TypeVar
 
 from bec_lib.client import BECClient, ServiceConfig
 from bec_lib.connector import MessageObject
@@ -7,12 +8,29 @@
 from bec_lib.messages import (
     BuiltinActorStateChangeNotification,
     BuiltinActorStateUpdatedNotification,
+    ScanInterlockModifyStateTableMessage,
+    ScanInterlockStateTableContent,
 )
 from bec_server.actors.actor import ActorBase
 from bec_server.actors.scan_interlock import ScanInterlockActor
 
 logger = bec_logger.logger
 
+ActorType = TypeVar("ActorType", bound=ActorBase)
+
+
+class ActorDict(dict):
+    def __setitem__(self, key: type[ActorType], value: tuple[ActorType, Thread, Event], /) -> None:
+        return super().__setitem__(key, value)
+
+    def __getitem__(self, key: type[ActorType], /) -> tuple[ActorType, Thread, Event]:
+        return super().__getitem__(key)
+
+    def get(  # type: ignore
+        self, key: type[ActorType], default: tuple[ActorType, Thread, Event] | None = None
+    ) -> tuple[ActorType, Thread, Event]:
+        return super().get(key)  # type: ignore
+
 
 class BuiltinActorManager:
     """A simple manager for builtin actors which are always available - only handles enabling and
@@ -25,12 +43,15 @@ def __init__(self, bootstrap_server: str) -> None:
             name="BuiltinActors",
         )
         self._client.start()
-        self._actors_threads_and_stops: dict[str, tuple[ActorBase, Thread, Event]] = {}
+        self._actors_threads_and_stops = ActorDict()
         self._builtin_actors = {cls.__name__: cls for cls in (ScanInterlockActor,)}
         self._start_all()
         self._client.connector.register(
             MessageEndpoints.builtin_actor_update_req_notif(), cb=self._on_state_changed
         )
+        self._client.connector.register(
+            MessageEndpoints.modify_interlock_table(), cb=self._modify_interlock_table
+        )
 
     def _ping_clients(self, actor_name: str):
         self._client.connector.send(
@@ -51,33 +72,68 @@ def _on_state_changed(self, msg_obj: MessageObject):
         self._ping_clients(msg.actor_name)
 
     def _start_all(self):
-        for actor_class in self._builtin_actors.values():
-            if self._client.builtin_actors.check_enabled(actor_class.__name__):
-                self._start_actor(actor_class)
+        for actor_class_name in self._builtin_actors:
+            if self._client.builtin_actors.check_enabled(actor_class_name):
+                self._start_actor(self._builtin_actors[actor_class_name])
 
     def _start_actor(self, actor_class: type[ActorBase]):
         name = actor_class.__name__
         logger.info(f"Starting {name}")
-        if name in self._actors_threads_and_stops:
+        if actor_class in self._actors_threads_and_stops:
             logger.warning(f"Actor {name} is already active!")
             return
         actor = actor_class(self._client, name=name, exec_id=name)
         t = Thread(target=actor.run)
-        self._actors_threads_and_stops[name] = (actor, t, actor.stop_event)
+        self._actors_threads_and_stops[actor_class] = (actor, t, actor.stop_event)
         t.start()
 
     def _stop_actor(self, actor_name: str):
         logger.info(f"Stopping {actor_name}")
-        if (entry := self._actors_threads_and_stops.get(actor_name)) is None:
+        actor_class = self._builtin_actors.get(actor_name)
+        if (entry := self._actors_threads_and_stops.get(actor_class)) is None:
             logger.warning(f"Actor {actor_name} is not active!")
             return
         actor, t, event = entry
         event.set()
         t.join()
-        del self._actors_threads_and_stops[actor_name]
+        del self._actors_threads_and_stops[actor_class]
         del actor
 
     def shutdown(self):
         for actor in self._actors_threads_and_stops:
             self._stop_actor(actor)
         self._client.shutdown()
+
+    # Actor specific management methods:
+    def _modify_interlock_table(self, msg_dict):
+        """Update the watched states for ScanInterlockActor - handled by the actor itself if it is
+        active, otherwise just the config in redis is updated."""
+        msg: ScanInterlockModifyStateTableMessage = msg_dict["data"]
+        if (ats := self._actors_threads_and_stops.get(ScanInterlockActor)) is not None:
+            actor, _, _ = ats
+            actor._on_state_modification(msg)
+        else:
+            states: ScanInterlockStateTableContent | None = self._client.connector.get(
+                MessageEndpoints.scan_interlock_states()
+            )
+            current_watched = states.states_watched if states is not None else {}
+            if msg.action == "add":
+                logger.info(f"Adding {msg.state_name} to the scan interlock actor")
+                current_watched[msg.state_name] = msg.status
+                self._client.connector.set(
+                    MessageEndpoints.scan_interlock_states(),
+                    ScanInterlockStateTableContent(states_watched=current_watched),
+                )
+            elif msg.action == "remove_all":
+                self._client.connector.set(
+                    MessageEndpoints.scan_interlock_states(),
+                    ScanInterlockStateTableContent(states_watched={}),
+                )
+            else:
+                logger.info(f"Removing {msg.state_name} from the scan interlock actor")
+                current_watched.pop(msg.state_name, None)
+                self._client.connector.set(
+                    MessageEndpoints.scan_interlock_states(),
+                    ScanInterlockStateTableContent(states_watched=current_watched),
+                )
+            self._ping_clients("ScanInterlockActor")

bec_server/bec_server/actors/scan_interlock.py

@@ -26,9 +26,6 @@ def __init__(self, client: BECClient, name: str, exec_id: str):
             self.state_table = {}
 
         super().__init__(client, name, exec_id)
-        self.client.connector.register(
-            MessageEndpoints.modify_interlock_table(), cb=self._on_state_modification
-        )
 
     def _ping_clients(self):
         logger.warning(self.name)
@@ -44,8 +41,7 @@ def _update_watched_states_in_redis(self):
         )
         self._ping_clients()
 
-    def _on_state_modification(self, msg_dict: dict):
-        msg: ScanInterlockModifyStateTableMessage = msg_dict["data"]
+    def _on_state_modification(self, msg: ScanInterlockModifyStateTableMessage):
         with self.state_table_lock:
             if msg.action == "add":
                 logger.info(f"Adding {msg.state_name} to the scan interlock actor")

bec_server/tests/tests_scan_server/test_builtin_actor_manager.py

@@ -29,20 +29,25 @@ def mocked_manager():
         mock_client_cls.return_value = mock_client
 
         manager = BuiltinActorManager("localhost:6379")
-        yield manager, mock_client
+        with patch.object(manager, "_builtin_actors", {"DummyActor": DummyActor}):
+            yield manager, mock_client
 
 
 def test_init_registers_callback(mocked_manager):
     manager, mock_client = mocked_manager
 
     mock_client.start.assert_called_once()
 
-    mock_client.connector.register.assert_called_once()
-    args, kwargs = mock_client.connector.register.call_args
+    assert mock_client.connector.register.call_count == 2
 
+    kwargs = mock_client.connector.register.call_args_list[0].kwargs
     assert "cb" in kwargs
     assert kwargs["cb"] == manager._on_state_changed
 
+    kwargs = mock_client.connector.register.call_args_list[1].kwargs
+    assert "cb" in kwargs
+    assert kwargs["cb"] == manager._modify_interlock_table
+
 
 def test_start_actor_starts_thread(mocked_manager):
     manager, _ = mocked_manager
@@ -53,9 +58,9 @@ def test_start_actor_starts_thread(mocked_manager):
 
         manager._start_actor(DummyActor)
 
-        assert "DummyActor" in manager._actors_threads_and_stops
+        assert DummyActor in manager._actors_threads_and_stops
 
-        actor, thread, stop_event = manager._actors_threads_and_stops["DummyActor"]
+        actor, thread, stop_event = manager._actors_threads_and_stops[DummyActor]
 
         assert isinstance(actor, DummyActor)
         assert thread == mock_thread
@@ -80,7 +85,7 @@ def test_stop_actor_sets_event_and_joins(mocked_manager):
     actor = DummyActor(None, "DummyActor", "DummyActor")
     mock_thread = MagicMock()
 
-    manager._actors_threads_and_stops["DummyActor"] = (actor, mock_thread, actor.stop_event)
+    manager._actors_threads_and_stops[DummyActor] = (actor, mock_thread, actor.stop_event)
 
     manager._stop_actor("DummyActor")
 

bec_server/tests/tests_scan_server/test_scan_interlock.py

@@ -52,7 +52,7 @@ def test_on_state_modification_remove_all(self, actor, mock_client):
         msg.status = None
 
         with patch("bec_server.actors.actor.BlStateActor.evaluate"):
-            actor._on_state_modification({"data": msg})
+            actor._on_state_modification(msg)
 
         assert actor.state_table == {}
         assert actor.state_cache == {}
@@ -63,7 +63,7 @@ def test_on_state_modification_add(self, actor, mock_client):
         actor._update_cache = MagicMock()
         mock_client.connector.register.reset_mock()
         msg = MagicMock(action="add", state_name="beam_ok", status="valid")
-        actor._on_state_modification({"data": msg})
+        actor._on_state_modification(msg)
         assert actor.state_table["beam_ok"] == "valid"
         mock_client.connector.register.assert_called_once()
         actor._update_cache.assert_called_once()
@@ -76,7 +76,7 @@ def test_on_state_modification_remove(self, actor, mock_client):
         msg = MagicMock(action="remove", state_name="beam_ok", status=None)
 
         with patch("bec_server.actors.actor.BlStateActor.evaluate"):
-            actor._on_state_modification({"data": msg})
+            actor._on_state_modification(msg)
 
         assert "beam_ok" not in actor.state_table
         assert "beam_ok" not in actor.state_cache
@@ -88,7 +88,7 @@ def test_on_state_modification_remove_missing(self, actor, mock_client):
         msg = MagicMock(action="remove", state_name="missing", status=None)
 
         with patch("bec_server.actors.actor.BlStateActor.evaluate"):
-            actor._on_state_modification({"data": msg})
+            actor._on_state_modification(msg)
 
         mock_client.connector.unregister.assert_not_called()