flwr-nightly 1.20.0.dev20250712__py3-none-any.whl → 1.20.0.dev20250715__py3-none-any.whl

flwr/server/serverapp/app.py

@@ -55,12 +55,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_
 
@@ -125,9 +125,9 @@ def run_serverapp(  # pylint: disable=R0914, disable=W0212, disable=R0915
             )
 
             # 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
@@ -207,10 +207,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:

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,
@@ -113,25 +112,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 +283,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/serverappio/serverappio_servicer.py

@@ -27,6 +27,7 @@ from flwr.common.inflatable import (
     UnexpectedObjectContentError,
     get_all_nested_objects,
     get_object_tree,
+    iterate_object_tree,
     no_object_id_recompute,
 )
 from flwr.common.logger import log
@@ -42,6 +43,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,
@@ -72,14 +83,6 @@ 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.linkstate import LinkState, LinkStateFactory
 from flwr.server.superlink.utils import abort_if
@@ -128,8 +131,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")
 
@@ -173,7 +176,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
             ],
@@ -181,8 +184,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")
 
@@ -209,12 +212,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
@@ -245,7 +242,9 @@ class ServerAppIoServicer(serverappio_pb2_grpc.ServerAppIoServicer):
 
             try:
                 msg_object_id = msg.metadata.message_id
-                descendants = store.get_message_descendant_ids(msg_object_id)
+                obj_tree = store.get_object_tree(msg_object_id)
+                descendants = [node.object_id for node in iterate_object_tree(obj_tree)]
+                descendants = descendants[:-1]  # Exclude the message itself
                 # Add mapping of message object ID to its descendants
                 objects_to_pull[msg_object_id] = ObjectIDs(object_ids=descendants)
             except NoObjectInStoreError as e:
@@ -253,7 +252,7 @@ class ServerAppIoServicer(serverappio_pb2_grpc.ServerAppIoServicer):
                 # Delete message ins from state
                 state.delete_messages(message_ins_ids={msg_object_id})
 
-        return PullResMessagesResponse(
+        return PullAppMessagesResponse(
             messages_list=messages_list, objects_to_pull=objects_to_pull
         )
 
@@ -287,11 +286,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()
 
@@ -301,7 +300,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()
@@ -317,7 +316,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),
@@ -327,11 +326,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()
@@ -347,7 +346,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
@@ -511,7 +510,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/utils.py

@@ -20,11 +20,10 @@ from typing import Optional, Union
 import grpc
 
 from flwr.common.constant import Status, SubStatus
-from flwr.common.inflatable import iterate_object_tree
 from flwr.common.typing import RunStatus
+from flwr.proto.appio_pb2 import PushAppMessagesRequest  # pylint: disable=E0611
 from flwr.proto.fleet_pb2 import PushMessagesRequest  # pylint: disable=E0611
 from flwr.proto.message_pb2 import ObjectIDs  # pylint: disable=E0611
-from flwr.proto.serverappio_pb2 import PushInsMessagesRequest  # pylint: disable=E0611
 from flwr.server.superlink.linkstate import LinkState
 from flwr.supercore.object_store import ObjectStore
 
@@ -77,7 +76,7 @@ def abort_if(
 
 
 def store_mapping_and_register_objects(
-    store: ObjectStore, request: Union[PushInsMessagesRequest, PushMessagesRequest]
+    store: ObjectStore, request: Union[PushAppMessagesRequest, PushMessagesRequest]
 ) -> dict[str, ObjectIDs]:
     """Store Message object to descendants mapping and preregister objects."""
     if not request.messages_list:
@@ -90,17 +89,10 @@ def store_mapping_and_register_objects(
     run_id = request.messages_list[0].metadata.run_id
 
     for object_tree in request.message_object_trees:
-        all_object_ids = [obj.object_id for obj in iterate_object_tree(object_tree)]
-        msg_object_id, descendant_ids = all_object_ids[-1], all_object_ids[:-1]
-        # Store mapping
-        store.set_message_descendant_ids(
-            msg_object_id=msg_object_id, descendant_ids=descendant_ids
-        )
-
         # Preregister
         object_ids_just_registered = store.preregister(run_id, object_tree)
         # Keep track of objects that need to be pushed
-        objects_to_push[msg_object_id] = ObjectIDs(
+        objects_to_push[object_tree.object_id] = ObjectIDs(
             object_ids=object_ids_just_registered
         )
 

flwr/supercore/grpc_health/__init__.py

@@ -0,0 +1,22 @@
+# 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.
+# ==============================================================================
+"""GRPC health servicers."""
+
+
+from .simple_health_servicer import SimpleHealthServicer
+
+__all__ = [
+    "SimpleHealthServicer",
+]

flwr/supercore/grpc_health/simple_health_servicer.py

@@ -0,0 +1,38 @@
+# 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.
+# ==============================================================================
+"""Simple gRPC health servicers."""
+
+
+import grpc
+
+# pylint: disable=E0611
+from grpc_health.v1.health_pb2 import HealthCheckRequest, HealthCheckResponse
+from grpc_health.v1.health_pb2_grpc import HealthServicer
+
+# pylint: enable=E0611
+
+
+class SimpleHealthServicer(HealthServicer):  # type: ignore
+    """A simple gRPC health servicer that always returns SERVING."""
+
+    def Check(
+        self, request: HealthCheckRequest, context: grpc.RpcContext
+    ) -> HealthCheckResponse:
+        """Return a HealthCheckResponse with SERVING status."""
+        return HealthCheckResponse(status=HealthCheckResponse.SERVING)
+
+    def Watch(self, request: HealthCheckRequest, context: grpc.RpcContext) -> None:
+        """Watch the health status (not implemented)."""
+        context.abort(grpc.StatusCode.UNIMPLEMENTED, "Watch is not implemented")

flwr/supercore/object_store/in_memory_object_store.py

@@ -20,7 +20,6 @@ from dataclasses import dataclass
 from typing import Optional
 
 from flwr.common.inflatable import (
-    get_object_children_ids_from_object_content,
     get_object_id,
     is_valid_sha256_hash,
     iterate_object_tree,
@@ -37,6 +36,7 @@ class ObjectEntry:
 
     content: bytes
     is_available: bool
+    child_object_ids: list[str]  # List of child object IDs
     ref_count: int  # Number of references (direct parents) to this object
     runs: set[int]  # Set of run IDs that used this object
 
@@ -70,6 +70,9 @@ class InMemoryObjectStore(ObjectStore):
                     self.store[obj_id] = ObjectEntry(
                         content=b"",  # Initially empty content
                         is_available=False,  # Initially not available
+                        child_object_ids=[  # List of child object IDs
+                            child.object_id for child in tree_node.children
+                        ],
                         ref_count=0,  # Reference count starts at 0
                         runs={run_id},  # Start with the current run ID
                     )
@@ -102,6 +105,32 @@ class InMemoryObjectStore(ObjectStore):
 
         return new_objects
 
+    def get_object_tree(self, object_id: str) -> ObjectTree:
+        """Get the object tree for a given object ID."""
+        with self.lock_store:
+            # Raise an exception if there's no object with the given ID
+            if not (object_entry := self.store.get(object_id)):
+                raise NoObjectInStoreError(
+                    f"Object with ID '{object_id}' was not pre-registered."
+                )
+
+            # Build the object trees of all children
+            try:
+                child_trees = [
+                    self.get_object_tree(child_id)
+                    for child_id in object_entry.child_object_ids
+                ]
+            except NoObjectInStoreError as e:
+                # Raise an error if any child object is missing
+                # This indicates an integrity issue
+                raise NoObjectInStoreError(
+                    f"Object tree for object ID '{object_id}' contains missing "
+                    "children. This may indicate a corrupted object store."
+                ) from e
+
+            # Create and return the ObjectTree for the current object
+            return ObjectTree(object_id=object_id, children=child_trees)
+
     def put(self, object_id: str, object_content: bytes) -> None:
         """Put an object into the store."""
         if self.verify:
@@ -128,29 +157,6 @@ class InMemoryObjectStore(ObjectStore):
             self.store[object_id].content = object_content
             self.store[object_id].is_available = True
 
-    def set_message_descendant_ids(
-        self, msg_object_id: str, descendant_ids: list[str]
-    ) -> None:
-        """Store the mapping from a ``Message`` object ID to the object IDs of its
-        descendants."""
-        with self.lock_msg_mapping:
-            self.msg_descendant_objects_mapping[msg_object_id] = descendant_ids
-
-    def get_message_descendant_ids(self, msg_object_id: str) -> list[str]:
-        """Retrieve the object IDs of all descendants of a given Message."""
-        with self.lock_msg_mapping:
-            if msg_object_id not in self.msg_descendant_objects_mapping:
-                raise NoObjectInStoreError(
-                    f"No message registered in Object Store with ID '{msg_object_id}'. "
-                    "Mapping to descendants could not be found."
-                )
-            return self.msg_descendant_objects_mapping[msg_object_id]
-
-    def delete_message_descendant_ids(self, msg_object_id: str) -> None:
-        """Delete the mapping from a ``Message`` object ID to its descendants."""
-        with self.lock_msg_mapping:
-            self.msg_descendant_objects_mapping.pop(msg_object_id, None)
-
     def get(self, object_id: str) -> Optional[bytes]:
         """Get an object from the store."""
         with self.lock_store:
@@ -177,10 +183,7 @@ class InMemoryObjectStore(ObjectStore):
                     self.run_objects_mapping[run_id].discard(object_id)
 
                 # Decrease the reference count of its children
-                children_ids = get_object_children_ids_from_object_content(
-                    object_entry.content
-                )
-                for child_id in children_ids:
+                for child_id in object_entry.child_object_ids:
                     self.store[child_id].ref_count -= 1
 
                     # Recursively try to delete the child object
@@ -205,9 +208,6 @@ class InMemoryObjectStore(ObjectStore):
                     # Delete the message object and its unreferenced descendants
                     self.delete(object_id)
 
-                    # Delete the message's descendants mapping
-                    self.delete_message_descendant_ids(object_id)
-
             # Remove the run from the mapping
             del self.run_objects_mapping[run_id]
 

flwr/supercore/object_store/object_store.py

@@ -60,6 +60,22 @@ class ObjectStore(abc.ABC):
             in the `ObjectStore`, or were preregistered but are not yet available.
         """
 
+    @abc.abstractmethod
+    def get_object_tree(self, object_id: str) -> ObjectTree:
+        """Get the object tree for a given object ID.
+
+        Parameters
+        ----------
+        object_id : str
+            The ID of the object for which to retrieve the object tree.
+
+        Returns
+        -------
+        ObjectTree
+            An ObjectTree representing the hierarchical structure of the object with
+            the given ID and its descendants.
+        """
+
     @abc.abstractmethod
     def put(self, object_id: str, object_content: bytes) -> None:
         """Put an object into the store.
@@ -126,46 +142,6 @@ class ObjectStore(abc.ABC):
         This method should remove all objects from the store.
         """
 
-    @abc.abstractmethod
-    def set_message_descendant_ids(
-        self, msg_object_id: str, descendant_ids: list[str]
-    ) -> None:
-        """Store the mapping from a ``Message`` object ID to the object IDs of its
-        descendants.
-
-        Parameters
-        ----------
-        msg_object_id : str
-            The object ID of the ``Message``.
-        descendant_ids : list[str]
-            A list of object IDs representing all descendant objects of the ``Message``.
-        """
-
-    @abc.abstractmethod
-    def get_message_descendant_ids(self, msg_object_id: str) -> list[str]:
-        """Retrieve the object IDs of all descendants of a given ``Message``.
-
-        Parameters
-        ----------
-        msg_object_id : str
-            The object ID of the ``Message``.
-
-        Returns
-        -------
-        list[str]
-            A list of object IDs of all descendant objects of the ``Message``.
-        """
-
-    @abc.abstractmethod
-    def delete_message_descendant_ids(self, msg_object_id: str) -> None:
-        """Delete the mapping from a ``Message`` object ID to its descendants.
-
-        Parameters
-        ----------
-        msg_object_id : str
-            The object ID of the ``Message``.
-        """
-
     @abc.abstractmethod
     def __contains__(self, object_id: str) -> bool:
         """Check if an object_id is in the store.

flwr/supernode/runtime/run_clientapp.py

@@ -50,8 +50,11 @@ from flwr.proto.clientappio_pb2 import (
     GetRunIdsWithPendingMessagesResponse,
     PullClientAppInputsRequest,
     PullClientAppInputsResponse,
+    PullMessageRequest,
+    PullMessageResponse,
     PushClientAppOutputsRequest,
     PushClientAppOutputsResponse,
+    PushMessageRequest,
     RequestTokenRequest,
     RequestTokenResponse,
 )
@@ -199,10 +202,14 @@ def pull_clientappinputs(
     masked_token = mask_string(token)
     log(INFO, "[flwr-clientapp] Pull `ClientAppInputs` for token %s", masked_token)
     try:
+        # Pull Message
+        res_msg: PullMessageResponse = stub.PullMessage(PullMessageRequest(token=token))
+        message = message_from_proto(res_msg.message)
+
+        # Pull Context, Run and (optional) FAB
         res: PullClientAppInputsResponse = stub.PullClientAppInputs(
             PullClientAppInputsRequest(token=token)
         )
-        message = message_from_proto(res.message)
         context = context_from_proto(res.context)
         run = run_from_proto(res.run)
         fab = fab_from_proto(res.fab) if res.fab else None
@@ -224,10 +231,13 @@ def push_clientappoutputs(
     proto_context = context_to_proto(context)
 
     try:
+
+        # Push Message
+        _ = stub.PushMessage(PushMessageRequest(token=token, message=proto_message))
+
+        # Push Context
         res: PushClientAppOutputsResponse = stub.PushClientAppOutputs(
-            PushClientAppOutputsRequest(
-                token=token, message=proto_message, context=proto_context
-            )
+            PushClientAppOutputsRequest(token=token, context=proto_context)
         )
         return res
     except grpc.RpcError as e:

flwr/supernode/servicer/clientappio/clientappio_servicer.py

@@ -39,8 +39,12 @@ from flwr.proto.clientappio_pb2 import (  # pylint: disable=E0401
     GetRunIdsWithPendingMessagesResponse,
     PullClientAppInputsRequest,
     PullClientAppInputsResponse,
+    PullMessageRequest,
+    PullMessageResponse,
     PushClientAppOutputsRequest,
     PushClientAppOutputsResponse,
+    PushMessageRequest,
+    PushMessageResponse,
     RequestTokenRequest,
     RequestTokenResponse,
 )
@@ -119,14 +123,12 @@ class ClientAppIoServicer(clientappio_pb2_grpc.ClientAppIoServicer):
             )
             raise RuntimeError("This line should never be reached.")
 
-        # Retrieve message, context, run and fab for this run
-        message = state.get_messages(run_ids=[run_id], is_reply=False)[0]
+        # Retrieve context, run and fab for this run
         context = cast(Context, state.get_context(run_id))
         run = cast(Run, state.get_run(run_id))
         fab = Fab(run.fab_hash, ffs.get(run.fab_hash)[0])  # type: ignore
 
         return PullClientAppInputsResponse(
-            message=message_to_proto(message),
             context=context_to_proto(context),
             run=run_to_proto(run),
             fab=fab_to_proto(fab),
@@ -150,8 +152,7 @@ class ClientAppIoServicer(clientappio_pb2_grpc.ClientAppIoServicer):
             )
             raise RuntimeError("This line should never be reached.")
 
-        # Save the message and context to the state
-        state.store_message(message_from_proto(request.message))
+        # Save the context to the state
         state.store_context(context_from_proto(request.context))
 
         # Remove the token to make the run eligible for processing
@@ -159,3 +160,45 @@ class ClientAppIoServicer(clientappio_pb2_grpc.ClientAppIoServicer):
         state.delete_token(run_id)
 
         return PushClientAppOutputsResponse()
+
+    def PullMessage(
+        self, request: PullMessageRequest, context: grpc.ServicerContext
+    ) -> PullMessageResponse:
+        """Pull one Message."""
+        # Initialize state and ffs connection
+        state = self.state_factory.state()
+
+        # Validate the token
+        run_id = state.get_run_id_by_token(request.token)
+        if run_id is None or not state.verify_token(run_id, request.token):
+            context.abort(
+                grpc.StatusCode.PERMISSION_DENIED,
+                "Invalid token.",
+            )
+            raise RuntimeError("This line should never be reached.")
+
+        # Retrieve message, context, run and fab for this run
+        message = state.get_messages(run_ids=[run_id], is_reply=False)[0]
+
+        return PullMessageResponse(message=message_to_proto(message))
+
+    def PushMessage(
+        self, request: PushMessageRequest, context: grpc.ServicerContext
+    ) -> PushMessageResponse:
+        """Push one Message."""
+        # Initialize state connection
+        state = self.state_factory.state()
+
+        # Validate the token
+        run_id = state.get_run_id_by_token(request.token)
+        if run_id is None or not state.verify_token(run_id, request.token):
+            context.abort(
+                grpc.StatusCode.PERMISSION_DENIED,
+                "Invalid token.",
+            )
+            raise RuntimeError("This line should never be reached.")
+
+        # Save the message and context to the state
+        state.store_message(message_from_proto(request.message))
+
+        return PushMessageResponse()