Secagg

def aggregate(
        self,
        *args,
        round_: int,
        total_sample_size: int,
        model_params: Dict[str, List[int]],
        encryption_factors: Dict[str, Union[List[int], None]] = {},
        num_expected_params: int = 1,
        **kwargs
) -> List[float]:
    """Aggregates given model parameters

    Args:
        round_: current training round number
        total_sample_size: sum of number of samples used by all nodes
        model_params: model parameters from the participating nodes
        encryption_factors: encryption factors from the participating nodes
        num_expected_params: number of decrypted parameters to decode from the model parameters

    Returns:
        Aggregated parameters

    Raises:
        FedbiomedSecureAggregationError: secure aggregation context not properly configured
    """

    if self._biprime is None or self._servkey is None:
        raise FedbiomedSecureAggregationError(
            f"{ErrorNumbers.FB417.value}: Can not aggregate parameters, one of Biprime or Servkey context is "
            f"not configured. Please setup secure aggregation before the aggregation.")

    if not self._biprime.status or not self._servkey.status:
        raise FedbiomedSecureAggregationError(
            f"{ErrorNumbers.FB417.value}: Can not aggregate parameters, one of Biprime or Servkey context is "
            f"not set properly")

    biprime = self._biprime.context["context"]["biprime"]
    key = self._servkey.context["context"]["server_key"]

    num_nodes = len(model_params)

    aggregate = functools.partial(self._secagg_crypter.aggregate,
                                  current_round=round_,
                                  num_nodes=num_nodes,
                                  key=key,
                                  total_sample_size=total_sample_size,
                                  biprime=biprime,
                                  clipping_range=self.clipping_range)

    return self._aggregate(
        model_params,
        encryption_factors,
        aggregate,
        num_expected_params,
    )

@classmethod
def load_state_breakpoint(
        cls,
        state: Dict
) -> 'SecureAggregation':
    """Create a `SecureAggregation` object from a saved state

    Args:
        state: saved state to restore in the created object

    Returns:
        The created `SecureAggregation` object
    """
    if state["attributes"]["_biprime"] is not None:
        state["attributes"]["_biprime"] = SecaggBiprimeContext. \
            load_state_breakpoint(state=state["attributes"]["_biprime"])

    if state["attributes"]["_servkey"] is not None:
        state["attributes"]["_servkey"] = SecaggServkeyContext. \
            load_state_breakpoint(state=state["attributes"]["_servkey"])

    return super().load_state_breakpoint(state)

def save_state_breakpoint(self) -> Dict[str, Any]:
    """Saves state of the secagg

    Returns:
        The secagg state to be saved
    """
    state = super().save_state_breakpoint()

    state["attributes"].update({
        "_biprime": self._biprime.save_state_breakpoint() if self._biprime is not None else None,
        "_servkey": self._servkey.save_state_breakpoint() if self._servkey is not None else None,
    })

    return state

def setup(self,
          parties: List[str],
          experiment_id: str,
          force: bool = False):
    """Setup secure aggregation instruments.

    Requires setting `parties` and `experiment_id` if they are not set in previous secagg
    setups. It is possible to execute without any argument if SecureAggregation
    has already `parties` and `experiment_id` defined. This feature provides researcher
    execute `secagg.setup()` if any connection issu#e

    Args:
        parties: Parties that participates secure aggregation
        experiment_id: The id of the experiment
        force: Forces secagg setup even context is already existing

    Returns:
        Status of setup

    Raises
        FedbiomedSecureAggregationError: Invalid argument type
    """
    super().setup(parties, experiment_id, force)

    if not self._biprime.status or force:
        self._biprime.setup()

    if not self._servkey.status or force:
        self._servkey.setup()

    return True

def train_arguments(self) -> Dict:
    """Gets train arguments for secagg train request

    Returns:
        Arguments that are going to be attached to the experiment.
    """
    arguments = super().train_arguments()
    arguments.update({
        'secagg_servkey_id': self._servkey.secagg_id if self._servkey is not None else None,
        'secagg_biprime_id': self._biprime.secagg_id if self._biprime is not None else None,
    })
    return arguments

def aggregate(
        self,
        *args,
        model_params: Dict[str, List[int]],
        total_sample_size: int,
        encryption_factors: Dict[str, Union[List[int], None]] = {},
        **kwargs
) -> List[float]:
    """Aggregates given model parameters

    Args:
        total_sample_size: sum of number of samples used by all nodes
        model_params: model parameters from the participating nodes
        encryption_factors: encryption factors from the participating nodes

    Returns:
        Aggregated parameters

    Raises:
        FedbiomedSecureAggregationError: secure aggregation context not properly configured
    """

    if self._dh is None:
        raise FedbiomedSecureAggregationError(
            f"{ErrorNumbers.FB417.value}: Can not aggregate parameters, Diffie Hellman context is "
            f"not configured. Please setup secure aggregation before the aggregation.")

    if not self._dh.status:
        raise FedbiomedSecureAggregationError(
            f"{ErrorNumbers.FB417.value}: Can not aggregate parameters, Diffie Hellman context is "
            f"not set properly")

    aggregate = functools.partial(
        self._secagg_crypter.aggregate,
        total_sample_size=total_sample_size,
        clipping_range=self.clipping_range)

    return self._aggregate(
        model_params,
        encryption_factors,
        aggregate,
    )

@classmethod
def load_state_breakpoint(
        cls,
        state: Dict
) -> 'SecureAggregation':
    """Create a `SecureAggregation` object from a saved state

    Args:
        state: saved state to restore in the created object

    Returns:
        The created `SecureAggregation` object
    """
    if state["attributes"]["_dh"] is not None:
        state["attributes"]["_dh"] = SecaggDHContext. \
            load_state_breakpoint(state=state["attributes"]["_dh"])

    return super().load_state_breakpoint(state)

