clean code and fix minor issues

Author: enriquetomasmb

nebula/addons/topologymanager.py

@@ -250,23 +250,6 @@ def generate_ring_topology(self, increase_convergence=False):
         """
         self.__ring_topology(increase_convergence=increase_convergence)
 
-    def generate_custom_topology(self, topology):
-        """
-        Sets the network topology to a custom topology provided by the user.
-
-        This method allows for the manual configuration of the network topology by directly assigning
-        the `topology` argument to the internal `self.topology` attribute.
-
-        Args:
-            topology (numpy.ndarray): A 2D array representing the custom network topology.
-                                      The array should have dimensions (n_nodes, n_nodes) where `n_nodes`
-                                      is the number of nodes in the network.
-
-        Returns:
-            None: The method modifies the internal `self.topology` to the provided custom topology.
-        """
-        self.topology = topology
-
     def generate_random_topology(self, probability):
         """
         Generates a random topology using Erdos-Renyi model with given probability.
@@ -281,62 +264,6 @@ def generate_random_topology(self, probability):
         self.topology = nx.to_numpy_array(random_graph, dtype=np.float32)
         np.fill_diagonal(self.topology, 0)  # No self-loops
 
-    def get_matrix_adjacency_from_neighbors(self, neighbors):
-        """
-        Generates an adjacency matrix from a list of neighbors.
-
-        This method constructs an adjacency matrix for the network based on the provided list of neighbors
-        for each node. A 1 in the matrix at position (i, j) indicates that node i is a neighbor of node j,
-        while a 0 indicates no connection.
-
-        Args:
-            neighbors (list of lists): A list of lists where each sublist contains the indices of the neighbors
-                                       for the corresponding node. The length of the outer list should be equal
-                                       to the number of nodes in the network (`self.n_nodes`).
-
-        Returns:
-            numpy.ndarray: A 2D adjacency matrix of shape (n_nodes, n_nodes), where n_nodes is the total number
-                           of nodes in the network. The matrix contains 1s where there is a connection and 0s
-                           where there is no connection.
-        """
-        matrix_adjacency = np.zeros((self.n_nodes, self.n_nodes), dtype=np.float32)
-        for i in range(self.n_nodes):
-            for j in range(self.n_nodes):
-                if i in neighbors[j]:
-                    matrix_adjacency[i, j] = 1
-        return matrix_adjacency
-
-    def get_topology(self):
-        """
-        Returns the network topology.
-
-        This method retrieves the current topology of the network. The behavior of the method depends on whether
-        the network is symmetric or asymmetric. For both cases in this implementation, it simply returns the
-        `self.topology`.
-
-        Returns:
-            numpy.ndarray: The current topology of the network as a 2D numpy array. The topology represents the
-                           connectivity between nodes, where a value of 1 indicates a connection and 0 indicates
-                           no connection between the nodes.
-        """
-        if self.b_symmetric:
-            return self.topology
-        else:
-            return self.topology
-
-    def get_nodes(self):
-        """
-        Returns the nodes in the network.
-
-        This method retrieves the current list of nodes in the network. Each node is represented by an array of
-        three values (such as coordinates or identifiers) in the `self.nodes` attribute.
-
-        Returns:
-            numpy.ndarray: A 2D numpy array representing the nodes in the network. Each row represents a node,
-                           and the columns may represent different properties (e.g., position, identifier, etc.).
-        """
-        return self.nodes
-
     @staticmethod
     def get_coordinates(random_geo=True):
         """
@@ -393,20 +320,6 @@ def update_nodes(self, config_participants):
         """
         self.nodes = config_participants
 
-    def get_node(self, node_idx):
-        """
-        Retrieves the node information based on the given index.
-
-        This method returns the details of a specific node from the `nodes` attribute using its index.
-
-        Parameters:
-            node_idx (int): The index of the node to retrieve from the `nodes` list.
-
-        Returns:
-            numpy.ndarray: A tuple or array containing the node's information at the given index.
-        """
-        return self.nodes[node_idx]
-
     def get_neighbors_string(self, node_idx):
         """
         Retrieves the neighbors of a given node as a string representation.

nebula/controller.py

@@ -90,6 +90,7 @@ async def get_least_memory_gpu():
     gpu_with_least_memory_index = None
 
     if importlib.util.find_spec("pynvml") is not None:
+        max_memory_used_percent = 50
         try:
             import pynvml
 

nebula/core/aggregation/aggregator.py

@@ -44,7 +44,6 @@ def __init__(self, config=None, engine=None):
         logging.info(f"[{self.__class__.__name__}] Starting Aggregator")
         self._federation_nodes = set()
         self._pending_models_to_aggregate = {}
-        self._pending_models_to_aggregate_lock = Locker(name="pending_models_to_aggregate_lock", async_lock=True)
         self._aggregation_done_lock = Locker(name="aggregation_done_lock", async_lock=True)
         self._aggregation_waiting_skip = asyncio.Event()
 
@@ -92,7 +91,7 @@ async def get_aggregation(self):
             await self.us.notify_if_all_updates_received()
             lock_task = asyncio.create_task(self._aggregation_done_lock.acquire_async(timeout=timeout))
             skip_task = asyncio.create_task(self._aggregation_waiting_skip.wait())
-            done, pending = await asyncio.wait(
+            done, _ = await asyncio.wait(
                 [lock_task, skip_task],
                 return_when=asyncio.FIRST_COMPLETED,
             )
@@ -131,13 +130,10 @@ async def get_aggregation(self):
         return aggregated_result
 
     def print_model_size(self, model):
-        total_params = 0
         total_memory = 0
 
         for _, param in model.items():
             num_params = param.numel()
-            total_params += num_params
-
             memory_usage = param.element_size() * num_params
             total_memory += memory_usage
 

nebula/core/aggregation/dualhistagg.py

@@ -376,7 +376,7 @@ async def update_neighbors(self, removed_neighbor_addr, neighbors, remove=False)
         updt_nei_event = UpdateNeighborEvent(removed_neighbor_addr, remove)
         asyncio.create_task(EventManager.get_instance().publish_node_event(updt_nei_event))
 
-    async def broadcast_models_include(self, age: AggregationEvent):
+    async def broadcast_models_include(self, aggregation_event: AggregationEvent):
         logging.info(f"🔄  Broadcasting MODELS_INCLUDED for round {self.get_round()}")
         message = self.cm.create_message(
             "federation", "federation_models_included", [str(arg) for arg in [self.get_round()]]
@@ -530,13 +530,6 @@ async def _waiting_model_updates(self):
         else:
             logging.error("Aggregation finished with no parameters")
 
-    def print_round_information(self):
-        print_msg_box(
-            msg=f"Round {self.round} of {self.total_rounds} started.",
-            indent=2,
-            title="Round information",
-        )
-
     def learning_cycle_finished(self):
         return not (self.round < self.total_rounds)
 
@@ -707,8 +700,8 @@ def __init__(
 
     async def _extended_learning_cycle(self):
         # Define the functionality of the aggregator node
-        await self.trainer.train()
         await self.trainer.test()
+        await self.trainer.train()
 
         self_update_event = UpdateReceivedEvent(
             self.trainer.get_model_parameters(), self.trainer.get_model_weight(), self.addr, self.round
@@ -770,8 +763,8 @@ async def _extended_learning_cycle(self):
         # Define the functionality of the trainer node
         logging.info("Waiting global update | Assign _waiting_global_update = True")
 
-        await self.trainer.train()
         await self.trainer.test()
+        await self.trainer.train()
 
         self_update_event = UpdateReceivedEvent(
             self.trainer.get_model_parameters(), self.trainer.get_model_weight(), self.addr, self.round, local=True