kaxil commented on code in PR #71072:
URL: https://github.com/apache/airflow/pull/71072#discussion_r4237597539
##########
airflow-core/src/airflow/jobs/scheduler_job_runner.py:
##########
@@ -2395,8 +2421,96 @@ def _create_dagruns_for_partitioned_asset_dags(self,
session: Session) -> set[st
)
).all()
if not pending_apdrs:
+ # Backlog fully drained: the next tick that re-crosses the cap is
a new episode,
+ # not a continuation of whatever was reported before.
+ self._partition_cap_backlog_reported = False
Review Comment:
With three or more schedulers an empty fetch is as ambiguous as a short one.
Each scheduler locks at most `cap` rows, so with a backlog of 900, A and C can
hold all of it while B's fetch comes back empty; B clears its flag, and its
next full fetch logs a fresh warning and writes another audit row for the same
episode, so the "at most one audit row per episode" comment at L402 doesn't
hold there. Running the unlocked count before clearing when the flag is set
would close it (a real drain to zero still resets), and it only costs a query
on empty ticks inside an episode. Not blocking, the extra row is noise only.
##########
airflow-core/tests/unit/jobs/test_scheduler_job.py:
##########
@@ -14345,6 +14415,617 @@ def
test_partition_cap_at_n_minus_one_leaves_one_pending(dag_maker: DagMaker, se
assert partition_dags == {"cap-consumer-n-minus-one"}
Review Comment:
This test now crosses the cap (3 pending, cap 2), so it commits an audit
`Log` row through the separate session, but it doesn't use
`clear_audit_log_rows`, so that row outlives the test. Adding the fixture here
keeps it isolated.
##########
airflow-core/src/airflow/jobs/scheduler_job_runner.py:
##########
@@ -2395,8 +2421,96 @@ def _create_dagruns_for_partitioned_asset_dags(self,
session: Session) -> set[st
)
).all()
if not pending_apdrs:
+ # Backlog fully drained: the next tick that re-crosses the cap is
a new episode,
+ # not a continuation of whatever was reported before.
+ self._partition_cap_backlog_reported = False
return set()
+ sorted_dag_ids: list[str] = []
+ if (
+ len(pending_apdrs) >= self._max_partition_dag_runs_per_loop
+ or self._partition_cap_backlog_reported
+ ):
+ # A full fetch alone can't tell us whether that's the entire
backlog or just
+ # this tick's slice of a larger one, and a short fetch can't
either while other
+ # schedulers hold row locks, so we only pay for this query on a
full fetch or
+ # while an episode is being tracked.
+ # Per-Dag counts across the *whole* backlog, not just this tick's
oldest-cap
+ # slice (`pending_apdrs`) — a Dag whose partitions haven't reached
the front of
+ # the FIFO queue yet would otherwise be missing from the log/audit
row until its
+ # turn comes up.
+ backlogs_per_dag: dict[str, int] = {
+ dag_id: count
+ for dag_id, count in session.execute(
+ select(AssetPartitionDagRun.target_dag_id, func.count())
+ .select_from(AssetPartitionDagRun)
+ .join(DagModel, DagModel.dag_id ==
AssetPartitionDagRun.target_dag_id)
+ .where(*_build_pending_partitioned_apdr_filters())
+ .group_by(AssetPartitionDagRun.target_dag_id)
+ )
+ }
+ backlog_total = sum(backlogs_per_dag.values())
+ # No SQL-side ORDER BY: it would have no dag_id tie-break, and
relying on dict
+ # insertion order mirroring DB row order across backends (SQLite
in tests vs
+ # Postgres/MySQL in production) isn't a contract worth trusting —
sort explicitly.
+ sorted_dag_ids = sorted(backlogs_per_dag, key=lambda dag_id:
(-backlogs_per_dag[dag_id], dag_id))
+ if backlog_total <= self._max_partition_dag_runs_per_loop:
+ self._partition_cap_backlog_reported = False
+ else:
+ backlog_total = len(pending_apdrs)
+
+ if backlog_total > self._max_partition_dag_runs_per_loop:
+ displayed_dag_ids =
sorted_dag_ids[:MAX_PARTITION_CAP_BACKLOG_DAG_IDS_LOGGED]
+ remaining_dag_id_count = len(sorted_dag_ids) -
len(displayed_dag_ids)
+ log_cap_reached = self.log.debug if
self._partition_cap_backlog_reported else self.log.warning
+ log_cap_reached(
+ (
+ "Reached the per-tick cap on pending partitioned Dag runs;
the remaining backlog "
+ "will be evaluated over subsequent scheduler ticks"
+ ),
+ cap=self._max_partition_dag_runs_per_loop,
+ backlog_total=backlog_total,
+ dag_ids=displayed_dag_ids,
+ )
+ # Edge-trigger the audit row: a persistent backlog re-hits this
branch every tick,
+ # and writing a `Log` row that often would flood the audit table.
Write it once per
+ # backlog episode; `_partition_cap_backlog_reported` is cleared
above, only when
+ # the fetch is empty or the count shows the backlog at or below
the cap.
+ if not self._partition_cap_backlog_reported:
+ # A separate, independently-committed session is required
here: the caller
+ # (`_create_dagruns_for_dags`) runs under
`@retry_db_transaction`, which rolls
+ # back *session* on a `DBAPIError` — sharing that transaction
would silently
+ # discard this audit row along with the rest of the tick's
work.
+ # Invariant: this must run before *session* makes any writes
this tick, or the
+ # new connection's commit can lock-contend with it on SQLite.
+ more_suffix = (
+ f", and {remaining_dag_id_count} more with fewer pending
runs"
+ if remaining_dag_id_count
+ else ""
+ )
+ try:
+ with create_session(scoped=False) as audit_session:
+ audit_session.add(
+ Log(
+ event="partition Dag run cap reached",
Review Comment:
Could this set `owner="scheduler", owner_display_name="Scheduler"` like the
`drain_completed` row in `_finalize_draining_dags`? Otherwise it shows up in
the audit log with an empty owner.
##########
airflow-core/src/airflow/jobs/scheduler_job_runner.py:
##########
@@ -2395,8 +2421,96 @@ def _create_dagruns_for_partitioned_asset_dags(self,
session: Session) -> set[st
)
).all()
if not pending_apdrs:
+ # Backlog fully drained: the next tick that re-crosses the cap is
a new episode,
+ # not a continuation of whatever was reported before.
+ self._partition_cap_backlog_reported = False
return set()
+ sorted_dag_ids: list[str] = []
+ if (
+ len(pending_apdrs) >= self._max_partition_dag_runs_per_loop
+ or self._partition_cap_backlog_reported
+ ):
+ # A full fetch alone can't tell us whether that's the entire
backlog or just
+ # this tick's slice of a larger one, and a short fetch can't
either while other
+ # schedulers hold row locks, so we only pay for this query on a
full fetch or
+ # while an episode is being tracked.
+ # Per-Dag counts across the *whole* backlog, not just this tick's
oldest-cap
+ # slice (`pending_apdrs`) — a Dag whose partitions haven't reached
the front of
+ # the FIFO queue yet would otherwise be missing from the log/audit
row until its
+ # turn comes up.
+ backlogs_per_dag: dict[str, int] = {
+ dag_id: count
+ for dag_id, count in session.execute(
+ select(AssetPartitionDagRun.target_dag_id, func.count())
+ .select_from(AssetPartitionDagRun)
+ .join(DagModel, DagModel.dag_id ==
AssetPartitionDagRun.target_dag_id)
+ .where(*_build_pending_partitioned_apdr_filters())
+ .group_by(AssetPartitionDagRun.target_dag_id)
+ )
+ }
+ backlog_total = sum(backlogs_per_dag.values())
+ # No SQL-side ORDER BY: it would have no dag_id tie-break, and
relying on dict
+ # insertion order mirroring DB row order across backends (SQLite
in tests vs
+ # Postgres/MySQL in production) isn't a contract worth trusting —
sort explicitly.
+ sorted_dag_ids = sorted(backlogs_per_dag, key=lambda dag_id:
(-backlogs_per_dag[dag_id], dag_id))
+ if backlog_total <= self._max_partition_dag_runs_per_loop:
+ self._partition_cap_backlog_reported = False
+ else:
+ backlog_total = len(pending_apdrs)
+
+ if backlog_total > self._max_partition_dag_runs_per_loop:
+ displayed_dag_ids =
sorted_dag_ids[:MAX_PARTITION_CAP_BACKLOG_DAG_IDS_LOGGED]
+ remaining_dag_id_count = len(sorted_dag_ids) -
len(displayed_dag_ids)
+ log_cap_reached = self.log.debug if
self._partition_cap_backlog_reported else self.log.warning
Review Comment:
The PR description still says this warns on every tick. With the debug
demotion it's one warning per episode per scheduler, so the description should
say that.
##########
airflow-core/src/airflow/jobs/scheduler_job_runner.py:
##########
@@ -2395,8 +2421,96 @@ def _create_dagruns_for_partitioned_asset_dags(self,
session: Session) -> set[st
)
).all()
if not pending_apdrs:
+ # Backlog fully drained: the next tick that re-crosses the cap is
a new episode,
+ # not a continuation of whatever was reported before.
+ self._partition_cap_backlog_reported = False
return set()
+ sorted_dag_ids: list[str] = []
+ if (
+ len(pending_apdrs) >= self._max_partition_dag_runs_per_loop
+ or self._partition_cap_backlog_reported
+ ):
+ # A full fetch alone can't tell us whether that's the entire
backlog or just
+ # this tick's slice of a larger one, and a short fetch can't
either while other
+ # schedulers hold row locks, so we only pay for this query on a
full fetch or
+ # while an episode is being tracked.
+ # Per-Dag counts across the *whole* backlog, not just this tick's
oldest-cap
+ # slice (`pending_apdrs`) — a Dag whose partitions haven't reached
the front of
+ # the FIFO queue yet would otherwise be missing from the log/audit
row until its
+ # turn comes up.
+ backlogs_per_dag: dict[str, int] = {
+ dag_id: count
+ for dag_id, count in session.execute(
+ select(AssetPartitionDagRun.target_dag_id, func.count())
+ .select_from(AssetPartitionDagRun)
+ .join(DagModel, DagModel.dag_id ==
AssetPartitionDagRun.target_dag_id)
+ .where(*_build_pending_partitioned_apdr_filters())
+ .group_by(AssetPartitionDagRun.target_dag_id)
+ )
+ }
+ backlog_total = sum(backlogs_per_dag.values())
+ # No SQL-side ORDER BY: it would have no dag_id tie-break, and
relying on dict
+ # insertion order mirroring DB row order across backends (SQLite
in tests vs
+ # Postgres/MySQL in production) isn't a contract worth trusting —
sort explicitly.
+ sorted_dag_ids = sorted(backlogs_per_dag, key=lambda dag_id:
(-backlogs_per_dag[dag_id], dag_id))
+ if backlog_total <= self._max_partition_dag_runs_per_loop:
+ self._partition_cap_backlog_reported = False
+ else:
+ backlog_total = len(pending_apdrs)
+
+ if backlog_total > self._max_partition_dag_runs_per_loop:
+ displayed_dag_ids =
sorted_dag_ids[:MAX_PARTITION_CAP_BACKLOG_DAG_IDS_LOGGED]
+ remaining_dag_id_count = len(sorted_dag_ids) -
len(displayed_dag_ids)
+ log_cap_reached = self.log.debug if
self._partition_cap_backlog_reported else self.log.warning
+ log_cap_reached(
+ (
+ "Reached the per-tick cap on pending partitioned Dag runs;
the remaining backlog "
+ "will be evaluated over subsequent scheduler ticks"
Review Comment:
This tells operators the backlog will be worked through, but the case this
warning matters most for is the opposite one: when the oldest `cap` pending
APDRs are waiting on upstream keys, a single scheduler fetches the same rows
every tick (the `continue` on an unsatisfied condition further down), and
nothing behind them is evaluated until those resolve. Could the log and the
audit `extra` say what actually happens, e.g. that only the oldest `cap`
pending partitioned Dag runs are evaluated per tick and newer ones wait behind
them? The audit text also says "evaluated" and "eligible", but it is written
before the evaluation loop runs and its count includes unsatisfied APDRs.
##########
airflow-core/tests/unit/jobs/test_scheduler_job.py:
##########
@@ -14345,6 +14415,617 @@ def
test_partition_cap_at_n_minus_one_leaves_one_pending(dag_maker: DagMaker, se
assert partition_dags == {"cap-consumer-n-minus-one"}
[email protected]_serialized_dag
[email protected]("clear_asset_partition_rows", "clear_audit_log_rows")
[email protected](
+ ("cap", "partition_keys", "expect_cap_hit"),
+ [
+ pytest.param(2, ["k1", "k2", "k3"], True, id="over-cap"),
+ pytest.param(3, ["k1", "k2"], False, id="under-cap"),
+ # Regression guard for the ``backlog_total > cap`` boundary:
``backlog_total == cap``
+ # must not be misreported as a backlog (that would be the "exactly
cap, nothing left"
+ # case wrongly reported as "more than cap, backlog remains").
+ pytest.param(3, ["k1", "k2", "k3"], False, id="exactly-cap"),
+ ],
+)
+def test_partition_cap_reporting(
+ dag_maker: DagMaker,
+ session: Session,
+ caplog,
+ cap: int,
+ partition_keys: list[str],
+ expect_cap_hit: bool,
+):
+ """Only a pending count strictly above the cap reports a backlog, via log
and audit `Log` row."""
+ suffix = f"{cap}-{len(partition_keys)}"
+ _make_n_satisfied_apdrs(
+ consumer_dag_id=f"cap-consumer-{suffix}",
+ asset=Asset(name=f"asset-cap-{suffix}"),
+ partition_keys=partition_keys,
+ session=session,
+ dag_maker=dag_maker,
+ )
+
+ runner = SchedulerJobRunner(
+ job=Job(job_type=SchedulerJobRunner.job_type),
executors=[MockExecutor(do_update=False)]
+ )
+ runner._max_partition_dag_runs_per_loop = cap
+ runner._create_dagruns_for_partitioned_asset_dags(session=session)
+
+ dag_id = f"cap-consumer-{suffix}"
+ audit_rows = session.scalars(select(Log).where(Log.event == "partition Dag
run cap reached")).all()
+ if expect_cap_hit:
+ assert {
+ "event": PARTITION_CAP_REACHED_EVENT,
+ "cap": cap,
+ "backlog_total": 3,
+ "dag_ids": [dag_id],
+ } in caplog
+ assert [row.event for row in audit_rows] == ["partition Dag run cap
reached"]
+ extra = audit_rows[0].extra
+ assert extra is not None
+ assert f"Affected dag_ids (highest pending count first): {dag_id}." in
extra
+ else:
+ assert PARTITION_CAP_REACHED_EVENT not in caplog
+ assert audit_rows == []
+
+
[email protected]_serialized_dag
[email protected]("clear_asset_partition_rows", "clear_audit_log_rows")
+def test_partition_cap_reporting_excludes_already_fired_decoy(dag_maker:
DagMaker, session: Session, caplog):
+ """
+ The ``backlog_total`` count query's ``WHERE`` clause must mirror the main
query's.
+
+ ``exactly-cap`` in :func:`test_partition_cap_reporting` cannot catch a
filter drift in the
+ count query (e.g. a dropped ``created_dag_run_id.is_(None)`` filter):
every row in that
+ table at that point is a genuine, unfired, non-stale APDR, so a looser
filter has nothing
+ extra to over-count. This test adds a decoy APDR that already fired
+ (``created_dag_run_id`` set) — a naive/unfiltered ``count()`` would
wrongly include it.
+ """
+ cap = 3
+ asset = Asset(name="asset-cap-decoy")
+ _make_n_satisfied_apdrs(
+ consumer_dag_id="cap-consumer-decoy",
+ asset=asset,
+ partition_keys=["k1", "k2", "k3"],
+ session=session,
+ dag_maker=dag_maker,
+ )
+ decoy = _produce_and_register_asset_event(
+ dag_id="asset-event-producer-decoy",
+ asset=asset,
+ partition_key="decoy",
+ session=session,
+ dag_maker=dag_maker,
+ )
+ fired_dag_run_id =
session.scalar(select(DagRun.id).order_by(DagRun.id.desc()).limit(1))
+ assert fired_dag_run_id is not None
+ decoy.created_dag_run_id = fired_dag_run_id
+ session.commit()
+
+ runner = SchedulerJobRunner(
+ job=Job(job_type=SchedulerJobRunner.job_type),
executors=[MockExecutor(do_update=False)]
+ )
+ runner._max_partition_dag_runs_per_loop = cap
+ runner._create_dagruns_for_partitioned_asset_dags(session=session)
+
+ assert PARTITION_CAP_REACHED_EVENT not in caplog
+ audit_rows = session.scalars(select(Log).where(Log.event == "partition Dag
run cap reached")).all()
+ assert audit_rows == []
+
+
[email protected]_serialized_dag
[email protected]("clear_asset_partition_rows", "clear_audit_log_rows")
[email protected]("attr", ["is_stale", "is_paused", "is_draining"])
+def test_partition_cap_reporting_excludes_inactive_dag_decoy(
+ attr: str, dag_maker: DagMaker, session: Session, caplog
+):
+ """
+ The ``backlog_total`` count query's ``WHERE`` clause must mirror the main
query's.
+
+ Mirrors :func:`test_partition_cap_reporting_excludes_already_fired_decoy`
for the fetch
+ query's Dag-state predicates: ``is_stale``, ``is_paused`` and
``is_draining`` must all be
+ ``False``. This test adds a decoy pending APDR whose target Dag became
inactive after the
+ APDR was written -- a count query missing that filter would wrongly
include it,
+ over-counting the backlog even though the fetch query never selects it.
+ """
+ cap = 3
+ suffix = attr.replace("_", "-")
+ _make_n_satisfied_apdrs(
+ consumer_dag_id=f"cap-consumer-inactive-decoy-{suffix}",
+ asset=Asset(name=f"asset-cap-inactive-decoy-{suffix}"),
+ partition_keys=["k1", "k2", "k3"],
+ session=session,
+ dag_maker=dag_maker,
+ )
+
+ inactive_consumer_dag_id = f"cap-consumer-inactive-decoy-{suffix}-inactive"
+ inactive_asset = Asset(name=f"asset-cap-inactive-decoy-{suffix}-inactive")
+ with dag_maker(
+ dag_id=inactive_consumer_dag_id,
+ schedule=PartitionedAssetTimetable(assets=inactive_asset,
default_partition_mapper=IdentityMapper()),
+ session=session,
+ ):
+ EmptyOperator(task_id="hi")
+ session.commit()
+
+ _produce_and_register_asset_event(
+ dag_id=f"asset-event-producer-inactive-decoy-{suffix}",
+ asset=inactive_asset,
+ partition_key="decoy",
+ session=session,
+ dag_maker=dag_maker,
+ )
+
+ # Flip the flag only after the APDR exists: AssetManager drops paused and
draining Dags
+ # before writing APDRs, so flag-first would produce no decoy at all.
+ dm = session.get(DagModel, inactive_consumer_dag_id)
+ assert dm is not None
+ setattr(dm, attr, True)
+ session.commit()
+
+ runner = SchedulerJobRunner(
+ job=Job(job_type=SchedulerJobRunner.job_type),
executors=[MockExecutor(do_update=False)]
+ )
+ runner._max_partition_dag_runs_per_loop = cap
+ runner._create_dagruns_for_partitioned_asset_dags(session=session)
+
+ assert PARTITION_CAP_REACHED_EVENT not in caplog
+ audit_rows = session.scalars(select(Log).where(Log.event == "partition Dag
run cap reached")).all()
+ assert audit_rows == []
+
+
[email protected]_serialized_dag
[email protected]("clear_asset_partition_rows", "clear_audit_log_rows")
+def test_partition_cap_reporting_lists_all_affected_dag_ids(dag_maker:
DagMaker, session: Session, caplog):
+ """
+ The log/audit ``dag_ids`` list surfaces every distinct dag_id in the
backlog, not just this
+ tick's oldest-cap slice.
+
+ Seven consumer Dags each get one satisfied APDR; with the cap at 6,
`pending_apdrs` (this
+ tick's oldest six) spans only six distinct dag_ids, but the seventh Dag's
APDR is still part
+ of the backlog `count()` sees — its dag_id must appear in both the log
event and the audit
+ row too, even though its Dag run won't be created until the next tick.
+ """
+ apdrs = _make_n_consumer_dags_each_with_one_pending_apdr(
+ name_prefix="many", n=7, session=session, dag_maker=dag_maker
+ )
+ base = timezone.utcnow()
+ for i, apdr in enumerate(apdrs):
+ apdr.created_at = base + timedelta(seconds=i)
+ session.commit()
+
+ runner = SchedulerJobRunner(
+ job=Job(job_type=SchedulerJobRunner.job_type),
executors=[MockExecutor(do_update=False)]
+ )
+ runner._max_partition_dag_runs_per_loop = 6
+ runner._create_dagruns_for_partitioned_asset_dags(session=session)
+
+ # The dag_ids list is sorted by (-count, dag_id); all seven Dags tie on
count (one pending
+ # APDR each), so the tie-break is dag_id ascending — a deterministic,
reproducible order,
+ # not just membership. All seven Dags are expected, including the seventh
whose APDR is
+ # deferred past this tick's cap.
+ expected_dag_ids = sorted(f"cap-consumer-many-{i:02d}" for i in range(1,
8))
+ assert {
+ "event": PARTITION_CAP_REACHED_EVENT,
+ "cap": 6,
+ "backlog_total": 7,
+ "dag_ids": expected_dag_ids,
+ } in caplog
+
+ audit_row = session.scalar(select(Log).where(Log.event == "partition Dag
run cap reached"))
+ assert audit_row is not None
+ assert audit_row.extra is not None
+ note_match = re.search(r"Affected dag_ids \(highest pending count first\):
(.+)\.$", audit_row.extra)
+ assert note_match is not None
+ assert note_match.group(1).split(", ") == expected_dag_ids
+ assert "and 1 more" not in audit_row.extra
+
+
[email protected]_serialized_dag
[email protected]("clear_asset_partition_rows", "clear_audit_log_rows")
+def test_partition_cap_reporting_truncates_dag_id_list_beyond_max(
+ dag_maker: DagMaker, session: Session, caplog
+):
+ """
+ The log/audit ``dag_ids`` list is capped at
``MAX_PARTITION_CAP_BACKLOG_DAG_IDS_LOGGED``,
+ ordered by descending per-dag pending count with a dag_id tie-break — not
alphabetically
+ truncated — and the audit row's prose notes how many dag_ids were dropped
by the cap.
+
+ Built via :func:`_make_n_consumer_dags_each_with_one_pending_apdr` rather
than
+ :func:`_make_n_satisfied_apdrs` because that helper's producer dag_id
numbering restarts at 1
+ on every call, colliding across these 21 consumer Dags — same issue
documented on
Review Comment:
This paragraph and the matching one on
`test_partition_cap_reporting_at_exactly_max_dag_ids_no_truncation` point at a
producer-collision explanation that now only lives on the helper's own
docstring. Since the helper documents it, both paragraphs can go.
--
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]
