This is an automated email from the ASF dual-hosted git repository. ashb pushed a commit to branch executor-ti-identity in repository https://gitbox.apache.org/repos/asf/airflow.git
commit 4fdd03ab392b13de837f601c081d97221790174d Author: Ash Berlin-Taylor <[email protected]> AuthorDate: Sun Sep 27 20:51:12 2026 +0100 Correlate executor task events by attempt UUID A delayed executor event must identify the attempt that produced it after clearing or retrying replaces the live task instance. Reconcile scheduler events by UUID and retain the retiring identity through restart completion and heartbeat cleanup. Core upgrades can retain older executor providers, so preserve their coordinate-keyed interfaces. Capture UUIDs before submission or adoption and translate legacy events from those snapshots without looking up the current row by coordinates. Retain accepted adoptions across empty event drains, including providers that track adopted work only in private maps. LocalExecutor opts into UUID keys end to end. Legacy providers cannot distinguish attempts that reuse the entire coordinate key; features that require that isolation must require the UUID capability until providers migrate their submission, adoption, result decoding and cleanup paths. --- airflow-core/docs/core-concepts/executor/index.rst | 12 ++ .../src/airflow/executors/base_executor.py | 107 ++++++++-- .../src/airflow/executors/local_executor.py | 13 +- .../src/airflow/executors/workloads/task.py | 3 +- .../src/airflow/executors/workloads/types.py | 5 +- .../src/airflow/jobs/scheduler_job_runner.py | 31 +-- .../versions/head/test_task_instances.py | 4 +- .../tests/unit/executors/test_base_executor.py | 232 +++++++++++++++++++-- .../tests/unit/executors/test_local_executor.py | 9 +- .../tests/unit/executors/test_workloads.py | 4 +- airflow-core/tests/unit/jobs/test_scheduler_job.py | 116 ++++++----- .../src/tests_common/test_utils/mock_executor.py | 5 +- .../executors/test_kubernetes_executor.py | 2 + 13 files changed, 420 insertions(+), 123 deletions(-) diff --git a/airflow-core/docs/core-concepts/executor/index.rst b/airflow-core/docs/core-concepts/executor/index.rst index cf21daafd3f..4a8d7faacff 100644 --- a/airflow-core/docs/core-concepts/executor/index.rst +++ b/airflow-core/docs/core-concepts/executor/index.rst @@ -336,6 +336,18 @@ The ``BaseExecutor`` class interface contains a set of attributes that Airflow c * ``is_single_threaded``: Whether or not the executor is single threaded. This is particularly relevant to what database backends are supported. Single threaded executors can run with any backend, including SQLite. * ``is_production``: Whether or not the executor should be used for production purposes. A UI message is displayed to users when they are using a non-production ready executor. * ``serve_logs``: Whether or not the executor supports serving logs, see :doc:`/administration-and-deployment/logging-monitoring/logging-tasks`. +* ``supports_task_instance_uuid``: Whether task submission, adoption, result decoding and cleanup all use task-instance UUIDs. Defaults to ``False`` to support existing provider releases. ``LocalExecutor`` enables this capability. + +UUID-capable executors use ``get_workload_key(workload)`` and ``get_task_key(ti)`` for executor bookkeeping. +``ExecuteTask.key`` retains its coordinate-based value for existing providers. Providers supporting older +Airflow releases must retain their older-core paths when adopting these helpers. + +The scheduler translates legacy coordinate events using UUIDs captured during submission or adoption; +it never resolves them against the current database row. Unknown or ambiguous identities are discarded. +Legacy executors cannot distinguish a delayed event from a new attempt that reuses the entire coordinate +key, including the try number. Features requiring isolation across such reuse must require +``supports_task_instance_uuid``. Updating dispatch keys alone does not satisfy this capability: remote +results and adopted tasks must retain the UUID of the submitted attempt. CLI ^^^ diff --git a/airflow-core/src/airflow/executors/base_executor.py b/airflow-core/src/airflow/executors/base_executor.py index 991365f17d8..83525eb6d87 100644 --- a/airflow-core/src/airflow/executors/base_executor.py +++ b/airflow-core/src/airflow/executors/base_executor.py @@ -24,8 +24,9 @@ import warnings from collections import defaultdict, deque from collections.abc import Sequence from dataclasses import dataclass, field -from functools import cached_property +from functools import cached_property, wraps from typing import TYPE_CHECKING, Any, ClassVar, cast +from uuid import UUID import pendulum @@ -38,13 +39,14 @@ from airflow.executors.executor_loader import ExecutorLoader from airflow.executors.workloads import WorkloadType from airflow.executors.workloads.callback import ExecuteCallback from airflow.executors.workloads.connection_test import TestConnection -from airflow.executors.workloads.task import ExecuteTask +from airflow.executors.workloads.task import ExecuteTask, TaskInstanceDTO from airflow.executors.workloads.types import state_class_for_key from airflow.models import Log from airflow.models.taskinstancekey import TaskInstanceKey from airflow.observability.metrics import stats_utils from airflow.utils.helpers import prune_dict from airflow.utils.log.logging_mixin import LoggingMixin +from airflow.utils.state import State PARALLELISM: int = conf.getint("core", "PARALLELISM") @@ -209,6 +211,9 @@ class BaseExecutor(LoggingMixin): supports_callbacks = _LegacyWorkloadFlag(WorkloadType.EXECUTE_CALLBACK) supports_connection_test = _LegacyWorkloadFlag(WorkloadType.TEST_CONNECTION) supports_multi_team: bool = False + # Opt in only when submission, adoption, result decoding and cleanup all carry UUIDs. + # Coordinate-only providers cannot distinguish attempts that reuse the entire key. + supports_task_instance_uuid: ClassVar[bool] = False sentry_integration: str = "" _legacy_warned: ClassVar[set[str]] = set() @@ -248,6 +253,31 @@ class BaseExecutor(LoggingMixin): def __init_subclass__(cls, **kwargs: Any) -> None: super().__init_subclass__(**kwargs) + if queue_workload := cls.__dict__.get("queue_workload"): + + @wraps(queue_workload) + def register_queued_task(self, workload, *args, **kwargs): + if isinstance(workload, ExecuteTask): + self.register_task(workload.ti) + return queue_workload(self, workload, *args, **kwargs) + + cls.queue_workload = register_queued_task # type: ignore[method-assign] + if try_adopt := cls.__dict__.get("try_adopt_task_instances"): + + @wraps(try_adopt) + def register_adopted_tasks(self, tis, *args, **kwargs): + keys = [(ti.id, self.get_task_key(ti)) for ti in tis] + for ti in tis: + self.register_task(ti) + rejected = try_adopt(self, tis, *args, **kwargs) + rejected_ids = {ti.id for ti in rejected} + for task_id, key in keys: + state, _ = self.event_buffer.get(key, (None, None)) + if task_id not in rejected_ids and state not in State.finished: + self.running.add(key) + return rejected + + cls.try_adopt_task_instances = register_adopted_tasks # type: ignore[method-assign] cls._legacy_warned = set() enabled: set[WorkloadType] = set() disabled: set[WorkloadType] = set() @@ -310,6 +340,7 @@ class BaseExecutor(LoggingMixin): ) self.running: set[WorkloadKey] = set() self.event_buffer: dict[WorkloadKey, EventBufferValueType] = {} + self._task_coordinates: dict[UUID, TaskInstanceKey] = {} self._task_event_logs: deque[Log] = deque() self.conf = ExecutorConf(team_name) @@ -382,10 +413,29 @@ class BaseExecutor(LoggingMixin): def log_task_event(self, *, event: str, extra: str, ti_key: WorkloadKey): """Add an event to the log table.""" - if not isinstance(ti_key, TaskInstanceKey): + if not isinstance(ti_key, (UUID, TaskInstanceKey)): self.log.debug("Skipping log_task_event for callback key %s (event=%s)", ti_key, event) return - self._task_event_logs.append(Log(event=event, task_instance=ti_key, extra=extra)) + coordinates = ti_key if isinstance(ti_key, TaskInstanceKey) else self._task_coordinates.get(ti_key) + if coordinates is None: + extra = f"Task instance {ti_key}: {extra}" + self._task_event_logs.append(Log(event=event, task_instance=coordinates, extra=extra)) + + def register_task(self, ti: TaskInstance | TaskInstanceDTO) -> None: + """Snapshot identity before provider submission or adoption can emit events.""" + self._task_coordinates[ti.id] = ti.key + + @classmethod + def get_task_key(cls, ti: TaskInstance | TaskInstanceDTO) -> UUID | TaskInstanceKey: + """Return the task key supported by this executor's provider contract.""" + return ti.id if cls.supports_task_instance_uuid else ti.key + + @classmethod + def get_workload_key(cls, workload: ExecutorWorkload) -> WorkloadKey: + """Return the executor key without changing the legacy workload property.""" + if isinstance(workload, ExecuteTask): + return cls.get_task_key(workload.ti) + return workload.key def queue_workload(self, workload: ExecutorWorkload, session: Session) -> None: if workload.type not in self.supported_workload_types: @@ -394,7 +444,9 @@ class BaseExecutor(LoggingMixin): f"Add WorkloadType.{workload.type.name} to supported_workload_types and implement handling " f"in _process_workloads()." ) - self.executor_queues[workload.type][workload.key] = workload + if isinstance(workload, ExecuteTask): + self.register_task(workload.ti) + self.executor_queues[workload.type][self.get_workload_key(workload)] = workload def _get_workloads_to_schedule(self, open_slots: int) -> list[tuple[WorkloadKey, ExecutorWorkload]]: """ @@ -437,12 +489,10 @@ class BaseExecutor(LoggingMixin): :return: True if the task is known to this executor """ task_queue = self.executor_queues.get(WorkloadType.EXECUTE_TASK, {}) - return ( - task_instance.id in task_queue - or task_instance.id in self.running - or task_instance.key in task_queue - or task_instance.key in self.running - ) + key = self.get_task_key(task_instance) + if isinstance(key, TaskInstanceKey) and self._task_coordinates.get(task_instance.id) != key: + return False + return key in task_queue or key in self.running def sync(self) -> None: """ @@ -641,11 +691,44 @@ class BaseExecutor(LoggingMixin): self.event_buffer = {} else: for key in list(self.event_buffer.keys()): - if not isinstance(key, TaskInstanceKey) or key.dag_id in dag_ids: + coordinates = self._task_coordinates.get(key) if isinstance(key, UUID) else None + if isinstance(key, TaskInstanceKey): + coordinates = key + if not isinstance(key, (UUID, TaskInstanceKey)) or ( + coordinates is not None and coordinates.dag_id in dag_ids + ): cleared_events[key] = self.event_buffer.pop(key) + task_queue = self.executor_queues.get(WorkloadType.EXECUTE_TASK, {}) + for task_id, coordinates in list(self._task_coordinates.items()): + if not any( + key in container + for key in (task_id, coordinates) + for container in (self.running, task_queue, self.event_buffer) + ): + del self._task_coordinates[task_id] + return cleared_events + def get_event_buffer_with_task_ids(self, dag_ids=None) -> dict[WorkloadKey, EventBufferValueType]: + """Drain events using submission/adoption UUIDs, never a current-row coordinate lookup.""" + task_ids: dict[TaskInstanceKey, UUID | None] = {} + for task_id, coordinates in self._task_coordinates.items(): + task_ids[coordinates] = None if coordinates in task_ids else task_id + events: dict[WorkloadKey, EventBufferValueType] = {} + for key, value in self.get_event_buffer(dag_ids).items(): + if isinstance(key, TaskInstanceKey): + captured_id = task_ids.get(key) + if captured_id is None: + self.log.warning( + "Discarding executor event without a unique captured task identity: %s", key + ) + continue + events[captured_id] = value + else: + events[key] = value + return events + def get_task_log(self, ti: TaskInstance, try_number: int) -> tuple[list[str], list[str]]: """ Return the task logs. diff --git a/airflow-core/src/airflow/executors/local_executor.py b/airflow-core/src/airflow/executors/local_executor.py index eb43093ffd7..74c92a4f31a 100644 --- a/airflow-core/src/airflow/executors/local_executor.py +++ b/airflow-core/src/airflow/executors/local_executor.py @@ -97,8 +97,9 @@ def _run_worker( with unread_messages: unread_messages.value -= 1 + key = LocalExecutor.get_workload_key(workload) if workload.running_state is not None: - output.put((workload.key, workload.running_state, None)) + output.put((key, workload.running_state, None)) try: BaseExecutor.run_workload( @@ -107,10 +108,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((workload.key, workload.success_state, None)) + output.put((key, workload.success_state, None)) except Exception as e: log.exception("Workload execution failed.", workload_type=type(workload).__name__) - output.put((workload.key, workload.failure_state, e)) + output.put((key, workload.failure_state, e)) class LocalExecutor(BaseExecutor): @@ -126,6 +127,7 @@ class LocalExecutor(BaseExecutor): is_mp_using_fork: bool supports_multi_team: bool = True + supports_task_instance_uuid = True serve_logs: bool = True # The connection-test supervisor uses ``signal.SIGALRM`` (via ``TimeoutPosix``) to bound hook # execution, so ``TEST_CONNECTION`` support requires a POSIX worker (LocalExecutor runs on the host). @@ -315,9 +317,10 @@ class LocalExecutor(BaseExecutor): def _process_workloads(self, workload_list): for workload in workload_list: self.activity_queue.put(workload) - removed = self.executor_queues[workload.type].pop(workload.key, None) + key = self.get_workload_key(workload) + removed = self.executor_queues[workload.type].pop(key, None) if not removed: - raise KeyError(f"Workload {workload.key} was not found in any queue") + raise KeyError(f"Workload {key} was not found in any queue") with self._unread_messages: self._unread_messages.value += len(workload_list) self._check_workers() diff --git a/airflow-core/src/airflow/executors/workloads/task.py b/airflow-core/src/airflow/executors/workloads/task.py index 56c69724038..5b1274d35b7 100644 --- a/airflow-core/src/airflow/executors/workloads/task.py +++ b/airflow-core/src/airflow/executors/workloads/task.py @@ -48,7 +48,6 @@ class TaskInstanceDTO(TaskInstance): external_executor_id: str | None = Field(default=None, exclude=True) executor_config: dict | None = Field(default=None, exclude=True) - # TODO: Task-SDK: Can we replace TaskInstanceKey with just the uuid across the codebase? @property def key(self) -> TaskInstanceKey: from airflow.models.taskinstancekey import TaskInstanceKey @@ -72,7 +71,7 @@ class ExecuteTask(BaseDagBundleWorkload): @property def key(self) -> TaskInstanceKey: - """Return the TaskInstanceKey for this workload.""" + """Return the coordinate key used by existing executor providers.""" return self.ti.key @property diff --git a/airflow-core/src/airflow/executors/workloads/types.py b/airflow-core/src/airflow/executors/workloads/types.py index 627f28ffb3f..6b882df5b42 100644 --- a/airflow-core/src/airflow/executors/workloads/types.py +++ b/airflow-core/src/airflow/executors/workloads/types.py @@ -19,6 +19,7 @@ from __future__ import annotations from typing import TYPE_CHECKING, TypeAlias +from uuid import UUID from airflow.models.callback import CallbackKey, ExecutorCallback from airflow.models.connection_test import ConnectionTestKey, ConnectionTestRequest, ConnectionTestState @@ -27,7 +28,7 @@ from airflow.models.taskinstancekey import TaskInstanceKey from airflow.utils.state import CallbackState, TaskInstanceState # Type aliases for workload keys and states (used by executor layer) -WorkloadKey: TypeAlias = TaskInstanceKey | CallbackKey | ConnectionTestKey +WorkloadKey: TypeAlias = UUID | TaskInstanceKey | CallbackKey | ConnectionTestKey WorkloadState: TypeAlias = TaskInstanceState | CallbackState | ConnectionTestState if TYPE_CHECKING: @@ -42,7 +43,7 @@ SchedulerWorkload: TypeAlias = TaskInstance | ExecutorCallback | ConnectionTestR def state_class_for_key( key: WorkloadKey, ) -> type[TaskInstanceState] | type[CallbackState] | type[ConnectionTestState]: - if isinstance(key, TaskInstanceKey): + if isinstance(key, (UUID, TaskInstanceKey)): return TaskInstanceState if isinstance(key, ConnectionTestKey): return ConnectionTestState diff --git a/airflow-core/src/airflow/jobs/scheduler_job_runner.py b/airflow-core/src/airflow/jobs/scheduler_job_runner.py index 741a7b4be6f..a952f164523 100644 --- a/airflow-core/src/airflow/jobs/scheduler_job_runner.py +++ b/airflow-core/src/airflow/jobs/scheduler_job_runner.py @@ -108,7 +108,6 @@ from airflow.models.log import resolve_team_name from airflow.models.pool import normalize_pool_name_for_stats from airflow.models.serialized_dag import SerializedDagModel from airflow.models.taskinstance import TaskInstance -from airflow.models.taskinstancekey import TaskInstanceKey from airflow.models.team import Team from airflow.models.trigger import TRIGGER_FAIL_REPR, Trigger, TriggerFailureReason, handle_event_submit from airflow.observability.metrics import stats_utils @@ -1469,16 +1468,14 @@ class SchedulerJobRunner(BaseJobRunner, LoggingMixin): `dag.test` execute DAGs with no scheduler, therefore it needs to handle the events pushed by the executors as well. """ - ti_event_keys: dict[tuple[str, str, str, int], list[TaskInstanceKey]] = defaultdict(list) - event_buffer = executor.get_event_buffer() + event_buffer = executor.get_event_buffer_with_task_ids() num_events = len(event_buffer) - tis_with_right_state: list[TaskInstanceKey] = [] + tis_with_right_state: list[UUID] = [] callback_keys_with_events: list[CallbackKey] = [] # Report execution - handle both task and callback events for key, (state, _) in event_buffer.items(): - if isinstance(key, TaskInstanceKey): - ti_event_keys[key.primary].append(key) + if isinstance(key, UUID): cls.logger().info("Received executor event with state %s for task instance %s", state, key) if state in ( TaskInstanceState.FAILED, @@ -1525,14 +1522,10 @@ class SchedulerJobRunner(BaseJobRunner, LoggingMixin): return len(event_buffer) # Check state of finished tasks - filter_for_tis = TI.filter_for_tis(tis_with_right_state) - if filter_for_tis is None: - cls._emit_executor_events_batch_metrics(num_events) - return len(event_buffer) asset_loader, alias_loader = _eager_load_dag_run_for_validation() query = ( select(TI) - .where(filter_for_tis) + .where(TI.id.in_(tis_with_right_state)) .options(selectinload(TI.dag_model)) .options(asset_loader) .options(alias_loader) @@ -1550,13 +1543,7 @@ class SchedulerJobRunner(BaseJobRunner, LoggingMixin): locked_query = with_row_locks(query, of=TI, session=session, skip_locked=True) tis: Iterator[TI] = session.scalars(locked_query.execution_options(populate_existing=True)) for ti in tis: - for event_key in ti_event_keys[ti.key.primary]: - if event_key.try_number != ti.try_number: - cls.logger().info("Ignoring executor event for a different attempt: %s", event_key) - event_buffer.pop(event_key) - buffer_key = ti.key - if buffer_key not in event_buffer: - continue + buffer_key = ti.id try_number = ti.try_number state, info = event_buffer.pop(buffer_key) @@ -1621,7 +1608,7 @@ class SchedulerJobRunner(BaseJobRunner, LoggingMixin): # All of this could also happen if the state is "running", # but that is handled by the scheduler detecting task instances without heartbeats. - ti_queued = ti.try_number == buffer_key.try_number and ti.state in ( + ti_queued = ti.state in ( TaskInstanceState.SCHEDULED, TaskInstanceState.QUEUED, TaskInstanceState.RUNNING, @@ -3338,7 +3325,7 @@ class SchedulerJobRunner(BaseJobRunner, LoggingMixin): executor.send_callback(request) finally: ti.set_state(TaskInstanceState.FAILED, session=session) - executor.fail(ti.key) + executor.fail(executor.get_task_key(ti)) def _reschedule_stuck_task(self, ti: TaskInstance, session: Session): filter_for_tis = TI.filter_for_tis([ti]) @@ -3560,6 +3547,7 @@ class SchedulerJobRunner(BaseJobRunner, LoggingMixin): TI.task_id, TI.run_id, TI.map_index, + TI.try_number, TI.state, TI.external_executor_id, ) @@ -3916,7 +3904,7 @@ class SchedulerJobRunner(BaseJobRunner, LoggingMixin): ) ) - failed_key = ti.key + failed_id, failed_coordinates = ti.id, ti.key if ti.state == TaskInstanceState.RESTARTING: ti.notify_failure(error=msg) ti.complete_restart(session=session) @@ -3932,6 +3920,7 @@ class SchedulerJobRunner(BaseJobRunner, LoggingMixin): ti.executor, ) continue + failed_key = failed_id if executor.supports_task_instance_uuid else failed_coordinates executor.change_state(failed_key, TaskInstanceState.FAILED, remove_running=True) stats.incr( "task_instances_without_heartbeats_killed", diff --git a/airflow-core/tests/unit/api_fastapi/execution_api/versions/head/test_task_instances.py b/airflow-core/tests/unit/api_fastapi/execution_api/versions/head/test_task_instances.py index c447149ac2c..a938461a636 100644 --- a/airflow-core/tests/unit/api_fastapi/execution_api/versions/head/test_task_instances.py +++ b/airflow-core/tests/unit/api_fastapi/execution_api/versions/head/test_task_instances.py @@ -1497,7 +1497,7 @@ class TestTIUpdateState: ) ti.hostname, ti.pid = "worker", 123 ti.try_number, ti.max_tries = 3, 0 - old_id, old_key = ti.id, ti.key + old_id = ti.id session.commit() if definition == "missing_dag": mocker.patch.object(DBDagBag, "get_dag_for_run", autospec=True, return_value=None) @@ -1525,7 +1525,7 @@ class TestTIUpdateState: ) assert response.status_code == 204 else: - executor.event_buffer[old_key] = executor_state, None + executor.event_buffer[old_id] = executor_state, None SchedulerJobRunner.process_executor_events(executor, None, DBDagBag(), session) session.commit() session.expunge_all() diff --git a/airflow-core/tests/unit/executors/test_base_executor.py b/airflow-core/tests/unit/executors/test_base_executor.py index 3ff2dd5a2cd..cad58e946a8 100644 --- a/airflow-core/tests/unit/executors/test_base_executor.py +++ b/airflow-core/tests/unit/executors/test_base_executor.py @@ -102,9 +102,11 @@ def test_get_event_buffer(): date = timezone.utcnow() try_number = 1 - key1 = TaskInstanceKey("my_dag1", "my_task1", date, try_number) - key2 = TaskInstanceKey("my_dag2", "my_task1", date, try_number) - key3 = TaskInstanceKey("my_dag2", "my_task2", date, try_number) + key1, key2, key3 = uuid4(), uuid4(), uuid4() + for key, dag_id in ((key1, "my_dag1"), (key2, "my_dag2"), (key3, "my_dag2")): + executor.register_task( + mock.Mock(spec=TaskInstance, id=key, key=TaskInstanceKey(dag_id, "task", date, try_number)) + ) state = State.SUCCESS executor.event_buffer[key1] = state, None executor.event_buffer[key2] = state, None @@ -120,7 +122,10 @@ def test_get_event_buffer_always_includes_callback_keys(): executor = BaseExecutor() date = timezone.utcnow() - ti_key = TaskInstanceKey("my_dag1", "my_task1", date, 1) + ti_key = uuid4() + executor.register_task( + mock.Mock(spec=TaskInstance, id=ti_key, key=TaskInstanceKey("my_dag1", "my_task1", date, 1)) + ) callback_key = CallbackKey(id="00000000-0000-0000-0000-000000000042") executor.event_buffer[ti_key] = State.SUCCESS, None @@ -132,12 +137,44 @@ def test_get_event_buffer_always_includes_callback_keys(): assert ti_key not in result +def test_get_event_buffer_retains_unregistered_task_until_unfiltered_drain(): + executor = BaseExecutor() + retired_id = uuid4() + executor.event_buffer[retired_id] = TaskInstanceState.SUCCESS, None + + assert executor.get_event_buffer(("some_dag",)) == {} + assert executor.get_event_buffer() == {retired_id: (TaskInstanceState.SUCCESS, None)} + + +def test_get_event_buffer_releases_metadata_for_discarded_tasks(): + executor = BaseExecutor() + ti = mock.Mock(spec=TaskInstance, id=uuid4(), key=TaskInstanceKey("dag", "task", "run", 1)) + executor.register_task(ti) + executor.running.add(ti.id) + executor.get_event_buffer() + assert ti.id in executor._task_coordinates + + executor.running.remove(ti.id) + executor.get_event_buffer() + + assert executor._task_coordinates == {} + + def test_log_task_event_branches_on_key_type(): executor = BaseExecutor() - ti_key = TaskInstanceKey("my_dag", "my_task", timezone.utcnow(), 1) + ti_key = uuid4() + executor.register_task( + mock.Mock( + spec=TaskInstance, + id=ti_key, + key=TaskInstanceKey("my_dag", "my_task", "run", 1), + ) + ) executor.log_task_event(event="task_event", extra="extra", ti_key=ti_key) assert len(executor._task_event_logs) == 1 + log = executor._task_event_logs[0] + assert (log.dag_id, log.task_id, log.run_id, log.try_number) == ("my_dag", "my_task", "run", 1) callback_key = CallbackKey(id=str(UUID("00000000-0000-0000-0000-000000000001"))) executor.log_task_event(event="callback_event", extra="extra", ti_key=callback_key) @@ -169,13 +206,9 @@ def test_state_methods_pick_callback_state_for_callback_key(method_name, expecte def test_fail_and_success(): executor = BaseExecutor() - date = timezone.utcnow() - try_number = 1 success_state = State.SUCCESS fail_state = State.FAILED - key1 = TaskInstanceKey("my_dag1", "my_task1", date, try_number) - key2 = TaskInstanceKey("my_dag2", "my_task1", date, try_number) - key3 = TaskInstanceKey("my_dag2", "my_task2", date, try_number) + key1, key2, key3 = uuid4(), uuid4(), uuid4() executor.fail(key1, fail_state) executor.fail(key2, fail_state) executor.success(key3, success_state) @@ -185,6 +218,161 @@ def test_fail_and_success(): assert len(executor.get_event_buffer()) == 3 [email protected] +def task_workload(): + return workloads.ExecuteTask( + ti=workloads.TaskInstanceDTO( + id=uuid4(), + dag_version_id=uuid4(), + dag_id="dag", + task_id="task", + run_id="run", + try_number=1, + pool_slots=1, + priority_weight=1, + queue="default", + ), + dag_rel_path="dag.py", + bundle_info=BundleInfo(name="bundle"), + token="", + log_path=None, + ) + + [email protected]("override_queue", [False, True]) [email protected]("terminal_state", [TaskInstanceState.SUCCESS, TaskInstanceState.FAILED]) [email protected]("next_try_number", [1, 2]) +def test_legacy_executor_dispatch_and_completion_keep_submitted_identity( + task_workload, override_queue, terminal_state, next_try_number +): + class LegacyExecutor(BaseExecutor): + def _process_workloads(self, items): + for workload in items: + key = workload.ti.key + del self.executor_queues[WorkloadType.EXECUTE_TASK][key] + self.running.add(key) + self.event_buffer[key] = (TaskInstanceState.QUEUED, "remote-id") + + class OverrideExecutor(LegacyExecutor): + def queue_workload(self, workload, session): + self.executor_queues[workload.type][workload.key] = workload + + executor = (OverrideExecutor if override_queue else LegacyExecutor)() + submitted_id, submitted_key = task_workload.ti.id, task_workload.ti.key + executor.queue_workload(task_workload, session=None) + assert executor.has_task(task_workload.ti) + executor.heartbeat() + assert executor.get_event_buffer_with_task_ids() == { + submitted_id: (TaskInstanceState.QUEUED, "remote-id") + } + assert executor.has_task(task_workload.ti) + + task_workload.ti.id = uuid4() + task_workload.ti.try_number = next_try_number + assert not executor.has_task(task_workload.ti) + getattr(executor, "success" if terminal_state == TaskInstanceState.SUCCESS else "fail")(submitted_key) + + assert executor.get_event_buffer_with_task_ids() == {submitted_id: (terminal_state, None)} + assert not executor.running + assert not executor._task_coordinates + + +def test_legacy_adoption_captures_identity_before_provider_emits_event(task_workload): + class AdoptingExecutor(BaseExecutor): + def try_adopt_task_instances(self, tis): + for ti in tis: + self.success(ti.key) + return [] + + executor = AdoptingExecutor() + adopted_id = task_workload.ti.id + assert executor.try_adopt_task_instances([task_workload.ti]) == [] + task_workload.ti.id = uuid4() + assert executor.get_event_buffer_with_task_ids() == {adopted_id: (TaskInstanceState.SUCCESS, None)} + assert not executor.running + + [email protected]("accepted", [False, True]) +def test_adoption_with_private_provider_tracking_survives_empty_event_drain(task_workload, accepted): + class AdoptingExecutor(BaseExecutor): + def try_adopt_task_instances(self, tis): + return [] if accepted else tis + + executor = AdoptingExecutor() + executor.try_adopt_task_instances([task_workload.ti]) + assert executor.get_event_buffer_with_task_ids() == {} + if not accepted: + assert not executor.running + assert not executor._task_coordinates + return + + assert executor.has_task(task_workload.ti) + executor.success(task_workload.ti.key) + assert executor.get_event_buffer_with_task_ids() == { + task_workload.ti.id: (TaskInstanceState.SUCCESS, None) + } + assert not executor.running + assert not executor._task_coordinates + + [email protected]("ambiguous", [False, True]) +def test_legacy_events_without_unique_captured_identity_are_dropped(task_workload, ambiguous): + executor = BaseExecutor() + if ambiguous: + executor.register_task(task_workload.ti) + executor.register_task(task_workload.ti.model_copy(update={"id": uuid4()})) + executor.success(task_workload.ti.key) + assert executor.get_event_buffer_with_task_ids() == {} + + +def test_legacy_event_dag_filter_retains_identity_until_consumed(task_workload): + executor = BaseExecutor() + executor.queue_workload(task_workload, session=None) + executor.executor_queues[WorkloadType.EXECUTE_TASK].clear() + executor.success(task_workload.ti.key) + assert executor.get_event_buffer_with_task_ids(["other_dag"]) == {} + assert executor.get_event_buffer_with_task_ids(["dag"]) == { + task_workload.ti.id: (TaskInstanceState.SUCCESS, None) + } + + [email protected]("next_try_number", [1, 2]) +def test_task_workloads_with_matching_coordinates_keep_distinct_attempts(next_try_number): + executor = LocalExecutor() + first = workloads.ExecuteTask( + ti=workloads.TaskInstanceDTO( + id=uuid4(), + dag_version_id=uuid4(), + dag_id="dag", + task_id="task", + run_id="run", + try_number=1, + pool_slots=1, + priority_weight=1, + queue="default", + ), + dag_rel_path="dag.py", + bundle_info=BundleInfo(name="bundle"), + token="", + log_path=None, + ) + second = first.model_copy( + update={"ti": first.ti.model_copy(update={"id": uuid4(), "try_number": next_try_number})} + ) + + executor.queue_workload(first, session=None) + executor.queue_workload(second, session=None) + + assert set(executor.executor_queues[WorkloadType.EXECUTE_TASK]) == {first.ti.id, second.ti.id} + executor.running.update((first.ti.id, second.ti.id)) + executor.success(first.ti.id) + assert executor.running == {second.ti.id} + assert executor.get_event_buffer() == {first.ti.id: (TaskInstanceState.SUCCESS, None)} + executor.executor_queues[WorkloadType.EXECUTE_TASK].clear() + assert not executor.has_task(first.ti) + assert executor.has_task(second.ti) + + @pytest.mark.parametrize( ("team_name", "expected_tags"), [ @@ -317,7 +505,7 @@ def setup_trigger_workloads(dag_maker, parallelism=None): for task_instance in dagrun.task_instances: workload = workloads.ExecuteTask.make(task_instance) - executor.executor_queues[WorkloadType.EXECUTE_TASK][task_instance.key] = workload + executor.executor_queues[WorkloadType.EXECUTE_TASK][task_instance.id] = workload return executor, dagrun @@ -355,7 +543,7 @@ def test_trigger_running_tasks(dag_maker): ti = dagrun.task_instances[0] workload = workloads.ExecuteTask.make(ti) - executor.executor_queues[WorkloadType.EXECUTE_TASK][ti.key] = workload + executor.executor_queues[WorkloadType.EXECUTE_TASK][ti.id] = workload executor.trigger_workloads(open_slots=10) @@ -379,7 +567,7 @@ def test_trigger_workloads_schedules_highest_priority_first(dag_maker): executor._process_workloads = mock.Mock(spec=lambda workloads: None) for task_instance in dagrun.task_instances: task_queue = executor.executor_queues[WorkloadType.EXECUTE_TASK] - task_queue[task_instance.key] = workloads.ExecuteTask.make(task_instance) + task_queue[task_instance.id] = workloads.ExecuteTask.make(task_instance) executor.trigger_workloads(open_slots=2) @@ -403,7 +591,7 @@ def test_debug_dump_with_populated_queues(caplog, dag_maker): for ti in dagrun.task_instances: workload = workloads.ExecuteTask.make(ti) - executor.executor_queues[WorkloadType.EXECUTE_TASK][ti.key] = workload + executor.executor_queues[WorkloadType.EXECUTE_TASK][ti.id] = workload with caplog.at_level(logging.INFO): executor.debug_dump() @@ -492,7 +680,7 @@ def test_has_task_does_not_vivify_executor_queue(): def test_unknown_workload_type_sorts_last_without_crashing(): executor = BaseExecutor() - known_key = TaskInstanceKey("d", "t", "r", 1, -1) + known_key = uuid4() known_workload = mock.Mock() known_workload.type = WorkloadType.EXECUTE_TASK known_workload.sort_key = 0 @@ -559,7 +747,7 @@ def test_running_retry_attempt_type(loop_duration, total_tries): def test_state_fail(): executor = BaseExecutor() - key = TaskInstanceKey("my_dag1", "my_task1", timezone.utcnow(), 1) + key = uuid4() executor.running.add(key) info = "info" executor.fail(key, info=info) @@ -569,7 +757,7 @@ def test_state_fail(): def test_state_success(): executor = BaseExecutor() - key = TaskInstanceKey("my_dag1", "my_task1", timezone.utcnow(), 1) + key = uuid4() executor.running.add(key) info = "info" executor.success(key, info=info) @@ -579,7 +767,7 @@ def test_state_success(): def test_state_queued(): executor = BaseExecutor() - key = TaskInstanceKey("my_dag1", "my_task1", timezone.utcnow(), 1) + key = uuid4() executor.running.add(key) info = "info" executor.queued(key, info=info) @@ -589,7 +777,7 @@ def test_state_queued(): def test_state_running(): executor = BaseExecutor() - key = TaskInstanceKey("my_dag1", "my_task1", timezone.utcnow(), 1) + key = uuid4() executor.running.add(key) info = "info" executor.running_state(key, info=info) @@ -655,13 +843,11 @@ def test_queued_connection_test_dispatched_to_process_workloads(): def test_connection_tests_prioritized_ahead_of_task_backlog(): """A task backlog must not starve short, user-interactive connection tests.""" executor = BaseExecutor() - for i in range(3): + for _ in range(3): task_workload = mock.Mock() task_workload.type = WorkloadType.EXECUTE_TASK task_workload.sort_key = 0 - executor.executor_queues[WorkloadType.EXECUTE_TASK][TaskInstanceKey("d", f"t{i}", "r", 1, -1)] = ( - task_workload - ) + executor.executor_queues[WorkloadType.EXECUTE_TASK][uuid4()] = task_workload conn_test = mock.Mock() conn_test.type = WorkloadType.TEST_CONNECTION conn_test.sort_key = 0 diff --git a/airflow-core/tests/unit/executors/test_local_executor.py b/airflow-core/tests/unit/executors/test_local_executor.py index 28adefd22aa..b438ec1291d 100644 --- a/airflow-core/tests/unit/executors/test_local_executor.py +++ b/airflow-core/tests/unit/executors/test_local_executor.py @@ -36,7 +36,6 @@ from airflow.executors.workloads.base import BundleInfo from airflow.executors.workloads.callback import CallbackDTO from airflow.executors.workloads.task import TaskInstanceDTO from airflow.models.callback import CallbackFetchMethod -from airflow.models.taskinstancekey import TaskInstanceKey from airflow.settings import Session from airflow.utils.state import State @@ -94,8 +93,8 @@ def _make_task_workload(): def _write_large_results_to_queue(result_queue, result_count, payload_size): payload = RuntimeError("x" * payload_size) - for index in range(result_count): - key = TaskInstanceKey("test_dag", f"test_task_{index}", "test_run") + for _ in range(result_count): + key = uuid7() result_queue.put((key, State.SUCCESS, payload)) @@ -237,8 +236,8 @@ class TestLocalExecutor: assert executor._unread_messages.value == 0 for ti in success_tis: - assert executor.event_buffer[ti.key][0] == State.SUCCESS - assert executor.event_buffer[fail_ti.key][0] == State.FAILED + assert executor.event_buffer[ti.id][0] == State.SUCCESS + assert executor.event_buffer[fail_ti.id][0] == State.FAILED @mock.patch("airflow.executors.local_executor.LocalExecutor.sync") @mock.patch("airflow.executors.base_executor.BaseExecutor.trigger_workloads") diff --git a/airflow-core/tests/unit/executors/test_workloads.py b/airflow-core/tests/unit/executors/test_workloads.py index cd31ea833c5..8530e25234d 100644 --- a/airflow-core/tests/unit/executors/test_workloads.py +++ b/airflow-core/tests/unit/executors/test_workloads.py @@ -20,7 +20,7 @@ from __future__ import annotations import dataclasses from pathlib import PurePosixPath from typing import get_args -from uuid import uuid4 +from uuid import UUID, uuid4 import jwt import pytest @@ -76,7 +76,7 @@ def test_workload_families_track_every_workload_type(): assert {schema.model_fields["type"].default for schema in schemas} == set(WorkloadType) assert _union_members(workloads.ExecutorWorkload) == schemas assert _union_members(workloads.All) == schemas | {RunTrigger} - assert _union_members(WorkloadKey) == keys + assert _union_members(WorkloadKey) == keys | {UUID} assert _union_members(WorkloadState) == states assert _union_members(SchedulerWorkload) == models diff --git a/airflow-core/tests/unit/jobs/test_scheduler_job.py b/airflow-core/tests/unit/jobs/test_scheduler_job.py index af339fa0e92..b4a0ceb1e38 100644 --- a/airflow-core/tests/unit/jobs/test_scheduler_job.py +++ b/airflow-core/tests/unit/jobs/test_scheduler_job.py @@ -546,7 +546,7 @@ class TestSchedulerJob: session.merge(ti1) session.commit() - executor.event_buffer[ti1.key] = State.FAILED, None + executor.event_buffer[ti1.id] = State.FAILED, None self.job_runner._process_executor_events(executor=executor, session=session) ti1.refresh_from_db(session=session) @@ -557,7 +557,7 @@ class TestSchedulerJob: ti1.state = State.SUCCESS session.merge(ti1) session.commit() - executor.event_buffer[ti1.key] = State.SUCCESS, None + executor.event_buffer[ti1.id] = State.SUCCESS, None self.job_runner._process_executor_events(executor=executor, session=session) ti1.refresh_from_db(session=session) @@ -659,7 +659,7 @@ class TestSchedulerJob: ti1.queued_by_job_id = -1 session.commit() elif bookkeeping == "still_running": - executor.running.add(retiring_key) + executor.running.add(retiring_id) if missing_definition == "dag": mocker.patch.object( job_runner.scheduler_dag_bag, "get_dag_for_run", autospec=True, return_value=None @@ -670,7 +670,7 @@ class TestSchedulerJob: ) # Simulate executor reporting task completion (this triggers the bug scenario) - executor.event_buffer[retiring_key] = executor_state, None + executor.event_buffer[retiring_id] = executor_state, None # Process the executor event with caplog.at_level(logging.INFO, logger="airflow.jobs.scheduler_job_runner"): @@ -701,7 +701,7 @@ class TestSchedulerJob: assert (history.task_instance_id, history.try_number, history.state) == (retiring_id, 4, State.FAILED) assert history.end_date is not None replacement_id = ti1.id - executor.event_buffer[retiring_key] = executor_state, None + executor.event_buffer[retiring_id] = executor_state, None job_runner._process_executor_events(executor=executor, session=session) ti1.refresh_from_db(session=session) assert (ti1.id, ti1.try_number, ti1.state) == (replacement_id, 5, None) @@ -715,7 +715,7 @@ class TestSchedulerJob: old_id = ti.id session.commit() executor = MockExecutor(do_update=False) - executor.event_buffer[ti.key] = event_state, "worker" + executor.event_buffer[ti.id] = event_state, "worker" runner = SchedulerJobRunner(Job(), executors=[executor]) runner._process_executor_events(executor=executor, session=session) @@ -735,20 +735,29 @@ class TestSchedulerJob: @pytest.mark.parametrize("event_state", [State.QUEUED, State.RUNNING, State.SUCCESS, State.FAILED]) @pytest.mark.parametrize("include_current", [False, True]) + @pytest.mark.parametrize("uuid_keys", [False, True]) def test_retired_attempt_events_do_not_modify_replacement( - self, dag_maker, session, event_state, include_current + self, dag_maker, session, event_state, include_current, uuid_keys, monkeypatch ): with dag_maker(dag_id="retired_attempt_event"): EmptyOperator(task_id="task") ti = dag_maker.create_dagrun().get_task_instance("task", session=session) - ti.try_number = 2 + ti.try_number = 1 + ti.state = State.RUNNING + retired_id = ti.id + monkeypatch.setattr(MockExecutor, "supports_task_instance_uuid", uuid_keys) + executor = MockExecutor(do_update=False) + executor.register_task(ti) + retired_key = executor.get_task_key(ti) + clear_task_instances([ti], session=session) + ti.complete_restart(session=session) ti.state = State.QUEUED ti.external_executor_id = "replacement" session.flush() - executor = MockExecutor(do_update=False) if include_current: - executor.event_buffer[ti.key] = State.RUNNING, "current_worker" - executor.event_buffer[ti.key.with_try_number(1)] = event_state, "retired_worker" + executor.register_task(ti) + executor.event_buffer[executor.get_task_key(ti)] = State.RUNNING, "current_worker" + executor.event_buffer[retired_key] = event_state, "retired_worker" runner = SchedulerJobRunner(Job(), executors=[executor]) runner._process_executor_events(executor=executor, session=session) @@ -756,6 +765,7 @@ class TestSchedulerJob: session.refresh(ti) assert ti.state == State.QUEUED + assert ti.id != retired_id assert ti.try_number == 2 assert ti.external_executor_id == ("current_worker" if include_current else "replacement") @@ -793,7 +803,7 @@ class TestSchedulerJob: session.merge(ti1) session.commit() - executor.event_buffer[ti1.key] = State.FAILED, None + executor.event_buffer[ti1.id] = State.FAILED, None self.job_runner._process_executor_events(executor=executor, session=session) ti1.refresh_from_db(session=session) @@ -804,7 +814,7 @@ class TestSchedulerJob: ti1.state = State.SUCCESS session.merge(ti1) session.commit() - executor.event_buffer[ti1.key] = State.SUCCESS, None + executor.event_buffer[ti1.id] = State.SUCCESS, None self.job_runner._process_executor_events(executor=executor, session=session) ti1.refresh_from_db(session=session) @@ -894,7 +904,7 @@ class TestSchedulerJob: session.merge(ti1) session.commit() - executor.event_buffer[ti1.key] = State.FAILED, None + executor.event_buffer[ti1.id] = State.FAILED, None self.job_runner._process_executor_events(executor=executor, session=session) ti1.refresh_from_db() @@ -975,7 +985,7 @@ class TestSchedulerJob: session.merge(ti1) session.commit() - executor.event_buffer[ti1.key] = State.FAILED, None + executor.event_buffer[ti1.id] = State.FAILED, None self.job_runner._process_executor_events(executor=executor, session=session) ti1.refresh_from_db() assert ti1.state == State.FAILED @@ -1012,14 +1022,13 @@ class TestSchedulerJob: session.merge(ti1) session.commit() - executor.event_buffer[ti1.key.with_try_number(1)] = State.SUCCESS, None + executor.event_buffer[uuid4()] = State.SUCCESS, None with caplog.at_level(logging.INFO, logger="airflow.jobs.scheduler_job_runner"): self.job_runner._process_executor_events(executor=executor, session=session) ti1.refresh_from_db(session=session) assert ti1.state == State.QUEUED self.job_runner.executor.callback_sink.send.assert_not_called() - assert any("Ignoring executor event for a different attempt" in rec.message for rec in caplog.records) # ti is queued by another scheduler - do not fail it ti1.state = State.QUEUED @@ -1027,7 +1036,7 @@ class TestSchedulerJob: session.merge(ti1) session.commit() - executor.event_buffer[ti1.key] = State.SUCCESS, None + executor.event_buffer[ti1.id] = State.SUCCESS, None self.job_runner._process_executor_events(executor=executor, session=session) ti1.refresh_from_db(session=session) @@ -1040,7 +1049,7 @@ class TestSchedulerJob: session.merge(ti1) session.commit() - executor.event_buffer[ti1.key] = State.SUCCESS, None + executor.event_buffer[ti1.id] = State.SUCCESS, None executor.has_task = mock.MagicMock(return_value=True) self.job_runner._process_executor_events(executor=executor, session=session) @@ -1086,7 +1095,7 @@ class TestSchedulerJob: session.merge(ti1) session.commit() - executor.event_buffer[ti1.key] = State.SUCCESS, None + executor.event_buffer[ti1.id] = State.SUCCESS, None executor.has_task = mock.MagicMock(return_value=False) mock_stats.incr.reset_mock() @@ -1103,7 +1112,7 @@ class TestSchedulerJob: session.merge(ti1) session.commit() - executor.event_buffer[ti1.key] = State.SUCCESS, None + executor.event_buffer[ti1.id] = State.SUCCESS, None mock_stats.incr.reset_mock() self.job_runner._process_executor_events(executor=executor, session=session) @@ -1148,7 +1157,7 @@ class TestSchedulerJob: session.merge(ti1) session.commit() - executor.event_buffer[ti1.key] = State.SUCCESS, None + executor.event_buffer[ti1.id] = State.SUCCESS, None executor.has_task = mock.MagicMock(return_value=False) mock_stats.incr.reset_mock() @@ -1165,7 +1174,7 @@ class TestSchedulerJob: session.merge(ti1) session.commit() - executor.event_buffer[ti1.key] = State.SUCCESS, None + executor.event_buffer[ti1.id] = State.SUCCESS, None mock_stats.incr.reset_mock() self.job_runner._process_executor_events(executor=executor, session=session) @@ -1199,13 +1208,12 @@ class TestSchedulerJob: session.merge(ti) session.commit() - executor.event_buffer[ti.key.with_try_number(1)] = State.RUNNING, "first_executor_id" - executor.event_buffer[ti.key.with_try_number(2)] = State.RUNNING, "second_executor_id" + executor.event_buffer[uuid4()] = State.RUNNING, "first_executor_id" + executor.event_buffer[ti.id] = State.RUNNING, "second_executor_id" with caplog.at_level(logging.INFO, logger="airflow.jobs.scheduler_job_runner"): self.job_runner._process_executor_events(executor=executor, session=session) - assert any("Ignoring executor event for a different attempt" in rec.message for rec in caplog.records) session.flush() ti.refresh_from_db(session=session) assert ti.external_executor_id == "second_executor_id" @@ -1266,7 +1274,7 @@ class TestSchedulerJob: session.merge(ti1) session.commit() - executor.event_buffer[ti1.key] = State.FAILED, None + executor.event_buffer[ti1.id] = State.FAILED, None # This should not raise DetachedInstanceError self.job_runner._process_executor_events(executor=executor, session=session) @@ -1338,7 +1346,7 @@ class TestSchedulerJob: session.merge(ti1) session.commit() - executor.event_buffer[ti1.key] = State.FAILED, None + executor.event_buffer[ti1.id] = State.FAILED, None callback = self.job_runner._schedule_dag_run(dr, session) session.flush() @@ -1464,7 +1472,7 @@ class TestSchedulerJob: session.merge(ti) session.flush() - executor.running.add(ti.key) + executor.running.add(ti.id) tis_without_heartbeats = self.job_runner._find_task_instances_without_heartbeats(session=session) assert len(tis_without_heartbeats) == 1 @@ -1589,7 +1597,7 @@ class TestSchedulerJob: self.job_runner._execute() for executor in self.job_runner.executors: - executor.get_event_buffer.assert_called_once() + executor.get_event_buffer_with_task_ids.assert_called_once() @patch("traceback.extract_stack") def test_executor_debug_dump(self, patch_traceback_extract_stack, mock_executors): @@ -5656,7 +5664,7 @@ class TestSchedulerJob: old_ti_id = ti.id executor = MockExecutor(do_update=False) - executor.event_buffer[ti.key] = TaskInstanceState.FAILED, None + executor.event_buffer[ti.id] = TaskInstanceState.FAILED, None scheduler_job = Job() self.job_runner = SchedulerJobRunner(job=scheduler_job, executors=[executor]) @@ -8731,9 +8739,9 @@ class TestSchedulerJob: ti.queued_by_job_id = scheduler_job.id session.flush() - executor.running.add(ti.key) # The executor normally does this during heartbeat. + executor.running.add(ti.id) # The executor normally does this during heartbeat. self.job_runner._find_and_purge_task_instances_without_heartbeats() - assert ti.key not in executor.running + assert ti.id not in executor.running executor.callback_sink.send.assert_called_once() callback_requests = executor.callback_sink.send.call_args.args @@ -9747,7 +9755,7 @@ class TestSchedulerJob: session.commit() # Executor reports task finished (FAILED) while TI still QUEUED -> external kill path - executor.event_buffer[ti.key] = State.FAILED, None + executor.event_buffer[ti.id] = State.FAILED, None self.job_runner._process_executor_events(executor=executor, session=session) @@ -9790,7 +9798,7 @@ class TestSchedulerJob: executor = MockExecutor(do_update=False) scheduler_job = Job() self.job_runner = SchedulerJobRunner(scheduler_job, executors=[executor]) - executor.event_buffer[ti.key] = State.FAILED, None + executor.event_buffer[ti.id] = State.FAILED, None self.job_runner._process_executor_events(executor=executor, session=session) @@ -9873,8 +9881,18 @@ class TestSchedulerJob: ("retries", "cleared_try", "next_try", "expected_max_tries"), [(0, 1, 2, 1), (2, 3, 4, 5)], ) + @pytest.mark.parametrize("uuid_keys", [False, True]) def test_heartbeat_timeout_completes_clear_with_fresh_retry_budget( - self, dag_maker, session, mocker, retries, cleared_try, next_try, expected_max_tries + self, + dag_maker, + session, + mocker, + retries, + cleared_try, + next_try, + expected_max_tries, + uuid_keys, + monkeypatch, ): with dag_maker(dag_id="hb_timeout_cleared", session=session): EmptyOperator( @@ -9886,6 +9904,7 @@ class TestSchedulerJob: ) dag_run = dag_maker.create_dagrun(run_id="test_run", state=DagRunState.RUNNING) + monkeypatch.setattr(MockExecutor, "supports_task_instance_uuid", uuid_keys) executor = MockExecutor(do_update=False) scheduler_job = Job() self.job_runner = SchedulerJobRunner(scheduler_job, executors=[executor]) @@ -9906,7 +9925,9 @@ class TestSchedulerJob: ti = session.merge(ti) session.commit() assert ti.state == TaskInstanceState.RESTARTING - executor.running.add(old_key) + executor.register_task(ti) + executor_key = executor.get_task_key(ti) + executor.running.add(executor_key) received = [] @@ -9928,8 +9949,9 @@ class TestSchedulerJob: assert ti.try_number == next_try assert ti.id != old_id assert ti.external_executor_id is None - assert old_key not in executor.running - assert executor.event_buffer == {old_key: (TaskInstanceState.FAILED, None)} + assert executor_key not in executor.running + assert executor.event_buffer == {executor_key: (TaskInstanceState.FAILED, None)} + assert executor.get_event_buffer_with_task_ids() == {old_id: (TaskInstanceState.FAILED, None)} requests = [call.args[0] for call in executor.callback_sink.send.call_args_list] assert len(requests) == 2 callback, email = requests @@ -9961,11 +9983,11 @@ class TestSchedulerJob: ti.try_number = 3 ti.max_tries = max_tries ti.external_executor_id = "retiring" - old_id, old_key = ti.id, ti.key + old_id = ti.id session.commit() executor = MockExecutor(do_update=False) runner = SchedulerJobRunner(Job(), executors=[executor]) - executor.running.add(old_key) + executor.running.add(old_id) if missing_definition == "dag": mocker.patch.object(runner.scheduler_dag_bag, "get_dag_for_run", autospec=True, return_value=None) else: @@ -9985,8 +10007,8 @@ class TestSchedulerJob: assert ti.id != old_id assert (ti.try_number, ti.state, ti.max_tries) == (4, None, expected_max_tries) assert ti.external_executor_id is None - assert old_key not in executor.running - assert executor.event_buffer == {old_key: (State.FAILED, None)} + assert old_id not in executor.running + assert executor.event_buffer == {old_id: (State.FAILED, None)} history = session.scalars( select(TaskInstanceHistory).where(TaskInstanceHistory.dag_id == ti.dag_id) @@ -10170,8 +10192,8 @@ class TestSchedulerJob: self.job_runner = SchedulerJobRunner(scheduler_job, executors=[mock_executor]) # Simulate executor reporting task as failed - executor_event = {ti.key: (TaskInstanceState.FAILED, None)} - mock_executor.get_event_buffer.return_value = executor_event + executor_event = {ti.id: (TaskInstanceState.FAILED, None)} + mock_executor.get_event_buffer_with_task_ids.return_value = executor_event # Process the executor events self.job_runner._process_executor_events(mock_executor, session) @@ -14578,8 +14600,8 @@ class TestSchedulerObservabilityMetrics: def test_executor_events_batch_metrics_emitted_on_success(self): """batch_size gauge and processed counter are emitted via the early-return path.""" # Empty event buffer → tis_with_right_state is empty → early return with num_events=0 - mock_executor = MagicMock() - mock_executor.get_event_buffer.return_value = {} + mock_executor = MagicMock(spec=BaseExecutor) + mock_executor.get_event_buffer_with_task_ids.return_value = {} with mock.patch("airflow.jobs.scheduler_job_runner.stats") as mock_stats: result = SchedulerJobRunner.process_executor_events( diff --git a/devel-common/src/tests_common/test_utils/mock_executor.py b/devel-common/src/tests_common/test_utils/mock_executor.py index 3ff7fcf8409..743778e4f86 100644 --- a/devel-common/src/tests_common/test_utils/mock_executor.py +++ b/devel-common/src/tests_common/test_utils/mock_executor.py @@ -43,6 +43,7 @@ class MockExecutor(BaseExecutor): """TestExecutor is used for unit testing purposes.""" supports_pickling = False + supports_task_instance_uuid = True mock_module_path = "mock.executor.path" mock_alias = "mock_executor" @@ -94,7 +95,7 @@ class MockExecutor(BaseExecutor): # for tests! def sort_by(item): key, workload = item - (dag_id, task_id, date, try_number, map_index) = key + (dag_id, task_id, date, try_number, map_index) = workload.ti.key # For workloads, use the task instance priority if available prio = getattr(workload.ti, "priority_weight", 1) if hasattr(workload, "ti") else 1 # Sort by priority (DESC), then date,task, try @@ -104,7 +105,7 @@ class MockExecutor(BaseExecutor): sorted_queue = sorted(task_queue.items(), key=sort_by) for key, workload in sorted_queue[:open_slots]: task_queue.pop(key) - state = self.mock_task_results[key] + state = self.mock_task_results[workload.ti.key] ti = TaskInstance.get_task_instance( task_id=workload.ti.task_id, run_id=workload.ti.run_id, diff --git a/providers/cncf/kubernetes/tests/unit/cncf/kubernetes/executors/test_kubernetes_executor.py b/providers/cncf/kubernetes/tests/unit/cncf/kubernetes/executors/test_kubernetes_executor.py index 364eeecb9c9..bfe2b94400c 100644 --- a/providers/cncf/kubernetes/tests/unit/cncf/kubernetes/executors/test_kubernetes_executor.py +++ b/providers/cncf/kubernetes/tests/unit/cncf/kubernetes/executors/test_kubernetes_executor.py @@ -3186,6 +3186,8 @@ class TestKubernetesExecutor: executor.kube_client = mock_kube_client executor.kube_scheduler = mock.MagicMock() ti.refresh_from_db() + if hasattr(executor, "register_task"): + executor.register_task(ti) executor.running.add(ti.key) # so we can verify it gets removed after revoke assert executor.has_task(task_instance=ti) executor.revoke_task(ti=ti)
