Implement on_killed support for tasks during DAG run timeout

airflow/api_connexion/schemas/dag_schema.py

@@ -111,6 +111,7 @@ class DAGDetailSchema(DAGSchema):
     template_searchpath = fields.String(dump_only=True)
     render_template_as_native_obj = fields.Boolean(dump_only=True)
     last_loaded = fields.DateTime(dump_only=True, data_key="last_parsed")
+    call_on_kill_on_dagrun_timeout = fields.Boolean(dump_only=True)
 
     @staticmethod
     def get_tags(obj: DAG):

airflow/jobs/scheduler_job_runner.py

@@ -1689,7 +1689,19 @@ def _schedule_dag_run(
                     .where(TI.state.in_(State.unfinished))
                 )
                 for task_instance in unfinished_task_instances:
+                    # If the DAG is set to kill the task instances on timeout, we call on_kill
+                    # otherwise we just skip the task instance
                     task_instance.state = TaskInstanceState.SKIPPED
+                    if dag.call_on_kill_on_dagrun_timeout:
+                        task = dag.get_task(task_instance.task_id)
+                        if hasattr(task, "on_kill"):
+                            try:
+                                task.on_kill()
+                                task_instance.state = TaskInstanceState.FAILED
+                            except Exception as e:
+                                self.log.error("Error when calling on_kill for task %s: %s", task_instance, e)
+                        else:
+                            self.log.info("Task %s does not have on_kill method, skipping", task_instance)
                     session.merge(task_instance)
                 session.flush()
                 self.log.info("Run %s of %s has timed-out", dag_run.run_id, dag_run.dag_id)

airflow/models/baseoperator.py

@@ -274,6 +274,7 @@ def partial(
     task_display_name: str | None | ArgNotSet = NOTSET,
     logger_name: str | None | ArgNotSet = NOTSET,
     allow_nested_operators: bool = True,
+    call_on_kill_on_dagrun_timeout: bool | ArgNotSet = NOTSET,
     **kwargs,
 ) -> OperatorPartial:
     from airflow.models.dag import DagContext
@@ -343,6 +344,7 @@ def partial(
         "task_display_name": task_display_name,
         "logger_name": logger_name,
         "allow_nested_operators": allow_nested_operators,
+        "call_on_kill_on_dagrun_timeout": call_on_kill_on_dagrun_timeout,
     }
 
     # Inject DAG-level default args into args provided to this function.
@@ -919,6 +921,7 @@ def __init__(
         task_display_name: str | None = None,
         logger_name: str | None = None,
         allow_nested_operators: bool = True,
+        call_on_kill_on_dagrun_timeout: bool = True,
         **kwargs,
     ):
         from airflow.models.dag import DagContext
@@ -961,6 +964,7 @@ def __init__(
         self.on_skipped_callback = on_skipped_callback
         self._pre_execute_hook = pre_execute
         self._post_execute_hook = post_execute
+        self.call_on_kill_on_dagrun_timeout = call_on_kill_on_dagrun_timeout
 
         if start_date and not isinstance(start_date, datetime):
             self.log.warning("start_date for %s isn't datetime.datetime", self)

airflow/models/dag.py

@@ -355,6 +355,7 @@ def _create_orm_dagrun(
     "auto_register": bool,
     "fail_stop": bool,
     "dag_display_name": str,
+    "call_on_kill_on_dagrun_timeout": bool,
 }
 
 
@@ -530,6 +531,7 @@ def __init__(
         auto_register: bool = True,
         fail_stop: bool = False,
         dag_display_name: str | None = None,
+        call_on_kill_on_dagrun_timeout: bool = True,
     ):
         from airflow.utils.task_group import TaskGroup
 
@@ -556,6 +558,7 @@ def __init__(
 
         self._dag_id = dag_id
         self._dag_display_property_value = dag_display_name
+        self.call_on_kill_on_dagrun_timeout = call_on_kill_on_dagrun_timeout
 
         self._max_active_tasks = max_active_tasks
         self._pickle_id: int | None = None
@@ -3227,6 +3230,7 @@ def dag(
     auto_register: bool = True,
     fail_stop: bool = False,
     dag_display_name: str | None = None,
+    call_on_kill_on_dagrun_timeout: bool = True,
 ) -> Callable[[Callable], Callable[..., DAG]]:
     """
     Python dag decorator which wraps a function into an Airflow DAG.
@@ -3280,6 +3284,7 @@ def factory(*args, **kwargs):
                 auto_register=auto_register,
                 fail_stop=fail_stop,
                 dag_display_name=dag_display_name,
+                call_on_kill_on_dagrun_timeout=call_on_kill_on_dagrun_timeout,
             ) as dag_obj:
                 # Set DAG documentation from function documentation if it exists and doc_md is not set.
                 if f.__doc__ and not dag_obj.doc_md:

airflow/models/mappedoperator.py

@@ -602,6 +602,10 @@ def on_skipped_callback(self) -> TaskStateChangeCallbackAttrType:
     def on_skipped_callback(self, value: TaskStateChangeCallbackAttrType) -> None:
         self.partial_kwargs["on_skipped_callback"] = value
 
+    @property
+    def call_on_kill_on_dagrun_timeout(self) -> bool:
+        return bool(self.partial_kwargs.get("call_on_kill_on_dagrun_timeout"))
+
     @property
     def run_as_user(self) -> str | None:
         return self.partial_kwargs.get("run_as_user")

airflow/serialization/schema.json

@@ -175,6 +175,7 @@
         "has_on_failure_callback":  { "type": "boolean" },
         "render_template_as_native_obj":  { "type": "boolean" },
         "tags": { "type": "array" },
+          "call_on_kill_on_dagrun_timeout": { "type": "boolean" },
         "_task_group": {"anyOf": [
           { "type": "null" },
           { "$ref": "#/definitions/task_group" }

airflow/serialization/serialized_objects.py

@@ -1604,6 +1604,10 @@ def serialize_dag(cls, dag: DAG) -> dict:
         try:
             serialized_dag = cls.serialize_to_json(dag, cls._decorated_fields)
             serialized_dag["_processor_dags_folder"] = DAGS_FOLDER
+
+            if dag.call_on_kill_on_dagrun_timeout:
+                serialized_dag["call_on_kill_on_dagrun_timeout"] = dag.call_on_kill_on_dagrun_timeout
+
             serialized_dag["tasks"] = [cls.serialize(task) for _, task in dag.task_dict.items()]
 
             dag_deps = [
@@ -1666,6 +1670,8 @@ def deserialize_dag(cls, encoded_dag: dict[str, Any]) -> SerializedDAG:
                 v = cls._deserialize_params_dict(v)
             elif k == "tags":
                 v = set(v)
+            elif k == "call_on_kill_on_dagrun_timeout":
+                v = cls.deserialize(v)
             # else use v as it is
 
             setattr(dag, k, v)

tests/api_connexion/endpoints/test_dag_endpoint.py

@@ -339,6 +339,7 @@ def test_should_respond_200(self, url_safe_serializer):
             "template_searchpath": None,
             "timetable_description": None,
             "timezone": UTC_JSON_REPR,
+            "call_on_kill_on_dagrun_timeout": True,
         }
         assert response.json == expected
 
@@ -401,6 +402,7 @@ def test_should_respond_200_with_asset_expression(self, url_safe_serializer):
             "template_searchpath": None,
             "timetable_description": None,
             "timezone": UTC_JSON_REPR,
+            "call_on_kill_on_dagrun_timeout": True,
         }
         assert response.json == expected
 
@@ -451,6 +453,7 @@ def test_should_response_200_with_doc_md_none(self, url_safe_serializer):
             "template_searchpath": None,
             "timetable_description": None,
             "timezone": UTC_JSON_REPR,
+            "call_on_kill_on_dagrun_timeout": True,
         }
         assert response.json == expected
 
@@ -501,6 +504,7 @@ def test_should_response_200_for_null_start_date(self, url_safe_serializer):
             "template_searchpath": None,
             "timetable_description": None,
             "timezone": UTC_JSON_REPR,
+            "call_on_kill_on_dagrun_timeout": True,
         }
         assert response.json == expected
 
@@ -560,6 +564,7 @@ def test_should_respond_200_serialized(self, url_safe_serializer):
             "template_searchpath": None,
             "timetable_description": None,
             "timezone": UTC_JSON_REPR,
+            "call_on_kill_on_dagrun_timeout": True,
         }
         response = self.client.get(
             f"/api/v1/dags/{self.dag_id}/details", environ_overrides={"REMOTE_USER": "test"}
@@ -620,6 +625,7 @@ def test_should_respond_200_serialized(self, url_safe_serializer):
             "template_searchpath": None,
             "timetable_description": None,
             "timezone": UTC_JSON_REPR,
+            "call_on_kill_on_dagrun_timeout": True,
         }
         expected.update({"last_parsed": response.json["last_parsed"]})
         assert response.json == expected

tests/api_connexion/schemas/test_dag_schema.py

@@ -194,6 +194,7 @@ def test_serialize_test_dag_detail_schema(url_safe_serializer):
         "end_date": None,
         "is_paused_upon_creation": None,
         "render_template_as_native_obj": False,
+        "call_on_kill_on_dagrun_timeout": True,
     }
     obj = schema.dump(dag)
     expected.update({"last_parsed": obj["last_parsed"]})
@@ -258,6 +259,7 @@ def test_serialize_test_dag_with_asset_schedule_detail_schema(url_safe_serialize
         "end_date": None,
         "is_paused_upon_creation": None,
         "render_template_as_native_obj": False,
+        "call_on_kill_on_dagrun_timeout": True,
     }
     obj = schema.dump(dag)
     expected.update({"last_parsed": obj["last_parsed"]})

tests/jobs/test_scheduler_job.py

@@ -2503,6 +2503,45 @@ def test_dagrun_timeout_fails_run(self, dag_maker):
         session.rollback()
         session.close()
 
+    def test_dagrun_timeout_triggers_on_kill(self, dag_maker, session):
+        """Test that when a DAG times out, call_on_kill_on_dagrun_timeout triggers on_kill for tasks."""
+        with dag_maker(
+            dag_id="test_dagrun_timeout",
+            dagrun_timeout=datetime.timedelta(seconds=1),
+            call_on_kill_on_dagrun_timeout=True,
+        ):
+            task = BashOperator(
+                task_id="task",
+                bash_command="sleep 10",
+            )
+
+        dr = dag_maker.create_dagrun(
+            state=DagRunState.RUNNING,
+            start_date=timezone.utcnow() - datetime.timedelta(seconds=2),
+        )
+        ti = dr.get_task_instance(task.task_id)
+        ti.set_state(TaskInstanceState.RUNNING, session)
+        session.merge(ti)
+        session.flush()
+
+        scheduler_job = Job()
+        self.job_runner = SchedulerJobRunner(job=scheduler_job, subdir=os.devnull)
+
+        scheduler_job.executor = MockExecutor(do_update=True)
+        self.job_runner.processor_agent = mock.Mock(spec=DagFileProcessorAgent)
+
+        self.job_runner._schedule_dag_run(dr, session)
+        session.flush()
+
+        # Check that the DAG run is now in a failed state
+        # This refresh here is causing an error.
+        # session.refresh(dr)
+        assert dr.state == DagRunState.FAILED
+
+        # Check that the task instance is now in a failed state
+        ti.refresh_from_db(session=session)
+        assert ti.state == TaskInstanceState.FAILED
+
     def test_dagrun_timeout_fails_run_and_update_next_dagrun(self, dag_maker):
         """
         Test that dagrun timeout fails run and update the next dagrun
@@ -4422,7 +4461,7 @@ def test_do_schedule_max_active_runs_dag_timed_out(self, dag_maker):
         run1 = session.merge(run1)
         session.refresh(run1)
         assert run1.state == State.FAILED
-        assert run1_ti.state == State.SKIPPED
+        assert run1_ti.state == State.FAILED
         session.flush()
         # Run relevant part of scheduling again to assert run2 has been scheduled
         self.job_runner._start_queued_dagruns(session)

tests/serialization/test_dag_serialization.py

@@ -230,6 +230,7 @@
         "dag_dependencies": [],
         "params": [],
         "tags": [],
+        "call_on_kill_on_dagrun_timeout": True,
     },
 }
 
@@ -1273,6 +1274,7 @@ def test_no_new_fields_added_to_base_operator(self):
             "wait_for_past_depends_before_skipping": False,
             "weight_rule": _DownstreamPriorityWeightStrategy(),
             "multiple_outputs": False,
+            "call_on_kill_on_dagrun_timeout": True,
         }, """
 !!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!