Jobs

@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) -> Tuple[Dict, Optional[Dict]]:
    """ Sends training request to nodes and waits for the responses

    Returns:
        A tuple of
          * training replies for this round
          * Dict of auxiliary variables, collating node-wise information, with
            format `{mod_name: {node_id: node_dict}}`.
    """

    # 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']),
        'secagg_servkey_id': self._secagg_arguments.get('secagg_servkey_id'),
        'secagg_biprime_id': self._secagg_arguments.get('secagg_biprime_id'),
        'secagg_random': self._secagg_arguments.get('secagg_random'),
        'secagg_clipping_range': self._secagg_arguments.get('secagg_clipping_range'),
        'command': 'train',
        'aggregator_args': {},
    }

    # Prepare optimizer auxiliary variables, if any.
    if self._do_training and self._optim_aux_var:
        aux_shared, aux_bynode = (
            self._prepare_agg_optimizer_aux_var(self._optim_aux_var, nodes=list(self._nodes))
        )
    else:
        aux_shared = {}
        aux_bynode = {}

    # 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['aux_vars'] = [aux_shared, aux_bynode.get(node, None)]
        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()

    training_replies = self._get_training_results(replies=replies,
                                                  errors=errors)

    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
    aux_vars = None
    if self._do_training:
        aux_vars = self._extract_received_optimizer_aux_var_from_round(training_replies)

    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(),
        'command': 'training-plan-status'
    })

    # 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