You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@airflow.apache.org by ka...@apache.org on 2025/04/26 09:34:38 UTC
(airflow) branch v3-0-test updated: Fall back when a task cannot be expanded (#49772)
This is an automated email from the ASF dual-hosted git repository.
kaxilnaik pushed a commit to branch v3-0-test
in repository https://gitbox.apache.org/repos/asf/airflow.git
The following commit(s) were added to refs/heads/v3-0-test by this push:
new 130cc4a1506 Fall back when a task cannot be expanded (#49772)
130cc4a1506 is described below
commit 130cc4a1506963944eb8f54d21af5f15d3810e35
Author: Tzu-ping Chung <ur...@gmail.com>
AuthorDate: Sat Apr 26 15:46:17 2025 +0800
Fall back when a task cannot be expanded (#49772)
(cherry picked from commit 19122e7be123b555e91fedb9088d206a5aa685af)
---
.../api_fastapi/core_api/services/ui/grid.py | 26 ++++-----
airflow-core/src/airflow/models/baseoperator.py | 2 +-
.../api_fastapi/core_api/routes/ui/test_grid.py | 67 +++++++++++++++++++++-
3 files changed, 79 insertions(+), 16 deletions(-)
diff --git a/airflow-core/src/airflow/api_fastapi/core_api/services/ui/grid.py b/airflow-core/src/airflow/api_fastapi/core_api/services/ui/grid.py
index 45ba4e42837..7865ffa168f 100644
--- a/airflow-core/src/airflow/api_fastapi/core_api/services/ui/grid.py
+++ b/airflow-core/src/airflow/api_fastapi/core_api/services/ui/grid.py
@@ -17,6 +17,7 @@
from __future__ import annotations
+import contextlib
from functools import cache
from operator import methodcaller
from typing import Callable
@@ -42,6 +43,8 @@ from airflow.models.baseoperator import BaseOperator as DBBaseOperator
from airflow.models.dag_version import DagVersion
from airflow.models.taskmap import TaskMap
from airflow.sdk import BaseOperator
+from airflow.sdk.definitions._internal.abstractoperator import NotMapped
+from airflow.sdk.definitions._internal.expandinput import NotFullyPopulated
from airflow.sdk.definitions.mappedoperator import MappedOperator
from airflow.sdk.definitions.taskgroup import MappedTaskGroup, TaskGroup
from airflow.serialization.serialized_objects import SerializedDAG
@@ -140,19 +143,14 @@ def get_child_task_map(parent_task_id: str, task_node_map: dict[str, dict[str, A
return [task_id for task_id, task_map in task_node_map.items() if task_map["parent_id"] == parent_task_id]
-def _get_total_task_count(
- run_id: str, task_count: list[int | MappedTaskGroup | MappedOperator], session: SessionDep
-) -> int:
- return sum(
- node
- if isinstance(node, int)
- else (
- DBBaseOperator.get_mapped_ti_count(node, run_id=run_id, session=session) or 0
- if isinstance(node, (MappedTaskGroup, MappedOperator))
- else node
- )
- for node in task_count
- )
+def _count_tis(node: int | MappedTaskGroup | MappedOperator, run_id: str, session: SessionDep) -> int:
+ if not isinstance(node, (MappedTaskGroup, MappedOperator)):
+ return node
+ with contextlib.suppress(NotFullyPopulated, NotMapped):
+ return DBBaseOperator.get_mapped_ti_count(node, run_id=run_id, session=session)
+ # If the downstream is not actually mapped, or we don't have information to
+ # determine the length yet, simply return 1 to represent the stand-in ti.
+ return 1
def fill_task_instance_summaries(
@@ -253,7 +251,7 @@ def fill_task_instance_summaries(
end_date=ti_end_date,
queued_dttm=ti_queued_dttm,
child_states=child_states,
- task_count=_get_total_task_count(run_id, task_node_map[task_id]["task_count"], session),
+ task_count=sum(_count_tis(n, run_id, session) for n in task_node_map[task_id]["task_count"]),
state=TaskInstanceState[overall_ti_state.upper()]
if overall_ti_state != "no_status"
else None,
diff --git a/airflow-core/src/airflow/models/baseoperator.py b/airflow-core/src/airflow/models/baseoperator.py
index 4c766007142..9b57d3d5251 100644
--- a/airflow-core/src/airflow/models/baseoperator.py
+++ b/airflow-core/src/airflow/models/baseoperator.py
@@ -621,7 +621,7 @@ class BaseOperator(TaskSDKBaseOperator):
raise NotImplementedError(f"Not implemented for {type(task)}")
# https://github.com/python/cpython/issues/86153
- # WHile we support Python 3.9 we can't rely on the type hint, we need to pass the type explicitly to
+ # While we support Python 3.9 we can't rely on the type hint, we need to pass the type explicitly to
# register.
@get_mapped_ti_count.register(TaskSDKAbstractOperator)
@classmethod
diff --git a/airflow-core/tests/unit/api_fastapi/core_api/routes/ui/test_grid.py b/airflow-core/tests/unit/api_fastapi/core_api/routes/ui/test_grid.py
index 4678967b8da..181d64b3a17 100644
--- a/airflow-core/tests/unit/api_fastapi/core_api/routes/ui/test_grid.py
+++ b/airflow-core/tests/unit/api_fastapi/core_api/routes/ui/test_grid.py
@@ -48,6 +48,7 @@ TASK_ID_3 = "task3"
TASK_ID_4 = "task4"
SUB_TASK_ID = "subtask"
MAPPED_TASK_ID = "mapped_task"
+MAPPED_TASK_ID_2 = "mapped_task_2"
TASK_GROUP_ID = "task_group"
INNER_TASK_GROUP = "inner_task_group"
INNER_TASK_GROUP_SUB_TASK = "inner_task_group_sub_task"
@@ -140,6 +141,31 @@ GRID_RUN_1 = {
"task_id": "task_group",
"try_number": 0,
},
+ {
+ "child_states": {
+ "deferred": 0,
+ "failed": 0,
+ "no_status": 0,
+ "queued": 0,
+ "removed": 0,
+ "restarting": 0,
+ "running": 0,
+ "scheduled": 0,
+ "skipped": 0,
+ "success": 1,
+ "up_for_reschedule": 0,
+ "up_for_retry": 0,
+ "upstream_failed": 0,
+ },
+ "end_date": None,
+ "note": None,
+ "queued_dttm": None,
+ "start_date": None,
+ "state": "success",
+ "task_count": 1,
+ "task_id": "mapped_task_2",
+ "try_number": 0,
+ },
{
"child_states": {
"deferred": 0,
@@ -331,6 +357,31 @@ GRID_RUN_2 = {
"task_id": "task_group",
"try_number": 0,
},
+ {
+ "child_states": {
+ "deferred": 0,
+ "failed": 0,
+ "no_status": 1,
+ "queued": 0,
+ "removed": 0,
+ "restarting": 0,
+ "running": 0,
+ "scheduled": 0,
+ "skipped": 0,
+ "success": 0,
+ "up_for_reschedule": 0,
+ "up_for_retry": 0,
+ "upstream_failed": 0,
+ },
+ "end_date": None,
+ "note": None,
+ "queued_dttm": None,
+ "start_date": None,
+ "state": None,
+ "task_count": 1,
+ "task_id": "mapped_task_2",
+ "try_number": 0,
+ },
{
"child_states": {
"deferred": 0,
@@ -517,6 +568,17 @@ STRUCTURE = {
"tooltip": "",
"type": "task",
},
+ {
+ "asset_condition_type": None,
+ "children": None,
+ "id": "mapped_task_2",
+ "is_mapped": True,
+ "label": "mapped_task_2",
+ "operator": "MockOperator",
+ "setup_teardown_type": None,
+ "tooltip": None,
+ "type": "task",
+ },
],
}
@@ -536,7 +598,7 @@ def setup(dag_maker, session=None):
# DAG 1
with dag_maker(dag_id=DAG_ID, serialized=True, session=session) as dag:
- EmptyOperator(task_id=TASK_ID)
+ task = EmptyOperator(task_id=TASK_ID)
@task_group
def mapped_task_group(arg1):
@@ -549,6 +611,9 @@ def setup(dag_maker, session=None):
with TaskGroup(group_id=INNER_TASK_GROUP):
MockOperator.partial(task_id=INNER_TASK_GROUP_SUB_TASK).expand(arg1=["a", "b"])
+ # Mapped but never expanded. API should not crash, but count this as one no-status ti.
+ MockOperator.partial(task_id=MAPPED_TASK_ID_2).expand(arg1=task.output)
+
triggered_by_kwargs = {"triggered_by": DagRunTriggeredByType.TEST}
logical_date = timezone.datetime(2024, 11, 30)
data_interval = dag.timetable.infer_manual_data_interval(run_after=logical_date)