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

flwr/common/inflatable_utils.py

@@ -15,10 +15,14 @@
 """InflatableObject utilities."""
 
 import concurrent.futures
+import os
 import random
 import threading
 import time
-from typing import Callable, Optional
+from collections.abc import Iterable, Iterator
+from typing import Callable, Optional, TypeVar
+
+from flwr.proto.message_pb2 import ObjectTree  # pylint: disable=E0611
 
 from .constant import (
     HEAD_BODY_DIVIDER,
@@ -30,6 +34,7 @@ from .constant import (
     PULL_MAX_TIME,
     PULL_MAX_TRIES_PER_OBJECT,
 )
+from .exit import add_exit_handler
 from .inflatable import (
     InflatableObject,
     UnexpectedObjectContentError,
@@ -37,12 +42,15 @@ from .inflatable import (
     get_object_head_values_from_object_content,
     get_object_id,
     is_valid_sha256_hash,
+    iterate_object_tree,
 )
 from .message import Message
 from .record import Array, ArrayRecord, ConfigRecord, MetricRecord, RecordDict
+from .record.arraychunk import ArrayChunk
 
 # Helper registry that maps names of classes to their type
 inflatable_class_registry: dict[str, type[InflatableObject]] = {
+    ArrayChunk.__qualname__: ArrayChunk,
     Array.__qualname__: Array,
     ArrayRecord.__qualname__: ArrayRecord,
     ConfigRecord.__qualname__: ConfigRecord,
@@ -51,6 +59,36 @@ inflatable_class_registry: dict[str, type[InflatableObject]] = {
     RecordDict.__qualname__: RecordDict,
 }
 
+T = TypeVar("T", bound=InflatableObject)
+
+
+# Allow thread pool executors to be shut down gracefully
+_thread_pool_executors: set[concurrent.futures.ThreadPoolExecutor] = set()
+_lock = threading.Lock()
+
+
+def _shutdown_thread_pool_executors() -> None:
+    """Shutdown all thread pool executors gracefully."""
+    with _lock:
+        for executor in _thread_pool_executors:
+            executor.shutdown(wait=False, cancel_futures=True)
+        _thread_pool_executors.clear()
+
+
+def _track_executor(executor: concurrent.futures.ThreadPoolExecutor) -> None:
+    """Track a thread pool executor for graceful shutdown."""
+    with _lock:
+        _thread_pool_executors.add(executor)
+
+
+def _untrack_executor(executor: concurrent.futures.ThreadPoolExecutor) -> None:
+    """Untrack a thread pool executor."""
+    with _lock:
+        _thread_pool_executors.discard(executor)
+
+
+add_exit_handler(_shutdown_thread_pool_executors)
+
 
 class ObjectUnavailableError(Exception):
     """Exception raised when an object has been pre-registered but is not yet
@@ -67,6 +105,13 @@ class ObjectIdNotPreregisteredError(Exception):
         super().__init__(f"Object with ID '{object_id}' could not be found.")
 
 
+def get_num_workers(max_concurrent: int) -> int:
+    """Get number of workers based on the number of CPU cores and the maximum
+    allowed."""
+    num_cores = os.cpu_count() or 1
+    return min(max_concurrent, num_cores)
+
+
 def push_objects(
     objects: dict[str, InflatableObject],
     push_object_fn: Callable[[str, bytes], None],
@@ -95,27 +140,73 @@ def push_objects(
     max_concurrent_pushes : int (default: MAX_CONCURRENT_PUSHES)
         The maximum number of concurrent pushes to perform.
     """
-    if object_ids_to_push is not None:
-        # Filter objects to push only those with IDs in the set
-        objects = {k: v for k, v in objects.items() if k in object_ids_to_push}
-
     lock = threading.Lock()
 
-    def push(obj_id: str) -> None:
+    def iter_dict_items() -> Iterator[tuple[str, bytes]]:
+        """Iterate over the dictionary items."""
+        for obj_id in list(objects.keys()):
+            # Skip the object if no need to push it
+            if object_ids_to_push is not None and obj_id not in object_ids_to_push:
+                continue
+
+            # Deflate the object content
+            object_content = objects[obj_id].deflate()
+            if not keep_objects:
+                with lock:
+                    del objects[obj_id]
+
+            yield obj_id, object_content
+
+    push_object_contents_from_iterable(
+        iter_dict_items(),
+        push_object_fn,
+        max_concurrent_pushes=max_concurrent_pushes,
+    )
+
+
+def push_object_contents_from_iterable(
+    object_contents: Iterable[tuple[str, bytes]],
+    push_object_fn: Callable[[str, bytes], None],
+    *,
+    max_concurrent_pushes: int = MAX_CONCURRENT_PUSHES,
+) -> None:
+    """Push multiple object contents to the servicer.
+
+    Parameters
+    ----------
+    object_contents : Iterable[tuple[str, bytes]]
+        An iterable of `(object_id, object_content)` pairs.
+        `object_id` is the object ID, and `object_content` is the object content.
+    push_object_fn : Callable[[str, bytes], None]
+        A function that takes an object ID and its content as bytes, and pushes
+        it to the servicer. This function should raise `ObjectIdNotPreregisteredError`
+        if the object ID is not pre-registered.
+    max_concurrent_pushes : int (default: MAX_CONCURRENT_PUSHES)
+        The maximum number of concurrent pushes to perform.
+    """
+
+    def push(args: tuple[str, bytes]) -> None:
         """Push a single object."""
-        object_content = objects[obj_id].deflate()
-        if not keep_objects:
-            with lock:
-                del objects[obj_id]
-        push_object_fn(obj_id, object_content)
+        obj_id, obj_content = args
+        # Push the object using the provided function
+        push_object_fn(obj_id, obj_content)
+
+    # Push all object contents concurrently
+    num_workers = get_num_workers(max_concurrent_pushes)
+    with concurrent.futures.ThreadPoolExecutor(max_workers=num_workers) as executor:
+        # Ensure that the thread pool executors are tracked for graceful shutdown
+        _track_executor(executor)
 
-    with concurrent.futures.ThreadPoolExecutor(
-        max_workers=max_concurrent_pushes
-    ) as executor:
-        list(executor.map(push, list(objects.keys())))
+        # Submit push tasks for each object content
+        executor.map(push, object_contents)  # Non-blocking map
 
+        # The context manager will block until all submitted tasks have completed
 
-def pull_objects(  # pylint: disable=too-many-arguments
+    # Remove the executor from the list of tracked executors
+    _untrack_executor(executor)
+
+
+def pull_objects(  # pylint: disable=too-many-arguments,too-many-locals
     object_ids: list[str],
     pull_object_fn: Callable[[str], bytes],
     *,
@@ -207,16 +298,20 @@ def pull_objects(  # pylint: disable=too-many-arguments
                 return
 
     # Submit all pull tasks concurrently
-    with concurrent.futures.ThreadPoolExecutor(
-        max_workers=max_concurrent_pulls
-    ) as executor:
-        futures = {
-            executor.submit(pull_with_retries, obj_id): obj_id for obj_id in object_ids
-        }
+    num_workers = get_num_workers(max_concurrent_pulls)
+    with concurrent.futures.ThreadPoolExecutor(max_workers=num_workers) as executor:
+        # Ensure that the thread pool executors are tracked for graceful shutdown
+        _track_executor(executor)
+
+        # Submit pull tasks for each object ID
+        executor.map(pull_with_retries, object_ids)  # Non-blocking map
+
+        # The context manager will block until all submitted tasks have completed
 
-        # Wait for completion
-        concurrent.futures.wait(futures)
+    # Remove the executor from the list of tracked executors
+    _untrack_executor(executor)
 
+    # If an error occurred during pulling, raise it
     if err_to_raise is not None:
         raise err_to_raise
 
@@ -339,3 +434,75 @@ def validate_object_content(content: bytes) -> None:
         raise UnexpectedObjectContentError(
             object_id=get_object_id(content), reason=str(err)
         ) from err
+
+
+def pull_and_inflate_object_from_tree(  # pylint: disable=R0913
+    object_tree: ObjectTree,
+    pull_object_fn: Callable[[str], bytes],
+    confirm_object_received_fn: Callable[[str], None],
+    *,
+    return_type: type[T] = InflatableObject,  # type: ignore
+    max_concurrent_pulls: int = MAX_CONCURRENT_PULLS,
+    max_time: Optional[float] = PULL_MAX_TIME,
+    max_tries_per_object: Optional[int] = PULL_MAX_TRIES_PER_OBJECT,
+    initial_backoff: float = PULL_INITIAL_BACKOFF,
+    backoff_cap: float = PULL_BACKOFF_CAP,
+) -> T:
+    """Pull and inflate the head object from the provided object tree.
+
+    Parameters
+    ----------
+    object_tree : ObjectTree
+        The object tree containing the object ID and its descendants.
+    pull_object_fn : Callable[[str], bytes]
+        A function that takes an object ID and returns the object content as bytes.
+    confirm_object_received_fn : Callable[[str], None]
+        A function to confirm that the object has been received.
+    return_type : type[T] (default: InflatableObject)
+        The type of the object to return. Must be a subclass of `InflatableObject`.
+    max_concurrent_pulls : int (default: MAX_CONCURRENT_PULLS)
+        The maximum number of concurrent pulls to perform.
+    max_time : Optional[float] (default: PULL_MAX_TIME)
+        The maximum time to wait for all pulls to complete. If `None`, waits
+        indefinitely.
+    max_tries_per_object : Optional[int] (default: PULL_MAX_TRIES_PER_OBJECT)
+        The maximum number of attempts to pull each object. If `None`, pulls
+        indefinitely until the object is available.
+    initial_backoff : float (default: PULL_INITIAL_BACKOFF)
+        The initial backoff time in seconds for retrying pulls after an
+        `ObjectUnavailableError`.
+    backoff_cap : float (default: PULL_BACKOFF_CAP)
+        The maximum backoff time in seconds. Backoff times will not exceed this value.
+
+    Returns
+    -------
+    T
+        An instance of the specified return type containing the inflated object.
+    """
+    # Pull the main object and all its descendants
+    pulled_object_contents = pull_objects(
+        [tree.object_id for tree in iterate_object_tree(object_tree)],
+        pull_object_fn,
+        max_concurrent_pulls=max_concurrent_pulls,
+        max_time=max_time,
+        max_tries_per_object=max_tries_per_object,
+        initial_backoff=initial_backoff,
+        backoff_cap=backoff_cap,
+    )
+
+    # Confirm that all objects were pulled
+    confirm_object_received_fn(object_tree.object_id)
+
+    # Inflate the main object
+    inflated_object = inflate_object_from_contents(
+        object_tree.object_id, pulled_object_contents, keep_object_contents=False
+    )
+
+    # Check if the inflated object is of the expected type
+    if not isinstance(inflated_object, return_type):
+        raise TypeError(
+            f"Expected object of type {return_type.__name__}, "
+            f"but got {type(inflated_object).__name__}."
+        )
+
+    return inflated_object

flwr/common/logger.py

@@ -132,13 +132,13 @@ if log_level := os.getenv("FLWR_LOG_LEVEL"):
     log_level = log_level.upper()
     try:
         is_debug = log_level == "DEBUG"
+        update_console_handler(level=log_level, timestamps=is_debug, colored=True)
         if is_debug:
             log(
                 WARN,
                 "DEBUG logs enabled. Do not use this in production, as it may expose "
                 "sensitive details.",
             )
-        update_console_handler(level=log_level, timestamps=is_debug, colored=True)
     except Exception:  # pylint: disable=broad-exception-caught
         # Alert user but don't raise exception
         log(

flwr/common/record/array.py

@@ -17,6 +17,7 @@
 
 from __future__ import annotations
 
+import json
 import sys
 from dataclasses import dataclass
 from io import BytesIO
@@ -24,11 +25,15 @@ from typing import TYPE_CHECKING, Any, cast, overload
 
 import numpy as np
 
-from flwr.proto.recorddict_pb2 import Array as ArrayProto  # pylint: disable=E0611
-
-from ..constant import SType
-from ..inflatable import InflatableObject, add_header_to_object_body, get_object_body
+from ..constant import MAX_ARRAY_CHUNK_SIZE, SType
+from ..inflatable import (
+    InflatableObject,
+    add_header_to_object_body,
+    get_object_body,
+    get_object_children_ids_from_object_content,
+)
 from ..typing import NDArray
+from .arraychunk import ArrayChunk
 
 if TYPE_CHECKING:
     import torch
@@ -252,16 +257,56 @@ class Array(InflatableObject):
         ndarray_deserialized = np.load(bytes_io, allow_pickle=False)
         return cast(NDArray, ndarray_deserialized)
 
+    @property
+    def children(self) -> dict[str, InflatableObject]:
+        """Return a dictionary of ArrayChunks with their Object IDs as keys."""
+        return dict(self.slice_array())
+
+    def slice_array(self) -> list[tuple[str, InflatableObject]]:
+        """Slice Array data and construct a list of ArrayChunks."""
+        # Return cached chunks if they exist
+        if "_chunks" in self.__dict__:
+            return cast(list[tuple[str, InflatableObject]], self.__dict__["_chunks"])
+
+        # Chunks are not children as some of them may be identical
+        chunks: list[tuple[str, InflatableObject]] = []
+        # memoryview allows for zero-copy slicing
+        data_view = memoryview(self.data)
+        for start in range(0, len(data_view), MAX_ARRAY_CHUNK_SIZE):
+            end = min(start + MAX_ARRAY_CHUNK_SIZE, len(data_view))
+            ac = ArrayChunk(data_view[start:end])
+            chunks.append((ac.object_id, ac))
+
+        # Cache the chunks for future use
+        self.__dict__["_chunks"] = chunks
+        return chunks
+
     def deflate(self) -> bytes:
         """Deflate the Array."""
-        array_proto = ArrayProto(
-            dtype=self.dtype,
-            shape=self.shape,
-            stype=self.stype,
-            data=self.data,
-        )
-
-        obj_body = array_proto.SerializeToString(deterministic=True)
+        array_metadata: dict[str, str | tuple[int, ...] | list[int]] = {}
+
+        # We want to record all object_id even if repeated
+        # it can happend that chunks carry the exact same data
+        # for example when the array has only zeros
+        children_list = self.slice_array()
+        # Let's not save the entire object_id but a mapping to those
+        # that will be carried in the object head
+        # (replace a long object_id with a single scalar)
+        unique_children = list(self.children.keys())
+        arraychunk_ids = [unique_children.index(ch_id) for ch_id, _ in children_list]
+
+        # The deflated Array carries everything but the data
+        # The `arraychunk_ids` will be used during Array inflation
+        # to rematerialize the data from ArrayChunk objects.
+        array_metadata = {
+            "dtype": self.dtype,
+            "shape": self.shape,
+            "stype": self.stype,
+            "arraychunk_ids": arraychunk_ids,
+        }
+
+        # Serialize metadata dict
+        obj_body = json.dumps(array_metadata).encode("utf-8")
         return add_header_to_object_body(object_body=obj_body, obj=self)
 
     @classmethod
@@ -276,26 +321,55 @@ class Array(InflatableObject):
             The deflated object content of the Array.
 
         children : Optional[dict[str, InflatableObject]] (default: None)
-            Must be ``None``. ``Array`` does not support child objects.
-            Providing any children will raise a ``ValueError``.
+            Must be ``None``. ``Array`` must have child objects.
+            Providing no children will raise a ``ValueError``.
 
         Returns
         -------
         Array
             The inflated Array.
         """
-        if children:
-            raise ValueError("`Array` objects do not have children.")
+        if children is None:
+            children = {}
 
         obj_body = get_object_body(object_content, cls)
-        proto_array = ArrayProto.FromString(obj_body)
-        return cls(
-            dtype=proto_array.dtype,
-            shape=tuple(proto_array.shape),
-            stype=proto_array.stype,
-            data=proto_array.data,
+
+        # Extract children IDs from head
+        children_ids = get_object_children_ids_from_object_content(object_content)
+        # Decode the Array body
+        array_metadata: dict[str, str | tuple[int, ...] | list[int]] = json.loads(
+            obj_body.decode(encoding="utf-8")
         )
 
+        # Verify children ids in body match those passed for inflation
+        chunk_ids_indices = cast(list[int], array_metadata["arraychunk_ids"])
+        # Convert indices back to IDs
+        chunk_ids = [children_ids[i] for i in chunk_ids_indices]
+        # Check consistency
+        unique_arrayschunks = set(chunk_ids)
+        children_obj_ids = set(children.keys())
+        if unique_arrayschunks != children_obj_ids:
+            raise ValueError(
+                "Unexpected set of `children`. "
+                f"Expected {unique_arrayschunks} but got {children_obj_ids}."
+            )
+
+        # Materialize Array with empty data
+        array = cls(
+            dtype=cast(str, array_metadata["dtype"]),
+            shape=cast(tuple[int], tuple(array_metadata["shape"])),
+            stype=cast(str, array_metadata["stype"]),
+            data=b"",
+        )
+
+        # Now inject data from chunks
+        buff = bytearray()
+        for ch_id in chunk_ids:
+            buff += cast(ArrayChunk, children[ch_id]).data
+
+        array.data = bytes(buff)
+        return array
+
     @property
     def object_id(self) -> str:
         """Get object ID."""
@@ -320,4 +394,9 @@ class Array(InflatableObject):
         if name in ("dtype", "shape", "stype", "data"):
             # Mark as dirty if any of the main attributes are set
             self.is_dirty = True
+            # Clear cached object ID
+            self.__dict__.pop("_object_id", None)
+            # Clear cached chunks if data is set
+            if name == "data":
+                self.__dict__.pop("_chunks", None)
         super().__setattr__(name, value)

flwr/common/record/arraychunk.py

@@ -0,0 +1,59 @@
+# 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.
+# 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.
+# ==============================================================================
+"""ArrayChunk."""
+
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+
+from ..inflatable import InflatableObject, add_header_to_object_body, get_object_body
+
+
+@dataclass
+class ArrayChunk(InflatableObject):
+    """ArrayChunk type."""
+
+    data: memoryview
+
+    def deflate(self) -> bytes:
+        """Deflate the ArrayChunk."""
+        return add_header_to_object_body(object_body=self.data, obj=self)
+
+    @classmethod
+    def inflate(
+        cls, object_content: bytes, children: dict[str, InflatableObject] | None = None
+    ) -> ArrayChunk:
+        """Inflate an ArrayChunk from bytes.
+
+        Parameters
+        ----------
+        object_content : bytes
+            The deflated object content of the ArrayChunk.
+
+        children : Optional[dict[str, InflatableObject]] (default: None)
+            Must be ``None``. ``ArrayChunk`` does not support child objects.
+            Providing any children will raise a ``ValueError``.
+
+        Returns
+        -------
+        ArrayChunk
+            The inflated ArrayChunk.
+        """
+        if children:
+            raise ValueError("`ArrayChunk` objects do not have children.")
+
+        obj_body = get_object_body(object_content, cls)
+        return cls(data=memoryview(obj_body))

flwr/common/retry_invoker.py

@@ -17,6 +17,7 @@
 
 import itertools
 import random
+import threading
 import time
 from collections.abc import Generator, Iterable
 from dataclasses import dataclass
@@ -319,8 +320,12 @@ class RetryInvoker:
 
 def _make_simple_grpc_retry_invoker() -> RetryInvoker:
     """Create a simple gRPC retry invoker."""
+    lock = threading.Lock()
+    system_healthy = threading.Event()
+    system_healthy.set()  # Initially, the connection is healthy
 
-    def _on_sucess(retry_state: RetryState) -> None:
+    def _on_success(retry_state: RetryState) -> None:
+        system_healthy.set()
         if retry_state.tries > 1:
             log(
                 INFO,
@@ -329,17 +334,11 @@ def _make_simple_grpc_retry_invoker() -> RetryInvoker:
                 retry_state.tries,
             )
 
-    def _on_backoff(retry_state: RetryState) -> None:
-        if retry_state.tries == 1:
-            log(WARN, "Connection attempt failed, retrying...")
-        else:
-            log(
-                WARN,
-                "Connection attempt failed, retrying in %.2f seconds",
-                retry_state.actual_wait,
-            )
+    def _on_backoff(_: RetryState) -> None:
+        system_healthy.clear()
 
     def _on_giveup(retry_state: RetryState) -> None:
+        system_healthy.clear()
         if retry_state.tries > 1:
             log(
                 WARN,
@@ -355,15 +354,35 @@ def _make_simple_grpc_retry_invoker() -> RetryInvoker:
             return False
         return True
 
+    def _wait(wait_time: float) -> None:
+        # Use a lock to prevent multiple gRPC calls from retrying concurrently,
+        # which is unnecessary since they are all likely to fail.
+        with lock:
+            # Log the wait time
+            log(
+                WARN,
+                "Connection attempt failed, retrying in %.2f seconds",
+                wait_time,
+            )
+
+            start = time.monotonic()
+            # Avoid sequential waits if the system is healthy
+            system_healthy.wait(wait_time)
+
+        remaining_time = wait_time - (time.monotonic() - start)
+        if remaining_time > 0:
+            time.sleep(remaining_time)
+
     return RetryInvoker(
         wait_gen_factory=lambda: exponential(max_delay=MAX_RETRY_DELAY),
         recoverable_exceptions=grpc.RpcError,
         max_tries=None,
         max_time=None,
-        on_success=_on_sucess,
+        on_success=_on_success,
         on_backoff=_on_backoff,
         on_giveup=_on_giveup,
         should_giveup=_should_giveup_fn,
+        wait_function=_wait,
     )
 
 

flwr/common/serde.py

@@ -19,7 +19,6 @@ from collections import OrderedDict
 from typing import Any, cast
 
 # pylint: disable=E0611
-from flwr.proto.clientappio_pb2 import ClientAppOutputCode, ClientAppOutputStatus
 from flwr.proto.fab_pb2 import Fab as ProtoFab
 from flwr.proto.message_pb2 import Context as ProtoContext
 from flwr.proto.message_pb2 import Message as ProtoMessage
@@ -653,33 +652,6 @@ def run_from_proto(run_proto: ProtoRun) -> typing.Run:
     return run
 
 
-# === ClientApp status messages ===
-
-
-def clientappstatus_to_proto(
-    status: typing.ClientAppOutputStatus,
-) -> ClientAppOutputStatus:
-    """Serialize `ClientAppOutputStatus` to ProtoBuf."""
-    code = ClientAppOutputCode.SUCCESS
-    if status.code == typing.ClientAppOutputCode.DEADLINE_EXCEEDED:
-        code = ClientAppOutputCode.DEADLINE_EXCEEDED
-    if status.code == typing.ClientAppOutputCode.UNKNOWN_ERROR:
-        code = ClientAppOutputCode.UNKNOWN_ERROR
-    return ClientAppOutputStatus(code=code, message=status.message)
-
-
-def clientappstatus_from_proto(
-    msg: ClientAppOutputStatus,
-) -> typing.ClientAppOutputStatus:
-    """Deserialize `ClientAppOutputStatus` from ProtoBuf."""
-    code = typing.ClientAppOutputCode.SUCCESS
-    if msg.code == ClientAppOutputCode.DEADLINE_EXCEEDED:
-        code = typing.ClientAppOutputCode.DEADLINE_EXCEEDED
-    if msg.code == ClientAppOutputCode.UNKNOWN_ERROR:
-        code = typing.ClientAppOutputCode.UNKNOWN_ERROR
-    return typing.ClientAppOutputStatus(code=code, message=msg.message)
-
-
 # === Run status ===
 
 

flwr/common/telemetry.py

@@ -181,6 +181,10 @@ class EventType(str, Enum):
     RUN_SUPERNODE_ENTER = auto()
     RUN_SUPERNODE_LEAVE = auto()
 
+    # CLI: `flower-superexec`
+    RUN_SUPEREXEC_ENTER = auto()
+    RUN_SUPEREXEC_LEAVE = auto()
+
 
 # Use the ThreadPoolExecutor with max_workers=1 to have a queue
 # and also ensure that telemetry calls are not blocking.