feature integrate sa training and policies

AlejandroAvilesSerrano · AlejandroAvilesSerrano · commit 60711e320468 · 2025-06-05T09:52:33.000+02:00
diff --git a/nebula/core/situationalawareness/awareness/sareasoner.py b/nebula/core/situationalawareness/awareness/sareasoner.py
@@ -315,11 +315,15 @@ async def _initialize_sa_components(self):
                 await sacomp.init()
 
     def _load_minimal_requirement_config(self):
-        self._config["situational_awareness"]["sa_reasoner"]["sa_network"]["addr"] = self._addr
-        self._config["situational_awareness"]["sa_reasoner"]["sa_network"]["sar"] = self
-        self._config["situational_awareness"]["sa_reasoner"]["sa_network"]["strict_topology"] = self._config[
-            "situational_awareness"
-        ]["strict_topology"]
+        #self._config["situational_awareness"]["sa_reasoner"]["sa_network"]["addr"] = self._addr
+        #self._config["situational_awareness"]["sa_reasoner"]["sa_network"]["sar"] = self
+        self._config["situational_awareness"]["sa_reasoner"]["sa_network"]["strict_topology"] = self._config["situational_awareness"]["strict_topology"]
+        
+        # SA Reasoner instance for all SA Reasoner Components
+        sar_components: dict = self._config["situational_awareness"]["sa_reasoner"]["sar_components"]
+        for sar_comp in sar_components.keys():
+            self._config["situational_awareness"]["sa_reasoner"][sar_comp]["sar"] = self
+            self._config["situational_awareness"]["sa_reasoner"][sar_comp]["addr"] = self._addr
 
     async def _set_minimal_requirements(self):
         if self._sa_components:
diff --git a/nebula/core/situationalawareness/awareness/satraining/satraining.py b/nebula/core/situationalawareness/awareness/satraining/satraining.py
@@ -0,0 +1,82 @@
+import asyncio
+import logging
+from nebula.core.utils.locker import Locker
+from nebula.core.situationalawareness.awareness.satraining.trainingpolicy.trainingpolicy import factory_training_policy
+from nebula.core.situationalawareness.awareness.sareasoner import SAMComponent
+from nebula.addons.functions import print_msg_box
+from nebula.core.situationalawareness.awareness.sareasoner import SAReasoner, SAMComponent
+from nebula.core.eventmanager import EventManager
+    
+RESTRUCTURE_COOLDOWN = 5    
+    
+class SATraining(SAMComponent):
+    """
+    SATraining is a Situational Awareness (SA) component responsible for enhancing
+    the training process in Distributed Federated Learning (DFL) environments
+    by leveraging context-awareness and environmental knowledge.
+
+    This component dynamically instantiates a training policy based on the configuration,
+    allowing the system to adapt training strategies depending on the local topology,
+    node behavior, or environmental constraints.
+
+    Attributes:
+        _config (dict): Configuration dictionary containing parameters and references.
+        _sar (SAReasoner): Reference to the shared situational reasoner.
+        _trainning_policy: Instantiated training policy strategy.
+    """
+    
+    def __init__(self, config):
+        """
+        Initialize the SATraining component with a given configuration.
+
+        Args:
+            config (dict): Configuration dictionary containing:
+                - 'addr': Node address.
+                - 'verbose': Verbosity flag.
+                - 'sar': Reference to the SAReasoner instance.
+                - 'training_policy': Training policy name to be used.
+        """
+        print_msg_box(
+            msg=f"Starting Training SA\nTraining policy: {training_policy}",
+            indent=2,
+            title="Training SA module",
+        )
+        self._config = config
+        self._sar: SAReasoner = self._config["sar"]
+        tp_config = {}
+        tp_config["addr"] = self._config["addr"]
+        tp_config["verbose"] = self._config["verbose"]
+        training_policy = self._config["training_policy"]
+        self._trainning_policy = factory_training_policy(training_policy, tp_config)
+
+    @property
+    def sar(self):
+        """
+        Returns the current instance of the SAReasoner.
+        """
+        return self._sar
+
+    @property
+    def tp(self):
+        """
+        Returns the currently active training policy instance.
+        """
+        return self._trainning_policy    
+
+    async def init(self):
+        """
+        Initialize the training policy with the current known neighbors from the SAReasoner.
+        This setup enables the policy to make informed decisions based on local topology.
+        """
+        config = {}
+        config["nodes"] = set(await self.sar.get_nodes_known(neighbors_only=True)) 
+        await self.tp.init(config)
+
+    async def sa_component_actions(self):
+        """
+        Periodically called action of the SA component to evaluate the current scenario.
+        This invokes the evaluation logic defined in the training policy to adapt behavior.
+        """
+        logging.info("SA Trainng evaluating current scenario")
+        asyncio.create_task(self.tp.get_evaluation_results())
+
diff --git a/nebula/core/situationalawareness/awareness/satraining/trainingpolicy/bpstrainingpolicy.py b/nebula/core/situationalawareness/awareness/satraining/trainingpolicy/bpstrainingpolicy.py
@@ -0,0 +1,38 @@
+from nebula.core.situationalawareness.awareness.satraining.trainingpolicy.trainingpolicy import TrainingPolicy
+from nebula.core.situationalawareness.awareness.suggestionbuffer import SuggestionBuffer
+from nebula.core.situationalawareness.awareness.sautils.sacommand import SACommand, factory_sa_command, SACommandAction, SACommandPRIO
+from nebula.core.nebulaevents import RoundEndEvent
+
+class BPSTrainingPolicy(TrainingPolicy):
+    
+    def __init__(self, config=None):
+        pass
+    
+    async def init(self, config):
+        await self.register_sa_agent()    
+
+    async def get_evaluation_results(self):
+        sac = factory_sa_command(
+            "connectivity",
+            SACommandAction.MAINTAIN_CONNECTIONS,
+            self, 
+            "",
+            SACommandPRIO.LOW,
+            False,
+            None,
+            None
+        )
+        await self.suggest_action(sac)
+        await self.notify_all_suggestions_done(RoundEndEvent)
+    
+    async def get_agent(self) -> str:
+        return "SATraining_BPSTP"
+
+    async def register_sa_agent(self):
+        await SuggestionBuffer.get_instance().register_event_agents(RoundEndEvent, self)
+    
+    async def suggest_action(self, sac : SACommand):
+        await SuggestionBuffer.get_instance().register_suggestion(RoundEndEvent, self, sac)
+    
+    async def notify_all_suggestions_done(self, event_type):
+        await SuggestionBuffer.get_instance().notify_all_suggestions_done_for_agent(self, event_type)
diff --git a/nebula/core/situationalawareness/awareness/satraining/trainingpolicy/fastreboot.py b/nebula/core/situationalawareness/awareness/satraining/trainingpolicy/fastreboot.py
@@ -0,0 +1,139 @@
+import logging
+from nebula.core.utils.locker import Locker
+from nebula.core.eventmanager import EventManager
+from nebula.core.nebulaevents import AggregationEvent, UpdateNeighborEvent
+from nebula.core.situationalawareness.awareness.satraining.trainingpolicy.trainingpolicy import TrainingPolicy
+from nebula.core.situationalawareness.awareness.suggestionbuffer import SuggestionBuffer
+from nebula.core.situationalawareness.awareness.sautils.sacommand import SACommand, SACommandAction, SACommandPRIO, factory_sa_command
+
+VANILLA_LEARNING_RATE = 1e-3
+FR_LEARNING_RATE = 1e-3
+MAX_ROUNDS = 20
+DEFAULT_WEIGHT_MODIFIER = 3
+
+
+class FastReboot(TrainingPolicy):
+    def __init__(
+        self,
+        config
+    ):
+        logging.info("🌐  Initializing FastReboot")
+        self._max_rounds = MAX_ROUNDS                           # Max rounds to be applied FastReboot
+        self._weight_mod_value = DEFAULT_WEIGHT_MODIFIER
+        self._default_lr = VANILLA_LEARNING_RATE                # Stable value for learning rate
+        self._upgrade_lr = FR_LEARNING_RATE                     # Increased value for learning rate
+        self._current_lr = VANILLA_LEARNING_RATE
+        self._verbose = config["verbose"]
+        
+        self._learning_rate_lock = Locker(name="learning_rate_lock", async_lock=True)
+        self._weight_modifier = {}
+        self._weight_modifier_lock = Locker(name="weight_modifier_lock", async_lock=True)
+
+        self._fr_in_progress = False
+        
+    async def init(self, config):
+        #await EventManager.get_instance().subscribe_node_event(UpdateNeighborEvent)
+        #await EventManager.get_instance().subscribe_node_event(AggregationEvent)
+        pass
+
+    async def get_evaluation_results(self):
+        pass
+
+    def __str__(self):
+        return "FRTS"
+
+    async def _get_current_learning_rate(self):
+        await self._learning_rate_lock.acquire_async()
+        lr = self._current_lr
+        await self._learning_rate_lock.release_async()
+        return lr
+
+    async def discard_fastreboot_for(self, addr):
+        await self._weight_modifier_lock.acquire_async()
+        try:
+            del self._weight_modifier[addr]
+        except KeyError:
+            pass
+        await self._weight_modifier_lock.release_async()
+
+    async def _set_learning_rate(self, lr):
+        await self._learning_rate_lock.acquire_async()
+        self._current_lr = lr
+        await self._learning_rate_lock.release_async()
+
+    async def add_fastReboot_addr(self, addr):
+        await self._weight_modifier_lock.acquire_async()
+        if addr not in self._weight_modifier:
+            self._fr_in_progress = True
+            wm = self._weight_mod_value
+            logging.info(
+                f"📝 Registering | FastReboot registered for source {addr} | round application: {self._max_rounds} | multiplier value: {wm}"
+            )
+            self._weight_modifier[addr] = (wm, 1)
+            await self._set_learning_rate(self._upgrade_lr)
+            current_lr = await self._get_current_learning_rate()
+            #TODO modify learning rate suggestion await self.nm.update_learning_rate(current_lr)
+        await self._weight_modifier_lock.release_async()
+
+    async def _remove_weight_modifier(self, addr):
+        logging.info(f"📝 Removing | FastReboot removed for source {addr}")
+        del self._weight_modifier[addr]
+
+    async def _weight_modifiers_empty(self):
+        await self._weight_modifier_lock.acquire_async()
+        empty = False if self._weight_modifier else True
+        await self._weight_modifier_lock.release_async()
+        return empty
+
+    async def apply_weight_strategy(self, updates: dict):
+        if await self._weight_modifiers_empty():
+            if self._fr_in_progress:
+                await self._end_fastreboot()
+            return
+        logging.info("🔄  Applying FastReboot Strategy...")
+        for addr, update in updates.items():
+            weightmodifier, rounds = await self._get_weight_modifier(addr)
+            if weightmodifier != 1:
+                logging.info(
+                    f"📝 Appliying FastReboot strategy | addr: {addr} | multiplier value: {weightmodifier}, rounds applied: {rounds}"
+                )
+                model, weight = update
+                updates.update({addr: (model, weight * weightmodifier)})
+        await self._update_weight_modifiers()
+
+    async def _update_weight_modifiers(self):
+        await self._weight_modifier_lock.acquire_async()
+        if self._weight_modifier:
+            logging.info("🔄  Update | weights being updated")
+            remove_addrs = []
+            for addr, (weight, rounds) in self._weight_modifier.items():
+                new_weight = weight - 1 / (rounds**2)
+                rounds = rounds + 1
+                if new_weight > 1 and rounds <= self._max_rounds:
+                    self._weight_modifier[addr] = (new_weight, rounds)
+                else:
+                    remove_addrs.append(addr)
+            for a in remove_addrs:
+                await self._remove_weight_modifier(a)
+        await self._weight_modifier_lock.release_async()
+
+    async def _end_fastreboot(self):
+        await self._weight_modifier_lock.acquire_async()
+        if not self._weight_modifier and await self._is_lr_modified():
+            logging.info("🔄  Finishing | FastReboot is completed")
+            self._fr_in_progress = False
+            await self._set_learning_rate(self._default_lr)
+            #TODO modify learning rate suggestion await self.nm.update_learning_rate(self._default_lr)
+        await self._weight_modifier_lock.release_async()
+
+    async def _get_weight_modifier(self, addr):
+        await self._weight_modifier_lock.acquire_async()
+        wm = self._weight_modifier.get(addr, (1, 0))
+        await self._weight_modifier_lock.release_async()
+        return wm
+
+    async def _is_lr_modified(self):
+        await self._learning_rate_lock.acquire_async()
+        mod = self._current_lr == self._upgrade_lr
+        await self._learning_rate_lock.release_async()
+        return mod
diff --git a/nebula/core/situationalawareness/awareness/satraining/trainingpolicy/htstrainingpolicy.py b/nebula/core/situationalawareness/awareness/satraining/trainingpolicy/htstrainingpolicy.py
@@ -0,0 +1,66 @@
+from nebula.core.situationalawareness.awareness.satraining.trainingpolicy.trainingpolicy import TrainingPolicy
+from nebula.core.situationalawareness.awareness.satraining.trainingpolicy.trainingpolicy import factory_training_policy
+from nebula.core.situationalawareness.awareness.satraining.trainingpolicy.trainingpolicy import TrainingPolicy
+import logging
+
+# "Hybrid Training Strategy"    (HTS)
+class HTSTrainingPolicy(TrainingPolicy):
+    """
+    Implements a Hybrid Training Strategy (HTS) that combines multiple training policies 
+    (e.g., QDS, FRTS) to collaboratively decide on the evaluation and potential pruning 
+    of neighbors in a decentralized federated learning scenario.
+    
+    Attributes:
+        TRAINING_POLICY (set): Names of training policy classes to instantiate and manage.
+    """
+    
+    TRAINING_POLICY = {
+        "qds",
+        "frts",
+    }
+    
+    def __init__(self, config):
+        """
+        Initializes the HTS policy with the node's address and verbosity level.
+        It creates instances of each sub-policy listed in TRAINING_POLICY.
+
+        Args:
+            config (dict): Configuration dictionary with keys:
+                - 'addr': Node's address
+                - 'verbose': Enable verbose logging
+        """
+        self._addr = config["addr"]
+        self._verbose = config["verbose"]
+        self._training_policies : set[TrainingPolicy] = set()
+        self._training_policies.add([factory_training_policy(x, config) for x in self.TRAINING_POLICY])
+        
+    def __str__(self):
+        return "HTS"    
+        
+    @property
+    def tps(self):
+        return self._training_policies  
+
+    async def init(self, config):
+        for tp in self.tps:
+            await tp.init(config)    
+
+    async def update_neighbors(self, node, remove=False):
+        pass
+    
+    async def get_evaluation_results(self):
+        """
+        Asynchronously calls the `get_evaluation_results` of each policy,
+        and logs the nodes each policy would remove.
+        
+        Returns:
+            None (future version may merge all evaluations).
+        """
+        nodes_to_remove = dict()
+        for tp in self.tps:
+            nodes_to_remove[tp] = await tp.get_evaluation_results()
+        
+        for tp, nodes in nodes_to_remove.items():
+            logging.info(f"Training Policy: {tp}, nodes to remove: {nodes}")
+            
+        return None
diff --git a/nebula/core/situationalawareness/awareness/satraining/trainingpolicy/qdstrainingpolicy.py b/nebula/core/situationalawareness/awareness/satraining/trainingpolicy/qdstrainingpolicy.py
diff --git a/nebula/core/situationalawareness/awareness/satraining/trainingpolicy/trainingpolicy.py b/nebula/core/situationalawareness/awareness/satraining/trainingpolicy/trainingpolicy.py
