flwr-nightly 1.19.0.dev20250527__py3-none-any.whl → 1.19.0.dev20250529__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, Union, cast
 
 import grpc
 import yaml
@@ -155,17 +155,14 @@ def run_superlink() -> 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)
+        # pylint: disable=unused-variable
+        auth_plugin, authz_plugin = _try_obtain_exec_auth_plugins(  # noqa: F841
+            Path(cfg_path), verify_tls_cert
+        )
+        # pylint: enable=unused-variable
         # 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)
@@ -258,6 +255,7 @@ def run_superlink() -> None:
                     args.ssl_certfile,
                     state_factory,
                     ffs_factory,
+                    objectstore_factory,
                     num_workers,
                 ),
                 daemon=True,
@@ -483,62 +481,58 @@ 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[Union[ExecAuthPlugin, ExecAuthzPlugin]]]],
+    ) -> Union[ExecAuthPlugin, ExecAuthzPlugin]:
+        section_cfg = config.get(section, {})
+        auth_plugin_name = section_cfg.get(yaml_key, "")
+        try:
+            plugins = loader()
+            plugin_cls = 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 = cast(
+        ExecAuthPlugin,
+        _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 = cast(
+        ExecAuthzPlugin,
+        _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 +630,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 +646,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/grid/grpc_grid.py

@@ -163,7 +163,7 @@ class GrpcGrid(Grid):
     def _check_message(self, message: Message) -> None:
         # Check if the message is valid
         if not (
-            message.metadata.message_id == ""
+            message.metadata.message_id != ""
             and message.metadata.reply_to_message_id == ""
             and message.metadata.ttl > 0
         ):
@@ -211,6 +211,7 @@ class GrpcGrid(Grid):
             # Populate metadata
             msg.metadata.__dict__["_run_id"] = run_id
             msg.metadata.__dict__["_src_node_id"] = self.node.node_id
+            msg.metadata.__dict__["_message_id"] = msg.object_id
             # Check message
             self._check_message(msg)
             # Convert to proto

flwr/server/grid/inmemory_grid.py

@@ -18,7 +18,7 @@
 import time
 from collections.abc import Iterable
 from typing import Optional, cast
-from uuid import UUID
+from uuid import uuid4
 
 from flwr.common import Message, RecordDict
 from flwr.common.constant import SUPERLINK_NODE_ID
@@ -56,7 +56,7 @@ class InMemoryGrid(Grid):
     def _check_message(self, message: Message) -> None:
         # Check if the message is valid
         if not (
-            message.metadata.message_id == ""
+            message.metadata.message_id != ""
             and message.metadata.reply_to_message_id == ""
             and message.metadata.ttl > 0
             and message.metadata.delivered_at == ""
@@ -111,6 +111,7 @@ class InMemoryGrid(Grid):
             # Populate metadata
             msg.metadata.__dict__["_run_id"] = cast(Run, self._run).run_id
             msg.metadata.__dict__["_src_node_id"] = self.node.node_id
+            msg.metadata.__dict__["_message_id"] = str(uuid4())
             # Check message
             self._check_message(msg)
             # Store in state
@@ -126,12 +127,12 @@ class InMemoryGrid(Grid):
         This method is used to collect messages from the SuperLink that correspond to a
         set of given message IDs.
         """
-        msg_ids = {UUID(msg_id) for msg_id in message_ids}
+        msg_ids = set(message_ids)
         # Pull Messages
         message_res_list = self.state.get_message_res(message_ids=msg_ids)
         # Get IDs of Messages these replies are for
         message_ins_ids_to_delete = {
-            UUID(msg_res.metadata.reply_to_message_id) for msg_res in message_res_list
+            msg_res.metadata.reply_to_message_id for msg_res in message_res_list
         }
         # Delete
         self.state.delete_messages(message_ins_ids=message_ins_ids_to_delete)

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

@@ -133,6 +133,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

@@ -16,7 +16,6 @@
 
 
 from typing import Optional
-from uuid import UUID
 
 from flwr.common import Message
 from flwr.common.constant import Status
@@ -52,6 +51,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 ObjectStore
+
+from ...utils import store_mapping_and_register_objects
 
 
 def create_node(
@@ -106,7 +108,9 @@ def pull_messages(
 
 
 def push_messages(
-    request: PushMessagesRequest, state: LinkState
+    request: PushMessagesRequest,
+    state: LinkState,
+    store: ObjectStore,
 ) -> PushMessagesResponse:
     """Push Messages handler."""
     # Convert Message from proto
@@ -122,12 +126,16 @@ def push_messages(
         raise InvalidRunStatusException(abort_msg)
 
     # Store Message in State
-    message_id: Optional[UUID] = state.store_message_res(message=msg)
+    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
@@ -123,9 +124,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/fleet/vce/vce_api.py

@@ -25,6 +25,7 @@ from pathlib import Path
 from queue import Empty, Queue
 from time import sleep
 from typing import Callable, Optional
+from uuid import uuid4
 
 from flwr.app.error import Error
 from flwr.client.client_app import ClientApp, ClientAppException, LoadClientAppError
@@ -134,6 +135,8 @@ def worker(
 
         finally:
             if out_mssg:
+                # Assign a message_id
+                out_mssg.metadata.__dict__["_message_id"] = str(uuid4())
                 # Store reply Messages in state
                 messageres_queue.put(out_mssg)
 

flwr/server/superlink/linkstate/in_memory_linkstate.py

@@ -21,7 +21,6 @@ from bisect import bisect_right
 from dataclasses import dataclass, field
 from logging import ERROR, WARNING
 from typing import Optional
-from uuid import UUID, uuid4
 
 from flwr.common import Context, Message, log, now
 from flwr.common.constant import (
@@ -76,15 +75,15 @@ class InMemoryLinkState(LinkState):  # pylint: disable=R0902,R0904
         self.run_ids: dict[int, RunRecord] = {}
         self.contexts: dict[int, Context] = {}
         self.federation_options: dict[int, ConfigRecord] = {}
-        self.message_ins_store: dict[UUID, Message] = {}
-        self.message_res_store: dict[UUID, Message] = {}
-        self.message_ins_id_to_message_res_id: dict[UUID, UUID] = {}
+        self.message_ins_store: dict[str, Message] = {}
+        self.message_res_store: dict[str, Message] = {}
+        self.message_ins_id_to_message_res_id: dict[str, str] = {}
 
         self.node_public_keys: set[bytes] = set()
 
         self.lock = threading.RLock()
 
-    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, is_reply_message=False)
@@ -112,12 +111,7 @@ class InMemoryLinkState(LinkState):  # pylint: disable=R0902,R0904
             )
             return None
 
-        # Create message_id
-        message_id = uuid4()
-
-        # Store Message
-        # pylint: disable-next=W0212
-        message.metadata._message_id = str(message_id)  # type: ignore
+        message_id = message.metadata.message_id
         with self.lock:
             self.message_ins_store[message_id] = message
 
@@ -153,7 +147,7 @@ class InMemoryLinkState(LinkState):  # pylint: disable=R0902,R0904
         return message_ins_list
 
     # pylint: disable=R0911
-    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, is_reply_message=True)
@@ -165,7 +159,7 @@ class InMemoryLinkState(LinkState):  # pylint: disable=R0902,R0904
         with self.lock:
             # Check if the Message it is replying to exists and is valid
             msg_ins_id = res_metadata.reply_to_message_id
-            msg_ins = self.message_ins_store.get(UUID(msg_ins_id))
+            msg_ins = self.message_ins_store.get(msg_ins_id)
 
             # Ensure that dst_node_id of original Message matches the src_node_id of
             # reply Message.
@@ -220,22 +214,17 @@ class InMemoryLinkState(LinkState):  # pylint: disable=R0902,R0904
             log(ERROR, "`metadata.run_id` is invalid")
             return None
 
-        # Create message_id
-        message_id = uuid4()
-
-        # Store Message
-        # pylint: disable-next=W0212
-        message.metadata._message_id = str(message_id)  # type: ignore
+        message_id = message.metadata.message_id
         with self.lock:
             self.message_res_store[message_id] = message
-            self.message_ins_id_to_message_res_id[UUID(msg_ins_id)] = message_id
+            self.message_ins_id_to_message_res_id[msg_ins_id] = message_id
 
         # Return the new message_id
         return 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."""
-        ret: dict[UUID, Message] = {}
+        ret: dict[str, Message] = {}
 
         with self.lock:
             current = time.time()
@@ -287,7 +276,7 @@ class InMemoryLinkState(LinkState):  # pylint: disable=R0902,R0904
 
         return list(ret.values())
 
-    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
@@ -304,9 +293,9 @@ class InMemoryLinkState(LinkState):  # pylint: disable=R0902,R0904
                     )
                     del self.message_res_store[message_res_id]
 
-    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."""
-        message_id_list: set[UUID] = set()
+        message_id_list: set[str] = set()
         with self.lock:
             for message_id, message in self.message_ins_store.items():
                 if message.metadata.run_id == run_id:

flwr/server/superlink/linkstate/linkstate.py

@@ -17,7 +17,6 @@
 
 import abc
 from typing import Optional
-from uuid import UUID
 
 from flwr.common import Context, Message
 from flwr.common.record import ConfigRecord
@@ -28,13 +27,13 @@ class LinkState(abc.ABC):  # pylint: disable=R0904
     """Abstract LinkState."""
 
     @abc.abstractmethod
-    def store_message_ins(self, message: Message) -> Optional[UUID]:
+    def store_message_ins(self, message: Message) -> Optional[str]:
         """Store one Message.
 
         Usually, the ServerAppIo API calls this to schedule instructions.
 
         Stores the value of the `message` in the link state and, if successful,
-        returns the `message_id` (UUID) of the `message`. If, for any reason,
+        returns the `message_id` (str) of the `message`. If, for any reason,
         storing the `message` fails, `None` is returned.
 
         Constraints
@@ -61,12 +60,12 @@ class LinkState(abc.ABC):  # pylint: disable=R0904
         """
 
     @abc.abstractmethod
-    def store_message_res(self, message: Message) -> Optional[UUID]:
+    def store_message_res(self, message: Message) -> Optional[str]:
         """Store one Message.
 
         Usually, the Fleet API calls this for Nodes returning results.
 
-        Stores the Message and, if successful, returns the `message_id` (UUID) of
+        Stores the Message and, if successful, returns the `message_id` (str) of
         the `message`. If storing the `message` fails, `None` is returned.
 
         Constraints
@@ -78,7 +77,7 @@ class LinkState(abc.ABC):  # pylint: disable=R0904
         """
 
     @abc.abstractmethod
-    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.
 
         This method is typically called by the ServerAppIo API to obtain
@@ -94,7 +93,7 @@ class LinkState(abc.ABC):  # pylint: disable=R0904
 
         Parameters
         ----------
-        message_ids : set[UUID]
+        message_ids : set[str]
             A set of Message IDs used to retrieve reply Messages responding to them.
 
         Returns
@@ -113,18 +112,18 @@ class LinkState(abc.ABC):  # pylint: disable=R0904
         """Calculate the number of reply Messages in store."""
 
     @abc.abstractmethod
-    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.
 
         Parameters
         ----------
-        message_ins_ids : set[UUID]
+        message_ins_ids : set[str]
             A set of Message IDs. For each ID in the set, the corresponding
             Message and its associated reply Message will be deleted.
         """
 
     @abc.abstractmethod
-    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."""
 
     @abc.abstractmethod