Aggregators

def aggregate(self, model_params: list, weights: list, *args, **kwargs) -> Dict:
    """
    Strategy to aggregate models

    Args:
        model_params: List of model parameters received from each node
        weights: Weight for each node-model-parameter set

    Raises:
        FedbiomedAggregatorError: If the method is not defined by inheritor
    """
    msg = ErrorNumbers.FB401.value + \
        ": aggregate method should be overloaded by the choosen strategy"
    logger.critical(msg)
    raise FedbiomedAggregatorError(msg)

def check_values(self, *args, **kwargs) -> True:
    return True

def create_aggregator_args(self, *args, **kwargs) -> Dict:
    """Returns aggregator arguments that are expecting by the nodes

    Args:
        args: ignored
        kwargs: ignored

    Returns:
        contains `Aggregator` parameters/argument that will be shared with the nodes 
    """
    return self._aggregator_args or {}

def load_state_breakpoint(self, state: Dict[str, Any], **kwargs) -> None:
    """
    use for breakpoints. load the aggregator state
    """
    if not isinstance(state["parameters"], Dict):
        self._aggregator_args = Serializer.load(state['parameters'])
    else:
        self._aggregator_args = state['parameters']

def save_state_breakpoint(
    self,
    breakpoint_path: Optional[str] = None,
    **aggregator_args_create: Any,
) -> Dict[str, Any]:
    """
    use for breakpoints. save the aggregator state
    """
    aggregator_args = self.create_aggregator_args(**aggregator_args_create)
    if aggregator_args:

        if self._aggregator_args is None:
            self._aggregator_args = {}
        self._aggregator_args.update(aggregator_args)

    if breakpoint_path:
        filename = self._save_arg_to_file(breakpoint_path, 'aggregator_args', uuid.uuid4(), self._aggregator_args)

    state = {
        "class": type(self).__name__,
        "module": self.__module__,
        "parameters": filename if breakpoint_path else self._aggregator_args
    }

    return state

def secure_aggregation(
        self,
        params: List[List[int]],
        encryption_factors: List[Dict[str, List[int]]],
        secagg_random: float,
        aggregation_round: int,
        total_sample_size: int,
        training_plan: 'BaseTrainingPlan'
):
    """ Apply aggregation for encrypted model parameters

    Args:
        params: List containing list of encrypted parameters of each node
        encryption_factors: List of encrypted integers to validate encryption
        secagg_random: Randomly generated float value to validate secure aggregation correctness
        aggregation_round: The round of the aggregation.
        total_sample_size: Sum of sample sizes used for training
        training_plan: Training plan instance used for the training.

    Returns:
        aggregated model parameters
    """

    # TODO: verify with secagg context number of parties
    num_nodes = len(params)

    # TODO: Use server key here
    key = -(len(params) * 10)

    # IMPORTANT = Keep this key for testing purposes
    key = -4521514305280526329525552501850970498079782904248225896786295610941010325354834129826500373412436986239012584207113747347251251180530850751209537684586944643780840182990869969844131477709433555348941386442841023261287875379985666260596635843322044109172782411303407030194453287409138194338286254652273563418119335656859169132074431378389356392955315045979603414700450628308979043208779867835835935403213000649039155952076869962677675951924910959437120608553858253906942559260892494214955907017206115207769238347962438107202114814163305602442458693305475834199715587932463252324681290310458316249381037969151400784780
    logger.info("Securely aggregating model parameters...")

    aggregate = functools.partial(self._secagg_crypter.aggregate,
                                  current_round=aggregation_round,
                                  num_nodes=num_nodes,
                                  key=key,
                                  total_sample_size=total_sample_size
                                  )
    # Validation
    encryption_factors = [f for k, f in encryption_factors.items()]
    validation: List[int] = aggregate(params=encryption_factors)

    if len(validation) != 1 or not math.isclose(validation[0], secagg_random, abs_tol=0.01):
        raise FedbiomedAggregatorError("Aggregation is failed due to incorrect decryption.")

    aggregated_params = aggregate(params=params)

    # Convert model params
    model = training_plan.get_model_wrapper_class()
    model_params = model.unflatten(aggregated_params)

    return model_params

def set_fds(self, fds: FederatedDataSet) -> FederatedDataSet:
    self._fds = fds
    return self._fds

def set_training_plan_type(self, training_plan_type: TrainingPlans) -> TrainingPlans:
    self._training_plan_type = training_plan_type
    return self._training_plan_type

def aggregate(
        self,
        model_params: Dict[str, Dict[str, Union['torch.Tensor', 'numpy.ndarray']]],
        weights: Dict[str, float],
        *args,
        **kwargs
) -> Mapping[str, Union['torch.Tensor', 'numpy.ndarray']]:
    """ Aggregates  local models sent by participating nodes into a global model, following Federated Averaging
    strategy.

    weights is a list of single-item dictionaries, each dictionary has the node id as key, and the weight as value.
    model_params is a list of single-item dictionaries, each dictionary has the node is as key,
    and a framework-specific representation of the model parameters as value.

    Args:
        model_params: contains each model layers
        weights: contains all weights of a given layer.

    Returns:
        Aggregated parameters
    """

    model_params_processed = []
    weights_processed = []

    for node_id, params in model_params.items():

        if node_id not in weights:
            raise FedbiomedAggregatorError(
                f"{ErrorNumbers.FB401.value}. Can not find corresponding calculated weight for the "
                f"node {node_id}. Aggregation is aborted."
            )

        weight = weights[node_id]
        model_params_processed.append(params)
        weights_processed.append(weight)

    if any([x < 0. or x > 1. for x in weights_processed]) or sum(weights_processed) == 0:
        raise FedbiomedAggregatorError(
            f"{ErrorNumbers.FB401.value}. Aggregation aborted due to sum of the weights is equal to 0 {weights}. "
            f"Sample sizes received from nodes might be corrupted."
        )

    agg_params = federated_averaging(model_params_processed, weights_processed)

    return agg_params

def aggregate(self,
              model_params: Dict,
              weights: Dict[str, float],
              global_model: Dict[str, Union[torch.Tensor, np.ndarray]],
              training_plan: BaseTrainingPlan,
              training_replies: Dict,
              n_updates: int = 1,
              n_round: int = 0,
              *args, **kwargs) -> Dict:
    """
    Aggregates local models coming from nodes into a global model, using SCAFFOLD algorithm (2nd option)
    [Scaffold: Stochastic Controlled Averaging for Federated Learning][https://arxiv.org/abs/1910.06378]

    Performed computations:
    -----------------------

    - Compute participating nodes' model update:
        * update_i = y_i - x
    - Compute aggregated model parameters:
        * x(+) = x - eta_g sum_S(update_i)
    - Update participating nodes' state:
        * c_i = delta_i + 1/(K*eta_i) * update_i
    - Update the global state and all nodes' correction state:
        * c = 1/N sum_{i=1}^n c_i
        * delta_i = (c_i - c)

    where, according to paper notations
        c_i: local state variable for node `i`
        c: global state variable
        delta_i: (c_i - c), correction state for node `i`
        eta_g: server's learning rate
        eta_i: node i's learning rate
        N: total number of node participating to federated learning
        S: number of nodes considered during current round (S<=N)
        K: number of updates done during the round (ie number of data batches).
        x: global model parameters
        y_i: node i 's local model parameters at the end of the round

    Args:
        model_params: list of models parameters received from nodes
        weights: weights depicting sample proportions available
            on each node. Unused for Scaffold.
        global_model: global model, ie aggregated model
        training_plan (BaseTrainingPlan): instance of TrainingPlan
        training_replies: Training replies from each node that participates in the current round
        n_updates: number of updates (number of batch performed). Defaults to 1.
        n_round: current round. Defaults to 0.

    Returns:
        Aggregated parameters, as a dict mapping weight names and values.

    Raises:
        FedbiomedAggregatorError: If no FederatedDataset is attached to this
            Scaffold instance, or if `node_ids` do not belong to the dataset
            attached to it.
    """
    # Gather the learning rates used by nodes, updating `self.nodes_lr`.
    self.set_nodes_learning_rate_after_training(training_plan, training_replies)
    # At round 0, initialize zero-valued correction states.
    if n_round == 0:
        self.init_correction_states(global_model)
    # Check that the input node_ids match known ones.
    if not set(model_params).issubset(self._fds.node_ids()):
        raise FedbiomedAggregatorError(
            "Received updates from nodes that are unknown to this aggregator."
        )
    # Compute the node-wise model update: (x^t - y_i^t).
    model_updates = {
        node_id: {
            key: (global_model[key] - local_value)
            for key, local_value in params.items()
        }
        for node_id, params in model_params.items()
    }
    # Update all Scaffold state variables.
    self.update_correction_states(model_updates, n_updates)
    # Compute and return the aggregated model parameters.
    global_new = {}  # type: Dict[str, Union[torch.Tensor, np.ndarray]]
    for key, val in global_model.items():
        upd = sum(model_updates[node_id][key] for node_id in model_params)
        global_new[key] = val - upd * (self.server_lr / len(model_params))
    return global_new

