retry downstream_trials test

Author: danabens

tests/integ/sagemaker/lineage/helpers.py

@@ -15,7 +15,8 @@
 
 import uuid
 from datetime import datetime
-
+from contextlib import contextmanager
+import time
 
 def name():
     return "lineage-integ-{}-{}".format(
@@ -30,3 +31,15 @@ def names():
         )
         for i in range(3)
     ]
+
+def retry(callable, num_attempts=8):
+    assert num_attempts >= 1
+    for i in range(num_attempts):
+        try:
+            return callable()
+        except Exception as ex:
+            if i == num_attempts - 1:
+                raise ex
+            print("Retrying", ex)
+            time.sleep(2 ** i)
+    assert False, "logic error in retry"

tests/integ/sagemaker/lineage/test_artifact.py

@@ -20,6 +20,7 @@
 import pytest
 
 from sagemaker.lineage import artifact
+from tests.integ.sagemaker.lineage.helpers import retry
 
 
 def test_create_delete(artifact_obj):
@@ -103,14 +104,20 @@ def test_list_by_type(artifact_objs, sagemaker_session):
 
 def test_downstream_trials(trial_associated_artifact, trial_obj, sagemaker_session):
     # allow trial components to index, 30 seconds max
-    for i in range(3):
-        time.sleep(10)
-        trials = trial_associated_artifact.downstream_trials(sagemaker_session=sagemaker_session)
-        if len(trials) > 0:
-            break
+    def validate():
+        for i in range(3):
+            time.sleep(10)
+            trials = trial_associated_artifact.downstream_trials(
+                sagemaker_session=sagemaker_session
+            )
+            logging.info(f"Found {len(trials)} downstream trials.")
+            if len(trials) > 0:
+                break
+
+        assert len(trials) == 1
+        assert trial_obj.trial_name in trials
 
-    assert len(trials) == 1
-    assert trial_obj.trial_name in trials
+    retry(validate, num_attempts=3)
 
 
 @pytest.mark.timeout(30)