From ea6a5f48656806507811e286e3f34be1fe5a1913 Mon Sep 17 00:00:00 2001 From: Taner Topal Date: Tue, 18 Aug 2026 17:47:49 +0200 Subject: [PATCH] fix(framework): Bound task log uploader shutdown --- framework/py/flwr/common/logger.py | 17 +++-- framework/py/flwr/common/logger_test.py | 98 +++++++++++++++++++++++++ 2 files changed, 109 insertions(+), 6 deletions(-) diff --git a/framework/py/flwr/common/logger.py b/framework/py/flwr/common/logger.py index 0566ec0c844c..2cbeef6b6a90 100644 --- a/framework/py/flwr/common/logger.py +++ b/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) diff --git a/framework/py/flwr/common/logger_test.py b/framework/py/flwr/common/logger_test.py index 8ef38f97280a..c87fa2d926d0 100644 --- a/framework/py/flwr/common/logger_test.py +++ b/framework/py/flwr/common/logger_test.py @@ -21,7 +21,13 @@ 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, @@ -29,9 +35,27 @@ 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