AIP-97 Pillar 2: classify infra task failures and refund the retry budget

[1fanwang]

Superseded by #10 — the whole AIP-97 design now lives in #10 as a single clean diff off the apache root (no stack, no add-then-remove churn). This PR is one step of the original incremental stack; kept for history. Review #10.

---

Why

When a worker dies from outside — pod eviction, OOM kill (SIGKILL), spot preemption, heartbeat loss — Airflow can't tell it apart from a real task crash. Two costs:

1. on_task_instance_failed receives error=None (or a stringified dict), so alerting and OpenLineage lump infra churn in with genuine bugs.
2. The infra attempt spends a retry, so retries=2 plus two evictions leaves nothing for the code's first real failure.

Airflow already refuses to spend a retry on infra for one case — _is_pre_execution_failure() requeues a pod that died while QUEUED "without consuming a task retry." It stops at the QUEUED boundary because once a task is RUNNING it can't tell eviction from a crash. This adds the missing signal and extends that behavior past the boundary.

Stacked on the AIP-97 foundation (#5). This is the Pillar-2 delta: carry structured failure context to the listener, and refund the retry when the failure is infra.

What

• The scheduler's state-mismatch branch (executor reported the task finished, but the TI never marked itself failed — the "killed externally" case) populates a TaskFailureInfo(source="infra", ...) and hands it to the failure path.
• on_task_instance_failed now carries failure_details, so listeners see structured origin instead of a flattened string.
• _maybe_refund_infra_attempt() bumps max_tries only when source=="infra", the [core] infra_failure_refund_retries flag is on, and under the max_infra_refunds cap. try_number stays monotonic — no retries==0 special case.

source is a closed set — user | infra | timeout; preemption is infra + infra_reason="preemption", matching Flyte ErrorKind, K8s DisruptionTarget, Spark exitCausedByApp.

On classification — two layers, and what this POC proves

Classification is two layers. The base discriminator is the state-mismatch signal (ti_queued): the executor reports the task finished while the TI still looks QUEUED/RUNNING, i.e. the worker died before self-reporting. Same signal _is_pre_execution_failure uses, extended past the QUEUED boundary. On a real deployment an app failure catches in the worker and PATCHes state=FAILED before the executor event, so it never lands here; a killed worker does. This is executor-agnostic and needs only that workers self-report — which they do once an api-server is up.

The base layer can't split a silent death that is genuinely infra (node evicted the pod) from one that is app-caused (app OOM). That lives in the executor — the K8s bridge (collect_pod_failure_details() → OOMKilled vs Evicted vs an app exit code), done-criterion #2, which supplies the precise infra_reason and excludes app OOM. This PR ships the contract + the base discriminator + the refund gate; the K8s bridge refines the residual.

The safety property — an app/user/timeout/unclassified failure never refunds — is pinned by the refund gate with unit tests, independent of the classifier.

Testing Done

Live e2e on a real multi-component deployment — airflow standalone (api-server on :8080 serving the Execution API + scheduler + dag-processor + triggerer), PostgreSQL 14, LocalExecutor, refund flag on (max_infra_refunds=3), one recording listener. Because the api-server is up, workers self-report — so an app failure and an eviction take different paths. Two DAGs, retries=1: app_fail_repro raises ValueError; evict_repro SIGKILLs its own running worker (eviction signature).

The split is clean on one deployment — DB task_instance final state:

dag	state	try_number	max_tries	outcome
app_fail_repro	failed	2	1	ValueError self-reported → retry consumed, not refunded
evict_repro	failed	5	4	evictions refunded 3× (capped) → survived 5 attempts

Raw listener + scheduler logs

# app failure — worker ran, raised, self-reported (clean worker path)
on_task_instance_failed dag=app_fail_repro try=1 max_tries=1 error_type=ValueError error='deliberate application bug' failure_details=None
on_task_instance_failed dag=app_fail_repro try=2 max_tries=1 error_type=ValueError error='deliberate application bug' failure_details=None
-> retries=1 spent by the bug; FAILED after 2 attempts; zero refund lines

# infra eviction — worker SIGKILLed, never self-reported (scheduler state-mismatch branch)
on_task_instance_failed dag=evict_repro try=1 max_tries=2 failure_details=TaskFailureInfo(source='infra', executor_kind='LocalExecutor', infra_reason='ExecutorReportedFailure', ...)
on_task_instance_failed dag=evict_repro try=2 max_tries=3 failure_details=TaskFailureInfo(source='infra', ...)
on_task_instance_failed dag=evict_repro try=3 max_tries=4 failure_details=TaskFailureInfo(source='infra', ...)
on_task_instance_failed dag=evict_repro try=4 max_tries=4 failure_details=TaskFailureInfo(source='infra', ...)   # capped
scheduler: AIP-97: infra failure (ExecutorReportedFailure) refunded attempt ...; max_tries now 2 (refund 1/3), user retry budget preserved
scheduler: ... max_tries now 3 (refund 2/3) ...
scheduler: ... max_tries now 4 (refund 3/3) ...

(An earlier sqlite-only harness with no api-server mis-read the app failure as infra — the worker couldn't self-report, so it fell through the same branch as an eviction. Standing up the api-server so the worker self-reports is the fix; the split is clean above.)

Unit — the safety gate. _maybe_refund_infra_attempt(): 6 deterministic cases — infra refunds once; None (app path), user, timeout don't; flag-off doesn't; the cap bounds it to 3. Plus 9 listener-type tests (construction, frozen, pluggy no-default dispatch). 15/15 pass.

Out of scope / follow-ups

• Per-executor classification (done-criterion Bump the npm_and_yarn group across 5 directories with 11 updates #2): the K8s bridge for the precise pod-reason split and app-OOM exclusion.
• Counter sharing: refunding max_tries shares the counter with app retries. A precise "infra attempts never count" semantic needs a separate infra-attempt counter — noted as an AIP open question.

Back-compat + opt-in (verified live)

• Legacy listeners keep working. A stock on_task_instance_failed(previous_state, task_instance, error) hookimpl (no failure_details) fires unchanged alongside one that declares the new param — pluggy dispatches by name. No migration for existing listeners.
• The refund is opt-in. Same DAG, same deployment, only [core] infra_failure_refund_retries differs: OFF (default) → infra kill spends the retry (max_tries=1, stock); ON → refunded to the cap (max_tries=4). Default behavior is exactly today's.
• Pillar 1 is additive even with Pillar 2 off — the listener still gets TaskFailureInfo when the refund is disabled, so a deployment can adopt the failure-source context without changing retry accounting.

Risk

Pillar 1 (the listener arg) is additive — failure_details=None at every existing call site; listeners that ignore it are unaffected. Pillar 2 (the refund) is opt-in, config-gated, capped, and unit-pinned to source=="infra" only. Classification is heuristic; unknown stays None and behaves as today.

[coderabbitai]

Important

Review skipped

Auto reviews are disabled on base/target branches other than the default branch.

Please check the settings in the CodeRabbit UI or the .coderabbit.yaml file in this repository. To trigger a single review, invoke the @coderabbitai review command.

⚙️ Run configuration

Configuration used: defaults

Review profile: CHILL

Plan: Pro

Run ID: 9f7bad37-3292-4565-baa1-238aca75023f

You can disable this status message by setting the reviews.review_status to false in the CodeRabbit configuration file.

Use the checkbox below for a quick retry:

• 🔍 Trigger review

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create PR with unit tests
• Commit unit tests in branch 1fanwang/aip97-green-e2e

---

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

_{Comment @coderabbitai help to get the list of available commands.}

[1fanwang]

Superseded by the split pillar #13 (infra refund) and the consolidated apache#66405. Closing; branch kept.