flwr-nightly 1.8.0.dev20240323__py3-none-any.whl → 1.8.0.dev20240328__py3-none-any.whl

flwr/server/compat/driver_client_proxy.py

@@ -19,7 +19,7 @@ import time
 from typing import List, Optional
 
 from flwr import common
-from flwr.common import MessageType, MessageTypeLegacy, RecordSet
+from flwr.common import DEFAULT_TTL, MessageType, MessageTypeLegacy, RecordSet
 from flwr.common import recordset_compat as compat
 from flwr.common import serde
 from flwr.proto import driver_pb2, node_pb2, task_pb2  # pylint: disable=E0611
@@ -129,8 +129,16 @@ class DriverClientProxy(ClientProxy):
                 ),
                 task_type=task_type,
                 recordset=serde.recordset_to_proto(recordset),
+                ttl=DEFAULT_TTL,
             ),
         )
+
+        # This would normally be recorded upon common.Message creation
+        # but this compatibility stack doesn't create Messages,
+        # so we need to inject `created_at` manually (needed for
+        # taskins validation by server.utils.validator)
+        task_ins.task.created_at = time.time()
+
         push_task_ins_req = driver_pb2.PushTaskInsRequest(  # pylint: disable=E1101
             task_ins_list=[task_ins]
         )
@@ -162,8 +170,24 @@ class DriverClientProxy(ClientProxy):
             )
             if len(task_res_list) == 1:
                 task_res = task_res_list[0]
+
+                # This will raise an Exception if task_res carries an `error`
+                validate_task_res(task_res=task_res)
+
                 return serde.recordset_from_proto(task_res.task.recordset)
 
             if timeout is not None and time.time() > start_time + timeout:
                 raise RuntimeError("Timeout reached")
             time.sleep(SLEEP_TIME)
+
+
+def validate_task_res(
+    task_res: task_pb2.TaskRes,  # pylint: disable=E1101
+) -> None:
+    """Validate if a TaskRes is empty or not."""
+    if not task_res.HasField("task"):
+        raise ValueError("Invalid TaskRes, field `task` missing")
+    if task_res.task.HasField("error"):
+        raise ValueError("Exception during client-side task execution")
+    if not task_res.task.HasField("recordset"):
+        raise ValueError("Invalid TaskRes, both `recordset` and `error` are missing")

flwr/server/driver/driver.py

@@ -18,7 +18,7 @@
 import time
 from typing import Iterable, List, Optional, Tuple
 
-from flwr.common import Message, Metadata, RecordSet
+from flwr.common import DEFAULT_TTL, Message, Metadata, RecordSet
 from flwr.common.serde import message_from_taskres, message_to_taskins
 from flwr.proto.driver_pb2 import (  # pylint: disable=E0611
     CreateRunRequest,
@@ -81,6 +81,7 @@ class Driver:
             and message.metadata.src_node_id == self.node.node_id
             and message.metadata.message_id == ""
             and message.metadata.reply_to_message == ""
+            and message.metadata.ttl > 0
         ):
             raise ValueError(f"Invalid message: {message}")
 
@@ -90,7 +91,7 @@ class Driver:
         message_type: str,
         dst_node_id: int,
         group_id: str,
