full loop tests; tweaks

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

@@ -18,6 +18,8 @@
 from logging import INFO
 from typing import Callable, Dict, List, Tuple, Union
 
+import ray
+
 from flwr.client.clientapp import ClientApp
 from flwr.common.context import Context
 from flwr.common.logger import log
@@ -157,3 +159,5 @@ async def process_message(
     async def terminate(self) -> None:
         """Terminate all actors in actor pool."""
         await self.pool.terminate_all_actors()
+        ray.shutdown()
+        log(INFO, "Terminated %s", self.__class__.__name__)

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

@@ -19,8 +19,6 @@
 from typing import Callable, Dict, Optional, Tuple, Union
 from unittest import IsolatedAsyncioTestCase
 
-import ray
-
 from flwr.client import Client, NumPyClient
 from flwr.client.clientapp import ClientApp
 from flwr.common import (
@@ -73,8 +71,6 @@ async def backend_build_process_and_termination(
 
     await backend.terminate()
 
-    ray.shutdown()
-
     return to_return
 
 

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

@@ -19,7 +19,7 @@
 import json
 import traceback
 from logging import DEBUG, ERROR, INFO
-from typing import Callable, Dict
+from typing import Callable, Dict, Optional
 
 from flwr.client.clientapp import ClientApp, load_client_app
 from flwr.client.node_state import NodeState
@@ -30,7 +30,6 @@
 
 from .backend import Backend, error_messages_backends, supported_backends
 
-TaskInsQueue = asyncio.Queue[TaskIns]
 NodeToPartitionMapping = Dict[int, int]
 
 
@@ -50,7 +49,7 @@ def _register_nodes(
 # pylint: disable=too-many-arguments,too-many-locals
 async def worker(
     app: Callable[[], ClientApp],
-    queue: TaskInsQueue,
+    queue: "asyncio.Queue[TaskIns]",
     node_states: Dict[int, NodeState],
     state_factory: StateFactory,
     nodes_mapping: NodeToPartitionMapping,
@@ -60,7 +59,7 @@ async def worker(
     state = state_factory.state()
     while True:
         try:
-            task_ins = await queue.get()
+            task_ins: TaskIns = await queue.get()
             node_id = task_ins.task.consumer.node_id
 
             # Register and retrive runstate
@@ -104,7 +103,7 @@ async def worker(
 
 
 async def generate_pull_requests(
-    queue: TaskInsQueue,
+    queue: "asyncio.Queue[TaskIns]",
     state_factory: StateFactory,
     nodes_mapping: NodeToPartitionMapping,
     f_stop: asyncio.Event,
@@ -132,7 +131,7 @@ async def run(
 ) -> None:
     """Run the VCE async."""
     # pylint: disable=fixme
-    queue: TaskInsQueue = asyncio.Queue(128)
+    queue: "asyncio.Queue[TaskIns]" = asyncio.Queue(128)
 
     # Build backend
     await backend.build()
@@ -150,7 +149,7 @@ async def run(
 
     # Produced task terminated, now cancel worker tasks
     for w_t in worker_tasks:
-        _ = w_t.cancel("Terminate on Simulation Engine shutdown.")
+        _ = w_t.cancel()
 
     # print('requested cancel')
     while not all(w_t.done() for w_t in worker_tasks):
@@ -172,12 +171,18 @@ def start_vce(
     state_factory: StateFactory,
     working_dir: str,
     f_stop: asyncio.Event,
+    existing_nodes_mapping: Optional[NodeToPartitionMapping] = None,
 ) -> None:
     """Start Fleet API with the VirtualClientEngine (VCE)."""
-    # Register SuperNodes
-    nodes_mapping = _register_nodes(
-        num_nodes=num_supernodes, state_factory=state_factory
-    )
+    if existing_nodes_mapping:
+        # Use mapping constructed externally. This also means nodes
+        # have previously being registered.
+        nodes_mapping = existing_nodes_mapping
+    else:
+        # Register SuperNodes
+        nodes_mapping = _register_nodes(
+            num_nodes=num_supernodes, state_factory=state_factory
+        )
 
     # Construct mapping of NodeStates
     node_states: Dict[int, NodeState] = {}

src/py/flwr/server/superlink/fleet/vce/vce_api_test.py

@@ -13,47 +13,185 @@
 # limitations under the License.
 # ==============================================================================
 """Test Fleet Simulation Engine API."""
-
 import asyncio
 import threading
+from itertools import cycle
+from math import pi
 from time import sleep
+from typing import Dict, Optional, Set
 from unittest import IsolatedAsyncioTestCase
+from uuid import UUID
+
+from flwr.client import Client, NumPyClient
+from flwr.client.clientapp import ClientApp
+from flwr.common import (
+    Config,
+    ConfigsRecord,
+    GetPropertiesIns,
+    Message,
+    Metadata,
+    Scalar,
+)
+from flwr.common.constant import MESSAGE_TYPE_GET_PROPERTIES
+from flwr.common.recordset_compat import getpropertiesins_to_recordset
+from flwr.common.serde import message_from_taskres, message_to_taskins
+from flwr.server.superlink.fleet.vce.vce_api import (
+    NodeToPartitionMapping,
+    _register_nodes,
+    start_vce,
+)
+from flwr.server.superlink.state import InMemoryState, StateFactory
+
+
+class DummyClient(NumPyClient):
+    """A dummy NumPyClient for tests."""
+
+    def get_properties(self, config: Config) -> Dict[str, Scalar]:
+        """Return properties by doing a simple calculation."""
+        result = float(config["factor"]) * pi
+
+        # store something in context
+        self.context.state.configs_records["result"] = ConfigsRecord({"result": result})
+        return {"result": result}
+
+
+def get_dummy_client(cid: str) -> Client:  # pylint: disable=unused-argument
+    """Return a DummyClient converted to Client type."""
+    return DummyClient().to_client()
+
+
+client_app = ClientApp(
+    client_fn=get_dummy_client,
+)
+
 
-from flwr.server.superlink.state import StateFactory
+def terminate_simulation(f_stop: asyncio.Event, sleep_duration: int) -> None:
+    """Set event to terminate Simulation Engine after `sleep_duration` seconds."""
+    sleep(sleep_duration)
+    f_stop.set()
 
-from . import start_vce
 
+def start_and_shutdown(
+    existing_state_factory: Optional[StateFactory] = None,
+    nodes_mapping: Optional[NodeToPartitionMapping] = None,
+    duration: int = 10,
+) -> None:
+    """Start Simulation Engine and terminate after specified number of seconds."""
+    f_stop = asyncio.Event()
+
+    # Initialize StateFactory
+    if nodes_mapping:
+        if existing_state_factory is None:
+            raise ValueError(
+                "If you specify a node mapping, you must pass a StateFactory."
+            )
+        state_factory = existing_state_factory
+    else:
+        state_factory = StateFactory(":flwr-in-memory-state:")
+
+    # Setup thread that will set the f_stop event, triggering the termination of all
+    # asyncio logic in the Simulation Engine. It will also terminate the Backend.
+    termination_th = threading.Thread(
+        target=terminate_simulation, args=(f_stop, duration)
+    )
+    termination_th.start()
+
+    start_vce(
+        num_supernodes=50,
+        client_app_module_name="vce_api_test:client_app",
+        backend_name="ray",
+        backend_config_json_stream="{}",  # an empty json stream (an empty config)
+        state_factory=state_factory,
+        working_dir="",
+        f_stop=f_stop,
+        existing_nodes_mapping=nodes_mapping,
+    )
+
+    # Trigger stop event
+    f_stop.set()
+
+    termination_th.join()
 
-class AsyncTestFleetSimulationEngine(IsolatedAsyncioTestCase):
-    """A basic class to test Fleet Simulation Enginge funcionality."""
+
+class AsyncTestFleetSimulationEngineRayBackend(IsolatedAsyncioTestCase):
+    """A basic class that enables testing asyncio functionalities."""
 
     def test_start_and_shutdown(self) -> None:
         """Start Simulation Engine Fleet and terminate it."""
-        f_stop = asyncio.Event()
+        start_and_shutdown()
+
+    # pylint: disable=too-many-locals
+    def test_start_and_shutdown_with_tasks_in_state(self) -> None:
+        """Run Simulation Engine with some TasksIns in State.
+
+        This test creates a few nodes and submits a few messages that need to be
+        executed by the Backend. In order for that to happen the asyncio
+        producer/consumer logic must function.
+        """
+        num_messages = 113
+        num_nodes = 59
 
-        # Initialize StateFactory
+        # Register a state and a run_id in it
+        run_id = 1234
         state_factory = StateFactory(":flwr-in-memory-state:")
+        state: InMemoryState = state_factory.state()  # type: ignore
+        state.run_ids.add(run_id)
 
-        superlink_th = threading.Thread(
-            target=start_vce,
-            args=(
-                50,
-                "",
-                "ray",
-                "{}",  # an empty json stream (represents an empty config)
-                state_factory,
-                "",
-                f_stop,
-            ),
-            daemon=False,
+        # Register a few nodes
+        nodes_mapping = _register_nodes(
+            num_nodes=num_nodes, state_factory=state_factory
         )
 
-        superlink_th.start()
+        # Artificially add TaskIns to state so they can be processed
+        # by the Simulation Engine logic
+        nodes_cycle = cycle(
+            nodes_mapping.keys()
+        )  # we have more messages than supernodes
+        task_ids: Set[UUID] = set()  # so we can retrieve them later
+        expected_results = {}
+        for i in range(num_messages):
+            dst_node_id = next(nodes_cycle)
+            # Construct a Message
+            mult_factor = 2024 + i
+            getproperties_ins = GetPropertiesIns(config={"factor": mult_factor})
+            recordset = getpropertiesins_to_recordset(getproperties_ins)
+            message = Message(
+                content=recordset,
+                metadata=Metadata(
+                    run_id=run_id,
+                    message_id="",
+                    group_id="",
+                    src_node_id=0,
+                    dst_node_id=dst_node_id,  # indicate destination node
+                    reply_to_message="",
+                    ttl="",
+                    message_type=MESSAGE_TYPE_GET_PROPERTIES,
+                ),
+            )
+            # Convert Message to TaskIns
+            taskins = message_to_taskins(message)
+            # Instert in state
+            task_id = state.store_task_ins(taskins)
+            if task_id:
+                # Add to UUID set
+                task_ids.add(task_id)
+                # Store expected output for check later on
+                expected_results[task_id] = mult_factor * pi
+
+        # Run
+        start_and_shutdown(state_factory, nodes_mapping)
+
+        # Get all TaskRes
+        task_res_list = state.get_task_res(task_ids=task_ids, limit=len(task_ids))
 
-        # Sleep for some time
-        sleep(10)
+        # Check results by first converting to Message
+        for task_res in task_res_list:
 
-        # Trigger stop event
-        f_stop.set()
+            message = message_from_taskres(task_res)
 
-        superlink_th.join()
+            # Verify message content is as expected
+            content = message.content
+            assert (
+                content.configs_records["getpropertiesres.properties"]["result"]
+                == expected_results[UUID(task_res.task.ancestry[0])]
+            )