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

flwr/server/grid/grpc_grid.py

@@ -22,31 +22,43 @@ 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_grpc_utils import (
-    make_pull_object_fn_grpc,
-    make_push_object_fn_grpc,
+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 remove_content_from_message
+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_to_proto, run_from_proto
 from flwr.common.typing import Run
+from flwr.proto.appio_pb2 import (  # pylint: disable=E0611
+    PullAppMessagesRequest,
+    PullAppMessagesResponse,
+    PushAppMessagesRequest,
+    PushAppMessagesResponse,
+)
 from flwr.proto.message_pb2 import (  # pylint: disable=E0611
     ConfirmMessageReceivedRequest,
 )
@@ -55,10 +67,6 @@ from flwr.proto.run_pb2 import GetRunRequest, GetRunResponse  # pylint: disable=
 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
 
@@ -215,37 +223,38 @@ class GrpcGrid(Grid):
         )
         return [node.node_id for node in res.nodes]
 
-    def _try_push_message(self, run_id: int, message: Message) -> str:
-        """Push one message and its associated objects."""
-        # Compute mapping of message descendants
-        all_objects = get_all_nested_objects(message)
-        msg_id = message.object_id
-        object_tree = get_object_tree(message)
+    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: PushInsMessagesResponse = self._stub.PushMessages(
-            PushInsMessagesRequest(
-                messages_list=[message_to_proto(remove_content_from_message(message))],
+        res: PushAppMessagesResponse = self._stub.PushMessages(
+            PushAppMessagesRequest(
+                messages_list=proto_messages,
                 run_id=run_id,
-                message_object_trees=[object_tree],
+                message_object_trees=object_trees,
             )
         )
 
         # Push objects
-        # 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_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
+        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.
@@ -257,16 +266,16 @@ class GrpcGrid(Grid):
         run_id = cast(Run, self._run).run_id
         message_ids: list[str] = []
         try:
-            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 message and its objects
-                with no_object_id_recompute():
-                    message_ids.append(self._try_push_message(run_id, msg))
+            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
@@ -294,24 +303,52 @@ class GrpcGrid(Grid):
         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=run_id,
                 )
             )
             # Pull Messages from store
             inflated_msgs: list[Message] = []
-            for msg_proto in res.messages_list:
+            for msg_proto, msg_tree in zip(res.messages_list, res.message_object_trees):
                 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,
