Jobs

def execute(self) -> Dict[str, FAReply]:
    """Send the FA request to all nodes and collect replies.

    Raises:
        FedbiomedError: If any node returns an error, or if no replies are received.

    Returns:
        A dictionary mapping node IDs to their `FAReply`.
    """
    fa_request = dict(
        researcher_id=self._researcher_id,
        experiment_id=self._experiment_id,
        stats=self._stats,
        fa_id=self._fa_id,
        stats_args=self._stats_args,
        dataset_schema=self._dataset_schema,
    )

    requests = MessagesByNode()
    for node in self._nodes:
        requests[node] = FARequest(
            **{
                **fa_request,
                "dataset_id": self._federated_dataset.data()[node]["dataset_id"],
            }
        )

    with self._reqs.send(requests, self._nodes, self._policies) as responses:
        errors: Dict[str, ErrorMessage] = responses.errors()
        replies: Dict[str, FAReply] = responses.replies()

    if errors:
        # Log errors encountered during the execution of the FA request
        for node_id, error in errors.items():
            logger.error(
                f"Node {node_id} analytics error [{error.errnum}]: {error.extra_msg}"
            )
        # Treat any node error as fatal
        raise FedbiomedError(
            f"FA request execution failed with errors from nodes: {', '.join(errors.keys())}"
        )

    # Ensure there is at least one successful reply
    if not replies:
        raise FedbiomedError("No successful replies received.")

    return replies

@abstractmethod
def execute(self) -> Any:
    """Payload of the job.

    Completes a request to the job's nodes and collects replies.

    Returns:
        values specific to the type of job
    """

def execute(self) -> Dict[str, PreprocReply]:
    """Executes preprocessing request

    Returns:
        A dictionary mapping node IDs to their respective preprocessing replies.
    """
    preproc_request = dict(
        researcher_id=self._researcher_id,
        experiment_id=self._experiment_id,
        preproc_type=self._preproc_type,
        preproc_step=self._preproc_step,
        preproc_id=self._preproc_id,
    )

    requests = MessagesByNode()
    for node in self._nodes:
        # Note: deepcopy is used to avoid mutation issues across nodes
        # may be later reconsidered depending on size and content of preproc_args
        preproc_args = deepcopy(self._preproc_args)

        requests[node] = PreprocRequest(
            **{
                **preproc_request,
                "dataset_id": self._federated_dataset.data()[node]["dataset_id"],
                "state_id": self._state_id.get(node),
                "preproc_args": preproc_args,
            }
        )

    with self._reqs.send(requests, self._nodes, self._policies) as responses:
        errors: Dict[str, ErrorMessage] = responses.errors()
        replies: Dict[str, PreprocReply] = responses.replies()

    if errors:
        # Handle errors appropriately (logging, raising exceptions, etc.)
        for node_id, error in errors.items():
            logger.error(
                "Error message received during preprocessing "
                f"in node_id={node_id}: {error.errnum}. {error.extra_msg}"
            )

    logger.debug(f"Replies are: {replies} for step {self._preproc_step}")

    return replies

def execute(
    self,
) -> Tuple[
    Dict[str, Dict[str, Any]],  # inner dicts are TrainReply dumps
    Union[Dict[str, Dict[str, AuxVar]], Dict[str, EncryptedAuxVar]],
]:
    """Sends training request to nodes and waits for the responses

    Returns:
        A tuple of
          * training replies for this round
          * node-wise optimizer auxiliary variables, as a dict with format
            `{node_name: encrypted_aux_var}` is secagg is used, and
            `{node_name: {module_name: module_aux_var}}` otherwise.
    """

    # Populate request message
    msg = {
        "researcher_id": self._researcher_id,
        "experiment_id": self._experiment_id,
        "training_args": self._training_args.dict(),
        "training": self._do_training,
        "model_args": self._model_args if self._model_args is not None else {},
        "round": self._round_,
        "training_plan": self._training_plan.source(),
        "training_plan_class": self._training_plan.__class__.__name__,
        "params": self._training_plan.get_model_params(
            exclude_buffers=not self._training_args.dict()[
                "share_persistent_buffers"
            ],
            local_params=self._training_plan.local_params,
        ),
        "secagg_arguments": self._secagg_arguments,
        "aggregator_args": {},
        "optim_aux_var": self._optim_aux_var,
    }

    # Loop over nodes, add node specific data and send train request
    messages = MessagesByNode()

    for node in self._nodes:
        msg["dataset_id"] = self._data.data()[node]["dataset_id"]
        msg["state_id"] = self._nodes_state_ids.get(node)

        # add aggregator parameters to message header
        msg["aggregator_args"] = (
            self._aggregator_args.get(node, {}) if self._aggregator_args else {}
        )

        self._log_round_info(node=node, training=self._do_training)

        messages.update({node: TrainRequest(**msg)})  # send request to node

    with self.RequestTimer(self._nodes) as timer:  # compute request time
        # Send training request
        with self._reqs.send(
            messages, self._nodes, self._policies
        ) as federated_req:
            errors = federated_req.errors()
            replies = federated_req.replies()

    # Loop over errors
    for node_id, error in errors.items():
        logger.warning(
            f"Error message received during training: {error.errnum}. {error.extra_msg}"
        )
        self._nodes.remove(node_id)

    training_replies = self._get_training_results(replies=replies)

    timing_results = self._get_timing_results(replies, timer)
    # `training_replies` can be empty if there wasnot any replies
    for node_id in replies:
        if training_replies.get(node_id):
            training_replies[node_id].update({"timing": timing_results[node_id]})

    # Extract aux variables from training replies.
    if self._do_training:
        aux_vars = self._extract_received_optimizer_aux_var_from_round(replies)
    else:
        aux_vars = {}
    return training_replies, aux_vars

def execute(self) -> Dict:
    """Requests the approval of the provided TrainingPlan.

    Returns:
        a dictionary of pairs (node_id: status), where status indicates to the researcher
        that the training plan has been correctly downloaded on the node side.
        Warning: status does not mean that the training plan is approved, only that it has been added
        to the "approval queue" on the node side.
    """
    return self._reqs.training_plan_approve(
        self._training_plan, self._description, self._nodes, self._policies
    )

def execute(self) -> Dict:
    """Checks whether model is approved or not.

    This method sends `training-plan-status` request to the nodes. It should be run before running experiment.
    So, researchers can find out if their model has been approved

    Returns:
        A dict of `Message` objects indexed by node ID, one for each job's nodes
    """

    message = TrainingPlanStatusRequest(
        **{
            "researcher_id": self._researcher_id,
            "experiment_id": self._experiment_id,
            "training_plan": self._training_plan.source(),
        }
    )

    # Send message to each node that has been found after dataset search request
    # TODO: add timer to compute request time
    with self._reqs.send(
        message, self._nodes, policies=self._policies
    ) as federated_req:
        replies = federated_req.replies()

        for node_id, reply in replies.items():
            if reply.success is True:
                if reply.approval_obligation is True:
                    if reply.status == TrainingPlanApprovalStatus.APPROVED.value:
                        logger.info(
                            f"Training plan has been approved by the node: {node_id}"
                        )
                    else:
                        logger.warning(
                            f"Training plan has NOT been approved by the node: {node_id}."
                            + f"Training plan status : {reply.status}"
                        )
                else:
                    logger.info(
                        f"Training plan approval is not required by the node: {node_id}"
                    )
            else:
                logger.warning(f"Node : {node_id} : {reply.msg}")

    # Get the nodes that haven't replied training-plan-status request
    non_replied_nodes = list(set(self._nodes) - set(replies.keys()))
    if non_replied_nodes:
        logger.warning(
            f"Request for checking training plan status hasn't been replied \
                         by the nodes: {non_replied_nodes}. You might get error \
                             while running your experiment. "
        )

    return replies