AIP-97 POC: classify ECS task failures as infra

[1fanwang]

Why

Review feedback on AIP-97 said the failure classification looks Kubernetes-only, and asked for a
contract on BaseExecutor that other executors can implement against instead of the untyped info
slot on change_state.

ECS is the cheapest counter-example, because it is not missing the data. It throws it away.
describe_tasks returns Task.stopCode, a closed enum naming who stopped the task, but
BotoTaskSchema parses only stoppedReason, and __handle_failed_workload logs that string and
then calls self.fail(task_key) with nothing attached. A reclaimed spot task and a task that raised
reach the scheduler looking identical.

What changed

BaseExecutor.fail() takes optional typed failure_kind and reason. That is the contract: the
failure hook every executor already calls, typed, additive, and it leaves info alone, since info
is polymorphic today and carries external_executor_id on queued and running.

ECS then parses stopCode and classifies through it. SpotInterruption is infra. A host that
stopped or terminated under a running task is infra, matched on the same stoppedReason text
EcsRunTaskOperator already treats as a distinct failure. EssentialContainerExited is application
and earns no refund.

The other five stay unclassified and behave as today. TerminationNotice and
InfrastructureHealth do read as infra, but AWS documents neither, so mapping them would be a guess.

Also fixes two leaked mock.patch(...).start() calls in the ECS tests. They never stopped, so
AwsEcsExecutor.fail stayed mocked for every test that ran after them.

Refs

• Parent: AIP-97: Disruption Readiness (Transparent Infra Failure Retries through Context Capture and Propagation) apache/airflow#66405
• Sibling executors on the same contract: AIP-97 POC: classify a lost Celery worker as infra #15 (Celery), AIP-97 POC: classify a missed heartbeat as infra (executor-agnostic) #14 (heartbeat, executor-agnostic)
• TaskStopCode enum: https://docs.aws.amazon.com/AmazonECS/latest/APIReference/API_Task.html
• SpotInterruption: https://docs.aws.amazon.com/AmazonECS/latest/developerguide/spot-interruption-errors.html

Testing Done

Committed coverage is 14 provider tests, run in CI. The chain below is evidence.

#	Scenario	Result
1	Spot interruption, refund on, retries=0	up_for_retry, max_tries 0->1
2	Spot interruption, refund off	failed, max_tries 0->0
3	Essential container exited, refund on	failed, max_tries 0->0
4	Schema field reverted (red control)	classification lost

1-3: ECS executor to refund, live Postgres 16

Real sync_running_workloads() with only the boto client replaced, real handle_failure, real
listener, real database. stoppedReason is verbatim from AWS's Fargate Spot example event.

$ docker run -d --name aip97-ecs-pg -e POSTGRES_PASSWORD=airflow -e POSTGRES_USER=airflow \
    -e POSTGRES_DB=airflow -p 55433:5432 postgres:16
$ export AIRFLOW_HOME=/tmp/aip97-ecs-e2e
$ export AIRFLOW__DATABASE__SQL_ALCHEMY_CONN="postgresql+psycopg2://airflow:airflow@localhost:55433/airflow"
$ export AIRFLOW__CORE__UNIT_TEST_MODE=False
$ uv run --project airflow-core airflow db migrate
Running stamp_revision  -> 02ca36b0235b
Database migration done!

$ uv run --project airflow-core pytest /tmp/aip97_ecs_e2e_evidence.py -q -s
[error] Airflow workload TaskInstanceKey(dag_id='aip97_ecs_True_SpotInterruption', task_id='op', ...)
        has failed a maximum of 1 times. Marking as failed. stop_code=SpotInterruption failure_kind=infra
[executor] stop_code='SpotInterruption' -> (<TaskFailureKind.INFRA: 'infra'>, 'SpotInterruption')
[db]       refund=True state=up_for_retry max_tries 0->1 (expected 1)
[listener] received=[(<TaskFailureKind.INFRA: 'infra'>, 'SpotInterruption')]

[executor] stop_code='SpotInterruption' -> (<TaskFailureKind.INFRA: 'infra'>, 'SpotInterruption')
[db]       refund=False state=failed max_tries 0->0 (expected 0)
[listener] received=[(<TaskFailureKind.INFRA: 'infra'>, 'SpotInterruption')]

