flwr-nightly 1.19.0.dev20250610__py3-none-any.whl → 1.19.0.dev20250612__py3-none-any.whl

flwr/server/grid/grpc_grid.py

@@ -28,7 +28,11 @@ from flwr.common.constant import (
     SUPERLINK_NODE_ID,
 )
 from flwr.common.grpc import create_channel, on_channel_state_change
-from flwr.common.inflatable import get_all_nested_objects
+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,
@@ -43,7 +47,9 @@ from flwr.common.message import 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.message_pb2 import ObjectIDs  # pylint: disable=E0611
+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
@@ -213,18 +219,15 @@ class GrpcGrid(Grid):
         """Push one message and its associated objects."""
         # Compute mapping of message descendants
         all_objects = get_all_nested_objects(message)
-        all_object_ids = list(all_objects.keys())
-        msg_id = all_object_ids[-1]  # Last object is the message itself
-        descendant_ids = all_object_ids[:-1]  # All but the last object are descendants
+        msg_id = message.object_id
+        object_tree = get_object_tree(message)
 
         # Call GrpcServerAppIoStub method
         res: PushInsMessagesResponse = self._stub.PushMessages(
             PushInsMessagesRequest(
                 messages_list=[message_to_proto(remove_content_from_message(message))],
                 run_id=run_id,
-                msg_to_descendant_mapping={
-                    msg_id: ObjectIDs(object_ids=descendant_ids)
-                },
+                message_object_trees=[object_tree],
             )
         )
 
@@ -262,7 +265,8 @@ class GrpcGrid(Grid):
                 # Check message
                 self._check_message(msg)
                 # Try pushing message and its objects
-                message_ids.append(self._try_push_message(run_id, msg))
+                with no_object_id_recompute():
+                    message_ids.append(self._try_push_message(run_id, msg))
 
         except grpc.RpcError as e:
             if e.code() == grpc.StatusCode.RESOURCE_EXHAUSTED:  # pylint: disable=E1101
@@ -308,6 +312,13 @@ class GrpcGrid(Grid):
                         run_id=run_id,
                     ),
                 )
+
+                # 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)
                 )

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

@@ -40,6 +40,8 @@ from flwr.proto.heartbeat_pb2 import (  # pylint: disable=E0611
     SendNodeHeartbeatResponse,
 )
 from flwr.proto.message_pb2 import (  # pylint: disable=E0611
+    ConfirmMessageReceivedRequest,
+    ConfirmMessageReceivedResponse,
     PullObjectRequest,
     PullObjectResponse,
     PushObjectRequest,
@@ -151,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)
@@ -167,6 +170,7 @@ 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)
@@ -219,3 +223,24 @@ class FleetServicer(fleet_pb2_grpc.FleetServicer):
             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)
+
+        return res

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

