make the handling of Unknown job status better by polling

Signed-off-by: prekshivyas <[email protected]>

Author: prekshivyas

nemo_run/core/execution/lepton.py

@@ -127,41 +127,39 @@ def move_data(
         job_id = response.metadata.id_
 
         start_time = time.time()
-        unknown_start_time = None
-        count = 0
 
         while True:
             if time.time() - start_time > timeout:
                 raise TimeoutError(f"Job {job_id} did not complete within {timeout} seconds.")
 
             current_job = client.job.get(job_id)
             current_job_status = current_job.status.state
-            if count > 0:
-                if (
-                    current_job_status == LeptonJobState.Completed
-                    or current_job_status == LeptonJobState.Failed
-                ):
-                    break
-                elif current_job_status == LeptonJobState.Unknown:
-                    if unknown_start_time is None:
-                        unknown_start_time = time.time()
-                        logging.warning(
-                            f"Job {job_id} entered Unknown state, giving it {unknowns_grace_period} seconds to recover..."
-                        )
-
-                    elif time.time() - unknown_start_time > unknowns_grace_period:
-                        logging.error(
-                            f"Job {job_id} has been in Unknown state for more than {unknowns_grace_period} seconds"
-                        )
-                        break
-                else:
-                    if unknown_start_time is not None:
+            if (
+                current_job_status == LeptonJobState.Completed
+                or current_job_status == LeptonJobState.Failed
+            ):
+                break
+            elif current_job_status == LeptonJobState.Unknown:
+                logging.warning(
+                    f"Job {job_id} entered Unknown state, checking for up to {unknowns_grace_period} seconds every 2 seconds..."
+                )
+                unknown_start_time = time.time()
+                recovered = False
+                while time.time() - unknown_start_time < unknowns_grace_period:
+                    time.sleep(2)
+                    current_job = client.job.get(job_id)
+                    current_job_status = current_job.status.state
+                    if current_job_status != LeptonJobState.Unknown:
                         logging.info(
                             f"Job {job_id} recovered from Unknown state to {current_job_status}"
                         )
-                        unknown_start_time = None
-
-            count += 1
+                        recovered = True
+                        break
+                if not recovered:
+                    logging.error(
+                        f"Job {job_id} has been in Unknown state for more than {unknowns_grace_period} seconds"
+                    )
+                    break
             time.sleep(poll_interval)
 
         if current_job_status != LeptonJobState.Completed: