Skip to content

fix: propagate rendered_map_index through the reschedule state path#68829

Open
kalluripradeep wants to merge 2 commits into
apache:mainfrom
kalluripradeep:fix/reschedule-rendered-map-index
Open

fix: propagate rendered_map_index through the reschedule state path#68829
kalluripradeep wants to merge 2 commits into
apache:mainfrom
kalluripradeep:fix/reschedule-rendered-map-index

Conversation

@kalluripradeep

Copy link
Copy Markdown
Contributor

What's the problem?

TIRescheduleStatePayload was the only intermediate-state payload missing rendered_map_index. As a result, when a mapped sensor sets context["map_index_template"] inside its poke() body and then raises AirflowRescheduleException to reschedule, the rendered_map_index value is never persisted to the database. The UI displays the raw integer map_index (e.g. 0, 1) on every up_for_reschedule row instead of the human-readable rendered label.

The final successful poke does show the rendered label (because the success path already serialises rendered_map_index via TaskState), so the column appears intermittently.

Fixes #67521

Root cause

The full data path for a reschedule is:

sensor.poke()
  ↓ raises AirflowRescheduleException
task_runner.run()
  ↓ except AirflowRescheduleException: builds RescheduleTask(…)   ← missing rendered_map_index
supervisor._handle_requests()
  ↓ client.task_instances.reschedule(id, msg)                      ← msg has no rendered_map_index
execution API PATCH /task-instances/{id}
  ↓ TIRescheduleStatePayload                                       ← model had no field
  ↓ query.values(state=…, next_method=None, next_kwargs=None)      ← _rendered_map_index never written

All other intermediate states (TIDeferredStatePayload, TIAwaitingInputStatePayload, TIRetryStatePayload, TITerminalStatePayload, TISuccessStatePayload) already carry rendered_map_index. The reschedule path was the only one left out.

Note: _render_map_index() is called in the except Exception: block (line ~1576 of task_runner.py) before re-raising AirflowRescheduleException, so ti.rendered_map_index is correctly populated by the time we build RescheduleTask. We only need to pass it through.

Changes

File Change
airflow-core/…/datamodels/taskinstance.py Add rendered_map_index: str | None = None to TIRescheduleStatePayload (server-side model)
task-sdk/…/api/datamodels/_generated.py Same field on the client-side generated model
task-sdk/…/execution_time/task_runner.py Pass rendered_map_index=ti.rendered_map_index when constructing RescheduleTask
task-sdk/…/execution_time/supervisor.py Set self._rendered_map_index = msg.rendered_map_index in both RescheduleTask handler branches
airflow-core/…/routes/task_instances.py Include _rendered_map_index=ti_patch_payload.rendered_map_index in the UP_FOR_RESCHEDULE query
airflow-core/…/versions/v2026_06_30.py Cadwyn migration AddRenderedMapIndexToReschedulePayload
airflow-core/…/versions/__init__.py Register the migration in the 2026-06-30 version bundle
task-sdk/tests/…/test_task_runner.py Regression test test_run_reschedule_includes_rendered_map_index

Testing

New regression test verifies that when a sensor sets map_index_template and rescheduled, the RescheduleTask message carries rendered_map_index:

def test_run_reschedule_includes_rendered_map_index(create_runtime_ti, mock_supervisor_comms):
    ...
    assert msg.rendered_map_index is not None

When a mapped sensor raises AirflowRescheduleException the rendered_map_index
(the human-readable label derived from map_index_template) was silently dropped,
leaving the DB column NULL on every UP_FOR_RESCHEDULE transition. All other
intermediate-state payloads (deferred, awaiting-input, retry, terminal) already
carry rendered_map_index; this fix closes the gap for the reschedule path.

Changes:
- Add `rendered_map_index` field to `TIRescheduleStatePayload` (server + client)
- Pass `ti.rendered_map_index` when constructing `RescheduleTask` in task_runner
- Store it in `_rendered_map_index` in both supervisor RescheduleTask handlers
- Write it to the DB in the `/task-instances/{id}/state` route handler
- Regenerate schema.json snapshot
- Add Cadwyn version migration (v2026-06-30) for backward compatibility
- Add regression test

Closes apache#67521
MyPy cannot infer the return type of `anyio.from_thread.run(request.form)`
without an explicit annotation. Add `FormData` to a TYPE_CHECKING block
(satisfying ruff TC002) and annotate both call sites with `# type: ignore[arg-type]`
to silence the incompatible-argument error that exists on main.
@kalluripradeep kalluripradeep force-pushed the fix/reschedule-rendered-map-index branch from a210c23 to 3dc8a05 Compare June 22, 2026 12:34
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Projects

None yet

Development

Successfully merging this pull request may close these issues.

[Airflow 3] map_index_template not rendered for sensors in reschedule mode (regression vs Airflow 2)

1 participant