Skip to content

Commit 2a6792d

Browse files
malthemalthe
authored
Fix race condition between triggerer and scheduler (#21316)
1 parent 2258e13 commit 2a6792d

2 files changed

Lines changed: 87 additions & 9 deletions

File tree

‎airflow/executors/base_executor.py‎

Lines changed: 31 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -17,7 +17,7 @@
1717
"""Base executor - this is the base class for all the implemented executors."""
1818
import sys
1919
from collections import OrderedDict
20-
from typing import Any, Dict, List, Optional, Set, Tuple, Union
20+
from typing import Any, Counter, Dict, List, Optional, Set, Tuple, Union
2121

2222
from airflow.configuration import conf
2323
from airflow.models.taskinstance import TaskInstance, TaskInstanceKey
@@ -29,6 +29,8 @@
2929

3030
NOT_STARTED_MESSAGE = "The executor should be started first!"
3131

32+
QUEUEING_ATTEMPTS = 5
33+
3234
# Command to execute - list of strings
3335
# the first element is always "airflow".
3436
# It should be result of TaskInstance.generate_command method.q
@@ -63,6 +65,7 @@ def __init__(self, parallelism: int = PARALLELISM):
6365
self.queued_tasks: OrderedDict[TaskInstanceKey, QueuedTaskInstanceType] = OrderedDict()
6466
self.running: Set[TaskInstanceKey] = set()
6567
self.event_buffer: Dict[TaskInstanceKey, EventBufferValueType] = {}
68+
self.attempts: Counter[TaskInstanceKey] = Counter()
6669

6770
def __repr__(self):
6871
return f"{self.__class__.__name__}(parallelism={self.parallelism})"
@@ -78,7 +81,7 @@ def queue_command(
7881
queue: Optional[str] = None,
7982
):
8083
"""Queues command to task"""
81-
if task_instance.key not in self.queued_tasks and task_instance.key not in self.running:
84+
if task_instance.key not in self.queued_tasks:
8285
self.log.info("Adding to queue: %s", command)
8386
self.queued_tasks[task_instance.key] = (command, priority, queue, task_instance)
8487
else:
@@ -183,9 +186,32 @@ def trigger_tasks(self, open_slots: int) -> None:
183186

184187
for _ in range(min((open_slots, len(self.queued_tasks)))):
185188
key, (command, _, queue, ti) = sorted_queue.pop(0)
186-
self.queued_tasks.pop(key)
187-
self.running.add(key)
188-
self.execute_async(key=key, command=command, queue=queue, executor_config=ti.executor_config)
189+
190+
# If a task makes it here but is still understood by the executor
191+
# to be running, it generally means that the task has been killed
192+
# externally and not yet been marked as failed.
193+
#
194+
# However, when a task is deferred, there is also a possibility of
195+
# a race condition where a task might be scheduled again during
196+
# trigger processing, even before we are able to register that the
197+
# deferred task has completed. In this case and for this reason,
198+
# we make a small number of attempts to see if the task has been
199+
# removed from the running set in the meantime.
200+
if key in self.running:
201+
attempt = self.attempts[key]
202+
if attempt < QUEUEING_ATTEMPTS - 1:
203+
self.attempts[key] = attempt + 1
204+
self.log.info("task %s is still running", key)
205+
continue
206+
207+
# We give up and remove the task from the queue.
208+
self.log.error("could not queue task %s (still running after %d attempts)", key, attempt)
209+
del self.attempts[key]
210+
del self.queued_tasks[key]
211+
else:
212+
del self.queued_tasks[key]
213+
self.running.add(key)
214+
self.execute_async(key=key, command=command, queue=queue, executor_config=ti.executor_config)
189215

190216
def change_state(self, key: TaskInstanceKey, state: str, info=None) -> None:
191217
"""

‎tests/executors/test_base_executor.py‎

Lines changed: 56 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -18,7 +18,9 @@
1818
from datetime import timedelta
1919
from unittest import mock
2020

21-
from airflow.executors.base_executor import BaseExecutor
21+
from pytest import mark
22+
23+
from airflow.executors.base_executor import QUEUEING_ATTEMPTS, BaseExecutor
2224
from airflow.models.baseoperator import BaseOperator
2325
from airflow.models.taskinstance import TaskInstanceKey
2426
from airflow.utils import timezone
@@ -57,7 +59,7 @@ def test_gauge_executor_metrics(mock_stats_gauge, mock_trigger_tasks, mock_sync)
5759
mock_stats_gauge.assert_has_calls(calls)
5860

5961

60-
def test_try_adopt_task_instances(dag_maker):
62+
def setup_dagrun(dag_maker):
6163
date = timezone.utcnow()
6264
start_date = date - timedelta(days=2)
6365

@@ -66,8 +68,58 @@ def test_try_adopt_task_instances(dag_maker):
6668
BaseOperator(task_id="task_2", start_date=start_date)
6769
BaseOperator(task_id="task_3", start_date=start_date)
6870

69-
dagrun = dag_maker.create_dagrun(execution_date=date)
70-
tis = dagrun.task_instances
71+
return dag_maker.create_dagrun(execution_date=date)
7172

73+
74+
def test_try_adopt_task_instances(dag_maker):
75+
dagrun = setup_dagrun(dag_maker)
76+
tis = dagrun.task_instances
7277
assert {ti.task_id for ti in tis} == {"task_1", "task_2", "task_3"}
7378
assert BaseExecutor().try_adopt_task_instances(tis) == tis
79+
80+
81+
def enqueue_tasks(executor, dagrun):
82+
for task_instance in dagrun.task_instances:
83+
executor.queue_command(task_instance, ["airflow"])
84+
85+
86+
def setup_trigger_tasks(dag_maker):
87+
dagrun = setup_dagrun(dag_maker)
88+
executor = BaseExecutor()
89+
executor.execute_async = mock.Mock()
90+
enqueue_tasks(executor, dagrun)
91+
return executor, dagrun
92+
93+
94+
@mark.parametrize("open_slots", [1, 2, 3])
95+
def test_trigger_queued_tasks(dag_maker, open_slots):
96+
executor, _ = setup_trigger_tasks(dag_maker)
97+
executor.trigger_tasks(open_slots)
98+
assert len(executor.execute_async.mock_calls) == open_slots
99+
100+
101+
@mark.parametrize("change_state_attempt", range(QUEUEING_ATTEMPTS + 2))
102+
def test_trigger_running_tasks(dag_maker, change_state_attempt):
103+
executor, dagrun = setup_trigger_tasks(dag_maker)
104+
open_slots = 100
105+
executor.trigger_tasks(open_slots)
106+
expected_calls = len(dagrun.task_instances) # initially `execute_async` called for each task
107+
assert len(executor.execute_async.mock_calls) == expected_calls
108+
109+
# All the tasks are now "running", so while we enqueue them again here,
110+
# they won't be executed again until the executor has been notified of a state change.
111+
enqueue_tasks(executor, dagrun)
112+
113+
for attempt in range(QUEUEING_ATTEMPTS + 2):
114+
# On the configured attempt, we notify the executor that the task has succeeded.
115+
if attempt == change_state_attempt:
116+
executor.change_state(dagrun.task_instances[0].key, State.SUCCESS)
117+
# If we have not exceeded QUEUEING_ATTEMPTS, we should expect an additional "execute" call
118+
if attempt < QUEUEING_ATTEMPTS:
119+
expected_calls += 1
120+
executor.trigger_tasks(open_slots)
121+
assert len(executor.execute_async.mock_calls) == expected_calls
122+
if change_state_attempt < QUEUEING_ATTEMPTS:
123+
assert len(executor.execute_async.mock_calls) == len(dagrun.task_instances) + 1
124+
else:
125+
assert len(executor.execute_async.mock_calls) == len(dagrun.task_instances)

0 commit comments

Comments
 (0)