kaxil commented on code in PR #76917:
URL: https://github.com/apache/airflow/pull/76917#discussion_r4550720886
##########
providers/common/ai/src/airflow/providers/common/ai/operators/agent.py:
##########
@@ -1260,19 +1366,68 @@ def _run_and_report_on_failure(
}
raise
+ def _decide_durable(self, context: Context) -> None:
+ """
+ Decide whether this run is durable, and log why when the default
turned it off.
+
+ ``durable=True`` and ``durable=False`` decide it themselves. The
default turns it on
+ where a retry can replay the journal and every step can be recorded,
and off, with the
+ reason in the task log, everywhere else.
+ """
+ if self.durable is not None:
+ self._durable = self.durable
+ return
+ if _contains_capability(self._declared_capabilities,
AirflowDurability):
+ # Inside a running task the capability opens a journal of its own
(current_journal).
+ self._durable = False
+ self.log.info(
+ "Durable execution is left to the AirflowDurability capability
the agent was given, "
+ "which deletes each agent run's steps as soon as that run
succeeds."
+ )
+ return
+ reason = self._why_not_durable_by_default(context)
+ self._durable = reason is None
+ if reason is not None:
+ self.log.info("Durable execution is off for this task: %s", reason)
+
+ def _why_not_durable_by_default(self, context: Context) -> str | None:
+ """Say why durable execution cannot be on by default for this task, or
return ``None``."""
+ # From 3.4 a retry policy can run a task with no retries again
(RetryAction.RESCHEDULE);
+ # before that it cannot, and before 3.3 BaseOperator has no
retry_policy.
+ if not self.retries and (not AIRFLOW_V_3_4_PLUS or self.retry_policy
is None):
+ return "the task has no retries, so no attempt would replay this
one."
+ if not AIRFLOW_V_3_3_PLUS:
+ return (
+ "it is on by default only from Airflow 3.3. Before that, pass
durable=True and set "
+ "[common.ai] durable_cache_path."
+ )
+ if "task_state_store" not in context:
+ return "the task's context has no task state store to keep the
journal in."
+ return self._durable_problem()
+
+ @property
+ def _durable_mode(self) -> str:
+ """How this run's durable execution was turned on, for the summary
line."""
+ return "on by default" if self.durable is None else "durable=True"
+
+ def _span_attributes(self, ti: Any) -> dict[str, Any]:
+ """Attributes for every span the agent emits: the task instance, and
whether the run is durable."""
+ return {**build_run_identity_attributes(ti),
"airflow.common_ai.durable": self._durable}
+
def _journal_scope(self) -> AbstractContextManager[DurableJournal | None]:
"""Run the code in the ``with`` block under the task's journal, when
the task is durable."""
return journal_scope(self._durable_journal) if self._durable_journal
is not None else nullcontext()
def execute(self, context: Context) -> Any:
# Coerced first so a bad rendered value fails before the expensive
setup below.
usage_limits = coerce_usage_limits(self.usage_limits)
+ self._decide_durable(context)
# A try that paused and then ended some other way (marked failed while
waiting, a failed
# request) leaves its transcript behind; a fresh run never reads it.
``.get``: a context
# built by hand in a unit test has no task state store, and nothing to
clean up.
if (
- not self.durable
+ not self._durable
Review Comment:
Real one, a transcript from a try before the upgrade stayed behind after a
durable-by-default retry succeeded. Fixed in 80f8e53baaf: every fresh run now
deletes it, durable or not. A durable run never reads that key, and the
non-durable resume doesn't go through `execute`.
`test_a_fresh_run_deletes_a_stale_transcript` now also runs with `retries=2`
and no `durable`, and that case failed before the change.
##########
providers/common/ai/src/airflow/providers/common/ai/operators/agent.py:
##########
@@ -1260,19 +1366,68 @@ def _run_and_report_on_failure(
}
raise
+ def _decide_durable(self, context: Context) -> None:
+ """
+ Decide whether this run is durable, and log why when the default
turned it off.
+
+ ``durable=True`` and ``durable=False`` decide it themselves. The
default turns it on
+ where a retry can replay the journal and every step can be recorded,
and off, with the
+ reason in the task log, everywhere else.
+ """
+ if self.durable is not None:
+ self._durable = self.durable
+ return
+ if _contains_capability(self._declared_capabilities,
AirflowDurability):
+ # Inside a running task the capability opens a journal of its own
(current_journal).
+ self._durable = False
+ self.log.info(
+ "Durable execution is left to the AirflowDurability capability
the agent was given, "
+ "which deletes each agent run's steps as soon as that run
succeeds."
+ )
+ return
+ reason = self._why_not_durable_by_default(context)
+ self._durable = reason is None
+ if reason is not None:
+ self.log.info("Durable execution is off for this task: %s", reason)
+
+ def _why_not_durable_by_default(self, context: Context) -> str | None:
+ """Say why durable execution cannot be on by default for this task, or
return ``None``."""
+ # From 3.4 a retry policy can run a task with no retries again
(RetryAction.RESCHEDULE);
+ # before that it cannot, and before 3.3 BaseOperator has no
retry_policy.
+ if not self.retries and (not AIRFLOW_V_3_4_PLUS or self.retry_policy
is None):
+ return "the task has no retries, so no attempt would replay this
one."
+ if not AIRFLOW_V_3_3_PLUS:
+ return (
+ "it is on by default only from Airflow 3.3. Before that, pass
durable=True and set "
+ "[common.ai] durable_cache_path."
+ )
+ if "task_state_store" not in context:
+ return "the task's context has no task state store to keep the
journal in."
+ return self._durable_problem()
+
+ @property
+ def _durable_mode(self) -> str:
+ """How this run's durable execution was turned on, for the summary
line."""
+ return "on by default" if self.durable is None else "durable=True"
+
+ def _span_attributes(self, ti: Any) -> dict[str, Any]:
+ """Attributes for every span the agent emits: the task instance, and
whether the run is durable."""
+ return {**build_run_identity_attributes(ti),
"airflow.common_ai.durable": self._durable}
+
def _journal_scope(self) -> AbstractContextManager[DurableJournal | None]:
"""Run the code in the ``with`` block under the task's journal, when
the task is durable."""
return journal_scope(self._durable_journal) if self._durable_journal
is not None else nullcontext()
def execute(self, context: Context) -> Any:
# Coerced first so a bad rendered value fails before the expensive
setup below.
usage_limits = coerce_usage_limits(self.usage_limits)
+ self._decide_durable(context)
Review Comment:
Done in 80f8e53baaf. Every test in `TestDurable` now runs twice, with
`durable=True` and with `retries=2` and `durable` left unset, the same way the
HITL review tests do. The approve test and the retry-after-a-failed-resume test
also check that success leaves no journal or saved answers behind (before 3.4
only the note that it asked stays).
--
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]