Reduce API server memory usage by eliminating SerializedDAG loads on task start

airflow-core/newsfragments/60803.significant.rst

@@ -0,0 +1 @@
+Move ``upstream_map_indexes`` computation from API server to Task SDK, reducing memory usage on task start by eliminating ``SerializedDAG`` loads.

airflow-core/src/airflow/api_fastapi/execution_api/datamodels/taskinstance.py

@@ -324,8 +324,6 @@ class TIRunContext(BaseModel):
     connections: Annotated[list[ConnectionResponse], Field(default_factory=list)]
     """Connections that can be accessed by the task instance."""
 
-    upstream_map_indexes: dict[str, int | list[int] | None] | None = None
-
     next_method: str | None = None
     """Method to call. Set when task resumes from a trigger."""
     next_kwargs: dict[str, Any] | str | None = None

airflow-core/src/airflow/api_fastapi/execution_api/routes/task_instances.py

@@ -31,7 +31,7 @@
 from fastapi import Body, HTTPException, Query, status
 from pydantic import JsonValue
 from sqlalchemy import func, or_, tuple_, update
-from sqlalchemy.engine import CursorResult, Row
+from sqlalchemy.engine import CursorResult
 from sqlalchemy.exc import NoResultFound, SQLAlchemyError
 from sqlalchemy.orm import joinedload
 from sqlalchemy.sql import select
@@ -64,14 +64,11 @@
 from airflow.models.asset import AssetActive
 from airflow.models.dag import DagModel
 from airflow.models.dagrun import DagRun as DR
-from airflow.models.expandinput import NotFullyPopulated
 from airflow.models.taskinstance import TaskInstance as TI, _stop_remaining_tasks
 from airflow.models.taskreschedule import TaskReschedule
 from airflow.models.trigger import Trigger
 from airflow.models.xcom import XComModel
 from airflow.serialization.definitions.assets import SerializedAsset, SerializedAssetUniqueKey
-from airflow.serialization.definitions.dag import SerializedDAG
-from airflow.task.trigger_rule import TriggerRule
 from airflow.utils.sqlalchemy import get_dialect_name
 from airflow.utils.state import DagRunState, TaskInstanceState, TerminalTIState
 
@@ -251,17 +248,6 @@ def ti_run(
             or 0
         )
 
-        if dag := dag_bag.get_dag_for_run(dag_run=dr, session=session):
-            upstream_map_indexes = dict(
-                _get_upstream_map_indexes(
-                    serialized_dag=dag,
-                    ti=ti,
-                    session=session,
-                )
-            )
-        else:
-            upstream_map_indexes = None
-
         context = TIRunContext(
             dag_run=dr,
             task_reschedule_count=task_reschedule_count,
@@ -271,7 +257,6 @@ def ti_run(
             connections=[],
             xcom_keys_to_clear=xcom_keys,
             should_retry=_is_eligible_to_retry(previous_state, ti.try_number, ti.max_tries),
-            upstream_map_indexes=upstream_map_indexes,
         )
 
         # Only set if they are non-null
@@ -287,47 +272,6 @@ def ti_run(
         )
 
 
-def _get_upstream_map_indexes(
-    *,
-    serialized_dag: SerializedDAG,
-    ti: TI | Row,
-    session: SessionDep,
-) -> Iterator[tuple[str, int | list[int] | None]]:
-    task = serialized_dag.get_task(ti.task_id)
-    for upstream_task in task.upstream_list:
-        map_indexes: int | list[int] | None
-        if (upstream_mapped_group := upstream_task.get_closest_mapped_task_group()) is None:
-            # regular tasks or non-mapped task groups
-            map_indexes = None
-        elif task.get_closest_mapped_task_group() is upstream_mapped_group:
-            # tasks in the same mapped task group hierarchy
-            map_indexes = ti.map_index
-        else:
-            # tasks not in the same mapped task group
-            # the upstream mapped task group should combine the return xcom as a list and return it
-            mapped_ti_count: int | None = None
-
-            try:
-                # First try: without resolving XCom
-                mapped_ti_count = upstream_mapped_group.get_parse_time_mapped_ti_count()
-            except NotFullyPopulated:
-                # Second try: resolve XCom for correct count
-                try:
-                    expand_input = upstream_mapped_group._expand_input
-                    mapped_ti_count = expand_input.get_total_map_length(ti.run_id, session=session)
-                except NotFullyPopulated:
-                    # For these trigger rules, unresolved map indexes are acceptable.
-                    # The success of the upstream task is not the main reason for triggering the current task.
-                    # Therefore, whether the upstream task is fully populated can be ignored.
-                    if task.trigger_rule != TriggerRule.ALL_SUCCESS:
-                        mapped_ti_count = None
-
-            # Compute map indexes if we have a valid count
-            map_indexes = list(range(mapped_ti_count)) if mapped_ti_count is not None else None
-
-        yield upstream_task.task_id, map_indexes
-
-
 @ti_id_router.patch(
     "/{task_instance_id}/state",
     status_code=status.HTTP_204_NO_CONTENT,

airflow-core/src/airflow/api_fastapi/execution_api/versions/__init__.py

@@ -33,11 +33,14 @@
     AddDagRunDetailEndpoint,
     MovePreviousRunEndpoint,
 )
-from airflow.api_fastapi.execution_api.versions.v2026_03_31 import ModifyDeferredTaskKwargsToJsonValue
+from airflow.api_fastapi.execution_api.versions.v2026_03_31 import (
+    ModifyDeferredTaskKwargsToJsonValue,
+    RemoveUpstreamMapIndexesField,
+)
 
 bundle = VersionBundle(
     HeadVersion(),
-    Version("2026-03-31", ModifyDeferredTaskKwargsToJsonValue),
+    Version("2026-03-31", ModifyDeferredTaskKwargsToJsonValue, RemoveUpstreamMapIndexesField),
     Version("2025-12-08", MovePreviousRunEndpoint, AddDagRunDetailEndpoint),
     Version("2025-11-07", AddPartitionKeyField),
     Version("2025-11-05", AddTriggeringUserNameField),

airflow-core/src/airflow/api_fastapi/execution_api/versions/v2026_03_31.py

@@ -19,9 +19,9 @@
 
 from typing import Any
 
-from cadwyn import VersionChange, schema
+from cadwyn import ResponseInfo, VersionChange, convert_response_to_previous_version_for, schema
 
-from airflow.api_fastapi.execution_api.datamodels.taskinstance import TIDeferredStatePayload
+from airflow.api_fastapi.execution_api.datamodels.taskinstance import TIDeferredStatePayload, TIRunContext
 
 
 class ModifyDeferredTaskKwargsToJsonValue(VersionChange):
@@ -33,3 +33,20 @@ class ModifyDeferredTaskKwargsToJsonValue(VersionChange):
         schema(TIDeferredStatePayload).field("trigger_kwargs").had(type=dict[str, Any] | str),
         schema(TIDeferredStatePayload).field("next_kwargs").had(type=dict[str, Any]),
     )
+
+
+class RemoveUpstreamMapIndexesField(VersionChange):
+    """Remove upstream_map_indexes field from TIRunContext - now computed by Task SDK."""
+
+    description = __doc__
+
+    instructions_to_migrate_to_previous_version = (
+        schema(TIRunContext)
+        .field("upstream_map_indexes")
+        .existed_as(type=dict[str, int | list[int] | None] | None),
+    )
+
+    @convert_response_to_previous_version_for(TIRunContext)  # type: ignore[arg-type]
+    def add_upstream_map_indexes_field(response: ResponseInfo) -> None:  # type: ignore[misc]
+        """Add upstream_map_indexes field with None for older API versions."""
+        response.body["upstream_map_indexes"] = None

airflow-core/tests/unit/api_fastapi/execution_api/versions/head/test_task_instances.py

@@ -202,7 +202,6 @@ def test_ti_run_state_to_running(
                 "partition_key": None,
             },
             "task_reschedule_count": 0,
-            "upstream_map_indexes": {},
             "max_tries": max_tries,
             "should_retry": should_retry,
             "variables": [],
@@ -249,10 +248,7 @@ def test_ti_run_state_to_running(
         assert response.status_code == 409
 
     def test_dynamic_task_mapping_with_parse_time_value(self, client, dag_maker):
-        """
-        Test that the Task Instance upstream_map_indexes is correctly fetched when to running the Task Instances
-        """
-
+        """Test that dynamic task mapping works correctly with parse-time values."""
         with dag_maker("test_dynamic_task_mapping_with_parse_time_value", serialized=True):
 
             @task_group
@@ -278,23 +274,6 @@ def task2():
             ti.set_state(State.QUEUED)
         dag_maker.session.flush()
 
-        # key: (task_id, map_index)
-        # value: result upstream_map_indexes ({task_id: map_indexes})
-        expected_upstream_map_indexes = {
-            # no upstream task for task_group_1.group_task_1
-            ("task_group_1.group1_task_1", 0): {},
-            ("task_group_1.group1_task_1", 1): {},
-            # the upstream task for task_group_1.group_task_2 is task_group_1.group_task_2
-            # since they are in the same task group, the upstream map index should be the same as the task
-            ("task_group_1.group1_task_2", 0): {"task_group_1.group1_task_1": 0},
-            ("task_group_1.group1_task_2", 1): {"task_group_1.group1_task_1": 1},
-            # the upstream task for task2 is the last tasks of task_group_1, which is
-            # task_group_1.group_task_2
-            # since they are not in the same task group, the upstream map index should include all the
-            # expanded tasks
-            ("task2", -1): {"task_group_1.group1_task_2": [0, 1]},
-        }
-
         for ti in dr.get_task_instances():
             response = client.patch(
                 f"/execution/task-instances/{ti.id}/run",
@@ -308,13 +287,9 @@ def task2():
             )
 
             assert response.status_code == 200
-            upstream_map_indexes = response.json()["upstream_map_indexes"]
-            assert upstream_map_indexes == expected_upstream_map_indexes[(ti.task_id, ti.map_index)]
 
-    def test_nested_mapped_task_group_upstream_indexes(self, client, dag_maker):
-        """
-        Test that upstream_map_indexes are correctly computed for tasks in nested mapped task groups.
-        """
+    def test_nested_mapped_task_group(self, client, dag_maker):
+        """Test that nested mapped task groups work correctly."""
         with dag_maker("test_nested_mapped_tg", serialized=True):
 
             @task
@@ -346,25 +321,11 @@ def expandable_task_group(param: str) -> None:
                 ti.set_state(State.QUEUED)
         dag_maker.session.flush()
 
-        # Expected upstream_map_indexes for each print_task instance
-        expected_upstream_map_indexes = {
-            ("expandable_task_group.inner_task_group.print_task", 0): {
-                "expandable_task_group.inner_task_group.alter_input": 0
-            },
-            ("expandable_task_group.inner_task_group.print_task", 1): {
-                "expandable_task_group.inner_task_group.alter_input": 1
-            },
-            ("expandable_task_group.inner_task_group.print_task", 2): {
-                "expandable_task_group.inner_task_group.alter_input": 2
-            },
-        }
-
         # Get only the expanded print_task instances (not the template)
         print_task_tis = [
             ti for ti in dr.get_task_instances() if "print_task" in ti.task_id and ti.map_index >= 0
         ]
 
-        # Test each print_task instance
         for ti in print_task_tis:
             response = client.patch(
                 f"/execution/task-instances/{ti.id}/run",
@@ -378,18 +339,9 @@ def expandable_task_group(param: str) -> None:
             )
 
             assert response.status_code == 200
-            upstream_map_indexes = response.json()["upstream_map_indexes"]
-            expected = expected_upstream_map_indexes[(ti.task_id, ti.map_index)]
-
-            assert upstream_map_indexes == expected, (
-                f"Task {ti.task_id}[{ti.map_index}] should have upstream_map_indexes {expected}, "
-                f"but got {upstream_map_indexes}"
-            )
 
     def test_dynamic_task_mapping_with_xcom(self, client: Client, dag_maker: DagMaker, session: Session):
-        """
-        Test that the Task Instance upstream_map_indexes is correctly fetched when to running the Task Instances with xcom
-        """
+        """Test that dynamic task mapping works correctly with XCom values."""
         from airflow.models.taskmap import TaskMap
 
         with dag_maker(session=session, serialized=True):
@@ -442,13 +394,10 @@ def task_3():
                 "start_date": "2024-09-30T12:00:00Z",
             },
         )
