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

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.
@@ -15,17 +15,37 @@
 """Experimental REST API server."""
 
 
+from __future__ import annotations
+
 import sys
+from collections.abc import Awaitable
+from typing import Callable, TypeVar, cast
+
+from google.protobuf.message import Message as GrpcMessage
 
 from flwr.common.constant import MISSING_EXTRA_REST
+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 import Ffs
+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 State
+from flwr.server.superlink.linkstate import LinkState, LinkStateFactory
 
 try:
     from starlette.applications import Starlette
@@ -38,125 +58,143 @@ except ModuleNotFoundError:
     sys.exit(MISSING_EXTRA_REST)
 
 
-async def create_node(request: Request) -> Response:
-    """Create Node."""
-    _check_headers(request.headers)
+GrpcRequest = TypeVar("GrpcRequest", bound=GrpcMessage)
+GrpcResponse = TypeVar("GrpcResponse", bound=GrpcMessage)
 
-    # Get the request body as raw bytes
-    create_node_request_bytes: bytes = await request.body()
+GrpcAsyncFunction = Callable[[GrpcRequest], Awaitable[GrpcResponse]]
+RestEndPoint = Callable[[Request], Awaitable[Response]]
 
-    # Deserialize ProtoBuf
-    create_node_request_proto = CreateNodeRequest()
-    create_node_request_proto.ParseFromString(create_node_request_bytes)
 
-    # Get state from app
-    state: State = app.state.STATE_FACTORY.state()
+def rest_request_response(
+    grpc_request_type: type[GrpcRequest],
+) -> Callable[[GrpcAsyncFunction[GrpcRequest, GrpcResponse]], RestEndPoint]:
+    """Convert an async gRPC-based function into a RESTful HTTP endpoint."""
 
-    # Handle message
-    create_node_response_proto = message_handler.create_node(
-        request=create_node_request_proto, state=state
-    )
+    def decorator(func: GrpcAsyncFunction[GrpcRequest, GrpcResponse]) -> RestEndPoint:
+        async def wrapper(request: Request) -> Response:
+            _check_headers(request.headers)
 
-    # Return serialized ProtoBuf
-    create_node_response_bytes = create_node_response_proto.SerializeToString()
-    return Response(
-        status_code=200,
-        content=create_node_response_bytes,
-        headers={"Content-Type": "application/protobuf"},
-    )
+            # Get the request body as raw bytes
+            grpc_req_bytes: bytes = await request.body()
 
+            # Deserialize ProtoBuf
+            grpc_req = grpc_request_type.FromString(grpc_req_bytes)
+            grpc_res = await func(grpc_req)
+            return Response(
+                status_code=200,
+                content=grpc_res.SerializeToString(),
+                headers={"Content-Type": "application/protobuf"},
+            )
 
-async def delete_node(request: Request) -> Response:
-    """Delete Node Id."""
-    _check_headers(request.headers)
+        return wrapper
 
-    # Get the request body as raw bytes
-    delete_node_request_bytes: bytes = await request.body()
+    return decorator
 
-    # Deserialize ProtoBuf
-    delete_node_request_proto = DeleteNodeRequest()
-    delete_node_request_proto.ParseFromString(delete_node_request_bytes)
 
+@rest_request_response(CreateNodeRequest)
+async def create_node(request: CreateNodeRequest) -> CreateNodeResponse:
+    """Create Node."""
     # Get state from app
-    state: State = app.state.STATE_FACTORY.state()
+    state: LinkState = cast(LinkStateFactory, app.state.STATE_FACTORY).state()
 
     # Handle message
-    delete_node_response_proto = message_handler.delete_node(
-        request=delete_node_request_proto, state=state
-    )
+    return message_handler.create_node(request=request, state=state)
+
 
-    # Return serialized ProtoBuf
-    delete_node_response_bytes = delete_node_response_proto.SerializeToString()
-    return Response(
-        status_code=200,
-        content=delete_node_response_bytes,
-        headers={"Content-Type": "application/protobuf"},
-    )
+@rest_request_response(DeleteNodeRequest)
+async def delete_node(request: DeleteNodeRequest) -> DeleteNodeResponse:
+    """Delete Node Id."""
+    # Get state from app
+    state: LinkState = cast(LinkStateFactory, app.state.STATE_FACTORY).state()
+
+    # Handle message
+    return message_handler.delete_node(request=request, state=state)
 
 
-async def pull_task_ins(request: Request) -> Response:
+@rest_request_response(PullTaskInsRequest)
+async def pull_task_ins(request: PullTaskInsRequest) -> PullTaskInsResponse:
     """Pull TaskIns."""
-    _check_headers(request.headers)
+    # Get state from app
+    state: LinkState = cast(LinkStateFactory, app.state.STATE_FACTORY).state()
 
-    # Get the request body as raw bytes
-    pull_task_ins_request_bytes: bytes = await request.body()
+    # Handle message
+    return message_handler.pull_task_ins(request=request, state=state)
 
-    # Deserialize ProtoBuf
-    pull_task_ins_request_proto = PullTaskInsRequest()
-    pull_task_ins_request_proto.ParseFromString(pull_task_ins_request_bytes)
 
+@rest_request_response(PullMessagesRequest)
+async def pull_message(request: PullMessagesRequest) -> PullMessagesResponse:
+    """Pull PullMessages."""
     # Get state from app
-    state: State = app.state.STATE_FACTORY.state()
+    state: LinkState = cast(LinkStateFactory, app.state.STATE_FACTORY).state()
 
     # Handle message
-    pull_task_ins_response_proto = message_handler.pull_task_ins(
-        request=pull_task_ins_request_proto,
-        state=state,
-    )
+    return message_handler.pull_messages(request=request, state=state)
 
-    # Return serialized ProtoBuf
-    pull_task_ins_response_bytes = pull_task_ins_response_proto.SerializeToString()
-    return Response(
-        status_code=200,
-        content=pull_task_ins_response_bytes,
-        headers={"Content-Type": "application/protobuf"},
-    )
 
-
-async def push_task_res(request: Request) -> Response:  # Check if token is needed here
+# Check if token is needed here
+@rest_request_response(PushTaskResRequest)
+async def push_task_res(request: PushTaskResRequest) -> PushTaskResResponse:
     """Push TaskRes."""
-    _check_headers(request.headers)
+    # Get state from app
+    state: LinkState = cast(LinkStateFactory, app.state.STATE_FACTORY).state()
 
-    # Get the request body as raw bytes
-    push_task_res_request_bytes: bytes = await request.body()
+    # Handle message
+    return message_handler.push_task_res(request=request, state=state)
+
+
+@rest_request_response(PushMessagesRequest)
+async def push_message(request: PushMessagesRequest) -> PushMessagesResponse:
+    """Pull PushMessages."""
+    # Get state from app
+    state: LinkState = cast(LinkStateFactory, app.state.STATE_FACTORY).state()
+
+    # Handle message
+    return message_handler.push_messages(request=request, state=state)
+
+
+@rest_request_response(PingRequest)
+async def ping(request: PingRequest) -> PingResponse:
+    """Ping."""
+    # Get state from app
+    state: LinkState = cast(LinkStateFactory, app.state.STATE_FACTORY).state()
+
+    # Handle message
+    return message_handler.ping(request=request, state=state)
 
-    # Deserialize ProtoBuf
-    push_task_res_request_proto = PushTaskResRequest()
-    push_task_res_request_proto.ParseFromString(push_task_res_request_bytes)
 
+@rest_request_response(GetRunRequest)
+async def get_run(request: GetRunRequest) -> GetRunResponse:
+    """GetRun."""
     # Get state from app
-    state: State = app.state.STATE_FACTORY.state()
+    state: LinkState = cast(LinkStateFactory, app.state.STATE_FACTORY).state()
 
     # Handle message
-    push_task_res_response_proto = message_handler.push_task_res(
-        request=push_task_res_request_proto,
-        state=state,
-    )
+    return message_handler.get_run(request=request, state=state)
+
 
-    # Return serialized ProtoBuf
-    push_task_res_response_bytes = push_task_res_response_proto.SerializeToString()
-    return Response(
-        status_code=200,
-        content=push_task_res_response_bytes,
-        headers={"Content-Type": "application/protobuf"},
-    )
+@rest_request_response(GetFabRequest)
+async def get_fab(request: GetFabRequest) -> GetFabResponse:
+    """GetRun."""
+    # Get ffs from app
+    ffs: Ffs = cast(FfsFactory, app.state.FFS_FACTORY).ffs()
+
+    # Get state from app
+    state: LinkState = cast(LinkStateFactory, app.state.STATE_FACTORY).state()
+
+    # Handle message
+    return message_handler.get_fab(request=request, ffs=ffs, state=state)
 
 
 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/pull-messages", pull_message, methods=["POST"]),
     Route("/api/v0/fleet/push-task-res", push_task_res, methods=["POST"]),
+    Route("/api/v0/fleet/push-messages", push_message, methods=["POST"]),
+    Route("/api/v0/fleet/ping", ping, methods=["POST"]),
+    Route("/api/v0/fleet/get-run", get_run, methods=["POST"]),
+    Route("/api/v0/fleet/get-fab", get_fab, methods=["POST"]),
 ]
 
 app: Starlette = Starlette(

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

@@ -14,6 +14,7 @@
 # ==============================================================================
 """Fleet Simulation Engine side."""
 
+
 from .vce_api import start_vce
 
 __all__ = [

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

@@ -14,18 +14,18 @@
 # ==============================================================================
 """Simulation Engine Backends."""
 
+
 import importlib
-from typing import Dict, Type
 
 from .backend import Backend, BackendConfig
 
 is_ray_installed = importlib.util.find_spec("ray") is not None
 
 # Mapping of supported backends
-supported_backends: Dict[str, Type[Backend]] = {}
+supported_backends: dict[str, type[Backend]] = {}
 
 # To log backend-specific error message when chosen backend isn't available
-error_messages_backends: Dict[str, str] = {}
+error_messages_backends: dict[str, str] = {}
 
 if is_ray_installed:
     from .raybackend import RayBackend
@@ -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

@@ -16,25 +16,25 @@
 
 
 from abc import ABC, abstractmethod
-from typing import Callable, Dict, Tuple
+from typing import Callable
 
 from flwr.client.client_app import ClientApp
 from flwr.common.context import Context
 from flwr.common.message import Message
 from flwr.common.typing import ConfigsRecordValues
 
-BackendConfig = Dict[str, Dict[str, ConfigsRecordValues]]
+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, app_fn: Callable[[], ClientApp]) -> 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,14 +54,13 @@ 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,
         context: Context,
-    ) -> Tuple[Message, Context]:
+    ) -> tuple[Message, Context]:
         """Submit a job to the backend."""

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

@@ -14,26 +14,26 @@
 # ==============================================================================
 """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
+
+import sys
+from logging import DEBUG, ERROR
+from typing import Callable, Optional, 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]]
+ClientResourcesDict = dict[str, Union[int, float]]
+ActorArgsDict = dict[str, Union[int, float, Callable[[], None]]]
 
 
 class RayBackend(Backend):
@@ -42,51 +42,24 @@ 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)
-
-        if not pathlib.Path(work_dir).exists():
-            raise ValueError(f"Specified work_dir {work_dir} does not exist.")
+        log(DEBUG, "Initialising: %s", self.__class__.__name__)
+        log(DEBUG, "Backend config: %s", backend_config)
 
-        # 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"
+        self.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 {}
-
-        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.
+        # Valide actor resources
+        self.actor_kwargs = self._validate_actor_arguments(config=backend_config)
+        self.pool: Optional[BasicActorPool] = None
 
-        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
+        self.app_fn: Optional[Callable[[], ClientApp]] = None
 
     def _validate_client_resources(self, config: BackendConfig) -> ClientResourcesDict:
         client_resources_config = config.get(self.client_resources_key)
@@ -109,7 +82,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,59 +90,105 @@ 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(
+                runtime_env={"env_vars": {"PYTHONPATH": ":".join(sys.path)}},
+                **ray_init_args,
+            )
+
     @property
     def num_workers(self) -> int:
         """Return number of actors in pool."""
-        return self.pool.num_actors
+        return self.pool.num_actors if self.pool else 0
 
     def is_worker_idle(self) -> bool:
         """Report whether the pool has idle actors."""
-        return self.pool.is_actor_available()
+        return self.pool.is_actor_available() if self.pool else False
 
-    async def build(self) -> None:
+    def build(self, app_fn: Callable[[], ClientApp]) -> 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)
+        # Create Actor Pool
+        try:
+            self.pool = BasicActorPool(
+                actor_type=ClientAppActor,
+                client_resources=self.client_resources,
+                actor_kwargs=self.actor_kwargs,
+            )
+        except Exception as ex:
+            raise ex
+
+        self.pool.add_actors_to_pool(self.pool.actors_capacity)
+        # Set ClientApp callable that ray actors will use
+        self.app_fn = app_fn
+        log(DEBUG, "Constructed ActorPool with: %i actors", self.pool.num_actors)
 
-    async def process_message(
+    def process_message(
         self,
-        app: Callable[[], ClientApp],
         message: Message,
         context: Context,
-    ) -> Tuple[Message, Context]:
+    ) -> tuple[Message, Context]:
         """Run ClientApp that process a given message.
 
         Return output message and updated context.
         """
-        partition_id = message.metadata.partition_id
+        partition_id = context.node_config[PARTITION_ID_KEY]
+
+        if self.pool is None:
+            raise ValueError("The actor pool is empty, unfit to process messages.")
+
+        if self.app_fn is None:
+            raise ValueError(
+                "Unspecified function to load a `ClientApp`. "
+                "Call the backend's `build()` method before processing messages."
+            )
 
         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),
+                (self.app_fn, 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()
+        if self.pool:
+            self.pool.terminate_all_actors()
         ray.shutdown()
-        log(INFO, "Terminated %s", self.__class__.__name__)
+        log(DEBUG, "Terminated %s", self.__class__.__name__)