flwr 1.21.0__py3-none-any.whl → 1.23.0__py3-none-any.whl

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

@@ -15,25 +15,31 @@
 """Fleet API gRPC request-response servicer."""
 
 
-from logging import DEBUG, INFO
+import threading
+from logging import DEBUG, ERROR, INFO
 
 import grpc
 from google.protobuf.json_format import MessageToDict
 
+from flwr.common.constant import PUBLIC_KEY_ALREADY_IN_USE_MESSAGE
 from flwr.common.inflatable import UnexpectedObjectContentError
 from flwr.common.logger import log
 from flwr.common.typing import InvalidRunStatusException
 from flwr.proto import fleet_pb2_grpc  # pylint: disable=E0611
 from flwr.proto.fab_pb2 import GetFabRequest, GetFabResponse  # pylint: disable=E0611
 from flwr.proto.fleet_pb2 import (  # pylint: disable=E0611
-    CreateNodeRequest,
-    CreateNodeResponse,
-    DeleteNodeRequest,
-    DeleteNodeResponse,
+    ActivateNodeRequest,
+    ActivateNodeResponse,
+    DeactivateNodeRequest,
+    DeactivateNodeResponse,
     PullMessagesRequest,
     PullMessagesResponse,
     PushMessagesRequest,
     PushMessagesResponse,
+    RegisterNodeFleetRequest,
+    RegisterNodeFleetResponse,
+    UnregisterNodeFleetRequest,
+    UnregisterNodeFleetResponse,
 )
 from flwr.proto.heartbeat_pb2 import (  # pylint: disable=E0611
     SendNodeHeartbeatRequest,
@@ -63,49 +69,137 @@ class FleetServicer(fleet_pb2_grpc.FleetServicer):
         state_factory: LinkStateFactory,
         ffs_factory: FfsFactory,
         objectstore_factory: ObjectStoreFactory,
+        enable_supernode_auth: bool,
     ) -> None:
         self.state_factory = state_factory
         self.ffs_factory = ffs_factory
         self.objectstore_factory = objectstore_factory
+        self.enable_supernode_auth = enable_supernode_auth
+        self.lock = threading.Lock()
 
-    def CreateNode(
-        self, request: CreateNodeRequest, context: grpc.ServicerContext
-    ) -> CreateNodeResponse:
-        """."""
-        log(
-            INFO,
-            "[Fleet.CreateNode] Request heartbeat_interval=%s",
-            request.heartbeat_interval,
-        )
-        log(DEBUG, "[Fleet.CreateNode] Request: %s", MessageToDict(request))
-        response = message_handler.create_node(
-            request=request,
-            state=self.state_factory.state(),
-        )
-        log(INFO, "[Fleet.CreateNode] Created node_id=%s", response.node.node_id)
-        log(DEBUG, "[Fleet.CreateNode] Response: %s", MessageToDict(response))
-        return response
+    def RegisterNode(
+        self, request: RegisterNodeFleetRequest, context: grpc.ServicerContext
+    ) -> RegisterNodeFleetResponse:
+        """Register a node."""
+        # Prevent registration when SuperNode authentication is enabled
+        if self.enable_supernode_auth:
+            log(ERROR, "SuperNode registration is disabled through Fleet API.")
+            context.abort(
+                grpc.StatusCode.FAILED_PRECONDITION,
+                "SuperNode authentication is enabled. "
+                "All SuperNodes must be registered via the CLI.",
+            )
 
