flwr-nightly 1.19.0.dev20250528__py3-none-any.whl → 1.19.0.dev20250530__py3-none-any.whl

flwr/proto/serverappio_pb2.py

@@ -20,15 +20,15 @@ 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
 
 
-DESCRIPTOR = _descriptor_pool.Default().AddSerializedFile(b'\n\x1c\x66lwr/proto/serverappio.proto\x12\nflwr.proto\x1a\x1a\x66lwr/proto/heartbeat.proto\x1a\x14\x66lwr/proto/log.proto\x1a\x15\x66lwr/proto/node.proto\x1a\x18\x66lwr/proto/message.proto\x1a\x14\x66lwr/proto/run.proto\x1a\x14\x66lwr/proto/fab.proto\"!\n\x0fGetNodesRequest\x12\x0e\n\x06run_id\x18\x01 \x01(\x04\"3\n\x10GetNodesResponse\x12\x1f\n\x05nodes\x18\x01 \x03(\x0b\x32\x10.flwr.proto.Node\"\x87\x02\n\x16PushInsMessagesRequest\x12*\n\rmessages_list\x18\x01 \x03(\x0b\x32\x13.flwr.proto.Message\x12\x0e\n\x06run_id\x18\x02 \x01(\x04\x12]\n\x17msg_to_children_mapping\x18\x03 \x03(\x0b\x32<.flwr.proto.PushInsMessagesRequest.MsgToChildrenMappingEntry\x1aR\n\x19MsgToChildrenMappingEntry\x12\x0b\n\x03key\x18\x01 \x01(\t\x12$\n\x05value\x18\x02 \x01(\x0b\x32\x15.flwr.proto.ObjectIDs:\x02\x38\x01\"\xcc\x01\n\x17PushInsMessagesResponse\x12\x13\n\x0bmessage_ids\x18\x01 \x03(\t\x12O\n\x0fobjects_to_push\x18\x02 \x03(\x0b\x32\x36.flwr.proto.PushInsMessagesResponse.ObjectsToPushEntry\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\"=\n\x16PullResMessagesRequest\x12\x13\n\x0bmessage_ids\x18\x01 \x03(\t\x12\x0e\n\x06run_id\x18\x02 \x01(\x04\"\xe3\x01\n\x17PullResMessagesResponse\x12*\n\rmessages_list\x18\x01 \x03(\x0b\x32\x13.flwr.proto.Message\x12O\n\x0fobjects_to_pull\x18\x02 \x03(\x0b\x32\x36.flwr.proto.PullResMessagesResponse.ObjectsToPullEntry\x1aK\n\x12ObjectsToPullEntry\x12\x0b\n\x03key\x18\x01 \x01(\t\x12$\n\x05value\x18\x02 \x01(\x0b\x32\x15.flwr.proto.ObjectIDs:\x02\x38\x01\"\x1c\n\x1aPullServerAppInputsRequest\"\x7f\n\x1bPullServerAppInputsResponse\x12$\n\x07\x63ontext\x18\x01 \x01(\x0b\x32\x13.flwr.proto.Context\x12\x1c\n\x03run\x18\x02 \x01(\x0b\x32\x0f.flwr.proto.Run\x12\x1c\n\x03\x66\x61\x62\x18\x03 \x01(\x0b\x32\x0f.flwr.proto.Fab\"S\n\x1bPushServerAppOutputsRequest\x12\x0e\n\x06run_id\x18\x01 \x01(\x04\x12$\n\x07\x63ontext\x18\x02 \x01(\x0b\x32\x13.flwr.proto.Context\"\x1e\n\x1cPushServerAppOutputsResponse2\xe4\x08\n\x0bServerAppIo\x12G\n\x08GetNodes\x12\x1b.flwr.proto.GetNodesRequest\x1a\x1c.flwr.proto.GetNodesResponse\"\x00\x12Y\n\x0cPushMessages\x12\".flwr.proto.PushInsMessagesRequest\x1a#.flwr.proto.PushInsMessagesResponse\"\x00\x12Y\n\x0cPullMessages\x12\".flwr.proto.PullResMessagesRequest\x1a#.flwr.proto.PullResMessagesResponse\"\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\x12h\n\x13PullServerAppInputs\x12&.flwr.proto.PullServerAppInputsRequest\x1a\'.flwr.proto.PullServerAppInputsResponse\"\x00\x12k\n\x14PushServerAppOutputs\x12\'.flwr.proto.PushServerAppOutputsRequest\x1a(.flwr.proto.PushServerAppOutputsResponse\"\x00\x12\\\n\x0fUpdateRunStatus\x12\".flwr.proto.UpdateRunStatusRequest\x1a#.flwr.proto.UpdateRunStatusResponse\"\x00\x12S\n\x0cGetRunStatus\x12\x1f.flwr.proto.GetRunStatusRequest\x1a .flwr.proto.GetRunStatusResponse\"\x00\x12G\n\x08PushLogs\x12\x1b.flwr.proto.PushLogsRequest\x1a\x1c.flwr.proto.PushLogsResponse\"\x00\x12_\n\x10SendAppHeartbeat\x12#.flwr.proto.SendAppHeartbeatRequest\x1a$.flwr.proto.SendAppHeartbeatResponse\"\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\x62\x06proto3')
+DESCRIPTOR = _descriptor_pool.Default().AddSerializedFile(b'\n\x1c\x66lwr/proto/serverappio.proto\x12\nflwr.proto\x1a\x1a\x66lwr/proto/heartbeat.proto\x1a\x14\x66lwr/proto/log.proto\x1a\x15\x66lwr/proto/node.proto\x1a\x18\x66lwr/proto/message.proto\x1a\x14\x66lwr/proto/run.proto\x1a\x14\x66lwr/proto/fab.proto\"!\n\x0fGetNodesRequest\x12\x0e\n\x06run_id\x18\x01 \x01(\x04\"3\n\x10GetNodesResponse\x12\x1f\n\x05nodes\x18\x01 \x03(\x0b\x32\x10.flwr.proto.Node\"\x8d\x02\n\x16PushInsMessagesRequest\x12*\n\rmessages_list\x18\x01 \x03(\x0b\x32\x13.flwr.proto.Message\x12\x0e\n\x06run_id\x18\x02 \x01(\x04\x12\x61\n\x19msg_to_descendant_mapping\x18\x03 \x03(\x0b\x32>.flwr.proto.PushInsMessagesRequest.MsgToDescendantMappingEntry\x1aT\n\x1bMsgToDescendantMappingEntry\x12\x0b\n\x03key\x18\x01 \x01(\t\x12$\n\x05value\x18\x02 \x01(\x0b\x32\x15.flwr.proto.ObjectIDs:\x02\x38\x01\"\xcc\x01\n\x17PushInsMessagesResponse\x12\x13\n\x0bmessage_ids\x18\x01 \x03(\t\x12O\n\x0fobjects_to_push\x18\x02 \x03(\x0b\x32\x36.flwr.proto.PushInsMessagesResponse.ObjectsToPushEntry\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\"=\n\x16PullResMessagesRequest\x12\x13\n\x0bmessage_ids\x18\x01 \x03(\t\x12\x0e\n\x06run_id\x18\x02 \x01(\x04\"\xe3\x01\n\x17PullResMessagesResponse\x12*\n\rmessages_list\x18\x01 \x03(\x0b\x32\x13.flwr.proto.Message\x12O\n\x0fobjects_to_pull\x18\x02 \x03(\x0b\x32\x36.flwr.proto.PullResMessagesResponse.ObjectsToPullEntry\x1aK\n\x12ObjectsToPullEntry\x12\x0b\n\x03key\x18\x01 \x01(\t\x12$\n\x05value\x18\x02 \x01(\x0b\x32\x15.flwr.proto.ObjectIDs:\x02\x38\x01\"\x1c\n\x1aPullServerAppInputsRequest\"\x7f\n\x1bPullServerAppInputsResponse\x12$\n\x07\x63ontext\x18\x01 \x01(\x0b\x32\x13.flwr.proto.Context\x12\x1c\n\x03run\x18\x02 \x01(\x0b\x32\x0f.flwr.proto.Run\x12\x1c\n\x03\x66\x61\x62\x18\x03 \x01(\x0b\x32\x0f.flwr.proto.Fab\"S\n\x1bPushServerAppOutputsRequest\x12\x0e\n\x06run_id\x18\x01 \x01(\x04\x12$\n\x07\x63ontext\x18\x02 \x01(\x0b\x32\x13.flwr.proto.Context\"\x1e\n\x1cPushServerAppOutputsResponse2\xe4\x08\n\x0bServerAppIo\x12G\n\x08GetNodes\x12\x1b.flwr.proto.GetNodesRequest\x1a\x1c.flwr.proto.GetNodesResponse\"\x00\x12Y\n\x0cPushMessages\x12\".flwr.proto.PushInsMessagesRequest\x1a#.flwr.proto.PushInsMessagesResponse\"\x00\x12Y\n\x0cPullMessages\x12\".flwr.proto.PullResMessagesRequest\x1a#.flwr.proto.PullResMessagesResponse\"\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\x12h\n\x13PullServerAppInputs\x12&.flwr.proto.PullServerAppInputsRequest\x1a\'.flwr.proto.PullServerAppInputsResponse\"\x00\x12k\n\x14PushServerAppOutputs\x12\'.flwr.proto.PushServerAppOutputsRequest\x1a(.flwr.proto.PushServerAppOutputsResponse\"\x00\x12\\\n\x0fUpdateRunStatus\x12\".flwr.proto.UpdateRunStatusRequest\x1a#.flwr.proto.UpdateRunStatusResponse\"\x00\x12S\n\x0cGetRunStatus\x12\x1f.flwr.proto.GetRunStatusRequest\x1a .flwr.proto.GetRunStatusResponse\"\x00\x12G\n\x08PushLogs\x12\x1b.flwr.proto.PushLogsRequest\x1a\x1c.flwr.proto.PushLogsResponse\"\x00\x12_\n\x10SendAppHeartbeat\x12#.flwr.proto.SendAppHeartbeatRequest\x1a$.flwr.proto.SendAppHeartbeatResponse\"\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\x62\x06proto3')
 
 _globals = globals()
 _builder.BuildMessageAndEnumDescriptors(DESCRIPTOR, _globals)
 _builder.BuildTopDescriptorsAndMessages(DESCRIPTOR, 'flwr.proto.serverappio_pb2', _globals)
 if _descriptor._USE_C_DESCRIPTORS == False:
   DESCRIPTOR._options = None
-  _globals['_PUSHINSMESSAGESREQUEST_MSGTOCHILDRENMAPPINGENTRY']._options = None
-  _globals['_PUSHINSMESSAGESREQUEST_MSGTOCHILDRENMAPPINGENTRY']._serialized_options = b'8\001'
+  _globals['_PUSHINSMESSAGESREQUEST_MSGTODESCENDANTMAPPINGENTRY']._options = None
+  _globals['_PUSHINSMESSAGESREQUEST_MSGTODESCENDANTMAPPINGENTRY']._serialized_options = b'8\001'
   _globals['_PUSHINSMESSAGESRESPONSE_OBJECTSTOPUSHENTRY']._options = None
   _globals['_PUSHINSMESSAGESRESPONSE_OBJECTSTOPUSHENTRY']._serialized_options = b'8\001'
   _globals['_PULLRESMESSAGESRESPONSE_OBJECTSTOPULLENTRY']._options = None
@@ -38,27 +38,27 @@ if _descriptor._USE_C_DESCRIPTORS == False:
   _globals['_GETNODESRESPONSE']._serialized_start=222
   _globals['_GETNODESRESPONSE']._serialized_end=273
   _globals['_PUSHINSMESSAGESREQUEST']._serialized_start=276
-  _globals['_PUSHINSMESSAGESREQUEST']._serialized_end=539
-  _globals['_PUSHINSMESSAGESREQUEST_MSGTOCHILDRENMAPPINGENTRY']._serialized_start=457
-  _globals['_PUSHINSMESSAGESREQUEST_MSGTOCHILDRENMAPPINGENTRY']._serialized_end=539
-  _globals['_PUSHINSMESSAGESRESPONSE']._serialized_start=542
-  _globals['_PUSHINSMESSAGESRESPONSE']._serialized_end=746
-  _globals['_PUSHINSMESSAGESRESPONSE_OBJECTSTOPUSHENTRY']._serialized_start=671
-  _globals['_PUSHINSMESSAGESRESPONSE_OBJECTSTOPUSHENTRY']._serialized_end=746
-  _globals['_PULLRESMESSAGESREQUEST']._serialized_start=748
-  _globals['_PULLRESMESSAGESREQUEST']._serialized_end=809
-  _globals['_PULLRESMESSAGESRESPONSE']._serialized_start=812
-  _globals['_PULLRESMESSAGESRESPONSE']._serialized_end=1039
-  _globals['_PULLRESMESSAGESRESPONSE_OBJECTSTOPULLENTRY']._serialized_start=964
-  _globals['_PULLRESMESSAGESRESPONSE_OBJECTSTOPULLENTRY']._serialized_end=1039
-  _globals['_PULLSERVERAPPINPUTSREQUEST']._serialized_start=1041
-  _globals['_PULLSERVERAPPINPUTSREQUEST']._serialized_end=1069
-  _globals['_PULLSERVERAPPINPUTSRESPONSE']._serialized_start=1071
-  _globals['_PULLSERVERAPPINPUTSRESPONSE']._serialized_end=1198
-  _globals['_PUSHSERVERAPPOUTPUTSREQUEST']._serialized_start=1200
-  _globals['_PUSHSERVERAPPOUTPUTSREQUEST']._serialized_end=1283
-  _globals['_PUSHSERVERAPPOUTPUTSRESPONSE']._serialized_start=1285
-  _globals['_PUSHSERVERAPPOUTPUTSRESPONSE']._serialized_end=1315
-  _globals['_SERVERAPPIO']._serialized_start=1318
-  _globals['_SERVERAPPIO']._serialized_end=2442
+  _globals['_PUSHINSMESSAGESREQUEST']._serialized_end=545
+  _globals['_PUSHINSMESSAGESREQUEST_MSGTODESCENDANTMAPPINGENTRY']._serialized_start=461
+  _globals['_PUSHINSMESSAGESREQUEST_MSGTODESCENDANTMAPPINGENTRY']._serialized_end=545
+  _globals['_PUSHINSMESSAGESRESPONSE']._serialized_start=548
+  _globals['_PUSHINSMESSAGESRESPONSE']._serialized_end=752
+  _globals['_PUSHINSMESSAGESRESPONSE_OBJECTSTOPUSHENTRY']._serialized_start=677
+  _globals['_PUSHINSMESSAGESRESPONSE_OBJECTSTOPUSHENTRY']._serialized_end=752
+  _globals['_PULLRESMESSAGESREQUEST']._serialized_start=754
+  _globals['_PULLRESMESSAGESREQUEST']._serialized_end=815
+  _globals['_PULLRESMESSAGESRESPONSE']._serialized_start=818
+  _globals['_PULLRESMESSAGESRESPONSE']._serialized_end=1045
+  _globals['_PULLRESMESSAGESRESPONSE_OBJECTSTOPULLENTRY']._serialized_start=970
+  _globals['_PULLRESMESSAGESRESPONSE_OBJECTSTOPULLENTRY']._serialized_end=1045
+  _globals['_PULLSERVERAPPINPUTSREQUEST']._serialized_start=1047
+  _globals['_PULLSERVERAPPINPUTSREQUEST']._serialized_end=1075
+  _globals['_PULLSERVERAPPINPUTSRESPONSE']._serialized_start=1077
+  _globals['_PULLSERVERAPPINPUTSRESPONSE']._serialized_end=1204
+  _globals['_PUSHSERVERAPPOUTPUTSREQUEST']._serialized_start=1206
+  _globals['_PUSHSERVERAPPOUTPUTSREQUEST']._serialized_end=1289
+  _globals['_PUSHSERVERAPPOUTPUTSRESPONSE']._serialized_start=1291
+  _globals['_PUSHSERVERAPPOUTPUTSRESPONSE']._serialized_end=1321
+  _globals['_SERVERAPPIO']._serialized_start=1324
+  _globals['_SERVERAPPIO']._serialized_end=2448
 # @@protoc_insertion_point(module_scope)

flwr/proto/serverappio_pb2.pyi

@@ -42,7 +42,7 @@ global___GetNodesResponse = GetNodesResponse
 class PushInsMessagesRequest(google.protobuf.message.Message):
     """PushMessages messages"""
     DESCRIPTOR: google.protobuf.descriptor.Descriptor
-    class MsgToChildrenMappingEntry(google.protobuf.message.Message):
+    class MsgToDescendantMappingEntry(google.protobuf.message.Message):
         DESCRIPTOR: google.protobuf.descriptor.Descriptor
         KEY_FIELD_NUMBER: builtins.int
         VALUE_FIELD_NUMBER: builtins.int
@@ -59,19 +59,19 @@ class PushInsMessagesRequest(google.protobuf.message.Message):
 
     MESSAGES_LIST_FIELD_NUMBER: builtins.int
     RUN_ID_FIELD_NUMBER: builtins.int
-    MSG_TO_CHILDREN_MAPPING_FIELD_NUMBER: builtins.int
+    MSG_TO_DESCENDANT_MAPPING_FIELD_NUMBER: builtins.int
     @property
     def messages_list(self) -> google.protobuf.internal.containers.RepeatedCompositeFieldContainer[flwr.proto.message_pb2.Message]: ...
     run_id: builtins.int
     @property
-    def msg_to_children_mapping(self) -> google.protobuf.internal.containers.MessageMap[typing.Text, flwr.proto.message_pb2.ObjectIDs]: ...
+    def msg_to_descendant_mapping(self) -> google.protobuf.internal.containers.MessageMap[typing.Text, flwr.proto.message_pb2.ObjectIDs]: ...
     def __init__(self,
         *,
         messages_list: typing.Optional[typing.Iterable[flwr.proto.message_pb2.Message]] = ...,
         run_id: builtins.int = ...,
-        msg_to_children_mapping: typing.Optional[typing.Mapping[typing.Text, flwr.proto.message_pb2.ObjectIDs]] = ...,
+        msg_to_descendant_mapping: typing.Optional[typing.Mapping[typing.Text, flwr.proto.message_pb2.ObjectIDs]] = ...,
         ) -> None: ...
-    def ClearField(self, field_name: typing_extensions.Literal["messages_list",b"messages_list","msg_to_children_mapping",b"msg_to_children_mapping","run_id",b"run_id"]) -> None: ...
+    def ClearField(self, field_name: typing_extensions.Literal["messages_list",b"messages_list","msg_to_descendant_mapping",b"msg_to_descendant_mapping","run_id",b"run_id"]) -> None: ...
 global___PushInsMessagesRequest = PushInsMessagesRequest
 
 class PushInsMessagesResponse(google.protobuf.message.Message):

