fedbiomed.common.training_plans

def add_dependency(self, dep: List[str]) -> None:
    """Add new dependencies to the TrainingPlan.

    These dependencies are used while creating a python module.

    Args:
        dep: Dependencies to add. Dependencies should be indicated as
            import statement strings, e.g. `"from torch import nn"`.
    """
    for val in dep:
        if val not in self._dependencies:
            self._dependencies.append(val)

def add_preprocess(
        self,
        method: Callable,
        process_type: ProcessTypes
    ) -> None:
    """Register a pre-processing method to be executed on training data.

    Args:
        method: Pre-processing method to be run before training.
        process_type: Type of pre-processing that will be run.
            The expected signature of `method` and the arguments
            passed to it depend on this parameter.
    """
    if not callable(method):
        msg = (
            f"{ErrorNumbers.FB605.value}: error while adding "
            "preprocess, `method` should be callable."
        )
        logger.critical(msg)
        raise FedbiomedTrainingPlanError(msg)
    if not isinstance(process_type, ProcessTypes):
        msg = (
            f"{ErrorNumbers.FB605.value}: error while adding "
            "preprocess, `process_type` should be an instance "
            "of `fedbiomed.common.constants.ProcessTypes`."
        )
        logger.critical(msg)
        raise FedbiomedTrainingPlanError(msg)
    # NOTE: this may be revised into a list rather than OrderedDict
    self.pre_processes[method.__name__] = {
        'method': method,
        'process_type': process_type
    }

def after_training_params(
        self,
        flatten: bool = False
) -> Union[Dict[str, Any], List[float]]:
    """Return the wrapped model's parameters for aggregation.

    This method returns a dict containing parameters that need to be
    reported back and aggregated in a federated learning setting.

    It may also implement post-processing steps to make these parameters
    suitable for sharing with the researcher after training - hence its
    being used over `get_model_params` at the end of training rounds.

    Returns:
        The trained parameters to aggregate.
    """

    # Get flatten model parameters
    if flatten:
        return self._model.flatten()

    return self.get_model_params()

def export_model(self, filename: str) -> None:
    """Export the wrapped model to a dump file.

    Args:
        filename: path to the file where the model will be saved.

    !!! info "Notes":
        This method is designed to save the model to a local dump
        file for easy re-use by the same user, possibly outside of
        Fed-BioMed. It is not designed to produce trustworthy data
        dumps and is not used to exchange models and their weights
        as part of the federated learning process.

        To save the model parameters for sharing as part of the FL process,
        use the `after_training_params` method (or `get_model_params` one
        outside of a training context) and export results using
        [`Serializer`][fedbiomed.common.serializer.Serializer].
    """
    self._model.export(filename)

def get_learning_rate(self) -> List[float]:
    raise FedbiomedTrainingPlanError("method not implemented")

def get_model_params(self) -> Dict[str, Any]:
    """Return a copy of the model's trainable weights.

    The type of data structure used to store weights depends on the actual
    framework of the wrapped model.

    Returns:
        Model weights, as a dictionary mapping parameters' names to their value.
    """
    return self._model.get_weights()

def import_model(self, filename: str) -> None:
    """Import and replace the wrapped model from a dump file.

    Args:
        filename: path to the file where the model has been exported.

    !!! info "Notes":
        This method is designed to load the model from a local dump
        file, that might not be in a trustworthy format. It should
        therefore only be used to re-load data exported locally and
        not received from someone else, including other FL peers.

        To load model parameters shared as part of the FL process, use the
        [`Serializer`][fedbiomed.common.serializer.Serializer] to read the
        network-exchanged file, and the `set_model_params` method to assign
        the loaded values into the wrapped model.
    """
    try:
        self._model.reload(filename)
    except FedbiomedModelError as exc:
        msg = (
            f"{ErrorNumbers.FB304.value}: failed to import a model from "
            f"a dump file: {exc}"
        )
        logger.critical(msg)
        raise FedbiomedTrainingPlanError(msg) from exc

def init_dependencies(self) -> List[str]:
    """Default method where dependencies are returned

    Returns:
        Empty list as default
    """
    return []

@abstractmethod
def model(self):
    """Gets model instance of the training plan"""

def optimizer_args(self) -> Dict:
    """Retrieves optimizer arguments (to be overridden
    by children classes)

    Returns:
        Empty dictionary: (to be overridden in children classes)
    """
    return {}

@abstractmethod
def post_init(
        self,
        model_args: Dict[str, Any],
        training_args: Dict[str, Any],
        aggregator_args: Optional[Dict[str, Any]] = None,
    ) -> None:
    """Process model, training and optimizer arguments.

    Args:
        model_args: Arguments defined to instantiate the wrapped model.
        training_args: Arguments that are used in training routines
            such as epoch, dry_run etc.
            Please see [`TrainingArgs`][fedbiomed.common.training_args.TrainingArgs]
        aggregator_args: Arguments managed by and shared with the
            researcher-side aggregator.
    """

def save_code(self, filepath: str) -> None:
    """Saves the class source/codes of the training plan class that is created byuser.

    Args:
        filepath: path to the destination file

    Raises:
        FedbiomedTrainingPlanError: raised when source of the model class cannot be assessed
        FedbiomedTrainingPlanError: raised when model file cannot be created/opened/edited
    """
    try:
        class_source = get_class_source(self.__class__)
    except FedbiomedError as exc:
        msg = f"{ErrorNumbers.FB605.value}: error while getting source of the model class: {exc}"
        logger.critical(msg)
        raise FedbiomedTrainingPlanError(msg) from exc

    # Preparing content of the module
    content = "\n".join(self._dependencies)
    content += "\n"
    content += class_source
    try:
        # should we write it in binary (for the sake of space optimization)?
        with open(filepath, "w", encoding="utf-8") as file:
            file.write(content)
        logger.debug("Model file has been saved: " + filepath)
    except PermissionError as exc:
        _msg = ErrorNumbers.FB605.value + f" : Unable to read {filepath} due to unsatisfactory privileges" + \
               ", can't write the model content into it"
        logger.critical(_msg)
        raise FedbiomedTrainingPlanError(_msg) from exc
    except MemoryError as exc:
        _msg = ErrorNumbers.FB605.value + f" : Can't write model file on {filepath}: out of memory!"
        logger.critical(_msg)
        raise FedbiomedTrainingPlanError(_msg) from exc
    except OSError as exc:
        _msg = ErrorNumbers.FB605.value + f" : Can't open file {filepath} to write model content"
        logger.critical(_msg)
        raise FedbiomedTrainingPlanError(_msg) from exc

def set_aggregator_args(self, aggregator_args: Dict[str, Any]):
    raise FedbiomedTrainingPlanError("method not implemented and needed")

def set_data_loaders(
        self,
        train_data_loader: Union[DataLoader, NPDataLoader, None],
        test_data_loader: Union[DataLoader, NPDataLoader, None]
    ) -> None:
    """Sets data loaders

    Args:
        train_data_loader: Data loader for training routine/loop
        test_data_loader: Data loader for validation routine
    """
    self.training_data_loader = train_data_loader
    self.testing_data_loader = test_data_loader

def set_dataset_path(self, dataset_path: str) -> None:
    """Dataset path setter for TrainingPlan

    Args:
        dataset_path: The path where data is saved on the node.
            This method is called by the node that executes the training.
    """
    self.dataset_path = dataset_path
    logger.debug(f"Dataset path has been set as {self.dataset_path}")

def set_model_params(self, params: Dict[str, Any]) -> None:
    """Assign new values to the model's trainable parameters.

    The type of data structure used to store weights depends on the actual
    framework of the wrapped model.

    Args:
        params: model weights, as a dictionary mapping parameters' names
            to their value.
    """
    return self._model.set_weights(params)

def testing_routine(
        self,
        metric: Optional[MetricTypes],
        metric_args: Dict[str, Any],
        history_monitor: Optional['HistoryMonitor'],
        before_train: bool
    ) -> None:
    """Evaluation routine, to be called once per round.

    !!! info "Note"
        If the training plan implements a `testing_step` method
        (the signature of which is func(data, target) -> metrics)
        then it will be used rather than the input metric.

    Args:
        metric: The metric used for validation.
            If None, use MetricTypes.ACCURACY.
        history_monitor: HistoryMonitor instance,
            used to record computed metrics and communicate them to
            the researcher (server).
        before_train: Whether the evaluation is being performed
            before local training occurs, of afterwards. This is merely
            reported back through `history_monitor`.
    """
    # TODO: Add preprocess option for testing_data_loader.
    if self.testing_data_loader is None:
        msg = f"{ErrorNumbers.FB605.value}: no validation dataset was set."
        logger.critical(msg)
        raise FedbiomedTrainingPlanError(msg)

    n_batches = len(self.testing_data_loader)
    n_samples = len(self.testing_data_loader.dataset)
    # Set up a batch-wise metrics-computation function.
    # Either use an optionally-implemented custom training routine.
    if hasattr(self, "testing_step"):
        evaluate = getattr(self, "testing_step")
        metric_name = "Custom"
    # Or use the provided `metric` (or its default value).
    else:
        if metric is None:
            metric = MetricTypes.ACCURACY
        metric_controller = Metrics()
        def evaluate(data, target):
            nonlocal metric, metric_args, metric_controller
            output = self._model.predict(data)
            if isinstance(target, torch.Tensor):
                target = target.numpy()
            return metric_controller.evaluate(
                target, output, metric=metric, **metric_args
            )
        metric_name = metric.name
    # Iterate over the validation dataset and run the defined routine.
    num_samples_observed_till_now: int = 0
    for idx, (data, target) in enumerate(self.testing_data_loader, 1):
        num_samples_observed_till_now += self._infer_batch_size(data)
        # Run the evaluation step; catch and raise exceptions.
        try:
            m_value = evaluate(data, target)
        except Exception as exc:
            msg = (
                f"{ErrorNumbers.FB605.value}: An error occurred "
                f"while computing the {metric_name} metric: {exc}"
            )
            logger.critical(msg)
            raise FedbiomedTrainingPlanError(msg) from exc
        # Log the computed value.
        logger.debug(
            f"Validation: Batch {idx}/{n_batches} "
            f"| Samples {num_samples_observed_till_now}/{n_samples} "
            f"| Metric[{metric_name}]: {m_value}"
        )
        # Further parse, and report it (provided a monitor is set).
        if history_monitor is not None:
            m_dict = self._create_metric_result_dict(m_value, metric_name)
            history_monitor.add_scalar(
                metric=m_dict,
                iteration=idx,
                epoch=None,
                test=True,
                test_on_local_updates=(not before_train),
                test_on_global_updates=before_train,
                total_samples=n_samples,
                batch_samples=num_samples_observed_till_now,
                num_batches=n_batches
            )

def training_data(self):
    """All subclasses must provide a training_data routine the purpose of this actual code is to detect
    that it has been provided

    Raises:
        FedbiomedTrainingPlanError: if called and not inherited
    """
    msg = f"{ErrorNumbers.FB303.value}: training_data must be implemented"
    logger.critical(msg)
    raise FedbiomedTrainingPlanError(msg)

@abstractmethod
def training_routine(
        self,
        history_monitor: Optional['HistoryMonitor'] = None,
        node_args: Optional[Dict[str, Any]] = None
    ) -> None:
    """Training routine, to be called once per round.

    Args:
        history_monitor: optional HistoryMonitor
            instance, recording training metadata.
        node_args: Command line arguments for node.
            These arguments can specify GPU use; however, this is not
            supported for scikit-learn models and thus will be ignored.
    """
    return None

def post_init(
        self,
        model_args: Dict[str, Any],
        training_args: Dict[str, Any],
        aggregator_args: Optional[Dict[str, Any]] = None,
    ) -> None:
    # get default values of Perceptron model (different from SGDClassifier model default values)
    perceptron_default_values = Perceptron().get_params()
    sgd_classifier_default_values = SGDClassifier().get_params()

    model_args["loss"] = "perceptron"
    super().post_init(model_args, training_args)

    # collect default values of Perceptron and set it to the model FedPerceptron
    model_hyperparameters = self._model.get_params()
    for hyperparameter_name, val in perceptron_default_values.items():
        if model_hyperparameters[hyperparameter_name] == sgd_classifier_default_values[hyperparameter_name]:
            # this means default parameter of SGDClassifier has not been changed by user
            self._model.set_params(**{hyperparameter_name: val})

def model(self) -> BaseEstimator:
    """Retrieve the wrapped scikit-learn model instance.

    Returns:
        Scikit-learn model instance
    """
    return self._model.model

def model_args(self) -> Dict[str, Any]:
    """Retrieve model arguments.

    Returns:
        Model arguments
    """
    return self._model.model_args

def post_init(
        self,
        model_args: Dict[str, Any],
        training_args: Dict[str, Any],
        aggregator_args: Optional[Dict[str, Any]] = None,
) -> None:
    """Process model, training and optimizer arguments.

    Args:
        model_args: Arguments defined to instantiate the wrapped model.
        training_args: Arguments that are used in training routines
            such as epoch, dry_run etc.
            Please see [`TrainingArgs`][fedbiomed.common.training_args.TrainingArgs]
        aggregator_args: Arguments managed by and shared with the
            researcher-side aggregator.
    """
    model_args.setdefault("verbose", 1)
    self._model.model_args = model_args
    self._aggregator_args = aggregator_args or {}

    self._training_args = training_args.pure_training_arguments()
    self._batch_maxnum = self._training_args.get('batch_maxnum', self._batch_maxnum)
    # Add dependencies
    self._configure_dependencies()
    # Override default model parameters based on `self._model_args`.
    params = {
        key: model_args.get(key, val)
        for key, val in self._model.get_params().items()
    }

    self._model.set_params(**params)
    # Set up additional parameters (normally created by `self._model.fit`).
    # TODO: raise error if
    self._model.set_init_params(model_args)

def set_data_loaders(
        self,
        train_data_loader: Union[DataLoader, NPDataLoader, None],
        test_data_loader: Union[DataLoader, NPDataLoader, None]
) -> None:
    """Sets data loaders

    Args:
        train_data_loader: Data loader for training routine/loop
        test_data_loader: Data loader for validation routine
    """
    args = (train_data_loader, test_data_loader)
    if not all(isinstance(data, NPDataLoader) for data in args):
        msg = (
            f"{ErrorNumbers.FB310.value}: SKLearnTrainingPlan expects "
            "NPDataLoader instances as training and testing data "
            f"loaders, but received {type(train_data_loader)} "
            f"and {type(test_data_loader)} respectively."
        )
        logger.error(msg)
        raise FedbiomedTrainingPlanError(msg)
    self.training_data_loader = train_data_loader
    self.testing_data_loader = test_data_loader

def testing_routine(
        self,
        metric: Optional[MetricTypes],
        metric_args: Dict[str, Any],
        history_monitor: Optional['HistoryMonitor'],
        before_train: bool
) -> None:
    """Evaluation routine, to be called once per round.

    !!! info "Note"
        If the training plan implements a `testing_step` method
        (the signature of which is func(data, target) -> metrics)
        then it will be used rather than the input metric.

    Args:
        metric: The metric used for validation.
            If None, use MetricTypes.ACCURACY.
        history_monitor: HistoryMonitor instance,
            used to record computed metrics and communicate them to
            the researcher (server).
        before_train: Whether the evaluation is being performed
            before local training occurs, of afterwards. This is merely
            reported back through `history_monitor`.
    """
    # Check that the testing data loader is of proper type.
    if not isinstance(self.testing_data_loader, NPDataLoader):
        msg = (
            f"{ErrorNumbers.FB310.value}: SKLearnTrainingPlan cannot be "
            "evaluated without a NPDataLoader as `testing_data_loader`."
        )
        logger.error(msg)
        raise FedbiomedTrainingPlanError(msg)
    # If required, make up for the lack of specifications regarding target
    # classification labels.
    if self._model.is_classification and not hasattr(self.model(), 'classes_'):
        classes = self._classes_from_concatenated_train_test()
        setattr(self.model(), 'classes_', classes)
    # If required, select the default metric (accuracy or mse).
    if metric is None:
        if self._model.is_classification:
            metric = MetricTypes.ACCURACY
        else:
            metric = MetricTypes.MEAN_SQUARE_ERROR
    # Delegate the actual evalation routine to the parent class.
    super().testing_routine(
        metric, metric_args, history_monitor, before_train
    )

def training_args(self) -> Dict[str, Any]:
    """Retrieve training arguments.

    Returns:
        Training arguments
    """
    return self._training_args

def training_routine(
        self,
        history_monitor: Optional['HistoryMonitor'] = None,
        node_args: Optional[Dict[str, Any]] = None
) -> None:
    """Training routine, to be called once per round.

    Args:
        history_monitor: optional HistoryMonitor
            instance, recording training metadata. Defaults to None.
        node_args: command line arguments for node.
            These arguments can specify GPU use; however, this is not
            supported for scikit-learn models and thus will be ignored.
    """
    if self._model is None:
        raise FedbiomedTrainingPlanError('model is None')

    # Run preprocesses
    self._preprocess()

    if not isinstance(self.training_data_loader, NPDataLoader):
        msg = (
            f"{ErrorNumbers.FB310.value}: SKLearnTrainingPlan cannot "
            "be trained without a NPDataLoader as `training_data_loader`."
        )
        logger.critical(msg)
        raise FedbiomedTrainingPlanError(msg)
    # Run preprocessing operations.
    self._preprocess()
    # Warn if GPU-use was expected (as it is not supported).
    if node_args is not None and node_args.get('gpu_only', False):
        logger.warning(
            'Node would like to force GPU usage, but sklearn training '
            'plan does not support it. Training on CPU.'
        )
    # Run the model-specific training routine.
    try:
        return self._training_routine(history_monitor)
    except Exception as exc:
        msg = (
            f"{ErrorNumbers.FB605.value}: error while fitting "
            f"the model: {exc}"
        )
        logger.critical(msg)
        raise FedbiomedTrainingPlanError(msg)

def type(self) -> TrainingPlans:
    """Getter for training plan type """
    return self.__type

def after_training_params(self, flatten: bool = False) -> Dict[str, torch.Tensor]:
    """Return the wrapped model's parameters for aggregation.

    This method returns a dict containing parameters that need to be
    reported back and aggregated in a federated learning setting.

    If the `postprocess` method exists (i.e. has been defined by end-users)
    it is called in the context of this method. DP-required adjustments are
    also set to happen as part of this method.

    Returns:
        The trained parameters to aggregate.
    """
    params = super().after_training_params()
    # Check whether postprocess method exists, and use it.
    if hasattr(self, 'postprocess'):
        logger.debug("running model.postprocess() method")
        try:
            params = self.postprocess(params)
        except Exception as exc:
            raise FedbiomedTrainingPlanError(
                f"{ErrorNumbers.FB605.value}: Error while running post-"
                f"process: {exc}"
            ) from exc
    # Run (optional) DP controller adjustments as well.
    params = self._dp_controller.after_training(params)

    if flatten:
        params = self._model.flatten()

    return params

def get_learning_rate(self) -> List[float]:
    """Gets learning rate from value set in optimizer.

    !!! warning
        This function gathers the base learning rate applied to the model weights,
        including alterations due to any LR scheduler. However, it does not catch
        any adaptive component, e.g. due to RMSProp, Adam or such.

    Returns:
        List[float]: list of single learning rate or multiple learning rates
            (as many as the number of the layers contained in the model)
    """
    if self._optimizer is None:
        raise FedbiomedTrainingPlanError(f"{ErrorNumbers.FB605.value}: Optimizer not found, please call "
                                         f"`init_optimizer beforehand")
    learning_rates = []
    params = self._optimizer.param_groups
    for param in params:
        learning_rates.append(param['lr'])
    return learning_rates

@abstractmethod
def init_model(self):
    """Abstract method where model should be defined."""

def init_optimizer(self):
    """Abstract method for declaring optimizer by default """
    try:
        self._optimizer = torch.optim.Adam(self._model.model.parameters(), **self._optimizer_args)
    except AttributeError as e:
        raise FedbiomedTrainingPlanError(f"{ErrorNumbers.FB605.value}: Invalid argument for default "
                                         f"optimizer Adam. Error: {e}")

    return self._optimizer

def initial_parameters(self) -> Dict:
    """Returns initial parameters without DP or training applied

    Returns:
        State dictionary of torch Module
    """
    return self._model.init_params

def model(self) -> torch.nn.Module:
    """Return the torch model wrapped by this instance."""
    return self._model.model

def model_args(self) -> Dict:
    """Retrieves model args

    Returns:
        Model arguments arguments
    """
    return self._model.model_args

def optimizer(self):
    return self._optimizer

def optimizer_args(self) -> Dict:
    """Retrieves optimizer arguments

    Returns:
        Optimizer arguments
    """
    self.update_optimizer_args()  # update `optimizer_args` (eg after training)
    return self._optimizer_args

def post_init(
        self,
        model_args: Dict[str, Any],
        training_args: TrainingArgs,
        aggregator_args: Optional[Dict[str, Any]] = None,
) -> None:
    """Process model, training and optimizer arguments.

    Args:
        model_args: Arguments defined to instantiate the wrapped model.
        training_args: Arguments that are used in training routines
            such as epoch, dry_run etc.
            Please see [`TrainingArgs`][fedbiomed.common.training_args.TrainingArgs]
        aggregator_args: Arguments managed by and shared with the
            researcher-side aggregator.

    Raises:
        FedbiomedTrainingPlanError: If the provided arguments do not
            match expectations, or if the optimizer, model and dependencies
            configuration goes wrong.
    """
    # Assign scalar attributes.
    self._optimizer_args = training_args.optimizer_arguments() or {}
    self._training_args = training_args.pure_training_arguments()
    self._use_gpu = self._training_args.get('use_gpu')
    self._batch_maxnum = self._training_args.get('batch_maxnum')
    self._log_interval = self._training_args.get('log_interval')
    self._epochs = self._training_args.get('epochs')
    self._num_updates = self._training_args.get('num_updates', 1)
    self._dry_run = self._training_args.get('dry_run')
    # Optionally set up differential privacy.
    self._dp_controller = DPController(training_args.dp_arguments() or None)
    # Add dependencies
    self._configure_dependencies()
    # Configure aggregator-related arguments
    # TODO: put fedprox mu inside strategy_args
    self._fedprox_mu = self._training_args.get('fedprox_mu')
    self.set_aggregator_args(aggregator_args or {})
    # Configure model and optimizer
    self._configure_model_and_optimizer(model_args)

def send_to_device(self,
                   to_send: Union[torch.Tensor, list, tuple, dict],
                   device: torch.device
                   ):
    """Send inputs to correct device for training.

    Recursively traverses lists, tuples and dicts until it meets a torch Tensor, then sends the Tensor
    to the specified device.

    Args:
        to_send: the data to be sent to the device.
        device: the device to send the data to.

    Raises:
       FedbiomedTrainingPlanError: when to_send is not the correct type
    """
    if isinstance(to_send, torch.Tensor):
        return to_send.to(device)
    elif isinstance(to_send, dict):
        return {key: self.send_to_device(val, device) for key, val in to_send.items()}
    elif isinstance(to_send, tuple):
        return tuple(self.send_to_device(d, device) for d in to_send)
    elif isinstance(to_send, list):
        return [self.send_to_device(d, device) for d in to_send]
    else:
        raise FedbiomedTrainingPlanError(f'{ErrorNumbers.FB310.value} cannot send data to device. '
                                         f'Data must be a torch Tensor or a list, tuple or dict '
                                         f'ultimately containing Tensors.')

