Skip to content

change: add KMS key option for Endpoint Configs #762

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 4 commits into from
Apr 23, 2019
Merged

change: add KMS key option for Endpoint Configs #762

Show file tree
Hide file tree
Changes from 2 commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
79f9376
change: add KMS key option for Endpoint Configs
laurenyu Apr 23, 2019
6b882a2
fix unit test
laurenyu Apr 23, 2019
472d7e9
address PR comments
laurenyu Apr 23, 2019
1eaedac
Merge branch 'master' into endpoint-config-kms
laurenyu Apr 23, 2019
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
10 changes: 7 additions & 3 deletions src/sagemaker/model.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -209,7 +209,7 @@ def compile(self, target_instance_family, input_shape, output_path, role,
return self

def deploy(self, initial_instance_count, instance_type, accelerator_type=None, endpoint_name=None,
update_endpoint=False, tags=None):
update_endpoint=False, tags=None, kms_key=None):
"""Deploy this ``Model`` to an ``Endpoint`` and optionally return a ``Predictor``.

Create a SageMaker ``Model`` and ``EndpointConfig``, and deploy an ``Endpoint`` from this ``Model``.
Expand All @@ -235,6 +235,8 @@ def deploy(self, initial_instance_count, instance_type, accelerator_type=None, e
If True, this will deploy a new EndpointConfig to an already existing endpoint and delete resources
corresponding to the previous EndpointConfig. If False, a new endpoint will be created. Default: False
tags(List[dict[str, str]]): The list of tags to attach to this specific endpoint.
kms_key (str): The KMS key that is used to encrypt the data on the storage volume attached
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Let's make it clear that this is the arn of the kms_key

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

done

to the instance hosting the endpoint.

Returns:
callable[string, sagemaker.session.Session] or None: Invocation of ``self.predictor_cls`` on
Expand Down Expand Up @@ -270,10 +272,12 @@ def deploy(self, initial_instance_count, instance_type, accelerator_type=None, e
initial_instance_count=initial_instance_count,
instance_type=instance_type,
accelerator_type=accelerator_type,
tags=tags)
tags=tags,
kms_key=kms_key)
self.sagemaker_session.update_endpoint(self.endpoint_name, endpoint_config_name)
else:
self.sagemaker_session.endpoint_from_production_variants(self.endpoint_name, [production_variant], tags)
self.sagemaker_session.endpoint_from_production_variants(self.endpoint_name, [production_variant],
tags, kms_key)

if self.predictor_cls:
return self.predictor_cls(self.endpoint_name, self.sagemaker_session)
Expand Down
29 changes: 21 additions & 8 deletions src/sagemaker/session.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -710,7 +710,7 @@ def wait_for_model_package(self, model_package_name, poll=5):
return desc

def create_endpoint_config(self, name, model_name, initial_instance_count, instance_type,
accelerator_type=None, tags=None):
accelerator_type=None, tags=None, kms_key=None):
"""Create an Amazon SageMaker endpoint configuration.

