You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@airflow.apache.org by "o-nikolas (via GitHub)" <gi...@apache.org> on 2025/11/05 00:27:23 UTC
[PR] Introduce multi-team task scheduling [airflow]
o-nikolas opened a new pull request, #57837:
URL: https://github.com/apache/airflow/pull/57837
Detect the teams that tasks belong to be resolving the relationship between task --> dag --> dag bundle --> team (we are unfortunately not allowed to add a team column to the task model).
Hide much of the new functionality behind the core.multi_team feature flag. This is to not affect the normal scheduling loop with the new queries for team-ness.
<!--
Licensed to the Apache Software Foundation (ASF) under one
or more contributor license agreements. See the NOTICE file
distributed with this work for additional information
regarding copyright ownership. The ASF licenses this file
to you under the Apache License, Version 2.0 (the
"License"); you may not use this file except in compliance
with the License. You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing,
software distributed under the License is distributed on an
"AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
KIND, either express or implied. See the License for the
specific language governing permissions and limitations
under the License.
-->
<!--
Thank you for contributing! Please make sure that your code changes
are covered with tests. And in case of new features or big changes
remember to adjust the documentation.
Feel free to ping committers for the review!
In case of an existing issue, reference it using one of the following:
closes: #ISSUE
related: #ISSUE
How to write a good git commit message:
http://chris.beams.io/posts/git-commit/
-->
<!-- Please keep an empty line above the dashes. -->
---
**^ Add meaningful description above**
Read the **[Pull Request Guidelines](https://github.com/apache/airflow/blob/main/contributing-docs/05_pull_requests.rst#pull-request-guidelines)** for more information.
In case of fundamental code changes, an Airflow Improvement Proposal ([AIP](https://cwiki.apache.org/confluence/display/AIRFLOW/Airflow+Improvement+Proposals)) is needed.
In case of a new dependency, check compliance with the [ASF 3rd Party License Policy](https://www.apache.org/legal/resolved.html#category-x).
In case of backwards incompatible changes please leave a note in a newsfragment file, named `{pr_number}.significant.rst` or `{issue_number}.significant.rst`, in [airflow-core/newsfragments](https://github.com/apache/airflow/tree/main/airflow-core/newsfragments).
--
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] AIP-67 multi-team: Introduce multi-team task scheduling [airflow]
Posted by "vincbeck (via GitHub)" <gi...@apache.org>.
vincbeck commented on code in PR #57837:
URL: https://github.com/apache/airflow/pull/57837#discussion_r2495046425
##########
airflow-core/src/airflow/jobs/job.py:
##########
@@ -139,6 +139,7 @@ def __init__(self, executor: BaseExecutor | None = None, heartrate=None, **kwarg
# Save init parameters as DB fields
self.heartbeat_failed = False
self.hostname = get_hostname()
+ # TODO[multi-team]: Make this multi-team?
Review Comment:
We want to keep it?
##########
airflow-core/src/airflow/jobs/scheduler_job_runner.py:
##########
@@ -270,6 +272,98 @@ def register_signals(self) -> ExitStack:
return resetter
+ def _get_team_names_for_dag_ids(
+ self, dag_ids: Collection[str], session: Session
+ ) -> dict[str, str | None]:
+ """
+ Batch query to resolve team names for multiple DAG IDs using the DAG > Bundle > Team relationship chain.
+
+ DAG IDs > DagModel (via dag_id) > DagBundleModel (via bundle_name) > Team
+
+ :param dag_ids: Collection of DAG IDs to resolve team names for
+ :param session: Database session for queries
+ :return: Dictionary mapping dag_id -> team_name (None if no team found)
+ """
+ if not dag_ids:
+ return {}
+
+ try:
+ # Query all team names for the given DAG IDs in a single query
+ query_results = session.execute(
+ select(DagModel.dag_id, Team.name)
+ .join(DagBundleModel.teams) # Join Team to DagBundleModel via association table
+ .join(
+ DagModel, DagModel.bundle_name == DagBundleModel.name
+ ) # Join DagBundleModel to DagModel
+ .where(DagModel.dag_id.in_(dag_ids))
+ ).all()
+
+ # Create mapping from results
+ dag_id_to_team_name = {dag_id: team_name for dag_id, team_name in query_results}
+
+ # Ensure all requested dag_ids are in the result (with None for those not found)
+ result = {dag_id: dag_id_to_team_name.get(dag_id) for dag_id in dag_ids}
+
+ self.log.debug(
+ "Resolved team names for %d DAGs: %s",
+ len([team for team in result.values() if team is not None]),
+ {dag_id: team for dag_id, team in result.items()},
+ )
+
+ return result
+
+ except Exception:
+ # Log the error, explicitly don't fail the scheduling loop
+ self.log.exception("Failed to resolve team names for DAG IDs: %s", list(dag_ids))
+ # Return dict with all None values to ensure graceful degradation
+ return {}
+
+ def _get_task_team_name(self, task_instance: TaskInstance, session: Session) -> str | None:
Review Comment:
Would it be worth it to use `_get_team_names_for_dag_ids` inside this function to avoid duplication?
--
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] AIP-67 multi-team: Introduce multi-team task scheduling [airflow]
Posted by "o-nikolas (via GitHub)" <gi...@apache.org>.
o-nikolas merged PR #57837:
URL: https://github.com/apache/airflow/pull/57837
--
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] AIP-67 multi-team: Introduce multi-team task scheduling [airflow]
Posted by "o-nikolas (via GitHub)" <gi...@apache.org>.
o-nikolas commented on code in PR #57837:
URL: https://github.com/apache/airflow/pull/57837#discussion_r2501072232
##########
airflow-core/src/airflow/jobs/scheduler_job_runner.py:
##########
@@ -270,6 +272,98 @@ def register_signals(self) -> ExitStack:
return resetter
+ def _get_team_names_for_dag_ids(
+ self, dag_ids: Collection[str], session: Session
+ ) -> dict[str, str | None]:
+ """
+ Batch query to resolve team names for multiple DAG IDs using the DAG > Bundle > Team relationship chain.
+
+ DAG IDs > DagModel (via dag_id) > DagBundleModel (via bundle_name) > Team
+
+ :param dag_ids: Collection of DAG IDs to resolve team names for
+ :param session: Database session for queries
+ :return: Dictionary mapping dag_id -> team_name (None if no team found)
+ """
+ if not dag_ids:
+ return {}
+
+ try:
+ # Query all team names for the given DAG IDs in a single query
+ query_results = session.execute(
+ select(DagModel.dag_id, Team.name)
+ .join(DagBundleModel.teams) # Join Team to DagBundleModel via association table
+ .join(
+ DagModel, DagModel.bundle_name == DagBundleModel.name
+ ) # Join DagBundleModel to DagModel
+ .where(DagModel.dag_id.in_(dag_ids))
+ ).all()
+
+ # Create mapping from results
+ dag_id_to_team_name = {dag_id: team_name for dag_id, team_name in query_results}
+
+ # Ensure all requested dag_ids are in the result (with None for those not found)
+ result = {dag_id: dag_id_to_team_name.get(dag_id) for dag_id in dag_ids}
+
+ self.log.debug(
+ "Resolved team names for %d DAGs: %s",
+ len([team for team in result.values() if team is not None]),
+ {dag_id: team for dag_id, team in result.items()},
+ )
+
+ return result
+
+ except Exception:
+ # Log the error, explicitly don't fail the scheduling loop
+ self.log.exception("Failed to resolve team names for DAG IDs: %s", list(dag_ids))
+ # Return dict with all None values to ensure graceful degradation
+ return {}
+
+ def _get_task_team_name(self, task_instance: TaskInstance, session: Session) -> str | None:
Review Comment:
Great idea! I'll do it
--
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] AIP-67 multi-team: Introduce multi-team task scheduling [airflow]
Posted by "o-nikolas (via GitHub)" <gi...@apache.org>.
o-nikolas commented on code in PR #57837:
URL: https://github.com/apache/airflow/pull/57837#discussion_r2501071186
##########
airflow-core/src/airflow/jobs/job.py:
##########
@@ -139,6 +139,7 @@ def __init__(self, executor: BaseExecutor | None = None, heartrate=None, **kwarg
# Save init parameters as DB fields
self.heartbeat_failed = False
self.hostname = get_hostname()
+ # TODO[multi-team]: Make this multi-team?
Review Comment:
I didn't end up needing it for tests, and I'm sure I'll remember if I do in the future. So I'll remove this one!
--
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