Implement on_killed support for tasks during DAG run timeout

[MRLab12]

---

• This is a PR to fix Test workflow issues in the previous PR.

This PR addresses issue #41036 by adding support for the on_killed callback on tasks that are still running when a DAG run reaches its timeout.

Key changes:

1. Added a new configuration option call_on_kill_on_dagrun_timeout to control whether tasks should be killed when a DAG run times out (default: True).
2. Updated logic in _schedule_dag_run to call task on_kill if call_on_kill_on_dagrun_timeout is enabled.

closes: #41036

Testing:

• Added unit test for the new timeout handling logic
• Verified that the new configuration option works as expected

[MRLab12]

After the last test fix I did there, I'm thinking what would the correct default behavior for this feature should be? Currently it's default True, but if we want to normally just skip the TI, then the default should be False.

[MRLab12]

@RNHTTR Checking in here, have you gotten a chance to look at these changes?

[RNHTTR]

Should this have a specific exception handling for AttributeError in the case where the operator doesn't have an on_kill method defined? In this case, I think we'd want to self.log.info("on_kill method does not exist for task %s: %s", task_instance, e).

What do you think?

[RNHTTR]

Or maybe something like...

if hasattr(task, 'on_kill'):
    try:
        ...
else:
    self.log.warn("Task %s was configured to execute its on_kill method after a DAG run timeout, but it does not have the method 'on_kill' method defined", task_instance)

[MRLab12]

If there can be a case of an operator without an on_kill, then definitely we should check for that. I think using hasattr to check for it before getting the task and using on_kill is better to prevent the AttributeError. I will make this change.

[MRLab12]

I had a conflict and made a mistake when pushing the commits again. After getting the commits pushed, the build is failing with the error: Please ask maintainer to assign the 'legacy api' label to the PR in order to continue . Is this correct?

[ashb]

Is there any context/previous discussion about this?

The reason I ask is that I'm wary of adding yet more options of ways things can run

[MRLab12]

@RNHTTR can help clarify, but basically this change would allow tasks to run on_kill when a dagrun reaches timeout. The use case given in the issue is that "Some users would like for externally running workloads (e.g. snowflake, emr, bigquery, etc etc) to stop executing when a DAG run times out." . With this flag the user will have more control on what behavior they want when there's a timeout.

[RNHTTR]

This all stems from the fairly controversial idea of setting running tasks to skipped if a DAG run reaches its timeout.

There are a couple paths forward:

1. Change the behavior after a dagrun_timeout is reached to mark running TIs as failed.
2. Proceed with this PR
3. Encourage users to use an on_skipped_callback for external systems that need to be shutdown when a task moves to the skipped state after a DAG run timeout.

[MRLab12]

@ashb @RNHTTR Checking in here. Do we want to move forward with this PR?

[RNHTTR]

IMO this would be a good change. It's somewhat counterintuitive that tasks will be marked as skipped when a Dagrun times out, and enabling on_kill on Dagrun timeout is intuitive.

[MRLab12]

@ashb Any thoughts on what's the best path forward?

Also, I mentioned in a comment here that there has been a significant change to the DAG properties. If we want to use this new call_on_kill_on_dagrun_timeout, where should it go now?

[MRLab12]

Looks like there has been an extensive change that was merged recently here that changed how the DAG properties are defined. I'm going through the changes in the PR, but would the changes from here go in Airflow's core rather than in the Task SDK?

Trying to fix the merge conflict and a bit confused in where the properties go now.

[RNHTTR]

@MRLab12 I think they definitely need to go into the Task SDK, because that's what users will interface with. Because the models inherit from their associated SDK objects (example), I don't think you need to duplicate them there unless you'll be explicitly writing code at the model level, which I don't think is necessary.

[github-actions]

This pull request has been automatically marked as stale because it has not had recent activity. It will be closed in 5 days if no further activity occurs. Thank you for your contributions.