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

flwr/common/inflatable_utils.py

@@ -0,0 +1,508 @@
+# 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.
+# ==============================================================================
+"""InflatableObject utilities."""
+
+import concurrent.futures
+import os
+import random
+import threading
+import time
+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,
+    HEAD_VALUE_DIVIDER,
+    MAX_CONCURRENT_PULLS,
+    MAX_CONCURRENT_PUSHES,
+    PULL_BACKOFF_CAP,
+    PULL_INITIAL_BACKOFF,
+    PULL_MAX_TIME,
+    PULL_MAX_TRIES_PER_OBJECT,
+)
+from .exit_handlers import add_exit_handler
+from .inflatable import (
+    InflatableObject,
+    UnexpectedObjectContentError,
+    _get_object_head,
+    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,
+    Message.__qualname__: Message,
+    MetricRecord.__qualname__: MetricRecord,
+    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
+    available."""
+
+    def __init__(self, object_id: str):
+        super().__init__(f"Object with ID '{object_id}' is not yet available.")
+
+
+class ObjectIdNotPreregisteredError(Exception):
+    """Exception raised when an object ID is not pre-registered."""
+
+    def __init__(self, object_id: str):
+        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],
+    *,
+    object_ids_to_push: Optional[set[str]] = None,
+    keep_objects: bool = False,
+    max_concurrent_pushes: int = MAX_CONCURRENT_PUSHES,
+) -> None:
+    """Push multiple objects to the servicer.
+
+    Parameters
+    ----------
+    objects : dict[str, InflatableObject]
+        A dictionary of objects to push, where keys are object IDs and values are
+        `InflatableObject` instances.
+    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.
+    object_ids_to_push : Optional[set[str]] (default: None)
+        A set of object IDs to push. If not provided, all objects will be pushed.
+    keep_objects : bool (default: False)
+        If `True`, the original objects will be kept in the `objects` dictionary
+        after pushing. If `False`, they will be removed from the dictionary to avoid
+        high memory usage.
+    max_concurrent_pushes : int (default: MAX_CONCURRENT_PUSHES)
+        The maximum number of concurrent pushes to perform.
+    """
+    lock = threading.Lock()
+
+    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."""
+        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)
+
+        # 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
+
+    # 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],
+    *,
+    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,
+) -> dict[str, bytes]:
+    """Pull multiple objects from the servicer.
+
+    Parameters
+    ----------
+    object_ids : list[str]
+        A list of object IDs to pull.
+    pull_object_fn : Callable[[str], bytes]
+        A function that takes an object ID and returns the object content as bytes.
+        The function should raise `ObjectUnavailableError` if the object is not yet
+        available, or `ObjectIdNotPreregisteredError` if the object ID is not
+        pre-registered.
+    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
+    -------
+    dict[str, bytes]
+        A dictionary where keys are object IDs and values are the pulled
+        object contents.
+    """
+    if max_tries_per_object is None:
+        max_tries_per_object = int(1e9)
+    if max_time is None:
+        max_time = float("inf")
+
+    results: dict[str, bytes] = {}
+    results_lock = threading.Lock()
+    err_to_raise: Optional[Exception] = None
+    early_stop = threading.Event()
+    start = time.monotonic()
+
+    def pull_with_retries(object_id: str) -> None:
+        """Attempt to pull a single object with retry and backoff."""
+        nonlocal err_to_raise
+        tries = 0
+        delay = initial_backoff
+
+        while not early_stop.is_set():
+            try:
+                object_content = pull_object_fn(object_id)
+                with results_lock:
+                    results[object_id] = object_content
+                return
+
+            except ObjectUnavailableError as err:
+                tries += 1
+                if (
+                    tries >= max_tries_per_object
+                    or time.monotonic() - start >= max_time
+                ):
+                    # Stop all work if one object exhausts retries
+                    early_stop.set()
+                    with results_lock:
+                        if err_to_raise is None:
+                            err_to_raise = err
+                    return
+
+                # Apply exponential backoff with ±20% jitter
+                sleep_time = delay * (1 + random.uniform(-0.2, 0.2))
+                early_stop.wait(sleep_time)
+                delay = min(delay * 2, backoff_cap)
+
+            except ObjectIdNotPreregisteredError as err:
+                # Permanent failure: object ID is invalid
+                early_stop.set()
+                with results_lock:
+                    if err_to_raise is None:
+                        err_to_raise = err
+                return
+
+    # Submit all pull tasks concurrently
+    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
+
+    # 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
+
+    return results
+
+
+def inflate_object_from_contents(
+    object_id: str,
+    object_contents: dict[str, bytes],
+    *,
+    keep_object_contents: bool = False,
+    objects: Optional[dict[str, InflatableObject]] = None,
+) -> InflatableObject:
+    """Inflate an object from object contents.
+
+    Parameters
+    ----------
+    object_id : str
+        The ID of the object to inflate.
+    object_contents : dict[str, bytes]
+        A dictionary mapping object IDs to their contents as bytes.
+        All descendant objects must be present in this dictionary.
+    keep_object_contents : bool (default: False)
+        If `True`, the object content will be kept in the `object_contents`
+        dictionary after inflation. If `False`, the object content will be
+        removed from the dictionary to save memory.
+    objects : Optional[dict[str, InflatableObject]] (default: None)
+        No need to provide this parameter. A dictionary to store already
+        inflated objects, mapping object IDs to their corresponding
+        `InflatableObject` instances.
+
+    Returns
+    -------
+    InflatableObject
+        The inflated object.
+    """
+    if objects is None:
+        # Initialize objects dictionary
+        objects = {}
+
+    if object_id in objects:
+        # If the object is already in the objects dictionary, return it
+        return objects[object_id]
+
+    # Extract object class and object_ids of children
+    object_content = object_contents[object_id]
+    obj_type, children_obj_ids, _ = get_object_head_values_from_object_content(
+        object_content=object_contents[object_id]
+    )
+
+    # Remove the object content from the dictionary to save memory
+    if not keep_object_contents:
+        del object_contents[object_id]
+
+    # Resolve object class
+    cls_type = inflatable_class_registry[obj_type]
+
+    # Inflate all children objects
+    children: dict[str, InflatableObject] = {}
+    for child_obj_id in children_obj_ids:
+        children[child_obj_id] = inflate_object_from_contents(
+            child_obj_id,
+            object_contents,
+            keep_object_contents=keep_object_contents,
+            objects=objects,
+        )
+
+    # Inflate object passing its children
+    obj = cls_type.inflate(object_content, children=children)
+    del object_content  # Free memory after inflation
+    objects[object_id] = obj
+    return obj
+
+
+def validate_object_content(content: bytes) -> None:
+    """Validate the deflated content of an InflatableObject."""
+    try:
+        # Check if there is a head-body divider
+        index = content.find(HEAD_BODY_DIVIDER)
+        if index == -1:
+            raise ValueError(
+                "Unexpected format for object content. Head and body "
+                "could not be split."
+            )
+
+        head = _get_object_head(content)
+
+        # check if the head has three parts:
+        # <object_type> <children_ids> <object_body_len>
+        head_decoded = head.decode(encoding="utf-8")
+        head_parts = head_decoded.split(HEAD_VALUE_DIVIDER)
+
+        if len(head_parts) != 3:
+            raise ValueError("Unexpected format for object head.")
+
+        obj_type, children_str, body_len = head_parts
+
+        # Check that children IDs are valid IDs
+        children = children_str.split(",")
+        for children_id in children:
+            if children_id and not is_valid_sha256_hash(children_id):
+                raise ValueError(
+                    f"Detected invalid object ID ({children_id}) in children."
+                )
+
+        # Check that object type is recognized
+        if obj_type not in inflatable_class_registry:
+            if obj_type != "CustomDataClass":  # to allow for the class in tests
+                raise ValueError(f"Object of type {obj_type} is not supported.")
+
+        # Check if the body length in the head matches that of the body
+        actual_body_len = len(content) - len(head) - len(HEAD_BODY_DIVIDER)
+        if actual_body_len != int(body_len):
+            raise ValueError(
+                f"Object content length expected {body_len} bytes but got "
+                f"{actual_body_len} bytes."
+            )
+
+    except ValueError as err:
+        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