flwr-nightly 1.22.0.dev20250917__py3-none-any.whl → 1.22.0.dev20250918__py3-none-any.whl

flwr/cli/new/new.py

@@ -41,6 +41,7 @@ class MlFramework(str, Enum):
     JAX = "JAX"
     MLX = "MLX"
     NUMPY = "NumPy"
+    XGBOOST = "XGBoost"
     FLOWERTUNE = "FlowerTune"
     BASELINE = "Flower Baseline"
     PYTORCH_LEGACY_API = "PyTorch (Legacy API, deprecated)"
@@ -247,6 +248,7 @@ def new(
             MlFramework.TENSORFLOW.value,
             MlFramework.SKLEARN.value,
             MlFramework.NUMPY.value,
+            MlFramework.XGBOOST.value,
             "pytorch_legacy_api",
         ]
         if framework_str in frameworks_with_tasks:

flwr/cli/new/templates/app/code/client.xgboost.py.tpl

@@ -0,0 +1,110 @@
+"""$project_name: A Flower / $framework_str app."""
+
+import warnings
+
+import numpy as np
+import xgboost as xgb
+from flwr.app import ArrayRecord, Context, Message, MetricRecord, RecordDict
+from flwr.clientapp import ClientApp
+from flwr.common.config import unflatten_dict
+
+from $import_name.task import load_data, replace_keys
+
+warnings.filterwarnings("ignore", category=UserWarning)
+
+
+# Flower ClientApp
+app = ClientApp()
+
+
+def _local_boost(bst_input, num_local_round, train_dmatrix):
+    # Update trees based on local training data.
+    for i in range(num_local_round):
+        bst_input.update(train_dmatrix, bst_input.num_boosted_rounds())
+
+    # Bagging: extract the last N=num_local_round trees for sever aggregation
+    bst = bst_input[
+        bst_input.num_boosted_rounds()
+        - num_local_round : bst_input.num_boosted_rounds()
+    ]
+    return bst
+
+
+@app.train()
+def train(msg: Message, context: Context) -> Message:
+    # Load model and data
+    partition_id = context.node_config["partition-id"]
+    num_partitions = context.node_config["num-partitions"]
+    train_dmatrix, _, num_train, _ = load_data(partition_id, num_partitions)
+
+    # Read from run config
+    num_local_round = context.run_config["local-epochs"]
+    # Flatted config dict and replace "-" with "_"
+    cfg = replace_keys(unflatten_dict(context.run_config))
+    params = cfg["params"]
+
+    global_round = msg.content["config"]["server-round"]
+    if global_round == 1:
+        # First round local training
+        bst = xgb.train(
+            params,
+            train_dmatrix,
+            num_boost_round=num_local_round,
+        )
+    else:
+        bst = xgb.Booster(params=params)
+        global_model = bytearray(msg.content["arrays"]["0"].numpy().tobytes())
+
+        # Load global model into booster
+        bst.load_model(global_model)
+
+        # Local training
+        bst = _local_boost(bst, num_local_round, train_dmatrix)
+
+    # Save model
+    local_model = bst.save_raw("json")
+    model_np = np.frombuffer(local_model, dtype=np.uint8)
+
+    # Construct reply message
+    # Note: we store the model as the first item in a list into ArrayRecord,
+    # which can be accessed using index ["0"].
+    model_record = ArrayRecord([model_np])
+    metrics = {
+        "num-examples": num_train,
+    }
+    metric_record = MetricRecord(metrics)
+    content = RecordDict({"arrays": model_record, "metrics": metric_record})
+    return Message(content=content, reply_to=msg)
+
+
+@app.evaluate()
+def evaluate(msg: Message, context: Context) -> Message:
+    # Load model and data
+    partition_id = context.node_config["partition-id"]
+    num_partitions = context.node_config["num-partitions"]
+    _, valid_dmatrix, _, num_val = load_data(partition_id, num_partitions)
+
+    # Load config
+    cfg = replace_keys(unflatten_dict(context.run_config))
+    params = cfg["params"]
+
+    # Load global model
+    bst = xgb.Booster(params=params)
+    global_model = bytearray(msg.content["arrays"]["0"].numpy().tobytes())
+    bst.load_model(global_model)
+
+    # Run evaluation
+    eval_results = bst.eval_set(
+        evals=[(valid_dmatrix, "valid")],
+        iteration=bst.num_boosted_rounds() - 1,
+    )
+    auc = float(eval_results.split("\t")[1].split(":")[1])
+
+    # Construct and return reply Message
+    metrics = {
+        "auc": auc,
+        "num-examples": num_val,
+    }
+    metric_record = MetricRecord(metrics)
+    content = RecordDict({"metrics": metric_record})
+    return Message(content=content, reply_to=msg)

flwr/cli/new/templates/app/code/server.xgboost.py.tpl

@@ -0,0 +1,56 @@
+"""$project_name: A Flower / $framework_str app."""
+
+import numpy as np
+import xgboost as xgb
+from flwr.app import ArrayRecord, Context
+from flwr.common.config import unflatten_dict
+from flwr.serverapp import Grid, ServerApp
+from flwr.serverapp.strategy import FedXgbBagging
+
+from $import_name.task import replace_keys
+
+# Create ServerApp
+app = ServerApp()
+
+
+@app.main()
+def main(grid: Grid, context: Context) -> None:
+    # Read run config
+    num_rounds = context.run_config["num-server-rounds"]
+    fraction_train = context.run_config["fraction-train"]
+    fraction_evaluate = context.run_config["fraction-evaluate"]
+    # Flatted config dict and replace "-" with "_"
+    cfg = replace_keys(unflatten_dict(context.run_config))
+    params = cfg["params"]
+
+    # Init global model
+    # Init with an empty object; the XGBooster will be created
+    # and trained on the client side.
+    global_model = b""
+    # Note: we store the model as the first item in a list into ArrayRecord,
+    # which can be accessed using index ["0"].
+    arrays = ArrayRecord([np.frombuffer(global_model, dtype=np.uint8)])
+
+    # Initialize FedXgbBagging strategy
+    strategy = FedXgbBagging(
+        fraction_train=fraction_train,
+        fraction_evaluate=fraction_evaluate,
+    )
+
+    # Start strategy, run FedXgbBagging for `num_rounds`
+    result = strategy.start(
+        grid=grid,
+        initial_arrays=arrays,
+        num_rounds=num_rounds,
+    )
+
+    # Save final model to disk
+    bst = xgb.Booster(params=params)
+    global_model = bytearray(result.arrays["0"].numpy().tobytes())
+
+    # Load global model into booster
+    bst.load_model(global_model)
+
+    # Save model
+    print("\nSaving final model to disk...")
+    bst.save_model("final_model.json")

flwr/cli/new/templates/app/code/task.xgboost.py.tpl

@@ -0,0 +1,67 @@
+"""$project_name: A Flower / $framework_str app."""
+
+import xgboost as xgb
+from flwr_datasets import FederatedDataset
+from flwr_datasets.partitioner import IidPartitioner
+
+
+def train_test_split(partition, test_fraction, seed):
+    """Split the data into train and validation set given split rate."""
+    train_test = partition.train_test_split(test_size=test_fraction, seed=seed)
+    partition_train = train_test["train"]
+    partition_test = train_test["test"]
+
+    num_train = len(partition_train)
+    num_test = len(partition_test)
+
+    return partition_train, partition_test, num_train, num_test
+
+
+def transform_dataset_to_dmatrix(data):
+    """Transform dataset to DMatrix format for xgboost."""
+    x = data["inputs"]
+    y = data["label"]
+    new_data = xgb.DMatrix(x, label=y)
+    return new_data
+
+
+fds = None  # Cache FederatedDataset
+
+
+def load_data(partition_id, num_clients):
+    """Load partition HIGGS data."""
+    # Only initialize `FederatedDataset` once
+    global fds
+    if fds is None:
+        partitioner = IidPartitioner(num_partitions=num_clients)
+        fds = FederatedDataset(
+            dataset="jxie/higgs",
+            partitioners={"train": partitioner},
+        )
+
+    # Load the partition for this `partition_id`
+    partition = fds.load_partition(partition_id, split="train")
+    partition.set_format("numpy")
+
+    # Train/test splitting
+    train_data, valid_data, num_train, num_val = train_test_split(
+        partition, test_fraction=0.2, seed=42
+    )
+
+    # Reformat data to DMatrix for xgboost
+    train_dmatrix = transform_dataset_to_dmatrix(train_data)
+    valid_dmatrix = transform_dataset_to_dmatrix(valid_data)
+
+    return train_dmatrix, valid_dmatrix, num_train, num_val
+
+
+def replace_keys(input_dict, match="-", target="_"):
+    """Recursively replace match string with target string in dictionary keys."""
+    new_dict = {}
+    for key, value in input_dict.items():
+        new_key = key.replace(match, target)
+        if isinstance(value, dict):
+            new_dict[new_key] = replace_keys(value, match, target)
+        else:
+            new_dict[new_key] = value
+    return new_dict

flwr/cli/new/templates/app/pyproject.xgboost.toml.tpl

@@ -0,0 +1,61 @@
+# =====================================================================
+# For a full TOML configuration guide, check the Flower docs:
+# https://flower.ai/docs/framework/how-to-configure-pyproject-toml.html
+# =====================================================================
+
+[build-system]
+requires = ["hatchling"]
+build-backend = "hatchling.build"
+
+[project]
+name = "$package_name"
+version = "1.0.0"
+description = ""
+license = "Apache-2.0"
+# Dependencies for your Flower App
+dependencies = [
+    "flwr[simulation]>=1.22.0",
+    "flwr-datasets>=0.5.0",
+    "xgboost>=2.0.0",
+]
+
+[tool.hatch.build.targets.wheel]
+packages = ["."]
+
+[tool.flwr.app]
+publisher = "$username"
+
+[tool.flwr.app.components]
+serverapp = "$import_name.server_app:app"
+clientapp = "$import_name.client_app:app"
+
+# Custom config values accessible via `context.run_config`
+[tool.flwr.app.config]
+num-server-rounds = 3
+fraction-train = 0.1
+fraction-evaluate = 0.1
+local-epochs = 1
+
+# XGBoost parameters
+params.objective = "binary:logistic"
+params.eta = 0.1 # Learning rate
+params.max-depth = 8
+params.eval-metric = "auc"
+params.nthread = 16
+params.num-parallel-tree = 1
+params.subsample = 1
+params.tree-method = "hist"
+
+# Default federation to use when running the app
+[tool.flwr.federations]
+default = "local-simulation"
+
+# Local simulation federation with 10 virtual SuperNodes
+[tool.flwr.federations.local-simulation]
+options.num-supernodes = 10
+
+# Remote federation example for use with SuperLink
+[tool.flwr.federations.remote-federation]
+address = "<SUPERLINK-ADDRESS>:<PORT>"
+insecure = true  # Remove this line to enable TLS
+# root-certificates = "<PATH/TO/ca.crt>"  # For TLS setup

flwr/common/exit/exit_code.py

@@ -45,6 +45,7 @@ class ExitCode:
     SUPERNODE_NODE_AUTH_KEYS_INVALID = 302
 
     # SuperExec-specific exit codes (400-499)
+    SUPEREXEC_INVALID_PLUGIN_CONFIG = 400
 
     # Common exit codes (600-699)
     COMMON_ADDRESS_INVALID = 600
@@ -112,6 +113,9 @@ EXIT_CODE_HELP = {
         "file and try again."
     ),
     # SuperExec-specific exit codes (400-499)
+    ExitCode.SUPEREXEC_INVALID_PLUGIN_CONFIG: (
+        "The YAML configuration for the SuperExec plugin is invalid."
+    ),
     # Common exit codes (600-699)
     ExitCode.COMMON_ADDRESS_INVALID: (
         "Please provide a valid URL, IPv4 or IPv6 address."

flwr/serverapp/strategy/__init__.py

@@ -27,6 +27,7 @@ from .fedmedian import FedMedian
 from .fedprox import FedProx
 from .fedtrimmedavg import FedTrimmedAvg
 from .fedxgb_bagging import FedXgbBagging
+from .fedxgb_cyclic import FedXgbCyclic
 from .fedyogi import FedYogi
 from .result import Result
 from .strategy import Strategy
@@ -42,6 +43,7 @@ __all__ = [
     "FedProx",
     "FedTrimmedAvg",
     "FedXgbBagging",
+    "FedXgbCyclic",
     "FedYogi",
     "Result",
     "Strategy",

flwr/serverapp/strategy/fedxgb_cyclic.py

@@ -0,0 +1,220 @@
+# Copyright 2025 Flower Labs GmbH. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Flower message-based FedXgbCyclic strategy."""
+
+
+from collections.abc import Iterable
+from logging import INFO
+from typing import Callable, Optional, cast
+
+from flwr.common import (
+    ArrayRecord,
+    ConfigRecord,
+    Message,
+    MessageType,
+    MetricRecord,
+    RecordDict,
+    log,
+)
+from flwr.server import Grid
+
+from .fedavg import FedAvg
+from .strategy_utils import sample_nodes
+
+
+# pylint: disable=line-too-long
+class FedXgbCyclic(FedAvg):
+    """Configurable FedXgbCyclic strategy implementation.
+
+    Parameters
+    ----------
+    fraction_train : float (default: 1.0)
+        Fraction of nodes used during training. In case `min_train_nodes`
+        is larger than `fraction_train * total_connected_nodes`, `min_train_nodes`
+        will still be sampled.
+    fraction_evaluate : float (default: 1.0)
+        Fraction of nodes used during validation. In case `min_evaluate_nodes`
+        is larger than `fraction_evaluate * total_connected_nodes`,
+        `min_evaluate_nodes` will still be sampled.
+    min_available_nodes : int (default: 2)
+        Minimum number of total nodes in the system.
+    weighted_by_key : str (default: "num-examples")
+        The key within each MetricRecord whose value is used as the weight when
+        computing weighted averages for both ArrayRecords and MetricRecords.
+    arrayrecord_key : str (default: "arrays")
+        Key used to store the ArrayRecord when constructing Messages.
+    configrecord_key : str (default: "config")
+        Key used to store the ConfigRecord when constructing Messages.
+    train_metrics_aggr_fn : Optional[callable] (default: None)
+        Function with signature (list[RecordDict], str) -> MetricRecord,
+        used to aggregate MetricRecords from training round replies.
+        If `None`, defaults to `aggregate_metricrecords`, which performs a weighted
+        average using the provided weight factor key.
+    evaluate_metrics_aggr_fn : Optional[callable] (default: None)
+        Function with signature (list[RecordDict], str) -> MetricRecord,
+        used to aggregate MetricRecords from training round replies.
+        If `None`, defaults to `aggregate_metricrecords`, which performs a weighted
+        average using the provided weight factor key.
+    """
+
+    # pylint: disable=too-many-arguments,too-many-positional-arguments
+    def __init__(
+        self,
+        fraction_train: float = 1.0,
+        fraction_evaluate: float = 1.0,
+        min_available_nodes: int = 2,
+        weighted_by_key: str = "num-examples",
+        arrayrecord_key: str = "arrays",
+        configrecord_key: str = "config",
+        train_metrics_aggr_fn: Optional[
+            Callable[[list[RecordDict], str], MetricRecord]
+        ] = None,
+        evaluate_metrics_aggr_fn: Optional[
+            Callable[[list[RecordDict], str], MetricRecord]
+        ] = None,
+    ) -> None:
+        super().__init__(
+            fraction_train=fraction_train,
+            fraction_evaluate=fraction_evaluate,
+            min_train_nodes=2,
+            min_evaluate_nodes=2,
+            min_available_nodes=min_available_nodes,
+            weighted_by_key=weighted_by_key,
+            arrayrecord_key=arrayrecord_key,
+            configrecord_key=configrecord_key,
+            train_metrics_aggr_fn=train_metrics_aggr_fn,
+            evaluate_metrics_aggr_fn=evaluate_metrics_aggr_fn,
+        )
+
+        self.registered_nodes: dict[int, int] = {}
+
+        if fraction_train not in (0.0, 1.0):
+            raise ValueError(
+                "fraction_train can only be set to 1.0 or 0.0 for FedXgbCyclic."
+            )
+        if fraction_evaluate not in (0.0, 1.0):
+            raise ValueError(
+                "fraction_evaluate can only be set to 1.0 or 0.0 for FedXgbCyclic."
+            )
+
+    def _reorder_nodes(self, node_ids: list[int]) -> list[int]:
+        """Re-order node ids based on registered nodes.
+
+        Each node ID is assigned a persistent index in `self.registered_nodes`
+        the first time it appears. The input list is then reordered according
+        to these stored indices, and the result is compacted into ascending
+        order (1..N) for the current call.
+        """
+        # Assign new indices to unknown nodes
+        next_index = max(self.registered_nodes.values(), default=0) + 1
+        for nid in node_ids:
+            if nid not in self.registered_nodes:
+                self.registered_nodes[nid] = next_index
+                next_index += 1
+
+        # Sort node_ids by their stored indices
+        sorted_by_index = sorted(node_ids, key=lambda x: self.registered_nodes[x])
+
+        # Compact re-map of indices just for this output list
+        unique_indices = sorted(self.registered_nodes[nid] for nid in sorted_by_index)
+        remap = {old: new for new, old in enumerate(unique_indices, start=1)}
+
+        # Build the result list ordered by compact indices
+        result_list = [
+            nid
+            for _, nid in sorted(
+                (remap[self.registered_nodes[nid]], nid) for nid in sorted_by_index
+            )
+        ]
+        return result_list
+
+    def _make_sampling(
+        self, grid: Grid, server_round: int, configure_type: str
+    ) -> list[int]:
+        """Sample nodes using the Grid."""
+        # Sample nodes
+        num_nodes = int(len(list(grid.get_node_ids())) * self.fraction_train)
+        sample_size = max(num_nodes, self.min_train_nodes)
+        node_ids, _ = sample_nodes(grid, self.min_available_nodes, sample_size)
+
+        # Re-order node_ids
+        node_ids = self._reorder_nodes(node_ids)
+
+        # Sample the clients sequentially given server_round
+        sampled_idx = (server_round - 1) % len(node_ids)
+        sampled_node_id = [node_ids[sampled_idx]]
+
+        log(
+            INFO,
+            f"{configure_type}: Sampled %s nodes (out of %s)",
+            len(sampled_node_id),
+            len(node_ids),
+        )
+        return sampled_node_id
+
+    def configure_train(
+        self, server_round: int, arrays: ArrayRecord, config: ConfigRecord, grid: Grid
+    ) -> Iterable[Message]:
+        """Configure the next round of federated training."""
+        # Sample one node
+        sampled_node_id = self._make_sampling(grid, server_round, "configure_train")
+
+        # Always inject current server round
+        config["server-round"] = server_round
+
+        # Construct messages
+        record = RecordDict(
+            {self.arrayrecord_key: arrays, self.configrecord_key: config}
+        )
+        return self._construct_messages(record, sampled_node_id, MessageType.TRAIN)
+
+    def aggregate_train(
+        self,
+        server_round: int,
+        replies: Iterable[Message],
+    ) -> tuple[Optional[ArrayRecord], Optional[MetricRecord]]:
+        """Aggregate ArrayRecords and MetricRecords in the received Messages."""
+        valid_replies, _ = self._check_and_log_replies(replies, is_train=True)
+
+        arrays, metrics = None, None
+        if valid_replies:
+            reply_contents = [msg.content for msg in valid_replies]
+            array_record_key = next(iter(reply_contents[0].array_records.keys()))
+
+            # Fetch the client model from current round as global model
+            arrays = cast(ArrayRecord, reply_contents[0][array_record_key])
+
+            # Aggregate MetricRecords
+            metrics = self.train_metrics_aggr_fn(
+                reply_contents,
+                self.weighted_by_key,
+            )
+        return arrays, metrics
+
+    def configure_evaluate(
+        self, server_round: int, arrays: ArrayRecord, config: ConfigRecord, grid: Grid
+    ) -> Iterable[Message]:
+        """Configure the next round of federated evaluation."""
+        # Sample one node
+        sampled_node_id = self._make_sampling(grid, server_round, "configure_evaluate")
+
+        # Always inject current server round
+        config["server-round"] = server_round
+
+        # Construct messages
+        record = RecordDict(
+            {self.arrayrecord_key: arrays, self.configrecord_key: config}
+        )
+        return self._construct_messages(record, sampled_node_id, MessageType.EVALUATE)

flwr/supercore/cli/flower_superexec.py

@@ -17,7 +17,9 @@
 
 import argparse
 from logging import INFO
-from typing import Optional
+from typing import Any, Optional
+
+import yaml
 
 from flwr.common import EventType, event
 from flwr.common.constant import ExecPluginType
@@ -26,6 +28,7 @@ from flwr.common.logger import log
 from flwr.proto.clientappio_pb2_grpc import ClientAppIoStub
 from flwr.proto.serverappio_pb2_grpc import ServerAppIoStub
 from flwr.proto.simulationio_pb2_grpc import SimulationIoStub
+from flwr.supercore.constant import EXEC_PLUGIN_SECTION
 from flwr.supercore.grpc_health import add_args_health
 from flwr.supercore.superexec.plugin import (
     ClientAppExecPlugin,
@@ -36,6 +39,7 @@ from flwr.supercore.superexec.plugin import (
 from flwr.supercore.superexec.run_superexec import run_superexec
 
 try:
+    from flwr.ee import add_ee_args_superexec
     from flwr.ee.constant import ExecEePluginType
     from flwr.ee.exec_plugin import get_ee_plugin_and_stub_class
 except ImportError:
@@ -54,6 +58,10 @@ except ImportError:
         """Get the EE plugin class and stub class based on the plugin type."""
         return None
 
+    # pylint: disable-next=unused-argument
+    def add_ee_args_superexec(parser: argparse.ArgumentParser) -> None:
+        """Add EE-specific arguments to the parser."""
+
 
 def flower_superexec() -> None:
     """Run `flower-superexec` command."""
@@ -70,12 +78,28 @@ def flower_superexec() -> None:
     # Trigger telemetry event
     event(EventType.RUN_SUPEREXEC_ENTER, {"plugin_type": args.plugin_type})
 
+    # Load plugin config from YAML file if provided
+    plugin_config = None
+    if plugin_config_path := getattr(args, "plugin_config", None):
+        try:
+            with open(plugin_config_path, encoding="utf-8") as file:
+                yaml_config: Optional[dict[str, Any]] = yaml.safe_load(file)
+                if yaml_config is None or EXEC_PLUGIN_SECTION not in yaml_config:
+                    raise ValueError(f"Missing '{EXEC_PLUGIN_SECTION}' section.")
+                plugin_config = yaml_config[EXEC_PLUGIN_SECTION]
+        except (FileNotFoundError, yaml.YAMLError, ValueError) as e:
+            flwr_exit(
+                ExitCode.SUPEREXEC_INVALID_PLUGIN_CONFIG,
+                f"Failed to load plugin config from '{plugin_config_path}': {e!r}",
+            )
+
     # Get the plugin class and stub class based on the plugin type
     plugin_class, stub_class = _get_plugin_and_stub_class(args.plugin_type)
     run_superexec(
         plugin_class=plugin_class,
         stub_class=stub_class,  # type: ignore
         appio_api_address=args.appio_api_address,
+        plugin_config=plugin_config,
         flwr_dir=args.flwr_dir,
         parent_pid=args.parent_pid,
         health_server_address=args.health_server_address,
@@ -122,6 +146,7 @@ def _parse_args() -> argparse.ArgumentParser:
         help="The PID of the parent process. When set, the process will terminate "
         "when the parent process exits.",
     )
+    add_ee_args_superexec(parser)
     add_args_health(parser)
     return parser
 

flwr/supercore/constant.py

@@ -0,0 +1,19 @@
+# Copyright 2025 Flower Labs GmbH. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Constants for Flower infrastructure."""
+
+
+# Top-level key in YAML config for exec plugin settings
+EXEC_PLUGIN_SECTION = "exec_plugin"

flwr/supercore/superexec/plugin/exec_plugin.py

@@ -17,7 +17,7 @@
 
 from abc import ABC, abstractmethod
 from collections.abc import Sequence
-from typing import Callable, Optional
+from typing import Any, Callable, Optional
 
 from flwr.common.typing import Run
 
@@ -69,3 +69,13 @@ class ExecPlugin(ABC):
            The ID of the run associated with the token, used for tracking or
            logging purposes.
         """
+
+    # This method is optional to implement
+    def load_config(self, yaml_config: dict[str, Any]) -> None:
+        """Load configuration from a YAML dictionary.
+
+        Parameters
+        ----------
+        yaml_config : dict[str, Any]
+            A dictionary representing the YAML configuration.
+        """

flwr/supercore/superexec/run_superexec.py

@@ -17,10 +17,10 @@
 
 import time
 from logging import WARN
-from typing import Optional, Union
+from typing import Any, Optional, Union
 
 from flwr.common.config import get_flwr_dir
-from flwr.common.exit import register_signal_handlers
+from flwr.common.exit import ExitCode, flwr_exit, register_signal_handlers
 from flwr.common.grpc import create_channel, on_channel_state_change
 from flwr.common.logger import log
 from flwr.common.retry_invoker import _make_simple_grpc_retry_invoker, _wrap_stub
@@ -47,6 +47,7 @@ def run_superexec(  # pylint: disable=R0913,R0914,R0917
         type[ClientAppIoStub], type[ServerAppIoStub], type[SimulationIoStub]
     ],
     appio_api_address: str,
+    plugin_config: Optional[dict[str, Any]] = None,
     flwr_dir: Optional[str] = None,
     parent_pid: Optional[int] = None,
     health_server_address: Optional[str] = None,
@@ -61,6 +62,9 @@ def run_superexec(  # pylint: disable=R0913,R0914,R0917
         The gRPC stub class for the AppIO API.
     appio_api_address : str
         The address of the AppIO API.
+    plugin_config : Optional[dict[str, Any]] (default: None)
+        The configuration dictionary for the plugin. If `None`, the plugin will use
+        its default configuration.
     flwr_dir : Optional[str] (default: None)
         The Flower directory.
     parent_pid : Optional[int] (default: None)
@@ -113,6 +117,16 @@ def run_superexec(  # pylint: disable=R0913,R0914,R0917
         get_run=get_run,
     )
 
+    # Load plugin configuration from file if provided
+    try:
+        if plugin_config is not None:
+            plugin.load_config(plugin_config)
+    except (KeyError, ValueError) as e:
+        flwr_exit(
+            code=ExitCode.SUPEREXEC_INVALID_PLUGIN_CONFIG,
+            message=f"Invalid plugin config: {e!r}",
+        )
+
     # Start the main loop
     try:
         while True:

{flwr_nightly-1.22.0.dev20250917.dist-info → flwr_nightly-1.22.0.dev20250918.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.3
 Name: flwr-nightly
-Version: 1.22.0.dev20250917
+Version: 1.22.0.dev20250918
 Summary: Flower: A Friendly Federated AI Framework
 License: Apache-2.0
 Keywords: Artificial Intelligence,Federated AI,Federated Analytics,Federated Evaluation,Federated Learning,Flower,Machine Learning

{flwr_nightly-1.22.0.dev20250917.dist-info → flwr_nightly-1.22.0.dev20250918.dist-info}/RECORD

@@ -18,7 +18,7 @@ flwr/cli/login/__init__.py,sha256=B1SXKU3HCQhWfFDMJhlC7FOl8UsvH4mxysxeBnrfyUE,80
 flwr/cli/login/login.py,sha256=RM1Jiv_VFm3oz4rTHSr3D87X90lW3WzErjBBU7WviWY,4309
 flwr/cli/ls.py,sha256=3YK7cpoImJ7PbjlP_JgYRQWz1GymX2q7Reu-mKJEpao,10957
 flwr/cli/new/__init__.py,sha256=QA1E2QtzPvFCjLTUHnFnJbufuFiGyT_0Y53Wpbvg1F0,790
-flwr/cli/new/new.py,sha256=GvXapfYMrUQMktn1qJ3bfxfXeK0IAsxsxHBIGfPg3sE,10535
+flwr/cli/new/new.py,sha256=nIuUrQSGDjI4kqnymlq-rOT0RU3AHwZrat3abqHhCwM,10598
 flwr/cli/new/templates/__init__.py,sha256=FpjWCfIySU2DB4kh0HOXLAjlZNNFDTVU4w3HoE2TzcI,725
 flwr/cli/new/templates/app/.gitignore.tpl,sha256=HZJcGQoxp7aUzaPg8Uqch3kNrIESwr9yjimDxJYgXVY,3104
 flwr/cli/new/templates/app/LICENSE.tpl,sha256=WNHhf_5RCaeuKWyq_K39vmp9F28LxKsB4SpomwSZ2L0,11357
@@ -39,6 +39,7 @@ flwr/cli/new/templates/app/code/client.pytorch.py.tpl,sha256=fYoh-dTu07LkqNYvwcx
 flwr/cli/new/templates/app/code/client.pytorch_legacy_api.py.tpl,sha256=fuxVmZpjHIueNy_aHWF81531vmi8DGu4CYjYDqmUwWo,1705
 flwr/cli/new/templates/app/code/client.sklearn.py.tpl,sha256=0qqEe-RRjkHGOH8gsD9e83ae-kyyYixhyBgzVHjYpzk,3500
 flwr/cli/new/templates/app/code/client.tensorflow.py.tpl,sha256=8o55KXpsbF_rv6o98ZNYJDCazjwMp_RPTaSzDfT7Qlw,2682
+flwr/cli/new/templates/app/code/client.xgboost.py.tpl,sha256=-ipRV8gqpbEg7Mht77Yyqs1viL-3JYSVZR47I7xeG4c,3493
 flwr/cli/new/templates/app/code/dataset.baseline.py.tpl,sha256=jbd_exHAk2-Blu_kVutjPO6a_dkJQWb232zxSeXIZ1k,1453
 flwr/cli/new/templates/app/code/flwr_tune/__init__.py,sha256=Xq5fEn5yZkw6HAJi10T_3HRBoqN5_5pNqJHY4wXvD5k,748
 flwr/cli/new/templates/app/code/flwr_tune/client_app.py.tpl,sha256=p6BzTbP-mXkFANiVC7iz3YlskOidWaLC341IJyrUotQ,2951
@@ -56,6 +57,7 @@ flwr/cli/new/templates/app/code/server.pytorch.py.tpl,sha256=epARqfcQ-EQsdZwaaaU
 flwr/cli/new/templates/app/code/server.pytorch_legacy_api.py.tpl,sha256=gvBsGA_Jg9kAH8xTxjzTjMcvBtciuccOwQFbO7ey8tU,916
 flwr/cli/new/templates/app/code/server.sklearn.py.tpl,sha256=ehQ5VRgBn92WeFl6kupwJnuxSNkKvE-EvKde6A9mNQo,1377
 flwr/cli/new/templates/app/code/server.tensorflow.py.tpl,sha256=2-WTOPd-ewdLd9QmSlflIH7ix7zxAzPEOZoyiPBOy8c,1010
+flwr/cli/new/templates/app/code/server.xgboost.py.tpl,sha256=fwtCRyCG2hDSH1zVMyZv7zA7wsdKNPfpugDSZjxCs5Q,1746
 flwr/cli/new/templates/app/code/strategy.baseline.py.tpl,sha256=YkHAgppUeD2BnBoGfVB6dEvBfjuIPGsU1gw4CiUi3qA,40
 flwr/cli/new/templates/app/code/task.huggingface.py.tpl,sha256=piBbY3Dg60bQnCg15uzMw0QiL5SDOYX4YhQouy-X2zI,3164
 flwr/cli/new/templates/app/code/task.jax.py.tpl,sha256=Fb0XgdTAQplM-ZCusI081XA9asO3gHptH772S-Xcyy8,1525
@@ -65,6 +67,7 @@ flwr/cli/new/templates/app/code/task.pytorch.py.tpl,sha256=RKA5lV6O6OnVKZ2r75pbz
 flwr/cli/new/templates/app/code/task.pytorch_legacy_api.py.tpl,sha256=XlJqA4Ix_PloO_zJLhjiN5vDj16w3I4CPVGdmbe8asE,3800
 flwr/cli/new/templates/app/code/task.sklearn.py.tpl,sha256=vHdhtMp0FHxbYafXyhDT9aKmmmA0Jvpx5Oum1Yu9lWY,1850
 flwr/cli/new/templates/app/code/task.tensorflow.py.tpl,sha256=impgWN7MfztmcWF4xh1llcZGsgTvrb1HD5ZE0t-8U08,1731
+flwr/cli/new/templates/app/code/task.xgboost.py.tpl,sha256=0xO8jQvrHuB1llVDopQPOmt5Hn6rBw8umzoNwiZZs-o,2135
 flwr/cli/new/templates/app/code/utils.baseline.py.tpl,sha256=YkHAgppUeD2BnBoGfVB6dEvBfjuIPGsU1gw4CiUi3qA,40
 flwr/cli/new/templates/app/pyproject.baseline.toml.tpl,sha256=mPIMPfneVYV03l8jWRgWZ0V5Kh_pJw-AMUvkhcKkmL8,3182
 flwr/cli/new/templates/app/pyproject.flowertune.toml.tpl,sha256=wqYW4bWcf12m0U2njR995lySSesFvnHB-eSkPWz-QdM,2501
@@ -76,6 +79,7 @@ flwr/cli/new/templates/app/pyproject.pytorch.toml.tpl,sha256=SE4H23OFkQbqNU64nYf
 flwr/cli/new/templates/app/pyproject.pytorch_legacy_api.toml.tpl,sha256=docQbs3MuRR-yT24lVz7N2sQL3Sj49EHuOCuRj_0djQ,1508
 flwr/cli/new/templates/app/pyproject.sklearn.toml.tpl,sha256=apauU_PUmLEbt2rjckKniEbzdRs1EnMri_qgtHtBJZ8,1484
 flwr/cli/new/templates/app/pyproject.tensorflow.toml.tpl,sha256=LQpDKJTEnRKj5Ygn5FkT44SxlnLVprkPlbrGaFf5Q50,1508
+flwr/cli/new/templates/app/pyproject.xgboost.toml.tpl,sha256=504pHibNRGFe-DLnzqHLYhKeF_n8BPMv0Xog5EfnZ0M,1661
 flwr/cli/pull.py,sha256=dHiMe6x8w8yRoFNKpjA-eiPD6eFiHz4Vah5HZrqNpuo,3364
 flwr/cli/run/__init__.py,sha256=RPyB7KbYTFl6YRiilCch6oezxrLQrl1kijV7BMGkLbA,790
 flwr/cli/run/run.py,sha256=ECa0kup9dn15O70H74QdgUsEaeErbzDqVX_U0zZO5IM,8173
@@ -127,7 +131,7 @@ flwr/common/event_log_plugin/__init__.py,sha256=ts3VAL3Fk6Grp1EK_1Qg_V-BfOof9F86
 flwr/common/event_log_plugin/event_log_plugin.py,sha256=4SkVa1Ic-sPlICJShBuggXmXDcQtWQ1KDby4kthFNF0,2064
 flwr/common/exit/__init__.py,sha256=8W7xaO1iw0vacgmQW7FTFbSh7csNv6XfsgIlnIbNF6U,978
 flwr/common/exit/exit.py,sha256=DcXJfbpW1g-pQJqSZmps-1MZydd7T7RaarghIf2e4tU,3636
-flwr/common/exit/exit_code.py,sha256=e8O71zIqVT1H84mNBeenTz7S39yPZSpZQm-xUenpzN4,5249
+flwr/common/exit/exit_code.py,sha256=Xa1NFGny2cefZ62kZZOfT8eii__PolMWCHxYmxoSQ2s,5416
 flwr/common/exit/exit_handler.py,sha256=uzDdWwhKgc1w5csZS52b86kjmEApmDZKwMn_X0zDZZo,2126
 flwr/common/exit/signal_handler.py,sha256=wqxykrwgmpFzmEMhpnlM7RtO0PnqIvYiSB1qYahZ5Sk,3710
 flwr/common/grpc.py,sha256=nHnFC7E84pZVTvd6BhcSYWnGd0jf8t5UmGea04qvilM,9806
@@ -333,7 +337,7 @@ flwr/server/workflow/secure_aggregation/secaggplus_workflow.py,sha256=DkayCsnlAy
 flwr/serverapp/__init__.py,sha256=ZujKNXULwhWYQhFnxOOT5Wi9MRq2JCWFhAAj7ouiQ78,884
 flwr/serverapp/dp_fixed_clipping.py,sha256=wbP4W7CaUHXdll8ZSVUnTBSEWrnWM00CGk63rOR-Q2s,12133
 flwr/serverapp/exception.py,sha256=5cuH-2AafvihzosWDdDjuMmHdDqZ1XxHvCqZXNBVklw,1334
-flwr/serverapp/strategy/__init__.py,sha256=MJNWeBWzpQDfqhhpND5LncxPVK91kUo_Yzu6IMFdLCc,1492
+flwr/serverapp/strategy/__init__.py,sha256=mt2l31EAQ9oSvBcQhk4Jj4SvTePmWzBHQxZqL1v0uhE,1552
 flwr/serverapp/strategy/dp_fixed_clipping.py,sha256=wbP4W7CaUHXdll8ZSVUnTBSEWrnWM00CGk63rOR-Q2s,12133
 flwr/serverapp/strategy/fedadagrad.py,sha256=fD65P6OEERa_pxq847e1UZpA083AcWR44XavYB0naGM,6343
 flwr/serverapp/strategy/fedadam.py,sha256=s3xPIqhopy6yPTeFxevSPnc7a6BcKnKsvo2AaO6Z_xs,7138
@@ -344,11 +348,11 @@ flwr/serverapp/strategy/fedopt.py,sha256=kqT0uV2IUE93O72XEVa1JJo61dcwbZEoT9KmYTj
 flwr/serverapp/strategy/fedprox.py,sha256=XkkLTk3XpXAj0QoAzHqAvcAlPjrNlX11ISAu5u2x6X8,7026
 flwr/serverapp/strategy/fedtrimmedavg.py,sha256=4-QxgAQGo_7vB_L7qDYy28d95OBt9MeDa92yaTRMHqk,7166
 flwr/serverapp/strategy/fedxgb_bagging.py,sha256=ktDjzov4y0BRecioq788umCEtcuwElou9olBizQKOnM,3282
+flwr/serverapp/strategy/fedxgb_cyclic.py,sha256=8H8WoLdG4Fy1_dtLLE4AYiidC-Cvaw2GxySfzAb7Xj0,8774
 flwr/serverapp/strategy/fedyogi.py,sha256=1Ripr4Hi2cdeTOLiFOXtMKvOxR3BsUQwc7bbTrXN4LM,6653
 flwr/serverapp/strategy/result.py,sha256=E0Hl2VLnZAgQJjE2GDoKsK7JX-kPPU2KXc47Axt6hGw,4295
 flwr/serverapp/strategy/strategy.py,sha256=8uJGGm1ROLZERQ_dkRS7Z_rs-yK6XCE0UxXtIdFiEWk,10789
 flwr/serverapp/strategy/strategy_utils.py,sha256=hiwS7k-Hx6_c4NZXoKpHucS5CBKb7f8GppXRBSMt3Us,10851
-flwr/serverapp/strategy/strategy_utils_tests.py,sha256=_adS23Lrv1QA6V_3oZ7P_csMd8RqDObFeIhOkFnNtTg,10690
 flwr/simulation/__init__.py,sha256=Gg6OsP1Z-ixc3-xxzvl7j7rz2Fijy9rzyEPpxgAQCeM,1556
 flwr/simulation/app.py,sha256=b_bDyZFwBf2zpKs37Vmd5cFJSzDRE0fL-8uqA0UkAv4,10393
 flwr/simulation/legacy_app.py,sha256=nMISQqW0otJL1-2Kfd94O6BLlGS2IEmEPKTM2WGKrIs,15861
@@ -361,7 +365,8 @@ flwr/simulation/simulationio_connection.py,sha256=mzS1C6EEREwQDPceDo30anAasmTDLb
 flwr/supercore/__init__.py,sha256=pqkFoow_E6UhbBlhmoD1gmTH-33yJRhBsIZqxRPFZ7U,755
 flwr/supercore/app_utils.py,sha256=K76Zt6R670b1hUmxOsNc1WUCVYvF7lejXPcCO9K0Q0g,1753
 flwr/supercore/cli/__init__.py,sha256=EDl2aO-fuQfxSbL-T1W9RAfA2N0hpWHmqX_GSwblJbQ,845
-flwr/supercore/cli/flower_superexec.py,sha256=kov4uEeihf7QEUAfHEgdEvsL_8nL_fzQI9EePnRM1Ww,5012
+flwr/supercore/cli/flower_superexec.py,sha256=JtqYrEWVu3BxLkjavsdohTOwvMwzuFqWP5j4Mo9dqsk,6155
+flwr/supercore/constant.py,sha256=F9kRjisedaZcoyGvUITSDmIG12QDSCpo2LlM_l-q6jM,820
 flwr/supercore/corestate/__init__.py,sha256=Vau6-L_JG5QzNqtCTa9xCKGGljc09wY8avZmIjSJemg,774
 flwr/supercore/corestate/corestate.py,sha256=rDAWWeG5DcpCyQso9Z3RhwL4zr2IroPlRMcDzqoSu8s,2328
 flwr/supercore/ffs/__init__.py,sha256=U3KXwG_SplEvchat27K0LYPoPHzh-cwwT_NHsGlYMt8,908
@@ -382,10 +387,10 @@ flwr/supercore/superexec/__init__.py,sha256=XKX208hZ6a9gZ4KT9kMqfpCtp_8VGxekzKFf
 flwr/supercore/superexec/plugin/__init__.py,sha256=GNwq8uNdE8RB7ywEFRAvKjLFzgS3YXgz39-HBGsemWw,1035
 flwr/supercore/superexec/plugin/base_exec_plugin.py,sha256=fL-Ufc9Dp56OhWOzNSJUc7HumbkuSDYqZKwde2opG4g,2074
 flwr/supercore/superexec/plugin/clientapp_exec_plugin.py,sha256=9FT6ufEqV5K9g4FaAB9lVDbIv-VCH5LcxT4YKy23roE,1035
-flwr/supercore/superexec/plugin/exec_plugin.py,sha256=w3jmtxdv7ov_EdAgifKcm4q8nV39e2Xna4sNjqClwOM,2447
+flwr/supercore/superexec/plugin/exec_plugin.py,sha256=4WtCQ4bsuFRlfCbg91ZcPAsX8htrCCo_fFh1DKo3cCQ,2764
 flwr/supercore/superexec/plugin/serverapp_exec_plugin.py,sha256=IwRzdPV-cSKwrP2krGh0De4IkAuxsmgK0WU6J-2GXqM,1035
 flwr/supercore/superexec/plugin/simulation_exec_plugin.py,sha256=upn5zE-YKkl_jTw8RzmeyQ58PU_UAlQ7CqnBXXdng8I,1060
-flwr/supercore/superexec/run_superexec.py,sha256=8hUlaVPVNnhePQ9OUgen4yy0fSGZAVggBGzm-33iJPw,6630
+flwr/supercore/superexec/run_superexec.py,sha256=JiwKq9s_WPpk0S9MSi1lIgMZU120NOZLf4GlObHzI_k,7217
 flwr/supercore/utils.py,sha256=ebuHMbeA8eXisX0oMPqBK3hk7uVnIE_yiqWVz8YbkpQ,1324
 flwr/superlink/__init__.py,sha256=GNSuJ4-N6Z8wun2iZNlXqENt5beUyzC0Gi_tN396bbM,707
 flwr/superlink/artifact_provider/__init__.py,sha256=pgZEcVPKRE874LSu3cgy0HbwSJBIpVy_HxQOmne4PAs,810
@@ -411,7 +416,7 @@ flwr/supernode/servicer/__init__.py,sha256=lucTzre5WPK7G1YLCfaqg3rbFWdNSb7ZTt-ca
 flwr/supernode/servicer/clientappio/__init__.py,sha256=7Oy62Y_oijqF7Dxi6tpcUQyOpLc_QpIRZ83NvwmB0Yg,813
 flwr/supernode/servicer/clientappio/clientappio_servicer.py,sha256=nIHRu38EWK-rpNOkcgBRAAKwYQQWFeCwu0lkO7OPZGQ,10239
 flwr/supernode/start_client_internal.py,sha256=Y9S1-QlO2WP6eo4JvWzIpfaCoh2aoE7bjEYyxNNnlyg,20777
-flwr_nightly-1.22.0.dev20250917.dist-info/METADATA,sha256=NGqW090jvJZ7dg66AQQr9jKb4SF7DdF78xwqFuhtFUE,14559
-flwr_nightly-1.22.0.dev20250917.dist-info/WHEEL,sha256=b4K_helf-jlQoXBBETfwnf4B04YC67LOev0jo4fX5m8,88
-flwr_nightly-1.22.0.dev20250917.dist-info/entry_points.txt,sha256=hxHD2ixb_vJFDOlZV-zB4Ao32_BQlL34ftsDh1GXv14,420
-flwr_nightly-1.22.0.dev20250917.dist-info/RECORD,,
+flwr_nightly-1.22.0.dev20250918.dist-info/METADATA,sha256=XdvpDYzZ_dWvMX49Pwi2gOUo0qLtrt466UABjDeyavg,14559
+flwr_nightly-1.22.0.dev20250918.dist-info/WHEEL,sha256=b4K_helf-jlQoXBBETfwnf4B04YC67LOev0jo4fX5m8,88
+flwr_nightly-1.22.0.dev20250918.dist-info/entry_points.txt,sha256=hxHD2ixb_vJFDOlZV-zB4Ao32_BQlL34ftsDh1GXv14,420
+flwr_nightly-1.22.0.dev20250918.dist-info/RECORD,,

flwr/serverapp/strategy/strategy_utils_tests.py

@@ -1,323 +0,0 @@
-# Copyright 2025 Flower Labs GmbH. All Rights Reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-# ==============================================================================
-"""Tests for message-based strategy utilities."""
-
-
-from collections import OrderedDict
-from unittest.mock import Mock
-
-import numpy as np
-import pytest
-from parameterized import parameterized
-
-from flwr.common import (
-    Array,
-    ArrayRecord,
-    ConfigRecord,
-    Message,
-    MetricRecord,
-    RecordDict,
-)
-from flwr.serverapp.exception import InconsistentMessageReplies
-
-from .strategy_utils import (
-    aggregate_arrayrecords,
-    aggregate_metricrecords,
-    config_to_str,
-    validate_message_reply_consistency,
-)
-
-
-def create_mock_reply(arrays: ArrayRecord, num_examples: float) -> Message:
-    """Create a mock reply Message with default keys."""
-    message = Mock(spec=Message)
-    message.content = RecordDict(
-        {"arrays": arrays, "metrics": MetricRecord({"num-examples": num_examples})}
-    )
-    message.has_error.side_effect = lambda: False
-    message.has_content.side_effect = lambda: True
-    return message
-
-
-def test_config_to_str() -> None:
-    """Test that items of types bytes are masked out."""
-    config = ConfigRecord({"a": 123, "b": [1, 2, 3], "c": b"bytes"})
-    expected_str = "{'a': 123, 'b': [1, 2, 3], 'c': <bytes>}"
-    assert config_to_str(config) == expected_str
-
-
-def test_arrayrecords_aggregation() -> None:
-    """Test aggregation of ArrayRecords."""
-    num_replies = 3
-    num_arrays = 4
-    weights = [0.25, 0.4, 0.35]
-    np_arrays = [
-        [np.random.randn(7, 3) for _ in range(num_arrays)] for _ in range(num_replies)
-    ]
-
-    avg_list = [
-        np.average([lst[i] for lst in np_arrays], axis=0, weights=weights)
-        for i in range(num_arrays)
-    ]
-
-    # Construct RecordDicts (mimicing replies)
-    records = [
-        RecordDict(
-            {
-                "arrays": ArrayRecord(np_arrays[i]),
-                "metrics": MetricRecord({"weight": weights[i]}),
-            }
-        )
-        for i in range(num_replies)
-    ]
-    # Execute aggregate
-    aggrd = aggregate_arrayrecords(records, weighting_metric_name="weight")
-
-    # Assert consistency
-    assert all(np.allclose(a, b) for a, b in zip(aggrd.to_numpy_ndarrays(), avg_list))
-    assert aggrd.object_id == ArrayRecord(avg_list).object_id
-
-
-def test_arrayrecords_aggregation_with_ndim_zero() -> None:
-    """Test aggregation of ArrayRecords with 0-dim arrays."""
-    num_replies = 3
-    weights = [0.25, 0.4, 0.35]
-    np_arrays = [np.array(np.random.randn()) for _ in range(num_replies)]
-
-    # For 0-dimensional arrays, we just compute the weighted average directly
-    avg_list = [np.average(np_arrays, axis=0, weights=weights)]
-
-    # Construct RecordDicts (mimicing replies)
-    records = [
-        RecordDict(
-            {
-                "arrays": ArrayRecord([np_arrays[i]]),
-                "metrics": MetricRecord({"weight": weights[i]}),
-            }
-        )
-        for i in range(num_replies)
-    ]
-    # Execute aggregate
-    aggrd = aggregate_arrayrecords(records, weighting_metric_name="weight")
-
-    # Assert consistency
-    assert np.isclose(aggrd.to_numpy_ndarrays()[0], avg_list[0])
-    assert aggrd.object_id == ArrayRecord([np.array(avg_list[0])]).object_id
-
-
-def test_metricrecords_aggregation() -> None:
-    """Test aggregation of MetricRecords."""
-    num_replies = 3
-    weights = [0.25, 0.4, 0.35]
-    metric_records = [
-        MetricRecord({"a": 1, "b": 2.0, "c": np.random.randn(3).tolist()})
-        for _ in range(num_replies)
-    ]
-
-    # Compute expected aggregated MetricRecord.
-    # For ease, we convert everything into numpy arrays, then aggregate
-    as_np_entries = [
-        {
-            k: np.array(v) if isinstance(v, (int, float, list)) else v
-            for k, v in record.items()
-        }
-        for record in metric_records
-    ]
-    avg_list = [
-        np.average(
-            [list(entries.values())[i] for entries in as_np_entries],
-            axis=0,
-            weights=weights,
-        ).tolist()
-        for i in range(len(as_np_entries[0]))
-    ]
-    expected_record = MetricRecord(dict(zip(as_np_entries[0].keys(), avg_list)))
-    expected_record["a"] = float(expected_record["a"])  # type: ignore
-    expected_record["b"] = float(expected_record["b"])  # type: ignore
-
-    # Construct RecordDicts (mimicing replies)
-    # Inject weighting factor
-    records = [
-        RecordDict(
-            {
-                "metrics": MetricRecord(
-                    record.__dict__["_data"] | {"weight": weights[i]}
-                ),
-            }
-        )
-        for i, record in enumerate(metric_records)
-    ]
-
-    # Execute aggregate
-    aggrd = aggregate_metricrecords(records, weighting_metric_name="weight")
-    # Assert
-    assert expected_record.object_id == aggrd.object_id
-
-
-@parameterized.expand(  # type: ignore
-    [
-        (
-            True,
-            RecordDict(
-                {
-                    "global-model": ArrayRecord([np.random.randn(7, 3)]),
-                    "metrics": MetricRecord({"weight": 0.123}),
-                }
-            ),
-        ),  # Compliant
-        (
-            False,
-            RecordDict(
-                {
-                    "global-model": ArrayRecord([np.random.randn(7, 3)]),
-                    "metrics": MetricRecord({"weight": [0.123]}),
-                }
-            ),
-        ),  # Weighting key is not a scalar (BAD)
-        (
-            False,
-            RecordDict(
-                {
-                    "global-model": ArrayRecord([np.random.randn(7, 3)]),
-                    "metrics": MetricRecord({"loss": 0.01}),
-                }
-            ),
-        ),  # No weighting key in MetricRecord (BAD)
-        (
-            False,
-            RecordDict({"global-model": ArrayRecord([np.random.randn(7, 3)])}),
-        ),  # No MetricsRecord (BAD)
-        (
-            False,
-            RecordDict(
-                {
-                    "global-model": ArrayRecord([np.random.randn(7, 3)]),
-                    "another-model": ArrayRecord([np.random.randn(7, 3)]),
-                }
-            ),
-        ),  # Two ArrayRecords (BAD)
-        (
-            False,
-            RecordDict(
-                {
-                    "global-model": ArrayRecord([np.random.randn(7, 3)]),
-                    "metrics": MetricRecord({"weight": 0.123}),
-                    "more-metrics": MetricRecord({"loss": 0.321}),
-                }
-            ),
-        ),  # Two MetricRecords (BAD)
-    ]
-)
-def test_consistency_of_replies_with_matching_keys(
-    is_valid: bool, recorddict: RecordDict
-) -> None:
-    """Test consistency in replies."""
-    # Create dummy records
-    records = [recorddict for _ in range(3)]
-
-    if not is_valid:
-        # Should raise InconsistentMessageReplies exception
-        with pytest.raises(InconsistentMessageReplies):
-            validate_message_reply_consistency(
-                records, weighted_by_key="weight", check_arrayrecord=True
-            )
-    else:
-        # Should not raise an exception
-        validate_message_reply_consistency(
-            records, weighted_by_key="weight", check_arrayrecord=True
-        )
-
-
-@parameterized.expand(  # type: ignore
-    [
-        (
-            [
-                RecordDict(
-                    {
-                        "global-model": ArrayRecord([np.random.randn(7, 3)]),
-                        "metrics": MetricRecord({"weight": 0.123}),
-                    }
-                ),
-                RecordDict(
-                    {
-                        "model": ArrayRecord([np.random.randn(7, 3)]),
-                        "metrics": MetricRecord({"weight": 0.123}),
-                    }
-                ),
-            ],
-        ),  # top-level keys don't match for ArrayRecords
-        (
-            [
-                RecordDict(
-                    {
-                        "global-model": ArrayRecord(
-                            OrderedDict({"a": Array(np.random.randn(7, 3))})
-                        ),
-                        "metrics": MetricRecord({"weight": 0.123}),
-                    }
-                ),
-                RecordDict(
-                    {
-                        "global-model": ArrayRecord(
-                            OrderedDict({"b": Array(np.random.randn(7, 3))})
-                        ),
-                        "metrics": MetricRecord({"weight": 0.123}),
-                    }
-                ),
-            ],
-        ),  # top-level keys match for ArrayRecords but not those for Arrays
-        (
-            [
-                RecordDict(
-                    {
-                        "global-model": ArrayRecord([np.random.randn(7, 3)]),
-                        "metrics": MetricRecord({"weight": 0.123}),
-                    }
-                ),
-                RecordDict(
-                    {
-                        "global-model": ArrayRecord([np.random.randn(7, 3)]),
-                        "my-metrics": MetricRecord({"weight": 0.123}),
-                    }
-                ),
-            ],
-        ),  # top-level keys don't match for MetricRecords
-        (
-            [
-                RecordDict(
-                    {
-                        "global-model": ArrayRecord([np.random.randn(7, 3)]),
-                        "metrics": MetricRecord({"weight": 0.123}),
-                    }
-                ),
-                RecordDict(
-                    {
-                        "global-model": ArrayRecord([np.random.randn(7, 3)]),
-                        "my-metrics": MetricRecord({"my-weights": 0.123}),
-                    }
-                ),
-            ],
-        ),  # top-level keys match for MetricRecords but not inner ones
-    ]
-)
-def test_consistency_of_replies_with_different_keys(
-    list_records: list[RecordDict],
-) -> None:
-    """Test consistency in replies when records don't have matching keys."""
-    # All test cases expect InconsistentMessageReplies exception to be raised
-    with pytest.raises(InconsistentMessageReplies):
-        validate_message_reply_consistency(
-            list_records, weighted_by_key="weight", check_arrayrecord=True
-        )