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

flwr/server/grid/grpc_grid.py

@@ -22,26 +22,51 @@ from typing import Optional, cast
 
 import grpc
 
-from flwr.common import Message, RecordDict
+from flwr.app.error import Error
+from flwr.common import Message, Metadata, RecordDict, now
 from flwr.common.constant import (
     SERVERAPPIO_API_DEFAULT_CLIENT_ADDRESS,
     SUPERLINK_NODE_ID,
+    ErrorCode,
+    MessageType,
 )
 from flwr.common.grpc import create_channel, on_channel_state_change
+from flwr.common.inflatable import (
+    InflatableObject,
+    get_all_nested_objects,
+    get_object_tree,
+    iterate_object_tree,
+    no_object_id_recompute,
+)
+from flwr.common.inflatable_protobuf_utils import (
+    make_pull_object_fn_protobuf,
+    make_push_object_fn_protobuf,
+)
+from flwr.common.inflatable_utils import (
+    ObjectUnavailableError,
+    inflate_object_from_contents,
+    pull_objects,
+    push_objects,
+)
 from flwr.common.logger import log, warn_deprecated_feature
+from flwr.common.message import make_message, remove_content_from_message
 from flwr.common.retry_invoker import _make_simple_grpc_retry_invoker, _wrap_stub
-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 Run
-from flwr.proto.message_pb2 import Message as ProtoMessage  # pylint: disable=E0611
+from flwr.proto.appio_pb2 import (  # pylint: disable=E0611
+    PullAppMessagesRequest,
+    PullAppMessagesResponse,
+    PushAppMessagesRequest,
+    PushAppMessagesResponse,
+)
+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
 from flwr.proto.serverappio_pb2 import (  # pylint: disable=E0611
     GetNodesRequest,
     GetNodesResponse,
-    PullResMessagesRequest,
-    PullResMessagesResponse,
-    PushInsMessagesRequest,
-    PushInsMessagesResponse,
 )
 from flwr.proto.serverappio_pb2_grpc import ServerAppIoStub  # pylint: disable=E0611
 
@@ -163,7 +188,7 @@ class GrpcGrid(Grid):
     def _check_message(self, message: Message) -> None:
         # Check if the message is valid
         if not (
-            message.metadata.message_id == ""
+            message.metadata.message_id != ""
             and message.metadata.reply_to_message_id == ""
             and message.metadata.ttl > 0
         ):
@@ -198,6 +223,39 @@ class GrpcGrid(Grid):
         )
         return [node.node_id for node in res.nodes]
 
+    def _try_push_messages(self, run_id: int, messages: Iterable[Message]) -> list[str]:
+        """Push all messages and its associated objects."""
+        # Prepare all Messages to be sent in a single request
+        proto_messages = []
+        object_trees = []
+        all_objects: dict[str, InflatableObject] = {}
+        for msg in messages:
+            proto_messages.append(message_to_proto(remove_content_from_message(msg)))
+            all_objects.update(get_all_nested_objects(msg))
+            object_trees.append(get_object_tree(msg))
+            del msg
+
+        # Call GrpcServerAppIoStub method
+        res: PushAppMessagesResponse = self._stub.PushMessages(
+            PushAppMessagesRequest(
+                messages_list=proto_messages,
+                run_id=run_id,
+                message_object_trees=object_trees,
+            )
+        )
+
+        # Push objects
+        push_objects(
+            all_objects,
+            push_object_fn=make_push_object_fn_protobuf(
+                push_object_protobuf=self._stub.PushObject,
+                node=self.node,
+                run_id=run_id,
+            ),
+            object_ids_to_push=set(res.objects_to_push),
+        )
+        return cast(list[str], res.message_ids)
+
     def push_messages(self, messages: Iterable[Message]) -> Iterable[str]:
         """Push messages to specified node IDs.
 
@@ -206,57 +264,106 @@ class GrpcGrid(Grid):
         """
         # Construct Messages
         run_id = cast(Run, self._run).run_id
-        message_proto_list: list[ProtoMessage] = []
-        for msg in messages:
-            # Populate metadata
-            msg.metadata.__dict__["_run_id"] = run_id
-            msg.metadata.__dict__["_src_node_id"] = self.node.node_id
-            # Check message
-            self._check_message(msg)
-            # Convert to proto
-            msg_proto = message_to_proto(msg)
-            # Add to list
-            message_proto_list.append(msg_proto)
-
+        message_ids: list[str] = []
         try:
-            # Call GrpcServerAppIoStub method
-            res: PushInsMessagesResponse = self._stub.PushMessages(
-                PushInsMessagesRequest(messages_list=message_proto_list, run_id=run_id)
-            )
-            if len([msg_id for msg_id in res.message_ids if msg_id]) != len(
-                message_proto_list
-            ):
-                log(
-                    WARNING,
-                    "Not all messages could be pushed to the SuperLink. The returned "
-                    "list has `None` for those messages (the order is preserved as "
-                    "passed to `push_messages`). This could be due to a malformed "
-                    "message.",
-                )
-            return list(res.message_ids)
+            with no_object_id_recompute():
+                for msg in messages:
+                    # Populate metadata
+                    msg.metadata.__dict__["_run_id"] = run_id
+                    msg.metadata.__dict__["_src_node_id"] = self.node.node_id
+                    msg.metadata.__dict__["_message_id"] = msg.object_id
+                    # Check message
+                    self._check_message(msg)
+                # Try pushing messages and their objects
+                message_ids = self._try_push_messages(run_id, messages)
+
         except grpc.RpcError as e:
             if e.code() == grpc.StatusCode.RESOURCE_EXHAUSTED:  # pylint: disable=E1101
                 log(ERROR, ERROR_MESSAGE_PUSH_MESSAGES_RESOURCE_EXHAUSTED)
                 return []
             raise
 
+        if None in message_ids:
+            log(
+                WARNING,
+                "Not all messages could be pushed to the SuperLink. The returned "
+                "list has `None` for those messages (the order is preserved as "
+                "passed to `push_messages`). This could be due to a malformed "
+                "message.",
+            )
+
+        return message_ids
+
     def pull_messages(self, message_ids: Iterable[str]) -> Iterable[Message]:
         """Pull messages based on message IDs.
 
         This method is used to collect messages from the SuperLink that correspond to a
         set of given message IDs.
         """
+        run_id = cast(Run, self._run).run_id
         try:
             # Pull Messages
