Added metric origin to output

Signed-off-by: Chaurasiya, Payal <[email protected]>
securefederatedai · Nov 21, 2024 · 4d42e48 · 4d42e48
1 parent d8a52b1
commit 4d42e48
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 0 deletions.
diff --git a/openfl/component/aggregator/aggregator.py b/openfl/component/aggregator/aggregator.py
@@ -976,6 +976,7 @@ def _end_of_round_check(self):
             # This is the place to check the memory usage of the aggregator
             memory_detail = get_memory_usage()
             memory_detail["round_number"] = self.round_number
+            memory_detail["metric_origin"] = "aggregator"
             self.memory_details.append(memory_detail)
 
         # Once all of the task results have been processed

diff --git a/openfl/component/collaborator/collaborator.py b/openfl/component/collaborator/collaborator.py
@@ -178,6 +178,7 @@ def run(self):
                     # This is the place to check the memory usage of the collaborator
                     memory_detail = get_memory_usage()
                     memory_detail["round_number"] = round_number
+                    memory_details["metric_origin"] = self.collaborator_name
                     memory_details.append(memory_detail)
         if self.log_memory_usage:
             self.logger.info(f"Publish memory usage: {memory_details}")