-        ttl: str,
+        ttl: float = DEFAULT_TTL,
     ) -> Message:
         """Create a new message with specified parameters.
 
@@ -110,10 +111,10 @@ class Driver:
         group_id : str
             The ID of the group to which this message is associated. In some settings,
             this is used as the FL round.
-        ttl : str
+        ttl : float (default: common.DEFAULT_TTL)
             Time-to-live for the round trip of this message, i.e., the time from sending
-            this message to receiving a reply. It specifies the duration for which the
-            message and its potential reply are considered valid.
+            this message to receiving a reply. It specifies in seconds the duration for
+            which the message and its potential reply are considered valid.
 
         Returns
         -------

flwr/server/superlink/driver/driver_servicer.py

@@ -15,6 +15,7 @@
 """Driver API servicer."""
 
 
+import time
 from logging import DEBUG, INFO
 from typing import List, Optional, Set
 from uuid import UUID
@@ -72,6 +73,11 @@ class DriverServicer(driver_pb2_grpc.DriverServicer):
         """Push a set of TaskIns."""
         log(DEBUG, "DriverServicer.PushTaskIns")
 
+        # Set pushed_at (timestamp in seconds)
+        pushed_at = time.time()
+        for task_ins in request.task_ins_list:
+            task_ins.task.pushed_at = pushed_at
+
         # Validate request
         _raise_if(len(request.task_ins_list) == 0, "`task_ins_list` must not be empty")
         for task_ins in request.task_ins_list:

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

@@ -15,7 +15,7 @@
 """Fleet API gRPC request-response servicer."""
 
 
-from logging import INFO
+from logging import DEBUG, INFO
 
 import grpc
 
@@ -26,6 +26,8 @@ from flwr.proto.fleet_pb2 import (  # pylint: disable=E0611
     CreateNodeResponse,
     DeleteNodeRequest,
     DeleteNodeResponse,
+    PingRequest,
+    PingResponse,
     PullTaskInsRequest,
     PullTaskInsResponse,
     PushTaskResRequest,
@@ -61,6 +63,14 @@ class FleetServicer(fleet_pb2_grpc.FleetServicer):
             state=self.state_factory.state(),
         )
 
+    def Ping(self, request: PingRequest, context: grpc.ServicerContext) -> PingResponse:
+        """."""
+        log(DEBUG, "FleetServicer.Ping")
+        return message_handler.ping(
+            request=request,
+            state=self.state_factory.state(),
+        )
+
     def PullTaskIns(
         self, request: PullTaskInsRequest, context: grpc.ServicerContext
     ) -> PullTaskInsResponse:

flwr/server/superlink/fleet/message_handler/message_handler.py

@@ -15,6 +15,7 @@
 """Fleet API message handlers."""
 
 
+import time
 from typing import List, Optional
 from uuid import UUID
 
@@ -23,6 +24,8 @@ from flwr.proto.fleet_pb2 import (  # pylint: disable=E0611
     CreateNodeResponse,
     DeleteNodeRequest,
     DeleteNodeResponse,
+    PingRequest,
+    PingResponse,
     PullTaskInsRequest,
     PullTaskInsResponse,
     PushTaskResRequest,
@@ -55,6 +58,14 @@ 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:
+    """."""
+    return PingResponse(success=True)
+
+
 def pull_task_ins(request: PullTaskInsRequest, state: State) -> PullTaskInsResponse:
     """Pull TaskIns handler."""
     # Get node_id if client node is not anonymous
@@ -77,6 +88,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)
 

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

@@ -20,7 +20,7 @@ from typing import Callable, Dict, List, Tuple, Union
 
 import ray
 
-from flwr.client.client_app import ClientApp, LoadClientAppError
+from flwr.client.client_app import ClientApp
 from flwr.common.context import Context
 from flwr.common.logger import log
 from flwr.common.message import Message
@@ -151,7 +151,6 @@ class RayBackend(Backend):
             )
 
             await future
-
             # Fetch result
             (
                 out_mssg,
@@ -160,13 +159,15 @@ class RayBackend(Backend):
 
             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
+            await self.pool.add_actor_back_to_pool(future)
+            raise ex
 
     async def terminate(self) -> None:
         """Terminate all actors in actor pool."""

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

@@ -14,9 +14,10 @@
 # ==============================================================================
 """Fleet Simulation Engine API."""
 
-
 import asyncio
 import json
+import sys
+import time
 import traceback
 from logging import DEBUG, ERROR, INFO, WARN
 from typing import Callable, Dict, List, Optional
@@ -24,6 +25,7 @@ from typing import Callable, Dict, List, Optional
 from flwr.client.client_app import ClientApp, LoadClientAppError
 from flwr.client.node_state import NodeState
 from flwr.common.logger import log
+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  # pylint: disable=E0611
@@ -59,6 +61,7 @@ async def worker(
     """Get TaskIns from queue and pass it to an actor in the pool to execute it."""
     state = state_factory.state()
     while True:
+        out_mssg = None
         try:
             task_ins: TaskIns = await queue.get()
             node_id = task_ins.task.consumer.node_id
@@ -82,24 +85,25 @@ async def worker(
                 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)
+            log(DEBUG, "Terminating async worker: %s", e)
             break
 
-        except LoadClientAppError as app_ex:
-            log(ERROR, "Async worker: %s", app_ex)
-            log(ERROR, traceback.format_exc())
-            raise
-
+        # 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
+            reason = str(type(ex)) + ":<'" + str(ex) + "'>"
+            error = Error(code=0, reason=reason)
+            out_mssg = message.create_error_reply(error=error)
+
+        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()
+                state.store_task_res(task_res)
 
 
 async def add_taskins_to_queue(
@@ -218,7 +222,7 @@ async def run(
         await 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
 def start_vce(
     backend_name: str,
     backend_config_json_stream: str,
@@ -300,12 +304,14 @@ def start_vce(
         """Instantiate a Backend."""
         return backend_type(backend_config, work_dir=app_dir)
 
-    log(INFO, "client_app_attr = %s", client_app_attr)
-
     # Load ClientApp if needed
     def _load() -> ClientApp:
 
         if client_app_attr:
+
+            if app_dir is not None:
+                sys.path.insert(0, app_dir)
+
             app: ClientApp = load_app(client_app_attr, LoadClientAppError)
 
             if not isinstance(app, ClientApp):
@@ -319,13 +325,23 @@ def start_vce(
 
     app_fn = _load
 
-    asyncio.run(
-        run(
-            app_fn,
-            backend_fn,
-            nodes_mapping,
-            state_factory,
-            node_states,
-            f_stop,
+    try:
+        # Test if ClientApp can be loaded
+        _ = app_fn()
+
+        # Run main simulation loop
+        asyncio.run(
+            run(
+                app_fn,
+                backend_fn,
+                nodes_mapping,
+                state_factory,
+                node_states,
+                f_stop,
+            )
         )
-    )
+    except LoadClientAppError as loadapp_ex:
+        f_stop.set()  # set termination event
+        raise loadapp_ex
+    except Exception as ex:
+        raise ex

flwr/server/superlink/state/in_memory_state.py

@@ -17,9 +17,9 @@
 
 import os
 import threading
-from datetime import datetime, timedelta
+import time
 from logging import ERROR
-from typing import Dict, List, Optional, Set
+from typing import Dict, List, Optional, Set, Tuple
 from uuid import UUID, uuid4
 
 from flwr.common import log, now
@@ -32,7 +32,8 @@ class InMemoryState(State):
     """In-memory State implementation."""
 
     def __init__(self) -> None:
-        self.node_ids: Set[int] = set()
+        # Map node_id to (online_until, ping_interval)
+        self.node_ids: Dict[int, Tuple[float, float]] = {}
         self.run_ids: Set[int] = set()
         self.task_ins_store: Dict[UUID, TaskIns] = {}
         self.task_res_store: Dict[UUID, TaskRes] = {}
@@ -50,15 +51,11 @@ class InMemoryState(State):
             log(ERROR, "`run_id` is invalid")
             return None
 
-        # Create task_id, created_at and ttl
+        # Create task_id
         task_id = uuid4()
-        created_at: datetime = now()
-        ttl: datetime = created_at + timedelta(hours=24)
 
         # Store TaskIns
         task_ins.task_id = str(task_id)
-        task_ins.task.created_at = created_at.isoformat()
-        task_ins.task.ttl = ttl.isoformat()
         with self.lock:
             self.task_ins_store[task_id] = task_ins
 
@@ -113,15 +110,11 @@ class InMemoryState(State):
             log(ERROR, "`run_id` is invalid")
             return None
 
-        # Create task_id, created_at and ttl
+        # Create task_id
         task_id = uuid4()
-        created_at: datetime = now()
-        ttl: datetime = created_at + timedelta(hours=24)
 
         # Store TaskRes
         task_res.task_id = str(task_id)
-        task_res.task.created_at = created_at.isoformat()
-        task_res.task.ttl = ttl.isoformat()
         with self.lock:
             self.task_res_store[task_id] = task_res
 
@@ -194,17 +187,21 @@ class InMemoryState(State):
         # Sample a random int64 as node_id
         node_id: int = int.from_bytes(os.urandom(8), "little", signed=True)
 
-        if node_id not in self.node_ids:
-            self.node_ids.add(node_id)
-            return node_id
+        with self.lock:
+            if node_id not in self.node_ids:
+                # Default ping interval is 30s
+                # TODO: change 1e9 to 30s  # pylint: disable=W0511
+                self.node_ids[node_id] = (time.time() + 1e9, 1e9)
+                return node_id
         log(ERROR, "Unexpected node registration failure.")
         return 0
 
     def delete_node(self, node_id: int) -> None:
         """Delete a client node."""
-        if node_id not in self.node_ids:
-            raise ValueError(f"Node {node_id} not found")
-        self.node_ids.remove(node_id)
+        with self.lock:
+            if node_id not in self.node_ids:
+                raise ValueError(f"Node {node_id} not found")
+            del self.node_ids[node_id]
 
     def get_nodes(self, run_id: int) -> Set[int]:
         """Return all available client nodes.
