Retry ValueError for airflow tests

This change adds a retry to import utils from airflow in case it trys to
load util from local module instead of airflow dependency

Author: navinsoni

tests/integ/test_airflow_config.py

@@ -13,6 +13,7 @@
 from __future__ import absolute_import
 
 import os
+import sys
 
 import pytest
 import numpy as np
@@ -42,7 +43,6 @@
 from sagemaker.sklearn import SKLearn
 from sagemaker.tensorflow import TensorFlow
 from sagemaker.utils import sagemaker_timestamp
-from sagemaker.workflow import airflow as sm_airflow
 from sagemaker.xgboost import XGBoost
 from tests.integ import datasets, DATA_DIR
 from tests.integ.record_set import prepare_record_set_from_local_files
@@ -54,16 +54,20 @@
     seconds_to_sleep=6,
 ):
     try:
-        from airflow import utils
-        from airflow import DAG
-        from airflow.providers.amazon.aws.operators.sagemaker import SageMakerTrainingOperator
-        from airflow.providers.amazon.aws.operators.sagemaker_transform import (
-            SageMakerTransformOperator,
-        )
+        import sagemaker.workflow.airflow as sm_airflow
+        import airflow.utils as utils
+        import airflow.DAG as DAG
+        import airflow.providers.amazon.aws.operators.sagemaker.SageMakerTrainingOperator as SageMakerTrainingOperator
+        import airflow.providers.amazon.aws.operators.sagemaker_transform.SageMakerTransformOperator as SageMakerTransformOperator
 
         break
     except ParsingError:
         pass
+    except ValueError as ve:
+        if "Unable to configure formatter" in str(ve):
+            print(f"Received: {ve}")
+        else:
+            raise ve
 
 PYTORCH_MNIST_DIR = os.path.join(DATA_DIR, "pytorch_mnist")
 PYTORCH_MNIST_SCRIPT = os.path.join(PYTORCH_MNIST_DIR, "mnist.py")