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 0b9bf0cd855 Stream remote task logs from S3 instead of reading whole
objects (#73517)
0b9bf0cd855 is described below
commit 0b9bf0cd855483fde15cb365e3e58fd4f42af483
Author: 장성태 <[email protected]>
AuthorDate: Sat Nov 7 02:49:50 2026 +0900
Stream remote task logs from S3 instead of reading whole objects (#73517)
Adds S3RemoteLogIO.stream() so that remote task logs stored in S3 are read
as a stream of lines instead of being materialised as one string per object,
and lets read() delegate to it.
---
.../providers/amazon/aws/log/s3_task_handler.py | 86 ++++++++++++++++++----
.../unit/amazon/aws/log/test_s3_task_handler.py | 84 +++++++++++++++++++++
2 files changed, 155 insertions(+), 15 deletions(-)
diff --git
a/providers/amazon/src/airflow/providers/amazon/aws/log/s3_task_handler.py
b/providers/amazon/src/airflow/providers/amazon/aws/log/s3_task_handler.py
index 32b0b1ee4dd..1bbe883f111 100644
--- a/providers/amazon/src/airflow/providers/amazon/aws/log/s3_task_handler.py
+++ b/providers/amazon/src/airflow/providers/amazon/aws/log/s3_task_handler.py
@@ -17,6 +17,7 @@
# under the License.
from __future__ import annotations
+import codecs
import inspect
import logging
import os
@@ -36,7 +37,19 @@ from airflow.utils.log.logging_mixin import LoggingMixin
if TYPE_CHECKING:
from airflow.models.taskinstance import TaskInstance
from airflow.sdk.types import RuntimeTaskInstanceProtocol as RuntimeTI
- from airflow.utils.log.file_task_handler import LogMessages, LogSourceInfo
+ from airflow.utils.log.file_task_handler import (
+ LogMessages,
+ LogResponse,
+ LogSourceInfo,
+ RawLogStream,
+ StreamingLogResponse,
+ )
+
+
+# Mirrors ``airflow.utils.log.file_task_handler.CHUNK_SIZE`` so that remote
reads are chunked the
+# same way local log reads are. Not imported from core to avoid coupling the
provider to a private
+# constant that older Airflow versions may not expose.
+CHUNK_SIZE = 1024 * 1024 * 5 # 5MB
@attrs.define
@@ -198,22 +211,65 @@ class S3RemoteLogIO(LoggingMixin): # noqa: D101
return False
return True
- def read(self, relative_path: str, ti: RuntimeTI | None) ->
tuple[LogSourceInfo, LogMessages | None]:
- logs: list[str] = []
- messages = []
+ def read(self, relative_path: str, ti: RuntimeTI | None) -> LogResponse:
+ messages, log_streams = self.stream(relative_path, ti)
+ if not log_streams:
+ return messages, None
+
+ # Exhaust each stream into a single string, for callers that cannot
consume a stream.
+ # The newline is added back per line so that the string is
byte-identical to the object.
+ return messages, ["".join(f"{line}\n" for line in log_stream) for
log_stream in log_streams]
+
+ def stream(self, relative_path: str, ti: RuntimeTI | None) ->
StreamingLogResponse:
+ messages: list[str] = []
bucket, prefix =
self.hook.parse_s3_url(s3url=os.path.join(self.remote_base, relative_path))
keys = self.hook.list_keys(bucket_name=bucket, prefix=prefix)
- if keys:
- keys = sorted(f"s3://{bucket}/{key}" for key in keys)
- if AIRFLOW_V_3_0_PLUS:
- messages = keys
- else:
- messages.append("Found logs in s3:")
- messages.extend(f" * {key}" for key in keys)
- for key in keys:
- logs.append(self.s3_read(key, return_error=True))
- return messages, logs
- return messages, None
+ if not keys:
+ return messages, []
+
+ keys = sorted(f"s3://{bucket}/{key}" for key in keys)
+ if AIRFLOW_V_3_0_PLUS:
+ messages = keys
+ else:
+ messages.append("Found logs in s3:")
+ messages.extend(f" * {key}" for key in keys)
+
+ # The objects are not fetched here: each generator issues its
GetObject on first use.
+ return messages, [self._get_log_stream(key) for key in keys]
+
+ def _get_log_stream(self, remote_log_location: str) -> RawLogStream:
+ r"""
+ Yield the lines of a remote log object without holding the whole
object in memory.
+
+ Lines are split exactly like
:func:`~airflow.utils.log.file_task_handler._stream_lines_by_chunk`
+ does for local logs: on ``"\n"`` only, with the trailing newline
stripped, so the records the
+ log reader sees are the same ones it saw when this method read the
object as a single string.
+
+ :param remote_log_location: the log's location in remote storage
+ :yield: Lines of the log object.
+ """
+ body = None
+ try:
+ body = self.hook.get_key(remote_log_location).get()["Body"]
+ # An incremental decoder keeps multi-byte characters intact across
chunk boundaries.
+ decoder = codecs.getincrementaldecoder("utf-8")(errors="replace")
+ buffer = ""
+ for chunk in body.iter_chunks(chunk_size=CHUNK_SIZE):
+ buffer += decoder.decode(chunk)
+ *lines, buffer = buffer.split("\n")
+ yield from lines
+ buffer += decoder.decode(b"", final=True)
+ if buffer:
+ yield from buffer.split("\n")
+ except Exception as error:
+ # Surface the failure as a log line instead of raising, so that a
single unreadable
+ # object does not hide the logs of the other objects of the same
task instance.
+ msg = f"Could not read logs from {remote_log_location} with error:
{error}"
+ self.log.exception(msg)
+ yield msg
+ finally:
+ if body is not None:
+ body.close()
class S3TaskHandler(FileTaskHandler, LoggingMixin):
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 85ef2325e64..e922400ffd7 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
@@ -289,6 +289,90 @@ class TestS3RemoteLogIO:
assert rec.levelno == logging.ERROR
assert rec.exc_info is not None
+ @staticmethod
+ def _expected_messages(*locations: str) -> list[str]:
+ """Source messages as ``stream()`` and ``read()`` format them, which
differs on Airflow 2."""
+ if AIRFLOW_V_3_0_PLUS:
+ return list(locations)
+ return ["Found logs in s3:", *(f" * {location}" for location in
locations)]
+
+ def test_stream_returns_no_streams_when_there_are_no_logs(self):
+ messages, log_streams = self.subject.stream("1.log", self.ti)
+
+ assert messages == []
+ assert log_streams == []
+
+ def test_read_returns_none_when_there_are_no_logs(self):
+ messages, logs = self.subject.read("1.log", self.ti)
+
+ assert messages == []
+ assert logs is None
+
+ def test_stream_yields_lines_without_the_trailing_newline(self):
+ self.conn.put_object(Bucket="bucket", Key=self.remote_log_key,
Body=b"Log line\nLine 2\n")
+
+ messages, log_streams = self.subject.stream("1.log", self.ti)
+
+ assert messages == self._expected_messages(self.remote_log_location)
+ assert list(log_streams[0]) == ["Log line", "Line 2"]
+
+ def test_stream_sorts_the_keys_of_the_task_instance(self):
+ self.conn.put_object(Bucket="bucket", Key=self.remote_log_key +
".trigger.log", Body=b"Line 3\n")
+ self.conn.put_object(Bucket="bucket", Key=self.remote_log_key,
Body=b"Line 1\n")
+
+ messages, log_streams = self.subject.stream("1.log", self.ti)
+
+ assert messages == self._expected_messages(
+ self.remote_log_location, self.remote_log_location + ".trigger.log"
+ )
+ assert [list(log_stream) for log_stream in log_streams] == [["Line
1"], ["Line 3"]]
+
+ def test_stream_does_not_fetch_the_object_until_it_is_consumed(self):
+ self.conn.put_object(Bucket="bucket", Key=self.remote_log_key,
Body=b"Log line\n")
+
+ with mock.patch.object(S3Hook, "get_key",
wraps=self.subject.hook.get_key) as get_key:
+ _, log_streams = self.subject.stream("1.log", self.ti)
+ get_key.assert_not_called()
+
+ assert list(log_streams[0]) == ["Log line"]
+ get_key.assert_called_once()
+
+ def test_read_does_not_load_the_object_as_a_single_string(self):
+ self.conn.put_object(Bucket="bucket", Key=self.remote_log_key,
Body=b"Log line\nLine 2\n")
+
+ with mock.patch.object(S3RemoteLogIO, "s3_read") as s3_read:
+ messages, logs = self.subject.read("1.log", self.ti)
+
+ s3_read.assert_not_called()
+ assert messages == self._expected_messages(self.remote_log_location)
+ assert logs == ["Log line\nLine 2\n"]
+
+ def test_stream_keeps_multibyte_characters_intact_across_chunks(self):
+ self.conn.put_object(
+ Bucket="bucket", Key=self.remote_log_key, Body="첫 번째 줄\n두 번째
줄\n".encode()
+ )
+
+ # A chunk size that is guaranteed to split multi-byte characters in
half.
+ with
mock.patch("airflow.providers.amazon.aws.log.s3_task_handler.CHUNK_SIZE", 3):
+ _, log_streams = self.subject.stream("1.log", self.ti)
+
+ assert list(log_streams[0]) == ["첫 번째 줄", "두 번째 줄"]
+
+ def test_stream_yields_an_error_line_when_the_object_cannot_be_read(self,
caplog):
+ self.conn.put_object(Bucket="bucket", Key=self.remote_log_key,
Body=b"Log line\n")
+ error = ClientError({"Error": {"Code": "403", "Message":
"Forbidden"}}, "GetObject")
+
+ with mock.patch.object(S3Hook, "get_key", side_effect=error):
+ _, log_streams = self.subject.stream("1.log", self.ti)
+ with caplog.at_level(logging.ERROR):
+ lines = list(log_streams[0])
+
+ assert len(lines) == 1
+ assert lines[0].startswith(f"Could not read logs from
{self.remote_log_location} with error:")
+ errors = [record for record in caplog.records if record.levelno ==
logging.ERROR]
+ assert len(errors) == 1
+ assert errors[0].exc_info is not None
+
def test_write(self):
with mock.patch.object(self.subject.log, "error") as mock_error:
self.subject.write("text", self.remote_log_location)
