diff --git a/airflow/config_templates/config.yml b/airflow/config_templates/config.yml index 8ef738f4f4ad6..3b6aa30870421 100644 --- a/airflow/config_templates/config.yml +++ b/airflow/config_templates/config.yml @@ -1961,6 +1961,15 @@ type: string example: ~ default: "15" + - name: min_active_runs_check_interval + description: | + The number of seconds after which the dag that has reached ``max_active_run`` + will be checked again for the need to create new DagRuns. Keeping this + number low will increase CPU usage. + version_added: ~ + type: integer + example: ~ + default: "10" - name: triggerer description: ~ options: diff --git a/airflow/config_templates/default_airflow.cfg b/airflow/config_templates/default_airflow.cfg index 520ab4442850d..f6f68283d3a18 100644 --- a/airflow/config_templates/default_airflow.cfg +++ b/airflow/config_templates/default_airflow.cfg @@ -984,6 +984,11 @@ dependency_detector = airflow.serialization.serialized_objects.DependencyDetecto # How often to check for expired trigger requests that have not run yet. trigger_timeout_check_interval = 15 +# The number of seconds after which the dag that has reached ``max_active_run`` +# will be checked again for the need to create new DagRuns. Keeping this +# number low will increase CPU usage. +min_active_runs_check_interval = 10 + [triggerer] # How many triggers a single Triggerer will run at once, by default. default_capacity = 1000 diff --git a/airflow/jobs/scheduler_job.py b/airflow/jobs/scheduler_job.py index cbda16eb49e02..f7338366e2b80 100644 --- a/airflow/jobs/scheduler_job.py +++ b/airflow/jobs/scheduler_job.py @@ -25,7 +25,7 @@ import time import warnings from collections import defaultdict -from datetime import timedelta +from datetime import datetime, timedelta from typing import Collection, DefaultDict, Dict, Iterator, List, Optional, Tuple from sqlalchemy import and_, func, not_, or_, text, tuple_ @@ -139,6 +139,8 @@ def __init__( self.dagbag = DagBag(dag_folder=self.subdir, read_dags_from_db=True, load_op_links=False) + self.last_check_time: Dict[str, datetime] = {} + if conf.getboolean('smart_sensor', 'use_smart_sensor'): compatible_sensors = set( map(lambda l: l.strip(), conf.get('smart_sensor', 'sensors_enabled').split(',')) @@ -884,11 +886,28 @@ def _get_next_dagruns_to_examine(self, state: DagRunState, session: Session): """Get Next DagRuns to Examine with retries""" return DagRun.next_dagruns_to_examine(state, session) + def _get_recently_checked_dags(self, current_time=None) -> List: + """Get dags for which `max_active_runs` has been reached recently""" + if current_time is None: + current_time = timezone.utcnow() + check_interval = timedelta( + seconds=conf.getint('scheduler', 'min_active_runs_check_interval', fallback=10) + ) + + skip_dags = [] + for dag_id in self.last_check_time: + if self.last_check_time[dag_id] + check_interval > current_time: + skip_dags.append(dag_id) + return skip_dags + @retry_db_transaction def _create_dagruns_for_dags(self, guard, session): """Find Dag Models needing DagRuns and Create Dag Runs with retries in case of OperationalError""" - query = DagModel.dags_needing_dagruns(session) - self._create_dag_runs(query.all(), session) + # Reduce the number of dags due to high CPU usage + skip_dags = self._get_recently_checked_dags() + self.log.debug("Skipping dags after max_active_runs has been reached: %s", skip_dags) + dag_models = DagModel.dags_needing_dagruns(session, skip_dags).all() + self._create_dag_runs(dag_models, session) # commit the session - Release the write lock on DagModel table. guard.commit() @@ -972,7 +991,7 @@ def _should_update_dag_next_dagruns(self, dag, dag_model: DagModel, total_active total_active_runs, dag.max_active_runs, ) - dag_model.next_dagrun_create_after = None + self.last_check_time[dag_model.dag_id] = timezone.utcnow() return False return True diff --git a/airflow/models/dag.py b/airflow/models/dag.py index 5cf1731dc6c8c..7385414b8cc8b 100644 --- a/airflow/models/dag.py +++ b/airflow/models/dag.py @@ -2438,9 +2438,7 @@ def bulk_write_to_db(cls, dags: Collection["DAG"], session=NEW_SESSION): data_interval = None else: data_interval = dag.get_run_data_interval(run) - if num_active_runs.get(dag.dag_id, 0) >= orm_dag.max_active_runs: - orm_dag.next_dagrun_create_after = None - else: + if num_active_runs.get(dag.dag_id, 0) < orm_dag.max_active_runs: orm_dag.calculate_dagrun_date_fields(dag, data_interval) for orm_tag in list(orm_dag.tags): @@ -2851,7 +2849,7 @@ def deactivate_deleted_dags(cls, alive_dag_filelocs: List[str], session=NEW_SESS continue @classmethod - def dags_needing_dagruns(cls, session: Session): + def dags_needing_dagruns(cls, session: Session, skip_dags=None): """ Return (and lock) a list of Dag objects that are due to create a new DagRun. @@ -2863,6 +2861,9 @@ def dags_needing_dagruns(cls, session: Session): # We limit so that _one_ scheduler doesn't try to do all the creation # of dag runs + if skip_dags is None: + skip_dags = [] + query = ( session.query(cls) .filter( @@ -2870,6 +2871,7 @@ def dags_needing_dagruns(cls, session: Session): cls.is_active == expression.true(), cls.has_import_errors == expression.false(), cls.next_dagrun_create_after <= func.now(), + cls.dag_id.notin_(skip_dags), ) .order_by(cls.next_dagrun_create_after) .limit(cls.NUM_DAGS_PER_DAGRUN_QUERY) diff --git a/tests/jobs/test_scheduler_job.py b/tests/jobs/test_scheduler_job.py index 707f587223f0a..e12960225242f 100644 --- a/tests/jobs/test_scheduler_job.py +++ b/tests/jobs/test_scheduler_job.py @@ -52,7 +52,7 @@ from airflow.utils.callback_requests import DagCallbackRequest from airflow.utils.file import list_py_file_paths from airflow.utils.session import create_session, provide_session -from airflow.utils.state import DagRunState, State, TaskInstanceState +from airflow.utils.state import DagRunState, State from airflow.utils.types import DagRunType from tests.test_utils.asserts import assert_queries_count from tests.test_utils.config import conf_vars, env_vars @@ -1255,7 +1255,6 @@ def test_queued_dagruns_stops_creating_when_max_active_is_reached(self, dag_make assert session.query(DagRun).count() == 10 assert session.query(DagRun.state).filter(DagRun.state == State.RUNNING).count() == 10 assert session.query(DagRun.state).filter(DagRun.state == State.QUEUED).count() == 0 - assert orm_dag.next_dagrun_create_after is None def test_dagrun_timeout_verify_max_active_runs(self, dag_maker): """ @@ -1287,7 +1286,6 @@ def test_dagrun_timeout_verify_max_active_runs(self, dag_maker): assert len(drs) == 1 dr = drs[0] - assert orm_dag.next_dagrun_create_after is None # But we should record the date of _what run_ it would be assert isinstance(orm_dag.next_dagrun, datetime.datetime) assert isinstance(orm_dag.next_dagrun_data_interval_start, datetime.datetime) @@ -2926,51 +2924,70 @@ def test_do_schedule_max_active_runs_task_removed(self, session, dag_maker): ti.refresh_from_db(session=session) assert ti.state == State.QUEUED - def test_more_runs_are_not_created_when_max_active_runs_is_reached(self, dag_maker, caplog): + def test_runs_are_created_after_max_active_runs_was_reached(self, dag_maker, session): """ - This tests that when max_active_runs is reached, _create_dag_runs doesn't create - more dagruns + Test that when creating runs once max_active_runs is reached the runs does not stick + """ + with conf_vars({('scheduler', 'min_active_runs_check_interval'): '0'}): + self.scheduler_job = SchedulerJob(subdir=os.devnull) + self.scheduler_job.executor = MockExecutor(do_update=True) + self.scheduler_job.processor_agent = mock.MagicMock(spec=DagFileProcessorAgent) + + with dag_maker(max_active_runs=1, session=session) as dag: + # Need to use something that doesn't immediately get marked as success by the scheduler + BashOperator(task_id='task', bash_command='true') + + dag_run = dag_maker.create_dagrun( + state=State.RUNNING, + session=session, + ) + + # Reach max_active_runs + for _ in range(3): + self.scheduler_job._do_scheduling(session) + + # Complete dagrun + # Add dag_run back in to the session (_do_scheduling does an expunge_all) + dag_run = session.merge(dag_run) + session.refresh(dag_run) + dag_run.get_task_instance(task_id='task', session=session).state = State.SUCCESS + + # create new run + for _ in range(3): + self.scheduler_job._do_scheduling(session) + + # Assert that new runs has created + dag_runs = DagRun.find(dag_id=dag.dag_id, session=session) + assert len(dag_runs) == 2 + + def test_more_runs_are_not_created_when_max_active_runs_is_reached(self, dag_maker, session): + """ + This tests that when max_active_runs is reached, _create_dag_runs doesn't create more dagruns """ - with dag_maker(max_active_runs=1): - DummyOperator(task_id='task') self.scheduler_job = SchedulerJob(subdir=os.devnull) self.scheduler_job.executor = MockExecutor(do_update=False) self.scheduler_job.processor_agent = mock.MagicMock(spec=DagFileProcessorAgent) - session = settings.Session() - assert session.query(DagRun).count() == 0 - dag_models = DagModel.dags_needing_dagruns(session).all() - self.scheduler_job._create_dag_runs(dag_models, session) - dr = session.query(DagRun).one() - dr.state == DagRunState.QUEUED - assert session.query(DagRun).count() == 1 - assert dag_maker.dag_model.next_dagrun_create_after is None - session.flush() - # dags_needing_dagruns query should not return any value - assert len(DagModel.dags_needing_dagruns(session).all()) == 0 - self.scheduler_job._create_dag_runs(dag_models, session) - assert session.query(DagRun).count() == 1 - assert dag_maker.dag_model.next_dagrun_create_after is None - assert dag_maker.dag_model.next_dagrun == DEFAULT_DATE - # set dagrun to success - dr = session.query(DagRun).one() - dr.state = DagRunState.SUCCESS - ti = dr.get_task_instance('task', session) - ti.state = TaskInstanceState.SUCCESS - session.merge(ti) - session.merge(dr) - session.flush() - # check that next_dagrun is set properly by Schedulerjob._update_dag_next_dagruns - self.scheduler_job._schedule_dag_run(dr, session) - session.flush() - assert len(DagModel.dags_needing_dagruns(session).all()) == 1 - # assert next_dagrun has been updated correctly - assert dag_maker.dag_model.next_dagrun == DEFAULT_DATE + timedelta(days=1) - # assert no dagruns is created yet - assert ( - session.query(DagRun).filter(DagRun.state.in_([DagRunState.RUNNING, DagRunState.QUEUED])).count() - == 0 + + with dag_maker(max_active_runs=1) as dag: + # Need to use something that doesn't immediately get marked as success by the scheduler + BashOperator(task_id='dummy1', bash_command='true') + + dag_maker.create_dagrun( + state=State.RUNNING, + session=session, ) + for _ in range(3): + self.scheduler_job._do_scheduling(session) + + # Assert that no more dagruns has been created + dag_runs = DagRun.find(dag_id=dag.dag_id, session=session) + assert len(dag_runs) == 1 + + # Assert that the next_dagrun hasn't been updated + dm = DagModel.get_current(dag_id=dag.dag_id, session=session) + assert dm.next_dagrun == DEFAULT_DATE + def test_max_active_runs_creation_phasing(self, dag_maker, session): """ Test that when creating runs once max_active_runs is reached that the runs come in the right order @@ -2991,40 +3008,42 @@ def complete_one_dagrun(): self.clean_db() - with dag_maker(max_active_runs=3, session=session) as dag: - # Need to use something that doesn't immediately get marked as success by the scheduler - BashOperator(task_id='task', bash_command='true') + with conf_vars({('scheduler', 'min_active_runs_check_interval'): '0'}): + with dag_maker(max_active_runs=3, session=session) as dag: + # Need to use something that doesn't immediately get marked as success by the scheduler + BashOperator(task_id='task', bash_command='true') - self.scheduler_job = SchedulerJob(subdir=os.devnull) - self.scheduler_job.executor = MockExecutor(do_update=True) - self.scheduler_job.processor_agent = mock.MagicMock(spec=DagFileProcessorAgent) + self.scheduler_job = SchedulerJob(subdir=os.devnull) + self.scheduler_job.executor = MockExecutor(do_update=True) + self.scheduler_job.processor_agent = mock.MagicMock(spec=DagFileProcessorAgent) - DagModel.dags_needing_dagruns(session).all() - for _ in range(3): - self.scheduler_job._do_scheduling(session) + DagModel.dags_needing_dagruns(session).all() + for _ in range(3): + self.scheduler_job._do_scheduling(session) - model: DagModel = session.query(DagModel).get(dag.dag_id) + model: DagModel = session.query(DagModel).get(dag.dag_id) - # Pre-condition - assert DagRun.active_runs_of_dags(session=session) == {'test_dag': 3} + # Pre-condition + assert DagRun.active_runs_of_dags(session=session) == {'test_dag': 3} - assert model.next_dagrun == timezone.DateTime(2016, 1, 3, tzinfo=UTC) - assert model.next_dagrun_create_after is None + assert model.next_dagrun == timezone.DateTime(2016, 1, 3, tzinfo=UTC) - complete_one_dagrun() + complete_one_dagrun() - assert DagRun.active_runs_of_dags(session=session) == {'test_dag': 3} + assert DagRun.active_runs_of_dags(session=session) == {'test_dag': 3} - for _ in range(5): - self.scheduler_job._do_scheduling(session) - complete_one_dagrun() - model: DagModel = session.query(DagModel).get(dag.dag_id) + for _ in range(5): + self.scheduler_job._do_scheduling(session) + complete_one_dagrun() + model: DagModel = session.query(DagModel).get(dag.dag_id) - expected_execution_dates = [datetime.datetime(2016, 1, d, tzinfo=timezone.utc) for d in range(1, 6)] - dagrun_execution_dates = [ - dr.execution_date for dr in session.query(DagRun).order_by(DagRun.execution_date).all() - ] - assert dagrun_execution_dates == expected_execution_dates + expected_execution_dates = [ + datetime.datetime(2016, 1, d, tzinfo=timezone.utc) for d in range(1, 6) + ] + dagrun_execution_dates = [ + dr.execution_date for dr in session.query(DagRun).order_by(DagRun.execution_date).all() + ] + assert dagrun_execution_dates == expected_execution_dates def test_do_schedule_max_active_runs_and_manual_trigger(self, dag_maker): """ @@ -3813,3 +3832,32 @@ def test_catchup_works_correctly(self, dag_maker): .filter(DagRun.execution_date != DEFAULT_DATE) # exclude the first run .scalar() ) > (timezone.utcnow() - timedelta(days=2)) + + def test_get_recently_checked_dags(self, dag_maker): + with conf_vars({('scheduler', 'min_active_runs_check_interval'): '10'}): + self.scheduler_job = SchedulerJob(subdir=os.devnull) + + self.scheduler_job.last_check_time = {} + assert self.scheduler_job._get_recently_checked_dags() == [] + + self.scheduler_job.last_check_time = {'test_dag': timezone.datetime(2016, 1, 1, 20, 59, 59)} + assert ( + self.scheduler_job._get_recently_checked_dags(timezone.datetime(2016, 1, 1, 21, 1, 0)) == [] + ) + + self.scheduler_job.last_check_time = {'test_dag': timezone.datetime(2016, 1, 1, 20, 59, 59)} + assert self.scheduler_job._get_recently_checked_dags(timezone.datetime(2016, 1, 1, 21, 0, 5)) == [ + 'test_dag' + ] + + self.scheduler_job.last_check_time = { + 'test_dag0': timezone.datetime(2016, 1, 1, 20, 59, 50), + 'test_dag1': timezone.datetime(2016, 1, 1, 20, 59, 51), + 'test_dag2': timezone.datetime(2016, 1, 1, 0, 0, 0), + 'test_dag3': timezone.datetime(2016, 1, 1, 20, 59, 55), + 'test_dag4': timezone.datetime(2015, 1, 1, 20, 59, 59), + 'test_dag5': timezone.datetime(2016, 1, 1, 20, 59, 59), + } + assert sorted( + self.scheduler_job._get_recently_checked_dags(timezone.datetime(2016, 1, 1, 21, 0, 0)) + ) == ['test_dag1', 'test_dag3', 'test_dag5'] diff --git a/tests/models/test_dag.py b/tests/models/test_dag.py index f7e40ef99cf70..514aba1dc0386 100644 --- a/tests/models/test_dag.py +++ b/tests/models/test_dag.py @@ -786,43 +786,6 @@ def test_bulk_write_to_db(self): for row in session.query(DagModel.last_parsed_time).all(): assert row[0] is not None - @parameterized.expand([State.RUNNING, State.QUEUED]) - def test_bulk_write_to_db_max_active_runs(self, state): - """ - Test that DagModel.next_dagrun_create_after is set to NULL when the dag cannot be created due to max - active runs being hit. - """ - dag = DAG(dag_id='test_scheduler_verify_max_active_runs', start_date=DEFAULT_DATE) - dag.max_active_runs = 1 - - DummyOperator(task_id='dummy', dag=dag, owner='airflow') - - session = settings.Session() - dag.clear() - DAG.bulk_write_to_db([dag], session) - - model = session.query(DagModel).get((dag.dag_id,)) - - assert model.next_dagrun == DEFAULT_DATE - assert model.next_dagrun_create_after == DEFAULT_DATE + timedelta(days=1) - - dr = dag.create_dagrun( - state=state, - execution_date=model.next_dagrun, - run_type=DagRunType.SCHEDULED, - session=session, - ) - assert dr is not None - DAG.bulk_write_to_db([dag]) - - model = session.query(DagModel).get((dag.dag_id,)) - # We signal "at max active runs" by saying this run is never eligible to be created - assert model.next_dagrun_create_after is None - # test that bulk_write_to_db again doesn't update next_dagrun_create_after - DAG.bulk_write_to_db([dag]) - model = session.query(DagModel).get((dag.dag_id,)) - assert model.next_dagrun_create_after is None - def test_bulk_write_to_db_has_import_error(self): """ Test that DagModel.has_import_error is set to false if no import errors.