This is an automated email from the ASF dual-hosted git repository.
ephraimbuddy pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/airflow.git
The following commit(s) were added to refs/heads/main by this push:
new 8e14c2960eb Cut the per-file cost of persisting Dag parse results
(#71621)
8e14c2960eb is described below
commit 8e14c2960ebdd3b5ccdf27a6e8ef17006a80124f
Author: Ephraim Anierobi <[email protected]>
AuthorDate: Wed Aug 26 22:56:56 2026 +0100
Cut the per-file cost of persisting Dag parse results (#71621)
* Cut the per-file cost of persisting Dag parse results
The Dag processor's manager persists parse results one file at a time, and
on
small deployments that process is the CPU bottleneck. Two costs were paid
there
for every parsed file.
Import errors were read by loading the whole import_error table, even for
files
that had none. The lookup is now scoped to the keys being written. The
table has
no index on (bundle_name, filename), so this saves materialising every row
rather than the scan itself.
The serialized Dag was hashed twice per write, once for change detection and
again in the model constructor, and hashing re-sorts and re-encodes the
whole
structure.
The compressed-storage test parametrization never ran, because the module
binds
that flag at import. It is made real here, since this change moves code into
that branch.
From @seanmuth's investigation of the Airflow 2-to-3 dag-processor CPU
regression: github.com/seanmuth/af2-af3-scheduler-cpu-repro. This is two of
the
smaller levers; the per-sweep result batching is left for follow-up.
* Say less around the bounded import-error lookup
Review feedback: the comments explained more than the code needed, and
falling back with `or` reads better than a None check for a value that is
always an md5 digest when it is set at all.
* Drop the per-call statement price the bounded lookup saves
The budget pinned in #71720 exists so a change to what persistence costs has
to move a number and account for it. This one does: a file that parses
cleanly no longer looks up import errors it has none of, so every call is a
statement cheaper.
* Pin what a file reporting an import error costs to persist
The budget covered only files that parse cleanly, which is the half of the
call this change makes cheaper. The half that still looks its errors up was
unpriced, so nothing would notice a round trip appearing there.
---
.../src/airflow/dag_processing/collection.py | 22 +++--
airflow-core/src/airflow/models/serialized_dag.py | 21 ++--
.../tests/unit/dag_processing/test_collection.py | 100 ++++++++++++++++++-
.../tests/unit/dag_processing/test_manager.py | 44 ++++++++-
.../tests/unit/models/test_serialized_dag.py | 106 +++++++++++++++++++++
5 files changed, 275 insertions(+), 18 deletions(-)
diff --git a/airflow-core/src/airflow/dag_processing/collection.py
b/airflow-core/src/airflow/dag_processing/collection.py
index 11e1099de7b..36b948abad6 100644
--- a/airflow-core/src/airflow/dag_processing/collection.py
+++ b/airflow-core/src/airflow/dag_processing/collection.py
@@ -382,16 +382,24 @@ def _update_dag_warnings(
def _update_import_errors(
files_parsed: set[tuple[str, str]],
- bundle_name: str,
import_errors: dict[tuple[str, str], str],
session: Session,
):
from airflow.listeners.listener import get_listener_manager
- # Check existing import errors BEFORE deleting, so we can determine if we
should update or create
- existing_import_error_files = set(
- session.execute(select(ParseImportError.bundle_name,
ParseImportError.filename))
- )
+ # Read before the delete below, so an error still present is updated
rather than recreated.
+ # Scoped to the keys being written: the table is not indexed on
(bundle_name, filename) so this
+ # still scans, but it no longer builds a row per import error the
deployment is carrying.
+ existing_import_error_files: set[tuple[str | None, str | None]] = set()
+ if import_errors:
+ existing_import_error_files = {
+ (bundle, filename)
+ for bundle, filename in session.execute(
+ select(ParseImportError.bundle_name,
ParseImportError.filename).where(
+ tuple_(ParseImportError.bundle_name,
ParseImportError.filename).in_(list(import_errors))
+ )
+ )
+ }
# Delete errors for files that were parsed but don't have errors in
import_errors
# (i.e., files that were successfully parsed without errors)
@@ -439,7 +447,7 @@ def _update_import_errors(
else:
import_error = ParseImportError(
filename=relative_fileloc,
- bundle_name=bundle_name,
+ bundle_name=bundle_name_,
timestamp=utcnow(),
stacktrace=stacktrace,
)
@@ -459,7 +467,6 @@ def _update_import_errors(
)
.values(
has_import_errors=True,
- bundle_name=bundle_name,
is_stale=True,
)
.execution_options(synchronize_session="fetch")
@@ -556,7 +563,6 @@ def update_dag_parsing_results_in_db(
try:
_update_import_errors(
files_parsed=files_parsed if files_parsed is not None else set(),
- bundle_name=bundle_name,
import_errors=import_errors,
session=session,
)
diff --git a/airflow-core/src/airflow/models/serialized_dag.py
b/airflow-core/src/airflow/models/serialized_dag.py
index 355c99f5291..d49c168013f 100644
--- a/airflow-core/src/airflow/models/serialized_dag.py
+++ b/airflow-core/src/airflow/models/serialized_dag.py
@@ -347,15 +347,20 @@ class SerializedDagModel(Base):
load_op_links = True
__table_args__ = (Index("idx_serialized_dag_dag_id_created_at", dag_id,
created_at),)
- def __init__(self, dag: LazyDeserializedDAG) -> None:
+ def __init__(self, dag: LazyDeserializedDAG, *, _dag_hash: str | None =
None) -> None:
+ """
+ Build a serialized Dag row.
+
+ :param _dag_hash: hash of ``dag.data``, when the caller has already
computed one. It has to
+ match the data as it stands, or the next parse misreads whether
the Dag changed.
+ """
self.dag_id = dag.dag_id
dag_data = dag.data
- self.dag_hash = SerializedDagModel.hash(dag_data)
-
- # partially ordered json data
- dag_data_json = json.dumps(dag_data, sort_keys=True).encode("utf-8")
+ self.dag_hash = _dag_hash or SerializedDagModel.hash(dag_data)
if _COMPRESS_SERIALIZED_DAGS:
+ # partially ordered json data
+ dag_data_json = json.dumps(dag_data,
sort_keys=True).encode("utf-8")
self._data = None
self._data_compressed = zlib.compress(dag_data_json)
else:
@@ -734,7 +739,7 @@ class SerializedDagModel(Base):
# This is for dynamic DAGs that the hashes changes often. We
should update
# the serialized dag, the dag_version and the dag_code instead of
a new version
# if the dag_version is not associated with any task instances
- new_serialized_dag = cls(dag)
+ new_serialized_dag = cls(dag, _dag_hash=new_dag_hash)
# Use direct UPDATE to avoid loading the full serialized DAG
result = session.execute(
@@ -788,8 +793,10 @@ class SerializedDagModel(Base):
if reused_deadline_data:
deadline_uuid_mapping = {str(uuid6.uuid7()): data for data in
reused_deadline_data.values()}
dag.data["dag"]["deadline"] = list(deadline_uuid_mapping.keys())
+ # The data just changed, so the hash computed above no longer
describes it.
+ new_dag_hash = cls.hash(dag.data)
- new_serialized_dag = cls(dag)
+ new_serialized_dag = cls(dag, _dag_hash=new_dag_hash)
new_serialized_dag.dag_version = dagv
session.add(new_serialized_dag)
diff --git a/airflow-core/tests/unit/dag_processing/test_collection.py
b/airflow-core/tests/unit/dag_processing/test_collection.py
index 8c21119305a..be36290abce 100644
--- a/airflow-core/tests/unit/dag_processing/test_collection.py
+++ b/airflow-core/tests/unit/dag_processing/test_collection.py
@@ -27,7 +27,7 @@ from unittest import mock
from unittest.mock import patch
import pytest
-from sqlalchemy import delete, func, inspect as sa_inspect, select
+from sqlalchemy import delete, event, func, inspect as sa_inspect, select
from sqlalchemy.exc import OperationalError, SAWarning
import airflow.dag_processing.collection
@@ -39,6 +39,7 @@ from airflow.dag_processing.collection import (
_get_latest_runs_stmt,
_get_latest_runs_stmt_partitioned,
_update_dag_tags,
+ _update_import_errors,
update_dag_parsing_results_in_db,
)
from airflow.exceptions import SerializationError
@@ -1467,6 +1468,103 @@ class TestUpdateDagParsingResults:
assert orm_dag.max_consecutive_failed_dag_runs == 6
[email protected]_test
+class TestUpdateImportErrors:
+ """Tests for the ``_update_import_errors`` helper."""
+
+ @pytest.fixture(autouse=True)
+ def clean_import_errors(self):
+ clear_db_import_errors()
+ yield
+ clear_db_import_errors()
+
+ @pytest.fixture
+ def import_error_statements(self, session):
+ """
+ Collect every SQL statement issued against the ``import_error`` table.
+
+ Matching on the bare table name would also catch statements naming
``dag.has_import_errors``,
+ so match the positions where the table itself can appear.
+ """
+ statements: list[str] = []
+
+ def _capture(conn, cursor, statement, parameters, context,
executemany):
+ lowered = statement.lower()
+ if any(f"{keyword} import_error" in lowered for keyword in
("from", "into", "update")):
+ statements.append(statement)
+
+ bind = session.get_bind()
+ event.listen(bind, "before_cursor_execute", _capture)
+ yield statements
+ event.remove(bind, "before_cursor_execute", _capture)
+
+ @staticmethod
+ def _selects(statements: list[str]) -> list[str]:
+ return [stmt for stmt in statements if
stmt.lower().lstrip().startswith("select")]
+
+ def test_no_lookup_when_there_are_no_import_errors(self, session,
import_error_statements):
+ session.add(ParseImportError(filename="broken.py",
bundle_name="testing", stacktrace="boom"))
+ session.flush()
+ import_error_statements.clear()
+
+ # files_parsed is empty so no DELETE runs either: on backends without
DELETE...RETURNING
+ # its synchronize_session fallback would emit a SELECT of its own and
muddy the assertion.
+ _update_import_errors(
+ files_parsed=set(),
+ import_errors={},
+ session=session,
+ )
+
+ assert self._selects(import_error_statements) == []
+
+ @patch.object(ParseImportError, "full_file_path", return_value="broken.py")
+ def test_existing_error_lookup_is_bounded(self, _mock_full_path, session,
import_error_statements):
+ session.add_all(
+ [
+ ParseImportError(filename="broken.py", bundle_name="testing",
stacktrace="old"),
+ ParseImportError(filename="untouched.py", bundle_name="other",
stacktrace="unrelated"),
+ ]
+ )
+ session.flush()
+ import_error_statements.clear()
+
+ _update_import_errors(
+ files_parsed={("testing", "broken.py")},
+ import_errors={("testing", "broken.py"): "new"},
+ session=session,
+ )
+
+ selects = self._selects(import_error_statements)
+ assert selects, "expected the existing-error lookup to run"
+ assert all("where" in stmt.lower() for stmt in selects), (
+ f"import_error must never be scanned unfiltered, got: {selects}"
+ )
+
+ rows = sorted(
+ (err.bundle_name, err.filename, err.stacktrace)
+ for err in session.scalars(select(ParseImportError))
+ )
+ assert rows == [
+ ("other", "untouched.py", "unrelated"),
+ ("testing", "broken.py", "new"),
+ ]
+
+ @patch.object(ParseImportError, "full_file_path", return_value="broken.py")
+ def test_new_errors_keep_their_own_bundle_name(self, _mock_full_path,
session):
+ _update_import_errors(
+ files_parsed=set(),
+ import_errors={
+ ("bundle_a", "a.py"): "error a",
+ ("bundle_b", "b.py"): "error b",
+ },
+ session=session,
+ )
+ session.flush()
+
+ rows = {(err.bundle_name, err.filename) for err in
session.scalars(select(ParseImportError))}
+ assert rows == {("bundle_a", "a.py"), ("bundle_b", "b.py")}
+
+
@pytest.mark.db_test
class TestUpdateDagTags:
@pytest.fixture(autouse=True)
diff --git a/airflow-core/tests/unit/dag_processing/test_manager.py
b/airflow-core/tests/unit/dag_processing/test_manager.py
index 88d835169b0..0c6fdbd1deb 100644
--- a/airflow-core/tests/unit/dag_processing/test_manager.py
+++ b/airflow-core/tests/unit/dag_processing/test_manager.py
@@ -261,10 +261,15 @@ def _statement_breakdown(counts: Counter[tuple[str,
str]]) -> str:
# Per persistence call, and per Dag in the file. A call leaves the serialized
Dag alone while the
# content is unchanged; once the hash has moved and [core]
min_serialized_dag_update_interval has
-# lapsed it rewrites it, which costs two more statements per Dag and nothing
extra per call.
-FIXED_PER_CALL = 10
+# lapsed it rewrites it, which costs two more statements per Dag and nothing
extra per call. The
+# per-call price is a file that parsed cleanly: one reporting import errors
also looks up whichever
+# of them are already recorded.
+FIXED_PER_CALL = 9
UNCHANGED_PER_DAG = 3
REWRITE_PER_DAG = 5
+# A file that failed to parse and so defines no Dags. Two of the five are
import_error SELECTs: the
+# bounded lookup, and the listener re-reading the row the update beside it
already had.
+IMPORT_ERROR_PER_CALL = 5
SWEEP_FILES = 4
# Calls the manager takes for that sweep: one per file today, 1 if a sweep is
ever batched.
@@ -3607,6 +3612,24 @@ class TestDagFileProcessorManager:
manager._collect_results()
return sum(counts.values())
+ @staticmethod
+ def _measure_import_error_call(session, rel_path: str) ->
Counter[tuple[str, str]]:
+ """Count one steady-state call for a file that fails to parse with its
error already recorded."""
+ files_parsed = {("testing", rel_path)}
+ recorded = {("testing", rel_path): "boom"}
+ update_dag_parsing_results_in_db(
+ "testing", None, [], recorded, 0.1, set(), session,
files_parsed=files_parsed
+ )
+ session.commit()
+
+ again = {("testing", rel_path): "boom again"}
+ with _count_statements(session) as counts:
+ update_dag_parsing_results_in_db(
+ "testing", None, [], again, 0.1, set(), session,
files_parsed=files_parsed
+ )
+ session.flush()
+ return counts
+
@pytest.mark.backend("postgres")
@pytest.mark.parametrize("n_dags", [1, 5])
@pytest.mark.parametrize(
@@ -3639,6 +3662,23 @@ class TestDagFileProcessorManager:
f"({FIXED_PER_CALL} per call + {per_dag} per
Dag).\n{_statement_breakdown(counts)}"
)
+ @pytest.mark.backend("postgres")
+ def test_a_file_reporting_an_import_error_stays_within_its_budget(
+ self, session, testing_dag_bundle, tmp_path
+ ):
+ """
+ The path a clean parse skips: a file that looks up the errors already
recorded for it.
+
+ Not comparable to ``FIXED_PER_CALL``, which is priced with Dags to
write; this file has none.
+ """
+ counts = self._measure_import_error_call(session, "broken.py")
+
+ total = sum(counts.values())
+ assert total == IMPORT_ERROR_PER_CALL, (
+ f"a file reporting one import error costs {total}, expected
{IMPORT_ERROR_PER_CALL}."
+ f"\n{_statement_breakdown(counts)}"
+ )
+
def test_a_sweep_pays_the_fixed_cost_once_per_call(self, session,
testing_dag_bundle, tmp_path):
"""
How a sweep scales with the number of persistence calls it takes.
diff --git a/airflow-core/tests/unit/models/test_serialized_dag.py
b/airflow-core/tests/unit/models/test_serialized_dag.py
index f3d720ddcc6..7609c57ba01 100644
--- a/airflow-core/tests/unit/models/test_serialized_dag.py
+++ b/airflow-core/tests/unit/models/test_serialized_dag.py
@@ -29,6 +29,7 @@ import pytest
from sqlalchemy import delete, func, select, update
import airflow.example_dags as example_dags_module
+import airflow.models.serialized_dag as serialized_dag_module
from airflow._shared.observability.metrics.base_stats_logger import StatsLogger
from airflow.dag_processing.dagbag import DagBag
from airflow.models.asset import AssetActive, AssetAliasModel, AssetModel
@@ -106,6 +107,9 @@ class TestSerializedDagModel:
db.clear_db_dags()
db.clear_db_runs()
db.clear_db_serialized_dags()
+ # The module reads the config once at import time, so conf_vars alone
leaves the already
+ # bound constant untouched and the compressed branch never runs.
+ monkeypatch.setattr(serialized_dag_module,
"_COMPRESS_SERIALIZED_DAGS", request.param)
with conf_vars({("core", "compress_serialized_dags"):
str(request.param)}):
yield
db.clear_db_serialized_dags()
@@ -1172,6 +1176,108 @@ class TestSerializedDagModel:
old_alert = session.scalar(select(DAM).where(DAM.serialized_dag_id ==
orig_serdag.id))
assert old_alert is not None
+ def test_write_dag_hashes_the_dag_once_when_inserting(self,
testing_dag_bundle, session):
+ """The hash computed for change detection is reused for the new row
rather than recomputed."""
+ dag_id = "test_hash_once_insert"
+ dag = DAG(dag_id=dag_id)
+ EmptyOperator(task_id="task1", dag=dag)
+ scheduler_dag = sync_dag_to_db(dag, session=session)
+
+ # Task instances on the current version force write_dag to insert a
new row.
+ scheduler_dag.create_dagrun(
+ run_id="test1",
+ run_after=DEFAULT_DATE,
+ state=DagRunState.QUEUED,
+ logical_date=DEFAULT_DATE,
+ data_interval=(DEFAULT_DATE, DEFAULT_DATE),
+ triggered_by=DagRunTriggeredByType.TEST,
+ run_type=DagRunType.MANUAL,
+ )
+ session.commit()
+
+ EmptyOperator(task_id="task2", dag=dag)
+ with mock.patch.object(SDM, "hash", wraps=SDM.hash) as hash_spy:
+ SDM.write_dag(LazyDeserializedDAG.from_dag(dag),
bundle_name="testing", session=session)
+ session.commit()
+
+ assert hash_spy.call_count == 1
+ assert session.scalar(select(func.count()).where(SDM.dag_id ==
dag_id)) == 2
+ # Expunge so ``data`` is read back from the row rather than the cache
``__init__`` populated.
+ session.expunge_all()
+ serdag = session.scalar(select(SDM).where(SDM.dag_id ==
dag_id).order_by(SDM.created_at.desc()))
+ assert serdag.dag_hash == SDM.hash(serdag.data)
+
+ def test_write_dag_hashes_the_dag_once_when_updating_in_place(self,
testing_dag_bundle, session):
+ """The in-place UPDATE branch reuses the hash too."""
+ dag_id = "test_hash_once_update"
+ dag = DAG(dag_id=dag_id)
+ EmptyOperator(task_id="task1", dag=dag)
+ sync_dag_to_db(dag, session=session)
+ session.commit()
+ original_hash = session.scalar(select(SDM.dag_hash).where(SDM.dag_id
== dag_id))
+
+ # No task instances exist for the current version, so a changed Dag
updates the row in place.
+ EmptyOperator(task_id="task2", dag=dag)
+ with mock.patch.object(SDM, "hash", wraps=SDM.hash) as hash_spy:
+ SDM.write_dag(LazyDeserializedDAG.from_dag(dag),
bundle_name="testing", session=session)
+ session.commit()
+
+ assert hash_spy.call_count == 1
+ assert session.scalar(select(func.count()).where(SDM.dag_id ==
dag_id)) == 1
+ # Expunge so ``data`` is read back from the row rather than the cache
``__init__`` populated.
+ session.expunge_all()
+ serdag = session.scalar(select(SDM).where(SDM.dag_id == dag_id))
+ # Without this the test would also pass if write_dag had done nothing
at all: one hash is
+ # computed for change detection either way and an unwritten row is
trivially consistent.
+ assert serdag.dag_hash != original_hash
+ assert serdag.dag_hash == SDM.hash(serdag.data)
+
+ def test_deadline_uuid_regeneration_keeps_hash_consistent(self,
testing_dag_bundle, session):
+ """Regenerated deadline UUIDs rewrite dag.data, so the stored hash
must describe the rewrite."""
+ dag_id = "test_hash_deadline_regenerated"
+
+ dag = DAG(
+ dag_id=dag_id,
+ deadline=DeadlineAlert(
+ reference=DeadlineReference.DAGRUN_QUEUED_AT,
+ interval=timedelta(minutes=5),
+ callback=AsyncCallback(empty_callback_for_deadline),
+ ),
+ )
+ EmptyOperator(task_id="task1", dag=dag)
+ scheduler_dag = sync_dag_to_db(dag, session=session)
+
+ # Task instances on the current version force write_dag down the
INSERT branch, where the
+ # reused deadline UUIDs are regenerated after the change-detection
hash was computed.
+ scheduler_dag.create_dagrun(
+ run_id="test1",
+ run_after=DEFAULT_DATE,
+ state=DagRunState.QUEUED,
+ logical_date=DEFAULT_DATE,
+ data_interval=(DEFAULT_DATE, DEFAULT_DATE),
+ triggered_by=DagRunTriggeredByType.TEST,
+ run_type=DagRunType.MANUAL,
+ )
+ session.commit()
+
+ session.expunge_all()
+ original = session.scalar(select(SDM).where(SDM.dag_id ==
dag_id).order_by(SDM.created_at.desc()))
+ original_deadlines = list(original.data["dag"]["deadline"])
+
+ # Non-deadline edit, so the deadline definitions still match and their
UUIDs are reused.
+ EmptyOperator(task_id="task2", dag=dag)
+ SDM.write_dag(LazyDeserializedDAG.from_dag(dag),
bundle_name="testing", session=session)
+ session.commit()
+
+ # Expunge so ``data`` is read back from the row rather than the cache
``__init__`` populated.
+ session.expunge_all()
+ serdag = session.scalar(select(SDM).where(SDM.dag_id ==
dag_id).order_by(SDM.created_at.desc()))
+ # Proves the reuse branch actually ran: it swaps the reused UUIDs for
fresh ones after the
+ # change-detection hash was taken. Without this the hash assertion
holds vacuously when
+ # write_dag does nothing.
+ assert serdag.data["dag"]["deadline"] != original_deadlines
+ assert serdag.dag_hash == SDM.hash(serdag.data)
+
def test_non_deadline_edit_preserves_alert_in_update_branch(self,
testing_dag_bundle, session):
"""UPDATE branch (no task instances): existing deadline_alert stays
linked after non-deadline edit."""
dag_id = "test_deadline_update_branch"
