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)