flwr-nightly 1.8.0.dev20240315__py3-none-any.whl → 1.15.0.dev20250114__py3-none-any.whl

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

@@ -1,4 +1,4 @@
-# Copyright 2020 Flower Labs GmbH. All Rights Reserved.
+# Copyright 2024 Flower Labs GmbH. All Rights Reserved.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
@@ -15,68 +15,175 @@
 """Fleet API gRPC request-response servicer."""
 
 
-from logging import INFO
+from logging import DEBUG, INFO
 
 import grpc
 
 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,
+    PingRequest,
+    PingResponse,
+    PullMessagesRequest,
+    PullMessagesResponse,
     PullTaskInsRequest,
     PullTaskInsResponse,
+    PushMessagesRequest,
+    PushMessagesResponse,
     PushTaskResRequest,
     PushTaskResResponse,
 )
+from flwr.proto.run_pb2 import GetRunRequest, GetRunResponse  # pylint: disable=E0611
+from flwr.server.superlink.ffs.ffs_factory import FfsFactory
 from flwr.server.superlink.fleet.message_handler import message_handler
-from flwr.server.superlink.state import StateFactory
+from flwr.server.superlink.linkstate import LinkStateFactory
+from flwr.server.superlink.utils import abort_grpc_context
 
 
 class FleetServicer(fleet_pb2_grpc.FleetServicer):
     """Fleet API servicer."""
 
-    def __init__(self, state_factory: StateFactory) -> None:
+    def __init__(
+        self, state_factory: LinkStateFactory, ffs_factory: FfsFactory
+    ) -> None:
         self.state_factory = state_factory
+        self.ffs_factory = ffs_factory
 
     def CreateNode(
         self, request: CreateNodeRequest, context: grpc.ServicerContext
     ) -> CreateNodeResponse:
         """."""
-        log(INFO, "FleetServicer.CreateNode")
-        return message_handler.create_node(
+        log(INFO, "[Fleet.CreateNode] Request ping_interval=%s", request.ping_interval)
+        log(DEBUG, "[Fleet.CreateNode] Request: %s", 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", response)
+        return response
 
     def DeleteNode(
         self, request: DeleteNodeRequest, context: grpc.ServicerContext
     ) -> DeleteNodeResponse:
         """."""
-        log(INFO, "FleetServicer.DeleteNode")
+        log(INFO, "[Fleet.DeleteNode] Delete node_id=%s", request.node.node_id)
+        log(DEBUG, "[Fleet.DeleteNode] Request: %s", request)
         return message_handler.delete_node(
             request=request,
             state=self.state_factory.state(),
         )
 
+    def Ping(self, request: PingRequest, context: grpc.ServicerContext) -> PingResponse:
+        """."""
+        log(DEBUG, "[Fleet.Ping] Request: %s", request)
+        return message_handler.ping(
+            request=request,
+            state=self.state_factory.state(),
+        )
+
     def PullTaskIns(
         self, request: PullTaskInsRequest, context: grpc.ServicerContext
     ) -> PullTaskInsResponse:
         """Pull TaskIns."""
-        log(INFO, "FleetServicer.PullTaskIns")
+        log(INFO, "[Fleet.PullTaskIns] node_id=%s", request.node.node_id)
+        log(DEBUG, "[Fleet.PullTaskIns] Request: %s", request)
         return message_handler.pull_task_ins(
             request=request,
             state=self.state_factory.state(),
         )
 
+    def PullMessages(
+        self, request: PullMessagesRequest, context: grpc.ServicerContext
+    ) -> PullMessagesResponse:
+        """Pull Messages."""
+        log(INFO, "[Fleet.PullMessages] node_id=%s", request.node.node_id)
+        log(DEBUG, "[Fleet.PullMessages] Request: %s", request)
+        return message_handler.pull_messages(
+            request=request,
+            state=self.state_factory.state(),
+        )
+
     def PushTaskRes(
         self, request: PushTaskResRequest, context: grpc.ServicerContext
     ) -> PushTaskResResponse:
         """Push TaskRes."""
-        log(INFO, "FleetServicer.PushTaskRes")
-        return message_handler.push_task_res(
-            request=request,
-            state=self.state_factory.state(),
-        )
+        if request.task_res_list:
+            log(
+                INFO,
+                "[Fleet.PushTaskRes] Push results from node_id=%s",
+                request.task_res_list[0].task.producer.node_id,
+            )
+        else:
+            log(INFO, "[Fleet.PushTaskRes] No task results to push")
+
+        try:
+            res = message_handler.push_task_res(
+                request=request,
+                state=self.state_factory.state(),
+            )
+        except InvalidRunStatusException as e:
+            abort_grpc_context(e.message, context)
+
+        return res
+
+    def PushMessages(
+        self, request: PushMessagesRequest, context: grpc.ServicerContext
+    ) -> PushMessagesResponse:
+        """Push Messages."""
+        if request.messages_list:
+            log(
+                INFO,
+                "[Fleet.PushMessages] Push results from node_id=%s",
+                request.messages_list[0].metadata.src_node_id,
+            )
+        else:
+            log(INFO, "[Fleet.PushMessages] No task results to push")
+
+        try:
+            res = message_handler.push_messages(
+                request=request,
+                state=self.state_factory.state(),
+            )
+        except InvalidRunStatusException as e:
+            abort_grpc_context(e.message, context)
+
+        return res
+
+    def GetRun(
+        self, request: GetRunRequest, context: grpc.ServicerContext
+    ) -> GetRunResponse:
+        """Get run information."""
+        log(INFO, "[Fleet.GetRun] Requesting `Run` for run_id=%s", request.run_id)
+
+        try:
+            res = message_handler.get_run(
+                request=request,
+                state=self.state_factory.state(),
+            )
+        except InvalidRunStatusException as e:
+            abort_grpc_context(e.message, context)
+
+        return res
+
+    def GetFab(
+        self, request: GetFabRequest, context: grpc.ServicerContext
+    ) -> GetFabResponse:
+        """Get FAB."""
+        log(INFO, "[Fleet.GetFab] Requesting FAB for fab_hash=%s", request.hash_str)
+        try:
+            res = message_handler.get_fab(
+                request=request,
+                ffs=self.ffs_factory.ffs(),
+                state=self.state_factory.state(),
+            )
+        except InvalidRunStatusException as e:
+            abort_grpc_context(e.message, context)
+
+        return res

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

@@ -0,0 +1,228 @@
+# Copyright 2024 Flower Labs GmbH. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Flower server interceptor."""
+
+
+import base64
+from collections.abc import Sequence
+from logging import INFO, WARNING
+from typing import Any, Callable, Optional, Union
+
+import grpc
+from cryptography.hazmat.primitives.asymmetric import ec
+
+from flwr.common.logger import log
+from flwr.common.secure_aggregation.crypto.symmetric_encryption import (
+    bytes_to_private_key,
+    bytes_to_public_key,
+    generate_shared_key,
+    verify_hmac,
+)
+from flwr.proto.fab_pb2 import GetFabRequest, GetFabResponse  # pylint: disable=E0611
+from flwr.proto.fleet_pb2 import (  # pylint: disable=E0611
+    CreateNodeRequest,
+    CreateNodeResponse,
+    DeleteNodeRequest,
+    DeleteNodeResponse,
+    PingRequest,
+    PingResponse,
+    PullTaskInsRequest,
+    PullTaskInsResponse,
+    PushTaskResRequest,
+    PushTaskResResponse,
+)
+from flwr.proto.node_pb2 import Node  # pylint: disable=E0611
+from flwr.proto.run_pb2 import GetRunRequest, GetRunResponse  # pylint: disable=E0611
+from flwr.server.superlink.linkstate import LinkStateFactory
+
+_PUBLIC_KEY_HEADER = "public-key"
+_AUTH_TOKEN_HEADER = "auth-token"
+
+Request = Union[
+    CreateNodeRequest,
+    DeleteNodeRequest,
+    PullTaskInsRequest,
+    PushTaskResRequest,
+    GetRunRequest,
+    PingRequest,
+    GetFabRequest,
+]
+
+Response = Union[
+    CreateNodeResponse,
+    DeleteNodeResponse,
+    PullTaskInsResponse,
+    PushTaskResResponse,
+    GetRunResponse,
+    PingResponse,
+    GetFabResponse,
+]
+
+
+def _get_value_from_tuples(
+    key_string: str, tuples: Sequence[tuple[str, Union[str, bytes]]]
+) -> bytes:
+    value = next((value for key, value in tuples if key == key_string), "")
+    if isinstance(value, str):
+        return value.encode()
+
+    return value
+
+
+class AuthenticateServerInterceptor(grpc.ServerInterceptor):  # type: ignore
+    """Server interceptor for node authentication."""
+
+    def __init__(self, state_factory: LinkStateFactory):
+        self.state_factory = state_factory
+        state = self.state_factory.state()
+
+        self.node_public_keys = state.get_node_public_keys()
+        if len(self.node_public_keys) == 0:
+            log(WARNING, "Authentication enabled, but no known public keys configured")
+
+        private_key = state.get_server_private_key()
+        public_key = state.get_server_public_key()
+
+        if private_key is None or public_key is None:
+            raise ValueError("Error loading authentication keys")
+
+        self.server_private_key = bytes_to_private_key(private_key)
+        self.encoded_server_public_key = base64.urlsafe_b64encode(public_key)
+
+    def intercept_service(
+        self,
+        continuation: Callable[[Any], Any],
+        handler_call_details: grpc.HandlerCallDetails,
+    ) -> grpc.RpcMethodHandler:
+        """Flower server interceptor authentication logic.
+
+        Intercept all unary calls from nodes and authenticate nodes by validating auth
+        metadata sent by the node. Continue RPC call if node is authenticated, else,
+        terminate RPC call by setting context to abort.
+        """
+        # 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._generic_auth_unary_method_handler(method_handler)
+
+    def _generic_auth_unary_method_handler(
+        self, method_handler: grpc.RpcMethodHandler
+    ) -> grpc.RpcMethodHandler:
+        def _generic_method_handler(
+            request: Request,
+            context: grpc.ServicerContext,
+        ) -> Response:
+            node_public_key_bytes = base64.urlsafe_b64decode(
+                _get_value_from_tuples(
+                    _PUBLIC_KEY_HEADER, context.invocation_metadata()
+                )
+            )
+            if node_public_key_bytes not in self.node_public_keys:
+                context.abort(grpc.StatusCode.UNAUTHENTICATED, "Access denied")
+
+            if isinstance(request, CreateNodeRequest):
+                response = self._create_authenticated_node(
+                    node_public_key_bytes, request, context
+                )
+                log(
+                    INFO,
+                    "AuthenticateServerInterceptor: Created node_id=%s",
+                    response.node.node_id,
+                )
+                return response
+
+            # Verify hmac value
+            hmac_value = base64.urlsafe_b64decode(
+                _get_value_from_tuples(
+                    _AUTH_TOKEN_HEADER, context.invocation_metadata()
+                )
+            )
+            public_key = bytes_to_public_key(node_public_key_bytes)
+
+            if not self._verify_hmac(public_key, request, hmac_value):
+                context.abort(grpc.StatusCode.UNAUTHENTICATED, "Access denied")
+
+            # Verify node_id
+            node_id = self.state_factory.state().get_node_id(node_public_key_bytes)
+
+            if not self._verify_node_id(node_id, request):
+                context.abort(grpc.StatusCode.UNAUTHENTICATED, "Access denied")
+
+            return method_handler.unary_unary(request, context)  # type: ignore
+
+        return grpc.unary_unary_rpc_method_handler(
+            _generic_method_handler,
+            request_deserializer=method_handler.request_deserializer,
+            response_serializer=method_handler.response_serializer,
+        )
+
+    def _verify_node_id(
+        self,
+        node_id: Optional[int],
+        request: Union[
+            DeleteNodeRequest,
+            PullTaskInsRequest,
+            PushTaskResRequest,
+            GetRunRequest,
+            PingRequest,
+            GetFabRequest,
+        ],
+    ) -> bool:
+        if node_id is None:
+            return False
+        if isinstance(request, PushTaskResRequest):
+            if len(request.task_res_list) == 0:
+                return False
+            return request.task_res_list[0].task.producer.node_id == node_id
+        if isinstance(request, GetRunRequest):
+            return node_id in self.state_factory.state().get_nodes(request.run_id)
+        return request.node.node_id == node_id
+
+    def _verify_hmac(
+        self, public_key: ec.EllipticCurvePublicKey, request: Request, hmac_value: bytes
+    ) -> bool:
+        shared_secret = generate_shared_key(self.server_private_key, public_key)
+        message_bytes = request.SerializeToString(deterministic=True)
+        return verify_hmac(shared_secret, message_bytes, hmac_value)
+
+    def _create_authenticated_node(
+        self,
+        public_key_bytes: bytes,
+        request: CreateNodeRequest,
+        context: grpc.ServicerContext,
+    ) -> CreateNodeResponse:
+        context.send_initial_metadata(
+            (
+                (
+                    _PUBLIC_KEY_HEADER,
+                    self.encoded_server_public_key,
+                ),
+            )
+        )
+        state = self.state_factory.state()
+        node_id = state.get_node_id(public_key_bytes)
+
+        # Handle `CreateNode` here instead of calling the default method handler
+        # Return previously assigned `node_id` for the provided `public_key`
+        if node_id is not None:
+            state.acknowledge_ping(node_id, request.ping_interval)
+            return CreateNodeResponse(node=Node(node_id=node_id, anonymous=False))
+
+        # 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)
+        state.set_node_public_key(node_id, public_key_bytes)
+        return CreateNodeResponse(node=Node(node_id=node_id, anonymous=False))

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

@@ -1,4 +1,4 @@
-# Copyright 2020 Flower Labs GmbH. All Rights Reserved.
+# Copyright 2024 Flower Labs GmbH. All Rights Reserved.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

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

@@ -1,4 +1,4 @@
-# Copyright 2020 Flower Labs GmbH. All Rights Reserved.
+# Copyright 2024 Flower Labs GmbH. All Rights Reserved.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
@@ -15,36 +15,61 @@
 """Fleet API message handlers."""
 
 
-from typing import List, Optional
+import time
+from typing import Optional
 from uuid import UUID
 
+from flwr.common.constant import Status
+from flwr.common.serde import (
+    fab_to_proto,
+    message_from_proto,
+    message_from_taskins,
+    message_to_proto,
+    message_to_taskres,
+    user_config_to_proto,
+)
+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,
+    PingRequest,
+    PingResponse,
+    PullMessagesRequest,
+    PullMessagesResponse,
     PullTaskInsRequest,
     PullTaskInsResponse,
+    PushMessagesRequest,
+    PushMessagesResponse,
     PushTaskResRequest,
     PushTaskResResponse,
     Reconnect,
 )
 from flwr.proto.node_pb2 import Node  # pylint: disable=E0611
+from flwr.proto.run_pb2 import (  # pylint: disable=E0611
+    GetRunRequest,
+    GetRunResponse,
+    Run,
+)
 from flwr.proto.task_pb2 import TaskIns, TaskRes  # pylint: disable=E0611
-from flwr.server.superlink.state import State
+from flwr.server.superlink.ffs.ffs import Ffs
+from flwr.server.superlink.linkstate import LinkState
+from flwr.server.superlink.utils import check_abort
 
 
 def create_node(
     request: CreateNodeRequest,  # pylint: disable=unused-argument
-    state: State,
+    state: LinkState,
 ) -> CreateNodeResponse:
     """."""
     # Create node
-    node_id = state.create_node()
+    node_id = state.create_node(ping_interval=request.ping_interval)
     return CreateNodeResponse(node=Node(node_id=node_id, anonymous=False))
 
 
-def delete_node(request: DeleteNodeRequest, state: State) -> DeleteNodeResponse:
+def delete_node(request: DeleteNodeRequest, state: LinkState) -> DeleteNodeResponse:
     """."""
     # Validate node_id
     if request.node.anonymous or request.node.node_id == 0:
@@ -55,14 +80,23 @@ def delete_node(request: DeleteNodeRequest, state: State) -> DeleteNodeResponse:
     return DeleteNodeResponse()
 
 
-def pull_task_ins(request: PullTaskInsRequest, state: State) -> PullTaskInsResponse:
+def ping(
+    request: PingRequest,  # pylint: disable=unused-argument
+    state: LinkState,  # pylint: disable=unused-argument
+) -> PingResponse:
+    """."""
+    res = state.acknowledge_ping(request.node.node_id, request.ping_interval)
+    return PingResponse(success=res)
+
+
+def pull_task_ins(request: PullTaskInsRequest, state: LinkState) -> PullTaskInsResponse:
     """Pull TaskIns handler."""
     # Get node_id if client node is not anonymous
     node = request.node  # pylint: disable=no-member
     node_id: Optional[int] = None if node.anonymous else node.node_id
 
     # Retrieve TaskIns from State
-    task_ins_list: List[TaskIns] = state.get_task_ins(node_id=node_id, limit=1)
+    task_ins_list: list[TaskIns] = state.get_task_ins(node_id=node_id, limit=1)
 
     # Build response
     response = PullTaskInsResponse(
@@ -71,12 +105,44 @@ def pull_task_ins(request: PullTaskInsRequest, state: State) -> PullTaskInsRespo
     return response
 
 
-def push_task_res(request: PushTaskResRequest, state: State) -> PushTaskResResponse:
+def pull_messages(
+    request: PullMessagesRequest, state: LinkState
+) -> PullMessagesResponse:
+    """Pull Messages handler."""
+    # Get node_id if client node is not anonymous
+    node = request.node  # pylint: disable=no-member
+    node_id: Optional[int] = None if node.anonymous else node.node_id
+
+    # Retrieve TaskIns from State
+    task_ins_list: list[TaskIns] = state.get_task_ins(node_id=node_id, limit=1)
+
+    # Convert to Messages
+    msg_proto = []
+    for task_ins in task_ins_list:
+        msg = message_from_taskins(task_ins)
+        msg_proto.append(message_to_proto(msg))
+
+    return PullMessagesResponse(messages_list=msg_proto)
+
+
+def push_task_res(request: PushTaskResRequest, state: LinkState) -> PushTaskResResponse:
     """Push TaskRes handler."""
     # pylint: disable=no-member
     task_res: TaskRes = request.task_res_list[0]
     # pylint: enable=no-member
 
+    # Abort if the run is not running
+    abort_msg = check_abort(
+        task_res.run_id,
+        [Status.PENDING, Status.STARTING, Status.FINISHED],
+        state,
+    )
+    if abort_msg:
+        raise InvalidRunStatusException(abort_msg)
+
+    # Set pushed_at (timestamp in seconds)
+    task_res.task.pushed_at = time.time()
+
     # Store TaskRes in State
     task_id: Optional[UUID] = state.store_task_res(task_res=task_res)
 
@@ -86,3 +152,81 @@ def push_task_res(request: PushTaskResRequest, state: State) -> PushTaskResRespo
         results={str(task_id): 0},
     )
     return response
+
+
+def push_messages(
+    request: PushMessagesRequest, state: LinkState
+) -> PushMessagesResponse:
+    """Push Messages handler."""
+    # Convert Message to TaskRes
+    msg = message_from_proto(message_proto=request.messages_list[0])
+    task_res = message_to_taskres(msg)
+
+    # Abort if the run is not running
+    abort_msg = check_abort(
+        task_res.run_id,
+        [Status.PENDING, Status.STARTING, Status.FINISHED],
+        state,
+    )
+    if abort_msg:
+        raise InvalidRunStatusException(abort_msg)
+
+    # Set pushed_at (timestamp in seconds)
+    task_res.task.pushed_at = time.time()
+
+    # Store TaskRes in State
+    message_id: Optional[UUID] = state.store_task_res(task_res=task_res)
+
+    # Build response
+    response = PushMessagesResponse(
+        reconnect=Reconnect(reconnect=5),
+        results={str(message_id): 0},
+    )
+    return response
+
+
+def get_run(request: GetRunRequest, state: LinkState) -> GetRunResponse:
+    """Get run information."""
+    run = state.get_run(request.run_id)
+
+    if run is None:
+        return GetRunResponse()
+
+    # Abort if the run is not running
+    abort_msg = check_abort(
+        request.run_id,
+        [Status.PENDING, Status.STARTING, Status.FINISHED],
+        state,
+    )
+    if abort_msg:
+        raise InvalidRunStatusException(abort_msg)
+
+    return GetRunResponse(
+        run=Run(
+            run_id=run.run_id,
+            fab_id=run.fab_id,
+            fab_version=run.fab_version,
+            override_config=user_config_to_proto(run.override_config),
+            fab_hash=run.fab_hash,
+        )
+    )
+
+
+def get_fab(
+    request: GetFabRequest, ffs: Ffs, state: LinkState  # pylint: disable=W0613
+) -> GetFabResponse:
+    """Get FAB."""
+    # Abort if the run is not running
+    abort_msg = check_abort(
+        request.run_id,
+        [Status.PENDING, Status.STARTING, Status.FINISHED],
+        state,
+    )
+    if abort_msg:
+        raise InvalidRunStatusException(abort_msg)
+
+    if result := ffs.get(request.hash_str):
+        fab = Fab(request.hash_str, result[0])
+        return GetFabResponse(fab=fab_to_proto(fab))
+
+    raise ValueError(f"Found no FAB with hash: {request.hash_str}")

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

@@ -1,4 +1,4 @@
-# Copyright 2020 Flower Labs GmbH. All Rights Reserved.
+# Copyright 2024 Flower Labs GmbH. All Rights Reserved.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.