o-nikolas commented on code in PR #75907:
URL: https://github.com/apache/airflow/pull/75907#discussion_r4407364423
##########
providers/amazon/src/airflow/providers/amazon/aws/executors/ecs/ecs_executor.py:
##########
@@ -424,83 +463,100 @@ def attempt_workload_runs(self):
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`` other than 1, the RunTask calls are made
in background threads, at
+ most that many at a time (0: no limit), 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
- try:
- run_task_response = self._run_task(task_key, cmd, queue,
exec_config)
- except NoCredentialsError:
+ 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
+ for future in [future for future in self._run_task_futures if
future.done()]:
+ self._handle_run_task_response(self._run_task_futures.pop(future),
future.result)
+
+ 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:
Review Comment:
Should we log something in this branch so the user has some indication of
what's happening? It might be noisy but silence doesn't quite seem right
either. Maybe debug level logging is the right middle ground?
##########
providers/amazon/src/airflow/providers/amazon/aws/executors/ecs/ecs_executor.py:
##########
@@ -143,6 +149,31 @@ def __init__(self, *args, **kwargs):
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 < 0:
+ raise ValueError(
+ f"{AllEcsConfigKeys.RUN_TASK_CONCURRENCY} must be 0 (no limit)
or more, "
+ f"got {self.run_task_concurrency}"
+ )
+ # With a concurrency of 1, RunTask calls are made one by one in the
scheduler loop. With 0, there is
+ # no limit: the executor never has more than ``parallelism`` workloads
to launch at once, and the
+ # pool only starts the threads that the calls in flight need.
+ 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 or
self.parallelism
+ self._run_task_pool = ThreadPoolExecutor(
+ max_workers=self._run_task_threads,
thread_name_prefix="ecs-run-task"
+ )
Review Comment:
On Airflow 2.11, which we still support, parallelism = 0 means unlimited and
is often used (if not the default in some cases?) which I don't think is
suitable for making this decision here. It makes this
`ThreadPoolExecutor(max_workers=0)` (which I'm not sure works?). On 3.x with a
large parallelism, it starts that many threads and connections. We might have
to re-think zero as unlimited, and parallelism is actually a good example of
this, we ran into issues with it being unlimited and instead we ask users to
make an informed decision themselves.
##########
providers/amazon/src/airflow/providers/amazon/aws/executors/ecs/ecs_executor.py:
##########
@@ -424,83 +463,100 @@ def attempt_workload_runs(self):
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`` other than 1, the RunTask calls are made
in background threads, at
+ most that many at a time (0: no limit), 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
- try:
- run_task_response = self._run_task(task_key, cmd, queue,
exec_config)
- except NoCredentialsError:
+ 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
+ for future in [future for future in self._run_task_futures if
future.done()]:
+ self._handle_run_task_response(self._run_task_futures.pop(future),
future.result)
Review Comment:
If credentials expire while several calls are in flight, each failed future
raises on its own sync. The first one marks the connection unhealthy. After the
reconnect, the next stale future marks it unhealthy again, so N failed calls
cause N reconnects. Could this loop handle all finished futures and re-raise
the credentials error once at the end?
##########
providers/amazon/src/airflow/providers/amazon/aws/executors/ecs/ecs_executor.py:
##########
@@ -570,17 +627,27 @@ def end(self, heartbeat_interval=10):
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)
Review Comment:
In `terminate()` here we shut the pool down before `self.end()`. `end()` →
`sync()` → `submit()` then raises RuntimeError: cannot schedule new futures
after shutdown for any pending workload that is due. That logs a traceback and
drops the workload.
##########
providers/amazon/src/airflow/providers/amazon/aws/executors/ecs/ecs_executor.py:
##########
@@ -570,17 +627,27 @@ def end(self, heartbeat_interval=10):
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)
+ for future, ecs_task in list(self._run_task_futures.items()):
+ del self._run_task_futures[future]
+ if not future.cancelled():
+ self._handle_run_task_response(ecs_task, future.result)
Review Comment:
if `_handle_run_task_response` raises here (credentials error, or the
"should never happen" EcsExecutorException), the StopTask loop at line 651 is
skipped for every running task. Could each result be handled in its own try, so
the StopTask calls still run?
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
