Introduce callbacks API (securefederatedai#1195)

* Get rid of kwargs Signed-off-by: Shah, Karan <[email protected]> * Use module-level logger Signed-off-by: Shah, Karan <[email protected]> * Reduce keras verbosity Signed-off-by: Shah, Karan <[email protected]> * Remove all log_metric and log_memory_usage traces; add callback hooks Signed-off-by: Shah, Karan <[email protected]> * Add `openfl.callbacks` module Signed-off-by: Shah, Karan <[email protected]> * Include round_num for task callbacks Signed-off-by: Shah, Karan <[email protected]> * Add tensordb to callbacks Signed-off-by: Shah, Karan <[email protected]> * No round_num on task callbacks Signed-off-by: Shah, Karan <[email protected]> * Remove task boundary callbacks Signed-off-by: Shah, Karan <[email protected]> * Remove tb/model_ckpt. Add memory_profiler Signed-off-by: Shah, Karan <[email protected]> * Restore psutil and tbX Signed-off-by: Shah, Karan <[email protected]> * Format code Signed-off-by: Shah, Karan <[email protected]> * Define default callbacks Signed-off-by: Shah, Karan <[email protected]> * Add write_logs for bwd compat Signed-off-by: Shah, Karan <[email protected]> * Add log_metric_callback for bwd compat Signed-off-by: Shah, Karan <[email protected]> * Migrate to module-level logger for collaborator Signed-off-by: Shah, Karan <[email protected]> * Review comments Signed-off-by: Shah, Karan <[email protected]> * Add metric_writer Signed-off-by: Shah, Karan <[email protected]> * Add collaborator side metric logging Signed-off-by: Shah, Karan <[email protected]> * Make log dirs on exp begin Signed-off-by: Shah, Karan <[email protected]> * Do not print use_tls Signed-off-by: Shah, Karan <[email protected]> * Assume reportable metric to be a scalar Signed-off-by: Shah, Karan <[email protected]> * Add aggregator side callbacks Signed-off-by: Shah, Karan <[email protected]> * do_task test returns mock dict Signed-off-by: Shah, Karan <[email protected]> * Consistency changes Signed-off-by: Shah, Karan <[email protected]> * Add documentation hooks Signed-off-by: Shah, Karan <[email protected]> * Update docstring Signed-off-by: Shah, Karan <[email protected]> * Update docs hook Signed-off-by: Shah, Karan <[email protected]> * Remove all traces of log_metric_callback and write_metric Signed-off-by: Shah, Karan <[email protected]> * Do on_round_begin if not time_to_quit Signed-off-by: Shah, Karan <[email protected]> --------- Signed-off-by: Shah, Karan <[email protected]>
refai06 · Dec 23, 2024 · c280f10 · c280f10
1 parent 4e021d2
commit c280f10
Show file tree

Hide file tree

Showing 37 changed files with 510 additions and 482 deletions.
diff --git a/docs/openfl.callbacks.rst b/docs/openfl.callbacks.rst
@@ -0,0 +1,16 @@
+``openfl.callbacks`` module
+===========================
+
+.. currentmodule:: openfl.callbacks
+
+.. automodule:: openfl.callbacks
+
+.. autosummary::
+   :toctree: _autosummary
+   :recursive:
+
+   Callback
+   CallbackList
+   LambdaCallback
+   MetricWriter
+   MemoryProfiler
diff --git a/docs/openfl.rst b/docs/openfl.rst
@@ -1,6 +1,6 @@
 .. currentmodule:: openfl
 
-Public API: ``openfl`` package
+API Reference: ``openfl``
 ===========================
 
 Subpackages
@@ -10,6 +10,7 @@ Subpackages
     :maxdepth: 1
 
     openfl.component
+    openfl.callbacks
     openfl.cryptography
     openfl.experimental
     openfl.databases

diff --git a/...orkspace/experimental/workflow/AggregatorBasedWorkflow/101_torch_cnn_mnist/plan/plan.yaml b/...orkspace/experimental/workflow/AggregatorBasedWorkflow/101_torch_cnn_mnist/plan/plan.yaml
@@ -6,8 +6,6 @@ aggregator :
   template : openfl.experimental.workflow.component.Aggregator
   settings :
     rounds_to_train     : 1
-    log_metric_callback :
-      template : src.utils.write_metric
 
 
 collaborator :

diff --git a/...-workspace/experimental/workflow/AggregatorBasedWorkflow/101_torch_cnn_mnist/src/utils.py b/...-workspace/experimental/workflow/AggregatorBasedWorkflow/101_torch_cnn_mnist/src/utils.py
diff --git a/...ce/experimental/workflow/AggregatorBasedWorkflow/102_aggregator_validation/plan/plan.yaml b/...ce/experimental/workflow/AggregatorBasedWorkflow/102_aggregator_validation/plan/plan.yaml
@@ -6,8 +6,6 @@ aggregator :
   template : openfl.experimental.workflow.component.Aggregator
   settings :
     rounds_to_train     : 1
-    log_metric_callback :
-      template : src.utils.write_metric
 
 
 collaborator :

diff --git a/...pace/experimental/workflow/AggregatorBasedWorkflow/102_aggregator_validation/src/utils.py b/...pace/experimental/workflow/AggregatorBasedWorkflow/102_aggregator_validation/src/utils.py
diff --git a/...fl-workspace/experimental/workflow/AggregatorBasedWorkflow/104_keras_mnist/plan/plan.yaml b/...fl-workspace/experimental/workflow/AggregatorBasedWorkflow/104_keras_mnist/plan/plan.yaml
@@ -6,8 +6,6 @@ aggregator :
   template : openfl.experimental.workflow.component.Aggregator
   settings :
     rounds_to_train     : 1
-    log_metric_callback :
-      template : src.utils.write_metric
 
 
 collaborator :

diff --git a/openfl-workspace/experimental/workflow/AggregatorBasedWorkflow/104_keras_mnist/src/utils.py b/openfl-workspace/experimental/workflow/AggregatorBasedWorkflow/104_keras_mnist/src/utils.py
diff --git a/...rimental/workflow/AggregatorBasedWorkflow/301_torch_cnn_mnist_watermarking/plan/plan.yaml b/...rimental/workflow/AggregatorBasedWorkflow/301_torch_cnn_mnist_watermarking/plan/plan.yaml
@@ -6,8 +6,6 @@ aggregator :
   template : openfl.experimental.workflow.component.Aggregator
   settings :
     rounds_to_train     : 1
-    log_metric_callback :
-      template : src.utils.write_metric
 
 
 collaborator :

diff --git a/...perimental/workflow/AggregatorBasedWorkflow/301_torch_cnn_mnist_watermarking/src/utils.py b/...perimental/workflow/AggregatorBasedWorkflow/301_torch_cnn_mnist_watermarking/src/utils.py
diff --git a/...orkflow/AggregatorBasedWorkflow/501_pytorch_tinyimagenet_transfer_learning/plan/plan.yaml b/...orkflow/AggregatorBasedWorkflow/501_pytorch_tinyimagenet_transfer_learning/plan/plan.yaml
@@ -6,8 +6,6 @@ aggregator :
   template : openfl.experimental.workflow.component.Aggregator
   settings :
     rounds_to_train     : 1
-    log_metric_callback :
-      template : src.utils.write_metric
 
 
 collaborator :

diff --git a/.../workflow/AggregatorBasedWorkflow/501_pytorch_tinyimagenet_transfer_learning/src/utils.py b/.../workflow/AggregatorBasedWorkflow/501_pytorch_tinyimagenet_transfer_learning/src/utils.py
diff --git a/openfl-workspace/experimental/workflow/AggregatorBasedWorkflow/vertical_fl/plan/plan.yaml b/openfl-workspace/experimental/workflow/AggregatorBasedWorkflow/vertical_fl/plan/plan.yaml
@@ -6,8 +6,6 @@ aggregator :
   template : openfl.experimental.workflow.component.aggregator.Aggregator
   settings :
     rounds_to_train     : 1
-    log_metric_callback :
-      template : src.utils.write_metric
 
 
 collaborator :

diff --git a/openfl-workspace/experimental/workflow/AggregatorBasedWorkflow/vertical_fl/src/utils.py b/openfl-workspace/experimental/workflow/AggregatorBasedWorkflow/vertical_fl/src/utils.py
diff --git a/...kspace/experimental/workflow/AggregatorBasedWorkflow/vertical_fl_two_party/plan/plan.yaml b/...kspace/experimental/workflow/AggregatorBasedWorkflow/vertical_fl_two_party/plan/plan.yaml
@@ -6,8 +6,6 @@ aggregator :
   template : openfl.experimental.workflow.component.aggregator.Aggregator
   settings :
     rounds_to_train     : 10
-    log_metric_callback :
-      template : src.utils.write_metric
 
 
 collaborator :

diff --git a/...orkspace/experimental/workflow/AggregatorBasedWorkflow/vertical_fl_two_party/src/utils.py b/...orkspace/experimental/workflow/AggregatorBasedWorkflow/vertical_fl_two_party/src/utils.py
diff --git a/openfl-workspace/torch_cnn_mnist_eden_compression/plan/plan.yaml b/openfl-workspace/torch_cnn_mnist_eden_compression/plan/plan.yaml
@@ -9,8 +9,6 @@ aggregator :
     best_state_path     : save/torch_cnn_mnist_best.pbuf
     last_state_path     : save/torch_cnn_mnist_last.pbuf
     rounds_to_train     : 10
-    log_metric_callback :
-      template : src.mnist_utils.write_metric
 
 
 collaborator :

diff --git a/openfl-workspace/torch_cnn_mnist_eden_compression/src/mnist_utils.py b/openfl-workspace/torch_cnn_mnist_eden_compression/src/mnist_utils.py
@@ -6,27 +6,11 @@
 from logging import getLogger
 
 import numpy as np
-from torch.utils.tensorboard import SummaryWriter
 from torchvision import datasets
 from torchvision import transforms
 
 logger = getLogger(__name__)
 
-writer = None
-
-
-def get_writer():
-    """Create global writer object."""
-    global writer
-    if not writer:
-        writer = SummaryWriter('./logs/cnn_mnist', flush_secs=5)
-
-
-def write_metric(node_name, task_name, metric_name, metric, round_number):
-    """Write metric callback."""
-    get_writer()
-    writer.add_scalar(f'{node_name}/{task_name}/{metric_name}', metric, round_number)
-
 
 def one_hot(labels, classes):
     """

diff --git a/openfl-workspace/torch_cnn_mnist_fed_eval/plan/plan.yaml b/openfl-workspace/torch_cnn_mnist_fed_eval/plan/plan.yaml
@@ -8,8 +8,6 @@ aggregator :
     init_state_path     : save/torch_cnn_mnist_init.pbuf
     best_state_path     : save/torch_cnn_mnist_best.pbuf
     last_state_path     : save/torch_cnn_mnist_last.pbuf
-    log_metric_callback :
-      template : src.mnist_utils.write_metric
 
 collaborator :
   defaults : plan/defaults/collaborator.yaml

diff --git a/openfl-workspace/torch_cnn_mnist_fed_eval/src/mnist_utils.py b/openfl-workspace/torch_cnn_mnist_fed_eval/src/mnist_utils.py
@@ -6,27 +6,11 @@
 from logging import getLogger
 
 import numpy as np
-from torch.utils.tensorboard import SummaryWriter
 from torchvision import datasets
 from torchvision import transforms
 
 logger = getLogger(__name__)
 
-writer = None
-
-
-def get_writer():
-    """Create global writer object."""
-    global writer
-    if not writer:
-        writer = SummaryWriter('./logs/cnn_mnist', flush_secs=5)
-
-
-def write_metric(node_name, task_name, metric_name, metric, round_number):
-    """Write metric callback."""
-    get_writer()
-    writer.add_scalar(f'{node_name}/{task_name}/{metric_name}', metric, round_number)
-
 
 def one_hot(labels, classes):
     """

diff --git a/openfl-workspace/torch_cnn_mnist_straggler_check/plan/plan.yaml b/openfl-workspace/torch_cnn_mnist_straggler_check/plan/plan.yaml
@@ -9,8 +9,6 @@ aggregator :
     best_state_path     : save/torch_cnn_mnist_best.pbuf
     last_state_path     : save/torch_cnn_mnist_last.pbuf
     rounds_to_train     : 6
-    log_metric_callback :
-      template : src.mnist_utils.write_metric
 
 
 collaborator :

diff --git a/openfl-workspace/torch_cnn_mnist_straggler_check/src/mnist_utils.py b/openfl-workspace/torch_cnn_mnist_straggler_check/src/mnist_utils.py
@@ -6,27 +6,11 @@
 from logging import getLogger
 
 import numpy as np
-from torch.utils.tensorboard import SummaryWriter
 from torchvision import datasets
 from torchvision import transforms
 
 logger = getLogger(__name__)
 
-writer = None
-
-
-def get_writer():
-    """Create global writer object."""
-    global writer
-    if not writer:
-        writer = SummaryWriter('./logs/cnn_mnist', flush_secs=5)
-
-
-def write_metric(node_name, task_name, metric_name, metric, round_number):
-    """Write metric callback."""
-    get_writer()
-    writer.add_scalar(f'{node_name}/{task_name}/{metric_name}', metric, round_number)
-
 
 def one_hot(labels, classes):
     """

diff --git a/openfl-workspace/torch_llm_horovod/plan/plan.yaml b/openfl-workspace/torch_llm_horovod/plan/plan.yaml
@@ -9,8 +9,6 @@ aggregator :
     best_state_path     : save/torch_llm_best.pbuf
     last_state_path     : save/torch_llm_last.pbuf
     rounds_to_train     : 5
-    log_metric_callback :
-      template : src.emotion_utils.write_metric
 
 
 collaborator :

diff --git a/openfl-workspace/torch_llm_horovod/src/emotion_utils.py b/openfl-workspace/torch_llm_horovod/src/emotion_utils.py
@@ -6,26 +6,10 @@
 from logging import getLogger
 
 from datasets import Dataset, load_dataset
-from torch.utils.tensorboard import SummaryWriter
 from transformers import AutoTokenizer, DataCollatorWithPadding
 
 logger = getLogger(__name__)
 
-writer = None
-
-
-def get_writer():
-    """Create global writer object."""
-    global writer
-    if not writer:
-        writer = SummaryWriter("./logs/llm", flush_secs=5)
-
-
-def write_metric(node_name, task_name, metric_name, metric, round_number):
-    """Write metric callback."""
-    get_writer()
-    writer.add_scalar(f"{node_name}/{task_name}/{metric_name}", metric, round_number)
-
 
 def get_emotion_dataset(tokenizer):
     dataset = load_dataset("dair-ai/emotion", cache_dir="dataset", revision="9ce6303")

diff --git a/openfl-workspace/workspace/plan/defaults/aggregator.yaml b/openfl-workspace/workspace/plan/defaults/aggregator.yaml
@@ -1,4 +1,3 @@
 template : openfl.component.Aggregator
 settings :
     db_store_rounds : 2
-    write_logs      : true