Repository: NVIDIA/NVFlare

Issue:
Logic error in InTimeAccumulateWeightedAggregator for multiple DXO
If  `aggregation_weights={"dxo1": {"client_0": 1.0, "client_1": 2.0}}`

and `expected_data_kind={"dxo1": DataKind.WEIGHT_DIFF, "dxo2": DataKind.WEIGHT_DIFF}`


https://github.com/NVIDIA/NVFlare/blob/main/nvflare/app_common/aggregators/intime_accumulate_model_aggregator.py#L85

Will set `aggregation_weights_dict["dxo2"]` to `{"dxo1": {"client_0": 1.0, "client_1": 2.0}}`

Which is not correct.


Discussion:
Hi @holgerroth ,

This is a logic error.
But what do you think might be a good solution:

1. We require users to specify the "complete" aggregation weight:
    means user need to specify both "dxo1" and "dxo2" in the aggregation weight,
    otherwise we raise an error

2. We still allow the weight to be like this way.
    For the "dxo2", we just take the "dxo1"'s aggregation weight.
 
3. We still allow the weight to be like this way. 
    For the "dxo2", we just set its aggregation weight to be `{}`.
I would say, we just apply clien1's aggregation weight to all dxo's in client1's collection, and so on for other clients.
The conclusion is:

1. If the aggregation_weights is provided as `{"client1": 1, "client2": 2}`, then it will be applied to all DXOs in the dxo collection
2. If the aggregation_weights is provided in a dict of dict format like `{"dxo1": {"client1": 1, "client2": 2}}`, then we will check all the dxo in expected_data_kind need to exist in this aggregation weights, otherwise we raise an exception.

I will modify the PR as requested.

Resolving patch:
diff --git a/nvflare/app_common/aggregators/dxo_aggregator.py b/nvflare/app_common/aggregators/dxo_aggregator.py
--- a/nvflare/app_common/aggregators/dxo_aggregator.py
+++ b/nvflare/app_common/aggregators/dxo_aggregator.py
@@ -13,6 +13,7 @@
 # limitations under the License.
 
 import logging
+from typing import Any, Dict, Optional
 
 from nvflare.apis.dxo import DXO, DataKind, MetaKey
 from nvflare.apis.fl_component import FLComponent
@@ -24,18 +25,18 @@
 class DXOAggregator(FLComponent):
     def __init__(
         self,
-        exclude_vars=None,
-        aggregation_weights=None,
+        exclude_vars: Optional[str] = None,
+        aggregation_weights: Optional[Dict[str, Any]] = None,
         expected_data_kind: DataKind = DataKind.WEIGHT_DIFF,
         name_postfix: str = "",
     ):
-        """Perform accumulated weighted aggregation for one kind of corresponding DXO from contributors
+        """Perform accumulated weighted aggregation for one kind of corresponding DXO from contributors.
 
         Args:
-            exclude_vars ([type], optional): regex to match excluded vars during aggregation. Defaults to None.
-            aggregation_weights ([type], optional): dictionary to map contributor name to its aggregation weights.
+            exclude_vars (str, optional): Regex to match excluded vars during aggregation. Defaults to None.
+            aggregation_weights (Dict[str, Any], optional): Aggregation weight for each contributor.
                                 Defaults to None.
-            expected_data_kind: Expected DataKind for this kind of DXO.
+            expected_data_kind (DataKind): Expected DataKind for this DXO.
             name_postfix: optional postfix to give to class name and show in logger output.
         """
         super().__init__()
diff --git a/nvflare/app_common/aggregators/intime_accumulate_model_aggregator.py b/nvflare/app_common/aggregators/intime_accumulate_model_aggregator.py
--- a/nvflare/app_common/aggregators/intime_accumulate_model_aggregator.py
+++ b/nvflare/app_common/aggregators/intime_accumulate_model_aggregator.py
@@ -12,6 +12,8 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+from typing import Any, Dict, Union
+
 from nvflare.apis.dxo import DXO, DataKind, from_shareable
 from nvflare.apis.fl_constant import ReservedKey, ReturnCode
 from nvflare.apis.fl_context import FLContext
@@ -21,22 +23,48 @@
 from nvflare.app_common.app_constant import AppConstants
 
 
+def _is_nested_aggregation_weights(aggregation_weights):
+    if not aggregation_weights:
+        return False
+    if not isinstance(aggregation_weights, dict):
+        return False
+    first_value = next(iter(aggregation_weights.items()))[1]
+    if not isinstance(first_value, dict):
+        return False
+    return True
+
+
+def _get_missing_keys(ref_dict: dict, dict_to_check: dict):
+    result = []
+    for k in ref_dict:
+        if k not in dict_to_check:
+            result.append(k)
+    return result
+
+
 class InTimeAccumulateWeightedAggregator(Aggregator):
-    def __init__(self, exclude_vars=None, aggregation_weights=None, expected_data_kind=DataKind.WEIGHT_DIFF):
+    def __init__(
+        self,
+        exclude_vars: Union[str, Dict[str, str], None] = None,
+        aggregation_weights: Union[Dict[str, Any], Dict[str, Dict[str, Any]], None] = None,
+        expected_data_kind: Union[DataKind, Dict[str, DataKind]] = DataKind.WEIGHT_DIFF,
+    ):
         """Perform accumulated weighted aggregation
         It parses the shareable and aggregates the contained DXO(s).
 
         Args:
-            exclude_vars ([type], optional): regex to match excluded vars during aggregation. Defaults to None.
-                                Can be one string or a dict of keys with regex strings corresponding to each aggregated
-                                DXO when processing a DXO of `DataKind.COLLECTION`.
-            aggregation_weights ([type], optional): dictionary to map contributor name to its aggregation weights.
-                                Defaults to None.
-                                Can be one dict or a dict of dicts corresponding to each aggregated DXO
-                                when processing DXO of `DataKind.COLLECTION`.
-            expected_data_kind: DataKind or dict of keys and matching DataKind entries
-                                when processing DXO of `DataKind.COLLECTION`.
-                                Only the keys in the dict will be processed.
+            exclude_vars (Union[str, Dict[str, str]], optional):
+                Regular expression string to match excluded vars during aggregation. Defaults to None.
+                Can be one string or a dict of {dxo_name: regex strings} corresponding to each aggregated DXO
+                when processing a DXO of `DataKind.COLLECTION`.
+            aggregation_weights (Union[Dict[str, Any], Dict[str, Dict[str, Any]]], optional):
+                Aggregation weight for each contributor. Defaults to None.
+                Can be one dict of {contrib_name: aggr_weight} or a dict of dicts corresponding to each aggregated DXO
+                when processing a DXO of `DataKind.COLLECTION`.
+            expected_data_kind (Union[DataKind, Dict[str, DataKind]]):
+                DataKind for DXO. Defaults to DataKind.WEIGHT_DIFF
+                Can be one DataKind or a dict of {dxo_name: DataKind} corresponding to each aggregated DXO
+                when processing a DXO of `DataKind.COLLECTION`. Only the keys in this dict will be processed.
         """
         super().__init__()
         self.logger.debug(f"exclude vars: {exclude_vars}")
@@ -49,22 +77,38 @@ def __init__(self, exclude_vars=None, aggregation_weights=None, expected_data_ki
         if isinstance(expected_data_kind, dict):
             for k, v in expected_data_kind.items():
                 if v not in [DataKind.WEIGHT_DIFF, DataKind.WEIGHTS]:
-                    raise ValueError(f"expected_data_kind[{k}] = {v} not in WEIGHT_DIFF or WEIGHTS")
+                    raise ValueError(
+                        f"expected_data_kind[{k}] = {v} is not {DataKind.WEIGHT_DIFF} or {DataKind.WEIGHTS}"
+                    )
             self.expected_data_kind = expected_data_kind
         else:
             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")
+                raise ValueError(
+                    f"expected_data_kind = {expected_data_kind} is not {DataKind.WEIGHT_DIFF} or {DataKind.WEIGHTS}"
+                )
             self.expected_data_kind = {self._single_dxo_key: expected_data_kind}
 
         # Check exclude_vars
+        if exclude_vars:
+            if not isinstance(exclude_vars, dict) and not isinstance(exclude_vars, str):
+                raise ValueError(
+                    f"exclude_vars = {exclude_vars} should be a regex string but got {type(exclude_vars)}."
+                )
+            if isinstance(exclude_vars, dict):
+                missing_keys = _get_missing_keys(expected_data_kind, exclude_vars)
+                if len(missing_keys) != 0:
+                    raise ValueError(
+                        "A dict exclude_vars should specify exclude_vars for every key in expected_data_kind. "
+                        f"But missed these keys: {missing_keys}"
+                    )
+
         exclude_vars_dict = dict()
         for k in self.expected_data_kind.keys():
             if isinstance(exclude_vars, dict):
                 if k in exclude_vars:
                     if not isinstance(exclude_vars[k], str):
                         raise ValueError(
-                            f"exclude_vars[{k}] = {exclude_vars[k]} not a string but {type(exclude_vars[k])}! "
-                            f"Expected type regex string."
+                            f"exclude_vars[{k}] = {exclude_vars[k]} should be a regex string but got {type(exclude_vars[k])}."
                         )
                     exclude_vars_dict[k] = exclude_vars[k]
             else:
@@ -75,6 +119,14 @@ def __init__(self, exclude_vars=None, aggregation_weights=None, expected_data_ki
         self.exclude_vars = exclude_vars_dict
 
         # Check aggregation weights
+        if _is_nested_aggregation_weights(aggregation_weights):
+            missing_keys = _get_missing_keys(expected_data_kind, aggregation_weights)
+            if len(missing_keys) != 0:
+                raise ValueError(
+                    "A dict of dict aggregation_weights should specify aggregation_weights "
+                    f"for every key in expected_data_kind. But missed these keys: {missing_keys}"
+                )
+
         aggregation_weights = aggregation_weights or {}
         aggregation_weights_dict = dict()
         for k in self.expected_data_kind.keys():
@@ -83,8 +135,6 @@ def __init__(self, exclude_vars=None, aggregation_weights=None, expected_data_ki
             else:
                 # assume same aggregation weights for each entry of DXO collection.
                 aggregation_weights_dict[k] = aggregation_weights
-        if self._single_dxo_key in self.expected_data_kind:
-            aggregation_weights_dict[self._single_dxo_key] = aggregation_weights
         self.aggregation_weights = aggregation_weights_dict
 
         # Set up DXO aggregators
diff --git a/nvflare/app_common/aggregators/weighted_aggregation_helper.py b/nvflare/app_common/aggregators/weighted_aggregation_helper.py
--- a/nvflare/app_common/aggregators/weighted_aggregation_helper.py
+++ b/nvflare/app_common/aggregators/weighted_aggregation_helper.py
@@ -13,20 +13,15 @@
 # limitations under the License.
 
 import re
+from typing import Optional
 
 
 class WeightedAggregationHelper(object):
-    def __init__(self, exclude_vars=None):
-        """Perform weighted aggregation
-        It computes
-        weighted_sum = sum(dict*weight)
-
-        in accept function
-        The aggregate function returns
-        weighted_sum / sum_of_weights
+    def __init__(self, exclude_vars: Optional[str] = None):
+        """Perform weighted aggregation.
 
         Args:
-            exclude_vars ([type], optional): regex to match excluded vars during aggregation. Defaults to None.
+            exclude_vars (str, optional): regex string to match excluded vars during aggregation. Defaults to None.
         """
         super().__init__()
         self.exclude_vars = re.compile(exclude_vars) if exclude_vars else None
@@ -40,7 +35,8 @@ def reset_stats(self):
         self.counts = {}
         self.history = []
 
-    def add(self, data, weight, contributor_name, contribution_round) -> bool:
+    def add(self, data, weight, contributor_name, contribution_round):
+        """Compute weighted sum and sum of weights."""
         for k, v in data.items():
             if self.exclude_vars is not None and self.exclude_vars.search(k):
                 continue
@@ -61,6 +57,7 @@ def add(self, data, weight, contributor_name, contribution_round) -> bool:
         )
 
     def get_result(self):
+        """Divide weighted sum by sum of weights."""
         aggregated_dict = {k: v / self.counts[k] for k, v in self.total.items()}
         self.reset_stats()
         return aggregated_dict


Test patch:
diff --git a/test/unit_test/app_common/test_accumulate_weighted_aggregator.py b/test/unit_test/app_common/test_accumulate_weighted_aggregator.py
--- a/test/unit_test/app_common/test_accumulate_weighted_aggregator.py
+++ b/test/unit_test/app_common/test_accumulate_weighted_aggregator.py
@@ -25,39 +25,40 @@
 from nvflare.app_common.app_constant import AppConstants
 
 
-class TestAggregator:
-    @pytest.mark.parametrize("aggregator", [AccumulateWeightedAggregator])
+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_accum_aggregator(self, aggregator, received, expected):
+    def test_aggregate(self, received, expected):
         aggregation_weights = {k: v["weight"] for k, v in received.items()}
-        agg = aggregator(aggregation_weights=aggregation_weights)
+        agg = AccumulateWeightedAggregator(aggregation_weights=aggregation_weights)
         fl_ctx_mgr = FLContextManager(engine=None, identity_name="", run_num=1, 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(
@@ -68,8 +69,6 @@ def test_accum_aggregator(self, aggregator, received, expected):
                 },
             )
 
-            fl_ctx = FLContext()
-            fl_ctx.set_prop(AppConstants.CURRENT_ROUND, 0)
             s = Shareable()
             s.set_peer_props({ReservedKey.IDENTITY_NAME: k})
             s.set_header(AppConstants.CONTRIBUTION_ROUND, 0)
@@ -78,21 +77,20 @@ def test_accum_aggregator(self, aggregator, received, expected):
         result = agg.aggregate(fl_ctx)
         np.testing.assert_allclose(result["DXO"]["data"]["var1"], expected["var1"])
 
-    @pytest.mark.parametrize("aggregator", [AccumulateWeightedAggregator])
-    @pytest.mark.parametrize("shape", [(4), (6, 6)])
+    @pytest.mark.parametrize("shape", [4, (6, 6)])
     @pytest.mark.parametrize("n_clients", [10, 50, 100])
-    def test_accum_aggregator_random(self, aggregator, shape, n_clients):
+    def test_aggregate_random(self, shape, n_clients):
         aggregation_weights = {f"client_{i}": random.random() for i in range(n_clients)}
-        agg = aggregator(aggregation_weights=aggregation_weights)
+        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=1, 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)
-            fl_ctx = FLContext()
-            fl_ctx.set_prop(AppConstants.CURRENT_ROUND, 0)
             s = Shareable()
             s.set_peer_props({ReservedKey.IDENTITY_NAME: client_name})
             s.set_header(AppConstants.CONTRIBUTION_ROUND, 0)
@@ -111,34 +109,21 @@ def test_accum_aggregator_random(self, aggregator, shape, n_clients):
         result_dxo = from_shareable(result)
         np.testing.assert_allclose(result_dxo.data["var1"], weighted_sum / sum_of_weights)
 
-    @pytest.mark.parametrize("aggregator", [AccumulateWeightedAggregator])
-    def test_accum_aggregator_accept(self, aggregator):
+    @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 = aggregator(aggregation_weights=aggregation_weights)
+        agg = AccumulateWeightedAggregator(aggregation_weights=aggregation_weights)
         client_name = "client_0"
         iter_number = 1
         fl_ctx_mgr = FLContextManager(engine=None, identity_name="", run_num=1, public_stickers={}, private_stickers={})
         fl_ctx_mgr.identity_name = client_name
-        weights = np.random.random((4))
+        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, 1)
-        fl_ctx.set_prop(AppConstants.CURRENT_ROUND, 1)
-        dxo = DXO(
-            DataKind.WEIGHT_DIFF,
-            data={"var1": weights},
-            meta={
-                MetaKey.NUM_STEPS_CURRENT_ROUND: iter_number,
-            },
-        )
-        assert (True) == agg.accept(dxo.update_shareable(s), fl_ctx)
-
-        client_name = "client_0"
-        iter_number = 1
-        fl_ctx_mgr.identity_name = client_name
-        weights = np.random.random((4))
+        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},
@@ -146,9 +131,4 @@ def test_accum_aggregator_accept(self, aggregator):
                 MetaKey.NUM_STEPS_CURRENT_ROUND: iter_number,
             },
         )
-        fl_ctx = FLContext()
-        s = Shareable()
-        s.set_peer_props({ReservedKey.IDENTITY_NAME: client_name})
-        s.set_header(AppConstants.CONTRIBUTION_ROUND, 1)
-        fl_ctx.set_prop(AppConstants.CURRENT_ROUND, 2)
-        assert (False) == agg.accept(dxo.update_shareable(s), fl_ctx)
+        assert agg.accept(dxo.update_shareable(s), fl_ctx) == expected
diff --git a/test/unit_test/app_common/test_in_time_accumulate_weighted_aggregator.py b/test/unit_test/app_common/test_in_time_accumulate_weighted_aggregator.py
--- a/test/unit_test/app_common/test_in_time_accumulate_weighted_aggregator.py
+++ b/test/unit_test/app_common/test_in_time_accumulate_weighted_aggregator.py
@@ -13,6 +13,7 @@
 # limitations under the License.
 
 import random
+import re
 
 import numpy as np
 import pytest
@@ -25,20 +26,138 @@
 from nvflare.app_common.app_constant import AppConstants
 
 
-class TestAggregator:
-    @pytest.mark.parametrize("aggregator", [InTimeAccumulateWeightedAggregator])
+class TestInTimeAccumulateWeightedAggregator:
+    @pytest.mark.parametrize(
+        "exclude_vars,aggregation_weights,expected_data_kind,error,error_msg",
+        [
+            (
+                2.0,
+                None,
+                DataKind.WEIGHT_DIFF,
+                ValueError,
+                f"exclude_vars = 2.0 should be a regex string but got {type(2.0)}.",
+            ),
+            (
+                {"dxo1": 3.0, "dxo2": ""},
+                None,
+                {"dxo1": DataKind.WEIGHT_DIFF, "dxo2": DataKind.WEIGHT_DIFF},
+                ValueError,
+                f"exclude_vars[dxo1] = 3.0 should be a regex string but got {type(3.0)}.",
+            ),
+            (None, None, DataKind.ANALYTIC, ValueError, "expected_data_kind = ANALYTIC is not WEIGHT_DIFF or WEIGHTS"),
+            (
+                None,
+                None,
+                {"dxo1": DataKind.WEIGHT_DIFF, "dxo2": DataKind.ANALYTIC},
+                ValueError,
+                "expected_data_kind[dxo2] = ANALYTIC is not WEIGHT_DIFF or WEIGHTS",
+            ),
+            (
+                None,
+                {"dxo1": {"client_0": 1.0, "client_1": 2.0}},
+                {"dxo1": DataKind.WEIGHT_DIFF, "dxo2": DataKind.WEIGHT_DIFF},
+                ValueError,
+                "A dict of dict aggregation_weights should specify aggregation_weights "
+                "for every key in expected_data_kind. But missed these keys: ['dxo2']",
+            ),
+            (
+                {"dxo2": ""},
+                None,
+                {"dxo1": DataKind.WEIGHT_DIFF, "dxo2": DataKind.WEIGHT_DIFF},
+                ValueError,
+                "A dict exclude_vars should specify exclude_vars for every key in expected_data_kind. "
+                "But missed these keys: ['dxo1']",
+            ),
+        ],
+    )
+    def test_invalid_create(self, exclude_vars, aggregation_weights, expected_data_kind, error, error_msg):
+        with pytest.raises(error, match=re.escape(error_msg)):
+            _ = InTimeAccumulateWeightedAggregator(
+                exclude_vars=exclude_vars,
+                aggregation_weights=aggregation_weights,
+                expected_data_kind=expected_data_kind,
+            )
+
+    @pytest.mark.parametrize(
+        "exclude_vars,aggregation_weights,expected_data_kind,expected_object",
+        [
+            (
+                None,
+                None,
+                DataKind.WEIGHTS,
+                InTimeAccumulateWeightedAggregator(
+                    exclude_vars=None, aggregation_weights=None, expected_data_kind=DataKind.WEIGHTS
+                ),
+            ),
+            (
+                "hello",
+                None,
+                {"dxo1": DataKind.WEIGHTS, "dxo2": DataKind.WEIGHT_DIFF},
+                InTimeAccumulateWeightedAggregator(
+                    exclude_vars={"dxo1": "hello", "dxo2": "hello"},
+                    aggregation_weights=None,
+                    expected_data_kind={"dxo1": DataKind.WEIGHTS, "dxo2": DataKind.WEIGHT_DIFF},
+                ),
+            ),
+            (
+                None,
+                {"client_0": 1.0, "client_1": 2.0},
+                {"dxo1": DataKind.WEIGHTS, "dxo2": DataKind.WEIGHT_DIFF},
+                InTimeAccumulateWeightedAggregator(
+                    exclude_vars=None,
+                    aggregation_weights={
+                        "dxo1": {"client_0": 1.0, "client_1": 2.0},
+                        "dxo2": {"client_0": 1.0, "client_1": 2.0},
+                    },
+                    expected_data_kind={"dxo1": DataKind.WEIGHTS, "dxo2": DataKind.WEIGHT_DIFF},
+                ),
+            ),
+        ],
+    )
+    def test_create(self, exclude_vars, aggregation_weights, expected_data_kind, expected_object):
+        result = InTimeAccumulateWeightedAggregator(
+            exclude_vars=exclude_vars, aggregation_weights=aggregation_weights, expected_data_kind=expected_data_kind
+        )
+        assert result.exclude_vars == expected_object.exclude_vars
+        assert result.aggregation_weights == expected_object.aggregation_weights
+        assert result.expected_data_kind == expected_object.expected_data_kind
+
+    @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 = InTimeAccumulateWeightedAggregator(aggregation_weights=aggregation_weights)
+        client_name = "client_0"
+        iter_number = 1
+        fl_ctx_mgr = FLContextManager(engine=None, identity_name="", run_num=1, 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
+
     @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([2.0, 3.0, 1.1, 0.1])},
-            ],
-            [
+            ),
+            (
                 {"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": 1, "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])}},
@@ -53,8 +172,8 @@ class TestAggregator:
                         ]
                     )
                 },
-            ],
-            [
+            ),
+            (
                 {
                     "client1": {"weight": 1.0, "iter_number": 2, "aggr_data": {"var1": np.array([2.0, 3.0, 1.1, 0.1])}},
                     "client2": {"weight": 1.0, "iter_number": 4, "aggr_data": {"var1": np.array([1.0, 1.0, 2.1, 0.5])}},
@@ -69,13 +188,15 @@ class TestAggregator:
                         ]
                     )
                 },
-            ],
+            ),
         ],
     )
-    def test_accum_aggregator(self, aggregator, received, expected):
+    def test_aggregate(self, received, expected):
         aggregation_weights = {k: v["weight"] for k, v in received.items()}
-        agg = aggregator(aggregation_weights=aggregation_weights)
+        agg = InTimeAccumulateWeightedAggregator(aggregation_weights=aggregation_weights)
         fl_ctx_mgr = FLContextManager(engine=None, identity_name="", run_num=1, 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(
@@ -86,8 +207,6 @@ def test_accum_aggregator(self, aggregator, received, expected):
                 },
             )
 
-            fl_ctx = FLContext()
-            fl_ctx.set_prop(AppConstants.CURRENT_ROUND, 0)
             s = Shareable()
             s.set_peer_props({ReservedKey.IDENTITY_NAME: k})
             s.set_header(AppConstants.CONTRIBUTION_ROUND, 0)
@@ -96,21 +215,20 @@ def test_accum_aggregator(self, aggregator, received, expected):
         result = agg.aggregate(fl_ctx)
         np.testing.assert_allclose(result["DXO"]["data"]["var1"], expected["var1"])
 
-    @pytest.mark.parametrize("aggregator", [InTimeAccumulateWeightedAggregator])
-    @pytest.mark.parametrize("shape", [(4), (6, 6)])
+    @pytest.mark.parametrize("shape", [4, (6, 6)])
     @pytest.mark.parametrize("n_clients", [10, 50, 100])
-    def test_accum_aggregator_random(self, aggregator, shape, n_clients):
+    def test_aggregate_random(self, shape, n_clients):
         aggregation_weights = {f"client_{i}": random.random() for i in range(n_clients)}
-        agg = aggregator(aggregation_weights=aggregation_weights)
+        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=1, 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)
-            fl_ctx = FLContext()
-            fl_ctx.set_prop(AppConstants.CURRENT_ROUND, 0)
             s = Shareable()
             s.set_peer_props({ReservedKey.IDENTITY_NAME: client_name})
             s.set_header(AppConstants.CONTRIBUTION_ROUND, 0)
