[workflow] Fix workflow recovery issue due to a bug of dynamic output
diff --git a/python/ray/workflow/tests/test_basic_workflows.py b/python/ray/workflow/tests/test_basic_workflows.py
index f5dc73f160e64..971c6f5724677 100644
--- a/python/ray/workflow/tests/test_basic_workflows.py
+++ b/python/ray/workflow/tests/test_basic_workflows.py
@@ -282,6 +282,29 @@ def f1(n):
     assert isinstance(err, ValueError)
 
 
+def test_dynamic_output(workflow_start_regular_shared):
+    @workflow.step
+    def exponential_fail(k, n):
+        if n > 0:
+            if n < 3:
+                raise Exception("Failed intentionally")
+            return exponential_fail.options(name=f"step_{n}").step(
+                k * 2, n - 1)
+        return k
+
+    # When workflow fails, the dynamic output should points to the
+    # latest successful step.
+    try:
+        exponential_fail.options(name="step_0").step(
+            3, 10).run(workflow_id="dynamic_output")
+    except Exception:
+        pass
+    from ray.workflow.workflow_storage import get_workflow_storage
+    wf_storage = get_workflow_storage(workflow_id="dynamic_output")
+    result = wf_storage.inspect_step("step_0")
+    assert result.output_step_id == "step_3"
+
+
 if __name__ == "__main__":
     import sys
     sys.exit(pytest.main(["-v", __file__]))
diff --git a/python/ray/workflow/workflow_storage.py b/python/ray/workflow/workflow_storage.py
index b0df708e08cd1..b6a67292dee02 100644
--- a/python/ray/workflow/workflow_storage.py
+++ b/python/ray/workflow/workflow_storage.py
@@ -134,6 +134,7 @@ def save_step_output(self, step_id: StepID, ret: Union[Workflow, Any], *,
             outer_most_step_id: See WorkflowStepContext.
         """
         tasks = []
+        dynamic_output_id = None
         if isinstance(ret, Workflow):
             # This workflow step returns a nested workflow.
             assert step_id != ret.step_id
@@ -154,9 +155,6 @@ def save_step_output(self, step_id: StepID, ret: Union[Workflow, Any], *,
                 # tasks.append(self._put(self._key_step_output(step_id), ret))
                 dynamic_output_id = step_id
                 # TODO (yic): Delete exception file
-                tasks.append(
-                    self._update_dynamic_output(outer_most_step_id,
-                                                dynamic_output_id))
             else:
                 assert ret is None
                 promise = serialization.dump_to_storage(
@@ -166,8 +164,17 @@ def save_step_output(self, step_id: StepID, ret: Union[Workflow, Any], *,
                 # tasks.append(
                 #     self._put(self._key_step_exception(step_id), exception))
 
+        # Finish checkpointing.
         asyncio_run(asyncio.gather(*tasks))
 
+        # NOTE: if we update the dynamic output before
+        # finishing checkpointing, then during recovery, the dynamic could
+        # would point to a checkpoint that does not exist.
+        if dynamic_output_id is not None:
+            asyncio_run(
+                self._update_dynamic_output(outer_most_step_id,
+                                            dynamic_output_id))
+
     def load_step_func_body(self, step_id: StepID) -> Callable:
         """Load the function body of the workflow step.
 

<!-- Thank you for your contribution! Please review https://github.com/ray-project/ray/blob/master/CONTRIBUTING.rst before opening a pull request. -->

<!-- Please add a reviewer to the assignee section when you create a PR. If you don't have the access to it, we will shortly find a reviewer and assign them to your PR. -->

## Why are these changes needed?

Workflow fails to recover to the latest checkpoint due to a bug of dynamic output in storage. For example, 

Step A returns Step B, and Step B returns step C. When step C failed, the workflow recovers from step A instead of step B. This is not a bug of recovery algorithm, but a bug of dynamic output: the latest dynamic output of the workflow does not get
updated, so it points to the checkpoint of step A instead of step B.

## Related issue number

<!-- For example: "Closes #1234" -->

## Checks

- [ ] I've run `scripts/format.sh` to lint the changes in this PR.
- [ ] I've included any doc changes needed for https://docs.ray.io/en/master/.
- [ ] I've made sure the tests are passing. Note that there might be a few flaky tests, see the recent failures at https://flakey-tests.ray.io/
- Testing Strategy
   - [ ] Unit tests
   - [ ] Release tests
   - [ ] This PR is not tested :(
