kaxil commented on code in PR #76862:
URL: https://github.com/apache/airflow/pull/76862#discussion_r4550715904
##########
providers/common/ai/src/airflow/providers/common/ai/durable/journal.py:
##########
Review Comment:
Done in 77aa7697e56, pulled both out with your names.
`_find_failed_steps(error)` also hands back the causes and the earliest
position a failing step had reached, since `fail()` needs both. No behaviour
change, the existing tests pass as they were.
##########
providers/common/ai/src/airflow/providers/common/ai/durable/journal.py:
##########
@@ -115,6 +128,8 @@ class JournalStep:
payload: Any = None
# Whether this step's replay is in ``stats.replayed``; see
:meth:`DurableRun.reserve`.
_replay_counted: bool = field(default=False, init=False, repr=False)
+ # The position of the last model step claimed before this one: the request
that chose a tool call.
+ _model_position: int | None = field(default=None, init=False, repr=False)
Review Comment:
Good call, renamed to `_choosing_model_position` in 77aa7697e56 and reworded
the comment to match.
##########
providers/common/ai/src/airflow/providers/common/ai/durable/journal.py:
##########
@@ -519,17 +583,47 @@ def pause(self) -> None:
log.warning("Durable: could not delete entries this attempt
replaced", positions=survived)
trusted_until = survived[0] if survived else self._position
self._diverge_from, self._model_at_fault = trusted_until, False
- self._save_meta(trusted_until, model_at_fault=False, paused=True)
+ # A pause is not a failure, and the task resumes in the same attempt,
so the count of
+ # attempts in a row a tool call failed carries over to a failure after
the resume.
+ self._save_meta(
+ trusted_until, model_at_fault=False, paused=True,
tool_failure=self._tool_failure or None
+ )
- def _save_meta(self, diverge_from: int | None, *, model_at_fault: bool,
paused: bool = False) -> None:
+ def _save_meta(
+ self,
+ diverge_from: int | None,
+ *,
+ model_at_fault: bool,
+ paused: bool = False,
+ tool_failure: Mapping[str, Any] | None = None,
+ ) -> None:
meta = {
"diverge_from": diverge_from,
"high_water": max(self._high_water, self._position),
"model_at_fault": model_at_fault,
"paused": paused,
+ "tool_failure": tool_failure,
}
self.journal.storage.save_step(build_run_meta_key(self.key), meta)
+ def _count_tool_failure(self, step: JournalStep | None) -> _ToolFailure |
None:
Review Comment:
It does count though: it returns the previous attempt's `attempts` plus one
when the call matches, or 1 when it doesn't. So I went with
`_next_tool_failure` in 77aa7697e56, which says it builds the next record.
Happy to switch if you still prefer `_compute_tool_failure`.
##########
providers/common/ai/tests/unit/common/ai/durable/test_journal.py:
##########
@@ -308,6 +309,189 @@ async def
test_a_step_without_a_fingerprint_does_not_replay_after_a_live_step(se
assert summary == "summary of v2"
+async def model_then_failing_tool(
Review Comment:
Renamed to `run_model_then_failing_tool`, `run_attempt_failing_in_tool` and
`collect_ask_model_again_attempts` in 77aa7697e56.
##########
providers/common/ai/src/airflow/providers/common/ai/durable/journal.py:
##########
@@ -80,6 +82,17 @@
StepKind = Literal["model", "tool", "other"]
"""What a step is, for the end-of-run summary and for adapters that treat
model steps specially."""
+# Attempts in a row a tool call may fail before a retry asks the model again
instead of replaying
+# its choice: the first retry runs the same call again, in case the failure
was transient.
Review Comment:
Shortened in 77aa7697e56, but kept "runs": "replays the call" reads like it
replays the recorded result, and the first retry actually runs the call again.
It's now "the first retry runs the same call again, in case the failure was
transient".
##########
providers/common/ai/src/airflow/providers/common/ai/durable/journal.py:
##########
@@ -461,25 +483,67 @@ def fail(self, error: BaseException, *, model_at_fault:
bool = False) -> None:
checks too many times. Those answers were recorded and would
replay into the same
failure on every retry, so the retry asks the model again from its
last request.
Ignored when a step raised ``error``: that step decides where the
retry runs live.
+
+ A tool call that raises ``error`` on two attempts in a row, with the
same arguments, is
+ treated as a failure the model caused as well: the first retry runs
the call again in
+ case its failure was transient, and the next asks the model again from
the request that
+ chose the call, rather than replay that choice into the same failure.
"""
failed_at = self._position
- raised_in_step = False
+ failed_steps: list[JournalStep] = []
+ causes: list[BaseException] = []
cause: BaseException | None = error
while cause is not None:
- positions = [position for raised, position in self._raised if
raised is cause]
- raised_in_step = raised_in_step or bool(positions)
- failed_at = min([failed_at, *positions])
+ causes.append(cause)
+ for raised, position, step in self._raised:
+ if raised is cause:
+ failed_steps.append(step)
+ failed_at = min(failed_at, position)
cause = cause.__cause__ or cause.__context__
self._load_meta()
# An attempt that was still replaying an earlier one when it failed,
and recorded
# nothing from the failure on, leaves that earlier attempt's steps as
they were.
diverge_from = failed_at if self.diverged or self._last_saved >=
failed_at else self._diverge_from
+ failed_step = min(failed_steps, key=lambda step: step.position,
default=None)
+ failures = self.journal._run_failures
+ # An agent run started inside a step that raised failed with the error
first and took it
+ # as its own: its tool call failed, or its model is asked again. This
run's model did not
+ # choose what failed. Only runs numbered under those steps count (a
tool can raise through
+ # an operation step it started); another run whose error is merely in
this one's context,
+ # such as a fallback agent started in an ``except``, does not.
Review Comment:
Shortened in 77aa7697e56. It moved into `_is_failure_owned_by_nested_run`
with the extraction, so the docstring says what owning means and the comment is
down to two lines about the fallback agent case. I didn't use "already
counted": the nested run may have asked its model again instead of counting a
tool failure, so the docstring says it took the error as its own.
##########
providers/common/ai/src/airflow/providers/common/ai/durable/journal.py:
##########
@@ -80,6 +82,17 @@
StepKind = Literal["model", "tool", "other"]
"""What a step is, for the end-of-run summary and for adapters that treat
model steps specially."""
+# Attempts in a row a tool call may fail before a retry asks the model again
instead of replaying
+# its choice: the first retry runs the same call again, in case the failure
was transient.
+_TOOL_FAILURES_BEFORE_ASKING_MODEL = 2
+
+
+class _ToolFailure(TypedDict):
+ """The tool call that failed a run, in the run's entry, and on how many
attempts in a row it did."""
+
+ call: list[int | str | None] # [position, name, fingerprint]
Review Comment:
Yeah, the positional list was confusing. It lives in the run's JSON entry
though, so a namedtuple would be written as a list and read back as a plain
list, and the comparison with the next attempt's call would need a conversion
on read. Went with a small `_ToolCall` TypedDict with `position`, `name` and
`fingerprint` keys instead, which reads back the same as it was written. In
24e4aa44f10.
--
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]
