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

flwr/server/superlink/simulation/simulationio_grpc.py

@@ -0,0 +1,65 @@
+# Copyright 2024 Flower Labs GmbH. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""SimulationIo gRPC API."""
+
+
+from logging import INFO
+from typing import Optional
+
+import grpc
+
+from flwr.common import GRPC_MAX_MESSAGE_LENGTH
+from flwr.common.logger import log
+from flwr.proto.simulationio_pb2_grpc import (  # pylint: disable=E0611
+    add_SimulationIoServicer_to_server,
+)
+from flwr.server.superlink.ffs.ffs_factory import FfsFactory
+from flwr.server.superlink.linkstate import LinkStateFactory
+
+from ..fleet.grpc_bidi.grpc_server import generic_create_grpc_server
+from .simulationio_servicer import SimulationIoServicer
+
+
+def run_simulationio_api_grpc(
+    address: str,
+    state_factory: LinkStateFactory,
+    ffs_factory: FfsFactory,
+    certificates: Optional[tuple[bytes, bytes, bytes]],
+) -> grpc.Server:
+    """Run SimulationIo API (gRPC, request-response)."""
+    # Create SimulationIo API gRPC server
+    simulationio_servicer: grpc.Server = SimulationIoServicer(
+        state_factory=state_factory,
+        ffs_factory=ffs_factory,
+    )
+    simulationio_add_servicer_to_server_fn = add_SimulationIoServicer_to_server
+    simulationio_grpc_server = generic_create_grpc_server(
+        servicer_and_add_fn=(
+            simulationio_servicer,
+            simulationio_add_servicer_to_server_fn,
+        ),
+        server_address=address,
+        max_message_length=GRPC_MAX_MESSAGE_LENGTH,
+        certificates=certificates,
+    )
+
+    log(
+        INFO,
+        "Flower Simulation Engine: Starting SimulationIo API on %s",
+        address,
+    )
+    simulationio_grpc_server.start()
+
+    return simulationio_grpc_server

flwr/server/superlink/simulation/simulationio_servicer.py

@@ -0,0 +1,132 @@
+# Copyright 2024 Flower Labs GmbH. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""SimulationIo API servicer."""
+
+import threading
+from logging import DEBUG, INFO
+
+import grpc
+from grpc import ServicerContext
+
+from flwr.common.constant import Status
+from flwr.common.logger import log
+from flwr.common.serde import (
+    context_from_proto,
+    context_to_proto,
+    fab_to_proto,
+    run_status_from_proto,
+    run_to_proto,
+)
+from flwr.common.typing import Fab, RunStatus
+from flwr.proto import simulationio_pb2_grpc
+from flwr.proto.log_pb2 import (  # pylint: disable=E0611
+    PushLogsRequest,
+    PushLogsResponse,
+)
+from flwr.proto.run_pb2 import (  # pylint: disable=E0611
+    UpdateRunStatusRequest,
+    UpdateRunStatusResponse,
+)
+from flwr.proto.simulationio_pb2 import (  # pylint: disable=E0611
+    PullSimulationInputsRequest,
+    PullSimulationInputsResponse,
+    PushSimulationOutputsRequest,
+    PushSimulationOutputsResponse,
+)
+from flwr.server.superlink.ffs.ffs_factory import FfsFactory
+from flwr.server.superlink.linkstate import LinkStateFactory
+
+
+class SimulationIoServicer(simulationio_pb2_grpc.SimulationIoServicer):
+    """SimulationIo API servicer."""
+
+    def __init__(
+        self, state_factory: LinkStateFactory, ffs_factory: FfsFactory
+    ) -> None:
+        self.state_factory = state_factory
+        self.ffs_factory = ffs_factory
+        self.lock = threading.RLock()
+
+    def PullSimulationInputs(
+        self, request: PullSimulationInputsRequest, context: ServicerContext
+    ) -> PullSimulationInputsResponse:
+        """Pull SimultionIo process inputs."""
+        log(DEBUG, "SimultionIoServicer.SimultionIoInputs")
+        # Init access to LinkState and Ffs
+        state = self.state_factory.state()
+        ffs = self.ffs_factory.ffs()
+
+        # Lock access to LinkState, preventing obtaining the same pending run_id
+        with self.lock:
+            # Attempt getting the run_id of a pending run
+            run_id = state.get_pending_run_id()
+            # If there's no pending run, return an empty response
+            if run_id is None:
+                return PullSimulationInputsResponse()
+
+            # Retrieve Context, Run and Fab for the run_id
+            serverapp_ctxt = state.get_serverapp_context(run_id)
+            run = state.get_run(run_id)
+            fab = None
+            if run and run.fab_hash:
+                if result := ffs.get(run.fab_hash):
+                    fab = Fab(run.fab_hash, result[0])
+            if run and fab and serverapp_ctxt:
+                # Update run status to STARTING
+                if state.update_run_status(run_id, RunStatus(Status.STARTING, "", "")):
+                    log(INFO, "Starting run %d", run_id)
+                    return PullSimulationInputsResponse(
+                        context=context_to_proto(serverapp_ctxt),
+                        run=run_to_proto(run),
+                        fab=fab_to_proto(fab),
+                    )
+
+        # Raise an exception if the Run or Fab is not found,
+        # or if the status cannot be updated to STARTING
+        raise RuntimeError(f"Failed to start run {run_id}")
+
+    def PushSimulationOutputs(
+        self, request: PushSimulationOutputsRequest, context: ServicerContext
+    ) -> PushSimulationOutputsResponse:
+        """Push Simulation process outputs."""
+        log(DEBUG, "SimultionIoServicer.PushSimulationOutputs")
+        state = self.state_factory.state()
+        state.set_serverapp_context(request.run_id, context_from_proto(request.context))
+        return PushSimulationOutputsResponse()
+
+    def UpdateRunStatus(
+        self, request: UpdateRunStatusRequest, context: grpc.ServicerContext
+    ) -> UpdateRunStatusResponse:
+        """Update the status of a run."""
+        log(DEBUG, "SimultionIoServicer.UpdateRunStatus")
+        state = self.state_factory.state()
+
+        # Update the run status
+        state.update_run_status(
+            run_id=request.run_id, new_status=run_status_from_proto(request.run_status)
+        )
+        return UpdateRunStatusResponse()
+
+    def PushLogs(
+        self, request: PushLogsRequest, context: grpc.ServicerContext
+    ) -> PushLogsResponse:
+        """Push logs."""
+        log(DEBUG, "ServerAppIoServicer.PushLogs")
+        state = self.state_factory.state()
+
+        # Add logs to LinkState
+        merged_logs = "".join(request.logs)
+        state.add_serverapp_log(request.run_id, merged_logs)
+        return PushLogsResponse()

flwr/simulation/__init__.py

@@ -18,6 +18,7 @@
 import importlib
 
 from flwr.simulation.run_simulation import run_simulation
+from flwr.simulation.simulationio_connection import SimulationIoConnection
 
 is_ray_installed = importlib.util.find_spec("ray") is not None
 
@@ -37,6 +38,7 @@ To install the necessary dependencies, install `flwr` with the `simulation` extr
 
 
 __all__ = [
+    "SimulationIoConnection",
     "run_simulation",
     "start_simulation",
 ]

flwr/simulation/app.py

@@ -36,7 +36,7 @@ from flwr.server.history import History
 from flwr.server.server import Server, init_defaults, run_fl
 from flwr.server.server_config import ServerConfig
 from flwr.server.strategy import Strategy
-from flwr.server.superlink.state.utils import generate_rand_int_from_bytes
+from flwr.server.superlink.linkstate.utils import generate_rand_int_from_bytes
 from flwr.simulation.ray_transport.ray_actor import (
     ClientAppActor,
     VirtualClientEngineActor,

flwr/simulation/ray_transport/ray_client_proxy.py

@@ -22,7 +22,7 @@ from typing import Optional
 from flwr import common
 from flwr.client import ClientFnExt
 from flwr.client.client_app import ClientApp
-from flwr.client.node_state import NodeState
+from flwr.client.run_info_store import DeprecatedRunInfoStore
 from flwr.common import DEFAULT_TTL, Message, Metadata, RecordSet
 from flwr.common.constant import (
     NUM_PARTITIONS_KEY,
@@ -65,7 +65,7 @@ class RayActorClientProxy(ClientProxy):
 
         self.app_fn = _load_app
         self.actor_pool = actor_pool
-        self.proxy_state = NodeState(
+        self.proxy_state = DeprecatedRunInfoStore(
             node_id=node_id,
             node_config={
                 PARTITION_ID_KEY: str(partition_id),

flwr/simulation/run_simulation.py

@@ -21,7 +21,6 @@ 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
@@ -29,69 +28,28 @@ from typing import Any, 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 import Context, EventType, RecordSet, event, log, now
 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.constant import RUN_ID_NUM_BYTES, Status
 from flwr.common.logger import (
     set_logger_propagation,
     update_console_handler,
-    warn_deprecated_feature,
     warn_deprecated_feature_with_example,
 )
-from flwr.common.typing import Run, UserConfig
+from flwr.common.typing import Run, RunStatus, UserConfig
 from flwr.server.driver import Driver, InMemoryDriver
-from flwr.server.run_serverapp import run as run_server_app
+from flwr.server.run_serverapp import run as _run
 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.server.superlink.linkstate import LinkStateFactory
+from flwr.server.superlink.linkstate.in_memory_linkstate import RunRecord
+from flwr.server.superlink.linkstate.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
-
-
 def _replace_keys(d: Any, match: str, target: str) -> Any:
     if isinstance(d, dict):
         return {
@@ -114,19 +72,6 @@ def run_simulation_from_cli() -> None:
         event_details={"backend": args.backend, "num-supernodes": args.num_supernodes},
     )
 
-    # Add warnings for deprecated server_app and client_app arguments
-    if args.server_app:
-        warn_deprecated_feature(
-            "The `--server-app` argument is deprecated. "
-            "Please use the `--app` argument instead."
-        )
-
-    if args.client_app:
-        warn_deprecated_feature(
-            "The `--client-app` argument is deprecated. "
-            "Use the `--app` argument instead."
-        )
-
     if args.enable_tf_gpu_growth:
         warn_deprecated_feature_with_example(
             "Passing `--enable-tf-gpu-growth` is deprecated.",
@@ -143,60 +88,39 @@ def run_simulation_from_cli() -> None:
         backend_config_dict = _replace_keys(backend_config_dict, match="-", target="_")
         log(DEBUG, "backend_config_dict: %s", backend_config_dict)
 
-    # 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)
+    app_path = Path(args.app)
+    if not app_path.is_dir():
+        log(ERROR, "--app is not a directory")
+        sys.exit("Simulation Engine cannot start.")
 
-        if config is None:
-            raise ValueError("Config extracted from FAB's pyproject.toml is not valid")
+    # Load pyproject.toml
+    config, errors, warnings = load_and_validate(
+        app_path / "pyproject.toml", check_module=False
+    )
+    if errors:
+        raise ValueError(errors)
 
-        # Get ClientApp and SeverApp components
-        app_components = config["tool"]["flwr"]["app"]["components"]
-        client_app_attr = app_components["clientapp"]
-        server_app_attr = app_components["serverapp"]
+    if warnings:
+        log(WARNING, warnings)
 
-        override_config = parse_config_args(
-            [args.run_config] if args.run_config else args.run_config
-        )
-        fused_config = get_fused_config_from_dir(app_path, override_config)
-        app_dir = args.app
-        is_app = True
+    if config is None:
+        raise ValueError("Config extracted from FAB's pyproject.toml is not valid")
 
-    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
+    # 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] if args.run_config else args.run_config
+    )
+    fused_config = get_fused_config_from_dir(app_path, override_config)
 
     # Create run
     run = Run(
@@ -213,13 +137,13 @@ def run_simulation_from_cli() -> None:
         num_supernodes=args.num_supernodes,
         backend_name=args.backend,
         backend_config=backend_config_dict,
-        app_dir=app_dir,
+        app_dir=args.app,
         run=run,
         enable_tf_gpu_growth=args.enable_tf_gpu_growth,
         delay_start=args.delay_start,
         verbose_logging=args.verbose,
         server_app_run_config=fused_config,
-        is_app=is_app,
+        is_app=True,
         exit_event=EventType.CLI_FLOWER_SIMULATION_LEAVE,
     )
 
@@ -310,6 +234,7 @@ def run_serverapp_th(
     f_stop: threading.Event,
     has_exception: threading.Event,
     enable_tf_gpu_growth: bool,
+    run_id: int,
 ) -> threading.Thread:
     """Run SeverApp in a thread."""
 
@@ -332,11 +257,20 @@ def run_serverapp_th(
                 log(INFO, "Enabling GPU growth for Tensorflow on the server thread.")
                 enable_gpu_growth()
 
+            # Initialize Context
+            context = Context(
+                run_id=run_id,
+                node_id=0,
+                node_config={},
+                state=RecordSet(),
+                run_config=_server_app_run_config,
+            )
+
             # Run ServerApp
-            run_server_app(
+            _run(
                 driver=_driver,
+                context=context,
                 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,
             )
@@ -389,7 +323,7 @@ def _main_loop(
 ) -> None:
     """Start ServerApp on a separate thread, then launch Simulation Engine."""
     # Initialize StateFactory
-    state_factory = StateFactory(":flwr-in-memory-state:")
+    state_factory = LinkStateFactory(":flwr-in-memory-state:")
 
     f_stop = threading.Event()
     # A Threading event to indicate if an exception was raised in the ServerApp thread
@@ -399,13 +333,21 @@ def _main_loop(
     try:
         # Register run
         log(DEBUG, "Pre-registering run with id %s", run.run_id)
-        state_factory.state().run_ids[run.run_id] = run  # type: ignore
+        init_status = RunStatus(Status.RUNNING, "", "")
+        state_factory.state().run_ids[run.run_id] = RunRecord(  # type: ignore
+            run=run,
+            status=init_status,
+            starting_at=now().isoformat(),
+            running_at=now().isoformat(),
+            finished_at="",
+        )
 
         if server_app_run_config is None:
             server_app_run_config = {}
 
         # Initialize Driver
-        driver = InMemoryDriver(run_id=run.run_id, state_factory=state_factory)
+        driver = InMemoryDriver(state_factory=state_factory)
+        driver.init_run(run_id=run.run_id)
 
         # Get and run ServerApp thread
         serverapp_th = run_serverapp_th(
@@ -417,6 +359,7 @@ def _main_loop(
             f_stop=f_stop,
             has_exception=server_app_thread_has_exception,
             enable_tf_gpu_growth=enable_tf_gpu_growth,
+            run_id=run.run_id,
         )
 
         # Buffer time so the `ServerApp` in separate thread is ready
@@ -566,20 +509,10 @@ def _parse_args_run_simulation() -> argparse.ArgumentParser:
     parser.add_argument(
         "--app",
         type=str,
-        default=None,
+        required=True,
         help="Path to a directory containing a FAB-like structure with a "
         "pyproject.toml.",
     )
-    parser.add_argument(
-        "--server-app",
-        help="(DEPRECATED: use --app instead) For example: `server:app` or "
-        "`project.package.module:wrapper.app`",
-    )
-    parser.add_argument(
-        "--client-app",
-        help="(DEPRECATED: use --app instead) For example: `client:app` or "
-        "`project.package.module:wrapper.app`",
-    )
     parser.add_argument(
         "--num-supernodes",
         type=int,
@@ -628,13 +561,6 @@ def _parse_args_run_simulation() -> argparse.ArgumentParser:
         help="When unset, only INFO, WARNING and ERROR log messages will be shown. "
         "If set, DEBUG-level logs will be displayed. ",
     )
-    parser.add_argument(
-        "--app-dir",
-        default="",
-        help="Add specified directory to the PYTHONPATH and load"
-        "ClientApp and ServerApp from there."
-        " Default: current working directory.",
-    )
     parser.add_argument(
         "--flwr-dir",
         default=None,

flwr/simulation/simulationio_connection.py

@@ -0,0 +1,86 @@
+# Copyright 2024 Flower Labs GmbH. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Flower SimulationIo connection."""
+
+
+from logging import DEBUG, WARNING
+from typing import Optional, cast
+
+import grpc
+
+from flwr.common.constant import SIMULATIONIO_API_DEFAULT_ADDRESS
+from flwr.common.grpc import create_channel
+from flwr.common.logger import log
+from flwr.proto.simulationio_pb2_grpc import SimulationIoStub  # pylint: disable=E0611
+
+
+class SimulationIoConnection:
+    """`SimulationIoConnection` provides an interface to the SimulationIo API.
+
+    Parameters
+    ----------
+    simulationio_service_address : str (default: "[::]:9094")
+        The address (URL, IPv6, IPv4) of the SuperLink SimulationIo API service.
+    root_certificates : Optional[bytes] (default: None)
+        The PEM-encoded root certificates as a byte string.
+        If provided, a secure connection using the certificates will be
+        established to an SSL-enabled Flower server.
+    """
+
+    def __init__(  # pylint: disable=too-many-arguments
+        self,
+        simulationio_service_address: str = SIMULATIONIO_API_DEFAULT_ADDRESS,
+        root_certificates: Optional[bytes] = None,
+    ) -> None:
+        self._addr = simulationio_service_address
+        self._cert = root_certificates
+        self._grpc_stub: Optional[SimulationIoStub] = None
+        self._channel: Optional[grpc.Channel] = None
+
+    @property
+    def _is_connected(self) -> bool:
+        """Check if connected to the SimulationIo API server."""
+        return self._channel is not None
+
+    @property
+    def _stub(self) -> SimulationIoStub:
+        """SimulationIo stub."""
+        if not self._is_connected:
+            self._connect()
+        return cast(SimulationIoStub, self._grpc_stub)
+
+    def _connect(self) -> None:
+        """Connect to the SimulationIo API."""
+        if self._is_connected:
+            log(WARNING, "Already connected")
+            return
+        self._channel = create_channel(
+            server_address=self._addr,
+            insecure=(self._cert is None),
+            root_certificates=self._cert,
+        )
+        self._grpc_stub = SimulationIoStub(self._channel)
+        log(DEBUG, "[SimulationIO] Connected to %s", self._addr)
+
+    def _disconnect(self) -> None:
+        """Disconnect from the SimulationIo API."""
+        if not self._is_connected:
+            log(DEBUG, "Already disconnected")
+            return
+        channel: grpc.Channel = self._channel
+        self._channel = None
+        self._grpc_stub = None
+        channel.close()
+        log(DEBUG, "[SimulationIO] Disconnected")