Add skip_on_exit_code support to EcsRunTaskOperator (#63274)

Allow users to specify exit codes that should raise an
AirflowSkipException (marking the task as skipped) via the new
`skip_on_exit_code` parameter. This is consistent with the existing
behavior in DockerOperator and KubernetesPodOperator.

Author: yunhobb

providers/amazon/src/airflow/providers/amazon/aws/operators/ecs.py

@@ -18,7 +18,7 @@
 from __future__ import annotations
 
 import re
-from collections.abc import Sequence
+from collections.abc import Container, Sequence
 from datetime import timedelta
 from functools import cached_property
 from time import sleep
@@ -39,7 +39,7 @@
 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, conf
+from airflow.providers.common.compat.sdk import AirflowException, AirflowSkipException, conf
 from airflow.utils.helpers import prune_dict
 
 if TYPE_CHECKING:
@@ -394,6 +394,9 @@ class EcsRunTaskOperator(EcsBaseOperator):
     :param deferrable: If True, the operator will wait asynchronously for the job to complete.
         This implies waiting for completion. This mode requires aiobotocore module to be installed.
         (default: False)
+    :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 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
@@ -461,6 +464,7 @@ def __init__(
         # Set the default waiter duration to 70 days (attempts*delay)
         # Airflow execution_timeout handles task timeout
         deferrable: bool = conf.getboolean("operators", "default_deferrable", fallback=False),
+        skip_on_exit_code: int | Container[int] | None = None,
         stop_task_on_failure: bool = True,
         **kwargs,
     ):
@@ -500,6 +504,13 @@ def __init__(
         self.waiter_delay = waiter_delay
         self.waiter_max_attempts = waiter_max_attempts
         self.deferrable = deferrable
+        self.skip_on_exit_code = (
+            skip_on_exit_code
+            if isinstance(skip_on_exit_code, Container)
+            else [skip_on_exit_code]
+            if skip_on_exit_code is not None
+            else []
+        )
         self.stop_task_on_failure = stop_task_on_failure
 
         if self._aws_logs_enabled() and not self.wait_for_completion:
@@ -763,15 +774,21 @@ def _check_success_task(self) -> None:
             containers = task["containers"]
             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
+                    else:
+                        exception_cls = AirflowException
+
                     if self.task_log_fetcher:
                         last_logs = "\n".join(
                             self.task_log_fetcher.get_last_log_messages(self.number_logs_exception)
                         )
-                        raise AirflowException(
+                        raise exception_cls(
                             f"This task is not in success state - last {self.number_logs_exception} "
                             f"logs from Cloudwatch:\n{last_logs}"
                         )
-                    raise AirflowException(f"This task is not in success state {task}")
+                    raise exception_cls(f"This task is not in success state {task}")
                 if container.get("lastStatus") == "PENDING":
                     raise AirflowException(f"This task is still pending {task}")
                 if "error" in container.get("reason", "").lower():

providers/amazon/tests/unit/amazon/aws/operators/test_ecs.py

@@ -38,7 +38,7 @@
 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, TaskDeferred
+from airflow.providers.common.compat.sdk import AirflowException, AirflowSkipException, TaskDeferred
 
 from unit.amazon.aws.utils.test_template_fields import validate_template_fields
 
@@ -603,6 +603,41 @@ def test_check_success_task_not_raises(self, client_mock):
         self.ecs._check_success_task()
         client_mock.describe_tasks.assert_called_once_with(cluster="c", tasks=["arn"])
 
+    @mock.patch.object(EcsBaseOperator, "client")
+    def test_check_success_task_raises_skip_exception(self, client_mock):
+        self.ecs.arn = "arn"
+        self.ecs.skip_on_exit_code = [2]
+        client_mock.describe_tasks.return_value = {
+            "tasks": [{"containers": [{"name": "container-name", "lastStatus": "STOPPED", "exitCode": 2}]}]
+        }
+        with pytest.raises(AirflowSkipException):
+            self.ecs._check_success_task()
+
+    @mock.patch.object(EcsBaseOperator, "client")
+    @mock.patch("airflow.providers.amazon.aws.utils.task_log_fetcher.AwsTaskLogFetcher")
+    def test_check_success_task_skip_exception_with_logs(self, log_fetcher_mock, client_mock):
+        self.ecs.arn = "arn"
+        self.ecs.skip_on_exit_code = [2]
+        self.ecs.task_log_fetcher = log_fetcher_mock
+        log_fetcher_mock.get_last_log_messages.return_value = ["log1", "log2"]
+        client_mock.describe_tasks.return_value = {
+            "tasks": [{"containers": [{"name": "container-name", "lastStatus": "STOPPED", "exitCode": 2}]}]
+        }
+        with pytest.raises(AirflowSkipException, match="This task is not in success state"):
+            self.ecs._check_success_task()
+
+    @mock.patch.object(EcsBaseOperator, "client")
+    def test_check_success_task_unmatched_exit_code_raises_airflow_exception(self, client_mock):
+        """Exit codes not in skip_on_exit_code raise AirflowException."""
+        self.ecs.arn = "arn"
+        self.ecs.skip_on_exit_code = [2]
+        client_mock.describe_tasks.return_value = {
+            "tasks": [{"containers": [{"name": "container-name", "lastStatus": "STOPPED", "exitCode": 1}]}]
+        }
+        with pytest.raises(AirflowException) as ctx:
+            self.ecs._check_success_task()
+        assert type(ctx.value) is AirflowException
+
     @pytest.mark.parametrize(
         ("launch_type", "tags"),
         [