Introduce callbacks API

openfl-workspace/workspace/plan/defaults/aggregator.yaml

@@ -1,4 +1,3 @@
 template : openfl.component.Aggregator
 settings :
     db_store_rounds : 2
-    write_logs      : true

openfl/callbacks/__init__.py

@@ -0,0 +1,6 @@
+# Copyright 2020-2024 Intel Corporation
+# SPDX-License-Identifier: Apache-2.0
+from openfl.callbacks.callback import Callback
+from openfl.callbacks.callback_list import CallbackList
+from openfl.callbacks.lambda_callback import LambdaCallback
+from openfl.callbacks.memory_profiler import MemoryProfiler

openfl/callbacks/callback.py

@@ -0,0 +1,57 @@
+# Copyright 2020-2024 Intel Corporation
+# SPDX-License-Identifier: Apache-2.0
+"""Callbacks API."""
+
+
+class Callback:
+    """Base class for callbacks.
+
+    Callbacks can be used to perform actions at different stages of the
+    Federated Learning process. To create a custom callback, subclass
+    `openfl.callbacks.Callback` and implement the necessary methods.
+
+    Callbacks can be triggered on the aggregator and collaborator side
+    for the following events:
+        * At the beginning of an experiment
+        * At the beginning of a round
+        * At the end of a round
+        * At the end of an experiment
+
+    Attributes:
+        params: Additional parameters saved for use within the callback.
+        tensor_db: The `TensorDB` instance of the respective participant.
+    """
+
+    def __init__(self):
+        self.params = None
+        self.tensor_db = None
+
+    def set_params(self, params):
+        self.params = params
+
+    def set_tensor_db(self, tensor_db):
+        self.tensor_db = tensor_db
+
+    def on_round_begin(self, round_num: int, logs=None):
+        """Callback function to be executed at the beginning of a round.
+
+        Subclasses need to implement actions to be taken here.
+        """
+
+    def on_round_end(self, round_num: int, logs=None):
+        """Callback function to be executed at the end of a round.
+
+        Subclasses need to implement actions to be taken here.
+        """
+
+    def on_experiment_begin(self, logs=None):
+        """Callback function to be executed at the beginning of an experiment.
+
+        Subclasses need to implement actions to be taken here.
+        """
+
+    def on_experiment_end(self, logs=None):
+        """Callback function to be executed at the end of an experiment.
+
+        Subclasses need to implement actions to be taken here.
+        """

openfl/callbacks/callback_list.py

@@ -0,0 +1,83 @@
+# Copyright 2020-2024 Intel Corporation
+# SPDX-License-Identifier: Apache-2.0
+from typing import List
+
+from openfl.callbacks.callback import Callback
+from openfl.callbacks.memory_profiler import MemoryProfiler
+
+
+class CallbackList(Callback):
+    """An ensemble of callbacks.
+
+    This class allows multiple callbacks to be used together, by sequentially
+    calling each callback's respective methods.
+
+    Attributes:
+        callbacks: A list of `openfl.callbacks.Callback` instances.
+        add_memory_profiler: If True, adds a `MemoryProfiler` callback to the list.
+        tensor_db: Optional `TensorDB` instance of the respective participant.
+            If provided, callbacks can access TensorDB for various actions.
+        params: Additional parameters saved for use within the callbacks.
+    """
+
+    def __init__(
+        self,
+        callbacks: List[Callback],
+        add_memory_profiler=False,
+        tensor_db=None,
+        **params,
+    ):
+        super().__init__()
+        self.callbacks = _flatten(callbacks) if callbacks else []
+
+        self._add_default_callbacks(add_memory_profiler)
+
+        self.set_tensor_db(tensor_db)
+        self.set_params(params)
+
+    def set_params(self, params):
+        self.params = params
+        if params:
+            for callback in self.callbacks:
+                callback.set_params(params)
+
+    def set_tensor_db(self, tensor_db):
+        self.tensor_db = tensor_db
+        if tensor_db:
+            for callback in self.callbacks:
+                callback.set_tensor_db(tensor_db)
+
+    def _add_default_callbacks(self, add_memory_profiler):
+        self._memory_profiler = None
+        for cb in self.callbacks:
+            if isinstance(cb, MemoryProfiler):
+                self._memory_profiler = cb
+
+        if add_memory_profiler and self._memory_profiler is None:
+            self._memory_profiler = MemoryProfiler()
+            self.callbacks.append(self._memory_profiler)
+
+    def on_round_begin(self, round_num: int, logs=None):
+        for callback in self.callbacks:
+            callback.on_round_begin(round_num, logs)
+
+    def on_round_end(self, round_num: int, logs=None):
+        for callback in self.callbacks:
+            callback.on_round_end(round_num, logs)
+
+    def on_experiment_begin(self, logs=None):
+        for callback in self.callbacks:
+            callback.on_experiment_begin(logs)
+
+    def on_experiment_end(self, logs=None):
+        for callback in self.callbacks:
+            callback.on_experiment_end(logs)
+
+
+def _flatten(l):
+    """Flatten a possibly-nested tree of lists."""
+    for elem in l:
+        if isinstance(elem, list):
+            yield from _flatten(elem)
+        else:
+            yield elem

