This is an automated email from the ASF dual-hosted git repository.
uranusjr 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 279565022b Consider mapped upstream in TriggerRuleDep (#22583)
279565022b is described below
commit 279565022b94d2dc8a1ae0f7dd6bf4d1dd10dbdc
Author: Tzu-ping Chung <[email protected]>
AuthorDate: Wed Apr 6 14:27:46 2022 +0800
Consider mapped upstream in TriggerRuleDep (#22583)
---
airflow/models/baseoperator.py | 2 +-
airflow/models/taskinstance.py | 30 +++++++++++++----
airflow/models/xcom.py | 4 ++-
airflow/models/xcom_arg.py | 15 ++++-----
airflow/ti_deps/deps/trigger_rule_dep.py | 41 +++++++++++++++++++----
tests/models/test_dagrun.py | 51 +++++++++++++++++++++++++++++
tests/models/test_taskinstance.py | 32 ++++++++++++++++++
tests/ti_deps/deps/test_trigger_rule_dep.py | 8 ++---
8 files changed, 156 insertions(+), 27 deletions(-)
diff --git a/airflow/models/baseoperator.py b/airflow/models/baseoperator.py
index 8f65f9c39a..1e65bb9635 100644
--- a/airflow/models/baseoperator.py
+++ b/airflow/models/baseoperator.py
@@ -837,7 +837,7 @@ class BaseOperator(AbstractOperator,
metaclass=BaseOperatorMeta):
f"'{dag.dag_id if dag else ''}.{task_id}'; received
'{trigger_rule}'."
)
- self.trigger_rule = trigger_rule
+ self.trigger_rule = TriggerRule(trigger_rule)
self.depends_on_past: bool = depends_on_past
self.ignore_first_depends_on_past = ignore_first_depends_on_past
self.wait_for_downstream = wait_for_downstream
diff --git a/airflow/models/taskinstance.py b/airflow/models/taskinstance.py
index 60ce957398..b9a1ebc148 100644
--- a/airflow/models/taskinstance.py
+++ b/airflow/models/taskinstance.py
@@ -823,17 +823,25 @@ class TaskInstance(Base, LoggingMixin):
self.operator = task.task_type
@provide_session
- def clear_xcom_data(self, session=NEW_SESSION):
- """
- Clears all XCom data from the database for the task instance
+ def clear_xcom_data(self, session: Session = NEW_SESSION):
+ """Clear all XCom data from the database for the task instance.
+
+ If the task is unmapped, all XComs matching this task ID in the same
DAG
+ run are removed. If the task is mapped, only the one with matching map
+ index is removed.
:param session: SQLAlchemy ORM Session
"""
self.log.debug("Clearing XCom data")
+ if self.map_index < 0:
+ map_index: Optional[int] = None
+ else:
+ map_index = self.map_index
XCom.clear(
dag_id=self.dag_id,
task_id=self.task_id,
run_id=self.run_id,
+ map_index=map_index,
session=session,
)
@@ -2279,6 +2287,7 @@ class TaskInstance(Base, LoggingMixin):
session: Session = NEW_SESSION,
*,
map_indexes: Optional[Union[int, Iterable[int]]] = None,
+ default: Any = None,
) -> Any:
"""Pull XComs that optionally meet certain criteria.
@@ -2303,7 +2312,8 @@ class TaskInstance(Base, LoggingMixin):
the specified task is mapped. If not, value from the one single task
instance is returned. If the task to pull is mapped, an iterator (not a
list) yielding XComs from mapped task instances is returned. In either
- case, *None* is returned if no matching XComs are found.
+ case, ``default`` (*None* if not specified) is returned if no matching
+ XComs are found.
When pulling multiple tasks (i.e. either ``task_id`` or ``map_index``
is
a non-str iterable), a list of matching XComs is returned. Elements in
@@ -2330,9 +2340,10 @@ class TaskInstance(Base, LoggingMixin):
if (task_ids is None or isinstance(task_ids, str)) and not
isinstance(map_indexes, Iterable):
first = query.with_entities(XCom.run_id, XCom.task_id,
XCom.map_index, XCom.value).first()
if first is None: # No matching XCom at all.
- return None
+ return default
if map_indexes is not None or first.map_index < 0:
return XCom.deserialize_value(first)
+
# We're pulling one specific mapped task. Add additional filters to
# make sure all XComs come from one task and run (for task_ids=None
# and include_prior_dates=True), and re-order by map index (reset
@@ -2343,7 +2354,14 @@ class TaskInstance(Base, LoggingMixin):
.order_by(None)
.order_by(XCom.map_index.asc())
)
- return (XCom.deserialize_value(r) for r in query)
+
+ def iter_xcom_values(query):
+ # The session passed to xcom_pull() may die before this is
+ # iterated through, so we need to bind to a new session.
+ for r in query.with_session(settings.Session()):
+ yield XCom.deserialize_value(r)
+
+ return iter_xcom_values(query)
# At this point either task_ids or map_indexes is explicitly
multi-value.
diff --git a/airflow/models/xcom.py b/airflow/models/xcom.py
index 9c0ce3b40c..d67160d1fa 100644
--- a/airflow/models/xcom.py
+++ b/airflow/models/xcom.py
@@ -103,7 +103,9 @@ class BaseXCom(Base, LoggingMixin):
self.value = self.orm_deserialize_value()
def __repr__(self):
- return f'<XCom "{self.key}" ({self.task_id} @ {self.run_id})>'
+ if self.map_index < 0:
+ return f'<XCom "{self.key}" ({self.task_id} @ {self.run_id})>'
+ return f'<XCom "{self.key}" ({self.task_id}[{self.map_index}] @
{self.run_id})>'
@overload
@classmethod
diff --git a/airflow/models/xcom_arg.py b/airflow/models/xcom_arg.py
index 7b6f28d5ca..7e4b274626 100644
--- a/airflow/models/xcom_arg.py
+++ b/airflow/models/xcom_arg.py
@@ -22,6 +22,7 @@ from airflow.models.taskmixin import DAGNode, DependencyMixin
from airflow.models.xcom import XCOM_RETURN_KEY
from airflow.utils.context import Context
from airflow.utils.edgemodifier import EdgeModifier
+from airflow.utils.types import NOTSET
if TYPE_CHECKING:
from airflow.models.operator import Operator
@@ -140,17 +141,13 @@ class XComArg(DependencyMixin):
Pull XCom value for the existing arg. This method is run during
``op.execute()``
in respectable context.
"""
- resolved_value =
context['ti'].xcom_pull(task_ids=[self.operator.task_id], key=str(self.key))
- if not resolved_value:
- if TYPE_CHECKING:
- assert self.operator.dag
+ result = context["ti"].xcom_pull(task_ids=self.operator.task_id,
key=str(self.key), default=NOTSET)
+ if result is NOTSET:
raise AirflowException(
- f'XComArg result from {self.operator.task_id} at
{self.operator.dag.dag_id} '
- f'with key="{self.key}"" is not found!'
+ f'XComArg result from {self.operator.task_id} at
{context["ti"].dag_id} '
+ f'with key="{self.key}" is not found!'
)
- resolved_value = resolved_value[0]
-
- return resolved_value
+ return result
@staticmethod
def apply_upstream_relationship(op: "Operator", arg: Any):
diff --git a/airflow/ti_deps/deps/trigger_rule_dep.py
b/airflow/ti_deps/deps/trigger_rule_dep.py
index f8383b4c7c..1a65467896 100644
--- a/airflow/ti_deps/deps/trigger_rule_dep.py
+++ b/airflow/ti_deps/deps/trigger_rule_dep.py
@@ -19,6 +19,8 @@
from collections import Counter
from typing import TYPE_CHECKING
+from sqlalchemy import func
+
from airflow.ti_deps.dep_context import DepContext
from airflow.ti_deps.deps.base_ti_dep import BaseTIDep
from airflow.utils.session import NEW_SESSION, provide_session
@@ -26,7 +28,9 @@ from airflow.utils.state import State
from airflow.utils.trigger_rule import TriggerRule as TR
if TYPE_CHECKING:
- from airflow.settings import SASession as Session
+ from sqlalchemy.orm import Session
+
+ from airflow.models.taskinstance import TaskInstance
class TriggerRuleDep(BaseTIDep):
@@ -40,7 +44,7 @@ class TriggerRuleDep(BaseTIDep):
IS_TASK_DEP = True
@staticmethod
- def _get_states_count_upstream_ti(ti, finished_tis):
+ def _get_states_count_upstream_ti(task, finished_tis):
"""
This function returns the states of the upstream tis for a specific ti
in order to determine
whether this ti can run in this iteration
@@ -48,7 +52,7 @@ class TriggerRuleDep(BaseTIDep):
:param ti: the ti that we want to calculate deps for
:param finished_tis: all the finished tasks of the dag_run
"""
- counter = Counter(task.state for task in finished_tis if task.task_id
in ti.task.upstream_task_ids)
+ counter = Counter(ti.state for ti in finished_tis if ti.task_id in
task.upstream_task_ids)
return (
counter.get(State.SUCCESS, 0),
counter.get(State.SKIPPED, 0),
@@ -69,7 +73,7 @@ class TriggerRuleDep(BaseTIDep):
return
# see if the task name is in the task upstream for our task
successes, skipped, failed, upstream_failed, done =
self._get_states_count_upstream_ti(
- ti=ti,
finished_tis=dep_context.ensure_finished_tis(ti.get_dagrun(session), session)
+ task=ti.task,
finished_tis=dep_context.ensure_finished_tis(ti.get_dagrun(session), session)
)
yield from self._evaluate_trigger_rule(
@@ -83,6 +87,31 @@ class TriggerRuleDep(BaseTIDep):
session=session,
)
+ @staticmethod
+ def _count_upstreams(ti: "TaskInstance", *, session: "Session"):
+ from airflow.models.taskinstance import TaskInstance
+
+ # Optimization: Don't need to hit the database if no upstreams are
mapped.
+ upstream_task_ids = ti.task.upstream_task_ids
+ if ti.task.dag and not any(ti.task.dag.get_task(tid).is_mapped for tid
in upstream_task_ids):
+ return len(upstream_task_ids)
+
+ # We don't naively count task instances because it is not guaranteed
+ # that all upstreams have been created in the database at this point.
+ # Instead, we look for already-expanded tasks, and add them to the raw
+ # task count without considering mapping.
+ mapped_tis_addition = (
+ session.query(func.count())
+ .filter(
+ TaskInstance.dag_id == ti.dag_id,
+ TaskInstance.run_id == ti.run_id,
+ TaskInstance.task_id.in_(upstream_task_ids),
+ TaskInstance.map_index > 0,
+ )
+ .scalar()
+ )
+ return len(upstream_task_ids) + mapped_tis_addition
+
@provide_session
def _evaluate_trigger_rule(
self,
@@ -112,8 +141,8 @@ class TriggerRuleDep(BaseTIDep):
:param session: database session
"""
task = ti.task
- upstream = len(task.upstream_task_ids)
- trigger_rule = task.trigger_rule
+ upstream = self._count_upstreams(ti, session=session)
+ trigger_rule: TR = task.trigger_rule
upstream_done = done >= upstream
upstream_tasks_state = {
"total": upstream,
diff --git a/tests/models/test_dagrun.py b/tests/models/test_dagrun.py
index 4697354db5..f2a42d40c8 100644
--- a/tests/models/test_dagrun.py
+++ b/tests/models/test_dagrun.py
@@ -981,3 +981,54 @@ def test_mapped_task_upstream_failed(dag_maker, session):
tis, _ = dr.update_state(execute_callbacks=False, session=session)
assert tis == []
assert dr.state == DagRunState.FAILED
+
+
+def test_mapped_task_all_finish_before_downstream(dag_maker, session):
+ result = None
+
+ with dag_maker(session=session) as dag:
+
+ @dag.task
+ def make_list():
+ return [1, 2]
+
+ @dag.task
+ def double(value):
+ return value * 2
+
+ @dag.task
+ def consumer(value):
+ nonlocal result
+ result = list(value)
+
+ consumer(value=double.expand(value=make_list()))
+
+ dr: DagRun = dag_maker.create_dagrun()
+
+ def _task_ids(tis):
+ return [ti.task_id for ti in tis]
+
+ # The first task is always make_list.
+ decision = dr.task_instance_scheduling_decisions(session=session)
+ assert _task_ids(decision.schedulable_tis) == ["make_list"]
+
+ # After make_list is run, double is expanded.
+ decision.schedulable_tis[0].run(verbose=False, session=session)
+ decision = dr.task_instance_scheduling_decisions(session=session)
+ assert _task_ids(decision.schedulable_tis) == ["double", "double"]
+
+ # Running just one of the mapped tis does not make downstream schedulable.
+ decision.schedulable_tis[0].run(verbose=False, session=session)
+ decision = dr.task_instance_scheduling_decisions(session=session)
+ assert _task_ids(decision.schedulable_tis) == ["double"]
+
+ # Downstream is schedulable after all mapped tis are run.
+ decision.schedulable_tis[0].run(verbose=False, session=session)
+ decision = dr.task_instance_scheduling_decisions(session=session)
+ assert _task_ids(decision.schedulable_tis) == ["consumer"]
+
+ # We should be able to get all values aggregated from mapped upstreams.
+ decision.schedulable_tis[0].run(verbose=False, session=session)
+ decision = dr.task_instance_scheduling_decisions(session=session)
+ assert decision.schedulable_tis == []
+ assert result == [2, 4]
diff --git a/tests/models/test_taskinstance.py
b/tests/models/test_taskinstance.py
index e8be37434c..38be50b12d 100644
--- a/tests/models/test_taskinstance.py
+++ b/tests/models/test_taskinstance.py
@@ -2585,3 +2585,35 @@ class TestMappedTaskInstanceReceiveValue:
with out.open() as f:
out_lines = [line.strip() for line in f]
assert out_lines == ["hello FOO", "goodbye FOO", "hello BAR", "goodbye
BAR"]
+
+
[email protected]("airflow.models.taskinstance.XCom.deserialize_value",
side_effect=XCom.deserialize_value)
+def
test_ti_xcom_pull_on_mapped_operator_return_lazy_iterator(mock_deserialize_value,
dag_maker, session):
+ """Ensure we access XCom lazily when pulling from a mapped operator."""
+ with dag_maker(dag_id="test_xcom", session=session):
+ task_1 = DummyOperator.partial(task_id="task_1").expand()
+ DummyOperator(task_id="task_2")
+
+ dagrun = dag_maker.create_dagrun()
+
+ ti_1_0 = dagrun.get_task_instance("task_1", session=session)
+ ti_1_0.map_index = 0
+ ti_1_1 = session.merge(TaskInstance(task_1, run_id=dagrun.run_id,
map_index=1, state=ti_1_0.state))
+ session.flush()
+
+ ti_1_0.xcom_push(key=XCOM_RETURN_KEY, value="a", session=session)
+ ti_1_1.xcom_push(key=XCOM_RETURN_KEY, value="b", session=session)
+
+ ti_2 = dagrun.get_task_instance("task_2", session=session)
+
+ # Simply pulling the joined XCom value should not deserialize.
+ joined = ti_2.xcom_pull("task_1", session=session)
+ assert mock_deserialize_value.call_count == 0
+
+ # Only when we go through the iterator does deserialization happen.
+ assert next(joined) == "a"
+ assert mock_deserialize_value.call_count == 1
+ assert next(joined) == "b"
+ assert mock_deserialize_value.call_count == 2
+ with pytest.raises(StopIteration):
+ next(joined)
diff --git a/tests/ti_deps/deps/test_trigger_rule_dep.py
b/tests/ti_deps/deps/test_trigger_rule_dep.py
index cf2cada840..14dc15dc3b 100644
--- a/tests/ti_deps/deps/test_trigger_rule_dep.py
+++ b/tests/ti_deps/deps/test_trigger_rule_dep.py
@@ -41,7 +41,7 @@ def get_task_instance(session, dag_maker):
task_id='test_task', trigger_rule=trigger_rule,
start_date=datetime(2015, 1, 1)
)
if upstream_task_ids:
- task.upstream_task_ids.update(upstream_task_ids)
+ [DummyOperator(task_id=task_id) for task_id in
upstream_task_ids] >> task
dr = dag_maker.create_dagrun()
ti = dr.task_instances[0]
ti.task = task
@@ -662,10 +662,10 @@ class TestTriggerRuleDep:
# check handling with cases that tasks are triggered from backfill
with no finished tasks
finished_tis = DepContext().ensure_finished_tis(ti_op2.dag_run,
session)
- assert get_states_count_upstream_ti(finished_tis=finished_tis,
ti=ti_op2) == (1, 0, 0, 0, 1)
+ assert get_states_count_upstream_ti(finished_tis=finished_tis,
task=op2) == (1, 0, 0, 0, 1)
finished_tis = dr.get_task_instances(state=State.finished,
session=session)
- assert get_states_count_upstream_ti(finished_tis=finished_tis,
ti=ti_op4) == (1, 0, 1, 0, 2)
- assert get_states_count_upstream_ti(finished_tis=finished_tis,
ti=ti_op5) == (2, 0, 1, 0, 3)
+ assert get_states_count_upstream_ti(finished_tis=finished_tis,
task=op4) == (1, 0, 1, 0, 2)
+ assert get_states_count_upstream_ti(finished_tis=finished_tis,
task=op5) == (2, 0, 1, 0, 3)
dr.update_state()
assert State.SUCCESS == dr.state
