kaxil commented on code in PR #75010:
URL: https://github.com/apache/airflow/pull/75010#discussion_r4323364825
##########
airflow-core/src/airflow/jobs/scheduler_job_runner.py:
##########
@@ -3994,6 +3994,11 @@ def _find_task_instances_without_heartbeats(self, *,
session: Session) -> list[T
session.execute(with_row_locks(run_query, of=DagRun,
session=session, skip_locked=True))
)
query = query.where(or_(~regional_expression(TI), tuple_(TI.dag_id,
TI.run_id).in_(locked_runs)))
+ # Bound the batch (oldest heartbeat first) so a mass heartbeat-loss
event does not purge every
+ # stale row, each with per-TI work under FOR UPDATE, in one
scheduler-stalling transaction.
+ query = query.order_by(TI.last_heartbeat_at)
+ if self.job.max_tis_per_query:
+ query = query.limit(self.job.max_tis_per_query)
Review Comment:
With the defaults this drains 16 TIs every 10s, since the purge runs on the
`task_instance_heartbeat_timeout_detection_interval` timer rather than once per
scheduler loop. Until a TI is purged it stays in `executor.running` (so it
counts against `slots_occupied`) and keeps its pool slot, so a deployment at
`parallelism = 1024` that loses a worker node with ~1000 running tasks now
needs about 63 passes (~10 minutes) to get that capacity back, where it used to
take one. Also, `handle_failure()` already commits per TI through
`save_to_db()`, so the scan's FOR UPDATE locks are released after the first TI
either way; what the limit mostly bounds is how long one timer call holds the
loop. Could a full batch schedule a follow-up pass on the next loop instead of
waiting the full interval, or use a larger dedicated cap like the `.limit(100)`
in the AWAITING_INPUT sweep?
##########
airflow-core/tests/unit/jobs/test_scheduler_job.py:
##########
@@ -918,6 +918,30 @@ def
test_heartbeat_scan_waits_for_a_locked_dagrun_only_for_regional_task_instanc
assert [row.id for row in
runner._find_task_instances_without_heartbeats(session=session)] == [ti.id]
+ def
test_find_task_instances_without_heartbeats_is_bounded_by_max_tis_per_query(self,
dag_maker, session):
+ """A mass heartbeat-loss event must be purged in bounded batches,
oldest heartbeat first."""
+ with dag_maker(serialized=True):
+ EmptyOperator(task_id="task")
+ runner = SchedulerJobRunner(Job(),
executors=[MockExecutor(do_update=False)])
+ session.add(runner.job)
+ session.flush()
+ runner.job.max_tis_per_query = 2
+
+ ids_oldest_first = []
+ for i in range(3):
+ dr = dag_maker.create_dagrun(run_id=f"run_{i}",
logical_date=DEFAULT_DATE + timedelta(days=i))
+ ti = dr.get_task_instance("task", session=session)
+ ti.state = State.RUNNING
+ ti.region_index = -1
+ ti.last_heartbeat_at = DEFAULT_DATE + timedelta(minutes=i)
Review Comment:
Heartbeat order matches insertion order here (and the uuid7 `id` order), so
this still passes with the `order_by` removed. Reversing it, e.g. `DEFAULT_DATE
- timedelta(minutes=i)` with the expected list built in heartbeat order, would
make the test pin oldest-first. A `max_tis_per_query = 0` case would cover the
unbounded branch too.
##########
airflow-core/src/airflow/jobs/scheduler_job_runner.py:
##########
@@ -3994,6 +3994,11 @@ def _find_task_instances_without_heartbeats(self, *,
session: Session) -> list[T
session.execute(with_row_locks(run_query, of=DagRun,
session=session, skip_locked=True))
)
query = query.where(or_(~regional_expression(TI), tuple_(TI.dag_id,
TI.run_id).in_(locked_runs)))
+ # Bound the batch (oldest heartbeat first) so a mass heartbeat-loss
event does not purge every
+ # stale row, each with per-TI work under FOR UPDATE, in one
scheduler-stalling transaction.
+ query = query.order_by(TI.last_heartbeat_at)
+ if self.job.max_tis_per_query:
Review Comment:
The config description says `0` means "use `[core] parallelism`", and
`_executable_task_instances_to_queued` (line 1306) does that; only the
`schedule_tis` chunking treats `0` as unbounded. Here `0` removes the limit
entirely, which is the one setting where this PR gives no protection. The
`max_tis_per_query` description in config.yml also still only mentions
scheduling. Could we map `0` to `self._parallelism` here, or at least document
the purge use and what `0` means for it?
##########
airflow-core/src/airflow/jobs/scheduler_job_runner.py:
##########
@@ -3994,6 +3994,11 @@ def _find_task_instances_without_heartbeats(self, *,
session: Session) -> list[T
session.execute(with_row_locks(run_query, of=DagRun,
session=session, skip_locked=True))
)
query = query.where(or_(~regional_expression(TI), tuple_(TI.dag_id,
TI.run_id).in_(locked_runs)))
+ # Bound the batch (oldest heartbeat first) so a mass heartbeat-loss
event does not purge every
+ # stale row, each with per-TI work under FOR UPDATE, in one
scheduler-stalling transaction.
+ query = query.order_by(TI.last_heartbeat_at)
Review Comment:
Ordering oldest-first gives the `continue` at line 4094 a new consequence. A
TI that hits it (no `dag_version_id` and no `DagVersion` for its Dag, e.g. an
adopted Airflow 2 TI whose Dag was never reparsed) stays RUNNING with its old
heartbeat, so it sorts to the front of every batch. With 16 of those, this
scheduler never purges anything else. Should that branch fail the TI without
the callback, the way a RESTARTING TI is already released?
--
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]
