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

flwr/server/superlink/serverappio/serverappio_servicer.py

@@ -26,6 +26,8 @@ from flwr.common.constant import SUPERLINK_NODE_ID, Status
 from flwr.common.inflatable import (
     UnexpectedObjectContentError,
     get_descendant_object_ids,
+    get_object_tree,
+    no_object_id_recompute,
 )
 from flwr.common.logger import log
 from flwr.common.serde import (
@@ -50,6 +52,8 @@ from flwr.proto.log_pb2 import (  # pylint: disable=E0611
     PushLogsResponse,
 )
 from flwr.proto.message_pb2 import (  # pylint: disable=E0611
+    ConfirmMessageReceivedRequest,
+    ConfirmMessageReceivedResponse,
     ObjectIDs,
     PullObjectRequest,
     PullObjectResponse,
@@ -107,14 +111,16 @@ class ServerAppIoServicer(serverappio_pb2_grpc.ServerAppIoServicer):
         """Get available nodes."""
         log(DEBUG, "ServerAppIoServicer.GetNodes")
 
-        # Init state
-        state: LinkState = self.state_factory.state()
+        # Init state and store
+        state = self.state_factory.state()
+        store = self.objectstore_factory.store()
 
         # Abort if the run is not running
         abort_if(
             request.run_id,
             [Status.PENDING, Status.STARTING, Status.FINISHED],
             state,
+            store,
             context,
         )
 
@@ -128,14 +134,16 @@ class ServerAppIoServicer(serverappio_pb2_grpc.ServerAppIoServicer):
         """Push a set of Messages."""
         log(DEBUG, "ServerAppIoServicer.PushMessages")
 
-        # Init state
-        state: LinkState = self.state_factory.state()
+        # Init state and store
+        state = self.state_factory.state()
+        store = self.objectstore_factory.store()
 
         # Abort if the run is not running
         abort_if(
             request.run_id,
             [Status.PENDING, Status.STARTING, Status.FINISHED],
             state,
+            store,
             context,
         )
 
@@ -146,8 +154,7 @@ class ServerAppIoServicer(serverappio_pb2_grpc.ServerAppIoServicer):
             detail="`messages_list` must not be empty",
         )
         message_ids: list[Optional[str]] = []
-        while request.messages_list:
-            message_proto = request.messages_list.pop(0)
+        for message_proto in request.messages_list:
             message = message_from_proto(message_proto=message_proto)
             validation_errors = validate_message(message, is_reply_message=False)
             _raise_if(
@@ -164,9 +171,6 @@ class ServerAppIoServicer(serverappio_pb2_grpc.ServerAppIoServicer):
             message_id: Optional[str] = state.store_message_ins(message=message)
             message_ids.append(message_id)
 
-        # Init store
-        store = self.objectstore_factory.store()
-
         # Store Message object to descendants mapping and preregister objects
         objects_to_push = store_mapping_and_register_objects(store, request=request)
 
@@ -183,10 +187,8 @@ class ServerAppIoServicer(serverappio_pb2_grpc.ServerAppIoServicer):
         """Pull a set of Messages."""
         log(DEBUG, "ServerAppIoServicer.PullMessages")
 
-        # Init state
-        state: LinkState = self.state_factory.state()
-
-        # Init store
+        # Init state and store
+        state = self.state_factory.state()
         store = self.objectstore_factory.store()
 
         # Abort if the run is not running
@@ -194,6 +196,7 @@ class ServerAppIoServicer(serverappio_pb2_grpc.ServerAppIoServicer):
             request.run_id,
             [Status.PENDING, Status.STARTING, Status.FINISHED],
             state,
+            store,
             context,
         )
 
@@ -205,14 +208,15 @@ class ServerAppIoServicer(serverappio_pb2_grpc.ServerAppIoServicer):
         # Register messages generated by LinkState in the Store for consistency
         for msg_res in messages_res:
             if msg_res.metadata.src_node_id == SUPERLINK_NODE_ID:
-                descendants = list(get_descendant_object_ids(msg_res))
-                message_obj_id = msg_res.metadata.message_id
+                with no_object_id_recompute():
+                    descendants = list(get_descendant_object_ids(msg_res))
+                    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(descendants + [message_obj_id])
+                store.preregister(request.run_id, get_object_tree(msg_res))
 
         # Delete the instruction Messages and their replies if found
         message_ins_ids_to_delete = {
@@ -328,14 +332,16 @@ class ServerAppIoServicer(serverappio_pb2_grpc.ServerAppIoServicer):
         """Push ServerApp process outputs."""
         log(DEBUG, "ServerAppIoServicer.PushServerAppOutputs")
 
-        # Init state
+        # Init state and store
         state = self.state_factory.state()
+        store = self.objectstore_factory.store()
 
         # Abort if the run is not running
         abort_if(
             request.run_id,
             [Status.PENDING, Status.STARTING, Status.FINISHED],
             state,
+            store,
             context,
         )
 
@@ -348,16 +354,23 @@ class ServerAppIoServicer(serverappio_pb2_grpc.ServerAppIoServicer):
         """Update the status of a run."""
         log(DEBUG, "ServerAppIoServicer.UpdateRunStatus")
 
-        # Init state
+        # Init state and store
         state = self.state_factory.state()
+        store = self.objectstore_factory.store()
 
         # Abort if the run is finished
-        abort_if(request.run_id, [Status.FINISHED], state, context)
+        abort_if(request.run_id, [Status.FINISHED], state, store, context)
 
         # Update the run status
         state.update_run_status(
             run_id=request.run_id, new_status=run_status_from_proto(request.run_status)
         )
+
+        # If the run is finished, delete the run from ObjectStore
+        if request.run_status.status == Status.FINISHED:
+            # Delete all objects related to the run
+            store.delete_objects_in_run(request.run_id)
+
         return UpdateRunStatusResponse()
 
     def PushLogs(
@@ -412,14 +425,16 @@ class ServerAppIoServicer(serverappio_pb2_grpc.ServerAppIoServicer):
         """Push an object to the ObjectStore."""
         log(DEBUG, "ServerAppIoServicer.PushObject")
 
-        # Init state
-        state: LinkState = self.state_factory.state()
+        # Init state and store
+        state = self.state_factory.state()
+        store = self.objectstore_factory.store()
 
         # Abort if the run is not running
         abort_if(
             request.run_id,
             [Status.PENDING, Status.STARTING, Status.FINISHED],
             state,
+            store,
             context,
         )
 
@@ -427,9 +442,6 @@ class ServerAppIoServicer(serverappio_pb2_grpc.ServerAppIoServicer):
             # Cancel insertion in ObjectStore
             context.abort(grpc.StatusCode.FAILED_PRECONDITION, "Unexpected node ID.")
 
-        # Init store
-        store = self.objectstore_factory.store()
-
         # Insert in store
         stored = False
         try:
@@ -449,14 +461,16 @@ class ServerAppIoServicer(serverappio_pb2_grpc.ServerAppIoServicer):
         """Pull an object from the ObjectStore."""
         log(DEBUG, "ServerAppIoServicer.PullObject")
 
-        # Init state
-        state: LinkState = self.state_factory.state()
+        # Init state and store
+        state = self.state_factory.state()
+        store = self.objectstore_factory.store()
 
         # Abort if the run is not running
         abort_if(
             request.run_id,
             [Status.PENDING, Status.STARTING, Status.FINISHED],
             state,
+            store,
             context,
         )
 
@@ -464,9 +478,6 @@ class ServerAppIoServicer(serverappio_pb2_grpc.ServerAppIoServicer):
             # Cancel insertion in ObjectStore
             context.abort(grpc.StatusCode.FAILED_PRECONDITION, "Unexpected node ID.")
 
-        # Init store
-        store = self.objectstore_factory.store()
-
         # Fetch from store
         content = store.get(request.object_id)
         if content is not None:
@@ -478,6 +489,31 @@ class ServerAppIoServicer(serverappio_pb2_grpc.ServerAppIoServicer):
             )
         return PullObjectResponse(object_found=False, object_available=False)
 
+    def ConfirmMessageReceived(
+        self, request: ConfirmMessageReceivedRequest, context: grpc.ServicerContext
+    ) -> ConfirmMessageReceivedResponse:
+        """Confirm message received."""
+        log(DEBUG, "ServerAppIoServicer.ConfirmMessageReceived")
+
+        # Init state and store
+        state = self.state_factory.state()
+        store = self.objectstore_factory.store()
+
+        # Abort if the run is not running
+        abort_if(
+            request.run_id,
+            [Status.PENDING, Status.STARTING, Status.FINISHED],
+            state,
+            store,
+            context,
+        )
+
+        # Delete the message object
+        store.delete(request.message_object_id)
+        store.delete_message_descendant_ids(request.message_object_id)
+
+        return ConfirmMessageReceivedResponse()
+
 
 def _raise_if(validation_error: bool, request_name: str, detail: str) -> None:
     """Raise a `ValueError` with a detailed message if a validation error occurs."""

flwr/server/superlink/simulation/simulationio_servicer.py

@@ -121,6 +121,7 @@ class SimulationIoServicer(simulationio_pb2_grpc.SimulationIoServicer):
             request.run_id,
             [Status.PENDING, Status.STARTING, Status.FINISHED],
             state,
+            None,
             context,
         )
 
@@ -135,7 +136,7 @@ class SimulationIoServicer(simulationio_pb2_grpc.SimulationIoServicer):
         state = self.state_factory.state()
 
         # Abort if the run is finished
-        abort_if(request.run_id, [Status.FINISHED], state, context)
+        abort_if(request.run_id, [Status.FINISHED], state, None, context)
 
         # Update the run status
         state.update_run_status(

flwr/server/superlink/utils.py

@@ -15,11 +15,12 @@
 """SuperLink utilities."""
 
 
-from typing import Union
+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.fleet_pb2 import PushMessagesRequest  # pylint: disable=E0611
 from flwr.proto.message_pb2 import ObjectIDs  # pylint: disable=E0611
@@ -39,6 +40,7 @@ def check_abort(
     run_id: int,
     abort_status_list: list[str],
     state: LinkState,
+    store: Optional[ObjectStore] = None,
 ) -> Union[str, None]:
     """Check if the status of the provided `run_id` is in `abort_status_list`."""
     run_status: RunStatus = state.get_run_status({run_id})[run_id]
@@ -49,6 +51,10 @@ def check_abort(
             msg += " Stopped by user."
         return msg
 
+    # Clear the objects of the run from the store if the run is finished
+    if store and run_status.status == Status.FINISHED:
+        store.delete_objects_in_run(run_id)
+
     return None
 
 
@@ -62,10 +68,11 @@ def abort_if(
     run_id: int,
     abort_status_list: list[str],
     state: LinkState,
+    store: Optional[ObjectStore],
     context: grpc.ServicerContext,
 ) -> None:
     """Abort context if status of the provided `run_id` is in `abort_status_list`."""
-    msg = check_abort(run_id, abort_status_list, state)
+    msg = check_abort(run_id, abort_status_list, state, store)
     abort_grpc_context(msg, context)
 
 
@@ -73,21 +80,27 @@ def store_mapping_and_register_objects(
     store: ObjectStore, request: Union[PushInsMessagesRequest, PushMessagesRequest]
 ) -> dict[str, ObjectIDs]:
     """Store Message object to descendants mapping and preregister objects."""
+    if not request.messages_list:
+        return {}
+
     objects_to_push: dict[str, ObjectIDs] = {}
-    for (
-        message_obj_id,
-        descendant_obj_ids,
-    ) in request.msg_to_descendant_mapping.items():
-        descendants = list(descendant_obj_ids.object_ids)
+
+    # Get run_id from the first message in the list
+    # All messages of a request should in the same run
+    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=message_obj_id, descendant_ids=descendants
+            msg_object_id=msg_object_id, descendant_ids=descendant_ids
         )
 
         # Preregister
-        object_ids_just_registered = store.preregister(descendants + [message_obj_id])
+        object_ids_just_registered = store.preregister(run_id, object_tree)
         # Keep track of objects that need to be pushed
-        objects_to_push[message_obj_id] = ObjectIDs(
+        objects_to_push[msg_object_id] = ObjectIDs(
             object_ids=object_ids_just_registered
         )
 

flwr/supercore/object_store/in_memory_object_store.py

@@ -15,44 +15,95 @@
 """Flower in-memory ObjectStore implementation."""
 
 
+import threading
+from dataclasses import dataclass
 from typing import Optional
 
-from flwr.common.inflatable import get_object_id, is_valid_sha256_hash
+from flwr.common.inflatable import (
+    get_object_children_ids_from_object_content,
+    get_object_id,
+    is_valid_sha256_hash,
+    iterate_object_tree,
+)
 from flwr.common.inflatable_utils import validate_object_content
+from flwr.proto.message_pb2 import ObjectTree  # pylint: disable=E0611
 
 from .object_store import NoObjectInStoreError, ObjectStore
 
 
+@dataclass
+class ObjectEntry:
+    """Data class representing an object entry in the store."""
+
+    content: bytes
+    is_available: bool
+    ref_count: int  # Number of references (direct parents) to this object
+    runs: set[int]  # Set of run IDs that used this object
+
+
 class InMemoryObjectStore(ObjectStore):
     """In-memory implementation of the ObjectStore interface."""
 
     def __init__(self, verify: bool = True) -> None:
         self.verify = verify
-        self.store: dict[str, bytes] = {}
-        # Mapping the Object ID of a message to the list of children object IDs
-        self.msg_children_objects_mapping: dict[str, list[str]] = {}
-
-    def preregister(self, object_ids: list[str]) -> list[str]:
+        self.store: dict[str, ObjectEntry] = {}
+        self.lock_store = threading.RLock()
+        # Mapping the Object ID of a message to the list of descendant object IDs
+        self.msg_descendant_objects_mapping: dict[str, list[str]] = {}
+        self.lock_msg_mapping = threading.RLock()
+        # Mapping each run ID to a set of object IDs that are used in that run
+        self.run_objects_mapping: dict[int, set[str]] = {}
+
+    def preregister(self, run_id: int, object_tree: ObjectTree) -> list[str]:
         """Identify and preregister missing objects."""
         new_objects = []
-        for obj_id in object_ids:
+        if run_id not in self.run_objects_mapping:
+            self.run_objects_mapping[run_id] = set()
+
+        for tree_node in iterate_object_tree(object_tree):
+            obj_id = tree_node.object_id
             # Verify object ID format (must be a valid sha256 hash)
             if not is_valid_sha256_hash(obj_id):
                 raise ValueError(f"Invalid object ID format: {obj_id}")
-            if obj_id not in self.store:
-                self.store[obj_id] = b""
-                new_objects.append(obj_id)
+            with self.lock_store:
+                if obj_id not in self.store:
+                    self.store[obj_id] = ObjectEntry(
+                        content=b"",  # Initially empty content
+                        is_available=False,  # Initially not available
+                        ref_count=0,  # Reference count starts at 0
+                        runs={run_id},  # Start with the current run ID
+                    )
+
+                    # Increment the reference count for all its children
+                    # Post-order traversal ensures that children are registered
+                    # before parents
+                    for child_node in tree_node.children:
+                        child_id = child_node.object_id
+                        self.store[child_id].ref_count += 1
+
+                    # Add the object ID to the run's mapping
+                    self.run_objects_mapping[run_id].add(obj_id)
+
+                    # Add to the list of new objects
+                    new_objects.append(obj_id)
+                else:
+                    # Object is in store, retrieve it
+                    obj_entry = self.store[obj_id]
+
+                    # Add to the list of new objects if not available
+                    if not obj_entry.is_available:
+                        new_objects.append(obj_id)
+
+                    # If the object is already registered but not in this run,
+                    # add the run ID to its runs
+                    if obj_id not in self.run_objects_mapping[run_id]:
+                        obj_entry.runs.add(run_id)
+                        self.run_objects_mapping[run_id].add(obj_id)
 
         return new_objects
 
     def put(self, object_id: str, object_content: bytes) -> None:
         """Put an object into the store."""
-        # Only allow adding the object if it has been preregistered
-        if object_id not in self.store:
-            raise NoObjectInStoreError(
-                f"Object with ID '{object_id}' was not pre-registered."
-            )
-
         if self.verify:
             # Verify object_id and object_content match
             object_id_from_content = get_object_id(object_content)
@@ -62,41 +113,117 @@ class InMemoryObjectStore(ObjectStore):
             # Validate object content
             validate_object_content(content=object_content)
 
-        # Return if object is already present in the store
-        if self.store[object_id] != b"":
-            return
+        with self.lock_store:
+            # Only allow adding the object if it has been preregistered
+            if object_id not in self.store:
+                raise NoObjectInStoreError(
+                    f"Object with ID '{object_id}' was not pre-registered."
+                )
+
+            # Return if object is already present in the store
+            if self.store[object_id].is_available:
+                return
 
-        self.store[object_id] = object_content
+            # Update the object entry in the store
+            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."""
-        self.msg_children_objects_mapping[msg_object_id] = descendant_ids
+        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."""
-        if msg_object_id not in self.msg_children_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_children_objects_mapping[msg_object_id]
+        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."""
-        return self.store.get(object_id)
+        with self.lock_store:
+            # Check if the object ID is pre-registered
+            if object_id not in self.store:
+                return None
+
+            # Return content (if not yet available, it will b"")
+            return self.store[object_id].content
 
     def delete(self, object_id: str) -> None:
-        """Delete an object from the store."""
-        if object_id in self.store:
-            del self.store[object_id]
+        """Delete an object and its unreferenced descendants from the store."""
+        with self.lock_store:
+            # If the object is not in the store, nothing to delete
+            if (object_entry := self.store.get(object_id)) is None:
+                return
+
+            # Delete the object if it has no references left
+            if object_entry.ref_count == 0:
+                del self.store[object_id]
+
+                # Remove the object from the run's mapping
+                for run_id in object_entry.runs:
+                    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:
+                    self.store[child_id].ref_count -= 1
+
+                    # Recursively try to delete the child object
+                    self.delete(child_id)
+
+    def delete_objects_in_run(self, run_id: int) -> None:
+        """Delete all objects that were registered in a specific run."""
+        with self.lock_store:
+            if run_id not in self.run_objects_mapping:
+                return
+            for object_id in list(self.run_objects_mapping[run_id]):
+                # Check if the object is still in the store
+                if (object_entry := self.store.get(object_id)) is None:
+                    continue
+
+                # Remove the run ID from the object's runs
+                object_entry.runs.discard(run_id)
+
+                # Only message objects are allowed to have a `ref_count` of 0,
+                # and every message object must have a `ref_count` of 0
+                if object_entry.ref_count == 0:
+                    # 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]
 
     def clear(self) -> None:
         """Clear the store."""
-        self.store.clear()
+        with self.lock_store:
+            self.store.clear()
+            self.msg_descendant_objects_mapping.clear()
+            self.run_objects_mapping.clear()
 
     def __contains__(self, object_id: str) -> bool:
         """Check if an object_id is in the store."""
-        return object_id in self.store
+        with self.lock_store:
+            return object_id in self.store
+
+    def __len__(self) -> int:
+        """Get the number of objects in the store."""
+        with self.lock_store:
+            return len(self.store)