amoghrajesh commented on code in PR #76862:
URL: https://github.com/apache/airflow/pull/76862#discussion_r4548499364
##########
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:
Hmm sounds like "this step's model position". But it is really the position
of the model request that chose this step. Can we rename to
`_choosing_model_position`.
##########
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:
nit: rename to `run_model_then_failing_tool`? Same foer
`attempt_failing_in_tool` and `asked_model_again`
##########
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:
```suggestion
# A nested run under a raising step already counted this error. Only
runs keyed under those steps count, not a fallback run with this error in its
context.
```
Too long, can we shorten it?
##########
providers/common/ai/src/airflow/providers/common/ai/durable/journal.py:
##########
Review Comment:
`fail()` now walks the cause chain, scans for nested runs, counts the
failure, logs, and picks where to diverge, all inline. It is the decision point
for every adapter, and it is now hard to change safely. Should we pull out
`_find_failed_steps(error)` and `_is_failure_owned_by_nested_run(failed_steps,
causes)` instead?
##########
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:
nit: call it `_compute_tool_failure` since no counting is involved?
##########
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:
```suggestion
# The first retry replays the call as is, in case the failure was transient.
```
Shorten to this?
--
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]
