Skip to content

Documentation for Aggregator Module

Aggregator

Bases: ABC

Source code in nebula/core/aggregation/aggregator.py 
 18
 19
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
class Aggregator(ABC):
    def __init__(self, config=None, engine=None):
        self.config = config
        self.engine: Engine = engine
        self._addr = config.participant["network_args"]["addr"]
        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()

        scenario = self.config.participant["scenario_args"]["federation"]
        self._update_storage = factory_update_handler(scenario, self, self._addr)

    def __str__(self):
        return self.__class__.__name__

    def __repr__(self):
        return self.__str__()

    @property
    def us(self):
        """Federation type UpdateHandler (e.g. DFL-UpdateHandler, CFL-UpdateHandler...)"""
        return self._update_storage

    @abstractmethod
    def run_aggregation(self, models):
        if len(models) == 0:
            logging.error("Trying to aggregate models when there are no models")
            return None

    async def init(self):
        await self.us.init(self.engine.rb.get_role_name(True))

    async def update_federation_nodes(self, federation_nodes: set):
        """
        Updates the current set of nodes expected to participate in the upcoming aggregation round.

        This method informs the update handler (`us`) about the new set of federation nodes,
        clears any pending models, and attempts to acquire the aggregation lock to prepare
        for model aggregation. If the aggregation process is already running, it releases the lock
        and tries again to ensure proper cleanup between rounds.

        Args:
            federation_nodes (set): A set of addresses representing the nodes expected to contribute
                                    updates for the next aggregation round.

        Raises:
            Exception: If the aggregation process is already running and the lock cannot be released.
        """
        await self.us.round_expected_updates(federation_nodes=federation_nodes)

        # If the aggregation lock is held, release it to prepare for the new round
        if self._aggregation_done_lock.locked():
            logging.info("🔄  update_federation_nodes | Aggregation lock is held, releasing for new round")
            try:
                await self._aggregation_done_lock.release_async()
            except Exception as e:
                logging.warning(f"🔄  update_federation_nodes | Error releasing aggregation lock: {e}")
                # If we can't release the lock, we might be in the middle of aggregation
                # In this case, we should wait a bit and try again
                await asyncio.sleep(0.1)
                if self._aggregation_done_lock.locked():
                    raise Exception("It is not possible to set nodes to aggregate when the aggregation is running.")

        # Now acquire the lock for the new round
        self._federation_nodes = federation_nodes
        self._pending_models_to_aggregate.clear()
        await self._aggregation_done_lock.acquire_async(
            timeout=self.config.participant["aggregator_args"]["aggregation_timeout"]
        )

    def get_nodes_pending_models_to_aggregate(self):
        return self._federation_nodes

    async def get_aggregation(self):
        """
        Handles the aggregation process for a training round.

        This method waits for all expected model updates from federation nodes or until a timeout occurs.
        It uses an asynchronous lock to coordinate access and includes an early exit mechanism if all
        updates are received before the timeout. Once the condition is satisfied, it releases the lock,
        collects the updates, identifies any missing nodes, and publishes an `AggregationEvent`.
        Finally, it runs the aggregation algorithm and returns the result.

        Returns:
            Any: The result of the aggregation process, as returned by `run_aggregation`.

        Raises:
            TimeoutError: If the aggregation lock is not acquired within the defined timeout.
            asyncio.CancelledError: If the aggregation lock acquisition is cancelled.
            Exception: For any other unexpected errors during the aggregation process.
        """
        lock_acquired = False
        lock_task = None
        skip_task = None
        try:
            timeout = self.config.participant["aggregator_args"]["aggregation_timeout"]
            logging.info(f"Aggregation timeout: {timeout} starts...")
            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(
                [lock_task, skip_task],
                return_when=asyncio.FIRST_COMPLETED,
            )

            if skip_task in done:
                logging.info("Skipping aggregation timeout, updates received before grace time")
                self._aggregation_waiting_skip.clear()
                if not lock_task.done():
                    lock_task.cancel()

            if lock_task in done:
                try:
                    await lock_task
                    lock_acquired = True
                except TimeoutError:
                    logging.info("🔄  get_aggregation | Timeout reached; aggregating received updates")
                except asyncio.CancelledError:
                    logging.info("🔄  get_aggregation | Lock acquisition was cancelled")

        except asyncio.CancelledError:
            logging.exception("🔄  get_aggregation | Aggregation wait was cancelled")
        except Exception as e:
            logging.exception(f"🔄  get_aggregation | Error acquiring lock: {e}")
        finally:
            for task in (lock_task, skip_task):
                if task is None:
                    continue
                if not task.done():
                    task.cancel()
                try:
                    await task
                except asyncio.CancelledError:
                    pass
                except TimeoutError:
                    pass
            if lock_acquired or self._aggregation_done_lock.locked():
                await self._aggregation_done_lock.release_async()

        await self.us.stop_notifying_updates()
        updates = await self.us.get_round_updates()
        if not updates:
            logging.info(f"🔄  get_aggregation | No updates has been received..resolving conflict to continue...")
            updates = {self._addr: await self.engine.resolve_missing_updates()}

        missing_nodes = await self.us.get_round_missing_nodes()
        if missing_nodes:
            logging.info(f"🔄  get_aggregation | Aggregation incomplete, missing models from: {missing_nodes}")
        else:
            logging.info("🔄  get_aggregation | All models accounted for, proceeding with aggregation.")

        await self.us.before_aggregation(updates, self._federation_nodes)

        agg_event = AggregationEvent(updates, self._federation_nodes, missing_nodes)
        await EventManager.get_instance().publish_node_event(agg_event)
        aggregated_result = self.run_aggregation(updates)
        return aggregated_result

    def print_model_size(self, model):
        total_memory = 0

        for _, param in model.items():
            num_params = param.numel()
            memory_usage = param.element_size() * num_params
            total_memory += memory_usage

        total_memory_in_mb = total_memory / (1024**2)
        logging.info(f"print_model_size | Model size: {total_memory_in_mb} MB")

    async def notify_all_updates_received(self):
        self._aggregation_waiting_skip.set()