[executor] stop_code='EssentialContainerExited' -> (<TaskFailureKind.APPLICATION: 'application'>, 'EssentialContainerExited')
[db]       refund=True state=failed max_tries 0->0 (expected 0)
[listener] received=[(<TaskFailureKind.APPLICATION: 'application'>, None)]

3 passed

db migrate stamping to the existing head is the runnable proof this adds no schema. Scenario 1
uses retries=0, so the refund is what rescues it at all.

4: red control, and the leaked-patch fix

Deleting only the one-line stop_code schema field, on otherwise identical source:

$ sed -i '' '/stop_code = fields.String(data_key="stopCode")/d' \
    providers/amazon/src/airflow/providers/amazon/aws/executors/ecs/boto_schema.py
$ uv run --project providers/amazon pytest \
    providers/amazon/tests/unit/amazon/aws/executors/ecs/test_ecs_executor.py \
    -k "TestEcsFailureClassification and (spot_interruption or boto_schema)" -q

[error] ... Marking as failed. stop_code=None failure_kind=None
E   AssertionError: assert task.stop_code == "SpotInterruption"
E   AssertionError: assert mock_executor.get_task_failure_info(task_key) == (...)
2 failed

stop_code=None failure_kind=None on a real spot interruption is the defect this fixes, in one
line. Restored, both pass.

The two new tests passed alone but failed in the full file, which surfaced the leaked patches:

$ uv run --project providers/amazon pytest providers/amazon/tests/unit/amazon/aws/executors/ecs/ -q
145 passed, 3 skipped

Evidence script (inline, so it can be re-run)

Not committed: it imports the amazon provider from a core test path, which core CI does not install.

from __future__ import annotations

import datetime as dt
import os
import sys
from unittest import mock

import pytest

from airflow._shared.state import TaskFailureKind
from airflow.listeners import hookimpl
from airflow.listeners.listener import get_listener_manager
from airflow.providers.amazon.aws.executors.ecs.ecs_executor import AwsEcsExecutor
from airflow.providers.amazon.aws.executors.ecs.utils import CONFIG_GROUP_NAME, AllEcsConfigKeys
from airflow.providers.standard.operators.empty import EmptyOperator
from airflow.utils.state import TaskInstanceState

from tests_common.test_utils.config import conf_vars

ARN = "arn:aws:ecs:us-east-1:111122223333:task/b99d40b3-5176-4f71-9a52-9dbd6f1cebef"

ECS_CONF: dict[tuple[str, str], str] = {
    (CONFIG_GROUP_NAME, AllEcsConfigKeys.REGION_NAME): "us-west-1",
    (CONFIG_GROUP_NAME, AllEcsConfigKeys.CLUSTER): "some-cluster",
    (CONFIG_GROUP_NAME, AllEcsConfigKeys.CONTAINER_NAME): "container-name",
    (CONFIG_GROUP_NAME, AllEcsConfigKeys.TASK_DEFINITION): "some-task-def",
    (CONFIG_GROUP_NAME, AllEcsConfigKeys.LAUNCH_TYPE): "FARGATE",
    (CONFIG_GROUP_NAME, AllEcsConfigKeys.PLATFORM_VERSION): "LATEST",
    (CONFIG_GROUP_NAME, AllEcsConfigKeys.ASSIGN_PUBLIC_IP): "False",
    (CONFIG_GROUP_NAME, AllEcsConfigKeys.SECURITY_GROUPS): "sg1,sg2",
    (CONFIG_GROUP_NAME, AllEcsConfigKeys.SUBNETS): "sub1,sub2",
    (CONFIG_GROUP_NAME, AllEcsConfigKeys.MAX_RUN_TASK_ATTEMPTS): "1",
}

# Verbatim from the AWS docs' Fargate Spot termination-notice example event:
# https://docs.aws.amazon.com/AmazonECS/latest/developerguide/fargate-capacity-providers.html
SPOT_STOPPED_REASON = "Your Spot Task was interrupted."
SPOT_STOP_CODE = "SpotInterruption"

received: list[tuple] = []


class _Listener:
    @hookimpl
    def on_task_instance_failed(self, previous_state, task_instance, error, failure_kind, reason):
        received.append((failure_kind, reason))


