This is an automated email from the ASF dual-hosted git repository.
o-nikolas pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/airflow.git
The following commit(s) were added to refs/heads/main by this push:
new 65fba6f1200 Add streaming task log support to BaseExecutor and
FileTaskHandler (#69299)
65fba6f1200 is described below
commit 65fba6f120069779a661f31f5a346079fd3a26d2
Author: Jason(Zhe-You) Liu <[email protected]>
AuthorDate: Wed Jul 8 05:01:24 2026 +0900
Add streaming task log support to BaseExecutor and FileTaskHandler (#69299)
Reading a running task's log through an executor materializes the whole
log in the API server before the bounded LogStreamAccumulator can bound
memory, so large logs spike the API server heap. This adds an interface
executors can implement to stream log lines lazily instead.
---
.../src/airflow/executors/base_executor.py | 14 ++++
.../src/airflow/utils/log/file_task_handler.py | 44 +++++++------
.../tests/unit/executors/test_base_executor.py | 8 +++
.../tests/unit/utils/log/test_file_task_handler.py | 74 ++++++++++++++++++++++
4 files changed, 120 insertions(+), 20 deletions(-)
diff --git a/airflow-core/src/airflow/executors/base_executor.py
b/airflow-core/src/airflow/executors/base_executor.py
index ef577233110..ebd850ad944 100644
--- a/airflow-core/src/airflow/executors/base_executor.py
+++ b/airflow-core/src/airflow/executors/base_executor.py
@@ -71,6 +71,7 @@ if TYPE_CHECKING:
from sqlalchemy.orm import Session
+ from airflow._shared.logging.remote import StreamingLogResponse
from airflow.api_fastapi.auth.tokens import JWTGenerator
from airflow.callbacks.base_callback_sink import BaseCallbackSink
from airflow.callbacks.callback_requests import CallbackRequest
@@ -559,6 +560,19 @@ class BaseExecutor(LoggingMixin):
"""
return [], []
+ def get_streaming_task_log(self, ti: TaskInstance, try_number: int) ->
StreamingLogResponse:
+ """
+ Return a streaming response for task logs.
+
+ Executors that don't implement this method raise
``NotImplementedError``; callers should
+ catch that and fall back to :meth:`get_task_log`.
+
+ :param ti: A TaskInstance object
+ :param try_number: current try_number to read log from
+ :return: StreamingLogResponse
+ """
+ raise NotImplementedError
+
def end(self) -> None: # pragma: no cover
"""Wait synchronously for the previously submitted job to complete."""
raise NotImplementedError
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 81f9a89c458..d22ce48fa08 100644
--- a/airflow-core/src/airflow/utils/log/file_task_handler.py
+++ b/airflow-core/src/airflow/utils/log/file_task_handler.py
@@ -23,7 +23,7 @@ import heapq
import io
import logging
import os
-from collections.abc import Callable, Generator, Iterator
+from collections.abc import Generator, Iterator
from contextlib import suppress
from datetime import datetime
from enum import Enum
@@ -557,27 +557,25 @@ class FileTaskHandler(logging.Handler):
)
raise RuntimeError(f"Unable to render log filename for {ti}. This
should never happen")
- def _get_executor_get_task_log(
- self, ti: TaskInstance | TaskInstanceHistory
- ) -> Callable[[TaskInstance | TaskInstanceHistory, int], tuple[list[str],
list[str]]]:
+ def _get_executor(self, ti: TaskInstance | TaskInstanceHistory) ->
BaseExecutor:
"""
- Get the get_task_log method from executor of current task instance.
+ Get the executor of current task instance.
Since there might be multiple executors, so we need to get the
executor of current task instance instead of getting from default executor.
:param ti: task instance object
- :return: get_task_log method of the executor
+ :return: executor of the task instance
"""
executor_name = ti.executor or self.DEFAULT_EXECUTOR_KEY
executor = self.executor_instances.get(executor_name)
- if executor is not None:
- return executor.get_task_log
+ if executor is None:
+ if executor_name == self.DEFAULT_EXECUTOR_KEY:
+ executor = ExecutorLoader.get_default_executor()
+ else:
+ executor = ExecutorLoader.load_executor(executor_name)
+ self.executor_instances[executor_name] = executor
- if executor_name == self.DEFAULT_EXECUTOR_KEY:
- self.executor_instances[executor_name] =
ExecutorLoader.get_default_executor()
- else:
- self.executor_instances[executor_name] =
ExecutorLoader.load_executor(executor_name)
- return self.executor_instances[executor_name].get_task_log
+ return executor
def _read(
self,
@@ -632,23 +630,29 @@ class FileTaskHandler(logging.Handler):
raise TypeError("Logs should be either a list of strings or a
generator of log lines.")
# Extend LogSourceInfo
source_list.extend(sources)
- has_k8s_exec_pod = False
+
+ has_executor_log = False
if ti.state == TaskInstanceState.RUNNING:
- executor_get_task_log = self._get_executor_get_task_log(ti)
- response = executor_get_task_log(ti, try_number)
- if response:
- sources, logs = response
+ executor = self._get_executor(ti)
+ try:
+ # check for streaming logs first
+ sources, executor_logs = executor.get_streaming_task_log(ti,
try_number)
+ except NotImplementedError:
+ # fallback to non-streaming logs if streaming not supported
+ sources, logs = executor.get_task_log(ti, try_number)
# make the logs stream-like compatible
executor_logs = [_get_compatible_log_stream(logs)]
+
if sources:
source_list.extend(sources)
- has_k8s_exec_pod = True
+ has_executor_log = True
+
if not (remote_logs and ti.state not in State.unfinished):
# when finished, if we have remote logs, no need to check local
worker_log_full_path = Path(self.local_base, worker_log_rel_path)
sources, local_logs = self._read_from_local(worker_log_full_path)
source_list.extend(sources)
- if ti.state in (TaskInstanceState.RUNNING, TaskInstanceState.DEFERRED)
and not has_k8s_exec_pod:
+ if ti.state in (TaskInstanceState.RUNNING, TaskInstanceState.DEFERRED)
and not has_executor_log:
sources, served_logs = self._read_from_logs_server(ti,
worker_log_rel_path)
source_list.extend(sources)
elif (ti.state not in State.unfinished or ti.state in
_STATES_WITH_COMPLETED_ATTEMPT) and not (
diff --git a/airflow-core/tests/unit/executors/test_base_executor.py
b/airflow-core/tests/unit/executors/test_base_executor.py
index cb646f7ce8d..8bc7cfd73db 100644
--- a/airflow-core/tests/unit/executors/test_base_executor.py
+++ b/airflow-core/tests/unit/executors/test_base_executor.py
@@ -78,6 +78,14 @@ def test_get_task_log():
assert executor.get_task_log(ti=ti, try_number=1) == ([], [])
+def test_get_streaming_task_log_not_implemented():
+ executor = BaseExecutor()
+ ti = TaskInstance(task=SerializedBaseOperator(task_id="dummy"),
dag_version_id=mock.MagicMock(spec=UUID))
+
+ with pytest.raises(NotImplementedError):
+ executor.get_streaming_task_log(ti=ti, try_number=1)
+
+
def test_serve_logs_default_value():
assert not BaseExecutor.serve_logs
diff --git a/airflow-core/tests/unit/utils/log/test_file_task_handler.py
b/airflow-core/tests/unit/utils/log/test_file_task_handler.py
index ceecffafa32..09bb0c23b51 100644
--- a/airflow-core/tests/unit/utils/log/test_file_task_handler.py
+++ b/airflow-core/tests/unit/utils/log/test_file_task_handler.py
@@ -21,6 +21,9 @@ from pathlib import Path
from unittest.mock import MagicMock, patch
from airflow.utils.log.file_task_handler import FileTaskHandler
+from airflow.utils.state import TaskInstanceState
+
+from tests_common.test_utils.file_task_handler import convert_list_to_stream,
extract_events
class TestFileTaskHandlerLogServer:
@@ -220,3 +223,74 @@ class TestFileTaskHandlerReadFromLocal:
assert len(sources) == 1
assert "through-symlink content" in self._drain(streams[0])
+
+
+class TestFileTaskHandlerExecutorLogs:
+ """Tests for executor log retrieval selection."""
+
+ @staticmethod
+ def _running_ti(executor_name: str) -> MagicMock:
+ ti = MagicMock()
+ ti.executor = executor_name
+ ti.state = TaskInstanceState.RUNNING
+ ti.try_number = 1
+ return ti
+
+ def test_running_task_prefers_streaming_executor_logs(self):
+ """Use executor streaming logs when the executor implements
streaming."""
+ handler = FileTaskHandler(base_log_folder="")
+ executor = MagicMock()
+ executor.get_streaming_task_log.return_value = (
+ ["streaming source"],
+ [convert_list_to_stream(["streaming log"])],
+ )
+ executor.get_task_log.return_value = (["legacy source"], ["legacy
log"])
+ handler.executor_instances = {"StreamingExecutor": executor}
+ ti = self._running_ti("StreamingExecutor")
+
+ with (
+ patch.object(handler, "_render_filename",
return_value="dag/run/task/1.log"),
+ patch.object(handler, "_read_remote_logs",
side_effect=NotImplementedError),
+ patch.object(handler, "_read_from_local", return_value=([], [])),
+ patch.object(handler, "_read_from_logs_server", return_value=([],
[])) as read_from_logs_server,
+ ):
+ logs, metadata = handler._read(ti=ti, try_number=1)
+
+ executor.get_streaming_task_log.assert_called_once_with(ti, 1)
+ executor.get_task_log.assert_not_called()
+ read_from_logs_server.assert_not_called()
+ assert extract_events(logs, skip_source_info=False) == [
+ "::group::Log message source details",
+ "streaming source",
+ "::endgroup::",
+ "streaming log",
+ ]
+ assert metadata == {"end_of_log": False, "log_pos": 1}
+
+ def test_running_task_falls_back_to_legacy_executor_logs(self):
+ """Use legacy executor logs when the executor doesn't implement
streaming."""
+ handler = FileTaskHandler(base_log_folder="")
+ executor = MagicMock()
+ executor.get_streaming_task_log.side_effect = NotImplementedError
+ executor.get_task_log.return_value = (["legacy source"], ["legacy
log"])
+ handler.executor_instances = {"LegacyExecutor": executor}
+ ti = self._running_ti("LegacyExecutor")
+
+ with (
+ patch.object(handler, "_render_filename",
return_value="dag/run/task/1.log"),
+ patch.object(handler, "_read_remote_logs",
side_effect=NotImplementedError),
+ patch.object(handler, "_read_from_local", return_value=([], [])),
+ patch.object(handler, "_read_from_logs_server", return_value=([],
[])) as read_from_logs_server,
+ ):
+ logs, metadata = handler._read(ti=ti, try_number=1)
+
+ executor.get_streaming_task_log.assert_called_once_with(ti, 1)
+ executor.get_task_log.assert_called_once_with(ti, 1)
+ read_from_logs_server.assert_not_called()
+ assert extract_events(logs, skip_source_info=False) == [
+ "::group::Log message source details",
+ "legacy source",
+ "::endgroup::",
+ "legacy log",
+ ]
+ assert metadata == {"end_of_log": False, "log_pos": 1}
