Avoid duplicate trigger-rule upstream-count queries per scheduling pass - #67672
Merged
Conversation
TriggerRuleDep runs a `SELECT task_id, count(*) ... GROUP BY task_id` per downstream task to size its upstream set, but only when an upstream is mapped. When many downstreams share the same mapped upstream, each issues an identical query within the same scheduling pass. Memoize the result on DepContext (one scheduling pass, same lifetime as finished_tis), keyed by (dag_id, run_id, frozenset of direct-upstream task_ids). Only the simple case is cached, where the predicate is exactly `task_id IN (upstream_ids)`; downstreams inside a mapped task group keep their own per-instance map-index query. The cache is cleared in _get_ready_tis when a mapped task expands and changes its instance count.
kaxil
added
the
full tests needed
_get_ready_tis already drops the trigger-rule upstream-count memo when a mapped task expands, but not when _revise_map_indexes_if_mapped grows an already-expanded task and adds new instances. A downstream evaluated later in the same pass could then read a stale, pre-grow count. Clear the memo in that case too, and add a regression test that drives _get_ready_tis with a fixed order so the count is recomputed after the grow (it asserts the query runs twice; fails without the clear).
Contributor
|
@kaxil static checks failing. Do we need this in 3.3.0? |
uranusjr
reviewed
Jun 23, 2026
uranusjr
approved these changes
Jun 23, 2026
Member
Author
It can wait for 3.3.1 |
are_dependencies_met rebuilds the DepContext with attrs.evolve for every UP_FOR_RESCHEDULE task instance, and UP_FOR_RESCHEDULE is in SCHEDULEABLE_STATES, so those instances reach TriggerRuleDep through the evolved context. attrs.evolve only carries over fields that __init__ accepts, so the memo's init=False meant each of them got a fresh empty dict: they neither read the memo nor warmed it for anything else. Reschedule-mode sensors fanned out behind a mapped upstream are the exact shape the memo collapses, so they were the one case opting out of it. Drop init=False (keeping the dict out of the repr instead), and route both invalidation sites through a DepContext method rather than having dagrun reach in and clear a dict only TriggerRuleDep writes. Also drop a redundant list() around _revise_map_indexes_if_mapped, which returns a list since apache#69565, and correct the staleness comment to name both invalidation points.
Merged
1 task
vatsrahul1001
added a commit
that referenced
this pull request
Aug 3, 2026
…ss (#67672) (#70826) TriggerRuleDep runs a `SELECT task_id, count(*) ... GROUP BY task_id` per downstream task to size its upstream set, but only when an upstream is mapped. When many downstreams share the same mapped upstream, each issues an identical query within the same scheduling pass. Memoize the result on DepContext (one scheduling pass, same lifetime as finished_tis), keyed by (dag_id, run_id, frozenset of direct-upstream task_ids). Only the simple case is cached, where the predicate is exactly `task_id IN (upstream_ids)`; downstreams inside a mapped task group keep their own per-instance map-index query. The cache is cleared in _get_ready_tis when a mapped task expands and changes its instance count. * Add newsfragment for trigger-rule upstream-count memoization * Clear upstream-count memo when a mapped task grows via revise _get_ready_tis already drops the trigger-rule upstream-count memo when a mapped task expands, but not when _revise_map_indexes_if_mapped grows an already-expanded task and adds new instances. A downstream evaluated later in the same pass could then read a stale, pre-grow count. Clear the memo in that case too, and add a regression test that drives _get_ready_tis with a fixed order so the count is recomputed after the grow (it asserts the query runs twice; fails without the clear). * Keep upstream-count memo across the UP_FOR_RESCHEDULE context rebuild are_dependencies_met rebuilds the DepContext with attrs.evolve for every UP_FOR_RESCHEDULE task instance, and UP_FOR_RESCHEDULE is in SCHEDULEABLE_STATES, so those instances reach TriggerRuleDep through the evolved context. attrs.evolve only carries over fields that __init__ accepts, so the memo's init=False meant each of them got a fresh empty dict: they neither read the memo nor warmed it for anything else. Reschedule-mode sensors fanned out behind a mapped upstream are the exact shape the memo collapses, so they were the one case opting out of it. Drop init=False (keeping the dict out of the repr instead), and route both invalidation sites through a DepContext method rather than having dagrun reach in and clear a dict only TriggerRuleDep writes. Also drop a redundant list() around _revise_map_indexes_if_mapped, which returns a list since #69565, and correct the staleness comment to name both invalidation points. (cherry picked from commit ecd4d95) Co-authored-by: Kaxil Naik <kaxilnaik@gmail.com>
vatsrahul1001
added a commit
that referenced
this pull request
Aug 5, 2026
…ss (#67672) (#70826) TriggerRuleDep runs a `SELECT task_id, count(*) ... GROUP BY task_id` per downstream task to size its upstream set, but only when an upstream is mapped. When many downstreams share the same mapped upstream, each issues an identical query within the same scheduling pass. Memoize the result on DepContext (one scheduling pass, same lifetime as finished_tis), keyed by (dag_id, run_id, frozenset of direct-upstream task_ids). Only the simple case is cached, where the predicate is exactly `task_id IN (upstream_ids)`; downstreams inside a mapped task group keep their own per-instance map-index query. The cache is cleared in _get_ready_tis when a mapped task expands and changes its instance count. * Add newsfragment for trigger-rule upstream-count memoization * Clear upstream-count memo when a mapped task grows via revise _get_ready_tis already drops the trigger-rule upstream-count memo when a mapped task expands, but not when _revise_map_indexes_if_mapped grows an already-expanded task and adds new instances. A downstream evaluated later in the same pass could then read a stale, pre-grow count. Clear the memo in that case too, and add a regression test that drives _get_ready_tis with a fixed order so the count is recomputed after the grow (it asserts the query runs twice; fails without the clear). * Keep upstream-count memo across the UP_FOR_RESCHEDULE context rebuild are_dependencies_met rebuilds the DepContext with attrs.evolve for every UP_FOR_RESCHEDULE task instance, and UP_FOR_RESCHEDULE is in SCHEDULEABLE_STATES, so those instances reach TriggerRuleDep through the evolved context. attrs.evolve only carries over fields that __init__ accepts, so the memo's init=False meant each of them got a fresh empty dict: they neither read the memo nor warmed it for anything else. Reschedule-mode sensors fanned out behind a mapped upstream are the exact shape the memo collapses, so they were the one case opting out of it. Drop init=False (keeping the dict out of the repr instead), and route both invalidation sites through a DepContext method rather than having dagrun reach in and clear a dict only TriggerRuleDep writes. Also drop a redundant list() around _revise_map_indexes_if_mapped, which returns a list since #69565, and correct the staleness comment to name both invalidation points. (cherry picked from commit ecd4d95) Co-authored-by: Kaxil Naik <kaxilnaik@gmail.com>
xBis7
pushed a commit
to xBis7/airflow
that referenced
this pull request
Aug 5, 2026
…ss (apache#67672) TriggerRuleDep runs a `SELECT task_id, count(*) ... GROUP BY task_id` per downstream task to size its upstream set, but only when an upstream is mapped. When many downstreams share the same mapped upstream, each issues an identical query within the same scheduling pass. Memoize the result on DepContext (one scheduling pass, same lifetime as finished_tis), keyed by (dag_id, run_id, frozenset of direct-upstream task_ids). Only the simple case is cached, where the predicate is exactly `task_id IN (upstream_ids)`; downstreams inside a mapped task group keep their own per-instance map-index query. The cache is cleared in _get_ready_tis when a mapped task expands and changes its instance count. * Add newsfragment for trigger-rule upstream-count memoization * Clear upstream-count memo when a mapped task grows via revise _get_ready_tis already drops the trigger-rule upstream-count memo when a mapped task expands, but not when _revise_map_indexes_if_mapped grows an already-expanded task and adds new instances. A downstream evaluated later in the same pass could then read a stale, pre-grow count. Clear the memo in that case too, and add a regression test that drives _get_ready_tis with a fixed order so the count is recomputed after the grow (it asserts the query runs twice; fails without the clear). * Keep upstream-count memo across the UP_FOR_RESCHEDULE context rebuild are_dependencies_met rebuilds the DepContext with attrs.evolve for every UP_FOR_RESCHEDULE task instance, and UP_FOR_RESCHEDULE is in SCHEDULEABLE_STATES, so those instances reach TriggerRuleDep through the evolved context. attrs.evolve only carries over fields that __init__ accepts, so the memo's init=False meant each of them got a fresh empty dict: they neither read the memo nor warmed it for anything else. Reschedule-mode sensors fanned out behind a mapped upstream are the exact shape the memo collapses, so they were the one case opting out of it. Drop init=False (keeping the dict out of the repr instead), and route both invalidation sites through a DepContext method rather than having dagrun reach in and clear a dict only TriggerRuleDep writes. Also drop a redundant list() around _revise_map_indexes_if_mapped, which returns a list since apache#69565, and correct the staleness comment to name both invalidation points. (cherry picked from commit ecd4d95) Backport note: on this branch `_revise_map_indexes_if_mapped` still returns a lazy iterator (apache#69565 is not backported), so the result is materialized with `list()` before the `if revised_tis:` truthiness check and extend. The upstream cleanup that removed `list()` does not apply here.
dabla
pushed a commit
to dabla/airflow
that referenced
this pull request
Aug 14, 2026
…ss (apache#67672) TriggerRuleDep runs a `SELECT task_id, count(*) ... GROUP BY task_id` per downstream task to size its upstream set, but only when an upstream is mapped. When many downstreams share the same mapped upstream, each issues an identical query within the same scheduling pass. Memoize the result on DepContext (one scheduling pass, same lifetime as finished_tis), keyed by (dag_id, run_id, frozenset of direct-upstream task_ids). Only the simple case is cached, where the predicate is exactly `task_id IN (upstream_ids)`; downstreams inside a mapped task group keep their own per-instance map-index query. The cache is cleared in _get_ready_tis when a mapped task expands and changes its instance count. * Add newsfragment for trigger-rule upstream-count memoization * Clear upstream-count memo when a mapped task grows via revise _get_ready_tis already drops the trigger-rule upstream-count memo when a mapped task expands, but not when _revise_map_indexes_if_mapped grows an already-expanded task and adds new instances. A downstream evaluated later in the same pass could then read a stale, pre-grow count. Clear the memo in that case too, and add a regression test that drives _get_ready_tis with a fixed order so the count is recomputed after the grow (it asserts the query runs twice; fails without the clear). * Keep upstream-count memo across the UP_FOR_RESCHEDULE context rebuild are_dependencies_met rebuilds the DepContext with attrs.evolve for every UP_FOR_RESCHEDULE task instance, and UP_FOR_RESCHEDULE is in SCHEDULEABLE_STATES, so those instances reach TriggerRuleDep through the evolved context. attrs.evolve only carries over fields that __init__ accepts, so the memo's init=False meant each of them got a fresh empty dict: they neither read the memo nor warmed it for anything else. Reschedule-mode sensors fanned out behind a mapped upstream are the exact shape the memo collapses, so they were the one case opting out of it. Drop init=False (keeping the dict out of the repr instead), and route both invalidation sites through a DepContext method rather than having dagrun reach in and clear a dict only TriggerRuleDep writes. Also drop a redundant list() around _revise_map_indexes_if_mapped, which returns a list since apache#69565, and correct the staleness comment to name both invalidation points.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
3 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
TriggerRuleDepsizes a task's upstream set with aSELECT task_id, count(*) ... GROUP BY task_id, and it only runs that when one of the upstreams is mapped. The problem: if several downstream tasks all depend on the same mapped upstream, each of them runs the identical query during the same scheduling pass. This caches the result on theDepContextfor the duration of a pass, so it runs once per distinct upstream set instead of once per downstream.Impact
The win scales with how many downstream tasks sit behind a mapped upstream. Measured on a DAG with one mapped upstream (5 instances) feeding 60 plain downstream tasks, counting just this query during one scheduling pass:
That count recurs every
update_statepass, so across a run it's roughlydownstreams x passesqueries collapsing topasses. Each query is individually cheap (sub-millisecond), so this is about cutting database round-trips, not query time. It happens in the per-run scheduling work, not the serialized critical section, so it's a latency win for mapping-heavy DAGs rather than a throughput-ceiling change.How it works
(dag_id, run_id, frozenset(upstream_task_ids)), stored onDepContext, which is already built once per_get_ready_tispass and reused for every TI in that pass (the same object that cachesfinished_tis).task_id IN (upstream_ids)and comes out the same for every downstream with the same upstreams. Tasks inside a mapped task group have per map-index predicates, so they keep running their own query.upstream_setupis still summed in the caller from the cached rows, so the setup count stays right per downstream.Staleness
It follows
finished_tis(a per-pass snapshot). Two things change a mapped task's instance count mid-pass and the memo is dropped for both: expansion of an unexpanded task, and_revise_map_indexes_if_mappedgrowing an already-expanded one (both in_get_ready_tis). So a downstream evaluated later in the same pass recomputes against the current count. State-only changes (a task finishing, or instances marked REMOVED) don't change the row count the memo is keyed on, so they need no invalidation.