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 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_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_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_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