flwr-nightly 1.19.0.dev20250527__py3-none-any.whl → 1.19.0.dev20250529__py3-none-any.whl

flwr/server/superlink/linkstate/sqlite_linkstate.py

@@ -24,7 +24,6 @@ import time
 from collections.abc import Sequence
 from logging import DEBUG, ERROR, WARNING
 from typing import Any, Optional, Union, cast
-from uuid import UUID, uuid4
 
 from flwr.common import Context, Message, Metadata, log, now
 from flwr.common.constant import (
@@ -251,19 +250,15 @@ class SqliteLinkState(LinkState):  # pylint: disable=R0904
 
         return result
 
-    def store_message_ins(self, message: Message) -> Optional[UUID]:
+    def store_message_ins(self, message: Message) -> Optional[str]:
         """Store one Message."""
         # Validate message
         errors = validate_message(message=message, is_reply_message=False)
         if any(errors):
             log(ERROR, errors)
             return None
-        # Create message_id
-        message_id = uuid4()
 
         # Store Message
-        # pylint: disable-next=W0212
-        message.metadata._message_id = str(message_id)  # type: ignore
         data = (message_to_dict(message),)
 
         # Convert values from uint64 to sint64 for SQLite
@@ -303,7 +298,7 @@ class SqliteLinkState(LinkState):  # pylint: disable=R0904
         # This may need to be changed in the future version with more integrity checks.
         self.query(query, data)
 
-        return message_id
+        return message.metadata.message_id
 
     def get_message_ins(self, node_id: int, limit: Optional[int]) -> list[Message]:
         """Get all Messages that have not been delivered yet."""
@@ -366,7 +361,7 @@ class SqliteLinkState(LinkState):  # pylint: disable=R0904
 
         return result
 
-    def store_message_res(self, message: Message) -> Optional[UUID]:
+    def store_message_res(self, message: Message) -> Optional[str]:
         """Store one Message."""
         # Validate message
         errors = validate_message(message=message, is_reply_message=True)
@@ -418,12 +413,7 @@ class SqliteLinkState(LinkState):  # pylint: disable=R0904
             )
             return None
 
-        # Create message_id
-        message_id = uuid4()
-
         # Store Message
-        # pylint: disable-next=W0212
-        message.metadata._message_id = str(message_id)  # type: ignore
         data = (message_to_dict(message),)
 
         # Convert values from uint64 to sint64 for SQLite
@@ -442,12 +432,12 @@ class SqliteLinkState(LinkState):  # pylint: disable=R0904
             log(ERROR, "`run` is invalid")
             return None
 
-        return message_id
+        return message.metadata.message_id
 
-    def get_message_res(self, message_ids: set[UUID]) -> list[Message]:
+    def get_message_res(self, message_ids: set[str]) -> list[Message]:
         """Get reply Messages for the given Message IDs."""
         # pylint: disable-msg=too-many-locals
-        ret: dict[UUID, Message] = {}
+        ret: dict[str, Message] = {}
 
         # Verify Message IDs
         current = time.time()
@@ -457,12 +447,12 @@ class SqliteLinkState(LinkState):  # pylint: disable=R0904
             WHERE message_id IN ({",".join(["?"] * len(message_ids))});
         """
         rows = self.query(query, tuple(str(message_id) for message_id in message_ids))
-        found_message_ins_dict: dict[UUID, Message] = {}
+        found_message_ins_dict: dict[str, Message] = {}
         for row in rows:
             convert_sint64_values_in_dict_to_uint64(
                 row, ["run_id", "src_node_id", "dst_node_id"]
             )
-            found_message_ins_dict[UUID(row["message_id"])] = dict_to_message(row)
+            found_message_ins_dict[row["message_id"]] = dict_to_message(row)
 
         ret = verify_message_ids(
             inquired_message_ids=message_ids,
@@ -551,7 +541,7 @@ class SqliteLinkState(LinkState):  # pylint: disable=R0904
         result: dict[str, int] = rows[0]
         return result["num"]
 
-    def delete_messages(self, message_ins_ids: set[UUID]) -> None:
+    def delete_messages(self, message_ins_ids: set[str]) -> None:
         """Delete a Message and its reply based on provided Message IDs."""
         if not message_ins_ids:
             return
@@ -577,7 +567,7 @@ class SqliteLinkState(LinkState):  # pylint: disable=R0904
             self.conn.execute(query_1, data)
             self.conn.execute(query_2, data)
 
-    def get_message_ids_from_run_id(self, run_id: int) -> set[UUID]:
+    def get_message_ids_from_run_id(self, run_id: int) -> set[str]:
         """Get all instruction Message IDs for the given run_id."""
         if self.conn is None:
             raise AttributeError("LinkState not initialized")
@@ -594,7 +584,7 @@ class SqliteLinkState(LinkState):  # pylint: disable=R0904
         with self.conn:
             rows = self.conn.execute(query, data).fetchall()
 
-        return {UUID(row["message_id"]) for row in rows}
+        return {row["message_id"] for row in rows}
 
     def create_node(self, heartbeat_interval: float) -> int:
         """Create, store in the link state, and return `node_id`."""

flwr/server/superlink/linkstate/utils.py

@@ -17,7 +17,7 @@
 
 from os import urandom
 from typing import Optional
-from uuid import UUID, uuid4
+from uuid import uuid4
 
 from flwr.common import ConfigRecord, Context, Error, Message, Metadata, now, serde
 from flwr.common.constant import (
@@ -273,7 +273,7 @@ def create_message_error_unavailable_res_message(
     )
 
 
-def create_message_error_unavailable_ins_message(reply_to_message_id: UUID) -> Message:
+def create_message_error_unavailable_ins_message(reply_to_message_id: str) -> Message:
     """Error to indicate that the enquired Message had expired before reply arrived or
     that it isn't found."""
     metadata = Metadata(
@@ -281,7 +281,7 @@ def create_message_error_unavailable_ins_message(reply_to_message_id: UUID) -> M
         message_id=str(uuid4()),
         src_node_id=SUPERLINK_NODE_ID,
         dst_node_id=SUPERLINK_NODE_ID,
-        reply_to_message_id=str(reply_to_message_id),
+        reply_to_message_id=reply_to_message_id,
         group_id="",  # Unknown
         message_type=MessageType.SYSTEM,
         created_at=now().timestamp(),
@@ -303,18 +303,18 @@ def message_ttl_has_expired(message_metadata: Metadata, current_time: float) ->
 
 
 def verify_message_ids(
-    inquired_message_ids: set[UUID],
-    found_message_ins_dict: dict[UUID, Message],
+    inquired_message_ids: set[str],
+    found_message_ins_dict: dict[str, Message],
     current_time: Optional[float] = None,
     update_set: bool = True,
-) -> dict[UUID, Message]:
+) -> dict[str, Message]:
     """Verify found Messages and generate error Messages for invalid ones.
 
     Parameters
     ----------
-    inquired_message_ids : set[UUID]
+    inquired_message_ids : set[str]
         Set of Message IDs for which to generate error Message if invalid.
-    found_message_ins_dict : dict[UUID, Message]
+    found_message_ins_dict : dict[str, Message]
         Dictionary containing all found Message indexed by their IDs.
     current_time : Optional[float] (default: None)
         The current time to check for expiration. If set to `None`, the current time
@@ -325,7 +325,7 @@ def verify_message_ids(
 
     Returns
     -------
-    dict[UUID, Message]
+    dict[str, Message]
         A dictionary of error Message indexed by the corresponding ID of the message
         they are a reply of.
     """
@@ -345,19 +345,19 @@ def verify_message_ids(
 
 
 def verify_found_message_replies(
-    inquired_message_ids: set[UUID],
-    found_message_ins_dict: dict[UUID, Message],
+    inquired_message_ids: set[str],
+    found_message_ins_dict: dict[str, Message],
     found_message_res_list: list[Message],
     current_time: Optional[float] = None,
     update_set: bool = True,
-) -> dict[UUID, Message]:
+) -> dict[str, Message]:
     """Verify found Message replies and generate error Message for invalid ones.
 
     Parameters
     ----------
-    inquired_message_ids : set[UUID]
+    inquired_message_ids : set[str]
         Set of Message IDs for which to generate error Message if invalid.
-    found_message_ins_dict : dict[UUID, Message]
+    found_message_ins_dict : dict[str, Message]
         Dictionary containing all found instruction Messages indexed by their IDs.
     found_message_res_list : dict[Message, Message]
         List of found Message to be verified.
@@ -370,13 +370,13 @@ def verify_found_message_replies(
 
     Returns
     -------
-    dict[UUID, Message]
+    dict[str, Message]
         A dictionary of Message indexed by the corresponding Message ID.
     """
-    ret_dict: dict[UUID, Message] = {}
+    ret_dict: dict[str, Message] = {}
     current = current_time if current_time else now().timestamp()
     for message_res in found_message_res_list:
-        message_ins_id = UUID(message_res.metadata.reply_to_message_id)
+        message_ins_id = message_res.metadata.reply_to_message_id
         if update_set:
             inquired_message_ids.remove(message_ins_id)
         # Check if the reply Message has expired
@@ -390,21 +390,21 @@ def verify_found_message_replies(
 
 
 def check_node_availability_for_in_message(
-    inquired_in_message_ids: set[UUID],
-    found_in_message_dict: dict[UUID, Message],
+    inquired_in_message_ids: set[str],
+    found_in_message_dict: dict[str, Message],
     node_id_to_online_until: dict[int, float],
     current_time: Optional[float] = None,
     update_set: bool = True,
-) -> dict[UUID, Message]:
+) -> dict[str, Message]:
     """Check node availability for given Message and generate error reply Message if
     unavailable. A Message error indicating node unavailability will be generated for
     each given Message whose destination node is offline or non-existent.
 
     Parameters
     ----------
-    inquired_in_message_ids : set[UUID]
+    inquired_in_message_ids : set[str]
         Set of Message IDs for which to check destination node availability.
-    found_in_message_dict : dict[UUID, Message]
+    found_in_message_dict : dict[str, Message]
         Dictionary containing all found Message indexed by their IDs.
     node_id_to_online_until : dict[int, float]
         Dictionary mapping node IDs to their online-until timestamps.
@@ -417,7 +417,7 @@ def check_node_availability_for_in_message(
 
     Returns
     -------
-    dict[UUID, Message]
+    dict[str, Message]
         A dictionary of error Message indexed by the corresponding Message ID.
     """
     ret_dict = {}

flwr/server/superlink/serverappio/serverappio_servicer.py

@@ -18,7 +18,6 @@
 import threading
 from logging import DEBUG, INFO
 from typing import Optional
-from uuid import UUID
 
 import grpc
 
@@ -81,6 +80,8 @@ from flwr.server.superlink.utils import abort_if
 from flwr.server.utils.validator import validate_message
 from flwr.supercore.object_store import ObjectStoreFactory
 
+from ..utils import store_mapping_and_register_objects
+
 
 class ServerAppIoServicer(serverappio_pb2_grpc.ServerAppIoServicer):
     """ServerAppIo API servicer."""
@@ -140,7 +141,7 @@ class ServerAppIoServicer(serverappio_pb2_grpc.ServerAppIoServicer):
             request_name="PushMessages",
             detail="`messages_list` must not be empty",
         )
-        message_ids: list[Optional[UUID]] = []
+        message_ids: list[Optional[str]] = []
         while request.messages_list:
             message_proto = request.messages_list.pop(0)
             message = message_from_proto(message_proto=message_proto)
@@ -156,13 +157,20 @@ class ServerAppIoServicer(serverappio_pb2_grpc.ServerAppIoServicer):
                 detail="`Message.metadata` has mismatched `run_id`",
             )
             # Store
-            message_id: Optional[UUID] = state.store_message_ins(message=message)
+            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)
+
         return PushInsMessagesResponse(
             message_ids=[
                 str(message_id) if message_id else "" for message_id in message_ids
-            ]
+            ],
+            objects_to_push=objects_to_push,
         )
 
     def PullMessages(
@@ -182,17 +190,14 @@ class ServerAppIoServicer(serverappio_pb2_grpc.ServerAppIoServicer):
             context,
         )
 
-        # Convert each message_id str to UUID
-        message_ids: set[UUID] = {
-            UUID(message_id) for message_id in request.message_ids
-        }
-
         # Read from state
-        messages_res: list[Message] = state.get_message_res(message_ids=message_ids)
+        messages_res: list[Message] = state.get_message_res(
+            message_ids=set(request.message_ids)
+        )
 
         # Delete the instruction Messages and their replies if found
         message_ins_ids_to_delete = {
-            UUID(msg_res.metadata.reply_to_message_id) for msg_res in messages_res
+            msg_res.metadata.reply_to_message_id for msg_res in messages_res
         }
 
         state.delete_messages(message_ins_ids=message_ins_ids_to_delete)

flwr/server/superlink/utils.py

@@ -21,7 +21,11 @@ import grpc
 
 from flwr.common.constant import Status, SubStatus
 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
+from flwr.proto.serverappio_pb2 import PushInsMessagesRequest  # pylint: disable=E0611
 from flwr.server.superlink.linkstate import LinkState
+from flwr.supercore.object_store import ObjectStore
 
 _STATUS_TO_MSG = {
     Status.PENDING: "Run is pending.",
@@ -63,3 +67,28 @@ def abort_if(
     """Abort context if status of the provided `run_id` is in `abort_status_list`."""
     msg = check_abort(run_id, abort_status_list, state)
     abort_grpc_context(msg, context)
+
+
+def store_mapping_and_register_objects(
+    store: ObjectStore, request: Union[PushInsMessagesRequest, PushMessagesRequest]
+) -> dict[str, ObjectIDs]:
+    """Store Message object to descendants mapping and preregister objects."""
+    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)
+        # Store mapping
+        store.set_message_descendant_ids(
+            msg_object_id=message_obj_id, descendant_ids=descendants
+        )
+
+        # Preregister
+        object_ids_just_registered = store.preregister(descendants + [message_obj_id])
+        # Keep track of objects that need to be pushed
+        objects_to_push[message_obj_id] = ObjectIDs(
+            object_ids=object_ids_just_registered
+        )
+
+    return objects_to_push

flwr/server/utils/validator.py

@@ -27,8 +27,8 @@ def validate_message(message: Message, is_reply_message: bool) -> list[str]:
     validation_errors = []
     metadata = message.metadata
 
-    if metadata.message_id != "":
-        validation_errors.append("non-empty `metadata.message_id`")
+    if metadata.message_id == "":
+        validation_errors.append("empty `metadata.message_id`")
 
     # Created/delivered/TTL/Pushed
     if (

flwr/supercore/object_store/in_memory_object_store.py

@@ -28,12 +28,27 @@ class InMemoryObjectStore(ObjectStore):
     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]:
+        """Identify and preregister missing objects."""
+        new_objects = []
+        for obj_id in object_ids:
+            # 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)
+
+        return new_objects
 
     def put(self, object_id: str, object_content: bytes) -> None:
         """Put an object into the store."""
-        # Verify object ID format (must be a valid sha256 hash)
-        if not is_valid_sha256_hash(object_id):
-            raise ValueError(f"Invalid object ID format: {object_id}")
+        # Only allow adding the object if it has been preregistered
+        if object_id not in self.store:
+            raise KeyError(f"Object with id {object_id} was not preregistered.")
 
         # Verify object_id and object_content match
         if self.verify:
@@ -42,11 +57,22 @@ class InMemoryObjectStore(ObjectStore):
                 raise ValueError(f"Object ID {object_id} does not match content hash")
 
         # Return if object is already present in the store
-        if object_id in self.store:
+        if self.store[object_id] != b"":
             return
 
         self.store[object_id] = object_content
 
+    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
+
+    def get_message_descendant_ids(self, msg_object_id: str) -> list[str]:
+        """Retrieve the object IDs of all descendants of a given Message."""
+        return self.msg_children_objects_mapping[msg_object_id]
+
     def get(self, object_id: str) -> Optional[bytes]:
         """Get an object from the store."""
         return self.store.get(object_id)

flwr/supercore/object_store/object_store.py

@@ -26,6 +26,23 @@ class ObjectStore(abc.ABC):
     delete objects identified by object IDs.
     """
 
+    @abc.abstractmethod
+    def preregister(self, object_ids: list[str]) -> list[str]:
+        """Identify and preregister missing objects in the `ObjectStore`.
+
+        Parameters
+        ----------
+        object_ids : list[str]
+            A list of object IDs to check against the store. Any object ID not already
+            present will be preregistered.
+
+        Returns
+        -------
+        list[str]
+            A list of object IDs that were not present in the `ObjectStore` and have now
+            been preregistered.
+        """
+
     @abc.abstractmethod
     def put(self, object_id: str, object_content: bytes) -> None:
         """Put an object into the store.
@@ -33,7 +50,7 @@ class ObjectStore(abc.ABC):
         Parameters
         ----------
         object_id : str
-            The object_id under which to store the object.
+            The object_id under which to store the object. Must be preregistered.
         object_content : bytes
             The deflated object to store.
         """
@@ -70,6 +87,36 @@ 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 __contains__(self, object_id: str) -> bool:
         """Check if an object_id is in the store.

flwr/superexec/exec_servicer.py

@@ -19,7 +19,6 @@ import time
 from collections.abc import Generator
 from logging import ERROR, INFO
 from typing import Any, Optional
-from uuid import UUID
 
 import grpc
 
@@ -163,7 +162,7 @@ class ExecServicer(exec_pb2_grpc.ExecServicer):
         )
 
         if update_success:
-            message_ids: set[UUID] = state.get_message_ids_from_run_id(request.run_id)
+            message_ids: set[str] = state.get_message_ids_from_run_id(request.run_id)
 
             # Delete Messages and their replies for the `run_id`
             state.delete_messages(message_ids)

{flwr_nightly-1.19.0.dev20250527.dist-info → flwr_nightly-1.19.0.dev20250529.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.3
 Name: flwr-nightly
-Version: 1.19.0.dev20250527
+Version: 1.19.0.dev20250529
 Summary: Flower: A Friendly Federated AI Framework
 License: Apache-2.0
 Keywords: Artificial Intelligence,Federated AI,Federated Analytics,Federated Evaluation,Federated Learning,Flower,Machine Learning