We are currently solving the following issue within our repository. Here is the issue text:
--- BEGIN ISSUE ---
HE aggregator has wrong check for data kind
This [block](https://github.com/NVIDIA/NVFlare/blob/52fa8fc989811526c50ecd4030cc89141b1b26cd/nvflare/app_common/homomorphic_encryption/he_intime_accumulate_model_aggregator.py#L104) should be replaced with 

```
        if dxo.data_kind != self.expected_data_kind:
            self.log_error(fl_ctx, "expected {self.expected_data_kind} type DXO only, skipping this shareable.")
            return False
```

_Originally posted by @holgerroth in https://github.com/NVIDIA/NVFlare/discussions/234#discussioncomment-2340297_

--- END ISSUE ---

Below are some code segments, each from a relevant file. One or more of these files may contain bugs.
--- BEGIN FILES ---
Path: `nvflare/app_common/homomorphic_encryption/he_intime_accumulate_model_aggregator.py`
Content:
```
1 # Copyright (c) 2021-2022, NVIDIA CORPORATION.  All rights reserved.
2 #
3 # Licensed under the Apache License, Version 2.0 (the "License");
4 # you may not use this file except in compliance with the License.
5 # You may obtain a copy of the License at
6 #
7 #     http://www.apache.org/licenses/LICENSE-2.0
8 #
9 # Unless required by applicable law or agreed to in writing, software
10 # distributed under the License is distributed on an "AS IS" BASIS,
11 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12 # See the License for the specific language governing permissions and
13 # limitations under the License.
14 
15 import re
16 import time
17 
18 import numpy as np
19 import tenseal as ts
20 
21 import nvflare.app_common.homomorphic_encryption.he_constant as he
22 from nvflare.apis.dxo import DXO, DataKind, MetaKey, from_shareable
23 from nvflare.apis.event_type import EventType
24 from nvflare.apis.fl_constant import ReservedKey, ReturnCode
25 from nvflare.apis.fl_context import FLContext
26 from nvflare.apis.shareable import Shareable
27 from nvflare.app_common.abstract.aggregator import Aggregator
28 from nvflare.app_common.app_constant import AppConstants
29 from nvflare.app_common.homomorphic_encryption.homomorphic_encrypt import (
30     count_encrypted_layers,
31     load_tenseal_context_from_workspace,
32 )
33 
34 
35 class HEInTimeAccumulateWeightedAggregator(Aggregator):
36     def __init__(
37         self,
38         exclude_vars=None,
39         aggregation_weights=None,
40         tenseal_context_file="server_context.tenseal",
41         weigh_by_local_iter=False,
42         expected_data_kind="WEIGHT_DIFF",
43         expected_algorithm=he.HE_ALGORITHM_CKKS,
44     ):
45         """In time aggregator for `Shareables` encrypted using homomorphic encryption (HE) with TenSEAL https://github.com/OpenMined/TenSEAL.
46 
47         Args:
48             exclude_vars ([list], optional): variable names that should be excluded from aggregation (use regular expression). Defaults to None.
49             aggregation_weights ([dict], optional): dictionary of client aggregation. Defaults to None.
50             tenseal_context_file (str, optional): [description]. Defaults to "server_context.tenseal".
51             weigh_by_local_iter (bool, optional): If true, multiply client weights on first in encryption space
52                                  (default: `False` which is recommended for HE, first multiply happens in `HEModelEncryptor`)].
53             expected_data_kind (str, optional): the data_kind this aggregator can process. Defaults to "WEIGHT_DIFF".
54             expected_algorithm ([str], optional): the HE algorithm it can process. Defaults to he.HE_ALGORITHM_CKKS.
55 
56         Raises:
57             ValueError: mismatched data_kind or HE algorithm
58         """
59         super().__init__()
60         self.tenseal_context = None
61         self.tenseal_context_file = tenseal_context_file
62         if expected_data_kind not in [DataKind.WEIGHT_DIFF, DataKind.WEIGHTS]:
63             raise ValueError(f"expected_data_kind={expected_data_kind} not in WEIGHT_DIFF or WEIGHTS")
64         self.expected_data_kind = expected_data_kind
65         self.expected_algorithm = expected_algorithm
66         if self.expected_algorithm != he.HE_ALGORITHM_CKKS:
67             raise ValueError(f"expected algorithm {self.expected_algorithm} not supported")
68         self.exclude_vars = re.compile(exclude_vars) if exclude_vars else None
69         self.aggregation_weights = aggregation_weights or {}
70         self.reset_stats()
71         self.weigh_by_local_iter = weigh_by_local_iter
72         self.logger.info(f"client weights control: {self.aggregation_weights}")
73         if not self.weigh_by_local_iter:
74             if self.aggregation_weights:
75                 self.logger.warning("aggregation_weights will be ignored if weigh_by_local_iter=False")
76             self.logger.info("Only divide by sum of local (weighted) iterations.")
77         self.warning_count = dict()
78         self.warning_limit = 0
79 
80     def handle_event(self, event_type: str, fl_ctx: FLContext):
81         if event_type == EventType.START_RUN:
82             self.tenseal_context = load_tenseal_context_from_workspace(self.tenseal_context_file, fl_ctx)
83         elif event_type == EventType.END_RUN:
84             self.tenseal_context = None
85 
86     def reset_stats(self):
87         self.total = dict()
88         self.counts = dict()
89         self.contribution_count = 0
90         self.history = list()
91         self.merged_encrypted_layers = dict()  # thread-safety is handled by workflow
92 
93     def accept(self, shareable: Shareable, fl_ctx: FLContext) -> bool:
94         """Accepts and adds the client updates to current average in HE encrypted space.
95 
96         Args:
97             shareable: a shareable from client
98             fl_ctx: FL Contenxt associated with this shareable
99 
100         Returns:
101             bool to indicate if this shareable is accepted.
102         """
103         dxo = from_shareable(shareable)
104         if dxo.data_kind != DataKind.WEIGHT_DIFF:
105             self.log_error(fl_ctx, "support WEIGHT_DIFF type DXO only, skipping this shareable.")
106             return False
107 
108         enc_algo = dxo.get_meta_prop(key=MetaKey.PROCESSED_ALGORITHM, default=None)
109         if enc_algo != self.expected_algorithm:
110             self.log_error(fl_ctx, "unsupported encryption algorithm {enc_algo}")
111             return False
112 
113         current_round = fl_ctx.get_prop(AppConstants.CURRENT_ROUND)
114         client_name = shareable.get_peer_prop(ReservedKey.IDENTITY_NAME, "?")
115         contribution_round = shareable.get_header(AppConstants.CONTRIBUTION_ROUND)
116 
117         rc = shareable.get_return_code()
118         if rc and rc != ReturnCode.OK:
119             self.log_debug(fl_ctx, f"Client {client_name} returned rc: {rc}. Disregarding contribution.")
120             return False
121 
122         self.log_debug(fl_ctx, f"current_round: {current_round}")
123 
124         if contribution_round != current_round:
125             self.log_debug(
126                 fl_ctx,
127                 "Discarded the contribution from {client_name} for round: {contribution_round}. Current round is: {current_round}",
128             )
129             return False
130 
131         start_time = time.time()
132 
133         for item in self.history:
134             if client_name == item["client_name"]:
135                 prev_round = item["round"]
136                 self.log_info(
137                     fl_ctx,
138                     f"discarding shareable from {client_name} at round: {contribution_round} as {prev_round} accepted already",
139                 )
140                 return False
141 
142         self.log_info(fl_ctx, f"Adding contribution from {client_name}.")
143 
144         n_iter = dxo.get_meta_prop(key=MetaKey.NUM_STEPS_CURRENT_ROUND)
145         if n_iter is None:
146             if self.warning_count.get(client_name, 0) <= self.warning_limit:
147                 self.log_warning(
148                     fl_ctx,
149                     f"NUM_STEPS_CURRENT_ROUND missing"
150                     f" from {client_name} and set to default value, 1.0. "
151                     f" This kind of message will show {self.warning_limit} times at most.",
152                 )
153                 if client_name in self.warning_count:
154                     self.warning_count[client_name] = self.warning_count[client_name] + 1
155                 else:
156                     self.warning_count[client_name] = 0
157             n_iter = 1.0
158         float_n_iter = np.float(n_iter)
159 
160         aggregation_weight = self.aggregation_weights.get(client_name)
161         if aggregation_weight is None:
162             aggregation_weight = 1.0
163 
164         aggr_data = dxo.data
165         encrypted_layers = dxo.get_meta_prop(MetaKey.PROCESSED_KEYS)
166         # TODO: test support of different encrypted layers for different clients!
167 
168         if encrypted_layers is None:
169             self.log_error(fl_ctx, "encrypted_layers is None!")
170             return False
171 
172         for k, v in aggr_data.items():
173             if self.exclude_vars is not None and self.exclude_vars.search(k):
174                 continue
175             if encrypted_layers[k]:
176                 if self.weigh_by_local_iter:
177                     weighted_value = ts.ckks_vector_from(self.tenseal_context, v) * (aggregation_weight * float_n_iter)
178                 else:
179                     weighted_value = ts.ckks_vector_from(self.tenseal_context, v)
180                 self.merged_encrypted_layers[k] = True  # any client can set this true
181             else:
182                 if self.weigh_by_local_iter:
183                     weighted_value = v * (aggregation_weight * float_n_iter)
184                 else:
185                     weighted_value = v
186                 if k not in self.merged_encrypted_layers:
187                     self.merged_encrypted_layers[k] = False  # only set False if no other client set it to True
188             current_total = self.total.get(k, None)
189             if current_total is None:
190                 self.total[k] = weighted_value
191                 self.counts[k] = n_iter
192             else:
193                 self.total[k] = current_total + weighted_value
194                 self.counts[k] = self.counts[k] + n_iter
195 
196         self.contribution_count += 1
197 
198         end_time = time.time()
199         n_encrypted, n_total = count_encrypted_layers(self.merged_encrypted_layers)
200         self.log_info(fl_ctx, f"{n_encrypted} of {n_total} layers encrypted")
201         self.log_info(fl_ctx, f"Round {current_round} adding {client_name} time is {end_time - start_time} seconds")
202 
203         self.history.append(
204             {
205                 "client_name": client_name,
206                 "round": contribution_round,
207                 "aggregation_weight": aggregation_weight,
208                 "n_iter": n_iter,
209             }
210         )
211         return True
212 
213     def aggregate(self, fl_ctx: FLContext) -> Shareable:
214         start_time = time.time()
215         current_round = fl_ctx.get_prop(AppConstants.CURRENT_ROUND)
216 
217         aggregated_dict = dict()
218         for k, v in self.total.items():
219             aggregated_dict[k] = v * (1.0 / self.counts[k])
220         end_time = time.time()
221         self.log_info(
222             fl_ctx,
223             f"Aggregated {self.contribution_count} contributions for round {current_round} time is {end_time - start_time} seconds",
224         )
225 
226         dxo = DXO(data_kind=self.expected_data_kind, data=aggregated_dict)
227         dxo.set_meta_prop(MetaKey.PROCESSED_KEYS, self.merged_encrypted_layers)
228         dxo.set_meta_prop(MetaKey.PROCESSED_ALGORITHM, self.expected_algorithm)
229         n_encrypted, n_total = count_encrypted_layers(self.merged_encrypted_layers)
230         self.log_info(fl_ctx, f"{n_encrypted} of {n_total} layers encrypted")
231 
232         fl_ctx.set_prop(AppConstants.DXO, dxo, private=True, sticky=False)
233 
234         self.reset_stats()  # only reset dictionary after adding merged_encrypted_layers to dictionary
235         return dxo.to_shareable()
236
```
--- END FILES ---

Please first localize the bug based on the issue statement, and then generate a patch according to the `git diff` format fenced by three backticks.

Here is an example:

```diff
diff --git a/examples/server_async.py b/examples/server_async.py
--- a/examples/server_async.py
+++ b/examples/server_async.py
@@ -313,4 +313,4 @@
 
 
 if __name__ == "__main__":
-    asyncio.run(run_async_server("."), debug=True)
+    asyncio.run(run_async_server(), debug=True)
diff --git a/examples/server_sync.py b/examples/server_sync.py
--- a/examples/server_sync.py
+++ b/examples/server_sync.py
@@ -313,5 +313,5 @@
 
 
 if __name__ == "__main__":
-    server = run_sync_server(".")
+    server = run_sync_server()
     server.shutdown()

```

diff --git a/nvflare/app_common/homomorphic_encryption/he_intime_accumulate_model_aggregator.py b/nvflare/app_common/homomorphic_encryption/he_intime_accumulate_model_aggregator.py
--- a/nvflare/app_common/homomorphic_encryption/he_intime_accumulate_model_aggregator.py
+++ b/nvflare/app_common/homomorphic_encryption/he_intime_accumulate_model_aggregator.py
@@ -101,8 +101,11 @@
             bool to indicate if this shareable is accepted.
         """
         dxo = from_shareable(shareable)
-        if dxo.data_kind != DataKind.WEIGHT_DIFF:
-            self.log_error(fl_ctx, "support WEIGHT_DIFF type DXO only, skipping this shareable.")
+        if dxo.data_kind != self.expected_data_kind:
+            self.log_error(
+                fl_ctx,
+                f"expected {self.expected_data_kind} type DXO only but received {dxo.data_kind}, skipping this shareable.",
+            )
             return False
 
         enc_algo = dxo.get_meta_prop(key=MetaKey.PROCESSED_ALGORITHM, default=None)

{"golden_diff": "diff --git a/nvflare/app_common/homomorphic_encryption/he_intime_accumulate_model_aggregator.py b/nvflare/app_common/homomorphic_encryption/he_intime_accumulate_model_aggregator.py\n--- a/nvflare/app_common/homomorphic_encryption/he_intime_accumulate_model_aggregator.py\n+++ b/nvflare/app_common/homomorphic_encryption/he_intime_accumulate_model_aggregator.py\n@@ -101,8 +101,11 @@\n             bool to indicate if this shareable is accepted.\n         \"\"\"\n         dxo = from_shareable(shareable)\n-        if dxo.data_kind != DataKind.WEIGHT_DIFF:\n-            self.log_error(fl_ctx, \"support WEIGHT_DIFF type DXO only, skipping this shareable.\")\n+        if dxo.data_kind != self.expected_data_kind:\n+            self.log_error(\n+                fl_ctx,\n+                f\"expected {self.expected_data_kind} type DXO only but received {dxo.data_kind}, skipping this shareable.\",\n+            )\n             return False\n \n         enc_algo = dxo.get_meta_prop(key=MetaKey.PROCESSED_ALGORITHM, default=None)\n", "issue": "HE aggregator has wrong check for data kind\nThis [block](https://github.com/NVIDIA/NVFlare/blob/52fa8fc989811526c50ecd4030cc89141b1b26cd/nvflare/app_common/homomorphic_encryption/he_intime_accumulate_model_aggregator.py#L104) should be replaced with \r\n\r\n```\r\n        if dxo.data_kind != self.expected_data_kind:\r\n            self.log_error(fl_ctx, \"expected {self.expected_data_kind} type DXO only, skipping this shareable.\")\r\n            return False\r\n```\r\n\r\n_Originally posted by @holgerroth in https://github.com/NVIDIA/NVFlare/discussions/234#discussioncomment-2340297_\n", "before_files": [{"content": "# Copyright (c) 2021-2022, NVIDIA CORPORATION.  All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nimport re\nimport time\n\nimport numpy as np\nimport tenseal as ts\n\nimport nvflare.app_common.homomorphic_encryption.he_constant as he\nfrom nvflare.apis.dxo import DXO, DataKind, MetaKey, from_shareable\nfrom nvflare.apis.event_type import EventType\nfrom nvflare.apis.fl_constant import ReservedKey, ReturnCode\nfrom nvflare.apis.fl_context import FLContext\nfrom nvflare.apis.shareable import Shareable\nfrom nvflare.app_common.abstract.aggregator import Aggregator\nfrom nvflare.app_common.app_constant import AppConstants\nfrom nvflare.app_common.homomorphic_encryption.homomorphic_encrypt import (\n    count_encrypted_layers,\n    load_tenseal_context_from_workspace,\n)\n\n\nclass HEInTimeAccumulateWeightedAggregator(Aggregator):\n    def __init__(\n        self,\n        exclude_vars=None,\n        aggregation_weights=None,\n        tenseal_context_file=\"server_context.tenseal\",\n        weigh_by_local_iter=False,\n        expected_data_kind=\"WEIGHT_DIFF\",\n        expected_algorithm=he.HE_ALGORITHM_CKKS,\n    ):\n        \"\"\"In time aggregator for `Shareables` encrypted using homomorphic encryption (HE) with TenSEAL https://github.com/OpenMined/TenSEAL.\n\n        Args:\n            exclude_vars ([list], optional): variable names that should be excluded from aggregation (use regular expression). Defaults to None.\n            aggregation_weights ([dict], optional): dictionary of client aggregation. Defaults to None.\n            tenseal_context_file (str, optional): [description]. Defaults to \"server_context.tenseal\".\n            weigh_by_local_iter (bool, optional): If true, multiply client weights on first in encryption space\n                                 (default: `False` which is recommended for HE, first multiply happens in `HEModelEncryptor`)].\n            expected_data_kind (str, optional): the data_kind this aggregator can process. Defaults to \"WEIGHT_DIFF\".\n            expected_algorithm ([str], optional): the HE algorithm it can process. Defaults to he.HE_ALGORITHM_CKKS.\n\n        Raises:\n            ValueError: mismatched data_kind or HE algorithm\n        \"\"\"\n        super().__init__()\n        self.tenseal_context = None\n        self.tenseal_context_file = tenseal_context_file\n        if expected_data_kind not in [DataKind.WEIGHT_DIFF, DataKind.WEIGHTS]:\n            raise ValueError(f\"expected_data_kind={expected_data_kind} not in WEIGHT_DIFF or WEIGHTS\")\n        self.expected_data_kind = expected_data_kind\n        self.expected_algorithm = expected_algorithm\n        if self.expected_algorithm != he.HE_ALGORITHM_CKKS:\n            raise ValueError(f\"expected algorithm {self.expected_algorithm} not supported\")\n        self.exclude_vars = re.compile(exclude_vars) if exclude_vars else None\n        self.aggregation_weights = aggregation_weights or {}\n        self.reset_stats()\n        self.weigh_by_local_iter = weigh_by_local_iter\n        self.logger.info(f\"client weights control: {self.aggregation_weights}\")\n        if not self.weigh_by_local_iter:\n            if self.aggregation_weights:\n                self.logger.warning(\"aggregation_weights will be ignored if weigh_by_local_iter=False\")\n            self.logger.info(\"Only divide by sum of local (weighted) iterations.\")\n        self.warning_count = dict()\n        self.warning_limit = 0\n\n    def handle_event(self, event_type: str, fl_ctx: FLContext):\n        if event_type == EventType.START_RUN:\n            self.tenseal_context = load_tenseal_context_from_workspace(self.tenseal_context_file, fl_ctx)\n        elif event_type == EventType.END_RUN:\n            self.tenseal_context = None\n\n    def reset_stats(self):\n        self.total = dict()\n        self.counts = dict()\n        self.contribution_count = 0\n        self.history = list()\n        self.merged_encrypted_layers = dict()  # thread-safety is handled by workflow\n\n    def accept(self, shareable: Shareable, fl_ctx: FLContext) -> bool:\n        \"\"\"Accepts and adds the client updates to current average in HE encrypted space.\n\n        Args:\n            shareable: a shareable from client\n            fl_ctx: FL Contenxt associated with this shareable\n\n        Returns:\n            bool to indicate if this shareable is accepted.\n        \"\"\"\n        dxo = from_shareable(shareable)\n        if dxo.data_kind != DataKind.WEIGHT_DIFF:\n            self.log_error(fl_ctx, \"support WEIGHT_DIFF type DXO only, skipping this shareable.\")\n            return False\n\n        enc_algo = dxo.get_meta_prop(key=MetaKey.PROCESSED_ALGORITHM, default=None)\n        if enc_algo != self.expected_algorithm:\n            self.log_error(fl_ctx, \"unsupported encryption algorithm {enc_algo}\")\n            return False\n\n        current_round = fl_ctx.get_prop(AppConstants.CURRENT_ROUND)\n        client_name = shareable.get_peer_prop(ReservedKey.IDENTITY_NAME, \"?\")\n        contribution_round = shareable.get_header(AppConstants.CONTRIBUTION_ROUND)\n\n        rc = shareable.get_return_code()\n        if rc and rc != ReturnCode.OK:\n            self.log_debug(fl_ctx, f\"Client {client_name} returned rc: {rc}. Disregarding contribution.\")\n            return False\n\n        self.log_debug(fl_ctx, f\"current_round: {current_round}\")\n\n        if contribution_round != current_round:\n            self.log_debug(\n                fl_ctx,\n                \"Discarded the contribution from {client_name} for round: {contribution_round}. Current round is: {current_round}\",\n            )\n            return False\n\n        start_time = time.time()\n\n        for item in self.history:\n            if client_name == item[\"client_name\"]:\n                prev_round = item[\"round\"]\n                self.log_info(\n                    fl_ctx,\n                    f\"discarding shareable from {client_name} at round: {contribution_round} as {prev_round} accepted already\",\n                )\n                return False\n\n        self.log_info(fl_ctx, f\"Adding contribution from {client_name}.\")\n\n        n_iter = dxo.get_meta_prop(key=MetaKey.NUM_STEPS_CURRENT_ROUND)\n        if n_iter is None:\n            if self.warning_count.get(client_name, 0) <= self.warning_limit:\n                self.log_warning(\n                    fl_ctx,\n                    f\"NUM_STEPS_CURRENT_ROUND missing\"\n                    f\" from {client_name} and set to default value, 1.0. \"\n                    f\" This kind of message will show {self.warning_limit} times at most.\",\n                )\n                if client_name in self.warning_count:\n                    self.warning_count[client_name] = self.warning_count[client_name] + 1\n                else:\n                    self.warning_count[client_name] = 0\n            n_iter = 1.0\n        float_n_iter = np.float(n_iter)\n\n        aggregation_weight = self.aggregation_weights.get(client_name)\n        if aggregation_weight is None:\n            aggregation_weight = 1.0\n\n        aggr_data = dxo.data\n        encrypted_layers = dxo.get_meta_prop(MetaKey.PROCESSED_KEYS)\n        # TODO: test support of different encrypted layers for different clients!\n\n        if encrypted_layers is None:\n            self.log_error(fl_ctx, \"encrypted_layers is None!\")\n            return False\n\n        for k, v in aggr_data.items():\n            if self.exclude_vars is not None and self.exclude_vars.search(k):\n                continue\n            if encrypted_layers[k]:\n                if self.weigh_by_local_iter:\n                    weighted_value = ts.ckks_vector_from(self.tenseal_context, v) * (aggregation_weight * float_n_iter)\n                else:\n                    weighted_value = ts.ckks_vector_from(self.tenseal_context, v)\n                self.merged_encrypted_layers[k] = True  # any client can set this true\n            else:\n                if self.weigh_by_local_iter:\n                    weighted_value = v * (aggregation_weight * float_n_iter)\n                else:\n                    weighted_value = v\n                if k not in self.merged_encrypted_layers:\n                    self.merged_encrypted_layers[k] = False  # only set False if no other client set it to True\n            current_total = self.total.get(k, None)\n            if current_total is None:\n                self.total[k] = weighted_value\n                self.counts[k] = n_iter\n            else:\n                self.total[k] = current_total + weighted_value\n                self.counts[k] = self.counts[k] + n_iter\n\n        self.contribution_count += 1\n\n        end_time = time.time()\n        n_encrypted, n_total = count_encrypted_layers(self.merged_encrypted_layers)\n        self.log_info(fl_ctx, f\"{n_encrypted} of {n_total} layers encrypted\")\n        self.log_info(fl_ctx, f\"Round {current_round} adding {client_name} time is {end_time - start_time} seconds\")\n\n        self.history.append(\n            {\n                \"client_name\": client_name,\n                \"round\": contribution_round,\n                \"aggregation_weight\": aggregation_weight,\n                \"n_iter\": n_iter,\n            }\n        )\n        return True\n\n    def aggregate(self, fl_ctx: FLContext) -> Shareable:\n        start_time = time.time()\n        current_round = fl_ctx.get_prop(AppConstants.CURRENT_ROUND)\n\n        aggregated_dict = dict()\n        for k, v in self.total.items():\n            aggregated_dict[k] = v * (1.0 / self.counts[k])\n        end_time = time.time()\n        self.log_info(\n            fl_ctx,\n            f\"Aggregated {self.contribution_count} contributions for round {current_round} time is {end_time - start_time} seconds\",\n        )\n\n        dxo = DXO(data_kind=self.expected_data_kind, data=aggregated_dict)\n        dxo.set_meta_prop(MetaKey.PROCESSED_KEYS, self.merged_encrypted_layers)\n        dxo.set_meta_prop(MetaKey.PROCESSED_ALGORITHM, self.expected_algorithm)\n        n_encrypted, n_total = count_encrypted_layers(self.merged_encrypted_layers)\n        self.log_info(fl_ctx, f\"{n_encrypted} of {n_total} layers encrypted\")\n\n        fl_ctx.set_prop(AppConstants.DXO, dxo, private=True, sticky=False)\n\n        self.reset_stats()  # only reset dictionary after adding merged_encrypted_layers to dictionary\n        return dxo.to_shareable()\n", "path": "nvflare/app_common/homomorphic_encryption/he_intime_accumulate_model_aggregator.py"}], "after_files": [{"content": "# Copyright (c) 2021-2022, NVIDIA CORPORATION.  All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nimport re\nimport time\n\nimport numpy as np\nimport tenseal as ts\n\nimport nvflare.app_common.homomorphic_encryption.he_constant as he\nfrom nvflare.apis.dxo import DXO, DataKind, MetaKey, from_shareable\nfrom nvflare.apis.event_type import EventType\nfrom nvflare.apis.fl_constant import ReservedKey, ReturnCode\nfrom nvflare.apis.fl_context import FLContext\nfrom nvflare.apis.shareable import Shareable\nfrom nvflare.app_common.abstract.aggregator import Aggregator\nfrom nvflare.app_common.app_constant import AppConstants\nfrom nvflare.app_common.homomorphic_encryption.homomorphic_encrypt import (\n    count_encrypted_layers,\n    load_tenseal_context_from_workspace,\n)\n\n\nclass HEInTimeAccumulateWeightedAggregator(Aggregator):\n    def __init__(\n        self,\n        exclude_vars=None,\n        aggregation_weights=None,\n        tenseal_context_file=\"server_context.tenseal\",\n        weigh_by_local_iter=False,\n        expected_data_kind=\"WEIGHT_DIFF\",\n        expected_algorithm=he.HE_ALGORITHM_CKKS,\n    ):\n        \"\"\"In time aggregator for `Shareables` encrypted using homomorphic encryption (HE) with TenSEAL https://github.com/OpenMined/TenSEAL.\n\n        Args:\n            exclude_vars ([list], optional): variable names that should be excluded from aggregation (use regular expression). Defaults to None.\n            aggregation_weights ([dict], optional): dictionary of client aggregation. Defaults to None.\n            tenseal_context_file (str, optional): [description]. Defaults to \"server_context.tenseal\".\n            weigh_by_local_iter (bool, optional): If true, multiply client weights on first in encryption space\n                                 (default: `False` which is recommended for HE, first multiply happens in `HEModelEncryptor`)].\n            expected_data_kind (str, optional): the data_kind this aggregator can process. Defaults to \"WEIGHT_DIFF\".\n            expected_algorithm ([str], optional): the HE algorithm it can process. Defaults to he.HE_ALGORITHM_CKKS.\n\n        Raises:\n            ValueError: mismatched data_kind or HE algorithm\n        \"\"\"\n        super().__init__()\n        self.tenseal_context = None\n        self.tenseal_context_file = tenseal_context_file\n        if expected_data_kind not in [DataKind.WEIGHT_DIFF, DataKind.WEIGHTS]:\n            raise ValueError(f\"expected_data_kind={expected_data_kind} not in WEIGHT_DIFF or WEIGHTS\")\n        self.expected_data_kind = expected_data_kind\n        self.expected_algorithm = expected_algorithm\n        if self.expected_algorithm != he.HE_ALGORITHM_CKKS:\n            raise ValueError(f\"expected algorithm {self.expected_algorithm} not supported\")\n        self.exclude_vars = re.compile(exclude_vars) if exclude_vars else None\n        self.aggregation_weights = aggregation_weights or {}\n        self.reset_stats()\n        self.weigh_by_local_iter = weigh_by_local_iter\n        self.logger.info(f\"client weights control: {self.aggregation_weights}\")\n        if not self.weigh_by_local_iter:\n            if self.aggregation_weights:\n                self.logger.warning(\"aggregation_weights will be ignored if weigh_by_local_iter=False\")\n            self.logger.info(\"Only divide by sum of local (weighted) iterations.\")\n        self.warning_count = dict()\n        self.warning_limit = 0\n\n    def handle_event(self, event_type: str, fl_ctx: FLContext):\n        if event_type == EventType.START_RUN:\n            self.tenseal_context = load_tenseal_context_from_workspace(self.tenseal_context_file, fl_ctx)\n        elif event_type == EventType.END_RUN:\n            self.tenseal_context = None\n\n    def reset_stats(self):\n        self.total = dict()\n        self.counts = dict()\n        self.contribution_count = 0\n        self.history = list()\n        self.merged_encrypted_layers = dict()  # thread-safety is handled by workflow\n\n    def accept(self, shareable: Shareable, fl_ctx: FLContext) -> bool:\n        \"\"\"Accepts and adds the client updates to current average in HE encrypted space.\n\n        Args:\n            shareable: a shareable from client\n            fl_ctx: FL Contenxt associated with this shareable\n\n        Returns:\n            bool to indicate if this shareable is accepted.\n        \"\"\"\n        dxo = from_shareable(shareable)\n        if dxo.data_kind != self.expected_data_kind:\n            self.log_error(\n                fl_ctx,\n                f\"expected {self.expected_data_kind} type DXO only but received {dxo.data_kind}, skipping this shareable.\",\n            )\n            return False\n\n        enc_algo = dxo.get_meta_prop(key=MetaKey.PROCESSED_ALGORITHM, default=None)\n        if enc_algo != self.expected_algorithm:\n            self.log_error(fl_ctx, \"unsupported encryption algorithm {enc_algo}\")\n            return False\n\n        current_round = fl_ctx.get_prop(AppConstants.CURRENT_ROUND)\n        client_name = shareable.get_peer_prop(ReservedKey.IDENTITY_NAME, \"?\")\n        contribution_round = shareable.get_header(AppConstants.CONTRIBUTION_ROUND)\n\n        rc = shareable.get_return_code()\n        if rc and rc != ReturnCode.OK:\n            self.log_debug(fl_ctx, f\"Client {client_name} returned rc: {rc}. Disregarding contribution.\")\n            return False\n\n        self.log_debug(fl_ctx, f\"current_round: {current_round}\")\n\n        if contribution_round != current_round:\n            self.log_debug(\n                fl_ctx,\n                \"Discarded the contribution from {client_name} for round: {contribution_round}. Current round is: {current_round}\",\n            )\n            return False\n\n        start_time = time.time()\n\n        for item in self.history:\n            if client_name == item[\"client_name\"]:\n                prev_round = item[\"round\"]\n                self.log_info(\n                    fl_ctx,\n                    f\"discarding shareable from {client_name} at round: {contribution_round} as {prev_round} accepted already\",\n                )\n                return False\n\n        self.log_info(fl_ctx, f\"Adding contribution from {client_name}.\")\n\n        n_iter = dxo.get_meta_prop(key=MetaKey.NUM_STEPS_CURRENT_ROUND)\n        if n_iter is None:\n            if self.warning_count.get(client_name, 0) <= self.warning_limit:\n                self.log_warning(\n                    fl_ctx,\n                    f\"NUM_STEPS_CURRENT_ROUND missing\"\n                    f\" from {client_name} and set to default value, 1.0. \"\n                    f\" This kind of message will show {self.warning_limit} times at most.\",\n                )\n                if client_name in self.warning_count:\n                    self.warning_count[client_name] = self.warning_count[client_name] + 1\n                else:\n                    self.warning_count[client_name] = 0\n            n_iter = 1.0\n        float_n_iter = np.float(n_iter)\n\n        aggregation_weight = self.aggregation_weights.get(client_name)\n        if aggregation_weight is None:\n            aggregation_weight = 1.0\n\n        aggr_data = dxo.data\n        encrypted_layers = dxo.get_meta_prop(MetaKey.PROCESSED_KEYS)\n        # TODO: test support of different encrypted layers for different clients!\n\n        if encrypted_layers is None:\n            self.log_error(fl_ctx, \"encrypted_layers is None!\")\n            return False\n\n        for k, v in aggr_data.items():\n            if self.exclude_vars is not None and self.exclude_vars.search(k):\n                continue\n            if encrypted_layers[k]:\n                if self.weigh_by_local_iter:\n                    weighted_value = ts.ckks_vector_from(self.tenseal_context, v) * (aggregation_weight * float_n_iter)\n                else:\n                    weighted_value = ts.ckks_vector_from(self.tenseal_context, v)\n                self.merged_encrypted_layers[k] = True  # any client can set this true\n            else:\n                if self.weigh_by_local_iter:\n                    weighted_value = v * (aggregation_weight * float_n_iter)\n                else:\n                    weighted_value = v\n                if k not in self.merged_encrypted_layers:\n                    self.merged_encrypted_layers[k] = False  # only set False if no other client set it to True\n            current_total = self.total.get(k, None)\n            if current_total is None:\n                self.total[k] = weighted_value\n                self.counts[k] = n_iter\n            else:\n                self.total[k] = current_total + weighted_value\n                self.counts[k] = self.counts[k] + n_iter\n\n        self.contribution_count += 1\n\n        end_time = time.time()\n        n_encrypted, n_total = count_encrypted_layers(self.merged_encrypted_layers)\n        self.log_info(fl_ctx, f\"{n_encrypted} of {n_total} layers encrypted\")\n        self.log_info(fl_ctx, f\"Round {current_round} adding {client_name} time is {end_time - start_time} seconds\")\n\n        self.history.append(\n            {\n                \"client_name\": client_name,\n                \"round\": contribution_round,\n                \"aggregation_weight\": aggregation_weight,\n                \"n_iter\": n_iter,\n            }\n        )\n        return True\n\n    def aggregate(self, fl_ctx: FLContext) -> Shareable:\n        start_time = time.time()\n        current_round = fl_ctx.get_prop(AppConstants.CURRENT_ROUND)\n\n        aggregated_dict = dict()\n        for k, v in self.total.items():\n            aggregated_dict[k] = v * (1.0 / self.counts[k])\n        end_time = time.time()\n        self.log_info(\n            fl_ctx,\n            f\"Aggregated {self.contribution_count} contributions for round {current_round} time is {end_time - start_time} seconds\",\n        )\n\n        dxo = DXO(data_kind=self.expected_data_kind, data=aggregated_dict)\n        dxo.set_meta_prop(MetaKey.PROCESSED_KEYS, self.merged_encrypted_layers)\n        dxo.set_meta_prop(MetaKey.PROCESSED_ALGORITHM, self.expected_algorithm)\n        n_encrypted, n_total = count_encrypted_layers(self.merged_encrypted_layers)\n        self.log_info(fl_ctx, f\"{n_encrypted} of {n_total} layers encrypted\")\n\n        fl_ctx.set_prop(AppConstants.DXO, dxo, private=True, sticky=False)\n\n        self.reset_stats()  # only reset dictionary after adding merged_encrypted_layers to dictionary\n        return dxo.to_shareable()\n", "path": "nvflare/app_common/homomorphic_encryption/he_intime_accumulate_model_aggregator.py"}]}