Send task logs from KubeExecutor tasks to pod stdout too

airflow/settings.py

@@ -389,8 +389,13 @@ def _session_maker(_engine):
     Session = scoped_session(NonScopedSession)
 
     # https://docs.sqlalchemy.org/en/20/core/pooling.html#using-connection-pools-with-multiprocessing-or-os-fork
-    os.register_at_fork(after_in_child=lambda: engine.dispose(close=False))
-    os.register_at_fork(after_in_child=lambda: async_engine.sync_engine.dispose(close=False))
+    def clean_in_fork():
+        if engine:
+            engine.dispose(close=False)
+        if async_engine:
+            async_engine.sync_engine.dispose(close=False)
+
+    os.register_at_fork(after_in_child=clean_in_fork)
 
 
 DEFAULT_ENGINE_ARGS = {
@@ -480,15 +485,16 @@ def prepare_engine_args(disable_connection_pool=False, pool_class=None):
     return engine_args
 
 
-def dispose_orm():
+def dispose_orm(do_log: bool = True):
     """Properly close pooled database connections."""
     global Session, engine, NonScopedSession
 
     _globals = globals()
-    if "engine" not in _globals and "Session" not in _globals:
+    if _globals.get("engine") is None and _globals.get("Session") is None:
         return
 
-    log.debug("Disposing DB connection pool (PID %s)", os.getpid())
+    if do_log:
+        log.debug("Disposing DB connection pool (PID %s)", os.getpid())
 
     if "Session" in _globals and Session is not None:
         from sqlalchemy.orm.session import close_all_sessions

task-sdk/src/airflow/sdk/execution_time/execute_workload.py

@@ -42,16 +42,19 @@ def execute_workload(input: str) -> None:
     from airflow.executors import workloads
     from airflow.sdk.execution_time.supervisor import supervise
     from airflow.sdk.log import configure_logging
+    from airflow.settings import dispose_orm
 
-    configure_logging(output=sys.stdout.buffer)
+    dispose_orm(do_log=False)
+
+    configure_logging(output=sys.stdout.buffer, enable_pretty_log=False)
 
     decoder = TypeAdapter[workloads.All](workloads.All)
     workload = decoder.validate_json(input)
 
     if not isinstance(workload, workloads.ExecuteTask):
-        raise ValueError(f"KubernetesExecutor does not know how to handle {type(workload)}")
+        raise ValueError(f"We do not know how to handle {type(workload)}")
 
-    log.info("Executing workload in Kubernetes", workload=workload)
+    log.info("Executing workload", workload=workload)
 
     supervise(
         # This is the "wrong" ti type, but it duck types the same. TODO: Create a protocol for this.
@@ -61,6 +64,9 @@ def execute_workload(input: str) -> None:
         token=workload.token,
         server=conf.get("core", "execution_api_server_url"),
         log_path=workload.log_path,
+        # Include the output of the task to stdout too, so that in process logs can be read from via the
+        # kubeapi as pod logs.
+        subprocess_logs_to_stdout=True,
     )
 
 

task-sdk/src/airflow/sdk/execution_time/supervisor.py

@@ -369,7 +369,10 @@ class WatchedSubprocess:
 
     selector: selectors.BaseSelector = attrs.field(factory=selectors.DefaultSelector)
 
-    log: FilteringBoundLogger
+    process_log: FilteringBoundLogger
+
+    subprocess_logs_to_stdout: bool = False
+    """Duplicate log messages to stdout, or only send them to ``self.process_log``."""
 
     @classmethod
     def start(
@@ -425,7 +428,7 @@ def start(
             stdin=feed_stdin,
             process=psutil.Process(pid),
             requests_fd=requests_fd,
-            log=logger,
+            process_log=logger,
             **constructor_kwargs,
         )
 
@@ -445,17 +448,22 @@ def _register_pipe_readers(self, stdout: socket, stderr: socket, requests: socke
         # 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(self.log, "stdout"))
+        target_loggers: tuple[FilteringBoundLogger, ...] = (self.process_log,)
+        if self.subprocess_logs_to_stdout:
+            target_loggers += (log,)
+        self.selector.register(
+            stdout, selectors.EVENT_READ, self._create_socket_handler(target_loggers, channel="stdout")
+        )
         self.selector.register(
             stderr,
             selectors.EVENT_READ,
-            self._create_socket_handler(self.log, "stderr", log_level=logging.ERROR),
+            self._create_socket_handler(target_loggers, channel="stderr", log_level=logging.ERROR),
         )
         self.selector.register(
             logs,
             selectors.EVENT_READ,
             make_buffered_socket_reader(
-                process_log_messages_from_subprocess(self.log), on_close=self._on_socket_closed
+                process_log_messages_from_subprocess(target_loggers), on_close=self._on_socket_closed
             ),
         )
         self.selector.register(
@@ -464,10 +472,10 @@ def _register_pipe_readers(self, stdout: socket, stderr: socket, requests: socke
             make_buffered_socket_reader(self.handle_requests(log), on_close=self._on_socket_closed),
         )
 
-    def _create_socket_handler(self, logger, channel, log_level=logging.INFO) -> Callable[[socket], bool]:
+    def _create_socket_handler(self, loggers, channel, log_level=logging.INFO) -> Callable[[socket], bool]:
         """Create a socket handler that forwards logs to a logger."""
         return make_buffered_socket_reader(
-            forward_to_log(logger.bind(chan=channel), level=log_level), on_close=self._on_socket_closed
+            forward_to_log(loggers, chan=channel, level=log_level), on_close=self._on_socket_closed
         )
 
     def _on_socket_closed(self):
@@ -746,7 +754,7 @@ def _upload_logs(self):
             if self._what
             else {}
         )
-        upload_to_remote(self.log, log_meta_dict)
+        upload_to_remote(self.process_log, log_meta_dict)
 
     def _monitor_subprocess(self):
         """
@@ -976,7 +984,9 @@ def cb(sock: socket):
     return cb
 
 
-def process_log_messages_from_subprocess(log: FilteringBoundLogger) -> Generator[None, bytes, None]:
+def process_log_messages_from_subprocess(
+    loggers: tuple[FilteringBoundLogger, ...],
+) -> Generator[None, bytes, None]:
     from structlog.stdlib import NAME_TO_LEVEL
 
     while True:
@@ -1003,21 +1013,27 @@ def process_log_messages_from_subprocess(log: FilteringBoundLogger) -> Generator
         if exc := event.pop("exception", None):
             # TODO: convert the dict back to a pretty stack trace
             event["error_detail"] = exc
-        log.log(NAME_TO_LEVEL[event.pop("level")], event.pop("event", None), **event)
+
+        level = NAME_TO_LEVEL[event.pop("level")]
+        msg = event.pop("event", None)
+        for target in loggers:
+            target.log(level, msg, **event)
 
 
-def forward_to_log(target_log: FilteringBoundLogger, level: int) -> Generator[None, bytes, None]:
+def forward_to_log(
+    target_loggers: tuple[FilteringBoundLogger, ...], chan: str, level: int
+) -> Generator[None, bytes, None]:
     while True:
         buf = yield
         line = bytes(buf)
         # Strip off new line
         line = line.rstrip()
         try:
             msg = line.decode("utf-8", errors="replace")
-            target_log.log(level, msg)
         except UnicodeDecodeError:
             msg = line.decode("ascii", errors="replace")
-            target_log.log(level, msg)
+        for log in target_loggers:
+            log.log(level, msg, chan=chan)
 
 
 def supervise(
@@ -1029,6 +1045,7 @@ def supervise(
     server: str | None = None,
     dry_run: bool = False,
     log_path: str | None = None,
+    subprocess_logs_to_stdout: bool = False,
     client: Client | None = None,
 ) -> int:
     """
@@ -1041,6 +1058,7 @@ def supervise(
     :param server: Base URL of the API server.
     :param dry_run: If True, execute without actual task execution (simulate run).
     :param log_path: Path to write logs, if required.
+    :param subprocess_logs_to_stdout: Should task logs also be sent to stdout via the main logger.
     :param client: Optional preconfigured client for communication with the server (Mostly for tests).
     :return: Exit code of the process.
     """
@@ -1081,6 +1099,7 @@ def supervise(
         client=client,
         logger=logger,
         bundle_info=bundle_info,
+        subprocess_logs_to_stdout=subprocess_logs_to_stdout,
     )
 
     exit_code = process.wait()

task-sdk/tests/task_sdk/execution_time/test_supervisor.py

@@ -516,7 +516,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(),
+            process_log=mocker.MagicMock(),
             id=TI_ID,
             pid=mock_process.pid,
             stdin=mocker.MagicMock(),
@@ -606,7 +606,7 @@ def test_overtime_handling(
         monkeypatch.setattr(ActivitySubprocess, "TASK_OVERTIME_THRESHOLD", overtime_threshold)
 
         mock_watched_subprocess = ActivitySubprocess(
-            log=mocker.MagicMock(),
+            process_log=mocker.MagicMock(),
             id=TI_ID,
             pid=12345,
             stdin=mocker.Mock(),
@@ -751,7 +751,7 @@ def mock_process(self, mocker):
     @pytest.fixture
     def watched_subprocess(self, mocker, mock_process):
         proc = ActivitySubprocess(
-            log=mocker.MagicMock(),
+            process_log=mocker.MagicMock(),
             id=TI_ID,
             pid=12345,
             stdin=mocker.Mock(),
@@ -937,7 +937,7 @@ class TestHandleRequest:
     def watched_subprocess(self, mocker):
         """Fixture to provide a WatchedSubprocess instance."""
         return ActivitySubprocess(
-            log=mocker.MagicMock(),
+            process_log=mocker.MagicMock(),
             id=TI_ID,
             pid=12345,
             stdin=BytesIO(),

tests/dag_processing/test_manager.py

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

tests/jobs/test_triggerer_job.py

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