diff --git a/airflow-core/src/airflow/executors/local_executor.py b/airflow-core/src/airflow/executors/local_executor.py index 24ab737f5ab20..6bc5d20cbfd4c 100644 --- a/airflow-core/src/airflow/executors/local_executor.py +++ b/airflow-core/src/airflow/executors/local_executor.py @@ -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. diff --git a/airflow-core/tests/unit/executors/test_local_executor.py b/airflow-core/tests/unit/executors/test_local_executor.py index 79987f9ab1cb4..8a52eef022a46 100644 --- a/airflow-core/tests/unit/executors/test_local_executor.py +++ b/airflow-core/tests/unit/executors/test_local_executor.py @@ -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"), [