flwr-nightly 1.10.0.dev20240707__py3-none-any.whl → 1.10.0.dev20240722__py3-none-any.whl

flwr/simulation/run_simulation.py

@@ -18,45 +18,157 @@ import argparse
 import asyncio
 import json
 import logging
+import sys
 import threading
 import traceback
+from argparse import Namespace
 from logging import DEBUG, ERROR, INFO, WARNING
+from pathlib import Path
 from time import sleep
-from typing import Optional
+from typing import List, Optional
 
+from flwr.cli.config_utils import load_and_validate
 from flwr.client import ClientApp
 from flwr.common import EventType, event, log
+from flwr.common.config import get_fused_config_from_dir, parse_config_args
+from flwr.common.constant import RUN_ID_NUM_BYTES
 from flwr.common.logger import set_logger_propagation, update_console_handler
-from flwr.common.typing import Run
+from flwr.common.typing import Run, UserConfig
 from flwr.server.driver import Driver, InMemoryDriver
-from flwr.server.run_serverapp import run
+from flwr.server.run_serverapp import run as run_server_app
 from flwr.server.server_app import ServerApp
 from flwr.server.superlink.fleet import vce
 from flwr.server.superlink.fleet.vce.backend.backend import BackendConfig
 from flwr.server.superlink.state import StateFactory
+from flwr.server.superlink.state.utils import generate_rand_int_from_bytes
 from flwr.simulation.ray_transport.utils import (
     enable_tf_gpu_growth as enable_gpu_growth,
 )
 
 
+def _check_args_do_not_interfere(args: Namespace) -> bool:
+    """Ensure decoupling of flags for different ways to start the simulation."""
+    mode_one_args = ["app", "run_config"]
+    mode_two_args = ["client_app", "server_app"]
+
+    def _resolve_message(conflict_keys: List[str]) -> str:
+        return ",".join([f"`--{key}`".replace("_", "-") for key in conflict_keys])
+
+    # When passing `--app`, `--app-dir` is ignored
+    if args.app and args.app_dir:
+        log(ERROR, "Either `--app` or `--app-dir` can be set, but not both.")
+        return False
+
+    if any(getattr(args, key) for key in mode_one_args):
+        if any(getattr(args, key) for key in mode_two_args):
+            log(
+                ERROR,
+                "Passing any of {%s} alongside with any of {%s}",
+                _resolve_message(mode_one_args),
+                _resolve_message(mode_two_args),
+            )
+            return False
+
+        if not args.app:
+            log(ERROR, "You need to pass --app")
+            return False
+
+        return True
+
+    # Ensure all args are set (required for the non-FAB mode of execution)
+    if not all(getattr(args, key) for key in mode_two_args):
+        log(
+            ERROR,
+            "Passing all of %s keys are required.",
+            _resolve_message(mode_two_args),
+        )
+        return False
+
+    return True
+
+
 # Entry point from CLI
+# pylint: disable=too-many-locals
 def run_simulation_from_cli() -> None:
     """Run Simulation Engine from the CLI."""
     args = _parse_args_run_simulation().parse_args()
 
+    # We are supporting two modes for the CLI entrypoint:
+    # 1) Running an app dir containing a `pyproject.toml`
+    # 2) Running any ClientApp and SeverApp w/o pyproject.toml being present
+    # For 2), some CLI args are compulsory, but they are not required for 1)
+    # We first do these checks
+    args_check_pass = _check_args_do_not_interfere(args)
+    if not args_check_pass:
+        sys.exit("Simulation Engine cannot start.")
+
+    run_id = (
+        generate_rand_int_from_bytes(RUN_ID_NUM_BYTES)
+        if args.run_id is None
+        else args.run_id
+    )
+    if args.app:
+        # Mode 1
+        app_path = Path(args.app)
+        if not app_path.is_dir():
+            log(ERROR, "--app is not a directory")
+            sys.exit("Simulation Engine cannot start.")
+
+        # Load pyproject.toml
+        config, errors, warnings = load_and_validate(
+            app_path / "pyproject.toml", check_module=False
+        )
+        if errors:
+            raise ValueError(errors)
+
+        if warnings:
+            log(WARNING, warnings)
+
+        if config is None:
+            raise ValueError("Config extracted from FAB's pyproject.toml is not valid")
+
+        # Get ClientApp and SeverApp components
+        app_components = config["tool"]["flwr"]["app"]["components"]
+        client_app_attr = app_components["clientapp"]
+        server_app_attr = app_components["serverapp"]
+
+        override_config = parse_config_args([args.run_config])
+        fused_config = get_fused_config_from_dir(app_path, override_config)
+        app_dir = args.app
+        is_app = True
+
+    else:
+        # Mode 2
+        client_app_attr = args.client_app
+        server_app_attr = args.server_app
+        override_config = {}
+        fused_config = None
+        app_dir = args.app_dir
+        is_app = False
+
+    # Create run
+    run = Run(
+        run_id=run_id,
+        fab_id="",
+        fab_version="",
+        override_config=override_config,
+    )
+
     # Load JSON config
     backend_config_dict = json.loads(args.backend_config)
 
     _run_simulation(
-        server_app_attr=args.server_app,
-        client_app_attr=args.client_app,
+        server_app_attr=server_app_attr,
+        client_app_attr=client_app_attr,
         num_supernodes=args.num_supernodes,
         backend_name=args.backend,
         backend_config=backend_config_dict,
-        app_dir=args.app_dir,
-        run_id=args.run_id,
+        app_dir=app_dir,
+        run=run,
         enable_tf_gpu_growth=args.enable_tf_gpu_growth,
         verbose_logging=args.verbose,
+        server_app_run_config=fused_config,
+        is_app=is_app,
     )
 
 
@@ -126,16 +238,25 @@ def run_simulation(
 def run_serverapp_th(
     server_app_attr: Optional[str],
     server_app: Optional[ServerApp],
+    server_app_run_config: UserConfig,
     driver: Driver,
     app_dir: str,
-    f_stop: asyncio.Event,
+    f_stop: threading.Event,
+    has_exception: threading.Event,
     enable_tf_gpu_growth: bool,
     delay_launch: int = 3,
 ) -> threading.Thread:
     """Run SeverApp in a thread."""
 
-    def server_th_with_start_checks(  # type: ignore
-        tf_gpu_growth: bool, stop_event: asyncio.Event, **kwargs
+    def server_th_with_start_checks(
+        tf_gpu_growth: bool,
+        stop_event: threading.Event,
+        exception_event: threading.Event,
+        _driver: Driver,
+        _server_app_dir: str,
+        _server_app_run_config: UserConfig,
+        _server_app_attr: Optional[str],
+        _server_app: Optional[ServerApp],
     ) -> None:
         """Run SeverApp, after check if GPU memory growth has to be set.
 
@@ -147,10 +268,18 @@ def run_serverapp_th(
                 enable_gpu_growth()
 
             # Run ServerApp
-            run(**kwargs)
+            run_server_app(
+                driver=_driver,
+                server_app_dir=_server_app_dir,
+                server_app_run_config=_server_app_run_config,
+                server_app_attr=_server_app_attr,
+                loaded_server_app=_server_app,
+            )
         except Exception as ex:  # pylint: disable=broad-exception-caught
             log(ERROR, "ServerApp thread raised an exception: %s", ex)
             log(ERROR, traceback.format_exc())
+            exception_event.set()
+            raise
         finally:
             log(DEBUG, "ServerApp finished running.")
             # Upon completion, trigger stop event if one was passed
@@ -160,71 +289,66 @@ def run_serverapp_th(
 
     serverapp_th = threading.Thread(
         target=server_th_with_start_checks,
-        args=(enable_tf_gpu_growth, f_stop),
-        kwargs={
-            "server_app_attr": server_app_attr,
-            "loaded_server_app": server_app,
-            "driver": driver,
-            "server_app_dir": app_dir,
-        },
+        args=(
+            enable_tf_gpu_growth,
+            f_stop,
+            has_exception,
+            driver,
+            app_dir,
+            server_app_run_config,
+            server_app_attr,
+            server_app,
+        ),
     )
     sleep(delay_launch)
     serverapp_th.start()
     return serverapp_th
 
 
-def _override_run_id(state: StateFactory, run_id_to_replace: int, run_id: int) -> None:
-    """Override the run_id of an existing Run."""
-    log(DEBUG, "Pre-registering run with id %s", run_id)
-    # Remove run
-    run_info: Run = state.state().run_ids.pop(run_id_to_replace)  # type: ignore
-    # Update with new run_id and insert back in state
-    run_info.run_id = run_id
-    state.state().run_ids[run_id] = run_info  # type: ignore
-
-
 # pylint: disable=too-many-locals
 def _main_loop(
     num_supernodes: int,
     backend_name: str,
     backend_config_stream: str,
     app_dir: str,
+    is_app: bool,
     enable_tf_gpu_growth: bool,
-    run_id: Optional[int] = None,
+    run: Run,
+    flwr_dir: Optional[str] = None,
     client_app: Optional[ClientApp] = None,
     client_app_attr: Optional[str] = None,
     server_app: Optional[ServerApp] = None,
     server_app_attr: Optional[str] = None,
+    server_app_run_config: Optional[UserConfig] = None,
 ) -> None:
-    """Launch SuperLink with Simulation Engine, then ServerApp on a separate thread.
-
-    Everything runs on the main thread or a separate one, depending on whether the main
-    thread already contains a running Asyncio event loop. This is the case if running
-    the Simulation Engine on a Jupyter/Colab notebook.
-    """
+    """Launch SuperLink with Simulation Engine, then ServerApp on a separate thread."""
     # Initialize StateFactory
     state_factory = StateFactory(":flwr-in-memory-state:")
 
-    f_stop = asyncio.Event()
+    f_stop = threading.Event()
+    # A Threading event to indicate if an exception was raised in the ServerApp thread
+    server_app_thread_has_exception = threading.Event()
     serverapp_th = None
     try:
-        # Create run (with empty fab_id and fab_version)
-        run_id_ = state_factory.state().create_run("", "")
+        # Register run
+        log(DEBUG, "Pre-registering run with id %s", run.run_id)
+        state_factory.state().run_ids[run.run_id] = run  # type: ignore
 
-        if run_id:
-            _override_run_id(state_factory, run_id_to_replace=run_id_, run_id=run_id)
-            run_id_ = run_id
+        if server_app_run_config is None:
+            server_app_run_config = {}
 
         # Initialize Driver
-        driver = InMemoryDriver(run_id=run_id_, state_factory=state_factory)
+        driver = InMemoryDriver(run_id=run.run_id, state_factory=state_factory)
 
         # Get and run ServerApp thread
         serverapp_th = run_serverapp_th(
             server_app_attr=server_app_attr,
             server_app=server_app,
+            server_app_run_config=server_app_run_config,
             driver=driver,
             app_dir=app_dir,
             f_stop=f_stop,
+            has_exception=server_app_thread_has_exception,
             enable_tf_gpu_growth=enable_tf_gpu_growth,
         )
 
@@ -237,8 +361,11 @@ def _main_loop(
             backend_name=backend_name,
             backend_config_json_stream=backend_config_stream,
             app_dir=app_dir,
+            is_app=is_app,
             state_factory=state_factory,
             f_stop=f_stop,
+            run=run,
+            flwr_dir=flwr_dir,
         )
 
     except Exception as ex:
@@ -253,6 +380,8 @@ def _main_loop(
         event(EventType.RUN_SUPERLINK_LEAVE)
         if serverapp_th:
             serverapp_th.join()
+            if server_app_thread_has_exception.is_set():
+                raise RuntimeError("Exception in ServerApp thread")
 
     log(DEBUG, "Stopping Simulation Engine now.")
 
@@ -266,10 +395,13 @@ def _run_simulation(
     backend_config: Optional[BackendConfig] = None,
     client_app_attr: Optional[str] = None,
     server_app_attr: Optional[str] = None,
+    server_app_run_config: Optional[UserConfig] = None,
     app_dir: str = "",
-    run_id: Optional[int] = None,
+    flwr_dir: Optional[str] = None,
+    run: Optional[Run] = None,
     enable_tf_gpu_growth: bool = False,
     verbose_logging: bool = False,
+    is_app: bool = False,
 ) -> None:
     r"""Launch the Simulation Engine.
 
@@ -298,20 +430,27 @@ def _run_simulation(
         parameters. Values supported in <value> are those included by
         `flwr.common.typing.ConfigsRecordValues`.
 
-    client_app_attr : str
+    client_app_attr : Optional[str]
         A path to a `ClientApp` module to be loaded: For example: `client:app` or
         `project.package.module:wrapper.app`."
 
-    server_app_attr : str
+    server_app_attr : Optional[str]
         A path to a `ServerApp` module to be loaded: For example: `server:app` or
         `project.package.module:wrapper.app`."
 
+    server_app_run_config : Optional[UserConfig]
+        Config dictionary that parameterizes the run config. It will be made accesible
+        to the ServerApp.
+
     app_dir : str
         Add specified directory to the PYTHONPATH and load `ClientApp` from there.
         (Default: current working directory.)
 
-    run_id : Optional[int]
-        An integer specifying the ID of the run started when running this function.
+    flwr_dir : Optional[str]
+        The path containing installed Flower Apps.
+
+    run : Optional[Run]
+        An object carrying details about the run.
 
     enable_tf_gpu_growth : bool (default: False)
         A boolean to indicate whether to enable GPU growth on the main thread. This is
@@ -324,6 +463,11 @@ def _run_simulation(
     verbose_logging : bool (default: False)
         When disabled, only INFO, WARNING and ERROR log messages will be shown. If
         enabled, DEBUG-level logs will be displayed.
+
+    is_app : bool (default: False)
+        A flag that indicates whether the simulation is running an app or not. This is
+        needed in order to attempt loading an app's pyproject.toml when nodes register
+        a context object.
     """
     if backend_config is None:
         backend_config = {}
@@ -336,8 +480,12 @@ def _run_simulation(
     if verbose_logging:
         update_console_handler(level=DEBUG, timestamps=True, colored=True)
     else:
-        backend_config["init_args"]["logging_level"] = WARNING
-        backend_config["init_args"]["log_to_driver"] = True
+        backend_config["init_args"]["logging_level"] = backend_config["init_args"].get(
+            "logging_level", WARNING
+        )
+        backend_config["init_args"]["log_to_driver"] = backend_config["init_args"].get(
+            "log_to_driver", True
+        )
 
     if enable_tf_gpu_growth:
         # Check that Backend config has also enabled using GPU growth
@@ -349,45 +497,47 @@ def _run_simulation(
     # Convert config to original JSON-stream format
     backend_config_stream = json.dumps(backend_config)
 
-    simulation_engine_th = None
+    # If no `Run` object is set, create one
+    if run is None:
+        run_id = generate_rand_int_from_bytes(RUN_ID_NUM_BYTES)
+        run = Run(run_id=run_id, fab_id="", fab_version="", override_config={})
+
     args = (
         num_supernodes,
         backend_name,
         backend_config_stream,
         app_dir,
+        is_app,
         enable_tf_gpu_growth,
-        run_id,
+        run,
+        flwr_dir,
         client_app,
         client_app_attr,
         server_app,
         server_app_attr,
+        server_app_run_config,
     )
     # Detect if there is an Asyncio event loop already running.
-    # If yes, run everything on a separate thread. In environments
-    # like Jupyter/Colab notebooks, there is an event loop present.
-    run_in_thread = False
+    # If yes, disable logger propagation. In environmnets
+    # like Jupyter/Colab notebooks, it's often better to do this.
+    asyncio_loop_running = False
     try:
         _ = (
             asyncio.get_running_loop()
         )  # Raises RuntimeError if no event loop is present
         log(DEBUG, "Asyncio event loop already running.")
 
-        run_in_thread = True
+        asyncio_loop_running = True
 
     except RuntimeError:
-        log(DEBUG, "No asyncio event loop running")
+        pass
 
     finally:
-        if run_in_thread:
+        if asyncio_loop_running:
             # Set logger propagation to False to prevent duplicated log output in Colab.
             logger = set_logger_propagation(logger, False)
-            log(DEBUG, "Starting Simulation Engine on a new thread.")
-            simulation_engine_th = threading.Thread(target=_main_loop, args=args)
-            simulation_engine_th.start()
-            simulation_engine_th.join()
-        else:
-            log(DEBUG, "Starting Simulation Engine on the main thread.")
-            _main_loop(*args)
+
+        _main_loop(*args)
 
 
 def _parse_args_run_simulation() -> argparse.ArgumentParser:
@@ -397,12 +547,10 @@ def _parse_args_run_simulation() -> argparse.ArgumentParser:
     )
     parser.add_argument(
         "--server-app",
-        required=True,
         help="For example: `server:app` or `project.package.module:wrapper.app`",
     )
     parser.add_argument(
         "--client-app",
-        required=True,
         help="For example: `client:app` or `project.package.module:wrapper.app`",
     )
     parser.add_argument(
@@ -411,6 +559,18 @@ def _parse_args_run_simulation() -> argparse.ArgumentParser:
         required=True,
         help="Number of simulated SuperNodes.",
     )
+    parser.add_argument(
+        "--app",
+        type=str,
+        default=None,
+        help="Path to a directory containing a FAB-like structure with a "
+        "pyproject.toml.",
+    )
+    parser.add_argument(
+        "--run-config",
+        default=None,
+        help="Override configuration key-value pairs.",
+    )
     parser.add_argument(
         "--backend",
         default="ray",
@@ -449,6 +609,17 @@ def _parse_args_run_simulation() -> argparse.ArgumentParser:
         "ClientApp and ServerApp from there."
         " Default: current working directory.",
     )
+    parser.add_argument(
+        "--flwr-dir",
+        default=None,
+        help="""The path containing installed Flower Apps.
+    By default, this value is equal to:
+
+        - `$FLWR_HOME/` if `$FLWR_HOME` is defined
+        - `$XDG_DATA_HOME/.flwr/` if `$XDG_DATA_HOME` is defined
+        - `$HOME/.flwr/` in all other cases
+    """,
+    )
     parser.add_argument(
         "--run-id",
         type=int,

flwr/superexec/app.py

@@ -24,6 +24,7 @@ import grpc
 
 from flwr.common import EventType, event, log
 from flwr.common.address import parse_address
+from flwr.common.config import parse_config_args
 from flwr.common.constant import SUPEREXEC_DEFAULT_ADDRESS
 from flwr.common.exit_handlers import register_exit_handlers
 from flwr.common.object_ref import load_app, validate
@@ -55,6 +56,7 @@ def run_superexec() -> None:
         address=address,
         executor=_load_executor(args),
         certificates=certificates,
+        config=parse_config_args([args.executor_config]),
     )
 
     grpc_servers = [superexec_server]
@@ -74,20 +76,25 @@ def _parse_args_run_superexec() -> argparse.ArgumentParser:
     parser = argparse.ArgumentParser(
         description="Start a Flower SuperExec",
     )
-    parser.add_argument(
-        "executor",
-        help="For example: `deployment:exec` or `project.package.module:wrapper.exec`.",
-    )
     parser.add_argument(
         "--address",
         help="SuperExec (gRPC) server address (IPv4, IPv6, or a domain name)",
         default=SUPEREXEC_DEFAULT_ADDRESS,
     )
+    parser.add_argument(
+        "--executor",
+        help="For example: `deployment:exec` or `project.package.module:wrapper.exec`.",
+        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 commas.",
+    )
     parser.add_argument(
         "--insecure",
         action="store_true",
@@ -126,11 +133,11 @@ def _try_obtain_certificates(
         return None
     # Check if certificates are provided
     if args.ssl_certfile and args.ssl_keyfile and args.ssl_ca_certfile:
-        if not Path.is_file(args.ssl_ca_certfile):
+        if not Path(args.ssl_ca_certfile).is_file():
             sys.exit("Path argument `--ssl-ca-certfile` does not point to a file.")
-        if not Path.is_file(args.ssl_certfile):
+        if not Path(args.ssl_certfile).is_file():
             sys.exit("Path argument `--ssl-certfile` does not point to a file.")
-        if not Path.is_file(args.ssl_keyfile):
+        if not Path(args.ssl_keyfile).is_file():
             sys.exit("Path argument `--ssl-keyfile` does not point to a file.")
         certificates = (
             Path(args.ssl_ca_certfile).read_bytes(),  # CA certificate