diff --git a/airflow/api_connexion/schemas/dag_schema.py b/airflow/api_connexion/schemas/dag_schema.py index 4fb1dd6ae4b39..f05eb3cca29ab 100644 --- a/airflow/api_connexion/schemas/dag_schema.py +++ b/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): diff --git a/airflow/jobs/scheduler_job_runner.py b/airflow/jobs/scheduler_job_runner.py index 15042b0d3f177..d3c2db5d24ff4 100644 --- a/airflow/jobs/scheduler_job_runner.py +++ b/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) diff --git a/airflow/models/baseoperator.py b/airflow/models/baseoperator.py index 514553e05a2dd..b062c78bf2934 100644 --- a/airflow/models/baseoperator.py +++ b/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) diff --git a/airflow/models/dag.py b/airflow/models/dag.py index d72eb0928cb46..29d38d238d9f7 100644 --- a/airflow/models/dag.py +++ b/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: diff --git a/airflow/models/mappedoperator.py b/airflow/models/mappedoperator.py index 8a9e790ea7fc6..234eb0e86a6ea 100644 --- a/airflow/models/mappedoperator.py +++ b/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") diff --git a/airflow/serialization/schema.json b/airflow/serialization/schema.json index fe1e63c490355..3d7160d0ecc93 100644 --- a/airflow/serialization/schema.json +++ b/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" } diff --git a/airflow/serialization/serialized_objects.py b/airflow/serialization/serialized_objects.py index 14528f62bddf1..20f9a7ff5ece8 100644 --- a/airflow/serialization/serialized_objects.py +++ b/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) diff --git a/tests/api_connexion/endpoints/test_dag_endpoint.py b/tests/api_connexion/endpoints/test_dag_endpoint.py index ecb0052a8a03b..4129e58c26aa3 100644 --- a/tests/api_connexion/endpoints/test_dag_endpoint.py +++ b/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 diff --git a/tests/api_connexion/schemas/test_dag_schema.py b/tests/api_connexion/schemas/test_dag_schema.py index 1576f9cf3d6ee..61a03eebd75c9 100644 --- a/tests/api_connexion/schemas/test_dag_schema.py +++ b/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"]}) diff --git a/tests/jobs/test_scheduler_job.py b/tests/jobs/test_scheduler_job.py index 311de0ce2b648..edb109e088a0c 100644 --- a/tests/jobs/test_scheduler_job.py +++ b/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) diff --git a/tests/serialization/test_dag_serialization.py b/tests/serialization/test_dag_serialization.py index b54ded3a55c6a..5694704e10e9d 100644 --- a/tests/serialization/test_dag_serialization.py +++ b/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, }, """ !!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!