This is an automated email from the ASF dual-hosted git repository.
eladkal pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/airflow.git
The following commit(s) were added to refs/heads/main by this push:
new 40832901c59 Improve LocalExecutor bookkeeping: correctly add tasks to
the `running` list (#74142)
40832901c59 is described below
commit 40832901c5928742369844bd0e00a1a9583b8daf
Author: Ash Berlin-Taylor <[email protected]>
AuthorDate: Mon Oct 5 06:55:48 2026 +0100
Improve LocalExecutor bookkeeping: correctly add tasks to the `running`
list (#74142)
* Improve LocalExecutor bookkeeping: correctly add tasks to the `running`
list
Since #73916 landed, BaseExecutor now keeps a UUID-to-coordinates map for
each
attempt and drops an entry once the key is no longer queued, running or in
the
event buffer. The issue was that LocalExecutor never added dispatched work
to
`running`, so the entry was dropped while the task was still executing. The
scheduler then received the final event with no coordinates ("Received
executor
event with state success for task instance <uuid> (coordinates=None)").
This PR fixes that, and addresses a few gotchas that could, in edge cases
around callbacks or post-task finalization, lead to a forever-dead/locked
worker slot.
Holding the key in `running` until the worker finishes needs LocalExecutor
to know when that is, and workers share one activity queue so the parent
cannot tell which worker took which workload. Results now carry the
worker's pid and START tells the parent which pid owns the key. That lets
it:
- release the slot only on a terminal result from the owning worker,
ignoring results from unknown pids or pids that don't own the key
- fail the workload when its worker dies, instead of leaving it in
`running` until the scheduler's heartbeat timeout
A resumed task reuses its key, so a per-key dispatch count stops the old
run's terminal result from releasing the new run's slot.
* fixup! Improve LocalExecutor bookkeeping: correctly add tasks to the
`running` list
* fixup! Improve LocalExecutor bookkeeping: correctly add tasks to the
`running` list
---
.../src/airflow/executors/local_executor.py | 75 +++-
.../tests/unit/executors/test_local_executor.py | 455 ++++++++++++++++++++-
.../executors/test_local_executor_check_workers.py | 5 +-
airflow-core/tests/unit/jobs/test_scheduler_job.py | 9 +-
4 files changed, 520 insertions(+), 24 deletions(-)
diff --git a/airflow-core/src/airflow/executors/local_executor.py
b/airflow-core/src/airflow/executors/local_executor.py
index 74c92a4f31a..9e1f1f8bbf8 100644
--- a/airflow-core/src/airflow/executors/local_executor.py
+++ b/airflow-core/src/airflow/executors/local_executor.py
@@ -38,6 +38,7 @@ import structlog
from airflow.executors.base_executor import BaseExecutor,
get_execution_api_server_url
from airflow.executors.workloads import WorkloadType
+from airflow.executors.workloads.types import state_class_for_key
# add logger to parameter of setproctitle to support logging
if sys.platform == "darwin":
@@ -49,7 +50,10 @@ else:
if TYPE_CHECKING:
from airflow.executors.workloads import ExecutorWorkload
- from airflow.executors.workloads.types import WorkloadResultType
+ from airflow.executors.workloads.types import WorkloadKey, WorkloadState
+ from airflow.models.taskinstance import TaskInstance
+
+ LocalResult = tuple[int, WorkloadKey, WorkloadState | None, Exception |
None]
def _get_executor_process_title_prefix(team_name: str | None) -> str:
@@ -65,7 +69,7 @@ def _get_executor_process_title_prefix(team_name: str | None)
-> str:
def _run_worker(
logger_name: str,
input: SimpleQueue[ExecutorWorkload | None],
- output: Queue[WorkloadResultType],
+ output: Queue[LocalResult],
unread_messages: multiprocessing.sharedctypes.Synchronized[int],
team_conf,
):
@@ -98,8 +102,7 @@ def _run_worker(
unread_messages.value -= 1
key = LocalExecutor.get_workload_key(workload)
- if workload.running_state is not None:
- output.put((key, workload.running_state, None))
+ output.put((os.getpid(), key, workload.running_state, None))
try:
BaseExecutor.run_workload(
@@ -108,10 +111,10 @@ def _run_worker(
proctitle=f"{_get_executor_process_title_prefix(team_conf.team_name)}
{workload.display_name}",
subprocess_logs_to_stdout=True,
)
- output.put((key, workload.success_state, None))
+ output.put((os.getpid(), key, workload.success_state, None))
except Exception as e:
log.exception("Workload execution failed.",
workload_type=type(workload).__name__)
- output.put((key, workload.failure_state, e))
+ output.put((os.getpid(), key, workload.failure_state, e))
class LocalExecutor(BaseExecutor):
@@ -136,12 +139,14 @@ class LocalExecutor(BaseExecutor):
)
activity_queue: SimpleQueue[ExecutorWorkload | None]
- result_queue: SimpleQueue[WorkloadResultType]
+ result_queue: SimpleQueue[LocalResult]
workers: dict[int, multiprocessing.Process]
_unread_messages: multiprocessing.sharedctypes.Synchronized[int]
def __init__(self, *args, **kwargs):
super().__init__(*args, **kwargs)
+ self._worker_tasks: dict[int, WorkloadKey] = {}
+ self._dispatch_counts: dict[WorkloadKey, int] = {}
# Resolve the start method at instantiation, not at import: the
component CLI entry may have
# set it via [<component>]/[core] mp_start_method before the executor
is created.
@@ -165,6 +170,8 @@ class LocalExecutor(BaseExecutor):
self.activity_queue = SimpleQueue()
self.result_queue = SimpleQueue()
self.workers = {}
+ 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)
@@ -177,10 +184,16 @@ class LocalExecutor(BaseExecutor):
self._spawn_workers_with_gc_freeze(self.parallelism)
def _check_workers(self):
+ self._read_results()
# Reap any dead workers
to_remove = set()
for pid, proc in self.workers.items():
if not proc.is_alive():
+ 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.
+ 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()
@@ -246,14 +259,21 @@ class LocalExecutor(BaseExecutor):
def sync(self) -> None:
"""Sync will get called periodically by the heartbeat method."""
- self._read_results()
self._check_workers()
def _read_results(self):
try:
while not self.result_queue.empty():
- key, state, exc = self.result_queue.get()
- self.change_state(key, state)
+ pid, key, state, exc = self.result_queue.get()
+ if pid not in self.workers or key not in self.running:
+ continue
+ if state is None or state == "running":
+ self._worker_tasks[pid] = key
+ 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]
+ self._finish_dispatch(key, state)
except (OSError, EOFError):
self.log.exception("Error reading from result queue")
@@ -316,11 +336,44 @@ class LocalExecutor(BaseExecutor):
def _process_workloads(self, workload_list):
for workload in workload_list:
- self.activity_queue.put(workload)
key = self.get_workload_key(workload)
+ self.activity_queue.put(workload)
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.
+ remaining = self._dispatch_counts.pop(key, 1) - 1
+ if remaining > 0:
+ self._dispatch_counts[key] = remaining
+ super().change_state(key, state, remove_running=remaining <= 0)
+
+ def _forget_workload(self, key: WorkloadKey) -> None:
+ self._dispatch_counts.pop(key, None)
+ self._worker_tasks = {
+ pid: task_key for pid, task_key in self._worker_tasks.items() if
task_key != key
+ }
+
+ def change_state(self, key, state, info=None, remove_running=True) -> None:
+ if remove_running:
+ self._forget_workload(key)
+ super().change_state(key, state, info=info,
remove_running=remove_running)
+
+ def fail_connection_test(self, key) -> None:
+ self._forget_workload(key)
+ super().fail_connection_test(key)
+
+ def revoke_task(self, *, ti: TaskInstance) -> None:
+ key = self.get_task_key(ti)
+ self.executor_queues[WorkloadType.EXECUTE_TASK].pop(key, None)
+ for pid, task_key in self._worker_tasks.items():
+ if task_key == key:
+ self._terminate_worker_process(self.workers[pid])
+ self._forget_workload(key)
+ self.running.discard(key)
diff --git a/airflow-core/tests/unit/executors/test_local_executor.py
b/airflow-core/tests/unit/executors/test_local_executor.py
index b89a7841e16..b5213b3c640 100644
--- a/airflow-core/tests/unit/executors/test_local_executor.py
+++ b/airflow-core/tests/unit/executors/test_local_executor.py
@@ -20,6 +20,8 @@ from __future__ import annotations
import gc
import multiprocessing
import os
+import signal
+import time
from pathlib import Path
from unittest import mock
@@ -27,10 +29,11 @@ import pytest
from kgb import spy_on
from uuid6 import uuid7
+import airflow.executors.local_executor as local_executor_module
from airflow._shared.timezones import timezone
from airflow.executors import workloads
from airflow.executors.base_executor import BaseExecutor, ExecutorConf,
get_execution_api_server_url
-from airflow.executors.local_executor import LocalExecutor
+from airflow.executors.local_executor import LocalExecutor, _run_worker
from airflow.executors.workloads import WorkloadType
from airflow.executors.workloads.base import BundleInfo
from airflow.executors.workloads.callback import CallbackDTO
@@ -92,11 +95,65 @@ def _make_task_workload():
)
-def _write_large_results_to_queue(result_queue, result_count, payload_size):
+def _write_large_results_to_queue(result_queue, activity_queue,
unread_messages, result_count, payload_size):
payload = RuntimeError("x" * payload_size)
for _ in range(result_count):
- key = uuid7()
- result_queue.put((key, State.SUCCESS, payload))
+ workload = activity_queue.get()
+ with unread_messages:
+ unread_messages.value -= 1
+ 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))
+
+
+def _make_workload(kind):
+ if kind == "task":
+ return _make_task_workload()
+ if kind == "callback":
+ return workloads.ExecuteCallback(
+ callback=CallbackDTO(
+ id=uuid7(),
+ fetch_method=CallbackFetchMethod.IMPORT_PATH,
+ data={"path": "test.func", "kwargs": {}},
+ ),
+ dag_rel_path="test.py",
+ bundle_info=BundleInfo(name="bundle"),
+ token="token",
+ log_path=None,
+ )
+ return workloads.TestConnection(
+ connection_test_id=uuid7(), connection_id="test", timeout=10,
token="token"
+ )
+
+
+def _hold_workload(workload, **kwargs):
+ Path(workload.token).touch()
+ signal.pause()
+
+
+def _run_blocking_worker(**kwargs):
+ with mock.patch.object(BaseExecutor, "run_workload", autospec=True,
side_effect=_hold_workload):
+ _run_worker(**kwargs)
+
+
+def _add_mock_worker(executor, mocker, pid):
+ proc = mocker.create_autospec(multiprocessing.Process, instance=True)
+ proc.pid = pid
+ proc.is_alive.return_value = True
+ executor.workers[pid] = proc
+ return proc
+
+
[email protected]
+def local_executor_with_mock_worker(mocker):
+ mocker.patch.object(LocalExecutor, "_spawn_workers_with_gc_freeze",
autospec=True)
+ mocker.patch.object(LocalExecutor, "_spawn_worker", autospec=True)
+ executor = LocalExecutor(parallelism=1)
+ executor.start()
+ proc = _add_mock_worker(executor, mocker, 12345)
+ yield executor, proc
+ executor.workers.clear()
+ executor.end()
class TestLocalExecutor:
@@ -347,7 +404,7 @@ class TestLocalExecutor:
assert proc.join.call_args_list == [mock.call(timeout=0.2),
mock.call(timeout=0.2)]
@pytest.mark.execution_timeout(10)
- def test_end_drains_result_queue_to_avoid_join_deadlock(self):
+ def test_end_drains_result_queue_to_avoid_join_deadlock(self, mocker):
# Pin the worker to "fork": the drain logic under test is
start-method-agnostic, but under the
# "forkserver" default (Python 3.14+ on Linux) each spawned worker
re-imports the whole airflow
# stack before it can write a result, which intermittently exceeds the
execution_timeout and
@@ -355,13 +412,24 @@ class TestLocalExecutor:
# immediately and reliably reproduces the full-result_queue scenario
this test guards.
ctx = multiprocessing.get_context("fork")
executor = LocalExecutor(parallelism=1)
- executor.activity_queue = ctx.SimpleQueue()
- executor.result_queue = ctx.SimpleQueue()
+ mocker.patch.object(executor, "_spawn_workers_with_gc_freeze",
autospec=True)
+ executor.start()
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)
proc = ctx.Process(
target=_write_large_results_to_queue,
- args=(executor.result_queue, result_count, payload_size),
+ args=(
+ executor.result_queue,
+ executor.activity_queue,
+ executor._unread_messages,
+ result_count,
+ payload_size,
+ ),
)
proc.start()
executor.workers = {proc.pid: proc}
@@ -369,6 +437,11 @@ class TestLocalExecutor:
executor.end()
assert len(executor.event_buffer) == result_count
+ assert set(executor.event_buffer) ==
{executor.get_task_key(workload.ti) for workload in submitted}
+ 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
@pytest.mark.parametrize(
("conf_values", "expected_server"),
@@ -508,6 +581,372 @@ class TestLocalExecutor:
executor.end()
+class TestLocalExecutorBookkeeping:
+ def test_dispatch_keeps_task_visible_without_a_worker_result(self, mocker):
+ mocker.patch.object(LocalExecutor, "_spawn_workers_with_gc_freeze",
autospec=True)
+ mocker.patch.object(LocalExecutor, "_check_workers", autospec=True)
+ executor = LocalExecutor(parallelism=1)
+ executor.start()
+ workload = _make_task_workload()
+ key = executor.get_task_key(workload.ti)
+ executor.queue_workload(workload,
session=mock.create_autospec(Session, instance=True))
+ try:
+ executor.heartbeat()
+ executor._drain_events_with_task_ids()
+
+ assert key in executor.running
+ assert executor.has_task(workload.ti)
+ assert executor.slots_available == 0
+ assert executor._task_coordinates[key] == workload.ti.key
+ finally:
+ executor.end()
+
+ def test_running_limits_later_heartbeats_and_reports_metrics(
+ self, local_executor_with_mock_worker, mocker
+ ):
+ executor, proc = local_executor_with_mock_worker
+ gauge = mocker.patch("airflow.executors.base_executor.stats.gauge",
autospec=True)
+ first, second = _make_task_workload(), _make_task_workload()
+ executor.queue_workload(first, session=mock.create_autospec(Session,
instance=True))
+ executor.heartbeat()
+ assert executor.slots_available == 0
+ executor.queue_workload(second, session=mock.create_autospec(Session,
instance=True))
+
+ 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:]}
+ assert metrics == {"executor.open_slots": 0, "executor.queued_tasks":
1, "executor.running_tasks": 1}
+
+ @pytest.mark.parametrize("kind", ["task", "callback", "connection"])
+ @pytest.mark.parametrize("succeeded", [True, False])
+ def test_start_retains_slot_and_terminal_clears_pid(
+ self, kind, succeeded, local_executor_with_mock_worker
+ ):
+ executor, proc = local_executor_with_mock_worker
+ workload = _make_workload(kind)
+ key = executor.get_workload_key(workload)
+ executor.queue_workload(workload,
session=mock.create_autospec(Session, instance=True))
+ executor.heartbeat()
+
+ executor.result_queue.put((proc.pid, key, workload.running_state,
None))
+ executor.sync()
+
+ assert executor._worker_tasks == {proc.pid: key}
+ assert key in executor.running
+ assert executor.slots_available == 0
+ if workload.running_state is None:
+ assert key not in executor.event_buffer
+ else:
+ assert executor.event_buffer[key] == (workload.running_state, None)
+ terminal = workload.success_state if succeeded else
workload.failure_state
+ executor.result_queue.put((proc.pid, key, terminal, None))
+ executor.sync()
+ assert executor.event_buffer[key] == (terminal, None)
+ assert not executor._worker_tasks
+ assert not executor._dispatch_counts
+ assert executor.slots_available == 1
+
+ def test_result_uses_original_submitted_uuid_after_dto_changes(self,
local_executor_with_mock_worker):
+ executor, proc = local_executor_with_mock_worker
+ workload = _make_task_workload()
+ 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()
+ workload.ti.id = uuid7()
+ workload.ti.try_number += 1
+ assert executor.get_workload_key(submitted) == key
+ executor.result_queue.put((proc.pid, key, None, None))
+ executor.result_queue.put((proc.pid, key, workload.success_state,
None))
+
+ executor.sync()
+ events, captured = executor._drain_events_with_task_ids()
+
+ assert events == {key: (workload.success_state, None)}
+ assert captured == {key: coordinates}
+ assert executor.slots_available == 1
+
+ def test_reaper_drains_start_sent_after_initial_poll(self,
local_executor_with_mock_worker):
+ executor, proc = local_executor_with_mock_worker
+ workload = _make_task_workload()
+ 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
+
+ def died_after_start():
+ executor.result_queue.put((proc.pid, key, None, None))
+ return False
+
+ proc.is_alive.side_effect = died_after_start
+ executor.sync()
+
+ assert executor.event_buffer[key] == (workload.failure_state, None)
+ assert not executor.running
+ assert not executor._worker_tasks
+ proc.close.assert_called_once()
+
+ def test_revoke_task_releases_slot_of_workload_lost_before_start(self,
local_executor_with_mock_worker):
+ executor, proc = local_executor_with_mock_worker
+ workload = _make_task_workload()
+ 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
+ proc.is_alive.return_value = False
+ executor.sync()
+ assert not executor.workers
+ assert key in executor.running
+
+ executor.revoke_task(ti=workload.ti)
+
+ assert not executor.running
+ assert not executor._dispatch_counts
+ assert executor.event_buffer == {}
+ assert executor.slots_available == 1
+
+ @pytest.mark.parametrize(
+ ("stage", "worker_terminated"),
+ [("queued", False), ("dispatched", False), ("started", True)],
+ )
+ def test_revoke_task_clears_workload_at_every_stage(
+ self, stage, worker_terminated, local_executor_with_mock_worker
+ ):
+ executor, proc = local_executor_with_mock_worker
+ workload = _make_task_workload()
+ key = executor.get_task_key(workload.ti)
+ executor.queue_workload(workload,
session=mock.create_autospec(Session, instance=True))
+ if stage != "queued":
+ executor.heartbeat()
+ if stage == "started":
+ executor.result_queue.put((proc.pid, key, None, None))
+ executor.sync()
+ assert executor._worker_tasks == {proc.pid: key}
+
+ executor.revoke_task(ti=workload.ti)
+
+ assert not executor.executor_queues[workload.type]
+ assert not executor.running
+ assert not executor._worker_tasks
+ assert not executor._dispatch_counts
+ assert executor.event_buffer == {}
+ assert proc.terminate.called is worker_terminated
+
+ @pytest.mark.parametrize("kind", ["task", "connection"])
+ def test_external_timeout_clears_pid_and_rejects_late_results(
+ self, kind, local_executor_with_mock_worker
+ ):
+ executor, proc = local_executor_with_mock_worker
+ workload = _make_workload(kind)
+ key = executor.get_workload_key(workload)
+ executor.queue_workload(workload,
session=mock.create_autospec(Session, instance=True))
+ executor.heartbeat()
+ executor.result_queue.put((proc.pid, key, workload.running_state,
None))
+ executor.sync()
+
+ if kind == "connection":
+ executor.fail_connection_test(key)
+ else:
+ executor.change_state(key, workload.failure_state,
remove_running=True)
+ executor.result_queue.put((proc.pid, key, workload.success_state,
None))
+ executor.sync()
+
+ assert not executor._worker_tasks
+ assert executor.slots_available == 1
+ expected_state = workload.running_state if kind == "connection" else
workload.failure_state
+ expected = {key: (expected_state, None)}
+ assert executor.event_buffer == expected
+ assert executor.workers[proc.pid] is proc
+
+ def test_one_worker_runs_workloads_back_to_back(self,
local_executor_with_mock_worker):
+ executor, proc = local_executor_with_mock_worker
+ first, second = _make_task_workload(), _make_task_workload()
+ first_key, second_key = executor.get_task_key(first.ti),
executor.get_task_key(second.ti)
+ for workload, key in ((first, first_key), (second, second_key)):
+ executor.queue_workload(workload,
session=mock.create_autospec(Session, instance=True))
+ executor.heartbeat()
+ executor.result_queue.put((proc.pid, key, None, None))
+ executor.sync()
+ assert executor._worker_tasks == {proc.pid: key}
+ executor.result_queue.put((proc.pid, key, workload.success_state,
None))
+ executor.sync()
+ assert not executor._worker_tasks
+ assert executor.event_buffer == {
+ first_key: (first.success_state, None),
+ second_key: (second.success_state, None),
+ }
+ assert executor.slots_available == 1
+
+ def test_redispatched_key_stays_tracked_after_previous_dispatch_finishes(
+ self, local_executor_with_mock_worker, mocker
+ ):
+ executor, first_proc = local_executor_with_mock_worker
+ second_proc = _add_mock_worker(executor, mocker, 54321)
+ workload = _make_task_workload()
+ key = executor.get_task_key(workload.ti)
+ session = mock.create_autospec(Session, instance=True)
+ executor.queue_workload(workload, session=session)
+ executor.heartbeat()
+ executor.queue_workload(workload, session=session)
+ executor._process_workloads([workload])
+ executor.result_queue.put((first_proc.pid, key, None, None))
+ executor.result_queue.put((first_proc.pid, key,
workload.success_state, None))
+ executor.result_queue.put((second_proc.pid, key, None, None))
+
+ executor.sync()
+
+ assert executor.event_buffer[key] == (workload.success_state, None)
+ assert executor.has_task(workload.ti)
+ assert executor._worker_tasks == {second_proc.pid: key}
+ executor.result_queue.put((second_proc.pid, key,
workload.failure_state, None))
+ executor.sync()
+ assert executor.event_buffer[key] == (workload.failure_state, None)
+ assert not executor.running
+ assert not executor._worker_tasks
+ assert not executor._dispatch_counts
+
+ def test_death_of_redispatched_workers_fails_key_after_last_dispatch(
+ self, local_executor_with_mock_worker, mocker
+ ):
+ executor, first_proc = local_executor_with_mock_worker
+ second_proc = _add_mock_worker(executor, mocker, 54321)
+ workload = _make_task_workload()
+ key = executor.get_task_key(workload.ti)
+ session = mock.create_autospec(Session, instance=True)
+ executor.queue_workload(workload, session=session)
+ executor.heartbeat()
+ executor.queue_workload(workload, session=session)
+ executor._process_workloads([workload])
+ executor.result_queue.put((first_proc.pid, key, None, None))
+ executor.result_queue.put((second_proc.pid, key, None, None))
+ executor.sync()
+ first_proc.is_alive.return_value = False
+
+ executor.sync()
+
+ assert key in executor.running
+ assert executor._worker_tasks == {second_proc.pid: key}
+ second_proc.is_alive.return_value = False
+ executor.sync()
+ assert executor.event_buffer[key] == (workload.failure_state, None)
+ assert not executor.running
+
+ def test_late_start_after_connection_test_reaped_is_ignored(self,
local_executor_with_mock_worker):
+ executor, proc = local_executor_with_mock_worker
+ workload = _make_workload("connection")
+ key = executor.get_workload_key(workload)
+ executor.queue_workload(workload,
session=mock.create_autospec(Session, instance=True))
+ executor.heartbeat()
+ executor.fail_connection_test(key)
+ executor.result_queue.put((proc.pid, key, workload.running_state,
None))
+ executor.result_queue.put((proc.pid, key, workload.success_state,
None))
+
+ executor.sync()
+
+ assert executor.event_buffer == {}
+ assert not executor._worker_tasks
+
+ def test_terminal_from_worker_that_does_not_own_the_key_is_ignored(
+ self, local_executor_with_mock_worker, mocker
+ ):
+ executor, owner = local_executor_with_mock_worker
+ other = _add_mock_worker(executor, mocker, 54321)
+ workload = _make_task_workload()
+ key = executor.get_task_key(workload.ti)
+ executor.queue_workload(workload,
session=mock.create_autospec(Session, instance=True))
+ executor.heartbeat()
+ executor.result_queue.put((owner.pid, key, None, None))
+ executor.result_queue.put((other.pid, key, workload.failure_state,
None))
+
+ executor.sync()
+
+ assert executor.event_buffer == {}
+ assert executor._worker_tasks == {owner.pid: key}
+ assert key in executor.running
+
+ def test_result_from_unknown_pid_is_ignored(self,
local_executor_with_mock_worker):
+ executor, proc = local_executor_with_mock_worker
+ workload = _make_task_workload()
+ key = executor.get_task_key(workload.ti)
+ executor.queue_workload(workload,
session=mock.create_autospec(Session, instance=True))
+ executor.heartbeat()
+ executor.result_queue.put((proc.pid + 1, key, None, None))
+
+ executor.sync()
+
+ assert not executor._worker_tasks
+ assert key in executor.running
+
+ def test_start_resets_bookkeeping_of_a_reused_executor(self, mocker):
+ mocker.patch.object(LocalExecutor, "_spawn_workers_with_gc_freeze",
autospec=True)
+ executor = LocalExecutor(parallelism=1)
+ key = TaskInstanceUuid(uuid7())
+ executor._worker_tasks[12345] = key
+ executor._dispatch_counts[key] = 1
+
+ executor.start()
+
+ try:
+ assert not executor._worker_tasks
+ assert not executor._dispatch_counts
+ finally:
+ executor.end()
+
+ @pytest.mark.parametrize("start_method", ["fork", "spawn"])
+ @pytest.mark.parametrize("kind", ["task", "callback", "connection"])
+ @pytest.mark.execution_timeout(30)
+ def test_actual_worker_death_after_start_releases_slot(self, start_method,
kind, mocker, tmp_path):
+ ctx = multiprocessing.get_context(start_method)
+ mocker.patch.object(
+ local_executor_module.multiprocessing,
+ "get_start_method",
+ autospec=True,
+ return_value=start_method,
+ )
+ mocker.patch.object(local_executor_module.multiprocessing, "Process",
new=ctx.Process)
+ mocker.patch.object(local_executor_module.multiprocessing, "Value",
new=ctx.Value)
+ mocker.patch.object(local_executor_module, "SimpleQueue",
new=ctx.SimpleQueue)
+ mocker.patch.object(local_executor_module, "_run_worker",
new=_run_blocking_worker)
+ executor = LocalExecutor(parallelism=1)
+ executor.start()
+ workload = _make_workload(kind)
+ marker = tmp_path / "entered"
+ workload.token = str(marker)
+ key = executor.get_workload_key(workload)
+ executor.queue_workload(workload,
session=mock.create_autospec(Session, instance=True))
+ try:
+ executor.heartbeat()
+ deadline = time.monotonic() + 10
+ while not marker.exists():
+ assert time.monotonic() < deadline
+ executor.sync()
+ time.sleep(0.01)
+ executor.sync()
+ pid, proc = next(iter(executor.workers.items()))
+ assert executor._worker_tasks == {pid: key}
+ proc.kill()
+ proc.join(timeout=1)
+
+ executor.sync()
+
+ assert executor.event_buffer[key] == (workload.failure_state, None)
+ assert executor.slots_available == 1
+ assert not executor._worker_tasks
+ assert not executor.workers
+ executor.result_queue.put((pid, key, workload.success_state, None))
+ executor.sync()
+ assert executor.event_buffer[key] == (workload.failure_state, None)
+ finally:
+ executor.terminate()
+ executor.end()
+
+
class TestLocalExecutorConnectionTestSupport:
def test_test_connection_is_supported(self):
executor = LocalExecutor()
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 557ff4bbcbf..98538f7b970 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
@@ -16,6 +16,7 @@
# under the License.
from __future__ import annotations
+from multiprocessing import SimpleQueue
from unittest.mock import MagicMock
import pytest
@@ -29,8 +30,10 @@ def setup_executor(monkeypatch):
executor.workers = {}
executor._unread_messages = MagicMock()
executor.activity_queue = MagicMock()
+ executor.result_queue = SimpleQueue()
monkeypatch.setattr(executor, "_spawn_worker", MagicMock())
- return executor
+ yield executor
+ executor.result_queue.close()
def test_no_workers_on_no_work(setup_executor):
diff --git a/airflow-core/tests/unit/jobs/test_scheduler_job.py
b/airflow-core/tests/unit/jobs/test_scheduler_job.py
index 6b487bf6347..eafe68660aa 100644
--- a/airflow-core/tests/unit/jobs/test_scheduler_job.py
+++ b/airflow-core/tests/unit/jobs/test_scheduler_job.py
@@ -4328,12 +4328,13 @@ class TestSchedulerJob:
ti.queued_dttm = timezone.utcnow() - timedelta(minutes=15)
session.commit()
+ class NoRevokeExecutor(BaseExecutor):
+ pass
+
assert "revoke_task" in BaseExecutor.__dict__
- # this is just verifying that LocalExecutor is good enough for this
test
- # in that it does not implement revoke_task
- assert "revoke_task" not in LocalExecutor.__dict__
+ assert "revoke_task" not in NoRevokeExecutor.__dict__
scheduler_job = Job()
- job_runner = SchedulerJobRunner(job=scheduler_job, num_runs=0,
executors=[LocalExecutor()])
+ job_runner = SchedulerJobRunner(job=scheduler_job, num_runs=0,
executors=[NoRevokeExecutor()])
job_runner._task_queued_timeout = 300
job_runner._handle_tasks_stuck_in_queued()
