flwr 1.18.0__py3-none-any.whl → 1.19.0__py3-none-any.whl

flwr/common/record/configrecord.py

@@ -15,12 +15,21 @@
 """ConfigRecord."""
 
 
+from __future__ import annotations
+
 from logging import WARN
-from typing import Optional, get_args
+from typing import cast, get_args
 
 from flwr.common.typing import ConfigRecordValues, ConfigScalar
 
+# pylint: disable=E0611
+from flwr.proto.recorddict_pb2 import ConfigRecord as ProtoConfigRecord
+from flwr.proto.recorddict_pb2 import ConfigRecordValue as ProtoConfigRecordValue
+
+# pylint: enable=E0611
+from ..inflatable import InflatableObject, add_header_to_object_body, get_object_body
 from ..logger import log
+from ..serde_utils import record_value_dict_from_proto, record_value_dict_to_proto
 from .typeddict import TypedDict
 
 
@@ -59,7 +68,7 @@ def _check_value(value: ConfigRecordValues) -> None:
         is_valid(value)
 
 
-class ConfigRecord(TypedDict[str, ConfigRecordValues]):
+class ConfigRecord(TypedDict[str, ConfigRecordValues], InflatableObject):
     """Config record.
 
     A :code:`ConfigRecord` is a Python dictionary designed to ensure that
@@ -111,7 +120,7 @@ class ConfigRecord(TypedDict[str, ConfigRecordValues]):
 
     def __init__(
         self,
-        config_dict: Optional[dict[str, ConfigRecordValues]] = None,
+        config_dict: dict[str, ConfigRecordValues] | None = None,
         keep_input: bool = True,
     ) -> None:
 
@@ -164,6 +173,52 @@ class ConfigRecord(TypedDict[str, ConfigRecordValues]):
 
         return num_bytes
 
+    def deflate(self) -> bytes:
+        """Deflate object."""
+        protos = record_value_dict_to_proto(
+            self,
+            [bool, int, float, str, bytes],
+            ProtoConfigRecordValue,
+        )
+        obj_body = ProtoConfigRecord(
+            items=[ProtoConfigRecord.Item(key=k, value=v) for k, v in protos.items()]
+        ).SerializeToString()
+        return add_header_to_object_body(object_body=obj_body, obj=self)
+
+    @classmethod
+    def inflate(
+        cls, object_content: bytes, children: dict[str, InflatableObject] | None = None
+    ) -> ConfigRecord:
+        """Inflate a ConfigRecord from bytes.
+
+        Parameters
+        ----------
+        object_content : bytes
+            The deflated object content of the ConfigRecord.
+
+        children : Optional[dict[str, InflatableObject]] (default: None)
+            Must be ``None``. ``ConfigRecord`` does not support child objects.
+            Providing any children will raise a ``ValueError``.
+
+        Returns
+        -------
+        ConfigRecord
+            The inflated ConfigRecord.
+        """
+        if children:
+            raise ValueError("`ConfigRecord` objects do not have children.")
+
+        obj_body = get_object_body(object_content, cls)
+        config_record_proto = ProtoConfigRecord.FromString(obj_body)
+        protos = {item.key: item.value for item in config_record_proto.items}
+        return ConfigRecord(
+            config_dict=cast(
+                dict[str, ConfigRecordValues],
+                record_value_dict_from_proto(protos),
+            ),
+            keep_input=False,
+        )
+
 
 class ConfigsRecord(ConfigRecord):
     """Deprecated class ``ConfigsRecord``, use ``ConfigRecord`` instead.
@@ -195,7 +250,7 @@ class ConfigsRecord(ConfigRecord):
 
     def __init__(
         self,
-        config_dict: Optional[dict[str, ConfigRecordValues]] = None,
+        config_dict: dict[str, ConfigRecordValues] | None = None,
         keep_input: bool = True,
     ):
         if not ConfigsRecord._warning_logged:

flwr/common/record/conversion_utils.py

@@ -17,7 +17,7 @@
 
 from ..logger import warn_deprecated_feature
 from ..typing import NDArray
-from .arrayrecord import Array
+from .array import Array
 
 WARN_DEPRECATED_MESSAGE = (
     "`array_from_numpy` is deprecated. Instead, use the `Array(ndarray)` class "

flwr/common/record/metricrecord.py

@@ -15,12 +15,21 @@
 """MetricRecord."""
 
 
+from __future__ import annotations
+
 from logging import WARN
-from typing import Optional, get_args
+from typing import cast, get_args
 
 from flwr.common.typing import MetricRecordValues, MetricScalar
 
+# pylint: disable=E0611
+from flwr.proto.recorddict_pb2 import MetricRecord as ProtoMetricRecord
+from flwr.proto.recorddict_pb2 import MetricRecordValue as ProtoMetricRecordValue
+
+# pylint: enable=E0611
+from ..inflatable import InflatableObject, add_header_to_object_body, get_object_body
 from ..logger import log
+from ..serde_utils import record_value_dict_from_proto, record_value_dict_to_proto
 from .typeddict import TypedDict
 
 
@@ -59,7 +68,7 @@ def _check_value(value: MetricRecordValues) -> None:
         is_valid(value)
 
 
-class MetricRecord(TypedDict[str, MetricRecordValues]):
+class MetricRecord(TypedDict[str, MetricRecordValues], InflatableObject):
     """Metric record.
 
     A :code:`MetricRecord` is a Python dictionary designed to ensure that
@@ -117,7 +126,7 @@ class MetricRecord(TypedDict[str, MetricRecordValues]):
 
     def __init__(
         self,
-        metric_dict: Optional[dict[str, MetricRecordValues]] = None,
+        metric_dict: dict[str, MetricRecordValues] | None = None,
         keep_input: bool = True,
     ) -> None:
         super().__init__(_check_key, _check_value)
@@ -143,6 +152,48 @@ class MetricRecord(TypedDict[str, MetricRecordValues]):
             num_bytes += len(k)
         return num_bytes
 
+    def deflate(self) -> bytes:
+        """Deflate object."""
+        protos = record_value_dict_to_proto(self, [float, int], ProtoMetricRecordValue)
+        obj_body = ProtoMetricRecord(
+            items=[ProtoMetricRecord.Item(key=k, value=v) for k, v in protos.items()]
+        ).SerializeToString()
+        return add_header_to_object_body(object_body=obj_body, obj=self)
+
+    @classmethod
+    def inflate(
+        cls, object_content: bytes, children: dict[str, InflatableObject] | None = None
+    ) -> MetricRecord:
+        """Inflate a MetricRecord from bytes.
+
+        Parameters
+        ----------
+        object_content : bytes
+            The deflated object content of the MetricRecord.
+
+        children : Optional[dict[str, InflatableObject]] (default: None)
+            Must be ``None``. ``MetricRecord`` does not support child objects.
+            Providing any children will raise a ``ValueError``.
+
+        Returns
+        -------
+        MetricRecord
+            The inflated MetricRecord.
+        """
+        if children:
+            raise ValueError("`MetricRecord` objects do not have children.")
+
+        obj_body = get_object_body(object_content, cls)
+        metric_record_proto = ProtoMetricRecord.FromString(obj_body)
+        protos = {item.key: item.value for item in metric_record_proto.items}
+        return cls(
+            metric_dict=cast(
+                dict[str, MetricRecordValues],
+                record_value_dict_from_proto(protos),
+            ),
+            keep_input=False,
+        )
+
 
 class MetricsRecord(MetricRecord):
     """Deprecated class ``MetricsRecord``, use ``MetricRecord`` instead.
@@ -174,7 +225,7 @@ class MetricsRecord(MetricRecord):
 
     def __init__(
         self,
-        metric_dict: Optional[dict[str, MetricRecordValues]] = None,
+        metric_dict: dict[str, MetricRecordValues] | None = None,
         keep_input: bool = True,
     ):
         if not MetricsRecord._warning_logged:

flwr/common/record/recorddict.py

@@ -17,10 +17,12 @@
 
 from __future__ import annotations
 
+import json
 from logging import WARN
 from textwrap import indent
 from typing import TypeVar, Union, cast
 
+from ..inflatable import InflatableObject, add_header_to_object_body, get_object_body
 from ..logger import log
 from .arrayrecord import ArrayRecord
 from .configrecord import ConfigRecord
@@ -97,7 +99,7 @@ class _SyncedDict(TypedDict[str, T]):
             )
 
 
