Add CloudFormation support for SageMaker Models (#3861)
* Create a formal interface for SM Cloudformation test configurations * Create SageMaker Models with CloudFormation * Utilize six for adding metaclass to TestConfig * Update SM backend to return Model objects instead of response objects
This commit is contained in:
Zach Churchill
GitHub
parent
0b11b0c716
commit
f6dda54a6c
4 changed files with 389 additions and 163 deletions
|
|
@ -1,4 +1,3 @@
|
|||
import json
|
||||
import boto3
|
||||
|
||||
import pytest
|
||||
|
|
@ -6,128 +5,58 @@ import sure # noqa
|
|||
from botocore.exceptions import ClientError
|
||||
|
||||
from moto import mock_cloudformation, mock_sagemaker
|
||||
from moto.sts.models import ACCOUNT_ID
|
||||
|
||||
|
||||
def _get_notebook_instance_template_string(
|
||||
resource_name="TestNotebook",
|
||||
instance_type="ml.c4.xlarge",
|
||||
role_arn="arn:aws:iam::{}:role/FakeRole".format(ACCOUNT_ID),
|
||||
include_outputs=True,
|
||||
):
|
||||
template = {
|
||||
"AWSTemplateFormatVersion": "2010-09-09",
|
||||
"Resources": {
|
||||
resource_name: {
|
||||
"Type": "AWS::SageMaker::NotebookInstance",
|
||||
"Properties": {"InstanceType": instance_type, "RoleArn": role_arn},
|
||||
},
|
||||
},
|
||||
}
|
||||
if include_outputs:
|
||||
template["Outputs"] = {
|
||||
"Arn": {"Value": {"Ref": resource_name}},
|
||||
"Name": {"Value": {"Fn::GetAtt": [resource_name, "NotebookInstanceName"]}},
|
||||
}
|
||||
return json.dumps(template)
|
||||
|
||||
|
||||
def _get_notebook_instance_lifecycle_config_template_string(
|
||||
resource_name="TestConfig", on_create=None, on_start=None, include_outputs=True,
|
||||
):
|
||||
template = {
|
||||
"AWSTemplateFormatVersion": "2010-09-09",
|
||||
"Resources": {
|
||||
resource_name: {
|
||||
"Type": "AWS::SageMaker::NotebookInstanceLifecycleConfig",
|
||||
"Properties": {},
|
||||
},
|
||||
},
|
||||
}
|
||||
if on_create is not None:
|
||||
template["Resources"][resource_name]["Properties"]["OnCreate"] = [
|
||||
{"Content": on_create}
|
||||
]
|
||||
if on_start is not None:
|
||||
template["Resources"][resource_name]["Properties"]["OnStart"] = [
|
||||
{"Content": on_start}
|
||||
]
|
||||
if include_outputs:
|
||||
template["Outputs"] = {
|
||||
"Arn": {"Value": {"Ref": resource_name}},
|
||||
"Name": {
|
||||
"Value": {
|
||||
"Fn::GetAtt": [
|
||||
resource_name,
|
||||
"NotebookInstanceLifecycleConfigName",
|
||||
]
|
||||
}
|
||||
},
|
||||
}
|
||||
return json.dumps(template)
|
||||
from .cloudformation_test_configs import (
|
||||
NotebookInstanceTestConfig,
|
||||
NotebookInstanceLifecycleConfigTestConfig,
|
||||
ModelTestConfig,
|
||||
)
|
||||
|
||||
|
||||
@mock_cloudformation
|
||||
@pytest.mark.parametrize(
|
||||
"stack_name,resource_name,template",
|
||||
"test_config",
|
||||
[
|
||||
(
|
||||
"test_sagemaker_notebook_instance",
|
||||
"TestNotebookInstance",
|
||||
_get_notebook_instance_template_string(
|
||||
resource_name="TestNotebookInstance", include_outputs=False
|
||||
),
|
||||
),
|
||||
(
|
||||
"test_sagemaker_notebook_instance_lifecycle_config",
|
||||
"TestNotebookInstanceLifecycleConfig",
|
||||
_get_notebook_instance_lifecycle_config_template_string(
|
||||
resource_name="TestNotebookInstanceLifecycleConfig",
|
||||
include_outputs=False,
|
||||
),
|
||||
),
|
||||
NotebookInstanceTestConfig(),
|
||||
NotebookInstanceLifecycleConfigTestConfig(),
|
||||
ModelTestConfig(),
|
||||
],
|
||||
)
|
||||
def test_sagemaker_cloudformation_create(stack_name, resource_name, template):
|
||||
def test_sagemaker_cloudformation_create(test_config):
|
||||
cf = boto3.client("cloudformation", region_name="us-east-1")
|
||||
cf.create_stack(StackName=stack_name, TemplateBody=template)
|
||||
|
||||
stack_name = "{}_stack".format(test_config.resource_name)
|
||||
cf.create_stack(
|
||||
StackName=stack_name,
|
||||
TemplateBody=test_config.get_cloudformation_template(include_outputs=False),
|
||||
)
|
||||
|
||||
provisioned_resource = cf.list_stack_resources(StackName=stack_name)[
|
||||
"StackResourceSummaries"
|
||||
][0]
|
||||
provisioned_resource["LogicalResourceId"].should.equal(resource_name)
|
||||
provisioned_resource["LogicalResourceId"].should.equal(test_config.resource_name)
|
||||
len(provisioned_resource["PhysicalResourceId"]).should.be.greater_than(0)
|
||||
|
||||
|
||||
@mock_cloudformation
|
||||
@mock_sagemaker
|
||||
@pytest.mark.parametrize(
|
||||
"stack_name,template,describe_function_name,name_parameter,arn_parameter",
|
||||
"test_config",
|
||||
[
|
||||
(
|
||||
"test_sagemaker_notebook_instance",
|
||||
_get_notebook_instance_template_string(),
|
||||
"describe_notebook_instance",
|
||||
"NotebookInstanceName",
|
||||
"NotebookInstanceArn",
|
||||
),
|
||||
(
|
||||
"test_sagemaker_notebook_instance_lifecycle_config",
|
||||
_get_notebook_instance_lifecycle_config_template_string(),
|
||||
"describe_notebook_instance_lifecycle_config",
|
||||
"NotebookInstanceLifecycleConfigName",
|
||||
"NotebookInstanceLifecycleConfigArn",
|
||||
),
|
||||
NotebookInstanceTestConfig(),
|
||||
NotebookInstanceLifecycleConfigTestConfig(),
|
||||
ModelTestConfig(),
|
||||
],
|
||||
)
|
||||
def test_sagemaker_cloudformation_get_attr(
|
||||
stack_name, template, describe_function_name, name_parameter, arn_parameter
|
||||
):
|
||||
def test_sagemaker_cloudformation_get_attr(test_config):
|
||||
cf = boto3.client("cloudformation", region_name="us-east-1")
|
||||
sm = boto3.client("sagemaker", region_name="us-east-1")
|
||||
|
||||
# Create stack and get description for output values
|
||||
cf.create_stack(StackName=stack_name, TemplateBody=template)
|
||||
stack_name = "{}_stack".format(test_config.resource_name)
|
||||
cf.create_stack(
|
||||
StackName=stack_name, TemplateBody=test_config.get_cloudformation_template()
|
||||
)
|
||||
stack_description = cf.describe_stacks(StackName=stack_name)["Stacks"][0]
|
||||
outputs = {
|
||||
output["OutputKey"]: output["OutputValue"]
|
||||
|
|
@ -135,62 +64,50 @@ def test_sagemaker_cloudformation_get_attr(
|
|||
}
|
||||
|
||||
# Using the describe function, ensure output ARN matches resource ARN
|
||||
resource_description = getattr(sm, describe_function_name)(
|
||||
**{name_parameter: outputs["Name"]}
|
||||
resource_description = getattr(sm, test_config.describe_function_name)(
|
||||
**{test_config.name_parameter: outputs["Name"]}
|
||||
)
|
||||
outputs["Arn"].should.equal(resource_description[arn_parameter])
|
||||
outputs["Arn"].should.equal(resource_description[test_config.arn_parameter])
|
||||
|
||||
|
||||
@mock_cloudformation
|
||||
@mock_sagemaker
|
||||
@pytest.mark.parametrize(
|
||||
"stack_name,template,describe_function_name,name_parameter,arn_parameter,error_message",
|
||||
"test_config,error_message",
|
||||
[
|
||||
(NotebookInstanceTestConfig(), "RecordNotFound"),
|
||||
(
|
||||
"test_sagemaker_notebook_instance",
|
||||
_get_notebook_instance_template_string(),
|
||||
"describe_notebook_instance",
|
||||
"NotebookInstanceName",
|
||||
"NotebookInstanceArn",
|
||||
"RecordNotFound",
|
||||
),
|
||||
(
|
||||
"test_sagemaker_notebook_instance_lifecycle_config",
|
||||
_get_notebook_instance_lifecycle_config_template_string(),
|
||||
"describe_notebook_instance_lifecycle_config",
|
||||
"NotebookInstanceLifecycleConfigName",
|
||||
"NotebookInstanceLifecycleConfigArn",
|
||||
NotebookInstanceLifecycleConfigTestConfig(),
|
||||
"Notebook Instance Lifecycle Config does not exist",
|
||||
),
|
||||
(ModelTestConfig(), "Could not find model"),
|
||||
],
|
||||
)
|
||||
def test_sagemaker_cloudformation_notebook_instance_delete(
|
||||
stack_name,
|
||||
template,
|
||||
describe_function_name,
|
||||
name_parameter,
|
||||
arn_parameter,
|
||||
error_message,
|
||||
):
|
||||
def test_sagemaker_cloudformation_notebook_instance_delete(test_config, error_message):
|
||||
cf = boto3.client("cloudformation", region_name="us-east-1")
|
||||
sm = boto3.client("sagemaker", region_name="us-east-1")
|
||||
|
||||
# Create stack and verify existence
|
||||
cf.create_stack(StackName=stack_name, TemplateBody=template)
|
||||
stack_name = "{}_stack".format(test_config.resource_name)
|
||||
cf.create_stack(
|
||||
StackName=stack_name, TemplateBody=test_config.get_cloudformation_template()
|
||||
)
|
||||
stack_description = cf.describe_stacks(StackName=stack_name)["Stacks"][0]
|
||||
outputs = {
|
||||
output["OutputKey"]: output["OutputValue"]
|
||||
for output in stack_description["Outputs"]
|
||||
}
|
||||
resource_description = getattr(sm, describe_function_name)(
|
||||
**{name_parameter: outputs["Name"]}
|
||||
resource_description = getattr(sm, test_config.describe_function_name)(
|
||||
**{test_config.name_parameter: outputs["Name"]}
|
||||
)
|
||||
outputs["Arn"].should.equal(resource_description[arn_parameter])
|
||||
outputs["Arn"].should.equal(resource_description[test_config.arn_parameter])
|
||||
|
||||
# Delete the stack and verify resource has also been deleted
|
||||
cf.delete_stack(StackName=stack_name)
|
||||
with pytest.raises(ClientError) as ce:
|
||||
getattr(sm, describe_function_name)(**{name_parameter: outputs["Name"]})
|
||||
getattr(sm, test_config.describe_function_name)(
|
||||
**{test_config.name_parameter: outputs["Name"]}
|
||||
)
|
||||
ce.value.response["Error"]["Message"].should.contain(error_message)
|
||||
|
||||
|
||||
|
|
@ -200,14 +117,16 @@ def test_sagemaker_cloudformation_notebook_instance_update():
|
|||
cf = boto3.client("cloudformation", region_name="us-east-1")
|
||||
sm = boto3.client("sagemaker", region_name="us-east-1")
|
||||
|
||||
test_config = NotebookInstanceTestConfig()
|
||||
|
||||
# Set up template for stack with initial and update instance types
|
||||
stack_name = "test_sagemaker_notebook_instance"
|
||||
stack_name = "{}_stack".format(test_config.resource_name)
|
||||
initial_instance_type = "ml.c4.xlarge"
|
||||
updated_instance_type = "ml.c4.4xlarge"
|
||||
initial_template_json = _get_notebook_instance_template_string(
|
||||
initial_template_json = test_config.get_cloudformation_template(
|
||||
instance_type=initial_instance_type
|
||||
)
|
||||
updated_template_json = _get_notebook_instance_template_string(
|
||||
updated_template_json = test_config.get_cloudformation_template(
|
||||
instance_type=updated_instance_type
|
||||
)
|
||||
|
||||
|
|
@ -219,10 +138,10 @@ def test_sagemaker_cloudformation_notebook_instance_update():
|
|||
for output in stack_description["Outputs"]
|
||||
}
|
||||
initial_notebook_name = outputs["Name"]
|
||||
notebook_instance_description = sm.describe_notebook_instance(
|
||||
NotebookInstanceName=initial_notebook_name,
|
||||
resource_description = getattr(sm, test_config.describe_function_name)(
|
||||
**{test_config.name_parameter: initial_notebook_name}
|
||||
)
|
||||
initial_instance_type.should.equal(notebook_instance_description["InstanceType"])
|
||||
initial_instance_type.should.equal(resource_description["InstanceType"])
|
||||
|
||||
# Update stack with new instance type and check attributes
|
||||
cf.update_stack(StackName=stack_name, TemplateBody=updated_template_json)
|
||||
|
|
@ -234,10 +153,10 @@ def test_sagemaker_cloudformation_notebook_instance_update():
|
|||
updated_notebook_name = outputs["Name"]
|
||||
updated_notebook_name.should.equal(initial_notebook_name)
|
||||
|
||||
notebook_instance_description = sm.describe_notebook_instance(
|
||||
NotebookInstanceName=updated_notebook_name,
|
||||
resource_description = getattr(sm, test_config.describe_function_name)(
|
||||
**{test_config.name_parameter: updated_notebook_name}
|
||||
)
|
||||
updated_instance_type.should.equal(notebook_instance_description["InstanceType"])
|
||||
updated_instance_type.should.equal(resource_description["InstanceType"])
|
||||
|
||||
|
||||
@mock_cloudformation
|
||||
|
|
@ -246,14 +165,16 @@ def test_sagemaker_cloudformation_notebook_instance_lifecycle_config_update():
|
|||
cf = boto3.client("cloudformation", region_name="us-east-1")
|
||||
sm = boto3.client("sagemaker", region_name="us-east-1")
|
||||
|
||||
test_config = NotebookInstanceLifecycleConfigTestConfig()
|
||||
|
||||
# Set up template for stack with initial and update instance types
|
||||
stack_name = "test_sagemaker_notebook_instance_lifecycle_config"
|
||||
stack_name = "{}_stack".format(test_config.resource_name)
|
||||
initial_on_create_script = "echo Hello World"
|
||||
updated_on_create_script = "echo Goodbye World"
|
||||
initial_template_json = _get_notebook_instance_lifecycle_config_template_string(
|
||||
initial_template_json = test_config.get_cloudformation_template(
|
||||
on_create=initial_on_create_script
|
||||
)
|
||||
updated_template_json = _get_notebook_instance_lifecycle_config_template_string(
|
||||
updated_template_json = test_config.get_cloudformation_template(
|
||||
on_create=updated_on_create_script
|
||||
)
|
||||
|
||||
|
|
@ -265,12 +186,12 @@ def test_sagemaker_cloudformation_notebook_instance_lifecycle_config_update():
|
|||
for output in stack_description["Outputs"]
|
||||
}
|
||||
initial_config_name = outputs["Name"]
|
||||
notebook_lifecycle_config_description = sm.describe_notebook_instance_lifecycle_config(
|
||||
NotebookInstanceLifecycleConfigName=initial_config_name,
|
||||
resource_description = getattr(sm, test_config.describe_function_name)(
|
||||
**{test_config.name_parameter: initial_config_name}
|
||||
)
|
||||
len(notebook_lifecycle_config_description["OnCreate"]).should.equal(1)
|
||||
len(resource_description["OnCreate"]).should.equal(1)
|
||||
initial_on_create_script.should.equal(
|
||||
notebook_lifecycle_config_description["OnCreate"][0]["Content"]
|
||||
resource_description["OnCreate"][0]["Content"]
|
||||
)
|
||||
|
||||
# Update stack with new instance type and check attributes
|
||||
|
|
@ -283,10 +204,63 @@ def test_sagemaker_cloudformation_notebook_instance_lifecycle_config_update():
|
|||
updated_config_name = outputs["Name"]
|
||||
updated_config_name.should.equal(initial_config_name)
|
||||
|
||||
notebook_lifecycle_config_description = sm.describe_notebook_instance_lifecycle_config(
|
||||
NotebookInstanceLifecycleConfigName=updated_config_name,
|
||||
resource_description = getattr(sm, test_config.describe_function_name)(
|
||||
**{test_config.name_parameter: updated_config_name}
|
||||
)
|
||||
len(notebook_lifecycle_config_description["OnCreate"]).should.equal(1)
|
||||
len(resource_description["OnCreate"]).should.equal(1)
|
||||
updated_on_create_script.should.equal(
|
||||
notebook_lifecycle_config_description["OnCreate"][0]["Content"]
|
||||
resource_description["OnCreate"][0]["Content"]
|
||||
)
|
||||
|
||||
|
||||
@mock_cloudformation
|
||||
@mock_sagemaker
|
||||
def test_sagemaker_cloudformation_model_update():
|
||||
cf = boto3.client("cloudformation", region_name="us-east-1")
|
||||
sm = boto3.client("sagemaker", region_name="us-east-1")
|
||||
|
||||
test_config = ModelTestConfig()
|
||||
|
||||
# Set up template for stack with initial and update instance types
|
||||
stack_name = "{}_stack".format(test_config.resource_name)
|
||||
image = "404615174143.dkr.ecr.us-east-2.amazonaws.com/kmeans:{}"
|
||||
initial_image_version = 1
|
||||
updated_image_version = 2
|
||||
initial_template_json = test_config.get_cloudformation_template(
|
||||
image=image.format(initial_image_version)
|
||||
)
|
||||
updated_template_json = test_config.get_cloudformation_template(
|
||||
image=image.format(updated_image_version)
|
||||
)
|
||||
|
||||
# Create stack with initial template and check attributes
|
||||
cf.create_stack(StackName=stack_name, TemplateBody=initial_template_json)
|
||||
stack_description = cf.describe_stacks(StackName=stack_name)["Stacks"][0]
|
||||
outputs = {
|
||||
output["OutputKey"]: output["OutputValue"]
|
||||
for output in stack_description["Outputs"]
|
||||
}
|
||||
inital_model_name = outputs["Name"]
|
||||
resource_description = getattr(sm, test_config.describe_function_name)(
|
||||
**{test_config.name_parameter: inital_model_name}
|
||||
)
|
||||
resource_description["PrimaryContainer"]["Image"].should.equal(
|
||||
image.format(initial_image_version)
|
||||
)
|
||||
|
||||
# Update stack with new instance type and check attributes
|
||||
cf.update_stack(StackName=stack_name, TemplateBody=updated_template_json)
|
||||
stack_description = cf.describe_stacks(StackName=stack_name)["Stacks"][0]
|
||||
outputs = {
|
||||
output["OutputKey"]: output["OutputValue"]
|
||||
for output in stack_description["Outputs"]
|
||||
}
|
||||
updated_notebook_name = outputs["Name"]
|
||||
updated_notebook_name.should_not.equal(inital_model_name)
|
||||
|
||||
resource_description = getattr(sm, test_config.describe_function_name)(
|
||||
**{test_config.name_parameter: updated_notebook_name}
|
||||
)
|
||||
resource_description["PrimaryContainer"]["Image"].should.equal(
|
||||
image.format(updated_image_version)
|
||||
)
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue