ashb commented on code in PR #23690:
URL: https://github.com/apache/airflow/pull/23690#discussion_r874926571
##########
airflow/executors/celery_executor.py:
##########
@@ -315,25 +320,23 @@ def sync(self) -> None:
self.log.debug("No task to query celery, skipping sync")
return
self.update_all_task_states()
+ self._check_for_stalled_tasks()
- if self.adopted_task_timeouts:
- self._check_for_stalled_adopted_tasks()
-
- def _check_for_stalled_adopted_tasks(self):
+ @provide_session
+ def _check_for_stalled_tasks(self, session: Session = NEW_SESSION):
Review Comment:
Refactoring for clarity: please split this in to two ("prviate") functions,
one that deals with adopted tasks, and another that deals with stalled.
##########
airflow/executors/celery_executor.py:
##########
@@ -469,6 +505,15 @@ def try_adopt_task_instances(self, tis:
List[TaskInstance]) -> List[TaskInstance
return not_adopted_tis
+ def _set_stalled_task_timeout(self, key: TaskInstanceKey, timeout:
datetime.timedelta) -> None:
+ if timeout:
+ self.stalled_task_timeouts[key] = utcnow() + timeout
+ self.stalled_task_timeouts = dict(
+ sorted(self.stalled_task_timeouts.items(), key=lambda item:
item[1])
+ )
Review Comment:
Hmmm, the performance implications of re-sorting this _every time_ we add
one task is not great (n^2 complexity)
We should look at doing this better -- either find a data structure that
stores this sorted on insert (a "heap") or only do this sorting once at the
end, not on every write.
🔋 https://docs.python.org/3/library/heapq.html
##########
airflow/executors/celery_executor.py:
##########
@@ -343,20 +346,53 @@ def _check_for_stalled_adopted_tasks(self):
# already finished, then it will be removed from this list -- so
# the only time it's still in this list is when it a) never made it
# to celery in the first place (i.e. race condition somewhere in
- # the dying executor) or b) a really long celery queue and it just
+ # the dying executor), b) celery lost the task before execution
+ # started, or c) a really long celery queue and it just
# hasn't started yet -- better cancel it and let the scheduler
# re-queue rather than have this task risk stalling for ever
timedout_keys.append(key)
- if timedout_keys:
- self.log.error(
- "Adopted tasks were still pending after %s, assuming they
never made it to celery and "
- "clearing:\n\t%s",
- self.task_adoption_timeout,
- "\n\t".join(repr(x) for x in timedout_keys),
+ if not timedout_keys:
+ return
+
+ self.log.error(
+ "Tasks were still pending after configured timeout (adopted: %s,
all: %s), "
+ "assuming they never made it to celery and clearing:\n\t%s",
Review Comment:
```suggestion
"assuming they never made it to celery and sending back to the
scheduler:\n\t%s",
```
"clearing" has a slightly different meaning which isn't what we do here.
##########
airflow/executors/celery_executor.py:
##########
@@ -343,20 +346,53 @@ def _check_for_stalled_adopted_tasks(self):
# already finished, then it will be removed from this list -- so
# the only time it's still in this list is when it a) never made it
# to celery in the first place (i.e. race condition somewhere in
- # the dying executor) or b) a really long celery queue and it just
+ # the dying executor), b) celery lost the task before execution
+ # started, or c) a really long celery queue and it just
# hasn't started yet -- better cancel it and let the scheduler
# re-queue rather than have this task risk stalling for ever
timedout_keys.append(key)
- if timedout_keys:
- self.log.error(
- "Adopted tasks were still pending after %s, assuming they
never made it to celery and "
- "clearing:\n\t%s",
- self.task_adoption_timeout,
- "\n\t".join(repr(x) for x in timedout_keys),
+ if not timedout_keys:
+ return
+
+ self.log.error(
+ "Tasks were still pending after configured timeout (adopted: %s,
all: %s), "
+ "assuming they never made it to celery and clearing:\n\t%s",
+ self.task_adoption_timeout,
+ self.stalled_task_timeout,
+ "\n\t".join(repr(x) for x in timedout_keys),
+ )
+
+ try:
+ filter_for_tis = TaskInstance.filter_for_tis(timedout_keys)
+ session.query(TaskInstance).filter(
+ filter_for_tis,
+ TaskInstance.state == State.QUEUED,
+ TaskInstance.queued_by_job_id == self.job_id,
+ ).update(
+ {
+ TaskInstance.state: State.SCHEDULED,
+ TaskInstance.queued_dttm: None,
+ TaskInstance.queued_by_job_id: None,
+ TaskInstance.external_executor_id: None,
+ },
+ synchronize_session=False,
)
- for key in timedout_keys:
- self.change_state(key, State.FAILED)
+ session.commit()
+ except Exception:
+ self.log.exception("Error clearing stalled tasks")
+ session.rollback()
+ return
+
+ for key in timedout_keys:
+ self.stalled_task_timeouts.pop(key, None)
+ self.running.discard(key)
+ celery_async_result = self.tasks.pop(key, None)
+ if celery_async_result:
+ try:
+ app.control.revoke(celery_async_result.task_id)
+ except Exception:
+ self.log.exception("Error revoking task instance %s from
celery", key)
Review Comment:
Do we know when this might raise exceptions? I'm wondering if logging the
stack trace for _every_ timedout task might be too much?
##########
airflow/config_templates/default_airflow.cfg:
##########
@@ -887,10 +887,17 @@ operation_timeout = 1.0
# or run in HA mode, it can adopt the orphan tasks launched by previous
SchedulerJob.
task_track_started = True
-# Time in seconds after which Adopted tasks are cleared by CeleryExecutor.
This is helpful to clear
-# stalled tasks.
+# Time in seconds after which adopted tasks which are queued in celery are
assumed to be stalled,
+# and are automatically cleared. This setting does the same thing as
``stalled_task_timeout`` but
+# applies specifically to adopted tasks only. When set to 0, the
``stalled_task_timeout`` setting
+# also applies to adopted tasks.
task_adoption_timeout = 600
+# Time in seconds after which tasks queued in celery are assumed to be
stalled, and are automatically
+# cleared. Adopted tasks will instead use the ``task_adoption_timeout``
setting if specified.
Review Comment:
```suggestion
# rescheduled. Adopted tasks will instead use the ``task_adoption_timeout``
setting if specified.
```
--
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]
