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

flwr/server/serverapp/app.py

@@ -0,0 +1,270 @@
+# 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 ServerApp process."""
+
+import argparse
+import sys
+from logging import DEBUG, ERROR, INFO, WARN
+from os.path import isfile
+from pathlib import Path
+from queue import Queue
+from time import sleep
+from typing import Optional
+
+from flwr.cli.config_utils import get_fab_metadata
+from flwr.cli.install import install_from_fab
+from flwr.common.config import (
+    get_flwr_dir,
+    get_fused_config_from_dir,
+    get_project_config,
+    get_project_dir,
+)
+from flwr.common.constant import Status, SubStatus
+from flwr.common.logger import (
+    log,
+    mirror_output_to_queue,
+    restore_output,
+    start_log_uploader,
+    stop_log_uploader,
+)
+from flwr.common.serde import (
+    context_from_proto,
+    context_to_proto,
+    fab_from_proto,
+    run_from_proto,
+    run_status_to_proto,
+)
+from flwr.common.typing import RunStatus
+from flwr.proto.run_pb2 import UpdateRunStatusRequest  # pylint: disable=E0611
+from flwr.proto.serverappio_pb2 import (  # pylint: disable=E0611
+    PullServerAppInputsRequest,
+    PullServerAppInputsResponse,
+    PushServerAppOutputsRequest,
+)
+from flwr.server.driver.grpc_driver import GrpcDriver
+from flwr.server.run_serverapp import run as run_
+
+
+def flwr_serverapp() -> None:
+    """Run process-isolated Flower ServerApp."""
+    # Capture stdout/stderr
+    log_queue: Queue[Optional[str]] = Queue()
+    mirror_output_to_queue(log_queue)
+
+    parser = argparse.ArgumentParser(
+        description="Run a Flower ServerApp",
+    )
+    parser.add_argument(
+        "--superlink",
+        type=str,
+        help="Address of SuperLink's DriverAPI",
+    )
+    parser.add_argument(
+        "--run-once",
+        action="store_true",
+        help="When set, this process will start a single ServerApp "
+        "for a pending Run. If no pending run the process will exit. ",
+    )
+    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(
+        "--insecure",
+        action="store_true",
+        help="Run the server without HTTPS, regardless of whether certificate "
+        "paths are provided. By default, the server runs with HTTPS enabled. "
+        "Use this flag only if you understand the risks.",
+    )
+    parser.add_argument(
+        "--root-certificates",
+        metavar="ROOT_CERT",
+        type=str,
+        help="Specifies the path to the PEM-encoded root certificate file for "
+        "establishing secure HTTPS connections.",
+    )
+    args = parser.parse_args()
+
+    log(INFO, "Starting Flower ServerApp")
+    certificates = _try_obtain_certificates(args)
+
+    log(
+        DEBUG,
+        "Staring isolated `ServerApp` connected to SuperLink DriverAPI at %s",
+        args.superlink,
+    )
+    run_serverapp(
+        superlink=args.superlink,
+        log_queue=log_queue,
+        run_once=args.run_once,
+        flwr_dir_=args.flwr_dir,
+        certificates=certificates,
+    )
+
+    # Restore stdout/stderr
+    restore_output()
+
+
+def _try_obtain_certificates(
+    args: argparse.Namespace,
+) -> Optional[bytes]:
+
+    if args.insecure:
+        if args.root_certificates is not None:
+            sys.exit(
+                "Conflicting options: The '--insecure' flag disables HTTPS, "
+                "but '--root-certificates' was also specified. Please remove "
+                "the '--root-certificates' option when running in insecure mode, "
+                "or omit '--insecure' to use HTTPS."
+            )
+        log(
+            WARN,
+            "Option `--insecure` was set. Starting insecure HTTP channel to %s.",
+            args.superlink,
+        )
+        root_certificates = None
+    else:
+        # Load the certificates if provided, or load the system certificates
+        if not isfile(args.root_certificates):
+            sys.exit("Path argument `--root-certificates` does not point to a file.")
+        root_certificates = Path(args.root_certificates).read_bytes()
+        log(
+            DEBUG,
+            "Starting secure HTTPS channel to %s "
+            "with the following certificates: %s.",
+            args.superlink,
+            args.root_certificates,
+        )
+    return root_certificates
+
+
+def run_serverapp(  # pylint: disable=R0914, disable=W0212
+    superlink: str,
+    log_queue: Queue[Optional[str]],
+    run_once: bool,
+    flwr_dir_: Optional[str] = None,
+    certificates: Optional[bytes] = None,
+) -> None:
+    """Run Flower ServerApp process."""
+    driver = GrpcDriver(
+        serverappio_service_address=superlink,
+        root_certificates=certificates,
+    )
+
+    # Resolve directory where FABs are installed
+    flwr_dir = get_flwr_dir(flwr_dir_)
+    log_uploader = None
+
+    while True:
+
+        try:
+            # Pull ServerAppInputs from LinkState
+            req = PullServerAppInputsRequest()
+            res: PullServerAppInputsResponse = driver._stub.PullServerAppInputs(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)
+
+            driver.init_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=driver._stub,
+            )
+
+            log(DEBUG, "ServerApp process starts FAB installation.")
+            install_from_fab(fab.content, flwr_dir=flwr_dir, skip_prompt=True)
+
+            fab_id, fab_version = get_fab_metadata(fab.content)
+
+            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
+            )
+
+            # Update run_config in context
+            context.run_config = server_app_run_config
+
+            log(
+                DEBUG,
+                "Flower 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, "", ""))
+            driver._stub.UpdateRunStatus(
+                UpdateRunStatusRequest(run_id=run.run_id, run_status=run_status_proto)
+            )
+
+            # Load and run the ServerApp with the Driver
+            updated_context = run_(
+                driver=driver,
+                server_app_dir=app_path,
+                server_app_attr=server_app_attr,
+                context=context,
+            )
+
+            # Send resulting context
+            context_proto = context_to_proto(updated_context)
+            out_req = PushServerAppOutputsRequest(
+                run_id=run.run_id, context=context_proto
+            )
+            _ = driver._stub.PushServerAppOutputs(out_req)
+
+            run_status = RunStatus(Status.FINISHED, SubStatus.COMPLETED, "")
+
+        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))
+
+        finally:
+            if run_status:
+                run_status_proto = run_status_to_proto(run_status)
+                driver._stub.UpdateRunStatus(
+                    UpdateRunStatusRequest(
+                        run_id=run.run_id, run_status=run_status_proto
+                    )
+                )
+
+            # Stop log uploader for this run
+            if log_uploader:
+                stop_log_uploader(log_queue, log_uploader)
+                log_uploader = None
+
+        # Stop the loop if `flwr-serverapp` is expected to process a single run
+        if run_once:
+            break

