Skip to content

Deferrable operator tasks do not call on_kill method when fail or restarted #36090

Closed

Description

@nirben82

Apache Airflow version

Other Airflow 2 version (please specify below)

What happened

When marking a task that is using a deferrable operators as failed or cleared, on_kill method is not called and the remote job is never stopped.

Task log output:

*** Found local files:
***   * /usr/local/google/home/nirben/airflow/logs/dag_id=dag_test/run_id=manual__2023-12-06T13:22:54.698489+00:00/task_id=bigquery_task/attempt=30.log
***   * /usr/local/google/home/nirben/airflow/logs/dag_id=dag_test/run_id=manual__2023-12-06T13:22:54.698489+00:00/task_id=bigquery_task/attempt=30.log.trigger.13880.log
[2023-12-06, 14:14:11 UTC] {taskinstance.py:1103} INFO - Dependencies all met for dep_context=non-requeueable deps ti=<TaskInstance: dag_test.bigquery_task manual__2023-12-06T13:22:54.698489+00:00 [queued]>
[2023-12-06, 14:14:11 UTC] {taskinstance.py:1103} INFO - Dependencies all met for dep_context=requeueable deps ti=<TaskInstance: dag_test.bigquery_task manual__2023-12-06T13:22:54.698489+00:00 [queued]>
[2023-12-06, 14:14:11 UTC] {taskinstance.py:1308} INFO - Starting attempt 30 of 30
[2023-12-06, 14:14:11 UTC] {taskinstance.py:1327} INFO - Executing <Task(BigQueryInsertJobOperator): bigquery_task> on 2023-12-06 13:22:54.698489+00:00
[2023-12-06, 14:14:11 UTC] {standard_task_runner.py:57} INFO - Started process 629728 to run task
[2023-12-06, 14:14:11 UTC] {standard_task_runner.py:84} INFO - Running: ['airflow', 'tasks', 'run', 'dag_test', 'bigquery_task', 'manual__2023-12-06T13:22:54.698489+00:00', '--job-id', '13896', '--raw', '--subdir', 'DAGS_FOLDER/dag_test.py', '--cfg-path', '/tmp/tmpxnqe4ysz']
[2023-12-06, 14:14:11 UTC] {standard_task_runner.py:85} INFO - Job 13896: Subtask bigquery_task
[2023-12-06, 14:14:11 UTC] {task_command.py:410} INFO - Running <TaskInstance: dag_test.bigquery_task manual__2023-12-06T13:22:54.698489+00:00 [running]> on host nirben-ws1.tlv.corp.google.com
[2023-12-06, 14:14:11 UTC] {taskinstance.py:1545} INFO - Exporting env vars: AIRFLOW_CTX_DAG_OWNER='some-owner' AIRFLOW_CTX_DAG_ID='dag_test' AIRFLOW_CTX_TASK_ID='bigquery_task' AIRFLOW_CTX_EXECUTION_DATE='2023-12-06T13:22:54.698489+00:00' AIRFLOW_CTX_TRY_NUMBER='30' AIRFLOW_CTX_DAG_RUN_ID='manual__2023-12-06T13:22:54.698489+00:00'
[2023-12-06, 14:14:11 UTC] {base.py:73} INFO - Using connection ID 'some-connection' for task execution.
[2023-12-06, 14:14:11 UTC] {bigquery.py:2799} INFO - Executing: {'query': {'query': 'SELECT * FROM table, 'useLegacySql': False, 'priority': 'batch', 'allowLargeResults': True, 'destinationTable': {'datasetId': 'datasetId', 'projectId': 'projectId', 'tableId': 'some_table_test'}, 'flattenResults': False, 'writeDisposition': 'WRITE_TRUNCATE', 'createDisposition': 'CREATE_IF_NEEDED'}}'
[2023-12-06, 14:14:11 UTC] {credentials_provider.py:353} INFO - Getting connection using `google.auth.default()` since no explicit credentials are provided.
[2023-12-06, 14:14:11 UTC] {logging_mixin.py:150} WARNING - /usr/local/google/home/nirben/venv/waze-data/lib/python3.8/site-packages/google/auth/_default.py:78 UserWarning: Your application has authenticated using end user credentials from Google Cloud SDK without a quota project. You might receive a "quota exceeded" or "API not enabled" error. See the following page for troubleshooting: https://cloud.google.com/docs/authentication/adc-troubleshooting/user-creds.
[2023-12-06, 14:14:12 UTC] {bigquery.py:1596} INFO - Inserting job airflow_dag_test_bigquery_task_2023_12_06T13_22_54_698489_00_00_5bf2e5098a664fd1d54ec9b9b75d077b
[2023-12-06, 14:14:13 UTC] {bigquery.py:51} INFO - Using the connection  some-connection .
[2023-12-06, 14:14:13 UTC] {taskinstance.py:1415} INFO - Pausing task as DEFERRED. dag_id=dag_test, task_id=bigquery_task, execution_date=20231206T132254, start_date=20231206T141411
[2023-12-06, 14:14:13 UTC] {local_task_job_runner.py:222} INFO - Task exited with return code 100 (task deferral)
[2023-12-06, 14:14:14 UTC] {base.py:73} INFO - Using connection ID 'google_cloud_pipelines' for task execution.
[2023-12-06, 14:14:15 UTC] {bigquery.py:93} INFO - Bigquery job status is running. Sleeping for 4.0 seconds.
[2023-12-06, 14:14:19 UTC] {bigquery.py:93} INFO - Bigquery job status is running. Sleeping for 4.0 seconds.
[2023-12-06, 14:14:21 UTC] {triggerer_job_runner.py:625} ERROR - Trigger cancelled; message=```

### What you think should happen instead

I think that the trigger should be stopped and task instance should continue the same [behavior](https://github.com/apache/airflow/blob/d2514b408cb98f792289a5d032aaf85fe605350d/airflow/models/taskinstance.py#L2452) that is done by any non-deferrable tasks.

### How to reproduce

1. Invoke a run of the below dag and after the task is in state `defer`, mark it as `failed` or clear it. 

    The task log ends with the below text and the job in BQ does not stop.
    ```
    [2023-12-06, 14:14:14 UTC] {base.py:73} INFO - Using connection ID 'google_cloud_pipelines' for task execution.
    [2023-12-06, 14:14:15 UTC] {bigquery.py:93} INFO - Bigquery job status is running. Sleeping for 4.0 seconds.
    [2023-12-06, 14:14:19 UTC] {bigquery.py:93} INFO - Bigquery job status is running. Sleeping for 4.0 seconds.
    [2023-12-06, 14:14:21 UTC] {triggerer_job_runner.py:625} ERROR - Trigger cancelled; message=
    ```

2. Change to `deferrable=False`,  Invoke a run of the dag and after the task is in state `running` and the job started in BQ, mark it as `failed` or clear it. 

    The task log ends with the below text and the job in BQ stops.
    ```
    [2023-12-06, 14:13:06 UTC] {bigquery.py:1596} INFO - Inserting job airflow_dag_test_bigquery_task_2023_12_06T13_22_54_698489_00_00_9238ba6363d9b8cca1e4a1fd26b7e985
    [2023-12-06, 14:13:36 UTC] {local_task_job_runner.py:291} WARNING - State of this instance has been externally set to failed. Terminating instance.
    [2023-12-06, 14:13:36 UTC] {process_utils.py:131} INFO - Sending Signals.SIGTERM to group 629540. PIDs of all processes in the group: [629540]
    [2023-12-06, 14:13:36 UTC] {process_utils.py:86} INFO - Sending the signal Signals.SIGTERM to group 629540
    [2023-12-06, 14:13:36 UTC] {taskinstance.py:1517} ERROR - Received SIGTERM. Terminating subprocesses.
    [2023-12-06, 14:13:36 UTC] {bigquery.py:1487} INFO - Attempting to cancel job : project-id, airflow_dag_test_bigquery_task_2023_12_06T13_22_54_698489_00_00_9238ba6363d9b8cca1e4a1fd26b7e985
    [2023-12-06, 14:13:37 UTC] {bigquery.py:1508} INFO - Waiting for canceled job project-id, airflow_dag_test_bigquery_task_2023_12_06T13_22_54_698489_00_00_9238ba6363d9b8cca1e4a1fd26b7e985 to finish.
    [2023-12-06, 14:13:43 UTC] {bigquery.py:1499} INFO - Job successfully canceled: project-id, airflow_dag_test_bigquery_task_2023_12_06T13_22_54_698489_00_00_9238ba6363d9b8cca1e4a1fd26b7e985
    [2023-12-06, 14:13:43 UTC] {process_utils.py:79} INFO - Process psutil.Process(pid=629540, status='terminated', exitcode=0, started='16:13:05') (629540) terminated with exit code 0
from datetime import datetime, timedelta
from airflow import DAG
from airflow.providers.google.cloud.operators.bigquery import BigQueryInsertJobOperator

default_args = {
    'owner': 'owner',
    'start_date': datetime(2023, 12, 5),
    'email_on_failure': False,
    'email_on_retry': False,
    'retries': 0,
    'retry_delay': timedelta(minutes=15)
}

dag = DAG('dag_test', default_args=default_args, catchup=False)

query_that_runs_for_a_few_minutes = """
SELECT * FROM some-large-table
"""

client_driver_events_task = BigQueryInsertJobOperator(
    task_id='bigquery_task',
    gcp_conn_id='google_cloud_pipelines',
    dag=dag,
    configuration={
        'query': {
            'query': query_that_runs_for_a_few_minutes.strip(),
            'useLegacySql': False,
            'priority': 'batch',
            "allowLargeResults": True,
            "destinationTable": {
                "datasetId": "datasetId",
                "projectId": "projectId",
                "tableId": "tableId"
            },
            'flattenResults': False,
            'writeDisposition': 'WRITE_TRUNCATE',
            'createDisposition': 'CREATE_IF_NEEDED'
        }
    },
    deferrable=True)

Operating System

Debian GNU

Versions of Apache Airflow Providers

apache-airflow==2.6.2
apache-airflow-providers-google==10.9.0

Deployment

Official Apache Airflow Helm Chart

Deployment details

No response

Anything else

No response

Are you willing to submit PR?

  • Yes I am willing to submit a PR!

Code of Conduct

Activity

  1. boring-cyborg commented on Dec 6, 2023

    @boring-cyborg

    Thanks for opening your first issue here! Be sure to follow the issue template! If you are willing to raise PR to address this issue please do so, no need to wait for approval.

  2. changed the title [-]Deferrable operator tasks don't call on_kill when fail or restarted[/-] [+]Deferrable operator tasks do not call `on_kill` method when fail or restarted[/+] on Dec 6, 2023
  3. tirkarthi commented on Dec 6, 2023

    @tirkarthi
    Contributor

    Related : #19929 . It seems this was reported earlier but got closed as stale issue. We would also find this feature useful to cancel remote jobs tracked by triggerer on clear/fail.

  4. removed
    needs-triagelabel for new issues that we didn't triage yet
    on Dec 11, 2023
  5. potiuk commented on Dec 11, 2023

    @potiuk
    Member

    Yes. Probably it should be changed. The problem is that this would have to bring the task from deferred state first to run it's on_kill method, but that would require someone to implement the logic for that.

  6. nirben82 commented on Dec 13, 2023

    @nirben82
    Author

    Thanks @potiuk.

    1. Do you know if someone is currently working on this?
    2. Could you please provide some more specifics on what needs to be implemented, so we can understand if our team can prioritize adding the logic?
    3. Is it also possible to update previous minor version 2.6? This is mainly so we can continue being aligned with Cloud Composer that is currently on 2.6.3.
  7. potiuk commented on Dec 13, 2023

    @potiuk
    Member

    1, Not that I am aware of. But you are asking random maintainer who might not know everything. In OSS everything is in the open so if this is not mentioned here, there is no other place where someone working on it would be recorded.

    1. It's really part of the task to design it. This one will likely require Airflow Improvement Proposal or at the very least discussion with the devlist outlining the proposal. Shortly you will need to add a whole new internal processing on how to make sure to instantiate the task in a worker when it is deferred to execute on_kiill in the worker. Currently on_kill is run when task is exiting from running state - so when it is already handled by the worker. Deferred task is handled by Triggerer by runing triggers, not real "task" code - just triggers that task deferred to. So what is needed is to decide who and how will schedule a task that should only execute on_kill method in the worker when the task is set to killed statis and is currently deferred to Trigger. This has not been discussed nor designed yet, so it has to start (and the person/people doing it must be capable of) with making desing and proposla how to implement it, lead the discussion in devlist and get approval on the proposal. So you really need someone not to execute precise instructions but somoene who will be able to propose and lead discussion on the whole design of it.

    2. No. This is against our versioning policies https://airflow.apache.org/docs/apache-airflow/stable/release-process.html - SemVer forbids adding tasks (this is a new feature), Also we only release new things in exactly one currently active branch. As of now we are voting on release candidate for 2.8.0 and the current main is 2.9.0dev0 so the first thing this feature might be available is 2.9.0 - but it woudl reqiure to start preparing the design and making proposal as soon as possible.

  8. added this to the Airflow 2.10.0 milestone on Mar 7, 2024
  9. pustoshilov-d commented on Apr 11, 2024

    @pustoshilov-d
    Contributor

    We're working with Slurm deferrable operator and waiting this feature so much 🙏

  10. potiuk commented on Apr 11, 2024

    @potiuk
    Member

    We're working with Slurm deferrable operator and waiting this feature so much 🙏

    Just a kind reminder: if you and your company wait for a feature badly, the most certain way to get it is to spend engineering effort and contribute it back. Other than that the only thing to do is to 🙏 that someone will do it. Also paying someone to contribute such feature works.

    This is how open-source project works.

  11. self-assigned this
    on Apr 12, 2024
  12. 33 remaining items

  13. phanikumv commented on Sep 16, 2024

    @phanikumv
    Contributor

    @sunank200 is occupied with other developments for Airflow 3. Would you or any one from the community be up for spending some engineering efforts on this @thesuperzapper ?

  14. removed this from the Airflow 2.11.0 milestone on Sep 18, 2024
  15. kaxil commented on Sep 18, 2024

    @kaxil
    Member

    @msumit is also interested in this afaik. Would be great if someone in the community is willing to take this on.

  16. msumit commented on Sep 19, 2024

    @msumit
    Contributor

    TLDR: I coded the fix for the same problem in a more generic way, which should solve this for all operators. I'll try to raise a diff by the EOW.

  17. thesuperzapper commented on Sep 20, 2024

    @thesuperzapper
    Contributor

    TLDR: I coded the fix for the same problem in a more generic way, which should solve this for all operators. I'll try to raise a diff by the EOW.

    @msumit whatever your solution is, make sure that it meets the following requirements:

    • setting a deferred task to failed/succeeded cancels the external job
    • clearing a deferred task cancels the external job (but not the newly scheduled the next try starts and deferred quickly)
    • shutting down the trigger should not cancel the external job of any deferred tasks running on it
  18. msumit commented on Sep 23, 2024

    @msumit
    Contributor

    Sorry for the bummer folks, I don't think the solution I was thinking of is generic enough, but here are 2 different approaches I could think of

    • There is a cleanup function in BaseTrigger, which I think has not been implemented by any other Trigger class. We can invoke that function when calling cancel of the async task, but yeah that function needs to be implemented in every Trigger class, very much like what @sunank200 did for a couple of classes.

    • We can send back the control to normal executors & call the on_kill method of the operators. This would work for most of the cases, but there are a couple of issues here as well

      • Many operator's "on_kill" function isn't deferred aware, which means it doesn't know that execution is being called from this function directly, and it may be using some information coming out of the execute function.
      • As we've to pass the task back to executors, the user would see strange task state changes, like from success/failure to scheduled & then running & then to the final state. Not sure if it's worth the effort to introduce some transient states like marking_success or being_cancelled etc.

    In short, we need an owner here who could either find a generic way or take the lead in making all trigger classes cancel aware.

  19. thesuperzapper commented on Sep 23, 2024

    @thesuperzapper
    Contributor

    @msumit we already have a generic workaround based on my code in: #36090 (comment)

    It's not perfect, it doesn't handle the case that the triggerer is force killed (power loss) at the exact same time that the user manually changes the state of the task, but that's a very rare situation (and also applies to non-deferred on_kill during a worker failures).

    It relies on catching the asyncio.CancelledError exception and canceling the external job from within the triggerer itself. But only if the current state of the task (read from DB) is not deferred, to avoid canceling when the triggerer is shut down normally.

    The maintainers of the providers really need to implement this workaround for all deferrable operators with on_kill, otherwise, we need a warning on each of the deferrable operator docs. Users do not expect this behavior, and leave their external jobs running accidentally.


    Airflow 3 Proposal:

    My proposal for a unified clean-up solution is to introduce a new TaskInstance state called pending cleanup.

    When an operator defines a "cleanup handler", the task would enter this state when it terminates for any reason, including success, failure, and clearing.

    The "cleanup handler" is guaranteed to run at least once, and would be passed the following information to help it decide what to do:

    1. Which state the task will go into after cleanup (success, failure, clear, up-for-retry)
    2. How the state changed (if it was manually set, failed naturally, or was cleared)
    3. The parameters of the task
    4. Any return xcom values (if the task succeeded)
    5. Runtime metadata set by the operator author using a new api (important if the task triggers an external job and gets an ID back, which could be required to cancel the external job)

    The question then becomes where to run the cleanup handler code. I think it makes sense to run it in the existing triggerers, as we can reuse the existing task serialization behavior, and know that it will run at least once.

  20. vikramkoka commented on Jan 9, 2025

    @vikramkoka
    Contributor

    Thank you for highlighting the issue and all the details around this.

    At the risk of complicating the solution further, the "cleanup handler" described above, seems identical to at least the intent for the "teardown task", described and implemented as part of the Setup / Teardown AIP. I believe that we may have even included very similar examples in that AIP.

    Was that considered in the solution proposal above?

  21. jayceslesar commented on Nov 13, 2025

    @jayceslesar
    Contributor

    Is this fixed with 3.0.0?

  22. ashb commented on Nov 14, 2025

    @ashb
    Member

    Is this fixed with 3.0.0?

    Nope

  23. amoghrajesh commented on Apr 21, 2026

    @amoghrajesh
    Contributor

    Taking a look at this one.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Metadata

Metadata

Assignees

Type

No type

Projects

No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions