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

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

@@ -15,19 +15,32 @@
 """Fleet Simulation Engine API."""
 
 
-import asyncio
 import json
+import threading
+import time
 import traceback
+from concurrent.futures import ThreadPoolExecutor
 from logging import DEBUG, ERROR, INFO, WARN
-from typing import Callable, Dict, List, Optional
+from pathlib import Path
+from queue import Empty, Queue
+from time import sleep
+from typing import Callable, Dict, Optional
 
-from flwr.client.client_app import ClientApp, LoadClientAppError
+from flwr.client.client_app import ClientApp, ClientAppException, LoadClientAppError
 from flwr.client.node_state import NodeState
+from flwr.client.supernode.app import _get_load_client_app_fn
+from flwr.common.constant import (
+    NUM_PARTITIONS_KEY,
+    PARTITION_ID_KEY,
+    PING_MAX_INTERVAL,
+    ErrorCode,
+)
 from flwr.common.logger import log
-from flwr.common.object_ref import load_app
+from flwr.common.message import Error
 from flwr.common.serde import message_from_taskins, message_to_taskres
-from flwr.proto.task_pb2 import TaskIns  # pylint: disable=E0611
-from flwr.server.superlink.state import StateFactory
+from flwr.common.typing import Run
+from flwr.proto.task_pb2 import TaskIns, TaskRes  # pylint: disable=E0611
+from flwr.server.superlink.state import State, StateFactory
 
 from .backend import Backend, error_messages_backends, supported_backends
 
@@ -41,39 +54,63 @@ def _register_nodes(
     nodes_mapping: NodeToPartitionMapping = {}
     state = state_factory.state()
     for i in range(num_nodes):
-        node_id = state.create_node()
+        node_id = state.create_node(ping_interval=PING_MAX_INTERVAL)
         nodes_mapping[node_id] = i
-    log(INFO, "Registered %i nodes", len(nodes_mapping))
+    log(DEBUG, "Registered %i nodes", len(nodes_mapping))
     return nodes_mapping
 
 
+def _register_node_states(
+    nodes_mapping: NodeToPartitionMapping,
+    run: Run,
+    app_dir: Optional[str] = None,
+) -> Dict[int, NodeState]:
+    """Create NodeState objects and pre-register the context for the run."""
+    node_states: Dict[int, NodeState] = {}
+    num_partitions = len(set(nodes_mapping.values()))
+    for node_id, partition_id in nodes_mapping.items():
+        node_states[node_id] = NodeState(
+            node_id=node_id,
+            node_config={
+                PARTITION_ID_KEY: partition_id,
+                NUM_PARTITIONS_KEY: num_partitions,
+            },
+        )
+
+        # Pre-register Context objects
+        node_states[node_id].register_context(
+            run_id=run.run_id, run=run, app_dir=app_dir
+        )
+
+    return node_states
+
+
 # pylint: disable=too-many-arguments,too-many-locals
-async def worker(
+def worker(
     app_fn: Callable[[], ClientApp],
-    queue: "asyncio.Queue[TaskIns]",
+    taskins_queue: "Queue[TaskIns]",
+    taskres_queue: "Queue[TaskRes]",
     node_states: Dict[int, NodeState],
-    state_factory: StateFactory,
-    nodes_mapping: NodeToPartitionMapping,
     backend: Backend,
+    f_stop: threading.Event,
 ) -> None:
     """Get TaskIns from queue and pass it to an actor in the pool to execute it."""
-    state = state_factory.state()
-    while True:
+    while not f_stop.is_set():
+        out_mssg = None
         try:
-            task_ins: TaskIns = await 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
-            node_states[node_id].register_context(run_id=task_ins.run_id)
+            # Retrieve context
             context = node_states[node_id].retrieve_context(run_id=task_ins.run_id)
 
             # Convert TaskIns to Message
             message = message_from_taskins(task_ins)
-            # Set partition_id
-            message.metadata.partition_id = nodes_mapping[node_id]
 
             # Let backend process message
-            out_mssg, updated_context = await backend.process_message(
+            out_mssg, updated_context = backend.process_message(
                 app_fn, message, context
             )
 
@@ -81,85 +118,74 @@ async def worker(
             node_states[node_id].update_context(
                 task_ins.run_id, context=updated_context
             )
-
-            # Convert to TaskRes
-            task_res = message_to_taskres(out_mssg)
-            # Store TaskRes in state
-            state.store_task_res(task_res)
-
-        except asyncio.CancelledError as e:
-            log(DEBUG, "Async worker: %s", e)
-            break
-
-        except LoadClientAppError as app_ex:
-            log(ERROR, "Async worker: %s", app_ex)
-            log(ERROR, traceback.format_exc())
-            raise
-
+        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)
             log(ERROR, traceback.format_exc())
-            break
 
+            if isinstance(ex, ClientAppException):
+                e_code = ErrorCode.CLIENT_APP_RAISED_EXCEPTION
+            elif isinstance(ex, LoadClientAppError):
+                e_code = ErrorCode.LOAD_CLIENT_APP_EXCEPTION
+            else:
+                e_code = ErrorCode.UNKNOWN
 
-async def add_taskins_to_queue(
-    queue: "asyncio.Queue[TaskIns]",
-    state_factory: StateFactory,
+            reason = str(type(ex)) + ":<'" + str(ex) + "'>"
+            out_mssg = message.create_error_reply(
+                error=Error(code=e_code, reason=reason)
+            )
+
+        finally:
+            if out_mssg:
+                # Convert to TaskRes
+                task_res = message_to_taskres(out_mssg)
+                # Store TaskRes in state
+                task_res.task.pushed_at = time.time()
+                taskres_queue.put(task_res)
+
+
+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.")
+def put_taskres_into_state(
+    state: State, queue: "Queue[TaskRes]", f_stop: threading.Event
+) -> None:
+    """Put TaskRes into State from a queue."""
+    while not f_stop.is_set():
+        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_api(
     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."""
-    queue: "asyncio.Queue[TaskIns]" = asyncio.Queue(128)
+    """Run the VCE."""
+    taskins_queue: "Queue[TaskIns]" = Queue()
+    taskres_queue: "Queue[TaskRes]" = Queue()
 
     try:
 
@@ -167,29 +193,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, queue, node_states, state_factory, nodes_mapping, backend
-                )
-            )
-            for _ in range(backend.num_workers)
-        ]
-        # Create producer (adds TaskIns into Queue)
-        producer = asyncio.create_task(
-            add_taskins_to_queue(
-                queue, state_factory, nodes_mapping, backend, worker_tasks, f_stop
-            )
+        state = state_factory.state()
+
+        extractor_th = threading.Thread(
+            target=add_taskins_to_queue,
+            args=(
+                state,
+                taskins_queue,
+                nodes_mapping,
+                f_stop,
+            ),
         )
+        extractor_th.start()
+
+        injector_th = threading.Thread(
+            target=put_taskres_into_state,
+            args=(
+                state,
+                taskres_queue,
+                f_stop,
+            ),
+        )
+        injector_th.start()
+
+        with ThreadPoolExecutor() as executor:
+            _ = [
+                executor.submit(
+                    worker,
+                    app_fn,
+                    taskins_queue,
+                    taskres_queue,
+                    node_states,
+                    backend,
+                    f_stop,
+                )
+                for _ in range(backend.num_workers)
+            ]
 
-        # Wait for producer to finish
-        # The producer runs 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(producer)
+        extractor_th.join()
+        injector_th.join()
 
     except Exception as ex:
 
@@ -204,26 +249,21 @@ 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
+# pylint: disable=too-many-arguments,unused-argument,too-many-locals,too-many-branches
+# pylint: disable=too-many-statements
 def start_vce(
     backend_name: str,
     backend_config_json_stream: str,
     app_dir: str,
-    f_stop: asyncio.Event,
+    is_app: bool,
+    f_stop: threading.Event,
+    run: Run,
+    flwr_dir: Optional[str] = None,
     client_app: Optional[ClientApp] = None,
     client_app_attr: Optional[str] = None,
     num_supernodes: Optional[int] = None,
@@ -259,6 +299,7 @@ def start_vce(
         # Use mapping constructed externally. This also means nodes
         # have previously being registered.
         nodes_mapping = existing_nodes_mapping
+    app_dir = str(Path(app_dir).absolute())
 
     if not state_factory:
         log(INFO, "A StateFactory was not supplied to the SimulationEngine.")
@@ -273,12 +314,12 @@ def start_vce(
         )
 
     # Construct mapping of NodeStates
-    node_states: Dict[int, NodeState] = {}
-    for node_id in nodes_mapping:
-        node_states[node_id] = NodeState()
+    node_states = _register_node_states(
+        nodes_mapping=nodes_mapping, run=run, app_dir=app_dir if is_app else None
+    )
 
     # Load backend config
-    log(INFO, "Supported backends: %s", list(supported_backends.keys()))
+    log(DEBUG, "Supported backends: %s", list(supported_backends.keys()))
     backend_config = json.loads(backend_config_json_stream)
 
     try:
@@ -298,20 +339,18 @@ def start_vce(
 
     def backend_fn() -> Backend:
         """Instantiate a Backend."""
-        return backend_type(backend_config, work_dir=app_dir)
-
-    log(INFO, "client_app_attr = %s", client_app_attr)
+        return backend_type(backend_config)
 
     # Load ClientApp if needed
     def _load() -> ClientApp:
 
         if client_app_attr:
-            app: ClientApp = load_app(client_app_attr, LoadClientAppError)
-
-            if not isinstance(app, ClientApp):
-                raise LoadClientAppError(
-                    f"Attribute {client_app_attr} is not of type {ClientApp}",
-                ) from None
+            app = _get_load_client_app_fn(
+                default_app_ref=client_app_attr,
+                app_path=app_dir,
+                flwr_dir=flwr_dir,
+                multi_app=False,
+            )(run.fab_id, run.fab_version)
 
         if client_app:
             app = client_app
@@ -319,8 +358,21 @@ def start_vce(
 
     app_fn = _load
 
-    asyncio.run(
-        run(
+    try:
+        # Test if ClientApp can be loaded
+        client_app = app_fn()
+
+        # Cache `ClientApp`
+        if client_app_attr:
+            # Now wrap the loaded ClientApp in a dummy function
+            # this prevent unnecesary low-level loading of ClientApp
+            def _load_client_app() -> ClientApp:
+                return client_app
+
+            app_fn = _load_client_app
+
+        # Run main simulation loop
+        run_api(
             app_fn,
             backend_fn,
             nodes_mapping,
@@ -328,4 +380,15 @@ def start_vce(
             node_states,
             f_stop,
         )
-    )
+    except LoadClientAppError as loadapp_ex:
+        f_stop_delay = 10
+        log(
+            ERROR,
+            "LoadClientAppError exception encountered. Terminating simulation in %is",
+            f_stop_delay,
+        )
+        time.sleep(f_stop_delay)
+        f_stop.set()  # set termination event
+        raise loadapp_ex
+    except Exception as ex:
+        raise ex

flwr/server/superlink/state/__init__.py

@@ -1,4 +1,4 @@
-# Copyright 2023 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.