us property

Federation type UpdateHandler (e.g. DFL-UpdateHandler, CFL-UpdateHandler...)

get_aggregation() async

Handles the aggregation process for a training round.

This method waits for all expected model updates from federation nodes or until a timeout occurs. It uses an asynchronous lock to coordinate access and includes an early exit mechanism if all updates are received before the timeout. Once the condition is satisfied, it releases the lock, collects the updates, identifies any missing nodes, and publishes an AggregationEvent. Finally, it runs the aggregation algorithm and returns the result.

Returns:

Name Type Description
Any

The result of the aggregation process, as returned by run_aggregation.

Raises:

Type Description
TimeoutError

If the aggregation lock is not acquired within the defined timeout.
CancelledError

If the aggregation lock acquisition is cancelled.
Exception

For any other unexpected errors during the aggregation process.
Source code in nebula/core/aggregation/aggregator.py 
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
async def get_aggregation(self):
    """
    Handles the aggregation process for a training round.

    This method waits for all expected model updates from federation nodes or until a timeout occurs.
    It uses an asynchronous lock to coordinate access and includes an early exit mechanism if all
    updates are received before the timeout. Once the condition is satisfied, it releases the lock,
    collects the updates, identifies any missing nodes, and publishes an `AggregationEvent`.
    Finally, it runs the aggregation algorithm and returns the result.

    Returns:
        Any: The result of the aggregation process, as returned by `run_aggregation`.

    Raises:
        TimeoutError: If the aggregation lock is not acquired within the defined timeout.
        asyncio.CancelledError: If the aggregation lock acquisition is cancelled.
        Exception: For any other unexpected errors during the aggregation process.
    """
    lock_acquired = False
    lock_task = None
    skip_task = None
    try:
        timeout = self.config.participant["aggregator_args"]["aggregation_timeout"]
        logging.info(f"Aggregation timeout: {timeout} starts...")
        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(
            [lock_task, skip_task],
            return_when=asyncio.FIRST_COMPLETED,
        )

        if skip_task in done:
            logging.info("Skipping aggregation timeout, updates received before grace time")
            self._aggregation_waiting_skip.clear()
            if not lock_task.done():
                lock_task.cancel()

        if lock_task in done:
            try:
                await lock_task
                lock_acquired = True
            except TimeoutError:
                logging.info("🔄  get_aggregation | Timeout reached; aggregating received updates")
            except asyncio.CancelledError:
                logging.info("🔄  get_aggregation | Lock acquisition was cancelled")

    except asyncio.CancelledError:
        logging.exception("🔄  get_aggregation | Aggregation wait was cancelled")
    except Exception as e:
        logging.exception(f"🔄  get_aggregation | Error acquiring lock: {e}")
    finally:
        for task in (lock_task, skip_task):
            if task is None:
                continue
            if not task.done():
                task.cancel()
            try:
                await task
            except asyncio.CancelledError:
                pass
            except TimeoutError:
                pass
        if lock_acquired or self._aggregation_done_lock.locked():
            await self._aggregation_done_lock.release_async()

    await self.us.stop_notifying_updates()
    updates = await self.us.get_round_updates()
    if not updates:
        logging.info(f"🔄  get_aggregation | No updates has been received..resolving conflict to continue...")
        updates = {self._addr: await self.engine.resolve_missing_updates()}

    missing_nodes = await self.us.get_round_missing_nodes()
    if missing_nodes:
        logging.info(f"🔄  get_aggregation | Aggregation incomplete, missing models from: {missing_nodes}")
    else:
        logging.info("🔄  get_aggregation | All models accounted for, proceeding with aggregation.")

    await self.us.before_aggregation(updates, self._federation_nodes)

    agg_event = AggregationEvent(updates, self._federation_nodes, missing_nodes)
    await EventManager.get_instance().publish_node_event(agg_event)
    aggregated_result = self.run_aggregation(updates)
    return aggregated_result

update_federation_nodes(federation_nodes) async

Updates the current set of nodes expected to participate in the upcoming aggregation round.

This method informs the update handler (us) about the new set of federation nodes, clears any pending models, and attempts to acquire the aggregation lock to prepare for model aggregation. If the aggregation process is already running, it releases the lock and tries again to ensure proper cleanup between rounds.

Parameters:

Name Type Description Default
federation_nodes set

A set of addresses representing the nodes expected to contribute updates for the next aggregation round.

 required

Raises:

Type Description
Exception

If the aggregation process is already running and the lock cannot be released.
Source code in nebula/core/aggregation/aggregator.py 
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
async def update_federation_nodes(self, federation_nodes: set):
    """
    Updates the current set of nodes expected to participate in the upcoming aggregation round.

    This method informs the update handler (`us`) about the new set of federation nodes,
    clears any pending models, and attempts to acquire the aggregation lock to prepare
    for model aggregation. If the aggregation process is already running, it releases the lock
    and tries again to ensure proper cleanup between rounds.

    Args:
        federation_nodes (set): A set of addresses representing the nodes expected to contribute
                                updates for the next aggregation round.

    Raises:
        Exception: If the aggregation process is already running and the lock cannot be released.
    """
    await self.us.round_expected_updates(federation_nodes=federation_nodes)

    # If the aggregation lock is held, release it to prepare for the new round
    if self._aggregation_done_lock.locked():
        logging.info("🔄  update_federation_nodes | Aggregation lock is held, releasing for new round")
        try:
            await self._aggregation_done_lock.release_async()
        except Exception as e:
            logging.warning(f"🔄  update_federation_nodes | Error releasing aggregation lock: {e}")
            # If we can't release the lock, we might be in the middle of aggregation
            # In this case, we should wait a bit and try again
            await asyncio.sleep(0.1)
            if self._aggregation_done_lock.locked():
                raise Exception("It is not possible to set nodes to aggregate when the aggregation is running.")

    # Now acquire the lock for the new round
    self._federation_nodes = federation_nodes
    self._pending_models_to_aggregate.clear()
    await self._aggregation_done_lock.acquire_async(
        timeout=self.config.participant["aggregator_args"]["aggregation_timeout"]
    )