-class RecordDict(TypedDict[str, RecordType]):
+class RecordDict(TypedDict[str, RecordType], InflatableObject):
     """RecordDict stores groups of arrays, metrics and configs.
 
     A :class:`RecordDict` is the unified mechanism by which arrays,
@@ -286,6 +288,72 @@ class RecordDict(TypedDict[str, RecordType]):
             )
         return self.config_records
 
+    @property
+    def children(self) -> dict[str, InflatableObject]:
+        """Return a dictionary of records with their Object IDs as keys."""
+        return {record.object_id: record for record in self.values()}
+
+    def deflate(self) -> bytes:
+        """Deflate the RecordDict."""
+        # record_name: record_object_id mapping
+        record_refs: dict[str, str] = {}
+
+        for record_name, record in self.items():
+            record_refs[record_name] = record.object_id
+
+        # Serialize references dict
+        object_body = json.dumps(record_refs).encode("utf-8")
+        return add_header_to_object_body(object_body=object_body, obj=self)
+
+    @classmethod
+    def inflate(
+        cls, object_content: bytes, children: dict[str, InflatableObject] | None = None
+    ) -> RecordDict:
+        """Inflate an RecordDict from bytes.
+
+        Parameters
+        ----------
+        object_content : bytes
+            The deflated object content of the RecordDict.
+        children : Optional[dict[str, InflatableObject]]  (default: None)
+            Dictionary of children InflatableObjects mapped to their Object IDs.
+            These children enable the full inflation of the RecordDict. Default is None.
+
+        Returns
+        -------
+        RecordDict
+            The inflated RecordDict.
+        """
+        if children is None:
+            children = {}
+
+        # Inflate mapping of record_names (keys in the RecordDict) to Record' object IDs
+        obj_body = get_object_body(object_content, cls)
+        record_refs: dict[str, str] = json.loads(obj_body.decode(encoding="utf-8"))
+
+        unique_records = set(record_refs.values())
+        children_obj_ids = set(children.keys())
+        if unique_records != children_obj_ids:
+            raise ValueError(
+                "Unexpected set of `children`. "
+                f"Expected {unique_records} but got {children_obj_ids}."
+            )
+
+        # Ensure children are one of the *Record objects exepecte in a RecordDict
+        if not all(
+            isinstance(ch, (ArrayRecord, ConfigRecord, MetricRecord))
+            for ch in children.values()
+        ):
+            raise ValueError(
+                "`Children` are expected to be of type `ArrayRecord`, "
+                "`ConfigRecord` or `MetricRecord`."
+            )
+
+        # Instantiate new RecordDict
+        return RecordDict(
+            {name: children[object_id] for name, object_id in record_refs.items()}  # type: ignore
+        )
+
 
 class RecordSet(RecordDict):
     """Deprecated class ``RecordSet``, use ``RecordDict`` instead.

flwr/common/recorddict_compat.py

@@ -111,12 +111,12 @@ def parameters_to_arrayrecord(parameters: Parameters, keep_input: bool) -> Array
         else:
             tensor = parameters.tensors.pop(0)
         ordered_dict[str(idx)] = Array(
-            data=tensor, dtype="", stype=tensor_type, shape=[]
+            data=tensor, dtype="", stype=tensor_type, shape=()
         )
 
     if num_arrays == 0:
         ordered_dict[EMPTY_TENSOR_KEY] = Array(
-            data=b"", dtype="", stype=tensor_type, shape=[]
+            data=b"", dtype="", stype=tensor_type, shape=()
         )
     return ArrayRecord(ordered_dict, keep_input=keep_input)
 

flwr/common/retry_invoker.py

@@ -25,10 +25,12 @@ from typing import Any, Callable, Optional, Union, cast
 
 import grpc
 
+from flwr.client.grpc_rere_client.grpc_adapter import GrpcAdapter
 from flwr.common.constant import MAX_RETRY_DELAY
 from flwr.common.logger import log
 from flwr.common.typing import RunNotRunningException
 from flwr.proto.clientappio_pb2_grpc import ClientAppIoStub
+from flwr.proto.fleet_pb2_grpc import FleetStub
 from flwr.proto.serverappio_pb2_grpc import ServerAppIoStub
 from flwr.proto.simulationio_pb2_grpc import SimulationIoStub
 
@@ -366,7 +368,9 @@ def _make_simple_grpc_retry_invoker() -> RetryInvoker:
 
 
 def _wrap_stub(
-    stub: Union[ServerAppIoStub, ClientAppIoStub, SimulationIoStub],
+    stub: Union[
+        ServerAppIoStub, ClientAppIoStub, SimulationIoStub, FleetStub, GrpcAdapter
+    ],
     retry_invoker: RetryInvoker,
 ) -> None:
     """Wrap a gRPC stub with a retry invoker."""