diederikperdok opened a new issue, #76320:
URL: https://github.com/apache/airflow/issues/76320
### Under which category would you file this issue?
Airflow Core
### Apache Airflow version
3.3.2
### What happened and how to reproduce it?
### The problem
If a Trigger's `on_kill` method takes more than 1 second to run, it doesn't
complete. This can leave external work running while the corresponding Airflow
task instance is no longer in the running state.
### Minimal reproducible example
Create a file `slow_kill.py` with contents:
```python
import asyncio
from collections.abc import AsyncIterator
from typing import Any
from airflow.sdk import BaseOperator, Context
from airflow.triggers.base import BaseTrigger, TriggerEvent
class SlowKillTrigger(BaseTrigger):
def __init__(
self,
duration_seconds: float = 60.0,
kill_duration_seconds: float = 1.5,
) -> None:
super().__init__()
self.duration_seconds = duration_seconds
self.kill_duration_seconds = kill_duration_seconds
def serialize(self) -> tuple[str, dict[str, Any]]:
return (
f"{self.__class__.__module__}.{self.__class__.__qualname__}",
{
"duration_seconds": self.duration_seconds,
"kill_duration_seconds": self.kill_duration_seconds,
},
)
async def run(self) -> AsyncIterator[TriggerEvent]:
self.log.info(
f"Trigger running for {self.duration_seconds:.1f}s."
)
await asyncio.sleep(self.duration_seconds)
yield TriggerEvent({"status": "completed"})
async def on_kill(self) -> None:
self.log.info(f"on_kill() STARTED — sleeping
{self.kill_duration_seconds}s...")
try:
# Simulate slow killing, e.g. databricks hook cancelling jobs
taking 100ms+ per job
# and queueing all cancellations on a single thread
await asyncio.sleep(self.kill_duration_seconds)
except asyncio.CancelledError:
self.log.error("on_kill() ABORTED before completing!")
raise
# This should be reached to successfully cancel external work
self.log.info("on_kill() FINISHED — external job successfully
cancelled!")
class SlowKillDeferrableOperator(BaseOperator):
def __init__(
self,
duration_seconds: float = 60.0,
kill_duration_seconds: float = 1.5,
**kwargs: Any,
) -> None:
super().__init__(**kwargs)
self.duration_seconds = duration_seconds
self.kill_duration_seconds = kill_duration_seconds
def execute(self, context: Context) -> None:
self.log.info(f"Deferring to SlowKillTrigger for
{self.duration_seconds}s.")
self.defer(
trigger=SlowKillTrigger(
duration_seconds=self.duration_seconds,
kill_duration_seconds=self.kill_duration_seconds,
),
method_name="execute_complete",
)
def execute_complete(self, context: Context, event: dict[str, Any]) ->
None:
self.log.info(f"Task completed after deferral: {event}")
```
and a file `repro_dag.py` with contents:
```python
from datetime import datetime
from airflow.sdk import DAG
from slow_kill import SlowKillDeferrableOperator
with DAG(
dag_id="example-slow-kill-does-not-complete",
start_date=datetime(2026, 1, 1),
schedule=None,
catchup=False,
) as dag:
task_fast = SlowKillDeferrableOperator(task_id="kill_0_1s",
kill_duration_seconds=0.1)
task_slow = SlowKillDeferrableOperator(task_id="kill_1_5s",
kill_duration_seconds=1.5)
```
and run Airflow from with the directory containing those files with
```bash
docker run --rm -it \
-p 8080:8080 \
-e AIRFLOW__CORE__LOAD_EXAMPLES=False \
-e AIRFLOW__CORE__SIMPLE_AUTH_MANAGER_ALL_ADMINS=True \
-v "$PWD/slow_kill.py:/opt/airflow/plugins/slow_kill.py" \
-v "$PWD/repro_dag.py:/opt/airflow/dags/repro_dag.py" \
apache/airflow:3.3.2 standalone
```
now go to the ui at `localhost:8080`, trigger the dag, wait a few seconds
for the tasks to get deferred and then mark the run as failed within 60s.
The logs of task `kill_0_1s` (with the fast `on_kill`) will show
```
on_kill() STARTED — sleeping 0.10s...
on_kill() FINISHED — external job successfully cancelled!
```
indicating that `on_kill` completed successfully.
and the logs of the task `kill_1_5s` (with the slow `on_kill`) show
```
on_kill() STARTED — sleeping 1.50s...
on_kill() ABORTED before completing!
```
indicating that `on_kill` did not finish.
### How this issue affects the Databricks provider
When marking a DAG with many `DatabricksSubmitRunOperator` tasks as failed,
we observed that for many of those tasks the external Databricks job kept on
running. Upon further inspection, it seems that the `on_kill` method for this
operator's trigger puts all job cancellation API calls on a single thread
across all task instances, so if there are many task instances to cancel,
`on_kill` takes >1s for some of them
The problem is not at all specific to Databricks though, as the above
example demonstrates.
### What should be done
I think the first question to be answered is: "Are `on_kill` implementations
allowed to take >1s?"
* If no, the only required airflow-core PR is a documentation PR to make
this requirement clear. The airflow-databricks-provider package should then
update its implementation to make sure `on_kill` always runs within a second,
even when many tasks are cancelled at the same time.
* If yes, some work on the triggerer is required, but at this time I don't
understand its inner workings well enough to make a proposal.
Personally, I think 1s is on the low side for cancelling external jobs that
may require one or more API calls, so `on_kill` should ideally be allowed to
take a bit longer.
### What you think should happen instead?
Either both `on_kill` method should succeed or the docs should clearly state
that `on_kill` may not take longer than 1 second (and when it does take >1s the
logs should clearly show that it failed because of that)
### Operating System
MacOS and Linux
### Deployment
None
### Apache Airflow Provider(s)
_No response_
### Versions of Apache Airflow Providers
_No response_
### Official Helm Chart version
Not Applicable
### Kubernetes Version
_No response_
### Helm Chart configuration
_No response_
### Docker Image customizations
_No response_
### Anything else?
_No response_
### Are you willing to submit PR?
- [ ] Yes I am willing to submit a PR!
### Code of Conduct
- [x] I agree to follow this project's [Code of
Conduct](https://github.com/apache/airflow/blob/main/CODE_OF_CONDUCT.md)
--
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]
