FedModelKit 0.5.0__py3-none-any.whl

FedModelKit/README.md

@@ -0,0 +1,25 @@
+### **Interface Directory**
+
+1. **`mlflow_interface.py`**
+    - This file defines the `FederatedModel` class, which encapsulates the local learner, the aggregation strategy, and their respective names. It facilitates the integration of these components into the MLflow tracking system. It also includes the `submit_fl_model` function for logging the model to MLflow.
+
+2. **`aggregator.py`**
+    - This file defines the `AggFactoryProtocol` and `AggProtocol` protocols. The `AggFactoryProtocol` is a protocol for the function creating the aggregator instance, while the `AggProtocol` defines the methods and attributes that an aggregator must implement.
+
+3. **`local_learner.py`**
+    - This file defines the `LLFactoryProtocol` and `LLProtocol` protocols. The `LLFactoryProtocol` is a protocol for the function creating the local learner instance, while the `LLProtocol` defines the methods and attributes that a local learner must implement. This file provides the structure for implementing custom local learners.
+
+4. **`default_create_functions.py`**
+    - This file provides default implementations for creating local learners and aggregators. The `default_create_local_learner` and `default_create_aggregator` functions.
+
+5. **`cli.py`**
+    - This file contains the `create_structure` function, which sets up the project structure for uploading models to MLflow. It creates necessary directories and files, including `utils.py`, `example.py`, and `README.md`.
+
+6. **`templates.py`**
+    - This file contains template functions for generating the content of `utils.py`, `example.py`, and `README.md`. It includes the `get_utils_template`, `get_main_template`, and `get_readme_template` functions.
+
+7. **`readme_template.txt`**
+    - This file contains the template for the `README.md` file that is generated by the `create_structure` function. It provides a detailed guide on how to use the package functionalities.
+
+8. **`src/utils.py`**
+    - This file contains utility functions for managing and converting Flower objects. It includes methods for serializing and deserializing model parameters and other helper functions.

FedModelKit/__init__.py

@@ -0,0 +1,17 @@
+'''from .interface import FederatedModel, submit_fl_model'''
+from .interface import FederatedModel
+from .local_learner import LLFactoryProtocol, LLProtocol
+from .aggregator import AggFactoryProtocol, AggProtocol
+from .default_create_functions import default_create_local_learner, default_create_aggregator
+from .cli import create_structure
+import flwr
+
+__all__ = ["FederatedModel",
+           "LLFactoryProtocol",
+           "LLProtocol",
+           "AggFactoryProtocol",
+           "AggProtocol",
+           "default_create_local_learner",
+           "default_create_aggregator",
+           "create_structure",
+           "flwr"]

FedModelKit/aggregator.py

@@ -0,0 +1,41 @@
+import flwr
+from typing import Protocol, runtime_checkable
+from typing import Protocol, Optional
+
+
+@runtime_checkable
+class AggProtocol(Protocol):
+    """
+    Defines the required structure for aggregation strategies in federated learning.
+    """
+
+    def aggregate_parameters(self, results: list[flwr.common.ParametersRecord], config: Optional[flwr.common.ConfigsRecord]=None) -> flwr.common.ParametersRecord: 
+        """
+        Aggregates model parameters from a list of clients.
+        Args:
+            results: List of `flwr.common.ParametersRecord` from clients.
+        Returns:
+            Aggregated `flwr.common.ParametersRecord`.
+        """
+        ...
+
+    def aggregate_metrics(self, results: list[flwr.common.MetricsRecord], config: Optional[flwr.common.ConfigsRecord]=None) -> flwr.common.MetricsRecord: 
+        """
+        Aggregates metrics from a list of clients.
+        Args:
+            results: List of `flwr.common.MetricsRecord` from clients.
+        Returns:
+            Aggregated `flwr.common.MetricsRecord`.
+        """
+        ...
+
+@runtime_checkable
+class AggFactoryProtocol(Protocol):
+    '''
+    Protocol for the function creating the aggregator instance
+    '''
+    def __call__(self) -> AggProtocol:
+        """
+        Creates and returns an instance of a class conforming to AggProtocol.
+        """
+        ...

FedModelKit/cli.py

@@ -0,0 +1,97 @@
+import os
+import sys
+import argparse
+from .templates import (
+    get_main_template,
+    get_task_template,
+    get_readme_template,
+    get_server_template,
+    get_client_template,
+    get_pyproject_template,
+    get_extern_pyproject_template,
+    get_uv_template,
+    get_init_template,
+    get_ds_template,
+    get_images,
+)
+
+def create_structure(exp_name: str = "new_experiment") -> None:
+    root_dir = os.getcwd()
+
+    structure = {
+        f"{exp_name}": {
+            "main.py": get_main_template(),
+            "pyproject.toml": get_pyproject_template(exp_name),
+            f"{exp_name}": {
+                "__init__.py": get_init_template(),
+                "task.py": get_task_template(),
+                "server_app.py": get_server_template(),
+                "client_app.py": get_client_template(),                
+            },
+        },
+        "pyproject.toml": get_extern_pyproject_template(exp_name),
+        "uv.lock": get_uv_template(),
+        "README.md": get_readme_template(),
+        "ds.ipynb": get_ds_template(),
+    }
+
+
+    for name, content in structure.items():
+        path = os.path.join(root_dir, name)
+        if isinstance(content, dict):  # If content is a folder
+            os.makedirs(path, exist_ok=True)
+            for name1, content1 in content.items():
+                nested_path = os.path.join(path, name1)
+                if isinstance(content1, dict):  # If nested content is a folder
+                    os.makedirs(nested_path, exist_ok=True)
+                    for name2, content2 in content1.items():
+                        nested_nested_path = os.path.join(nested_path, name2)
+                        with open(nested_nested_path, "w") as f:
+                            f.write(content2)
+                else:
+                    with open(nested_path, "w") as f:
+                        f.write(content1)
+        else:  # Single file at the root
+            with open(path, "w") as f:
+                f.write(content)
+    # Create images directory and copy images
+    get_images(root_dir)
+
+    print(f"Project structure for {exp_name} created successfully in {root_dir}.")
+
+
+'''def main():
+    if len(sys.argv) < 2:
+        print("Usage: msi <command> [options]")
+        sys.exit(1)
+
+    command = sys.argv[1]
+
+    if command == "init":
+        structure_type = sys.argv[2] if len(sys.argv) > 2 else "default"
+        create_structure(structure_type)
+    else:
+        print(f"Unknown command: {command}")
+        sys.exit(1)
+'''
+
+def main():
+    if len(sys.argv) < 2:
+        print("Usage: fmk init -n/--name [project_name]")
+        sys.exit(1)
+    
+    parser = argparse.ArgumentParser(prog='fmk', description='FedModelKit CLI for managing federated learning projects')
+    subparsers = parser.add_subparsers(dest='command')
+
+
+    # init subcommand
+    init_parser = subparsers.add_parser('init', help='Initialize the project')
+    init_parser.add_argument('-n', '--name', required=True, help='Set the experiment title')
+
+    args = parser.parse_args()
+
+    if args.command == 'init':
+        print(f"Initializing project with title: {args.name}")
+        
+        # Create the project structure
+        create_structure(args.name)

FedModelKit/default_create_functions.py

@@ -0,0 +1,114 @@
+# Create your custom function defining a PyTorch-based model and returning its instance
+def default_create_local_learner():
+    import pandas as pd
+    from torch import nn, optim
+    from torch.utils.data import DataLoader, Dataset
+    import torch
+    import flwr
+
+    from .src.utils import Utils # Dependency script stored inside the src directory
+    
+    class DefaultLocalLearner(nn.Module):
+            def __init__(self, input_size: int) -> None:
+                super(DefaultLocalLearner, self).__init__()
+                self.linear = nn.Linear(input_size, 3)
+                self.softmax = nn.Softmax(dim=1)
+                self.loss_fn = nn.CrossEntropyLoss()
+                self.optimizer = optim.Adam(self.parameters(), lr=0.001)
+
+            def _forward(self, x):
+                x = self.linear(x)
+                return self.softmax(x)
+
+            def get_parameters(self) -> flwr.common.ParametersRecord:
+                return Utils.pytorch_to_parameter_record(self.state_dict())
+
+            def set_parameters(self, parameters: flwr.common.ParametersRecord) -> None:
+                self.load_state_dict(Utils.parameters_to_pytorch_state_dict(parameters))
+
+            def prepare_data(self, data: pd.DataFrame) -> None:
+                class IrisDataset(Dataset):
+                    def __init__(self, dataframe: pd.DataFrame) -> None:
+                        self.dataframe = dataframe
+                        self.dataframe.loc[:, "class"] = dataframe.loc[:, "class"].replace(
+                            {"Iris-setosa": 0, "Iris-versicolor": 1, "Iris-virginica": 2}
+                        )
+
+                    def __getitem__(self, idx: int):
+                        x = self.dataframe.iloc[idx, :-1].to_numpy("float32")
+                        y = torch.tensor(self.dataframe.iloc[idx, -1], dtype=torch.long)
+                        return x, y
+
+                    def __len__(self) -> int:
+                        return len(self.dataframe)
+
+                dataset = IrisDataset(data)
+                dataloader = DataLoader(dataset, batch_size=32)
+                self.dataloader = dataloader
+
+            def train_round(self) -> flwr.common.MetricsRecord: # type: ignore
+                for batch in self.dataloader:
+                    x, y = batch
+                    self.optimizer.zero_grad()
+                    y_hat = self._forward(x)
+                    loss = self.loss_fn(y_hat, y)
+                    loss.backward()
+                    self.optimizer.step()
+
+                return flwr.common.MetricsRecord({"loss": loss.item()}) # type: ignore
+
+            def evaluate(self) -> flwr.common.MetricsRecord:
+                correct = 0
+                total = 0
+                with torch.no_grad():
+                    for batch in self.dataloader:
+                        x, y = batch
+                        y_hat = self._forward(x)
+                        _, predicted = torch.max(y_hat, 1)
+                        total += y.size(0)
+                        correct += (predicted == y).sum().item()
+                return flwr.common.MetricsRecord({"accuracy": correct / total})
+    
+    return DefaultLocalLearner(4)
+
+
+# Create your custom function defining an aggregation strategy and returning its instance
+def default_create_aggregator():
+    from collections import OrderedDict
+    import numpy as np
+    import flwr
+    from typing import Optional
+
+    from .src.utils import Utils # Dependency script stored inside the src directory
+
+    class DefaultAggregator:
+
+        def aggregate_parameters(self, results: list[flwr.common.ParametersRecord], config: Optional[flwr.common.ConfigsRecord]=None
+            ) -> flwr.common.ParametersRecord:
+                parameters = [Utils.parameters_to_dict(param) for param in results]
+                keys = parameters[0].keys()
+                result = OrderedDict()
+                for key in keys:
+                    # Init array
+                    this_array: np.ndarray = np.zeros_like(parameters[0][key])
+                    for p in parameters:
+                        this_array += p[key]
+                    result[key] = this_array / len(results)
+                return Utils.dict_to_parameter_record(result)
+
+        def aggregate_metrics(self, results: list[flwr.common.MetricsRecord], config: Optional[flwr.common.ConfigsRecord]=None) -> flwr.common.MetricsRecord:
+                keys = results[0].keys()
+                result = OrderedDict()
+                for key in keys:
+                    # Init array
+                    cumsum = 0.0
+                    for m in results:
+                        if not isinstance(m[key], (int, float)):
+                            raise ValueError(
+                                f"flwr.common.MetricsRecord value type not supported: {type(m[key])}"
+                            )
+                        cumsum += m[key]  # type: ignore
+                    result[key] = cumsum / len(results)
+                return flwr.common.MetricsRecord(result)  # type: ignore
+    
+    return DefaultAggregator()

FedModelKit/interface.py

