Add *Record ProtoBuf messages and corresponding serde functions.

src/proto/flwr/proto/recordset.proto

@@ -0,0 +1,36 @@
+// Copyright 2024 Flower Labs GmbH. All Rights Reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+// ==============================================================================
+
+syntax = "proto3";
+
+package flwr.proto;
+
+import "flwr/proto/task.proto";
+
+message Array {
+  string dtype = 1;
+  repeated int32 shape = 2;
+  string stype = 3;
+  bytes data = 4;
+}
+
+message ParametersRecord {
+  repeated string data_keys = 1;
+  repeated Array data_values = 2;
+}
+
+message MetricsRecord { map<string, Value> data = 1; }
+
+message ConfigsRecord { map<string, Value> data = 1; }

src/py/flwr/common/serde.py

@@ -15,8 +15,12 @@
 """ProtoBuf serialization and deserialization."""
 
 
-from typing import Any, Dict, List, MutableMapping, cast
+from typing import Any, Dict, List, MutableMapping, OrderedDict, cast
 
+from flwr.proto.recordset_pb2 import Array as ArrayProto
+from flwr.proto.recordset_pb2 import ConfigsRecord as ConfigsRecordProto
+from flwr.proto.recordset_pb2 import MetricsRecord as MetricsRecordProto
+from flwr.proto.recordset_pb2 import ParametersRecord as ParametersRecordProto
 from flwr.proto.task_pb2 import Value  # pylint: disable=E0611
 from flwr.proto.transport_pb2 import (  # pylint: disable=E0611
     ClientMessage,
@@ -29,6 +33,9 @@
 )
 
 from . import typing
+from .configsrecord import ConfigsRecord
+from .metricsrecord import MetricsRecord
+from .parametersrecord import Array, ParametersRecord
 
 #  === ServerMessage message ===
 
@@ -573,3 +580,77 @@ def named_values_from_proto(
 ) -> Dict[str, typing.Value]:
     """Deserialize named values from ProtoBuf."""
     return {name: value_from_proto(value) for name, value in named_values_proto.items()}
+
+
+# === Record messages ===
+
+
+def array_to_proto(array: Array) -> ArrayProto:
+    """Serialize Array to ProtoBuf."""
+    return ArrayProto(**vars(array))
+
+
+def array_from_proto(array_proto: ArrayProto) -> Array:
+    """Deserialize Array from ProtoBuf."""
+    return Array(
+        dtype=array_proto.dtype,
+        shape=list(array_proto.shape),
+        stype=array_proto.stype,
+        data=array_proto.data,
+    )
+
+
+def parameters_record_to_proto(record: ParametersRecord) -> ParametersRecordProto:
+    """Serialize ParametersRecord to ProtoBuf."""
+    return ParametersRecordProto(
+        data_keys=record.data.keys(),
+        data_values=map(array_to_proto, record.data.values()),
+    )
+
+
+def parameters_record_from_proto(
+    record_proto: ParametersRecordProto,
+) -> ParametersRecord:
+    """Deserialize ParametersRecord from ProtoBuf."""
+    return ParametersRecord(
+        array_dict=OrderedDict(
+            zip(record_proto.data_keys, map(array_from_proto, record_proto.data_values))
+        ),
+        keep_input=False,
+    )
+
+
+def metrics_record_to_proto(record: MetricsRecord) -> MetricsRecordProto:
+    """Serialize MetricsRecord to ProtoBuf."""
+    return MetricsRecordProto(
+        data=named_values_to_proto(cast(Dict[str, typing.Value], record.data))
+    )
+
+
+def metrics_record_from_proto(record_proto: MetricsRecordProto) -> MetricsRecord:
+    """Deserialize MetricsRecord from ProtoBuf."""
+    return MetricsRecord(
+        metrics_dict=cast(
+            Dict[str, typing.MetricsRecordValues],
+            named_values_from_proto(record_proto.data),
+        ),
+        keep_input=False,
+    )
+
+
+def configs_record_to_proto(record: ConfigsRecord) -> ConfigsRecordProto:
+    """Serialize ConfigsRecord to ProtoBuf."""
+    return ConfigsRecordProto(
+        data=named_values_to_proto(cast(Dict[str, typing.Value], record.data))
+    )
+
+
+def configs_record_from_proto(record_proto: ConfigsRecordProto) -> ConfigsRecord:
+    """Deserialize ConfigsRecord from ProtoBuf."""
+    return ConfigsRecord(
+        configs_dict=cast(
+            Dict[str, typing.ConfigsRecordValues],
+            named_values_from_proto(record_proto.data),
+        ),
+        keep_input=False,
+    )

src/py/flwr/common/serde_test.py

