From d9fbcfbb6661e0bd88ef1f64d19eb4e266786e56 Mon Sep 17 00:00:00 2001 From: LIU ZHE YOU Date: Thu, 23 Jul 2026 08:13:37 +0000 Subject: [PATCH 1/2] Skip log accumulator disk spill when reading finished task logs --- .../airflow/utils/log/file_task_handler.py | 12 +++++++ .../tests/unit/utils/log/test_log_reader.py | 4 +-- .../tests/unit/utils/test_log_handlers.py | 35 ++++++++++++++++--- 3 files changed, 45 insertions(+), 6 deletions(-) diff --git a/airflow-core/src/airflow/utils/log/file_task_handler.py b/airflow-core/src/airflow/utils/log/file_task_handler.py index d22ce48fa0897..f017cd0baa09c 100644 --- a/airflow-core/src/airflow/utils/log/file_task_handler.py +++ b/airflow-core/src/airflow/utils/log/file_task_handler.py @@ -683,6 +683,18 @@ def _read( TaskInstanceState.DEFERRED, ) + if end_of_log: + # No continuation token will be issued, so the total line count is not + # needed: stream straight through instead of paying the accumulator's + # eager full drain to memory/disk. Memory stays bounded by the merge + # heap in _interleave_logs either way. + if metadata and "log_pos" in metadata: + out_stream = islice(out_stream, metadata["log_pos"], None) + else: + # first time reading log, add messages before interleaved log stream + out_stream = chain(header, out_stream) + return out_stream, {"end_of_log": end_of_log} + with LogStreamAccumulator(out_stream, HEAP_DUMP_SIZE) as stream_accumulator: log_pos = stream_accumulator.total_lines out_stream = stream_accumulator.stream diff --git a/airflow-core/tests/unit/utils/log/test_log_reader.py b/airflow-core/tests/unit/utils/log/test_log_reader.py index 1c65cb1806b14..01ba3aa6b566f 100644 --- a/airflow-core/tests/unit/utils/log/test_log_reader.py +++ b/airflow-core/tests/unit/utils/log/test_log_reader.py @@ -136,7 +136,7 @@ def test_test_read_log_chunks_should_read_one_try(self): ) assert logs[2].event == "::endgroup::" assert logs[3].event == "try_number=1." - assert metadata == {"end_of_log": True, "log_pos": 1} + assert metadata == {"end_of_log": True} def test_test_read_log_chunks_should_read_latest_files(self): task_log_reader = TaskLogReader() @@ -151,7 +151,7 @@ def test_test_read_log_chunks_should_read_latest_files(self): ) assert logs[2].event == "::endgroup::" assert logs[3].event == f"try_number={ti.try_number}." - assert metadata == {"end_of_log": True, "log_pos": 1} + assert metadata == {"end_of_log": True} def test_test_test_read_log_stream_should_read_one_try(self): task_log_reader = TaskLogReader() diff --git a/airflow-core/tests/unit/utils/test_log_handlers.py b/airflow-core/tests/unit/utils/test_log_handlers.py index 89cde6f32ff7e..01fdb30d7a741 100644 --- a/airflow-core/tests/unit/utils/test_log_handlers.py +++ b/airflow-core/tests/unit/utils/test_log_handlers.py @@ -474,7 +474,7 @@ def test__read_when_local(self, mock_read_local, create_task_instance): log_handler_output_stream, metadata = fth._read(ti=local_log_file_read, try_number=1) mock_read_local.assert_called_with(path) assert extract_events(log_handler_output_stream) == ["the log"] - assert metadata == {"end_of_log": True, "log_pos": 1} + assert metadata == {"end_of_log": True} @patch("airflow.utils.log.file_task_handler.FileTaskHandler._read_from_local") def test__read_when_local_respects_log_pos_metadata(self, mock_read_local, create_task_instance): @@ -503,7 +503,34 @@ def test__read_when_local_respects_log_pos_metadata(self, mock_read_local, creat # Should resume from the third line only. assert extract_events(log_handler_output_stream) == ["line 3"] - assert metadata == {"end_of_log": True, "log_pos": 3} + assert metadata == {"end_of_log": True} + + @pytest.mark.parametrize("ti_state", [TaskInstanceState.SUCCESS, TaskInstanceState.FAILED]) + @patch("airflow.utils.log.file_task_handler.LogStreamAccumulator") + @patch("airflow.utils.log.file_task_handler.FileTaskHandler._read_from_local") + def test__read_end_of_log_skips_accumulator( + self, mock_read_local, mock_accumulator, ti_state, create_task_instance + ): + """A terminal-state read issues no continuation token, so the stream must pass + through without the accumulator's eager full drain to memory/disk.""" + mock_read_local.return_value = ( + ["the messages"], + [convert_list_to_stream(["line 1", "line 2", "line 3"])], + ) + ti = create_task_instance( + dag_id="dag_for_testing_local_log_read", + task_id="task_for_testing_local_log_read", + run_type=DagRunType.SCHEDULED, + logical_date=DEFAULT_DATE, + state=ti_state, + ) + fth = FileTaskHandler("") + + log_handler_output_stream, metadata = fth._read(ti=ti, try_number=1) + + assert extract_events(log_handler_output_stream) == ["line 1", "line 2", "line 3"] + assert metadata == {"end_of_log": True} + mock_accumulator.assert_not_called() @patch("airflow.utils.log.file_task_handler.FileTaskHandler._read_from_local") def test_read_respects_log_pos_metadata(self, mock_read_local, create_task_instance): @@ -528,7 +555,7 @@ def test_read_respects_log_pos_metadata(self, mock_read_local, create_task_insta # Should resume from the third line only. assert extract_events(log_handler_output_stream) == ["line 3"] - assert metadata == {"end_of_log": True, "log_pos": 3} + assert metadata == {"end_of_log": True} def test__read_from_local(self, tmp_path): """Tests the behavior of method _read_from_local""" @@ -623,7 +650,7 @@ def test__read_served_logs_checked_when_done_and_no_local_or_remote_logs( else: fth._read_from_logs_server.assert_not_called() assert extract_events(logs, False) == expected_logs - assert metadata == {"end_of_log": True, "log_pos": 3} + assert metadata == {"end_of_log": True} @pytest.mark.parametrize("is_tih", [False, True]) def test_read_served_logs(self, is_tih, create_task_instance): From b0d23766ca87bb34f0089b608528767d25b96c5a Mon Sep 17 00:00:00 2001 From: LIU ZHE YOU Date: Fri, 31 Jul 2026 11:32:57 +0000 Subject: [PATCH 2/2] Version-gate remote log handler tests for terminal-state metadata Finished-task reads drop the log_pos continuation token from 3.4.0 onward, but the provider test matrix still runs against released cores (3.0-3.3) that emit it. Gate the S3, GCS, and WASB task-handler metadata assertions on AIRFLOW_V_3_4_0_PLUS so both the source run and the older-core compatibility runs pass. --- .../tests_common/test_utils/version_compat.py | 2 ++ .../unit/amazon/aws/log/test_s3_task_handler.py | 16 +++++++++++++--- .../google/cloud/log/test_gcs_task_handler.py | 16 +++++++++++++--- .../azure/log/test_wasb_task_handler.py | 11 +++++++++-- 4 files changed, 37 insertions(+), 8 deletions(-) diff --git a/devel-common/src/tests_common/test_utils/version_compat.py b/devel-common/src/tests_common/test_utils/version_compat.py index 7eb25dec2b3cb..904c57160464c 100644 --- a/devel-common/src/tests_common/test_utils/version_compat.py +++ b/devel-common/src/tests_common/test_utils/version_compat.py @@ -42,6 +42,7 @@ def get_base_airflow_version_tuple() -> tuple[int, int, int]: AIRFLOW_V_3_2_PLUS = get_base_airflow_version_tuple() >= (3, 2, 0) AIRFLOW_V_3_2_2_PLUS = get_base_airflow_version_tuple() >= (3, 2, 2) AIRFLOW_V_3_3_PLUS = get_base_airflow_version_tuple() >= (3, 3, 0) +AIRFLOW_V_3_4_0_PLUS = get_base_airflow_version_tuple() >= (3, 4, 0) if AIRFLOW_V_3_1_PLUS: from airflow.sdk import PokeReturnValue, timezone @@ -64,6 +65,7 @@ def get_base_airflow_version_tuple() -> tuple[int, int, int]: "AIRFLOW_V_3_1_PLUS", "AIRFLOW_V_3_2_PLUS", "AIRFLOW_V_3_3_PLUS", + "AIRFLOW_V_3_4_0_PLUS", "NOTSET", "XCOM_RETURN_KEY", "ArgNotSet", diff --git a/providers/amazon/tests/unit/amazon/aws/log/test_s3_task_handler.py b/providers/amazon/tests/unit/amazon/aws/log/test_s3_task_handler.py index 0576756c1965f..68a902601a778 100644 --- a/providers/amazon/tests/unit/amazon/aws/log/test_s3_task_handler.py +++ b/providers/amazon/tests/unit/amazon/aws/log/test_s3_task_handler.py @@ -39,7 +39,11 @@ from tests_common.test_utils.dag import sync_dag_to_db from tests_common.test_utils.db import clear_db_dag_bundles, clear_db_dags, clear_db_runs from tests_common.test_utils.taskinstance import create_task_instance -from tests_common.test_utils.version_compat import AIRFLOW_V_3_0_PLUS, AIRFLOW_V_3_2_2_PLUS +from tests_common.test_utils.version_compat import ( + AIRFLOW_V_3_0_PLUS, + AIRFLOW_V_3_2_2_PLUS, + AIRFLOW_V_3_4_0_PLUS, +) try: from airflow.sdk.timezone import datetime @@ -418,7 +422,10 @@ def test_read(self): assert log[5].event == "Line 2" assert log[6].event == "Log line 3" assert log[7].event == "Line 4" - assert metadata == {"end_of_log": True, "log_pos": 4} + if AIRFLOW_V_3_4_0_PLUS: + assert metadata == {"end_of_log": True} + else: + assert metadata == {"end_of_log": True, "log_pos": 4} elif AIRFLOW_V_3_0_PLUS: log = list(log) assert log[0].event == "::group::Log message source details" @@ -443,7 +450,10 @@ def test_read_when_s3_log_missing(self): if AIRFLOW_V_3_0_PLUS: log = list(log) assert len(log) == 2 - assert metadata == {"end_of_log": True, "log_pos": 0} + if AIRFLOW_V_3_4_0_PLUS: + assert metadata == {"end_of_log": True} + else: + assert metadata == {"end_of_log": True, "log_pos": 0} else: assert len(log) == 1 assert len(log) == len(metadata) diff --git a/providers/google/tests/unit/google/cloud/log/test_gcs_task_handler.py b/providers/google/tests/unit/google/cloud/log/test_gcs_task_handler.py index 08e3a45399022..bf7261d7182dd 100644 --- a/providers/google/tests/unit/google/cloud/log/test_gcs_task_handler.py +++ b/providers/google/tests/unit/google/cloud/log/test_gcs_task_handler.py @@ -35,7 +35,11 @@ from tests_common.test_utils.config import conf_vars from tests_common.test_utils.db import clear_db_dags, clear_db_runs -from tests_common.test_utils.version_compat import AIRFLOW_V_3_0_PLUS, AIRFLOW_V_3_2_2_PLUS +from tests_common.test_utils.version_compat import ( + AIRFLOW_V_3_0_PLUS, + AIRFLOW_V_3_2_2_PLUS, + AIRFLOW_V_3_4_0_PLUS, +) if TYPE_CHECKING: from pathlib import Path @@ -465,7 +469,10 @@ def test_should_read_logs_from_remote( assert logs[1].event == expected_gs_uri assert logs[2].event == "::endgroup::" assert logs[3].event == "CONTENT" - assert metadata == {"end_of_log": True, "log_pos": 1} + if AIRFLOW_V_3_4_0_PLUS: + assert metadata == {"end_of_log": True} + else: + assert metadata == {"end_of_log": True, "log_pos": 1} elif AIRFLOW_V_3_0_PLUS: logs = list(logs) assert logs[0].event == "::group::Log message source details" @@ -503,7 +510,10 @@ def test_should_read_from_local_on_logs_read_error(self, mock_blob, mock_client, assert log[1].event == expected_gs_uri assert log[2].event == f"{self.gcs_task_handler.local_base}/1.log" assert log[3].event == "::endgroup::" - assert metadata == {"end_of_log": True, "log_pos": 0} + if AIRFLOW_V_3_4_0_PLUS: + assert metadata == {"end_of_log": True} + else: + assert metadata == {"end_of_log": True, "log_pos": 0} elif AIRFLOW_V_3_0_PLUS: log = list(log) assert log[0].event == "::group::Log message source details" diff --git a/providers/microsoft/azure/tests/unit/microsoft/azure/log/test_wasb_task_handler.py b/providers/microsoft/azure/tests/unit/microsoft/azure/log/test_wasb_task_handler.py index 57d21664c131f..3f9b36c6534a9 100644 --- a/providers/microsoft/azure/tests/unit/microsoft/azure/log/test_wasb_task_handler.py +++ b/providers/microsoft/azure/tests/unit/microsoft/azure/log/test_wasb_task_handler.py @@ -33,7 +33,11 @@ from tests_common.test_utils.config import conf_vars from tests_common.test_utils.db import clear_db_dags, clear_db_runs -from tests_common.test_utils.version_compat import AIRFLOW_V_3_0_PLUS, AIRFLOW_V_3_2_2_PLUS +from tests_common.test_utils.version_compat import ( + AIRFLOW_V_3_0_PLUS, + AIRFLOW_V_3_2_2_PLUS, + AIRFLOW_V_3_4_0_PLUS, +) pytestmark = pytest.mark.db_test @@ -127,7 +131,10 @@ def test_wasb_read(self, mock_hook_cls, ti): ) assert logs[2].event == "::endgroup::" assert logs[3].event == "Log line" - assert metadata == {"end_of_log": True, "log_pos": 1} + if AIRFLOW_V_3_4_0_PLUS: + assert metadata == {"end_of_log": True} + else: + assert metadata == {"end_of_log": True, "log_pos": 1} elif AIRFLOW_V_3_0_PLUS: logs = list(logs) assert logs[0].event == "::group::Log message source details"