Lee-W commented on code in PR #76862:
URL: https://github.com/apache/airflow/pull/76862#discussion_r4548745366
##########
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:
should we make this a namedtuple? this looks a bit confusing
##########
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.
+ inside = tuple(f"{self.key}.{step.position}." for step in failed_steps)
+ failed_inside = any(
+ seen is cause and key.startswith(inside) for seen, key in failures
for cause in causes
+ )
+ tool_failure = None if failed_inside else
self._count_tool_failure(failed_step)
+ # The model request a retry asks again, if the failure is the model's.
+ ask_model_from: int | None = None
+ if failed_step is None:
+ if model_at_fault:
+ # Unlike a step that raised, the failure is known to come
after every recorded
+ # step, so this holds even for an attempt that was still
replaying.
+ ask_model_from = self._last_model_position
+ elif (
+ tool_failure is not None
+ and tool_failure["attempts"] >= _TOOL_FAILURES_BEFORE_ASKING_MODEL
+ and failed_step._model_position is not None
+ ):
+ ask_model_from = failed_step._model_position
+ log.warning(
+ "Durable: the same tool call failed on consecutive attempts; a
retry asks the model "
+ "again from the request that chose it",
+ position=failed_step.position,
+ step=failed_step.name,
+ attempts=tool_failure["attempts"],
+ )
+ # The count starts over with the model's next answer, even one
whose call cannot be
+ # told from this one (no fingerprint), so that call too is retried
once as it is.
+ tool_failure = None
+ if failed_inside or tool_failure is not None or ask_model_from is not
None:
+ failures.append((error, self.key))
rerun_model = False
- if model_at_fault and not raised_in_step and (last_model :=
self._last_model_position) is not None:
- # Unlike a step that raised, the failure is known to come after
every recorded step,
- # so this holds even for an attempt that was still replaying.
- if diverge_from is None or last_model <= diverge_from:
- diverge_from, rerun_model = last_model, True
+ if ask_model_from is not None and (diverge_from is None or
ask_model_from <= diverge_from):
Review Comment:
Once a retry asks the model again, the run entry keeps `diverge_from` and
`model_at_fault=True` even after that attempt succeeds, because nothing resets
them on success. A later attempt of the same journal (e.g. an Airflow retry
after a post-run failure) then re-runs the model and tools live instead of
replaying the successful attempt. In a probe, attempt 4 after a successful
attempt 3 made 2 model calls and 1 tool call; base made none.
Could we reset the marker after a successful run, writing only when a marker
is set, and add a test that one more attempt after the ask-again success makes
no live calls?
--
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]
