flwr-nightly 1.15.0.dev20250107__py3-none-any.whl → 1.15.0.dev20250109__py3-none-any.whl

flwr/proto/fleet_pb2.pyi

@@ -3,6 +3,7 @@
 isort:skip_file
 """
 import builtins
+import flwr.proto.message_pb2
 import flwr.proto.node_pb2
 import flwr.proto.task_pb2
 import google.protobuf.descriptor
@@ -169,6 +170,89 @@ class PushTaskResResponse(google.protobuf.message.Message):
     def ClearField(self, field_name: typing_extensions.Literal["reconnect",b"reconnect","results",b"results"]) -> None: ...
 global___PushTaskResResponse = PushTaskResResponse
 
+class PullMessagesRequest(google.protobuf.message.Message):
+    """PullMessages messages"""
+    DESCRIPTOR: google.protobuf.descriptor.Descriptor
+    NODE_FIELD_NUMBER: builtins.int
+    MESSAGE_IDS_FIELD_NUMBER: builtins.int
+    @property
+    def node(self) -> flwr.proto.node_pb2.Node: ...
+    @property
+    def message_ids(self) -> google.protobuf.internal.containers.RepeatedScalarFieldContainer[typing.Text]: ...
+    def __init__(self,
+        *,
+        node: typing.Optional[flwr.proto.node_pb2.Node] = ...,
+        message_ids: typing.Optional[typing.Iterable[typing.Text]] = ...,
+        ) -> None: ...
+    def HasField(self, field_name: typing_extensions.Literal["node",b"node"]) -> builtins.bool: ...
+    def ClearField(self, field_name: typing_extensions.Literal["message_ids",b"message_ids","node",b"node"]) -> None: ...
+global___PullMessagesRequest = PullMessagesRequest
+
+class PullMessagesResponse(google.protobuf.message.Message):
+    DESCRIPTOR: google.protobuf.descriptor.Descriptor
+    RECONNECT_FIELD_NUMBER: builtins.int
+    MESSAGES_LIST_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]: ...
+    def __init__(self,
+        *,
+        reconnect: typing.Optional[global___Reconnect] = ...,
+        messages_list: typing.Optional[typing.Iterable[flwr.proto.message_pb2.Message]] = ...,
+        ) -> 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","reconnect",b"reconnect"]) -> None: ...
+global___PullMessagesResponse = PullMessagesResponse
+
+class PushMessagesRequest(google.protobuf.message.Message):
+    """PushMessages messages"""
+    DESCRIPTOR: google.protobuf.descriptor.Descriptor
+    NODE_FIELD_NUMBER: builtins.int
+    MESSAGES_LIST_FIELD_NUMBER: builtins.int
+    @property
+    def node(self) -> flwr.proto.node_pb2.Node: ...
+    @property
+    def messages_list(self) -> google.protobuf.internal.containers.RepeatedCompositeFieldContainer[flwr.proto.message_pb2.Message]: ...
+    def __init__(self,
+        *,
+        node: typing.Optional[flwr.proto.node_pb2.Node] = ...,
+        messages_list: typing.Optional[typing.Iterable[flwr.proto.message_pb2.Message]] = ...,
+        ) -> None: ...
+    def HasField(self, field_name: typing_extensions.Literal["node",b"node"]) -> builtins.bool: ...
+    def ClearField(self, field_name: typing_extensions.Literal["messages_list",b"messages_list","node",b"node"]) -> None: ...
+global___PushMessagesRequest = PushMessagesRequest
+
+class PushMessagesResponse(google.protobuf.message.Message):
+    DESCRIPTOR: google.protobuf.descriptor.Descriptor
+    class ResultsEntry(google.protobuf.message.Message):
+        DESCRIPTOR: google.protobuf.descriptor.Descriptor
+        KEY_FIELD_NUMBER: builtins.int
+        VALUE_FIELD_NUMBER: builtins.int
+        key: typing.Text
+        value: builtins.int
+        def __init__(self,
+            *,
+            key: typing.Text = ...,
+            value: builtins.int = ...,
+            ) -> None: ...
+        def ClearField(self, field_name: typing_extensions.Literal["key",b"key","value",b"value"]) -> None: ...
+
+    RECONNECT_FIELD_NUMBER: builtins.int
+    RESULTS_FIELD_NUMBER: builtins.int
+    @property
+    def reconnect(self) -> global___Reconnect: ...
+    @property
+    def results(self) -> google.protobuf.internal.containers.ScalarMap[typing.Text, builtins.int]: ...
+    def __init__(self,
+        *,
+        reconnect: typing.Optional[global___Reconnect] = ...,
+        results: typing.Optional[typing.Mapping[typing.Text, builtins.int]] = ...,
+        ) -> None: ...
+    def HasField(self, field_name: typing_extensions.Literal["reconnect",b"reconnect"]) -> builtins.bool: ...
+    def ClearField(self, field_name: typing_extensions.Literal["reconnect",b"reconnect","results",b"results"]) -> None: ...
+global___PushMessagesResponse = PushMessagesResponse
+
 class Reconnect(google.protobuf.message.Message):
     DESCRIPTOR: google.protobuf.descriptor.Descriptor
     RECONNECT_FIELD_NUMBER: builtins.int

flwr/proto/fleet_pb2_grpc.py

@@ -36,11 +36,21 @@ class FleetStub(object):
                 request_serializer=flwr_dot_proto_dot_fleet__pb2.PullTaskInsRequest.SerializeToString,
                 response_deserializer=flwr_dot_proto_dot_fleet__pb2.PullTaskInsResponse.FromString,
                 )
+        self.PullMessages = channel.unary_unary(
+                '/flwr.proto.Fleet/PullMessages',
+                request_serializer=flwr_dot_proto_dot_fleet__pb2.PullMessagesRequest.SerializeToString,
+                response_deserializer=flwr_dot_proto_dot_fleet__pb2.PullMessagesResponse.FromString,
+                )
         self.PushTaskRes = channel.unary_unary(
                 '/flwr.proto.Fleet/PushTaskRes',
                 request_serializer=flwr_dot_proto_dot_fleet__pb2.PushTaskResRequest.SerializeToString,
                 response_deserializer=flwr_dot_proto_dot_fleet__pb2.PushTaskResResponse.FromString,
                 )
+        self.PushMessages = channel.unary_unary(
+                '/flwr.proto.Fleet/PushMessages',
+                request_serializer=flwr_dot_proto_dot_fleet__pb2.PushMessagesRequest.SerializeToString,
+                response_deserializer=flwr_dot_proto_dot_fleet__pb2.PushMessagesResponse.FromString,
+                )
         self.GetRun = channel.unary_unary(
                 '/flwr.proto.Fleet/GetRun',
                 request_serializer=flwr_dot_proto_dot_run__pb2.GetRunRequest.SerializeToString,
@@ -83,6 +93,12 @@ class FleetServicer(object):
         context.set_details('Method not implemented!')
         raise NotImplementedError('Method not implemented!')
 
+    def PullMessages(self, request, context):
+        """Missing associated documentation comment in .proto file."""
+        context.set_code(grpc.StatusCode.UNIMPLEMENTED)
+        context.set_details('Method not implemented!')
+        raise NotImplementedError('Method not implemented!')
+
     def PushTaskRes(self, request, context):
         """Complete one or more tasks, if possible
 
