Re: [PR] Set the state of a dag run's task instances with one flush, not one per task instance [airflow]

["rjgoyln (via GitHub)" <gi...@apache.org>]

rjgoyln commented on code in PR #72185:
URL: https://github.com/apache/airflow/pull/72185#discussion_r4187024501


##########
airflow-core/tests/unit/api/common/test_mark_tasks.py:
##########
@@ -61,6 +64,75 @@ def test_set_dag_run_state_to_failed(dag_maker: DagMaker[SerializedDAG]):
     assert "teardown" not in task_dict
 
 
+def test_set_dag_run_state_to_failed_skips_pending_tis_in_bulk(dag_maker: DagMaker[SerializedDAG]):
+    """Pending TIs that never started are skipped by one UPDATE; one that did keeps its start_date."""
+    with dag_maker("TEST_DAG_1") as dag:
+        EmptyOperator(task_id="never_started")
+        EmptyOperator(task_id="never_started_2")
+        EmptyOperator(task_id="up_for_retry")
+    dr = dag_maker.create_dagrun()
+    session = dag_maker.session
+    started_at = timezone.datetime(2024, 1, 1, 12, 0, 0)
+    for ti in dr.get_task_instances(session=session):
+        if ti.task_id == "up_for_retry":
+            ti.state = TaskInstanceState.UP_FOR_RETRY
+            ti.start_date = started_at
+    session.flush()
+
+    updated_tis, killed_tis = set_dag_run_state_to_failed(
+        dag=dag, run_id=dr.run_id, commit=True, session=session
+    )

Review Comment:
   One small test gap: this test also passes on `main`, since the per-TI `set_state` path produces the same states, dates, and durations.
   
   Would it make sense to use `assert_queries_count(5)` here to pin the single flush and catch a regression to per-TI flushing?
   
   ```suggestion
       with assert_queries_count(5):
           updated_tis, killed_tis = set_dag_run_state_to_failed(
               dag=dag, run_id=dr.run_id, commit=True, session=session
           )
   ```
   



##########
airflow-core/tests/unit/api/common/test_mark_tasks.py:
##########
@@ -61,6 +64,75 @@ def test_set_dag_run_state_to_failed(dag_maker: DagMaker[SerializedDAG]):
     assert "teardown" not in task_dict
 
 
+def test_set_dag_run_state_to_failed_skips_pending_tis_in_bulk(dag_maker: DagMaker[SerializedDAG]):
+    """Pending TIs that never started are skipped by one UPDATE; one that did keeps its start_date."""
+    with dag_maker("TEST_DAG_1") as dag:
+        EmptyOperator(task_id="never_started")
+        EmptyOperator(task_id="never_started_2")
+        EmptyOperator(task_id="up_for_retry")
+    dr = dag_maker.create_dagrun()
+    session = dag_maker.session
+    started_at = timezone.datetime(2024, 1, 1, 12, 0, 0)
+    for ti in dr.get_task_instances(session=session):
+        if ti.task_id == "up_for_retry":
+            ti.state = TaskInstanceState.UP_FOR_RETRY
+            ti.start_date = started_at
+    session.flush()
+
+    updated_tis, killed_tis = set_dag_run_state_to_failed(
+        dag=dag, run_id=dr.run_id, commit=True, session=session
+    )
+    session.flush()
+
+    assert killed_tis == []
+    assert {ti.task_id for ti in updated_tis} == {"never_started", "never_started_2", "up_for_retry"}
+    # The returned objects reflect the new state without a reload...
+    assert all(ti.state == TaskInstanceState.SKIPPED for ti in updated_tis)
+    # ...and so does the database.
+    session.expire_all()
+    task_dict = {ti.task_id: ti for ti in dr.get_task_instances(session=session)}
+    for task_id in ("never_started", "never_started_2"):
+        ti = task_dict[task_id]
+        assert ti.state == TaskInstanceState.SKIPPED
+        assert ti.start_date is not None
+        assert ti.end_date == ti.start_date
+        assert ti.duration == 0
+    ti = task_dict["up_for_retry"]
+    assert ti.state == TaskInstanceState.SKIPPED
+    assert ti.start_date == started_at
+    assert ti.end_date is not None
+    assert ti.duration == pytest.approx((ti.end_date - started_at).total_seconds())
+
+
+def test_set_dag_run_state_to_failed_mapped_task_only_fails_running_map_indexes(
+    dag_maker: DagMaker[SerializedDAG],
+):
+    """Only the running map indexes fail; the pending siblings of the same mapped task are skipped."""
+    with dag_maker("TEST_DAG_1") as dag:
+        MockOperator.partial(task_id="mapped").expand(arg2=[1, 2, 3])

Review Comment:
   Could we also pin the state of finished siblings here? I noticed that on main, they can be overwritten as well: a finished success sibling becomes failed when marking the run failed, and a failed sibling becomes success when marking it successful. With this change, they keep their existing state as expected.
   
   Would it make sense to add a fourth map index that is already success and assert that it remains success?



-- 
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: commits-unsubscribe@airflow.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org