Propagate logs to stdout when in k8s executor pod

airflow/cli/commands/task_command.py

@@ -23,6 +23,7 @@
 import json
 import logging
 import os
+import sys
 import textwrap
 from contextlib import contextmanager, redirect_stderr, redirect_stdout, suppress
 from typing import Generator, Union
@@ -43,6 +44,7 @@
 from airflow.models.dag import DAG
 from airflow.models.dagrun import DagRun
 from airflow.models.operator import needs_expansion
+from airflow.settings import IS_K8S_EXECUTOR_POD
 from airflow.ti_deps.dep_context import DepContext
 from airflow.ti_deps.dependencies_deps import SCHEDULER_QUEUED_DEPS
 from airflow.typing_compat import Literal
@@ -283,38 +285,52 @@ def _extract_external_executor_id(args) -> str | None:
 
 
 @contextmanager
-def _capture_task_logs(ti: TaskInstance) -> Generator[None, None, None]:
+def _move_task_handlers_to_root(ti: TaskInstance) -> Generator[None, None, None]:
     """
-    Manage logging context for a task run.
+    Move handlers for task logging to root logger.
 
-    - Replace the root logger configuration with the airflow.task configuration
-      so we can capture logs from any custom loggers used in the task.
+    We want anything logged during task run to be propagated to task log handlers.
+    If running in a k8s executor pod, also keep the stream handler on root logger
+    so that logs are still emitted to stdout.
+    """
+    # nothing to do
+    if not ti.log.handlers or settings.DONOT_MODIFY_HANDLERS:
+        yield
+        return
+
+    # Move task handlers to root and reset task logger and restore original logger settings after exit.
+    # If k8s executor, we need to ensure that root logger has a console handler, so that
+    # task logs propagate to stdout (this is how webserver retrieves them while task is running).
+    root_logger = logging.getLogger()
+    console_handler = next((h for h in root_logger.handlers if h.name == "console"), None)
+    with LoggerMutationHelper(root_logger), LoggerMutationHelper(ti.log) as task_helper:
+        task_helper.move(root_logger)
+        if IS_K8S_EXECUTOR_POD:
+            if console_handler and console_handler not in root_logger.handlers:
+                root_logger.addHandler(console_handler)
+        yield
 
-    - Redirect stdout and stderr to the task instance log, as INFO and WARNING
-      level messages, respectively.
 
+@contextmanager
+def _redirect_stdout_to_ti_log(ti: TaskInstance) -> Generator[None, None, None]:
     """
-    modify = not settings.DONOT_MODIFY_HANDLERS
-    if modify:
-        root_logger, task_logger = logging.getLogger(), logging.getLogger("airflow.task")
+    Redirect stdout to ti logger.
 
-        orig_level = root_logger.level
-        root_logger.setLevel(task_logger.level)
-        orig_handlers = root_logger.handlers.copy()
-        root_logger.handlers[:] = task_logger.handlers
+    Redirect stdout and stderr to the task instance log as INFO and WARNING
+    level messages, respectively.
 
-    try:
+    If stdout already redirected (possible when task running with option
+    `--local`), don't redirect again.
+    """
+    # if sys.stdout is StreamLogWriter, it means we already redirected
+    # likely before forking in LocalTaskJob
+    if not isinstance(sys.stdout, StreamLogWriter):
         info_writer = StreamLogWriter(ti.log, logging.INFO)
         warning_writer = StreamLogWriter(ti.log, logging.WARNING)
-
         with redirect_stdout(info_writer), redirect_stderr(warning_writer):
             yield
-
-    finally:
-        if modify:
-            # Restore the root logger to its original state.
-            root_logger.setLevel(orig_level)
-            root_logger.handlers[:] = orig_handlers
+    else:
+        yield
 
 
 class TaskCommandMarker:
@@ -366,12 +382,6 @@ def task_run(args, dag=None):
 
     settings.MASK_SECRETS_IN_LOGS = True
 
-    # IMPORTANT, have to re-configure ORM with the NullPool, otherwise, each "run" command may leave
-    # behind multiple open sleeping connections while heartbeating, which could
-    # easily exceed the database connection limit when
-    # processing hundreds of simultaneous tasks.
-    settings.reconfigure_orm(disable_connection_pool=True)
-
     get_listener_manager().hook.on_starting(component=TaskCommandMarker())
 
     if args.pickle:
@@ -390,11 +400,19 @@ def task_run(args, dag=None):
 
     log.info("Running %s on host %s", ti, hostname)
 
+    # IMPORTANT, have to re-configure ORM with the NullPool, otherwise, each "run" command may leave
+    # behind multiple open sleeping connections while heartbeating, which could
+    # easily exceed the database connection limit when
+    # processing hundreds of simultaneous tasks.
+    # this should be last thing before running, to reduce likelihood of an open session
+    # which can cause trouble if running process in a fork.
+    settings.reconfigure_orm(disable_connection_pool=True)
+
     try:
         if args.interactive:
             _run_task_by_selected_method(args, dag, ti)
         else:
-            with _capture_task_logs(ti):
+            with _move_task_handlers_to_root(ti), _redirect_stdout_to_ti_log(ti):
                 _run_task_by_selected_method(args, dag, ti)
     finally:
         try:
@@ -643,3 +661,53 @@ def task_clear(args):
         include_subdags=not args.exclude_subdags,
         include_parentdag=not args.exclude_parentdag,
     )
+
+
+class LoggerMutationHelper:
+    """
+    Helper for moving and resetting handlers and other logger attrs.
+
+    :meta private:
+    """
+
+    def __init__(self, logger):
+        self.handlers = logger.handlers[:]
+        self.level = logger.level
+        self.propagate = logger.propagate
+        self.source_logger = logger
+
+    def apply(self, logger, replace=True):
+        """
+        Set ``logger`` with attrs stored on instance.
+
+        If ``logger`` is root logger, don't change propagate.
+        """
+        if replace:
+            logger.handlers[:] = self.handlers
+        else:
+            for h in self.handlers:
+                if h not in logger.handlers:
+                    logger.addHandler(h)
+        logger.level = self.level
+        if logger is not logging.getLogger():
+            logger.propagate = self.propagate
+
+    def move(self, logger, replace=True):
+        """
+        Replace ``logger`` attrs with those from source.
+
+        :param logger: target logger
+        :param replace: if True, remove all handlers from target first; otherwise add if not present.
+        """
+        self.apply(logger, replace=replace)
+        self.source_logger.propagate = True
+        self.source_logger.handlers[:] = []
+
+    def reset(self):
+        self.apply(self.source_logger)
+
+    def __enter__(self):
+        return self
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        self.reset()

airflow/settings.py

@@ -624,6 +624,8 @@ def initialize():
     executor_constants.CELERY_KUBERNETES_EXECUTOR,
     executor_constants.LOCAL_KUBERNETES_EXECUTOR,
 }
+IS_K8S_EXECUTOR_POD = bool(os.environ.get("AIRFLOW_IS_K8S_EXECUTOR_POD", ""))
+"""Will be True if running in kubernetes executor pod."""
 
 HIDE_SENSITIVE_VAR_CONN_FIELDS = conf.getboolean("core", "hide_sensitive_var_conn_fields")
 

airflow/utils/log/file_task_handler.py

@@ -30,13 +30,13 @@
 from airflow.exceptions import RemovedInAirflow3Warning
 from airflow.utils.context import Context
 from airflow.utils.helpers import parse_template_string, render_template_to_string
+from airflow.utils.log.logging_mixin import SetContextPropagate
 from airflow.utils.log.non_caching_file_handler import NonCachingFileHandler
 from airflow.utils.session import create_session
 from airflow.utils.state import State
 
 if TYPE_CHECKING:
     from airflow.models import TaskInstance
-    from airflow.utils.log.logging_mixin import SetContextPropagate
 
 
 class FileTaskHandler(logging.Handler):
@@ -62,19 +62,29 @@ def __init__(self, base_log_folder: str, filename_template: str | None = None):
                 # handler, not the one that calls super()__init__.
                 stacklevel=(2 if type(self) == FileTaskHandler else 3),
             )
+        self.maintain_propagate: bool = False
+        """If true, overrides default behavior of setting propagate=False
+        :meta private:
+        """
 
     def set_context(self, ti: TaskInstance) -> None | SetContextPropagate:
         """
         Provide task_instance context to airflow task handler.
 
+        Generally speaking returns None.  But if attr `maintain_propagate` has
+        been set to propagate, then returns sentinel MAINTAIN_PROPAGATE. This
+        has the effect of overriding the default behavior to set `propagate`
+        to False whenever set_context is called.  At time of writing, this
+        functionality is only used in unit testing.
+
         :param ti: task instance object
         """
         local_loc = self._init_file(ti)
         self.handler = NonCachingFileHandler(local_loc, encoding="utf-8")
         if self.formatter:
             self.handler.setFormatter(self.formatter)
         self.handler.setLevel(self.level)
-        return None
+        return SetContextPropagate.MAINTAIN_PROPAGATE if self.maintain_propagate else None
 
     def emit(self, record):
         if self.handler:
@@ -92,16 +102,17 @@ def _render_filename(self, ti: TaskInstance, try_number: int) -> str:
         with create_session() as session:
             dag_run = ti.get_dagrun(session=session)
             template = dag_run.get_log_template(session=session).filename
-        str_tpl, jinja_tpl = parse_template_string(template)
+            str_tpl, jinja_tpl = parse_template_string(template)
 
-        if jinja_tpl:
-            if hasattr(ti, "task"):
-                context = ti.get_template_context()
-            else:
-                context = Context(ti=ti, ts=dag_run.logical_date.isoformat())
-            context["try_number"] = try_number
-            return render_template_to_string(jinja_tpl, context)
-        elif str_tpl:
+            if jinja_tpl:
+                if hasattr(ti, "task"):
+                    context = ti.get_template_context(session=session)
+                else:
+                    context = Context(ti=ti, ts=dag_run.logical_date.isoformat())
+                context["try_number"] = try_number
+                return render_template_to_string(jinja_tpl, context)
+
+        if str_tpl:
             try:
                 dag = ti.task.dag
             except AttributeError:  # ti.task is not always set.

airflow/utils/log/logging_mixin.py

@@ -26,6 +26,8 @@
 from logging import Handler, Logger, StreamHandler
 from typing import IO, cast
 
+from airflow.settings import IS_K8S_EXECUTOR_POD
+
 # 7-bit C1 ANSI escape sequences
 ANSI_ESCAPE = re.compile(r"\x1B[@-_][0-?]*[ -/]*[@-~]")
 
@@ -165,9 +167,10 @@ def isatty(self):
 
 class RedirectStdHandler(StreamHandler):
     """
-    This class is like a StreamHandler using sys.stderr/stdout, but always uses
+    This class is like a StreamHandler using sys.stderr/stdout, but uses
     whatever sys.stderr/stderr is currently set to rather than the value of
-    sys.stderr/stdout at handler construction time.
+    sys.stderr/stdout at handler construction time, except when running a
+    task in a kubernetes executor pod.
     """
 
     def __init__(self, stream):
@@ -179,13 +182,17 @@ def __init__(self, stream):
         self._use_stderr = True
         if "stdout" in stream:
             self._use_stderr = False
-
+            self._orig_stream = sys.stdout
+        else:
+            self._orig_stream = sys.stderr
         # StreamHandler tries to set self.stream
         Handler.__init__(self)
 
     @property
     def stream(self):
         """Returns current stream."""
+        if IS_K8S_EXECUTOR_POD:
+            return self._orig_stream
         if self._use_stderr:
             return sys.stderr