add wait for trial component to be created to improve test reliability

Author: danabens

tests/helpers.py

@@ -16,6 +16,7 @@
 import uuid
 import boto3
 import sys
+import logging
 
 
 def name():
@@ -145,3 +146,20 @@ def wait_for_job(job_name, get_job, status_field):
                 sys.stdout.write(".")
                 sys.stdout.flush()
                 time.sleep(30)
+
+
+def wait_for_trial_component(sagemaker_client, training_job_name=None, trial_component_name=None):
+    # wait for trial component to be created
+    if training_job_name and not trial_component_name:
+        # assume tj
+        trial_component_name = training_job_name + "-aws-training-job"
+
+    # wait for the trial component to be created from the training job, usually < 5s
+    with timeout(minutes=15):
+        while True:
+            try:
+                sagemaker_client.describe_trial_component(TrialComponentName=trial_component_name)
+                break
+            except sagemaker_client.exceptions.ResourceNotFound:
+                logging.info("Trial component %s not created yet.", trial_component_name)
+                time.sleep(5)

tests/integ/test_tracker.py

@@ -15,11 +15,12 @@
 
 import pytest
 
-from tests.helpers import name
+from tests.helpers import name, wait_for_trial_component
 from smexperiments import tracker, trial_component, _utils
 
 
 def test_load_trial_component(trial_component_obj, sagemaker_boto_client):
+    wait_for_trial_component(sagemaker_boto_client, trial_component_name=trial_component_obj.trial_component_name)
     tracker_obj = tracker.Tracker.load(
         trial_component_name=trial_component_obj.trial_component_name, sagemaker_boto_client=sagemaker_boto_client
     )
@@ -29,6 +30,7 @@ def test_load_trial_component(trial_component_obj, sagemaker_boto_client):
 
 @pytest.mark.docker
 def test_load_training_job(training_job_name, sagemaker_boto_client):
+    wait_for_trial_component(sagemaker_boto_client, training_job_name=training_job_name)
     tracker_obj = tracker.Tracker.load(training_job_name=training_job_name, sagemaker_boto_client=sagemaker_boto_client)
     assert tracker_obj
     assert tracker_obj.trial_component.trial_component_name == training_job_name + "-aws-training-job"