flwr-nightly 1.19.0.dev20250604__tar.gz → 1.19.0.dev20250606__tar.gz

{flwr_nightly-1.19.0.dev20250604 → flwr_nightly-1.19.0.dev20250606}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.3
 Name: flwr-nightly
-Version: 1.19.0.dev20250604
+Version: 1.19.0.dev20250606
 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.dev20250604 → flwr_nightly-1.19.0.dev20250606}/py/flwr/client/grpc_rere_client/connection.py

@@ -31,11 +31,17 @@ from flwr.common.constant import HEARTBEAT_CALL_TIMEOUT, HEARTBEAT_DEFAULT_INTER
 from flwr.common.grpc import create_channel, on_channel_state_change
 from flwr.common.heartbeat import HeartbeatSender
 from flwr.common.inflatable_grpc_utils import (
+    make_pull_object_fn_grpc,
+    make_push_object_fn_grpc,
     pull_object_from_servicer,
     push_object_to_servicer,
 )
 from flwr.common.logger import log
-from flwr.common.message import Message, get_message_to_descendant_id_mapping
+from flwr.common.message import (
+    Message,
+    get_message_to_descendant_id_mapping,
+    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,
@@ -265,9 +271,11 @@ def grpc_request_response(  # pylint: disable=R0913,R0914,R0915,R0917
                 Message,
                 pull_object_from_servicer(
                     object_id=message_proto.metadata.message_id,
-                    stub=stub,
-                    node=node,
-                    run_id=message_proto.metadata.run_id,
+                    pull_object_fn=make_pull_object_fn_grpc(
+                        pull_object_grpc=stub.PullObject,
+                        node=node,
+                        run_id=message_proto.metadata.run_id,
+                    ),
                 ),
             )
 
@@ -305,7 +313,7 @@ def grpc_request_response(  # pylint: disable=R0913,R0914,R0915,R0917
             return
 
         # Serialize Message
-        message_proto = message_to_proto(message=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,
@@ -318,9 +326,11 @@ def grpc_request_response(  # pylint: disable=R0913,R0914,R0915,R0917
             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_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_nightly-1.19.0.dev20250606/py/flwr/common/inflatable_grpc_utils.py

@@ -0,0 +1,227 @@
+# 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 utils."""
+
+
+from time import sleep
+from typing import Callable, Optional
+
+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 import (
+    InflatableObject,
+    get_object_head_values_from_object_content,
+    get_object_id,
+)
+from .message import Message
+from .record import Array, ArrayRecord, ConfigRecord, MetricRecord, RecordDict
+
+# Helper registry that maps names of classes to their type
+inflatable_class_registry: dict[str, type[InflatableObject]] = {
+    Array.__qualname__: Array,
+    ArrayRecord.__qualname__: ArrayRecord,
+    ConfigRecord.__qualname__: ConfigRecord,
+    Message.__qualname__: Message,
+    MetricRecord.__qualname__: MetricRecord,
+    RecordDict.__qualname__: RecordDict,
+}
+
+
+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 push_object_to_servicer(
+    obj: InflatableObject,
+    push_object_fn: Callable[[str, bytes], None],
+    object_ids_to_push: Optional[set[str]] = None,
+) -> set[str]:
+    """Recursively deflate an object and push it to the servicer.
+
+    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
+    ----------
+    obj : InflatableObject
+        The object to push.
+    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 object ID of the given object is not in this
+        set, the object will not be pushed.
+
+    Returns
+    -------
+    set[str]
+        A set of object IDs that were pushed to the servicer.
+    """
+    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, push_object_fn, 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:
+        # The function may raise an error if the object ID is not pre-registered
+        push_object_fn(object_id, object_content)
+        pushed_object_ids.add(object_id)
+
+    return pushed_object_ids
+
+
+def pull_object_from_servicer(
+    object_id: str,
+    pull_object_fn: Callable[[str], bytes],
+) -> InflatableObject:
+    """Recursively inflate an object by pulling it from the servicer.
+
+    Parameters
+    ----------
+    object_id : str
+        The ID of the object 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.
+
+    Returns
+    -------
+    InflatableObject
+        The pulled object.
+    """
+    # Pull object
+    while True:
+        try:
+            # The function may raise an error if the object ID is not pre-registered
+            object_content: bytes = pull_object_fn(object_id)
+            break  # Exit loop if object is successfully pulled
+        except ObjectUnavailableError:
+            sleep(0.1)  # Retry after a short delay
+
+    # Extract object class and object_ids of children
+    obj_type, children_obj_ids, _ = get_object_head_values_from_object_content(
+        object_content=object_content
+    )
+    # Resolve object class
+    cls_type = inflatable_class_registry[obj_type]
+
+    # Pull 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, pull_object_fn
+        )
+
+    # Inflate object passing its children
+    return cls_type.inflate(object_content, children=children)
+
+
+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

{flwr_nightly-1.19.0.dev20250604 → flwr_nightly-1.19.0.dev20250606}/py/flwr/common/message.py

@@ -426,6 +426,17 @@ def make_message(
     return Message(metadata=metadata, content=content, error=error)  # type: ignore
 
 
+def remove_content_from_message(message: Message) -> Message:
+    """Return a copy of the Message but with an empty RecordDict as content.
+
+    If message has no content, it returns itself.
+    """
+    if message.has_error():
+        return message
+
+    return make_message(metadata=message.metadata, content=RecordDict())
+
+
 def _limit_reply_ttl(
     current: float, reply_ttl: float | None, reply_to: Message
 ) -> float:

{flwr_nightly-1.19.0.dev20250604 → flwr_nightly-1.19.0.dev20250606}/py/flwr/server/compat/app_utils.py

@@ -16,6 +16,7 @@
 
 
 import threading
+from typing import Any, Callable
 
 from flwr.common.typing import RunNotRunningException
 
@@ -80,36 +81,57 @@ def _update_client_manager(
     """Update the nodes list in the client manager."""
     # Loop until the grid is disconnected
     registered_nodes: dict[int, GridClientProxy] = {}
+    lock = threading.RLock()
+
+    def update_registered_nodes() -> None:
+        with lock:
+            all_node_ids = set(grid.get_node_ids())
+            dead_nodes = set(registered_nodes).difference(all_node_ids)
+            new_nodes = all_node_ids.difference(registered_nodes)
+
+            # Unregister dead nodes
+            for node_id in dead_nodes:
+                client_proxy = registered_nodes[node_id]
+                client_manager.unregister(client_proxy)
+                del registered_nodes[node_id]
+
+            # Register new nodes
+            for node_id in new_nodes:
+                client_proxy = GridClientProxy(
+                    node_id=node_id,
+                    grid=grid,
+                    run_id=grid.run.run_id,
+                )
+                if client_manager.register(client_proxy):
+                    registered_nodes[node_id] = client_proxy
+                else:
+                    raise RuntimeError("Could not register node.")
+
+    # Get the wrapped method of ClientManager instance
+    def get_wrapped_method(method_name: str) -> Callable[..., Any]:
+        original_method = getattr(client_manager, method_name)
+
+        def wrapped_method(*args: Any, **kwargs: Any) -> Any:
+            # Update registered nodes before calling the original method
+            update_registered_nodes()
+            return original_method(*args, **kwargs)
+
+        return wrapped_method
+
+    # Wrap the ClientManager
+    for method_name in ["num_available", "all", "sample"]:
+        setattr(client_manager, method_name, get_wrapped_method(method_name))
+
+    c_done.set()
+
     while not f_stop.is_set():
+        # Sleep for 5 seconds
+        if not f_stop.is_set():
+            f_stop.wait(5)
+
         try:
-            all_node_ids = set(grid.get_node_ids())
+            # Update registered nodes
+            update_registered_nodes()
         except RunNotRunningException:
             f_stop.set()
             break
-        dead_nodes = set(registered_nodes).difference(all_node_ids)
-        new_nodes = all_node_ids.difference(registered_nodes)
-
-        # Unregister dead nodes
-        for node_id in dead_nodes:
-            client_proxy = registered_nodes[node_id]
-            client_manager.unregister(client_proxy)
-            del registered_nodes[node_id]
-
-        # Register new nodes
-        for node_id in new_nodes:
-            client_proxy = GridClientProxy(
-                node_id=node_id,
-                grid=grid,
-                run_id=grid.run.run_id,
-            )
-            if client_manager.register(client_proxy):
-                registered_nodes[node_id] = client_proxy
-            else:
-                raise RuntimeError("Could not register node.")
-
-        # Flag first pass for nodes registration is completed
-        c_done.set()
-
-        # Sleep for 3 seconds
-        if not f_stop.is_set():
-            f_stop.wait(3)

{flwr_nightly-1.19.0.dev20250604 → flwr_nightly-1.19.0.dev20250606}/py/flwr/server/grid/grpc_grid.py

@@ -29,11 +29,16 @@ from flwr.common.constant import (
 )
 from flwr.common.grpc import create_channel, on_channel_state_change
 from flwr.common.inflatable_grpc_utils import (
+    make_pull_object_fn_grpc,
+    make_push_object_fn_grpc,
     pull_object_from_servicer,
     push_object_to_servicer,
 )
 from flwr.common.logger import log, warn_deprecated_feature
-from flwr.common.message import get_message_to_descendant_id_mapping
+from flwr.common.message import (
+    get_message_to_descendant_id_mapping,
+    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
@@ -210,7 +215,7 @@ class GrpcGrid(Grid):
         # Call GrpcServerAppIoStub method
         res: PushInsMessagesResponse = self._stub.PushMessages(
             PushInsMessagesRequest(
-                messages_list=[message_to_proto(message)],
+                messages_list=[message_to_proto(remove_content_from_message(message))],
                 run_id=run_id,
                 msg_to_descendant_mapping=descendants_mapping,
             )
@@ -224,9 +229,11 @@ class GrpcGrid(Grid):
             # Push only object that are not in the store
             push_object_to_servicer(
                 message,
-                self._stub,
-                node=self.node,
-                run_id=run_id,
+                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,9 +296,11 @@ class GrpcGrid(Grid):
 
                 message = pull_object_from_servicer(
                     msg_proto.metadata.message_id,
-                    self._stub,
-                    node=self.node,
-                    run_id=run_id,
+                    pull_object_fn=make_pull_object_fn_grpc(
+                        pull_object_grpc=self._stub.PullObject,
+                        node=self.node,
+                        run_id=run_id,
+                    ),
                 )
                 inflated_msgs.append(cast(Message, message))
 

{flwr_nightly-1.19.0.dev20250604 → flwr_nightly-1.19.0.dev20250606}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "poetry.core.masonry.api"
 
 [tool.poetry]
 name = "flwr-nightly"
-version = "1.19.0.dev20250604"
+version = "1.19.0.dev20250606"
 description = "Flower: A Friendly Federated AI Framework"
 license = "Apache-2.0"
 authors = ["The Flower Authors <hello@flower.ai>"]

flwr_nightly-1.19.0.dev20250604/py/flwr/common/inflatable_grpc_utils.py

@@ -1,122 +0,0 @@
-# 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 utils."""
-
-
-from time import sleep
-from typing import Optional, Union
-
-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,
-    PushObjectRequest,
-    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 .message import Message
-from .record import Array, ArrayRecord, ConfigRecord, MetricRecord, RecordDict
-
-# Helper registry that maps names of classes to their type
-inflatable_class_registry: dict[str, type[InflatableObject]] = {
-    Array.__qualname__: Array,
-    ArrayRecord.__qualname__: ArrayRecord,
-    ConfigRecord.__qualname__: ConfigRecord,
-    Message.__qualname__: Message,
-    MetricRecord.__qualname__: MetricRecord,
-    RecordDict.__qualname__: RecordDict,
-}
-
-
-def push_object_to_servicer(
-    obj: InflatableObject,
-    stub: Union[FleetStub, ServerAppIoStub, GrpcAdapter],
-    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.
-
-    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.
-    """
-    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_from_servicer(
-    object_id: str,
-    stub: Union[FleetStub, ServerAppIoStub, GrpcAdapter],
-    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)
-        )
-        object_available = object_proto.object_available
-        object_content = object_proto.object_content
-        sleep(0.1)
-
-    # Extract object class and object_ids of children
-    obj_type, children_obj_ids, _ = get_object_head_values_from_object_content(
-        object_content=object_content
-    )
-    # Resolve object class
-    cls_type = inflatable_class_registry[obj_type]
-
-    # Pull 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
-        )
-
-    # Inflate object passing its children
-    return cls_type.inflate(object_content, children=children)