-    def DeleteNode(
-        self, request: DeleteNodeRequest, context: grpc.ServicerContext
-    ) -> DeleteNodeResponse:
-        """."""
-        log(INFO, "[Fleet.DeleteNode] Delete node_id=%s", request.node.node_id)
-        log(DEBUG, "[Fleet.DeleteNode] Request: %s", MessageToDict(request))
-        return message_handler.delete_node(
-            request=request,
-            state=self.state_factory.state(),
-        )
+        try:
+            response = message_handler.register_node(
+                request=request,
+                state=self.state_factory.state(),
+            )
+            log(DEBUG, "[Fleet.RegisterNode] Registered node_id=%s", response.node_id)
+            return response
+        except ValueError:
+            # Public key already in use
+            # This should NEVER happen due to the public keys should be automatically
+            # generated and unique for each SuperNode instance.
+            log(
+                ERROR,
+                "[Fleet.RegisterNode] Registration failed: %s",
+                PUBLIC_KEY_ALREADY_IN_USE_MESSAGE,
+            )
+            context.abort(
+                grpc.StatusCode.FAILED_PRECONDITION, PUBLIC_KEY_ALREADY_IN_USE_MESSAGE
+            )
+
+        raise RuntimeError  # Make mypy happy
+
+    def ActivateNode(
+        self, request: ActivateNodeRequest, context: grpc.ServicerContext
+    ) -> ActivateNodeResponse:
+        """Activate a node."""
+        try:
+            response = message_handler.activate_node(
+                request=request,
+                state=self.state_factory.state(),
+            )
+            log(INFO, "[Fleet.ActivateNode] Activated node_id=%s", response.node_id)
+            return response
+        except message_handler.InvalidHeartbeatIntervalError:
+            # Heartbeat interval is invalid
+            log(ERROR, "[Fleet.ActivateNode] Invalid heartbeat interval")
+            context.abort(
+                grpc.StatusCode.INVALID_ARGUMENT, "Invalid heartbeat interval"
+            )
+        except ValueError as e:
+            log(ERROR, "[Fleet.ActivateNode] Activation failed: %s", str(e))
+            context.abort(grpc.StatusCode.FAILED_PRECONDITION, str(e))
+
+        raise RuntimeError  # Make mypy happy
+
+    def DeactivateNode(
+        self, request: DeactivateNodeRequest, context: grpc.ServicerContext
+    ) -> DeactivateNodeResponse:
+        """Deactivate a node."""
+        try:
+            response = message_handler.deactivate_node(
+                request=request,
+                state=self.state_factory.state(),
+            )
+            log(INFO, "[Fleet.DeactivateNode] Deactivated node_id=%s", request.node_id)
+            return response
+        except ValueError as e:
+            log(ERROR, "[Fleet.DeactivateNode] Deactivation failed: %s", str(e))
+            context.abort(grpc.StatusCode.FAILED_PRECONDITION, str(e))
+
+        raise RuntimeError  # Make mypy happy
+
+    def UnregisterNode(
+        self, request: UnregisterNodeFleetRequest, context: grpc.ServicerContext
+    ) -> UnregisterNodeFleetResponse:
+        """Unregister a node."""
+        # Prevent unregistration when SuperNode authentication is enabled
+        if self.enable_supernode_auth:
+            log(ERROR, "SuperNode unregistration is disabled through Fleet API.")
+            context.abort(
+                grpc.StatusCode.FAILED_PRECONDITION,
+                "SuperNode authentication is enabled. "
+                "All SuperNodes must be unregistered via the CLI.",
+            )
+
+        try:
+            response = message_handler.unregister_node(
+                request=request,
+                state=self.state_factory.state(),
+            )
+            log(
+                DEBUG, "[Fleet.UnregisterNode] Unregistered node_id=%s", request.node_id
+            )
+            return response
+        except ValueError as e:
+            log(
+                ERROR,
+                "[Fleet.UnregisterNode] Unregistration failed: %s",
+                str(e),
+            )
+            context.abort(grpc.StatusCode.FAILED_PRECONDITION, str(e))
+            raise RuntimeError from None  # Make mypy happy
 
     def SendNodeHeartbeat(
         self, request: SendNodeHeartbeatRequest, context: grpc.ServicerContext
     ) -> SendNodeHeartbeatResponse:
         """."""
         log(DEBUG, "[Fleet.SendNodeHeartbeat] Request: %s", MessageToDict(request))
-        return message_handler.send_node_heartbeat(
-            request=request,
-            state=self.state_factory.state(),
-        )
+        try:
+            return message_handler.send_node_heartbeat(
+                request=request,
+                state=self.state_factory.state(),
+            )
+        except message_handler.InvalidHeartbeatIntervalError:
+            # Heartbeat interval is invalid
+            log(ERROR, "[Fleet.SendNodeHeartbeat] Invalid heartbeat interval")
+            context.abort(
+                grpc.StatusCode.INVALID_ARGUMENT, "Invalid heartbeat interval"
+            )
+        raise RuntimeError  # Make mypy happy
 
     def PullMessages(
         self, request: PullMessagesRequest, context: grpc.ServicerContext
@@ -183,7 +277,7 @@ class FleetServicer(fleet_pb2_grpc.FleetServicer):
         """Push an object to the ObjectStore."""
         log(
             DEBUG,
-            "[ServerAppIoServicer.PushObject] Push Object with object_id=%s",
+            "[Fleet.PushObject] Push Object with object_id=%s",
             request.object_id,
         )
 
@@ -208,7 +302,7 @@ class FleetServicer(fleet_pb2_grpc.FleetServicer):
         """Pull an object from the ObjectStore."""
         log(
             DEBUG,
-            "[ServerAppIoServicer.PullObject] Pull Object with object_id=%s",
+            "[Fleet.PullObject] Pull Object with object_id=%s",
             request.object_id,
         )
 

flwr/server/superlink/fleet/grpc_rere/{server_interceptor.py → node_auth_server_interceptor.py}

@@ -16,7 +16,7 @@
 
 
 import datetime
-from typing import Any, Callable, Optional, cast
+from typing import Any, Callable, cast
 
 import grpc
 from google.protobuf.message import Message as GrpcMessage
@@ -29,15 +29,12 @@ from flwr.common.constant import (
     TIMESTAMP_HEADER,
     TIMESTAMP_TOLERANCE,
 )
-from flwr.common.secure_aggregation.crypto.symmetric_encryption import (
-    bytes_to_public_key,
-    verify_signature,
-)
 from flwr.proto.fleet_pb2 import (  # pylint: disable=E0611
-    CreateNodeRequest,
-    CreateNodeResponse,
+    ActivateNodeRequest,
+    RegisterNodeFleetRequest,
 )
 from flwr.server.superlink.linkstate import LinkStateFactory
+from flwr.supercore.primitives.asymmetric import bytes_to_public_key, verify_signature
 
 MIN_TIMESTAMP_DIFF = -SYSTEM_TIME_TOLERANCE
 MAX_TIMESTAMP_DIFF = TIMESTAMP_TOLERANCE + SYSTEM_TIME_TOLERANCE
@@ -53,22 +50,17 @@ def _unary_unary_rpc_terminator(
     return grpc.unary_unary_rpc_method_handler(terminate)
 
 
-class AuthenticateServerInterceptor(grpc.ServerInterceptor):  # type: ignore
+class NodeAuthServerInterceptor(grpc.ServerInterceptor):  # type: ignore
     """Server interceptor for node authentication.
 
     Parameters
     ----------
     state_factory : LinkStateFactory
         A factory for creating new instances of LinkState.
-    auto_auth : bool (default: False)
-        If True, nodes are authenticated without requiring their public keys to be
-        pre-stored in the LinkState. If False, only nodes with pre-stored public keys
-        can be authenticated.
     """
 
-    def __init__(self, state_factory: LinkStateFactory, auto_auth: bool = False):
+    def __init__(self, state_factory: LinkStateFactory):
         self.state_factory = state_factory
-        self.auto_auth = auto_auth
 
     def intercept_service(  # pylint: disable=too-many-return-statements
         self,
@@ -85,7 +77,6 @@ class AuthenticateServerInterceptor(grpc.ServerInterceptor):  # type: ignore
         if not handler_call_details.method.startswith("/flwr.proto.Fleet/"):
             return continuation(handler_call_details)
 
-        state = self.state_factory.state()
         metadata_dict = dict(handler_call_details.invocation_metadata)
 
         # Retrieve info from the metadata
@@ -96,11 +87,6 @@ class AuthenticateServerInterceptor(grpc.ServerInterceptor):  # type: ignore
         except KeyError:
             return _unary_unary_rpc_terminator("Missing authentication metadata")
 
-        if not self.auto_auth:
-            # Abort the RPC call if the node public key is not found
-            if node_pk_bytes not in state.get_node_public_keys():
-                return _unary_unary_rpc_terminator("Public key not recognized")
-
         # Verify the signature
         node_pk = bytes_to_public_key(node_pk_bytes)
         if not verify_signature(node_pk, timestamp_iso.encode("ascii"), signature):
@@ -113,50 +99,40 @@ class AuthenticateServerInterceptor(grpc.ServerInterceptor):  # type: ignore
         if not MIN_TIMESTAMP_DIFF < time_diff.total_seconds() < MAX_TIMESTAMP_DIFF:
             return _unary_unary_rpc_terminator("Invalid timestamp")
 
-        # Continue the RPC call
-        expected_node_id = state.get_node_id(node_pk_bytes)
-        if not handler_call_details.method.endswith("CreateNode"):
-            # All calls, except for `CreateNode`, must provide a public key that is
-            # already mapped to a `node_id` (in `LinkState`)
-            if expected_node_id is None:
-                return _unary_unary_rpc_terminator("Invalid node ID")
-        # One of the method handlers in
+        # Continue the RPC call: One of the method handlers in
         # `flwr.server.superlink.fleet.grpc_rere.fleet_server.FleetServicer`
         method_handler: grpc.RpcMethodHandler = continuation(handler_call_details)
-        return self._wrap_method_handler(
-            method_handler, expected_node_id, node_pk_bytes
-        )
+        return self._wrap_method_handler(method_handler, node_pk_bytes)
 
     def _wrap_method_handler(
         self,
         method_handler: grpc.RpcMethodHandler,
-        expected_node_id: Optional[int],
-        node_public_key: bytes,
+        expected_public_key: bytes,
     ) -> grpc.RpcMethodHandler:
         def _generic_method_handler(
             request: GrpcMessage,
             context: grpc.ServicerContext,
         ) -> GrpcMessage:
-            # Verify the node ID
-            if not isinstance(request, CreateNodeRequest):
-                try:
-                    if request.node.node_id != expected_node_id:  # type: ignore
-                        raise ValueError
-                except (AttributeError, ValueError):
-                    context.abort(grpc.StatusCode.UNAUTHENTICATED, "Invalid node ID")
+            # Note: This function runs in a different thread
+            # than the `intercept_service` function.
+
+            # Retrieve the public key
+            if isinstance(request, (RegisterNodeFleetRequest, ActivateNodeRequest)):
+                actual_public_key = request.public_key
+            else:
+                if hasattr(request, "node"):
+                    node_id = request.node.node_id
+                else:
+                    node_id = request.node_id  # type: ignore[attr-defined]
+                actual_public_key = self.state_factory.state().get_node_public_key(
+                    node_id
+                )
+
+            # Verify the public key
+            if actual_public_key != expected_public_key:
+                context.abort(grpc.StatusCode.UNAUTHENTICATED, "Invalid node ID")
 
             response: GrpcMessage = method_handler.unary_unary(request, context)
-
-            # Set the public key after a successful CreateNode request
-            if isinstance(response, CreateNodeResponse):
-                state = self.state_factory.state()
-                try:
-                    state.set_node_public_key(response.node.node_id, node_public_key)
-                except ValueError as e:
-                    # Remove newly created node if setting the public key fails
-                    state.delete_node(response.node.node_id)
-                    context.abort(grpc.StatusCode.UNAUTHENTICATED, str(e))
-
             return response
 
         return grpc.unary_unary_rpc_method_handler(

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

@@ -18,7 +18,12 @@ from logging import ERROR
 from typing import Optional
 
 from flwr.common import Message, log
-from flwr.common.constant import Status
+from flwr.common.constant import (
+    HEARTBEAT_MAX_INTERVAL,
+    HEARTBEAT_MIN_INTERVAL,
+    NOOP_FLWR_AID,
+    Status,
+)
 from flwr.common.inflatable import UnexpectedObjectContentError
 from flwr.common.serde import (
     fab_to_proto,
@@ -29,15 +34,19 @@ from flwr.common.serde import (
 from flwr.common.typing import Fab, InvalidRunStatusException
 from flwr.proto.fab_pb2 import GetFabRequest, GetFabResponse  # pylint: disable=E0611
 from flwr.proto.fleet_pb2 import (  # pylint: disable=E0611
-    CreateNodeRequest,
-    CreateNodeResponse,
-    DeleteNodeRequest,
-    DeleteNodeResponse,
+    ActivateNodeRequest,
+    ActivateNodeResponse,
+    DeactivateNodeRequest,
+    DeactivateNodeResponse,
     PullMessagesRequest,
     PullMessagesResponse,
     PushMessagesRequest,
     PushMessagesResponse,
     Reconnect,
+    RegisterNodeFleetRequest,
+    RegisterNodeFleetResponse,
+    UnregisterNodeFleetRequest,
+    UnregisterNodeFleetResponse,
 )
 from flwr.proto.heartbeat_pb2 import (  # pylint: disable=E0611
     SendNodeHeartbeatRequest,
@@ -51,7 +60,6 @@ from flwr.proto.message_pb2 import (  # pylint: disable=E0611
     PushObjectRequest,
     PushObjectResponse,
 )
-from flwr.proto.node_pb2 import Node  # pylint: disable=E0611
 from flwr.proto.run_pb2 import (  # pylint: disable=E0611
     GetRunRequest,
     GetRunResponse,
@@ -64,25 +72,52 @@ from flwr.supercore.object_store import NoObjectInStoreError, ObjectStore
 from flwr.supercore.object_store.utils import store_mapping_and_register_objects
 
 
-def create_node(
-    request: CreateNodeRequest,  # pylint: disable=unused-argument
+class InvalidHeartbeatIntervalError(Exception):
+    """Invalid heartbeat interval exception."""
+
+
+def register_node(
+    request: RegisterNodeFleetRequest,
     state: LinkState,
-) -> CreateNodeResponse:
-    """."""
-    # Create node
-    node_id = state.create_node(heartbeat_interval=request.heartbeat_interval)
-    return CreateNodeResponse(node=Node(node_id=node_id))
+) -> RegisterNodeFleetResponse:
+    """Register a node (Fleet API only)."""
+    node_id = state.create_node(NOOP_FLWR_AID, request.public_key, 0)
+    return RegisterNodeFleetResponse(node_id=node_id)
 
 
-def delete_node(request: DeleteNodeRequest, state: LinkState) -> DeleteNodeResponse:
-    """."""
-    # Validate node_id
-    if request.node.node_id == 0:  # i.e. unset `node_id`
-        return DeleteNodeResponse()
+def activate_node(
+    request: ActivateNodeRequest,
+    state: LinkState,
+) -> ActivateNodeResponse:
+    """Activate a node."""
+    node_id = state.get_node_id_by_public_key(request.public_key)
+    if node_id is None:
+        raise ValueError("No SuperNode found with the given public key.")
+    _validate_heartbeat_interval(request.heartbeat_interval)
+    if not state.activate_node(node_id, request.heartbeat_interval):
+        raise ValueError(f"SuperNode with node ID {node_id} could not be activated.")
+    return ActivateNodeResponse(node_id=node_id)
+
+
+def deactivate_node(
+    request: DeactivateNodeRequest,
+    state: LinkState,
+) -> DeactivateNodeResponse:
+    """Deactivate a node."""
+    if not state.deactivate_node(request.node_id):
+        raise ValueError(
+            f"SuperNode with node ID {request.node_id} could not be deactivated."
+        )
+    return DeactivateNodeResponse()
+
 
-    # Update state
-    state.delete_node(node_id=request.node.node_id)
-    return DeleteNodeResponse()
+def unregister_node(
+    request: UnregisterNodeFleetRequest,
+    state: LinkState,
+) -> UnregisterNodeFleetResponse:
+    """Unregister a node (Fleet API only)."""
+    state.delete_node(NOOP_FLWR_AID, request.node_id)
+    return UnregisterNodeFleetResponse()
 
 
 def send_node_heartbeat(
@@ -90,6 +125,7 @@ def send_node_heartbeat(
     state: LinkState,  # pylint: disable=unused-argument
 ) -> SendNodeHeartbeatResponse:
     """."""
+    _validate_heartbeat_interval(request.heartbeat_interval)
     res = state.acknowledge_node_heartbeat(
         request.node.node_id, request.heartbeat_interval
     )
@@ -208,7 +244,7 @@ def get_fab(
         raise InvalidRunStatusException(abort_msg)
 
     if result := ffs.get(request.hash_str):
-        fab = Fab(request.hash_str, result[0])
+        fab = Fab(request.hash_str, result[0], result[1])
         return GetFabResponse(fab=fab_to_proto(fab))
 
     raise ValueError(f"Found no FAB with hash: {request.hash_str}")
@@ -284,3 +320,12 @@ def confirm_message_received(
     store.delete(request.message_object_id)
 
     return ConfirmMessageReceivedResponse()
+
+
+def _validate_heartbeat_interval(interval: float) -> None:
+    """Raise if heartbeat interval is out of bounds."""
+    if not HEARTBEAT_MIN_INTERVAL <= interval <= HEARTBEAT_MAX_INTERVAL:
+        raise InvalidHeartbeatIntervalError(
+            f"Heartbeat interval {interval} is out of bounds "
+            f"[{HEARTBEAT_MIN_INTERVAL}, {HEARTBEAT_MAX_INTERVAL}]."
+        )

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

@@ -25,14 +25,18 @@ from google.protobuf.message import Message as GrpcMessage
 from flwr.common.exit import ExitCode, flwr_exit
 from flwr.proto.fab_pb2 import GetFabRequest, GetFabResponse  # pylint: disable=E0611
 from flwr.proto.fleet_pb2 import (  # pylint: disable=E0611
-    CreateNodeRequest,
-    CreateNodeResponse,
-    DeleteNodeRequest,
-    DeleteNodeResponse,
+    ActivateNodeRequest,
+    ActivateNodeResponse,
+    DeactivateNodeRequest,
+    DeactivateNodeResponse,
     PullMessagesRequest,
     PullMessagesResponse,
     PushMessagesRequest,
     PushMessagesResponse,
+    RegisterNodeFleetRequest,
+    RegisterNodeFleetResponse,
+    UnregisterNodeFleetRequest,
+    UnregisterNodeFleetResponse,
 )
 from flwr.proto.heartbeat_pb2 import (  # pylint: disable=E0611
     SendNodeHeartbeatRequest,
@@ -69,6 +73,8 @@ GrpcResponse = TypeVar("GrpcResponse", bound=GrpcMessage)
 GrpcAsyncFunction = Callable[[GrpcRequest], Awaitable[GrpcResponse]]
 RestEndPoint = Callable[[Request], Awaitable[Response]]
 
+routes = []
+
 
 def rest_request_response(
     grpc_request_type: type[GrpcRequest],
@@ -76,6 +82,7 @@ def rest_request_response(
     """Convert an async gRPC-based function into a RESTful HTTP endpoint."""
 
     def decorator(func: GrpcAsyncFunction[GrpcRequest, GrpcResponse]) -> RestEndPoint:
+
         async def wrapper(request: Request) -> Response:
             _check_headers(request.headers)
 
@@ -91,33 +98,64 @@ def rest_request_response(
                 headers={"Content-Type": "application/protobuf"},
             )
 
+        # Register route
+        path = f"/api/v0/fleet/{func.__name__.replace('_', '-')}"
+        routes.append(Route(path, wrapper, methods=["POST"]))
         return wrapper
 
     return decorator
 
 
-@rest_request_response(CreateNodeRequest)
-async def create_node(request: CreateNodeRequest) -> CreateNodeResponse:
-    """Create Node."""
+@rest_request_response(RegisterNodeFleetRequest)
+async def register_node(
+    request: RegisterNodeFleetRequest,
+) -> RegisterNodeFleetResponse:
+    """Register a node (Fleet API only)."""
+    # Get state from app
+    state: LinkState = cast(LinkStateFactory, app.state.STATE_FACTORY).state()
+
+    # Handle message
+    return message_handler.register_node(request=request, state=state)
+
+
+@rest_request_response(ActivateNodeRequest)
+async def activate_node(
+    request: ActivateNodeRequest,
+) -> ActivateNodeResponse:
+    """Activate a node."""
     # Get state from app
     state: LinkState = cast(LinkStateFactory, app.state.STATE_FACTORY).state()
 
     # Handle message
-    return message_handler.create_node(request=request, state=state)
+    return message_handler.activate_node(request=request, state=state)
 
 
-@rest_request_response(DeleteNodeRequest)
-async def delete_node(request: DeleteNodeRequest) -> DeleteNodeResponse:
-    """Delete Node Id."""
+@rest_request_response(DeactivateNodeRequest)
+async def deactivate_node(
+    request: DeactivateNodeRequest,
+) -> DeactivateNodeResponse:
+    """Deactivate a node."""
     # Get state from app
     state: LinkState = cast(LinkStateFactory, app.state.STATE_FACTORY).state()
 
     # Handle message
-    return message_handler.delete_node(request=request, state=state)
+    return message_handler.deactivate_node(request=request, state=state)
+
+
+@rest_request_response(UnregisterNodeFleetRequest)
+async def unregister_node(
+    request: UnregisterNodeFleetRequest,
+) -> UnregisterNodeFleetResponse:
+    """Unregister a node (Fleet API only)."""
+    # Get state from app
+    state: LinkState = cast(LinkStateFactory, app.state.STATE_FACTORY).state()
+
+    # Handle message
+    return message_handler.unregister_node(request=request, state=state)
 
 
 @rest_request_response(PullMessagesRequest)
-async def pull_message(request: PullMessagesRequest) -> PullMessagesResponse:
+async def pull_messages(request: PullMessagesRequest) -> PullMessagesResponse:
     """Pull PullMessages."""
     # Get state from app
     state: LinkState = cast(LinkStateFactory, app.state.STATE_FACTORY).state()
@@ -128,7 +166,7 @@ async def pull_message(request: PullMessagesRequest) -> PullMessagesResponse:
 
 
 @rest_request_response(PushMessagesRequest)
-async def push_message(request: PushMessagesRequest) -> PushMessagesResponse:
+async def push_messages(request: PushMessagesRequest) -> PushMessagesResponse:
     """Pull PushMessages."""
     # Get state from app
     state: LinkState = cast(LinkStateFactory, app.state.STATE_FACTORY).state()
@@ -212,23 +250,6 @@ async def confirm_message_received(
     )
 
 
-routes = [
-    Route("/api/v0/fleet/create-node", create_node, methods=["POST"]),
-    Route("/api/v0/fleet/delete-node", delete_node, methods=["POST"]),
-    Route("/api/v0/fleet/pull-messages", pull_message, methods=["POST"]),
-    Route("/api/v0/fleet/push-messages", push_message, methods=["POST"]),
-    Route("/api/v0/fleet/pull-object", pull_object, methods=["POST"]),
-    Route("/api/v0/fleet/push-object", push_object, methods=["POST"]),
-    Route("/api/v0/fleet/send-node-heartbeat", send_node_heartbeat, methods=["POST"]),
-    Route("/api/v0/fleet/get-run", get_run, methods=["POST"]),
-    Route("/api/v0/fleet/get-fab", get_fab, methods=["POST"]),
-    Route(
-        "/api/v0/fleet/confirm-message-received",
-        confirm_message_received,
-        methods=["POST"],
-    ),
-]
-
 app: Starlette = Starlette(
     debug=False,
     routes=routes,

flwr/server/superlink/fleet/vce/backend/backend.py

@@ -18,7 +18,7 @@
 from abc import ABC, abstractmethod
 from typing import Callable
 
-from flwr.client.client_app import ClientApp
+from flwr.clientapp.client_app import ClientApp
 from flwr.common.context import Context
 from flwr.common.message import Message
 from flwr.common.typing import ConfigRecordValues

flwr/server/superlink/fleet/vce/backend/raybackend.py

@@ -21,7 +21,7 @@ from typing import Callable, Optional, Union
 
 import ray
 
-from flwr.client.client_app import ClientApp
+from flwr.clientapp.client_app import ClientApp
 from flwr.common.constant import PARTITION_ID_KEY
 from flwr.common.context import Context
 from flwr.common.logger import log