Optimization/messages (#32)

* messages optimized

* fix: solved logging error when the experiment is finished

* fix: websockets on monitor and layout

* make checl

Author: FerTV

docs/_prebuilt/developerguide.md

@@ -707,10 +707,7 @@ The new aggregator must inherit from the **Aggregator** class. You can use **Fed
                 logging.info(
                     f"🔄  include_model_in_buffer | Broadcasting MODELS_INCLUDED for round {self.engine.get_round()}"
                 )
-                message = self.cm.mm.generate_federation_message(
-                    nebula_pb2.FederationMessage.Action.FEDERATION_MODELS_INCLUDED,
-                    [self.engine.get_round()],
-                )
+                message = self.cm.create_message("federation", "federation_models_included", [str(arg) for arg in [self.engine.get_round()]])
                 await self.cm.send_message_to_neighbors(message)
 
             return

nebula/core/aggregation/aggregator.py

@@ -3,7 +3,6 @@
 from abc import ABC, abstractmethod
 from functools import partial
 
-from nebula.core.pb import nebula_pb2
 from nebula.core.utils.locker import Locker
 
 
@@ -195,9 +194,12 @@ async def include_model_in_buffer(self, model, weight, source=None, round=None,
             logging.info(
                 f"🔄  include_model_in_buffer | Broadcasting MODELS_INCLUDED for round {self.engine.get_round()}"
             )
-            message = self.cm.mm.generate_federation_message(
-                nebula_pb2.FederationMessage.Action.FEDERATION_MODELS_INCLUDED,
-                [self.engine.get_round()],
+            # message = self.cm.mm.generate_federation_message(
+            #     nebula_pb2.FederationMessage.Action.FEDERATION_MODELS_INCLUDED,
+            #     [self.engine.get_round()],
+            # )
+            message = self.cm.create_message(
+                "federation", "federation_models_included", [str(arg) for arg in [self.engine.get_round()]]
             )
             await self.cm.send_message_to_neighbors(message)
 

nebula/core/engine.py

@@ -8,9 +8,8 @@
 from nebula.addons.functions import print_msg_box
 from nebula.addons.reporter import Reporter
 from nebula.core.aggregation.aggregator import create_aggregator, create_target_aggregator
-from nebula.core.eventmanager import EventManager, event_handler
+from nebula.core.eventmanager import EventManager
 from nebula.core.network.communications import CommunicationsManager
-from nebula.core.pb import nebula_pb2
 from nebula.core.utils.locker import Locker
 
 logging.getLogger("requests").setLevel(logging.WARNING)
@@ -140,22 +139,18 @@ def __init__(
         self._reporter = Reporter(config=self.config, trainer=self.trainer, cm=self.cm)
 
         self._event_manager = EventManager(
-            default_callbacks=[
-                self._discovery_discover_callback,
-                self._control_alive_callback,
-                self._connection_connect_callback,
-                self._connection_disconnect_callback,
-                self._federation_ready_callback,
-                self._start_federation_callback,
-                self._federation_models_included_callback,
-            ]
+            # default_callbacks=[
+            #     self._discovery_discover_callback,
+            #     self._control_alive_callback,
+            #     self._connection_connect_callback,
+            #     self._connection_disconnect_callback,
+            #     # self._federation_ready_callback,
+            #     # self._start_federation_callback,
+            #     # self._federation_models_included_callback,
+            # ]
         )
 
-        # Register additional callbacks
-        self._event_manager.register_callback(
-            self._reputation_callback,
-            # ... add more callbacks here
-        )
+        self.register_message_events_callbacks()
 
     @property
     def cm(self):
@@ -207,7 +202,26 @@ def get_federation_setup_lock(self):
     def get_round_lock(self):
         return self.round_lock
 
-    @event_handler(nebula_pb2.DiscoveryMessage, nebula_pb2.DiscoveryMessage.Action.DISCOVER)
+    def register_message_events_callbacks(self):
+        me_dict = self.cm.get_messages_events()
+        message_events = [
+            (message_name, message_action)
+            for (message_name, message_actions) in me_dict.items()
+            for message_action in message_actions
+        ]
+        logging.info(f"{message_events}")
+        for event_type, action in message_events:
+            callback_name = f"_{event_type}_{action}_callback"
+            logging.info(f"Searching callback named: {callback_name}")
+            method = getattr(self, callback_name, None)
+
+            if callable(method):
+                self.event_manager.subscribe((event_type, action), method)
+
+    async def trigger_event(self, message_event):
+        logging.info(f"Publishing MessageEvent: {message_event.message_type}")
+        await self.event_manager.publish(message_event)
+
     async def _discovery_discover_callback(self, source, message):
         logging.info(
             f"🔍  handle_discovery_message | Trigger | Received discovery message from {source} (network propagation)"
@@ -231,7 +245,6 @@ async def _discovery_discover_callback(self, source, message):
                         f"🔍  Invalid geolocation received from {source}: latitude={message.latitude}, longitude={message.longitude}"
                     )
 
-    @event_handler(nebula_pb2.ControlMessage, nebula_pb2.ControlMessage.Action.ALIVE)
     async def _control_alive_callback(self, source, message):
         logging.info(f"🔧  handle_control_message | Trigger | Received alive message from {source}")
         current_connections = await self.cm.get_addrs_current_connections(myself=True)
@@ -243,38 +256,27 @@ async def _control_alive_callback(self, source, message):
         else:
             logging.error(f"❗️  Connection {source} not found in connections...")
 
-    @event_handler(nebula_pb2.ConnectionMessage, nebula_pb2.ConnectionMessage.Action.CONNECT)
     async def _connection_connect_callback(self, source, message):
         logging.info(f"🔗  handle_connection_message | Trigger | Received connection message from {source}")
         current_connections = await self.cm.get_addrs_current_connections(myself=True)
         if source not in current_connections:
             logging.info(f"🔗  handle_connection_message | Trigger | Connecting to {source}")
             await self.cm.connect(source, direct=True)
 
-    @event_handler(nebula_pb2.ConnectionMessage, nebula_pb2.ConnectionMessage.Action.DISCONNECT)
     async def _connection_disconnect_callback(self, source, message):
         logging.info(f"🔗  handle_connection_message | Trigger | Received disconnection message from {source}")
         await self.cm.disconnect(source, mutual_disconnection=False)
 
-    @event_handler(
-        nebula_pb2.FederationMessage,
-        nebula_pb2.FederationMessage.Action.FEDERATION_READY,
-    )
-    async def _federation_ready_callback(self, source, message):
+    async def _federation_federation_ready_callback(self, source, message):
         logging.info(f"📝  handle_federation_message | Trigger | Received ready federation message from {source}")
         if self.config.participant["device_args"]["start"]:
             logging.info(f"📝  handle_federation_message | Trigger | Adding ready connection {source}")
             await self.cm.add_ready_connection(source)
 
-    @event_handler(
-        nebula_pb2.FederationMessage,
-        nebula_pb2.FederationMessage.Action.FEDERATION_START,
-    )
-    async def _start_federation_callback(self, source, message):
+    async def _federation_federation_start_callback(self, source, message):
         logging.info(f"📝  handle_federation_message | Trigger | Received start federation message from {source}")
         await self.create_trainer_module()
 
-    @event_handler(nebula_pb2.FederationMessage, nebula_pb2.FederationMessage.Action.REPUTATION)
     async def _reputation_callback(self, source, message):
         malicious_nodes = message.arguments  # List of malicious nodes
         if self.with_reputation:
@@ -287,11 +289,7 @@ async def _reputation_callback(self, source, message):
                         malicious_nodes,
                     )
 
-    @event_handler(
-        nebula_pb2.FederationMessage,
-        nebula_pb2.FederationMessage.Action.FEDERATION_MODELS_INCLUDED,
-    )
-    async def _federation_models_included_callback(self, source, message):
+    async def _federation_federation_models_included_callback(self, source, message):
         logging.info(f"📝  handle_federation_message | Trigger | Received aggregation finished message from {source}")
         try:
             await self.cm.get_connections_lock().acquire_async()
@@ -346,7 +344,8 @@ async def deploy_federation(self):
                 while not await self.cm.check_federation_ready():
                     await asyncio.sleep(1)
                 logging.info("Sending FEDERATION_START to neighbors...")
-                message = self.cm.mm.generate_federation_message(nebula_pb2.FederationMessage.Action.FEDERATION_START)
+                # message = self.cm.mm.generate_federation_message(nebula_pb2.FederationMessage.Action.FEDERATION_START)
+                message = self.cm.create_message("federation", "federation_start")
                 await self.cm.send_message_to_neighbors(message)
                 await self.get_federation_ready_lock().release_async()
                 await self.create_trainer_module()
@@ -355,7 +354,8 @@ async def deploy_federation(self):
 
         else:
             logging.info("Sending FEDERATION_READY to neighbors...")
-            message = self.cm.mm.generate_federation_message(nebula_pb2.FederationMessage.Action.FEDERATION_READY)
+            # message = self.cm.mm.generate_federation_message(nebula_pb2.FederationMessage.Action.FEDERATION_READY)
+            message = self.cm.create_message("federation", "federation_ready")
             await self.cm.send_message_to_neighbors(message)
             logging.info("💤  Waiting until receiving the start signal from the start node")
 
@@ -454,6 +454,9 @@ async def _waiting_model_updates(self):
         else:
             logging.error("Aggregation finished with no parameters")
 
+    def learning_cycle_finished(self):
+        return not (self.round < self.total_rounds)
+
     async def _learning_cycle(self):
         while self.round is not None and self.round < self.total_rounds:
             print_msg_box(
@@ -488,8 +491,6 @@ async def _learning_cycle(self):
         # End of the learning cycle
         self.trainer.on_learning_cycle_end()
         await self.trainer.test()
-        self.round = None
-        self.total_rounds = None
         print_msg_box(
             msg="Federated Learning process has been completed.",
             indent=2,
@@ -507,8 +508,7 @@ async def _learning_cycle(self):
         while not self.cm.check_finished_experiment():
             await asyncio.sleep(1)
 
-        # Enable loggin info
-        logging.getLogger().disabled = True
+        await asyncio.sleep(5)
 
         # Kill itself
         if self.config.participant["scenario_args"]["deployment"] == "docker":
@@ -564,9 +564,10 @@ def reputation_calculation(self, aggregated_models_weights):
 
     async def send_reputation(self, malicious_nodes):
         logging.info(f"Sending REPUTATION to the rest of the topology: {malicious_nodes}")
-        message = self.cm.mm.generate_federation_message(
-            nebula_pb2.FederationMessage.Action.REPUTATION, malicious_nodes
-        )
+        # message = self.cm.mm.generate_federation_message(
+        #     nebula_pb2.FederationMessage.Action.REPUTATION, malicious_nodes
+        # )
+        message = self.cm.create_message("federation", "reputation", arguments=[str(arg) for arg in (malicious_nodes)])
         await self.cm.send_message_to_neighbors(message)
 
 

nebula/core/eventmanager.py

@@ -4,6 +4,8 @@
 from collections import defaultdict
 from functools import wraps
 
+from nebula.core.network.messages import MessageEvent
+
 
 def event_handler(message_type, action):
     """Decorator for registering an event handler."""
@@ -33,6 +35,30 @@ class EventManager:
     def __init__(self, default_callbacks=None):
         self._event_callbacks = defaultdict(list)
         self._register_default_callbacks(default_callbacks or [])
+        self._subscribers: dict[tuple[str, str], list] = {}
+
+    def subscribe(self, event_type: tuple[str, str], callback: callable):
+        """Register a callback for a specific event type."""
+        if event_type not in self._subscribers:
+            self._subscribers[event_type] = []
+        self._subscribers[event_type].append(callback)
+        logging.info(f"EventManager | Subscribed callback for event: {event_type}")
+
+    async def publish(self, message_event: MessageEvent):
+        """Trigger all callbacks registered for a specific event type."""
+        event_type = message_event.message_type
+        if event_type not in self._subscribers:
+            logging.error(f"EventManager | No subscribers for event: {event_type}")
+            return
+
+        for callback in self._subscribers[event_type]:
+            try:
+                logging.info(
+                    f"EventManager | Triggering callback for event: {event_type}, from source: {message_event.source}"
+                )
+                await callback(message_event.source, message_event.message)
+            except Exception as e:
+                logging.exception(f"EventManager | Error in callback for event {event_type}: {e}")
 
     def _register_default_callbacks(self, default_callbacks):
         """Registers default callbacks for events."""

nebula/core/network/actions.py

@@ -0,0 +1,71 @@
+from enum import Enum
+
+from nebula.core.pb import nebula_pb2
+
+
+class ConnectionAction(Enum):
+    CONNECT = nebula_pb2.ConnectionMessage.Action.CONNECT
+    DISCONNECT = nebula_pb2.ConnectionMessage.Action.DISCONNECT
+
+
+class FederationAction(Enum):
+    FEDERATION_START = nebula_pb2.FederationMessage.Action.FEDERATION_START
+    REPUTATION = nebula_pb2.FederationMessage.Action.REPUTATION
+    FEDERATION_MODELS_INCLUDED = nebula_pb2.FederationMessage.Action.FEDERATION_MODELS_INCLUDED
+    FEDERATION_READY = nebula_pb2.FederationMessage.Action.FEDERATION_READY
+
+
+class DiscoveryAction(Enum):
+    DISCOVER = nebula_pb2.DiscoveryMessage.Action.DISCOVER
+    REGISTER = nebula_pb2.DiscoveryMessage.Action.REGISTER
+    DEREGISTER = nebula_pb2.DiscoveryMessage.Action.DEREGISTER
+
+
+class ControlAction(Enum):
+    ALIVE = nebula_pb2.ControlMessage.Action.ALIVE
+    OVERHEAD = nebula_pb2.ControlMessage.Action.OVERHEAD
+    MOBILITY = nebula_pb2.ControlMessage.Action.MOBILITY
+    RECOVERY = nebula_pb2.ControlMessage.Action.RECOVERY
+    WEAK_LINK = nebula_pb2.ControlMessage.Action.WEAK_LINK
+
+
+ACTION_CLASSES = {
+    "connection": ConnectionAction,
+    "federation": FederationAction,
+    "discovery": DiscoveryAction,
+    "control": ControlAction,
+}
+
+
+def get_action_name_from_value(message_type: str, action_value: int) -> str:
+    # Obtener el Enum correspondiente al tipo de mensaje
+    enum_class = ACTION_CLASSES.get(message_type)
+    if not enum_class:
+        raise ValueError(f"Unknown message type: {message_type}")
+
+    # Buscar el nombre de la acción a partir del valor
+    for action in enum_class:
+        if action.value == action_value:
+            return action.name.lower()  # Convertimos a lowercase para mantener el formato "late_connect"
+
+    raise ValueError(f"Unknown action value {action_value} for message type {message_type}")
+
+
+def get_actions_names(message_type: str):
+    message_actions = ACTION_CLASSES.get(message_type)
+    if not message_actions:
+        raise ValueError(f"Invalid message type: {message_type}")
+
+    return [action.name.lower() for action in message_actions]
+
+
+def factory_message_action(message_type: str, action: str):
+    message_actions = ACTION_CLASSES.get(message_type)
+
+    if message_actions:
+        normalized_action = action.upper()
+        enum_action = message_actions[normalized_action]
+        # logging.info(f"Message action: {enum_action}, value: {enum_action.value}")
+        return enum_action.value
+    else:
+        return None