flwr-nightly 1.20.0.dev20250714__tar.gz → 1.20.0.dev20250716__tar.gz

{flwr_nightly-1.20.0.dev20250714 → flwr_nightly-1.20.0.dev20250716}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.3
 Name: flwr-nightly
-Version: 1.20.0.dev20250714
+Version: 1.20.0.dev20250716
 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
@@ -34,6 +34,7 @@ Provides-Extra: simulation
 Requires-Dist: click (<8.2.0)
 Requires-Dist: cryptography (>=44.0.1,<45.0.0)
 Requires-Dist: grpcio (>=1.62.3,<2.0.0,!=1.65.0)
+Requires-Dist: grpcio-health-checking (>=1.62.3,<2.0.0)
 Requires-Dist: iterators (>=0.0.2,<0.0.3)
 Requires-Dist: numpy (>=1.26.0,<3.0.0)
 Requires-Dist: pathspec (>=0.12.1,<0.13.0)

{flwr_nightly-1.20.0.dev20250714 → flwr_nightly-1.20.0.dev20250716}/py/flwr/client/grpc_rere_client/connection.py

@@ -31,11 +31,12 @@ from flwr.common.heartbeat import HeartbeatSender
 from flwr.common.inflatable import (
     get_all_nested_objects,
     get_object_tree,
+    iterate_object_tree,
     no_object_id_recompute,
 )
-from flwr.common.inflatable_grpc_utils import (
-    make_pull_object_fn_grpc,
-    make_push_object_fn_grpc,
+from flwr.common.inflatable_protobuf_utils import (
+    make_pull_object_fn_protobuf,
+    make_push_object_fn_protobuf,
 )
 from flwr.common.inflatable_utils import (
     inflate_object_from_contents,
@@ -273,10 +274,11 @@ def grpc_request_response(  # pylint: disable=R0913,R0914,R0915,R0917
         if message_proto:
             msg_id = message_proto.metadata.message_id
             run_id = message_proto.metadata.run_id
+            object_tree = response.message_object_trees[0]
             all_object_contents = pull_objects(
-                list(response.objects_to_pull[msg_id].object_ids) + [msg_id],
-                pull_object_fn=make_pull_object_fn_grpc(
-                    pull_object_grpc=stub.PullObject,
+                [tree.object_id for tree in iterate_object_tree(object_tree)],
+                pull_object_fn=make_pull_object_fn_protobuf(
+                    pull_object_protobuf=stub.PullObject,
                     node=node,
                     run_id=run_id,
                 ),
@@ -328,8 +330,8 @@ def grpc_request_response(  # pylint: disable=R0913,R0914,R0915,R0917
                 )
                 push_objects(
                     all_objects,
-                    push_object_fn=make_push_object_fn_grpc(
-                        push_object_grpc=stub.PushObject,
+                    push_object_fn=make_push_object_fn_protobuf(
+                        push_object_protobuf=stub.PushObject,
                         node=node,
                         run_id=message.metadata.run_id,
                     ),

{flwr_nightly-1.20.0.dev20250714 → flwr_nightly-1.20.0.dev20250716}/py/flwr/client/rest_client/connection.py

@@ -14,6 +14,7 @@
 # ==============================================================================
 """Contextmanager for a REST request-response channel to the Flower server."""
 
+
 from collections.abc import Iterator
 from contextlib import contextmanager
 from logging import DEBUG, ERROR, INFO, WARN
@@ -30,11 +31,12 @@ from flwr.common.heartbeat import HeartbeatSender
 from flwr.common.inflatable import (
     get_all_nested_objects,
     get_object_tree,
+    iterate_object_tree,
     no_object_id_recompute,
 )
-from flwr.common.inflatable_rest_utils import (
-    make_pull_object_fn_rest,
-    make_push_object_fn_rest,
+from flwr.common.inflatable_protobuf_utils import (
+    make_pull_object_fn_protobuf,
+    make_push_object_fn_protobuf,
 )
 from flwr.common.inflatable_utils import (
     inflate_object_from_contents,
@@ -333,10 +335,11 @@ def http_request_response(  # pylint: disable=R0913,R0914,R0915,R0917
                 return res
 
             try:
+                object_tree = res.message_object_trees[0]
                 all_object_contents = pull_objects(
-                    list(res.objects_to_pull[msg_id].object_ids) + [msg_id],
-                    pull_object_fn=make_pull_object_fn_rest(
-                        pull_object_rest=fn,
+                    [tree.object_id for tree in iterate_object_tree(object_tree)],
+                    pull_object_fn=make_pull_object_fn_protobuf(
+                        pull_object_protobuf=fn,
                         node=node,
                         run_id=run_id,
                     ),
@@ -413,8 +416,8 @@ def http_request_response(  # pylint: disable=R0913,R0914,R0915,R0917
                 try:
                     push_objects(
                         all_objects,
-                        push_object_fn=make_push_object_fn_rest(
-                            push_object_rest=fn,
+                        push_object_fn=make_push_object_fn_protobuf(
+                            push_object_protobuf=fn,
                             node=node,
                             run_id=message_proto.metadata.run_id,
                         ),

{flwr_nightly-1.20.0.dev20250714 → flwr_nightly-1.20.0.dev20250716}/py/flwr/common/grpc.py

@@ -23,6 +23,9 @@ from logging import DEBUG, ERROR
 from typing import Any, Callable, Optional
 
 import grpc
+from grpc_health.v1.health_pb2_grpc import add_HealthServicer_to_server
+
+from flwr.supercore.grpc_health import SimpleHealthServicer
 
 from .address import is_port_in_use
 from .logger import log
@@ -98,7 +101,7 @@ def valid_certificates(certificates: tuple[bytes, bytes, bytes]) -> bool:
     return is_valid
 
 
-def generic_create_grpc_server(  # pylint: disable=too-many-arguments,R0917
+def generic_create_grpc_server(  # pylint: disable=too-many-arguments, R0914, R0917
     servicer_and_add_fn: tuple[Any, AddServicerToServerFn],
     server_address: str,
     max_concurrent_workers: int = 1000,
@@ -106,6 +109,7 @@ def generic_create_grpc_server(  # pylint: disable=too-many-arguments,R0917
     keepalive_time_ms: int = 210000,
     certificates: Optional[tuple[bytes, bytes, bytes]] = None,
     interceptors: Optional[Sequence[grpc.ServerInterceptor]] = None,
+    health_servicer: Optional[Any] = None,
 ) -> grpc.Server:
     """Create a gRPC server with a single servicer.
 
@@ -153,6 +157,10 @@ def generic_create_grpc_server(  # pylint: disable=too-many-arguments,R0917
             * server private key.
     interceptors : Optional[Sequence[grpc.ServerInterceptor]] (default: None)
         A list of gRPC interceptors.
+    health_servicer : Optional[Any] (default: None)
+        An optional health servicer to add to the server. If provided, it should be an
+        instance of a class that inherits the `HealthServicer` class.
+        If None is provided, `SimpleHealthServicer` will be used by default.
 
     Returns
     -------
@@ -203,6 +211,9 @@ def generic_create_grpc_server(  # pylint: disable=too-many-arguments,R0917
     )
     add_servicer_to_server_fn(servicer, server)
 
+    # Enable health service
+    add_HealthServicer_to_server(health_servicer or SimpleHealthServicer(), server)
+
     if certificates is not None:
         if not valid_certificates(certificates):
             sys.exit(1)

flwr_nightly-1.20.0.dev20250714/py/flwr/common/inflatable_grpc_utils.py → flwr_nightly-1.20.0.dev20250716/py/flwr/common/inflatable_protobuf_utils.py

@@ -28,8 +28,8 @@ from flwr.proto.node_pb2 import Node  # pylint: disable=E0611
 from .inflatable_utils import ObjectIdNotPreregisteredError, ObjectUnavailableError
 
 
-def make_pull_object_fn_grpc(
-    pull_object_grpc: Callable[[PullObjectRequest], PullObjectResponse],
+def make_pull_object_fn_protobuf(
+    pull_object_protobuf: Callable[[PullObjectRequest], PullObjectResponse],
     node: Node,
     run_id: int,
 ) -> Callable[[str], bytes]:
@@ -37,8 +37,9 @@ def make_pull_object_fn_grpc(
 
     Parameters
     ----------
-    pull_object_grpc : Callable[[PullObjectRequest], PullObjectResponse]
-        The gRPC function to pull objects, e.g., `FleetStub.PullObject`.
+    pull_object_protobuf : Callable[[PullObjectRequest], PullObjectResponse]
+        A callable that takes a `PullObjectRequest` and returns a `PullObjectResponse`.
+        This function is typically backed by a gRPC client stub.
     node : Node
         The node making the request.
     run_id : int
@@ -54,7 +55,7 @@ def make_pull_object_fn_grpc(
 
     def pull_object_fn(object_id: str) -> bytes:
         request = PullObjectRequest(node=node, run_id=run_id, object_id=object_id)
-        response: PullObjectResponse = pull_object_grpc(request)
+        response: PullObjectResponse = pull_object_protobuf(request)
         if not response.object_found:
             raise ObjectIdNotPreregisteredError(object_id)
         if not response.object_available:
@@ -64,8 +65,8 @@ def make_pull_object_fn_grpc(
     return pull_object_fn
 
 
-def make_push_object_fn_grpc(
-    push_object_grpc: Callable[[PushObjectRequest], PushObjectResponse],
+def make_push_object_fn_protobuf(
+    push_object_protobuf: Callable[[PushObjectRequest], PushObjectResponse],
     node: Node,
     run_id: int,
 ) -> Callable[[str, bytes], None]:
@@ -73,8 +74,9 @@ def make_push_object_fn_grpc(
 
     Parameters
     ----------
-    push_object_grpc : Callable[[PushObjectRequest], PushObjectResponse]
-        The gRPC function to push objects, e.g., `FleetStub.PushObject`.
+    push_object_protobuf : Callable[[PushObjectRequest], PushObjectResponse]
+        A callable that takes a `PushObjectRequest` and returns a `PushObjectResponse`.
+        This function is typically backed by a gRPC client stub.
     node : Node
         The node making the request.
     run_id : int
@@ -92,7 +94,7 @@ def make_push_object_fn_grpc(
         request = PushObjectRequest(
             node=node, run_id=run_id, object_id=object_id, object_content=object_content
         )
-        response: PushObjectResponse = push_object_grpc(request)
+        response: PushObjectResponse = push_object_protobuf(request)
         if not response.stored:
             raise ObjectIdNotPreregisteredError(object_id)
 

flwr_nightly-1.20.0.dev20250716/py/flwr/proto/fleet_pb2.py

@@ -0,0 +1,57 @@
+# -*- coding: utf-8 -*-
+# Generated by the protocol buffer compiler.  DO NOT EDIT!
+# source: flwr/proto/fleet.proto
+# Protobuf Python Version: 4.25.1
+"""Generated protocol buffer code."""
+from google.protobuf import descriptor as _descriptor
+from google.protobuf import descriptor_pool as _descriptor_pool
+from google.protobuf import symbol_database as _symbol_database
+from google.protobuf.internal import builder as _builder
+# @@protoc_insertion_point(imports)
+
+_sym_db = _symbol_database.Default()
+
+
+from flwr.proto import heartbeat_pb2 as flwr_dot_proto_dot_heartbeat__pb2
+from flwr.proto import node_pb2 as flwr_dot_proto_dot_node__pb2
+from flwr.proto import run_pb2 as flwr_dot_proto_dot_run__pb2
+from flwr.proto import fab_pb2 as flwr_dot_proto_dot_fab__pb2
+from flwr.proto import message_pb2 as flwr_dot_proto_dot_message__pb2
+
+
+DESCRIPTOR = _descriptor_pool.Default().AddSerializedFile(b'\n\x16\x66lwr/proto/fleet.proto\x12\nflwr.proto\x1a\x1a\x66lwr/proto/heartbeat.proto\x1a\x15\x66lwr/proto/node.proto\x1a\x14\x66lwr/proto/run.proto\x1a\x14\x66lwr/proto/fab.proto\x1a\x18\x66lwr/proto/message.proto\"/\n\x11\x43reateNodeRequest\x12\x1a\n\x12heartbeat_interval\x18\x01 \x01(\x01\"4\n\x12\x43reateNodeResponse\x12\x1e\n\x04node\x18\x01 \x01(\x0b\x32\x10.flwr.proto.Node\"3\n\x11\x44\x65leteNodeRequest\x12\x1e\n\x04node\x18\x01 \x01(\x0b\x32\x10.flwr.proto.Node\"\x14\n\x12\x44\x65leteNodeResponse\"J\n\x13PullMessagesRequest\x12\x1e\n\x04node\x18\x01 \x01(\x0b\x32\x10.flwr.proto.Node\x12\x13\n\x0bmessage_ids\x18\x02 \x03(\t\"\xa2\x01\n\x14PullMessagesResponse\x12(\n\treconnect\x18\x01 \x01(\x0b\x32\x15.flwr.proto.Reconnect\x12*\n\rmessages_list\x18\x02 \x03(\x0b\x32\x13.flwr.proto.Message\x12\x34\n\x14message_object_trees\x18\x03 \x03(\x0b\x32\x16.flwr.proto.ObjectTree\"\x97\x01\n\x13PushMessagesRequest\x12\x1e\n\x04node\x18\x01 \x01(\x0b\x32\x10.flwr.proto.Node\x12*\n\rmessages_list\x18\x02 \x03(\x0b\x32\x13.flwr.proto.Message\x12\x34\n\x14message_object_trees\x18\x03 \x03(\x0b\x32\x16.flwr.proto.ObjectTree\"\xcb\x02\n\x14PushMessagesResponse\x12(\n\treconnect\x18\x01 \x01(\x0b\x32\x15.flwr.proto.Reconnect\x12>\n\x07results\x18\x02 \x03(\x0b\x32-.flwr.proto.PushMessagesResponse.ResultsEntry\x12L\n\x0fobjects_to_push\x18\x03 \x03(\x0b\x32\x33.flwr.proto.PushMessagesResponse.ObjectsToPushEntry\x1a.\n\x0cResultsEntry\x12\x0b\n\x03key\x18\x01 \x01(\t\x12\r\n\x05value\x18\x02 \x01(\r:\x02\x38\x01\x1aK\n\x12ObjectsToPushEntry\x12\x0b\n\x03key\x18\x01 \x01(\t\x12$\n\x05value\x18\x02 \x01(\x0b\x32\x15.flwr.proto.ObjectIDs:\x02\x38\x01\"\x1e\n\tReconnect\x12\x11\n\treconnect\x18\x01 \x01(\x04\x32\xca\x06\n\x05\x46leet\x12M\n\nCreateNode\x12\x1d.flwr.proto.CreateNodeRequest\x1a\x1e.flwr.proto.CreateNodeResponse\"\x00\x12M\n\nDeleteNode\x12\x1d.flwr.proto.DeleteNodeRequest\x1a\x1e.flwr.proto.DeleteNodeResponse\"\x00\x12\x62\n\x11SendNodeHeartbeat\x12$.flwr.proto.SendNodeHeartbeatRequest\x1a%.flwr.proto.SendNodeHeartbeatResponse\"\x00\x12S\n\x0cPullMessages\x12\x1f.flwr.proto.PullMessagesRequest\x1a .flwr.proto.PullMessagesResponse\"\x00\x12S\n\x0cPushMessages\x12\x1f.flwr.proto.PushMessagesRequest\x1a .flwr.proto.PushMessagesResponse\"\x00\x12\x41\n\x06GetRun\x12\x19.flwr.proto.GetRunRequest\x1a\x1a.flwr.proto.GetRunResponse\"\x00\x12\x41\n\x06GetFab\x12\x19.flwr.proto.GetFabRequest\x1a\x1a.flwr.proto.GetFabResponse\"\x00\x12M\n\nPushObject\x12\x1d.flwr.proto.PushObjectRequest\x1a\x1e.flwr.proto.PushObjectResponse\"\x00\x12M\n\nPullObject\x12\x1d.flwr.proto.PullObjectRequest\x1a\x1e.flwr.proto.PullObjectResponse\"\x00\x12q\n\x16\x43onfirmMessageReceived\x12).flwr.proto.ConfirmMessageReceivedRequest\x1a*.flwr.proto.ConfirmMessageReceivedResponse\"\x00\x62\x06proto3')
+
+_globals = globals()
+_builder.BuildMessageAndEnumDescriptors(DESCRIPTOR, _globals)
+_builder.BuildTopDescriptorsAndMessages(DESCRIPTOR, 'flwr.proto.fleet_pb2', _globals)
+if _descriptor._USE_C_DESCRIPTORS == False:
+  DESCRIPTOR._options = None
+  _globals['_PUSHMESSAGESRESPONSE_RESULTSENTRY']._options = None
+  _globals['_PUSHMESSAGESRESPONSE_RESULTSENTRY']._serialized_options = b'8\001'
+  _globals['_PUSHMESSAGESRESPONSE_OBJECTSTOPUSHENTRY']._options = None
+  _globals['_PUSHMESSAGESRESPONSE_OBJECTSTOPUSHENTRY']._serialized_options = b'8\001'
+  _globals['_CREATENODEREQUEST']._serialized_start=159
+  _globals['_CREATENODEREQUEST']._serialized_end=206
+  _globals['_CREATENODERESPONSE']._serialized_start=208
+  _globals['_CREATENODERESPONSE']._serialized_end=260
+  _globals['_DELETENODEREQUEST']._serialized_start=262
+  _globals['_DELETENODEREQUEST']._serialized_end=313
+  _globals['_DELETENODERESPONSE']._serialized_start=315
+  _globals['_DELETENODERESPONSE']._serialized_end=335
+  _globals['_PULLMESSAGESREQUEST']._serialized_start=337
+  _globals['_PULLMESSAGESREQUEST']._serialized_end=411
+  _globals['_PULLMESSAGESRESPONSE']._serialized_start=414
+  _globals['_PULLMESSAGESRESPONSE']._serialized_end=576
+  _globals['_PUSHMESSAGESREQUEST']._serialized_start=579
+  _globals['_PUSHMESSAGESREQUEST']._serialized_end=730
+  _globals['_PUSHMESSAGESRESPONSE']._serialized_start=733
+  _globals['_PUSHMESSAGESRESPONSE']._serialized_end=1064
+  _globals['_PUSHMESSAGESRESPONSE_RESULTSENTRY']._serialized_start=941
+  _globals['_PUSHMESSAGESRESPONSE_RESULTSENTRY']._serialized_end=987
+  _globals['_PUSHMESSAGESRESPONSE_OBJECTSTOPUSHENTRY']._serialized_start=989
+  _globals['_PUSHMESSAGESRESPONSE_OBJECTSTOPUSHENTRY']._serialized_end=1064
+  _globals['_RECONNECT']._serialized_start=1066
+  _globals['_RECONNECT']._serialized_end=1096
+  _globals['_FLEET']._serialized_start=1099
+  _globals['_FLEET']._serialized_end=1941
+# @@protoc_insertion_point(module_scope)

{flwr_nightly-1.20.0.dev20250714 → flwr_nightly-1.20.0.dev20250716}/py/flwr/proto/fleet_pb2.pyi

@@ -78,38 +78,23 @@ global___PullMessagesRequest = PullMessagesRequest
 
 class PullMessagesResponse(google.protobuf.message.Message):
     DESCRIPTOR: google.protobuf.descriptor.Descriptor
-    class ObjectsToPullEntry(google.protobuf.message.Message):
-        DESCRIPTOR: google.protobuf.descriptor.Descriptor
-        KEY_FIELD_NUMBER: builtins.int
-        VALUE_FIELD_NUMBER: builtins.int
-        key: typing.Text
-        @property
-        def value(self) -> flwr.proto.message_pb2.ObjectIDs: ...
-        def __init__(self,
-            *,
-            key: typing.Text = ...,
-            value: typing.Optional[flwr.proto.message_pb2.ObjectIDs] = ...,
-            ) -> None: ...
-        def HasField(self, field_name: typing_extensions.Literal["value",b"value"]) -> builtins.bool: ...
-        def ClearField(self, field_name: typing_extensions.Literal["key",b"key","value",b"value"]) -> None: ...
-
     RECONNECT_FIELD_NUMBER: builtins.int
     MESSAGES_LIST_FIELD_NUMBER: builtins.int
-    OBJECTS_TO_PULL_FIELD_NUMBER: builtins.int
+    MESSAGE_OBJECT_TREES_FIELD_NUMBER: builtins.int
     @property
     def reconnect(self) -> global___Reconnect: ...
     @property
     def messages_list(self) -> google.protobuf.internal.containers.RepeatedCompositeFieldContainer[flwr.proto.message_pb2.Message]: ...
     @property
-    def objects_to_pull(self) -> google.protobuf.internal.containers.MessageMap[typing.Text, flwr.proto.message_pb2.ObjectIDs]: ...
+    def message_object_trees(self) -> google.protobuf.internal.containers.RepeatedCompositeFieldContainer[flwr.proto.message_pb2.ObjectTree]: ...
     def __init__(self,
         *,
         reconnect: typing.Optional[global___Reconnect] = ...,
         messages_list: typing.Optional[typing.Iterable[flwr.proto.message_pb2.Message]] = ...,
-        objects_to_pull: typing.Optional[typing.Mapping[typing.Text, flwr.proto.message_pb2.ObjectIDs]] = ...,
+        message_object_trees: typing.Optional[typing.Iterable[flwr.proto.message_pb2.ObjectTree]] = ...,
         ) -> None: ...
     def HasField(self, field_name: typing_extensions.Literal["reconnect",b"reconnect"]) -> builtins.bool: ...
-    def ClearField(self, field_name: typing_extensions.Literal["messages_list",b"messages_list","objects_to_pull",b"objects_to_pull","reconnect",b"reconnect"]) -> None: ...
+    def ClearField(self, field_name: typing_extensions.Literal["message_object_trees",b"message_object_trees","messages_list",b"messages_list","reconnect",b"reconnect"]) -> None: ...
 global___PullMessagesResponse = PullMessagesResponse
 
 class PushMessagesRequest(google.protobuf.message.Message):

{flwr_nightly-1.20.0.dev20250714 → flwr_nightly-1.20.0.dev20250716}/py/flwr/server/grid/grpc_grid.py

@@ -33,9 +33,9 @@ from flwr.common.inflatable import (
     get_object_tree,
     no_object_id_recompute,
 )
-from flwr.common.inflatable_grpc_utils import (
-    make_pull_object_fn_grpc,
-    make_push_object_fn_grpc,
+from flwr.common.inflatable_protobuf_utils import (
+    make_pull_object_fn_protobuf,
+    make_push_object_fn_protobuf,
 )
 from flwr.common.inflatable_utils import (
     inflate_object_from_contents,
@@ -240,8 +240,8 @@ class GrpcGrid(Grid):
             # Push only object that are not in the store
             push_objects(
                 all_objects,
-                push_object_fn=make_push_object_fn_grpc(
-                    push_object_grpc=self._stub.PushObject,
+                push_object_fn=make_push_object_fn_protobuf(
+                    push_object_protobuf=self._stub.PushObject,
                     node=self.node,
                     run_id=run_id,
                 ),
@@ -308,8 +308,8 @@ class GrpcGrid(Grid):
                 msg_id = msg_proto.metadata.message_id
                 all_object_contents = pull_objects(
                     list(res.objects_to_pull[msg_id].object_ids) + [msg_id],
-                    pull_object_fn=make_pull_object_fn_grpc(
-                        pull_object_grpc=self._stub.PullObject,
+                    pull_object_fn=make_pull_object_fn_protobuf(
+                        pull_object_protobuf=self._stub.PullObject,
                         node=self.node,
                         run_id=run_id,
                     ),

{flwr_nightly-1.20.0.dev20250714 → flwr_nightly-1.20.0.dev20250716}/py/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_nightly-1.20.0.dev20250714 → flwr_nightly-1.20.0.dev20250716}/py/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
@@ -211,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
@@ -247,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:
@@ -513,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_nightly-1.20.0.dev20250714 → flwr_nightly-1.20.0.dev20250716}/py/flwr/server/superlink/utils.py

@@ -20,7 +20,6 @@ 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
@@ -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_nightly-1.20.0.dev20250716/py/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_nightly-1.20.0.dev20250716/py/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_nightly-1.20.0.dev20250714 → flwr_nightly-1.20.0.dev20250716}/py/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]