uranusjr commented on code in PR #71072:
URL: https://github.com/apache/airflow/pull/71072#discussion_r4439020211
##########
airflow-core/src/airflow/jobs/scheduler_job_runner.py:
##########
@@ -2509,8 +2535,106 @@ def _create_dagruns_for_partitioned_asset_dags(self,
session: Session) -> set[st
)
).all()
if not pending_apdrs:
+ # An empty fetch is as ambiguous as a short one: other schedulers
may hold every
+ # row lock. Only an unlocked count showing the backlog at or below
the cap ends the
+ # episode, otherwise the next full fetch would re-warn and
re-write the audit row.
+ if self._partition_cap_backlog_reported:
+ pending_total = session.scalar(
+ select(func.count())
+ .select_from(AssetPartitionDagRun)
+ .join(DagModel, DagModel.dag_id ==
AssetPartitionDagRun.target_dag_id)
+ .where(*_build_pending_partitioned_apdr_filters())
+ )
+ if (pending_total or 0) <=
self._max_partition_dag_runs_per_loop:
+ self._partition_cap_backlog_reported = False
return set()
+ sorted_dag_ids: list[str] = []
+ backlogs_per_dag: dict[str, int] = {}
+ 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 (the empty-fetch branch above
follows the
+ # same rule).
+ # 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 = {
+ 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]
+ 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;
only the oldest "
+ "pending runs up to the cap are evaluated per tick, newer
ones wait behind them"
+ ),
+ 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
+ # an unlocked 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.
+ try:
+ with create_session(scoped=False) as audit_session:
+ audit_session.add_all(
+ Log(
+ event="partition Dag run cap reached",
+ dag_id=dag_id,
+ owner="scheduler",
+ owner_display_name="Scheduler",
+ extra=(
+ f"The scheduler fetched the oldest
{len(pending_apdrs)} pending "
+ f"partitioned Dag runs this tick, reaching
the internal per-tick cap "
+ f"of
{self._max_partition_dag_runs_per_loop}; these are not "
+ f"necessarily satisfied. This Dag has
{backlogs_per_dag[dag_id]} "
+ f"pending partitioned Dag runs. Only the
oldest pending runs up to "
+ f"the cap are evaluated per tick, and
newer ones wait behind them. "
+ f"A total of {backlog_total} partitioned
Dag runs are currently "
+ f"pending across all Dags."
+ ),
+ )
+ for dag_id in displayed_dag_ids
+ )
+ except SQLAlchemyError:
+ # Purely observational — must not fail or retry the tick's
actual DagRun
+ # creation work. Leave the flag False so the next tick
retries the write.
+ self.log.warning("Failed to write the partition Dag run
cap audit Log row", exc_info=True)
+ else:
+ self._partition_cap_backlog_reported = True
Review Comment:
These rows are now per Dag, but the edge trigger is still global. A Dag that
joins the backlog after the first tick of an episode never gets a row for the
whole episode. Is this intended? This would potentially be hours for a
steady-state backlog. Its audit page then shows nothing while its runs are held
behind the cap, which is the situation this change is meant to surface.
--
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]
