flwr 1.15.2__py3-none-any.whl → 1.17.0__py3-none-any.whl

flwr/server/{driver/grpc_driver.py → grid/grpc_grid.py}

@@ -1,4 +1,4 @@
-# Copyright 2024 Flower Labs GmbH. All Rights Reserved.
+# Copyright 2025 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,24 +12,23 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 # ==============================================================================
-"""Flower gRPC Driver."""
+"""Flower gRPC Grid."""
 
 
 import time
-import warnings
 from collections.abc import Iterable
-from logging import DEBUG, WARNING
+from logging import DEBUG, ERROR, WARNING
 from typing import Optional, cast
 
 import grpc
 
-from flwr.common import DEFAULT_TTL, Message, Metadata, RecordSet
+from flwr.common import Message, RecordDict
 from flwr.common.constant import (
     SERVERAPPIO_API_DEFAULT_CLIENT_ADDRESS,
     SUPERLINK_NODE_ID,
 )
 from flwr.common.grpc import create_channel, on_channel_state_change
-from flwr.common.logger import log
+from flwr.common.logger import log, warn_deprecated_feature
 from flwr.common.retry_invoker import _make_simple_grpc_retry_invoker, _wrap_stub
 from flwr.common.serde import message_from_proto, message_to_proto, run_from_proto
 from flwr.common.typing import Run
@@ -46,18 +45,39 @@ from flwr.proto.serverappio_pb2 import (  # pylint: disable=E0611
 )
 from flwr.proto.serverappio_pb2_grpc import ServerAppIoStub  # pylint: disable=E0611
 
-from .driver import Driver
+from .grid import Grid
 
-ERROR_MESSAGE_DRIVER_NOT_CONNECTED = """
-[flwr-serverapp] Error: Not connected.
+ERROR_MESSAGE_PUSH_MESSAGES_RESOURCE_EXHAUSTED = """
 
-Call `connect()` on the `GrpcDriverStub` instance before calling any of the other
-`GrpcDriverStub` methods.
+[Grid.push_messages] gRPC error occurred:
+
+The 2GB gRPC limit has been reached. Consider reducing the number of messages pushed
+at once, or push messages individually, for example:
+
+> msgs = [msg1, msg2, msg3]
+> msg_ids = []
+> for msg in msgs:
+>     msg_id = grid.push_messages([msg])
+>     msg_ids.extend(msg_id)
+"""
+
+ERROR_MESSAGE_PULL_MESSAGES_RESOURCE_EXHAUSTED = """
+
+[Grid.pull_messages] gRPC error occurred:
+
+The 2GB gRPC limit has been reached. Consider reducing the number of messages pulled
+at once, or pull messages individually, for example:
+
+> msgs_ids = [msg_id1, msg_id2, msg_id3]
+> msgs = []
+> for msg_id in msg_ids:
+>     msg = grid.pull_messages([msg_id])
+>     msgs.extend(msg)
 """
 
 
-class GrpcDriver(Driver):
-    """`GrpcDriver` provides an interface to the ServerAppIo API.
+class GrpcGrid(Grid):
+    """`GrpcGrid` provides an interface to the ServerAppIo API.
 
     Parameters
     ----------
@@ -69,6 +89,8 @@ class GrpcDriver(Driver):
         established to an SSL-enabled Flower server.
     """
 
+    _deprecation_warning_logged = False
+
     def __init__(  # pylint: disable=too-many-arguments
         self,
         serverappio_service_address: str = SERVERAPPIO_API_DEFAULT_CLIENT_ADDRESS,
@@ -81,6 +103,7 @@ class GrpcDriver(Driver):
         self._channel: Optional[grpc.Channel] = None
         self.node = Node(node_id=SUPERLINK_NODE_ID)
         self._retry_invoker = _make_simple_grpc_retry_invoker()
+        super().__init__()
 
     @property
     def _is_connected(self) -> bool:
@@ -140,18 +163,15 @@ class GrpcDriver(Driver):
     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 == ""
+            message.metadata.message_id == ""
+            and message.metadata.reply_to_message_id == ""
             and message.metadata.ttl > 0
         ):
             raise ValueError(f"Invalid message: {message}")
 
     def create_message(  # pylint: disable=too-many-arguments,R0917
         self,
-        content: RecordSet,
+        content: RecordDict,
         message_type: str,
         dst_node_id: int,
         group_id: str,
@@ -162,30 +182,17 @@ 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.
         """
-        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,
+        if not GrpcGrid._deprecation_warning_logged:
+            GrpcGrid._deprecation_warning_logged = True
+            warn_deprecated_feature(
+                "`Driver.create_message` / `Grid.create_message` is deprecated."
+                "Use `Message` constructor instead."
             )
+        return Message(content, dst_node_id, message_type, ttl=ttl, group_id=group_id)
 
-        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]:
+    def get_node_ids(self) -> Iterable[int]:
         """Get node IDs."""
-        # Call GrpcDriverStub method
+        # Call GrpcServerAppIoStub method
         res: GetNodesResponse = self._stub.GetNodes(
             GetNodesRequest(run_id=cast(Run, self._run).run_id)
         )
@@ -198,21 +205,40 @@ class GrpcDriver(Driver):
         to the node specified in `dst_node_id`.
         """
         # Construct Messages
+        run_id = cast(Run, self._run).run_id
         message_proto_list: list[ProtoMessage] = []
         for msg in messages:
+            # Populate metadata
+            msg.metadata.__dict__["_run_id"] = run_id
+            msg.metadata.__dict__["_src_node_id"] = self.node.node_id
             # Check message
             self._check_message(msg)
             # Convert to proto
             msg_proto = message_to_proto(msg)
             # Add to list
             message_proto_list.append(msg_proto)
-        # Call GrpcDriverStub method
-        res: PushInsMessagesResponse = self._stub.PushMessages(
-            PushInsMessagesRequest(
-                messages_list=message_proto_list, run_id=cast(Run, self._run).run_id
+
+        try:
+            # Call GrpcServerAppIoStub method
+            res: PushInsMessagesResponse = self._stub.PushMessages(
+                PushInsMessagesRequest(messages_list=message_proto_list, run_id=run_id)
             )
-        )
-        return list(res.message_ids)
+            if len([msg_id for msg_id in res.message_ids if msg_id]) != len(
+                message_proto_list
+            ):
+                log(
+                    WARNING,
+                    "Not all messages could be pushed to the SuperLink. The returned "
+                    "list has `None` for those messages (the order is preserved as "
+                    "passed to `push_messages`). This could be due to a malformed "
+                    "message.",
+                )
+            return list(res.message_ids)
+        except grpc.RpcError as e:
+            if e.code() == grpc.StatusCode.RESOURCE_EXHAUSTED:  # pylint: disable=E1101
+                log(ERROR, ERROR_MESSAGE_PUSH_MESSAGES_RESOURCE_EXHAUSTED)
+                return []
+            raise
 
     def pull_messages(self, message_ids: Iterable[str]) -> Iterable[Message]:
         """Pull messages based on message IDs.
@@ -220,16 +246,22 @@ class GrpcDriver(Driver):
         This method is used to collect messages from the SuperLink that correspond to a
         set of given message IDs.
         """
-        # Pull Messages
-        res: PullResMessagesResponse = self._stub.PullMessages(
-            PullResMessagesRequest(
-                message_ids=message_ids,
-                run_id=cast(Run, self._run).run_id,
+        try:
+            # Pull Messages
+            res: PullResMessagesResponse = self._stub.PullMessages(
+                PullResMessagesRequest(
+                    message_ids=message_ids,
+                    run_id=cast(Run, self._run).run_id,
+                )
             )
-        )
-        # Convert Message from Protobuf representation
-        msgs = [message_from_proto(msg_proto) for msg_proto in res.messages_list]
-        return msgs
+            # Convert Message from Protobuf representation
+            msgs = [message_from_proto(msg_proto) for msg_proto in res.messages_list]
+            return msgs
+        except grpc.RpcError as e:
+            if e.code() == grpc.StatusCode.RESOURCE_EXHAUSTED:  # pylint: disable=E1101
+                log(ERROR, ERROR_MESSAGE_PULL_MESSAGES_RESOURCE_EXHAUSTED)
+                return []
+            raise
 
     def send_and_receive(
         self,
@@ -253,7 +285,7 @@ class GrpcDriver(Driver):
             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}
+                {msg.metadata.reply_to_message_id for msg in res_msgs}
             )
             if len(msg_ids) == 0:
                 break

flwr/server/{driver/inmemory_driver.py → grid/inmemory_grid.py}

@@ -1,4 +1,4 @@
-# Copyright 2024 Flower Labs GmbH. All Rights Reserved.
+# Copyright 2025 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,36 +12,37 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 # ==============================================================================
-"""Flower in-memory Driver."""
+"""Flower in-memory Grid."""
 
 
 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 import Message, RecordDict
 from flwr.common.constant import SUPERLINK_NODE_ID
-from flwr.common.serde import message_from_taskres, message_to_taskins
+from flwr.common.logger import warn_deprecated_feature
 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
+from .grid import Grid
 
 
-class InMemoryDriver(Driver):
-    """`InMemoryDriver` class provides an interface to the ServerAppIo API.
+class InMemoryGrid(Grid):
+    """`InMemoryGrid` class provides an interface to the ServerAppIo API.
 
     Parameters
     ----------
     state_factory : StateFactory
-        A StateFactory embedding a state that this driver can interface with.
+        A StateFactory embedding a state that this grid can interface with.
     pull_interval : float (default=0.1)
         Sleep duration between calls to `pull_messages`.
     """
 
+    _deprecation_warning_logged = False
+
     def __init__(
         self,
         state_factory: LinkStateFactory,
@@ -55,11 +56,10 @@ class InMemoryDriver(Driver):
     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 == ""
+            message.metadata.message_id == ""
+            and message.metadata.reply_to_message_id == ""
             and message.metadata.ttl > 0
+            and message.metadata.delivered_at == ""
         ):
             raise ValueError(f"Invalid message: {message}")
 
@@ -77,7 +77,7 @@ class InMemoryDriver(Driver):
 
     def create_message(  # pylint: disable=too-many-arguments,R0917
         self,
-        content: RecordSet,
+        content: RecordDict,
         message_type: str,
         dst_node_id: int,
         group_id: str,
@@ -88,30 +88,17 @@ class InMemoryDriver(Driver):
         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,
+        if not InMemoryGrid._deprecation_warning_logged:
+            InMemoryGrid._deprecation_warning_logged = True
+            warn_deprecated_feature(
+                "`Driver.create_message` / `Grid.create_message` is deprecated."
+                "Use `Message` constructor instead."
             )
-        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]:
+        return Message(content, dst_node_id, message_type, ttl=ttl, group_id=group_id)
+
+    def get_node_ids(self) -> Iterable[int]:
         """Get node IDs."""
-        return list(self.state.get_nodes(cast(Run, self._run).run_id))
+        return 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.
@@ -119,18 +106,19 @@ class InMemoryDriver(Driver):
         This method takes an iterable of messages and sends each message
         to the node specified in `dst_node_id`.
         """
-        task_ids: list[str] = []
+        msg_ids: list[str] = []
         for msg in messages:
+            # Populate metadata
+            msg.metadata.__dict__["_run_id"] = cast(Run, self._run).run_id
+            msg.metadata.__dict__["_src_node_id"] = self.node.node_id
             # Check message
             self._check_message(msg)
-            # Convert Message to TaskIns
-            taskins = message_to_taskins(msg)
             # Store in state
-            task_id = self.state.store_task_ins(taskins)
-            if task_id:
-                task_ids.append(str(task_id))
+            msg_id = self.state.store_message_ins(msg)
+            if msg_id:
+                msg_ids.append(str(msg_id))
 
-        return task_ids
+        return msg_ids
 
     def pull_messages(self, message_ids: Iterable[str]) -> Iterable[Message]:
         """Pull messages based on message IDs.
@@ -139,17 +127,16 @@ class InMemoryDriver(Driver):
         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
+        # Pull Messages
+        message_res_list = self.state.get_message_res(message_ids=msg_ids)
+        # Get IDs of Messages these replies are for
+        message_ins_ids_to_delete = {
+            UUID(msg_res.metadata.reply_to_message_id) for msg_res in message_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
+        # Delete
+        self.state.delete_messages(message_ins_ids=message_ins_ids_to_delete)
+
+        return message_res_list
 
     def send_and_receive(
         self,
@@ -173,7 +160,7 @@ class InMemoryDriver(Driver):
             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}
+                {msg.metadata.reply_to_message_id for msg in res_msgs}
             )
             if len(msg_ids) == 0:
                 break

flwr/server/run_serverapp.py

@@ -15,26 +15,25 @@
 """Run ServerApp."""
 
 
-from logging import DEBUG, ERROR
+from logging import DEBUG
 from typing import Optional
 
-from flwr.common import Context, EventType, event
-from flwr.common.exit_handlers import register_exit_handlers
+from flwr.common import Context
 from flwr.common.logger import log
 from flwr.common.object_ref import load_app
 
-from .driver import Driver
+from .grid import Grid
 from .server_app import LoadServerAppError, ServerApp
 
 
 def run(
-    driver: Driver,
+    grid: Grid,
     context: Context,
     server_app_dir: str,
     server_app_attr: Optional[str] = None,
     loaded_server_app: Optional[ServerApp] = None,
 ) -> Context:
-    """Run ServerApp with a given Driver."""
+    """Run ServerApp with a given Grid."""
     if not (server_app_attr is None) ^ (loaded_server_app is None):
         raise ValueError(
             "Either `server_app_attr` or `loaded_server_app` should be set "
@@ -60,17 +59,7 @@ def run(
     server_app = _load()
 
     # Call ServerApp
-    server_app(driver=driver, context=context)
+    server_app(grid=grid, context=context)
 
     log(DEBUG, "ServerApp finished running.")
     return context
-
-
-def run_server_app() -> None:
-    """Run Flower server app."""
-    event(EventType.RUN_SERVER_APP_ENTER)
-    log(
-        ERROR,
-        "The command `flower-server-app` has been replaced by `flwr run`.",
-    )
-    register_exit_handlers(event_type=EventType.RUN_SERVER_APP_LEAVE)