Skip to content

Commit

Permalink
fix: Retry ValueError for airflow tests (#3566)
Browse files Browse the repository at this point in the history
This change adds a retry to import utils from airflow in case it trys to
load util from local module instead of airflow dependency
  • Loading branch information
navinsoni committed Jan 5, 2023
1 parent 11ff7c0 commit d4203da
Showing 1 changed file with 7 additions and 2 deletions.
9 changes: 7 additions & 2 deletions tests/integ/test_airflow_config.py
Original file line number Diff line number Diff line change
Expand Up @@ -42,7 +42,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
Expand All @@ -54,7 +53,8 @@
seconds_to_sleep=6,
):
try:
from airflow import utils
import sagemaker.workflow.airflow as sm_airflow
import airflow.utils as utils
from airflow import DAG
from airflow.providers.amazon.aws.operators.sagemaker import SageMakerTrainingOperator
from airflow.providers.amazon.aws.operators.sagemaker_transform import (
Expand All @@ -64,6 +64,11 @@
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")
Expand Down

0 comments on commit d4203da

Please sign in to comment.