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

flwr/common/exit/exit.py

@@ -37,13 +37,13 @@ def flwr_exit(
 ) -> NoReturn:
     """Handle application exit with an optional message.
 
-    The exit message logged and displayed will follow this structure:
+    The exit message logged and displayed will follow this structure::
 
-    >>> Exit Code: <code>
-    >>> <message>
-    >>> <short-help-message>
-    >>>
-    >>> For more information, visit: <help-page-url>
+        Exit Code: <code>
+        <message>
+        <short-help-message>
+
+        For more information, visit: <help-page-url>
 
     - `<code>`: The unique exit code representing the termination reason.
     - `<message>`: Optional context or additional information about the exit.

flwr/common/exit_handlers.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.
@@ -30,6 +30,7 @@ SIGNAL_TO_EXIT_CODE: dict[int, int] = {
     signal.SIGINT: ExitCode.GRACEFUL_EXIT_SIGINT,
     signal.SIGTERM: ExitCode.GRACEFUL_EXIT_SIGTERM,
 }
+registered_exit_handlers: list[Callable[[], None]] = []
 
 # SIGQUIT is not available on Windows
 if hasattr(signal, "SIGQUIT"):
@@ -41,6 +42,7 @@ def register_exit_handlers(
     exit_message: Optional[str] = None,
     grpc_servers: Optional[list[Server]] = None,
     bckg_threads: Optional[list[Thread]] = None,
+    exit_handlers: Optional[list[Callable[[], None]]] = None,
 ) -> None:
     """Register exit handlers for `SIGINT`, `SIGTERM` and `SIGQUIT` signals.
 
@@ -56,8 +58,12 @@ def register_exit_handlers(
     bckg_threads: Optional[List[Thread]] (default: None)
         An optional list of threads that need to be gracefully
         terminated before exiting.
+    exit_handlers: Optional[List[Callable[[], None]]] (default: None)
+        An optional list of exit handlers to be called before exiting.
+        Additional exit handlers can be added using `add_exit_handler`.
     """
     default_handlers: dict[int, Callable[[int, FrameType], None]] = {}
+    registered_exit_handlers.extend(exit_handlers or [])
 
     def graceful_exit_handler(signalnum: int, _frame: FrameType) -> None:
         """Exit handler to be registered with `signal.signal`.
@@ -68,6 +74,9 @@ def register_exit_handlers(
         # Reset to default handler
         signal.signal(signalnum, default_handlers[signalnum])  # type: ignore
 
+        for handler in registered_exit_handlers:
+            handler()
+
         if grpc_servers is not None:
             for grpc_server in grpc_servers:
                 grpc_server.stop(grace=1)
@@ -87,3 +96,24 @@ def register_exit_handlers(
     for sig in SIGNAL_TO_EXIT_CODE:
         default_handler = signal.signal(sig, graceful_exit_handler)  # type: ignore
         default_handlers[sig] = default_handler  # type: ignore
+
+
+def add_exit_handler(exit_handler: Callable[[], None]) -> None:
+    """Add an exit handler to be called on graceful exit.
+
+    This function allows you to register additional exit handlers
+    that will be executed when the application exits gracefully,
+    if `register_exit_handlers` was called.
+
+    Parameters
+    ----------
+    exit_handler : Callable[[], None]
+        A callable that takes no arguments and performs cleanup or
+        other actions before the application exits.
+
+    Notes
+    -----
+    This method is not thread-safe, and it allows you to add the
+    same exit handler multiple times.
+    """
+    registered_exit_handlers.append(exit_handler)

flwr/common/grpc.py

@@ -1,4 +1,4 @@
-# Copyright 2022 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.

flwr/common/heartbeat.py

@@ -0,0 +1,165 @@
+# 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.
+# ==============================================================================
+"""Heartbeat sender."""
+
+
+import random
+import threading
+from typing import Callable, Union
+
+import grpc
+
+# pylint: disable=E0611
+from flwr.proto.heartbeat_pb2 import SendAppHeartbeatRequest
+from flwr.proto.serverappio_pb2_grpc import ServerAppIoStub
+from flwr.proto.simulationio_pb2_grpc import SimulationIoStub
+
+# pylint: enable=E0611
+from .constant import (
+    HEARTBEAT_BASE_MULTIPLIER,
+    HEARTBEAT_CALL_TIMEOUT,
+    HEARTBEAT_DEFAULT_INTERVAL,
+    HEARTBEAT_RANDOM_RANGE,
+)
+from .retry_invoker import RetryInvoker, exponential
+
+
+class HeartbeatFailure(Exception):
+    """Exception raised when a heartbeat fails."""
+
+
+class HeartbeatSender:
+    """Periodically send heartbeat signals to a server in a background thread.
+
+    This class uses the provided `heartbeat_fn` to send heartbeats. If a heartbeat
+    attempt fails, it will be retried using an exponential backoff strategy.
+
+    Parameters
+    ----------
+    heartbeat_fn : Callable[[], bool]
+        Function used to send a heartbeat signal. It should return True if the heartbeat
+        succeeds, or False if it fails. Any internal exceptions (e.g., gRPC errors)
+        should be handled within this function to ensure boolean return values.
+    """
+
+    def __init__(
+        self,
+        heartbeat_fn: Callable[[], bool],
+    ) -> None:
+        self.heartbeat_fn = heartbeat_fn
+        self._stop_event = threading.Event()
+        self._thread = threading.Thread(target=self._run, daemon=True)
+        self._retry_invoker = RetryInvoker(
+            lambda: exponential(max_delay=20),
+            HeartbeatFailure,  # The only exception we want to retry on
+            max_tries=None,
+            max_time=None,
+            # Allow the stop event to interrupt the wait
+            wait_function=self._stop_event.wait,  # type: ignore
+        )
+
+    def start(self) -> None:
+        """Start the heartbeat sender."""
+        if self._thread.is_alive():
+            raise RuntimeError("Heartbeat sender is already running.")
+        if self._stop_event.is_set():
+            raise RuntimeError("Cannot start a stopped heartbeat sender.")
+        self._thread.start()
+
+    def stop(self) -> None:
+        """Stop the heartbeat sender."""
+        if not self._thread.is_alive():
+            raise RuntimeError("Heartbeat sender is not running.")
+        self._stop_event.set()
+        self._thread.join()
+
+    @property
+    def is_running(self) -> bool:
+        """Return True if the heartbeat sender is running, False otherwise."""
+        return self._thread.is_alive() and not self._stop_event.is_set()
+
+    def _run(self) -> None:
+        """Periodically send heartbeats until stopped."""
+        while not self._stop_event.is_set():
+            # Attempt to send a heartbeat with retry on failure
+            self._retry_invoker.invoke(self._heartbeat)
+
+            # Calculate the interval for the next heartbeat
+            # Formula: next_interval = (interval - timeout) * random.uniform(0.7, 0.9)
+            rd = random.uniform(*HEARTBEAT_RANDOM_RANGE)
+            next_interval: float = HEARTBEAT_DEFAULT_INTERVAL - HEARTBEAT_CALL_TIMEOUT
+            next_interval *= HEARTBEAT_BASE_MULTIPLIER + rd
+
+            # Wait for the calculated interval or exit early if stopped
+            self._stop_event.wait(next_interval)
+
+    def _heartbeat(self) -> None:
+        """Send a single heartbeat and raise an exception if it fails.
+
+        Call the provided `heartbeat_fn`. If the function returns False,
+        a `HeartbeatFailure` exception is raised to trigger the retry mechanism.
+        """
+        if not self._stop_event.is_set():
+            if not self.heartbeat_fn():
+                raise HeartbeatFailure
+
+
+def get_grpc_app_heartbeat_fn(
+    stub: Union[ServerAppIoStub, SimulationIoStub],
+    run_id: int,
+    *,
+    failure_message: str,
+) -> Callable[[], bool]:
+    """Get the function to send a heartbeat to gRPC endpoint.
+
+    This function is for app heartbeats only. It is not used for node heartbeats.
+
+    Parameters
+    ----------
+    stub : Union[ServerAppIoStub, SimulationIoStub]
+        gRPC stub to send the heartbeat.
+    run_id : int
+        The run ID to use in the heartbeat request.
+    failure_message : str
+        Error message to raise if the heartbeat fails.
+
+    Returns
+    -------
+    Callable[[], bool]
+        Function that sends a heartbeat to the gRPC endpoint.
+    """
+    # Construct the heartbeat request
+    req = SendAppHeartbeatRequest(
+        run_id=run_id, heartbeat_interval=HEARTBEAT_DEFAULT_INTERVAL
+    )
+
+    def fn() -> bool:
+        # Call ServerAppIo API
+        try:
+            res = stub.SendAppHeartbeat(req)
+        except grpc.RpcError as e:
+            status_code = e.code()
+            if status_code == grpc.StatusCode.UNAVAILABLE:
+                return False
+            if status_code == grpc.StatusCode.DEADLINE_EXCEEDED:
+                return False
+            raise
+
+        # Check if not successful
+        if not res.success:
+            raise RuntimeError(failure_message)
+        return True
+
+    return fn

flwr/common/inflatable.py

@@ -0,0 +1,290 @@
+# 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 base class."""
+
+
+from __future__ import annotations
+
+import hashlib
+import threading
+from collections.abc import Iterator
+from contextlib import contextmanager
+from typing import TypeVar, cast
+
+from flwr.proto.message_pb2 import ObjectTree  # pylint: disable=E0611
+
+from .constant import HEAD_BODY_DIVIDER, HEAD_VALUE_DIVIDER
+
+
+class UnexpectedObjectContentError(Exception):
+    """Exception raised when the content of an object does not conform to the expected
+    structure for an InflatableObject (i.e., head, body, and values within the head)."""
+
+    def __init__(self, object_id: str, reason: str):
+        super().__init__(
+            f"Object with ID '{object_id}' has an unexpected structure. {reason}"
+        )
+
+
+_ctx = threading.local()
+
+
+def _is_recompute_enabled() -> bool:
+    """Check if recomputing object IDs is enabled."""
+    return getattr(_ctx, "recompute_object_id_enabled", True)
+
+
+def _get_computed_object_ids() -> set[str]:
+    """Get the set of computed object IDs."""
+    return getattr(_ctx, "computed_object_ids", set())
+
+
+@contextmanager
+def no_object_id_recompute() -> Iterator[None]:
+    """Context manager to disable recomputing object IDs."""
+    old_value = _is_recompute_enabled()
+    old_set = _get_computed_object_ids()
+    _ctx.recompute_object_id_enabled = False
+    _ctx.computed_object_ids = set()
+    try:
+        yield
+    finally:
+        _ctx.recompute_object_id_enabled = old_value
+        _ctx.computed_object_ids = old_set
+
+
+class InflatableObject:
+    """Base class for inflatable objects."""
+
+    def deflate(self) -> bytes:
+        """Deflate object."""
+        raise NotImplementedError()
+
+    @classmethod
+    def inflate(
+        cls, object_content: bytes, children: dict[str, InflatableObject] | None = None
+    ) -> InflatableObject:
+        """Inflate the object from bytes.
+
+        Parameters
+        ----------
+        object_content : bytes
+            The deflated object content.
+
+        children : Optional[dict[str, InflatableObject]] (default: None)
+            Dictionary of children InflatableObjects mapped to their object IDs. These
+            childrens enable the full inflation of the parent InflatableObject.
+
+        Returns
+        -------
+        InflatableObject
+            The inflated object.
+        """
+        raise NotImplementedError()
+
+    @property
+    def object_id(self) -> str:
+        """Get object_id."""
+        # If recomputing object ID is disabled and the object ID is already computed,
+        # return the cached object ID.
+        if (
+            not _is_recompute_enabled()
+            and (obj_id := self.__dict__.get("_object_id"))
+            in _get_computed_object_ids()
+        ):
+            return cast(str, obj_id)
+
+        if self.is_dirty or "_object_id" not in self.__dict__:
+            obj_id = get_object_id(self.deflate())
+            self.__dict__["_object_id"] = obj_id
+
+            # If recomputing object ID is disabled, add the object ID to the set of
+            # computed object IDs to avoid recomputing it within the context.
+            if not _is_recompute_enabled():
+                _get_computed_object_ids().add(obj_id)
+        return cast(str, self.__dict__["_object_id"])
+
+    @property
+    def children(self) -> dict[str, InflatableObject] | None:
+        """Get all child objects as a dictionary or None if there are no children."""
+        return None
+
+    @property
+    def is_dirty(self) -> bool:
+        """Check if the object is dirty after the last deflation.
+
+        An object is considered dirty if its content has changed since the last its
+        object ID was computed.
+        """
+        return True
+
+
+T = TypeVar("T", bound=InflatableObject)
+
+
+def get_object_id(object_content: bytes) -> str:
+    """Return a SHA-256 hash of the (deflated) object content."""
+    return hashlib.sha256(object_content).hexdigest()
+
+
+def get_object_body(object_content: bytes, cls: type[T]) -> bytes:
+    """Return object body but raise an error if object type doesn't match class name."""
+    class_name = cls.__qualname__
+    object_type = get_object_type_from_object_content(object_content)
+    if not object_type == class_name:
+        raise ValueError(
+            f"Class name ({class_name}) and object type "
+            f"({object_type}) do not match."
+        )
+
+    # Return object body
+    return _get_object_body(object_content)
+
+
+def add_header_to_object_body(object_body: bytes, obj: InflatableObject) -> bytes:
+    """Add header to object content."""
+    # Construct header
+    header = f"%s{HEAD_VALUE_DIVIDER}%s{HEAD_VALUE_DIVIDER}%d" % (
+        obj.__class__.__qualname__,  # Type of object
+        ",".join((obj.children or {}).keys()),  # IDs of child objects
+        len(object_body),  # Length of object body
+    )
+
+    # Concatenate header and object body
+    ret = bytearray()
+    ret.extend(header.encode(encoding="utf-8"))
+    ret.extend(HEAD_BODY_DIVIDER)
+    ret.extend(object_body)
+    return bytes(ret)
+
+
+def _get_object_head(object_content: bytes) -> bytes:
+    """Return object head from object content."""
+    index = object_content.find(HEAD_BODY_DIVIDER)
+    return object_content[:index]
+
+
+def _get_object_body(object_content: bytes) -> bytes:
+    """Return object body from object content."""
+    index = object_content.find(HEAD_BODY_DIVIDER)
+    return object_content[index + len(HEAD_BODY_DIVIDER) :]
+
+
+def is_valid_sha256_hash(object_id: str) -> bool:
+    """Check if the given string is a valid SHA-256 hash.
+
+    Parameters
+    ----------
+    object_id : str
+        The string to check.
+
+    Returns
+    -------
+    bool
+        ``True`` if the string is a valid SHA-256 hash, ``False`` otherwise.
+    """
+    if len(object_id) != 64:
+        return False
+    try:
+        # If base 16 int conversion succeeds, it's a valid hexadecimal str
+        int(object_id, 16)
+        return True
+    except ValueError:
+        return False
+
+
+def get_object_type_from_object_content(object_content: bytes) -> str:
+    """Return object type from bytes."""
+    return get_object_head_values_from_object_content(object_content)[0]
+
+
+def get_object_children_ids_from_object_content(object_content: bytes) -> list[str]:
+    """Return object children IDs from bytes."""
+    return get_object_head_values_from_object_content(object_content)[1]
+
+
+def get_object_body_len_from_object_content(object_content: bytes) -> int:
+    """Return length of the object body."""
+    return get_object_head_values_from_object_content(object_content)[2]
+
+
+def get_object_head_values_from_object_content(
+    object_content: bytes,
+) -> tuple[str, list[str], int]:
+    """Return object type and body length from object content.
+
+    Parameters
+    ----------
+    object_content : bytes
+        The deflated object content.
+
+    Returns
+    -------
+    tuple[str, list[str], int]
+        A tuple containing:
+        - The object type as a string.
+        - A list of child object IDs as strings.
+        - The length of the object body as an integer.
+    """
+    head = _get_object_head(object_content).decode(encoding="utf-8")
+    obj_type, children_str, body_len = head.split(HEAD_VALUE_DIVIDER)
+    children_ids = children_str.split(",") if children_str else []
+    return obj_type, children_ids, int(body_len)
+
+
+def get_descendant_object_ids(obj: InflatableObject) -> set[str]:
+    """Get a set of object IDs of all descendants."""
+    descendants = set(get_all_nested_objects(obj).keys())
+    # Exclude Object ID of parent object
+    descendants.discard(obj.object_id)
+    return descendants
+
+
+def get_all_nested_objects(obj: InflatableObject) -> dict[str, InflatableObject]:
+    """Get a dictionary of all nested objects, including the object itself.
+
+    Each key in the dictionary is an object ID, and the entries are ordered by post-
+    order traversal, i.e., child objects appear before their respective parents.
+    """
+    ret: dict[str, InflatableObject] = {}
+    if children := obj.children:
+        for child in children.values():
+            ret.update(get_all_nested_objects(child))
+
+    ret[obj.object_id] = obj
+
+    return ret
+
+
+def get_object_tree(obj: InflatableObject) -> ObjectTree:
+    """Get a tree representation of the InflatableObject."""
+    tree_children = []
+    if children := obj.children:
+        for child in children.values():
+            tree_children.append(get_object_tree(child))
+    return ObjectTree(object_id=obj.object_id, children=tree_children)
+
+
+def iterate_object_tree(
+    tree: ObjectTree,
+) -> Iterator[ObjectTree]:
+    """Iterate over the object tree and yield object IDs.
+
+    This function performs a post-order traversal of the tree, yielding the object ID of
+    each node after all its children have been yielded.
+    """
+    for child in tree.children:
+        yield from iterate_object_tree(child)
+    yield tree

flwr/common/inflatable_grpc_utils.py

@@ -0,0 +1,99 @@
+# 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 gRPC utils."""
+
+
+from typing import Callable
+
+from flwr.proto.message_pb2 import (  # pylint: disable=E0611
+    PullObjectRequest,
+    PullObjectResponse,
+    PushObjectRequest,
+    PushObjectResponse,
+)
+from flwr.proto.node_pb2 import Node  # pylint: disable=E0611
+
+from .inflatable_utils import ObjectIdNotPreregisteredError, ObjectUnavailableError
+
+
+def make_pull_object_fn_grpc(
+    pull_object_grpc: Callable[[PullObjectRequest], PullObjectResponse],
+    node: Node,
+    run_id: int,
+) -> Callable[[str], bytes]:
+    """Create a pull object function that uses gRPC to pull objects.
+
+    Parameters
+    ----------
+    pull_object_grpc : Callable[[PullObjectRequest], PullObjectResponse]
+        The gRPC function to pull objects, e.g., `FleetStub.PullObject`.
+    node : Node
+        The node making the request.
+    run_id : int
+        The run ID for the current operation.
+
+    Returns
+    -------
+    Callable[[str], bytes]
+        A function that takes an object ID and returns the object content as bytes.
+        The function raises `ObjectIdNotPreregisteredError` if the object ID is not
+        pre-registered, or `ObjectUnavailableError` if the object is not yet available.
+    """
+
+    def pull_object_fn(object_id: str) -> bytes:
+        request = PullObjectRequest(node=node, run_id=run_id, object_id=object_id)
+        response: PullObjectResponse = pull_object_grpc(request)
+        if not response.object_found:
+            raise ObjectIdNotPreregisteredError(object_id)
+        if not response.object_available:
+            raise ObjectUnavailableError(object_id)
+        return response.object_content
+
+    return pull_object_fn
+
+
+def make_push_object_fn_grpc(
+    push_object_grpc: Callable[[PushObjectRequest], PushObjectResponse],
+    node: Node,
+    run_id: int,
+) -> Callable[[str, bytes], None]:
+    """Create a push object function that uses gRPC to push objects.
+
+    Parameters
+    ----------
+    push_object_grpc : Callable[[PushObjectRequest], PushObjectResponse]
+        The gRPC function to push objects, e.g., `FleetStub.PushObject`.
+    node : Node
+        The node making the request.
+    run_id : int
+        The run ID for the current operation.
+
+    Returns
+    -------
+    Callable[[str, bytes], None]
+        A function that takes an object ID and its content as bytes, and pushes it
+        to the servicer. The function raises `ObjectIdNotPreregisteredError` if
+        the object ID is not pre-registered.
+    """
+
+    def push_object_fn(object_id: str, object_content: bytes) -> None:
+        request = PushObjectRequest(
+            node=node, run_id=run_id, object_id=object_id, object_content=object_content
+        )
+        response: PushObjectResponse = push_object_grpc(request)
+        if not response.stored:
+            raise ObjectIdNotPreregisteredError(object_id)
+
+    return push_object_fn