This is an automated email from the ASF dual-hosted git repository.
o-nikolas pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/airflow.git
The following commit(s) were added to refs/heads/main by this push:
new d12021da679 Allow ECS exit codes to fail tasks without retrying
(#73885)
d12021da679 is described below
commit d12021da679d57a0eb1da65fca9ded9a0ed154a8
Author: Jingi Kim <[email protected]>
AuthorDate: Sat Oct 24 00:43:19 2026 +0900
Allow ECS exit codes to fail tasks without retrying (#73885)
Allow Dag authors to stop retries for known application failures reported
by an ECS container. For example, with retries=2 and fail_on_exit_code=2, an
invalid-input exit code fails the Airflow task after its first attempt, while
other exit codes retain the normal retry behavior.
---
providers/amazon/docs/operators/ecs.rst | 36 +++++++++++++++
.../airflow/providers/amazon/aws/operators/ecs.py | 38 +++++++++++++--
.../tests/unit/amazon/aws/operators/test_ecs.py | 54 +++++++++++++++++++++-
3 files changed, 124 insertions(+), 4 deletions(-)
diff --git a/providers/amazon/docs/operators/ecs.rst
b/providers/amazon/docs/operators/ecs.rst
index ee46902373b..8e136e10626 100644
--- a/providers/amazon/docs/operators/ecs.rst
+++ b/providers/amazon/docs/operators/ecs.rst
@@ -173,6 +173,42 @@ all replace the task definition's value rather than adding
to it. See
`ContainerOverride
<https://docs.aws.amazon.com/AmazonECS/latest/APIReference/API_ContainerOverride.html>`__
in the Amazon ECS API reference for the full list of fields.
+Control retries for container exit codes
+""""""""""""""""""""""""""""""""""""""""
+
+Use ``fail_on_exit_code`` when a known application failure should fail the
Airflow task
+without consuming its remaining retries. Other failures continue to use the
task's
+normal retry settings. For example, if your application uses exit code ``2``
for an
+invalid input that will not succeed on retry:
+
+.. code-block:: python
+
+ EcsRunTaskOperator(
+ task_id="run_application",
+ cluster="application_cluster",
+ task_definition="application_task",
+ overrides={},
+ retries=2,
+ fail_on_exit_code=2,
+ )
+
+The option accepts an integer or a container of integers. Its default is
``None``,
+which preserves the existing retry behavior. Exit code ``0`` remains
successful;
+a missing exit code does not match this option.
+
+If any stopped container explicitly reports a matching non-zero code, the task
fails
+without retrying, even if another container reports a different failure or a
code in
+``skip_on_exit_code``. This result does not depend on container order. When
the same
+code appears in both options, ``fail_on_exit_code`` takes precedence.
Task-level
+startup failures, host termination checks, and AWS API errors retain their
existing
+handling.
+
+This behavior applies to synchronous execution and resumed deferrable
execution.
+The option has no effect when ``wait_for_completion=False`` and
``deferrable=False``:
+the operator returns before checking exit codes. On Airflow versions supporting
+``retry_policy``, a matching code raises ``AirflowFailException``, which
bypasses
+retry policies, including a policy that would otherwise retry the task.
+
Stream logs to AWS CloudWatch
"""""""""""""""""""""""""""""
diff --git a/providers/amazon/src/airflow/providers/amazon/aws/operators/ecs.py
b/providers/amazon/src/airflow/providers/amazon/aws/operators/ecs.py
index aad4678df5a..a82931b1263 100644
--- a/providers/amazon/src/airflow/providers/amazon/aws/operators/ecs.py
+++ b/providers/amazon/src/airflow/providers/amazon/aws/operators/ecs.py
@@ -39,7 +39,12 @@ from airflow.providers.amazon.aws.utils import
validate_execute_complete_event
from airflow.providers.amazon.aws.utils.identifiers import generate_uuid
from airflow.providers.amazon.aws.utils.mixins import aws_template_fields
from airflow.providers.amazon.aws.utils.task_log_fetcher import
AwsTaskLogFetcher
-from airflow.providers.common.compat.sdk import AirflowException,
AirflowSkipException, conf
+from airflow.providers.common.compat.sdk import (
+ AirflowException,
+ AirflowFailException,
+ AirflowSkipException,
+ conf,
+)
from airflow.utils.helpers import prune_dict
if TYPE_CHECKING:
@@ -404,6 +409,10 @@ class EcsRunTaskOperator(EcsBaseOperator):
:param skip_on_exit_code: If task exits with this exit code, leave the task
in ``skipped`` state (default: None). If set to ``None``, any non-zero
exit code will be treated as a failure. Can be an int or a container
of ints.
+ :param fail_on_exit_code: Non-zero exit codes that fail the task without
retrying.
+ Can be an int or a container of ints (default: None). A matching code
reported
+ by a stopped container takes precedence over ``skip_on_exit_code`` and
other
+ container failures. Missing exit codes do not match this option.
:param do_xcom_push: If True, the operator will push the ECS task ARN to
XCom with key 'ecs_task_arn'.
Additionally, if logs are fetched, the last log message will be pushed
to XCom with the key 'return_value'. (default: False)
:param stop_task_on_failure: If True, attempt to stop the ECS task if the
Airflow task fails
@@ -472,6 +481,7 @@ class EcsRunTaskOperator(EcsBaseOperator):
# Airflow execution_timeout handles task timeout
deferrable: bool = conf.getboolean("operators", "default_deferrable",
fallback=False),
skip_on_exit_code: int | Container[int] | None = None,
+ fail_on_exit_code: int | Container[int] | None = None,
stop_task_on_failure: bool = True,
**kwargs,
):
@@ -516,6 +526,13 @@ class EcsRunTaskOperator(EcsBaseOperator):
else []
)
self.stop_task_on_failure = stop_task_on_failure
+ self.fail_on_exit_code = (
+ fail_on_exit_code
+ if isinstance(fail_on_exit_code, Container)
+ else [fail_on_exit_code]
+ if fail_on_exit_code is not None
+ else []
+ )
@staticmethod
def _get_ecs_task_id(task_arn: str | None) -> str | None:
@@ -787,11 +804,26 @@ class EcsRunTaskOperator(EcsBaseOperator):
f" {task.get('stoppedReason', '')}"
)
containers = task["containers"]
+ if self.fail_on_exit_code:
+ failed_container = next(
+ (
+ container
+ for container in containers
+ if container.get("lastStatus") == "STOPPED"
+ and container.get("exitCode") not in (None, 0)
+ and container["exitCode"] in self.fail_on_exit_code
+ ),
+ None,
+ )
+ if failed_container is not None:
+ containers = [failed_container]
for container in containers:
if container.get("lastStatus") == "STOPPED" and
container.get("exitCode", 1) != 0:
exit_code = container.get("exitCode", 1)
- if exit_code in self.skip_on_exit_code:
- exception_cls: type[AirflowException] =
AirflowSkipException
+ if container.get("exitCode") is not None and exit_code in
self.fail_on_exit_code:
+ exception_cls: type[AirflowException] =
AirflowFailException
+ elif exit_code in self.skip_on_exit_code:
+ exception_cls = AirflowSkipException
else:
exception_cls = AirflowException
diff --git a/providers/amazon/tests/unit/amazon/aws/operators/test_ecs.py
b/providers/amazon/tests/unit/amazon/aws/operators/test_ecs.py
index b53ca45d9c6..6887c392ea7 100644
--- a/providers/amazon/tests/unit/amazon/aws/operators/test_ecs.py
+++ b/providers/amazon/tests/unit/amazon/aws/operators/test_ecs.py
@@ -40,10 +40,16 @@ from airflow.providers.amazon.aws.operators.ecs import (
from airflow.providers.amazon.aws.triggers.ecs import TaskDoneTrigger
from airflow.providers.amazon.aws.utils.task_log_fetcher import
AwsTaskLogFetcher
from airflow.providers.amazon.version_compat import NOTSET
-from airflow.providers.common.compat.sdk import AirflowException,
AirflowSkipException, TaskDeferred
+from airflow.providers.common.compat.sdk import (
+ AirflowException,
+ AirflowFailException,
+ AirflowSkipException,
+ TaskDeferred,
+)
from unit.amazon.aws.utils.test_template_fields import validate_template_fields
+TASK_ARN =
"arn:aws:ecs:us-east-1:123456789012:task/probe/00000000000000000000000000000001"
CLUSTER_NAME = "test_cluster"
CONTAINER_NAME = "e1ed7aac-d9b2-4315-8726-d2432bf11868"
TASK_ID = "d8c67b3c-ac87-4ffe-a847-4785bc3a8b55"
@@ -665,6 +671,52 @@ class TestEcsRunTaskOperator(EcsBaseTestCase):
self.ecs._check_success_task()
assert type(ctx.value) is AirflowException
+ @pytest.mark.parametrize(
+ ("fail_codes", "skip_codes", "exit_codes", "expected"),
+ [
+ pytest.param(2, None, [2], AirflowFailException, id="int_match"),
+ pytest.param([2, 3], None, [3], AirflowFailException,
id="list_match"),
+ pytest.param(2, None, [1], AirflowException, id="unmatched_code"),
+ pytest.param(1, None, [None], AirflowException,
id="missing_exit_code"),
+ pytest.param(2, 2, [2], AirflowFailException, id="fail_over_skip"),
+ pytest.param(2, 3, [3, 2], AirflowFailException,
id="later_container_over_skip"),
+ ],
+ )
+ @mock.patch.object(EcsBaseOperator, "client")
+ def test_check_success_task_fail_on_exit_code(
+ self, client_mock, fail_codes, skip_codes, exit_codes, expected
+ ):
+ self.set_up_operator(fail_on_exit_code=fail_codes,
skip_on_exit_code=skip_codes)
+ self.ecs.arn = "arn"
+ containers = [
+ {"lastStatus": "STOPPED", **({"exitCode": code} if code is not
None else {})}
+ for code in exit_codes
+ ]
+ client_mock.describe_tasks.return_value = {"tasks": [{"containers":
containers}]}
+
+ with pytest.raises(expected) as raised:
+ self.ecs._check_success_task()
+
+ assert type(raised.value) is expected
+ client_mock.describe_tasks.assert_called_once_with(cluster="c",
tasks=["arn"])
+
+ @mock.patch.object(EcsBaseOperator, "client")
+ def test_check_success_task_fail_on_exit_code_preserves_log_tail(self,
client_mock):
+ self.set_up_operator(fail_on_exit_code=2)
+ self.ecs.arn = "arn"
+ self.ecs.task_log_fetcher = mock.Mock(spec=AwsTaskLogFetcher)
+ self.ecs.task_log_fetcher.get_last_log_messages.return_value =
["COPY_PROBE_123"]
+ client_mock.describe_tasks.return_value = {
+ "tasks": [{"containers": [{"lastStatus": "STOPPED", "exitCode":
2}]}]
+ }
+
+ with pytest.raises(AirflowFailException, match="COPY_PROBE_123"):
+ self.ecs._check_success_task()
+
+
self.ecs.task_log_fetcher.get_last_log_messages.assert_called_once_with(
+ self.ecs.number_logs_exception
+ )
+
@pytest.mark.parametrize(
("launch_type", "tags"),
[
