This is an automated email from the ASF dual-hosted git repository.
o-nikolas pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/airflow.git
The following commit(s) were added to refs/heads/main by this push:
new 0abd1616919 Run team tasks on the team executor when named by class or
path (#75362)
0abd1616919 is described below
commit 0abd1616919bd23ac16c872cdbb5012ebfe58645
Author: Niko Oliveira <[email protected]>
AuthorDate: Thu Oct 22 16:21:34 2026 -0700
Run team tasks on the team executor when named by class or path (#75362)
With multi-team enabled, a team task that names its executor by class
name or module path (e.g. executor="CeleryExecutor") is documented to
run on its team's executor of that class. Only the global alias should
reach the global executor ("Aliasing Executors Across Teams" in the
multi-team docs).
The scheduler took the first matching executor instead, and global
executors are always listed before team ones. So these tasks silently
ran on the shared global executor, outside the team boundary, and
anything resolving executors the documented way (such as reading live
task logs) looked in the wrong place.
---
.../src/airflow/jobs/scheduler_job_runner.py | 48 +++++-------
airflow-core/tests/unit/jobs/test_scheduler_job.py | 85 +++++++++++++++++++---
2 files changed, 94 insertions(+), 39 deletions(-)
diff --git a/airflow-core/src/airflow/jobs/scheduler_job_runner.py
b/airflow-core/src/airflow/jobs/scheduler_job_runner.py
index 5a56f7bfc33..198ef4bdc0a 100644
--- a/airflow-core/src/airflow/jobs/scheduler_job_runner.py
+++ b/airflow-core/src/airflow/jobs/scheduler_job_runner.py
@@ -71,7 +71,7 @@ from airflow.callbacks.callback_requests import (
)
from airflow.configuration import conf
from airflow.dag_processing.bundles.base import BundleUsageTrackingManager
-from airflow.exceptions import DagNotFound, TaskNotFound
+from airflow.exceptions import DagNotFound, TaskNotFound,
UnknownExecutorException
from airflow.executors import workloads
from airflow.executors.executor_loader import ExecutorLoader
from airflow.executors.workloads.types import TaskInstanceUuid
@@ -4560,40 +4560,32 @@ class SchedulerJobRunner(BaseJobRunner, LoggingMixin):
:param team_name: Optional pre-resolved team name. If NOTSET and
multi-team is enabled,
will query the database to resolve team name. None
indicates global team.
"""
- executor = None
if self._multi_team:
# Use provided team_name if available, otherwise query the database
if team_name is NOTSET:
team_name = self._get_workload_team_name(workload, session)
else:
team_name = None
- # If there is no executor set on the workload fetch the default
(either globally or for the team)
- if workload.get_executor_name() is None:
- if not team_name:
- # No team is specified, use the global default executor
- executor = self.executor
- else:
- # We do have a team, use the default executor for that team
- for _executor in self.executors:
- # First executor that resolves should be the default for
that team
- if _executor.team_name == team_name:
- executor = _executor
- break
- else:
- # No executor found for that team, fall back to global
default
- executor = self.executor
+
+ executor: BaseExecutor | None
+ executor_name = workload.get_executor_name()
+ if executor_name is None and not team_name:
+ # No team and no executor named on the workload, so it's always
the global default.
+ executor = self.executor
else:
- # An executor is specified on the workload (as a str), so we need
to find it in the list of executors
- for _executor in self.executors:
- if _executor.name and workload.get_executor_name() in (
- _executor.name.alias,
- _executor.name.module_path,
- _executor.name.module_path.split(".")[-1],
- ):
- # The executor must either match the team or be global
(i.e. team_name is None)
- if team_name and _executor.team_name == team_name or
_executor.team_name is None:
- executor = _executor
- break
+ # The resolver below only sees names, so keep a mapping from each
name to its executor.
+ executors_by_name: dict[ExecutorName, BaseExecutor] = {
+ _executor.name: _executor for _executor in self.executors if
_executor.name
+ }
+ try:
+ resolved_name = ExecutorLoader.resolve_executor_name(
+ executor_name, team_name, executor_names=executors_by_name
+ )
+ except UnknownExecutorException:
+ # The name matches nothing this team may use, e.g. another
team's executor. Warned below.
+ executor = None
+ else:
+ executor = executors_by_name[resolved_name]
if executor is not None:
self.log.debug(
diff --git a/airflow-core/tests/unit/jobs/test_scheduler_job.py
b/airflow-core/tests/unit/jobs/test_scheduler_job.py
index 84bb06cc768..9f7d83b9787 100644
--- a/airflow-core/tests/unit/jobs/test_scheduler_job.py
+++ b/airflow-core/tests/unit/jobs/test_scheduler_job.py
@@ -363,6 +363,12 @@ def make_pool_stats(
}
+def assign_executor_to_team(executor: MagicMock, team_name: str) -> None:
+ # ExecutorLoader.load_executor sets the team on both; the scheduler reads
it from the name.
+ executor.team_name = team_name
+ executor.name.team_name = team_name
+
+
@patch.dict(
ExecutorLoader.executors, {MOCK_EXECUTOR:
f"{MockExecutor.__module__}.{MockExecutor.__qualname__}"}
)
@@ -2392,8 +2398,8 @@ class TestSchedulerJob:
session.add_all([bundle1, bundle2])
session.flush()
- mock_executors[0].team_name = "team_a"
- mock_executors[1].team_name = "team_b"
+ assign_executor_to_team(mock_executors[0], "team_a")
+ assign_executor_to_team(mock_executors[1], "team_b")
with dag_maker(dag_id="dag_a", bundle_name="bundle_a",
session=session):
op1 = EmptyOperator(task_id="task_a_default") # No explicit
executor - should use team's default
@@ -4653,6 +4659,8 @@ class TestSchedulerJob:
dag_run = dag_maker.create_dagrun(run_id="test_run",
state=DagRunState.RUNNING)
mock_executor = MagicMock()
+ mock_executor.name =
ExecutorName(module_path="mock.executor.module.path")
+ mock_executor.team_name = None
scheduler_job = Job()
ti = dag_run.get_task_instance(task_id="task", session=session)
@@ -11950,7 +11958,7 @@ class TestSchedulerJob:
session.flush()
# Configure one executor to be team-specific
- mock_executors[1].team_name = "team_a"
+ assign_executor_to_team(mock_executors[1], "team_a")
with dag_maker(dag_id="dag_a", bundle_name="bundle_a",
session=session):
task = EmptyOperator(task_id="test_task") # No explicit executor
@@ -12016,7 +12024,7 @@ class TestSchedulerJob:
session.flush()
# Configure executor for the team
- mock_executors[1].team_name = "team_a"
+ assign_executor_to_team(mock_executors[1], "team_a")
with dag_maker(dag_id="dag_a", bundle_name="bundle_a",
session=session):
task = EmptyOperator(task_id="test_task",
executor="secondary_exec")
@@ -12050,7 +12058,7 @@ class TestSchedulerJob:
session.flush()
# Configure one executor for the team, but keep default as global
- mock_executors[1].team_name = "team_a"
+ assign_executor_to_team(mock_executors[1], "team_a")
with dag_maker(dag_id="dag_a", bundle_name="bundle_a",
session=session):
task = EmptyOperator(task_id="test_task", executor="default_exec")
# Global executor
@@ -12085,7 +12093,7 @@ class TestSchedulerJob:
session.flush()
# Configure executors for different teams
- mock_executors[1].team_name = "team_b" # Different team!
+ assign_executor_to_team(mock_executors[1], "team_b") # Different team!
with dag_maker(dag_id="dag_a", bundle_name="bundle_a",
session=session): # DAG belongs to team_a
task = EmptyOperator(
@@ -12144,7 +12152,7 @@ class TestSchedulerJob:
session.add(bundle)
session.flush()
- mock_executors[1].team_name = "team_a"
+ assign_executor_to_team(mock_executors[1], "team_a")
with dag_maker(dag_id="dag_a", bundle_name="bundle_a",
session=session):
task = EmptyOperator(task_id="test_task")
@@ -12193,6 +12201,61 @@ class TestSchedulerJob:
# Should match by classname (last component of module_path) and return
the global executor
assert result == mock_executors[0]
+ @conf_vars({("core", "multi_team"): "true"})
+ @pytest.mark.parametrize(
+ ("task_executor", "expected_index"),
+ [
+ pytest.param("CeleryExecutor", 1, id="classname-resolves-to-team"),
+ pytest.param(
+
"airflow.providers.celery.executors.celery_executor.CeleryExecutor",
+ 1,
+ id="module-path-resolves-to-team",
+ ),
+ pytest.param("team_celery_exec", 1, id="team-alias"),
+ pytest.param("global_celery_exec", 0, id="global-alias"),
+ ],
+ )
+ def
test_multi_team_try_to_load_executor_prefers_team_over_global_with_same_name(
+ self, dag_maker, mock_executors, session, task_executor, expected_index
+ ):
+ """A team task naming an executor that exists both globally and for
its team gets the team one.
+
+ Mirrors
``global_celery_exec:CeleryExecutor;team_a=team_celery_exec:CeleryExecutor``
from the
+ multi-team docs. Global executors come first in ``self.executors``, so
only the global alias
+ should reach the global executor.
+ """
+ clear_db_teams()
+ clear_db_dag_bundles()
+
+ team = Team(name="team_a")
+ session.add(team)
+ session.flush()
+
+ bundle = DagBundleModel(name="bundle_a")
+ bundle.teams.append(team)
+ session.add(bundle)
+ session.flush()
+
+ celery_module_path =
"airflow.providers.celery.executors.celery_executor.CeleryExecutor"
+ mock_executors[0].name = ExecutorName(alias="global_celery_exec",
module_path=celery_module_path)
+ mock_executors[1].name = ExecutorName(
+ alias="team_celery_exec", module_path=celery_module_path,
team_name="team_a"
+ )
+ mock_executors[1].team_name = "team_a"
+
+ with dag_maker(dag_id="dag_a", bundle_name="bundle_a",
session=session):
+ task = EmptyOperator(task_id="test_task", executor=task_executor)
+
+ dr = dag_maker.create_dagrun()
+ ti = dr.get_task_instance(task.task_id, session=session)
+
+ scheduler_job = Job()
+ self.job_runner = SchedulerJobRunner(job=scheduler_job)
+
+ result = self.job_runner._try_to_load_executor(ti, session=session)
+
+ assert result == mock_executors[expected_index]
+
@conf_vars({("core", "multi_team"): "true"})
def test_multi_team_scheduling_loop_batch_optimization(self, dag_maker,
mock_executors, session):
"""Test that the scheduling loop uses batch team resolution
optimization."""
@@ -12211,8 +12274,8 @@ class TestSchedulerJob:
session.add_all([bundle1, bundle2])
session.flush()
- mock_executors[0].team_name = "team_a"
- mock_executors[1].team_name = "team_b"
+ assign_executor_to_team(mock_executors[0], "team_a")
+ assign_executor_to_team(mock_executors[1], "team_b")
with dag_maker(dag_id="dag_a", bundle_name="bundle_a",
session=session):
EmptyOperator(task_id="task_a")
@@ -12261,8 +12324,8 @@ class TestSchedulerJob:
session.add_all([bundle1, bundle2])
session.flush()
- mock_executors[0].team_name = "team_a"
- mock_executors[1].team_name = "team_b"
+ assign_executor_to_team(mock_executors[0], "team_a")
+ assign_executor_to_team(mock_executors[1], "team_b")
with dag_maker(dag_id="dag_a", bundle_name="bundle_a",
session=session):
EmptyOperator(task_id="task_a")
