ccoliu opened a new pull request, #76201: URL: https://github.com/apache/airflow/pull/76201
`LLMBatchOperator` / `@task.llm_batch` and `invoke_durably()` (Strands) derive a run-stable identity key from the Dag, task and run ids and `map_index`. Every iteration of a task loop shares all four, so iterations shared one batch state record and one result file (see #76198): - same prompts every iteration: only iteration 1 was submitted, later iterations returned its results; - different prompts: each iteration cancelled the previous (finished) batch, and all iterations wrote the same `<key16>.jsonl`, so earlier iterations' results were overwritten; - different prompts with `on_stale_state="fail"`: iteration 2 failed; - `invoke_durably()` with a non-default storage: iteration 2 resumed iteration 1's agent session. This adds the loop's trajectory (`loop.region_id`) and iteration (`loop.index`) to the key when the task runs inside a loop: - Retries of one iteration, and clearing an iteration within the same trajectory, keep the same key, so reattaching to an in-flight batch still works. - A rerun that replaces later iterations starts a new trajectory and gets a new key, so it does not reattach to the replaced iteration's batch. - Outside a loop, the key is unchanged, so batches in flight across the upgrade are still found. Note this differs from what I proposed in #76198 (using `ti.region_id` / `ti.region_index`): on 3.4.0rc2 an ordinary mapped task outside any loop also has a non-zero `ti.region_id`, so that would have changed the key for every mapped task. The loop context is only present inside a loop. `DurableStorage` (the pre-3.3 journal) uses the same identity scheme but is not changed: task loops need Airflow 3.4, where the journal uses the task state store instead. Testing: - Unit tests for the key (unchanged outside a loop, distinct per iteration and trajectory, shared by retries), an operator test where two iterations with the same prompts each submit and land their own batch, and a Strands test mirroring the existing mapped-instance one. The operator and Strands tests fail without the fix. - End to end on `apache/airflow:3.4.0rc2` with the reproduction Dags from #76198 (offline fake batch adapter): before, 1 submission for 3 iterations, cancelled batches, overwritten results and a failing `on_stale_state="fail"` run; after, 3 submissions per Dag, no cancellations, one result file per iteration, all runs succeed. `strands` is not in the provider's test dependencies, so `test_strands.py` is skipped in CI; I ran it locally with `strands-agents` installed (17 passed). closes: #76198 --- ##### Was generative AI tooling used to co-author this PR? - [X] Yes (please specify the tool below) Generated-by: Claude Code (Claude Opus 5.5) following [the guidelines] -- 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]