@@ -0,0 +1,130 @@
+'''import mlflow
+from mlflow import MlflowClient
+from mlflow.pyfunc.model import PythonModel
+from mlflow.pyfunc import log_model'''
+
+from .local_learner import LLFactoryProtocol, LLProtocol
+from .aggregator import AggFactoryProtocol, AggProtocol
+from .default_create_functions import default_create_local_learner, default_create_aggregator
+
+from typing import Literal
+import os
+import warnings
+
+
+class FederatedModel():
+    """
+    A class for creating a federated learning model and its aggregator. If no aggregator is provided
+    the PlainAVGAggregator is passed as default
+
+    This class encapsulates the local learner, the aggregation strategy, and their respective names.
+    It is designed to facilitate the integration of these components into the MLflow tracking system.
+
+    Attributes:
+        aggregator (AggFactoryProtocol): The aggregation strategy for federated learning.
+        aggregator_name (str): The name of the aggregation strategy (default is "PlainAVG").
+        local_learner (LLFactoryProtocol): The local learner factory for creating model instances.
+        model_name (str): The name of the model to be registered in MLflow.
+    """
+    def __init__(self, create_local_learner: LLFactoryProtocol = default_create_local_learner,
+                  model_name: str = "Default_iris_model", 
+                  create_aggregator: AggFactoryProtocol = default_create_aggregator,
+                  aggregator_name: str = "PlainAVG"
+                  ) -> None:
+        self.create_aggregator = create_aggregator
+        self.aggregator_name = aggregator_name
+        self.create_local_learner = create_local_learner
+        self.model_name = model_name
+
+
+
+'''def submit_fl_model(model: FederatedModel,
+                    platform_url: str,
+                    username: str,
+                    password: str,
+                    experiment_name: str,
+                    disease: Literal["AML", "SCD"],
+                    trained: bool) -> dict:
+    """
+    Submit the model and aggregation strategy to MLflow.
+
+    This function logs the model and aggregator as artifacts, registers the model in MLflow,
+    and ensures a new run is created in the specified experiment.
+
+    Args:
+        model (FederatedModel): The federated learning model to be submitted.
+        platform_url (str): The URL of the MLflow tracking server.
+        username (str): The username for authenticating with the MLflow tracking server.
+        password (str): The password for authenticating with the MLflow tracking server.
+        experiment_name (str): The name of the MLflow experiment to use or create.
+        disease (Literal["AML", "SCD"]): The use case for the model.
+        trained (bool): Indicates whether the model is already trained.
+
+    Raises:
+        RuntimeError: If the username or password is not provided.
+    """
+    # Ignore warnings about input example
+    warnings.filterwarnings("ignore")
+
+    # Check if username and password are provided
+    if not username or not password:
+        raise RuntimeError("Username and password must be provided.")
+    # Check if the use case is valid
+    assert disease in ["AML", "SCD"], "Disease must be either 'AML' or 'SCD'"
+
+    os.environ['MLFLOW_TRACKING_USERNAME'] = username
+    os.environ['MLFLOW_TRACKING_PASSWORD'] = password
+    os.environ["MLFLOW_TRACKING_INSECURE_TLS"] = "true"
+
+    MLFLOW_URL = platform_url
+    mlflow.set_tracking_uri(MLFLOW_URL)
+
+    # Test model and aggregator protocol
+    model_instance = model.create_local_learner()
+    aggregator_instance = model.create_aggregator()
+    assert isinstance(model.create_aggregator, AggFactoryProtocol), "create_aggregator function does not conform to AggFactoryProtocol"
+    assert isinstance(model.create_local_learner, LLFactoryProtocol), "create_local_learner function does not conform to LLFactoryProtocol"
+    assert isinstance(model_instance, LLProtocol), "Local learner does not conform to LLProtocol" 
+    assert isinstance(aggregator_instance, AggProtocol), "Aggregator instance does not conform to AggProtocol"
+    
+    # Ensure the experiment exists or create it
+    experiment = mlflow.get_experiment_by_name(experiment_name)
+    if experiment is None:
+        experiment_id = mlflow.create_experiment(experiment_name)
+    else:
+        experiment_id = experiment.experiment_id
+
+
+    # Start a new run
+    with mlflow.start_run(experiment_id=experiment_id) as run:
+        mlflow.set_tag("use_case", disease)
+        mlflow.set_tag("trained", str(trained))
+        if os.path.isdir("./src"):
+            model_info = log_model(
+                artifact_path="model",
+                python_model=model,
+                registered_model_name=model.model_name,
+                code_paths=["src"],
+            )
+        else:
+            model_info = log_model(
+                artifact_path="model",
+                python_model=model,
+                registered_model_name=model.model_name,
+            )
+    mlflow_client = MlflowClient(tracking_uri=MLFLOW_URL)
+    model_meta = mlflow_client.get_latest_versions(model.model_name, stages=["None"])
+    version = model_meta[0].version
+    # mlflow_client.update_model_version(model.model_name, version, description)
+    tags = {"use_case": disease, "trained": str(trained)}
+    for key, value in tags.items():
+        mlflow_client.set_model_version_tag(model.model_name, version, key, value)
+    # mlflow_client.set_experiment_tag(experiment_id, "use_case", disease)
+    # mlflow_client.set_experiment_tag(experiment_id, "Trained", str(trained))
+    return {
+        "detail": f"Model '{model.model_name}' registered.",
+        "model_uuid": model_info.model_uuid,
+        "run_id": model_info.run_id,
+        "model_uri": model_info.model_uri,
+        }'''
+

