flwr 1.13.0__py3-none-any.whl → 1.14.0__py3-none-any.whl

flwr/server/superlink/fleet/grpc_bidi/flower_service_servicer.py

@@ -18,6 +18,7 @@ Relevant knowledge for reading this modules code:
 - https://github.com/grpc/grpc/blob/master/doc/statuscodes.md
 """
 
+
 import uuid
 from collections.abc import Iterator
 from typing import Callable

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

@@ -20,6 +20,7 @@ from logging import DEBUG, INFO
 import grpc
 
 from flwr.common.logger import log
+from flwr.common.typing import InvalidRunStatusException
 from flwr.proto import fleet_pb2_grpc  # pylint: disable=E0611
 from flwr.proto.fab_pb2 import GetFabRequest, GetFabResponse  # pylint: disable=E0611
 from flwr.proto.fleet_pb2 import (  # pylint: disable=E0611
@@ -38,6 +39,7 @@ from flwr.proto.run_pb2 import GetRunRequest, GetRunResponse  # pylint: disable=
 from flwr.server.superlink.ffs.ffs_factory import FfsFactory
 from flwr.server.superlink.fleet.message_handler import message_handler
 from flwr.server.superlink.linkstate import LinkStateFactory
+from flwr.server.superlink.utils import abort_grpc_context
 
 
 class FleetServicer(fleet_pb2_grpc.FleetServicer):
@@ -105,27 +107,45 @@ class FleetServicer(fleet_pb2_grpc.FleetServicer):
             )
         else:
             log(INFO, "[Fleet.PushTaskRes] No task results to push")
-        return message_handler.push_task_res(
-            request=request,
-            state=self.state_factory.state(),
-        )
+
+        try:
+            res = message_handler.push_task_res(
+                request=request,
+                state=self.state_factory.state(),
+            )
+        except InvalidRunStatusException as e:
+            abort_grpc_context(e.message, context)
+
+        return res
 
     def GetRun(
         self, request: GetRunRequest, context: grpc.ServicerContext
     ) -> GetRunResponse:
         """Get run information."""
         log(INFO, "[Fleet.GetRun] Requesting `Run` for run_id=%s", request.run_id)
-        return message_handler.get_run(
-            request=request,
-            state=self.state_factory.state(),
-        )
+
+        try:
+            res = message_handler.get_run(
+                request=request,
+                state=self.state_factory.state(),
+            )
+        except InvalidRunStatusException as e:
+            abort_grpc_context(e.message, context)
+
+        return res
 
     def GetFab(
         self, request: GetFabRequest, context: grpc.ServicerContext
     ) -> GetFabResponse:
         """Get FAB."""
         log(INFO, "[Fleet.GetFab] Requesting FAB for fab_hash=%s", request.hash_str)
-        return message_handler.get_fab(
-            request=request,
-            ffs=self.ffs_factory.ffs(),
-        )
+        try:
+            res = message_handler.get_fab(
+                request=request,
+                ffs=self.ffs_factory.ffs(),
+                state=self.state_factory.state(),
+            )
+        except InvalidRunStatusException as e:
+            abort_grpc_context(e.message, context)
+
+        return res

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

@@ -45,7 +45,7 @@ from flwr.proto.fleet_pb2 import (  # pylint: disable=E0611
 )
 from flwr.proto.node_pb2 import Node  # pylint: disable=E0611
 from flwr.proto.run_pb2 import GetRunRequest, GetRunResponse  # pylint: disable=E0611
-from flwr.server.superlink.linkstate import LinkState
+from flwr.server.superlink.linkstate import LinkStateFactory
 
 _PUBLIC_KEY_HEADER = "public-key"
 _AUTH_TOKEN_HEADER = "auth-token"
@@ -84,15 +84,16 @@ def _get_value_from_tuples(
 class AuthenticateServerInterceptor(grpc.ServerInterceptor):  # type: ignore
     """Server interceptor for node authentication."""
 
-    def __init__(self, state: LinkState):
-        self.state = state
+    def __init__(self, state_factory: LinkStateFactory):
+        self.state_factory = state_factory
+        state = self.state_factory.state()
 
         self.node_public_keys = state.get_node_public_keys()
         if len(self.node_public_keys) == 0:
             log(WARNING, "Authentication enabled, but no known public keys configured")
 
-        private_key = self.state.get_server_private_key()
-        public_key = self.state.get_server_public_key()
+        private_key = state.get_server_private_key()
+        public_key = state.get_server_public_key()
 
         if private_key is None or public_key is None:
             raise ValueError("Error loading authentication keys")
@@ -154,7 +155,7 @@ class AuthenticateServerInterceptor(grpc.ServerInterceptor):  # type: ignore
                 context.abort(grpc.StatusCode.UNAUTHENTICATED, "Access denied")
 
             # Verify node_id
-            node_id = self.state.get_node_id(node_public_key_bytes)
+            node_id = self.state_factory.state().get_node_id(node_public_key_bytes)
 
             if not self._verify_node_id(node_id, request):
                 context.abort(grpc.StatusCode.UNAUTHENTICATED, "Access denied")
@@ -186,7 +187,7 @@ class AuthenticateServerInterceptor(grpc.ServerInterceptor):  # type: ignore
                 return False
             return request.task_res_list[0].task.producer.node_id == node_id
         if isinstance(request, GetRunRequest):
-            return node_id in self.state.get_nodes(request.run_id)
+            return node_id in self.state_factory.state().get_nodes(request.run_id)
         return request.node.node_id == node_id
 
     def _verify_hmac(
@@ -210,17 +211,17 @@ class AuthenticateServerInterceptor(grpc.ServerInterceptor):  # type: ignore
                 ),
             )
         )
-
-        node_id = self.state.get_node_id(public_key_bytes)
+        state = self.state_factory.state()
+        node_id = state.get_node_id(public_key_bytes)
 
         # Handle `CreateNode` here instead of calling the default method handler
         # Return previously assigned `node_id` for the provided `public_key`
         if node_id is not None:
-            self.state.acknowledge_ping(node_id, request.ping_interval)
+            state.acknowledge_ping(node_id, request.ping_interval)
             return CreateNodeResponse(node=Node(node_id=node_id, anonymous=False))
 
         # No `node_id` exists for the provided `public_key`
         # Handle `CreateNode` here instead of calling the default method handler
         # Note: the innermost `CreateNode` method will never be called
-        node_id = self.state.create_node(request.ping_interval, public_key_bytes)
+        node_id = state.create_node(request.ping_interval, public_key_bytes)
         return CreateNodeResponse(node=Node(node_id=node_id, anonymous=False))

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

@@ -19,8 +19,9 @@ import time
 from typing import Optional
 from uuid import UUID
 
+from flwr.common.constant import Status
 from flwr.common.serde import fab_to_proto, user_config_to_proto
-from flwr.common.typing import Fab
+from flwr.common.typing import Fab, InvalidRunStatusException
 from flwr.proto.fab_pb2 import GetFabRequest, GetFabResponse  # pylint: disable=E0611
 from flwr.proto.fleet_pb2 import (  # pylint: disable=E0611
     CreateNodeRequest,
@@ -44,6 +45,7 @@ from flwr.proto.run_pb2 import (  # pylint: disable=E0611
 from flwr.proto.task_pb2 import TaskIns, TaskRes  # pylint: disable=E0611
 from flwr.server.superlink.ffs.ffs import Ffs
 from flwr.server.superlink.linkstate import LinkState
+from flwr.server.superlink.utils import check_abort
 
 
 def create_node(
@@ -98,6 +100,15 @@ def push_task_res(request: PushTaskResRequest, state: LinkState) -> PushTaskResR
     task_res: TaskRes = request.task_res_list[0]
     # pylint: enable=no-member
 
+    # Abort if the run is not running
+    abort_msg = check_abort(
+        task_res.run_id,
+        [Status.PENDING, Status.STARTING, Status.FINISHED],
+        state,
+    )
+    if abort_msg:
+        raise InvalidRunStatusException(abort_msg)
+
     # Set pushed_at (timestamp in seconds)
     task_res.task.pushed_at = time.time()
 
@@ -112,15 +123,22 @@ def push_task_res(request: PushTaskResRequest, state: LinkState) -> PushTaskResR
     return response
 
 
-def get_run(
-    request: GetRunRequest, state: LinkState  # pylint: disable=W0613
-) -> GetRunResponse:
+def get_run(request: GetRunRequest, state: LinkState) -> GetRunResponse:
     """Get run information."""
     run = state.get_run(request.run_id)
 
     if run is None:
         return GetRunResponse()
 
+    # Abort if the run is not running
+    abort_msg = check_abort(
+        request.run_id,
+        [Status.PENDING, Status.STARTING, Status.FINISHED],
+        state,
+    )
+    if abort_msg:
+        raise InvalidRunStatusException(abort_msg)
+
     return GetRunResponse(
         run=Run(
             run_id=run.run_id,
@@ -133,9 +151,18 @@ def get_run(
 
 
 def get_fab(
-    request: GetFabRequest, ffs: Ffs  # pylint: disable=W0613
+    request: GetFabRequest, ffs: Ffs, state: LinkState  # pylint: disable=W0613
 ) -> GetFabResponse:
     """Get FAB."""
+    # Abort if the run is not running
+    abort_msg = check_abort(
+        request.run_id,
+        [Status.PENDING, Status.STARTING, Status.FINISHED],
+        state,
+    )
+    if abort_msg:
+        raise InvalidRunStatusException(abort_msg)
+
     if result := ffs.get(request.hash_str):
         fab = Fab(request.hash_str, result[0])
         return GetFabResponse(fab=fab_to_proto(fab))

flwr/server/superlink/fleet/rest_rere/rest_api.py

@@ -154,8 +154,11 @@ async def get_fab(request: GetFabRequest) -> GetFabResponse:
     # 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)
+    return message_handler.get_fab(request=request, ffs=ffs, state=state)
 
 
 routes = [

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,6 +14,7 @@
 # ==============================================================================
 """Simulation Engine Backends."""
 
+
 import importlib
 
 from .backend import Backend, BackendConfig

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

@@ -14,6 +14,7 @@
 # ==============================================================================
 """Ray backend for the Fleet API using the Simulation Engine."""
 
+
 import sys
 from logging import DEBUG, ERROR
 from typing import Callable, Optional, Union

flwr/server/superlink/linkstate/in_memory_linkstate.py

@@ -265,41 +265,15 @@ class InMemoryLinkState(LinkState):  # pylint: disable=R0902,R0904
             for task_res in task_res_found:
                 task_res.task.delivered_at = delivered_at
 
-            # Cleanup
-            self._force_delete_tasks_by_ids(set(ret.keys()))
-
         return list(ret.values())
 
-    def delete_tasks(self, task_ids: set[UUID]) -> None:
-        """Delete all delivered TaskIns/TaskRes pairs."""
-        task_ins_to_be_deleted: set[UUID] = set()
-        task_res_to_be_deleted: set[UUID] = set()
-
-        with self.lock:
-            for task_ins_id in task_ids:
-                # Find the task_id of the matching task_res
-                for task_res_id, task_res in self.task_res_store.items():
-                    if UUID(task_res.task.ancestry[0]) != task_ins_id:
-                        continue
-                    if task_res.task.delivered_at == "":
-                        continue
-
-                    task_ins_to_be_deleted.add(task_ins_id)
-                    task_res_to_be_deleted.add(task_res_id)
-
-            for task_id in task_ins_to_be_deleted:
-                del self.task_ins_store[task_id]
-                del self.task_ins_id_to_task_res_id[task_id]
-            for task_id in task_res_to_be_deleted:
-                del self.task_res_store[task_id]
-
-    def _force_delete_tasks_by_ids(self, task_ids: set[UUID]) -> None:
-        """Delete tasks based on a set of TaskIns IDs."""
-        if not task_ids:
+    def delete_tasks(self, task_ins_ids: set[UUID]) -> None:
+        """Delete TaskIns/TaskRes pairs based on provided TaskIns IDs."""
+        if not task_ins_ids:
             return
 
         with self.lock:
-            for task_id in task_ids:
+            for task_id in task_ins_ids:
                 # Delete TaskIns
                 if task_id in self.task_ins_store:
                     del self.task_ins_store[task_id]
@@ -308,6 +282,16 @@ class InMemoryLinkState(LinkState):  # pylint: disable=R0902,R0904
                     task_res_id = self.task_ins_id_to_task_res_id.pop(task_id)
                     del self.task_res_store[task_res_id]
 
+    def get_task_ids_from_run_id(self, run_id: int) -> set[UUID]:
+        """Get all TaskIns IDs for the given run_id."""
+        task_id_list: set[UUID] = set()
+        with self.lock:
+            for task_id, task_ins in self.task_ins_store.items():
+                if task_ins.run_id == run_id:
+                    task_id_list.add(task_id)
+
+        return task_id_list
+
     def num_task_ins(self) -> int:
         """Calculate the number of task_ins in store.
 
@@ -446,6 +430,13 @@ class InMemoryLinkState(LinkState):  # pylint: disable=R0902,R0904
         """Retrieve `server_public_key` in urlsafe bytes."""
         return self.server_public_key
 
+    def clear_supernode_auth_keys_and_credentials(self) -> None:
+        """Clear stored `node_public_keys` and credentials in the link state if any."""
+        with self.lock:
+            self.server_private_key = None
+            self.server_public_key = None
+            self.node_public_keys.clear()
+
     def store_node_public_keys(self, public_keys: set[bytes]) -> None:
         """Store a set of `node_public_keys` in the link state."""
         with self.lock:

flwr/server/superlink/linkstate/linkstate.py

@@ -139,8 +139,19 @@ class LinkState(abc.ABC):  # pylint: disable=R0904
         """
 
     @abc.abstractmethod
-    def delete_tasks(self, task_ids: set[UUID]) -> None:
-        """Delete all delivered TaskIns/TaskRes pairs."""
+    def delete_tasks(self, task_ins_ids: set[UUID]) -> None:
+        """Delete TaskIns/TaskRes pairs based on provided TaskIns IDs.
+
+        Parameters
+        ----------
+        task_ins_ids : set[UUID]
+            A set of TaskIns IDs. For each ID in the set, the corresponding
+            TaskIns and its associated TaskRes will be deleted.
+        """
+
+    @abc.abstractmethod
+    def get_task_ids_from_run_id(self, run_id: int) -> set[UUID]:
+        """Get all TaskIns IDs for the given run_id."""
 
     @abc.abstractmethod
     def create_node(
@@ -273,6 +284,10 @@ class LinkState(abc.ABC):  # pylint: disable=R0904
     def get_server_public_key(self) -> Optional[bytes]:
         """Retrieve `server_public_key` in urlsafe bytes."""
 
+    @abc.abstractmethod
+    def clear_supernode_auth_keys_and_credentials(self) -> None:
+        """Clear stored `node_public_keys` and credentials in the link state if any."""
+
     @abc.abstractmethod
     def store_node_public_keys(self, public_keys: set[bytes]) -> None:
         """Store a set of `node_public_keys` in the link state."""

flwr/server/superlink/linkstate/sqlite_linkstate.py

@@ -14,12 +14,12 @@
 # ==============================================================================
 """SQLite based implemenation of the link state."""
 
+
 # pylint: disable=too-many-lines
 
 import json
 import re
 import sqlite3
-import threading
 import time
 from collections.abc import Sequence
 from logging import DEBUG, ERROR, WARNING
@@ -183,7 +183,6 @@ class SqliteLinkState(LinkState):  # pylint: disable=R0904
         """
         self.database_path = database_path
         self.conn: Optional[sqlite3.Connection] = None