@@ -44,6 +44,8 @@ from flwr.proto.heartbeat_pb2 import (  # pylint: disable=E0611
     SendNodeHeartbeatResponse,
 )
 from flwr.proto.message_pb2 import (  # pylint: disable=E0611
+    ConfirmMessageReceivedRequest,
+    ConfirmMessageReceivedResponse,
     ObjectIDs,
     PullObjectRequest,
     PullObjectResponse,
@@ -146,6 +148,7 @@ def push_messages(
         msg.metadata.run_id,
         [Status.PENDING, Status.STARTING, Status.FINISHED],
         state,
+        store,
     )
     if abort_msg:
         raise InvalidRunStatusException(abort_msg)
@@ -165,7 +168,9 @@ def push_messages(
     return response
 
 
-def get_run(request: GetRunRequest, state: LinkState) -> GetRunResponse:
+def get_run(
+    request: GetRunRequest, state: LinkState, store: ObjectStore
+) -> GetRunResponse:
     """Get run information."""
     run = state.get_run(request.run_id)
 
@@ -177,6 +182,7 @@ def get_run(request: GetRunRequest, state: LinkState) -> GetRunResponse:
         request.run_id,
         [Status.PENDING, Status.STARTING, Status.FINISHED],
         state,
+        store,
     )
     if abort_msg:
         raise InvalidRunStatusException(abort_msg)
@@ -193,7 +199,7 @@ def get_run(request: GetRunRequest, state: LinkState) -> GetRunResponse:
 
 
 def get_fab(
-    request: GetFabRequest, ffs: Ffs, state: LinkState  # pylint: disable=W0613
+    request: GetFabRequest, ffs: Ffs, state: LinkState, store: ObjectStore
 ) -> GetFabResponse:
     """Get FAB."""
     # Abort if the run is not running
@@ -201,6 +207,7 @@ def get_fab(
         request.run_id,
         [Status.PENDING, Status.STARTING, Status.FINISHED],
         state,
+        store,
     )
     if abort_msg:
         raise InvalidRunStatusException(abort_msg)
@@ -220,6 +227,7 @@ def push_object(
         request.run_id,
         [Status.PENDING, Status.STARTING, Status.FINISHED],
         state,
+        store,
     )
     if abort_msg:
         raise InvalidRunStatusException(abort_msg)
@@ -245,6 +253,7 @@ def pull_object(
         request.run_id,
         [Status.PENDING, Status.STARTING, Status.FINISHED],
         state,
+        store,
     )
     if abort_msg:
         raise InvalidRunStatusException(abort_msg)
@@ -259,3 +268,25 @@ def pull_object(
             object_content=content,
         )
     return PullObjectResponse(object_found=False, object_available=False)
+
+
+def confirm_message_received(
+    request: ConfirmMessageReceivedRequest,
+    state: LinkState,
+    store: ObjectStore,
+) -> ConfirmMessageReceivedResponse:
+    """Confirm message received handler."""
+    abort_msg = check_abort(
+        request.run_id,
+        [Status.PENDING, Status.STARTING, Status.FINISHED],
+        state,
+        store,
+    )
+    if abort_msg:
+        raise InvalidRunStatusException(abort_msg)
+
+    # 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

@@ -38,6 +38,14 @@ 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 import Ffs
 from flwr.server.superlink.ffs.ffs_factory import FfsFactory
@@ -131,6 +139,28 @@ async def push_message(request: PushMessagesRequest) -> PushMessagesResponse:
     return message_handler.push_messages(request=request, state=state, store=store)
 
 
+@rest_request_response(PullObjectRequest)
+async def pull_object(request: PullObjectRequest) -> PullObjectResponse:
+    """Pull PullObject."""
+    # Get state from app
+    state: LinkState = cast(LinkStateFactory, app.state.STATE_FACTORY).state()
+    store: ObjectStore = cast(ObjectStoreFactory, app.state.OBJECTSTORE_FACTORY).store()
+
+    # Handle message
+    return message_handler.pull_object(request=request, state=state, store=store)
+
+
+@rest_request_response(PushObjectRequest)
+async def push_object(request: PushObjectRequest) -> PushObjectResponse:
+    """Pull PushObject."""
+    # Get state from app
+    state: LinkState = cast(LinkStateFactory, app.state.STATE_FACTORY).state()
+    store: ObjectStore = cast(ObjectStoreFactory, app.state.OBJECTSTORE_FACTORY).store()
+
+    # Handle message
+    return message_handler.push_object(request=request, state=state, store=store)
+
+
 @rest_request_response(SendNodeHeartbeatRequest)
 async def send_node_heartbeat(
     request: SendNodeHeartbeatRequest,
@@ -148,9 +178,10 @@ async def get_run(request: GetRunRequest) -> GetRunResponse:
     """GetRun."""
     # Get state from app
     state: LinkState = cast(LinkStateFactory, app.state.STATE_FACTORY).state()
+    store: ObjectStore = cast(ObjectStoreFactory, app.state.OBJECTSTORE_FACTORY).store()
 
     # Handle message
-    return message_handler.get_run(request=request, state=state)
+    return message_handler.get_run(request=request, state=state, store=store)
 
 
 @rest_request_response(GetFabRequest)
@@ -161,9 +192,25 @@ async def get_fab(request: GetFabRequest) -> GetFabResponse:
 
     # Get state from app
     state: LinkState = cast(LinkStateFactory, app.state.STATE_FACTORY).state()
+    store: ObjectStore = cast(ObjectStoreFactory, app.state.OBJECTSTORE_FACTORY).store()
+
+    # Handle message
+    return message_handler.get_fab(request=request, ffs=ffs, state=state, store=store)
+
+
+@rest_request_response(ConfirmMessageReceivedRequest)
+async def confirm_message_received(
+    request: ConfirmMessageReceivedRequest,
+) -> ConfirmMessageReceivedResponse:
+    """Confirm message received."""
+    # Get state from app
+    state: LinkState = cast(LinkStateFactory, app.state.STATE_FACTORY).state()
+    store: ObjectStore = cast(ObjectStoreFactory, app.state.OBJECTSTORE_FACTORY).store()
 
     # Handle message
-    return message_handler.get_fab(request=request, ffs=ffs, state=state)
+    return message_handler.confirm_message_received(
+        request=request, state=state, store=store
+    )
 
 
 routes = [
@@ -171,9 +218,16 @@ routes = [
     Route("/api/v0/fleet/delete-node", delete_node, methods=["POST"]),
     Route("/api/v0/fleet/pull-messages", pull_message, methods=["POST"]),
     Route("/api/v0/fleet/push-messages", push_message, methods=["POST"]),
+    Route("/api/v0/fleet/pull-object", pull_object, methods=["POST"]),
+    Route("/api/v0/fleet/push-object", push_object, methods=["POST"]),
     Route("/api/v0/fleet/send-node-heartbeat", send_node_heartbeat, methods=["POST"]),
     Route("/api/v0/fleet/get-run", get_run, methods=["POST"]),
     Route("/api/v0/fleet/get-fab", get_fab, methods=["POST"]),
+    Route(
+        "/api/v0/fleet/confirm-message-received",
+        confirm_message_received,
+        methods=["POST"],
+    ),
 ]
 
 app: Starlette = Starlette(

flwr/server/superlink/linkstate/in_memory_linkstate.py

@@ -18,6 +18,7 @@
 import threading
 import time
 from bisect import bisect_right
+from collections import defaultdict
 from dataclasses import dataclass, field
 from logging import ERROR, WARNING
 from typing import Optional
@@ -79,6 +80,9 @@ class InMemoryLinkState(LinkState):  # pylint: disable=R0902,R0904
         self.message_res_store: dict[str, Message] = {}
         self.message_ins_id_to_message_res_id: dict[str, str] = {}
 
+        # Map flwr_aid to run_ids for O(1) reverse index lookup
+        self.flwr_aid_to_run_ids: dict[str, set[int]] = defaultdict(set)
+
         self.node_public_keys: set[bytes] = set()
 
         self.lock = threading.RLock()
@@ -398,6 +402,7 @@ class InMemoryLinkState(LinkState):  # pylint: disable=R0902,R0904
         fab_hash: Optional[str],
         override_config: UserConfig,
         federation_options: ConfigRecord,
+        flwr_aid: Optional[str],
     ) -> int:
         """Create a new run for the specified `fab_hash`."""
         # Sample a random int64 as run_id
@@ -421,9 +426,13 @@ class InMemoryLinkState(LinkState):  # pylint: disable=R0902,R0904
                             sub_status="",
                             details="",
                         ),
+                        flwr_aid=flwr_aid if flwr_aid else "",
                     ),
                 )
                 self.run_ids[run_id] = run_record
+                # Add run_id to the flwr_aid_to_run_ids mapping if flwr_aid is provided
+                if flwr_aid:
+                    self.flwr_aid_to_run_ids[flwr_aid].add(run_id)
 
                 # Record federation options. Leave empty if not passed
                 self.federation_options[run_id] = federation_options
@@ -451,9 +460,15 @@ class InMemoryLinkState(LinkState):  # pylint: disable=R0902,R0904
         with self.lock:
             return self.node_public_keys.copy()
 
-    def get_run_ids(self) -> set[int]:
-        """Retrieve all run IDs."""
+    def get_run_ids(self, flwr_aid: Optional[str]) -> set[int]:
+        """Retrieve all run IDs if `flwr_aid` is not specified.
+
+        Otherwise, retrieve all run IDs for the specified `flwr_aid`.
+        """
         with self.lock:
+            if flwr_aid is not None:
+                # Return run IDs for the specified flwr_aid
+                return set(self.flwr_aid_to_run_ids.get(flwr_aid, ()))
             return set(self.run_ids.keys())
 
     def _check_and_tag_inactive_run(self, run_ids: set[int]) -> None:

flwr/server/superlink/linkstate/linkstate.py

@@ -164,12 +164,16 @@ class LinkState(abc.ABC):  # pylint: disable=R0904
         fab_hash: Optional[str],
         override_config: UserConfig,
         federation_options: ConfigRecord,
