Repository navigation
Bound LLM operator usage_limits across task retries - #74396
Merged
Lee-W merged 5 commits intoOct 11, 2026
Merged
Conversation
Lee-W
marked this pull request as ready for review
October 7, 2026 08:13
Lee-W
force-pushed
the
llm-operator-usage-budget-across-retries
branch
6 times, most recently
from
October 9, 2026 07:39
93a7380 to
78c0754
Compare
kaxil
approved these changes
Oct 10, 2026
On Airflow 3.3+, AgentOperator already carries its spend across retries through the task state store, but LLMOperator and its subclasses counted usage_limits per attempt, so every retry got a fresh budget and a failing LLM task could spend retries x the configured limit. Applying the same cumulative budget to the LLM operators closes that gap, and sharing one implementation keeps the two operator families from drifting apart.
Clearing the persisted usage after a resumed review no longer coerces usage_limits or reads the task instance, so a usage_limits template that no longer renders to a valid value cannot fail a task the reviewer already approved. Remove UsageBudgetMixin._get_usage_budget; LLMOperator reads the task instance itself only when usage_limits is set.
The LLM* operators now count usage_limits across the whole task instance on Airflow 3.3+, but llm.rst still said that applied to AgentOperator only. Update the request and token limit bullets, note that clearing a running task instance keeps the old spend, and add an unreleased changelog note for the LLM* operators. The 0.11.0 AgentOperator note is restored to what shipped in providers-common-ai/0.11.0rc1. Also state in the LLMOperator docstring that nothing may fail the attempt once the usage budget is cleared, since a retry would start with a fresh budget. Subclasses that override execute or execute_complete are affected.
The build-budget test only checked the budget's type, so a wrong task state store or a wrong max_tries (which would reset the budget on every retry) went unnoticed. It now saves through the budget and asserts the record lands in the context's accessor with max_tries == 3. test_agent.py had its own copy of the task_state_store_accessor fixture from conftest.py; use the shared fixture and drop the copy.
Lee-W
force-pushed
the
llm-operator-usage-budget-across-retries
branch
from
October 11, 2026 06:57
78c0754 to
62a6a7a
Compare
The test saves through the budget, and save() imports NEVER_EXPIRE, which only exists on Airflow 3.3+. Older versions swallow the ImportError, so nothing reaches the store and the test fails there.
This was referenced Oct 14, 2026
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
On Airflow 3.3+, AgentOperator already carries its spend across retries through the task state store, but LLMOperator and its subclasses counted usage_limits per attempt, so every retry got a fresh budget and a failing LLM task could spend retries x the configured limit. Applying the same cumulative budget to the LLM operators closes that gap, and sharing one implementation keeps the two operator families from drifting apart.
Was generative AI tooling used to co-author this PR?
Generated-by: [Claude] following the guidelines
{pr_number}.significant.rst, in airflow-core/newsfragments. You can add this file in a follow-up commit after the PR is created so you know the PR number.