flwr/server/strategy/fedadam.py

@@ -170,8 +170,18 @@ class FedAdam(FedOpt):
             for x, y in zip(self.v_t, delta_t)
         ]
 
+        # Compute the bias-corrected learning rate, `eta_norm` for improving convergence
+        # in the early rounds of FL training. This `eta_norm` is `\alpha_t` in Kingma &
+        # Ba, 2014 (http://arxiv.org/abs/1412.6980) "Adam: A Method for Stochastic
+        # Optimization" in the formula line right before Section 2.1.
+        eta_norm = (
+            self.eta
+            * np.sqrt(1 - np.power(self.beta_2, server_round + 1.0))
+            / (1 - np.power(self.beta_1, server_round + 1.0))
+        )
+
         new_weights = [
-            x + self.eta * y / (np.sqrt(z) + self.tau)
+            x + eta_norm * y / (np.sqrt(z) + self.tau)
             for x, y, z in zip(self.current_weights, self.m_t, self.v_t)
         ]
 

flwr/server/superlink/driver/__init__.py

@@ -12,4 +12,4 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 # ==============================================================================
-"""Flower driver service."""
+"""Flower ServerAppIo service."""

flwr/server/superlink/driver/{driver_grpc.py → serverappio_grpc.py}

@@ -12,7 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 # ==============================================================================
-"""Driver gRPC API."""
+"""ServerAppIo gRPC API."""
 
 from logging import INFO
 from typing import Optional
@@ -21,37 +21,40 @@ import grpc
 
 from flwr.common import GRPC_MAX_MESSAGE_LENGTH
 from flwr.common.logger import log
-from flwr.proto.driver_pb2_grpc import (  # pylint: disable=E0611
-    add_DriverServicer_to_server,
+from flwr.proto.serverappio_pb2_grpc import (  # pylint: disable=E0611
+    add_ServerAppIoServicer_to_server,
 )
 from flwr.server.superlink.ffs.ffs_factory import FfsFactory
-from flwr.server.superlink.state import StateFactory
+from flwr.server.superlink.linkstate import LinkStateFactory
 
 from ..fleet.grpc_bidi.grpc_server import generic_create_grpc_server
-from .driver_servicer import DriverServicer
+from .serverappio_servicer import ServerAppIoServicer
 
 
-def run_driver_api_grpc(
+def run_serverappio_api_grpc(
     address: str,
-    state_factory: StateFactory,
+    state_factory: LinkStateFactory,
     ffs_factory: FfsFactory,
     certificates: Optional[tuple[bytes, bytes, bytes]],
 ) -> grpc.Server:
-    """Run Driver API (gRPC, request-response)."""
-    # Create Driver API gRPC server
-    driver_servicer: grpc.Server = DriverServicer(
+    """Run ServerAppIo API (gRPC, request-response)."""
+    # Create ServerAppIo API gRPC server
+    serverappio_servicer: grpc.Server = ServerAppIoServicer(
         state_factory=state_factory,
         ffs_factory=ffs_factory,
     )
-    driver_add_servicer_to_server_fn = add_DriverServicer_to_server
-    driver_grpc_server = generic_create_grpc_server(
-        servicer_and_add_fn=(driver_servicer, driver_add_servicer_to_server_fn),
+    serverappio_add_servicer_to_server_fn = add_ServerAppIoServicer_to_server
+    serverappio_grpc_server = generic_create_grpc_server(
+        servicer_and_add_fn=(
+            serverappio_servicer,
+            serverappio_add_servicer_to_server_fn,
+        ),
         server_address=address,
         max_message_length=GRPC_MAX_MESSAGE_LENGTH,
         certificates=certificates,
     )
 
-    log(INFO, "Flower ECE: Starting Driver API (gRPC-rere) on %s", address)
-    driver_grpc_server.start()
+    log(INFO, "Flower ECE: Starting ServerAppIo API (gRPC-rere) on %s", address)
+    serverappio_grpc_server.start()
 
-    return driver_grpc_server
+    return serverappio_grpc_server

flwr/server/superlink/driver/{driver_servicer.py → serverappio_servicer.py}

@@ -12,62 +12,80 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 # ==============================================================================
-"""Driver API servicer."""
+"""ServerAppIo API servicer."""
 
 
+import threading
 import time
-from logging import DEBUG
+from logging import DEBUG, INFO
 from typing import Optional
 from uuid import UUID
 
 import grpc
 
+from flwr.common import ConfigsRecord
+from flwr.common.constant import Status
 from flwr.common.logger import log
 from flwr.common.serde import (
+    context_from_proto,
+    context_to_proto,
     fab_from_proto,
     fab_to_proto,
+    run_status_from_proto,
+    run_to_proto,
     user_config_from_proto,
-    user_config_to_proto,
-)
-from flwr.common.typing import Fab
-from flwr.proto import driver_pb2_grpc  # pylint: disable=E0611
-from flwr.proto.driver_pb2 import (  # pylint: disable=E0611
-    GetNodesRequest,
-    GetNodesResponse,
-    PullTaskResRequest,
-    PullTaskResResponse,
-    PushTaskInsRequest,
-    PushTaskInsResponse,
 )
+from flwr.common.typing import Fab, RunStatus
+from flwr.proto import serverappio_pb2_grpc  # pylint: disable=E0611
 from flwr.proto.fab_pb2 import GetFabRequest, GetFabResponse  # pylint: disable=E0611
+from flwr.proto.log_pb2 import (  # pylint: disable=E0611
+    PushLogsRequest,
+    PushLogsResponse,
+)
 from flwr.proto.node_pb2 import Node  # pylint: disable=E0611
 from flwr.proto.run_pb2 import (  # pylint: disable=E0611
     CreateRunRequest,
     CreateRunResponse,
     GetRunRequest,
     GetRunResponse,
-    Run,
+    UpdateRunStatusRequest,
+    UpdateRunStatusResponse,
+)
+from flwr.proto.serverappio_pb2 import (  # pylint: disable=E0611
+    GetNodesRequest,
+    GetNodesResponse,
+    PullServerAppInputsRequest,
+    PullServerAppInputsResponse,
+    PullTaskResRequest,
+    PullTaskResResponse,
+    PushServerAppOutputsRequest,
+    PushServerAppOutputsResponse,
+    PushTaskInsRequest,
+    PushTaskInsResponse,
 )
 from flwr.proto.task_pb2 import TaskRes  # pylint: disable=E0611
 from flwr.server.superlink.ffs.ffs import Ffs
 from flwr.server.superlink.ffs.ffs_factory import FfsFactory
-from flwr.server.superlink.state import State, StateFactory
+from flwr.server.superlink.linkstate import LinkState, LinkStateFactory
 from flwr.server.utils.validator import validate_task_ins_or_res
 
 
-class DriverServicer(driver_pb2_grpc.DriverServicer):
-    """Driver API servicer."""
+class ServerAppIoServicer(serverappio_pb2_grpc.ServerAppIoServicer):
+    """ServerAppIo API servicer."""
 
-    def __init__(self, state_factory: StateFactory, ffs_factory: FfsFactory) -> None:
+    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 GetNodes(
         self, request: GetNodesRequest, context: grpc.ServicerContext
     ) -> GetNodesResponse:
         """Get available nodes."""
-        log(DEBUG, "DriverServicer.GetNodes")
-        state: State = self.state_factory.state()
+        log(DEBUG, "ServerAppIoServicer.GetNodes")
+        state: LinkState = self.state_factory.state()
         all_ids: set[int] = state.get_nodes(request.run_id)
         nodes: list[Node] = [
             Node(node_id=node_id, anonymous=False) for node_id in all_ids
@@ -78,8 +96,8 @@ class DriverServicer(driver_pb2_grpc.DriverServicer):
         self, request: CreateRunRequest, context: grpc.ServicerContext
     ) -> CreateRunResponse:
         """Create run ID."""
-        log(DEBUG, "DriverServicer.CreateRun")
-        state: State = self.state_factory.state()
+        log(DEBUG, "ServerAppIoServicer.CreateRun")
+        state: LinkState = self.state_factory.state()
         if request.HasField("fab"):
             fab = fab_from_proto(request.fab)
             ffs: Ffs = self.ffs_factory.ffs()
@@ -95,6 +113,7 @@ class DriverServicer(driver_pb2_grpc.DriverServicer):
             request.fab_version,
             fab_hash,
             user_config_from_proto(request.override_config),
+            ConfigsRecord(),
         )
         return CreateRunResponse(run_id=run_id)
 
@@ -102,7 +121,7 @@ class DriverServicer(driver_pb2_grpc.DriverServicer):
         self, request: PushTaskInsRequest, context: grpc.ServicerContext
     ) -> PushTaskInsResponse:
         """Push a set of TaskIns."""
-        log(DEBUG, "DriverServicer.PushTaskIns")
+        log(DEBUG, "ServerAppIoServicer.PushTaskIns")
 
         # Set pushed_at (timestamp in seconds)
         pushed_at = time.time()
@@ -116,7 +135,7 @@ class DriverServicer(driver_pb2_grpc.DriverServicer):
             _raise_if(bool(validation_errors), ", ".join(validation_errors))
 
         # Init state
-        state: State = self.state_factory.state()
+        state: LinkState = self.state_factory.state()
 
         # Store each TaskIns
         task_ids: list[Optional[UUID]] = []
@@ -132,17 +151,20 @@ class DriverServicer(driver_pb2_grpc.DriverServicer):
         self, request: PullTaskResRequest, context: grpc.ServicerContext
     ) -> PullTaskResResponse:
         """Pull a set of TaskRes."""
-        log(DEBUG, "DriverServicer.PullTaskRes")
+        log(DEBUG, "ServerAppIoServicer.PullTaskRes")
 
         # Convert each task_id str to UUID
         task_ids: set[UUID] = {UUID(task_id) for task_id in request.task_ids}
 
         # Init state
-        state: State = self.state_factory.state()
+        state: LinkState = self.state_factory.state()
 
         # Register callback
         def on_rpc_done() -> None:
-            log(DEBUG, "DriverServicer.PullTaskRes callback: delete TaskIns/TaskRes")
+            log(
+                DEBUG,
+                "ServerAppIoServicer.PullTaskRes callback: delete TaskIns/TaskRes",
+            )
 
             if context.is_active():
                 return
@@ -164,10 +186,10 @@ class DriverServicer(driver_pb2_grpc.DriverServicer):
         self, request: GetRunRequest, context: grpc.ServicerContext
     ) -> GetRunResponse:
         """Get run information."""
-        log(DEBUG, "DriverServicer.GetRun")
+        log(DEBUG, "ServerAppIoServicer.GetRun")
 
         # Init state
-        state: State = self.state_factory.state()
+        state: LinkState = self.state_factory.state()
 
         # Retrieve run information
         run = state.get_run(request.run_id)
@@ -175,21 +197,13 @@ class DriverServicer(driver_pb2_grpc.DriverServicer):
         if run is None:
             return GetRunResponse()
 
-        return GetRunResponse(
-            run=Run(
-                run_id=run.run_id,
-                fab_id=run.fab_id,
-                fab_version=run.fab_version,
-                override_config=user_config_to_proto(run.override_config),
-                fab_hash=run.fab_hash,
-            )
-        )
+        return GetRunResponse(run=run_to_proto(run))
 
     def GetFab(
         self, request: GetFabRequest, context: grpc.ServicerContext
     ) -> GetFabResponse:
         """Get FAB from Ffs."""
-        log(DEBUG, "DriverServicer.GetFab")
+        log(DEBUG, "ServerAppIoServicer.GetFab")
 
         ffs: Ffs = self.ffs_factory.ffs()
         if result := ffs.get(request.hash_str):
@@ -198,6 +212,78 @@ class DriverServicer(driver_pb2_grpc.DriverServicer):
 
         raise ValueError(f"Found no FAB with hash: {request.hash_str}")
 