+        flwr_aid: Optional[str],
     ) -> int:
         """Create a new run for the specified `fab_hash`."""
 
     @abc.abstractmethod
-    def get_run_ids(self) -> set[int]:
-        """Retrieve all run IDs."""
+    def get_run_ids(self, flwr_aid: Optional[str]) -> set[int]:
+        """Retrieve all run IDs if `flwr_aid` is not specified.
+
+        Otherwise, retrieve all run IDs for the specified `flwr_aid`.
+        """
 
     @abc.abstractmethod
     def get_run(self, run_id: int) -> Optional[Run]:

flwr/server/superlink/linkstate/sqlite_linkstate.py

@@ -102,7 +102,8 @@ CREATE TABLE IF NOT EXISTS run(
     finished_at           TEXT,
     sub_status            TEXT,
     details               TEXT,
-    federation_options    BLOB
+    federation_options    BLOB,
+    flwr_aid              TEXT
 );
 """
 
@@ -719,6 +720,7 @@ class SqliteLinkState(LinkState):  # pylint: disable=R0904
         fab_hash: Optional[str],
         override_config: UserConfig,
         federation_options: ConfigRecord,
+        flwr_aid: Optional[str],
     ) -> int:
         """Create a new run for the specified `fab_id` and `fab_version`."""
         # Sample a random int64 as run_id
@@ -735,8 +737,8 @@ class SqliteLinkState(LinkState):  # pylint: disable=R0904
                 "INSERT INTO run "
                 "(run_id, active_until, heartbeat_interval, fab_id, fab_version, "
                 "fab_hash, override_config, federation_options, pending_at, "
-                "starting_at, running_at, finished_at, sub_status, details) "
-                "VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?);"
+                "starting_at, running_at, finished_at, sub_status, details, flwr_aid) "
+                "VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?);"
             )
             override_config_json = json.dumps(override_config)
             data = [
@@ -754,6 +756,7 @@ class SqliteLinkState(LinkState):  # pylint: disable=R0904
                 "",
                 "",
                 "",
+                flwr_aid or "",
             ]
             self.query(query, tuple(data))
             return uint64_run_id
@@ -782,10 +785,18 @@ class SqliteLinkState(LinkState):  # pylint: disable=R0904
         result: set[bytes] = {row["public_key"] for row in rows}
         return result
 
-    def get_run_ids(self) -> set[int]:
-        """Retrieve all run IDs."""
-        query = "SELECT run_id FROM run;"
-        rows = self.query(query)
+    def get_run_ids(self, flwr_aid: Optional[str]) -> set[int]:
+        """Retrieve all run IDs if `flwr_aid` is not specified.
+
+        Otherwise, retrieve all run IDs for the specified `flwr_aid`.
+        """
+        if flwr_aid:
+            rows = self.query(
+                "SELECT run_id FROM run WHERE flwr_aid = ?;",
+                (flwr_aid,),
+            )
+        else:
+            rows = self.query("SELECT run_id FROM run;", ())
         return {convert_sint64_to_uint64(row["run_id"]) for row in rows}
 
     def _check_and_tag_inactive_run(self, run_ids: set[int]) -> None:
@@ -836,6 +847,7 @@ class SqliteLinkState(LinkState):  # pylint: disable=R0904
                     sub_status=row["sub_status"],
                     details=row["details"],
                 ),
+                flwr_aid=row["flwr_aid"],
             )
         log(ERROR, "`run_id` does not exist.")
         return None