-                    ),
-                )
+                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(

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
@@ -55,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_
 
@@ -107,11 +108,6 @@ 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
@@ -119,13 +115,21 @@ def run_serverapp(  # pylint: disable=R0914, disable=W0212, disable=R0915
     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
@@ -205,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:
@@ -236,7 +238,7 @@ def run_serverapp(  # pylint: disable=R0914, disable=W0212, disable=R0915
                 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(
@@ -244,6 +246,14 @@ def run_serverapp(  # pylint: disable=R0914, disable=W0212, disable=R0915
                     )
                 )
 
+            # 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

@@ -41,6 +41,7 @@ from flwr.proto.fleet_pb2 import (  # pylint: disable=E0611
 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,
 )
@@ -101,4 +102,11 @@ class GrpcAdapterServicer(grpcadapter_pb2_grpc.GrpcAdapterServicer, FleetService
             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

@@ -48,10 +48,10 @@ from flwr.proto.message_pb2 import (  # pylint: disable=E0611
     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
 
 

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)

flwr/server/superlink/fleet/message_handler/message_handler.py

@@ -46,7 +46,6 @@ from flwr.proto.heartbeat_pb2 import (  # pylint: disable=E0611
 from flwr.proto.message_pb2 import (  # pylint: disable=E0611
     ConfirmMessageReceivedRequest,
     ConfirmMessageReceivedResponse,
-    ObjectIDs,
     PullObjectRequest,
     PullObjectResponse,
     PushObjectRequest,
@@ -58,12 +57,11 @@ from flwr.proto.run_pb2 import (  # pylint: disable=E0611
     GetRunResponse,
     Run,
 )
-from flwr.server.superlink.ffs.ffs import Ffs
 from flwr.server.superlink.linkstate import LinkState
 from flwr.server.superlink.utils import check_abort
+from flwr.supercore.ffs import Ffs
 from flwr.supercore.object_store import NoObjectInStoreError, ObjectStore
-
-from ...utils import store_mapping_and_register_objects
+from flwr.supercore.object_store.utils import store_mapping_and_register_objects
 
 
 def create_node(
@@ -113,25 +111,22 @@ def pull_messages(
 
     # Convert to Messages
     msg_proto = []
-    objects_to_pull: dict[str, ObjectIDs] = {}
+    trees = []
     for msg in message_list:
         try:
-            msg_proto.append(message_to_proto(msg))
-
+            # Retrieve Message object tree from ObjectStore
             msg_object_id = msg.metadata.message_id
-            descendants = store.get_message_descendant_ids(msg_object_id)
-            # Include the object_id of the message itself
-            objects_to_pull[msg_object_id] = ObjectIDs(
-                object_ids=descendants + [msg_object_id]
-            )
+            obj_tree = store.get_object_tree(msg_object_id)
+
+            # Add Message and its object tree to the response
+            msg_proto.append(message_to_proto(msg))
+            trees.append(obj_tree)
         except NoObjectInStoreError as e:
             log(ERROR, e.message)
             # Delete message ins from state
             state.delete_messages(message_ins_ids={msg_object_id})
 
-    return PullMessagesResponse(
-        messages_list=msg_proto, objects_to_pull=objects_to_pull
-    )
+    return PullMessagesResponse(messages_list=msg_proto, message_object_trees=trees)
 
 
 def push_messages(
@@ -287,6 +282,5 @@ def confirm_message_received(
 
     # Delete the message object
     store.delete(request.message_object_id)
-    store.delete_message_descendant_ids(request.message_object_id)
 
     return ConfirmMessageReceivedResponse()

flwr/server/superlink/fleet/rest_rere/rest_api.py

@@ -47,10 +47,9 @@ from flwr.proto.message_pb2 import (  # pylint: disable=E0611
     PushObjectResponse,
 )
 from flwr.proto.run_pb2 import GetRunRequest, GetRunResponse  # pylint: disable=E0611
-from flwr.server.superlink.ffs.ffs import Ffs
-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 LinkState, LinkStateFactory
+from flwr.supercore.ffs import Ffs, FfsFactory
 from flwr.supercore.object_store import ObjectStore, ObjectStoreFactory
 
 try:

flwr/server/superlink/serverappio/serverappio_grpc.py

@@ -26,8 +26,8 @@ from flwr.common.logger import log
 from flwr.proto.serverappio_pb2_grpc import (  # pylint: disable=E0611
     add_ServerAppIoServicer_to_server,
 )
-from flwr.server.superlink.ffs.ffs_factory import FfsFactory
 from flwr.server.superlink.linkstate import LinkStateFactory
+from flwr.supercore.ffs import FfsFactory
 from flwr.supercore.object_store import ObjectStoreFactory
 
 from .serverappio_servicer import ServerAppIoServicer

flwr/server/superlink/serverappio/serverappio_servicer.py

@@ -42,6 +42,16 @@ from flwr.common.serde import (
 )
 from flwr.common.typing import Fab, RunStatus
 from flwr.proto import serverappio_pb2_grpc  # pylint: disable=E0611
+from flwr.proto.appio_pb2 import (  # pylint: disable=E0611
+    PullAppInputsRequest,
+    PullAppInputsResponse,
+    PullAppMessagesRequest,
+    PullAppMessagesResponse,
+    PushAppMessagesRequest,
+    PushAppMessagesResponse,
+    PushAppOutputsRequest,
+    PushAppOutputsResponse,
+)
 from flwr.proto.fab_pb2 import GetFabRequest, GetFabResponse  # pylint: disable=E0611
 from flwr.proto.heartbeat_pb2 import (  # pylint: disable=E0611
     SendAppHeartbeatRequest,
@@ -54,7 +64,6 @@ from flwr.proto.log_pb2 import (  # pylint: disable=E0611
 from flwr.proto.message_pb2 import (  # pylint: disable=E0611
     ConfirmMessageReceivedRequest,
     ConfirmMessageReceivedResponse,
-    ObjectIDs,
     PullObjectRequest,
     PullObjectResponse,
     PushObjectRequest,
@@ -72,23 +81,13 @@ from flwr.proto.run_pb2 import (  # pylint: disable=E0611
 from flwr.proto.serverappio_pb2 import (  # pylint: disable=E0611
     GetNodesRequest,
     GetNodesResponse,
-    PullResMessagesRequest,
-    PullResMessagesResponse,
-    PullServerAppInputsRequest,
-    PullServerAppInputsResponse,
-    PushInsMessagesRequest,
-    PushInsMessagesResponse,
-    PushServerAppOutputsRequest,
-    PushServerAppOutputsResponse,
 )
-from flwr.server.superlink.ffs.ffs import Ffs
-from flwr.server.superlink.ffs.ffs_factory import FfsFactory
 from flwr.server.superlink.linkstate import LinkState, LinkStateFactory
 from flwr.server.superlink.utils import abort_if
 from flwr.server.utils.validator import validate_message
+from flwr.supercore.ffs import Ffs, FfsFactory
 from flwr.supercore.object_store import NoObjectInStoreError, ObjectStoreFactory
-
-from ..utils import store_mapping_and_register_objects
+from flwr.supercore.object_store.utils import store_mapping_and_register_objects
 
 
 class ServerAppIoServicer(serverappio_pb2_grpc.ServerAppIoServicer):
@@ -129,8 +128,8 @@ class ServerAppIoServicer(serverappio_pb2_grpc.ServerAppIoServicer):
         return GetNodesResponse(nodes=nodes)
 
     def PushMessages(
-        self, request: PushInsMessagesRequest, context: grpc.ServicerContext
-    ) -> PushInsMessagesResponse:
+        self, request: PushAppMessagesRequest, context: grpc.ServicerContext
+    ) -> PushAppMessagesResponse:
         """Push a set of Messages."""
         log(DEBUG, "ServerAppIoServicer.PushMessages")
 
@@ -174,7 +173,7 @@ class ServerAppIoServicer(serverappio_pb2_grpc.ServerAppIoServicer):
         # Store Message object to descendants mapping and preregister objects
         objects_to_push = store_mapping_and_register_objects(store, request=request)
 
-        return PushInsMessagesResponse(
+        return PushAppMessagesResponse(
             message_ids=[
                 str(message_id) if message_id else "" for message_id in message_ids
             ],
@@ -182,8 +181,8 @@ class ServerAppIoServicer(serverappio_pb2_grpc.ServerAppIoServicer):
         )
 
     def PullMessages(  # pylint: disable=R0914
-        self, request: PullResMessagesRequest, context: grpc.ServicerContext
-    ) -> PullResMessagesResponse:
+        self, request: PullAppMessagesRequest, context: grpc.ServicerContext
+    ) -> PullAppMessagesResponse:
         """Pull a set of Messages."""
         log(DEBUG, "ServerAppIoServicer.PullMessages")
 
@@ -210,12 +209,6 @@ class ServerAppIoServicer(serverappio_pb2_grpc.ServerAppIoServicer):
             if msg_res.metadata.src_node_id == SUPERLINK_NODE_ID:
                 with no_object_id_recompute():
                     all_objects = get_all_nested_objects(msg_res)
-                    descendants = list(all_objects.keys())[:-1]
-                    message_obj_id = msg_res.metadata.message_id
-                    # Store mapping
-                    store.set_message_descendant_ids(
-                        msg_object_id=message_obj_id, descendant_ids=descendants
-                    )
                     # Preregister
                     store.preregister(request.run_id, get_object_tree(msg_res))
                     # Store objects
@@ -231,7 +224,7 @@ class ServerAppIoServicer(serverappio_pb2_grpc.ServerAppIoServicer):
 
         # Convert Messages to proto
         messages_list = []
-        objects_to_pull: dict[str, ObjectIDs] = {}
+        trees = []
         while messages_res:
             msg = messages_res.pop(0)
 
@@ -242,20 +235,20 @@ class ServerAppIoServicer(serverappio_pb2_grpc.ServerAppIoServicer):
                     request_name="PullMessages",
                     detail="`message.metadata` has mismatched `run_id`",
                 )
-            messages_list.append(message_to_proto(msg))
 
             try:
                 msg_object_id = msg.metadata.message_id
-                descendants = store.get_message_descendant_ids(msg_object_id)
-                # Add mapping of message object ID to its descendants
-                objects_to_pull[msg_object_id] = ObjectIDs(object_ids=descendants)
+                obj_tree = store.get_object_tree(msg_object_id)
+                # Add message and object tree to the response
+                messages_list.append(message_to_proto(msg))
+                trees.append(obj_tree)
             except NoObjectInStoreError as e:
                 log(ERROR, e.message)
                 # Delete message ins from state
                 state.delete_messages(message_ins_ids={msg_object_id})
 
-        return PullResMessagesResponse(
-            messages_list=messages_list, objects_to_pull=objects_to_pull
+        return PullAppMessagesResponse(
+            messages_list=messages_list, message_object_trees=trees
         )
 
     def GetRun(
@@ -288,11 +281,11 @@ class ServerAppIoServicer(serverappio_pb2_grpc.ServerAppIoServicer):
 
         raise ValueError(f"Found no FAB with hash: {request.hash_str}")
 
-    def PullServerAppInputs(
-        self, request: PullServerAppInputsRequest, context: grpc.ServicerContext
-    ) -> PullServerAppInputsResponse:
+    def PullAppInputs(
+        self, request: PullAppInputsRequest, context: grpc.ServicerContext
+    ) -> PullAppInputsResponse:
         """Pull ServerApp process inputs."""
-        log(DEBUG, "ServerAppIoServicer.PullServerAppInputs")
+        log(DEBUG, "ServerAppIoServicer.PullAppInputs")
         # Init access to LinkState
         state = self.state_factory.state()
 
@@ -302,7 +295,7 @@ class ServerAppIoServicer(serverappio_pb2_grpc.ServerAppIoServicer):
             run_id = state.get_pending_run_id()
             # If there's no pending run, return an empty response
             if run_id is None:
-                return PullServerAppInputsResponse()
+                return PullAppInputsResponse()
 
             # Init access to Ffs
             ffs = self.ffs_factory.ffs()
@@ -318,7 +311,7 @@ class ServerAppIoServicer(serverappio_pb2_grpc.ServerAppIoServicer):
                 # Update run status to STARTING
                 if state.update_run_status(run_id, RunStatus(Status.STARTING, "", "")):
                     log(INFO, "Starting run %d", run_id)
-                    return PullServerAppInputsResponse(
+                    return PullAppInputsResponse(
                         context=context_to_proto(serverapp_ctxt),
                         run=run_to_proto(run),
                         fab=fab_to_proto(fab),
@@ -328,11 +321,11 @@ class ServerAppIoServicer(serverappio_pb2_grpc.ServerAppIoServicer):
         # or if the status cannot be updated to STARTING
         raise RuntimeError(f"Failed to start run {run_id}")
 
-    def PushServerAppOutputs(
-        self, request: PushServerAppOutputsRequest, context: grpc.ServicerContext
-    ) -> PushServerAppOutputsResponse:
+    def PushAppOutputs(
+        self, request: PushAppOutputsRequest, context: grpc.ServicerContext
+    ) -> PushAppOutputsResponse:
         """Push ServerApp process outputs."""
-        log(DEBUG, "ServerAppIoServicer.PushServerAppOutputs")
+        log(DEBUG, "ServerAppIoServicer.PushAppOutputs")
 
         # Init state and store
         state = self.state_factory.state()
@@ -348,7 +341,7 @@ class ServerAppIoServicer(serverappio_pb2_grpc.ServerAppIoServicer):
         )
 
         state.set_serverapp_context(request.run_id, context_from_proto(request.context))
-        return PushServerAppOutputsResponse()
+        return PushAppOutputsResponse()
 
     def UpdateRunStatus(
         self, request: UpdateRunStatusRequest, context: grpc.ServicerContext
@@ -512,7 +505,6 @@ class ServerAppIoServicer(serverappio_pb2_grpc.ServerAppIoServicer):
 
         # Delete the message object
         store.delete(request.message_object_id)
-        store.delete_message_descendant_ids(request.message_object_id)
 
         return ConfirmMessageReceivedResponse()
 

flwr/server/superlink/simulation/simulationio_grpc.py

@@ -26,8 +26,8 @@ from flwr.common.logger import log
 from flwr.proto.simulationio_pb2_grpc import (  # pylint: disable=E0611
     add_SimulationIoServicer_to_server,
 )
-from flwr.server.superlink.ffs.ffs_factory import FfsFactory
 from flwr.server.superlink.linkstate import LinkStateFactory
+from flwr.supercore.ffs import FfsFactory
 
 from .simulationio_servicer import SimulationIoServicer
 

flwr/server/superlink/simulation/simulationio_servicer.py

@@ -56,9 +56,9 @@ from flwr.proto.simulationio_pb2 import (  # pylint: disable=E0611
     PushSimulationOutputsRequest,
     PushSimulationOutputsResponse,
 )
-from flwr.server.superlink.ffs.ffs_factory import FfsFactory
 from flwr.server.superlink.linkstate import LinkStateFactory
 from flwr.server.superlink.utils import abort_if
+from flwr.supercore.ffs import FfsFactory
 
 
 class SimulationIoServicer(simulationio_pb2_grpc.SimulationIoServicer):