-        assert response.json()["upstream_map_indexes"] == {"tg.task_2": [0, 1, 2, 3, 4, 5]}
+        assert response.status_code == 200
 
     def test_dynamic_task_mapping_with_all_success_trigger_rule(self, dag_maker: DagMaker, session: Session):
-        """
-        Test that the Task Instance upstream_map_indexes is not populuated but
-        the downstream task should not be run.
-        """
+        """Test that with ALL_SUCCESS trigger rule and skipped upstream, downstream should not run."""
 
         with dag_maker(session=session, serialized=True):
 
@@ -504,10 +453,7 @@ def task_4():
     def test_dynamic_task_mapping_with_non_all_success_trigger_rule(
         self, client: Client, dag_maker: DagMaker, session: Session, trigger_rule: TriggerRule
     ):
-        """
-        Test that the Task Instance upstream_map_indexes is not populuated but
-        the downstream task should still be run due to trigger rule.
-        """
+        """Test that with non-ALL_SUCCESS trigger rule, downstream task should still run."""
 
         with dag_maker(session=session, serialized=True):
 
@@ -564,7 +510,7 @@ def task_4():
                 "start_date": "2024-09-30T12:00:00Z",
             },
         )
-        assert response.json()["upstream_map_indexes"] == {"tg.task_2": None}
+        assert response.status_code == 200
 
     def test_next_kwargs_still_encoded(self, client, session, create_task_instance, time_machine):
         instant_str = "2024-09-30T12:00:00Z"
@@ -615,7 +561,6 @@ def test_next_kwargs_still_encoded(self, client, session, create_task_instance,
         assert response.json() == {
             "dag_run": mock.ANY,
             "task_reschedule_count": 0,
-            "upstream_map_indexes": {},
             "max_tries": 0,
             "should_retry": False,
             "variables": [],
@@ -687,7 +632,6 @@ def test_next_kwargs_determines_start_date_update(self, client, session, create_
         assert response.json() == {
             "dag_run": mock.ANY,
             "task_reschedule_count": 0,
-            "upstream_map_indexes": {},
             "max_tries": 0,
             "should_retry": False,
             "variables": [],