Repository: apache/airflow

Issue:
Backfill mode with mapped tasks: "Failed to populate all mapping metadata"
### Apache Airflow version

2.3.3

### What happened

I was backfilling some DAGs that use dynamic tasks when I got an exception like the following:

```
Traceback (most recent call last):
  File "/opt/conda/envs/production/bin/airflow", line 11, in <module>
    sys.exit(main())
  File "/opt/conda/envs/production/lib/python3.9/site-packages/airflow/__main__.py", line 38, in main
    args.func(args)
  File "/opt/conda/envs/production/lib/python3.9/site-packages/airflow/cli/cli_parser.py", line 51, in command
    return func(*args, **kwargs)
  File "/opt/conda/envs/production/lib/python3.9/site-packages/airflow/utils/cli.py", line 99, in wrapper
    return f(*args, **kwargs)
  File "/opt/conda/envs/production/lib/python3.9/site-packages/airflow/cli/commands/dag_command.py", line 107, in dag_backfill
    dag.run(
  File "/opt/conda/envs/production/lib/python3.9/site-packages/airflow/models/dag.py", line 2288, in run
    job.run()
  File "/opt/conda/envs/production/lib/python3.9/site-packages/airflow/jobs/base_job.py", line 244, in run
    self._execute()
  File "/opt/conda/envs/production/lib/python3.9/site-packages/airflow/utils/session.py", line 71, in wrapper
    return func(*args, session=session, **kwargs)
  File "/opt/conda/envs/production/lib/python3.9/site-packages/airflow/jobs/backfill_job.py", line 847, in _execute
    self._execute_dagruns(
  File "/opt/conda/envs/production/lib/python3.9/site-packages/airflow/utils/session.py", line 68, in wrapper
    return func(*args, **kwargs)
  File "/opt/conda/envs/production/lib/python3.9/site-packages/airflow/jobs/backfill_job.py", line 737, in _execute_dagruns
    processed_dag_run_dates = self._process_backfill_task_instances(
  File "/opt/conda/envs/production/lib/python3.9/site-packages/airflow/utils/session.py", line 68, in wrapper
    return func(*args, **kwargs)
  File "/opt/conda/envs/production/lib/python3.9/site-packages/airflow/jobs/backfill_job.py", line 612, in _process_backfill_task_instances
    for node, run_id, new_mapped_tis, max_map_index in self._manage_executor_state(
  File "/opt/conda/envs/production/lib/python3.9/site-packages/airflow/jobs/backfill_job.py", line 270, in _manage_executor_state
    new_tis, num_mapped_tis = node.expand_mapped_task(ti.run_id, session=session)
  File "/opt/conda/envs/production/lib/python3.9/site-packages/airflow/models/mappedoperator.py", line 614, in expand_mapped_task
    operator.mul, self._resolve_map_lengths(run_id, session=session).values()
  File "/opt/conda/envs/production/lib/python3.9/site-packages/airflow/models/mappedoperator.py", line 600, in _resolve_map_lengths
    raise RuntimeError(f"Failed to populate all mapping metadata; missing: {keys}")
RuntimeError: Failed to populate all mapping metadata; missing: 'x'
```

Digging further, it appears this always happens if the task used as input to an `.expand` raises an Exception.  Airflow doesn't handle this exception gracefully like it does with exceptions in "normal" tasks, which can lead to other errors from deeper within Airflow.  This also means that since this is not a "typical" failure case, things like `--rerun-failed-tasks` do not work as expected.

### What you think should happen instead

Airflow should fail gracefully if exceptions are raised in dynamic task generators.

### How to reproduce

```
#!/usr/bin/env python3

import datetime
import logging

from airflow.decorators import dag, task


logger = logging.getLogger(__name__)


@dag(
    schedule_interval='@daily',
    start_date=datetime.datetime(2022, 8, 12),
    default_args={
        'retries': 5,
        'retry_delay': 5.0,
    },
)
def test_backfill():
    @task
    def get_tasks(ti=None):
        logger.info(f'{ti.try_number=}')
        if ti.try_number < 3:
            raise RuntimeError('')
        return ['a', 'b', 'c']

    @task
    def do_stuff(x=None, ti=None):
        logger.info(f'do_stuff: {x=}, {ti.try_number=}')
        if ti.try_number < 3:
            raise RuntimeError('')

    do_stuff.expand(x=do_stuff.expand(x=get_tasks()))
    do_stuff() >> do_stuff()  # this works as expected


dag = test_backfill()


if __name__ == '__main__':
    dag.cli()
```
```
airflow dags backfill test_backfill -s 2022-08-05 -e 2022-08-07 --rerun-failed-tasks
```
You can repeat the `backfill` command multiple times to slowly make progress through the DAG.  Things will eventually succeed (assuming the exception that triggers this bug stops being raised), but obviously this is a pain when trying to backfill a non-trivial number of DAG Runs.

### Operating System

CentOS Stream 8

### Versions of Apache Airflow Providers

None

### Deployment

Other

### Deployment details

Standalone

### Anything else

I was able to reproduce this both with SQLite + `SequentialExecutor` as well as with Postgres + `LocalExecutor`.

I haven't yet been able to reproduce this outside of `backfill` mode.

Possibly related since they mention the same exception text:
* #23533
* #23642

### Are you willing to submit PR?

- [ ] Yes I am willing to submit a PR!

### Code of Conduct

- [X] I agree to follow this project's [Code of Conduct](https://github.com/apache/airflow/blob/main/CODE_OF_CONDUCT.md)



Discussion:
After #25661 this may cause a different error. I will need to look into this further (not now).

To clarify, the downstream (mapped task) will never run correctly in any scenarios, since if the upstream raises an exception, there’s nothing the task can be expanded into. But the scheduler should handle this more gracefully.

Resolving patch:
diff --git a/airflow/models/mappedoperator.py b/airflow/models/mappedoperator.py
--- a/airflow/models/mappedoperator.py
+++ b/airflow/models/mappedoperator.py
@@ -623,7 +623,17 @@ def expand_mapped_task(self, run_id: str, *, session: Session) -> Tuple[Sequence
         from airflow.models.taskinstance import TaskInstance
         from airflow.settings import task_instance_mutation_hook
 
-        total_length = self._get_specified_expand_input().get_total_map_length(run_id, session=session)
+        total_length: Optional[int]
+        try:
+            total_length = self._get_specified_expand_input().get_total_map_length(run_id, session=session)
+        except NotFullyPopulated as e:
+            self.log.info(
+                "Cannot expand %r for run %s; missing upstream values: %s",
+                self,
+                run_id,
+                sorted(e.missing),
+            )
+            total_length = None
 
         state: Optional[TaskInstanceState] = None
         unmapped_ti: Optional[TaskInstance] = (
@@ -643,23 +653,32 @@ def expand_mapped_task(self, run_id: str, *, session: Session) -> Tuple[Sequence
         if unmapped_ti:
             # The unmapped task instance still exists and is unfinished, i.e. we
             # haven't tried to run it before.
-            if total_length < 1:
-                # If the upstream maps this to a zero-length value, simply marked the
-                # unmapped task instance as SKIPPED (if needed).
+            if total_length is None:
+                # If the map length cannot be calculated (due to unavailable
+                # upstream sources), fail the unmapped task.
+                unmapped_ti.state = TaskInstanceState.UPSTREAM_FAILED
+                indexes_to_map: Iterable[int] = ()
+            elif total_length < 1:
+                # If the upstream maps this to a zero-length value, simply mark
+                # the unmapped task instance as SKIPPED (if needed).
                 self.log.info(
                     "Marking %s as SKIPPED since the map has %d values to expand",
                     unmapped_ti,
                     total_length,
                 )
                 unmapped_ti.state = TaskInstanceState.SKIPPED
+                indexes_to_map = ()
             else:
                 # Otherwise convert this into the first mapped index, and create
                 # TaskInstance for other indexes.
                 unmapped_ti.map_index = 0
                 self.log.debug("Updated in place to become %s", unmapped_ti)
                 all_expanded_tis.append(unmapped_ti)
+                indexes_to_map = range(1, total_length)
             state = unmapped_ti.state
-            indexes_to_map = range(1, total_length)
+        elif not total_length:
+            # Nothing to fixup.
+            indexes_to_map = ()
         else:
             # Only create "missing" ones.
             current_max_mapping = (
@@ -682,17 +701,21 @@ def expand_mapped_task(self, run_id: str, *, session: Session) -> Tuple[Sequence
             ti.refresh_from_task(self)  # session.merge() loses task information.
             all_expanded_tis.append(ti)
 
+        # Coerce the None case to 0 -- these two are almost treated identically,
+        # except the unmapped ti (if exists) is marked to different states.
+        total_expanded_ti_count = total_length or 0
+
         # Set to "REMOVED" any (old) TaskInstances with map indices greater
         # than the current map value
         session.query(TaskInstance).filter(
             TaskInstance.dag_id == self.dag_id,
             TaskInstance.task_id == self.task_id,
             TaskInstance.run_id == run_id,
-            TaskInstance.map_index >= total_length,
+            TaskInstance.map_index >= total_expanded_ti_count,
         ).update({TaskInstance.state: TaskInstanceState.REMOVED})
 
         session.flush()
-        return all_expanded_tis, total_length
+        return all_expanded_tis, total_expanded_ti_count - 1
 
     def prepare_for_execution(self) -> "MappedOperator":
         # Since a mapped operator cannot be used for execution, and an unmapped


Test patch:
diff --git a/tests/jobs/test_backfill_job.py b/tests/jobs/test_backfill_job.py
--- a/tests/jobs/test_backfill_job.py
+++ b/tests/jobs/test_backfill_job.py
@@ -46,6 +46,7 @@
 from airflow.utils.session import create_session
 from airflow.utils.state import DagRunState, State, TaskInstanceState
 from airflow.utils.timeout import timeout
+from airflow.utils.trigger_rule import TriggerRule
 from airflow.utils.types import DagRunType
 from tests.models import TEST_DAGS_FOLDER
 from tests.test_utils.db import (
@@ -1640,7 +1641,7 @@ def test_backfilling_dags(self, dag_id, executor_name, session):
             assert ti.end_date is not None
 
     def test_mapped_dag_pre_existing_tis(self, dag_maker, session):
-        """If the DagRun already some mapped TIs, ensure that we re-run them successfully"""
+        """If the DagRun already has some mapped TIs, ensure that we re-run them successfully"""
         from airflow.decorators import task
         from airflow.operators.python import PythonOperator
 
@@ -1729,3 +1730,34 @@ def on_change_state(key, state, info=None):
                 dag_id=dr.dag_id, task_id='make_arg_lists', run_id='test', try_number=1, map_index=-1
             ),
         }
+
+    def test_mapped_dag_unexpandable(self, dag_maker, session):
+        with dag_maker(session=session) as dag:
+
+            @dag.task
+            def get_things():
+                return [1, 2]
+
+            @dag.task
+            def this_fails() -> None:
+                raise RuntimeError("sorry!")
+
+            @dag.task(trigger_rule=TriggerRule.ALL_DONE)
+            def consumer(a, b):
+                print(a, b)
+
+            consumer.expand(a=get_things(), b=this_fails())
+
+        executor = MockExecutor()
+        when = timezone.datetime(2022, 1, 1)
+        BackfillJob(dag=dag, start_date=when, end_date=when, donot_pickle=True, executor=executor).run()
+
+        (dr,) = DagRun.find(dag_id=dag.dag_id, execution_date=when, session=session)
+        assert dr.state == DagRunState.FAILED
+
+        # Check that every task has a start and end date
+        tis = {(ti.task_id, ti.map_index): ti for ti in dr.task_instances}
+        assert len(tis) == 3
+        tis[("get_things", -1)].state == TaskInstanceState.SUCCESS
+        tis[("this_fails", -1)].state == TaskInstanceState.FAILED
+        tis[("consumer", -1)].state == TaskInstanceState.UPSTREAM_FAILED
diff --git a/tests/models/test_dagrun.py b/tests/models/test_dagrun.py
--- a/tests/models/test_dagrun.py
+++ b/tests/models/test_dagrun.py
@@ -1329,7 +1329,8 @@ def test_ti_scheduling_mapped_zero_length(dag_maker, session):
     assert indices == [(-1, TaskInstanceState.SKIPPED)]
 
 
-def test_mapped_task_upstream_failed(dag_maker, session):
+@pytest.mark.parametrize("trigger_rule", [TriggerRule.ALL_DONE, TriggerRule.ALL_SUCCESS])
+def test_mapped_task_upstream_failed(dag_maker, session, trigger_rule):
     from airflow.operators.python import PythonOperator
 
     with dag_maker(session=session) as dag:
@@ -1341,7 +1342,11 @@ def make_list():
         def consumer(*args):
             print(repr(args))
 
-        PythonOperator.partial(task_id='consumer', python_callable=consumer).expand(op_args=make_list())
+        PythonOperator.partial(
+            task_id="consumer",
+            trigger_rule=trigger_rule,
+            python_callable=consumer,
+        ).expand(op_args=make_list())
 
     dr = dag_maker.create_dagrun()
     _, make_list_ti = sorted(dr.task_instances, key=lambda ti: ti.task_id)
diff --git a/tests/models/test_taskinstance.py b/tests/models/test_taskinstance.py
--- a/tests/models/test_taskinstance.py
+++ b/tests/models/test_taskinstance.py
@@ -3120,8 +3120,8 @@ def show(value):
         emit_ti.run()
 
         show_task = dag.get_task("show")
-        mapped_tis, num = show_task.expand_mapped_task(dag_run.run_id, session=session)
-        assert num == len(mapped_tis) == len(upstream_return)
+        mapped_tis, max_map_index = show_task.expand_mapped_task(dag_run.run_id, session=session)
+        assert max_map_index + 1 == len(mapped_tis) == len(upstream_return)
 
         for ti in sorted(mapped_tis, key=operator.attrgetter("map_index")):
             ti.refresh_from_task(show_task)
@@ -3154,8 +3154,8 @@ def show(number, letter):
             ti.run()
 
         show_task = dag.get_task("show")
-        mapped_tis, num = show_task.expand_mapped_task(dag_run.run_id, session=session)
-        assert len(mapped_tis) == 6 == num
+        mapped_tis, max_map_index = show_task.expand_mapped_task(dag_run.run_id, session=session)
+        assert max_map_index + 1 == len(mapped_tis) == 6
 
         for ti in sorted(mapped_tis, key=operator.attrgetter("map_index")):
             ti.refresh_from_task(show_task)
@@ -3193,8 +3193,8 @@ def show(a, b):
 
         show_task = dag.get_task("show")
         assert show_task.parse_time_mapped_ti_count is None
-        mapped_tis, num = show_task.expand_mapped_task(dag_run.run_id, session=session)
-        assert num == len(mapped_tis) == 4
+        mapped_tis, max_map_index = show_task.expand_mapped_task(dag_run.run_id, session=session)
+        assert max_map_index + 1 == len(mapped_tis) == 4
 
         for ti in sorted(mapped_tis, key=operator.attrgetter("map_index")):
             ti.refresh_from_task(show_task)
@@ -3217,9 +3217,9 @@ def show(a, b):
 
         show_task = dag.get_task("show")
         assert show_task.parse_time_mapped_ti_count == 6
-        mapped_tis, num = show_task.expand_mapped_task(dag_run.run_id, session=session)
+        mapped_tis, max_map_index = show_task.expand_mapped_task(dag_run.run_id, session=session)
         assert len(mapped_tis) == 0  # Expanded at parse!
-        assert num == 6
+        assert max_map_index == 5
 
         tis = (
             session.query(TaskInstance)
@@ -3265,8 +3265,8 @@ def cmds():
             ti.run()
 
         bash_task = dag.get_task("dynamic.bash")
-        mapped_bash_tis, num = bash_task.expand_mapped_task(dag_run.run_id, session=session)
-        assert num == 2 * 2
+        mapped_bash_tis, max_map_index = bash_task.expand_mapped_task(dag_run.run_id, session=session)
+        assert max_map_index == 3  # 2 * 2 mapped tasks.
         for ti in sorted(mapped_bash_tis, key=operator.attrgetter("map_index")):
             ti.refresh_from_task(bash_task)
             ti.run()