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

flwr/common/record/arrayrecord.py

@@ -12,38 +12,31 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 # ==============================================================================
-"""ArrayRecord and Array."""
+"""ArrayRecord."""
 
 
 from __future__ import annotations
 
 import gc
+import json
 import sys
 from collections import OrderedDict
-from dataclasses import dataclass
-from io import BytesIO
 from logging import WARN
 from typing import TYPE_CHECKING, Any, cast, overload
 
 import numpy as np
 
-from ..constant import GC_THRESHOLD, SType
+from ..constant import GC_THRESHOLD
+from ..inflatable import InflatableObject, add_header_to_object_body, get_object_body
 from ..logger import log
 from ..typing import NDArray
+from .array import Array
 from .typeddict import TypedDict
 
 if TYPE_CHECKING:
     import torch
 
 
-def _raise_array_init_error() -> None:
-    raise TypeError(
-        f"Invalid arguments for {Array.__qualname__}. Expected either a "
-        "PyTorch tensor, a NumPy ndarray, or explicit"
-        " dtype/shape/stype/data values."
-    )
-
-
 def _raise_array_record_init_error() -> None:
     raise TypeError(
         f"Invalid arguments for {ArrayRecord.__qualname__}. Expected either "
@@ -52,217 +45,6 @@ def _raise_array_record_init_error() -> None:
     )
 
 
-@dataclass
-class Array:
-    """Array type.
-
-    A dataclass containing serialized data from an array-like or tensor-like object
-    along with metadata about it. The class can be initialized in one of three ways:
-
-    1. By specifying explicit values for `dtype`, `shape`, `stype`, and `data`.
-    2. By providing a NumPy ndarray (via the `ndarray` argument).
-    3. By providing a PyTorch tensor (via the `torch_tensor` argument).
-
-    In scenarios (2)-(3), the `dtype`, `shape`, `stype`, and `data` are automatically
-    derived from the input. In scenario (1), these fields must be specified manually.
-
-    Parameters
-    ----------
-    dtype : Optional[str] (default: None)
-        A string representing the data type of the serialized object (e.g. `"float32"`).
-        Only required if you are not passing in a ndarray or a tensor.
-
-    shape : Optional[list[int]] (default: None)
-        A list representing the shape of the unserialized array-like object. Only
-        required if you are not passing in a ndarray or a tensor.
-
-    stype : Optional[str] (default: None)
-        A string indicating the serialization mechanism used to generate the bytes in
-        `data` from an array-like or tensor-like object. Only required if you are not
-        passing in a ndarray or a tensor.
-
-    data : Optional[bytes] (default: None)
-        A buffer of bytes containing the data. Only required if you are not passing in
-        a ndarray or a tensor.
-
-    ndarray : Optional[NDArray] (default: None)
-        A NumPy ndarray. If provided, the `dtype`, `shape`, `stype`, and `data`
-        fields are derived automatically from it.
-
-    torch_tensor : Optional[torch.Tensor] (default: None)
-        A PyTorch tensor. If provided, it will be **detached and moved to CPU**
-        before conversion, and the `dtype`, `shape`, `stype`, and `data` fields
-        will be derived automatically from it.
-
-    Examples
-    --------
-    Initializing by specifying all fields directly:
-
-    >>> arr1 = Array(
-    >>>     dtype="float32",
-    >>>     shape=[3, 3],
-    >>>     stype="numpy.ndarray",
-    >>>     data=b"serialized_data...",
-    >>> )
-
-    Initializing with a NumPy ndarray:
-
-    >>> import numpy as np
-    >>> arr2 = Array(np.random.randn(3, 3))
-
-    Initializing with a PyTorch tensor:
-
-    >>> import torch
-    >>> arr3 = Array(torch.randn(3, 3))
-    """
-
-    dtype: str
-    shape: list[int]
-    stype: str
-    data: bytes
-
-    @overload
-    def __init__(  # noqa: E704
-        self, dtype: str, shape: list[int], stype: str, data: bytes
-    ) -> None: ...
-
-    @overload
-    def __init__(self, ndarray: NDArray) -> None: ...  # noqa: E704
-
-    @overload
-    def __init__(self, torch_tensor: torch.Tensor) -> None: ...  # noqa: E704
-
-    def __init__(  # pylint: disable=too-many-arguments, too-many-locals
-        self,
-        *args: Any,
-        dtype: str | None = None,
-        shape: list[int] | None = None,
-        stype: str | None = None,
-        data: bytes | None = None,
-        ndarray: NDArray | None = None,
-        torch_tensor: torch.Tensor | None = None,
-    ) -> None:
-        # Determine the initialization method and validate input arguments.
-        # Support three initialization formats:
-        # 1. Array(dtype: str, shape: list[int], stype: str, data: bytes)
-        # 2. Array(ndarray: NDArray)
-        # 3. Array(torch_tensor: torch.Tensor)
-
-        # Initialize all arguments
-        # If more than 4 positional arguments are provided, raise an error.
-        if len(args) > 4:
-            _raise_array_init_error()
-        all_args = [None] * 4
-        for i, arg in enumerate(args):
-            all_args[i] = arg
-        init_method: str | None = None  # Track which init method is being used
-
-        # Try to assign a value to all_args[index] if it's not already set.
-        # If an initialization method is provided, update init_method.
-        def _try_set_arg(index: int, arg: Any, method: str) -> None:
-            # Skip if arg is None
-            if arg is None:
-                return
-            # Raise an error if all_args[index] is already set
-            if all_args[index] is not None:
-                _raise_array_init_error()
-            # Raise an error if a different initialization method is already set
-            nonlocal init_method
-            if init_method is not None and init_method != method:
-                _raise_array_init_error()
-            # Set init_method and all_args[index]
-            if init_method is None:
-                init_method = method
-            all_args[index] = arg
-
-        # Try to set keyword arguments in all_args
-        _try_set_arg(0, dtype, "direct")
-        _try_set_arg(1, shape, "direct")
-        _try_set_arg(2, stype, "direct")
-        _try_set_arg(3, data, "direct")
-        _try_set_arg(0, ndarray, "ndarray")
-        _try_set_arg(0, torch_tensor, "torch_tensor")
-
-        # Check if all arguments are correctly set
-        all_args = [arg for arg in all_args if arg is not None]
-
-        # Handle direct field initialization
-        if not init_method or init_method == "direct":
-            if (
-                len(all_args) == 4  # pylint: disable=too-many-boolean-expressions
-                and isinstance(all_args[0], str)
-                and isinstance(all_args[1], list)
-                and all(isinstance(i, int) for i in all_args[1])
-                and isinstance(all_args[2], str)
-                and isinstance(all_args[3], bytes)
-            ):
-                self.dtype, self.shape, self.stype, self.data = all_args
-                return
-
-        # Handle NumPy array
-        if not init_method or init_method == "ndarray":
-            if len(all_args) == 1 and isinstance(all_args[0], np.ndarray):
-                self.__dict__.update(self.from_numpy_ndarray(all_args[0]).__dict__)
-                return
-
-        # Handle PyTorch tensor
-        if not init_method or init_method == "torch_tensor":
-            if (
-                len(all_args) == 1
-                and "torch" in sys.modules
-                and isinstance(all_args[0], sys.modules["torch"].Tensor)
-            ):
-                self.__dict__.update(self.from_torch_tensor(all_args[0]).__dict__)
-                return
-
-        _raise_array_init_error()
-
-    @classmethod
-    def from_numpy_ndarray(cls, ndarray: NDArray) -> Array:
-        """Create Array from NumPy ndarray."""
-        assert isinstance(
-            ndarray, np.ndarray
-        ), f"Expected NumPy ndarray, got {type(ndarray)}"
-        buffer = BytesIO()
-        # WARNING: NEVER set allow_pickle to true.
-        # Reason: loading pickled data can execute arbitrary code
-        # Source: https://numpy.org/doc/stable/reference/generated/numpy.save.html
-        np.save(buffer, ndarray, allow_pickle=False)
-        data = buffer.getvalue()
-        return Array(
-            dtype=str(ndarray.dtype),
-            shape=list(ndarray.shape),
-            stype=SType.NUMPY,
-            data=data,
-        )
-
-    @classmethod
-    def from_torch_tensor(cls, tensor: torch.Tensor) -> Array:
-        """Create Array from PyTorch tensor."""
-        if not (torch := sys.modules.get("torch")):
-            raise RuntimeError(
-                f"PyTorch is required to use {cls.from_torch_tensor.__name__}"
-            )
-
-        assert isinstance(
-            tensor, torch.Tensor
-        ), f"Expected PyTorch Tensor, got {type(tensor)}"
-        return cls.from_numpy_ndarray(tensor.detach().cpu().numpy())
-
-    def numpy(self) -> NDArray:
-        """Return the array as a NumPy array."""
-        if self.stype != SType.NUMPY:
-            raise TypeError(
-                f"Unsupported serialization type for numpy conversion: '{self.stype}'"
-            )
-        bytes_io = BytesIO(self.data)
-        # WARNING: NEVER set allow_pickle to true.
-        # Reason: loading pickled data can execute arbitrary code
-        # Source: https://numpy.org/doc/stable/reference/generated/numpy.load.html
-        ndarray_deserialized = np.load(bytes_io, allow_pickle=False)
-        return cast(NDArray, ndarray_deserialized)
-
-
 def _check_key(key: str) -> None:
     """Check if key is of expected type."""
     if not isinstance(key, str):
@@ -276,7 +58,7 @@ def _check_value(value: Array) -> None:
         )
 
 
-class ArrayRecord(TypedDict[str, Array]):
+class ArrayRecord(TypedDict[str, Array], InflatableObject):
     """Array record.
 
     A typed dictionary (``str`` to :class:`Array`) that can store named arrays,
@@ -315,33 +97,33 @@ class ArrayRecord(TypedDict[str, Array]):
 
     Examples
     --------
-    Initializing an empty ArrayRecord:
+    Initializing an empty ArrayRecord::
 
-    >>> record = ArrayRecord()
+        record = ArrayRecord()
 
-    Initializing with a dictionary of :class:`Array`:
+    Initializing with a dictionary of :class:`Array`::
 
-    >>> arr = Array("float32", [5, 5], "numpy.ndarray", b"serialized_data...")
-    >>> record = ArrayRecord({"weight": arr})
+        arr = Array("float32", [5, 5], "numpy.ndarray", b"serialized_data...")
+        record = ArrayRecord({"weight": arr})
 
-    Initializing with a list of NumPy arrays:
+    Initializing with a list of NumPy arrays::
 
-    >>> import numpy as np
-    >>> arr1 = np.random.randn(3, 3)
-    >>> arr2 = np.random.randn(2, 2)
-    >>> record = ArrayRecord([arr1, arr2])
+        import numpy as np
+        arr1 = np.random.randn(3, 3)
+        arr2 = np.random.randn(2, 2)
+        record = ArrayRecord([arr1, arr2])
 
-    Initializing with a PyTorch model state_dict:
+    Initializing with a PyTorch model state_dict::
 
-    >>> import torch.nn as nn
-    >>> model = nn.Linear(10, 5)
-    >>> record = ArrayRecord(model.state_dict())
+        import torch.nn as nn
+        model = nn.Linear(10, 5)
+        record = ArrayRecord(model.state_dict())
 
-    Initializing with a TensorFlow model weights (a list of NumPy arrays):
+    Initializing with a TensorFlow model weights (a list of NumPy arrays)::
 
-    >>> import tensorflow as tf
-    >>> model = tf.keras.Sequential([tf.keras.layers.Dense(5, input_shape=(10,))])
-    >>> record = ArrayRecord(model.get_weights())
+        import tensorflow as tf
+        model = tf.keras.Sequential([tf.keras.layers.Dense(5, input_shape=(10,))])
+        record = ArrayRecord(model.get_weights())
     """
 
     @overload
@@ -470,7 +252,7 @@ class ArrayRecord(TypedDict[str, Array]):
         record = ArrayRecord()
         for k, v in array_dict.items():
             record[k] = Array(
-                dtype=v.dtype, shape=list(v.shape), stype=v.stype, data=v.data
+                dtype=v.dtype, shape=tuple(v.shape), stype=v.stype, data=v.data
             )
         if not keep_input:
             array_dict.clear()
@@ -585,6 +367,102 @@ class ArrayRecord(TypedDict[str, Array]):
 
         return num_bytes
 
+    @property
+    def children(self) -> dict[str, InflatableObject]:
+        """Return a dictionary of Arrays with their Object IDs as keys."""
+        return {arr.object_id: arr for arr in self.values()}
+
+    def deflate(self) -> bytes:
+        """Deflate the ArrayRecord."""
+        # array_name: array_object_id mapping
+        array_refs: dict[str, str] = {}
+
+        for array_name, array in self.items():
+            array_refs[array_name] = array.object_id
+
+        # Serialize references dict
+        object_body = json.dumps(array_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
+    ) -> ArrayRecord:
+        """Inflate an ArrayRecord from bytes.
+
+        Parameters
+        ----------
+        object_content : bytes
+            The deflated object content of the ArrayRecord.
+        children : Optional[dict[str, InflatableObject]] (default: None)
+            Dictionary of children InflatableObjects mapped to their Object IDs.
+            These children enable the full inflation of the ArrayRecord.
+
+        Returns
+        -------
+        ArrayRecord
+            The inflated ArrayRecord.
+        """
+        if children is None:
+            children = {}
+
+        # Inflate mapping of array_names (keys in the ArrayRecord) to Arrays' object IDs
+        obj_body = get_object_body(object_content, cls)
+        array_refs: dict[str, str] = json.loads(obj_body.decode(encoding="utf-8"))
+
+        unique_arrays = set(array_refs.values())
+        children_obj_ids = set(children.keys())
+        if unique_arrays != children_obj_ids:
+            raise ValueError(
+                "Unexpected set of `children`. "
+                f"Expected {unique_arrays} but got {children_obj_ids}."
+            )
+
+        # Ensure children are of type Array
+        if not all(isinstance(arr, Array) for arr in children.values()):
+            raise ValueError("`Children` are expected to be of type `Array`.")
+
+        # Instantiate new ArrayRecord
+        return ArrayRecord(
+            OrderedDict(
+                {name: children[object_id] for name, object_id in array_refs.items()}
+            )
+        )
+
+    @property
+    def object_id(self) -> str:
+        """Get object ID."""
+        ret = super().object_id
+        self.is_dirty = False  # Reset dirty flag
+        return ret
+
+    @property
+    def is_dirty(self) -> bool:
+        """Check if the object is dirty after the last deflation."""
+        if "_is_dirty" not in self.__dict__:
+            self.__dict__["_is_dirty"] = True
+
+        if not self.__dict__["_is_dirty"]:
+            if any(v.is_dirty for v in self.values()):
+                # If any Array is dirty, mark the record as dirty
+                self.__dict__["_is_dirty"] = True
+        return cast(bool, self.__dict__["_is_dirty"])
+
+    @is_dirty.setter
+    def is_dirty(self, value: bool) -> None:
+        """Set the dirty flag."""
+        self.__dict__["_is_dirty"] = value
+
+    def __setitem__(self, key: str, value: Array) -> None:
+        """Set item and mark the record as dirty."""
+        self.is_dirty = True  # Mark as dirty when setting an item
+        super().__setitem__(key, value)
+
+    def __delitem__(self, key: str) -> None:
+        """Delete item and mark the record as dirty."""
+        self.is_dirty = True  # Mark as dirty when deleting an item
+        super().__delitem__(key)
+
 
 class ParametersRecord(ArrayRecord):
     """Deprecated class ``ParametersRecord``, use ``ArrayRecord`` instead.

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,8 +68,8 @@ def _check_value(value: ConfigRecordValues) -> None:
         is_valid(value)
 
 
-class ConfigRecord(TypedDict[str, ConfigRecordValues]):
-    """Configs record.
+class ConfigRecord(TypedDict[str, ConfigRecordValues], InflatableObject):
+    """Config record.
 
     A :code:`ConfigRecord` is a Python dictionary designed to ensure that
     each key-value pair adheres to specified data types. A :code:`ConfigRecord`
@@ -90,18 +99,18 @@ class ConfigRecord(TypedDict[str, ConfigRecordValues]):
     encourage you to use a :code:`ArrayRecord` instead if these are of high
     dimensionality.
 
-    Let's see some examples of how to construct a :code:`ConfigRecord` from scratch:
+    Let's see some examples of how to construct a :code:`ConfigRecord` from scratch::
+
+        from flwr.common import ConfigRecord
 
-    >>> from flwr.common import ConfigRecord
-    >>>
-    >>> # A `ConfigRecord` is a specialized Python dictionary
-    >>> record = ConfigRecord({"lr": 0.1, "batch-size": 128})
-    >>> # You can add more content to an existing record
-    >>> record["compute-average"] = True
-    >>> # It also supports lists
-    >>> record["loss-fn-coefficients"] = [0.4, 0.25, 0.35]
-    >>> # And string values (among other types)
-    >>> record["path-to-S3"] = "s3://bucket_name/folder1/fileA.json"
+        # A `ConfigRecord` is a specialized Python dictionary
+        record = ConfigRecord({"lr": 0.1, "batch-size": 128})
+        # You can add more content to an existing record
+        record["compute-average"] = True
+        # It also supports lists
+        record["loss-fn-coefficients"] = [0.4, 0.25, 0.35]
+        # And string values (among other types)
+        record["path-to-S3"] = "s3://bucket_name/folder1/fileA.json"
 
     Just like the other types of records in a :code:`flwr.common.RecordDict`, types are
     enforced. If you need to add a custom data structure or object, we recommend to
@@ -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

@@ -1,4 +1,4 @@
-# Copyright 2024 Flower Labs GmbH. All Rights Reserved.
+# Copyright 2025 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.
@@ -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 "