flwr-nightly 1.10.0.dev20240710__py3-none-any.whl → 1.10.0.dev20240712__py3-none-any.whl

flwr/server/server_app.py

@@ -17,8 +17,11 @@
 
 from typing import Callable, Optional
 
-from flwr.common import Context, RecordSet
-from flwr.common.logger import warn_preview_feature
+from flwr.common import Context
+from flwr.common.logger import (
+    warn_deprecated_feature_with_example,
+    warn_preview_feature,
+)
 from flwr.server.strategy import Strategy
 
 from .client_manager import ClientManager
@@ -26,7 +29,20 @@ from .compat import start_driver
 from .driver import Driver
 from .server import Server
 from .server_config import ServerConfig
-from .typing import ServerAppCallable
+from .typing import ServerAppCallable, ServerFn
+
+SERVER_FN_USAGE_EXAMPLE = """
+
+        def server_fn(context: Context):
+            server_config = ServerConfig(num_rounds=3)
+            strategy = FedAvg()
+            return ServerAppComponents(
+                strategy=strategy,
+                server_config=server_config,
+        )
+
+        app = ServerApp(server_fn=server_fn)
+"""
 
 
 class ServerApp:
@@ -36,13 +52,15 @@ class ServerApp:
     --------
     Use the `ServerApp` with an existing `Strategy`:
 
-    >>> server_config = ServerConfig(num_rounds=3)
-    >>> strategy = FedAvg()
+    >>> def server_fn(context: Context):
+    >>>     server_config = ServerConfig(num_rounds=3)
+    >>>     strategy = FedAvg()
+    >>>     return ServerAppComponents(
+    >>>         strategy=strategy,
+    >>>         server_config=server_config,
+    >>>     )
     >>>
-    >>> app = ServerApp(
-    >>>     server_config=server_config,
-    >>>     strategy=strategy,
-    >>> )
+    >>> app = ServerApp(server_fn=server_fn)
 
     Use the `ServerApp` with a custom main function:
 
@@ -53,23 +71,52 @@ class ServerApp:
     >>>    print("ServerApp running")
     """
 
+    # pylint: disable=too-many-arguments
     def __init__(
         self,
         server: Optional[Server] = None,
         config: Optional[ServerConfig] = None,
         strategy: Optional[Strategy] = None,
         client_manager: Optional[ClientManager] = None,
+        server_fn: Optional[ServerFn] = None,
     ) -> None:
+        if any([server, config, strategy, client_manager]):
+            warn_deprecated_feature_with_example(
+                deprecation_message="Passing either `server`, `config`, `strategy` or "
+                "`client_manager` directly to the ServerApp "
+                "constructor is deprecated.",
+                example_message="Pass `ServerApp` arguments wrapped "
+                "in a `flwr.server.ServerAppComponents` object that gets "
+                "returned by a function passed as the `server_fn` argument "
+                "to the `ServerApp` constructor. For example: ",
+                code_example=SERVER_FN_USAGE_EXAMPLE,
+            )
+
+            if server_fn:
+                raise ValueError(
+                    "Passing `server_fn` is incompatible with passing the "
+                    "other arguments (now deprecated) to ServerApp. "
+                    "Use `server_fn` exclusively."
+                )
+
         self._server = server
         self._config = config
         self._strategy = strategy
         self._client_manager = client_manager
+        self._server_fn = server_fn
         self._main: Optional[ServerAppCallable] = None
 
     def __call__(self, driver: Driver, context: Context) -> None:
         """Execute `ServerApp`."""
         # Compatibility mode
         if not self._main:
+            if self._server_fn:
+                # Execute server_fn()
+                components = self._server_fn(context)
+                self._server = components.server
+                self._config = components.config
+                self._strategy = components.strategy
+                self._client_manager = components.client_manager
             start_driver(
                 server=self._server,
                 config=self._config,
@@ -80,7 +127,6 @@ class ServerApp:
             return
 
         # New execution mode
-        context = Context(state=RecordSet(), run_config={})
         self._main(driver, context)
 
     def main(self) -> Callable[[ServerAppCallable], ServerAppCallable]:

flwr/server/serverapp_components.py

@@ -0,0 +1,52 @@
+# 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.
+# ==============================================================================
+"""ServerAppComponents for the ServerApp."""
+
+
+from dataclasses import dataclass
+from typing import Optional
+
+from .client_manager import ClientManager
+from .server import Server
+from .server_config import ServerConfig
+from .strategy import Strategy
+
+
+@dataclass
+class ServerAppComponents:  # pylint: disable=too-many-instance-attributes
+    """Components to construct a ServerApp.
+
+    Parameters
+    ----------
+    server : Optional[Server] (default: None)
+        A server implementation, either `flwr.server.Server` or a subclass
+        thereof. If no instance is provided, one will be created internally.
+    config : Optional[ServerConfig] (default: None)
+        Currently supported values are `num_rounds` (int, default: 1) and
+        `round_timeout` in seconds (float, default: None).
+    strategy : Optional[Strategy] (default: None)
+        An implementation of the abstract base class
+        `flwr.server.strategy.Strategy`. If no strategy is provided, then
+        `flwr.server.strategy.FedAvg` will be used.
+    client_manager : Optional[ClientManager] (default: None)
+        An implementation of the class `flwr.server.ClientManager`. If no
+        implementation is provided, then `flwr.server.SimpleClientManager`
+        will be used.
+    """
+
+    server: Optional[Server] = None
+    config: Optional[ServerConfig] = None
+    strategy: Optional[Strategy] = None
+    client_manager: Optional[ClientManager] = None

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

@@ -33,8 +33,8 @@ class Backend(ABC):
         """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

