mirror of
https://github.com/bec-project/bec_widgets.git
synced 2026-03-05 00:12:49 +01:00
278 lines
10 KiB
Python
278 lines
10 KiB
Python
from __future__ import annotations
|
|
|
|
import functools
|
|
import traceback
|
|
import types
|
|
from contextlib import contextmanager
|
|
from typing import TYPE_CHECKING, Callable, TypeVar
|
|
|
|
from bec_lib.client import BECClient
|
|
from bec_lib.endpoints import MessageEndpoints
|
|
from bec_lib.logger import bec_logger
|
|
from bec_lib.utils.import_utils import lazy_import
|
|
from qtpy.QtCore import QTimer
|
|
from qtpy.QtWidgets import QApplication
|
|
from redis.exceptions import RedisError
|
|
|
|
from bec_widgets.cli.rpc.rpc_register import RPCRegister
|
|
from bec_widgets.utils import BECDispatcher
|
|
from bec_widgets.utils.bec_connector import BECConnector
|
|
from bec_widgets.utils.error_popups import ErrorPopupUtility
|
|
from bec_widgets.widgets.containers.main_window.main_window import BECMainWindow
|
|
|
|
if TYPE_CHECKING: # pragma: no cover
|
|
from bec_lib import messages
|
|
from qtpy.QtCore import QObject
|
|
else:
|
|
messages = lazy_import("bec_lib.messages")
|
|
logger = bec_logger.logger
|
|
|
|
|
|
T = TypeVar("T")
|
|
|
|
|
|
@contextmanager
|
|
def rpc_exception_hook(err_func):
|
|
"""This context replaces the popup message box for error display with a specific hook"""
|
|
# get error popup utility singleton
|
|
popup = ErrorPopupUtility()
|
|
# save current setting
|
|
old_exception_hook = popup.custom_exception_hook
|
|
|
|
# install err_func, if it is a callable
|
|
# IMPORTANT, Keep self here, because this method is overwriting the custom_exception_hook
|
|
# of the ErrorPopupUtility (popup instance) class.
|
|
def custom_exception_hook(self, exc_type, value, tb, **kwargs):
|
|
err_func({"error": popup.get_error_message(exc_type, value, tb)})
|
|
|
|
popup.custom_exception_hook = types.MethodType(custom_exception_hook, popup)
|
|
|
|
try:
|
|
yield popup
|
|
finally:
|
|
# restore state of error popup utility singleton
|
|
popup.custom_exception_hook = old_exception_hook
|
|
|
|
|
|
class RPCServer:
|
|
|
|
client: BECClient
|
|
|
|
def __init__(
|
|
self,
|
|
gui_id: str,
|
|
dispatcher: BECDispatcher | None = None,
|
|
client: BECClient | None = None,
|
|
config=None,
|
|
gui_class_id: str = "bec",
|
|
) -> None:
|
|
self.status = messages.BECStatus.BUSY
|
|
self.dispatcher = BECDispatcher(config=config) if dispatcher is None else dispatcher
|
|
self.client = self.dispatcher.client if client is None else client
|
|
self.client.start()
|
|
self.gui_id = gui_id
|
|
# register broadcast callback
|
|
self.rpc_register = RPCRegister()
|
|
self.rpc_register.add_callback(self.broadcast_registry_update)
|
|
|
|
self.dispatcher.connect_slot(
|
|
self.on_rpc_update, MessageEndpoints.gui_instructions(self.gui_id)
|
|
)
|
|
|
|
# Setup QTimer for heartbeat
|
|
self._heartbeat_timer = QTimer()
|
|
self._heartbeat_timer.timeout.connect(self.emit_heartbeat)
|
|
self._heartbeat_timer.start(200)
|
|
self._registry_update_callbacks = []
|
|
self._broadcasted_data = {}
|
|
|
|
self.status = messages.BECStatus.RUNNING
|
|
logger.success(f"Server started with gui_id: {self.gui_id}")
|
|
|
|
def on_rpc_update(self, msg: dict, metadata: dict):
|
|
request_id = metadata.get("request_id")
|
|
if request_id is None:
|
|
logger.error("Received RPC instruction without request_id")
|
|
return
|
|
logger.debug(f"Received RPC instruction: {msg}, metadata: {metadata}")
|
|
with rpc_exception_hook(functools.partial(self.send_response, request_id, False)):
|
|
try:
|
|
obj = self.get_object_from_config(msg["parameter"])
|
|
method = msg["action"]
|
|
args = msg["parameter"].get("args", [])
|
|
kwargs = msg["parameter"].get("kwargs", {})
|
|
res = self.run_rpc(obj, method, args, kwargs)
|
|
except Exception:
|
|
content = traceback.format_exc()
|
|
logger.error(f"Error while executing RPC instruction: {content}")
|
|
self.send_response(request_id, False, {"error": content})
|
|
else:
|
|
logger.debug(f"RPC instruction executed successfully: {res}")
|
|
self.send_response(request_id, True, {"result": res})
|
|
|
|
def send_response(self, request_id: str, accepted: bool, msg: dict):
|
|
self.client.connector.set_and_publish(
|
|
MessageEndpoints.gui_instruction_response(request_id),
|
|
messages.RequestResponseMessage(accepted=accepted, message=msg),
|
|
expire=60,
|
|
)
|
|
|
|
def get_object_from_config(self, config: dict):
|
|
gui_id = config.get("gui_id")
|
|
obj = self.rpc_register.get_rpc_by_id(gui_id)
|
|
if obj is None:
|
|
raise ValueError(f"Object with gui_id {gui_id} not found")
|
|
return obj
|
|
|
|
def run_rpc(self, obj, method, args, kwargs):
|
|
# Run with rpc registry broadcast, but only once
|
|
with RPCRegister.delayed_broadcast():
|
|
logger.debug(f"Running RPC instruction: {method} with args: {args}, kwargs: {kwargs}")
|
|
method_obj = getattr(obj, method)
|
|
# check if the method accepts args and kwargs
|
|
if not callable(method_obj):
|
|
if not args:
|
|
res = method_obj
|
|
else:
|
|
setattr(obj, method, args[0])
|
|
res = None
|
|
else:
|
|
res = method_obj(*args, **kwargs)
|
|
|
|
if isinstance(res, list):
|
|
res = [self.serialize_object(obj) for obj in res]
|
|
elif isinstance(res, dict):
|
|
res = {key: self.serialize_object(val) for key, val in res.items()}
|
|
else:
|
|
res = self.serialize_object(res)
|
|
return res
|
|
|
|
def serialize_object(self, obj: T) -> None | dict | T:
|
|
"""
|
|
Serialize all BECConnector objects.
|
|
|
|
Args:
|
|
obj: The object to be serialized.
|
|
|
|
Returns:
|
|
None | dict | T: The serialized object or None if the object is not a BECConnector.
|
|
"""
|
|
if not isinstance(obj, BECConnector):
|
|
return obj
|
|
# Respect RPC = False
|
|
if getattr(obj, "RPC", True) is False:
|
|
return None
|
|
return self._serialize_bec_connector(obj, wait=True)
|
|
|
|
def emit_heartbeat(self) -> None:
|
|
"""
|
|
Emit a heartbeat message to the GUI server.
|
|
This method is called periodically to indicate that the server is still running.
|
|
"""
|
|
logger.trace(f"Emitting heartbeat for {self.gui_id}")
|
|
try:
|
|
self.client.connector.set(
|
|
MessageEndpoints.gui_heartbeat(self.gui_id),
|
|
messages.StatusMessage(name=self.gui_id, status=self.status, info={}),
|
|
expire=10,
|
|
)
|
|
except RedisError as exc:
|
|
logger.error(f"Error while emitting heartbeat: {exc}")
|
|
|
|
def broadcast_registry_update(self, connections: dict) -> None:
|
|
"""
|
|
Broadcast the registry update to all the callbacks.
|
|
This method is called whenever the registry is updated.
|
|
"""
|
|
data = {}
|
|
for key, val in connections.items():
|
|
if not isinstance(val, BECConnector):
|
|
continue
|
|
if not getattr(val, "RPC", True):
|
|
continue
|
|
data[key] = self._serialize_bec_connector(val)
|
|
if self._broadcasted_data == data:
|
|
return
|
|
self._broadcasted_data = data
|
|
|
|
logger.debug(f"Broadcasting registry update: {data} for {self.gui_id}")
|
|
self.client.connector.xadd(
|
|
MessageEndpoints.gui_registry_state(self.gui_id),
|
|
msg_dict={"data": messages.GUIRegistryStateMessage(state=data)},
|
|
max_size=1,
|
|
)
|
|
|
|
def _serialize_bec_connector(self, connector: BECConnector, wait=False) -> dict:
|
|
"""
|
|
Create the serialization dict for a single BECConnector.
|
|
|
|
Args:
|
|
connector (BECConnector): The BECConnector to serialize.
|
|
wait (bool): If True, wait until the object is registered in the RPC register.
|
|
|
|
Returns:
|
|
dict: The serialized BECConnector object.
|
|
"""
|
|
|
|
config_dict = connector.config.model_dump()
|
|
config_dict["parent_id"] = getattr(connector, "parent_id", None)
|
|
|
|
try:
|
|
parent = connector.parent()
|
|
if isinstance(parent, BECMainWindow):
|
|
container_proxy = parent.gui_id
|
|
else:
|
|
container_proxy = None
|
|
except Exception:
|
|
container_proxy = None
|
|
|
|
if wait:
|
|
while not self.rpc_register.object_is_registered(connector):
|
|
QApplication.processEvents()
|
|
|
|
widget_class = getattr(connector, "rpc_widget_class", None)
|
|
if not widget_class:
|
|
widget_class = connector.__class__.__name__
|
|
|
|
return {
|
|
"gui_id": connector.gui_id,
|
|
"object_name": connector.object_name or connector.__class__.__name__,
|
|
"widget_class": widget_class,
|
|
"config": config_dict,
|
|
"container_proxy": container_proxy,
|
|
"__rpc__": True,
|
|
}
|
|
|
|
@staticmethod
|
|
def _get_becwidget_ancestor(widget: QObject) -> BECConnector | None:
|
|
"""
|
|
Traverse up the parent chain to find the nearest BECConnector.
|
|
Returns None if none is found.
|
|
"""
|
|
|
|
parent = widget.parent()
|
|
while parent is not None:
|
|
if isinstance(parent, BECConnector):
|
|
return parent
|
|
parent = parent.parent()
|
|
return None
|
|
|
|
# Suppose clients register callbacks to receive updates
|
|
def add_registry_update_callback(self, cb: Callable) -> None:
|
|
"""
|
|
Add a callback to be called whenever the registry is updated.
|
|
The specified callback is called whenever the registry is updated.
|
|
|
|
Args:
|
|
cb (Callable): The callback to be added. It should accept a dictionary of all the
|
|
registered RPC objects as an argument.
|
|
"""
|
|
self._registry_update_callbacks.append(cb)
|
|
|
|
def shutdown(self): # TODO not sure if needed when cleanup is done at level of BECConnector
|
|
self.status = messages.BECStatus.IDLE
|
|
self._heartbeat_timer.stop()
|
|
self.emit_heartbeat()
|
|
logger.info("Succeded in shutting down CLI server")
|
|
self.client.shutdown()
|