diff --git a/airflow-core/docs/migrations-ref.rst b/airflow-core/docs/migrations-ref.rst index b2eec9366d6ae..900ab447c50af 100644 --- a/airflow-core/docs/migrations-ref.rst +++ b/airflow-core/docs/migrations-ref.rst @@ -39,7 +39,10 @@ Here's the list of all the Database Migrations that are executed via when you ru +-------------------------+------------------+-------------------+--------------------------------------------------------------+ | Revision ID | Revises ID | Airflow Version | Description | +=========================+==================+===================+==============================================================+ -| ``b6a9c2e7d410`` (head) | ``f8c2a1d94e03`` | ``3.4.0`` | Add draining state to DagModel. | +| ``c7d4e8f1a203`` (head) | ``b6a9c2e7d410`` | ``3.4.0`` | Add workload_run_id to task_instance and | +| | | | task_instance_history. | ++-------------------------+------------------+-------------------+--------------------------------------------------------------+ +| ``b6a9c2e7d410`` | ``f8c2a1d94e03`` | ``3.4.0`` | Add draining state to DagModel. | +-------------------------+------------------+-------------------+--------------------------------------------------------------+ | ``f8c2a1d94e03`` | ``8d3f1a6b2c47`` | ``3.4.0`` | Add team_name and bundle_names scope columns to job table. | +-------------------------+------------------+-------------------+--------------------------------------------------------------+ diff --git a/airflow-core/src/airflow/executors/base_executor.py b/airflow-core/src/airflow/executors/base_executor.py index 991365f17d8a1..cb61f4afb6434 100644 --- a/airflow-core/src/airflow/executors/base_executor.py +++ b/airflow-core/src/airflow/executors/base_executor.py @@ -87,8 +87,8 @@ def get_execution_api_server_url(conf_source: AirflowConfigParser | ExecutorConf from airflow.models.taskinstance import TaskInstance # Event_buffer dict value type - # Tuple of: state, info - EventBufferValueType = tuple[str | None, Any] + # Tuple of: state, info, workload_run_id (optional; None when unknown/legacy) + EventBufferValueType = tuple[Any, Any, str | None] log = logging.getLogger(__name__) @@ -309,6 +309,10 @@ def __init__(self, parallelism: int = PARALLELISM, team_name: str | None = None) dict ) self.running: set[WorkloadKey] = set() + # FIFO of workload_run_id values per TaskInstanceKey, appended on queue and + # consumed when a terminal SUCCESS/FAILED event is recorded. Survives a + # re-enqueue of the same key so a stale defer-exit SUCCESS keeps its old id. + self._workload_run_ids: dict[TaskInstanceKey, deque[str]] = defaultdict(deque) self.event_buffer: dict[WorkloadKey, EventBufferValueType] = {} self._task_event_logs: deque[Log] = deque() self.conf = ExecutorConf(team_name) @@ -388,14 +392,90 @@ def log_task_event(self, *, event: str, extra: str, ti_key: WorkloadKey): self._task_event_logs.append(Log(event=event, task_instance=ti_key, extra=extra)) def queue_workload(self, workload: ExecutorWorkload, session: Session) -> None: + + if isinstance(workload, workloads.ExecuteTask): + ti = workload.ti + self.queued_tasks[ti.key] = workload + run_id = getattr(ti, "workload_run_id", None) + if isinstance(run_id, str): + self._workload_run_ids[ti.key].append(run_id) + elif isinstance(workload, workloads.ExecuteCallback): + if not self.supports_callbacks: + raise NotImplementedError( + f"{type(self).__name__} does not support ExecuteCallback workloads. " + f"Set supports_callbacks = True and implement callback handling in _process_workloads(). " + f"See LocalExecutor or CeleryExecutor for reference implementation." + ) + self.queued_callbacks[workload.key] = workload + elif isinstance(workload, workloads.TestConnection): + if not self.supports_connection_test: + raise NotImplementedError( + f"{type(self).__name__} does not support TestConnection workloads. " + f"Set supports_connection_test = True and implement connection test handling " + f"in _process_workloads(). See LocalExecutor for reference implementation." + ) + self.queued_connection_tests[workload.key] = workload + else: + raise ValueError( + f"Un-handled workload type {type(workload).__name__!r} in {type(self).__name__}. " + f"Workload must be one of: ExecuteTask, ExecuteCallback, TestConnection." + if workload.type not in self.supported_workload_types: raise NotImplementedError( f"{type(self).__name__} does not support {workload.type.value} workloads. " 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 + def _pop_workload_run_id(self, key: WorkloadKey) -> str | None: + """Consume the oldest queued workload_run_id for a task key, if any.""" + if not isinstance(key, TaskInstanceKey): + return None + pending = self._workload_run_ids.get(key) + if not pending: + return None + run_id = pending.popleft() + if not pending: + del self._workload_run_ids[key] + return run_id + + @staticmethod + def unpack_event(value: EventBufferValueType | tuple[Any, ...]) -> tuple[Any, Any, str | None]: + """Normalize event buffer values to ``(state, info, workload_run_id)``.""" + if len(value) >= 3: + return value[0], value[1], value[2] + return value[0], value[1], None + + def record_event( + self, + key: WorkloadKey, + state: WorkloadState, + info=None, + *, + workload_run_id: str | None = None, + consume_run_id: bool | None = None, + ) -> None: + """ + Write an executor event, optionally attaching a workload_run_id. + + Terminal task SUCCESS/FAILED events consume the next queued run id when + one was not supplied explicitly. Intermediate states (QUEUED/RUNNING) + must not consume it — those events are drained separately. + """ + from airflow.utils.state import TaskInstanceState + + if consume_run_id is None: + consume_run_id = isinstance(key, TaskInstanceKey) and state in ( + TaskInstanceState.SUCCESS, + TaskInstanceState.FAILED, + ) + popped_run_id = self._pop_workload_run_id(key) if consume_run_id else None + if workload_run_id is None: + workload_run_id = popped_run_id + self.event_buffer[key] = (state, info, workload_run_id) + def _get_workloads_to_schedule(self, open_slots: int) -> list[tuple[WorkloadKey, ExecutorWorkload]]: """ Select and return the next batch of workloads to schedule, respecting priority policy. @@ -571,7 +651,14 @@ def order_queued_tasks_by_priority(self) -> list: # TODO: This should not be using `TaskInstanceState` here, this is just "did the process complete, or did # it die". It is possible for the task itself to finish with success, but the state of the task to be set # to FAILED. By using TaskInstanceState enum here it confuses matters! - def change_state(self, key: WorkloadKey, state: WorkloadState, info=None, remove_running=True) -> None: + def change_state( + self, + key: WorkloadKey, + state: WorkloadState, + info=None, + remove_running=True, + workload_run_id: str | None = None, + ) -> None: """ Change state of the task. @@ -579,6 +666,8 @@ def change_state(self, key: WorkloadKey, state: WorkloadState, info=None, remove :param state: State to set for the task. :param info: Executor information for the task instance :param remove_running: Whether or not to remove the TI key from running set + :param workload_run_id: Invocation id for the finished worker; when omitted for + terminal task events, the next id queued for this key is consumed. """ self.log.debug("Changing state: %s", key) if remove_running: @@ -586,25 +675,27 @@ def change_state(self, key: WorkloadKey, state: WorkloadState, info=None, remove self.running.remove(key) except KeyError: self.log.debug("Could not find key: %s", key) - self.event_buffer[key] = state, info + self.record_event(key, state, info, workload_run_id=workload_run_id) - def fail(self, key: WorkloadKey, info=None) -> None: + def fail(self, key: WorkloadKey, info=None, workload_run_id: str | None = None) -> None: """ Set fail state for the event. :param info: Executor information for the task instance :param key: Unique key for the task instance + :param workload_run_id: Invocation id for the finished worker, if known """ - self.change_state(key, state_class_for_key(key).FAILED, info) + self.change_state(key, state_class_for_key(key).FAILED, info, workload_run_id=workload_run_id) - def success(self, key: WorkloadKey, info=None) -> None: + def success(self, key: WorkloadKey, info=None, workload_run_id: str | None = None) -> None: """ Set success state for the event. :param info: Executor information for the task instance :param key: Unique key for the task instance + :param workload_run_id: Invocation id for the finished worker, if known """ - self.change_state(key, state_class_for_key(key).SUCCESS, info) + self.change_state(key, state_class_for_key(key).SUCCESS, info, workload_run_id=workload_run_id) def queued(self, key: WorkloadKey, info=None) -> None: """ diff --git a/airflow-core/src/airflow/executors/local_executor.py b/airflow-core/src/airflow/executors/local_executor.py index eb43093ffd7e2..a4efca3682d86 100644 --- a/airflow-core/src/airflow/executors/local_executor.py +++ b/airflow-core/src/airflow/executors/local_executor.py @@ -98,7 +98,7 @@ def _run_worker( unread_messages.value -= 1 if workload.running_state is not None: - output.put((workload.key, workload.running_state, None)) + output.put((workload.key, workload.running_state, None, None)) try: BaseExecutor.run_workload( @@ -107,10 +107,12 @@ 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)) + run_id = getattr(getattr(workload, "ti", None), "workload_run_id", None) + output.put((workload.key, workload.success_state, None, run_id)) except Exception as e: log.exception("Workload execution failed.", workload_type=type(workload).__name__) - output.put((workload.key, workload.failure_state, e)) + run_id = getattr(getattr(workload, "ti", None), "workload_run_id", None) + output.put((workload.key, workload.failure_state, e, run_id)) class LocalExecutor(BaseExecutor): @@ -250,8 +252,15 @@ def sync(self) -> None: def _read_results(self): try: while not self.result_queue.empty(): - key, state, exc = self.result_queue.get() - self.change_state(key, state) + result = self.result_queue.get() + # Support legacy 3-tuples (key, state, exc) and current 4-tuples + # (key, state, exc, workload_run_id) from workers / tests. + if len(result) == 3: + key, state, _exc = result + workload_run_id = None + else: + key, state, _exc, workload_run_id = result + self.change_state(key, state, workload_run_id=workload_run_id) except (OSError, EOFError): self.log.exception("Error reading from result queue") diff --git a/airflow-core/src/airflow/executors/workloads/task.py b/airflow-core/src/airflow/executors/workloads/task.py index 56c6972403885..7044a5cf6602c 100644 --- a/airflow-core/src/airflow/executors/workloads/task.py +++ b/airflow-core/src/airflow/executors/workloads/task.py @@ -46,6 +46,7 @@ class TaskInstanceDTO(TaskInstance): priority_weight: int external_executor_id: str | None = Field(default=None, exclude=True) + workload_run_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? diff --git a/airflow-core/src/airflow/executors/workloads/types.py b/airflow-core/src/airflow/executors/workloads/types.py index 627f28ffb3fed..d00d5d89fad8a 100644 --- a/airflow-core/src/airflow/executors/workloads/types.py +++ b/airflow-core/src/airflow/executors/workloads/types.py @@ -32,7 +32,8 @@ if TYPE_CHECKING: # Type alias for executor workload results (used by executor implementations) - WorkloadResultType: TypeAlias = tuple[WorkloadKey, WorkloadState, Exception | None] + # key, state, exception, optional workload_run_id from the finished invocation + WorkloadResultType: TypeAlias = tuple[WorkloadKey, WorkloadState, Exception | None, str | None] # Type alias for scheduler workloads (ORM models that can be routed to executors) # Must be outside TYPE_CHECKING for use in function signatures diff --git a/airflow-core/src/airflow/jobs/scheduler_job_runner.py b/airflow-core/src/airflow/jobs/scheduler_job_runner.py index ad78a4fbcedcc..dd4de97fd8e42 100644 --- a/airflow-core/src/airflow/jobs/scheduler_job_runner.py +++ b/airflow-core/src/airflow/jobs/scheduler_job_runner.py @@ -1081,6 +1081,9 @@ def _executable_task_instances_to_queued(self, max_tis: int, session: Session) - "state": TaskInstanceState.QUEUED, "queued_dttm": timezone.utcnow(), "queued_by_job_id": self.job.id, + # Per-invocation id so a stale executor SUCCESS from a previous worker + # (e.g. defer exit) cannot be matched to a later enqueue of the same key. + "workload_run_id": random_db_uuid(), } # Pre-assign external_executor_id atomically with the QUEUED state so it @@ -1116,25 +1119,26 @@ def _executable_task_instances_to_queued(self, max_tis: int, session: Session) - .execution_options(synchronize_session=False) ) + # Always read workload_run_id (and external_executor_id when pre-assigned) + # back onto in-memory objects so ExecuteTask.make carries them through + # make_transient. Use RETURNING on PostgreSQL; SELECT elsewhere. + returning_cols = [TI.id, TI.workload_run_id] if pre_assign_executors: - # Read the DB-generated UUIDs back onto the in-memory objects so the - # workload DTO carries them through to send_workload_to_executor (the - # objects are about to be detached by make_transient). Use RETURNING - # where supported (PostgreSQL); fall back to a SELECT for MySQL and - # SQLite (RETURNING requires SQLite 3.35+ which isn't guaranteed). - if get_dialect_name(session) == "postgresql": - result = session.execute(queued_update.returning(TI.id, TI.external_executor_id)) - id_map = {row[0]: row[1] for row in result} - else: - session.execute(queued_update) - id_rows = session.execute( - select(TI.id, TI.external_executor_id).where(filter_for_tis) - ).all() - id_map = {row[0]: row[1] for row in id_rows} - for ti in executable_tis: - ti.external_executor_id = id_map.get(ti.id) + returning_cols.append(TI.external_executor_id) + + if get_dialect_name(session) == "postgresql": + result = session.execute(queued_update.returning(*returning_cols)) + rows = list(result) else: session.execute(queued_update) + rows = list(session.execute(select(*returning_cols).where(filter_for_tis)).all()) + + workload_run_id_map = {row[0]: row[1] for row in rows} + external_id_map = {row[0]: row[2] for row in rows} if pre_assign_executors else {} + for ti in executable_tis: + ti.workload_run_id = workload_run_id_map.get(ti.id) + if pre_assign_executors: + ti.external_executor_id = external_id_map.get(ti.id) for ti in executable_tis: ti.emit_state_change_metric(TaskInstanceState.QUEUED) @@ -1412,8 +1416,14 @@ def process_executor_events( tis_with_right_state: list[TaskInstanceKey] = [] callback_keys_with_events: list[CallbackKey] = [] + def _unpack_event(value: tuple) -> tuple[Any, Any, str | None]: + if len(value) >= 3: + return value[0], value[1], value[2] + return value[0], value[1], None + # Report execution - handle both task and callback events - for key, (state, _) in event_buffer.items(): + for key, event_value in event_buffer.items(): + state, _, _ = _unpack_event(event_value) if isinstance(key, TaskInstanceKey): existing_try = ti_primary_key_to_try_number_map.get(key.primary) if existing_try is not None and existing_try != key.try_number: @@ -1447,7 +1457,7 @@ def process_executor_events( # Handle callback state events for callback_id in callback_keys_with_events: - state, info = event_buffer.pop(callback_id) + state, info, _ = _unpack_event(event_buffer.pop(callback_id)) callback = session.get(Callback, UUID(str(callback_id))) if not callback: # This should not normally happen - we just received an event for this callback. @@ -1510,7 +1520,7 @@ def process_executor_events( ti.state, job_id, ) - state, info = event_buffer.pop(buffer_key) + state, info, event_workload_run_id = _unpack_event(event_buffer.pop(buffer_key)) if state in (TaskInstanceState.QUEUED, TaskInstanceState.RUNNING): ti.external_executor_id = info @@ -1559,9 +1569,8 @@ def process_executor_events( # from the worker exit after defer() has not been processed yet - should not fail it. # 4) the trigger already put the TI back to queued (resume after defer) but the executor success # from the worker exit after defer() has not been processed yet - should not fail it. - - # All of this could also happen if the state is "running", - # but that is handled by the scheduler detecting task instances without heartbeats. + # 5) the resumed attempt is already RUNNING when the stale defer-exit SUCCESS arrives — + # workload_run_id mismatch identifies this without relying on next_method/state heuristics. ti_queued = ti.try_number == buffer_key.try_number and ti.state in ( TaskInstanceState.SCHEDULED, @@ -1569,12 +1578,18 @@ def process_executor_events( TaskInstanceState.RUNNING, TaskInstanceState.RESTARTING, ) + stale_workload_run = ( + event_workload_run_id is not None + and ti.workload_run_id is not None + and event_workload_run_id != ti.workload_run_id + ) ti_requeued = ( ti.queued_by_job_id != job_id # Another scheduler has queued this task again or executor.has_task(ti) # This scheduler has this task already + or stale_workload_run or ( - # Resume-after-defer: trigger moved TI to scheduled or queued (next_method set) - # before we saw the executor success from the defer exit for the same try_number. + # Defense in depth for older events without workload_run_id: resume-after-defer + # while next_method is still set (SCHEDULED/QUEUED variants from #66431/#68741). ti.state in (TaskInstanceState.SCHEDULED, TaskInstanceState.QUEUED) and state == TaskInstanceState.SUCCESS and ti.next_method is not None @@ -3545,6 +3560,7 @@ def adopt_or_reset_orphaned_tasks(self, *, session: Session = NEW_SESSION) -> in ti.state = None ti.queued_by_job_id = None ti.external_executor_id = None + ti.workload_run_id = None ti.clear_next_method_args() for ti in set(tis_to_adopt_or_reset) - set(to_reset): diff --git a/airflow-core/src/airflow/migrations/versions/0135_3_4_0_add_workload_run_id_to_task_instance.py b/airflow-core/src/airflow/migrations/versions/0135_3_4_0_add_workload_run_id_to_task_instance.py new file mode 100644 index 0000000000000..ac6bb3ec9b246 --- /dev/null +++ b/airflow-core/src/airflow/migrations/versions/0135_3_4_0_add_workload_run_id_to_task_instance.py @@ -0,0 +1,56 @@ +# +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. + +""" +Add workload_run_id to task_instance and task_instance_history. + +Per-invocation UUID generated on every scheduler enqueue so executor +completion events can be matched to the attempt that produced them. +This prevents stale SUCCESS from a defer-exit worker from failing a +resumed attempt that shares the same TaskInstanceKey. + +Revision ID: c7d4e8f1a203 +Revises: b6a9c2e7d410 +Create Date: 2026-09-18 00:00:00.000000 + +""" + +from __future__ import annotations + +import sqlalchemy as sa +from alembic import op + +revision = "c7d4e8f1a203" +down_revision = "b6a9c2e7d410" +branch_labels = None +depends_on = None +airflow_version = "3.4.0" + + +def upgrade(): + """Add workload_run_id to task_instance and task_instance_history.""" + for table in ("task_instance", "task_instance_history"): + with op.batch_alter_table(table, schema=None) as batch_op: + batch_op.add_column(sa.Column("workload_run_id", sa.Text(), nullable=True)) + + +def downgrade(): + """Remove workload_run_id from task_instance and task_instance_history.""" + for table in ("task_instance", "task_instance_history"): + with op.batch_alter_table(table, schema=None) as batch_op: + batch_op.drop_column("workload_run_id") diff --git a/airflow-core/src/airflow/models/taskinstance.py b/airflow-core/src/airflow/models/taskinstance.py index d9c3f8cab9595..31e17d2b6acf5 100644 --- a/airflow-core/src/airflow/models/taskinstance.py +++ b/airflow-core/src/airflow/models/taskinstance.py @@ -443,6 +443,7 @@ def clear_task_instances( ti.max_tries = max(ti.max_tries, ti.try_number) ti.state = None ti.external_executor_id = None + ti.workload_run_id = None ti.clear_next_method_args() # Match DagVersion to latest serialized DAG when running on the latest version. if use_latest_version: @@ -660,6 +661,10 @@ class TaskInstance(Base, LoggingMixin, BaseWorkload): external_executor_id: Mapped[str | None] = mapped_column(Text(), nullable=True) + # Per-invocation id generated on every enqueue. Used to ignore stale executor + # SUCCESS events from a previous worker exit (e.g. defer) after the TI was resumed. + workload_run_id: Mapped[str | None] = mapped_column(Text(), nullable=True) + # The trigger to resume on if we are in state DEFERRED trigger_id: Mapped[int | None] = mapped_column(Integer, nullable=True) diff --git a/airflow-core/src/airflow/models/taskinstancehistory.py b/airflow-core/src/airflow/models/taskinstancehistory.py index 1f741507bd5aa..0763a1c807b1c 100644 --- a/airflow-core/src/airflow/models/taskinstancehistory.py +++ b/airflow-core/src/airflow/models/taskinstancehistory.py @@ -104,6 +104,7 @@ class TaskInstanceHistory(Base): context_carrier: Mapped[dict | None] = mapped_column(MutableDict.as_mutable(ExtendedJSON), nullable=True) external_executor_id: Mapped[str | None] = mapped_column(Text(), nullable=True) + workload_run_id: Mapped[str | None] = mapped_column(Text(), nullable=True) trigger_id: Mapped[int | None] = mapped_column(Integer, nullable=True) trigger_timeout: Mapped[DateTime | None] = mapped_column(DateTime, nullable=True) next_method: Mapped[str | None] = mapped_column(String(1000), nullable=True) diff --git a/airflow-core/src/airflow/utils/db.py b/airflow-core/src/airflow/utils/db.py index c29af4a971c47..bbcb70223ebba 100644 --- a/airflow-core/src/airflow/utils/db.py +++ b/airflow-core/src/airflow/utils/db.py @@ -117,7 +117,7 @@ class MappedClassProtocol(Protocol): "3.1.8": "509b94a1042d", "3.2.0": "1d6611b6ab7c", "3.3.0": "d2f4e1b3c5a7", - "3.4.0": "b6a9c2e7d410", + "3.4.0": "c7d4e8f1a203", } # Prefix used to identify tables holding data moved during migration. diff --git a/airflow-core/tests/unit/executors/test_base_executor.py b/airflow-core/tests/unit/executors/test_base_executor.py index 3ff2dd5a2cd4f..60c5aab8c26af 100644 --- a/airflow-core/tests/unit/executors/test_base_executor.py +++ b/airflow-core/tests/unit/executors/test_base_executor.py @@ -163,7 +163,7 @@ def test_state_methods_pick_callback_state_for_callback_key(method_name, expecte getattr(executor, method_name)(callback_key) - assert executor.event_buffer[callback_key] == (expected_state, None) + assert executor.event_buffer[callback_key] == (expected_state, None, None) def test_fail_and_success(): @@ -557,6 +557,20 @@ def test_running_retry_attempt_type(loop_duration, total_tries): assert a.tries_after_min == 1 +def test_success_consumes_queued_workload_run_id_fifo(): + """Re-enqueue of the same key keeps earlier run ids for stale events.""" + executor = BaseExecutor() + key = TaskInstanceKey("my_dag1", "my_task1", "run1", 1) + executor._workload_run_ids[key].append("run-a") + executor._workload_run_ids[key].append("run-b") + + executor.success(key) + assert executor.event_buffer[key] == (TaskInstanceState.SUCCESS, None, "run-a") + + executor.success(key) + assert executor.event_buffer[key] == (TaskInstanceState.SUCCESS, None, "run-b") + + def test_state_fail(): executor = BaseExecutor() key = TaskInstanceKey("my_dag1", "my_task1", timezone.utcnow(), 1) @@ -564,7 +578,7 @@ def test_state_fail(): info = "info" executor.fail(key, info=info) assert not executor.running - assert executor.event_buffer[key] == (TaskInstanceState.FAILED, info) + assert executor.event_buffer[key] == (TaskInstanceState.FAILED, info, None) def test_state_success(): @@ -574,7 +588,7 @@ def test_state_success(): info = "info" executor.success(key, info=info) assert not executor.running - assert executor.event_buffer[key] == (TaskInstanceState.SUCCESS, info) + assert executor.event_buffer[key] == (TaskInstanceState.SUCCESS, info, None) def test_state_queued(): @@ -584,7 +598,7 @@ def test_state_queued(): info = "info" executor.queued(key, info=info) assert not executor.running - assert executor.event_buffer[key] == (TaskInstanceState.QUEUED, info) + assert executor.event_buffer[key] == (TaskInstanceState.QUEUED, info, None) def test_state_running(): @@ -595,7 +609,7 @@ def test_state_running(): executor.running_state(key, info=info) # Running state should not remove a command as running assert executor.running - assert executor.event_buffer[key] == (TaskInstanceState.RUNNING, info) + assert executor.event_buffer[key] == (TaskInstanceState.RUNNING, info, None) def test_repr(): diff --git a/airflow-core/tests/unit/executors/test_local_executor.py b/airflow-core/tests/unit/executors/test_local_executor.py index 28adefd22aaf4..6cb2692d5cfff 100644 --- a/airflow-core/tests/unit/executors/test_local_executor.py +++ b/airflow-core/tests/unit/executors/test_local_executor.py @@ -96,7 +96,8 @@ 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") - result_queue.put((key, State.SUCCESS, payload)) + # (key, state, exc, workload_run_id) — matches LocalExecutor worker output + result_queue.put((key, State.SUCCESS, payload, None)) class TestLocalExecutor: diff --git a/airflow-core/tests/unit/executors/test_workloads.py b/airflow-core/tests/unit/executors/test_workloads.py index cd31ea833c5cb..64be99ad43a0c 100644 --- a/airflow-core/tests/unit/executors/test_workloads.py +++ b/airflow-core/tests/unit/executors/test_workloads.py @@ -225,6 +225,7 @@ def test_workload_ti_round_trips_through_sdk_generated_model(): dumped = ti.model_dump(mode="json") assert "external_executor_id" not in dumped + assert "workload_run_id" not in dumped assert "executor_config" not in dumped # Executor-side scheduling fields stay on the workload wire (older workers # deserialize the workload with a model that requires them) but are not @@ -283,6 +284,7 @@ def _make_mock_ti( ti.context_carrier = None ti.hostname = None ti.external_executor_id = None + ti.workload_run_id = None ti.dag_model.bundle_name = "test-bundle" ti.dag_model.relative_fileloc = "dags/test_dag.py" diff --git a/airflow-core/tests/unit/jobs/test_scheduler_job.py b/airflow-core/tests/unit/jobs/test_scheduler_job.py index 05d6a76c7cce1..217cc98ed61ab 100644 --- a/airflow-core/tests/unit/jobs/test_scheduler_job.py +++ b/airflow-core/tests/unit/jobs/test_scheduler_job.py @@ -1037,6 +1037,67 @@ def test_process_executor_events_stale_success_when_queued_after_defer( tags={"dag_id": dag_id, "task_id": ti1.task_id}, ) + @pytest.mark.parametrize( + "ti_state", + [State.SCHEDULED, State.QUEUED, State.RUNNING], + ) + @mock.patch("airflow.jobs.scheduler_job_runner.TaskCallbackRequest") + @mock.patch("airflow._shared.observability.metrics.stats._get_backend") + def test_process_executor_events_stale_success_mismatched_workload_run_id( + self, mock_get_backend, mock_task_callback, dag_maker, ti_state + ): + """ + Stale defer-exit SUCCESS carries an older workload_run_id than the resumed attempt. + + Must not treat as state mismatch for SCHEDULED, QUEUED, or RUNNING — even when + next_method is already cleared (the RUNNING variant of #72716). + """ + mock_stats = mock.MagicMock(spec=StatsLogger) + mock_get_backend.return_value = mock_stats + dag_id = f"test_stale_success_workload_run_id_{ti_state}" + task_id_1 = "dummy_task" + + session = settings.Session() + with dag_maker(dag_id=dag_id, fileloc="/test_path1/"): + task1 = EmptyOperator(task_id=task_id_1) + ti1 = dag_maker.create_dagrun().get_task_instance(task1.task_id) + + executor = MockExecutor(do_update=False) + mock_task_callback.return_value = mock.MagicMock() + scheduler_job = Job() + session.add(scheduler_job) + session.flush() + self.job_runner = SchedulerJobRunner(scheduler_job, executors=[executor]) + + ti1.state = ti_state + ti1.next_method = None + ti1.queued_by_job_id = scheduler_job.id + ti1.try_number = 1 + ti1.workload_run_id = "current-run-id" + session.merge(ti1) + session.commit() + + # Event from the previous (defer-exit) invocation + executor.event_buffer[ti1.key] = State.SUCCESS, None, "stale-defer-run-id" + executor.has_task = mock.MagicMock(return_value=False) + mock_stats.incr.reset_mock() + + self.job_runner._process_executor_events(executor=executor, session=session) + ti1.refresh_from_db(session=session) + assert ti1.state == ti_state + self.job_runner.executor.callback_sink.send.assert_not_called() + mock_stats.incr.assert_called_once_with("scheduler.executor_events.processed", count=1) + + # Matching run id with no other requeue signal is still an external kill. + executor.event_buffer[ti1.key] = State.SUCCESS, None, "current-run-id" + mock_stats.incr.reset_mock() + + self.job_runner._process_executor_events(executor=executor, session=session) + mock_stats.incr.assert_any_call( + "scheduler.tasks.killed_externally", + tags={"dag_id": dag_id, "task_id": ti1.task_id}, + ) + @mock.patch("airflow.jobs.scheduler_job_runner.TaskCallbackRequest") @mock.patch("airflow._shared.observability.metrics.stats._get_backend") def test_process_executor_events_multiple_try_numbers_warns( @@ -3235,6 +3296,30 @@ class PreAssigningExecutor(MockExecutor): session.rollback() + def test_executable_task_instances_to_queued_sets_workload_run_id(self, dag_maker, session): + """workload_run_id is written for every TI on enqueue, for all executors.""" + dag_id = "SchedulerJobTest.test_executable_sets_workload_run_id" + session = settings.Session() + with dag_maker(dag_id=dag_id, start_date=DEFAULT_DATE, session=session): + EmptyOperator(task_id="task_a") + + self.job_runner = SchedulerJobRunner(job=Job(), executors=[MockExecutor()]) + + dr = dag_maker.create_dagrun() + ti = dr.get_task_instance("task_a", session=session) + ti.state = State.SCHEDULED + session.flush() + + returned_tis = self.job_runner._executable_task_instances_to_queued(max_tis=32, session=session) + assert len(returned_tis) == 1 + assert returned_tis[0].workload_run_id is not None + assert UUID(returned_tis[0].workload_run_id), "is valid uuid" + + db_value = session.scalar(select(TaskInstance.workload_run_id).where(TaskInstance.id == ti.id)) + assert db_value == returned_tis[0].workload_run_id + + session.rollback() + @pytest.mark.parametrize("state", [State.FAILED, State.SUCCESS]) def test_enqueue_task_instances_sets_ti_state_to_None_if_dagrun_in_finish_state(self, state, dag_maker): """This tests that task instances whose dagrun is in finished state are not queued""" diff --git a/airflow-core/tests/unit/models/test_taskinstance.py b/airflow-core/tests/unit/models/test_taskinstance.py index 7401c7781f3cc..00d1ba0de2b5d 100644 --- a/airflow-core/tests/unit/models/test_taskinstance.py +++ b/airflow-core/tests/unit/models/test_taskinstance.py @@ -2580,6 +2580,7 @@ def test_refresh_from_db(self, create_task_instance): "executor": "some_executor", "executor_config": {"Some": {"extra": "information"}}, "external_executor_id": "some_executor_id", + "workload_run_id": "some-run-id", "trigger_timeout": None, "trigger_id": None, "next_kwargs": None, 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 3ff7fcf8409e2..c3bc5ead9e7dc 100644 --- a/devel-common/src/tests_common/test_utils/mock_executor.py +++ b/devel-common/src/tests_common/test_utils/mock_executor.py @@ -122,11 +122,16 @@ def terminate(self): def end(self): self.sync() - def change_state(self, key, state, info=None, remove_running=False): - super().change_state(key, state, info=info, remove_running=remove_running) + def change_state(self, key, state, info=None, remove_running=False, workload_run_id=None): + try: + super().change_state( + key, state, info=info, remove_running=remove_running, workload_run_id=workload_run_id + ) + except TypeError: + super().change_state(key, state, info=info, remove_running=remove_running) # The normal event buffer is cleared after reading, we want to keep # a list of all events for testing - self.sorted_tasks.append((key, (state, info))) + self.sorted_tasks.append((key, (state, info, workload_run_id))) def mock_task_fail(self, dag_id, task_id, run_id: str, try_number=1): """ diff --git a/providers/celery/src/airflow/providers/celery/executors/celery_executor.py b/providers/celery/src/airflow/providers/celery/executors/celery_executor.py index ae5a824ac7a0d..71a051413f721 100644 --- a/providers/celery/src/airflow/providers/celery/executors/celery_executor.py +++ b/providers/celery/src/airflow/providers/celery/executors/celery_executor.py @@ -236,7 +236,7 @@ def _send_workloads(self, workload_tuples_to_send: Sequence[WorkloadInCelery]): self.workload_publish_retries.pop(key, None) if isinstance(result, ExceptionWithTraceback): self.log.error("%s: %s\n%s\n", CELERY_SEND_ERR_MSG_HEADER, result.exception, result.traceback) - self.event_buffer[key] = (TaskInstanceState.FAILED, None) + self._emit_task_event(key, TaskInstanceState.FAILED, None) elif result is not None: result.backend = cached_celery_backend self.running.add(key) @@ -245,7 +245,19 @@ def _send_workloads(self, workload_tuples_to_send: Sequence[WorkloadInCelery]): # Store the Celery task_id (workload execution ID) in the event buffer. This will get "overwritten" if the task # has another event, but that is fine, because the only other events are success/failed at # which point we don't need the ID anymore anyway. - self.event_buffer[key] = (TaskInstanceState.QUEUED, result.task_id) + self._emit_task_event(key, TaskInstanceState.QUEUED, result.task_id, consume_run_id=False) + + def _emit_task_event(self, key, state, info=None, *, consume_run_id: bool | None = None) -> None: + """Write an executor event; compatible with cores that lack ``record_event``.""" + record_event = getattr(self, "record_event", None) + if callable(record_event): + if consume_run_id is None: + record_event(key, state, info) + else: + record_event(key, state, info, consume_run_id=consume_run_id) + return + # Older BaseExecutor: event buffer is a plain (state, info[, workload_run_id]) tuple. + self.event_buffer[key] = (state, info, None) def _send_workloads_to_celery(self, workload_tuples_to_send: Sequence[WorkloadInCelery]): from airflow.providers.celery.executors.celery_executor_utils import send_workload_to_executor @@ -289,12 +301,25 @@ def update_all_workload_states(self) -> None: self.log.debug("Inquiries completed.") for key, async_result in list(self.workloads.items()): - state, info = state_and_info_by_celery_task_id.get(async_result.task_id) + state, info, *_ = state_and_info_by_celery_task_id.get(async_result.task_id) if state: self.update_task_state(cast("TaskInstanceKey", key), state, info) - def change_state(self, key: WorkloadKey, state: WorkloadState, info=None, remove_running=True) -> None: - super().change_state(key, state, info, remove_running=remove_running) + def change_state( + self, + key: WorkloadKey, + state: WorkloadState, + info=None, + remove_running=True, + workload_run_id: str | None = None, + ) -> None: + try: + super().change_state( + key, state, info, remove_running=remove_running, workload_run_id=workload_run_id + ) + except TypeError: + # Older BaseExecutor.change_state has no workload_run_id kwarg (provider compat matrix). + super().change_state(key, state, info, remove_running=remove_running) self.workloads.pop(key, None) def update_task_state(self, key: TaskInstanceKey, state: str, info: Any) -> None: @@ -360,7 +385,7 @@ def try_adopt_task_instances(self, tis: Sequence[TaskInstance]) -> Sequence[Task adopted = [] cached_celery_backend = next(iter(celery_tasks.values()))[0].backend - for celery_task_id, (state, info) in states_by_celery_task_id.items(): + for celery_task_id, (state, info, *_) in states_by_celery_task_id.items(): result, ti = celery_tasks[celery_task_id] result.backend = cached_celery_backend if isinstance(result.result, BaseException): diff --git a/providers/celery/src/airflow/providers/celery/executors/celery_executor_utils.py b/providers/celery/src/airflow/providers/celery/executors/celery_executor_utils.py index 39d52deb1dd13..ff4359661f709 100644 --- a/providers/celery/src/airflow/providers/celery/executors/celery_executor_utils.py +++ b/providers/celery/src/airflow/providers/celery/executors/celery_executor_utils.py @@ -556,7 +556,7 @@ def _prepare_state_and_info_by_task_dict( else: state = celery_states.PENDING info = None - state_info[task_id] = state, info + state_info[task_id] = state, info, None return state_info def _get_many_using_multiprocessing( @@ -581,5 +581,5 @@ def _get_many_using_multiprocessing( state_or_exception.traceback, ) else: - states_and_info_by_task_id[task_id] = state_or_exception, info + states_and_info_by_task_id[task_id] = state_or_exception, info, None return states_and_info_by_task_id diff --git a/providers/celery/tests/integration/celery/test_celery_executor.py b/providers/celery/tests/integration/celery/test_celery_executor.py index d5c977aa11ba2..1c05f1d307859 100644 --- a/providers/celery/tests/integration/celery/test_celery_executor.py +++ b/providers/celery/tests/integration/celery/test_celery_executor.py @@ -413,7 +413,7 @@ def test_should_support_kv_backend(self, mock_mget, caplog): assert set(mget_args[0]) == {b"celery-task-meta-456", b"celery-task-meta-123"} mock_mget.assert_called_once_with(mock.ANY) - assert result == {"123": ("SUCCESS", None), "456": ("PENDING", None)} + assert result == {"123": ("SUCCESS", None, None), "456": ("PENDING", None, None)} assert caplog.messages == ["Fetched 2 state(s) for 2 task(s)"] @mock.patch("celery.backends.database.DatabaseBackend.ResultSession") @@ -440,7 +440,7 @@ def test_should_support_db_backend(self, mock_session, caplog): ] ) - assert result == {"123": ("SUCCESS", None), "456": ("PENDING", None)} + assert result == {"123": ("SUCCESS", None, None), "456": ("PENDING", None, None)} assert caplog.messages == ["Fetched 2 state(s) for 2 task(s)"] @mock.patch("celery.backends.database.DatabaseBackend.ResultSession") @@ -474,7 +474,7 @@ def test_should_retry_db_backend(self, mock_session, caplog): ] ) assert mock_retry_db_result.call_count == 2 - assert result == {"123": ("SUCCESS", None), "456": ("PENDING", None)} + assert result == {"123": ("SUCCESS", None, None), "456": ("PENDING", None, None)} assert caplog.messages == [ "Failed operation _query_task_cls_from_db_backend. Retrying 2 more times.", "Fetched 2 state(s) for 2 task(s)", @@ -499,5 +499,5 @@ def test_should_support_base_backend(self, caplog): ] ) - assert result == {"123": ("SUCCESS", None), "456": ("PENDING", None)} + assert result == {"123": ("SUCCESS", None, None), "456": ("PENDING", None, None)} assert caplog.messages == ["Fetched 2 state(s) for 2 task(s)"] diff --git a/providers/cncf/kubernetes/src/airflow/providers/cncf/kubernetes/executors/kubernetes_executor.py b/providers/cncf/kubernetes/src/airflow/providers/cncf/kubernetes/executors/kubernetes_executor.py index 3216ed232b8dc..06e0602b666ae 100644 --- a/providers/cncf/kubernetes/src/airflow/providers/cncf/kubernetes/executors/kubernetes_executor.py +++ b/providers/cncf/kubernetes/src/airflow/providers/cncf/kubernetes/executors/kubernetes_executor.py @@ -365,11 +365,33 @@ def execute_async( queue, ) - self.event_buffer[key] = (TaskInstanceState.QUEUED, self.scheduler_job_id) + self._emit_task_event(key, TaskInstanceState.QUEUED, self.scheduler_job_id, consume_run_id=False) job = KubernetesJob(key, command, kube_executor_config, pod_template_file, coordinator_kube_image) self.pod_launch_attempts[key] = _PodLaunchAttempt(job=job) self.task_queue.put(job) + def _emit_task_event(self, key, state, info=None, *, consume_run_id: bool | None = None) -> None: + """Write an executor event; compatible with cores that lack ``record_event``.""" + record_event = getattr(self, "record_event", None) + if callable(record_event): + if consume_run_id is None: + record_event(key, state, info) + else: + record_event(key, state, info, consume_run_id=consume_run_id) + return + self.event_buffer[key] = (state, info, None) + + def queue_workload(self, workload: workloads.All, session: Session | None) -> None: + from airflow.executors import workloads + + if not isinstance(workload, workloads.ExecuteTask): + raise RuntimeError(f"{type(self)} cannot handle workloads of type {type(workload)}") + ti = workload.ti + self.queued_tasks[ti.key] = workload + workload_run_id = getattr(ti, "workload_run_id", None) + if isinstance(workload_run_id, str) and hasattr(self, "_workload_run_ids"): + self._workload_run_ids[ti.key].append(workload_run_id) + def _process_workloads(self, workloads: Sequence[workloads.All]) -> None: from airflow.executors.workloads import ExecuteTask @@ -731,7 +753,7 @@ def _change_state( return if state == TaskInstanceState.RUNNING: - self.event_buffer[key] = state, None + self._emit_task_event(key, state, None, consume_run_id=False) return if self.kube_config.delete_worker_pods: @@ -805,7 +827,7 @@ def _change_state( if state is None: state = self._get_task_instance_state(key, session=session) - self.event_buffer[key] = state, termination_reason + self._emit_task_event(key, state, termination_reason) def _get_task_instance_state(self, key: TaskInstanceKey, *, session: Session) -> TaskInstanceState | None: """Look up the current task instance state from the metadata database."""