-        self.lock = threading.RLock()
 
     def initialize(self, log_queries: bool = False) -> list[tuple[str]]:
         """Create tables if they don't exist yet.
@@ -216,7 +215,6 @@ class SqliteLinkState(LinkState):  # pylint: disable=R0904
         cur.execute(SQL_CREATE_TABLE_PUBLIC_KEY)
         cur.execute(SQL_CREATE_INDEX_ONLINE_UNTIL)
         res = cur.execute("SELECT name FROM sqlite_schema;")
-
         return res.fetchall()
 
     def query(
@@ -569,9 +567,6 @@ class SqliteLinkState(LinkState):  # pylint: disable=R0904
         data: list[Any] = [delivered_at] + task_res_ids
         self.query(query, data)
 
-        # Cleanup
-        self._force_delete_tasks_by_ids(set(ret.keys()))
-
         return list(ret.values())
 
     def num_task_ins(self) -> int:
@@ -595,68 +590,50 @@ class SqliteLinkState(LinkState):  # pylint: disable=R0904
         result: dict[str, int] = rows[0]
         return result["num"]
 
-    def delete_tasks(self, task_ids: set[UUID]) -> None:
-        """Delete all delivered TaskIns/TaskRes pairs."""
-        ids = list(task_ids)
-        if len(ids) == 0:
-            return None
+    def delete_tasks(self, task_ins_ids: set[UUID]) -> None:
+        """Delete TaskIns/TaskRes pairs based on provided TaskIns IDs."""
+        if not task_ins_ids:
+            return
+        if self.conn is None:
+            raise AttributeError("LinkState not initialized")
 
-        placeholders = ",".join([f":id_{index}" for index in range(len(task_ids))])
-        data = {f"id_{index}": str(task_id) for index, task_id in enumerate(task_ids)}
+        placeholders = ",".join(["?"] * len(task_ins_ids))
+        data = tuple(str(task_id) for task_id in task_ins_ids)
 
-        # 1. Query: Delete task_ins which have a delivered task_res
+        # Delete task_ins
         query_1 = f"""
             DELETE FROM task_ins
