Skip to content

Unify executor workload queues - #63491

Merged
kaxil merged 22 commits into
apache:mainfrom
anishgirianish:refactor-workload-queue
Sep 21, 2026
Merged

kaxil merged 22 commits into
apache:mainfrom
anishgirianish:refactor-workload-queue

Conversation

@anishgirianish

@anishgirianish anishgirianish commented Mar 12, 2026 •

Copy link
Copy Markdown
Contributor

Was generative AI tooling used to co-author this PR?
  • Yes (please specify the tool below)

Summary

Refactors BaseExecutor workload queue management so that adding a new workload type is a small.

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

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

@boring-cyborg boring-cyborg Bot added area:Executors-core LocalExecutor & SequentialExecutor area:providers provider:amazon AWS/Amazon - related issues provider:celery provider:cncf-kubernetes Kubernetes (k8s) provider related issues provider:edge Edge Executor / Worker (AIP-69) / edge3 labels Mar 12, 2026
Comment thread airflow-core/src/airflow/executors/base_executor.py Outdated
Comment thread airflow-core/src/airflow/executors/base_executor.py

@ferruzzi ferruzzi left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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

Comment thread airflow-core/src/airflow/executors/workloads/base.py Outdated
Comment thread airflow-core/src/airflow/executors/base_executor.py Outdated
Comment thread airflow-core/src/airflow/executors/base_executor.py
Comment thread airflow-core/src/airflow/executors/base_executor.py Outdated
Comment thread airflow-core/tests/unit/executors/test_base_executor.py
Comment thread airflow-core/src/airflow/executors/base_executor.py
@anishgirianish
anishgirianish force-pushed the refactor-workload-queue branch 3 times, most recently from aee94fb to 8997ee4 Compare March 13, 2026 06:02
@anishgirianish
anishgirianish marked this pull request as draft March 13, 2026 07:43
@anishgirianish
anishgirianish force-pushed the refactor-workload-queue branch 6 times, most recently from 11ee7ef to 249b014 Compare March 14, 2026 04:40
@anishgirianish
anishgirianish marked this pull request as ready for review March 14, 2026 05:45
Comment thread airflow-core/src/airflow/executors/base_executor.py Outdated
Comment thread airflow-core/src/airflow/executors/base_executor.py
Comment thread airflow-core/src/airflow/executors/workloads/base.py Outdated
Comment thread airflow-core/src/airflow/executors/base_executor.py Outdated
Comment thread airflow-core/newsfragments/63491.significant.rst
@anishgirianish

Copy link
Copy Markdown
Contributor Author

@kaxil thank you very much for the review. I have addressed all the feedbacks, Would like to request you for another round when ever you get a chance.

thank you

Comment thread airflow-core/src/airflow/executors/base_executor.py Outdated
Comment thread airflow-core/tests/unit/executors/test_base_executor.py Outdated
Comment thread airflow-core/src/airflow/executors/base_executor.py Outdated
Comment thread airflow-core/src/airflow/executors/workloads/base.py Outdated
Comment thread airflow-core/newsfragments/63491.significant.rst Outdated
…add a note to clearify the change in queued task priority in base executer
The new API ships in Airflow 3.4, not 3.3 as assumed when the guards
were written, so providers failed to import against released 3.3.0.
The rebase onto main brought in apache#70942, which flipped queued-task
dispatch to highest priority_weight first, but the unified sort_key
still sorted ascending and would have silently reverted that fix.
@anishgirianish

Copy link
Copy Markdown
Contributor Author

@kaxil thank you so much for the re-review and the merge. I really appreciate it. thank you

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

area:Executors-core LocalExecutor & SequentialExecutor area:providers provider:amazon AWS/Amazon - related issues provider:celery provider:cncf-kubernetes Kubernetes (k8s) provider related issues provider:edge Edge Executor / Worker (AIP-69) / edge3 ready for maintainer review Set after triaging when all criteria pass.

Projects

None yet

Development

Successfully merging this pull request may close these issues.

9 participants