flwr-nightly 1.17.0.dev20250317__py3-none-any.whl → 1.17.0.dev20250319__py3-none-any.whl

flwr/common/constant.py

@@ -61,6 +61,7 @@ PING_CALL_TIMEOUT = 5
 PING_BASE_MULTIPLIER = 0.8
 PING_RANDOM_RANGE = (-0.1, 0.1)
 PING_MAX_INTERVAL = 1e300
+PING_PATIENCE = 2
 
 # IDs
 RUN_ID_NUM_BYTES = 8
@@ -120,6 +121,9 @@ TIMESTAMP_HEADER = "flwr-timestamp"
 TIMESTAMP_TOLERANCE = 10  # General tolerance for timestamp verification
 SYSTEM_TIME_TOLERANCE = 5  # Allowance for system time drift
 
+# Constants for ParametersRecord
+GC_THRESHOLD = 200_000_000  # 200 MB
+
 
 class MessageType:
     """Message type."""
@@ -163,6 +167,7 @@ class ErrorCode:
     CLIENT_APP_RAISED_EXCEPTION = 2
     MESSAGE_UNAVAILABLE = 3
     REPLY_MESSAGE_UNAVAILABLE = 4
+    NODE_UNAVAILABLE = 5
 
     def __new__(cls) -> ErrorCode:
         """Prevent instantiation."""

flwr/common/logger.py

@@ -250,9 +250,9 @@ def warn_deprecated_feature_with_example(
     log(
         WARN,
         """FEATURE UPDATE: %s
-        ------------------------------------------------------------
+            ------------------------------------------------------------
         %s
-        ------------------------------------------------------------
+            ------------------------------------------------------------
         """,
         example_message,
         code_example,

flwr/common/record/parametersrecord.py

@@ -17,22 +17,36 @@
 
 from __future__ import annotations
 
+import gc
+import sys
 from collections import OrderedDict
 from dataclasses import dataclass
 from io import BytesIO
-from typing import Any, cast, overload
+from typing import TYPE_CHECKING, Any, cast, overload
 
 import numpy as np
 
-from ..constant import SType
+from ..constant import GC_THRESHOLD, SType
 from ..typing import NDArray
 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 "
-        "NumPy ndarray, or explicit dtype/shape/stype/data values."
+        "PyTorch tensor, a NumPy ndarray, or explicit"
+        " dtype/shape/stype/data values."
+    )
+
+
+def _raise_parameters_record_init_error() -> None:
+    raise TypeError(
+        f"Invalid arguments for {ParametersRecord.__qualname__}. Expected either "
+        "a list of NumPy ndarrays, a PyTorch state_dict, or a dictionary of Arrays. "
+        "The `keep_input` argument is keyword-only."
     )
 
 
@@ -41,37 +55,43 @@ 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 two ways:
+    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 scenario (2), the `dtype`, `shape`, `stype`, and `data` are automatically
+    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.
+        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.
+        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.
+        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.
+        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:
@@ -87,6 +107,11 @@ class Array:
 
     >>> 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
@@ -102,6 +127,9 @@ class Array:
     @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,
@@ -110,11 +138,13 @@ class Array:
         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 two initialization formats:
+        # 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.
@@ -149,6 +179,7 @@ class Array:
         _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]
@@ -172,6 +203,16 @@ class Array:
                 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
@@ -193,6 +234,19 @@ class Array:
             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:
@@ -223,103 +277,293 @@ def _check_value(value: Array) -> None:
 class ParametersRecord(TypedDict[str, Array]):
     r"""Parameters record.
 
-    A dataclass storing named Arrays in order. This means that it holds entries as an
-    OrderedDict[str, Array]. ParametersRecord objects can be viewed as an equivalent to
-    PyTorch's state_dict, but holding serialised tensors instead. A
-    :code:`ParametersRecord`  is one of the types of records that a
-    `flwr.common.RecordSet <flwr.common.RecordSet.html#recordset>`_ supports and
-    can therefore be used to construct :code:`common.Message` objects.
+    A typed dictionary (``str`` to :class:`Array`) that can store named parameters
+    as serialized tensors. Internally, this behaves similarly to an
+    ``OrderedDict[str, Array]``. A ``ParametersRecord`` can be viewed as an
+    equivalent to PyTorch's ``state_dict``, but it holds arrays in serialized form.
+
+    This object is one of the record types supported by :class:`RecordSet` and can
+    therefore be stored in the ``content`` of a :class:`Message` or the ``state``
+    of a :class:`Context`.
+
+    This class can be instantiated in multiple ways:
+
+    1. By providing nothing (empty container).
+    2. By providing a dictionary of :class:`Array` (via the ``array_dict`` argument).
+    3. By providing a list of NumPy ``ndarray`` (via the ``numpy_ndarrays`` argument).
+    4. By providing a PyTorch ``state_dict`` (via the ``torch_state_dict`` argument).
 
     Parameters
     ----------
-    array_dict : Optional[OrderedDict[str, Array]]
-        A dictionary that stores serialized array-like or tensor-like objects.
-    keep_input : bool (default: False)
-        A boolean indicating whether parameters should be deleted from the input
-        dictionary immediately after adding them to the record. If False, the
-        dictionary passed to `set_parameters()` will be empty once exiting from that
-        function. This is the desired behaviour when working with very large
-        models/tensors/arrays. However, if you plan to continue working with your
-        parameters after adding it to the record, set this flag to True. When set
-        to True, the data is duplicated in memory.
+    array_dict : Optional[OrderedDict[str, Array]] (default: None)
+        An existing dictionary containing named :class:`Array` instances. If
+        provided, these entries will be used directly to populate the record.
+    numpy_ndarrays : Optional[list[NDArray]] (default: None)
+        A list of NumPy arrays. Each array will be automatically converted
+        into an :class:`Array` and stored in this record with generated keys.
+    torch_state_dict : Optional[OrderedDict[str, torch.Tensor]] (default: None)
+        A PyTorch ``state_dict`` (``str`` keys to ``torch.Tensor`` values). Each
+        tensor will be converted into an :class:`Array` and stored in this record.
+    keep_input : bool (default: True)
+        If ``False``, entries from the input are removed after being added to
+        this record to free up memory. If ``True``, the input remains unchanged.
+        Regardless of this value, no duplicate memory is used if the input is a
+        dictionary of :class:`Array`, i.e., ``array_dict``.
 
     Examples
     --------
-    The usage of :code:`ParametersRecord` is envisioned for storing data arrays (e.g.
-    parameters of a machine learning model). These first need to be serialized into
-    a :code:`flwr.common.Array` data structure.
+    Initializing an empty ParametersRecord:
+
+    >>> p_record = ParametersRecord()
+
+    Initializing with a dictionary of :class:`Array`:
+
+    >>> arr = Array("float32", [5, 5], "numpy.ndarray", b"serialized_data...")
+    >>> p_record = ParametersRecord({"weight": arr})
 
-    Let's see some examples:
+    Initializing with a list of NumPy arrays:
 
     >>> import numpy as np
-    >>> from flwr.common import ParametersRecord
-    >>>
-    >>> # Let's create a simple NumPy array
-    >>> arr_np = np.random.randn(3, 3)
-    >>>
-    >>> # If we print it
-    >>> array([[-1.84242409, -1.01539537, -0.46528405],
-    >>>      [ 0.32991896,  0.55540414,  0.44085534],
-    >>>      [-0.10758364,  1.97619858, -0.37120501]])
-    >>>
-    >>> # Let's create an Array out of it
-    >>> arr = Array(arr_np)
-    >>>
-    >>> # If we print it you'll see (note the binary data)
-    >>> Array(dtype='float64', shape=[3,3], stype='numpy.ndarray', data=b'@\x99\x18...')
-    >>>
-    >>> # Adding it to a ParametersRecord:
-    >>> p_record = ParametersRecord({"my_array": arr})
-
-    Now that the NumPy array is embedded into a :code:`ParametersRecord` it could be
-    sent if added as part of a :code:`common.Message` or it could be saved as a
-    persistent state of a :code:`ClientApp` via its context. Regardless of the usecase,
-    we will sooner or later want to recover the array in its original NumPy
-    representation. For the example above, where the array was serialized using the
-    built-in utility function, deserialization can be done as follows:
-
-    >>> # Use the Array's built-in method
-    >>> arr_np_d = arr.numpy()
-    >>>
-    >>> # If printed, it will show the exact same data as above:
-    >>> array([[-1.84242409, -1.01539537, -0.46528405],
-    >>>      [ 0.32991896,  0.55540414,  0.44085534],
-    >>>      [-0.10758364,  1.97619858, -0.37120501]])
-
-    If you need finer control on how your arrays are serialized and deserialized, you
-    can construct :code:`Array` objects directly like this:
-
-    >>> from flwr.common import Array
-    >>> # Serialize your array and construct Array object
-    >>> arr = Array(
-    >>>         data=ndarray.tobytes(),
-    >>>         dtype=str(ndarray.dtype),
-    >>>         stype="",  # Could be used in a deserialization function
-    >>>         shape=list(ndarray.shape),
-    >>>       )
-    >>>
-    >>> # Then you can deserialize it like this
-    >>> arr_np_d = np.frombuffer(
-    >>>             buffer=array.data,
-    >>>             dtype=array.dtype,
-    >>>            ).reshape(array.shape)
-
-    Note that different arrays (e.g. from PyTorch, Tensorflow) might require different
-    serialization mechanism. Howerver, they often support a conversion to NumPy,
-    therefore allowing to use the same or similar steps as in the example above.
+    >>> arr1 = np.random.randn(3, 3)
+    >>> arr2 = np.random.randn(2, 2)
+    >>> p_record = ParametersRecord([arr1, arr2])
+
+    Initializing with a PyTorch model state_dict:
+
+    >>> import torch.nn as nn
+    >>> model = nn.Linear(10, 5)
+    >>> p_record = ParametersRecord(model.state_dict())
+
+    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,))])
+    >>> p_record = ParametersRecord(model.get_weights())
     """
 
