kaxil commented on code in PR #63491:
URL: https://github.com/apache/airflow/pull/63491#discussion_r3740488002
##########
airflow-core/src/airflow/executors/base_executor.py:
##########
@@ -472,6 +491,29 @@ def trigger_tasks(self, open_slots: int) -> None:
if workload_list:
self._process_workloads(workload_list)
+ def trigger_tasks(self, open_slots: int) -> None:
Review Comment:
Following up on my April comment about `trigger_tasks` overrides: the shim
runs the opposite direction from what's needed. It forwards `trigger_tasks` to
`trigger_workloads`, which helps a *caller*, but `heartbeat()` calls
`self.trigger_workloads(open_slots)` directly, so a subclass that *overrides*
`trigger_tasks` is never invoked, and nothing warns.
`order_queued_tasks_by_priority` has the same shape. Before this PR,
`_get_workloads_to_schedule` called `self.order_queued_tasks_by_priority()`, so
an override was honored. Now the sort is inlined and the shim is never called.
Detecting either name in `cls.__dict__` from `__init_subclass__` and
warning, the way the `supports_callbacks` branch does, would at least make the
break loud instead of silent.
##########
airflow-core/src/airflow/executors/base_executor.py:
##########
@@ -210,6 +211,30 @@ def jwt_generator(self) -> JWTGenerator:
return generator
+ def __init_subclass__(cls, **kwargs: Any) -> None:
+ super().__init_subclass__(**kwargs)
+ cls._legacy_warned = set()
+ legacy_flag = cls.__dict__.get("supports_callbacks")
Review Comment:
Two things in this block.
First, `supports_connection_test` doesn't get the treatment
`supports_callbacks` gets. It shipped in 3.3.0 as a public `BaseExecutor`
attribute right alongside `supports_callbacks`, and it's removed here with no
property shim and no `__init_subclass__` branch. A 3.3-era executor declaring
`supports_connection_test = True` still imports fine, but
`supported_workload_types` stays at the default `frozenset({EXECUTE_TASK})`, so
the scheduler fails every connection test routed to it with "Executor 'X' does
not support connection testing", with nothing to explain why.
Second, the synthesis overwrites rather than unions:
```python
cls.supported_workload_types = frozenset({WorkloadType.EXECUTE_TASK,
WorkloadType.EXECUTE_CALLBACK})
```
`LocalExecutor` declares three types including `TEST_CONNECTION`. For `class
MyLocal(LocalExecutor): supports_callbacks = True`,
`"supported_workload_types"` is not in `MyLocal.__dict__` (it's on the parent),
so the guard passes and `TEST_CONNECTION` is dropped from a subclass whose only
change was using the deprecated spelling. `cls.supported_workload_types |
{WorkloadType.EXECUTE_CALLBACK}` would cover that half.
##########
airflow-core/src/airflow/executors/base_executor.py:
##########
@@ -210,6 +211,30 @@ def jwt_generator(self) -> JWTGenerator:
return generator
+ def __init_subclass__(cls, **kwargs: Any) -> None:
+ super().__init_subclass__(**kwargs)
+ cls._legacy_warned = set()
+ legacy_flag = cls.__dict__.get("supports_callbacks")
+ if legacy_flag is True:
+ warnings.warn(
Review Comment:
These warnings won't reach the executors they're aimed at.
`RemovedInAirflow4Warning` subclasses `DeprecationWarning`, and the only
filter re-enabling those is scoped to modules matching `airflow`:
```python
# configuration.py:58
warnings.filterwarnings(action="default", category=DeprecationWarning,
module="airflow")
```
With `stacklevel=2` the warning is attributed to the subclass definition
site, so for an executor in something like `mycompany.executors` it falls
through to CPython's default `ignore::DeprecationWarning`. The newsfragment
says legacy `supports_callbacks` attributes "are still honored ... while
emitting a deprecation warning", but an out-of-tree executor sees nothing
across the whole 3.4 line and then breaks at the 4.0 removal. Deduping once per
class via `_legacy_warned` narrows it further.
Pairing each `warnings.warn` with a `log.warning` would land it in scheduler
logs regardless of filter state.
##########
airflow-core/src/airflow/executors/base_executor.py:
##########
@@ -262,66 +333,45 @@ def log_task_event(self, *, event: str, extra: str,
ti_key: WorkloadKey):
return
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."
+ def queue_workload(self, workload: QueueableWorkload, session: Session) ->
None:
+ if workload.type not in self.supported_workload_types:
+ raise NotImplementedError(
+ f"{type(self).__name__} does not support {workload.type!r}
workloads. "
+ f"Add {workload.type!r} 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]]:
+ def _get_workloads_to_schedule(self, open_slots: int) ->
list[tuple[WorkloadKey, QueueableWorkload]]:
Review Comment:
This changes connection-test admission, so the "no behavioural change"
framing in the description isn't quite accurate.
Before, `heartbeat()` ran two independent budgets:
`trigger_tasks(open_slots)`, then `trigger_connection_tests()` gated on
`slots_available`, which subtracts every queued item. Now there's a single
budget and `TEST_CONNECTION` sorts last.
It diverges in both directions:
- parallelism 32, 40 queued tasks, 1 connection test. Before: 32 tasks
dispatched, then `slots_available = 32 - 0 - 8 - 0 - 1 = 23`, so the test ran
in the same heartbeat. Now it sorts to index 40 and the `[:32]` slice drops it.
Under a sustained backlog it never runs and the reaper times it out.
- parallelism 10, 3 tasks, 20 connection tests. Before: `slots_available =
10 - 0 - 0 - 0 - 20`, so zero tests ran. Now 3 tasks and 7 tests run.
Losing the old over-subscription is an improvement. The ordering is the part
worth another look: connection tests are short and user-interactive, and
`LocalExecutor` is the only in-tree executor supporting them, so a task backlog
starving them is the whole feature. Either move `TEST_CONNECTION` ahead of
`EXECUTE_TASK` in `_workload_type_priority_order` or reserve a small budget for
it, and either way the description needs amending.
##########
airflow-core/src/airflow/executors/base_executor.py:
##########
@@ -223,9 +248,9 @@ def __init__(self, parallelism: int = PARALLELISM,
team_name: str | None = None)
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, QueueableWorkload]
once the deprecated
+ # queued_tasks / queued_callbacks compat properties are removed.
+ self.executor_queues: dict[WorkloadType, dict[WorkloadKey,
QueueableWorkload]] = defaultdict(dict)
Review Comment:
The annotation says `dict` but this is a `defaultdict`, and the difference
is load-bearing. `queue_workload` (`:343`), `fail_connection_test` (`:416`) and
both compat getters index directly and rely on auto-vivification, while
`has_task` (`:385`) defensively uses `.get(..., {})`. That inconsistency inside
one class is the tell.
It leaks downstream. `kubernetes_executor.py:413`'s `if
self.executor_queues:` is permanently truthy after the first task, because the
outer key survives while the inner dict empties, so that debug line now fires
on every `sync()`. `debug_dump` has the inverse problem: iterating only
vivified keys means an idle executor prints no queue lines at all, where before
it always printed the counts.
It also traps anyone migrating off the compat property, since
`self.executor_queues = {}` is legal per the annotation and `KeyError`s on the
next `queue_workload`. Either annotate it as a `defaultdict`/`MutableMapping`
and mention it in the newsfragment, or drop the `defaultdict` and use
`setdefault` at the write sites.
##########
airflow-core/src/airflow/executors/workloads/base.py:
##########
@@ -32,6 +33,32 @@
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.
Review Comment:
Step 3 sends the next author somewhere that doesn't resolve.
`QueueableWorkload` is defined in `workloads/types.py` under `if
TYPE_CHECKING`, isn't imported by `workloads/__init__.py`, and isn't in its
`__all__`, so `airflow.executors.workloads.QueueableWorkload` raises
`AttributeError`. It also isn't a discriminated union, it's a bare `X | Y | Z`,
and `queue_workload` never consults it at runtime, it checks `workload.type not
in self.supported_workload_types`.
The two unions that do gate deserialization aren't mentioned: `All` and
`ExecutorWorkload` in `workloads/__init__.py`, both carrying
`Field(discriminator="type")`. `ExecutorWorkload` is what Celery's
`TypeAdapter` decodes with, so a fourth workload type added by following this
checklist would queue fine and then fail validation on every dequeue, which is
the failure this comment exists to prevent.
Worth asking separately: `ExecutorWorkload` already has identical membership
to the new alias, and `base_executor.py` still uses it in `run_workload`. Is
the second alias earning its keep, or could the queue-facing signatures just
use `ExecutorWorkload`?
##########
airflow-core/src/airflow/executors/workloads/base.py:
##########
@@ -168,3 +195,15 @@ def running_state(self) -> WorkloadState | None:
no intermediate state is emitted.
"""
return None
+
+ @property
+ def sort_key(self) -> int:
Review Comment:
This default lands on `BaseDagBundleWorkload`, but `TestConnection` extends
`BaseWorkloadSchema` directly and so doesn't inherit it. That's why there's a
second copy at `connection_test.py:50`.
`_get_workloads_to_schedule` sorts on `item[1].sort_key` across everything
in the queues, so the next queueable workload following `TestConnection`'s
pattern raises `AttributeError` in the scheduler's sort. Moving the default up
to `BaseWorkloadSchema` drops the duplicate and closes that in one edit.
`sort_key` is arguably a fourth entry for the checklist above too.
##########
providers/edge3/tests/unit/edge3/executors/test_edge_executor.py:
##########
@@ -63,7 +63,7 @@ def get_test_executor(self, pool_slots=1):
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]
Review Comment:
Not about this line, but it's the only edge3 line in the diff, so it's the
nearest place to raise it.
`edge_executor.py:372` in `revoke_task` still calls
`self.queued_tasks.pop(ti.key, None)`. The byte-identical line got an
`AIRFLOW_V_3_4_PLUS` branch in celery (`celery_executor.py:411`) and
cncf-kubernetes (`kubernetes_executor.py:1001`), but edge3's source wasn't
touched, only this test.
`revoke_task` is a live 3.x path, so on 3.4 Airflow's own bundled provider
emits `RemovedInAirflow4Warning` from the scheduler's revoke path, which is the
noise the once-per-class throttle exists to reduce. `RemovedInAirflow4Warning`
isn't in `forbidden_warnings`, so CI won't catch it either.
(`edge_executor.py:97` is the same pattern but sits in `_process_tasks`, which
is Airflow 2 only.)
##########
providers/amazon/tests/unit/amazon/aws/executors/batch/test_batch_executor.py:
##########
@@ -202,16 +206,19 @@ def test_execute(self, mock_executor):
mock_executor.batch.submit_job.assert_called_once()
assert len(mock_executor.active_workers) == 1
- @pytest.mark.skipif(not AIRFLOW_V_3_0_PLUS, reason="Test requires Airflow
3+")
+ @pytest.mark.skipif(not AIRFLOW_V_3_4_PLUS, reason="Test requires Airflow
3.4+")
Review Comment:
This gate moved from `AIRFLOW_V_3_0_PLUS` to `AIRFLOW_V_3_4_PLUS` in the
same PR that added the pre-3.4 `else` branches these tests were covering, so
the back-compat arm now runs on no CI leg. The provider compatibility matrix
runs unit tests on 2.11.1, 3.0.6, 3.1.8, 3.2.2 and 3.3.0, all above the old
gate and all below the new one.
Same flip in `test_lambda_executor.py:133/186/239`,
`test_batch_executor.py:281/351` and `test_ecs_executor.py:421`. Net effect:
Lambda and Batch have no sub-3.4 coverage of the `del self.queued_*` branches,
ECS keeps only its callback path, and `KubernetesExecutor._process_workloads`
has no test on any version.
This PR already contains the pattern that avoids it:
`test_gauge_executor_metrics` branches the mock target on the version rather
than skipping. Doing that here keeps both arms covered.
Separately, the dispatch half of `TEST_CONNECTION` is unasserted. The
surviving tests cover `queue_workload` accept/reject and the scheduler-side
enqueue, and `test_trigger_connection_tests_skipped_when_not_supported` was
deleted without a replacement, so nothing checks that a queued `TestConnection`
reaches `_process_workloads`.
##########
airflow-core/newsfragments/63491.significant.rst:
##########
@@ -0,0 +1,23 @@
+Deprecate ``BaseExecutor.queued_tasks``, ``queued_callbacks``,
``supports_callbacks``, ``trigger_tasks``, and
``order_queued_tasks_by_priority``
Review Comment:
Two gaps here.
The connection-test half of the refactor isn't listed.
`supports_connection_test`, `queued_connection_tests` and
`trigger_connection_tests()` were all public `BaseExecutor` surface in 3.3.0
and are gone in this PR, so a custom-executor author reading this won't find
the `executor_queues[WorkloadType.TEST_CONNECTION]` and
`supported_workload_types` replacements.
And the `order_queued_tasks_by_priority` line points at
`_get_workloads_to_schedule`, a private method whose return shape differs, as
the shim's own docstring says: the old one returned all queued tasks,
tasks-only and untruncated, while the replacement includes other workload types
and truncates to open slots. Following that line literally means depending on a
private symbol and silently changing behaviour. Either promote a public
replacement or carry the caveat into the note.
--
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]
