Fix local retries in Cloud to always run in-process

CHANGELOG.md

@@ -19,7 +19,8 @@ These changes are available in the [master branch](https://github.com/PrefectHQ/
 
 ### Fixes
 
-- Fixes issue with Docker storage not respecting user-provided image names - [#1335](https://github.com/PrefectHQ/prefect/pull/1335)
+- Fix issue with Docker storage not respecting user-provided image names - [#1335](https://github.com/PrefectHQ/prefect/pull/1335)
+- Fix issue with local retries in Cloud not always running in-process - [#1348](https://github.com/PrefectHQ/prefect/pull/1348)
 
 ### Deprecations
 

src/prefect/engine/cloud/task_runner.py

@@ -239,6 +239,7 @@ def run(
         Returns:
             - `State` object representing the final post-run state of the Task
         """
+        context = context or {}
         end_state = super().run(
             state=state,
             upstream_states=upstream_states,
@@ -253,6 +254,14 @@ def run(
                 (end_state.start_time - pendulum.now("utc")).total_seconds(), 0
             )
             time.sleep(naptime)
+
+            # currently required as context has reset to its original state
+            task_run_info = self.client.get_task_run_info(
+                flow_run_id=context.get("flow_run_id", ""),
+                task_id=context.get("task_id", ""),
+                map_index=context.get("map_index"),
+            )
+            context.update(task_run_version=task_run_info.version)  # type: ignore
             return self.run(
                 state=end_state,
                 upstream_states=upstream_states,

tests/engine/cloud/test_cloud_task_runner.py

@@ -622,15 +622,19 @@ def noop():
             global_list.append(0)
             raise ValueError("oops")
 
+    client.get_task_run_info.side_effect = [MagicMock(version=i) for i in range(4, 7)]
     res = CloudTaskRunner(task=noop).run(
+        context={"task_run_version": 1},
         state=None,
         upstream_states={},
         executor=prefect.engine.executors.LocalExecutor(),
     )
 
     ## assertions
     assert res.is_successful()
-    assert client.get_task_run_info.call_count == 0  # never called
+    assert client.get_task_run_info.call_count == 1  # called once on the retry
     assert (
         client.set_task_run_state.call_count == 5
     )  # Pending -> Running -> Failed -> Retrying -> Running -> Success
+    versions = [call[1]["version"] for call in client.set_task_run_state.call_args_list]
+    assert versions == [1, 2, 3, 4, 5]