Removes the next_dagrun_create_after reset

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:

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

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
 

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,13 +2861,17 @@ 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(
                 cls.is_paused == expression.false(),
                 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)

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']