fix(framework): Bound task log uploader shutdown

framework/py/flwr/common/logger.py

@@ -37,7 +37,7 @@
 from flwr.proto.node_pb2 import Node  # pylint: disable=E0611
 from flwr.proto.runtime_pb2_grpc import RuntimeStub  # pylint: disable=E0611
 
-from .constant import LOG_UPLOAD_INTERVAL
+from .constant import LOG_UPLOAD_INTERVAL, TASK_WORKER_CALL_TIMEOUT
 
 # Create logger
 LOGGER_NAME = "flwr"
@@ -408,15 +408,20 @@ def _log_uploader(
                 logs=msgs,
             )
             try:
-                stub.PushLogs(req)
+                stub.PushLogs(req, timeout=TASK_WORKER_CALL_TIMEOUT)
                 msgs.clear()
             except grpc.RpcError as e:
-                # Ignore minor network errors
+                # Ignore minor network errors. A deadline leaves delivery ambiguous:
+                # the server may already have appended this non-idempotent batch, so
+                # discard it instead of risking duplicate logs on a retry.
                 # pylint: disable-next=no-member
-                if e.code() != grpc.StatusCode.UNAVAILABLE:
-                    raise e
+                status_code = e.code()
+                if status_code == grpc.StatusCode.DEADLINE_EXCEEDED:
+                    msgs.clear()
+                elif status_code != grpc.StatusCode.UNAVAILABLE:
+                    raise
 
-        if exit_flag:
+        if exit_flag and not msgs:
             break
 
         time.sleep(LOG_UPLOAD_INTERVAL)

framework/py/flwr/common/logger_test.py

@@ -21,17 +21,41 @@
 from logging.handlers import TimedRotatingFileHandler
 from pathlib import Path
 from queue import Queue
+from unittest.mock import Mock
 
+import grpc
+
+from flwr.proto.log_pb2 import PushLogsRequest  # pylint: disable=E0611
+
+from .constant import TASK_WORKER_CALL_TIMEOUT
 from .logger import (
     FLOWER_LOGGER,
     configure_superlink_log_file,
     console_handler,
     flush_logs,
     mirror_output_to_queue,
     restore_output,
+    start_log_uploader,
+    stop_log_uploader,
 )
 
 
+class _DeadlineExceededError(grpc.RpcError):  # type: ignore[misc]
+    """gRPC error reporting an expired call deadline."""
+
+    def code(self) -> grpc.StatusCode:
+        """Return the gRPC status code."""
+        return grpc.StatusCode.DEADLINE_EXCEEDED
+
+
+class _UnavailableError(grpc.RpcError):  # type: ignore[misc]
+    """gRPC error reporting an unavailable endpoint."""
+
+    def code(self) -> grpc.StatusCode:
+        """Return the gRPC status code."""
+        return grpc.StatusCode.UNAVAILABLE
+
+
 def test_mirror_output_to_queue() -> None:
     """Test that stdout and stderr are mirrored to the provided queue."""
     # Prepare
@@ -116,6 +140,80 @@ def test_flush_logs_returns_false_when_queue_does_not_drain() -> None:
     assert not log_queue.empty()
 
 
+def test_log_uploader_uses_bounded_rpc() -> None:
+    """Task log uploads must not block executor shutdown indefinitely."""
+    log_queue: Queue[str | None] = Queue()
+    log_queue.put("Test message")
+    stub = Mock()
+
+    uploader = start_log_uploader(log_queue, node_id=1, run_id=2, stub=stub)
+    stop_log_uploader(log_queue, uploader, timeout=1.0)
+
+    assert not uploader.is_alive()
+    assert stub.PushLogs.call_args.kwargs["timeout"] == TASK_WORKER_CALL_TIMEOUT
+
+
+def test_log_uploader_does_not_retry_after_deadline_expiry() -> None:
+    """An ambiguously delivered timed-out batch must not be duplicated."""
+    log_queue: Queue[str | None] = Queue()
+    log_queue.put("Timed-out message")
+    first_attempt_finished = threading.Event()
+    upload_succeeded = threading.Event()
+    requests: list[PushLogsRequest] = []
+
+    def push_logs(request: PushLogsRequest, **_kwargs: object) -> None:
+        requests.append(request)
+        if len(requests) == 1:
+            first_attempt_finished.set()
+            raise _DeadlineExceededError
+        upload_succeeded.set()
+
+    stub = Mock()
+    stub.PushLogs.side_effect = push_logs
+    uploader = start_log_uploader(log_queue, node_id=1, run_id=2, stub=stub)
+
+    assert first_attempt_finished.wait(timeout=1.0)
+    log_queue.put("Next message")
+    assert upload_succeeded.wait(timeout=2.0)
+    stop_log_uploader(log_queue, uploader, timeout=1.0)
+
+    assert not uploader.is_alive()
+    assert len(requests) == 2
+    assert list(requests[0].logs) == ["Timed-out message"]
+    assert list(requests[1].logs) == ["Next message"]
+
+
+def test_log_uploader_retries_retained_batch_during_stop() -> None:
+    """The stop sentinel must not discard a retained final log batch."""
+    log_queue: Queue[str | None] = Queue()
+    log_queue.put("Final message")
+    first_attempt_started = threading.Event()
+    release_first_attempt = threading.Event()
+    requests: list[PushLogsRequest] = []
+
+    def push_logs(request: PushLogsRequest, **_kwargs: object) -> None:
+        requests.append(request)
+        if len(requests) == 1:
+            first_attempt_started.set()
+            assert release_first_attempt.wait(timeout=1.0)
+        if len(requests) < 3:
+            raise _UnavailableError
+
+    stub = Mock()
+    stub.PushLogs.side_effect = push_logs
+    uploader = start_log_uploader(log_queue, node_id=1, run_id=2, stub=stub)
+    assert first_attempt_started.wait(timeout=1.0)
+
+    # Queue the stop sentinel while the first upload still owns the retained batch.
+    log_queue.put(None)
+    release_first_attempt.set()
+    uploader.join(timeout=2.0)
+
+    assert not uploader.is_alive()
+    assert len(requests) == 3
+    assert requests[0] == requests[1] == requests[2]
+
+
 def test_configure_superlink_log_file(tmp_path: Path) -> None:
     """Test configuring timed file rotation for SuperLink logs."""
     # Prepare