Truncate stack trace to DAG user code for exceptions raised during execution (#20731)

(cherry picked from commit 7ea0f76)

Author: malthe

airflow/models/taskinstance.py

@@ -26,9 +26,12 @@
 from collections import defaultdict
 from datetime import datetime, timedelta
 from functools import partial
+from inspect import currentframe
 from tempfile import NamedTemporaryFile
+from types import FrameType
 from typing import IO, TYPE_CHECKING, Any, Iterable, List, NamedTuple, Optional, Tuple, Union
 from urllib.parse import quote
+from weakref import WeakKeyDictionary
 
 import dill
 import jinja2
@@ -115,6 +118,8 @@
     from airflow.models.baseoperator import BaseOperator
     from airflow.models.dag import DAG, DagModel, DagRun
 
+_EXECUTION_FRAME_MAPPING: "WeakKeyDictionary[Operator, FrameType]" = WeakKeyDictionary()
+
 
 @contextlib.contextmanager
 def set_current_context(context: Context) -> None:
@@ -1409,7 +1414,18 @@ def _execute_task_with_callbacks(self, context):
         """Prepare Task for Execution"""
         from airflow.models.renderedtifields import RenderedTaskInstanceFields
 
+        parent_pid = os.getpid()
+
         def signal_handler(signum, frame):
+            pid = os.getpid()
+
+            # If a task forks during execution (from DAG code) for whatever
+            # reason, we want to make sure that we react to the signal only in
+            # the process that we've spawned ourselves (referred to here as the
+            # parent process).
+            if pid != parent_pid:
+                os._exit(1)
+                return
             self.log.error("Received SIGTERM. Terminating subprocesses.")
             self.task.on_kill()
             raise AirflowException("Task received SIGTERM signal")
@@ -1480,7 +1496,6 @@ def _execute_task(self, context, task_copy):
         # If the task has been deferred and is being executed due to a trigger,
         # then we need to pick the right method to come back to, otherwise
         # we go for the default execute
-        execute_callable = task_copy.execute
         if self.next_method:
             # __fail__ is a special signal value for next_method that indicates
             # this task was scheduled specifically to fail.
@@ -1494,29 +1509,35 @@ def _execute_task(self, context, task_copy):
             execute_callable = getattr(task_copy, self.next_method)
             if self.next_kwargs:
                 execute_callable = partial(execute_callable, **self.next_kwargs)
+        else:
+            execute_callable = task_copy.execute
         # If a timeout is specified for the task, make it fail
         # if it goes beyond
-        if task_copy.execution_timeout:
-            # If we are coming in with a next_method (i.e. from a deferral),
-            # calculate the timeout from our start_date.
-            if self.next_method:
-                timeout_seconds = (
-                    task_copy.execution_timeout - (timezone.utcnow() - self.start_date)
-                ).total_seconds()
+        try:
+            if task_copy.execution_timeout:
+                # If we are coming in with a next_method (i.e. from a deferral),
+                # calculate the timeout from our start_date.
+                if self.next_method:
+                    timeout_seconds = (
+                        task_copy.execution_timeout - (timezone.utcnow() - self.start_date)
+                    ).total_seconds()
+                else:
+                    timeout_seconds = task_copy.execution_timeout.total_seconds()
+                try:
+                    # It's possible we're already timed out, so fast-fail if true
+                    if timeout_seconds <= 0:
+                        raise AirflowTaskTimeout()
+                    # Run task in timeout wrapper
+                    with timeout(timeout_seconds):
+                        result = execute_callable(context=context)
+                except AirflowTaskTimeout:
+                    task_copy.on_kill()
+                    raise
             else:
-                timeout_seconds = task_copy.execution_timeout.total_seconds()
-            try:
-                # It's possible we're already timed out, so fast-fail if true
-                if timeout_seconds <= 0:
-                    raise AirflowTaskTimeout()
-                # Run task in timeout wrapper
-                with timeout(timeout_seconds):
-                    result = execute_callable(context=context)
-            except AirflowTaskTimeout:
-                task_copy.on_kill()
-                raise
-        else:
-            result = execute_callable(context=context)
+                result = execute_callable(context=context)
+        except:  # noqa: E722
+            _EXECUTION_FRAME_MAPPING[task_copy] = currentframe()
+            raise
         # If the task returns a result, push an XCom containing it
         if task_copy.do_xcom_push and result is not None:
             self.xcom_push(key=XCOM_RETURN_KEY, value=result)
@@ -1718,7 +1739,15 @@ def handle_failure(
 
         if error:
             if isinstance(error, Exception):
-                self.log.error("Task failed with exception", exc_info=error)
+                execution_frame = _EXECUTION_FRAME_MAPPING.get(self.task)
+                tb = error.__traceback__
+                while tb is not None:
+                    if tb.tb_frame is execution_frame:
+                        tb = tb.tb_next
+                        break
+                    tb = tb.tb_next
+                tb = tb or error.__traceback__
+                self.log.error("Task failed with exception", exc_info=(type(error), error, tb))
             else:
                 self.log.error("%s", error)
             # external monitoring process provides pickle file so _run_raw_task

airflow/task/task_runner/standard_task_runner.py

@@ -73,23 +73,29 @@ def _start_by_fork(self):
             # [1:] - remove "airflow" from the start of the command
             args = parser.parse_args(self._command[1:])
 
+            # We prefer the job_id passed on the command-line because at this time, the
+            # task instance may not have been updated.
+            job_id = getattr(args, "job_id", self._task_instance.job_id)
             self.log.info('Running: %s', self._command)
-            self.log.info('Job %s: Subtask %s', self._task_instance.job_id, self._task_instance.task_id)
+            self.log.info('Job %s: Subtask %s', job_id, self._task_instance.task_id)
 
             proc_title = "airflow task runner: {0.dag_id} {0.task_id} {0.execution_date_or_run_id}"
-            if hasattr(args, "job_id"):
+            if job_id is not None:
                 proc_title += " {0.job_id}"
             setproctitle(proc_title.format(args))
 
             try:
                 args.func(args, dag=self.dag)
                 return_code = 0
-            except Exception:
+            except Exception as exc:
                 return_code = 1
-                self.log.exception(
-                    "Failed to execute job %s for task %s",
-                    self._task_instance.job_id,
+
+                self.log.error(
+                    "Failed to execute job %s for task %s (%s; %r)",
+                    job_id,
                     self._task_instance.task_id,
+                    exc,
+                    os.getpid(),
                 )
             finally:
                 # Explicitly flush any pending exception to Sentry if enabled

tests/models/test_taskinstance.py

@@ -19,8 +19,10 @@
 import datetime
 import os
 import signal
+import sys
 import urllib
 from tempfile import NamedTemporaryFile
+from traceback import format_exception
 from typing import List, Optional, Union, cast
 from unittest import mock
 from unittest.mock import call, mock_open, patch
@@ -1843,6 +1845,27 @@ def fail():
             pass  # expected
         assert State.UP_FOR_RETRY == ti.state
 
+    def test_stacktrace_on_failure_starts_with_task_execute_method(self, dag_maker):
+        def fail():
+            raise AirflowException("maybe this will pass?")
+
+        with dag_maker(dag_id='test_retries_on_other_exceptions'):
+            task = PythonOperator(
+                task_id='test_raise_other_exception',
+                python_callable=fail,
+                retries=1,
+            )
+        ti = dag_maker.create_dagrun(execution_date=timezone.utcnow()).task_instances[0]
+        ti.task = task
+        with patch.object(TI, "log") as log, pytest.raises(AirflowException):
+            ti.run()
+        assert len(log.error.mock_calls) == 1
+        assert log.error.call_args[0] == ("Task failed with exception",)
+        exc_info = log.error.call_args[1]["exc_info"]
+        filename = exc_info[2].tb_frame.f_code.co_filename
+        formatted_exc = format_exception(*exc_info)
+        assert sys.modules[PythonOperator.__module__].__file__ == filename, "".join(formatted_exc)
+
     def _env_var_check_callback(self):
         assert 'test_echo_env_variables' == os.environ['AIRFLOW_CTX_DAG_ID']
         assert 'hive_in_python_op' == os.environ['AIRFLOW_CTX_TASK_ID']

tests/task/task_runner/test_standard_task_runner.py

@@ -17,8 +17,10 @@
 # under the License.
 import logging
 import os
+import re
 import time
 from logging.config import dictConfig
+from tempfile import NamedTemporaryFile
 from unittest import mock
 
 import psutil
@@ -40,11 +42,13 @@
 
 DEFAULT_DATE = timezone.datetime(2016, 1, 1)
 
+TASK_FORMAT = '{{%(filename)s:%(lineno)d}} %(levelname)s - %(message)s'
+
 LOGGING_CONFIG = {
     'version': 1,
     'disable_existing_loggers': False,
     'formatters': {
-        'airflow.task': {'format': '[%(asctime)s] {{%(filename)s:%(lineno)d}} %(levelname)s - %(message)s'},
+        'airflow.task': {'format': TASK_FORMAT},
     },
     'handlers': {
         'console': {
@@ -197,19 +201,23 @@ def test_on_kill(self):
         dag = dagbag.dags.get('test_on_kill')
         task = dag.get_task('task1')
 
-        with create_session() as session:
+        with create_session() as session, NamedTemporaryFile("w", delete=False) as f:
             dag.create_dagrun(
                 run_id="test",
+                data_interval=(DEFAULT_DATE, DEFAULT_DATE),
                 state=State.RUNNING,
-                execution_date=DEFAULT_DATE,
                 start_date=DEFAULT_DATE,
                 session=session,
             )
-            ti = TaskInstance(task=task, execution_date=DEFAULT_DATE)
+            ti = TaskInstance(task=task, run_id="test")
             job1 = LocalTaskJob(task_instance=ti, ignore_ti_state=True)
             session.commit()
+            ti.refresh_from_task(task)
 
             runner = StandardTaskRunner(job1)
+            handler = logging.StreamHandler(f)
+            handler.setFormatter(logging.Formatter(TASK_FORMAT))
+            runner.log.addHandler(handler)
             runner.start()
 
             with timeout(seconds=3):
@@ -232,6 +240,14 @@ def test_on_kill(self):
             logging.info(f"Terminating processes {processes} belonging to {runner_pgid} group")
             runner.terminate()
             session.close()  # explicitly close as `create_session`s commit will blow up otherwise
+            with open(f.name) as g:
+                logged = g.read()
+            os.unlink(f.name)
+
+        ti.refresh_from_db()
+        assert re.findall(r'ERROR - Failed to execute job (\S+) for task (\S+)', logged) == [
+            (str(ti.job_id), ti.task_id)
+        ], logged
 
         logging.info("Waiting for the on kill killed file to appear")
         with timeout(seconds=4):