@@ -92,6 +108,12 @@ class FleetServicer(object):
         context.set_details('Method not implemented!')
         raise NotImplementedError('Method not implemented!')
 
+    def PushMessages(self, request, context):
+        """Missing associated documentation comment in .proto file."""
+        context.set_code(grpc.StatusCode.UNIMPLEMENTED)
+        context.set_details('Method not implemented!')
+        raise NotImplementedError('Method not implemented!')
+
     def GetRun(self, request, context):
         """Missing associated documentation comment in .proto file."""
         context.set_code(grpc.StatusCode.UNIMPLEMENTED)
@@ -128,11 +150,21 @@ def add_FleetServicer_to_server(servicer, server):
                     request_deserializer=flwr_dot_proto_dot_fleet__pb2.PullTaskInsRequest.FromString,
                     response_serializer=flwr_dot_proto_dot_fleet__pb2.PullTaskInsResponse.SerializeToString,
             ),
+            'PullMessages': grpc.unary_unary_rpc_method_handler(
+                    servicer.PullMessages,
+                    request_deserializer=flwr_dot_proto_dot_fleet__pb2.PullMessagesRequest.FromString,
+                    response_serializer=flwr_dot_proto_dot_fleet__pb2.PullMessagesResponse.SerializeToString,
+            ),
             'PushTaskRes': grpc.unary_unary_rpc_method_handler(
                     servicer.PushTaskRes,
                     request_deserializer=flwr_dot_proto_dot_fleet__pb2.PushTaskResRequest.FromString,
                     response_serializer=flwr_dot_proto_dot_fleet__pb2.PushTaskResResponse.SerializeToString,
             ),
+            'PushMessages': grpc.unary_unary_rpc_method_handler(
+                    servicer.PushMessages,
+                    request_deserializer=flwr_dot_proto_dot_fleet__pb2.PushMessagesRequest.FromString,
+                    response_serializer=flwr_dot_proto_dot_fleet__pb2.PushMessagesResponse.SerializeToString,
+            ),
             'GetRun': grpc.unary_unary_rpc_method_handler(
                     servicer.GetRun,
                     request_deserializer=flwr_dot_proto_dot_run__pb2.GetRunRequest.FromString,
@@ -221,6 +253,23 @@ class Fleet(object):
             options, channel_credentials,
             insecure, call_credentials, compression, wait_for_ready, timeout, metadata)
 
+    @staticmethod
+    def PullMessages(request,
+            target,
+            options=(),
+            channel_credentials=None,
+            call_credentials=None,
+            insecure=False,
+            compression=None,
+            wait_for_ready=None,
+            timeout=None,
+            metadata=None):
+        return grpc.experimental.unary_unary(request, target, '/flwr.proto.Fleet/PullMessages',
+            flwr_dot_proto_dot_fleet__pb2.PullMessagesRequest.SerializeToString,
+            flwr_dot_proto_dot_fleet__pb2.PullMessagesResponse.FromString,
+            options, channel_credentials,
+            insecure, call_credentials, compression, wait_for_ready, timeout, metadata)
+
     @staticmethod
     def PushTaskRes(request,
             target,
@@ -238,6 +287,23 @@ class Fleet(object):
             options, channel_credentials,
             insecure, call_credentials, compression, wait_for_ready, timeout, metadata)
 
+    @staticmethod
+    def PushMessages(request,
+            target,
+            options=(),
+            channel_credentials=None,
+            call_credentials=None,
+            insecure=False,
+            compression=None,
+            wait_for_ready=None,
+            timeout=None,
+            metadata=None):
+        return grpc.experimental.unary_unary(request, target, '/flwr.proto.Fleet/PushMessages',
+            flwr_dot_proto_dot_fleet__pb2.PushMessagesRequest.SerializeToString,
+            flwr_dot_proto_dot_fleet__pb2.PushMessagesResponse.FromString,
+            options, channel_credentials,
+            insecure, call_credentials, compression, wait_for_ready, timeout, metadata)
+
     @staticmethod
     def GetRun(request,
             target,

flwr/proto/fleet_pb2_grpc.pyi

@@ -30,6 +30,10 @@ class FleetStub:
     HTTP API path: /api/v1/fleet/pull-task-ins
     """
 
+    PullMessages: grpc.UnaryUnaryMultiCallable[
+        flwr.proto.fleet_pb2.PullMessagesRequest,
+        flwr.proto.fleet_pb2.PullMessagesResponse]
+
     PushTaskRes: grpc.UnaryUnaryMultiCallable[
         flwr.proto.fleet_pb2.PushTaskResRequest,
         flwr.proto.fleet_pb2.PushTaskResResponse]
@@ -38,6 +42,10 @@ class FleetStub:
     HTTP API path: /api/v1/fleet/push-task-res
     """
 
+    PushMessages: grpc.UnaryUnaryMultiCallable[
+        flwr.proto.fleet_pb2.PushMessagesRequest,
+        flwr.proto.fleet_pb2.PushMessagesResponse]
+
     GetRun: grpc.UnaryUnaryMultiCallable[
         flwr.proto.run_pb2.GetRunRequest,
         flwr.proto.run_pb2.GetRunResponse]
@@ -78,6 +86,12 @@ class FleetServicer(metaclass=abc.ABCMeta):
         """
         pass
 
+    @abc.abstractmethod
+    def PullMessages(self,
+        request: flwr.proto.fleet_pb2.PullMessagesRequest,
+        context: grpc.ServicerContext,
+    ) -> flwr.proto.fleet_pb2.PullMessagesResponse: ...
+
     @abc.abstractmethod
     def PushTaskRes(self,
         request: flwr.proto.fleet_pb2.PushTaskResRequest,
@@ -89,6 +103,12 @@ class FleetServicer(metaclass=abc.ABCMeta):
         """
         pass
 
+    @abc.abstractmethod
+    def PushMessages(self,
+        request: flwr.proto.fleet_pb2.PushMessagesRequest,
+        context: grpc.ServicerContext,
+    ) -> flwr.proto.fleet_pb2.PushMessagesResponse: ...
+
     @abc.abstractmethod
     def GetRun(self,
         request: flwr.proto.run_pb2.GetRunRequest,

flwr/server/app.py

@@ -263,11 +263,10 @@ def run_superlink() -> None:
     # Obtain certificates
     certificates = try_obtain_server_certificates(args, args.fleet_api_type)
 
-    user_auth_config = _try_obtain_user_auth_config(args)
     auth_plugin: Optional[ExecAuthPlugin] = None
-    # user_auth_config is None only if the args.user_auth_config is not provided
-    if user_auth_config is not None:
-        auth_plugin = _try_obtain_exec_auth_plugin(user_auth_config)
+    # 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))
 
     # Initialize StateFactory
     state_factory = LinkStateFactory(args.database)
@@ -584,21 +583,20 @@ def _try_setup_node_authentication(
         )
 
 
-def _try_obtain_user_auth_config(args: argparse.Namespace) -> Optional[dict[str, Any]]:
-    if getattr(args, "user_auth_config", None) is not None:
-        with open(args.user_auth_config, encoding="utf-8") as file:
-            config: dict[str, Any] = yaml.safe_load(file)
-            return config
-    return None
+def _try_obtain_exec_auth_plugin(config_path: Path) -> Optional[ExecAuthPlugin]:
+    # Load YAML file
+    with config_path.open("r", encoding="utf-8") as file:
+        config: dict[str, Any] = yaml.safe_load(file)
 
-
-def _try_obtain_exec_auth_plugin(config: dict[str, Any]) -> Optional[ExecAuthPlugin]:
+    # Load authentication configuration
     auth_config: dict[str, Any] = config.get("authentication", {})
     auth_type: str = auth_config.get(AUTH_TYPE, "")
+
+    # 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(config=auth_config)
+        return auth_plugin_class(user_auth_config_path=config_path)
     except KeyError:
         if auth_type != "":
             sys.exit(

flwr/server/superlink/driver/serverappio_servicer.py

@@ -118,8 +118,9 @@ class ServerAppIoServicer(serverappio_pb2_grpc.ServerAppIoServicer):
             ffs: Ffs = self.ffs_factory.ffs()
             fab_hash = ffs.put(fab.content, {})
             _raise_if(
-                fab_hash != fab.hash_str,
-                f"FAB ({fab.hash_str}) hash from request doesn't match contents",
+                validation_error=fab_hash != fab.hash_str,
+                request_name="CreateRun",
+                detail=f"FAB ({fab.hash_str}) hash from request doesn't match contents",
             )
         else:
             fab_hash = ""
@@ -155,12 +156,22 @@ class ServerAppIoServicer(serverappio_pb2_grpc.ServerAppIoServicer):
             task_ins.task.pushed_at = pushed_at
 
         # Validate request
-        _raise_if(len(request.task_ins_list) == 0, "`task_ins_list` must not be empty")
+        _raise_if(
+            validation_error=len(request.task_ins_list) == 0,
+            request_name="PushTaskIns",
+            detail="`task_ins_list` must not be empty",
+        )
         for task_ins in request.task_ins_list:
             validation_errors = validate_task_ins_or_res(task_ins)
-            _raise_if(bool(validation_errors), ", ".join(validation_errors))
             _raise_if(
-                request.run_id != task_ins.run_id, "`task_ins` has mismatched `run_id`"
+                validation_error=bool(validation_errors),
+                request_name="PushTaskIns",
+                detail=", ".join(validation_errors),
+            )
+            _raise_if(
+                validation_error=request.run_id != task_ins.run_id,
+                request_name="PushTaskIns",
+                detail="`task_ins` has mismatched `run_id`",
             )
 
         # Store each TaskIns
@@ -199,7 +210,9 @@ class ServerAppIoServicer(serverappio_pb2_grpc.ServerAppIoServicer):
         # Validate request
         for task_res in task_res_list:
             _raise_if(
-                request.run_id != task_res.run_id, "`task_res` has mismatched `run_id`"
+                validation_error=request.run_id != task_res.run_id,
+                request_name="PullTaskRes",
+                detail="`task_res` has mismatched `run_id`",
             )
 
         # Delete the TaskIns/TaskRes pairs if TaskRes is found
@@ -344,6 +357,7 @@ class ServerAppIoServicer(serverappio_pb2_grpc.ServerAppIoServicer):
         return GetRunStatusResponse(run_status_dict=run_status_dict)
 
 
-def _raise_if(validation_error: bool, detail: str) -> None:
+def _raise_if(validation_error: bool, request_name: str, detail: str) -> None:
+    """Raise a `ValueError` with a detailed message if a validation error occurs."""
     if validation_error:
-        raise ValueError(f"Malformed PushTaskInsRequest: {detail}")
+        raise ValueError(f"Malformed {request_name}: {detail}")

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

@@ -30,8 +30,12 @@ from flwr.proto.fleet_pb2 import (  # pylint: disable=E0611
     DeleteNodeResponse,
     PingRequest,
     PingResponse,
+    PullMessagesRequest,
+    PullMessagesResponse,
     PullTaskInsRequest,
     PullTaskInsResponse,
+    PushMessagesRequest,
+    PushMessagesResponse,
     PushTaskResRequest,
     PushTaskResResponse,
 )
@@ -95,6 +99,12 @@ class FleetServicer(fleet_pb2_grpc.FleetServicer):
             state=self.state_factory.state(),
         )
 
+    def PullMessages(
+        self, request: PullMessagesRequest, context: grpc.ServicerContext
+    ) -> PullMessagesResponse:
+        """Pull Messages."""
+        return PullMessagesResponse()
+
     def PushTaskRes(
         self, request: PushTaskResRequest, context: grpc.ServicerContext
     ) -> PushTaskResResponse:
@@ -118,6 +128,12 @@ class FleetServicer(fleet_pb2_grpc.FleetServicer):
 
         return res
 
+    def PushMessages(
+        self, request: PushMessagesRequest, context: grpc.ServicerContext
+    ) -> PushMessagesResponse:
+        """Push Messages."""
+        return PushMessagesResponse()
+
     def GetRun(
         self, request: GetRunRequest, context: grpc.ServicerContext
     ) -> GetRunResponse:

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

@@ -223,5 +223,6 @@ class AuthenticateServerInterceptor(grpc.ServerInterceptor):  # type: ignore
         # No `node_id` exists for the provided `public_key`
         # Handle `CreateNode` here instead of calling the default method handler
         # Note: the innermost `CreateNode` method will never be called
-        node_id = state.create_node(request.ping_interval, public_key_bytes)
+        node_id = state.create_node(request.ping_interval)
+        state.set_node_public_key(node_id, public_key_bytes)
         return CreateNodeResponse(node=Node(node_id=node_id, anonymous=False))

flwr/server/superlink/linkstate/in_memory_linkstate.py

@@ -62,6 +62,7 @@ class InMemoryLinkState(LinkState):  # pylint: disable=R0902,R0904
         # Map node_id to (online_until, ping_interval)
         self.node_ids: dict[int, tuple[float, float]] = {}
         self.public_key_to_node_id: dict[bytes, int] = {}
+        self.node_id_to_public_key: dict[int, bytes] = {}
 
         # Map run_id to RunRecord
         self.run_ids: dict[int, RunRecord] = {}
@@ -306,9 +307,7 @@ class InMemoryLinkState(LinkState):  # pylint: disable=R0902,R0904
         """
         return len(self.task_res_store)
 
-    def create_node(
-        self, ping_interval: float, public_key: Optional[bytes] = None
-    ) -> int:
+    def create_node(self, ping_interval: float) -> int:
         """Create, store in the link state, and return `node_id`."""
         # Sample a random int64 as node_id
         node_id = generate_rand_int_from_bytes(NODE_ID_NUM_BYTES)
@@ -318,33 +317,18 @@ class InMemoryLinkState(LinkState):  # pylint: disable=R0902,R0904
                 log(ERROR, "Unexpected node registration failure.")
                 return 0
 
-            if public_key is not None:
-                if (
-                    public_key in self.public_key_to_node_id
-                    or node_id in self.public_key_to_node_id.values()
-                ):
-                    log(ERROR, "Unexpected node registration failure.")
-                    return 0
-
-                self.public_key_to_node_id[public_key] = node_id
-
             self.node_ids[node_id] = (time.time() + ping_interval, ping_interval)
             return node_id
 
-    def delete_node(self, node_id: int, public_key: Optional[bytes] = None) -> None:
+    def delete_node(self, node_id: int) -> None:
         """Delete a node."""
         with self.lock:
             if node_id not in self.node_ids:
                 raise ValueError(f"Node {node_id} not found")
 
-            if public_key is not None:
-                if (
-                    public_key not in self.public_key_to_node_id
-                    or node_id not in self.public_key_to_node_id.values()
-                ):
-                    raise ValueError("Public key or node_id not found")
-
-                del self.public_key_to_node_id[public_key]
+            # Remove node ID <> public key mappings
+            if pk := self.node_id_to_public_key.pop(node_id, None):
+                del self.public_key_to_node_id[pk]
 
             del self.node_ids[node_id]
 
@@ -366,6 +350,26 @@ class InMemoryLinkState(LinkState):  # pylint: disable=R0902,R0904
                 if online_until > current_time
             }
 
+    def set_node_public_key(self, node_id: int, public_key: bytes) -> None:
+        """Set `public_key` for the specified `node_id`."""
+        with self.lock:
+            if node_id not in self.node_ids:
+                raise ValueError(f"Node {node_id} not found")
+
+            if public_key in self.public_key_to_node_id:
+                raise ValueError("Public key already in use")
+
+            self.public_key_to_node_id[public_key] = node_id
+            self.node_id_to_public_key[node_id] = public_key
+
+    def get_node_public_key(self, node_id: int) -> Optional[bytes]:
+        """Get `public_key` for the specified `node_id`."""
+        with self.lock:
+            if node_id not in self.node_ids:
+                raise ValueError(f"Node {node_id} not found")
+
+            return self.node_id_to_public_key.get(node_id)
+
     def get_node_id(self, node_public_key: bytes) -> Optional[int]:
         """Retrieve stored `node_id` filtered by `node_public_keys`."""
         return self.public_key_to_node_id.get(node_public_key)

flwr/server/superlink/linkstate/linkstate.py

@@ -154,13 +154,11 @@ class LinkState(abc.ABC):  # pylint: disable=R0904
         """Get all TaskIns IDs for the given run_id."""
 
     @abc.abstractmethod
-    def create_node(
-        self, ping_interval: float, public_key: Optional[bytes] = None
-    ) -> int:
+    def create_node(self, ping_interval: float) -> int:
         """Create, store in the link state, and return `node_id`."""
 
     @abc.abstractmethod
-    def delete_node(self, node_id: int, public_key: Optional[bytes] = None) -> None:
+    def delete_node(self, node_id: int) -> None:
         """Remove `node_id` from the link state."""
 
     @abc.abstractmethod
@@ -173,6 +171,14 @@ class LinkState(abc.ABC):  # pylint: disable=R0904
         an empty `Set` MUST be returned.
         """
 
+    @abc.abstractmethod
+    def set_node_public_key(self, node_id: int, public_key: bytes) -> None:
+        """Set `public_key` for the specified `node_id`."""
+
+    @abc.abstractmethod
+    def get_node_public_key(self, node_id: int) -> Optional[bytes]:
+        """Get `public_key` for the specified `node_id`."""
+
     @abc.abstractmethod
     def get_node_id(self, node_public_key: bytes) -> Optional[int]:
         """Retrieve stored `node_id` filtered by `node_public_keys`."""