flwr-nightly 1.8.0.dev20240315__py3-none-any.whl → 1.15.0.dev20250115__py3-none-any.whl

flwr/server/driver/grpc_driver.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,20 +12,30 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 # ==============================================================================
-"""Flower driver service client."""
+"""Flower gRPC Driver."""
 
 
-from logging import DEBUG, ERROR, WARNING
-from typing import Optional
+import time
+import warnings
+from collections.abc import Iterable
+from logging import DEBUG, WARNING
+from typing import Optional, cast
 
 import grpc
 
-from flwr.common import EventType, event
+from flwr.common import DEFAULT_TTL, Message, Metadata, RecordSet
+from flwr.common.constant import (
+    SERVERAPPIO_API_DEFAULT_CLIENT_ADDRESS,
+    SUPERLINK_NODE_ID,
+)
 from flwr.common.grpc import create_channel
 from flwr.common.logger import log
-from flwr.proto.driver_pb2 import (  # pylint: disable=E0611
-    CreateRunRequest,
-    CreateRunResponse,
+from flwr.common.retry_invoker import _make_simple_grpc_retry_invoker, _wrap_stub
+from flwr.common.serde import message_from_taskres, message_to_taskins, run_from_proto
+from flwr.common.typing import Run
+from flwr.proto.node_pb2 import Node  # pylint: disable=E0611
+from flwr.proto.run_pb2 import GetRunRequest, GetRunResponse  # pylint: disable=E0611
+from flwr.proto.serverappio_pb2 import (  # pylint: disable=E0611
     GetNodesRequest,
     GetNodesResponse,
     PullTaskResRequest,
@@ -33,97 +43,226 @@ from flwr.proto.driver_pb2 import (  # pylint: disable=E0611
     PushTaskInsRequest,
     PushTaskInsResponse,
 )
-from flwr.proto.driver_pb2_grpc import DriverStub  # pylint: disable=E0611
+from flwr.proto.serverappio_pb2_grpc import ServerAppIoStub  # pylint: disable=E0611
+from flwr.proto.task_pb2 import TaskIns  # pylint: disable=E0611
 
-DEFAULT_SERVER_ADDRESS_DRIVER = "[::]:9091"
+from .driver import Driver
 
 ERROR_MESSAGE_DRIVER_NOT_CONNECTED = """
 [Driver] Error: Not connected.
 
-Call `connect()` on the `GrpcDriver` instance before calling any of the other
-`GrpcDriver` methods.
+Call `connect()` on the `GrpcDriverStub` instance before calling any of the other
+`GrpcDriverStub` methods.
 """
 
 
-class GrpcDriver:
-    """`GrpcDriver` provides access to the gRPC Driver API/service."""
+class GrpcDriver(Driver):
+    """`GrpcDriver` provides an interface to the ServerAppIo API.
+
+    Parameters
+    ----------
+    serverappio_service_address : str (default: "[::]:9091")
+        The address (URL, IPv6, IPv4) of the SuperLink ServerAppIo 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,
-        driver_service_address: str = DEFAULT_SERVER_ADDRESS_DRIVER,
+        serverappio_service_address: str = SERVERAPPIO_API_DEFAULT_CLIENT_ADDRESS,
         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:
+        self._addr = serverappio_service_address
+        self._cert = root_certificates
+        self._run: Optional[Run] = None
+        self._grpc_stub: Optional[ServerAppIoStub] = None
+        self._channel: Optional[grpc.Channel] = None
+        self.node = Node(node_id=SUPERLINK_NODE_ID)
+        self._retry_invoker = _make_simple_grpc_retry_invoker()
+
+    @property
+    def _is_connected(self) -> bool:
+        """Check if connected to the ServerAppIo API server."""
+        return self._channel is not None
+
+    def _connect(self) -> None:
+        """Connect to the ServerAppIo API.
+
+        This will not call GetRun.
+        """
+        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 = ServerAppIoStub(self._channel)
+        _wrap_stub(self._grpc_stub, self._retry_invoker)
+        log(DEBUG, "[Driver] Connected to %s", self._addr)
 
-    def disconnect(self) -> None:
-        """Disconnect from the Driver API."""
-        event(EventType.DRIVER_DISCONNECT)
-        if self.channel is None or self.stub is None:
+    def _disconnect(self) -> None:
+        """Disconnect from the ServerAppIo API."""
+        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("`GrpcDriver` 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("`GrpcDriver` 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("`GrpcDriver` 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("`GrpcDriver` instance not connected")
-
-        # Call Driver API
-        res: PullTaskResResponse = self.stub.PullTaskRes(request=req)
-        return res
+    def set_run(self, run_id: int) -> None:
+        """Set the run."""
+        # Get the run info
+        req = GetRunRequest(run_id=run_id)
+        res: GetRunResponse = self._stub.GetRun(req)
+        if not res.HasField("run"):
+            raise RuntimeError(f"Cannot find the run with ID: {run_id}")
+        self._run = run_from_proto(res.run)
+
+    @property
+    def run(self) -> Run:
+        """Run information."""
+        return Run(**vars(self._run))
+
+    @property
+    def _stub(self) -> ServerAppIoStub:
+        """ServerAppIo stub."""
+        if not self._is_connected:
+            self._connect()
+        return cast(ServerAppIoStub, self._grpc_stub)
+
+    def _check_message(self, message: Message) -> None:
+        # Check if the message is valid
+        if not (
+            # Assume self._run being initialized
+            message.metadata.run_id == cast(Run, self._run).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,R0917
+        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.
+        """
+        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=cast(Run, self._run).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."""
+        # Call GrpcDriverStub method
+        res: GetNodesResponse = self._stub.GetNodes(
+            GetNodesRequest(run_id=cast(Run, self._run).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`.
+        """
+        # 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 GrpcDriverStub method
+        res: PushTaskInsResponse = self._stub.PushTaskIns(
+            PushTaskInsRequest(
+                task_ins_list=task_ins_list, run_id=cast(Run, self._run).run_id
+            )
+        )
+        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.
+        """
+        # Pull TaskRes
+        res: PullTaskResResponse = self._stub.PullTaskRes(
+            PullTaskResRequest(
+                node=self.node, task_ids=message_ids, run_id=cast(Run, self._run).run_id
+            )
+        )
+        # 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 `connect` was called before
+        if not self._is_connected:
+            return
+        # Disconnect
+        self._disconnect()

flwr/server/driver/inmemory_driver.py

@@ -0,0 +1,183 @@
+# 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 in-memory Driver."""
+
+
+import time
+import warnings
+from collections.abc import Iterable
+from typing import Optional, cast
+from uuid import UUID
+
+from flwr.common import DEFAULT_TTL, Message, Metadata, RecordSet
+from flwr.common.constant import SUPERLINK_NODE_ID
+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.linkstate import LinkStateFactory
+
+from .driver import Driver
+
+
+class InMemoryDriver(Driver):
+    """`InMemoryDriver` class provides an interface to the ServerAppIo API.
+
+    Parameters
+    ----------
+    state_factory : StateFactory
+        A StateFactory embedding a state that this driver can interface with.
+    pull_interval : float (default=0.1)
+        Sleep duration between calls to `pull_messages`.
+    """
+
+    def __init__(
+        self,
+        state_factory: LinkStateFactory,
+        pull_interval: float = 0.1,
+    ) -> None:
+        self._run: Optional[Run] = None
+        self.state = state_factory.state()
+        self.pull_interval = pull_interval
+        self.node = Node(node_id=SUPERLINK_NODE_ID)
+
+    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
+            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 set_run(self, run_id: int) -> None:
+        """Initialize the run."""
+        run = self.state.get_run(run_id)
+        if run is None:
+            raise RuntimeError(f"Cannot find the run with ID: {run_id}")
+        self._run = run
+
+    @property
+    def run(self) -> Run:
+        """Run ID."""
+        return Run(**vars(cast(Run, self._run)))
+
+    def create_message(  # pylint: disable=too-many-arguments,R0917
+        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.
+        """
+        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=cast(Run, self._run).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."""
+        return list(self.state.get_nodes(cast(Run, self._run).run_id))
+
+    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`.
+        """
+        task_ids: list[str] = []
+        for msg in messages:
+            # Check message
+            self._check_message(msg)
+            # Convert Message to TaskIns
+            taskins = message_to_taskins(msg)
+            # Store in state
+            taskins.task.pushed_at = time.time()
+            task_id = self.state.store_task_ins(taskins)
+            if task_id:
+                task_ids.append(str(task_id))
+
+        return 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.
+        """
+        msg_ids = {UUID(msg_id) for msg_id in message_ids}
+        # Pull TaskRes
+        task_res_list = self.state.get_task_res(task_ids=msg_ids)
+        # Delete tasks in state
+        # Delete the TaskIns/TaskRes pairs if TaskRes is found
+        task_ins_ids_to_delete = {
+            UUID(task_res.task.ancestry[0]) for task_res in task_res_list
+        }
+        self.state.delete_tasks(task_ins_ids=task_ins_ids_to_delete)
+        # Convert TaskRes to Message
+        msgs = [message_from_taskres(taskres) for taskres in 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(self.pull_interval)
+        return ret

flwr/server/history.py

@@ -17,7 +17,6 @@
 
 import pprint
 from functools import reduce
-from typing import Dict, List, Tuple
 
 from flwr.common.typing import Scalar
 
@@ -26,11 +25,11 @@ class History:
     """History class for training and/or evaluation metrics collection."""
 
     def __init__(self) -> None:
-        self.losses_distributed: List[Tuple[int, float]] = []
-        self.losses_centralized: List[Tuple[int, float]] = []
-        self.metrics_distributed_fit: Dict[str, List[Tuple[int, Scalar]]] = {}
-        self.metrics_distributed: Dict[str, List[Tuple[int, Scalar]]] = {}
-        self.metrics_centralized: Dict[str, List[Tuple[int, Scalar]]] = {}
+        self.losses_distributed: list[tuple[int, float]] = []
+        self.losses_centralized: list[tuple[int, float]] = []
+        self.metrics_distributed_fit: dict[str, list[tuple[int, Scalar]]] = {}
+        self.metrics_distributed: dict[str, list[tuple[int, Scalar]]] = {}
+        self.metrics_centralized: dict[str, list[tuple[int, Scalar]]] = {}
 
     def add_loss_distributed(self, server_round: int, loss: float) -> None:
         """Add one loss entry (from distributed evaluation)."""
@@ -41,7 +40,7 @@ class History:
         self.losses_centralized.append((server_round, loss))
 
     def add_metrics_distributed_fit(
-        self, server_round: int, metrics: Dict[str, Scalar]
+        self, server_round: int, metrics: dict[str, Scalar]
     ) -> None:
         """Add metrics entries (from distributed fit)."""
         for key in metrics:
@@ -52,7 +51,7 @@ class History:
             self.metrics_distributed_fit[key].append((server_round, metrics[key]))
 
     def add_metrics_distributed(
-        self, server_round: int, metrics: Dict[str, Scalar]
+        self, server_round: int, metrics: dict[str, Scalar]
     ) -> None:
         """Add metrics entries (from distributed evaluation)."""
         for key in metrics:
@@ -63,7 +62,7 @@ class History:
             self.metrics_distributed[key].append((server_round, metrics[key]))
 
     def add_metrics_centralized(
-        self, server_round: int, metrics: Dict[str, Scalar]
+        self, server_round: int, metrics: dict[str, Scalar]
     ) -> None:
         """Add metrics entries (from centralized evaluation)."""
         for key in metrics:
@@ -91,32 +90,32 @@ class History:
         """
         rep = ""
         if self.losses_distributed:
-            rep += "History (loss, distributed):\n" + pprint.pformat(
-                reduce(
-                    lambda a, b: a + b,
-                    [
-                        f"\tround {server_round}: {loss}\n"
-                        for server_round, loss in self.losses_distributed
-                    ],
-                )
+            rep += "History (loss, distributed):\n" + reduce(
+                lambda a, b: a + b,
+                [
+                    f"\tround {server_round}: {loss}\n"
+                    for server_round, loss in self.losses_distributed
+                ],
             )
         if self.losses_centralized:
-            rep += "History (loss, centralized):\n" + pprint.pformat(
-                reduce(
-                    lambda a, b: a + b,
-                    [
-                        f"\tround {server_round}: {loss}\n"
-                        for server_round, loss in self.losses_centralized
-                    ],
-                )
+            rep += "History (loss, centralized):\n" + reduce(
+                lambda a, b: a + b,
+                [
+                    f"\tround {server_round}: {loss}\n"
+                    for server_round, loss in self.losses_centralized
+                ],
             )
         if self.metrics_distributed_fit:
-            rep += "History (metrics, distributed, fit):\n" + pprint.pformat(
-                self.metrics_distributed_fit
+            rep += (
+                "History (metrics, distributed, fit):\n"
+                + pprint.pformat(self.metrics_distributed_fit)
+                + "\n"
             )
         if self.metrics_distributed:
-            rep += "History (metrics, distributed, evaluate):\n" + pprint.pformat(
-                self.metrics_distributed
+            rep += (
+                "History (metrics, distributed, evaluate):\n"
+                + pprint.pformat(self.metrics_distributed)
+                + "\n"
             )
         if self.metrics_centralized:
             rep += "History (metrics, centralized):\n" + pprint.pformat(