-    def __init__(
+    @overload
+    def __init__(self) -> None: ...  # noqa: E704
+
+    @overload
+    def __init__(  # noqa: E704
+        self, array_dict: OrderedDict[str, Array], *, keep_input: bool = True
+    ) -> None: ...
+
+    @overload
+    def __init__(  # noqa: E704
+        self, numpy_ndarrays: list[NDArray], *, keep_input: bool = True
+    ) -> None: ...
+
+    @overload
+    def __init__(  # noqa: E704
+        self,
+        torch_state_dict: OrderedDict[str, torch.Tensor],
+        *,
+        keep_input: bool = True,
+    ) -> None: ...
+
+    def __init__(  # pylint: disable=too-many-arguments
         self,
+        *args: Any,
+        numpy_ndarrays: list[NDArray] | None = None,
+        torch_state_dict: OrderedDict[str, torch.Tensor] | None = None,
         array_dict: OrderedDict[str, Array] | None = None,
-        keep_input: bool = False,
+        keep_input: bool = True,
     ) -> None:
         super().__init__(_check_key, _check_value)
-        if array_dict:
-            for k in list(array_dict.keys()):
-                self[k] = array_dict[k]
-                if not keep_input:
-                    del array_dict[k]
+
+        # Determine the initialization method and validates input arguments.
+        # Support the following initialization formats:
+        # 1. cls(array_dict: OrderedDict[str, Array], keep_input: bool)
+        # 2. cls(numpy_ndarrays: list[NDArray], keep_input: bool)
+        # 3. cls(torch_state_dict: dict[str, torch.Tensor], keep_input: bool)
+
+        # Init the argument
+        if len(args) > 1:
+            _raise_parameters_record_init_error()
+        arg = args[0] if args else None
+        init_method: str | None = None  # Track which init method is being used
+
+        # Try to assign a value to arg if it's not already set.
+        # If an initialization method is provided, update init_method.
+        def _try_set_arg(_arg: Any, method: str) -> None:
+            # Skip if _arg is None
+            if _arg is None:
+                return
+            nonlocal arg, init_method
+            # Raise an error if arg is already set
+            if arg is not None:
+                _raise_parameters_record_init_error()
+            # Raise an error if a different initialization method is already set
+            if init_method is not None:
+                _raise_parameters_record_init_error()
+            # Set init_method and arg
+            if init_method is None:
+                init_method = method
+            arg = _arg
+
+        # Try to set keyword arguments
+        _try_set_arg(array_dict, "array_dict")
+        _try_set_arg(numpy_ndarrays, "numpy_ndarrays")
+        _try_set_arg(torch_state_dict, "state_dict")
+
+        # If no arguments are provided, return and keep self empty
+        if arg is None:
+            return
+
+        # Handle dictionary of Arrays
+        if not init_method or init_method == "array_dict":
+            # Type check the input
+            if (
+                isinstance(arg, dict)
+                and all(isinstance(k, str) for k in arg.keys())
+                and all(isinstance(v, Array) for v in arg.values())
+            ):
+                array_dict = cast(OrderedDict[str, Array], arg)
+                converted = self.from_array_dict(array_dict, keep_input=keep_input)
+                self.__dict__.update(converted.__dict__)
+                return
+
+        # Handle NumPy ndarrays
+        if not init_method or init_method == "numpy_ndarrays":
+            # Type check the input
+            # pylint: disable-next=not-an-iterable
+            if isinstance(arg, list) and all(isinstance(v, np.ndarray) for v in arg):
+                numpy_ndarrays = cast(list[NDArray], arg)
+                converted = self.from_numpy_ndarrays(
+                    numpy_ndarrays, keep_input=keep_input
+                )
+                self.__dict__.update(converted.__dict__)
+                return
+
+        # Handle PyTorch state_dict
+        if not init_method or init_method == "state_dict":
+            # Type check the input
+            if (
+                (torch := sys.modules.get("torch")) is not None
+                and isinstance(arg, dict)
+                and all(isinstance(k, str) for k in arg.keys())
+                and all(isinstance(v, torch.Tensor) for v in arg.values())
+            ):
+                torch_state_dict = cast(
+                    OrderedDict[str, torch.Tensor], arg  # type: ignore
+                )
+                converted = self.from_torch_state_dict(
+                    torch_state_dict, keep_input=keep_input
+                )
+                self.__dict__.update(converted.__dict__)
+                return
+
+        _raise_parameters_record_init_error()
+
+    @classmethod
+    def from_array_dict(
+        cls,
+        array_dict: OrderedDict[str, Array],
+        *,
+        keep_input: bool = True,
+    ) -> ParametersRecord:
+        """Create ParametersRecord from a dictionary of :class:`Array`."""
+        record = ParametersRecord()
+        for k, v in array_dict.items():
+            record[k] = Array(
+                dtype=v.dtype, shape=list(v.shape), stype=v.stype, data=v.data
+            )
+        if not keep_input:
+            array_dict.clear()
+        return record
+
+    @classmethod
+    def from_numpy_ndarrays(
+        cls,
+        ndarrays: list[NDArray],
+        *,
+        keep_input: bool = True,
+    ) -> ParametersRecord:
+        """Create ParametersRecord from a list of NumPy ``ndarray``."""
+        record = ParametersRecord()
+        total_serialized_bytes = 0
+
+        for i in range(len(ndarrays)):  # pylint: disable=C0200
+            record[str(i)] = Array.from_numpy_ndarray(ndarrays[i])
+
+            if not keep_input:
+                # Remove the reference
+                ndarrays[i] = None  # type: ignore
+                total_serialized_bytes += len(record[str(i)].data)
+
+                # If total serialized data exceeds the threshold, trigger GC
+                if total_serialized_bytes > GC_THRESHOLD:
+                    total_serialized_bytes = 0
+                    gc.collect()
+
+        if not keep_input:
+            # Clear the entire list to remove all references and force GC
+            ndarrays.clear()
+            gc.collect()
+        return record
+
+    @classmethod
+    def from_torch_state_dict(
+        cls,
+        state_dict: OrderedDict[str, torch.Tensor],
+        *,
+        keep_input: bool = True,
+    ) -> ParametersRecord:
+        """Create ParametersRecord from PyTorch ``state_dict``."""
+        if "torch" not in sys.modules:
+            raise RuntimeError(
+                f"PyTorch is required to use {cls.from_torch_state_dict.__name__}"
+            )
+
+        record = ParametersRecord()
+
+        for k in list(state_dict.keys()):
+            v = state_dict[k] if keep_input else state_dict.pop(k)
+            record[k] = Array.from_numpy_ndarray(v.detach().cpu().numpy())
+
+        return record
+
+    def to_numpy_ndarrays(self, *, keep_input: bool = True) -> list[NDArray]:
+        """Return the ParametersRecord as a list of NumPy ``ndarray``."""
+        if keep_input:
+            return [v.numpy() for v in self.values()]
+
+        # Clear the record and return the list of NumPy arrays
+        ret: list[NDArray] = []
+        total_serialized_bytes = 0
+        for k in list(self.keys()):
+            arr = self.pop(k)
+            ret.append(arr.numpy())
+            total_serialized_bytes += len(arr.data)
+            del arr
+
+            # If total serialized data exceeds the threshold, trigger GC
+            if total_serialized_bytes > GC_THRESHOLD:
+                total_serialized_bytes = 0
+                gc.collect()
+
+        if not keep_input:
+            # Force GC
+            gc.collect()
+        return ret
+
+    def to_torch_state_dict(
+        self, *, keep_input: bool = True
+    ) -> OrderedDict[str, torch.Tensor]:
+        """Return the ParametersRecord as a PyTorch ``state_dict``."""
+        if not (torch := sys.modules.get("torch")):
+            raise RuntimeError(
+                f"PyTorch is required to use {self.to_torch_state_dict.__name__}"
+            )
+
+        state_dict = OrderedDict()
+
+        for k in list(self.keys()):
+            arr = self[k] if keep_input else self.pop(k)
+            state_dict[k] = torch.from_numpy(arr.numpy())
+
+        return state_dict
 
     def count_bytes(self) -> int:
         """Return number of Bytes stored in this object.

flwr/server/__init__.py

@@ -21,7 +21,8 @@ from .app import start_server as start_server
 from .client_manager import ClientManager as ClientManager
 from .client_manager import SimpleClientManager as SimpleClientManager
 from .compat import LegacyContext as LegacyContext
-from .driver import Driver as Driver
+from .grid import Driver as Driver
+from .grid import Grid as Grid
 from .history import History as History
 from .server import Server as Server
 from .server_app import ServerApp as ServerApp
@@ -31,6 +32,7 @@ from .serverapp_components import ServerAppComponents as ServerAppComponents
 __all__ = [
     "ClientManager",
     "Driver",
+    "Grid",
     "History",
     "LegacyContext",
     "Server",

flwr/server/app.py

@@ -79,13 +79,13 @@ from .history import History
 from .server import Server, init_defaults, run_fl
 from .server_config import ServerConfig
 from .strategy import Strategy
-from .superlink.driver.serverappio_grpc import run_serverappio_api_grpc
 from .superlink.ffs.ffs_factory import FfsFactory
 from .superlink.fleet.grpc_adapter.grpc_adapter_servicer import GrpcAdapterServicer
 from .superlink.fleet.grpc_bidi.grpc_server import start_grpc_server
 from .superlink.fleet.grpc_rere.fleet_servicer import FleetServicer
 from .superlink.fleet.grpc_rere.server_interceptor import AuthenticateServerInterceptor
 from .superlink.linkstate import LinkStateFactory
+from .superlink.serverappio.serverappio_grpc import run_serverappio_api_grpc
 from .superlink.simulation.simulationio_grpc import run_simulationio_api_grpc
 
 DATABASE = ":flwr-in-memory-state:"

flwr/server/compat/__init__.py

@@ -15,10 +15,10 @@
 """Flower ServerApp compatibility package."""
 
 
-from .app import start_driver as start_driver
+from .app import start_grid as start_grid
 from .legacy_context import LegacyContext as LegacyContext
 
 __all__ = [
     "LegacyContext",
-    "start_driver",
+    "start_grid",
 ]