@@ -129,44 +247,55 @@ def test_accum_aggregator_random(self, aggregator, shape, n_clients):
         result_dxo = from_shareable(result)
         np.testing.assert_allclose(result_dxo.data["var1"], weighted_sum / sum_of_weights)
 
-    @pytest.mark.parametrize("aggregator", [InTimeAccumulateWeightedAggregator])
-    def test_accum_aggregator_accept(self, aggregator):
-        aggregation_weights = {f"client_{i}": random.random() for i in range(2)}
-        agg = aggregator(aggregation_weights=aggregation_weights)
-        client_name = "client_0"
-        iter_number = 1
+    @pytest.mark.parametrize("num_dxo", [1, 2, 3])
+    @pytest.mark.parametrize("shape", [4, (6, 6)])
+    @pytest.mark.parametrize("n_clients", [10, 50, 100])
+    def test_aggregate_random_dxos(self, num_dxo, shape, n_clients):
+        dxo_names = [f"dxo_{i}" for i in range(num_dxo)]
+        client_names = [f"client_{i}" for i in range(n_clients)]
+        aggregation_weights = {
+            dxo_name: {client_name: random.random() for client_name in client_names} for dxo_name in dxo_names
+        }
+        agg = InTimeAccumulateWeightedAggregator(
+            aggregation_weights=aggregation_weights,
+            expected_data_kind={dxo_name: DataKind.WEIGHT_DIFF for dxo_name in dxo_names},
+        )
+        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=1, 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, 1)
-        fl_ctx.set_prop(AppConstants.CURRENT_ROUND, 1)
-        dxo = DXO(
-            DataKind.WEIGHT_DIFF,
-            data={"var1": weights},
-            meta={
-                MetaKey.NUM_STEPS_CURRENT_ROUND: iter_number,
-            },
-        )
-        assert (True) == agg.accept(dxo.update_shareable(s), fl_ctx)
+        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
 
-        client_name = "client_0"
-        iter_number = 1
-        fl_ctx_mgr.identity_name = client_name
-        weights = np.random.random((4))
-        dxo = DXO(
-            DataKind.WEIGHT_DIFF,
-            data={"var1": weights},
-            meta={
-                MetaKey.NUM_STEPS_CURRENT_ROUND: iter_number,
-            },
-        )
-        fl_ctx = FLContext()
-        s = Shareable()
-        s.set_peer_props({ReservedKey.IDENTITY_NAME: client_name})
-        s.set_header(AppConstants.CONTRIBUTION_ROUND, 1)
-        fl_ctx.set_prop(AppConstants.CURRENT_ROUND, 2)
-        assert (False) == agg.accept(dxo.update_shareable(s), fl_ctx)
+            dxo_collection_data = {}
+            for dxo_name in dxo_names:
+                values = np.random.random(shape)
+                dxo = DXO(
+                    data_kind=DataKind.WEIGHT_DIFF,
+                    data={"var1": values},
+                    meta={
+                        MetaKey.NUM_STEPS_CURRENT_ROUND: iter_number,
+                    },
+                )
+                dxo_collection_data[dxo_name] = dxo
+                weighted_sum[dxo_name] = (
+                    weighted_sum[dxo_name] + values * iter_number * aggregation_weights[dxo_name][client_name]
+                )
+                sum_of_weights[dxo_name] = (
+                    sum_of_weights[dxo_name] + iter_number * aggregation_weights[dxo_name][client_name]
+                )
+
+            dxo_collection = DXO(data_kind=DataKind.COLLECTION, data=dxo_collection_data)
+            s = Shareable()
+            s.set_peer_props({ReservedKey.IDENTITY_NAME: client_name})
+            s.set_header(AppConstants.CONTRIBUTION_ROUND, 0)
+            agg.accept(dxo_collection.update_shareable(s), fl_ctx)
+
+        result = agg.aggregate(fl_ctx)
+        result_dxo = from_shareable(result)
+        for dxo_name in dxo_names:
+            np.testing.assert_allclose(
+                result_dxo.data[dxo_name].data["var1"], weighted_sum[dxo_name] / sum_of_weights[dxo_name]
+            )