From 1aae23096b086f6a57294813ea0d6f4cf10a443b Mon Sep 17 00:00:00 2001 From: Ephraim Anierobi Date: Mon, 11 Oct 2021 22:52:14 +0100 Subject: [PATCH 1/3] Fix catchup by limiting queued dagrun creation using max_active_runs Currently, when catchup is True, we create a lot of dagruns limited by max_queued_runs_per_dag setting. This is not efficient as some dagruns takes longer to run. This PR brings back the old behaviour of not creating dagruns once max_active_runs is reached thereby solving the catchup issue. Now, the dagruns appears as though they were created in running state improve code and add more tests fixup! improve code and add more tests Add information about removing of max_queued_runs_per_dag add method for active runs remove comment deduplicate dag_ids Update airflow/models/dagrun.py Co-authored-by: Tzu-ping Chung Update airflow/jobs/scheduler_job.py Co-authored-by: Tzu-ping Chung Update airflow/jobs/scheduler_job.py Co-authored-by: Tzu-ping Chung --- UPDATING.md | 4 + airflow/config_templates/config.yml | 8 -- airflow/config_templates/default_airflow.cfg | 4 - airflow/jobs/scheduler_job.py | 60 +++++----- airflow/models/dag.py | 21 ++-- airflow/models/dagrun.py | 18 ++- tests/jobs/test_scheduler_job.py | 110 ++++++++++++++++--- tests/models/test_dag.py | 24 ++-- 8 files changed, 180 insertions(+), 69 deletions(-) diff --git a/UPDATING.md b/UPDATING.md index 322b88c606af2..4102af2838d2c 100644 --- a/UPDATING.md +++ b/UPDATING.md @@ -90,6 +90,10 @@ Now if you resolve a ``Param`` without a default and don't pass a value, you wil ```python Param().resolve() # raises TypeError ``` +### `max_queued_runs_per_dag` configuration has been removed + +The `max_queued_runs_per_dag` configuration option in `[core]` section has been removed. Previously, this controlled the number of queued dagrun +the scheduler can create in a dag. Now, the maximum number is controlled internally by the DAG's `max_active_runs` ## Airflow 2.2.0 diff --git a/airflow/config_templates/config.yml b/airflow/config_templates/config.yml index edc5b63f591f6..2136e536c8b41 100644 --- a/airflow/config_templates/config.yml +++ b/airflow/config_templates/config.yml @@ -195,14 +195,6 @@ type: string example: ~ default: "16" - - name: max_queued_runs_per_dag - description: | - The maximum number of queued dagruns for a single DAG. The scheduler will not create more DAG runs - if it reaches the limit. This is not configurable at the DAG level. - version_added: 2.1.4 - type: string - example: ~ - default: "16" - name: load_examples description: | Whether to load the DAG examples that ship with Airflow. It's good to diff --git a/airflow/config_templates/default_airflow.cfg b/airflow/config_templates/default_airflow.cfg index 41998eaa4d80f..d03a06bcca9b2 100644 --- a/airflow/config_templates/default_airflow.cfg +++ b/airflow/config_templates/default_airflow.cfg @@ -131,10 +131,6 @@ dags_are_paused_at_creation = True # which is defaulted as ``max_active_runs_per_dag``. max_active_runs_per_dag = 16 -# The maximum number of queued dagruns for a single DAG. The scheduler will not create more DAG runs -# if it reaches the limit. This is not configurable at the DAG level. -max_queued_runs_per_dag = 16 - # Whether to load the DAG examples that ship with Airflow. It's good to # get started, but you probably want to set this to ``False`` in a production # environment diff --git a/airflow/jobs/scheduler_job.py b/airflow/jobs/scheduler_job.py index a161c671b48d3..be959477b8f28 100644 --- a/airflow/jobs/scheduler_job.py +++ b/airflow/jobs/scheduler_job.py @@ -839,30 +839,19 @@ def _create_dag_runs(self, dag_models: Collection[DagModel], session: Session) - existing_dagruns = ( session.query(DagRun.dag_id, DagRun.execution_date).filter(existing_dagruns_filter).all() ) - max_queued_dagruns = conf.getint('core', 'max_queued_runs_per_dag') - queued_runs_of_dags = defaultdict( + active_runs_of_dags = defaultdict( int, - session.query(DagRun.dag_id, func.count('*')) - .filter( # We use `list` here because SQLA doesn't accept a set - # We use set to avoid duplicate dag_ids - DagRun.dag_id.in_(list({dm.dag_id for dm in dag_models})), - DagRun.state == State.QUEUED, - ) - .group_by(DagRun.dag_id) - .all(), + DagRun.active_runs_of_dags(dag_ids=(dm.dag_id for dm in dag_models), session=session), ) for dag_model in dag_models: - # Lets quickly check if we have exceeded the number of queued dagruns per dags - total_queued = queued_runs_of_dags[dag_model.dag_id] - if total_queued >= max_queued_dagruns: - continue dag = self.dagbag.get_dag(dag_model.dag_id, session=session) if not dag: self.log.error("DAG '%s' not found in serialized_dag table", dag_model.dag_id) continue + dag_hash = self.dagbag.dags_hash.get(dag.dag_id) data_interval = dag.get_next_data_interval(dag_model) @@ -885,12 +874,28 @@ def _create_dag_runs(self, dag_models: Collection[DagModel], session: Session) - dag_hash=dag_hash, creating_job_id=self.id, ) - queued_runs_of_dags[dag_model.dag_id] += 1 - dag_model.calculate_dagrun_date_fields(dag, data_interval) - + active_runs_of_dags[dag.dag_id] += 1 + self._update_dag_next_dagruns(dag, dag_model, active_runs_of_dags[dag.dag_id]) # TODO[HA]: Should we do a session.flush() so we don't have to keep lots of state/object in # memory for larger dags? or expunge_all() + def _update_dag_next_dagruns(self, dag, dag_model: DagModel, total_active_runs) -> None: + """ + Update the next_dagrun, next_dagrun_data_interval_start/end + and next_dagrun_create_after for this dag. + """ + if total_active_runs >= dag_model.max_active_runs: + self.log.info( + "DAG %s is at (or above) max_active_runs (%d of %d), not creating any more runs", + dag_model.dag_id, + total_active_runs, + dag_model.max_active_runs, + ) + dag_model.next_dagrun_create_after = None + else: + data_interval = dag.get_next_data_interval(dag_model) + dag_model.calculate_dagrun_date_fields(dag, data_interval) + def _start_queued_dagruns( self, session: Session, @@ -899,15 +904,10 @@ def _start_queued_dagruns( dag_runs = self._get_next_dagruns_to_examine(State.QUEUED, session) active_runs_of_dags = defaultdict( - lambda: 0, - session.query(DagRun.dag_id, func.count('*')) - .filter( # We use `list` here because SQLA doesn't accept a set - # We use set to avoid duplicate dag_ids - DagRun.dag_id.in_(list({dr.dag_id for dr in dag_runs})), - DagRun.state == State.RUNNING, - ) - .group_by(DagRun.dag_id) - .all(), + int, + DagRun.active_runs_of_dags( + (dr.dag_id for dr in dag_runs), only_running=True, session=session + ), ) def _update_state(dag: DAG, dag_run: DagRun): @@ -958,6 +958,7 @@ def _schedule_dag_run( if not dag: self.log.error("Couldn't find dag %s in DagBag/DB!", dag_run.dag_id) return 0 + dag_model = DM.get_dagmodel(dag.dag_id, session) if ( dag_run.start_date @@ -976,6 +977,9 @@ def _schedule_dag_run( session.merge(task_instance) session.flush() self.log.info("Run %s of %s has timed-out", dag_run.run_id, dag_run.dag_id) + active_runs = dag.get_num_active_runs(only_running=False, session=session) + # Work out if we should allow creating a new DagRun now? + self._update_dag_next_dagruns(dag, dag_model, active_runs) callback_to_execute = DagCallbackRequest( full_filepath=dag.fileloc, @@ -997,6 +1001,10 @@ def _schedule_dag_run( self._verify_integrity_if_dag_changed(dag_run=dag_run, session=session) # TODO[HA]: Rename update_state -> schedule_dag_run, ?? something else? schedulable_tis, callback_to_run = dag_run.update_state(session=session, execute_callbacks=False) + if dag_run.state in State.finished: + active_runs = dag.get_num_active_runs(only_running=False, session=session) + # Work out if we should allow creating a new DagRun now? + self._update_dag_next_dagruns(dag, dag_model, active_runs) # This will do one query per dag run. We "could" build up a complex # query to update all the TIs across all the execution dates and dag diff --git a/airflow/models/dag.py b/airflow/models/dag.py index 149606b992a12..0b7fdf59de05c 100644 --- a/airflow/models/dag.py +++ b/airflow/models/dag.py @@ -1138,7 +1138,7 @@ def get_active_runs(self): return active_dates @provide_session - def get_num_active_runs(self, external_trigger=None, session=None): + def get_num_active_runs(self, external_trigger=None, only_running=True, session=None): """ Returns the number of active "running" dag runs @@ -1148,11 +1148,11 @@ def get_num_active_runs(self, external_trigger=None, session=None): :return: number greater than 0 for active dag runs """ # .count() is inefficient - query = ( - session.query(func.count()) - .filter(DagRun.dag_id == self.dag_id) - .filter(DagRun.state == State.RUNNING) - ) + query = session.query(func.count()).filter(DagRun.dag_id == self.dag_id) + if only_running: + query = query.filter(DagRun.state == State.RUNNING) + else: + query = query.filter(DagRun.state.in_({State.RUNNING, State.QUEUED})) if external_trigger is not None: query = query.filter( @@ -2423,6 +2423,10 @@ def bulk_write_to_db(cls, dags: Collection["DAG"], session=None): ) most_recent_runs = {run.dag_id: run for run in most_recent_runs_iter} + # Get number of active dagruns for all dags we are processing as a single query. + + num_active_runs = DagRun.active_runs_of_dags(dag_ids=existing_dag_ids, session=session) + filelocs = [] for orm_dag in sorted(orm_dags, key=lambda d: d.dag_id): @@ -2451,7 +2455,10 @@ def bulk_write_to_db(cls, dags: Collection["DAG"], session=None): data_interval = None else: data_interval = dag.get_run_data_interval(run) - orm_dag.calculate_dagrun_date_fields(dag, data_interval) + if num_active_runs.get(dag.dag_id, 0) >= orm_dag.max_active_runs: + orm_dag.next_dagrun_create_after = None + else: + orm_dag.calculate_dagrun_date_fields(dag, data_interval) for orm_tag in list(orm_dag.tags): if orm_tag.name not in set(dag.tags): diff --git a/airflow/models/dagrun.py b/airflow/models/dagrun.py index 2b651c5f5873a..1d5326581295c 100644 --- a/airflow/models/dagrun.py +++ b/airflow/models/dagrun.py @@ -17,7 +17,7 @@ # under the License. import warnings from datetime import datetime -from typing import TYPE_CHECKING, Any, Iterable, List, NamedTuple, Optional, Tuple, Union +from typing import TYPE_CHECKING, Any, Dict, Iterable, List, NamedTuple, Optional, Tuple, Union from sqlalchemy import ( Boolean, @@ -207,6 +207,22 @@ def refresh_from_db(self, session: Session = None): self.id = dr.id self.state = dr.state + @classmethod + @provide_session + def active_runs_of_dags(cls, dag_ids=None, only_running=False, session=None) -> Dict[str, int]: + """Get the number of active dag runs for each dag.""" + query = session.query(cls.dag_id, func.count('*')) + if dag_ids is not None: + # 'set' called to avoid duplicate dag_ids, but converted back to 'list' + # because SQLAlchemy doesn't accept a set here. + query = query.filter(cls.dag_id.in_(list(set(dag_ids)))) + if only_running: + query = query.filter(cls.state == State.RUNNING) + else: + query = query.filter(cls.state.in_([State.RUNNING, State.QUEUED])) + query = query.group_by(cls.dag_id) + return {dag_id: count for dag_id, count in query.all()} + @classmethod def next_dagruns_to_examine( cls, diff --git a/tests/jobs/test_scheduler_job.py b/tests/jobs/test_scheduler_job.py index 8b00b2711f8d4..42b6619f0c8e2 100644 --- a/tests/jobs/test_scheduler_job.py +++ b/tests/jobs/test_scheduler_job.py @@ -49,7 +49,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 State +from airflow.utils.state import DagRunState, State, TaskInstanceState 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 @@ -1082,9 +1082,9 @@ def test_cleanup_methods_all_called(self, mock_processor_agent): self.scheduler_job.executor.end.assert_called_once() mock_processor_agent.return_value.end.reset_mock(side_effect=True) - def test_theres_limit_to_queued_dagruns_in_a_dag(self, dag_maker): - """This tests that there's limit to the number of queued dagrun scheduler can create in a dag""" - with dag_maker() as dag: + def test_queued_dagruns_stops_creating_when_max_active_is_reached(self, dag_maker): + """This tests that queued dagruns stops creating once max_active_runs is reached""" + with dag_maker(max_active_runs=10) as dag: DummyOperator(task_id='mytask') session = settings.Session() @@ -1099,13 +1099,20 @@ def test_theres_limit_to_queued_dagruns_in_a_dag(self, dag_maker): assert orm_dag is not None for _ in range(20): self.scheduler_job._create_dag_runs([orm_dag], session) - assert session.query(DagRun).count() == 16 + drs = session.query(DagRun).all() + assert len(drs) == 10 - with conf_vars({('core', 'max_queued_runs_per_dag'): '5'}): - clear_db_runs() - for i in range(20): - self.scheduler_job._create_dag_runs([orm_dag], session) - assert session.query(DagRun).count() == 5 + for dr in drs: + dr.state = State.RUNNING + session.merge(dr) + session.commit() + assert session.query(DagRun.state).filter(DagRun.state == State.RUNNING).count() == 10 + for _ in range(20): + self.scheduler_job._create_dag_runs([orm_dag], session) + 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): """ @@ -1137,9 +1144,7 @@ def test_dagrun_timeout_verify_max_active_runs(self, dag_maker): assert len(drs) == 1 dr = drs[0] - # This should have a value since we control max_active_runs - # by DagRun State. - assert orm_dag.next_dagrun_create_after + 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) @@ -1219,6 +1224,40 @@ def test_dagrun_timeout_fails_run(self, dag_maker): session.rollback() session.close() + def test_dagrun_timeout_fails_run_and_update_next_dagrun(self, dag_maker): + """ + Test that dagrun timeout fails run and update the next dagrun + """ + session = settings.Session() + with dag_maker( + max_active_runs=1, + dag_id='test_scheduler_fail_dagrun_timeout', + dagrun_timeout=datetime.timedelta(seconds=60), + ): + DummyOperator(task_id='dummy') + + dr = dag_maker.create_dagrun(start_date=timezone.utcnow() - datetime.timedelta(days=1)) + # check that next_dagrun is dr.execution_date + dag_maker.dag_model.next_dagrun == dr.execution_date + self.scheduler_job = SchedulerJob(subdir=os.devnull) + self.scheduler_job.dagbag = dag_maker.dagbag + + # Mock that processor_agent is started + self.scheduler_job.processor_agent = mock.Mock() + self.scheduler_job.processor_agent.send_callback_to_execute = mock.Mock() + + self.scheduler_job._schedule_dag_run(dr, session) + session.flush() + session.refresh(dr) + assert dr.state == State.FAILED + # check that next_dagrun has been updated by Schedulerjob._update_dag_next_dagruns + assert dag_maker.dag_model.next_dagrun == dr.execution_date + timedelta(days=1) + # check that no running/queued runs yet + assert ( + session.query(DagRun).filter(DagRun.state.in_([DagRunState.RUNNING, DagRunState.QUEUED])).count() + == 0 + ) + @pytest.mark.parametrize( "state, expected_callback_msg", [(State.SUCCESS, "success"), (State.FAILED, "task_failure")] ) @@ -2738,6 +2777,51 @@ 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): + """ + 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 + ) + def test_do_schedule_max_active_runs_and_manual_trigger(self, dag_maker): """ Make sure that when a DAG is already at max_active_runs, that manually triggered diff --git a/tests/models/test_dag.py b/tests/models/test_dag.py index dfa1a0a1ddd3c..5bea2b7d48ccd 100644 --- a/tests/models/test_dag.py +++ b/tests/models/test_dag.py @@ -726,7 +726,7 @@ def test_bulk_write_to_db(self): clear_db_dags() dags = [DAG(f'dag-bulk-sync-{i}', start_date=DEFAULT_DATE, tags=["test-dag"]) for i in range(0, 4)] - with assert_queries_count(4): + with assert_queries_count(5): DAG.bulk_write_to_db(dags) with create_session() as session: assert {'dag-bulk-sync-0', 'dag-bulk-sync-1', 'dag-bulk-sync-2', 'dag-bulk-sync-3'} == { @@ -743,14 +743,14 @@ def test_bulk_write_to_db(self): assert row[0] is not None # Re-sync should do fewer queries - with assert_queries_count(3): + with assert_queries_count(4): DAG.bulk_write_to_db(dags) - with assert_queries_count(3): + with assert_queries_count(4): DAG.bulk_write_to_db(dags) # Adding tags for dag in dags: dag.tags.append("test-dag2") - with assert_queries_count(4): + with assert_queries_count(5): DAG.bulk_write_to_db(dags) with create_session() as session: assert {'dag-bulk-sync-0', 'dag-bulk-sync-1', 'dag-bulk-sync-2', 'dag-bulk-sync-3'} == { @@ -769,7 +769,7 @@ def test_bulk_write_to_db(self): # Removing tags for dag in dags: dag.tags.remove("test-dag") - with assert_queries_count(4): + with assert_queries_count(5): DAG.bulk_write_to_db(dags) with create_session() as session: assert {'dag-bulk-sync-0', 'dag-bulk-sync-1', 'dag-bulk-sync-2', 'dag-bulk-sync-3'} == { @@ -785,7 +785,8 @@ def test_bulk_write_to_db(self): for row in session.query(DagModel.last_parsed_time).all(): assert row[0] is not None - def test_bulk_write_to_db_max_active_runs(self): + @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. @@ -805,7 +806,7 @@ def test_bulk_write_to_db_max_active_runs(self): assert model.next_dagrun_create_after == DEFAULT_DATE + timedelta(days=1) dr = dag.create_dagrun( - state=State.RUNNING, + state=state, execution_date=model.next_dagrun, run_type=DagRunType.SCHEDULED, session=session, @@ -814,9 +815,12 @@ def test_bulk_write_to_db_max_active_runs(self): DAG.bulk_write_to_db([dag]) model = session.query(DagModel).get((dag.dag_id,)) - assert model.next_dagrun == DEFAULT_DATE + timedelta(days=1) - # Next dagrun after is not None because the dagrun would be in queued state - assert model.next_dagrun_create_after is not None + # 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_sync_to_db(self): dag = DAG( From 38ba3569e680ce56192e8288585b1a268523fc87 Mon Sep 17 00:00:00 2001 From: Ephraim Anierobi Date: Wed, 20 Oct 2021 10:23:08 +0100 Subject: [PATCH 2/3] fixup! Fix catchup by limiting queued dagrun creation using max_active_runs --- airflow/jobs/scheduler_job.py | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/airflow/jobs/scheduler_job.py b/airflow/jobs/scheduler_job.py index be959477b8f28..2827de8551b98 100644 --- a/airflow/jobs/scheduler_job.py +++ b/airflow/jobs/scheduler_job.py @@ -905,9 +905,7 @@ def _start_queued_dagruns( active_runs_of_dags = defaultdict( int, - DagRun.active_runs_of_dags( - (dr.dag_id for dr in dag_runs), only_running=True, session=session - ), + DagRun.active_runs_of_dags((dr.dag_id for dr in dag_runs), only_running=True, session=session), ) def _update_state(dag: DAG, dag_run: DagRun): From d8f024106307da93b755ae2a7331022bd9c7ee21 Mon Sep 17 00:00:00 2001 From: Jed Cunningham <66968678+jedcunningham@users.noreply.github.com> Date: Wed, 20 Oct 2021 11:44:52 -0600 Subject: [PATCH 3/3] Update UPDATING.md --- UPDATING.md | 1 + 1 file changed, 1 insertion(+) diff --git a/UPDATING.md b/UPDATING.md index 4102af2838d2c..ab2e81bfb6306 100644 --- a/UPDATING.md +++ b/UPDATING.md @@ -90,6 +90,7 @@ Now if you resolve a ``Param`` without a default and don't pass a value, you wil ```python Param().resolve() # raises TypeError ``` + ### `max_queued_runs_per_dag` configuration has been removed The `max_queued_runs_per_dag` configuration option in `[core]` section has been removed. Previously, this controlled the number of queued dagrun