Fix map_index_template not rendered for sensors in reschedule mode - #67867
Fix map_index_template not rendered for sensors in reschedule mode#67867GayathriSrividya wants to merge 1 commit into
Conversation
b606b2d to
7cd80e8
Compare
|
Did you check #67521 (comment) ? Were you able to reproduce the behaviour? Could you post screenshots and/or videos of before/after |
7cd80e8 to
d9ef5af
Compare
I checked that comment and re-verified the behavior. Yes, I was able to reproduce the issue on an affected release ref: Fresh clone at main Main includes rendered_map_index in the reschedule path and does not show the same issue. If you want, I can also push a short reproduction note to the PR description with exact commands used. |
potiuk
added
the
ready for maintainer review
Sure, please post the screenshots when you do |
kaxil
removed
the
ready for maintainer review
Reproduction script: from __future__ import annotations
from datetime import datetime
from airflow.sdk import DAG, task
from airflow.sdk.bases.sensor import PokeReturnValue
with DAG(
dag_id="sensor_reschedule_map_index_demo",
start_date=datetime(2025, 1, 1),
schedule=None,
catchup=False,
):
@task
def fan_out() -> list[dict]:
return [{"id": "alpha"}, {"id": "beta"}, {"id": "gamma"}]
@task.sensor(
poke_interval=10,
timeout=600,
mode="reschedule",
map_index_template="{{ map_index_template }}",
)
def poll(item: dict) -> PokeReturnValue:
from airflow.sdk import get_current_context
ctx = get_current_context()
ctx["map_index_template"] = f"poll_{item['id']}"
return PokeReturnValue(is_done=False)
poll.expand(item=fan_out())What I did to reproduce: I ran two Airflow 3.1.2 instances locally at the same time: one without the fix on Before — While the sensors are waiting to re-poke, the Map Index column shows After — The same task instances show Why this happens: When a sensor reschedules itself, Airflow was not saving the rendered map index to the database. Other paths already persist it, but the reschedule path was missing it. This fix adds that in three places:
|
| ] | ||
| reschedule_date: UtcDateTime | ||
| end_date: UtcDateTime | ||
| rendered_map_index: str | None = None |
There was a problem hiding this comment.
The PR description lists an AddRenderedMapIndexToReschedulePayload migration in v2026_06_30.py, but I don't see it in the diff. New payload fields need a Cadwyn VersionChange so older API version schemas stay accurate -- AddRenderedMapIndexField in v2025_04_28.py is the precedent for exactly this field on the other payloads. Since 2026-06-30 isn't released yet, schema(TIRescheduleStatePayload).field("rendered_map_index").didnt_exist in that file should do it. (Static checks being green doesn't cover this btw: the check-execution-api-versions prek hook skips the schema diff when run with --all-files.)
| updated_state = TaskInstanceState.UP_FOR_RESCHEDULE | ||
| query = query.values(state=updated_state, next_method=None, next_kwargs=None) | ||
| reschedule_values: dict[str, Any] = {"state": updated_state, "next_method": None, "next_kwargs": None} | ||
| if ti_patch_payload.rendered_map_index is not None: |
There was a problem hiding this comment.
Nothing exercises this write. Can you extend test_ti_update_state_to_reschedule in airflow-core/tests/unit/api_fastapi/execution_api/versions/head/test_task_instances.py to send rendered_map_index in the payload and assert it lands on the TI? The new SDK test only checks the message construction on the worker side.
| reschedule_date=reschedule.reschedule_date, end_date=datetime.now(tz=timezone.utc) | ||
| reschedule_date=reschedule.reschedule_date, | ||
| end_date=datetime.now(tz=timezone.utc), | ||
| rendered_map_index=ti.rendered_map_index, |
There was a problem hiding this comment.
Worth stating in the description: on main this doesn't change what users see. #57208 (shipped in 3.2.0) already persists the label on this path: the except Exception block above re-renders and sends SetRenderedMapIndex (AirflowRescheduleException is an Exception), and the supervisor writes it to the DB before the reschedule call goes out. That's why the issue reproduces on your 3.1.2 build but not on main, as kevinhongzl found. Carrying the value in the reschedule payload itself is still reasonable for consistency with the other state payloads, but #67521 as reported only affects 3.0.x-3.1.x, so actually fixing affected users is a backport question rather than this merge.
| sensor.poke = Mock(return_value=False) | ||
|
|
||
| date1 = timezone.utcnow() | ||
| time_machine.move_to(date1, tick=False) |
There was a problem hiding this comment.
date1 and the time freeze aren't used by any assertion. Either assert msg.end_date == date1 (the freeze makes that deterministic) or drop these two lines and the time_machine fixture.
When a sensor is rescheduled, the rendered_map_index was not being persisted to the database. The TIRescheduleStatePayload was missing the rendered_map_index field that all other terminal state payloads have (retry, succeed, skip). The reschedule route handler creates a brand-new UPDATE query for the reschedule-specific columns, discarding the initial one — so the field had to be explicitly added to that new query as well. Changes: - Add rendered_map_index to TIRescheduleStatePayload server-side datamodel - Update reschedule route handler to write rendered_map_index to the DB - Add Cadwyn version migration AddRenderedMapIndexToReschedulePayload - Update _generated.py (task-sdk client model) to include the field - Pass rendered_map_index in RescheduleTask message from task runner - Save _rendered_map_index from RescheduleTask in supervisor - Add test for rendered_map_index in reschedule mode closes: apache#67521
ae8ac46 to
a63ce3c
Compare
|
Thanks for the detailed review @kaxil. This PR has accumulated several rebases and merge commits and I've lost track of the version migration conflict resolution. I'll close this and open a clean PR with all the feedback addressed properly (Cadwyn migration, server-side test, dead test code fix, and the backport to v3-1-test). Sorry for the noise. |
3 participants
closes: #67521
When a sensor is configured with `map_index_template` and runs in reschedule mode, the rendered label was not being persisted to the database on each reschedule. As a result, the task instance's `rendered_map_index` column was left NULL, causing the UI to display the raw numeric `map_index` (e.g. "0", "1") instead of the human-readable rendered label.
Root cause: `TIRescheduleStatePayload` was missing the `rendered_map_index` field that all other terminal/retry state payloads have (`TIRetryStatePayload`, `TaskState`). The reschedule route handler builds a brand-new `UPDATE` query for the reschedule-specific columns (discarding the initial shared one), so the field had to be explicitly included in that new query as well.
Changes: