change: add KMS key option for Endpoint Configs (#762)

Author: laurenyu

src/sagemaker/model.py

@@ -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``.
@@ -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 ARN of the KMS key that is used to encrypt the data on the
+                storage volume attached to the instance hosting the endpoint.
 
         Returns:
             callable[string, sagemaker.session.Session] or None: Invocation of ``self.predictor_cls`` on
@@ -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)

src/sagemaker/session.py

@@ -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
@@ -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):
@@ -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:
@@ -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)

tests/integ/kms_utils.py

@@ -101,13 +101,16 @@ def _add_role_to_policy(kms_client,
                                   Policy=KEY_POLICY.format(id=POLICY_NAME, principal=principal))
 
 
-def get_or_create_kms_key(kms_client,
-                          account_id,
+def get_or_create_kms_key(sagemaker_session,
                           role_arn=None,
                           alias=KEY_ALIAS,
                           sagemaker_role='SageMakerRole'):
+    kms_client = sagemaker_session.boto_session.client('kms')
     kms_key_arn = _get_kms_key_arn(kms_client, alias)
 
+    sts_client = sagemaker_session.boto_session.client('sts')
+    account_id = sts_client.get_caller_identity()['Account']
+
     if kms_key_arn is None:
         return _create_kms_key(kms_client, account_id, role_arn, sagemaker_role, alias)
 

tests/integ/test_mxnet_train.py

@@ -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
 
 
@@ -78,7 +79,7 @@ def test_deploy_model(mxnet_training_job, sagemaker_session, mxnet_full_version)
         assert 'Could not find model' in str(exception.value)
 
 
-def test_deploy_model_with_tags(mxnet_training_job, sagemaker_session, mxnet_full_version):
+def test_deploy_model_with_tags_and_kms(mxnet_training_job, sagemaker_session, mxnet_full_version):
     endpoint_name = 'test-mxnet-deploy-model-{}'.format(sagemaker_timestamp())
 
     with timeout_and_delete_endpoint_by_name(endpoint_name, sagemaker_session):
@@ -88,8 +89,11 @@ def test_deploy_model_with_tags(mxnet_training_job, sagemaker_session, mxnet_ful
         model = MXNetModel(model_data, 'SageMakerRole', entry_point=script_path,
                            py_version=PYTHON_VERSION, sagemaker_session=sagemaker_session,
                            framework_version=mxnet_full_version)
+
         tags = [{'Key': 'TagtestKey', 'Value': 'TagtestValue'}]
-        model.deploy(1, 'ml.m4.xlarge', endpoint_name=endpoint_name, tags=tags)
+        kms_key_arn = get_or_create_kms_key(sagemaker_session)
+
+        model.deploy(1, 'ml.m4.xlarge', endpoint_name=endpoint_name, tags=tags, kms_key=kms_key_arn)
 
         returned_model = sagemaker_session.describe_model(EndpointName=model.name)
         returned_model_tags = sagemaker_session.list_tags(ResourceArn=returned_model['ModelArn'])['Tags']
@@ -107,6 +111,7 @@ def test_deploy_model_with_tags(mxnet_training_job, sagemaker_session, mxnet_ful
         assert endpoint_tags == tags
         assert production_variants[0]['InstanceType'] == 'ml.m4.xlarge'
         assert production_variants[0]['InitialInstanceCount'] == 1
+        assert endpoint_config['KmsKeyId'] == kms_key_arn
 
 
 def test_deploy_model_with_update_endpoint(mxnet_training_job, sagemaker_session, mxnet_full_version):

tests/integ/test_transformer.py

@@ -50,10 +50,7 @@ def test_transform_mxnet(sagemaker_session, mxnet_full_version):
     transform_input = mx.sagemaker_session.upload_data(path=transform_input_path,
                                                        key_prefix=transform_input_key_prefix)
 
-    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)
+    kms_key_arn = get_or_create_kms_key(sagemaker_session)
 
     transformer = _create_transformer_and_transform_job(mx, transform_input, kms_key_arn)
     with timeout_and_delete_model_with_transformer(transformer, sagemaker_session,

tests/unit/test_estimator.py

@@ -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,

tests/unit/test_model.py

@@ -165,6 +165,7 @@ def test_deploy(sagemaker_session, tmpdir):
           'InstanceType': INSTANCE_TYPE,
           'InitialInstanceCount': 1,
           'VariantName': 'AllTraffic'}],
+        None,
         None)
 
 
@@ -180,6 +181,7 @@ def test_deploy_endpoint_name(sagemaker_session, tmpdir):
           'InstanceType': INSTANCE_TYPE,
           'InitialInstanceCount': 55,
           'VariantName': 'AllTraffic'}],
+        None,
         None)
 
 
@@ -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())
@@ -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())
@@ -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,