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

flwr/server/compat/{driver_client_proxy.py → grid_client_proxy.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,26 +12,26 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 # ==============================================================================
-"""Flower ClientProxy implementation for Driver API."""
+"""Flower ClientProxy implementation using Grid."""
 
 
 from typing import Optional
 
 from flwr import common
-from flwr.common import Message, MessageType, MessageTypeLegacy, RecordSet
-from flwr.common import recordset_compat as compat
+from flwr.common import Message, MessageType, MessageTypeLegacy, RecordDict
+from flwr.common import recorddict_compat as compat
 from flwr.server.client_proxy import ClientProxy
 
-from ..driver.driver import Driver
+from ..grid.grid import Grid
 
 
-class DriverClientProxy(ClientProxy):
-    """Flower client proxy which delegates work using the Driver API."""
+class GridClientProxy(ClientProxy):
+    """Flower client proxy which delegates work using Grid."""
 
-    def __init__(self, node_id: int, driver: Driver, run_id: int):
+    def __init__(self, node_id: int, grid: Grid, run_id: int):
         super().__init__(str(node_id))
         self.node_id = node_id
-        self.driver = driver
+        self.grid = grid
         self.run_id = run_id
 
     def get_properties(
@@ -41,14 +41,14 @@ class DriverClientProxy(ClientProxy):
         group_id: Optional[int],
     ) -> common.GetPropertiesRes:
         """Return client's properties."""
-        # Ins to RecordSet
-        out_recordset = compat.getpropertiesins_to_recordset(ins)
+        # Ins to RecordDict
+        out_recorddict = compat.getpropertiesins_to_recorddict(ins)
         # Fetch response
-        in_recordset = self._send_receive_recordset(
-            out_recordset, MessageTypeLegacy.GET_PROPERTIES, timeout, group_id
+        in_recorddict = self._send_receive_recorddict(
+            out_recorddict, MessageTypeLegacy.GET_PROPERTIES, timeout, group_id
         )
-        # RecordSet to Res
-        return compat.recordset_to_getpropertiesres(in_recordset)
+        # RecordDict to Res
+        return compat.recorddict_to_getpropertiesres(in_recorddict)
 
     def get_parameters(
         self,
@@ -57,40 +57,40 @@ class DriverClientProxy(ClientProxy):
         group_id: Optional[int],
     ) -> common.GetParametersRes:
         """Return the current local model parameters."""
-        # Ins to RecordSet
-        out_recordset = compat.getparametersins_to_recordset(ins)
+        # Ins to RecordDict
+        out_recorddict = compat.getparametersins_to_recorddict(ins)
         # Fetch response
-        in_recordset = self._send_receive_recordset(
-            out_recordset, MessageTypeLegacy.GET_PARAMETERS, timeout, group_id
+        in_recorddict = self._send_receive_recorddict(
+            out_recorddict, MessageTypeLegacy.GET_PARAMETERS, timeout, group_id
         )
-        # RecordSet to Res
-        return compat.recordset_to_getparametersres(in_recordset, False)
+        # RecordDict to Res
+        return compat.recorddict_to_getparametersres(in_recorddict, False)
 
     def fit(
         self, ins: common.FitIns, timeout: Optional[float], group_id: Optional[int]
     ) -> common.FitRes:
         """Train model parameters on the locally held dataset."""
-        # Ins to RecordSet
-        out_recordset = compat.fitins_to_recordset(ins, keep_input=True)
+        # Ins to RecordDict
+        out_recorddict = compat.fitins_to_recorddict(ins, keep_input=True)
         # Fetch response
-        in_recordset = self._send_receive_recordset(
-            out_recordset, MessageType.TRAIN, timeout, group_id
+        in_recorddict = self._send_receive_recorddict(
+            out_recorddict, MessageType.TRAIN, timeout, group_id
         )
-        # RecordSet to Res
-        return compat.recordset_to_fitres(in_recordset, keep_input=False)
+        # RecordDict to Res
+        return compat.recorddict_to_fitres(in_recorddict, keep_input=False)
 
     def evaluate(
         self, ins: common.EvaluateIns, timeout: Optional[float], group_id: Optional[int]
     ) -> common.EvaluateRes:
         """Evaluate model parameters on the locally held dataset."""
-        # Ins to RecordSet
-        out_recordset = compat.evaluateins_to_recordset(ins, keep_input=True)
+        # Ins to RecordDict
+        out_recorddict = compat.evaluateins_to_recorddict(ins, keep_input=True)
         # Fetch response
-        in_recordset = self._send_receive_recordset(
-            out_recordset, MessageType.EVALUATE, timeout, group_id
+        in_recorddict = self._send_receive_recorddict(
+            out_recorddict, MessageType.EVALUATE, timeout, group_id
         )
-        # RecordSet to Res
-        return compat.recordset_to_evaluateres(in_recordset)
+        # RecordDict to Res
+        return compat.recorddict_to_evaluateres(in_recorddict)
 
     def reconnect(
         self,
@@ -101,17 +101,17 @@ class DriverClientProxy(ClientProxy):
         """Disconnect and (optionally) reconnect later."""
         return common.DisconnectRes(reason="")  # Nothing to do here (yet)
 
-    def _send_receive_recordset(
+    def _send_receive_recorddict(
         self,
-        recordset: RecordSet,
+        recorddict: RecordDict,
         message_type: str,
         timeout: Optional[float],
         group_id: Optional[int],
-    ) -> RecordSet:
+    ) -> RecordDict:
 
         # Create message
-        message = self.driver.create_message(
-            content=recordset,
+        message = Message(
+            content=recorddict,
             message_type=message_type,
             dst_node_id=self.node_id,
             group_id=str(group_id) if group_id else "",
@@ -119,7 +119,7 @@ class DriverClientProxy(ClientProxy):
         )
 
         # Send message and wait for reply
-        messages = list(self.driver.send_and_receive(messages=[message]))
+        messages = list(self.grid.send_and_receive(messages=[message]))
 
         # A single reply is expected
         if len(messages) != 1:

flwr/server/fleet_event_log_interceptor.py

@@ -0,0 +1,94 @@
+# 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.
+# 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 Fleet API event log interceptor."""
+
+
+from typing import Any, Callable, cast
+
+import grpc
+from google.protobuf.message import Message as GrpcMessage
+
+from flwr.common.event_log_plugin.event_log_plugin import EventLogWriterPlugin
+from flwr.common.typing import LogEntry
+
+
+class FleetEventLogInterceptor(grpc.ServerInterceptor):  # type: ignore
+    """Fleet API interceptor for logging events."""
+
+    def __init__(self, log_plugin: EventLogWriterPlugin) -> None:
+        self.log_plugin = log_plugin
+
+    def intercept_service(
+        self,
+        continuation: Callable[[Any], Any],
+        handler_call_details: grpc.HandlerCallDetails,
+    ) -> grpc.RpcMethodHandler:
+        """Flower Fleet API server interceptor logging logic.
+
+        Intercept all unary-unary calls from users and log the event. Continue RPC call
+        if event logger is enabled on the SuperLink, else, terminate RPC call by setting
+        context to abort.
+        """
+        # One of the method handlers in
+        # `flwr.server.superlink.fleet.grpc_rere.fleet_servicer.FleetServicer`
+        method_handler: grpc.RpcMethodHandler = continuation(handler_call_details)
+        method_name: str = handler_call_details.method
+        return self._generic_event_log_unary_method_handler(method_handler, method_name)
+
+    def _generic_event_log_unary_method_handler(
+        self, method_handler: grpc.RpcMethodHandler, method_name: str
+    ) -> grpc.RpcMethodHandler:
+        def _generic_method_handler(
+            request: GrpcMessage,
+            context: grpc.ServicerContext,
+        ) -> GrpcMessage:
+            log_entry: LogEntry
+            # Log before call
+            log_entry = self.log_plugin.compose_log_before_event(
+                request=request,
+                context=context,
+                user_info=None,
+                method_name=method_name,
+            )
+            self.log_plugin.write_log(log_entry)
+
+            call = method_handler.unary_unary
+            unary_response, error = None, None
+            try:
+                unary_response = cast(GrpcMessage, call(request, context))
+            except BaseException as e:
+                error = e
+                raise
+            finally:
+                log_entry = self.log_plugin.compose_log_after_event(
+                    request=request,
+                    context=context,
+                    user_info=None,
+                    method_name=method_name,
+                    response=unary_response or error,
+                )
+                self.log_plugin.write_log(log_entry)
+            return unary_response
+
+        if method_handler.unary_unary:
+            message_handler = grpc.unary_unary_rpc_method_handler
+        else:
+            # If the method type is not `unary_unary` raise an error
+            raise NotImplementedError("This RPC method type is not supported.")
+        return message_handler(
+            _generic_method_handler,
+            request_deserializer=method_handler.request_deserializer,
+            response_serializer=method_handler.response_serializer,
+        )

flwr/server/{driver → grid}/__init__.py

@@ -1,4 +1,4 @@
-# Copyright 2022 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,15 +12,16 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 # ==============================================================================
-"""Flower driver SDK."""
+"""Flower grid SDK."""
 
 
-from .driver import Driver
-from .grpc_driver import GrpcDriver
-from .inmemory_driver import InMemoryDriver
+from .grid import Driver, Grid
+from .grpc_grid import GrpcGrid
+from .inmemory_grid import InMemoryGrid
 
 __all__ = [
     "Driver",
-    "GrpcDriver",
-    "InMemoryDriver",
+    "Grid",
+    "GrpcGrid",
+    "InMemoryGrid",
 ]

flwr/server/{driver/driver.py → grid/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,32 +12,32 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 # ==============================================================================
-"""Driver (abstract base class)."""
+"""Grid (abstract base class)."""
 
 
 from abc import ABC, abstractmethod
 from collections.abc import Iterable
 from typing import Optional
 
-from flwr.common import Message, RecordSet
+from flwr.common import Message, RecordDict
 from flwr.common.typing import Run
 
 
-class Driver(ABC):
-    """Abstract base Driver class for the ServerAppIo API."""
+class Grid(ABC):
+    """Abstract base class Grid to send/receive messages."""
 
     @abstractmethod
     def set_run(self, run_id: int) -> None:
         """Request a run to the SuperLink with a given `run_id`.
 
-        If a Run with the specified `run_id` exists, a local Run
+        If a ``Run`` with the specified ``run_id`` exists, a local ``Run``
         object will be created. It enables further functionality
-        in the driver, such as sending `Messages`.
+        in the grid, such as sending ``Message``s.
 
         Parameters
         ----------
         run_id : int
-            The `run_id` of the Run this Driver object operates in.
+            The ``run_id`` of the ``Run`` this ``Grid`` object operates in.
         """
 
     @property
@@ -48,7 +48,7 @@ class Driver(ABC):
     @abstractmethod
     def create_message(  # pylint: disable=too-many-arguments,R0917
         self,
-        content: RecordSet,
+        content: RecordDict,
         message_type: str,
         dst_node_id: int,
         group_id: str,
@@ -56,12 +56,12 @@ class Driver(ABC):
     ) -> 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.
+        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
+        content : RecordDict
             The content for the new message. This holds records that are to be sent
             to the destination node.
         message_type : str
@@ -71,12 +71,12 @@ class Driver(ABC):
             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.
+            this is used as the federated learning 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.
+            the default TTL (i.e., ``common.DEFAULT_TTL``) will be used.
 
         Returns
         -------
@@ -93,7 +93,7 @@ class Driver(ABC):
         """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`.
+        to the node specified in ``dst_node_id``.
 
         Parameters
         ----------
@@ -154,8 +154,37 @@ class Driver(ABC):
 
         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
+        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`.
+        which is not affected by ``timeout``.
         """
+
+
+class Driver(Grid):
+    """Deprecated abstract base class ``Driver``, use ``Grid`` instead.
+
+    This class is provided solely for backward compatibility with legacy
+    code that previously relied on the ``Driver`` class. It has been deprecated
+    in favor of the updated abstract base class ``Grid``, which now encompasses
+    all communication-related functionality and improvements between the
+    ServerApp and the SuperLink.
+
+    .. warning::
+        ``Driver`` is deprecated and will be removed in a future release.
+        Use `Grid` in the signature of your ServerApp.
+
+    Examples
+    --------
+    Legacy (deprecated) usage::
+
+        @app.main()
+        def main(driver: Driver, context: Context) -> None:
+            ...
+
+    Updated usage::
+
+        @app.main()
+        def main(grid: Grid, context: Context) -> None:
+            ...
+    """

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."
             )
-
-        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)
+        return Message(content, dst_node_id, message_type, ttl=ttl, group_id=group_id)
 
     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,30 +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
-            )
-        )
-        if len([msg_id for msg_id in res.message_ids if msg_id]) != len(
-            list(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.",
+
+        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.
@@ -229,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,
@@ -262,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