ephraimbuddy commented on code in PR #76420:
URL: https://github.com/apache/airflow/pull/76420#discussion_r4467806724
##########
airflow-core/src/airflow/models/dagrun.py:
##########
@@ -2378,12 +2378,10 @@ def _check_for_removed_or_restored_tasks(
producer_contexts=coordinates.producer_contexts(ti),
)
except NotFullyPopulated:
- # Not all upstreams finished, so we can't tell what should
be here. Remove everything.
+ # Not all upstreams finished, so the count isn't known
yet. Keep the expanded
+ # instances: _revise_map_indexes_if_mapped never brings
back a removed index, but
+ # trims or extends pending ones once the upstreams finish.
if ti.region_index >= 0:
- self.log.debug(
- "Removing the unmapped TI '%s' as the mapping
can't be resolved yet", ti
- )
- ti.state = TaskInstanceState.REMOVED
continue
Review Comment:
Done
##########
airflow-core/tests/unit/jobs/test_scheduler_job.py:
##########
@@ -6419,6 +6420,33 @@ def
test_verify_integrity_updates_dag_version_for_waiting_tasks(self, dag_maker,
assert downstream_ti.state == State.SCHEDULED
assert downstream_ti.dag_version_id == latest_version_id
+ def
test_verify_integrity_keeps_mapped_tis_while_cleared_producer_reruns(self,
dag_maker, session):
+ with dag_maker("test_cleared_mapped_producer", session=session):
+ size = EmptyOperator(task_id="size")
+ MockOperator.partial(task_id="mapped").expand(arg1=size.output)
+ dr = dag_maker.create_dagrun(state=DagRunState.RUNNING)
+ assert dr.bundle_version is None
+ size_ti = dr.get_task_instance("size", session=session)
+ size_ti.state = TaskInstanceState.SUCCESS
+ push_mapped_length(size_ti, [1, 2], session=session)
+ dr.task_instance_scheduling_decisions(session=session)
+ for ti in dr.get_task_instances(session=session):
+ ti.state = TaskInstanceState.SUCCESS
+ session.flush()
+ clear_task_instances(dr.get_task_instances(session=session), session)
+
+ with dag_maker("test_cleared_mapped_producer", session=session):
+ size = EmptyOperator(task_id="size")
+ MockOperator.partial(task_id="mapped").expand(arg1=size.output)
+ EmptyOperator(task_id="added")
+
+ self.job_runner = SchedulerJobRunner(job=Job())
+ self.job_runner._schedule_dag_run(dr, session)
+
+ tis = dr.get_task_instances(session=session)
+ assert "added" in {ti.task_id for ti in tis}
+ assert {ti.map_index: ti.state for ti in tis if ti.task_id ==
"mapped"} == {0: None, 1: None}
Review Comment:
Done
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
