LOGMLE - Fix argo metaflow retry integration

metaflow/cli_components/step_cmd.py

@@ -162,6 +162,17 @@ def step(
             retry_count,
         )
     else:
+        latest_done_attempt = task.flow_datastore.get_latest_done_attempt(
+            run_id=run_id,
+            step_name=step_name,
+            task_id=task_id
+        )
+        retry_count = latest_done_attempt + 1
+        echo_always(f"{latest_done_attempt=}")
+        echo_always(f"{retry_count=}")
+        # Not sure what are the side effects to this.
+        if retry_count >= max_user_code_retries:
+            max_user_code_retries = retry_count
         task.run_step(
             step_name,
             run_id,

metaflow/datastore/flow_datastore.py

@@ -67,6 +67,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) -> 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=-1)  # default=-1, when no successful done_attempts found.
+
     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, int(latest_done_attempt), mode="w"
     )
 
     try: