flwr-nightly 1.13.0.dev20241021__py3-none-any.whl → 1.13.0.dev20241106__py3-none-any.whl

flwr/server/app.py

@@ -17,12 +17,14 @@
 import argparse
 import csv
 import importlib.util
+import subprocess
 import sys
 import threading
 from collections.abc import Sequence
-from logging import INFO, WARN
+from logging import DEBUG, INFO, WARN
 from os.path import isfile
 from pathlib import Path
+from time import sleep
 from typing import Optional
 
 import grpc
@@ -35,13 +37,17 @@ 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.config import get_flwr_dir
+from flwr.common.config import get_flwr_dir, parse_config_args
 from flwr.common.constant import (
-    DRIVER_API_DEFAULT_ADDRESS,
+    EXEC_API_DEFAULT_ADDRESS,
     FLEET_API_GRPC_BIDI_DEFAULT_ADDRESS,
     FLEET_API_GRPC_RERE_DEFAULT_ADDRESS,
     FLEET_API_REST_DEFAULT_ADDRESS,
+    ISOLATION_MODE_PROCESS,
+    ISOLATION_MODE_SUBPROCESS,
     MISSING_EXTRA_REST,
+    SERVERAPPIO_API_DEFAULT_ADDRESS,
+    SIMULATIONIO_API_DEFAULT_ADDRESS,
     TRANSPORT_TYPE_GRPC_ADAPTER,
     TRANSPORT_TYPE_GRPC_RERE,
     TRANSPORT_TYPE_REST,
@@ -56,13 +62,16 @@ from flwr.proto.fleet_pb2_grpc import (  # pylint: disable=E0611
     add_FleetServicer_to_server,
 )
 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
 from .server import Server, init_defaults, run_fl
 from .server_config import ServerConfig
 from .strategy import Strategy
-from .superlink.driver.driver_grpc import run_driver_api_grpc
+from .superlink.driver.serverappio_grpc import run_serverappio_api_grpc
 from .superlink.ffs.ffs_factory import FfsFactory
 from .superlink.fleet.grpc_adapter.grpc_adapter_servicer import GrpcAdapterServicer
 from .superlink.fleet.grpc_bidi.grpc_server import (
@@ -71,7 +80,8 @@ from .superlink.fleet.grpc_bidi.grpc_server import (
 )
 from .superlink.fleet.grpc_rere.fleet_servicer import FleetServicer
 from .superlink.fleet.grpc_rere.server_interceptor import AuthenticateServerInterceptor
-from .superlink.state import StateFactory
+from .superlink.linkstate import LinkStateFactory
+from .superlink.simulation.simulationio_grpc import run_simulationio_api_grpc
 
 DATABASE = ":flwr-in-memory-state:"
 BASE_DIR = get_flwr_dir() / "superlink" / "ffs"
@@ -198,125 +208,169 @@ def start_server(  # pylint: disable=too-many-arguments,too-many-locals
 
 # pylint: disable=too-many-branches, too-many-locals, too-many-statements
 def run_superlink() -> None:
-    """Run Flower SuperLink (Driver API and Fleet API)."""
+    """Run Flower SuperLink (ServerAppIo API and Fleet API)."""
     args = _parse_args_run_superlink().parse_args()
 
     log(INFO, "Starting Flower SuperLink")
 
     event(EventType.RUN_SUPERLINK_ENTER)
 
-    # Parse IP address
-    driver_address, _, _ = _format_address(args.driver_api_address)
+    # Parse IP addresses
+    serverappio_address, _, _ = _format_address(args.serverappio_api_address)
+    exec_address, _, _ = _format_address(args.exec_api_address)
+    simulationio_address, _, _ = _format_address(args.simulationio_api_address)
 
     # Obtain certificates
     certificates = _try_obtain_certificates(args)
 
     # Initialize StateFactory
-    state_factory = StateFactory(args.database)
+    state_factory = LinkStateFactory(args.database)
 
     # Initialize FfsFactory
     ffs_factory = FfsFactory(args.storage_dir)
 
-    # Start Driver API
-    driver_server: grpc.Server = run_driver_api_grpc(
-        address=driver_address,
+    # Start Exec API
+    executor = load_executor(args)
+    exec_server: grpc.Server = run_exec_api_grpc(
+        address=exec_address,
         state_factory=state_factory,
         ffs_factory=ffs_factory,
+        executor=executor,
         certificates=certificates,
+        config=parse_config_args(
+            [args.executor_config] if args.executor_config else args.executor_config
+        ),
     )
+    grpc_servers = [exec_server]
 
-    grpc_servers = [driver_server]
-    bckg_threads = []
-    if not args.fleet_api_address:
-        if args.fleet_api_type in [
-            TRANSPORT_TYPE_GRPC_RERE,
-            TRANSPORT_TYPE_GRPC_ADAPTER,
-        ]:
-            args.fleet_api_address = FLEET_API_GRPC_RERE_DEFAULT_ADDRESS
-        elif args.fleet_api_type == TRANSPORT_TYPE_REST:
-            args.fleet_api_address = FLEET_API_REST_DEFAULT_ADDRESS
-
-    fleet_address, host, port = _format_address(args.fleet_api_address)
-
-    num_workers = args.fleet_api_num_workers
-    if num_workers != 1:
-        log(
-            WARN,
-            "The Fleet API currently supports only 1 worker. "
-            "You have specified %d workers. "
-            "Support for multiple workers will be added in future releases. "
-            "Proceeding with a single worker.",
-            args.fleet_api_num_workers,
-        )
-        num_workers = 1
-
-    # Start Fleet API
-    if args.fleet_api_type == TRANSPORT_TYPE_REST:
-        if (
-            importlib.util.find_spec("requests")
-            and importlib.util.find_spec("starlette")
-            and importlib.util.find_spec("uvicorn")
-        ) is None:
-            sys.exit(MISSING_EXTRA_REST)
-
-        _, ssl_certfile, ssl_keyfile = (
-            certificates if certificates is not None else (None, None, None)
-        )
+    # Determine Exec plugin
+    # If simulation is used, don't start ServerAppIo and Fleet APIs
+    sim_exec = isinstance(executor, SimulationEngine)
 
-        fleet_thread = threading.Thread(
-            target=_run_fleet_api_rest,
-            args=(
-                host,
-                port,
-                ssl_keyfile,
-                ssl_certfile,
-                state_factory,
-                ffs_factory,
-                num_workers,
-            ),
-        )
-        fleet_thread.start()
-        bckg_threads.append(fleet_thread)
-    elif args.fleet_api_type == TRANSPORT_TYPE_GRPC_RERE:
-        maybe_keys = _try_setup_node_authentication(args, certificates)
-        interceptors: Optional[Sequence[grpc.ServerInterceptor]] = None
-        if maybe_keys is not None:
-            (
-                node_public_keys,
-                server_private_key,
-                server_public_key,
-            ) = maybe_keys
-            state = state_factory.state()
-            state.store_node_public_keys(node_public_keys)
-            state.store_server_private_public_key(
-                private_key_to_bytes(server_private_key),
-                public_key_to_bytes(server_public_key),
-            )
-            log(
-                INFO,
-                "Node authentication enabled with %d known public keys",
-                len(node_public_keys),
-            )
-            interceptors = [AuthenticateServerInterceptor(state)]
+    bckg_threads = []
 
-        fleet_server = _run_fleet_api_grpc_rere(
-            address=fleet_address,
+    if sim_exec:
+        simulationio_server: grpc.Server = run_simulationio_api_grpc(
+            address=simulationio_address,
             state_factory=state_factory,
             ffs_factory=ffs_factory,
             certificates=certificates,
-            interceptors=interceptors,
         )
-        grpc_servers.append(fleet_server)
-    elif args.fleet_api_type == TRANSPORT_TYPE_GRPC_ADAPTER:
-        fleet_server = _run_fleet_api_grpc_adapter(
-            address=fleet_address,
+        grpc_servers.append(simulationio_server)
+
+    else:
+        # Start ServerAppIo API
+        serverappio_server: grpc.Server = run_serverappio_api_grpc(
+            address=serverappio_address,
             state_factory=state_factory,
             ffs_factory=ffs_factory,
             certificates=certificates,
         )
-        grpc_servers.append(fleet_server)
-    else:
-        raise ValueError(f"Unknown fleet_api_type: {args.fleet_api_type}")
+        grpc_servers.append(serverappio_server)
+
+        # Start Fleet API
+        if not args.fleet_api_address:
+            if args.fleet_api_type in [
+                TRANSPORT_TYPE_GRPC_RERE,
+                TRANSPORT_TYPE_GRPC_ADAPTER,
+            ]:
+                args.fleet_api_address = FLEET_API_GRPC_RERE_DEFAULT_ADDRESS
+            elif args.fleet_api_type == TRANSPORT_TYPE_REST:
+                args.fleet_api_address = FLEET_API_REST_DEFAULT_ADDRESS
+
+        fleet_address, host, port = _format_address(args.fleet_api_address)
+
+        num_workers = args.fleet_api_num_workers
+        if num_workers != 1:
+            log(
+                WARN,
+                "The Fleet API currently supports only 1 worker. "
+                "You have specified %d workers. "
+                "Support for multiple workers will be added in future releases. "
+                "Proceeding with a single worker.",
+                args.fleet_api_num_workers,
+            )
+            num_workers = 1
+
+        if args.fleet_api_type == TRANSPORT_TYPE_REST:
+            if (
+                importlib.util.find_spec("requests")
+                and importlib.util.find_spec("starlette")
+                and importlib.util.find_spec("uvicorn")
+            ) is None:
+                sys.exit(MISSING_EXTRA_REST)
+
+            _, ssl_certfile, ssl_keyfile = (
+                certificates if certificates is not None else (None, None, None)
+            )
+
+            fleet_thread = threading.Thread(
+                target=_run_fleet_api_rest,
+                args=(
+                    host,
+                    port,
+                    ssl_keyfile,
+                    ssl_certfile,
+                    state_factory,
+                    ffs_factory,
+                    num_workers,
+                ),
+            )
+            fleet_thread.start()
+            bckg_threads.append(fleet_thread)
+        elif args.fleet_api_type == TRANSPORT_TYPE_GRPC_RERE:
+            maybe_keys = _try_setup_node_authentication(args, certificates)
+            interceptors: Optional[Sequence[grpc.ServerInterceptor]] = None
+            if maybe_keys is not None:
+                (
+                    node_public_keys,
+                    server_private_key,
+                    server_public_key,
+                ) = maybe_keys
+                state = state_factory.state()
+                state.store_node_public_keys(node_public_keys)
+                state.store_server_private_public_key(
+                    private_key_to_bytes(server_private_key),
+                    public_key_to_bytes(server_public_key),
+                )
+                log(
+                    INFO,
+                    "Node authentication enabled with %d known public keys",
+                    len(node_public_keys),
+                )
+                interceptors = [AuthenticateServerInterceptor(state)]
+
+            fleet_server = _run_fleet_api_grpc_rere(
+                address=fleet_address,
+                state_factory=state_factory,
+                ffs_factory=ffs_factory,
+                certificates=certificates,
+                interceptors=interceptors,
+            )
+            grpc_servers.append(fleet_server)
+        elif args.fleet_api_type == TRANSPORT_TYPE_GRPC_ADAPTER:
+            fleet_server = _run_fleet_api_grpc_adapter(
+                address=fleet_address,
+                state_factory=state_factory,
+                ffs_factory=ffs_factory,
+                certificates=certificates,
+            )
+            grpc_servers.append(fleet_server)
+        else:
+            raise ValueError(f"Unknown fleet_api_type: {args.fleet_api_type}")
+
+        if args.isolation == ISOLATION_MODE_SUBPROCESS:
+            # Scheduler thread
+            scheduler_th = threading.Thread(
+                target=_flwr_serverapp_scheduler,
+                args=(
+                    state_factory,
+                    args.serverappio_api_address,
+                    args.ssl_ca_certfile,
+                ),
+            )
+            scheduler_th.start()
+            bckg_threads.append(scheduler_th)
 
     # Graceful shutdown
     register_exit_handlers(
@@ -331,7 +385,46 @@ def run_superlink() -> None:
             for thread in bckg_threads:
                 if not thread.is_alive():
                     sys.exit(1)
-        driver_server.wait_for_termination(timeout=1)
+        exec_server.wait_for_termination(timeout=1)
+
+
+def _flwr_serverapp_scheduler(
+    state_factory: LinkStateFactory,
+    serverappio_api_address: str,
+    ssl_ca_certfile: Optional[str],
+) -> None:
+    log(DEBUG, "Started flwr-serverapp scheduler thread.")
+
+    state = state_factory.state()
+
+    # Periodically check for a pending run in the LinkState
+    while True:
+        sleep(3)
+        pending_run_id = state.get_pending_run_id()
+
+        if pending_run_id:
+
+            log(
+                INFO,
+                "Launching `flwr-serverapp` subprocess. Connects to SuperLink on %s",
+                serverappio_api_address,
+            )
+            # Start ServerApp subprocess
+            command = [
+                "flwr-serverapp",
+                "--superlink",
+                serverappio_api_address,
+            ]
+            if ssl_ca_certfile:
+                command.append("--root-certificates")
+                command.append(ssl_ca_certfile)
+            else:
+                command.append("--insecure")
+
+            subprocess.Popen(  # pylint: disable=consider-using-with
+                command,
+                text=True,
+            )
 
 
 def _format_address(address: str) -> tuple[str, str, int]:
@@ -489,7 +582,7 @@ def _try_obtain_certificates(
 
 def _run_fleet_api_grpc_rere(
     address: str,
-    state_factory: StateFactory,
+    state_factory: LinkStateFactory,
     ffs_factory: FfsFactory,
     certificates: Optional[tuple[bytes, bytes, bytes]],
     interceptors: Optional[Sequence[grpc.ServerInterceptor]] = None,
@@ -517,7 +610,7 @@ def _run_fleet_api_grpc_rere(
 
 def _run_fleet_api_grpc_adapter(
     address: str,
-    state_factory: StateFactory,
+    state_factory: LinkStateFactory,
     ffs_factory: FfsFactory,
     certificates: Optional[tuple[bytes, bytes, bytes]],
 ) -> grpc.Server:
@@ -548,11 +641,11 @@ def _run_fleet_api_rest(
     port: int,
     ssl_keyfile: Optional[str],
     ssl_certfile: Optional[str],
-    state_factory: StateFactory,
+    state_factory: LinkStateFactory,
     ffs_factory: FfsFactory,
     num_workers: int,
 ) -> None:
-    """Run Driver API (REST-based)."""
+    """Run ServerAppIo API (REST-based)."""
     try:
         import uvicorn
 
@@ -579,14 +672,16 @@ def _run_fleet_api_rest(
 
 
 def _parse_args_run_superlink() -> argparse.ArgumentParser:
-    """Parse command line arguments for both Driver API and Fleet API."""
+    """Parse command line arguments for both ServerAppIo API and Fleet API."""
     parser = argparse.ArgumentParser(
         description="Start a Flower SuperLink",
     )
 
     _add_args_common(parser=parser)
-    _add_args_driver_api(parser=parser)
+    _add_args_serverappio_api(parser=parser)
     _add_args_fleet_api(parser=parser)
+    _add_args_exec_api(parser=parser)
+    _add_args_simulationio_api(parser=parser)
 
     return parser
 
@@ -618,6 +713,19 @@ def _add_args_common(parser: argparse.ArgumentParser) -> None:
         "to create a secure connection.",
         type=str,
     )
+    parser.add_argument(
+        "--isolation",
+        default=ISOLATION_MODE_SUBPROCESS,
+        required=False,
+        choices=[
+            ISOLATION_MODE_SUBPROCESS,
+            ISOLATION_MODE_PROCESS,
+        ],
+        help="Isolation mode when running a `ServerApp` (`subprocess` by default, "
+        "possible values: `subprocess`, `process`). Use `subprocess` to configure "
+        "SuperLink to run a `ServerApp` in a subprocess. Use `process` to indicate "
+        "that a separate independent process gets created outside of SuperLink.",
+    )
     parser.add_argument(
         "--database",
         help="A string representing the path to the database "
@@ -650,11 +758,11 @@ def _add_args_common(parser: argparse.ArgumentParser) -> None:
     )
 
 
-def _add_args_driver_api(parser: argparse.ArgumentParser) -> None:
+def _add_args_serverappio_api(parser: argparse.ArgumentParser) -> None:
     parser.add_argument(
-        "--driver-api-address",
-        help="Driver API (gRPC) server address (IPv4, IPv6, or a domain name).",
-        default=DRIVER_API_DEFAULT_ADDRESS,
+        "--serverappio-api-address",
+        help="ServerAppIo API (gRPC) server address (IPv4, IPv6, or a domain name).",
+        default=SERVERAPPIO_API_DEFAULT_ADDRESS,
     )
 
 
@@ -681,3 +789,37 @@ def _add_args_fleet_api(parser: argparse.ArgumentParser) -> None:
         type=int,
         help="Set the number of concurrent workers for the Fleet API server.",
     )
+
+
+def _add_args_exec_api(parser: argparse.ArgumentParser) -> None:
+    """Add command line arguments for Exec API."""
+    parser.add_argument(
+        "--exec-api-address",
+        help="Exec API server address (IPv4, IPv6, or a domain name)",
+        default=EXEC_API_DEFAULT_ADDRESS,
+    )
+    parser.add_argument(
+        "--executor",
+        help="For example: `deployment:exec` or `project.package.module:wrapper.exec`. "
+        "The default is `flwr.superexec.deployment:executor`",
+        default="flwr.superexec.deployment:executor",
+    )
+    parser.add_argument(
+        "--executor-dir",
+        help="The directory for the executor.",
+        default=".",
+    )
+    parser.add_argument(
+        "--executor-config",
+        help="Key-value pairs for the executor config, separated by spaces. "
+        "For example:\n\n`--executor-config 'verbose=true "
+        'root-certificates="certificates/superlink-ca.crt"\'`',
+    )
+
+
+def _add_args_simulationio_api(parser: argparse.ArgumentParser) -> None:
+    parser.add_argument(
+        "--simulationio-api-address",
+        help="SimulationIo API (gRPC) server address (IPv4, IPv6, or a domain name).",
+        default=SIMULATIONIO_API_DEFAULT_ADDRESS,
+    )

flwr/server/driver/driver.py

@@ -24,7 +24,21 @@ from flwr.common.typing import Run
 
 
 class Driver(ABC):
-    """Abstract base Driver class for the Driver API."""
+    """Abstract base Driver class for the ServerAppIo API."""
+
+    @abstractmethod
+    def init_run(self, run_id: int) -> None:
+        """Request a run to the SuperLink with a given `run_id`.
+
+        If a Run with the specified `run_id` exists, a local Run
+        object will be created. It enables further functionality
+        in the driver, such as sending `Messages`.
+
+        Parameters
+        ----------
+        run_id : int
+            The `run_id` of the Run this Driver object operates in.
+        """
 
     @property
     @abstractmethod