LOGMLE - Fix argo metaflow retry integration

metaflow/cli_components/step_cmd.py

@@ -102,7 +102,7 @@ def step(
     input_paths_filename=None,
     split_index=None,
     opt_namespace=None,
-    retry_count=None,
+    retry_count=None,  # TODO remove this from interface, look for effecting codepaths
     max_user_code_retries=None,
     clone_only=None,
     clone_run_id=None,
@@ -155,6 +155,12 @@ 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,13 @@ 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

@@ -21,7 +21,6 @@ def _read_file(path):
 
     # these env vars are set by mflog.mflog_env
     pathspec = os.environ["MF_PATHSPEC"]
-    attempt = os.environ["MF_ATTEMPT"]
     ds_type = os.environ["MF_DATASTORE"]
     ds_root = os.environ.get("MF_DATASTORE_ROOT")
     paths = (os.environ["MFLOG_STDOUT"], os.environ["MFLOG_STDERR"])
@@ -37,8 +36,10 @@ 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: