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)
mock_pipeline = Mock(name="pipeline_model")
mock_pipeline.deploy = Mock(name="model_deploy")
auto_ml.best_candidate = Mock(name="best_candidate",
return_value=CANDIDATE_DICT)
auto_ml.create_model = Mock(name="create_model",
return_value=mock_pipeline)
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,
inference_response_keys=None,
)
auto_ml.create_model.assert_called_once()
auto_ml.create_model.assert_called_with(
name=JOB_NAME,
sagemaker_session=sagemaker_session,
candidate=CANDIDATE_DICT,
inference_response_keys=None,
vpc_config=VPC_CONFIG,
enable_network_isolation=True,
model_kms_key=OUTPUT_KMS_KEY,
predictor_cls=RealTimePredictor,
)
mock_pipeline.deploy.assert_called_once()
mock_pipeline.deploy.assert_called_with(
initial_instance_count=INSTANCE_COUNT,
instance_type=INSTANCE_TYPE,
endpoint_name=JOB_NAME,
tags=TAGS,
wait=False,
update_endpoint=True,
)
def test_deploy(sagemaker_session, candidate_mock):
auto_ml = AutoML(
role=ROLE, target_attribute_name=TARGET_ATTRIBUTE_NAME, sagemaker_session=sagemaker_session
)
mock_pipeline = Mock(name="pipeline_model")
mock_pipeline.deploy = Mock(name="model_deploy")
auto_ml.best_candidate = Mock(name="best_candidate", return_value=CANDIDATE_DICT)
auto_ml.create_model = Mock(name="create_model", return_value=mock_pipeline)
auto_ml.deploy(
initial_instance_count=INSTANCE_COUNT,
instance_type=INSTANCE_TYPE,
sagemaker_session=sagemaker_session,
)
auto_ml.create_model.assert_called_once()
mock_pipeline.deploy.assert_called_once()
def test_create_model(sagemaker_session):
auto_ml = AutoML(
role=ROLE, target_attribute_name=TARGET_ATTRIBUTE_NAME, sagemaker_session=sagemaker_session
)
pipeline_model = auto_ml.create_model(
name=JOB_NAME,
sagemaker_session=sagemaker_session,
candidate=CLASSIFICATION_CANDIDATE_DICT,
vpc_config=VPC_CONFIG,
enable_network_isolation=True,
model_kms_key=None,
predictor_cls=None,
inference_response_keys=None,
)
assert isinstance(pipeline_model, PipelineModel)
