flwr 1.20.0__py3-none-any.whl → 1.21.0__py3-none-any.whl

flwr/server/serverapp/app.py

@@ -16,13 +16,12 @@
 
 
 import argparse
-import gc
 from logging import DEBUG, ERROR, INFO
 from pathlib import Path
 from queue import Queue
-from time import sleep
 from typing import Optional
 
+from flwr.app.exception import AppExitException
 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
@@ -35,10 +34,11 @@ from flwr.common.config import (
 )
 from flwr.common.constant import (
     SERVERAPPIO_API_DEFAULT_CLIENT_ADDRESS,
+    ExecPluginType,
     Status,
     SubStatus,
 )
-from flwr.common.exit import ExitCode, flwr_exit
+from flwr.common.exit import ExitCode, add_exit_handler, flwr_exit
 from flwr.common.heartbeat import HeartbeatSender, get_grpc_app_heartbeat_fn
 from flwr.common.logger import (
     log,
@@ -62,8 +62,12 @@ from flwr.proto.appio_pb2 import (  # pylint: disable=E0611
     PushAppOutputsRequest,
 )
 from flwr.proto.run_pb2 import UpdateRunStatusRequest  # pylint: disable=E0611
+from flwr.proto.serverappio_pb2_grpc import ServerAppIoStub
 from flwr.server.grid.grpc_grid import GrpcGrid
 from flwr.server.run_serverapp import run as run_
+from flwr.supercore.app_utils import start_parent_process_monitor
+from flwr.supercore.superexec.plugin import ServerAppExecPlugin
+from flwr.supercore.superexec.run_superexec import run_with_deprecation_warning
 
 
 def flwr_serverapp() -> None:
@@ -74,14 +78,27 @@ def flwr_serverapp() -> None:
 
     args = _parse_args_run_flwr_serverapp().parse_args()
 
-    log(INFO, "Start `flwr-serverapp` process")
-
     if not args.insecure:
         flwr_exit(
             ExitCode.COMMON_TLS_NOT_SUPPORTED,
             "`flwr-serverapp` does not support TLS yet.",
         )
 
+    # Disallow long-running `flwr-serverapp` processes
+    if args.token is None:
+        run_with_deprecation_warning(
+            cmd="flwr-serverapp",
+            plugin_type=ExecPluginType.SERVER_APP,
+            plugin_class=ServerAppExecPlugin,
+            stub_class=ServerAppIoStub,
+            appio_api_address=args.serverappio_api_address,
+            flwr_dir=args.flwr_dir,
+            parent_pid=args.parent_pid,
+            warn_run_once=args.run_once,
+        )
+        return
+
+    log(INFO, "Start `flwr-serverapp` process")
     log(
         DEBUG,
         "`flwr-serverapp` will attempt to connect to SuperLink's "
@@ -91,177 +108,177 @@ def flwr_serverapp() -> None:
     run_serverapp(
         serverappio_api_address=args.serverappio_api_address,
         log_queue=log_queue,
-        run_once=args.run_once,
+        token=args.token,
         flwr_dir=args.flwr_dir,
         certificates=None,
+        parent_pid=args.parent_pid,
     )
 
     # Restore stdout/stderr
     restore_output()
 
 
-def run_serverapp(  # pylint: disable=R0914, disable=W0212, disable=R0915
+def run_serverapp(  # pylint: disable=R0913, R0914, R0915, R0917, W0212
     serverappio_api_address: str,
     log_queue: Queue[Optional[str]],
-    run_once: bool,
+    token: str,
     flwr_dir: Optional[str] = None,
     certificates: Optional[bytes] = None,
+    parent_pid: Optional[int] = None,
 ) -> None:
     """Run Flower ServerApp process."""
+    # Monitor the main process in case of SIGKILL
+    if parent_pid is not None:
+        start_parent_process_monitor(parent_pid)
+
     # Resolve directory where FABs are installed
     flwr_dir_ = get_flwr_dir(flwr_dir)
     log_uploader = None
-    success = True
     hash_run_id = None
     run_status = None
     heartbeat_sender = None
     grid = None
     context = None
-    while True:
+    exit_code = ExitCode.SUCCESS
 
-        try:
-            # Initialize the GrpcGrid
-            grid = GrpcGrid(
-                serverappio_service_address=serverappio_api_address,
-                root_certificates=certificates,
-            )
+    def on_exit() -> None:
+        # Stop heartbeat sender
+        if heartbeat_sender:
+            heartbeat_sender.stop()
+
+        # Stop log uploader for this run and upload final logs
+        if log_uploader:
+            stop_log_uploader(log_queue, log_uploader)
 
-            # Pull ServerAppInputs from LinkState
-            req = PullAppInputsRequest()
-            log(DEBUG, "[flwr-serverapp] Pull ServerAppInputs")
-            res: PullAppInputsResponse = grid._stub.PullAppInputs(req)
-            if not res.HasField("run"):
-                sleep(3)
-                run_status = None
-                continue
-
-            context = context_from_proto(res.context)
-            run = run_from_proto(res.run)
-            fab = fab_from_proto(res.fab)
-
-            hash_run_id = get_sha256_hash(run.run_id)
-
-            grid.set_run(run.run_id)
-
-            # Start log uploader for this run
-            log_uploader = start_log_uploader(
-                log_queue=log_queue,
-                node_id=0,
-                run_id=run.run_id,
-                stub=grid._stub,
+        # Update run status
+        if run_status and grid:
+            run_status_proto = run_status_to_proto(run_status)
+            grid._stub.UpdateRunStatus(
+                UpdateRunStatusRequest(run_id=run.run_id, run_status=run_status_proto)
             )
 
-            log(DEBUG, "[flwr-serverapp] Start FAB installation.")
-            install_from_fab(fab.content, flwr_dir=flwr_dir_, skip_prompt=True)
+        # Close the Grpc connection
+        if grid:
+            grid.close()
 
-            fab_id, fab_version = get_fab_metadata(fab.content)
+    add_exit_handler(on_exit)
 
-            app_path = str(
-                get_project_dir(fab_id, fab_version, fab.hash_str, flwr_dir_)
-            )
-            config = get_project_config(app_path)
+    try:
+        # Initialize the GrpcGrid
+        grid = GrpcGrid(
+            serverappio_service_address=serverappio_api_address,
+            root_certificates=certificates,
+        )
 
-            # Obtain server app reference and the run config
-            server_app_attr = config["tool"]["flwr"]["app"]["components"]["serverapp"]
-            server_app_run_config = get_fused_config_from_dir(
-                Path(app_path), run.override_config
-            )
+        # Pull ServerAppInputs from LinkState
+        req = PullAppInputsRequest(token=token)
+        log(DEBUG, "[flwr-serverapp] Pull ServerAppInputs")
+        res: PullAppInputsResponse = grid._stub.PullAppInputs(req)
+        context = context_from_proto(res.context)
+        run = run_from_proto(res.run)
+        fab = fab_from_proto(res.fab)
 
-            # Update run_config in context
-            context.run_config = server_app_run_config
+        hash_run_id = get_sha256_hash(run.run_id)
 
-            log(
-                DEBUG,
-                "[flwr-serverapp] Will load ServerApp `%s` in %s",
-                server_app_attr,
-                app_path,
-            )
+        grid.set_run(run.run_id)
 
-            # Change status to Running
-            run_status_proto = run_status_to_proto(RunStatus(Status.RUNNING, "", ""))
-            grid._stub.UpdateRunStatus(
-                UpdateRunStatusRequest(run_id=run.run_id, run_status=run_status_proto)
-            )
+        # Start log uploader for this run
+        log_uploader = start_log_uploader(
+            log_queue=log_queue,
+            node_id=0,
+            run_id=run.run_id,
+            stub=grid._stub,
+        )
 
-            event(
-                EventType.FLWR_SERVERAPP_RUN_ENTER,
-                event_details={"run-id-hash": hash_run_id},
-            )
+        log(DEBUG, "[flwr-serverapp] Start FAB installation.")
+        install_from_fab(fab.content, flwr_dir=flwr_dir_, skip_prompt=True)
 
-            # Set up heartbeat sender
-            heartbeat_fn = get_grpc_app_heartbeat_fn(
-                grid._stub,
-                run.run_id,
-                failure_message="Heartbeat failed unexpectedly. The SuperLink could "
-                "not find the provided run ID, or the run status is invalid.",
-            )
-            heartbeat_sender = HeartbeatSender(heartbeat_fn)
-            heartbeat_sender.start()
-
-            # Load and run the ServerApp with the Grid
-            updated_context = run_(
-                grid=grid,
-                server_app_dir=app_path,
-                server_app_attr=server_app_attr,
-                context=context,
-            )
+        fab_id, fab_version = get_fab_metadata(fab.content)
 
-            # Send resulting context
-            context_proto = context_to_proto(updated_context)
-            log(DEBUG, "[flwr-serverapp] Will push ServerAppOutputs")
-            out_req = PushAppOutputsRequest(run_id=run.run_id, context=context_proto)
-            _ = grid._stub.PushAppOutputs(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 heartbeat sender
-            if heartbeat_sender:
-                heartbeat_sender.stop()
-                heartbeat_sender = None
-
-            # Stop log uploader for this run and upload final logs
-            if log_uploader:
-                stop_log_uploader(log_queue, log_uploader)
-                log_uploader = None
-
-            # Update run status
-            if run_status and grid:
-                run_status_proto = run_status_to_proto(run_status)
-                grid._stub.UpdateRunStatus(
-                    UpdateRunStatusRequest(
-                        run_id=run.run_id, run_status=run_status_proto
-                    )
-                )
-
-            # Close the Grpc connection
-            if grid:
-                grid.close()
-
-            # Clean up the Context
-            context = None
-            gc.collect()
-
-            event(
-                EventType.FLWR_SERVERAPP_RUN_LEAVE,
-                event_details={"run-id-hash": hash_run_id, "success": success},
-            )
+        app_path = str(get_project_dir(fab_id, fab_version, fab.hash_str, flwr_dir_))
+        config = get_project_config(app_path)
+
+        # Obtain server app reference and the run config
+        server_app_attr = config["tool"]["flwr"]["app"]["components"]["serverapp"]
+        server_app_run_config = get_fused_config_from_dir(
+            Path(app_path), run.override_config
+        )
 
-        # Stop the loop if `flwr-serverapp` is expected to process a single run
-        if run_once:
-            break
+        # Update run_config in context
+        context.run_config = server_app_run_config
+
+        log(
+            DEBUG,
+            "[flwr-serverapp] Will load ServerApp `%s` in %s",
+            server_app_attr,
+            app_path,
+        )
+
+        # Change status to Running
+        run_status_proto = run_status_to_proto(RunStatus(Status.RUNNING, "", ""))
+        grid._stub.UpdateRunStatus(
+            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},
+        )
+
+        # Set up heartbeat sender
+        heartbeat_fn = get_grpc_app_heartbeat_fn(
+            grid._stub,
+            run.run_id,
+            failure_message="Heartbeat failed unexpectedly. The SuperLink could "
+            "not find the provided run ID, or the run status is invalid.",
+        )
+        heartbeat_sender = HeartbeatSender(heartbeat_fn)
+        heartbeat_sender.start()
+
+        # Load and run the ServerApp with the Grid
+        updated_context = run_(
+            grid=grid,
+            server_app_dir=app_path,
+            server_app_attr=server_app_attr,
+            context=context,
+        )
+
+        # Send resulting context
+        context_proto = context_to_proto(updated_context)
+        log(DEBUG, "[flwr-serverapp] Will push ServerAppOutputs")
+        out_req = PushAppOutputsRequest(
+            token=token, run_id=run.run_id, context=context_proto
+        )
+        _ = grid._stub.PushAppOutputs(out_req)
+
+        run_status = RunStatus(Status.FINISHED, SubStatus.COMPLETED, "")
+
+    # Raised when the run is already stopped by the user
+    except RunNotRunningException:
+        log(INFO, "")
+        log(INFO, "Run ID %s stopped.", run.run_id)
+        log(INFO, "")
+        run_status = None
+        # No need to update the exit code since this is expected behavior
+
+    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))
+
+        # Set exit code
+        exit_code = ExitCode.SERVERAPP_EXCEPTION  # General exit code
+        if isinstance(ex, AppExitException):
+            exit_code = ex.exit_code
+
+    flwr_exit(
+        code=exit_code,
+        event_type=EventType.FLWR_SERVERAPP_RUN_LEAVE,
+        event_details={
+            "run-id-hash": hash_run_id,
+            "success": exit_code == ExitCode.SUCCESS,
+        },
+    )
 
 
 def _parse_args_run_flwr_serverapp() -> argparse.ArgumentParser:
@@ -276,11 +293,5 @@ def _parse_args_run_flwr_serverapp() -> argparse.ArgumentParser:
         help="Address of SuperLink's ServerAppIo API (IPv4, IPv6, or a domain name)."
         f"By default, it is set to {SERVERAPPIO_API_DEFAULT_CLIENT_ADDRESS}.",
     )
-    parser.add_argument(
-        "--run-once",
-        action="store_true",
-        help="When set, this process will start a single ServerApp for a pending Run. "
-        "If there is no pending Run, the process will exit.",
-    )
     add_args_flwr_app_common(parser=parser)
     return parser

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

@@ -161,6 +161,7 @@ class RayBackend(Backend):
                 "Call the backend's `build()` method before processing messages."
             )
 
+        future = None
         try:
             # Submit a task to the pool
             future = self.pool.submit(
@@ -183,7 +184,8 @@ class RayBackend(Backend):
                 self.__class__.__name__,
             )
             # add actor back into pool
-            self.pool.add_actor_back_to_pool(future)
+            if future is not None:
+                self.pool.add_actor_back_to_pool(future)
             raise ex
 
     def terminate(self) -> None:

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

@@ -23,7 +23,6 @@ from concurrent.futures import ThreadPoolExecutor
 from logging import DEBUG, ERROR, INFO, WARN
 from pathlib import Path
 from queue import Empty, Queue
-from time import sleep
 from typing import Callable, Optional
 from uuid import uuid4
 
@@ -153,7 +152,7 @@ def add_messages_to_queue(
             message_ins_list = state.get_message_ins(node_id=node_id, limit=1)
             for msg in message_ins_list:
                 queue.put(msg)
-        sleep(0.1)
+        f_stop.wait(0.1)
 
 
 def put_message_into_state(
@@ -182,6 +181,7 @@ def run_api(
     messageins_queue: Queue[Message] = Queue()
     messageres_queue: Queue[Message] = Queue()
 
+    backend = None
     try:
 
         # Instantiate backend
@@ -236,16 +236,16 @@ def run_api(
         log(ERROR, traceback.format_exc())
         log(WARN, "Stopping Simulation Engine.")
 
-        # Manually trigger stopping event
-        f_stop.set()
-
         # Raise exception
         raise RuntimeError("Simulation Engine crashed.") from ex
 
     finally:
+        # Manually trigger stopping event
+        f_stop.set()
 
         # Terminate backend
-        backend.terminate()
+        if backend is not None:
+            backend.terminate()
 
 
 # pylint: disable=too-many-arguments,unused-argument,too-many-locals,too-many-branches

flwr/server/superlink/linkstate/in_memory_linkstate.py

@@ -15,6 +15,7 @@
 """In-memory LinkState implementation."""
 
 
+import secrets
 import threading
 import time
 from bisect import bisect_right
@@ -25,6 +26,7 @@ from typing import Optional
 
 from flwr.common import Context, Message, log, now
 from flwr.common.constant import (
+    FLWR_APP_TOKEN_LENGTH,
     HEARTBEAT_MAX_INTERVAL,
     HEARTBEAT_PATIENCE,
     MESSAGE_TTL_TOLERANCE,
@@ -80,6 +82,11 @@ class InMemoryLinkState(LinkState):  # pylint: disable=R0902,R0904
         self.message_res_store: dict[str, Message] = {}
         self.message_ins_id_to_message_res_id: dict[str, str] = {}
 
+        # Store run ID to token mapping and token to run ID mapping
+        self.token_store: dict[int, str] = {}
+        self.token_to_run_id: dict[str, int] = {}
+        self.lock_token_store = threading.Lock()
+
         # Map flwr_aid to run_ids for O(1) reverse index lookup
         self.flwr_aid_to_run_ids: dict[str, set[int]] = defaultdict(set)
 
@@ -678,3 +685,30 @@ class InMemoryLinkState(LinkState):  # pylint: disable=R0902,R0904
             index = bisect_right(run.logs, (after_timestamp, ""))
             latest_timestamp = run.logs[-1][0] if index < len(run.logs) else 0.0
             return "".join(log for _, log in run.logs[index:]), latest_timestamp
+
+    def create_token(self, run_id: int) -> Optional[str]:
+        """Create a token for the given run ID."""
+        token = secrets.token_hex(FLWR_APP_TOKEN_LENGTH)  # Generate a random token
+        with self.lock_token_store:
+            if run_id in self.token_store:
+                return None  # Token already created for this run ID
+            self.token_store[run_id] = token
+            self.token_to_run_id[token] = run_id
+        return token
+
+    def verify_token(self, run_id: int, token: str) -> bool:
+        """Verify a token for the given run ID."""
+        with self.lock_token_store:
+            return self.token_store.get(run_id) == token
+
+    def delete_token(self, run_id: int) -> None:
+        """Delete the token for the given run ID."""
+        with self.lock_token_store:
+            token = self.token_store.pop(run_id, None)
+            if token is not None:
+                self.token_to_run_id.pop(token, None)
+
+    def get_run_id_by_token(self, token: str) -> Optional[int]:
+        """Get the run ID associated with a given token."""
+        with self.lock_token_store:
+            return self.token_to_run_id.get(token)

flwr/server/superlink/linkstate/linkstate.py

@@ -21,9 +21,10 @@ from typing import Optional
 from flwr.common import Context, Message
 from flwr.common.record import ConfigRecord
 from flwr.common.typing import Run, RunStatus, UserConfig
+from flwr.supercore.corestate import CoreState
 
 
-class LinkState(abc.ABC):  # pylint: disable=R0904
+class LinkState(CoreState):  # pylint: disable=R0904
     """Abstract LinkState."""
 
     @abc.abstractmethod

flwr/server/superlink/linkstate/sqlite_linkstate.py

@@ -19,6 +19,7 @@
 
 import json
 import re
+import secrets
 import sqlite3
 import time
 from collections.abc import Sequence
@@ -27,6 +28,7 @@ from typing import Any, Optional, Union, cast
 
 from flwr.common import Context, Message, Metadata, log, now
 from flwr.common.constant import (
+    FLWR_APP_TOKEN_LENGTH,
     HEARTBEAT_MAX_INTERVAL,
     HEARTBEAT_PATIENCE,
     MESSAGE_TTL_TOLERANCE,
@@ -163,6 +165,13 @@ CREATE TABLE IF NOT EXISTS message_res(
 );
 """
 
+SQL_CREATE_TABLE_TOKEN_STORE = """
+CREATE TABLE IF NOT EXISTS token_store (
+    run_id                  INTEGER PRIMARY KEY,
+    token                   TEXT UNIQUE NOT NULL
+);
+"""
+
 DictOrTuple = Union[tuple[Any, ...], dict[str, Any]]
 
 
@@ -212,6 +221,7 @@ class SqliteLinkState(LinkState):  # pylint: disable=R0904
         cur.execute(SQL_CREATE_TABLE_MESSAGE_RES)
         cur.execute(SQL_CREATE_TABLE_NODE)
         cur.execute(SQL_CREATE_TABLE_PUBLIC_KEY)
+        cur.execute(SQL_CREATE_TABLE_TOKEN_STORE)
         cur.execute(SQL_CREATE_INDEX_ONLINE_UNTIL)
         res = cur.execute("SELECT name FROM sqlite_schema;")
         return res.fetchall()
@@ -1138,6 +1148,41 @@ class SqliteLinkState(LinkState):  # pylint: disable=R0904
 
         return message_ins
 
+    def create_token(self, run_id: int) -> Optional[str]:
+        """Create a token for the given run ID."""
+        token = secrets.token_hex(FLWR_APP_TOKEN_LENGTH)  # Generate a random token
+        query = "INSERT INTO token_store (run_id, token) VALUES (:run_id, :token);"
+        data = {"run_id": convert_uint64_to_sint64(run_id), "token": token}
+        try:
+            self.query(query, data)
+        except sqlite3.IntegrityError:
+            return None  # Token already created for this run ID
+        return token
+
+    def verify_token(self, run_id: int, token: str) -> bool:
+        """Verify a token for the given run ID."""
+        query = "SELECT token FROM token_store WHERE run_id = :run_id;"
+        data = {"run_id": convert_uint64_to_sint64(run_id)}
+        rows = self.query(query, data)
+        if not rows:
+            return False
+        return cast(str, rows[0]["token"]) == token
+
+    def delete_token(self, run_id: int) -> None:
+        """Delete the token for the given run ID."""
+        query = "DELETE FROM token_store WHERE run_id = :run_id;"
+        data = {"run_id": convert_uint64_to_sint64(run_id)}
+        self.query(query, data)
+
+    def get_run_id_by_token(self, token: str) -> Optional[int]:
+        """Get the run ID associated with a given token."""
+        query = "SELECT run_id FROM token_store WHERE token = :token;"
+        data = {"token": token}
+        rows = self.query(query, data)
+        if not rows:
+            return None
+        return convert_sint64_to_uint64(rows[0]["run_id"])
+
 
 def dict_factory(
     cursor: sqlite3.Cursor,

flwr/server/superlink/serverappio/serverappio_grpc.py

@@ -58,7 +58,7 @@ def run_serverappio_api_grpc(
         certificates=certificates,
     )
 
-    log(INFO, "Flower ECE: Starting ServerAppIo API (gRPC-rere) on %s", address)
+    log(INFO, "Flower Deployment Runtime: Starting ServerAppIo API on %s", address)
     serverappio_grpc_server.start()
 
     return serverappio_grpc_server