@@ -153,12 +153,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)
+        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,
@@ -172,17 +172,16 @@ class RayBackend(Backend):
 
         try:
             # Submit a task to the pool
-            future = await self.pool.submit(
+            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
 
@@ -193,11 +192,11 @@ class RayBackend(Backend):
                 self.__class__.__name__,
             )
             # add actor back into pool
-            await self.pool.add_actor_back_to_pool(future)
+            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(DEBUG, "Terminated %s", self.__class__.__name__)

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

@@ -14,14 +14,18 @@
 # ==============================================================================
 """Fleet Simulation Engine API."""
 
-import asyncio
+
 import json
 import sys
+import threading
 import time
 import traceback
+from concurrent.futures import ThreadPoolExecutor
 from logging import DEBUG, ERROR, INFO, WARN
 from pathlib import Path
-from typing import Callable, Dict, List, Optional
+from queue import Empty, Queue
+from time import sleep
+from typing import Callable, Dict, Optional
 
 from flwr.client.client_app import ClientApp, ClientAppException, LoadClientAppError
 from flwr.client.node_state import NodeState
@@ -31,7 +35,7 @@ from flwr.common.message import Error
 from flwr.common.object_ref import load_app
 from flwr.common.serde import message_from_taskins, message_to_taskres
 from flwr.proto.task_pb2 import TaskIns, TaskRes  # pylint: disable=E0611
-from flwr.server.superlink.state import StateFactory
+from flwr.server.superlink.state import State, StateFactory
 
 from .backend import Backend, error_messages_backends, supported_backends
 
