This is an automated email from the ASF dual-hosted git repository.
jason810496 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 402ab11a4b0 Avoid loading timetables for dags in Execution API
endpoints (#75765)
402ab11a4b0 is described below
commit 402ab11a4b08fa4027e5b88dbf9022d69809a4b3
Author: Ash Berlin-Taylor <[email protected]>
AuthorDate: Wed Oct 28 08:02:48 2026 +0000
Avoid loading timetables for dags in Execution API endpoints (#75765)
* Avoid loading timetables for dags in Execution API endpoints
For a few reasons/places now in the Exec API we now need the serialized
dag, but we don't need the timetable (which could involve loading) plugins
or user code.
This should let the ExecAPI run without user code if a deployment admin
wishes to try that.
* Stop eagerly decoding custom weight_rules
`weight_rule` on SerializedDag is already a property that decodes/inflates
strings into the object
* Fix reviews
* fixup! Avoid loading timetables for dags in Execution API endpoints
---------
Co-authored-by: Jason(Zhe-You) Liu
<[email protected]>
---
.../src/airflow/jobs/scheduler_job_runner.py | 9 ++++++-
.../src/airflow/serialization/definitions/dag.py | 29 ++++++++++++++++++++--
.../airflow/serialization/serialized_objects.py | 28 +++++----------------
airflow-core/tests/unit/models/test_dag.py | 6 ++---
.../unit/serialization/test_dag_serialization.py | 25 ++++++++++++++++---
5 files changed, 65 insertions(+), 32 deletions(-)
diff --git a/airflow-core/src/airflow/jobs/scheduler_job_runner.py
b/airflow-core/src/airflow/jobs/scheduler_job_runner.py
index a8ccf1d8145..f26ef79e393 100644
--- a/airflow-core/src/airflow/jobs/scheduler_job_runner.py
+++ b/airflow-core/src/airflow/jobs/scheduler_job_runner.py
@@ -415,7 +415,14 @@ class SchedulerJobRunner(BaseJobRunner, LoggingMixin):
if not serdag:
return None
serdag.load_op_links = False
- return serdag.dag
+ dag = serdag.dag
+
+ # Vivify lazily loaded things now
+ _ = dag.timetable
+ for task in dag.tasks:
+ _ = task.weight_rule
+
+ return dag
except Exception:
self.log.exception("Failed to deserialize DAG '%s'", dag_id)
return None
diff --git a/airflow-core/src/airflow/serialization/definitions/dag.py
b/airflow-core/src/airflow/serialization/definitions/dag.py
index 64ceeff7ccc..5b13823fcc2 100644
--- a/airflow-core/src/airflow/serialization/definitions/dag.py
+++ b/airflow-core/src/airflow/serialization/definitions/dag.py
@@ -37,6 +37,7 @@ from airflow.exceptions import (
AirflowException,
DagNotPartitionedError,
DagVersionNotFound,
+ DeserializationError,
InvalidPartitionKeyError,
NodeNotFound,
TaskNotFound,
@@ -50,9 +51,14 @@ from airflow.models.deadline import Deadline
from airflow.models.deadline_alert import DeadlineAlert as DeadlineAlertModel
from airflow.models.task_coordinates import public_map_index_expression
from airflow.models.tasklog import LogTemplate
-from airflow.serialization.decoders import decode_deadline_alert,
resolve_deadline_alert_interval
+from airflow.serialization.decoders import (
+ decode_deadline_alert,
+ decode_timetable,
+ resolve_deadline_alert_interval,
+)
from airflow.serialization.definitions.deadline import DeadlineAlertFields,
SerializedReferenceModels
from airflow.serialization.definitions.param import SerializedParamsDict
+from airflow.serialization.helpers import TimetableNotRegistered
from airflow.timetables.base import DagRunInfo, DataInterval, TimeRestriction
from airflow.utils.helpers import prune_dict
from airflow.utils.session import NEW_SESSION, provide_session
@@ -142,7 +148,9 @@ class SerializedDAG:
# These are set dynamically during deserialization.
task_dict: dict[str, SerializedOperator] = attrs.field(init=False)
task_group: SerializedTaskGroup = attrs.field(init=False)
- timetable: Timetable = attrs.field(init=False)
+ _timetable: Timetable = attrs.field(init=False)
+ # Decoded on first access, so loading a Dag does not need the timetable's
plugin.
+ _encoded_timetable: dict[str, Any] | None = attrs.field(init=False,
default=None)
timezone: FixedTimezone | Timezone = attrs.field(init=False)
# Only on serialized dag.
@@ -160,6 +168,23 @@ class SerializedDAG:
def __repr__(self) -> str:
return f"<SerializedDAG: {self.dag_id}>"
+ @property
+ def timetable(self) -> Timetable:
+ if (encoded := self._encoded_timetable) is not None:
+ try:
+ self._timetable = decode_timetable(encoded)
+ self._encoded_timetable = None
+ except TimetableNotRegistered:
+ raise
+ except Exception as err:
+ raise DeserializationError(self.dag_id) from err
+ return self._timetable
+
+ @timetable.setter
+ def timetable(self, value: Timetable) -> None:
+ self._timetable = value
+ self._encoded_timetable = None
+
@classmethod
def get_serialized_fields(cls) -> frozenset[str]:
return frozenset(
diff --git a/airflow-core/src/airflow/serialization/serialized_objects.py
b/airflow-core/src/airflow/serialization/serialized_objects.py
index 73fc26e0114..a6ceadcf1e8 100644
--- a/airflow-core/src/airflow/serialization/serialized_objects.py
+++ b/airflow-core/src/airflow/serialization/serialized_objects.py
@@ -238,19 +238,6 @@ def _encode_priority_weight_strategy(var:
PriorityWeightStrategy | str) -> str:
return importable_string
-def _decode_priority_weight_strategy(var: str) -> PriorityWeightStrategy:
- """
- Decode a previously serialized priority weight strategy.
-
- In this version, we only store the importable string, so we just need to
get the class
- from the dictionary of registered classes and instantiate it with no
parameters.
- """
- priority_weight_strategy_class =
_get_registered_priority_weight_strategy(var)
- if priority_weight_strategy_class is None:
- raise _PriorityWeightStrategyNotRegistered(var)
- return priority_weight_strategy_class()
-
-
# Builtin exceptions a BASE_EXC_SER node can rebuild into. A user defined
subclass
# (e.g. a custom KeyError) serializes to a name absent here and won't
round-trip --
# unchanged from before, since builtins never held it either.
@@ -1223,8 +1210,8 @@ class OperatorSerialization(DAGNode, BaseSerialization):
elif k == "_on_failure_fail_dagrun":
k = "on_failure_fail_dagrun"
elif k == "weight_rule":
+ # Decoded by SerializedBaseOperator.weight_rule, so loading a
Dag does not need the plugin.
k = "_weight_rule"
- v = _decode_priority_weight_strategy(v)
elif k == "retry_exponential_backoff":
if isinstance(v, bool):
v = 2.0 if v else 0
@@ -1877,9 +1864,7 @@ class DagSerialization(BaseSerialization):
# Value structure matches exactly
pass
elif k == "timetable":
- v = decode_timetable(v)
- elif k == "weight_rule":
- v = _decode_priority_weight_strategy(v)
+ k = "_encoded_timetable"
elif k in cls._decorated_fields:
v = cls.deserialize(v)
elif k == "params":
@@ -2211,22 +2196,21 @@ class DagSerialization(BaseSerialization):
)
try:
deserialized = cls.from_dict(copy.deepcopy(serialized_obj))
+ timetable = deserialized.timetable
except Exception as e:
cause = e.__cause__ if isinstance(e, DeserializationError) and
e.__cause__ else e
raise DeserializationError(
dag_id, f"Dag {dag_id!r} cannot be deserialized:
{type(cause).__name__}: {cause}"
) from e
- if (
- limit := deserialized.timetable.active_runs_limit
- ) is not None and deserialized.max_active_runs > limit:
+ if (limit := timetable.active_runs_limit) is not None and
deserialized.max_active_runs > limit:
raise DeserializationError(
dag_id,
f"Dag {dag_id!r} sets max_active_runs
{deserialized.max_active_runs}, "
- f"but {type(deserialized.timetable).__name__} allows at most
{limit}",
+ f"but {type(timetable).__name__} allows at most {limit}",
)
if (
deserialized.catchup
- and deserialized.timetable.can_be_scheduled
+ and timetable.can_be_scheduled
and not (deserialized.start_date or "start_date" in
deserialized.default_args)
):
raise DeserializationError(dag_id, f"Dag {dag_id!r} sets catchup
but no start_date")
diff --git a/airflow-core/tests/unit/models/test_dag.py
b/airflow-core/tests/unit/models/test_dag.py
index 949b2d0d82b..54b3571bd5f 100644
--- a/airflow-core/tests/unit/models/test_dag.py
+++ b/airflow-core/tests/unit/models/test_dag.py
@@ -3992,12 +3992,10 @@ def test_iter_dagrun_infos_between_error(caplog):
"airflow.serialization.encoders.find_registered_custom_timetable",
_get_registered_timetable,
),
+ caplog.at_level(logging.ERROR),
):
scheduler_dag = create_scheduler_dag(dag)
-
- iterator = scheduler_dag.iter_dagrun_infos_between(earliest=start,
latest=end)
- with caplog.at_level(logging.ERROR):
- infos = list(iterator)
+ infos = list(scheduler_dag.iter_dagrun_infos_between(earliest=start,
latest=end))
# The second timetable.next_dagrun_info() call raises an exception, so
only the first result is returned.
assert infos == [DagRunInfo.interval(start, end)]
diff --git a/airflow-core/tests/unit/serialization/test_dag_serialization.py
b/airflow-core/tests/unit/serialization/test_dag_serialization.py
index 0caec6fa2ba..e6c12262321 100644
--- a/airflow-core/tests/unit/serialization/test_dag_serialization.py
+++ b/airflow-core/tests/unit/serialization/test_dag_serialization.py
@@ -93,6 +93,7 @@ from airflow.serialization.definitions.param import
SerializedParam
from airflow.serialization.definitions.xcom_arg import SchedulerPlainXComArg
from airflow.serialization.encoders import ensure_serialized_asset
from airflow.serialization.enums import DagAttributeTypes, Encoding
+from airflow.serialization.helpers import TimetableNotRegistered
from airflow.serialization.json_schema import load_dag_schema_dict
from airflow.serialization.serialized_objects import (
BaseSerialization,
@@ -1215,8 +1216,11 @@ class TestStringifiedDAGs:
"you have a top level database access that disrupted the session. "
"Please check the airflow best practices documentation."
)
- with pytest.raises(ValueError, match=message):
- DagSerialization.from_dict(serialized)
+ dag = DagSerialization.from_dict(serialized)
+
+ assert dag.dag_id == "simple_dag"
+ with pytest.raises(TimetableNotRegistered, match=message):
+ dag.timetable
@pytest.mark.parametrize(
("val", "expected"),
@@ -5236,14 +5240,29 @@ class TestWeightRule:
return_value=RegisteredPriorityWeightStrategy,
) as mock_get_registered_priority_weight_strategy:
serop =
OperatorSerialization.deserialize(OperatorSerialization.serialize(sdkop))
+ weight_rule = serop.weight_rule
- assert serop.weight_rule == RegisteredPriorityWeightStrategy()
+ assert weight_rule == RegisteredPriorityWeightStrategy()
assert mock_get_registered_priority_weight_strategy.mock_calls == [
mock.call("unit.serialization.test_dag_serialization.RegisteredPriorityWeightStrategy"),
mock.call("unit.serialization.test_dag_serialization.RegisteredPriorityWeightStrategy"),
mock.call("unit.serialization.test_dag_serialization.RegisteredPriorityWeightStrategy"),
]
+ def test_unregistered_custom_deserializes_until_read(self):
+ sdkop = BaseOperator(task_id="custom",
weight_rule=RegisteredPriorityWeightStrategy())
+ with mock.patch(
+
"airflow.serialization.serialized_objects._get_registered_priority_weight_strategy",
+ return_value=RegisteredPriorityWeightStrategy,
+ ):
+ serialized = OperatorSerialization.serialize(sdkop)
+
+ serop = OperatorSerialization.deserialize(serialized)
+
+ assert serop.task_id == "custom"
+ with pytest.raises(ValueError, match="Unknown priority strategy"):
+ serop.weight_rule
+
def test_invalid(self):
op = BaseOperator(task_id="should_fail", weight_rule="no rule")
with pytest.raises(ValueError, match="Unknown priority strategy"):
