hussein-awala opened a new pull request, #71967:
URL: https://github.com/apache/airflow/pull/71967
## Suggested PR title
Improve scheduler performance for large mapped task expansion
---
## Draft PR body
### Problem
When a mapped task expands to thousands of task instances in a single shot,
scheduler load can spike in one pass.
Current behavior has two expensive patterns:
1. **Bursty same-pass dependency evaluation**
- In `DagRun._get_ready_tis`, newly expanded mapped task instances are
immediately appended to `additional_tis` and dependency-checked in the same
scheduler pass.
- For very large expansions, this creates a sudden CPU/DB burst that can
starve other Dags.
2. **Per-row mapped TI creation overhead**
- `TaskMap.expand_mapped_task` and `DagRun._revise_map_indexes_if_mapped`
create missing mapped task instances one-by-one via ORM objects, even when
`task_instance_mutation_hook` is noop.
- This adds Python/ORM overhead proportional to mapped cardinality.
### What this PR changes
#### 1) Defer part of expanded mapped TIs to later scheduler passes
- **File:** `airflow-core/src/airflow/models/dagrun.py`
- **Method:** `DagRun._get_ready_tis`
- **Change:** Bound how many newly expanded mapped TIs are immediately
dependency-checked in the same pass.
- Budget source: `[scheduler] max_tis_per_query`
- Fallback when `<= 0`: `[core] parallelism`
- Expanded TIs beyond budget are already persisted, and are picked up in
later passes.
#### 2) Bulk-insert fast path when mutation hook is noop
- **File:** `airflow-core/src/airflow/models/taskmap.py`
- **Method:** `TaskMap.expand_mapped_task`
- **Change:** If `task_instance_mutation_hook.is_noop is True`, create
missing mapped TIs with `bulk_insert_mappings` using
`TaskInstance.insert_mapping`.
- **File:** `airflow-core/src/airflow/models/dagrun.py`
- **Method:** `DagRun._revise_map_indexes_if_mapped`
- **Change:** Same noop-only bulk insert fast path for map-index growth
during revision.
- Non-noop hook behavior remains unchanged (falls back to existing per-TI
path so custom hook semantics are preserved).
### Why this helps
- Reduces scheduler latency spikes during very large mapped expansions.
- Improves fairness by preventing a single expansion from dominating one
scheduler heartbeat.
- Reduces Python/ORM overhead in mapped TI fan-out when the mutation hook is
noop.
### Tests added/updated
- `airflow-core/tests/unit/models/test_dagrun.py`
- `test_mapped_expansion_defers_some_tis_to_later_scheduler_pass`
- Verifies expanded mapped TIs can be deferred to later pass under low
scheduler budget.
- `airflow-core/tests/unit/models/test_mappedoperator.py`
- `test_expand_mapped_task_uses_bulk_insert_when_mutation_hook_is_noop`
- Verifies noop hook uses bulk insert path and expected mapped indexes
are created.
- Existing behavior check still passing:
- `test_mapped_length_increase_at_runtime_adds_additional_tis`
- `test_expand_mapped_task_task_instance_mutation_hook` (classic +
taskflow)
### Scope / non-goals
This PR does not yet change:
- stale-index cleanup strategy in `TaskMap.expand_mapped_task` (still row
iteration)
- DagRun-wide TI materialization strategy
- mapped trigger-rule query complexity
Those can be follow-up optimizations.
---
### Was generative AI tooling used to co-author this PR?
- [X] Yes — Codex (GPT-5)
Generated-by: Codex (GPT-5) following [the
guidelines](https://github.com/apache/airflow/blob/main/contributing-docs/05_pull_requests.rst#gen-ai-assisted-contributions)
--
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]
