Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

retry for failure in the queue #39398

Draft
wants to merge 1 commit into
base: main
Choose a base branch
from
+30 −7
Draft

retry for failure in the queue #39398

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Jump to
Jump to file
Failed to load files.
Diff view
Diff view
9 changes: 9 additions & 0 deletions airflow/executors/base_executor.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -322,6 +322,15 @@ def change_state(
self.log.debug("Could not find key: %s", key)
self.event_buffer[key] = state, info

def failed_in_queue(self, key: TaskInstanceKey, info=None) -> None:
"""
Set failed in queue state for the event.

:param info: Executor information for the task instance
:param key: Unique key for the task instance
"""
self.change_state(key, TaskInstanceState.FAILED_IN_QUEUE, info)

def fail(self, key: TaskInstanceKey, info=None) -> None:
"""
Set fail state for the event.
Expand Down
12 changes: 9 additions & 3 deletions airflow/jobs/backfill_job_runner.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -104,6 +104,7 @@ class _DagRunTaskStatus:
skipped: set[TaskInstanceKey] = attr.ib(factory=set)
succeeded: set[TaskInstanceKey] = attr.ib(factory=set)
failed: set[TaskInstanceKey] = attr.ib(factory=set)
failed_in_queue: set[TaskInstanceKey] = attr.ib(factory=set)
not_ready: set[TaskInstanceKey] = attr.ib(factory=set)
deadlocked: set[TaskInstance] = attr.ib(factory=set)
active_runs: set[DagRun] = attr.ib(factory=set)
Expand Down Expand Up @@ -208,6 +209,10 @@ def _update_counters(self, ti_status: _DagRunTaskStatus, session: Session) -> No
ti_status.failed.add(ti_key)
ti_status.running.pop(ti_key)
continue
if ti.state == TaskInstanceState.FAILED_IN_QUEUE:
self.log.error("Task instance %s failed in queue", ti)
ti_status.failed_in_queue.add(ti_key)
ti_status.queued.pop(ti_key)
# special case: if the task needs to run again put it back
if ti.state == TaskInstanceState.UP_FOR_RETRY:
self.log.warning("Task instance %s is up for retry", ti)
Expand Down Expand Up @@ -448,7 +453,7 @@ def _task_instances_for_dag_run(
def _log_progress(self, ti_status: _DagRunTaskStatus) -> None:
self.log.info(
"[backfill progress] | finished run %s of %s | tasks waiting: %s | succeeded: %s | "
"running: %s | failed: %s | skipped: %s | deadlocked: %s | not ready: %s",
"running: %s | failed: %s | skipped: %s | deadlocked: %s | not ready: %s | failed in queue: %s",
ti_status.finished_runs,
ti_status.total_runs,
len(ti_status.to_run),
Expand All @@ -458,6 +463,7 @@ def _log_progress(self, ti_status: _DagRunTaskStatus) -> None:
len(ti_status.skipped),
len(ti_status.deadlocked),
len(ti_status.not_ready),
len(ti_status.failed_in_queue),
)

self.log.debug("Finished dag run loop iteration. Remaining tasks %s", ti_status.to_run.values())
Expand Down Expand Up @@ -521,7 +527,7 @@ def _per_task_process(key, ti: TaskInstance, session):

if self.rerun_failed_tasks:
# Rerun failed tasks or upstreamed failed tasks
if ti.state in (TaskInstanceState.FAILED, TaskInstanceState.UPSTREAM_FAILED):
if ti.state in (TaskInstanceState.FAILED, TaskInstanceState.UPSTREAM_FAILED, TaskInstanceState.FAILED_IN_QUEUE):
self.log.error("Task instance %s with state %s", ti, ti.state)
if key in ti_status.running:
ti_status.running.pop(key)
Expand All @@ -532,7 +538,7 @@ def _per_task_process(key, ti: TaskInstance, session):
ti_status.active_runs.add(ti.dag_run)
else:
# Default behaviour which works for subdag.
if ti.state in (TaskInstanceState.FAILED, TaskInstanceState.UPSTREAM_FAILED):
if ti.state in (TaskInstanceState.FAILED, TaskInstanceState.UPSTREAM_FAILED,TaskInstanceState.FAILED_IN_QUEUE):
self.log.error("Task instance %s with state %s", ti, ti.state)
ti_status.failed.add(key)
ti_status.to_run.pop(key)
Expand Down
1 change: 1 addition & 0 deletions airflow/jobs/scheduler_job_runner.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -698,6 +698,7 @@ def _process_executor_events(self, session: Session) -> int:
TaskInstanceState.SUCCESS,
TaskInstanceState.QUEUED,
TaskInstanceState.RUNNING,
TaskInstanceState.FAILED_IN_QUEUE,
):
tis_with_right_state.append(ti_key)

Expand Down
2 changes: 1 addition & 1 deletion airflow/models/dag.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2381,7 +2381,7 @@ def clear(

state: list[TaskInstanceState] = []
if only_failed:
state += [TaskInstanceState.FAILED, TaskInstanceState.UPSTREAM_FAILED]
state += [TaskInstanceState.FAILED, TaskInstanceState.UPSTREAM_FAILED, TaskInstanceState.FAILED_IN_QUEUE]
if only_running:
# Yes, having `+=` doesn't make sense, but this was the existing behaviour
state += [TaskInstanceState.RUNNING]
Expand Down
1 change: 1 addition & 0 deletions airflow/models/taskinstance.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -409,6 +409,7 @@ def _stop_remaining_tasks(*, task_instance: TaskInstance | TaskInstancePydantic,
if ti.task_id == task_instance.task_id or ti.state in (
TaskInstanceState.SUCCESS,
TaskInstanceState.FAILED,
TaskInstanceState.FAILED_IN_QUEUE,
):
continue
task = task_instance.task.dag.task_dict[ti.task_id]
Expand Down
2 changes: 1 addition & 1 deletion airflow/operators/subdag.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -151,7 +151,7 @@ def _reset_dag_run_and_task_instances(self, dag_run: DagRun, execution_date: dat
select(TaskInstance)
.where(TaskInstance.dag_id == self.subdag.dag_id)
.where(TaskInstance.execution_date == execution_date)
.where(TaskInstance.state.in_((TaskInstanceState.FAILED, TaskInstanceState.UPSTREAM_FAILED)))
.where(TaskInstance.state.in_((TaskInstanceState.FAILED, TaskInstanceState.UPSTREAM_FAILED, TaskInstanceState.FAILED_IN_QUEUE)))
)

for task_instance in failed_task_instances:
Expand Down
2 changes: 1 addition & 1 deletion airflow/providers/celery/executors/celery_executor.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -452,7 +452,7 @@ def cleanup_stuck_queued_tasks(self, tis: list[TaskInstance]) -> list[str]:
for ti in tis:
readable_tis.append(repr(ti))
task_instance_key = ti.key
self.fail(task_instance_key, None)
self.failed_in_queue(task_instance_key, None)
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This one is backward incompatible change with Airflow 2.7, 2.8 and 2.9.
Maybe it is also a good idea to made separate PR in Core and and after in Provider(s)

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

ah i see. let me do that @Taragolis

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

can you explain me the logic behind splitting it into separate PR @Taragolis ?

celery_async_result = self.tasks.pop(task_instance_key, None)
if celery_async_result:
try:
Expand Down
2 changes: 2 additions & 0 deletions airflow/ti_deps/deps/trigger_rule_dep.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -48,6 +48,7 @@ class _UpstreamTIStates(NamedTuple):
success: int
skipped: int
failed: int
failed_in_queue: int
upstream_failed: int
removed: int
done: int
Expand Down Expand Up @@ -77,6 +78,7 @@ def calculate(cls, finished_upstreams: Iterator[TaskInstance]) -> _UpstreamTISta
success=counter.get(TaskInstanceState.SUCCESS, 0),
skipped=counter.get(TaskInstanceState.SKIPPED, 0),
failed=counter.get(TaskInstanceState.FAILED, 0),
failed_in_queue=counter.get(TaskInstanceState.FAILED_IN_QUEUE,0),
upstream_failed=counter.get(TaskInstanceState.UPSTREAM_FAILED, 0),
removed=counter.get(TaskInstanceState.REMOVED, 0),
done=sum(counter.values()),
Expand Down
6 changes: 5 additions & 1 deletion airflow/utils/state.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -52,6 +52,7 @@ class TaskInstanceState(str, Enum):
SUCCESS = "success" # Task completed
RESTARTING = "restarting" # External request to restart (e.g. cleared when running)
FAILED = "failed" # Task errored out
FAILED_IN_QUEUE = "failed_in_queue" # Task failed due to queue timeout
UP_FOR_RETRY = "up_for_retry" # Task failed but has retries left
UP_FOR_RESCHEDULE = "up_for_reschedule" # A waiting `reschedule` sensor
UPSTREAM_FAILED = "upstream_failed" # One or more upstream deps failed
Expand Down Expand Up @@ -95,6 +96,7 @@ class State:
SUCCESS = TaskInstanceState.SUCCESS
RUNNING = TaskInstanceState.RUNNING
FAILED = TaskInstanceState.FAILED
FAILED_IN_QUEUE = TaskInstanceState.FAILED_IN_QUEUE

# These are TaskState only
NONE = None
Expand Down Expand Up @@ -135,6 +137,7 @@ class State:
TaskInstanceState.SUCCESS: "green",
TaskInstanceState.RESTARTING: "violet",
TaskInstanceState.FAILED: "red",
TaskInstanceState.FAILED_IN_QUEUE: "coral",
TaskInstanceState.UP_FOR_RETRY: "gold",
TaskInstanceState.UP_FOR_RESCHEDULE: "turquoise",
TaskInstanceState.UPSTREAM_FAILED: "orange",
Expand All @@ -161,6 +164,7 @@ def color_fg(cls, state):
[
TaskInstanceState.SUCCESS,
TaskInstanceState.FAILED,
TaskInstanceState.FAILED_IN_QUEUE,
TaskInstanceState.SKIPPED,
TaskInstanceState.UPSTREAM_FAILED,
TaskInstanceState.REMOVED,
Expand Down Expand Up @@ -193,7 +197,7 @@ def color_fg(cls, state):
"""

failed_states: frozenset[TaskInstanceState] = frozenset(
[TaskInstanceState.FAILED, TaskInstanceState.UPSTREAM_FAILED]
[TaskInstanceState.FAILED, TaskInstanceState.UPSTREAM_FAILED, TaskInstanceState.FAILED_IN_QUEUE]
)
"""
A list of states indicating that a task or dag is a failed state.
Expand Down