This is an automated email from the ASF dual-hosted git repository.
vincbeck 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 acb2208f15f Let the ECS executor start tasks in parallel (#75907)
acb2208f15f is described below
commit acb2208f15fd7b1acc9dff33aa48455611da8235
Author: Vincent <[email protected]>
AuthorDate: Fri Oct 30 08:41:37 2026 -0400
Let the ECS executor start tasks in parallel (#75907)
The ECS executor calls RunTask for one task at a time, inside the
scheduler loop, and waits for each answer. A call takes about 0.7 s, so a
scheduler starts at most about 1.4 tasks per second however many slots
are free, and its loop stalls while it works through a backlog:
scheduling, heartbeats and executor events all wait, for up to 13 minutes
in a test with 10,000 tasks.
A throttled RunTask call no longer uses one of the task's attempts, since
the throttling says nothing about the task; parallel calls make it more
likely.
---
providers/amazon/docs/executors/ecs-executor.rst | 4 +
providers/amazon/provider.yaml | 9 +
.../amazon/aws/executors/ecs/ecs_executor.py | 230 ++++++++++++------
.../providers/amazon/aws/executors/ecs/utils.py | 3 +
.../airflow/providers/amazon/get_provider_info.py | 7 +
.../amazon/aws/executors/ecs/test_ecs_executor.py | 265 ++++++++++++++++++++-
.../observability/metrics/metrics_template.yaml | 6 +
7 files changed, 453 insertions(+), 71 deletions(-)
diff --git a/providers/amazon/docs/executors/ecs-executor.rst
b/providers/amazon/docs/executors/ecs-executor.rst
index ce66b3884ba..67d3f2165b1 100644
--- a/providers/amazon/docs/executors/ecs-executor.rst
+++ b/providers/amazon/docs/executors/ecs-executor.rst
@@ -111,6 +111,10 @@ Optional config options:
fails to start (i.e. ECS API failures, container failures etc.)
- CHECK_HEALTH_ON_STARTUP - Whether or not to check the ECS Executor
health on startup
+- RUN_TASK_CONCURRENCY - The maximum number of ECS ``run_task`` calls
+ the ECS Executor makes in parallel, in background threads. Defaults
+ to 1: tasks are launched one at a time, and the scheduler loop waits
+ for each call.
For a more detailed description of available options, including type
hints and examples, see the ``config_templates`` folder in the Amazon
diff --git a/providers/amazon/provider.yaml b/providers/amazon/provider.yaml
index 72669ee719e..190a90162c2 100644
--- a/providers/amazon/provider.yaml
+++ b/providers/amazon/provider.yaml
@@ -1511,6 +1511,15 @@ config:
type: integer
example: "3"
default: "3"
+ run_task_concurrency:
+ description: |
+ The maximum number of ECS RunTask calls the ECS Executor makes in
parallel. With 1, tasks are
+ launched one at a time and the scheduler loop waits for each call.
With more, the calls run in
+ background threads and the scheduler loop does not wait for them.
+ version_added: "9.40.0"
+ type: integer
+ example: "16"
+ default: "1"
run_task_kwargs:
description: |
A JSON string containing arguments to provide the ECS `run_task` API
(see url above).
diff --git
a/providers/amazon/src/airflow/providers/amazon/aws/executors/ecs/ecs_executor.py
b/providers/amazon/src/airflow/providers/amazon/aws/executors/ecs/ecs_executor.py
index 77cd60b704b..35c6a3b4afc 100644
---
a/providers/amazon/src/airflow/providers/amazon/aws/executors/ecs/ecs_executor.py
+++
b/providers/amazon/src/airflow/providers/amazon/aws/executors/ecs/ecs_executor.py
@@ -25,11 +25,14 @@ from __future__ import annotations
import time
import warnings
-from collections import defaultdict, deque
-from collections.abc import Sequence
+from collections import deque
+from collections.abc import Callable, Sequence
+from concurrent.futures import Future, ThreadPoolExecutor
from copy import deepcopy
+from functools import partial
from typing import TYPE_CHECKING, TypeAlias
+from botocore.config import Config
from botocore.exceptions import ClientError, NoCredentialsError
from airflow.exceptions import AirflowProviderDeprecationWarning
@@ -83,6 +86,10 @@ INVALID_CREDENTIALS_EXCEPTIONS = [
"InvalidClientTokenId",
"UnrecognizedClientException",
]
+THROTTLING_EXCEPTIONS = ["ThrottlingException"]
+# Retries on top of botocore's own. Past this number, a throttled RunTask call
uses one of the workload's
+# ``max_run_task_attempts``, so a workload cannot stay pending for as long as
the throttling lasts.
+MAX_THROTTLED_RUN_TASK_RETRIES = 5
class AwsEcsExecutor(BaseExecutor):
@@ -143,6 +150,28 @@ class AwsEcsExecutor(BaseExecutor):
self.container_name = self.conf.get(CONFIG_GROUP_NAME,
AllEcsConfigKeys.CONTAINER_NAME)
self.attempts_since_last_successful_connection = 0
+ # ``ExecutorConf`` has no ``getint`` before Airflow 3.2.
+ self.run_task_concurrency = int(
+ self.conf.get(
+ CONFIG_GROUP_NAME,
+ AllEcsConfigKeys.RUN_TASK_CONCURRENCY,
+
fallback=CONFIG_DEFAULTS[AllEcsConfigKeys.RUN_TASK_CONCURRENCY],
+ )
+ )
+ if self.run_task_concurrency < 1:
+ raise ValueError(
+ f"{AllEcsConfigKeys.RUN_TASK_CONCURRENCY} must be 1 or more,
got {self.run_task_concurrency}"
+ )
+ # With a concurrency of 1, RunTask calls are made one by one in the
scheduler loop.
+ self._run_task_threads = 0
+ self._run_task_pool: ThreadPoolExecutor | None = None
+ if self.run_task_concurrency > 1:
+ self._run_task_threads = self.run_task_concurrency
+ self._run_task_pool = ThreadPoolExecutor(
+ max_workers=self._run_task_threads,
thread_name_prefix="ecs-run-task"
+ )
+ self._run_task_futures: dict[Future, EcsQueuedTask] = {}
+
self.load_ecs_connection(check_connection=False)
self.IS_BOTO_CONNECTION_HEALTHY = False
@@ -271,7 +300,15 @@ class AwsEcsExecutor(BaseExecutor):
fallback=CONFIG_DEFAULTS[AllEcsConfigKeys.AWS_CONN_ID],
)
region_name = self.conf.get(CONFIG_GROUP_NAME,
AllEcsConfigKeys.REGION_NAME, fallback=None)
- self.ecs = EcsHook(aws_conn_id=aws_conn_id,
region_name=region_name).conn
+ hook = EcsHook(aws_conn_id=aws_conn_id, region_name=region_name)
+ # Each RunTask thread, and the scheduler thread, needs a connection of
its own: past the client's pool
+ # size (10 by default), the extra connections are closed with a
warning when a burst of calls is over,
+ # and the next burst opens them again.
+ config = hook.config
+ connections = self._run_task_threads + 1
+ if config.max_pool_connections < connections:
+ config = config.merge(Config(max_pool_connections=connections))
+ self.ecs = hook.get_client_type(region_name=hook.region_name,
config=config)
self.attempts_since_last_successful_connection += 1
self.last_connection_reload = timezone.utcnow()
@@ -424,83 +461,120 @@ class AwsEcsExecutor(BaseExecutor):
called again in the next heart-beat.
If the launch type is FARGATE, this will run the tasks on new AWS
Fargate instances.
+
+ With ``run_task_concurrency`` above 1, the RunTask calls are made in
background threads, at most
+ that many at a time, and the scheduler loop does not wait for them:
each call is handled by the
+ first sync that finds it finished.
"""
- queue_len = len(self.pending_workloads)
- failure_reasons = defaultdict(int)
- for _ in range(queue_len):
+ for _ in range(len(self.pending_workloads)):
ecs_task = self.pending_workloads.popleft()
- task_key = ecs_task.key
- cmd = ecs_task.command
- queue = ecs_task.queue
- exec_config = ecs_task.executor_config
- attempt_number = ecs_task.attempt_number
- failure_reasons = []
if timezone.utcnow() < ecs_task.next_attempt_time:
self.pending_workloads.append(ecs_task)
continue
+ run_task = partial(
+ self._run_task, ecs_task.key, ecs_task.command,
ecs_task.queue, ecs_task.executor_config
+ )
+ if self._run_task_pool is None:
+ self._handle_run_task_response(ecs_task, run_task)
+ else:
+ self._run_task_futures[self._run_task_pool.submit(run_task)] =
ecs_task
+ # All the calls made with expired credentials fail together: raising
on the first one would leave
+ # the others to fail one sync at a time, each forcing its own
reconnection.
+ credentials_error: ClientError | NoCredentialsError | None = None
+ for future in [future for future in self._run_task_futures if
future.done()]:
try:
- run_task_response = self._run_task(task_key, cmd, queue,
exec_config)
- except NoCredentialsError:
+
self._handle_run_task_response(self._run_task_futures.pop(future),
future.result)
+ except (ClientError, NoCredentialsError) as error:
+ credentials_error = credentials_error or error
+ if credentials_error is not None:
+ raise credentials_error
+
+ def _handle_run_task_response(self, ecs_task: EcsQueuedTask, get_response:
Callable[[], dict]) -> None:
+ """Get the RunTask response for ``ecs_task``, then track the ECS task,
or retry or fail the workload."""
+ task_key = ecs_task.key
+ cmd = ecs_task.command
+ queue = ecs_task.queue
+ exec_config = ecs_task.executor_config
+ attempt_number = ecs_task.attempt_number
+ failure_reasons = []
+ try:
+ run_task_response = get_response()
+ except NoCredentialsError:
+ self.pending_workloads.append(ecs_task)
+ raise
+ except ClientError as e:
+ error_code = e.response["Error"]["Code"]
+ if error_code in INVALID_CREDENTIALS_EXCEPTIONS:
self.pending_workloads.append(ecs_task)
raise
- except ClientError as e:
- error_code = e.response["Error"]["Code"]
- if error_code in INVALID_CREDENTIALS_EXCEPTIONS:
- self.pending_workloads.append(ecs_task)
- raise
- failure_reasons.append(str(e))
- except Exception as e:
- # Failed to even get a response back from the Boto3 API or
something else went
- # wrong. For any possible failure we want to add the
exception reasons to the
- # failure list so that it is logged to the user and most
importantly the task is
- # added back to the pending list to be retried later.
- failure_reasons.append(str(e))
+ if (
+ error_code in THROTTLING_EXCEPTIONS
+ and ecs_task.throttled_attempts <
MAX_THROTTLED_RUN_TASK_RETRIES
+ ):
+ # ECS is asking to slow down; that says nothing about the
workload, so keep its attempts.
+ ecs_task.throttled_attempts += 1
+ delay =
calculate_next_attempt_delay(ecs_task.throttled_attempts, exponent_base=2)
+ self.log.info(
+ "ECS throttled the RunTask call for %s, retrying in %s
(retry %s of %s)",
+ task_key,
+ delay,
+ ecs_task.throttled_attempts,
+ MAX_THROTTLED_RUN_TASK_RETRIES,
+ )
+ ecs_task.next_attempt_time = timezone.utcnow() + delay
+ self.pending_workloads.append(ecs_task)
+ return
+ failure_reasons.append(str(e))
+ except Exception as e:
+ # Failed to even get a response back from the Boto3 API or
something else went
+ # wrong. For any possible failure we want to add the exception
reasons to the
+ # failure list so that it is logged to the user and most
importantly the task is
+ # added back to the pending list to be retried later.
+ failure_reasons.append(str(e))
+ else:
+ # We got a response back, check if there were failures. If so, add
them to the
+ # failures list so that it is logged to the user and most
importantly the task
+ # is added back to the pending list to be retried later.
+ if run_task_response["failures"]:
+ failure_reasons.extend([f["reason"] for f in
run_task_response["failures"]])
+
+ if failure_reasons:
+ # Make sure the number of attempts does not exceed
max_run_task_attempts
+ if int(attempt_number) < int(self.max_run_task_attempts):
+ ecs_task.attempt_number += 1
+ ecs_task.next_attempt_time = timezone.utcnow() +
calculate_next_attempt_delay(attempt_number)
+ self.pending_workloads.append(ecs_task)
else:
- # We got a response back, check if there were failures. If so,
add them to the
- # failures list so that it is logged to the user and most
importantly the task
- # is added back to the pending list to be retried later.
- if run_task_response["failures"]:
- failure_reasons.extend([f["reason"] for f in
run_task_response["failures"]])
-
- if failure_reasons:
- # Make sure the number of attempts does not exceed
max_run_task_attempts
- if int(attempt_number) < int(self.max_run_task_attempts):
- ecs_task.attempt_number += 1
- ecs_task.next_attempt_time = timezone.utcnow() +
calculate_next_attempt_delay(
- attempt_number
- )
- self.pending_workloads.append(ecs_task)
- else:
- reasons_str = ", ".join(failure_reasons)
- self.log.error(
- "ECS workload %s has failed a maximum of %s times.
Marking as failed. Reasons: %s",
- task_key,
- attempt_number,
- reasons_str,
- )
- self.log_task_event(
- event="ecs task submit failure",
- ti_key=task_key,
- extra=(
- f"Task could not be queued after {attempt_number}
attempts. "
- f"Marking as failed. Reasons: {reasons_str}"
- ),
- )
- self.fail(task_key)
- elif not run_task_response["tasks"]:
- self.log.error("ECS RunTask Response: %s", run_task_response)
+ reasons_str = ", ".join(failure_reasons)
+ self.log.error(
+ "ECS workload %s has failed a maximum of %s times. Marking
as failed. Reasons: %s",
+ task_key,
+ attempt_number,
+ reasons_str,
+ )
self.log_task_event(
event="ecs task submit failure",
- extra=f"ECS RunTask Response: {run_task_response}",
ti_key=task_key,
+ extra=(
+ f"Task could not be queued after {attempt_number}
attempts. "
+ f"Marking as failed. Reasons: {reasons_str}"
+ ),
)
- raise EcsExecutorException(
- "No failures and no ECS tasks provided in response. This
should never happen."
- )
- else:
- task = run_task_response["tasks"][0]
- self.active_workers.add_task(task, task_key, queue, cmd,
exec_config, attempt_number)
- self.running_state(task_key, task.task_arn)
+ self.fail(task_key)
+ elif not run_task_response["tasks"]:
+ self.log.error("ECS RunTask Response: %s", run_task_response)
+ self.log_task_event(
+ event="ecs task submit failure",
+ extra=f"ECS RunTask Response: {run_task_response}",
+ ti_key=task_key,
+ )
+ raise EcsExecutorException(
+ "No failures and no ECS tasks provided in response. This
should never happen."
+ )
+ else:
+ task = run_task_response["tasks"][0]
+ self.active_workers.add_task(task, task_key, queue, cmd,
exec_config, attempt_number)
+ self.running_state(task_key, task.task_arn)
def _run_task(
self,
@@ -517,7 +591,8 @@ class AwsEcsExecutor(BaseExecutor):
section of the JSON request before calling Boto3's "run_task" function.
"""
run_task_kwargs = self._run_task_kwargs(task_id, cmd, queue,
exec_config)
- boto_run_task = self.ecs.run_task(**run_task_kwargs)
+ with Stats.timer("ecs_executor.run_task.duration",
tags=prune_dict({"team_name": self.team_name})):
+ boto_run_task = self.ecs.run_task(**run_task_kwargs)
run_task_response = BotoRunTaskSchema().load(boto_run_task)
return run_task_response
@@ -570,17 +645,34 @@ class AwsEcsExecutor(BaseExecutor):
try:
while True:
self.sync()
- if not self.active_workers:
+ if not self.active_workers and not self._run_task_futures:
break
time.sleep(heartbeat_interval)
except Exception:
# We catch any and all exceptions because otherwise they would
bubble
# up and kill the scheduler process.
self.log.exception("Failed to end %s", self.__class__.__name__)
+ if self._run_task_pool is not None:
+ self._run_task_pool.shutdown(wait=False, cancel_futures=True)
def terminate(self):
"""Kill all ECS processes by calling Boto3's StopTask API."""
try:
+ if self._run_task_pool is not None:
+ # A RunTask call still in flight may start an ECS task after
the StopTask calls below:
+ # wait for the calls already started, and drop the ones not
started yet.
+ self._run_task_pool.shutdown(wait=True, cancel_futures=True)
+ # The pool no longer takes calls: ``end()`` below makes them
in the scheduler loop.
+ self._run_task_pool = None
+ for future, ecs_task in list(self._run_task_futures.items()):
+ del self._run_task_futures[future]
+ if future.cancelled():
+ continue
+ # One failure must not keep the other ECS tasks from being
tracked, then stopped.
+ try:
+ self._handle_run_task_response(ecs_task, future.result)
+ except Exception:
+ self.log.exception("Failed to handle the RunTask
response for %s", ecs_task.key)
for arn in self.active_workers.get_all_arns():
self.ecs.stop_task(
cluster=self.cluster, task=arn, reason="Airflow Executor
received a SIGTERM"
diff --git
a/providers/amazon/src/airflow/providers/amazon/aws/executors/ecs/utils.py
b/providers/amazon/src/airflow/providers/amazon/aws/executors/ecs/utils.py
index fa36377abfd..6d49c4030ba 100644
--- a/providers/amazon/src/airflow/providers/amazon/aws/executors/ecs/utils.py
+++ b/providers/amazon/src/airflow/providers/amazon/aws/executors/ecs/utils.py
@@ -61,6 +61,7 @@ CONFIG_DEFAULTS = {
"assign_public_ip": "False",
"platform_version": "LATEST",
"check_health_on_startup": "True",
+ "run_task_concurrency": "1",
}
@@ -74,6 +75,7 @@ class EcsQueuedTask:
executor_config: ExecutorConfigType
attempt_number: int
next_attempt_time: datetime.datetime
+ throttled_attempts: int = 0
@dataclass
@@ -106,6 +108,7 @@ class AllEcsConfigKeys(RunTaskKwargsConfigKeys):
CHECK_HEALTH_ON_STARTUP = "check_health_on_startup"
MAX_RUN_TASK_ATTEMPTS = "max_run_task_attempts"
REGION_NAME = "region_name"
+ RUN_TASK_CONCURRENCY = "run_task_concurrency"
RUN_TASK_KWARGS = "run_task_kwargs"
diff --git a/providers/amazon/src/airflow/providers/amazon/get_provider_info.py
b/providers/amazon/src/airflow/providers/amazon/get_provider_info.py
index 76bfd3ec6e2..71a338b87c5 100644
--- a/providers/amazon/src/airflow/providers/amazon/get_provider_info.py
+++ b/providers/amazon/src/airflow/providers/amazon/get_provider_info.py
@@ -1554,6 +1554,13 @@ def get_provider_info():
"example": "3",
"default": "3",
},
+ "run_task_concurrency": {
+ "description": "The maximum number of ECS RunTask
calls the ECS Executor makes in parallel. With 1, tasks are\nlaunched one at a
time and the scheduler loop waits for each call. With more, the calls run
in\nbackground threads and the scheduler loop does not wait for them.\n",
+ "version_added": "9.40.0",
+ "type": "integer",
+ "example": "16",
+ "default": "1",
+ },
"run_task_kwargs": {
"description": "A JSON string containing arguments to
provide the ECS `run_task` API (see url above).\n",
"version_added": "8.10",
diff --git
a/providers/amazon/tests/unit/amazon/aws/executors/ecs/test_ecs_executor.py
b/providers/amazon/tests/unit/amazon/aws/executors/ecs/test_ecs_executor.py
index 9b5d2e96108..12b816588f5 100644
--- a/providers/amazon/tests/unit/amazon/aws/executors/ecs/test_ecs_executor.py
+++ b/providers/amazon/tests/unit/amazon/aws/executors/ecs/test_ecs_executor.py
@@ -22,14 +22,17 @@ import json
import logging
import os
import re
+import threading
import time
from collections.abc import Callable
+from concurrent.futures import Future, wait
from functools import partial
from unittest import mock
from unittest.mock import MagicMock, patch
from uuid import UUID, uuid4
import pytest
+import time_machine
import yaml
from botocore.exceptions import ClientError
from inflection import camelize
@@ -49,6 +52,7 @@ from airflow.providers.amazon.aws.executors.ecs.ecs_executor
import (
from airflow.providers.amazon.aws.executors.ecs.utils import (
CONFIG_DEFAULTS,
EcsExecutorTask,
+ EcsQueuedTask,
_recursive_flatten_dict,
parse_assign_public_ip,
)
@@ -1406,6 +1410,262 @@ class TestAwsEcsExecutor:
mock_timer.assert_called_once_with("ecs_executor.adopt_task_instances.duration",
tags=expected_tags)
[email protected]
+def concurrent_executor(set_env_vars):
+ with conf_vars({(CONFIG_GROUP_NAME,
AllEcsConfigKeys.RUN_TASK_CONCURRENCY): "2"}):
+ executor = AwsEcsExecutor()
+ executor.IS_BOTO_CONNECTION_HEALTHY = True
+ executor.ecs = mock.Mock(spec=executor.ecs)
+ yield executor
+ if executor._run_task_pool is not None:
+ executor._run_task_pool.shutdown(wait=True, cancel_futures=True)
+
+
+class TestRunTaskConcurrency:
+ """RunTask calls made in background threads when ``run_task_concurrency``
is above 1."""
+
+ def test_scheduler_loop_does_not_wait_for_run_task(self,
concurrent_executor):
+ started = threading.Semaphore(0)
+ release = threading.Event()
+
+ def run_task(**kwargs):
+ started.release()
+ release.wait(timeout=10)
+ task_number =
kwargs["overrides"]["containerOverrides"][0]["command"][-1]
+ task = {
+ "taskArn": f"{ARN1}-{task_number}",
+ "lastStatus": "",
+ "desiredStatus": "",
+ "containers": [{"name": "some-ecs-container"}],
+ }
+ return {"tasks": [task], "failures": []}
+
+ concurrent_executor.ecs.run_task.side_effect = run_task
+ keys = [TaskInstanceKey("dag", f"task_{i}", "run", 1, -1) for i in
range(3)]
+ for i, key in enumerate(keys):
+ concurrent_executor.execute_async(key, ["airflow", "tasks",
str(i)])
+
+ concurrent_executor.attempt_workload_runs()
+
+ assert started.acquire(timeout=10)
+ assert started.acquire(timeout=10)
+ assert not started.acquire(timeout=0.2), "more RunTask calls in flight
than run_task_concurrency"
+ assert not concurrent_executor.pending_workloads
+ assert len(concurrent_executor._run_task_futures) == 3
+ assert len(concurrent_executor.active_workers) == 0
+
+ release.set()
+ wait(list(concurrent_executor._run_task_futures), timeout=10)
+ concurrent_executor.attempt_workload_runs()
+
+ assert not concurrent_executor._run_task_futures
+ assert set(concurrent_executor.active_workers.get_all_task_keys()) ==
set(keys)
+ assert set(concurrent_executor.active_workers.get_all_arns()) ==
{f"{ARN1}-{i}" for i in range(3)}
+
+ def test_credentials_error_is_raised_once_for_all_finished_calls(self,
concurrent_executor):
+ release = threading.Event()
+
+ def run_task(**kwargs):
+ release.wait(timeout=10)
+ raise ClientError({"Error": {"Code": "ExpiredTokenException",
"Message": "expired"}}, "RunTask")
+
+ concurrent_executor.ecs.run_task.side_effect = run_task
+ for i in range(2):
+ concurrent_executor.execute_async(
+ TaskInstanceKey("dag", f"task_{i}", "run", 1, -1), ["airflow",
"tasks", str(i)]
+ )
+ concurrent_executor.attempt_workload_runs()
+ release.set()
+ wait(list(concurrent_executor._run_task_futures), timeout=10)
+
+ with pytest.raises(ClientError, match="ExpiredTokenException"):
+ concurrent_executor.attempt_workload_runs()
+
+ assert not concurrent_executor._run_task_futures
+ assert len(concurrent_executor.pending_workloads) == 2
+
+ def test_end_waits_for_in_flight_run_task_calls(self, concurrent_executor):
+ in_flight = Future()
+ concurrent_executor._run_task_futures[in_flight] =
mock.Mock(spec=EcsQueuedTask)
+
+ def finish_run_task_on_second_sync():
+ if concurrent_executor.sync.call_count == 2:
+ del concurrent_executor._run_task_futures[in_flight]
+
+ concurrent_executor.sync = mock.create_autospec(
+ concurrent_executor.sync,
side_effect=finish_run_task_on_second_sync
+ )
+
+ concurrent_executor.end(heartbeat_interval=0)
+
+ assert concurrent_executor.sync.call_count == 2
+
+ def test_terminate_stops_the_tasks_of_run_task_calls_in_flight(self,
concurrent_executor):
+ started = threading.Semaphore(0)
+ release = threading.Event()
+
+ def run_task(**kwargs):
+ started.release()
+ release.wait(timeout=10)
+ task_number =
kwargs["overrides"]["containerOverrides"][0]["command"][-1]
+ task = {
+ "taskArn": f"{ARN1}-{task_number}",
+ "lastStatus": "",
+ "desiredStatus": "",
+ "containers": [{"name": "some-ecs-container"}],
+ }
+ return {"tasks": [task], "failures": []}
+
+ concurrent_executor.ecs.run_task.side_effect = run_task
+ for i in range(3):
+ concurrent_executor.execute_async(
+ TaskInstanceKey("dag", f"task_{i}", "run", 1, -1), ["airflow",
"tasks", str(i)]
+ )
+ concurrent_executor.attempt_workload_runs()
+ # Two calls run (run_task_concurrency is 2), the third waits for a
thread.
+ assert started.acquire(timeout=10)
+ assert started.acquire(timeout=10)
+ threading.Timer(0.2, release.set).start()
+
+ with mock.patch.object(concurrent_executor, "end"):
+ concurrent_executor.terminate()
+
+ assert concurrent_executor.ecs.run_task.call_count == 2
+ assert not concurrent_executor._run_task_futures
+ stopped = {call.kwargs["task"] for call in
concurrent_executor.ecs.stop_task.call_args_list}
+ assert stopped == {f"{ARN1}-0", f"{ARN1}-1"}
+
+ def
test_terminate_stops_the_tasks_started_when_another_response_fails(self,
concurrent_executor):
+ task = {
+ "taskArn": ARN1,
+ "lastStatus": "",
+ "desiredStatus": "",
+ "containers": [{"name": "some-ecs-container"}],
+ }
+ release = threading.Event()
+
+ def run_task(**kwargs):
+ release.wait(timeout=10)
+ if kwargs["overrides"]["containerOverrides"][0]["command"][-1] ==
"0":
+ # Neither tasks nor failures: handling this response raises an
EcsExecutorException.
+ return {"tasks": [], "failures": []}
+ return {"tasks": [task], "failures": []}
+
+ concurrent_executor.ecs.run_task.side_effect = run_task
+ for i in range(2):
+ concurrent_executor.execute_async(
+ TaskInstanceKey("dag", f"task_{i}", "run", 1, -1), ["airflow",
"tasks", str(i)]
+ )
+ concurrent_executor.attempt_workload_runs()
+ release.set()
+ wait(list(concurrent_executor._run_task_futures), timeout=10)
+
+ with mock.patch.object(concurrent_executor, "end"):
+ concurrent_executor.terminate()
+
+ concurrent_executor.ecs.stop_task.assert_called_once_with(
+ cluster=concurrent_executor.cluster, task=ARN1, reason="Airflow
Executor received a SIGTERM"
+ )
+
+ def
test_run_task_calls_are_made_in_the_scheduler_loop_after_terminate(self,
concurrent_executor):
+ task = {
+ "taskArn": ARN1,
+ "lastStatus": "",
+ "desiredStatus": "",
+ "containers": [{"name": "some-ecs-container"}],
+ }
+ concurrent_executor.ecs.run_task.return_value = {"tasks": [task],
"failures": []}
+ with mock.patch.object(concurrent_executor, "end"):
+ concurrent_executor.terminate()
+ key = TaskInstanceKey("dag", "task", "run", 1, -1)
+ concurrent_executor.execute_async(key, ["airflow", "tasks", "run"])
+
+ concurrent_executor.attempt_workload_runs()
+
+ assert not concurrent_executor._run_task_futures
+ assert concurrent_executor.active_workers.get_all_task_keys() == [key]
+
+ @pytest.mark.parametrize("concurrency", ["0", "-1"])
+ def test_run_task_concurrency_below_one_is_rejected(self, set_env_vars,
concurrency):
+ with conf_vars({(CONFIG_GROUP_NAME,
AllEcsConfigKeys.RUN_TASK_CONCURRENCY): concurrency}):
+ with pytest.raises(
+ ValueError, match=f"run_task_concurrency must be 1 or more,
got {concurrency}"
+ ):
+ AwsEcsExecutor()
+
+ @pytest.mark.parametrize(
+ ("concurrency", "connection_config", "expected_pool_size"),
+ [
+ pytest.param("16", {}, 17,
id="one-per-thread-and-for-the-scheduler-thread"),
+ pytest.param("16", {"max_pool_connections": 50}, 50,
id="larger-connection-setting-kept"),
+ ],
+ )
+ def test_client_has_a_connection_for_each_thread(
+ self, set_env_vars, monkeypatch, concurrency, connection_config,
expected_pool_size
+ ):
+ config_kwargs = {"connect_timeout": 7, **connection_config}
+ monkeypatch.setenv(
+ "AIRFLOW_CONN_AWS_DEFAULT",
+ json.dumps({"conn_type": "aws", "extra": {"config_kwargs":
config_kwargs}}),
+ )
+ with conf_vars({(CONFIG_GROUP_NAME,
AllEcsConfigKeys.RUN_TASK_CONCURRENCY): concurrency}):
+ executor = AwsEcsExecutor(parallelism=32)
+
+ assert executor.ecs.meta.config.max_pool_connections ==
expected_pool_size
+ assert executor.ecs.meta.config.connect_timeout == 7
+
+
+class TestRunTaskThrottling:
+ @pytest.mark.parametrize("concurrency", ["1", "2"])
+ @time_machine.travel("2026-10-27 12:00:00", tick=False)
+ def test_throttled_run_task_is_retried_without_using_an_attempt(self,
set_env_vars, concurrency):
+ with conf_vars({(CONFIG_GROUP_NAME,
AllEcsConfigKeys.RUN_TASK_CONCURRENCY): concurrency}):
+ executor = AwsEcsExecutor()
+ executor.ecs = mock.Mock(spec=executor.ecs)
+ executor.ecs.run_task.side_effect = ClientError(
+ {"Error": {"Code": "ThrottlingException", "Message": "Rate
exceeded"}}, "RunTask"
+ )
+ key = TaskInstanceKey("dag", "task", "run", 1, -1)
+ executor.execute_async(key, ["airflow", "tasks", "run"])
+
+ executor.attempt_workload_runs()
+ if executor._run_task_pool is not None:
+ wait(list(executor._run_task_futures), timeout=10)
+ executor.attempt_workload_runs()
+ executor._run_task_pool.shutdown(wait=True)
+
+ assert executor.ecs.run_task.call_count == 1
+ [ecs_task] = executor.pending_workloads
+ assert ecs_task.attempt_number == 1
+ assert ecs_task.throttled_attempts == 1
+ assert ecs_task.next_attempt_time == timezone.utcnow() +
dt.timedelta(seconds=2)
+ assert key not in executor.event_buffer
+
+ @time_machine.travel("2026-10-27 12:00:00", tick=False)
+ def test_throttled_run_task_retries_back_off_then_use_attempts(self,
mock_executor):
+ mock_executor.ecs.run_task.side_effect = ClientError(
+ {"Error": {"Code": "ThrottlingException", "Message": "Rate
exceeded"}}, "RunTask"
+ )
+ key = TaskInstanceKey("dag", "task", "run", 1, -1)
+ mock_executor.execute_async(key, ["airflow", "tasks", "run"])
+ delays = []
+ for _ in range(ecs_executor.MAX_THROTTLED_RUN_TASK_RETRIES + 1):
+ mock_executor.pending_workloads[0].next_attempt_time =
timezone.utcnow()
+ mock_executor.attempt_workload_runs()
+
delays.append((mock_executor.pending_workloads[0].next_attempt_time -
timezone.utcnow()).seconds)
+
+ assert delays == [2, 4, 8, 16, 32, 4]
+ assert mock_executor.pending_workloads[0].attempt_number == 2
+
+ @mock.patch.object(ecs_executor, "Stats")
+ def test_run_task_duration_is_timed(self, stats_mock, mock_executor):
+ mock_executor.execute_async(TaskInstanceKey("dag", "task", "run", 1,
-1), ["airflow", "tasks", "run"])
+
+ mock_executor.attempt_workload_runs()
+
+
stats_mock.timer.assert_called_once_with("ecs_executor.run_task.duration",
tags={})
+
+
class TestEcsExecutorConfig:
@pytest.fixture
def assign_subnets(self):
@@ -1537,12 +1797,13 @@ class TestEcsExecutorConfig:
found_keys = {convert_camel_to_snake(key): key for key in
task_kwargs.keys()}
for expected_key, expected_value_raw in CONFIG_DEFAULTS.items():
- # conn_id, max_run_task_attempts, and check_health_on_startup are
used by the executor,
- # but are not expected to appear in the task_kwargs.
+ # conn_id, max_run_task_attempts, check_health_on_startup and
run_task_concurrency are used by
+ # the executor, but are not expected to appear in the task_kwargs.
if expected_key in [
AllEcsConfigKeys.AWS_CONN_ID,
AllEcsConfigKeys.MAX_RUN_TASK_ATTEMPTS,
AllEcsConfigKeys.CHECK_HEALTH_ON_STARTUP,
+ AllEcsConfigKeys.RUN_TASK_CONCURRENCY,
]:
assert expected_key not in found_keys.keys()
else:
diff --git
a/shared/observability/src/airflow_shared/observability/metrics/metrics_template.yaml
b/shared/observability/src/airflow_shared/observability/metrics/metrics_template.yaml
index 4a70cc9eb1b..5758848964f 100644
---
a/shared/observability/src/airflow_shared/observability/metrics/metrics_template.yaml
+++
b/shared/observability/src/airflow_shared/observability/metrics/metrics_template.yaml
@@ -961,6 +961,12 @@ metrics:
legacy_name: "-"
name_variables: []
+ - name: "ecs_executor.run_task.duration"
+ description: "Milliseconds taken by one ECS RunTask call of the AWS ECS
Executor"
+ type: "timer"
+ legacy_name: "-"
+ name_variables: []
+
- name: "lambda_executor.adopt_task_instances.duration"
description: "Milliseconds taken to adopt the task instances in the AWS
Lambda Executor"
type: "timer"