The endpoint configuration identifies the Amazon SageMaker model (created using the
Expand Down Expand Up @@ -738,12 +738,21 @@ def create_endpoint_config(self, name, model_name, initial_instance_count, insta

tags = tags or []

self.sagemaker_client.create_endpoint_config(
EndpointConfigName=name,
ProductionVariants=[production_variant(model_name, instance_type, initial_instance_count,
accelerator_type=accelerator_type)],
Tags=tags
)
request = {
'EndpointConfigName': name,
'ProductionVariants': [
production_variant(model_name, instance_type, initial_instance_count,
accelerator_type=accelerator_type)
],
}

if tags is not None:
request['Tags'] = tags

if kms_key is not None:
request['KmsKeyId'] = kms_key

self.sagemaker_client.create_endpoint_config(**request)
return name

def create_endpoint(self, endpoint_name, config_name, tags=None, wait=True):
Expand Down Expand Up @@ -1032,13 +1041,15 @@ def endpoint_from_model_data(self, model_s3_location, deployment_image, initial_
self.create_endpoint(endpoint_name=name, config_name=name, wait=wait)
return name

def endpoint_from_production_variants(self, name, production_variants, tags=None, wait=True):
def endpoint_from_production_variants(self, name, production_variants, tags=None, kms_key=None, wait=True):
"""Create an SageMaker ``Endpoint`` from a list of production variants.

Args:
name (str): The name of the ``Endpoint`` to create.
production_variants (list[dict[str, str]]): The list of production variants to deploy.
tags (list[dict[str, str]]): A list of key-value pairs for tagging the endpoint (default: None).
kms_key (str): The KMS key that is used to encrypt the data on the storage volume attached
to the instance hosting the endpoint.
wait (bool): Whether to wait for the endpoint deployment to complete before returning (default: True).

Returns:
Expand All @@ -1050,6 +1061,8 @@ def endpoint_from_production_variants(self, name, production_variants, tags=None
config_options = {'EndpointConfigName': name, 'ProductionVariants': production_variants}
if tags:
config_options['Tags'] = tags
if kms_key:
config_options['KmsKeyId'] = kms_key

self.sagemaker_client.create_endpoint_config(**config_options)
return self.create_endpoint(endpoint_name=name, config_name=name, tags=tags, wait=wait)
Expand Down
25 changes: 25 additions & 0 deletions tests/integ/test_mxnet_train.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -23,6 +23,7 @@
from sagemaker.mxnet.model import MXNetModel
from sagemaker.utils import sagemaker_timestamp
from tests.integ import DATA_DIR, PYTHON_VERSION, TRAINING_DEFAULT_TIMEOUT_MINUTES
from tests.integ.kms_utils import get_or_create_kms_key
from tests.integ.timeout import timeout, timeout_and_delete_endpoint_by_name


Expand Down Expand Up @@ -109,6 +110,30 @@ def test_deploy_model_with_tags(mxnet_training_job, sagemaker_session, mxnet_ful
assert production_variants[0]['InitialInstanceCount'] == 1


def test_deploy_model_with_kms_key(mxnet_training_job, sagemaker_session, mxnet_full_version):
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Should we just add this to an existing test? They are already taking so long to run use so much resource.

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

done

endpoint_name = 'test-mxnet-deploy-model-{}'.format(sagemaker_timestamp())

with timeout_and_delete_endpoint_by_name(endpoint_name, sagemaker_session):
desc = sagemaker_session.sagemaker_client.describe_training_job(TrainingJobName=mxnet_training_job)
model_data = desc['ModelArtifacts']['S3ModelArtifacts']
script_path = os.path.join(DATA_DIR, 'mxnet_mnist', 'mnist.py')
model = MXNetModel(model_data, 'SageMakerRole', entry_point=script_path,
py_version=PYTHON_VERSION, sagemaker_session=sagemaker_session,
framework_version=mxnet_full_version)

sts_client = sagemaker_session.boto_session.client('sts')
account_id = sts_client.get_caller_identity()['Account']
kms_client = sagemaker_session.boto_session.client('kms')
kms_key_arn = get_or_create_kms_key(kms_client, account_id)
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Optional - I think it will make this code cleaner if get_or_create_kms_key just take sageamker_session as the argument.

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

done


model.deploy(1, 'ml.m4.xlarge', endpoint_name=endpoint_name, kms_key=kms_key_arn)

endpoint = sagemaker_session.describe_endpoint(EndpointName=endpoint_name)
endpoint_config = sagemaker_session.describe_endpoint_config(EndpointConfigName=endpoint['EndpointConfigName'])

assert endpoint_config['KmsKeyId'] == kms_key_arn


def test_deploy_model_with_update_endpoint(mxnet_training_job, sagemaker_session, mxnet_full_version):
endpoint_name = 'test-mxnet-deploy-model-{}'.format(sagemaker_timestamp())

Expand Down
3 changes: 2 additions & 1 deletion tests/unit/test_estimator.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -902,7 +902,8 @@ def test_fit_deploy_keep_tags(sagemaker_session):
job_name = estimator._current_job_name
sagemaker_session.endpoint_from_production_variants.assert_called_with(job_name,
variant,
tags)
tags,
None)

sagemaker_session.create_model.assert_called_with(
ANY,
Expand Down
27 changes: 25 additions & 2 deletions tests/unit/test_model.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -165,6 +165,7 @@ def test_deploy(sagemaker_session, tmpdir):
'InstanceType': INSTANCE_TYPE,
'InitialInstanceCount': 1,
'VariantName': 'AllTraffic'}],
None,
None)


Expand All @@ -180,6 +181,7 @@ def test_deploy_endpoint_name(sagemaker_session, tmpdir):
'InstanceType': INSTANCE_TYPE,
'InitialInstanceCount': 55,
'VariantName': 'AllTraffic'}],
None,
None)


Expand All @@ -196,7 +198,8 @@ def test_deploy_tags(sagemaker_session, tmpdir):
'InstanceType': INSTANCE_TYPE,
'InitialInstanceCount': 1,
'VariantName': 'AllTraffic'}],
tags)
tags,
None)


@patch('sagemaker.fw_utils.tar_and_upload_dir', MagicMock())
Expand All @@ -213,9 +216,28 @@ def test_deploy_accelerator_type(tfo, time, sagemaker_session):
'InitialInstanceCount': 1,
'VariantName': 'AllTraffic',
'AcceleratorType': ACCELERATOR_TYPE}],
None,
None)


@patch('sagemaker.fw_utils.tar_and_upload_dir', MagicMock())
@patch('tarfile.open')
@patch('time.strftime', return_value=TIMESTAMP)
def test_deploy_kms_key(tfo, time, sagemaker_session):
key = 'some-key-arn'
model = DummyFrameworkModel(sagemaker_session)
model.deploy(instance_type=INSTANCE_TYPE, initial_instance_count=1, kms_key=key)
sagemaker_session.endpoint_from_production_variants.assert_called_with(
MODEL_NAME,
[{'InitialVariantWeight': 1,
'ModelName': MODEL_NAME,
'InstanceType': INSTANCE_TYPE,
'InitialInstanceCount': 1,
'VariantName': 'AllTraffic'}],
None,
key)


@patch('sagemaker.session.Session')
@patch('sagemaker.local.LocalSession')
@patch('sagemaker.fw_utils.tar_and_upload_dir', MagicMock())
Expand Down Expand Up @@ -246,7 +268,8 @@ def test_deploy_update_endpoint(sagemaker_session, tmpdir):
initial_instance_count=INSTANCE_COUNT,
instance_type=INSTANCE_TYPE,
accelerator_type=ACCELERATOR_TYPE,
tags=None
tags=None,
kms_key=None,
)
config_name = sagemaker_session.create_endpoint_config(
name=model.name,
Expand Down