def check_values(self, n_updates: int, training_plan: BaseTrainingPlan) -> True:
    """Check if all values/parameters are correct and have been set before using aggregator.

    Raise an error otherwise.

    This can prove useful if user has set wrong hyperparameter values, so that user will
    have errors before performing first round of training

    Args:
        n_updates: number of updates. Must be non-zero and an integer.
        training_plan: training plan. used for checking if optimizer is SGD, otherwise,
            triggers warning.

    Raises:
        FedbiomedAggregatorError: triggered if `num_updates` entry is missing (needed for Scaffold aggregator)
        FedbiomedAggregatorError: triggered if any of the learning rate(s) equals 0
        FedbiomedAggregatorError: triggered if number of updates equals 0 or is not an integer
        FedbiomedAggregatorError: triggered if [FederatedDataset][fedbiomed.researcher.datasets.FederatedDataset]
            has not been set.
    """
    if n_updates is None:
        raise FedbiomedAggregatorError("Cannot perform Scaffold: missing 'num_updates' entry in the training_args")
    elif n_updates <= 0 or int(n_updates) != float(n_updates):
        raise FedbiomedAggregatorError(
            "n_updates should be a positive non zero integer, but got "
            f"n_updates: {n_updates} in SCAFFOLD aggregator"
        )
    if self._fds is None:
        raise FedbiomedAggregatorError(
            "Federated Dataset not provided, but needed for Scaffold. Please use setter `set_fds()`."
        )

    return True

def create_aggregator_args(
    self,
    global_model: Dict[str, Union[torch.Tensor, np.ndarray]],
    node_ids: Collection[str]
) -> Dict[str, Dict[str, Any]]:
    """Return correction states that are to be sent to the nodes.

    Args:
        global_model: parameters of the global model, formatted as a dict
            mapping weight tensors to their names.
        node_ids: identifiers of the nodes that are to receive messages.

    Returns:
        Aggregator arguments to share with the nodes for the next round
    """
    # Optionally initialize states, and verify that nodes are known.
    if not self.nodes_deltas:
        self.init_correction_states(global_model)
    if not set(node_ids).issubset(self._fds.node_ids()):
        raise FedbiomedAggregatorError(
            "Scaffold cannot create aggregator args for nodes that are not"
            "covered by its attached FederatedDataset."
        )

    aggregator_dat = {}
    for node_id in node_ids:
        # If a node was late-added to the FederatedDataset, create states.
        if node_id not in self.nodes_deltas:
            zeros = {key: initialize(val)[1] for key, val in self.global_state.items()}
            self.nodes_deltas[node_id] = zeros
            self.nodes_states[node_id] = copy.deepcopy(zeros)
        # Add information for the current node to the message dicts.
        aggregator_dat[node_id] = {
            'aggregator_name': self.aggregator_name,
            'aggregator_correction': self.nodes_deltas[node_id]
        }

    return aggregator_dat

def init_correction_states(
    self,
    global_model: Dict[str, Union[torch.Tensor, np.ndarray]],
) -> None:
    """Initialize Scaffold state variables.

    Args:
        global_model: parameters of the global model, formatted as a dict
            mapping weight tensors to their names.

    Raises:
        FedbiomedAggregatorError: if no FederatedDataset is attached to
            this aggregator.
    """
    # Gather node ids from the attached FederatedDataset.
    if self._fds is None:
        raise FedbiomedAggregatorError(
            "Cannot initialize correction states: Scaffold aggregator does "
            "not have a FederatedDataset attached."
        )
    node_ids = self._fds.node_ids()
    # Initialize nodes states with zero scalars, that will be summed into actual tensors.
    init_params = {key: initialize(tensor)[1] for key, tensor in global_model.items()}
    self.nodes_deltas = {node_id: copy.deepcopy(init_params) for node_id in node_ids}
    self.nodes_states = copy.deepcopy(self.nodes_deltas)
    self.global_state = init_params

def load_state_breakpoint(self, state: Dict[str, Any] = None):
    super().load_state_breakpoint(state)

    self.server_lr = self._aggregator_args['server_lr']

    # loading global state
    global_state_filename = self._aggregator_args['global_state_filename']
    self.global_state = Serializer.load(global_state_filename)

    for node_id in self._aggregator_args['nodes']:
        self.nodes_deltas[node_id] = self._aggregator_args[node_id]['aggregator_correction']

    self.nodes_states = copy.deepcopy(self.nodes_deltas)

def save_state_breakpoint(
    self,
    breakpoint_path: str,
    global_model: Mapping[str, Union[torch.Tensor, np.ndarray]]
) -> Dict[str, Any]:
    # adding aggregator parameters to the breakpoint that wont be sent to nodes
    self._aggregator_args['server_lr'] = self.server_lr

    # saving global state variable into a file
    filename = os.path.join(breakpoint_path, f"global_state_{uuid.uuid4()}.mpk")
    Serializer.dump(self.global_state, filename)
    self._aggregator_args['global_state_filename'] = filename

    self._aggregator_args["nodes"] = self._fds.node_ids()
    # adding aggregator parameters that will be sent to nodes afterwards
    return super().save_state_breakpoint(
        breakpoint_path, global_model=global_model, node_ids=self._fds.node_ids()
    )

def set_nodes_learning_rate_after_training(
    self,
    training_plan: BaseTrainingPlan,
    training_replies: Dict,
) -> Dict[str, List[float]]:
    """Gets back learning rate of optimizer from Node (if learning rate scheduler is used)

    Args:
        training_plan: training plan instance
        training_replies: training replies that must contain am `optimizer_args`
            entry and a learning rate

    Raises:
        FedbiomedAggregatorError: raised when setting learning rate has been unsuccessful

    Returns:
        Dict[str, List[float]]: dictionary mapping node_id and a list of float, as many as
            the number of layers contained in the model (in Pytroch, each layer can have a specific learning rate).
    """

    n_model_layers = len(training_plan.get_model_params(
        only_trainable=False,
        exclude_buffers=True)
    )
    for node_id in self._fds.node_ids():
        lrs: Dict[str, float] = {}

        node = training_replies.get(node_id, None)
        if node is not None:
            lrs = training_replies[node_id]["optimizer_args"].get('lr')

        if node is None or lrs is None:
            # fall back to default value if no lr information was provided
            lrs = training_plan.optimizer().get_learning_rate()

        if len(lrs) != n_model_layers:
            raise FedbiomedAggregatorError(
                "Error when setting node learning rate for SCAFFOLD: cannot extract node learning rate."
            )

        self.nodes_lr[node_id] = lrs
    return self.nodes_lr

def set_training_plan_type(self, training_plan_type: TrainingPlans) -> TrainingPlans:
    """
    Overrides `set_training_plan_type` from parent class.
    Checks the training plan type, and if it is SKlearnTrainingPlan,
    raises an error. Otherwise, calls parent method.

    Args:
        training_plan_type (TrainingPlans): training_plan type

    Raises:
        FedbiomedAggregatorError: raised if training_plan type has been set to SKLearn training plan

    Returns:
        TrainingPlans: training plan type
    """
    if training_plan_type == TrainingPlans.SkLearnTrainingPlan:
        raise FedbiomedAggregatorError("Aggregator SCAFFOLD not implemented for SKlearn")
    training_plan_type = super().set_training_plan_type(training_plan_type)

    # TODO: trigger a warning if user is trying to use scaffold with something else than SGD
    return training_plan_type

def update_correction_states(
    self,
    model_updates: Dict[str, Dict[str, Union[np.ndarray, torch.Tensor]]],
    n_updates: int,
) -> None:
    """Update all Scaffold state variables based on node-wise model updates.

    Performed computations:
    ----------------------

    - Update participating nodes' state:
        * c_i = delta_i + 1/(K*eta_i) * update_i
    - Update the global state and all nodes' correction state:
        * c = 1/N sum_{i=1}^n c_i
        * delta_i = (c_i - c)

    Args:
        model_updates: node-wise model weight updates.
        n_updates: number of local optimization steps.
    """
    # Update the node-wise states for participating nodes:
    # c_i^{t+1} = delta_i^t + (x^t - y_i^t) / (M * eta)
    for node_id, updates in model_updates.items():
        d_i = self.nodes_deltas[node_id]
        for (key, val) in updates.items():
            if self.nodes_lr[node_id].get(key) is not None:
                self.nodes_states[node_id].update(
                    {
                    key: d_i[key] + val / (self.nodes_lr[node_id][key] * n_updates)
                     }
                )
    # Update the global state: c^{t+1} = average(c_i^{t+1})
    for key in self.global_state:
        self.global_state[key] = 0
        for state in self.nodes_states.values():
            if state.get(key) is not None:
                self.global_state[key] = (
                    sum(state[key] for state in self.nodes_states.values())
                        / len(self.nodes_states)
                    )

    # Compute the new node-wise correction states:
    # delta_i^{t+1} = c_i^{t+1} - c^{t+1}
    self.nodes_deltas = {
        node_id: {
            key: val - self.global_state[key] for key, val in state.items()
        }
        for node_id, state in self.nodes_states.items()
    }