[Resolve OOM When Reading Large Logs in Webserver] Refactor to Use K-Way Merge for Log Streams Instead of Sorting Entire Log Records

airflow/api_connexion/endpoints/log_endpoint.py

@@ -115,11 +115,14 @@ def get_log(
     # return_type would be either the above two or None
     logs: Any
     if return_type == "application/json" or return_type is None:  # default
-        logs, metadata = task_log_reader.read_log_chunks(ti, task_try_number, metadata)
-        logs = logs[0] if task_try_number is not None else logs
+        hosts, log_streams, metadata = task_log_reader.read_log_chunks(ti, task_try_number, metadata)
+        host = f"{hosts[0] or ''}\n"
+        logs = log_streams[0] if task_try_number is not None else log_streams
         # we must have token here, so we can safely ignore it
         token = URLSafeSerializer(key).dumps(metadata)  # type: ignore[assignment]
-        return logs_schema.dump(LogResponseObject(continuation_token=token, content=logs))
+        return logs_schema.dump(
+            LogResponseObject(continuation_token=token, content=host + "\n".join(log for log in logs))
+        )
     # text/plain. Stream
     logs = task_log_reader.read_log_stream(ti, task_try_number, metadata)
 

airflow/api_fastapi/core_api/routes/public/log.py

@@ -18,13 +18,14 @@
 from __future__ import annotations
 
 import textwrap
-from typing import Any
+from collections.abc import Iterable
 
-from fastapi import HTTPException, Request, Response, status
+from fastapi import HTTPException, Request, status
 from itsdangerous import BadSignature, URLSafeSerializer
 from pydantic import PositiveInt
 from sqlalchemy.orm import joinedload
 from sqlalchemy.sql import select
+from starlette.responses import StreamingResponse
 
 from airflow.api_fastapi.common.db.common import SessionDep
 from airflow.api_fastapi.common.headers import HeaderAcceptJsonOrText
@@ -135,12 +136,17 @@ def get_log(
         except TaskNotFound:
             pass
 
-    logs: Any
+    log_streams: list[Iterable[str]]
+    log_stream: Iterable[str]
     if accept == Mimetype.JSON or accept == Mimetype.ANY:  # default
-        logs, metadata = task_log_reader.read_log_chunks(ti, try_number, metadata)
+        hosts, log_streams, metadata = task_log_reader.read_log_chunks(ti, try_number, metadata)
         # we must have token here, so we can safely ignore it
         token = URLSafeSerializer(request.app.state.secret_key).dumps(metadata)  # type: ignore[assignment]
-        return TaskInstancesLogResponse(continuation_token=token, content=str(logs[0])).model_dump()
+        host = f"{hosts[0] or ''}\n"
+        log_stream = log_streams[0]
+        return TaskInstancesLogResponse(
+            continuation_token=token, content=host + "\n".join(log for log in log_stream)
+        ).model_dump()
     # text/plain. Stream
-    logs = task_log_reader.read_log_stream(ti, try_number, metadata)
-    return Response(media_type=accept, content="".join(list(logs)))
+    log_stream = task_log_reader.read_log_stream(ti, try_number, metadata)
+    return StreamingResponse(log_stream, media_type=Mimetype.TEXT.value)

airflow/executors/base_executor.py

@@ -21,7 +21,7 @@
 import logging
 import sys
 from collections import defaultdict, deque
-from collections.abc import Sequence
+from collections.abc import Generator, Sequence
 from dataclasses import dataclass, field
 from typing import TYPE_CHECKING, Any, Optional
 
@@ -544,13 +544,20 @@ def execute_async(
         """
         raise NotImplementedError()
 
-    def get_task_log(self, ti: TaskInstance, try_number: int) -> tuple[list[str], list[str]]:
+    def get_task_log(
+        self, ti: TaskInstance, try_number: int
+    ) -> (
+        tuple[list[str], list[Generator[tuple[pendulum.DateTime | None, int, str], None, None]], int]
+        | tuple[list[str], list[str]]
+    ):
         """
         Return the task logs.
 
         :param ti: A TaskInstance object
         :param try_number: current try_number to read log from
-        :return: tuple of logs and messages
+        :return:
+            - old interface: Tuple of messages and list of log lines.
+            - new interface: Tuple of messages, parsed log streams, total size of logs.
         """
         return [], []