+    def PullServerAppInputs(
+        self, request: PullServerAppInputsRequest, context: grpc.ServicerContext
+    ) -> PullServerAppInputsResponse:
+        """Pull ServerApp process inputs."""
+        log(DEBUG, "ServerAppIoServicer.PullServerAppInputs")
+        # 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 PullServerAppInputsResponse()
+
+            # 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 PullServerAppInputsResponse(
+                        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 PushServerAppOutputs(
+        self, request: PushServerAppOutputsRequest, context: grpc.ServicerContext
+    ) -> PushServerAppOutputsResponse:
+        """Push ServerApp process outputs."""
+        log(DEBUG, "ServerAppIoServicer.PushServerAppOutputs")
+        state = self.state_factory.state()
+        state.set_serverapp_context(request.run_id, context_from_proto(request.context))
+        return PushServerAppOutputsResponse()
+
+    def UpdateRunStatus(
+        self, request: UpdateRunStatusRequest, context: grpc.ServicerContext
+    ) -> UpdateRunStatusResponse:
+        """Update the status of a run."""
+        log(DEBUG, "ControlServicer.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()
+
 
 def _raise_if(validation_error: bool, detail: str) -> None:
     if validation_error:

flwr/server/superlink/fleet/grpc_adapter/grpc_adapter_servicer.py

@@ -48,7 +48,7 @@ from flwr.proto.grpcadapter_pb2 import MessageContainer  # pylint: disable=E0611
 from flwr.proto.run_pb2 import GetRunRequest, GetRunResponse  # pylint: disable=E0611
 from flwr.server.superlink.ffs.ffs_factory import FfsFactory
 from flwr.server.superlink.fleet.message_handler import message_handler
-from flwr.server.superlink.state import StateFactory
+from flwr.server.superlink.linkstate import LinkStateFactory
 
 T = TypeVar("T", bound=GrpcMessage)
 
@@ -77,7 +77,9 @@ def _handle(
 class GrpcAdapterServicer(grpcadapter_pb2_grpc.GrpcAdapterServicer):
     """Fleet API via GrpcAdapter servicer."""
 
-    def __init__(self, state_factory: StateFactory, ffs_factory: FfsFactory) -> None:
+    def __init__(
+        self, state_factory: LinkStateFactory, ffs_factory: FfsFactory
+    ) -> None:
         self.state_factory = state_factory
         self.ffs_factory = ffs_factory
 

flwr/server/superlink/fleet/grpc_bidi/grpc_server.py

@@ -30,7 +30,7 @@ from flwr.proto.transport_pb2_grpc import (  # pylint: disable=E0611
     add_FlowerServiceServicer_to_server,
 )
 from flwr.server.client_manager import ClientManager
-from flwr.server.superlink.driver.driver_servicer import DriverServicer
+from flwr.server.superlink.driver.serverappio_servicer import ServerAppIoServicer
 from flwr.server.superlink.fleet.grpc_adapter.grpc_adapter_servicer import (
     GrpcAdapterServicer,
 )
@@ -161,7 +161,7 @@ def generic_create_grpc_server(  # pylint: disable=too-many-arguments,R0917
         tuple[FleetServicer, AddServicerToServerFn],
         tuple[GrpcAdapterServicer, AddServicerToServerFn],
         tuple[FlowerServiceServicer, AddServicerToServerFn],
-        tuple[DriverServicer, AddServicerToServerFn],
+        tuple[ServerAppIoServicer, AddServicerToServerFn],
     ],
     server_address: str,
     max_concurrent_workers: int = 1000,