def _stopped_task_json(stop_code, stopped_reason):
    return {
        "taskArn": ARN,
        "desiredStatus": "STOPPED",
        "lastStatus": "STOPPED",
        "startedAt": dt.datetime.now(),
        "stoppedReason": stopped_reason,
        "stopCode": stop_code,
        "containers": [
            {
                "containerArn": "c1",
                "name": "test_container",
                "lastStatus": "STOPPED",
                "exitCode": 137,
                "reason": stopped_reason,
            }
        ],
    }


def _build_executor(task_key, stop_code, stopped_reason):
    """A real AwsEcsExecutor with only the boto client replaced."""
    from airflow.providers.amazon.aws.executors.ecs.utils import EcsExecutorTask

    executor = AwsEcsExecutor()
    executor.ecs = mock.MagicMock()
    executor.max_run_task_attempts = "1"
    running_task = EcsExecutorTask(
        task_arn=ARN, last_status="RUNNING", desired_status="RUNNING", containers=[{"exit_code": 0}]
    )
    executor.active_workers.add_task(running_task, task_key, "default", ["airflow"], {}, 1)
    executor.ecs.describe_tasks.return_value = {
        "tasks": [_stopped_task_json(stop_code, stopped_reason)],
        "failures": [],
    }
    return executor


@pytest.mark.db_test
@pytest.mark.parametrize(
    ("refund_enabled", "stop_code", "stopped_reason", "expect_kind", "expect_state", "expect_max_tries"),
    [
        (
            "True",
            SPOT_STOP_CODE,
            SPOT_STOPPED_REASON,
            TaskFailureKind.INFRA,
            TaskInstanceState.UP_FOR_RETRY,
            1,
        ),
        ("False", SPOT_STOP_CODE, SPOT_STOPPED_REASON, TaskFailureKind.INFRA, TaskInstanceState.FAILED, 0),
        (
            "True",
            "EssentialContainerExited",
            "Essential container in task exited",
            TaskFailureKind.APPLICATION,
            TaskInstanceState.FAILED,
            0,
        ),
    ],
    ids=["spot_infra_refunded", "spot_infra_flag_off", "application_not_refunded"],
)
def test_ecs_failure_end_to_end(
    session,
    dag_maker,
    refund_enabled,
    stop_code,
    stopped_reason,
    expect_kind,
    expect_state,
    expect_max_tries,
):
    """ECS spot interruption -> executor classification -> scheduler read -> refund, on real Postgres."""
    received.clear()
    lm = get_listener_manager()
    lm.clear()
    lm.add_listener(_Listener())

    with dag_maker(dag_id=f"aip97_ecs_{refund_enabled}_{stop_code}", session=session):
        EmptyOperator(task_id="op", retries=0)
    dr = dag_maker.create_dagrun()
    ti = dr.task_instances[0]
    ti.task = dag_maker.dag.get_task("op")
    ti.state = TaskInstanceState.RUNNING
    ti.try_number = 1
    session.merge(ti)
    session.commit()

    pre_max_tries = ti.max_tries
    with conf_vars(ECS_CONF):
        executor = _build_executor(ti.key, stop_code, stopped_reason)
        executor.sync_running_workloads()
    classified = executor.get_task_failure_info(ti.key)
    print(f"\n[executor] stop_code={stop_code!r} -> {classified}")
    assert classified is not None, "executor did not classify the failure"
    failure_kind, reason = classified
    assert failure_kind == expect_kind

    with conf_vars({("core", "infra_failure_refund_retries"): refund_enabled}):
        ti.handle_failure(
            error="ECS task stopped",
            session=session,
            failure_kind=failure_kind,
            reason=reason if failure_kind == TaskFailureKind.INFRA else None,
        )
    session.commit()
    session.refresh(ti)

    print(
        f"[db]       refund={refund_enabled} state={ti.state} "
        f"max_tries {pre_max_tries}->{ti.max_tries} (expected {expect_max_tries})"
    )
    print(f"[listener] received={received}")

    assert ti.state == expect_state
    assert ti.max_tries == expect_max_tries
    assert received
    assert received[0][0] == expect_kind
    lm.clear()


if __name__ == "__main__":
    sys.exit(pytest.main([os.path.abspath(__file__), "-p", "no:cacheprovider", "-q", "-s"]))