This is an automated email from the ASF dual-hosted git repository.
kaxil pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/airflow.git
The following commit(s) were added to refs/heads/main by this push:
new 4b488e2bb01 Unify executor workload queues (#63491)
4b488e2bb01 is described below
commit 4b488e2bb01235aa2e54e5635f389fa13094a854
Author: Anish Giri <[email protected]>
AuthorDate: Mon Sep 21 14:33:07 2026 -0500
Unify executor workload queues (#63491)
Refactors `BaseExecutor` workload queue management so that adding a new
workload type is a small.
## Problem
Adding a new workload type (like ExecuteCallback or TestConnection)
required touching ~6 places in BaseExecutor: a new queue dict, a new
`supports_*` flag,slots calculation, an isinstance branch in queue_workload, a
dedicated scheduling method, and isinstance branches in dequeue/trigger logic.
Each provider executor that overrode `queue_workload` also needed updating.
This made extending the executor interface unnecessarily painful.
## What this does
Replaces the per-type queue dicts and boolean capability flags with three
primitives:
- `executor_queues`: one `defaultdict(dict)` keyed by `WorkloadType`,
replacing `queued_tasks`, `queued_callbacks` and `queued_connection_tests`.
- `supported_workload_types`: a `frozenset[WorkloadType]` replacing
`supports_callbacks` and `supports_connection_test`.
- `WORKLOAD_TYPE_PRIORITY` plus a per-schema `sort_key`: the type registry
decides ordering across types, the schema decides ordering within a type.
`queue_workload` in the base class is now generic: validate the type, store
by key. `trigger_tasks` and `trigger_connection_tests` are merged into
`trigger_workloads`, which drains all queue types through a single
per-heartbeat budget.
## Adding a new workload type after this refactor
1. Declare the `WorkloadType` member and place it in
`_workload_type_priority_order`.
2. Add the schema to the `All` and `ExecutorWorkload` discriminated unions,
and override `sort_key` if it needs ordering within its group.
3. Add a branch to `BaseExecutor.run_workload` and to `state_class_for_key`.
4. Add the type to `supported_workload_types` on supporting executors and
handle it in `_process_workloads`.
---
airflow-core/newsfragments/63491.significant.rst | 60 +++
.../src/airflow/executors/base_executor.py | 312 +++++++++------
.../src/airflow/executors/local_executor.py | 16 +-
.../src/airflow/executors/workloads/__init__.py | 4 +-
.../src/airflow/executors/workloads/base.py | 41 ++
.../src/airflow/executors/workloads/callback.py | 4 +-
.../airflow/executors/workloads/connection_test.py | 4 +-
.../src/airflow/executors/workloads/task.py | 9 +-
.../src/airflow/executors/workloads/types.py | 8 +-
.../src/airflow/jobs/scheduler_job_runner.py | 6 +-
.../tests/unit/executors/test_base_executor.py | 441 ++++++++++++++++++---
.../tests/unit/executors/test_local_executor.py | 17 +-
.../tests/unit/executors/test_workloads.py | 54 ++-
airflow-core/tests/unit/jobs/test_scheduler_job.py | 53 ++-
.../src/tests_common/test_utils/mock_executor.py | 15 +-
.../aws/executors/aws_lambda/lambda_executor.py | 46 ++-
.../amazon/aws/executors/batch/batch_executor.py | 69 ++--
.../amazon/aws/executors/ecs/ecs_executor.py | 46 ++-
.../src/airflow/providers/amazon/version_compat.py | 2 +
.../executors/aws_lambda/test_lambda_executor.py | 29 +-
.../aws/executors/batch/test_batch_executor.py | 32 +-
.../amazon/aws/executors/ecs/test_ecs_executor.py | 20 +-
.../providers/celery/executors/celery_executor.py | 34 +-
.../celery/executors/celery_kubernetes_executor.py | 5 +-
.../src/airflow/providers/celery/version_compat.py | 2 +
.../integration/celery/test_celery_executor.py | 2 +-
.../unit/celery/executors/test_celery_executor.py | 8 +-
.../kubernetes/executors/kubernetes_executor.py | 34 +-
.../executors/local_kubernetes_executor.py | 13 +-
.../providers/cncf/kubernetes/version_compat.py | 2 +
.../executors/test_kubernetes_executor.py | 42 +-
.../executors/test_local_kubernetes_executor.py | 19 +-
.../providers/edge3/executors/edge_executor.py | 9 +-
.../src/airflow/providers/edge3/version_compat.py | 2 +
.../unit/edge3/executors/test_edge_executor.py | 2 +-
35 files changed, 1110 insertions(+), 352 deletions(-)
diff --git a/airflow-core/newsfragments/63491.significant.rst
b/airflow-core/newsfragments/63491.significant.rst
new file mode 100644
index 00000000000..d81d2951d43
--- /dev/null
+++ b/airflow-core/newsfragments/63491.significant.rst
@@ -0,0 +1,60 @@
+Unify executor workload queues under ``BaseExecutor.executor_queues``
+
+Executor workload state is now stored on the unified
``BaseExecutor.executor_queues`` mapping
+keyed by ``WorkloadType``, and scheduling is driven by ``trigger_workloads``.
The previous
+per-type attributes and entrypoints ``queued_tasks``, ``queued_callbacks``,
``supports_callbacks``,
+``trigger_tasks``, and ``order_queued_tasks_by_priority`` are kept as
backward-compatible shims
+that emit ``RemovedInAirflow4Warning`` and will be removed in Airflow 4.0.
+
+The connection-test surface introduced in Airflow 3.3
(``supports_connection_test``,
+``queued_connection_tests``, and ``trigger_connection_tests()``) is folded
into the same
+mechanism: connection tests are queued in
``executor_queues[WorkloadType.TEST_CONNECTION]``
+and dispatched by ``trigger_workloads`` alongside other workload types. A
deprecated
+``supports_connection_test`` compat attribute remains;
``queued_connection_tests`` and
+``trigger_connection_tests()`` are removed without a shim.
+
+**Migration:**
+
+- Replace ``executor.queued_tasks`` with
``executor.executor_queues[WorkloadType.EXECUTE_TASK]``.
+- Replace ``executor.queued_callbacks`` with
``executor.executor_queues[WorkloadType.EXECUTE_CALLBACK]``.
+- Replace ``executor.queued_connection_tests`` with
+ ``executor.executor_queues[WorkloadType.TEST_CONNECTION]``.
+- Replace ``supports_callbacks = True`` / ``supports_connection_test = True``
class declarations
+ by adding ``WorkloadType.EXECUTE_CALLBACK`` /
``WorkloadType.TEST_CONNECTION`` to the
+ ``supported_workload_types`` frozenset.
+- Replace ``executor.trigger_tasks(open_slots)`` and
``executor.trigger_connection_tests()``
+ with ``executor.trigger_workloads(open_slots)``.
+- ``order_queued_tasks_by_priority()`` has no direct public replacement:
workload ordering is
+ applied internally by ``trigger_workloads``. The deprecated shim now
forwards to the private
+ ``_get_workloads_to_schedule(open_slots)`` helper, whose return shape
differs from the old
+ method — it includes *all* queued workload types (not just tasks) and
truncates the result
+ to the number of currently open slots. Subclasses that need a custom
ordering should not
+ depend on the private helper; open an issue describing the use case instead.
+
+**Behavioural changes:**
+
+- Workload admission now uses a single per-heartbeat budget instead of
independent budgets for
+ tasks and connection tests. Connection tests can no longer over-subscribe
``parallelism``,
+ and they are prioritized *ahead of* tasks (they are short-lived and
user-interactive), so a
+ sustained task backlog no longer starves them.
+- ``executor_queues`` is a ``collections.defaultdict`` and the
per-``WorkloadType`` inner
+ queues auto-vivify on access; do not replace it with a plain ``dict``.
+
+Legacy ``supports_callbacks`` / ``supports_connection_test`` usage on
out-of-tree executors is
+still honored, with a deprecation warning in each case:
+
+- Class attributes: ``BaseExecutor.__init_subclass__`` folds ``= True`` / ``=
False``
+ declarations into ``supported_workload_types`` (adding or removing the type
while preserving
+ everything else inherited from the parent class) and drops the plain bool so
the compat
+ attribute keeps working on the subclass.
+- Class-level reads (``SomeExecutor.supports_callbacks``) and instance reads
still return a
+ ``bool`` derived from ``supported_workload_types``.
+- Instance assignment (``self.supports_callbacks = True`` in ``__init__``)
still works and toggles
+ the corresponding entry in that instance's ``supported_workload_types``.
+
+Subclasses that override the no-longer-invoked ``trigger_tasks``,
``trigger_connection_tests``
+or ``order_queued_tasks_by_priority`` methods also get a warning at
class-definition time.
+
+Deprecation warnings for these compat entrypoints are emitted at most once per
executor class to
+avoid flooding scheduler heartbeat logs, and each is paired with a
``log.warning`` so it is
+visible to out-of-tree executors regardless of Python warning-filter
configuration.
diff --git a/airflow-core/src/airflow/executors/base_executor.py
b/airflow-core/src/airflow/executors/base_executor.py
index ef5fafded8b..991365f17d8 100644
--- a/airflow-core/src/airflow/executors/base_executor.py
+++ b/airflow-core/src/airflow/executors/base_executor.py
@@ -20,6 +20,7 @@ from __future__ import annotations
import logging
import sys
+import warnings
from collections import defaultdict, deque
from collections.abc import Sequence
from dataclasses import dataclass, field
@@ -31,8 +32,10 @@ import pendulum
from airflow._shared.observability.metrics import stats
from airflow.cli.cli_config import DefaultHelpParser
from airflow.configuration import conf
+from airflow.exceptions import RemovedInAirflow4Warning
from airflow.executors import workloads
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
@@ -80,7 +83,6 @@ if TYPE_CHECKING:
from airflow.executors.executor_utils import ExecutorName
from airflow.executors.workloads import ExecutorWorkload
from airflow.executors.workloads.types import WorkloadKey, WorkloadState
- from airflow.models.callback import CallbackKey
from airflow.models.connection_test import ConnectionTestKey
from airflow.models.taskinstance import TaskInstance
@@ -92,6 +94,39 @@ if TYPE_CHECKING:
log = logging.getLogger(__name__)
+def _warn_deprecated_executor_usage(message: str) -> None:
+ warnings.warn(message, RemovedInAirflow4Warning, stacklevel=3)
+ log.warning(message)
+
+
+class _LegacyWorkloadFlag:
+ """Deprecated bool accessor for one ``supported_workload_types`` member,
usable on the class and on instances."""
+
+ def __init__(self, workload_type: WorkloadType) -> None:
+ self.workload_type = workload_type
+
+ def __set_name__(self, owner: type, name: str) -> None:
+ self.name = name
+
+ @property
+ def deprecation_message(self) -> str:
+ return (
+ f"{self.name} is deprecated. "
+ f"Use WorkloadType.{self.workload_type.name} in
supported_workload_types instead."
+ )
+
+ def __get__(self, obj: BaseExecutor | None, owner: type[BaseExecutor] |
None = None) -> bool:
+ target = obj if obj is not None else owner
+ if target is None:
+ raise TypeError(f"{self.name} must be accessed through a
BaseExecutor class or instance")
+ target._warn_legacy_property(self.name, self.deprecation_message)
+ return self.workload_type in target.supported_workload_types
+
+ def __set__(self, obj: BaseExecutor, value: bool) -> None:
+ obj._warn_legacy_property(self.name, self.deprecation_message)
+ obj._set_workload_type_supported(self.workload_type, value)
+
+
@dataclass
class RunningRetryAttemptType:
"""
@@ -170,13 +205,14 @@ class BaseExecutor(LoggingMixin):
"""
supports_ad_hoc_ti_run: bool = False
- supports_callbacks: bool = False
+ supported_workload_types: frozenset[WorkloadType] =
frozenset({WorkloadType.EXECUTE_TASK})
+ supports_callbacks = _LegacyWorkloadFlag(WorkloadType.EXECUTE_CALLBACK)
+ supports_connection_test =
_LegacyWorkloadFlag(WorkloadType.TEST_CONNECTION)
supports_multi_team: bool = False
- # The connection-test supervisor uses ``signal.SIGALRM`` (via
``TimeoutPosix``)
- # to bound hook execution. Executors that opt in must run on POSIX systems.
- supports_connection_test: bool = False
sentry_integration: str = ""
+ _legacy_warned: ClassVar[set[str]] = set()
+
is_local: bool = False
is_production: bool = True
@@ -210,6 +246,47 @@ class BaseExecutor(LoggingMixin):
return generator
+ def __init_subclass__(cls, **kwargs: Any) -> None:
+ super().__init_subclass__(**kwargs)
+ cls._legacy_warned = set()
+ enabled: set[WorkloadType] = set()
+ disabled: set[WorkloadType] = set()
+ for flag in ("supports_callbacks", "supports_connection_test"):
+ value = cls.__dict__.get(flag)
+ if not isinstance(value, bool):
+ continue
+ workload_type = vars(BaseExecutor)[flag].workload_type
+ verb, preposition = ("Add", "to") if value else ("Remove", "from")
+ _warn_deprecated_executor_usage(
+ f"{cls.__name__}: setting `{flag} = {value}` as a class
attribute is deprecated. "
+ f"{verb} `WorkloadType.{workload_type.name}` {preposition}
`supported_workload_types` "
+ f"instead.",
+ )
+ (enabled if value else disabled).add(workload_type)
+ delattr(cls, flag)
+ if (enabled or disabled) and "supported_workload_types" not in
cls.__dict__:
+ cls.supported_workload_types = (cls.supported_workload_types |
enabled) - disabled
+ legacy_override_replacements = {
+ "trigger_tasks": "trigger_workloads",
+ "trigger_connection_tests": "trigger_workloads",
+ "order_queued_tasks_by_priority": "_get_workloads_to_schedule",
+ }
+ for legacy_method, replacement in legacy_override_replacements.items():
+ if legacy_method in cls.__dict__:
+ _warn_deprecated_executor_usage(
+ f"{cls.__name__} overrides `{legacy_method}`, which
BaseExecutor no longer "
+ f"calls: the override will not be invoked during
scheduling. Override "
+ f"`{replacement}` instead.",
+ )
+
+ @classmethod
+ def _warn_legacy_property(cls, prop_name: str, message: str) -> None:
+ if prop_name in cls._legacy_warned:
+ return
+ cls._legacy_warned.add(prop_name)
+ warnings.warn(message, RemovedInAirflow4Warning, stacklevel=3)
+ log.warning(message)
+
def __init__(self, parallelism: int = PARALLELISM, team_name: str | None =
None):
stats.initialize(
factory=stats_utils.get_stats_factory(),
@@ -223,9 +300,14 @@ class BaseExecutor(LoggingMixin):
self.parallelism: int = parallelism
self.team_name: str | None = team_name
- self.queued_tasks: dict[TaskInstanceKey, workloads.ExecuteTask] = {}
- self.queued_callbacks: dict[CallbackKey, workloads.ExecuteCallback] =
{}
- self.queued_connection_tests: dict[ConnectionTestKey,
workloads.TestConnection] = {}
+ # TODO(airflow 4.0): flatten to dict[WorkloadKey, ExecutorWorkload]
once the deprecated
+ # queued_tasks / queued_callbacks compat properties are removed.
+ # The defaultdict is load-bearing: queue_workload,
fail_connection_test and the compat
+ # properties rely on auto-vivification of per-type queues, so do not
replace it with a
+ # plain dict.
+ self.executor_queues: defaultdict[WorkloadType, dict[WorkloadKey,
ExecutorWorkload]] = defaultdict(
+ dict
+ )
self.running: set[WorkloadKey] = set()
self.event_buffer: dict[WorkloadKey, EventBufferValueType] = {}
self._task_event_logs: deque[Log] = deque()
@@ -252,6 +334,49 @@ class BaseExecutor(LoggingMixin):
_repr += ")"
return _repr
+ @property
+ def queued_tasks(self) -> dict:
+ """Backward-compat property: delegates to
``executor_queues[WorkloadType.EXECUTE_TASK]``."""
+ self._warn_legacy_property(
+ "queued_tasks",
+ "queued_tasks is deprecated. Use
executor_queues[WorkloadType.EXECUTE_TASK] instead.",
+ )
+ return self.executor_queues[WorkloadType.EXECUTE_TASK]
+
+ @queued_tasks.setter
+ def queued_tasks(self, value: dict) -> None:
+ """Backward-compat setter: writes through to
``executor_queues[WorkloadType.EXECUTE_TASK]``."""
+ self._warn_legacy_property(
+ "queued_tasks",
+ "queued_tasks is deprecated. Use
executor_queues[WorkloadType.EXECUTE_TASK] instead.",
+ )
+ self.executor_queues[WorkloadType.EXECUTE_TASK] = value
+
+ @property
+ def queued_callbacks(self) -> dict:
+ """Backward-compat property: delegates to
``executor_queues[WorkloadType.EXECUTE_CALLBACK]``."""
+ self._warn_legacy_property(
+ "queued_callbacks",
+ "queued_callbacks is deprecated. Use
executor_queues[WorkloadType.EXECUTE_CALLBACK] instead.",
+ )
+ return self.executor_queues[WorkloadType.EXECUTE_CALLBACK]
+
+ @queued_callbacks.setter
+ def queued_callbacks(self, value: dict) -> None:
+ """Backward-compat setter: writes through to
``executor_queues[WorkloadType.EXECUTE_CALLBACK]``."""
+ self._warn_legacy_property(
+ "queued_callbacks",
+ "queued_callbacks is deprecated. Use
executor_queues[WorkloadType.EXECUTE_CALLBACK] instead.",
+ )
+ self.executor_queues[WorkloadType.EXECUTE_CALLBACK] = value
+
+ def _set_workload_type_supported(self, workload_type: WorkloadType,
supported: bool) -> None:
+ # Assign on the instance so the class-level frozenset shared by all
instances is untouched.
+ if supported:
+ self.supported_workload_types = self.supported_workload_types |
{workload_type}
+ else:
+ self.supported_workload_types = self.supported_workload_types -
{workload_type}
+
def start(self): # pragma: no cover
"""Executors may need to get things started."""
@@ -263,55 +388,34 @@ class BaseExecutor(LoggingMixin):
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
- 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 _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.
- Priority Policy: Callbacks are scheduled before tasks (callbacks
complete existing work).
- Callbacks are processed in FIFO order. Tasks are sorted by
priority_weight (higher priority first).
+ Workloads are sorted by ``WORKLOAD_TYPE_PRIORITY`` (priority assigned
by workload type) first,
+ then by ``sort_key`` within the same priority. Lower priority values
are scheduled first;
+ within the same priority, lower ``sort_key`` values come first
(``sort_key=0`` gives FIFO).
:param open_slots: Number of available execution slots
"""
- workloads_to_schedule: list[tuple[WorkloadKey, ExecutorWorkload]] = []
-
- if self.queued_callbacks:
- for key, workload in self.queued_callbacks.items():
- if len(workloads_to_schedule) >= open_slots:
- break
- workloads_to_schedule.append((key, workload))
-
- if open_slots > len(workloads_to_schedule) and self.queued_tasks:
- for task_key, task_workload in
self.order_queued_tasks_by_priority():
- if len(workloads_to_schedule) >= open_slots:
- break
- workloads_to_schedule.append((task_key, task_workload))
-
- return workloads_to_schedule
+ all_workloads: list[tuple[WorkloadKey, ExecutorWorkload]] = [
+ (key, workload) for queue in self.executor_queues.values() for
key, workload in queue.items()
+ ]
+ all_workloads.sort(
+ key=lambda item: (
+ workloads.WORKLOAD_TYPE_PRIORITY.get(item[1].type,
len(workloads.WORKLOAD_TYPE_PRIORITY)),
+ item[1].sort_key,
+ )
+ )
+ return all_workloads[: max(0, open_slots)]
def _process_workloads(self, workload_items: Sequence[ExecutorWorkload])
-> None:
"""
@@ -332,10 +436,11 @@ class BaseExecutor(LoggingMixin):
:param task_instance: TaskInstance
:return: True if the task is known to this executor
"""
+ task_queue = self.executor_queues.get(WorkloadType.EXECUTE_TASK, {})
return (
- task_instance.id in self.queued_tasks
+ task_instance.id in task_queue
or task_instance.id in self.running
- or task_instance.key in self.queued_tasks
+ or task_instance.key in task_queue
or task_instance.key in self.running
)
@@ -351,34 +456,18 @@ class BaseExecutor(LoggingMixin):
open_slots = self.parallelism - len(self.running)
num_running_workloads = len(self.running)
- num_queued_workloads = (
- len(self.queued_tasks) + len(self.queued_callbacks) +
len(self.queued_connection_tests)
- )
+ num_queued_workloads = sum(len(q) for q in
self.executor_queues.values())
self._emit_metrics(open_slots, num_running_workloads,
num_queued_workloads)
- self.trigger_tasks(open_slots)
-
- self.trigger_connection_tests()
+ self.trigger_workloads(open_slots)
# Calling child class sync method
self.log.debug("Calling the %s sync method", self.__class__)
self.sync()
- def trigger_connection_tests(self) -> None:
- """Process queued connection tests, respecting available slot
capacity."""
- if not self.supports_connection_test or not
self.queued_connection_tests:
- return
-
- available = self.slots_available
- if available <= 0:
- return
-
- tests_to_run = list(self.queued_connection_tests.values())[:available]
- self._process_workloads(tests_to_run)
-
def fail_connection_test(self, key: ConnectionTestKey) -> None:
"""Drop a connection-test workload from in-memory queues (called by
the reaper)."""
- self.queued_connection_tests.pop(key, None)
+ self.executor_queues[WorkloadType.TEST_CONNECTION].pop(key, None)
self.running.discard(key)
def _get_metric_name(self, metric_base_name: str) -> str:
@@ -426,30 +515,11 @@ class BaseExecutor(LoggingMixin):
tags=prune_dict({"status": "running", "executor_class_name": name,
"team_name": self.team_name}),
)
- def order_queued_tasks_by_priority(self) -> list[tuple[TaskInstanceKey,
workloads.ExecuteTask]]:
- """
- Orders the queued tasks by priority, highest ``priority_weight`` first.
-
- Consumers take workloads from the front of this list, so the highest
priority
- tasks must come first for them to be scheduled before the rest.
-
- :return: List of workloads from the queued_tasks according to the
priority.
- """
- if not self.queued_tasks:
- return []
-
- # V3 + new executor that supports workloads
- return sorted(
- self.queued_tasks.items(),
- key=lambda x: x[1].ti.priority_weight,
- reverse=True,
- )
-
- def trigger_tasks(self, open_slots: int) -> None:
+ def trigger_workloads(self, open_slots: int) -> None:
"""
- Initiate async execution of queued workloads (tasks and callbacks), up
to the number of available slots.
+ Initiate async execution of queued workloads, up to the number of
available slots.
- Callbacks are prioritized over tasks to complete existing work before
starting new work.
+ Workloads are scheduled according to their ``WORKLOAD_TYPE_PRIORITY``
and ``sort_key``.
:param open_slots: Number of open slots
"""
@@ -475,6 +545,29 @@ class BaseExecutor(LoggingMixin):
if workload_list:
self._process_workloads(workload_list)
+ def trigger_tasks(self, open_slots: int) -> None:
+ """Backward-compat shim: forwards to :meth:`trigger_workloads`."""
+ self._warn_legacy_property(
+ "trigger_tasks",
+ "trigger_tasks is deprecated, use trigger_workloads instead.",
+ )
+ self.trigger_workloads(open_slots)
+
+ def order_queued_tasks_by_priority(self) -> list:
+ """
+ Backward-compat shim: forwards to :meth:`_get_workloads_to_schedule`.
+
+ Note: the return shape differs from the original method. The old
implementation returned
+ *all* queued tasks, tasks-only and untruncated. This shim iterates
every queue in
+ ``executor_queues`` (so callbacks and other workload types are
included) and truncates the
+ result to the number of currently open slots.
+ """
+ self._warn_legacy_property(
+ "order_queued_tasks_by_priority",
+ "order_queued_tasks_by_priority is deprecated, use
_get_workloads_to_schedule instead.",
+ )
+ return self._get_workloads_to_schedule(self.parallelism -
len(self.running))
+
# 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!
@@ -615,37 +708,24 @@ class BaseExecutor(LoggingMixin):
@property
def slots_available(self):
- """Number of new workloads (tasks, callbacks, and connection tests)
this executor instance can accept."""
- return (
- self.parallelism
- - len(self.running)
- - len(self.queued_tasks)
- - len(self.queued_callbacks)
- - len(self.queued_connection_tests)
- )
+ """Number of new workloads this executor instance can accept."""
+ return self.parallelism - self.slots_occupied
@property
def slots_occupied(self):
- """Number of workloads (tasks, callbacks, and connection tests) this
executor instance is currently managing."""
- return (
- len(self.running)
- + len(self.queued_tasks)
- + len(self.queued_callbacks)
- + len(self.queued_connection_tests)
- )
+ """Number of workloads this executor instance is currently managing."""
+ return len(self.running) + sum(len(q) for q in
self.executor_queues.values())
def debug_dump(self):
"""Get called in response to SIGUSR2 by the scheduler."""
- self.log.info(
- "executor.queued_tasks (%d)\n\t%s",
- len(self.queued_tasks),
- "\n\t".join(map(repr, self.queued_tasks.items())),
- )
- self.log.info(
- "executor.queued_callbacks (%d)\n\t%s",
- len(self.queued_callbacks),
- "\n\t".join(map(repr, self.queued_callbacks.items())),
- )
+ for workload_type in WorkloadType:
+ queue = self.executor_queues.get(workload_type, {})
+ self.log.info(
+ "executor.queued[%s] (%d)\n\t%s",
+ workload_type,
+ len(queue),
+ "\n\t".join(map(repr, queue.items())),
+ )
self.log.info("executor.running (%d)\n\t%s", len(self.running),
"\n\t".join(map(repr, self.running)))
self.log.info(
"executor.event_buffer (%d)\n\t%s",
diff --git a/airflow-core/src/airflow/executors/local_executor.py
b/airflow-core/src/airflow/executors/local_executor.py
index 24ab737f5ab..eb43093ffd7 100644
--- a/airflow-core/src/airflow/executors/local_executor.py
+++ b/airflow-core/src/airflow/executors/local_executor.py
@@ -37,6 +37,7 @@ from typing import TYPE_CHECKING
import structlog
from airflow.executors.base_executor import BaseExecutor,
get_execution_api_server_url
+from airflow.executors.workloads import WorkloadType
# add logger to parameter of setproctitle to support logging
if sys.platform == "darwin":
@@ -126,8 +127,11 @@ class LocalExecutor(BaseExecutor):
supports_multi_team: bool = True
serve_logs: bool = True
- supports_callbacks: bool = True
- supports_connection_test: 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).
+ supported_workload_types: frozenset[WorkloadType] = frozenset(
+ {WorkloadType.EXECUTE_TASK, WorkloadType.EXECUTE_CALLBACK,
WorkloadType.TEST_CONNECTION}
+ )
activity_queue: SimpleQueue[ExecutorWorkload | None]
result_queue: SimpleQueue[WorkloadResultType]
@@ -311,13 +315,7 @@ class LocalExecutor(BaseExecutor):
def _process_workloads(self, workload_list):
for workload in workload_list:
self.activity_queue.put(workload)
- # A valid workload will exist in exactly one of these dicts.
- # One pop will succeed, the others will return None gracefully.
- removed = (
- self.queued_tasks.pop(workload.key, None)
- or self.queued_callbacks.pop(workload.key, None)
- or self.queued_connection_tests.pop(workload.key, None)
- )
+ removed = self.executor_queues[workload.type].pop(workload.key,
None)
if not removed:
raise KeyError(f"Workload {workload.key} was not found in any
queue")
with self._unread_messages:
diff --git a/airflow-core/src/airflow/executors/workloads/__init__.py
b/airflow-core/src/airflow/executors/workloads/__init__.py
index 8e4a96ae48d..57f28d2f949 100644
--- a/airflow-core/src/airflow/executors/workloads/__init__.py
+++ b/airflow-core/src/airflow/executors/workloads/__init__.py
@@ -22,7 +22,7 @@ from typing import Annotated
from pydantic import Field
-from airflow.executors.workloads.base import BaseWorkload, BundleInfo
+from airflow.executors.workloads.base import WORKLOAD_TYPE_PRIORITY,
BaseWorkload, BundleInfo, WorkloadType
from airflow.executors.workloads.callback import CallbackFetchMethod,
ExecuteCallback
from airflow.executors.workloads.connection_test import TestConnection
from airflow.executors.workloads.task import ExecuteTask, TaskInstanceDTO
@@ -52,4 +52,6 @@ __all__ = [
"TaskInstance",
"TaskInstanceDTO",
"TestConnection",
+ "WORKLOAD_TYPE_PRIORITY",
+ "WorkloadType",
]
diff --git a/airflow-core/src/airflow/executors/workloads/base.py
b/airflow-core/src/airflow/executors/workloads/base.py
index 43e47065b66..64cd238b83a 100644
--- a/airflow-core/src/airflow/executors/workloads/base.py
+++ b/airflow-core/src/airflow/executors/workloads/base.py
@@ -21,6 +21,7 @@ from __future__ import annotations
import os
from abc import ABC, abstractmethod
from collections.abc import Hashable
+from enum import Enum
from typing import TYPE_CHECKING, Any, ClassVar
from pydantic import BaseModel, ConfigDict, Field
@@ -32,6 +33,34 @@ if TYPE_CHECKING:
from airflow.executors.workloads.types import WorkloadState
+class WorkloadType(str, Enum):
+ """Central registry of executor workload types."""
+
+ EXECUTE_TASK = "ExecuteTask"
+ EXECUTE_CALLBACK = "ExecuteCallback"
+ TEST_CONNECTION = "TestConnection"
+
+
+# Central executor priority registry: tuple is ordered from highest priority
to lowest.
+#
+# Connection tests are short-lived and user-interactive, so they sort ahead of
tasks:
+# otherwise a sustained task backlog would starve them until the reaper times
them out.
+#
+# Adding a ``WorkloadType`` member: the aliases are enforced by
+# ``test_workload_families_track_every_workload_type``; the ``isinstance``
chains in
+# ``BaseExecutor.run_workload``, ``state_class_for_key`` and
+# ``SchedulerJobRunner.process_executor_events`` must be extended by hand.
+_workload_type_priority_order = (
+ WorkloadType.EXECUTE_CALLBACK,
+ WorkloadType.TEST_CONNECTION,
+ WorkloadType.EXECUTE_TASK,
+)
+
+WORKLOAD_TYPE_PRIORITY: dict[WorkloadType, int] = {
+ name: idx for idx, name in enumerate(_workload_type_priority_order)
+}
+
+
class BaseWorkload:
"""
Mixin for ORM models that can be scheduled as workloads.
@@ -96,6 +125,18 @@ class BaseWorkloadSchema(BaseModel):
valid_for=valid_for,
)
+ @property
+ def sort_key(self) -> int:
+ """
+ Return the sort key for ordering workloads within the same priority.
+
+ The default of ``0`` gives FIFO behaviour (Python's stable sort
preserves
+ insertion order among equal keys). Override in subclasses that need
+ priority ordering within their priority group — for example,
``ExecuteTask`` returns
+ ``-self.ti.priority_weight`` so that higher-weight tasks are scheduled
first.
+ """
+ return 0
+
class BaseDagBundleWorkload(BaseWorkloadSchema, ABC):
"""Base class for Workloads that are associated with a DAG bundle."""
diff --git a/airflow-core/src/airflow/executors/workloads/callback.py
b/airflow-core/src/airflow/executors/workloads/callback.py
index 29d6234123a..2f4ef1f09a8 100644
--- a/airflow-core/src/airflow/executors/workloads/callback.py
+++ b/airflow-core/src/airflow/executors/workloads/callback.py
@@ -26,7 +26,7 @@ from uuid import UUID
import structlog
from pydantic import BaseModel, Field, field_validator
-from airflow.executors.workloads.base import BaseDagBundleWorkload, BundleInfo
+from airflow.executors.workloads.base import BaseDagBundleWorkload,
BundleInfo, WorkloadType
from airflow.utils.state import CallbackState
if TYPE_CHECKING:
@@ -75,7 +75,7 @@ class ExecuteCallback(BaseDagBundleWorkload):
callback: CallbackDTO
- type: Literal["ExecuteCallback"] = Field(init=False,
default="ExecuteCallback")
+ type: Literal[WorkloadType.EXECUTE_CALLBACK] = Field(init=False,
default=WorkloadType.EXECUTE_CALLBACK)
token_scope: ClassVar[str] = "callback"
diff --git a/airflow-core/src/airflow/executors/workloads/connection_test.py
b/airflow-core/src/airflow/executors/workloads/connection_test.py
index d3b411af076..5667cfd20a7 100644
--- a/airflow-core/src/airflow/executors/workloads/connection_test.py
+++ b/airflow-core/src/airflow/executors/workloads/connection_test.py
@@ -23,7 +23,7 @@ from typing import TYPE_CHECKING, Literal
from pydantic import Field
-from airflow.executors.workloads.base import BaseWorkloadSchema
+from airflow.executors.workloads.base import BaseWorkloadSchema, WorkloadType
from airflow.models.connection_test import ConnectionTestKey,
ConnectionTestState
if TYPE_CHECKING:
@@ -39,7 +39,7 @@ class TestConnection(BaseWorkloadSchema):
queue: str | None = None
team_name: str | None = None
- type: Literal["TestConnection"] = Field(init=False,
default="TestConnection")
+ type: Literal[WorkloadType.TEST_CONNECTION] = Field(init=False,
default=WorkloadType.TEST_CONNECTION)
@property
def key(self) -> ConnectionTestKey:
diff --git a/airflow-core/src/airflow/executors/workloads/task.py
b/airflow-core/src/airflow/executors/workloads/task.py
index 3099fe1d774..56c69724038 100644
--- a/airflow-core/src/airflow/executors/workloads/task.py
+++ b/airflow-core/src/airflow/executors/workloads/task.py
@@ -24,7 +24,7 @@ from typing import TYPE_CHECKING, Literal
from pydantic import Field
from airflow.api_fastapi.execution_api.datamodels.taskinstance import
TaskInstance
-from airflow.executors.workloads.base import BaseDagBundleWorkload, BundleInfo
+from airflow.executors.workloads.base import BaseDagBundleWorkload,
BundleInfo, WorkloadType
from airflow.utils.state import TaskInstanceState
if TYPE_CHECKING:
@@ -68,13 +68,18 @@ class ExecuteTask(BaseDagBundleWorkload):
ti: TaskInstanceDTO
sentry_integration: str = ""
- type: Literal["ExecuteTask"] = Field(init=False, default="ExecuteTask")
+ type: Literal[WorkloadType.EXECUTE_TASK] = Field(init=False,
default=WorkloadType.EXECUTE_TASK)
@property
def key(self) -> TaskInstanceKey:
"""Return the TaskInstanceKey for this workload."""
return self.ti.key
+ @property
+ def sort_key(self) -> int:
+ """Return the negated task priority weight so the ascending sort
dispatches the highest ``priority_weight`` first."""
+ return -self.ti.priority_weight
+
@property
def display_name(self) -> str:
"""Return the task instance ID as a display name."""
diff --git a/airflow-core/src/airflow/executors/workloads/types.py
b/airflow-core/src/airflow/executors/workloads/types.py
index 3e5d0b06f10..627f28ffb3f 100644
--- a/airflow-core/src/airflow/executors/workloads/types.py
+++ b/airflow-core/src/airflow/executors/workloads/types.py
@@ -26,11 +26,11 @@ from airflow.models.taskinstance import TaskInstance
from airflow.models.taskinstancekey import TaskInstanceKey
from airflow.utils.state import CallbackState, TaskInstanceState
-if TYPE_CHECKING:
- # Type aliases for workload keys and states (used by executor layer)
- WorkloadKey: TypeAlias = TaskInstanceKey | CallbackKey | ConnectionTestKey
- WorkloadState: TypeAlias = TaskInstanceState | CallbackState |
ConnectionTestState
+# Type aliases for workload keys and states (used by executor layer)
+WorkloadKey: TypeAlias = TaskInstanceKey | CallbackKey | ConnectionTestKey
+WorkloadState: TypeAlias = TaskInstanceState | CallbackState |
ConnectionTestState
+if TYPE_CHECKING:
# Type alias for executor workload results (used by executor
implementations)
WorkloadResultType: TypeAlias = tuple[WorkloadKey, WorkloadState,
Exception | None]
diff --git a/airflow-core/src/airflow/jobs/scheduler_job_runner.py
b/airflow-core/src/airflow/jobs/scheduler_job_runner.py
index 089c78f3827..ad78a4fbced 100644
--- a/airflow-core/src/airflow/jobs/scheduler_job_runner.py
+++ b/airflow-core/src/airflow/jobs/scheduler_job_runner.py
@@ -1443,7 +1443,7 @@ class SchedulerJobRunner(BaseJobRunner, LoggingMixin):
if state in (CallbackState.FAILED, CallbackState.SUCCESS):
callback_keys_with_events.append(key)
else:
- cls.logger().error("Unknown workload key type in event buffer:
%r", key)
+ raise TypeError(f"Unknown workload key type in event buffer:
{key!r}")
# Handle callback state events
for callback_id in callback_keys_with_events:
@@ -4156,7 +4156,7 @@ class SchedulerJobRunner(BaseJobRunner, LoggingMixin):
ct.result_message = reason
self.log.warning("Failing connection test %s: %s", ct.id,
reason)
continue
- if not executor.supports_connection_test:
+ if workloads.WorkloadType.TEST_CONNECTION not in
executor.supported_workload_types:
exec_name = executor.name
name = ct.executor or (exec_name and (exec_name.alias or
exec_name.module_path))
reason = f"Executor '{name}' does not support connection
testing"
@@ -4227,7 +4227,7 @@ class SchedulerJobRunner(BaseJobRunner, LoggingMixin):
)
key = ConnectionTestKey(id=str(ct.id))
for executor in self.executors:
- if executor.supports_connection_test:
+ if workloads.WorkloadType.TEST_CONNECTION in
executor.supported_workload_types:
executor.fail_connection_test(key)
session.flush()
diff --git a/airflow-core/tests/unit/executors/test_base_executor.py
b/airflow-core/tests/unit/executors/test_base_executor.py
index 1d8a2dfa329..3ff2dd5a2cd 100644
--- a/airflow-core/tests/unit/executors/test_base_executor.py
+++ b/airflow-core/tests/unit/executors/test_base_executor.py
@@ -18,6 +18,7 @@
from __future__ import annotations
import logging
+import warnings
from datetime import timedelta
from pathlib import Path
from textwrap import dedent
@@ -34,9 +35,11 @@ from airflow._shared.timezones import timezone
from airflow.callbacks.callback_requests import CallbackRequest
from airflow.cli.cli_config import DefaultHelpParser, GroupCommand
from airflow.cli.cli_parser import AirflowHelpFormatter
+from airflow.exceptions import RemovedInAirflow4Warning
from airflow.executors import workloads
from airflow.executors.base_executor import BaseExecutor,
RunningRetryAttemptType
from airflow.executors.local_executor import LocalExecutor
+from airflow.executors.workloads import WorkloadType
from airflow.executors.workloads.base import BundleInfo
from airflow.executors.workloads.callback import CallbackDTO
from airflow.models.callback import CallbackFetchMethod, CallbackKey
@@ -194,10 +197,10 @@ def test_fail_and_success():
],
)
@mock.patch("airflow.executors.base_executor.BaseExecutor.sync")
[email protected]("airflow.executors.base_executor.BaseExecutor.trigger_tasks")
[email protected]("airflow.executors.base_executor.BaseExecutor.trigger_workloads")
@mock.patch("airflow.executors.base_executor.stats.gauge")
def test_gauge_executor_metrics_single_executor(
- mock_stats_gauge, mock_trigger_tasks, mock_sync, team_name, expected_tags
+ mock_stats_gauge, mock_trigger_workloads, mock_sync, team_name,
expected_tags
):
executor = BaseExecutor(team_name=team_name)
executor.heartbeat()
@@ -215,13 +218,13 @@ def test_gauge_executor_metrics_single_executor(
[(LocalExecutor, "LocalExecutor")],
)
@mock.patch("airflow.executors.local_executor.LocalExecutor.sync")
[email protected]("airflow.executors.base_executor.BaseExecutor.trigger_tasks")
[email protected]("airflow.executors.base_executor.BaseExecutor.trigger_workloads")
@mock.patch("airflow.executors.base_executor.stats.gauge")
@mock.patch("airflow.executors.base_executor.ExecutorLoader.get_executor_names")
def test_gauge_executor_metrics_with_multiple_executors(
mock_get_executor_names,
mock_stats_gauge,
- mock_trigger_tasks,
+ mock_trigger_workloads,
mock_local_sync,
executor_class,
executor_name,
@@ -303,7 +306,7 @@ def test_try_adopt_task_instances(dag_maker):
assert BaseExecutor().try_adopt_task_instances(tis) == tis
-def setup_trigger_tasks(dag_maker, parallelism=None):
+def setup_trigger_workloads(dag_maker, parallelism=None):
dagrun = setup_dagrun(dag_maker)
if parallelism:
executor = BaseExecutor(parallelism=parallelism)
@@ -314,21 +317,21 @@ def setup_trigger_tasks(dag_maker, parallelism=None):
for task_instance in dagrun.task_instances:
workload = workloads.ExecuteTask.make(task_instance)
- executor.queued_tasks[task_instance.key] = workload
+ executor.executor_queues[WorkloadType.EXECUTE_TASK][task_instance.key]
= workload
return executor, dagrun
@pytest.mark.db_test
def test_trigger_queued_tasks(dag_maker):
- """Test that trigger_tasks() calls _process_workloads() when there are
queued workloads."""
- executor, dagrun = setup_trigger_tasks(dag_maker)
+ """Test that trigger_workloads() calls _process_workloads() when there are
queued workloads."""
+ executor, dagrun = setup_trigger_workloads(dag_maker)
# Verify tasks are queued
- assert len(executor.queued_tasks) == 3
+ assert len(executor.executor_queues[WorkloadType.EXECUTE_TASK]) == 3
- # Call trigger_tasks with enough slots
- executor.trigger_tasks(open_slots=10)
+ # Call trigger_workloads with enough slots
+ executor.trigger_workloads(open_slots=10)
executor._process_workloads.assert_called_once()
@@ -339,10 +342,10 @@ def test_trigger_queued_tasks(dag_maker):
@pytest.mark.db_test
def test_trigger_running_tasks(dag_maker):
- """Test that trigger_tasks() works when tasks are re-queued."""
- executor, dagrun = setup_trigger_tasks(dag_maker)
+ """Test that trigger_workloads() works when tasks are re-queued."""
+ executor, dagrun = setup_trigger_workloads(dag_maker)
- executor.trigger_tasks(open_slots=10)
+ executor.trigger_workloads(open_slots=10)
executor._process_workloads.assert_called_once()
# Reset mock for second call
@@ -352,20 +355,20 @@ def test_trigger_running_tasks(dag_maker):
ti = dagrun.task_instances[0]
workload = workloads.ExecuteTask.make(ti)
- executor.queued_tasks[ti.key] = workload
+ executor.executor_queues[WorkloadType.EXECUTE_TASK][ti.key] = workload
- executor.trigger_tasks(open_slots=10)
+ executor.trigger_workloads(open_slots=10)
# Verify _process_workloads was called again
executor._process_workloads.assert_called_once()
@pytest.mark.db_test
-def test_trigger_tasks_schedules_highest_priority_first(dag_maker):
+def test_trigger_workloads_schedules_highest_priority_first(dag_maker):
"""When there are fewer open slots than queued tasks, the lowest priority
ones wait."""
date = timezone.utcnow()
- with dag_maker("test_trigger_tasks_priority_order"):
+ with dag_maker("test_trigger_workloads_priority_order"):
BaseOperator(task_id="low", priority_weight=1)
BaseOperator(task_id="medium", priority_weight=5)
BaseOperator(task_id="high", priority_weight=10)
@@ -375,9 +378,10 @@ def
test_trigger_tasks_schedules_highest_priority_first(dag_maker):
executor = BaseExecutor()
executor._process_workloads = mock.Mock(spec=lambda workloads: None)
for task_instance in dagrun.task_instances:
- executor.queued_tasks[task_instance.key] =
workloads.ExecuteTask.make(task_instance)
+ task_queue = executor.executor_queues[WorkloadType.EXECUTE_TASK]
+ task_queue[task_instance.key] =
workloads.ExecuteTask.make(task_instance)
- executor.trigger_tasks(open_slots=2)
+ executor.trigger_workloads(open_slots=2)
scheduled = [workload.ti.task_id for workload in
executor._process_workloads.call_args[0][0]]
assert scheduled == ["high", "medium"]
@@ -387,17 +391,122 @@ def test_debug_dump(caplog):
executor = BaseExecutor()
with caplog.at_level(logging.INFO):
executor.debug_dump()
- assert "executor.queued" in caplog.text
assert "executor.running" in caplog.text
assert "executor.event_buffer" in caplog.text
[email protected]_test
+def test_debug_dump_with_populated_queues(caplog, dag_maker):
+ """Test debug_dump outputs queued workloads when queues are populated."""
+ executor = BaseExecutor()
+ dagrun = setup_dagrun(dag_maker)
+
+ for ti in dagrun.task_instances:
+ workload = workloads.ExecuteTask.make(ti)
+ executor.executor_queues[WorkloadType.EXECUTE_TASK][ti.key] = workload
+
+ with caplog.at_level(logging.INFO):
+ executor.debug_dump()
+
+ queued_msgs = [m for m in caplog.messages if "executor.queued" in m]
+ assert queued_msgs, "Expected at least one 'executor.queued' log message"
+ assert "executor.running" in caplog.text
+ assert "executor.event_buffer" in caplog.text
+
+
+def test_debug_dump_idle_prints_all_queue_counts(caplog):
+ """An idle executor still dumps a (zero) count line for every workload
type."""
+ executor = BaseExecutor()
+ with caplog.at_level(logging.INFO):
+ executor.debug_dump()
+
+ queued_msgs = [m for m in caplog.messages if "executor.queued" in m]
+ assert len(queued_msgs) == len(WorkloadType)
+
+
def test_base_executor_cannot_send_callback():
executor = BaseExecutor()
with pytest.raises(ValueError, match="Callback sink is not ready"):
executor.send_callback(mock.Mock(spec=CallbackRequest))
+def test_queued_tasks_setter_emits_warning_and_writes_through():
+ executor = BaseExecutor()
+ new_queue = {"k": "v"}
+ with pytest.warns(RemovedInAirflow4Warning, match="queued_tasks is
deprecated"):
+ executor.queued_tasks = new_queue # type: ignore[misc]
+ assert executor.executor_queues[WorkloadType.EXECUTE_TASK] is new_queue
+
+
+def test_queued_callbacks_setter_emits_warning_and_writes_through():
+ executor = BaseExecutor()
+ new_queue = {"k": "v"}
+ with pytest.warns(RemovedInAirflow4Warning, match="queued_callbacks is
deprecated"):
+ executor.queued_callbacks = new_queue # type: ignore[misc]
+ assert executor.executor_queues[WorkloadType.EXECUTE_CALLBACK] is new_queue
+
+
[email protected](
+ ("flag", "workload_type"),
+ [
+ ("supports_callbacks", WorkloadType.EXECUTE_CALLBACK),
+ ("supports_connection_test", WorkloadType.TEST_CONNECTION),
+ ],
+)
+def test_supports_flag_setter_emits_warning_and_toggles_workload_type(flag,
workload_type):
+ executor = BaseExecutor()
+ with pytest.warns(RemovedInAirflow4Warning, match=f"{flag} is deprecated"):
+ setattr(executor, flag, True)
+ assert workload_type in executor.supported_workload_types
+ assert workload_type not in BaseExecutor.supported_workload_types
+
+ setattr(executor, flag, False)
+ assert workload_type not in executor.supported_workload_types
+ assert WorkloadType.EXECUTE_TASK in executor.supported_workload_types
+
+
+def test_trigger_tasks_shim_emits_warning_and_forwards():
+ executor = BaseExecutor()
+ with mock.patch.object(executor, "trigger_workloads") as mocked:
+ with pytest.warns(RemovedInAirflow4Warning, match="trigger_tasks is
deprecated"):
+ executor.trigger_tasks(7)
+ mocked.assert_called_once_with(7)
+
+
+def test_order_queued_tasks_by_priority_shim_emits_warning_and_forwards():
+ executor = BaseExecutor()
+ with mock.patch.object(executor, "_get_workloads_to_schedule",
return_value=[]) as mocked:
+ with pytest.warns(RemovedInAirflow4Warning,
match="order_queued_tasks_by_priority is deprecated"):
+ executor.order_queued_tasks_by_priority()
+ mocked.assert_called_once()
+
+
+def test_has_task_does_not_vivify_executor_queue():
+ executor = BaseExecutor()
+ ti = mock.Mock(spec=TaskInstance)
+ ti.id = "id-1"
+ ti.key = TaskInstanceKey("d", "t", "r", 1, -1)
+ assert executor.has_task(ti) is False
+ assert WorkloadType.EXECUTE_TASK not in executor.executor_queues
+
+
+def test_unknown_workload_type_sorts_last_without_crashing():
+ executor = BaseExecutor()
+ known_key = TaskInstanceKey("d", "t", "r", 1, -1)
+ known_workload = mock.Mock()
+ known_workload.type = WorkloadType.EXECUTE_TASK
+ known_workload.sort_key = 0
+ unknown_workload = mock.Mock()
+ unknown_workload.type = "SomeFutureType"
+ unknown_workload.sort_key = 0
+ executor.executor_queues[WorkloadType.EXECUTE_TASK][known_key] =
known_workload
+ executor.executor_queues["SomeFutureType"]["unk"] = unknown_workload #
type: ignore[index]
+
+ scheduled = executor._get_workloads_to_schedule(open_slots=10)
+
+ assert [w for _, w in scheduled] == [known_workload, unknown_workload]
+
+
@skip_if_force_lowest_dependencies_marker
def test_parser_and_formatter_class():
executor = BaseExecutor(42)
@@ -496,12 +605,12 @@ def test_repr():
assert repr(executor) == "BaseExecutor(parallelism=10, team_name='teamA')"
-def test_supports_connection_test_default_value():
- assert not BaseExecutor.supports_connection_test
+def test_test_connection_not_supported_by_default():
+ assert WorkloadType.TEST_CONNECTION not in
BaseExecutor.supported_workload_types
def test_queue_connection_test_workload_rejected_by_default():
- """BaseExecutor (supports_connection_test=False) rejects TestConnection
workloads."""
+ """BaseExecutor (no TEST_CONNECTION in supported_workload_types) rejects
TestConnection workloads."""
executor = BaseExecutor()
wl = workloads.TestConnection.make(
connection_test_id=uuid4(),
@@ -513,9 +622,10 @@ def
test_queue_connection_test_workload_rejected_by_default():
def test_queue_connection_test_workload_accepted_when_supported():
- """An executor with supports_connection_test=True accepts TestConnection
workloads."""
+ """An executor that supports TEST_CONNECTION accepts TestConnection
workloads."""
executor = LocalExecutor()
- executor.queued_connection_tests.clear()
+ queue = executor.executor_queues[WorkloadType.TEST_CONNECTION]
+ queue.clear()
wl = workloads.TestConnection.make(
connection_test_id=uuid4(),
connection_id="test_conn",
@@ -523,11 +633,45 @@ def
test_queue_connection_test_workload_accepted_when_supported():
team_name="team_a",
)
executor.queue_workload(wl, session=mock.MagicMock(spec=Session))
- assert len(executor.queued_connection_tests) == 1
- assert executor.queued_connection_tests[wl.key] is wl
+ assert len(queue) == 1
+ assert queue[wl.key] is wl
assert wl.team_name == "team_a"
+def test_queued_connection_test_dispatched_to_process_workloads():
+ """A queued TestConnection workload reaches _process_workloads via
trigger_workloads."""
+ executor = LocalExecutor()
+ wl = workloads.TestConnection.make(
+ connection_test_id=uuid4(),
+ connection_id="test_conn",
+ timeout=60,
+ )
+ executor.queue_workload(wl, session=mock.MagicMock(spec=Session))
+ with mock.patch.object(executor, "_process_workloads") as mock_process:
+ executor.trigger_workloads(executor.parallelism)
+ mock_process.assert_called_once_with([wl])
+
+
+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):
+ 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
+ )
+ conn_test = mock.Mock()
+ conn_test.type = WorkloadType.TEST_CONNECTION
+ conn_test.sort_key = 0
+
executor.executor_queues[WorkloadType.TEST_CONNECTION][ConnectionTestKey(id="ct")]
= conn_test
+
+ scheduled = [w for _, w in
executor._get_workloads_to_schedule(open_slots=2)]
+
+ assert scheduled[0] is conn_test
+
+
@mock.patch(
"airflow.sdk.execution_time.connection_test_supervisor.supervise_connection_test",
autospec=True,
@@ -556,17 +700,6 @@ def
test_run_workload_passes_team_name_to_connection_test_supervisor(mock_superv
)
-def test_trigger_connection_tests_skipped_when_not_supported():
- """trigger_connection_tests is a no-op when supports_connection_test is
False."""
- executor = BaseExecutor()
- executor.queued_connection_tests[ConnectionTestKey(id="dummy")] =
mock.MagicMock(
- spec=workloads.TestConnection
- )
- with mock.patch.object(executor, "_process_workloads") as mock_process:
- executor.trigger_connection_tests()
- mock_process.assert_not_called()
-
-
@mock.patch.dict("os.environ", {}, clear=True)
class TestExecutorConf:
"""Test ExecutorConf shim class that provides team-specific configuration
access."""
@@ -700,11 +833,11 @@ class TestExecutorConf:
class TestCallbackSupport:
def test_supports_callbacks_flag_default_false(self):
executor = BaseExecutor()
- assert executor.supports_callbacks is False
+ assert WorkloadType.EXECUTE_CALLBACK not in
executor.supported_workload_types
@pytest.mark.db_test
def test_queue_callback_without_support_raises_error(self, dag_maker,
session):
- executor = BaseExecutor() # supports_callbacks = False by default
+ executor = BaseExecutor() # EXECUTE_CALLBACK not in
supported_workload_types by default
callback_data = CallbackDTO(
id="12345678-1234-5678-1234-567812345678",
fetch_method=CallbackFetchMethod.IMPORT_PATH,
@@ -718,13 +851,15 @@ class TestCallbackSupport:
log_path="test.log",
)
- with pytest.raises(NotImplementedError, match="does not support
ExecuteCallback"):
+ with pytest.raises(NotImplementedError, match="does not support
ExecuteCallback workloads"):
executor.queue_workload(callback_workload, session)
@pytest.mark.db_test
def test_queue_workload_with_execute_callback(self, dag_maker, session):
executor = BaseExecutor()
- executor.supports_callbacks = True # Enable for this test
+ executor.supported_workload_types = frozenset(
+ {WorkloadType.EXECUTE_TASK, WorkloadType.EXECUTE_CALLBACK}
+ )
callback_data = CallbackDTO(
id="12345678-1234-5678-1234-567812345678",
fetch_method=CallbackFetchMethod.IMPORT_PATH,
@@ -740,13 +875,15 @@ class TestCallbackSupport:
executor.queue_workload(callback_workload, session)
- assert len(executor.queued_callbacks) == 1
- assert callback_workload.key in executor.queued_callbacks
+ assert len(executor.executor_queues[WorkloadType.EXECUTE_CALLBACK]) ==
1
+ assert callback_workload.key in
executor.executor_queues[WorkloadType.EXECUTE_CALLBACK]
@pytest.mark.db_test
def test_get_workloads_prioritizes_callbacks(self, dag_maker, session):
executor = BaseExecutor()
- executor.supports_callbacks = True # Enable for this test
+ executor.supported_workload_types = frozenset(
+ {WorkloadType.EXECUTE_TASK, WorkloadType.EXECUTE_CALLBACK}
+ )
dagrun = setup_dagrun(dag_maker)
callback_data = CallbackDTO(
id="12345678-1234-5678-1234-567812345678",
@@ -773,6 +910,218 @@ class TestCallbackSupport:
assert isinstance(first_workload, workloads.ExecuteCallback) # Assert
callback comes first
+class TestBackwardCompatProperties:
+ """Tests for the backward-compat properties (queued_tasks,
queued_callbacks, supports_callbacks)."""
+
+ @pytest.fixture(autouse=True)
+ def _reset_legacy_warned(self):
+ BaseExecutor._legacy_warned = set()
+ yield
+ BaseExecutor._legacy_warned = set()
+
+ def test_queued_tasks_delegates_to_executor_queues(self):
+ executor = BaseExecutor()
+ executor.executor_queues[WorkloadType.EXECUTE_TASK]["key1"] =
"workload1"
+
+ with pytest.warns(DeprecationWarning, match="queued_tasks is
deprecated"):
+ result = executor.queued_tasks
+
+ assert result is executor.executor_queues[WorkloadType.EXECUTE_TASK]
+ assert "key1" in result
+
+ def test_queued_callbacks_delegates_to_executor_queues(self):
+ executor = BaseExecutor()
+ executor.executor_queues[WorkloadType.EXECUTE_CALLBACK]["cb1"] =
"callback1"
+
+ with pytest.warns(DeprecationWarning, match="queued_callbacks is
deprecated"):
+ result = executor.queued_callbacks
+
+ assert result is
executor.executor_queues[WorkloadType.EXECUTE_CALLBACK]
+ assert "cb1" in result
+
+ def test_supports_callbacks_delegates_to_supported_workload_types(self):
+ executor = BaseExecutor()
+
+ with pytest.warns(DeprecationWarning, match="supports_callbacks is
deprecated"):
+ assert executor.supports_callbacks is False
+
+ executor.supported_workload_types = frozenset(
+ {WorkloadType.EXECUTE_TASK, WorkloadType.EXECUTE_CALLBACK}
+ )
+ with warnings.catch_warnings():
+ warnings.simplefilter("error", RemovedInAirflow4Warning)
+ assert executor.supports_callbacks is True
+
+ def
test_supports_connection_test_delegates_to_supported_workload_types(self):
+ executor = BaseExecutor()
+
+ with pytest.warns(DeprecationWarning, match="supports_connection_test
is deprecated"):
+ assert executor.supports_connection_test is False
+
+ executor.supported_workload_types = frozenset(
+ {WorkloadType.EXECUTE_TASK, WorkloadType.TEST_CONNECTION}
+ )
+ with warnings.catch_warnings():
+ warnings.simplefilter("error", RemovedInAirflow4Warning)
+ assert executor.supports_connection_test is True
+
+ def test_warning_emitted_once_per_class(self, recwarn):
+ executor = BaseExecutor()
+ for _ in range(5):
+ _ = executor.queued_tasks
+ legacy = [w for w in recwarn.list if "queued_tasks is deprecated" in
str(w.message)]
+ assert len(legacy) == 1
+
+ def test_warning_independent_per_subclass(self, recwarn):
+ class ExecutorA(BaseExecutor):
+ pass
+
+ class ExecutorB(BaseExecutor):
+ pass
+
+ _ = ExecutorA().queued_tasks
+ _ = ExecutorA().queued_tasks
+ _ = ExecutorB().queued_tasks
+ legacy = [w for w in recwarn.list if "queued_tasks is deprecated" in
str(w.message)]
+ assert len(legacy) == 2
+
+ def test_queued_tasks_dict_operations(self):
+ """Verify dict operations through the backward-compat property work
correctly."""
+ executor = BaseExecutor()
+ executor.executor_queues[WorkloadType.EXECUTE_TASK]["k1"] = "w1"
+ executor.executor_queues[WorkloadType.EXECUTE_TASK]["k2"] = "w2"
+
+ with pytest.warns(DeprecationWarning, match="queued_tasks is
deprecated"):
+ qt = executor.queued_tasks
+
+ # All standard dict operations should work on the returned reference
+ assert len(qt) == 2
+ assert "k1" in qt
+ qt.pop("k1")
+ assert len(executor.executor_queues[WorkloadType.EXECUTE_TASK]) == 1
+
+
+class TestLegacySupportsCallbacksShim:
+ """Subclasses declaring legacy ``supports_callbacks = True`` must still
receive callbacks."""
+
+ def test_legacy_flag_synthesises_supported_workload_types(self):
+ with pytest.warns(RemovedInAirflow4Warning, match="supports_callbacks
= True"):
+
+ class LegacyExecutor(BaseExecutor):
+ supports_callbacks = True
+
+ assert LegacyExecutor.supported_workload_types == frozenset(
+ {WorkloadType.EXECUTE_TASK, WorkloadType.EXECUTE_CALLBACK}
+ )
+
+ def test_explicit_supported_workload_types_wins(self):
+ explicit = frozenset({WorkloadType.EXECUTE_TASK})
+ with pytest.warns(RemovedInAirflow4Warning, match="supports_callbacks
= True"):
+
+ class MixedExecutor(BaseExecutor):
+ supports_callbacks = True
+ supported_workload_types = explicit
+
+ assert MixedExecutor.supported_workload_types is explicit
+
+ def test_modern_subclass_emits_no_warning(self, recwarn):
+ class ModernExecutor(BaseExecutor):
+ supported_workload_types = frozenset({WorkloadType.EXECUTE_TASK,
WorkloadType.EXECUTE_CALLBACK})
+
+ legacy_warnings = [w for w in recwarn.list if "supports_callbacks =
True" in str(w.message)]
+ assert legacy_warnings == []
+ assert WorkloadType.EXECUTE_CALLBACK in
ModernExecutor.supported_workload_types
+
+ def test_legacy_false_removes_inherited_workload_type(self):
+ """``= False`` on a parent that declared the type must opt out, not
silently inherit it."""
+ assert WorkloadType.TEST_CONNECTION in
LocalExecutor.supported_workload_types
+ with pytest.warns(RemovedInAirflow4Warning,
match="supports_connection_test = False"):
+
+ class Restricted(LocalExecutor):
+ supports_connection_test = False
+
+ assert WorkloadType.TEST_CONNECTION not in
Restricted.supported_workload_types
+ assert Restricted.supported_workload_types ==
LocalExecutor.supported_workload_types - {
+ WorkloadType.TEST_CONNECTION
+ }
+ assert "supports_connection_test" not in vars(Restricted)
+
+ @pytest.mark.parametrize(
+ ("executor_cls", "expected"),
+ [
+ pytest.param(BaseExecutor, False, id="base-false"),
+ pytest.param(LocalExecutor, True, id="local-true"),
+ ],
+ )
+ def test_class_level_read_returns_bool(self, executor_cls, expected):
+ """``ExecutorCls.supports_callbacks`` must be a bool, not a truthy
property object."""
+ executor_cls._legacy_warned = set()
+ with pytest.warns(RemovedInAirflow4Warning, match="supports_callbacks
is deprecated"):
+ assert executor_cls.supports_callbacks is expected
+
+ def test_instance_assignment_on_legacy_subclass_writes_through(self):
+ """Once the class-body bool is folded in, ``self.flag = X`` must still
reach the descriptor."""
+ with pytest.warns(RemovedInAirflow4Warning, match="supports_callbacks
= False"):
+
+ class LegacyOptOut(LocalExecutor):
+ supports_callbacks = False
+
+ executor = LegacyOptOut()
+ assert WorkloadType.EXECUTE_CALLBACK not in
executor.supported_workload_types
+
+ with pytest.warns(RemovedInAirflow4Warning, match="supports_callbacks
is deprecated"):
+ executor.supports_callbacks = True
+
+ assert executor.supports_callbacks is True
+ assert WorkloadType.EXECUTE_CALLBACK in
executor.supported_workload_types
+ assert "supports_callbacks" not in vars(executor)
+ assert WorkloadType.EXECUTE_CALLBACK not in
LegacyOptOut.supported_workload_types
+
+ def
test_legacy_connection_test_flag_synthesises_supported_workload_types(self):
+ with pytest.warns(RemovedInAirflow4Warning,
match="supports_connection_test = True"):
+
+ class LegacyExecutor(BaseExecutor):
+ supports_connection_test = True
+
+ assert LegacyExecutor.supported_workload_types == frozenset(
+ {WorkloadType.EXECUTE_TASK, WorkloadType.TEST_CONNECTION}
+ )
+
+ def test_both_legacy_flags_synthesise_together(self):
+ with pytest.warns(RemovedInAirflow4Warning):
+
+ class LegacyExecutor(BaseExecutor):
+ supports_callbacks = True
+ supports_connection_test = True
+
+ assert LegacyExecutor.supported_workload_types == frozenset(
+ {WorkloadType.EXECUTE_TASK, WorkloadType.EXECUTE_CALLBACK,
WorkloadType.TEST_CONNECTION}
+ )
+
+ def test_legacy_flag_unions_with_inherited_workload_types(self):
+ """A legacy flag on a subclass must not drop workload types inherited
from the parent."""
+ with pytest.warns(RemovedInAirflow4Warning, match="supports_callbacks
= True"):
+
+ class LegacyLocal(LocalExecutor):
+ supports_callbacks = True
+
+ assert LegacyLocal.supported_workload_types >=
LocalExecutor.supported_workload_types
+ assert WorkloadType.TEST_CONNECTION in
LegacyLocal.supported_workload_types
+
+ @pytest.mark.parametrize(
+ ("legacy_method", "replacement"),
+ [
+ ("trigger_tasks", "trigger_workloads"),
+ ("trigger_connection_tests", "trigger_workloads"),
+ ("order_queued_tasks_by_priority", "_get_workloads_to_schedule"),
+ ],
+ )
+ def test_legacy_method_override_warns_at_class_definition(self,
legacy_method, replacement):
+ """Overrides of methods BaseExecutor no longer calls must warn instead
of breaking silently."""
+ with pytest.warns(RemovedInAirflow4Warning, match=f"overrides
`{legacy_method}`"):
+ type("OverridingExecutor", (BaseExecutor,), {legacy_method: lambda
self, *args: None})
+
+
class TestExecuteCallbackWorkload:
@pytest.mark.parametrize(
("path", "kwargs", "dag_rel_path", "bundle_path", "expect_success",
"error_contains"),
diff --git a/airflow-core/tests/unit/executors/test_local_executor.py
b/airflow-core/tests/unit/executors/test_local_executor.py
index 79987f9ab1c..28adefd22aa 100644
--- a/airflow-core/tests/unit/executors/test_local_executor.py
+++ b/airflow-core/tests/unit/executors/test_local_executor.py
@@ -31,6 +31,7 @@ from airflow._shared.timezones import timezone
from airflow.executors import workloads
from airflow.executors.base_executor import BaseExecutor, ExecutorConf,
get_execution_api_server_url
from airflow.executors.local_executor import LocalExecutor
+from airflow.executors.workloads import WorkloadType
from airflow.executors.workloads.base import BundleInfo
from airflow.executors.workloads.callback import CallbackDTO
from airflow.executors.workloads.task import TaskInstanceDTO
@@ -223,7 +224,7 @@ class TestLocalExecutor:
)
# Process queued workloads to trigger worker spawning
- executor._process_workloads(list(executor.queued_tasks.values()))
+
executor._process_workloads(list(executor.executor_queues[WorkloadType.EXECUTE_TASK].values()))
executor.end()
@@ -240,9 +241,9 @@ class TestLocalExecutor:
assert executor.event_buffer[fail_ti.key][0] == State.FAILED
@mock.patch("airflow.executors.local_executor.LocalExecutor.sync")
- @mock.patch("airflow.executors.base_executor.BaseExecutor.trigger_tasks")
+
@mock.patch("airflow.executors.base_executor.BaseExecutor.trigger_workloads")
@mock.patch("airflow.executors.base_executor.stats.gauge")
- def test_gauge_executor_metrics(self, mock_stats_gauge,
mock_trigger_tasks, mock_sync):
+ def test_gauge_executor_metrics(self, mock_stats_gauge,
mock_trigger_workloads, mock_sync):
executor = LocalExecutor()
executor.heartbeat()
calls = [
@@ -508,9 +509,9 @@ class TestLocalExecutor:
class TestLocalExecutorConnectionTestSupport:
- def test_supports_connection_test_flag_is_true(self):
+ def test_test_connection_is_supported(self):
executor = LocalExecutor()
- assert executor.supports_connection_test is True
+ assert WorkloadType.TEST_CONNECTION in
executor.supported_workload_types
class TestLocalExecutorCallbackSupport:
@@ -520,7 +521,7 @@ class TestLocalExecutorCallbackSupport:
def test_supports_callbacks_flag_is_true(self):
executor = LocalExecutor()
- assert executor.supports_callbacks is True
+ assert WorkloadType.EXECUTE_CALLBACK in
executor.supported_workload_types
@skip_non_fork_mp_start
def test_process_callback_workload_queue_management(self):
@@ -542,9 +543,9 @@ class TestLocalExecutorCallbackSupport:
executor.start()
try:
- executor.queued_callbacks[callback_workload.key] =
callback_workload
+
executor.executor_queues[WorkloadType.EXECUTE_CALLBACK][callback_workload.key]
= callback_workload
executor._process_workloads([callback_workload])
- assert len(executor.queued_callbacks) == 0
+ assert
len(executor.executor_queues[WorkloadType.EXECUTE_CALLBACK]) == 0
# We can't easily verify worker execution without running the
worker,
# but we can verify the helper is called via mock
diff --git a/airflow-core/tests/unit/executors/test_workloads.py
b/airflow-core/tests/unit/executors/test_workloads.py
index 7c639c43594..cd31ea833c5 100644
--- a/airflow-core/tests/unit/executors/test_workloads.py
+++ b/airflow-core/tests/unit/executors/test_workloads.py
@@ -19,6 +19,7 @@ from __future__ import annotations
import dataclasses
from pathlib import PurePosixPath
+from typing import get_args
from uuid import uuid4
import jwt
@@ -26,13 +27,58 @@ import pytest
from airflow.api_fastapi.auth.tokens import JWTGenerator
from airflow.executors import workloads
-from airflow.executors.workloads import TaskInstance, TaskInstanceDTO, base as
workloads_base
-from airflow.executors.workloads.base import BaseWorkloadSchema, BundleInfo
+from airflow.executors.workloads import TaskInstance, TaskInstanceDTO,
WorkloadType, base as workloads_base
+from airflow.executors.workloads.base import WORKLOAD_TYPE_PRIORITY,
BaseWorkloadSchema, BundleInfo
from airflow.executors.workloads.callback import CallbackDTO,
CallbackFetchMethod, ExecuteCallback
from airflow.executors.workloads.task import ExecuteTask
-from airflow.executors.workloads.types import state_class_for_key
-from airflow.models.callback import CallbackKey
+from airflow.executors.workloads.trigger import RunTrigger
+from airflow.executors.workloads.types import (
+ SchedulerWorkload,
+ WorkloadKey,
+ WorkloadState,
+ state_class_for_key,
+)
+from airflow.models.callback import CallbackKey, ExecutorCallback
+from airflow.models.connection_test import ConnectionTestKey,
ConnectionTestRequest, ConnectionTestState
+from airflow.models.taskinstance import TaskInstance as TaskInstanceModel
+from airflow.models.taskinstancekey import TaskInstanceKey
from airflow.sdk.api.datamodels._generated import TaskInstance as
GeneratedTaskInstance
+from airflow.utils.state import CallbackState, TaskInstanceState
+
+# One row per WorkloadType: (schema, key, state enum, ORM model).
+WORKLOAD_FAMILIES: dict[WorkloadType, tuple[type, type, type, type]] = {
+ WorkloadType.EXECUTE_TASK: (ExecuteTask, TaskInstanceKey,
TaskInstanceState, TaskInstanceModel),
+ WorkloadType.EXECUTE_CALLBACK: (ExecuteCallback, CallbackKey,
CallbackState, ExecutorCallback),
+ # Referenced via the package so pytest does not collect ``TestConnection``
as a test class.
+ WorkloadType.TEST_CONNECTION: (
+ workloads.TestConnection,
+ ConnectionTestKey,
+ ConnectionTestState,
+ ConnectionTestRequest,
+ ),
+}
+
+
+def _union_members(alias) -> set[type]:
+ union = get_args(alias)[0] if hasattr(alias, "__metadata__") else alias
+ return set(get_args(union))
+
+
+def test_workload_families_track_every_workload_type():
+ assert set(WORKLOAD_FAMILIES) == set(WorkloadType)
+ assert set(WORKLOAD_TYPE_PRIORITY) == set(WorkloadType)
+
+ schemas = {row[0] for row in WORKLOAD_FAMILIES.values()}
+ keys = {row[1] for row in WORKLOAD_FAMILIES.values()}
+ states = {row[2] for row in WORKLOAD_FAMILIES.values()}
+ models = {row[3] for row in WORKLOAD_FAMILIES.values()}
+
+ 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(WorkloadState) == states
+ assert _union_members(SchedulerWorkload) == models
def test_task_instance_alias_keeps_backwards_compat():
diff --git a/airflow-core/tests/unit/jobs/test_scheduler_job.py
b/airflow-core/tests/unit/jobs/test_scheduler_job.py
index 4d1a540feda..d3ed94e70e8 100644
--- a/airflow-core/tests/unit/jobs/test_scheduler_job.py
+++ b/airflow-core/tests/unit/jobs/test_scheduler_job.py
@@ -63,6 +63,7 @@ from airflow.executors.executor_constants import MOCK_EXECUTOR
from airflow.executors.executor_loader import ExecutorLoader
from airflow.executors.executor_utils import ExecutorName
from airflow.executors.local_executor import LocalExecutor
+from airflow.executors.workloads import WorkloadType
from airflow.jobs.job import Job, run_job
from airflow.jobs.scheduler_job_runner import SCHEDULER_DAG_CACHE_SIZE,
SchedulerJobRunner
from airflow.models.asset import (
@@ -798,6 +799,15 @@ class TestSchedulerJob:
callback_lookups = [c for c in spy_get.call_args_list if c.args
and c.args[0] is Callback]
assert callback_lookups == []
+ def test_process_executor_events_raises_on_unknown_key_type(self, session):
+ """An unrecognised key must fail loudly, matching run_workload and
state_class_for_key."""
+ executor = MockExecutor(do_update=False)
+ self.job_runner = SchedulerJobRunner(Job(), executors=[executor])
+ executor.event_buffer["not-a-workload-key"] =
(TaskInstanceState.SUCCESS, None)
+
+ with pytest.raises(TypeError, match="Unknown workload key type in
event buffer"):
+ self.job_runner._process_executor_events(executor=executor,
session=session)
+
@mock.patch("airflow.jobs.scheduler_job_runner.TaskCallbackRequest")
@mock.patch("airflow._shared.observability.metrics.stats._get_backend")
def test_process_executor_event_missing_dag(
@@ -13508,7 +13518,7 @@ def scheduler_job_runner_for_connection_tests(session):
executor.name = ExecutorName(
module_path="airflow.executors.local_executor.LocalExecutor",
alias="LocalExecutor"
)
- executor.queued_connection_tests.clear()
+ executor.executor_queues[WorkloadType.TEST_CONNECTION].clear()
yield _make_scheduler_runner_for_connection_tests([executor])
session.execute(delete(ConnectionTestRequest))
session.commit()
@@ -13534,7 +13544,14 @@ class TestDispatchConnectionTests:
session.expire_all()
ct = session.get(ConnectionTestRequest, ct.id)
assert ct.state == ConnectionTestState.QUEUED
- assert
len(scheduler_job_runner_for_connection_tests.executor.queued_connection_tests)
== 1
+ assert (
+ len(
+
scheduler_job_runner_for_connection_tests.executor.executor_queues[
+ WorkloadType.TEST_CONNECTION
+ ]
+ )
+ == 1
+ )
@mock.patch.dict(
os.environ,
@@ -13596,7 +13613,7 @@ class TestDispatchConnectionTests:
runner._enqueue_connection_tests(session=session)
- queued = list(runner.executor.queued_connection_tests.values())
+ queued =
list(runner.executor.executor_queues[WorkloadType.TEST_CONNECTION].values())
assert len(queued) == 1
assert queued[0].team_name == expected_workload_team
@@ -13635,7 +13652,7 @@ class TestDispatchConnectionTests:
):
"""Failure message names the executor that was tried, not 'no
executor'."""
unsupporting_executor = BaseExecutor()
- unsupporting_executor.supports_connection_test = False
+ unsupporting_executor.supported_workload_types =
frozenset({WorkloadType.EXECUTE_TASK})
unsupporting_executor.name = ExecutorName(
module_path="airflow.executors.base_executor.BaseExecutor",
alias="celery"
)
@@ -13783,11 +13800,11 @@ class TestDispatchConnectionTests:
executor_a = LocalExecutor()
executor_a.name = ExecutorName(module_path="path.to.ExecutorA",
alias="executor_a")
- executor_a.queued_connection_tests.clear()
+ executor_a.executor_queues[WorkloadType.TEST_CONNECTION].clear()
executor_b = LocalExecutor()
executor_b.name = ExecutorName(module_path="path.to.ExecutorB",
alias="executor_b")
- executor_b.queued_connection_tests.clear()
+ executor_b.executor_queues[WorkloadType.TEST_CONNECTION].clear()
runner = _make_scheduler_runner_for_connection_tests([executor_a,
executor_b])
@@ -13797,8 +13814,8 @@ class TestDispatchConnectionTests:
runner._enqueue_connection_tests(session=session)
- assert len(executor_b.queued_connection_tests) == 1
- assert len(executor_a.queued_connection_tests) == 0
+ assert len(executor_b.executor_queues[WorkloadType.TEST_CONNECTION])
== 1
+ assert len(executor_a.executor_queues[WorkloadType.TEST_CONNECTION])
== 0
@mock.patch.dict(
os.environ,
@@ -13814,11 +13831,11 @@ class TestDispatchConnectionTests:
executor_a = LocalExecutor()
executor_a.name = ExecutorName(module_path="path.to.ExecutorA",
alias="executor_a")
- executor_a.queued_connection_tests.clear()
+ executor_a.executor_queues[WorkloadType.TEST_CONNECTION].clear()
executor_b = LocalExecutor()
executor_b.name = ExecutorName(module_path="path.to.ExecutorB",
alias="executor_b")
- executor_b.queued_connection_tests.clear()
+ executor_b.executor_queues[WorkloadType.TEST_CONNECTION].clear()
runner = _make_scheduler_runner_for_connection_tests([executor_a,
executor_b])
@@ -13830,8 +13847,8 @@ class TestDispatchConnectionTests:
runner._enqueue_connection_tests(session=session)
- assert len(executor_b.queued_connection_tests) == 1
- assert len(executor_a.queued_connection_tests) == 0
+ assert len(executor_b.executor_queues[WorkloadType.TEST_CONNECTION])
== 1
+ assert len(executor_a.executor_queues[WorkloadType.TEST_CONNECTION])
== 0
def test_dispatch_executor_matched_by_class_name(self, session):
"""When executor is specified by class name only, the matching
executor is selected."""
@@ -13840,11 +13857,11 @@ class TestDispatchConnectionTests:
executor_a = LocalExecutor()
executor_a.name = ExecutorName(module_path="path.to.ExecutorA",
alias="executor_a")
- executor_a.queued_connection_tests.clear()
+ executor_a.executor_queues[WorkloadType.TEST_CONNECTION].clear()
executor_b = LocalExecutor()
executor_b.name = ExecutorName(module_path="path.to.ExecutorB",
alias="executor_b")
- executor_b.queued_connection_tests.clear()
+ executor_b.executor_queues[WorkloadType.TEST_CONNECTION].clear()
runner = _make_scheduler_runner_for_connection_tests([executor_a,
executor_b])
@@ -13854,8 +13871,8 @@ class TestDispatchConnectionTests:
runner._enqueue_connection_tests(session=session)
- assert len(executor_b.queued_connection_tests) == 1
- assert len(executor_a.queued_connection_tests) == 0
+ assert len(executor_b.executor_queues[WorkloadType.TEST_CONNECTION])
== 1
+ assert len(executor_a.executor_queues[WorkloadType.TEST_CONNECTION])
== 0
@mock.patch.dict(
os.environ,
@@ -13869,7 +13886,9 @@ class TestDispatchConnectionTests:
):
"""When the resolved executor does not support connection tests, the
test is failed gracefully."""
executor = scheduler_job_runner_for_connection_tests.executor
- executor.supports_connection_test = False
+ executor.supported_workload_types = frozenset(
+ {WorkloadType.EXECUTE_TASK, WorkloadType.EXECUTE_CALLBACK}
+ )
ct = ConnectionTestRequest(conn_type="test_type",
connection_id="test_conn")
session.add(ct)
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 c7a2f263152..3ff7fcf8409 100644
--- a/devel-common/src/tests_common/test_utils/mock_executor.py
+++ b/devel-common/src/tests_common/test_utils/mock_executor.py
@@ -30,6 +30,11 @@ from airflow.models.taskinstancekey import TaskInstanceKey
from airflow.utils.session import create_session
from airflow.utils.state import State
+from tests_common.test_utils.version_compat import AIRFLOW_V_3_4_PLUS
+
+if AIRFLOW_V_3_4_PLUS:
+ from airflow.executors.workloads import WorkloadType
+
if TYPE_CHECKING:
from airflow.executors import workloads
@@ -79,7 +84,11 @@ class MockExecutor(BaseExecutor):
return
with create_session() as session:
- self.history.append(list(self.queued_tasks.values()))
+ if AIRFLOW_V_3_4_PLUS:
+ task_queue = self.executor_queues[WorkloadType.EXECUTE_TASK]
+ else:
+ task_queue = self.queued_tasks
+ self.history.append(list(task_queue.values()))
# Create a stable/predictable sort order for events in self.history
# for tests!
@@ -92,9 +101,9 @@ class MockExecutor(BaseExecutor):
return -prio, date, dag_id, task_id, map_index, try_number
open_slots = self.parallelism - len(self.running)
- sorted_queue = sorted(self.queued_tasks.items(), key=sort_by)
+ sorted_queue = sorted(task_queue.items(), key=sort_by)
for key, workload in sorted_queue[:open_slots]:
- self.queued_tasks.pop(key)
+ task_queue.pop(key)
state = self.mock_task_results[key]
ti = TaskInstance.get_task_instance(
task_id=workload.ti.task_id,
diff --git
a/providers/amazon/src/airflow/providers/amazon/aws/executors/aws_lambda/lambda_executor.py
b/providers/amazon/src/airflow/providers/amazon/aws/executors/aws_lambda/lambda_executor.py
index 5d30d5c84a2..3656fa1338f 100644
---
a/providers/amazon/src/airflow/providers/amazon/aws/executors/aws_lambda/lambda_executor.py
+++
b/providers/amazon/src/airflow/providers/amazon/aws/executors/aws_lambda/lambda_executor.py
@@ -42,10 +42,19 @@ from
airflow.providers.amazon.aws.executors.utils.exponential_backoff_retry impo
)
from airflow.providers.amazon.aws.hooks.lambda_function import LambdaHook
from airflow.providers.amazon.aws.hooks.sqs import SqsHook
-from airflow.providers.amazon.version_compat import AIRFLOW_V_3_0_PLUS,
AIRFLOW_V_3_3_PLUS
+from airflow.providers.amazon.version_compat import (
+ AIRFLOW_V_3_1_PLUS,
+ AIRFLOW_V_3_3_PLUS,
+ AIRFLOW_V_3_4_PLUS,
+)
from airflow.providers.common.compat.sdk import AirflowException, Stats,
timezone
from airflow.utils.helpers import prune_dict
+if AIRFLOW_V_3_4_PLUS:
+ from airflow.executors.workloads.base import WorkloadType
+
+ _SUPPORTED_WORKLOAD_TYPES = frozenset({WorkloadType.EXECUTE_TASK,
WorkloadType.EXECUTE_CALLBACK})
+
if TYPE_CHECKING:
from sqlalchemy.orm import Session
@@ -72,15 +81,11 @@ class AwsLambdaExecutor(BaseExecutor):
"""
supports_multi_team: bool = True
-
- if AIRFLOW_V_3_3_PLUS:
+ if AIRFLOW_V_3_4_PLUS:
+ supported_workload_types: frozenset[WorkloadType] =
_SUPPORTED_WORKLOAD_TYPES
+ elif AIRFLOW_V_3_3_PLUS:
supports_callbacks: bool = True
- if TYPE_CHECKING and AIRFLOW_V_3_0_PLUS:
- # In the v3 path, we store workloads, not commands as strings.
- # TODO: TaskSDK: move this type change into BaseExecutor.
- queued_tasks: dict[WorkloadKey, workloads.All] # type:
ignore[assignment]
-
def __init__(self, *args, **kwargs):
super().__init__(*args, **kwargs)
self.pending_workloads: deque = deque()
@@ -230,17 +235,14 @@ class AwsLambdaExecutor(BaseExecutor):
except Exception:
self.log.exception("An error occurred while syncing workloads.")
- # TODO: Remove this once the minimum supported version is 3.2+, and defer
to BaseExecutor.queue_workload.
- def queue_workload(self, workload: workloads.All, session: Session | None)
-> None:
- from airflow.executors import workloads
+ if not AIRFLOW_V_3_1_PLUS:
- if isinstance(workload, workloads.ExecuteTask):
+ 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)}")
self.queued_tasks[workload.ti.key] = workload
- return
- if AIRFLOW_V_3_3_PLUS and isinstance(workload,
workloads.ExecuteCallback):
- self.queued_callbacks[workload.callback.key] = workload
- return
- raise RuntimeError(f"{type(self)} cannot handle workloads of type
{type(workload)}")
def _process_workloads(self, workload_items: Sequence[workloads.All]) ->
None:
from airflow.executors import workloads
@@ -255,7 +257,10 @@ class AwsLambdaExecutor(BaseExecutor):
queue = workload.ti.queue
executor_config = workload.ti.executor_config or {}
- del self.queued_tasks[key]
+ if AIRFLOW_V_3_4_PLUS:
+ del self.executor_queues[WorkloadType.EXECUTE_TASK][key]
+ else:
+ del self.queued_tasks[key]
self.execute_async(
key=key,
@@ -275,7 +280,10 @@ class AwsLambdaExecutor(BaseExecutor):
if isinstance(workload.callback.data, dict) and "queue" in
workload.callback.data:
queue = workload.callback.data["queue"]
- del self.queued_callbacks[key]
+ if AIRFLOW_V_3_4_PLUS:
+ del
self.executor_queues[WorkloadType.EXECUTE_CALLBACK][key]
+ else:
+ del self.queued_callbacks[key]
self.execute_async(
key=key,
diff --git
a/providers/amazon/src/airflow/providers/amazon/aws/executors/batch/batch_executor.py
b/providers/amazon/src/airflow/providers/amazon/aws/executors/batch/batch_executor.py
index 454ba38c0a4..def7a275353 100644
---
a/providers/amazon/src/airflow/providers/amazon/aws/executors/batch/batch_executor.py
+++
b/providers/amazon/src/airflow/providers/amazon/aws/executors/batch/batch_executor.py
@@ -33,10 +33,20 @@ from
airflow.providers.amazon.aws.executors.utils.exponential_backoff_retry impo
exponential_backoff_retry,
)
from airflow.providers.amazon.aws.hooks.batch_client import BatchClientHook
-from airflow.providers.amazon.version_compat import AIRFLOW_V_3_0_PLUS,
AIRFLOW_V_3_3_PLUS
+from airflow.providers.amazon.version_compat import (
+ AIRFLOW_V_3_0_PLUS,
+ AIRFLOW_V_3_1_PLUS,
+ AIRFLOW_V_3_3_PLUS,
+ AIRFLOW_V_3_4_PLUS,
+)
from airflow.providers.common.compat.sdk import AirflowException, Stats,
timezone
from airflow.utils.helpers import merge_dicts, prune_dict
+if AIRFLOW_V_3_4_PLUS:
+ from airflow.executors.workloads.base import WorkloadType
+
+ _SUPPORTED_WORKLOAD_TYPES = frozenset({WorkloadType.EXECUTE_TASK,
WorkloadType.EXECUTE_CALLBACK})
+
if TYPE_CHECKING:
from sqlalchemy.orm import Session
@@ -91,17 +101,14 @@ class AwsBatchExecutor(BaseExecutor):
"""
supports_multi_team: bool = True
- if AIRFLOW_V_3_3_PLUS:
+ if AIRFLOW_V_3_4_PLUS:
+ supported_workload_types: frozenset[WorkloadType] =
_SUPPORTED_WORKLOAD_TYPES
+ elif AIRFLOW_V_3_3_PLUS:
supports_callbacks: bool = True
# AWS only allows a maximum number of JOBs in the describe_jobs function
DESCRIBE_JOBS_BATCH_SIZE = 99
- if TYPE_CHECKING and AIRFLOW_V_3_0_PLUS:
- # In the v3 path, we store workloads, not commands as strings.
- # TODO: TaskSDK: move this type change into BaseExecutor
- queued_tasks: dict[TaskInstanceKey, workloads.All] # type:
ignore[assignment]
-
def __init__(self, *args, **kwargs):
super().__init__(*args, **kwargs)
self.active_workers = BatchJobCollection()
@@ -132,28 +139,29 @@ class AwsBatchExecutor(BaseExecutor):
fallback=CONFIG_DEFAULTS[AllBatchConfigKeys.MAX_SUBMIT_JOB_ATTEMPTS],
)
- def queue_workload(self, workload: workloads.All, session: Session | None)
-> None:
- from airflow.executors import workloads
+ if not AIRFLOW_V_3_1_PLUS:
- if isinstance(workload, workloads.ExecuteTask):
+ 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)}")
self.queued_tasks[workload.ti.key] = workload
- return
- if AIRFLOW_V_3_3_PLUS and isinstance(workload,
workloads.ExecuteCallback):
- self.queued_callbacks[workload.callback.key] = workload
- return
- raise RuntimeError(f"{type(self)} cannot handle workloads of type
{type(workload)}")
def _process_workloads(self, workload_items: Sequence[workloads.All]) ->
None:
from airflow.executors import workloads
- for w in workload_items:
- if isinstance(w, workloads.ExecuteTask):
- task_command = [w]
- task_key = w.ti.key
- queue = w.ti.queue
- executor_config = w.ti.executor_config or {}
+ for workload in workload_items:
+ if isinstance(workload, workloads.ExecuteTask):
+ task_command = [workload]
+ task_key = workload.ti.key
+ queue = workload.ti.queue
+ executor_config = workload.ti.executor_config or {}
- del self.queued_tasks[task_key]
+ if AIRFLOW_V_3_4_PLUS:
+ del
self.executor_queues[WorkloadType.EXECUTE_TASK][task_key]
+ else:
+ del self.queued_tasks[task_key]
self.execute_async(
key=task_key,
command=task_command, # type: ignore[arg-type]
@@ -161,18 +169,21 @@ class AwsBatchExecutor(BaseExecutor):
executor_config=executor_config,
)
self.running.add(task_key)
- elif AIRFLOW_V_3_3_PLUS and isinstance(w,
workloads.ExecuteCallback):
- callback_command = [w]
- callback_key = w.callback.key
+ elif AIRFLOW_V_3_3_PLUS and isinstance(workload,
workloads.ExecuteCallback):
+ callback_command = [workload]
+ callback_key = workload.callback.key
queue = None
- if isinstance(w.callback.data, dict) and "queue" in
w.callback.data:
- queue = w.callback.data["queue"]
+ if isinstance(workload.callback.data, dict) and "queue" in
workload.callback.data:
+ queue = workload.callback.data["queue"]
- del self.queued_callbacks[callback_key]
+ if AIRFLOW_V_3_4_PLUS:
+ del
self.executor_queues[WorkloadType.EXECUTE_CALLBACK][callback_key]
+ else:
+ del self.queued_callbacks[callback_key]
self.execute_async(key=callback_key, command=callback_command,
queue=queue) # type: ignore[arg-type]
self.running.add(callback_key)
else:
- raise RuntimeError(f"{type(self)} cannot handle workloads of
type {type(w)}")
+ raise RuntimeError(f"{type(self)} cannot handle workloads of
type {type(workload)}")
def check_health(self):
"""Make a test API call to check the health of the Batch Executor."""
diff --git
a/providers/amazon/src/airflow/providers/amazon/aws/executors/ecs/ecs_executor.py
b/providers/amazon/src/airflow/providers/amazon/aws/executors/ecs/ecs_executor.py
index cba439ddc62..fa96034e033 100644
---
a/providers/amazon/src/airflow/providers/amazon/aws/executors/ecs/ecs_executor.py
+++
b/providers/amazon/src/airflow/providers/amazon/aws/executors/ecs/ecs_executor.py
@@ -48,11 +48,19 @@ from
airflow.providers.amazon.aws.executors.utils.exponential_backoff_retry impo
exponential_backoff_retry,
)
from airflow.providers.amazon.aws.hooks.ecs import EcsHook
-from airflow.providers.amazon.version_compat import AIRFLOW_V_3_0_PLUS,
AIRFLOW_V_3_3_PLUS
+from airflow.providers.amazon.version_compat import (
+ AIRFLOW_V_3_0_PLUS,
+ AIRFLOW_V_3_1_PLUS,
+ AIRFLOW_V_3_3_PLUS,
+ AIRFLOW_V_3_4_PLUS,
+)
from airflow.providers.common.compat.sdk import AirflowException, Stats,
timezone
from airflow.utils.helpers import merge_dicts, prune_dict
from airflow.utils.state import State
+if AIRFLOW_V_3_4_PLUS:
+ from airflow.executors.workloads.base import WorkloadType
+
if TYPE_CHECKING:
from sqlalchemy.orm import Session
@@ -101,17 +109,16 @@ class AwsEcsExecutor(BaseExecutor):
supports_multi_team: bool = True
- if AIRFLOW_V_3_3_PLUS:
+ if AIRFLOW_V_3_4_PLUS:
+ supported_workload_types: frozenset[WorkloadType] = frozenset(
+ {WorkloadType.EXECUTE_TASK, WorkloadType.EXECUTE_CALLBACK}
+ )
+ elif AIRFLOW_V_3_3_PLUS:
supports_callbacks: bool = True
# AWS limits the maximum number of ARNs in the describe_tasks function.
DESCRIBE_TASKS_BATCH_SIZE = 99
- if TYPE_CHECKING and AIRFLOW_V_3_0_PLUS:
- # In the v3 path, we store workloads, not commands as strings.
- # TODO: TaskSDK: move this type change into BaseExecutor
- queued_tasks: dict[TaskInstanceKey, workloads.All] # type:
ignore[assignment]
-
def __init__(self, *args, **kwargs):
super().__init__(*args, **kwargs)
self.active_workers: EcsTaskCollection = EcsTaskCollection()
@@ -147,17 +154,14 @@ class AwsEcsExecutor(BaseExecutor):
fallback=CONFIG_DEFAULTS[AllEcsConfigKeys.MAX_RUN_TASK_ATTEMPTS],
)
- # TODO: Remove this once the minimum supported version is 3.3+, and defer
to BaseExecutor.queue_workload.
- def queue_workload(self, workload: workloads.All, session: Session | None)
-> None:
- from airflow.executors import workloads
+ if not AIRFLOW_V_3_1_PLUS:
- if isinstance(workload, workloads.ExecuteTask):
+ 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)}")
self.queued_tasks[workload.ti.key] = workload
- return
- if AIRFLOW_V_3_3_PLUS and isinstance(workload,
workloads.ExecuteCallback):
- self.queued_callbacks[workload.callback.key] = workload
- return
- raise RuntimeError(f"{type(self)} cannot handle workloads of type
{type(workload)}")
def _process_workloads(self, workload_items: Sequence[workloads.All]) ->
None:
""":sphinx-autoapi-skip:."""
@@ -173,7 +177,10 @@ class AwsEcsExecutor(BaseExecutor):
queue = workload.ti.queue
executor_config = workload.ti.executor_config or {}
- del self.queued_tasks[key]
+ if AIRFLOW_V_3_4_PLUS:
+ del self.executor_queues[WorkloadType.EXECUTE_TASK][key]
+ else:
+ del self.queued_tasks[key]
self.execute_async(key=key, command=command, queue=queue,
executor_config=executor_config)
self.running.add(key)
@@ -181,7 +188,10 @@ class AwsEcsExecutor(BaseExecutor):
command = [workload]
key = workload.callback.key
- del self.queued_callbacks[key]
+ if AIRFLOW_V_3_4_PLUS:
+ del
self.executor_queues[WorkloadType.EXECUTE_CALLBACK][key]
+ else:
+ del self.queued_callbacks[key]
self.execute_async(key=key, command=command, queue=None)
self.running.add(key)
diff --git a/providers/amazon/src/airflow/providers/amazon/version_compat.py
b/providers/amazon/src/airflow/providers/amazon/version_compat.py
index 4cb19eb3642..fa6faf2ba41 100644
--- a/providers/amazon/src/airflow/providers/amazon/version_compat.py
+++ b/providers/amazon/src/airflow/providers/amazon/version_compat.py
@@ -41,6 +41,7 @@ AIRFLOW_V_3_1_PLUS: bool = get_base_airflow_version_tuple()
>= (3, 1, 0)
AIRFLOW_V_3_1_1_PLUS: bool = get_base_airflow_version_tuple() >= (3, 1, 1)
AIRFLOW_V_3_1_8_PLUS: bool = get_base_airflow_version_tuple() >= (3, 1, 8)
AIRFLOW_V_3_3_PLUS: bool = get_base_airflow_version_tuple() >= (3, 3, 0)
+AIRFLOW_V_3_4_PLUS: bool = get_base_airflow_version_tuple() >= (3, 4, 0)
try:
from airflow.sdk.definitions._internal.types import NOTSET, ArgNotSet
@@ -60,6 +61,7 @@ __all__ = [
"AIRFLOW_V_3_1_1_PLUS",
"AIRFLOW_V_3_1_8_PLUS",
"AIRFLOW_V_3_3_PLUS",
+ "AIRFLOW_V_3_4_PLUS",
"NOTSET",
"ArgNotSet",
"is_arg_set",
diff --git
a/providers/amazon/tests/unit/amazon/aws/executors/aws_lambda/test_lambda_executor.py
b/providers/amazon/tests/unit/amazon/aws/executors/aws_lambda/test_lambda_executor.py
index 6613b8e39d7..0da44f151cd 100644
---
a/providers/amazon/tests/unit/amazon/aws/executors/aws_lambda/test_lambda_executor.py
+++
b/providers/amazon/tests/unit/amazon/aws/executors/aws_lambda/test_lambda_executor.py
@@ -36,7 +36,12 @@ from airflow.version import version as airflow_version_str
from tests_common.test_utils.compat import timezone
from tests_common.test_utils.config import conf_vars
-from tests_common.test_utils.version_compat import AIRFLOW_V_3_0_PLUS,
AIRFLOW_V_3_1_PLUS, AIRFLOW_V_3_3_PLUS
+from tests_common.test_utils.version_compat import (
+ AIRFLOW_V_3_0_PLUS,
+ AIRFLOW_V_3_1_PLUS,
+ AIRFLOW_V_3_3_PLUS,
+ AIRFLOW_V_3_4_PLUS,
+)
airflow_version = VersionInfo(*map(int, airflow_version_str.split(".")[:3]))
@@ -145,13 +150,22 @@ class TestAwsLambdaExecutor:
ser_workload = json.dumps({"test_key": "test_value"})
workload.model_dump_json.return_value = ser_workload
+ if AIRFLOW_V_3_4_PLUS:
+ from airflow.executors.workloads.base import WorkloadType
+
+ workload.type = WorkloadType.EXECUTE_TASK
+ workload.key = airflow_key
+ task_queue =
mock_executor.executor_queues[WorkloadType.EXECUTE_TASK]
+ else:
+ task_queue = mock_executor.queued_tasks
+
mock_executor.queue_workload(workload, mock.Mock())
- assert mock_executor.queued_tasks[workload.ti.key] == workload
+ assert task_queue[workload.ti.key] == workload
assert len(mock_executor.pending_workloads) == 0
assert len(mock_executor.running) == 0
mock_executor._process_workloads([workload])
- assert len(mock_executor.queued_tasks) == 0
+ assert len(task_queue) == 0
assert len(mock_executor.running) == 1
assert workload.ti.key in mock_executor.running
assert len(mock_executor.pending_workloads) == 1
@@ -189,6 +203,10 @@ class TestAwsLambdaExecutor:
workload.callback = mock.Mock()
workload.callback.key = callback_id
workload.callback.data = {}
+ if AIRFLOW_V_3_4_PLUS:
+ from airflow.executors.workloads.base import WorkloadType
+
+ workload.type = WorkloadType.EXECUTE_CALLBACK
ser_workload = json.dumps({"test_key": "test_value"})
workload.model_dump_json.return_value = ser_workload
@@ -235,9 +253,14 @@ class TestAwsLambdaExecutor:
callback_id = mock_airflow_key()
workload = mock.Mock(spec=ExecuteCallback)
+ workload.key = callback_id
workload.callback = mock.Mock()
workload.callback.key = callback_id
workload.callback.data = {"queue": "fast-queue"}
+ if AIRFLOW_V_3_4_PLUS:
+ from airflow.executors.workloads.base import WorkloadType
+
+ workload.type = WorkloadType.EXECUTE_CALLBACK
ser_workload = json.dumps({"test_key": "test_value"})
workload.model_dump_json.return_value = ser_workload
diff --git
a/providers/amazon/tests/unit/amazon/aws/executors/batch/test_batch_executor.py
b/providers/amazon/tests/unit/amazon/aws/executors/batch/test_batch_executor.py
index 50a90757eae..d3f23f71cd6 100644
---
a/providers/amazon/tests/unit/amazon/aws/executors/batch/test_batch_executor.py
+++
b/providers/amazon/tests/unit/amazon/aws/executors/batch/test_batch_executor.py
@@ -50,7 +50,12 @@ from airflow.version import version as airflow_version_str
from tests_common import RUNNING_TESTS_AGAINST_AIRFLOW_PACKAGES
from tests_common.test_utils.config import conf_vars
-from tests_common.test_utils.version_compat import AIRFLOW_V_3_0_PLUS,
AIRFLOW_V_3_1_PLUS, AIRFLOW_V_3_3_PLUS
+from tests_common.test_utils.version_compat import (
+ AIRFLOW_V_3_0_PLUS,
+ AIRFLOW_V_3_1_PLUS,
+ AIRFLOW_V_3_3_PLUS,
+ AIRFLOW_V_3_4_PLUS,
+)
airflow_version = VersionInfo(*map(int, airflow_version_str.split(".")[:3]))
ARN1 = "arn1"
@@ -232,15 +237,24 @@ class TestAwsBatchExecutor:
ser_workload = json.dumps({"test_key": "test_value"})
workload.model_dump_json.return_value = ser_workload
+ if AIRFLOW_V_3_4_PLUS:
+ from airflow.executors.workloads.base import WorkloadType
+
+ workload.type = WorkloadType.EXECUTE_TASK
+ workload.key = workload.ti.key
+ task_queue =
mock_executor.executor_queues[WorkloadType.EXECUTE_TASK]
+ else:
+ task_queue = mock_executor.queued_tasks
+
mock_executor.queue_workload(workload, mock.Mock())
mock_executor.batch.submit_job.return_value = {"jobId": ARN1,
"jobName": "some-job-name"}
- assert mock_executor.queued_tasks[workload.ti.key] == workload
+ assert task_queue[workload.ti.key] == workload
assert len(mock_executor.pending_jobs) == 0
assert len(mock_executor.running) == 0
mock_executor._process_workloads([workload])
- assert len(mock_executor.queued_tasks) == 0
+ assert len(task_queue) == 0
assert len(mock_executor.running) == 1
assert workload.ti.key in mock_executor.running
assert len(mock_executor.pending_jobs) == 1
@@ -295,6 +309,11 @@ class TestAwsBatchExecutor:
workload = mock.Mock(spec=ExecuteCallback)
workload.callback = mock.Mock()
workload.callback.key = mock_airflow_key()
+ workload.key = workload.callback.key
+ if AIRFLOW_V_3_4_PLUS:
+ from airflow.executors.workloads.base import WorkloadType
+
+ workload.type = WorkloadType.EXECUTE_CALLBACK
ser_workload = json.dumps({"test_key": "test_value"})
workload.model_dump_json.return_value = ser_workload
@@ -355,7 +374,7 @@ class TestAwsBatchExecutor:
@pytest.mark.skipif(not AIRFLOW_V_3_3_PLUS, reason="Test requires Airflow
3.3+")
@mock.patch("airflow.providers.amazon.aws.executors.batch.batch_executor.AwsBatchExecutor.running_state")
- def test_task_sdk_callback_with_queue(self, mock_airflow_key,
mock_executor):
+ def test_task_sdk_callback_with_queue(self, running_state_mock,
mock_airflow_key, mock_executor):
"""Test task sdk execution for callbacks with queue from end-to-end."""
from airflow.executors.workloads import ExecuteCallback
@@ -363,6 +382,11 @@ class TestAwsBatchExecutor:
workload.callback = mock.Mock()
workload.callback.key = mock_airflow_key()
workload.callback.data = {"queue": "fast-queue"}
+ workload.key = workload.callback.key
+ if AIRFLOW_V_3_4_PLUS:
+ from airflow.executors.workloads.base import WorkloadType
+
+ workload.type = WorkloadType.EXECUTE_CALLBACK
mock_executor.queue_workload(workload, mock.Mock())
diff --git
a/providers/amazon/tests/unit/amazon/aws/executors/ecs/test_ecs_executor.py
b/providers/amazon/tests/unit/amazon/aws/executors/ecs/test_ecs_executor.py
index a600257bdc8..0d1f9012fbd 100644
--- a/providers/amazon/tests/unit/amazon/aws/executors/ecs/test_ecs_executor.py
+++ b/providers/amazon/tests/unit/amazon/aws/executors/ecs/test_ecs_executor.py
@@ -59,7 +59,12 @@ from airflow.version import version as airflow_version_str
from tests_common import RUNNING_TESTS_AGAINST_AIRFLOW_PACKAGES
from tests_common.test_utils.config import conf_vars
-from tests_common.test_utils.version_compat import AIRFLOW_V_3_0_PLUS,
AIRFLOW_V_3_1_PLUS, AIRFLOW_V_3_3_PLUS
+from tests_common.test_utils.version_compat import (
+ AIRFLOW_V_3_0_PLUS,
+ AIRFLOW_V_3_1_PLUS,
+ AIRFLOW_V_3_3_PLUS,
+ AIRFLOW_V_3_4_PLUS,
+)
airflow_version = VersionInfo(*map(int, airflow_version_str.split(".")[:3]))
@@ -426,6 +431,15 @@ class TestAwsEcsExecutor:
ser_workload = json.dumps({"test_key": "test_value"})
workload.model_dump_json.return_value = ser_workload
+ if AIRFLOW_V_3_4_PLUS:
+ from airflow.executors.workloads.base import WorkloadType
+
+ workload.type = WorkloadType.EXECUTE_TASK
+ workload.key = workload.ti.key
+ task_queue =
mock_executor.executor_queues[WorkloadType.EXECUTE_TASK]
+ else:
+ task_queue = mock_executor.queued_tasks
+
mock_executor.queue_workload(workload, mock.Mock())
mock_executor.ecs.run_task.return_value = {
@@ -440,11 +454,11 @@ class TestAwsEcsExecutor:
"failures": [],
}
- assert mock_executor.queued_tasks[workload.ti.key] == workload
+ assert task_queue[workload.ti.key] == workload
assert len(mock_executor.pending_workloads) == 0
assert len(mock_executor.running) == 0
mock_executor._process_workloads([workload])
- assert len(mock_executor.queued_tasks) == 0
+ assert len(task_queue) == 0
assert len(mock_executor.running) == 1
assert workload.ti.key in mock_executor.running
assert len(mock_executor.pending_workloads) == 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 ba3af24415a..ae5a824ac7a 100644
--- a/providers/celery/src/airflow/providers/celery/executors/celery_executor.py
+++ b/providers/celery/src/airflow/providers/celery/executors/celery_executor.py
@@ -42,11 +42,16 @@ from airflow.executors.base_executor import BaseExecutor
from airflow.providers.celery.executors import (
celery_executor_utils as _celery_executor_utils, # noqa: F401 # Needed to
register Celery tasks at worker startup, see #63043.
)
-from airflow.providers.celery.version_compat import AIRFLOW_V_3_0_PLUS,
AIRFLOW_V_3_2_PLUS
+from airflow.providers.celery.version_compat import AIRFLOW_V_3_2_PLUS,
AIRFLOW_V_3_4_PLUS
from airflow.providers.common.compat.sdk import AirflowTaskTimeout, Stats
from airflow.utils.helpers import prune_dict
from airflow.utils.state import TaskInstanceState
+if AIRFLOW_V_3_4_PLUS:
+ from airflow.executors.workloads.base import WorkloadType
+
+ _SUPPORTED_WORKLOAD_TYPES = frozenset({WorkloadType.EXECUTE_TASK,
WorkloadType.EXECUTE_CALLBACK})
+
log = logging.getLogger(__name__)
@@ -107,7 +112,10 @@ class CeleryExecutor(BaseExecutor):
"""
supports_ad_hoc_ti_run: bool = True
- supports_callbacks: bool = True
+ if AIRFLOW_V_3_4_PLUS:
+ supported_workload_types: frozenset[WorkloadType] =
_SUPPORTED_WORKLOAD_TYPES
+ else:
+ supports_callbacks: bool = True
sentry_integration: str =
"sentry_sdk.integrations.celery.CeleryIntegration"
pre_assigns_external_executor_id: ClassVar[bool] = True
@@ -115,11 +123,6 @@ class CeleryExecutor(BaseExecutor):
supports_sentry: bool = True
supports_multi_team: bool = True
- if TYPE_CHECKING:
- if AIRFLOW_V_3_0_PLUS:
- # TODO: TaskSDK: move this type change into BaseExecutor
- queued_tasks: dict[WorkloadKey, workloads.All] # type:
ignore[assignment]
-
def __init__(self, *args, **kwargs):
super().__init__(*args, **kwargs)
@@ -220,10 +223,16 @@ class CeleryExecutor(BaseExecutor):
)
self.workload_publish_retries[key] = retries + 1
continue
- if key in self.queued_tasks:
- self.queued_tasks.pop(key)
+ if AIRFLOW_V_3_4_PLUS:
+ if key in self.executor_queues.get(WorkloadType.EXECUTE_TASK,
{}):
+ self.executor_queues[WorkloadType.EXECUTE_TASK].pop(key)
+ else:
+
self.executor_queues[WorkloadType.EXECUTE_CALLBACK].pop(key, None)
else:
- self.queued_callbacks.pop(key, None)
+ if key in self.queued_tasks:
+ self.queued_tasks.pop(key)
+ else:
+ self.queued_callbacks.pop(key, None)
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)
@@ -399,7 +408,10 @@ class CeleryExecutor(BaseExecutor):
except Exception:
self.log.exception("Error revoking task instance %s from
celery", ti.key)
self.running.discard(ti.key)
- self.queued_tasks.pop(ti.key, None)
+ if AIRFLOW_V_3_4_PLUS:
+ self.executor_queues[WorkloadType.EXECUTE_TASK].pop(ti.key, None)
+ else:
+ self.queued_tasks.pop(ti.key, None)
@staticmethod
def get_cli_commands() -> list[GroupCommand]:
diff --git
a/providers/celery/src/airflow/providers/celery/executors/celery_kubernetes_executor.py
b/providers/celery/src/airflow/providers/celery/executors/celery_kubernetes_executor.py
index 9f0d5656890..64cee6f7c4b 100644
---
a/providers/celery/src/airflow/providers/celery/executors/celery_kubernetes_executor.py
+++
b/providers/celery/src/airflow/providers/celery/executors/celery_kubernetes_executor.py
@@ -25,7 +25,8 @@ from deprecated import deprecated
from airflow.exceptions import AirflowProviderDeprecationWarning
from airflow.executors.base_executor import BaseExecutor
-from airflow.providers.celery.executors.celery_executor import
AIRFLOW_V_3_0_PLUS, CeleryExecutor
+from airflow.providers.celery.executors.celery_executor import CeleryExecutor
# noqa: TC001
+from airflow.providers.celery.version_compat import AIRFLOW_V_3_0_PLUS
from airflow.providers.common.compat.sdk import conf
from airflow.utils.providers_configuration_loader import
providers_configuration_loaded
@@ -149,7 +150,7 @@ class CeleryKubernetesExecutor(BaseExecutor):
@property
def slots_occupied(self):
"""Number of tasks this executor instance is currently managing."""
- return len(self.running) + len(self.queued_tasks)
+ return self.celery_executor.slots_occupied +
self.kubernetes_executor.slots_occupied
def queue_command(
self,
diff --git a/providers/celery/src/airflow/providers/celery/version_compat.py
b/providers/celery/src/airflow/providers/celery/version_compat.py
index 8841c057bed..0e871000159 100644
--- a/providers/celery/src/airflow/providers/celery/version_compat.py
+++ b/providers/celery/src/airflow/providers/celery/version_compat.py
@@ -31,6 +31,7 @@ AIRFLOW_V_3_1_PLUS = get_base_airflow_version_tuple() >= (3,
1, 0)
AIRFLOW_V_3_1_9_PLUS = get_base_airflow_version_tuple() >= (3, 1, 9)
AIRFLOW_V_3_2_PLUS = get_base_airflow_version_tuple() >= (3, 2, 0)
AIRFLOW_V_3_3_PLUS = get_base_airflow_version_tuple() >= (3, 3, 0)
+AIRFLOW_V_3_4_PLUS = get_base_airflow_version_tuple() >= (3, 4, 0)
__all__ = [
"AIRFLOW_V_3_0_PLUS",
@@ -38,4 +39,5 @@ __all__ = [
"AIRFLOW_V_3_1_9_PLUS",
"AIRFLOW_V_3_2_PLUS",
"AIRFLOW_V_3_3_PLUS",
+ "AIRFLOW_V_3_4_PLUS",
]
diff --git a/providers/celery/tests/integration/celery/test_celery_executor.py
b/providers/celery/tests/integration/celery/test_celery_executor.py
index f068f93375c..d5c977aa11b 100644
--- a/providers/celery/tests/integration/celery/test_celery_executor.py
+++ b/providers/celery/tests/integration/celery/test_celery_executor.py
@@ -239,7 +239,7 @@ class TestCeleryExecutor:
)
executor.queue_workload(w, session=None)
- executor.trigger_tasks(open_slots=10)
+ executor.trigger_workloads(open_slots=10)
for _ in range(20):
num_tasks = len(executor.workloads.keys())
if num_tasks == 2:
diff --git
a/providers/celery/tests/unit/celery/executors/test_celery_executor.py
b/providers/celery/tests/unit/celery/executors/test_celery_executor.py
index c4641ae2493..9db8475aef2 100644
--- a/providers/celery/tests/unit/celery/executors/test_celery_executor.py
+++ b/providers/celery/tests/unit/celery/executors/test_celery_executor.py
@@ -52,6 +52,7 @@ from tests_common.test_utils.version_compat import (
AIRFLOW_V_3_1_PLUS,
AIRFLOW_V_3_2_PLUS,
AIRFLOW_V_3_3_PLUS,
+ AIRFLOW_V_3_4_PLUS,
)
try:
@@ -193,9 +194,12 @@ class TestCeleryExecutor:
assert FAKE_EXCEPTION_MSG in caplog.text, caplog.record_tuples
@mock.patch("airflow.providers.celery.executors.celery_executor.CeleryExecutor.sync")
-
@mock.patch("airflow.providers.celery.executors.celery_executor.CeleryExecutor.trigger_tasks")
+ @mock.patch(
+ "airflow.providers.celery.executors.celery_executor.CeleryExecutor."
+ + ("trigger_workloads" if AIRFLOW_V_3_4_PLUS else "trigger_tasks")
+ )
@mock.patch(f"{stats_reference}.gauge")
- def test_gauge_executor_metrics(self, mock_stats_gauge,
mock_trigger_tasks, mock_sync):
+ def test_gauge_executor_metrics(self, mock_stats_gauge, mock_trigger,
mock_sync):
executor = celery_executor.CeleryExecutor()
executor.heartbeat()
calls = [
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 c9f597e13ad..3216ed232b8 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
@@ -59,13 +59,16 @@ from
airflow.providers.cncf.kubernetes.kubernetes_helper_functions import (
annotations_to_key,
)
from airflow.providers.cncf.kubernetes.pod_generator import PodGenerator
-from airflow.providers.cncf.kubernetes.version_compat import AIRFLOW_V_3_0_PLUS
+from airflow.providers.cncf.kubernetes.version_compat import AIRFLOW_V_3_4_PLUS
from airflow.providers.common.compat.sdk import Stats, conf
from airflow.utils.helpers import prune_dict
from airflow.utils.log.logging_mixin import remove_escape_codes
from airflow.utils.session import NEW_SESSION, provide_session
from airflow.utils.state import TaskInstanceState
+if AIRFLOW_V_3_4_PLUS:
+ from airflow.executors.workloads.base import WorkloadType
+
if TYPE_CHECKING:
from collections.abc import Sequence
from multiprocessing.managers import SyncManager
@@ -105,11 +108,6 @@ class KubernetesExecutor(BaseExecutor):
supports_ad_hoc_ti_run: bool = True
supports_multi_team: bool = True
- if TYPE_CHECKING and AIRFLOW_V_3_0_PLUS:
- # In the v3 path, we store workloads, not commands as strings.
- # TODO: TaskSDK: move this type change into BaseExecutor
- queued_tasks: dict[TaskInstanceKey, workloads.All] # type:
ignore[assignment]
-
def __init__(self, *args, **kwargs):
super().__init__(*args, **kwargs)
@@ -372,14 +370,6 @@ class KubernetesExecutor(BaseExecutor):
self.pod_launch_attempts[key] = _PodLaunchAttempt(job=job)
self.task_queue.put(job)
- 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
-
def _process_workloads(self, workloads: Sequence[workloads.All]) -> None:
from airflow.executors.workloads import ExecuteTask
@@ -394,7 +384,10 @@ class KubernetesExecutor(BaseExecutor):
queue = w.ti.queue
executor_config = w.ti.executor_config or {}
- del self.queued_tasks[key]
+ if AIRFLOW_V_3_4_PLUS:
+ del self.executor_queues[WorkloadType.EXECUTE_TASK][key]
+ else:
+ del self.queued_tasks[key]
self.execute_async(key=key, command=command, queue=queue,
executor_config=executor_config)
self.running.add(key)
@@ -416,7 +409,10 @@ class KubernetesExecutor(BaseExecutor):
if self.running:
self.log.debug("self.running: %s", self.running)
- if self.queued_tasks:
+ if AIRFLOW_V_3_4_PLUS:
+ if any(self.executor_queues.values()):
+ self.log.debug("self.queued: %s", self.executor_queues)
+ elif self.queued_tasks:
self.log.debug("self.queued: %s", self.queued_tasks)
self.kube_scheduler.sync()
@@ -1000,8 +996,12 @@ class KubernetesExecutor(BaseExecutor):
if TYPE_CHECKING:
assert self.kube_client
assert self.kube_scheduler
+
self.running.discard(ti.key)
- self.queued_tasks.pop(ti.key, None)
+ if AIRFLOW_V_3_4_PLUS:
+ self.executor_queues[WorkloadType.EXECUTE_TASK].pop(ti.key, None)
+ else:
+ self.queued_tasks.pop(ti.key, None)
pod_combined_search_str_to_pod_map =
self.get_pod_combined_search_str_to_pod_map()
# Build the pod selector
base_label_selector = f"dag_id={ti.dag_id},task_id={ti.task_id}"
diff --git
a/providers/cncf/kubernetes/src/airflow/providers/cncf/kubernetes/executors/local_kubernetes_executor.py
b/providers/cncf/kubernetes/src/airflow/providers/cncf/kubernetes/executors/local_kubernetes_executor.py
index a13423ccbf6..97505c441ec 100644
---
a/providers/cncf/kubernetes/src/airflow/providers/cncf/kubernetes/executors/local_kubernetes_executor.py
+++
b/providers/cncf/kubernetes/src/airflow/providers/cncf/kubernetes/executors/local_kubernetes_executor.py
@@ -18,7 +18,7 @@
from __future__ import annotations
from collections.abc import Sequence
-from typing import TYPE_CHECKING, Any
+from typing import TYPE_CHECKING
from deprecated import deprecated
@@ -96,15 +96,6 @@ class LocalKubernetesExecutor(BaseExecutor):
def _task_event_logs(self, value):
"""Not implemented for hybrid executors."""
- @property
- def queued_tasks(self) -> dict[TaskInstanceKey, Any]:
- """Return queued tasks from local and kubernetes executor."""
- return self.local_executor.queued_tasks |
self.kubernetes_executor.queued_tasks
-
- @queued_tasks.setter
- def queued_tasks(self, value) -> None:
- """Not implemented for hybrid executors."""
-
@property # type: ignore[override]
def running(self) -> set[TaskInstanceKey]:
"""Return running tasks from local and kubernetes executor."""
@@ -145,7 +136,7 @@ class LocalKubernetesExecutor(BaseExecutor):
@property
def slots_occupied(self):
"""Number of tasks this executor instance is currently managing."""
- return len(self.running) + len(self.queued_tasks)
+ return self.local_executor.slots_occupied +
self.kubernetes_executor.slots_occupied
def queue_command(
self,
diff --git
a/providers/cncf/kubernetes/src/airflow/providers/cncf/kubernetes/version_compat.py
b/providers/cncf/kubernetes/src/airflow/providers/cncf/kubernetes/version_compat.py
index dab50860d19..6a2f0ddebf2 100644
---
a/providers/cncf/kubernetes/src/airflow/providers/cncf/kubernetes/version_compat.py
+++
b/providers/cncf/kubernetes/src/airflow/providers/cncf/kubernetes/version_compat.py
@@ -36,6 +36,7 @@ AIRFLOW_V_3_0_PLUS = get_base_airflow_version_tuple() >= (3,
0, 0)
AIRFLOW_V_3_1_PLUS = get_base_airflow_version_tuple() >= (3, 1, 0)
AIRFLOW_V_3_2_PLUS = get_base_airflow_version_tuple() >= (3, 2, 0)
AIRFLOW_V_3_3_PLUS = get_base_airflow_version_tuple() >= (3, 3, 0)
+AIRFLOW_V_3_4_PLUS = get_base_airflow_version_tuple() >= (3, 4, 0)
__all__ = [
@@ -43,4 +44,5 @@ __all__ = [
"AIRFLOW_V_3_1_PLUS",
"AIRFLOW_V_3_2_PLUS",
"AIRFLOW_V_3_3_PLUS",
+ "AIRFLOW_V_3_4_PLUS",
]
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 0c1e03cdeea..15b95ab4d1c 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
@@ -70,6 +70,7 @@ from tests_common.test_utils.version_compat import (
AIRFLOW_V_3_1_PLUS,
AIRFLOW_V_3_2_PLUS,
AIRFLOW_V_3_3_PLUS,
+ AIRFLOW_V_3_4_PLUS,
)
try:
@@ -1633,9 +1634,12 @@ class TestKubernetesExecutor:
@mock.patch("airflow.providers.cncf.kubernetes.executors.kubernetes_executor.KubeConfig")
@mock.patch("airflow.providers.cncf.kubernetes.executors.kubernetes_executor.KubernetesExecutor.sync")
- @mock.patch("airflow.executors.base_executor.BaseExecutor.trigger_tasks")
+ @mock.patch(
+ "airflow.executors.base_executor.BaseExecutor."
+ + ("trigger_workloads" if AIRFLOW_V_3_4_PLUS else "trigger_tasks")
+ )
@mock.patch(f"{stats_reference}.gauge")
- def test_gauge_executor_metrics(self, mock_stats_gauge,
mock_trigger_tasks, mock_sync, mock_kube_config):
+ def test_gauge_executor_metrics(self, mock_stats_gauge, mock_trigger,
mock_sync, mock_kube_config):
executor = self.kubernetes_executor
executor.heartbeat()
calls = [
@@ -1657,6 +1661,40 @@ class TestKubernetesExecutor:
]
mock_stats_gauge.assert_has_calls(calls)
+ @pytest.mark.skipif(not AIRFLOW_V_3_0_PLUS, reason="Test requires Airflow
3+")
+ @mock.patch(
+
"airflow.providers.cncf.kubernetes.executors.kubernetes_executor.KubernetesExecutor.execute_async"
+ )
+ def test_process_workloads(self, mock_execute_async):
+ """Test that _process_workloads dequeues an ExecuteTask and hands it
to execute_async."""
+ from airflow.executors.workloads import ExecuteTask
+
+ executor = self.kubernetes_executor
+ key = TaskInstanceKey("dag", "task", "run_id", 1, -1)
+ workload = mock.Mock(spec=ExecuteTask)
+ workload.ti = mock.Mock()
+ workload.ti.key = key
+ workload.ti.queue = "default"
+ workload.ti.executor_config = None
+
+ if AIRFLOW_V_3_4_PLUS:
+ from airflow.executors.workloads.base import WorkloadType
+
+ workload.type = WorkloadType.EXECUTE_TASK
+ workload.key = key
+ task_queue = executor.executor_queues[WorkloadType.EXECUTE_TASK]
+ else:
+ task_queue = executor.queued_tasks
+ task_queue[key] = workload
+
+ executor._process_workloads([workload])
+
+ assert len(task_queue) == 0
+ assert key in executor.running
+ mock_execute_async.assert_called_once_with(
+ key=key, command=[workload], queue="default", executor_config={}
+ )
+
@mock.patch("airflow.providers.cncf.kubernetes.executors.kubernetes_executor_utils.KubernetesJobWatcher")
@mock.patch("airflow.providers.cncf.kubernetes.kube_client.get_kube_client")
def test_invalid_executor_config(self, mock_get_kube_client,
mock_kubernetes_job_watcher):
diff --git
a/providers/cncf/kubernetes/tests/unit/cncf/kubernetes/executors/test_local_kubernetes_executor.py
b/providers/cncf/kubernetes/tests/unit/cncf/kubernetes/executors/test_local_kubernetes_executor.py
index b8ec857d6be..e35596822ac 100644
---
a/providers/cncf/kubernetes/tests/unit/cncf/kubernetes/executors/test_local_kubernetes_executor.py
+++
b/providers/cncf/kubernetes/tests/unit/cncf/kubernetes/executors/test_local_kubernetes_executor.py
@@ -51,21 +51,18 @@ class TestLocalKubernetesExecutor:
def test_cli_commands_vended(self):
assert LocalKubernetesExecutor.get_cli_commands()
- def test_queued_tasks(self):
+ def test_slots_occupied_sums_children_without_deprecation(self):
local_executor_mock = mock.MagicMock()
k8s_executor_mock = mock.MagicMock()
- local_kubernetes_executor =
LocalKubernetesExecutor(local_executor_mock, k8s_executor_mock)
-
- local_queued_tasks = {("dag_id", "task_id", "2020-08-30", 1):
"queued_command"}
- k8s_queued_tasks = {("dag_id_2", "task_id_2", "2020-08-30", 2):
"queued_command"}
+ local_executor_mock.slots_occupied = 3
+ k8s_executor_mock.slots_occupied = 2
+ local_executor_mock.running = {("dag_id", "task_id", "2020-08-30", 1)}
+ k8s_executor_mock.running = set()
- local_executor_mock.queued_tasks = local_queued_tasks
- k8s_executor_mock.queued_tasks = k8s_queued_tasks
-
- expected_queued_tasks = {**local_queued_tasks, **k8s_queued_tasks}
+ local_kubernetes_executor =
LocalKubernetesExecutor(local_executor_mock, k8s_executor_mock)
- assert local_kubernetes_executor.queued_tasks == expected_queued_tasks
- assert len(local_kubernetes_executor.queued_tasks) == 2
+ assert local_kubernetes_executor.slots_occupied == 5
+ assert "queued_tasks" not in {c[0] for c in
local_executor_mock.method_calls}
def test_running(self):
local_executor_mock = mock.MagicMock()
diff --git
a/providers/edge3/src/airflow/providers/edge3/executors/edge_executor.py
b/providers/edge3/src/airflow/providers/edge3/executors/edge_executor.py
index 0b5b9c9e4dd..6990b0097b7 100644
--- a/providers/edge3/src/airflow/providers/edge3/executors/edge_executor.py
+++ b/providers/edge3/src/airflow/providers/edge3/executors/edge_executor.py
@@ -40,11 +40,15 @@ from airflow.providers.edge3.models.types import (
build_callback_run_id,
is_callback_execute,
)
+from airflow.providers.edge3.version_compat import AIRFLOW_V_3_4_PLUS
from airflow.utils.db import DBLocks, create_global_lock
from airflow.utils.helpers import prune_dict
from airflow.utils.session import NEW_SESSION, provide_session
from airflow.utils.state import TaskInstanceState
+if AIRFLOW_V_3_4_PLUS:
+ from airflow.executors.workloads.base import WorkloadType
+
if TYPE_CHECKING:
from sqlalchemy.orm import Session
@@ -399,7 +403,10 @@ class EdgeExecutor(BaseExecutor):
"""
# Remove from executor's internal state
self.running.discard(ti.key)
- self.queued_tasks.pop(ti.key, None)
+ if AIRFLOW_V_3_4_PLUS:
+ self.executor_queues[WorkloadType.EXECUTE_TASK].pop(ti.key, None)
+ else:
+ self.queued_tasks.pop(ti.key, None)
if ti.key in self.last_reported_state:
del self.last_reported_state[ti.key]
diff --git a/providers/edge3/src/airflow/providers/edge3/version_compat.py
b/providers/edge3/src/airflow/providers/edge3/version_compat.py
index 0f3b2b445c1..41f58928638 100644
--- a/providers/edge3/src/airflow/providers/edge3/version_compat.py
+++ b/providers/edge3/src/airflow/providers/edge3/version_compat.py
@@ -35,9 +35,11 @@ def get_base_airflow_version_tuple() -> tuple[int, int, int]:
AIRFLOW_V_3_1_PLUS = get_base_airflow_version_tuple() >= (3, 1, 0)
AIRFLOW_V_3_2_PLUS = get_base_airflow_version_tuple() >= (3, 2, 0)
AIRFLOW_V_3_3_PLUS = get_base_airflow_version_tuple() >= (3, 3, 0)
+AIRFLOW_V_3_4_PLUS = get_base_airflow_version_tuple() >= (3, 4, 0)
__all__ = [
"AIRFLOW_V_3_1_PLUS",
"AIRFLOW_V_3_2_PLUS",
"AIRFLOW_V_3_3_PLUS",
+ "AIRFLOW_V_3_4_PLUS",
]
diff --git a/providers/edge3/tests/unit/edge3/executors/test_edge_executor.py
b/providers/edge3/tests/unit/edge3/executors/test_edge_executor.py
index fd08a503e55..463c515cc3f 100644
--- a/providers/edge3/tests/unit/edge3/executors/test_edge_executor.py
+++ b/providers/edge3/tests/unit/edge3/executors/test_edge_executor.py
@@ -65,7 +65,7 @@ class TestEdgeExecutor:
ti.dag_run.run_id = key.run_id
ti.dag_run.start_date = datetime(2021, 1, 1)
executor = EdgeExecutor()
- executor.queued_tasks = {key: [None, None, None, ti]}
+ executor.queued_tasks[key] = [None, None, None, ti]
return (executor, key)
