Add streaming task log support to BaseExecutor and FileTaskHandler

airflow-core/src/airflow/executors/base_executor.py

@@ -71,6 +71,7 @@ def get_execution_api_server_url(conf_source: AirflowConfigParser | ExecutorConf
 
     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 @@ def get_task_log(self, ti: TaskInstance, try_number: int) -> tuple[list[str], li
         """
         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

airflow-core/src/airflow/utils/log/file_task_handler.py

@@ -23,7 +23,7 @@
 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 @@ def _render_filename(
             )
         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 @@ def _read(
                 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 (

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
 

airflow-core/tests/unit/utils/log/test_file_task_handler.py

@@ -21,6 +21,9 @@
 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 @@ def test_handles_base_log_folder_that_is_itself_a_symlink(self, tmp_path):
 
         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}