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

flwr/client/grpc_rere_client/connection.py

@@ -15,47 +15,61 @@
 """Contextmanager for a gRPC request-response channel to the Flower server."""
 
 
-import random
-import threading
 from collections.abc import Iterator, Sequence
 from contextlib import contextmanager
 from copy import copy
-from logging import ERROR
+from logging import DEBUG, ERROR
 from pathlib import Path
 from typing import Callable, Optional, Union, cast
 
 import grpc
 from cryptography.hazmat.primitives.asymmetric import ec
 
-from flwr.client.heartbeat import start_ping_loop
+from flwr.app.metadata import Metadata
 from flwr.client.message_handler.message_handler import validate_out_message
 from flwr.common import GRPC_MAX_MESSAGE_LENGTH
-from flwr.common.constant import (
-    PING_BASE_MULTIPLIER,
-    PING_CALL_TIMEOUT,
-    PING_DEFAULT_INTERVAL,
-    PING_RANDOM_RANGE,
-)
+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,
+    get_object_tree,
+    no_object_id_recompute,
+)
+from flwr.common.inflatable_grpc_utils import (
+    make_pull_object_fn_grpc,
+    make_push_object_fn_grpc,
+)
+from flwr.common.inflatable_utils import (
+    inflate_object_from_contents,
+    pull_objects,
+    push_objects,
+)
 from flwr.common.logger import log
-from flwr.common.message import Message, Metadata
-from flwr.common.retry_invoker import RetryInvoker
+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,
 )
-from flwr.common.serde import message_from_proto, message_to_proto, run_from_proto
+from flwr.common.serde import message_to_proto, run_from_proto
 from flwr.common.typing import Fab, Run, RunNotRunningException
 from flwr.proto.fab_pb2 import GetFabRequest, GetFabResponse  # pylint: disable=E0611
 from flwr.proto.fleet_pb2 import (  # pylint: disable=E0611
     CreateNodeRequest,
     DeleteNodeRequest,
-    PingRequest,
-    PingResponse,
     PullMessagesRequest,
     PullMessagesResponse,
     PushMessagesRequest,
+    PushMessagesResponse,
 )
 from flwr.proto.fleet_pb2_grpc import FleetStub  # pylint: disable=E0611
+from flwr.proto.heartbeat_pb2 import (  # pylint: disable=E0611
+    SendNodeHeartbeatRequest,
+    SendNodeHeartbeatResponse,
+)
+from flwr.proto.message_pb2 import (  # pylint: disable=E0611
+    ConfirmMessageReceivedRequest,
+)
 from flwr.proto.node_pb2 import Node  # pylint: disable=E0611
 from flwr.proto.run_pb2 import GetRunRequest, GetRunResponse  # pylint: disable=E0611
 
