Repository: NVIDIA/NVFlare

Issue:
Deprecate `AccumulateWeightedAggregator`
`AccumulateWeightedAggregator` should be deprecated. We have a more memory-efficient and more flexible version with `InTimeAccumulateWeightedAggregator` (e.g. support for DXO collections).

Examples using `AccumulateWeightedAggregator` should be updated.


Discussion:
@YuanTingHsieh @IsaacYangSLA, any opinions on this one?

Resolving patch:
diff --git a/nvflare/app_common/aggregators/accumulate_model_aggregator.py b/nvflare/app_common/aggregators/accumulate_model_aggregator.py
--- a/nvflare/app_common/aggregators/accumulate_model_aggregator.py
+++ b/nvflare/app_common/aggregators/accumulate_model_aggregator.py
@@ -12,197 +12,12 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-import re
+from .intime_accumulate_model_aggregator import InTimeAccumulateWeightedAggregator
 
-import numpy as np
 
-from nvflare.apis.dxo import DXO, DataKind, MetaKey, from_shareable
-from nvflare.apis.fl_constant import ReservedKey, ReturnCode
-from nvflare.apis.fl_context import FLContext
-from nvflare.apis.shareable import Shareable
-from nvflare.app_common.abstract.aggregator import Aggregator
-from nvflare.app_common.app_constant import AppConstants
-
-
-class _AccuItem(object):
-    def __init__(self, client, data, steps):
-        self.client = client
-        self.data = data
-        self.steps = steps
-
-
-class AccumulateWeightedAggregator(Aggregator):
-    def __init__(self, exclude_vars=None, aggregation_weights=None, expected_data_kind="WEIGHT_DIFF"):
-        """Fed average aggregator.
-
-        This aggregator performs weighted arithmetic average among received shareables from clients.
-
-        Args:
-            exclude_vars (list, optional): if not specified (None), all layers are included;
-                    if list of variable/layer names, only specified variables are excluded;
-                    if string containing regular expression (e.g. "conv"), only matched variables are being excluded.
-                    Defaults to None.
-            aggregation_weights (dict, optional): a mapping from client names to weights. Defaults to None.
-            expected_data_kind (str, optional): the data_kind this aggregator can process. Defaults to "WEIGHT_DIFF".
-
-        Raises:
-            ValueError: if data_kind is neither WEIGHT_DIFF nor WEIGHTS
-        """
-        super().__init__()
-        self.exclude_vars = re.compile(exclude_vars) if exclude_vars else None
-        self.aggregation_weights = aggregation_weights or {}
-        if expected_data_kind not in [DataKind.WEIGHT_DIFF, DataKind.WEIGHTS]:
-            raise ValueError(f"expected_data_kind={expected_data_kind} not in WEIGHT_DIFF or WEIGHTS")
-        self.expected_data_kind = expected_data_kind
-        self.logger.debug(f"aggregation weights control: {aggregation_weights}")
-
-        self.accumulator = []
-
-        self.warning_count = dict()
-        self.warning_limit = 10
-
-    def accept(self, shareable: Shareable, fl_ctx: FLContext) -> bool:
-        try:
-            dxo = from_shareable(shareable)
-        except:
-            self.log_exception(fl_ctx, "shareable data is not a valid DXO")
-            return False
-
-        if dxo.data_kind not in (DataKind.WEIGHT_DIFF, DataKind.WEIGHTS):
-            self.log_error(fl_ctx, "cannot handle data kind {}".format(dxo.data_kind))
-            return False
-
-        if dxo.data_kind != self.expected_data_kind:
-            self.log_error(fl_ctx, "expect {} but got {}".format(self.expected_data_kind, dxo.data_kind))
-            return False
-
-        processed_algorithm = dxo.get_meta_prop(MetaKey.PROCESSED_ALGORITHM)
-        if processed_algorithm is not None:
-            self.log_error(fl_ctx, f"unable to accept shareable processed by {processed_algorithm}")
-            return False
-
-        current_round = fl_ctx.get_prop(AppConstants.CURRENT_ROUND)
-        self.log_debug(fl_ctx, f"current_round: {current_round}")
-        client_name = shareable.get_peer_prop(key=ReservedKey.IDENTITY_NAME, default="?")
-        contribution_round = shareable.get_header(AppConstants.CONTRIBUTION_ROUND)
-
-        rc = shareable.get_return_code()
-        if rc and rc != ReturnCode.OK:
-            self.log_info(fl_ctx, f"Client {client_name} returned rc: {rc}. Disregarding contribution.")
-            return False
-
-        data = dxo.data
-        if data is None:
-            self.log_error(fl_ctx, "no data to aggregate")
-            return False
-
-        n_iter = dxo.get_meta_prop(MetaKey.NUM_STEPS_CURRENT_ROUND)
-
-        if contribution_round == current_round:
-            if not self._client_in_accumulator(client_name):
-                self.accumulator.append(_AccuItem(client_name, data, n_iter))
-                accepted = True
-            else:
-                self.log_info(
-                    fl_ctx,
-                    f"Discarded: Current round: {current_round} contributions already include client: {client_name}",
-                )
-                accepted = False
-        else:
-            self.log_info(
-                fl_ctx,
-                "Discarded the contribution from {} for round: {}. Current round is: {}".format(
-                    client_name, contribution_round, current_round
-                ),
-            )
-            accepted = False
-        return accepted
-
-    def _client_in_accumulator(self, client_name):
-        return any(client_name == item.client for item in self.accumulator)
-
-    def aggregate(self, fl_ctx: FLContext) -> Shareable:
-        """Aggregate model variables.
-
-        This function is not thread-safe.
-
-
-        Args:
-            fl_ctx (FLContext): System-wide FL Context
-
-        Returns:
-            Shareable: Return True to indicates the current model is the best model so far.
-        """
-        current_round = fl_ctx.get_prop(AppConstants.CURRENT_ROUND)
-        self.log_info(fl_ctx, "aggregating {} updates at round {}".format(len(self.accumulator), current_round))
-
-        # TODO: What if AppConstants.GLOBAL_MODEL is None?
-        acc_vars = [set(acc.data.keys()) for acc in self.accumulator]
-        acc_vars = set.union(*acc_vars) if acc_vars else acc_vars
-        # update vars that are not in exclude pattern
-        vars_to_aggregate = (
-            [g_var for g_var in acc_vars if not self.exclude_vars.search(g_var)] if self.exclude_vars else acc_vars
+class AccumulateWeightedAggregator(InTimeAccumulateWeightedAggregator):
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self.logger.warning(
+            "'AccumulateWeightedAggregator' was deprecated, please use " "'InTimeAccumulateWeightedAggregator'"
         )
-
-        clients_with_messages = []
-        aggregated_model = {}
-        for v_name in vars_to_aggregate:
-            n_local_iters, np_vars = [], []
-            for item in self.accumulator:
-                client_name = item.client
-                data = item.data
-                n_iter = item.steps
-
-                if n_iter is None:
-                    if self.warning_count.get(client_name, 0) <= self.warning_limit:
-                        self.log_warning(
-                            fl_ctx,
-                            f"NUM_STEPS_CURRENT_ROUND missing in meta of shareable"
-                            f" from {client_name} and set to default value, 1.0. "
-                            f" This kind of message will show {self.warning_limit} times at most.",
-                        )
-                        if client_name in self.warning_count:
-                            self.warning_count[client_name] = self.warning_count[client_name] + 1
-                        else:
-                            self.warning_count[client_name] = 0
-                    n_iter = 1.0
-                if v_name not in data.keys():
-                    continue  # this acc doesn't have the variable from client
-                float_n_iter = float(n_iter)
-                n_local_iters.append(float_n_iter)
-                aggregation_weight = self.aggregation_weights.get(client_name)
-                if aggregation_weight is None:
-                    if self.warning_count.get(client_name, 0) <= self.warning_limit:
-                        self.log_warning(
-                            fl_ctx,
-                            f"Aggregation_weight missing for {client_name} and set to default value, 1.0"
-                            f" This kind of message will show {self.warning_limit} times at most.",
-                        )
-                        if client_name in self.warning_count:
-                            self.warning_count[client_name] = self.warning_count[client_name] + 1
-                        else:
-                            self.warning_count[client_name] = 0
-                    aggregation_weight = 1.0
-
-                weighted_value = data[v_name] * float_n_iter * aggregation_weight
-                if client_name not in clients_with_messages:
-                    if client_name in self.aggregation_weights.keys():
-                        self.log_debug(fl_ctx, f"Client {client_name} use weight {aggregation_weight} for aggregation.")
-                    else:
-                        self.log_debug(
-                            fl_ctx,
-                            f"Client {client_name} not defined in the aggregation weight list. Use default value 1.0",
-                        )
-                    clients_with_messages.append(client_name)
-                np_vars.append(weighted_value)
-            if not n_local_iters:
-                continue  # all acc didn't receive the variable from clients
-            new_val = np.sum(np_vars, axis=0) / np.sum(n_local_iters)
-            aggregated_model[v_name] = new_val
-
-        self.accumulator.clear()
-
-        self.log_debug(fl_ctx, f"Model after aggregation: {aggregated_model}")
-
-        dxo = DXO(data_kind=self.expected_data_kind, data=aggregated_model)
-        return dxo.to_shareable()


Test patch:
diff --git a/tests/integration_test/data/apps/tf/config/config_fed_server.json b/tests/integration_test/data/apps/tf/config/config_fed_server.json
--- a/tests/integration_test/data/apps/tf/config/config_fed_server.json
+++ b/tests/integration_test/data/apps/tf/config/config_fed_server.json
@@ -20,7 +20,7 @@
     },
     {
       "id": "aggregator",
-      "path": "nvflare.app_common.aggregators.accumulate_model_aggregator.AccumulateWeightedAggregator",
+      "path": "nvflare.app_common.aggregators.intime_accumulate_model_aggregator.InTimeAccumulateWeightedAggregator",
       "args": {
         "expected_data_kind": "WEIGHTS"
       }
diff --git a/tests/unit_test/app_common/aggregators/accumulate_weighted_aggregator_test.py b/tests/unit_test/app_common/aggregators/accumulate_weighted_aggregator_test.py
deleted file mode 100644
--- a/tests/unit_test/app_common/aggregators/accumulate_weighted_aggregator_test.py
+++ /dev/null
@@ -1,140 +0,0 @@
-# Copyright (c) 2021-2022, NVIDIA CORPORATION.  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.
-
-import random
-
-import numpy as np
-import pytest
-
-from nvflare.apis.dxo import DXO, DataKind, MetaKey, from_shareable
-from nvflare.apis.fl_constant import ReservedKey
-from nvflare.apis.fl_context import FLContext, FLContextManager
-from nvflare.apis.shareable import Shareable
-from nvflare.app_common.aggregators.accumulate_model_aggregator import AccumulateWeightedAggregator
-from nvflare.app_common.app_constant import AppConstants
-
-
-class TestAccumulateWeightedAggregator:
-    @pytest.mark.parametrize(
-        "received,expected",
-        [
-            (
-                {"client1": {"weight": 0.5, "iter_number": 1, "aggr_data": {"var1": np.array([2.0, 3.0, 1.1, 0.1])}}},
-                {"var1": np.array([1.0, 1.5, 0.55, 0.05])},
-            ),
-            (
-                {"client1": {"weight": 1.0, "iter_number": 1, "aggr_data": {"var1": np.array([2.0, 3.0, 1.1, 0.1])}}},
-                {"var1": np.array([2.0, 3.0, 1.1, 0.1])},
-            ),
-            (
-                {
-                    "client1": {"weight": 0.5, "iter_number": 3, "aggr_data": {"var1": np.array([2.0, 3.0, 1.1, 0.1])}},
-                    "client2": {"weight": 1.0, "iter_number": 1, "aggr_data": {"var1": np.array([1.0, 1.0, 2.1, 0.5])}},
-                },
-                {"var1": np.array([1.0, 1.375, 0.9375, 0.1625])},
-            ),
-            (
-                {
-                    "client1": {"weight": 1.0, "iter_number": 5, "aggr_data": {"var1": np.array([2.0, 3.0, 1.1, 0.1])}},
-                    "client2": {"weight": 0.5, "iter_number": 4, "aggr_data": {"var1": np.array([1.0, 1.0, 2.1, 0.5])}},
-                },
-                {"var1": np.array([4.0 / 3, 17.0 / 9, 9.7 / 9, 1.0 / 6])},
-            ),
-        ],
-    )
-    def test_aggregate(self, received, expected):
-        aggregation_weights = {k: v["weight"] for k, v in received.items()}
-        agg = AccumulateWeightedAggregator(aggregation_weights=aggregation_weights)
-        fl_ctx_mgr = FLContextManager(
-            engine=None, identity_name="", run_num="unit_test", public_stickers={}, private_stickers={}
-        )
-        fl_ctx = FLContext()
-        fl_ctx.set_prop(AppConstants.CURRENT_ROUND, 0)
-        for k, v in received.items():
-            fl_ctx_mgr.identity_name = k
-            dxo = DXO(
-                DataKind.WEIGHT_DIFF,
-                data=v["aggr_data"],
-                meta={
-                    MetaKey.NUM_STEPS_CURRENT_ROUND: v["iter_number"],
-                },
-            )
-
-            s = Shareable()
-            s.set_peer_props({ReservedKey.IDENTITY_NAME: k})
-            s.set_header(AppConstants.CONTRIBUTION_ROUND, 0)
-            agg.accept(dxo.update_shareable(s), fl_ctx)
-
-        result = agg.aggregate(fl_ctx)
-        np.testing.assert_allclose(result["DXO"]["data"]["var1"], expected["var1"])
-
-    @pytest.mark.parametrize("shape", [4, (6, 6)])
-    @pytest.mark.parametrize("n_clients", [10, 50, 100])
-    def test_aggregate_random(self, shape, n_clients):
-        aggregation_weights = {f"client_{i}": random.random() for i in range(n_clients)}
-        agg = AccumulateWeightedAggregator(aggregation_weights=aggregation_weights)
-        weighted_sum = np.zeros(shape)
-        sum_of_weights = 0
-        fl_ctx_mgr = FLContextManager(
-            engine=None, identity_name="", run_num="unit_test", public_stickers={}, private_stickers={}
-        )
-        fl_ctx = FLContext()
-        fl_ctx.set_prop(AppConstants.CURRENT_ROUND, 0)
-        for client_name in aggregation_weights:
-            iter_number = random.randint(1, 50)
-            fl_ctx_mgr.identity_name = client_name
-            weights = np.random.random(shape)
-            s = Shareable()
-            s.set_peer_props({ReservedKey.IDENTITY_NAME: client_name})
-            s.set_header(AppConstants.CONTRIBUTION_ROUND, 0)
-            dxo = DXO(
-                DataKind.WEIGHT_DIFF,
-                data={"var1": weights},
-                meta={
-                    MetaKey.NUM_STEPS_CURRENT_ROUND: iter_number,
-                },
-            )
-            weighted_sum = weighted_sum + weights * iter_number * aggregation_weights[client_name]
-            sum_of_weights = sum_of_weights + iter_number
-            agg.accept(dxo.update_shareable(s), fl_ctx)
-
-        result = agg.aggregate(fl_ctx)
-        result_dxo = from_shareable(result)
-        np.testing.assert_allclose(result_dxo.data["var1"], weighted_sum / sum_of_weights)
-
-    @pytest.mark.parametrize("current_round,contribution_round,expected", [(1, 1, True), (2, 1, False)])
-    def test_accept(self, current_round, contribution_round, expected):
-        aggregation_weights = {f"client_{i}": random.random() for i in range(2)}
-        agg = AccumulateWeightedAggregator(aggregation_weights=aggregation_weights)
-        client_name = "client_0"
-        iter_number = 1
-        fl_ctx_mgr = FLContextManager(
-            engine=None, identity_name="", run_num="unit_test", public_stickers={}, private_stickers={}
-        )
-        fl_ctx_mgr.identity_name = client_name
-        weights = np.random.random(4)
-
-        fl_ctx = FLContext()
-        s = Shareable()
-        s.set_peer_props({ReservedKey.IDENTITY_NAME: client_name})
-        s.set_header(AppConstants.CONTRIBUTION_ROUND, contribution_round)
-        fl_ctx.set_prop(AppConstants.CURRENT_ROUND, current_round)
-        dxo = DXO(
-            DataKind.WEIGHT_DIFF,
-            data={"var1": weights},
-            meta={
-                MetaKey.NUM_STEPS_CURRENT_ROUND: iter_number,
-            },
-        )
-        assert agg.accept(dxo.update_shareable(s), fl_ctx) == expected
diff --git a/tests/unit_test/app_common/aggregators/in_time_accumulate_weighted_aggregator_test.py b/tests/unit_test/app_common/aggregators/in_time_accumulate_weighted_aggregator_test.py
--- a/tests/unit_test/app_common/aggregators/in_time_accumulate_weighted_aggregator_test.py
+++ b/tests/unit_test/app_common/aggregators/in_time_accumulate_weighted_aggregator_test.py
@@ -20,7 +20,7 @@
 
 from nvflare.apis.dxo import DXO, DataKind, MetaKey, from_shareable
 from nvflare.apis.fl_constant import ReservedKey
-from nvflare.apis.fl_context import FLContext, FLContextManager
+from nvflare.apis.fl_context import FLContext
 from nvflare.apis.shareable import Shareable
 from nvflare.app_common.aggregators.intime_accumulate_model_aggregator import InTimeAccumulateWeightedAggregator
 from nvflare.app_common.app_constant import AppConstants
@@ -128,10 +128,6 @@ def test_accept(self, current_round, contribution_round, expected):
         agg = InTimeAccumulateWeightedAggregator(aggregation_weights=aggregation_weights)
         client_name = "client_0"
         iter_number = 1
-        fl_ctx_mgr = FLContextManager(
-            engine=None, identity_name="", run_num="unit_test", public_stickers={}, private_stickers={}
-        )
-        fl_ctx_mgr.identity_name = client_name
         weights = np.random.random(4)
 
         fl_ctx = FLContext()
@@ -196,13 +192,9 @@ def test_accept(self, current_round, contribution_round, expected):
     def test_aggregate(self, received, expected):
         aggregation_weights = {k: v["weight"] for k, v in received.items()}
         agg = InTimeAccumulateWeightedAggregator(aggregation_weights=aggregation_weights)
-        fl_ctx_mgr = FLContextManager(
-            engine=None, identity_name="", run_num="unit_test", public_stickers={}, private_stickers={}
-        )
         fl_ctx = FLContext()
         fl_ctx.set_prop(AppConstants.CURRENT_ROUND, 0)
         for k, v in received.items():
-            fl_ctx_mgr.identity_name = k
             dxo = DXO(
                 DataKind.WEIGHT_DIFF,
                 data=v["aggr_data"],
@@ -226,14 +218,10 @@ def test_aggregate_random(self, shape, n_clients):
         agg = InTimeAccumulateWeightedAggregator(aggregation_weights=aggregation_weights)
         weighted_sum = np.zeros(shape)
         sum_of_weights = 0
-        fl_ctx_mgr = FLContextManager(
-            engine=None, identity_name="", run_num="unit_test", public_stickers={}, private_stickers={}
-        )
         fl_ctx = FLContext()
         fl_ctx.set_prop(AppConstants.CURRENT_ROUND, 0)
         for client_name in aggregation_weights:
             iter_number = random.randint(1, 50)
-            fl_ctx_mgr.identity_name = client_name
             weights = np.random.random(shape)
             s = Shareable()
             s.set_peer_props({ReservedKey.IDENTITY_NAME: client_name})
@@ -268,14 +256,10 @@ def test_aggregate_random_dxos(self, num_dxo, shape, n_clients):
         )
         weighted_sum = {dxo_name: np.zeros(shape) for dxo_name in dxo_names}
         sum_of_weights = {dxo_name: 0 for dxo_name in dxo_names}
-        fl_ctx_mgr = FLContextManager(
-            engine=None, identity_name="", run_num="unit_test", public_stickers={}, private_stickers={}
-        )
         fl_ctx = FLContext()
         fl_ctx.set_prop(AppConstants.CURRENT_ROUND, 0)
         for client_name in client_names:
             iter_number = random.randint(1, 50)
-            fl_ctx_mgr.identity_name = client_name
 
             dxo_collection_data = {}
             for dxo_name in dxo_names: