flwr-nightly 1.15.0.dev20250108__py3-none-any.whl → 1.15.0.dev20250110__py3-none-any.whl

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

@@ -18,7 +18,8 @@
 import argparse
 import csv
 import importlib.util
-import subprocess
+import multiprocessing
+import multiprocessing.context
 import sys
 import threading
 from collections.abc import Sequence
@@ -59,6 +60,7 @@ from flwr.common.constant import (
     TRANSPORT_TYPE_REST,
 )
 from flwr.common.exit_handlers import register_exit_handlers
+from flwr.common.grpc import generic_create_grpc_server
 from flwr.common.logger import log, warn_deprecated_feature
 from flwr.common.secure_aggregation.crypto.symmetric_encryption import (
     private_key_to_bytes,
@@ -68,6 +70,8 @@ from flwr.proto.fleet_pb2_grpc import (  # pylint: disable=E0611
     add_FleetServicer_to_server,
 )
 from flwr.proto.grpcadapter_pb2_grpc import add_GrpcAdapterServicer_to_server
+from flwr.server.serverapp.app import flwr_serverapp
+from flwr.simulation.app import flwr_simulation
 from flwr.superexec.app import load_executor
 from flwr.superexec.exec_grpc import run_exec_api_grpc
 
@@ -79,10 +83,7 @@ from .strategy import Strategy
 from .superlink.driver.serverappio_grpc import run_serverappio_api_grpc
 from .superlink.ffs.ffs_factory import FfsFactory
 from .superlink.fleet.grpc_adapter.grpc_adapter_servicer import GrpcAdapterServicer
-from .superlink.fleet.grpc_bidi.grpc_server import (
-    generic_create_grpc_server,
-    start_grpc_server,
-)
+from .superlink.fleet.grpc_bidi.grpc_server import start_grpc_server
 from .superlink.fleet.grpc_rere.fleet_servicer import FleetServicer
 from .superlink.fleet.grpc_rere.server_interceptor import AuthenticateServerInterceptor
 from .superlink.linkstate import LinkStateFactory
@@ -292,7 +293,7 @@ def run_superlink() -> None:
     # Determine Exec plugin
     # If simulation is used, don't start ServerAppIo and Fleet APIs
     sim_exec = executor.__class__.__qualname__ == "SimulationEngine"
-    bckg_threads = []
+    bckg_threads: list[threading.Thread] = []
 
     if sim_exec:
         simulationio_server: grpc.Server = run_simulationio_api_grpc(
@@ -360,6 +361,7 @@ def run_superlink() -> None:
                     ffs_factory,
                     num_workers,
                 ),
+                daemon=True,
             )
             fleet_thread.start()
             bckg_threads.append(fleet_thread)
@@ -426,6 +428,7 @@ def run_superlink() -> None:
                 address,
                 cmd,
             ),
+            daemon=True,
         )
         scheduler_th.start()
         bckg_threads.append(scheduler_th)
@@ -434,16 +437,24 @@ def run_superlink() -> None:
     register_exit_handlers(
         event_type=EventType.RUN_SUPERLINK_LEAVE,
         grpc_servers=grpc_servers,
-        bckg_threads=bckg_threads,
     )
 
-    # Block
-    while True:
-        if bckg_threads:
-            for thread in bckg_threads:
-                if not thread.is_alive():
-                    sys.exit(1)
-        exec_server.wait_for_termination(timeout=1)
+    # Block until a thread exits prematurely
+    while all(thread.is_alive() for thread in bckg_threads):
+        sleep(0.1)
+
+    # Exit if any thread has exited prematurely
+    sys.exit(1)
+
+
+def _run_flwr_command(args: list[str]) -> None:
+    sys.argv = args
+    if args[0] == "flwr-serverapp":
+        flwr_serverapp()
+    elif args[0] == "flwr-simulation":
+        flwr_simulation()
+    else:
+        raise ValueError(f"Unknown command: {args[0]}")
 
 
 def _flwr_scheduler(
@@ -453,15 +464,18 @@ def _flwr_scheduler(
     cmd: str,
 ) -> None:
     log(DEBUG, "Started %s scheduler thread.", cmd)
-
     state = state_factory.state()
+    run_id_to_proc: dict[int, multiprocessing.context.SpawnProcess] = {}
+
+    # Use the "spawn" start method for multiprocessing.
+    mp_spawn_context = multiprocessing.get_context("spawn")
 
     # Periodically check for a pending run in the LinkState
     while True:
-        sleep(3)
+        sleep(0.1)
         pending_run_id = state.get_pending_run_id()
 
-        if pending_run_id:
+        if pending_run_id and pending_run_id not in run_id_to_proc:
 
             log(
                 INFO,
@@ -478,10 +492,18 @@ def _flwr_scheduler(
                 "--insecure",
             ]
 
-            subprocess.Popen(  # pylint: disable=consider-using-with
-                command,
-                text=True,
+            proc = mp_spawn_context.Process(
+                target=_run_flwr_command, args=(command,), daemon=True
             )
+            proc.start()
+
+            # Store the process
+            run_id_to_proc[pending_run_id] = proc
+
+        # Clean up finished processes
+        for run_id, proc in list(run_id_to_proc.items()):
+            if not proc.is_alive():
+                del run_id_to_proc[run_id]
 
 
 def _format_address(address: str) -> tuple[str, str, int]:

flwr/server/superlink/driver/serverappio_grpc.py

@@ -21,6 +21,7 @@ from typing import Optional
 import grpc
 
 from flwr.common import GRPC_MAX_MESSAGE_LENGTH
+from flwr.common.grpc import generic_create_grpc_server
 from flwr.common.logger import log
 from flwr.proto.serverappio_pb2_grpc import (  # pylint: disable=E0611
     add_ServerAppIoServicer_to_server,
@@ -28,7 +29,6 @@ from flwr.proto.serverappio_pb2_grpc import (  # pylint: disable=E0611
 from flwr.server.superlink.ffs.ffs_factory import FfsFactory
 from flwr.server.superlink.linkstate import LinkStateFactory
 
-from ..fleet.grpc_bidi.grpc_server import generic_create_grpc_server
 from .serverappio_servicer import ServerAppIoServicer
 
 

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_bidi/grpc_server.py

@@ -15,49 +15,19 @@
 """Implements utility function to create a gRPC server."""
 
 
-import concurrent.futures
-import sys
-from collections.abc import Sequence
-from logging import ERROR
-from typing import Any, Callable, Optional, Union
+from typing import Optional
 
 import grpc
 
 from flwr.common import GRPC_MAX_MESSAGE_LENGTH
-from flwr.common.address import is_port_in_use
-from flwr.common.logger import log
+from flwr.common.grpc import generic_create_grpc_server
 from flwr.proto.transport_pb2_grpc import (  # pylint: disable=E0611
     add_FlowerServiceServicer_to_server,
 )
 from flwr.server.client_manager import ClientManager
-from flwr.server.superlink.driver.serverappio_servicer import ServerAppIoServicer
-from flwr.server.superlink.fleet.grpc_adapter.grpc_adapter_servicer import (
-    GrpcAdapterServicer,
-)
 from flwr.server.superlink.fleet.grpc_bidi.flower_service_servicer import (
     FlowerServiceServicer,
 )
-from flwr.server.superlink.fleet.grpc_rere.fleet_servicer import FleetServicer
-
-INVALID_CERTIFICATES_ERR_MSG = """
-    When setting any of root_certificate, certificate, or private_key,
-    all of them need to be set.
-"""
-
-AddServicerToServerFn = Callable[..., Any]
-
-
-def valid_certificates(certificates: tuple[bytes, bytes, bytes]) -> bool:
-    """Validate certificates tuple."""
-    is_valid = (
-        all(isinstance(certificate, bytes) for certificate in certificates)
-        and len(certificates) == 3
-    )
-
-    if not is_valid:
-        log(ERROR, INVALID_CERTIFICATES_ERR_MSG)
-
-    return is_valid
 
 
 def start_grpc_server(  # pylint: disable=too-many-arguments,R0917
@@ -154,136 +124,3 @@ def start_grpc_server(  # pylint: disable=too-many-arguments,R0917
     server.start()
 
     return server
-
-
-def generic_create_grpc_server(  # pylint: disable=too-many-arguments,R0917
-    servicer_and_add_fn: Union[
-        tuple[FleetServicer, AddServicerToServerFn],
-        tuple[GrpcAdapterServicer, AddServicerToServerFn],
-        tuple[FlowerServiceServicer, AddServicerToServerFn],
-        tuple[ServerAppIoServicer, AddServicerToServerFn],
-    ],
-    server_address: str,
-    max_concurrent_workers: int = 1000,
-    max_message_length: int = GRPC_MAX_MESSAGE_LENGTH,
-    keepalive_time_ms: int = 210000,
-    certificates: Optional[tuple[bytes, bytes, bytes]] = None,
-    interceptors: Optional[Sequence[grpc.ServerInterceptor]] = None,
-) -> grpc.Server:
-    """Create a gRPC server with a single servicer.
-
-    Parameters
-    ----------
-    servicer_and_add_fn : tuple
-        A tuple holding a servicer implementation and a matching
-        add_Servicer_to_server function.
-    server_address : str
-        Server address in the form of HOST:PORT e.g. "[::]:8080"
-    max_concurrent_workers : int
-        Maximum number of clients the server can process before returning
-        RESOURCE_EXHAUSTED status (default: 1000)
-    max_message_length : int
-        Maximum message length that the server can send or receive.
-        Int valued in bytes. -1 means unlimited. (default: GRPC_MAX_MESSAGE_LENGTH)
-    keepalive_time_ms : int
-        Flower uses a default gRPC keepalive time of 210000ms (3 minutes 30 seconds)
-        because some cloud providers (for example, Azure) agressively clean up idle
-        TCP connections by terminating them after some time (4 minutes in the case
-        of Azure). Flower does not use application-level keepalive signals and relies
-        on the assumption that the transport layer will fail in cases where the
-        connection is no longer active. `keepalive_time_ms` can be used to customize
-        the keepalive interval for specific environments. The default Flower gRPC
-        keepalive of 210000 ms (3 minutes 30 seconds) ensures that Flower can keep
-        the long running streaming connection alive in most environments. The actual
-        gRPC default of this setting is 7200000 (2 hours), which results in dropped
-        connections in some cloud environments.
-
-        These settings are related to the issue described here:
-        - https://github.com/grpc/proposal/blob/master/A8-client-side-keepalive.md
-        - https://github.com/grpc/grpc/blob/master/doc/keepalive.md
-        - https://grpc.io/docs/guides/performance/
-
-        Mobile Flower clients may choose to increase this value if their server
-        environment allows long-running idle TCP connections.
-        (default: 210000)
-    certificates : Tuple[bytes, bytes, bytes] (default: None)
-        Tuple containing root certificate, server certificate, and private key to
-        start a secure SSL-enabled server. The tuple is expected to have three bytes
-        elements in the following order:
-
-            * CA certificate.
-            * server certificate.
-            * server private key.
-    interceptors : Optional[Sequence[grpc.ServerInterceptor]] (default: None)
-        A list of gRPC interceptors.
-
-    Returns
-    -------
-    server : grpc.Server
-        A non-running instance of a gRPC server.
-    """
-    # Check if port is in use
-    if is_port_in_use(server_address):
-        sys.exit(f"Port in server address {server_address} is already in use.")
-
-    # Deconstruct tuple into servicer and function
-    servicer, add_servicer_to_server_fn = servicer_and_add_fn
-
-    # Possible options:
-    # https://github.com/grpc/grpc/blob/v1.43.x/include/grpc/impl/codegen/grpc_types.h
-    options = [
-        # Maximum number of concurrent incoming streams to allow on a http2
-        # connection. Int valued.
-        ("grpc.max_concurrent_streams", max(100, max_concurrent_workers)),
-        # Maximum message length that the channel can send.
-        # Int valued, bytes. -1 means unlimited.
-        ("grpc.max_send_message_length", max_message_length),
-        # Maximum message length that the channel can receive.
-        # Int valued, bytes. -1 means unlimited.
-        ("grpc.max_receive_message_length", max_message_length),
-        # The gRPC default for this setting is 7200000 (2 hours). Flower uses a
-        # customized default of 210000 (3 minutes and 30 seconds) to improve
-        # compatibility with popular cloud providers. Mobile Flower clients may
-        # choose to increase this value if their server environment allows
-        # long-running idle TCP connections.
-        ("grpc.keepalive_time_ms", keepalive_time_ms),
-        # Setting this to zero will allow sending unlimited keepalive pings in between
-        # sending actual data frames.
-        ("grpc.http2.max_pings_without_data", 0),
-        # Is it permissible to send keepalive pings from the client without
-        # any outstanding streams. More explanation here:
-        # https://github.com/adap/flower/pull/2197
-        ("grpc.keepalive_permit_without_calls", 0),
-    ]
-
-    server = grpc.server(
-        concurrent.futures.ThreadPoolExecutor(max_workers=max_concurrent_workers),
-        # Set the maximum number of concurrent RPCs this server will service before
-        # returning RESOURCE_EXHAUSTED status, or None to indicate no limit.
-        maximum_concurrent_rpcs=max_concurrent_workers,
-        options=options,
-        interceptors=interceptors,
-    )
-    add_servicer_to_server_fn(servicer, server)
-
-    if certificates is not None:
-        if not valid_certificates(certificates):
-            sys.exit(1)
-
-        root_certificate_b, certificate_b, private_key_b = certificates
-
-        server_credentials = grpc.ssl_server_credentials(
-            ((private_key_b, certificate_b),),
-            root_certificates=root_certificate_b,
-            # A boolean indicating whether or not to require clients to be
-            # authenticated. May only be True if root_certificates is not None.
-            # We are explicitly setting the current gRPC default to document
-            # the option. For further reference see:
-            # https://grpc.github.io/grpc/python/grpc.html#create-server-credentials
-            require_client_auth=False,
-        )
-        server.add_secure_port(server_address, server_credentials)
-    else:
-        server.add_insecure_port(server_address)
-
-    return server

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`."""