flwr/server/app.py

@@ -27,7 +27,7 @@ from collections.abc import Sequence
 from logging import DEBUG, INFO, WARN
 from pathlib import Path
 from time import sleep
-from typing import Any, Optional
+from typing import Any, Callable, Optional, TypeVar
 
 import grpc
 import yaml
@@ -85,6 +85,7 @@ from .superlink.simulation.simulationio_grpc import run_simulationio_api_grpc
 
 DATABASE = ":flwr-in-memory-state:"
 BASE_DIR = get_flwr_dir() / "superlink" / "ffs"
+P = TypeVar("P", ExecAuthPlugin, ExecAuthzPlugin)
 
 
 try:
@@ -151,21 +152,16 @@ def run_superlink() -> None:
     verify_tls_cert = not getattr(args, "disable_oidc_tls_cert_verification", None)
 
     auth_plugin: Optional[ExecAuthPlugin] = None
-    authz_plugin: Optional[ExecAuthzPlugin] = None  # pylint: disable=unused-variable
+    authz_plugin: Optional[ExecAuthzPlugin] = None
     event_log_plugin: Optional[EventLogWriterPlugin] = None
     # Load the auth plugin if the args.user_auth_config is provided
     if cfg_path := getattr(args, "user_auth_config", None):
-        auth_plugin = _try_obtain_exec_auth_plugin(Path(cfg_path), verify_tls_cert)
+        auth_plugin, authz_plugin = _try_obtain_exec_auth_plugins(
+            Path(cfg_path), verify_tls_cert
+        )
         # Enable event logging if the args.enable_event_log is True
         if args.enable_event_log:
             event_log_plugin = _try_obtain_exec_event_log_writer_plugin()
-        # Enable authorization if the args.enable_authorization is True
-        if args.enable_authorization:
-            # pylint: disable=unused-variable
-            authz_plugin = _try_obtain_exec_authz_plugin(  # noqa: F841
-                Path(cfg_path), verify_tls_cert
-            )
-            # pylint: enable=unused-variable
 
     # Initialize StateFactory
     state_factory = LinkStateFactory(args.database)
@@ -188,6 +184,7 @@ def run_superlink() -> None:
             [args.executor_config] if args.executor_config else args.executor_config
         ),
         auth_plugin=auth_plugin,
+        authz_plugin=authz_plugin,
         event_log_plugin=event_log_plugin,
     )
     grpc_servers = [exec_server]
@@ -258,6 +255,7 @@ def run_superlink() -> None:
                     args.ssl_certfile,
                     state_factory,
                     ffs_factory,
+                    objectstore_factory,
                     num_workers,
                 ),
                 daemon=True,
@@ -483,62 +481,50 @@ def _try_load_public_keys_node_authentication(
     return node_public_keys
 
 
-def _try_obtain_exec_auth_plugin(
+def _try_obtain_exec_auth_plugins(
     config_path: Path, verify_tls_cert: bool
-) -> Optional[ExecAuthPlugin]:
+) -> tuple[ExecAuthPlugin, ExecAuthzPlugin]:
+    """Obtain Exec API authentication and authorization plugins."""
     # Load YAML file
     with config_path.open("r", encoding="utf-8") as file:
         config: dict[str, Any] = yaml.safe_load(file)
 
-    # Load authentication configuration
-    auth_config: dict[str, Any] = config.get("authentication", {})
-    auth_type: str = auth_config.get(AUTH_TYPE_YAML_KEY, "")
-
-    # Load authentication plugin
-    try:
-        all_plugins: dict[str, type[ExecAuthPlugin]] = get_exec_auth_plugins()
-        auth_plugin_class = all_plugins[auth_type]
-        return auth_plugin_class(
-            user_auth_config_path=config_path, verify_tls_cert=verify_tls_cert
-        )
-    except KeyError:
-        if auth_type != "":
-            sys.exit(
-                f'Authentication type "{auth_type}" is not supported. '
-                "Please provide a valid authentication type in the configuration."
+    def _load_plugin(
+        section: str, yaml_key: str, loader: Callable[[], dict[str, type[P]]]
+    ) -> P:
+        section_cfg = config.get(section, {})
+        auth_plugin_name = section_cfg.get(yaml_key, "")
+        try:
+            plugins: dict[str, type[P]] = loader()
+            plugin_cls: type[P] = plugins[auth_plugin_name]
+            return plugin_cls(
+                user_auth_config_path=config_path, verify_tls_cert=verify_tls_cert
             )
-        sys.exit("No authentication type is provided in the configuration.")
-    except NotImplementedError:
-        sys.exit("No authentication plugins are currently supported.")
-
-
-def _try_obtain_exec_authz_plugin(
-    config_path: Path, verify_tls_cert: bool
-) -> Optional[ExecAuthzPlugin]:
-    # Load YAML file
-    with config_path.open("r", encoding="utf-8") as file:
-        config: dict[str, Any] = yaml.safe_load(file)
+        except KeyError:
+            if auth_plugin_name:
+                sys.exit(
+                    f"{yaml_key}: {auth_plugin_name} is not supported. "
+                    f"Please provide a valid {section} type in the configuration."
+                )
+            sys.exit(f"No {section} type is provided in the configuration.")
+        except NotImplementedError:
+            sys.exit(f"No {section} plugins are currently supported.")
 
-    # Load authentication configuration
-    authz_config: dict[str, Any] = config.get("authorization", {})
-    authz_type: str = authz_config.get(AUTHZ_TYPE_YAML_KEY, "")
+    # Load authentication plugin
+    auth_plugin = _load_plugin(
+        section="authentication",
+        yaml_key=AUTH_TYPE_YAML_KEY,
+        loader=get_exec_auth_plugins,
+    )
 
     # Load authorization plugin
-    try:
-        all_plugins: dict[str, type[ExecAuthzPlugin]] = get_exec_authz_plugins()
-        authz_plugin_class = all_plugins[authz_type]
-        return authz_plugin_class(
-            user_authz_config_path=config_path, verify_tls_cert=verify_tls_cert
-        )
-    except KeyError:
-        if authz_type != "":
-            sys.exit(
-                f'Authentication type "{authz_type}" is not supported. '
-                "Please provide a valid authorization type in the configuration."
-            )
-        sys.exit("No authorization type is provided in the configuration.")
-    except NotImplementedError:
-        sys.exit("No authorization plugins are currently supported.")
+    authz_plugin = _load_plugin(
+        section="authorization",
+        yaml_key=AUTHZ_TYPE_YAML_KEY,
+        loader=get_exec_authz_plugins,
+    )
+
+    return auth_plugin, authz_plugin
 
 
 def _try_obtain_exec_event_log_writer_plugin() -> Optional[EventLogWriterPlugin]:
@@ -636,6 +622,7 @@ def _run_fleet_api_rest(
     ssl_certfile: Optional[str],
     state_factory: LinkStateFactory,
     ffs_factory: FfsFactory,
+    objectstore_factory: ObjectStoreFactory,
     num_workers: int,
 ) -> None:
     """Run ServerAppIo API (REST-based)."""
@@ -651,6 +638,7 @@ def _run_fleet_api_rest(
     # See: https://www.starlette.io/applications/#accessing-the-app-instance
     fast_api_app.state.STATE_FACTORY = state_factory
     fast_api_app.state.FFS_FACTORY = ffs_factory
+    fast_api_app.state.OBJECTSTORE_FACTORY = objectstore_factory
 
     uvicorn.run(
         app="flwr.server.superlink.fleet.rest_rere.rest_api:app",

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

@@ -114,6 +114,7 @@ class FleetServicer(fleet_pb2_grpc.FleetServicer):
         return message_handler.pull_messages(
             request=request,
             state=self.state_factory.state(),
+            store=self.objectstore_factory.store(),
         )
 
     def PushMessages(
@@ -133,6 +134,7 @@ class FleetServicer(fleet_pb2_grpc.FleetServicer):
             res = message_handler.push_messages(
                 request=request,
                 state=self.state_factory.state(),
+                store=self.objectstore_factory.store(),
             )
         except InvalidRunStatusException as e:
             abort_grpc_context(e.message, context)

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

@@ -14,10 +14,10 @@
 # ==============================================================================
 """Fleet API message handlers."""
 
-
+from logging import ERROR
 from typing import Optional
 
-from flwr.common import Message
+from flwr.common import Message, log
 from flwr.common.constant import Status
 from flwr.common.serde import (
     fab_to_proto,
@@ -42,6 +42,7 @@ from flwr.proto.heartbeat_pb2 import (  # pylint: disable=E0611
     SendNodeHeartbeatRequest,
     SendNodeHeartbeatResponse,
 )
+from flwr.proto.message_pb2 import ObjectIDs  # pylint: disable=E0611
 from flwr.proto.node_pb2 import Node  # pylint: disable=E0611
 from flwr.proto.run_pb2 import (  # pylint: disable=E0611
     GetRunRequest,
@@ -51,6 +52,9 @@ from flwr.proto.run_pb2 import (  # pylint: disable=E0611
 from flwr.server.superlink.ffs.ffs import Ffs
 from flwr.server.superlink.linkstate import LinkState
 from flwr.server.superlink.utils import check_abort
+from flwr.supercore.object_store import NoObjectInStoreError, ObjectStore
+
+from ...utils import store_mapping_and_register_objects
 
 
 def create_node(
@@ -86,7 +90,9 @@ def send_node_heartbeat(
 
 
 def pull_messages(
-    request: PullMessagesRequest, state: LinkState
+    request: PullMessagesRequest,
+    state: LinkState,
+    store: ObjectStore,
 ) -> PullMessagesResponse:
     """Pull Messages handler."""
     # Get node_id if client node is not anonymous
@@ -98,14 +104,31 @@ def pull_messages(
 
     # Convert to Messages
     msg_proto = []
+    objects_to_pull: dict[str, ObjectIDs] = {}
     for msg in message_list:
-        msg_proto.append(message_to_proto(msg))
-
-    return PullMessagesResponse(messages_list=msg_proto)
+        try:
+            msg_proto.append(message_to_proto(msg))
+
+            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]
+            )
+        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
+    )
 
 
 def push_messages(
-    request: PushMessagesRequest, state: LinkState
+    request: PushMessagesRequest,
+    state: LinkState,
+    store: ObjectStore,
 ) -> PushMessagesResponse:
     """Push Messages handler."""
     # Convert Message from proto
@@ -123,10 +146,14 @@ def push_messages(
     # Store Message in State
     message_id: Optional[str] = state.store_message_res(message=msg)
 
+    # Store Message object to descendants mapping and preregister objects
+    objects_to_push = store_mapping_and_register_objects(store, request=request)
+
     # Build response
     response = PushMessagesResponse(
         reconnect=Reconnect(reconnect=5),
         results={str(message_id): 0},
+        objects_to_push=objects_to_push,
     )
     return response
 

flwr/server/superlink/fleet/rest_rere/rest_api.py

@@ -43,6 +43,7 @@ from flwr.server.superlink.ffs.ffs import Ffs
 from flwr.server.superlink.ffs.ffs_factory import FfsFactory
 from flwr.server.superlink.fleet.message_handler import message_handler
 from flwr.server.superlink.linkstate import LinkState, LinkStateFactory
+from flwr.supercore.object_store import ObjectStore, ObjectStoreFactory
 
 try:
     from starlette.applications import Starlette
@@ -113,9 +114,10 @@ async def pull_message(request: PullMessagesRequest) -> PullMessagesResponse:
     """Pull PullMessages."""
     # Get state from app
     state: LinkState = cast(LinkStateFactory, app.state.STATE_FACTORY).state()
+    store: ObjectStore = cast(ObjectStoreFactory, app.state.STATE_FACTORY).store()
 
     # Handle message
-    return message_handler.pull_messages(request=request, state=state)
+    return message_handler.pull_messages(request=request, state=state, store=store)
 
 
 @rest_request_response(PushMessagesRequest)
@@ -123,9 +125,10 @@ async def push_message(request: PushMessagesRequest) -> PushMessagesResponse:
     """Pull PushMessages."""
     # Get state from app
     state: LinkState = cast(LinkStateFactory, app.state.STATE_FACTORY).state()
+    store: ObjectStore = cast(ObjectStoreFactory, app.state.STATE_FACTORY).store()
 
     # Handle message
-    return message_handler.push_messages(request=request, state=state)
+    return message_handler.push_messages(request=request, state=state, store=store)
 
 
 @rest_request_response(SendNodeHeartbeatRequest)

flwr/server/superlink/linkstate/utils.py

@@ -17,7 +17,6 @@
 
 from os import urandom
 from typing import Optional
-from uuid import uuid4
 
 from flwr.common import ConfigRecord, Context, Error, Message, Metadata, now, serde
 from flwr.common.constant import (
@@ -246,7 +245,7 @@ def create_message_error_unavailable_res_message(
     ttl = max(ins_metadata.ttl - (current_time - ins_metadata.created_at), 0)
     metadata = Metadata(
         run_id=ins_metadata.run_id,
-        message_id=str(uuid4()),
+        message_id="",
         src_node_id=SUPERLINK_NODE_ID,
         dst_node_id=SUPERLINK_NODE_ID,
         reply_to_message_id=ins_metadata.message_id,
@@ -256,7 +255,7 @@ def create_message_error_unavailable_res_message(
         ttl=ttl,
     )
 
-    return make_message(
+    msg = make_message(
         metadata=metadata,
         error=Error(
             code=(
@@ -271,6 +270,8 @@ def create_message_error_unavailable_res_message(
             ),
         ),
     )
+    msg.metadata.__dict__["_message_id"] = msg.object_id
+    return msg
 
 
 def create_message_error_unavailable_ins_message(reply_to_message_id: str) -> Message:
@@ -278,7 +279,7 @@ def create_message_error_unavailable_ins_message(reply_to_message_id: str) -> Me
     that it isn't found."""
     metadata = Metadata(
         run_id=0,  # Unknown
-        message_id=str(uuid4()),
+        message_id="",
         src_node_id=SUPERLINK_NODE_ID,
         dst_node_id=SUPERLINK_NODE_ID,
         reply_to_message_id=reply_to_message_id,
@@ -288,13 +289,15 @@ def create_message_error_unavailable_ins_message(reply_to_message_id: str) -> Me
         ttl=0,
     )
 
-    return make_message(
+    msg = make_message(
         metadata=metadata,
         error=Error(
             code=ErrorCode.MESSAGE_UNAVAILABLE,
             reason=MESSAGE_UNAVAILABLE_ERROR_REASON,
         ),
     )
+    msg.metadata.__dict__["_message_id"] = msg.object_id
+    return msg
 
 
 def message_ttl_has_expired(message_metadata: Metadata, current_time: float) -> bool:

flwr/server/superlink/serverappio/serverappio_servicer.py

@@ -16,14 +16,14 @@
 
 
 import threading
-from logging import DEBUG, INFO
+from logging import DEBUG, ERROR, INFO
 from typing import Optional
 
 import grpc
 
 from flwr.common import Message
 from flwr.common.constant import SUPERLINK_NODE_ID, Status
-from flwr.common.inflatable import check_body_len_consistency
+from flwr.common.inflatable import check_body_len_consistency, get_desdendant_object_ids
 from flwr.common.logger import log
 from flwr.common.serde import (
     context_from_proto,
@@ -47,6 +47,7 @@ from flwr.proto.log_pb2 import (  # pylint: disable=E0611
     PushLogsResponse,
 )
 from flwr.proto.message_pb2 import (  # pylint: disable=E0611
+    ObjectIDs,
     PullObjectRequest,
     PullObjectResponse,
     PushObjectRequest,
@@ -78,7 +79,9 @@ from flwr.server.superlink.ffs.ffs_factory import FfsFactory
 from flwr.server.superlink.linkstate import LinkState, LinkStateFactory
 from flwr.server.superlink.utils import abort_if
 from flwr.server.utils.validator import validate_message
-from flwr.supercore.object_store import ObjectStoreFactory
+from flwr.supercore.object_store import NoObjectInStoreError, ObjectStoreFactory
+
+from ..utils import store_mapping_and_register_objects
 
 
 class ServerAppIoServicer(serverappio_pb2_grpc.ServerAppIoServicer):
@@ -158,10 +161,17 @@ 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)
+
         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(
@@ -173,6 +183,9 @@ class ServerAppIoServicer(serverappio_pb2_grpc.ServerAppIoServicer):
         # Init state
         state: LinkState = self.state_factory.state()
 
+        # Init store
+        store = self.objectstore_factory.store()
+
         # Abort if the run is not running
         abort_if(
             request.run_id,
@@ -186,6 +199,18 @@ class ServerAppIoServicer(serverappio_pb2_grpc.ServerAppIoServicer):
             message_ids=set(request.message_ids)
         )
 
+        # 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_desdendant_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])
+
         # Delete the instruction Messages and their replies if found
         message_ins_ids_to_delete = {
             msg_res.metadata.reply_to_message_id for msg_res in messages_res
@@ -195,6 +220,7 @@ class ServerAppIoServicer(serverappio_pb2_grpc.ServerAppIoServicer):
 
         # Convert Messages to proto
         messages_list = []
+        objects_to_pull: dict[str, ObjectIDs] = {}
         while messages_res:
             msg = messages_res.pop(0)
 
@@ -207,7 +233,21 @@ class ServerAppIoServicer(serverappio_pb2_grpc.ServerAppIoServicer):
                 )
             messages_list.append(message_to_proto(msg))
 
-        return PullResMessagesResponse(messages_list=messages_list)
+            try:
+                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]
+                )
+            except NoObjectInStoreError as e:
+                log(ERROR, e.message)
+                # Delete message ins from state
+                state.delete_messages(message_ins_ids={msg_object_id})
+
+        return PullResMessagesResponse(
+            messages_list=messages_list, objects_to_pull=objects_to_pull
+        )
 
     def GetRun(
         self, request: GetRunRequest, context: grpc.ServicerContext

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/supercore/object_store/__init__.py

@@ -14,10 +14,11 @@
 # ==============================================================================
 """Flower ObjectStore."""
 
-from .object_store import ObjectStore
+from .object_store import NoObjectInStoreError, ObjectStore
 from .object_store_factory import ObjectStoreFactory
 
 __all__ = [
+    "NoObjectInStoreError",
     "ObjectStore",
     "ObjectStoreFactory",
 ]