def test_sagemaker_apply_delete_model_fail():
orig = botocore.client.BaseClient._make_api_call
yatai_service = create_yatai_service_mock()
sagemaker_deployment_pb = generate_sagemaker_deployment_pb()
deployment_operator = SageMakerDeploymentOperator(yatai_service)
def fail_delete_model(self, operation_name, kwarg):
if operation_name == 'DeleteModel':
raise ClientError(
{
'Error': {
'Code': 'ValidationException',
'Message': 'failed message'
}
},
'DeleteModel',
)
elif operation_name == 'CreateEndpoint':
raise ClientError({}, 'CreateEndpoint')
else:
return orig(self, operation_name, kwarg)
with patch('botocore.client.BaseClient._make_api_call',
new=fail_delete_model):
result = deployment_operator.add(sagemaker_deployment_pb)
assert result.status.status_code == Status.INTERNAL
assert result.status.error_message.startswith(
'Failed to cleanup sagemaker model')
def test_sagemaker_apply_success():
yatai_service = create_yatai_service_mock()
sagemaker_deployment_pb = generate_sagemaker_deployment_pb()
deployment_operator = SageMakerDeploymentOperator(yatai_service)
result_pb = deployment_operator.add(sagemaker_deployment_pb)
assert result_pb.status.status_code == Status.OK
assert result_pb.deployment.name == TEST_DEPLOYMENT_NAME
def test_sagemaker_apply_fail_not_local_repo():
yatai_service = create_yatai_service_mock(repo_storage_type=BentoUri.UNSET)
sagemaker_deployment_pb = generate_sagemaker_deployment_pb()
deployment_operator = SageMakerDeploymentOperator(yatai_service)
result_pb = deployment_operator.add(sagemaker_deployment_pb)
assert result_pb.status.status_code == Status.INTERNAL
assert result_pb.status.error_message.startswith('BentoML currently not support')
def test_sagemaker_update_deployment_with_new_bento_service_tag():
mocked_yatai_service = create_yatai_service_mock()
mocked_sagemaker_deployment_pb = generate_sagemaker_deployment_pb()
deployment_operator = SageMakerDeploymentOperator(mocked_yatai_service)
deployment_operator.add(mocked_sagemaker_deployment_pb)
mocked_sagemaker_deployment_pb_with_new_bento_service_tag = (
generate_sagemaker_deployment_pb())
mocked_sagemaker_deployment_pb_with_new_bento_service_tag.spec.bento_version = (
'NEW_BENTO_VERSION')
update_sagemaker_deployment_result = deployment_operator.update(
mocked_sagemaker_deployment_pb_with_new_bento_service_tag,
mocked_sagemaker_deployment_pb,
)
assert update_sagemaker_deployment_result.status.status_code == Status.OK
def get_deployment_operator(yatai_service, deployment_pb):
operator = deployment_pb.spec.operator
if operator == DeploymentSpec.AWS_SAGEMAKER:
from bentoml.yatai.deployment.sagemaker.operator import (
SageMakerDeploymentOperator, )
return SageMakerDeploymentOperator(yatai_service)
elif operator == DeploymentSpec.AWS_LAMBDA:
from bentoml.yatai.deployment.aws_lambda.operator import (
AwsLambdaDeploymentOperator, )
return AwsLambdaDeploymentOperator(yatai_service)
elif operator == DeploymentSpec.GCP_FUNCTION:
raise NotImplementedError(
"GCP Function deployment operator is not supported in current version of "
"BentoML")
elif operator == DeploymentSpec.KUBERNETES:
raise NotImplementedError(
"Kubernetes deployment operator is not supported in current version of "
"BentoML")
elif operator == DeploymentSpec.CUSTOM:
raise NotImplementedError(
"Custom deployment operator is not supported in current version of BentoML"
)
else:
raise YataiDeploymentException("DeployOperator must be set")
def get_deployment_operator(yatai_service, deployment_pb):
operator = deployment_pb.spec.operator
if operator == DeploymentSpec.AWS_SAGEMAKER:
from bentoml.yatai.deployment.sagemaker.operator import (
SageMakerDeploymentOperator, )
return SageMakerDeploymentOperator(yatai_service)
elif operator == DeploymentSpec.AWS_LAMBDA:
from bentoml.yatai.deployment.aws_lambda.operator import (
AwsLambdaDeploymentOperator, )
return AwsLambdaDeploymentOperator(yatai_service)
elif operator == DeploymentSpec.AZURE_FUNCTIONS:
from bentoml.yatai.deployment.azure_functions.operator import (
AzureFunctionsDeploymentOperator, )
return AzureFunctionsDeploymentOperator(yatai_service)
elif operator == DeploymentSpec.AWS_EC2:
from bentoml.yatai.deployment.aws_ec2.operator import AwsEc2DeploymentOperator
return AwsEc2DeploymentOperator(yatai_service)
elif operator == DeploymentSpec.CUSTOM:
raise NotImplementedError(
"Custom deployment operator is not supported in current version of BentoML"
)
else:
raise YataiDeploymentException("DeployOperator must be set")
def test_sagemaker_apply_duplicate_endpoint():
orig = botocore.client.BaseClient._make_api_call
yatai_service = create_yatai_service_mock()
sagemaker_deployment_pb = generate_sagemaker_deployment_pb()
deployment_operator = SageMakerDeploymentOperator(yatai_service)
deployment_operator.add(sagemaker_deployment_pb)
endpoint_name = '{ns}-{name}'.format(
ns=TEST_DEPLOYMENT_NAMESPACE, name=TEST_DEPLOYMENT_NAME
)
expect_value = 'Endpoint {} already exists'.format(endpoint_name.replace('_', '-'))
def mock_ok_return(self, op_name, kwargs):
if op_name == 'CreateModel' or op_name == 'CreateEndpointConfig':
return ''
else:
return orig(self, op_name, kwargs)
with patch('botocore.client.BaseClient._make_api_call', new=mock_ok_return):
with pytest.raises(ValueError) as error:
deployment_operator.add(sagemaker_deployment_pb)
assert str(error.value) == expect_value