[WIP] feature: support argo retry

metaflow/cli_components/step_cmd.py

@@ -155,6 +155,13 @@ def step(
         ctx.obj.monitor,
         ubf_context,
     )
+
+    latest_done_attempt = task.flow_datastore.get_latest_done_attempt(
+        run_id=run_id, step_name=step_name, task_id=task_id
+    )
+    if latest_done_attempt:
+        retry_count = latest_done_attempt + 1
+
     if clone_only:
         task.clone_only(
             step_name,

metaflow/datastore/flow_datastore.py

@@ -1,5 +1,6 @@
 import itertools
 import json
+from typing import Optional
 
 from .. import metaflow_config
 
@@ -67,6 +68,14 @@ def __init__(
     def datastore_root(self):
         return self._storage_impl.datastore_root
 
+    def get_latest_done_attempt(self, run_id, step_name, task_id) -> Optional[int]:
+        t_datastores = self.get_task_datastores(
+            pathspecs=[f"{run_id}/{step_name}/{task_id}"], include_prior=True
+        )
+        return max(
+            [t.attempt for t in t_datastores], default=None
+        )  # returns default, if this was a first attempt.
+
     def get_task_datastores(
         self,
         run_id=None,

metaflow/mflog/save_logs.py

@@ -37,8 +37,12 @@ def print_clean(line, **kwargs):
     flow_datastore = FlowDataStore(
         flow_name, None, storage_impl=storage_impl, ds_root=ds_root
     )
+    # Use inferred attempt - to save task_stdout.log and task_stderr.log
+    latest_done_attempt = flow_datastore.get_latest_done_attempt(
+        run_id=run_id, step_name=step_name, task_id=task_id
+    )
     task_datastore = flow_datastore.get_task_datastore(
-        run_id, step_name, task_id, int(attempt), mode="w"
+        run_id, step_name, task_id, latest_done_attempt, mode="w"
     )
 
     try: