Repository: NVIDIA/NVFlare

Issue:
Server process getting killed after 3 rounds
I succeeded deploy my FL system on NVFlare but I got notice that the server could never get over 4th round. I have tried increased the heart beat time out but still the whole system could not get up to 4th round.

```
2022-01-27 21:47:23,106 - ScatterAndGather - INFO - [run=1, wf=scatter_gather_ctl]: Round 3 started.
2022-01-27 21:47:23,106 - ScatterAndGather - INFO - [run=1, wf=scatter_gather_ctl]: scheduled task train
2022-01-27 21:47:23,748 - FederatedServer - INFO - Fetch task requested from client: client3 (1e49f27c-5cd1-470a-824e-e7cff18ba7a5)
2022-01-27 21:47:23,748 - ServerRunner - INFO - [run=1, wf=scatter_gather_ctl, peer=client3, peer_run=1]: got task request from client
2022-01-27 21:47:23,749 - ServerRunner - INFO - [run=1, wf=scatter_gather_ctl, peer=client3, peer_run=1]: assigned task to client: name=train, id=0d13e935-9560-4e02-9f49-0c315b93b0bc
2022-01-27 21:47:23,749 - FederatedServer - INFO - Fetch task requested from client: client4 (d0c87b41-716f-4b61-8eec-8e926a942edd)
2022-01-27 21:47:23,749 - ServerRunner - INFO - [run=1, wf=scatter_gather_ctl, peer=client3, peer_run=1, task_name=train, task_id=0d13e935-9560-4e02-9f49-0c315b93b0bc]: sent task assignment to client
2022-01-27 21:47:23,913 - FederatedServer - INFO - Return task:train to client:client3 --- (1e49f27c-5cd1-470a-824e-e7cff18ba7a5) 
2022-01-27 21:47:23,967 - FederatedServer - INFO - Fetch task requested from client: client2 (9a54296a-934c-4afc-80d7-e9f05d57cda1)
2022-01-27 21:47:23,967 - ServerRunner - INFO - [run=1, wf=scatter_gather_ctl, peer=client4, peer_run=1]: got task request from client
2022-01-27 21:47:24,025 - ServerRunner - INFO - [run=1, wf=scatter_gather_ctl, peer=client4, peer_run=1]: assigned task to client: name=train, id=9c091627-5edb-4c06-bc6e-dc1e434067f8
2022-01-27 21:47:24,025 - ServerRunner - INFO - [run=1, wf=scatter_gather_ctl, peer=client4, peer_run=1, task_name=train, task_id=9c091627-5edb-4c06-bc6e-dc1e434067f8]: sent task assignment to client
2022-01-27 21:47:24,236 - FederatedServer - INFO - Return task:train to client:client4 --- (d0c87b41-716f-4b61-8eec-8e926a942edd) 
2022-01-27 21:47:24,288 - ServerRunner - INFO - [run=1, wf=scatter_gather_ctl, peer=client2, peer_run=1]: got task request from client
2022-01-27 21:47:24,352 - ServerRunner - INFO - [run=1, wf=scatter_gather_ctl, peer=client2, peer_run=1]: assigned task to client: name=train, id=55a4fa40-def5-4c1a-8b0f-fb693f0c524f
2022-01-27 21:47:24,353 - ServerRunner - INFO - [run=1, wf=scatter_gather_ctl, peer=client2, peer_run=1, task_name=train, task_id=55a4fa40-def5-4c1a-8b0f-fb693f0c524f]: sent task assignment to client
2022-01-27 21:47:24,574 - FederatedServer - INFO - Return task:train to client:client2 --- (9a54296a-934c-4afc-80d7-e9f05d57cda1) 
2022-01-27 21:47:26,769 - FederatedServer - INFO - Fetch task requested from client: client1 (01d07381-5eac-4415-9cc1-26d5d32c5fa8)
2022-01-27 21:47:26,769 - ServerRunner - INFO - [run=1, wf=scatter_gather_ctl, peer=client1, peer_run=1]: got task request from client
2022-01-27 21:47:26,769 - ServerRunner - INFO - [run=1, wf=scatter_gather_ctl, peer=client1, peer_run=1]: assigned task to client: name=train, id=0814d10c-d567-422b-a898-362a414c2de8
2022-01-27 21:47:26,770 - ServerRunner - INFO - [run=1, wf=scatter_gather_ctl, peer=client1, peer_run=1, task_name=train, task_id=0814d10c-d567-422b-a898-362a414c2de8]: sent task assignment to client
2022-01-27 21:47:26,944 - FederatedServer - INFO - Return task:train to client:client1 --- (01d07381-5eac-4415-9cc1-26d5d32c5fa8) 
start fl because process of 3101356 does not exist
new pid 3103341
2022-01-27 21:47:36,987 - FederatedServer - INFO - starting secure server at imtl-fl:8765
deployed FL server trainer.
2022-01-27 21:47:36,994 - FedAdminServer - INFO - Starting Admin Server imtl-fl on Port 8766
2022-01-27 21:47:36,994 - root - INFO - Server started
2022-01-27 21:47:40,645 - ClientManager - INFO - Re-activate dead client:client2 with token: 9a54296a-934c-4afc-80d7-e9f05d57cda1
2022-01-27 21:47:40,647 - FederatedServer - INFO - Fetch task requested from client: client2 (9a54296a-934c-4afc-80d7-e9f05d57cda1)
2022-01-27 21:47:40,648 - FederatedServer - INFO - server has no current run - asked client to end the run
2022-01-27 21:47:40,648 - FederatedServer - INFO - Return task:__end_run__ to client:client2 --- (9a54296a-934c-4afc-80d7-e9f05d57cda1) 
2022-01-27 21:48:06,726 - ClientManager - INFO - Re-activate dead client:client4 with token: d0c87b41-716f-4b61-8eec-8e926a942edd
2022-01-27 21:48:10,679 - FederatedServer - INFO - Fetch task requested from client: client4 (d0c87b41-716f-4b61-8eec-8e926a942edd)
2022-01-27 21:48:10,679 - FederatedServer - INFO - server has no current run - asked client to end the run
2022-01-27 21:48:10,679 - FederatedServer - INFO - Return task:__end_run__ to client:client4 --- (d0c87b41-716f-4b61-8eec-8e926a942edd) 
2022-01-27 21:48:14,767 - ClientManager - INFO - Re-activate dead client:client3 with token: 1e49f27c-5cd1-470a-824e-e7cff18ba7a5
2022-01-27 21:48:15,685 - FederatedServer - INFO - Fetch task requested from client: client3 (1e49f27c-5cd1-470a-824e-e7cff18ba7a5)
2022-01-27 21:48:15,686 - FederatedServer - INFO - server has no current run - asked client to end the run
2022-01-27 21:48:15,686 - FederatedServer - INFO - Return task:__end_run__ to client:client3 --- (1e49f27c-5cd1-470a-824e-e7cff18ba7a5) 
2022-01-27 21:48:23,364 - ClientManager - INFO - Re-activate dead client:client1 with token: 01d07381-5eac-4415-9cc1-26d5d32c5fa8
2022-01-27 21:48:25,714 - FederatedServer - INFO - Fetch task requested from client: client1 (01d07381-5eac-4415-9cc1-26d5d32c5fa8)
2022-01-27 21:48:25,714 - FederatedServer - INFO - server has no current run - asked client to end the run
2022-01-27 21:48:25,714 - FederatedServer - INFO - Return task:__end_run__ to client:client1 --- (01d07381-5eac-4415-9cc1-26d5d32c5fa8) 
```

Discussion:

Patch:
diff --git a/docs/faq.rst b/docs/faq.rst
index 14e075b9..382f1a75 100644
--- a/docs/faq.rst
+++ b/docs/faq.rst
@@ -296,6 +296,11 @@ Server related questions
     When restarting the FL server, you can find the previous training round number from the previous log. Then you can choose to
     train from scratch or continuously using previous training model.
 
+#. Why does my FL server keep crashing after a certain round?
+
+    Check that the amount of memory being consumed is not increasing in a way that it exceeds the available resources.
+    If the process consumes too much memory, the operating system may kill it.
+
 #. Does the federated learning server need a GPU?
 
     No, there is no need to have GPU on the server side for the FL server to deploy. However, certain handlers may require
diff --git a/docs/quickstart.rst b/docs/quickstart.rst
index 7375c96c..1778e167 100644
--- a/docs/quickstart.rst
+++ b/docs/quickstart.rst
@@ -70,6 +70,22 @@ In one last terminal, start the admin:
 This will launch a command prompt where you can input admin commands to control and monitor many aspects of
 the FL process. Log in by entering ``admin`` for both the username and password.
 
+Custom Code in Example Apps
+===========================
+
+There are several ways to make :ref:`custom code <custom_code>` available to clients when using NVIDIA FLARE. Most
+hello-* examples use a custom folder within the FL application. Note that using a custom folder in the app needs to be
+:ref:`allowed <troubleshooting_byoc>` when using secure provisioning. By default, this option is disabled in the secure
+mode. POC mode, however, will work with custom code by default.
+
+In contrast, the `CIFAR-10 <https://github.com/NVIDIA/NVFlare/tree/main/examples/cifar10>`_,
+`prostate segmentation <https://github.com/NVIDIA/NVFlare/tree/main/examples/prostate>`_,
+and `BraTS18 segmentation <https://github.com/NVIDIA/NVFlare/tree/main/examples/brats18>`_ examples assume that the
+learner code is already installed on the client's system and
+available in the PYTHONPATH. Hence, the app folders do not include the custom code there. The PYTHONPATH is
+set in the ``run_poc.sh`` or ``run_secure.sh`` scripts of the example. Running these scripts as described in the README
+will make the learner code available to the clients.
+
 .. _example_apps:
 
 Example Apps for NVIDIA FLARE
@@ -88,3 +104,5 @@ The following quickstart guides walk you through some of these examples:
    examples/hello_tf2
    examples/hello_cross_val
    Federated Learning with CIFAR-10 <https://github.com/NVIDIA/NVFlare/tree/main/examples/cifar10>
+
+For the complete collection of example applications, see https://github.com/NVIDIA/NVFlare/tree/main/examples.
diff --git a/docs/user_guide/application.rst b/docs/user_guide/application.rst
index 874f7624..a1ca909e 100644
--- a/docs/user_guide/application.rst
+++ b/docs/user_guide/application.rst
@@ -64,6 +64,9 @@ Example:
     task_result_filters, "What filters to apply to data leaving client, :ref:`filters`"
     components, All of the Components to use
 
+
+.. _custom_code:
+
 ***********
 Custom code
 ***********
@@ -93,6 +96,8 @@ the client config should have the following in order to configure it as an Execu
 
 Please follow :ref:`quickstart:Quickstart` to learn more.
 
+.. _troubleshooting_byoc:
+
 Troubleshooting BYOC
 ====================
 There is an ``enable_byoc`` flag for each participant that can be set at provisioning, and if that is disabled, even if
diff --git a/nvflare/apis/dxo.py b/nvflare/apis/dxo.py
index 8e006143..5b606f31 100644
--- a/nvflare/apis/dxo.py
+++ b/nvflare/apis/dxo.py
@@ -25,6 +25,7 @@ class DataKind(object):
     METRICS = "METRICS"
     MODEL = "MODEL"
     ANALYTIC = "ANALYTIC"
+    COLLECTION = "COLLECTION"  # Dict of DXO objects
 
 
 class MetaKey(object):
diff --git a/nvflare/app_common/aggregators/dxo_aggregator.py b/nvflare/app_common/aggregators/dxo_aggregator.py
new file mode 100644
index 00000000..44dc3ce7
--- /dev/null
+++ b/nvflare/app_common/aggregators/dxo_aggregator.py
@@ -0,0 +1,163 @@
+# 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 logging
+
+from nvflare.apis.dxo import DXO, DataKind, MetaKey
+from nvflare.apis.fl_component import FLComponent
+from nvflare.apis.fl_context import FLContext
+from nvflare.app_common.aggregators.weighted_aggregation_helper import WeightedAggregationHelper
+from nvflare.app_common.app_constant import AppConstants
+
+
+class DXOAggregator(FLComponent):
+    def __init__(
+        self,
+        exclude_vars=None,
+        aggregation_weights=None,
+        expected_data_kind: DataKind = DataKind.WEIGHT_DIFF,
+        name_postfix: str = "",
+    ):
+        """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.
+                                Defaults to None.
+            expected_data_kind: Expected DataKind for this kind of DXO.
+            name_postfix: optional postfix to give to class name and show in logger output.
+        """
+        super().__init__()
+        self.expected_data_kind = expected_data_kind
+        self.aggregation_weights = aggregation_weights or {}
+        self.logger.debug(f"aggregation weights control: {aggregation_weights}")
+
+        self.aggregation_helper = WeightedAggregationHelper(exclude_vars=exclude_vars)
+
+        self.warning_count = {}
+        self.warning_limit = 10
+
+        if name_postfix:
+            self._name += name_postfix
+            self.logger = logging.getLogger(self._name)
+
+    def reset_aggregation_helper(self):
+        if self.aggregation_helper:
+            self.aggregation_helper.reset_stats()
+
+    def accept(self, dxo: DXO, contributor_name, contribution_round, fl_ctx: FLContext) -> bool:
+        """Store DXO and update aggregator's internal state
+        Args:
+            dxo: information from contributor
+            contributor_name: name of the contributor
+            contribution_round: round of the contribution
+            fl_ctx: context provided by workflow
+        Returns:
+            The boolean to indicate if DXO is accepted.
+        """
+
+        if not isinstance(dxo, DXO):
+            self.log_error(fl_ctx, f"Expected DXO but got {type(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, "expected {} 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 DXO 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}")
+
+        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:
+            self.log_warning(
+                fl_ctx,
+                f"discarding DXO from {contributor_name} at round: "
+                f"{contribution_round}. Current round is: {current_round}",
+            )
+            return False
+
+        for item in self.aggregation_helper.get_history():
+            if contributor_name == item["contributor_name"]:
+                prev_round = item["round"]
+                self.log_warning(
+                    fl_ctx,
+                    f"discarding DXO from {contributor_name} at round: "
+                    f"{contribution_round} as {prev_round} accepted already",
+                )
+                return False
+
+        if n_iter is None:
+            if self.warning_count.get(contributor_name, 0) <= self.warning_limit:
+                self.log_warning(
+                    fl_ctx,
+                    f"NUM_STEPS_CURRENT_ROUND missing in meta of DXO"
+                    f" from {contributor_name} and set to default value, 1.0. "
+                    f" This kind of message will show {self.warning_limit} times at most.",
+                )
+                if contributor_name in self.warning_count:
+                    self.warning_count[contributor_name] = self.warning_count[contributor_name] + 1
+                else:
+                    self.warning_count[contributor_name] = 0
+            n_iter = 1.0
+        float_n_iter = float(n_iter)
+        aggregation_weight = self.aggregation_weights.get(contributor_name)
+        if aggregation_weight is None:
+            if self.warning_count.get(contributor_name, 0) <= self.warning_limit:
+                self.log_warning(
+                    fl_ctx,
+                    f"Aggregation_weight missing for {contributor_name} and set to default value, 1.0"
+                    f" This kind of message will show {self.warning_limit} times at most.",
+                )
+                if contributor_name in self.warning_count:
+                    self.warning_count[contributor_name] = self.warning_count[contributor_name] + 1
+                else:
+                    self.warning_count[contributor_name] = 0
+            aggregation_weight = 1.0
+
+        # aggregate
+        self.aggregation_helper.add(data, aggregation_weight * float_n_iter, contributor_name, contribution_round)
+        self.log_debug(fl_ctx, "End accept")
+        return True
+
+    def aggregate(self, fl_ctx: FLContext) -> DXO:
+        """Called when workflow determines to generate DXO to send back to contributors
+        Args:
+            fl_ctx (FLContext): context provided by workflow
+        Returns:
+            DXO: the weighted mean of accepted DXOs from contributors
+        """
+
+        self.log_debug(fl_ctx, "Start aggregation")
+        current_round = fl_ctx.get_prop(AppConstants.CURRENT_ROUND)
+        self.log_info(fl_ctx, f"aggregating {self.aggregation_helper.get_len()} update(s) at round {current_round}")
+        self.log_debug(fl_ctx, f"complete history {self.aggregation_helper.get_len()}")
+        aggregated_dict = self.aggregation_helper.get_result()
+        self.log_debug(fl_ctx, "End aggregation")
+
+        dxo = DXO(data_kind=self.expected_data_kind, data=aggregated_dict)
+        return dxo
diff --git a/nvflare/app_common/aggregators/intime_accumulate_model_aggregator.py b/nvflare/app_common/aggregators/intime_accumulate_model_aggregator.py
index 511f5ba0..ba3e0631 100644
--- a/nvflare/app_common/aggregators/intime_accumulate_model_aggregator.py
+++ b/nvflare/app_common/aggregators/intime_accumulate_model_aggregator.py
@@ -12,181 +12,173 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-import re
-
-from nvflare.apis.dxo import DXO, DataKind, MetaKey, from_shareable
+from nvflare.apis.dxo import DXO, DataKind, 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.aggregators.dxo_aggregator import DXOAggregator
 from nvflare.app_common.app_constant import AppConstants
 
 
 class InTimeAccumulateWeightedAggregator(Aggregator):
     def __init__(self, exclude_vars=None, aggregation_weights=None, expected_data_kind=DataKind.WEIGHT_DIFF):
