You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@airflow.apache.org by je...@apache.org on 2021/10/20 21:26:54 UTC
[airflow] 09/10: Fix catchup by limiting queued dagrun creation
using max_active_runs (#18897)
This is an automated email from the ASF dual-hosted git repository.
jedcunningham pushed a commit to branch v2-2-test
in repository https://gitbox.apache.org/repos/asf/airflow.git
commit 69048f788983ae5fef5a1eacb8d1118ba92ae47a
Author: Ephraim Anierobi <sp...@gmail.com>
AuthorDate: Wed Oct 20 21:33:45 2021 +0100
Fix catchup by limiting queued dagrun creation using max_active_runs (#18897)
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 behavior 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
(cherry picked from commit 05eea002fbd6dfc97ffc207f2a7f73ccf49aacde)
---
UPDATING.md | 5 ++
airflow/config_templates/config.yml | 8 --
airflow/config_templates/default_airflow.cfg | 4 -
airflow/jobs/scheduler_job.py | 58 +++++++-------
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, 179 insertions(+), 69 deletions(-)
diff --git a/UPDATING.md b/UPDATING.md
index 9c0c4fd..443afe9 100644
--- a/UPDATING.md
+++ b/UPDATING.md
@@ -91,6 +91,11 @@ Now if you resolve a ``Param`` without a default and don't pass a value, you wil
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
### `worker_log_server_port` configuration has been moved to the ``logging`` section.
diff --git a/airflow/config_templates/config.yml b/airflow/config_templates/config.yml
index 0e58f26..7aefbe0 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 41998ea..d03a06b 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 17fc55c..bf2410a 100644
--- a/airflow/jobs/scheduler_job.py
+++ b/airflow/jobs/scheduler_job.py
@@ -848,30 +848,19 @@ class SchedulerJob(BaseJob):
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)
@@ -894,12 +883,28 @@ class SchedulerJob(BaseJob):
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,
@@ -908,15 +913,8 @@ class SchedulerJob(BaseJob):
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):
@@ -967,6 +965,7 @@ class SchedulerJob(BaseJob):
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
@@ -985,6 +984,9 @@ class SchedulerJob(BaseJob):
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,
@@ -1006,6 +1008,10 @@ class SchedulerJob(BaseJob):
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 b372cfe..55be0ec 100644
--- a/airflow/models/dag.py
+++ b/airflow/models/dag.py
@@ -1138,7 +1138,7 @@ class DAG(LoggingMixin):
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 @@ class DAG(LoggingMixin):
: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(
@@ -2425,6 +2425,10 @@ class DAG(LoggingMixin):
)
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):
@@ -2453,7 +2457,10 @@ class DAG(LoggingMixin):
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 2b651c5..1d53265 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,
@@ -208,6 +208,22 @@ class DagRun(Base, LoggingMixin):
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,
state: DagRunState,
diff --git a/tests/jobs/test_scheduler_job.py b/tests/jobs/test_scheduler_job.py
index 11a2fca..16b2921 100644
--- a/tests/jobs/test_scheduler_job.py
+++ b/tests/jobs/test_scheduler_job.py
@@ -49,7 +49,7 @@ from airflow.utils import timezone
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
@@ -1105,9 +1105,9 @@ class TestSchedulerJob:
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()
@@ -1122,13 +1122,20 @@ class TestSchedulerJob:
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):
"""
@@ -1160,9 +1167,7 @@ class TestSchedulerJob:
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)
@@ -1242,6 +1247,40 @@ class TestSchedulerJob:
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")]
)
@@ -2756,6 +2795,51 @@ class TestSchedulerJob:
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 dfa1a0a..5bea2b7 100644
--- a/tests/models/test_dag.py
+++ b/tests/models/test_dag.py
@@ -726,7 +726,7 @@ class TestDag(unittest.TestCase):
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 @@ class TestDag(unittest.TestCase):
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 @@ class TestDag(unittest.TestCase):
# 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 @@ class TestDag(unittest.TestCase):
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 @@ class TestDag(unittest.TestCase):
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 @@ class TestDag(unittest.TestCase):
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(