- Pull out magic numbers to constants

Signed-off-by: Shailesh Pant <[email protected]>

Author: ishaileshpant

openfl/component/aggregator/aggregator.py

@@ -10,6 +10,7 @@
 from typing import List, Optional
 
 import openfl.callbacks as callbacks_module
+from openfl.component.aggregator import constants
 from openfl.component.aggregator.straggler_handling import CutoffTimePolicy, StragglerPolicy
 from openfl.databases import PersistentTensorDB, TensorDB
 from openfl.interface.aggregation_functions import WeightedAverage
@@ -74,10 +75,10 @@ def __init__(
         assigner,
         use_delta_updates=True,
         straggler_handling_policy: StragglerPolicy = CutoffTimePolicy,
-        rounds_to_train=256,
+        rounds_to_train=constants.DEFAULT_ROUNDS_TO_TRAIN,
         single_col_cert_common_name=None,
         compression_pipeline=None,
-        db_store_rounds=1,
+        db_store_rounds=constants.DEFAULT_DB_STORE_ROUNDS,
         initial_tensor_dict=None,
         log_memory_usage=False,
         write_logs=False,
@@ -100,13 +101,13 @@ def __init__(
             assigner: Assigner object.
             straggler_handling_policy (optional): Straggler handling policy.
             rounds_to_train (int, optional): Number of rounds to train.
-                Defaults to 256.
+                Defaults to constants.DEFAULT_ROUNDS_TO_TRAIN.
             single_col_cert_common_name (str, optional): Common name for single
                 collaborator certificate. Defaults to None.
             compression_pipeline (optional): Compression pipeline. Defaults to
                 NoCompressionPipeline.
             db_store_rounds (int, optional): Rounds to store in TensorDB.
-                Defaults to 1.
+                Defaults to constants.DEFAULT_DB_STORE_ROUNDS.
             initial_tensor_dict (dict, optional): Initial tensor dictionary.
             callbacks: List of callbacks to be used during the experiment.
         """
@@ -120,15 +121,17 @@ def __init__(
                 "provide proper Public Key Infrastructure (PKI) security. "
                 "Please use this mode with caution."
             )
-        # FIXME: "" instead of None is for protobuf compatibility.
-        self.single_col_cert_common_name = single_col_cert_common_name or ""
+        # FIXME: using DEFAULT_CERT_COMMON_NAME for protobuf compatibility.
+        self.single_col_cert_common_name = (
+            single_col_cert_common_name or constants.DEFAULT_CERT_COMMON_NAME
+        )
 
         self.straggler_handling_policy = straggler_handling_policy()
 
         self.rounds_to_train = rounds_to_train
         self.assigner = assigner
         if self.assigner.is_task_group_evaluation():
-            self.rounds_to_train = 1
+            self.rounds_to_train = constants.EVALUATION_ROUNDS
             logger.info(f"For evaluation tasks setting rounds_to_train = {self.rounds_to_train}")
 
         self._end_of_round_check_done = [False] * rounds_to_train
@@ -143,7 +146,7 @@ def __init__(
 
         self.tensor_db = TensorDB()
         if persist_checkpoint:
-            persistent_db_path = persistent_db_path or "tensor.db"
+            persistent_db_path = persistent_db_path or constants.DEFAULT_PERSISTENT_DB_PATH
             logger.info(
                 "Persistent checkpoint is enabled, setting persistent db at path %s",
                 persistent_db_path,
@@ -590,7 +593,8 @@ def get_aggregated_tensor(
 
         tensor_key = TensorKey(tensor_name, self.uuid, round_number, report, tags)
         tensor_name, origin, round_number, report, tags = tensor_key
-
+        # TODO: This is a temporary fix. The tags should be updated in the
+        #  TensorDB.
         if "aggregated" in tags and "delta" in tags and round_number != 0:
             agg_tensor_key = TensorKey(tensor_name, origin, round_number, report, ("aggregated",))
         else:

openfl/component/aggregator/constants.py

@@ -0,0 +1,16 @@
+# Constants for Aggregator component
+
+# Default number of rounds to train during federated learning
+DEFAULT_ROUNDS_TO_TRAIN = 256
+
+# Number of rounds to train during evaluation mode
+EVALUATION_ROUNDS = 1
+
+# Default number of rounds to store in the TensorDB
+DEFAULT_DB_STORE_ROUNDS = 1
+
+# Default persistent database path if not provided
+DEFAULT_PERSISTENT_DB_PATH = "tensor.db"
+
+# Default certificate common name to satisfy protobuf compatibility when none is provided
+DEFAULT_CERT_COMMON_NAME = ""