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

flwr/server/driver/grpc_driver.py

@@ -23,7 +23,7 @@ from typing import Optional, cast
 import grpc
 
 from flwr.common import DEFAULT_TTL, Message, Metadata, RecordSet
-from flwr.common.constant import DRIVER_API_DEFAULT_ADDRESS
+from flwr.common.constant import SERVERAPPIO_API_DEFAULT_ADDRESS
 from flwr.common.grpc import create_channel
 from flwr.common.logger import log
 from flwr.common.serde import (
@@ -32,7 +32,9 @@ from flwr.common.serde import (
     user_config_from_proto,
 )
 from flwr.common.typing import Run
-from flwr.proto.driver_pb2 import (  # pylint: disable=E0611
+from flwr.proto.node_pb2 import Node  # pylint: disable=E0611
+from flwr.proto.run_pb2 import GetRunRequest, GetRunResponse  # pylint: disable=E0611
+from flwr.proto.serverappio_pb2 import (  # pylint: disable=E0611
     GetNodesRequest,
     GetNodesResponse,
     PullTaskResRequest,
@@ -40,9 +42,7 @@ from flwr.proto.driver_pb2 import (  # pylint: disable=E0611
     PushTaskInsRequest,
     PushTaskInsResponse,
 )
-from flwr.proto.driver_pb2_grpc import DriverStub  # pylint: disable=E0611
-from flwr.proto.node_pb2 import Node  # pylint: disable=E0611
-from flwr.proto.run_pb2 import GetRunRequest, GetRunResponse  # pylint: disable=E0611
+from flwr.proto.serverappio_pb2_grpc import ServerAppIoStub  # pylint: disable=E0611
 from flwr.proto.task_pb2 import TaskIns  # pylint: disable=E0611
 
 from .driver import Driver
@@ -56,14 +56,12 @@ Call `connect()` on the `GrpcDriverStub` instance before calling any of the othe
 
 
 class GrpcDriver(Driver):
-    """`GrpcDriver` provides an interface to the Driver API.
+    """`GrpcDriver` provides an interface to the ServerAppIo API.
 
     Parameters
     ----------
-    run_id : int
-        The identifier of the run.
-    driver_service_address : str (default: "[::]:9091")
-        The address (URL, IPv6, IPv4) of the SuperLink Driver API service.
+    serverappio_service_address : str (default: "[::]:9091")
+        The address (URL, IPv6, IPv4) of the SuperLink ServerAppIo 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
@@ -72,25 +70,23 @@ class GrpcDriver(Driver):
 
     def __init__(  # pylint: disable=too-many-arguments
         self,
-        run_id: int,
-        driver_service_address: str = DRIVER_API_DEFAULT_ADDRESS,
+        serverappio_service_address: str = SERVERAPPIO_API_DEFAULT_ADDRESS,
         root_certificates: Optional[bytes] = None,
     ) -> None:
-        self._run_id = run_id
-        self._addr = driver_service_address
+        self._addr = serverappio_service_address
         self._cert = root_certificates
         self._run: Optional[Run] = None
-        self._grpc_stub: Optional[DriverStub] = None
+        self._grpc_stub: Optional[ServerAppIoStub] = None
         self._channel: Optional[grpc.Channel] = None
         self.node = Node(node_id=0, anonymous=True)
 
     @property
     def _is_connected(self) -> bool:
-        """Check if connected to the Driver API server."""
+        """Check if connected to the ServerAppIo API server."""
         return self._channel is not None
 
     def _connect(self) -> None:
-        """Connect to the Driver API.
+        """Connect to the ServerAppIo API.
 
         This will not call GetRun.
         """
@@ -102,11 +98,11 @@ class GrpcDriver(Driver):
             insecure=(self._cert is None),
             root_certificates=self._cert,
         )
-        self._grpc_stub = DriverStub(self._channel)
+        self._grpc_stub = ServerAppIoStub(self._channel)
         log(DEBUG, "[Driver] Connected to %s", self._addr)
 
     def _disconnect(self) -> None:
-        """Disconnect from the Driver API."""
+        """Disconnect from the ServerAppIo API."""
         if not self._is_connected:
             log(DEBUG, "Already disconnected")
             return
@@ -116,15 +112,17 @@ class GrpcDriver(Driver):
         channel.close()
         log(DEBUG, "[Driver] Disconnected")
 
-    def _init_run(self) -> None:
+    def init_run(self, run_id: int) -> None:
+        """Initialize the run."""
         # Check if is initialized
         if self._run is not None:
             return
+
         # Get the run info
-        req = GetRunRequest(run_id=self._run_id)
+        req = GetRunRequest(run_id=run_id)
         res: GetRunResponse = self._stub.GetRun(req)
         if not res.HasField("run"):
-            raise RuntimeError(f"Cannot find the run with ID: {self._run_id}")
+            raise RuntimeError(f"Cannot find the run with ID: {run_id}")
         self._run = Run(
             run_id=res.run.run_id,
             fab_id=res.run.fab_id,
@@ -136,21 +134,20 @@ class GrpcDriver(Driver):
     @property
     def run(self) -> Run:
         """Run information."""
-        self._init_run()
         return Run(**vars(self._run))
 
     @property
-    def _stub(self) -> DriverStub:
-        """Driver stub."""
+    def _stub(self) -> ServerAppIoStub:
+        """ServerAppIo stub."""
         if not self._is_connected:
             self._connect()
-        return cast(DriverStub, self._grpc_stub)
+        return cast(ServerAppIoStub, self._grpc_stub)
 
     def _check_message(self, message: Message) -> None:
         # Check if the message is valid
         if not (
             # Assume self._run being initialized
-            message.metadata.run_id == self._run_id
+            message.metadata.run_id == cast(Run, self._run).run_id
             and message.metadata.src_node_id == self.node.node_id
             and message.metadata.message_id == ""
             and message.metadata.reply_to_message == ""
@@ -171,7 +168,6 @@ class GrpcDriver(Driver):
         This method constructs a new `Message` with given content and metadata.
         The `run_id` and `src_node_id` will be set automatically.
         """
-        self._init_run()
         if ttl:
             warnings.warn(
                 "A custom TTL was set, but note that the SuperLink does not enforce "
@@ -182,7 +178,7 @@ class GrpcDriver(Driver):
 
         ttl_ = DEFAULT_TTL if ttl is None else ttl
         metadata = Metadata(
-            run_id=self._run_id,
+            run_id=cast(Run, self._run).run_id,
             message_id="",  # Will be set by the server
             src_node_id=self.node.node_id,
             dst_node_id=dst_node_id,
@@ -195,10 +191,9 @@ class GrpcDriver(Driver):
 
     def get_node_ids(self) -> list[int]:
         """Get node IDs."""
-        self._init_run()
         # Call GrpcDriverStub method
         res: GetNodesResponse = self._stub.GetNodes(
-            GetNodesRequest(run_id=self._run_id)
+            GetNodesRequest(run_id=cast(Run, self._run).run_id)
         )
         return [node.node_id for node in res.nodes]
 
@@ -208,7 +203,6 @@ class GrpcDriver(Driver):
         This method takes an iterable of messages and sends each message
         to the node specified in `dst_node_id`.
         """
-        self._init_run()
         # Construct TaskIns
         task_ins_list: list[TaskIns] = []
         for msg in messages:
@@ -230,7 +224,6 @@ class GrpcDriver(Driver):
         This method is used to collect messages from the SuperLink that correspond to a
         set of given message IDs.
         """
-        self._init_run()
         # Pull TaskRes
         res: PullTaskResResponse = self._stub.PullTaskRes(
             PullTaskResRequest(node=self.node, task_ids=message_ids)

flwr/server/driver/inmemory_driver.py

@@ -25,18 +25,16 @@ from flwr.common import DEFAULT_TTL, Message, Metadata, RecordSet
 from flwr.common.serde import message_from_taskres, message_to_taskins
 from flwr.common.typing import Run
 from flwr.proto.node_pb2 import Node  # pylint: disable=E0611
-from flwr.server.superlink.state import StateFactory
+from flwr.server.superlink.linkstate import LinkStateFactory
 
 from .driver import Driver
 
 
 class InMemoryDriver(Driver):
-    """`InMemoryDriver` class provides an interface to the Driver API.
+    """`InMemoryDriver` class provides an interface to the ServerAppIo API.
 
     Parameters
     ----------
-    run_id : int
-        The identifier of the run.
     state_factory : StateFactory
         A StateFactory embedding a state that this driver can interface with.
     pull_interval : float (default=0.1)
@@ -45,18 +43,15 @@ class InMemoryDriver(Driver):
 
     def __init__(
         self,
-        run_id: int,
-        state_factory: StateFactory,
+        state_factory: LinkStateFactory,
         pull_interval: float = 0.1,
     ) -> None:
-        self._run_id = run_id
         self._run: Optional[Run] = None
         self.state = state_factory.state()
         self.pull_interval = pull_interval
         self.node = Node(node_id=0, anonymous=True)
 
     def _check_message(self, message: Message) -> None:
-        self._init_run()
         # Check if the message is valid
         if not (
             message.metadata.run_id == cast(Run, self._run).run_id
@@ -67,19 +62,18 @@ class InMemoryDriver(Driver):
         ):
             raise ValueError(f"Invalid message: {message}")
 
-    def _init_run(self) -> None:
+    def init_run(self, run_id: int) -> None:
         """Initialize the run."""
         if self._run is not None:
             return
-        run = self.state.get_run(self._run_id)
+        run = self.state.get_run(run_id)
         if run is None:
-            raise RuntimeError(f"Cannot find the run with ID: {self._run_id}")
+            raise RuntimeError(f"Cannot find the run with ID: {run_id}")
         self._run = run
 
     @property
     def run(self) -> Run:
         """Run ID."""
-        self._init_run()
         return Run(**vars(cast(Run, self._run)))
 
     def create_message(  # pylint: disable=too-many-arguments,R0917
@@ -95,7 +89,6 @@ class InMemoryDriver(Driver):
         This method constructs a new `Message` with given content and metadata.
         The `run_id` and `src_node_id` will be set automatically.
         """
-        self._init_run()
         if ttl:
             warnings.warn(
                 "A custom TTL was set, but note that the SuperLink does not enforce "
@@ -119,7 +112,6 @@ class InMemoryDriver(Driver):
 
     def get_node_ids(self) -> list[int]:
         """Get node IDs."""
-        self._init_run()
         return list(self.state.get_nodes(cast(Run, self._run).run_id))
 
     def push_messages(self, messages: Iterable[Message]) -> Iterable[str]:

flwr/server/run_serverapp.py

@@ -31,10 +31,9 @@ from flwr.common.config import (
     get_project_config,
     get_project_dir,
 )
-from flwr.common.constant import DRIVER_API_DEFAULT_ADDRESS
+from flwr.common.constant import SERVERAPPIO_API_DEFAULT_ADDRESS
 from flwr.common.logger import log, update_console_handler, warn_deprecated_feature
 from flwr.common.object_ref import load_app
-from flwr.common.typing import UserConfig
 from flwr.proto.fab_pb2 import GetFabRequest, GetFabResponse  # pylint: disable=E0611
 from flwr.proto.run_pb2 import (  # pylint: disable=E0611
     CreateRunRequest,
@@ -48,11 +47,11 @@ from .server_app import LoadServerAppError, ServerApp
 
 def run(
     driver: Driver,
+    context: Context,
     server_app_dir: str,
-    server_app_run_config: UserConfig,
     server_app_attr: Optional[str] = None,
     loaded_server_app: Optional[ServerApp] = None,
-) -> None:
+) -> Context:
     """Run ServerApp with a given Driver."""
     if not (server_app_attr is None) ^ (loaded_server_app is None):
         raise ValueError(
@@ -78,15 +77,11 @@ def run(
 
     server_app = _load()
 
-    # Initialize Context
-    context = Context(
-        node_id=0, node_config={}, state=RecordSet(), run_config=server_app_run_config
-    )
-
     # Call ServerApp
     server_app(driver=driver, context=context)
 
     log(DEBUG, "ServerApp finished running.")
+    return context
 
 
 # pylint: disable-next=too-many-branches,too-many-statements,too-many-locals
@@ -111,18 +106,18 @@ def run_server_app() -> None:
             "app by executing `flwr new` and following the prompt."
         )
 
-    if args.server != DRIVER_API_DEFAULT_ADDRESS:
+    if args.server != SERVERAPPIO_API_DEFAULT_ADDRESS:
         warn = "Passing flag --server is deprecated. Use --superlink instead."
         warn_deprecated_feature(warn)
 
-        if args.superlink != DRIVER_API_DEFAULT_ADDRESS:
+        if args.superlink != SERVERAPPIO_API_DEFAULT_ADDRESS:
             # if `--superlink` also passed, then
             # warn user that this argument overrides what was passed with `--server`
             log(
                 WARN,
                 "Both `--server` and `--superlink` were passed. "
-                "`--server` will be ignored. Connecting to the Superlink Driver API "
-                "at %s.",
+                "`--server` will be ignored. Connecting to the "
+                "SuperLink ServerAppIo API at %s.",
                 args.superlink,
             )
         else:
@@ -175,11 +170,11 @@ def run_server_app() -> None:
     if app_path is None:
         # User provided `--run-id`, but not `app_dir`
         driver = GrpcDriver(
-            run_id=args.run_id,
-            driver_service_address=args.superlink,
+            serverappio_service_address=args.superlink,
             root_certificates=root_certificates,
         )
         flwr_dir = get_flwr_dir(args.flwr_dir)
+        driver.init_run(args.run_id)
         run_ = driver.run
         if not run_.fab_hash:
             raise ValueError("FAB hash not provided.")
@@ -193,12 +188,12 @@ def run_server_app() -> None:
 
         app_path = str(get_project_dir(fab_id, fab_version, run_.fab_hash, flwr_dir))
         config = get_project_config(app_path)
+        run_id = run_.run_id
     else:
         # User provided `app_dir`, but not `--run-id`
         # Create run if run_id is not provided
         driver = GrpcDriver(
-            run_id=0,  # Will be overwritten
-            driver_service_address=args.superlink,
+            serverappio_service_address=args.superlink,
             root_certificates=root_certificates,
         )
         # Load config from the project directory
@@ -208,8 +203,9 @@ def run_server_app() -> None:
         # Create run
         req = CreateRunRequest(fab_id=fab_id, fab_version=fab_version)
         res: CreateRunResponse = driver._stub.CreateRun(req)  # pylint: disable=W0212
-        # Overwrite driver._run_id
-        driver._run_id = res.run_id  # pylint: disable=W0212
+        # Fetch full `Run` using `run_id`
+        driver.init_run(res.run_id)  # pylint: disable=W0212
+        run_id = res.run_id
 
     # Obtain server app reference and the run config
     server_app_attr = config["tool"]["flwr"]["app"]["components"]["serverapp"]
@@ -225,11 +221,20 @@ def run_server_app() -> None:
         root_certificates,
     )
 
+    # Initialize Context
+    context = Context(
+        run_id=run_id,
+        node_id=0,
+        node_config={},
+        state=RecordSet(),
+        run_config=server_app_run_config,
+    )
+
     # Run the ServerApp with the Driver
     run(
         driver=driver,
+        context=context,
         server_app_dir=app_path,
-        server_app_run_config=server_app_run_config,
         server_app_attr=server_app_attr,
     )
 
@@ -272,13 +277,14 @@ def _parse_args_run_server_app() -> argparse.ArgumentParser:
     )
     parser.add_argument(
         "--server",
-        default=DRIVER_API_DEFAULT_ADDRESS,
+        default=SERVERAPPIO_API_DEFAULT_ADDRESS,
         help="Server address",
     )
     parser.add_argument(
         "--superlink",
-        default=DRIVER_API_DEFAULT_ADDRESS,
-        help="SuperLink Driver API (gRPC-rere) address (IPv4, IPv6, or a domain name)",
+        default=SERVERAPPIO_API_DEFAULT_ADDRESS,
+        help="SuperLink ServerAppIo API (gRPC-rere) address "
+        "(IPv4, IPv6, or a domain name)",
     )
     parser.add_argument(
         "--run-id",

flwr/server/{superlink/state → serverapp}/__init__.py

@@ -12,17 +12,11 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 # ==============================================================================
-"""Flower server state."""
+"""Flower AppIO service."""
 
 
-from .in_memory_state import InMemoryState as InMemoryState
-from .sqlite_state import SqliteState as SqliteState
-from .state import State as State
-from .state_factory import StateFactory as StateFactory
+from .app import flwr_serverapp as flwr_serverapp
 
 __all__ = [
-    "InMemoryState",
-    "SqliteState",
-    "State",
-    "StateFactory",
+    "flwr_serverapp",
 ]