You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@airflow.apache.org by as...@apache.org on 2025/02/26 16:58:16 UTC
(airflow) 01/01: Re-enable uploading to remote blob sotrage from Trigger and executors for the beta
This is an automated email from the ASF dual-hosted git repository.
ash pushed a commit to branch task-sdk-upload-remote-logs-from-supervisor
in repository https://gitbox.apache.org/repos/asf/airflow.git
commit c1796c4c93196f006879a27d6905b675742a220a
Author: Ash Berlin-Taylor <as...@apache.org>
AuthorDate: Wed Feb 26 16:37:47 2025 +0000
Re-enable uploading to remote blob sotrage from Trigger and executors for the beta
The way this achieved is the opitomy of hacky, but it works well enough for
S3, GCS and WASB task handlers.
It won't work for the "non-blob" based logging backends such as
CloudWatchLogs, and it only works with the way the 3 blob-based log backends
are currently implemented.
I want remote logs to be working for most people for the Beta, and I will
follow up with a better fix (that doesn't rely on internal implementation
details and works for log "streams" too.)
---
airflow/jobs/triggerer_job_runner.py | 22 +++++++++----
airflow/models/trigger.py | 5 ++-
.../src/airflow/sdk/execution_time/supervisor.py | 32 +++++++++++++-----
task_sdk/src/airflow/sdk/log.py | 38 +++++++++++++++++++++-
4 files changed, 81 insertions(+), 16 deletions(-)
diff --git a/airflow/jobs/triggerer_job_runner.py b/airflow/jobs/triggerer_job_runner.py
index 0e82cd1c193..b630f66ee13 100644
--- a/airflow/jobs/triggerer_job_runner.py
+++ b/airflow/jobs/triggerer_job_runner.py
@@ -234,6 +234,15 @@ class TriggerLoggingFactory:
self.bound_logger = logger
return logger
+ def upload_to_remote(self):
+ from airflow.sdk.log import upload_to_remote
+
+ if not hasattr(self, "bound_logger"):
+ # Never actually called, nothing to do
+ return
+
+ upload_to_remote(self.bound_logger)
+
@attrs.define(kw_only=True)
class TriggerRunnerSupervisor(WatchedSubprocess):
@@ -299,7 +308,8 @@ class TriggerRunnerSupervisor(WatchedSubprocess):
self.cancelling_triggers.discard(id)
# Remove logger from the cache, and since structlog doesn't have an explicit close method, we
# only need to remove the last reference to it to close the open FH
- self.logger_cache.pop(id, None)
+ if factory := self.logger_cache.pop(id, None):
+ factory.upload_to_remote()
return
raise ValueError(f"Unknown message type {type(msg)}")
@@ -460,10 +470,8 @@ class TriggerRunnerSupervisor(WatchedSubprocess):
self.cancelling_triggers.update(cancel_trigger_ids)
self._send(messages.CancelTriggers(ids=cancel_trigger_ids))
- def _register_pipe_readers(
- self, logger: FilteringBoundLogger, stdout: socket, stderr: socket, requests: socket, logs: socket
- ):
- super()._register_pipe_readers(logger, stdout, stderr, requests, logs)
+ def _register_pipe_readers(self, stdout: socket, stderr: socket, requests: socket, logs: socket):
+ super()._register_pipe_readers(stdout, stderr, requests, logs)
# We want to handle logging differently here, so un-register the one our parent class created
self.selector.unregister(logs)
@@ -488,7 +496,9 @@ class TriggerRunnerSupervisor(WatchedSubprocess):
def get_logger(trigger_id: int) -> WrappedLogger:
# TODO: Is a separate dict worth it, or should we make `self.running_triggers` a dict?
- return self.logger_cache[trigger_id](processors)
+ if factory := self.logger_cache.get(trigger_id):
+ return factory(processors)
+ return fallback_log
# We need to look at the json, pull out the
while True:
diff --git a/airflow/models/trigger.py b/airflow/models/trigger.py
index 478eb89a675..db7c764ada4 100644
--- a/airflow/models/trigger.py
+++ b/airflow/models/trigger.py
@@ -246,7 +246,10 @@ class Trigger(Base):
handle_event_submit(event, task_instance=task_instance, session=session)
# Send an event to assets
- trigger = session.scalars(select(cls).where(cls.id == trigger_id)).one()
+ trigger = session.scalars(select(cls).where(cls.id == trigger_id)).one_or_one()
+ if trigger is None:
+ # Already deleted for some reason
+ return
for asset in trigger.assets:
AssetManager.register_asset_change(
asset=asset.to_public(), session=session, extra={"from_trigger": True}
diff --git a/task_sdk/src/airflow/sdk/execution_time/supervisor.py b/task_sdk/src/airflow/sdk/execution_time/supervisor.py
index fada6105bce..8bc867a3f5b 100644
--- a/task_sdk/src/airflow/sdk/execution_time/supervisor.py
+++ b/task_sdk/src/airflow/sdk/execution_time/supervisor.py
@@ -316,6 +316,8 @@ class WatchedSubprocess:
selector: selectors.BaseSelector = attrs.field(factory=selectors.DefaultSelector)
+ log: FilteringBoundLogger
+
@classmethod
def start(
cls,
@@ -363,17 +365,17 @@ class WatchedSubprocess:
# other end of the pair open
cls._close_unused_sockets(child_stdin, child_stdout, child_stderr, child_comms, child_logs)
+ logger = logger or cast("FilteringBoundLogger", structlog.get_logger(logger_name="task").bind())
proc = cls(
pid=pid,
stdin=feed_stdin,
process=psutil.Process(pid),
requests_fd=requests_fd,
+ log=logger,
**constructor_kwargs,
)
- logger = logger or cast("FilteringBoundLogger", structlog.get_logger(logger_name="task").bind())
proc._register_pipe_readers(
- logger=logger,
stdout=read_stdout,
stderr=read_stderr,
requests=read_msgs,
@@ -382,26 +384,24 @@ class WatchedSubprocess:
return proc
- def _register_pipe_readers(
- self, logger: FilteringBoundLogger, stdout: socket, stderr: socket, requests: socket, logs: socket
- ):
+ def _register_pipe_readers(self, stdout: socket, stderr: socket, requests: socket, logs: socket):
"""Register handlers for subprocess communication channels."""
# self.selector is a way of registering a handler/callback to be called when the given IO channel has
# activity to read on (https://www.man7.org/linux/man-pages/man2/select.2.html etc, but better
# alternatives are used automatically) -- this is a way of having "event-based" code, but without
# needing full async, to read and process output from each socket as it is received.
- self.selector.register(stdout, selectors.EVENT_READ, self._create_socket_handler(logger, "stdout"))
+ self.selector.register(stdout, selectors.EVENT_READ, self._create_socket_handler(self.log, "stdout"))
self.selector.register(
stderr,
selectors.EVENT_READ,
- self._create_socket_handler(logger, "stderr", log_level=logging.ERROR),
+ self._create_socket_handler(self.log, "stderr", log_level=logging.ERROR),
)
self.selector.register(
logs,
selectors.EVENT_READ,
make_buffered_socket_reader(
- process_log_messages_from_subprocess(logger), on_close=self._on_socket_closed
+ process_log_messages_from_subprocess(self.log), on_close=self._on_socket_closed
),
)
self.selector.register(
@@ -658,8 +658,24 @@ class ActivitySubprocess(WatchedSubprocess):
self.client.task_instances.finish(
id=self.id, state=self.final_state, when=datetime.now(tz=timezone.utc)
)
+
+ # Now at the last possible moment, when all logs and comms with the subprocess has finished, lets
+ # upload the remote logs
+ self._upload_logs()
+
return self._exit_code
+ def _upload_logs(self):
+ """
+ Upload all log files found to the remote storage.
+
+ We upload logs from here after the task has finished to give us the best possible chance of logs being
+ uploaded in case the task task.
+ """
+ from airflow.sdk.log import upload_to_remote
+
+ upload_to_remote(self.log)
+
def _monitor_subprocess(self):
"""
Monitor the subprocess until it exits.
diff --git a/task_sdk/src/airflow/sdk/log.py b/task_sdk/src/airflow/sdk/log.py
index 0bb06934c94..9396e0c7c01 100644
--- a/task_sdk/src/airflow/sdk/log.py
+++ b/task_sdk/src/airflow/sdk/log.py
@@ -31,7 +31,7 @@ import msgspec
import structlog
if TYPE_CHECKING:
- from structlog.typing import EventDict, ExcInfo, Processor
+ from structlog.typing import EventDict, ExcInfo, FilteringBoundLogger, Processor
__all__ = [
@@ -456,3 +456,39 @@ def init_log_file(local_relative_path: str) -> Path:
log.warning("OSError while changing ownership of the log file. %s", e)
return full_path
+
+
+def upload_to_remote(logger: FilteringBoundLogger):
+ # We haven't yet switched the Remote log handlers over, they are still wired up in providers as
+ # logging.Handlers (but we should re-write most of them to just be the upload and read instead of full
+ # variants.) In the mean time, lets just create the right handler directly
+ from airflow.configuration import conf
+ from airflow.utils.log.file_task_handler import FileTaskHandler
+ from airflow.utils.log.log_reader import TaskLogReader
+
+ raw_logger = getattr(logger, "_logger")
+
+ if not raw_logger or not hasattr(raw_logger, "_file"):
+ logger.warning("Unable to find log file, logger was of unexpected type", type=type(logger))
+ return
+
+ fname = raw_logger._file.name
+ base_log_folder = conf.get("logging", "base_log_folder")
+ relative_path = Path(fname).relative_to(base_log_folder)
+
+ handler = TaskLogReader().log_handler
+ if not isinstance(handler, FileTaskHandler):
+ logger.warning(
+ "Airflow core logging is not using a FileTaskHandler, can't upload logs to remote",
+ handler=type(handler),
+ )
+ return
+
+ # This is a _monstrosity_, and super fragile, but we don't want to do the base FileTaskHandler
+ # set_context() which opens a real FH again. (And worse, in some cases it _truncates_ the file too). This
+ # is just for the first Airflow 3 betas, but we will re-write a better remote log interface that isn't
+ # tied to being a logging Handler.
+ handler.log_relative_path = relative_path.as_posix() # type: ignore[attr-defined]
+ handler.upload_on_close = True # type: ignore[attr-defined]
+
+ handler.close()