flwr-nightly 1.10.0.dev20240707__py3-none-any.whl → 1.11.0.dev20240724__py3-none-any.whl

flwr/proto/exec_pb2.pyi

@@ -3,7 +3,9 @@
 isort:skip_file
 """
 import builtins
+import flwr.proto.transport_pb2
 import google.protobuf.descriptor
+import google.protobuf.internal.containers
 import google.protobuf.message
 import typing
 import typing_extensions
@@ -12,13 +14,51 @@ DESCRIPTOR: google.protobuf.descriptor.FileDescriptor
 
 class StartRunRequest(google.protobuf.message.Message):
     DESCRIPTOR: google.protobuf.descriptor.Descriptor
+    class OverrideConfigEntry(google.protobuf.message.Message):
+        DESCRIPTOR: google.protobuf.descriptor.Descriptor
+        KEY_FIELD_NUMBER: builtins.int
+        VALUE_FIELD_NUMBER: builtins.int
+        key: typing.Text
+        @property
+        def value(self) -> flwr.proto.transport_pb2.Scalar: ...
+        def __init__(self,
+            *,
+            key: typing.Text = ...,
+            value: typing.Optional[flwr.proto.transport_pb2.Scalar] = ...,
+            ) -> None: ...
+        def HasField(self, field_name: typing_extensions.Literal["value",b"value"]) -> builtins.bool: ...
+        def ClearField(self, field_name: typing_extensions.Literal["key",b"key","value",b"value"]) -> None: ...
+
+    class FederationConfigEntry(google.protobuf.message.Message):
+        DESCRIPTOR: google.protobuf.descriptor.Descriptor
+        KEY_FIELD_NUMBER: builtins.int
+        VALUE_FIELD_NUMBER: builtins.int
+        key: typing.Text
+        @property
+        def value(self) -> flwr.proto.transport_pb2.Scalar: ...
+        def __init__(self,
+            *,
+            key: typing.Text = ...,
+            value: typing.Optional[flwr.proto.transport_pb2.Scalar] = ...,
+            ) -> None: ...
+        def HasField(self, field_name: typing_extensions.Literal["value",b"value"]) -> builtins.bool: ...
+        def ClearField(self, field_name: typing_extensions.Literal["key",b"key","value",b"value"]) -> None: ...
+
     FAB_FILE_FIELD_NUMBER: builtins.int
+    OVERRIDE_CONFIG_FIELD_NUMBER: builtins.int
+    FEDERATION_CONFIG_FIELD_NUMBER: builtins.int
     fab_file: builtins.bytes
+    @property
+    def override_config(self) -> google.protobuf.internal.containers.MessageMap[typing.Text, flwr.proto.transport_pb2.Scalar]: ...
+    @property
+    def federation_config(self) -> google.protobuf.internal.containers.MessageMap[typing.Text, flwr.proto.transport_pb2.Scalar]: ...
     def __init__(self,
         *,
         fab_file: builtins.bytes = ...,
+        override_config: typing.Optional[typing.Mapping[typing.Text, flwr.proto.transport_pb2.Scalar]] = ...,
+        federation_config: typing.Optional[typing.Mapping[typing.Text, flwr.proto.transport_pb2.Scalar]] = ...,
         ) -> None: ...
-    def ClearField(self, field_name: typing_extensions.Literal["fab_file",b"fab_file"]) -> None: ...
+    def ClearField(self, field_name: typing_extensions.Literal["fab_file",b"fab_file","federation_config",b"federation_config","override_config",b"override_config"]) -> None: ...
 global___StartRunRequest = StartRunRequest
 
 class StartRunResponse(google.protobuf.message.Message):

flwr/proto/run_pb2.py

@@ -12,19 +12,24 @@ from google.protobuf.internal import builder as _builder
 _sym_db = _symbol_database.Default()
 
 
+from flwr.proto import transport_pb2 as flwr_dot_proto_dot_transport__pb2
 
 
-DESCRIPTOR = _descriptor_pool.Default().AddSerializedFile(b'\n\x14\x66lwr/proto/run.proto\x12\nflwr.proto\":\n\x03Run\x12\x0e\n\x06run_id\x18\x01 \x01(\x12\x12\x0e\n\x06\x66\x61\x62_id\x18\x02 \x01(\t\x12\x13\n\x0b\x66\x61\x62_version\x18\x03 \x01(\t\"\x1f\n\rGetRunRequest\x12\x0e\n\x06run_id\x18\x01 \x01(\x12\".\n\x0eGetRunResponse\x12\x1c\n\x03run\x18\x01 \x01(\x0b\x32\x0f.flwr.proto.Runb\x06proto3')
+DESCRIPTOR = _descriptor_pool.Default().AddSerializedFile(b'\n\x14\x66lwr/proto/run.proto\x12\nflwr.proto\x1a\x1a\x66lwr/proto/transport.proto\"\xc3\x01\n\x03Run\x12\x0e\n\x06run_id\x18\x01 \x01(\x12\x12\x0e\n\x06\x66\x61\x62_id\x18\x02 \x01(\t\x12\x13\n\x0b\x66\x61\x62_version\x18\x03 \x01(\t\x12<\n\x0foverride_config\x18\x04 \x03(\x0b\x32#.flwr.proto.Run.OverrideConfigEntry\x1aI\n\x13OverrideConfigEntry\x12\x0b\n\x03key\x18\x01 \x01(\t\x12!\n\x05value\x18\x02 \x01(\x0b\x32\x12.flwr.proto.Scalar:\x02\x38\x01\"\x1f\n\rGetRunRequest\x12\x0e\n\x06run_id\x18\x01 \x01(\x12\".\n\x0eGetRunResponse\x12\x1c\n\x03run\x18\x01 \x01(\x0b\x32\x0f.flwr.proto.Runb\x06proto3')
 
 _globals = globals()
 _builder.BuildMessageAndEnumDescriptors(DESCRIPTOR, _globals)
 _builder.BuildTopDescriptorsAndMessages(DESCRIPTOR, 'flwr.proto.run_pb2', _globals)
 if _descriptor._USE_C_DESCRIPTORS == False:
   DESCRIPTOR._options = None
-  _globals['_RUN']._serialized_start=36
-  _globals['_RUN']._serialized_end=94
-  _globals['_GETRUNREQUEST']._serialized_start=96
-  _globals['_GETRUNREQUEST']._serialized_end=127
-  _globals['_GETRUNRESPONSE']._serialized_start=129
-  _globals['_GETRUNRESPONSE']._serialized_end=175
+  _globals['_RUN_OVERRIDECONFIGENTRY']._options = None
+  _globals['_RUN_OVERRIDECONFIGENTRY']._serialized_options = b'8\001'
+  _globals['_RUN']._serialized_start=65
+  _globals['_RUN']._serialized_end=260
+  _globals['_RUN_OVERRIDECONFIGENTRY']._serialized_start=187
+  _globals['_RUN_OVERRIDECONFIGENTRY']._serialized_end=260
+  _globals['_GETRUNREQUEST']._serialized_start=262
+  _globals['_GETRUNREQUEST']._serialized_end=293
+  _globals['_GETRUNRESPONSE']._serialized_start=295
+  _globals['_GETRUNRESPONSE']._serialized_end=341
 # @@protoc_insertion_point(module_scope)

flwr/proto/run_pb2.pyi

@@ -3,7 +3,9 @@
 isort:skip_file
 """
 import builtins
+import flwr.proto.transport_pb2
 import google.protobuf.descriptor
+import google.protobuf.internal.containers
 import google.protobuf.message
 import typing
 import typing_extensions
@@ -12,19 +14,38 @@ DESCRIPTOR: google.protobuf.descriptor.FileDescriptor
 
 class Run(google.protobuf.message.Message):
     DESCRIPTOR: google.protobuf.descriptor.Descriptor
+    class OverrideConfigEntry(google.protobuf.message.Message):
+        DESCRIPTOR: google.protobuf.descriptor.Descriptor
+        KEY_FIELD_NUMBER: builtins.int
+        VALUE_FIELD_NUMBER: builtins.int
+        key: typing.Text
+        @property
+        def value(self) -> flwr.proto.transport_pb2.Scalar: ...
+        def __init__(self,
+            *,
+            key: typing.Text = ...,
+            value: typing.Optional[flwr.proto.transport_pb2.Scalar] = ...,
+            ) -> None: ...
+        def HasField(self, field_name: typing_extensions.Literal["value",b"value"]) -> builtins.bool: ...
+        def ClearField(self, field_name: typing_extensions.Literal["key",b"key","value",b"value"]) -> None: ...
+
     RUN_ID_FIELD_NUMBER: builtins.int
     FAB_ID_FIELD_NUMBER: builtins.int
     FAB_VERSION_FIELD_NUMBER: builtins.int
+    OVERRIDE_CONFIG_FIELD_NUMBER: builtins.int
     run_id: builtins.int
     fab_id: typing.Text
     fab_version: typing.Text
+    @property
+    def override_config(self) -> google.protobuf.internal.containers.MessageMap[typing.Text, flwr.proto.transport_pb2.Scalar]: ...
     def __init__(self,
         *,
         run_id: builtins.int = ...,
         fab_id: typing.Text = ...,
         fab_version: typing.Text = ...,
+        override_config: typing.Optional[typing.Mapping[typing.Text, flwr.proto.transport_pb2.Scalar]] = ...,
         ) -> None: ...
-    def ClearField(self, field_name: typing_extensions.Literal["fab_id",b"fab_id","fab_version",b"fab_version","run_id",b"run_id"]) -> None: ...
+    def ClearField(self, field_name: typing_extensions.Literal["fab_id",b"fab_id","fab_version",b"fab_version","override_config",b"override_config","run_id",b"run_id"]) -> None: ...
 global___Run = Run
 
 class GetRunRequest(google.protobuf.message.Message):

flwr/proto/task_pb2.py

@@ -14,21 +14,20 @@ _sym_db = _symbol_database.Default()
 
 from flwr.proto import node_pb2 as flwr_dot_proto_dot_node__pb2
 from flwr.proto import recordset_pb2 as flwr_dot_proto_dot_recordset__pb2
-from flwr.proto import transport_pb2 as flwr_dot_proto_dot_transport__pb2
 from flwr.proto import error_pb2 as flwr_dot_proto_dot_error__pb2
 
 
-DESCRIPTOR = _descriptor_pool.Default().AddSerializedFile(b'\n\x15\x66lwr/proto/task.proto\x12\nflwr.proto\x1a\x15\x66lwr/proto/node.proto\x1a\x1a\x66lwr/proto/recordset.proto\x1a\x1a\x66lwr/proto/transport.proto\x1a\x16\x66lwr/proto/error.proto\"\x89\x02\n\x04Task\x12\"\n\x08producer\x18\x01 \x01(\x0b\x32\x10.flwr.proto.Node\x12\"\n\x08\x63onsumer\x18\x02 \x01(\x0b\x32\x10.flwr.proto.Node\x12\x12\n\ncreated_at\x18\x03 \x01(\x01\x12\x14\n\x0c\x64\x65livered_at\x18\x04 \x01(\t\x12\x11\n\tpushed_at\x18\x05 \x01(\x01\x12\x0b\n\x03ttl\x18\x06 \x01(\x01\x12\x10\n\x08\x61ncestry\x18\x07 \x03(\t\x12\x11\n\ttask_type\x18\x08 \x01(\t\x12(\n\trecordset\x18\t \x01(\x0b\x32\x15.flwr.proto.RecordSet\x12 \n\x05\x65rror\x18\n \x01(\x0b\x32\x11.flwr.proto.Error\"\\\n\x07TaskIns\x12\x0f\n\x07task_id\x18\x01 \x01(\t\x12\x10\n\x08group_id\x18\x02 \x01(\t\x12\x0e\n\x06run_id\x18\x03 \x01(\x12\x12\x1e\n\x04task\x18\x04 \x01(\x0b\x32\x10.flwr.proto.Task\"\\\n\x07TaskRes\x12\x0f\n\x07task_id\x18\x01 \x01(\t\x12\x10\n\x08group_id\x18\x02 \x01(\t\x12\x0e\n\x06run_id\x18\x03 \x01(\x12\x12\x1e\n\x04task\x18\x04 \x01(\x0b\x32\x10.flwr.proto.Taskb\x06proto3')
+DESCRIPTOR = _descriptor_pool.Default().AddSerializedFile(b'\n\x15\x66lwr/proto/task.proto\x12\nflwr.proto\x1a\x15\x66lwr/proto/node.proto\x1a\x1a\x66lwr/proto/recordset.proto\x1a\x16\x66lwr/proto/error.proto\"\x89\x02\n\x04Task\x12\"\n\x08producer\x18\x01 \x01(\x0b\x32\x10.flwr.proto.Node\x12\"\n\x08\x63onsumer\x18\x02 \x01(\x0b\x32\x10.flwr.proto.Node\x12\x12\n\ncreated_at\x18\x03 \x01(\x01\x12\x14\n\x0c\x64\x65livered_at\x18\x04 \x01(\t\x12\x11\n\tpushed_at\x18\x05 \x01(\x01\x12\x0b\n\x03ttl\x18\x06 \x01(\x01\x12\x10\n\x08\x61ncestry\x18\x07 \x03(\t\x12\x11\n\ttask_type\x18\x08 \x01(\t\x12(\n\trecordset\x18\t \x01(\x0b\x32\x15.flwr.proto.RecordSet\x12 \n\x05\x65rror\x18\n \x01(\x0b\x32\x11.flwr.proto.Error\"\\\n\x07TaskIns\x12\x0f\n\x07task_id\x18\x01 \x01(\t\x12\x10\n\x08group_id\x18\x02 \x01(\t\x12\x0e\n\x06run_id\x18\x03 \x01(\x12\x12\x1e\n\x04task\x18\x04 \x01(\x0b\x32\x10.flwr.proto.Task\"\\\n\x07TaskRes\x12\x0f\n\x07task_id\x18\x01 \x01(\t\x12\x10\n\x08group_id\x18\x02 \x01(\t\x12\x0e\n\x06run_id\x18\x03 \x01(\x12\x12\x1e\n\x04task\x18\x04 \x01(\x0b\x32\x10.flwr.proto.Taskb\x06proto3')
 
 _globals = globals()
 _builder.BuildMessageAndEnumDescriptors(DESCRIPTOR, _globals)
 _builder.BuildTopDescriptorsAndMessages(DESCRIPTOR, 'flwr.proto.task_pb2', _globals)
 if _descriptor._USE_C_DESCRIPTORS == False:
   DESCRIPTOR._options = None
-  _globals['_TASK']._serialized_start=141
-  _globals['_TASK']._serialized_end=406
-  _globals['_TASKINS']._serialized_start=408
-  _globals['_TASKINS']._serialized_end=500
-  _globals['_TASKRES']._serialized_start=502
-  _globals['_TASKRES']._serialized_end=594
+  _globals['_TASK']._serialized_start=113
+  _globals['_TASK']._serialized_end=378
+  _globals['_TASKINS']._serialized_start=380
+  _globals['_TASKINS']._serialized_end=472
+  _globals['_TASKRES']._serialized_start=474
+  _globals['_TASKRES']._serialized_end=566
 # @@protoc_insertion_point(module_scope)

flwr/server/__init__.py

@@ -28,6 +28,7 @@ from .run_serverapp import run_server_app as run_server_app
 from .server import Server as Server
 from .server_app import ServerApp as ServerApp
 from .server_config import ServerConfig as ServerConfig
+from .serverapp_components import ServerAppComponents as ServerAppComponents
 
 __all__ = [
     "ClientManager",
@@ -36,6 +37,7 @@ __all__ = [
     "LegacyContext",
     "Server",
     "ServerApp",
+    "ServerAppComponents",
     "ServerConfig",
     "SimpleClientManager",
     "run_server_app",

flwr/server/compat/legacy_context.py

@@ -18,7 +18,7 @@
 from dataclasses import dataclass
 from typing import Optional
 
-from flwr.common import Context, RecordSet
+from flwr.common import Context
 
 from ..client_manager import ClientManager, SimpleClientManager
 from ..history import History
@@ -35,9 +35,9 @@ class LegacyContext(Context):
     client_manager: ClientManager
     history: History
 
-    def __init__(
+    def __init__(  # pylint: disable=too-many-arguments
         self,
-        state: RecordSet,
+        context: Context,
         config: Optional[ServerConfig] = None,
         strategy: Optional[Strategy] = None,
         client_manager: Optional[ClientManager] = None,
@@ -52,4 +52,5 @@ class LegacyContext(Context):
         self.strategy = strategy
         self.client_manager = client_manager
         self.history = History()
-        super().__init__(state)
+
+        super().__init__(**vars(context))

flwr/server/driver/grpc_driver.py

@@ -16,19 +16,21 @@
 
 import time
 import warnings
-from logging import DEBUG, ERROR, WARNING
-from typing import Iterable, List, Optional, Tuple, cast
+from logging import DEBUG, WARNING
+from typing import Iterable, List, Optional, cast
 
 import grpc
 
 from flwr.common import DEFAULT_TTL, EventType, Message, Metadata, RecordSet, event
 from flwr.common.grpc import create_channel
 from flwr.common.logger import log
-from flwr.common.serde import message_from_taskres, message_to_taskins
+from flwr.common.serde import (
+    message_from_taskres,
+    message_to_taskins,
+    user_config_from_proto,
+)
 from flwr.common.typing import Run
 from flwr.proto.driver_pb2 import (  # pylint: disable=E0611
-    CreateRunRequest,
-    CreateRunResponse,
     GetNodesRequest,
     GetNodesResponse,
     PullTaskResRequest,
@@ -53,167 +55,103 @@ Call `connect()` on the `GrpcDriverStub` instance before calling any of the othe
 """
 
 
-class GrpcDriverStub:
-    """`GrpcDriverStub` provides access to the gRPC Driver API/service.
+class GrpcDriver(Driver):
+    """`GrpcDriver` provides an interface to the Driver API.
 
     Parameters
     ----------
-    driver_service_address : Optional[str]
-        The IPv4 or IPv6 address of the Driver API server.
-        Defaults to `"[::]:9091"`.
+    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.
     root_certificates : Optional[bytes] (default: None)
         The PEM-encoded root certificates as a byte string.
         If provided, a secure connection using the certificates will be
         established to an SSL-enabled Flower server.
     """
 
-    def __init__(
+    def __init__(  # pylint: disable=too-many-arguments
         self,
+        run_id: int,
         driver_service_address: str = DEFAULT_SERVER_ADDRESS_DRIVER,
         root_certificates: Optional[bytes] = None,
     ) -> None:
-        self.driver_service_address = driver_service_address
-        self.root_certificates = root_certificates
-        self.channel: Optional[grpc.Channel] = None
-        self.stub: Optional[DriverStub] = None
+        self._run_id = run_id
+        self._addr = driver_service_address
+        self._cert = root_certificates
+        self._run: Optional[Run] = None
+        self._grpc_stub: Optional[DriverStub] = None
+        self._channel: Optional[grpc.Channel] = None
+        self.node = Node(node_id=0, anonymous=True)
 
-    def is_connected(self) -> bool:
-        """Return True if connected to the Driver API server, otherwise False."""
-        return self.channel is not None
+    @property
+    def _is_connected(self) -> bool:
+        """Check if connected to the Driver API server."""
+        return self._channel is not None
+
+    def _connect(self) -> None:
+        """Connect to the Driver API.
 
-    def connect(self) -> None:
-        """Connect to the Driver API."""
+        This will not call GetRun.
+        """
         event(EventType.DRIVER_CONNECT)
-        if self.channel is not None or self.stub is not None:
+        if self._is_connected:
             log(WARNING, "Already connected")
             return
-        self.channel = create_channel(
-            server_address=self.driver_service_address,
-            insecure=(self.root_certificates is None),
-            root_certificates=self.root_certificates,
+        self._channel = create_channel(
+            server_address=self._addr,
+            insecure=(self._cert is None),
+            root_certificates=self._cert,
         )
-        self.stub = DriverStub(self.channel)
-        log(DEBUG, "[Driver] Connected to %s", self.driver_service_address)
+        self._grpc_stub = DriverStub(self._channel)
+        log(DEBUG, "[Driver] Connected to %s", self._addr)
 
-    def disconnect(self) -> None:
+    def _disconnect(self) -> None:
         """Disconnect from the Driver API."""
         event(EventType.DRIVER_DISCONNECT)
-        if self.channel is None or self.stub is None:
+        if not self._is_connected:
             log(DEBUG, "Already disconnected")
             return
-        channel = self.channel
-        self.channel = None
-        self.stub = None
+        channel: grpc.Channel = self._channel
+        self._channel = None
+        self._grpc_stub = None
         channel.close()
         log(DEBUG, "[Driver] Disconnected")
 
-    def create_run(self, req: CreateRunRequest) -> CreateRunResponse:
-        """Request for run ID."""
-        # Check if channel is open
-        if self.stub is None:
-            log(ERROR, ERROR_MESSAGE_DRIVER_NOT_CONNECTED)
-            raise ConnectionError("`GrpcDriverStub` instance not connected")
-
-        # Call Driver API
-        res: CreateRunResponse = self.stub.CreateRun(request=req)
-        return res
-
-    def get_run(self, req: GetRunRequest) -> GetRunResponse:
-        """Get run information."""
-        # Check if channel is open
-        if self.stub is None:
-            log(ERROR, ERROR_MESSAGE_DRIVER_NOT_CONNECTED)
-            raise ConnectionError("`GrpcDriverStub` instance not connected")
-
-        # Call gRPC Driver API
-        res: GetRunResponse = self.stub.GetRun(request=req)
-        return res
-
-    def get_nodes(self, req: GetNodesRequest) -> GetNodesResponse:
-        """Get client IDs."""
-        # Check if channel is open
-        if self.stub is None:
-            log(ERROR, ERROR_MESSAGE_DRIVER_NOT_CONNECTED)
-            raise ConnectionError("`GrpcDriverStub` instance not connected")
-
-        # Call gRPC Driver API
-        res: GetNodesResponse = self.stub.GetNodes(request=req)
-        return res
-
-    def push_task_ins(self, req: PushTaskInsRequest) -> PushTaskInsResponse:
-        """Schedule tasks."""
-        # Check if channel is open
-        if self.stub is None:
-            log(ERROR, ERROR_MESSAGE_DRIVER_NOT_CONNECTED)
-            raise ConnectionError("`GrpcDriverStub` instance not connected")
-
-        # Call gRPC Driver API
-        res: PushTaskInsResponse = self.stub.PushTaskIns(request=req)
-        return res
-
-    def pull_task_res(self, req: PullTaskResRequest) -> PullTaskResResponse:
-        """Get task results."""
-        # Check if channel is open
-        if self.stub is None:
-            log(ERROR, ERROR_MESSAGE_DRIVER_NOT_CONNECTED)
-            raise ConnectionError("`GrpcDriverStub` instance not connected")
-
-        # Call Driver API
-        res: PullTaskResResponse = self.stub.PullTaskRes(request=req)
-        return res
-
-
-class GrpcDriver(Driver):
-    """`Driver` class provides an interface to the Driver API.
-
-    Parameters
-    ----------
-    run_id : int
-        The identifier of the run.
-    stub : Optional[GrpcDriverStub] (default: None)
-        The ``GrpcDriverStub`` instance used to communicate with the SuperLink.
-        If None, an instance connected to "[::]:9091" will be created.
-    """
-
-    def __init__(  # pylint: disable=too-many-arguments
-        self,
-        run_id: int,
-        stub: Optional[GrpcDriverStub] = None,
-    ) -> None:
-        self._run_id = run_id
-        self._run: Optional[Run] = None
-        self.stub = stub if stub is not None else GrpcDriverStub()
-        self.node = Node(node_id=0, anonymous=True)
+    def _init_run(self) -> None:
+        # Check if is initialized
+        if self._run is not None:
+            return
+        # Get the run info
+        req = GetRunRequest(run_id=self._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}")
+        self._run = Run(
+            run_id=res.run.run_id,
+            fab_id=res.run.fab_id,
+            fab_version=res.run.fab_version,
+            override_config=user_config_from_proto(res.run.override_config),
+        )
 
     @property
     def run(self) -> Run:
         """Run information."""
-        self._get_stub_and_run_id()
-        return Run(**vars(cast(Run, self._run)))
+        self._init_run()
+        return Run(**vars(self._run))
 
-    def _get_stub_and_run_id(self) -> Tuple[GrpcDriverStub, int]:
-        # Check if is initialized
-        if self._run is None:
-            # Connect
-            if not self.stub.is_connected():
-                self.stub.connect()
-            # Get the run info
-            req = GetRunRequest(run_id=self._run_id)
-            res = self.stub.get_run(req)
-            if not res.HasField("run"):
-                raise RuntimeError(f"Cannot find the run with ID: {self._run_id}")
-            self._run = Run(
-                run_id=res.run.run_id,
-                fab_id=res.run.fab_id,
-                fab_version=res.run.fab_version,
-            )
-
-        return self.stub, self._run.run_id
+    @property
+    def _stub(self) -> DriverStub:
+        """Driver stub."""
+        if not self._is_connected:
+            self._connect()
+        return cast(DriverStub, self._grpc_stub)
 
     def _check_message(self, message: Message) -> None:
         # Check if the message is valid
         if not (
-            message.metadata.run_id == cast(Run, self._run).run_id
+            # Assume self._run being initialized
+            message.metadata.run_id == self._run_id
             and message.metadata.src_node_id == self.node.node_id
             and message.metadata.message_id == ""
             and message.metadata.reply_to_message == ""
@@ -234,7 +172,7 @@ 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.
         """
-        _, run_id = self._get_stub_and_run_id()
+        self._init_run()
         if ttl:
             warnings.warn(
                 "A custom TTL was set, but note that the SuperLink does not enforce "
@@ -245,7 +183,7 @@ class GrpcDriver(Driver):
 
         ttl_ = DEFAULT_TTL if ttl is None else ttl
         metadata = Metadata(
-            run_id=run_id,
+            run_id=self._run_id,
             message_id="",  # Will be set by the server
             src_node_id=self.node.node_id,
             dst_node_id=dst_node_id,
@@ -258,9 +196,11 @@ class GrpcDriver(Driver):
 
     def get_node_ids(self) -> List[int]:
         """Get node IDs."""
-        stub, run_id = self._get_stub_and_run_id()
+        self._init_run()
         # Call GrpcDriverStub method
-        res = stub.get_nodes(GetNodesRequest(run_id=run_id))
+        res: GetNodesResponse = self._stub.GetNodes(
+            GetNodesRequest(run_id=self._run_id)
+        )
         return [node.node_id for node in res.nodes]
 
     def push_messages(self, messages: Iterable[Message]) -> Iterable[str]:
@@ -269,7 +209,7 @@ class GrpcDriver(Driver):
         This method takes an iterable of messages and sends each message
         to the node specified in `dst_node_id`.
         """
-        stub, _ = self._get_stub_and_run_id()
+        self._init_run()
         # Construct TaskIns
         task_ins_list: List[TaskIns] = []
         for msg in messages:
@@ -280,7 +220,9 @@ class GrpcDriver(Driver):
             # Add to list
             task_ins_list.append(taskins)
         # Call GrpcDriverStub method
-        res = stub.push_task_ins(PushTaskInsRequest(task_ins_list=task_ins_list))
+        res: PushTaskInsResponse = self._stub.PushTaskIns(
+            PushTaskInsRequest(task_ins_list=task_ins_list)
+        )
         return list(res.task_ids)
 
     def pull_messages(self, message_ids: Iterable[str]) -> Iterable[Message]:
@@ -289,9 +231,9 @@ class GrpcDriver(Driver):
         This method is used to collect messages from the SuperLink that correspond to a
         set of given message IDs.
         """
-        stub, _ = self._get_stub_and_run_id()
+        self._init_run()
         # Pull TaskRes
-        res = stub.pull_task_res(
+        res: PullTaskResResponse = self._stub.PullTaskRes(
             PullTaskResRequest(node=self.node, task_ids=message_ids)
         )
         # Convert TaskRes to Message
@@ -331,7 +273,7 @@ class GrpcDriver(Driver):
     def close(self) -> None:
         """Disconnect from the SuperLink if connected."""
         # Check if `connect` was called before
-        if not self.stub.is_connected():
+        if not self._is_connected:
             return
         # Disconnect
-        self.stub.disconnect()
+        self._disconnect()

flwr/server/run_serverapp.py

@@ -22,13 +22,22 @@ from pathlib import Path
 from typing import Optional
 
 from flwr.common import Context, EventType, RecordSet, event
-from flwr.common.config import get_flwr_dir, get_project_config, get_project_dir
+from flwr.common.config import (
+    get_flwr_dir,
+    get_fused_config,
+    get_project_config,
+    get_project_dir,
+)
 from flwr.common.logger import log, update_console_handler, warn_deprecated_feature
 from flwr.common.object_ref import load_app
-from flwr.proto.driver_pb2 import CreateRunRequest  # pylint: disable=E0611
+from flwr.common.typing import UserConfig
+from flwr.proto.driver_pb2 import (  # pylint: disable=E0611
+    CreateRunRequest,
+    CreateRunResponse,
+)
 
 from .driver import Driver
-from .driver.grpc_driver import GrpcDriver, GrpcDriverStub
+from .driver.grpc_driver import GrpcDriver
 from .server_app import LoadServerAppError, ServerApp
 
 ADDRESS_DRIVER_API = "0.0.0.0:9091"
@@ -37,6 +46,7 @@ ADDRESS_DRIVER_API = "0.0.0.0:9091"
 def run(
     driver: Driver,
     server_app_dir: str,
+    server_app_run_config: UserConfig,
     server_app_attr: Optional[str] = None,
     loaded_server_app: Optional[ServerApp] = None,
 ) -> None:
@@ -47,9 +57,6 @@ def run(
             "but not both."
         )
 
-    if server_app_dir is not None:
-        sys.path.insert(0, str(Path(server_app_dir).absolute()))
-
     # Load ServerApp if needed
     def _load() -> ServerApp:
         if server_app_attr:
@@ -69,7 +76,9 @@ def run(
     server_app = _load()
 
     # Initialize Context
-    context = Context(state=RecordSet())
+    context = Context(
+        node_id=0, node_config={}, state=RecordSet(), run_config=server_app_run_config
+    )
 
     # Call ServerApp
     server_app(driver=driver, context=context)
@@ -144,22 +153,29 @@ def run_server_app() -> None:  # pylint: disable=too-many-branches
             "For more details, use: ``flower-server-app -h``"
         )
 
-    stub = GrpcDriverStub(
-        driver_service_address=args.superlink, root_certificates=root_certificates
-    )
+    # Initialize GrpcDriver
     if args.run_id is not None:
         # User provided `--run-id`, but not `server-app`
-        run_id = args.run_id
+        driver = GrpcDriver(
+            run_id=args.run_id,
+            driver_service_address=args.superlink,
+            root_certificates=root_certificates,
+        )
     else:
         # User provided `server-app`, but not `--run-id`
         # Create run if run_id is not provided
-        stub.connect()
+        driver = GrpcDriver(
+            run_id=0,  # Will be overwritten
+            driver_service_address=args.superlink,
+            root_certificates=root_certificates,
+        )
+        # Create run
         req = CreateRunRequest(fab_id=args.fab_id, fab_version=args.fab_version)
-        res = stub.create_run(req)
-        run_id = res.run_id
+        res: CreateRunResponse = driver._stub.CreateRun(req)  # pylint: disable=W0212
+        # Overwrite driver._run_id
+        driver._run_id = res.run_id  # pylint: disable=W0212
 
-    # Initialize GrpcDriver
-    driver = GrpcDriver(run_id=run_id, stub=stub)
+    server_app_run_config = {}
 
     # Dynamically obtain ServerApp path based on run_id
     if args.run_id is not None:
@@ -168,7 +184,8 @@ def run_server_app() -> None:  # pylint: disable=too-many-branches
         run_ = driver.run
         server_app_dir = str(get_project_dir(run_.fab_id, run_.fab_version, flwr_dir))
         config = get_project_config(server_app_dir)
-        server_app_attr = config["flower"]["components"]["serverapp"]
+        server_app_attr = config["tool"]["flwr"]["app"]["components"]["serverapp"]
+        server_app_run_config = get_fused_config(run_, flwr_dir)
     else:
         # User provided `server-app`, but not `--run-id`
         server_app_dir = str(Path(args.dir).absolute())
@@ -182,7 +199,12 @@ def run_server_app() -> None:  # pylint: disable=too-many-branches
     )
 
     # Run the ServerApp with the Driver
-    run(driver=driver, server_app_dir=server_app_dir, server_app_attr=server_app_attr)
+    run(
+        driver=driver,
+        server_app_dir=server_app_dir,
+        server_app_run_config=server_app_run_config,
+        server_app_attr=server_app_attr,
+    )
 
     # Clean up
     driver.close()