FedModelKit/local_learner.py

@@ -0,0 +1,61 @@
+from typing import Protocol, runtime_checkable
+
+import pandas as pd
+import flwr
+
+@runtime_checkable
+class LLProtocol(Protocol):
+    """
+    Protocol for defining a local learner in a federated learning setup.
+    """
+
+    def prepare_data(self, data: pd.DataFrame) -> None:
+        """
+        Prepares input data for training or evaluation.
+        Args:
+            data: Input data as a pandas DataFrame.
+        """
+        ...
+
+    def train_round(self) -> flwr.common.MetricsRecord:
+        """
+        Trains the model and returns metrics.
+        Returns:
+            flwr.common.MetricsRecord: Performance metrics after training.
+        """
+        ...
+
+    def get_parameters(self) -> flwr.common.ParametersRecord:
+        """
+        Retrieves the current model parameters.
+        Returns:
+            flwr.common.ParametersRecord: Current model parameters.
+        """
+        ...
+
+    def set_parameters(self, parameters: flwr.common.ParametersRecord) -> None:
+        """
+        Sets the model parameters.
+        Args:
+            parameters: A flwr.common.ParametersRecord containing the parameters to set.
+        """
+        ...
+
+    def evaluate(self) -> flwr.common.MetricsRecord:
+        """
+        Evaluates the model and returns performance metrics.
+        Returns:
+            flwr.common.MetricsRecord: Evaluation metrics.
+        """
+        ...
+
+@runtime_checkable
+class LLFactoryProtocol(Protocol):
+    '''
+    Protocol for the function creating the local learner instance
+    '''
+    def __call__(self) -> LLProtocol:
+            """
+            Creates and returns an instance of a class conforming to LLProtocol.
+            """
+            ...

FedModelKit/src/utils.py

@@ -0,0 +1,65 @@
+from collections import OrderedDict
+import flwr
+
+import numpy as np
+
+
+class Utils:
+    import flwr
+    @staticmethod
+    def _basic_array_deserialisation(array: flwr.common.Array) -> flwr.common.NDArray:
+        return np.frombuffer(buffer=array.data, dtype=array.dtype).reshape(array.shape)
+
+    @staticmethod
+    def parameters_to_dict(params_record: flwr.common.ParametersRecord) -> OrderedDict:
+        state_dict = OrderedDict()
+        for k, v in params_record.items():
+            state_dict[k] = Utils._basic_array_deserialisation(v)
+
+        return state_dict
+
+    @staticmethod
+    def _ndarray_to_array(ndarray: flwr.common.NDArray) -> flwr.common.Array:
+        """Represent NumPy ndarray as Array."""
+        return flwr.common.Array(
+            data=ndarray.tobytes(),
+            dtype=str(ndarray.dtype),
+            stype="numpy.ndarray.tobytes",
+            shape=list(ndarray.shape),
+        )
+
+    @staticmethod
+    def dict_to_parameter_record(
+        parameters: OrderedDict["str", flwr.common.NDArray],
+    ) -> flwr.common.ParametersRecord:
+        state_dict = OrderedDict()
+        for k, v in parameters.items():
+            state_dict[k] = Utils._ndarray_to_array(v)
+
+        return flwr.common.ParametersRecord(state_dict)
+
+    @staticmethod
+    def pytorch_to_parameter_record(
+        state_dict: dict,
+    ) -> flwr.common.ParametersRecord:
+        """Serialise your PyTorch model."""
+        transformed_state_dict = OrderedDict()
+
+        for k, v in state_dict.items():
+            transformed_state_dict[k] = Utils._ndarray_to_array(v.numpy())
+
+        return flwr.common.ParametersRecord(transformed_state_dict)
+
+    @staticmethod
+    def parameters_to_pytorch_state_dict(
+        params_record: flwr.common.ParametersRecord,
+    ) -> dict:
+        # Make sure to import locally torch as it is only available in the server
+        import torch
+
+        """Reconstruct PyTorch state_dict from its serialised representation."""
+        state_dict = {}
+        for k, v in params_record.items():
+            state_dict[k] = torch.tensor(Utils._basic_array_deserialisation(v))
+
+        return state_dict

