def test_deploy(sagemaker_session, candidate_mock):
auto_ml = AutoML(role=ROLE,
target_attribute_name=TARGET_ATTRIBUTE_NAME,
sagemaker_session=sagemaker_session)
auto_ml.best_candidate = Mock(name="best_candidate",
return_value=CANDIDATE_DICT)
auto_ml._deploy_inference_pipeline = Mock("_deploy_inference_pipeline",
return_value=None)
auto_ml.deploy(
initial_instance_count=INSTANCE_COUNT,
instance_type=INSTANCE_TYPE,
sagemaker_session=sagemaker_session,
)
auto_ml._deploy_inference_pipeline.assert_called_once()
auto_ml._deploy_inference_pipeline.assert_called_with(
candidate_mock.containers,
initial_instance_count=INSTANCE_COUNT,
instance_type=INSTANCE_TYPE,
name=None,
sagemaker_session=sagemaker_session,
endpoint_name=None,
tags=None,
wait=True,
update_endpoint=False,
vpc_config=None,
enable_network_isolation=False,
model_kms_key=None,
predictor_cls=None,
)
def test_deploy_optional_args(candidate_estimator, sagemaker_session,
candidate_mock):
candidate_estimator.return_value = candidate_mock
auto_ml = AutoML(role=ROLE,
target_attribute_name=TARGET_ATTRIBUTE_NAME,
sagemaker_session=sagemaker_session)
auto_ml._deploy_inference_pipeline = Mock("_deploy_inference_pipeline",
return_value=None)
auto_ml.deploy(
initial_instance_count=INSTANCE_COUNT,
instance_type=INSTANCE_TYPE,
candidate=CANDIDATE_DICT,
sagemaker_session=sagemaker_session,
name=JOB_NAME,
endpoint_name=JOB_NAME,
tags=TAGS,
wait=False,
update_endpoint=True,
vpc_config=VPC_CONFIG,
enable_network_isolation=True,
model_kms_key=OUTPUT_KMS_KEY,
predictor_cls=RealTimePredictor,
)
auto_ml._deploy_inference_pipeline.assert_called_once()
auto_ml._deploy_inference_pipeline.assert_called_with(
candidate_mock.containers,
initial_instance_count=INSTANCE_COUNT,
instance_type=INSTANCE_TYPE,
name=JOB_NAME,
sagemaker_session=sagemaker_session,
endpoint_name=JOB_NAME,
tags=TAGS,
wait=False,
update_endpoint=True,
vpc_config=VPC_CONFIG,
enable_network_isolation=True,
model_kms_key=OUTPUT_KMS_KEY,
predictor_cls=RealTimePredictor,
)
candidate_estimator.assert_called_with(CANDIDATE_DICT,
sagemaker_session=sagemaker_session)
