You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@airflow.apache.org by "steveahnahn (via GitHub)" <gi...@apache.org> on 2026/02/04 21:48:48 UTC
[PR] remove N+1 db queries for team names [airflow]
steveahnahn opened a new pull request, #61471:
URL: https://github.com/apache/airflow/pull/61471
## Summary
Optimize multi-team executor resolution in the scheduler by batching dag_id -> team_name lookups and passing the resolved team_name into executor selection, eliminating per-task DB queries when core.multi_team=true.
With multi-team enabled, the scheduler can process many task instances per loop. Previously, team resolution could trigger repeated database lookups (effectively N queries for N task instances) when selecting executors. This PR reduces that overhead by resolving team names once per unique DAG ID and reusing the results.
- Adds a batched lookup helper to resolve team names for multiple DAG IDs in one query
- Preserves legacy behavior when core.multi_team=false (no additional queries / logic).
- Add test_multi_team_executor_to_tis_batch_optimization to assert _executor_to_tis() performs a single batched query and does not call _get_task_team_name().
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] remove N+1 db queries for team names [airflow]
Posted by "steveahnahn (via GitHub)" <gi...@apache.org>.
steveahnahn commented on code in PR #61471:
URL: https://github.com/apache/airflow/pull/61471#discussion_r2771999424
##########
airflow-core/src/airflow/jobs/scheduler_job_runner.py:
##########
@@ -2925,7 +2931,11 @@ def _purge_task_instances_without_heartbeats(
request,
)
self.job.executor.send_callback(request)
- if (executor := self._try_to_load_executor(ti, session)) is None:
+ if (
+ executor := self._try_to_load_executor(
+ ti, session, team_name=dag_id_to_team_name.get(ti.dag_id, NOTSET)
+ )
+ ) is None:
Review Comment:
Addressed in 7ac6892
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] remove N+1 db queries for team names [airflow]
Posted by "o-nikolas (via GitHub)" <gi...@apache.org>.
o-nikolas commented on PR #61471:
URL: https://github.com/apache/airflow/pull/61471#issuecomment-3874410403
Thanks for the PR! This was a TODO on my end to update everything on the none core scheduling path (adoption, stuck tasks, etc) to batching, so I appreciate this!
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] remove N+1 db queries for team names [airflow]
Posted by "uranusjr (via GitHub)" <gi...@apache.org>.
uranusjr commented on code in PR #61471:
URL: https://github.com/apache/airflow/pull/61471#discussion_r2767642712
##########
airflow-core/src/airflow/jobs/scheduler_job_runner.py:
##########
@@ -2881,6 +2882,11 @@ def _find_task_instances_without_heartbeats(self, *, session: Session) -> list[T
def _purge_task_instances_without_heartbeats(
self, task_instances_without_heartbeats: list[TI], *, session: Session
) -> None:
+ dag_id_to_team_name: dict[str, str | None] = {}
+ if conf.getboolean("core", "multi_team"):
+ unique_dag_ids = {ti.dag_id for ti in task_instances_without_heartbeats}
+ dag_id_to_team_name = self._get_team_names_for_dag_ids(unique_dag_ids, session)
Review Comment:
```suggestion
if conf.getboolean("core", "multi_team"):
unique_dag_ids = {ti.dag_id for ti in task_instances_without_heartbeats}
dag_id_to_team_name = self._get_team_names_for_dag_ids(unique_dag_ids, session)
else:
dag_id_to_team_name = {}
```
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] remove N+1 db queries for team names [airflow]
Posted by "uranusjr (via GitHub)" <gi...@apache.org>.
uranusjr merged PR #61471:
URL: https://github.com/apache/airflow/pull/61471
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] remove N+1 db queries for team names [airflow]
Posted by "uranusjr (via GitHub)" <gi...@apache.org>.
uranusjr commented on code in PR #61471:
URL: https://github.com/apache/airflow/pull/61471#discussion_r2767640350
##########
airflow-core/src/airflow/jobs/scheduler_job_runner.py:
##########
@@ -2925,7 +2931,11 @@ def _purge_task_instances_without_heartbeats(
request,
)
self.job.executor.send_callback(request)
- if (executor := self._try_to_load_executor(ti, session)) is None:
+ if (
+ executor := self._try_to_load_executor(
+ ti, session, team_name=dag_id_to_team_name.get(ti.dag_id, NOTSET)
+ )
+ ) is None:
Review Comment:
Let’s split this to a separate `=` statement and then `if executor is None`, this is too long IMO.
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] remove N+1 db queries for team names [airflow]
Posted by "steveahnahn (via GitHub)" <gi...@apache.org>.
steveahnahn commented on code in PR #61471:
URL: https://github.com/apache/airflow/pull/61471#discussion_r2771999736
##########
airflow-core/src/airflow/jobs/scheduler_job_runner.py:
##########
@@ -2881,6 +2882,11 @@ def _find_task_instances_without_heartbeats(self, *, session: Session) -> list[T
def _purge_task_instances_without_heartbeats(
self, task_instances_without_heartbeats: list[TI], *, session: Session
) -> None:
+ dag_id_to_team_name: dict[str, str | None] = {}
+ if conf.getboolean("core", "multi_team"):
+ unique_dag_ids = {ti.dag_id for ti in task_instances_without_heartbeats}
+ dag_id_to_team_name = self._get_team_names_for_dag_ids(unique_dag_ids, session)
Review Comment:
Addressed in 7ac6892, thank you!
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org