kaxil commented on code in PR #63491:
URL: https://github.com/apache/airflow/pull/63491#discussion_r4004970210
##########
airflow-core/src/airflow/executors/base_executor.py:
##########
@@ -263,65 +398,44 @@ def log_task_event(self, *, event: str, extra: str,
ti_key: WorkloadKey):
self._task_event_logs.append(Log(event=event, task_instance=ti_key,
extra=extra))
def queue_workload(self, workload: ExecutorWorkload, session: Session) ->
None:
- if isinstance(workload, workloads.ExecuteTask):
- ti = workload.ti
- self.queued_tasks[ti.key] = workload
- 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:
+ def _process_workloads(self, workloads: Sequence[ExecutorWorkload]) ->
None:
Review Comment:
Renaming this parameter off `workload_items` breaks keyword calls into a
documented extension point (`core-concepts/executor/index.rst:313` lists
`_process_workloads` as one), since 3.3.1 and the merge base both named it
`workload_items` and an out-of-tree
`super()._process_workloads(workload_items=...)` now raises `TypeError`. It
also shadows the module-level `workloads` import that
`_get_workloads_to_schedule` uses just above, and splits the name three ways
in-tree while celery, ecs, batch and lambda all stay on `workload_items`.
##########
airflow-core/tests/unit/executors/test_base_executor.py:
##########
@@ -773,6 +910,181 @@ def test_get_workloads_prioritizes_callbacks(self,
dag_maker, session):
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_does_not_synthesise(self, recwarn):
+ class OptedOutExecutor(BaseExecutor):
+ supports_callbacks = False
+
+ legacy_warnings = [w for w in recwarn.list if "supports_callbacks =
True" in str(w.message)]
+ assert legacy_warnings == []
+ assert WorkloadType.EXECUTE_CALLBACK not in
OptedOutExecutor.supported_workload_types
Review Comment:
This passes for the wrong reason. `OptedOutExecutor` derives from bare
`BaseExecutor`, whose default `supported_workload_types` never held
`EXECUTE_CALLBACK`, so the assertion holds even with `__init_subclass__`
deleted outright. `test_legacy_flag_unions_with_inherited_workload_types` below
picks `LocalExecutor` as the parent for exactly this reason, and the `= False`
case needs that same parent before it binds to anything.
##########
airflow-core/src/airflow/executors/base_executor.py:
##########
@@ -210,6 +216,47 @@ def jwt_generator(self) -> JWTGenerator:
return generator
+ def __init_subclass__(cls, **kwargs: Any) -> None:
+ super().__init_subclass__(**kwargs)
+ cls._legacy_warned = set()
+ legacy_workload_types: set[WorkloadType] = set()
+ if cls.__dict__.get("supports_callbacks") is True:
+ _warn_deprecated_executor_usage(
+ f"{cls.__name__}: setting `supports_callbacks = True` as a
class attribute is "
+ f"deprecated. Add `WorkloadType.EXECUTE_CALLBACK` to
`supported_workload_types` "
+ f"instead.",
+ )
+ legacy_workload_types.add(WorkloadType.EXECUTE_CALLBACK)
+ if cls.__dict__.get("supports_connection_test") is True:
Review Comment:
Three unwarned breaks fall out of turning these two released attributes into
properties. Class-level reads flip: `BaseExecutor.supports_callbacks` is now a
truthy `property` object where 3.3.1 gave a plain `False`
(`base_executor.py:172` at that tag), and `property.__get__(None, owner)` never
runs the getter, so `if executor_cls.supports_callbacks:` inverts with no
warning at all. The `= False` opt-out is dropped, since this check only matches
`is True` while 3.3.1's `LocalExecutor` declared both flags `True`, so `class
Restricted(LocalExecutor): supports_connection_test = False` keeps
TEST_CONNECTION in the inherited frozenset and `scheduler_job_runner.py:4098`
now dispatches the connection tests 3.3.x refused with "does not support
connection testing". And the setters added last round never fire for exactly
those subclasses: a retained legacy bool sits earlier in the MRO, so
`self.supports_callbacks = X` writes a plain instance attribute, and a
legacy-`False` subclass enabling a
t runtime reads back `True` while `supported_workload_types` stays
`{EXECUTE_TASK}` and `queue_workload` then raises. A descriptor whose `__get__`
resolves from `owner.supported_workload_types` when `obj is None`, plus
handling `is False` here and dropping the legacy attribute once it is folded
in, closes all three.
##########
airflow-core/newsfragments/63491.significant.rst:
##########
@@ -0,0 +1,52 @@
+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
read-only
+``supports_connection_test`` compat property remains;
``queued_connection_tests`` and
Review Comment:
Two bits of this drifted from the latest commit. This line still calls
`supports_connection_test` read-only, but `@supports_callbacks.setter` and
`@supports_connection_test.setter` landed in the same push, so instance
assignment works again. Line 47 also still names only `trigger_tasks` and
`order_queued_tasks_by_priority` as the overrides that warn at class-definition
time, while `trigger_connection_tests` joined `legacy_override_replacements` in
that same commit.
##########
airflow-core/src/airflow/executors/workloads/base.py:
##########
@@ -32,6 +33,42 @@
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 new workload type is a six-place change that must stay in sync:
Review Comment:
The checklist has come up short twice now, so rather than a seventh and
eighth bullet, here is the closed set, derived from the tree rather than by
reading it.
`WorkloadType` has three members, and four parallel three-member families
track it: the schemas (`ExecuteTask`, `ExecuteCallback`, `TestConnection`), the
keys (`TaskInstanceKey`, `CallbackKey`, `ConnectionTestKey`), the states
(`TaskInstanceState`, `CallbackState`, `ConnectionTestState`), and the ORM
models (`TaskInstance`, `ExecutorCallback`, `ConnectionTestRequest`). Adding a
fourth workload type touches ten places: the enum,
`_workload_type_priority_order`, the five aliases (`All` and `ExecutorWorkload`
in `workloads/__init__.py`; `WorkloadKey`, `WorkloadState` and
`SchedulerWorkload` in `workloads/types.py`), and three dispatch sites.
The dispatch sites are `BaseExecutor.run_workload`, ending in `raise
ValueError` at `base_executor.py:845`; `state_class_for_key`, ending in `raise
TypeError`; and `SchedulerJobRunner.process_executor_events`, whose key-type
chain ends at `scheduler_job_runner.py:1446` in `cls.logger().error("Unknown
workload key type in event buffer: %r", key)`.
That last one is worth fixing whatever you decide about the checklist,
because it is the only one of the three that fails quietly: a fourth type would
queue, dispatch and run, then lose its terminal events to a log line while both
sibling sites raise. (`_process_executor_events` at :1349 is only the wrapper;
the chain lives in the classmethod.)
Since the list has been wrong twice, a test asserting each family's
membership equals `set(WorkloadType)` would keep this honest without anyone
maintaining prose.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
