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

flwr/server/app.py

@@ -14,6 +14,7 @@
 # ==============================================================================
 """Flower server app."""
 
+
 import argparse
 import csv
 import importlib.util
@@ -24,9 +25,10 @@ from collections.abc import Sequence
 from logging import DEBUG, INFO, WARN
 from pathlib import Path
 from time import sleep
-from typing import Optional
+from typing import Any, Optional
 
 import grpc
+import yaml
 from cryptography.exceptions import UnsupportedAlgorithm
 from cryptography.hazmat.primitives.asymmetric import ec
 from cryptography.hazmat.primitives.serialization import (
@@ -37,8 +39,10 @@ from cryptography.hazmat.primitives.serialization import (
 from flwr.common import GRPC_MAX_MESSAGE_LENGTH, EventType, event
 from flwr.common.address import parse_address
 from flwr.common.args import try_obtain_server_certificates
+from flwr.common.auth_plugin import ExecAuthPlugin
 from flwr.common.config import get_flwr_dir, parse_config_args
 from flwr.common.constant import (
+    AUTH_TYPE,
     CLIENT_OCTET,
     EXEC_API_DEFAULT_SERVER_ADDRESS,
     FLEET_API_GRPC_BIDI_DEFAULT_ADDRESS,
@@ -66,7 +70,6 @@ from flwr.proto.fleet_pb2_grpc import (  # pylint: disable=E0611
 from flwr.proto.grpcadapter_pb2_grpc import add_GrpcAdapterServicer_to_server
 from flwr.superexec.app import load_executor
 from flwr.superexec.exec_grpc import run_exec_api_grpc
-from flwr.superexec.simulation import SimulationEngine
 
 from .client_manager import ClientManager
 from .history import History
@@ -89,6 +92,19 @@ DATABASE = ":flwr-in-memory-state:"
 BASE_DIR = get_flwr_dir() / "superlink" / "ffs"
 
 
+try:
+    from flwr.ee import add_ee_args_superlink, get_exec_auth_plugins
+except ImportError:
+
+    # pylint: disable-next=unused-argument
+    def add_ee_args_superlink(parser: argparse.ArgumentParser) -> None:
+        """Add EE-specific arguments to the parser."""
+
+    def get_exec_auth_plugins() -> dict[str, type[ExecAuthPlugin]]:
+        """Return all Exec API authentication plugins."""
+        raise NotImplementedError("No authentication plugins are currently supported.")
+
+
 def start_server(  # pylint: disable=too-many-arguments,too-many-locals
     *,
     server_address: str = FLEET_API_GRPC_BIDI_DEFAULT_ADDRESS,
@@ -247,6 +263,12 @@ def run_superlink() -> None:
     # Obtain certificates
     certificates = try_obtain_server_certificates(args, args.fleet_api_type)
 
+    user_auth_config = _try_obtain_user_auth_config(args)
+    auth_plugin: Optional[ExecAuthPlugin] = None
+    # user_auth_config is None only if the args.user_auth_config is not provided
+    if user_auth_config is not None:
+        auth_plugin = _try_obtain_exec_auth_plugin(user_auth_config)
+
     # Initialize StateFactory
     state_factory = LinkStateFactory(args.database)
 
@@ -264,13 +286,13 @@ def run_superlink() -> None:
         config=parse_config_args(
             [args.executor_config] if args.executor_config else args.executor_config
         ),
+        auth_plugin=auth_plugin,
     )
     grpc_servers = [exec_server]
 
     # Determine Exec plugin
     # If simulation is used, don't start ServerAppIo and Fleet APIs
-    sim_exec = isinstance(executor, SimulationEngine)
-
+    sim_exec = executor.__class__.__qualname__ == "SimulationEngine"
     bckg_threads = []
 
     if sim_exec:
@@ -278,7 +300,7 @@ def run_superlink() -> None:
             address=simulationio_address,
             state_factory=state_factory,
             ffs_factory=ffs_factory,
-            certificates=certificates,
+            certificates=None,  # SimulationAppIo API doesn't support SSL yet
         )
         grpc_servers.append(simulationio_server)
 
@@ -352,6 +374,7 @@ def run_superlink() -> None:
                     server_public_key,
                 ) = maybe_keys
                 state = state_factory.state()
+                state.clear_supernode_auth_keys_and_credentials()
                 state.store_node_public_keys(node_public_keys)
                 state.store_server_private_public_key(
                     private_key_to_bytes(server_private_key),
@@ -362,7 +385,7 @@ def run_superlink() -> None:
                     "Node authentication enabled with %d known public keys",
                     len(node_public_keys),
                 )
-                interceptors = [AuthenticateServerInterceptor(state)]
+                interceptors = [AuthenticateServerInterceptor(state_factory)]
 
             fleet_server = _run_fleet_api_grpc_rere(
                 address=fleet_address,
@@ -389,6 +412,9 @@ def run_superlink() -> None:
         io_address = (
             f"{CLIENT_OCTET}:{_port}" if _octet == SERVER_OCTET else serverappio_address
         )
+        address_arg = (
+            "--simulationio-api-address" if sim_exec else "--serverappio-api-address"
+        )
         address = simulationio_address if sim_exec else io_address
         cmd = "flwr-simulation" if sim_exec else "flwr-serverapp"
 
@@ -397,6 +423,7 @@ def run_superlink() -> None:
             target=_flwr_scheduler,
             args=(
                 state_factory,
+                address_arg,
                 address,
                 cmd,
             ),
@@ -422,6 +449,7 @@ def run_superlink() -> None:
 
 def _flwr_scheduler(
     state_factory: LinkStateFactory,
+    io_api_arg: str,
     io_api_address: str,
     cmd: str,
 ) -> None:
@@ -446,7 +474,7 @@ def _flwr_scheduler(
             command = [
                 cmd,
                 "--run-once",
-                "--serverappio-api-address",
+                io_api_arg,
                 io_api_address,
                 "--insecure",
             ]
@@ -556,6 +584,32 @@ def _try_setup_node_authentication(
         )
 
 
+def _try_obtain_user_auth_config(args: argparse.Namespace) -> Optional[dict[str, Any]]:
+    if getattr(args, "user_auth_config", None) is not None:
+        with open(args.user_auth_config, encoding="utf-8") as file:
+            config: dict[str, Any] = yaml.safe_load(file)
+            return config
+    return None
+
+
+def _try_obtain_exec_auth_plugin(config: dict[str, Any]) -> Optional[ExecAuthPlugin]:
+    auth_config: dict[str, Any] = config.get("authentication", {})
+    auth_type: str = auth_config.get(AUTH_TYPE, "")
+    try:
+        all_plugins: dict[str, type[ExecAuthPlugin]] = get_exec_auth_plugins()
+        auth_plugin_class = all_plugins[auth_type]
+        return auth_plugin_class(config=auth_config)
+    except KeyError:
+        if auth_type != "":
+            sys.exit(
+                f'Authentication type "{auth_type}" is not supported. '
+                "Please provide a valid authentication type in the configuration."
+            )
+        sys.exit("No authentication type is provided in the configuration.")
+    except NotImplementedError:
+        sys.exit("No authentication plugins are currently supported.")
+
+
 def _run_fleet_api_grpc_rere(
     address: str,
     state_factory: LinkStateFactory,
@@ -654,6 +708,7 @@ def _parse_args_run_superlink() -> argparse.ArgumentParser:
     )
 
     _add_args_common(parser=parser)
+    add_ee_args_superlink(parser=parser)
     _add_args_serverappio_api(parser=parser)
     _add_args_fleet_api(parser=parser)
     _add_args_exec_api(parser=parser)

flwr/server/compat/app_utils.py

@@ -17,6 +17,8 @@
 
 import threading
 
+from flwr.common.typing import RunNotRunningException
+
 from ..client_manager import ClientManager
 from ..compat.driver_client_proxy import DriverClientProxy
 from ..driver import Driver
@@ -74,7 +76,11 @@ def _update_client_manager(
     # Loop until the driver is disconnected
     registered_nodes: dict[int, DriverClientProxy] = {}
     while not f_stop.is_set():
-        all_node_ids = set(driver.get_node_ids())
+        try:
+            all_node_ids = set(driver.get_node_ids())
+        except RunNotRunningException:
+            f_stop.set()
+            break
         dead_nodes = set(registered_nodes).difference(all_node_ids)
         new_nodes = all_node_ids.difference(registered_nodes)
 

flwr/server/driver/grpc_driver.py

@@ -14,19 +14,20 @@
 # ==============================================================================
 """Flower gRPC Driver."""
 
+
 import time
 import warnings
 from collections.abc import Iterable
-from logging import DEBUG, INFO, WARN, WARNING
-from typing import Any, Optional, cast
+from logging import DEBUG, WARNING
+from typing import Optional, cast
 
 import grpc
 
 from flwr.common import DEFAULT_TTL, Message, Metadata, RecordSet
-from flwr.common.constant import MAX_RETRY_DELAY, SERVERAPPIO_API_DEFAULT_CLIENT_ADDRESS
+from flwr.common.constant import SERVERAPPIO_API_DEFAULT_CLIENT_ADDRESS
 from flwr.common.grpc import create_channel
 from flwr.common.logger import log
-from flwr.common.retry_invoker import RetryInvoker, RetryState, exponential
+from flwr.common.retry_invoker import _make_simple_grpc_retry_invoker, _wrap_stub
 from flwr.common.serde import message_from_taskres, message_to_taskins, run_from_proto
 from flwr.common.typing import Run
 from flwr.proto.node_pb2 import Node  # pylint: disable=E0611
@@ -203,7 +204,9 @@ class GrpcDriver(Driver):
             task_ins_list.append(taskins)
         # Call GrpcDriverStub method
         res: PushTaskInsResponse = self._stub.PushTaskIns(
-            PushTaskInsRequest(task_ins_list=task_ins_list)
+            PushTaskInsRequest(
+                task_ins_list=task_ins_list, run_id=cast(Run, self._run).run_id
+            )
         )
         return list(res.task_ids)
 
@@ -215,7 +218,9 @@ class GrpcDriver(Driver):
         """
         # Pull TaskRes
         res: PullTaskResResponse = self._stub.PullTaskRes(
-            PullTaskResRequest(node=self.node, task_ids=message_ids)
+            PullTaskResRequest(
+                node=self.node, task_ids=message_ids, run_id=cast(Run, self._run).run_id
+            )
         )
         # Convert TaskRes to Message
         msgs = [message_from_taskres(taskres) for taskres in res.task_res_list]
@@ -258,60 +263,3 @@ class GrpcDriver(Driver):
             return
         # Disconnect
         self._disconnect()
-
-
-def _make_simple_grpc_retry_invoker() -> RetryInvoker:
-    """Create a simple gRPC retry invoker."""
-
-    def _on_sucess(retry_state: RetryState) -> None:
-        if retry_state.tries > 1:
-            log(
-                INFO,
-                "Connection successful after %.2f seconds and %s tries.",
-                retry_state.elapsed_time,
-                retry_state.tries,
-            )
-
-    def _on_backoff(retry_state: RetryState) -> None:
-        if retry_state.tries == 1:
-            log(WARN, "Connection attempt failed, retrying...")
-        else:
-            log(
-                WARN,
-                "Connection attempt failed, retrying in %.2f seconds",
-                retry_state.actual_wait,
-            )
-
-    def _on_giveup(retry_state: RetryState) -> None:
-        if retry_state.tries > 1:
-            log(
-                WARN,
-                "Giving up reconnection after %.2f seconds and %s tries.",
-                retry_state.elapsed_time,
-                retry_state.tries,
-            )
-
-    return RetryInvoker(
-        wait_gen_factory=lambda: exponential(max_delay=MAX_RETRY_DELAY),
-        recoverable_exceptions=grpc.RpcError,
-        max_tries=None,
-        max_time=None,
-        on_success=_on_sucess,
-        on_backoff=_on_backoff,
-        on_giveup=_on_giveup,
-        should_giveup=lambda e: e.code() != grpc.StatusCode.UNAVAILABLE,  # type: ignore
-    )
-
-
-def _wrap_stub(stub: ServerAppIoStub, retry_invoker: RetryInvoker) -> None:
-    """Wrap the gRPC stub with a retry invoker."""
-
-    def make_lambda(original_method: Any) -> Any:
-        return lambda *args, **kwargs: retry_invoker.invoke(
-            original_method, *args, **kwargs
-        )
-
-    for method_name in vars(stub):
-        method = getattr(stub, method_name)
-        if callable(method):
-            setattr(stub, method_name, make_lambda(method))

flwr/server/driver/inmemory_driver.py

@@ -142,7 +142,11 @@ class InMemoryDriver(Driver):
         # Pull TaskRes
         task_res_list = self.state.get_task_res(task_ids=msg_ids)
         # Delete tasks in state
-        self.state.delete_tasks(msg_ids)
+        # 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
+        }
+        self.state.delete_tasks(task_ins_ids=task_ins_ids_to_delete)
         # Convert TaskRes to Message
         msgs = [message_from_taskres(taskres) for taskres in task_res_list]
         return msgs

flwr/server/run_serverapp.py

@@ -15,12 +15,12 @@
 """Run ServerApp."""
 
 
-import sys
 from logging import DEBUG, ERROR
 from typing import Optional
 
-from flwr.common import Context
-from flwr.common.logger import log, warn_unsupported_feature
+from flwr.common import Context, EventType, event
+from flwr.common.exit_handlers import register_exit_handlers
+from flwr.common.logger import log
 from flwr.common.object_ref import load_app
 
 from .driver import Driver
@@ -66,12 +66,11 @@ def run(
     return context
 
 
-# pylint: disable-next=too-many-branches,too-many-statements,too-many-locals
 def run_server_app() -> None:
     """Run Flower server app."""
-    warn_unsupported_feature(
-        "The command `flower-server-app` is deprecated and no longer in use. "
-        "Use the `flwr-serverapp` exclusively instead."
+    event(EventType.RUN_SERVER_APP_ENTER)
+    log(
+        ERROR,
+        "The command `flower-server-app` has been replaced by `flwr run`.",
     )
-    log(ERROR, "`flower-server-app` used.")
-    sys.exit()
+    register_exit_handlers(event_type=EventType.RUN_SERVER_APP_LEAVE)

flwr/server/serverapp/app.py

@@ -14,6 +14,7 @@
 # ==============================================================================
 """Flower ServerApp process."""
 
+
 import argparse
 import sys
 from logging import DEBUG, ERROR, INFO
@@ -24,6 +25,7 @@ from typing import Optional
 
 from flwr.cli.config_utils import get_fab_metadata
 from flwr.cli.install import install_from_fab
+from flwr.cli.utils import get_sha256_hash
 from flwr.common.args import add_args_flwr_app_common
 from flwr.common.config import (
     get_flwr_dir,
@@ -50,7 +52,8 @@ from flwr.common.serde import (
     run_from_proto,
     run_status_to_proto,
 )
-from flwr.common.typing import RunStatus
+from flwr.common.telemetry import EventType, event
+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 +99,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,
@@ -112,7 +115,8 @@ def run_serverapp(  # pylint: disable=R0914, disable=W0212
     # Resolve directory where FABs are installed
     flwr_dir_ = get_flwr_dir(flwr_dir)
     log_uploader = None
-
+    success = True
+    hash_run_id = None
     while True:
 
         try:
@@ -128,6 +132,8 @@ def run_serverapp(  # pylint: disable=R0914, disable=W0212
             run = run_from_proto(res.run)
             fab = fab_from_proto(res.fab)
 
+            hash_run_id = get_sha256_hash(run.run_id)
+
             driver.set_run(run.run_id)
 
             # Start log uploader for this run
@@ -170,6 +176,11 @@ def run_serverapp(  # pylint: disable=R0914, disable=W0212
                 UpdateRunStatusRequest(run_id=run.run_id, run_status=run_status_proto)
             )
 
+            event(
+                EventType.FLWR_SERVERAPP_RUN_ENTER,
+                event_details={"run-id-hash": hash_run_id},
+            )
+
             # Load and run the ServerApp with the Driver
             updated_context = run_(
                 driver=driver,
@@ -186,11 +197,18 @@ def run_serverapp(  # pylint: disable=R0914, disable=W0212
             _ = driver._stub.PushServerAppOutputs(out_req)
 
             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
+            success = False
 
         except Exception as ex:  # pylint: disable=broad-exception-caught
             exc_entity = "ServerApp"
             log(ERROR, "%s raised an exception", exc_entity, exc_info=ex)
             run_status = RunStatus(Status.FINISHED, SubStatus.FAILED, str(ex))
+            success = False
 
         finally:
             # Stop log uploader for this run and upload final logs
@@ -206,6 +224,10 @@ def run_serverapp(  # pylint: disable=R0914, disable=W0212
                         run_id=run.run_id, run_status=run_status_proto
                     )
                 )
+            event(
+                EventType.FLWR_SERVERAPP_RUN_LEAVE,
+                event_details={"run-id-hash": hash_run_id, "success": success},
+            )
 
         # Stop the loop if `flwr-serverapp` is expected to process a single run
         if run_once:
@@ -230,12 +252,5 @@ def _parse_args_run_flwr_serverapp() -> argparse.ArgumentParser:
         help="When set, this process will start a single ServerApp for a pending Run. "
         "If there is no pending Run, the process will exit.",
     )
-    parser.add_argument(
-        "--root-certificates",
-        metavar="ROOT_CERT",
-        type=str,
-        help="Specifies the path to the PEM-encoded root certificate file for "
-        "establishing secure HTTPS connections.",
-    )
     add_args_flwr_app_common(parser=parser)
     return parser

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:
@@ -133,9 +159,9 @@ class ServerAppIoServicer(serverappio_pb2_grpc.ServerAppIoServicer):
         for task_ins in request.task_ins_list:
             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()
+            _raise_if(
+                request.run_id != task_ins.run_id, "`task_ins` has mismatched `run_id`"
+            )
 
         # Store each TaskIns
         task_ids: list[Optional[UUID]] = []
@@ -153,33 +179,35 @@ 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)
+        # Validate request
+        for task_res in task_res_list:
+            _raise_if(
+                request.run_id != task_res.run_id, "`task_res` has mismatched `run_id`"
+            )
+
+        # 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 +283,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()
 
@@ -263,9 +302,14 @@ class ServerAppIoServicer(serverappio_pb2_grpc.ServerAppIoServicer):
         self, request: UpdateRunStatusRequest, context: grpc.ServicerContext
     ) -> UpdateRunStatusResponse:
         """Update the status of a run."""
-        log(DEBUG, "ControlServicer.UpdateRunStatus")
+        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 +328,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(),
         )