@@ -78,10 +92,10 @@ def grpc_request_response(  # pylint: disable=R0913,R0914,R0915,R0917
     tuple[
         Callable[[], Optional[Message]],
         Callable[[Message], None],
-        Optional[Callable[[], Optional[int]]],
-        Optional[Callable[[], None]],
-        Optional[Callable[[int], Run]],
-        Optional[Callable[[str, int], Fab]],
+        Callable[[], Optional[int]],
+        Callable[[], None],
+        Callable[[int], Run],
+        Callable[[str, int], Fab],
     ]
 ]:
     """Primitives for request/response-based interaction with a server.
@@ -151,8 +165,6 @@ def grpc_request_response(  # pylint: disable=R0913,R0914,R0915,R0917
     stub = adapter_cls(channel)
     metadata: Optional[Metadata] = None
     node: Optional[Node] = None
-    ping_thread: Optional[threading.Thread] = None
-    ping_stop_event = threading.Event()
 
     def _should_giveup_fn(e: Exception) -> bool:
         if e.code() == grpc.StatusCode.PERMISSION_DENIED:  # type: ignore
@@ -165,46 +177,58 @@ def grpc_request_response(  # pylint: disable=R0913,R0914,R0915,R0917
     # If the status code is PERMISSION_DENIED, additionally raise RunNotRunningException
     retry_invoker.should_giveup = _should_giveup_fn
 
+    # Wrap stub
+    _wrap_stub(stub, retry_invoker)
     ###########################################################################
-    # ping/create_node/delete_node/receive/send/get_run functions
+    # send_node_heartbeat/create_node/delete_node/receive/send/get_run functions
     ###########################################################################
 
-    def ping() -> None:
+    def send_node_heartbeat() -> bool:
         # Get Node
         if node is None:
             log(ERROR, "Node instance missing")
-            return
+            return False
 
-        # Construct the ping request
-        req = PingRequest(node=node, ping_interval=PING_DEFAULT_INTERVAL)
+        # Construct the heartbeat request
+        req = SendNodeHeartbeatRequest(
+            node=node, heartbeat_interval=HEARTBEAT_DEFAULT_INTERVAL
+        )
 
         # Call FleetAPI
-        res: PingResponse = stub.Ping(req, timeout=PING_CALL_TIMEOUT)
+        try:
+            res: SendNodeHeartbeatResponse = stub.SendNodeHeartbeat(
+                req, timeout=HEARTBEAT_CALL_TIMEOUT
+            )
+        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 success
         if not res.success:
-            raise RuntimeError("Ping failed unexpectedly.")
+            raise RuntimeError(
+                "Heartbeat failed unexpectedly. The SuperLink does not "
+                "recognize this SuperNode."
+            )
+        return True
 
-        # Wait
-        rd = random.uniform(*PING_RANDOM_RANGE)
-        next_interval: float = PING_DEFAULT_INTERVAL - PING_CALL_TIMEOUT
-        next_interval *= PING_BASE_MULTIPLIER + rd
-        if not ping_stop_event.is_set():
-            ping_stop_event.wait(next_interval)
+    heartbeat_sender = HeartbeatSender(send_node_heartbeat)
 
     def create_node() -> Optional[int]:
         """Set create_node."""
         # Call FleetAPI
-        create_node_request = CreateNodeRequest(ping_interval=PING_DEFAULT_INTERVAL)
-        create_node_response = retry_invoker.invoke(
-            stub.CreateNode,
-            request=create_node_request,
+        create_node_request = CreateNodeRequest(
+            heartbeat_interval=HEARTBEAT_DEFAULT_INTERVAL
         )
+        create_node_response = stub.CreateNode(request=create_node_request)
 
-        # Remember the node and the ping-loop thread
-        nonlocal node, ping_thread
+        # Remember the node and start the heartbeat sender
+        nonlocal node
         node = cast(Node, create_node_response.node)
-        ping_thread = start_ping_loop(ping, ping_stop_event)
+        heartbeat_sender.start()
         return node.node_id
 
     def delete_node() -> None:
@@ -215,12 +239,12 @@ def grpc_request_response(  # pylint: disable=R0913,R0914,R0915,R0917
             log(ERROR, "Node instance missing")
             return
 
-        # Stop the ping-loop thread
-        ping_stop_event.set()
+        # Stop the heartbeat sender
+        heartbeat_sender.stop()
 
         # Call FleetAPI
         delete_node_request = DeleteNodeRequest(node=node)
-        retry_invoker.invoke(stub.DeleteNode, request=delete_node_request)
+        stub.DeleteNode(request=delete_node_request)
 
         # Cleanup
         node = None
@@ -234,9 +258,7 @@ def grpc_request_response(  # pylint: disable=R0913,R0914,R0915,R0917
 
         # Request instructions (message) from server
         request = PullMessagesRequest(node=node)
-        response: PullMessagesResponse = retry_invoker.invoke(
-            stub.PullMessages, request=request
-        )
+        response: PullMessagesResponse = stub.PullMessages(request=request)
 
         # Get the current Messages
         message_proto = (
@@ -250,7 +272,33 @@ def grpc_request_response(  # pylint: disable=R0913,R0914,R0915,R0917
             message_proto = None
 
         # Construct the Message
-        in_message = message_from_proto(message_proto) if message_proto else None
+        in_message: Optional[Message] = None
+
+        if message_proto:
+            msg_id = message_proto.metadata.message_id
+            run_id = message_proto.metadata.run_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=run_id,
+                ),
+            )
+
+            # Confirm that the message has been received
+            stub.ConfirmMessageReceived(
+                ConfirmMessageReceivedRequest(
+                    node=node, run_id=run_id, message_object_id=msg_id
+                )
+            )
+
+            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
+            in_message.metadata.__dict__["_message_id"] = msg_id
 
         # Remember `metadata` of the in message
         nonlocal metadata
@@ -272,15 +320,43 @@ def grpc_request_response(  # pylint: disable=R0913,R0914,R0915,R0917
             log(ERROR, "No current message")
             return
 
+        # Set message_id
+        message.metadata.__dict__["_message_id"] = message.object_id
         # Validate out message
         if not validate_out_message(message, metadata):
             log(ERROR, "Invalid out message")
             return
 
-        # Serialize Message
-        message_proto = message_to_proto(message=message)
-        request = PushMessagesRequest(node=node, messages_list=[message_proto])
-        _ = retry_invoker.invoke(stub.PushMessages, request)
+        with no_object_id_recompute():
+            # Get all nested objects
+            all_objects = get_all_nested_objects(message)
+            object_tree = get_object_tree(message)
+
+            # Serialize Message
+            message_proto = message_to_proto(
+                message=remove_content_from_message(message)
+            )
+            request = PushMessagesRequest(
+                node=node,
+                messages_list=[message_proto],
+                message_object_trees=[object_tree],
+            )
+            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_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))
 
         # Cleanup
         metadata = None
@@ -288,10 +364,7 @@ def grpc_request_response(  # pylint: disable=R0913,R0914,R0915,R0917
     def get_run(run_id: int) -> Run:
         # Call FleetAPI
         get_run_request = GetRunRequest(node=node, run_id=run_id)
-        get_run_response: GetRunResponse = retry_invoker.invoke(
-            stub.GetRun,
-            request=get_run_request,
-        )
+        get_run_response: GetRunResponse = stub.GetRun(request=get_run_request)
 
         # Return fab_id and fab_version
         return run_from_proto(get_run_response.run)
@@ -299,10 +372,7 @@ def grpc_request_response(  # pylint: disable=R0913,R0914,R0915,R0917
     def get_fab(fab_hash: str, run_id: int) -> Fab:
         # Call FleetAPI
         get_fab_request = GetFabRequest(node=node, hash_str=fab_hash, run_id=run_id)
-        get_fab_response: GetFabResponse = retry_invoker.invoke(
-            stub.GetFab,
-            request=get_fab_request,
-        )
+        get_fab_response: GetFabResponse = stub.GetFab(request=get_fab_request)
 
         return Fab(get_fab_response.fab.hash_str, get_fab_response.fab.content)
 

flwr/client/grpc_rere_client/grpc_adapter.py

@@ -38,8 +38,6 @@ from flwr.proto.fleet_pb2 import (  # pylint: disable=E0611
     CreateNodeResponse,
     DeleteNodeRequest,
     DeleteNodeResponse,
-    PingRequest,
-    PingResponse,
     PullMessagesRequest,
     PullMessagesResponse,
     PushMessagesRequest,
@@ -47,6 +45,18 @@ from flwr.proto.fleet_pb2 import (  # pylint: disable=E0611
 )
 from flwr.proto.grpcadapter_pb2 import MessageContainer  # pylint: disable=E0611
 from flwr.proto.grpcadapter_pb2_grpc import GrpcAdapterStub
+from flwr.proto.heartbeat_pb2 import (  # pylint: disable=E0611
+    SendNodeHeartbeatRequest,
+    SendNodeHeartbeatResponse,
+)
+from flwr.proto.message_pb2 import (  # pylint: disable=E0611
+    ConfirmMessageReceivedRequest,
+    ConfirmMessageReceivedResponse,
+    PullObjectRequest,
+    PullObjectResponse,
+    PushObjectRequest,
+    PushObjectResponse,
+)
 from flwr.proto.run_pb2 import GetRunRequest, GetRunResponse  # pylint: disable=E0611
 
 T = TypeVar("T", bound=GrpcMessage)
@@ -120,11 +130,11 @@ class GrpcAdapter:
         """."""
         return self._send_and_receive(request, DeleteNodeResponse, **kwargs)
 
-    def Ping(  # pylint: disable=C0103
-        self, request: PingRequest, **kwargs: Any
-    ) -> PingResponse:
+    def SendNodeHeartbeat(  # pylint: disable=C0103
+        self, request: SendNodeHeartbeatRequest, **kwargs: Any
+    ) -> SendNodeHeartbeatResponse:
         """."""
-        return self._send_and_receive(request, PingResponse, **kwargs)
+        return self._send_and_receive(request, SendNodeHeartbeatResponse, **kwargs)
 
     def PullMessages(  # pylint: disable=C0103
         self, request: PullMessagesRequest, **kwargs: Any
@@ -149,3 +159,21 @@ class GrpcAdapter:
     ) -> GetFabResponse:
         """."""
         return self._send_and_receive(request, GetFabResponse, **kwargs)
+
+    def PushObject(  # pylint: disable=C0103
+        self, request: PushObjectRequest, **kwargs: Any
+    ) -> PushObjectResponse:
+        """."""
+        return self._send_and_receive(request, PushObjectResponse, **kwargs)
+
+    def PullObject(  # pylint: disable=C0103
+        self, request: PullObjectRequest, **kwargs: Any
+    ) -> PullObjectResponse:
+        """."""
+        return self._send_and_receive(request, PullObjectResponse, **kwargs)
+
+    def ConfirmMessageReceived(  # pylint: disable=C0103
+        self, request: ConfirmMessageReceivedRequest, **kwargs: Any
+    ) -> ConfirmMessageReceivedResponse:
+        """."""
+        return self._send_and_receive(request, ConfirmMessageReceivedResponse, **kwargs)

flwr/client/message_handler/message_handler.py

@@ -164,7 +164,7 @@ def validate_out_message(out_message: Message, in_message_metadata: Metadata) ->
     in_meta = in_message_metadata
     if (  # pylint: disable-next=too-many-boolean-expressions
         out_meta.run_id == in_meta.run_id
-        and out_meta.message_id == ""  # This will be generated by the server
+        and out_meta.message_id == out_message.object_id  # Should match the object id
         and out_meta.src_node_id == in_meta.dst_node_id
         and out_meta.dst_node_id == in_meta.src_node_id
         and out_meta.reply_to_message_id == in_meta.message_id

flwr/client/mod/comms_mods.py

@@ -32,14 +32,17 @@ def message_size_mod(
 
     This mod logs the size in bytes of the message being transmited.
     """
