flwr-nightly 1.19.0.dev20250605__py3-none-any.whl → 1.19.0.dev20250607__py3-none-any.whl

flwr/client/grpc_rere_client/connection.py

@@ -30,16 +30,16 @@ from flwr.common import GRPC_MAX_MESSAGE_LENGTH
 from flwr.common.constant import HEARTBEAT_CALL_TIMEOUT, HEARTBEAT_DEFAULT_INTERVAL
 from flwr.common.grpc import create_channel, on_channel_state_change
 from flwr.common.heartbeat import HeartbeatSender
+from flwr.common.inflatable import get_all_nested_objects
 from flwr.common.inflatable_grpc_utils import (
-    pull_object_from_servicer,
-    push_object_to_servicer,
+    inflate_object_from_contents,
+    make_pull_object_fn_grpc,
+    make_push_object_fn_grpc,
+    pull_objects,
+    push_objects,
 )
 from flwr.common.logger import log
-from flwr.common.message import (
-    Message,
-    get_message_to_descendant_id_mapping,
-    remove_content_from_message,
-)
+from flwr.common.message import Message, remove_content_from_message
 from flwr.common.retry_invoker import RetryInvoker, _wrap_stub
 from flwr.common.secure_aggregation.crypto.symmetric_encryption import (
     generate_key_pairs,
@@ -60,6 +60,7 @@ from flwr.proto.heartbeat_pb2 import (  # pylint: disable=E0611
     SendNodeHeartbeatRequest,
     SendNodeHeartbeatResponse,
 )
+from flwr.proto.message_pb2 import ObjectIDs  # pylint: disable=E0611
 from flwr.proto.node_pb2 import Node  # pylint: disable=E0611
 from flwr.proto.run_pb2 import GetRunRequest, GetRunResponse  # pylint: disable=E0611
 
@@ -265,21 +266,21 @@ def grpc_request_response(  # pylint: disable=R0913,R0914,R0915,R0917
         in_message: Optional[Message] = None
 
         if message_proto:
-            in_message = cast(
-                Message,
-                pull_object_from_servicer(
-                    object_id=message_proto.metadata.message_id,
-                    stub=stub,
+            msg_id = message_proto.metadata.message_id
+            all_object_contents = pull_objects(
+                list(response.objects_to_pull[msg_id].object_ids) + [msg_id],
+                pull_object_fn=make_pull_object_fn_grpc(
+                    pull_object_grpc=stub.PullObject,
                     node=node,
                     run_id=message_proto.metadata.run_id,
                 ),
             )
-
-        if in_message:
+            in_message = cast(
+                Message, inflate_object_from_contents(msg_id, all_object_contents)
+            )
             # The deflated message doesn't contain the message_id (its own object_id)
             # Inject
-            # pylint: disable-next=W0212
-            in_message.metadata._message_id = message_proto.metadata.message_id  # type: ignore
+            in_message.metadata.__dict__["_message_id"] = msg_id
 
         # Remember `metadata` of the in message
         nonlocal metadata
@@ -308,23 +309,30 @@ def grpc_request_response(  # pylint: disable=R0913,R0914,R0915,R0917
             log(ERROR, "Invalid out message")
             return
 
+        # Get all nested objects
+        all_objects = get_all_nested_objects(message)
+        all_object_ids = list(all_objects.keys())
+        msg_id = all_object_ids[-1]  # Last object is the message itself
+        descendant_ids = all_object_ids[:-1]  # All but the last object are descendants
+
         # Serialize Message
         message_proto = message_to_proto(message=remove_content_from_message(message))
-        descendants_mapping = get_message_to_descendant_id_mapping(message)
         request = PushMessagesRequest(
             node=node,
             messages_list=[message_proto],
-            msg_to_descendant_mapping=descendants_mapping,
+            msg_to_descendant_mapping={msg_id: ObjectIDs(object_ids=descendant_ids)},
         )
         response: PushMessagesResponse = stub.PushMessages(request=request)
 
         if response.objects_to_push:
             objs_to_push = set(response.objects_to_push[message.object_id].object_ids)
-            push_object_to_servicer(
-                message,
-                stub,
-                node,
-                run_id=message.metadata.run_id,
+            push_objects(
+                all_objects,
+                push_object_fn=make_push_object_fn_grpc(
+                    push_object_grpc=stub.PushObject,
+                    node=node,
+                    run_id=message.metadata.run_id,
+                ),
                 object_ids_to_push=objs_to_push,
             )
             log(DEBUG, "Pushed %s objects to servicer.", len(objs_to_push))

flwr/common/constant.py

@@ -138,6 +138,17 @@ HEAD_VALUE_DIVIDER = " "
 # Constants for serialization
 INT64_MAX_VALUE = 9223372036854775807  # (1 << 63) - 1
 
+# Constants for `flwr-serverapp` and `flwr-clientapp` CLI commands
+FLWR_APP_TOKEN_LENGTH = 128  # Length of the token used
+
+# Constants for object pushing and pulling
+MAX_CONCURRENT_PUSHES = 8  # Default maximum number of concurrent pushes
+MAX_CONCURRENT_PULLS = 8  # Default maximum number of concurrent pulls
+PULL_MAX_TIME = 7200  # Default maximum time to wait for pulling objects
+PULL_MAX_TRIES_PER_OBJECT = 500  # Default maximum number of tries to pull an object
+PULL_INITIAL_BACKOFF = 1  # Initial backoff time for pulling objects
+PULL_BACKOFF_CAP = 10  # Maximum backoff time for pulling objects
+
 
 class MessageType:
     """Message type."""

flwr/common/inflatable.py

@@ -25,6 +25,16 @@ from .constant import HEAD_BODY_DIVIDER, HEAD_VALUE_DIVIDER
 from .logger import log
 
 
+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}"
+        )
+
+
 class InflatableObject:
     """Base class for inflatable objects."""
 
@@ -117,12 +127,14 @@ def add_header_to_object_body(object_body: bytes, obj: InflatableObject) -> byte
 
 def _get_object_head(object_content: bytes) -> bytes:
     """Return object head from object content."""
-    return object_content.split(HEAD_BODY_DIVIDER, 1)[0]
+    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."""
-    return object_content.split(HEAD_BODY_DIVIDER, 1)[1]
+    index = object_content.find(HEAD_BODY_DIVIDER)
+    return object_content[index + len(HEAD_BODY_DIVIDER) :]
 
 
 def is_valid_sha256_hash(object_id: str) -> bool:
@@ -197,21 +209,25 @@ def get_object_head_values_from_object_content(
     return obj_type, children_ids, int(body_len)
 
 
-def _get_descendants_object_ids_recursively(obj: InflatableObject) -> set[str]:
+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
 
-    descendants: set[str] = set()
-    if children := obj.children:
-        for child in children.values():
-            descendants |= _get_descendants_object_ids_recursively(child)
 
-    descendants.add(obj.object_id)
+def get_all_nested_objects(obj: InflatableObject) -> dict[str, InflatableObject]:
+    """Get a dictionary of all nested objects, including the object itself.
 
-    return descendants
+    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
 
-def get_desdendant_object_ids(obj: InflatableObject) -> set[str]:
-    """Get a set of object IDs of all descendants."""
-    descendants = _get_descendants_object_ids_recursively(obj)
-    # Exclude Object ID of parent object
-    descendants.discard(obj.object_id)
-    return descendants
+    return ret

flwr/common/inflatable_grpc_utils.py

@@ -15,11 +15,12 @@
 """InflatableObject utils."""
 
 
-from time import sleep
-from typing import Optional, Union
+import concurrent.futures
+import random
+import threading
+import time
+from typing import Callable, Optional
 
-from flwr.client.grpc_rere_client.grpc_adapter import GrpcAdapter
-from flwr.proto.fleet_pb2_grpc import FleetStub  # pylint: disable=E0611
 from flwr.proto.message_pb2 import (  # pylint: disable=E0611
     PullObjectRequest,
     PullObjectResponse,
@@ -27,13 +28,16 @@ from flwr.proto.message_pb2 import (  # pylint: disable=E0611
     PushObjectResponse,
 )
 from flwr.proto.node_pb2 import Node  # pylint: disable=E0611
-from flwr.proto.serverappio_pb2_grpc import ServerAppIoStub  # pylint: disable=E0611
 
-from .inflatable import (
-    InflatableObject,
-    get_object_head_values_from_object_content,
-    get_object_id,
+from .constant import (
+    MAX_CONCURRENT_PULLS,
+    MAX_CONCURRENT_PUSHES,
+    PULL_BACKOFF_CAP,
+    PULL_INITIAL_BACKOFF,
+    PULL_MAX_TIME,
+    PULL_MAX_TRIES_PER_OBJECT,
 )
+from .inflatable import InflatableObject, get_object_head_values_from_object_content
 from .message import Message
 from .record import Array, ArrayRecord, ConfigRecord, MetricRecord, RecordDict
 
@@ -48,75 +52,311 @@ inflatable_class_registry: dict[str, type[InflatableObject]] = {
 }
 
 
-def push_object_to_servicer(
-    obj: InflatableObject,
-    stub: Union[FleetStub, ServerAppIoStub, GrpcAdapter],
+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 make_pull_object_fn_grpc(
+    pull_object_grpc: Callable[[PullObjectRequest], PullObjectResponse],
     node: Node,
     run_id: int,
-    object_ids_to_push: Optional[set[str]] = None,
-) -> set[str]:
-    """Recursively deflate an object and push it to the servicer.
+) -> Callable[[str], bytes]:
+    """Create a pull object function that uses gRPC to pull objects.
 
-    Objects with the same ID are not pushed twice. If `object_ids_to_push` is set,
-    only objects with those IDs are pushed. It returns the set of pushed object
-    IDs.
+    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.
     """
-    pushed_object_ids: set[str] = set()
-    # Push children if it has any
-    if children := obj.children:
-        for child in children.values():
-            pushed_object_ids |= push_object_to_servicer(
-                child, stub, node, run_id, object_ids_to_push
-            )
-
-    # Deflate object and push
-    object_content = obj.deflate()
-    object_id = get_object_id(object_content)
-    # Push always if no object set is specified, or if the object is in the set
-    if object_ids_to_push is None or object_id in object_ids_to_push:
-        _: PushObjectResponse = stub.PushObject(
-            PushObjectRequest(
-                node=node,
-                run_id=run_id,
-                object_id=object_id,
-                object_content=object_content,
-            )
-        )
-        pushed_object_ids.add(object_id)
 
-    return pushed_object_ids
+    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 pull_object_from_servicer(
-    object_id: str,
-    stub: Union[FleetStub, ServerAppIoStub, GrpcAdapter],
+
+def make_push_object_fn_grpc(
+    push_object_grpc: Callable[[PushObjectRequest], PushObjectResponse],
     node: Node,
     run_id: int,
-) -> InflatableObject:
-    """Recursively inflate an object by pulling it from the servicer."""
-    # Pull object
-    object_available = False
-    while not object_available:
-        object_proto: PullObjectResponse = stub.PullObject(
-            PullObjectRequest(node=node, run_id=run_id, object_id=object_id)
+) -> 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
         )
-        object_available = object_proto.object_available
-        object_content = object_proto.object_content
-        sleep(0.1)
+        response: PushObjectResponse = push_object_grpc(request)
+        if not response.stored:
+            raise ObjectIdNotPreregisteredError(object_id)
+
+    return push_object_fn
+
+
+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.
+    """
+    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:
+        """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)
+
+    with concurrent.futures.ThreadPoolExecutor(
+        max_workers=max_concurrent_pushes
+    ) as executor:
+        list(executor.map(push, list(objects.keys())))
+
+
+def pull_objects(  # pylint: disable=too-many-arguments
+    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
+    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
+        }
+
+        # Wait for completion
+        concurrent.futures.wait(futures)
+
+    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_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]
 
-    # Pull all children objects
+    # Inflate all children objects
     children: dict[str, InflatableObject] = {}
-    for child_object_id in children_obj_ids:
-        children[child_object_id] = pull_object_from_servicer(
-            child_object_id, stub, node, run_id
+    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
-    return cls_type.inflate(object_content, children=children)
+    obj = cls_type.inflate(object_content, children=children)
+    del object_content  # Free memory after inflation
+    objects[object_id] = obj
+    return obj

flwr/common/inflatable_utils.py

@@ -0,0 +1,75 @@
+# 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."""
+
+
+from .constant import HEAD_BODY_DIVIDER, HEAD_VALUE_DIVIDER
+from .inflatable import (
+    UnexpectedObjectContentError,
+    _get_object_head,
+    get_object_id,
+    is_valid_sha256_hash,
+)
+from .inflatable_grpc_utils import inflatable_class_registry
+
+
+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

flwr/common/message.py

@@ -32,7 +32,7 @@ from .constant import MESSAGE_TTL_TOLERANCE
 from .inflatable import (
     InflatableObject,
     add_header_to_object_body,
-    get_desdendant_object_ids,
+    get_descendant_object_ids,
     get_object_body,
     get_object_children_ids_from_object_content,
 )
@@ -524,6 +524,6 @@ def get_message_to_descendant_id_mapping(message: Message) -> dict[str, ObjectID
     """Construct a mapping between message object_id and that of its descendants."""
     return {
         message.object_id: ObjectIDs(
-            object_ids=list(get_desdendant_object_ids(message))
+            object_ids=list(get_descendant_object_ids(message))
         )
     }

flwr/server/grid/grpc_grid.py

@@ -28,18 +28,20 @@ from flwr.common.constant import (
     SUPERLINK_NODE_ID,
 )
 from flwr.common.grpc import create_channel, on_channel_state_change
+from flwr.common.inflatable import get_all_nested_objects
 from flwr.common.inflatable_grpc_utils import (
-    pull_object_from_servicer,
-    push_object_to_servicer,
+    inflate_object_from_contents,
+    make_pull_object_fn_grpc,
+    make_push_object_fn_grpc,
+    pull_objects,
+    push_objects,
 )
 from flwr.common.logger import log, warn_deprecated_feature
-from flwr.common.message import (
-    get_message_to_descendant_id_mapping,
-    remove_content_from_message,
-)
+from flwr.common.message import remove_content_from_message
 from flwr.common.retry_invoker import _make_simple_grpc_retry_invoker, _wrap_stub
 from flwr.common.serde import message_to_proto, run_from_proto
 from flwr.common.typing import Run
+from flwr.proto.message_pb2 import ObjectIDs  # pylint: disable=E0611
 from flwr.proto.node_pb2 import Node  # pylint: disable=E0611
 from flwr.proto.run_pb2 import GetRunRequest, GetRunResponse  # pylint: disable=E0611
 from flwr.proto.serverappio_pb2 import (  # pylint: disable=E0611
@@ -208,28 +210,34 @@ class GrpcGrid(Grid):
     def _try_push_message(self, run_id: int, message: Message) -> str:
         """Push one message and its associated objects."""
         # Compute mapping of message descendants
-        descendants_mapping = get_message_to_descendant_id_mapping(message)
+        all_objects = get_all_nested_objects(message)
+        all_object_ids = list(all_objects.keys())
+        msg_id = all_object_ids[-1]  # Last object is the message itself
+        descendant_ids = all_object_ids[:-1]  # All but the last object are descendants
 
         # Call GrpcServerAppIoStub method
         res: PushInsMessagesResponse = self._stub.PushMessages(
             PushInsMessagesRequest(
                 messages_list=[message_to_proto(remove_content_from_message(message))],
                 run_id=run_id,
-                msg_to_descendant_mapping=descendants_mapping,
+                msg_to_descendant_mapping={
+                    msg_id: ObjectIDs(object_ids=descendant_ids)
+                },
             )
         )
 
         # Push objects
-        msg_id = res.message_ids[0]
         # If Message was added to the LinkState correctly
         if msg_id is not None:
             obj_ids_to_push = set(res.objects_to_push[msg_id].object_ids)
             # Push only object that are not in the store
-            push_object_to_servicer(
-                message,
-                self._stub,
-                node=self.node,
-                run_id=run_id,
+            push_objects(
+                all_objects,
+                push_object_fn=make_push_object_fn_grpc(
+                    push_object_grpc=self._stub.PushObject,
+                    node=self.node,
+                    run_id=run_id,
+                ),
                 object_ids_to_push=obj_ids_to_push,
             )
         return msg_id
@@ -289,14 +297,20 @@ class GrpcGrid(Grid):
             # Pull Messages from store
             inflated_msgs: list[Message] = []
             for msg_proto in res.messages_list:
-
-                message = pull_object_from_servicer(
-                    msg_proto.metadata.message_id,
-                    self._stub,
-                    node=self.node,
-                    run_id=run_id,
+                msg_id = msg_proto.metadata.message_id
+                all_object_contents = pull_objects(
+                    list(res.objects_to_pull[msg_id].object_ids) + [msg_id],
+                    pull_object_fn=make_pull_object_fn_grpc(
+                        pull_object_grpc=self._stub.PullObject,
+                        node=self.node,
+                        run_id=run_id,
+                    ),
+                )
+                message = cast(
+                    Message, inflate_object_from_contents(msg_id, all_object_contents)
                 )
-                inflated_msgs.append(cast(Message, message))
+                message.metadata.__dict__["_message_id"] = msg_id
+                inflated_msgs.append(message)
 
             return inflated_msgs
 

flwr/server/superlink/serverappio/serverappio_servicer.py

@@ -23,7 +23,7 @@ import grpc
 
 from flwr.common import Message
 from flwr.common.constant import SUPERLINK_NODE_ID, Status
-from flwr.common.inflatable import check_body_len_consistency, get_desdendant_object_ids
+from flwr.common.inflatable import check_body_len_consistency, get_descendant_object_ids
 from flwr.common.logger import log
 from flwr.common.serde import (
     context_from_proto,
@@ -202,7 +202,7 @@ class ServerAppIoServicer(serverappio_pb2_grpc.ServerAppIoServicer):
         # Register messages generated by LinkState in the Store for consistency
         for msg_res in messages_res:
             if msg_res.metadata.src_node_id == SUPERLINK_NODE_ID:
-                descendants = list(get_desdendant_object_ids(msg_res))
+                descendants = list(get_descendant_object_ids(msg_res))
                 message_obj_id = msg_res.metadata.message_id
                 # Store mapping
                 store.set_message_descendant_ids(

flwr/supernode/nodestate/in_memory_nodestate.py

@@ -15,12 +15,14 @@
 """In-memory NodeState implementation."""
 
 
+import secrets
 from collections.abc import Sequence
 from dataclasses import dataclass
 from threading import Lock
 from typing import Optional
 
 from flwr.common import Context, Message
+from flwr.common.constant import FLWR_APP_TOKEN_LENGTH
 from flwr.common.typing import Run
 
 from .nodestate import NodeState
@@ -34,7 +36,7 @@ class MessageEntry:
     is_retrieved: bool = False
 
 
-class InMemoryNodeState(NodeState):
+class InMemoryNodeState(NodeState):  # pylint: disable=too-many-instance-attributes
     """In-memory NodeState implementation."""
 
     def __init__(self) -> None:
@@ -49,6 +51,9 @@ class InMemoryNodeState(NodeState):
         # Store run ID to Context mapping
         self.ctx_store: dict[int, Context] = {}
         self.lock_ctx_store = Lock()
+        # Store run ID to token mapping
+        self.token_store: dict[int, str] = {}
+        self.lock_token_store = Lock()
 
     def set_node_id(self, node_id: Optional[int]) -> None:
         """Set the node ID."""
@@ -148,3 +153,38 @@ class InMemoryNodeState(NodeState):
         """Retrieve a context by its run ID."""
         with self.lock_ctx_store:
             return self.ctx_store.get(run_id)
+
+    def get_run_ids_with_pending_messages(self) -> Sequence[int]:
+        """Retrieve run IDs that have at least one pending message."""
+        # Collect run IDs from messages
+        with self.lock_msg_store:
+            ret = {
+                entry.message.metadata.run_id
+                for entry in self.msg_store.values()
+                if entry.message.metadata.reply_to_message_id == ""
+                and not entry.is_retrieved
+            }
+
+        # Remove run IDs that have tokens stored (indicating they are in progress)
+        with self.lock_token_store:
+            ret -= set(self.token_store.keys())
+            return list(ret)
+
+    def create_token(self, run_id: int) -> str:
+        """Create a token for the given run ID."""
+        token = secrets.token_hex(FLWR_APP_TOKEN_LENGTH)  # Generate a random token
+        with self.lock_token_store:
+            if run_id in self.token_store:
+                raise ValueError("Token already created for this run ID")
+            self.token_store[run_id] = token
+        return token
+
+    def verify_token(self, run_id: int, token: str) -> bool:
+        """Verify a token for the given run ID."""
+        with self.lock_token_store:
+            return self.token_store.get(run_id) == token
+
+    def delete_token(self, run_id: int) -> None:
+        """Delete the token for the given run ID."""
+        with self.lock_token_store:
+            self.token_store.pop(run_id, None)

flwr/supernode/nodestate/nodestate.py

@@ -155,3 +155,58 @@ class NodeState(ABC):
         Optional[Context]
             The `Context` instance if found, otherwise None.
         """
+
+    @abstractmethod
+    def get_run_ids_with_pending_messages(self) -> Sequence[int]:
+        """Retrieve run IDs that have at least one pending message.
+
+        Run IDs that are currently in progress (i.e., those associated with tokens)
+        will not be returned, even if they have pending messages.
+
+        Returns
+        -------
+        Sequence[int]
+            Sequence of run IDs with pending messages.
+        """
+
+    @abstractmethod
+    def create_token(self, run_id: int) -> str:
+        """Create a token for the given run ID.
+
+        Parameters
+        ----------
+        run_id : int
+            The ID of the run for which to create a token.
+
+        Returns
+        -------
+        str
+            A unique token associated with the run ID.
+        """
+
+    @abstractmethod
+    def verify_token(self, run_id: int, token: str) -> bool:
+        """Verify a token for the given run ID.
+
+        Parameters
+        ----------
+        run_id : int
+            The ID of the run for which to verify the token.
+        token : str
+            The token to verify.
+
+        Returns
+        -------
+        bool
+            True if the token is valid for the run ID, False otherwise.
+        """
+
+    @abstractmethod
+    def delete_token(self, run_id: int) -> None:
+        """Delete the token for the given run ID.
+
+        Parameters
+        ----------
+        run_id : int
+            The ID of the run for which to delete the token.
+        """

{flwr_nightly-1.19.0.dev20250605.dist-info → flwr_nightly-1.19.0.dev20250607.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.3
 Name: flwr-nightly
-Version: 1.19.0.dev20250605
+Version: 1.19.0.dev20250607
 Summary: Flower: A Friendly Federated AI Framework
 License: Apache-2.0
 Keywords: Artificial Intelligence,Federated AI,Federated Analytics,Federated Evaluation,Federated Learning,Flower,Machine Learning

{flwr_nightly-1.19.0.dev20250605.dist-info → flwr_nightly-1.19.0.dev20250607.dist-info}/RECORD

@@ -84,7 +84,7 @@ flwr/client/grpc_adapter_client/__init__.py,sha256=RQWP5mFPROLHKgombiRvPXVWSoVrQ
 flwr/client/grpc_adapter_client/connection.py,sha256=aj5tTYyE8z2hQLXPPydsJiz8gBDIWLUhfWvqYkAL1L4,3966
 flwr/client/grpc_rere_client/__init__.py,sha256=i7iS0Lt8B7q0E2L72e4F_YrKm6ClRKnd71PNA6PW2O0,752
 flwr/client/grpc_rere_client/client_interceptor.py,sha256=zFaVHw6AxeNO-7eCKKb-RxrPa7zbM5Z-2-1Efc4adQY,2451
-flwr/client/grpc_rere_client/connection.py,sha256=XwhT-yrWLHHegXpiAu1YuICLQ1t-KkOh6pNYJ9R4FEE,13358
+flwr/client/grpc_rere_client/connection.py,sha256=nk4xSjBKG_mvoZsCVkkGLPPc6juxBId7RNcckgA6VSE,13994
 flwr/client/grpc_rere_client/grpc_adapter.py,sha256=JvMZ7vCFTaTEo6AzKYh3zDmeQAU7VSjdysbC6t3ufWg,6351
 flwr/client/message_handler/__init__.py,sha256=0lyljDVqre3WljiZbPcwCCf8GiIaSVI_yo_ylEyPwSE,719
 flwr/client/message_handler/message_handler.py,sha256=X9SXX6et97Lw9_DGD93HKsEBGNjXClcFgc_5aLK0oiU,6541
@@ -108,7 +108,7 @@ flwr/common/args.py,sha256=-aX_jVnSaDrJR2KZ8Wq0Y3dQHII4R4MJtJOIXzVUA0c,5417
 flwr/common/auth_plugin/__init__.py,sha256=3rzPkVLn9WyB5n7HLk1XGDw3SLCqRWAU1_CnglcWPfw,970
 flwr/common/auth_plugin/auth_plugin.py,sha256=yEIapwim14tn99oGhkQKun_WifoY8hscAtU8Z6fcGXE,4796
 flwr/common/config.py,sha256=glcZDjco-amw1YfQcYTFJ4S1pt9APoexT-mf1QscuHs,13960
-flwr/common/constant.py,sha256=lQCel7YanrEdFzy7pHBejhrQ7ZIylJrhBD9pXY02qv0,7472
+flwr/common/constant.py,sha256=A8rWHZZUH4Rxbx-bCKW5pNJ_BG4W9Xfw4fLevNhZ8l0,8077
 flwr/common/context.py,sha256=Be8obQR_OvEDy1OmshuUKxGRQ7Qx89mf5F4xlhkR10s,2407
 flwr/common/date.py,sha256=1ZT2cRSpC2DJqprOVTLXYCR_O2_OZR0zXO_brJ3LqWc,1554
 flwr/common/differential_privacy.py,sha256=FdlpdpPl_H_2HJa8CQM1iCUGBBQ5Dc8CzxmHERM-EoE,6148
@@ -122,10 +122,11 @@ flwr/common/exit/exit_code.py,sha256=PNEnCrZfOILjfDAFu5m-2YWEJBrk97xglq4zCUlqV7E
 flwr/common/exit_handlers.py,sha256=IaqJ60fXZuu7McaRYnoYKtlbH9t4Yl9goNExKqtmQbs,4304
 flwr/common/grpc.py,sha256=manTaHaPiyYngUq1ErZvvV2B2GxlXUUUGRy3jc3TBIQ,9798
 flwr/common/heartbeat.py,sha256=SyEpNDnmJ0lni0cWO67rcoJVKasCLmkNHm3dKLeNrLU,5749
-flwr/common/inflatable.py,sha256=9yPsSFOfNM2OIb15JQ6-wY5kblwXiC5zNX-tsp2ZwW0,7017
-flwr/common/inflatable_grpc_utils.py,sha256=YGP8oJRfnkwvY6segWH1DUf_ljDIkku7-2zH66tv3HA,4337
+flwr/common/inflatable.py,sha256=b8Br0TKB_yMijbICPVD4gbGMr4l0iXtH9s8MztTry4A,7717
+flwr/common/inflatable_grpc_utils.py,sha256=M982VQ8x1wh4QnLXR-IKWdtGuOFACvIsU6lYbbs1TI8,13040
+flwr/common/inflatable_utils.py,sha256=hPwap-P9pxHcxo-EtxtrnKNiFA2HsNGJaKX2b002PnE,2892
 flwr/common/logger.py,sha256=JbRf6E2vQxXzpDBq1T8IDUJo_usu3gjWEBPQ6uKcmdg,13049
-flwr/common/message.py,sha256=ZdH35PznYhIzfNKPzHTL1YS5z-flTpHnjdNpQR6B8x0,19953
+flwr/common/message.py,sha256=xAL7iZN5-n-xPQpgoSFvxNrzs8fmiiPfoU0DjNQEhRw,19953
 flwr/common/object_ref.py,sha256=p3SfTeqo3Aj16SkB-vsnNn01zswOPdGNBitcbRnqmUk,9134
 flwr/common/parameter.py,sha256=UVw6sOgehEFhFs4uUCMl2kfVq1PD6ncmWgPLMsZPKPE,2095
 flwr/common/pyproject.py,sha256=2SU6yJW7059SbMXgzjOdK1GZRWO6AixDH7BmdxbMvHI,1386
@@ -236,7 +237,7 @@ flwr/server/criterion.py,sha256=G4e-6B48Pc7d5rmGVUpIzNKb6UF88O3VmTRuUltgjzM,1061
 flwr/server/fleet_event_log_interceptor.py,sha256=AkL7Y5d3xm2vRhL3ahmEVVoOvAP7PA7dRgB-je4v-Ys,3774
 flwr/server/grid/__init__.py,sha256=aWZHezoR2UGMJISB_gPMCm2N_2GSbm97A3lAp7ruhRQ,888
 flwr/server/grid/grid.py,sha256=naGCYt5J6dnmUvrcGkdNyKPe3MBd-0awGm1ALmgahqY,6625
-flwr/server/grid/grpc_grid.py,sha256=ymLaqWYYkv0uWrYvxavMKEyB0TVRMSPi1tJhmvhe7_g,12392
+flwr/server/grid/grpc_grid.py,sha256=rVzPh0BGQOI3Dzhd8upBXNLEA5bKy82eBZkU4a1qFJY,13257
 flwr/server/grid/inmemory_grid.py,sha256=RjejYT-d-hHuTs1KSs_5wvOdAWKLus8w5_UAcnGt4iw,6168
 flwr/server/history.py,sha256=cCkFhBN4GoHsYYNk5GG1Y089eKJh2DH_ZJbYPwLaGyk,5026
 flwr/server/run_serverapp.py,sha256=v0p6jXj2dFxlRUdoEeF1mnaFd9XRQi6dZCflPY6d3qI,2063
@@ -303,7 +304,7 @@ flwr/server/superlink/linkstate/sqlite_linkstate.py,sha256=sHJPK1w0tP0m2WCXH2F9l
 flwr/server/superlink/linkstate/utils.py,sha256=IeLh7iGRCHU5MEWOl7iriaSE4L__8GWOa2OleXadK5M,15444
 flwr/server/superlink/serverappio/__init__.py,sha256=Fy4zJuoccZe5mZSEIpOmQvU6YeXFBa1M4eZuXXmJcn8,717
 flwr/server/superlink/serverappio/serverappio_grpc.py,sha256=6-FUUt0GiLcBPljj8bBrUNeAITUoDQOLzaMihKo52hg,2326
-flwr/server/superlink/serverappio/serverappio_servicer.py,sha256=qInBXn7xcnNUNIXj_BkjoWfZd96By55gbTsp4onwfDQ,17290
+flwr/server/superlink/serverappio/serverappio_servicer.py,sha256=GElRmXPxkrbnqD9lMVvJJHQF5hYJODL8uhAvvI97ItQ,17290
 flwr/server/superlink/simulation/__init__.py,sha256=Ry8DrNaZCMcQXvUc4FoCN2m3dvUQgWjasfp015o3Ec4,718
 flwr/server/superlink/simulation/simulationio_grpc.py,sha256=0l0F-UjYEk6W7HZmI28PbJQLFxSi_vBHRkdchgdaSMQ,2224
 flwr/server/superlink/simulation/simulationio_servicer.py,sha256=aJezU8RSJswcmWm7Eoy0BqsU13jrcfuFwX3ljm-cORM,7719
@@ -348,8 +349,8 @@ flwr/supernode/cli/__init__.py,sha256=JuEMr0-s9zv-PEWKuLB9tj1ocNfroSyNJ-oyv7ati9
 flwr/supernode/cli/flower_supernode.py,sha256=ly2AQhbla2sufDaMsENaEALDEd0a4CS4D0eUrUOkHzY,8778
 flwr/supernode/cli/flwr_clientapp.py,sha256=KfVUO20ZMnUDSGZTJ9I1KkMawFsRV6kdRUmGIRNbg_8,2812
 flwr/supernode/nodestate/__init__.py,sha256=CyLLObbmmVgfRO88UCM0VMait1dL57mUauUDfuSHsbU,976
-flwr/supernode/nodestate/in_memory_nodestate.py,sha256=4ZiLA45fMi2bJgmfDNLtiv-gVNru95Bi48xBy7xtatA,5212
-flwr/supernode/nodestate/nodestate.py,sha256=SgblnKtqzTHRiODwg4QUREw1-uYPQrLzoeTBlROHf_0,4571
+flwr/supernode/nodestate/in_memory_nodestate.py,sha256=vfHgG7dxgVLo0ar4ltLOveEm7YVvXxdj0Cac4CE85qQ,6903
+flwr/supernode/nodestate/nodestate.py,sha256=Ro60vlm4xUm6FGNV440DrlzZ6rPsem-ULmJg0eAFmhU,6034
 flwr/supernode/nodestate/nodestate_factory.py,sha256=UYTDCcwK_baHUmkzkJDxL0UEqvtTfOMlQRrROMCd0Xo,1430
 flwr/supernode/runtime/__init__.py,sha256=JQdqd2EMTn-ORMeTvewYYh52ls0YKP68jrps1qioxu4,718
 flwr/supernode/runtime/run_clientapp.py,sha256=cvWSby7u31u97QapWHxJM-Wer6F1k6mbbD-d1gxwxZA,7962
@@ -357,7 +358,7 @@ flwr/supernode/servicer/__init__.py,sha256=lucTzre5WPK7G1YLCfaqg3rbFWdNSb7ZTt-ca
 flwr/supernode/servicer/clientappio/__init__.py,sha256=vJyOjO2FXZ2URbnthmdsgs6948wbYfdq1L1V8Um-Lr8,895
 flwr/supernode/servicer/clientappio/clientappio_servicer.py,sha256=LmzkxtNQBn5vVrHc0Bhq2WqaK6-LM2v4kfLBN0PiNNM,8522
 flwr/supernode/start_client_internal.py,sha256=5CwTNV-XmIhwR1jv3G7aQAXGhf6OFWS6U-vmxY1iKGA,16984
-flwr_nightly-1.19.0.dev20250605.dist-info/METADATA,sha256=gsyzTuVl8GpQD7gk1RPBCbNvdugpyl08gCr1sC0kUHc,15910
-flwr_nightly-1.19.0.dev20250605.dist-info/WHEEL,sha256=b4K_helf-jlQoXBBETfwnf4B04YC67LOev0jo4fX5m8,88
-flwr_nightly-1.19.0.dev20250605.dist-info/entry_points.txt,sha256=jNpDXGBGgs21RqUxelF_jwGaxtqFwm-MQyfz-ZqSjrA,367
-flwr_nightly-1.19.0.dev20250605.dist-info/RECORD,,
+flwr_nightly-1.19.0.dev20250607.dist-info/METADATA,sha256=YhYf05eYMhWNR6GdMIdEh39UZLZMmlX3x0MifwkcURo,15910
+flwr_nightly-1.19.0.dev20250607.dist-info/WHEEL,sha256=b4K_helf-jlQoXBBETfwnf4B04YC67LOev0jo4fX5m8,88
+flwr_nightly-1.19.0.dev20250607.dist-info/entry_points.txt,sha256=jNpDXGBGgs21RqUxelF_jwGaxtqFwm-MQyfz-ZqSjrA,367
+flwr_nightly-1.19.0.dev20250607.dist-info/RECORD,,