You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@airflow.apache.org by on...@apache.org on 2024/03/18 18:29:35 UTC
(airflow) branch main updated: `ECSExecutor` API Retry bug fix (#38118)
This is an automated email from the ASF dual-hosted git repository.
onikolas 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 111245af3f `ECSExecutor` API Retry bug fix (#38118)
111245af3f is described below
commit 111245af3fe0bb1292ea83df2dfe0e268e6b848a
Author: Syed Hussain <10...@users.noreply.github.com>
AuthorDate: Mon Mar 18 11:29:29 2024 -0700
`ECSExecutor` API Retry bug fix (#38118)
* Fix bug in ECS Executor regarding API retries. Add unit test
---
.../amazon/aws/executors/ecs/ecs_executor.py | 22 ++-
.../amazon/aws/executors/ecs/test_ecs_executor.py | 196 ++++++++++++++++++---
2 files changed, 187 insertions(+), 31 deletions(-)
diff --git a/airflow/providers/amazon/aws/executors/ecs/ecs_executor.py b/airflow/providers/amazon/aws/executors/ecs/ecs_executor.py
index 6167ddcbf4..87b9abe435 100644
--- a/airflow/providers/amazon/aws/executors/ecs/ecs_executor.py
+++ b/airflow/providers/amazon/aws/executors/ecs/ecs_executor.py
@@ -243,20 +243,17 @@ class AwsEcsExecutor(BaseExecutor):
task_key = self.active_workers.arn_to_key[task.task_arn]
# Mark finished tasks as either a success/failure.
- if task_state == State.FAILED:
- self.fail(task_key)
+ if task_state == State.FAILED or task_state == State.REMOVED:
self.__log_container_failures(task_arn=task.task_arn)
- elif task_state == State.SUCCESS:
- self.success(task_key)
- elif task_state == State.REMOVED:
self.__handle_failed_task(task.task_arn, task.stopped_reason)
- if task_state in (State.FAILED, State.SUCCESS):
+ elif task_state == State.SUCCESS:
self.log.debug(
"Airflow task %s marked as %s after running on ECS Task (arn) %s",
task_key,
task_state,
task.task_arn,
)
+ self.success(task_key)
self.active_workers.pop_by_key(task_key)
def __describe_tasks(self, task_arns):
@@ -289,7 +286,14 @@ class AwsEcsExecutor(BaseExecutor):
)
def __handle_failed_task(self, task_arn: str, reason: str):
- """If an API failure occurs, the task is rescheduled."""
+ """
+ If an API failure occurs, the task is rescheduled.
+
+ This function will determine whether the task has been attempted the appropriate number
+ of times, and determine whether the task should be marked failed or not. The task will
+ be removed active_workers, and marked as FAILED, or set into pending_tasks depending on
+ how many times it has been retried.
+ """
task_key = self.active_workers.arn_to_key[task_arn]
task_info = self.active_workers.info_by_key(task_key)
task_cmd = task_info.cmd
@@ -305,7 +309,6 @@ class AwsEcsExecutor(BaseExecutor):
self.__class__.MAX_RUN_TASK_ATTEMPTS,
task_arn,
)
- self.active_workers.increment_failure_count(task_key)
self.pending_tasks.append(
EcsQueuedTask(
task_key,
@@ -322,8 +325,8 @@ class AwsEcsExecutor(BaseExecutor):
task_key,
failure_count,
)
- self.active_workers.pop_by_key(task_key)
self.fail(task_key)
+ self.active_workers.pop_by_key(task_key)
def attempt_task_runs(self):
"""
@@ -346,6 +349,7 @@ class AwsEcsExecutor(BaseExecutor):
attempt_number = ecs_task.attempt_number
_failure_reasons = []
if timezone.utcnow() < ecs_task.next_attempt_time:
+ self.pending_tasks.append(ecs_task)
continue
try:
run_task_response = self._run_task(task_key, cmd, queue, exec_config)
diff --git a/tests/providers/amazon/aws/executors/ecs/test_ecs_executor.py b/tests/providers/amazon/aws/executors/ecs/test_ecs_executor.py
index c4c7c6c9c7..8762480821 100644
--- a/tests/providers/amazon/aws/executors/ecs/test_ecs_executor.py
+++ b/tests/providers/amazon/aws/executors/ecs/test_ecs_executor.py
@@ -598,6 +598,109 @@ class TestAwsEcsExecutor:
== caplog.messages[0]
)
+ @mock.patch.object(ecs_executor, "calculate_next_attempt_delay", return_value=dt.timedelta(seconds=0))
+ def test_task_retry_on_api_failure_all_tasks_fail(self, _, mock_executor, caplog):
+ """
+ Test API failure retries.
+ """
+ AwsEcsExecutor.MAX_RUN_TASK_ATTEMPTS = "2"
+ airflow_keys = ["TaskInstanceKey1", "TaskInstanceKey2"]
+ airflow_commands = [mock.Mock(spec=list), mock.Mock(spec=list)]
+
+ mock_executor.execute_async(airflow_keys[0], airflow_commands[0])
+ mock_executor.execute_async(airflow_keys[1], airflow_commands[1])
+ assert len(mock_executor.pending_tasks) == 2
+ caplog.set_level("WARNING")
+
+ describe_tasks = [
+ {
+ "taskArn": ARN1,
+ "desiredStatus": "STOPPED",
+ "lastStatus": "FAILED",
+ "startedAt": dt.datetime.now(),
+ "stoppedReason": "Task marked as FAILED",
+ "containers": [
+ {
+ "name": "some-ecs-container",
+ "lastStatus": "STOPPED",
+ "exitCode": 100,
+ }
+ ],
+ },
+ {
+ "taskArn": ARN2,
+ "desiredStatus": "STOPPED",
+ "lastStatus": "FAILED",
+ "stoppedReason": "Task marked as REMOVED",
+ "containers": [
+ {
+ "name": "some-ecs-container",
+ "lastStatus": "STOPPED",
+ "exitCode": 100,
+ }
+ ],
+ },
+ ]
+ run_tasks = [
+ {
+ "taskArn": ARN1,
+ "lastStatus": "",
+ "desiredStatus": "",
+ "containers": [{"name": "some-ecs-container"}],
+ },
+ {
+ "taskArn": ARN2,
+ "lastStatus": "",
+ "desiredStatus": "",
+ "containers": [{"name": "some-ecs-container"}],
+ },
+ ]
+ mock_executor.ecs.run_task.side_effect = [
+ {"tasks": [run_tasks[0]], "failures": []},
+ {"tasks": [run_tasks[1]], "failures": []},
+ ]
+ mock_executor.ecs.describe_tasks.side_effect = [{"tasks": describe_tasks, "failures": []}]
+
+ mock_executor.attempt_task_runs()
+
+ for i in range(2):
+ RUN_TASK_KWARGS["overrides"]["containerOverrides"][0]["command"] = airflow_commands[i]
+ assert mock_executor.ecs.run_task.call_args_list[i].kwargs == RUN_TASK_KWARGS
+
+ assert len(mock_executor.pending_tasks) == 0
+ assert len(mock_executor.active_workers.get_all_arns()) == 2
+
+ mock_executor.sync_running_tasks()
+ for i in range(2):
+ assert (
+ f"Airflow task {airflow_keys[i]} failed due to {describe_tasks[i]['stoppedReason']}. Failure 1 out of 2"
+ in caplog.messages[i]
+ )
+
+ caplog.clear()
+ mock_executor.ecs.run_task.call_args_list.clear()
+
+ mock_executor.ecs.run_task.side_effect = [
+ {"tasks": [run_tasks[0]], "failures": []},
+ {"tasks": [run_tasks[1]], "failures": []},
+ ]
+ mock_executor.ecs.describe_tasks.side_effect = [{"tasks": describe_tasks, "failures": []}]
+
+ mock_executor.attempt_task_runs()
+
+ mock_executor.attempt_task_runs()
+
+ for i in range(2):
+ RUN_TASK_KWARGS["overrides"]["containerOverrides"][0]["command"] = airflow_commands[i]
+ assert mock_executor.ecs.run_task.call_args_list[i].kwargs == RUN_TASK_KWARGS
+
+ mock_executor.sync_running_tasks()
+ for i in range(2):
+ assert (
+ f"Airflow task {airflow_keys[i]} has failed a maximum of 2 times. Marking as failed"
+ in caplog.messages[i]
+ )
+
@mock.patch.object(BaseExecutor, "fail")
@mock.patch.object(BaseExecutor, "success")
def test_sync(self, success_mock, fail_mock, mock_executor):
@@ -629,6 +732,7 @@ class TestAwsEcsExecutor:
@mock.patch.object(BaseExecutor, "success")
def test_failed_sync(self, success_mock, fail_mock, mock_executor):
"""Test success and failure states."""
+ AwsEcsExecutor.MAX_RUN_TASK_ATTEMPTS = "1"
self._mock_sync(mock_executor, State.FAILED)
mock_executor.sync()
@@ -636,22 +740,23 @@ class TestAwsEcsExecutor:
# Task is not stored in active workers.
assert len(mock_executor.active_workers) == 0
- # Task is immediately succeeded.
+ # Task is immediately failed.
fail_mock.assert_called_once()
success_mock.assert_not_called()
- @mock.patch.object(BaseExecutor, "fail")
@mock.patch.object(BaseExecutor, "success")
+ @mock.patch.object(BaseExecutor, "fail")
def test_removed_sync(self, fail_mock, success_mock, mock_executor):
- """A removed task will increment failure count but call neither fail() nor success()."""
+ """A removed task will be treated as a failed task."""
+ AwsEcsExecutor.MAX_RUN_TASK_ATTEMPTS = "1"
self._mock_sync(mock_executor, expected_state=State.REMOVED, set_task_state=State.REMOVED)
- task_instance_key = mock_executor.active_workers.arn_to_key[ARN1]
mock_executor.sync_running_tasks()
- assert ARN1 in mock_executor.active_workers.get_all_arns()
- assert mock_executor.active_workers.key_to_failure_counts[task_instance_key] == 2
- fail_mock.assert_not_called()
+ # Task is not stored in active workers.
+ assert len(mock_executor.active_workers) == 0
+ # Task is immediately failed.
+ fail_mock.assert_called_once()
success_mock.assert_not_called()
@mock.patch.object(BaseExecutor, "fail")
@@ -698,19 +803,27 @@ class TestAwsEcsExecutor:
],
}
- # Call sync_running_tasks 2 times with failures.
+ # Call sync_running_tasks and attempt_task_runs 2 times with failures.
for _ in range(2):
mock_executor.sync_running_tasks()
- # Ensure task arn is not removed from active.
+ # Ensure task gets removed from active_workers.
+ assert ARN1 not in mock_executor.active_workers.get_all_arns()
+ # Ensure task gets back on the pending_tasks queue
+ assert len(mock_executor.pending_tasks) == 1
+ keys = [task.key for task in mock_executor.pending_tasks]
+ assert task_key in keys
+
+ mock_executor.attempt_task_runs()
+ assert len(mock_executor.pending_tasks) == 0
assert ARN1 in mock_executor.active_workers.get_all_arns()
- # Task is neither failed nor succeeded.
- fail_mock.assert_not_called()
- success_mock.assert_not_called()
+ # Task is neither failed nor succeeded.
+ fail_mock.assert_not_called()
+ success_mock.assert_not_called()
- # run_task failed twice, and passed once
- assert mock_executor.ecs.run_task.call_count == 3
+ # run_task failed twice, and passed 3 times
+ assert mock_executor.ecs.run_task.call_count == 5
# describe_tasks failed 2 times so far
assert mock_executor.ecs.describe_tasks.call_count == 2
@@ -731,27 +844,60 @@ class TestAwsEcsExecutor:
@mock.patch.object(BaseExecutor, "fail")
@mock.patch.object(BaseExecutor, "success")
- def test_failed_sync_api(self, success_mock, fail_mock, mock_executor):
+ @mock.patch.object(ecs_executor, "calculate_next_attempt_delay", return_value=dt.timedelta(seconds=0))
+ def test_failed_sync_api(self, _, success_mock, fail_mock, mock_executor):
"""Test what happens when ECS sync fails for certain tasks repeatedly."""
- self._mock_sync(mock_executor)
- mock_executor.ecs.describe_tasks.return_value = {
+ airflow_key = "test-key"
+ airflow_cmd = mock.Mock(spec=list)
+ mock_executor.execute_async(airflow_key, airflow_cmd)
+ assert len(mock_executor.pending_tasks) == 1
+
+ run_task_ret_val = {
+ "taskArn": ARN1,
+ "desiredStatus": "STOPPED",
+ "lastStatus": "RUNNING",
+ "containers": [
+ {
+ "name": "some-ecs-container",
+ "lastStatus": "STOPPED",
+ "exitCode": 0,
+ }
+ ],
+ }
+ mock_executor.ecs.run_task.return_value = {"tasks": [run_task_ret_val], "failures": []}
+ describe_tasks_ret_value = {
"tasks": [],
"failures": [
{"arn": ARN1, "reason": "Sample Failure", "detail": "UnitTest Failure - Please ignore"}
],
}
+ mock_executor.ecs.describe_tasks.return_value = describe_tasks_ret_value
+ mock_executor.attempt_task_runs()
+ assert len(mock_executor.pending_tasks) == 0
+ assert len(mock_executor.active_workers.get_all_arns()) == 1
+ task_key = mock_executor.active_workers.arn_to_key[ARN1]
# Call Sync 2 times with failures. The task can only fail MAX_RUN_TASK_ATTEMPTS times.
for check_count in range(1, int(AwsEcsExecutor.MAX_RUN_TASK_ATTEMPTS)):
mock_executor.sync_running_tasks()
assert mock_executor.ecs.describe_tasks.call_count == check_count
- # Ensure task arn is not removed from active.
- assert ARN1 in mock_executor.active_workers.get_all_arns()
+ # Ensure task gets removed from active_workers.
+ assert ARN1 not in mock_executor.active_workers.get_all_arns()
+ # Ensure task gets back on the pending_tasks queue
+ assert len(mock_executor.pending_tasks) == 1
+ keys = [task.key for task in mock_executor.pending_tasks]
+ assert task_key in keys
# Task is neither failed nor succeeded.
fail_mock.assert_not_called()
success_mock.assert_not_called()
+ mock_executor.attempt_task_runs()
+
+ assert len(mock_executor.pending_tasks) == 0
+ assert len(mock_executor.active_workers.get_all_arns()) == 1
+ assert ARN1 in mock_executor.active_workers.get_all_arns()
+ task_key = mock_executor.active_workers.arn_to_key[ARN1]
# Last call should fail the task.
mock_executor.sync_running_tasks()
@@ -878,6 +1024,7 @@ class TestAwsEcsExecutor:
set_task_state=TaskInstanceState.RUNNING,
) -> None:
"""Mock ECS to the expected state."""
+ executor.pending_tasks.clear()
self._add_mock_task(executor, ARN1, set_task_state)
response_task_json = {
@@ -938,9 +1085,13 @@ class TestAwsEcsExecutor:
}
mock_executor.ecs.describe_tasks.return_value = {"tasks": [test_response_task_json], "failures": []}
mock_executor.sync_running_tasks()
- assert mock_executor.active_workers.tasks["arn1"].get_task_state() == expected_status
- # The task is not removed from active_workers in these states
- assert len(mock_executor.active_workers) == 1
+ if expected_status != State.REMOVED:
+ assert mock_executor.active_workers.tasks["arn1"].get_task_state() == expected_status
+ # The task is not removed from active_workers in these states
+ assert len(mock_executor.active_workers) == 1
+ else:
+ # The task is removed from active_workers in this state
+ assert len(mock_executor.active_workers) == 0
def test_update_running_tasks_success(self, mock_executor):
self._add_mock_task(mock_executor, ARN1)
@@ -967,6 +1118,7 @@ class TestAwsEcsExecutor:
mock_success_function.assert_called_once()
def test_update_running_tasks_failed(self, mock_executor, caplog):
+ AwsEcsExecutor.MAX_RUN_TASK_ATTEMPTS = "1"
caplog.set_level(logging.WARNING)
self._add_mock_task(mock_executor, ARN1)
test_response_task_json = {