@@ -214,17 +211,32 @@ class InMemoryState(State):
         If the provided `run_id` does not exist or has no matching nodes,
         an empty `Set` MUST be returned.
         """
-        if run_id not in self.run_ids:
-            return set()
-        return self.node_ids
+        with self.lock:
+            if run_id not in self.run_ids:
+                return set()
+            current_time = time.time()
+            return {
+                node_id
+                for node_id, (online_until, _) in self.node_ids.items()
+                if online_until > current_time
+            }
 
     def create_run(self) -> int:
         """Create one run."""
         # Sample a random int64 as run_id
-        run_id: int = int.from_bytes(os.urandom(8), "little", signed=True)
+        with self.lock:
+            run_id: int = int.from_bytes(os.urandom(8), "little", signed=True)
 
-        if run_id not in self.run_ids:
-            self.run_ids.add(run_id)
-            return run_id
+            if run_id not in self.run_ids:
+                self.run_ids.add(run_id)
+                return run_id
         log(ERROR, "Unexpected run creation failure.")
         return 0
+
+    def acknowledge_ping(self, node_id: int, ping_interval: float) -> bool:
+        """Acknowledge a ping received from a node, serving as a heartbeat."""
+        with self.lock:
+            if node_id in self.node_ids:
+                self.node_ids[node_id] = (time.time() + ping_interval, ping_interval)
+                return True
+        return False

flwr/server/superlink/state/sqlite_state.py

@@ -18,7 +18,7 @@
 import os
 import re
 import sqlite3
-from datetime import datetime, timedelta
+import time
 from logging import DEBUG, ERROR
 from typing import Any, Dict, List, Optional, Set, Tuple, Union, cast
 from uuid import UUID, uuid4
@@ -33,10 +33,16 @@ from .state import State
 
 SQL_CREATE_TABLE_NODE = """
 CREATE TABLE IF NOT EXISTS node(
-    node_id INTEGER UNIQUE
+    node_id         INTEGER UNIQUE,
+    online_until    REAL,
+    ping_interval   REAL
 );
 """
 
+SQL_CREATE_INDEX_ONLINE_UNTIL = """
+CREATE INDEX IF NOT EXISTS idx_online_until ON node (online_until);
+"""
+
 SQL_CREATE_TABLE_RUN = """
 CREATE TABLE IF NOT EXISTS run(
     run_id INTEGER UNIQUE
@@ -52,9 +58,10 @@ CREATE TABLE IF NOT EXISTS task_ins(
     producer_node_id        INTEGER,
     consumer_anonymous      BOOLEAN,
     consumer_node_id        INTEGER,
-    created_at              TEXT,
+    created_at              REAL,
     delivered_at            TEXT,
-    ttl                     TEXT,
+    pushed_at               REAL,
+    ttl                     REAL,
     ancestry                TEXT,
     task_type               TEXT,
     recordset               BLOB,
@@ -72,9 +79,10 @@ CREATE TABLE IF NOT EXISTS task_res(
     producer_node_id        INTEGER,
     consumer_anonymous      BOOLEAN,
     consumer_node_id        INTEGER,
-    created_at              TEXT,
+    created_at              REAL,
     delivered_at            TEXT,
-    ttl                     TEXT,
+    pushed_at               REAL,
+    ttl                     REAL,
     ancestry                TEXT,
     task_type               TEXT,
     recordset               BLOB,
@@ -82,7 +90,7 @@ CREATE TABLE IF NOT EXISTS task_res(
 );
 """
 
-DictOrTuple = Union[Tuple[Any], Dict[str, Any]]
+DictOrTuple = Union[Tuple[Any, ...], Dict[str, Any]]
 
 
 class SqliteState(State):
@@ -123,6 +131,7 @@ class SqliteState(State):
         cur.execute(SQL_CREATE_TABLE_TASK_INS)
         cur.execute(SQL_CREATE_TABLE_TASK_RES)
         cur.execute(SQL_CREATE_TABLE_NODE)
+        cur.execute(SQL_CREATE_INDEX_ONLINE_UNTIL)
         res = cur.execute("SELECT name FROM sqlite_schema;")
 
         return res.fetchall()
@@ -185,15 +194,11 @@ class SqliteState(State):
             log(ERROR, errors)
             return None
 
-        # Create task_id, created_at and ttl
+        # Create task_id
         task_id = uuid4()
-        created_at: datetime = now()
-        ttl: datetime = created_at + timedelta(hours=24)
 
         # Store TaskIns
         task_ins.task_id = str(task_id)
-        task_ins.task.created_at = created_at.isoformat()
-        task_ins.task.ttl = ttl.isoformat()
         data = (task_ins_to_dict(task_ins),)
         columns = ", ".join([f":{key}" for key in data[0]])
         query = f"INSERT INTO task_ins VALUES({columns});"
@@ -320,15 +325,11 @@ class SqliteState(State):
             log(ERROR, errors)
             return None
 
-        # Create task_id, created_at and ttl
+        # Create task_id
         task_id = uuid4()
-        created_at: datetime = now()
-        ttl: datetime = created_at + timedelta(hours=24)
 
         # Store TaskIns
         task_res.task_id = str(task_id)
-        task_res.task.created_at = created_at.isoformat()
-        task_res.task.ttl = ttl.isoformat()
         data = (task_res_to_dict(task_res),)
         columns = ", ".join([f":{key}" for key in data[0]])
         query = f"INSERT INTO task_res VALUES({columns});"
@@ -472,9 +473,14 @@ class SqliteState(State):
         # Sample a random int64 as node_id
         node_id: int = int.from_bytes(os.urandom(8), "little", signed=True)
 
-        query = "INSERT INTO node VALUES(:node_id);"
+        query = (
+            "INSERT INTO node (node_id, online_until, ping_interval) VALUES (?, ?, ?)"
+        )
+
         try:
-            self.query(query, {"node_id": node_id})
+            # Default ping interval is 30s
+            # TODO: change 1e9 to 30s  # pylint: disable=W0511
+            self.query(query, (node_id, time.time() + 1e9, 1e9))
         except sqlite3.IntegrityError:
             log(ERROR, "Unexpected node registration failure.")
             return 0
@@ -499,8 +505,8 @@ class SqliteState(State):
             return set()
 
         # Get nodes
-        query = "SELECT * FROM node;"
-        rows = self.query(query)
+        query = "SELECT node_id FROM node WHERE online_until > ?;"
+        rows = self.query(query, (time.time(),))
         result: Set[int] = {row["node_id"] for row in rows}
         return result
 
@@ -519,6 +525,17 @@ class SqliteState(State):
         log(ERROR, "Unexpected run creation failure.")
         return 0
 
+    def acknowledge_ping(self, node_id: int, ping_interval: float) -> bool:
+        """Acknowledge a ping received from a node, serving as a heartbeat."""
+        # Update `online_until` and `ping_interval` for the given `node_id`
+        query = "UPDATE node SET online_until = ?, ping_interval = ? WHERE node_id = ?;"
+        try:
+            self.query(query, (time.time() + ping_interval, ping_interval, node_id))
+            return True
+        except sqlite3.IntegrityError:
+            log(ERROR, "`node_id` does not exist.")
+            return False
+
 
 def dict_factory(
     cursor: sqlite3.Cursor,
@@ -544,6 +561,7 @@ def task_ins_to_dict(task_msg: TaskIns) -> Dict[str, Any]:
         "consumer_node_id": task_msg.task.consumer.node_id,
         "created_at": task_msg.task.created_at,
         "delivered_at": task_msg.task.delivered_at,
+        "pushed_at": task_msg.task.pushed_at,
         "ttl": task_msg.task.ttl,
         "ancestry": ",".join(task_msg.task.ancestry),
         "task_type": task_msg.task.task_type,
@@ -564,6 +582,7 @@ def task_res_to_dict(task_msg: TaskRes) -> Dict[str, Any]:
         "consumer_node_id": task_msg.task.consumer.node_id,
         "created_at": task_msg.task.created_at,
         "delivered_at": task_msg.task.delivered_at,
+        "pushed_at": task_msg.task.pushed_at,
         "ttl": task_msg.task.ttl,
         "ancestry": ",".join(task_msg.task.ancestry),
         "task_type": task_msg.task.task_type,
@@ -592,6 +611,7 @@ def dict_to_task_ins(task_dict: Dict[str, Any]) -> TaskIns:
             ),
             created_at=task_dict["created_at"],
             delivered_at=task_dict["delivered_at"],
+            pushed_at=task_dict["pushed_at"],
             ttl=task_dict["ttl"],
             ancestry=task_dict["ancestry"].split(","),
             task_type=task_dict["task_type"],
@@ -621,6 +641,7 @@ def dict_to_task_res(task_dict: Dict[str, Any]) -> TaskRes:
             ),
             created_at=task_dict["created_at"],
             delivered_at=task_dict["delivered_at"],
+            pushed_at=task_dict["pushed_at"],
             ttl=task_dict["ttl"],
             ancestry=task_dict["ancestry"].split(","),
             task_type=task_dict["task_type"],

flwr/server/superlink/state/state.py

@@ -152,3 +152,22 @@ class State(abc.ABC):
     @abc.abstractmethod
     def create_run(self) -> int:
         """Create one run."""
+
+    @abc.abstractmethod
+    def acknowledge_ping(self, node_id: int, ping_interval: float) -> bool:
+        """Acknowledge a ping received from a node, serving as a heartbeat.
+
+        Parameters
+        ----------
+        node_id : int
+            The `node_id` from which the ping was received.
+        ping_interval : float
+            The interval (in seconds) from the current timestamp within which the next
+            ping from this node must be received. This acts as a hard deadline to ensure
+            an accurate assessment of the node's availability.
+
+        Returns
+        -------
+        is_acknowledged : bool
+            True if the ping is successfully acknowledged; otherwise, False.
+        """