-            res: PullResMessagesResponse = self._stub.PullMessages(
-                PullResMessagesRequest(
+            res: PullAppMessagesResponse = self._stub.PullMessages(
+                PullAppMessagesRequest(
                     message_ids=message_ids,
-                    run_id=cast(Run, self._run).run_id,
+                    run_id=run_id,
                 )
             )
-            # Convert Message from Protobuf representation
-            msgs = [message_from_proto(msg_proto) for msg_proto in res.messages_list]
-            return msgs
+            # Pull Messages from store
+            inflated_msgs: list[Message] = []
+            for msg_proto, msg_tree in zip(res.messages_list, res.message_object_trees):
+                msg_id = msg_proto.metadata.message_id
+                try:
+                    all_object_contents = pull_objects(
+                        object_ids=[
+                            tree.object_id for tree in iterate_object_tree(msg_tree)
+                        ],
+                        pull_object_fn=make_pull_object_fn_protobuf(
+                            pull_object_protobuf=self._stub.PullObject,
+                            node=self.node,
+                            run_id=run_id,
+                        ),
+                    )
+                except ObjectUnavailableError as e:
+                    # An ObjectUnavailableError indicates that the object is not yet
+                    # available. If this point has been reached, it means that the
+                    # Grid has tried to pull the object for the maximum number of times
+                    # or for the maximum time allowed, so we return an inflated message
+                    # with an error
+                    inflated_msgs.append(
+                        make_message(
+                            metadata=Metadata(
+                                run_id=run_id,
+                                message_id="",
+                                src_node_id=self.node.node_id,
+                                dst_node_id=self.node.node_id,
+                                message_type=MessageType.SYSTEM,
+                                group_id="",
+                                ttl=0,
+                                reply_to_message_id=msg_proto.metadata.reply_to_message_id,
+                                created_at=now().timestamp(),
+                            ),
+                            error=Error(
+                                code=ErrorCode.MESSAGE_UNAVAILABLE, reason=(str(e))
+                            ),
+                        )
+                    )
+                    continue
+
+                # Confirm that the message has been received
+                self._stub.ConfirmMessageReceived(
+                    ConfirmMessageReceivedRequest(
+                        node=self.node, run_id=run_id, message_object_id=msg_id
+                    )
+                )
+                message = cast(
+                    Message, inflate_object_from_contents(msg_id, all_object_contents)
+                )
+                message.metadata.__dict__["_message_id"] = msg_id
+                inflated_msgs.append(message)
+
+            return inflated_msgs
+
         except grpc.RpcError as e:
             if e.code() == grpc.StatusCode.RESOURCE_EXHAUSTED:  # pylint: disable=E1101
                 log(ERROR, ERROR_MESSAGE_PULL_MESSAGES_RESOURCE_EXHAUSTED)

flwr/server/grid/inmemory_grid.py

@@ -18,7 +18,7 @@
 import time
 from collections.abc import Iterable
 from typing import Optional, cast
-from uuid import UUID
+from uuid import uuid4
 
 from flwr.common import Message, RecordDict
 from flwr.common.constant import SUPERLINK_NODE_ID
@@ -56,7 +56,7 @@ class InMemoryGrid(Grid):
     def _check_message(self, message: Message) -> None:
         # Check if the message is valid
         if not (
-            message.metadata.message_id == ""
+            message.metadata.message_id != ""
             and message.metadata.reply_to_message_id == ""
             and message.metadata.ttl > 0
             and message.metadata.delivered_at == ""
@@ -111,6 +111,7 @@ class InMemoryGrid(Grid):
             # Populate metadata
             msg.metadata.__dict__["_run_id"] = cast(Run, self._run).run_id
             msg.metadata.__dict__["_src_node_id"] = self.node.node_id
+            msg.metadata.__dict__["_message_id"] = str(uuid4())
             # Check message
             self._check_message(msg)
             # Store in state
@@ -126,12 +127,12 @@ class InMemoryGrid(Grid):
         This method is used to collect messages from the SuperLink that correspond to a
         set of given message IDs.
         """
-        msg_ids = {UUID(msg_id) for msg_id in message_ids}
+        msg_ids = set(message_ids)
         # Pull Messages
         message_res_list = self.state.get_message_res(message_ids=msg_ids)
         # Get IDs of Messages these replies are for
         message_ins_ids_to_delete = {
-            UUID(msg_res.metadata.reply_to_message_id) for msg_res in message_res_list
+            msg_res.metadata.reply_to_message_id for msg_res in message_res_list
         }
         # Delete
         self.state.delete_messages(message_ins_ids=message_ins_ids_to_delete)

flwr/server/serverapp/app.py

@@ -16,6 +16,7 @@
 
 
 import argparse
+import gc
 from logging import DEBUG, ERROR, INFO
 from pathlib import Path
 from queue import Queue
@@ -38,6 +39,7 @@ from flwr.common.constant import (
     SubStatus,
 )
 from flwr.common.exit import ExitCode, flwr_exit
+from flwr.common.heartbeat import HeartbeatSender, get_grpc_app_heartbeat_fn
 from flwr.common.logger import (
     log,
     mirror_output_to_queue,
@@ -54,12 +56,12 @@ from flwr.common.serde import (
 )
 from flwr.common.telemetry import EventType, event
 from flwr.common.typing import RunNotRunningException, RunStatus
-from flwr.proto.run_pb2 import UpdateRunStatusRequest  # pylint: disable=E0611
-from flwr.proto.serverappio_pb2 import (  # pylint: disable=E0611
-    PullServerAppInputsRequest,
-    PullServerAppInputsResponse,
-    PushServerAppOutputsRequest,
+from flwr.proto.appio_pb2 import (  # pylint: disable=E0611
+    PullAppInputsRequest,
+    PullAppInputsResponse,
+    PushAppOutputsRequest,
 )
+from flwr.proto.run_pb2 import UpdateRunStatusRequest  # pylint: disable=E0611
 from flwr.server.grid.grpc_grid import GrpcGrid
 from flwr.server.run_serverapp import run as run_
 
@@ -106,24 +108,28 @@ def run_serverapp(  # pylint: disable=R0914, disable=W0212, disable=R0915
     certificates: Optional[bytes] = None,
 ) -> None:
     """Run Flower ServerApp process."""
-    grid = GrpcGrid(
-        serverappio_service_address=serverappio_api_address,
-        root_certificates=certificates,
-    )
-
     # Resolve directory where FABs are installed
     flwr_dir_ = get_flwr_dir(flwr_dir)
     log_uploader = None
     success = True
     hash_run_id = None
     run_status = None
+    heartbeat_sender = None
+    grid = None
+    context = None
     while True:
 
         try:
+            # Initialize the GrpcGrid
+            grid = GrpcGrid(
+                serverappio_service_address=serverappio_api_address,
+                root_certificates=certificates,
+            )
+
             # Pull ServerAppInputs from LinkState
-            req = PullServerAppInputsRequest()
+            req = PullAppInputsRequest()
             log(DEBUG, "[flwr-serverapp] Pull ServerAppInputs")
-            res: PullServerAppInputsResponse = grid._stub.PullServerAppInputs(req)
+            res: PullAppInputsResponse = grid._stub.PullAppInputs(req)
             if not res.HasField("run"):
                 sleep(3)
                 run_status = None
@@ -182,6 +188,16 @@ def run_serverapp(  # pylint: disable=R0914, disable=W0212, disable=R0915
                 event_details={"run-id-hash": hash_run_id},
             )
 
+            # Set up heartbeat sender
+            heartbeat_fn = get_grpc_app_heartbeat_fn(
+                grid._stub,
+                run.run_id,
+                failure_message="Heartbeat failed unexpectedly. The SuperLink could "
+                "not find the provided run ID, or the run status is invalid.",
+            )
+            heartbeat_sender = HeartbeatSender(heartbeat_fn)
+            heartbeat_sender.start()
+
             # Load and run the ServerApp with the Grid
             updated_context = run_(
                 grid=grid,
@@ -193,10 +209,8 @@ def run_serverapp(  # pylint: disable=R0914, disable=W0212, disable=R0915
             # Send resulting context
             context_proto = context_to_proto(updated_context)
             log(DEBUG, "[flwr-serverapp] Will push ServerAppOutputs")
-            out_req = PushServerAppOutputsRequest(
-                run_id=run.run_id, context=context_proto
-            )
-            _ = grid._stub.PushServerAppOutputs(out_req)
+            out_req = PushAppOutputsRequest(run_id=run.run_id, context=context_proto)
+            _ = grid._stub.PushAppOutputs(out_req)
 
             run_status = RunStatus(Status.FINISHED, SubStatus.COMPLETED, "")
         except RunNotRunningException:
@@ -213,19 +227,33 @@ def run_serverapp(  # pylint: disable=R0914, disable=W0212, disable=R0915
             success = False
 
         finally:
+            # Stop heartbeat sender
+            if heartbeat_sender:
+                heartbeat_sender.stop()
+                heartbeat_sender = None
+
             # Stop log uploader for this run and upload final logs
             if log_uploader:
                 stop_log_uploader(log_queue, log_uploader)
                 log_uploader = None
 
             # Update run status
-            if run_status:
+            if run_status and grid:
                 run_status_proto = run_status_to_proto(run_status)
                 grid._stub.UpdateRunStatus(
                     UpdateRunStatusRequest(
                         run_id=run.run_id, run_status=run_status_proto
                     )
                 )
+
+            # Close the Grpc connection
+            if grid:
+                grid.close()
+
+            # Clean up the Context
+            context = None
+            gc.collect()
+
             event(
                 EventType.FLWR_SERVERAPP_RUN_LEAVE,
                 event_details={"run-id-hash": hash_run_id, "success": success},

flwr/server/superlink/fleet/grpc_adapter/grpc_adapter_servicer.py

@@ -35,11 +35,16 @@ from flwr.proto.fab_pb2 import GetFabRequest  # pylint: disable=E0611
 from flwr.proto.fleet_pb2 import (  # pylint: disable=E0611
     CreateNodeRequest,
     DeleteNodeRequest,
-    PingRequest,
     PullMessagesRequest,
     PushMessagesRequest,
 )
 from flwr.proto.grpcadapter_pb2 import MessageContainer  # pylint: disable=E0611
+from flwr.proto.heartbeat_pb2 import SendNodeHeartbeatRequest  # pylint: disable=E0611
+from flwr.proto.message_pb2 import (  # pylint: disable=E0611
+    ConfirmMessageReceivedRequest,
+    PullObjectRequest,
+    PushObjectRequest,
+)
 from flwr.proto.run_pb2 import GetRunRequest  # pylint: disable=E0611
 
 from ..grpc_rere.fleet_servicer import FleetServicer
@@ -81,8 +86,10 @@ class GrpcAdapterServicer(grpcadapter_pb2_grpc.GrpcAdapterServicer, FleetService
             return _handle(request, context, CreateNodeRequest, self.CreateNode)
         if request.grpc_message_name == DeleteNodeRequest.__qualname__:
             return _handle(request, context, DeleteNodeRequest, self.DeleteNode)
-        if request.grpc_message_name == PingRequest.__qualname__:
-            return _handle(request, context, PingRequest, self.Ping)
+        if request.grpc_message_name == SendNodeHeartbeatRequest.__qualname__:
+            return _handle(
+                request, context, SendNodeHeartbeatRequest, self.SendNodeHeartbeat
+            )
         if request.grpc_message_name == GetRunRequest.__qualname__:
             return _handle(request, context, GetRunRequest, self.GetRun)
         if request.grpc_message_name == GetFabRequest.__qualname__:
@@ -91,4 +98,15 @@ class GrpcAdapterServicer(grpcadapter_pb2_grpc.GrpcAdapterServicer, FleetService
             return _handle(request, context, PullMessagesRequest, self.PullMessages)
         if request.grpc_message_name == PushMessagesRequest.__qualname__:
             return _handle(request, context, PushMessagesRequest, self.PushMessages)
+        if request.grpc_message_name == PushObjectRequest.__qualname__:
+            return _handle(request, context, PushObjectRequest, self.PushObject)
+        if request.grpc_message_name == PullObjectRequest.__qualname__:
+            return _handle(request, context, PullObjectRequest, self.PullObject)
+        if request.grpc_message_name == ConfirmMessageReceivedRequest.__qualname__:
+            return _handle(
+                request,
+                context,
+                ConfirmMessageReceivedRequest,
+                self.ConfirmMessageReceived,
+            )
         raise ValueError(f"Invalid grpc_message_name: {request.grpc_message_name}")

flwr/server/superlink/fleet/grpc_rere/fleet_servicer.py

@@ -20,6 +20,7 @@ from logging import DEBUG, INFO
 import grpc
 from google.protobuf.json_format import MessageToDict
 
+from flwr.common.inflatable import UnexpectedObjectContentError
 from flwr.common.logger import log
 from flwr.common.typing import InvalidRunStatusException
 from flwr.proto import fleet_pb2_grpc  # pylint: disable=E0611
@@ -29,34 +30,53 @@ from flwr.proto.fleet_pb2 import (  # pylint: disable=E0611
     CreateNodeResponse,
     DeleteNodeRequest,
     DeleteNodeResponse,
-    PingRequest,
-    PingResponse,
     PullMessagesRequest,
     PullMessagesResponse,
     PushMessagesRequest,
     PushMessagesResponse,
 )
+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
-from flwr.server.superlink.ffs.ffs_factory import FfsFactory
 from flwr.server.superlink.fleet.message_handler import message_handler
 from flwr.server.superlink.linkstate import LinkStateFactory
 from flwr.server.superlink.utils import abort_grpc_context
+from flwr.supercore.ffs import FfsFactory
+from flwr.supercore.object_store import ObjectStoreFactory
 
 
 class FleetServicer(fleet_pb2_grpc.FleetServicer):
     """Fleet API servicer."""
 
     def __init__(
-        self, state_factory: LinkStateFactory, ffs_factory: FfsFactory
+        self,
+        state_factory: LinkStateFactory,
+        ffs_factory: FfsFactory,
+        objectstore_factory: ObjectStoreFactory,
     ) -> None:
         self.state_factory = state_factory
         self.ffs_factory = ffs_factory
+        self.objectstore_factory = objectstore_factory
 
     def CreateNode(
         self, request: CreateNodeRequest, context: grpc.ServicerContext
     ) -> CreateNodeResponse:
         """."""
-        log(INFO, "[Fleet.CreateNode] Request ping_interval=%s", request.ping_interval)
+        log(
+            INFO,
+            "[Fleet.CreateNode] Request heartbeat_interval=%s",
+            request.heartbeat_interval,
+        )
         log(DEBUG, "[Fleet.CreateNode] Request: %s", MessageToDict(request))
         response = message_handler.create_node(
             request=request,
@@ -77,10 +97,12 @@ class FleetServicer(fleet_pb2_grpc.FleetServicer):
             state=self.state_factory.state(),
         )
 
-    def Ping(self, request: PingRequest, context: grpc.ServicerContext) -> PingResponse:
+    def SendNodeHeartbeat(
+        self, request: SendNodeHeartbeatRequest, context: grpc.ServicerContext
+    ) -> SendNodeHeartbeatResponse:
         """."""
-        log(DEBUG, "[Fleet.Ping] Request: %s", MessageToDict(request))
-        return message_handler.ping(
+        log(DEBUG, "[Fleet.SendNodeHeartbeat] Request: %s", MessageToDict(request))
+        return message_handler.send_node_heartbeat(
             request=request,
             state=self.state_factory.state(),
         )
@@ -94,6 +116,7 @@ class FleetServicer(fleet_pb2_grpc.FleetServicer):
         return message_handler.pull_messages(
             request=request,
             state=self.state_factory.state(),
+            store=self.objectstore_factory.store(),
         )
 
     def PushMessages(
@@ -113,6 +136,7 @@ class FleetServicer(fleet_pb2_grpc.FleetServicer):
             res = message_handler.push_messages(
                 request=request,
                 state=self.state_factory.state(),
+                store=self.objectstore_factory.store(),
             )
         except InvalidRunStatusException as e:
             abort_grpc_context(e.message, context)
@@ -129,6 +153,7 @@ class FleetServicer(fleet_pb2_grpc.FleetServicer):
             res = message_handler.get_run(
                 request=request,
                 state=self.state_factory.state(),
+                store=self.objectstore_factory.store(),
             )
         except InvalidRunStatusException as e:
             abort_grpc_context(e.message, context)
@@ -145,6 +170,75 @@ class FleetServicer(fleet_pb2_grpc.FleetServicer):
                 request=request,
                 ffs=self.ffs_factory.ffs(),
                 state=self.state_factory.state(),
+                store=self.objectstore_factory.store(),
+            )
+        except InvalidRunStatusException as e:
+            abort_grpc_context(e.message, context)
+
+        return res
+
+    def PushObject(
+        self, request: PushObjectRequest, context: grpc.ServicerContext
+    ) -> PushObjectResponse:
+        """Push an object to the ObjectStore."""
+        log(
+            DEBUG,
+            "[ServerAppIoServicer.PushObject] Push Object with object_id=%s",
+            request.object_id,
+        )
+
+        try:
+            # Insert in Store
+            res = message_handler.push_object(
+                request=request,
+                state=self.state_factory.state(),
+                store=self.objectstore_factory.store(),
+            )
+        except InvalidRunStatusException as e:
+            abort_grpc_context(e.message, context)
+        except UnexpectedObjectContentError as e:
+            # Object content is not valid
+            context.abort(grpc.StatusCode.FAILED_PRECONDITION, str(e))
+
+        return res
+
+    def PullObject(
+        self, request: PullObjectRequest, context: grpc.ServicerContext
+    ) -> PullObjectResponse:
+        """Pull an object from the ObjectStore."""
+        log(
+            DEBUG,
+            "[ServerAppIoServicer.PullObject] Pull Object with object_id=%s",
+            request.object_id,
+        )
+
+        try:
+            # Fetch from store
+            res = message_handler.pull_object(
+                request=request,
+                state=self.state_factory.state(),
+                store=self.objectstore_factory.store(),
+            )
+        except InvalidRunStatusException as e:
+            abort_grpc_context(e.message, context)
+
+        return res
+
+    def ConfirmMessageReceived(
+        self, request: ConfirmMessageReceivedRequest, context: grpc.ServicerContext
+    ) -> ConfirmMessageReceivedResponse:
+        """Confirm message received."""
+        log(
+            DEBUG,
+            "[Fleet.ConfirmMessageReceived] Message with ID '%s' has been received",
+            request.message_object_id,
+        )
+
+        try:
+            res = message_handler.confirm_message_received(
+                request=request,
+                state=self.state_factory.state(),
+                store=self.objectstore_factory.store(),
             )
         except InvalidRunStatusException as e:
             abort_grpc_context(e.message, context)

flwr/server/superlink/fleet/grpc_rere/server_interceptor.py

@@ -81,12 +81,9 @@ class AuthenticateServerInterceptor(grpc.ServerInterceptor):  # type: ignore
         metadata sent by the node. Continue RPC call if node is authenticated, else,
         terminate RPC call by setting context to abort.
         """
-        # Filter out non-Fleet service calls
+        # Only apply to Fleet service
         if not handler_call_details.method.startswith("/flwr.proto.Fleet/"):
-            return _unary_unary_rpc_terminator(
-                "This request should be sent to a different service.",
-                grpc.StatusCode.FAILED_PRECONDITION,
-            )
+            return continuation(handler_call_details)
 
         state = self.state_factory.state()
         metadata_dict = dict(handler_call_details.invocation_metadata)