kaxil commented on code in PR #76862:
URL: https://github.com/apache/airflow/pull/76862#discussion_r4550718226
##########
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:
Good catch, reproduced it: with the journal kept after the run (as
`AgentOperator` does until its post-run XCom work succeeds), attempt 4 asked
the model again and re-ran the query. Fixed in 24e4aa44f10.
`DurableRun.complete()` now takes the note down when a run succeeds and the
task keeps the journal, and writes nothing when there is no note. It also
deletes entries an earlier attempt left where this one recorded nothing, as
`pause()` does. The new test runs one more attempt after the ask-again success
and checks it makes no model or tool calls and no store writes.
--
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]
