bec/bec_widgets
0
0
Fork 0
mirror of https://github.com/bec-project/bec_widgets.git synced 2025-07-13 19:21:50 +02:00
Code Activity

refactor(rpc_reference): refactor rpc reference tracking

Browse Source
This commit is contained in:
appel_c
2025-03-19 16:55:00 +01:00
committed by wyzula-jan
parent f3d3c9425d
commit bd5e251ee9
4 changed files with 124 additions and 64 deletions
Download Patch File Download Diff File Expand all files Collapse all files

118
bec_widgets/cli/client_utils.py
View File

@ -11,6 +11,7 @@ import subprocess
import threading import threading
import time import time
from contextlib import contextmanager from contextlib import contextmanager
from threading import Lock
from typing import TYPE_CHECKING, Any from typing import TYPE_CHECKING, Any
from bec_lib.endpoints import MessageEndpoints from bec_lib.endpoints import MessageEndpoints
@ -202,6 +203,7 @@ class BECGuiClient(RPCBase):
def __init__(self, **kwargs) -> None: def __init__(self, **kwargs) -> None:
super().__init__(**kwargs) super().__init__(**kwargs)
self._lock = Lock()
self._default_dock_name = "bec" self._default_dock_name = "bec"
self._auto_updates_enabled = True self._auto_updates_enabled = True
self._auto_updates = None self._auto_updates = None
@ -211,7 +213,7 @@ class BECGuiClient(RPCBase):
self._gui_started_event = threading.Event() self._gui_started_event = threading.Event()
self._process = None self._process = None
self._process_output_processing_thread = None self._process_output_processing_thread = None
self._exposed_dock_areas = [] self._exposed_widgets = []
self._registry_state = {} self._registry_state = {}
self._ipython_registry = {} self._ipython_registry = {}
self.available_widgets = AvailableWidgetsNamespace() self.available_widgets = AvailableWidgetsNamespace()
@ -312,6 +314,10 @@ class BECGuiClient(RPCBase):
def kill_server(self) -> None: def kill_server(self) -> None:
"""Kill the GUI server.""" """Kill the GUI server."""
self._top_level.clear() self._top_level.clear()
# Unregister the registry state
self._client.connector.unregister(
MessageEndpoints.gui_registry_state(self._gui_id), cb=self._handle_registry_update
)
self._killed = True self._killed = True
if self._gui_started_timer is not None: if self._gui_started_timer is not None:
@ -416,6 +422,7 @@ class BECGuiClient(RPCBase):
return self._start_server(wait=wait) return self._start_server(wait=wait)
def _handle_registry_update(self, msg: StreamMessage) -> None: def _handle_registry_update(self, msg: StreamMessage) -> None:
# with self._lock:
self._registry_state = msg["data"].state self._registry_state = msg["data"].state
self._update_dynamic_namespace() self._update_dynamic_namespace()
@ -433,39 +440,51 @@ class BECGuiClient(RPCBase):
with wait_for_server(self): with wait_for_server(self):
rpc_client = RPCBase(gui_id=f"{self._gui_id}:window", parent=self) rpc_client = RPCBase(gui_id=f"{self._gui_id}:window", parent=self)
rpc_client._run_rpc("hide") # pylint: disable=protected-access rpc_client._run_rpc("hide") # pylint: disable=protected-access
# because of the registry callbacks, we may have
# dock areas that are already killed, but not yet
# removed from the registry state
if not self._killed: if not self._killed:
for window in self._top_level.values(): for window in self._top_level.values():
window.hide() window.hide()
def _update_dynamic_namespace(self): def _update_dynamic_namespace(self):
"""Update the dynamic name space""" """Update the dynamic name space"""
self._clear_exposed_dock_areas() # First we update the name space based on the new registry state
self._cleanup_ipython_registry()
self._add_registry_to_namespace() self._add_registry_to_namespace()
# Then we clear the ipython registry from old objects
def _clear_exposed_dock_areas(self): self._cleanup_ipython_registry()
"""Clear the exposed dock areas"""
self._top_level.clear()
for widget_id in self._exposed_dock_areas:
delattr(self, widget_id)
self._exposed_dock_areas.clear()
def _cleanup_ipython_registry(self): def _cleanup_ipython_registry(self):
"""Cleanup the ipython registry""" """Cleanup the ipython registry"""
remove_ids = [] names_in_registry = list(self._ipython_registry.keys())
for widget_id in self._ipython_registry: remove_ids = list(set(names_in_registry) - set(self._exposed_widgets))
if widget_id not in self._registry_state:
remove_ids.append(widget_id)
for widget_id in remove_ids: for widget_id in remove_ids:
self._ipython_registry.pop(widget_id) self._ipython_registry.pop(widget_id)
self._cleanup_rpc_references_on_rpc_base(remove_ids)
# Clear the exposed widgets
self._exposed_widgets.clear()
def _cleanup_rpc_references_on_rpc_base(self, remove_ids: list[str]) -> None:
"""Cleanup the rpc references on the RPCBase object"""
if not remove_ids:
return
for widget in self._ipython_registry.values():
to_delete = []
for attr_name, gui_id in widget._rpc_references.items():
if gui_id in remove_ids:
to_delete.append(attr_name)
for attr_name in to_delete:
if hasattr(widget, attr_name):
delattr(widget, attr_name)
if attr_name.startswith("elements."):
delattr(widget.elements, attr_name.split(".")[1])
widget._rpc_references.pop(attr_name)
def _set_dynamic_attributes(self, obj: object, name: str, value: Any) -> None: def _set_dynamic_attributes(self, obj: object, name: str, value: Any) -> None:
"""Add an object to the namespace""" """Add an object to the namespace"""
setattr(obj, name, value) setattr(obj, name, value)
def _update_rpc_references(self, widget: RPCBase, name: str, gui_id: str) -> None:
"""Update the RPC references"""
widget._rpc_references[name] = gui_id
def _add_registry_to_namespace(self) -> None: def _add_registry_to_namespace(self) -> None:
"""Add registry to namespace""" """Add registry to namespace"""
# Add dock areas # Add dock areas
@ -475,33 +494,51 @@ class BECGuiClient(RPCBase):
if state["widget_class"] == "BECDockArea" if state["widget_class"] == "BECDockArea"
] ]
for state in dock_area_states: for state in dock_area_states:
# obj is an RPC reference to the RPCBase object dock_area_ref = self._add_widget(state, self)
dock_area_obj = self._add_widget(state, self) dock_area = self._ipython_registry.get(dock_area_ref._gui_id)
self._set_dynamic_attributes(self, dock_area_obj.widget_name, dock_area_obj) if not hasattr(dock_area, "elements"):
self._set_dynamic_attributes(dock_area, "elements", WidgetNameSpace())
self._set_dynamic_attributes(self, dock_area.widget_name, dock_area_ref)
# Keep track of rpc references on RPCBase object
self._update_rpc_references(self, dock_area.widget_name, dock_area_ref._gui_id)
# Add dock_area to the top level # Add dock_area to the top level
self._top_level[dock_area_obj.widget_name] = dock_area_obj self._top_level[dock_area_ref.widget_name] = dock_area_ref
self._exposed_dock_areas.append(dock_area_obj.widget_name) self._exposed_widgets.append(dock_area_ref._gui_id)
# Add docks # Add docks
dock_states = [ dock_states = [
state state
for state in self._registry_state.values() for state in self._registry_state.values()
if state["config"].get("parent_id", "") == dock_area_obj._gui_id if state["config"].get("parent_id", "") == dock_area_ref._gui_id
] ]
for state in dock_states: for state in dock_states:
dock_obj = self._add_widget(state, dock_area_obj) dock_ref = self._add_widget(state, dock_area)
self._set_dynamic_attributes(dock_area_obj, dock_obj.widget_name, dock_obj) dock = self._ipython_registry.get(dock_ref._gui_id)
self._set_dynamic_attributes(dock_area, dock_ref.widget_name, dock_ref)
# Keep track of rpc references on RPCBase object
self._update_rpc_references(dock_area, dock_ref.widget_name, dock_ref._gui_id)
# Keep track of exposed docks
self._exposed_widgets.append(dock_ref._gui_id)
# Add widgets # Add widgets
widget_states = [ widget_states = [
state state
for state in self._registry_state.values() for state in self._registry_state.values()
if state["config"].get("parent_id", "") == dock_obj._gui_id if state["config"].get("parent_id", "") == dock_ref._gui_id
] ]
for state in widget_states: for state in widget_states:
widget = self._add_widget(state, dock_obj) widget_ref = self._add_widget(state, dock)
self._set_dynamic_attributes(dock_obj, widget.widget_name, widget) self._set_dynamic_attributes(dock, widget_ref.widget_name, widget_ref)
self._set_dynamic_attributes(dock_area_obj.elements, widget.widget_name, widget) self._set_dynamic_attributes(
dock_area.elements, widget_ref.widget_name, widget_ref
)
# Keep track of rpc references on RPCBase object
self._update_rpc_references(
dock_area, f"elements.{widget_ref.widget_name}", widget_ref._gui_id
)
self._update_rpc_references(dock, widget_ref.widget_name, widget_ref._gui_id)
# Keep track of exposed widgets
self._exposed_widgets.append(widget_ref._gui_id)
def _add_widget(self, state: dict, parent: object) -> RPCReference: def _add_widget(self, state: dict, parent: object) -> RPCReference:
"""Add a widget to the namespace """Add a widget to the namespace
@ -519,9 +556,6 @@ class BECGuiClient(RPCBase):
self._ipython_registry[gui_id] = widget self._ipython_registry[gui_id] = widget
else: else:
widget = obj widget = obj
if widget_class == client.BECDockArea:
# Add elements to dynamic namespace
self._set_dynamic_attributes(widget, "elements", WidgetNameSpace())
obj = RPCReference(registry=self._ipython_registry, gui_id=gui_id) obj = RPCReference(registry=self._ipython_registry, gui_id=gui_id)
return obj return obj
@ -599,11 +633,17 @@ if __name__ == "__main__": # pragma: no cover
from bec_lib.client import BECClient from bec_lib.client import BECClient
from bec_lib.service_config import ServiceConfig from bec_lib.service_config import ServiceConfig
config = ServiceConfig() try:
bec_client = BECClient(config) config = ServiceConfig()
bec_client.start() bec_client = BECClient(config)
bec_client.start()
# Test the client_utils.py module # Test the client_utils.py module
gui = BECGuiClient() gui = BECGuiClient()
gui.start()
print(gui.window_list) gui.start(wait=True)
print(gui.window_list)
gui.new()
time.sleep(10)
finally:
gui.kill_server()

12
bec_widgets/cli/rpc/rpc_base.py
View File

@ -1,5 +1,6 @@
from __future__ import annotations from __future__ import annotations
import inspect
import threading import threading
import uuid import uuid
from functools import wraps from functools import wraps
@ -35,6 +36,13 @@ def rpc_call(func):
def wrapper(self, *args, **kwargs): def wrapper(self, *args, **kwargs):
# we could rely on a strict type check here, but this is more flexible # we could rely on a strict type check here, but this is more flexible
# moreover, it would anyway crash for objects... # moreover, it would anyway crash for objects...
caller_frame = inspect.currentframe().f_back
while caller_frame:
if "jedi" in caller_frame.f_globals:
# Jedi module is present, likely tab completion
return None # func(*args, **kwargs)
caller_frame = caller_frame.f_back
out = [] out = []
for arg in args: for arg in args:
if hasattr(arg, "name"): if hasattr(arg, "name"):
@ -120,7 +128,7 @@ class RPCBase:
self._msg_wait_event = threading.Event() self._msg_wait_event = threading.Event()
self._rpc_response = None self._rpc_response = None
super().__init__() super().__init__()
# print(f"RPCBase: {self._gui_id}") self._rpc_references: dict[str, str] = {}
def __repr__(self): def __repr__(self):
type_ = type(self) type_ = type(self)
@ -171,7 +179,7 @@ class RPCBase:
parameter={"args": args, "kwargs": kwargs, "gui_id": self._gui_id}, parameter={"args": args, "kwargs": kwargs, "gui_id": self._gui_id},
metadata={"request_id": request_id}, metadata={"request_id": request_id},
) )
print(f"running and rpc {method}")
# pylint: disable=protected-access # pylint: disable=protected-access
receiver = self._root._gui_id receiver = self._root._gui_id
if wait_for_rpc_response: if wait_for_rpc_response:

12
bec_widgets/cli/rpc/rpc_register.py
View File

@ -1,5 +1,6 @@
from __future__ import annotations from __future__ import annotations
from contextlib import contextmanager
from functools import wraps from functools import wraps
from threading import Lock from threading import Lock
from typing import TYPE_CHECKING, Callable from typing import TYPE_CHECKING, Callable
@ -17,6 +18,17 @@ if TYPE_CHECKING: # pragma: no cover
logger = bec_logger.logger logger = bec_logger.logger
@contextmanager
def rpc_register_broadcast(rpc_register):
"""
Context manager to broadcast updates to the RPCRegister whenever a new RPC object is added or removed.
"""
try:
yield rpc_register
finally:
rpc_register.broadcast()
def broadcast_update(func): def broadcast_update(func):
""" """
Decorator to broadcast updates to the RPCRegister whenever a new RPC object is added or removed. Decorator to broadcast updates to the RPCRegister whenever a new RPC object is added or removed.

46
bec_widgets/cli/server.py
View File

@ -15,8 +15,7 @@ from bec_lib.utils.import_utils import lazy_import
from qtpy.QtCore import Qt, QTimer from qtpy.QtCore import Qt, QTimer
from redis.exceptions import RedisError from redis.exceptions import RedisError
from bec_widgets.cli.rpc import rpc_register from bec_widgets.cli.rpc.rpc_register import RPCRegister, rpc_register_broadcast
from bec_widgets.cli.rpc.rpc_register import RPCRegister
from bec_widgets.qt_utils.error_popups import ErrorPopupUtility from bec_widgets.qt_utils.error_popups import ErrorPopupUtility
from bec_widgets.utils import BECDispatcher from bec_widgets.utils import BECDispatcher
from bec_widgets.utils.bec_connector import BECConnector from bec_widgets.utils.bec_connector import BECConnector
@ -70,8 +69,6 @@ class BECWidgetsCLIServer:
# register broadcast callback # register broadcast callback
self.rpc_register = RPCRegister() self.rpc_register = RPCRegister()
self.rpc_register.add_callback(self.broadcast_registry_update) self.rpc_register.add_callback(self.broadcast_registry_update)
self.gui = gui_class(parent=None, name=gui_class_id, gui_id=gui_class_id)
# self.rpc_register.add_rpc(self.gui)
self.dispatcher.connect_slot( self.dispatcher.connect_slot(
self.on_rpc_update, MessageEndpoints.gui_instructions(self.gui_id) self.on_rpc_update, MessageEndpoints.gui_instructions(self.gui_id)
@ -83,6 +80,8 @@ class BECWidgetsCLIServer:
self._heartbeat_timer.start(200) self._heartbeat_timer.start(200)
self.status = messages.BECStatus.RUNNING self.status = messages.BECStatus.RUNNING
with rpc_register_broadcast(self.rpc_register):
self.gui = gui_class(parent=None, name=gui_class_id, gui_id=gui_class_id)
logger.success(f"Server started with gui_id: {self.gui_id}") logger.success(f"Server started with gui_id: {self.gui_id}")
# Create initial object -> BECFigure or BECDockArea # Create initial object -> BECFigure or BECDockArea
@ -118,30 +117,31 @@ class BECWidgetsCLIServer:
return obj return obj
def run_rpc(self, obj, method, args, kwargs): def run_rpc(self, obj, method, args, kwargs):
logger.debug(f"Running RPC instruction: {method} with args: {args}, kwargs: {kwargs}") with rpc_register_broadcast(self.rpc_register):
method_obj = getattr(obj, method) logger.debug(f"Running RPC instruction: {method} with args: {args}, kwargs: {kwargs}")
# check if the method accepts args and kwargs method_obj = getattr(obj, method)
if not callable(method_obj): # check if the method accepts args and kwargs
if not args: if not callable(method_obj):
res = method_obj if not args:
res = method_obj
else:
setattr(obj, method, args[0])
res = None
else: else:
setattr(obj, method, args[0]) res = method_obj(*args, **kwargs)
res = None
else:
res = method_obj(*args, **kwargs)
if isinstance(res, list): if isinstance(res, list):
res = [self.serialize_object(obj) for obj in res] res = [self.serialize_object(obj) for obj in res]
elif isinstance(res, dict): elif isinstance(res, dict):
res = {key: self.serialize_object(val) for key, val in res.items()} res = {key: self.serialize_object(val) for key, val in res.items()}
else: else:
res = self.serialize_object(res) res = self.serialize_object(res)
return res return res
def serialize_object(self, obj): def serialize_object(self, obj):
if isinstance(obj, BECConnector): if isinstance(obj, BECConnector):
config = {} # obj.config.model_dump() config = obj.config.model_dump()
config["parent_id"] = obj.parent_id config["parent_id"] = obj.parent_id # add parent_id to config
return { return {
"gui_id": obj.gui_id, "gui_id": obj.gui_id,
"name": ( "name": (