ccoliu opened a new issue, #76198:
URL: https://github.com/apache/airflow/issues/76198
### Under which category would you file this issue?
Providers
### Apache Airflow version
3.4.0rc2
### What happened and how to reproduce it?
`LLMBatchOperator` (and `@task.llm_batch`) keys its run-stable batch state
on `compute_identity_key(dag_id, task_id, run_id, map_index)`. Inside a task
loop (`@task_group(...).loop(...)`), `map_index` is `-1` for every iteration,
so all iterations of the same task get the same key. The key is meant to be
shared by retries of one task instance, but iterations are different task
instances, and they end up sharing one state file under `_airflow_batch_state/`
and one result file (`<key16>.jsonl`).
Running a 3-iteration loop with an offline fake batch adapter (full Dag
below, no API key needed):
**1. Same prompt every iteration: only iteration 1 is submitted.** The
provider received a single submission (`batch-9205026e`), and all three
iterations returned it. A "regenerate until a judge accepts" loop gets the
first answer every time and never converges.
```
iteration 1: batch_id 'batch-9205026e'
iteration 2: batch_id 'batch-9205026e'
iteration 3: batch_id 'batch-9205026e'
```
**2. Different prompt per iteration (default
`on_stale_state="cancel_and_resubmit"`): earlier results are overwritten.**
Each iteration submits a new batch, but first cancels the previous iteration's
already-completed batch and logs a misleading warning ("If this happens on
every retry of a @task.llm_batch task, the decorated function is not
deterministic across attempts..."). All three iterations return the same
`result_uri` (`file:///tmp/fakellm_results/71c3dc4cd0d5f91c.jsonl`), which
afterwards only contains iteration 3's output:
```
{"custom_id": "71c3dc4cd0d5f91c-0", "index": 0, "status": "success",
"output": "answer to 'Write slogan version 3' from batch-dc9fdbff", ...}
```
So a task outside the loop reading the loop's results (or anyone following
iteration 1's `result_uri`) silently gets iteration 3's data.
**3. Different prompt per iteration with `on_stale_state="fail"`: iteration
2 fails.**
```
LLMBatchStaleStateError: Recorded batch 'batch-e3d1e2c9' for this task
instance no longer matches its input (the prompts or other request content
changed since the recorded batch); on_stale_state='fail' rejects it ...
```
`invoke_durably()` (Strands) derives its default `session_id` from the same
key, so it likely has the same problem when used with a custom (non
task-state-store) `storage`. I have not reproduced that part.
<details><summary>Reproduction Dag (fake batch adapter, runs
offline)</summary>
Requires a connection `fake_llm` of type `pydanticai`.
```python
from __future__ import annotations
import json
import uuid
from pathlib import Path
from airflow.providers.common.ai.batch.base import (
BatchAdapter,
BatchState,
ExtractedOutput,
RawResultItem,
SubmitResult,
)
from airflow.providers.common.ai.batch.dispatch import register_adapter
from airflow.providers.common.ai.operators.llm_batch import LLMBatchOperator
from airflow.sdk import dag, task, task_group
STORE = Path("/tmp/fakellm")
class FakeLLMAdapter(BatchAdapter):
"""Offline batch "provider": every submit is a new batch whose answers
name that batch."""
name = "fakellm"
conn_types = frozenset({"pydanticai"})
max_requests = 100
max_payload_bytes = 10_000_000
allows_per_request_model = False
def validate_requests(self, requests, *, model, output_spec,
idempotency_key, **kwargs):
return None
def submit(self, requests, *, model, idempotency_key, input_fingerprint,
output_spec, **kwargs):
STORE.mkdir(exist_ok=True)
batch_id = f"batch-{uuid.uuid4().hex[:8]}"
prompts = [r["prompt"] for r in requests]
(STORE / f"{batch_id}.json").write_text(json.dumps({"key16":
idempotency_key, "prompts": prompts}))
with (STORE / "submissions.log").open("a") as f:
f.write(f"{batch_id} {prompts}\n")
return SubmitResult(batch_id=batch_id, provider_input_ref=None)
def get_batch(self, batch_id):
return BatchState(status="completed", counts={"succeeded": 1},
error_message=None)
def cancel_batch(self, batch_id):
with (STORE / "cancellations.log").open("a") as f:
f.write(f"{batch_id}\n")
def iter_results(self, batch_id):
data = json.loads((STORE / f"{batch_id}.json").read_text())
return iter(
[
RawResultItem(
custom_id=f"{data['key16']}-{i}",
index=i,
provider_status="success",
model="echo",
usage=None,
finish_reason="stop",
error=None,
raw=f"answer to {prompt!r} from {batch_id}",
)
for i, prompt in enumerate(data["prompts"])
]
)
def build_output_directive(self, spec):
return {}
def extract_output(self, raw, spec):
return ExtractedOutput(kind="text", text=raw.raw)
def find_orphaned_batch(self, idempotency_key, input_fingerprint,
not_before):
return None
register_adapter(FakeLLMAdapter)
COMMON = dict(
llm_conn_id="fake_llm",
model_id="fakellm:echo",
result_path="file:///tmp/fakellm_results",
poll_interval=30,
deferrable=False,
)
def make_dag(dag_id, *, changing_prompt,
on_stale_state="cancel_and_resubmit"):
@dag(dag_id=dag_id, schedule=None, catchup=False)
def _dag():
@task_group
def generate():
@task
def build_prompts(*, loop):
if changing_prompt:
return [f"Write slogan version {loop.index}"]
return ["Write a catchy slogan"]
LLMBatchOperator(
task_id="llm", requests=build_prompts(),
on_stale_state=on_stale_state, **COMMON
)
@task
def report(results):
for i, r in enumerate(results, start=1):
print(f"iteration {i}: {r}")
loop_ = generate.loop(max_iterations=3)
report(loop_["llm"].output)
return _dag()
make_dag("llm_batch_same_prompt", changing_prompt=False)
make_dag("llm_batch_changing_prompt", changing_prompt=True)
make_dag("llm_batch_changing_prompt_fail", changing_prompt=True,
on_stale_state="fail")
```
</details>
### What you think should happen instead?
Each loop iteration should get its own identity, while retries of the same
iteration keep sharing one (so reattaching to an in-flight batch after a crash
still works).
On Airflow 3.4+, the runtime task instance exposes `ti.region_id` and
`ti.region_index`. I checked on 3.4.0rc2: inside a loop they are the loop's
region and the iteration number (`region_index=1`, `2`, ... while `map_index`
stays `-1`), and outside any region `region_id` is the all-zero UUID with
`region_index == map_index`. Including `region_id` and `region_index` in
`compute_identity_key` only when `region_id` is not the all-zero UUID would
separate iterations (and mapped instances inside an iteration), while leaving
the key unchanged for tasks outside a loop and on older Airflow versions, so
in-flight batches are not orphaned by an upgrade. That also gives each
iteration its own `<key16>.jsonl` result file.
### Operating System
Debian (official `apache/airflow:3.4.0rc2` image)
### Deployment
Other Docker-based deployment
### Apache Airflow Provider(s)
common-ai
### Versions of Apache Airflow Providers
`apache-airflow-providers-common-ai==1.0.0` (same code on `main`)
### Official Helm Chart version
Not Applicable
### Kubernetes Version
_No response_
### Helm Chart configuration
_No response_
### Docker Image customizations
_No response_
### Anything else?
Searched open/closed issues and PRs for `llm_batch loop`,
`compute_identity_key`, and `LLMBatchOperator iteration`; no existing report.
cc @Lee-W (author of the batch mode, #72938) @kaxil
### Are you willing to submit PR?
- [x] Yes I am willing to submit a PR!
### Code of Conduct
- [x] I agree to follow this project's [Code of
Conduct](https://github.com/apache/airflow/blob/main/CODE_OF_CONDUCT.md)
--
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]
