diff --git a/tests/integ/test_auto_ml.py b/tests/integ/test_auto_ml.py index 50ff31ba8f..a473c12100 100644 --- a/tests/integ/test_auto_ml.py +++ b/tests/integ/test_auto_ml.py @@ -26,7 +26,6 @@ ROLE = "SageMakerRole" PREFIX = "sagemaker/beta-automl-xgboost" -HOSTING_INSTANCE_TYPE = "ml.c4.xlarge" AUTO_ML_INSTANCE_TYPE = "ml.m5.2xlarge" INSTANCE_COUNT = 1 RESOURCE_POOLS = [{"InstanceType": AUTO_ML_INSTANCE_TYPE, "PoolSize": INSTANCE_COUNT}] @@ -216,7 +215,7 @@ def test_best_candidate(sagemaker_session): tests.integ.test_region() in tests.integ.NO_AUTO_ML_REGIONS, reason="AutoML is not supported in the region yet.", ) -def test_deploy_best_candidate(sagemaker_session): +def test_deploy_best_candidate(sagemaker_session, cpu_instance_type): auto_ml_utils.create_auto_ml_job_if_not_exist(sagemaker_session) auto_ml = AutoML( @@ -229,7 +228,7 @@ def test_deploy_best_candidate(sagemaker_session): auto_ml.deploy( candidate=best_candidate, initial_instance_count=INSTANCE_COUNT, - instance_type=HOSTING_INSTANCE_TYPE, + instance_type=cpu_instance_type, endpoint_name=endpoint_name, ) @@ -244,7 +243,7 @@ def test_deploy_best_candidate(sagemaker_session): tests.integ.test_region() in tests.integ.NO_AUTO_ML_REGIONS, reason="AutoML is not supported in the region yet.", ) -def test_candidate_estimator_default_rerun_and_deploy(sagemaker_session): +def test_candidate_estimator_default_rerun_and_deploy(sagemaker_session, cpu_instance_type): auto_ml_utils.create_auto_ml_job_if_not_exist(sagemaker_session) auto_ml = AutoML( @@ -261,7 +260,7 @@ def test_candidate_estimator_default_rerun_and_deploy(sagemaker_session): candidate_estimator.fit(inputs) auto_ml.deploy( initial_instance_count=INSTANCE_COUNT, - instance_type=HOSTING_INSTANCE_TYPE, + instance_type=cpu_instance_type, candidate=candidate, endpoint_name=endpoint_name, ) @@ -277,7 +276,7 @@ def test_candidate_estimator_default_rerun_and_deploy(sagemaker_session): tests.integ.test_region() in tests.integ.NO_AUTO_ML_REGIONS, reason="AutoML is not supported in the region yet.", ) -def test_candidate_estimator_rerun_with_optional_args(sagemaker_session): +def test_candidate_estimator_rerun_with_optional_args(sagemaker_session, cpu_instance_type): auto_ml_utils.create_auto_ml_job_if_not_exist(sagemaker_session) auto_ml = AutoML( @@ -294,7 +293,7 @@ def test_candidate_estimator_rerun_with_optional_args(sagemaker_session): candidate_estimator.fit(inputs, encrypt_inter_container_traffic=True) auto_ml.deploy( initial_instance_count=INSTANCE_COUNT, - instance_type=HOSTING_INSTANCE_TYPE, + instance_type=cpu_instance_type, candidate=candidate, endpoint_name=endpoint_name, )