Skip to content

Fix teardown TaskFlow tasks failing on reuse and with expand_kwargs - #74782

Open
surawut-jirasaktavee wants to merge 1 commit into
apache:mainfrom
surawut-jirasaktavee:fix-teardown-decorator-reuse
Open

surawut-jirasaktavee wants to merge 1 commit into
apache:mainfrom
surawut-jirasaktavee:fix-teardown-decorator-reuse

Conversation

@surawut-jirasaktavee

Copy link
Copy Markdown
Contributor

A @teardown TaskFlow task cannot be used twice in a Dag, and mapping one with
expand_kwargs() gives it the wrong trigger rule.

@teardown
def cleanup(x=None): ...

with DAG("example"):
    cleanup()
    cleanup()                            # ValueError: Trigger rule not configurable for teardown tasks.
    cleanup.override(task_id="c2")()     # same ValueError

Cause. __call__ and expand() wrote trigger_rule into the decorator's own
kwargs. The next call, and any .override() copy, saw that stored value as user
input and failed the check meant for users. __call__ also popped
on_failure_fail_dagrun from the same kwargs, so a value set through
.override(on_failure_fail_dagrun=True) applied only to the first call.

expand_kwargs() had no teardown handling at all. A teardown mapped with it got
all_success instead of all_done_setup_success, so it was skipped when an
upstream task failed.

Fix. Both paths now work on a copy of the decorator's kwargs. The teardown
check and trigger rule are applied in one helper, called from __call__ and from
_expand(), which serves both expand() and expand_kwargs(). The error message
is unchanged.

Tests. I added five tests to task-sdk/tests/task_sdk/definitions/decorators/test_setup_teardown.py.
All five fail without the fix.

  • Reusing a teardown after a first __call__ or expand(), with the second use through .override(task_id=...).
  • on_failure_fail_dagrun set through .override() is kept on a second call.
  • expand_kwargs() sets all_done_setup_success.
  • expand_kwargs() rejects a user-supplied trigger_rule on a teardown, as expand() does.

Checks run locally:

  • task-sdk/tests/task_sdk/definitions/decorators/ and task-sdk/tests/task_sdk/bases/test_decorator.py: 137 passed.
  • breeze testing task-sdk-tests, the core tests with and without the database, the provider tests with the database, and breeze build-docs all pass.
  • In Breeze, 23 setup/teardown, override and expand tests in providers/standard/tests/unit/standard/decorators/ pass.
  • prek passes for both the pre-commit and the manual stages.
  • The full non-database provider run did not complete locally because of xdist collection errors in unrelated providers; I'm relying on CI for that part.

closes: #74756


Was generative AI tooling used to co-author this PR?
  • Yes (please specify the tool below)

Generated-by: Claude Code (Opus 5.5) following the guidelines

A @teardown function could be used only once. The first call stored the
teardown trigger rule in the decorator's own kwargs, so a later call,
.override() or .expand() read it as a user-supplied trigger rule and
raised ValueError. The same shared state dropped an on_failure_fail_dagrun
override after the first task was created.

Teardowns mapped with .expand_kwargs() never got the teardown trigger
rule and fell back to all_success. When an upstream task failed, the
teardown was skipped and the resources it should have cleaned up were
left behind.

@fat-catTW fat-catTW left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thanks for the work!

This branch has not been deployed

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

Projects

None yet

Development

Successfully merging this pull request may close these issues.

@teardown task raises ValueError when called twice, and gets the wrong trigger rule with expand_kwargs

2 participants