Repository: NVIDIA/NVFlare

Issue:
Error when saving HE encrypted MONAI model
Error reported by @LSnyd

I was able to reproduce in 2.0.16

```
2022-03-11 18:00:37,429 - HEInTimeAccumulateWeightedAggregator - INFO - [run=1, wf=scatter_and_gather, peer=site-1, peer_run=1, peer_rc=OK, task_name=train, task_id=843b287a-b255-450f-8eb3-05a96acd6e5d]: Round 0 adding site-1 time is 0.9321484565734863 seconds
2022-03-11 18:00:37,429 - ScatterAndGather - INFO - [run=1, wf=scatter_and_gather, peer=site-1, peer_run=1, peer_rc=OK, task_name=train, task_id=843b287a-b255-450f-8eb3-05a96acd6e5d]: Contribution from site-1 ACCEPTED by the aggregator.
Traceback (most recent call last):
2022-03-11 18:00:37,429 - ScatterAndGather - INFO - [run=1, wf=scatter_and_gather, peer=site-1, peer_run=1, peer_rc=OK, task_name=train, task_id=843b287a-b255-450f-8eb3-05a96acd6e5d]: Result of unknown task train sent to aggregator.
2022-03-11 18:00:37,429 - ServerRunner - INFO - [run=1, wf=scatter_and_gather, peer=site-1, peer_run=1, peer_rc=OK, task_name=train, task_id=843b287a-b255-450f-8eb3-05a96acd6e5d]: finished processing client result by scatter_and_gather
  File "/home/lisa/miniconda3/envs/FL_monai_nvflare/lib/python3.8/site-packages/nvflare/app_common/workflows/scatter_and_gather.py", line 211, in control_flow
    self.persistor.save(self._global_weights, fl_ctx)
  File "/home/lisa/miniconda3/envs/FL_monai_nvflare/lib/python3.8/site-packages/nvflare/app_common/abstract/model_persistor.py", line 29, in save
    self.save_model(learnable, fl_ctx)
  File "/home/lisa/miniconda3/envs/FL_monai_nvflare/lib/python3.8/site-packages/nvflare/app_common/pt/pt_file_model_persistor.py", line 229, in save_model
    self.save_model_file(self._ckpt_save_path)
2022-03-11 18:00:37,446 - ServerRunner - INFO - [run=1, wf=scatter_and_gather, peer=site-3, peer_run=1]: got task request from client
  File "/home/lisa/miniconda3/envs/FL_monai_nvflare/lib/python3.8/site-packages/nvflare/app_common/pt/pt_file_model_persistor.py", line 224, in save_model_file
    torch.save(save_dict, save_path)
2022-03-11 18:00:37,447 - ServerRunner - INFO - [run=1, wf=scatter_and_gather, peer=site-3, peer_run=1]: no task currently for client - asked client to try again later
  File "/home/lisa/miniconda3/envs/FL_monai_nvflare/lib/python3.8/site-packages/torch/serialization.py", line 379, in save
    _save(obj, opened_zipfile, pickle_module, pickle_protocol)
2022-03-11 18:00:37,447 - FederatedServer - INFO - Return task:__try_again__ to client:site-3 --- (efeeb23c-a6fb-4925-939a-5459c3bb64e2) 
  File "/home/lisa/miniconda3/envs/FL_monai_nvflare/lib/python3.8/site-packages/torch/serialization.py", line 484, in _save
    pickler.dump(obj)
TypeError: cannot pickle '_tenseal_cpp.CKKSVector' object
```

_Originally posted by @LSnyd in https://github.com/NVIDIA/NVFlare/discussions/234#discussioncomment-2340731_


Discussion:

Resolving patch:
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 @@ def accept(self, shareable: Shareable, fl_ctx: FLContext) -> bool:
             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)
diff --git a/nvflare/app_common/homomorphic_encryption/he_model_decryptor.py b/nvflare/app_common/homomorphic_encryption/he_model_decryptor.py
--- a/nvflare/app_common/homomorphic_encryption/he_model_decryptor.py
+++ b/nvflare/app_common/homomorphic_encryption/he_model_decryptor.py
@@ -130,16 +130,16 @@ def _process(self, shareable: Shareable, fl_ctx: FLContext):
         self.log_info(fl_ctx, "Running decryption...")
         dxo = from_shareable(shareable)
 
-        encrypted_algo = dxo.get_meta_prop(key=MetaKey.PROCESSED_ALGORITHM, default=None)
-        if encrypted_algo != he.HE_ALGORITHM_CKKS:
-            self.log_error(fl_ctx, "shareable is not HE CKKS encrypted")
-            return shareable
-
         encrypted_layers = dxo.get_meta_prop(key=MetaKey.PROCESSED_KEYS, default=None)
         if not encrypted_layers:
             self.log_warning(fl_ctx, "dxo does not contain PROCESSED_KEYS (do nothing)")
             return shareable
 
+        encrypted_algo = dxo.get_meta_prop(key=MetaKey.PROCESSED_ALGORITHM, default=None)
+        if encrypted_algo != he.HE_ALGORITHM_CKKS:
+            self.log_error(fl_ctx, "shareable is not HE CKKS encrypted")
+            return shareable
+
         n_encrypted, n_total = count_encrypted_layers(encrypted_layers)
         self.log_info(fl_ctx, f"{n_encrypted} of {n_total} layers encrypted")
         decrypted_params = self.decryption(
diff --git a/nvflare/app_common/homomorphic_encryption/he_model_shareable_generator.py b/nvflare/app_common/homomorphic_encryption/he_model_shareable_generator.py
--- a/nvflare/app_common/homomorphic_encryption/he_model_shareable_generator.py
+++ b/nvflare/app_common/homomorphic_encryption/he_model_shareable_generator.py
@@ -122,6 +122,10 @@ def _shareable_to_learnable(self, shareable: Shareable, fl_ctx: FLContext) -> Mo
             )
         elif dxo.data_kind == DataKind.WEIGHTS:
             weights = dxo.data
+            for v_name in weights.keys():
+                if encrypt_layers[v_name]:
+                    self.log_info(fl_ctx, f"serialize encrypted {dxo.data_kind}: {v_name}")
+                    weights[v_name] = weights[v_name].serialize()
             base_model[ModelLearnableKey.WEIGHTS] = weights
         else:
             raise NotImplementedError(f"data type {dxo.data_kind} not supported!")


Test patch: