flwr-nightly 1.9.0.dev20240420__py3-none-any.whl → 1.9.0.dev20240422__py3-none-any.whl

flwr/client/supernode/app.py

@@ -28,12 +28,11 @@ def run_supernode() -> None:
 
     event(EventType.RUN_SUPERNODE_ENTER)
 
-    args = _parse_args_run_supernode().parse_args()
+    _ = _parse_args_run_supernode().parse_args()
 
     log(
         DEBUG,
-        "Flower will load ClientApp `%s`",
-        getattr(args, "client-app"),
+        "Flower SuperNode starting...",
     )
 
     # Graceful shutdown
@@ -48,7 +47,16 @@ def _parse_args_run_supernode() -> argparse.ArgumentParser:
         description="Start a Flower SuperNode",
     )
 
-    parse_args_run_client_app(parser=parser)
+    parser.add_argument(
+        "client-app",
+        nargs="?",
+        default="",
+        help="For example: `client:app` or `project.package.module:wrapper.app`. "
+        "This is optional and serves as the default ClientApp to be loaded when "
+        "the ServerApp does not specify `fab_id` and `fab_version`. "
+        "If not provided, defaults to an empty string.",
+    )
+    _parse_args_common(parser)
 
     return parser
 
@@ -59,6 +67,10 @@ def parse_args_run_client_app(parser: argparse.ArgumentParser) -> None:
         "client-app",
         help="For example: `client:app` or `project.package.module:wrapper.app`",
     )
+    _parse_args_common(parser)
+
+
+def _parse_args_common(parser: argparse.ArgumentParser) -> None:
     parser.add_argument(
         "--insecure",
         action="store_true",

flwr/server/compat/app.py

@@ -29,7 +29,7 @@ from flwr.server.server import Server, init_defaults, run_fl
 from flwr.server.server_config import ServerConfig
 from flwr.server.strategy import Strategy
 
-from ..driver import Driver
+from ..driver import Driver, GrpcDriver
 from .app_utils import start_update_client_manager_thread
 
 DEFAULT_SERVER_ADDRESS_DRIVER = "[::]:9091"
@@ -114,7 +114,7 @@ def start_driver(  # pylint: disable=too-many-arguments, too-many-locals
         # Create the Driver
         if isinstance(root_certificates, str):
             root_certificates = Path(root_certificates).read_bytes()
-        driver = Driver(
+        driver = GrpcDriver(
             driver_service_address=address, root_certificates=root_certificates
         )
 

flwr/server/compat/driver_client_proxy.py

@@ -25,7 +25,7 @@ from flwr.common import serde
 from flwr.proto import driver_pb2, node_pb2, task_pb2  # pylint: disable=E0611
 from flwr.server.client_proxy import ClientProxy
 
-from ..driver.driver import GrpcDriverHelper
+from ..driver.grpc_driver import GrpcDriverHelper
 
 SLEEP_TIME = 1
 

flwr/server/driver/__init__.py

@@ -16,7 +16,10 @@
 
 
 from .driver import Driver
+from .grpc_driver import GrpcDriver, GrpcDriverHelper
 
 __all__ = [
     "Driver",
+    "GrpcDriver",
+    "GrpcDriverHelper",
 ]

flwr/server/driver/driver.py

@@ -1,4 +1,4 @@
-# Copyright 2022 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,180 +12,19 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 # ==============================================================================
-"""Flower driver service client."""
+"""Driver (abstract base class)."""
 
-import time
-import warnings
-from logging import DEBUG, ERROR, WARNING
-from typing import Iterable, List, Optional, Tuple
 
-import grpc
+from abc import ABC, abstractmethod
+from typing import Iterable, List, Optional
 
-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.proto.driver_pb2 import (  # pylint: disable=E0611
-    CreateRunRequest,
-    CreateRunResponse,
-    GetNodesRequest,
-    GetNodesResponse,
-    PullTaskResRequest,
-    PullTaskResResponse,
-    PushTaskInsRequest,
-    PushTaskInsResponse,
-)
-from flwr.proto.driver_pb2_grpc import DriverStub  # pylint: disable=E0611
-from flwr.proto.node_pb2 import Node  # pylint: disable=E0611
-from flwr.proto.task_pb2 import TaskIns  # pylint: disable=E0611
+from flwr.common import Message, RecordSet
 
-DEFAULT_SERVER_ADDRESS_DRIVER = "[::]:9091"
 
-ERROR_MESSAGE_DRIVER_NOT_CONNECTED = """
-[Driver] Error: Not connected.
-
-Call `connect()` on the `GrpcDriverHelper` instance before calling any of the other
-`GrpcDriverHelper` methods.
-"""
-
-
-class GrpcDriverHelper:
-    """`GrpcDriverHelper` provides access to the gRPC Driver API/service."""
-
-    def __init__(
-        self,
-        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
-
-    def connect(self) -> None:
-        """Connect to the Driver API."""
-        event(EventType.DRIVER_CONNECT)
-        if self.channel is not None or self.stub is not None:
-            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.stub = DriverStub(self.channel)
-        log(DEBUG, "[Driver] Connected to %s", self.driver_service_address)
-
-    def disconnect(self) -> None:
-        """Disconnect from the Driver API."""
-        event(EventType.DRIVER_DISCONNECT)
-        if self.channel is None or self.stub is None:
-            log(DEBUG, "Already disconnected")
-            return
-        channel = self.channel
-        self.channel = None
-        self.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("`GrpcDriverHelper` instance not connected")
-
-        # Call Driver API
-        res: CreateRunResponse = self.stub.CreateRun(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("`GrpcDriverHelper` 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("`GrpcDriverHelper` 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("`GrpcDriverHelper` instance not connected")
-
-        # Call Driver API
-        res: PullTaskResResponse = self.stub.PullTaskRes(request=req)
-        return res
-
-
-class Driver:
-    """`Driver` class 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"`.
-    certificates : bytes (default: None)
-        Tuple containing root certificate, server certificate, and private key
-        to start a secure SSL-enabled server. The tuple is expected to have
-        three bytes elements in the following order:
-
-            * CA certificate.
-            * server certificate.
-            * server private key.
-    """
-
-    def __init__(
-        self,
-        driver_service_address: str = DEFAULT_SERVER_ADDRESS_DRIVER,
-        root_certificates: Optional[bytes] = None,
-    ) -> None:
-        self.addr = driver_service_address
-        self.root_certificates = root_certificates
-        self.grpc_driver_helper: Optional[GrpcDriverHelper] = None
-        self.run_id: Optional[int] = None
-        self.node = Node(node_id=0, anonymous=True)
-
-    def _get_grpc_driver_helper_and_run_id(self) -> Tuple[GrpcDriverHelper, int]:
-        # Check if the GrpcDriverHelper is initialized
-        if self.grpc_driver_helper is None or self.run_id is None:
-            # Connect and create run
-            self.grpc_driver_helper = GrpcDriverHelper(
-                driver_service_address=self.addr,
-                root_certificates=self.root_certificates,
-            )
-            self.grpc_driver_helper.connect()
-            res = self.grpc_driver_helper.create_run(CreateRunRequest())
-            self.run_id = res.run_id
-        return self.grpc_driver_helper, self.run_id
-
-    def _check_message(self, message: Message) -> None:
-        # Check if the message is valid
-        if not (
-            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 == ""
-            and message.metadata.ttl > 0
-        ):
-            raise ValueError(f"Invalid message: {message}")
+class Driver(ABC):
+    """Abstract base Driver class for the Driver API."""
 
+    @abstractmethod
     def create_message(  # pylint: disable=too-many-arguments
         self,
         content: RecordSet,
@@ -223,35 +62,12 @@ class Driver:
         message : Message
             A new `Message` instance with the specified content and metadata.
         """
-        _, run_id = self._get_grpc_driver_helper_and_run_id()
-        if ttl:
-            warnings.warn(
-                "A custom TTL was set, but note that the SuperLink does not enforce "
-                "the TTL yet. The SuperLink will start enforcing the TTL in a future "
-                "version of Flower.",
-                stacklevel=2,
-            )
-
-        ttl_ = DEFAULT_TTL if ttl is None else ttl
-        metadata = Metadata(
-            run_id=run_id,
-            message_id="",  # Will be set by the server
-            src_node_id=self.node.node_id,
-            dst_node_id=dst_node_id,
-            reply_to_message="",
-            group_id=group_id,
-            ttl=ttl_,
-            message_type=message_type,
-        )
-        return Message(metadata=metadata, content=content)
 
+    @abstractmethod
     def get_node_ids(self) -> List[int]:
         """Get node IDs."""
-        grpc_driver_helper, run_id = self._get_grpc_driver_helper_and_run_id()
-        # Call GrpcDriverHelper method
-        res = grpc_driver_helper.get_nodes(GetNodesRequest(run_id=run_id))
-        return [node.node_id for node in res.nodes]
 
+    @abstractmethod
     def push_messages(self, messages: Iterable[Message]) -> Iterable[str]:
         """Push messages to specified node IDs.
 
@@ -269,22 +85,8 @@ class Driver:
             An iterable of IDs for the messages that were sent, which can be used
             to pull replies.
         """
-        grpc_driver_helper, _ = self._get_grpc_driver_helper_and_run_id()
-        # Construct TaskIns
-        task_ins_list: List[TaskIns] = []
-        for msg in messages:
-            # Check message
-            self._check_message(msg)
-            # Convert Message to TaskIns
-            taskins = message_to_taskins(msg)
-            # Add to list
-            task_ins_list.append(taskins)
-        # Call GrpcDriverHelper method
-        res = grpc_driver_helper.push_task_ins(
-            PushTaskInsRequest(task_ins_list=task_ins_list)
-        )
-        return list(res.task_ids)
 
+    @abstractmethod
     def pull_messages(self, message_ids: Iterable[str]) -> Iterable[Message]:
         """Pull messages based on message IDs.
 
@@ -301,15 +103,8 @@ class Driver:
         messages : Iterable[Message]
             An iterable of messages received.
         """
-        grpc_driver, _ = self._get_grpc_driver_helper_and_run_id()
-        # Pull TaskRes
-        res = grpc_driver.pull_task_res(
-            PullTaskResRequest(node=self.node, task_ids=message_ids)
-        )
-        # Convert TaskRes to Message
-        msgs = [message_from_taskres(taskres) for taskres in res.task_res_list]
-        return msgs
 
+    @abstractmethod
     def send_and_receive(
         self,
         messages: Iterable[Message],
@@ -343,28 +138,3 @@ class Driver:
         replies for all sent messages. A message remains valid until its TTL,
         which is not affected by `timeout`.
         """
-        # Push messages
-        msg_ids = set(self.push_messages(messages))
-
-        # Pull messages
-        end_time = time.time() + (timeout if timeout is not None else 0.0)
-        ret: List[Message] = []
-        while timeout is None or time.time() < end_time:
-            res_msgs = self.pull_messages(msg_ids)
-            ret.extend(res_msgs)
-            msg_ids.difference_update(
-                {msg.metadata.reply_to_message for msg in res_msgs}
-            )
-            if len(msg_ids) == 0:
-                break
-            # Sleep
-            time.sleep(3)
-        return ret
-
-    def close(self) -> None:
-        """Disconnect from the SuperLink if connected."""
-        # Check if GrpcDriverHelper is initialized
-        if self.grpc_driver_helper is None:
-            return
-        # Disconnect
-        self.grpc_driver_helper.disconnect()

flwr/server/driver/grpc_driver.py

@@ -0,0 +1,306 @@
+# Copyright 2022 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 gRPC Driver."""
+
+import time
+import warnings
+from logging import DEBUG, ERROR, WARNING
+from typing import Iterable, List, Optional, Tuple
+
+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.proto.driver_pb2 import (  # pylint: disable=E0611
+    CreateRunRequest,
+    CreateRunResponse,
+    GetNodesRequest,
+    GetNodesResponse,
+    PullTaskResRequest,
+    PullTaskResResponse,
+    PushTaskInsRequest,
+    PushTaskInsResponse,
+)
+from flwr.proto.driver_pb2_grpc import DriverStub  # pylint: disable=E0611
+from flwr.proto.node_pb2 import Node  # pylint: disable=E0611
+from flwr.proto.task_pb2 import TaskIns  # pylint: disable=E0611
+
+from .driver import Driver
+
+DEFAULT_SERVER_ADDRESS_DRIVER = "[::]:9091"
+
+ERROR_MESSAGE_DRIVER_NOT_CONNECTED = """
+[Driver] Error: Not connected.
+
+Call `connect()` on the `GrpcDriverHelper` instance before calling any of the other
+`GrpcDriverHelper` methods.
+"""
+
+
+class GrpcDriverHelper:
+    """`GrpcDriverHelper` provides access to the gRPC Driver API/service."""
+
+    def __init__(
+        self,
+        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
+
+    def connect(self) -> None:
+        """Connect to the Driver API."""
+        event(EventType.DRIVER_CONNECT)
+        if self.channel is not None or self.stub is not None:
+            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.stub = DriverStub(self.channel)
+        log(DEBUG, "[Driver] Connected to %s", self.driver_service_address)
+
+    def disconnect(self) -> None:
+        """Disconnect from the Driver API."""
+        event(EventType.DRIVER_DISCONNECT)
+        if self.channel is None or self.stub is None:
+            log(DEBUG, "Already disconnected")
+            return
+        channel = self.channel
+        self.channel = None
+        self.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("`GrpcDriverHelper` instance not connected")
+
+        # Call Driver API
+        res: CreateRunResponse = self.stub.CreateRun(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("`GrpcDriverHelper` 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("`GrpcDriverHelper` 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("`GrpcDriverHelper` 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
+    ----------
+    driver_service_address : Optional[str]
+        The IPv4 or IPv6 address of the Driver API server.
+        Defaults to `"[::]:9091"`.
+    certificates : bytes (default: None)
+        Tuple containing root certificate, server certificate, and private key
+        to start a secure SSL-enabled server. The tuple is expected to have
+        three bytes elements in the following order:
+
+            * CA certificate.
+            * server certificate.
+            * server private key.
+    """
+
+    def __init__(
+        self,
+        driver_service_address: str = DEFAULT_SERVER_ADDRESS_DRIVER,
+        root_certificates: Optional[bytes] = None,
+    ) -> None:
+        self.addr = driver_service_address
+        self.root_certificates = root_certificates
+        self.grpc_driver_helper: Optional[GrpcDriverHelper] = None
+        self.run_id: Optional[int] = None
+        self.node = Node(node_id=0, anonymous=True)
+
+    def _get_grpc_driver_helper_and_run_id(self) -> Tuple[GrpcDriverHelper, int]:
+        # Check if the GrpcDriverHelper is initialized
+        if self.grpc_driver_helper is None or self.run_id is None:
+            # Connect and create run
+            self.grpc_driver_helper = GrpcDriverHelper(
+                driver_service_address=self.addr,
+                root_certificates=self.root_certificates,
+            )
+            self.grpc_driver_helper.connect()
+            res = self.grpc_driver_helper.create_run(CreateRunRequest())
+            self.run_id = res.run_id
+        return self.grpc_driver_helper, self.run_id
+
+    def _check_message(self, message: Message) -> None:
+        # Check if the message is valid
+        if not (
+            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 == ""
+            and message.metadata.ttl > 0
+        ):
+            raise ValueError(f"Invalid message: {message}")
+
+    def create_message(  # pylint: disable=too-many-arguments
+        self,
+        content: RecordSet,
+        message_type: str,
+        dst_node_id: int,
+        group_id: str,
+        ttl: Optional[float] = None,
+    ) -> Message:
+        """Create a new message with specified parameters.
+
+        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_grpc_driver_helper_and_run_id()
+        if ttl:
+            warnings.warn(
+                "A custom TTL was set, but note that the SuperLink does not enforce "
+                "the TTL yet. The SuperLink will start enforcing the TTL in a future "
+                "version of Flower.",
+                stacklevel=2,
+            )
+
+        ttl_ = DEFAULT_TTL if ttl is None else ttl
+        metadata = Metadata(
+            run_id=run_id,
+            message_id="",  # Will be set by the server
+            src_node_id=self.node.node_id,
+            dst_node_id=dst_node_id,
+            reply_to_message="",
+            group_id=group_id,
+            ttl=ttl_,
+            message_type=message_type,
+        )
+        return Message(metadata=metadata, content=content)
+
+    def get_node_ids(self) -> List[int]:
+        """Get node IDs."""
+        grpc_driver_helper, run_id = self._get_grpc_driver_helper_and_run_id()
+        # Call GrpcDriverHelper method
+        res = grpc_driver_helper.get_nodes(GetNodesRequest(run_id=run_id))
+        return [node.node_id for node in res.nodes]
+
+    def push_messages(self, messages: Iterable[Message]) -> Iterable[str]:
+        """Push messages to specified node IDs.
+
+        This method takes an iterable of messages and sends each message
+        to the node specified in `dst_node_id`.
+        """
+        grpc_driver_helper, _ = self._get_grpc_driver_helper_and_run_id()
+        # Construct TaskIns
+        task_ins_list: List[TaskIns] = []
+        for msg in messages:
+            # Check message
+            self._check_message(msg)
+            # Convert Message to TaskIns
+            taskins = message_to_taskins(msg)
+            # Add to list
+            task_ins_list.append(taskins)
+        # Call GrpcDriverHelper method
+        res = grpc_driver_helper.push_task_ins(
+            PushTaskInsRequest(task_ins_list=task_ins_list)
+        )
+        return list(res.task_ids)
+
+    def pull_messages(self, message_ids: Iterable[str]) -> Iterable[Message]:
+        """Pull messages based on message IDs.
+
+        This method is used to collect messages from the SuperLink that correspond to a
+        set of given message IDs.
+        """
+        grpc_driver, _ = self._get_grpc_driver_helper_and_run_id()
+        # Pull TaskRes
+        res = grpc_driver.pull_task_res(
+            PullTaskResRequest(node=self.node, task_ids=message_ids)
+        )
+        # Convert TaskRes to Message
+        msgs = [message_from_taskres(taskres) for taskres in res.task_res_list]
+        return msgs
+
+    def send_and_receive(
+        self,
+        messages: Iterable[Message],
+        *,
+        timeout: Optional[float] = None,
+    ) -> Iterable[Message]:
+        """Push messages to specified node IDs and pull the reply messages.
+
+        This method sends a list of messages to their destination node IDs and then
+        waits for the replies. It continues to pull replies until either all replies are
+        received or the specified timeout duration is exceeded.
+        """
+        # Push messages
+        msg_ids = set(self.push_messages(messages))
+
+        # Pull messages
+        end_time = time.time() + (timeout if timeout is not None else 0.0)
+        ret: List[Message] = []
+        while timeout is None or time.time() < end_time:
+            res_msgs = self.pull_messages(msg_ids)
+            ret.extend(res_msgs)
+            msg_ids.difference_update(
+                {msg.metadata.reply_to_message for msg in res_msgs}
+            )
+            if len(msg_ids) == 0:
+                break
+            # Sleep
+            time.sleep(3)
+        return ret
+
+    def close(self) -> None:
+        """Disconnect from the SuperLink if connected."""
+        # Check if GrpcDriverHelper is initialized
+        if self.grpc_driver_helper is None:
+            return
+        # Disconnect
+        self.grpc_driver_helper.disconnect()

flwr/server/run_serverapp.py

@@ -25,7 +25,7 @@ from flwr.common import Context, EventType, RecordSet, event
 from flwr.common.logger import log, update_console_handler
 from flwr.common.object_ref import load_app
 
-from .driver.driver import Driver
+from .driver import Driver, GrpcDriver
 from .server_app import LoadServerAppError, ServerApp
 
 
@@ -128,13 +128,13 @@ def run_server_app() -> None:
     server_app_dir = args.dir
     server_app_attr = getattr(args, "server-app")
 
-    # Initialize Driver
-    driver = Driver(
+    # Initialize GrpcDriver
+    driver = GrpcDriver(
         driver_service_address=args.server,
         root_certificates=root_certificates,
     )
 
-    # Run the Server App with the Driver
+    # Run the ServerApp with the Driver
     run(driver=driver, server_app_dir=server_app_dir, server_app_attr=server_app_attr)
 
     # Clean up

flwr/simulation/run_simulation.py

@@ -29,7 +29,7 @@ import grpc
 from flwr.client import ClientApp
 from flwr.common import EventType, event, log
 from flwr.common.typing import ConfigsRecordValues
-from flwr.server.driver.driver import Driver
+from flwr.server.driver import Driver, GrpcDriver
 from flwr.server.run_serverapp import run
 from flwr.server.server_app import ServerApp
 from flwr.server.superlink.driver.driver_grpc import run_driver_api_grpc
@@ -204,7 +204,7 @@ def _main_loop(
     serverapp_th = None
     try:
         # Initialize Driver
-        driver = Driver(
+        driver = GrpcDriver(
             driver_service_address=driver_api_address,
             root_certificates=None,
         )

{flwr_nightly-1.9.0.dev20240420.dist-info → flwr_nightly-1.9.0.dev20240422.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: flwr-nightly
-Version: 1.9.0.dev20240420
+Version: 1.9.0.dev20240422
 Summary: Flower: A Friendly Federated Learning Framework
 Home-page: https://flower.ai
 License: Apache-2.0

{flwr_nightly-1.9.0.dev20240420.dist-info → flwr_nightly-1.9.0.dev20240422.dist-info}/RECORD

@@ -51,7 +51,7 @@ flwr/client/numpy_client.py,sha256=u76GWAdHmJM88Agm2EgLQSvO8Jnk225mJTk-_TmPjFE,1
 flwr/client/rest_client/__init__.py,sha256=ThwOnkMdzxo_UuyTI47Q7y9oSpuTgNT2OuFvJCfuDiw,735
 flwr/client/rest_client/connection.py,sha256=ZxTFVDXlONqKTX6uYgxshoEWqzqVcQ8QQ2hKS93oLM8,11302
 flwr/client/supernode/__init__.py,sha256=D5swXxemuRbA2rB_T9B8LwJW-_PucXwmlFQQerwIUv0,793
-flwr/client/supernode/app.py,sha256=JXRZ76JdyAkhfaEEqsMiONWVQ0bn8YqzZg9oHC4Qfko,3436
+flwr/client/supernode/app.py,sha256=gauvN8elkIy0vuT0GxT7MmkuBRY74ckZfpxejE7dduM,3861
 flwr/client/typing.py,sha256=c9EvjlEjasxn1Wqx6bGl6Xg6vM1gMFfmXht-E2i5J-k,1006
 flwr/common/__init__.py,sha256=dHOptgKxna78CEQLD5Yu0QIsoSgpIIw5AhIUZCHDWAU,3721
 flwr/common/address.py,sha256=iTAN9jtmIGMrWFnx9XZQl45ZEtQJVZZLYPRBSNVARGI,1882
@@ -124,16 +124,16 @@ flwr/server/app.py,sha256=FriloRrkDHTlB5G7EBn6sH4v5GhiYFf_ZhbdROgjKbY,24199
 flwr/server/client_manager.py,sha256=T8UDSRJBVD3fyIDI7NTAA-NA7GPrMNNgH2OAF54RRxE,6127
 flwr/server/client_proxy.py,sha256=4G-oTwhb45sfWLx2uZdcXD98IZwdTS6F88xe3akCdUg,2399
 flwr/server/compat/__init__.py,sha256=VxnJtJyOjNFQXMNi9hIuzNlZM5n0Hj1p3aq_Pm2udw4,892
-flwr/server/compat/app.py,sha256=3Skh76Rg80B4oME1dJOhZvn9eTfVmTNIQ0jCiZ6CzeQ,5271
+flwr/server/compat/app.py,sha256=BhF3DySbvKkOIyNXnB1rwZhw8cC8yK_w91Fku8HmC_w,5287
 flwr/server/compat/app_utils.py,sha256=S-M4sGIiZPXXgKFLjlbFP2yN7d-oIj6DaiJNPIZ2z3A,3503
-flwr/server/compat/driver_client_proxy.py,sha256=XbiGCBGIyZbEaY4Ayu70DwOQHNCVOFBQOeRu8AxfGOw,7365
+flwr/server/compat/driver_client_proxy.py,sha256=5XWroBrtA8MrQ5xQjgsju5RauMxNPshYLS_EtONEL1I,7370
 flwr/server/compat/legacy_context.py,sha256=D2s7PvQoDnTexuRmf1uG9Von7GUj4Qqyr7qLklSlKAM,1766
 flwr/server/criterion.py,sha256=ypbAexbztzGUxNen9RCHF91QeqiEQix4t4Ih3E-42MM,1061
-flwr/server/driver/__init__.py,sha256=TTW0bqKrbv8l2Nem2UbLsugW__ahkxULpRqA4eWpu98,766
-flwr/server/driver/abc_driver.py,sha256=t9SSSDlo9wT_y2Nl7waGYMTm2VlkvK3_bOb7ggPPlho,5090
-flwr/server/driver/driver.py,sha256=STP6X_uDKaK9yB5YqX1Ch4-xcCwaQPe260OhMsPsgiM,13928
+flwr/server/driver/__init__.py,sha256=bbVL5pyA0Y2HcUK4s5U0B4epI-BuUFyEJbchew_8tJY,862
+flwr/server/driver/driver.py,sha256=t9SSSDlo9wT_y2Nl7waGYMTm2VlkvK3_bOb7ggPPlho,5090
+flwr/server/driver/grpc_driver.py,sha256=U5zfI3uYPUBaoOe4JI32t3dvCoSDacZ6EE0g9B8tKbU,11418
 flwr/server/history.py,sha256=hDsoBaA4kUa6d1yvDVXuLluBqOBKSm0_fVDtUtYJkmg,5121
-flwr/server/run_serverapp.py,sha256=3hoXa57T4L1vOWVWPSSdZ_UyRO-uTwUIrhha6TJAXMg,5592
+flwr/server/run_serverapp.py,sha256=3FqKVdFJ280dOVQQ63fu3kL7yNg_4ggtx2H7ljSBT1c,5604
 flwr/server/server.py,sha256=UnBRlI6AGTj0nKeRtEQ3IalM3TJmggMKXhDyn8yKZNk,17664
 flwr/server/server_app.py,sha256=KgAT_HqsfseTLNnfX2ph42PBbVqQ0lFzvYrT90V34y0,4402
 flwr/server/server_config.py,sha256=CZaHVAsMvGLjpWVcLPkiYxgJN4xfIyAiUrCI3fETKY4,1349
@@ -204,9 +204,9 @@ flwr/simulation/ray_transport/__init__.py,sha256=FsaAnzC4cw4DqoouBCix6496k29jACk
 flwr/simulation/ray_transport/ray_actor.py,sha256=_wv2eP7qxkCZ-6rMyYWnjLrGPBZRxjvTPjaVk8zIaQ4,19367
 flwr/simulation/ray_transport/ray_client_proxy.py,sha256=oDu4sEPIOu39vrNi-fqDAe10xtNUXMO49bM2RWfRcyw,6738
 flwr/simulation/ray_transport/utils.py,sha256=TYdtfg1P9VfTdLMOJlifInGpxWHYs9UfUqIv2wfkRLA,2392
-flwr/simulation/run_simulation.py,sha256=HiIH6aa_v56NfKQN5ZBd94NyVfaZNyFs43_kItYsQXU,15685
-flwr_nightly-1.9.0.dev20240420.dist-info/LICENSE,sha256=z8d0m5b2O9McPEK1xHG_dWgUBT6EfBDz6wA0F7xSPTA,11358
-flwr_nightly-1.9.0.dev20240420.dist-info/METADATA,sha256=xxYZp-8KwhbZqeNsEFC-ngcjxMxaG-7IgMBw164-oBQ,15260
-flwr_nightly-1.9.0.dev20240420.dist-info/WHEEL,sha256=FMvqSimYX_P7y0a7UY-_Mc83r5zkBZsCYPm7Lr0Bsq4,88
-flwr_nightly-1.9.0.dev20240420.dist-info/entry_points.txt,sha256=DBrrf685V2W9NbbchQwvuqBEpj5ik8tMZNoZg_W2bZY,363
-flwr_nightly-1.9.0.dev20240420.dist-info/RECORD,,
+flwr/simulation/run_simulation.py,sha256=nxXNv3r8ODImd5o6f0sa_w5L0I08LD2Udw2OTXStRnQ,15694
+flwr_nightly-1.9.0.dev20240422.dist-info/LICENSE,sha256=z8d0m5b2O9McPEK1xHG_dWgUBT6EfBDz6wA0F7xSPTA,11358
+flwr_nightly-1.9.0.dev20240422.dist-info/METADATA,sha256=2g_AiXLNJzV4x9RNTWo1h1LjzMpUdhUQ8uNAPPxqlv8,15260
+flwr_nightly-1.9.0.dev20240422.dist-info/WHEEL,sha256=FMvqSimYX_P7y0a7UY-_Mc83r5zkBZsCYPm7Lr0Bsq4,88
+flwr_nightly-1.9.0.dev20240422.dist-info/entry_points.txt,sha256=DBrrf685V2W9NbbchQwvuqBEpj5ik8tMZNoZg_W2bZY,363
+flwr_nightly-1.9.0.dev20240422.dist-info/RECORD,,

flwr/server/driver/abc_driver.py

@@ -1,140 +0,0 @@
-# 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.
-# ==============================================================================
-"""Driver (abstract base class)."""
-
-
-from abc import ABC, abstractmethod
-from typing import Iterable, List, Optional
-
-from flwr.common import Message, RecordSet
-
-
-class Driver(ABC):
-    """Abstract base Driver class for the Driver API."""
-
-    @abstractmethod
-    def create_message(  # pylint: disable=too-many-arguments
-        self,
-        content: RecordSet,
-        message_type: str,
-        dst_node_id: int,
-        group_id: str,
-        ttl: Optional[float] = None,
-    ) -> Message:
-        """Create a new message with specified parameters.
-
-        This method constructs a new `Message` with given content and metadata.
-        The `run_id` and `src_node_id` will be set automatically.
-
-        Parameters
-        ----------
-        content : RecordSet
-            The content for the new message. This holds records that are to be sent
-            to the destination node.
-        message_type : str
-            The type of the message, defining the action to be executed on
-            the receiving end.
-        dst_node_id : int
-            The ID of the destination node to which the message is being sent.
-        group_id : str
-            The ID of the group to which this message is associated. In some settings,
-            this is used as the FL round.
-        ttl : Optional[float] (default: None)
-            Time-to-live for the round trip of this message, i.e., the time from sending
-            this message to receiving a reply. It specifies in seconds the duration for
-            which the message and its potential reply are considered valid. If unset,
-            the default TTL (i.e., `common.DEFAULT_TTL`) will be used.
-
-        Returns
-        -------
-        message : Message
-            A new `Message` instance with the specified content and metadata.
-        """
-
-    @abstractmethod
-    def get_node_ids(self) -> List[int]:
-        """Get node IDs."""
-
-    @abstractmethod
-    def push_messages(self, messages: Iterable[Message]) -> Iterable[str]:
-        """Push messages to specified node IDs.
-
-        This method takes an iterable of messages and sends each message
-        to the node specified in `dst_node_id`.
-
-        Parameters
-        ----------
-        messages : Iterable[Message]
-            An iterable of messages to be sent.
-
-        Returns
-        -------
-        message_ids : Iterable[str]
-            An iterable of IDs for the messages that were sent, which can be used
-            to pull replies.
-        """
-
-    @abstractmethod
-    def pull_messages(self, message_ids: Iterable[str]) -> Iterable[Message]:
-        """Pull messages based on message IDs.
-
-        This method is used to collect messages from the SuperLink
-        that correspond to a set of given message IDs.
-
-        Parameters
-        ----------
-        message_ids : Iterable[str]
-            An iterable of message IDs for which reply messages are to be retrieved.
-
-        Returns
-        -------
-        messages : Iterable[Message]
-            An iterable of messages received.
-        """
-
-    @abstractmethod
-    def send_and_receive(
-        self,
-        messages: Iterable[Message],
-        *,
-        timeout: Optional[float] = None,
-    ) -> Iterable[Message]:
-        """Push messages to specified node IDs and pull the reply messages.
-
-        This method sends a list of messages to their destination node IDs and then
-        waits for the replies. It continues to pull replies until either all
-        replies are received or the specified timeout duration is exceeded.
-
-        Parameters
-        ----------
-        messages : Iterable[Message]
-            An iterable of messages to be sent.
-        timeout : Optional[float] (default: None)
-            The timeout duration in seconds. If specified, the method will wait for
-            replies for this duration. If `None`, there is no time limit and the method
-            will wait until replies for all messages are received.
-
-        Returns
-        -------
-        replies : Iterable[Message]
-            An iterable of reply messages received from the SuperLink.
-
-        Notes
-        -----
-        This method uses `push_messages` to send the messages and `pull_messages`
-        to collect the replies. If `timeout` is set, the method may not return
-        replies for all sent messages. A message remains valid until its TTL,
-        which is not affected by `timeout`.
-        """