def set_aggregator_args(self, aggregator_args: Dict[str, Any]):
    """Handles and loads aggregators arguments sent through MQTT and
    file exchanged system. If sent through file exchanged system, loads the arguments.

    Args:
        aggregator_args (Dict[str, Any]): dictionary mapping aggregator argument name with its value (eg
        'aggregator_correction' with correction states)
    """
    self.aggregator_name = aggregator_args.get('aggregator_name', self.aggregator_name)
    # FIXME: this is too specific to Scaffold. Should be redesigned, or handled
    # by an aggregator handler that contains all keys for all strategies
    # implemented in fedbiomed
    # here we ae loading all args that have been sent from file exchange system
    for arg_name, aggregator_arg in aggregator_args.items():
        if arg_name == 'aggregator_correction':
            if not isinstance(aggregator_arg, dict):
                raise FedbiomedTrainingPlanError(
                    f"{ErrorNumbers.FB309.value}: TorchTrainingPlan received "
                    "invalid 'aggregator_correction' aggregator args."
                )
            self.correction_state = aggregator_arg

def testing_routine(
        self,
        metric: Optional[MetricTypes],
        metric_args: Dict[str, Any],
        history_monitor: Optional['HistoryMonitor'],
        before_train: bool
) -> None:
    """Evaluation routine, to be called once per round.

    !!! info "Note"
        If the training plan implements a `testing_step` method
        (the signature of which is func(data, target) -> metrics)
        then it will be used rather than the input metric.

    Args:
        metric: The metric used for validation.
            If None, use MetricTypes.ACCURACY.
        history_monitor: HistoryMonitor instance,
            used to record computed metrics and communicate them to
            the researcher (server).
        before_train: Whether the evaluation is being performed
            before local training occurs, of afterwards. This is merely
            reported back through `history_monitor`.
    """
    if not isinstance(self.model(), torch.nn.Module):
        msg = (
            f"{ErrorNumbers.FB320.value}: model should be a torch "
            f"nn.Module, but is of type {type(self.model())}"
        )
        logger.critical(msg)
        raise FedbiomedTrainingPlanError(msg)
    try:
        self.model().eval()  # pytorch switch for model inference-mode TODO should be removed
        with torch.no_grad():
            super().testing_routine(
                metric, metric_args, history_monitor, before_train
            )
    finally:
        self.model().train()  # restore training behaviors

def training_args(self) -> Dict:
    """Retrieves training args

    Returns:
        Training arguments
    """
    return self._training_args

@abstractmethod
def training_data(self):
    """Abstract method to return training data"""

def training_routine(self,
                     history_monitor: Any = None,
                     node_args: Union[dict, None] = None,
                     ) -> int:
    # FIXME: add betas parameters for ADAM solver + momentum for SGD
    # FIXME 2: remove parameters specific for validation specified in the
    # training routine
    """Training routine procedure.

    End-user should define;

    - a `training_data()` function defining how sampling / handling data in node's dataset is done. It should
        return a generator able to output tuple (batch_idx, (data, targets)) that is iterable for each batch.
    - a `training_step()` function defining how cost is computed. It should output loss values for backpropagation.

    Args:
        history_monitor: Monitor handler for real-time feed. Defined by the Node and can't be overwritten
        node_args: command line arguments for node. Can include:
            - `gpu (bool)`: propose use a GPU device if any is available. Default False.
            - `gpu_num (Union[int, None])`: if not None, use the specified GPU device instead of default
                GPU device if this GPU device is available. Default None.
            - `gpu_only (bool)`: force use of a GPU device if any available, even if researcher
                doesn't request for using a GPU. Default False.
    Returns:
        Total number of samples observed during the training.
    """

    #self.model().train()  # pytorch switch for training
    self._model.init_training()
    # set correct type for node args
    node_args = {} if not isinstance(node_args, dict) else node_args

    # send all model to device, ensures having all the requested tensors
    self._set_device(self._use_gpu, node_args)
    self._model.send_to_device(self._device)

    # Run preprocess when everything is ready before the training
    self._preprocess()

    # # initial aggregated model parameters
    # self._init_params = deepcopy(list(self.model().parameters()))

    # DP actions
    self._model.model, self._optimizer, self.training_data_loader = self._dp_controller.before_training(
        self.model(), self._optimizer, self.training_data_loader
    )

    # set number of training loop iterations
    iterations_accountant = MiniBatchTrainingIterationsAccountant(self)

    # Training loop iterations
    for epoch in iterations_accountant.iterate_epochs():
        training_data_iter: Iterator = iter(self.training_data_loader)

        for batch_idx in iterations_accountant.iterate_batches():
            # retrieve data and target
            data, target = next(training_data_iter)

            # update accounting for number of observed samples
            batch_size = self._infer_batch_size(data)
            iterations_accountant.increment_sample_counters(batch_size)

            # handle training on accelerator devices
            data, target = self.send_to_device(data, self._device), self.send_to_device(target, self._device)

            # train this batch
            corrected_loss, loss = self._train_over_batch(data, target)

            # Reporting
            if iterations_accountant.should_log_this_batch():
                # Retrieve reporting information: semantics differ whether num_updates or epochs were specified
                num_samples, num_samples_max = iterations_accountant.reporting_on_num_samples()
                num_iter, num_iter_max = iterations_accountant.reporting_on_num_iter()
                epoch_to_report = iterations_accountant.reporting_on_epoch()

                logger.debug('Train {}| '
                             'Iteration {}/{} | '
                             'Samples {}/{} ({:.0f}%)\tLoss: {:.6f}'.format(
                    f'Epoch: {epoch_to_report} ' if epoch_to_report is not None else '',
                    num_iter,
                    num_iter_max,
                    num_samples,
                    num_samples_max,
                    100. * num_iter / num_iter_max,
                    loss.item()))

                # Send scalar values via general/feedback topic
                if history_monitor is not None:
                    # the researcher only sees the average value of samples observed until now
                    history_monitor.add_scalar(metric={'Loss': loss.item()},
                                               iteration=num_iter,
                                               epoch=epoch_to_report,
                                               train=True,
                                               num_samples_trained=num_samples,
                                               num_batches=num_iter_max,
                                               total_samples=num_samples_max,
                                               batch_samples=batch_size)

            # Handle dry run mode
            if self._dry_run:
                self._model.to(self._device_init)
                torch.cuda.empty_cache()
                return iterations_accountant.num_samples_observed_in_total

    # release gpu usage as much as possible though:
    # - it should be done by deleting the object
    # - and some gpu memory remains used until process (cuda kernel ?) finishes

    self._model.send_to_device(self._device_init)
    torch.cuda.empty_cache()

    return iterations_accountant.num_samples_observed_in_total

@abstractmethod
def training_step(self):
    """Abstract method, all subclasses must provide a training_step."""

def type(self) -> TrainingPlans.TorchTrainingPlan:
    """ Gets training plan type"""
    return self.__type

def update_optimizer_args(self) -> Dict:
    """
    Updates `_optimizer_args` variable. Can prove useful
    to retrieve optimizer parameters after having trained a
    model, parameters which may have changed during training (eg learning rate).

    Updated arguments:
     - learning_rate

    Returns:
        Dict: updated `_optimizer_args`
    """
    if self._optimizer_args is None:
        self._optimizer_args = {}
    self._optimizer_args['lr'] = self.get_learning_rate()
    return self._optimizer_args