flwr 1.20.0__py3-none-any.whl → 1.22.0__py3-none-any.whl

flwr/serverapp/strategy/strategy.py

@@ -0,0 +1,285 @@
+# 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 strategy."""
+
+
+import io
+import time
+from abc import ABC, abstractmethod
+from collections.abc import Iterable
+from logging import INFO
+from typing import Callable, Optional
+
+from flwr.common import ArrayRecord, ConfigRecord, Message, MetricRecord, log
+from flwr.server import Grid
+
+from .result import Result
+from .strategy_utils import log_strategy_start_info
+
+
+class Strategy(ABC):
+    """Abstract base class for server strategy implementations."""
+
+    @abstractmethod
+    def configure_train(
+        self, server_round: int, arrays: ArrayRecord, config: ConfigRecord, grid: Grid
+    ) -> Iterable[Message]:
+        """Configure the next round of training.
+
+        Parameters
+        ----------
+        server_round : int
+            The current round of federated learning.
+        arrays : ArrayRecord
+            Current global ArrayRecord (e.g. global model) to be sent to client
+            nodes for training.
+        config : ConfigRecord
+            Configuration to be sent to clients nodes for training.
+        grid : Grid
+            The Grid instance used for node sampling and communication.
+
+        Returns
+        -------
+        Iterable[Message]
+            An iterable of messages to be sent to selected client nodes for training.
+        """
+
+    @abstractmethod
+    def aggregate_train(
+        self,
+        server_round: int,
+        replies: Iterable[Message],
+    ) -> tuple[Optional[ArrayRecord], Optional[MetricRecord]]:
+        """Aggregate training results from client nodes.
+
+        Parameters
+        ----------
+        server_round : int
+            The current round of federated learning, starting from 1.
+        replies : Iterable[Message]
+            Iterable of reply messages received from client nodes after training.
+            Each message contains ArrayRecords and MetricRecords that get aggregated.
+
+        Returns
+        -------
+        tuple[Optional[ArrayRecord], Optional[MetricRecord]]
+            A tuple containing:
+            - ArrayRecord: Aggregated ArrayRecord, or None if aggregation failed
+            - MetricRecord: Aggregated MetricRecord, or None if aggregation failed
+        """
+
+    @abstractmethod
+    def configure_evaluate(
+        self, server_round: int, arrays: ArrayRecord, config: ConfigRecord, grid: Grid
+    ) -> Iterable[Message]:
+        """Configure the next round of evaluation.
+
+        Parameters
+        ----------
+        server_round : int
+            The current round of federated learning.
+        arrays : ArrayRecord
+            Current global ArrayRecord (e.g. global model) to be sent to client
+            nodes for evaluation.
+        config : ConfigRecord
+            Configuration to be sent to clients nodes for evaluation.
+        grid : Grid
+            The Grid instance used for node sampling and communication.
+
+        Returns
+        -------
+        Iterable[Message]
+            An iterable of messages to be sent to selected client nodes for evaluation.
+        """
+
+    @abstractmethod
+    def aggregate_evaluate(
+        self,
+        server_round: int,
+        replies: Iterable[Message],
+    ) -> Optional[MetricRecord]:
+        """Aggregate evaluation metrics from client nodes.
+
+        Parameters
+        ----------
+        server_round : int
+            The current round of federated learning.
+        replies : Iterable[Message]
+            Iterable of reply messages received from client nodes after evaluation.
+            MetricRecords in the messages are aggregated.
+
+        Returns
+        -------
+        Optional[MetricRecord]
+            Aggregated evaluation metrics from all participating clients,
+            or None if aggregation failed.
+        """
+
+    @abstractmethod
+    def summary(self) -> None:
+        """Log summary configuration of the strategy."""
+
+    # pylint: disable=too-many-arguments, too-many-positional-arguments, too-many-locals
+    def start(
+        self,
+        grid: Grid,
+        initial_arrays: ArrayRecord,
+        num_rounds: int = 3,
+        timeout: float = 3600,
+        train_config: Optional[ConfigRecord] = None,
+        evaluate_config: Optional[ConfigRecord] = None,
+        evaluate_fn: Optional[
+            Callable[[int, ArrayRecord], Optional[MetricRecord]]
+        ] = None,
+    ) -> Result:
+        """Execute the federated learning strategy.
+
+        Runs the complete federated learning workflow for the specified number of
+        rounds, including training, evaluation, and optional centralized evaluation.
+
+        Parameters
+        ----------
+        grid : Grid
+            The Grid instance used to send/receive Messages from nodes executing a
+            ClientApp.
+        initial_arrays : ArrayRecord
+            Initial model parameters (arrays) to be used for federated learning.
+        num_rounds : int (default: 3)
+            Number of federated learning rounds to execute.
+        timeout : float (default: 3600)
+            Timeout in seconds for waiting for node responses.
+        train_config : ConfigRecord, optional
+            Configuration to be sent to nodes during training rounds.
+            If unset, an empty ConfigRecord will be used.
+        evaluate_config : ConfigRecord, optional
+            Configuration to be sent to nodes during evaluation rounds.
+            If unset, an empty ConfigRecord will be used.
+        evaluate_fn : Callable[[int, ArrayRecord], Optional[MetricRecord]], optional
+            Optional function for centralized evaluation of the global model. Takes
+            server round number and array record, returns a MetricRecord or None. If
+            provided, will be called before the first round and after each round.
+            Defaults to None.
+
+        Returns
+        -------
+        Results
+            Results containing final model arrays and also training metrics, evaluation
+            metrics and global evaluation metrics (if provided) from all rounds.
+        """
+        log(INFO, "Starting %s strategy:", self.__class__.__name__)
+        log_strategy_start_info(
+            num_rounds, initial_arrays, train_config, evaluate_config
+        )
+        self.summary()
+        log(INFO, "")
+
+        # Initialize if None
+        train_config = ConfigRecord() if train_config is None else train_config
+        evaluate_config = ConfigRecord() if evaluate_config is None else evaluate_config
+        result = Result()
+
+        t_start = time.time()
+        # Evaluate starting global parameters
+        if evaluate_fn:
+            res = evaluate_fn(0, initial_arrays)
+            log(INFO, "Initial global evaluation results: %s", res)
+            if res is not None:
+                result.evaluate_metrics_serverapp[0] = res
+
+        arrays = initial_arrays
+
+        for current_round in range(1, num_rounds + 1):
+            log(INFO, "")
+            log(INFO, "[ROUND %s/%s]", current_round, num_rounds)
+
+            # -----------------------------------------------------------------
+            # --- TRAINING (CLIENTAPP-SIDE) -----------------------------------
+            # -----------------------------------------------------------------
+
+            # Call strategy to configure training round
+            # Send messages and wait for replies
+            train_replies = grid.send_and_receive(
+                messages=self.configure_train(
+                    current_round,
+                    arrays,
+                    train_config,
+                    grid,
+                ),
+                timeout=timeout,
+            )
+
+            # Aggregate train
+            agg_arrays, agg_train_metrics = self.aggregate_train(
+                current_round,
+                train_replies,
+            )
+
+            # Log training metrics and append to history
+            if agg_arrays is not None:
+                result.arrays = agg_arrays
+                arrays = agg_arrays
+            if agg_train_metrics is not None:
+                log(INFO, "\t└──> Aggregated MetricRecord: %s", agg_train_metrics)
+                result.train_metrics_clientapp[current_round] = agg_train_metrics
+
+            # -----------------------------------------------------------------
+            # --- EVALUATION (CLIENTAPP-SIDE) ---------------------------------
+            # -----------------------------------------------------------------
+
+            # Call strategy to configure evaluation round
+            # Send messages and wait for replies
+            evaluate_replies = grid.send_and_receive(
+                messages=self.configure_evaluate(
+                    current_round,
+                    arrays,
+                    evaluate_config,
+                    grid,
+                ),
+                timeout=timeout,
+            )
+
+            # Aggregate evaluate
+            agg_evaluate_metrics = self.aggregate_evaluate(
+                current_round,
+                evaluate_replies,
+            )
+
+            # Log training metrics and append to history
+            if agg_evaluate_metrics is not None:
+                log(INFO, "\t└──> Aggregated MetricRecord: %s", agg_evaluate_metrics)
+                result.evaluate_metrics_clientapp[current_round] = agg_evaluate_metrics
+
+            # -----------------------------------------------------------------
+            # --- EVALUATION (SERVERAPP-SIDE) ---------------------------------
+            # -----------------------------------------------------------------
+
+            # Centralized evaluation
+            if evaluate_fn:
+                log(INFO, "Global evaluation")
+                res = evaluate_fn(current_round, arrays)
+                log(INFO, "\t└──> MetricRecord: %s", res)
+                if res is not None:
+                    result.evaluate_metrics_serverapp[current_round] = res
+
+        log(INFO, "")
+        log(INFO, "Strategy execution finished in %.2fs", time.time() - t_start)
+        log(INFO, "")
+        log(INFO, "Final results:")
+        log(INFO, "")
+        for line in io.StringIO(str(result)):
+            log(INFO, "\t%s", line.strip("\n"))
+        log(INFO, "")
+
+        return result

flwr/serverapp/strategy/strategy_utils.py

@@ -0,0 +1,299 @@
+# 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 strategy utilities."""
+
+
+import json
+import random
+from collections import OrderedDict
+from logging import INFO
+from time import sleep
+from typing import Optional, cast
+
+import numpy as np
+
+from flwr.common import (
+    Array,
+    ArrayRecord,
+    ConfigRecord,
+    MetricRecord,
+    NDArray,
+    RecordDict,
+    log,
+)
+from flwr.server import Grid
+
+from ..exception import InconsistentMessageReplies
+
+
+def config_to_str(config: ConfigRecord) -> str:
+    """Convert a ConfigRecord to a string representation masking bytes."""
+    content = ", ".join(
+        f"'{k}': {'<bytes>' if isinstance(v, bytes) else v}" for k, v in config.items()
+    )
+    return f"{{{content}}}"
+
+
+def log_strategy_start_info(
+    num_rounds: int,
+    arrays: ArrayRecord,
+    train_config: Optional[ConfigRecord],
+    evaluate_config: Optional[ConfigRecord],
+) -> None:
+    """Log information about the strategy start."""
+    log(INFO, "\t├── Number of rounds: %d", num_rounds)
+    log(
+        INFO,
+        "\t├── ArrayRecord (%.2f MB)",
+        sum(len(array.data) for array in arrays.values()) / (1024**2),
+    )
+    log(
+        INFO,
+        "\t├── ConfigRecord (train): %s",
+        config_to_str(train_config) if train_config else "(empty!)",
+    )
+    log(
+        INFO,
+        "\t├── ConfigRecord (evaluate): %s",
+        config_to_str(evaluate_config) if evaluate_config else "(empty!)",
+    )
+
+
+def aggregate_arrayrecords(
+    records: list[RecordDict], weighting_metric_name: str
+) -> ArrayRecord:
+    """Perform weighted aggregation all ArrayRecords using a specific key."""
+    # Retrieve weighting factor from MetricRecord
+    weights: list[float] = []
+    for record in records:
+        # Get the first (and only) MetricRecord in the record
+        metricrecord = next(iter(record.metric_records.values()))
+        # Because replies have been checked for consistency,
+        # we can safely cast the weighting factor to float
+        w = cast(float, metricrecord[weighting_metric_name])
+        weights.append(w)
+
+    # Average
+    total_weight = sum(weights)
+    weight_factors = [w / total_weight for w in weights]
+
+    # Perform weighted aggregation
+    aggregated_np_arrays: dict[str, NDArray] = {}
+
+    for record, weight in zip(records, weight_factors):
+        for record_item in record.array_records.values():
+            # aggregate in-place
+            for key, value in record_item.items():
+                if key not in aggregated_np_arrays:
+                    aggregated_np_arrays[key] = value.numpy() * weight
+                else:
+                    aggregated_np_arrays[key] += value.numpy() * weight
+
+    return ArrayRecord(
+        OrderedDict({k: Array(np.asarray(v)) for k, v in aggregated_np_arrays.items()})
+    )
+
+
+def aggregate_metricrecords(
+    records: list[RecordDict], weighting_metric_name: str
+) -> MetricRecord:
+    """Perform weighted aggregation all MetricRecords using a specific key."""
+    # Retrieve weighting factor from MetricRecord
+    weights: list[float] = []
+    for record in records:
+        # Get the first (and only) MetricRecord in the record
+        metricrecord = next(iter(record.metric_records.values()))
+        # Because replies have been checked for consistency,
+        # we can safely cast the weighting factor to float
+        w = cast(float, metricrecord[weighting_metric_name])
+        weights.append(w)
+
+    # Average
+    total_weight = sum(weights)
+    weight_factors = [w / total_weight for w in weights]
+
+    aggregated_metrics = MetricRecord()
+    for record, weight in zip(records, weight_factors):
+        for record_item in record.metric_records.values():
+            # aggregate in-place
+            for key, value in record_item.items():
+                if key == weighting_metric_name:
+                    # We exclude the weighting key from the aggregated MetricRecord
+                    continue
+                if key not in aggregated_metrics:
+                    if isinstance(value, list):
+                        aggregated_metrics[key] = [v * weight for v in value]
+                    else:
+                        aggregated_metrics[key] = value * weight
+                else:
+                    if isinstance(value, list):
+                        current_list = cast(list[float], aggregated_metrics[key])
+                        aggregated_metrics[key] = [
+                            curr + val * weight
+                            for curr, val in zip(current_list, value)
+                        ]
+                    else:
+                        current_value = cast(float, aggregated_metrics[key])
+                        aggregated_metrics[key] = current_value + value * weight
+
+    return aggregated_metrics
+
+
+def sample_nodes(
+    grid: Grid, min_available_nodes: int, sample_size: int
+) -> tuple[list[int], list[int]]:
+    """Sample the specified number of nodes using the Grid.
+
+    Parameters
+    ----------
+    grid : Grid
+        The grid object.
+    min_available_nodes : int
+        The minimum number of available nodes to sample from.
+    sample_size : int
+        The number of nodes to sample.
+
+    Returns
+    -------
+    tuple[list[int], list[int]]
+        A tuple containing the sampled node IDs and the list
+        of all connected node IDs.
+    """
+    sampled_nodes = []
+
+    # Ensure min_available_nodes is at least as large as sample_size
+    min_available_nodes = max(min_available_nodes, sample_size)
+
+    # wait for min_available_nodes to be online
+    while len(all_nodes := list(grid.get_node_ids())) < min_available_nodes:
+        log(
+            INFO,
+            "Waiting for nodes to connect: %d connected (minimum required: %d).",
+            len(all_nodes),
+            min_available_nodes,
+        )
+        sleep(1)
+
+    # Sample nodes
+    sampled_nodes = random.sample(all_nodes, sample_size)
+
+    return sampled_nodes, all_nodes
+
+
+# pylint: disable=too-many-return-statements
+def validate_message_reply_consistency(
+    replies: list[RecordDict], weighted_by_key: str, check_arrayrecord: bool
+) -> None:
+    """Validate that replies contain exactly one ArrayRecord and one MetricRecord, and
+    that the MetricRecord includes a weight factor key.
+
+    These checks ensure that Message-based strategies behave consistently with
+    *Ins/*Res-based strategies.
+    """
+    # Checking for ArrayRecord consistency
+    if check_arrayrecord:
+        if any(len(msg.array_records) != 1 for msg in replies):
+            raise InconsistentMessageReplies(
+                reason="Expected exactly one ArrayRecord in replies. "
+                "Skipping aggregation."
+            )
+
+        # Ensure all key are present in all ArrayRecords
+        record_key = next(iter(replies[0].array_records.keys()))
+        all_keys = set(replies[0][record_key].keys())
+        if any(set(msg.get(record_key, {}).keys()) != all_keys for msg in replies[1:]):
+            raise InconsistentMessageReplies(
+                reason="All ArrayRecords must have the same keys for aggregation. "
+                "This condition wasn't met. Skipping aggregation."
+            )
+
+    # Checking for MetricRecord consistency
+    if any(len(msg.metric_records) != 1 for msg in replies):
+        raise InconsistentMessageReplies(
+            reason="Expected exactly one MetricRecord in replies, but found more. "
+            "Skipping aggregation."
+        )
+
+    # Ensure all key are present in all MetricRecords
+    record_key = next(iter(replies[0].metric_records.keys()))
+    all_keys = set(replies[0][record_key].keys())
+    if any(set(msg.get(record_key, {}).keys()) != all_keys for msg in replies[1:]):
+        raise InconsistentMessageReplies(
+            reason="All MetricRecords must have the same keys for aggregation. "
+            "This condition wasn't met. Skipping aggregation."
+        )
+
+    # Verify the weight factor key presence in all MetricRecords
+    if weighted_by_key not in all_keys:
+        raise InconsistentMessageReplies(
+            reason=f"Missing required key `{weighted_by_key}` in the MetricRecord of "
+            "reply messages. Cannot average ArrayRecords and MetricRecords. Skipping "
+            "aggregation."
+        )
+
+    # Check that it is not a list
+    if any(isinstance(msg[record_key][weighted_by_key], list) for msg in replies):
+        raise InconsistentMessageReplies(
+            reason=f"Key `{weighted_by_key}` in the MetricRecord of reply messages "
+            "must be a single value (int or float), but a list was found. Skipping "
+            "aggregation."
+        )
+
+
+def aggregate_bagging(
+    bst_prev_org: bytes,
+    bst_curr_org: bytes,
+) -> bytes:
+    """Conduct bagging aggregation for given trees."""
+    if bst_prev_org == b"":
+        return bst_curr_org
+
+    # Get the tree numbers
+    tree_num_prev, _ = _get_tree_nums(bst_prev_org)
+    _, paral_tree_num_curr = _get_tree_nums(bst_curr_org)
+
+    bst_prev = json.loads(bytearray(bst_prev_org))
+    bst_curr = json.loads(bytearray(bst_curr_org))
+
+    previous_model = bst_prev["learner"]["gradient_booster"]["model"]
+    previous_model["gbtree_model_param"]["num_trees"] = str(
+        tree_num_prev + paral_tree_num_curr
+    )
+    iteration_indptr = previous_model["iteration_indptr"]
+    previous_model["iteration_indptr"].append(
+        iteration_indptr[-1] + paral_tree_num_curr
+    )
+
+    # Aggregate new trees
+    trees_curr = bst_curr["learner"]["gradient_booster"]["model"]["trees"]
+    for tree_count in range(paral_tree_num_curr):
+        trees_curr[tree_count]["id"] = tree_num_prev + tree_count
+        previous_model["trees"].append(trees_curr[tree_count])
+        previous_model["tree_info"].append(0)
+
+    bst_prev_bytes = bytes(json.dumps(bst_prev), "utf-8")
+
+    return bst_prev_bytes
+
+
+def _get_tree_nums(xgb_model_org: bytes) -> tuple[int, int]:
+    xgb_model = json.loads(bytearray(xgb_model_org))
+
+    # Access model parameters
+    model_param = xgb_model["learner"]["gradient_booster"]["model"][
+        "gbtree_model_param"
+    ]
+    # Return the number of trees and the number of parallel trees
+    return int(model_param["num_trees"]), int(model_param["num_parallel_tree"])