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._servkey is None:
        raise FedbiomedSecureAggregationError(
            f"{ErrorNumbers.FB417.value}: Can not aggregate parameters, Servkey context is "
            f"not configured. Please setup secure aggregation before the aggregation."
        )

    if 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"
        )

    key = self._servkey.context["server_key"]
    biprime = self._servkey.context["biprime"]
    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"]["_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(
        {
            "_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._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
            ),
        }
    )
    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 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 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
    request = self._REQUEST_DELETE(
        researcher_id=self._researcher_id,
        secagg_id=self._secagg_id,
        element=self._element.value,
        experiment_id=self._experiment_id,
    )

    _ = self._launch_request(request)

    status = self._secagg_manager.remove(self._secagg_id, self.experiment_id)
    if status:
        logger.debug(
            f"Context element successfully deleted for researcher_id='{environ['ID']}' "
            f"secagg_id='{self._secagg_id}'"
        )
    else:
        logger.error(
            f"{ErrorNumbers.FB415.value}: No such context element secagg_id={self._secagg_id} "
            f"on researcher researcher_id='{environ['ID']}'"
        )

    return True

@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

@abstractmethod
def secagg_round(
    self,
    request: Message,
) -> bool:
    """Negotiate secagg context element action with defined parties.

    Args:
        msg: message sent to the parties during the round

    Returns:
        True if secagg context element action could be done for all parties, False if at least
            one of the parties could not do the context element action.
    """

def setup(self) -> Dict:
    """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.
    """
    context = self._secagg_manager.get(self._secagg_id, self._experiment_id)
    if context and matching_parties_servkey(context, self._parties):
        logger.info(
            f"{ErrorNumbers.FB415.value}: secagg context for {self._secagg_id} exists"
        )
        self._context = context['context']
    else:
        request = self._REQUEST_SETUP(
            researcher_id=self._researcher_id,
            secagg_id=self._secagg_id,
            element=self._element.value,
            experiment_id=self._experiment_id,
            parties=self._parties,
        )

        self._context = self.secagg_round(request)

    self._status = True

    return self._status

def secagg_round(
    self,
    request: Message,
) -> Tuple[dict, dict[str, bool]]:
    """Negotiate secagg context element action with defined parties for DH key exchange.

    Args:
        request: message sent to the parties during the round

    Returns:
        A tuple of
            - a dict containing the context describing this secagg context element
            - a dict where key/values are node ids/boolean with success status for
                secagg on each party

    """

    _ = self._launch_request(request)
    context = {}
    return self._register(context)

def secagg_round(
    self,
    request: Message,
) -> Tuple[dict, dict[str, bool]]:
    """Negotiate secagg context element action with defined parties for secret sharing.

    Args:
        request: message sent to the parties during the round

    Returns:
        A tuple of
            - a dict containing the context describing this secagg context element
            - a dict where key/values are node ids/boolean with success status
                for secagg on each party
    """

    replies = self._launch_request(request)
    servkey: int = AdditiveShares([AdditiveShare(rep.share) for
        rep in replies.values()]).reconstruct()

    biprime = get_default_biprime()
    context = {'server_key': -servkey, 'biprime': biprime}
    return self._register(context)

@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