-        """Perform accumulated weighted aggregation.
-
-        It computes
-        weighted_sum = sum(shareable*n_iteration*aggregation_weights) and
-        sum_of_weights = sum(n_iteration)
-
-        in accept function
-        The aggregate function returns
-        weighted_sum / sum_of_weights
+        """Perform accumulated weighted aggregation
+        It parses the shareable and aggregates the contained DXO(s).
 
         Args:
-            exclude_vars (str, optional): regex to match excluded vars during aggregation. Defaults to None.
-            aggregation_weights (dict, optional): dictionary to map client name to its aggregation weights. Defaults to None.
-            expected_data_kind (DataKind, optional): the data_kind this aggregator can process.  Defaults to WEIGHT_DIFF
+            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.
         """
         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"exclude vars: {exclude_vars}")
         self.logger.debug(f"aggregation weights control: {aggregation_weights}")
-        self.reset_stats()
-        self.warning_count = {}
-        self.warning_limit = 10
-        self.total = dict()
-        self.counts = dict()
-        self.history = list()
-
-    def reset_stats(self):
-        self.total = {}
-        self.counts = {}
-        self.history = []
+        self.logger.debug(f"expected data kind: {expected_data_kind}")
+
+        self._single_dxo_key = ""
+
+        # Check expected data kind
+        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")
+            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")
+            self.expected_data_kind = {self._single_dxo_key: expected_data_kind}
+
+        # Check exclude_vars
+        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."
+                        )
+                    exclude_vars_dict[k] = exclude_vars[k]
+            else:
+                # assume same exclude vars for each entry of DXO collection.
+                exclude_vars_dict[k] = exclude_vars
+        if self._single_dxo_key in self.expected_data_kind:
+            exclude_vars_dict[self._single_dxo_key] = exclude_vars
+        self.exclude_vars = exclude_vars_dict
+
+        # Check aggregation weights
+        aggregation_weights = aggregation_weights or {}
+        aggregation_weights_dict = dict()
+        for k in self.expected_data_kind.keys():
+            if k in aggregation_weights:
+                aggregation_weights_dict[k] = aggregation_weights[k]
+            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
+        self.dxo_aggregators = dict()
+        for k in self.expected_data_kind.keys():
+            self.dxo_aggregators.update(
+                {
+                    k: DXOAggregator(
+                        exclude_vars=self.exclude_vars[k],
+                        aggregation_weights=self.aggregation_weights[k],
+                        expected_data_kind=self.expected_data_kind[k],
+                        name_postfix=k,
+                    )
+                }
+            )
 
     def accept(self, shareable: Shareable, fl_ctx: FLContext) -> bool:
-        """Store shareable and update aggregator's internal state.
+        """Store shareable and update aggregator's internal state
 
         Args:
-            shareable: information from client
+            shareable: information from contributor
             fl_ctx: context provided by workflow
 
         Returns:
             The first boolean indicates if this shareable is accepted.
-            The second bollean indicates if aggregate can be called.
+            The second boolean indicates if aggregate can be called.
         """
         try:
             dxo = from_shareable(shareable)
-        except:
+        except BaseException:
             self.log_exception(fl_ctx, "shareable data is not a valid DXO")
             return False
         assert isinstance(dxo, DXO)
 
-        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}")
+        if dxo.data_kind not in (DataKind.WEIGHT_DIFF, DataKind.WEIGHTS, DataKind.COLLECTION):
+            self.log_error(
+                fl_ctx,
+                f"cannot handle data kind {dxo.data_kind}, "
+                f"expecting DataKind.WEIGHT_DIFF, DataKind.WEIGHTS, or DataKind.COLLECTION.",
+            )
             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="?")
+        contributor_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.")
+            self.log_warning(fl_ctx, f"Contributor {contributor_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
+        # Accept expected DXO(s) in shareable
+        n_accepted = 0
+        for key in self.expected_data_kind.keys():
+            if key == self._single_dxo_key:  # expecting a single DXO
+                sub_dxo = dxo
+            else:  # expecting a collection of DXOs
+                sub_dxo = dxo.data.get(key)
+            if not isinstance(sub_dxo, DXO):
+                self.log_warning(fl_ctx, f"Collection does not contain DXO for key {key} but {type(sub_dxo)}.")
+                continue
 
-        n_iter = dxo.get_meta_prop(MetaKey.NUM_STEPS_CURRENT_ROUND)
-        if contribution_round != current_round:
-            self.log_info(
-                fl_ctx,
-                f"discarding shareable from {client_name} at round: {contribution_round}. Current round is: {current_round}",
+            accepted = self.dxo_aggregators[key].accept(
+                dxo=sub_dxo, contributor_name=contributor_name, contribution_round=contribution_round, fl_ctx=fl_ctx
             )
-            return False
-
-        for item in self.history:
-            if client_name == item["client_name"]:
-                prev_round = item["round"]
-                self.log_info(
-                    fl_ctx,
-                    f"discarding shareable from {client_name} at round: {contribution_round} as {prev_round} accepted already",
-                )
+            if not accepted:
                 return False
-
-        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
-        float_n_iter = 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
-
-        for k, v in data.items():
-            if self.exclude_vars is not None and self.exclude_vars.search(k):
-                continue
-            weighted_value = v * aggregation_weight * float_n_iter
-            current_total = self.total.get(k, None)
-            if current_total is None:
-                self.total[k] = weighted_value
-                self.counts[k] = n_iter
             else:
-                self.total[k] = current_total + weighted_value
-                self.counts[k] = self.counts[k] + n_iter
-        self.history.append(
-            {
-                "client_name": client_name,
-                "round": contribution_round,
-                "aggregation_weight": aggregation_weight,
-                "n_iter": n_iter,
-            }
-        )
-        self.log_debug(fl_ctx, "End accept")
-        return True
+                n_accepted += 1
+
+        if n_accepted > 0:
+            return True
+        else:
+            self.log_warning(fl_ctx, f"Did not accept any DXOs from {contributor_name} in round {contribution_round}!")
+            return False
 
     def aggregate(self, fl_ctx: FLContext) -> Shareable:
-        """Called when workflow determines to generate shareable to send back to clients.
+        """Called when workflow determines to generate shareable to send back to contributors
 
         Args:
             fl_ctx (FLContext): context provided by workflow
 
         Returns:
-            Shareable: the weighted mean of accepted shareables from clients
+            Shareable: the weighted mean of accepted shareables from contributors
         """
+
         self.log_debug(fl_ctx, "Start aggregation")
-        current_round = fl_ctx.get_prop(AppConstants.CURRENT_ROUND)
-        self.log_info(fl_ctx, f"aggregating {len(self.history)} update(s) at round {current_round}")
-        self.log_debug(fl_ctx, f"complete history {self.history}")
-        aggregated_dict = {k: v / self.counts[k] for k, v in self.total.items()}
-        self.reset_stats()
-        self.log_debug(fl_ctx, "End aggregation")
-
-        dxo = DXO(data_kind=self.expected_data_kind, data=aggregated_dict)
-        return dxo.to_shareable()
+        result_dxo_dict = dict()
+        # Aggregate the expected DXO(s)
+        for key in self.expected_data_kind.keys():
+            aggregated_dxo = self.dxo_aggregators[key].aggregate(fl_ctx)
+            if key == self._single_dxo_key:  # return single DXO with aggregation results
+                return aggregated_dxo.to_shareable()
+            self.log_info(fl_ctx, f"Aggregated contributions matching key '{key}'.")
+            result_dxo_dict.update({key: aggregated_dxo})
+        # return collection of DXOs with aggregation results
+        collection_dxo = DXO(data_kind=DataKind.COLLECTION, data=result_dxo_dict)
+        return collection_dxo.to_shareable()
diff --git a/nvflare/app_common/aggregators/weighted_aggregation_helper.py b/nvflare/app_common/aggregators/weighted_aggregation_helper.py
new file mode 100644
index 00000000..d2cc51bf
--- /dev/null
+++ b/nvflare/app_common/aggregators/weighted_aggregation_helper.py
@@ -0,0 +1,72 @@
+# 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 re
+
+
+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
+
+        Args:
+            exclude_vars ([type], optional): regex to match excluded vars during aggregation. Defaults to None.
+        """
+        super().__init__()
+        self.exclude_vars = re.compile(exclude_vars) if exclude_vars else None
+        self.reset_stats()
+        self.total = dict()
+        self.counts = dict()
+        self.history = list()
+
+    def reset_stats(self):
+        self.total = {}
+        self.counts = {}
+        self.history = []
+
+    def add(self, data, weight, contributor_name, contribution_round) -> bool:
+        for k, v in data.items():
+            if self.exclude_vars is not None and self.exclude_vars.search(k):
+                continue
+            weighted_value = v * weight
+            current_total = self.total.get(k, None)
+            if current_total is None:
+                self.total[k] = weighted_value
+                self.counts[k] = weight
+            else:
+                self.total[k] = current_total + weighted_value
+                self.counts[k] = self.counts[k] + weight
+        self.history.append(
+            {
+                "contributor_name": contributor_name,
+                "round": contribution_round,
+                "weight": weight,
+            }
+        )
+
+    def get_result(self):
+        aggregated_dict = {k: v / self.counts[k] for k, v in self.total.items()}
+        self.reset_stats()
+        return aggregated_dict
+
+    def get_history(self):
+        return self.history
+
+    def get_len(self):
+        return len(self.get_history())


Test patch:
diff --git a/test/test_aggregator.py b/test/test_aggregator.py
index 633512be..8571fc10 100644
--- a/test/test_aggregator.py
+++ b/test/test_aggregator.py
@@ -20,16 +20,13 @@ 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 ReservedHeaderKey, Shareable
+from nvflare.apis.shareable import Shareable
 from nvflare.app_common.aggregators.accumulate_model_aggregator import AccumulateWeightedAggregator
-from nvflare.app_common.aggregators.intime_accumulate_model_aggregator import InTimeAccumulateWeightedAggregator
-
-# from nvflare.app_common.app_constant import AppConstants, AppShareableKey, AppShareableValue, ShareableContentType
 from nvflare.app_common.app_constant import AppConstants
 
 
 class TestAggregator:
-    @pytest.mark.parametrize("aggregator", [AccumulateWeightedAggregator, InTimeAccumulateWeightedAggregator])
+    @pytest.mark.parametrize("aggregator", [AccumulateWeightedAggregator])
     @pytest.mark.parametrize(
         "received,expected",
         [
@@ -81,7 +78,7 @@ class TestAggregator:
         result = agg.aggregate(fl_ctx)
         np.testing.assert_allclose(result["DXO"]["data"]["var1"], expected["var1"])
 
-    @pytest.mark.parametrize("aggregator", [AccumulateWeightedAggregator, InTimeAccumulateWeightedAggregator])
+    @pytest.mark.parametrize("aggregator", [AccumulateWeightedAggregator])
     @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):
@@ -114,7 +111,7 @@ class TestAggregator:
         result_dxo = from_shareable(result)
         np.testing.assert_allclose(result_dxo.data["var1"], weighted_sum / sum_of_weights)
 
-    @pytest.mark.parametrize("aggregator", [AccumulateWeightedAggregator, InTimeAccumulateWeightedAggregator])
+    @pytest.mark.parametrize("aggregator", [AccumulateWeightedAggregator])
     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)
diff --git a/test/test_intime_aggregator.py b/test/test_intime_aggregator.py
new file mode 100644
index 00000000..6ab83515
--- /dev/null
+++ b/test/test_intime_aggregator.py
@@ -0,0 +1,172 @@
+# 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.intime_accumulate_model_aggregator import InTimeAccumulateWeightedAggregator
+from nvflare.app_common.app_constant import AppConstants
+
+
+class TestAggregator:
+    @pytest.mark.parametrize("aggregator", [InTimeAccumulateWeightedAggregator])
+    @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])}},
+                },
+                {
+                    "var1": np.array(
+                        [
+                            (0.5 * 2.0 + 1.0 * 1.0) / (0.5 + 1),
+                            (0.5 * 3.0 + 1.0 * 1.0) / (0.5 + 1),
+                            (0.5 * 1.1 + 1.0 * 2.1) / (0.5 + 1),
+                            (0.5 * 0.1 + 1.0 * 0.5) / (0.5 + 1),
+                        ]
+                    )
+                },
+            ],
+            [
+                {
+                    "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])}},
+                },
+                {
+                    "var1": np.array(
+                        [
+                            (2 * 2.0 + 4 * 1.0) / (2 + 4),
+                            (2 * 3.0 + 4 * 1.0) / (2 + 4),
+                            (2 * 1.1 + 4 * 2.1) / (2 + 4),
+                            (2 * 0.1 + 4 * 0.5) / (2 + 4),
+                        ]
+                    )
+                },
+            ],
+        ],
+    )
+    def test_accum_aggregator(self, aggregator, received, expected):
+        aggregation_weights = {k: v["weight"] for k, v in received.items()}
+        agg = aggregator(aggregation_weights=aggregation_weights)
+        fl_ctx_mgr = FLContextManager(engine=None, identity_name="", run_num=1, public_stickers={}, private_stickers={})
+        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"],
+                },
+            )
+
+            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)
+            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("aggregator", [InTimeAccumulateWeightedAggregator])
+    @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):
+        aggregation_weights = {f"client_{i}": random.random() for i in range(n_clients)}
+        agg = aggregator(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={})
+        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)
+            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 * aggregation_weights[client_name])
+            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("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
+        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)
+
+        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)