diff --git a/airflow-core/newsfragments/74588.bugfix.rst b/airflow-core/newsfragments/74588.bugfix.rst new file mode 100644 index 0000000000000..fba91ca3988d9 --- /dev/null +++ b/airflow-core/newsfragments/74588.bugfix.rst @@ -0,0 +1 @@ +Fix ``LocalExecutor`` stalling all queued tasks when a worker is killed while holding the queue read lock (e.g. by the OOM killer). Each worker now has its own input queue. diff --git a/airflow-core/src/airflow/executors/local_executor.py b/airflow-core/src/airflow/executors/local_executor.py index 9e1f1f8bbf8dc..977b6eda4cf19 100644 --- a/airflow-core/src/airflow/executors/local_executor.py +++ b/airflow-core/src/airflow/executors/local_executor.py @@ -26,9 +26,7 @@ from __future__ import annotations import contextlib -import ctypes import multiprocessing -import multiprocessing.sharedctypes import os import sys from multiprocessing import Queue, SimpleQueue @@ -70,7 +68,6 @@ def _run_worker( logger_name: str, input: SimpleQueue[ExecutorWorkload | None], output: Queue[LocalResult], - unread_messages: multiprocessing.sharedctypes.Synchronized[int], team_conf, ): import signal @@ -97,10 +94,6 @@ def _run_worker( # Received poison pill, no more tasks to run return - # Decrement this as soon as we pick up a message off the queue - with unread_messages: - unread_messages.value -= 1 - key = LocalExecutor.get_workload_key(workload) output.put((os.getpid(), key, workload.running_state, None)) @@ -138,15 +131,21 @@ class LocalExecutor(BaseExecutor): {WorkloadType.EXECUTE_TASK, WorkloadType.EXECUTE_CALLBACK, WorkloadType.TEST_CONNECTION} ) - activity_queue: SimpleQueue[ExecutorWorkload | None] result_queue: SimpleQueue[LocalResult] workers: dict[int, multiprocessing.Process] - _unread_messages: multiprocessing.sharedctypes.Synchronized[int] + # One input queue per worker. A SimpleQueue read lock shared by all workers wedges the whole executor + # if the worker holding it is SIGKILLed (e.g. by the OOM killer), so no two workers share a queue. + _worker_queues: dict[int, SimpleQueue[ExecutorWorkload | None]] + # Workers handed a workload that have not reported START yet. Kept apart from _worker_tasks so a worker + # that dies before START still leaves the key for the scheduler's stuck-in-queued revoke, as before. + _dispatched: dict[int, WorkloadKey] def __init__(self, *args, **kwargs): super().__init__(*args, **kwargs) self._worker_tasks: dict[int, WorkloadKey] = {} self._dispatch_counts: dict[WorkloadKey, int] = {} + self._worker_queues: dict[int, SimpleQueue[ExecutorWorkload | None]] = {} + self._dispatched: dict[int, WorkloadKey] = {} # Resolve the start method at instantiation, not at import: the component CLI entry may have # set it via []/[core] mp_start_method before the executor is created. @@ -167,17 +166,13 @@ def start(self) -> None: # We delay opening these queues until the start method mostly for unit tests. ExecutorLoader caches # instances, so each test reusues the same instance! (i.e. test 1 runs, closes the queues, then test 2 # comes back and gets the same LocalExecutor instance, so we have to open new here.) - self.activity_queue = SimpleQueue() self.result_queue = SimpleQueue() self.workers = {} + self._worker_queues = {} + self._dispatched = {} self._worker_tasks.clear() self._dispatch_counts.clear() - # Mypy sees this value as `SynchronizedBase[c_uint]`, but that isn't the right runtime type behaviour - # (it looks like an int to python) - - self._unread_messages = multiprocessing.Value(ctypes.c_uint) - if self.is_mp_using_fork: # This creates the maximum number of worker processes (parallelism) at once # to minimize gc freeze/unfreeze cycles when using fork in multiprocessing @@ -185,50 +180,41 @@ def start(self) -> None: def _check_workers(self): self._read_results() - # Reap any dead workers - to_remove = set() - for pid, proc in self.workers.items(): + # Reap any dead workers. Each worker owns its input queue, so a dead one cannot block the others. + for pid, proc in list(self.workers.items()): if not proc.is_alive(): + # Drain what the worker wrote before dying (e.g. its START) so its key is attributed to it. self._read_results() - # A worker killed between dequeue and START has no entry here; the scheduler's - # stuck-in-queued handling releases that workload through revoke_task. + self._dispatched.pop(pid, None) + # A worker killed before START keeps its key in _dispatched; the scheduler's stuck-in-queued + # revoke releases it, as before. if (key := self._worker_tasks.pop(pid, None)) is not None: self._finish_dispatch(key, state_class_for_key(key).FAILED) - to_remove.add(pid) proc.close() + del self.workers[pid] + del self._worker_queues[pid] - if to_remove: - self.workers = {pid: proc for pid, proc in self.workers.items() if pid not in to_remove} - - with self._unread_messages: - num_outstanding = self._unread_messages.value - - if num_outstanding <= 0 or self.activity_queue.empty(): - # Nothing to do. Future enhancement if someone wants: shut down workers that have been idle for N - # seconds - return - - # If we're using spawn in multiprocessing (default on macOS now) to start tasks, this can get called a - # via `sync()` a few times before the spawned process actually starts picking up messages. Try not to - # create too much - if num_outstanding and len(self.workers) < self.parallelism: - if self.is_mp_using_fork: - # This creates the maximum number of worker processes at once - # to minimize gc freeze/unfreeze cycles when using fork in multiprocessing - self._spawn_workers_with_gc_freeze(self.parallelism - len(self.workers)) - else: - # This only creates one worker, which is fine as we call this directly after putting a message on - # activity_queue in execute_async when using spawn in multiprocessing - self._spawn_worker() - - def _spawn_worker(self): + def _pick_worker(self) -> int: + self._check_workers() + idle = next( + (pid for pid in self.workers if pid not in self._worker_tasks and pid not in self._dispatched), + None, + ) + if idle is not None: + return idle + if len(self.workers) < self.parallelism: + return self._spawn_worker() + # ponytail: all workers busy yet a slot is free (a stale result not read yet); queue behind one. + return next(iter(self.workers)) + + def _spawn_worker(self) -> int: + queue: SimpleQueue[ExecutorWorkload | None] = SimpleQueue() p = multiprocessing.Process( target=_run_worker, kwargs={ "logger_name": self.log.name, - "input": self.activity_queue, + "input": queue, "output": self.result_queue, - "unread_messages": self._unread_messages, "team_conf": self.conf, }, ) @@ -236,6 +222,8 @@ def _spawn_worker(self): if TYPE_CHECKING: assert p.pid # Since we've called start self.workers[p.pid] = p + self._worker_queues[p.pid] = queue + return p.pid def _spawn_workers_with_gc_freeze(self, spawn_number): """ @@ -265,14 +253,21 @@ def _read_results(self): try: while not self.result_queue.empty(): pid, key, state, exc = self.result_queue.get() - if pid not in self.workers or key not in self.running: + if pid not in self.workers: + continue + if key not in self.running: + # Late result for a key already finished or revoked; a terminal one frees its worker. + if state is not None and state != "running" and self._dispatched.get(pid) == key: + del self._dispatched[pid] continue if state is None or state == "running": self._worker_tasks[pid] = key + self._dispatched.pop(pid, None) if state is not None: self.change_state(key, state, remove_running=False) - elif self._worker_tasks.get(pid) == key: - del self._worker_tasks[pid] + elif self._worker_tasks.get(pid) == key or self._dispatched.get(pid) == key: + self._worker_tasks.pop(pid, None) + self._dispatched.pop(pid, None) self._finish_dispatch(key, state) except (OSError, EOFError): self.log.exception("Error reading from result queue") @@ -287,10 +282,10 @@ def end(self) -> None: # We can't tell which proc will pick which close message up, so we send all the messages, and then # wait on all the procs - for proc in self.workers.values(): - # Send the shutdown message once for each alive worker + for pid, proc in self.workers.items(): + # Send the shutdown message to each alive worker's own queue if proc.is_alive(): - self.activity_queue.put(None) + self._worker_queues[pid].put(None) # To prevent deadlock, we should consume results from result_queue while waiting for processes to join. # Otherwise, a worker blocked on putting results into a full result_queue pipe will never exit, @@ -313,7 +308,8 @@ def end(self) -> None: with contextlib.suppress(ValueError): proc.close() - self.activity_queue.close() + for queue in self._worker_queues.values(): + queue.close() self.result_queue.close() def terminate(self): @@ -337,15 +333,14 @@ def _terminate_worker_process(self, proc: multiprocessing.Process) -> None: def _process_workloads(self, workload_list): for workload in workload_list: key = self.get_workload_key(workload) - self.activity_queue.put(workload) + pid = self._pick_worker() + self._worker_queues[pid].put(workload) + self._dispatched[pid] = key removed = self.executor_queues[workload.type].pop(key, None) if not removed: raise KeyError(f"Workload {key} was not found in any queue") self.running.add(key) self._dispatch_counts[key] = self._dispatch_counts.get(key, 0) + 1 - with self._unread_messages: - self._unread_messages.value += len(workload_list) - self._check_workers() def _finish_dispatch(self, key: WorkloadKey, state: WorkloadState) -> None: # A resumed attempt reuses its key, so the previous dispatch can finish while the next one is live. @@ -359,6 +354,7 @@ def _forget_workload(self, key: WorkloadKey) -> None: self._worker_tasks = { pid: task_key for pid, task_key in self._worker_tasks.items() if task_key != key } + self._dispatched = {pid: task_key for pid, task_key in self._dispatched.items() if task_key != key} def change_state(self, key, state, info=None, remove_running=True) -> None: if remove_running: diff --git a/airflow-core/tests/unit/executors/test_local_executor.py b/airflow-core/tests/unit/executors/test_local_executor.py index 9f5d196d6cd94..ef75f5283d9d4 100644 --- a/airflow-core/tests/unit/executors/test_local_executor.py +++ b/airflow-core/tests/unit/executors/test_local_executor.py @@ -22,6 +22,7 @@ import os import signal import time +from multiprocessing import SimpleQueue from pathlib import Path from unittest import mock @@ -95,12 +96,10 @@ def _make_task_workload(): ) -def _write_large_results_to_queue(result_queue, activity_queue, unread_messages, result_count, payload_size): +def _write_large_results_to_queue(result_queue, worker_queue, result_count, payload_size): payload = RuntimeError("x" * payload_size) for _ in range(result_count): - workload = activity_queue.get() - with unread_messages: - unread_messages.value -= 1 + workload = worker_queue.get() key = LocalExecutor.get_workload_key(workload) result_queue.put((os.getpid(), key, workload.running_state, None)) result_queue.put((os.getpid(), key, State.SUCCESS, payload)) @@ -141,6 +140,7 @@ def _add_mock_worker(executor, mocker, pid): proc.pid = pid proc.is_alive.return_value = True executor.workers[pid] = proc + executor._worker_queues[pid] = SimpleQueue() return proc @@ -199,20 +199,10 @@ def test_executor_lazy_worker_spawning(self, mock_freeze, mock_unfreeze): executor.start() try: - # No workers should be pre-spawned + # No workers should be pre-spawned; they are created on dispatch assert len(executor.workers) == 0 mock_freeze.assert_not_called() mock_unfreeze.assert_not_called() - - # Simulate a queued message so _check_workers spawns one worker on demand - with executor._unread_messages: - executor._unread_messages.value = 1 - executor.activity_queue.put(None) # poison pill so the worker exits cleanly - executor._check_workers() - - assert len(executor.workers) == 1 - # gc.freeze is still not used for non-fork - mock_freeze.assert_not_called() finally: executor.end() @@ -291,7 +281,6 @@ def fake_run_workload(workload, **kwargs): # By that time Queues are already shutdown so we cannot check if they are empty assert len(executor.running) == 0 - assert executor._unread_messages.value == 0 for ti in success_tis: assert executor.event_buffer[TaskInstanceUuid(ti.id)][0] == State.SUCCESS @@ -342,29 +331,29 @@ def test_clean_stop_on_signal(self): def test_end_drains_results_while_joining_workers(self): executor = LocalExecutor(parallelism=1) - executor.activity_queue = mock.MagicMock() executor.result_queue = mock.MagicMock() proc = mock.MagicMock(spec=multiprocessing.Process) proc.is_alive.side_effect = [True, True, True, False] executor.workers = {1: proc} + executor._worker_queues = {1: mock.MagicMock()} with mock.patch.object(executor, "_read_results") as mock_read_results: executor.end() - executor.activity_queue.put.assert_called_once_with(None) + executor._worker_queues[1].put.assert_called_once_with(None) assert proc.join.call_args_list == [mock.call(timeout=0.05), mock.call(timeout=0.05)] assert mock_read_results.call_count == 3 proc.close.assert_called_once() - executor.activity_queue.close.assert_called_once() + executor._worker_queues[1].close.assert_called_once() executor.result_queue.close.assert_called_once() def test_end_terminates_workers_and_closes_resources_on_interrupt(self): executor = LocalExecutor(parallelism=1) - executor.activity_queue = mock.MagicMock() executor.result_queue = mock.MagicMock() proc = mock.MagicMock(spec=multiprocessing.Process) proc.is_alive.side_effect = [True, True] executor.workers = {1: proc} + executor._worker_queues = {1: mock.MagicMock()} with ( mock.patch.object(executor, "_read_results", side_effect=[KeyboardInterrupt, None]), @@ -375,7 +364,7 @@ def test_end_terminates_workers_and_closes_resources_on_interrupt(self): mock_terminate_worker_process.assert_called_once_with(proc) proc.close.assert_called_once() - executor.activity_queue.close.assert_called_once() + executor._worker_queues[1].close.assert_called_once() executor.result_queue.close.assert_called_once() def test_terminate_joins_worker_after_sigterm(self): @@ -417,22 +406,20 @@ def test_end_drains_result_queue_to_avoid_join_deadlock(self, mocker): result_count = 8 payload_size = 128 * 1024 submitted = [_make_task_workload() for _ in range(result_count)] - for workload in submitted: - executor.queue_workload(workload, session=mock.create_autospec(Session, instance=True)) - with mock.patch.object(executor, "_check_workers", autospec=True): - executor._process_workloads(submitted) + worker_queue = ctx.SimpleQueue() proc = ctx.Process( target=_write_large_results_to_queue, - args=( - executor.result_queue, - executor.activity_queue, - executor._unread_messages, - result_count, - payload_size, - ), + args=(executor.result_queue, worker_queue, result_count, payload_size), ) proc.start() + # Register the fake worker before dispatch so _process_workloads routes every workload to it + # instead of spawning a real worker that nothing would ever stop. executor.workers = {proc.pid: proc} + executor._worker_queues = {proc.pid: worker_queue} + for workload in submitted: + executor.queue_workload(workload, session=mock.create_autospec(Session, instance=True)) + with mock.patch.object(executor, "_check_workers", autospec=True): + executor._process_workloads(submitted) executor.end() @@ -441,7 +428,40 @@ def test_end_drains_result_queue_to_avoid_join_deadlock(self, mocker): assert all(state == State.SUCCESS for state, _ in executor.event_buffer.values()) assert not executor.running assert not executor._worker_tasks - assert executor._unread_messages.value == 0 + + @skip_non_fork_mp_start + @pytest.mark.execution_timeout(30) + def test_workload_runs_after_worker_holding_read_lock_is_killed(self, mocker): + """ + A worker SIGKILLed while blocked reading its input must not wedge the executor. + + Regression for #74575: with one shared SimpleQueue, the worker blocked in get() holds the queue's + read lock, so killing it left the lock held and every later workload stuck in queued. + """ + mocker.patch.object(BaseExecutor, "run_workload", autospec=True, return_value=0) + executor = LocalExecutor(parallelism=2) + executor.start() + try: + # Let the workers come up and block on their input. Kill them all so whichever one holds the + # shared read lock on the old queue is certainly dead (the OOM-killer case). + time.sleep(0.5) + for pid, proc in list(executor.workers.items()): + os.kill(pid, signal.SIGKILL) + proc.join(timeout=5) + + workload = _make_task_workload() + key = executor.get_task_key(workload.ti) + executor.queue_workload(workload, session=mock.create_autospec(Session, instance=True)) + + deadline = time.monotonic() + 20 + while key not in executor.event_buffer and time.monotonic() < deadline: + executor.heartbeat() + time.sleep(0.05) + + assert key in executor.event_buffer + assert executor.event_buffer[key][0] == State.SUCCESS + finally: + executor.end() @pytest.mark.parametrize( ("conf_values", "expected_server"), @@ -568,7 +588,6 @@ def test_multiple_team_executors_isolation(self): try: # Verify each executor has its own queues - assert team_a_executor.activity_queue is not team_b_executor.activity_queue assert team_a_executor.result_queue is not team_b_executor.result_queue # Verify each executor has its own workers dict @@ -583,9 +602,6 @@ def test_multiple_team_executors_isolation(self): assert len(team_a_executor.workers) == 0 assert len(team_b_executor.workers) == 0 - # Verify each executor has its own unread_messages counter - assert team_a_executor._unread_messages is not team_b_executor._unread_messages - # Verify each has correct team config assert team_a_executor.conf.team_name == "team_a" assert team_b_executor.conf.team_name == "team_b" @@ -647,7 +663,6 @@ def test_running_limits_later_heartbeats_and_reports_metrics( executor.heartbeat() assert executor.running == {executor.get_task_key(first.ti)} - assert executor._unread_messages.value == 1 assert second in executor.executor_queues[second.type].values() assert executor.has_task(first.ti) metrics = {call.args[0]: call.kwargs["value"] for call in gauge.call_args_list[-3:]} @@ -688,7 +703,7 @@ def test_result_uses_original_submitted_uuid_after_dto_changes(self, local_execu key, coordinates = executor.get_task_key(workload.ti), workload.ti.key executor.queue_workload(workload, session=mock.create_autospec(Session, instance=True)) executor.heartbeat() - submitted = executor.activity_queue.get() + submitted = executor._worker_queues[proc.pid].get() workload.ti.id = uuid7() workload.ti.try_number += 1 assert executor.get_workload_key(submitted) == key @@ -708,8 +723,7 @@ def test_reaper_drains_start_sent_after_initial_poll(self, local_executor_with_m key = executor.get_task_key(workload.ti) executor.queue_workload(workload, session=mock.create_autospec(Session, instance=True)) executor.heartbeat() - executor.activity_queue.get() - executor._unread_messages.value = 0 + executor._worker_queues[proc.pid].get() def died_after_start(): executor.result_queue.put((proc.pid, key, None, None)) @@ -729,8 +743,7 @@ def test_revoke_task_releases_slot_of_workload_lost_before_start(self, local_exe key = executor.get_task_key(workload.ti) executor.queue_workload(workload, session=mock.create_autospec(Session, instance=True)) executor.heartbeat() - executor.activity_queue.get() - executor._unread_messages.value = 0 + executor._worker_queues[proc.pid].get() proc.is_alive.return_value = False executor.sync() assert not executor.workers diff --git a/airflow-core/tests/unit/executors/test_local_executor_check_workers.py b/airflow-core/tests/unit/executors/test_local_executor_check_workers.py index 98538f7b970d3..b4ea40878c0e9 100644 --- a/airflow-core/tests/unit/executors/test_local_executor_check_workers.py +++ b/airflow-core/tests/unit/executors/test_local_executor_check_workers.py @@ -24,116 +24,48 @@ from airflow.executors.local_executor import LocalExecutor -@pytest.fixture(autouse=True) -def setup_executor(monkeypatch): +@pytest.fixture +def executor(): executor = LocalExecutor(parallelism=2) executor.workers = {} - executor._unread_messages = MagicMock() - executor.activity_queue = MagicMock() + executor._worker_queues = {} executor.result_queue = SimpleQueue() - monkeypatch.setattr(executor, "_spawn_worker", MagicMock()) yield executor executor.result_queue.close() -def test_no_workers_on_no_work(setup_executor): - executor = setup_executor - executor._unread_messages.value = 0 - executor.activity_queue.empty.return_value = True - executor._check_workers() - executor._spawn_worker.assert_not_called() - assert executor.workers == {} +def _worker(executor, pid, alive): + proc = MagicMock() + proc.is_alive.return_value = alive + executor.workers[pid] = proc + executor._worker_queues[pid] = MagicMock() + return proc -def test_all_workers_alive(setup_executor): - executor = setup_executor - proc1 = MagicMock() - proc1.is_alive.return_value = True - proc2 = MagicMock() - proc2.is_alive.return_value = True - executor.workers = {1: proc1, 2: proc2} - executor._unread_messages.value = 0 - executor.activity_queue.empty.return_value = True +def test_alive_workers_are_kept(executor): + proc1 = _worker(executor, 1, alive=True) + proc2 = _worker(executor, 2, alive=True) executor._check_workers() proc1.close.assert_not_called() proc2.close.assert_not_called() - assert len(executor.workers) == 2 - - -def test_some_workers_dead(setup_executor): - executor = setup_executor - proc1 = MagicMock() - proc1.is_alive.return_value = False - proc2 = MagicMock() - proc2.is_alive.return_value = True - executor.workers = {1: proc1, 2: proc2} - executor._unread_messages.value = 0 - executor.activity_queue.empty.return_value = True - executor._check_workers() - proc1.close.assert_called_once() - proc2.close.assert_not_called() - assert executor.workers == {2: proc2} - - -def test_all_workers_dead(setup_executor): - executor = setup_executor - proc1 = MagicMock() - proc1.is_alive.return_value = False - proc2 = MagicMock() - proc2.is_alive.return_value = False - executor.workers = {1: proc1, 2: proc2} - executor._unread_messages.value = 0 - executor.activity_queue.empty.return_value = True - executor._check_workers() - proc1.close.assert_called_once() - proc2.close.assert_called_once() - assert executor.workers == {} + assert set(executor.workers) == {1, 2} + assert set(executor._worker_queues) == {1, 2} -def test_outstanding_messages_and_empty_queue(setup_executor): - executor = setup_executor - executor._unread_messages.value = 1 - executor.activity_queue.empty.return_value = True +def test_dead_worker_is_reaped_with_its_queue(executor): + dead = _worker(executor, 1, alive=False) + alive = _worker(executor, 2, alive=True) executor._check_workers() - executor._spawn_worker.assert_not_called() - - -def test_spawn_worker_when_needed(setup_executor): - executor = setup_executor - executor._unread_messages.value = 1 - executor.activity_queue.empty.return_value = False - executor.workers = {} - executor._check_workers() - executor._spawn_worker.assert_called() - - -def test_no_spawn_if_parallelism_reached(setup_executor): - executor = setup_executor - executor._unread_messages.value = 2 - executor.activity_queue.empty.return_value = False - proc1 = MagicMock() - proc1.is_alive.return_value = True - proc2 = MagicMock() - proc2.is_alive.return_value = True - executor.workers = {1: proc1, 2: proc2} - executor._check_workers() - executor._spawn_worker.assert_not_called() - + dead.close.assert_called_once() + alive.close.assert_not_called() + assert executor.workers == {2: alive} + assert set(executor._worker_queues) == {2} -def test_spawn_worker_when_we_have_parallelism_left(setup_executor): - executor = setup_executor - # Simulate 4 running workers - running_workers = {} - for i in range(4): - proc = MagicMock() - proc.is_alive.return_value = True - running_workers[i] = proc - executor.workers = running_workers - executor.parallelism = 5 # Allow more workers if needed - # Simulate 4 pending tasks (equal to running workers) - executor._unread_messages.value = 4 - executor.activity_queue.empty.return_value = False - executor._spawn_worker.reset_mock() +def test_dead_worker_before_start_leaves_key_for_revoke(executor): + _worker(executor, 1, alive=False) + executor._dispatched[1] = "key" + executor._finish_dispatch = MagicMock() executor._check_workers() - executor._spawn_worker.assert_called() + assert 1 not in executor._dispatched + executor._finish_dispatch.assert_not_called() # scheduler's stuck-in-queued revoke releases it