diff --git a/airflow/jobs/backfill_job_runner.py b/airflow/jobs/backfill_job_runner.py index c99cae2d21a25..4a78890d3b557 100644 --- a/airflow/jobs/backfill_job_runner.py +++ b/airflow/jobs/backfill_job_runner.py @@ -618,7 +618,7 @@ def _per_task_process(key, ti: TaskInstance, session): "Not scheduling since DAG max_active_tasks limit is reached." ) - if task.max_active_tis_per_dag: + if task.max_active_tis_per_dag is not None: num_running_task_instances_in_task = DAG.get_num_task_instances( dag_id=self.dag_id, task_ids=[task.task_id], @@ -631,6 +631,20 @@ def _per_task_process(key, ti: TaskInstance, session): "Not scheduling since Task concurrency limit is reached." ) + if task.max_active_tis_per_dagrun is not None: + num_running_task_instances_in_task_dagrun = DAG.get_num_task_instances( + dag_id=self.dag_id, + run_id=ti.run_id, + task_ids=[task.task_id], + states=self.STATES_COUNT_AS_RUNNING, + session=session, + ) + + if num_running_task_instances_in_task_dagrun >= task.max_active_tis_per_dagrun: + raise TaskConcurrencyLimitReached( + "Not scheduling since Task concurrency per DAG run limit is reached." + ) + _per_task_process(key, ti, session) session.commit() except (NoAvailablePoolSlot, DagConcurrencyLimitReached, TaskConcurrencyLimitReached) as e: diff --git a/airflow/jobs/scheduler_job_runner.py b/airflow/jobs/scheduler_job_runner.py index 706d980253a85..aae98373b621f 100644 --- a/airflow/jobs/scheduler_job_runner.py +++ b/airflow/jobs/scheduler_job_runner.py @@ -25,10 +25,11 @@ import sys import time import warnings -from collections import defaultdict +from collections import Counter +from dataclasses import dataclass from datetime import datetime, timedelta from pathlib import Path -from typing import TYPE_CHECKING, Collection, DefaultDict, Iterator +from typing import TYPE_CHECKING, Collection, Iterable, Iterator from sqlalchemy import and_, func, not_, or_, text from sqlalchemy.exc import OperationalError @@ -85,6 +86,29 @@ DM = DagModel +@dataclass +class ConcurrencyMap: + """ + Dataclass to represent concurrency maps + + It contains a map from (dag_id, task_id) to # of task instances, a map from (dag_id, task_id) + to # of task instances in the given state list and a map from (dag_id, run_id, task_id) + to # of task instances in the given state list in each DAG run. + """ + + dag_active_tasks_map: dict[str, int] + task_concurrency_map: dict[tuple[str, str], int] + task_dagrun_concurrency_map: dict[tuple[str, str, str], int] + + @classmethod + def from_concurrency_map(cls, mapping: dict[tuple[str, str, str], int]) -> ConcurrencyMap: + instance = cls(Counter(), Counter(), Counter(mapping)) + for (d, r, t), c in mapping.items(): + instance.dag_active_tasks_map[d] += c + instance.task_concurrency_map[(d, t)] += c + return instance + + def _is_parent_process() -> bool: """ Whether this is a parent process. @@ -231,28 +255,21 @@ def is_alive(self, grace_multiplier: float | None = None) -> bool: < scheduler_health_check_threshold ) - def __get_concurrency_maps( - self, states: list[TaskInstanceState], session: Session - ) -> tuple[DefaultDict[str, int], DefaultDict[tuple[str, str], int]]: + def __get_concurrency_maps(self, states: Iterable[TaskInstanceState], session: Session) -> ConcurrencyMap: """ Get the concurrency maps. :param states: List of states to query for - :return: A map from (dag_id, task_id) to # of task instances and - a map from (dag_id, task_id) to # of task instances in the given state list + :return: Concurrency map """ - ti_concurrency_query: list[tuple[str, str, int]] = ( - session.query(TI.task_id, TI.dag_id, func.count("*")) + ti_concurrency_query: list[tuple[str, str, str, int]] = ( + session.query(TI.task_id, TI.run_id, TI.dag_id, func.count("*")) .filter(TI.state.in_(states)) - .group_by(TI.task_id, TI.dag_id) - ).all() - dag_map: DefaultDict[str, int] = defaultdict(int) - task_map: DefaultDict[tuple[str, str], int] = defaultdict(int) - for result in ti_concurrency_query: - task_id, dag_id, count = result - dag_map[dag_id] += count - task_map[(dag_id, task_id)] = count - return dag_map, task_map + .group_by(TI.task_id, TI.run_id, TI.dag_id) + ) + return ConcurrencyMap.from_concurrency_map( + {(dag_id, run_id, task_id): count for task_id, run_id, dag_id, count in ti_concurrency_query} + ) def _executable_task_instances_to_queued(self, max_tis: int, session: Session) -> list[TI]: """ @@ -263,6 +280,8 @@ def _executable_task_instances_to_queued(self, max_tis: int, session: Session) - - DAG max_active_tasks - executor state - priority + - max active tis per DAG + - max active tis per DAG run :param max_tis: Maximum number of TIs to queue in this loop. :return: list[airflow.models.TaskInstance] @@ -304,11 +323,7 @@ def _executable_task_instances_to_queued(self, max_tis: int, session: Session) - starved_pools = {pool_name for pool_name, stats in pools.items() if stats["open"] <= 0} # dag_id to # of running tasks and (dag_id, task_id) to # of running tasks. - dag_active_tasks_map: DefaultDict[str, int] - task_concurrency_map: DefaultDict[tuple[str, str], int] - dag_active_tasks_map, task_concurrency_map = self.__get_concurrency_maps( - states=list(EXECUTION_STATES), session=session - ) + concurrency_map = self.__get_concurrency_maps(states=EXECUTION_STATES, session=session) # Number of tasks that cannot be scheduled because of no open slot in pool num_starving_tasks_total = 0 @@ -316,14 +331,16 @@ def _executable_task_instances_to_queued(self, max_tis: int, session: Session) - # dag and task ids that can't be queued because of concurrency limits starved_dags: set[str] = set() starved_tasks: set[tuple[str, str]] = set() + starved_tasks_task_dagrun_concurrency: set[tuple[str, str, str]] = set() - pool_num_starving_tasks: DefaultDict[str, int] = defaultdict(int) + pool_num_starving_tasks: dict[str, int] = Counter() for loop_count in itertools.count(start=1): num_starved_pools = len(starved_pools) num_starved_dags = len(starved_dags) num_starved_tasks = len(starved_tasks) + num_starved_tasks_task_dagrun_concurrency = len(starved_tasks_task_dagrun_concurrency) # Get task instances associated with scheduled # DagRuns which are not backfilled, in the given states, @@ -347,7 +364,14 @@ def _executable_task_instances_to_queued(self, max_tis: int, session: Session) - query = query.filter(not_(TI.dag_id.in_(starved_dags))) if starved_tasks: - task_filter = tuple_in_condition((TaskInstance.dag_id, TaskInstance.task_id), starved_tasks) + task_filter = tuple_in_condition((TI.dag_id, TI.task_id), starved_tasks) + query = query.filter(not_(task_filter)) + + if starved_tasks_task_dagrun_concurrency: + task_filter = tuple_in_condition( + (TI.dag_id, TI.run_id, TI.task_id), + starved_tasks_task_dagrun_concurrency, + ) query = query.filter(not_(task_filter)) query = query.limit(max_tis) @@ -439,7 +463,7 @@ def _executable_task_instances_to_queued(self, max_tis: int, session: Session) - # reached. dag_id = task_instance.dag_id - current_active_tasks_per_dag = dag_active_tasks_map[dag_id] + current_active_tasks_per_dag = concurrency_map.dag_active_tasks_map[dag_id] max_active_tasks_per_dag_limit = task_instance.dag_model.max_active_tasks self.log.info( "DAG %s has %s/%s running and queued tasks", @@ -481,7 +505,7 @@ def _executable_task_instances_to_queued(self, max_tis: int, session: Session) - ).max_active_tis_per_dag if task_concurrency_limit is not None: - current_task_concurrency = task_concurrency_map[ + current_task_concurrency = concurrency_map.task_concurrency_map[ (task_instance.dag_id, task_instance.task_id) ] @@ -494,10 +518,35 @@ def _executable_task_instances_to_queued(self, max_tis: int, session: Session) - starved_tasks.add((task_instance.dag_id, task_instance.task_id)) continue + task_dagrun_concurrency_limit: int | None = None + if serialized_dag.has_task(task_instance.task_id): + task_dagrun_concurrency_limit = serialized_dag.get_task( + task_instance.task_id + ).max_active_tis_per_dagrun + + if task_dagrun_concurrency_limit is not None: + current_task_dagrun_concurrency = concurrency_map.task_dagrun_concurrency_map[ + (task_instance.dag_id, task_instance.run_id, task_instance.task_id) + ] + + if current_task_dagrun_concurrency >= task_dagrun_concurrency_limit: + self.log.info( + "Not executing %s since the task concurrency per DAG run for" + " this task has been reached.", + task_instance, + ) + starved_tasks_task_dagrun_concurrency.add( + (task_instance.dag_id, task_instance.run_id, task_instance.task_id) + ) + continue + executable_tis.append(task_instance) open_slots -= task_instance.pool_slots - dag_active_tasks_map[dag_id] += 1 - task_concurrency_map[(task_instance.dag_id, task_instance.task_id)] += 1 + concurrency_map.dag_active_tasks_map[dag_id] += 1 + concurrency_map.task_concurrency_map[(task_instance.dag_id, task_instance.task_id)] += 1 + concurrency_map.task_dagrun_concurrency_map[ + (task_instance.dag_id, task_instance.run_id, task_instance.task_id) + ] += 1 pool_stats["open"] = open_slots @@ -507,6 +556,7 @@ def _executable_task_instances_to_queued(self, max_tis: int, session: Session) - len(starved_pools) > num_starved_pools or len(starved_dags) > num_starved_dags or len(starved_tasks) > num_starved_tasks + or len(starved_tasks_task_dagrun_concurrency) > num_starved_tasks_task_dagrun_concurrency ) if is_done or not found_new_filters: @@ -816,13 +866,13 @@ def _update_dag_run_state_for_paused_dags(self, session: Session = NEW_SESSION) paused_runs = ( session.query(DagRun) .join(DagRun.dag_model) - .join(TaskInstance) + .join(TI) .filter( DagModel.is_paused == expression.true(), DagRun.state == DagRunState.RUNNING, DagRun.run_type != DagRunType.BACKFILL_JOB, ) - .having(DagRun.last_scheduling_decision <= func.max(TaskInstance.updated_at)) + .having(DagRun.last_scheduling_decision <= func.max(TI.updated_at)) .group_by(DagRun) ) for dag_run in paused_runs: @@ -1079,8 +1129,7 @@ def _create_dag_runs(self, dag_models: Collection[DagModel], session: Session) - .all() ) - active_runs_of_dags = defaultdict( - int, + active_runs_of_dags = Counter( DagRun.active_runs_of_dags(dag_ids=(dm.dag_id for dm in dag_models), session=session), ) @@ -1237,8 +1286,7 @@ def _start_queued_dagruns(self, session: Session) -> None: """Find DagRuns in queued state and decide moving them to running state.""" dag_runs = self._get_next_dagruns_to_examine(DagRunState.QUEUED, session) - active_runs_of_dags = defaultdict( - int, + active_runs_of_dags = Counter( DagRun.active_runs_of_dags((dr.dag_id for dr in dag_runs), only_running=True, session=session), ) @@ -1533,10 +1581,10 @@ def check_trigger_timeouts(self, session: Session = NEW_SESSION) -> None: or execution timeout has passed, so they can be marked as failed. """ num_timed_out_tasks = ( - session.query(TaskInstance) + session.query(TI) .filter( - TaskInstance.state == TaskInstanceState.DEFERRED, - TaskInstance.trigger_timeout < timezone.utcnow(), + TI.state == TaskInstanceState.DEFERRED, + TI.trigger_timeout < timezone.utcnow(), ) .update( # We have to schedule these to fail themselves so it doesn't @@ -1599,7 +1647,7 @@ def _find_zombies(self) -> None: ) @staticmethod - def _generate_zombie_message_details(ti: TaskInstance): + def _generate_zombie_message_details(ti: TI): zombie_message_details = { "DAG Id": ti.dag_id, "Task Id": ti.task_id, diff --git a/airflow/models/baseoperator.py b/airflow/models/baseoperator.py index 1760bd1f420ab..d81d45dc7e92a 100644 --- a/airflow/models/baseoperator.py +++ b/airflow/models/baseoperator.py @@ -214,6 +214,7 @@ def partial( weight_rule: str = DEFAULT_WEIGHT_RULE, sla: timedelta | None = None, max_active_tis_per_dag: int | None = None, + max_active_tis_per_dagrun: int | None = None, on_execute_callback: None | TaskStateChangeCallback | list[TaskStateChangeCallback] = None, on_failure_callback: None | TaskStateChangeCallback | list[TaskStateChangeCallback] = None, on_success_callback: None | TaskStateChangeCallback | list[TaskStateChangeCallback] = None, @@ -274,6 +275,7 @@ def partial( partial_kwargs.setdefault("weight_rule", weight_rule) partial_kwargs.setdefault("sla", sla) partial_kwargs.setdefault("max_active_tis_per_dag", max_active_tis_per_dag) + partial_kwargs.setdefault("max_active_tis_per_dagrun", max_active_tis_per_dagrun) partial_kwargs.setdefault("on_execute_callback", on_execute_callback) partial_kwargs.setdefault("on_failure_callback", on_failure_callback) partial_kwargs.setdefault("on_retry_callback", on_retry_callback) @@ -578,6 +580,8 @@ class derived from this one results in the creation of a task object, :param run_as_user: unix username to impersonate while running the task :param max_active_tis_per_dag: When set, a task will be able to limit the concurrent runs across execution_dates. + :param max_active_tis_per_dagrun: When set, a task will be able to limit the concurrent + task instances per DAG run. :param executor_config: Additional task-level configuration parameters that are interpreted by a specific executor. Parameters are namespaced by the name of executor. @@ -729,6 +733,7 @@ def __init__( run_as_user: str | None = None, task_concurrency: int | None = None, max_active_tis_per_dag: int | None = None, + max_active_tis_per_dagrun: int | None = None, executor_config: dict | None = None, do_xcom_push: bool = True, inlets: Any | None = None, @@ -872,6 +877,7 @@ def __init__( ) max_active_tis_per_dag = task_concurrency self.max_active_tis_per_dag: int | None = max_active_tis_per_dag + self.max_active_tis_per_dagrun: int | None = max_active_tis_per_dagrun self.do_xcom_push = do_xcom_push self.doc_md = doc_md diff --git a/airflow/models/dag.py b/airflow/models/dag.py index 8b259cd3e825a..9888e5cbd8d04 100644 --- a/airflow/models/dag.py +++ b/airflow/models/dag.py @@ -2789,7 +2789,10 @@ def bulk_write_to_db( orm_dag.description = dag.description orm_dag.max_active_tasks = dag.max_active_tasks orm_dag.max_active_runs = dag.max_active_runs - orm_dag.has_task_concurrency_limits = any(t.max_active_tis_per_dag is not None for t in dag.tasks) + orm_dag.has_task_concurrency_limits = any( + t.max_active_tis_per_dag is not None or t.max_active_tis_per_dagrun is not None + for t in dag.tasks + ) orm_dag.schedule_interval = dag.schedule_interval orm_dag.timetable_description = dag.timetable.description orm_dag.processor_subdir = processor_subdir @@ -2990,12 +2993,13 @@ def deactivate_stale_dags(expiration_date, session=NEW_SESSION): @staticmethod @provide_session - def get_num_task_instances(dag_id, task_ids=None, states=None, session=NEW_SESSION) -> int: + def get_num_task_instances(dag_id, run_id=None, task_ids=None, states=None, session=NEW_SESSION) -> int: """ Returns the number of task instances in the given DAG. :param session: ORM session :param dag_id: ID of the DAG to get the task concurrency of + :param run_id: ID of the DAG run to get the task concurrency of :param task_ids: A list of valid task IDs for the given DAG :param states: A list of states to filter by if supplied :return: The number of running tasks @@ -3003,6 +3007,10 @@ def get_num_task_instances(dag_id, task_ids=None, states=None, session=NEW_SESSI qry = session.query(func.count(TaskInstance.task_id)).filter( TaskInstance.dag_id == dag_id, ) + if run_id: + qry = qry.filter( + TaskInstance.run_id == run_id, + ) if task_ids: qry = qry.filter( TaskInstance.task_id.in_(task_ids), diff --git a/airflow/models/dagrun.py b/airflow/models/dagrun.py index 101bda0a471a6..edb0ec78ac44c 100644 --- a/airflow/models/dagrun.py +++ b/airflow/models/dagrun.py @@ -553,6 +553,7 @@ def should_schedule(self) -> bool: bool(self.tis) and all(not t.task.depends_on_past for t in self.tis) and all(t.task.max_active_tis_per_dag is None for t in self.tis) + and all(t.task.max_active_tis_per_dagrun is None for t in self.tis) and all(t.state != TaskInstanceState.DEFERRED for t in self.tis) ) diff --git a/airflow/models/mappedoperator.py b/airflow/models/mappedoperator.py index a10fd10cdc04b..345329ef6d4a1 100644 --- a/airflow/models/mappedoperator.py +++ b/airflow/models/mappedoperator.py @@ -451,6 +451,10 @@ def sla(self) -> datetime.timedelta | None: def max_active_tis_per_dag(self) -> int | None: return self.partial_kwargs.get("max_active_tis_per_dag") + @property + def max_active_tis_per_dagrun(self) -> int | None: + return self.partial_kwargs.get("max_active_tis_per_dagrun") + @property def resources(self) -> Resources | None: return self.partial_kwargs.get("resources") diff --git a/airflow/models/taskinstance.py b/airflow/models/taskinstance.py index c2ed7b786864c..b02076c0767df 100644 --- a/airflow/models/taskinstance.py +++ b/airflow/models/taskinstance.py @@ -2479,18 +2479,17 @@ def xcom_pull( return LazyXComAccess.build_from_xcom_query(query) @provide_session - def get_num_running_task_instances(self, session: Session) -> int: + def get_num_running_task_instances(self, session: Session, same_dagrun=False) -> int: """Return Number of running TIs from the DB""" # .count() is inefficient - return ( - session.query(func.count()) - .filter( - TaskInstance.dag_id == self.dag_id, - TaskInstance.task_id == self.task_id, - TaskInstance.state == State.RUNNING, - ) - .scalar() + num_running_task_instances_query = session.query(func.count()).filter( + TaskInstance.dag_id == self.dag_id, + TaskInstance.task_id == self.task_id, + TaskInstance.state == State.RUNNING, ) + if same_dagrun: + num_running_task_instances_query.filter(TaskInstance.run_id == self.run_id) + return num_running_task_instances_query.scalar() def init_run_context(self, raw: bool = False) -> None: """Sets the log context.""" diff --git a/airflow/ti_deps/deps/task_concurrency_dep.py b/airflow/ti_deps/deps/task_concurrency_dep.py index 5b5f4f515acf0..1f1416214c7a4 100644 --- a/airflow/ti_deps/deps/task_concurrency_dep.py +++ b/airflow/ti_deps/deps/task_concurrency_dep.py @@ -30,13 +30,22 @@ class TaskConcurrencyDep(BaseTIDep): @provide_session def _get_dep_statuses(self, ti, session, dep_context): - if ti.task.max_active_tis_per_dag is None: + if ti.task.max_active_tis_per_dag is None and ti.task.max_active_tis_per_dagrun is None: yield self._passing_status(reason="Task concurrency is not set.") return - if ti.get_num_running_task_instances(session) >= ti.task.max_active_tis_per_dag: + if ( + ti.task.max_active_tis_per_dag is not None + and ti.get_num_running_task_instances(session) >= ti.task.max_active_tis_per_dag + ): yield self._failing_status(reason="The max task concurrency has been reached.") return - else: - yield self._passing_status(reason="The max task concurrency has not been reached.") + if ( + ti.task.max_active_tis_per_dagrun is not None + and ti.get_num_running_task_instances(session, same_dagrun=True) + >= ti.task.max_active_tis_per_dagrun + ): + yield self._failing_status(reason="The max task concurrency per run has been reached.") return + yield self._passing_status(reason="The max task concurrency has not been reached.") + return diff --git a/tests/conftest.py b/tests/conftest.py index cfae8468959a3..f169e2b62e793 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -726,6 +726,7 @@ def create_dag( dag_id="dag", task_id="op1", max_active_tis_per_dag=16, + max_active_tis_per_dagrun=None, pool="default_pool", executor_config={}, trigger_rule="all_done", @@ -741,6 +742,7 @@ def create_dag( op = EmptyOperator( task_id=task_id, max_active_tis_per_dag=max_active_tis_per_dag, + max_active_tis_per_dagrun=max_active_tis_per_dagrun, executor_config=executor_config, on_success_callback=on_success_callback, on_execute_callback=on_execute_callback, diff --git a/tests/jobs/test_backfill_job.py b/tests/jobs/test_backfill_job.py index 7001f6f1b917f..164d5af1df20b 100644 --- a/tests/jobs/test_backfill_job.py +++ b/tests/jobs/test_backfill_job.py @@ -21,6 +21,7 @@ import json import logging import threading +from collections import defaultdict from unittest import mock from unittest.mock import patch @@ -376,6 +377,88 @@ def test_backfill_respect_max_active_tis_per_dag_limit(self, mock_log, dag_maker assert 0 == times_dag_concurrency_limit_reached_in_debug assert times_task_concurrency_limit_reached_in_debug > 0 + @pytest.mark.parametrize("with_max_active_tis_per_dag", [False, True]) + @patch("airflow.jobs.backfill_job_runner.BackfillJobRunner.log") + def test_backfill_respect_max_active_tis_per_dagrun_limit( + self, mock_log, dag_maker, with_max_active_tis_per_dag + ): + max_active_tis_per_dag = 3 + max_active_tis_per_dagrun = 2 + kwargs = {"max_active_tis_per_dagrun": max_active_tis_per_dagrun} + if with_max_active_tis_per_dag: + kwargs["max_active_tis_per_dag"] = max_active_tis_per_dag + + with dag_maker(dag_id="test_backfill_respect_max_active_tis_per_dag_limit", schedule="@daily") as dag: + EmptyOperator.partial(task_id="task1", **kwargs).expand_kwargs([{"x": i} for i in range(10)]) + + dag_maker.create_dagrun(state=None) + + executor = MockExecutor() + + job = Job(executor=executor) + job_runner = BackfillJobRunner( + job=job, + dag=dag, + start_date=DEFAULT_DATE, + end_date=DEFAULT_DATE + datetime.timedelta(days=7), + ) + + run_job(job=job, execute_callable=job_runner._execute) + + assert len(executor.history) > 0 + + task_concurrency_limit_reached_at_least_once = False + + def get_running_tis_per_dagrun(running_tis): + running_tis_per_dagrun_dict = defaultdict(int) + for running_ti in running_tis: + running_tis_per_dagrun_dict[running_ti[3].dag_run.id] += 1 + return running_tis_per_dagrun_dict + + num_running_task_instances = 0 + for running_task_instances in executor.history: + if with_max_active_tis_per_dag: + assert len(running_task_instances) <= max_active_tis_per_dag + running_tis_per_dagrun_dict = get_running_tis_per_dagrun(running_task_instances) + assert all( + [ + num_running_tis <= max_active_tis_per_dagrun + for num_running_tis in running_tis_per_dagrun_dict.values() + ] + ) + num_running_task_instances += len(running_task_instances) + task_concurrency_limit_reached_at_least_once = ( + task_concurrency_limit_reached_at_least_once + or any( + [ + num_running_tis == max_active_tis_per_dagrun + for num_running_tis in running_tis_per_dagrun_dict.values() + ] + ) + ) + + assert 80 == num_running_task_instances # (7 backfill run + 1 manual run ) * 10 mapped task per run + assert task_concurrency_limit_reached_at_least_once + + times_dag_concurrency_limit_reached_in_debug = self._times_called_with( + mock_log.debug, + DagConcurrencyLimitReached, + ) + + times_pool_limit_reached_in_debug = self._times_called_with( + mock_log.debug, + NoAvailablePoolSlot, + ) + + times_task_concurrency_limit_reached_in_debug = self._times_called_with( + mock_log.debug, + TaskConcurrencyLimitReached, + ) + + assert 0 == times_pool_limit_reached_in_debug + assert 0 == times_dag_concurrency_limit_reached_in_debug + assert times_task_concurrency_limit_reached_in_debug > 0 + @patch("airflow.jobs.backfill_job_runner.BackfillJobRunner.log") def test_backfill_respect_dag_concurrency_limit(self, mock_log, dag_maker): dag = self._get_dummy_dag(dag_maker, dag_id="test_backfill_respect_concurrency_limit") diff --git a/tests/jobs/test_scheduler_job.py b/tests/jobs/test_scheduler_job.py index 4870a2aa132e9..fafc5ee2b4a18 100644 --- a/tests/jobs/test_scheduler_job.py +++ b/tests/jobs/test_scheduler_job.py @@ -1264,6 +1264,68 @@ def test_find_executable_task_instances_not_enough_task_concurrency_for_first(se session.rollback() + def test_find_executable_task_instances_task_concurrency_per_dagrun_for_first(self, dag_maker): + scheduler_job = Job() + self.job_runner = SchedulerJobRunner(job=scheduler_job, subdir=os.devnull) + session = settings.Session() + + dag_id = "SchedulerJobTest.test_find_executable_task_instances_task_concurrency_per_dagrun_for_first" + + with dag_maker(dag_id=dag_id): + op1a = EmptyOperator(task_id="dummy1-a", priority_weight=2, max_active_tis_per_dagrun=1) + op1b = EmptyOperator(task_id="dummy1-b", priority_weight=1) + dr1 = dag_maker.create_dagrun(run_type=DagRunType.SCHEDULED) + dr2 = dag_maker.create_dagrun_after(dr1, run_type=DagRunType.SCHEDULED) + + ti1a = dr1.get_task_instance(op1a.task_id, session) + ti1b = dr1.get_task_instance(op1b.task_id, session) + ti2a = dr2.get_task_instance(op1a.task_id, session) + ti1a.state = State.RUNNING + ti1b.state = State.SCHEDULED + ti2a.state = State.SCHEDULED + session.flush() + + # Schedule ti with higher priority, + # because it's running in a different DAG run with 0 active tis + res = self.job_runner._executable_task_instances_to_queued(max_tis=1, session=session) + assert 1 == len(res) + assert res[0].key == ti2a.key + + session.rollback() + + def test_find_executable_task_instances_not_enough_task_concurrency_per_dagrun_for_first(self, dag_maker): + scheduler_job = Job() + self.job_runner = SchedulerJobRunner(job=scheduler_job, subdir=os.devnull) + session = settings.Session() + + dag_id = ( + "SchedulerJobTest" + ".test_find_executable_task_instances_not_enough_task_concurrency_per_dagrun_for_first" + ) + + with dag_maker(dag_id=dag_id): + op1a = EmptyOperator.partial( + task_id="dummy1-a", priority_weight=2, max_active_tis_per_dagrun=1 + ).expand_kwargs([{"inputs": 1}, {"inputs": 2}]) + op1b = EmptyOperator(task_id="dummy1-b", priority_weight=1) + dr = dag_maker.create_dagrun(run_type=DagRunType.SCHEDULED) + + ti1a0 = dr.get_task_instance(op1a.task_id, session, map_index=0) + ti1a1 = dr.get_task_instance(op1a.task_id, session, map_index=1) + ti1b = dr.get_task_instance(op1b.task_id, session) + ti1a0.state = State.RUNNING + ti1a1.state = State.SCHEDULED + ti1b.state = State.SCHEDULED + session.flush() + + # Schedule ti with lower priority, + # because the one with higher priority is limited by a concurrency limit + res = self.job_runner._executable_task_instances_to_queued(max_tis=1, session=session) + assert 1 == len(res) + assert res[0].key == ti1b.key + + session.rollback() + def test_find_executable_task_instances_negative_open_pool_slots(self, dag_maker): """ Pools with negative open slots should not block other pools. @@ -1419,7 +1481,9 @@ def test_critical_section_enqueue_task_instances(self, dag_maker): session.flush() assert State.RUNNING == dr1.state - assert 2 == DAG.get_num_task_instances(dag_id, dag.task_ids, states=[State.RUNNING], session=session) + assert 2 == DAG.get_num_task_instances( + dag_id, task_ids=dag.task_ids, states=[State.RUNNING], session=session + ) # create second dag run dr2 = dag_maker.create_dagrun_after(dr1, run_type=DagRunType.SCHEDULED) @@ -1440,7 +1504,7 @@ def test_critical_section_enqueue_task_instances(self, dag_maker): ti3.refresh_from_db() ti4.refresh_from_db() assert 3 == DAG.get_num_task_instances( - dag_id, dag.task_ids, states=[State.RUNNING, State.QUEUED], session=session + dag_id, task_ids=dag.task_ids, states=[State.RUNNING, State.QUEUED], session=session ) assert State.RUNNING == ti1.state assert State.RUNNING == ti2.state @@ -4127,6 +4191,54 @@ def test_dag_file_processor_process_task_instances_with_max_active_tis_per_dag( session.refresh(ti) assert ti.state == State.SCHEDULED + @pytest.mark.parametrize( + "state,start_date,end_date", + [ + [State.NONE, None, None], + [ + State.UP_FOR_RETRY, + timezone.utcnow() - datetime.timedelta(minutes=30), + timezone.utcnow() - datetime.timedelta(minutes=15), + ], + [ + State.UP_FOR_RESCHEDULE, + timezone.utcnow() - datetime.timedelta(minutes=30), + timezone.utcnow() - datetime.timedelta(minutes=15), + ], + ], + ) + def test_dag_file_processor_process_task_instances_with_max_active_tis_per_dagrun( + self, state, start_date, end_date, dag_maker + ): + """ + Test if _process_task_instances puts the right task instances into the + mock_list. + """ + with dag_maker(dag_id="test_scheduler_process_execute_task_with_max_active_tis_per_dagrun"): + BashOperator(task_id="dummy", max_active_tis_per_dagrun=2, bash_command="echo Hi") + + scheduler_job = Job() + self.job_runner = SchedulerJobRunner(job=scheduler_job, subdir=os.devnull) + + self.job_runner.processor_agent = mock.MagicMock() + + dr = dag_maker.create_dagrun( + run_type=DagRunType.SCHEDULED, + ) + assert dr is not None + + with create_session() as session: + ti = dr.get_task_instances(session=session)[0] + ti.state = state + ti.start_date = start_date + ti.end_date = end_date + + self.job_runner._schedule_dag_run(dr, session) + assert session.query(TaskInstance).filter_by(state=State.SCHEDULED).count() == 1 + + session.refresh(ti) + assert ti.state == State.SCHEDULED + @pytest.mark.parametrize( "state, start_date, end_date", [ diff --git a/tests/models/test_dag.py b/tests/models/test_dag.py index 571dde87a62b9..b4ff9a0bed848 100644 --- a/tests/models/test_dag.py +++ b/tests/models/test_dag.py @@ -461,18 +461,22 @@ def test_get_num_task_instances(self): session.merge(ti4) session.commit() - assert 0 == DAG.get_num_task_instances(test_dag_id, ["fakename"], session=session) - assert 4 == DAG.get_num_task_instances(test_dag_id, [test_task_id], session=session) - assert 4 == DAG.get_num_task_instances(test_dag_id, ["fakename", test_task_id], session=session) - assert 1 == DAG.get_num_task_instances(test_dag_id, [test_task_id], states=[None], session=session) + assert 0 == DAG.get_num_task_instances(test_dag_id, task_ids=["fakename"], session=session) + assert 4 == DAG.get_num_task_instances(test_dag_id, task_ids=[test_task_id], session=session) + assert 4 == DAG.get_num_task_instances( + test_dag_id, task_ids=["fakename", test_task_id], session=session + ) + assert 1 == DAG.get_num_task_instances( + test_dag_id, task_ids=[test_task_id], states=[None], session=session + ) assert 2 == DAG.get_num_task_instances( - test_dag_id, [test_task_id], states=[State.RUNNING], session=session + test_dag_id, task_ids=[test_task_id], states=[State.RUNNING], session=session ) assert 3 == DAG.get_num_task_instances( - test_dag_id, [test_task_id], states=[None, State.RUNNING], session=session + test_dag_id, task_ids=[test_task_id], states=[None, State.RUNNING], session=session ) assert 4 == DAG.get_num_task_instances( - test_dag_id, [test_task_id], states=[None, State.QUEUED, State.RUNNING], session=session + test_dag_id, task_ids=[test_task_id], states=[None, State.QUEUED, State.RUNNING], session=session ) session.close() diff --git a/tests/models/test_taskinstance.py b/tests/models/test_taskinstance.py index 2d86cdfde5ffe..fb91982c802dd 100644 --- a/tests/models/test_taskinstance.py +++ b/tests/models/test_taskinstance.py @@ -288,6 +288,19 @@ def test_requeue_over_max_active_tis_per_dag(self, create_task_instance): ti.run() assert ti.state == State.NONE + def test_requeue_over_max_active_tis_per_dagrun(self, create_task_instance): + ti = create_task_instance( + dag_id="test_requeue_over_max_active_tis_per_dagrun", + task_id="test_requeue_over_max_active_tis_per_dagrun_op", + max_active_tis_per_dagrun=0, + max_active_runs=1, + max_active_tasks=2, + dagrun_state=State.QUEUED, + ) + + ti.run() + assert ti.state == State.NONE + def test_requeue_over_pool_concurrency(self, create_task_instance, test_pool): ti = create_task_instance( dag_id="test_requeue_over_pool_concurrency", diff --git a/tests/serialization/test_dag_serialization.py b/tests/serialization/test_dag_serialization.py index 2d9a5fa72ee60..b025a61ae574b 100644 --- a/tests/serialization/test_dag_serialization.py +++ b/tests/serialization/test_dag_serialization.py @@ -1189,6 +1189,7 @@ def test_no_new_fields_added_to_base_operator(self): "ignore_first_depends_on_past": True, "inlets": [], "max_active_tis_per_dag": None, + "max_active_tis_per_dagrun": None, "max_retry_delay": None, "on_execute_callback": None, "on_failure_callback": None, diff --git a/tests/ti_deps/deps/test_task_concurrency.py b/tests/ti_deps/deps/test_task_concurrency.py index 5d208a650c906..aa6c8e116c1ad 100644 --- a/tests/ti_deps/deps/test_task_concurrency.py +++ b/tests/ti_deps/deps/test_task_concurrency.py @@ -20,6 +20,8 @@ from datetime import datetime from unittest.mock import Mock +import pytest + from airflow.models import DAG from airflow.models.baseoperator import BaseOperator from airflow.ti_deps.dep_context import DepContext @@ -30,24 +32,25 @@ class TestTaskConcurrencyDep: def _get_task(self, **kwargs): return BaseOperator(task_id="test_task", dag=DAG("test_dag"), **kwargs) - def test_not_task_concurrency(self): - task = self._get_task(start_date=datetime(2016, 1, 1)) - dep_context = DepContext() - ti = Mock(task=task, execution_date=datetime(2016, 1, 1)) - assert TaskConcurrencyDep().is_met(ti=ti, dep_context=dep_context) - - def test_not_reached_concurrency(self): - task = self._get_task(start_date=datetime(2016, 1, 1), max_active_tis_per_dag=1) - dep_context = DepContext() - ti = Mock(task=task, execution_date=datetime(2016, 1, 1)) - ti.get_num_running_task_instances = lambda x: 0 - assert TaskConcurrencyDep().is_met(ti=ti, dep_context=dep_context) - - def test_reached_concurrency(self): - task = self._get_task(start_date=datetime(2016, 1, 1), max_active_tis_per_dag=2) + @pytest.mark.parametrize( + "kwargs, num_running_tis, is_task_concurrency_dep_met", + [ + ({}, None, True), + ({"max_active_tis_per_dag": 1}, 0, True), + ({"max_active_tis_per_dag": 2}, 1, True), + ({"max_active_tis_per_dag": 2}, 2, False), + ({"max_active_tis_per_dagrun": 2}, 1, True), + ({"max_active_tis_per_dagrun": 2}, 2, False), + ({"max_active_tis_per_dag": 2, "max_active_tis_per_dagrun": 2}, 1, True), + ({"max_active_tis_per_dag": 1, "max_active_tis_per_dagrun": 2}, 1, False), + ({"max_active_tis_per_dag": 2, "max_active_tis_per_dagrun": 1}, 1, False), + ({"max_active_tis_per_dag": 1, "max_active_tis_per_dagrun": 1}, 1, False), + ], + ) + def test_concurrency(self, kwargs, num_running_tis, is_task_concurrency_dep_met): + task = self._get_task(start_date=datetime(2016, 1, 1), **kwargs) dep_context = DepContext() ti = Mock(task=task, execution_date=datetime(2016, 1, 1)) - ti.get_num_running_task_instances = lambda x: 1 - assert TaskConcurrencyDep().is_met(ti=ti, dep_context=dep_context) - ti.get_num_running_task_instances = lambda x: 2 - assert not TaskConcurrencyDep().is_met(ti=ti, dep_context=dep_context) + if num_running_tis is not None: + ti.get_num_running_task_instances.return_value = num_running_tis + assert TaskConcurrencyDep().is_met(ti=ti, dep_context=dep_context) == is_task_concurrency_dep_met