def save_state_breakpoint(self) -> Dict[str, Any]:
    """Saves state of the secagg

    Returns:
        The secagg state to be saved
    """
    state = super().save_state_breakpoint()

    state["attributes"].update({
        "_dh": self._dh.save_state_breakpoint() if self._dh is not None else None,
    })

    return state

def setup(self,
          parties: List[str],
          experiment_id: str,
          force: bool = False) -> bool:
    """Setup secure aggregation instruments.

    Requires setting `parties` and `experiment_id` if they are not set in previous secagg
    setups. It is possible to execute without any argument if SecureAggregation
    has already `parties` and `experiment_id` defined. This feature provides researcher
    execute `secagg.setup()` if any connection issue

    Args:
        parties: Parties that participates secure aggregation
        experiment_id: The id of the experiment
        force: Forces secagg setup even if context is already existing

    Returns:
        Status of setup

    Raises
        FedbiomedSecureAggregationError: Invalid argument type
    """

    parties = list(
        filter(lambda x: x != environ["ID"], parties)
    )

    super().setup(parties, experiment_id, force)

    if not self._dh.status or force:
        self._dh.setup()

    return self._dh.status

def train_arguments(self) -> Dict:
    """Gets train arguments for secagg train request

    Returns:
        Arguments that are going to be attached to the experiment.
    """
    arguments = super().train_arguments()
    arguments.update({
        'secagg_dh_id': self._dh.secagg_id if self._dh is not None else None,
    })
    return arguments

def delete(self) -> bool:
    """Delete secagg context element on defined parties.

    Returns:
        True if secagg context element could be deleted for all parties, False if at least
            one of the parties could not delete context element.
    """
    self._status = False
    self._context = None
    msg = ResearcherMessages.format_outgoing_message({
        'researcher_id': self._researcher_id,
        'secagg_id': self._secagg_id,
        'element': self._element.value,
        'experiment_id': self._experiment_id,
        'command': 'secagg-delete',
    })
    return self._secagg_round(msg, False, self._delete_payload)

@staticmethod
def load_state_breakpoint(
        state: Dict[str, Any]
) -> 'SecaggContext':

    """
    Method for loading secagg state from breakpoint state

    Args:
        state: The state that will be loaded
    """

    # Get class
    cls = getattr(importlib.import_module(state["module"]), state["class"])

    # Validate experiment id
    spec = get_method_spec(cls)
    if 'experiment_id' in spec:
        secagg = cls(**state["arguments"])
    else:
        state["arguments"].pop('experiment_id')
        secagg = cls(**state["arguments"])

    for key, value in state["attributes"].items():
        setattr(secagg, key, value)

    return secagg

def save_state_breakpoint(self) -> Dict[str, Any]:
    """Method for saving secagg state for saving breakpoints

    Returns:
        The state of the secagg
    """
    # `_v` and `_requests` dont need to be savec (properly initiated in constructor)
    state = {
        "class": type(self).__name__,
        "module": self.__module__,
        "arguments": {
            "secagg_id": self._secagg_id,
            "parties": self._parties,
            "experiment_id": self._experiment_id,

        },
        "attributes": {
            "_status": self._status,
            "_context": self._context,
            "_researcher_id": self._researcher_id,
        }
    }
    return state

def set_experiment_id(self, experiment_id: Union[str, None]) -> None:
    """Setter for secagg context element experiment_id

    Args:
        experiment_id: ID of the experiment to which this secagg context element is attached.

    Raises:
        FedbiomedSecaggError: bad argument type or value
    """

    if not isinstance(experiment_id, (str, type(None))):
        errmess = f'{ErrorNumbers.FB415.value}: bad parameter `experiment_id` must be a str or None if the ' \
                  f'context is set for biprime.'
        logger.error(errmess)
        raise FedbiomedSecaggError(errmess)

    self._experiment_id = experiment_id

def setup(self) -> bool:
    """Setup secagg context element on defined parties.

    Returns:
        True if secagg context element could be setup for all parties, False if at least
            one of the parties could not setup context element.
    """
    msg = ResearcherMessages.format_outgoing_message({
        'researcher_id': self._researcher_id,
        'secagg_id': self._secagg_id,
        'element': self._element.value,
        'experiment_id': self._experiment_id,
        'parties': self._parties,
        'command': 'secagg',
    })

    return self._secagg_round(msg, True, self._create_payload)

@classmethod
def load_state_breakpoint(
        cls,
        state: Dict
) -> 'SecureAggregation':
    """Create a `SecureAggregation` object from a saved state

    Args:
        state: saved state to restore in the created object

    Returns:
        The created `SecureAggregation` object
    """
    secagg = cls(
        scheme=SecureAggregationSchemes(state['arguments']['scheme'])
    )

    for name, val in state["attributes_states"].items():

        _sub_cls = getattr(importlib.import_module(val['module']), val["class"])
        instance = _sub_cls.load_state_breakpoint(val)
        setattr(secagg, name, instance)

    return secagg

@abstractmethod
def save_state_breakpoint(self) -> Dict[str, Any]:
    """Saves state of the secagg

    This methods also save states of `__secagg` which provides
    a single entry point for secagg schemes

    Returns:
        The secagg state to be saved
    """

    state = {
        "class": type(self).__name__,
        "module": self.__module__,
        "arguments": {
            'scheme': self.__scheme.value,
        },
        "attributes": {},
        "attributes_states": {
            "_SecureAggregation__secagg": self.__secagg.save_state_breakpoint()
        }
    }

    return state