openfl/callbacks/lambda_callback.py

@@ -0,0 +1,38 @@
+# Copyright 2020-2024 Intel Corporation
+# SPDX-License-Identifier: Apache-2.0
+from openfl.callbacks.callback import Callback
+
+
+class LambdaCallback(Callback):
+    """Custom on-the-fly callbacks.
+
+    This callback can be constructed with functions that will be called
+    at the appropriate time during the life-cycle of a Federated Learning experiment.
+    Certain callbacks may expect positional arguments, for example:
+
+    * on_round_begin: expects `round_num` as a positional argument.
+    * on_round_end: expects `round_num` as a positional argument.
+
+    Args:
+        on_round_begin: called at the beginning of every round.
+        on_round_end: called at the end of every round.
+        on_experiment_begin: called at the beginning of an experiment.
+        on_experiment_end: called at the end of an experiment.
+    """
+
+    def __init__(
+        self,
+        on_round_begin=None,
+        on_round_end=None,
+        on_experiment_begin=None,
+        on_experiment_end=None,
+    ):
+        super().__init__()
+        if on_round_begin is not None:
+            self.on_round_begin = on_round_begin
+        if on_round_end is not None:
+            self.on_round_end = on_round_end
+        if on_experiment_begin is not None:
+            self.on_experiment_begin = on_experiment_begin
+        if on_experiment_end is not None:
+            self.on_experiment_end = on_experiment_end

openfl/callbacks/memory_profiler.py

@@ -0,0 +1,64 @@
+# Copyright 2020-2024 Intel Corporation
+# SPDX-License-Identifier: Apache-2.0
+"""Memory Profiler callback."""
+
+import json
+import logging
+import os
+
+import psutil
+
+from openfl.callbacks.callback import Callback
+
+logger = logging.getLogger(__name__)
+
+
+class MemoryProfiler(Callback):
+    """Profile memory usage of the current process at the end of each round.
+
+    Attributes:
+        log_dir: If set, writes logs as lines of JSON.
+    """
+
+    def __init__(self, log_dir: str = "./logs/"):
+        super().__init__()
+        self.log_dir = None
+        if log_dir:
+            os.makedirs(log_dir, exist_ok=True)
+            self.log_dir = log_dir
+
+    def on_round_end(self, round_num: int, logs=None):
+        origin = self.params["origin"]
+
+        info = _get_memory_usage()
+        info["round_number"] = round_num
+        info["origin"] = origin
+
+        logger.info(f"Round {round_num}: Memory usage: {info}")
+        if self.log_dir:
+            with open(os.path.join(self.log_dir, f"{origin}_memory_usage.json"), "a") as f:
+                f.write(json.dumps(info) + "\n")
+
+
+def _get_memory_usage() -> dict:
+    process = psutil.Process(os.getpid())
+    virtual_memory = psutil.virtual_memory()
+    swap_memory = psutil.swap_memory()
+    info = {
+        "process_memory": round(process.memory_info().rss / (1024**2), 2),
+        "virtual_memory/total": round(virtual_memory.total / (1024**2), 2),
+        "virtual_memory/available": round(virtual_memory.available / (1024**2), 2),
+        "virtual_memory/percent": virtual_memory.percent,
+        "virtual_memory/used": round(virtual_memory.used / (1024**2), 2),
+        "virtual_memory/free": round(virtual_memory.free / (1024**2), 2),
+        "virtual_memory/active": round(virtual_memory.active / (1024**2), 2),
+        "virtual_memory/inactive": round(virtual_memory.inactive / (1024**2), 2),
+        "virtual_memory/buffers": round(virtual_memory.buffers / (1024**2), 2),
+        "virtual_memory/cached": round(virtual_memory.cached / (1024**2), 2),
+        "virtual_memory/shared": round(virtual_memory.shared / (1024**2), 2),
+        "swap_memory/total": round(swap_memory.total / (1024**2), 2),
+        "swap_memory/used": round(swap_memory.used / (1024**2), 2),
+        "swap_memory/free": round(swap_memory.free / (1024**2), 2),
+        "swap_memory/percent": swap_memory.percent,
+    }
+    return info