Unify executor workload queues

[anishgirianish]

---

Was generative AI tooling used to co-author this PR?

• Yes (please specify the tool below)

Summary

Refactors executor workload queue management for extensibility. No behavioral change , scheduling order, slot accounting, and all provider executors work identically to before.

Follows the direction proposed by @ferruzzi #62343 (comment).

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 simple primitives:

• executor_queues: a single defaultdict(dict) keyed by workload type string (e.g. "ExecuteTask","ExecuteCallback") instead of separatequeued_tasks / queued_callbacks dicts
• supported_workload_types: a frozenset of type strings instead of individual supports_callbacks booleans
• WORKLOAD_TYPE_PRIORITY + sort_key / queue_key : properties on each workload schema that control scheduling priority and queue indexing

The base class queue_workload is now generic: validate the type, store by key. Four provider executors (K8s, ECS, Batch, Lambda) no longer need their own queue_workload overrides. trigger_tasks becomes trigger_workloads since it handles all workload types now.

Adding a new workload type after this refactor

1. Define queue_key and sort_key on the workload schema
2. Add the type string to supported_workload_types on supporting executors
3. Handle the type in _process_workloads, done

No changes needed in BaseExecutor itself.

---

• Read the Pull Request Guidelines for more information. Note: commit author/co-author name and email in commits become permanently public when merged.
• For fundamental code changes, an Airflow Improvement Proposal (AIP) is needed.
• When adding dependency, check compliance with the ASF 3rd Party License Policy.
• For significant user-facing changes create newsfragment: {pr_number}.significant.rst, in airflow-core/newsfragments. You can add this file in a follow-up commit after the PR is created so you know the PR number.

---

Important

🛠️ Maintainer triage note for @anishgirianish · by @potiuk · 2026-07-08 15:51 UTC

Some review feedback from @eladkal is waiting on you (2 unresolved threads):

The ball is in your court — you've been assigned to this PR. Reply or push a fix in each thread, then mark them resolved. See the Pull Request quality criteria.

_{Automated triage — may be imperfect; a maintainer takes the next look.}

[ferruzzi]

Made a real quick pass and left some comments and questions, I'll try to get a more thorough one tomorrow.

[anishgirianish]

Both of Ash's blocking comments have been replied to and I think he's had a ton of time to reply and several pings at this point. I'm going to dismiss his blocking change request and he can open a new one if he wants. That said, I think there is a merge conflict that you mixed up, so I left a comment to look at. If it's intentional, then I'll need you to explain it to me, but I suspect it's just a merge conflict gone wrong.

Hi @ferruzzi @potiuk thank you so much for re-review. It was actually jumbled up while fixing conflicts and thank you for catching it. I have updated the pr, would like to request you for your re-review thanks.

[ferruzzi]

LGTM

[kaxil]

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.

[kaxil]

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:

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.

[kaxil]

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:

# 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.

[kaxil]

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.

[kaxil]

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 KeyErrors 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.

[kaxil]

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?

[kaxil]

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.

[kaxil]

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.)

[kaxil]

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.

[kaxil]

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.