Retry on RecursionError

AgnostiqHQ · Nov 30, 2023 · f81129f · f81129f
1 parent b24e01b
commit f81129f
Show file tree

Hide file tree

Showing 2 changed files with 26 additions and 20 deletions.
diff --git a/covalent/_results_manager/results_manager.py b/covalent/_results_manager/results_manager.py
@@ -491,19 +491,23 @@ def get_result(
         The Result object from the Covalent server
 
     """
+    max_attempts = int(os.getenv("COVALENT_GET_RESULT_RETRIES", 10))
+    num_attempts = 0
+    while num_attempts < max_attempts:
+        try:
+            return _get_result_multistage(
+                dispatch_id=dispatch_id,
+                wait=wait,
+                dispatcher_addr=dispatcher_addr,
+                status_only=status_only,
+                results_dir=results_dir,
+                workflow_output=workflow_output,
+                intermediate_outputs=intermediate_outputs,
+                sublattice_results=sublattice_results,
+                qelectron_db=qelectron_db,
+            )
 
-    try:
-        return _get_result_multistage(
-            dispatch_id=dispatch_id,
-            wait=wait,
-            dispatcher_addr=dispatcher_addr,
-            status_only=status_only,
-            results_dir=results_dir,
-            workflow_output=workflow_output,
-            intermediate_outputs=intermediate_outputs,
-            sublattice_results=sublattice_results,
-            qelectron_db=qelectron_db,
-        )
-
-    except RecursionError as re:
-        app_log.error(re)
+        except RecursionError as re:
+            app_log.error(re)
+            num_attempts += 1
+    raise RuntimeError("Timed out waiting for result. Please retry or check dispatch.")
diff --git a/tests/covalent_tests/results_manager_tests/results_manager_test.py b/tests/covalent_tests/results_manager_tests/results_manager_test.py
@@ -265,11 +265,13 @@ def test_get_result_RecursionError(mocker):
     """Check exception handing for RecursionError."""
 
     dispatch_id = "test_get_result_RecursionError"
-    with mocker.patch(
-        "covalent._results_manager.results_manager._get_result_multistage",
-        side_effect=RecursionError("from test_get_result_RecursionError"),
-    ):
-        assert get_result(dispatch_id, wait=True) is None
+
+    with pytest.raises(RuntimeError):
+        with mocker.patch(
+            "covalent._results_manager.results_manager._get_result_multistage",
+            side_effect=RecursionError("from test_get_result_RecursionError"),
+        ):
+            get_result(dispatch_id, wait=True)
 
 
 def test_get_status_only(mocker):