catalin/moto
1
0
Fork 0
Code Issues Pull requests Projects Releases Wiki Activity

Add CloudFormation support for SageMaker Notebook Instances (#3845)

Browse source 
* Create SageMaker Notebook Instances with CloudFormation

* Implement attributes for SageMaker notebook instance in Cloudformation

* Delete SageMaker Notebook Instances with CloudFormation

* Update SageMaker Notebook Instances with CloudFormation

* Factor out template into function for SageMaker notebook instance tests
This commit is contained in:
Zach Churchill 2021-04-09 13:54:00 -04:00 • committed by GitHub
commit d45233fa00
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
3 changed files with 221 additions and 2 deletions
Download patch file Download diff file Expand all files Collapse all files

155
tests/test_sagemaker/test_sagemaker_cloudformation.py Normal file
View file

@ -0,0 +1,155 @@
import json
import boto3
import pytest
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(
instance_type="ml.c4.xlarge",
role_arn="arn:aws:iam::{}:role/FakeRole".format(ACCOUNT_ID),
include_outputs=True,
):
template = {
"AWSTemplateFormatVersion": "2010-09-09",
"Resources": {
"TestNotebookInstance": {
"Type": "AWS::SageMaker::NotebookInstance",
"Properties": {"InstanceType": instance_type, "RoleArn": role_arn},
},
},
}
if include_outputs:
template["Outputs"] = {
"NotebookInstanceArn": {"Value": {"Ref": "TestNotebookInstance"}},
"NotebookInstanceName": {
"Value": {
"Fn::GetAtt": ["TestNotebookInstance", "NotebookInstanceName"]
},
},
}
return json.dumps(template)
@mock_cloudformation
def test_sagemaker_cloudformation_create_notebook_instance():
cf = boto3.client("cloudformation", region_name="us-east-1")
stack_name = "test_sagemaker_notebook_instance"
template = _get_notebook_instance_template_string(include_outputs=False)
cf.create_stack(StackName=stack_name, TemplateBody=template)
provisioned_resource = cf.list_stack_resources(StackName=stack_name)[
"StackResourceSummaries"
][0]
provisioned_resource["LogicalResourceId"].should.equal("TestNotebookInstance")
len(provisioned_resource["PhysicalResourceId"]).should.be.greater_than(0)
@mock_cloudformation
@mock_sagemaker
def test_sagemaker_cloudformation_notebook_instance_get_attr():
cf = boto3.client("cloudformation", region_name="us-east-1")
sm = boto3.client("sagemaker", region_name="us-east-1")
stack_name = "test_sagemaker_notebook_instance"
template = _get_notebook_instance_template_string()
cf.create_stack(StackName=stack_name, TemplateBody=template)
stack_description = cf.describe_stacks(StackName=stack_name)["Stacks"][0]
outputs = {
output["OutputKey"]: output["OutputValue"]
for output in stack_description["Outputs"]
}
notebook_instance_name = outputs["NotebookInstanceName"]
notebook_instance_arn = outputs["NotebookInstanceArn"]
notebook_instance_description = sm.describe_notebook_instance(
NotebookInstanceName=notebook_instance_name,
)
notebook_instance_arn.should.equal(
notebook_instance_description["NotebookInstanceArn"]
)
@mock_cloudformation
@mock_sagemaker
def test_sagemaker_cloudformation_notebook_instance_delete():
cf = boto3.client("cloudformation", region_name="us-east-1")
sm = boto3.client("sagemaker", region_name="us-east-1")
# Create stack with notebook instance and verify existence
stack_name = "test_sagemaker_notebook_instance"
template = _get_notebook_instance_template_string()
cf.create_stack(StackName=stack_name, TemplateBody=template)
stack_description = cf.describe_stacks(StackName=stack_name)["Stacks"][0]
outputs = {
output["OutputKey"]: output["OutputValue"]
for output in stack_description["Outputs"]
}
notebook_instance = sm.describe_notebook_instance(
NotebookInstanceName=outputs["NotebookInstanceName"],
)
outputs["NotebookInstanceArn"].should.equal(
notebook_instance["NotebookInstanceArn"]
)
# Delete the stack and verify notebook instance has also been deleted
# TODO replace exception check with `list_notebook_instances` method when implemented
cf.delete_stack(StackName=stack_name)
with pytest.raises(ClientError) as ce:
sm.describe_notebook_instance(
NotebookInstanceName=outputs["NotebookInstanceName"]
)
ce.value.response["Error"]["Message"].should.contain("RecordNotFound")
@mock_cloudformation
@mock_sagemaker
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")
# Set up template for stack with initial and update instance types
stack_name = "test_sagemaker_notebook_instance"
initial_instance_type = "ml.c4.xlarge"
updated_instance_type = "ml.c4.4xlarge"
initial_template_json = _get_notebook_instance_template_string(
instance_type=initial_instance_type
)
updated_template_json = _get_notebook_instance_template_string(
instance_type=updated_instance_type
)
# 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"]
}
initial_notebook_name = outputs["NotebookInstanceName"]
notebook_instance_description = sm.describe_notebook_instance(
NotebookInstanceName=initial_notebook_name,
)
initial_instance_type.should.equal(notebook_instance_description["InstanceType"])
# 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["NotebookInstanceName"]
updated_notebook_name.should.equal(initial_notebook_name)
notebook_instance_description = sm.describe_notebook_instance(
NotebookInstanceName=updated_notebook_name,
)
updated_instance_type.should.equal(notebook_instance_description["InstanceType"])