Enable trigger logging in webserver

airflow/api_connexion/endpoints/log_endpoint.py

@@ -21,6 +21,7 @@
 from flask import Response, request
 from itsdangerous.exc import BadSignature
 from itsdangerous.url_safe import URLSafeSerializer
+from sqlalchemy.orm import joinedload
 from sqlalchemy.orm.session import Session
 
 from airflow.api_connexion import security
@@ -73,9 +74,10 @@ def get_log(
         metadata["download_logs"] = False
 
     task_log_reader = TaskLogReader()
+
     if not task_log_reader.supports_read:
         raise BadRequest("Task log handler does not support read logs.")
-    ti = (
+    query = (
         session.query(TaskInstance)
         .filter(
             TaskInstance.task_id == task_id,
@@ -84,8 +86,10 @@ def get_log(
             TaskInstance.map_index == map_index,
         )
         .join(TaskInstance.dag_run)
-        .one_or_none()
+        .options(joinedload("trigger"))
+        .options(joinedload("trigger.triggerer_job"))
     )
+    ti = query.one_or_none()
     if ti is None:
         metadata["end_of_log"] = True
         raise NotFound(title="TaskInstance not found")

airflow/cli/cli_parser.py

@@ -2072,6 +2072,7 @@ class GroupCommand(NamedTuple):
             ARG_LOG_FILE,
             ARG_CAPACITY,
             ARG_VERBOSE,
+            ARG_SKIP_SERVE_LOGS,
         ),
     ),
     ActionCommand(

airflow/cli/commands/task_command.py

@@ -44,6 +44,7 @@
 from airflow.models.dag import DAG
 from airflow.models.dagrun import DagRun
 from airflow.models.operator import needs_expansion
+from airflow.models.taskinstance import TaskReturnCode
 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
@@ -58,6 +59,7 @@
     suppress_logs_and_warning,
 )
 from airflow.utils.dates import timezone
+from airflow.utils.log.file_task_handler import _set_task_deferred_context_var
 from airflow.utils.log.logging_mixin import StreamLogWriter
 from airflow.utils.log.secrets_masker import RedactedIO
 from airflow.utils.net import get_hostname
@@ -182,7 +184,7 @@ def _get_ti(
     return ti, dr_created
 
 
-def _run_task_by_selected_method(args, dag: DAG, ti: TaskInstance) -> None:
+def _run_task_by_selected_method(args, dag: DAG, ti: TaskInstance) -> None | TaskReturnCode:
     """
     Runs the task based on a mode.
 
@@ -193,11 +195,11 @@ def _run_task_by_selected_method(args, dag: DAG, ti: TaskInstance) -> None:
     - by executor
     """
     if args.local:
-        _run_task_by_local_task_job(args, ti)
+        return _run_task_by_local_task_job(args, ti)
     elif args.raw:
-        _run_raw_task(args, ti)
+        return _run_raw_task(args, ti)
     else:
-        _run_task_by_executor(args, dag, ti)
+        return _run_task_by_executor(args, dag, ti)
 
 
 def _run_task_by_executor(args, dag, ti):
@@ -239,7 +241,7 @@ def _run_task_by_executor(args, dag, ti):
     executor.end()
 
 
-def _run_task_by_local_task_job(args, ti):
+def _run_task_by_local_task_job(args, ti) -> TaskReturnCode | None:
     """Run LocalTaskJob, which monitors the raw task execution process."""
     run_job = LocalTaskJob(
         task_instance=ti,
@@ -254,11 +256,14 @@ def _run_task_by_local_task_job(args, ti):
         external_executor_id=_extract_external_executor_id(args),
     )
     try:
-        run_job.run()
+        ret = run_job.run()
 
     finally:
         if args.shut_down_logging:
             logging.shutdown()
+    with suppress(ValueError):
+        return TaskReturnCode(ret)
+    return None
 
 
 RAW_TASK_UNSUPPORTED_OPTION = [
@@ -269,9 +274,9 @@ def _run_task_by_local_task_job(args, ti):
 ]
 
 
-def _run_raw_task(args, ti: TaskInstance) -> None:
+def _run_raw_task(args, ti: TaskInstance) -> None | TaskReturnCode:
     """Runs the main task handling code."""
-    ti._run_raw_task(
+    return ti._run_raw_task(
         mark_success=args.mark_success,
         job_id=args.job_id,
         pool=args.pool,
@@ -407,18 +412,21 @@ def task_run(args, dag=None):
     # 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)
-
+    task_return_code = None
     try:
         if args.interactive:
-            _run_task_by_selected_method(args, dag, ti)
+            task_return_code = _run_task_by_selected_method(args, dag, ti)
         else:
             with _move_task_handlers_to_root(ti), _redirect_stdout_to_ti_log(ti):
-                _run_task_by_selected_method(args, dag, ti)
+                task_return_code = _run_task_by_selected_method(args, dag, ti)
+                if task_return_code == TaskReturnCode.DEFERRED:
+                    _set_task_deferred_context_var()
     finally:
         try:
             get_listener_manager().hook.before_stopping(component=TaskCommandMarker())
         except Exception:
             pass
+    return task_return_code
 
 
 @cli_utils.action_cli(check_db=False)

airflow/cli/commands/triggerer_command.py

@@ -18,14 +18,35 @@
 from __future__ import annotations
 
 import signal
+from contextlib import contextmanager
+from functools import partial
+from multiprocessing import Process
+from typing import Generator
 
 import daemon
 from daemon.pidfile import TimeoutPIDLockFile
 
 from airflow import settings
+from airflow.configuration import conf
 from airflow.jobs.triggerer_job import TriggererJob
 from airflow.utils import cli as cli_utils
 from airflow.utils.cli import setup_locations, setup_logging, sigint_handler, sigquit_handler
+from airflow.utils.serve_logs import serve_logs
+
+
+@contextmanager
+def _serve_logs(skip_serve_logs: bool = False) -> Generator[None, None, None]:
+    """Starts serve_logs sub-process"""
+    sub_proc = None
+    if skip_serve_logs is False:
+        port = conf.getint("logging", "trigger_log_server_port", fallback=8794)
+        sub_proc = Process(target=partial(serve_logs, port=port))
+        sub_proc.start()
+    try:
+        yield
+    finally:
+        if sub_proc:
+            sub_proc.terminate()
 
 
 @cli_utils.action_cli
@@ -44,18 +65,18 @@ def triggerer(args):
             stdout_handle.truncate(0)
             stderr_handle.truncate(0)
 
-            ctx = daemon.DaemonContext(
+            daemon_context = daemon.DaemonContext(
                 pidfile=TimeoutPIDLockFile(pid, -1),
                 files_preserve=[handle],
                 stdout=stdout_handle,
                 stderr=stderr_handle,
                 umask=int(settings.DAEMON_UMASK, 8),
             )
-            with ctx:
+            with daemon_context, _serve_logs(args.skip_serve_logs):
                 job.run()
-
     else:
         signal.signal(signal.SIGINT, sigint_handler)
         signal.signal(signal.SIGTERM, sigint_handler)
         signal.signal(signal.SIGQUIT, sigquit_handler)
-        job.run()
+        with _serve_logs(args.skip_serve_logs):
+            job.run()

airflow/config_templates/config.yml

@@ -788,6 +788,24 @@ logging:
       type: string
       example: ~
       default: "8793"
+    trigger_log_server_port:
+      description: |
+        Port to serve logs from for triggerer.  See worker_log_server_port description
+        for more info.
+      version_added: 2.6.0
+      type: string
+      example: ~
+      default: "8794"
+    interleave_timestamp_parser:
+      description: |
+        We must parse timestamps to interleave logs between trigger and task.  To do so,
+        we need to parse timestamps in log files. In case your log format is non-standard,
+        you may provide import path to callable which takes a string log line and returns
+        the timestamp (datetime.datetime compatible).
+      version_added: 2.6.0
+      type: string
+      example: path.to.my_func
+      default: ~
 metrics:
   description: |
     StatsD (https://github.com/etsy/statsd) integration settings.

airflow/config_templates/default_airflow.cfg

@@ -434,6 +434,17 @@ extra_logger_names =
 # visible from the main web server to connect into the workers.
 worker_log_server_port = 8793
 
+# Port to serve logs from for triggerer.  See worker_log_server_port description
+# for more info.
+trigger_log_server_port = 8794
+
+# We must parse timestamps to interleave logs between trigger and task.  To do so,
+# we need to parse timestamps in log files. In case your log format is non-standard,
+# you may provide import path to callable which takes a string log line and returns
+# the timestamp (datetime.datetime compatible).
+# Example: interleave_timestamp_parser = path.to.my_func
+# interleave_timestamp_parser =
+
 [metrics]
 
 # StatsD (https://github.com/etsy/statsd) integration settings.

airflow/example_dags/example_time_delta_sensor_async.py

@@ -36,6 +36,6 @@
     catchup=False,
     tags=["example"],
 ) as dag:
-    wait = TimeDeltaSensorAsync(task_id="wait", delta=datetime.timedelta(seconds=10))
+    wait = TimeDeltaSensorAsync(task_id="wait", delta=datetime.timedelta(seconds=30))
     finish = EmptyOperator(task_id="finish")
     wait >> finish

airflow/executors/base_executor.py

@@ -355,14 +355,15 @@ def execute_async(
         """
         raise NotImplementedError()
 
-    def get_task_log(self, ti: TaskInstance, log: str = "") -> None | str | tuple[str, dict[str, bool]]:
+    def get_task_log(self, ti: TaskInstance) -> tuple[list[str], list[str]]:
         """
         This method can be implemented by any child class to return the task logs.
 
         :param ti: A TaskInstance object
         :param log: log str
         :return: logs or tuple of logs and meta dict
         """
+        return [], []
 
     def end(self) -> None:  # pragma: no cover
         """Wait synchronously for the previously submitted job to complete."""

airflow/executors/celery_kubernetes_executor.py

@@ -141,11 +141,11 @@ def queue_task_instance(
             cfg_path=cfg_path,
         )
 
-    def get_task_log(self, ti: TaskInstance, log: str = "") -> None | str | tuple[str, dict[str, bool]]:
+    def get_task_log(self, ti: TaskInstance) -> tuple[list[str], list[str]]:
         """Fetch task log from Kubernetes executor"""
         if ti.queue == self.kubernetes_executor.kubernetes_queue:
-            return self.kubernetes_executor.get_task_log(ti=ti, log=log)
-        return None
+            return self.kubernetes_executor.get_task_log(ti=ti)
+        return [], []
 
     def has_task(self, task_instance: TaskInstance) -> bool:
         """

airflow/executors/kubernetes_executor.py

@@ -781,14 +781,16 @@ def _get_pod_namespace(ti: TaskInstance):
             namespace = pod_override.metadata.namespace
         return namespace or conf.get("kubernetes_executor", "namespace", fallback="default")
 
-    def get_task_log(self, ti: TaskInstance, log: str = "") -> str | tuple[str, dict[str, bool]]:
-
+    def get_task_log(self, ti: TaskInstance) -> tuple[list[str], list[str]]:
+        messages = []
+        log = []
         try:
+            from airflow.kubernetes.kube_client import get_kube_client
             from airflow.kubernetes.pod_generator import PodGenerator
 
             client = get_kube_client()
 
-            log += f"*** Trying to get logs (last 100 lines) from worker pod {ti.hostname} ***\n\n"
+            messages.append(f"Trying to get logs (last 100 lines) from worker pod {ti.hostname}")
             selector = PodGenerator.build_selector_for_k8s_executor_pod(
                 dag_id=ti.dag_id,
                 task_id=ti.task_id,
@@ -816,13 +818,10 @@ def get_task_log(self, ti: TaskInstance, log: str = "") -> str | tuple[str, dict
             )
 
             for line in res:
-                log += line.decode()
-
-            return log
-
-        except Exception as f:
-            log += f"*** Unable to fetch logs from worker pod {ti.hostname} ***\n{str(f)}\n\n"
-            return log, {"end_of_log": True}
+                log.append(line.decode())
+        except Exception as e:
+            messages.append(f"Reading from k8s pod logs failed: {str(e)}")
+        return messages, ["\n".join(log)]
 
     def try_adopt_task_instances(self, tis: Sequence[TaskInstance]) -> Sequence[TaskInstance]:
         tis_to_flush = [ti for ti in tis if not ti.queued_by_job_id]

airflow/executors/local_kubernetes_executor.py

@@ -142,12 +142,11 @@ def queue_task_instance(
             cfg_path=cfg_path,
         )
 
-    def get_task_log(self, ti: TaskInstance, log: str = "") -> None | str | tuple[str, dict[str, bool]]:
+    def get_task_log(self, ti: TaskInstance) -> tuple[list[str], list[str]]:
         """Fetch task log from kubernetes executor"""
         if ti.queue == self.kubernetes_executor.kubernetes_queue:
-            return self.kubernetes_executor.get_task_log(ti=ti, log=log)
-
-        return None
+            return self.kubernetes_executor.get_task_log(ti=ti)
+        return [], []
 
     def has_task(self, task_instance: TaskInstance) -> bool:
         """

airflow/jobs/base_job.py

@@ -249,14 +249,15 @@ def run(self):
         """Starts the job."""
         Stats.incr(self.__class__.__name__.lower() + "_start", 1, 1)
         # Adding an entry in the DB
+        ret = None
         with create_session() as session:
             self.state = State.RUNNING
             session.add(self)
             session.commit()
             make_transient(self)
 
             try:
-                self._execute()
+                ret = self._execute()
                 # In case of max runs or max duration
                 self.state = State.SUCCESS
             except SystemExit:
@@ -272,6 +273,7 @@ def run(self):
                 session.commit()
 
         Stats.incr(self.__class__.__name__.lower() + "_end", 1, 1)
+        return ret
 
     def _execute(self):
         raise NotImplementedError("This method needs to be overridden")