fedbiomed.node.round

def download_aggregator_args(self) -> Tuple[bool, str]:
    """Retrieves aggregator arguments, that are sent through file exchange service

    Returns:
        Tuple[bool, str]: a tuple containing:
            a bool that indicates the success of operation
            a string containing the error message
    """
    # download heavy aggregator args (if any)

    if self.aggregator_args is not None:

        for arg_name, aggregator_arg in self.aggregator_args.items():
            if isinstance(aggregator_arg, dict):
                url = aggregator_arg.get('url', False)

                if any((url, arg_name)):
                    # if both `filename` and `arg_name` fields are present, it means that parameters
                    # should be retrieved using file
                    # exchanged system
                    success, param_path, error_msg = self.download_file(url, f"{arg_name}_{uuid.uuid4()}.mpk")
                    if not success:
                        return success, error_msg
                    else:
                        # FIXME: should we load parameters here or in the training plan
                        self.aggregator_args[arg_name] = {'param_path': param_path,
                                                          # 'params': training_plan.load(param_path,
                                                          # update_model=True)
                                                          }
                    self.aggregator_args[arg_name] = Serializer.load(param_path)
        return True, ''
    else:
        return True, "no file downloads required for aggregator args"

def download_file(self, url: str, file_path: str) -> Tuple[bool, str, str]:
    """Downloads file from file exchange system

    Args:
        url (str): url used to download file
        file_path (str): file path used to store the downloaded content

    Returns:
        Tuple[bool, str, str]: tuple that contains:
            bool that indicates the success of the download
            str that returns the complete path file
            str containing the error message (if any). Returns empty
            string if operation successful.
    """

    status, params_path = self.repository.download_file(url, file_path)

    if (status != 200) or params_path is None:

        error_message = f"Cannot download param file: {url}"
        return False, '', error_message
    else:
        return True, params_path, ''

def initialize_validate_training_arguments(self) -> None:
    """Validates and separates training argument for experiment round"""

    self.training_arguments = TrainingArgs(self.training_kwargs, only_required=False)
    self.testing_arguments = self.training_arguments.testing_arguments()
    self.loader_arguments = self.training_arguments.loader_arguments()

def run_model_training(
        self,
        secagg_arguments: Union[Dict, None] = None,
) -> Dict[str, Any]:
    """This method downloads training plan file; then runs the training of a model
    and finally uploads model params to the file repository

    Args:
        secagg_arguments:
            - secagg_servkey_id: Secure aggregation Servkey context id. None means that the parameters
                are not going to be encrypted
            - secagg_biprime_id: Secure aggregation Biprime context ID.
            - secagg_random: Float value to validate secure aggregation on the researcher side

    Returns:
        Returns the corresponding node message, training reply instance
    """
    is_failed = False

    # Validate secagg status. Raises error if the training request is compatible with
    # secure aggregation settings

    secagg_arguments = {} if secagg_arguments is None else secagg_arguments
    self._use_secagg = self._configure_secagg(
        secagg_servkey_id=secagg_arguments.get('secagg_servkey_id'),
        secagg_biprime_id=secagg_arguments.get('secagg_biprime_id'),
        secagg_random=secagg_arguments.get('secagg_random')
    )

    # Initialize and validate requested experiment/training arguments
    try:
        self.initialize_validate_training_arguments()
    except FedbiomedUserInputError as e:
        return self._send_round_reply(success=False, message=repr(e))
    except Exception as e:
        msg = 'Unexpected error while validating training argument'
        logger.debug(f"{msg}: {repr(e)}")
        return self._send_round_reply(success=False, message=f'{msg}. Please contact system provider')
    try:
        # module name cannot contain dashes
        import_module = 'training_plan_' + str(uuid.uuid4().hex)
        status, _ = self.repository.download_file(self.training_plan_url,
                                                  import_module + '.py')

        if status != 200:
            error_message = "Cannot download training plan file: " + self.training_plan_url
            return self._send_round_reply(success=False, message=error_message)
        else:
            if environ["TRAINING_PLAN_APPROVAL"]:
                approved, training_plan_ = self.tp_security_manager.check_training_plan_status(
                    os.path.join(environ["TMP_DIR"], import_module + '.py'),
                    TrainingPlanApprovalStatus.APPROVED)

                if not approved:
                    error_message = f'Requested training plan is not approved by the node: {environ["NODE_ID"]}'
                    return self._send_round_reply(success=False, message=error_message)
                else:
                    logger.info(f'Training plan has been approved by the node {training_plan_["name"]}')

        if not is_failed:

            success, params_path, error_msg = self.download_file(self.params_url, f"my_model_{uuid.uuid4()}.mpk")
            if success:
                # retrieving aggregator args
                success, error_msg = self.download_aggregator_args()

            if not success:
                return self._send_round_reply(success=False, message=error_msg)

    except Exception as e:
        is_failed = True
        # FIXME: this will trigger if model is not approved by node
        error_message = f"Cannot download training plan files: {repr(e)}"
        return self._send_round_reply(success=False, message=error_message)

    # import module, declare the training plan, load parameters
    try:
        sys.path.insert(0, environ['TMP_DIR'])
        module = importlib.import_module(import_module)
        train_class = getattr(module, self.training_plan_class)
        self.training_plan = train_class()
        sys.path.pop(0)
    except Exception as e:
        error_message = f"Cannot instantiate training plan object: {repr(e)}"
        return self._send_round_reply(success=False, message=error_message)

    try:
        self.training_plan.post_init(model_args=self.model_arguments,
                                     training_args=self.training_arguments,
                                     aggregator_args=self.aggregator_args)
    except Exception as e:
        error_message = f"Can't initialize training plan with the arguments: {repr(e)}"
        return self._send_round_reply(success=False, message=error_message)

    # import model params into the training plan instance
    try:
        params = Serializer.load(params_path)["model_weights"]
        self.training_plan.set_model_params(params)
    except Exception as e:
        error_message = f"Cannot initialize model parameters: {repr(e)}"
        return self._send_round_reply(success=False, message=error_message)

    # Split training and validation data
    try:
        self._set_training_testing_data_loaders()
    except FedbiomedError as e:
        error_message = f"Can not create validation/train data: {repr(e)}"
        return self._send_round_reply(success=False, message=error_message)
    except Exception as e:
        error_message = f"Undetermined error while creating data for training/validation. Can not create " \
                        f"validation/train data: {repr(e)}"
        return self._send_round_reply(success=False, message=error_message)

    # Validation Before Training
    if self.testing_arguments.get('test_on_global_updates', False) is not False:

        # Last control to make sure validation data loader is set.
        if self.training_plan.testing_data_loader is not None:
            try:
                self.training_plan.testing_routine(metric=self.testing_arguments.get('test_metric', None),
                                                   metric_args=self.testing_arguments.get('test_metric_args', {}),
                                                   history_monitor=self.history_monitor,
                                                   before_train=True)
            except FedbiomedError as e:
                logger.error(f"{ErrorNumbers.FB314}: During the validation phase on global parameter updates; "
                             f"{repr(e)}")
            except Exception as e:
                logger.error(f"Undetermined error during the testing phase on global parameter updates: "
                             f"{repr(e)}")
        else:
            logger.error(f"{ErrorNumbers.FB314}: Can not execute validation routine due to missing testing dataset"
                         f"Please make sure that `test_ratio` has been set correctly")

    # If training is activated.
    if self.training:
        if self.training_plan.training_data_loader is not None:
            try:
                results = {}
                rtime_before = time.perf_counter()
                ptime_before = time.process_time()
                self.training_plan.training_routine(history_monitor=self.history_monitor,
                                                    node_args=self.node_args)
                rtime_after = time.perf_counter()
                ptime_after = time.process_time()
            except Exception as e:
                error_message = f"Cannot train model in round: {repr(e)}"
                return self._send_round_reply(success=False, message=error_message)

        # Validation after training
        if self.testing_arguments.get('test_on_local_updates', False) is not False:

            if self.training_plan.testing_data_loader is not None:
                try:
                    self.training_plan.testing_routine(metric=self.testing_arguments.get('test_metric', None),
                                                       metric_args=self.testing_arguments.get('test_metric_args',
                                                                                              {}),
                                                       history_monitor=self.history_monitor,
                                                       before_train=False)
                except FedbiomedError as e:
                    logger.error(
                        f"{ErrorNumbers.FB314.value}: During the validation phase on local parameter updates; "
                        f"{repr(e)}")
                except Exception as e:
                    logger.error(f"Undetermined error during the validation phase on local parameter updates"
                                 f"{repr(e)}")
            else:
                logger.error(
                    f"{ErrorNumbers.FB314.value}: Can not execute validation routine due to missing testing "
                    f"dataset please make sure that test_ratio has been set correctly")

        sample_size = len(self.training_plan.training_data_loader.dataset)

        results["encrypted"] = False
        model_weights = self.training_plan.after_training_params(flatten=self._use_secagg)
        if self._use_secagg:
            logger.info("Encrypting model parameters. This process can take some time depending on model size.")

            encrypt = functools.partial(
                self._secagg_crypter.encrypt,
                num_nodes=len(self._servkey["parties"]) - 1,  # -1: don't count researcher
                current_round=self._round,
                key=self._servkey["context"]["server_key"],
                biprime=self._biprime["context"]["biprime"],
                weight=sample_size,
                clipping_range=secagg_arguments.get('secagg_clipping_range')
            )
            model_weights = encrypt(params=model_weights)
            results["encrypted"] = True
            results["encryption_factor"] = encrypt(params=[secagg_arguments["secagg_random"]])
            logger.info("Encryption is completed!")

        results['researcher_id'] = self.researcher_id
        results['job_id'] = self.job_id
        results['model_weights'] = model_weights
        results['node_id'] = environ['NODE_ID']
        results['optimizer_args'] = self.training_plan.optimizer_args()

        try:
            # TODO: add validation status to these results?
            # Dump the results to a msgpack file.
            filename = os.path.join(environ["TMP_DIR"], f"node_params_{uuid.uuid4()}.mpk")
            Serializer.dump(results, filename)
            # Upload that file to the remote repository.
            res = self.repository.upload_file(filename)
            logger.info("results uploaded successfully ")
        except Exception as exc:
            return self._send_round_reply(success=False, message=f"Cannot upload results: {exc}")

        # end : clean the namespace
        try:
            del self.training_plan
            del import_module
        except Exception as e:
            logger.debug(f'Exception raise while deleting training plan instance: {repr(e)}')

        return self._send_round_reply(success=True,
                                      timing={'rtime_training': rtime_after - rtime_before,
                                              'ptime_training': ptime_after - ptime_before},
                                      params_url=res['file'],
                                      sample_size=sample_size)
    else:
        # Only for validation
        return self._send_round_reply(success=True)