flwr-nightly 1.14.0.dev20241204__py3-none-any.whl → 1.14.0.dev20241214__py3-none-any.whl

flwr/server/serverapp/app.py

@@ -14,6 +14,7 @@
 # ==============================================================================
 """Flower ServerApp process."""
 
+
 import argparse
 import sys
 from logging import DEBUG, ERROR, INFO
@@ -50,7 +51,7 @@ from flwr.common.serde import (
     run_from_proto,
     run_status_to_proto,
 )
-from flwr.common.typing import RunStatus
+from flwr.common.typing import RunNotRunningException, RunStatus
 from flwr.proto.run_pb2 import UpdateRunStatusRequest  # pylint: disable=E0611
 from flwr.proto.serverappio_pb2 import (  # pylint: disable=E0611
     PullServerAppInputsRequest,
@@ -96,7 +97,7 @@ def flwr_serverapp() -> None:
     restore_output()
 
 
-def run_serverapp(  # pylint: disable=R0914, disable=W0212
+def run_serverapp(  # pylint: disable=R0914, disable=W0212, disable=R0915
     serverappio_api_address: str,
     log_queue: Queue[Optional[str]],
     run_once: bool,
@@ -187,6 +188,12 @@ def run_serverapp(  # pylint: disable=R0914, disable=W0212
 
             run_status = RunStatus(Status.FINISHED, SubStatus.COMPLETED, "")
 
+        except RunNotRunningException:
+            log(INFO, "")
+            log(INFO, "Run ID %s stopped.", run.run_id)
+            log(INFO, "")
+            run_status = None
+
         except Exception as ex:  # pylint: disable=broad-exception-caught
             exc_entity = "ServerApp"
             log(ERROR, "%s raised an exception", exc_entity, exc_info=ex)

flwr/server/strategy/dpfedavg_fixed.py

@@ -17,6 +17,7 @@
 Paper: arxiv.org/pdf/1710.06963.pdf
 """
 
+
 from typing import Optional, Union
 
 from flwr.common import EvaluateIns, EvaluateRes, FitIns, FitRes, Parameters, Scalar

flwr/server/superlink/driver/serverappio_grpc.py

@@ -14,6 +14,7 @@
 # ==============================================================================
 """ServerAppIo gRPC API."""
 
+
 from logging import INFO
 from typing import Optional
 

flwr/server/superlink/driver/serverappio_servicer.py

@@ -32,6 +32,7 @@ from flwr.common.serde import (
     fab_from_proto,
     fab_to_proto,
     run_status_from_proto,
+    run_status_to_proto,
     run_to_proto,
     user_config_from_proto,
 )
@@ -48,6 +49,8 @@ from flwr.proto.run_pb2 import (  # pylint: disable=E0611
     CreateRunResponse,
     GetRunRequest,
     GetRunResponse,
+    GetRunStatusRequest,
+    GetRunStatusResponse,
     UpdateRunStatusRequest,
     UpdateRunStatusResponse,
 )
@@ -67,6 +70,7 @@ from flwr.proto.task_pb2 import TaskRes  # pylint: disable=E0611
 from flwr.server.superlink.ffs.ffs import Ffs
 from flwr.server.superlink.ffs.ffs_factory import FfsFactory
 from flwr.server.superlink.linkstate import LinkState, LinkStateFactory
+from flwr.server.superlink.utils import abort_if
 from flwr.server.utils.validator import validate_task_ins_or_res
 
 
@@ -85,7 +89,18 @@ class ServerAppIoServicer(serverappio_pb2_grpc.ServerAppIoServicer):
     ) -> GetNodesResponse:
         """Get available nodes."""
         log(DEBUG, "ServerAppIoServicer.GetNodes")
+
+        # Init state
         state: LinkState = self.state_factory.state()
+
+        # Abort if the run is not running
+        abort_if(
+            request.run_id,
+            [Status.PENDING, Status.STARTING, Status.FINISHED],
+            state,
+            context,
+        )
+
         all_ids: set[int] = state.get_nodes(request.run_id)
         nodes: list[Node] = [
             Node(node_id=node_id, anonymous=False) for node_id in all_ids
@@ -123,6 +138,17 @@ class ServerAppIoServicer(serverappio_pb2_grpc.ServerAppIoServicer):
         """Push a set of TaskIns."""
         log(DEBUG, "ServerAppIoServicer.PushTaskIns")
 
+        # Init state
+        state: LinkState = self.state_factory.state()
+
+        # Abort if the run is not running
+        abort_if(
+            request.run_id,
+            [Status.PENDING, Status.STARTING, Status.FINISHED],
+            state,
+            context,
+        )
+
         # Set pushed_at (timestamp in seconds)
         pushed_at = time.time()
         for task_ins in request.task_ins_list:
@@ -134,9 +160,6 @@ class ServerAppIoServicer(serverappio_pb2_grpc.ServerAppIoServicer):
             validation_errors = validate_task_ins_or_res(task_ins)
             _raise_if(bool(validation_errors), ", ".join(validation_errors))
 
-        # Init state
-        state: LinkState = self.state_factory.state()
-
         # Store each TaskIns
         task_ids: list[Optional[UUID]] = []
         for task_ins in request.task_ins_list:
@@ -153,33 +176,29 @@ class ServerAppIoServicer(serverappio_pb2_grpc.ServerAppIoServicer):
         """Pull a set of TaskRes."""
         log(DEBUG, "ServerAppIoServicer.PullTaskRes")
 
-        # Convert each task_id str to UUID
-        task_ids: set[UUID] = {UUID(task_id) for task_id in request.task_ids}
-
         # Init state
         state: LinkState = self.state_factory.state()
 
-        # Register callback
-        def on_rpc_done() -> None:
-            log(
-                DEBUG,
-                "ServerAppIoServicer.PullTaskRes callback: delete TaskIns/TaskRes",
-            )
-
-            if context.is_active():
-                return
-            if context.code() != grpc.StatusCode.OK:
-                return
-
-            # Delete delivered TaskIns and TaskRes
-            state.delete_tasks(task_ids=task_ids)
+        # Abort if the run is not running
+        abort_if(
+            request.run_id,
+            [Status.PENDING, Status.STARTING, Status.FINISHED],
+            state,
+            context,
+        )
 
-        context.add_callback(on_rpc_done)
+        # Convert each task_id str to UUID
+        task_ids: set[UUID] = {UUID(task_id) for task_id in request.task_ids}
 
         # Read from state
         task_res_list: list[TaskRes] = state.get_task_res(task_ids=task_ids)
 
-        context.set_code(grpc.StatusCode.OK)
+        # Delete the TaskIns/TaskRes pairs if TaskRes is found
+        task_ins_ids_to_delete = {
+            UUID(task_res.task.ancestry[0]) for task_res in task_res_list
+        }
+        state.delete_tasks(task_ins_ids=task_ins_ids_to_delete)
+
         return PullTaskResResponse(task_res_list=task_res_list)
 
     def GetRun(
@@ -255,7 +274,18 @@ class ServerAppIoServicer(serverappio_pb2_grpc.ServerAppIoServicer):
     ) -> PushServerAppOutputsResponse:
         """Push ServerApp process outputs."""
         log(DEBUG, "ServerAppIoServicer.PushServerAppOutputs")
+
+        # Init state
         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 PushServerAppOutputsResponse()
 
@@ -264,8 +294,13 @@ class ServerAppIoServicer(serverappio_pb2_grpc.ServerAppIoServicer):
     ) -> UpdateRunStatusResponse:
         """Update the status of a run."""
         log(DEBUG, "ServerAppIoServicer.UpdateRunStatus")
+
+        # Init state
         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)
@@ -284,6 +319,21 @@ class ServerAppIoServicer(serverappio_pb2_grpc.ServerAppIoServicer):
         state.add_serverapp_log(request.run_id, merged_logs)
         return PushLogsResponse()
 
+    def GetRunStatus(
+        self, request: GetRunStatusRequest, context: grpc.ServicerContext
+    ) -> GetRunStatusResponse:
+        """Get the status of a run."""
+        log(DEBUG, "ServerAppIoServicer.GetRunStatus")
+        state = self.state_factory.state()
+
+        # Get run status from LinkState
+        run_statuses = state.get_run_status(set(request.run_ids))
+        run_status_dict = {
+            run_id: run_status_to_proto(run_status)
+            for run_id, run_status in run_statuses.items()
+        }
+        return GetRunStatusResponse(run_status_dict=run_status_dict)
+
 
 def _raise_if(validation_error: bool, detail: str) -> None:
     if validation_error:

flwr/server/superlink/ffs/disk_ffs.py

@@ -14,6 +14,7 @@
 # ==============================================================================
 """Disk based Flower File Storage."""
 
+
 import hashlib
 import json
 from pathlib import Path

flwr/server/superlink/fleet/grpc_adapter/grpc_adapter_servicer.py

@@ -158,4 +158,5 @@ class GrpcAdapterServicer(grpcadapter_pb2_grpc.GrpcAdapterServicer):
         return message_handler.get_fab(
             request=request,
             ffs=self.ffs_factory.ffs(),
+            state=self.state_factory.state(),
         )

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/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()
 
@@ -121,6 +132,15 @@ def get_run(
     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 +153,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.
 

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(

flwr/server/superlink/linkstate/sqlite_linkstate.py

@@ -14,6 +14,7 @@
 # ==============================================================================
 """SQLite based implemenation of the link state."""
 
+
 # pylint: disable=too-many-lines
 
 import json
@@ -566,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:
@@ -592,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