fix: Add retry for airflow ParsingError (#3430)

Author: mufaddal-rohawala

tests/integ/test_airflow_config.py

@@ -16,10 +16,9 @@
 
 import pytest
 import numpy as np
-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
+from configparser import ParsingError
+from sagemaker.utils import retries
+
 from six.moves.urllib.parse import urlparse
 
 import tests.integ
@@ -49,6 +48,23 @@
 from tests.integ.record_set import prepare_record_set_from_local_files
 from tests.integ.timeout import timeout
 
+for _ in retries(
+    max_retry_count=10,  # 10*6 = 1min
+    exception_message_prefix="airflow import ",
+    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,
+        )
+
+        break
+    except ParsingError:
+        pass
+
 PYTORCH_MNIST_DIR = os.path.join(DATA_DIR, "pytorch_mnist")
 PYTORCH_MNIST_SCRIPT = os.path.join(PYTORCH_MNIST_DIR, "mnist.py")
 AIRFLOW_CONFIG_TIMEOUT_IN_SECONDS = 10