Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
1 change: 1 addition & 0 deletions airflow-core/src/airflow/executors/local_executor.py
Original file line number Diff line number Diff line change
Expand Up @@ -310,6 +310,7 @@ def _terminate_worker_process(self, proc: multiprocessing.Process) -> None:

def _process_workloads(self, workload_list):
for workload in workload_list:
self._read_results()
self.activity_queue.put(workload)
# A valid workload will exist in exactly one of these dicts.
# One pop will succeed, the others will return None gracefully.
Expand Down
61 changes: 61 additions & 0 deletions airflow-core/tests/unit/executors/test_local_executor.py
Original file line number Diff line number Diff line change
Expand Up @@ -369,6 +369,67 @@ def test_end_drains_result_queue_to_avoid_join_deadlock(self):

assert len(executor.event_buffer) == result_count

def test_process_workloads_drains_results_before_put(self):
executor = LocalExecutor(parallelism=1)
executor.activity_queue = mock.MagicMock()
executor.result_queue = mock.MagicMock()
executor._unread_messages = mock.MagicMock()
executor._check_workers = mock.MagicMock()
workload = mock.MagicMock()
workload.key = TaskInstanceKey("test_dag", "test_task", "test_run")
executor.queued_tasks = {workload.key: workload}

call_order = []
orig_read_results = executor._read_results
orig_put = executor.activity_queue.put

def tracking_read_results():
call_order.append("_read_results")
return orig_read_results()

def tracking_put(*args, **kwargs):
call_order.append("put")
return orig_put(*args, **kwargs)

executor._read_results = tracking_read_results
executor.activity_queue.put = tracking_put

executor._process_workloads([workload])

assert call_order == ["_read_results", "put"]

def test_process_workloads_drains_results_for_multiple_workloads(self):
executor = LocalExecutor(parallelism=1)
executor.activity_queue = mock.MagicMock()
executor.result_queue = mock.MagicMock()
executor._unread_messages = mock.MagicMock()
executor._check_workers = mock.MagicMock()
workloads_list = []
for i in range(3):
workload = mock.MagicMock()
workload.key = TaskInstanceKey("test_dag", f"test_task_{i}", "test_run")
workloads_list.append(workload)
executor.queued_tasks[workload.key] = workload

call_order = []
orig_read_results = executor._read_results
orig_put = executor.activity_queue.put

def tracking_read_results():
call_order.append("_read_results")
return orig_read_results()

def tracking_put(*args, **kwargs):
call_order.append("put")
return orig_put(*args, **kwargs)

executor._read_results = tracking_read_results
executor.activity_queue.put = tracking_put

executor._process_workloads(workloads_list)

assert call_order == ["_read_results", "put", "_read_results", "put", "_read_results", "put"]

@pytest.mark.parametrize(
("conf_values", "expected_server"),
[
Expand Down