From b815f8483120b1c25e4ba33fef7c297e8b4ee817 Mon Sep 17 00:00:00 2001
From: Payal Chaurasiya <payal.chaurasiya@intel.com>
Date: Tue, 10 Dec 2024 21:26:09 +0530
Subject: [PATCH] Remove log_metric_callback and write_metric from aggregator

Remove `log_metric_callback` and its usage from `Aggregator` class in `openfl/component/aggregator/aggregator.py`.

* Remove the import of `write_metric`.
* Remove the `log_metric_callback` attribute from the `Aggregator` class.
* Remove the initialization and usage of `log_metric_callback` in the `Aggregator` class constructor.
* Remove the conditional logging logic that uses `log_metric_callback` in the `Aggregator` class methods.
---
 openfl/component/aggregator/aggregator.py | 31 +----------------------
 1 file changed, 1 insertion(+), 30 deletions(-)

diff --git a/openfl/component/aggregator/aggregator.py b/openfl/component/aggregator/aggregator.py
index 1e34aa7e92a..e872bb7c766 100644
--- a/openfl/component/aggregator/aggregator.py
+++ b/openfl/component/aggregator/aggregator.py
@@ -15,7 +15,7 @@
 from openfl.pipelines import NoCompressionPipeline, TensorCodec
 from openfl.protocols import base_pb2, utils
 from openfl.utilities import TaskResultKey, TensorKey, change_tags
-from openfl.utilities.logs import get_memory_usage, write_memory_usage_to_file, write_metric
+from openfl.utilities.logs import get_memory_usage, write_memory_usage_to_file
 
 
 class Aggregator:
@@ -39,7 +39,6 @@ class Aggregator:
         db_store_rounds* (int): Rounds to store in TensorDB.
         logger: Object for logging.
         write_logs (bool): Flag to enable log writing.
-        log_metric_callback: Callback for logging metrics.
         best_model_score (optional): Score of the best model. Defaults to
             None.
         metric_queue (queue.Queue): Queue for metrics.
@@ -78,7 +77,6 @@ def __init__(
         db_store_rounds=1,
         write_logs=False,
         log_memory_usage=False,
-        log_metric_callback=None,
         initial_tensor_dict=None,
     ):
         """Initializes the Aggregator.
@@ -106,8 +104,6 @@ def __init__(
                 Defaults to 1.
             write_logs (bool, optional): Whether to write logs. Defaults to
                 False.
-            log_metric_callback (optional): Callback for log metric. Defaults
-                to None.
             **kwargs: Additional keyword arguments.
         """
         self.logger = getLogger(__name__)
@@ -147,13 +143,6 @@ def __init__(
 
         # Gathered together logging-related objects
         self.write_logs = write_logs
-        self.log_metric_callback = log_metric_callback
-
-        if self.write_logs:
-            self.log_metric = write_metric
-            if self.log_metric_callback:
-                self.log_metric = log_metric_callback
-                self.logger.info("Using custom log metric: %s", self.log_metric)
 
         self.best_model_score = None
         self.metric_queue = queue.Queue()
@@ -664,14 +653,6 @@ def send_local_task_results(
                 }
                 self.metric_queue.put(metrics)
                 self.logger.metric("%s", str(metrics))
-                if self.write_logs:
-                    self.log_metric(
-                        collaborator_name,
-                        task_name,
-                        tensor_key.tensor_name,
-                        float(value),
-                        round_number,
-                    )
 
             task_results.append(tensor_key)
 
@@ -717,9 +698,7 @@ def _process_named_tensor(self, named_tensor, collaborator_name):
 
         Returns:
             tensor_key (TensorKey): Tensor key.
-                The tensorkey extracted from the protobuf.
             nparray (np.array): Numpy array.
-                The numpy array associated with the returned tensorkey.
         """
         raw_bytes = named_tensor.data_bytes
         metadata = [
@@ -946,14 +925,6 @@ def _compute_validation_related_task_metrics(self, task_name):
 
                 self.metric_queue.put(metrics)
                 self.logger.metric("%s", metrics)
-                if self.write_logs:
-                    self.log_metric(
-                        "aggregator",
-                        task_name,
-                        tensor_key.tensor_name,
-                        float(agg_results),
-                        round_number,
-                    )
 
                 # FIXME: Configurable logic for min/max criteria in saving best.
                 if "validate_agg" in tags: