Skip log accumulator disk spill when reading finished task logs

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

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()

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):

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",

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)

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"

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"