Re-enable uploading to remote blob storage from Trigger and executors for the beta (apache#47115)

The way this achieved is the epitome of hackyness, 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.)

Author: ashb

airflow/jobs/triggerer_job_runner.py

@@ -234,6 +234,15 @@ def __call__(self, processors: Iterable[structlog.typing.Processor]) -> WrappedL
         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 @@ def _handle_request(self, msg: ToTriggerSupervisor, log: FilteringBoundLogger) -
                 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 @@ def update_triggers(self, requested_trigger_ids: set[int]):
             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 @@ def _process_log_messages_from_subprocess(self) -> Generator[None, bytes, None]:
 
         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:

airflow/models/trigger.py

@@ -246,7 +246,10 @@ def submit_event(cls, trigger_id, event: events.TriggerEvent, session: Session =
             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_none()
+        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}

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 @@ def start(
         # 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 @@ def start(
 
         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 @@ def wait(self) -> int:
             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.

task_sdk/src/airflow/sdk/log.py

@@ -31,7 +31,7 @@
 import structlog
 
 if TYPE_CHECKING:
-    from structlog.typing import EventDict, ExcInfo, Processor
+    from structlog.typing import EventDict, ExcInfo, FilteringBoundLogger, Processor
 
 
 __all__ = [
@@ -456,3 +456,44 @@ 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
+
+    fh = raw_logger._file
+    fname = fh.name
+
+    if fh.fileno() == 1 or not isinstance(fname, str):
+        # Logging to stdout, or something odd about this logger, don't try to upload!
+        return
+    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()

task_sdk/tests/execution_time/test_supervisor.py

@@ -505,6 +505,7 @@ def test_heartbeat_failures_handling(self, monkeypatch, mocker, captured_logs, t
         mock_kill = mocker.patch("airflow.sdk.execution_time.supervisor.WatchedSubprocess.kill")
 
         proc = ActivitySubprocess(
+            log=mocker.MagicMock(),
             id=TI_ID,
             pid=mock_process.pid,
             stdin=mocker.MagicMock(),
@@ -594,6 +595,7 @@ def test_overtime_handling(
         monkeypatch.setattr(ActivitySubprocess, "TASK_OVERTIME_THRESHOLD", overtime_threshold)
 
         mock_watched_subprocess = ActivitySubprocess(
+            log=mocker.MagicMock(),
             id=TI_ID,
             pid=12345,
             stdin=mocker.Mock(),
@@ -735,6 +737,7 @@ def mock_process(self, mocker):
     @pytest.fixture
     def watched_subprocess(self, mocker, mock_process):
         proc = ActivitySubprocess(
+            log=mocker.MagicMock(),
             id=TI_ID,
             pid=12345,
             stdin=mocker.Mock(),
@@ -918,6 +921,7 @@ class TestHandleRequest:
     def watched_subprocess(self, mocker):
         """Fixture to provide a WatchedSubprocess instance."""
         return ActivitySubprocess(
+            log=mocker.MagicMock(),
             id=TI_ID,
             pid=12345,
             stdin=BytesIO(),

tests/dag_processing/test_manager.py

@@ -137,6 +137,7 @@ def mock_processor(self) -> DagFileProcessorProcess:
         proc.create_time.return_value = time.time()
         proc.wait.return_value = 0
         ret = DagFileProcessorProcess(
+            log=MagicMock(),
             id=uuid7(),
             pid=1234,
             process=proc,

tests/jobs/test_triggerer_job.py

@@ -149,6 +149,7 @@ def builder(job=None):
 
         process = mocker.Mock(spec=psutil.Process, pid=10 * job.id + 1)
         proc = TriggerRunnerSupervisor(
+            log=mocker.Mock(),
             id=job.id,
             job=job,
             pid=process.pid,