Instrument metrics reporting to Redis

.github/workflows/integration_tests.yaml

@@ -1,4 +1,4 @@
-name: tests
+name: integration-tests
 
 on:
   push:
@@ -33,7 +33,7 @@ on:
       - '**.md'
 
 jobs:
-  python-tests:
+  integration-tests:
     runs-on: ubuntu-latest
     steps:
       - uses: actions/[email protected]

.github/workflows/unit_tests.yaml

@@ -1,4 +1,4 @@
-name: tests
+name: unit-tests
 
 on:
   push:
@@ -33,7 +33,7 @@ on:
       - '**.md'
 
 jobs:
-  python-tests:
+  unit-tests:
     runs-on: ubuntu-latest
     steps:
       - uses: actions/[email protected]

florist/api/monitoring/metrics.py

@@ -0,0 +1,57 @@
+"""Classes for the instrumentation of metrics reporting from clients and servers."""
+import json
+from logging import DEBUG
+from typing import Any, Dict, Optional
+
+import redis
+from fl4health.reporting.metrics import DateTimeEncoder, MetricsReporter
+from flwr.common.logger import log
+
+
+class RedisMetricsReporter(MetricsReporter):  # type: ignore
+    """Save the metrics to a Redis instance while it records them."""
+
+    def __init__(
+        self,
+        redis_connection: redis.client.Redis,
+        run_id: Optional[str] = None,
+    ):
+        """
+        Init an instance of RedisMetricsReporter.
+
+        :param redis_connection: (redis.client.Redis) the connection object to a Redis. Should be the output
+            of redis.Redis(host=host, port=port)
+        :param run_id: (Optional[str]) the identifier for the run which these metrics are from.
+            It will be used as the name of the object in Redis. Optional, default is a random UUID.
+        """
+        super().__init__(run_id)
+        self.redis_connection = redis_connection
+
+    def add_to_metrics(self, data: Dict[str, Any]) -> None:
+        """
+        Add a dictionary of data into the main metrics dictionary.
+
+        At the end, dumps the current state of the metrics to Redis.
+
+        :param data: (Dict[str, Any]) Data to be added to the metrics dictionary via .update().
+        """
+        super().add_to_metrics(data)
+        self.dump()
+
+    def add_to_metrics_at_round(self, fl_round: int, data: Dict[str, Any]) -> None:
+        """
+        Add a dictionary of data into the metrics dictionary for a specific FL round.
+
+        At the end, dumps the current state of the metrics to Redis.
+
+        :param fl_round: (int) the FL round these metrics are from.
+        :param data: (Dict[str, Any]) Data to be added to the round's metrics dictionary via .update().
+        """
+        super().add_to_metrics_at_round(fl_round, data)
+        self.dump()
+
+    def dump(self) -> None:
+        """Dump the current metrics to Redis under the run_id name."""
+        encoded_metrics = json.dumps(self.metrics, cls=DateTimeEncoder)
+        log(DEBUG, f"Dumping metrics to redis at key '{self.run_id}': {encoded_metrics}")
+        self.redis_connection.set(self.run_id, encoded_metrics)

florist/tests/api/monitoring/test_metrics.py

@@ -0,0 +1,60 @@
+import datetime
+import json
+from unittest.mock import Mock
+
+from fl4health.reporting.metrics import DateTimeEncoder
+from freezegun import freeze_time
+
+from florist.api.monitoring.metrics import RedisMetricsReporter
+
+
+@freeze_time("2012-12-11 10:09:08")
+def test_add_to_metrics() -> None:
+    mock_redis_connection = Mock()
+    test_run_id = "123"
+    test_data = {"test": "data", "date": datetime.datetime.now()}
+
+    redis_metric_reporter = RedisMetricsReporter(mock_redis_connection, test_run_id)
+    redis_metric_reporter.add_to_metrics(test_data)
+
+    mock_redis_connection.set.assert_called_once_with(test_run_id, json.dumps(test_data, cls=DateTimeEncoder))
+
+
+@freeze_time("2012-12-11 10:09:08")
+def test_add_to_metrics_at_round() -> None:
+    mock_redis_connection = Mock()
+    test_run_id = "123"
+    test_data = {"test": "data", "date": datetime.datetime.now()}
+    test_round = 2
+
+    redis_metric_reporter = RedisMetricsReporter(mock_redis_connection, test_run_id)
+    redis_metric_reporter.add_to_metrics_at_round(test_round, test_data)
+
+    expected_data = {
+        "rounds": {
+            str(test_round): test_data,
+        }
+    }
+    mock_redis_connection.set.assert_called_once_with(test_run_id, json.dumps(expected_data, cls=DateTimeEncoder))
+
+
+@freeze_time("2012-12-11 10:09:08")
+def test_dump() -> None:
+    mock_redis_connection = Mock()
+    test_run_id = "123"
+    test_data = {"test": "data", "date": datetime.datetime.now()}
+    test_round = 2
+
+    redis_metric_reporter = RedisMetricsReporter(mock_redis_connection, test_run_id)
+    redis_metric_reporter.add_to_metrics(test_data)
+    redis_metric_reporter.add_to_metrics_at_round(test_round, test_data)
+    redis_metric_reporter.dump()
+
+    expected_data = {
+        **test_data,
+        "rounds": {
+            str(test_round): test_data,
+        },
+    }
+    assert mock_redis_connection.set.call_args_list[2][0][0] == test_run_id
+    assert mock_redis_connection.set.call_args_list[2][0][1] == json.dumps(expected_data, cls=DateTimeEncoder)

pyproject.toml

@@ -34,6 +34,7 @@ mypy = "^1.7.0"
 ruff = "^0.2.0"
 pip-audit = "^2.7.1"
 nbqa = {extras = ["toolchain"], version = "^1.7.1"}
+freezegun = "^1.4.0"
 
 [tool.poetry.group.docs]
 optional = true