flwr-nightly 1.8.0.dev20240314__py3-none-any.whl → 1.11.0.dev20240813__py3-none-any.whl

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

@@ -0,0 +1,214 @@
+# 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 logging import WARNING
+from typing import Any, Callable, Optional, Sequence, Tuple, 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.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.state import State
+
+_PUBLIC_KEY_HEADER = "public-key"
+_AUTH_TOKEN_HEADER = "auth-token"
+
+Request = Union[
+    CreateNodeRequest,
+    DeleteNodeRequest,
+    PullTaskInsRequest,
+    PushTaskResRequest,
+    GetRunRequest,
+    PingRequest,
+]
+
+Response = Union[
+    CreateNodeResponse,
+    DeleteNodeResponse,
+    PullTaskInsResponse,
+    PushTaskResResponse,
+    GetRunResponse,
+    PingResponse,
+]
+
+
+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 client authentication."""
+
+    def __init__(self, state: State):
+        self.state = state
+
+        self.client_public_keys = state.get_client_public_keys()
+        if len(self.client_public_keys) == 0:
+            log(WARNING, "Authentication enabled, but no known public keys configured")
+
+        private_key = self.state.get_server_private_key()
+        public_key = self.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 clients and authenticate clients by validating
+        auth metadata sent by the client. Continue RPC call if client 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:
+            client_public_key_bytes = base64.urlsafe_b64decode(
+                _get_value_from_tuples(
+                    _PUBLIC_KEY_HEADER, context.invocation_metadata()
+                )
+            )
+            if client_public_key_bytes not in self.client_public_keys:
+                context.abort(grpc.StatusCode.UNAUTHENTICATED, "Access denied")
+
+            if isinstance(request, CreateNodeRequest):
+                return self._create_authenticated_node(
+                    client_public_key_bytes, request, context
+                )
+
+            # Verify hmac value
+            hmac_value = base64.urlsafe_b64decode(
+                _get_value_from_tuples(
+                    _AUTH_TOKEN_HEADER, context.invocation_metadata()
+                )
+            )
+            public_key = bytes_to_public_key(client_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.get_node_id(client_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,
+        ],
+    ) -> 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.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)
+        return verify_hmac(shared_secret, request.SerializeToString(True), 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,
+                ),
+            )
+        )
+
+        node_id = self.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:
+            self.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 = self.state.create_node(request.ping_interval, 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,14 +15,18 @@
 """Fleet API message handlers."""
 
 
+import time
 from typing import List, Optional
 from uuid import UUID
 
+from flwr.common.serde import user_config_to_proto
 from flwr.proto.fleet_pb2 import (  # pylint: disable=E0611
     CreateNodeRequest,
     CreateNodeResponse,
     DeleteNodeRequest,
     DeleteNodeResponse,
+    PingRequest,
+    PingResponse,
     PullTaskInsRequest,
     PullTaskInsResponse,
     PushTaskResRequest,
@@ -30,6 +34,11 @@ from flwr.proto.fleet_pb2 import (  # pylint: disable=E0611
     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
 
@@ -40,7 +49,7 @@ def create_node(
 ) -> 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))
 
 
@@ -55,6 +64,15 @@ def delete_node(request: DeleteNodeRequest, state: State) -> DeleteNodeResponse:
     return DeleteNodeResponse()
 
 
+def ping(
+    request: PingRequest,  # pylint: disable=unused-argument
+    state: State,  # 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: State) -> PullTaskInsResponse:
     """Pull TaskIns handler."""
     # Get node_id if client node is not anonymous
@@ -77,6 +95,9 @@ def push_task_res(request: PushTaskResRequest, state: State) -> PushTaskResRespo
     task_res: TaskRes = request.task_res_list[0]
     # pylint: enable=no-member
 
+    # 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 +107,22 @@ def push_task_res(request: PushTaskResRequest, state: State) -> PushTaskResRespo
         results={str(task_id): 0},
     )
     return response
+
+
+def get_run(
+    request: GetRunRequest, state: State  # pylint: disable=W0613
+) -> GetRunResponse:
+    """Get run information."""
+    run = state.get_run(request.run_id)
+
+    if run is None:
+        return GetRunResponse()
+
+    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),
+        )
+    )

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.

flwr/server/superlink/fleet/rest_rere/rest_api.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.
@@ -21,9 +21,11 @@ from flwr.common.constant import MISSING_EXTRA_REST
 from flwr.proto.fleet_pb2 import (  # pylint: disable=E0611
     CreateNodeRequest,
     DeleteNodeRequest,
+    PingRequest,
     PullTaskInsRequest,
     PushTaskResRequest,
 )
+from flwr.proto.run_pb2 import GetRunRequest  # pylint: disable=E0611
 from flwr.server.superlink.fleet.message_handler import message_handler
 from flwr.server.superlink.state import State
 
@@ -152,11 +154,67 @@ async def push_task_res(request: Request) -> Response:  # Check if token is need
     )
 
 
+async def ping(request: Request) -> Response:
+    """Ping."""
+    _check_headers(request.headers)
+
+    # Get the request body as raw bytes
+    ping_request_bytes: bytes = await request.body()
+
+    # Deserialize ProtoBuf
+    ping_request_proto = PingRequest()
+    ping_request_proto.ParseFromString(ping_request_bytes)
+
+    # Get state from app
+    state: State = app.state.STATE_FACTORY.state()
+
+    # Handle message
+    ping_response_proto = message_handler.ping(request=ping_request_proto, state=state)
+
+    # Return serialized ProtoBuf
+    ping_response_bytes = ping_response_proto.SerializeToString()
+    return Response(
+        status_code=200,
+        content=ping_response_bytes,
+        headers={"Content-Type": "application/protobuf"},
+    )
+
+
+async def get_run(request: Request) -> Response:
+    """GetRun."""
+    _check_headers(request.headers)
+
+    # Get the request body as raw bytes
+    get_run_request_bytes: bytes = await request.body()
+
+    # Deserialize ProtoBuf
+    get_run_request_proto = GetRunRequest()
+    get_run_request_proto.ParseFromString(get_run_request_bytes)
+
+    # Get state from app
+    state: State = app.state.STATE_FACTORY.state()
+
+    # Handle message
+    get_run_response_proto = message_handler.get_run(
+        request=get_run_request_proto, state=state
+    )
+
+    # Return serialized ProtoBuf
+    get_run_response_bytes = get_run_response_proto.SerializeToString()
+    return Response(
+        status_code=200,
+        content=get_run_response_bytes,
+        headers={"Content-Type": "application/protobuf"},
+    )
+
+
 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-task-ins", pull_task_ins, methods=["POST"]),
     Route("/api/v0/fleet/push-task-res", push_task_res, methods=["POST"]),
+    Route("/api/v0/fleet/ping", ping, methods=["POST"]),
+    Route("/api/v0/fleet/get-run", get_run, methods=["POST"]),
 ]
 
 app: Starlette = Starlette(

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

@@ -38,7 +38,7 @@ else:
 
     To install the necessary dependencies, install `flwr` with the `simulation` extra:
 
-        pip install -U flwr["simulation"]
+        pip install -U "flwr[simulation]"
     """
 
 

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

@@ -29,12 +29,12 @@ BackendConfig = Dict[str, Dict[str, ConfigsRecordValues]]
 class Backend(ABC):
     """Abstract base class for a Simulation Engine Backend."""
 
-    def __init__(self, backend_config: BackendConfig, work_dir: str) -> None:
+    def __init__(self, backend_config: BackendConfig) -> None:
         """Construct a backend."""
 
     @abstractmethod
-    async def build(self) -> None:
-        """Build backend asynchronously.
+    def build(self) -> None:
+        """Build backend.
 
         Different components need to be in place before workers in a backend are ready
         to accept jobs. When this method finishes executing, the backend should be fully
@@ -54,11 +54,11 @@ class Backend(ABC):
         """Report whether a backend worker is idle and can therefore run a ClientApp."""
 
     @abstractmethod
-    async def terminate(self) -> None:
+    def terminate(self) -> None:
         """Terminate backend."""
 
     @abstractmethod
-    async def process_message(
+    def process_message(
         self,
         app: Callable[[], ClientApp],
         message: Message,

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

@@ -14,26 +14,24 @@
 # ==============================================================================
 """Ray backend for the Fleet API using the Simulation Engine."""
 
-import pathlib
-from logging import ERROR, INFO
-from typing import Callable, Dict, List, Tuple, Union
+from logging import DEBUG, ERROR
+from typing import Callable, Dict, Tuple, Union
 
 import ray
 
-from flwr.client.client_app import ClientApp, LoadClientAppError
+from flwr.client.client_app import ClientApp
+from flwr.common.constant import PARTITION_ID_KEY
 from flwr.common.context import Context
 from flwr.common.logger import log
 from flwr.common.message import Message
-from flwr.simulation.ray_transport.ray_actor import (
-    BasicActorPool,
-    ClientAppActor,
-    init_ray,
-)
+from flwr.common.typing import ConfigsRecordValues
+from flwr.simulation.ray_transport.ray_actor import BasicActorPool, ClientAppActor
 from flwr.simulation.ray_transport.utils import enable_tf_gpu_growth
 
 from .backend import Backend, BackendConfig
 
 ClientResourcesDict = Dict[str, Union[int, float]]
+ActorArgsDict = Dict[str, Union[int, float, Callable[[], None]]]
 
 
 class RayBackend(Backend):
@@ -42,52 +40,28 @@ class RayBackend(Backend):
     def __init__(
         self,
         backend_config: BackendConfig,
-        work_dir: str,
     ) -> None:
         """Prepare RayBackend by initialising Ray and creating the ActorPool."""
-        log(INFO, "Initialising: %s", self.__class__.__name__)
-        log(INFO, "Backend config: %s", backend_config)
+        log(DEBUG, "Initialising: %s", self.__class__.__name__)
+        log(DEBUG, "Backend config: %s", backend_config)
 
-        if not pathlib.Path(work_dir).exists():
-            raise ValueError(f"Specified work_dir {work_dir} does not exist.")
-
-        # Init ray and append working dir if needed
-        runtime_env = (
-            self._configure_runtime_env(work_dir=work_dir) if work_dir else None
-        )
-        init_ray(runtime_env=runtime_env)
+        # Initialise ray
+        self.init_args_key = "init_args"
+        self.init_ray(backend_config)
 
         # Validate client resources
         self.client_resources_key = "client_resources"
+        client_resources = self._validate_client_resources(config=backend_config)
 
         # Create actor pool
-        use_tf = backend_config.get("tensorflow", False)
-        actor_kwargs = {"on_actor_init_fn": enable_tf_gpu_growth} if use_tf else {}
+        actor_kwargs = self._validate_actor_arguments(config=backend_config)
 
-        client_resources = self._validate_client_resources(config=backend_config)
         self.pool = BasicActorPool(
             actor_type=ClientAppActor,
             client_resources=client_resources,
             actor_kwargs=actor_kwargs,
         )
 
-    def _configure_runtime_env(self, work_dir: str) -> Dict[str, Union[str, List[str]]]:
-        """Return list of files/subdirectories to exclude relative to work_dir.
-
-        Without this, Ray will push everything to the Ray Cluster.
-        """
-        runtime_env: Dict[str, Union[str, List[str]]] = {"working_dir": work_dir}
-
-        excludes = []
-        path = pathlib.Path(work_dir)
-        for p in path.rglob("*"):
-            # Exclude files need to be relative to the working_dir
-            if p.is_file() and not str(p).endswith(".py"):
-                excludes.append(str(p.relative_to(path)))
-        runtime_env["excludes"] = excludes
-
-        return runtime_env
-
     def _validate_client_resources(self, config: BackendConfig) -> ClientResourcesDict:
         client_resources_config = config.get(self.client_resources_key)
         client_resources: ClientResourcesDict = {}
@@ -109,7 +83,7 @@ class RayBackend(Backend):
         else:
             client_resources = {"num_cpus": 2, "num_gpus": 0.0}
             log(
-                INFO,
+                DEBUG,
                 "`%s` not specified in backend config. Applying default setting: %s",
                 self.client_resources_key,
                 client_resources,
@@ -117,6 +91,29 @@ class RayBackend(Backend):
 
         return client_resources
 
+    def _validate_actor_arguments(self, config: BackendConfig) -> ActorArgsDict:
+        actor_args_config = config.get("actor", False)
+        actor_args: ActorArgsDict = {}
+        if actor_args_config:
+            use_tf = actor_args.get("tensorflow", False)
+            if use_tf:
+                actor_args["on_actor_init_fn"] = enable_tf_gpu_growth
+        return actor_args
+
+    def init_ray(self, backend_config: BackendConfig) -> None:
+        """Intialises Ray if not already initialised."""
+        if not ray.is_initialized():
+            ray_init_args: Dict[
+                str,
+                ConfigsRecordValues,
+            ] = {}
+
+            if backend_config.get(self.init_args_key):
+                for k, v in backend_config[self.init_args_key].items():
+                    ray_init_args[k] = v
+
+            ray.init(**ray_init_args)
+
     @property
     def num_workers(self) -> int:
         """Return number of actors in pool."""
@@ -126,12 +123,12 @@ class RayBackend(Backend):
         """Report whether the pool has idle actors."""
         return self.pool.is_actor_available()
 
-    async def build(self) -> None:
+    def build(self) -> None:
         """Build pool of Ray actors that this backend will submit jobs to."""
-        await self.pool.add_actors_to_pool(self.pool.actors_capacity)
-        log(INFO, "Constructed ActorPool with: %i actors", self.pool.num_actors)
+        self.pool.add_actors_to_pool(self.pool.actors_capacity)
+        log(DEBUG, "Constructed ActorPool with: %i actors", self.pool.num_actors)
 
-    async def process_message(
+    def process_message(
         self,
         app: Callable[[], ClientApp],
         message: Message,
@@ -141,35 +138,35 @@ class RayBackend(Backend):
 
         Return output message and updated context.
         """
-        partition_id = message.metadata.partition_id
+        partition_id = context.node_config[PARTITION_ID_KEY]
 
         try:
-            # Submite a task to the pool
-            future = await self.pool.submit(
+            # Submit a task to the pool
+            future = self.pool.submit(
                 lambda a, a_fn, mssg, cid, state: a.run.remote(a_fn, mssg, cid, state),
                 (app, message, str(partition_id), context),
             )
 
-            await future
-
             # Fetch result
             (
                 out_mssg,
                 updated_context,
-            ) = await self.pool.fetch_result_and_return_actor_to_pool(future)
+            ) = self.pool.fetch_result_and_return_actor_to_pool(future)
 
             return out_mssg, updated_context
 
-        except LoadClientAppError as load_ex:
+        except Exception as ex:
             log(
                 ERROR,
                 "An exception was raised when processing a message by %s",
                 self.__class__.__name__,
             )
-            raise load_ex
+            # add actor back into pool
+            self.pool.add_actor_back_to_pool(future)
+            raise ex
 
-    async def terminate(self) -> None:
+    def terminate(self) -> None:
         """Terminate all actors in actor pool."""
-        await self.pool.terminate_all_actors()
+        self.pool.terminate_all_actors()
         ray.shutdown()
-        log(INFO, "Terminated %s", self.__class__.__name__)
+        log(DEBUG, "Terminated %s", self.__class__.__name__)