@@ -15,14 +15,29 @@
 """(De-)serialization tests."""
 
 
-from typing import Dict, Union, cast
+from typing import Dict, OrderedDict, Union, cast
 
 from flwr.common import typing
 from flwr.proto import transport_pb2 as pb2  # pylint: disable=E0611
-
+from flwr.proto.recordset_pb2 import Array as ArrayProto
+from flwr.proto.recordset_pb2 import ConfigsRecord as ConfigsRecordProto
+from flwr.proto.recordset_pb2 import MetricsRecord as MetricsRecordProto
+from flwr.proto.recordset_pb2 import ParametersRecord as ParametersRecordProto
+
+from .configsrecord import ConfigsRecord
+from .metricsrecord import MetricsRecord
+from .parametersrecord import Array, ParametersRecord
 from .serde import (
+    array_from_proto,
+    array_to_proto,
+    configs_record_from_proto,
+    configs_record_to_proto,
+    metrics_record_from_proto,
+    metrics_record_to_proto,
     named_values_from_proto,
     named_values_to_proto,
+    parameters_record_from_proto,
+    parameters_record_to_proto,
     scalar_from_proto,
     scalar_to_proto,
     status_from_proto,
@@ -157,3 +172,71 @@ def test_named_values_serialization_deserialization() -> None:
                 assert elm1 == elm2
         else:
             assert expected == actual
+
+
+def test_array_serialization_deserialization() -> None:
+    """Test serialization and deserialization of Array."""
+    # Prepare
+    original = Array(dtype="float", shape=[2, 2], stype="dense", data=b"1234")
+
+    # Execute
+    proto = array_to_proto(original)
+    deserialized = array_from_proto(proto)
+
+    # Assert
+    assert isinstance(proto, ArrayProto)
+    assert original == deserialized
+
+
+def test_parameters_record_serialization_deserialization() -> None:
+    """Test serialization and deserialization of ParametersRecord."""
+    # Prepare
+    original = ParametersRecord(
+        array_dict=OrderedDict(
+            [
+                ("k1", Array(dtype="float", shape=[2, 2], stype="dense", data=b"1234")),
+                ("k2", Array(dtype="int", shape=[3], stype="sparse", data=b"567")),
+            ]
+        ),
+        keep_input=False,
+    )
+
+    # Execute
+    proto = parameters_record_to_proto(original)
+    deserialized = parameters_record_from_proto(proto)
+
+    # Assert
+    assert isinstance(proto, ParametersRecordProto)
+    assert original.data == deserialized.data
+
+
+def test_metrics_record_serialization_deserialization() -> None:
+    """Test serialization and deserialization of MetricsRecord."""
+    # Prepare
+    original = MetricsRecord(
+        metrics_dict={"accuracy": 0.95, "loss": 0.1}, keep_input=False
+    )
+
+    # Execute
+    proto = metrics_record_to_proto(original)
+    deserialized = metrics_record_from_proto(proto)
+
+    # Assert
+    assert isinstance(proto, MetricsRecordProto)
+    assert original.data == deserialized.data
+
+
+def test_configs_record_serialization_deserialization() -> None:
+    """Test serialization and deserialization of ConfigsRecord."""
+    # Prepare
+    original = ConfigsRecord(
+        configs_dict={"learning_rate": 0.01, "batch_size": 32}, keep_input=False
+    )
+
+    # Execute
+    proto = configs_record_to_proto(original)
+    deserialized = configs_record_from_proto(proto)
+
+    # Assert
+    assert isinstance(proto, ConfigsRecordProto)
+    assert original.data == deserialized.data

src/py/flwr/proto/recordset_pb2.pyi

@@ -0,0 +1,104 @@
+"""
+@generated by mypy-protobuf.  Do not edit manually!
+isort:skip_file
+"""
+import builtins
+import flwr.proto.task_pb2
+import google.protobuf.descriptor
+import google.protobuf.internal.containers
+import google.protobuf.message
+import typing
+import typing_extensions
+
+DESCRIPTOR: google.protobuf.descriptor.FileDescriptor
+
+class Array(google.protobuf.message.Message):
+    DESCRIPTOR: google.protobuf.descriptor.Descriptor
+    DTYPE_FIELD_NUMBER: builtins.int
+    SHAPE_FIELD_NUMBER: builtins.int
+    STYPE_FIELD_NUMBER: builtins.int
+    DATA_FIELD_NUMBER: builtins.int
+    dtype: typing.Text
+    @property
+    def shape(self) -> google.protobuf.internal.containers.RepeatedScalarFieldContainer[builtins.int]: ...
+    stype: typing.Text
+    data: builtins.bytes
+    def __init__(self,
+        *,
+        dtype: typing.Text = ...,
+        shape: typing.Optional[typing.Iterable[builtins.int]] = ...,
+        stype: typing.Text = ...,
+        data: builtins.bytes = ...,
+        ) -> None: ...
+    def ClearField(self, field_name: typing_extensions.Literal["data",b"data","dtype",b"dtype","shape",b"shape","stype",b"stype"]) -> None: ...
+global___Array = Array
+
+class ParametersRecord(google.protobuf.message.Message):
+    DESCRIPTOR: google.protobuf.descriptor.Descriptor
+    DATA_KEYS_FIELD_NUMBER: builtins.int
+    DATA_VALUES_FIELD_NUMBER: builtins.int
+    @property
+    def data_keys(self) -> google.protobuf.internal.containers.RepeatedScalarFieldContainer[typing.Text]: ...
+    @property
+    def data_values(self) -> google.protobuf.internal.containers.RepeatedCompositeFieldContainer[global___Array]: ...
+    def __init__(self,
+        *,
+        data_keys: typing.Optional[typing.Iterable[typing.Text]] = ...,
+        data_values: typing.Optional[typing.Iterable[global___Array]] = ...,
+        ) -> None: ...
+    def ClearField(self, field_name: typing_extensions.Literal["data_keys",b"data_keys","data_values",b"data_values"]) -> None: ...
+global___ParametersRecord = ParametersRecord
+
+class MetricsRecord(google.protobuf.message.Message):
+    DESCRIPTOR: google.protobuf.descriptor.Descriptor
+    class DataEntry(google.protobuf.message.Message):
+        DESCRIPTOR: google.protobuf.descriptor.Descriptor
+        KEY_FIELD_NUMBER: builtins.int
+        VALUE_FIELD_NUMBER: builtins.int
+        key: typing.Text
+        @property
+        def value(self) -> flwr.proto.task_pb2.Value: ...
+        def __init__(self,
+            *,
+            key: typing.Text = ...,
+            value: typing.Optional[flwr.proto.task_pb2.Value] = ...,
+            ) -> None: ...
+        def HasField(self, field_name: typing_extensions.Literal["value",b"value"]) -> builtins.bool: ...
+        def ClearField(self, field_name: typing_extensions.Literal["key",b"key","value",b"value"]) -> None: ...
+
+    DATA_FIELD_NUMBER: builtins.int
+    @property
+    def data(self) -> google.protobuf.internal.containers.MessageMap[typing.Text, flwr.proto.task_pb2.Value]: ...
+    def __init__(self,
+        *,
+        data: typing.Optional[typing.Mapping[typing.Text, flwr.proto.task_pb2.Value]] = ...,
+        ) -> None: ...
+    def ClearField(self, field_name: typing_extensions.Literal["data",b"data"]) -> None: ...
+global___MetricsRecord = MetricsRecord
+
+class ConfigsRecord(google.protobuf.message.Message):
+    DESCRIPTOR: google.protobuf.descriptor.Descriptor
+    class DataEntry(google.protobuf.message.Message):
+        DESCRIPTOR: google.protobuf.descriptor.Descriptor
+        KEY_FIELD_NUMBER: builtins.int
+        VALUE_FIELD_NUMBER: builtins.int
+        key: typing.Text
+        @property
+        def value(self) -> flwr.proto.task_pb2.Value: ...
+        def __init__(self,
+            *,
+            key: typing.Text = ...,
+            value: typing.Optional[flwr.proto.task_pb2.Value] = ...,
+            ) -> None: ...
+        def HasField(self, field_name: typing_extensions.Literal["value",b"value"]) -> builtins.bool: ...
+        def ClearField(self, field_name: typing_extensions.Literal["key",b"key","value",b"value"]) -> None: ...
+
+    DATA_FIELD_NUMBER: builtins.int
+    @property
+    def data(self) -> google.protobuf.internal.containers.MessageMap[typing.Text, flwr.proto.task_pb2.Value]: ...
+    def __init__(self,
+        *,
+        data: typing.Optional[typing.Mapping[typing.Text, flwr.proto.task_pb2.Value]] = ...,
+        ) -> None: ...
+    def ClearField(self, field_name: typing_extensions.Literal["data",b"data"]) -> None: ...
+global___ConfigsRecord = ConfigsRecord

src/py/flwr/proto/recordset_pb2_grpc.py

@@ -0,0 +1,4 @@
+# Generated by the gRPC Python protocol compiler plugin. DO NOT EDIT!
+"""Client and server classes corresponding to protobuf-defined services."""
+import grpc
+

src/py/flwr/proto/recordset_pb2_grpc.pyi

@@ -0,0 +1,4 @@
+"""
+@generated by mypy-protobuf.  Do not edit manually!
+isort:skip_file
+"""