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

flwr/client/app.py

@@ -37,6 +37,8 @@ from flwr.common import GRPC_MAX_MESSAGE_LENGTH, Context, EventType, Message, ev
 from flwr.common.address import parse_address
 from flwr.common.constant import (
     CLIENTAPPIO_API_DEFAULT_ADDRESS,
+    ISOLATION_MODE_PROCESS,
+    ISOLATION_MODE_SUBPROCESS,
     MISSING_EXTRA_REST,
     RUN_ID_NUM_BYTES,
     TRANSPORT_TYPE_GRPC_ADAPTER,
@@ -52,18 +54,15 @@ from flwr.common.retry_invoker import RetryInvoker, RetryState, exponential
 from flwr.common.typing import Fab, Run, UserConfig
 from flwr.proto.clientappio_pb2_grpc import add_ClientAppIoServicer_to_server
 from flwr.server.superlink.fleet.grpc_bidi.grpc_server import generic_create_grpc_server
-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 .clientapp.clientappio_servicer import ClientAppInputs, ClientAppIoServicer
 from .grpc_adapter_client.connection import grpc_adapter
 from .grpc_client.connection import grpc_connection
 from .grpc_rere_client.connection import grpc_request_response
 from .message_handler.message_handler import handle_control_message
-from .node_state import NodeState
 from .numpy_client import NumPyClient
-
-ISOLATION_MODE_SUBPROCESS = "subprocess"
-ISOLATION_MODE_PROCESS = "process"
+from .run_info_store import DeprecatedRunInfoStore
 
 
 def _check_actionable_client(
@@ -364,8 +363,8 @@ def start_client_internal(
         on_backoff=_on_backoff,
     )
 
-    # NodeState gets initialized when the first connection is established
-    node_state: Optional[NodeState] = None
+    # DeprecatedRunInfoStore gets initialized when the first connection is established
+    run_info_store: Optional[DeprecatedRunInfoStore] = None
 
     runs: dict[int, Run] = {}
 
@@ -382,7 +381,7 @@ def start_client_internal(
             receive, send, create_node, delete_node, get_run, get_fab = conn
 
             # Register node when connecting the first time
-            if node_state is None:
+            if run_info_store is None:
                 if create_node is None:
                     if transport not in ["grpc-bidi", None]:
                         raise NotImplementedError(
@@ -391,7 +390,7 @@ def start_client_internal(
                         )
                     # gRPC-bidi doesn't have the concept of node_id,
                     # so we set it to -1
-                    node_state = NodeState(
+                    run_info_store = DeprecatedRunInfoStore(
                         node_id=-1,
                         node_config={},
                     )
@@ -402,7 +401,7 @@ def start_client_internal(
                     )  # pylint: disable=not-callable
                     if node_id is None:
                         raise ValueError("Node registration failed")
-                    node_state = NodeState(
+                    run_info_store = DeprecatedRunInfoStore(
                         node_id=node_id,
                         node_config=node_config,
                     )
@@ -461,7 +460,7 @@ def start_client_internal(
                     run.fab_id, run.fab_version = fab_id, fab_version
 
                     # Register context for this run
-                    node_state.register_context(
+                    run_info_store.register_context(
                         run_id=run_id,
                         run=run,
                         flwr_path=flwr_path,
@@ -469,7 +468,7 @@ def start_client_internal(
                     )
 
                     # Retrieve context for this run
-                    context = node_state.retrieve_context(run_id=run_id)
+                    context = run_info_store.retrieve_context(run_id=run_id)
                     # Create an error reply message that will never be used to prevent
                     # the used-before-assignment linting error
                     reply_message = message.create_error_reply(
@@ -542,7 +541,7 @@ def start_client_internal(
                             # Raise exception, crash process
                             raise ex
 
-                        # Don't update/change NodeState
+                        # Don't update/change DeprecatedRunInfoStore
 
                         e_code = ErrorCode.CLIENT_APP_RAISED_EXCEPTION
                         # Ex fmt: "<class 'ZeroDivisionError'>:<'division by zero'>"
@@ -567,7 +566,7 @@ def start_client_internal(
                         )
                     else:
                         # No exception, update node state
-                        node_state.update_context(
+                        run_info_store.update_context(
                             run_id=run_id,
                             context=context,
                         )

flwr/client/node_state_tests.py

@@ -17,7 +17,7 @@
 
 from typing import cast
 
-from flwr.client.node_state import NodeState
+from flwr.client.run_info_store import DeprecatedRunInfoStore
 from flwr.common import ConfigsRecord, Context
 from flwr.proto.task_pb2 import TaskIns  # pylint: disable=E0611
 
@@ -34,32 +34,31 @@ def _run_dummy_task(context: Context) -> Context:
 
 
 def test_multirun_in_node_state() -> None:
-    """Test basic NodeState logic."""
+    """Test basic DeprecatedRunInfoStore logic."""
     # Tasks to perform
     tasks = [TaskIns(run_id=run_id) for run_id in [0, 1, 1, 2, 3, 2, 1, 5]]
     # the "tasks" is to count how many times each run is executed
     expected_values = {0: "1", 1: "1" * 3, 2: "1" * 2, 3: "1", 5: "1"}
 
-    # NodeState
-    node_state = NodeState(node_id=0, node_config={})
+    node_info_store = DeprecatedRunInfoStore(node_id=0, node_config={})
 
     for task in tasks:
         run_id = task.run_id
 
         # Register
-        node_state.register_context(run_id=run_id)
+        node_info_store.register_context(run_id=run_id)
 
         # Get run state
-        context = node_state.retrieve_context(run_id=run_id)
+        context = node_info_store.retrieve_context(run_id=run_id)
 
         # Run "task"
         updated_state = _run_dummy_task(context)
 
         # Update run state
-        node_state.update_context(run_id=run_id, context=updated_state)
+        node_info_store.update_context(run_id=run_id, context=updated_state)
 
     # Verify values
-    for run_id, run_info in node_state.run_infos.items():
+    for run_id, run_info in node_info_store.run_infos.items():
         assert (
             run_info.context.state.configs_records["counter"]["count"]
             == expected_values[run_id]

flwr/client/{node_state.py → run_info_store.py}

@@ -1,4 +1,4 @@
-# Copyright 2023 Flower Labs GmbH. All Rights Reserved.
+# 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.
@@ -12,7 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 # ==============================================================================
-"""Node state."""
+"""Deprecated Run Info Store."""
 
 
 from dataclasses import dataclass
@@ -36,7 +36,7 @@ class RunInfo:
     initial_run_config: UserConfig
 
 
-class NodeState:
+class DeprecatedRunInfoStore:
     """State of a node where client nodes execute runs."""
 
     def __init__(

flwr/client/supernode/app.py

@@ -31,6 +31,8 @@ from flwr.common import EventType, event
 from flwr.common.config import parse_config_args
 from flwr.common.constant import (
     FLEET_API_GRPC_RERE_DEFAULT_ADDRESS,
+    ISOLATION_MODE_PROCESS,
+    ISOLATION_MODE_SUBPROCESS,
     TRANSPORT_TYPE_GRPC_ADAPTER,
     TRANSPORT_TYPE_GRPC_RERE,
     TRANSPORT_TYPE_REST,
@@ -38,11 +40,7 @@ from flwr.common.constant import (
 from flwr.common.exit_handlers import register_exit_handlers
 from flwr.common.logger import log, warn_deprecated_feature
 
-from ..app import (
-    ISOLATION_MODE_PROCESS,
-    ISOLATION_MODE_SUBPROCESS,
-    start_client_internal,
-)
+from ..app import start_client_internal
 from ..clientapp.utils import get_load_client_app_fn
 
 
@@ -200,10 +198,10 @@ def _parse_args_run_supernode() -> argparse.ArgumentParser:
             ISOLATION_MODE_SUBPROCESS,
             ISOLATION_MODE_PROCESS,
         ],
-        help="Isolation mode when running `ClientApp` (optional, possible values: "
-        "`subprocess`, `process`). By default, `ClientApp` runs in the same process "
+        help="Isolation mode when running a `ClientApp` (optional, possible values: "
+        "`subprocess`, `process`). By default, a `ClientApp` runs in the same process "
         "that executes the SuperNode. Use `subprocess` to configure SuperNode to run "
-        "`ClientApp` in a subprocess. Use `process` to indicate that a separate "
+        "a `ClientApp` in a subprocess. Use `process` to indicate that a separate "
         "independent process gets created outside of SuperNode.",
     )
     parser.add_argument(

flwr/common/constant.py

@@ -40,15 +40,14 @@ TRANSPORT_TYPES = [
 # Addresses
 # SuperNode
 CLIENTAPPIO_API_DEFAULT_ADDRESS = "0.0.0.0:9094"
-# SuperExec
-EXEC_API_DEFAULT_ADDRESS = "0.0.0.0:9093"
 # SuperLink
 DRIVER_API_DEFAULT_ADDRESS = "0.0.0.0:9091"
 FLEET_API_GRPC_RERE_DEFAULT_ADDRESS = "0.0.0.0:9092"
 FLEET_API_GRPC_BIDI_DEFAULT_ADDRESS = (
     "[::]:8080"  # IPv6 to keep start_server compatible
 )
-FLEET_API_REST_DEFAULT_ADDRESS = "0.0.0.0:9093"
+FLEET_API_REST_DEFAULT_ADDRESS = "0.0.0.0:9095"
+EXEC_API_DEFAULT_ADDRESS = "0.0.0.0:9093"
 
 # Constants for ping
 PING_DEFAULT_INTERVAL = 30
@@ -84,6 +83,10 @@ GRPC_ADAPTER_METADATA_MESSAGE_QUALNAME_KEY = "grpc-message-qualname"
 # Message TTL
 MESSAGE_TTL_TOLERANCE = 1e-1
 
+# Isolation modes
+ISOLATION_MODE_SUBPROCESS = "subprocess"
+ISOLATION_MODE_PROCESS = "process"
+
 
 class MessageType:
     """Message type."""
@@ -129,3 +132,28 @@ class ErrorCode:
     def __new__(cls) -> ErrorCode:
         """Prevent instantiation."""
         raise TypeError(f"{cls.__name__} cannot be instantiated.")
+
+
+class Status:
+    """Run status."""
+
+    PENDING = "pending"
+    STARTING = "starting"
+    RUNNING = "running"
+    FINISHED = "finished"
+
+    def __new__(cls) -> Status:
+        """Prevent instantiation."""
+        raise TypeError(f"{cls.__name__} cannot be instantiated.")
+
+
+class SubStatus:
+    """Run sub-status."""
+
+    COMPLETED = "completed"
+    FAILED = "failed"
+    STOPPED = "stopped"
+
+    def __new__(cls) -> SubStatus:
+        """Prevent instantiation."""
+        raise TypeError(f"{cls.__name__} cannot be instantiated.")

flwr/common/typing.py

@@ -218,6 +218,15 @@ class Run:
     override_config: UserConfig
 
 
+@dataclass
+class RunStatus:
+    """Run status information."""
+
+    status: str
+    sub_status: str
+    details: str
+
+
 @dataclass
 class Fab:
     """Fab file representation."""

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,16 +37,20 @@ 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,
     TRANSPORT_TYPE_GRPC_ADAPTER,
     TRANSPORT_TYPE_GRPC_RERE,
     TRANSPORT_TYPE_REST,
+    Status,
 )
 from flwr.common.exit_handlers import register_exit_handlers
 from flwr.common.logger import log
@@ -52,10 +58,13 @@ from flwr.common.secure_aggregation.crypto.symmetric_encryption import (
     private_key_to_bytes,
     public_key_to_bytes,
 )
+from flwr.common.typing import RunStatus
 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_superexec_api_grpc
 
 from .client_manager import ClientManager
 from .history import History
@@ -71,7 +80,7 @@ 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
 
 DATABASE = ":flwr-in-memory-state:"
 BASE_DIR = get_flwr_dir() / "superlink" / "ffs"
@@ -205,14 +214,15 @@ def run_superlink() -> None:
 
     event(EventType.RUN_SUPERLINK_ENTER)
 
-    # Parse IP address
+    # Parse IP addresses
     driver_address, _, _ = _format_address(args.driver_api_address)
+    exec_address, _, _ = _format_address(args.exec_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)
@@ -224,8 +234,9 @@ def run_superlink() -> None:
         ffs_factory=ffs_factory,
         certificates=certificates,
     )
-
     grpc_servers = [driver_server]
+
+    # Start Fleet API
     bckg_threads = []
     if not args.fleet_api_address:
         if args.fleet_api_type in [
@@ -250,7 +261,6 @@ def run_superlink() -> None:
         )
         num_workers = 1
 
-    # Start Fleet API
     if args.fleet_api_type == TRANSPORT_TYPE_REST:
         if (
             importlib.util.find_spec("requests")
@@ -318,6 +328,26 @@ def run_superlink() -> None:
     else:
         raise ValueError(f"Unknown fleet_api_type: {args.fleet_api_type}")
 
+    # Start Exec API
+    exec_server: grpc.Server = run_superexec_api_grpc(
+        address=exec_address,
+        executor=load_executor(args),
+        certificates=certificates,
+        config=parse_config_args(
+            [args.executor_config] if args.executor_config else args.executor_config
+        ),
+    )
+    grpc_servers.append(exec_server)
+
+    if args.isolation == ISOLATION_MODE_SUBPROCESS:
+        # Scheduler thread
+        scheduler_th = threading.Thread(
+            target=_flwr_serverapp_scheduler,
+            args=(state_factory, args.driver_api_address),
+        )
+        scheduler_th.start()
+        bckg_threads.append(scheduler_th)
+
     # Graceful shutdown
     register_exit_handlers(
         event_type=EventType.RUN_SUPERLINK_LEAVE,
@@ -334,6 +364,47 @@ def run_superlink() -> None:
         driver_server.wait_for_termination(timeout=1)
 
 
+def _flwr_serverapp_scheduler(
+    state_factory: LinkStateFactory, driver_api_address: 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:
+
+            # Set run as starting
+            state.update_run_status(
+                run_id=pending_run_id, new_status=RunStatus(Status.STARTING, "", "")
+            )
+            log(
+                INFO,
+                "Launching `flwr-serverapp` subprocess with run-id %d. "
+                "Connects to SuperLink on %s",
+                pending_run_id,
+                driver_api_address,
+            )
+            # Start ServerApp subprocess
+            command = [
+                "flwr-serverapp",
+                "--superlink",
+                driver_api_address,
+                "--run-id",
+                str(pending_run_id),
+            ]
+            subprocess.run(
+                command,
+                stdout=None,
+                stderr=None,
+                check=True,
+            )
+
+
 def _format_address(address: str) -> tuple[str, str, int]:
     parsed_address = parse_address(address)
     if not parsed_address:
@@ -489,7 +560,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 +588,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,7 +619,7 @@ 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:
@@ -587,6 +658,7 @@ def _parse_args_run_superlink() -> argparse.ArgumentParser:
     _add_args_common(parser=parser)
     _add_args_driver_api(parser=parser)
     _add_args_fleet_api(parser=parser)
+    _add_args_exec_api(parser=parser)
 
     return parser
 
@@ -618,6 +690,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 "
@@ -681,3 +766,29 @@ 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"\'`',
+    )

flwr/server/driver/inmemory_driver.py

@@ -25,7 +25,7 @@ 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
 
@@ -46,7 +46,7 @@ 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

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",
 ]

flwr/server/serverapp/app.py

@@ -0,0 +1,78 @@
+# 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
+from logging import DEBUG, INFO
+from typing import Optional
+
+from flwr.common.logger import log
+from flwr.server.driver.grpc_driver import GrpcDriver
+
+
+def flwr_serverapp() -> None:
+    """Run process-isolated Flower ServerApp."""
+    log(INFO, "Starting Flower ServerApp")
+
+    parser = argparse.ArgumentParser(
+        description="Run a Flower ServerApp",
+    )
+    parser.add_argument(
+        "--superlink",
+        type=str,
+        help="Address of SuperLink's DriverAPI",
+    )
+    parser.add_argument(
+        "--run-id",
+        type=int,
+        required=False,
+        help="Id of the Run this process should start. If not supplied, this "
+        "function will request a pending run to the LinkState.",
+    )
+    args = parser.parse_args()
+
+    log(
+        DEBUG,
+        "Staring isolated `ServerApp` connected to SuperLink DriverAPI at %s "
+        "for run-id %s",
+        args.superlink,
+        args.run_id,
+    )
+    run_serverapp(superlink=args.superlink, run_id=args.run_id)
+
+
+def run_serverapp(  # pylint: disable=R0914
+    superlink: str,
+    run_id: Optional[int] = None,
+) -> None:
+    """Run Flower ServerApp process.
+
+    Parameters
+    ----------
+    superlink : str
+        Address of SuperLink
+    run_id : Optional[int] (default: None)
+        Unique identifier of a Run registered at the LinkState. If not supplied,
+        this function will request a pending run to the LinkState.
+    """
+    _ = GrpcDriver(
+        run_id=run_id if run_id else 0,
+        driver_service_address=superlink,
+        root_certificates=None,
+    )
+
+    # Then, GetServerInputs
+
+    # Then, run ServerApp

flwr/server/superlink/driver/driver_grpc.py

@@ -25,7 +25,7 @@ from flwr.proto.driver_pb2_grpc import (  # pylint: disable=E0611
     add_DriverServicer_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
@@ -33,7 +33,7 @@ from .driver_servicer import DriverServicer
 
 def run_driver_api_grpc(
     address: str,
-    state_factory: StateFactory,
+    state_factory: LinkStateFactory,
     ffs_factory: FfsFactory,
     certificates: Optional[tuple[bytes, bytes, bytes]],
 ) -> grpc.Server: