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

flwr/client/app.py

@@ -52,15 +52,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
+from .run_info_store import DeprecatedRunInfoStore
 
 ISOLATION_MODE_SUBPROCESS = "subprocess"
 ISOLATION_MODE_PROCESS = "process"
@@ -364,8 +364,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 +382,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 +391,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 +402,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 +461,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 +469,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 +542,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 +567,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/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

flwr/server/app.py

@@ -35,9 +35,10 @@ 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,
@@ -56,6 +57,8 @@ 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 +74,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 +208,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 +228,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 +255,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 +322,17 @@ 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)
+
     # Graceful shutdown
     register_exit_handlers(
         event_type=EventType.RUN_SUPERLINK_LEAVE,
@@ -489,7 +504,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 +532,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 +563,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 +602,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
 
@@ -681,3 +697,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,20 @@
+# 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."""
+
+
+def flwr_serverapp() -> None:
+    """Run process-isolated Flower ServerApp."""
+    raise NotImplementedError()

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:

flwr/server/superlink/driver/driver_servicer.py

@@ -51,14 +51,16 @@ from flwr.proto.run_pb2 import (  # pylint: disable=E0611
 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."""
 
-    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
 
@@ -67,7 +69,7 @@ class DriverServicer(driver_pb2_grpc.DriverServicer):
     ) -> GetNodesResponse:
         """Get available nodes."""
         log(DEBUG, "DriverServicer.GetNodes")
-        state: State = self.state_factory.state()
+        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
@@ -79,7 +81,7 @@ class DriverServicer(driver_pb2_grpc.DriverServicer):
     ) -> CreateRunResponse:
         """Create run ID."""
         log(DEBUG, "DriverServicer.CreateRun")
-        state: State = self.state_factory.state()
+        state: LinkState = self.state_factory.state()
         if request.HasField("fab"):
             fab = fab_from_proto(request.fab)
             ffs: Ffs = self.ffs_factory.ffs()
@@ -116,7 +118,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]] = []
@@ -138,7 +140,7 @@ class DriverServicer(driver_pb2_grpc.DriverServicer):
         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:
@@ -167,7 +169,7 @@ class DriverServicer(driver_pb2_grpc.DriverServicer):
         log(DEBUG, "DriverServicer.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)

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_rere/fleet_servicer.py

@@ -37,13 +37,15 @@ from flwr.proto.fleet_pb2 import (  # 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
 
 
 class FleetServicer(fleet_pb2_grpc.FleetServicer):
     """Fleet 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
 

flwr/server/superlink/fleet/grpc_rere/server_interceptor.py

@@ -45,7 +45,7 @@ from flwr.proto.fleet_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.server.superlink.state import State
+from flwr.server.superlink.linkstate import LinkState
 
 _PUBLIC_KEY_HEADER = "public-key"
 _AUTH_TOKEN_HEADER = "auth-token"
@@ -84,7 +84,7 @@ def _get_value_from_tuples(
 class AuthenticateServerInterceptor(grpc.ServerInterceptor):  # type: ignore
     """Server interceptor for node authentication."""
 
-    def __init__(self, state: State):
+    def __init__(self, state: LinkState):
         self.state = state
 
         self.node_public_keys = state.get_node_public_keys()

flwr/server/superlink/fleet/message_handler/message_handler.py

@@ -43,12 +43,12 @@ from flwr.proto.run_pb2 import (  # pylint: disable=E0611
 )
 from flwr.proto.task_pb2 import TaskIns, TaskRes  # pylint: disable=E0611
 from flwr.server.superlink.ffs.ffs import Ffs
-from flwr.server.superlink.state import State
+from flwr.server.superlink.linkstate import LinkState
 
 
 def create_node(
     request: CreateNodeRequest,  # pylint: disable=unused-argument
-    state: State,
+    state: LinkState,
 ) -> CreateNodeResponse:
     """."""
     # Create node
@@ -56,7 +56,7 @@ def create_node(
     return CreateNodeResponse(node=Node(node_id=node_id, anonymous=False))
 
 
-def delete_node(request: DeleteNodeRequest, state: State) -> DeleteNodeResponse:
+def delete_node(request: DeleteNodeRequest, state: LinkState) -> DeleteNodeResponse:
     """."""
     # Validate node_id
     if request.node.anonymous or request.node.node_id == 0:
@@ -69,14 +69,14 @@ def delete_node(request: DeleteNodeRequest, state: State) -> DeleteNodeResponse:
 
 def ping(
     request: PingRequest,  # pylint: disable=unused-argument
-    state: State,  # pylint: disable=unused-argument
+    state: LinkState,  # pylint: disable=unused-argument
 ) -> PingResponse:
     """."""
     res = state.acknowledge_ping(request.node.node_id, request.ping_interval)
     return PingResponse(success=res)
 
 
-def pull_task_ins(request: PullTaskInsRequest, state: State) -> PullTaskInsResponse:
+def pull_task_ins(request: PullTaskInsRequest, state: LinkState) -> PullTaskInsResponse:
     """Pull TaskIns handler."""
     # Get node_id if client node is not anonymous
     node = request.node  # pylint: disable=no-member
@@ -92,7 +92,7 @@ def pull_task_ins(request: PullTaskInsRequest, state: State) -> PullTaskInsRespo
     return response
 
 
-def push_task_res(request: PushTaskResRequest, state: State) -> PushTaskResResponse:
+def push_task_res(request: PushTaskResRequest, state: LinkState) -> PushTaskResResponse:
     """Push TaskRes handler."""
     # pylint: disable=no-member
     task_res: TaskRes = request.task_res_list[0]
@@ -113,7 +113,7 @@ def push_task_res(request: PushTaskResRequest, state: State) -> PushTaskResRespo
 
 
 def get_run(
-    request: GetRunRequest, state: State  # pylint: disable=W0613
+    request: GetRunRequest, state: LinkState  # pylint: disable=W0613
 ) -> GetRunResponse:
     """Get run information."""
     run = state.get_run(request.run_id)

flwr/server/superlink/fleet/rest_rere/rest_api.py

@@ -40,7 +40,7 @@ from flwr.proto.fleet_pb2 import (  # pylint: disable=E0611
 from flwr.proto.run_pb2 import GetRunRequest, GetRunResponse  # pylint: disable=E0611
 from flwr.server.superlink.ffs.ffs import Ffs
 from flwr.server.superlink.fleet.message_handler import message_handler
-from flwr.server.superlink.state import State
+from flwr.server.superlink.linkstate import LinkState
 
 try:
     from starlette.applications import Starlette
@@ -90,7 +90,7 @@ def rest_request_response(
 async def create_node(request: CreateNodeRequest) -> CreateNodeResponse:
     """Create Node."""
     # Get state from app
-    state: State = app.state.STATE_FACTORY.state()
+    state: LinkState = app.state.STATE_FACTORY.state()
 
     # Handle message
     return message_handler.create_node(request=request, state=state)
@@ -100,7 +100,7 @@ async def create_node(request: CreateNodeRequest) -> CreateNodeResponse:
 async def delete_node(request: DeleteNodeRequest) -> DeleteNodeResponse:
     """Delete Node Id."""
     # Get state from app
-    state: State = app.state.STATE_FACTORY.state()
+    state: LinkState = app.state.STATE_FACTORY.state()
 
     # Handle message
     return message_handler.delete_node(request=request, state=state)
@@ -110,7 +110,7 @@ async def delete_node(request: DeleteNodeRequest) -> DeleteNodeResponse:
 async def pull_task_ins(request: PullTaskInsRequest) -> PullTaskInsResponse:
     """Pull TaskIns."""
     # Get state from app
-    state: State = app.state.STATE_FACTORY.state()
+    state: LinkState = app.state.STATE_FACTORY.state()
 
     # Handle message
     return message_handler.pull_task_ins(request=request, state=state)
@@ -121,7 +121,7 @@ async def pull_task_ins(request: PullTaskInsRequest) -> PullTaskInsResponse:
 async def push_task_res(request: PushTaskResRequest) -> PushTaskResResponse:
     """Push TaskRes."""
     # Get state from app
-    state: State = app.state.STATE_FACTORY.state()
+    state: LinkState = app.state.STATE_FACTORY.state()
 
     # Handle message
     return message_handler.push_task_res(request=request, state=state)
@@ -131,7 +131,7 @@ async def push_task_res(request: PushTaskResRequest) -> PushTaskResResponse:
 async def ping(request: PingRequest) -> PingResponse:
     """Ping."""
     # Get state from app
-    state: State = app.state.STATE_FACTORY.state()
+    state: LinkState = app.state.STATE_FACTORY.state()
 
     # Handle message
     return message_handler.ping(request=request, state=state)
@@ -141,7 +141,7 @@ async def ping(request: PingRequest) -> PingResponse:
 async def get_run(request: GetRunRequest) -> GetRunResponse:
     """GetRun."""
     # Get state from app
-    state: State = app.state.STATE_FACTORY.state()
+    state: LinkState = app.state.STATE_FACTORY.state()
 
     # Handle message
     return message_handler.get_run(request=request, state=state)