ephraimbuddy commented on issue #30146: URL: https://github.com/apache/airflow/issues/30146#issuecomment-1473307192
On further debugging, I realized why the failure callback was not being called when tasks are terminated externally. It's because the state of the task instance was already marked as failed so it runs this: https://github.com/apache/airflow/blob/b65dbaaf3f21ea5396da121bbfa7f895d0ab8516/airflow/models/taskinstance.py#L1471-L1475 instead of https://github.com/apache/airflow/blob/b65dbaaf3f21ea5396da121bbfa7f895d0ab8516/airflow/models/taskinstance.py#L1476-L1479 We can use flag to signify when we have a sigterm but it needs testing to ensure the behaviour is desired due to this: > It seems incorrect to me that assuming SIGTERM is the same thing being marked failed externally. As SIGTERM can be received for other reason, like the pod-preemption when running on k8s, and for such cases we want to go through the normal failure handling code. Here's my implementation which you can test and give me feedback: ```diff commit d9b03d60323381d97e6df55d28c309ec1ae3962b Author: Ephraim Anierobi <[email protected]> Date: Fri Mar 17 08:21:01 2023 +0100 Fix failure callback when task is terminated externally The previous fix caused the callbacks to be called even on task retries. On debugging, I found that the AirflowException was actually caught but because the task instance was already in a failed state, the failure handler was not called. To fix this, I introduced a flag to signify that sigterm occurred so that it's handled properly. diff --git a/airflow/exceptions.py b/airflow/exceptions.py index e6ef9bd4e1..4bf946fd8e 100644 --- a/airflow/exceptions.py +++ b/airflow/exceptions.py @@ -29,12 +29,6 @@ if TYPE_CHECKING: from airflow.models import DagRun -class AirflowTermSignal(Exception): - """Raise when we receive a TERM signal""" - - status_code = HTTPStatus.INTERNAL_SERVER_ERROR - - class AirflowException(Exception): """ Base class for all Airflow's errors. diff --git a/airflow/models/taskinstance.py b/airflow/models/taskinstance.py index ffcdfb5651..ae65d434db 100644 --- a/airflow/models/taskinstance.py +++ b/airflow/models/taskinstance.py @@ -77,7 +77,6 @@ from airflow.exceptions import ( AirflowSensorTimeout, AirflowSkipException, AirflowTaskTimeout, - AirflowTermSignal, DagRunNotFound, RemovedInAirflow3Warning, TaskDeferralError, @@ -539,6 +538,8 @@ class TaskInstance(Base, LoggingMixin): self.raw = False # can be changed when calling 'run' self.test_mode = False + # signify task received sigterm + self.received_sigterm = False @staticmethod def insert_mapping(run_id: str, task: Operator, map_index: int) -> dict[str, Any]: @@ -567,6 +568,7 @@ class TaskInstance(Base, LoggingMixin): # correctly config the ti log self._log = logging.getLogger("airflow.task") self.test_mode = False # can be changed when calling 'run' + self.received_sigterm = False # signify task received sigterm @hybrid_property def try_number(self): @@ -1468,7 +1470,7 @@ class TaskInstance(Base, LoggingMixin): # for case when task is marked as success/failed externally # or dagrun timed out and task is marked as skipped # current behavior doesn't hit the callbacks - if self.state in State.finished: + if self.state in State.finished and not self.received_sigterm: self.clear_next_method_args() session.merge(self) session.commit() @@ -1535,7 +1537,9 @@ class TaskInstance(Base, LoggingMixin): os._exit(1) return self.log.error("Received SIGTERM. Terminating subprocesses.") - raise AirflowTermSignal("Task received SIGTERM signal") + self.task.on_kill() + self.received_sigterm = True + raise AirflowException("Task received SIGTERM signal") signal.signal(signal.SIGTERM, signal_handler) @@ -1574,15 +1578,9 @@ class TaskInstance(Base, LoggingMixin): # Execute the task with set_current_context(context): - try: - result = self._execute_task(context, task_orig) - # Run post_execute callback - self.task.post_execute(context=context, result=result) - except AirflowTermSignal: - self.task.on_kill() - if self.task.on_failure_callback: - self._run_finished_callback(self.task.on_failure_callback, context, "on_failure") - raise AirflowException("Task received SIGTERM signal") + result = self._execute_task(context, task_orig) + # Run post_execute callback + self.task.post_execute(context=context, result=result) Stats.incr(f"operator_successes_{self.task.task_type}", 1, 1) Stats.incr( diff --git a/tests/models/test_taskinstance.py b/tests/models/test_taskinstance.py index 50cac05296..b9774aaeff 100644 --- a/tests/models/test_taskinstance.py +++ b/tests/models/test_taskinstance.py @@ -518,6 +518,32 @@ class TestTaskInstance: ti.refresh_from_db() ti.state == state + def test_task_sigterm_call_on_failure_callback_only_once(self, dag_maker, caplog): + """ + Test that ensures on_failure_callback is called once on sigterm + """ + + def task_function(ti): + os.kill(ti.pid, signal.SIGTERM) + + with dag_maker("test_mark_failure_2"): + task_ = PythonOperator( + task_id="test_on_failure", + python_callable=task_function, + retries=0, + retry_delay=datetime.timedelta(seconds=2), + on_failure_callback=lambda context: context["ti"].log.info("on_failure_callback called"), + ) + + dr = dag_maker.create_dagrun() + ti = dr.task_instances[0] + ti.task = task_ + with pytest.raises(AirflowException): + ti.run() + ti.refresh_from_db() + assert ti.state == State.FAILED + assert caplog.text.count("on_failure_callback called") == 1 + @pytest.mark.parametrize( "state, exception, retries", [ ``` What I did was to revert the previous commit and introduced a flag -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected]
