break(framework): Generate error replies for unavailable nodes in Strategy.start()

framework/py/flwr/serverapp/strategy/dp_fixed_clipping.py

@@ -552,4 +552,4 @@ def validate_replies(
         if strict:
             return False
 
-    return True
+    return num_replies_with_content > 0

framework/py/flwr/serverapp/strategy/strategy.py

@@ -26,7 +26,7 @@
 
 from ..grid import Grid
 from .result import Result
-from .strategy_utils import log_strategy_start_info
+from .strategy_utils import log_strategy_start_info, send_and_receive
 
 
 class Strategy(ABC):
@@ -158,7 +158,8 @@ def start(
         num_rounds : int (default: 3)
             Maximum number of federated learning rounds to execute.
         timeout : float (default: 3600)
-            Timeout in seconds for waiting for node responses.
+            Timeout in seconds for waiting for node responses. Outstanding responses
+            from nodes that become unavailable are replaced with error replies.
         train_config : ConfigRecord, optional
             Configuration to be sent to nodes during training rounds.
             If unset, an empty ConfigRecord will be used.
@@ -215,7 +216,8 @@ def start(
 
             # Call strategy to configure training round
             # Send messages and wait for replies
-            train_replies = grid.send_and_receive(
+            train_replies = send_and_receive(
+                grid=grid,
                 messages=self.configure_train(
                     current_round,
                     arrays,
@@ -245,7 +247,8 @@ def start(
 
             # Call strategy to configure evaluation round
             # Send messages and wait for replies
-            evaluate_replies = grid.send_and_receive(
+            evaluate_replies = send_and_receive(
+                grid=grid,
                 messages=self.configure_evaluate(
                     current_round,
                     arrays,

framework/py/flwr/serverapp/strategy/strategy_utils.py

@@ -17,14 +17,25 @@
 
 import json
 import random
+from collections.abc import Iterable
 from logging import INFO
-from time import sleep
+from time import monotonic, sleep
 from typing import cast
 
 import numpy as np
 
-from flwr.app import Array, ArrayRecord, ConfigRecord, MetricRecord, RecordDict
+from flwr.app import (
+    Array,
+    ArrayRecord,
+    ConfigRecord,
+    Error,
+    Message,
+    MetricRecord,
+    RecordDict,
+)
+from flwr.app.message import remove_content_from_message
 from flwr.common import NDArray
+from flwr.common.constant import ErrorCode
 from flwr.supercore import log
 
 from ..exception import InconsistentMessageReplies
@@ -185,6 +196,61 @@ def sample_nodes(
     return sampled_nodes, all_nodes
 
 
+def send_and_receive(
+    grid: Grid, messages: Iterable[Message], timeout: float
+) -> list[Message]:
+    """Send messages and collect replies, generating errors for unavailable nodes.
+
+    Parameters
+    ----------
+    grid : Grid
+        The grid used to send messages, pull replies, and check node availability.
+    messages : Iterable[Message]
+        Messages to send to the selected nodes.
+    timeout : float
+        Maximum time in seconds to wait for replies after sending the messages.
+
+    Returns
+    -------
+    list[Message]
+        Received replies and generated error replies for unavailable nodes.
+    """
+    messages = list(messages)
+    sent_ids = {msg_id for msg_id in grid.push_messages(messages) if msg_id}
+    pending = {
+        msg.metadata.message_id: remove_content_from_message(msg)
+        for msg in messages
+        if msg.metadata.message_id in sent_ids
+    }
+    del messages
+
+    deadline = monotonic() + timeout
+    replies: list[Message] = []
+    while pending and monotonic() < deadline:
+        pulled = list(grid.pull_messages(pending))
+        replies.extend(pulled)
+        for reply in pulled:
+            pending.pop(reply.metadata.reply_to_message_id, None)
+        if not pending:
+            break
+
+        available_nodes = set(grid.get_node_ids())
+        for msg_id, msg in list(pending.items()):
+            if msg.metadata.dst_node_id not in available_nodes:
+                error = Error(
+                    ErrorCode.NODE_UNAVAILABLE, "The destination node is unavailable."
+                )
+                replies.append(Message(error, reply_to=msg))
+                del pending[msg_id]
+
+        remaining = deadline - monotonic()
+        if not pending or remaining <= 0:
+            break
+        sleep(min(1.0, remaining))
+
+    return replies
+
+
 # pylint: disable=too-many-return-statements
 def validate_message_reply_consistency(
     replies: list[RecordDict], weighted_by_key: str, check_arrayrecord: bool

framework/py/flwr/serverapp/strategy/strategy_utils_test.py

@@ -22,12 +22,16 @@
 from parameterized import parameterized
 
 from flwr.app import Array, ArrayRecord, ConfigRecord, Message, MetricRecord, RecordDict
+from flwr.common.constant import ErrorCode
 from flwr.serverapp.exception import InconsistentMessageReplies
+from flwr.serverapp.grid import Grid
+from flwr.supercore.task_identity import TaskIdentity
 
 from .strategy_utils import (
     aggregate_arrayrecords,
     aggregate_metricrecords,
     config_to_str,
+    send_and_receive,
     validate_message_reply_consistency,
 )
 
@@ -50,6 +54,50 @@ def test_config_to_str() -> None:
     assert config_to_str(config) == expected_str
 
 
+@pytest.mark.parametrize("available_nodes", [[], [2]])
+def test_send_and_receive_with_missing_reply(
+    monkeypatch: pytest.MonkeyPatch, available_nodes: list[int]
+) -> None:
+    """Keep replies, generating errors for unavailable nodes or timing out otherwise."""
+    monkeypatch.setattr(TaskIdentity, "_task_id", 123)
+    monkeypatch.setattr(TaskIdentity, "_run_id", 456)
+    monkeypatch.setattr(TaskIdentity, "_node_id", 789)
+    messages: list[Message] = [
+        Message(RecordDict(), dst_node_id=node_id, message_type="train")
+        for node_id in [1, 2, 2]
+    ]
+    for i, msg in enumerate(messages):
+        msg.metadata.__dict__["_message_id"] = str(i)
+    reply = Message(RecordDict(), reply_to=messages[0])
+    grid = Mock(spec=Grid)
+    grid.push_messages.return_value = [msg.metadata.message_id for msg in messages]
+    grid.pull_messages.return_value = [reply]
+    grid.get_node_ids.return_value = available_nodes
+    sleep_mock = Mock()
+    monkeypatch.setattr("flwr.serverapp.strategy.strategy_utils.sleep", sleep_mock)
+    monkeypatch.setattr(
+        "flwr.serverapp.strategy.strategy_utils.monotonic",
+        Mock(side_effect=[0.0, 0.0, 0.5, 1.0]),
+    )
+
+    replies = send_and_receive(grid, iter(messages), timeout=1.0)
+    assert replies[0] is reply
+    grid.push_messages.assert_called_once_with(messages)
+    grid.pull_messages.assert_called_once()
+    if available_nodes:
+        assert replies == [reply]
+        sleep_mock.assert_called_once_with(0.5)
+    else:
+        assert len(replies) == len(messages)
+        for error_reply, msg in zip(replies[1:], messages[1:], strict=True):
+            assert error_reply.has_error()
+            assert error_reply.error.code == ErrorCode.NODE_UNAVAILABLE
+            assert error_reply.error.reason == "The destination node is unavailable."
+            assert error_reply.metadata.reply_to_message_id == msg.metadata.message_id
+            assert error_reply.metadata.src_node_id == msg.metadata.dst_node_id
+        sleep_mock.assert_not_called()
+
+
 def test_arrayrecords_aggregation() -> None:
     """Test aggregation of ArrayRecords."""
     num_replies = 3