-            WHERE delivered_at != ''
-            AND task_id IN (
-                SELECT ancestry
-                FROM task_res
-                WHERE ancestry IN ({placeholders})
-                AND delivered_at != ''
-            );
+            WHERE task_id IN ({placeholders});
         """
 
-        # 2. Query: Delete delivered task_res to be run after 1. Query
+        # Delete task_res
         query_2 = f"""
             DELETE FROM task_res
-            WHERE ancestry IN ({placeholders})
-            AND delivered_at != '';
+            WHERE ancestry IN ({placeholders});
         """
 
-        if self.conn is None:
-            raise AttributeError("LinkState not intitialized")
-
         with self.conn:
             self.conn.execute(query_1, data)
             self.conn.execute(query_2, data)
 
-        return None
-
-    def _force_delete_tasks_by_ids(self, task_ids: set[UUID]) -> None:
-        """Delete tasks based on a set of TaskIns IDs."""
-        if not task_ids:
-            return
+    def get_task_ids_from_run_id(self, run_id: int) -> set[UUID]:
+        """Get all TaskIns IDs for the given run_id."""
         if self.conn is None:
             raise AttributeError("LinkState not initialized")
 
-        placeholders = ",".join([f":id_{index}" for index in range(len(task_ids))])
-        data = {f"id_{index}": str(task_id) for index, task_id in enumerate(task_ids)}
-
-        # Delete task_ins
-        query_1 = f"""
-            DELETE FROM task_ins
-            WHERE task_id IN ({placeholders});
+        query = """
+            SELECT task_id
+            FROM task_ins
+            WHERE run_id = :run_id;
         """
 
-        # Delete task_res
-        query_2 = f"""
-            DELETE FROM task_res
-            WHERE ancestry IN ({placeholders});
-        """
+        sint64_run_id = convert_uint64_to_sint64(run_id)
+        data = {"run_id": sint64_run_id}
 
         with self.conn:
-            self.conn.execute(query_1, data)
-            self.conn.execute(query_2, data)
+            rows = self.conn.execute(query, data).fetchall()
+
+        return {UUID(row["task_id"]) for row in rows}
 
     def create_node(
         self, ping_interval: float, public_key: Optional[bytes] = None
@@ -784,8 +761,6 @@ class SqliteLinkState(LinkState):  # pylint: disable=R0904
                 "federation_options, pending_at, starting_at, running_at, finished_at, "
                 "sub_status, details) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?);"
             )
-            if fab_hash:
-                fab_id, fab_version = "", ""
             override_config_json = json.dumps(override_config)
             data = [
                 sint64_run_id,
@@ -843,6 +818,12 @@ class SqliteLinkState(LinkState):  # pylint: disable=R0904
             public_key = None
         return public_key
 
+    def clear_supernode_auth_keys_and_credentials(self) -> None:
+        """Clear stored `node_public_keys` and credentials in the link state if any."""
+        queries = ["DELETE FROM public_key;", "DELETE FROM credential;"]
+        for query in queries:
+            self.query(query)
+
     def store_node_public_keys(self, public_keys: set[bytes]) -> None:
         """Store a set of `node_public_keys` in the link state."""
         query = "INSERT INTO public_key (public_key) VALUES (?)"

flwr/server/superlink/simulation/simulationio_servicer.py

@@ -14,6 +14,7 @@
 # ==============================================================================
 """SimulationIo API servicer."""
 
+
 import threading
 from logging import DEBUG, INFO
 
@@ -28,6 +29,7 @@ from flwr.common.serde import (
     context_to_proto,
     fab_to_proto,
     run_status_from_proto,
+    run_status_to_proto,
     run_to_proto,
 )
 from flwr.common.typing import Fab, RunStatus
@@ -39,6 +41,8 @@ from flwr.proto.log_pb2 import (  # pylint: disable=E0611
 from flwr.proto.run_pb2 import (  # pylint: disable=E0611
     GetFederationOptionsRequest,
     GetFederationOptionsResponse,
+    GetRunStatusRequest,
+    GetRunStatusResponse,
     UpdateRunStatusRequest,
     UpdateRunStatusResponse,
 )
@@ -50,6 +54,7 @@ from flwr.proto.simulationio_pb2 import (  # pylint: disable=E0611
 )
 from flwr.server.superlink.ffs.ffs_factory import FfsFactory
 from flwr.server.superlink.linkstate import LinkStateFactory
+from flwr.server.superlink.utils import abort_if
 
 
 class SimulationIoServicer(simulationio_pb2_grpc.SimulationIoServicer):
@@ -106,6 +111,15 @@ class SimulationIoServicer(simulationio_pb2_grpc.SimulationIoServicer):
         """Push Simulation process outputs."""
         log(DEBUG, "SimultionIoServicer.PushSimulationOutputs")
         state = self.state_factory.state()
+
+        # Abort if the run is not running
+        abort_if(
+            request.run_id,
+            [Status.PENDING, Status.STARTING, Status.FINISHED],
+            state,
+            context,
+        )
+
         state.set_serverapp_context(request.run_id, context_from_proto(request.context))
         return PushSimulationOutputsResponse()
 
@@ -116,12 +130,31 @@ class SimulationIoServicer(simulationio_pb2_grpc.SimulationIoServicer):
         log(DEBUG, "SimultionIoServicer.UpdateRunStatus")
         state = self.state_factory.state()
 
+        # Abort if the run is finished
+        abort_if(request.run_id, [Status.FINISHED], state, context)
+
         # Update the run status
         state.update_run_status(
             run_id=request.run_id, new_status=run_status_from_proto(request.run_status)
         )
         return UpdateRunStatusResponse()
 
+    def GetRunStatus(
+        self, request: GetRunStatusRequest, context: ServicerContext
+    ) -> GetRunStatusResponse:
+        """Get status of requested runs."""
+        log(DEBUG, "SimultionIoServicer.GetRunStatus")
+        state = self.state_factory.state()
+
+        statuses = state.get_run_status(set(request.run_ids))
+
+        return GetRunStatusResponse(
+            run_status_dict={
+                run_id: run_status_to_proto(status)
+                for run_id, status in statuses.items()
+            }
+        )
+
     def PushLogs(
         self, request: PushLogsRequest, context: grpc.ServicerContext
     ) -> PushLogsResponse: