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_delete_model_fail(
mock_init_project, mock_docker_push, mock_docker_build, mock_check_output
):
orig = botocore.client.BaseClient._make_api_call
yatai_service = create_yatai_service_mock()
sagemaker_deployment_pb = generate_sagemaker_deployment_pb()
deployment_operator = SageMakerDeploymentOperator()
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.apply(sagemaker_deployment_pb, yatai_service)
assert result.status.status_code == Status.NOT_FOUND
assert result.status.error_message.startswith(
'Failed to clean up model after unsuccessfully'
)
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_apply_success(
mock_init_project, mock_docker_push, mock_docker_build, mock_check_output
):
yatai_service = create_yatai_service_mock()
sagemaker_deployment_pb = generate_sagemaker_deployment_pb()
deployment_operator = SageMakerDeploymentOperator()
result_pb = deployment_operator.apply(sagemaker_deployment_pb, yatai_service)
assert result_pb.status.status_code == Status.OK
assert result_pb.deployment.name == TEST_DEPLOYMENT_NAME
def test_sagemaker_apply_fail_not_local_repo(
mock_init_project, mock_docker_push, mock_docker_build, mock_check_output
):
yatai_service = create_yatai_service_mock(repo_storage_type=BentoUri.UNSET)
sagemaker_deployment_pb = generate_sagemaker_deployment_pb()
deployment_operator = SageMakerDeploymentOperator()
result_pb = deployment_operator.apply(sagemaker_deployment_pb, yatai_service)
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.deployment.sagemaker import SageMakerDeploymentOperator
return SageMakerDeploymentOperator(yatai_service)
elif operator == DeploymentSpec.AWS_LAMBDA:
from bentoml.deployment.aws_lambda 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 test_sagemaker_apply_create_model_fail(
mock_init_project, mock_docker_push, mock_docker_build, mock_check_output
):
yatai_service = create_yatai_service_mock()
sagemaker_deployment_pb = generate_sagemaker_deployment_pb()
deployment_operator = SageMakerDeploymentOperator()
orig = botocore.client.BaseClient._make_api_call
def fail_create_model_random(self, operation_name, kwarg):
if operation_name == 'CreateModel':
raise ClientError({'Error': {'Code': 'Random'}}, 'CreateModel')
else:
return orig(self, operation_name, kwarg)
with patch(
'botocore.client.BaseClient._make_api_call', new=fail_create_model_random
):
failed_result = deployment_operator.apply(
sagemaker_deployment_pb, yatai_service
)
assert failed_result.status.status_code == Status.INTERNAL
assert failed_result.status.error_message.startswith(
'Failed to create model for SageMaker Deployment'
)
def fail_create_model_validation(self, operation_name, kwarg):
if operation_name == 'CreateModel':
raise ClientError(
{'Error': {'Code': 'ValidationException', 'Message': 'failed message'}},
'CreateModel',
)
else:
return orig(self, operation_name, kwarg)
with patch(
'botocore.client.BaseClient._make_api_call', new=fail_create_model_validation
):
result = deployment_operator.apply(sagemaker_deployment_pb, yatai_service)
assert result.status.status_code == Status.NOT_FOUND
assert result.status.error_message.startswith(
'Failed to create model for SageMaker'
)
def get_deployment_operator(deployment_pb):
operator = deployment_pb.spec.operator
if operator == DeploymentOperator.AWS_SAGEMAKER:
from bentoml.deployment.sagemaker import SageMakerDeploymentOperator
return SageMakerDeploymentOperator()
elif operator == DeploymentOperator.AWS_LAMBDA:
pass
elif operator == DeploymentOperator.GCP_FUNCTION:
pass
elif operator == DeploymentOperator.KUBERNETES:
pass
elif operator == DeploymentOperator.CUSTOM:
pass
else:
raise BentoMLDeploymentException("DeployOperator must be set")
def get_deployment_operator(deployment_pb):
operator = deployment_pb.spec.operator
if operator == deployment_pb2.AWS_SAGEMAKER:
from bentoml.deployment.sagemaker import SageMakerDeploymentOperator
return SageMakerDeploymentOperator()
elif operator == deployment_pb2.AWS_LAMBDA:
pass
elif operator == deployment_pb2.GCP_FUNCTION:
raise NotImplementedError(
"GCP function deployment operator is not implemented")
elif operator == deployment_pb2.KUBERNETES:
raise NotImplementedError(
"Kubernetes deployment operator is not implemented")
elif operator == deployment_pb2.CUSTOM:
raise NotImplementedError(
"Custom deployment operator is not implemented")
else:
raise BentoMLDeploymentException("DeployOperator must be set")
def get_deployment_operator(deployment_pb):
operator = deployment_pb.spec.operator
if operator == DeploymentSpec.AWS_SAGEMAKER:
from bentoml.deployment.sagemaker import SageMakerDeploymentOperator
return SageMakerDeploymentOperator()
elif operator == DeploymentSpec.AWS_LAMBDA:
from bentoml.deployment.serverless.aws_lambda import AwsLambdaDeploymentOperator
return AwsLambdaDeploymentOperator()
elif operator == DeploymentSpec.GCP_FUNCTION:
from bentoml.deployment.serverless.gcp_function import (
GcpFunctionDeploymentOperator, )
return GcpFunctionDeploymentOperator()
elif operator == DeploymentSpec.KUBERNETES:
raise NotImplementedError(
"Kubernetes deployment operator is not implemented")
elif operator == DeploymentSpec.CUSTOM:
raise NotImplementedError(
"Custom deployment operator is not implemented")
else:
raise BentoMLDeploymentException("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