-    message_size_in_bytes = 0
+    # Log the size of the incoming message in bytes
+    total_bytes = sum(record.count_bytes() for record in msg.content.values())
+    log(INFO, "Incoming message size: %i bytes", total_bytes)
 
-    for record in msg.content.values():
-        message_size_in_bytes += record.count_bytes()
+    # Call the next layer
+    msg = call_next(msg, ctxt)
 
-    log(INFO, "Message size: %i bytes", message_size_in_bytes)
-
-    return call_next(msg, ctxt)
+    # Log the size of the outgoing message in bytes
+    total_bytes = sum(record.count_bytes() for record in msg.content.values())
+    log(INFO, "Outgoing message size: %i bytes", total_bytes)
+    return msg
 
 
 def arrays_size_mod(
@@ -50,25 +53,41 @@ def arrays_size_mod(
     This mod logs the number of array elements transmitted in ``ArrayRecord`` objects
     of the message as well as their sizes in bytes.
     """
-    model_size_stats = {}
-    arrays_size_in_bytes = 0
+    # Log the ArrayRecord size statistics and the total size in the incoming message
+    array_record_size_stats = _get_array_record_size_stats(msg)
+    total_bytes = sum(stat["bytes"] for stat in array_record_size_stats.values())
+    if array_record_size_stats:
+        log(INFO, "Incoming `ArrayRecord` size statistics:")
+        log(INFO, array_record_size_stats)
+    log(INFO, "Total array elements received: %i bytes", total_bytes)
+
+    msg = call_next(msg, ctxt)
+
+    # Log the ArrayRecord size statistics and the total size in the outgoing message
+    array_record_size_stats = _get_array_record_size_stats(msg)
+    total_bytes = sum(stat["bytes"] for stat in array_record_size_stats.values())
+    if array_record_size_stats:
+        log(INFO, "Outgoing `ArrayRecord` size statistics:")
+        log(INFO, array_record_size_stats)
+    log(INFO, "Total array elements sent: %i bytes", total_bytes)
+    return msg
+
+
+def _get_array_record_size_stats(
+    msg: Message,
+) -> dict[str, dict[str, int]]:
+    """Get `ArrayRecord` size statistics from the message."""
+    array_record_size_stats = {}
     for record_name, arr_record in msg.content.array_records.items():
         arr_record_bytes = arr_record.count_bytes()
-        arrays_size_in_bytes += arr_record_bytes
         element_count = 0
         for array in arr_record.values():
             element_count += (
                 int(np.prod(array.shape)) if array.shape else array.numpy().size
             )
 
-        model_size_stats[f"{record_name}"] = {
+        array_record_size_stats[record_name] = {
             "elements": element_count,
             "bytes": arr_record_bytes,
         }
-
-    if model_size_stats:
-        log(INFO, model_size_stats)
-
-    log(INFO, "Total array elements transmitted: %i bytes", arrays_size_in_bytes)
-
-    return call_next(msg, ctxt)
+    return array_record_size_stats