@@ -52,18 +56,21 @@ def _register_nodes(
 
 
 # pylint: disable=too-many-arguments,too-many-locals
-async def worker(
+def worker(
     app_fn: Callable[[], ClientApp],
-    taskins_queue: "asyncio.Queue[TaskIns]",
-    taskres_queue: "asyncio.Queue[TaskRes]",
+    taskins_queue: "Queue[TaskIns]",
+    taskres_queue: "Queue[TaskRes]",
     node_states: Dict[int, NodeState],
     backend: Backend,
+    f_stop: threading.Event,
 ) -> None:
     """Get TaskIns from queue and pass it to an actor in the pool to execute it."""
-    while True:
+    while not f_stop.is_set():
         out_mssg = None
         try:
-            task_ins: TaskIns = await taskins_queue.get()
+            # Fetch from queue with timeout. We use a timeout so
+            # the stopping event can be evaluated even when the queue is empty.
+            task_ins: TaskIns = taskins_queue.get(timeout=1.0)
             node_id = task_ins.task.consumer.node_id
 
             # Register and retrieve runstate
@@ -74,7 +81,7 @@ async def worker(
             message = message_from_taskins(task_ins)
 
             # Let backend process message
-            out_mssg, updated_context = await backend.process_message(
+            out_mssg, updated_context = backend.process_message(
                 app_fn, message, context
             )
 
@@ -82,11 +89,9 @@ async def worker(
             node_states[node_id].update_context(
                 task_ins.run_id, context=updated_context
             )
-
-        except asyncio.CancelledError as e:
-            log(DEBUG, "Terminating async worker: %s", e)
-            break
-
+        except Empty:
+            # An exception raised if queue.get times out
+            pass
         # Exceptions aren't raised but reported as an error message
         except Exception as ex:  # pylint: disable=broad-exception-caught
             log(ERROR, ex)
@@ -110,83 +115,48 @@ async def worker(
                 task_res = message_to_taskres(out_mssg)
                 # Store TaskRes in state
                 task_res.task.pushed_at = time.time()
-                await taskres_queue.put(task_res)
+                taskres_queue.put(task_res)
 
 
-async def add_taskins_to_queue(
-    queue: "asyncio.Queue[TaskIns]",
-    state_factory: StateFactory,
+def add_taskins_to_queue(
+    state: State,
+    queue: "Queue[TaskIns]",
     nodes_mapping: NodeToPartitionMapping,
-    backend: Backend,
-    consumers: List["asyncio.Task[None]"],
-    f_stop: asyncio.Event,
+    f_stop: threading.Event,
 ) -> None:
-    """Retrieve TaskIns and add it to the queue."""
-    state = state_factory.state()
-    num_initial_consumers = len(consumers)
+    """Put TaskIns in a queue from State."""
     while not f_stop.is_set():
         for node_id in nodes_mapping.keys():
-            task_ins = state.get_task_ins(node_id=node_id, limit=1)
-            if task_ins:
-                await queue.put(task_ins[0])
-
-        # Count consumers that are running
-        num_active = sum(not (cc.done()) for cc in consumers)
-
-        # Alert if number of consumers decreased by half
-        if num_active < num_initial_consumers // 2:
-            log(
-                WARN,
-                "Number of active workers has more than halved: (%i/%i active)",
-                num_active,
-                num_initial_consumers,
-            )
+            task_ins_list = state.get_task_ins(node_id=node_id, limit=1)
+            for task_ins in task_ins_list:
+                queue.put(task_ins)
+        sleep(0.1)
 
-        # Break if consumers died
-        if num_active == 0:
-            raise RuntimeError("All workers have died. Ending Simulation.")
 
-        # Log some stats
-        log(
-            DEBUG,
-            "Simulation Engine stats: "
-            "Active workers: (%i/%i) | %s (%i workers) | Tasks in queue: %i)",
-            num_active,
-            num_initial_consumers,
-            backend.__class__.__name__,
-            backend.num_workers,
-            queue.qsize(),
-        )
-        await asyncio.sleep(1.0)
-    log(DEBUG, "Async producer: Stopped pulling from StateFactory.")
-
-
-async def put_taskres_into_state(
-    queue: "asyncio.Queue[TaskRes]",
-    state_factory: StateFactory,
-    f_stop: asyncio.Event,
+def put_taskres_into_state(
+    state: State, queue: "Queue[TaskRes]", f_stop: threading.Event
 ) -> None:
-    """Remove TaskRes from queue and add into State."""
-    state = state_factory.state()
+    """Put TaskRes into State from a queue."""
     while not f_stop.is_set():
-        if queue.qsize():
-            task_res = await queue.get()
-            state.store_task_res(task_res)
-        else:
-            await asyncio.sleep(0.1)
+        try:
+            taskres = queue.get(timeout=1.0)
+            state.store_task_res(taskres)
+        except Empty:
+            # queue is empty when timeout was triggered
+            pass
 
 
-async def run(
+def run(
     app_fn: Callable[[], ClientApp],
     backend_fn: Callable[[], Backend],
     nodes_mapping: NodeToPartitionMapping,
     state_factory: StateFactory,
     node_states: Dict[int, NodeState],
-    f_stop: asyncio.Event,
+    f_stop: threading.Event,
 ) -> None:
-    """Run the VCE async."""
-    taskins_queue: "asyncio.Queue[TaskIns]" = asyncio.Queue(128)
-    taskres_queue: "asyncio.Queue[TaskRes]" = asyncio.Queue(128)
+    """Run the VCE."""
+    taskins_queue: "Queue[TaskIns]" = Queue()
+    taskres_queue: "Queue[TaskRes]" = Queue()
 
     try:
 
@@ -194,42 +164,48 @@ async def run(
         backend = backend_fn()
 
         # Build backend
-        await backend.build()
+        backend.build()
 
         # Add workers (they submit Messages to Backend)
-        worker_tasks = [
-            asyncio.create_task(
-                worker(
-                    app_fn,
-                    taskins_queue,
-                    taskres_queue,
-                    node_states,
-                    backend,
-                )
-            )
-            for _ in range(backend.num_workers)
-        ]
-        # Create producer (adds TaskIns into Queue)
-        taskins_producer = asyncio.create_task(
-            add_taskins_to_queue(
+        state = state_factory.state()
+
+        extractor_th = threading.Thread(
+            target=add_taskins_to_queue,
+            args=(
+                state,
                 taskins_queue,
-                state_factory,
                 nodes_mapping,
-                backend,
-                worker_tasks,
                 f_stop,
-            )
+            ),
         )
+        extractor_th.start()
 
-        taskres_consumer = asyncio.create_task(
-            put_taskres_into_state(taskres_queue, state_factory, f_stop)
+        injector_th = threading.Thread(
+            target=put_taskres_into_state,
+            args=(
+                state,
+                taskres_queue,
+                f_stop,
+            ),
         )
+        injector_th.start()
 
-        # Wait for asyncio taks pulling/pushing TaskIns/TaskRes.
-        # These run forever until f_stop is set or until
-        # all worker (consumer) coroutines are completed. Workers
-        # also run forever and only end if an exception is raised.
-        await asyncio.gather(*(taskins_producer, taskres_consumer))
+        with ThreadPoolExecutor() as executor:
+            _ = [
+                executor.submit(
+                    worker,
+                    app_fn,
+                    taskins_queue,
+                    taskres_queue,
+                    node_states,
+                    backend,
+                    f_stop,
+                )
+                for _ in range(backend.num_workers)
+            ]
+
+        extractor_th.join()
+        injector_th.join()
 
     except Exception as ex:
 
@@ -244,18 +220,9 @@ async def run(
         raise RuntimeError("Simulation Engine crashed.") from ex
 
     finally:
-        # Produced task terminated, now cancel worker tasks
-        for w_t in worker_tasks:
-            _ = w_t.cancel()
-
-        while not all(w_t.done() for w_t in worker_tasks):
-            log(DEBUG, "Terminating async workers...")
-            await asyncio.sleep(0.5)
-
-        await asyncio.gather(*[w_t for w_t in worker_tasks if not w_t.done()])
 
         # Terminate backend
-        await backend.terminate()
+        backend.terminate()
 
 
 # pylint: disable=too-many-arguments,unused-argument,too-many-locals,too-many-branches
@@ -264,7 +231,7 @@ def start_vce(
     backend_name: str,
     backend_config_json_stream: str,
     app_dir: str,
-    f_stop: asyncio.Event,
+    f_stop: threading.Event,
     client_app: Optional[ClientApp] = None,
     client_app_attr: Optional[str] = None,
     num_supernodes: Optional[int] = None,
@@ -317,7 +284,9 @@ def start_vce(
     # Construct mapping of NodeStates
     node_states: Dict[int, NodeState] = {}
     for node_id, partition_id in nodes_mapping.items():
-        node_states[node_id] = NodeState(partition_id=partition_id)
+        node_states[node_id] = NodeState(
+            node_id=node_id, node_config={}, partition_id=partition_id
+        )
 
     # Load backend config
     log(DEBUG, "Supported backends: %s", list(supported_backends.keys()))
@@ -368,15 +337,13 @@ def start_vce(
         _ = app_fn()
 
         # Run main simulation loop
-        asyncio.run(
-            run(
-                app_fn,
-                backend_fn,
-                nodes_mapping,
-                state_factory,
-                node_states,
-                f_stop,
-            )
+        run(
+            app_fn,
+            backend_fn,
+            nodes_mapping,
+            state_factory,
+            node_states,
+            f_stop,
         )
     except LoadClientAppError as loadapp_ex:
         f_stop_delay = 10

flwr/server/typing.py

@@ -20,6 +20,8 @@ from typing import Callable
 from flwr.common import Context
 
 from .driver import Driver
+from .serverapp_components import ServerAppComponents
 
 ServerAppCallable = Callable[[Driver, Context], None]
 Workflow = Callable[[Driver, Context], None]
+ServerFn = Callable[[Context], ServerAppComponents]

flwr/simulation/ray_transport/ray_actor.py

@@ -14,7 +14,6 @@
 # ==============================================================================
 """Ray-based Flower Actor and ActorPool implementation."""
 
-import asyncio
 import threading
 from abc import ABC
 from logging import DEBUG, ERROR, WARNING
@@ -411,9 +410,7 @@ class BasicActorPool:
         self.client_resources = client_resources
 
         # Queue of idle actors
-        self.pool: "asyncio.Queue[Type[VirtualClientEngineActor]]" = asyncio.Queue(
-            maxsize=1024
-        )
+        self.pool: List[VirtualClientEngineActor] = []
         self.num_actors = 0
 
         # Resolve arguments to pass during actor init
@@ -427,38 +424,37 @@ class BasicActorPool:
         # Figure out how many actors can be created given the cluster resources
         # and the resources the user indicates each VirtualClient will need
         self.actors_capacity = pool_size_from_resources(client_resources)
-        self._future_to_actor: Dict[Any, Type[VirtualClientEngineActor]] = {}
+        self._future_to_actor: Dict[Any, VirtualClientEngineActor] = {}
 
     def is_actor_available(self) -> bool:
         """Return true if there is an idle actor."""
-        return self.pool.qsize() > 0
+        return len(self.pool) > 0
 
-    async def add_actors_to_pool(self, num_actors: int) -> None:
+    def add_actors_to_pool(self, num_actors: int) -> None:
         """Add actors to the pool.
 
         This method may be executed also if new resources are added to your Ray cluster
         (e.g. you add a new node).
         """
         for _ in range(num_actors):
-            await self.pool.put(self.create_actor_fn())  # type: ignore
+            self.pool.append(self.create_actor_fn())  # type: ignore
         self.num_actors += num_actors
 
-    async def terminate_all_actors(self) -> None:
+    def terminate_all_actors(self) -> None:
         """Terminate actors in pool."""
         num_terminated = 0
-        while self.pool.qsize():
-            actor = await self.pool.get()
+        for actor in self.pool:
             actor.terminate.remote()  # type: ignore
             num_terminated += 1
 
         log(DEBUG, "Terminated %i actors", num_terminated)
 
-    async def submit(
+    def submit(
         self, actor_fn: Any, job: Tuple[ClientAppFn, Message, str, Context]
     ) -> Any:
         """On idle actor, submit job and return future."""
         # Remove idle actor from pool
-        actor = await self.pool.get()
+        actor = self.pool.pop()
         # Submit job to actor
         app_fn, mssg, cid, context = job
         future = actor_fn(actor, app_fn, mssg, cid, context)
@@ -467,18 +463,18 @@ class BasicActorPool:
         self._future_to_actor[future] = actor
         return future
 
-    async def add_actor_back_to_pool(self, future: Any) -> None:
+    def add_actor_back_to_pool(self, future: Any) -> None:
         """Ad actor assigned to run future back into the pool."""
         actor = self._future_to_actor.pop(future)
-        await self.pool.put(actor)
+        self.pool.append(actor)
 
-    async def fetch_result_and_return_actor_to_pool(
+    def fetch_result_and_return_actor_to_pool(
         self, future: Any
     ) -> Tuple[Message, Context]:
         """Pull result given a future and add actor back to pool."""
-        # Get actor that ran job
-        await self.add_actor_back_to_pool(future)
         # Retrieve result for object store
         # Instead of doing ray.get(future) we await it
-        _, out_mssg, updated_context = await future
+        _, out_mssg, updated_context = ray.get(future)
+        # Get actor that ran job
+        self.add_actor_back_to_pool(future)
         return out_mssg, updated_context