kaxil commented on code in PR #64294:
URL: https://github.com/apache/airflow/pull/64294#discussion_r3124814684
##########
airflow-core/src/airflow/config_templates/config.yml:
##########
@@ -2568,6 +2568,17 @@ scheduler:
type: integer
default: "20"
see_also: ":ref:`scheduler:ha:tunables`"
+ max_new_dagruns_per_loop_to_schedule:
+ description: |
+ How many NEW dagruns should be scheduled and examined (and locked)
when scheduling
+ and queuing tasks.
+ If set, select `max_dagruns_per_loop_to_schedule` old dagruns (that
have been
+ examined before)
+ And `max_new_dagruns_per_loop_to_schedule` new dagruns (that have not
yet been examined).
+ example: ~
+ version_added: 3.2.2
Review Comment:
`version_added: 3.2.2` doesn't match a real release -- main is 3.3.0 and
there's no planned 3.2.2. Should be `3.3.0`.
##########
airflow-core/src/airflow/config_templates/config.yml:
##########
@@ -2568,6 +2568,17 @@ scheduler:
type: integer
default: "20"
see_also: ":ref:`scheduler:ha:tunables`"
+ max_new_dagruns_per_loop_to_schedule:
+ description: |
+ How many NEW dagruns should be scheduled and examined (and locked)
when scheduling
+ and queuing tasks.
+ If set, select `max_dagruns_per_loop_to_schedule` old dagruns (that
have been
+ examined before)
+ And `max_new_dagruns_per_loop_to_schedule` new dagruns (that have not
yet been examined).
+ example: ~
Review Comment:
This description is hard to parse -- the wrap across "examined before)" /
"And `max_new...`" reads as a run-on sentence, and "NEW" is capitalized
arbitrarily. It also doesn't mention that enabling this increases the total
dagruns locked per loop to the sum of both limits, which is load users should
know about before turning it on.
Something like:
```
When set > 0, the scheduler runs a second query per loop that fetches up to
this
many dagruns that have never been examined (``last_scheduling_decision IS
NULL``),
in addition to the ``max_dagruns_per_loop_to_schedule`` already-examined
ones.
This prevents starvation of older dagruns when large batches of new dagruns
are
created at once (for example, via ``TriggerDagRunOperator``). Note that the
total
number of dagruns locked and scheduled per loop becomes the sum of both
limits.
```
##########
airflow-core/src/airflow/models/dagrun.py:
##########
@@ -615,7 +629,7 @@ def active_runs_of_dags(
@classmethod
@retry_db_transaction
- def get_running_dag_runs_to_examine(cls, session: Session) ->
ScalarResult[DagRun]:
+ def get_running_dag_runs_to_examine(cls, session: Session) ->
Sequence[DagRun]:
Review Comment:
Two things on this signature change:
1. The docstring below wasn't updated for the new two-query behavior (fetch
already-examined + new dagruns separately).
2. Changing the return type from `ScalarResult[DagRun]` to
`Sequence[DagRun]` is a breaking change for any downstream caller (custom jobs,
plugins, tests outside this repo) that does `.all()` on the result --
`ScalarResult` has `.all()`, a plain list/Sequence doesn't need it. Worth
calling out in the docstring and adding a newsfragment.
This PR is also missing a newsfragment for the new
`max_new_dagruns_per_loop_to_schedule` config option.
##########
airflow-core/src/airflow/models/dagrun.py:
##########
@@ -628,27 +642,61 @@ def get_running_dag_runs_to_examine(cls, session:
Session) -> ScalarResult[DagRu
from airflow.models.backfill import BackfillDagRun
from airflow.models.dag import DagModel
- query = (
- select(cls)
- .with_hint(cls, "USE INDEX (idx_dag_run_running_dags)",
dialect_name="mysql")
- .where(cls.state == DagRunState.RUNNING)
- .join(DagModel, DagModel.dag_id == cls.dag_id)
- .join(BackfillDagRun, BackfillDagRun.dag_run_id == DagRun.id,
isouter=True)
- .where(
- DagModel.is_paused == false(),
- DagModel.is_stale == false(),
- )
- .order_by(
- nulls_first(cast("ColumnElement[Any]",
BackfillDagRun.sort_ordinal), session=session),
- nulls_first(cast("ColumnElement[Any]",
cls.last_scheduling_decision), session=session),
- cls.run_after,
+ def _get_dagrun_query(
+ filters: list[ColumnElement[bool]], order_by:
list[SQLColumnExpression[Any]], limit: int
+ ):
+ return (
+ select(DagRun)
+ .with_hint(DagRun, "USE INDEX (idx_dag_run_running_dags)",
dialect_name="mysql")
+ .where(DagRun.state == DagRunState.RUNNING)
+ .join(DagModel, DagModel.dag_id == cls.dag_id)
+ .join(BackfillDagRun, BackfillDagRun.dag_run_id == DagRun.id,
isouter=True)
+ .where(*filters)
+ .order_by(*order_by)
+ .limit(limit)
)
- .limit(cls.DEFAULT_DAGRUNS_TO_EXAMINE)
+
+ filters = [
+ DagRun.run_after <= func.now(),
+ DagModel.is_paused == false(),
+ DagModel.is_stale == false(),
+ ]
+
+ order = [
+ nulls_first(cast("ColumnElement[Any]",
BackfillDagRun.sort_ordinal), session=session),
+ nulls_first(cast("ColumnElement[Any]",
DagRun.last_scheduling_decision), session=session),
+ DagRun.run_after,
+ ]
+
+ new_dagruns_to_examine = max(cls.DEFAULT_NEW_DAGRUNS_TO_EXAMINE, 0)
+ dagruns_to_examine = cls.DEFAULT_DAGRUNS_TO_EXAMINE
+
+ query = _get_dagrun_query(
+ filters=filters
+ if new_dagruns_to_examine == 0
+ else [*filters, DagRun.last_scheduling_decision.is_not(None)],
Review Comment:
The inline ternary inside the kwarg is hard to follow. You already have the
two-branch logic below (`if new_dagruns_to_examine > 0:` builds a second
query), so it's cleaner to build this query's filters once before calling:
```python
old_filters = (
[*filters, DagRun.last_scheduling_decision.is_not(None)]
if new_dagruns_to_examine > 0
else filters
)
query = _get_dagrun_query(filters=old_filters, order_by=order,
limit=dagruns_to_examine)
```
Reads top-to-bottom and avoids the inverted `== 0` / `is_not(None)` pairing.
##########
airflow-core/src/airflow/models/dagrun.py:
##########
@@ -628,27 +642,61 @@ def get_running_dag_runs_to_examine(cls, session:
Session) -> ScalarResult[DagRu
from airflow.models.backfill import BackfillDagRun
from airflow.models.dag import DagModel
- query = (
- select(cls)
- .with_hint(cls, "USE INDEX (idx_dag_run_running_dags)",
dialect_name="mysql")
- .where(cls.state == DagRunState.RUNNING)
- .join(DagModel, DagModel.dag_id == cls.dag_id)
- .join(BackfillDagRun, BackfillDagRun.dag_run_id == DagRun.id,
isouter=True)
- .where(
- DagModel.is_paused == false(),
- DagModel.is_stale == false(),
- )
- .order_by(
- nulls_first(cast("ColumnElement[Any]",
BackfillDagRun.sort_ordinal), session=session),
- nulls_first(cast("ColumnElement[Any]",
cls.last_scheduling_decision), session=session),
- cls.run_after,
+ def _get_dagrun_query(
+ filters: list[ColumnElement[bool]], order_by:
list[SQLColumnExpression[Any]], limit: int
+ ):
+ return (
+ select(DagRun)
+ .with_hint(DagRun, "USE INDEX (idx_dag_run_running_dags)",
dialect_name="mysql")
+ .where(DagRun.state == DagRunState.RUNNING)
+ .join(DagModel, DagModel.dag_id == cls.dag_id)
+ .join(BackfillDagRun, BackfillDagRun.dag_run_id == DagRun.id,
isouter=True)
+ .where(*filters)
+ .order_by(*order_by)
+ .limit(limit)
)
- .limit(cls.DEFAULT_DAGRUNS_TO_EXAMINE)
+
+ filters = [
+ DagRun.run_after <= func.now(),
+ DagModel.is_paused == false(),
+ DagModel.is_stale == false(),
+ ]
+
+ order = [
+ nulls_first(cast("ColumnElement[Any]",
BackfillDagRun.sort_ordinal), session=session),
+ nulls_first(cast("ColumnElement[Any]",
DagRun.last_scheduling_decision), session=session),
+ DagRun.run_after,
+ ]
+
+ new_dagruns_to_examine = max(cls.DEFAULT_NEW_DAGRUNS_TO_EXAMINE, 0)
+ dagruns_to_examine = cls.DEFAULT_DAGRUNS_TO_EXAMINE
+
+ query = _get_dagrun_query(
+ filters=filters
+ if new_dagruns_to_examine == 0
+ else [*filters, DagRun.last_scheduling_decision.is_not(None)],
+ order_by=order,
+ limit=dagruns_to_examine,
)
- query = query.where(DagRun.run_after <= func.now())
+ result: Sequence[DagRun] = (
+ session.scalars(with_row_locks(query, of=cls, session=session,
skip_locked=True)).unique().all()
+ )
+
+ if new_dagruns_to_examine > 0:
+ new_dagruns_query = _get_dagrun_query(
+ filters=[*filters, DagRun.last_scheduling_decision.is_(None)],
+ order_by=order,
+ limit=new_dagruns_to_examine,
+ )
+ new_dagruns: Sequence[DagRun] = (
+ session.scalars(with_row_locks(new_dagruns_query, of=cls,
session=session, skip_locked=True))
+ .unique()
+ .all()
+ )
+
+ result = [*result, *new_dagruns]
Review Comment:
Minor: `result.extend(new_dagruns)` avoids the extra copy from `[*result,
*new_dagruns]`. Not a correctness issue, just a style/memory nit.
##########
airflow-core/tests/unit/models/test_dagrun.py:
##########
@@ -993,6 +993,121 @@ def test_wait_for_downstream(self, dag_maker, session,
prev_ti_state, is_ti_sche
schedulable_tis = [ti.task_id for ti in decision.schedulable_tis]
assert (upstream.task_id in schedulable_tis) == is_ti_schedulable
+ def
test_get_running_dag_runs_ignores_new_dagruns_to_examine_when_smaller_than_0(
Review Comment:
Test name says `..._when_smaller_than_0` but the body sets
`DEFAULT_NEW_DAGRUNS_TO_EXAMINE = 0`, not a negative value. Either rename to
`..._when_equal_to_0` (matches what's tested), or parametrize over `[0, -1]` so
both paths are actually exercised -- the
`max(cls.DEFAULT_NEW_DAGRUNS_TO_EXAMINE, 0)` clamp on line 671 deserves
coverage for the negative case too.
##########
airflow-core/tests/unit/models/test_dagrun.py:
##########
@@ -993,6 +993,121 @@ def test_wait_for_downstream(self, dag_maker, session,
prev_ti_state, is_ti_sche
schedulable_tis = [ti.task_id for ti in decision.schedulable_tis]
assert (upstream.task_id in schedulable_tis) == is_ti_schedulable
+ def
test_get_running_dag_runs_ignores_new_dagruns_to_examine_when_smaller_than_0(
+ self, session, dag_maker, monkeypatch
+ ):
+
+ monkeypatch.setattr(DagRun, "DEFAULT_NEW_DAGRUNS_TO_EXAMINE", 0)
+
+ def create_dagruns(
+ last_scheduling_decision: datetime.datetime | None = None,
+ count: int = 20,
+ ):
+ dagrun = dag_maker.create_dagrun(
+ run_type=DagRunType.SCHEDULED,
+ state=State.RUNNING,
+ run_after=datetime.datetime(2024, 1, 1),
+ )
+ dagrun.last_scheduling_decision = last_scheduling_decision
+ session.merge(dagrun)
+ for _ in range(count - 1):
+ dagrun = dag_maker.create_dagrun_after(
+ dagrun,
+ run_type=DagRunType.SCHEDULED,
+ state=State.RUNNING,
+ run_after=datetime.datetime(2024, 1, 1),
+ )
+
+ dagrun.last_scheduling_decision = last_scheduling_decision
+ session.merge(dagrun)
+
+ with dag_maker(
+ dag_id="dummy_dag",
+ schedule=datetime.timedelta(days=1),
+ start_date=datetime.datetime(2024, 1, 1),
+ session=session,
+ ):
+ EmptyOperator(task_id="dummy_task")
+
+ create_dagruns(None, 10)
+
+ with dag_maker(
+ dag_id="dummy_dag2",
+ schedule=datetime.timedelta(days=1),
+ start_date=datetime.datetime(2024, 1, 1),
+ session=session,
+ ):
+ EmptyOperator(task_id="dummy_task2")
+
+ create_dagruns(func.now(), 20)
Review Comment:
Passing `func.now()` as `last_scheduling_decision` is a SQL expression, not
a Python datetime. It works because SQLAlchemy flushes it as a deferred UPDATE,
but it's confusing -- the type annotation on `create_dagruns` above is
`datetime.datetime | None`, and a reader has to know SQLAlchemy internals to
square the two. A literal like `datetime.datetime(2024, 1, 2,
tzinfo=datetime.timezone.utc)` or `timezone.utcnow()` would be clearer and
deterministic (the SQL `now()` also makes the test sensitive to DB clock drift
/ fixture timing).
Also: `create_dagruns` is duplicated between the two new tests with
identical bodies. Worth lifting to a module-level helper or a fixture.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
