robosim.plugins_manager

  1"""
  2plugins_manager.py A list of plugins and the orchestration of events between them.
  3Import level: 4
  4"""
  5from __future__ import annotations
  6from queue import Empty
  7from microecs import EntityId, World, QueryResult
  8from microspec import Protocol
  9
 10from robosim.utils import logger
 11from robosim.network import ObjToSimChannel
 12from robosim.components import HasChannel
 13from robosim.plugin import Plugin, Event, Message, ResponseKind
 14from robosim.traits import Serializable
 15
 16class PluginsManager(list, Serializable):
 17    """
 18    A small wrapper on top of all loaded plugins. The order of plugins determine priority in case >=2 target the same
 19    behavior, e.g. manual move plugin vs automatic trajectory plugin. Also routes tcp messages to&from each plugin.
 20    Order of calling:
 21    [start tick] -> [tcp recv] -> io_handler -> _update_ecs_data_from_events ->
 22                 -> [physics system] -> response_handler -> [tcp send] -> [end tick]
 23
 24    """
 25
 26    REGISTERED_PLUGINS: dict[str, type[Plugin]] = {}
 27
 28    def __init__(self, plugin_names: list[str]):
 29        self.plugins: list[Plugin] = [PluginsManager.REGISTERED_PLUGINS[name]() for name in plugin_names]
 30        list.__init__(self, self.plugins)
 31        Serializable.__init__(self)
 32        self.plugin_names = plugin_names
 33        self.all_endpoints = [e.name for p in self.plugins for e in p.endpoints]
 34        assert len(set(self.all_endpoints)) == len(self.all_endpoints), f"Duplicates: {self.all_endpoints}"
 35        self._endpoint_name_to_plugin = {e.name: plugin for plugin in self.plugins for e in plugin.endpoints}
 36        self._protocol = Protocol(endpoints={e.name: e for p in self.plugins for e in p.endpoints}) # formal validation
 37        self._events_this_tick: dict[EntityId, list[Event]] = {}
 38
 39    def on_after_physics(self, world: World):
 40        """Calls all the plugins after the physics loop. All events must be responded to after this phase"""
 41        [plugin.on_after_physics(world) for plugin in self.plugins]
 42
 43    def on_before_physics(self, world: World):
 44        """Calls all the plugins before the physics loop"""
 45        [plugin.on_before_physics(world) for plugin in self.plugins]
 46
 47    def io_handler(self, world: World):
 48        """Called at the beginning of each render tick. Sets events_this_tick {eid -> [events]} and modifies ECS data"""
 49        entities = world.query(HasChannel)
 50        messages = self._route_messages_to_plugins(entities)
 51
 52        self._events_this_tick.clear()
 53        n_events = 0
 54        for plugin in self.plugins:
 55            # first, we go through all the messages of this plugin (if any) as route_messages_to_plugins gave us
 56            for message in messages.get(plugin, []):
 57                plg_response = plugin.on_message_receive(world, entity_id=message.target, message=message)
 58                # for physics responses, the even is further relayed to the ECS and processed by physics system.
 59                if plg_response.kind == ResponseKind.Physics:
 60                    event_data = {"payload": plg_response.payload}
 61                else:
 62                    event_data = {"response": plg_response.payload}
 63                event = Event(source=plugin, target=message.target, message=message, **event_data)
 64                self._events_this_tick.setdefault(message.target, []).append(event)
 65                n_events += 1
 66
 67            for event in plugin.on_tick(world):
 68                self._events_this_tick.setdefault(event.target, []).append(event)
 69                n_events += 1
 70
 71        logger.log_every_s(f"Processing {n_events} across {len(self._events_this_tick)} ents this tick", "TRACE", True)
 72        self._update_ecs_data_from_events(world)
 73
 74    def responses_handler(self, world: World):
 75        """Called at the end of each render tick to ensure all the received messages have their responses"""
 76        for entity_events in self._events_this_tick.values():
 77            for event in entity_events:
 78                # events w/o a message are generated via on_tick and they are internal (e.g. wind or default moves)
 79                if event.message is None:
 80                    continue
 81                # respond to this unresponded-yet event (e.g. from on_message_receive)
 82                if event.response is None:
 83                    plg_response = event.source.on_message_response(world, event)
 84                    assert plg_response.kind in (ResponseKind.Ok, ResponseKind.Error), plg_response
 85                    assert plg_response.payload is not None, plg_response
 86                    event.response = plg_response.payload
 87                event.message.channel.sim2tcp.put(event.response)
 88
 89    def to_dict(self) -> dict:
 90        plugin_type_to_name = {plg_type: plg_name for plg_name, plg_type in PluginsManager.REGISTERED_PLUGINS.items()}
 91        return {plugin_type_to_name[type(plugin)]: plugin.to_dict() for plugin in self.plugins}
 92
 93    @staticmethod
 94    def from_dict(state: dict) -> PluginsManager:
 95        plugin_names = []
 96        for plugin_name in state.keys():
 97            if plugin_name not in PluginsManager.REGISTERED_PLUGINS:
 98                logger.warning(f"Plugin: '{plugin_name}' not in {list(PluginsManager.REGISTERED_PLUGINS)}. Skipping.")
 99                continue
100            plugin_names.append(plugin_name)
101
102        res = PluginsManager(plugin_names=plugin_names)
103        for plugin_name, plugin in zip(res.plugin_names, res.plugins):
104            plugin.load_state_dict(state[plugin_name])
105        return res
106
107    def _route_messages_to_plugins(self, entities: QueryResult) -> dict[Plugin, list[Message]]:
108        res: dict[Plugin, list[Message]] = {}
109
110        channel: ObjToSimChannel
111        for eid, channel in zip(map(lambda x: int(x), entities.entity_ids), map(lambda x: x.item(), entities.channel)):
112            try:
113                msg: dict = channel.tcp2sim.get_nowait()
114            except Empty:
115                continue
116
117            cmd = msg.pop("cmd", None)
118            if (plugin := self._endpoint_name_to_plugin.get(cmd)) is not None:
119                # validate the payload using microspec formal validation.
120                if (err := self._protocol.validate_endpoint(endpoint=cmd, data=msg)) is not None:
121                    logger.error(err)
122                    channel.sim2tcp.put({"error": err.error})
123                else: # the plugin message is validated, so no syntax sanitization is needed inside the plugin
124                    res.setdefault(plugin, []).append(Message(target=eid, cmd=cmd, payload=msg, channel=channel))
125            else:
126                logger.error(err := f"Unknown message: {msg} ({cmd=}) from channel: {channel.index} (entity id: {eid})")
127                channel.sim2tcp.put({"error": err})
128        return res
129
130    def _update_ecs_data_from_events(self, world: World):
131        for entity_id, entity_events in self._events_this_tick.items():
132            for event in reversed(entity_events):
133                if event.payload is not None: # the first event targeting this entity wins this tick
134                    entity = world.get_entity(entity_id)
135                    for key, value in event.payload.items():
136                        setattr(entity, key, value) # set the data as-is from the event to ECS data
137                    event.applied = True
138                    break
139
140    def __repr__(self):
141        plugins = "".join(f"\n  - {type(s).__name__}. Endpoints: {s.endpoints}" for s in self.plugins)
142        return f"[PluginsManager]\n- Endpoints: {len(self.all_endpoints)}\n- Plugins ({len(self.plugins)}):{plugins}"