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, arg_name)

                    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, to_params=True)
                                                          }
        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 + str(uuid.uuid4()) + '.pt')

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

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

    # 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=str(e))
    except Exception as e:
        msg = 'Unexpected error while validating training argument'
        logger.debug(f"{msg}: {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, 'my_model_')
            if success:
                # retrieving arggegator 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: {str(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: {str(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: {e}"
        return self._send_round_reply(success=False, message=error_message)

    # import model params into the training plan instance
    try:

        self.training_plan.load(params_path, to_params=False)
    except Exception as e:
        error_message = f"Cannot initialize model parameters: f{str(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: {str(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: {str(e)}"
        return self._send_round_reply(success=False, message=error_message)

    # training_kwargs_with_history = dict(history_monitor=self.history_monitor,
    #                                     node_args=self.node_args,
    #                                     aggregator_args=self.aggregator_args)
    # training_kwargs_print = {key:value for key, value in training_kwargs_with_history.items() if key != 'aggregator_args'}
    # logger.info(f'training with arguments {training_kwargs_print}')

    # 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"{str(e)}")
            except Exception as e:
                logger.error(f"Undetermined error during the testing phase on global parameter updates: "
                             f"{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: {str(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"{str(e)}")
                except Exception as e:
                    logger.error(f"Undetermined error during the validation phase on local parameter updates"
                                 f"{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")

        # Upload results
        results['researcher_id'] = self.researcher_id
        results['job_id'] = self.job_id
        results['model_params'] = self.training_plan.after_training_params()
        results['node_id'] = environ['NODE_ID']
        results['optimizer_args'] = self.training_plan.optimizer_args()
        try:
            # TODO : should validation status code but not yet returned
            # by upload_file
            filename = os.path.join(environ['TMP_DIR'], 'node_params_' + str(uuid.uuid4()) + '.pt')
            self.training_plan.save(filename, results)
            res = self.repository.upload_file(filename)
            logger.info("results uploaded successfully ")


        except Exception as e:
            is_failed = True
            error_message = f"Cannot upload results: {str(e)}"
            return self._send_round_reply(success=False, message=error_message)

        # 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: {e}')
            pass

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