Render map_index_template during task running state

airflow-core/src/airflow/api_fastapi/execution_api/routes/task_instances.py

@@ -692,6 +692,43 @@ def ti_put_rtif(
     return {"message": "Rendered task instance fields successfully set"}
 
 
+@ti_id_router.patch(
+    "/{task_instance_id}/rendered-map-index",
+    status_code=status.HTTP_204_NO_CONTENT,
+    responses={
+        status.HTTP_404_NOT_FOUND: {"description": "Task Instance not found"},
+        status.HTTP_422_UNPROCESSABLE_ENTITY: {"description": "Invalid rendered_map_index value"},
+    },
+)
+def ti_patch_rendered_map_index(
+    task_instance_id: UUID,
+    rendered_map_index: Annotated[str, Body()],
+    session: SessionDep,
+):
+    """Update rendered_map_index for a task instance, sent by the worker during task execution."""
+    ti_id_str = str(task_instance_id)
+    bind_contextvars(ti_id=ti_id_str)
+
+    if not rendered_map_index:
+        log.error("rendered_map_index cannot be empty")
+        raise HTTPException(
+            status_code=status.HTTP_422_UNPROCESSABLE_ENTITY,
+            detail="rendered_map_index cannot be empty",
+        )
+
+    log.debug("Updating rendered_map_index", length=len(rendered_map_index))
+
+    query = update(TI).where(TI.id == ti_id_str).values(rendered_map_index=rendered_map_index)
+    result = session.execute(query)
+
+    if result.rowcount == 0:
+        log.error("Task Instance not found")
+        raise HTTPException(
+            status_code=status.HTTP_404_NOT_FOUND,
+            detail="Task Instance not found",
+        )
+
+
 @ti_id_router.get(
     "/{task_instance_id}/previous-successful-dagrun",
     status_code=status.HTTP_200_OK,

airflow-core/tests/unit/api_fastapi/execution_api/versions/head/test_task_instances.py

@@ -2479,3 +2479,59 @@ def test_ti_run_with_null_conf(self, client, session, create_task_instance):
         assert response.status_code == 200, f"Response: {response.text}"
         context = response.json()
         assert context["dag_run"]["conf"] is None
+
+
+class TestTIPatchRenderedMapIndex:
+    def setup_method(self):
+        clear_db_runs()
+
+    def teardown_method(self):
+        clear_db_runs()
+
+    def test_ti_patch_rendered_map_index(self, client, session, create_task_instance):
+        """Test updating rendered_map_index for a task instance."""
+        ti = create_task_instance(
+            task_id="test_ti_patch_rendered_map_index",
+            state=State.RUNNING,
+            session=session,
+        )
+        session.commit()
+
+        rendered_map_index = "custom_label_123"
+        response = client.patch(
+            f"/execution/task-instances/{ti.id}/rendered-map-index",
+            json=rendered_map_index,
+        )
+
+        assert response.status_code == 204
+        assert response.text == ""
+
+        session.expire_all()
+        ti = session.get(TaskInstance, ti.id)
+        assert ti.rendered_map_index == rendered_map_index
+
+    def test_ti_patch_rendered_map_index_not_found(self, client, session):
+        """Test 404 error when task instance does not exist."""
+        fake_id = str(uuid4())
+        response = client.patch(
+            f"/execution/task-instances/{fake_id}/rendered-map-index",
+            json="test",
+        )
+
+        assert response.status_code == 404
+
+    def test_ti_patch_rendered_map_index_empty_string(self, client, session, create_task_instance):
+        """Test that empty string is accepted (clears the rendered_map_index)."""
+        ti = create_task_instance(
+            task_id="test_ti_patch_rendered_map_index_empty",
+            state=State.RUNNING,
+            session=session,
+        )
+        session.commit()
+
+        response = client.patch(
+            f"/execution/task-instances/{ti.id}/rendered-map-index",
+            json="",
+        )
+
+        assert response.status_code == 422

airflow-core/tests/unit/dag_processing/test_processor.py

@@ -1803,6 +1803,7 @@ def get_type_names(union_type):
             "CreateHITLDetailPayload",
             "UpdateHITLDetail",
             "GetHITLDetailResponse",
+            "SetRenderedMapIndex",
         }
 
         in_task_runner_but_not_in_dag_processing_process = {

airflow-core/tests/unit/jobs/test_triggerer_job.py

@@ -1213,6 +1213,7 @@ def get_type_names(union_type):
             "TriggerDagRun",
             "ResendLoggingFD",
             "CreateHITLDetailPayload",
+            "SetRenderedMapIndex",
         }
 
         in_task_but_not_in_trigger_runner = {

task-sdk/src/airflow/sdk/api/client.py

@@ -271,6 +271,11 @@ def set_rtif(self, id: uuid.UUID, body: dict[str, str]) -> OKResponse:
         # decouple from the server response string
         return OKResponse(ok=True)
 
+    def set_rendered_map_index(self, id: uuid.UUID, rendered_map_index: str) -> OKResponse:
+        """Set rendered_map_index for a task instance via the API server."""
+        self.client.patch(f"task-instances/{id}/rendered-map-index", json=rendered_map_index)
+        return OKResponse(ok=True)
+
     def get_previous_successful_dagrun(self, id: uuid.UUID) -> PrevSuccessfulDagRunResponse:
         """
         Get the previous successful dag run for a given task instance.

task-sdk/src/airflow/sdk/execution_time/comms.py

@@ -792,6 +792,13 @@ class SetRenderedFields(BaseModel):
     type: Literal["SetRenderedFields"] = "SetRenderedFields"
 
 
+class SetRenderedMapIndex(BaseModel):
+    """Payload for setting rendered_map_index for a task instance."""
+
+    rendered_map_index: str
+    type: Literal["SetRenderedMapIndex"] = "SetRenderedMapIndex"
+
+
 class TriggerDagRun(TriggerDAGRunPayload):
     dag_id: str
     run_id: Annotated[str, Field(title="Dag Run Id")]
@@ -934,6 +941,7 @@ class MaskSecret(BaseModel):
     | RescheduleTask
     | RetryTask
     | SetRenderedFields
+    | SetRenderedMapIndex
     | SetXCom
     | SkipDownstreamTasks
     | SucceedTask

task-sdk/src/airflow/sdk/execution_time/supervisor.py

@@ -103,6 +103,7 @@
     RetryTask,
     SentFDs,
     SetRenderedFields,
+    SetRenderedMapIndex,
     SetXCom,
     SkipDownstreamTasks,
     StartupDetails,
@@ -1263,6 +1264,8 @@ def _handle_request(self, msg: ToSupervisor, log: FilteringBoundLogger, req_id:
             self.client.variables.set(msg.key, msg.value, msg.description)
         elif isinstance(msg, SetRenderedFields):
             self.client.task_instances.set_rtif(self.id, msg.rendered_fields)
+        elif isinstance(msg, SetRenderedMapIndex):
+            self.client.task_instances.set_rendered_map_index(self.id, msg.rendered_map_index)
         elif isinstance(msg, GetAssetByName):
             asset_resp = self.client.assets.get(name=msg.name)
             if isinstance(asset_resp, AssetResponse):

task-sdk/src/airflow/sdk/execution_time/task_runner.py

@@ -80,6 +80,7 @@
     RetryTask,
     SentFDs,
     SetRenderedFields,
+    SetRenderedMapIndex,
     SkipDownstreamTasks,
     StartupDetails,
     SucceedTask,
@@ -811,6 +812,19 @@ def _prepare(ti: RuntimeTaskInstance, log: Logger, context: Context) -> ToSuperv
         # so that we do not call the API unnecessarily
         SUPERVISOR_COMMS.send(msg=SetRenderedFields(rendered_fields=rendered_fields))
 
+    # Try to render map_index_template early with available context (will be re-rendered after execution)
+    # This provides a partial label during task execution for templates using pre-execution context
+    # If rendering fails here, we suppress the error since it will be re-rendered after execution
+    try:
+        if rendered_map_index := _render_map_index(context, ti=ti, log=log):
+            ti.rendered_map_index = rendered_map_index
+            log.debug("Sending early rendered map index", length=len(rendered_map_index))
+            SUPERVISOR_COMMS.send(msg=SetRenderedMapIndex(rendered_map_index=rendered_map_index))
+    except Exception:
+        log.debug(
+            "Early rendering of map_index_template failed, will retry after task execution", exc_info=True
+        )
+
     _validate_task_inlets_and_outlets(ti=ti, log=log)
 
     try:
@@ -927,10 +941,20 @@ def _on_term(signum, frame):
 
                 # If the task failed, swallow rendering error so it doesn't mask the main error.
                 with contextlib.suppress(jinja2.TemplateSyntaxError, jinja2.UndefinedError):
+                    previous_rendered_map_index = ti.rendered_map_index
                     ti.rendered_map_index = _render_map_index(context, ti=ti, log=log)
+                    # Send update only if value changed (e.g., user set context variables during execution)
+                    if ti.rendered_map_index and ti.rendered_map_index != previous_rendered_map_index:
+                        SUPERVISOR_COMMS.send(
+                            msg=SetRenderedMapIndex(rendered_map_index=ti.rendered_map_index)
+                        )
                 raise
             else:  # If the task succeeded, render normally to let rendering error bubble up.
+                previous_rendered_map_index = ti.rendered_map_index
                 ti.rendered_map_index = _render_map_index(context, ti=ti, log=log)
+                # Send update only if value changed (e.g., user set context variables during execution)
+                if ti.rendered_map_index and ti.rendered_map_index != previous_rendered_map_index:
+                    SUPERVISOR_COMMS.send(msg=SetRenderedMapIndex(rendered_map_index=ti.rendered_map_index))
 
         _push_xcom_if_needed(result, ti, log)
 
@@ -1322,9 +1346,10 @@ def _render_map_index(context: Context, ti: RuntimeTaskInstance, log: Logger) ->
     """Render named map index if the Dag author defined map_index_template at the task level."""
     if (template := context.get("map_index_template")) is None:
         return None
+    log.debug("Rendering map_index_template", template_length=len(template))
     jinja_env = ti.task.dag.get_template_env()
     rendered_map_index = jinja_env.from_string(template).render(context)
-    log.info("Map index rendered as %s", rendered_map_index)
+    log.debug("Map index rendered", length=len(rendered_map_index))
     return rendered_map_index
 
 

task-sdk/tests/task_sdk/api/test_client.py

@@ -448,6 +448,24 @@ def handle_request(request: httpx.Request) -> httpx.Response:
 
         assert result == OKResponse(ok=True)
 
+    def test_taskinstance_set_rendered_map_index_success(self):
+        TI_ID = uuid6.uuid7()
+        rendered_map_index = "Label: task_1"
+
+        def handle_request(request: httpx.Request) -> httpx.Response:
+            if request.url.path == f"/task-instances/{TI_ID}/rendered-map-index":
+                actual_body = json.loads(request.read())
+                assert request.method == "PATCH"
+                # Body should be the string directly, not wrapped in JSON
+                assert actual_body == rendered_map_index
+                return httpx.Response(status_code=204)
+            return httpx.Response(status_code=400, json={"detail": "Bad Request"})
+
+        client = make_client(transport=httpx.MockTransport(handle_request))
+        result = client.task_instances.set_rendered_map_index(id=TI_ID, rendered_map_index=rendered_map_index)
+
+        assert result == OKResponse(ok=True)
+
     def test_get_count_basic(self):
         """Test basic get_count functionality with just dag_id."""
 

task-sdk/tests/task_sdk/execution_time/test_supervisor.py

@@ -104,6 +104,7 @@
     RetryTask,
     SentFDs,
     SetRenderedFields,
+    SetRenderedMapIndex,
     SetXCom,
     SkipDownstreamTasks,
     SucceedTask,
@@ -1604,6 +1605,15 @@ class RequestTestCase:
         ),
         test_id="set_rtif",
     ),
+    RequestTestCase(
+        message=SetRenderedMapIndex(rendered_map_index="Label: task_1"),
+        client_mock=ClientMock(
+            method_path="task_instances.set_rendered_map_index",
+            args=(TI_ID, "Label: task_1"),
+            response=OKResponse(ok=True),
+        ),
+        test_id="set_rendered_map_index",
+    ),
     RequestTestCase(
         message=SucceedTask(
             end_date=timezone.parse("2024-10-31T12:00:00Z"), rendered_map_index="test success task"

task-sdk/tests/task_sdk/execution_time/test_task_runner.py

@@ -1149,6 +1149,28 @@ def test_execute_failed_task_with_rendered_map_index(create_runtime_ti, mock_sup
     assert ti.rendered_map_index == "Hello! test_run"
 
 
+def test_rendered_map_index_updates_sent_progressively(create_runtime_ti, mock_supervisor_comms):
+    """Test that rendered_map_index is rendered and potentially updated after execution."""
+
+    def test_function(ti):
+        # Simulate setting a context variable during execution
+        ti.xcom_push(key="execution_result", value="completed")
+        return "test function"
+
+    task = PythonOperator(
+        task_id="test_task",
+        python_callable=test_function,
+        map_index_template="Label: {{ task.task_id }}",
+    )
+
+    ti = create_runtime_ti(task=task, dag_id="dag_with_progressive_map_index")
+
+    run(ti, ti.get_template_context(), log=mock.MagicMock())
+
+    # Verify that rendered_map_index is set (existing behavior)
+    assert ti.rendered_map_index == "Label: test_task"
+
+
 class TestRuntimeTaskInstance:
     def test_get_context_without_ti_context_from_server(self, mocked_parse, make_ti_context):
         """Test get_template_context without ti_context_from_server."""