This is an automated email from the ASF dual-hosted git repository.
kaxilnaik 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 cfce5734b21 Remove unmap method from scheduler-side (#54816)
cfce5734b21 is described below
commit cfce5734b214a5d235c9d9410fe8f54dc7e8a94b
Author: Kaxil Naik <[email protected]>
AuthorDate: Sat Aug 23 23:49:10 2025 +0100
Remove unmap method from scheduler-side (#54816)
Remove unnecessary unmap() functionality from server-side operators to
simplify scheduler architecture and eliminate synthetic operator creation.
Key changes:
- Remove unmap() method from MappedOperator class
- Update TaskInstance.fetch_handle_failure_context() to use original task
directly
- Remove unmap() call from SerializedBaseOperator.get_extra_links()
- Update related tests to verify serialization without unmap functionality
The scheduler no longer needs to 'unmap' operators since:
- Callbacks are handled by DAG processor, not scheduler
- Email settings and fail-fast logic work with original task
- Extra links work consistently between regular and mapped operators
This eliminates the TODO comment about moving runtime unmap to task runner
and provides cleaner separation between scheduler and execution concerns.
Includes pre-commit formatting fixes applied by ruff and ruff-format.
---
airflow-core/src/airflow/models/mappedoperator.py | 24 -------------
airflow-core/src/airflow/models/taskinstance.py | 13 ++------
.../airflow/serialization/serialized_objects.py | 6 +---
.../unit/serialization/test_dag_serialization.py | 39 ----------------------
.../unit/serialization/test_serialized_objects.py | 23 -------------
5 files changed, 4 insertions(+), 101 deletions(-)
diff --git a/airflow-core/src/airflow/models/mappedoperator.py
b/airflow-core/src/airflow/models/mappedoperator.py
index b7df995291e..bda98c8b5a5 100644
--- a/airflow-core/src/airflow/models/mappedoperator.py
+++ b/airflow-core/src/airflow/models/mappedoperator.py
@@ -438,30 +438,6 @@ class MappedOperator(DAGNode):
def expand_start_trigger_args(self, *, context: Context) ->
StartTriggerArgs | None:
raise NotImplementedError
- def unmap(self, resolve: None) -> SerializedBaseOperator:
- """
- Get the "normal" Operator after applying the current mapping.
-
- The *resolve* argument is never used and should always be *None*. It
- exists only to match the signature of the non-serialized
implementation.
-
- The return value is a SerializedBaseOperator that "looks like" the
- actual unmapping result.
-
- :meta private:
- """
- # After a mapped operator is serialized, there's no real way to
actually
- # unmap it since we've lost access to the underlying operator class.
- # This tries its best to simply "forward" all the attributes on this
- # mapped operator to a new SerializedBaseOperator instance.
- sop = SerializedBaseOperator(task_id=self.task_id, params=self.params,
_airflow_from_mapped=True)
- for partial_attr, value in self.partial_kwargs.items():
- setattr(sop, partial_attr, value)
- SerializedBaseOperator.populate_operator(sop, self.operator_class)
- if self.dag is not None: # For Mypy; we only serialize tasks in a DAG
so the check always satisfies.
- SerializedBaseOperator.set_task_dag_references(sop, self.dag)
- return sop
-
@functools.singledispatch
def get_mapped_ti_count(task: DAGNode, run_id: str, *, session: Session) ->
int:
diff --git a/airflow-core/src/airflow/models/taskinstance.py
b/airflow-core/src/airflow/models/taskinstance.py
index 220b68d612a..608aadd1d78 100644
--- a/airflow-core/src/airflow/models/taskinstance.py
+++ b/airflow-core/src/airflow/models/taskinstance.py
@@ -1610,18 +1610,13 @@ class TaskInstance(Base, LoggingMixin):
# only mark task instance as FAILED if the next task instance
# try_number exceeds the max_tries ... or if force_fail is truthy
- task: SerializedBaseOperator | None = None
- try:
- if (orig_task := getattr(ti, "task", None)) and context:
- # TODO (GH-52141): Move runtime unmap into task runner.
- task = orig_task.unmap((context, session))
- except Exception:
- cls.logger().error("Unable to unmap task to determine if we need
to send an alert email")
+ # Use the original task directly - scheduler only needs to check email
settings
+ # Actual callbacks are handled by the DAG processor, not the scheduler
+ task = getattr(ti, "task", None)
if not ti.is_eligible_to_retry():
ti.state = TaskInstanceState.FAILED
email_for_state = operator.attrgetter("email_on_failure")
- callbacks = task.on_failure_callback if task else None
if task and fail_fast:
_stop_remaining_tasks(task_instance=ti, session=session)
@@ -1634,7 +1629,6 @@ class TaskInstance(Base, LoggingMixin):
ti.state = State.UP_FOR_RETRY
email_for_state = operator.attrgetter("email_on_retry")
- callbacks = task.on_retry_callback if task else None
try:
get_listener_manager().hook.on_task_instance_failed(
@@ -1647,7 +1641,6 @@ class TaskInstance(Base, LoggingMixin):
"ti": ti,
"email_for_state": email_for_state,
"task": task,
- "callbacks": callbacks,
"context": context,
}
diff --git a/airflow-core/src/airflow/serialization/serialized_objects.py
b/airflow-core/src/airflow/serialization/serialized_objects.py
index 6a8476ba7a3..f3116c59414 100644
--- a/airflow-core/src/airflow/serialization/serialized_objects.py
+++ b/airflow-core/src/airflow/serialization/serialized_objects.py
@@ -110,7 +110,6 @@ if TYPE_CHECKING:
from airflow.ti_deps.deps.base_ti_dep import BaseTIDep
from airflow.timetables.base import DagRunInfo, DataInterval, Timetable
from airflow.triggers.base import BaseEventTrigger
- from airflow.typing_compat import Self
HAS_KUBERNETES: bool
try:
@@ -1310,7 +1309,7 @@ class SerializedBaseOperator(DAGNode, BaseSerialization):
link = self.operator_extra_link_dict.get(name) or
self.global_operator_extra_link_dict.get(name)
if not link:
return None
- return link.get_link(self.unmap(None), ti_key=ti.key) # type:
ignore[arg-type] # TODO: GH-52141 - BaseOperatorLink.get_link expects
BaseOperator but receives SerializedBaseOperator
+ return link.get_link(self, ti_key=ti.key) # type: ignore[arg-type] #
TODO: GH-52141 - BaseOperatorLink.get_link expects BaseOperator but receives
SerializedBaseOperator
@property
def task_type(self) -> str:
@@ -1759,9 +1758,6 @@ class SerializedBaseOperator(DAGNode, BaseSerialization):
def get_serialized_fields(self):
return BaseOperator.get_serialized_fields()
- def unmap(self, resolve: None) -> Self:
- return self
-
def _iter_all_mapped_downstreams(self) -> Iterator[MappedOperator |
MappedTaskGroup]:
"""
Return mapped nodes that are direct dependencies of the current task.
diff --git a/airflow-core/tests/unit/serialization/test_dag_serialization.py
b/airflow-core/tests/unit/serialization/test_dag_serialization.py
index 3420bae225a..ed3fc7306a9 100644
--- a/airflow-core/tests/unit/serialization/test_dag_serialization.py
+++ b/airflow-core/tests/unit/serialization/test_dag_serialization.py
@@ -62,7 +62,6 @@ from airflow.models.mappedoperator import MappedOperator
from airflow.models.xcom import XCOM_RETURN_KEY, XComModel
from airflow.providers.cncf.kubernetes.pod_generator import PodGenerator
from airflow.providers.standard.operators.bash import BashOperator
-from airflow.providers.standard.sensors.bash import BashSensor
from airflow.sdk import AssetAlias, BaseHook, teardown
from airflow.sdk.bases.decorator import DecoratedOperator
from airflow.sdk.bases.operator import BaseOperator
@@ -2718,44 +2717,6 @@ def test_operator_expand_kwargs_xcomarg_serde(strict):
assert xcom_arg.operator is serialized_dag.task_dict["op1"]
-def test_operator_expand_deserialized_unmap():
- """Unmap a deserialized mapped operator should be similar to deserializing
an non-mapped operator."""
- normal = BashOperator(task_id="a", bash_command=[1, 2],
executor_config={"a": "b"})
- mapped = BashOperator.partial(task_id="a", executor_config={"a":
"b"}).expand(bash_command=[1, 2])
-
- ser_mapped = BaseSerialization.serialize(mapped)
- deser_mapped = BaseSerialization.deserialize(ser_mapped)
- deser_mapped.dag = None
-
- ser_normal = BaseSerialization.serialize(normal)
- deser_normal = BaseSerialization.deserialize(ser_normal)
- deser_normal.dag = None
- unmapped_deser_mapped = deser_mapped.unmap(None)
-
- assert type(unmapped_deser_mapped) is type(deser_normal) is
SerializedBaseOperator
- assert unmapped_deser_mapped.task_id == deser_normal.task_id == "a"
- assert unmapped_deser_mapped.executor_config ==
deser_normal.executor_config == {"a": "b"}
-
-
[email protected]_test
-def test_sensor_expand_deserialized_unmap():
- """Unmap a deserialized mapped sensor should be similar to deserializing a
non-mapped sensor"""
- dag = DAG(dag_id="hello", schedule=None, start_date=None)
- with dag:
- normal = BashSensor(task_id="a", bash_command=[1, 2],
mode="reschedule")
- mapped = BashSensor.partial(task_id="b",
mode="reschedule").expand(bash_command=[1, 2])
- ser_mapped = SerializedBaseOperator.serialize(mapped)
- deser_mapped = SerializedBaseOperator.deserialize(ser_mapped)
- deser_mapped.dag = dag
- deser_unmapped = deser_mapped.unmap(None)
- ser_normal = SerializedBaseOperator.serialize(normal)
- deser_normal = SerializedBaseOperator.deserialize(ser_normal)
- comps = set(BashSensor._comps)
- comps.remove("task_id")
- comps.remove("dag_id")
- assert all(getattr(deser_unmapped, c, None) == getattr(deser_normal, c,
None) for c in comps)
-
-
def test_task_resources_serde():
"""
Test task resources serialization/deserialization.
diff --git a/airflow-core/tests/unit/serialization/test_serialized_objects.py
b/airflow-core/tests/unit/serialization/test_serialized_objects.py
index 4ac943192c1..7aededd12bd 100644
--- a/airflow-core/tests/unit/serialization/test_serialized_objects.py
+++ b/airflow-core/tests/unit/serialization/test_serialized_objects.py
@@ -494,29 +494,6 @@ def test_backcompat_deserialize_connection(conn_uri):
assert deserialized.get_uri() == conn_uri
[email protected]_test
-def test_serialized_mapped_operator_unmap(dag_maker):
- from airflow.serialization.serialized_objects import SerializedDAG
-
- from tests_common.test_utils.mock_operators import MockOperator
-
- with dag_maker(dag_id="dag") as dag:
- MockOperator(task_id="task1", arg1="x")
- MockOperator.partial(task_id="task2").expand(arg1=["a", "b"])
-
- serialized_dag = SerializedDAG.from_dict(SerializedDAG.to_dict(dag))
- assert serialized_dag.dag_id == "dag"
-
- serialized_task1 = serialized_dag.get_task("task1")
- assert serialized_task1.dag is serialized_dag
-
- serialized_task2 = serialized_dag.get_task("task2")
- assert serialized_task2.dag is serialized_dag
-
- serialized_unmapped_task = serialized_task2.unmap(None)
- assert serialized_unmapped_task.dag is serialized_dag
-
-
def test_ser_of_asset_event_accessor():
# todo: (Airflow 3.0) we should force reserialization on upgrade
d = OutletEventAccessors()
