kaxil opened a new pull request, #76863: URL: https://github.com/apache/airflow/pull/76863
A durable agent task looks up every step in the task state store before running it, so a retry can replay what the previous attempt recorded. It did that on a first attempt too, where nothing can have been recorded: one blocking Execution API call per step, plus a lookup of the run's entry and a search past the last step at cleanup. With durable execution about to become the default for agent tasks, every agent run would pay for those lookups. This PR skips them for any agent run that no earlier attempt started, and deletes everything a successful task leaves in one `delete_steps` call. A first attempt now reads back only the list of agent runs (and, in `AgentOperator`, the run id). A retry still reads the journal back and replays as before. ## Measurements A scripted pydantic-ai `FunctionModel` agent (no LLM) against a local `airflow standalone` with Postgres, one task at a time. Each model request but the last calls a tool returning about 1 KB, so N requests make 2N-1 journal steps. Counts are every message the task process sent its supervisor, each of which is one Execution API call. Times are `execute()`, median of three runs, on one laptop, and vary by a few ms per step between runs. "Before" is the tip of the branch below this one. | Model requests (journal steps) | Execution API calls, before → after | Journal reads | Journal writes / deletes | Added per journal step, before → after | Of which waiting on journal calls | |---|---|---|---|---|---| | 10 (19) | 72 → **50** | 23 → **2** | 21 / 22 → 21 / 21 | 26.9 → 24.3 ms | 23.5 → 21.0 ms | | 50 (99) | 312 → **210** | 103 → **2** | 101 / 102 → 101 / 101 | 22.9 → 21.0 ms | 20.2 → 18.3 ms | | 200 (399) | 1,212 → **810** | 403 → **2** | 401 / 402 → 401 / 401 | 26.6 → 21.6 ms | 20.4 → 16.1 ms | The same agent with `durable=False` makes 7 calls whatever its length and takes 0.08, 0.21 and 1.07 s. A mapped agent task with 100 instances of 5 model requests each makes 4,200 → **3,000** calls across all instances (1,300 → 200 reads). The median instance's `execute()` went from 4.30 to 3.55 s (0.62 s without durable), but the mapped task, 16 instances at a time, took about as long overall (43 → 44 s, 24 s without durable). I did not measure where that time goes. The bytes written to the metadata DB don't change: 22 KB, 115 KB and 466 KB of journal for the three run lengths, and 1.03 MB for the mapped task. What the journal still costs is one write per step, plus one delete per step when the task succeeds. The measurement Dag counts by wrapping the task runner's supervisor channel; it is not part of this PR. ## Design rationale **The skip is decided per agent run.** A run's key goes onto the stored run list before any of its steps is written, so a run missing from the list has nothing to look up. That covers a first attempt, and also a run that a retry reaches for the first time because the earlier attempt failed before starting it. That only holds if the list really is written first. The list write is best effort like every other store write, and in the first version of this change a failed one left steps behind that the next retry would not look up. A step or the run's entry is now written only once its run is on the stored list. A failed list write is retried before the next write, and if it fails again the step is left unrecorded and runs again on retry, as when any other write fails. This costs no extra calls when the writes succeed. **`try_number == 1` would save the last two reads, and breaks deferred tasks.** A task that defers and resumes keeps its try number, and the resumed run has to read the journal back. **Cleanup makes one `delete_steps` call.** Every run's steps, their entries, the run list and the run id go into one `delete_steps` call, so the ObjectStorage journal (Airflow before 3.3) is rewritten once instead of once per run plus once per run-list update. The task state store API has no multi-key delete, and `clear()` would also delete the task's own keys in the same namespace, so on 3.3+ that call still makes one request per key. A first attempt also stops deleting run entries it never wrote. ## Large tool results Every step over `[state_store] max_value_storage_bytes` (64 KB by default) logs a warning from the Task SDK and is stored anyway. Once durable execution is the default, an agent with large tool results will log that warning on every such step. I'd keep the warning. It tells the user that the journal is putting large values in the metadata DB and names the setting that moves them out (`[workers] state_store_backend`). The alternative, leaving large results out of the journal, would make a retry run those tool calls again. The durable execution page now mentions the warning and the setting. ## Known issues - A task retried across the upgrade from 1.0.0 or 1.1.0 runs an agent run's steps live once if, under 1.x, that run's run-list write failed while its steps were written. Both failures are needed together. - One request per key at cleanup remains until the task state store API has a multi-key or prefix delete. ## Retries still replay The same scripted agent, with a tool that fails on try 1 (`retries=1`), on this branch. Try 1 read back two keys (the run id and the run list), recorded 7 steps and failed:  Try 2 read the journal back, replayed those 7 steps and ran only the failed tool call and the final model request:  -- 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]
