Add map index to task logs for mapped task runs.

CHANGELOG.md

@@ -13,6 +13,7 @@ These changes are available in the [master branch](https://github.com/PrefectHQ/
 - Added serializer for `RemoteDaskEnvironment` - [#2369](https://github.com/PrefectHQ/prefect/issues/2369)
 - `server start` CLI command now defaults to image build based on current Prefect installation version - [#2375](https://github.com/PrefectHQ/prefect/issues/2375)
 - Add option to set `executor_kwargs` on `KubernetesJobEnvironment` and `FargateTaskEnvironment` - [#2258](https://github.com/PrefectHQ/prefect/issues/2258)
+- Add map index to task logs for mapped task runs - [#2402](https://github.com/PrefectHQ/prefect/pull/2402)
 - Agents can now register themselves with Cloud for better management - [#2312](https://github.com/PrefectHQ/prefect/issues/2312)
 - Adding support for `environment`, `secrets`, and `mountPoints` via configurable `containerDefinitions` to the Fargate Agent - [#2397](https://github.com/PrefectHQ/prefect/pull/2397)
 
@@ -36,6 +37,7 @@ These changes are available in the [master branch](https://github.com/PrefectHQ/
 
 - [Nelson Cornet](https://github.com/sk4la)
 - [Braun Reyes](https://github.com/braunreyes)
+- [Fraznist](https://github.com/Fraznist)
 
 ## 0.10.4 <Badge text="beta" type="success"/>
 

src/prefect/engine/task_runner.py

@@ -169,6 +169,13 @@ def initialize_run(  # type: ignore
             task_slug=self.task.slug,
         )
         context.setdefault("checkpointing", config.flows.checkpointing)
+
+        map_index = context.get("map_index", None)
+        if isinstance(map_index, int):
+            self.task.logger = prefect.utilities.logging.get_logger(
+                "Task: {}[{}]".format(self.task.name, map_index)
+            )
+
         context.update(logger=self.task.logger)
 
         return TaskRunnerInitializeResult(state=state, context=context)

tests/engine/test_task_runner.py

@@ -2265,3 +2265,24 @@ def run(self):
 
     logs = [r.message for r in caplog.records]
     assert "TEST_HERE" not in logs
+
+
+def test_task_runner_logs_map_index_for_mapped_tasks(caplog):
+    class MyTask(Task):
+        def run(self):
+            map_index = prefect.context.get("map_index")
+            self.logger.info("{}".format(map_index))
+            print(map_index)
+
+    task = MyTask()
+    edge = Edge(Task(), task, mapped=True)
+    new_state = TaskRunner(task=task).run(
+        state=None, upstream_states={edge: Success(result=Result(list(range(10))))}
+    )
+
+    logs = [r.message for r in caplog.records if "prefect.Task:" in r.message]
+    task_name = task.name
+    for line in logs:
+        msg = line.split("INFO")[1]
+        logged_map_index = msg[-1]
+        assert msg.count(logged_map_index) == 2