FedModelKit/templates/client_app_template.py

@@ -0,0 +1,118 @@
+
+import numpy as np
+import pandas as pd
+import pickle
+from pathlib import Path
+from flwr.client import ClientApp
+from flwr.common import Message, Context
+from flwr.common.record import RecordSet, MetricsRecord, ConfigsRecord
+from sklearn.preprocessing import OneHotEncoder
+import FedModelKit as msi
+
+from model_example import create_local_learner #type: ignore[import]
+from load_data import load_data #type: ignore[import]
+
+
+# Initialize the Flower ClientApp
+app = ClientApp()
+
+@app.query()
+def query(msg: Message, ctx: Context) -> Message:
+    """
+    Query function to be executed by the Flower client. This function handles the
+    initial configuration sent by the server.
+    """
+
+    # Retrieve the configuration sent by the server
+    fancy_config = msg.content.configs_records['fancy_config']
+
+    # Load the client split data using the load_data function
+    data = load_data(fancy_config['num_clients'], fancy_config['client_id'])
+
+    # Instantiate the federated model
+    federated_model = msi.FederatedModel(create_local_learner=create_local_learner, model_name='simple_lr')
+
+    # Store the local learner and the data split in the context
+    # To store in context other objects, you can use ctx.state.<object_name> = <object>
+    ctx.state.local_learner = federated_model.create_local_learner()
+    ctx.state.data = data
+    ctx.state.undestand = "che succede"
+
+    return msg.create_reply(RecordSet())
+
+@app.train()
+def train(msg: Message, ctx: Context):
+    """
+    Train function to be executed by the Flower client.
+    This function handles the training of the local model using the data provided.
+    """
+
+    # Retrieve the local learner and the client split from the context
+    local_learner = ctx.state.local_learner
+    data = ctx.state.data
+
+    # Retrieve configuration sent by the server - example
+    #fancy_config = msg.content.configs_records['fancy_config']
+    #local_epochs = fancy_config['local_epochs']
+
+    # Retrieve the model parameters sent by the server
+    fancy_parameters = msg.content.parameters_records['fancy_model']
+    local_learner.set_parameters(fancy_parameters)    
+
+    # Prepare the data using the local learner
+    local_learner.prepare_data(data)
+
+    # Perform local training and obtain training metrics
+    train_metrics = local_learner.train_round()
+
+    # Retrieve the trained model parameters
+    new_parameters_records = local_learner.get_parameters()
+    assert ctx.state.undestand.startswith("che"), "The context state is not being stored correctly"
+
+    # Construct a reply message carrying updated model parameters and generated metrics
+    reply_content = RecordSet()
+    reply_content.parameters_records['fancy_model_returned'] = new_parameters_records
+    reply_content.metrics_records['train_metrics'] = train_metrics
+
+    # Store the metrics and the local learner in the context for future reference
+    ctx.state.metrics_records['prev'] = train_metrics
+    ctx.state.local_learner =  local_learner
+
+    # Return the reply message to the server
+    return msg.create_reply(reply_content)
+
+@app.evaluate()
+def eval(msg: Message, ctx: Context):
+    """
+    Evaluate function to be executed by the Flower client.
+    This function handles the evaluation of the local model using the data provided.
+    """
+
+    # Retrieve the local learner and the client split from the context
+    local_learner = ctx.state.local_learner
+    data = ctx.state.data
+
+    # Retrieve configuration sent by the server - example
+    #fancy_config = msg.content.configs_records['fancy_config']
+    #local_epochs = fancy_config['local_epochs']
+
+    # Retrieve the model parameters sent by the server
+    fancy_parameters = msg.content.parameters_records['fancy_model']
+    local_learner.set_parameters(fancy_parameters)
+
+    # Prepare the data using the local learner
+    local_learner.prepare_data(data)
+
+    # Evaluate the model and obtain evaluation metrics
+    eval_metrics = local_learner.evaluate()
+
+    # Construct a reply message with evaluation metrics
+    reply_content = RecordSet()
+    reply_content.metrics_records['eval_metrics'] = eval_metrics
+
+    # Store the metrics and the local learner in the context for future reference
+    ctx.state.metrics_records['prev'] = eval_metrics
+    ctx.state.local_learner =  local_learner
+
+    # Return the reply message to the server
+    return msg.create_reply(reply_content)