kaxil commented on code in PR #74396:
URL: https://github.com/apache/airflow/pull/74396#discussion_r4237634405
##########
providers/common/ai/docs/changelog.rst:
##########
@@ -29,18 +29,21 @@ Changelog
......
.. note::
- On Airflow >= 3.3, ``AgentOperator``'s ``usage_limits`` counts usage across
every
- attempt of the task instance combined -- initial run, every retry, and every
HITL
- regeneration add to one running total (including the implicit
``request_limit=50``
- default), instead of each attempt starting fresh. Scale each limit by
``retries + 1``,
- or set ``usage_limits=None``, to keep the old per-attempt headroom. On every
Airflow
+ On Airflow >= 3.3, the ``usage_limits`` of ``AgentOperator`` and the
``LLM*`` operators
Review Comment:
`providers-common-ai/0.11.0rc1` was already tagged without this change, so
the 0.11.0 section isn't the right place for it. Could the AgentOperator note
stay as it shipped, with a new note for the `LLM*` operators just under the
`Changelog` header (per the note to contributors at the top)? Merging them also
makes the durable-replay and failed-attempt `usage` XCom sentences read as if
they apply to the `LLM*` operators, which they don't.
##########
providers/common/ai/docs/operators/llm.rst:
##########
@@ -144,23 +144,25 @@ When the limit is hit pydantic-ai raises
``UsageLimitExceeded``, which
propagates to Airflow as a task failure, so Airflow's standard retry policy
applies on top.
-For the ``LLM*`` operator family, every limit here bounds a single agent
-*run*, not a task: each Airflow task retry re-renders ``usage_limits`` and
-starts a fresh count. A ``cost_limit`` of ``Decimal("0.50")`` caps one run,
-so it is not a bound on what the task spends in total.
-
-For ``AgentOperator`` on Airflow >= 3.3, it is the other way around: setting
-``usage_limits`` makes every limit bound the whole task instance -- the
-initial run, every retry, and every HITL regeneration all add to one running
-total, so a ``cost_limit`` of ``Decimal("0.50")`` bounds what the task spends
-across every attempt combined (best-effort, like the rest of ``cost_limit``
-above -- a SIGKILL loses the spend of the attempt it kills, and this is not
-a hard guarantee). See :ref:`howto/operator:agent` for how that budget is
-persisted and how to reset it, and for what a HITL regeneration shares. On
-Airflow < 3.3, and whenever ``usage_limits`` is ``None``, ``AgentOperator``
-behaves like the ``LLM*`` family: each attempt starts a fresh count. Within
-one attempt, its HITL regenerations share the initial run's count whenever
-``usage_limits`` is set, on every Airflow version.
+On Airflow >= 3.3, setting ``usage_limits`` makes every limit bound the whole
+task instance, for the ``LLM*`` family (``LLMOperator``, ``LLMBranchOperator``,
+``LLMSQLQueryOperator``, ``LLMSchemaCompareOperator``,
``LLMFileAnalysisOperator`` and
+the matching ``@task.*`` decorators) the same way it does for
``AgentOperator``:
+the initial run and every retry add to one running total, including the
implicit
+``request_limit=50``. A pause for approval, and a retry after a reviewer
rejects the
+output, share the same total. The total is cleared only once the task succeeds
-- a
+reviewer approval, or a ``LLMBranchOperator`` skip, counts as success -- and a
+cleared task instance starts a fresh one. A ``cost_limit`` of
``Decimal("0.50")``
Review Comment:
Clearing a running task instance does not bump `max_tries`, so the restarted
attempt keeps the old spend. agent.rst mentions this; can we add the same note
here? Also, the `request_limit` and `input_tokens_limit` /
`output_tokens_limit` bullets at lines 116-119 still say the per-task-instance
count only applies to `AgentOperator`.
##########
providers/common/ai/src/airflow/providers/common/ai/mixins/usage_budget.py:
##########
@@ -0,0 +1,164 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+
+from __future__ import annotations
+
+from collections.abc import Callable
+from typing import TYPE_CHECKING, Any, Protocol
+
+from airflow.providers.common.ai.utils.usage import coerce_usage_limits
+from airflow.providers.common.ai.utils.usage_budget import (
+ TaskStateStoreUsageBudget,
+ copy_run_usage,
+ subtract_run_usage,
+)
+from airflow.providers.common.compat.version_compat import AIRFLOW_V_3_3_PLUS
+
+if TYPE_CHECKING:
+ from logging import Logger
+
+ from pydantic_ai import Agent
+ from pydantic_ai.agent import AgentRunResult
+ from pydantic_ai.usage import RunUsage, UsageLimits
+
+ from airflow.sdk import Context
+
+
+class UsageBudgetHostProtocol(Protocol):
+ """Protocol for the operator attributes and methods
:class:`UsageBudgetMixin` relies on."""
+
+ log: Logger
+ usage_limits: UsageLimits | dict[str, Any] | None
+ _usage_budget: TaskStateStoreUsageBudget | None
+
+ def run_agent_sync(
+ self, agent: Agent[Any, Any], user_prompt: Any, **run_kwargs: Any
+ ) -> AgentRunResult[Any]: ...
+
+ def _build_usage_budget(
+ self, context: Context, usage_limits: UsageLimits | None, *, ti: Any
+ ) -> TaskStateStoreUsageBudget | None: ...
+
+ def _get_usage_budget(
+ self, context: Context, usage_limits: UsageLimits | None
+ ) -> TaskStateStoreUsageBudget | None: ...
+
+ def _settle_tracked_usage(self) -> None: ...
+
+
+class UsageBudgetMixin:
+ """
+ Mixin that bounds ``usage_limits`` across every attempt of a task instance.
+
+ On Airflow >= 3.3 with ``usage_limits`` set, cumulative usage is persisted
in the
+ task state store (see
:class:`~airflow.providers.common.ai.utils.usage_budget.TaskStateStoreUsageBudget`)
+ so a task retry resumes from the usage earlier attempts already spent
instead of starting a
+ fresh count.
+
+ Operators that use this mixin must provide ``usage_limits``, ``log`` and
+ ``run_agent_sync`` (see
:class:`~airflow.providers.common.ai.mixins.cancellable_run.CancellableAgentRunMixin`).
+ """
+
+ _usage_budget: TaskStateStoreUsageBudget | None = None
+
+ def _build_usage_budget(
+ self, context: Context, usage_limits: UsageLimits | None, *, ti: Any
+ ) -> TaskStateStoreUsageBudget | None:
+ """
+ Return the cross-attempt usage-budget accessor, or ``None`` when it
should not apply.
+
+ Gated like ``_build_durable_storage``: only on Airflow >= 3.3, where
the task
+ state store survives retries. Also gated on ``usage_limits is not
None`` --
+ with ``usage_limits=None`` turning this on would silently impose
pydantic-ai's
+ default ``request_limit=50`` across every attempt of every operator on
+ 3.3+, which nobody asked for.
+
+ To switch the gate off in a test, patch ``AIRFLOW_V_3_3_PLUS`` in this
module.
+ """
+ if not (AIRFLOW_V_3_3_PLUS and usage_limits is not None):
+ return None
+ return TaskStateStoreUsageBudget(context["task_state_store"],
max_tries=ti.max_tries)
+
+ def _get_usage_budget(
+ self, context: Context, usage_limits: UsageLimits | None
+ ) -> TaskStateStoreUsageBudget | None:
+ """Return the usage budget for ``context``, reading the task instance
only when a budget can apply."""
+ # Contexts built without a ``task_instance`` stay valid for runs that
have no ``usage_limits``.
+ ti = context["task_instance"] if usage_limits is not None else None
+ return self._build_usage_budget(context, usage_limits, ti=ti)
+
+ def _settle_tracked_usage(self) -> None:
+ """Do nothing; override to run after a tracked run and before its
usage is persisted."""
+
+ def _run_tracked(
+ self: UsageBudgetHostProtocol,
+ agent: Agent[Any, Any],
+ prompt: Any,
+ *,
+ run_usage: RunUsage,
+ before_run: Callable[[], None] | None = None,
+ **run_kwargs: Any,
+ ) -> tuple[AgentRunResult[Any], RunUsage]:
+ """
+ Run the agent, persisting cumulative usage after every attempt
(success or failure).
+
+ ``run_usage`` is the (possibly cross-attempt) cumulative total the run
is seeded
+ with; the returned ``RunUsage`` is what this attempt alone added to it.
+ ``before_run`` runs after the snapshot is taken, so anything it
credits to
+ ``run_usage`` never shows up in this attempt's delta.
+ """
+ base = copy_run_usage(run_usage)
+ try:
+ if before_run is not None:
+ before_run()
+ result = self.run_agent_sync(agent, prompt, usage=run_usage,
**run_kwargs)
+ finally:
+ self._settle_tracked_usage()
+ if self._usage_budget:
+ self._usage_budget.save(run_usage)
+ return result, subtract_run_usage(run_usage, base)
+
+ def _log_cumulative_usage(self: UsageBudgetHostProtocol, run_usage:
RunUsage) -> None:
+ """Log the usage accumulated across all attempts of this task
instance."""
+ self.log.info(
+ "Cumulative usage across attempts: requests=%s, tool_calls=%s,
input_tokens=%s, "
+ "output_tokens=%s, total_tokens=%s",
+ run_usage.requests,
+ run_usage.tool_calls,
+ run_usage.input_tokens,
+ run_usage.output_tokens,
+ run_usage.total_tokens,
+ )
+ if run_usage.cost is not None:
+ self.log.info(
+ "Cumulative cost across attempts: $%s (USD, best-effort)",
+ format(run_usage.cost, "f"),
+ )
+
+ def _clear_usage_budget(self: UsageBudgetHostProtocol, context: Context)
-> None:
+ """
+ Drop the persisted usage so the next task instance starts with a fresh
budget.
+
+ A resumed operator (``execute_complete``) is a new instance whose
``_usage_budget``
+ was never set, so the budget is rebuilt from ``usage_limits`` in that
case.
+ """
+ budget = self._usage_budget
+ if budget is None:
+ budget = self._get_usage_budget(context,
coerce_usage_limits(self.usage_limits))
Review Comment:
On a resumed instance this coerces `self.usage_limits` and reads
`context["task_instance"]` before `_build_usage_budget` ever reaches the 3.3
check. Templates are re-rendered on resume, so if `usage_limits` comes from a
Variable that was emptied or mistyped while the review was open, a task the
reviewer just approved fails with a `ValueError`. On Airflow < 3.3 there's
nothing to clear at all. Deleting the key doesn't need the limits, so could
this check `AIRFLOW_V_3_3_PLUS and self.usage_limits is not None` and skip the
coerce? That would also let `_get_usage_budget` go, since its only job is
tolerating test contexts that have no `task_instance`.
##########
providers/common/ai/tests/unit/common/ai/mixins/test_usage_budget.py:
##########
@@ -0,0 +1,183 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+from __future__ import annotations
+
+import logging
+from decimal import Decimal
+from unittest import mock
+from unittest.mock import MagicMock
+
+import pytest
+from pydantic_ai.usage import RunUsage, UsageLimits
+
+from airflow.providers.common.ai.mixins.usage_budget import UsageBudgetMixin
+from airflow.providers.common.ai.utils.usage_budget import
TaskStateStoreUsageBudget
+
+MIXIN_MODULE = "airflow.providers.common.ai.mixins.usage_budget"
+
+
+class _Host(UsageBudgetMixin):
+ def __init__(self, usage_limits=None):
+ self.usage_limits = usage_limits
+ self.log = logging.getLogger(__name__)
+ self.run_agent_sync = MagicMock()
+
+
+class TestBuildUsageBudget:
+ @mock.patch(f"{MIXIN_MODULE}.AIRFLOW_V_3_3_PLUS", True)
+ def test_builds_budget_from_task_state_store_and_max_tries(self):
+ context = {"task_state_store": MagicMock()}
+ ti = MagicMock(max_tries=3)
+
+ budget = _Host()._build_usage_budget(context,
UsageLimits(request_limit=5), ti=ti)
+
+ assert isinstance(budget, TaskStateStoreUsageBudget)
Review Comment:
This passes even if the wrong store or the wrong `max_tries` is passed
through, and a wrong `max_tries` would quietly reset the budget on every retry.
Could it save through the budget and assert the record lands in
`context["task_state_store"]` with `max_tries == 3`?
##########
providers/common/ai/src/airflow/providers/common/ai/operators/llm.py:
##########
@@ -364,12 +369,41 @@ def execute_complete(
decision: dict[str, Any] | None = None,
) -> Any:
"""Resume after human review and restore the Pydantic model for XCom
consumers."""
+ output = self._complete_review(context, generated_output, event,
decision)
+ self._clear_usage_budget(context)
Review Comment:
If a subclass validates after `super().execute_complete()`, the budget is
already cleared when the validation fails, so the retry gets a fresh budget. I
think that is why `LLMSQLQueryOperator` and `LLMSchemaCompareOperator` moved to
`_complete_review` here. Subclasses outside this repo that override
`execute_complete` the way these two used to will hit the same thing. Can the
class docstring say that nothing should raise after `_clear_usage_budget` is
called?
##########
providers/common/ai/tests/unit/common/ai/conftest.py:
##########
@@ -99,3 +102,48 @@ def _make(output, *, cost=None):
return mock_result
return _make
+
+
[email protected]
+def task_state_store_accessor():
Review Comment:
This is a line-for-line copy of `_make_task_state_store_accessor` in
`operators/test_agent.py`. Could test_agent.py use this fixture instead, so
there's one copy to keep in sync?
--
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]
